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,20 +34,20 @@
     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
 # loop, and now we polish the reconstruction with a lower learning rate
 # and larger minibatch
 for loss in recon.optimize(50, lr=0.005, batch_size=50):
     print(model.report())
-    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 @@
 # 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()

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

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

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)

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

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

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