Merge pull request #77 from cdtools-developers/savenload

Improve saving and loading of in-progress reconstructions

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()
 ```

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
 -----------------
 

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()
 

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()
     

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()

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);"
    ]
   },

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);"
    ]
   },

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()

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')

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()

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)

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()

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()

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

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()

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):

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()

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

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()

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

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

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.]))

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)