From 83b1ed5d27a31c0fa1a9ab19ccf052500ef89d60 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Mon, 13 Apr 2026 17:29:10 +0200
Subject: [PATCH] Add a script to demonstrate loading from a saved
 reconstruction and add the obj_view_crop and units arguments back

---
 examples/view_gold_ball_ptycho.py  | 11 +++++++++++
 src/cdtools/models/base.py         |  5 +++--
 src/cdtools/models/fancy_ptycho.py | 10 +++++++++-
 3 files changed, 23 insertions(+), 3 deletions(-)
 create mode 100644 examples/view_gold_ball_ptycho.py

diff --git a/examples/view_gold_ball_ptycho.py b/examples/view_gold_ball_ptycho.py
new file mode 100644
index 0000000..86a0c7b
--- /dev/null
+++ b/examples/view_gold_ball_ptycho.py
@@ -0,0 +1,11 @@
+import cdtools
+from matplotlib import pyplot as plt
+
+model = cdtools.models.FancyPtycho.from_results_h5(
+    'example_reconstructions/gold_balls.h5',
+    obj_view_crop=260, # How far in to crop from the edge
+    units='um', # The units to display in
+)
+
+model.inspect()
+plt.show()
diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py
index 7cd0411..6527920 100644
--- a/src/cdtools/models/base.py
+++ b/src/cdtools/models/base.py
@@ -255,7 +255,7 @@ class CDIModel(t.nn.Module):
 
 
     @classmethod
-    def from_results_h5(cls, filename):
+    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.
@@ -272,7 +272,8 @@ class CDIModel(t.nn.Module):
             A fully reconstructed model with all parameters, buffers, and
             training metadata restored.
         """
-        return cls.from_results_dict(h5_to_nested_dict(filename))
+        return cls.from_results_dict(
+            h5_to_nested_dict(filename), *args, **kwargs)
 
 
     @contextmanager
diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py
index 325860c..6f53e13 100644
--- a/src/cdtools/models/fancy_ptycho.py
+++ b/src/cdtools/models/fancy_ptycho.py
@@ -149,6 +149,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:
@@ -1224,7 +1225,12 @@ class FancyPtycho(CDIModel):
 
 
     @classmethod
-    def from_results_dict(cls, results_dict):
+    def from_results_dict(
+        cls,
+        results_dict,
+        obj_view_crop=0,
+        units='um',    
+    ):
         """Reconstructs a FancyPtycho model from a results dictionary.
 
         Parameters
@@ -1280,6 +1286,8 @@ class FancyPtycho(CDIModel):
             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