Factorized sampler: a composable Step-based inference spine

.gitignore

@@ -148,6 +148,9 @@ dmypy.json
 # .DS_Store files for Mac OS
 .DS_Store
 
+# Claude Code personal notes
+CLAUDE.local.md
+
 # directories with temporary files
 **/tmp_files/*
 

compatibility/remove_domain_window_factor.py

@@ -27,24 +27,32 @@ def main():
     if path.suffix == ".hdf5":
         with h5py.File(path, "r+") as f:
             settings = ast.literal_eval(f.attrs["settings"])
-            try:
-                if "domain" in settings:
-                    del settings["domain"]["window_factor"]
-                elif "domain_dict" in settings:
-                    del settings["domain_dict"]["window_factor"]
+            domain_settings = settings.get("domain", settings.get("domain_dict"))
+            if domain_settings is not None and _remove_window_factor(domain_settings):
                 f.attrs["settings"] = str(settings)
-            except KeyError:
+            else:
                 print("Dataset is already in correct format.")
 
     elif path.suffix == ".pt":
-        d, _ = torch_load_with_fallback(args.checkpoint)
+        d, _ = torch_load_with_fallback(path)
 
-        try:
-            del d["metadata"]["dataset_settings"]["domain"]["window_factor"]
+        domain_settings = d["metadata"]["dataset_settings"]["domain"]
+        if _remove_window_factor(domain_settings):
             torch.save(d, path)
-        except KeyError:
+        else:
             print("Dataset is already in correct format.")
 
 
+def _remove_window_factor(domain_settings: dict) -> bool:
+    """Delete `window_factor` from the domain settings, including a multibanded
+    domain's nested base domain. Returns whether anything was removed."""
+    removed = False
+    for settings in (domain_settings, domain_settings.get("base_domain", {})):
+        if "window_factor" in settings:
+            del settings["window_factor"]
+            removed = True
+    return removed
+
+
 if __name__ == "__main__":
     main()

dingo/core/dataset.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 from typing import List, Optional, Union
 import ast
 import h5py
@@ -93,7 +94,7 @@ class DingoDataset:
 
     def __init__(
         self,
-        file_name: Optional[str] = None,
+        file_name: Optional[Union[str, Path]] = None,
         dictionary: Optional[dict] = None,
         data_keys: Optional[List] = None,
         leave_on_disk_keys: Optional[list] = None,
@@ -104,7 +105,7 @@ def __init__(
 
         Parameters
         ----------
-        file_name : str
+        file_name : str or Path
             HDF5 file containing a dataset
         dictionary : dict
             Contains settings and data entries. The data keys should be the same as
@@ -135,7 +136,7 @@ def __init__(
         elif dictionary is not None:
             self.from_dictionary(dictionary)
 
-    def to_file(self, file_name: str, mode: str = "w"):
+    def to_file(self, file_name: Union[str, Path], mode: str = "w"):
         print("Saving dataset to " + str(file_name))
         save_dict = {
             k: v
@@ -149,7 +150,7 @@ def to_file(self, file_name: str, mode: str = "w"):
             if self.dataset_type:
                 f.attrs["dataset_type"] = self.dataset_type
 
-    def from_file(self, file_name: str):
+    def from_file(self, file_name: Union[str, Path]):
         print(f"Loading dataset from {str(file_name)}.")
         if self._leave_on_disk_keys:
             print(f"Omitting data keys {self._leave_on_disk_keys}.")