From fd17ae30ba5333d64926b71a69c89a38dbfc4cf7 Mon Sep 17 00:00:00 2001 From: allevitan Date: Thu, 19 Mar 2026 17:14:08 +0100 Subject: [PATCH] Update tutorial to use loss instance attribute pattern Updates tutorial_simple_ptycho.py to match simple_ptycho.py: removes the def loss() method and instead assigns self.loss and self.loss_normalizer as instance attributes in __init__. Updates tutorial.rst accordingly: adds the loss assignment to the __init__ code block with an explanation, and removes def loss() from the forward model section. Co-Authored-By: Claude Sonnet 4.6 --- docs/source/tutorial.rst | 15 +++++++++------ examples/tutorial_simple_ptycho.py | 7 ++++--- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst index 8a2fa86..48a2ac7 100644 --- a/docs/source/tutorial.rst +++ b/docs/source/tutorial.rst @@ -251,7 +251,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 = tools.losses.amplitude_mse + 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 +272,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. 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. + Initialization from Dataset +++++++++++++++++++++++++++ @@ -347,7 +353,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 +363,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 diff --git a/examples/tutorial_simple_ptycho.py b/examples/tutorial_simple_ptycho.py index 67544d7..cd3c870 100644 --- a/examples/tutorial_simple_ptycho.py +++ b/examples/tutorial_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): @@ -102,9 +106,6 @@ class SimplePtycho(CDIModel): 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 = [