diff --git a/src/cdtools/datasets/polarized_ptycho_2d_dataset.py b/src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py similarity index 57% rename from src/cdtools/datasets/polarized_ptycho_2d_dataset.py rename to src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py index ecb3cfc..391fd19 100644 --- a/src/cdtools/datasets/polarized_ptycho_2d_dataset.py +++ b/src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py @@ -1,40 +1,48 @@ -import numpy as np import torch as t import h5py import pathlib from cdtools.datasets import Ptycho2DDataset from cdtools.tools import data as cdtdata -from cdtools.tools import plotting -__all__ = ['PolarizedPtycho2DDataset'] +__all__ = ['PolarizationSweptPtycho2DDataset'] -class PolarizedPtycho2DDataset(Ptycho2DDataset): +class PolarizationSweptPtycho2DDataset(Ptycho2DDataset): """The standard dataset for a 2D ptychography scan - Subclasses datasets.Ptycho2DDataset + Subclasses datasets.CDataset This class loads and saves 2D ptychography scan data from .cxi files. It should save and load files compatible with most reconstruction programs, although it is only tested against SHARP. """ - - def __init__(self, translations, polarizer_angles, analyzer_angles, patterns, axes=None, *args, **kwargs): + def __init__(self, translations, patterns, polarization_indices, + polarization_states, + *args, **kwargs): """The __init__ function allows construction from python objects. The detector_geometry dictionary is defined to have the entries defined by the outputs of data.get_detector_geometry. + Note that the created dataset object will not copy the data in the + patterns parameter in order to avoid doubling the memory requiement + for large datasets. Parameters ---------- - translations : array (nx3 t.tensor when polarized=True) + translations : array An nx3 array containing the probe translations at each scan point patterns : array - An nxmxl array containing the full stack of measured diffraction patterns + An nxmxl array containing the full stack of measured diffraction + patterns + polarization_indices : array(int) + A length-n array containing the index of the polarization state + attached to each pattern + polarization_states : array + An mx2 array, where m is the number of polarization indices, + encoding the polarzation state associated with each index axes : list(str) - A - of names for the axes of the probe translations + A list of names for the axes of the probe translations entry_info : dict A dictionary containing the entry_info metadata sample_info : dict @@ -50,15 +58,16 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): background : array An initial guess for the not-previously-subtracted detector background + intensities : array + A list of measured shot-to-shot intensities """ - super(PolarizedPtycho2DDataset,self).__init__(translations, patterns, - *args, **kwargs) - - - self.polarizer = t.tensor(polarizer_angles, dtype=t.float32) - self.analyzer = t.tensor(analyzer_angles, dtype=t.float32) + super().__init__(translations, patterns, + *args, **kwargs) + self.polarization_indices = t.tensor(polarization_indices, dtype=t.int32) + self.polarization_states = t.tensor(polarization_states, dtype=t.float32) + def _load(self, index): """ Internal function to load data @@ -96,7 +105,7 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): The output pattern or stack of output patterns """ return ((index, self.translations[index], - self.polarizer[index], self.analyzer[index]), + self.polarization_indices[index]), self.patterns[index]) def to(self, *args, **kwargs): @@ -107,19 +116,18 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): Accepts the same parameters as torch.Tensor.to """ - super(PolarizedPtycho2DDataset, self).to(*args, **kwargs) - self.polarizer = self.polarizer.to(*args, **kwargs) - self.analyzer = self.analyzer.to(*args, **kwargs) + super().to(*args, **kwargs) + self.polarization_states = self.polarization_states.to(*args, **kwargs) # 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, cut_zeros=True): - """Generates a new PolarizedPtycho2DDataset from a .cxi file directly + """Generates a new PolarizationSweptPtycho2DDataset from a .cxi file - This generates a new PolarizedPtycho2DDataset from a .cxi file storing - a 2D ptychography scan. + This generates a new PolarizationSweptPtycho2DDataset from a .cxi file + storing the 2D ptychography scan. Parameters ---------- @@ -130,7 +138,7 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): Returns ------- - dataset : PolarizedPtycho2DDataset + dataset : PolarizationSweptPtycho2DDataset The constructed dataset object """ @@ -146,20 +154,20 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): dataset.__class__ = cls # Now, we save out the polarizer and analyzer states - polarizer = cdtdata.get_shot_to_shot_info(cxi_file, 'polarizer_angle') - analyzer = cdtdata.get_shot_to_shot_info(cxi_file, 'analyzer_angle') + polarization_indices = cdtdata.get_shot_to_shot_info(cxi_file, 'polarization_indices') + polarization_states = cdtdata.get_entry_info(cxi_file, 'polarization_states') - dataset.analyzer = t.tensor(analyzer, dtype=t.float32) - dataset.polarizer = t.tensor(polarizer, dtype=t.float32) + dataset.polarization_indices = t.tensor(polarization_indices, dtype=t.int32) + dataset.polarization_states = t.tensor(polarization_states, dtype=t.complex64) return dataset def to_cxi(self, cxi_file, polarized=False): - """Saves out a PolarizedPtycho2DDataset as a .cxi file + """Saves out a PolarizationSweptPtycho2DDataset as a .cxi file This function saves all the compatible information in a - PolarizedPtycho2DDataset object into a .cxi file. This saved .cxi file - should be compatible with any standard .cxi file based + PolarizationSweptPtycho2DDataset object into a .cxi file. This saved + .cxi file should be compatible with any standard .cxi file based reconstruction tool, such as SHARP. Parameters @@ -174,45 +182,10 @@ class PolarizedPtycho2DDataset(Ptycho2DDataset): return self.to_cxi(f) # This saves the translations, patterns, etc. - super(PolarizedPtycho2DDataset, self).to_cxi(cxi_file) - - # Now, we save out the polarizer and analyzer states - cdtdata.add_shot_to_shot_info(cxi_file, self.polarizer, - 'polarizer_angle') - cdtdata.add_shot_to_shot_info(cxi_file, self.analyzer, - 'analyzer_angle') - - def inspect(self, logarithmic=True, units='um'): - """Launches an interactive plot for perusing the data - - This launches an interactive plotting tool in matplotlib that - shows the spatial map constructed from the integrated intensity - at each position on the left, next to a panel on the right that - can display a base-10 log plot of the detector readout at each - position. - """ - - - def get_images(idx): - inputs, output = self[idx] - meas_data = output.detach().cpu().numpy() - if hasattr(self, 'mask') and self.mask is not None: - mask = self.mask.detach().cpu().numpy() - else: - mask = 1 - - if logarithmic: - return np.log(meas_data) / np.log(10) * mask - else: - return meas_data * mask - - translations = self.translations.detach().cpu().numpy() - nanomap_values = (self.mask.to(t.float32) * self.patterns).sum(dim=(1,2)).detach().cpu().numpy() - - if logarithmic: - cbar_title='Log Base 10 of Diffraction Intensity' - else: - 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) + super().to_cxi(cxi_file) + # Now, we save out the polarization states + cdtdata.add_shot_to_shot_info(cxi_file, self.polarization_indices, + 'polarization_indices') + cdtdata.add_entry_info(cxi_file, self.polarization_states, + 'polarization_states') diff --git a/src/cdtools/models/polarization_swept_ptycho.py b/src/cdtools/models/polarization_swept_ptycho.py new file mode 100644 index 0000000..b834c97 --- /dev/null +++ b/src/cdtools/models/polarization_swept_ptycho.py @@ -0,0 +1,610 @@ +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 + +__all__ = ['FancyPtycho'] + +class PolarizationSweptPtycho(CDIModel): + + def __init__(self, wavelength, detector_geometry, + probe_basis, + probe_guess, + obj_guess, + polarization_states, + 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, + fourier_probe=False, + loss='amplitude mse', + units='um', + simulate_probe_translation=False, + simulate_finite_pixels=False, + dtype=t.float32, + obj_view_crop=0 + ): + + super(PolarizationSweptPtycho, self).__init__() + self.register_buffer('wavelength', + t.tensor(wavelength, dtype=dtype)) + self.store_detector_geometry(detector_geometry, + dtype=dtype) + + self.register_buffer('min_translation', + t.tensor(min_translation, dtype=dtype)) + + self.register_buffer('probe_basis', + t.tensor(probe_basis, dtype=dtype)) + + self.detector_slice = copy(detector_slice) + self.register_buffer('surface_normal', + t.tensor(surface_normal, dtype=dtype)) + if saturation is None: + self.saturation = None + else: + self.register_buffer('saturation', + t.tensor(saturation, dtype=dtype)) + # Not sure how to make this a buffer... + self.units = units + + self.fourier_probe = fourier_probe + + if mask is None: + self.mask = None + else: + self.register_buffer('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) + + # We rescale the probe here so it learns at the same rate as the + # object + if probe_guess.dim() > 2: + probe_norm = 1 * t.max(t.abs(probe_guess[0])) + else: + probe_norm = 1 * t.max(t.abs(probe_guess)) + self.register_buffer('probe_norm', probe_norm.to(dtype)) + + self.probe = t.nn.Parameter(probe_guess / self.probe_norm) + self.obj = t.nn.Parameter(obj_guess) + + + self.obj_view_slice = np.s_[obj_view_crop:-obj_view_crop, + obj_view_crop:-obj_view_crop] + + if background is None: + if detector_slice is not None: + dummy_det = t.empty([s//oversampling + for s in self.probe.shape[-2:]]) + shape = dummy_det[self.detector_slice].shape + #shape = self.probe[0][self.detector_slice].shape + else: + shape = [s//oversampling for s in self.probe.shape[-2:]] + background = 1e-6 * t.ones(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.register_buffer('translation_scale', + t.tensor(translation_scale, dtype=dtype)) + + if probe_support is None: + probe_support = t.ones_like(self.probe[0], dtype=t.bool) + self.register_buffer('probe_support', + t.tensor(probe_support, dtype=t.bool)) + + self.oversampling = oversampling + + self.simulate_probe_translation = simulate_probe_translation + + if simulate_probe_translation: + Is = t.arange(self.probe.shape[-2], dtype=dtype) + Js = t.arange(self.probe.shape[-1], dtype=dtype) + Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js)) + + I_phase = 2 * np.pi* Is * self.oversampling + J_phase = 2 * np.pi* Js * self.oversampling + self.register_buffer('I_phase', I_phase) + self.register_buffer('J_phase', J_phase) + + + self.simulate_finite_pixels = simulate_finite_pixels + + self.polarization_states = t.nn.Parameter( + t.tensor(polarization_states, dtype=dtype)) + # by default, don't optimize this, but I think it might be + # interesting to try it because the polarization states + # are not very pure + self.polarization_states.requires_grad = False + + + # 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, + probe_size=None, + randomize_ang=0, + padding=0, + n_modes=1, + translation_scale=1, + saturation=None, + probe_support_radius=None, + probe_fourier_crop=None, + propagation_distance=None, + scattering_mode=None, + oversampling=1, + auto_center=False, + opt_for_fft=False, + fourier_probe=False, + loss='amplitude mse', + units='um', + simulate_probe_translation=False, + simulate_finite_pixels=False, + obj_view_crop=None, + obj_padding=200 + ): + + 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[:] + polarization_states = dataset.polarization_states + + 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=obj_padding) + + 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) + + if probe_fourier_crop is not None: + probe = tools.propagators.far_field(probe) + probe = probe[probe_fourier_crop:-probe_fourier_crop, + probe_fourier_crop:-probe_fourier_crop] + probe = tools.propagators.inverse_far_field(probe) + + # 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) + # The probe gets one extra dimension, so we can simulate one probe + # per polarization state + n_states = polarization_states.shape[0] + probe = t.stack([probe] * n_states, dim=0) + + # Looks like an identity matrix + obj_base = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5)) + obj_top = t.stack([obj_base, obj_base*0], dim=0) + obj_bottom = t.stack([obj_base*0, obj_base], dim=0) + obj = t.stack([obj_top, obj_bottom], dim=0) + + if obj_view_crop is None: + obj_view_crop = min(probe.shape[-2], probe.shape[-1]) // 2 + if obj_view_crop < 0: + obj_view_crop += min(probe.shape[-2], probe.shape[-1]) // 2 + obj_view_crop += obj_padding + + det_geo = dataset.detector_geometry + + translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5) + + Ws = t.ones(len(dataset)) + + if hasattr(dataset, 'intensities') and dataset.intensities is not None: + Ws *= (dataset.intensities.to(dtype=Ws.dtype)[:,...] + / t.mean(dataset.intensities)) + + 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, + polarization_states=polarization_states, + 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, + fourier_probe=fourier_probe, + oversampling=oversampling, + loss=loss, units=units, + simulate_probe_translation=simulate_probe_translation, + simulate_finite_pixels=simulate_finite_pixels, + obj_view_crop=obj_view_crop) + + + def interaction(self, index, translations, polarization_indices, *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]) + # This restricts the basis probes to stay within the probe support + basis_prs = self.probe * self.probe_support[..., :, :] + print(basis_prs.shape) + # Now we expand the polarization states explicitly. + # pol_basis_prs has dimensions n_states x 2 x n_modes x m x l + pol_basis_prs = ( + self.polarization_states[...,None,None,None] + * basis_prs[...,None,:,:,:]) + print(pol_basis_prs.shape) + # For a Fourier-space probe, we take an IFT + if self.fourier_probe: + basis_prs = tools.propagators.inverse_far_field(basis_prs) + + # So, first we get the appropriate probe for each shot. + # prs is now n_frames x 2 x n_modes x m x l + prs = pol_basis_prs[polarization_indices] + print(prs.shape) + # Now we construct the probes for each shot from the basis probes + if self.weights is not None: + Ws = self.weights[index] + # And then we multiply by the weights, along the 0th dimension + prs = Ws[..., None, None, None, None] * basis_prs + + if self.simulate_probe_translation: + det_pix_trans = tools.interactions.translations_to_pixel( + self.det_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,None,:,:] + + # We automatically rescale the probe to match the background size, + # which allows us to do stuff like let the object be super-resolution, + # while restricting the probe to the detector resolution but still + # doing an explicit real-space limitation of the probe + padding = [self.oversampling * self.background.shape[-2] - prs.shape[-2], + self.oversampling * self.background.shape[-1] - prs.shape[-1]] + + if any([p != 0 for p in padding]): # For probe_fourier_crop != 0. + padding = [padding[-1]//2, padding[-1]-padding[-1]//2, + padding[-2]//2, padding[-2]-padding[-2]//2] + prs = tools.propagators.far_field(prs) + prs = t.nn.functional.pad(prs, padding) + prs = tools.propagators.inverse_far_field(prs) + # Now we actually do the interaction, using the sinc subpixel + # translation model as per usual + print('hi') + print(prs[...,0,:,:,:].shape) + print(self.obj[:,0].shape) + exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( + prs[...,0,:,:,:], self.obj[:,0], pix_trans, + shift_probe=True, multiple_modes=True) + exit_waves = exit_waves + ( + self.probe_norm * tools.interactions.ptycho_2D_sinc( + prs[...,1,:,:,:], self.obj[:,1], pix_trans, + shift_probe=True, multiple_modes=True) + ) + print(exit_waves.shape) + # After the object, no point in treating the polarization modes + # any differently from the object/probe modes + print(exit_waves.flatten(start_dim=1,end_dim=-3).shape) + print('sup') + return exit_waves.flatten(start_dim=1,end_dim=-3) + + + 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, + simulate_finite_pixels=self.simulate_finite_pixels) + + + # Note: No "loss" function is defined here, because it is added + # dynamically during object creation in __init__ + + 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} + + + 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=self.get_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_rhos(self): + # If this is the general unified mode model + if self.weights.dim() >= 2: + Ws = self.weights.detach().cpu().numpy() + rhos_out = np.matmul(np.swapaxes(Ws, 1, 2), Ws.conj()) + return rhos_out + # This is the purely incoherent case + 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 + probe = self.probe.detach().cpu().numpy() + ortho_probes = analysis.orthogonalize_probes(probe) + self.probe.data = t.as_tensor( + ortho_probes, + device=self.probe.device, + dtype=self.probe.dtype) + + + plot_list = [ + ('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), + ('% 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[self.obj_view_slice], + fig=fig, + basis=self.probe_basis, + units=self.units)), + ('Object Phase', + lambda self, fig: p.plot_phase( + self.obj[self.obj_view_slice], + 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)) + ] + +# 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() + probe = self.probe.detach().cpu().numpy() + probe = probe * self.probe_norm.detach().cpu().numpy() + obj = self.obj.detach().cpu().numpy() + background = self.background.detach().cpu().numpy()**2 + weights = self.weights.detach().cpu().numpy() + oversampling = self.oversampling + wavelength = self.wavelength.cpu().numpy() + + return {'basis': basis, 'translation': translations, + 'probe': probe, 'obj': obj, + 'background': background, + 'oversampling': oversampling, + 'weights': weights, 'wavelength': wavelength}