mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 05:22:41 +02:00
87 lines
4.0 KiB
Python
87 lines
4.0 KiB
Python
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
|