mllam/neural-lam

Add a high-level Python API so notebooks and scripts don't have to shell out to the CLI

Open

#707 opened on Jul 16, 2026

View on GitHub
 (1 comment) (0 reactions) (0 assignees)Python (276 forks)auto 404
enhancementhelp wanted

Repository metrics

Stars
 (282 stars)
PR merge metrics
 (PR metrics pending)

Description

The only way into training/evaluation is the CLI (python -m neural_lam.train_model). There is no neural_lam/api.py and no console script, so anything that wants to train from Python has to shell out and then guess where the outputs landed.

Two things make that painful:

  • main() is one ~490-line, argparse-bound function (neural_lam/train_model.py:76-565), so none of the machinery is callable without building a CLI invocation.
  • The on-disk layout is written in train_model.py:508-520 (<runs_root>/<run_name>/checkpoints/{min_val_loss,last}.ckpt), and eval plots land under the logger's save_dir. Consumers have to re-implement those paths by globbing, and nothing keeps their copy in sync with the writer.

This is tricky for e.g. the DANRA tutorial (#577) to prevent from rotting. Working on that notebook inspired this issues.

I suggest to:

  • factor main() into build_parser() + run(args, *, config=None, datastore=None) + a thin main(), so the CLI and a Python API share one code path;
  • add neural_lam/api.py with train(...) / evaluate(...) returning a Run handle exposing checkpoint_path, plot_dir, example_plots and rmse_plot, all computed from the same constants the writer uses, so reader and writer cannot drift apart;
  • have API keywords default to None and inherit the parser defaults, so the API stays a faithful mirror of the CLI.

The config/datastore injection seam also unlocks a fast drift gate: with an in-memory datastore (tests/dummy_datastore.py::DummyDatastore) a tests/test_api.py can train + evaluate fully offline in ~10s and run on every PR, instead of relying on a e.g. 3 min notebook job that only runs post-merge.

Contributor guide