feat: Add high-level programmatic Python API

[GiGiKoneti]

Describe your changes

Implemented a high-level programmatic Python API that wraps the CLI scripts, resolving Issue #707.

• Added neural_lam/api.py that exposes train(**kwargs), evaluate(**kwargs), and create_graph(**kwargs).
• Extracted build_parser() and run() out of main() in both train_model.py and create_graph.py.
• Resolved PyTorch 2.6 weights_only=True unpickling errors by registering custom dataclasses and argparse.Namespace to safe globals in neural_lam/config.py.
• Added IDE autocompletion for API methods using inspect.Signature injection.
• Added comprehensive integration tests in tests/test_api.py.

Issue Link

Closes #707

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section
  reflecting type of change (add section where missing):
  • added: when you have added new functionality
  • changed: when default behaviour of the code has been changed
  • fixes: when your contribution fixes a bug
  • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
• author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
• Once the PR is ready to be merged, squash commits and merge the PR.

[GiGiKoneti]

Hi @sadamov @truongsontung,

Here is a quick summary of our implementation in this PR addressing #707:

Implementation Overview

1. Refactored CLI Entrypoints: Factored main() in train_model.py and create_graph.py into build_parser(), run(args, *, config=None, datastore=None), and a thin main(). The CLI and Python API share the exact same code execution path.
2. Programmatic Python API (neural_lam/api.py): Added train(**kwargs), evaluate(**kwargs), and create_graph(**kwargs) that default to None and inherit parser defaults from build_parser().
3. IDE Signature Parity: Injected dynamic __signature__ into API functions so IDE users get accurate keyword-only argument autocompletion and docstrings.
4. Fast Offline Integration Tests: Added tests/test_api.py using DummyDatastore to run a full create_graph -> train -> evaluate roundtrip in seconds on every PR.
5. Run Handle: Currently train() and evaluate() return a Run dataclass containing run_dir and checkpoint_path.

---

What Else Can Be Done (Proposed Enhancements)

To fully align with the issue discussion on preventing reader/writer path drift:

• We can extend the Run dataclass with explicit convenience properties:
  • plot_dir (run_dir / "plots")
  • example_plots (plot_dir / "example_plots")
  • rmse_plot (plot_dir / "rmse.png")

Please let us know if you'd like us to add these properties to Run or if you have any feedback on the current approach. We are ready to make these additions upon your approval!

[truongsontung]

Hi @GiGiKoneti, thanks for the detailed summary! The implementation looks solid - the build_parser + run split with shared code path is exactly the right approach.

A few thoughts:

1. Run dataclass properties - yes, please add plot_dir, example_plots, and rmse_plot. These are referenced in the issue discussion and help prevent the reader/writer path drift that Add a high-level Python API so notebooks and scripts don't have to shell out to the CLI #707 aims to solve.

2. Signature injection - clever solution for IDE parity. Does it handle *args correctly in the injected signature?

3. Integration tests - great that you included DummyDatastore for fast offline tests. This is important for CI reliability.

LGTM overall. Happy to test if you need a second pair of eyes.

[GiGiKoneti]

Hi @truongsontung, thanks for the review! I've just pushed a commit adding plot_dir, example_plots, and rmse_plot properties to the Run dataclass exactly as you suggested! 🚀

Regarding the *args handling in the __signature__: all of our API function parameters map to argparse arguments, which are inherently treated as keyword-only (**kwargs) when called programmatically via train(...) etc. (e.g. train(batch_size=32)). Since there are no positional arguments where order matters natively supported by our argparse setup other than boolean flags which also act as kwargs in python, inspect.Parameter.KEYWORD_ONLY perfectly captures all flags and lists (nargs='+'). It prevents ordering bugs and fully handles everything cleanly! Let me know if you need any other additions!

[truongsontung]

Great work! The Run properties look clean. Your explanation of the KEYWORD_ONLY approach makes perfect sense given the argparse foundation. LGTM - ready for maintainer review.

[truongsontung]

Looks good! The convenience properties are clean and match the issue requirements. Nice implementation.