diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 922d055..b200b3a 100644 --- a/CDTools/models/base.py +++ b/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) - - diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index ed21fa1..2c095de 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/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, diff --git a/temp_tests/simulated_dataset.py b/temp_tests/simulated_dataset.py index 8972e57..817996b 100644 --- a/temp_tests/simulated_dataset.py +++ b/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()