From 766a7d1adf038f893e35dcf82d497d357521017b Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 13 Oct 2025 14:20:38 +0200 Subject: [PATCH 1/8] Move to better supported strategy for avoiding circular imports --- src/cdtools/models/base.py | 17 +---------------- src/cdtools/reconstructors/adam.py | 9 +++++++-- src/cdtools/reconstructors/base.py | 9 +++++++-- src/cdtools/reconstructors/lbfgs.py | 10 ++++++++-- src/cdtools/reconstructors/sgd.py | 10 ++++++++-- 5 files changed, 31 insertions(+), 24 deletions(-) diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index 635cdfb..5d3ac9a 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -40,6 +40,7 @@ import time from scipy import io from contextlib import contextmanager from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch +from cdtools.reconstructors import AdamReconstructor, LBFGSReconstructor, SGDReconstructor from cdtools.datasets import CDataset from typing import List, Union, Tuple import os @@ -378,13 +379,7 @@ class CDIModel(t.nn.Module): # cdtools.reconstructors.AdamReconstructor to perform reconstructions # without creating a new reconstructor each time we update the # hyperparameters. - # - # The only way to do this is to make the Adam reconstructor an - # attribute of the model. But since the Adam reconstructor also - # depends on CDIModel, this seems to give rise to a circular import - # error unless we import cdtools.reconstructors within this method: if not hasattr(self, 'reconstructor'): - from cdtools.reconstructors import AdamReconstructor self.reconstructor = AdamReconstructor(model=self, dataset=dataset, subset=subset) @@ -449,12 +444,7 @@ class CDIModel(t.nn.Module): # cdtools.reconstructors.LBFGSReconstructor as an attribute to run # reconstructions without generating new reconstructors each time # CDIModel.LBFGS_optimize is called. - # - # Since the LBFGS reconstructor also depends on CDIModel, a circular - # import error arises unless we import cdtools.reconstructors within - # this method: if not hasattr(self, 'reconstructor'): - from cdtools.reconstructors import LBFGSReconstructor self.reconstructor = LBFGSReconstructor(model=self, dataset=dataset, subset=subset) @@ -524,12 +514,7 @@ class CDIModel(t.nn.Module): # cdtools.reconstructors.SGDReconstructor as an attribute to run # reconstructions without generating new reconstructors each time # CDIModel.SGD_optimize is called. - # - # Since the SGD reconstructor also depends on CDIModel, a circular - # import error arises unless we import cdtools.reconstructors within - # this method: if not hasattr(self, 'reconstructor'): - from cdtools.reconstructors import SGDReconstructor self.reconstructor = SGDReconstructor(model=self, dataset=dataset, subset=subset) diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py index 6eecc9c..4c54f1f 100644 --- a/src/cdtools/reconstructors/adam.py +++ b/src/cdtools/reconstructors/adam.py @@ -6,12 +6,17 @@ The Reconstructor class is designed to resemble so-called 'Trainer' classes that (in the language of the AI/ML folks) handles the 'training' of a model given some dataset and optimizer. """ +from __future__ import annotations +from typing import TYPE_CHECKING + import torch as t -from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset -from cdtools.models import CDIModel from typing import Tuple, List, Union from cdtools.reconstructors import Reconstructor +if TYPE_CHECKING: + from cdtools.models import CDIModel + from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset + __all__ = ['AdamReconstructor'] diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py index fd2cbb2..1937661 100644 --- a/src/cdtools/reconstructors/base.py +++ b/src/cdtools/reconstructors/base.py @@ -8,16 +8,21 @@ the 'training' of a model given some dataset and optimizer. The subclasses of Reconstructor are required to implement their own data loaders and optimizer adjusters """ +from __future__ import annotations +from typing import TYPE_CHECKING import torch as t from torch.utils import data as td import threading import queue import time -from cdtools.datasets import CDataset -from cdtools.models import CDIModel from typing import List, Union +if TYPE_CHECKING: + from cdtools.models import CDIModel + from cdtools.datasets import CDataset + + __all__ = ['Reconstructor'] diff --git a/src/cdtools/reconstructors/lbfgs.py b/src/cdtools/reconstructors/lbfgs.py index 8bbd7ae..60c6f97 100644 --- a/src/cdtools/reconstructors/lbfgs.py +++ b/src/cdtools/reconstructors/lbfgs.py @@ -6,12 +6,18 @@ The Reconstructor class is designed to resemble so-called 'Trainer' classes that (in the language of the AI/ML folks) handles the 'training' of a model given some dataset and optimizer. """ +from __future__ import annotations +from typing import TYPE_CHECKING + import torch as t -from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset -from cdtools.models import CDIModel from typing import List, Union from cdtools.reconstructors import Reconstructor +if TYPE_CHECKING: + from cdtools.models import CDIModel + from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset + + __all__ = ['LBFGSReconstructor'] diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py index a8f2ebc..5b7bf4e 100644 --- a/src/cdtools/reconstructors/sgd.py +++ b/src/cdtools/reconstructors/sgd.py @@ -6,12 +6,18 @@ The Reconstructor class is designed to resemble so-called 'Trainer' classes that (in the language of the AI/ML folks) handles the 'training' of a model given some dataset and optimizer. """ +from __future__ import annotations +from typing import TYPE_CHECKING + import torch as t -from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset -from cdtools.models import CDIModel from typing import List, Union from cdtools.reconstructors import Reconstructor +if TYPE_CHECKING: + from cdtools.models import CDIModel + from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset + + __all__ = ['SGDReconstructor'] From a089dc3b7266ff6120f50e17638a8f8ba6e57ce2 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 13 Oct 2025 14:26:07 +0200 Subject: [PATCH 2/8] Switch to a simpler pattern for the CDIModel._optimize functions that preserves the old behavior when the old pattern is used --- src/cdtools/models/base.py | 99 ++++++++++++++++++-------------------- 1 file changed, 48 insertions(+), 51 deletions(-) diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index 5d3ac9a..df347b6 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -375,25 +375,24 @@ class CDIModel(t.nn.Module): only the calculation speed. """ - # We want to have model.Adam_optimize call AND store - # cdtools.reconstructors.AdamReconstructor to perform reconstructions - # without creating a new reconstructor each time we update the - # hyperparameters. - if not hasattr(self, 'reconstructor'): - self.reconstructor = AdamReconstructor(model=self, - dataset=dataset, - subset=subset) + reconstructor = AdamReconstructor( + model=self, + dataset=dataset, + subset=subset, + ) # Run some reconstructions - return self.reconstructor.optimize(iterations=iterations, - batch_size=batch_size, - lr=lr, - betas=betas, - schedule=schedule, - amsgrad=amsgrad, - regularization_factor=regularization_factor, # noqa - thread=thread, - calculation_width=calculation_width) + return reconstructor.optimize( + iterations=iterations, + batch_size=batch_size, + lr=lr, + betas=betas, + schedule=schedule, + amsgrad=amsgrad, + regularization_factor=regularization_factor, # noqa + thread=thread, + calculation_width=calculation_width, + ) def LBFGS_optimize(self, iterations: int, @@ -440,24 +439,23 @@ class CDIModel(t.nn.Module): round of gradient accumulation. Does not affect the result, only the calculation speed. """ - # We want to have model.LBFGS_optimize store - # cdtools.reconstructors.LBFGSReconstructor as an attribute to run - # reconstructions without generating new reconstructors each time - # CDIModel.LBFGS_optimize is called. - if not hasattr(self, 'reconstructor'): - self.reconstructor = LBFGSReconstructor(model=self, - dataset=dataset, - subset=subset) + reconstructor = LBFGSReconstructor( + model=self, + dataset=dataset, + subset=subset, + ) # Run some reconstructions - return self.reconstructor.optimize(iterations=iterations, - lr=lr, - history_size=history_size, - regularization_factor=regularization_factor, # noqa - thread=thread, - calculation_width=calculation_width, - line_search_fn=line_search_fn) - + return reconstructor.optimize( + iterations=iterations, + lr=lr, + history_size=history_size, + regularization_factor=regularization_factor, # noqa + thread=thread, + calculation_width=calculation_width, + line_search_fn=line_search_fn, + ) + def SGD_optimize(self, iterations: int, dataset: CDataset, @@ -510,26 +508,25 @@ class CDIModel(t.nn.Module): round of gradient accumulation. """ - # We want to have model.SGD_optimize store - # cdtools.reconstructors.SGDReconstructor as an attribute to run - # reconstructions without generating new reconstructors each time - # CDIModel.SGD_optimize is called. - if not hasattr(self, 'reconstructor'): - self.reconstructor = SGDReconstructor(model=self, - dataset=dataset, - subset=subset) + reconstructor = SGDReconstructor( + model=self, + dataset=dataset, + subset=subset, + ) # Run some reconstructions - return self.reconstructor.optimize(iterations=iterations, - batch_size=batch_size, - lr=lr, - momentum=momentum, - dampening=dampening, - weight_decay=weight_decay, - nesterov=nesterov, - regularization_factor=regularization_factor, # noqa - thread=thread, - calculation_width=calculation_width) + return reconstructor.optimize( + iterations=iterations, + batch_size=batch_size, + lr=lr, + momentum=momentum, + dampening=dampening, + weight_decay=weight_decay, + nesterov=nesterov, + regularization_factor=regularization_factor, # noqa + thread=thread, + calculation_width=calculation_width, + ) def report(self): From aa684f27eb447ae7efd56bed0652943c682cca6c Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 13 Oct 2025 19:20:42 +0200 Subject: [PATCH 3/8] Change the pattern for Reconstructor so that the optimizer is defined at object creation, and move the dataloader creation logic to the base optimize() function as it was reused in all subclasses --- src/cdtools/reconstructors/adam.py | 45 ++++++++++------ src/cdtools/reconstructors/base.py | 83 +++++++++++++++++++++++------ src/cdtools/reconstructors/lbfgs.py | 33 ++++++------ src/cdtools/reconstructors/sgd.py | 39 +++++++------- 4 files changed, 133 insertions(+), 67 deletions(-) diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py index 4c54f1f..ac11ee5 100644 --- a/src/cdtools/reconstructors/adam.py +++ b/src/cdtools/reconstructors/adam.py @@ -51,10 +51,12 @@ class AdamReconstructor(Reconstructor): dataset: Ptycho2DDataset, subset: List[int] = None): - super().__init__(model, dataset, subset) - # Define the optimizer for use in this subclass - self.optimizer = t.optim.Adam(self.model.parameters()) + optimizer = t.optim.Adam(model.parameters()) + + super().__init__(model, dataset, optimizer, subset=subset) + + def adjust_optimizer(self, lr: int = 0.005, @@ -79,11 +81,13 @@ class AdamReconstructor(Reconstructor): param_group['betas'] = betas param_group['amsgrad'] = amsgrad + def optimize(self, iterations: int, batch_size: int = 15, lr: float = 0.005, betas: Tuple[float] = (0.9, 0.999), + custom_data_loader = None, schedule: bool = False, amsgrad: bool = False, regularization_factor: Union[float, List[float]] = None, @@ -99,6 +103,12 @@ class AdamReconstructor(Reconstructor): (formerly `CDIModel.AD_optimize`) to run a round of reconstruction once the dataloader and optimizer hyperparameters have been set up. + + The `batch_size` parameter sets the batch size for the default + dataloader. If a custom data loader is desired, it can be passed + in to the `custom_data_loader` argument, which will override the + `batch_size` and `shuffle` parameters + Parameters ---------- @@ -115,6 +125,9 @@ class AdamReconstructor(Reconstructor): schedule : bool Optional, create a learning rate scheduler (torch.optim.lr_scheduler._LRScheduler). + custom_data_loader : t.utils.data.DataLoader + Optional, a custom DataLoader to use. If set, will override + batch_size. amsgrad : bool Optional, whether to use the AMSGrad variant of this algorithm. regularization_factor : float or list(float) @@ -138,18 +151,13 @@ class AdamReconstructor(Reconstructor): f'{regularization_factor}, and schedule = {schedule}.\n' ) - # 1) The subset statement is contained in Reconstructor.__init__ - - # 2) Set up / re-initialize the data laoder - self.setup_dataloader(batch_size=batch_size, shuffle=shuffle) - - # 3) The optimizer is created in self.__init__, but the - # hyperparameters need to be set up with self.adjust_optimizer + # The optimizer is created in self.__init__, but the + # hyperparameters need to be set up with self.adjust_optimizer self.adjust_optimizer(lr=lr, betas=betas, amsgrad=amsgrad) - # 4) Set up the scheduler + # Set up the scheduler if schedule: self.scheduler = \ t.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, @@ -158,8 +166,13 @@ class AdamReconstructor(Reconstructor): else: self.scheduler = None - # 5) This is analagous to making a call to CDIModel.AD_optimize - return super(AdamReconstructor, self).optimize(iterations, - regularization_factor, - thread, - calculation_width) + # Now, we run the optimize routine defined in the base class + return super(AdamReconstructor, self).optimize( + iterations, + batch_size=batch_size, + custom_data_loader=custom_data_loader, + regularization_factor=regularization_factor, + thread=thread, + calculation_width=calculation_width, + shuffle=shuffle, + ) diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py index 1937661..bae992b 100644 --- a/src/cdtools/reconstructors/base.py +++ b/src/cdtools/reconstructors/base.py @@ -40,6 +40,8 @@ class Reconstructor: Model for CDI/ptychography reconstruction dataset: CDataset The dataset to reconstruct against + optimizer: torch.optim.Optimizer + The optimizer to use for the reconstruction subset : list(int) or int Optional, a pattern index or list of pattern indices to use @@ -55,31 +57,32 @@ class Reconstructor: def __init__(self, model: CDIModel, dataset: CDataset, + optimizer: t.optim.Optimizer, subset: Union[int, List[int]] = None): + # Store parameters as attributes of Reconstructor - self.subset = subset - - # Initialize attributes that must be defined by the subclasses - self.optimizer = None - self.scheduler = None - self.data_loader = None - - # Store the original model self.model = model + self.optimizer = optimizer - # Store the dataset + # Store the dataset, clipping it to a subset if needed if subset is not None: # if subset is just one pattern, turn into a list for convenience if isinstance(subset, int): subset = [subset] dataset = td.Subset(dataset, subset) + self.dataset = dataset + # Initialize attributes that must be defined by the subclasses + self.scheduler = None + self.data_loader = None + + def setup_dataloader(self, batch_size: int = None, shuffle: bool = True): """ - Sets up / re-initializes the dataloader. + Sets up or re-initializes the dataloader. Parameters ---------- @@ -96,6 +99,7 @@ class Reconstructor: else: self.data_loader = td.Dataloader(self.dataset) + def adjust_optimizer(self, **kwargs): """ Change hyperparameters for the utilized optimizer. @@ -105,7 +109,8 @@ class Reconstructor: """ raise NotImplementedError() - def _run_epoch(self, + + def run_epoch(self, stop_event: threading.Event = None, regularization_factor: Union[float, List[float]] = None, calculation_width: int = 10): @@ -133,6 +138,18 @@ class Reconstructor: diffraction pattern intensity """ + # Setting this as an explicit catch makes me feel more comfortable + # exposing it as a public method. This way a user won't be confused + # if they try to use this directly + if self.data_loader is None: + raise RuntimeError( + 'No data loader was defined. Please run ' + 'Reconstructor.setup_dataloader() before running ' + 'Reconstructor.run_epoch(), or use Reconstructor.optimize(), ' + 'which does it automatically.' + ) + + # Initialize some tracking variables normalization = 0 loss = 0 @@ -212,9 +229,12 @@ class Reconstructor: def optimize(self, iterations: int, + batch_size: int = 1, + custom_data_loader = None, regularization_factor: Union[float, List[float]] = None, thread: bool = True, - calculation_width: int = 10): + calculation_width: int = 10, + shuffle=True): """ Runs a round of reconstruction using the provided optimizer @@ -233,10 +253,25 @@ class Reconstructor: the plots. This behavior can be turned off by setting the keyword argument 'thread' to False. + The `batch_size` parameter sets the batch size for the default + dataloader. If a custom data loader is desired, it can be passed + in to the `custom_data_loader` argument, which will override the + `batch_size` and `shuffle` parameters + + Please see `AdamReconstructor.optimize()` for an example of how to + override this function when designing a subclass + Parameters ---------- iterations : int How many epochs of the algorithm to run. + batch_size : int + Optional, the batch size to use. Default is 1. This is typically + overridden by subclasses with an appropriate default for the + specific optimizer. + custom_data_loader : torch.utils.data.DataLoader + Optional, a custom DataLoader to use. Will override batch_size + if set. regularization_factor : float or list(float) Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method. @@ -247,6 +282,10 @@ class Reconstructor: Default 10, how many translations to pass through at once for each round of gradient accumulation. This does not affect the result, but may affect the calculation speed. + shuffle : bool + Optional, enable/disable shuffling of the dataset. This option + is intended for diagnostic purposes and should be left as True. + Yields ------ @@ -255,6 +294,11 @@ class Reconstructor: diffraction pattern intensity. """ + if custom_data_loader is None: + self.setup_dataloader(batch_size=batch_size, shuffle=shuffle) + else: + self.data_loader = custom_data_loader + # We store the current optimizer as a model parameter so that # it can be saved and loaded for checkpointing self.current_optimizer = self.optimizer @@ -270,8 +314,10 @@ class Reconstructor: yield float('nan') continue - yield self._run_epoch(regularization_factor=regularization_factor, # noqa - calculation_width=calculation_width) + yield self.run_epoch( + regularization_factor=regularization_factor, # noqa + calculation_width=calculation_width, + ) # But if we do want to thread, it's annoying: else: @@ -282,9 +328,12 @@ class Reconstructor: def target(): try: result_queue.put( - self._run_epoch(stop_event=stop_event, - regularization_factor=regularization_factor, # noqa - calculation_width=calculation_width)) + self.run_epoch( + stop_event=stop_event, + regularization_factor=regularization_factor, # noqa + calculation_width=calculation_width, + ) + ) except Exception as e: # If something bad happens, put the exception into the # result queue diff --git a/src/cdtools/reconstructors/lbfgs.py b/src/cdtools/reconstructors/lbfgs.py index 60c6f97..8907782 100644 --- a/src/cdtools/reconstructors/lbfgs.py +++ b/src/cdtools/reconstructors/lbfgs.py @@ -53,10 +53,16 @@ class LBFGSReconstructor(Reconstructor): dataset: Ptycho2DDataset, subset: List[int] = None): - super().__init__(model, dataset, subset) - # Define the optimizer for use in this subclass - self.optimizer = t.optim.LBFGS(self.model.parameters()) + optimizer = t.optim.LBFGS(model.parameters()) + + super().__init__( + model, + dataset, + optimizer, + subset=subset, + ) + def adjust_optimizer(self, lr: int = 0.005, @@ -121,20 +127,17 @@ class LBFGSReconstructor(Reconstructor): round of gradient accumulation. Does not affect the result, only the calculation speed. """ - # 1) The subset statement is contained in Reconstructor.__init__ - # 2) Set up / re-initialize the data loader. For LBFGS, we load - # all the data at once. - self.setup_dataloader(batch_size=len(self.dataset)) - - # 3) The optimizer is created in self.__init__, but the - # hyperparameters need to be set up with self.adjust_optimizer + # The optimizer is created in self.__init__, but the + # hyperparameters need to be set up with self.adjust_optimizer self.adjust_optimizer(lr=lr, history_size=history_size, line_search_fn=line_search_fn) - # 4) This is analagous to making a call to CDIModel.AD_optimize - return super(LBFGSReconstructor, self).optimize(iterations, - regularization_factor, - thread, - calculation_width) + # Now, we run the optimize routine defined in the base class + return super(LBFGSReconstructor, self).optimize( + iterations, + batch_size=len(self.dataset), + regularization_factor=regularization_factor, + thread=thread, + calculation_width=calculation_width) diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py index 5b7bf4e..cb1e26b 100644 --- a/src/cdtools/reconstructors/sgd.py +++ b/src/cdtools/reconstructors/sgd.py @@ -49,11 +49,17 @@ class SGDReconstructor(Reconstructor): dataset: Ptycho2DDataset, subset: List[int] = None): - super().__init__(model, dataset, subset) - # Define the optimizer for use in this subclass - self.optimizer = t.optim.SGD(self.model.parameters()) + optimizer = t.optim.SGD(model.parameters()) + super().__init__( + model, + dataset, + optimizer, + subset=subset, + ) + + def adjust_optimizer(self, lr: int = 0.005, momentum: float = 0, @@ -88,7 +94,7 @@ class SGDReconstructor(Reconstructor): def optimize(self, iterations: int, - batch_size: int = None, + batch_size: int = 15, lr: float = 2e-7, momentum: float = 0, dampening: float = 0, @@ -139,25 +145,20 @@ class SGDReconstructor(Reconstructor): Optional, enable/disable shuffling of the dataset. This option is intended for diagnostic purposes and should be left as True. """ - # 1) The subset statement is contained in Reconstructor.__init__ - # 2) Set up / re-initialize the data laoder - if batch_size is not None: - self.setup_dataloader(batch_size=batch_size, shuffle=shuffle) - else: - # Use default torch dataloader parameters - self.setup_dataloader(batch_size=1, shuffle=False) - - # 3) The optimizer is created in self.__init__, but the - # hyperparameters need to be set up with self.adjust_optimizer + # The optimizer is created in self.__init__, but the + # hyperparameters need to be set up with self.adjust_optimizer self.adjust_optimizer(lr=lr, momentum=momentum, dampening=dampening, weight_decay=weight_decay, nesterov=nesterov) - # 4) This is analagous to making a call to CDIModel.AD_optimize - return super(SGDReconstructor, self).optimize(iterations, - regularization_factor, - thread, - calculation_width) + # Now, we run the optimize routine defined in the base class + return super(SGDReconstructor, self).optimize( + iterations, + batch_size=batch_size, + regularization_factor=regularization_factor, + thread=thread, + calculation_width=calculation_width, + ) From 3c8de5dc192072fdc5e0753d00183e3c31cca41f Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 16 Oct 2025 14:16:12 +0200 Subject: [PATCH 4/8] Update the example codes --- examples/fancy_ptycho.py | 20 ++++++++++++++++---- examples/gold_ball_ptycho.py | 11 +++++++---- examples/gold_ball_split.py | 9 ++++++--- examples/simple_ptycho.py | 11 ++++++++++- 4 files changed, 39 insertions(+), 12 deletions(-) diff --git a/examples/fancy_ptycho.py b/examples/fancy_ptycho.py index 1cf5855..067f78d 100644 --- a/examples/fancy_ptycho.py +++ b/examples/fancy_ptycho.py @@ -19,19 +19,31 @@ 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 +recon = cdtools.reconstructors.AdamReconstructor(model, dataset) + # The learning rate parameter sets the alpha for Adam. # The beta parameters are (0.9, 0.999) by default # The batch size sets the minibatch size -for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10): +for loss in recon.optimize(50, lr=0.02, batch_size=10): print(model.report()) # Plotting is expensive, so we only do it every tenth epoch if model.epoch % 10 == 0: model.inspect(dataset) # It's common to chain several different reconstruction loops. Here, we -# started with an aggressive refinement to find the probe, and now we -# polish the reconstruction with a lower learning rate and larger minibatch -for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50): +# started with an aggressive refinement to find the probe in the previous +# loop, and now we polish the reconstruction with a lower learning rate +# and larger minibatch +for loss in recon.optimize(50, lr=0.005, batch_size=50): print(model.report()) if model.epoch % 10 == 0: model.inspect(dataset) diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 4747666..4971975 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -29,6 +29,7 @@ model = cdtools.models.FancyPtycho.from_dataset( probe_fourier_crop=pad ) + # This is a trick that my grandmother taught me, to combat the raster grid # pathology: we randomze the our initial guess of the probe positions. # The units here are pixels in the object array. @@ -42,17 +43,20 @@ device = 'cuda' model.to(device=device) dataset.get_as(device=device) +# Create the reconstructor +recon = cdtools.reconstructors.AdamReconstructor(model, dataset) + # This will save out the intermediate results if an exception is thrown # during the reconstruction with model.save_on_exception( 'example_reconstructions/gold_balls_earlyexit.h5', dataset): - for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50): + for loss in recon.optimize(20, lr=0.005, batch_size=50): print(model.report()) if model.epoch % 10 == 0: model.inspect(dataset) - for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100): + for loss in recon.optimize(50, lr=0.002, batch_size=100): print(model.report()) if model.epoch % 10 == 0: model.inspect(dataset) @@ -64,8 +68,7 @@ with model.save_on_exception( # Setting schedule=True automatically lowers the learning rate if # the loss fails to improve after 10 epochs - for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100, - schedule=True): + for loss in recon.optimize(100, lr=0.001, batch_size=100, schedule=True): print(model.report()) if model.epoch % 10 == 0: model.inspect(dataset) diff --git a/examples/gold_ball_split.py b/examples/gold_ball_split.py index 8624b9f..9fc5b08 100644 --- a/examples/gold_ball_split.py +++ b/examples/gold_ball_split.py @@ -36,15 +36,18 @@ for label, dataset in zip(labels, datasets): model.to(device=device) dataset.get_as(device=device) + # Create the reconstructor + recon = cdtools.reconstructors.AdamReconstructor(model, dataset) + # For batched reconstructions like this, there's no need to live-plot # the progress - for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50): + for loss in recon.optimize(20, lr=0.005, batch_size=50): print(model.report()) - for loss in model.Adam_optimize(50, dataset, lr=0.002, batch_size=100): + for loss in recon.optimize(50, lr=0.002, batch_size=100): print(model.report()) - for loss in model.Adam_optimize(100, dataset, lr=0.001, batch_size=100, + for loss in recon.optimize(100, lr=0.001, batch_size=100, schedule=True): print(model.report()) diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index 726af94..217d96b 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -1,3 +1,12 @@ +""" +Runs a very simple reconstruction using the SimplePtycho model, which was +designed to be an easy introduction to show how the models are made and used. + +For a more realistic example of how to use cdtools for real-world data, +look at fancy_ptycho.py and gold_ball_ptycho.py, both of which use the +more powerful FancyPtycho model and include more information on how to +correct for common sources of error. +""" import cdtools from matplotlib import pyplot as plt @@ -13,7 +22,7 @@ device = 'cuda' model.to(device=device) dataset.get_as(device=device) -# We run the actual reconstruction +# We run the reconstruction for loss in model.Adam_optimize(100, dataset, batch_size=10): # We print a quick report of the optimization status print(model.report()) From 8328b852d393a164aa89abb2660ae242fdfd21a7 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 16 Oct 2025 14:43:04 +0200 Subject: [PATCH 5/8] 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, From f022b5e2c4ba90f9876d842b5b7407aece1cf2f5 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 16 Oct 2025 15:58:28 +0200 Subject: [PATCH 6/8] Get the tests passing again after the change of the optimization functions in the model classes --- tests/models/test_fancy_ptycho.py | 9 +++++++-- tests/test_reconstructors.py | 16 +++++++++------- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py index 4182c93..8bc4d87 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -56,8 +56,8 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): dataset, n_modes=3, oversampling=2, - exponentiate_obj=True, dm_rank=2, + exponentiate_obj=True, probe_support_radius=120, propagation_distance=5e-3, units='mm', @@ -70,7 +70,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): model.to(device=reconstruction_device) dataset.get_as(device=reconstruction_device) - for loss in model.Adam_optimize(70, dataset, lr=0.02, batch_size=10): + for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10): print(model.report()) if show_plot and model.epoch % 10 == 0: model.inspect(dataset) @@ -79,6 +79,11 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): print(model.report()) if show_plot and model.epoch % 10 == 0: model.inspect(dataset) + + for loss in model.Adam_optimize(25, dataset, lr=0.001, batch_size=50): + print(model.report()) + if show_plot and model.epoch % 10 == 0: + model.inspect(dataset) model.tidy_probes() diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py index 3995b6f..2e30828 100644 --- a/tests/test_reconstructors.py +++ b/tests/test_reconstructors.py @@ -91,7 +91,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): ' reconstruction_device,', reconstruction_device) t.manual_seed(0) - for i, iterations in enumerate(epoch_tup): + # We only need to test the first loop to ensure it's identical + for i, iterations in enumerate(epoch_tup[:1]): for loss in model.Adam_optimize(iterations, dataset, lr=lr_tup[i], @@ -106,14 +107,15 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): model.inspect(dataset) model.compare(dataset) - # Ensure equivalency between the model reconstructions - assert np.allclose(model_recon.loss_history[-1], model.loss_history[-1]) + # Ensure equivalency between the model reconstructions during the first + # pass, where they should be identical + assert np.allclose(model_recon.loss_history[:epoch_tup[0]], model.loss_history[:epoch_tup[0]]) # Ensure reconstructions have reached a certain loss tolerance. This just # comes from running a reconstruction when it was working well and # choosing a rough value. If it triggers this assertion error, something # changed to make the final quality worse! - assert model.loss_history[-1] < 0.0001 + assert model_recon.loss_history[-1] < 0.0001 @pytest.mark.slow @@ -184,7 +186,7 @@ def test_LBFGS_RPI(optical_data_ss_cxi, print('Running reconstruction using CDIModel.LBFGS_optimize.' + 'optimize on provided reconstruction_device,', reconstruction_device) t.manual_seed(0) - for i, iterations in enumerate(epoch_tup): + for i, iterations in enumerate(epoch_tup[:1]): for loss in model.LBFGS_optimize(iterations, dataset, lr=0.4, @@ -198,12 +200,12 @@ def test_LBFGS_RPI(optical_data_ss_cxi, model.compare(dataset) # Check loss equivalency between the two reconstructions - assert np.allclose(model.loss_history[-1], model_recon.loss_history[-1]) + assert np.allclose(model.loss_history[:epoch_tup[0]], model_recon.loss_history[:epoch_tup[0]]) # The final loss when testing this was 2.28607e-3. Based on this, we set # a threshold of 2.3e-3 for the tested loss. If this value has been # exceeded, the reconstructions have gotten worse. - assert model.loss_history[-1] < 0.0023 + assert model_recon.loss_history[-1] < 0.0023 @pytest.mark.slow From 9d633074d8fa509087a48f26668ec3aa7f70d60c Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Fri, 17 Oct 2025 18:37:38 -0300 Subject: [PATCH 7/8] Apply suggestions from code review Co-authored-by: Dayne Yoshiki Sasaki <37006268+yoshikisd@users.noreply.github.com> --- src/cdtools/reconstructors/adam.py | 2 +- src/cdtools/reconstructors/base.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py index ac11ee5..e298bdb 100644 --- a/src/cdtools/reconstructors/adam.py +++ b/src/cdtools/reconstructors/adam.py @@ -87,7 +87,7 @@ class AdamReconstructor(Reconstructor): batch_size: int = 15, lr: float = 0.005, betas: Tuple[float] = (0.9, 0.999), - custom_data_loader = None, + custom_data_loader: t.utils.data.DataLoader = None, schedule: bool = False, amsgrad: bool = False, regularization_factor: Union[float, List[float]] = None, diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py index 5177b79..cf3de8c 100644 --- a/src/cdtools/reconstructors/base.py +++ b/src/cdtools/reconstructors/base.py @@ -232,7 +232,7 @@ class Reconstructor: def optimize(self, iterations: int, batch_size: int = 1, - custom_data_loader = None, + custom_data_loader: torch.utils.data.DataLoader = None, regularization_factor: Union[float, List[float]] = None, thread: bool = True, calculation_width: int = 10, From 7d390157199b90fefddea691a6d5eaba737ba9c9 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Fri, 17 Oct 2025 23:41:08 +0200 Subject: [PATCH 8/8] response to dayne's review --- tests/test_reconstructors.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py index 2e30828..132a8c2 100644 --- a/tests/test_reconstructors.py +++ b/tests/test_reconstructors.py @@ -16,7 +16,8 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): 2) We are only using the single-GPU dataloading method 3) Ensure `recon.model` points to the original `model` 4) Reconstructions performed by `Adam.optimize` and - `model.Adam_optimize` calls produce identical results. + `model.Adam_optimize` calls produce identical results when + run over one round of optimization. 5) The quality of the reconstruction remains below a specified threshold. 5) Ensure that the FancyPtycho model works fine and dandy with the @@ -129,7 +130,8 @@ def test_LBFGS_RPI(optical_data_ss_cxi, hyperparameters 2) Ensure `recon.model` points to the original `model` 3) Reconstructions performed by `LBFGS.optimize` and - `model.LBFGS_optimize` calls produce identical results. + `model.LBFGS_optimize` calls produce identical results when + run over one round of reconstruction. 4) The quality of the reconstruction remains below a specified threshold. 5) Ensure that the RPI model works fine and dandy with the @@ -216,7 +218,8 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): hyperparameters 3) Ensure `recon.model` points to the original `model` 4) Reconstructions performed by `SGD.optimize` and - `model.SGD_optimize` calls produce identical results. + `model.SGD_optimize` calls produce identical results + when run over one round of reconstruction. 5) The quality of the reconstruction remains below a specified threshold. 5) Ensure that the FancyPtycho model works fine and dandy with the