From 7d46ef62a8b2ec66ff2b27db42ee613052dbbd13 Mon Sep 17 00:00:00 2001
From: MarcelMB <marcel@solaris>
Date: Tue, 10 Mar 2026 15:19:11 -0700
Subject: [PATCH] Fix bundle loading: handle tuple YAML, add behavior_fps to
 BehaviorConfig

---
 placecell/config.py                     | 20 ++++++++++++++++-
 placecell/config/pcell_config_WL27.yaml | 29 +++++++++++++++++++++++++
 placecell/dataset/base.py               | 13 +++++++++--
 3 files changed, 59 insertions(+), 3 deletions(-)
 create mode 100644 placecell/config/pcell_config_WL27.yaml

diff --git a/placecell/config.py b/placecell/config.py
index d6ff448..eddfcb9 100644
--- a/placecell/config.py
+++ b/placecell/config.py
@@ -260,6 +260,11 @@ class BehaviorConfig(BaseModel):
         ...,
         description="Analysis type: 'arena' for 2D open-field, 'maze' for 1D arm analysis.",
     )
+    behavior_fps: float = Field(
+        20.0,
+        gt=0.0,
+        description="Behavior camera sampling rate in frames per second.",
+    )
     speed_threshold: float = Field(
         25.0,
         description="Minimum running speed to keep events (mm/s).",
@@ -500,10 +505,23 @@ class AnalysisConfig(BaseModel):
     @classmethod
     def from_yaml(cls, path: str | Path) -> "AnalysisConfig":
         """Load from a YAML file."""
+        loader = yaml.SafeLoader
+        # Old bundles were saved with !!python/tuple tags; handle them safely.
+        loader.add_constructor(
+            "tag:yaml.org,2002:python/tuple",
+            lambda l, n: tuple(l.construct_sequence(n)),
+        )
         with open(path) as f:
-            data = yaml.safe_load(f)
+            data = yaml.load(f, Loader=loader)
         return cls(**data)
 
+    def to_yaml(self, path: str | Path) -> None:
+        """Save to a YAML file."""
+        import json
+        data = json.loads(self.model_dump_json())
+        with open(path, "w") as f:
+            yaml.dump(data, f, default_flow_style=False)
+
     def with_data_overrides(self, data_cfg: BaseDataConfig) -> "AnalysisConfig":
         """Create a new config with data-specific overrides applied.
 
diff --git a/placecell/config/pcell_config_WL27.yaml b/placecell/config/pcell_config_WL27.yaml
new file mode 100644
index 0000000..8b04fd3
--- /dev/null
+++ b/placecell/config/pcell_config_WL27.yaml
@@ -0,0 +1,29 @@
+neural:
+  fps: 20.0
+  oasis:
+    g: [1.60, -0.63]  # AR(2) coefficients (required, usually overridden by data config)
+    baseline: p10
+    penalty: 0.8  # Sparsity penalty (higher = fewer events). Default 0.
+    s_min: 0  # Minimum event size threshold. Default 0.
+  trace_name: C_lp
+
+behavior:
+  type: arena
+  speed_threshold: 10.0  # mm/s
+  speed_window_frames: 5
+  jump_threshold_mm: 100  # Max plausible frame-to-frame displacement (mm). Larger = tracking error.
+  spatial_map_2d:
+    bins: 50
+    min_occupancy: 0.025 # Minimum occupancy (in seconds) to include a bin in spatial map
+    spatial_sigma: 3  # Gaussian smoothing (in bins) for occupancy and rate maps
+    n_shuffles: 1000
+    random_seed: 1
+    event_threshold_sigma: 0  # Sigma multiplier for event amplitude threshold in trajectory visualization
+    p_value_threshold: 0.05 # P-value threshold. Units with p < threshold pass.
+    min_shift_seconds: 20  # Minimum circular shift (seconds) for shuffle test. 0 = no minimum.
+    si_weight_mode: amplitude  # SI weight mode: 'amplitude' (event amplitudes) or 'binary' (event counts)
+    place_field_threshold: 0.35  # Fraction of peak rate for place field boundary (red contour). Applied to smoothed rate map.
+    place_field_min_bins: 5  # Minimum contiguous bins for a place field component. Smaller disconnected regions are discarded.
+    place_field_seed_percentile: 95  # Percentile of shuffled rates for seed detection (Guo et al. 2023).
+    n_split_blocks: 10  # Number of temporal blocks for interleaved stability splitting.
+    block_shifts: [0]  # Block boundary shifts (fractions of block width). [0, 0.5] for two arrangements.
diff --git a/placecell/dataset/base.py b/placecell/dataset/base.py
index 1617d7c..ddc8a1e 100644
--- a/placecell/dataset/base.py
+++ b/placecell/dataset/base.py
@@ -449,8 +449,17 @@ def save_bundle(self, path: str | Path, *, save_figures: bool = True) -> Path:
         }
         (path / "metadata.json").write_text(json.dumps(meta, indent=2))
 
-        # Config
-        self.cfg.to_yaml(path / "config.yaml")
+        # Config — propagate behavior_fps from data_cfg so bundles are self-contained
+        cfg_to_save = self.cfg
+        if self.data_cfg is not None and hasattr(self.data_cfg, "behavior_fps"):
+            cfg_to_save = self.cfg.model_copy(
+                update={
+                    "behavior": self.cfg.behavior.model_copy(
+                        update={"behavior_fps": self.data_cfg.behavior_fps}
+                    )
+                }
+            )
+        cfg_to_save.to_yaml(path / "config.yaml")
 
         # Spatial arrays
         spatial_kw: dict[str, np.ndarray] = {}