make some changes relating to oversampling in FancyPtycho

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')

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:

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)