Get all tests passing by updating the loss thresholds for the new normalizations, and add the new normalization to all models

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.)

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)

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')
 

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(

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 = [

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

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

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