diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index 881759f..157d704 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -237,7 +237,10 @@ class Bragg2DPtycho(CDIModel): # TODO: probably doesn't support non-float-32 dtypes self.register_buffer('universal_propagator', universal_propagator) - + + # We register a loss function and an appropriate normalization + self.loss = tools.losses.amplitude_mse + self.loss_normalizer = tools.losses.AmplitudeMSENormalizer() @classmethod @@ -534,10 +537,6 @@ class Bragg2DPtycho(CDIModel): ) - def loss(self, sim_data, real_data, mask=None): - return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) - - def sim_to_dataset(self, args_list): # In the future, potentially add more control # over what metadata is saved (names, etc.) diff --git a/src/cdtools/models/multislice_2d_ptycho.py b/src/cdtools/models/multislice_2d_ptycho.py index d6663be..9ff1e61 100644 --- a/src/cdtools/models/multislice_2d_ptycho.py +++ b/src/cdtools/models/multislice_2d_ptycho.py @@ -155,6 +155,10 @@ class Multislice2DPtycho(CDIModel): self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, self.bandlimit) + # We register a loss function and an appropriate normalization + self.loss = tools.losses.amplitude_mse + self.loss_normalizer = tools.losses.AmplitudeMSENormalizer() + @classmethod def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True, probe_support_radius=None, panel_plot_mode=False, plot_level=1): @@ -415,11 +419,6 @@ class Multislice2DPtycho(CDIModel): oversampling=self.oversampling) - def loss(self, sim_data, real_data, mask=None): - return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) - #return tools.losses.poisson_nll(real_data, sim_data, mask=mask,eps=0.5) - - def to(self, *args, **kwargs): super(Multislice2DPtycho, self).to(*args, **kwargs) self.wavelength = self.wavelength.to(*args,**kwargs) diff --git a/src/cdtools/models/multislice_ptycho.py b/src/cdtools/models/multislice_ptycho.py index 3e79b22..9940a41 100644 --- a/src/cdtools/models/multislice_ptycho.py +++ b/src/cdtools/models/multislice_ptycho.py @@ -164,14 +164,16 @@ class MultislicePtycho(CDIModel): self.register_buffer('simulate_finite_pixels', t.as_tensor(simulate_finite_pixels, dtype=bool)) - + # Here we set the appropriate loss function if (loss.lower().strip() == 'amplitude mse' or loss.lower().strip() == 'amplitude_mse'): self.loss = tools.losses.amplitude_mse + self.loss_normalizer = tools.losses.AmplitudeMSENormalizer() elif (loss.lower().strip() == 'poisson nll' or loss.lower().strip() == 'poisson_nll'): self.loss = tools.losses.poisson_nll + self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer() else: raise KeyError('Specified loss function not supported') diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index 3d8ac6c..2e642e0 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -146,6 +146,10 @@ class RPI(CDIModel): self.register_buffer('prop_dir', t.as_tensor([0, 0, 1], dtype=dtype)) + # Define the loss + self.loss = tools.losses.amplitude_mse + self.loss_normalizer = tools.losses.AmplitudeMSENormalizer() + @classmethod def from_dataset( diff --git a/src/cdtools/models/simple_ptycho.py b/src/cdtools/models/simple_ptycho.py index 734ba20..57c88e5 100644 --- a/src/cdtools/models/simple_ptycho.py +++ b/src/cdtools/models/simple_ptycho.py @@ -41,6 +41,10 @@ class SimplePtycho(CDIModel): self.probe = t.nn.Parameter(probe_guess / self.probe_norm) self.obj = t.nn.Parameter(obj_guess) + # We register a loss function and an appropriate normalization + self.loss = tools.losses.amplitude_mse + self.loss_normalizer = tools.losses.AmplitudeMSENormalizer() + @classmethod def from_dataset(cls, dataset): @@ -99,12 +103,10 @@ class SimplePtycho(CDIModel): def forward_propagator(self, wavefields): return tools.propagators.far_field(wavefields) + def measurement(self, wavefields): return tools.measurements.intensity(wavefields) - def loss(self, real_data, sim_data): - return tools.losses.amplitude_mse(real_data, sim_data) - # This lists all the plots to display on a call to model.inspect() plot_list = [ diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py index 79a5caf..19953d7 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -124,7 +124,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): plt.close('all') # If this fails, the reconstruction has gotten worse - assert model.loss_history[-1] < 0.0013 + assert model.loss_history[-1] < 0.38 @pytest.mark.slow @@ -165,4 +165,4 @@ def test_near_field_ptycho(near_field_ptycho_cxi, reconstruction_device, show_pl plt.close('all') # If this fails, the reconstruction has gotten worse - assert model.loss_history[-1] < 0.005 + assert model.loss_history[-1] < 3.9 diff --git a/tests/models/test_simple_ptycho.py b/tests/models/test_simple_ptycho.py index 225e463..bcc5e46 100644 --- a/tests/models/test_simple_ptycho.py +++ b/tests/models/test_simple_ptycho.py @@ -30,4 +30,4 @@ def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): plt.close('all') # If this fails, the reconstruction got worse - assert model.loss_history[-1] < 0.013 + assert model.loss_history[-1] < 6.5 diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py index 3f9245b..fc394d6 100644 --- a/tests/test_reconstructors.py +++ b/tests/test_reconstructors.py @@ -122,7 +122,7 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): # comes from running a reconstruction when it was working well and # choosing a rough value. If it triggers this assertion error, something # changed to make the final quality worse! - assert model_recon.loss_history[-1] < 0.0001 + assert model_recon.loss_history[-1] < 0.09 @pytest.mark.slow @@ -217,7 +217,7 @@ def test_LBFGS_RPI(optical_data_ss_cxi, # The final loss when testing this was 2.28607e-3. Based on this, we set # a threshold of 2.3e-3 for the tested loss. If this value has been # exceeded, the reconstructions have gotten worse. - assert model_recon.loss_history[-1] < 0.0023 + assert model_recon.loss_history[-1] < 0.14 @pytest.mark.slow @@ -330,4 +330,4 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): # The final loss when testing this was 7.12188e-4. Based on this, we set # a threshold of 7.2e-4 for the tested loss. If this value has been # exceeded, the reconstructions have gotten worse. - assert model.loss_history[-1] < 0.00072 + assert model.loss_history[-1] < 0.65