From 540cc997b5f91170f56de9ddb920b9c2ef1e7e6b Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Fri, 20 Aug 2021 11:52:23 -0400 Subject: [PATCH 1/6] . --- CDTools/models/base.py | 116 ++++++++-------- CDTools/models/polarized_fancy_ptycho.py | 108 +++++++++------ temp_tests/simulated_dataset.py | 169 ++++++++++++++--------- 3 files changed, 233 insertions(+), 160 deletions(-) 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() From 8493a3fb4113524b13fbe2143d3965ab1f05a790 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Fri, 20 Aug 2021 17:19:28 -0400 Subject: [PATCH 2/6] . --- CDTools/models/polarized_fancy_ptycho.py | 6 +- CDTools/tools/interactions/interactions.py | 104 ++++++++++----------- 2 files changed, 55 insertions(+), 55 deletions(-) diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index 2c095de..e81545e 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/CDTools/models/polarized_fancy_ptycho.py @@ -78,12 +78,11 @@ class PolarizedFancyPtycho(FancyPtycho): # 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 = t.cat((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) - + print('probe', type(probe), probe.shape) 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) @@ -91,6 +90,7 @@ class PolarizedFancyPtycho(FancyPtycho): obj = model.obj.detach() obj = t.stack((obj, obj), dim=-3) obj = t.stack((obj, obj), dim=-4) + print('object', type(obj), obj.shape) model.obj.data = obj # tensor vs tensor.data return model diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py index d2d8667..7f60942 100644 --- a/CDTools/tools/interactions/interactions.py +++ b/CDTools/tools/interactions/interactions.py @@ -17,14 +17,14 @@ __all__ = ['translations_to_pixel', 'pixel_to_translations', def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.])): """Takes real space translations and outputs them in pixel space - + This works for any 2D ptychography geometry. It takes in A set of translations in (x,y) space and outputs the same translations - in internal pixel units perpendicular to the detector. - + in internal pixel units perpendicular to the detector. + It uses information on the wavefield basis and, if defined, the sample normal, to perform the conversion. - + The assumed geometry is incoming radiation with a wavevector parallel to the +z axis, [0,0,1]. The default sample orientation has a surface normal parallel to this direction @@ -33,7 +33,7 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1. ---------- basis : torch.Tensor The real space basis the wavefields are defined in - translations : torch.Tensor + translations : torch.Tensor A Jx3 stack of real-space translations, or a single translation surface_normal : torch.Tensor Optional, the sample's surface normal @@ -68,18 +68,18 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1. return pixel_translations[0] else: return pixel_translations - + def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,0,1])): """Takes pixel-space translations and outputs them in real space - + This works for any 2D ptychography geometry. It takes in A set of internal pixel unit translations in (i,j) space and outputs the same translations real (x,y) space - + It uses information on the wavefield basis and, if defined, the sample normal, to perform the conversion. - + The assumed geometry is incoming radiation with a wavevector parallel to the +z axis, [0,0,1]. The default sample orientation has a surface normal parallel to this direction. Because of this, the z direction @@ -96,7 +96,7 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0, Returns ------- - real_translations : torch.Tensor + real_translations : torch.Tensor A Jx3 stack of real-space translations, or a single translation """ projection_1 = t.Tensor([[1,0,0], @@ -129,7 +129,7 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0, def project_translations_to_sample(sample_basis, translations): """Takes real space translations and outputs them in pixels in a sample basis - + This projection function is designed for the Bragg2DPtycho class. More broadly, it works to take a set of translations in the lab frame and convert each one into two values. First, an (i,j) value in pixels @@ -145,7 +145,7 @@ def project_translations_to_sample(sample_basis, translations): relative amount the probe needs to be propagated to reach any given location), a positive motion along the z-axis of the probe forming optics will lead to a negative propagation distance. - + The assumed geometry is incoming radiation with a wavevector parallel to the +z axis, [0,0,1]. @@ -153,7 +153,7 @@ def project_translations_to_sample(sample_basis, translations): ---------- sample_basis : torch.Tensor The real space basis the wavefields are defined in - translations : torch.Tensor + translations : torch.Tensor A Jx3 stack of real-space translations, or a single translation Returns @@ -171,7 +171,7 @@ def project_translations_to_sample(sample_basis, translations): # Then we calculate a matrix which can do the projection - + propagation_dir = t.Tensor(np.array([0,0,1])).to( device=surface_normal.device, dtype=surface_normal.dtype) @@ -179,13 +179,13 @@ def project_translations_to_sample(sample_basis, translations): I = t.eye(3).to( device=surface_normal.device, dtype=surface_normal.dtype) - + # Here we're setting up a matrix-vector equation mat*answer=input # At some point ger will need to be replaced by outer, but for now # outer many places still don't have new enough versions of torch. mat = t.cat((I - t.ger(propagation_dir,propagation_dir), surface_normal.unsqueeze(0))) - + # And we invert the matrix to do the projection projector = t.pinverse(mat)[:,:3].to(device=translations.device, dtype=translations.dtype) @@ -200,7 +200,7 @@ def project_translations_to_sample(sample_basis, translations): device=translations.device, dtype=translations.dtype) - + sample_projection = t.mm(basis_vectors_inv, projector).t() prop_projection = t.mm(propagation_dir_inv, projector).t() @@ -218,9 +218,9 @@ def project_translations_to_sample(sample_basis, translations): return pixel_translations[0], propagations[0] else: return pixel_translations, propagations - - + + def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj=False): """Returns a stack of exit waves without accounting for subpixel shifts @@ -229,15 +229,15 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj dimension as the translation index and the final dimensions corresponding to the detector. The exit waves are calculated by shifting the probe by the rounded value of the translation - + If multiple_modes is set to False, any additional dimensions in the ptycho_2D_round function will be assumed to correspond to the translation index. If multiple_modes is set to true, the (-4th) dimension of the probe will always be assumed to be defining a set of (P) incoherently mixing modes to be broadcast all translation indices. If any additional dimensions closer to the start exist, they will be assumed to be translation indices - - + + Parameters ---------- probe : torch.Tensor @@ -251,7 +251,7 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj Returns ------- - exit_waves : torch.Tensor + exit_waves : torch.Tensor An (N)x(P)xMxL tensor of the calculated exit waves """ @@ -260,9 +260,9 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj translations = translations[None,:] single_translation = True - + integer_translations = t.round(translations).to(dtype=t.int32) - + if upsample_obj: selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2]//2, tr[1]:tr[1]+probe.shape[-1]//2] @@ -292,7 +292,7 @@ def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj def ptycho_2D_linear(probe, obj, translations, shift_probe=True): """Returns a stack of exit waves accounting for subpixel shifts - + This function returns a collection of exit waves, with the first dimension as the translation index and the final dimensions corresponding to the detector. The exit waves are calculated by @@ -322,7 +322,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True): if translations.dim() == 1: translations = translations[None,:] single_translation = True - + # Separate the translations into a part that chooses the window # And a part that defines the windowing function integer_translations = t.floor(translations) @@ -342,15 +342,15 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True): sel01 = t.cat((probe[:,-1:],probe[:,:-1]),dim=1) sel10 = t.cat((probe[-1:,:],probe[:-1,:]),dim=0) sel11 = t.cat((sel01[-1:,:],sel01[:-1,:]),dim=0) - + selection = sel00 * (1-sp[0])*(1-sp[1]) + \ sel10 * sp[0]*(1-sp[1]) + \ sel01 * (1-sp[0])*sp[1] + \ sel11 * sp[0]*sp[1] - + obj_slice = obj[tr[0]:tr[0]+probe.shape[0], tr[1]:tr[1]+probe.shape[1]] - + exit_waves.append(selection * obj_slice) else: for tr, sp in zip(integer_translations, @@ -359,16 +359,16 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True): # Here we subpixel shift the object by (-i,-j) after # slicing out the correct translation of the probe # - + sel00 = obj[tr[0]:tr[0]+probe.shape[0], tr[1]:tr[1]+probe.shape[1]] - + sel01 = obj[tr[0]:tr[0]+probe.shape[0], tr[1]+1:tr[1]+1+probe.shape[1]] - + sel10 = obj[tr[0]+1:tr[0]+1+probe.shape[0], tr[1]:tr[1]+probe.shape[1]] - + sel11 = obj[tr[0]+1:tr[0]+1+probe.shape[0], tr[1]+1:tr[1]+1+probe.shape[1]] @@ -387,7 +387,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True): def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True, polarized=False, polarizer=None, analyzer=None): """Returns a stack of exit waves accounting for subpixel shifts - + This function returns a collection of exit waves, with the first dimension as the translation index and the final dimensions corresponding to the detector. The exit waves are calculated by @@ -427,7 +427,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi if translations.dim() == 1: translations = translations[None,:] single_translation = True - + # Separate the translations into a part that chooses the window # And a part that defines the windowing function integer_translations = t.floor(translations) @@ -480,7 +480,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi else: raise NotImplementedError('Object shift not yet implemented') - + print('ptyvho 2d sinc', output.shape) if single_translation: return output[0] else: @@ -489,7 +489,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, padding=10): """Returns a stack of exit waves accounting for subpixel shifts - + This function returns a collection of exit waves, with the first dimension as the translation index and the final dimensions corresponding to the detector. The exit waves are calculated by @@ -505,7 +505,7 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad on the input wavefield, and the first two indexes index differences from that pixel. It is easier to interpret the resulting matrix though if the latter two indices index locations in the output plane. NOTE: I believe - this change has now been made + this change has now been made Parameters ---------- @@ -529,17 +529,17 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad if translations.dim() == 1: translations = translations[None,:] single_translation = True - + # Separate the translations into a part that chooses the window # And a part that defines the windowing function integer_translations = t.floor(translations) subpixel_translations = translations - integer_translations integer_translations = integer_translations.to(dtype=t.int32) - + exit_waves = [] B = s_matrix.shape[0]//2 - + if shift_probe: i = t.arange(probe.shape[-2]) - probe.shape[-2]//2 j = t.arange(probe.shape[-1]) - probe.shape[-1]//2 @@ -548,14 +548,14 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad J = 2 * np.pi * J.to(t.float32) / probe.shape[-1] I = I.to(dtype=probe.dtype,device=probe.device) J = J.to(dtype=probe.dtype,device=probe.device) - + for tr, sp in zip(integer_translations, subpixel_translations): fft_probe = t.fft.fftshift(t.fft.fft2(probe), dim=(-1,-2)) shifted_fft_probe = fft_probe * t.exp(1j*(-sp[0]*I - sp[1]*J)) shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe, dim=(-1,-2))) - + s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[-2]+2*B, tr[1]:tr[1]+probe.shape[-1]+2*B] @@ -564,14 +564,14 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad device=s_matrix_slice.device, dtype=s_matrix_slice.dtype) - + for i in range(s_matrix.shape[0]): for j in range(s_matrix.shape[1]): output [i:i+probe.shape[-2],j:j+probe.shape[-1]] += \ shifted_probe * s_matrix_slice[i,j,i:i+probe.shape[-2],j:j+probe.shape[-1]] exit_waves.append(output) - + else: raise NotImplementedError('Object shift not yet implemented') @@ -579,11 +579,11 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad return exit_waves[0] else: return t.stack(exit_waves) - + def RPI_interaction(probe, obj): """Returns an exit wave from a high-res probe and a low-res obj - + In this interaction, the probe and object arrays are assumed to cover the same physical region of space, but with the probe array sampling that region of space more finely. Thus, to do the interaction, the object @@ -593,7 +593,7 @@ def RPI_interaction(probe, obj): method and is not commonly used elsewhere. This also works with object functions that have an extra first dimension - for an incoherently mixing model. + for an incoherently mixing model. Parameters @@ -610,7 +610,7 @@ def RPI_interaction(probe, obj): """ # TODO: The upsampling only works for arrays of even dimension! - + # The far-field propagator is just a 2D FFT but with an fftshift fftobj = propagators.far_field(obj) # We calculate the padding that we need to do the upsampling @@ -618,7 +618,7 @@ def RPI_interaction(probe, obj): pad0r = probe.shape[-2] - obj.shape[-2] - pad0l pad1l = (probe.shape[-1] - obj.shape[-1])//2 pad1r = probe.shape[-1] - obj.shape[-1] - pad1l - + if obj.dim() == 2: fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad0l, pad0r)) elif obj.dim() == 3: @@ -626,7 +626,7 @@ def RPI_interaction(probe, obj): fftobj, (pad1l, pad1r, pad0l, pad0r, 0,0)) else: raise NotImplementedError('RPI interaction with obj of dimension higher than 4 (including complex dimension) is not supported.') - + # Again, just an inverse FFT but with an fftshift upsampled_obj = propagators.inverse_far_field(fftobj) From 53742d024574d908c073dea15955487e5298198d Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Fri, 20 Aug 2021 21:45:27 -0400 Subject: [PATCH 3/6] . --- CDTools/models/base.py | 12 +++++++++++- CDTools/models/polarized_fancy_ptycho.py | 8 ++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index b200b3a..48a33d4 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -451,6 +451,16 @@ class CDIModel(t.nn.Module): """ + print('base models inspect: checking the object') + a = self.obj.detach() + def saveobj(a, filename): + a = np.abs(a) + plt.imshow(a) + plt.savefig('filename') + f = ['base_a.png', 'base_b.png', 'base_c.png', 'base_d.png'] + comp = [a[i, j, :, :] for i, j in zip([0, 0, 1, 1], [0, 1, 0, 1])] + for i in range(4): + saveobj(comp[i], f[i]) first_update = False if update and hasattr(self, 'figs') and self.figs: figs = self.figs @@ -585,7 +595,7 @@ class CDIModel(t.nn.Module): 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') axes[1].set_title('Measured') diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index e81545e..93ae931 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/CDTools/models/polarized_fancy_ptycho.py @@ -92,6 +92,13 @@ class PolarizedFancyPtycho(FancyPtycho): obj = t.stack((obj, obj), dim=-4) print('object', type(obj), obj.shape) model.obj.data = obj + print('polarized fancy ptycho from datset obj') + a = obj.detach() + plt.imshow(np.real(a[0, 0, :, :])) + plt.show() + plt.imshow(np.real(a[0, 1, :, :])) + plt.show() + # tensor vs tensor.data return model @@ -130,6 +137,7 @@ class PolarizedFancyPtycho(FancyPtycho): shift_probe=True, multiple_modes=True, polarized=True) analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer) + # print('POLARIZED FANCY PTYCHO INTERACTION OBJ') return analyzed_exit_waves From aadd720f1e309191e95bd6a52cedaf67b1b48f61 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Fri, 20 Aug 2021 23:46:24 -0400 Subject: [PATCH 4/6] . --- CDTools/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 48a33d4..9d8e9eb 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -456,7 +456,7 @@ class CDIModel(t.nn.Module): def saveobj(a, filename): a = np.abs(a) plt.imshow(a) - plt.savefig('filename') + plt.savefig(filename) f = ['base_a.png', 'base_b.png', 'base_c.png', 'base_d.png'] comp = [a[i, j, :, :] for i, j in zip([0, 0, 1, 1], [0, 1, 0, 1])] for i in range(4): From d0bca7633b39f2adb6208b30b92ffb9935baf6ef Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Tue, 24 Aug 2021 01:26:11 -0400 Subject: [PATCH 5/6] . --- CDTools/models/polarized_fancy_ptycho.py | 153 ++++++++++++++++++++- CDTools/tools/polarization/polarization.py | 55 ++++---- 2 files changed, 177 insertions(+), 31 deletions(-) diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index 93ae931..2331dcc 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/CDTools/models/polarized_fancy_ptycho.py @@ -73,10 +73,6 @@ class PolarizedFancyPtycho(FancyPtycho): else: x = -1j - # 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.cat((probe, probe * x), dim=-3) probe_max = t.max(t.abs(probe)) @@ -104,7 +100,153 @@ class PolarizedFancyPtycho(FancyPtycho): polarizers = [tools.polarization.generate_linear_polarizer(i * 45) for i in range(3)] - # WHAT IS INDEX? + @classmethod + def from_dataset2(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, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um'): + + wavelength = dataset.wavelength + det_basis = dataset.detector_geometry['basis'] + det_shape = dataset[0][1].shape + distance = dataset.detector_geometry['distance'] + + # always do this on the cpu + get_as_args = dataset.get_as_args + dataset.get_as(device='cpu') + + # We include the *extras to make this work even with datasets, like + # polarization dependent datasets, that might toss out extra inputs + (indices, translations, polarizer, analyzer), patterns = dataset[:] + + dataset.get_as(*get_as_args[0], **get_as_args[1]) + + # Set to none to avoid issues with things outside the detector + if auto_center: + center = tools.image_processing.centroid(t.sum(patterns, dim=0)) + else: + center = None + + if left_polarized: + x = 1j + else: + x = -1j + + + # Then, generate the probe geometry from the dataset + ewg = tools.initializers.exit_wave_geometry + probe_basis, probe_shape, det_slice = ewg(det_basis, + det_shape, + wavelength, + distance, + center=center, + padding=padding, + opt_for_fft=opt_for_fft, + oversampling=oversampling) + + probe_shape = t.stack((2, probe_shape), dim=-3) + if hasattr(dataset, 'sample_info') and \ + dataset.sample_info is not None and \ + 'orientation' in dataset.sample_info: + surface_normal = dataset.sample_info['orientation'][2] + else: + surface_normal = np.array([0., 0., 1.]) + + # If this information is supplied when the function is called, + # then we override the information in the .cxi file + if scattering_mode in {'t', 'transmission'}: + surface_normal = np.array([0., 0., 1.]) + elif scattering_mode in {'r', 'reflection'}: + outgoing_dir = np.cross(det_basis[:, 0], det_basis[:, 1]) + outgoing_dir /= np.linalg.norm(outgoing_dir) + surface_normal = outgoing_dir + np.array([0., 0., 1.]) + surface_normal /= -np.linalg.norm(surface_normal) + + # Next generate the object geometry from the probe geometry and + # the translations + pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal) + + obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200) + + if hasattr(dataset, 'background') and dataset.background is not None: + background = t.sqrt(dataset.background) + else: + background = None + + # Finally, initialize the probe and object using this information + 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) + + # Now we initialize all the subdominant probe modes + 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) + # probe = t.stack([tools.propagators.far_field(probe),] + probe_stack) + probe_x, probe_y = probe, probe * x + probe = t.stact((probe_x, probe_y), dim=-3) + + a = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) + b = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) + c = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) + d = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) + + ab = t.stack((a, b), dim=-3) + cd = t.stack((c, d), dim=-3) + obj = t.stack((ab, cd), dim=-4) + det_geo = dataset.detector_geometry + + translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5) + + if dm_rank is not None and dm_rank != 0: + if dm_rank > n_modes: + raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.') + elif dm_rank == -1: + # dm_rank == -1 is defined to mean full-rank + dm_rank = n_modes + + Ws = t.zeros(len(dataset), dm_rank, n_modes, dtype=t.complex64) + # Start with as close to the identity matrix as possible, + # cutting of when we hit the specified maximum rank + for i in range(0, dm_rank): + Ws[:, i, i] = 1 + else: + # dm_rank == None or dm_rank = 0 triggers a special case where + # a standard incoherent multi-mode model is used. This is the + # default, because it is so common. + # In this case, we define a set of weights which only has one index + Ws = t.ones(len(dataset)) + + if hasattr(dataset, 'mask') and dataset.mask is not None: + mask = dataset.mask.to(t.bool) + else: + mask = None + + if probe_support_radius is not None: + probe_support = t.zeros(probe[0].shape, dtype=t.bool) + xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]] + xs = xs - np.mean(xs) + ys = ys - np.mean(ys) + Rs = np.sqrt(xs**2 + ys**2) + + probe_support[Rs < probe_support_radius] = 1 + probe = probe * probe_support[None, :, :] + + else: + probe_support = None + + return cls(wavelength, det_geo, probe_basis, probe, obj, + detector_slice=det_slice, + surface_normal=surface_normal, + min_translation=min_translation, + translation_offsets=translation_offsets, + weights=Ws, mask=mask, background=background, + translation_scale=translation_scale, + saturation=saturation, + probe_support=probe_support, + oversampling=oversampling, + loss=loss, units=units) + + + def interaction(self, index, translations, polarizer, analyzer, test=False): # Step 1 is to convert the translations for each position into a @@ -137,7 +279,6 @@ class PolarizedFancyPtycho(FancyPtycho): shift_probe=True, multiple_modes=True, polarized=True) analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer) - # print('POLARIZED FANCY PTYCHO INTERACTION OBJ') return analyzed_exit_waves diff --git a/CDTools/tools/polarization/polarization.py b/CDTools/tools/polarization/polarization.py index 7014e20..f3bc04c 100644 --- a/CDTools/tools/polarization/polarization.py +++ b/CDTools/tools/polarization/polarization.py @@ -14,7 +14,7 @@ __all__ = ['apply_linear_polarizer', 'apply_circular_polarizer', 'apply_jones_matrix', 'generate_linear_polarizer', - 'generate_phase_retarder'] + 'generate_birefringent_obj'] # Abe - split these into two functions @@ -36,10 +36,10 @@ def generate_linear_polarizer(pol_angle): cd = t.stack((c, d), dim=-1) jones_matrices = t.stack((ab, cd), dim=-2) if single_angle: - return jones_matrices[0].to(dtype=t.cfloat) + return jones_matrices[0].to(dtype=t.cfloat) else: return jones_matrices.to(dtype=t.cfloat) - + def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True): """ @@ -56,7 +56,7 @@ def apply_linear_polarizer(probe, polarizer, multiple_modes=True, transpose=True Returns: -------- linearly polarized probe: t.Tensor - (N)(P)x2x1xMxL + (N)(P)x2x1xMxL """ jones_matrices = generate_linear_polarizer(polarizer) return apply_jones_matrix(probe, jones_matrices, transpose=transpose, multiple_modes=multiple_modes) @@ -75,19 +75,19 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True) probe: t.Tensor A (N)(P)x2xMxL tensor representing the probe jones_matrix: t.tensor - (N)x2x2x(M)x(L) + (N)x2x2x(M)x(L) Returns: -------- a probe with the jones matrix applied: t.Tensor - (N)(P)x2xMxL + (N)(P)x2xMxL """ if transpose: if jones_matrix.dim() < 4: jones_matrix = jones_matrix[..., None, None] if multiple_modes: - jones_matrix = jones_matrix.unsqueeze(-5) + jones_matrix = jones_matrix.unsqueeze(-5) probe = probe[..., None, :, :] # if jones matrices do not differ from pattern to pattern if probe.dim() > jones_matrix.dim(): @@ -96,19 +96,19 @@ def apply_jones_matrix(probe, jones_matrix, transpose=True, multiple_modes=True) elif jones_matrix.dim() > probe.dim(): probe = probe.unsqueeze(0) # print('apply jonesmatrix: probe', probe.shape, 'matrix:', jones_matrix) - jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4) + jones_matrix = jones_matrix.transpose(-1, -3).transpose(-2, -4) probe = probe.transpose(-1, -3).transpose(-2, -4) output = t.matmul(jones_matrix, probe).transpose(-2, -4).transpose(-1, -3).squeeze(-3) - + else: raise NotImplementedError - + return output def apply_phase_retardance(probe, phase_shift, multiple_modes=True): """ - Shifts the y-component of the field wrt the x-component by a given phase shift + Shifts the y-component of the field wrt the x-component by a given phase shift Parameters: ---------- @@ -120,7 +120,7 @@ def apply_phase_retardance(probe, phase_shift, multiple_modes=True): Returns: -------- probe: t.Tensor - (...)x2x1xMxL + (...)x2x1xMxL """ theta = t.as_tensor(phase_shift, dtype=t.float32) theta = t.deg2rad(theta) @@ -140,11 +140,11 @@ def apply_circular_polarizer(probe, left_polarized=True, multiple_modes=True): A (...)x2xMxL tensor representing the probe left_polarizd: bool True for the left-polarization, False for the right - + Returns: -------- circularly polarized probe: t.Tensor - (...)x2xMxL + (...)x2xMxL """ probe = probe.to(dtype=t.cfloat) if left_polarized: @@ -166,7 +166,7 @@ def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True): Returns: -------- polarized probe: t.Tensor - (...)x2x1xMxL + (...)x2x1xMxL """ probe = probe.to(dtype=t.cfloat) theta = math.radians(fast_axis_angle) @@ -174,7 +174,7 @@ def apply_quarter_wave_plate(probe, fast_axis_angle, multiple_modes=True): jones_matrix = exponent* t.tensor([[(cos(theta))**2 + 1j * (sin(theta))**2, (1 - 1j) * sin(theta) * cos(theta)], [(1 - 1j) * sin(theta) * cos(theta), (sin(theta))**2 + 1j * (cos(theta))**2]]).to(dtype=t.cfloat) out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes) - return out + return out def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True): """ @@ -188,7 +188,7 @@ def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True): Returns: -------- polarized probe: t.Tensor - (...)x2x1xMxL + (...)x2x1xMxL """ probe = probe.to(dtype=t.cfloat) theta = math.radians(fast_axis_angle) @@ -196,19 +196,24 @@ def apply_half_wave_plate(probe, fast_axis_angle, multiple_modes=True): jones_matrix = exponent * t.tensor([[(cos(theta))**2 - (sin(theta))**2, 2 * sin(theta) * cos(theta)], [2 * sin(theta) * cos(theta), (sin(theta))**2 - (cos(theta))**2]]).to(dtype=t.cfloat) out = apply_jones_matrix(probe, jones_matrix, multiple_modes=multiple_modes) - return out - -def generate_phase_retarder(fast_axis=0, phase=0): - phase = t.as_tensor(phase).to(dtype=t.float32) - phase = t.deg2rad(phase) - def coord_rot(angle): + return out + +def generate_birefringent_obj(fast_axis=90, phase_ret=10, atten_fast=1, atten_ret=1, global_phase=0): + def to_rad(angle): angle = t.as_tensor(angle, dtype=t.float32) angle = t.deg2rad(angle) + return angle + + fast_axis = to_rad(fast_axis) + phase_ret = to_rad(phase_ret) + global_phase = to_rad(global_phase) + + def coord_rot(angle): a = t.stack((t.cos(angle), t.sin(angle)), dim=-1) b = t.stack((-t.sin(angle), t.cos(angle)), dim=-1) return t.stack((a, b), dim=-2).to(dtype=t.cfloat) r1 = coord_rot(-fast_axis) r2 = coord_rot(fast_axis) - p = t.as_tensor([[1, 0], [0, t.exp(phase*1j)]], dtype=t.cfloat) - return t.matmul(r1, t.matmul(p, r2)) \ No newline at end of file + p = t.exp(global_phase * 1j) * t.as_tensor([[atten_fast, 0], [0, atten_ret * t.exp(phase_ret*1j)]], dtype=t.cfloat) + return t.matmul(r1, t.matmul(p, r2)) From e7c15273fa16fbe4660afd7623f24722671feb24 Mon Sep 17 00:00:00 2001 From: Anastasiia Kutakh Date: Tue, 24 Aug 2021 11:31:35 -0400 Subject: [PATCH 6/6] . --- CDTools/models/polarized_fancy_ptycho.py | 2 +- CDTools/tools/plotting/plotting.py | 86 +++++++++++++----------- 2 files changed, 46 insertions(+), 42 deletions(-) diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index 2331dcc..e16b664 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/CDTools/models/polarized_fancy_ptycho.py @@ -62,7 +62,7 @@ class PolarizedFancyPtycho(FancyPtycho): @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') + 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', left_polarized=True) # Mutate the class to its subclass diff --git a/CDTools/tools/plotting/plotting.py b/CDTools/tools/plotting/plotting.py index b5c65a8..f1451b2 100644 --- a/CDTools/tools/plotting/plotting.py +++ b/CDTools/tools/plotting/plotting.py @@ -17,7 +17,11 @@ from matplotlib import ticker, patheffects __all__ = ['colorize', 'plot_amplitude', 'plot_phase', 'plot_colorized', 'plot_translations', 'get_units_factor', 'plot_nanomap', 'plot_real', 'plot_imag', - 'plot_nanomap_with_images'] + 'plot_nanomap_with_images', + 'polarized_plot_component_amplitudes', + 'polarized_plot_phase_ret', + 'polarized_plot_global_phases', + 'polarized_plot_ellipses'] def colorize(z): @@ -84,7 +88,7 @@ def get_units_factor(units): def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, **kwargs): """Plots an image with a colorbar and on an appropriate spatial grid - + If a figure is given explicitly, it will clear that existing figure and plot over it. Otherwise, it will generate a new figure. @@ -94,7 +98,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', Finally, if a function is passed to the plot_func argument, this function will be called on each slice of data before it is plotted. This is used internally to enable the plot_real, plot_image, plot_phase, etc. functions. - + Parameters ---------- @@ -120,7 +124,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', used_fig : matplotlib.figure.Figure The figure object that was actually plotted to. """ - + # convert to numpy if isinstance(im, t.Tensor): # If final dimension is 2, assume it is a complex array. If not, @@ -142,7 +146,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', title = plt.gca().get_title() fig.clear() - + # If im only has two dimensions, this reshape will add a leading # dimension, and update will be called on index 0. If it has 3 or more # dimensions, then all the leading dimensions will be compressed into @@ -151,9 +155,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', reshaped_im = im.reshape(-1,s[-2],s[-1]) num_images = reshaped_im.shape[0] fig.plot_idx = idx % num_images - + to_plot = plot_func(reshaped_im[fig.plot_idx]) - + #Plot in a basis if it exists, otherwise dont if basis is not None: if isinstance(basis,t.Tensor): @@ -181,7 +185,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', plt.xlabel('j (pixels)') plt.ylabel('i (pixels)') - + plt.title(title) if len(im.shape) >= 3: @@ -194,14 +198,14 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', result = make_plot(0) update = make_plot - - + + def on_action(event): if not hasattr(event, 'button'): event.button = None if not hasattr(event, 'key'): event.key = None - + if event.key == 'up' or event.button == 'up': update(fig.plot_idx - 1) elif event.key == 'down' or event.button == 'down': @@ -217,9 +221,9 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', fig.my_callbacks = [] fig.my_callbacks.append(fig.canvas.mpl_connect('key_press_event',on_action)) fig.my_callbacks.append(fig.canvas.mpl_connect('scroll_event',on_action)) - + return result - + def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs): """Plots the real part of a complex array with dimensions NxM @@ -256,7 +260,7 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_ return plot_image(im, plot_func=plot_func, fig=fig, basis=basis, units=units, cmap=cmap, cmap_label=cmap_label, **kwargs) - + def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs): @@ -304,7 +308,7 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', If a basis is explicitly passed, the image will be plotted in real-space coordinates. - + Parameters ---------- im : array @@ -520,12 +524,12 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr def plot_nanomap_with_images(translations, get_image_func, values=None, mask=None, basis=None, fig=None, nanomap_units='$\\mu$m', image_units='$\\mu$m', convention='probe', image_title='Image', image_colorbar_title='Image Amplitude', nanomap_colorbar_title='Integrated Intensity', cmap='viridis', **kwargs): """Plots a nanomap, with an image or stack of images for each point - + In many situations, ptychography data or the output of ptychography reconstructions is formatted as a set of images associated with various points in real space. This function is designed to allow for browsing through this kind of data, by making it possible to visualize a - + """ # This should pull heavily from the dataset.inspect function @@ -558,11 +562,11 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch s0 /= 4 # A rough value to make the size work out s = np.ones(translations.shape[0]) * s0 - + s[idx] *= 4 return s - + def update_colorbar(im): # # This solves the problem of the colorbar being changed @@ -571,29 +575,29 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non if hasattr(im, 'norecurse') and im.norecurse: im.norecurse=False return - + im.norecurse=True # This is needed to update the colorbar # only change limits if array contains multiple values if np.min(im.get_array()) != np.max(im.get_array()): im.set_clim(vmin=np.min(im.get_array()), vmax=np.max(im.get_array())) - + # # The meatiest part of this program, here we just go through and # set up the plot how we want it # - + # First we set up the left-hand plot, which shows an overview map axes[0].set_title('Relative Displacement Map') - + translations = translations.detach().cpu().numpy() if convention.lower() != 'probe': translations = translations * -1 - + s = calculate_sizes(0) - + nanomap_units_factor = get_units_factor(nanomap_units) nanomap = axes[0].scatter(nanomap_units_factor * translations[:,0], nanomap_units_factor * translations[:,1], @@ -614,7 +618,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # where the colorbar should have been to avoid stretching the # nanomap plot, while still not showing the (now useless) colorbar. cb1.remove() - + # Now we set up the second plot, which shows the individual # diffraction patterns axes[1].set_title(image_title) @@ -634,7 +638,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # This fails if the basis is not rectangular basis_norm = np.linalg.norm(np_basis, axis = 0) basis_norm = basis_norm * get_units_factor(image_units) - + extent = [0, example_im.shape[-1]*basis_norm[1], 0, example_im.shape[-2]*basis_norm[0]] else: @@ -655,9 +659,9 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non axes[1].text_box.set_path_effects( [patheffects.Stroke(linewidth=2, foreground='black'), patheffects.Normal()]) - + meas = axes[1].imshow(im, extent=extent, cmap=cmap) - + cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal', format='%.2e', ticks=ticker.LinearLocator(numticks=5), @@ -665,7 +669,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non cb2.ax.tick_params(labelrotation=20) cb2.ax.set_title(image_colorbar_title, size="medium", pad=5) cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas)) - + # This function handles all the updating, except for moving the # slider value. This is done because the slider widget is # ultimately responsible for triggering an update, so all other @@ -675,7 +679,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # We have to explicitly make it an integer because the slider will # output floats (even if they are still integer-valued) idx = int(idx) - + # Get the new data for this index im = get_image_func(idx) if len(im.shape) >= 3: @@ -686,22 +690,22 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non axes[1].image_idx = im_idx axes[1].text_box.set_text(str(im_idx)) im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx] - + # Now we resize the nanomap to show the new selection axes[0].collections[0].set_sizes(calculate_sizes(idx)) - + # And we update the data in the image as well ax_im = axes[1].images[-1] ax_im.set_data(im) update_colorbar(ax_im) - + # # Now we define the functions to handle various kinds of events # that can be thrown our way # - + # We start by creating the slider here, so it can be used # by the update hooks. slider = Slider(axslider, 'Image #', 0, translations.shape[0]-1, valstep=1, valfmt="%d") @@ -715,7 +719,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # while the mouse is within the image display im = im.reshape(-1,im.shape[-2],im.shape[-1]) im_idx = axes[1].image_idx - + if event.key == 'up' or event.button == 'up' \ or event.key == 'left': im_idx = (im_idx - 1) % im.shape[0] @@ -733,7 +737,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non event.button = None if not hasattr(event, 'key'): event.key = None - + if event.key == 'up' or event.button == 'up' or event.key == 'left': idx = slider.val - 1 elif event.key == 'down' or event.button == 'down' or event.key == 'right': @@ -742,7 +746,7 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # This prevents errors from being thrown on irrelevant key # or mouse input return - + # Handle the wraparound and trigger the update idx = int(idx) % translations.shape[0] slider.set_val(idx) @@ -753,16 +757,16 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # for example, scroll events that happen over the nanomap if event.mouseevent.button == 1: slider.set_val(event.ind[0]) - + # Here we connect the various update functions cid1 = fig.canvas.mpl_connect('pick_event',on_pick) cid2 = fig.canvas.mpl_connect('key_press_event',on_action) cid3 = fig.canvas.mpl_connect('scroll_event',on_action) # It's so dumb that matplotlib doesn't automatically track this for you - fig.nanomap_cids = [cid1,cid2,cid3] + fig.nanomap_cids = [cid1,cid2,cid3] slider.on_changed(update) - + # Throw an extra update into the mix just to get rid of any things # (like the nanomap dot sizes) that otherwise would change on the # first update