From 08a77ce5c626f129dc0e9672d7486027d9992971 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Sun, 22 Dec 2024 15:08:23 +0100 Subject: [PATCH 1/3] Add the object view crop to Bragg2DPtycho for parity, and add a small fix to obj_view_crop which limits the view expansion to the actual object array size --- src/cdtools/models/bragg_2d_ptycho.py | 39 +++++++++++++++++++++------ src/cdtools/models/fancy_ptycho.py | 13 ++++++--- 2 files changed, 41 insertions(+), 11 deletions(-) diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index 473c35f..bfbe7db 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -79,6 +79,7 @@ class Bragg2DPtycho(CDIModel): lens=False, units='um', dtype=t.float32, + obj_view_crop=0, ): # We need the detector geometry @@ -147,6 +148,19 @@ class Bragg2DPtycho(CDIModel): self.probe = t.nn.Parameter(probe_guess / self.probe_norm) self.obj = t.nn.Parameter(obj_guess) + # NOTE: I think it makes sense to protect against obj_view_crop + # being zero or below, because there is nothing else to show outside + # the object array. No reason to throw an error if, e.g., the user + # asks for a big padding which goes outside of the actual object array. + # Just show the full array. + print(obj_view_crop) + if obj_view_crop > 0: + self.obj_view_slice = np.s_[obj_view_crop:-obj_view_crop, + obj_view_crop:-obj_view_crop] + else: + self.obj_view_slice = np.s_[:,:] + + if probe_support is None: probe_support = t.ones_like(self.probe[0], dtype=t.bool) self.register_buffer('probe_support', @@ -241,6 +255,7 @@ class Bragg2DPtycho(CDIModel): correct_tilt=True, lens=False, obj_padding=200, + obj_view_crop=None, units='um', ): wavelength = dataset.wavelength @@ -353,9 +368,16 @@ class Bragg2DPtycho(CDIModel): probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)] probe = t.stack([probe,] + probe_stack) - obj = t.exp(1j*(randomize_ang * (t.rand(obj_size)-0.5))) - + + pfc = (probe_fourier_crop if probe_fourier_crop else 0) + if obj_view_crop is None: + obj_view_crop = min(probe.shape[-2], probe.shape[-1]) // 2 + pfc + if obj_view_crop < 0: + obj_view_crop += min(probe.shape[-2], probe.shape[-1]) // 2 + pfc + + obj_view_crop += obj_padding + det_geo = dataset.detector_geometry translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) @@ -401,6 +423,7 @@ class Bragg2DPtycho(CDIModel): propagate_probe=propagate_probe, correct_tilt=correct_tilt, lens=lens, + obj_view_crop=obj_view_crop, units=units, ) @@ -566,21 +589,21 @@ class Bragg2DPtycho(CDIModel): )), ('Object Amplitude, Surface Normal View', lambda self, fig: p.plot_amplitude( - self.obj, + self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units, )), ('Object Phase, Surface Normal View', lambda self, fig: p.plot_phase( - self.obj, + self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, units=self.units, )), ('Object Amplitude, Beam View', lambda self, fig: p.plot_amplitude( - self.obj, + self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, view_basis=beam_basis, @@ -588,7 +611,7 @@ class Bragg2DPtycho(CDIModel): )), ('Object Phase, Beam View', lambda self, fig: p.plot_phase( - self.obj, + self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, view_basis=beam_basis, @@ -596,7 +619,7 @@ class Bragg2DPtycho(CDIModel): )), ('Object Amplitude, Detector View', lambda self, fig: p.plot_amplitude( - self.obj, + self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, view_basis=self.det_basis, @@ -604,7 +627,7 @@ class Bragg2DPtycho(CDIModel): )), ('Object Phase, Detector View', lambda self, fig: p.plot_phase( - self.obj, + self.obj[self.obj_view_slice], fig=fig, basis=self.obj_basis, view_basis=self.det_basis, diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index ce5c9c4..7e3ce89 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -103,9 +103,16 @@ class FancyPtycho(CDIModel): self.probe = t.nn.Parameter(probe_guess / self.probe_norm) self.obj = t.nn.Parameter(obj_guess) - - self.obj_view_slice = np.s_[obj_view_crop:-obj_view_crop, - obj_view_crop:-obj_view_crop] + # NOTE: I think it makes sense to protect against obj_view_crop + # being zero or below, because there is nothing else to show outside + # the object array. No reason to throw an error if, e.g., the user + # asks for a big padding which goes outside of the actual object array. + # Just show the full array. + if obj_view_crop > 0: + self.obj_view_slice = np.s_[obj_view_crop:-obj_view_crop, + obj_view_crop:-obj_view_crop] + else: + self.obj_view_slice = np.s_[:,:] # TODO: perhaps not working anymore for fourier cropped probes if background is None: From a22ae1b4768207abfba49ff3997a08138aaa3037 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Sun, 22 Dec 2024 16:59:30 +0100 Subject: [PATCH 2/3] Fix the issue with oversampling in bragg 2d ptycho when correct_tilt=True, and fix one other issue with object sizing in bragg 2d ptycho with oversampling --- src/cdtools/models/bragg_2d_ptycho.py | 23 +++++++++++++---------- src/cdtools/models/fancy_ptycho.py | 3 ++- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index bfbe7db..fc17aed 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -153,7 +153,7 @@ class Bragg2DPtycho(CDIModel): # the object array. No reason to throw an error if, e.g., the user # asks for a big padding which goes outside of the actual object array. # Just show the full array. - print(obj_view_crop) + if obj_view_crop > 0: self.obj_view_slice = np.s_[obj_view_crop:-obj_view_crop, obj_view_crop:-obj_view_crop] @@ -169,8 +169,8 @@ class Bragg2DPtycho(CDIModel): if background is None: raise NotImplementedError('Issues with this due to probe fourier padding') - background = 1e-6 * t.ones(self.probe[0].shape, - dtype=t.float32) + shape = [s//oversampling for s in self.probe[0]] + background = 1e-6 * t.ones(shape, dtype=t.float32) self.background = t.nn.Parameter(background) @@ -204,10 +204,11 @@ class Bragg2DPtycho(CDIModel): tools.propagators.generate_high_NA_k_intensity_map( self.obj_basis, self.get_detector_geometry()['basis'] / oversampling, - self.background.shape, + [oversampling * d for d in self.background.shape], self.get_detector_geometry()['distance'], self.wavelength,dtype=t.float32, lens=lens) + self.register_buffer('k_map', t.as_tensor(k_map, dtype=dtype)) self.register_buffer('intensity_map', @@ -326,7 +327,7 @@ class Bragg2DPtycho(CDIModel): obj_size, min_translation = tools.initializers.calc_object_setup( - det_shape, + [s * oversampling for s in det_shape], pix_translations, padding=obj_padding, ) @@ -496,11 +497,13 @@ class Bragg2DPtycho(CDIModel): def measurement(self, wavefields): - return tools.measurements.quadratic_background(wavefields, - self.background, - measurement=tools.measurements.incoherent_sum, - saturation=self.saturation, - oversampling=self.oversampling) + return tools.measurements.quadratic_background( + wavefields, + self.background, + measurement=tools.measurements.incoherent_sum, + saturation=self.saturation, + oversampling=self.oversampling, + ) def loss(self, sim_data, real_data, mask=None): diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 7e3ce89..46e7db6 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -505,6 +505,7 @@ class FancyPtycho(CDIModel): shift_probe=True, multiple_modes=True, probe_support=self.probe_support) + return exit_waves @@ -522,7 +523,7 @@ class FancyPtycho(CDIModel): self.background, measurement=tools.measurements.incoherent_sum, saturation=self.saturation, - oversampling=int(self.oversampling), + oversampling=self.oversampling, simulate_finite_pixels=self.simulate_finite_pixels, ) From c03d4284f2ebc39e345acbb86d111fcd1bf79ea3 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 23 Dec 2024 10:10:40 +0100 Subject: [PATCH 3/3] Fix the issue with obj view padding in Bragg2DPtycho not working with a probe fourier crop --- src/cdtools/models/bragg_2d_ptycho.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index fc17aed..bc10ab4 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -371,11 +371,17 @@ class Bragg2DPtycho(CDIModel): obj = t.exp(1j*(randomize_ang * (t.rand(obj_size)-0.5))) - pfc = (probe_fourier_crop if probe_fourier_crop else 0) + pfc = (probe_fourier_crop if probe_fourier_crop else [0,0]) if obj_view_crop is None: - obj_view_crop = min(probe.shape[-2], probe.shape[-1]) // 2 + pfc + obj_view_crop = min( + probe.shape[-2] // 2 + pfc[0], + probe.shape[-1] // 2 + pfc[1] + ) if obj_view_crop < 0: - obj_view_crop += min(probe.shape[-2], probe.shape[-1]) // 2 + pfc + obj_view_crop += min( + probe.shape[-2] // 2 + pfc[0], + probe.shape[-1] // 2 + pfc[1] + ) obj_view_crop += obj_padding