From d8fade9f327ce23170ddf81f7bdc6f580f336f3a Mon Sep 17 00:00:00 2001 From: allevitan Date: Mon, 13 Apr 2026 14:34:05 +0200 Subject: [PATCH] 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 | 5 +- examples/tutorial_simple_ptycho.py | 3 +- src/cdtools/models/bragg_2d_ptycho.py | 5 +- src/cdtools/models/fancy_ptycho.py | 5 +- src/cdtools/models/multislice_2d_ptycho.py | 5 +- src/cdtools/models/multislice_ptycho.py | 5 +- src/cdtools/models/rpi.py | 5 +- src/cdtools/models/simple_ptycho.py | 3 +- src/cdtools/tools/losses/losses.py | 64 ++++++++++++++-------- tests/test_reconstructors.py | 1 + tests/tools/test_losses.py | 40 +++++++++++++- 11 files changed, 102 insertions(+), 39 deletions(-) diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst index 48a2ac7..58f3464 100644 --- a/docs/source/tutorial.rst +++ b/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 diff --git a/examples/tutorial_simple_ptycho.py b/examples/tutorial_simple_ptycho.py index cd3c870..e5ff284 100644 --- a/examples/tutorial_simple_ptycho.py +++ b/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() diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index 4edb4ab..43f6360 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/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') diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 573295d..ea1ddd7 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/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') diff --git a/src/cdtools/models/multislice_2d_ptycho.py b/src/cdtools/models/multislice_2d_ptycho.py index 1679dd1..318b0fc 100644 --- a/src/cdtools/models/multislice_2d_ptycho.py +++ b/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') diff --git a/src/cdtools/models/multislice_ptycho.py b/src/cdtools/models/multislice_ptycho.py index bdb0de7..03b1b52 100644 --- a/src/cdtools/models/multislice_ptycho.py +++ b/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') diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index eb88b67..4bce45a 100644 --- a/src/cdtools/models/rpi.py +++ b/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') diff --git a/src/cdtools/models/simple_ptycho.py b/src/cdtools/models/simple_ptycho.py index 57c88e5..eb99802 100644 --- a/src/cdtools/models/simple_ptycho.py +++ b/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() diff --git a/src/cdtools/tools/losses/losses.py b/src/cdtools/tools/losses/losses.py index bb65858..e8b8c19 100644 --- a/src/cdtools/tools/losses/losses.py +++ b/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 diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py index 378330e..2602de8 100644 --- a/tests/test_reconstructors.py +++ b/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 diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py index 35e98c8..836413c 100644 --- a/tests/tools/test_losses.py +++ b/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))