Add COSMO example notebook for onboarding (#587)

[info-gallary]

Describe your changes

This PR adds a new onboarding example notebook docs/notebooks/COSMO_example.ipynb that demonstrates the full Neural-LAM workflow (Datastore → Graph → Training → Evaluation → Visualization) using a COSMO-structured setup.

Key Improvements:

• Onboarding Focused: Designed as a "getting started" guide that provides a complete overview of the library's capabilities.
• CPU-Friendly: Uses synthetic, COSMO-structured Zarr data and reduced model dimensions, allowing the entire pipeline to be verified on standard hardware without a GPU or massive dataset downloads.
• Repo Alignment: Moves documentation notebooks into the docs/notebooks/ directory to follow the project's organizational patterns.
• Deduplication: Removes the generic HelloWorld.ipynb to avoid overlap with existing DANRA notebook work.

Motivation:
The maintainers requested a COSMO-specific example to serve as the primary onboarding documentation, replacing the previous general "hello world" implementation.

Dependencies:

• neural-lam
• mllam-data-prep (demonstrated in config)
• weather-model-graphs (demonstrated in graph construction)
• Standard stack: numpy, xarray, pandas, torch, pyyaml

Issue Link

closes #587

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 (No README changes required for this doc update)
• I have added tests that prove my fix is effective or that my feature works (The notebook serves as an integration test)
• I have given the PR a name that clearly describes the change, written in imperative form.
• 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

• 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:
  • added: COSMO_example.ipynb notebook to documentation for onboarding

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)

[info-gallary]

Hi @sadamov , I’ve updated this PR to a COSMO example notebook as discussed. The notebook and changelog have been revised accordingly. Please let me know if any further changes are needed.

[sadamov]

I think there was a misunderstanding here: the notebook for cosmo is already created. There are however some outstanding issues that need to be adressed, as mentioned in:

#69 (comment)

@info-gallary great, you will face one main obstacle. Currently the cosmo data is hosted on the ETH research collection and is multiple 100GB in size. so for a hello world example we will need to host a tiny piece of that data on a more accessible storage option, similar to DANRA. You might want to convert this .md file to a jupyter notebook. Here is my md file for cosmo: https://github.com/joeloskarsson/neural-lam-dev/blob/research/docs/reproduce_paper_sample.md

also #69 is mostly talking about danra, so you shouldn't close that issue with this PR.

So to rephrase, the issue about the COSMO hello world example is: how to bring the notebook from the research branch into main and get it to run in a reasonable amout of time. The preparation and hosting of a readuced cosmo sample is very much the core work. I do understand that hosting data is not simple/cheap, but it is required here. Since this PR is not solving the issues I'll mark it as draft for now.
If you would like to start working on the example data hosting I suggest to contact @leifdenby as he has some options from what I understand 😉

[info-gallary]

Thanks for the clarification, that makes sense. I understand now that the core missing work is preparing and hosting a reduced COSMO sample so the notebook can actually run in a reasonable time, rather than only converting the existing notebook into main. I’ll avoid closing #69 from this PR.

I’d still be interested in helping with the reduced sample data preparation/hosting side. I’ll reach out to @leifdenby to better understand the available hosting options and what would be most useful here.

I think there was a misunderstanding here: the notebook for cosmo is already created. There are however some outstanding issues that need to be adressed, as mentioned in:

#69 (comment)

@info-gallary great, you will face one main obstacle. Currently the cosmo data is hosted on the ETH research collection and is multiple 100GB in size. so for a hello world example we will need to host a tiny piece of that data on a more accessible storage option, similar to DANRA. You might want to convert this .md file to a jupyter notebook. Here is my md file for cosmo: https://github.com/joeloskarsson/neural-lam-dev/blob/research/docs/reproduce_paper_sample.md

also #69 is mostly talking about danra, so you shouldn't close that issue with this PR.

So to rephrase, the issue about the COSMO hello world example is: how to bring the notebook from the research branch into main and get it to run in a reasonable amout of time. The preparation and hosting of a readuced cosmo sample is very much the core work. I do understand that hosting data is not simple/cheap, but it is required here. Since this PR is not solving the issues I'll mark it as draft for now. If you would like to start working on the example data hosting I suggest to contact @leifdenby as he has some options from what I understand 😉

[sadamov]

Okay I organized the hello_world issue and PRs:

• Add DANRA tutorial notebook with pytest-nbmake #69 now exclusively talks about DANRA -> assigned to @Sharkyii
• Add DANRA tutorial notebook with pytest-nbmake (#69) #577 is the official and only PR for the DANRA notebook
• Add COSMO example notebook for onboarding #587 is a new issue only targeting COSMO -> would like to assign to @info-gallary and @PalaniappanR02 to collaborate
• Add COSMO example notebook for onboarding (#587) #392 is the existing COSMO notebook that now closes Add COSMO example notebook for onboarding #587

If there was some oversight let me know, tried my best to look through all previous comms.

[GiGiKoneti]

Hi @sadamov and @leifdenby!

I would love to help take this PR over the line by preparing and hosting the reduced COSMO sample dataset.

If you can share a tiny slice of the real COSMO Zarr data (or point me to where I can access a small subset), I can structure it to match the notebook config and:

1. Package it into a lightweight, CPU-friendly format.
2. Host it on Hugging Face Datasets or coordinate with Leif to upload it to the ECMWF object store.
3. Update the notebook to download and run on this real sample out-of-the-box.

Let me know if this direction works, and I can start on this immediately!

[sadamov]

@GiGiKoneti let's give the current assigness a few days to respond.

1. the cosmo data is in the eth research collection, the research notebook linked in the issue body contains the links
2. Leif most likely will only see your message on slack. In a previous message he said hosting on DMI S3 bucket should be okay for small sizes

[sadamov]

@GiGiKoneti It looks like the original assignees are not active anylonger. Please comment here: #587 if you are still interested in finishing this work.

[info-gallary]

Hi @sadamov, yes, I’m still interested in contributing. Apologies for the delayed response, I was tied up with other commitments and missed the update. I’m happy to collaborate with @GiGiKoneti and help with the COSMO notebook integration/testing once the reduced dataset is available on the DMI S3 storage. Please let me know where I can be most useful.

[GiGiKoneti]

Hi @sadamov, yes, I’m still interested in contributing. Apologies for the delayed response, I was tied up with other commitments and missed the update. I’m happy to collaborate with @GiGiKoneti and help with the COSMO notebook integration/testing once the reduced dataset is available on the DMI S3 storage. Please let me know where I can be most useful.

I have the reduced COSMO Zarr sample and graph packaged and ready to link once the S3 storage is live. Once uploaded, it would be great to have your help validating the notebook execution and runtime benchmarks across CPU/GPU environments.

[sadamov]

Heads-up: we discussed the notebook format in the dev meeting and ran a poll in Slack afterwards, marimo won over Jupyter. Going forward notebooks in this repo are marimo .py files under docs/notebooks/, no .ipynb in git anymore.

For this PR that means: convert the notebook with marimo convert <name>.ipynb -o <name>.py, drop the .ipynb, and check it runs top to bottom with python docs/notebooks/<name>.py (marimo notebooks are plain python scripts, so no kernel is needed and ruff lints them like the rest of the code). marimo edit <name>.py if you want to tidy up cells after the conversion.

Sorry for the churn, I know this lands mid-review. Let me know if smth is unclear or if you would rather I push the conversion to your branch.

Specifically here: COSMO_example.ipynb to COSMO_example.py. Since #577 sets up the notebook CI, I suggest to sync with @Sharkyii so both notebooks run the same way.