Upload hf_train_1kb.py with huggingface_hub

hf_train_1kb.py

@@ -0,0 +1,677 @@
+#!/usr/bin/env python3
+"""
+hf_train_1kb.py
+------------------
+HF Jobs version of modal_pipeline_1kb.py. Combines prepare_data() and
+train_model() into one script, since HF Jobs runs a single command to
+completion rather than exposing persistent, independently-callable
+Functions the way Modal does -- there's no direct equivalent of
+`train_model.remote()` to call repeatedly against a warm environment.
+
+All the actual logic (blacklist handling, bedpe parsing with the chr1
+fix, span-fit window placement, anchor-weighted masked loss) is copied
+verbatim from modal_pipeline_1kb.py, which was already tested against
+synthetic data -- this file only changes HOW it's invoked (plain
+function calls + argparse, not Modal decorators) and WHERE data lives
+(a mounted bucket path instead of a Modal Volume path).
+
+Expects two mounted volumes (see submit_hf_job.py):
+  /code  -- read-only dataset repo containing src/ and configs/
+  /data  -- read-write bucket, plays the role all three Modal volumes
+            (hic-data, hic-cache, hic-runs) played: raw/ downloads,
+            cache/windows_1kb/ per-window .npz cache, runs/ checkpoints.
+
+Run (inside the job, via submit_hf_job.py):
+    python hf_train_1kb.py --model bilstm
+    python hf_train_1kb.py --model bilstm --exclude-track RAD21
+    python hf_train_1kb.py --prepare-only
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import gzip
+import hashlib
+import json
+import math
+import random
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import yaml
+
+
+def _cfg(config_dir: Path, config_name: str):
+    with open(config_dir / config_name) as f:
+        return yaml.safe_load(f)
+
+
+def _track_fingerprint(track_names: list[str]) -> str:
+    h = hashlib.sha1("|".join(sorted(track_names)).encode()).hexdigest()
+    return h[:8]
+
+
+# ── Blacklist -- verbatim from modal_pipeline_1kb.py ────────────────────────
+
+def load_blacklist(bed_gz_path: Path) -> dict:
+    regions: dict = {}
+    with gzip.open(bed_gz_path, "rt") as fh:
+        for line in fh:
+            if line.startswith("#") or line.startswith("track"):
+                continue
+            p = line.rstrip("\n").split("\t")
+            if len(p) < 3:
+                continue
+            chrom = p[0] if p[0].startswith("chr") else f"chr{p[0]}"
+            try:
+                start, end = int(p[1]), int(p[2])
+            except ValueError:
+                continue
+            regions.setdefault(chrom, []).append((start, end))
+    for chrom in regions:
+        regions[chrom].sort()
+    return regions
+
+
+def overlaps_blacklist(chrom: str, start: int, end: int, blacklist: dict) -> bool:
+    for (bl_start, bl_end) in blacklist.get(chrom, []):
+        if start < bl_end and end > bl_start:
+            return True
+    return False
+
+
+# ── Loop parsing -- verbatim, includes the chr1 header-detection fix ───────
+
+def parse_encode_bedpe_gz(gz_path) -> list[dict]:
+    loops = []
+    with gzip.open(gz_path, "rt") as fh:
+        for line in fh:
+            if line.startswith("#"):
+                continue
+            p = line.strip().split("\t")
+            if len(p) < 6:
+                continue
+            c1, c2 = p[0], p[3]
+            try:
+                x1, x2, y1, y2 = int(p[1]), int(p[2]), int(p[4]), int(p[5])
+            except ValueError:
+                continue
+            if c1 != c2:
+                continue
+            chrom = c1 if c1.startswith("chr") else f"chr{c1}"
+            loops.append({"chrom": chrom, "mid_x": (x1 + x2) // 2, "mid_y": (y1 + y2) // 2})
+    return loops
+
+
+# ── Window placement -- verbatim ────────────────────────────────────────────
+
+def place_window_loop_centered(loop, window_size, flank_margin, jitter, rng, sizes, blacklist,
+                                max_attempts: int = 5):
+    chrom = loop["chrom"]
+    mid_x, mid_y = loop["mid_x"], loop["mid_y"]
+    span = abs(mid_y - mid_x)
+    max_span = window_size - 2 * flank_margin
+    if span > max_span:
+        return None, "span_too_large"
+    clen = sizes.get(chrom)
+    if clen is None:
+        return None, "unknown_chrom"
+    true_center = (mid_x + mid_y) // 2
+    for _ in range(max_attempts):
+        center = true_center + rng.randint(-jitter, jitter)
+        start = center - window_size // 2
+        end = start + window_size
+        if start < 0 or end > clen:
+            continue
+        if not (start + flank_margin <= mid_x <= end - flank_margin):
+            continue
+        if not (start + flank_margin <= mid_y <= end - flank_margin):
+            continue
+        if overlaps_blacklist(chrom, start, end, blacklist):
+            continue
+        return (start, end), None
+    return None, "placement_failed"
+
+
+def sample_background_window(chroms, window_size, sizes, blacklist, rng, max_attempts=20):
+    chrom_list = list(chroms)
+    for _ in range(max_attempts):
+        chrom = rng.choice(chrom_list)
+        clen = sizes.get(chrom)
+        if clen is None or clen <= window_size:
+            continue
+        start = rng.randint(0, clen - window_size)
+        end = start + window_size
+        if overlaps_blacklist(chrom, start, end, blacklist):
+            continue
+        return chrom, start, end
+    return None
+
+
+# ── Losses -- verbatim ───────────────────────────────────────────────────────
+
+def masked_huber_loss(pred, target, valid_mask, delta: float = 0.5):
+    import torch
+    diff = pred - target
+    abs_diff = diff.abs()
+    quadratic = torch.clamp(abs_diff, max=delta)
+    linear = abs_diff - quadratic
+    per_pixel = (0.5 * quadratic ** 2 + delta * linear) * valid_mask
+    denom = valid_mask.sum().clamp(min=1.0)
+    return per_pixel.sum() / denom
+
+
+def masked_mse_loss(pred, target, valid_mask):
+    per_pixel = (pred - target) ** 2 * valid_mask
+    denom = valid_mask.sum().clamp(min=1.0)
+    return per_pixel.sum() / denom
+
+
+def build_anchor_weight_map(n_bins: int, anchor_r, anchor_c, boost_radius: int = 5,
+                              boost_factor: float = 10.0):
+    w = np.ones((n_bins, n_bins), dtype=np.float32)
+    if anchor_r is not None and anchor_c is not None and anchor_r >= 0 and anchor_c >= 0:
+        r0, r1 = max(0, anchor_r - boost_radius), min(n_bins, anchor_r + boost_radius + 1)
+        c0, c1 = max(0, anchor_c - boost_radius), min(n_bins, anchor_c + boost_radius + 1)
+        w[r0:r1, c0:c1] = boost_factor
+    return w
+
+
+def masked_weighted_huber_loss(pred, target, valid_mask, weight_map, delta: float = 0.5):
+    import torch
+    diff = pred - target
+    abs_diff = diff.abs()
+    quadratic = torch.clamp(abs_diff, max=delta)
+    linear = abs_diff - quadratic
+    per_pixel = (0.5 * quadratic ** 2 + delta * linear) * valid_mask * weight_map
+    denom = (valid_mask * weight_map).sum().clamp(min=1.0)
+    return per_pixel.sum() / denom
+
+
+# ── Data prep ────────────────────────────────────────────────────────────────
+
+def resolve_file(rel_path: str, shared_raw, local_raw: Path, url, genomic_io) -> Path:
+    """Returns a usable local path for `rel_path`, preferring a file that
+    already exists over downloading it.
+
+    Lookup order:
+      1. shared_raw/rel_path -- the read-only data bucket (e.g. the existing
+         3d-genomics-prediction bucket, which already holds every bigWig and
+         both .hic files)
+      2. local_raw/rel_path  -- downloaded by an earlier run
+      3. download from `url` into local_raw/rel_path
+
+    This exists because the first HF Jobs run re-downloaded ~90GB from
+    ENCODE into a fresh bucket while identical files sat unused in an
+    existing one. Checking before fetching turns that into a no-op.
+    Raises if missing everywhere AND no url is given, rather than returning
+    a path that doesn't exist and failing later inside pyBigWig/hicstraw
+    with a much less obvious message.
+    """
+    if shared_raw is not None:
+        candidate = Path(shared_raw) / rel_path
+        if candidate.exists() and candidate.stat().st_size > 0:
+            print(f"  [shared] {rel_path}")
+            return candidate
+    candidate = local_raw / rel_path
+    if candidate.exists() and candidate.stat().st_size > 0:
+        print(f"  [cache] {rel_path}")
+        return candidate
+    if not url:
+        raise FileNotFoundError(
+            f"'{rel_path}' not found in shared raw dir ({shared_raw}) or local raw dir "
+            f"({local_raw}), and no download url was configured for it.")
+    genomic_io.download_file(url, candidate)
+    return candidate
+
+
+def prepare_data(cfg: dict, data_dir: Path, genomic_io, genome_constants,
+                  shared_raw=None) -> None:
+    exp = cfg["experiment"]
+    raw = data_dir / "raw"
+    raw.mkdir(parents=True, exist_ok=True)
+    manifest_path = data_dir / "loop_manifest_1kb.csv"
+
+    track_name_sets = {cl: sorted(cl_cfg["tracks"].keys()) for cl, cl_cfg in cfg["cell_lines"].items()}
+    reference_tracks = next(iter(track_name_sets.values()))
+    for cl, names in track_name_sets.items():
+        if names != reference_tracks:
+            raise ValueError(f"Track set mismatch: '{cl}' has {names}, expected {reference_tracks}.")
+    track_fp = _track_fingerprint(reference_tracks)
+
+    fingerprint_path = data_dir / "manifest_fingerprint_1kb.txt"
+    needed_fp = (f"{sorted(cfg['cell_lines'].keys())}|{track_fp}|"
+                 f"{exp['bin_size']}|{exp['window_size']}|{exp['flank_margin']}")
+
+    blacklist = {}
+    if cfg.get("blacklist_filename"):
+        bl_path = resolve_file(cfg["blacklist_filename"], shared_raw, raw,
+                                cfg.get("blacklist_url"), genomic_io)
+        blacklist = load_blacklist(bl_path)
+        n_regions = sum(len(v) for v in blacklist.values())
+        print(f"[prepare] Blacklist: {n_regions} regions across {len(blacklist)} chroms.")
+
+    for cl_name, cl_cfg in cfg["cell_lines"].items():
+        print(f"\n[prepare] {cl_name}: resolving {len(cl_cfg['tracks'])} tracks, loop calls, "
+              f"and the intact .hic file...")
+        for track_name, entry in cl_cfg["tracks"].items():
+            resolve_file(entry["path"], shared_raw, raw, entry.get("url"), genomic_io)
+        resolve_file(cl_cfg["loop_filename"], shared_raw, raw, cl_cfg.get("loop_url"), genomic_io)
+        # hic_path is config-driven now -- the bucket names .hic files by
+        # accession (HepG2/ENCFF807IRK.hic), so the old hardcoded
+        # f"{cell_line}.hic" convention doesn't match it. Falls back to that
+        # convention when a config doesn't specify hic_path.
+        resolve_file(cl_cfg.get("hic_path", f"{cl_name}.hic"), shared_raw, raw,
+                      cl_cfg.get("hic_url"), genomic_io)
+        print(f"[prepare] {cl_name}: done.")
+
+    if manifest_path.exists():
+        old_fp = fingerprint_path.read_text().strip() if fingerprint_path.exists() else None
+        if old_fp == needed_fp:
+            print(f"\n[prepare] Manifest exists for {needed_fp}. Downloads verified above; "
+                  f"skipping window rebuild.")
+            return
+        print(f"\n[prepare] Config changed ({old_fp} -> {needed_fp}). Rebuilding manifest.")
+        manifest_path.unlink()
+
+    sizes = genome_constants.chrom_sizes_hg38()
+    bin_size = exp["bin_size"]; window_size = exp["window_size"]
+    jitter = exp["jitter"]; flank_margin = exp["flank_margin"]
+    n_bins = window_size // bin_size
+    train_chroms = set(exp["train_chroms"]); test_chroms = set(exp["test_chroms"])
+    rng = random.Random(exp["seed"])
+
+    rows = []
+    n_span_rejected = 0; n_placement_failed = 0
+    for cl_name, cl_cfg in cfg["cell_lines"].items():
+        loop_file = resolve_file(cl_cfg["loop_filename"], shared_raw, raw,
+                                  cl_cfg.get("loop_url"), genomic_io)
+        loops = parse_encode_bedpe_gz(loop_file)
+        print(f"\n[prepare] {cl_name}: {len(loops)} intra-chromosomal loops parsed.")
+
+        for split, chroms, n_target in [
+            ("train", train_chroms, exp["train_loops_per_cell_line"]),
+            ("test", test_chroms, exp["test_loops_per_cell_line"]),
+        ]:
+            eligible = [l for l in loops if l["chrom"] in chroms]
+            rng.shuffle(eligible)
+            placed = 0
+            for loop in eligible:
+                if placed >= n_target:
+                    break
+                result, reason = place_window_loop_centered(
+                    loop, window_size, flank_margin, jitter, rng, sizes, blacklist)
+                if result is None:
+                    if reason == "span_too_large":
+                        n_span_rejected += 1
+                    else:
+                        n_placement_failed += 1
+                    continue
+                start, end = result
+                rows.append({
+                    "split": split, "cell_line": cl_name, "window_type": "loop",
+                    "chrom": loop["chrom"], "start": start, "end": end,
+                    "bin_size": bin_size, "n_bins": n_bins,
+                    "loop_mid_x": loop["mid_x"], "loop_mid_y": loop["mid_y"],
+                    "loop_span_bp": abs(loop["mid_y"] - loop["mid_x"]),
+                    "track_fingerprint": track_fp,
+                })
+                placed += 1
+            status = "WARNING: only" if placed < n_target else " "
+            print(f"  {status} {placed}/{n_target} loop windows: {cl_name}/{split}")
+
+        n_bg_target = exp.get("background_windows_per_cell_line", 0)
+        placed_bg = 0
+        for _ in range(n_bg_target * 3):
+            if placed_bg >= n_bg_target:
+                break
+            result = sample_background_window(train_chroms, window_size, sizes, blacklist, rng)
+            if result is None:
+                continue
+            chrom, start, end = result
+            rows.append({
+                "split": "train", "cell_line": cl_name, "window_type": "background",
+                "chrom": chrom, "start": start, "end": end,
+                "bin_size": bin_size, "n_bins": n_bins,
+                "loop_mid_x": "", "loop_mid_y": "", "loop_span_bp": "",
+                "track_fingerprint": track_fp,
+            })
+            placed_bg += 1
+        print(f"  {placed_bg}/{n_bg_target} background windows: {cl_name}/train")
+
+    print(f"\n[prepare] Rejected: {n_span_rejected} loops (span too large), "
+          f"{n_placement_failed} (blacklist/bounds after retries).")
+
+    fieldnames = ["split", "cell_line", "window_type", "chrom", "start", "end",
+                  "bin_size", "n_bins", "loop_mid_x", "loop_mid_y", "loop_span_bp",
+                  "track_fingerprint"]
+    with open(manifest_path, "w", newline="") as fh:
+        w = csv.DictWriter(fh, fieldnames=fieldnames)
+        w.writeheader(); w.writerows(rows)
+    fingerprint_path.write_text(needed_fp)
+    print(f"\n[prepare] Done: {len(rows)} windows -> {manifest_path}")
+
+
+# ── Training ─────────────────────────────────────────────────────────────────
+
+def train_model(model_name: str, cfg: dict, data_dir: Path, src_dir: Path,
+                 exclude_track: str = "", shared_raw=None) -> dict:
+    import torch
+    from torch.utils.data import DataLoader, Dataset
+    sys.path.insert(0, str(src_dir))
+    from genomic_io import read_bigwig, _get_hicfile
+    from oe_extraction import get_oe_matrix, MAX_OE
+
+    tr_cfg = cfg["training"]
+    raw = data_dir / "raw"
+
+    def _resolve_existing(rel_path: str) -> str:
+        """Same shared-then-local lookup as resolve_file(), minus the
+        download fallback -- by training time prepare_data() has already
+        ensured every file exists somewhere, so a miss here is a bug worth
+        surfacing immediately rather than silently re-fetching 40GB mid-run."""
+        if shared_raw is not None:
+            cand = Path(shared_raw) / rel_path
+            if cand.exists():
+                return str(cand)
+        return str(raw / rel_path)
+
+    train_seed = int(tr_cfg.get("seed", cfg.get("experiment", {}).get("seed", 42)))
+    torch.manual_seed(train_seed)
+    torch.cuda.manual_seed_all(train_seed)
+    print(f"[train] seed={train_seed}")
+
+    track_names = sorted(next(iter(cfg["cell_lines"].values()))["tracks"].keys())
+    if exclude_track:
+        if exclude_track not in track_names:
+            raise ValueError(f"exclude_track='{exclude_track}' not in {track_names}")
+        track_names = [t for t in track_names if t != exclude_track]
+        print(f"[train] LEAVE-ONE-TRACK-OUT: excluding '{exclude_track}'")
+    n_tracks = len(track_names)
+    track_fp = _track_fingerprint(track_names)
+    max_oe = float(tr_cfg.get("max_oe", MAX_OE))
+    print(f"[train] {n_tracks} tracks: {track_names}")
+
+    class LoopDataset(Dataset):
+        def __init__(self, manifest_path, split, normalizer=None, cache_dir=None, clip_pct=99.0):
+            full = pd.read_csv(manifest_path)
+            self.manifest = full[full["split"] == split].reset_index(drop=True)
+            self.normalizer = normalizer; self.clip_pct = clip_pct
+            self.cache_root = Path(cache_dir) if cache_dir else None
+            if self.cache_root:
+                self.cache_root.mkdir(parents=True, exist_ok=True)
+
+        def __len__(self):
+            return len(self.manifest)
+
+        def _paths(self, cl):
+            cl_tracks = cfg["cell_lines"][cl]["tracks"]
+            return {name: _resolve_existing(cl_tracks[name]["path"]) for name in track_names} | \
+                   {"hic": _resolve_existing(cfg["cell_lines"][cl].get("hic_path", f"{cl}.hic"))}
+
+        def _cache(self, row):
+            if self.cache_root is None:
+                return None
+            bs = int(row["bin_size"])
+            return self.cache_root / (f"{row['cell_line']}_{row['chrom']}_{row['start']}_"
+                                        f"{row['end']}_log2oe_v2_{bs}_{track_fp}.npz")
+
+        def get_tracks(self, idx):
+            row = self.manifest.iloc[idx]
+            p = self._paths(row["cell_line"]); n = int(row["n_bins"])
+            ch, s, e = row["chrom"], int(row["start"]), int(row["end"])
+            arrays = [read_bigwig(p[name], ch, s, e, n, self.clip_pct) for name in track_names]
+            return np.stack(arrays, axis=0)
+
+        def __getitem__(self, idx):
+            row = self.manifest.iloc[idx]; cl = row["cell_line"]
+            ch, s, e = row["chrom"], int(row["start"]), int(row["end"])
+            bs = int(row["bin_size"]); cf = self._cache(row)
+            if cf is not None and cf.exists():
+                data = np.load(cf)
+                tracks, log2_oe, valid_mask = data["tracks"], data["log2_oe"], data["valid_mask"]
+            else:
+                tracks = self.get_tracks(idx)
+                hf = _get_hicfile(self._paths(cl)["hic"])
+                _, log2_oe, valid_mask = get_oe_matrix(hf, ch, s, e, bs, max_oe=max_oe)
+                if cf is not None:
+                    np.savez_compressed(cf, tracks=tracks.astype(np.float32),
+                                         log2_oe=log2_oe, valid_mask=valid_mask)
+            if self.normalizer is not None:
+                s_n = self.normalizer[cl]
+                tracks = (tracks - s_n["mean"][:, None]) / s_n["std"][:, None]
+            if row["window_type"] == "loop" and pd.notna(row.get("loop_mid_x")):
+                anchor_r = (int(row["loop_mid_x"]) - s) // bs
+                anchor_c = (int(row["loop_mid_y"]) - s) // bs
+            else:
+                anchor_r = anchor_c = -1
+            return {"tracks": torch.from_numpy(tracks).float(),
+                    "target": torch.from_numpy(log2_oe).float(),
+                    "valid_mask": torch.from_numpy(valid_mask.astype(np.float32)),
+                    "anchor_r": anchor_r, "anchor_c": anchor_c, "cell_line": cl}
+
+    manifest = data_dir / "loop_manifest_1kb.csv"
+    run_dir = data_dir / "runs" / f"1kb_{model_name}_{track_fp}"
+    run_dir.mkdir(parents=True, exist_ok=True)
+    cache_dir = data_dir / "cache" / "windows_1kb"
+
+    norm_path = data_dir / "runs" / f"normalizer_1kb_{track_fp}.json"
+    config_cls = set(cfg["cell_lines"].keys())
+    need_refit = True
+    if norm_path.exists():
+        with open(norm_path) as f:
+            normalizer = json.load(f)
+        if set(normalizer.keys()) == config_cls:
+            need_refit = False
+            print(f"[train] Loaded normalizer for {config_cls}.")
+
+    if need_refit:
+        train_df = pd.read_csv(manifest)
+        train_df = train_df[train_df["split"] == "train"].reset_index(drop=True)
+        raw_ds = LoopDataset(manifest, "train", normalizer=None, cache_dir=cache_dir)
+        normalizer = {}
+        for cl in train_df["cell_line"].unique():
+            cl_pos = train_df[train_df["cell_line"] == cl].index.tolist()
+            n_s = min(80, len(cl_pos))
+            sample_idx = np.linspace(0, len(cl_pos) - 1, n_s).astype(int)
+            sums = np.zeros(n_tracks); sq_sums = np.zeros(n_tracks); count = 0
+            for si in sample_idx:
+                t = raw_ds.get_tracks(cl_pos[si])
+                sums += t.sum(axis=1); sq_sums += (t ** 2).sum(axis=1); count += t.shape[1]
+            mean = sums / count
+            std = np.sqrt(np.clip(sq_sums / count - mean ** 2, 1e-8, None))
+            normalizer[cl] = {"mean": mean.tolist(), "std": std.tolist()}
+            print(f" {cl}: " + " ".join(f"{n}={m:.3f}" for n, m in zip(track_names, mean)))
+        run_dir.parent.mkdir(parents=True, exist_ok=True)
+        with open(norm_path, "w") as f:
+            json.dump(normalizer, f, indent=2)
+
+    norm_np = {cl: {"mean": np.array(v["mean"], dtype=np.float32),
+                     "std": np.array(v["std"], dtype=np.float32)} for cl, v in normalizer.items()}
+
+    train_ds = LoopDataset(manifest, "train", normalizer=norm_np, cache_dir=cache_dir)
+    test_ds = LoopDataset(manifest, "test", normalizer=norm_np, cache_dir=cache_dir)
+    print(f"[train] {len(train_ds)} train, {len(test_ds)} test windows")
+
+    bs = tr_cfg["batch_size"]
+    train_loader = DataLoader(train_ds, batch_size=bs, shuffle=True, num_workers=0, pin_memory=True)
+    test_loader = DataLoader(test_ds, batch_size=bs, shuffle=False, num_workers=0, pin_memory=True)
+    n_bins = int(train_ds.manifest["n_bins"].iloc[0])
+
+    device = "cuda"
+    if model_name == "cnn":
+        from model_cnn import CNNHiCModel
+        model = CNNHiCModel(n_tracks=n_tracks, n_bins=n_bins).to(device)
+        wd = float(tr_cfg["weight_decay"]); clip_norm = None
+    elif model_name == "unet":
+        from model_unet import UNetHiCModel
+        model = UNetHiCModel(n_tracks=n_tracks, n_bins=n_bins).to(device)
+        wd = float(tr_cfg["weight_decay"]); clip_norm = None
+    elif model_name == "transformer":
+        from model_transformer import TransformerHiCModel
+        dropout = float(tr_cfg.get("transformer_dropout", 0.3))
+        model = TransformerHiCModel(n_tracks=n_tracks, n_bins=n_bins, dropout=dropout).to(device)
+        wd = float(tr_cfg.get("transformer_weight_decay", 1e-3))
+        clip_norm = float(tr_cfg.get("transformer_grad_clip_norm", 0.5))
+    elif model_name == "bilstm":
+        from model_bilstm import BiLSTMHiCModel
+        dropout = float(tr_cfg.get("bilstm_dropout", 0.2))
+        model = BiLSTMHiCModel(n_tracks=n_tracks, n_bins=n_bins, dropout=dropout).to(device)
+        wd = float(tr_cfg.get("bilstm_weight_decay", 1e-5))
+        clip_norm = float(tr_cfg.get("bilstm_grad_clip_norm", 1.0))
+    else:
+        raise ValueError(f"Unknown model: {model_name}")
+
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"[train] {model_name}: {n_params:,} params")
+
+    epochs = tr_cfg["epochs"][model_name]
+    patience = tr_cfg.get("early_stopping_patience", 10)
+    warmup = tr_cfg.get("transformer_warmup_steps", 0) if model_name == "transformer" else 0
+    total_steps = len(train_loader) * epochs
+    optimizer = torch.optim.Adam(model.parameters(), lr=float(tr_cfg["lr"]), weight_decay=wd)
+
+    def lr_lambda(step):
+        if warmup <= 0:
+            return 1.0
+        if step < warmup:
+            return (step + 1) / warmup
+        progress = min(1.0, (step - warmup) / max(1, total_steps - warmup))
+        return max(0.05, 0.5 * (1 + math.cos(math.pi * progress)))
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+    log_path = run_dir / "training_log.csv"
+    ckpt_path = run_dir / "best_model.pt"
+    best_test_loss = float("inf"); epochs_no_improve = 0; stopped_at = None
+    huber_delta = float(tr_cfg.get("huber_delta", 0.5))
+    anchor_boost_radius = int(tr_cfg.get(f"{model_name}_anchor_boost_radius",
+                                          tr_cfg.get("anchor_boost_radius", 5)))
+    anchor_boost_factor = float(tr_cfg.get(f"{model_name}_anchor_boost_factor",
+                                            tr_cfg.get("anchor_boost_factor", 10.0)))
+    print(f"[train] anchor-weighted loss: radius={anchor_boost_radius} factor={anchor_boost_factor}")
+
+    fields = ["epoch", "train_huber", "train_huber_weighted", "train_mse", "test_huber",
+              "test_mse", "epoch_seconds", "lr"]
+    with open(log_path, "w", newline="") as log_fh:
+        writer = csv.DictWriter(log_fh, fieldnames=fields)
+        writer.writeheader()
+        for epoch in range(1, epochs + 1):
+            t0 = time.time()
+            model.train()
+            tr_hub = tr_hub_weighted = tr_mse = n_tr = 0; last_lr = None
+            for batch in train_loader:
+                tracks = batch["tracks"].to(device, non_blocking=True)
+                target = batch["target"].to(device, non_blocking=True)
+                valid = batch["valid_mask"].to(device, non_blocking=True)
+                anchors_r = batch["anchor_r"].numpy(); anchors_c = batch["anchor_c"].numpy()
+                weight_maps = np.stack([
+                    build_anchor_weight_map(n_bins, r if r >= 0 else None, c if c >= 0 else None,
+                                             boost_radius=anchor_boost_radius,
+                                             boost_factor=anchor_boost_factor)
+                    for r, c in zip(anchors_r, anchors_c)])
+                weight = torch.from_numpy(weight_maps).float().to(device, non_blocking=True)
+
+                optimizer.zero_grad()
+                pred = model(tracks)
+                loss = masked_weighted_huber_loss(pred, target, valid, weight, delta=huber_delta)
+                loss.backward()
+                if clip_norm:
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), clip_norm)
+                optimizer.step(); scheduler.step()
+                last_lr = scheduler.get_last_lr()[0]
+                tr_hub_weighted += loss.item()
+                tr_hub += masked_huber_loss(pred, target, valid, delta=huber_delta).item()
+                tr_mse += masked_mse_loss(pred, target, valid).item()
+                n_tr += 1
+            tr_hub /= max(1, n_tr); tr_hub_weighted /= max(1, n_tr); tr_mse /= max(1, n_tr)
+
+            model.eval(); te_hub = te_mse = n_te = 0
+            with torch.no_grad():
+                for batch in test_loader:
+                    tracks = batch["tracks"].to(device, non_blocking=True)
+                    target = batch["target"].to(device, non_blocking=True)
+                    valid = batch["valid_mask"].to(device, non_blocking=True)
+                    pred = model(tracks)
+                    te_hub += masked_huber_loss(pred, target, valid, delta=huber_delta).item()
+                    te_mse += masked_mse_loss(pred, target, valid).item()
+                    n_te += 1
+            te_hub /= max(1, n_te); te_mse /= max(1, n_te)
+
+            dt = time.time() - t0
+            print(f"  epoch {epoch:3d}/{epochs}  tr_hub={tr_hub:.4f} tr_hub_w={tr_hub_weighted:.4f} "
+                  f"tr_mse={tr_mse:.4f}  te_hub={te_hub:.4f} te_mse={te_mse:.4f}  "
+                  f"lr={last_lr:.2e}  ({dt:.1f}s)")
+            writer.writerow({"epoch": epoch, "train_huber": tr_hub,
+                              "train_huber_weighted": tr_hub_weighted, "train_mse": tr_mse,
+                              "test_huber": te_hub, "test_mse": te_mse,
+                              "epoch_seconds": dt, "lr": last_lr})
+            log_fh.flush()
+
+            if te_hub < best_test_loss:
+                best_test_loss = te_hub; epochs_no_improve = 0
+                torch.save({"model_name": model_name, "n_bins": n_bins, "n_tracks": n_tracks,
+                            "track_names": track_names, "epoch": epoch,
+                            "model_state_dict": model.state_dict(),
+                            "test_huber": te_hub, "test_mse": te_mse}, ckpt_path)
+            else:
+                epochs_no_improve += 1
+                if epochs_no_improve >= patience:
+                    stopped_at = epoch
+                    print(f"[train] Early stop epoch {epoch}")
+                    break
+
+    return {"model": model_name, "n_tracks": n_tracks, "best_test_huber": best_test_loss,
+            "best_epoch": int(torch.load(ckpt_path, map_location="cpu", weights_only=False)["epoch"]),
+            "stopped_early_at": stopped_at}
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="bilstm", choices=["cnn", "unet", "transformer", "bilstm"])
+    p.add_argument("--config-name", default="encode_sources_1kb.yaml")
+    p.add_argument("--exclude-track", default="")
+    p.add_argument("--prepare-only", action="store_true")
+    p.add_argument("--code-dir", default="/code")
+    p.add_argument("--data-dir", default="/data")
+    p.add_argument("--shared-raw-dir", default="",
+                    help="Read-only mount holding the already-uploaded raw data "
+                         "(e.g. /shared for the 3d-genomics-prediction bucket). Files found "
+                         "here are used as-is; anything missing falls back to downloading "
+                         "into --data-dir/raw. Empty = download everything, the old behavior.")
+    args = p.parse_args()
+
+    code_dir = Path(args.code_dir)
+    data_dir = Path(args.data_dir)
+    shared_raw = Path(args.shared_raw_dir) if args.shared_raw_dir else None
+    src_dir = code_dir / "src"
+    sys.path.insert(0, str(src_dir))
+    import genomic_io
+    import genome_constants
+
+    cfg = _cfg(code_dir / "configs", args.config_name)
+
+    print("=" * 65)
+    print(f"HepG2/HCT116 1kb loop prediction (HF Jobs) -- config: {args.config_name}")
+    if shared_raw:
+        print(f"shared read-only raw data: {shared_raw}")
+    print("=" * 65)
+    print("\n[Step 1] Preparing data...")
+    prepare_data(cfg, data_dir, genomic_io, genome_constants, shared_raw=shared_raw)
+    print("[Step 1] Done.")
+    if args.prepare_only:
+        return
+
+    print(f"\n[Step 2] Training {args.model}"
+          + (f" (excluding {args.exclude_track})" if args.exclude_track else "") + "...")
+    r = train_model(args.model, cfg, data_dir, src_dir, args.exclude_track, shared_raw=shared_raw)
+    s = f" (early@{r['stopped_early_at']})" if r["stopped_early_at"] else ""
+    print(f"\n<- {args.model.upper()}{s}: best_test_huber={r['best_test_huber']:.4f}  "
+          f"epoch={r['best_epoch']}")
+
+
+if __name__ == "__main__":
+    main()