mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 21:42:39 +02:00
One more try to fix the off axis near field propagation
This commit is contained in:
Abe Levitan
@@ -167,6 +167,8 @@ class Bragg2DPtycho(CDIModel):
|
||||
else:
|
||||
self.k_map = None
|
||||
self.intensity_map = None
|
||||
|
||||
self.prop_dir = t.Tensor([0,0,1]).to(dtype=t.float32)
|
||||
|
||||
|
||||
@classmethod
|
||||
@@ -356,7 +358,8 @@ class Bragg2DPtycho(CDIModel):
|
||||
for j in range(translations.size()[0]):
|
||||
if self.propagate_probe:
|
||||
propagator = ggasp(pr.shape, self.probe_basis, self.wavelength,
|
||||
t.Tensor([0,0,0*props[j]]),
|
||||
t.Tensor([0,0,props[j]]),
|
||||
propagation_vector=self.prop_dir,
|
||||
dtype=pr.dtype,device=pr.device, propagate_along_offset=True)
|
||||
prop_pr = tools.propagators.near_field(pr, propagator)
|
||||
|
||||
@@ -435,6 +438,7 @@ class Bragg2DPtycho(CDIModel):
|
||||
self.probe_support = self.probe_support.to(*args,**kwargs)
|
||||
self.obj_support = self.obj_support.to(*args,**kwargs)
|
||||
self.surface_normal = self.surface_normal.to(*args, **kwargs)
|
||||
self.prop_dir = self.prop_dir.to(*args, **kwargs)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -306,9 +306,13 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
|
||||
# Define this as complex so the square root properly gives
|
||||
# k>k0 components imaginary frequencies
|
||||
k0 = np.complex128((2*np.pi/wavelength))
|
||||
|
||||
propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * z)
|
||||
|
||||
# Properly accuount for evanescent waves
|
||||
if z >=0:
|
||||
propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * z)
|
||||
else:
|
||||
propagator = np.exp(1j*np.conj(np.sqrt(k0**2 - Ki**2 - Kj**2)) * z)
|
||||
|
||||
if remove_z_phase:
|
||||
propagator *= np.exp(-1j * k0 * z)
|
||||
|
||||
@@ -353,6 +357,35 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
|
||||
vector will be set equal to the offset vector. This overrides the
|
||||
propagation_vector option
|
||||
|
||||
Note that, unlike in the case of the simple angular spectrum propagator,
|
||||
the direction of "forward propagation" is defined by the offset vector.
|
||||
Therefore, in the simple case of a perpendicular offset, there will be
|
||||
no difference between using an offset vector or the negative of the
|
||||
offset vector. This is because, for the light propagation problem to
|
||||
be well posed, the assumption must be made that light only passes through
|
||||
the plane of the known wavefield in one direction. Mathematically, this
|
||||
corresponds to a choice of uniform phase objects either accumulating
|
||||
positive or negative phase. In the simple propagation case, there is
|
||||
no ambiguity introduced by always choosing the light field to propagate
|
||||
along the positive z direction. In the general case, there is no equivalent
|
||||
obvious choice - thus, the light is always assumed to pass through the
|
||||
initial plane travelling in the direction of the final plane.
|
||||
|
||||
Practically, if one wants to simulate inverse propagation, there are then
|
||||
two possible approaches. First, one can use the inverse_near_field
|
||||
function, which simulates the inverse propagation problem and therefore
|
||||
will naturally simulate propagation in the opposite direction. Second,
|
||||
one can explicitly include a propagation_vector argument, which overrides
|
||||
the offset vector in defining the direction in which light passes through
|
||||
the input plane. However, in this case, the resulting light field will have
|
||||
the overall phase accumulation due to propagation along the propagation
|
||||
vector removed, which may not be the intended behavior. However, this is
|
||||
not recommended, as inverse propagation will tend to magnify evanescent
|
||||
waves - it is therefore preferable (unless there is a specific need to
|
||||
account for evanescent waves properly) to use the inverse near field
|
||||
propagator
|
||||
|
||||
|
||||
Parameters
|
||||
----------
|
||||
shape : array
|
||||
@@ -420,7 +453,6 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
|
||||
|
||||
# This may have a sign error - must be checked
|
||||
phase_mask = np.exp(1j * np.tensordot(offset_vector,K_xyz,axes=1))
|
||||
|
||||
|
||||
# Next, we apply a shift to the k-space vectors which sets up
|
||||
# propagation such that a uniform phase object will propagate along the
|
||||
@@ -434,13 +466,15 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
|
||||
perpendicular_dir = np.cross(basis[:,1],basis[:,0])
|
||||
perpendicular_dir /= np.linalg.norm(perpendicular_dir)
|
||||
offset_perpendicular = np.dot(perpendicular_dir, offset_vector)
|
||||
|
||||
|
||||
|
||||
k0 = 2*np.pi/wavelength
|
||||
|
||||
sign_correction = 1
|
||||
|
||||
# Only implement the shift if the flag is set to True
|
||||
if propagation_vector is not None:
|
||||
propagation_vector = propagation_vector / np.linalg.norm(propagation_vector)
|
||||
|
||||
prop_perpendicular = np.dot(perpendicular_dir, propagation_vector)
|
||||
prop_parallel = propagation_vector - perpendicular_dir \
|
||||
* prop_perpendicular
|
||||
@@ -451,28 +485,39 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
|
||||
# a special case
|
||||
k_offset = np.array([0,0,0])
|
||||
else:
|
||||
k_offset = prop_parallel * k0
|
||||
k_offset = prop_parallel * k0 / np.linalg.norm(propagation_vector)
|
||||
|
||||
K_xyz = K_xyz + k_offset[:,None,None]
|
||||
|
||||
# There apparently is a sign correction that I need to apply
|
||||
sign_correction = np.sign(np.dot(perpendicular_dir,propagation_vector))
|
||||
#sign_correction = np.sign(np.dot(perpendicular_dir,propagation_vector))
|
||||
sign_correction = np.sign(np.dot(offset_vector,propagation_vector))
|
||||
|
||||
K_xyz = K_xyz + k_offset[:,None,None] * sign_correction
|
||||
|
||||
# we also need to remove the z-dependence on the phase
|
||||
# This time, though, the z-dependence actually has to do with
|
||||
# the out of plane component of k at the central offset. Normally
|
||||
# this is 0, so the z-component is just k0, but not in this case
|
||||
# I need to understand this better I think
|
||||
# We only need one case here, unlike with the propagator, because
|
||||
# k_offset will always be less than k0
|
||||
phase_mask *= np.exp(-1j * np.sqrt(k0**2 - np.linalg.norm(k_offset)**2)
|
||||
* offset_perpendicular)
|
||||
* sign_correction
|
||||
* np.abs(offset_perpendicular))
|
||||
|
||||
|
||||
# Redefine this as complex so the square root properly gives
|
||||
# k>k0 components imaginary frequencies
|
||||
k0 = np.complex128(k0)
|
||||
|
||||
|
||||
# Finally, generate the propagator!
|
||||
propagator = np.exp(1j*np.sqrt(k0**2 - np.linalg.norm(K_xyz,axis=0)**2)
|
||||
* offset_perpendicular)
|
||||
# Must have cases to ensure that evanescent waves decay instead of grow
|
||||
if sign_correction > 0:
|
||||
propagator = np.exp(1j*np.sqrt(k0**2 - np.linalg.norm(K_xyz,axis=0)**2)
|
||||
* sign_correction * np.abs(offset_perpendicular))
|
||||
else:
|
||||
propagator = np.exp(-1j * np.conj(np.sqrt(k0**2 -
|
||||
np.linalg.norm(K_xyz,axis=0)**2))
|
||||
* np.abs(offset_perpendicular))
|
||||
propagator *= phase_mask
|
||||
|
||||
|
||||
@@ -519,8 +564,20 @@ def inverse_near_field(wavefront, angular_spectrum_propagator):
|
||||
using the supplied angular spectrum propagator, which is a premade
|
||||
phase mask.
|
||||
|
||||
It propagates the wave using the conjugate of the supplied phase mask,
|
||||
which corresponds to the inverse propagation problem.
|
||||
It propagates the wave using the complex conjugate of the supplied
|
||||
phase mask. This corresponds to propagation backward across the original
|
||||
propagation region - however, the treatment of evanescent waves is such
|
||||
that evanescent waves will decay both during the forward propagation and
|
||||
inverse propagation. This is done for reasons of numerical stability,
|
||||
as the choice to magnify evanescent waves during the inverse propagation
|
||||
process will quickly lead to magnification of any small amount of noise at
|
||||
frequencies larger than k_0, and in most typical situations will even
|
||||
lead to overflow of the floating point range. If evanescent waves need
|
||||
to be treated appropriately for any reason, it is recommended to use the
|
||||
"magnify_evanescent" option in the appropriate helper function used to
|
||||
generate the propagation phase mask. In this case, evanescent waves will
|
||||
be magnified both when used with the forward and inverse near field
|
||||
functions
|
||||
|
||||
|
||||
Parameters
|
||||
|
||||
@@ -14,7 +14,7 @@ from CDTools.datasets import Ptycho2DDataset
|
||||
from matplotlib import pyplot as plt
|
||||
from scipy.spatial.transform import Rotation
|
||||
from datetime import datetime
|
||||
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
def load_raw_image_stack(filename):
|
||||
# The resulting data is an array of (exposure, image-i, image-j),
|
||||
@@ -28,6 +28,32 @@ def load_raw_image_stack(filename):
|
||||
# One of the directions seems to be flipped
|
||||
return rawdata.reshape(numshots,130,128)[:,:128,:][:,:,::-1].copy()
|
||||
|
||||
def load_metadata(filename):
|
||||
return ET.parse(filename).getroot()
|
||||
|
||||
def get_scan_shape(metadata):
|
||||
sp = metadata.find("scan_parameters[@mode='acquire']")
|
||||
# print(sp)
|
||||
# exit()
|
||||
shape_x = int(sp.find('scan_resolution_x').text)
|
||||
shape_y = int(sp.find('scan_resolution_y').text)
|
||||
return [shape_x,shape_y]
|
||||
|
||||
def get_camera_length(metadata):
|
||||
iomm = metadata.find('iom_measurements')
|
||||
ncl = iomm.find('nominal_camera_length')
|
||||
return float(ncl.text)
|
||||
|
||||
def get_scan_steps(metadata):
|
||||
shape = get_scan_shape(metadata)
|
||||
iomm = metadata.find('iom_measurements')
|
||||
fov = iomm.find('full_scan_field_of_view')
|
||||
xfov = float(fov.find('x').text)
|
||||
yfov = float(fov.find('y').text)
|
||||
return np.array([xfov,yfov]) / np.array(shape)
|
||||
|
||||
|
||||
|
||||
def gen_scan_grid(shape, step):
|
||||
ys, xs = np.mgrid[:shape[0],:shape[1]]
|
||||
xs = xs * step[0]
|
||||
@@ -49,29 +75,50 @@ def generate_dataset(translations, patterns, detector_geometry, electron_energy)
|
||||
wavelength = calculate_wavelength(electron_energy)
|
||||
print(wavelength)
|
||||
return Ptycho2DDataset(translations, patterns, wavelength=wavelength, detector_geometry=det_geo)
|
||||
|
||||
|
||||
|
||||
|
||||
data_folder = '/media/Data Bank/ptychography_firsttry/out_of_focus_58Mx_1ms_reso80x80_ss1'
|
||||
image_filename = 'scan_x80_y80.raw'
|
||||
save_filename = 'test_defocus.cxi'
|
||||
save_filename = 'test_defocus_newcalibration.cxi'
|
||||
metadata_filename = 'out_of_focus_58Mx_1ms_reso80x80_ss1.xml'
|
||||
|
||||
|
||||
#data_folder = '/media/Data Bank/ptychography_firsttry/acquisition_3'
|
||||
#image_filename = 'scan_x80_y80.raw'
|
||||
#save_filename = 'test_acq3.cxi'
|
||||
#save_filename = 'test_acq3_newcalibration.cxi'
|
||||
#metadata_filename = 'acquisition_3.xml'
|
||||
|
||||
scan_shape = 80
|
||||
|
||||
metadata = load_metadata(data_folder + '/' + metadata_filename)
|
||||
|
||||
scan_shape = get_scan_shape(metadata)
|
||||
#scan_shape = 80
|
||||
|
||||
# These are reasonable initial guesses, until we get calibration data
|
||||
scan_step = 0.2e-10 #Angstrom
|
||||
pixel_pitches = [150e-6,150e-6]
|
||||
detector_distance = 100e-3 # mm
|
||||
#scan_step = 0.2e-10 #Angstrom, old value from manual measurement
|
||||
scan_steps = get_scan_steps(metadata)
|
||||
|
||||
# This is something I can calculate from the detector length
|
||||
# A good calibration is to assume that the pixel size is 0.2276 mm and
|
||||
# the detector distance is equal to the nomninal camera length
|
||||
camera_length = get_camera_length(metadata)
|
||||
detector_distance = camera_length
|
||||
pixel_pitches = [0.231e-3,0.231e-3] # best guess near length=0.230
|
||||
# pixel_pitches = [0.2276e-3,0.2276e-3] # best overall average
|
||||
|
||||
# old manual calibration
|
||||
#pixel_pitches = [150e-6,150e-6]
|
||||
#detector_distance = 100e-3 # mm
|
||||
#print([pp / detector_distance for pp in pixel_pitches])
|
||||
#exit()
|
||||
|
||||
electron_energy = 200 * 1.602e-16 # Joules
|
||||
# Important question: Check which side the images fill in from
|
||||
|
||||
data = load_raw_image_stack(data_folder + '/' + image_filename)
|
||||
#data[:,30:-30,30:-30] = 0 # For HAADF
|
||||
scan_points = gen_scan_grid([scan_shape,scan_shape],[scan_step, scan_step])
|
||||
scan_points = gen_scan_grid(scan_shape,scan_steps)
|
||||
|
||||
det_geo = generate_detector_geometry(detector_distance, pixel_pitches)
|
||||
dataset = generate_dataset(scan_points[1:], data[1:], det_geo, electron_energy)
|
||||
|
||||
@@ -178,7 +178,7 @@ def test_near_field():
|
||||
asp = propagators.generate_angular_spectrum_propagator(
|
||||
E0.shape,(1.5e-9,1e-9),wavelength,z,remove_z_phase=True,
|
||||
dtype=t.float64)
|
||||
|
||||
|
||||
|
||||
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
|
||||
Ez_t = cmath.torch_to_complex(Ez_t)
|
||||
@@ -266,9 +266,21 @@ def test_generalized_near_field():
|
||||
Rshear = np.array([[1,shear,0],
|
||||
[0,1,0],
|
||||
[0,0,1]])
|
||||
|
||||
|
||||
# This tests an inversion of the axes
|
||||
Rinv = np.array([[-1,0,0],
|
||||
[0,-1,0],
|
||||
[0,0,-1]])
|
||||
|
||||
# This tests a reflection about the y-z plane
|
||||
Rrefl = np.array([[-1,0,0],
|
||||
[0,1,0],
|
||||
[0,0,-1]])
|
||||
|
||||
# This tests a shearing and a rotation together
|
||||
Rall = np.matmul(Rboth,Rshear)
|
||||
Rall = np.matmul(Rrefl,np.matmul(Rboth,Rshear))
|
||||
|
||||
|
||||
# And we make some propagation vectors to test:
|
||||
|
||||
@@ -277,22 +289,27 @@ def test_generalized_near_field():
|
||||
|
||||
# This checks that it's not sensitive to the magnitude
|
||||
z_dir_large = np.array([0,0,10])
|
||||
|
||||
|
||||
# And finally some offset vectors
|
||||
|
||||
# This checks straight ahead
|
||||
z_offset = np.array([0,0,z])
|
||||
|
||||
|
||||
# This checks with an offset in x and y
|
||||
shear_offset = np.array([0.1*z,-0.03*z,z])
|
||||
|
||||
# This checks with an offset in x and y, with negative z
|
||||
shear_back_offset = np.array([0.1*z,-0.03*z,-z])
|
||||
|
||||
|
||||
|
||||
rot_mats = [I,I,I,Rboth, Rboth,Rboth, Rall, Rall, Rall]
|
||||
offset_vecs = [z_offset]*8 + [shear_offset]
|
||||
rot_mats = [Rrefl,I,Rinv, Rboth, Rboth,Rboth, Rall, Rall, Rall, I, Rall]
|
||||
offset_vecs = [z_offset]*8 + [shear_offset] + [shear_back_offset]*2
|
||||
propagation_vecs = ['perp','offset',z_dir,
|
||||
'perp','offset',z_dir_large,
|
||||
'perp','offset',z_dir_large]
|
||||
purposes = ['standard']*3 + ['both-rot']*3 + ['shear-rot']*3
|
||||
'perp','offset',z_dir_large,
|
||||
z_dir, z_dir_large]
|
||||
purposes = ['standard']*3 + ['both-rot']*3 + ['shear-rot']*3 + ['backward']*2
|
||||
|
||||
for purpose,rot_mat,offset_vec, propagation_vec \
|
||||
in zip(purposes,rot_mats,offset_vecs,propagation_vecs):
|
||||
@@ -325,14 +342,14 @@ def test_generalized_near_field():
|
||||
|
||||
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
|
||||
Ez_t = cmath.torch_to_complex(Ez_t)
|
||||
|
||||
# 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)):
|
||||
# plt.close('all')
|
||||
# plt.imshow(np.abs(Ez))
|
||||
# plt.figure()
|
||||
# plt.imshow(np.abs(Ez_t))
|
||||
# plt.show()
|
||||
if not np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez)):
|
||||
plt.close('all')
|
||||
plt.imshow(np.angle(Ez))
|
||||
plt.figure()
|
||||
plt.imshow(np.angle(Ez_t))
|
||||
plt.show()
|
||||
|
||||
assert np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez))
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user