Merge branch 'master' of github.mit.edu:Scattering/CDTools

CDTools/models/simple_ptycho.py

@@ -7,7 +7,7 @@ from CDTools.tools import plotting as p
 from copy import copy
 from torch.utils import data as torchdata
 from matplotlib import pyplot as plt
-
+import numpy as np
 
 class SimplePtycho(CDIModel):
 
@@ -45,6 +45,7 @@ class SimplePtycho(CDIModel):
         self.obj = t.nn.Parameter(obj_guess.to(t.float32))
 
 
+
     @classmethod
     def from_dataset(cls, dataset):
         wavelength = dataset.wavelength
@@ -138,14 +139,10 @@ class SimplePtycho(CDIModel):
 
 
     def inspect(self):
-        p.plot_amplitude(self.probe, basis=self.probe_basis)
-        plt.title('Probe Amplitude')
-        p.plot_phase(self.probe, basis=self.probe_basis)
-        plt.title('Probe Phase')
-        p.plot_amplitude(self.obj, basis=self.probe_basis)
-        plt.title('Object Amplitude')
-        p.plot_phase(self.obj, basis=self.probe_basis)
-        plt.title('Object Phase')
+        p.plot_amplitude(self.probe, basis=self.probe_basis, title = 'Probe Amplitude')
+        p.plot_phase(self.probe, basis=self.probe_basis, title = 'Probe Phase')
+        p.plot_amplitude(self.obj, basis=self.probe_basis, title = 'Object Amplitude')
+        p.plot_phase(self.obj, basis=self.probe_basis, title = 'Object Phase')
 
 
     def save_results(self):
@@ -153,8 +150,8 @@ class SimplePtycho(CDIModel):
         probe = probe * self.probe_norm.detach().cpu().numpy()
         obj = tools.cmath.torch_to_complex(self.obj.detach().cpu())
         return {'probe':probe,'obj':obj}
-        
-    
+
+
     def ePIE(self, iterations, dataset, beta = 1.0):
         """Runs an ePIE reconstruction as described in `Maiden et al. (2017) <https://www.osapublishing.org/optica/abstract.cfm?uri=optica-4-7-736>`_.
         Optional parameters are:
@@ -172,13 +169,13 @@ class SimplePtycho(CDIModel):
             mask=None
 
         def probe_update(exit_wave, exit_wave_corrected, probe, object, translation):
-            return probe+tools.cmath.cmult(beta*(tools.cmath.cconj(object)/t.max(tools.cmath.cabssq(object)))[translation[0]:translation[0]+probe_shape[0],translation[1]:translation[1]+probe_shape[1]], \
-            exit_wave_corrected-exit_wave)
+            new_probe = probe + tools.cmath.cmult(beta * tools.cmath.cconj(object[translation])/(self.probe_norm*t.max(tools.cmath.cabssq(object))), exit_wave_corrected-exit_wave)
+            return new_probe
 
         def object_update(exit_wave, exit_wave_corrected, probe, object, translation):
-            object[translation[0]:translation[0]+probe_shape[0],translation[1]:translation[1]+probe_shape[1]]\
-            +=tools.cmath.cmult(tools.cmath.cconj(probe)/t.max(tools.cmath.cabssq(probe)),exit_wave_corrected-exit_wave)
-            return object
+            new_object = object.clone()
+            new_object[translation] = object[translation] + tools.cmath.cmult(beta * tools.cmath.cconj(probe)/(self.probe_norm*t.max(tools.cmath.cabssq(probe))), exit_wave_corrected-exit_wave)
+            return new_object
 
         with t.no_grad():
             data_loader = torchdata.DataLoader(dataset, shuffle=True)
@@ -186,15 +183,33 @@ class SimplePtycho(CDIModel):
             for it in range(iterations):
                 loss = []
                 for (i, [translations]), [patterns] in data_loader:
-                    probe = self.probe.clone()
-                    object = self.obj.clone()
-                    exit_wave = self.interaction(i, translations)
-                    exit_wave_corrected = exit_wave.clone()
-                    exit_wave_corrected[self.detector_slice] = tools.projectors.modulus(self.forward_propagator(exit_wave)[self.detector_slice], patterns, mask = mask)
+                    probe = self.probe.data.clone()
+                    object = self.obj.data.clone()
 
-                    integer_translations = t.round(translations).to(dtype=t.int32)
-                    self.probe.data = probe_update(exit_wave, exit_wave_corrected, probe, object, integer_translations)
-                    self.obj.data = object_update(exit_wave, exit_wave_corrected, probe, object, integer_translations)
+                    exit_wave = self.interaction(i, translations).clone()
+                    # Apply modulus constraint
+                    exit_wave_corrected = exit_wave.clone()
+                    exit_wave_corrected = self.forward_propagator(exit_wave_corrected.clone())
+                    exit_wave_corrected[self.detector_slice] = tools.projectors.modulus(exit_wave_corrected.clone()[self.detector_slice], patterns, mask = mask)
+                    exit_wave_corrected = self.backward_propagator(exit_wave_corrected.clone())
+
+                    # Calculate the section of the object wavefunction to be modified
+                    pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
+                                                                         translations)
+                    pix_trans -= self.min_translation
+
+                    pix_trans = t.round(pix_trans).to(dtype=t.int32).numpy()
+
+                    object_slice = np.s_[pix_trans[0]:
+                                      pix_trans[0]+probe_shape[0],
+                                      pix_trans[1]:
+                                      pix_trans[1]+probe_shape[1]]
+
+                    # Apply probe and object updates
+                    self.probe.data = probe_update(exit_wave, exit_wave_corrected, probe, object, object_slice)
+                    self.obj.data = object_update(exit_wave, exit_wave_corrected, probe, object, object_slice)
+
+                    # Calculate loss
                     loss.append(self.loss(self.measurement(self.interaction(i, translations)), patterns))
 
-                yield t.mean(t.Tensor(loss)).cpu().numpy()
+                yield it, t.mean(t.Tensor(loss)).cpu().numpy()

CDTools/tools/data.py

@@ -30,7 +30,7 @@ __all__ = ['get_entry_info',
 #
 # I will put here some thoughts about how to load data into this program.
 #
-# 
+#
 # The reconstructions should have the ability to generate datasets.
 # So you could write a reconstruction engine and then it would be
 # able to simulate data directly in the engine for you to use as a
@@ -41,7 +41,7 @@ __all__ = ['get_entry_info',
 # file into an h5py object. This file could host the simple cxi file
 # browser, perhaps. But I think the reality is that we need individual
 # loaders for each kind of experiment. Perhaps we could put some basic
-# reuseable tools for inspecting cxi-type h5 files in this file. 
+# reuseable tools for inspecting cxi-type h5 files in this file.
 #
 #
 # Then, there can be some more sophisticated tools that load data for
@@ -65,7 +65,7 @@ __all__ = ['get_entry_info',
 # whatever), and a diffraction pattern. They would also have a "setup"
 # attribute, or "metadata", or whatever you'd want to call it, that contain
 # the various fixed experimental parameters (energy, distance, etc.)
-# 
+#
 # And I think the cxi visualizer should really go into it's own script,
 # because it's not a reuseable component.
 #
@@ -84,7 +84,7 @@ def get_entry_info(cxi_file):
     String type metadata is read out as a string, and datetime metadata
     is converted to python datetime objects if the string is properly
     formatted.
-    
+
     Args:
         cxi_file (h5py.File) : a file object to be read
 
@@ -112,7 +112,7 @@ def get_entry_info(cxi_file):
 
 def get_sample_info(cxi_file):
     """Returns a dictionary with the basic metadata from the cxi file's entry_1/sample_1 attribute
-    
+
     Args:
         cxi_file (h5py.File) : a file object to be read
 
@@ -122,7 +122,7 @@ def get_sample_info(cxi_file):
     """
     if 'entry_1/sample_1' not in cxi_file:
         return None
-    
+
     s1 = cxi_file['entry_1/sample_1']
     metadata_attrs = ['name','description','unit_cell_group']
 
@@ -135,7 +135,7 @@ def get_sample_info(cxi_file):
                 metadata[attr] = str(s1[attr][()].decode())
             except AttributeError as e:
                 metadata[attr] = str(np.array(s1[attr][:])[0].decode())
-    
+
     float_attrs = ['concentration',
                    'mass',
                    'temperature',
@@ -149,23 +149,23 @@ def get_sample_info(cxi_file):
         metadata['unit_cell'] = np.array(s1['unit_cell']).astype(np.float32)
 
     # TODO: Add my nonstandard "surface normal" attribute here
-    
+
     # TODO: I should add the sample geometry as a valid metadata that can
     # be copied over
-    
+
     # Check if the metadata is empty
     if metadata == {}:
         metadata = None
-    
+
     return metadata
 
 
 def get_wavelength(cxi_file):
     """Returns the wavelength of the source defined in the cxi file object, in m
-    
+
     Args:
         cxi_file (h5py.File) : a file object to be read
-    
+
     Returns:
         np.float32 : The wavelength of the source defined in the cxi file
     """
@@ -177,7 +177,7 @@ def get_wavelength(cxi_file):
         wavelength = 1.9864459e-25 / energy
     else:
         raise KeyError('Neither Wavelength or Energy Defined in provided .cxi File')
-    
+
     return wavelength
 
 
@@ -189,16 +189,16 @@ def get_detector_geometry(cxi_file):
     outputs includes the sample to detector distance, the corner location
     of the detector, and the basis vectors defining the detector. It can
     only handle detectors defined as rectangular grids of pixels.
-    
+
     The distance and corner_location values are technically overdetermining
     the detector location, but for many experiments (particularly
-    transmission experiments), the distance is needed and the exact 
+    transmission experiments), the distance is needed and the exact
     corner location is not. If the corner location is not reported in
     the cxi file, no attempt will be made to calculate it.
 
     Args:
         cxi_file (h5py.File) : a file object to be read
-    
+
     Returns:
         distance (np.float32) : The sample to detector distance, in m
         basis_vectors (np.array) : The basis vectors for the detector
@@ -210,6 +210,8 @@ def get_detector_geometry(cxi_file):
 
     if 'detector_1/basis_vectors' in i1:
         basis_vectors = np.array(d1['basis_vectors'])
+        if basis_vectors.shape == (2,3):
+            basis_vectors = basis_vectors.T
     else:
         # This whole thing just to account for all the ways people can
         # implicitly define the x or y pixel size for a detector. I've
@@ -231,7 +233,7 @@ def get_detector_geometry(cxi_file):
             raise KeyError('Detector pixel size not defined in file.')
         basis_vectors = np.array([[0,-y_pixel_size,0],
                                   [-x_pixel_size,0,0]]).transpose()
-        
+
     try:
         distance = np.float32(d1['distance'])
     except:
@@ -240,7 +242,7 @@ def get_detector_geometry(cxi_file):
         corner_position = np.array(d1['corner_position'])
     except:
         corner_position = None
-    
+
     # Don't pretend to calculate corner position from distance if it's
     # if it's not defined, but do calculate distance from corner position
     # if distance is not defined. If neither is defined, then raise
@@ -250,7 +252,7 @@ def get_detector_geometry(cxi_file):
                                    basis_vectors[:,1])
         detector_normal /= np.linalg.norm(detector_normal)
         distance = np.linalg.norm(np.dot(corner_position, detector_normal))
-        
+
     if distance is None and corner_position is not None:
         raise KeyError('Neither sample to detector distance or corner position is defined in file.')
 
@@ -268,14 +270,14 @@ def get_mask(cxi_file):
     pixel, with the exception of pixels marked exactly as 0x00001000,
     which is defined to mean that the pixel has signal above the
     background. These pixels are treated as on pixels
-    
+
     Args:
         cxi_file (h5py.File) : a file object to be read
-    
+
     Returns:
         np.array : An array storing the mask from the cxi file
     """
-    
+
     i1 = cxi_file['entry_1/instrument_1']
     if 'detector_1/mask' in i1:
         mask = np.array(i1['detector_1/mask']).astype(np.uint32)
@@ -299,7 +301,7 @@ def get_dark(cxi_file):
 
     Args:
         cxi_file (h5py.File) : a file object to be read
-    
+
     Returns:
         np.array : An array storing the dark image
     """
@@ -314,7 +316,7 @@ def get_dark(cxi_file):
     return darks
 
 
-    
+
 def get_data(cxi_file, cut_zeroes = True):
     """Returns an array with the full stack of detector data defined in the cxi file object
 
@@ -324,19 +326,19 @@ def get_data(cxi_file, cut_zeroes = True):
     to all the required locations.
 
     It will return the data array in whatever shape it's defined in.
-    
+
     It will also read out the axes attribute of the data into a list
     of strings
-    
+
     Args:
         cxi_file (h5py.File) : a file object to be read
-    
+
     Returns:
         np.array : An array storing the data defined in the cxi file
         list : A list of the axes defined in the axes attribute, if any
     """
     # Possible locations for the data
-    # 
+    #
     # entry_1/detector_1/data
     if 'entry_1/data_1/data' in cxi_file:
         pull_from = 'entry_1/data_1/data'
@@ -348,7 +350,7 @@ def get_data(cxi_file, cut_zeroes = True):
 
     if cut_zeroes:
         data[data < 0] = 0
-        
+
     if 'axes' in cxi_file[pull_from].attrs:
         axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':')
         axes = [axis.strip().lower() for axis in axes]
@@ -365,15 +367,15 @@ def get_ptycho_translations(cxi_file):
     It negates the translations, because the CXI file format is designed
     to specify translations of the samples and the CDTools code specifies
     translations of the optics.
-    
+
     Args:
         cxi_file (h5py.File) : a file object to be read
-    
+
     Returns:
         np.array : An array storing the translations defined in the cxi file
         list : A list of the axes defined in the axes attribute, if any
     """
-    
+
     if 'entry_1/data_1/translation' in cxi_file:
         pull_from = 'entry_1/data_1/translation'
     elif 'entry_1/sample_1/geometry_1/translation' in cxi_file:
@@ -385,7 +387,7 @@ def get_ptycho_translations(cxi_file):
 
     translations = -np.array(cxi_file[pull_from]).astype(np.float32)
     return translations
-    
+
 
 
 #
@@ -397,10 +399,10 @@ def create_cxi(filename):
     """Creates a new cxi file with a single entry group
 
     Args:
-        filename (str) : The path at which to create the file    
+        filename (str) : The path at which to create the file
     """
     file_obj = h5py.File(filename,'w')
-    file_obj.create_dataset('cxi_version', data=160)  
+    file_obj.create_dataset('cxi_version', data=160)
     file_obj.create_dataset('number_of_entries',data=1)
     e1f = file_obj.create_group('entry_1')
     return file_obj
@@ -408,7 +410,7 @@ def create_cxi(filename):
 
 def add_entry_info(cxi_file, metadata):
     """Adds a dictionary of entry metadata to the entry_1 group of a cxi file object
-    
+
     Args:
         cxi_file (h5py.File) : The file to add the info to
         metadata (dict) : A dictionary containing all the metadata to be stored
@@ -433,7 +435,7 @@ def add_sample_info(cxi_file, metadata):
     """Adds a dictionary of entry metadata to the entry_1/sample_1 group of a cxi file object
 
     This function will create the sample_1 attribute if it doesn't already exist
-    
+
     Args:
         cxi_file (h5py.File) : The file to add the info to
         metadata (dict) : A dictionary containing all the metadata to be stored
@@ -441,7 +443,7 @@ def add_sample_info(cxi_file, metadata):
     if 'entry_1/sample_1' not in cxi_file:
         cxi_file['entry_1'].create_group('sample_1')
     s1 = cxi_file['entry_1/sample_1']
-        
+
     for key, value in metadata.items():
         if isinstance(value,(str,bytes)):
             s1[key] = np.string_(value)
@@ -454,11 +456,11 @@ def add_sample_info(cxi_file, metadata):
         elif isinstance(value, t.Tensor):
             asnumpy = value.detach().cpu().numpy()
             s1.create_dataset(key, data=asnumpy)
-    
+
 
 def add_source(cxi_file, wavelength):
     """Adds the entry_1/source_1 group to a cxi file object
-    
+
     It stores the energy and wavelength attributes in the source_1 group,
     given a wavelength to define them from.
 
@@ -479,7 +481,7 @@ def add_source(cxi_file, wavelength):
 
 def add_detector(cxi_file, distance, basis, corner=None):
     """Adds the entry_1/instrument_1/detector_1 group to a cxi file object
-    
+
     It will define all the relevant parameters - distance, pixel size,
     detector basis, and corner position (if relevant) based on the provided
     information
@@ -489,7 +491,7 @@ def add_detector(cxi_file, distance, basis, corner=None):
         distance (float) : The sample to detector distance
         basis (array_like) : The detector basis
         corner (array_like) : Optional, the corner position of the detector
-    
+
     """
     if 'entry_1/instrument_1' not in cxi_file:
         cxi_file['entry_1'].create_group('instrument_1')
@@ -513,18 +515,18 @@ def add_detector(cxi_file, distance, basis, corner=None):
 def add_mask(cxi_file, mask):
     """Adds the specified mask to the cxi file
 
-    It places the mask into the mask dataset under 
-    entry_1/instrument_1/detector_1. The internal mask is defined 
+    It places the mask into the mask dataset under
+    entry_1/instrument_1/detector_1. The internal mask is defined
     simply as a 1 for an "on" pixel and a 0 for an "off" pixel, and
     the saved mask is exactly the opposite. This is simpler than the
     most general mask allowed by the cxi file format but it captures the
     distinction between pixels to be used and pixels not to be used.
-    
+
     Args:
-        cxi_file (h5py.File) : The file to add the mask to    
+        cxi_file (h5py.File) : The file to add the mask to
         mask (array_like) : The mask to save out to the file
     """
-    
+
     if 'entry_1/instrument_1' not in cxi_file:
         cxi_file['entry_1'].create_group('instrument_1')
     i1 = cxi_file['entry_1/instrument_1']
@@ -533,7 +535,7 @@ def add_mask(cxi_file, mask):
     d1 = i1['detector_1']
     if isinstance(mask, t.Tensor):
         mask = mask.detach().cpu().numpy()
-    
+
     mask_to_save = np.zeros(mask.shape).astype(np.uint32)
     mask_to_save[mask == 0] = 1
     d1.create_dataset('mask',data=mask_to_save)
@@ -542,11 +544,11 @@ def add_mask(cxi_file, mask):
 def add_dark(cxi_file, dark):
     """Adds the specified dark image to a cxi file
 
-    It places the dark image data into the data_dark dataset under 
+    It places the dark image data into the data_dark dataset under
     entry_1/instrument_1/detector_1.
-    
+
     Args:
-        cxi_file (h5py.File) : The file to add the mask to    
+        cxi_file (h5py.File) : The file to add the mask to
         dark (array_like) : The dark image(s) to save out to the file
     """
     if 'entry_1/instrument_1' not in cxi_file:
@@ -559,13 +561,13 @@ def add_dark(cxi_file, dark):
         dark = dark.detach().cpu().numpy()
 
     d1.create_dataset('data_dark',data=dark)
-    
-    
+
+
 def add_data(cxi_file, data, axes=None):
     """Adds the specified data to the cxi file
-    
+
     It will add the data unchanged to the file, placing it in two spots:
-    
+
     1) The entry_1/instrument_1/detector_1/data path
     2) A softlink at entry_1/data_1/data
 
@@ -577,7 +579,7 @@ def add_data(cxi_file, data, axes=None):
     if 'entry_1/data_1' not in cxi_file:
         cxi_file['entry_1'].create_group('data_1')
     data1 = cxi_file['entry_1/data_1']
-    
+
     if 'entry_1/instrument_1' not in cxi_file:
         cxi_file['entry_1'].create_group('instrument_1')
     i1 = cxi_file['entry_1/instrument_1']
@@ -601,17 +603,17 @@ def add_data(cxi_file, data, axes=None):
 
 def add_ptycho_translations(cxi_file, translations):
     """Adds the specified translations to the cxi file
-    
+
     It will add the translations to the file, negating them to conform to
     the standard in cxi files that the translations refer to the object's
     translation.
-    
+
     It will generally store them in 3 places:
 
     1) The entry_1/sample_1/geometry_1/translation path
     2) A softlink at entry_1/data_1/translation
     3) A softlink at entry_1/instrument_1/detector_1/translation
-    
+
     Args:
         cxi_file (h5py.File) : The file to add the translations to
         translations (array_like) : The translations to be saved
@@ -624,11 +626,11 @@ def add_ptycho_translations(cxi_file, translations):
     if 'geometry_1' not in s1:
         s1.create_group('geometry_1')
     g1 = s1['geometry_1']
-    
+
     if 'entry_1/data_1' not in cxi_file:
         cxi_file['entry_1'].create_group('data_1')
     data1 = cxi_file['entry_1/data_1']
-    
+
     if 'entry_1/instrument_1' not in cxi_file:
         cxi_file['entry_1'].create_group('instrument_1')
     i1 = cxi_file['entry_1/instrument_1']
@@ -643,7 +645,7 @@ def add_ptycho_translations(cxi_file, translations):
     # accounting for the different definition between cxi files and
     # CDTools
     translations = -translations
-    
+
     g1.create_dataset('translation', data=translations)
     data1['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')
     det1['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation')

CDTools/tools/projectors.py

@@ -22,7 +22,7 @@ def modulus(wavefront, intensities, mask = None):
         wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts
         intensities (torch.Tensor): The measured diffraction pattern(s) stored as an JxNxM stack of real tensors
         mask (torch.Tensor) : Mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1
-    
+
     Returns:
         torch.Tensor : The JxNxMx2 propagated wavefield with corrected intensities
     """
@@ -55,10 +55,9 @@ def support(wavefront, support):
 
     Args:
         wavefront (torch.Tensor) : The JxNxMx2 stack of complex propagated wavefronts
-        support (torch.Tensor) : An NxM support, with 1s within the support and 0s outside 
-    
+        support (torch.Tensor) : An NxM support, with 1s within the support and 0s outside
+
     Returns:
         torch.Tensor : The JxNxMx2 wavefield with the support mask applied
     """
     return wavefront * support.to(wavefront.dtype)[...,None]
-

examples/generate_data.py

@@ -0,0 +1,86 @@
+from __future__ import division, print_function, absolute_import
+
+import CDTools
+from CDTools.tools.initializers import gaussian
+from CDTools.tools.propagators import far_field
+from CDTools.tools.measurements import intensity
+
+from CDTools.tools.cmath import *
+from CDTools.tools import interactions
+import h5py
+import numpy as np
+from matplotlib import pyplot as plt
+import torch as t
+import scipy.misc
+
+
+
+def simulate_pattern(pixel_translation_step = (10, 10), probe = None, intensity_err = True,
+                    random_noise = True, translation_err = True, background_noise = True):
+    """
+    Generates diffraction data from scipy's face image with a gaussian phase.
+    Can simulate different experimental errors, including:
+        * Varying probe intensities over time
+        * Nanopositioner location uncertainties
+        * Background detector noise
+        * Random noise
+        * Uncentered probe
+
+    We can do this by simulating the data collection process with a large
+    probe array, and changing the translation by a random (integer) amount.
+    Then, we can scale the intensity by a random amount (close to 1),
+    replace pixels with some scaled version of their current value to simulate
+    random noise, then add in horizontal and vertical detector background noise.
+    Finally, we slice the simulated diffraction patterns to uncenter the probe
+    on the final patterns [TO DO].
+
+    Args:
+        pixel_translation_step (array-like) : a tuple containing (y pixel translation, x pixel translation)
+        probe (np.ndarray) : a two-dimensional array representing the probe wavefunction. Defaults to a gaussian probe
+        intensity_err (Boolean) : Defaults to True, to simulate varying probe intensities over time
+        random_noise (Boolean) : Defaults to True, indicating adding random noise to the diffraction patterns
+        translation_err (Boolean) : Defaults to True, to simulate nanopositioner errors
+        background_noise (Boolean) : [TO DO]
+    Returns:
+        patterns (t.Tensor) : NxMx2 array of simulated diffraction patterns with error
+        translations (t.Tensor) : N*Mx2 array of instrument translations
+    """
+    obj = scipy.misc.face()
+    # Convert to grayscale
+    obj = 0.2989 * obj[:,:,0] + 0.5870 * obj[:,:,1] + 0.1140 * obj[:,:,2]
+    phase = gaussian(obj.shape, [500,500], amplitude=1, center = None)
+    obj = obj*np.exp(1j*torch_to_complex(phase))
+
+    if probe is None:
+        probe = torch_to_complex(gaussian([512,512], [5,5], amplitude=1, center = None))
+
+    obj_shape = np.array(obj.shape)
+    probe_shape = np.array(probe.shape)
+
+    translation_range = (obj_shape-probe_shape)//np.array(pixel_translation_step)*np.array(pixel_translation_step)
+    translations = np.mgrid[0:translation_range[0]+pixel_translation_step[0]:pixel_translation_step[0], 0:translation_range[1]+pixel_translation_step[1]:pixel_translation_step[1]]
+    translations = translations.T.reshape((translations.T.shape[0]*translations.T.shape[1], 2))
+    translations = t.tensor(translations, dtype = t.float32)
+    ideal_translations = t.tensor(translations, dtype = t.float32)
+    obj = complex_to_torch(obj)
+    probe = complex_to_torch(probe)
+
+    if translation_err:
+        err = t.normal(mean=t.zeros(translations.shape), std=t.ones(translations.shape)/2)
+        translations += err
+        # Ensure that the new translations don't go off the object area
+        translations[translations<0] = 0
+        translation_range = t.tensor(translation_range, dtype = t.float32)
+        translations[...,0][translations[...,0] > translation_range[0]] = translation_range[0]
+        translations[...,1][translations[...,1] > translation_range[1]] = translation_range[1]
+
+    patterns = intensity(far_field(interactions.ptycho_2D_round(probe, obj, translations)))
+
+    if intensity_err:
+        patterns *= t.normal(mean = t.ones(patterns.shape), std = t.ones(patterns.shape)/32)
+
+    if background_noise:
+        patterns += t.normal(mean = t.zeeros(patterns.shape), std = t.ones(patterns.shape)*t.max(patterns)/32)
+
+
+    return patterns, ideal_translations

examples/simple_ptycho.py

@@ -12,8 +12,8 @@ import pickle
 
 #filename = '../../Downloads/114429_p.cxi'
 #filename = '../../../Projects/CSX_3_19/cxis/processed/114429_p.cxi'
-#filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi'
-filename = '../../../Downloads/AuBalls_700ms_30nmStep_3_3SS_filter.cxi'
+filename = '../../../Projects/CSX_3_19/cxis/processed/115145_p.cxi'
+#filename = '../../Desktop/Reconstructions/114429_p.cxi'
 
 
 with h5py.File(filename,'r') as f:
@@ -24,16 +24,19 @@ model = CDTools.models.SimplePtycho.from_dataset(dataset)
 
 # Uncomment these to use on the CPU
 # default is CPU with 32-bit floats
-model.to(device='cuda')
+model.to(device='cpu')
 #dataset.to(device='cuda')
-dataset.get_as(device='cuda')
+dataset.get_as(device='cpu')
 
+for i, loss in enumerate(model.ePIE(10, dataset)):
+    print(i, loss)
 
-for loss in model.Adam_optimize(1, dataset):
-    print(loss)
+for i, loss in enumerate(model.Adam_optimize(10, dataset)):
+    print(i, loss)
+
+model.inspect()
 
 #with open('test_results.pickle', 'wb') as f:
 #    pickle.dump(model.save_results(),f)
 
-model.inspect()
 plt.show()