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

@@ -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

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