From a9e0157707338cd747fd4b1a3af0062e86fd20c3 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Fri, 1 Dec 2023 15:58:18 +0100 Subject: [PATCH] Get tests to past and removeopt_for_fft --- docs/source/tutorial.rst | 3 +- examples/simple_ptycho_model.py | 3 +- src/cdtools/datasets/__init__.py | 2 +- .../polarization_swept_ptycho_2d_dataset.py | 19 +- src/cdtools/models/base.py | 2 +- src/cdtools/models/bragg_2d_ptycho.py | 3 +- src/cdtools/models/fancy_ptycho.py | 2 - src/cdtools/models/fastccd_ptycho.py | 2 - src/cdtools/models/multimode_rpi.py | 9 +- src/cdtools/models/multislice_2d_ptycho.py | 1 - .../models/polarization_swept_ptycho.py | 68 +++---- src/cdtools/models/polarized_fancy_ptycho.py | 6 +- src/cdtools/models/rpi.py | 5 +- src/cdtools/models/s_matrix_ptycho.py | 1 - .../time_resolved_ptycho_calibration.py | 3 +- src/cdtools/models/time_resolved_rpi.py | 10 +- src/cdtools/tools/analysis/analysis.py | 2 - src/cdtools/tools/data/data.py | 6 + .../tools/initializers/initializers.py | 14 +- src/cdtools/tools/propagators/propagators.py | 8 +- tests/test_datasets.py | 180 ------------------ tests/tools/test_analysis.py | 10 +- tests/tools/test_image_processing.py | 8 +- tests/tools/test_initializers.py | 19 +- tests/tools/test_plotting.py | 8 +- tests/tools/test_propagators.py | 10 +- 26 files changed, 96 insertions(+), 308 deletions(-) diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst index 1397cab..1297c2e 100644 --- a/docs/source/tutorial.rst +++ b/docs/source/tutorial.rst @@ -358,8 +358,7 @@ To initialize the object from a dataset, we need to start by extracting the rele probe_basis, probe_shape, det_slice = ewg(det_basis, det_shape, wavelength, - distance, - opt_for_fft=False) + distance) # Next generate the object geometry pix_translations = tools.interactions.translations_to_pixel( diff --git a/examples/simple_ptycho_model.py b/examples/simple_ptycho_model.py index 14d97da..8f96984 100644 --- a/examples/simple_ptycho_model.py +++ b/examples/simple_ptycho_model.py @@ -39,8 +39,7 @@ class SimplePtycho(CDIModel): probe_basis, probe_shape, det_slice = ewg(det_basis, det_shape, wavelength, - distance, - opt_for_fft=False) + distance) # Next generate the object geometry pix_translations = tools.interactions.translations_to_pixel( diff --git a/src/cdtools/datasets/__init__.py b/src/cdtools/datasets/__init__.py index b5ab33d..afcedf3 100644 --- a/src/cdtools/datasets/__init__.py +++ b/src/cdtools/datasets/__init__.py @@ -29,7 +29,7 @@ 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__ = ['CDataset','Ptycho2DDataset','PolarizationSwepttycho2DDataset'] +__all__ = ['CDataset','Ptycho2DDataset','PolarizationSweptPtycho2DDataset'] from cdtools.datasets.base import CDataset from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset diff --git a/src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py b/src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py index 391fd19..32bff52 100644 --- a/src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py +++ b/src/cdtools/datasets/polarization_swept_ptycho_2d_dataset.py @@ -66,7 +66,7 @@ class PolarizationSweptPtycho2DDataset(Ptycho2DDataset): *args, **kwargs) self.polarization_indices = t.tensor(polarization_indices, dtype=t.int32) - self.polarization_states = t.tensor(polarization_states, dtype=t.float32) + self.polarization_states = t.tensor(polarization_states, dtype=t.complex64) def _load(self, index): @@ -154,15 +154,18 @@ class PolarizationSweptPtycho2DDataset(Ptycho2DDataset): dataset.__class__ = cls # Now, we save out the polarizer and analyzer states - polarization_indices = cdtdata.get_shot_to_shot_info(cxi_file, 'polarization_indices') - polarization_states = cdtdata.get_entry_info(cxi_file, 'polarization_states') + polarization_indices = cdtdata.get_shot_to_shot_info( + cxi_file, 'polarization_indices') + polarization_states = dataset.entry_info['polarization_states'] - dataset.polarization_indices = t.tensor(polarization_indices, dtype=t.int32) - dataset.polarization_states = t.tensor(polarization_states, dtype=t.complex64) + 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): + def to_cxi(self, cxi_file): """Saves out a PolarizationSweptPtycho2DDataset as a .cxi file This function saves all the compatible information in a @@ -187,5 +190,5 @@ class PolarizationSweptPtycho2DDataset(Ptycho2DDataset): # 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') + cdtdata.add_entry_info(cxi_file, {'polarization_states': + self.polarization_states}) diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index bbe2476..b143c40 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -284,7 +284,7 @@ class CDIModel(t.nn.Module): # If we don't want to run in a different thread, this is easy if not thread: for it in range(iterations): - yield run_iteration() + yield run_epoch() # But if we do want to thread, it's annoying: else: diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index db9093c..7ff602b 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -209,7 +209,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): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -236,7 +236,6 @@ class Bragg2DPtycho(CDIModel): distance, center=center, padding=padding, - opt_for_fft=opt_for_fft, oversampling=oversampling) # now we grab the sample surface normal if hasattr(dataset, 'sample_info') and \ diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 9d883c9..5a3384e 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -177,7 +177,6 @@ class FancyPtycho(CDIModel): scattering_mode=None, oversampling=1, auto_center=False, - opt_for_fft=False, fourier_probe=False, loss='amplitude mse', units='um', @@ -216,7 +215,6 @@ class FancyPtycho(CDIModel): distance, center=center, padding=padding, - opt_for_fft=opt_for_fft, oversampling=oversampling) if hasattr(dataset, 'sample_info') and \ diff --git a/src/cdtools/models/fastccd_ptycho.py b/src/cdtools/models/fastccd_ptycho.py index 88eca59..a887bd6 100644 --- a/src/cdtools/models/fastccd_ptycho.py +++ b/src/cdtools/models/fastccd_ptycho.py @@ -160,7 +160,6 @@ class FastCCDPtycho(CDIModel): scattering_mode=None, oversampling=1, auto_center=False, - opt_for_fft=False, fourier_probe=False, loss='amplitude mse', units='um', @@ -197,7 +196,6 @@ class FastCCDPtycho(CDIModel): distance, center=center, padding=padding, - opt_for_fft=opt_for_fft, oversampling=oversampling) if hasattr(dataset, 'sample_info') and \ diff --git a/src/cdtools/models/multimode_rpi.py b/src/cdtools/models/multimode_rpi.py index 5e6ab55..c022154 100644 --- a/src/cdtools/models/multimode_rpi.py +++ b/src/cdtools/models/multimode_rpi.py @@ -4,16 +4,12 @@ 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 +from scipy.ndimage import binary_dilation import numpy as np from copy import copy __all__ = ['MultimodeRPI'] - - -__all__ = ['RPI'] - class MultimodeRPI(CDIModel): @property @@ -99,7 +95,7 @@ class MultimodeRPI(CDIModel): @classmethod - def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0): + def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', weight_matrix=False, probe_threshold=0): raise NotImplementedError() wavelength = dataset.wavelength @@ -128,7 +124,6 @@ class MultimodeRPI(CDIModel): distance, center=center, padding=padding, - opt_for_fft=opt_for_fft, oversampling=oversampling) if not isinstance(probe,t.Tensor): diff --git a/src/cdtools/models/multislice_2d_ptycho.py b/src/cdtools/models/multislice_2d_ptycho.py index 7100ad0..b1455f6 100644 --- a/src/cdtools/models/multislice_2d_ptycho.py +++ b/src/cdtools/models/multislice_2d_ptycho.py @@ -181,7 +181,6 @@ class Multislice2DPtycho(CDIModel): distance, center=center, padding=padding, - opt_for_fft=False, oversampling=oversampling) diff --git a/src/cdtools/models/polarization_swept_ptycho.py b/src/cdtools/models/polarization_swept_ptycho.py index b834c97..7deb572 100644 --- a/src/cdtools/models/polarization_swept_ptycho.py +++ b/src/cdtools/models/polarization_swept_ptycho.py @@ -117,7 +117,7 @@ class PolarizationSweptPtycho(CDIModel): t.tensor(translation_scale, dtype=dtype)) if probe_support is None: - probe_support = t.ones_like(self.probe[0], dtype=t.bool) + probe_support = t.ones(self.probe.shape[-2:], dtype=t.bool) self.register_buffer('probe_support', t.tensor(probe_support, dtype=t.bool)) @@ -139,7 +139,7 @@ class PolarizationSweptPtycho(CDIModel): self.simulate_finite_pixels = simulate_finite_pixels self.polarization_states = t.nn.Parameter( - t.tensor(polarization_states, dtype=dtype)) + t.tensor(polarization_states, dtype=t.complex64)) # by default, don't optimize this, but I think it might be # interesting to try it because the polarization states # are not very pure @@ -172,7 +172,6 @@ class PolarizationSweptPtycho(CDIModel): scattering_mode=None, oversampling=1, auto_center=False, - opt_for_fft=False, fourier_probe=False, loss='amplitude mse', units='um', @@ -212,7 +211,6 @@ class PolarizationSweptPtycho(CDIModel): distance, center=center, padding=padding, - opt_for_fft=opt_for_fft, oversampling=oversampling) if hasattr(dataset, 'sample_info') and \ @@ -298,7 +296,7 @@ class PolarizationSweptPtycho(CDIModel): mask = None if probe_support_radius is not None: - probe_support = t.zeros(probe[0].shape, dtype=t.bool) + probe_support = t.zeros(probe.shape[-2:], dtype=t.bool) xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]] xs = xs - np.mean(xs) ys = ys - np.mean(ys) @@ -344,15 +342,14 @@ class PolarizationSweptPtycho(CDIModel): 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) @@ -360,13 +357,13 @@ class PolarizationSweptPtycho(CDIModel): # 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 - + prs = Ws[..., None, None, None, None] * prs#basis_prs + if self.simulate_probe_translation: det_pix_trans = tools.interactions.translations_to_pixel( self.det_basis, @@ -394,23 +391,28 @@ class PolarizationSweptPtycho(CDIModel): 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) + # So, this stacks up all the output modes in lexicographic order + exit_waves = t.cat( + [self.probe_norm * tools.interactions.ptycho_2D_sinc( + prs[...,0,:,:,:], self.obj[0,0], pix_trans, + shift_probe=True, multiple_modes=True), + self.probe_norm * tools.interactions.ptycho_2D_sinc( + prs[...,0,:,:,:], self.obj[1,0], pix_trans, + shift_probe=True, multiple_modes=True)], + dim=1) + exit_waves = exit_waves + t.cat( + [self.probe_norm * tools.interactions.ptycho_2D_sinc( + prs[...,1,:,:,:], self.obj[0,1], pix_trans, + shift_probe=True, multiple_modes=True), + self.probe_norm * tools.interactions.ptycho_2D_sinc( + prs[...,1,:,:,:], self.obj[1,1], pix_trans, + shift_probe=True, multiple_modes=True)], + dim=1) + + return exit_waves def forward_propagator(self, wavefields): @@ -519,14 +521,12 @@ class PolarizationSweptPtycho(CDIModel): 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) + ortho_probes = np.empty(probe.shape, dtype=probe.dtype) + for idx in range(probe.shape[0]): + ortho_probes[idx] = analysis.orthogonalize_probes(probe[idx]) + self.probe.data = t.as_tensor( ortho_probes, device=self.probe.device, @@ -572,13 +572,13 @@ class PolarizationSweptPtycho(CDIModel): lambda self: len(self.weights.shape) >= 2), ('Object Amplitude', lambda self, fig: p.plot_amplitude( - self.obj[self.obj_view_slice], + self.obj[np.s_[:,:] + 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], + self.obj[np.s_[:,:] + self.obj_view_slice], fig=fig, basis=self.probe_basis, units=self.units)), diff --git a/src/cdtools/models/polarized_fancy_ptycho.py b/src/cdtools/models/polarized_fancy_ptycho.py index 4d309fa..af6999f 100644 --- a/src/cdtools/models/polarized_fancy_ptycho.py +++ b/src/cdtools/models/polarized_fancy_ptycho.py @@ -61,7 +61,7 @@ class PolarizedFancyPtycho(FancyPtycho): 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): + 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, loss='amplitude mse', units='um', left_polarized=True): # When using this method, remember to pass through the inputs model = FancyPtycho.from_dataset( @@ -78,7 +78,6 @@ class PolarizedFancyPtycho(FancyPtycho): scattering_mode=scattering_mode, oversampling=oversampling, auto_center=auto_center, - opt_for_fft=opt_for_fft, loss=loss, units=units) @@ -122,7 +121,7 @@ class PolarizedFancyPtycho(FancyPtycho): polarizers = [tools.polarization.generate_linear_polarizer(i * 45) for i in range(3)] @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'): + 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, loss='amplitude mse', units='um'): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -159,7 +158,6 @@ class PolarizedFancyPtycho(FancyPtycho): distance, center=center, padding=padding, - opt_for_fft=opt_for_fft, oversampling=oversampling) probe_shape = t.stack((2, probe_shape), dim=-3) diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index 0d588a6..b2246aa 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -4,7 +4,7 @@ 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 +from scipy.ndimage import binary_dilation import numpy as np from copy import copy import time @@ -139,7 +139,7 @@ class RPI(CDIModel): @classmethod - def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0, simulate_finite_pixels=False): + def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', weight_matrix=False, probe_threshold=0, simulate_finite_pixels=False): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -167,7 +167,6 @@ class RPI(CDIModel): distance, center=center, padding=padding, - opt_for_fft=opt_for_fft, oversampling=oversampling) if not isinstance(probe,t.Tensor): diff --git a/src/cdtools/models/s_matrix_ptycho.py b/src/cdtools/models/s_matrix_ptycho.py index 5926be4..adb7b46 100644 --- a/src/cdtools/models/s_matrix_ptycho.py +++ b/src/cdtools/models/s_matrix_ptycho.py @@ -140,7 +140,6 @@ class SMatrixPtycho(CDIModel): wavelength, distance, padding=padding, - opt_for_fft=False, oversampling=oversampling) if propagation_distances is None: diff --git a/src/cdtools/models/time_resolved_ptycho_calibration.py b/src/cdtools/models/time_resolved_ptycho_calibration.py index a071892..1461050 100644 --- a/src/cdtools/models/time_resolved_ptycho_calibration.py +++ b/src/cdtools/models/time_resolved_ptycho_calibration.py @@ -125,7 +125,7 @@ class TimeResolvedPtychoCalibration(CDIModel): @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): + 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, loss='amplitude mse', units='um', simulate_probe_translation=False): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -156,7 +156,6 @@ class TimeResolvedPtychoCalibration(CDIModel): distance, center=center, padding=padding, - opt_for_fft=opt_for_fft, oversampling=oversampling) if hasattr(dataset, 'sample_info') and \ diff --git a/src/cdtools/models/time_resolved_rpi.py b/src/cdtools/models/time_resolved_rpi.py index 265053a..146826b 100644 --- a/src/cdtools/models/time_resolved_rpi.py +++ b/src/cdtools/models/time_resolved_rpi.py @@ -4,15 +4,12 @@ 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 +from scipy.ndimage import binary_dilation import numpy as np from copy import copy -__all__ = ['MultimodeRPI'] - - -__all__ = ['RPI'] +__all__ = ['TimeResolvedRPI'] class TimeResolvedRPI(CDIModel): @@ -89,7 +86,7 @@ class TimeResolvedRPI(CDIModel): @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): + 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', probe_threshold=0): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -117,7 +114,6 @@ class TimeResolvedRPI(CDIModel): distance, center=center, padding=padding, - opt_for_fft=opt_for_fft, oversampling=oversampling) if not isinstance(probe,t.Tensor): diff --git a/src/cdtools/tools/analysis/analysis.py b/src/cdtools/tools/analysis/analysis.py index 984312b..5d61874 100644 --- a/src/cdtools/tools/analysis/analysis.py +++ b/src/cdtools/tools/analysis/analysis.py @@ -764,11 +764,9 @@ def calc_fidelity(fields_1, fields_2, dims=2): fields_1 = t.as_tensor(fields_1) fields_2 = t.as_tensor(fields_2) - mult = fields_1.unsqueeze(-dims-2) * fields_2.unsqueeze(-dims-1).conj() sumdims = tuple(d - dims for d in range(dims)) mat = t.sum(mult,dim=sumdims) - # Because I think this is the nuclear norm squared, I would like to swap # Out the definition for this, but I need to test it before swapping. # It also probably makes sense to implement sqrt_fidelity separately diff --git a/src/cdtools/tools/data/data.py b/src/cdtools/tools/data/data.py index 105c8a9..e1e0e70 100644 --- a/src/cdtools/tools/data/data.py +++ b/src/cdtools/tools/data/data.py @@ -76,6 +76,12 @@ def get_entry_info(cxi_file): metadata[attr] = dateutil.parser.parse(str(e1[attr][()].decode())) except ValueError: metadata[attr] = str(e1[attr][()].decode()) + + array_attrs = ['polarization_states'] + for attr in array_attrs: + if attr in e1: + metadata[attr] = np.array(e1[attr]) + return metadata diff --git a/src/cdtools/tools/initializers/initializers.py b/src/cdtools/tools/initializers/initializers.py index 367bcfc..860d1cd 100644 --- a/src/cdtools/tools/initializers/initializers.py +++ b/src/cdtools/tools/initializers/initializers.py @@ -21,12 +21,7 @@ __all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian', 'RPI_spectral_init', 'generate_subdominant_modes'] -# I need to remove a dependency on scipy.fftpack, so I did this. -# I should replace this with a real function -def next_fast_len(x): - return x - -def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0, oversampling=1): +def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, padding=0, oversampling=1): """Returns an exit wave basis and shape, as well as a detector slice for the given detector geometry It takes in the parameters for a given detector - the basis defining @@ -48,8 +43,6 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, The sample-detector distance, in m center : torch.Tensor If defined, the location of the zero frequency pixel - opt_for_fft : bool - Default is true, whether to increase detector size to improve fft performance padding : int Default is 0, the size of an extra border of nonphysical pixels around the detector oversampling : int @@ -85,11 +78,6 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, # In some edge cases this shape can be smaller than the detector shape full_shape = t.max(full_shape, det_shape) - - - if opt_for_fft: - full_shape = t.as_tensor([next_fast_len(dim) for dim in full_shape], - dtype=t.int32) # Then, generate a slice that pops the actual detector from the full # detector shape diff --git a/src/cdtools/tools/propagators/propagators.py b/src/cdtools/tools/propagators/propagators.py index 0f8ffa4..5d10fba 100644 --- a/src/cdtools/tools/propagators/propagators.py +++ b/src/cdtools/tools/propagators/propagators.py @@ -540,6 +540,12 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o # adding to get the full 3D wavevectors. Note that we convert to complex # before the square root to appropriately map negative numbers to # complex frequencies + + # The correct algorithm for sqrt(1-x**2) - 1 is + # - x**2 / (sqrt(1-x**2) + 1) + # This one doesn't have numerical stability issues near x=0. + # I need to figure out how to apply that knowledge here. + K = K_ip + perpendicular_dir[:,None,None] \ * t.sqrt(t.complex(K_oop_squared,t.zeros_like(K_oop_squared))) @@ -558,7 +564,7 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o K_m_K_0 = K - K_0[:,None,None] # We actually calculate the phase mask phase_mask = t.tensordot(offset_vector,K_m_K_0, dims=1) - + # And we take the conjugate if needed, which makes sure that the # imaginary part is always positive (supressing evanescent waves) if propagation_vector is not None \ diff --git a/tests/test_datasets.py b/tests/test_datasets.py index f4441f7..d1b5f82 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -276,183 +276,3 @@ def test_Ptycho2DDataset_get_as(ptycho_cxi_1): assert t.allclose(pattern.to(device='cpu'), t.tensor(expected['data'][3,:,:])) - -# -# And we now test the derived class for polarization-dependent ptycho -# - - -def test_PolarizedPtycho2DDataset_init(): - entry_info = {'start_time': datetime.datetime.now(), - 'title' : 'A simple test'} - sample_info = {'name': 'A test sample', - 'mass' : 3.4, - 'unit_cell' : np.array([1,1,1,87,84.5,90])} - wavelength = 1e-9 - detector_geometry = {'distance': 0.7, - 'basis': np.array([[0,-30e-6,0], - [-20e-6,0,0]]).transpose(), - 'corner': np.array((2550e-6,3825e-6,0.3))} - mask = np.ones((256,256)) - patterns = np.random.rand(20,256,256) - translations = np.random.rand(20,3) - analyzer = np.random.rand(20) - polarizer = np.random.rand(20) - - dataset = PolarizedPtycho2DDataset(translations, polarizer, - analyzer, patterns, - entry_info=entry_info, - sample_info=sample_info, - wavelength=wavelength, - detector_geometry=detector_geometry, - mask=mask) - - assert t.all(t.eq(dataset.mask,t.BoolTensor(mask))) - assert dataset.entry_info == entry_info - assert dataset.sample_info == sample_info - assert dataset.wavelength == wavelength - assert dataset.detector_geometry == detector_geometry - assert t.allclose(dataset.patterns, t.as_tensor(patterns)) - assert t.allclose(dataset.translations, t.as_tensor(translations)) - assert t.allclose(dataset.analyzer, t.as_tensor(analyzer, dtype=t.float32)) - assert t.allclose(dataset.polarizer, t.as_tensor(polarizer, dtype=t.float32)) - - -def test_PolarizedPtycho2DDataset_from_cxi(polarized_ptycho_cxi): - cxi, expected = polarized_ptycho_cxi - dataset = PolarizedPtycho2DDataset.from_cxi(cxi) - - # The entry metadata loaded - for key in expected['entry metadata']: - assert dataset.entry_info[key] == expected['entry metadata'][key] - - # Don't test for fidelity since this is tested in the data, just test - # that it is loaded - if expected['sample info'] is None: - assert dataset.sample_info is None - else: - assert dataset.sample_info is not None - - assert np.isclose(dataset.wavelength,expected['wavelength']) - - # Just check one of the loaded attributes - assert np.isclose(dataset.detector_geometry['distance'], - expected['detector']['distance']) - # Check that the other ones are loaded but not for fidelity - assert 'basis' in dataset.detector_geometry - if expected['detector']['corner'] is not None: - assert 'corner' in dataset.detector_geometry - - if expected['mask'] is not None: - assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask)) - - if expected['dark'] is not None: - assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32), - dataset.background)) - - - assert t.allclose(t.tensor(expected['data']),dataset.patterns) - assert t.allclose(t.tensor(expected['translations']),dataset.translations) - assert t.allclose(t.tensor(expected['analyzer_angle']),dataset.analyzer) - assert t.allclose(t.tensor(expected['polarizer_angle']),dataset.polarizer) - - -def test_PolarizedPtycho2DDataset_to_cxi(polarized_ptycho_cxi, tmp_path): - cxi, expected = polarized_ptycho_cxi - print('loading dataset') - dataset = PolarizedPtycho2DDataset.from_cxi(cxi) - print('dataset mask is type', dataset.mask.dtype) - with cdtdata.create_cxi(tmp_path / 'test_PolarizedPtycho2DDataset_to_cxi.cxi') as f: - dataset.to_cxi(f) - - # Now we have to check that all the stuff was written - with h5py.File(tmp_path / 'test_PolarizedPtycho2DDataset_to_cxi.cxi', 'r') as f: - read_dataset = PolarizedPtycho2DDataset.from_cxi(f) - - assert dataset.entry_info == read_dataset.entry_info - - if dataset.sample_info is None: - assert read_dataset.sample_info is None - else: - assert read_dataset.sample_info is not None - - assert np.isclose(dataset.wavelength, read_dataset.wavelength) - - - # Just check one of the loaded attributes - assert np.isclose(dataset.detector_geometry['distance'], - read_dataset.detector_geometry['distance']) - # Check that the other ones are loaded but not for fidelity - assert 'basis' in read_dataset.detector_geometry - if dataset.detector_geometry['corner'] is not None: - assert 'corner' in read_dataset.detector_geometry - - - - if dataset.mask is not None: - assert t.all(t.eq(dataset.mask,read_dataset.mask)) - - if dataset.background is not None: - assert t.all(t.eq(dataset.background, read_dataset.background)) - - assert t.allclose(dataset.patterns, read_dataset.patterns) - assert t.allclose(dataset.translations, read_dataset.translations) - assert t.allclose(dataset.analyzer, read_dataset.analyzer) - assert t.allclose(dataset.analyzer, read_dataset.analyzer) - - -def test_PolarizedPtycho2DDataset_to(polarized_ptycho_cxi): - dataset = PolarizedPtycho2DDataset.from_cxi(polarized_ptycho_cxi[0]) - - dataset.to(dtype=t.float64) - assert dataset.mask.dtype == t.bool - assert dataset.patterns.dtype == t.float64 - assert dataset.translations.dtype == t.float64 - assert dataset.analyzer.dtype == t.float64 - assert dataset.polarizer.dtype == t.float64 - # If cuda is available, check that moving the mask to CUDA works. - if t.cuda.is_available(): - dataset.to(device='cuda:0') - assert dataset.mask.device == t.device('cuda:0') - assert dataset.background.device == t.device('cuda:0') - assert dataset.patterns.device == t.device('cuda:0') - assert dataset.translations.device == t.device('cuda:0') - assert dataset.analyzer.device == t.device('cuda:0') - assert dataset.polarizer.device == t.device('cuda:0') - - - -def test_PolarizedPtycho2DDataset_ops(polarized_ptycho_cxi): - cxi, expected = polarized_ptycho_cxi - dataset = PolarizedPtycho2DDataset.from_cxi(cxi) - dataset.get_as('cpu') - - assert len(dataset) == expected['data'].shape[0] - (idx, translation, polarizer, analyzer), pattern = dataset[3] - assert idx == 3 - assert t.allclose(translation, t.tensor(expected['translations'][3,:])) - assert t.allclose(pattern, t.tensor(expected['data'][3,:,:])) - assert t.allclose(analyzer, t.tensor(expected['analyzer_angle'][3])) - assert t.allclose(polarizer, t.tensor(expected['polarizer_angle'][3])) - - -def test_PolarizedPtycho2DDataset_get_as(polarized_ptycho_cxi): - cxi, expected = polarized_ptycho_cxi - dataset = PolarizedPtycho2DDataset.from_cxi(cxi) - if t.cuda.is_available(): - dataset.get_as('cuda:0') - assert len(dataset) == expected['data'].shape[0] - - (idx, translation, polarizer, analyzer), pattern = dataset[3] - assert str(translation.device) == 'cuda:0' - assert str(pattern.device) == 'cuda:0' - - assert idx == 3 - assert t.allclose(translation.to(device='cpu'), - t.tensor(expected['translations'][3,:])) - assert t.allclose(pattern.to(device='cpu'), - t.tensor(expected['data'][3,:,:])) - assert t.allclose(polarizer.to(device='cpu'), - t.tensor(expected['polarizer_angle'][3])) - assert t.allclose(analyzer.to(device='cpu'), - t.tensor(expected['analyzer_angle'][3])) diff --git a/tests/tools/test_analysis.py b/tests/tools/test_analysis.py index 15f086a..e367483 100644 --- a/tests/tools/test_analysis.py +++ b/tests/tools/test_analysis.py @@ -367,18 +367,20 @@ def test_calc_fidelity(): dm_1 = t.tensordot(dm_1.transpose(0,1), dm_1.conj(), dims=1).numpy() dm_2 = t.reshape(fields_2, (3,-1)) dm_2 = t.tensordot(dm_2.transpose(0,1), dm_2.conj(), dims=1).numpy() - - inner_mat = la.sqrtm(np.dot(np.dot(la.sqrtm(dm_1),dm_2),la.sqrtm(dm_1))) + + sqrt_dm_1 = la.sqrtm(dm_1).astype(dm_1.dtype) + inner_mat = la.sqrtm(np.dot(np.dot(sqrt_dm_1,dm_2), sqrt_dm_1)) + inner_mat = inner_mat.astype(dm_1.dtype) #la.sqrtm doubles the precision fidelity = t.as_tensor(np.abs(np.trace(inner_mat))**2) - + assert t.isclose(fidelity, analysis.calc_fidelity(fields_1, fields_2)) + # Check that it reduces to the overlap for coherent fields fields_1 = t.rand(1,30,17, dtype=t.complex128) fields_2 = t.rand(1,30,17, dtype=t.complex128) assert t.isclose(t.abs(t.sum(fields_1*fields_2.conj()))**2, analysis.calc_fidelity(fields_1, fields_2)) - # Checking that it works with extra dimensions fields_1 = t.rand(3,3,30,17, dtype=t.complex128) fields_2 = t.rand(3,1,30,17, dtype=t.complex128) diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py index 0929003..57795fb 100644 --- a/tests/tools/test_image_processing.py +++ b/tests/tools/test_image_processing.py @@ -7,13 +7,13 @@ from scipy import ndimage def test_centroid(): # Test single im im = t.rand((30,40)) - sp_centroid = ndimage.measurements.center_of_mass(im.numpy()) + sp_centroid = ndimage.center_of_mass(im.numpy()) centroid = image_processing.centroid(im) assert t.allclose(centroid, t.Tensor(sp_centroid)) # Test stack o' ims ims = t.rand((5,30,40)) - sp_centroids = [ndimage.measurements.center_of_mass(im.numpy()) + sp_centroids = [ndimage.center_of_mass(im.numpy()) for im in ims] centroids = image_processing.centroid(ims) assert t.allclose(centroids, t.Tensor(sp_centroids)) @@ -22,14 +22,14 @@ def test_centroid(): def test_centroid_sq(): # Test single im im = t.rand((30,40)) - sp_centroid = ndimage.measurements.center_of_mass(im.numpy()**2) + sp_centroid = ndimage.center_of_mass(im.numpy()**2) centroid = image_processing.centroid_sq(im) assert t.allclose(centroid, t.Tensor(sp_centroid)) # Test complex with multiple ims ims = t.rand((5,30,40)) + 1j * t.rand((5,30,40)) np_ims = ims.numpy() - sp_centroids = [ndimage.measurements.center_of_mass(np.abs(im)**2) + sp_centroids = [ndimage.center_of_mass(np.abs(im)**2) for im in np_ims] centroids = image_processing.centroid_sq(ims, comp=True) assert t.allclose(centroids, t.Tensor(np.array(sp_centroids))) diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index 907ebd9..d4e9495 100644 --- a/tests/tools/test_initializers.py +++ b/tests/tools/test_initializers.py @@ -12,28 +12,16 @@ def test_exit_wave_geometry(): distance = 1. rs_basis, full_shape, det_slice = \ initializers.exit_wave_geometry(basis, shape, wavelength, - distance, opt_for_fft=False) + distance) assert full_shape == shape assert t.ones(full_shape)[det_slice].shape == shape assert t.allclose(rs_basis[0,1],t.Tensor([-8.928571428571428e-07])) assert t.allclose(rs_basis[1,0],t.Tensor([-4.5662100456621004e-07])) - - # Then test it's expanding functionality for a non-optimal array - rs_basis, full_shape, det_slice = \ - initializers.exit_wave_geometry(basis, shape, wavelength, - distance, opt_for_fft=True) - exp_shape = t.Size([75,60]) - assert full_shape == exp_shape - assert t.ones(full_shape)[det_slice].shape == shape - assert t.allclose(rs_basis[0,1],t.Tensor([-8.333333333333333e-07])) - assert t.allclose(rs_basis[1,0],t.Tensor([-4.444444444444444e-07])) - # Then test it's padding function rs_basis, full_shape, det_slice = \ initializers.exit_wave_geometry(basis, shape, wavelength, - distance, opt_for_fft=False, - padding=2) + distance, padding=2) exp_shape = t.Size([77,60]) assert full_shape == exp_shape assert t.ones(full_shape)[det_slice].shape == shape @@ -43,8 +31,7 @@ def test_exit_wave_geometry(): center = t.Tensor([20,42]) rs_basis, full_shape, det_slice = \ initializers.exit_wave_geometry(basis, shape, wavelength, - distance, center=center, - opt_for_fft=False) + distance, center=center) exp_shape = t.Size([105,84]) assert full_shape == exp_shape assert t.ones(full_shape)[det_slice].shape == shape diff --git a/tests/tools/test_plotting.py b/tests/tools/test_plotting.py index caba38c..a2a50b8 100644 --- a/tests/tools/test_plotting.py +++ b/tests/tools/test_plotting.py @@ -2,18 +2,18 @@ from cdtools.tools import plotting from cdtools.tools import initializers import numpy as np import torch as t -import scipy.misc +import scipy.datasets import matplotlib.pyplot as plt def test_plot_amplitude(show_plot): # Test with tensor - im = t.as_tensor(scipy.misc.ascent(),dtype=t.complex128) + im = t.as_tensor(scipy.datasets.ascent(),dtype=t.complex128) plotting.plot_amplitude(im, basis = np.array([[1,1], [1,1], [0,0]]), title = 'Test Amplitude') if show_plot: plt.show() # Test with numpy array - im = scipy.misc.ascent().astype(np.complex128) + im = scipy.datasets.ascent().astype(np.complex128) plotting.plot_amplitude(im, title = 'Test Amplitude') if show_plot: plt.show() @@ -35,7 +35,7 @@ def test_plot_phase(show_plot): def test_plot_colorize(show_plot): # Test with tensor gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]) - im = gaussian * t.as_tensor(scipy.misc.ascent(), dtype=t.complex64) + im = gaussian * t.as_tensor(scipy.datasets.ascent(), dtype=t.complex64) plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[1,1], [1,1], [0,0]])) if show_plot: plt.show() diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index ca4c734..2e1b1df 100644 --- a/tests/tools/test_propagators.py +++ b/tests/tools/test_propagators.py @@ -5,7 +5,7 @@ from cdtools.tools import image_processing import numpy as np import torch as t import pytest -import scipy.misc +import scipy.datasets from scipy import stats from matplotlib import pyplot as plt @@ -13,7 +13,7 @@ from matplotlib import pyplot as plt @pytest.fixture(scope='module') def exit_waves_1(): # Import scipy test image and add a random phase - obj = scipy.misc.ascent()[0:64,0:64].astype(np.complex128) + obj = scipy.datasets.ascent()[0:64,0:64].astype(np.complex128) arr = np.random.random_sample((64,64)) obj *= (arr+(1-arr**2)**.5*1j) obj = t.as_tensor(obj) @@ -53,7 +53,7 @@ def test_generate_high_NA_k_intensity_map(): distance = 1#6e-3 rs_basis, full_shape, det_slice = \ initializers.exit_wave_geometry(basis, shape, wavelength, - distance, opt_for_fft=False) + distance) k_map, intensity_map = propagators.generate_high_NA_k_intensity_map( rs_basis, basis, shape, distance, wavelength, @@ -84,7 +84,7 @@ def test_generate_high_NA_k_intensity_map(): distance = 0.01#6e-3 rs_basis, full_shape, det_slice = \ initializers.exit_wave_geometry(basis, shape, wavelength, - distance, opt_for_fft=False) + distance) rs_basis_tilted = rs_basis.clone() rs_basis_tilted[2,1] = rs_basis_tilted[0,1] @@ -271,7 +271,7 @@ def test_near_field(): asp = propagators.generate_angular_spectrum_propagator( E0.shape, spacing, wavelength, z) - asp[10,10].backward() + t.real(asp[10,10]).backward() assert z.grad != 0 assert spacing.grad[0] != 0 assert wavelength.grad !=0