Migrating from TraInf¶

COATI is an independent package. It does not import from, write to or require a local TraInf checkout. The original files were copied from the current working tree, including existing research changes, rather than assuming the committed public mirror represented the current algorithm.

Previous workflow	COATI workflow
One working directory per experiment	One reusable package; separate output directories
`Args.py` parses process arguments	`Config` works in scripts and notebooks
NPZ paths embedded in configuration	`TemporalData` adapters and explicit time order
`sys.path` manipulation	Installable `coati` import
Checkpoint files alone	`Result`, input snapshots, configuration and metadata
Plot script tied to a dataset	`result.plot()` for the first two coordinates

Preserved engine¶

The neural field, loss assembly, optimizer, scheduler, adaptive blur, kernel projections, mass loss, dual updates and training discretization remain in coati._core. Supported engine modules remain available for advanced use; unused MMD matching and custom Riemannian metric support have been removed. The reusable downstreamSrc analysis modules are copied unchanged into coati.downstream; run python -m coati.downstream.cli --help or coati-downstream --help. Biological analyses are not executed in this validation. Flow matching, legacy plotting and biological helpers remain low-level modules; this release’s high-level tutorials exercise dynamic OT and synchronized OT. Optional legacy plotting imports may require the legacy extra.

The copied core includes the following packaging changes. The intentional unbalanced secondary-energy correction is described separately below:

from src.* becomes package-relative imports.
Remove the unused from Args import get_args import in Epoch.py.
Allow an explicit args.device; device="auto" follows the original selection.
Return zero gradient in the 100-step diagnostic for disabled scalar/non-differentiable losses. The original crashes in that diagnostic when synchronized density is disabled; training loss and optimizer updates are unchanged.

For retained W2/Sinkhorn matching, normalization formulas, neural architectures and optimizer update rules are preserved. Unbalanced synchronized kinetic energy has the explicit correction documented below. evaluation.py retains its secondary evaluation behavior. The wrapper checks inputs, stages numeric NPZ files, creates output directories, protects existing results and adds a reloadable result interface.

The public synchronized default now adds a paired FiLM training step before calling this preserved engine. It freezes the resulting map and passes it through the existing T_model interface. map_type="mlp" uses an ordinary MLP instead. Explicit map_type="kernel" retains the original public default and is used by the original-engine equivalence checks. The automatic map’s standardization and checkpoint selection are documented in Cross-modal maps; they do not modify the trajectory engine.

Reproducing a previous experiment¶

Transfer all numerical settings, use the same data, row order, numeric times, normalization dictionaries, model/map, seed and device. Public defaults intentionally target small CPU examples; replacing a production Args.py with the default Config() is not equivalent.

Advanced engine arguments can be provided through Config(extra={...}). Data paths, times, dimensions and output locations are owned by the public adapter and derived from its inputs. Direct engine entry remains available as coati._core.Training.run_training for unsupported experimental features such as support-point workflows.

Provenance¶

provenance/core_sha256.json records the original core file hashes. scripts/verify_equivalence.py compares original and packaged training with the same configuration and CPU seed, using read-only original imports with bytecode writing disabled. Private notes, personal configuration files, datasets and repository history are not copied.

W2-only matching cleanup¶

MMD matching (MMD_loss, mmd_loss, sync_mmd_loss) and the custom metric_primal / metric_secondary interfaces have been removed, including coati._core.RiemannMetric. Distribution matching retains the existing W2-based Sinkhorn implementation (p=2), with the same blur and weighting. Balanced Euclidean secondary energy retains 0.5 * sum(||y[k+1]-y[k]||²) / dt, averaged over cells. Unbalanced secondary energy uses the evolving particle masses as described below. Model evaluation diagnostics remain available.

Old configurations with mmd_loss=false and sync_mmd_loss=false load normally; these obsolete keys are discarded. An enabled MMD option is rejected rather than silently changing the requested objective. Remove custom metric arguments from calls to fit or the low-level training functions.

Vector-field usability API¶

New train_vector_field, predict_trajectory, plot_trajectory and compute_velocity names make training, integration, plotting and field evaluation explicit. Existing fit/predict/plot calls remain supported. The new multimodal/balanced controls resolve into the original sync_loss/unbalancedModel fields, with documented presets only on explicit mode selection. Clear numerical parameter aliases resolve into existing keys. No _core training, loss, map or integration code changes are part of this update.

AnnData embedding plots estimate a local linear projection of learned primary-space velocities. They do not change training and do not fit preprocessing implicitly. See velocity plotting for row alignment, units and limitations.

Unbalanced secondary kinetic energy¶

On 2026-10-31, COATI’s unbalanced synchronized objective was corrected to weight each mapped displacement by the particle mass at the start of that solver interval:

\[ E_Y = \sum_{k,i} \exp(\mathrm{lnw}_{k,i}) \frac{\|Y_i^{k+1}-Y_i^k\|^2}{2(t_{k+1}-t_k)}. \]

This uses the actual solver time differences and the same left-endpoint convention as primary Euler integration. Particle masses already start at 1 / batch_size: sum the weighted contributions without another particle mean or per-time mass normalization. Gradients pass through both the mapped positions and log masses. The external secondary_weight and energy_weight multipliers are unchanged. Balanced synchronized energy retains its original equal-weight formula, and single-modality training is unaffected.

This is an intentional correction to the copied research objective, rather than a packaging-only change. The original TraInf repositories are not modified. Existing checkpoints and their predictions do not change; new training uses the corrected objective. Previously executed unbalanced multimodal tutorial outputs remain historical results, not results of retraining with this correction.

Synthetic data preparation¶

Toy generators, Gaussian comparison helpers, reference checkpoints and example settings have moved to examples/ in the COATI source repository. They are not included in the installed Python package. Training and inference are unchanged.

Use python examples/run_toy.py ... instead of coati toy .... Repository notebooks import examples/toy_data.py as datasets after adding the examples directory to the import path. Replace ct.datasets.paired(...) with datasets.paired(...) in your own repository examples, or supply your prepared TemporalData/AnnData directly to the public training API.

The original Gaussian-reference notebook loads the same archived NPZ files and weights from examples/data/gaussian_reference/; it does not regenerate data. The separate examples/generate_gaussian_data.py script reproduces the original primary, observed secondary and map-support arrays when explicitly requested.