From 54ddb078e19822790a4bf61c4d7abc4063299471 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 8 Jun 2022 11:01:16 -0700 Subject: [PATCH] make some changes relating to oversampling in FancyPtycho --- CDTools/datasets/ptycho_2d_dataset.py | 22 ++++++++-------- CDTools/models/fancy_ptycho.py | 12 +++++---- CDTools/models/rpi.py | 36 +++++++++++++++++++++++++-- 3 files changed, 53 insertions(+), 17 deletions(-) diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index d684477..c10864e 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -123,7 +123,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, cut_zeros=True): + def from_cxi(cls, cxi_file, cut_zeros=True, load_patterns=True): """Generates a new Ptycho2DDataset from a .cxi file directly This generates a new Ptycho2DDataset from a .cxi file storing @@ -144,7 +144,7 @@ class Ptycho2DDataset(CDataset): # 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, load_patterns=load_patterns) # Generate a base dataset dataset = CDataset.from_cxi(cxi_file) @@ -152,17 +152,19 @@ class Ptycho2DDataset(CDataset): dataset.__class__ = cls # Load the data that is only relevant for this class - 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') - - dataset.axes = axes - if dataset.mask is None: - dataset.mask = t.ones(dataset.patterns.shape[-2:]).to(dtype=t.bool) + if load_patterns: + patterns, axes = cdtdata.get_data(cxi_file, cut_zeros=cut_zeros) + 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') + + dataset.axes = axes + + if dataset.mask is None: + dataset.mask = t.ones(dataset.patterns.shape[-2:]).to(dtype=t.bool) try: intensities = cdtdata.get_shot_to_shot_info(cxi_file, 'intensities') diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index ca4996c..c43a63a 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -77,12 +77,14 @@ class FancyPtycho(CDIModel): if background is None: if detector_slice is not None: - shape = self.probe[0][self.detector_slice].shape + dummy_det = t.empty([s//oversampling + for s in self.probe[0].shape]) + shape = dummy_det[self.detector_slice].shape + #shape = self.probe[0][self.detector_slice].shape else: - shape = self.probe[0].shape - background = 1e-6 * t.ones([s//oversampling for s in shape], - dtype=t.float32) - + shape = [s//oversampling for s in self.probe[0]] + background = 1e-6 * t.ones(shape, dtype=t.float32) + self.background = t.nn.Parameter(background) if weights is None: diff --git a/CDTools/models/rpi.py b/CDTools/models/rpi.py index a2e0e51..3703546 100644 --- a/CDTools/models/rpi.py +++ b/CDTools/models/rpi.py @@ -7,6 +7,7 @@ from CDTools.tools import initializers from scipy.ndimage.morphology import binary_dilation import numpy as np from copy import copy +import time __all__ = ['RPI'] @@ -245,7 +246,36 @@ class RPI(CDIModel): obj_support=obj_support, oversampling=oversampling, weight_matrix=weight_matrix) + @classmethod + def from_calibration(cls, calibration, obj_size=None, n_modes=1, saturation=None): + + wavelength = calibration['wavelength'] + probe_basis = t.as_tensor(calibration['basis']) + probe = t.as_tensor(calibration['probe']) + if 'background' in calibration: + background = t.as_tensor(calibration['background']) + else: + background = t.zeros_like(probe.real) + if 'mask' in calibration: + mask = t.as_tensor(calibration['mask']) + else: + mask = t.ones_like(probe.real, dtype=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 = (np.array(probe.shape[-2:]) // 2).astype(int) + + obj_guess = t.exp(2j * np.pi * t.rand([n_modes,]+obj_size)) + + det_geo = {'distance': 1, + 'basis': wavelength / probe_basis} # Is this even right? + + return cls(wavelength, det_geo, probe_basis, + probe, obj_guess, + background=background, mask=mask) + + def random_init(self, pattern): scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2) self.obj.data = scale * t.exp( @@ -306,7 +336,8 @@ class RPI(CDIModel): def forward_propagator(self, wavefields): - return tools.propagators.far_field(wavefields) + p = tools.propagators.far_field(wavefields) + return p def backward_propagator(self, wavefields): @@ -317,12 +348,13 @@ class RPI(CDIModel): # 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, + m = tools.measurements.quadratic_background(wavefields, self.background, detector_slice=self.detector_slice, measurement=tools.measurements.incoherent_sum, saturation=self.saturation, oversampling=self.oversampling) + return m def loss(self, sim_data, real_data, mask=None): return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)