From 936b203aac5eb93062bb906a0a999a87e06cba58 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 28 Dec 2020 22:09:20 -0500 Subject: [PATCH] Fix the issue with odd-sized probes & objs in RPI functions, still not fully tested --- CDTools/tools/initializers/initializers.py | 14 ++++++++------ CDTools/tools/interactions/interactions.py | 14 ++++++++------ examples/transmission_RPI.py | 6 ++++-- tests/tools/test_initializers.py | 4 ++++ tests/tools/test_interactions.py | 19 +++++++++++++++++++ 5 files changed, 43 insertions(+), 14 deletions(-) diff --git a/CDTools/tools/initializers/initializers.py b/CDTools/tools/initializers/initializers.py index 15a6927..68487ef 100644 --- a/CDTools/tools/initializers/initializers.py +++ b/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] diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py index 2f5259b..7b7e5ad 100644 --- a/CDTools/tools/interactions/interactions.py +++ b/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.') diff --git a/examples/transmission_RPI.py b/examples/transmission_RPI.py index c28ee26..9c49978 100644 --- a/examples/transmission_RPI.py +++ b/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])): diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index be8e7f1..bb3a85e 100644 --- a/tests/tools/test_initializers.py +++ b/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) diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py index 765dbfc..e4f1177 100644 --- a/tests/tools/test_interactions.py +++ b/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 +