mllam/neural-lam

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

Open

#707 创建于 2026年7月16日

在 GitHub 查看
 (1 评论) (0 反应) (0 负责人)Python (276 fork)auto 404
enhancementhelp wanted

仓库指标

Star
 (282 star)
PR 合并指标
 (PR 指标待抓取)

描述

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.

贡献者指南