cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-11 22:12:38 +02:00
Code Issues Packages Projects Releases Wiki Activity

Bring all the tools to the point where they pass the tests

Browse Source
This commit is contained in:
Abe Levitan
2021-06-18 14:01:39 -04:00
parent a3765e9097
commit c2b8f2fe1b
19 changed files with 247 additions and 507 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/tools/__init__.py
-1
View File
@@ -22,7 +22,6 @@ from CDTools.tools import data
from CDTools.tools import image_processing
from CDTools.tools import initializers
from CDTools.tools import plotting
from CDTools.tools import projectors
from CDTools.tools import interactions
from CDTools.tools import propagators
from CDTools.tools import measurements
CDTools/tools/analysis/analysis.py
+18 -30
View File
@@ -170,21 +170,21 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
# First, we normalize the probe intensity to a fixed value.
probe_np = False
if isinstance(probe, np.ndarray):
probe = t.Tensor(probe).to(t.complex64)
probe = t.as_tensor(probe, dtype=t.complex64)
probe_np = True
obj_np = False
if isinstance(obj, np.ndarray):
obj = t.Tensor(obj).to(t.complex64)
obj = t.as_tensor(obj,dtype=t.complex64)
obj_np = True
# If this is a single probe and not a stack of probes
if len(probe.shape) == 3:
if len(probe.shape) == 2:
probe = probe[None,...]
single_probe = True
else:
single_probe = False
normalization = t.sqrt(t.sum(t.abs(probe[0])**2) / (len(probe[0].view(-1))/2))
normalization = t.sqrt(t.sum(t.abs(probe[0])**2) / (len(probe[0].view(-1))))
probe = probe / normalization
obj = obj * normalization
@@ -198,8 +198,8 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
# Need to check if this is actually working and, if not, why not
center_freq = ip.centroid(t.abs(t.fft.fftshift(t.fft.fft2(probe[0]),
dim=(-1,-2)))**2)
center_freq -= (t.tensor(probe[0].shape[:-1]) // 2).to(t.float32)
center_freq /= t.tensor(probe[0].shape[:-1]).to(t.float32)
center_freq -= t.div(t.tensor(probe[0].shape,dtype=t.float32),2,rounding_mode='floor')
center_freq /= t.as_tensor(probe[0].shape,dtype=t.float32)
Is, Js = np.mgrid[:probe[0].shape[0],:probe[0].shape[1]]
probe_phase_ramp = t.exp(2j * np.pi *
@@ -214,11 +214,11 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
# Then, we set them to consistent absolute phases
obj_angle = t.angle(t.sum(obj[obj_slice],dim=(0,1)))
obj_angle = t.angle(t.sum(obj[obj_slice]))
obj = obj * t.exp(-1j*obj_angle)
for i in range(probe.shape[0]):
probe_angle = t.angle(t.sum(probe[i],dim=(0,1)))
probe_angle = t.angle(t.sum(probe[i]))
probe[i] = probe[i] * t.exp(-1j*probe_angle)
if single_probe:
@@ -265,16 +265,17 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
A list of standardized objects, for further processing
"""
# This should be cleaned up so it accepts anything array_like
probe_np = False
if isinstance(probes[0], np.ndarray):
probes = [t.Tensor(probe).to(t.complex64) for probe in probes]
probes = [t.as_tensor(probe,dtype=t.complex64) for probe in probes]
probe_np = True
obj_np = False
if isinstance(objects[0], np.ndarray):
objects = [t.Tensor(obj).to(t.complex64) for obj in objects]
objects = [t.as_tensor(obj,dtype=t.complex64) for obj in objects]
obj_np = True
obj_shape = np.min(np.array([obj.shape[:-1] for obj in objects]),axis=0)
obj_shape = np.min(np.array([obj.shape for obj in objects]),axis=0)
objects = [obj[:obj_shape[0],:obj_shape[1]] for obj in objects]
if obj_slice is None:
@@ -357,10 +358,10 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
obj_np = False
if isinstance(objects[0], np.ndarray):
objects = [t.Tensor(obj).to(t.complex64) for obj in objects]
objects = [t.as_tensor(obj, dtype=t.complex64) for obj in objects]
obj_np = True
if isinstance(synth_obj, np.ndarray):
synth_obj = t.Tensor(synth_obj).to(t.complex64)
synth_obj = t.as_tensor(synth_obj, dtype=t.complex64)
if isinstance(basis, t.Tensor):
basis = basis.detach().cpu().numpy()
@@ -433,18 +434,12 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
im_np = False
if isinstance(im1, np.ndarray):
im1 = t.Tensor(im1)
im1 = t.as_tensor(im1)
im_np = True
if isinstance(im2, np.ndarray):
im2 = t.Tensor(im2)
im2 = t.as_tensor(im2)
im_np = True
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
@@ -500,22 +495,15 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
im_np = False
if isinstance(im1, np.ndarray):
im1 = t.Tensor(im1)
im1 = t.as_tensor(im1)
im_np = True
if isinstance(im2, np.ndarray):
im2 = t.Tensor(im2)
im2 = t.as_tensor(im2)
im_np = True
if isinstance(basis, np.ndarray):
basis = t.tensor(basis)
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
CDTools/tools/image_processing/image_processing.py
+28 -56
View File
@@ -147,14 +147,6 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
# using an FFT with upsampling by a factor of resolution in reciprocal
# space
#
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2))
# Not sure if this is more or less stable than just the correlation
@@ -171,14 +163,14 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
window_size = 15
shift_zero = tuple(-search_around + t.tensor([window_size,window_size]))
cor_window = t.roll(cor, shift_zero, dims=(0,1))[:2*window_size,:2*window_size]
cor_window = t.roll(cor, shift_zero, dims=(-2,-1))[...,:2*window_size,:2*window_size]
# Now we upsample this window
cor_window_fft = t.fft.fftshift(t.fft.fft2(cor_window),dim=(-2,-1))
upsampled = t.zeros(tuple(t.tensor(cor_window_fft.shape)[:-1] * resolution) + (2,),
upsampled = t.zeros(tuple(t.tensor(cor_window_fft.shape) * resolution),
dtype=cor.dtype,device=cor.device)
upsampled[:2*window_size,:2*window_size] = cor_window_fft
upsampled[...,:2*window_size,:2*window_size] = cor_window_fft
upsampled = t.roll(upsampled,(-window_size,-window_size),dims=(0,1))
upsampled = t.roll(t.abs(t.fft.ifft2(upsampled))**2,
(-window_size*resolution,-window_size*resolution),
@@ -186,10 +178,14 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
# And we extract the shift from the window
sh = t.tensor(upsampled.shape).to(device=upsampled.device)
cormax = t.tensor([t.argmax(upsampled) // sh[1],
t.argmax(upsampled) % sh[1]]).to(device=upsampled.device)
subpixel_shift = ((cormax + sh // 2) % sh - sh//2).to(dtype=upsampled.dtype)
sh = t.as_tensor(upsampled.shape, device=upsampled.device)
cormax = t.as_tensor([t.div(t.argmax(upsampled), sh[1],
rounding_mode='floor'),
t.argmax(upsampled) % sh[1]],
device=upsampled.device)
sh_over_2 = t.div(sh,2,rounding_mode='floor')
subpixel_shift = ((cormax + sh_over_2) % sh - sh_over_2).to(dtype=upsampled.dtype)
return search_around.to(device=upsampled.device, dtype=upsampled.dtype) + \
subpixel_shift / resolution
@@ -215,13 +211,6 @@ def find_pixel_shift(im1, im2):
shift : torch.Tensor
The integer-valued shift (i,j) that best maps im1 onto im2
"""
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2))
# Not sure if this is more or less stable than just the correlation
@@ -229,10 +218,12 @@ def find_pixel_shift(im1, im2):
cor = t.abs(t.fft.ifft2(cor_fft / t.abs(cor_fft)))
sh = t.tensor(cor.shape).to(device=im1.device)
cormax = t.tensor([t.argmax(cor) // sh[1],
sh = t.as_tensor(cor.shape,device=im1.device)
cormax = t.tensor([t.div(t.argmax(cor),sh[1],rounding_mode='floor'),
t.argmax(cor) % sh[1]]).to(device=im1.device)
return (cormax + sh // 2) % sh - sh//2
sh_over_2 = t.div(sh,2,rounding_mode='floor')
return (cormax + sh_over_2) % sh - sh_over_2
@@ -292,52 +283,33 @@ def convolve_1d(image, kernel, dim=0, fftshift_kernel=True):
The convolved image
"""
complex_things = 2
im_complex = True
if image.shape[-1] != 2:
image = t.stack((image,t.zeros_like(image)),dim=-1)
complex_things -= 1
im_complex = False
if kernel.shape[-1] != 2:
kernel = t.stack((kernel,t.zeros_like(kernel)),dim=-1)
complex_things -= 1
if fftshift_kernel:
kernel = t.fft.ifftshift(kernel,dim=(-2,-1))
kernel = t.fft.ifftshift(kernel,dim=(-1,))
# If the image wasn't originally complex, and the dimension
# was passed with the nexative-indexing convention
if not im_complex and dim < 0:
dim = dim-1
# We have to transpose the relevant dimension to -2 before using the fft,
# which expects to operate on the final non-complex dimension
trans_im = t.transpose(image, dim, -2)
# We have to transpose the relevant dimension to -1 before using the fft,
# which expects to operate on the final dimension
trans_im = t.transpose(image, dim, -1)
# Take a correlation
fft_im = t.fft.fft(trans_im)
fft_kernel = t.fft.fft(kernel)
trans_conv = t.fft.ifft(fft_im * fft_kernel)
conv_im = t.transpose(trans_conv, dim, -2)
conv_im = t.transpose(trans_conv, dim, -1)
# If nothing was input as complex, the result should be returned as real
if complex_things == 0:
return conv_im[...,0]
else:
return conv_im
return conv_im
def fourier_upsample(ims):
upsampled = t.zeros(ims.shape[:-3]+(2*ims.shape[-3],2*ims.shape[-2])+(2,),
upsampled = t.zeros(ims.shape[:-2]+(2*ims.shape[-2],2*ims.shape[-1]),
dtype=ims.dtype,
device=ims.device)
left = [ims.shape[-3]//2,ims.shape[-2]//2]
right = [ims.shape[-3]//2+ims.shape[-3],
ims.shape[-2]//2+ims.shape[-2]]
left = [ims.shape[-2]//2,ims.shape[-1]//2]
right = [ims.shape[-2]//2+ims.shape[-2],
ims.shape[-1]//2+ims.shape[-1]]
upsampled[...,left[0]:right[0],left[1]:right[1],:] = propagators.far_field(ims)
upsampled[...,left[0]:right[0],left[1]:right[1]] = propagators.far_field(ims)
return propagators.inverse_far_field(upsampled)
CDTools/tools/initializers/initializers.py
+18 -16
View File
@@ -60,8 +60,8 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
The slice corresponding to the physical detector
"""
det_shape = t.tensor(tuple(det_shape)).to(t.int32)
det_basis = t.tensor(det_basis)
det_shape = t.as_tensor(tuple(det_shape), dtype=t.int32)
det_basis = t.as_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:
@@ -70,7 +70,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
center = t.div(det_shape,2,rounding_mode='floor')# // 2
else:
center = t.tensor(center).to(t.int32)
center = t.as_tensor(center, dtype=t.int32)
# Then, calculate the required detector size from the centering
# This is a bit opaque but was worth doing accurately
@@ -83,7 +83,8 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
if opt_for_fft:
full_shape = t.tensor([next_fast_len(dim) for dim in full_shape]).to(t.int32)
full_shape = t.as_tensor([next_fast_len(dim) for dim in full_shape],
dtype=t.int32)
# Then, generate a slice that pops the actual detector from the full
# detector shape
@@ -205,7 +206,7 @@ def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]):
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 t.tensor(amplitude*result).to(t.complex64)
return t.as_tensor(amplitude*result,dtype=t.complex64)
@@ -314,6 +315,8 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
The complex-style tensor storing the probe guess
"""
# NOTE: I don't love the way np and torch are mixed here, I think this
# function deserves some love.
# to use the mask or not?
intensities = np.zeros([dim // oversampling for dim in shape])
@@ -342,15 +345,15 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
probe_guess[center[0], center[1]-1],
probe_guess[center[0], center[1]+1]])
probe_guess = t.tensor(probe_guess).to(dtype=t.complex64)
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
if propagation_distance is not None:
# First generate the propagation array
probe_shape = t.tensor(tuple(probe_guess.shape))
probe_shape = t.as_tensor(tuple(probe_guess.shape))
# Start by recalculating the probe basis from the given information
det_basis = t.tensor(dataset.detector_geometry['basis'])
det_basis = t.as_tensor(dataset.detector_geometry['basis'])
basis_dirs = det_basis / t.norm(det_basis, dim=0)
distance = dataset.detector_geometry['distance']
probe_basis = basis_dirs * dataset.wavelength * distance / \
@@ -360,7 +363,6 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
probe_spacing = t.norm(probe_basis,dim=0).numpy()
probe_shape = probe_shape.numpy().astype(np.int32)
#assert 0
# And generate the propagator
AS_prop = generate_angular_spectrum_propagator(probe_shape, probe_spacing, dataset.wavelength, propagation_distance)
@@ -384,10 +386,10 @@ def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, backgroun
if probe.dim() == 4:
probe = orthogonalize_probes(probe)[0]
pad0l = (probe.shape[-3] - obj_shape[0])//2
pad0r = probe.shape[-3] - obj_shape[0] - pad0l
pad1l = (probe.shape[-2] - obj_shape[1])//2
pad1r = probe.shape[-2] - obj_shape[1] - pad1l
pad0l = (probe.shape[-2] - obj_shape[0])//2
pad0r = probe.shape[-2] - obj_shape[0] - pad0l
pad1l = (probe.shape[-1] - obj_shape[1])//2
pad1r = probe.shape[-1] - obj_shape[1] - pad1l
def a_dagger(im):
im = t.tensor(im.reshape(obj_shape)).to(dtype=t.complex64)
@@ -433,13 +435,13 @@ def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, backgroun
# Now we set the overall scale and relative weights of the guess
scale_factor = np.sqrt(np.sum(np_pattern) /
t.sum(cmath.cabssq(probe)).numpy())
t.sum(t.abs(probe)**2).numpy())
relative_weights = eigval / np.sum(eigval**2)
z0 = z0 * (scale_factor * relative_weights[:,None,None])
# Now we have to normalize the modes by their eigenvalues
return cmath.complex_to_torch(z0).to(dtype=t.float32)
return t.as_tensor(z0, dtype=t.complex64)
def generate_subdominant_modes(dominant_mode, n_modes, circular=True):
CDTools/tools/interactions/interactions.py
+16 -16
View File
@@ -434,9 +434,9 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
exit_waves = []
if shift_probe:
i = t.arange(probe.shape[-2],device=probe.device,dtype=t.float32) \
- probe.shape[-3]//2
j = t.arange(probe.shape[-1],device=probe.device,dtype=t.float32) \
- probe.shape[-2]//2
j = t.arange(probe.shape[-1],device=probe.device,dtype=t.float32) \
- probe.shape[-1]//2
I,J = t.meshgrid(i,j)
I = 2 * np.pi * I / probe.shape[-2]
J = 2 * np.pi * J / probe.shape[-1]
@@ -445,19 +445,18 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
-subpixel_translations[:,1,None,None]*J))
fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2))
if multiple_modes:
# if the probe dimension is 4, then this hasn't yet been broadcast
# over the translation dimensions
if probe.dim == 3: # Multi-mode probe
shifted_fft_probe = fft_probe * phase_masks[:,None,:,:]
else:
shifted_fft_probe = fft_probe * phase_masks
shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
dim=(-1,-2)))
if multiple_modes:
# if the probe dimension is 4, then this hasn't yet been broadcast
# over the translation dimensions
print('p',probe.shape)
print('fftp',fft_probe.shape)
print('sp',shifted_probe.shape)
print('sel',selections.shape)
if probe.dim == 3: # Multi-mode probe
output = shifted_probe * selections[:,None,:,:]
else:
output = shifted_probe * selections
@@ -465,6 +464,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
else:
raise NotImplementedError('Object shift not yet implemented')
print(output.shape)
if single_translation:
return output[0]
else:
@@ -599,14 +599,14 @@ def RPI_interaction(probe, obj):
# The far-field propagator is just a 2D FFT but with an fftshift
fftobj = propagators.far_field(obj)
# We calculate the padding that we need to do the upsampling
pad0l = (probe.shape[-3] - obj.shape[-3])//2
pad0r = probe.shape[-3] - obj.shape[-3] - pad0l
pad1l = (probe.shape[-2] - obj.shape[-2])//2
pad1r = probe.shape[-2] - obj.shape[-2] - pad1l
pad0l = (probe.shape[-2] - obj.shape[-2])//2
pad0r = probe.shape[-2] - obj.shape[-2] - pad0l
pad1l = (probe.shape[-1] - obj.shape[-1])//2
pad1r = probe.shape[-1] - obj.shape[-1] - pad1l
if obj.dim() == 3:
if obj.dim() == 2:
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad0l, pad0r))
elif obj.dim() == 4:
elif obj.dim() == 3:
fftobj = t.nn.functional.pad(
fftobj, (pad1l, pad1r, pad0l, pad0r, 0,0))
else:
@@ -615,7 +615,7 @@ def RPI_interaction(probe, obj):
# Again, just an inverse FFT but with an fftshift
upsampled_obj = propagators.inverse_far_field(fftobj)
if obj.dim() == 4:
if obj.dim() == 3:
return probe[None,...] * upsampled_obj
else:
return probe * upsampled_obj
CDTools/tools/measurements/measurements.py
+9 -105
View File
@@ -16,8 +16,7 @@ from torch.nn.functional import avg_pool2d
# intensity pattern on a detector
#
__all__ = ['intensity', 'incoherent_sum', 'density_matrix',
'quadratic_background']
__all__ = ['intensity', 'incoherent_sum', 'quadratic_background']
def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
@@ -30,7 +29,7 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
Parameters
----------
wavefield : torch.Tensor
A JxMxNx2 stack of complex wavefields
A JxMxN stack of complex-valued wavefields
detector_slice : slice
Optional, a slice or tuple of slices defining a section of the simulation to return
saturation : float
@@ -44,7 +43,7 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
A real MxN array storing the wavefield's intensities
"""
output = t.abs(wavefield)**2
# Now we apply oversampling
if oversampling != 1:
if wavefield.dim() == 2:
@@ -64,101 +63,6 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
return output + epsilon
else:
return t.clamp(output + epsilon,0,saturation)
def density_matrix(wavefields, density_matrix, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
"""Returns the intensities associated with a given density matrix state
The essential idea is that the most general description of a light field
at the detector plane will consist of a density matrix state. Here, that
low rank density matrix state is encoded as a set of basis wavefields
and a density matrix in that basis.
For computational efficiency, the density matrix is coded in an unusual
format. The density matrix formally is a complex Hermetian matrix,
which also happens to be positive definite. Here, we store it as a
real-valued matrix, where the upper triangle corresponds to the real
part of the elements in the upper triangle, and the lower triangle
corresponds to the imaginary parts. The elements on the diagonal are
purely real, and are stored as they are.
As with other multi-mode measurement functions, the modes are stored in
the first index, and the index of the diffraction pattern in the stack
of diffraction patterns is the second index. The stack-direction index
can be omitted if only a single pattern needs to be simulated
It is important to note that this method does not inforce the positive
definiteness of the density matrix, this it is possible for negative
values of intensity to appear if the underlying density matrix passed
to this method is not positive definite
Parameters
----------
wavefields : torch.Tensor
An Lx(Jx)MxNx2 stack of complex wavefields
density_matrix : torch.Tensor
A (Jx)LxL stack of real-valued representations of density matrices, as per above
saturation : float
Optional, a maximum saturation value to clamp the resulting intensities to
oversampling : int
Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
Returns
-------
sim_patterns : torch.Tensor
A real Lx(Jx)MxN array storing the incoherently summed intensities
"""
#if wavefields.dim() == 4:
# wavefields.unsqueeze(1)
# single_frame = True
#elif wavefields.dim() == 5:
# single_frame=False
output = t.zeros(wavefields.shape[1:-1],
dtype=wavefields.dtype,
device=wavefields.device)
# flat is better than nested, but simple is better than complex...
for (i,j) in ((i,j) for i in range(density_matrix.shape[-2])
for j in range(density_matrix.shape[-1])):
if i == j: # diagonal
output += density_matrix[...,i,j,None,None] \
* t.abs(wavefields[i])**2
if i < j: # upper triangle, real part
output += 2 * density_matrix[...,i,j,None,None] \
* (wavefields[i,...,0] * wavefields[j,...,0]
+ wavefields[i,...,1] * wavefields[j,...,1])
if i > j: # lower triangle, imaginary part
# We pull the i,jth element from the density matrix,
# but this correponds to wavefield j and wavefield i,
# unlike above where it was wavefield i and j (swapped order).
# We also get the one negative sign because this is the imaginary
# part
output += 2 * density_matrix[...,i,j,None,None] \
* (wavefields[j,...,0] * wavefields[i,...,1]
- wavefields[j,...,1] * wavefields[i,...,0])
# Now we apply oversampling
if oversampling != 1:
if wavefields.dim() == 4:
output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
else:
output = avg_pool2d(output, oversampling)
# Then we grab the detector slice
if detector_slice is not None:
if wavefields.dim() == 4:
output = output[detector_slice]
else:
output = output[(np.s_[:],) + detector_slice]
if saturation is None:
return t.clamp(output,min=0) + epsilon
else:
return t.clamp(output + epsilon,0,saturation)
def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
@@ -168,15 +72,15 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
the wavefields. If a detector slice is given, the returned array
will only include that slice from the simulated wavefronts.
The (-4th) index is the set of incoherently adding patterns, and any
The (-3) index is the set of incoherently adding patterns, and any
indexes further to the front correspond to the set of diffraction patterns
to meaasure. The (-3rd) and (-2nd) indices are the wavefield, and the final
to measure. The (-2) and (-1) indices are the wavefield, and the final
index is the complex index
Parameters
----------
wavefields : torch.Tensor
An LxJxMxNx2 stack of complex wavefields
An LxJxMxNx stack of complex wavefields
detector_slice : slice
Optional, a slice or tuple of slices defining a section of the simulation to return
saturation : float
@@ -194,14 +98,14 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
# Now we apply oversampling
if oversampling != 1:
if wavefields.dim() == 4:
if wavefields.dim() == 3:
output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
else:
output = avg_pool2d(output, oversampling)
# Then we grab the detector slice
if detector_slice is not None:
if wavefields.dim() == 4:
if wavefields.dim() == 3:
output = output[detector_slice]
else:
output = output[(np.s_[:],) + detector_slice]
@@ -223,7 +127,7 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
Parameters
----------
wavefield : torch.Tensor
A JxMxNx2 stack of complex wavefields
A JxMxN stack of complex-valued wavefields
background : torch.Tensor
An tensor storing the square root of the detector background
detector_slice : slice
CDTools/tools/projectors/__init__.py
-3
View File
@@ -1,3 +0,0 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools.projectors.projectors import *
CDTools/tools/projectors/projectors.py
-79
View File
@@ -1,79 +0,0 @@
"""This module contains various projection functions
These functions are useful when defining declarative algorithms to run
alongside the automatic differentiation ones, for comparison or in a
situation where they might be needed.
"""
from __future__ import division, print_function, absolute_import
import torch as t
__all__ = ['modulus', 'support']
def modulus(wavefront, intensities, mask = None):
"""Implements the modulus constraint in torch
This accepts a tensor representing the propagated simulated wavefront(s),
where the last dimension represents the real and imaginary components of
the propagated wavefield(s). It projects the modulus of the diffraction
pattern onto the modulus of the simulated wavefield.
It assumes that the wavefront is stored in an array
[i,j] where i corresponds to the y-axis and j corresponds to the
x-axis, with the origin following the CS standard of being in the
upper right.
Parameters
----------
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
A mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1
Returns
-------
projected : torch.Tensor
The JxNxMx2 projected wavefield with corrected intensities
"""
# Calculate amplitudes from intensities
amplitudes = t.sqrt(intensities)
# Normalize wavefront so the complex elements have modulus one
wavefront_mag = t.abs(wavefront)
projected = wavefront * (amplitudes / wavefront_mag)[...,None]
# Replace amplitude of wavefront with measured amplitude
if mask is not None:
selection = mask == 0
# Apply the mask to replace unmasked pixels in the original wavefront
projected = projected.masked_scatter(selection, wavefront.masked_select(selection))
return projected
def support(wavefront, support):
"""Implements the support constraint in torch
This accepts a torch tensor representing (a) simulated wavefield(s),
where the last dimension represents the real and imaginary components of
the propagated wavefield(s). It projects the support of the imaged object
onto the simulated wavefront via a support mask.
It assumes that the wavefront is stored in an array
[i,j] where i corresponds to the y-axis and j corresponds to the
x-axis, with the origin following the CS standard of being in the
upper right.
Parameters
----------
wavefront : torch.Tensor
The JxNxMx2 stack of complex propagated wavefronts
support : torch.Tensor
An NxM support, with 1s within the support and 0s outside
Returns
-------
projected : torch.Tensor
The JxNxMx2 wavefield with the support mask applied
"""
return wavefront * support.to(wavefront.dtype)[...,None]
CDTools/tools/propagators/propagators.py
+13 -12
View File
@@ -288,27 +288,28 @@ def high_NA_far_field(wavefront, k_map, intensity_map=None):
# np.ones_like(k_map[0,:-1,:-1,0].cpu().numpy()))
#plt.show()
def process_wavefield_stack(low_NA_wavefield):
real_output = grid_sample(low_NA_wavefield[None,:,:,:,0],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
imag_output = grid_sample(low_NA_wavefield[None,:,:,:,1],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
# grid_sample doesn't work on complex-valued wavefields
real_output = grid_sample(low_NA_wavefield[None,:,:,:].real,k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
imag_output = grid_sample(low_NA_wavefield[None,:,:,:].imag,k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
result = t.stack((real_output[0,:,:,:],imag_output[0,:,:,:]),dim=3)
result = real_output[0,:,:,:] + 1j * imag_output[0,:,:,:]
if intensity_map is not None:
result = result * intensity_map[None,:,:,None]
result = result * intensity_map[None,:,:]
return result
original_dim = wavefront.dim()
if original_dim == 2:
result = process_wavefield_stack(low_NA_wavefield[None,:,:])
return result[0,:,:]
if original_dim == 3:
result = process_wavefield_stack(low_NA_wavefield[None,:,:,:])
return result[0,:,:,:]
if original_dim == 4:
result = process_wavefield_stack(low_NA_wavefield)
return result
if original_dim == 5:
if original_dim == 4:
result = []
for i in range(low_NA_wavefield.size()[0]):
result.append(process_wavefield_stack(low_NA_wavefield[i,:,:,:,:]))
result.append(process_wavefield_stack(low_NA_wavefield[i,:,:,:]))
return t.stack(result)
else:
raise IndexError('Wavefield had incorrect number of dimensions')
@@ -358,7 +359,7 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
"""
ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0]).numpy()
kj = 2 * np.pi * t.fftfreq(shape[1],spacing[1]).numpy()
kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1]).numpy()
Kj, Ki = np.meshgrid(kj,ki)
# Define this as complex so the square root properly gives
@@ -381,7 +382,7 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
# Take the conjugate explicitly here instead of negating
# the previous expression to ensure that complex frequencies
# get mapped to values <1 instead of >1
propagator = complex_to_torch(np.conj(propagator))
propagator = t.as_tensor(np.conj(propagator))
return propagator.to(*args, **kwargs)
@@ -585,7 +586,7 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
# Take the conjugate explicitly here instead of negating
# the previous expression to ensure that complex frequencies
# get mapped to values <1 instead of >1
propagator = complex_to_torch(np.conj(propagator))
propagator = t.as_tensor(np.conj(propagator))
return propagator.to(**kwargs)
tests/test_datasets.py
+1 -1
View File
@@ -28,7 +28,7 @@ def test_CDataset_init():
dataset = CDataset(entry_info, sample_info,
wavelength, detector_geometry, mask)
assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(np.bool))))
assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(bool))))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
tests/tools/test_analysis.py
+17 -20
View File
@@ -30,7 +30,7 @@ def test_orthogonalize_probes():
ortho_probes = analysis.orthogonalize_probes(probes)
# test that it also works on torch tensors
ortho_probes_t = cmath.torch_to_complex(analysis.orthogonalize_probes(cmath.complex_to_torch(probes)))
ortho_probes_t = analysis.orthogonalize_probes(t.as_tensor(probes)).numpy()
# This tests for orthogonality
for p1,p2 in combinations(ortho_probes,2):
@@ -80,13 +80,12 @@ def test_standardize():
# Start by making a probe and object that should meet the standardization
# conditions
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01))
probe = cmath.torch_to_complex(probe)
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2))
probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
assert np.isclose(0,np.angle(np.sum(probe)))
assert np.angle(np.sum(probe)) < 1e-7
obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
@@ -105,9 +104,9 @@ def test_standardize():
assert np.allclose(obj, s_obj)
# Test that it works on torch tensors
s_probe, s_obj = analysis.standardize(cmath.complex_to_torch(test_probe).to(t.float32), cmath.complex_to_torch(test_obj).to(t.float32))
s_probe = cmath.torch_to_complex(s_probe)
s_obj = cmath.torch_to_complex(s_obj)
s_probe, s_obj = analysis.standardize(t.as_tensor(test_probe,dtype=t.complex64), t.as_tensor(test_obj,dtype=t.complex64))
s_probe = s_probe.numpy()
s_obj = s_obj.numpy()
assert np.allclose(probe, s_probe)
assert np.allclose(obj, s_obj)
@@ -146,7 +145,6 @@ def test_standardize():
from matplotlib import pyplot as plt
def test_synthesize_reconstructions():
# I can only really test for a lack of failures, so I think my plan
# will be to create a dataset that just needs to be added and see that
@@ -154,13 +152,12 @@ def test_synthesize_reconstructions():
# Start by making a probe and object that should meet the standardization
# conditions
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01))
probe = cmath.torch_to_complex(probe)
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2))
probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
assert np.isclose(0,np.angle(np.sum(probe)))
assert np.abs(np.angle(np.sum(probe))) < 1e-7
obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
@@ -203,14 +200,14 @@ def test_calc_consistency_prtf():
assert np.allclose(prtf, 0.7)
# Check that it also works with torch input
t_synth_obj = cmath.complex_to_torch(synth_obj)
t_obj_stack = [cmath.complex_to_torch(obj) for obj in obj_stack]
t_synth_obj = t.as_tensor(synth_obj)
t_obj_stack = [t.as_tensor(obj) for obj in obj_stack]
freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, basis, nbins=30)
assert np.allclose(prtf.numpy(), 0.7)
# And also when the basis is in torch
t_synth_obj = cmath.complex_to_torch(synth_obj)
t_obj_stack = [cmath.complex_to_torch(obj) for obj in obj_stack]
t_synth_obj = t.as_tensor(synth_obj)
t_obj_stack = [t.as_tensor(obj) for obj in obj_stack]
freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, t.Tensor(basis), nbins=30)
assert np.allclose(prtf.numpy(), 0.7)
@@ -238,11 +235,11 @@ def test_calc_deconvolved_cross_correlation():
assert np.allclose(test_cor, np_cor)
# test with pytorch inputs
obj1_t = cmath.complex_to_torch(obj1)
obj2_t = cmath.complex_to_torch(obj2)
obj1_t = t.as_tensor(obj1)
obj2_t = t.as_tensor(obj2)
test_cor_t = analysis.calc_deconvolved_cross_correlation(obj1_t,obj2_t, im_slice=np.s_[:,:])
assert np.allclose(cmath.torch_to_complex(test_cor_t), np_cor)
assert np.allclose(test_cor_t.numpy(), np_cor)
@@ -294,8 +291,8 @@ def test_calc_frc():
assert np.allclose(threshold, test_threshold)
# try again with complex
obj1_torch = cmath.complex_to_torch(obj1)
obj2_torch = cmath.complex_to_torch(obj2)
obj1_torch = t.as_tensor(obj1)
obj2_torch = t.as_tensor(obj2)
basis_torch = t.tensor(basis)
test_bins_t, test_frc_t, test_threshold_t = analysis.calc_frc(obj1_torch,
tests/tools/test_image_processing.py
+19 -15
View File
@@ -32,8 +32,8 @@ def test_centroid_sq():
assert t.allclose(centroid, t.Tensor(sp_centroid))
# Test complex with multiple ims
ims = t.rand((5,30,40,2))
np_ims = cmath.torch_to_complex(ims)
ims = t.rand((5,30,40)) + 1j * t.rand((5,30,40))
np_ims = ims.numpy()
sp_centroids = [ndimage.measurements.center_of_mass(np.abs(im)**2)
for im in np_ims]
centroids = image_processing.centroid_sq(ims, comp=True)
@@ -51,12 +51,12 @@ def test_sinc_subpixel_shift():
Ys,Xs = np.meshgrid(xs,xs)
sinc_im = np.sinc(Xs-0.3) * np.sinc(Ys-0.6)
torch_im = cmath.complex_to_torch(im)
torch_im = t.as_tensor(im)
test_im = image_processing.sinc_subpixel_shift(torch_im,(0.3,0.6))
# The fidelity isn't great due to the FFT-based approach, so we need
# a pretty relaxed condition
assert np.max(np.abs(sinc_im - cmath.torch_to_complex(test_im))) < 0.005
assert np.max(np.abs(sinc_im - test_im.numpy())) < 0.005
def test_find_pixel_shift():
@@ -69,13 +69,13 @@ def test_find_pixel_shift():
# Test a real and complex im
big_im = t.rand((30,70))
im1 = t.stack((big_im[:-5,10:],t.zeros_like(big_im[:-5,10:])),dim=-1)
im1 = big_im[:-5,10:].to(dtype=t.complex64)
im2 = big_im[5:,:-10]
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,-10]))
assert t.all(image_processing.find_pixel_shift(im2,im1) == t.LongTensor([-5,10]))
# Test two complex ims
big_im = t.rand((45,45,2))
big_im = t.rand((45,45)) + 1j * t.rand((45,45))
im1 = big_im[:-5,:-4]
im2 = big_im[5:,4:]
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,4]))
@@ -83,13 +83,14 @@ def test_find_pixel_shift():
def test_find_subpixel_shift():
# We can do this by creating a test probe and a test object
test_probe = t.rand((70,70,2))
test_obj = t.ones((300,300,2))
test_probe = t.rand((70,70)) + 1j * t.rand((70,70))
test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
shift = t.tensor((0.8,0.75))
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift)
print(im.shape)
retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0,0), resolution=50)
# tolerance of 0.03 on this measurement
assert t.all(t.abs(shift - retrieved_shift) < 0.03)
@@ -98,8 +99,8 @@ def test_find_subpixel_shift():
def test_find_shift():
# We can do this by creating a test probe and a test object
test_probe = t.rand((200,200,2))
test_obj = t.ones((300,300,2))
test_probe = t.rand((200,200)) + 1j * t.rand((200,200))
test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
shift = t.tensor((0.8,0.75))
@@ -111,14 +112,14 @@ def test_find_shift():
def test_convolve_1d():
from matplotlib import pyplot as plt
test_image = np.random.rand(400,300)
#test_image = np.hstack((np.ones((400,150)),np.zeros((400,150))))
xs = np.linspace(-100,100,300)
kernel = 1/(1+xs**2)
# First, we test with everything real, dim=1
convolved = image_processing.convolve_1d(t.Tensor(test_image),t.Tensor(kernel),dim=1)
convolved = image_processing.convolve_1d(t.as_tensor(test_image),
t.as_tensor(kernel),dim=1)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1))
assert np.allclose(convolved.numpy(),np_result)
@@ -128,13 +129,16 @@ def test_convolve_1d():
kernel = 1/(1+xs**2)
# Then with dim=0, and a non-fftshifted kernel
convolved = image_processing.convolve_1d(t.Tensor(test_image),t.Tensor(np.fft.ifftshift(kernel)), fftshift_kernel=False)
convolved = image_processing.convolve_1d(t.as_tensor(test_image),
t.as_tensor(np.fft.ifftshift(kernel)),
fftshift_kernel=False)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0))
assert np.allclose(convolved.numpy(),np_result)
# And finally with complex input
convolved = cmath.torch_to_complex(image_processing.convolve_1d(cmath.complex_to_torch(test_image),cmath.complex_to_torch(kernel)))
convolved = image_processing.convolve_1d(t.as_tensor(test_image,dtype=t.complex64),
t.as_tensor(kernel,dtype=t.complex64)).numpy()
np_result = np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0)
assert np.allclose(convolved,np_result)
tests/tools/test_initializers.py
+16 -17
View File
@@ -77,11 +77,12 @@ def test_gaussian():
# Generate gaussian as a numpy array (square array)
shape = [10, 10]
sigma = [2.5, 2.5]
center = ((shape[0]-1)/2, (shape[1]-1)/2)
y, x = np.mgrid[:shape[0], :shape[1]]
np_result = 10*np.exp(-0.5*((x-center[1])/sigma[1])**2
-0.5*((y-center[0])/sigma[0])**2)
init_result = cmath.torch_to_complex(initializers.gaussian([10, 10], [2.5, 2.5], amplitude=10))
init_result = initializers.gaussian(shape, sigma, amplitude=10).numpy()
assert np.allclose(init_result, np_result)
# Generate gaussian as a numpy array (rectangular array)
@@ -91,7 +92,7 @@ def test_gaussian():
y, x = np.mgrid[:shape[0], :shape[1]]
np_result = np.exp(-0.5*((x-center[1])/sigma[1])**2
-0.5*((y-center[0])/sigma[0])**2)
init_result = cmath.torch_to_complex(initializers.gaussian(shape, sigma))
init_result = initializers.gaussian(shape, sigma).numpy()
assert np.allclose(init_result, np_result)
# Generate gaussian with curvature
@@ -104,8 +105,8 @@ def test_gaussian():
-0.5*((y-center[0])/sigma[0])**2)
np_result *= np.exp(0.5j*curvature[1]*(x-center[1])**2
+0.5j*curvature[0]*(y-center[0])**2)
init_result = cmath.torch_to_complex(initializers.gaussian(shape, sigma,
center=center, curvature=curvature, amplitude=10))
init_result = initializers.gaussian(shape, sigma, center=center,
curvature=curvature, amplitude=10).numpy()
assert np.allclose(init_result, np_result)
@@ -146,8 +147,7 @@ def test_gaussian_probe(ptycho_cxi_1):
normalization_1 = normalization / np.sum(np.abs(np_probe)**2)
probe = initializers.gaussian_probe(dataset, basis, shape, sigma)
probe = cmath.torch_to_complex(probe)
probe = initializers.gaussian_probe(dataset, basis, shape, sigma).numpy()
assert np.allclose(probe, normalization_1*np_probe)
# And then a propagated probe
@@ -162,8 +162,7 @@ def test_gaussian_probe(ptycho_cxi_1):
normalization_2 = normalization / np.sum(np.abs(np_probe)**2)
probe = initializers.gaussian_probe(dataset, basis, shape, sigma,
propagation_distance=z)
probe = cmath.torch_to_complex(probe)
propagation_distance=z).numpy()
assert np.allclose(probe, normalization_2*np_probe)
@@ -184,10 +183,10 @@ def test_SHARP_style_probe(ptycho_cxi_1):
distance)
probe = initializers.SHARP_style_probe(dataset, shape, det_slice)
assert probe.shape == t.Size([256,256,2])
assert probe.shape == t.Size([256,256])
probe = initializers.SHARP_style_probe(dataset, shape, det_slice, propagation_distance=20e-6)
assert probe.shape == t.Size([256,256,2])
assert probe.shape == t.Size([256,256])
def test_RPI_spectral_init():
@@ -199,24 +198,24 @@ def test_RPI_spectral_init():
probe = np.random.rand(230,253).astype(np.complex64)
obj_shape = [37,53]
mask = t.Tensor(np.random.rand(*pattern.shape) > 0.04)
background = t.Tensor(np.random.rand(*pattern.shape) .astype(np.float32)* 0.05)
background = t.as_tensor(np.random.rand(*pattern.shape),dtype=t.float32) * 0.05
probe = cmath.complex_to_torch(probe)
pattern = t.Tensor(pattern)
probe = t.as_tensor(probe)
pattern = t.as_tensor(pattern)
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape)
assert list(obj.shape) == [1]+obj_shape+[2]
assert list(obj.shape) == [1]+obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, mask=mask)
assert list(obj.shape) == [2]+obj_shape+[2]
assert list(obj.shape) == [2]+obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, background=background)
assert list(obj.shape) == [2]+obj_shape+[2]
assert list(obj.shape) == [2]+obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, mask=mask,
background=background)
assert list(obj.shape) == [2]+obj_shape+[2]
assert list(obj.shape) == [2]+obj_shape
tests/tools/test_interactions.py
+31 -31
View File
@@ -136,16 +136,16 @@ def test_ptycho_2D_round(random_probe, random_obj):
random_obj[tr[0]:tr[0]+random_probe.shape[0],
tr[1]:tr[1]+random_probe.shape[1]] for
tr in np.round(translations).astype(int)]
exit_waves_t = interactions.ptycho_2D_round(cmath.complex_to_torch(random_probe),
cmath.complex_to_torch(random_obj),
t.tensor(translations))
assert np.allclose(cmath.torch_to_complex(exit_waves_t), exit_waves_np)
exit_waves_t = interactions.ptycho_2D_round(t.as_tensor(random_probe),
t.as_tensor(random_obj),
t.as_tensor(translations))
assert np.allclose(exit_waves_t.numpy(), exit_waves_np)
# Test the single wave case
exit_wave_t = interactions.ptycho_2D_round(cmath.complex_to_torch(random_probe),
cmath.complex_to_torch(random_obj),
t.tensor(translations[0]))
assert np.allclose(cmath.torch_to_complex(exit_wave_t), exit_waves_np[0])
exit_wave_t = interactions.ptycho_2D_round(t.as_tensor(random_probe),
t.as_tensor(random_obj),
t.as_tensor(translations[0]))
assert np.allclose(exit_wave_t.numpy(), exit_waves_np[0])
@@ -157,14 +157,14 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
translation = np.array([46.7,53.2])
exit_waves_probe = interactions.ptycho_2D_linear(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.tensor(translations),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.as_tensor(translations),
shift_probe=True)
exit_wave_probe = interactions.ptycho_2D_linear(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.tensor(translation),
shift_probe=True)
@@ -173,14 +173,14 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
exit_waves_obj = interactions.ptycho_2D_linear(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.tensor(translations),
shift_probe=False)
exit_wave_obj = interactions.ptycho_2D_linear(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.tensor(translation),
shift_probe=False)
@@ -188,7 +188,7 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
assert t.allclose(exit_waves_obj[0],exit_wave_obj)
# For the shifted probe, we should find 4 pixels with intensity
exit_waves_probe = cmath.torch_to_complex(exit_waves_probe)[0]
exit_waves_probe = t.as_tensor(exit_waves_probe)[0]
probe_shift = np.array([[0.3*0.8,0.3*0.2],
[0.7*0.8,0.7*0.2]])
@@ -198,7 +198,7 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
assert np.allclose(probe_shift * obj_section, exit_section)
# For the shifted obj, we should find one pixel with intensity
exit_waves_obj = cmath.torch_to_complex(exit_waves_obj)[0]
exit_waves_obj = t.as_tensor(exit_waves_obj)[0]
obj_shift = np.array([[0.3*0.8,0.3*0.2],
[0.7*0.8,0.7*0.2]])
obj_section = random_obj[128+46:128+48,
@@ -218,15 +218,15 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
translation = np.array([46.7,53.2])
exit_waves_probe = interactions.ptycho_2D_sinc(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.tensor(translations),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.as_tensor(translations),
shift_probe=True)
exit_wave_probe = interactions.ptycho_2D_sinc(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.tensor(translation),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.as_tensor(translation),
shift_probe=True)
# Check that the outputs match
@@ -246,7 +246,7 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
53:53+256]
exit_wave_np = sinc_shifted_probe * obj_section
exit_wave_torch = cmath.torch_to_complex(exit_wave_probe)
exit_wave_torch = exit_wave_probe.numpy()
# The fidelity isn't great due to the FFT-based approach, so we need
# a pretty relaxed condition
@@ -258,8 +258,8 @@ def test_RPI_interaction(random_probe, random_obj):
random_obj1 = random_obj[:79,:68]
random_probe1 = random_probe
t_random_obj1 = cmath.complex_to_torch(random_obj1)
t_random_probe1 = cmath.complex_to_torch(random_probe1)
t_random_obj1 = t.as_tensor(random_obj1)
t_random_probe1 = t.as_tensor(random_probe1)
t_output1 = interactions.RPI_interaction(t_random_probe1, t_random_obj1)
obj1_fourier = fftshift(fft.fft2(ifftshift(random_obj1), norm='ortho'))
@@ -271,13 +271,13 @@ def test_RPI_interaction(random_probe, random_obj):
output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups),
norm='ortho'))
assert np.allclose(cmath.torch_to_complex(t_output1), output1)
assert np.allclose(t.as_tensor(t_output1), output1)
random_obj2 = np.stack([random_obj[:64,:89]]*3)
random_probe2 = random_probe[3:,5:]
t_random_obj2 = cmath.complex_to_torch(random_obj2)
t_random_probe2 = cmath.complex_to_torch(random_probe2)
t_random_obj2 = t.as_tensor(random_obj2)
t_random_probe2 = t.as_tensor(random_probe2)
t_output2 = interactions.RPI_interaction(t_random_probe2, t_random_obj2)
obj2_fourier = fftshift(fft.fft2(ifftshift(random_obj2), norm='ortho'))
tests/tools/test_losses.py
+3 -3
View File
@@ -15,7 +15,7 @@ def test_amplitude_mse():
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(np.bool)
mask = (np.random.rand(100,100) > 0.1).astype(bool)
# First, test without a mask
np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2)
@@ -37,7 +37,7 @@ def test_intensity_mse():
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(np.bool)
mask = (np.random.rand(100,100) > 0.1).astype(bool)
# First, test without a mask
@@ -60,7 +60,7 @@ def test_poisson_nll():
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(np.bool)
mask = (np.random.rand(100,100) > 0.1).astype(bool)
# First, test without a mask
tests/tools/test_measurements.py
+23 -23
View File
@@ -7,24 +7,24 @@ import pytest
def test_intensity():
wavefields = t.rand((5,10,10,2))
wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10))
epsilon=1e-6
np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + epsilon
np_result = np.abs(t.as_tensor(wavefields))**2 + epsilon
assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon),
t.tensor(np_result))
t.as_tensor(np_result))
# Test single field case
assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon),
t.tensor(np_result[0]))
t.as_tensor(np_result[0]))
det_slice = np.s_[3:,5:8]
assert t.allclose(measurements.intensity(wavefields,det_slice,epsilon=epsilon),
t.tensor(np_result[(np.s_[:],)+det_slice]))
t.as_tensor(np_result[(np.s_[:],)+det_slice]))
# Test single field case
assert t.allclose(measurements.intensity(wavefields[0],det_slice,epsilon=epsilon),
t.tensor(np_result[0][det_slice]))
t.as_tensor(np_result[0][det_slice]))
# With oversampling on
@@ -35,34 +35,34 @@ def test_intensity():
# With multiple fields
assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon, oversampling=2),
t.tensor(np_oversampling_result,))
t.as_tensor(np_oversampling_result,))
# With a single field
assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon, oversampling=2),
t.tensor(np_oversampling_result[0],))
t.as_tensor(np_oversampling_result[0],))
def test_incoherent_sum():
# With no explicit slice given
wavefields = t.rand((5,4,10,10,2))
wavefields = t.rand((5,4,10,10)) + 1j * t.rand((5,4,10,10))
epsilon=1e-6
np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0) + epsilon
np_result = np.sum(np.abs(wavefields.numpy())**2,axis=-3) + epsilon
assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon),
t.tensor(np_result))
t.as_tensor(np_result))
# Test single field case
assert t.allclose(measurements.incoherent_sum(wavefields[:,0],epsilon=epsilon),
t.tensor(np_result[0]))
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon),
t.as_tensor(np_result[0]))
# With a slice given
det_slice = np.s_[3:,5:8]
assert t.allclose(measurements.incoherent_sum(wavefields,det_slice,epsilon=epsilon),
t.tensor(np_result[(np.s_[:],)+det_slice]))
t.as_tensor(np_result[(np.s_[:],)+det_slice]))
# Test single field case
assert t.allclose(measurements.incoherent_sum(wavefields[:,0],det_slice,epsilon=epsilon),
t.tensor(np_result[0][det_slice]))
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],det_slice,epsilon=epsilon),
t.as_tensor(np_result[0][det_slice]))
# With oversampling on
np_oversampling_result = (np_result[:,::2,::2] + \
@@ -72,19 +72,19 @@ def test_incoherent_sum():
# With multiple fields
assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon, oversampling=2),
t.tensor(np_oversampling_result,))
t.as_tensor(np_oversampling_result,))
# With a single field
assert t.allclose(measurements.incoherent_sum(wavefields[:,0],epsilon=epsilon, oversampling=2),
t.tensor(np_oversampling_result[0],))
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon, oversampling=2),
t.as_tensor(np_oversampling_result[0],))
def test_quadratic_background():
# test with intensity
wavefields = t.rand((5,10,10,2))
wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10))
epsilon=1e-6
background = t.rand((10,10))
np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + background.numpy()**2 + epsilon
np_result = np.abs(wavefields.numpy())**2 + background.numpy()**2 + epsilon
det_slice = np.s_[3:,5:8]
result = measurements.quadratic_background(wavefields,background[det_slice],
@@ -95,8 +95,8 @@ def test_quadratic_background():
# test with incoherent sum but no slice and no stack
wavefields = t.rand((4,10,10,2))
np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0)
wavefields = t.rand((4,10,10)) + 1j * t.rand((4,10,10))
np_result = np.sum(np.abs(wavefields.numpy())**2,axis=0)
np_result += background.numpy()**2
result = measurements.quadratic_background(wavefields, background,
epsilon=epsilon,
tests/tools/test_plotting.py
+4 -5
View File
@@ -10,7 +10,7 @@ import matplotlib.pyplot as plt
def test_plot_amplitude(show_plot):
# Test with tensor
im = cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64))
im = t.as_tensor(scipy.misc.ascent(),dtype=t.complex128)
plotting.plot_amplitude(im, basis = np.array([[1,1], [1,1], [0,0]]), title = 'Test Amplitude')
if show_plot:
plt.show()
@@ -30,7 +30,7 @@ def test_plot_phase(show_plot):
plt.show()
# Test with numpy array
im = cmath.torch_to_complex(initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]))
im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]).numpy()
plotting.plot_phase(im, title = 'Test Phase', basis = np.array([[1,1], [1,1], [0,0]]))
if show_plot:
plt.show()
@@ -38,14 +38,13 @@ def test_plot_phase(show_plot):
def test_plot_colorize(show_plot):
# Test with tensor
gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
im = cmath.cmult(gaussian, cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64)))
im = gaussian * t.as_tensor(scipy.misc.ascent(), dtype=t.complex64)
plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[1,1], [1,1], [0,0]]))
if show_plot:
plt.show()
# Test with numpy array
gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
im = cmath.torch_to_complex(cmath.cmult(gaussian, cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64))))
im = im.numpy()
plotting.plot_colorized(im, title = 'Test Colorize')
if show_plot:
plt.show()
tests/tools/test_projectors.py
-38
View File
@@ -1,38 +0,0 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools import projectors
import numpy as np
import torch as t
def test_modulus():
# Create a complex array with random modulus and known phase
np_result = np.sqrt(6) * (1 + 1j) * np.random.rand(10,10)
projection_intensity = t.from_numpy(np.abs(np_result)**2).to(t.float32)
original_wavefront = cmath.complex_to_torch((1+1j) * np.random.rand(10,10)).to(t.float32)
# Test without masks
torch_result = projectors.modulus(original_wavefront,projection_intensity)
assert np.allclose(cmath.torch_to_complex(torch_result),np_result)
# Test with mask
mask = t.ones((10,10,2), dtype = t.uint8)
mask[5]*=0
np_result[5] = cmath.torch_to_complex(original_wavefront[5])
torch_result = projectors.modulus(original_wavefront,projection_intensity, mask=mask)
print(np_result[5])
print(cmath.torch_to_complex(torch_result)[5])
assert np.allclose(cmath.torch_to_complex(torch_result),np_result)
def test_support():
# Define a mask as uint8 to make sure it works when support is not
# the same type as the wavefield
support = t.zeros((10,10)).to(t.uint8)
# some masked, some unmasked
support[:3,:3] = 1
np_result = np.zeros((10,10)).astype(np.complex128)
np_result[:3,:3] = 1 + 1j
assert(np.allclose(cmath.torch_to_complex(projectors.support(t.ones((10,10,2)), support)), np_result))
tests/tools/test_propagators.py
+31 -36
View File
@@ -18,26 +18,26 @@ def exit_waves_1():
obj = scipy.misc.ascent()[0:64,0:64].astype(np.complex128)
arr = np.random.random_sample((64,64))
obj *= (arr+(1-arr**2)**.5*1j)
obj = cmath.complex_to_torch(obj)
obj = t.as_tensor(obj)
# Construct wavefront from image
probe = initializers.gaussian([64, 64], [5, 5], amplitude=1e3)
return cmath.cmult(probe,obj)
return probe * obj
def test_far_field(exit_waves_1):
# Far field diffraction patterns calculated by numpy with zero frequency in center
np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(cmath.torch_to_complex(exit_waves_1)),norm='ortho'))
np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho'))
assert(np.allclose(np_result, cmath.torch_to_complex(propagators.far_field(exit_waves_1))))
assert(np.allclose(np_result, propagators.far_field(exit_waves_1).numpy()))
def test_inverse_far_field(exit_waves_1):
# We want the inverse far field to map back to the exit waves with no intensity corrections
# Far field result for exit waves calculated with numpy
far_field_np_result = cmath.complex_to_torch(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(cmath.torch_to_complex(exit_waves_1)),norm='ortho')))
far_field_np_result = t.as_tensor(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho')))
assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
@@ -66,14 +66,14 @@ def test_generate_high_NA_k_intensity_map():
j = (np.arange(573) - 270)
Is,Js = np.meshgrid(i,j,indexing='ij')
wavefield = ((np.abs(Is) < 20) * (np.abs(Js) < 25)).astype(np.complex128)
t_wavefield = cmath.complex_to_torch(wavefield).to(dtype=t.float32)
t_wavefield = t.as_tensor(wavefield, dtype=t.complex64)
high_NA_propagated = propagators.high_NA_far_field(
t_wavefield, k_map, intensity_map=intensity_map)
low_NA_propagated = propagators.far_field(t_wavefield)
low_NA = cmath.torch_to_complex(low_NA_propagated)
high_NA = cmath.torch_to_complex(high_NA_propagated)
low_NA = low_NA_propagated.numpy()
high_NA = high_NA_propagated.numpy()
# Checking first that for a low-NA propagation they give the same result
# 1e-4 tolerance seems to be reasonable in this comparison given my
@@ -99,8 +99,8 @@ def test_generate_high_NA_k_intensity_map():
t_wavefield, k_map, intensity_map=intensity_map)
low_NA_propagated = propagators.far_field(t_wavefield)
low_NA = cmath.torch_to_complex(low_NA_propagated)
high_NA = cmath.torch_to_complex(high_NA_propagated)
low_NA = low_NA_propagated.numpy()
high_NA = high_NA_propagated.numpy()
#plt.close('all')
#plt.imshow(np.abs(low_NA))
@@ -120,7 +120,7 @@ def test_generate_high_NA_k_intensity_map():
# I believe this works, but I still would like to get a second method for
# simulating at least one diffraction pattern as an independent check
assert 0
#assert 0
@@ -156,11 +156,10 @@ def test_near_field():
# First we check it normally
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.float64)
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
# Check for at least 10^-3 relative accuracy in this scenario
assert np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez))
@@ -168,8 +167,7 @@ def test_near_field():
Emz = np.conj(Ez)
Emz_t = propagators.inverse_near_field(cmath.complex_to_torch(E0),asp)
Emz_t = cmath.torch_to_complex(Emz_t)
Emz_t = propagators.inverse_near_field(t.as_tensor(E0),asp).numpy()
# Again, 10^-3 is about all the accuracy we can expect
assert np.max(np.abs(Emz-Emz_t)) < 1e-3 * np.max(np.abs(Emz))
@@ -177,11 +175,10 @@ def test_near_field():
# Then, we check it with the phase correction
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,remove_z_phase=True,
dtype=t.float64)
dtype=t.complex128)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
# Check for at least 10^-3 relative accuracy in this scenario
assert np.max(np.abs(Ez_nozphase-Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase))
@@ -189,8 +186,7 @@ def test_near_field():
Emz = np.conj(Ez_nozphase)
Emz_t = propagators.inverse_near_field(cmath.complex_to_torch(E0),asp)
Emz_t = cmath.torch_to_complex(Emz_t)
Emz_t = propagators.inverse_near_field(t.as_tensor(E0),asp).numpy()
# Again, 10^-3 is about all the accuracy we can expect
assert np.max(np.abs(Emz-Emz_t)) < 1e-3 * np.max(np.abs(Emz))
@@ -337,22 +333,22 @@ def test_generalized_near_field():
E0 = get_E(Xs,Ys,Zs_0, correct=False)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=False)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,dtype=t.float64)
E0.shape,new_basis,wavelength,offset_vec,dtype=t.complex128)
elif str(propagation_vec) == 'offset':
E0 = get_E(Xs,Ys,Zs_0, correct=True)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,
dtype=t.float64, propagate_along_offset=True)
dtype=t.complex128, propagate_along_offset=True)
else:
E0 = get_E(Xs,Ys,Zs_0, correct=True)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,
dtype=t.float64, propagation_vector=propagation_vec)
dtype=t.complex128, propagation_vector=propagation_vec)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
# Check for at least 10^-3 relative accuracy in this scenario
if not np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez)):
#if True:
@@ -367,8 +363,7 @@ def test_generalized_near_field():
assert np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez))
Em0_t = propagators.inverse_near_field(cmath.complex_to_torch(Ez),asp)
Em0_t = cmath.torch_to_complex(Em0_t)
Em0_t = propagators.inverse_near_field(t.as_tensor(Ez),asp).numpy()
# Again, 10^-3 is about all the accuracy we can expect
assert np.max(np.abs(E0-Em0_t)) < 1e-3 * np.max(np.abs(E0))
@@ -387,9 +382,8 @@ def test_generalized_near_field():
E0 = get_E(Xs,Ys,Zs_0, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape, new_basis, wavelength,offset_vec,
dtype=t.float64, propagation_vector=propagation_vec)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
dtype=t.complex128, propagation_vector=propagation_vec)
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
for Rrand in Rrands:
@@ -399,9 +393,9 @@ def test_generalized_near_field():
rot_prop = np.dot(Rrand, propagation_vec)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape, rot_basis, wavelength, rot_offset,
dtype=t.float64, propagation_vector=rot_prop)
Ez_rot_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_rot_t = cmath.torch_to_complex(Ez_rot_t)
dtype=t.complex128, propagation_vector=rot_prop)
Ez_rot_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
assert np.max(np.abs(Ez_t-Ez_rot_t)) < 1e-3 * np.max(np.abs(Ez_t))
@@ -420,9 +414,10 @@ def test_inverse_near_field():
E0 = np.exp(-Rs**2 / w0**2)
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.float64)
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
E0 = cmath.complex_to_torch(E0)
E0 = t.as_tensor(E0,dtype=t.complex128)
E_prop = propagators.near_field(E0,asp)
E_backprop = propagators.inverse_near_field(E_prop, asp)