Enabling training on interior and boundary datastore

[leifdenby]

@sadamov and @joeloskarsson have been doing great work on #93 and #90 to enable training and inference where the domain interior and boundary come from different datastores.

I am making this issue to start trawling through changes on https://github.com/joeloskarsson/neural-lam-dev/tree/research that relate to this but are not currently included in the branches that the above PRs currently contain. The aim here is to have a complete list of changes and built a plan to get #93 and #90 merged.

• bugfix for boundary forcing time sampling: a2bcf23
• feat: dynamic time deltas for boundary forcing: e58e334
• feat: allow for using forecast data on boundary: fb820e2
• feat: writing eval predictions to zarr: 39eb6e2 (appears to use implementation discussed on Add writing to zarr dataset for eval-mode of trained models #104)
• bugfix restoration of trainer state - e7d11c9, I think this reverts b1c7ef8 hack for checkpoint loading without setting opt/sched state
• bugfix: ensure that eval actually uses defined eval_split - 819b47e (MR created Use correct --eval data split #139)
• bugfix for clamping of predictions during finetuning - 7985f2c already fixed on main in Fix nans in gradient from inverse softplus #123
• bugfix for boundary datastore being None - f783f6d
• bugfix for scaling of step_diff - 2e46ecc
• feat: allow for separate embedding for interior/boudnary (I think) - 9468fbb
• feat: handling of gaps in training data ("swiss cheese strategy") - 20a5464 and ba9df3a
• feat: cosine annealing learning rate scheduler - 20a5464
• feat: check-pointing based on validation loss - e3a10fb
• feat: time subsampling on boundary vs interior domain - 18a16d4
• feat: optional gradient checkpointing during roll-out - 5f4705b
• feat: cli arg for number of HPC nodes - d95c1bf
• feat: boundary time-delta encoding - 284a954
• bugfix: indexing error when extracting time from batch - 6672756

[sadamov]

@joeloskarsson can you add your items here please

[joeloskarsson]

Crucial things to merge

These are the things that are crucial to merge to have a "realistic boundary setup", as in our paper. In some categories:

Graph

In WMG

• Separate g2m connectivity arguments for boundary and interior grid nodes (joeloskarsson/weather-model-graphs@974ad07)
  • This is absolutely necessary to be able to use wmg to create graphs without having disconnected nodes.
• Decoding mask, to only include a subset of the grid nodes (the interior grid nodes) in m2g (Add a decoding mask option to only include subset of grid nodes in m2g weather-model-graphs#34)
• Allow for filtering created mesh nodes to convex hull of all grid nodes.
  • Crucial, as without this mesh nodes in g2m will easily end up with degree 0, creating all kinds of problems.
  • code: joeloskarsson/weather-model-graphs@c824c23
  • Issue: Add an option to crop mesh nodes to convex hull of grid coordinates weather-model-graphs#40
• Add check that no nodes in g2m have degree 0
  • Not strictly crucial, but this is a very good idea as graphs wont work otherwise. Could be done together with above.
  • code: https://github.com/joeloskarsson/weather-model-graphs/blob/33d937a0a04848eb0a3d1d86ebcd4b63e4ec4eda/src/weather_model_graphs/create/base.py#L199-L207
  • Issue: Assert that all nodes are in g2m weather-model-graphs#42
• Fix bug where nodes could be dropped when filtering edges (joeloskarsson/weather-model-graphs@6556dbb#diff-72daea2494478edf48e319a44c42b839b9a6175a03dd1b2fa97ca8db212a79caL300-L304). They were dropped if they ended up with no edges after filtering, but this messes up the node indexing.
  • Fixed in Fix the bug with edgeless nodes being dropped weather-model-graphs#51
• Make sure grid is only connected to bottom level in hierarchical graph joeloskarsson/weather-model-graphs@04c8052
• Make computation of number of leaf nodes in mesh hierarchy more robust (joeloskarsson/weather-model-graphs@a2ec24a)

In neural-lam

• Change main graph creation script. Remove create_graph.py and replace it with build_rectangular_graph.py, which calls wmg.
  • This is the switch to using wmg in neural-lam so corresponds to Add graph creation functionality using weather-model-graphs #83 and also solves Adapt graph generation script to general limited areas #4.
  • Code: https://github.com/joeloskarsson/neural-lam-dev/blob/research/neural_lam/build_rectangular_graph.py Can probably be used without much cleanup.
  • Will be fixed in Add create_graph_with_wmg.py CLI using weather-model-graphs #596
• Move all edge_index reindexing from InteractionNetwork class to utils.load_graph
  • This is crucial to get only forecasting interior to work correctly.
  • Code: Should correspond to joeloskarsson@95bebe7#diff-d4abab0d198cc8e511f26170efe7665fbbd20ffe41dd89dd6e6729441e37f348, but we don't need the tri_graph case.
  • Should also remove the notice

    neural-lam/neural_lam/models/base_graph_model.py

    Lines 22 to 23 in 78a7b1f

    	# NOTE: (IMPORTANT!) mesh nodes MUST have the first
    	# num_mesh_nodes indices,

    , as that will no longer be necessary.
• Rescale graph features by maximum coordinate. When switching from create_graph.py to wmg this is not done before saving the graph features.
  • code: https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/models/base_graph_model.py#L42-L54
  • It would maybe be better to do this as the graph is loaded, rather than in a model class.
  • As part of this, to make this work we need to implement the __imul__ and __itruediv__ methods in the BufferList, to enable direct standardization. Code: https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/utils.py#L37-L44
• Upgrade from requiring python 3.9 to python 3.10. As wmg requires python 3.10 this is necessary to work with neural-lam and wmg together. [Maint] update tests to use recent python versions #173

Boundary data from separate datastore

• Make reference to boundary datastore config part of the neural-lam config.
  • See changes to https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/config.py#L117 and https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/config.py#L169
• Remove boundary_mask from datastore.
  • We no longer have a boundary mask, as all grid points in the interior datastore are in the interior.
  • We need to instead introduce a way to easily get the number of grid points in a datastore: https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/datastore/mdp.py#L505 This should probably be already in the base datastore, rather than in the MDP datastore.
• Rework all input data handling in model classes to account for separate interior and boundary datasets.
  • We get an additional boundary tensor as part of each batch (https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/weather_dataset.py#L967) so the models must handle this correctly.
  • Code mainly in: joeloskarsson@315aa0f
  • This is pretty much the same thing as Separate interior state and boundary forcing to only predict state #93 , and the solution is based on the things I started there, but we might want to open a new PR instead of building on that one.
• (Highly related to above) Introduce separate interior and boundary MLPs in BaseGraphModel.
  • For code see https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/models/base_graph_model.py#L79-L100 and how they are used in the forward pass.
  • When working on this we had the --shared_grid_embedder option, for sharing embedder MLP between interior and boundary. However, there are so few situations (in particular https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/models/base_graph_model.py#L90) where this is applicable that I am not even sure if we should keep it.
  • Working on this, it's probably easiest to just merge the full base_graph_model from research, and then make any edits necessary. In particular we do not need the forcing time delta embedding.
• Rework WeatherDataset to handle boundary forcing from separate datastore.
  • This work is started in the PR Feat/boundary dataloader #90, but also here it might be a good idea to make a new PR.
  • The changes are so substantial that the best way to start merging this might be to start from the WeatherDataset class in https://github.com/joeloskarsson/neural-lam-dev/blob/research/neural_lam/weather_dataset.py and then try to polish things up.
  • Note that we have these arguments for temporal subsampling in there, but we ended up never using them (always set to 1) and I don't think all the code works correctly with that. We should drop that functionality as we merge this.
  • A key part of making this work was figuring out all the time offsets and how many samples we actually get in the dataset. There are some crucial helper functions that go with this, but are in the utils file: https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/utils.py#L464 and https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/utils.py#L538
  • With all these changes the dataset class can also handle forecast data correctly. That is needed for evaluation (see below).

Visualization

• Use pcolormesh for plotting of more general geometry.
  • What we do here is to make a general plot_on_axis function that can then be used whenever a spatial field should be plotted. This is good to reduce duplicate code in the plotting.
  • See https://github.com/joeloskarsson/neural-lam-dev/blob/research/neural_lam/vis.py#L77 for code.
  • We didn't plot the boundary here yet, but that would be nice as well. Tracked in Plotting Boundaries in train_model.py --eval not just interior #108.
  • PR: 166 plotting with correct projection config using pcolormesh #168

Evaluation

I suppose it can be discussed if this is absolutely crucial, but for a "realistic setting" we really need to be able to have forecasts as boundary forcing.

• Properly handle forecast data in MDP datastore
  • There were some things that did not work correctly with forecast data in the datastore, that needed to be fixed.
  • This should be all the code for this: joeloskarsson@b83354e
• Allow for loading data statistics from another (third!) datastore.
  • This is necessary if you train using a one datasource (e.g. reanalysis) as boundary forcing, but want to run evaluation using another (e.g. forecasts). Then we can not take the data statistics from our new datastore, as that might give a different range of values than the model has been trained to work with. So we then need to use the data statistics from the same datastore as used during training, and rescale the new data using these.
  • The way we solved this was to introduce a third reference to a "stats" datastore in the config that only the data stats would be loaded from.
  • Code: joeloskarsson@6a68d32 ( + joeloskarsson@c9d9707 I suppose)

I'll continue with some "nice to have things" below 😄

[joeloskarsson]

Nice to have things

These should probably not be prioritized, but would be nice additions. Mostly good to have them written down here for if/when we want to look at it in the future.

• Number of sanity steps argument
  • joeloskarsson@3816170
  • Can be very nice when debugging training, to not have to always wait on sanity validation steps.
• Subset the initialization times to run evaluation on.
  • Currently implemented through a wrapper-class https://github.com/joeloskarsson/neural-lam-dev/blob/research/neural_lam/weather_dataset.py#L1075, but only works for 00 and 12 UTC. Should be generalized if we want to merge this features.
  • This is useful as we often want to train on all possible init times, but only eval at some (e.g. benchmarks like weatherbench only consider forecasts started at 00/12 during eval.
• Change how the --restore_opt option sets up model training to better reflect the intended workflow in pytorch geometric
  • This is a bit messy now, and I think this is a much nicer solution.
  • code: joeloskarsson@e7d11c9
• Allow for separate hidden dims for grid and mesh, using a --hidden_dim_grid option, and linear layers projecting between these hidden dimensionalities in BaseGraphModel.
  • Note that # Bottom mesh level is first embedded to hidden dim of grid in base_hi_graph_model
  • Code: joeloskarsson@9468fbb
  • As this is a small architectural change I think some consideration is warranted for how the user can choose this architecture or the "standard one" with the same hidden dimensionality everywhere.
• Save all evaluation metrics in a single object, and upload to logging service.
  • We used a pickled xarray.Dataset for this, but other options are possible.
  • This is very good if you later want to plot your evaluation results in some specific way, e.g. for a publication.
  • Code: https://github.com/joeloskarsson/neural-lam-dev/blob/a44d432ed1c5c0f0af65e5913a137d63ae160b1d/neural_lam/models/ar_model.py#L836-L868
• Output and save all evaluation forecasts to zarr (Output predictions as zarr dataset #89 )
  • The code we used for this is just Add writing to zarr dataset for eval-mode of trained models #104 with some small modifications. I think that we can just continue on that PR and sort out the small issues there.
• Introduce a --plot_vars option to specify which variables should be plotted as examples.
  • Code: joeloskarsson@ae7ae4b
  • Very much needed when we work with models with lots of variables.
• Introduce (optional) gradient checkpointing during rollout
  • Code: joeloskarsson@5f4705b
• Substantially improve graph plotting in 3d using plotly
  • Many new options, earth texture, etc.
  • As this is pretty independent from everything else, it is likely very easy to merge :)
  • Code: https://github.com/joeloskarsson/neural-lam-dev/blob/research/neural_lam/plot_graph.py
• Triangular graphs
  • Code: 🔺-graphs joeloskarsson/neural-lam-dev#13
  • Triangular graph creation is done through separate code within neural-lam, and does not use wmg. Not neccesarily how we want to do it, but could be an option for first getting this to work.
  • A major consideration for this is that the triangular graph code uses some (very few) calls to the graphcast code. There are also some modified versions of graphcast functions. The question is if we want a graphcast dependency and how to handle that (likely a separate dependency group). One should note that with a dependency to graphcast there is also a dependency to jax.
  • Also licenses should be considered if we opt for just including small parts of the graphcast code directly here, instead of having a dependency.
• Add option to exclude certain time-intervals from dataset.
  • This is the so-called "swiss-cheese strategy", that pretty much splits one dataset into multiple ones, leaving holes in the excluded intervals, and then cocatenates them back. This is useful if you for example have some NaNs in your data at some date, but the way this is implemented right now is not the most efficient (it can create a lot of very similar dataset instances).
  • Code: joeloskarsson@20a5464

[joeloskarsson]

One addition here is that mllam/mllam-data-prep#72 would also be needed for mdp in order to properly use forecast data (i.e. including forcing of interest) as input through mdp.

[Jayant-kernel]

@leifdenby @sadamov @joeloskarsson I went through issue #138 and the linked work for #89 (including PR #104).

I’d like to help on #89 (Output predictions as zarr dataset) by building on the existing PR/discussion, especially around the multi-GPU zarr writing path and making it merge-ready.

If that sounds good, I can first review PR #104 in detail, summarize what is still pending, and then open a follow-up PR (or contribute directly to the existing branch/approach, depending on your preference).

[sahilkr31]

@joeloskarsson @sadamov @leifdenby

Hi,

I’ve gone through the discussion in #138 and the linked PRs, particularly around boundary handling and the evaluation updates.

From the evaluation section, it seems important to ensure forecast data is handled correctly in the MDP datastore and that normalization statistics can be loaded from a separate datastore when boundary and interior data come from different sources.

Before proposing any changes, I’d like to trace the current evaluation pipeline to understand how forecast data and normalization stats are currently applied, and where the separation should occur.

Would it make sense for me to focus on clarifying and improving this part of the evaluation logic? I’m happy to align with the current merge direction and start with a well-scoped piece.

[Jayant-kernel]

PR #242 addresses Issue #108 (plotting boundaries / crop_to_interior extent) and adds a test. #242