diff --git a/README.md b/README.md index 8381ff10..2796c5bd 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 4c101c90..6f356cd0 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 22bc2c1f..78d9f413 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 58f3464b..b13d72cc 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 2e9db133..2420e568 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 d6dfc748..14d75f8b 100644 --- a/examples/fancy_ptycho.py +++ b/examples/fancy_ptycho.py @@ -34,7 +34,7 @@ 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 @@ # 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 f5cab0fe..b0dc8764 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 0e9adad6..8d5343f5 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 aeb510c1..ff0c54f9 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -50,15 +50,15 @@ # 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 @@ # 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 fde19c6a..b73b508e 100644 --- a/examples/gold_ball_split.py +++ b/examples/gold_ball_split.py @@ -54,4 +54,4 @@ 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 47da91de..55a90866 100644 --- a/examples/near_field_ptycho.py +++ b/examples/near_field_ptycho.py @@ -35,21 +35,21 @@ 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 41c0c6c2..dfd78c10 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -23,14 +23,14 @@ 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 32c02d55..c972db22 100644 --- a/examples/transmission_RPI.py +++ b/examples/transmission_RPI.py @@ -31,19 +31,19 @@ # 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 4e95ec63..eed45148 100644 --- a/examples/tutorial_finale.py +++ b/examples/tutorial_finale.py @@ -17,9 +17,9 @@ 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 e5ff284f..685e6824 100644 --- a/examples/tutorial_simple_ptycho.py +++ b/examples/tutorial_simple_ptycho.py @@ -132,7 +132,7 @@ def measurement(self, wavefields): }, ] - 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 @@ def save_results(self, dataset): } 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 00000000..86a0c7b6 --- /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 5b990a04..4761b6ed 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -36,7 +36,8 @@ 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 @@ def save_results(self): 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 @@ def save_to_h5(self, filename, *args): 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 43f6360d..04f43697 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -83,6 +83,7 @@ def __init__( obj_view_crop=0, panel_plot_mode=False, plot_level=1, + translations=None, ): # We need the detector geometry @@ -176,7 +177,7 @@ def __init__( 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 @@ def __init__( 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 @@ def from_dataset( units=units, panel_plot_mode=panel_plot_mode, plot_level=plot_level, + translations=translations, ) @@ -590,11 +596,25 @@ def sim_to_dataset(self, args_list): 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 @@ def corrected_translations(self,dataset): 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 @@ def save_results(self, dataset): # 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 ea1ddd7b..50def1ff 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -46,6 +46,7 @@ def __init__(self, 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 @@ def __init__(self, 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 @@ def __init__(self, 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 @@ def __init__(self, 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, @@ -398,6 +404,7 @@ def from_dataset(cls, dataset.patterns.shape[-2:], dtype=t.float32) if probe_fourier_crop is not None: + initial_probe_shape = np.array(probe.shape) probe = tools.propagators.far_field(probe) probe = probe[probe_fourier_crop : probe.shape[-2] - probe_fourier_crop, @@ -405,7 +412,7 @@ def from_dataset(cls, - probe_fourier_crop] probe = tools.propagators.inverse_far_field(probe) - scale_factor = np.array(det_shape) / np.array(probe.shape) + scale_factor = initial_probe_shape / np.array(probe.shape) probe_basis = obj_basis * scale_factor[None,:] else: probe_basis = obj_basis.clone() @@ -530,6 +537,7 @@ def from_dataset(cls, inv_angular_spectrum_propagator=inv_angular_spectrum_propagator, panel_plot_mode=panel_plot_mode, plot_level=plot_level, + translations=translations, ) @@ -717,9 +725,18 @@ def sim_to_dataset(self, args_list, calculation_width=None): 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( @@ -899,7 +916,7 @@ def get_probe_intensities(self): 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, @@ -933,7 +950,7 @@ def get_probes(idx): **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), @@ -948,10 +965,14 @@ def plot_illumination_intensity(self, fig, dataset): 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', @@ -1085,7 +1106,7 @@ def plot_translations_and_originals(self, fig, dataset): { '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', @@ -1095,7 +1116,7 @@ def plot_translations_and_originals(self, fig, dataset): { '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', @@ -1114,8 +1135,8 @@ def plot_translations_and_originals(self, fig, dataset): { '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, @@ -1146,8 +1167,7 @@ def plot_translations_and_originals(self, fig, dataset): {'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', @@ -1156,8 +1176,7 @@ def plot_translations_and_originals(self, fig, dataset): {'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)', @@ -1166,8 +1185,7 @@ def plot_translations_and_originals(self, fig, dataset): {'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)', @@ -1176,7 +1194,7 @@ def plot_translations_and_originals(self, fig, dataset): ] - 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. @@ -1185,8 +1203,11 @@ def save_results(self, dataset): # 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() @@ -1209,3 +1230,71 @@ def save_results(self, dataset): } 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 03b1b520..35473511 100644 --- a/src/cdtools/models/multislice_ptycho.py +++ b/src/cdtools/models/multislice_ptycho.py @@ -40,6 +40,7 @@ def __init__(self, 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 @@ def __init__(self, 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 @@ def __init__(self, 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 @@ def from_dataset(cls, obj_view_crop=obj_view_crop, panel_plot_mode=panel_plot_mode, plot_level=plot_level, + translations=translations, ) @@ -591,9 +597,18 @@ def sim_to_dataset(self, args_list, calculation_width=None): 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 @@ def tidy_probes(self): 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 @@ def get_probes(idx): 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 @@ def get_probes(idx): 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 @@ def get_probes(idx): 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 @@ def save_results(self, dataset): # 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 4bce45ae..6e94103b 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -111,10 +111,11 @@ def __init__( # 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 @@ def sim_to_dataset(self, args_list): ] + @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 eb998026..596f70b1 100644 --- a/src/cdtools/models/simple_ptycho.py +++ b/src/cdtools/models/simple_ptycho.py @@ -133,9 +133,9 @@ def measurement(self, wavefields): }, ] - 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 @@ def save_results(self, dataset): } 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 23e2334e..25e5e0cd 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 abace7e9..15f2f1cb 100644 --- a/tests/models/test_simple_ptycho.py +++ b/tests/models/test_simple_ptycho.py @@ -9,6 +9,52 @@ 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)