[RFC/Design]: Support for multiple datastores

[leifdenby]

I have been thinking about how to generalise #635 to allow for using sources of data that would describe say multiple point-wise observations and/or regular gridded data. Although I am not suggesting that we alter #635 as-is, with which the will be support for a datastore (which used both for input and output of the model) and datastore_boundary (only used for input), I thought it might be useful to think about how that work could be evolved on down the line. With #625 Neural-LAM configuration will assume a fixed set of datastore roles, such as a main datastore and optionally a boundary datastore. This does not scale well to configurations that combine multiple data sources, for example regional reanalysis, global boundary data, synoptic observations, or satellite-derived inputs.

The aim is to replace role-specific datastore config keys with a single datastores mapping. Each datastore should be named by the user and explicitly declare how it is used by the model.

Proposal

Introduce a datastores dictionary in NeuralLAMConfig, keyed by datastore name. Each entry should contain the datastore selection fields plus a new use section defining which data categories are used as model inputs and outputs.

Suggested config shape:

datastores:
  danra:
    use:
      input: [state, static, forcing]
      output: [state]
    kind: mdp
    config_path: danra.datastore.yaml

  era5:
    use:
      input: [state, forcing]
      output: []
    kind: mdp
    config_path: era5.datastore.yaml

Suggested dataclasses:

@dataclasses.dataclass
class DatastoreUse:
    input: list[str]
    output: list[str]

@dataclasses.dataclass
class DatastoreSelection:
    use: DatastoreUse
    kind: str
    config_path: str

@dataclasses.dataclass
class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard):
    datastores: Dict[str, DatastoreSelection]
    training: TrainingConfig = dataclasses.field(default_factory=TrainingConfig)

Motivation

This makes datastore roles explicit and extensible. The datastore name describes the data source, while use describes how the data source contributes to the model, and the graph connectivity that will be required to be made.

This supports configurations such as:

# LAM trained on regional reanalysis with global boundary data
datastores:
  danra:
    use:
      input: [state, static, forcing]
      output: [state]
    kind: mdp
    config_path: danra.datastore.yaml

  era5:
    use:
      input: [state, forcing]
      output: []
    kind: mdp
    config_path: era5.datastore.yaml

# LAM trained with regional reanalysis, global reanalysis, and synoptic observations
datastores:
  danra:
    use:
      input: [state, static, forcing]
      output: [state]
    kind: mdp
    config_path: danra.datastore.yaml

  synop:
    use:
      input: [forcing]
      output: []
    kind: mdp
    config_path: synop.datastore.yaml

  era5:
    use:
      input: [state, forcing]
      output: []
    kind: mdp
    config_path: era5.datastore.yaml

# Global reanalysis with satellite-derived inputs
datastores:
  era5:
    use:
      input: [state, static, forcing]
      output: [state]
    kind: mdp
    config_path: era5.datastore.yaml

  amsua:
    use:
      input: [forcing]
      output: []
    kind: mdp
    config_path: amsua.datastore.yaml

[sadamov]

@leifdenby Agree with the dict-of-named-datastores direction. #656 ships the bare dict shape as an alternative to #635 (same main base, not stacked). Per-sample return shape is unchanged from current main, so #656 does not require @joeloskarsson's model-side work in lockstep. The boundary tensor and per-source ForecastBatch land in his follow-up.

I propose to extend your schema in three places. None of these are in #656; each lands in its own follow-up so we do not declare schema fields that have no runtime effect:

datastores:
  interior:
    kind: mdp
    config_path: danra.datastore.yaml
    inputs:
      state: [t2m, u10m, v10m, ...]
      forcing: [toa_rad, ...]
      static: [lsm, orography]
    outputs:
      state: [t2m, u10m, v10m, ...]   # prognostic
      diagnostic: [tp, cape]          # predict-only
  boundary:
    kind: mdp
    config_path: era5.datastore.yaml
    inputs:
      forcing: [...]

1. Variable include-lists per category. Reuses the same zarrs across input-output experiments. Follow-up PR on my side; lands together with Joel's model-side adapter so feature dims agree.

2. Diagnostic outputs via set membership. Derive prognostic/diagnostic from outputs.state vs outputs.diagnostic rather than from category == "state". Orthogonal to the schema and the filter work, separate follow-up - what do you think?

3. Variable name collisions across datastores (e.g. synop's t2m and interior's t2m). Validate at config load and raise, pointing at mdp's dim_mapping.name_format for new builds or a small zarr-python snippet (zarr.open_group(path, mode="r+")["{category}_feature"][:] = new_names) for existing zarrs. Lands together with the filter follow-up since the validator only matters once outputs: are declared.

To be clear: #656 and #635 are alternatives, exactly one should land. Both can land on their own, but the follow-ups behave differently: with #635 the model-side wiring is an independent later PR (the 5th tuple element is harmlessly ignored until then). With #656 my filtering follow-up and @joeloskarsson's model-side adapter must land together since they share feature dims. Plotting in #636 stays on #635 for now and rebases onto whichever lands.

@joeloskarsson does the dict shape in #656 match what you had in mind on the model side? Or are you opposed to this theoretical abstraction layer, because it means our paper approach is never in main.

[joeloskarsson]

My opinion is generally to just merge the explicit interior + (optional) boundary datastore now first, and do these kinds of generalizations later. Potentially we could make it so that these two specific datastores are not hard-coded, but we make it general enough that potentially other ones could be added. But I would not want to introduce some input/output fields in the config immediately, as this seems like it would require additional work on logic and slow things down.

In general I am not sure if I understand the design intention here. In particular, the notions of state/forcing/static have specific meaning in the codebase now, and I don't get the motivation to move away from these. At the moment we always have:

Tensor	Input	Output
State(Prognostic)	✓	✓
Forcing	✓
Static	✓
Diagnostic (not used atm, could be added)		✓

This spans all possibilities, and one can assign data to these different categories by placing it in the corresponding tensor. Following this logic there is never a need to specify what should be inputs and outputs and something like

era5:
    use:
      input: [state, forcing]
      output: []

does not really make sense. But maybe I should understand the proposal here as wanting to move away from these pre-defined notions, and instead have datastores containing arbitrary arrays? But I guess I don't see the shape that would take, as it would have consequences far beyond the datastores and e.g. also on the model side. What is lacking in these pre-defined notions that needs generalization in this way?

[leifdenby]

@joeloskarsson: My opinion is generally to just merge the explicit interior + (optional) boundary datastore now first, and do these kinds of generalizations later.

Yes, I think that is right approach too, but I thought it would be useful to see if we could tweak the config now so that people won't have to change them.

the notions of state/forcing/static have specific meaning in the codebase now, and I don't get the motivation to move away from these

yes, I agree. We have (somewhat implicit I would argue) definitions of state, static and forcing and their use. By this what I mean is that currently:

1. a datastore must provide state variables and they are used for both input and output.
2. if a datastore contains static data this will be used for input
3. if a datastore contains forcing data this will be used for input

(as your table clearly lays out)

What I was trying to see was whether it might be a good idea to explicitly state in the config how a get category of data (state/static/forcing/prognostic) is actually used. And by used I was thinking of what becomes an input to the model and what should be include in the input of the model, and thereby defines what graph connectivity that will be needed from the grid nodes from different datastores and to the mesh, and back again from the mesh nodes and to the grid nodes.

Reading your reply though @joeloskarsson I think actually we already have enough with the input/output-use-and-graph-connectivity be implied by the presence of different categories in each datastore. Maybe this just needs to be explicitly written down somewhere. What I also like about this is that the use of the data in a given datastore is given within datastore config itself. It does mean that we have to relax the requirement that a datastore has a state category of data. Maybe what we need is to ensure that at least one of the datastores defined contains state category of data?

With that in mind I have adjusted one of my examples above for what I would propose:

# LAM trained with regional reanalysis, global reanalysis, and synoptic
# observations, which is trained to use ERA5 for input only, and DANRA and SYNOP
# for both input and output
datastores:
  danra:
    kind: mdp
    # contains [state, static, forcing] categories, which will be used for both input and output of the model
    config_path: danra.datastore.yaml

  synop:
    kind: mdp
    # contains [state] category, so will be used for both input and output of the model
    config_path: synop.datastore.yaml

  era5:
    kind: mdp
    # contains [static, forcing] category, so will be used for input only of the model
    config_path: era5.datastore.yaml

For the DANRA + ERA5 LAM training the config of #635 could we be adjusted as:

# LAM trained with regional reanalysis, global reanalysis, and synoptic
# observations, which is trained to use ERA5 for input only, and DANRA and SYNOP
# for both input and output
datastores:
  danra:
    kind: mdp
    # contains [state, static, forcing] categories, which will be used for both input and output of the model
    config_path: danra.datastore.yaml

  era5:
    kind: mdp
    # contains [static, forcing] category, so will be used for input only of the model
    config_path: era5.datastore.yaml

This means that there is no explicit notion of a "boundary" datastore. Given the state/static/forcing categories of data, and the presence of these categories of data explicitly defining the input/output and graph structure, I think we don't think we need the notion of a "boundary" datastore. Said a different way: a datastore (like the era5 one above) which contains no state data will result in no edges connecting from the mesh to the grid nodes from this datastore.

@sadamov: I propose to extend your schema in three places. None of these are in #656; each lands in its own follow-up so we do not declare schema fields that have no runtime effect:

I would rather not put the requirement for use of listing variables in this config as well. To me that introduces duplication with the datastore config where the variables to use are also named. If we go the route of the input/output and graph connectivity being implied through the presence of the different categories of data I don't think we need the other suggestions here :)

[sadamov]

my extensions were heavily inspired by the weathergenerator, which I think does a rather good job with the config yaml files to run experiments. So to be concrete: I'll make #656 a draft for now, but keep it in case we go more flexible in the future (e.g. sparse obs datastores). @leifdenby can you please add your ideas directly as a review to #635, to make sure that I don't develop smth further outside of DMI's requirements. Especially the ForecastBatch contract also in combination with #637
Also keep in mind the boundary plotting requires some sort of "boundary" notion from the datastore.

[joeloskarsson]

I agree with everything from @leifdenby's comment above, and think that the examples given there are exactly what I had in mind. They generalize the multiple-datastores only to the point of having different names for them, but nothing more (since if we assume the state/static/forcing meaning, more is not needed). But yes, we should make this meaning of state/static/forcing explicit in documentation.

@sadamov the plotting is a good point. While for the interior state we could assume that all state variables in all datastores (but we only need to implement this for 1 now) should be plotted, we somehow have to configure which datastore is the boundary. I guess this is not so different from having to configure (state, boundary forcing) pairs to plot together, just that we also have to specify which datastore the boundary forcing should come from. I see a few solutions:

1. We could just restrict everything to <= 2 datastores now, and these have to be in the order interior, boundary. This would also simplify things on the model side, with the separate encoding MLPs.
2. We could have some form of role parameter associated with each datastore:

datastores:
  danra:
    kind: mdp
    config_path: danra.datastore.yaml
    role: interior

  era5:
    kind: mdp
    config_path: era5.datastore.yaml
    role: boundary_forcing

but this feels a bit like duplicate information to what the datastores contain (state/forcing).
3. We could configure this just as part of the plotting config. So that one has to specify which datastore to take boundary forcing variables from to include in plots.

I think both 1 and 3 are ok solutions. 2 might require some more thinking.