mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 05:22:41 +02:00
265 lines
11 KiB
Python
265 lines
11 KiB
Python
from __future__ import division, print_function, absolute_import
|
|
import numpy as np
|
|
import torch as t
|
|
|
|
__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
|
|
'gaussian_probe']
|
|
|
|
from CDTools.tools import cmath
|
|
from CDTools.tools.propagators import inverse_far_field
|
|
from scipy.fftpack import next_fast_len
|
|
import numpy as np
|
|
|
|
|
|
def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0):
|
|
"""Returns an exit wave basis and shape, as well as a detector slice for the given detector geometry
|
|
|
|
It takes in the parameters for a given detector - the basis defining
|
|
the pixel pitch and the shape, as well as the wavelength and propagation
|
|
distance. Optionally, it accepts a defined "center", the zero-frequency
|
|
pixel's location. It then will automatically define a larger detector
|
|
if necessary, define the exit wave basis associated with a far-field
|
|
diffraction experiment, and return that basis, shape, and detector slice
|
|
|
|
Args:
|
|
det_basis (torch.Tensor) : The detector basis, as defined elsewhere
|
|
det_shape (torch.Size) : The (i,j) shape of the detector
|
|
wavelength (float) : The wavelength of light for the experiment, in m
|
|
distance (float) : The sample-detector distance, in m
|
|
center (torch.Tensor) : If defined, the location of the zero frequency pixel
|
|
opt_for_fft (bool) : Default is true, whether to increase detector size to improve fft performance
|
|
padding (int) : Default is 0, an extra border to allow for subpixel shifting later
|
|
|
|
Returns:
|
|
torch.Tensor : The exit wave basis
|
|
torch.Tensor : The exit wave's shape
|
|
tuple(slice) : The slice corresponding to the physical detector
|
|
"""
|
|
|
|
det_shape = t.Tensor(tuple(det_shape)).to(t.int32)
|
|
det_basis = t.Tensor(det_basis)
|
|
# First, set the center if it's not already specified
|
|
# This definition matches the center pixel of an fftshifted array
|
|
if center is None:
|
|
center = det_shape // 2
|
|
else:
|
|
center = t.Tensor(center).to(t.int32)
|
|
|
|
# Then, calculate the required detector size from the centering
|
|
# This is a bit opaque but was worth doing accurately
|
|
min_left = center * 2
|
|
min_right = (det_shape - center) * 2 - 1
|
|
full_shape = t.max(min_left,min_right).to(t.int32) + 2 * padding
|
|
|
|
# In some edge cases this shape can be smaller than the detector shape
|
|
full_shape = t.max(full_shape, det_shape)
|
|
|
|
if opt_for_fft:
|
|
full_shape = t.Tensor([next_fast_len(dim) for dim in full_shape]).to(t.int32)
|
|
|
|
# Then, generate a slice that pops the actual detector from the full
|
|
# detector shape
|
|
full_center = full_shape // 2
|
|
det_slice = np.s_[int(full_center[0]-center[0]):
|
|
int(full_center[0]-center[0]+det_shape[0]),
|
|
int(full_center[1]-center[1]):
|
|
int(full_center[1]-center[1]+det_shape[1])]
|
|
|
|
|
|
# Finally, generate the basis for the exit wave in real space
|
|
# I believe this calculation is incorrect for non-rectangular
|
|
# detectors, because the real space basis should be related to the
|
|
# dual of the original basis. Leaving this for now since
|
|
# non-rectangular detectors are not a pressing concern.
|
|
basis_dirs = det_basis / t.norm(det_basis, dim=0)
|
|
real_space_basis = basis_dirs * wavelength * distance / \
|
|
(full_shape.to(t.float32) * t.norm(det_basis,dim=0))
|
|
|
|
# Finally, convert the shape back to a torch.Size
|
|
full_shape = t.Size([dim for dim in full_shape])
|
|
|
|
|
|
return real_space_basis, full_shape, det_slice
|
|
|
|
|
|
def calc_object_setup(probe_shape, translations, padding=0):
|
|
"""Returns an object shape and minimum pixel translation
|
|
|
|
Based on the given pixel-space translations, it will calculate the
|
|
required size for an object array and calculate the pixel translation
|
|
that corresponds to a shift by (0,0) of the probe.
|
|
|
|
Optionally a small extra border can be defined via the padding
|
|
attribute. If this is done, the calculated pixel translation will
|
|
correspond to (padding,padding)
|
|
|
|
Args:
|
|
probe_shape (torch.Size) : The size of the probe array
|
|
translations (torch.Tensor) : Jx2 stack of pixel-valued (i,j) translations
|
|
padding (int) : Optional, the size of an extra border to include
|
|
Returns:
|
|
torch.Size : required size of object array
|
|
torch.Tensor : minimum pixel-valued translation
|
|
"""
|
|
|
|
# First we look at the translations to find the minimum translation
|
|
# and the range of translations
|
|
min_translation = t.min(translations, dim=0)[0]
|
|
translation_range = t.max(translations, dim=0)[0] - min_translation
|
|
|
|
# Calculate the required shape
|
|
translation_range = t.ceil(translation_range).numpy().astype(np.int32)
|
|
shape = translation_range + np.array(probe_shape) + 2 * padding
|
|
shape = t.Size(shape)
|
|
|
|
# And the minimum translation
|
|
min_translation = min_translation - padding
|
|
|
|
return shape, min_translation
|
|
|
|
|
|
|
|
def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]):
|
|
"""Returns an array with a centered Gaussian
|
|
|
|
Takes in the shape, amplitude, and standard deviation of a gaussian
|
|
and returns a complex torch tensor (trailing dimension is 2) with
|
|
values corresponding to a two-dimensional gaussian function
|
|
|
|
Note that [0, 0] is taken to be at the upper left corner of the array.
|
|
Default is centered at ((shape[0]-1)/2, (shape[1]-1)/2)) because x and y are zero-indexed.
|
|
By default, the phase is uniformly 0, however a curvature can be
|
|
specified to simulate a probe that has been propagated a known distance
|
|
from it's focal point. The curvature is implemented by adding a quadratic
|
|
phase phi = exp(i*curvature/2 r^2) to the Gaussian
|
|
|
|
Args:
|
|
shape (array_like) : A 1x2 array-like object specifying the dimensions of the output array in the form (i shape, j shape)
|
|
amplitude (float or int): The amplitude the gaussian to simulate
|
|
sigma (array_like): A 1x2 array-like object specifying the i- and j- standard deviation of the gaussian in the form (i stdev, j stdev)
|
|
center (array_like) : Optional 1x2 array-like object specifying the location of the center of the gaussian (i center, j center)
|
|
curvature (array_like) : Optional complex part to add to the gaussian coefficient
|
|
|
|
Returns:
|
|
torch.Tensor : The complex-style tensor storing the Gaussian
|
|
"""
|
|
if center is None:
|
|
center = ((shape[0]-1)/2, (shape[1]-1)/2)
|
|
|
|
i, j = np.mgrid[:shape[0], :shape[1]]
|
|
isq = (i - center[0])**2
|
|
jsq = (j - center[1])**2
|
|
result = np.exp((1j*curvature[0] / 2 - 1 / (2 * sigma[0]**2)) * isq + \
|
|
(1j*curvature[1] / 2 - 1 / (2 * sigma[1]**2)) * jsq)
|
|
return cmath.complex_to_torch(amplitude*result)
|
|
|
|
|
|
|
|
def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0):
|
|
"""Initializes a gaussian probe based on experimental parameters
|
|
|
|
This function generates a gaussian probe initialization which has a
|
|
total fluence matching the order of magnitude of the intensity in
|
|
the observed dataset, provided the object function is of order 1.
|
|
|
|
The initialization is done using parameters defined in physical units,
|
|
such as sigma (in meters) and the propagation distance (in meters).
|
|
The internal conversion to pixel space is done with a provided probe
|
|
basis and probe shape.
|
|
|
|
TODO: Should be updated to accept a mask
|
|
|
|
Sigma can be provided either as a scalar for a uniform beam, or as
|
|
an iterable of length 2 with [sigma_i, sigma_j] being the components
|
|
of sigma in the directions parallel to the i and j basis vectors of
|
|
the probe basis
|
|
|
|
Args:
|
|
dataset (Ptycho_2D_Dataset) : The dataset whose intensity we want to match
|
|
basis (array_like) : The real space basis for exit waves in our experiment
|
|
shape (array_like): The shape of the simulated real space arrays
|
|
sigma (array_like): The standard deviation of the probe at it's focus
|
|
propagation_distance (float) : Optional, a distance to propagate the gaussian from it's focus
|
|
|
|
Returns:
|
|
torch.Tensor : The complex-style tensor storing the Gaussian
|
|
"""
|
|
# First, we want to generate the parameters (sigma and curvature) for the
|
|
# propagated gaussian. Ignore the purely z-dependent phases
|
|
wavelength = dataset.wavelength
|
|
z = propagation_distance # for shorthand
|
|
sigma = np.array(sigma)
|
|
k = 2 * np.pi / wavelength
|
|
zr = k * sigma**2
|
|
sigmaz = sigma * np.sqrt(1 + (z / zr)**2)
|
|
curvature = -k * z / (z**2 + zr**2)
|
|
|
|
# The conversion must then be done to pixel space
|
|
sigma_pix = sigmaz / np.array([np.linalg.norm(basis[:,0]),
|
|
np.linalg.norm(basis[:,1])])
|
|
curvature_pix = curvature * np.array([np.linalg.norm(basis[:,0]),
|
|
np.linalg.norm(basis[:,1])])**2
|
|
|
|
# Then we can generate the gaussian array
|
|
probe = gaussian(shape, sigma=sigma_pix, curvature=curvature_pix)
|
|
|
|
# Finally, we should calculate the average pattern intensity from the
|
|
# dataset and normalize the gaussian probe. This should be done by
|
|
avg_intensities = [t.sum(dataset[idx][1]) for idx in range(len(dataset))]
|
|
avg_intensity = t.mean(t.Tensor(avg_intensities))
|
|
probe_intensity = t.sum(cmath.cabssq(probe))
|
|
return avg_intensity / probe_intensity * probe
|
|
|
|
|
|
def SHARP_style_probe(dataset, shape, det_slice):
|
|
"""Generates a SHARP style probe guess from a dataset
|
|
|
|
What we call the "SHARP" style probe guess is to take a mean of all
|
|
the diffraction patterns and use that as an initial guess of the
|
|
Fourier space distribution of the probe. We set all the phases to
|
|
zero, which would for many simple beams (like a zone plate) generate
|
|
a first guess of the probe that is very close to the focal spot of
|
|
the probe beam.
|
|
|
|
If the probe is simulated in higher resolution than the detector,
|
|
a common occurence, these undefined pixels are set to zero for the
|
|
purposes of defining the guess
|
|
|
|
We make a small tweak to this procedure to lower the central pixel of
|
|
the probe generated this way, which can often overwhelm the rest of the
|
|
probe if there is significant noise on the detector
|
|
|
|
Args:
|
|
dataset (Ptycho_2D_Dataset) : The dataset to work from
|
|
shape (torch.Size) : The size of the probe array to simulate
|
|
det_slice (slice) : A slice or tuple of slices corresponding to the detector region in Fourier space
|
|
"""
|
|
|
|
|
|
intensities = np.zeros(shape)
|
|
for params, im in dataset:
|
|
intensities[det_slice] += im.cpu().numpy()
|
|
intensities /= len(dataset)
|
|
|
|
probe_fft = cmath.complex_to_torch(np.sqrt(intensities))
|
|
|
|
probe_guess = cmath.torch_to_complex(inverse_far_field(probe_fft))
|
|
|
|
# Now we remove the central pixel
|
|
|
|
center = np.array(probe_guess.shape) // 2
|
|
|
|
|
|
# I had to remove this because it put some intensity outside of
|
|
# the detector region that caused issues
|
|
|
|
probe_guess[center[0], center[1]]=np.mean([
|
|
probe_guess[center[0]-1, center[1]],
|
|
probe_guess[center[0]+1, center[1]],
|
|
probe_guess[center[0], center[1]-1],
|
|
probe_guess[center[0], center[1]+1]])
|
|
|
|
return cmath.complex_to_torch(probe_guess)
|
|
|
|
|