diff --git a/README.md b/README.md index 8381ff1..2796c5b 100644 --- a/README.md +++ b/README.md @@ -18,10 +18,10 @@ for loss in model.Adam_optimize(10, dataset): print(model.report()) # Save the results -model.save_to_h5('ptycho_results.h5', dataset) +model.save_to_h5('ptycho_results.h5') # And look at them! -model.inspect(dataset) # See the reconstructed object, probe, etc. +model.inspect() # See the reconstructed object, probe, etc. model.compare(dataset) # See how the simulated and measured patterns compare plt.show() ``` diff --git a/docs/source/examples.rst b/docs/source/examples.rst index 4c101c9..6f356cd 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -31,7 +31,7 @@ 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 the end of every epoch, to allow some monitoring code to be run. Inside the loop, :code:`model.inspect(dataset)` is called every epoch to live-update a set of plots showing the current state of the model parameters. +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. Inside the loop, :code:`model.inspect()` is called every epoch to live-update a set of plots showing the current state of the model parameters. Finally, :code:`model.compare(dataset)` is called to show how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset. @@ -69,7 +69,7 @@ We use this pattern, instead of the simpler call to :code:`model.Adam_optimize() 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. :code:`model.LBFGS_optimize()`). -Note also the use of :code:`min_interval=10` in the calls to :code:`model.inspect(dataset)`. Because generating plots can be expensive, passing a minimum interval (in seconds) prevents excessive replots. Finally, the call to :code:`model.inspect(dataset, replot_all=True)` at the end of the script reopens any plot windows that the user may have closed during the reconstruction, so that all results are visible at the end. +Note also the use of :code:`min_interval=10` in the calls to :code:`model.inspect()`. Because generating plots can be expensive, passing a minimum interval (in seconds) prevents excessive replots. Finally, the call to :code:`model.inspect(replot_all=True)` at the end of the script reopens any plot windows that the user may have closed during the reconstruction, so that all results are visible at the end. Gold Ball Ptycho @@ -86,6 +86,16 @@ Note also the use of :code:`model.save_on_exception` and :code:`model.save_to_h5 Finally, note that there are several small adjustments made to the script to counteract particular sources of error that are present in this dataset, for example the raster grid pathology caused by the scan pattern used. Also note that not every mixin is needed every time - in this case, we turn off optimization of the :code:`weights` parameter. +View Gold Ball Ptycho Results +----------------------------- + +This script shows how to load and view a saved ptychography reconstruction. + +.. literalinclude:: ../../examples/view_gold_ball_ptycho.py + +Note that :code:`obj_view_crop` and :code:`units` are directly set when loading from the saved reconstruction, because this information purely refers to the settings of the viewer in :code:`model.inspect()` and is not saved with the reconstruction. + + Near-Field Ptycho ----------------- diff --git a/docs/source/intro.rst b/docs/source/intro.rst index 22bc2c1..78d9f41 100644 --- a/docs/source/intro.rst +++ b/docs/source/intro.rst @@ -29,10 +29,10 @@ CDTools is an open source python library for ptychography and CDI reconstruction print(model.report()) # Save the results - model.save_to_h5('ptycho_results.h5', dataset) + model.save_to_h5('ptycho_results.h5') # And look at them! - model.inspect(dataset) # See the reconstructed object, probe, etc. + model.inspect() # See the reconstructed object, probe, etc. model.compare(dataset) # See how the simulated and measured patterns compare plt.show() diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst index 58f3464..b13d72c 100644 --- a/docs/source/tutorial.rst +++ b/docs/source/tutorial.rst @@ -408,8 +408,8 @@ In this case, we've made use of the convenience plotting functions defined in :c More advanced models like :code:`FancyPtycho` also define a :code:`plot_panel_list`, which groups related plots together into multi-subplot figures. The :code:`panel_plot_mode` argument (passed at construction time) controls whether these panels are rendered as combined multi-subplot figures or as individual windows. For a simple model like :code:`SimplePtycho`, :code:`plot_list` is sufficient. -Saving -++++++ +Saving and Loading +++++++++++++++++++ By default, a function :code:`model.save_results()` is defined, which returns a python dictionary with an entry, :code:`'state_dict'`, containing all the registered parameters and buffers in the model. It also contains a basic record of the model's training history. This function is used internally by :code:`model.save_to_h5()`, as well as all other convenience functions for saving results. @@ -439,9 +439,44 @@ Sometimes, it is also useful to return a more user-friendly version of the resul return {**base_results, **results} -However, it is perfectly possible to write a new ptychography model without overriding :code:`model.save_results()` +However, it is perfectly possible to write a new ptychography model without overriding :code:`model.save_results()`. + +Sometimes, it is useful to be able to load this saved reconstruction back into a cdtools model, either to continue a reconstruction from it or just to quickly view the results using the standard :code:`model.inspect()` function. For this purpose, we can override the function :code:`model.from_results_dict()`, which loads a model from the exact dictionary produced by :code:`model.save_results()`. This is also called internally by :code:`model.from_results_h5()`, which loads a model directly from a saved .h5 file. + +.. code-block:: python + + @classmethod + def from_results_dict(cls, results_dict): + """Reconstructs a SimplePtycho model from a results dictionary. + + Parameters + ---------- + results_dict : dict + The dictionary returned by save_results(), as loaded from an h5 file + or produced directly in memory. + + Returns + ------- + model : SimplePtycho + A fully reconstructed model with all parameters, buffers, and + training metadata restored. + """ + sd = results_dict['state_dict'] + + model = cls( + wavelength=sd['wavelength'], + probe_basis=sd['probe_basis'], + probe_guess=sd['probe'], # normalized; probe_norm restored by _load_results_dict + obj_guess=sd['obj'], + min_translation=sd['min_translation'], + ) + model._load_results_dict(results_dict) + return model + +Here, we first directly load the model by initializing the object using the main parameters stored in the results which are needed to properly run through the model initialization. Then, we use the private method :code:`model._load_results_dict(results_dict)` to load the standard information - like the current epoch, loss history, and so forth, as well as to populate each element of the state dict from the saved state_dict dictionary - information such as the :code:`probe_norm`. + +With these functions, it is now possible to easily and quickly save and load the reconstructions produced by our new model! - Testing +++++++ @@ -466,10 +501,10 @@ We can test this model with a simple script, in examples/tutorial_finale.py. By dataset.get_as(device='mps')#cuda') for loss in model.Adam_optimize(10, dataset): - model.inspect(dataset) + model.inspect() print(model.report()) - model.inspect(dataset) + model.inspect() model.compare(dataset) plt.show() diff --git a/examples/example_reconstructions/gold_balls.h5 b/examples/example_reconstructions/gold_balls.h5 index 2e9db13..2420e56 100644 Binary files a/examples/example_reconstructions/gold_balls.h5 and b/examples/example_reconstructions/gold_balls.h5 differ diff --git a/examples/fancy_ptycho.py b/examples/fancy_ptycho.py index d6dfc74..14d75f8 100644 --- a/examples/fancy_ptycho.py +++ b/examples/fancy_ptycho.py @@ -34,7 +34,7 @@ for loss in recon.optimize(50, lr=0.02, batch_size=10): print(model.report()) # Because plotting can be expensive, setting a minimum plotting interval # (in seconds) can avoid excessive replots. - model.inspect(dataset, min_interval=10) + model.inspect(min_interval=10) # It's common to chain several different reconstruction loops. Here, we # started with an aggressive refinement to find the probe in the previous @@ -42,12 +42,12 @@ for loss in recon.optimize(50, lr=0.02, batch_size=10): # and larger minibatch for loss in recon.optimize(50, lr=0.005, batch_size=50): print(model.report()) - model.inspect(dataset, min_interval=10) + model.inspect(min_interval=10) # This orthogonalizes the recovered probe modes model.tidy_probes() # Setting replot_all will reopen any windows which were closed earlier -model.inspect(dataset, replot_all=True) +model.inspect(replot_all=True) model.compare(dataset) plt.show() diff --git a/examples/fancy_ptycho_inline.ipynb b/examples/fancy_ptycho_inline.ipynb index f5cab0f..b0dc876 100644 --- a/examples/fancy_ptycho_inline.ipynb +++ b/examples/fancy_ptycho_inline.ipynb @@ -66,7 +66,7 @@ "# Workaround reconstruction pattern for interactive plotting in jupyter:\n", "# First, a standalone cell to plot the current model state\n", "\n", - "model.inspect(dataset, replot_all=True);" + "model.inspect(replot_all=True);" ] }, { @@ -96,7 +96,7 @@ "source": [ "# Save out the results\n", "\n", - "model.save_to_h5('lab_ptycho_reconstruction.h5', dataset)" + "model.save_to_h5('lab_ptycho_reconstruction.h5');" ] }, { @@ -114,7 +114,7 @@ "model.tidy_probes()\n", "\n", "# Final plotting\n", - "model.inspect(dataset)\n", + "model.inspect()\n", "model.compare(dataset);" ] }, diff --git a/examples/fancy_ptycho_interactive.ipynb b/examples/fancy_ptycho_interactive.ipynb index 0e9adad..8d5343f 100644 --- a/examples/fancy_ptycho_interactive.ipynb +++ b/examples/fancy_ptycho_interactive.ipynb @@ -54,7 +54,7 @@ "# Then, create a reconstructor object and view the initialized model\n", "\n", "recon = cdtools.reconstructors.AdamReconstructor(model, dataset)\n", - "model.inspect(dataset);" + "model.inspect();" ] }, { @@ -72,12 +72,12 @@ "while model.epoch < 50:\n", " for loss in recon.optimize(1, lr=0.02, batch_size=10):\n", " print(model.report())\n", - " model.inspect(dataset, min_interval=10)\n", + " model.inspect(min_interval=10)\n", "\n", "while model.epoch < 100:\n", " for loss in recon.optimize(1, lr=0.005, batch_size=10):\n", " print(model.report())\n", - " model.inspect(dataset, min_interval=10)" + " model.inspect(min_interval=10)" ] }, { @@ -89,7 +89,7 @@ "source": [ "# Save out the results\n", "\n", - "model.save_to_h5('lab_ptycho_reconstruction.h5', dataset)" + "model.save_to_h5('lab_ptycho_reconstruction.h5')" ] }, { @@ -107,7 +107,7 @@ "model.tidy_probes()\n", "\n", "# Final plotting\n", - "model.inspect(dataset)\n", + "model.inspect()\n", "model.compare(dataset);" ] }, diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index aeb510c..ff0c54f 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -50,15 +50,15 @@ 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): + 'example_reconstructions/gold_balls_earlyexit.h5'): for loss in recon.optimize(20, lr=0.005, batch_size=50): print(model.report()) - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) for loss in recon.optimize(50, lr=0.002, batch_size=100): print(model.report()) - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) # We can often reset our guess of the probe positions once we have a # good guess of probe and object, but in this case it causes the @@ -69,14 +69,14 @@ with model.save_on_exception( # the loss fails to improve after 10 epochs for loss in recon.optimize(100, lr=0.001, batch_size=100, schedule=True): print(model.report()) - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) model.tidy_probes() # This saves the final result -model.save_to_h5('example_reconstructions/gold_balls.h5', dataset) +model.save_to_h5('example_reconstructions/gold_balls.h5') -model.inspect(dataset, replot_all=True) +model.inspect(replot_all=True) model.compare(dataset) plt.show() diff --git a/examples/gold_ball_split.py b/examples/gold_ball_split.py index fde19c6..b73b508 100644 --- a/examples/gold_ball_split.py +++ b/examples/gold_ball_split.py @@ -54,4 +54,4 @@ for label, dataset in zip(labels, datasets): model.tidy_probes() - model.save_to_h5(f'example_reconstructions/gold_balls_{label}.h5', dataset) + model.save_to_h5(f'example_reconstructions/gold_balls_{label}.h5') diff --git a/examples/near_field_ptycho.py b/examples/near_field_ptycho.py index 47da91d..55a9086 100644 --- a/examples/near_field_ptycho.py +++ b/examples/near_field_ptycho.py @@ -35,21 +35,21 @@ if t.cuda.is_available(): model.to(device='cuda') dataset.get_as(device='cuda') -model.inspect(dataset) +model.inspect() recon = cdtools.reconstructors.AdamReconstructor(model, dataset) for loss in recon.optimize(100, lr=0.04, batch_size=10): print(model.report()) - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) for loss in recon.optimize(50, lr=0.005, batch_size=50): print(model.report()) - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) # This orthogonalizes the recovered probe modes model.tidy_probes() -model.inspect(dataset, replot_all=True) +model.inspect(replot_all=True) model.compare(dataset) plt.show() diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index 41c0c6c..dfd78c1 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -23,14 +23,14 @@ if t.cuda.is_available(): model.to(device='cuda') dataset.get_as(device='cuda') -model.inspect(dataset) +model.inspect() # 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()) # And liveplot the updates to the model as they happen - model.inspect(dataset) + model.inspect() # We open a comparison of the simulated and measured data model.compare(dataset) diff --git a/examples/transmission_RPI.py b/examples/transmission_RPI.py index 32c02d5..c972db2 100644 --- a/examples/transmission_RPI.py +++ b/examples/transmission_RPI.py @@ -31,19 +31,19 @@ if t.cuda.is_available(): # The regularization is an L2 regularizer that empirically helps accelerate # convergence for loss in model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05]): - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) print(model.report()) # Now we use the regularizer to damp all but the top modes for loss in model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1]): - model.inspect(dataset, min_interval=5) + model.inspect(min_interval=5) print(model.report()) # Save results to an h5 file -model.save_to_h5('example_reconstructions/transmission_RPI.h5', dataset) +model.save_to_h5('example_reconstructions/transmission_RPI.h5') # Finally, we plot the results -model.inspect(dataset, replot_all=True) +model.inspect(replot_all=True) model.compare(dataset) plt.show() diff --git a/examples/tutorial_finale.py b/examples/tutorial_finale.py index 4e95ec6..eed4514 100644 --- a/examples/tutorial_finale.py +++ b/examples/tutorial_finale.py @@ -17,9 +17,9 @@ if t.cuda.is_available(): dataset.get_as(device='cuda') for loss in model.Adam_optimize(10, dataset): - model.inspect(dataset) + model.inspect() print(model.report()) -model.inspect(dataset) +model.inspect() model.compare(dataset) plt.show() diff --git a/examples/tutorial_simple_ptycho.py b/examples/tutorial_simple_ptycho.py index e5ff284..685e682 100644 --- a/examples/tutorial_simple_ptycho.py +++ b/examples/tutorial_simple_ptycho.py @@ -132,7 +132,7 @@ class SimplePtycho(CDIModel): }, ] - def save_results(self, dataset): + def save_results(self, dataset=None): # This will save out everything needed to recreate the object # in the same state, but it's not the best formatted. base_results = super().save_results() @@ -152,3 +152,32 @@ class SimplePtycho(CDIModel): } return {**base_results, **results} + + + @classmethod + def from_results_dict(cls, results_dict): + """Reconstructs a SimplePtycho model from a results dictionary. + + Parameters + ---------- + results_dict : dict + The dictionary returned by save_results(), as loaded from an h5 file + or produced directly in memory. + + Returns + ------- + model : SimplePtycho + A fully reconstructed model with all parameters, buffers, and + training metadata restored. + """ + sd = results_dict['state_dict'] + + model = cls( + wavelength=sd['wavelength'], + probe_basis=sd['probe_basis'], + probe_guess=sd['probe'], # normalized; probe_norm restored by _load_results_dict + obj_guess=sd['obj'], + min_translation=sd['min_translation'], + ) + model._load_results_dict(results_dict) + return model diff --git a/examples/view_gold_ball_ptycho.py b/examples/view_gold_ball_ptycho.py new file mode 100644 index 0000000..86a0c7b --- /dev/null +++ b/examples/view_gold_ball_ptycho.py @@ -0,0 +1,11 @@ +import cdtools +from matplotlib import pyplot as plt + +model = cdtools.models.FancyPtycho.from_results_h5( + 'example_reconstructions/gold_balls.h5', + obj_view_crop=260, # How far in to crop from the edge + units='um', # The units to display in +) + +model.inspect() +plt.show() diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index 5b990a0..4761b6e 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -36,7 +36,8 @@ from matplotlib import ticker import numpy as np import time from contextlib import contextmanager -from cdtools.tools.data import nested_dict_to_h5, nested_dict_to_numpy, nested_dict_to_torch +import cdtools +from cdtools.tools.data import nested_dict_to_h5, nested_dict_to_numpy, nested_dict_to_torch, h5_to_nested_dict from cdtools.reconstructors import AdamReconstructor, LBFGSReconstructor, SGDReconstructor from cdtools.datasets import CDataset from typing import List, Union, Tuple @@ -196,6 +197,8 @@ class CDIModel(t.nn.Module): getattr(self.loss, '__name__', None) or getattr(self.loss.func, '__name__', str(self.loss)) ), + 'model_class': type(self).__name__, + 'cdtools_version': cdtools.__version__, } @@ -210,7 +213,71 @@ class CDIModel(t.nn.Module): Accepts any additional args that model.save_results needs, for this model """ return nested_dict_to_h5(filename, self.save_results(*args)) - + + + def _load_results_dict(self, results_dict): + """Restores model state and training metadata from a results dictionary. + + This is the kernel used by from_results_dict implementations. It loads + the state_dict (all parameters and buffers) and restores loss history, + epoch count, and training history. + + Parameters + ---------- + results_dict : dict + The dictionary returned by save_results(), as loaded from an h5 file + or produced directly in memory. + """ + state_dict = nested_dict_to_torch(results_dict['state_dict']) + self.load_state_dict(state_dict, strict=False) + self.loss_history = list(results_dict.get('loss_history', [])) + self.epoch = int(results_dict.get('epoch', 0)) + self.training_history = str(results_dict.get('training_history', '')) + + + @classmethod + def from_results_dict(cls, results_dict): + """Reconstructs a model from the dictionary returned by save_results(). + + Must be implemented by each subclass. The base class raises + NotImplementedError. + + Parameters + ---------- + results_dict : dict + The dictionary returned by save_results(), as loaded from an h5 file + or produced directly in memory. + + Returns + ------- + model : CDIModel subclass + A fully reconstructed model with all parameters, buffers, and + training metadata restored. + """ + raise NotImplementedError() + + + @classmethod + def from_results_h5(cls, filename, *args, **kwargs): + """Reconstructs a model directly from a saved .h5 result file. + + Reads the file into a dictionary and delegates to cls.from_results_dict. + Subclasses inherit this method and only need to implement from_results_dict. + + Parameters + ---------- + filename : str or Path + Path to the .h5 file saved by save_to_h5. + + Returns + ------- + model : CDIModel subclass + A fully reconstructed model with all parameters, buffers, and + training metadata restored. + """ + return cls.from_results_dict( + h5_to_nested_dict(filename), *args, **kwargs) + @contextmanager def save_on_exit(self, filename, *args, exception_filename=None): diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index 43f6360..04f4369 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -83,6 +83,7 @@ class Bragg2DPtycho(CDIModel): obj_view_crop=0, panel_plot_mode=False, plot_level=1, + translations=None, ): # We need the detector geometry @@ -176,7 +177,7 @@ class Bragg2DPtycho(CDIModel): shape = [s//oversampling for s in self.probe[0]] background = 1e-6 * t.ones(shape, dtype=t.float32) - self.background = t.nn.Parameter(background) + self.background = t.nn.Parameter(t.as_tensor(background, dtype=dtype)) if weights is None: self.weights = None @@ -255,7 +256,11 @@ class Bragg2DPtycho(CDIModel): self.loss_normalizer = tools.losses.IntensityMSENormalizer() else: raise KeyError('Specified loss function not supported') - + + if translations is not None: + self.register_buffer('original_translations', + t.as_tensor(translations, dtype=dtype)) + @classmethod def from_dataset( @@ -473,6 +478,7 @@ class Bragg2DPtycho(CDIModel): units=units, panel_plot_mode=panel_plot_mode, plot_level=plot_level, + translations=translations, ) @@ -590,11 +596,25 @@ class Bragg2DPtycho(CDIModel): mask=mask) - def corrected_translations(self,dataset): - translations = dataset.translations.to(dtype=self.probe.real.dtype, - device=self.probe.device) - t_offset = tools.interactions.pixel_to_translations(self.obj_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal) - return translations + t_offset + def corrected_translations(self, dataset=None): + if dataset is not None: + translations = dataset.translations.to( + dtype=self.probe.real.dtype, device=self.probe.device) + elif (hasattr(self, 'original_translations') and + self.original_translations is not None): + translations = self.original_translations.to( + dtype=self.probe.real.dtype, device=self.probe.device) + else: + raise ValueError( + 'Must provide a dataset or have original_translations stored ' + 'internally (via from_dataset or from_results_dict).') + if self.translation_offsets is not None: + t_offset = tools.interactions.pixel_to_translations( + self.obj_basis, + self.translation_offsets * self.translation_scale, + surface_normal=self.surface_normal) + return translations + t_offset + return translations plot_list = [ @@ -679,13 +699,48 @@ class Bragg2DPtycho(CDIModel): units=self.units, )}, {'title': 'Corrected Translations', - 'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)}, + 'plot_func': lambda self, fig: p.plot_translations(self.corrected_translations(), fig=fig, units=self.units)}, {'title': 'Background', 'plot_func': lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)}, ] - def save_results(self, dataset): + @classmethod + def from_results_dict(cls, results_dict, obj_view_crop=0, units='um'): + sd = results_dict['state_dict'] + translation_offsets = sd.get('translation_offsets') + model = cls( + wavelength=sd['wavelength'], + detector_geometry={ + 'basis': sd['det_basis'], + 'distance': sd.get('det_distance'), + 'corner': sd.get('det_corner'), + }, + obj_basis=sd['obj_basis'], + probe_guess=sd['probe'], + obj_guess=sd['obj'], + min_translation=sd.get('min_translation', np.array([0., 0.])), + probe_basis=sd.get('probe_basis'), + median_propagation=sd.get('median_propagation', 0.0), + background=sd['background'], + translation_offsets=translation_offsets, + mask=sd.get('mask'), + weights=sd.get('weights'), + translation_scale=float(sd.get('translation_scale', 1.0)), + saturation=sd.get('saturation'), + oversampling=int(sd.get('oversampling', 1)), + propagate_probe=bool(sd.get('propagate_probe', True)), + correct_tilt=bool(sd.get('correct_tilt', True)), + loss=results_dict.get('loss_function', 'amplitude mse'), + obj_view_crop=obj_view_crop, + units=units, + translations=sd.get('original_translations'), + ) + model._load_results_dict(results_dict) + return model + + + def save_results(self, dataset=None): # This will save out everything needed to recreate the object # in the same state, but it's not the best formatted. For example, # "background" stores the square root of the background, etc. @@ -694,8 +749,11 @@ class Bragg2DPtycho(CDIModel): # We also save out the main results in a more readable format obj_basis = self.obj_basis.detach().cpu().numpy() probe_basis = self.probe_basis.detach().cpu().numpy() - translations=self.corrected_translations(dataset).detach().cpu().numpy() - original_translations = dataset.translations.detach().cpu().numpy() + translations = self.corrected_translations(dataset).detach().cpu().numpy() + if dataset is not None: + original_translations = dataset.translations.detach().cpu().numpy() + else: + original_translations = self.original_translations.detach().cpu().numpy() probe = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() obj = self.obj.detach().cpu().numpy() diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 3c70f86..50def1f 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -46,6 +46,7 @@ class FancyPtycho(CDIModel): inv_angular_spectrum_propagator=None, panel_plot_mode=True, plot_level=2, + translations=None, ): super(FancyPtycho, self).__init__(panel_plot_mode=panel_plot_mode, @@ -149,6 +150,7 @@ class FancyPtycho(CDIModel): obj_view_crop:-obj_view_crop] else: self.obj_view_slice = np.s_[:,:] + # TODO: perhaps not working anymore for fourier cropped probes if background is None: @@ -156,7 +158,7 @@ class FancyPtycho(CDIModel): shape = [s//oversampling for s in self.probe[0]] background = 1e-6 * t.ones(shape, dtype=t.float32) - self.background = t.nn.Parameter(background) + self.background = t.nn.Parameter(t.as_tensor(background, dtype=dtype)) if weights is None: self.weights = None @@ -233,6 +235,10 @@ class FancyPtycho(CDIModel): else: raise KeyError('Specified loss function not supported') + if translations is not None: + self.register_buffer('original_translations', + t.as_tensor(translations, dtype=dtype)) + @classmethod def from_dataset(cls, @@ -531,6 +537,7 @@ class FancyPtycho(CDIModel): inv_angular_spectrum_propagator=inv_angular_spectrum_propagator, panel_plot_mode=panel_plot_mode, plot_level=plot_level, + translations=translations, ) @@ -718,9 +725,18 @@ class FancyPtycho(CDIModel): mask=mask) - def corrected_translations(self, dataset): - translations = dataset.translations.to( - dtype=t.float32, device=self.probe.device) + def corrected_translations(self, dataset=None): + if dataset is not None: + translations = dataset.translations.to( + dtype=t.float32, device=self.probe.device) + elif (hasattr(self, 'original_translations') and + self.original_translations is not None): + translations = self.original_translations.to( + dtype=t.float32, device=self.probe.device) + else: + raise ValueError( + 'Must provide a dataset or have original_translations stored ' + 'internally (via from_dataset or from_results_dict).') if (hasattr(self, 'translation_offsets') and self.translation_offsets is not None): t_offset = tools.interactions.pixel_to_translations( @@ -900,7 +916,7 @@ class FancyPtycho(CDIModel): return probe_intensities - def plot_wavefront_variation(self, dataset, fig=None, mode='amplitude', **kwargs): + def plot_wavefront_variation(self, dataset=None, fig=None, mode='amplitude', **kwargs): def get_probes(idx): basis_prs = self.probe * self.probe_support[..., :, :] prs = t.sum(self.weights[idx, :, :, None, None] * basis_prs, @@ -934,7 +950,7 @@ class FancyPtycho(CDIModel): **kwargs), - def plot_illumination_intensity(self, fig, dataset): + def plot_illumination_intensity(self, fig, dataset=None): """Plots the probe intensity nanomap. Only used to make a plot for the plot list.""" p.plot_nanomap( self.corrected_translations(dataset), @@ -949,10 +965,14 @@ class FancyPtycho(CDIModel): plt.gca().set_aspect('equal') - def plot_translations_and_originals(self, fig, dataset): + def plot_translations_and_originals(self, fig, dataset=None): """Only used to make a plot for the plot list.""" + if dataset is not None: + original_translations = dataset.translations + else: + original_translations = self.original_translations p.plot_translations( - dataset.translations, + original_translations, fig=fig, units=self.units, label='original translations', @@ -1086,7 +1106,7 @@ class FancyPtycho(CDIModel): { 'title': 'Illumination Intensity', 'subplot': (0,1), - 'plot_func': lambda self, fig, dataset: self.plot_illumination_intensity(fig, dataset), + 'plot_func': lambda self, fig: self.plot_illumination_intensity(fig), }, { 'title': 'Detector Background', @@ -1096,7 +1116,7 @@ class FancyPtycho(CDIModel): { 'title': 'Corrected Translations', 'subplot': (0,2), - 'plot_func': lambda self, fig, dataset: self.plot_translations_and_originals(fig, dataset), + 'plot_func': lambda self, fig: self.plot_translations_and_originals(fig), }, { 'title': 'Loss History', @@ -1115,8 +1135,8 @@ class FancyPtycho(CDIModel): { 'title': '% of Power in Top Mode', 'subplot': (0,0), - 'plot_func': lambda self, fig, dataset: p.plot_nanomap( - self.corrected_translations(dataset), + 'plot_func': lambda self, fig: p.plot_nanomap( + self.corrected_translations(), 100 * t.stack([ analysis.calc_mode_power_fractions( self.probe.data, @@ -1147,8 +1167,7 @@ class FancyPtycho(CDIModel): {'title': 'Per-Exposure Probe Intensity', 'plot_level': 3, 'figure_size': (8,5.3), - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', @@ -1157,8 +1176,7 @@ class FancyPtycho(CDIModel): {'title': 'Per-Exposure Probe Amplitudes', 'plot_level': 3, 'figure_size': (8,5.3), - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', @@ -1167,8 +1185,7 @@ class FancyPtycho(CDIModel): {'title': 'Per-Exposure Probe Phases', 'plot_level': 3, 'figure_size': (8,5.3), - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', @@ -1177,7 +1194,7 @@ class FancyPtycho(CDIModel): ] - def save_results(self, dataset): + def save_results(self, dataset=None): # This will save out everything needed to recreate the object # in the same state, but it's not the best formatted. For example, # "background" stores the square root of the background, etc. @@ -1186,8 +1203,11 @@ class FancyPtycho(CDIModel): # We also save out the main results in a more readable format obj_basis = self.obj_basis.detach().cpu().numpy() probe_basis = self.probe_basis.detach().cpu().numpy() - translations=self.corrected_translations(dataset).detach().cpu().numpy() - original_translations = dataset.translations.detach().cpu().numpy() + translations = self.corrected_translations(dataset).detach().cpu().numpy() + if dataset is not None: + original_translations = dataset.translations.detach().cpu().numpy() + else: + original_translations = self.original_translations.detach().cpu().numpy() probe = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() obj = self.obj.detach().cpu().numpy() @@ -1210,3 +1230,71 @@ class FancyPtycho(CDIModel): } return {**base_results, **results} + + + @classmethod + def from_results_dict( + cls, + results_dict, + obj_view_crop=0, + units='um', + ): + """Reconstructs a FancyPtycho model from a results dictionary. + + Parameters + ---------- + results_dict : dict + The dictionary returned by save_results(), as loaded from an h5 file + or produced directly in memory. + + Returns + ------- + model : FancyPtycho + A fully reconstructed model with all parameters, buffers, and + training metadata restored. + """ + sd = results_dict['state_dict'] + + # For optional Parameters (translation_offsets, weights, qe_mask, etc.), + # we pass the saved values directly so they get registered as + # Parameters/buffers before _load_results_dict overwrites them with the + # exact saved state via load_state_dict. + translation_offsets = sd.get('translation_offsets') + + model = cls( + wavelength=sd['wavelength'], + detector_geometry={ + 'basis': sd['det_basis'], + 'distance': sd.get('det_distance'), + 'corner': sd.get('det_corner'), + }, + obj_basis=sd['obj_basis'], + probe_guess=sd['probe'], # normalized; probe_norm restored by _load_results_dict + obj_guess=sd['obj'], + surface_normal=sd.get('surface_normal', np.array([0., 0., 1.])), + min_translation=sd.get('min_translation', np.array([0., 0.])), + background=sd['background'], # sqrt form; restored exactly by _load_results_dict + probe_basis=sd.get('probe_basis'), + translation_offsets=translation_offsets, # overwritten by _load_results_dict + probe_fourier_shifts=sd.get('probe_fourier_shifts'), # overwritten by _load_results_dict + mask=sd.get('mask'), + weights=sd.get('weights'), # overwritten by _load_results_dict + qe_mask=sd.get('qe_mask'), # overwritten by _load_results_dict + saturation=sd.get('saturation'), + translation_scale=float(sd.get('translation_scale', 1.0)), + oversampling=int(sd.get('oversampling', 1)), + fourier_probe=bool(sd.get('fourier_probe', False)), + loss=results_dict.get('loss_function', 'amplitude mse'), + simulate_probe_translation=bool(sd.get('simulate_probe_translation', False)), + simulate_finite_pixels=bool(sd.get('simulate_finite_pixels', False)), + exponentiate_obj=bool(sd.get('exponentiate_obj', False)), + phase_only=bool(sd.get('phase_only', False)), + near_field=bool(sd.get('near_field', False)), + angular_spectrum_propagator=sd.get('angular_spectrum_propagator'), + inv_angular_spectrum_propagator=sd.get('inv_angular_spectrum_propagator'), + translations=sd.get('original_translations'), + obj_view_crop=obj_view_crop, + units=units, + ) + model._load_results_dict(results_dict) + return model diff --git a/src/cdtools/models/multislice_ptycho.py b/src/cdtools/models/multislice_ptycho.py index 03b1b52..3547351 100644 --- a/src/cdtools/models/multislice_ptycho.py +++ b/src/cdtools/models/multislice_ptycho.py @@ -40,6 +40,7 @@ class MultislicePtycho(CDIModel): obj_view_crop=0, panel_plot_mode=False, plot_level=1, + translations=None, ): super(MultislicePtycho, self).__init__(panel_plot_mode=panel_plot_mode, @@ -113,7 +114,7 @@ class MultislicePtycho(CDIModel): shape = [s//oversampling for s in self.probe[0]] background = 1e-6 * t.ones(shape, dtype=t.float32) - self.background = t.nn.Parameter(background) + self.background = t.nn.Parameter(t.as_tensor(background, dtype=dtype)) if weights is None: self.weights = None @@ -182,6 +183,10 @@ class MultislicePtycho(CDIModel): else: raise KeyError('Specified loss function not supported') + if translations is not None: + self.register_buffer('original_translations', + t.as_tensor(translations, dtype=dtype)) + @classmethod def from_dataset(cls, @@ -420,6 +425,7 @@ class MultislicePtycho(CDIModel): obj_view_crop=obj_view_crop, panel_plot_mode=panel_plot_mode, plot_level=plot_level, + translations=translations, ) @@ -591,9 +597,18 @@ class MultislicePtycho(CDIModel): mask=mask) - def corrected_translations(self, dataset): - translations = dataset.translations.to( - dtype=t.float32, device=self.probe.device) + def corrected_translations(self, dataset=None): + if dataset is not None: + translations = dataset.translations.to( + dtype=t.float32, device=self.probe.device) + elif (hasattr(self, 'original_translations') and + self.original_translations is not None): + translations = self.original_translations.to( + dtype=t.float32, device=self.probe.device) + else: + raise ValueError( + 'Must provide a dataset or have original_translations stored ' + 'internally (via from_dataset or from_results_dict).') if (hasattr(self, 'translation_offsets') and self.translation_offsets is not None): t_offset = tools.interactions.pixel_to_translations( @@ -711,7 +726,7 @@ class MultislicePtycho(CDIModel): self.weights.data = S[:,:,None] * (Vh / probe_sqrt_intensities) - def plot_wavefront_variation(self, dataset, fig=None, mode='amplitude', **kwargs): + def plot_wavefront_variation(self, dataset=None, fig=None, mode='amplitude', **kwargs): def get_probes(idx): basis_prs = self.probe * self.probe_support[..., :, :] prs = t.sum(self.weights[idx, :, :, None, None] * basis_prs, @@ -761,24 +776,21 @@ class MultislicePtycho(CDIModel): plot_list = [ {'title': '', - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='root_sum_intensity', image_title='Root Summed Probe Intensities', image_colorbar_title='Square Root of Intensity'), 'condition': lambda self: len(self.weights.shape) >= 2}, {'title': '', - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='amplitude', image_title='Probe Amplitudes (scroll to view modes)', image_colorbar_title='Probe Amplitude'), 'condition': lambda self: len(self.weights.shape) >= 2}, {'title': '', - 'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation( - dataset, + 'plot_func': lambda self, fig: self.plot_wavefront_variation( fig=fig, mode='phase', image_title='Probe Phases (scroll to view modes)', @@ -814,8 +826,8 @@ class MultislicePtycho(CDIModel): fig=fig), 'condition': lambda self: len(self.weights.shape) >= 2}, {'title': '% of Power in Top Mode', - 'plot_func': lambda self, fig, dataset: p.plot_nanomap( - self.corrected_translations(dataset), + 'plot_func': lambda self, fig: p.plot_nanomap( + self.corrected_translations(), 100 * t.stack([ analysis.calc_mode_power_fractions( self.probe.data, @@ -884,13 +896,51 @@ class MultislicePtycho(CDIModel): cmap='cividis'), 'condition': lambda self: self.exponentiate_obj}, {'title': 'Corrected Translations', - 'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)}, + 'plot_func': lambda self, fig: p.plot_translations(self.corrected_translations(), fig=fig, units=self.units)}, {'title': 'Background', 'plot_func': lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)}, ] - def save_results(self, dataset): + @classmethod + def from_results_dict(cls, results_dict, obj_view_crop=0, units='um'): + sd = results_dict['state_dict'] + translation_offsets = sd.get('translation_offsets') + model = cls( + wavelength=sd['wavelength'], + detector_geometry={ + 'basis': sd['det_basis'], + 'distance': sd.get('det_distance'), + 'corner': sd.get('det_corner'), + }, + obj_basis=sd['obj_basis'], + probe_guess=sd['probe'], + obj_guess=sd['obj'], + interslice_propagator=sd['interslice_propagator'], + surface_normal=sd.get('surface_normal', np.array([0., 0., 1.])), + min_translation=sd.get('min_translation', np.array([0., 0.])), + background=sd['background'], + probe_basis=sd.get('probe_basis'), + translation_offsets=translation_offsets, + mask=sd.get('mask'), + weights=sd.get('weights'), + translation_scale=float(sd.get('translation_scale', 1.0)), + saturation=sd.get('saturation'), + oversampling=int(sd.get('oversampling', 1)), + fourier_probe=bool(sd.get('fourier_probe', False)), + simulate_probe_translation=bool(sd.get('simulate_probe_translation', False)), + simulate_finite_pixels=bool(sd.get('simulate_finite_pixels', False)), + exponentiate_obj=bool(sd.get('exponentiate_obj', False)), + loss=results_dict.get('loss_function', 'amplitude mse'), + obj_view_crop=obj_view_crop, + units=units, + translations=sd.get('original_translations'), + ) + model._load_results_dict(results_dict) + return model + + + def save_results(self, dataset=None): # This will save out everything needed to recreate the object # in the same state, but it's not the best formatted. For example, # "background" stores the square root of the background, etc. @@ -899,8 +949,11 @@ class MultislicePtycho(CDIModel): # We also save out the main results in a more readable format obj_basis = self.obj_basis.detach().cpu().numpy() probe_basis = self.probe_basis.detach().cpu().numpy() - translations=self.corrected_translations(dataset).detach().cpu().numpy() - original_translations = dataset.translations.detach().cpu().numpy() + translations = self.corrected_translations(dataset).detach().cpu().numpy() + if dataset is not None: + original_translations = dataset.translations.detach().cpu().numpy() + else: + original_translations = self.original_translations.detach().cpu().numpy() probe = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() obj = self.obj.detach().cpu().numpy() diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index 4bce45a..6e94103 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -111,10 +111,11 @@ class RPI(CDIModel): # We always use multi-modes to store the object, so we convert it # if we just get a single 2D array as an input + obj_guess = t.as_tensor(obj_guess, dtype=complex_dtype) if obj_guess.dim() == 2: obj_guess = obj_guess[None, :, :] - - self.obj = t.nn.Parameter(t.as_tensor(obj_guess, dtype=complex_dtype)) + + self.obj = t.nn.Parameter(obj_guess) self.weights = t.nn.Parameter( t.eye(probe.shape[0], dtype=complex_dtype)) @@ -632,6 +633,32 @@ class RPI(CDIModel): ] + @classmethod + def from_results_dict(cls, results_dict, units='um'): + sd = results_dict['state_dict'] + model = cls( + wavelength=sd['wavelength'], + detector_geometry={ + 'basis': sd['det_basis'], + 'distance': sd.get('det_distance'), + 'corner': sd.get('det_corner'), + }, + probe_basis=sd['probe_basis'], + probe=sd['probe'], + obj_guess=sd['obj'], + background=sd.get('background'), + mask=sd.get('mask'), + saturation=sd.get('saturation'), + oversampling=int(sd.get('oversampling', 1)), + exponentiate_obj=bool(sd.get('exponentiate_obj', False)), + phase_only=bool(sd.get('phase_only', False)), + loss=results_dict.get('loss_function', 'amplitude mse'), + units=units, + ) + model._load_results_dict(results_dict) + return model + + def save_results(self, dataset=None): # dataset is set as a kwarg here because it isn't needed, but the # common pattern is to pass a dataset. This makes it okay if one diff --git a/src/cdtools/models/simple_ptycho.py b/src/cdtools/models/simple_ptycho.py index eb99802..596f70b 100644 --- a/src/cdtools/models/simple_ptycho.py +++ b/src/cdtools/models/simple_ptycho.py @@ -133,9 +133,9 @@ class SimplePtycho(CDIModel): }, ] - def save_results(self, dataset): + def save_results(self, dataset=None): # This will save out everything needed to recreate the object - # in the same state, but it's not the best formatted. + # in the same state, but it's not the best formatted. base_results = super().save_results() # So we also save out the main results in a more useable format @@ -153,3 +153,32 @@ class SimplePtycho(CDIModel): } return {**base_results, **results} + + + @classmethod + def from_results_dict(cls, results_dict): + """Reconstructs a SimplePtycho model from a results dictionary. + + Parameters + ---------- + results_dict : dict + The dictionary returned by save_results(), as loaded from an h5 file + or produced directly in memory. + + Returns + ------- + model : SimplePtycho + A fully reconstructed model with all parameters, buffers, and + training metadata restored. + """ + sd = results_dict['state_dict'] + + model = cls( + wavelength=sd['wavelength'], + probe_basis=sd['probe_basis'], + probe_guess=sd['probe'], # normalized; probe_norm restored by _load_results_dict + obj_guess=sd['obj'], + min_translation=sd['min_translation'], + ) + model._load_results_dict(results_dict) + return model diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py index 23e2334..25e5e0c 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -170,3 +170,99 @@ def test_near_field_ptycho(near_field_ptycho_cxi, reconstruction_device, show_pl # If this fails, the reconstruction has gotten worse assert model.loss_history[-1] < 18 + + +def test_fancy_ptycho_from_results_dict(lab_ptycho_cxi, tmp_path): + dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi) + + t.manual_seed(42) + model = cdtools.models.FancyPtycho.from_dataset( + dataset, + n_modes=2, + ) + + # Verify original_translations is stored after from_dataset + assert hasattr(model, 'original_translations') + assert model.original_translations is not None + + # Run a few epochs to get non-trivial state + for loss in model.Adam_optimize(5, dataset, batch_size=10): + pass + + # Test from_results_dict with in-memory dict (no dataset argument needed) + results_dict = model.save_results() + loaded_model = cdtools.models.FancyPtycho.from_results_dict(results_dict) + + # Test from_results_h5 via a temporary file + h5_path = str(tmp_path / 'fancy_ptycho_test.h5') + model.save_to_h5(h5_path) + loaded_model_h5 = cdtools.models.FancyPtycho.from_results_h5(h5_path) + + # Verify training metadata is restored + assert loaded_model.epoch == model.epoch + assert loaded_model.loss_history == model.loss_history + + # Verify original_translations round-trips correctly + assert t.allclose( + loaded_model.original_translations, + model.original_translations, + ) + + # Verify all parameters and buffers are restored exactly + original_sd = model.state_dict() + loaded_sd = loaded_model.state_dict() + loaded_h5_sd = loaded_model_h5.state_dict() + for key in original_sd: + assert t.allclose(original_sd[key].float(), loaded_sd[key].float()), \ + f'from_results_dict: state_dict mismatch for key {key}' + assert t.allclose(original_sd[key].float(), loaded_h5_sd[key].float()), \ + f'from_results_h5: state_dict mismatch for key {key}' + + # Verify forward pass produces identical output + (indices, translations), patterns = dataset[:5] + with t.no_grad(): + original_out = model(indices, translations) + loaded_out = loaded_model(indices, translations) + loaded_h5_out = loaded_model_h5(indices, translations) + + assert t.allclose(original_out, loaded_out), \ + 'from_results_dict: forward pass output mismatch' + assert t.allclose(original_out, loaded_h5_out), \ + 'from_results_h5: forward pass output mismatch' + + +def test_fancy_ptycho_from_results_dict_with_missing_keys(lab_ptycho_cxi): + dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi) + + t.manual_seed(42) + model = cdtools.models.FancyPtycho.from_dataset(dataset, n_modes=1) + results_dict = model.save_results() + + # Strip top-level training metadata + for key in ('loss_history', 'epoch', 'training_history'): + results_dict.pop(key, None) + + # Strip defaultable state_dict keys + sd_keys_to_strip = ( + 'exponentiate_obj', 'phase_only', 'near_field', 'fourier_probe', + 'simulate_probe_translation', 'simulate_finite_pixels', + 'translation_scale', 'oversampling', 'surface_normal', 'min_translation', + ) + for key in sd_keys_to_strip: + results_dict['state_dict'].pop(key, None) + + loaded = cdtools.models.FancyPtycho.from_results_dict(results_dict) + + assert loaded.loss_history == [] + assert loaded.epoch == 0 + assert loaded.training_history == '' + assert bool(loaded.exponentiate_obj) == False + assert bool(loaded.phase_only) == False + assert bool(loaded.near_field) == False + assert bool(loaded.fourier_probe) == False + assert bool(loaded.simulate_probe_translation) == False + assert bool(loaded.simulate_finite_pixels) == False + assert float(loaded.translation_scale) == 1.0 + assert int(loaded.oversampling) == 1 + assert t.allclose(loaded.surface_normal, t.tensor([0., 0., 1.])) + assert t.allclose(loaded.min_translation, t.tensor([0., 0.])) diff --git a/tests/models/test_simple_ptycho.py b/tests/models/test_simple_ptycho.py index abace7e..15f2f1c 100644 --- a/tests/models/test_simple_ptycho.py +++ b/tests/models/test_simple_ptycho.py @@ -9,6 +9,52 @@ import cdtools t.manual_seed(0) +def test_simple_ptycho_from_results_dict(lab_ptycho_cxi, tmp_path): + dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi) + + t.manual_seed(42) + model = cdtools.models.SimplePtycho.from_dataset(dataset) + + # Run a few epochs to get non-trivial state + for loss in model.Adam_optimize(5, dataset, batch_size=10): + pass + + # Test from_results_dict with in-memory dict (no dataset argument needed) + results_dict = model.save_results() + loaded_model = cdtools.models.SimplePtycho.from_results_dict(results_dict) + + # Test from_results_h5 via a temporary file + h5_path = str(tmp_path / 'simple_ptycho_test.h5') + model.save_to_h5(h5_path) + loaded_model_h5 = cdtools.models.SimplePtycho.from_results_h5(h5_path) + + # Verify training metadata is restored + assert loaded_model.epoch == model.epoch + assert loaded_model.loss_history == model.loss_history + + # Verify all parameters and buffers are restored exactly + original_sd = model.state_dict() + loaded_sd = loaded_model.state_dict() + loaded_h5_sd = loaded_model_h5.state_dict() + for key in original_sd: + assert t.allclose(original_sd[key].float(), loaded_sd[key].float()), \ + f'from_results_dict: state_dict mismatch for key {key}' + assert t.allclose(original_sd[key].float(), loaded_h5_sd[key].float()), \ + f'from_results_h5: state_dict mismatch for key {key}' + + # Verify forward pass produces identical output + (indices, translations), patterns = dataset[:5] + with t.no_grad(): + original_out = model(indices, translations) + loaded_out = loaded_model(indices, translations) + loaded_h5_out = loaded_model_h5(indices, translations) + + assert t.allclose(original_out, loaded_out), \ + 'from_results_dict: forward pass output mismatch' + assert t.allclose(original_out, loaded_h5_out), \ + 'from_results_h5: forward pass output mismatch' + + @pytest.mark.slow def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot): dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)