Sakthigsjhy/MagBridge-Battery

1"""2MagBridge-Battery v1.0 — minimal loader example.3 4Run from the bundle root:5    python load_example.py6 7Requires: pandas, pyarrow.8Licensed under Apache-2.0 (see LICENSE-CODE).9"""10from __future__ import annotations11 12import json13from pathlib import Path14 15import numpy as np16import pandas as pd17 18 19def load_bundle(base: Path) -> tuple[pd.DataFrame, dict, dict, dict]:20    """Load all shards, both split files, and the manifest.21 22    Returns23    -------24    df : pd.DataFrame25        Concatenated shards. One row per sample. Signal columns hold length-10026        arrays.27    primary_split : dict28        Cell-disjoint, leakage-free split. Keys include 'train_samples',29        'val_samples', 'test_samples', and 'split_guarantee'.30    optimistic_split : dict31        Intentionally leaky baseline. Do not use for reporting; see its32        'warning' field.33    manifest : dict34        Provenance, hashes, bridge config.35    """36    shards = sorted((base / "data").glob("shard_*.parquet"))37    if not shards:38        raise FileNotFoundError(f"No shards found under {base / 'data'}")39    df = pd.concat([pd.read_parquet(s) for s in shards], ignore_index=True)40 41    primary_split = json.loads((base / "splits" / "by_cell_primary.json").read_text())42    optimistic_split = json.loads((base / "splits" / "by_record_optimistic_baseline.json").read_text())43    manifest = json.loads((base / "manifest.json").read_text())44 45    return df, primary_split, optimistic_split, manifest46 47 48def apply_split(df: pd.DataFrame, split: dict) -> dict[str, pd.DataFrame]:49    """Slice df into train/val/test using a split dict."""50    by_id = df.set_index("sample_id", drop=False)51    out = {}52    for subset in ("train", "val", "test"):53        ids = split[f"{subset}_samples"]54        out[subset] = by_id.loc[by_id.index.intersection(ids)].reset_index(drop=True)55    return out56 57 58def stack_signals(df: pd.DataFrame, channels: list[str] | None = None) -> np.ndarray:59    """Stack signal columns into a (N, T, C) numpy array.60 61    Default channels: the six signal channels. ``time_norm`` is omitted because62    it is constant across samples (a fixed reference grid) and adds no63    per-sample information.64    """65    if channels is None:66        channels = ["B_s1Y", "B_s1Z", "B_s2Y", "B_s2Z", "B_s1C5", "B_s2C6"]67    arrays = [np.stack(df[c].values) for c in channels]  # each (N, T)68    return np.stack(arrays, axis=-1)  # (N, T, C)69 70 71def main() -> None:72    base = Path(__file__).parent.resolve()73 74    df, primary, optimistic, manifest = load_bundle(base)75 76    print(f"Dataset: {manifest['dataset_name']} v{manifest['dataset_version']}")77    print(f"Schema: {manifest['schema_version']}")78    print(f"Generated: {manifest['generated_at_utc']}")79    print(f"Total samples loaded: {len(df)}")80    print()81 82    splits = apply_split(df, primary)83    print("Primary (cell-disjoint) split:")84    for name, sub in splits.items():85        print(f"  {name:5s}: {len(sub):5d} samples")86    print()87    print("Split guarantee:")88    print(f"  {primary['split_guarantee'][:120]}...")89    print()90 91    # Show that the optimistic split is shipped with a clear warning92    print("Optimistic split warning (first 120 chars):")93    print(f"  {optimistic['warning'][:120]}...")94    print(f"  leakage_stats: {optimistic.get('leakage_stats', {})}")95    print()96 97    # Stack train signals into a tensor98    X_train = stack_signals(splits["train"])99    print(f"Train signal tensor shape: {X_train.shape}  (N, T, C)")100    print(f"  channels = ['B_s1Y','B_s1Z','B_s2Y','B_s2Z','B_s1C5','B_s2C6']")101 102    # Show that Regime-B has missing SOH by design103    regime_b = df[df["anomaly_subtype"] == "low_voltage_regime_B"]104    print()105    print(f"Regime-B samples: {len(regime_b)}; SOH missing: {regime_b['soh'].isna().sum()} (expected = all)")106 107 108if __name__ == "__main__":109    main()110