Add a high-level Python API so notebooks and scripts don't have to shell out to the CLI
#707 opened on Jul 16, 2026
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'ssave_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()intobuild_parser()+run(args, *, config=None, datastore=None)+ a thinmain(), so the CLI and a Python API share one code path; - add
neural_lam/api.pywithtrain(...)/evaluate(...)returning aRunhandle exposingcheckpoint_path,plot_dir,example_plotsandrmse_plot, all computed from the same constants the writer uses, so reader and writer cannot drift apart; - have API keywords default to
Noneand 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.