Update the amplitude MSE and intensity MSE loss to act equivalently - currently, the intensity MSE loss was by default a mean, but amplitude was by default a sum. Both functions now have a flag, use_sum=False, which is False by default. This changes the default behavior of amplitude_mse. All models, including the tutorial version of simple_ptycho, are updated accordingly and test coverage was added.

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
@@ -252,7 +253,7 @@ There is no requirement for what the arguments to the initialization function of
        self.obj = t.nn.Parameter(obj_guess)
 
        # We register a loss function and an appropriate normalization
-       self.loss = tools.losses.amplitude_mse
+       self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
        self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
 
 
@@ -272,7 +273,7 @@ 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. The loss function is stored as an instance attribute rather than defined as a method, which allows it to be swapped out at construction time. The normalizer is a stateful object that accumulates statistics over the first epoch and uses them to convert the raw summed loss into a normalized mean value. Here we use :code:`amplitude_mse` and its paired :code:`AmplitudeMSENormalizer`, which computes the mean squared error between the square roots of the simulated and measured intensities.
+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

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
@@ -42,7 +43,7 @@ class SimplePtycho(CDIModel):
         self.obj = t.nn.Parameter(obj_guess)
 
         # We register a loss function and an appropriate normalization
-        self.loss = tools.losses.amplitude_mse
+        self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
         self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
 
 

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
@@ -242,7 +243,7 @@ class Bragg2DPtycho(CDIModel):
         # 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'):
@@ -250,7 +251,7 @@ class Bragg2DPtycho(CDIModel):
             self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer()
         elif (loss.lower().strip() == 'intensity mse'
                 or loss.lower().strip() == 'intensity_mse'):
-            self.loss = tools.losses.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/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,7 +220,7 @@ class FancyPtycho(CDIModel):
         # 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'):
@@ -227,7 +228,7 @@ class FancyPtycho(CDIModel):
             self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer()
         elif (loss.lower().strip() == 'intensity mse'
                 or loss.lower().strip() == 'intensity_mse'):
-            self.loss = tools.losses.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
@@ -159,7 +160,7 @@ class Multislice2DPtycho(CDIModel):
         # 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'):
@@ -167,7 +168,7 @@ class Multislice2DPtycho(CDIModel):
             self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer()
         elif (loss.lower().strip() == 'intensity mse'
                 or loss.lower().strip() == 'intensity_mse'):
-            self.loss = tools.losses.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_ptycho.py

@@ -1,3 +1,4 @@
+from functools import partial
 import torch as t
 from cdtools.models import CDIModel
 from cdtools.datasets import Ptycho2DDataset
@@ -168,7 +169,7 @@ class MultislicePtycho(CDIModel):
         # 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'):
@@ -176,7 +177,7 @@ class MultislicePtycho(CDIModel):
             self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer()
         elif (loss.lower().strip() == 'intensity mse'
                 or loss.lower().strip() == 'intensity_mse'):
-            self.loss = tools.losses.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
@@ -150,7 +151,7 @@ class RPI(CDIModel):
         # 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'):
@@ -158,7 +159,7 @@ class RPI(CDIModel):
             self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer()
         elif (loss.lower().strip() == 'intensity mse'
                 or loss.lower().strip() == 'intensity_mse'):
-            self.loss = tools.losses.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/simple_ptycho.py

@@ -1,3 +1,4 @@
+from functools import partial
 import torch as t
 from cdtools.models import CDIModel
 from cdtools import tools
@@ -42,7 +43,7 @@ class SimplePtycho(CDIModel):
         self.obj = t.nn.Parameter(obj_guess)
 
         # We register a loss function and an appropriate normalization
-        self.loss = tools.losses.amplitude_mse
+        self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
         self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
 
 

src/cdtools/tools/losses/losses.py

@@ -18,22 +18,16 @@ __all__ = [
 ]
 
 
-def amplitude_mse(intensities, sim_intensities, mask=None):
+def amplitude_mse(intensities, sim_intensities, mask=None, use_sum=False):
     """ Returns the mean squared error of a simulated dataset's amplitudes
 
     Calculates the mean squared error between a given set of 
     measured diffraction intensities and a simulated set.
 
-
     This function calculates the mean squared error between their
     associated amplitudes. Because this is not well defined for negative
     numbers, make sure that all the intensities are >0 before using this
     loss.
-    
-    Note that this is actually, by defauly, a sum-squared error. In this
-    case, it is intended to be used with the loss normalization strategy
-    in the base CDIModel class, which works well if the minibatch size
-    is not fixed.
 
     It can accept intensity and simulated intensity tensors of any shape
     as long as their shapes match, and the provided mask array can be
@@ -41,6 +35,12 @@ def amplitude_mse(intensities, sim_intensities, mask=None):
 
     This is empirically the most useful loss function for most cases where
     a photon counting detector cannot be used.
+    
+    Note that, when used with the AmplitudeMSENormalizer, this function
+    should be called with use_sum=True, in order to return the sum-squared
+    error rather than the mean-squared error. This allows for the
+    AmplitudeMSENormalizer to properly weight the loss arising from minibatches
+    which may not have equal length.
 
     Parameters
     ----------
@@ -51,7 +51,7 @@ def amplitude_mse(intensities, sim_intensities, mask=None):
     mask : torch.Tensor
         A mask with ones for pixels to include and zeros for pixels to exclude
     use_sum : bool
-        Default is True. If set to True, actually performs the sum squared error
+        Default is False. If set to True, actually performs the sum squared error
 
     Returns
     -------
@@ -64,13 +64,20 @@ def amplitude_mse(intensities, sim_intensities, mask=None):
     # with all the errors working off of the same inputs
 
     if mask is None:
-        return t.sum((t.sqrt(sim_intensities) -
-                      t.sqrt(intensities))**2)
+        if use_sum:
+            return t.sum((t.sqrt(sim_intensities) -
+                          t.sqrt(intensities))**2)
+        else:
+            return t.mean((t.sqrt(sim_intensities) -
+                          t.sqrt(intensities))**2)
     else:
         masked_intensities = intensities.masked_select(mask)
-        return t.sum((t.sqrt(sim_intensities.masked_select(mask)) -
-                      t.sqrt(masked_intensities))**2)
-
+        if use_sum:
+            return t.sum((t.sqrt(sim_intensities.masked_select(mask)) -
+                          t.sqrt(masked_intensities))**2)
+        else:
+            return t.mean((t.sqrt(sim_intensities.masked_select(mask)) -
+                          t.sqrt(masked_intensities))**2)
 
 class AmplitudeMSENormalizer(object):
     """ Normalizer for the amplitude MSE loss, used with recon.optimize
@@ -115,7 +122,7 @@ class AmplitudeMSENormalizer(object):
         return loss / self.num_pix
                       
     
-def intensity_mse(intensities, sim_intensities, mask=None):
+def intensity_mse(intensities, sim_intensities, mask=None, use_sum=False):
     """ Returns the mean squared error of a simulated dataset's intensities
 
     Calculates the summed mean squared error between a given set of 
@@ -129,6 +136,12 @@ def intensity_mse(intensities, sim_intensities, mask=None):
     
     This is rarely a good loss function for ptychography, but can occasionally
     be useful.
+    
+    Note that, when used with the IntensityMSENormalizer, this function
+    should be called with use_sum=True, in order to return the sum-squared
+    error rather than the mean-squared error. This allows for the
+    IntensityMSENormalizer to properly weight the loss arising from minibatches
+    which may not have equal length.
 
     Parameters
     ----------
@@ -138,6 +151,8 @@ def intensity_mse(intensities, sim_intensities, mask=None):
         A tensor of simulated detector intensities
     mask : torch.Tensor
         A mask with ones for pixels to include and zeros for pixels to exclude
+    use_sum : bool
+        Default is False. If set to True, actually performs the sum squared error
 
     Returns
     -------
@@ -146,13 +161,18 @@ def intensity_mse(intensities, sim_intensities, mask=None):
 
     """
     if mask is None:
-        return t.sum((sim_intensities - intensities)**2) \
-            / intensities.view(-1).shape[0]
+        if use_sum:
+            return t.sum((sim_intensities - intensities)**2)
+        else:
+            return t.mean((sim_intensities - intensities)**2)
     else:
-        masked_intensities = intensities.masked_select(mask)
-        return t.sum((sim_intensities.masked_select(mask) -
-                      masked_intensities)**2) \
-                      / masked_intensities.shape[0]
+        if use_sum:
+            return t.sum((sim_intensities.masked_select(mask) -
+                          intensities.masked_select(mask))**2)
+        else:
+            return t.mean((sim_intensities.masked_select(mask) -
+                          intensities.masked_select(mask))**2)
+            
 
 
 class IntensityMSENormalizer(object):
@@ -309,8 +329,8 @@ class SimplePoissonNLLNormalizer(object):
 
     2. **Normalization scaling**: Divides by 0.5 times the count of non-zero
        pixels in the measured patterns. This is because, roughly, each non-zero
-       pixel is expected to contribute to the Poisson NLL, if Poisson noise were
-       the only relevant source of noise in the data.
+       pixel is expected to contribute 0.5 to the Poisson NLL, if Poisson noise
+       were the only relevant source of noise in the data.
 
     The normalizer is stateful: it completes its accumulation phase on the
     first epoch by processing all patterns in the data, then applies the

tests/test_reconstructors.py

@@ -5,6 +5,7 @@ import cdtools
 import torch as t
 import numpy as np
 import pickle
+from matplotlib import pyplot as plt
 from copy import deepcopy
 
 

tests/tools/test_losses.py

@@ -21,15 +21,34 @@ def test_amplitude_mse():
     # First, test without a mask
     np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2)
     # np_result /=  data.size
-    torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim))
+    torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim),
+                                        use_sum=True)
     assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
 
     # Then, test with a mask
     np_result = np.sum(mask * (np.sqrt(data) - np.sqrt(sim))**2)
     # np_result /=  np.count_nonzero(mask * np.ones_like(data))
-    torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask))
+    torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim),
+                                        mask=t.from_numpy(mask), use_sum=True)
     assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
 
+    # Now, test the version with use_sum=False, the default
+
+    # First, test without a mask
+    np_result = np.mean((np.sqrt(data) - np.sqrt(sim))**2)
+    # np_result /=  data.size
+    torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim))
+    assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
+
+    # Then, test with a mask. Note that with a mask, the masked pixels
+    # should not contribute to the denominator for the mean.
+    np_result = np.sum(mask * (np.sqrt(data) - np.sqrt(sim))**2)
+    np_result /=  np.count_nonzero(mask * np.ones_like(data))
+    torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim),
+                                        mask=t.from_numpy(mask), use_sum=False)
+    assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
+
+
 
 def test_intensity_mse():
     # Make some fake data
@@ -39,6 +58,20 @@ def test_intensity_mse():
     # and define a simple mask that needs to be broadcast
     mask = (np.random.rand(100, 100) > 0.1).astype(bool)
 
+    # First, test without a mask
+    np_result = np.sum((data - sim)**2)
+    torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim),
+                                        use_sum=True)
+    assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
+
+    # Then, test with a mask
+    np_result = np.sum(mask * (data - sim)**2)
+    torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim),
+                                        mask=t.from_numpy(mask), use_sum=True)
+    assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
+    
+    # Now, test the version with use_sum=False, the default
+
     # First, test without a mask
     np_result = np.sum((data - sim)**2)
     np_result /= data.size
@@ -48,7 +81,8 @@ def test_intensity_mse():
     # Then, test with a mask
     np_result = np.sum(mask * (data - sim)**2)
     np_result /= np.count_nonzero(mask * np.ones_like(data))
-    torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask))
+    torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim),
+                                        mask=t.from_numpy(mask), use_sum=False)
     assert np.isclose(np_result, np.take(torch_result.numpy(), 0))