diff --git a/CDTools/models/bragg_2d_ptycho.py b/CDTools/models/bragg_2d_ptycho.py index a1298bb..b0e827b 100644 --- a/CDTools/models/bragg_2d_ptycho.py +++ b/CDTools/models/bragg_2d_ptycho.py @@ -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) diff --git a/CDTools/tools/propagators.py b/CDTools/tools/propagators.py index 1321bde..5af98b3 100644 --- a/CDTools/tools/propagators.py +++ b/CDTools/tools/propagators.py @@ -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 diff --git a/converters/TITAN_stem.py b/converters/TITAN_stem.py index d5ac537..3c06f79 100644 --- a/converters/TITAN_stem.py +++ b/converters/TITAN_stem.py @@ -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) diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index 37351eb..40abea7 100644 --- a/tests/tools/test_propagators.py +++ b/tests/tools/test_propagators.py @@ -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))