From 7bf9cedfa9d47458ab5d60ce4e2632842eccf535 Mon Sep 17 00:00:00 2001 From: Kevin Zhou <86838731+kevzho@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:36:35 -0400 Subject: [PATCH] Add reusable echo embedding replacement manifest path --- scripts/build_echo_hubert_manifest.py | 19 +++++ src/primed_ai/data/echo_hubert_manifest.py | 83 ++++++++++++++++++++++ tests/test_echo_hubert.py | 60 ++++++++++++++++ 3 files changed, 162 insertions(+) diff --git a/scripts/build_echo_hubert_manifest.py b/scripts/build_echo_hubert_manifest.py index 68eb6dd..e263984 100644 --- a/scripts/build_echo_hubert_manifest.py +++ b/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) diff --git a/src/primed_ai/data/echo_hubert_manifest.py b/src/primed_ai/data/echo_hubert_manifest.py index eb5c7cf..2e1b495 100644 --- a/src/primed_ai/data/echo_hubert_manifest.py +++ b/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 = ( diff --git a/tests/test_echo_hubert.py b/tests/test_echo_hubert.py index ef2b1c2..cf64fe4 100644 --- a/tests/test_echo_hubert.py +++ b/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)