Original Gaussian experiment

Reproduce the original GaussToy/GaussianBump/BalancedSync experiment using copied observations and the original frozen learned map. The NPZ and checkpoint files in examples/synthetic/data/gaussian_reference/ retain their original SHA-256 hashes. No endpoint clouds are regenerated.

from pathlib import Path
import sys

example_dir = next(parent / "examples" for parent in (Path.cwd(), *Path.cwd().parents)
                   if (parent / "examples" / "synthetic" / "toy_data.py").is_file())
if str(example_dir) not in sys.path:
    sys.path.insert(0, str(example_dir))

import uuid
import numpy as np
import torch
import matplotlib.pyplot as plt
import coati as ct
from synthetic.gaussian_reference import load_example, replay_reference, plot_reference

%matplotlib inline
example = load_example()
print("Primary snapshots:", [x.shape for x in example.data.primary])
print("Secondary snapshots:", [x.shape for x in example.data.secondary])
Primary snapshots: [(1000, 2), (2000, 2)]
Secondary snapshots: [(1000, 3), (2000, 3)]

Replay the archived result

Load the original vector-field checkpoint and integrate it with COATI on the original visualization grid. The reported errors compare this replay with the original saved trajectory tensors.

_, original_prediction, replay_error = replay_reference()
print("Maximum replay errors:", replay_error)
plot_reference(example, original_prediction)
plt.show()
Maximum replay errors: {'primary': 0.0, 'secondary': 1.1920928955078125e-07}
../_images/0196a0048e258e053a86110fd3bb0c17bbff17f16538255b88ce75832ed3abfd.png

Train with the original data and settings

Use the same copied observations, frozen learned map and reference configuration, directly from scratch. Checkpoint replay reproduces the archived trajectory; refitting creates a new optimization run and need not reproduce every floating-point value of the historical run.

run_dir = Path("outputs") / ("gaussian-reference-" + uuid.uuid4().hex[:8])
result = ct.train_vector_field(
    example.data, example.config, T_model=example.mapping, output_dir=run_dir,
)
print("Model:", run_dir)
Model: outputs/gaussian-reference-7c416252

Plot the new fit

Use the original integration method, time grid, seeded path selection and plotting layout.

prediction = result.predict_trajectory(
    times=example.times, method=example.settings["prediction"]["method"],
)
with torch.no_grad():
    prediction["secondary"] = example.mapping(torch.tensor(prediction["primary"])).numpy()
plot_reference(example, prediction)
plt.show()
../_images/c7d589e1804feee535fe6698387325d0b6f2498c39a98783b72578bcc973c57d.png

Check the difference from the archived run

These arraywise errors describe reproducibility relative to the saved checkpoint, rather than biological or distributional accuracy.

for space in ("primary", "secondary"):
    delta = prediction[space] - original_prediction[space]
    print(space, "maximum difference:", float(np.abs(delta).max()),
          "RMSE:", float(np.sqrt(np.mean(delta ** 2))))
primary maximum difference: 0.35370585322380066 RMSE: 0.04401407390832901
secondary maximum difference: 0.6980593204498291 RMSE: 0.04108712449669838

Run the same experiment from Python:

python -m examples.synthetic.run_gaussian_reference --output outputs/gaussian-reference

Replay the original checkpoint without refitting:

python -m examples.synthetic.run_gaussian_reference --output outputs/gaussian-replay --replay

The synthetic data tutorials provide balanced and unbalanced sync comparisons.