generate diffraction data with errors added for debugging on gpu

2026-09-11 22:12:38 +02:00 · 2019-05-14 01:38:15 -04:00
parent 4896d7354a
commit e7c0e4771c
1 changed files with 86 additions and 0 deletions
@@ -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