Add reusable echo embedding replacement manifest path

scripts/build_echo_hubert_manifest.py

@@ -96,6 +96,15 @@ def parse_args() -> argparse.Namespace:
     ]:
         join.add_argument(f"--{name}")
 
+    replace_echo = subparsers.add_parser(
+        "replace-echo", help="Replace echo embeddings in an existing joined manifest."
+    )
+    replace_echo.add_argument("--base-manifest", type=Path, required=True)
+    replace_echo.add_argument("--echo", type=Path, required=True)
+    replace_echo.add_argument("--manifest", type=Path, required=True)
+    replace_echo.add_argument("--metadata-csv", type=Path, required=True)
+    replace_echo.add_argument("--summary-json", type=Path, required=True)
+
     all_steps = subparsers.add_parser("all", help="Run all build steps.")
     all_steps.add_argument("--echo-input", type=Path, required=True)
     all_steps.add_argument("--hubert-csv", type=Path, required=True)
@@ -154,6 +163,7 @@ def main() -> None:
         build_echo_study_embeddings,
         build_joined_manifest,
         convert_hubert_csv_to_parquet,
+        replace_manifest_echo_embeddings,
     )
 
     if args.command == "build-echo":
@@ -191,6 +201,15 @@ def main() -> None:
             **cohort_kwargs(args),
         )
         print_summary(summary)
+    elif args.command == "replace-echo":
+        _, summary = replace_manifest_echo_embeddings(
+            base_manifest_path=args.base_manifest,
+            echo_embeddings_path=args.echo,
+            manifest_path=args.manifest,
+            metadata_csv_path=args.metadata_csv,
+            summary_json_path=args.summary_json,
+        )
+        print_summary(summary)
     elif args.command == "all":
         build_echo_study_embeddings(args.echo_input, args.echo_output, max_clips=args.max_clips)
         convert_hubert_csv_to_parquet(args.hubert_csv, args.ecg_output, chunksize=args.chunksize)

src/primed_ai/data/echo_hubert_manifest.py

@@ -494,6 +494,89 @@ def build_joined_manifest(
     return manifest, summary
 
 
+def replace_manifest_echo_embeddings(
+    *,
+    base_manifest_path: Path,
+    echo_embeddings_path: Path,
+    manifest_path: Path,
+    metadata_csv_path: Path,
+    summary_json_path: Path,
+) -> tuple[pd.DataFrame, dict[str, Any]]:
+    """Replace EchoJEPA vectors while retaining a joined manifest's cohort and ECG data.
+
+    This supports a controlled pooling ablation: labels, subject-level splits, and ECG
+    vectors are held fixed while only the echo representation changes.  It also avoids
+    re-reading the much larger source HuBERT parquet when a valid joined manifest exists.
+    """
+    base = pd.read_parquet(base_manifest_path)
+    echo = pd.read_parquet(echo_embeddings_path)
+
+    base_required = {"subject_id", "echo_study_id", "ecg_embedding"}
+    echo_required = {
+        "subject_id",
+        "echo_study_id",
+        "n_echo_clips",
+        "echo_embedding",
+        "echo_model",
+    }
+    missing_base = sorted(base_required - set(base.columns))
+    missing_echo = sorted(echo_required - set(echo.columns))
+    if missing_base:
+        raise ValueError(f"Base manifest missing required columns: {missing_base}")
+    if missing_echo:
+        raise ValueError(f"Echo embeddings missing required columns: {missing_echo}")
+
+    echo_columns = [
+        "subject_id",
+        "echo_study_id",
+        "n_echo_clips",
+        "echo_embedding",
+        "echo_model",
+    ]
+    if "n_echo_clips_retained" in echo.columns:
+        echo_columns.append("n_echo_clips_retained")
+
+    replacement_columns = [
+        "n_echo_clips",
+        "n_echo_clips_retained",
+        "echo_embedding",
+        "echo_model",
+        "has_echo_embedding",
+    ]
+    base = base.drop(columns=[c for c in replacement_columns if c in base.columns])
+    manifest = base.merge(
+        echo[echo_columns],
+        on=["subject_id", "echo_study_id"],
+        how="left",
+        validate="many_to_one",
+    )
+    manifest["has_echo_embedding"] = manifest["echo_embedding"].notna()
+    manifest["has_ecg_embedding"] = manifest["ecg_embedding"].notna()
+
+    manifest_path.parent.mkdir(parents=True, exist_ok=True)
+    metadata_csv_path.parent.mkdir(parents=True, exist_ok=True)
+    summary_json_path.parent.mkdir(parents=True, exist_ok=True)
+    manifest.to_parquet(manifest_path, index=False)
+
+    metadata_columns = [
+        "subject_id", "echo_study_id", "ecg_study_id", "lvef", "ef_le_40", "split",
+        "sex", "age", "race", "n_echo_clips", "n_echo_clips_retained",
+        "has_echo_embedding", "has_ecg_embedding", "echo_model", "ecg_model",
+    ]
+    manifest[[c for c in metadata_columns if c in manifest]].to_csv(metadata_csv_path, index=False)
+
+    summary = summarize_manifest(manifest)
+    summary["outputs"] = {
+        "manifest_path": str(manifest_path),
+        "metadata_csv_path": str(metadata_csv_path),
+        "summary_json_path": str(summary_json_path),
+    }
+    summary["base_manifest_path"] = str(base_manifest_path)
+    summary["echo_embeddings_path"] = str(echo_embeddings_path)
+    summary_json_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
+    return manifest, summary
+
+
 def summarize_manifest(manifest: pd.DataFrame) -> dict[str, Any]:
     split_counts = manifest["split"].value_counts(dropna=False).to_dict()
     leakage = (

tests/test_echo_hubert.py

@@ -2,6 +2,7 @@
 
 import json
 
+import numpy as np
 import pandas as pd
 import pytest
 
@@ -16,6 +17,7 @@
     build_echo_study_embeddings,
     build_joined_manifest,
     convert_hubert_csv_to_parquet,
+    replace_manifest_echo_embeddings,
 )
 
 
@@ -228,3 +230,61 @@ def test_subject_split_leakage_raises(tmp_path):
             tmp_path / "metadata.csv",
             tmp_path / "summary.json",
         )
+
+
+def test_replace_manifest_echo_embeddings_preserves_base_cohort_and_ecg(tmp_path):
+    base_path = tmp_path / "base.parquet"
+    echo_path = tmp_path / "clip_echo.parquet"
+    manifest_path = tmp_path / "clip_manifest.parquet"
+    metadata_path = tmp_path / "clip_metadata.csv"
+    summary_path = tmp_path / "clip_summary.json"
+
+    pd.DataFrame(
+        {
+            "subject_id": [1, 2],
+            "echo_study_id": [10, 20],
+            "ecg_study_id": [100, 200],
+            "lvef": [35.0, 60.0],
+            "ef_le_40": [True, False],
+            "split": ["train", "val"],
+            "n_echo_clips": [5, 4],
+            "echo_embedding": [[1.0, 2.0], [3.0, 4.0]],
+            "echo_model": ["pooled", "pooled"],
+            "ecg_embedding": [[0.1, 0.2], [0.3, 0.4]],
+            "ecg_model": ["hubert", "hubert"],
+        }
+    ).to_parquet(base_path, index=False)
+    pd.DataFrame(
+        {
+            "subject_id": [1, 2],
+            "echo_study_id": [10, 20],
+            "n_echo_clips": [5, 4],
+            "n_echo_clips_retained": [2, 2],
+            "echo_embedding": [
+                [[1.0, 2.0], [3.0, 4.0]],
+                [[5.0, 6.0], [7.0, 8.0]],
+            ],
+            "echo_model": ["clip", "clip"],
+        }
+    ).to_parquet(echo_path, index=False)
+
+    manifest, summary = replace_manifest_echo_embeddings(
+        base_manifest_path=base_path,
+        echo_embeddings_path=echo_path,
+        manifest_path=manifest_path,
+        metadata_csv_path=metadata_path,
+        summary_json_path=summary_path,
+    )
+
+    assert np.asarray(manifest.loc[0, "ecg_embedding"]).tolist() == [0.1, 0.2]
+    assert np.asarray(manifest.loc[1, "ecg_embedding"]).tolist() == [0.3, 0.4]
+    assert np.allclose(
+        np.stack(manifest.loc[0, "echo_embedding"]),
+        [[1.0, 2.0], [3.0, 4.0]],
+    )
+    assert manifest["n_echo_clips_retained"].tolist() == [2, 2]
+    assert manifest["has_echo_embedding"].tolist() == [True, True]
+    assert summary["n_with_both_embeddings"] == 2
+    assert summary["base_manifest_path"] == str(base_path)
+    assert "echo_embedding" not in pd.read_csv(metadata_path).columns
+    assert json.loads(summary_path.read_text())["echo_embeddings_path"] == str(echo_path)