# 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 and weighted MSE + soft Dice losses (both 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 (S2F about 7 h, U-Net about 2.5 h each on an RTX A6000; resumable), then their evaluation
python -m experiments.run_all --steps train,eval_runs --variants s2f_full,unet,unet_gan,s2f_no_gan,s2f_no_substrate,s2f_no_force
```

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 (and finished phases) are skipped and interrupted runs resume from the phase's `last_checkpoint.pth`, so the command can simply be re-run.
- Run time on an RTX A6000: S2F about 3 min per epoch, i.e. about 7 h per variant (100 + 40 epochs); U-Net about 2.5 h. On two GPUs the 6 core variants take about 15 h (`s2f_no_force` reuses phase 1 of `s2f_full`).
- Progress is logged to `experiments/runs/<variant>_seed<seed>/train.log` (validation pixel r and WFM-RME every epoch) and `phase*/history.json`.

For a quick check that everything runs, add `--epochs 1 --max-batches 2` (1 epoch per phase; add `--allow-cpu` without a GPU).

## Training protocol (all variants)

Every variant is trained from scratch on `alltune_train` (280 pairs), batch size 2, 1024 × 1024 px, the augmentations of `data/augmentations.py`, Adam (β = 0.5/0.999) with cosine decay after 10 % linear warm-up, in two phases that mirror the staged training of the released model:

| phase | epochs | reconstruction loss | other terms | lr | starts from |
|---|---|---|---|---|---|
| 1 localisation | 100 | 100 × weighted MSE + soft Dice (pos_weight 3, threshold 0.2, Dice 0.5) | 1 × adversarial | 1e-4 | random weights |
| 2 magnitude | 40 | 10 × per-image WFM-RME | 1 × adversarial + 100 × MSE of mean intensity | 1e-5 | best phase-1 generator |

Each phase uses a new discriminator and optimiser and keeps the checkpoint with the lowest validation loss of that phase on `alltune_val` (81 pairs); the final model is the best phase-2 checkpoint. WFM-RME alone does not learn to localise the forces from random weights (a first protocol that used it from the start gave pixel r ≈ 0), which is also why the released model was first trained with the weighted MSE + Dice loss. 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 in phase 2; phase 1 is identical to `s2f_full` and reused) |
| `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 from earlier checkpoints, with substrate channels added only in its last two stages (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`.