Fix the issue with odd-sized probes & objs in RPI functions, still not fully tested

CDTools/tools/initializers/initializers.py

@@ -380,12 +380,14 @@ def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, backgroun
     if probe.dim() == 4:
         probe = orthogonalize_probes(probe)[0]
     
-    pad0 = (probe.shape[-3] - obj_shape[0])//2
-    pad1 = (probe.shape[-2] - obj_shape[1])//2
-
+    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
+    
     def a_dagger(im):
         im = cmath.complex_to_torch(im.reshape(obj_shape)).to(dtype=t.float32)
-        im = inverse_far_field(pad(far_field(im), (0,0,pad1,pad1,pad0,pad0)))
+        im = inverse_far_field(pad(far_field(im), (0,0,pad1l,pad1r,pad0l,pad0r)))
         exit_wave = cmath.cmult(probe,im)
         farfield = cmath.torch_to_complex(far_field(exit_wave))
         return farfield.ravel()
@@ -395,8 +397,8 @@ def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, backgroun
         im = inverse_far_field(measured)
         multiplied = cmath.cmult(cmath.cconj(probe), im)
         backplane = far_field(multiplied)
-        clipped = backplane[pad0:pad0+obj_shape[0],
-                            pad1:pad1+obj_shape[1],:]
+        clipped = backplane[pad0l:pad0l+obj_shape[0],
+                            pad1l:pad1l+obj_shape[1],:]
         return cmath.torch_to_complex(inverse_far_field(clipped)).ravel()
 
     patsize = pattern.shape[0]*pattern.shape[1]

CDTools/tools/interactions/interactions.py

@@ -560,14 +560,16 @@ 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
-    pad0 = (probe.shape[-3] - obj.shape[-3])//2
-    pad1 = (probe.shape[-2] - obj.shape[-2])//2
-    
+    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
+        
     if obj.dim() == 3:
-        fftobj = t.nn.functional.pad(fftobj, (0, 0, pad1, pad1, pad0, pad0))
+        fftobj = t.nn.functional.pad(fftobj, (0, 0, pad1l, pad1r, pad0l, pad0r))
     elif obj.dim() == 4:
-        fftobj = t.nn.functional.pad(fftobj,
-                                     (0, 0, pad1, pad1, pad0, pad0, 0, 0))
+        fftobj = t.nn.functional.pad(
+            fftobj, (0, 0, pad1l, pad1r, pad0l, pad0r, 0, 0))
     else:
         raise NotImplementedError('RPI interaction with obj of dimension higher than 4 (including complex dimension) is not supported.')
         

examples/transmission_RPI.py

@@ -18,10 +18,12 @@ background = ptycho_results['background']
 
 dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(ss_filename)
 
+
 # Next, we create a ptychography model from the dataset
 # Note that we explicitly as for two incoherent probe modes
 model = CDTools.models.RPI.from_dataset(dataset, probe, [900,900],
-                                        background=background, n_modes=2)
+                                        background=background, n_modes=2,
+                                        initialization='random')
 
 
 # Let's do this reconstruction on the GPU, shall we? 
@@ -35,7 +37,7 @@ for i, loss in enumerate(model.LBFGS_optimize(30, dataset, lr=0.4, regularizatio
     #model.inspect(dataset)
     print(i,loss)
     
-model.inspect(dataset)
+#model.inspect(dataset)
     
 # Now we use the regularizer to damp all but the top modes
 for i, loss in enumerate(model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1])):

tests/tools/test_initializers.py

@@ -191,3 +191,7 @@ def test_SHARP_style_probe(ptycho_cxi_1):
     assert probe.shape == t.Size([256,256,2])
 
 
+def test_RPI_spectral_init():
+    # Figure out a good way to test this
+    assert 0
+    RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, background=None)

tests/tools/test_interactions.py

@@ -225,5 +225,24 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
     assert np.max(np.abs(exit_wave_np-exit_wave_torch)) < 0.005
 
 
+def test_RPI_interaction(random_probe, random_obj):
+
+    random_obj1 = cmath.complex_to_torch(random_obj[:79,:68])
+    random_probe1 = cmath.complex_to_torch(random_probe)
+    output = interactions.RPI_interaction(random_probe1, random_obj1)
+
+
+    random_obj1 = cmath.complex_to_torch(random_obj[:256,:256])
+    random_probe1 = cmath.complex_to_torch(random_probe)
+    output = interactions.RPI_interaction(random_probe1, random_obj1)
 
     
+    random_obj1 = cmath.complex_to_torch(random_obj[:42,:103])
+    random_probe1 = cmath.complex_to_torch(random_probe)
+    output = interactions.RPI_interaction(random_probe1, random_obj1)
+
+    # Need to actually test against a numpy implementation
+    print(random_probe.shape)
+    print(random_obj.shape)
+    assert 0
+