From 1b5ef2271cea38f2b2bf18643ce9ceb08223d3d3 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 26 Oct 2021 13:11:11 -0400 Subject: [PATCH] Buncha stuff, sorry I wasn't pushing --- CDTools/datasets/ptycho_2d_dataset.py | 9 +-- CDTools/models/base.py | 11 ++-- CDTools/models/bragg_2d_ptycho.py | 42 ++++++++++---- CDTools/models/fancy_ptycho.py | 67 +++++++++++++++++----- CDTools/models/multislice_2d_ptycho.py | 2 +- CDTools/tools/data/data.py | 6 +- CDTools/tools/interactions/interactions.py | 8 +-- CDTools/tools/plotting/plotting.py | 6 +- 8 files changed, 105 insertions(+), 46 deletions(-) diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index f2c1c76..c7a9661 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -188,7 +188,7 @@ class Ptycho2DDataset(CDataset): cdtdata.add_ptycho_translations(cxi_file, self.translations) - def inspect(self, logarithmic=True, units='um'): + def inspect(self, logarithmic=True, units='um', log_offset=1): """Launches an interactive plot for perusing the data This launches an interactive plotting tool in matplotlib that @@ -208,7 +208,7 @@ class Ptycho2DDataset(CDataset): mask = 1 if logarithmic: - return np.log(meas_data) / np.log(10) * mask + return np.log(meas_data + log_offset) / np.log(10) * mask else: return meas_data * mask @@ -229,9 +229,10 @@ class Ptycho2DDataset(CDataset): # nanomap_values = (self.mask * self.patterns).sum(dim=(1,2)).detach().cpu().numpy() if logarithmic: - cbar_title='Log Base 10 of Diffraction Intensity' + cbar_title = ('Log Base 10 of Diffraction Intensity + %0.2f' + % log_offset) else: - cbar_title='Diffraction Intensity' + cbar_title = 'Diffraction Intensity' plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title) diff --git a/CDTools/models/base.py b/CDTools/models/base.py index 922d055..8eabaa6 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -55,7 +55,7 @@ class CDIModel(t.nn.Module): def __init__(self): super(CDIModel,self).__init__() - self.iteration_count = 0 + self.loss_train = [] def from_dataset(self, dataset): raise NotImplementedError() @@ -183,9 +183,8 @@ class CDIModel(t.nn.Module): if scheduler is not None: scheduler.step(loss) - self.latest_loss = loss + self.loss_train.append(loss) self.latest_iteration_time = time.time() - t0 - self.iteration_count += 1 return loss if thread: @@ -412,10 +411,10 @@ class CDIModel(t.nn.Module): report : str A string with basic info on the latest iteration """ - if hasattr(self, 'latest_loss'): - return 'Iteration ' + str(self.iteration_count) + \ + if hasattr(self, 'latest_iteration_time'): + return 'Iteration ' + str(len(self.loss_train)) + \ ' completed in %0.2f s with loss ' %\ - self.latest_iteration_time + str(self.latest_loss) + self.latest_iteration_time + str(self.loss_train[-1]) else: return 'No reconstruction iterations performed yet!' diff --git a/CDTools/models/bragg_2d_ptycho.py b/CDTools/models/bragg_2d_ptycho.py index 37188cd..b02cdce 100644 --- a/CDTools/models/bragg_2d_ptycho.py +++ b/CDTools/models/bragg_2d_ptycho.py @@ -54,6 +54,16 @@ __all__ = ['Bragg2DPtycho'] class Bragg2DPtycho(CDIModel): + # Needed to do the real/complex split + #@property + #def obj(self): + # return t.complex(self.obj_real, self.obj_imag) + + #@property + #def probe(self): + # return t.complex(self.probe_real, self.probe_imag) + + def __init__(self, wavelength, detector_geometry, probe_basis, probe_guess, obj_guess, detector_slice=None, @@ -124,7 +134,13 @@ class Bragg2DPtycho(CDIModel): # you look at the phase map probe_guess[probe_guess == 0] = 0 else: - self.probe_support = t.ones(self.probe[0].shape, dtype=t.bool) + self.probe_support = t.ones(probe_guess[0].shape, dtype=t.bool) + + #self.probe_real = t.nn.Parameter(probe_guess.real / self.probe_norm) + #self.probe_imag = t.nn.Parameter(probe_guess.imag / self.probe_norm) + + #self.obj_real = t.nn.Parameter(obj_guess.real ) + #self.obj_imag = t.nn.Parameter(obj_guess.imag) self.probe = t.nn.Parameter(probe_guess / self.probe_norm) self.obj = t.nn.Parameter(obj_guess) @@ -164,7 +180,7 @@ class Bragg2DPtycho(CDIModel): # recall that here we always want the shape of the detector # before it's cut down by the detector slice to match the # physical detector region - probe_shape = self.probe[0] + probe_shape = self.probe[0].shape self.k_map, self.intensity_map = \ tools.propagators.generate_high_NA_k_intensity_map( @@ -192,7 +208,7 @@ class Bragg2DPtycho(CDIModel): @classmethod - def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True,correct_tilt=True, lens=False, opt_for_fft=False): + def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True, correct_tilt=True, lens=False, opt_for_fft=False): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -309,12 +325,16 @@ class Bragg2DPtycho(CDIModel): if probe_support_radius is not None: probe_support = t.zeros(probe[0].shape, dtype=t.bool) - p_cent = np.array(probe[0].shape).astype(int) // 2 - psr = int(probe_support_radius) - probe_support[p_cent[0]-psr:p_cent[0]+psr, - p_cent[1]-psr:p_cent[1]+psr] = 1 + 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; + probe_support = None # Here we need to implement a simple condition to choose whether # to propagate the probe or not @@ -352,7 +372,7 @@ class Bragg2DPtycho(CDIModel): pix_trans += self.translation_scale * self.translation_offsets[index] Ws = self.weights[index] - prs = Ws[...,None,None,None] * self.probe + prs = Ws[...,None,None,None] * self.probe * self.probe_support[...,:,:] # Now we need to propagate each of the probes @@ -504,8 +524,10 @@ class Bragg2DPtycho(CDIModel): obj = self.obj.detach().cpu().numpy() background = self.background.detach().cpu().numpy()**2 weights = self.weights.detach().cpu().numpy() + losses = np.array(self.loss_train) return {'basis':basis, 'translation':translations, 'probe':probe,'obj':obj, 'background':background, - 'weights':weights} + 'weights':weights, + 'losses':losses} diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 84b4b48..7a8553a 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -17,14 +17,23 @@ class FancyPtycho(CDIModel): def __init__(self, wavelength, detector_geometry, probe_basis, - probe_guess, obj_guess, + probe_guess, + obj_guess, detector_slice=None, surface_normal=t.tensor([0., 0., 1.], dtype=t.float32), min_translation=t.tensor([0, 0], dtype=t.float32), - background=None, translation_offsets=None, mask=None, - weights=None, translation_scale=1, saturation=None, - probe_support=None, oversampling=1, - loss='amplitude mse', units='um'): + background=None, + translation_offsets=None, + mask=None, + weights=None, + translation_scale=1, + saturation=None, + probe_support=None, + oversampling=1, + fourier_probe=False, + loss='amplitude mse', + units='um', + ): super(FancyPtycho, self).__init__() self.wavelength = t.tensor(wavelength) @@ -45,12 +54,13 @@ class FancyPtycho(CDIModel): self.saturation = saturation self.units = units + self.fourier_probe = fourier_probe if mask is None: self.mask = mask else: self.mask = t.tensor(mask, dtype=t.bool) - + probe_guess = t.tensor(probe_guess, dtype=t.complex64) obj_guess = t.tensor(obj_guess, dtype=t.complex64) @@ -60,7 +70,7 @@ class FancyPtycho(CDIModel): 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) self.obj = t.nn.Parameter(obj_guess) @@ -117,7 +127,25 @@ class FancyPtycho(CDIModel): @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, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um'): + 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, + scattering_mode=None, + oversampling=1, + auto_center=False, + opt_for_fft=False, + fourier_probe=False, + loss='amplitude mse', + units='um' + ): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -188,8 +216,12 @@ class FancyPtycho(CDIModel): # 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)] + + # For a Fourier space probe + if fourier_probe: + probe = tools.propagators.far_field(probe) + probe = t.stack([probe, ] + probe_stack) - # probe = t.stack([tools.propagators.far_field(probe),] + probe_stack) obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) @@ -243,6 +275,7 @@ class FancyPtycho(CDIModel): translation_scale=translation_scale, saturation=saturation, probe_support=probe_support, + fourier_probe=fourier_probe, oversampling=oversampling, loss=loss, units=units) @@ -267,6 +300,10 @@ class FancyPtycho(CDIModel): # This restricts the basis probes to stay within the probe support basis_prs = self.probe * self.probe_support[..., :, :] + # For a Fourier-space probe, we take an IFT + if self.fourier_probe: + basis_prs = tools.propagators.inverse_far_field(basis_prs) + # Now we construct the probes for each shot from the basis probes Ws = self.weights[index] if len(self.weights[0].shape) == 0: @@ -510,10 +547,14 @@ class FancyPtycho(CDIModel): ('', lambda self, fig, dataset: self.plot_wavefront_variation(dataset, fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', image_colorbar_title='Probe Phase'), lambda self: len(self.weights.shape) >= 2), - ('Basis Probe Amplitudes (scroll to view modes)', - lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)), - ('Basis Probe Phases (scroll to view modes)', - lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis, units=self.units)), + ('Basis Probe Fourier Space Amplitudes', + lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), + ('Basis Probe Fourier Space Phases', + lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), + ('Basis Probe Real Space Amplitudes', + lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)), + ('Basis Probe Real Space Phases', + lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)), ('Average Density Matrix Amplitudes', lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()), axis=0), fig=fig), lambda self: len(self.weights.shape) >= 2), diff --git a/CDTools/models/multislice_2d_ptycho.py b/CDTools/models/multislice_2d_ptycho.py index 2843e3d..45abeb5 100644 --- a/CDTools/models/multislice_2d_ptycho.py +++ b/CDTools/models/multislice_2d_ptycho.py @@ -567,7 +567,7 @@ class Multislice2DPtycho(CDIModel): # Needs to be updated to allow for plotting to an existing figure - plot_list = [ + plot_list = [ ('Probe Fourier Space Amplitude', lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)), ('Probe Fourier Space Phase', diff --git a/CDTools/tools/data/data.py b/CDTools/tools/data/data.py index ea59ac3..cd2da1a 100644 --- a/CDTools/tools/data/data.py +++ b/CDTools/tools/data/data.py @@ -66,9 +66,9 @@ def get_entry_info(cxi_file): 'program_name'] metadata = {attr: str(e1[attr][()].decode()) for attr in metadata_attrs if attr in e1} - datetime_attrs = ['start_time', - 'end_time'] - for attr in datetime_attrs: + + + for attr in metadata_attrs: if attr in e1: try: metadata[attr] = dateutil.parser.parse(str(e1[attr][()].decode())) diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py index 1a70400..c50e095 100644 --- a/CDTools/tools/interactions/interactions.py +++ b/CDTools/tools/interactions/interactions.py @@ -627,13 +627,7 @@ def RPI_interaction(probe, obj): 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: - fftobj = t.nn.functional.pad( - 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.') + fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad0l, pad0r)) # Again, just an inverse FFT but with an fftshift upsampled_obj = propagators.inverse_far_field(fftobj) diff --git a/CDTools/tools/plotting/plotting.py b/CDTools/tools/plotting/plotting.py index b5c65a8..bee128c 100644 --- a/CDTools/tools/plotting/plotting.py +++ b/CDTools/tools/plotting/plotting.py @@ -82,7 +82,7 @@ def get_units_factor(units): return factor -def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, **kwargs): +def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, interpolation=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 @@ -112,6 +112,8 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', Default is 'viridis', the colormap to plot with cmap_label : str What to label the colorbar when plotting + interpolation : str + What interpolation to use for imshow \\**kwargs All other args are passed to fig.add_subplot(111, \\**kwargs) @@ -169,7 +171,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', else: extent=None - plt.imshow(to_plot, cmap = cmap, extent = extent) + plt.imshow(to_plot, cmap = cmap, extent = extent, interpolation=interpolation) cbar = plt.colorbar() if cmap_label is not None: cbar.set_label(cmap_label)