From 8328b852d393a164aa89abb2660ae242fdfd21a7 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 16 Oct 2025 14:43:04 +0200 Subject: [PATCH] Update the docs to include the reconstructors class and discuss their use in the examples --- docs/source/examples.rst | 10 ++++++++-- docs/source/index.rst | 1 + docs/source/reconstructors.rst | 5 +++++ examples/fancy_ptycho.py | 14 +++++--------- src/cdtools/reconstructors/base.py | 18 ++++++++++-------- 5 files changed, 29 insertions(+), 19 deletions(-) create mode 100644 docs/source/reconstructors.rst diff --git a/docs/source/examples.rst b/docs/source/examples.rst index 517f1b6..a3ef41e 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -31,9 +31,9 @@ When reading this script, note the basic workflow. After the data is loaded, a m Next, the model is moved to the GPU using the :code:`model.to` function. Any device understood by :code:`torch.Tensor.to` can be specified here. The next line is a bit more subtle - the dataset is told to move patterns to the GPU before passing them to the model using the :code:`dataset.get_as` function. This function does not move the stored patterns to the GPU. If there is sufficient GPU memory, the patterns can also be pre-moved to the GPU using :code:`dataset.to`, but the speedup is empirically quite small. -Once the device is selected, a reconstruction is run using :code:`model.Adam_optimize`. This is a generator function which will yield at every epoch, to allow some monitoring code to be run. +Once the device is selected, a reconstruction is run using :code:`model.Adam_optimize`. This is a generator function which will yield at the end of every epoch, to allow some monitoring code to be run. -Finally, the results can be studied using :code:`model.inspect(dataet)`, which creates or updates a set of plots showing the current state of the model parameters. :code:`model.compare(dataset)` is also called, which shows how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset. +Finally, the results can be studied using :code:`model.inspect(dataset)`, which creates or updates a set of plots showing the current state of the model parameters. :code:`model.compare(dataset)` is also called, which shows how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset. Fancy Ptycho @@ -63,6 +63,12 @@ By default, FancyPtycho will also optimize over the following model parameters, These corrections can be turned off (on) by calling :code:`model..requires_grad = False #(True)`. +Note as well two other changes that are made in this script, when compared to `simple_ptycho.py`. First, a `Reconstructor` object is explicitly created, in this case an `AdamReconstructor`. This object stores a model, dataset, and pytorch optimizer. It is then used to orchestrate the later reconstruction using a call to `Reconstructor.optimize()`. + +We use this pattern, instead of the simpler call to `model.Adam_optimize()`, because having the reconstructor store the optimizer as well as the model and dataset allows the moment estimates to persist between multiple rounds of optimization. This leads to the second change: In this script, we run two optimization loops. The first loop aggressively refines the probe, with a low minibatch size and a high learning rate. The second loop has a smaller learning rate and a larger batch size, which allow for a more precise final estimation of the object. + +In this case, we used one reconstructor, but it is possible to create additional reconstructors to zero out all the persistant information in the optimizer, if desired, or even to instantiate multiple reconstructors on the same model with different optimization algorithms (e.g. `model.LBFGS_optimize()`). + Gold Ball Ptycho ---------------- diff --git a/docs/source/index.rst b/docs/source/index.rst index c6fcb8d..b8cd288 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -9,6 +9,7 @@ general datasets models + reconstructors tools/index indices_tables diff --git a/docs/source/reconstructors.rst b/docs/source/reconstructors.rst new file mode 100644 index 0000000..2e78a74 --- /dev/null +++ b/docs/source/reconstructors.rst @@ -0,0 +1,5 @@ +Reconstructors +============== + +.. automodule:: cdtools.reconstructors + :members: diff --git a/examples/fancy_ptycho.py b/examples/fancy_ptycho.py index 067f78d..ebee68f 100644 --- a/examples/fancy_ptycho.py +++ b/examples/fancy_ptycho.py @@ -19,15 +19,11 @@ device = 'cuda' model.to(device=device) dataset.get_as(device=device) -# Now, to do the reconstruction we will use the more flexible pattern of -# creating an explicit reconstructor. This is what is used behind the -# scenes by the convenience functions model._optimize. For a long -# reconstruction script with multiple steps, it is better to create the -# reconstructor explicitly. -# -# The reconstructor will store the model and dataset and create an appropriate -# optimizer. This allows the optimizer to persist, along with e.g. estimates -# of the moments of individual parameters between loops +# For this script, we use a slightly different pattern where we explicitly +# create a `Reconstructor` class to orchestrate the reconstruction. The +# reconstructor will store the model and dataset and create an appropriate +# optimizer. This allows the optimizer to persist between loops, along with +# e.g. estimates of the moments of individual parameters recon = cdtools.reconstructors.AdamReconstructor(model, dataset) # The learning rate parameter sets the alpha for Adam. diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py index bae992b..5177b79 100644 --- a/src/cdtools/reconstructors/base.py +++ b/src/cdtools/reconstructors/base.py @@ -45,14 +45,16 @@ class Reconstructor: subset : list(int) or int Optional, a pattern index or list of pattern indices to use - Important attributes: - - **model** -- Always points to the core model used. - - **optimizer** -- A `torch.optim.Optimizer` that must be defined when - initializing the Reconstructor subclass. - - **scheduler** -- A `torch.optim.lr_scheduler` that may be defined during - the `optimize` method. - - **data_loader** -- A torch.utils.data.DataLoader that is defined by - calling the `setup_dataloader` method. + Attributes + ---------- + model : CDIModel + Points to the core model used. + optimizer : torch.optim.Optimizer + Must be defined when initializing the Reconstructor subclass. + scheduler : torch.optim.lr_scheduler, optional + May be defined during the ``optimize`` method. + data_loader : torch.utils.data.DataLoader + Defined by calling the ``setup_dataloader`` method. """ def __init__(self, model: CDIModel,