Remove the need to use dataset for model.inspect() and model.save_results(), and update the examples accordingly. Docs and ipython notebooks not yet updated

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/gold_ball_ptycho.py

@@ -54,11 +54,11 @@ with model.save_on_exception(
     
     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

@@ -34,21 +34,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

@@ -30,7 +30,7 @@ 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,7 +17,7 @@ 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)

examples/tutorial_simple_ptycho.py

@@ -130,7 +130,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()

src/cdtools/models/fancy_ptycho.py

@@ -907,7 +907,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,
@@ -941,7 +941,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),
@@ -956,10 +956,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',
@@ -1093,7 +1097,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',
@@ -1103,7 +1107,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',
@@ -1122,8 +1126,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,
@@ -1154,8 +1158,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',
@@ -1164,8 +1167,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)',
@@ -1174,8 +1176,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)',