Make the loading code for fancyptycho a bit more robust

examples/gold_ball_ptycho.py

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

src/cdtools/models/base.py

@@ -229,10 +229,10 @@ class CDIModel(t.nn.Module):
             or produced directly in memory.
         """
         state_dict = nested_dict_to_torch(results_dict['state_dict'])
-        self.load_state_dict(state_dict)
-        self.loss_history = list(results_dict['loss_history'])
-        self.epoch = int(results_dict['epoch'])
-        self.training_history = str(results_dict['training_history'])
+        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

src/cdtools/models/fancy_ptycho.py

@@ -1252,7 +1252,6 @@ class FancyPtycho(CDIModel):
             A fully reconstructed model with all parameters, buffers, and
             training metadata restored.
         """
-        import numpy as np
         sd = results_dict['state_dict']
 
         # For optional Parameters (translation_offsets, weights, qe_mask, etc.),
@@ -1271,8 +1270,8 @@ class FancyPtycho(CDIModel):
             obj_basis=sd['obj_basis'],
             probe_guess=sd['probe'],   # normalized; probe_norm restored by _load_results_dict
             obj_guess=sd['obj'],
-            surface_normal=sd['surface_normal'],
-            min_translation=sd['min_translation'],
+            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
@@ -1281,15 +1280,15 @@ class FancyPtycho(CDIModel):
             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['translation_scale']),
-            oversampling=int(sd['oversampling']),
-            fourier_probe=bool(sd['fourier_probe']),
+            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['simulate_probe_translation']),
-            simulate_finite_pixels=bool(sd['simulate_finite_pixels']),
-            exponentiate_obj=bool(sd['exponentiate_obj']),
-            phase_only=bool(sd['phase_only']),
-            near_field=bool(sd['near_field']),
+            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'),

tests/models/test_fancy_ptycho.py

@@ -229,3 +229,40 @@ def test_fancy_ptycho_from_results_dict(lab_ptycho_cxi, tmp_path):
         '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.]))