Unable to use UniformWeights config option for area_weights

[veech]

What happened?

I'm unable to create a graph using UniformArea weights option. Perhaps I'm missing something but looking at the code here, it looks like it should be an option.

Here is my config override for the graph:

defaults:
  - hardware: example
  - graph: multi_scale

graph:
  attributes:
    nodes:
      area_weight:
        _target_: anemoi.graphs.nodes.attributes.UniformWeights
        norm: unit-max
        fill_value: 0
...

And the errors:

Traceback (most recent call last):
  File "/DIR/anemoi/.venv/lib/python3.12/site-packages/anemoi/training/train/train.py", line 499, in main
    AnemoiTrainer(config).train()
    ^^^^^^^^^^^^^^^^^^^^^
  File "/DIR/anemoi/.venv/lib/python3.12/site-packages/anemoi/training/train/train.py", line 70, in __init__
    self.config = BaseSchema(**config)
                  ^^^^^^^^^^^^^^^^^^^^
  File "/DIR/anemoi/.venv/lib/python3.12/site-packages/pydantic/main.py", line 253, in __init__
    validated_self = self.__pydantic_validator__.validate_python(data, self_instance=self)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
pydantic_core._pydantic_core.ValidationError: 26 validation errors for BaseSchema
graph.nodes.data.attributes.area_weight.PlanarAreaWeightSchema.fill_value
  Extra inputs are not permitted [type=extra_forbidden, input_value=0, input_type=int]
    For further information visit https://errors.pydantic.dev/2.11/v/extra_forbidden
graph.nodes.data.attributes.area_weight.SphericalAreaWeightSchema._target_
  Input should be 'anemoi.graphs.nodes.attributes.SphericalAreaWeights' [type=literal_error, input_value='anemoi.graphs.nodes.attributes.UniformWeights', input_type=str]
    For further information visit https://errors.pydantic.dev/2.11/v/literal_error
graph.nodes.data.attributes.area_weight.CutOutMaskSchema._target_
  Input should be 'anemoi.graphs.nodes.attributes.CutOutMask' [type=literal_error, input_value='anemoi.graphs.nodes.attributes.UniformWeights', input_type=str]
    For further information visit https://errors.pydantic.dev/2.11/v/literal_error
...
...

Perhaps I'm doing something wrong? Are there more config options that I need to change?

I'm attempting to initialize a model using 1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2 from the ARCO ERA5 dataset and I was told that I should try using UniformWeights until #223 makes it through.

What are the steps to reproduce the bug?

Execute a training run using the above config overrides for graph.

Version

anemoi-training==0.4.0.post17, anemoi-graphs==0.5.2.post11

Platform (OS and architecture)

Linux 6.11.0-25-generic #25~24.04.1-Ubuntu SMP PREEMPT_DYNAMIC

Relevant log output

Accompanying data

No response

Organisation

No response

[veech]

I attempted to create the graph manually via anemoi-graphs with the exact default multi-scale graph that is in the default training config with my UniformWeights changes and that seemed to work fine:

nodes:
  data:
    node_builder:
      _target_: anemoi.graphs.nodes.AnemoiDatasetNodes
      dataset: ${dataset}
    attributes:
      area_weights:
        _target_: anemoi.graphs.nodes.attributes.area_weights.UniformWeights
        norm: unit-max
  
  hidden:
    node_builder:
      _target_: anemoi.graphs.nodes.TriNodes
      resolution: 5
    area_weights:
      _target_: anemoi.graphs.nodes.attributes.area_weights.UniformWeights
      norm: unit-max

edges:
  # Encoder connections
  - source_name: data
    target_name: hidden
    edge_builders:
      - _target_: anemoi.graphs.edges.CutOffEdges
        cutoff_factor: 0.6
        source_mask_attr_name: null
        target_mask_attr_name: null
    attributes:
      edge_length:
        _target_: anemoi.graphs.edges.attributes.EdgeLength
        norm: unit-std
      edge_dirs:
        _target_: anemoi.graphs.edges.attributes.EdgeDirection
        norm: unit-std

  # Processor connections
  - source_name: hidden
    target_name: hidden
    edge_builders:
      - _target_: anemoi.graphs.edges.MultiScaleEdges
        x_hops: 1
        scale_resolutions: 5
        source_mask_attr_name: null
        target_mask_attr_name: null
    attributes:
      edge_length:
        _target_: anemoi.graphs.edges.attributes.EdgeLength
        norm: unit-std
      edge_dirs:
        _target_: anemoi.graphs.edges.attributes.EdgeDirection
        norm: unit-std

  # Decoder connections
  - source_name: hidden
    target_name: data
    edge_builders:
      - _target_: anemoi.graphs.edges.KNNEdges
        num_nearest_neighbours: 3
        source_mask_attr_name: null
        target_mask_attr_name: null
    attributes:
      edge_length:
        _target_: anemoi.graphs.edges.attributes.EdgeLength
        norm: unit-std
      edge_dirs:
        _target_: anemoi.graphs.edges.attributes.EdgeDirection
        norm: unit-std

post_processors: []

Results:

2025-05-09 10:36:03 INFO Cleaning graph.
2025-05-09 10:36:03 INFO _grid_reference_distance deleted from graph.
2025-05-09 10:36:03 INFO _dataset deleted from graph.
2025-05-09 10:36:03 INFO _grid_reference_distance deleted from graph.
2025-05-09 10:36:03 INFO _node_ordering deleted from graph.
2025-05-09 10:36:03 INFO _resolutions deleted from graph.
2025-05-09 10:36:03 INFO _nx_graph deleted from graph.
2025-05-09 10:36:03 INFO Graph saved at graphs/graph-equi.pt.

📦 Path       : graphs/graph-equi.pt
💽 Size       : 2.4 MiB (2564880)

🪩  Nodes summary

   Nodes name │ Num. nodes │   Attributes │ Attribute dim │ Min. latitude │ Max. latitude │ Min. longitude │ Max. longitude
   ───────────┼────────────┼──────────────┼───────────────┼───────────────┼───────────────┼────────────────┼───────────────
   data       │       8192 │ area_weights │             1 │      -88.5937 │       88.5937 │              0 │        357.188
   hidden     │      10242 │              │             0 │           -90 │            90 │       -178.606 │            180
   ───────────┴────────────┴──────────────┴───────────────┴───────────────┴───────────────┴────────────────┴───────────────


🌐  Edges summary

   Source │ Target │ Num. edges │ Isolated Source │ Isolated Target │ Attribute dim │                     Attributes
   ───────┼────────┼────────────┼─────────────────┼─────────────────┼───────────────┼───────────────────────────────
   data   │ hidden │      12816 │               0 │            1396 │             3 │ edge_length(1D), edge_dirs(2D)
   hidden │ hidden │      81840 │               0 │               0 │             3 │ edge_length(1D), edge_dirs(2D)
   hidden │ data   │      24576 │              12 │               0 │             3 │ edge_length(1D), edge_dirs(2D)
   ───────┴────────┴────────────┴─────────────────┴─────────────────┴───────────────┴───────────────────────────────


📊 Attribute distributions

   Type │ Source          │         Name │         Dtype │      Min. │         Mean │    Max. │ Std. dev.
   ─────┼─────────────────┼──────────────┼───────────────┼───────────┼──────────────┼─────────┼──────────
   Node │ data            │ area_weights │ torch.float32 │         1 │            1 │       1 │         0
   Edge │ data-->hidden   │  edge_length │ torch.float32 │ 0.0549464 │      2.84916 │ 4.25287 │         1
   Edge │ data-->hidden   │    edge_dirs │ torch.float32 │  -1.41419 │ -4.10759e-08 │ 1.41419 │         1
   Edge │ hidden-->hidden │  edge_length │ torch.float32 │  0.720712 │      1.14287 │ 13.0884 │         1
   Edge │ hidden-->hidden │    edge_dirs │ torch.float32 │  -1.41421 │  3.08802e-10 │ 1.41421 │         1
   Edge │ hidden-->data   │  edge_length │ torch.float32 │ 0.0389353 │      2.79869 │ 4.86443 │         1
   Edge │ hidden-->data   │    edge_dirs │ torch.float32 │   -1.4142 │  0.000417567 │  1.4142 │         1
   ─────┴─────────────────┴──────────────┴───────────────┴───────────┴──────────────┴─────────┴──────────

🔋 Graph ready.

Perhaps there is a mismatch in the hydra config between training and graph?