Switch to a more sensible normalization strategy for loss metrics

docs/source/tutorial.rst

@@ -196,6 +196,7 @@ Once again, we start with the basic skeleton
 
 .. code-block:: python
 
+   from functools import partial
    import torch as t
    from cdtools.models import CDIModel
    from cdtools import tools
@@ -251,7 +252,11 @@ There is no requirement for what the arguments to the initialization function of
        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 = partial(tools.losses.amplitude_mse, use_sum=True)
+       self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
+
+
 The first thing to notice about this model is that all the fixed, geometric information is stored with the :code:`module.register_buffer()` function. This is what makes it possible to move all the relevant tensors between devices using a single call to :code:`module.to()`, for example. It stores thetensor as an object attribute, but it also registers it so that pytorch is aware that this attribute helps to encode the state of the model.
 
 The supporting information we need is the wavelength of the illumination, the basis of the probe array in real space, and an offset to define the zero point of the translation. 
@@ -268,6 +273,8 @@ The Adam optimizer is designed so that the learning rate sets the maximum stepsi
 
 This is important to remember when adding additional error models. Rescaling all the parameters to have a typical amplitude near 1 is the best way to get well-behaved reconstructions.
 
+The final two lines assign a loss function and its associated normalizer. Here we use :code:`amplitude_mse` and its paired :code:`AmplitudeMSENormalizer`. The normalization function :code:`amplitude_mse` computes the the mean squared error between the square roots of the simulated and measured intensities. In this case, we call it with the :code:`use_sum=True` flag, which will actually calculate a sum-square error, which will be normalized afterward to a mean-squared error by the normalizer, :code:`AmplitudeMSENormalizer`. This pattern is used to ensure that the losses from minibatches with different sizes are properly weighted.
+
 
 Initialization from Dataset
 +++++++++++++++++++++++++++
@@ -347,7 +354,7 @@ Here, we take input in the form of an index and a translation. Note that this in
 
 We start by mapping the translation, given in real space, into pixel coordinates. Then, we use an "off-the-shelf" interaction model - :code:`ptycho_2d_round`, which models a standard 2D ptychography interaction, but rounds the translations to the nearest whole pixel (does not attempt subpixel translations).
 
-The next three definitions amount to just choosing an off-the-shelf function to simulate each step in the chain.
+The next two definitions amount to just choosing an off-the-shelf function to simulate each step in the chain.
 
 .. code-block:: python
 
@@ -357,11 +364,8 @@ The next three definitions amount to just choosing an off-the-shelf function to
     def measurement(self, wavefields):
         return tools.measurements.intensity(wavefields)
 
-    def loss(self, sim_data, real_data):
-        return tools.losses.amplitude_mse(real_data, sim_data)
-
 
-The forward propagator maps the exit wave to the wave at the surface of the detector, here using a far-field propagator. The measurement maps that exit wave to a measured pixel value, and the loss defines a loss function to attempt to minimize. The loss function we've chosen - the amplitude mean squared error - is the most reliable one, and can also easily be overridden by an end user.
+The forward propagator maps the exit wave to the wave at the surface of the detector, here using a far-field propagator. The measurement maps that wavefield to a measured pixel value. The loss function was already assigned in :code:`__init__` as described above.
 
 
 Plotting

examples/near_field_ptycho.py

@@ -26,7 +26,8 @@
     near_field=True,
     propagation_distance=3.65e-3, # 3.65 downstream from focus
     units='um', # Set the units for the live plots
-    obj_view_crop=-35,
+    obj_view_crop=-35, # Expand the view for the live plots
+    loss="poisson_nll", # Best option for photon-counting detectors
     panel_plot_mode=True, # Set to False to get individual figures
 )
 

examples/tutorial_simple_ptycho.py

@@ -1,3 +1,4 @@
+from functools import partial
 import torch as t
 from cdtools.models import CDIModel
 from cdtools import tools
@@ -41,6 +42,10 @@ def __init__(
         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 = partial(tools.losses.amplitude_mse, use_sum=True)
+        self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
+
 
     @classmethod
     def from_dataset(cls, dataset):
@@ -102,9 +107,6 @@ def forward_propagator(self, 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 = [

src/cdtools/models/bragg_2d_ptycho.py

@@ -1,3 +1,4 @@
+from functools import partial
 import torch as t
 from cdtools.models import CDIModel
 from cdtools.datasets import Ptycho2DDataset
@@ -76,6 +77,7 @@ def __init__(
             propagate_probe=True,
             correct_tilt=True,
             lens=False,
+            loss='amplitude mse',
             units='um',
             dtype=t.float32,
             obj_view_crop=0,
@@ -237,7 +239,22 @@ def __init__(
         # TODO: probably doesn't support non-float-32 dtypes
         self.register_buffer('universal_propagator',
                              universal_propagator)
-
+
+        # Here we set the appropriate loss function
+        if (loss.lower().strip() == 'amplitude mse'
+                or loss.lower().strip() == 'amplitude_mse'):
+            self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
+            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()
+        elif (loss.lower().strip() == 'intensity mse'
+                or loss.lower().strip() == 'intensity_mse'):
+            self.loss = partial(tools.losses.intensity_mse, use_sum=True)
+            self.loss_normalizer = tools.losses.IntensityMSENormalizer()
+        else:
+            raise KeyError('Specified loss function not supported')
 
 
     @classmethod
@@ -257,6 +274,7 @@ def from_dataset(
             propagate_probe=True,
             correct_tilt=True,
             lens=False,
+            loss='amplitude mse',
             obj_padding=200,
             obj_view_crop=None,
             units='um',
@@ -450,6 +468,7 @@ def from_dataset(
                    propagate_probe=propagate_probe,
                    correct_tilt=correct_tilt,
                    lens=lens,
+                   loss=loss,
                    obj_view_crop=obj_view_crop,
                    units=units,
                    panel_plot_mode=panel_plot_mode,
@@ -534,10 +553,6 @@ def measurement(self, wavefields):
         )
 
 
-    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/fancy_ptycho.py

@@ -1,3 +1,4 @@
+from functools import partial
 import torch as t
 from cdtools.models import CDIModel
 from cdtools.datasets import Ptycho2DDataset
@@ -219,10 +220,16 @@ def __init__(self,
         # 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 = partial(tools.losses.amplitude_mse, use_sum=True)
+            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()
+        elif (loss.lower().strip() == 'intensity mse'
+                or loss.lower().strip() == 'intensity_mse'):
+            self.loss = partial(tools.losses.intensity_mse, use_sum=True)
+            self.loss_normalizer = tools.losses.IntensityMSENormalizer()
         else:
             raise KeyError('Specified loss function not supported')
 

src/cdtools/models/multislice_2d_ptycho.py

@@ -1,3 +1,4 @@
+from functools import partial
 import torch as t
 from cdtools.models import CDIModel
 from cdtools.datasets import Ptycho2DDataset
@@ -46,6 +47,7 @@ def __init__(self,
                  fourier_probe=False,
                  prevent_aliasing=True,
                  phase_only=False,
+                 loss='amplitude mse',
                  units='um',
                  panel_plot_mode=False,
                  plot_level=1,
@@ -155,9 +157,25 @@ def __init__(self,
 
         self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, self.bandlimit)
 
+        # Here we set the appropriate loss function
+        if (loss.lower().strip() == 'amplitude mse'
+                or loss.lower().strip() == 'amplitude_mse'):
+            self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
+            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()
+        elif (loss.lower().strip() == 'intensity mse'
+                or loss.lower().strip() == 'intensity_mse'):
+            self.loss = partial(tools.losses.intensity_mse, use_sum=True)
+            self.loss_normalizer = tools.losses.IntensityMSENormalizer()
+        else:
+            raise KeyError('Specified loss function not supported')
+
 
     @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):
+    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, loss='amplitude_mse'):
 
         wavelength = dataset.wavelength
         det_basis = dataset.detector_geometry['basis']
@@ -301,7 +319,9 @@ def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n
                    phase_only=phase_only,
                    prevent_aliasing=prevent_aliasing,
                    panel_plot_mode=panel_plot_mode,
-                   plot_level=plot_level)
+                   plot_level=plot_level,
+                   loss=loss,
+        )
 
 
     def interaction(self, index, translations):
@@ -415,11 +435,6 @@ def measurement(self, wavefields):
                             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

@@ -1,3 +1,4 @@
+from functools import partial
 import torch as t
 from cdtools.models import CDIModel
 from cdtools.datasets import Ptycho2DDataset
@@ -164,14 +165,20 @@ def __init__(self,
 
         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 = partial(tools.losses.amplitude_mse, use_sum=True)
+            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()
+        elif (loss.lower().strip() == 'intensity mse'
+                or loss.lower().strip() == 'intensity_mse'):
+            self.loss = partial(tools.losses.intensity_mse, use_sum=True)
+            self.loss_normalizer = tools.losses.IntensityMSENormalizer()
         else:
             raise KeyError('Specified loss function not supported')
 

src/cdtools/models/rpi.py

@@ -1,3 +1,4 @@
+from functools import partial
 import torch as t
 from cdtools.models import CDIModel
 from cdtools import tools
@@ -54,6 +55,7 @@ def __init__(
             exponentiate_obj=False,
             phase_only=False,
             propagation_distance=0,
+            loss='amplitude mse',
             units='um',
             dtype=t.float32,
             panel_plot_mode=True,
@@ -146,6 +148,22 @@ def __init__(
         self.register_buffer('prop_dir',
                              t.as_tensor([0, 0, 1], dtype=dtype))
 
+        # Here we set the appropriate loss function
+        if (loss.lower().strip() == 'amplitude mse'
+                or loss.lower().strip() == 'amplitude_mse'):
+            self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
+            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()
+        elif (loss.lower().strip() == 'intensity mse'
+                or loss.lower().strip() == 'intensity_mse'):
+            self.loss = partial(tools.losses.intensity_mse, use_sum=True)
+            self.loss_normalizer = tools.losses.IntensityMSENormalizer()
+        else:
+            raise KeyError('Specified loss function not supported')
+
 
     @classmethod
     def from_dataset(
@@ -164,6 +182,7 @@ def from_dataset(
             exponentiate_obj=False,
             phase_only=False,
             probe_threshold=0,
+            loss='amplitude mse',
             dtype=t.float32,
             panel_plot_mode=True,
             plot_level=1,
@@ -259,7 +278,9 @@ def from_dataset(
                          phase_only=phase_only,
                          weight_matrix=weight_matrix,
                          panel_plot_mode=panel_plot_mode,
-                         plot_level=plot_level)
+                         plot_level=plot_level,
+                         loss=loss,
+        )
 
         # I don't love this pattern, where I do the "real" obj initialization
         # after creating the rpi object. But, I chose this so that I could
@@ -291,6 +312,7 @@ def from_calibration(
             dtype=t.float32,
             panel_plot_mode=True,
             plot_level=1,
+            loss='amplitude_mse',
     ):
 
         complex_dtype = (t.ones([1], dtype=dtype) +
@@ -336,6 +358,7 @@ def from_calibration(
             phase_only=phase_only,
             panel_plot_mode=panel_plot_mode,
             plot_level=plot_level,
+            loss=loss,
         )
 
         rpi_object.init_obj(initialization)

src/cdtools/models/simple_ptycho.py

@@ -1,3 +1,4 @@
+from functools import partial
 import torch as t
 from cdtools.models import CDIModel
 from cdtools import tools
@@ -41,6 +42,10 @@ def __init__(
         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 = partial(tools.losses.amplitude_mse, use_sum=True)
+        self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
+
 
     @classmethod
     def from_dataset(cls, dataset):
@@ -99,12 +104,10 @@ def interaction(self, index, translations):
     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 = [