experiments: dataset audit, evaluation, baselines and ablations

.gitignore

@@ -11,3 +11,6 @@ S2FApp/ckp/spheroid/ckp_spheroid_FN.pth
 S2FApp/ckp/single_cell/ckp_singlecell_GN.pth
 S2FApp/ckp/spheroid/ckp_spheroid_GN.pth
 S2FApp/*.docx
+
+# experiments: trained ablation checkpoints (large)
+experiments/runs/

experiments/README.md

@@ -0,0 +1,113 @@
+# S2F experiments: data audit, evaluation, baselines and ablations
+
+This folder holds the code for every number reported in the paper beyond the training of the released checkpoints:
+- the dataset audit and split manifests;
+- the evaluation of the released models;
+- the baselines and ablation study.
+
+Run every command from the repository root (`Shape2Force/`); a GPU is used automatically when available.
+
+```
+experiments/
+  layout.py           dataset layout, per-set metadata (Young's modulus, TFM settings), manifest helpers
+  prepare_dataset.py  audit legacy data vs dataset/, copy missing images, write splits/*.csv
+  manifest_data.py    data loader driven by a split manifest (same preprocessing as data/cell_dataset.py)
+  ablation_models.py  plain U-Net and S2F generator variants
+  losses.py           per-image WFM-RME reconstruction loss (as used for the released model)
+  configs.py          shared training protocol + one entry per variant
+  train.py            train one variant
+  evaluate.py         evaluate released checkpoints / baselines / trained runs on all manifests
+  run_all.py          run the steps above in order
+  splits/             split manifests (tracked in git)
+  runs/               training outputs (git-ignored)
+  results/            audit and evaluation tables
+```
+
+## Data
+
+`dataset/` (git-ignored, Zenodo record 10.5281/zenodo.18457879) contains:
+- `singlecell/train` (236), `singlecell/test` (46);
+- `spheroid/train` (129), `spheroid/test` (28).
+
+The released single-cell model was trained on a larger split (`all_tune`: 280 training / 81 validation pairs, 1, 6 and 10 kPa). The paper also reports independent test sets. `prepare_dataset.py` reconstructs all of this from the original data tree and adds the missing images to `dataset/`:
+
+```
+dataset/singlecell_extra/<set>/<sample>/      all_tune images that are not in singlecell/ (1 and 10 kPa)
+dataset/singlecell_external/<set>/<sample>/   independent test sets (DMSO, Y-27632, TGF-beta, fibronectin, ...)
+```
+
+Each manifest in `splits/` lists one sample per row:
+- the sample folder relative to `dataset/`;
+- condition, Young's modulus and TFM settings: colour limit of the heatmap (`ythreshold_Pa`, 500 Pa for single cells and 200 Pa for spheroids) and pixel-size setting;
+- the substrate setting the original code fed to the network (`substrate_legacy`) and the nominal one (`substrate_nominal`);
+- the MD5 of the bright-field image.
+
+## Steps
+
+```bash
+# 1. audit only: what is in dataset/ vs the original tree, all_tune composition, leakage check
+python -m experiments.run_all --legacy-root <path>/cellforce/dataset --steps audit
+
+# 2. complete dataset/ and write splits/*.csv (then commit experiments/splits/)
+python -m experiments.run_all --legacy-root <path>/cellforce/dataset --steps prepare
+
+# 3. released checkpoints, substrate sweep, input-free baseline, image-only stage-1 model
+python -m experiments.run_all --steps eval_released --stage1 <path>/outputs/checkpoints/last/last_checkpoint_final.pth
+
+# 4. baselines / ablations (about 7-8 h each on an 11 GB GPU; resumable), then their evaluation
+python -m experiments.run_all --steps train,eval_runs --variants s2f_full,unet,s2f_no_gan,mask_only,s2f_no_substrate
+```
+
+A single run: `python -m experiments.train --variant unet --seed 0`.
+
+### GPUs and run time
+
+- `--steps train` runs one job per GPU in parallel (default `--gpus 0,1`). Each job sees only its own GPU (`CUDA_VISIBLE_DEVICES`), and the next variant starts as soon as a GPU is free.
+- A model is never split across GPUs: with batch size 2, that would harm batch-norm/GAN training.
+- `train.py` refuses to run without CUDA unless `--allow-cpu` is given.
+- Augmentation runs in `--num-workers` DataLoader processes (default 2 per job; use 0 if RAM is short).
+- Finished runs are skipped and interrupted runs resume from `last_checkpoint.pth`, so the command can simply be re-run.
+- Run time on an 11 GB GPU (2080 Ti class) is about 4-5 min per epoch, i.e. about 7-8 h per variant (100 epochs). On two GPUs:
+  - the 5 core variants take about 16-24 h;
+  - all 14 variants take about 55 h.
+- Progress is logged to `experiments/runs/<variant>_seed<seed>/train.log` and `history.json`.
+
+For a quick check that everything runs, add `--epochs 1 --max-batches 2` (and `--allow-cpu` without a GPU).
+
+## Training protocol (all variants)
+
+Every variant is trained from scratch on `alltune_train` (280 pairs) with the same settings:
+- 100 epochs, batch size 2, 1024 × 1024 px, the augmentations of `data/augmentations.py`;
+- Adam (lr 1e-4, β = 0.5/0.999), cosine decay after 10 % linear warm-up;
+- loss = 10 × per-image WFM-RME + 1 × adversarial (PatchGAN, BCE) + 100 × MSE of mean intensity.
+
+The checkpoint with the lowest reconstruction loss on `alltune_val` (81 pairs) is kept. Variants change one component each (`configs.py`):
+
+| variant | change |
+|---|---|
+| `s2f_full` | none (from-scratch counterpart of the released model) |
+| `unet` | plain U-Net (Ronneberger et al. 2015), no attention modules, no discriminator |
+| `unet_gan` | plain U-Net + PatchGAN (pix2pix-style) |
+| `mask_only` | binary cell mask instead of bright field (needs `cell_mask.tif` in each sample folder) |
+| `s2f_no_gan` | no discriminator (λ_GAN = 0) |
+| `s2f_no_force` | no stress-consistency term (λ_force = 0) |
+| `s2f_gan_0.1`, `s2f_gan_5` | λ_GAN = 0.1 / 5 |
+| `s2f_no_substrate` | bright field only (no substrate channels) |
+| `s2f_nominal_substrate` | substrate channels from the nominal Young's modulus of each image |
+| `s2f_no_multiscale` | no multi-scale input module |
+| `s2f_bridge_no_cbam` | bottleneck without CBAM |
+| `s2f_bridge_plain` | bottleneck without CBAM, Global Context Module and hierarchical attention |
+| `s2f_no_attention_gates` | skip connections without attention gates |
+
+The released checkpoint was obtained by staged fine-tuning (see Methods). `evaluate.py` reports it alongside the from-scratch variants.
+
+## Metrics
+
+`evaluate.py` computes every metric per image and reports the mean and SD across images:
+- MSE, MS-SSIM and pixel-wise Pearson r;
+- per-image WFM-RME;
+- mean traction stress in Pa (heatmap value × colour limit);
+- the fraction of saturated ground-truth pixels;
+- inference time.
+
+Its outputs are `results/eval/per_image.csv`, `summary.csv` and `model_params.csv`.

experiments/ablation_models.py

@@ -0,0 +1,74 @@
+"""
+Generators for the baseline / ablation study. All variants output tanh in [-1, 1], like the
+released S2F generator, so they share one training and evaluation path.
+
+  arch='s2f'   S2FGenerator (models/s2f_model.py) with optional switches:
+               use_multi_scale_input  (False -> single 3x3 input conv)
+               bridge                 'cbam' (paper) | 'cnn' (no CBAM) | 'plain' (dilated conv block only:
+                                      no CBAM, no Global Context Module, no hierarchical attention)
+               attention_gates        False -> skip connections pass encoder features unchanged
+  arch='unet'  plain U-Net (Ronneberger et al., 2015): 4 levels, 64-1024 channels, double 3x3 conv + BN +
+               ReLU, transposed-conv upsampling, concatenation skips; no attention modules.
+"""
+import torch
+import torch.nn as nn
+
+from models.blocks import ResidualBlock
+from models.s2f_model import PatchGANDiscriminator, S2FGenerator
+
+
+class _PassSkip(nn.Module):
+    """Replaces an attention gate: returns the encoder skip feature unchanged."""
+    def forward(self, g, x):
+        return x
+
+
+def _double_conv(i, o):
+    return nn.Sequential(nn.Conv2d(i, o, 3, padding=1, bias=False), nn.BatchNorm2d(o), nn.ReLU(inplace=True),
+                         nn.Conv2d(o, o, 3, padding=1, bias=False), nn.BatchNorm2d(o), nn.ReLU(inplace=True))
+
+
+class UNet(nn.Module):
+    def __init__(self, in_channels=1, out_channels=1, base=64):
+        super().__init__()
+        c = [base, base * 2, base * 4, base * 8, base * 16]
+        self.downs = nn.ModuleList([_double_conv(in_channels, c[0])] + [_double_conv(c[i], c[i + 1]) for i in range(3)])
+        self.pool = nn.MaxPool2d(2)
+        self.bottleneck = _double_conv(c[3], c[4])
+        self.ups = nn.ModuleList([nn.ConvTranspose2d(c[i + 1], c[i], 2, stride=2) for i in reversed(range(4))])
+        self.decs = nn.ModuleList([_double_conv(c[i] * 2, c[i]) for i in reversed(range(4))])
+        self.out = nn.Sequential(nn.Conv2d(c[0], out_channels, 1), nn.Tanh())
+
+    def forward(self, x):
+        skips = []
+        for d in self.downs:
+            x = d(x); skips.append(x); x = self.pool(x)
+        x = self.bottleneck(x)
+        for up, dec, s in zip(self.ups, self.decs, reversed(skips)):
+            x = dec(torch.cat([up(x), s], dim=1))
+        return self.out(x)
+
+
+def build_generator(arch="s2f", in_channels=3, use_multi_scale_input=True, bridge="cbam",
+                    attention_gates=True, **_):
+    if arch == "unet":
+        return UNet(in_channels=in_channels)
+    if arch != "s2f":
+        raise ValueError(f"unknown arch {arch}")
+    g = S2FGenerator(in_channels=in_channels, out_channels=1, img_size=1024,
+                     bridge_type="cnn" if bridge in ("cnn", "plain") else "cbam",
+                     use_multi_scale_input=use_multi_scale_input)
+    if bridge == "plain":
+        g.bridge = nn.Sequential(nn.Conv2d(512, 1024, 3, padding=2, dilation=2), nn.BatchNorm2d(1024),
+                                 nn.ReLU(inplace=True), ResidualBlock(1024, 1024))
+    if not attention_gates:
+        g.att4, g.att3, g.att2, g.att1 = _PassSkip(), _PassSkip(), _PassSkip(), _PassSkip()
+    return g
+
+
+def build_discriminator(in_channels=3):
+    return PatchGANDiscriminator(in_channels=in_channels + 1, ndf=64, n_layers=3)
+
+
+def n_params(module):
+    return sum(p.numel() for p in module.parameters())

experiments/configs.py

@@ -0,0 +1,66 @@
+"""
+Training protocol shared by every baseline / ablation run, and the variants.
+
+All variants are trained FROM SCRATCH with the same protocol, split, number of epochs and
+checkpoint-selection rule, so that differences are attributable to the ablated component only.
+(The released single-cell checkpoint was obtained by staged fine-tuning, see the Methods; the
+'s2f_full' variant is its from-scratch counterpart under this protocol.)
+"""
+import copy
+
+PROTOCOL = dict(
+    train_manifest="alltune_train",      # 280 pairs, released model's training split
+    val_manifest="alltune_val",          # 81 pairs, used for checkpoint selection only
+    epochs=100,
+    batch_size=2,
+    img_size=1024,
+    augment=True,                        # data/augmentations.AdvancedAugmentations
+    g_lr=1e-4, d_lr=1e-4, beta1=0.5, beta2=0.999,
+    warmup_frac=0.1,                     # cosine schedule with linear warm-up
+    lambda_rec=10.0,                     # per-image WFM-RME
+    lambda_gan=1.0,                      # vanilla (BCE) PatchGAN loss
+    lambda_force=100.0,                  # MSE between mean predicted and mean target intensity
+    arch="s2f", in_channels=3, use_multi_scale_input=True, bridge="cbam", attention_gates=True,
+    substrate_mode="legacy",             # see experiments/manifest_data.py
+    input_mode="bf",
+    seed=0,
+)
+
+VARIANTS = {
+    # reference
+    "s2f_full":               dict(),
+    # baselines
+    "unet":                   dict(arch="unet", lambda_gan=0.0),          # plain U-Net, no attention, no discriminator
+    "unet_gan":               dict(arch="unet"),                          # pix2pix-style U-Net + PatchGAN
+    "mask_only":              dict(input_mode="mask"),                    # binary cell mask instead of bright field
+    # loss ablations
+    "s2f_no_gan":             dict(lambda_gan=0.0),
+    "s2f_no_force":           dict(lambda_force=0.0),
+    "s2f_gan_0.1":            dict(lambda_gan=0.1),
+    "s2f_gan_5":              dict(lambda_gan=5.0),
+    # input ablations
+    "s2f_no_substrate":       dict(in_channels=1, substrate_mode="none"),
+    "s2f_nominal_substrate":  dict(substrate_mode="nominal"),
+    # architecture ablations
+    "s2f_no_multiscale":      dict(use_multi_scale_input=False),
+    "s2f_bridge_no_cbam":     dict(bridge="cnn"),
+    "s2f_bridge_plain":       dict(bridge="plain"),
+    "s2f_no_attention_gates": dict(attention_gates=False),
+}
+
+# Recommended order if GPU time is limited (about 7-8 h per run of 100 epochs on an 11 GB GPU)
+PRIORITY = ["s2f_full", "unet", "s2f_no_gan", "mask_only", "s2f_no_substrate",
+            "s2f_no_multiscale", "s2f_bridge_plain", "s2f_no_force", "s2f_no_attention_gates",
+            "s2f_bridge_no_cbam", "s2f_nominal_substrate", "unet_gan", "s2f_gan_0.1", "s2f_gan_5"]
+
+
+def get_config(variant, **overrides):
+    if variant not in VARIANTS:
+        raise KeyError(f"unknown variant '{variant}'. Options: {', '.join(VARIANTS)}")
+    cfg = copy.deepcopy(PROTOCOL)
+    cfg.update(VARIANTS[variant])
+    cfg.update({k: v for k, v in overrides.items() if v is not None})
+    cfg["variant"] = variant
+    if cfg["substrate_mode"] == "none":
+        cfg["in_channels"] = 1
+    return cfg

experiments/evaluate.py

@@ -0,0 +1,202 @@
+"""
+Evaluate released checkpoints, baselines and trained ablation runs on every split manifest.
+
+    python -m experiments.evaluate --released --meanmap               # released models + input-free baseline
+    python -m experiments.evaluate --released --sweep                  # + substrate-setting sensitivity
+    python -m experiments.evaluate --stage1 <path>/last_checkpoint_final.pth   # image-only stage-1 model
+    python -m experiments.evaluate --runs                              # every experiments/runs/*/best_checkpoint.pth
+
+Metrics are computed PER IMAGE (mean and SD are then taken across images):
+    mse, ms_ssim, pixel_r (Pearson), rme (per-image WFM-RME), mean_gt / mean_pred (heatmap units),
+    mean_gt_Pa / mean_pred_Pa (x colour-scale limit of the set, 500 Pa single cells, 200 Pa spheroids),
+    gt_frac_saturated (pixels >= 0.98 of the colour limit), inference time.
+Outputs (experiments/results/eval/ by default; existing rows of other models are kept):
+    per_image.csv, summary.csv
+"""
+import argparse
+import glob
+import os
+import time
+
+import numpy as np
+import pandas as pd
+import torch
+
+from experiments.ablation_models import build_generator, n_params
+from experiments.layout import DEFAULT_DATASET_ROOT, RESULTS_DIR, RUNS_DIR, S2F_ROOT, list_manifests
+from experiments.losses import PerImageWFMRMELoss
+from experiments.manifest_data import make_loader
+from models.s2f_model import create_s2f_model, create_settings_channels
+from utils.metrics import TorchMetricsWrapper
+from utils.substrate_settings import compute_settings_normalization
+
+CONFIG_PATH = os.path.join(S2F_ROOT, "config", "substrate_settings.json")
+SWEEP_KEYS = ["PDMS_1kPa_3um", "fibroblasts_PDMS", "PDMS_10kPa_3um"]
+
+
+def released_models(device):
+    out = {}
+    g, _ = create_s2f_model(in_channels=3, model_type="s2f")
+    g.load_state_dict(torch.load(os.path.join(S2F_ROOT, "ckp/single_cell/ckp_singlecell_GN.pth"), map_location="cpu",
+                                 weights_only=False)["generator_state_dict"], strict=True)
+    out["released_singlecell"] = dict(model=g.to(device).eval(), kind="singlecell", in_channels=3, tanh=True)
+    g, _ = create_s2f_model(in_channels=1, model_type="s2f")
+    g.load_state_dict(torch.load(os.path.join(S2F_ROOT, "ckp/spheroid/ckp_spheroid_GN.pth"), map_location="cpu",
+                                 weights_only=False)["generator_state_dict"], strict=True)
+    out["released_spheroid_GN"] = dict(model=g.to(device).eval(), kind="spheroid", in_channels=1, tanh=True)
+    fn = os.path.join(S2F_ROOT, "ckp/spheroid/ckp_spheroid_FN.pth")
+    if os.path.exists(fn):
+        g, _ = create_s2f_model(in_channels=1, model_type="s2f_spheroid")
+        g.load_state_dict(torch.load(fn, map_location="cpu", weights_only=False)["model_state_dict"], strict=True)
+        g.set_output_mode(use_tanh=False)
+        out["supervised_spheroid_FN"] = dict(model=g.to(device).eval(), kind="spheroid", in_channels=1, tanh=False)
+    return out
+
+
+def stage1_model(path, device):
+    g, _ = create_s2f_model(in_channels=1, model_type="s2f")
+    sd = torch.load(path, map_location="cpu", weights_only=False)["generator_state_dict"]
+    g.load_state_dict({k.replace("denoise_module", "initial_conv"): v for k, v in sd.items()}, strict=True)
+    return dict(model=g.to(device).eval(), kind="singlecell", in_channels=1, tanh=True)
+
+
+def run_models(device, pattern):
+    out = {}
+    for p in sorted(glob.glob(pattern)):
+        ck = torch.load(p, map_location="cpu", weights_only=False)
+        cfg = ck["config"]
+        g = build_generator(**cfg); g.load_state_dict(ck["generator_state_dict"], strict=True)
+        name = os.path.basename(os.path.dirname(p))
+        out[name] = dict(model=g.to(device).eval(), kind="singlecell", in_channels=cfg["in_channels"], tanh=True,
+                         substrate_mode=cfg["substrate_mode"], input_mode=cfg["input_mode"],
+                         epoch=ck.get("epoch"), val_loss=ck.get("val_loss"))
+    return out
+
+
+def mean_map(manifest, dataset_root):
+    acc, n = None, 0
+    for b in make_loader(manifest, dataset_root, batch_size=8, substrate_mode="none"):
+        acc = b[1].sum(0, keepdim=True) if acc is None else acc + b[1].sum(0, keepdim=True); n += b[1].size(0)
+    return acc / n
+
+
+def evaluate_one(name, spec, manifest, dataset_root, device, mw, norm, substrate_mode, batch, mean_pred=None):
+    input_mode = spec.get("input_mode", "bf")
+    use_settings = spec.get("in_channels", 1) == 3
+    sm = substrate_mode if use_settings else "none"
+    forced = None
+    if sm and sm.startswith("force:"):
+        forced, sm = sm.split(":", 1)[1], "reference"
+    try:
+        loader = make_loader(manifest, dataset_root, batch_size=batch, substrate_mode=sm, input_mode=input_mode)
+    except FileNotFoundError as e:
+        print(f"  [skip] {name} on {manifest}: {e}"); return []
+    rows_meta, rec, out, k = loader.dataset.rows, PerImageWFMRMELoss(), [], 0
+    with torch.no_grad():
+        for b in loader:
+            x, y = b[0].to(device), b[1].to(device)
+            t0 = time.time()
+            if mean_pred is not None:
+                p = mean_pred.to(device).expand_as(y)
+            else:
+                if use_settings:
+                    meta = b[4]
+                    if forced:
+                        meta = dict(meta, substrate=[forced] * x.size(0))
+                    x = torch.cat([x, create_settings_channels(meta, norm, device, x.shape, config_path=CONFIG_PATH)], 1)
+                p = spec["model"](x)
+                p = (p + 1) / 2 if spec.get("tanh", True) else p
+            dt = (time.time() - t0) / x.size(0)
+            for i in range(y.size(0)):
+                r = rows_meta[k]; k += 1
+                t, q = y[i:i + 1], p[i:i + 1]
+                yt = float(r.get("ythreshold_Pa") or (200 if "spheroid" in manifest else 500))
+                out.append({"model": name, "manifest": manifest, "substrate_mode": substrate_mode if use_settings else "n/a",
+                            "sample_dir": r["sample_dir"], "set": r.get("set", ""), "condition": r.get("condition", ""),
+                            "mse": mw.compute_mse(t, q), "ms_ssim": mw.compute_ms_ssim(t, q),
+                            "pixel_r": float(np.corrcoef(t.flatten().cpu().numpy(), q.flatten().cpu().numpy())[0, 1]),
+                            "rme": float(rec(q, t)), "mean_gt": float(t.mean()), "mean_pred": float(q.mean()),
+                            "mean_gt_Pa": float(t.mean()) * yt, "mean_pred_Pa": float(q.mean()) * yt,
+                            "gt_max": float(t.max()), "gt_frac_saturated": float((t >= 0.98).float().mean()),
+                            "seconds_per_image": dt})
+    return out
+
+
+def summarise(df):
+    g = df.groupby(["model", "manifest", "substrate_mode"])
+    s = g.agg(n=("mse", "size"), **{f"{m}_{f}": (m, f) for m in ["mse", "ms_ssim", "pixel_r", "rme", "mean_gt_Pa",
+                                                                   "mean_pred_Pa", "gt_frac_saturated"] for f in ["mean", "std"]},
+              seconds_per_image=("seconds_per_image", "mean")).reset_index()
+    s["pred_over_gt"] = s["mean_pred_Pa_mean"] / s["mean_gt_Pa_mean"]
+    return s
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--released", action="store_true", help="released single-cell + spheroid checkpoints")
+    ap.add_argument("--sweep", action="store_true", help="released single-cell model with 1 / 6 / 10 kPa settings")
+    ap.add_argument("--meanmap", action="store_true", help="input-free baseline (mean training map)")
+    ap.add_argument("--stage1", default=None, help="path to last_checkpoint_final.pth (image-only stage 1)")
+    ap.add_argument("--runs", nargs="?", const=os.path.join(RUNS_DIR, "*", "best_checkpoint.pth"), default=None,
+                    help="glob of trained run checkpoints (default experiments/runs/*/best_checkpoint.pth)")
+    ap.add_argument("--manifests", default="", help="comma-separated manifest names (default: all in experiments/splits)")
+    ap.add_argument("--dataset-root", default=DEFAULT_DATASET_ROOT)
+    ap.add_argument("--out", default=os.path.join(RESULTS_DIR, "eval"))
+    ap.add_argument("--batch", type=int, default=2)
+    ap.add_argument("--device", default=None)
+    a = ap.parse_args()
+    device = torch.device(a.device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    mans = [m for m in (a.manifests.split(",") if a.manifests else list_manifests()) if m]
+    sc_mans = [m for m in mans if not m.startswith("spheroid") and m != "alltune_train" and m != "zenodo_v1_train"]
+    sph_mans = [m for m in mans if m.startswith("spheroid") and m != "spheroid_train"]
+    print(f"device {device}; single-cell manifests {sc_mans}; spheroid manifests {sph_mans}")
+    mw, norm = TorchMetricsWrapper(device=device), compute_settings_normalization(config_path=CONFIG_PATH)
+
+    models = {}
+    if a.released:
+        models.update(released_models(device))
+    if a.stage1:
+        models["stage1_image_only"] = stage1_model(a.stage1, device)
+    if a.runs:
+        models.update(run_models(device, a.runs))
+    rows, params = [], []
+    for name, spec in models.items():
+        params.append({"model": name, "n_params_M": n_params(spec["model"]) / 1e6, "epoch": spec.get("epoch"),
+                       "val_loss": spec.get("val_loss")})
+        targets = sc_mans if spec["kind"] == "singlecell" else sph_mans
+        modes = [spec.get("substrate_mode", "legacy")]
+        if name == "released_singlecell":
+            modes = ["legacy", "nominal"] + ([f"force:{k}" for k in SWEEP_KEYS] if a.sweep else [])
+        for m in targets:
+            for sm in modes:
+                print(f"  {name:28s} {m:28s} {sm}", flush=True)
+                rows += evaluate_one(name, spec, m, a.dataset_root, device, mw, norm, sm, a.batch)
+    if a.meanmap:
+        for train_m, targets, name in [("alltune_train", sc_mans, "baseline_mean_map_singlecell"),
+                                       ("spheroid_train", sph_mans, "baseline_mean_map_spheroid")]:
+            if train_m not in list_manifests():
+                continue
+            M = mean_map(train_m, a.dataset_root)
+            for m in targets:
+                print(f"  {name:28s} {m}", flush=True)
+                rows += evaluate_one(name, {"in_channels": 1}, m, a.dataset_root, device, mw, norm, "n/a", a.batch, mean_pred=M)
+    os.makedirs(a.out, exist_ok=True)
+    new = pd.DataFrame(rows)
+    pi = os.path.join(a.out, "per_image.csv")
+    if os.path.exists(pi) and len(new):           # keep rows of models not re-evaluated now
+        old = pd.read_csv(pi)
+        new = pd.concat([old[~old.model.isin(new.model.unique())], new], ignore_index=True)
+    if len(new):
+        new.to_csv(pi, index=False)
+        summarise(new).to_csv(os.path.join(a.out, "summary.csv"), index=False)
+    pp = os.path.join(a.out, "model_params.csv")
+    if params:
+        P = pd.DataFrame(params)
+        if os.path.exists(pp):
+            old = pd.read_csv(pp); P = pd.concat([old[~old.model.isin(P.model)], P], ignore_index=True)
+        P.to_csv(pp, index=False)
+    print(f"wrote {pi}, summary.csv, model_params.csv in {a.out}")
+
+
+if __name__ == "__main__":
+    main()

experiments/layout.py

@@ -0,0 +1,150 @@
+"""
+Dataset layout, per-set metadata and manifest helpers for the S2F experiments.
+
+Two dataset trees are involved:
+
+* the **repository dataset** ``S2F/dataset/`` (git-ignored; same content as the Zenodo deposit):
+      singlecell/{train,test}/<sample>/      236 / 46 fibroblasts on 6 kPa collagen gels
+      spheroid/{train,test}/<sample>/        129 / 28 LS174T spheroids
+  ``experiments/prepare_dataset.py --write`` adds the remaining sets used in the paper:
+      singlecell_extra/<set>/<sample>/       training/validation images of the released model
+                                             that are not in singlecell/ (1 and 10 kPa)
+      singlecell_external/<set>/<sample>/    independent test sets (drugs, fibronectin, ...)
+
+* the **legacy tree** used to develop the model (``cellforce/dataset`` on the workstation), whose
+  ``single_cell/all_tune/{train,test}`` folder (280 / 81 pairs) is the split of the released
+  single-cell checkpoint.
+
+Splits are defined by manifest CSV files in ``experiments/splits/`` (tracked in git), one row per
+sample, so that every experiment is reproducible from the repository dataset alone.
+"""
+import csv
+import hashlib
+import os
+
+S2F_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+# Override with environment variables if needed (e.g. dataset on another disk)
+DEFAULT_DATASET_ROOT = os.environ.get("S2F_DATASET_ROOT", os.path.join(S2F_ROOT, "dataset"))
+SPLITS_DIR = os.environ.get("S2F_SPLITS_DIR", os.path.join(S2F_ROOT, "experiments", "splits"))
+RUNS_DIR = os.environ.get("S2F_RUNS_DIR", os.path.join(S2F_ROOT, "experiments", "runs"))
+RESULTS_DIR = os.environ.get("S2F_RESULTS_DIR", os.path.join(S2F_ROOT, "experiments", "results"))
+
+BF_SUFFIX = "BF_001.tif"
+HM_SUFFIX = "_gray.jpg"
+MASK_NAMES = ("cell_mask.tif", "cell_mask.png", "mask.tif", "mask.png")  # optional, for the mask-only baseline
+
+# TFM reconstruction settings (from the MATLAB settings files of each set).
+# ythreshold_Pa = upper limit of the colour scale used to export the heatmaps -> grey 1.0 == ythreshold_Pa.
+# tfm_pixelsize = value of Settings.PixelSize used in the TFM reconstruction.
+FIBRO_6KPA = dict(cell_type="HFF", coating="collagen I", young_Pa=6000, tfm_pixelsize=3.0769, ythreshold_Pa=500)
+
+# Legacy sets: key -> (relative path under legacy root, role, clean name, metadata)
+# role: 'alltune'   split of the released single-cell model (train 280 / val 81)
+#       'published' already in S2F/dataset
+#       'source'    pool from which all_tune was assembled (1 / 10 kPa)
+#       'external'  independent test set
+#       'union'     duplicate collection, audited only
+LEGACY_SETS = {
+    "all_tune/train": ("single_cell/all_tune/train", "alltune", None, {}),
+    "all_tune/test": ("single_cell/all_tune/test", "alltune", None, {}),
+    "fibro_split/train": ("single_cell/fibro_split/train", "published", "fibro_6kPa_train", dict(FIBRO_6KPA, condition="6 kPa collagen")),
+    "fibro_split/test": ("single_cell/fibro_split/test", "published", "fibro_6kPa_setVI", dict(FIBRO_6KPA, condition="6 kPa collagen")),
+    "1kpa_new/train": ("single_cell/1kpa_new/train", "source", "fibro_1kPa_train", dict(FIBRO_6KPA, young_Pa=1000, condition="1 kPa collagen")),
+    "1kpa_new/test": ("single_cell/1kpa_new/test", "source", "fibro_1kPa_setI", dict(FIBRO_6KPA, young_Pa=1000, condition="1 kPa collagen")),
+    "10kpa_new/train": ("single_cell/10kpa_new/train", "source", "fibro_10kPa_train", dict(FIBRO_6KPA, young_Pa=10000, condition="10 kPa collagen")),
+    "10kpa_new/test": ("single_cell/10kpa_new/test", "source", "fibro_10kPa_setI", dict(FIBRO_6KPA, young_Pa=10000, condition="10 kPa collagen")),
+    "1kPa_SET_III": ("single_cell/last_new_1kpa_10kpa/Heatmaps_01kPa_SET_III_for_Kaveh", "external", "fibro_1kPa_setIII", dict(FIBRO_6KPA, young_Pa=1000, condition="1 kPa collagen")),
+    "DMSO_SET_II": ("single_cell/last_new_1kpa_10kpa/Heatmaps_DMSO_SET_II_for_Kaveh", "external", "fibro_DMSO_setII", dict(FIBRO_6KPA, condition="DMSO")),
+    "Y27632_SET_II": ("single_cell/last_new_1kpa_10kpa/Heatmaps_Y27632_SET_II_for_Kaveh", "external", "fibro_Y27632_setII", dict(FIBRO_6KPA, condition="Y-27632")),
+    "Control_SET_I": ("single_cell/Drugs/Fibros_Control_SET_I_for_Kaveh", "external", "fibro_DMSO_setI", dict(FIBRO_6KPA, condition="DMSO")),
+    "Fibronectin_SET_I": ("single_cell/Drugs/Fibros_Fibronectin_SET_I_for_Kaveh", "external", "fibro_fibronectin_setI", dict(FIBRO_6KPA, coating="fibronectin", condition="fibronectin")),
+    "TGFB_SET_I": ("single_cell/Drugs/Fibros_TGFB_SETI_for_Kaveh", "external", "fibro_TGFb_setI", dict(FIBRO_6KPA, condition="TGF-beta")),
+    "Y27632_SET_I": ("single_cell/Drugs/Fibros_Y27632_SET_I_for_Kaveh", "external", "fibro_Y27632_setI", dict(FIBRO_6KPA, condition="Y-27632")),
+    "U2OS/test": ("single_cell/u2os/test", "external", "u2os_6kPa", dict(cell_type="U2OS", coating="collagen I", young_Pa=6000, tfm_pixelsize=6.1538, ythreshold_Pa=500, condition="U2OS")),
+    "Drugs/all": ("single_cell/Drugs/all", "union", None, {}),
+    "spheroid/train": ("spheroid/train", "published", "spheroid_train", dict(cell_type="LS174T spheroid", coating="collagen I", ythreshold_Pa=200, condition="spheroid")),
+    "spheroid/test": ("spheroid/test", "published", "spheroid_val", dict(cell_type="LS174T spheroid", coating="collagen I", ythreshold_Pa=200, condition="spheroid")),
+}
+
+# Published repository dataset (S2F/dataset), relative paths
+PUBLISHED_SETS = {
+    "singlecell/train": "singlecell/train",
+    "singlecell/test": "singlecell/test",
+    "spheroid/train": "spheroid/train",
+    "spheroid/test": "spheroid/test",
+}
+
+MANIFEST_FIELDS = ["sample_dir", "split", "set", "condition", "cell_type", "coating", "young_Pa",
+                   "tfm_pixelsize", "ythreshold_Pa", "legacy_folder", "legacy_substrate_code",
+                   "substrate_legacy", "substrate_nominal", "bf_md5"]
+
+
+def legacy_substrate_code(folder_name):
+    """Substrate code exactly as the original training code derived it (folder_name[:4]; 'kpan*' -> [:5])."""
+    code = folder_name[:4]
+    return folder_name[:5] if code.startswith("kpan") else code
+
+
+def legacy_code_to_config_key(code):
+    """S2F config key that encodes the same (pixel size, Young's modulus) as the original code's substrate code."""
+    return {"kpa0": "PDMS_1kPa", "kpa1": "PDMS_10kPa", "kpan0": "PDMS_1kPa_3um",
+            "kpan1": "PDMS_10kPa_3um", "u2os": "U2OS_PDMS"}.get(code, "fibroblasts_PDMS")
+
+
+def nominal_config_key(young_Pa, tfm_pixelsize, cell_type="HFF"):
+    """S2F config key for the nominal Young's modulus and TFM pixel-size setting of a sample."""
+    if cell_type == "U2OS":
+        return "U2OS_PDMS"
+    if young_Pa in (None, ""):
+        return None
+    young_Pa, px = float(young_Pa), float(tfm_pixelsize or 3.0769)
+    if young_Pa == 6000:
+        return "fibroblasts_PDMS"
+    suffix = "_3um" if abs(px - 3.0769) < 1e-3 else ""
+    return {1000: "PDMS_1kPa", 10000: "PDMS_10kPa"}.get(int(young_Pa), "fibroblasts_PDMS") + suffix
+
+
+def sample_folders(root):
+    """Yield (folder_name, bf_path, hm_path) for every direct sub-folder holding a BF image and a gray heatmap."""
+    if not os.path.isdir(root):
+        return
+    for sub in sorted(os.listdir(root)):
+        p = os.path.join(root, sub)
+        if not os.path.isdir(p):
+            continue
+        files = os.listdir(p)
+        bf = [f for f in files if f.endswith(BF_SUFFIX)]
+        hm = [f for f in files if f.endswith(HM_SUFFIX)]
+        if bf and hm:
+            yield sub, os.path.join(p, bf[0]), os.path.join(p, hm[0])
+
+
+def md5(path, chunk=1 << 20):
+    h = hashlib.md5()
+    with open(path, "rb") as f:
+        for b in iter(lambda: f.read(chunk), b""):
+            h.update(b)
+    return h.hexdigest()
+
+
+def read_manifest(path):
+    with open(path, newline="") as f:
+        return list(csv.DictReader(f))
+
+
+def write_manifest(path, rows):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, "w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=MANIFEST_FIELDS, extrasaction="ignore")
+        w.writeheader()
+        w.writerows(rows)
+
+
+def manifest_path(name):
+    return os.path.join(SPLITS_DIR, f"{name}.csv")
+
+
+def list_manifests():
+    if not os.path.isdir(SPLITS_DIR):
+        return []
+    return sorted(os.path.splitext(f)[0] for f in os.listdir(SPLITS_DIR) if f.endswith(".csv"))

experiments/losses.py

@@ -0,0 +1,23 @@
+"""
+Reconstruction loss used for the released single-cell model.
+
+Per-image weighted force-magnitude relative error (WFM-RME), as in the manuscript and as in the
+loss that trained the released checkpoint (Archive/src/utils/losses.WFMRelativeMagnitudeLoss):
+the weight of each pixel is its ground-truth magnitude divided by the mean ground-truth magnitude
+*of the same image* (eps = 1e-6).  Note: utils.metrics.WFMRMELoss normalises by the mean over the
+whole batch instead; the two coincide only for batch size 1.
+"""
+import torch
+import torch.nn as nn
+
+
+class PerImageWFMRMELoss(nn.Module):
+    def __init__(self, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, pred, target):
+        t, p = target[:, 0], pred[:, 0]                       # (B, H, W) magnitudes
+        rel = (p - t).abs() / (t + self.eps)
+        w = t / (t.mean(dim=(1, 2), keepdim=True) + self.eps)
+        return (rel * w).mean()

experiments/manifest_data.py

@@ -0,0 +1,92 @@
+"""
+Manifest-based data loading. Uses exactly the preprocessing of the S2F package
+(``data.cell_dataset.load_image`` -> grayscale, resize to 1024, /255; ``ImageDataset``;
+``AdvancedAugmentations``), so results are comparable with the released models.
+
+substrate_mode (single-cell only) decides which substrate setting builds the two substrate channels:
+  'legacy'    what the original training code fed (derived from the all_tune folder name) - default,
+              reproduces the conditions under which the released model was trained
+  'nominal'   the true Young's modulus / TFM pixel-size setting of each sample
+  'reference' the 6 kPa reference setting (fibroblasts_PDMS) for every sample
+  'none'      no metadata (models without substrate channels)
+input_mode:
+  'bf'        bright-field image (default)
+  'mask'      binary cell mask (file cell_mask.tif/png in the sample folder) - mask-only baseline
+"""
+import os
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+from data.cell_dataset import ImageDataset, load_image
+from experiments.layout import DEFAULT_DATASET_ROOT, MASK_NAMES, manifest_path, read_manifest
+
+
+def _find(folder, suffix=None, names=None):
+    for f in sorted(os.listdir(folder)):
+        if (suffix and f.endswith(suffix)) or (names and f in names):
+            return os.path.join(folder, f)
+    return None
+
+
+def load_manifest_pairs(manifest, dataset_root=DEFAULT_DATASET_ROOT, img_size=1024,
+                        substrate_mode="legacy", input_mode="bf"):
+    rows = read_manifest(manifest if os.path.isfile(manifest) else manifest_path(manifest))
+    pairs = []
+    for r in rows:
+        d = os.path.join(dataset_root, r["sample_dir"])
+        bf = _find(d, suffix="BF_001.tif")
+        hm = _find(d, suffix="_gray.jpg")
+        if bf is None or hm is None:
+            raise FileNotFoundError(f"missing BF/heatmap in {d}")
+        if input_mode == "mask":
+            mp = _find(d, names=MASK_NAMES)
+            if mp is None:
+                raise FileNotFoundError(f"input_mode='mask' but no {MASK_NAMES} in {d}")
+            x = (load_image(mp, (img_size, img_size)) > 0.5).astype(np.float32)
+        else:
+            x = load_image(bf, (img_size, img_size))
+        y = load_image(hm, (img_size, img_size))
+        if substrate_mode == "none":
+            pairs.append((x, y, (0, 0)))
+            continue
+        sub = {"legacy": r.get("substrate_legacy"), "nominal": r.get("substrate_nominal"),
+               "reference": "fibroblasts_PDMS"}[substrate_mode] or "fibroblasts_PDMS"
+        pairs.append((x, y, (0, 0), {"folder_name": os.path.basename(d), "substrate": sub,
+                                     "root_folder": os.path.dirname(r["sample_dir"]), "set": r.get("set", "")}))
+    return pairs, rows
+
+
+def _worker_init(worker_id):
+    """Distinct, reproducible numpy seed per worker (augmentations use np.random). Module-level for Windows spawn."""
+    np.random.seed(torch.initial_seed() % 2 ** 32)
+
+
+class _MaskDataset(ImageDataset):
+    """Re-binarise the input after augmentation (intensity augmentations are meaningless for masks)."""
+    def __getitem__(self, idx):
+        out = list(super().__getitem__(idx))
+        out[0] = (out[0] > 0.5).float()
+        return tuple(out)
+
+
+def make_loader(manifest, dataset_root=DEFAULT_DATASET_ROOT, batch_size=2, img_size=1024, augment=False,
+                shuffle=False, substrate_mode="legacy", input_mode="bf", seed=None, num_workers=0):
+    pairs, rows = load_manifest_pairs(manifest, dataset_root, img_size, substrate_mode, input_mode)
+    transform = None
+    if augment:
+        from data.augmentations import AdvancedAugmentations
+        transform = AdvancedAugmentations((img_size, img_size))
+    cls = _MaskDataset if input_mode == "mask" else ImageDataset
+    ds = cls(pairs, transform=transform, channel_first=True, threshold=0.0,
+             return_metadata=substrate_mode != "none")
+    ds.rows = rows
+    g = None
+    if seed is not None:
+        g = torch.Generator(); g.manual_seed(seed)
+    if num_workers > 0:   # workers run the augmentations so the GPU is not left waiting
+        return DataLoader(ds, batch_size=batch_size, shuffle=shuffle, generator=g, num_workers=num_workers,
+                          persistent_workers=True, pin_memory=torch.cuda.is_available(),
+                          worker_init_fn=_worker_init)
+    return DataLoader(ds, batch_size=batch_size, shuffle=shuffle, generator=g, num_workers=0)

experiments/prepare_dataset.py

@@ -0,0 +1,237 @@
+"""
+Audit the legacy dataset tree against the repository dataset and (optionally) complete the
+repository dataset with every image used in the paper, plus split manifests.
+
+Step 1 - audit only (nothing is copied):
+    python -m experiments.prepare_dataset --legacy-root C:/Users/Kaveh/Documents/cellforce/dataset
+
+Step 2 - copy missing images into S2F/dataset and write the manifests:
+    python -m experiments.prepare_dataset --legacy-root C:/Users/Kaveh/Documents/cellforce/dataset --write
+
+What the audit answers (written to experiments/results/dataset_audit/):
+  counts.csv                 number of samples per legacy and repository set
+  published_vs_legacy.csv    is S2F/dataset/singlecell (236/46) identical to legacy fibro_split? spheroids?
+  alltune_composition.csv    where every image of all_tune/train (280) and all_tune/test (81) comes from
+  missing_from_published.csv all_tune images that are NOT in S2F/dataset (i.e. not in the Zenodo v1 deposit)
+  overlap_matrix.csv         shared bright-field images between every pair of sets (leakage check)
+  substrate_codes.csv        substrate code the original code derived from each all_tune folder name
+  report.txt                 human-readable summary
+Images are matched by the MD5 of the bright-field file, so renamed/copied folders are still recognised.
+
+With --write:
+  dataset/singlecell_extra/<set>/<folder>/      all_tune images missing from dataset/singlecell
+  dataset/singlecell_external/<set>/<folder>/   independent test sets
+  experiments/splits/*.csv                       manifests (commit these to git):
+      alltune_train.csv, alltune_val.csv        split of the released single-cell model (280 / 81)
+      zenodo_v1_train.csv, zenodo_v1_test.csv   published 6 kPa split (236 / 46)
+      spheroid_train.csv, spheroid_val.csv      spheroid split (129 / 28)
+      test_<set>.csv                            one per independent test set
+"""
+import argparse
+import collections
+import csv
+import os
+import shutil
+
+from experiments.layout import (DEFAULT_DATASET_ROOT, LEGACY_SETS, PUBLISHED_SETS, RESULTS_DIR,
+                                legacy_code_to_config_key, legacy_substrate_code, manifest_path,
+                                md5, nominal_config_key, sample_folders, write_manifest)
+
+
+def index_tree(root, rel):
+    """List samples of one set: [{folder, bf, hm, dir, md5}]."""
+    out = []
+    for folder, bf, hm in sample_folders(os.path.join(root, rel)):
+        out.append({"folder": folder, "bf": bf, "hm": hm, "dir": os.path.dirname(bf), "md5": md5(bf)})
+    return out
+
+
+def discover_unlisted(root, listed_rel, max_depth=4):
+    """Directories under root that contain sample folders but are not in LEGACY_SETS (so nothing is missed)."""
+    found = []
+    listed = {os.path.normpath(os.path.join(root, r)) for r in listed_rel}
+    base_depth = os.path.normpath(root).count(os.sep)
+    for dp, dn, _ in os.walk(root):
+        if os.path.normpath(dp).count(os.sep) - base_depth >= max_depth:
+            dn[:] = []
+            continue
+        if os.path.normpath(dp) in listed:
+            dn[:] = []
+            continue
+        n = sum(1 for _ in sample_folders(dp))
+        if n:
+            found.append((os.path.relpath(dp, root), n))
+            dn[:] = []
+    return found
+
+
+def row_for(sample_dir, split, set_name, meta, legacy_folder, bf_md5):
+    code = legacy_substrate_code(legacy_folder) if legacy_folder else ""
+    return {"sample_dir": sample_dir.replace(os.sep, "/"), "split": split, "set": set_name,
+            "condition": meta.get("condition", ""), "cell_type": meta.get("cell_type", ""),
+            "coating": meta.get("coating", ""), "young_Pa": meta.get("young_Pa", ""),
+            "tfm_pixelsize": meta.get("tfm_pixelsize", ""), "ythreshold_Pa": meta.get("ythreshold_Pa", ""),
+            "legacy_folder": legacy_folder, "legacy_substrate_code": code,
+            "substrate_legacy": legacy_code_to_config_key(code) if meta.get("cell_type") != "LS174T spheroid" else "",
+            "substrate_nominal": (nominal_config_key(meta.get("young_Pa"), meta.get("tfm_pixelsize"), meta.get("cell_type"))
+                                  if meta.get("cell_type") != "LS174T spheroid" else "") or "",
+            "bf_md5": bf_md5}
+
+
+def copy_sample(src_dir, dst_dir, expected_md5):
+    if os.path.isdir(dst_dir):
+        bf = [f for f in os.listdir(dst_dir) if f.endswith("BF_001.tif")]
+        if bf and md5(os.path.join(dst_dir, bf[0])) == expected_md5:
+            return "exists"
+        raise RuntimeError(f"{dst_dir} exists with different content; refusing to overwrite")
+    os.makedirs(dst_dir)
+    for f in os.listdir(src_dir):
+        if os.path.isfile(os.path.join(src_dir, f)):
+            shutil.copy2(os.path.join(src_dir, f), os.path.join(dst_dir, f))
+    return "copied"
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--legacy-root", required=True, help="old dataset folder (contains single_cell/ and spheroid/)")
+    ap.add_argument("--dataset-root", default=DEFAULT_DATASET_ROOT, help="repository dataset (default: S2F/dataset)")
+    ap.add_argument("--write", action="store_true", help="copy missing images and write split manifests")
+    ap.add_argument("--out", default=os.path.join(RESULTS_DIR, "dataset_audit"))
+    a = ap.parse_args()
+    os.makedirs(a.out, exist_ok=True)
+    rep = []
+
+    def say(s=""):
+        print(s); rep.append(s)
+
+    # ---------------- index ----------------
+    legacy = {k: index_tree(a.legacy_root, v[0]) for k, v in LEGACY_SETS.items()}
+    pub = {k: index_tree(a.dataset_root, v) for k, v in PUBLISHED_SETS.items()}
+    say(f"legacy root : {a.legacy_root}\ndataset root: {a.dataset_root}\n")
+    with open(os.path.join(a.out, "counts.csv"), "w", newline="") as f:
+        w = csv.writer(f); w.writerow(["tree", "set", "path", "n_samples", "n_unique_bf"])
+        for k, v in legacy.items():
+            w.writerow(["legacy", k, LEGACY_SETS[k][0], len(v), len({s['md5'] for s in v})])
+            say(f"  legacy    {k:20s} n={len(v):4d}")
+        for k, v in pub.items():
+            w.writerow(["repository", k, PUBLISHED_SETS[k], len(v), len({s['md5'] for s in v})])
+            say(f"  repository {k:19s} n={len(v):4d}")
+    unlisted = discover_unlisted(a.legacy_root, [v[0] for v in LEGACY_SETS.values()])
+    if unlisted:
+        say("\nFolders with samples in the legacy tree that are NOT listed in experiments/layout.py:")
+        for rel, n in unlisted:
+            say(f"  {rel}  ({n} samples)")
+    if not legacy["all_tune/train"]:
+        say("\nERROR: single_cell/all_tune/train not found under --legacy-root."); open(os.path.join(a.out, "report.txt"), "w").write("\n".join(rep)); return
+
+    # ---------------- published vs legacy ----------------
+    say("\nRepository dataset vs legacy:")
+    rows = []
+    for p, l in [("singlecell/train", "fibro_split/train"), ("singlecell/test", "fibro_split/test"),
+                 ("spheroid/train", "spheroid/train"), ("spheroid/test", "spheroid/test")]:
+        P, Lg = {s["md5"] for s in pub[p]}, {s["md5"] for s in legacy[l]}
+        r = {"repository_set": p, "legacy_set": l, "n_repository": len(P), "n_legacy": len(Lg),
+             "shared": len(P & Lg), "only_repository": len(P - Lg), "only_legacy": len(Lg - P)}
+        rows.append(r); say(f"  {p:18s} vs {l:18s}: shared {r['shared']}, only repo {r['only_repository']}, only legacy {r['only_legacy']}")
+    with open(os.path.join(a.out, "published_vs_legacy.csv"), "w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0])); w.writeheader(); w.writerows(rows)
+
+    # ---------------- all_tune composition ----------------
+    by_md5_pub = {s["md5"]: (k, s) for k, v in pub.items() for s in v}
+    src_keys = [k for k, v in LEGACY_SETS.items() if v[1] in ("published", "source", "external")]
+    by_md5_leg = collections.defaultdict(list)
+    for k in src_keys:
+        for s in legacy[k]:
+            by_md5_leg[s["md5"]].append((k, s))
+    comp, missing, codes = [], [], collections.Counter()
+    for split_key, split in [("all_tune/train", "train"), ("all_tune/test", "val")]:
+        for s in legacy[split_key]:
+            in_pub = by_md5_pub.get(s["md5"])
+            srcs = [k for k, _ in by_md5_leg.get(s["md5"], [])]
+            src = next((k for k in srcs if LEGACY_SETS[k][1] != "external"), srcs[0] if srcs else "UNMATCHED")
+            code = legacy_substrate_code(s["folder"])
+            codes[(split, code)] += 1
+            r = {"alltune_split": split, "alltune_folder": s["folder"], "source_legacy_set": src,
+                 "all_matching_sets": ";".join(srcs), "in_repository_dataset": in_pub[0] if in_pub else "",
+                 "legacy_substrate_code": code, "bf_md5": s["md5"]}
+            comp.append(r)
+            if not in_pub:
+                missing.append(r)
+    for name, data in [("alltune_composition.csv", comp), ("missing_from_published.csv", missing)]:
+        with open(os.path.join(a.out, name), "w", newline="") as f:
+            w = csv.DictWriter(f, fieldnames=list(comp[0])); w.writeheader(); w.writerows(data)
+    with open(os.path.join(a.out, "substrate_codes.csv"), "w", newline="") as f:
+        w = csv.writer(f); w.writerow(["alltune_split", "legacy_substrate_code", "n", "config_key_fed_to_model"])
+        for (sp, c), n in sorted(codes.items()):
+            w.writerow([sp, c, n, legacy_code_to_config_key(c)])
+    say("\nall_tune composition (split x source set):")
+    cc = collections.Counter((r["alltune_split"], r["source_legacy_set"]) for r in comp)
+    for (sp, src), n in sorted(cc.items()):
+        say(f"  {sp:5s}  {src:20s} {n:4d}")
+    say(f"\nall_tune images NOT in the repository dataset (Zenodo v1): {len(missing)}")
+    for (sp, src), n in sorted(collections.Counter((r['alltune_split'], r['source_legacy_set']) for r in missing).items()):
+        say(f"  {sp:5s}  {src:20s} {n:4d}")
+    say("\nSubstrate code the original code fed for all_tune folders (folder_name[:4]):")
+    for (sp, c), n in sorted(codes.items()):
+        say(f"  {sp:5s}  '{c}' -> {legacy_code_to_config_key(c):18s} {n:4d}")
+
+    # ---------------- overlaps / leakage ----------------
+    names = list(legacy)
+    sets = {k: {s["md5"] for s in legacy[k]} for k in names}
+    with open(os.path.join(a.out, "overlap_matrix.csv"), "w", newline="") as f:
+        w = csv.writer(f); w.writerow(["set"] + names)
+        for x in names:
+            w.writerow([x] + [len(sets[x] & sets[y]) for y in names])
+    tv = sets["all_tune/train"] | sets["all_tune/test"]
+    say("\nLeakage check - independent test sets vs all_tune train+val (must be 0):")
+    for k, v in LEGACY_SETS.items():
+        if v[1] == "external":
+            say(f"  {k:20s} {len(sets[k] & tv):3d}   (of {len(sets[k])})")
+    say("Validation-set content - legacy 'test' folders inside all_tune/test (val):")
+    for k in ["fibro_split/test", "1kpa_new/test", "10kpa_new/test"]:
+        say(f"  {k:20s} {len(sets[k] & sets['all_tune/test']):3d} of {len(sets[k])} in val; {len(sets[k] & sets['all_tune/train']):3d} in train")
+
+    # ---------------- write ----------------
+    if a.write:
+        say("\nCopying and writing manifests ...")
+        n_copied = collections.Counter()
+        alltune_rows = {"train": [], "val": []}
+        for r, s in zip(comp, legacy["all_tune/train"] + legacy["all_tune/test"]):
+            split, src = r["alltune_split"], r["source_legacy_set"]
+            meta = LEGACY_SETS[src][3] if src in LEGACY_SETS else {}
+            if r["in_repository_dataset"]:
+                pk, ps = by_md5_pub[s["md5"]]
+                sample_dir = os.path.join(PUBLISHED_SETS[pk], ps["folder"])
+            else:
+                clean = LEGACY_SETS[src][2] if src in LEGACY_SETS and LEGACY_SETS[src][2] else "unmatched"
+                sample_dir = os.path.join("singlecell_extra", clean, s["folder"])
+                n_copied[copy_sample(s["dir"], os.path.join(a.dataset_root, sample_dir), s["md5"])] += 1
+            alltune_rows[split].append(row_for(sample_dir, split, src, meta, s["folder"], s["md5"]))
+        write_manifest(manifest_path("alltune_train"), alltune_rows["train"])
+        write_manifest(manifest_path("alltune_val"), alltune_rows["val"])
+        for pk, man, lk in [("singlecell/train", "zenodo_v1_train", "fibro_split/train"),
+                            ("singlecell/test", "zenodo_v1_test", "fibro_split/test"),
+                            ("spheroid/train", "spheroid_train", "spheroid/train"),
+                            ("spheroid/test", "spheroid_val", "spheroid/test")]:
+            meta = LEGACY_SETS[lk][3]
+            write_manifest(manifest_path(man), [row_for(os.path.join(PUBLISHED_SETS[pk], s["folder"]), man.split("_")[-1],
+                                                        lk, meta, s["folder"], s["md5"]) for s in pub[pk]])
+        for k, (rel, role, clean, meta) in LEGACY_SETS.items():
+            if role != "external" or not legacy[k]:
+                continue
+            rows_ = []
+            for s in legacy[k]:
+                sample_dir = os.path.join("singlecell_external", clean, s["folder"])
+                n_copied[copy_sample(s["dir"], os.path.join(a.dataset_root, sample_dir), s["md5"])] += 1
+                rows_.append(row_for(sample_dir, "test", k, meta, s["folder"], s["md5"]))
+            write_manifest(manifest_path(f"test_{clean}"), rows_)
+        say(f"  sample folders: {dict(n_copied)}")
+        say(f"  manifests written to experiments/splits/ - commit them to git")
+    else:
+        say("\n(audit only - rerun with --write to copy the missing images and write experiments/splits/*.csv)")
+    open(os.path.join(a.out, "report.txt"), "w").write("\n".join(rep))
+    print(f"\nReport: {os.path.join(a.out, 'report.txt')}")
+
+
+if __name__ == "__main__":
+    main()

experiments/run_all.py

@@ -0,0 +1,104 @@
+"""
+Run the revision experiments end to end (cross-platform; run from the S2F repository root).
+
+    python -m experiments.run_all --legacy-root C:/Users/Kaveh/Documents/cellforce/dataset --steps audit
+    python -m experiments.run_all --legacy-root ... --steps prepare,eval_released
+    python -m experiments.run_all --steps train,eval_runs --variants s2f_full,unet,s2f_no_gan
+    python -m experiments.run_all --legacy-root ... --stage1 C:/.../outputs/checkpoints/last/last_checkpoint_final.pth --steps eval_released
+
+steps:
+  audit          dataset audit only (no copying)                          -> experiments/results/dataset_audit/
+  prepare        copy missing images into dataset/ + write experiments/splits/*.csv
+  eval_released  released checkpoints (+ substrate sweep, mean-map baseline, optional stage-1)
+                                                                          -> experiments/results/eval/
+  train          train the listed variants (default: configs.PRIORITY), skipping finished runs.
+                 One job per GPU in parallel (--gpus 0,1); the next variant starts as soon as a GPU is free.
+                 Logs: experiments/runs/<variant>_seed<seed>/train.log
+  eval_runs      evaluate every experiments/runs/*/best_checkpoint.pth
+"""
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+
+from experiments.configs import PRIORITY, get_config
+from experiments.layout import RUNS_DIR, S2F_ROOT
+
+
+def sh(args):
+    print("\n$ " + " ".join(args), flush=True)
+    r = subprocess.run([sys.executable, "-m"] + args, cwd=S2F_ROOT)
+    if r.returncode:
+        sys.exit(r.returncode)
+
+
+def finished(variant, seed):
+    h = os.path.join(RUNS_DIR, f"{variant}_seed{seed}", "history.json")
+    if not os.path.exists(h):
+        return False
+    return len(json.load(open(h))) >= get_config(variant)["epochs"]
+
+
+def gpu_queue(queue, gpus, num_workers, extra=()):
+    """Run training jobs in parallel, one per GPU (each job sees only its GPU as cuda:0)."""
+    running, failed = {}, []
+    queue = list(queue)
+    while queue or running:
+        for gpu in gpus:
+            if gpu not in running and queue:
+                v, seed = queue.pop(0)
+                run_dir = os.path.join(RUNS_DIR, f"{v}_seed{seed}")
+                os.makedirs(run_dir, exist_ok=True)
+                log = open(os.path.join(run_dir, "train.log"), "a")
+                env = dict(os.environ, CUDA_VISIBLE_DEVICES=gpu, PYTHONUNBUFFERED="1")
+                cmd = [sys.executable, "-m", "experiments.train", "--variant", v, "--seed", seed, "--resume",
+                       "--num-workers", num_workers] + list(extra)
+                running[gpu] = (subprocess.Popen(cmd, cwd=S2F_ROOT, env=env, stdout=log, stderr=subprocess.STDOUT), v, seed, log)
+                print(f"[{time.strftime('%H:%M')}] GPU {gpu}: started {v} seed {seed}  (log {run_dir}/train.log)", flush=True)
+        time.sleep(10 if extra else 30)
+        for gpu, (proc, v, seed, log) in list(running.items()):
+            if proc.poll() is not None:
+                log.close(); del running[gpu]
+                status = "finished" if proc.returncode == 0 else f"FAILED (exit {proc.returncode}, see train.log)"
+                if proc.returncode:
+                    failed.append((v, seed))
+                print(f"[{time.strftime('%H:%M')}] GPU {gpu}: {v} seed {seed} {status}", flush=True)
+    if failed:
+        print(f"failed jobs: {failed}")
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--legacy-root", default=None)
+    ap.add_argument("--stage1", default=None, help="path to last_checkpoint_final.pth")
+    ap.add_argument("--steps", default="audit")
+    ap.add_argument("--variants", default=",".join(PRIORITY))
+    ap.add_argument("--seeds", default="0")
+    ap.add_argument("--gpus", default="0,1", help="GPU ids; one training job per GPU at a time")
+    ap.add_argument("--num-workers", default="2", help="DataLoader workers per training job")
+    ap.add_argument("--extra-train-args", default="", help="passed to experiments.train, e.g. '--epochs 1 --max-batches 2 --allow-cpu' for a smoke test")
+    a = ap.parse_args()
+    steps = a.steps.split(",")
+    if ("audit" in steps or "prepare" in steps) and not a.legacy_root:
+        sys.exit("--legacy-root is required for audit/prepare")
+    if "audit" in steps:
+        sh(["experiments.prepare_dataset", "--legacy-root", a.legacy_root])
+    if "prepare" in steps:
+        sh(["experiments.prepare_dataset", "--legacy-root", a.legacy_root, "--write"])
+    if "eval_released" in steps:
+        cmd = ["experiments.evaluate", "--released", "--sweep", "--meanmap"]
+        if a.stage1:
+            cmd += ["--stage1", a.stage1]
+        sh(cmd)
+    if "train" in steps:
+        queue = [(v, s) for v in a.variants.split(",") for s in a.seeds.split(",") if not finished(v, int(s))]
+        print(f"{len(queue)} training jobs on GPUs {a.gpus}: {queue}")
+        gpu_queue(queue, a.gpus.split(","), a.num_workers, a.extra_train_args.split())
+    if "eval_runs" in steps:
+        sh(["experiments.evaluate", "--runs"])
+
+
+if __name__ == "__main__":
+    main()

experiments/train.py

@@ -0,0 +1,211 @@
+"""
+Train one baseline / ablation variant with the shared protocol (experiments/configs.py).
+
+    python -m experiments.train --variant s2f_full
+    python -m experiments.train --variant unet --seed 1
+    python -m experiments.train --variant s2f_no_gan --resume          # continue from last_checkpoint.pth
+    python -m experiments.train --variant unet --epochs 1 --max-batches 5   # smoke test
+
+Outputs in experiments/runs/<variant>_seed<seed>/:
+    config.json            exact configuration used
+    history.json           per-epoch losses (each generator loss term separately), lr, time
+    best_checkpoint.pth    generator weights at the lowest validation reconstruction loss
+    last_checkpoint.pth    full state for --resume
+    visualizations/        predictions on the first validation batch every 10 epochs
+The training loop follows training/s2f_trainer.py (same losses, label convention and update order);
+it additionally skips the discriminator when lambda_gan = 0, records every loss term, and seeds all RNGs.
+"""
+import argparse
+import json
+import math
+import os
+import random
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from experiments.ablation_models import build_discriminator, build_generator, n_params
+from experiments.configs import get_config
+from experiments.layout import DEFAULT_DATASET_ROOT, RUNS_DIR, S2F_ROOT
+from experiments.losses import PerImageWFMRMELoss
+from experiments.manifest_data import make_loader
+from models.s2f_model import create_settings_channels
+from utils.substrate_settings import compute_settings_normalization
+
+CONFIG_PATH = os.path.join(S2F_ROOT, "config", "substrate_settings.json")
+
+
+def seed_everything(seed):
+    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)
+
+
+def cosine_with_warmup(optimizer, warmup_steps, total_steps):
+    """Same schedule as diffusers.optimization.get_cosine_schedule_with_warmup (num_cycles=0.5)."""
+    def f(step):
+        if step < warmup_steps:
+            return step / max(1, warmup_steps)
+        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
+        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, f)
+
+
+def add_settings(x, batch, cfg, norm, device):
+    if cfg["substrate_mode"] == "none":
+        return x
+    meta = batch[4]
+    return torch.cat([x, create_settings_channels(meta, norm, device, x.shape, config_path=CONFIG_PATH)], dim=1)
+
+
+def save_vis(G, val_loader, cfg, norm, device, path):
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    batch = next(iter(val_loader))
+    x = add_settings(batch[0].to(device), batch, cfg, norm, device)
+    with torch.no_grad():
+        p = (G(x) + 1) / 2
+    n = min(2, x.size(0))
+    fig, ax = plt.subplots(3, n, figsize=(4 * n, 12), squeeze=False)
+    for i in range(n):
+        ax[0, i].imshow(x[i, 0].cpu(), cmap="gray"); ax[1, i].imshow(p[i, 0].cpu(), cmap="jet", vmin=0, vmax=1)
+        ax[2, i].imshow(batch[1][i, 0], cmap="jet", vmin=0, vmax=1)
+        for r in range(3):
+            ax[r, i].axis("off")
+    fig.tight_layout(); fig.savefig(path, dpi=100); plt.close(fig)
+
+
+def train(cfg, run_dir, dataset_root=DEFAULT_DATASET_ROOT, device=None, resume=False, max_batches=None,
+          num_workers=2, allow_cpu=False):
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    if str(device).startswith("cpu") and not allow_cpu:
+        raise SystemExit("CUDA is not available - refusing to train on CPU (pass --allow-cpu for a smoke test)")
+    if str(device).startswith("cuda"):
+        torch.backends.cudnn.benchmark = True
+        print(f"GPU: {torch.cuda.get_device_name(torch.device(device))}  "
+              f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', 'all')})")
+    os.makedirs(os.path.join(run_dir, "visualizations"), exist_ok=True)
+    seed_everything(cfg["seed"])
+    train_loader = make_loader(cfg["train_manifest"], dataset_root, cfg["batch_size"], cfg["img_size"],
+                               augment=cfg["augment"], shuffle=True, substrate_mode=cfg["substrate_mode"],
+                               input_mode=cfg["input_mode"], seed=cfg["seed"], num_workers=num_workers)
+    if num_workers > 0:                      # check that worker processes work on this machine
+        try:
+            next(iter(train_loader))
+        except RuntimeError as e:
+            print(f"DataLoader workers failed ({str(e).splitlines()[0][:120]}); falling back to num_workers=0")
+            train_loader = make_loader(cfg["train_manifest"], dataset_root, cfg["batch_size"], cfg["img_size"],
+                                       augment=cfg["augment"], shuffle=True, substrate_mode=cfg["substrate_mode"],
+                                       input_mode=cfg["input_mode"], seed=cfg["seed"], num_workers=0)
+        seed_everything(cfg["seed"])
+    val_loader = make_loader(cfg["val_manifest"], dataset_root, cfg["batch_size"], cfg["img_size"],
+                             substrate_mode=cfg["substrate_mode"], input_mode=cfg["input_mode"])
+    norm = compute_settings_normalization(config_path=CONFIG_PATH)
+    use_gan = cfg["lambda_gan"] > 0
+    G = build_generator(**cfg).to(device)
+    D = build_discriminator(cfg["in_channels"]).to(device) if use_gan else None
+    cfg = dict(cfg, n_params_generator=n_params(G), n_params_discriminator=n_params(D) if D else 0,
+               n_train=len(train_loader.dataset), n_val=len(val_loader.dataset), device=str(device),
+               torch_version=torch.__version__)
+    json.dump(cfg, open(os.path.join(run_dir, "config.json"), "w"), indent=2)
+    print(f"[{cfg['variant']}] G {cfg['n_params_generator']/1e6:.2f} M params, D {cfg['n_params_discriminator']/1e6:.2f} M; "
+          f"train {cfg['n_train']}, val {cfg['n_val']}, device {device}")
+
+    g_opt = torch.optim.Adam(G.parameters(), lr=cfg["g_lr"], betas=(cfg["beta1"], cfg["beta2"]))
+    d_opt = torch.optim.Adam(D.parameters(), lr=cfg["d_lr"], betas=(cfg["beta1"], cfg["beta2"])) if use_gan else None
+    steps = len(train_loader) * cfg["epochs"]
+    g_sch = cosine_with_warmup(g_opt, int(steps * cfg["warmup_frac"]), steps)
+    d_sch = cosine_with_warmup(d_opt, int(steps * cfg["warmup_frac"]), steps) if use_gan else None
+    rec_loss, bce = PerImageWFMRMELoss(), nn.BCEWithLogitsLoss()
+    hist, best, start = [], float("inf"), 0
+
+    last_path = os.path.join(run_dir, "last_checkpoint.pth")
+    if resume and os.path.exists(last_path):
+        ck = torch.load(last_path, map_location="cpu", weights_only=False)
+        G.load_state_dict(ck["generator_state_dict"]); g_opt.load_state_dict(ck["g_optimizer"]); g_sch.load_state_dict(ck["g_scheduler"])
+        if use_gan:
+            D.load_state_dict(ck["discriminator_state_dict"]); d_opt.load_state_dict(ck["d_optimizer"]); d_sch.load_state_dict(ck["d_scheduler"])
+        hist, best, start = ck["history"], ck["best_val_loss"], ck["epoch"] + 1
+        print(f"resumed at epoch {start}")
+
+    for epoch in range(start, cfg["epochs"]):
+        t0 = time.time(); G.train(); D.train() if use_gan else None
+        acc = dict(g=0.0, rec=0.0, gan=0.0, force=0.0, d=0.0); nb = 0
+        for bi, batch in enumerate(train_loader):
+            if max_batches and bi >= max_batches:
+                break
+            x = add_settings(batch[0].to(device), batch, cfg, norm, device)
+            y = batch[1].to(device)
+            fake = G(x); p = (fake + 1) / 2
+            l_rec = rec_loss(p, y)
+            l_force = F.mse_loss(p.flatten(1).mean(1), y.flatten(1).mean(1))
+            l_gan = torch.zeros((), device=device)
+            if use_gan:
+                pf = D(torch.cat([x, fake], 1)); l_gan = bce(pf, torch.ones_like(pf))
+            l_g = cfg["lambda_rec"] * l_rec + cfg["lambda_gan"] * l_gan + cfg["lambda_force"] * l_force
+            g_opt.zero_grad(); l_g.backward(); g_opt.step()
+            l_d = torch.zeros(())
+            if use_gan:
+                pr = D(torch.cat([x, y * 2 - 1], 1)); pf = D(torch.cat([x, fake.detach()], 1))
+                l_d = 0.5 * (bce(pr, torch.ones_like(pr)) + bce(pf, torch.zeros_like(pf)))
+                d_opt.zero_grad(); l_d.backward(); d_opt.step(); d_sch.step()
+            g_sch.step()
+            for k, v in zip(acc, [l_g, l_rec, l_gan, l_force, l_d]):
+                acc[k] += float(v)
+            nb += 1
+        G.eval(); v_rec = v_mse = 0.0; nv = 0
+        with torch.no_grad():
+            for bi, batch in enumerate(val_loader):
+                if max_batches and bi >= max_batches:
+                    break
+                x = add_settings(batch[0].to(device), batch, cfg, norm, device); y = batch[1].to(device)
+                p = (G(x) + 1) / 2
+                v_rec += float(rec_loss(p, y)); v_mse += float(F.mse_loss(p, y)); nv += 1
+        h = {"epoch": epoch, **{f"train_{k}": v / max(nb, 1) for k, v in acc.items()},
+             "val_rec": v_rec / max(nv, 1), "val_mse": v_mse / max(nv, 1),
+             "g_lr": g_opt.param_groups[0]["lr"], "epoch_seconds": time.time() - t0}
+        hist.append(h)
+        is_best = h["val_rec"] < best
+        best = min(best, h["val_rec"])
+        print(f"epoch {epoch:3d}  G {h['train_g']:.4f} (rec {h['train_rec']:.4f} gan {h['train_gan']:.4f} "
+              f"force {h['train_force']:.5f})  D {h['train_d']:.4f}  val_rec {h['val_rec']:.4f}  "
+              f"{'*' if is_best else ' '}  {h['epoch_seconds']:.0f}s", flush=True)
+        json.dump(hist, open(os.path.join(run_dir, "history.json"), "w"), indent=1)
+        if is_best:
+            torch.save({"epoch": epoch, "generator_state_dict": G.state_dict(), "config": cfg,
+                        "val_loss": h["val_rec"], "history": hist}, os.path.join(run_dir, "best_checkpoint.pth"))
+        state = {"epoch": epoch, "generator_state_dict": G.state_dict(), "g_optimizer": g_opt.state_dict(),
+                 "g_scheduler": g_sch.state_dict(), "history": hist, "best_val_loss": best, "config": cfg}
+        if use_gan:
+            state.update(discriminator_state_dict=D.state_dict(), d_optimizer=d_opt.state_dict(), d_scheduler=d_sch.state_dict())
+        torch.save(state, last_path)
+        if epoch % 10 == 0 or epoch == cfg["epochs"] - 1:
+            save_vis(G, val_loader, cfg, norm, device, os.path.join(run_dir, "visualizations", f"epoch_{epoch:03d}.png"))
+    print(f"done. best val_rec {best:.4f}; checkpoints in {run_dir}")
+    return run_dir
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--variant", required=True)
+    ap.add_argument("--seed", type=int, default=None)
+    ap.add_argument("--epochs", type=int, default=None)
+    ap.add_argument("--g-lr", type=float, default=None); ap.add_argument("--d-lr", type=float, default=None)
+    ap.add_argument("--batch-size", type=int, default=None)
+    ap.add_argument("--dataset-root", default=DEFAULT_DATASET_ROOT)
+    ap.add_argument("--run-dir", default=None)
+    ap.add_argument("--device", default=None)
+    ap.add_argument("--resume", action="store_true")
+    ap.add_argument("--max-batches", type=int, default=None, help="limit batches per epoch (smoke test only)")
+    ap.add_argument("--num-workers", type=int, default=2, help="DataLoader worker processes for augmentation")
+    ap.add_argument("--allow-cpu", action="store_true", help="allow training without CUDA (smoke tests)")
+    a = ap.parse_args()
+    cfg = get_config(a.variant, seed=a.seed, epochs=a.epochs, g_lr=a.g_lr, d_lr=a.d_lr, batch_size=a.batch_size)
+    run_dir = a.run_dir or os.path.join(RUNS_DIR, f"{a.variant}_seed{cfg['seed']}")
+    train(cfg, run_dir, a.dataset_root, a.device, a.resume, a.max_batches, a.num_workers, a.allow_cpu)
+
+
+if __name__ == "__main__":
+    main()