.

CDTools/models/base.py

@@ -11,7 +11,7 @@ simulate_to_dataset
     Creates a CDataset from the simulation defined in the model
 save_results
     Saves out a dictionary with the recovered parameters
-    
+
 
 Simulation
 ----------
@@ -56,7 +56,7 @@ class CDIModel(t.nn.Module):
     def __init__(self):
         super(CDIModel,self).__init__()
         self.iteration_count = 0
-        
+
     def from_dataset(self, dataset):
         raise NotImplementedError()
 
@@ -79,11 +79,11 @@ class CDIModel(t.nn.Module):
 
     def forward(self, *args):
         """The complete forward model
-        
+
         This model relies on composing the interaction, forward propagator,
         and measurement functions which are required to be defined by all
         subclasses. It therefore should not be redefined by the subclasses.
-        
+
         The arguments to this function, for any given subclass, will be
         the same as the arguments to the interaction function.
         """
@@ -99,7 +99,7 @@ class CDIModel(t.nn.Module):
 
     def simulate_to_dataset(self, args_list):
         raise NotImplementedError()
-    
+
     def save_results(self):
         raise NotImplementedError()
 
@@ -107,10 +107,10 @@ class CDIModel(t.nn.Module):
                     scheduler=None, regularization_factor=None, thread=True,
                     calculation_width=10):
         """Runs a round of reconstruction using the provided optimizer
-        
+
         This is the basic automatic differentiation reconstruction tool
         which all the other, algorithm-specific tools, use.
-        
+
         Like all the other optimization routines, it is defined as a
         generator function which yields the average loss each epoch.
 
@@ -135,7 +135,7 @@ class CDIModel(t.nn.Module):
         normalization = 0
         for inputs, patterns in data_loader:
             normalization += t.sum(patterns).cpu().numpy()
-            
+
         def run_iteration(stop_event=None):
             loss = 0
             N = 0
@@ -144,7 +144,7 @@ class CDIModel(t.nn.Module):
                 N += 1
                 def closure():
                     optimizer.zero_grad()
-                    
+
                     input_chunks = [[inp[i:i + calculation_width]
                                      for inp in inputs]
                                     for i in range(0, len(inputs[0]),
@@ -152,7 +152,7 @@ class CDIModel(t.nn.Module):
                     pattern_chunks = [patterns[i:i + calculation_width]
                                       for i in range(0, len(inputs[0]),
                                                      calculation_width)]
-                    
+
                     total_loss = 0
                     for inp, pats in zip(input_chunks, pattern_chunks):
                         # This is just used to allow graceful exit when
@@ -168,7 +168,7 @@ class CDIModel(t.nn.Module):
                             loss = self.loss(pats,sim_patterns)
 
                         loss.backward()
-                    
+
                         total_loss += loss.detach()
 
                     if regularization_factor is not None \
@@ -178,7 +178,7 @@ class CDIModel(t.nn.Module):
                     return total_loss
 
                 loss += optimizer.step(closure).detach().cpu().numpy()
-                
+
             loss /= normalization
             if scheduler is not None:
                 scheduler.step(loss)
@@ -198,7 +198,7 @@ class CDIModel(t.nn.Module):
                     # If something bad happens, put the exception into the
                     # result queue
                     result_queue.put(e)
-                    
+
         for it in range(iterations):
             if thread:
                 calc = threading.Thread(target=target, name='calculator', daemon=True)
@@ -206,10 +206,10 @@ class CDIModel(t.nn.Module):
                     calc.start()
                     while calc.is_alive():
                         if hasattr(self, 'figs'):
-                            self.figs[0].canvas.start_event_loop(0.01) 
+                            self.figs[0].canvas.start_event_loop(0.01)
                         else:
                             calc.join()
-                            
+
                 except KeyboardInterrupt as e:
                     stop_event.set()
                     print('\nAsking execution thread to stop cleanly - please be patient.')
@@ -233,7 +233,7 @@ class CDIModel(t.nn.Module):
                       regularization_factor=None, thread=True,
                       calculation_width=10):
         """Runs a round of reconstruction using the Adam optimizer
-        
+
         This is generally accepted to be the most robust algorithm for use
         with ptychography. Like all the other optimization routines,
         it is defined as a generator function, which yields the average
@@ -267,7 +267,7 @@ class CDIModel(t.nn.Module):
             if type(subset) == type(1):
                 subset = [subset]
             dataset = torchdata.Subset(dataset, subset)
-            
+
         # Make a dataloader
         data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
                                            shuffle=True)
@@ -290,17 +290,17 @@ class CDIModel(t.nn.Module):
                                 calculation_width=calculation_width)
 
 
-    def LBFGS_optimize(self, iterations, dataset, 
+    def LBFGS_optimize(self, iterations, dataset,
                        lr=0.1,history_size=2, subset=None,
                        regularization_factor=None, thread=True,
                        calculation_width=10):
         """Runs a round of reconstruction using the L-BFGS optimizer
-        
+
         This algorithm is often less stable that Adam, however in certain
         situations or geometries it can be shockingly efficient. Like all
         the other optimization routines, it is defined as a generator
         function which yields the average loss each epoch.
-        
+
         Note: There is no batch size, because it is a usually a bad idea to use
         LBFGS on anything but all the data at onece
 
@@ -320,14 +320,14 @@ class CDIModel(t.nn.Module):
             Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
         thread : bool
             Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
-        
+
         """
         if subset is not None:
             # if just one pattern, turn into a list for convenience
             if type(subset) == type(1):
                 subset = [subset]
             dataset = torchdata.Subset(dataset, subset)
-        
+
         # Make a dataloader. This basically does nothing but load all the
         # data at once
         data_loader = torchdata.DataLoader(dataset, batch_size=len(dataset))
@@ -338,7 +338,7 @@ class CDIModel(t.nn.Module):
                                   lr = lr, history_size=history_size)
         #optimizer = MyLBFGS(self.parameters(),
         #                    lr = lr, history_size=history_size)
-        
+
         return self.AD_optimize(iterations, data_loader, optimizer,
                                 regularization_factor=regularization_factor,
                                 thread=thread,
@@ -350,7 +350,7 @@ class CDIModel(t.nn.Module):
                      nesterov=False, subset=None, regularization_factor=None,
                      thread=True, calculation_width=10):
         """Runs a round of reconstruction using the SGDoptimizer
-        
+
         This algorithm is often less stable that Adam, but it is simpler
         and is the basic workhorse of gradience descent.
 
@@ -382,7 +382,7 @@ class CDIModel(t.nn.Module):
             if type(subset) == type(1):
                 subset = [subset]
             dataset = torchdata.Subset(dataset, subset)
-            
+
         # Make a dataloader
         if batch_size is not None:
             data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
@@ -418,28 +418,28 @@ class CDIModel(t.nn.Module):
                   self.latest_iteration_time + str(self.latest_loss)
         else:
             return 'No reconstruction iterations performed yet!'
-        
+
     # By default, the plot_list is empty
     plot_list = []
-    
-    
+
+
     def inspect(self, dataset=None, update=True):
         """Plots all the plots defined in the model's plot_list attribute
-  
+
         If update is set to True, it will update any previously plotted set
         of plots, if one exists, and then redraw them. Otherwise, it will
         plot a new set, and any subsequent updates will update the new set
-        
+
         Optionally, a dataset can be passed, which then will plot any
         registered plots which need to incorporate some information from
         the dataset (such as geometry or a comparison with measured data).
-        
+
         Plots can be registered in any subclass by defining the plot_list
         attribute. This should be a list of tuples in the following format:
-        ( 'Plot Title', function_to_generate_plot(self), 
+        ( 'Plot Title', function_to_generate_plot(self),
         function_to_determine_whether_to_plot(self))
 
-        Where the third element in the tuple (a function that returns 
+        Where the third element in the tuple (a function that returns
         True if the plot is relevant) is not required.
 
         Parameters
@@ -448,9 +448,9 @@ class CDIModel(t.nn.Module):
             Optional, a dataset matched to the model type
         update : bool
             Default True, whether to update existing plots or plot new ones
-        
+
         """
-                
+
         first_update = False
         if update and hasattr(self, 'figs') and self.figs:
             figs = self.figs
@@ -476,7 +476,7 @@ class CDIModel(t.nn.Module):
             plotter = plots[1]
 
             if figs is None:
-                fig = plt.figure()                
+                fig = plt.figure()
                 self.figs.append(fig)
             else:
                 fig = figs[idx]
@@ -484,13 +484,13 @@ class CDIModel(t.nn.Module):
             try:
                 plotter(self,fig)
                 plt.title(name)
-                            
+
             except TypeError as e:
                 if dataset is not None:
                     try:
                         plotter(self, fig, dataset)
                         plt.title(name)
-                        
+
                     except (IndexError, KeyError, AttributeError, np.linalg.LinAlgError) as e:
                         pass
 
@@ -498,15 +498,15 @@ class CDIModel(t.nn.Module):
                 pass
 
             idx += 1
-            
+
             if update:
                 plt.draw()
                 fig.canvas.start_event_loop(0.001)
-                
+
         if first_update:
             plt.pause(0.05 * len(self.figs))
 
-            
+
     def save_figures(self, prefix='', extension='.eps'):
         """Saves all currently open inspection figures.
 
@@ -520,7 +520,7 @@ class CDIModel(t.nn.Module):
         By default, the files will be named by the figure titles as defined
         in the plot_list. Files can be saved with any extension suported by
         matplotlib.pyplot.savefig.
-        
+
         Parameters
         ----------
         prefix : str
@@ -528,7 +528,7 @@ class CDIModel(t.nn.Module):
         extention : strategy
             Default is .eps, the file extension to save with.
         """
-        
+
         if hasattr(self, 'figs') and self.figs:
             figs = self.figs
         else:
@@ -538,21 +538,21 @@ class CDIModel(t.nn.Module):
             fig.savefig(prefix + fig.axes[0].get_title() + extension,
                         bbox_inches = 'tight')
 
-            
-    def compare(self, dataset):
+
+    def compare(self, dataset, logarithmic=False):
         """Opens a tool for comparing simulated and measured diffraction patterns
-        
+
         Parameters
         ----------
         dataset : CDataset
             A dataset containing the simulated diffraction patterns to compare against
         """
-        
+
         fig, axes = plt.subplots(1,3,figsize=(12,5.3))
         fig.tight_layout(rect=[0.02, 0.09, 0.98, 0.96])
         axslider = plt.axes([0.15,0.06,0.75,0.03])
-        
-        
+
+
         def update_colorbar(im):
             # If the update brought the colorbar out of whack
             # (say, from clicking back in the navbar)
@@ -564,7 +564,7 @@ class CDIModel(t.nn.Module):
             if hasattr(im, 'norecurse') and im.norecurse:
                 im.norecurse=False
                 return
-            
+
             im.norecurse=True
             im.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array()))
 
@@ -572,7 +572,7 @@ class CDIModel(t.nn.Module):
             idx = int(idx) % len(dataset)
             fig.pattern_idx = idx
             updating = True if len(axes[0].images) >= 1 else False
-            
+
             inputs, output = dataset[idx]
             sim_data = self.forward(*inputs).detach().cpu().numpy()
             sim_data = sim_data
@@ -581,6 +581,10 @@ class CDIModel(t.nn.Module):
                 mask = self.mask.detach().cpu().numpy()
             else:
                 mask = 1
+
+            if logarithmic:
+                sim_data =np.log(sim_data)/np.log(10)
+                meas_data = np.log(meas_data)/np.log(10)
                 
             if not updating:
                 axes[0].set_title('Simulated')
@@ -600,7 +604,7 @@ class CDIModel(t.nn.Module):
                 cb3 = plt.colorbar(diff, ax=axes[2], orientation='horizontal',format='%.2e',ticks=ticker.LinearLocator(numticks=5),pad=0.1,fraction=0.1)
                 cb3.ax.tick_params(labelrotation=20)
                 cb3.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(diff))
-                
+
             else:
 
                 sim = axes[0].images[-1]
@@ -614,8 +618,8 @@ class CDIModel(t.nn.Module):
                 diff = axes[2].images[-1]
                 diff.set_data((sim_data-meas_data) * mask)
                 update_colorbar(diff)
-                
-                
+
+
         # This is dumb but the slider doesn't work unless a reference to it is
         # kept somewhere...
         self.slider = Slider(axslider, 'Pattern #', 0, len(dataset)-1, valstep=1, valfmt="%d")
@@ -626,7 +630,7 @@ class CDIModel(t.nn.Module):
                 event.button = None
             if not hasattr(event, 'key'):
                 event.key = None
-                
+
             if event.key == 'up' or event.button == 'up':
                 update(fig.pattern_idx - 1)
             elif event.key == 'down' or event.button == 'down':
@@ -637,5 +641,3 @@ class CDIModel(t.nn.Module):
         fig.canvas.mpl_connect('key_press_event',on_action)
         fig.canvas.mpl_connect('scroll_event',on_action)
         update(0)
-        
-

CDTools/models/polarized_fancy_ptycho.py

@@ -21,13 +21,13 @@ class PolarizedFancyPtycho(FancyPtycho):
                  detector_slice=None,
                  surface_normal=np.array([0.,0.,1.]),
                  min_translation = t.Tensor([0,0]),
-                 background = None, translation_offsets=None, 
+                 background = None, translation_offsets=None,
                  polarizer_offsets=None, analyzer_offsets=None,
                  polarizer_scale=1, analyzer_scale=1,  mask=None,
                  weights = None, translation_scale = 1, saturation=None,
                  probe_support = None, obj_support=None, oversampling=1,
                  loss='amplitude mse',units='um'):
-        
+
         super(FancyPtycho, self).__init__(wavelength, detector_geometry,
                  probe_basis,
                  probe_guess, obj_guess,
@@ -48,29 +48,55 @@ class PolarizedFancyPtycho(FancyPtycho):
             self.analyzer_offsets = None
         else:
             self.analyzer_offsets = t.nn.Parameter(t.tensor(analyzer_offsets).to(dtype=t.float32)) / analyzer_scale
-        
+
+        self.polarizer = polarizer
+        self.analyzer = analyzer
+        probe_guess = t.tensor(probe_guess, dtype=t.complex64)
+        if probe_guess.dim() > 4:
+            self.probe_norm = 1 * t.max(t.abs(probe_guess[0]))
+        else:
+            self.probe_norm = 1 * t.max(t.abs(probe_guess))
+
+        self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
+
     @classmethod
     def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', left_polarized=True):
-        
+
         model = FancyPtycho.from_dataset(dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um')
 
 
-        # Mutate the class to its subclass 
+        # Mutate the class to its subclass
         model.__class__ = cls
 
         if left_polarized:
             x = 1j
         else:
             x = -1j
-        model.probe.data = t.stack((model.probe.data.to(dtype=t.cfloat), x * model.probe.data.to(dtype=t.cfloat)), dim=-3)
-        obj = t.stack((model.obj.data, model.obj.data), dim=-3)
-        model.obj.data = t.stack((obj, obj), dim=-4)        
 
+        # if probe_size is None:
+        #     probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
+        # else:
+        #     probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
+        probe = model.probe.detach()
+
+        probe = t.stack((probe, probe * x), dim=-3)
+        probe_max = t.max(t.abs(probe))
+        probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)]
+        probe = t.stack([probe, ] + probe_stack)
+
+        model.probe.data = probe
+        # obj = t.stack((model.obj.data, model.obj.data), dim=-3)
+        # model.obj.data = t.stack((obj.data, obj.data), dim=-4)
+        # obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
+        obj = model.obj.detach()
+        obj = t.stack((obj, obj), dim=-3)
+        obj = t.stack((obj, obj), dim=-4)
+        model.obj.data = obj
         # tensor vs tensor.data
         return model
 
     polarizers = [tools.polarization.generate_linear_polarizer(i * 45) for i in range(3)]
-    
+
     # WHAT IS INDEX?
     def interaction(self, index, translations, polarizer, analyzer, test=False):
 
@@ -84,9 +110,9 @@ class PolarizedFancyPtycho(FancyPtycho):
         if self.translation_offsets is not None:
             pix_trans += self.translation_scale * self.translation_offsets[index]
 
-             
+
         # This restricts the basis probes to stay within the probe support
-        basis_prs = self.probe * self.probe_support[...,:,:] # This makes no sense 
+        basis_prs = self.probe * self.probe_support[...,:,:] # This makes no sense
         # self.probe is an Nx2xXxY stach of probes
 
         # Now we construct the probes for each shot from the basis probes
@@ -106,7 +132,7 @@ class PolarizedFancyPtycho(FancyPtycho):
         analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer)
 
         return analyzed_exit_waves
-    
+
 
     def vectorial_wavefields(wavefields, func, *args, **kwargs):
         wavefields_x = wavefields[..., 0, :, :, :]
@@ -122,7 +148,7 @@ class PolarizedFancyPtycho(FancyPtycho):
     def backward_propagator(self, wavefields):
         return tools.propagators.inverse_far_field(wavefields)
 
-    
+
     def measurement(self, wavefields):
         wavefields_x = wavefields[..., 0, :, :]
         wavefields_y = wavefields[..., 1, :, :]
@@ -145,15 +171,15 @@ class PolarizedFancyPtycho(FancyPtycho):
 
     # Note: No "loss" function is defined here, because it is added
     # dynamically during object creation in __init__
-    
+
     def to(self, *args, **kwargs):
         super(PolarizedFancyPtycho, self).to(*args, **kwargs)
 
-        
+
     def sim_to_dataset(self, args_list):
         # In the future, potentially add more control
         # over what metadata is saved (names, etc.)
-        
+
         # First, I need to gather all the relevant data
         # that needs to be added to the dataset
         entry_info = {'program_name': 'CDTools',
@@ -166,16 +192,16 @@ class PolarizedFancyPtycho(FancyPtycho):
         ysurfacevec = np.cross(surface_normal, xsurfacevec)
         ysurfacevec /= np.linalg.norm(ysurfacevec)
         orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
-        
+
         sample_info = {'description': 'A simulated sample',
                        'orientation': orientation}
 
-        
+
         detector_geometry = self.detector_geometry
         mask = self.mask
         wavelength = self.wavelength
         indices, translations = args_list
-        
+
         # Then we simulate the results
         data = self.forward(indices, translations)
 
@@ -187,14 +213,14 @@ class PolarizedFancyPtycho(FancyPtycho):
                                  detector_geometry=detector_geometry,
                                  mask=mask)
 
-    
+
     def corrected_translations(self, dataset):
         translations = dataset.translations.to(dtype=t.float32,device=self.probe.device)
         t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
         return translations + t_offset
 
 
-    
+
     def get_rhos(self):
         # If this is the general unified mode model
         if self.weights.dim() >= 2:
@@ -205,18 +231,18 @@ class PolarizedFancyPtycho(FancyPtycho):
         else:
             return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
                             dtype=np.complex64)
-        
+
     def tidy_probes(self, normalization=1, normalize=False):
         """Tidies up the probes
-        
+
         What we want to do here is use all the information on all the probes
         to calculate a natural basis for the experiment, and update all the
         density matrices to operate in that updated basis
-        
+
         """
-        
+
         # First we treat the purely incoherent case
-        
+
         # I don't love this pattern of using an if statement with a return
         # to catch this case, but because it's so much simpler than the
         # unified mode case I think it's appropriate
@@ -228,11 +254,11 @@ class PolarizedFancyPtycho(FancyPtycho):
             return
 
         # This is for the unified mode case
-        
+
         # Note to future: We could probably do this more cleanly with an
         # SVD directly on the Ws matrix, instead of an eigendecomposition
         # of the rho matrix.
-        
+
         rhos = self.get_rhos()
         overall_rho = np.mean(rhos,axis=0)
         probe = self.probe.detach().cpu().numpy()
@@ -248,7 +274,7 @@ class PolarizedFancyPtycho(FancyPtycho):
         ortho_probes *= np.sqrt(normalization)
 
         dm_rank = self.weights.shape[1]
-        
+
         new_Ws = []
         for rho in new_rhos:
             # These are returned from smallest to largest - we want to keep
@@ -263,14 +289,14 @@ class PolarizedFancyPtycho(FancyPtycho):
             # when there are thousands of individual matrices to transform
             # every time this is called.
             w = np.maximum(w,0)
-            
+
             new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
-            
+
         new_Ws = np.array(new_Ws)
 
         self.weights.data = t.as_tensor(new_Ws,
             dtype=self.weights.dtype,device=self.weights.device)
-        
+
         self.probe.data = t.as_tensor(ortho_probes,
             device=self.probe.device,dtype=self.probe.dtype)
 
@@ -287,32 +313,32 @@ class PolarizedFancyPtycho(FancyPtycho):
                 return np.sum(np.abs(ortho_probes.detach().cpu().numpy())**2,axis=0)
             if mode.lower() == 'phase':
                 return np.angle(ortho_probes.detach().cpu().numpy())
-            
+
         probe_matrix = np.zeros([self.probe.shape[0]]*2,
                                 dtype=np.complex64)
         np_probes = self.probe.detach().cpu().numpy()
         for i in range(probe_matrix.shape[0]):
             for j in range(probe_matrix.shape[0]):
                 probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj())
-        
+
 
         weights = self.weights.detach().cpu().numpy()
-        
+
         probe_intensities = np.sum(np.tensordot(weights,probe_matrix,axes=1)*
                                    weights.conj(),axis=2)
 
         # Imaginary part is already essentially zero up to rounding error
         probe_intensities = np.real(probe_intensities)
-        
+
         values = np.sum(probe_intensities,axis=1)
         if mode.lower() == 'amplitude' or mode.lower() == 'root_sum_intensity':
             cmap = 'viridis'
         else:
             cmap = 'twilight'
-            
+
         p.plot_nanomap_with_images(self.corrected_translations(dataset), get_probes, values=values, fig=fig, units=self.units, basis=self.probe_basis, nanomap_colorbar_title='Total Probe Intensity',cmap=cmap,**kwargs),
 
-        
+
     plot_list = [
         ('',
          lambda self, fig, dataset: self.plot_wavefront_variation(dataset,fig=fig,mode='root_sum_intensity',image_title='Root Summed Probe Intensities',image_colorbar_title='Square Root of Intensity'),
@@ -333,7 +359,7 @@ class PolarizedFancyPtycho(FancyPtycho):
         ('% Power in Top Mode (only accurate after tidy_probes)',
          lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig,units=self.units),
          lambda self: len(self.weights.shape) >=2),
-        ('Object Amplitude', 
+        ('Object Amplitude',
          lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis,units=self.units)),
         ('Object Phase',
          lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis,units=self.units)),
@@ -343,7 +369,7 @@ class PolarizedFancyPtycho(FancyPtycho):
          lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
     ]
 
-    
+
     def save_results(self, dataset):
         basis = self.probe_basis.detach().cpu().numpy()
         translations = self.corrected_translations(dataset).detach().cpu().numpy()
@@ -352,7 +378,7 @@ class PolarizedFancyPtycho(FancyPtycho):
         obj = self.obj.detach().cpu().numpy()
         background = self.background.detach().cpu().numpy()**2
         weights = self.weights.detach().cpu().numpy()
-        
+
         return {'basis':basis, 'translation':translations,
                 'probe':probe,'obj':obj,
                 'background':background,

temp_tests/simulated_dataset.py

@@ -1,7 +1,8 @@
 import numpy as np
 import torch as t
+from scipy import misc
 from CDTools.models import PolarizedFancyPtycho
-#from CDTools.datasets import Polarized2DDataset
+from CDTools.datasets import PolarizedPtycho2DDataset
 import CDTools
 from CDTools.tools import polarization
 from CDTools import tools
@@ -11,77 +12,121 @@ from PIL import Image
 # upolad 4 different images representing 4 components of the object
 # and 2 gaaussian functionas corresponding to the probe components
 
-a = np.asarray(Image.open('a.jpg'))
-b = np.asarray(Image.open('b.jpg'))
-c = np.asarray(Image.open('c.jpg'))
-d = np.asarray(Image.open('d.jpg'))
+f = misc.ascent()
+x , y = np.shape(f)
+aa = f[:x//2, :y//2]
+bb = f[:x//2, -y//2:]
+cc = f[-x//2:, :y//2]
+dd = f[-x//2:, -y//2:]
 
-#a = np.dot(a[..., :3], [.3, 6., .1])
+print(1)
 
-def simulate_dataset(probe_size, obj_size, num_patt):
-    translations = []
-    xs, ys = np.mgrid[:num_patt, :num_patt]
-    for x, y in zip(xs, ys):
-        translations.append((x*10e-3, y*10e-3))
-
-    translations = t.as_tensor(translations, dtype=t.float32)
-    a = t.as_tensor(a, dtype=t.cfloat)
-    a = t.tensordot(a, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1],[0]))[:obj_size, :obj_size]
-    
-    probe = tools.initializers.gaussian(np.array([probe_size, probe_size]), 50)
-    wavefields = tools.interactions.ptycho_2D_sinc(probe, obj, translations)
-    patterns = tools.propagators.far_field(wavefront)
-    patterns = np(patterns)
-    translations = np(t.cat((translations, t.zeros(num_patt)), dim=-1))
-
-    # needs to be stored as a cxi file
-    dataset = CDTools.datasets.Ptycho2DDataset.from_cxi('simulated_dataset.cxi')
-    dataset.detector_geometry = None
-
-def simulate polarized_datset(probe_size, obj_size, num_patt):
+def simulate_polarized_dataset(probe_size, obj_size, num_patt, a, b, c, d):
     a, b, c, d = t.as_tensor(a, dtype=t.cfloat), t.as_tensor(b, dtype=t.cfloat), t.as_tensor(c, dtype=t.cfloat), t.as_tensor(d, dtype=t.cfloat)
-    a = t.tensordot(a, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1],[0]))[:obj_size, :obj_size]
-    b = t.tensordot(b, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
-    c = t.tensordot(c, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
-    d = t.tensordot(d, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
-    
+    # a = t.tensordot(a, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1],[0]))[:obj_size, :obj_size]
+    # b = t.tensordot(b, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
+    # c = t.tensordot(c, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
+    # d = t.tensordot(d, t.tensor([.3, .6, .1], dtype=t.cfloat), dims=([-1], [0]))[:obj_size, :obj_size]
+
     translations = []
     xs, ys = np.mgrid[:num_patt, :num_patt]
+    xs, ys = np.ravel(xs), np.ravel(ys)
     for x, y in zip(xs, ys):
-        translations.append((x*10e-3, y*10e-3))
+        for j in range(9):
+            translations.append((5*x, 5*y))
     translations = t.as_tensor(translations, dtype=t.float32)
-    
-    obj = t.stack((t.stack((a, c), dim=0), t.stack((b, d), dim=0)), dim=-3)
-    probe = tools.initializers.gaussian(np.array([probe_size, probe_size]), 50)
+    obj = t.stack((t.stack((a, b), dim=-3), t.stack((c, d), dim=-3)), dim=-4)
+
+    # for i, j in zip([0, 0, 1, 1], [0, 1, 0, 1]):
+    #     plt.imshow(np.real(obj[i, j, ...]))
+    #     plt.show()
+    probe = tools.initializers.gaussian(np.array([probe_size, probe_size]), np.array((2, 2)))
     probe = t.stack((probe, probe), dim=-3)
-    probe = polarization.apply_circular_polarizer(probe)
+    probe = polarization.apply_circular_polarizer(probe, multiple_modes=False)
+    polarizers = polarization.generate_linear_polarizer(t.tensor([0, 45, 90]))
+    num_transl = len(translations)
+    pol_probes = [polarization.apply_jones_matrix(probe, polarizers[i], multiple_modes=False) for i in range(3)]
+    analyzers = t.stack(([polarizers[i % 3] for i in range(num_transl)]), dim=0)
+    analyzer = t.tensor([(i % 3) * 45 for i in range(num_transl)])
+    polarizer = t.tensor([(i // 3) % 3 * 45 for i in range(num_transl)])
+    polarized_probes = t.stack(([pol_probes[(i // 3) % 3] for i in range(num_transl)]), dim=0)
+    polarized_wavefields = t.stack([tools.interactions.ptycho_2D_sinc(polarized_probes[i], obj, translations[i], polarized=True, multiple_modes=False) for i in range(num_transl)])
+    wavefields = tools.propagators.far_field(polarized_wavefields)
+    wavefields = polarization.apply_jones_matrix(wavefields, analyzers, multiple_modes=False)
+    patterns = t.abs(wavefields[:, 0, :, :])**2 + t.abs(wavefields[:, 1, :, :])**2
+    detector_basis = t.transpose(t.tensor([[0, -4.8e-6, 0], [-4.8e-6, 0, 0]]), 0, 1)
+    det_shape = t.Size((obj_size, obj_size))
+    wavelength = 532e-9
+    real_basis = tools.initializers.exit_wave_geometry(detector_basis, det_shape, wavelength, 2.5e-2)[0]
+    # print('f', type(real_basis), translations.shape)
+    real_translations = tools.interactions.pixel_to_translations(real_basis, translations)
+    # print(type(real_translations))
+    patterns = patterns.numpy()
+    detector_geometry = {
+                         'corner': np.array([probe_size*4.8e-6/2, probe_size*4.8e-6/2, 2.5e-2]),
+                         'basis': np.array([[0, -4.8e-6, 0], [-4.8e-6, 0, 0]]).transpose(),
+                         'distance': 2.5e-2
+                        }
 
-    selections = tools.interactions.ptycho_2D_sinc(t.ones(2, probe_size, probe_size).to(dtype=t.cfloat), obj, translations, polarized=True)
-    polarizers = [polarization.generate_linear_polarizer(i * 45) for i in range(3)]
-    
-    pol_probes = [polarization.apply_jones_matrix(probe, polarizers[i]) for i in range(3)]
-    
+    return PolarizedPtycho2DDataset(real_translations, polarizer, analyzer, patterns,
+                                    axes=("x", "y"), detector_geometry=detector_geometry, wavelength=wavelength)
 
-    analyzer = t.stack(([polaryzers[i % 3] for i in range(num_patt)]), dim=0)
-    # probes = t.stack(([probes[i // 3] for i in num_patt]), dim=0)
-    wavefields = t.as_tensor([tools.interactions.ptycho_2D_sinc(pol_probes[i], obj, translations, polarized=True) for i in range(3)]).to(dtype=t.cfloat)
+# print(1)
+dataset = simulate_polarized_dataset(100, 500, 10, aa, bb, cc, dd)
+# print(2)
+dataset.inspect()
+# print(3)
+plt.show()
+model = PolarizedFancyPtycho.from_dataset(dataset, propagation_distance=1e-3)
+model.inspect()
+# plt.show()
+model.compare(dataset)
+plt.show()
 
-    wf = t.empty(1, 2, obj_size, obj_size)
-    for i in range(num_patt):
-        for j in range(3):
-            pol_channel = t.stack(([wavefileds[j] for k in range(3)]), dim=0)
-        wf = t.cat((wf, pol_channel), dim=0)
-    
-    pol_wavefieds = polarization.apply_jones_matrix(wf, analyzer)
+# for loss in model.Adam_optimize(400, dataset, batch_size=5, lr=0.002, schedule=True):
+#     # And we liveplot the updates to the model as they happen
+#     print(model.report())
+#     model.inspect(dataset)
+#     model.save_figures(prefix='simulated', extension='png')
+#     res = model.save_results(dataset)
+#     np.save('simulated_dataset.npy', res)
 
-    patterns = tools.propagators.far_field(pol_wavefieds)
+dataset = simulate_polarized_dataset(50, 100, 10, aa, bb, cc, dd)
 
-    translations = np(t.cat((translations, t.zeros(num_patt)), dim=-1))
-    patterns = np(patterns)
-    dataset.detector_geometry = None
-    # needs to be stored in a cxi file
-    dataset = CDTools.datasets.FancyPtycho2DDataset.from_cxi('polarized_simulated_dataset.cxi')
-    dataset.inspect()
-
-    
-    model = tools.models.PolarizedFancyPtycho.from_dataset(dataset)
+res = np.load('simulated_dataset.npy', allow_pickle=True)
+res = res[()]
+print(type(res))
+Ws = t.ones(len(dataset))
+ewg = CDTools.tools.initializers.exit_wave_geometry
+probe_basis, probe_shape, det_slice = ewg(res['basis'],
+                                          dataset[0][1].shape,
+                                          dataset.wavelength,
+                                          dataset.detector_geometry['distance'],
+                                          center=None,
+                                          padding=0)
+print('object', res['obj'].shape)
+obj = res['obj'][..., 200:275, 200:275]
+# print('obj', obj)
+a = obj[0, 0, :, :]
+b = obj[0, 1, :, :]
+c = obj[1, 0, :, :]
+d = obj[1, 1, :, :]
+for i in [a, b, c, d]:
+    plt.imshow(np.real(i))
+    # plt.imshow(np.log(np))
+    plt.colorbar()
+    plt.show()
+print(np.allclose(np.real(a), np.real(b)))
+print('a', a[..., 10:20, 10:20])
+print('b', b[..., 10:20, 10:20])
+models = [CDTools.models.FancyPtycho(dataset.wavelength, dataset.detector_geometry, probe_basis,
+                                     res['probe'], component, surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
+                                     min_translation=t.tensor([0, 0], dtype=t.float32),
+                                     background=t.tensor(res['background']), translation_offsets=None, mask=None,
+                                     weights=Ws, translation_scale=1, saturation=None,
+                                     probe_support=None, oversampling=1,
+                                     loss='amplitude mse', units='um') for component in [a, b, c, d]]
+# print('model is created')
+# for model in models:
+#     model.inspect()
+#     plt.show()