Limited Area Modeling (LAM)

AnemoI Graphs brings another level of flexibility to the user by allowing the definition of limited area graphs.

Define hidden nodes in area of interest

The user can use a regional dataset to define the data nodes over the region of interest. Then, it can define the hidden nodes only over the region of interest using any of the LimitedArea_____Nodes classes.

nodes:
  data:
    node_builder:
      _target_: anemoi.graphs.nodes.ZarrDatasetNodes
      dataset:
        cutout:
          - dataset: regional-dataset.zarr
            thinning: 25
          - dataset: /path/to/global-dataset.zarr
        adjust: all
        min_distance_km: 10
    attributes:
      cutout_mask:
        _target_: anemoi.graphs.nodes.attributes.CutOutMask
  hidden:
    node_builder:
      _target_: anemoi.graphs.nodes.LimitedAreaTriNodes
      resolution: 5
      reference_node_name: data
      mask_attr_name: cutout_mask
    attributes: ...

edges: ...

Cut out regional dataset into a global dataset

In this case, the user may want to include boundary forcings to the region of interest. AnemoI Graphs allows the user to use two datasets to build the data nodes, combining nodes from the LAM dataset and the global dataset (as boundary forcings). The class ZarrDatasetNodes allows this functionality:

nodes:
  data:
    node_builder:
      _target_: anemoi.graphs.nodes.ZarrDatasetNodes
      dataset:
        cutout:
          - dataset: regional-dataset.zarr
            thinning: 25
          - dataset: /path/to/global-dataset.zarr
        adjust: all
        min_distance_km: 10
    attributes: ...
  hidden: ...

edges: ...

The ZarrDatasetNodes supports an optional thinning argument which can be used to sampling points from the regional dataset to reduce computation during development stage.

In addition, this node builder class will create an additional node attribute with a mask showing which node correspond to each of the two datasets.

>>> graph
HeteroData(
   data={
      x=[40320, 2],
      node_type='ZarrDatasetNodes',
      area_weight=[40320, 1],
      cutout_mask=[40320, 1],
   }
)

Define hidden nodes over region of interest

Once the data nodes are defined, the user can define the hidden nodes only over the region of interest. In this case, the area of interest is defined by the data nodes masked by the cutout attribute.

nodes:
  data: ...
  hidden:
    node_builder:
      _target_: anemoi.graphs.nodes.LimitedAreaTriNodes
      resolution: 5
      reference_node_name: data
      mask_attr_name: cutout_mask

edges: ...

>>> graph
HeteroData(
   data={
      x=[40320, 2],
      node_type='ZarrDatasetNodes',
      area_weight=[40320, 1],
      cutout_mask=[40320, 1],
   },
   hidden={
      x=[10242, 2],
      node_type='TriNodes',
   }
)

Adding edges

The user may define the edges using the same configuration as for the global graphs.

nodes:
  data: ...
  hidden: ...

edges:
  # A) Encoder connections
  - source_name: data
    target_name: hidden
    edge_builders:
    - _target_: anemoi.graphs.edges.CutOffEdges
      cutoff_factor: 0.7
  # B) Decoder connections
  - source_name: hidden
    target_name: data
    target_mask_attr_name: cutout
    edge_builders:
    - _target_: anemoi.graphs.edges.KNNEdges
      num_nearest_neighbours: 3
 # C) Processor connections
  - source_name: hidden
    target_name: hidden
    edge_builders:
    - _target_: anemoi.graphs.edges.MultiScaleEdges
      x_hops: 1