From 0d91272766d6604441b9d579efb662be31e19f16 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 15 Feb 2022 15:16:07 -0500 Subject: [PATCH] Create a base version of the models --- .../datasets/polarized_ptycho_2d_dataset.py | 9 +- CDTools/datasets/ptycho_2d_dataset.py | 7 +- CDTools/models/__init__.py | 4 +- CDTools/models/base.py | 1 - CDTools/models/fancy_ptycho.py | 9 +- CDTools/models/polarized_fancy_ptycho.py | 1 + .../time_resolved_ptycho_calibration.py | 504 ++++++++++++++++++ CDTools/models/time_resolved_rpi.py | 302 +++++++++++ CDTools/tools/analysis/analysis.py | 1 - CDTools/tools/data/data.py | 7 +- 10 files changed, 828 insertions(+), 17 deletions(-) create mode 100644 CDTools/models/time_resolved_ptycho_calibration.py create mode 100644 CDTools/models/time_resolved_rpi.py diff --git a/CDTools/datasets/polarized_ptycho_2d_dataset.py b/CDTools/datasets/polarized_ptycho_2d_dataset.py index 1063afc..e34350f 100644 --- a/CDTools/datasets/polarized_ptycho_2d_dataset.py +++ b/CDTools/datasets/polarized_ptycho_2d_dataset.py @@ -115,7 +115,7 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): # It sucks that I can't reuse the base factory method here, # perhaps there is a way but I couldn't figure it out. @classmethod - def from_cxi(cls, cxi_file): + def from_cxi(cls, cxi_file, cut_zeros=True): """Generates a new PolarizedPtycho2DDataset from a .cxi file directly This generates a new PolarizedPtycho2DDataset from a .cxi file storing @@ -125,19 +125,22 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): ---------- file : str, pathlib.Path, or h5py.File The .cxi file to load from + cut_zeros : bool + Default True, whether to set all negative data to zero Returns ------- dataset : PolarizedPtycho2DDataset The constructed dataset object """ + # If a bare string is passed if isinstance(cxi_file, str) or isinstance(cxi_file, pathlib.Path): with h5py.File(cxi_file, 'r') as f: - return cls.from_cxi(f) + return cls.from_cxi(f, cut_zeros=cut_zeros) # Generate a base dataset - dataset = Ptycho2DDataset.from_cxi(cxi_file) + dataset = Ptycho2DDataset.from_cxi(cxi_file, cut_zeros=cut_zeros) # Mutate the class to this subclass (PolarizedPtycho2DDataset) dataset.__class__ = cls diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index 8959163..9cd0d71 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -118,7 +118,7 @@ class Ptycho2DDataset(CDataset): # It sucks that I can't reuse the base factory method here, # perhaps there is a way but I couldn't figure it out. @classmethod - def from_cxi(cls, cxi_file): + def from_cxi(cls, cxi_file, cut_zeros=True): """Generates a new Ptycho2DDataset from a .cxi file directly This generates a new Ptycho2DDataset from a .cxi file storing @@ -128,6 +128,8 @@ class Ptycho2DDataset(CDataset): ---------- file : str, pathlib.Path, or h5py.File The .cxi file to load from + cut_zeros : bool + Default True, whether to set all negative data to zero Returns ------- @@ -145,11 +147,10 @@ class Ptycho2DDataset(CDataset): dataset.__class__ = cls # Load the data that is only relevant for this class - patterns, axes = cdtdata.get_data(cxi_file) + patterns, axes = cdtdata.get_data(cxi_file, cut_zeros=cut_zeros) translations = cdtdata.get_ptycho_translations(cxi_file) # And now re-do the stuff from __init__ dataset.translations = t.tensor(translations, dtype=t.float32) - dataset.patterns = t.as_tensor(patterns) if dataset.patterns.dtype == t.float64: raise NotImplementedError('64-bit floats are not supported and precision will not be retained in reconstructions! Please explicitly convert your data to 32-bit or submit a pull request') diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 90f9e37..5e005a1 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -22,7 +22,7 @@ defining a new ptychography model before attempting to do so. # I don't believe that __all__ really needed, but it's nice to define it # to be explicit that import * is safe -__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'PolarizedFancyPtycho', 'Bragg2DPtycho', 'Multislice2DPtycho', 'RPI'] +__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'PolarizedFancyPtycho', 'Bragg2DPtycho', 'Multislice2DPtycho', 'RPI', 'TimeResolvedPtychoCalibration', 'TimeResolvedRPI'] from CDTools.models.base import CDIModel from CDTools.models.simple_ptycho import SimplePtycho @@ -32,6 +32,8 @@ from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho from CDTools.models.rpi import RPI from CDTools.models.multimode_rpi import MultimodeRPI +from CDTools.models.time_resolved_ptycho_calibration import TimeResolvedPtychoCalibration +from CDTools.models.time_resolved_rpi import TimeResolvedRPI # Still needs to be updated for the new complex numbers #from CDTools.models.s_matrix_ptycho import SMatrixPtycho diff --git a/CDTools/models/base.py b/CDTools/models/base.py index bca2e3a..84d232d 100644 --- a/CDTools/models/base.py +++ b/CDTools/models/base.py @@ -586,7 +586,6 @@ class CDIModel(t.nn.Module): inputs, output = dataset[idx] sim_data = self.forward(*inputs).detach().cpu().numpy() - sim_data = sim_data meas_data = output.detach().cpu().numpy() if hasattr(self, 'mask') and self.mask is not None: mask = self.mask.detach().cpu().numpy() diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index b768012..02b817e 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -67,12 +67,11 @@ class FancyPtycho(CDIModel): if background is None: if detector_slice is not None: - background = 1e-6 * t.ones( - self.probe[0][self.detector_slice].shape, - dtype=t.float32) + shape = self.probe[0][self.detector_slice].shape else: - background = 1e-6 * t.ones(self.probe[0].shape, - dtype=t.float32) + shape = self.probe[0].shape + background = 1e-6 * t.ones([s//oversampling for s in shape], + dtype=t.float32) self.background = t.nn.Parameter(background) diff --git a/CDTools/models/polarized_fancy_ptycho.py b/CDTools/models/polarized_fancy_ptycho.py index 302b3e9..6d307a2 100644 --- a/CDTools/models/polarized_fancy_ptycho.py +++ b/CDTools/models/polarized_fancy_ptycho.py @@ -294,6 +294,7 @@ class PolarizedFancyPtycho(FancyPtycho): else: raise NotImplementedError('Unstable Modes not Implemented for polarized light') + pol_probes = polarization.apply_linear_polarizer(prs, polarizer) exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( diff --git a/CDTools/models/time_resolved_ptycho_calibration.py b/CDTools/models/time_resolved_ptycho_calibration.py new file mode 100644 index 0000000..eecadc4 --- /dev/null +++ b/CDTools/models/time_resolved_ptycho_calibration.py @@ -0,0 +1,504 @@ +import torch as t +from CDTools.models import CDIModel +from CDTools.datasets import Ptycho2DDataset +from CDTools import tools +from CDTools.tools import plotting as p +from CDTools.tools import analysis +from matplotlib import pyplot as plt +from datetime import datetime +import numpy as np +from scipy import linalg as sla +from copy import copy + +# +# Basic points: +# Just one probe mode, no need to overcomplicate things. +# Weights is only a list of numbers, no matrices or anything like that +# Mandatory probe support in Fourier space +# +# When loading from a dataset, we need information on the zone plate geometry +# + +__all__ = ['TimeResolvedPtychoCalibration'] + + +class TimeResolvedPtychoCalibration(CDIModel): + + def __init__(self, wavelength, detector_geometry, + probe_basis, + probe_guess, obj_guess, + fourier_times, probe_fourier_support, + times, time_dependence, frame_delays, + 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, + oversampling=1, + loss='amplitude mse', units='um', + simulate_probe_translation=False): + + super(TimeResolvedPtychoCalibration, self).__init__() + self.wavelength = t.tensor(wavelength) + self.detector_geometry = copy(detector_geometry) + det_geo = self.detector_geometry + if 'distance' in det_geo: + det_geo['distance'] = t.tensor(det_geo['distance'], dtype=t.float32) + if 'basis' in det_geo: + det_geo['basis'] = t.tensor(det_geo['basis'], dtype=t.float32) + if 'corner' in det_geo: + det_geo['corner'] = t.tensor(det_geo['corner'], dtype=t.float32) + + self.min_translation = t.tensor(min_translation) + + self.probe_basis = t.tensor(probe_basis) + self.detector_slice = copy(detector_slice) + self.surface_normal = t.tensor(surface_normal) + + self.saturation = saturation + self.units = units + + 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) + + 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) + + if background is None: + if detector_slice is not None: + background = 1e-6 * t.ones( + self.probe[self.detector_slice].shape, + dtype=t.float32) + else: + background = 1e-6 * t.ones(self.probe[0].shape, + dtype=t.float32) + + self.background = t.nn.Parameter(background) + + if weights is None: + self.weights = None + else: + self.weights = t.nn.Parameter(t.tensor(weights, + dtype=t.float32)) + + if translation_offsets is None: + self.translation_offsets = None + else: + t_o = t.tensor(translation_offsets, dtype=t.float32) + t_o = t_o / translation_scale + self.translation_offsets = t.nn.Parameter(t_o) + + self.translation_scale = translation_scale + + self.probe_fourier_support = probe_fourier_support + self.fourier_times = fourier_times + self.times = times + self.time_dependence = t.nn.Parameter(time_dependence) + self.frame_delays = frame_delays + + self.oversampling = oversampling + + self.simulate_probe_translation = simulate_probe_translation + if simulate_probe_translation: + Is = t.arange(self.probe.shape[-2], dtype=t.float32) + Js = t.arange(self.probe.shape[-1], dtype=t.float32) + Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js)) + self.I_phase = 2 * np.pi* Is + self.J_phase = 2 * np.pi* Js + + # Here we set the appropriate loss function + if (loss.lower().strip() == 'amplitude mse' + or loss.lower().strip() == 'amplitude_mse'): + self.loss = tools.losses.amplitude_mse + elif (loss.lower().strip() == 'poisson nll' + or loss.lower().strip() == 'poisson_nll'): + self.loss = tools.losses.poisson_nll + else: + raise KeyError('Specified loss function not supported') + + + @classmethod + def from_dataset(cls, dataset, zp_geometry, time_window, n_times, n_frames, randomize_ang=0, padding=0, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um', simulate_probe_translation=False): + + 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, *extras), 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 + + # 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) + + 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 + probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling) + + probe = tools.propagators.far_field(probe) + + obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) + + det_geo = dataset.detector_geometry + + translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5) + # 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 + + # What do we know about the zp? + # delta_r, N, and beamstop ratio. It's probably best, though, + # to just read diameter, beamstop_diameter and focal_length directly + # because that is the most general even if the optic isn't truly + # a zone plate. + zp_distance = zp_geometry['focal_length'] + probe_size = probe_basis * t.as_tensor(probe_shape, dtype=t.float32) + pinv_basis = t.tensor(np.linalg.pinv(probe_size).transpose()).to(t.float32) + zone_plate_basis = pinv_basis * wavelength * zp_distance + zone_plate_steps = t.sum(zone_plate_basis,axis=1) + # This may very well mix up x & y and fail on non-square detectors + probe_fourier_support = t.zeros(probe.shape, dtype=t.bool) + xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]] + xs = zone_plate_steps[1] * (xs - np.mean(xs)) + ys = zone_plate_steps[0] * (ys - np.mean(ys)) + Rs = np.sqrt(xs**2 + ys**2) + + distances = np.sqrt(zp_distance**2 + xs**2 + ys**2) + times = (distances - t.min(distances)) / 2.99792e8 + + # This sets the support of the probe and also restricts the + # timing matrix so it only considers times that are actually + # in the window defined by the probe fourier support + probe_fourier_support[Rs < zp_geometry['diameter']/2] = 1 + times[Rs > zp_geometry['diameter']/2] = 0 + times[Rs > zp_geometry['diameter']/2] = t.max(times) + + probe_fourier_support[Rs < zp_geometry['beamstop_diameter']/2] = 0 + times[Rs < zp_geometry['beamstop_diameter']/2] = t.max(times) + times[Rs < zp_geometry['beamstop_diameter']/2] = t.min(times) + + fourier_times = times - t.min(times) + probe = probe * probe_fourier_support + + # This is now the time axis for the probe's envelope + times = t.linspace(0, time_window, n_times+1) + time_dependence = t.ones(n_times, dtype=t.complex64) + frame_delays = t.linspace(0, t.max(fourier_times) + time_window, n_frames+2) + frame_delays -= time_window + frame_delays = frame_delays[1:-1] + + return cls(wavelength, det_geo, probe_basis, probe, obj, + fourier_times, probe_fourier_support, + times, time_dependence, frame_delays, + 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, + oversampling=oversampling, + loss=loss, units=units, + simulate_probe_translation=simulate_probe_translation) + + + def interaction(self, index, translations, *args): + # The *args is included so that this can work even when given, say, + # a polarized ptycho dataset that might spit out more inputs. + + # Step 1 is to convert the translations for each position into a + # value in pixels + pix_trans = tools.interactions.translations_to_pixel( + self.probe_basis, + translations, + surface_normal=self.surface_normal) + pix_trans -= self.min_translation + # We then add on any recovered translation offset, if they exist + if self.translation_offsets is not None: + pix_trans += (self.translation_scale * + self.translation_offsets[index]) + + probes = self.get_probes(space='real') + Ws = self.weights[index] + + # This might not work well + prs = Ws[...,None,None,None] * probes + #prs = t.sum(Ws[..., None, None, None] * probes, axis=-3) + #print(prs.shape) + + if self.simulate_probe_translation: + det_pix_trans = tools.interactions.translations_to_pixel( + self.detector_geometry['basis'], + translations, + surface_normal=self.surface_normal) + + probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] * + self.I_phase[None,...] + + det_pix_trans[:,1,None,None] * + self.J_phase[None,...])) + prs = prs * probe_masks[...,None,:,:] + + + # Now we actually do the interaction, using the sinc subpixel + # translation model as per usual + exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( + prs, self.obj, pix_trans, + shift_probe=True, multiple_modes=True) + return exit_waves + + + def forward_propagator(self, wavefields): + return tools.propagators.far_field(wavefields) + + + def backward_propagator(self, wavefields): + return tools.propagators.inverse_far_field(wavefields) + + + def measurement(self, wavefields): + return tools.measurements.quadratic_background( + wavefields, + self.background, + detector_slice=self.detector_slice, + measurement=tools.measurements.incoherent_sum, + saturation=self.saturation, + oversampling=self.oversampling) + + + # Note: No "loss" function is defined here, because it is added + # dynamically during object creation in __init__ + + def to(self, *args, **kwargs): + super(TimeResolvedPtychoCalibration, self).to(*args, **kwargs) + self.wavelength = self.wavelength.to(*args, **kwargs) + # move the detector geometry too + det_geo = self.detector_geometry + if 'distance' in det_geo: + det_geo['distance'] = det_geo['distance'].to(*args, **kwargs) + if 'basis' in det_geo: + det_geo['basis'] = det_geo['basis'].to(*args, **kwargs) + if 'corner' in det_geo: + det_geo['corner'] = det_geo['corner'].to(*args, **kwargs) + + if self.mask is not None: + self.mask = self.mask.to(*args, **kwargs) + + if self.simulate_probe_translation: + self.I_phase = self.I_phase.to(*args, **kwargs) + self.J_phase = self.J_phase.to(*args, **kwargs) + + self.min_translation = self.min_translation.to(*args, **kwargs) + self.probe_basis = self.probe_basis.to(*args, **kwargs) + self.probe_norm = self.probe_norm.to(*args, **kwargs) + self.probe_fourier_support = self.probe_fourier_support.to(*args, + **kwargs) + self.fourier_times = self.fourier_times.to(*args, **kwargs) + self.times = self.times.to(*args, **kwargs) + self.surface_normal = self.surface_normal.to(*args, **kwargs) + + + def sim_to_dataset(self, args_list, calculation_width=None): + # 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', + 'instrument_n': 'Simulated Data', + 'start_time': datetime.now()} + + surface_normal = self.surface_normal.detach().cpu().numpy() + xsurfacevec = np.cross(np.array([0., 1., 0.]), surface_normal) + xsurfacevec /= np.linalg.norm(xsurfacevec) + 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 + + data = [] + len(indices) + if calculation_width is None: + calculation_width = len(indices) + index_chunks = [indices[i:i + calculation_width] + for i in range(0, len(indices), + calculation_width)] + translation_chunks = [translations[i:i + calculation_width] + for i in range(0, len(indices), + calculation_width)] + + + # Then we simulate the results + data = [self.forward(idx, trans).detach() + for idx, trans in zip(index_chunks, translation_chunks)] + + data = t.cat(data, dim=0) + # And finally, we make the dataset + return Ptycho2DDataset( + translations, data, + entry_info=entry_info, + sample_info=sample_info, + wavelength=wavelength, + detector_geometry=detector_geometry, + mask=mask) + + + def corrected_translations(self, dataset): + translations = dataset.translations.to( + dtype=t.float32, device=self.probe.device) + if (hasattr(self, 'translation_offsets') and + self.translation_offsets is not None): + 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 + else: + return translations + + + def get_probes(self, space='real'): + # This is the part where I need to create the probe modes using the + # time-dependent stuff + + # This restricts the basis probes to stay within the probe support + optic_mask = self.probe * self.probe_fourier_support + + probes = t.zeros([len(self.frame_delays)] + list(self.probe.shape), + dtype=self.probe.dtype, device=self.probe.device) + for i, delay in enumerate(self.frame_delays): + indices = t.bucketize(self.fourier_times, self.times + delay) + clamped_indices = t.clamp(indices-1, max=len(self.time_dependence)-1) + illumination = t.take(self.time_dependence, clamped_indices) + illumination[indices==0] = 0 + illumination[indices==len(self.time_dependence)+1] = 0 + probes[i] = illumination * optic_mask + + if space.lower()=='real': + return tools.propagators.inverse_far_field(probes) + elif space.lower()=='fourier' or space.lower()=='reciprocal': + return probes + + + plot_list = [ + ('Probe Amplitudes (scroll to view modes)', + lambda self, fig: p.plot_amplitude(self.get_probes(space='real'), fig=fig, basis=self.probe_basis, units=self.units)), + ('Probe Phases (scroll to view modes)', + lambda self, fig: p.plot_phase(self.get_probes(space='real'), fig=fig, basis=self.probe_basis, units=self.units)), + ('Fourier Probe Amplitudes (scroll to view modes)', + lambda self, fig: p.plot_amplitude(self.get_probes(space='fourier'), fig=fig, basis=self.probe_basis, units=self.units)), + ('Fourier Probe Phases (scroll to view modes)', + lambda self, fig: p.plot_phase(self.get_probes(space='fourier'), fig=fig, basis=self.probe_basis, units=self.units)), + ('Optic Amplitude', + lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)), + ('Optic Phase', + lambda self, fig: p.plot_phase(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), + ('% 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', + 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)), + ('Corrected Translations', + lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)), + ('Background', + lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)), + ('Time Structure', + lambda self, fig: (plt.figure(fig.number) and (plt.clf() or True) and plt.plot(self.time_dependence.real.detach().cpu().numpy()) and plt.plot(self.time_dependence.imag.detach().cpu().numpy()))) + ] + +# def plot_errors(self, dataset): + + + + def save_results(self, dataset): + basis = self.probe_basis.detach().cpu().numpy() + translations = self.corrected_translations(dataset).detach().cpu().numpy() + optic = self.probe.detach().cpu().numpy() + optic = optic * self.probe_norm.detach().cpu().numpy() + time_dependence = self.time_dependence.detach().cpu().numpy() + times = self.times.detach().cpu().numpy() + fourier_times = self.fourier_times.detach().cpu().numpy() + probes = self.get_probes().detach().cpu().numpy() + obj = self.obj.detach().cpu().numpy() + background = self.background.detach().cpu().numpy()**2 + weights = self.weights.detach().cpu().numpy() + + return {'basis': basis, 'translation': translations, + 'probes': probes, 'optic': optic, + 'times': times, 'time_dependence': time_dependence, + 'fourier_times': fourier_times, + 'obj': obj, + 'background': background, + 'weights': weights, + } diff --git a/CDTools/models/time_resolved_rpi.py b/CDTools/models/time_resolved_rpi.py new file mode 100644 index 0000000..c8b52fb --- /dev/null +++ b/CDTools/models/time_resolved_rpi.py @@ -0,0 +1,302 @@ +import torch as t +from CDTools.models import CDIModel +from CDTools import tools +from CDTools.tools import plotting as p +from CDTools.tools.interactions import RPI_interaction +from CDTools.tools import initializers +from scipy.ndimage.morphology import binary_dilation +import numpy as np +from copy import copy + +__all__ = ['MultimodeRPI'] + + + +__all__ = ['RPI'] + +class TimeResolvedRPI(CDIModel): + + @property + def obj(self): + return t.complex(self.obj_real, self.obj_imag) + + + def __init__(self, wavelength, detector_geometry, probe_basis, + probe, obj_guess, framerate, detector_slice=None, + background=None, mask=None, saturation=None, + obj_support=None, oversampling=1): + + super(TimeResolvedRPI, self).__init__() + + self.wavelength = t.tensor(wavelength) + self.framerate = framerate + self.detector_geometry = copy(detector_geometry) + + det_geo = self.detector_geometry + if hasattr(det_geo, 'distance'): + det_geo['distance'] = t.tensor(det_geo['distance']) + if hasattr(det_geo, 'basis'): + det_geo['basis'] = t.tensor(det_geo['basis']) + if hasattr(det_geo, 'corner'): + det_geo['corner'] = t.tensor(det_geo['corner']) + + self.probe_basis = t.tensor(probe_basis) + + scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1], + probe.shape[-2]/obj_guess.shape[-2]]) + self.obj_basis = self.probe_basis * scale_factor + self.detector_slice = detector_slice + + # Maybe something to include in a bit + # self.surface_normal = t.tensor(surface_normal) + + self.saturation = saturation + + if mask is None: + self.mask = mask + else: + self.mask = t.tensor(mask, dtype=t.bool) + + + self.probe = t.tensor(probe, dtype=t.complex64) + + obj_guess = t.tensor(obj_guess, dtype=t.complex64) + + self.obj_real = t.nn.Parameter(obj_guess.real) + self.obj_imag = t.nn.Parameter(obj_guess.imag) + + # Wait for LBFGS to be updated for complex-valued parameters + # self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + + if background is None: + if detector_slice is not None: + background = 1e-6 * t.ones( + self.probe[0][self.detector_slice].shape, + dtype=t.float32) + else: + background = 1e-6 * t.ones(self.probe[0].shape, + dtype=t.float32) + + self.background = t.tensor(background, dtype=t.float32) + + if obj_support is not None: + self.obj_support = obj_support + self.obj.data = self.obj * obj_support[None, ...] + else: + self.obj_support = t.ones_like(self.obj[0, ...]) + + self.oversampling = oversampling + + + @classmethod + def from_dataset(cls, dataset, probe, framerate, obj_size=None, background=None, mask=None, padding=0, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, probe_threshold=0): + + 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 only need the patterns here, not the inputs associated with them. + _, 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 + + # 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) + + if not isinstance(probe,t.Tensor): + probe = t.as_tensor(probe) + + if background is None and hasattr(dataset, 'background') \ + and dataset.background is not None: + background = t.sqrt(dataset.background) + elif background is not None: + background = t.sqrt(t.Tensor(background).to(dtype=t.float32)) + + det_geo = dataset.detector_geometry + + # If no mask is given, but one exists in the dataset, load it. + if mask is None and hasattr(dataset, 'mask') \ + and dataset.mask is not None: + mask = dataset.mask.to(t.bool) + + # Now we initialize the object + if obj_size is None: + # This is a standard size for a well-matched probe and detector + obj_size = list((np.array(probe_shape) // 2).astype(int)) + + # I think something to do with the fact that the object is defined + # on a coarser grid needs to be accounted for here that is not + # accounted for yet + scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2) + n_modes = (probe.shape[0] - 1) // framerate + 1 + + obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size)) + + + probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0)) + probe_fft = tools.propagators.far_field(probe_intensity) + pad0l = (probe.shape[-2] - obj_size[-2])//2 + pad0r = probe.shape[-2] - obj_size[-2] - pad0l + pad1l = (probe.shape[-1] - obj_size[-1])//2 + pad1r = probe.shape[-1] - obj_size[-1] - pad1l + probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r] + probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft)) + + obj_support = probe_lr > t.max(probe_lr) * probe_threshold + obj_support = t.as_tensor(binary_dilation(obj_support)) + + return cls(wavelength, det_geo, probe_basis, + probe, obj_guess, framerate, detector_slice=det_slice, + background=background, mask=mask, saturation=saturation, + obj_support=obj_support, oversampling=oversampling) + + + def random_init(self, pattern): + scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2) + self.obj.data = scale * t.exp( + 2j * np.pi * t.rand(self.obj.shape)).to( + dtype=self.obj.dtype, device=self.obj.device) + + + # Needs work + def interaction(self, index, *args): + # including *args allows this to work with all sorts of datasets + # that might include other information in with the index in their + # "input" parameters (such as translations for a ptychography dataset). + # This makes it seamless to use such a dataset even though those + # extra arguments will not be used. + + + all_exit_waves = [] + + # Mix the probes with the weight matrix + prs = self.probe + + for i in range(self.probe.shape[0]): + obj_frame = i // self.framerate + pr = prs[i] + exit_waves = RPI_interaction(pr, + self.obj_support * self.obj[obj_frame]) + all_exit_waves.append(exit_waves.unsqueeze(0)) + + # This creates a bunch of modes generated from all possible combos + # of the probe and object modes all strung out along the first index + + output = t.cat(all_exit_waves) + + # If we have multiple indexes input, we unsqueeze and repeat the stack + # of wavefields enough times to simulate each requested index. This + # seems silly, but it enables (for example) one to do a reconstruction + # from a set of diffraction patterns that are all known to be from the + # same object. + try: + # will fail if index has no length, for example when index + # is just an int. In this case, we just do nothing instead + output = output.unsqueeze(0).repeat(1,len(index),1,1,1) + except TypeError: + pass + return output + + + def forward_propagator(self, wavefields): + return tools.propagators.far_field(wavefields) + + + def backward_propagator(self, wavefields): + return tools.propagators.inverse_far_field(wavefields) + + + def measurement(self, wavefields): + # Here I'm taking advantage of an undocumented feature in the + # incoherent_sum measurement function where it will work with + # a 4D wavefield array as well as a 5D array. + return tools.measurements.quadratic_background(wavefields, + self.background, + detector_slice=self.detector_slice, + measurement=tools.measurements.incoherent_sum, + saturation=self.saturation, + oversampling=self.oversampling) + + def loss(self, sim_data, real_data, mask=None): + return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) + #return tools.losses.poisson_nll(real_data, sim_data, mask=mask) + + def regularizer(self, factors): + return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \ + + factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2) + + def to(self, *args, **kwargs): + super(TimeResolvedRPI, self).to(*args, **kwargs) + self.wavelength = self.wavelength.to(*args,**kwargs) + # move the detector geometry too + det_geo = self.detector_geometry + if hasattr(det_geo, 'distance'): + det_geo['distance'] = det_geo['distance'].to(*args,**kwargs) + if hasattr(det_geo, 'basis'): + det_geo['basis'] = det_geo['basis'].to(*args,**kwargs) + if hasattr(det_geo, 'corner'): + det_geo['corner'] = det_geo['corner'].to(*args,**kwargs) + + if self.mask is not None: + self.mask = self.mask.to(*args, **kwargs) + + self.probe = self.probe.to(*args,**kwargs) + self.probe_basis = self.probe_basis.to(*args,**kwargs) + self.obj_basis = self.obj_basis.to(*args,**kwargs) + self.obj_support = self.obj_support.to(*args,**kwargs) + self.background = self.background.to(*args, **kwargs) + + # Maybe include in a bit + #self.surface_normal = self.surface_normal.to(*args, **kwargs) + + def sim_to_dataset(self, args_list): + raise NotImplementedError('No sim to dataset yet, sorry!') + + plot_list = [ + ('Probe Amplitudes', + lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)), + ('Object Amplitudes', + lambda self, fig: p.plot_amplitude(self.obj, fig=fig, + basis=self.obj_basis)), + ('Object Phases', + lambda self, fig: p.plot_phase(self.obj, fig=fig, + basis=self.obj_basis)) + ] + + + def save_results(self, dataset=None, full_obj=False): + # dataset is set as a kwarg here because it isn't needed, but the + # common pattern is to pass a dataset. This makes it okay if one + # continues to use that standard pattern + probe_basis = self.probe_basis.detach().cpu().numpy() + obj_basis = self.obj_basis.detach().cpu().numpy() + probe = self.probe.detach().cpu().numpy() + # Provide the option to save out the subdominant objects or + # just the dominant one + if full_obj: + obj = self.obj.detach().cpu().numpy() + else: + obj = self.obj[0].detach().cpu().numpy() + background = self.background.detach().cpu().numpy()**2 + + return {'probe_basis': probe_basis, 'obj_basis': obj_basis, + 'probe': probe,'obj': obj, + 'background': background} + diff --git a/CDTools/tools/analysis/analysis.py b/CDTools/tools/analysis/analysis.py index feefb11..e28d2d2 100644 --- a/CDTools/tools/analysis/analysis.py +++ b/CDTools/tools/analysis/analysis.py @@ -546,7 +546,6 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'): weights=F2.detach().cpu().numpy()) n_pix, bins = np.histogram(Rs, bins=nbins, range=frc_range) - n_pix = n_pix / 4 # This is for an apodized image, apodized with a hann window frc = np.abs(numerator) / np.sqrt(denominator_F1*denominator_F2) diff --git a/CDTools/tools/data/data.py b/CDTools/tools/data/data.py index ea59ac3..bb97a0f 100644 --- a/CDTools/tools/data/data.py +++ b/CDTools/tools/data/data.py @@ -317,7 +317,7 @@ def get_dark(cxi_file): return darks -def get_data(cxi_file, cut_zeroes = True): +def get_data(cxi_file, cut_zeros = True): """Returns an array with the full stack of detector data defined in the cxi file object This function will make sure to check all the various places that it's @@ -333,6 +333,8 @@ def get_data(cxi_file, cut_zeroes = True): ---------- cxi_file : h5py.File A file object to be read + cut_zeros : bool + Default True, whether to set all negative data to zero Returns ------- @@ -351,9 +353,8 @@ def get_data(cxi_file, cut_zeroes = True): raise KeyError('Data is not defined within cxi file') data = cxi_file[pull_from][:] - # Use maximum in-place to avoid allocating any more memory than is needed - if cut_zeroes: + if cut_zeros: np.maximum(data,0,data) if 'axes' in cxi_file[pull_from].attrs: