NeCast / README.md

Redraw the card figures in the NeCast colour palette (NeCast purple)

325dda3 verified 1 day ago

6.7 kB

---
license: cc-by-4.0
library_name: pytorch
tags:
- population-genetics
- demographic-inference
- effective-population-size
- coalescent
datasets:
- Larrivhan/NeCast-simulations
---
# NeCast

NeCast estimates the size history of a population from the unphased genotypes of 32 diploid
individuals in one window of 30 Mb of accessible sequence, together with the recombination and
mutation maps estimated for that window. It returns log10 Ne on 64 time intervals, evenly spaced
in log time between 50 and 500,000 generations ago, and analyses a window in seconds on a CPU.
This repository holds the network used in the NeCast paper, its training record, and the other
networks the paper compares it with.

## Files

| path | contents | size |
|---|---|---|
| `final/checkpoint.pt` | weights of the network used in the paper | 4.2 MB |
| `final/training_state.pt` | the state saved at the end of training | 16.7 MB |
| `final/PLAN.json` | the full training configuration, input checksums and the lists of training, validation and test windows | 30.6 MB |
| `final/DEVELOPMENT.json` | training and validation error at every epoch | small |
| `final/test_predictions.npz` | predictions and truth for the 500 test windows | 0.6 MB |
| `final/report.json`, `DONE.json`, `run_DONE.json` | run summaries with checksums | small |
| `final/INTAKE_VERIFY.json`, `compshare_sha256.txt` | checksums of the copied files against the training machine | small |
| `final/training_log/` | the training command and logs | 0.8 MB |
| `final/snapshots/epoch06/`, `epoch09/` | states saved after epochs 6 and 9 | 47 MB each |
| `training_contrasts/` | 7 networks trained on 8,000 or 88,000 windows | 0.36 GB |
| `ablations/` | 5 networks trained on 12,500 windows | 0.25 GB |
| `figures/` | the figure on this card | small |
| `SHA256SUMS` | a checksum for every file except this card | small |

## The network

Five streams read the same sites. A site encoder reads the genotypes at every segregating site
with each site's position on the recombination and mutation maps, so that distances count
expected crossovers and mutations. Three summaries come from the same sites: the folded
site-frequency spectrum, linkage disequilibrium at seven lags measured in segregating sites,
and the spectra of identity-by-state tracts over the 496 pairs of individuals. The fifth stream
is the posterior of PSMC+ (Cousins et al. 2024), decoded for each individual with both maps and
pooled over the 32. The network outputs a correction for each of the 64 intervals, added to a
baseline curve: the window's Watterson estimate of log10 Ne plus a fixed offset per interval.

## Training

- **Data.** `Larrivhan/NeCast-simulations`: 89,996 training windows; 500 further training
  windows to choose the epoch; 500 development windows as the test set.
- **Settings.** 12 epochs on four RTX 4090 GPUs (50.3 h, 201.2 GPU-hours), width 128, seed
  20260728. Loss: mean squared error of log10 Ne over the 64 intervals. AdamW with learning rate
  3 × 10⁻⁴, cosine decay to 5% of it over the 12 epochs, weight decay 10⁻⁴, gradient-norm
  clipping at 1, one update per 32 windows. The last epoch had the lowest validation error, and
  its checkpoint is `final/checkpoint.pt`.
- **Test error.** 0.287: the root mean squared error of log10 Ne over all 64 intervals of the
  500 test windows taken together, from `final/test_predictions.npz`.

![Training and test error](https://huggingface.co/Larrivhan/NeCast/resolve/main/figures/model_training.png)

**a**, Training error and validation error after each of the 12 epochs, from
`final/DEVELOPMENT.json`; the validation windows are the 500 further training windows. **b**,
Test error in each of the 64 intervals (5-interval running mean), with the error over all
intervals as a dotted line. **c**, The test windows at the 10th, 50th and 90th percentiles of
per-window error: truth (black step) and NeCast (purple).

## Running it

The network reads inputs built by the NeCast code (summary statistics, PSMC+ posteriors and the
site encoder's input), and its checkpoint is loaded by that code. The code is being prepared for
release; its link will be added here. `final/checkpoint.pt` is a plain PyTorch file:
`torch.load(path, map_location="cpu", weights_only=True)` returns a dictionary whose
`state_dict` holds the weights. The inputs the network read for every simulated and real
window in the paper are in `Larrivhan/NeCast-simulations`, `Larrivhan/NeCast-benchmark` and
`Larrivhan/NeCast-1000G`.

## Training contrasts (`training_contrasts/`)
Each network was trained for 12 epochs with the same settings and scored on the same 500 test
windows. `results/` has the paired comparisons (`paired_20260925.json`), every prediction
(`eval_all_20260925.npz`) and a summary (`RESULTS.md`).

| folder | training windows | trained on | inputs |
|---|---|---|---|
| `full_mix_8k` | 8,000 | the full training mix | all five streams |
| `no_errors_8k` | 8,000 | windows without errors, constant maps | all five streams |
| `errors_constant_maps_8k` | 8,000 | windows with errors, constant maps | all five streams |
| `no_templates_8k` | 8,000 | the full mix without published human histories | all five streams |
| `site_encoder_only_8k` | 8,000 | the full training mix | site encoder only |
| `summaries_only_8k` | 8,000 | the full training mix | the three summaries only |
| `summaries_only_88k` | 88,000 | the full training mix | the three summaries only |

`no_errors_8k` against `errors_constant_maps_8k` gives the effect of training with errors
(Fig. 3d of the paper); `errors_constant_maps_8k` against `full_mix_8k` that of training with
varying maps (Fig. 4a).

## Ablations (`ablations/`)

Each network was trained once on the same 12,500 windows of the full mix for 12 epochs and
scored on the 500 test windows. `results/ablation_paired.json` has the paired comparisons behind
Extended Data Fig. 1 (`ablation_paired.py` recomputes them).

| folder | inputs |
|---|---|
| `full_12500` | all five streams |
| `no_posterior_12500` | without the PSMC+ posterior |
| `no_summaries_12500` | without the three summaries |
| `summaries_only_12500` | the three summaries alone |
| `converged_posterior_12500` | all five streams, with a PSMC+ posterior whose prior was refitted five times instead of once |

Inside a run folder, `PLAN.json` holds the configuration, `DEVELOPMENT.json` the per-epoch
errors, `<inputs>.final.pt` the final state, and `<inputs>/checkpoint.pt` the weights with
`mechanism_evaluation.npz` (the test predictions).

## Licence

The weights and files here are released under CC BY 4.0.

## Citation

The NeCast paper is in preparation; its reference will be added here.

	---
	license: cc-by-4.0
	library_name: pytorch
	tags:
	- population-genetics
	- demographic-inference
	- effective-population-size
	- coalescent
	datasets:
	- Larrivhan/NeCast-simulations
	---

	# NeCast

	NeCast estimates the size history of a population from the unphased genotypes of 32 diploid
	individuals in one window of 30 Mb of accessible sequence, together with the recombination and
	mutation maps estimated for that window. It returns log10 Ne on 64 time intervals, evenly spaced
	in log time between 50 and 500,000 generations ago, and analyses a window in seconds on a CPU.
	This repository holds the network used in the NeCast paper, its training record, and the other
	networks the paper compares it with.

	## Files

	\| path \| contents \| size \|
	\|---\|---\|---\|
	\| `final/checkpoint.pt` \| weights of the network used in the paper \| 4.2 MB \|
	\| `final/training_state.pt` \| the state saved at the end of training \| 16.7 MB \|
	\| `final/PLAN.json` \| the full training configuration, input checksums and the lists of training, validation and test windows \| 30.6 MB \|
	\| `final/DEVELOPMENT.json` \| training and validation error at every epoch \| small \|
	\| `final/test_predictions.npz` \| predictions and truth for the 500 test windows \| 0.6 MB \|
	\| `final/report.json`, `DONE.json`, `run_DONE.json` \| run summaries with checksums \| small \|
	\| `final/INTAKE_VERIFY.json`, `compshare_sha256.txt` \| checksums of the copied files against the training machine \| small \|
	\| `final/training_log/` \| the training command and logs \| 0.8 MB \|
	\| `final/snapshots/epoch06/`, `epoch09/` \| states saved after epochs 6 and 9 \| 47 MB each \|
	\| `training_contrasts/` \| 7 networks trained on 8,000 or 88,000 windows \| 0.36 GB \|
	\| `ablations/` \| 5 networks trained on 12,500 windows \| 0.25 GB \|
	\| `figures/` \| the figure on this card \| small \|
	\| `SHA256SUMS` \| a checksum for every file except this card \| small \|

	## The network

	Five streams read the same sites. A site encoder reads the genotypes at every segregating site
	with each site's position on the recombination and mutation maps, so that distances count
	expected crossovers and mutations. Three summaries come from the same sites: the folded
	site-frequency spectrum, linkage disequilibrium at seven lags measured in segregating sites,
	and the spectra of identity-by-state tracts over the 496 pairs of individuals. The fifth stream
	is the posterior of PSMC+ (Cousins et al. 2024), decoded for each individual with both maps and
	pooled over the 32. The network outputs a correction for each of the 64 intervals, added to a
	baseline curve: the window's Watterson estimate of log10 Ne plus a fixed offset per interval.

	## Training

	- Data. `Larrivhan/NeCast-simulations`: 89,996 training windows; 500 further training
	windows to choose the epoch; 500 development windows as the test set.
	- Settings. 12 epochs on four RTX 4090 GPUs (50.3 h, 201.2 GPU-hours), width 128, seed
	20260728. Loss: mean squared error of log10 Ne over the 64 intervals. AdamW with learning rate
	3 × 10⁻⁴, cosine decay to 5% of it over the 12 epochs, weight decay 10⁻⁴, gradient-norm
	clipping at 1, one update per 32 windows. The last epoch had the lowest validation error, and
	its checkpoint is `final/checkpoint.pt`.
	- Test error. 0.287: the root mean squared error of log10 Ne over all 64 intervals of the
	500 test windows taken together, from `final/test_predictions.npz`.

	![Training and test error](https://huggingface.co/Larrivhan/NeCast/resolve/main/figures/model_training.png)

	a, Training error and validation error after each of the 12 epochs, from
	`final/DEVELOPMENT.json`; the validation windows are the 500 further training windows. b,
	Test error in each of the 64 intervals (5-interval running mean), with the error over all
	intervals as a dotted line. c, The test windows at the 10th, 50th and 90th percentiles of
	per-window error: truth (black step) and NeCast (purple).

	## Running it

	The network reads inputs built by the NeCast code (summary statistics, PSMC+ posteriors and the
	site encoder's input), and its checkpoint is loaded by that code. The code is being prepared for
	release; its link will be added here. `final/checkpoint.pt` is a plain PyTorch file:
	`torch.load(path, map_location="cpu", weights_only=True)` returns a dictionary whose
	`state_dict` holds the weights. The inputs the network read for every simulated and real
	window in the paper are in `Larrivhan/NeCast-simulations`, `Larrivhan/NeCast-benchmark` and
	`Larrivhan/NeCast-1000G`.

	## Training contrasts (`training_contrasts/`)

	Each network was trained for 12 epochs with the same settings and scored on the same 500 test
	windows. `results/` has the paired comparisons (`paired_20260925.json`), every prediction
	(`eval_all_20260925.npz`) and a summary (`RESULTS.md`).

	\| folder \| training windows \| trained on \| inputs \|
	\|---\|---\|---\|---\|
	\| `full_mix_8k` \| 8,000 \| the full training mix \| all five streams \|
	\| `no_errors_8k` \| 8,000 \| windows without errors, constant maps \| all five streams \|
	\| `errors_constant_maps_8k` \| 8,000 \| windows with errors, constant maps \| all five streams \|
	\| `no_templates_8k` \| 8,000 \| the full mix without published human histories \| all five streams \|
	\| `site_encoder_only_8k` \| 8,000 \| the full training mix \| site encoder only \|
	\| `summaries_only_8k` \| 8,000 \| the full training mix \| the three summaries only \|
	\| `summaries_only_88k` \| 88,000 \| the full training mix \| the three summaries only \|

	`no_errors_8k` against `errors_constant_maps_8k` gives the effect of training with errors
	(Fig. 3d of the paper); `errors_constant_maps_8k` against `full_mix_8k` that of training with
	varying maps (Fig. 4a).

	## Ablations (`ablations/`)

	Each network was trained once on the same 12,500 windows of the full mix for 12 epochs and
	scored on the 500 test windows. `results/ablation_paired.json` has the paired comparisons behind
	Extended Data Fig. 1 (`ablation_paired.py` recomputes them).

	\| folder \| inputs \|
	\|---\|---\|
	\| `full_12500` \| all five streams \|
	\| `no_posterior_12500` \| without the PSMC+ posterior \|
	\| `no_summaries_12500` \| without the three summaries \|
	\| `summaries_only_12500` \| the three summaries alone \|
	\| `converged_posterior_12500` \| all five streams, with a PSMC+ posterior whose prior was refitted five times instead of once \|

	Inside a run folder, `PLAN.json` holds the configuration, `DEVELOPMENT.json` the per-epoch
	errors, `<inputs>.final.pt` the final state, and `<inputs>/checkpoint.pt` the weights with
	`mechanism_evaluation.npz` (the test predictions).

	## Licence

	The weights and files here are released under CC BY 4.0.

	## Citation

	The NeCast paper is in preparation; its reference will be added here.