From 14114f2bd4a799ee4b8346dc9db2ab59853da4f1 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 4 Aug 2021 22:32:58 -0400 Subject: [PATCH] Add tests to confirm that that the generation of angular spectrum propagators works with AD (and fix gasp so it indeed is) --- CDTools/tools/propagators/propagators.py | 18 +++++++++++------- tests/tools/test_propagators.py | 16 +++++++++++++++- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/CDTools/tools/propagators/propagators.py b/CDTools/tools/propagators/propagators.py index 2d517ab..1bec5f7 100644 --- a/CDTools/tools/propagators/propagators.py +++ b/CDTools/tools/propagators/propagators.py @@ -366,9 +366,9 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r basis[1,0] = -spacing[0] basis[0,1] = -spacing[1] # And similarly, the offset is just z along the z direction - offset = t.zeros([3], dtype=basis.dtype) + #offset = t.tensor([0,0,z], dtype=basis.dtype) + offset = t.zeros(3, dtype=basis.dtype) offset[2] = z - # And we call the generalized function! propagator = generate_generalized_angular_spectrum_propagator(shape, basis, @@ -500,20 +500,24 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o # Now, we need to generate the out-of-plane direction, so we can # expand these Ks to the full Ks in 3D reciprocal space. - - # THis is broken down into steps to avoid floating point underflow + + # This is broken down into 2 steps to avoid floating point underflow # which was a real problem that showed up for electron ptycho b1_dir = basis[:,0] / t.linalg.norm(basis[:,0]) b2_dir = basis[:,1] / t.linalg.norm(basis[:,1]) perpendicular_dir = t.cross(b1_dir,b2_dir) - perpendicular_dir /= t.linalg.norm(perpendicular_dir) + # Note that we cannot use in-place operations if we want to be able to + # use automatic differentiation successfully + perpendicular_dir = perpendicular_dir / t.linalg.norm(perpendicular_dir) # We set the sign of the propagation direction appropriately if propagation_vector is not None: - perpendicular_dir *= t.sign(t.dot(perpendicular_dir,propagation_vector)) + perpendicular_dir = perpendicular_dir \ + * t.sign(t.dot(perpendicular_dir,propagation_vector)) else: pass - perpendicular_dir *= t.sign(t.dot(perpendicular_dir,offset_vector)) + perpendicular_dir = perpendicular_dir * \ + t.sign(t.dot(perpendicular_dir,offset_vector)) # Then, if we have a propagation vector, we shift the in-plane # components of all the pixels to be centered around the in-plane diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index abe5e9c..f45a9b9 100644 --- a/tests/tools/test_propagators.py +++ b/tests/tools/test_propagators.py @@ -256,7 +256,7 @@ def test_near_field(): # 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)) - # Finally, we check that the bandlimiting at least does something + # Then, we check that the bandlimiting at least does something asp = propagators.generate_angular_spectrum_propagator( E0.shape,(1.5e-9,1e-9),wavelength,z,remove_z_phase=True, dtype=t.complex128, bandlimit=0.3) @@ -266,6 +266,20 @@ def test_near_field(): assert asp[130,0] != 0 assert asp[0,175] != 0 + # Then, we check that automatic differentiation works + z = t.tensor([z],requires_grad=True) + spacing = t.tensor((1.5e-9,1e-9), requires_grad=True) + wavelength = t.tensor([wavelength], requires_grad=True) + + asp = propagators.generate_angular_spectrum_propagator( + E0.shape, spacing, wavelength, z) + + asp[10,10].backward() + assert z.grad != 0 + assert spacing.grad[0] != 0 + assert wavelength.grad !=0 + + def test_generalized_near_field(): # The strategy is to compare the propagation of a gaussian beam to