DynaCLR gut-cell training: per-cell Z centering, cell split, gut configs

applications/dynaclr/configs/collections/zuben_gut/conform_parquet.py

@@ -0,0 +1,64 @@
+"""Step 2: conform Zuben's cell-index parquet to the canonical schema (end-to-end).
+
+Single authoritative step, run after the v3 stores are preprocessed (Step 1):
+
+1. Repoint ``store_path`` from the v2 originals to the v3 copies (same basename).
+2. Fill required/derived columns absent from the source
+   (``tracks_path``, ``microscope``, ``T_shape``, ``C_shape``).
+3. ``preprocess_cell_index`` fills the ``norm_*`` columns from the v3 ``.zattrs``.
+4. Seed ``z_focus = z`` LAST — gut has no ``focus_slice`` zattrs, so the bbox-center
+   ``z`` IS the focus. This must come after step 3 because ``preprocess_cell_index``
+   also writes ``z_focus`` (as NaN here, since there is no focus_slice) and would
+   otherwise clobber it. The datamodule centers the Z window on ``z_focus``.
+
+Run::
+
+    uv run --no-sync python applications/dynaclr/configs/collections/zuben_gut/conform_parquet.py
+"""
+
+import os
+
+import pandas as pd
+
+from viscy_data.cell_index import preprocess_cell_index, read_cell_index, write_cell_index
+
+SRC = "/hpc/projects/jacobo_group/zuben/proj/gutCellClassifier/data/dynaclr_cell_index_bbox_center.parquet"
+V3_ROOT = "/hpc/projects/organelle_phenotyping/datasets/zuben_gut_development"
+OUT = "/hpc/projects/jacobo_group/collab/ed/dynaclr/dynaclr_cell_index_gut_v1.parquet"
+N_CHANNELS = 4
+
+
+def main() -> None:
+    """Conform the parquet, fill norm stats, and seed z_focus from the bbox-center z."""
+    df = pd.read_parquet(SRC)
+    n_rows_in = len(df)
+
+    # Repoint store_path to the v3 copies (same basename).
+    df["store_path"] = df["store_path"].astype(str).map(lambda p: f"{V3_ROOT}/{os.path.basename(p)}")
+
+    # Required columns absent from the source parquet.
+    df["tracks_path"] = ""  # ignored by ExperimentRegistry.from_cell_index
+    df["microscope"] = ""
+    df["T_shape"] = 1  # static: single timepoint
+    df["C_shape"] = N_CHANNELS
+
+    write_cell_index(df, OUT)
+
+    # Fill norm_* from the v3 .zattrs (writes z_focus as NaN — no focus_slice).
+    preprocess_cell_index(OUT, focus_channel="nuclear")
+
+    # Seed z_focus = z LAST so it survives preprocess_cell_index.
+    out_df = read_cell_index(OUT)
+    out_df["z_focus"] = out_df["z"].astype("float32")
+    write_cell_index(out_df, OUT)
+
+    check = read_cell_index(OUT)
+    print("# Step 2 — conform parquet\n")
+    print(f"- rows in: **{n_rows_in}**, rows out: **{len(check)}**")
+    print(f"- output: `{OUT}`")
+    print(f"- norm_mean NaN: **{int(check['norm_mean'].isna().sum())}**")
+    print(f"- z_focus NaN: **{int(check['z_focus'].isna().sum())}** (should be 0)")
+
+
+if __name__ == "__main__":
+    main()

applications/dynaclr/configs/collections/zuben_gut/gen_convert_v3.py

@@ -0,0 +1,92 @@
+"""Generate per-store biahub-concatenate configs to copy Zuben's 25 v2 gut stores to v3.
+
+Each of Zuben's stores is one gut = one experiment = a single position ``A/1/0`` with the
+four channels ``[nuclear, septate, brush_border, SuH]``. We convert each store to its own
+v3 store under ``OUTPUT_ROOT`` preserving the original basename and the ``A/1/0`` layout, so
+the cell-index parquet's ``(store_path, well, fov)`` mapping stays valid — only the
+``store_path`` directory changes.
+
+Emits, per store, ``configs/<stem>.yml`` and a single ``submit_all.sh`` that runs
+``biahub concatenate`` (conda env ``biautils``) for every store as SLURM jobs.
+
+Run::
+
+    uv run --no-sync python applications/dynaclr/configs/collections/zuben_gut/gen_convert_v3.py
+    bash <this_dir>/convert_v3/submit_all.sh
+"""
+
+import os
+from pathlib import Path
+
+import pandas as pd
+
+PARQUET = "/hpc/projects/jacobo_group/zuben/proj/gutCellClassifier/data/dynaclr_cell_index_bbox_center.parquet"
+OUTPUT_ROOT = "/hpc/projects/organelle_phenotyping/datasets/zuben_gut_development"
+CHANNELS = ["nuclear", "septate", "brush_border", "SuH"]
+CHUNKS_CZYX = [1, 16, 256, 256]
+SHARDS_RATIO = [1, 1, 4, 8, 8]
+OME_VERSION = "0.5"
+BIAHUB_ENV = "biautils"
+
+HERE = Path(__file__).resolve().parent
+WORKDIR = HERE / "convert_v3"
+CONFIGS_DIR = WORKDIR / "configs"
+
+
+def _yaml_config(store_path: str) -> str:
+    import yaml
+
+    cfg = {
+        # One input FOV glob per data path: this store's single position A/1/0.
+        "concat_data_paths": [f"{store_path}/A/1/*"],
+        "time_indices": "all",
+        # List-of-lists: one channel list per data path (here, one path).
+        "channel_names": [list(CHANNELS)],
+        "X_slice": "all",
+        "Y_slice": "all",
+        "Z_slice": "all",
+        "chunks_czyx": list(CHUNKS_CZYX),
+        "shards_ratio": list(SHARDS_RATIO),
+        "output_ome_zarr_version": OME_VERSION,
+    }
+    return yaml.safe_dump(cfg, default_flow_style=False, sort_keys=False)
+
+
+def main() -> None:
+    """Emit one biahub-concatenate config per store plus a SLURM submit driver."""
+    df = pd.read_parquet(PARQUET, columns=["store_path"])
+    stores = sorted(df["store_path"].astype(str).unique())
+
+    CONFIGS_DIR.mkdir(parents=True, exist_ok=True)
+    submit_lines = [
+        "#!/bin/bash",
+        "set -euo pipefail",
+        "# Auto-generated by gen_convert_v3.py. Submits one biahub concatenate per store.",
+        f"mkdir -p {OUTPUT_ROOT}",
+        "",
+    ]
+
+    for store in stores:
+        stem = os.path.basename(store)  # e.g. AAY6_sox21a_d0_63x_gut1.zarr
+        cfg_path = CONFIGS_DIR / f"{stem.replace('.zarr', '')}.yml"
+        cfg_path.write_text(_yaml_config(store))
+        out_path = f"{OUTPUT_ROOT}/{stem}"
+        submit_lines.append(f'echo "=== {stem} ==="')
+        submit_lines.append(
+            f"conda run -n {BIAHUB_ENV} biahub concatenate "
+            f'-c "{cfg_path}" -o "{out_path}" -m -sb "{WORKDIR / "sbatch_overrides.sh"}"'
+        )
+        submit_lines.append("")
+
+    (WORKDIR / "submit_all.sh").write_text("\n".join(submit_lines) + "\n")
+    (WORKDIR / "sbatch_overrides.sh").write_text(
+        "#!/bin/bash\n#SBATCH --partition=cpu\n#SBATCH --cpus-per-task=4\n#SBATCH --mem-per-cpu=32G\n"
+    )
+
+    print(f"Generated {len(stores)} configs under {CONFIGS_DIR}")
+    print(f"Submit driver: {WORKDIR / 'submit_all.sh'}")
+    print(f"Output root:   {OUTPUT_ROOT}")
+
+
+if __name__ == "__main__":
+    main()

applications/dynaclr/configs/collections/zuben_gut/gen_preprocess.py

@@ -0,0 +1,56 @@
+"""Generate + submit `viscy preprocess` SLURM jobs for the 25 v3 gut stores.
+
+Writes per-channel normalization stats (fov/dataset/timepoint) into each store's
+``.zattrs``. One SLURM job per store (each store is a single large FOV).
+
+Run::
+
+    uv run --no-sync python applications/dynaclr/configs/collections/zuben_gut/gen_preprocess.py
+    bash <this_dir>/preprocess/submit_all.sh
+"""
+
+import glob
+from pathlib import Path
+
+V3_ROOT = "/hpc/projects/organelle_phenotyping/datasets/zuben_gut_development"
+REPO = "/hpc/mydata/eduardo.hirata/repos/viscy"
+VENV = f"{REPO}/.venv-dynaclr"
+HERE = Path(__file__).resolve().parent
+WORKDIR = HERE / "preprocess"
+
+JOB_TEMPLATE = """#!/bin/bash
+#SBATCH --job-name=pp_{stem}
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --partition=cpu
+#SBATCH --cpus-per-task=32
+#SBATCH --mem-per-cpu=4G
+#SBATCH --time=02:00:00
+#SBATCH --output={workdir}/slurm_pp_{stem}_%j.out
+
+export PYTHONNOUSERSITE=1
+export UV_PROJECT_ENVIRONMENT={venv}
+
+uv run --project "{repo}" --package dynaclr \\
+    viscy preprocess --data_path "{store}" \\
+    --channel_names=-1 --num_workers 32 --block_size 32
+"""
+
+
+def main() -> None:
+    """Emit one `viscy preprocess` SLURM job per v3 store plus a submit driver."""
+    WORKDIR.mkdir(parents=True, exist_ok=True)
+    stores = sorted(glob.glob(f"{V3_ROOT}/*.zarr"))
+    submit = ["#!/bin/bash", "set -euo pipefail", ""]
+    for store in stores:
+        stem = Path(store).name.replace(".zarr", "")
+        job = WORKDIR / f"pp_{stem}.sh"
+        job.write_text(JOB_TEMPLATE.format(stem=stem, workdir=WORKDIR, venv=VENV, repo=REPO, store=store))
+        submit.append(f"sbatch {job}")
+    (WORKDIR / "submit_all.sh").write_text("\n".join(submit) + "\n")
+    print(f"Generated {len(stores)} preprocess jobs under {WORKDIR}")
+    print(f"Submit: bash {WORKDIR / 'submit_all.sh'}")
+
+
+if __name__ == "__main__":
+    main()

applications/dynaclr/configs/collections/zuben_gut/verify_parquet_zarr.py

@@ -0,0 +1,80 @@
+"""Step 0: verify Zuben's cell-index parquet is consistent with its zarr stores.
+
+Checks, for every store referenced in the parquet:
+- the store opens and exposes the expected channel names,
+- each (well, fov) referenced resolves to a real position,
+- cell centroids fit inside the FOV with room for the requested patch half-width.
+
+Run::
+
+    uv run --no-sync python applications/dynaclr/configs/collections/zuben_gut/verify_parquet_zarr.py
+"""
+
+import sys
+
+import pandas as pd
+from iohub import open_ome_zarr
+
+PARQUET = "/hpc/projects/jacobo_group/zuben/proj/gutCellClassifier/data/dynaclr_cell_index_bbox_center.parquet"
+EXPECTED_CHANNELS = ["nuclear", "septate", "brush_border", "SuH"]
+YX_PATCH_SIZE = (256, 256)  # extraction patch used by the bag-of-channels config
+
+
+def main() -> int:
+    """Check every store opens, channels match, positions resolve, and report border-OOB cells."""
+    df = pd.read_parquet(PARQUET)
+    y_half = YX_PATCH_SIZE[0] // 2
+    x_half = YX_PATCH_SIZE[1] // 2
+
+    problems: list[str] = []
+    n_cells_out_of_bounds = 0
+
+    for store_path, store_group in df.groupby("store_path", observed=True):
+        store_path = str(store_path)
+        try:
+            with open_ome_zarr(store_path, mode="r") as plate:
+                channels = list(plate.channel_names)
+                if channels != EXPECTED_CHANNELS:
+                    problems.append(f"{store_path}: channels {channels} != {EXPECTED_CHANNELS}")
+                positions = {name for name, _ in plate.positions()}
+        except Exception as exc:  # noqa: BLE001 - surface any open failure
+            problems.append(f"{store_path}: failed to open ({exc})")
+            continue
+
+        for (well, fov), fov_group in store_group.groupby(["well", "fov"], observed=True):
+            pos_key = f"{well}/{fov}"
+            if pos_key not in positions:
+                problems.append(f"{store_path}: position {pos_key} not found")
+                continue
+            y_shape = int(fov_group["Y_shape"].iloc[0])
+            x_shape = int(fov_group["X_shape"].iloc[0])
+            y = fov_group["y"].to_numpy()
+            x = fov_group["x"].to_numpy()
+            oob = (y < y_half) | (y > y_shape - y_half) | (x < x_half) | (x > x_shape - x_half)
+            n_cells_out_of_bounds += int(oob.sum())
+
+    n_stores = df["store_path"].nunique()
+    n_cells = df["cell_id"].nunique()
+
+    print("# Step 0 — parquet↔zarr consistency\n")
+    print(f"- stores checked: **{n_stores}**")
+    print(f"- unique cells: **{n_cells}**")
+    print(f"- expected channels: `{EXPECTED_CHANNELS}`")
+    print(
+        f"- cells whose {YX_PATCH_SIZE} patch would fall out of bounds: "
+        f"**{n_cells_out_of_bounds}** "
+        f"({100 * n_cells_out_of_bounds / len(df):.1f}% of rows)"
+    )
+
+    if problems:
+        print(f"\n**{len(problems)} problem(s):**")
+        for p in problems[:50]:
+            print(f"  - {p}")
+        return 1
+
+    print("\n**OK** — all stores open, channels match, positions resolve.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

applications/dynaclr/configs/training/DynaCLR-3D/DynaCLR-3D-Gut-BagOfChannels.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+# DynaCLR-3D-Gut-BagOfChannels (Phase 1) — Zuben gut cells, bag-of-channels SimCLR.
+#
+# New run:
+#   sbatch applications/dynaclr/configs/training/DynaCLR-3D/DynaCLR-3D-Gut-BagOfChannels.sh
+# Resume:
+#   CKPT_PATH=.../last.ckpt WANDB_RUN_ID=<id> sbatch .../DynaCLR-3D-Gut-BagOfChannels.sh
+
+#SBATCH --job-name=dynaclr_gut_boc
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=2
+#SBATCH --gpus-per-node=2
+#SBATCH --partition=gpu
+#SBATCH --cpus-per-task=15
+#SBATCH --mem-per-cpu=8G
+#SBATCH --time=2-00:00:00
+
+# Set WORKSPACE_DIR to YOUR clone of the repo before submitting, e.g.
+#   WORKSPACE_DIR=/hpc/mydata/<you>/repos/VisCy sbatch <this>.sh
+# MODEL_ROOT is where checkpoints/configs are written (defaults to your clone's
+# models/ dir; override to a shared project path if desired).
+WORKSPACE_DIR="${WORKSPACE_DIR:?Set WORKSPACE_DIR to your repo clone path}"
+
+export PROJECT="DynaCLR-3D-Gut-BagOfChannels"
+export RUN_NAME="gut-3d-z24-cellz-64-ntxent-t0p2-self"
+export CONFIGS="applications/dynaclr/configs/training/DynaCLR-3D/DynaCLR-3D-Gut-BagOfChannels.yml"
+export MODEL_ROOT="${MODEL_ROOT:-${WORKSPACE_DIR}/models}"
+# Point at the dynaclr-pinned venv in your clone (avoids the shared .venv sync race).
+export UV_PROJECT_ENVIRONMENT="${UV_PROJECT_ENVIRONMENT:-${WORKSPACE_DIR}/.venv-dynaclr}"
+
+# W&B writes sample images to a temp dir under $TMPDIR before upload. On some
+# SLURM nodes the default /tmp is per-job and gets swept mid-run, causing
+# `FileNotFoundError: .../wandb-media/*.png` at validation image logging. Pin
+# TMPDIR + WANDB_DIR to persistent paths we create so they can't disappear.
+export TMPDIR="${TMPDIR:-${MODEL_ROOT}/tmp}"
+export WANDB_DIR="${WANDB_DIR:-${MODEL_ROOT}/wandb}"
+mkdir -p "$TMPDIR" "$WANDB_DIR"
+
+# The shared trainer recipe logs to the `computational_imaging` W&B entity. Set
+# WANDB_ENTITY to your own entity to log there instead (EXTRA_ARGS overrides the
+# recipe). Leave unset to keep the default.
+if [ -n "${WANDB_ENTITY:-}" ]; then
+  export EXTRA_ARGS="${EXTRA_ARGS:-} --trainer.logger.init_args.entity=${WANDB_ENTITY}"
+fi
+
+# Absolute path (SLURM spools this script, so $(dirname "$0") would break).
+source "${WORKSPACE_DIR}/applications/dynaclr/configs/training/slurm/train.sh"