feat: focus area

[icedoom888]

Description

Introduce a focus_area attribute to plot callbacks to focus the plot on a specific area.

Spawns from: #782

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.

---

📚 Documentation preview 📚: https://anemoi-training--799.org.readthedocs.build/en/799/

---

📚 Documentation preview 📚: https://anemoi-graphs--799.org.readthedocs.build/en/799/

---

📚 Documentation preview 📚: https://anemoi-models--799.org.readthedocs.build/en/799/

[JPXKQX]

Hi Alberto, I have reviewed the changes. The implementation looks very clean now! I left some minor comments. I also have a more general question: how should the focus_area be specified in multiple datasets? I see different options:

• You can only instantiate one focus_area, which is applied to all datasets.
• Add a new dataset_names key under the focus_area field to specify in which datasets the focus area should be applied to.
• A dict-like focus area similar to training_loss, validation_metrics, ... with the dataset_name as key.

[icedoom888]

Hi Alberto, I have reviewed the changes. The implementation looks very clean now! I left some minor comments. I also have a more general question: how should the focus_area be specified in multiple datasets? I see different options:

• You can only instantiate one focus_area, which is applied to all datasets.
• Add a new dataset_names key under the focus_area field to specify in which datasets the focus area should be applied to.
• A dict-like focus area similar to training_loss, validation_metrics, ... with the dataset_name as key.

@JPXKQX in the current code:

if focus_areas is None:
    focus_areas = [None] * len(self.dataset_names)

# Focus area for plotting
self.focus_masks = [
    build_spatial_mask(
        node_attribute_name=fa.get("mask_attr_name", None) if fa is not None else None,
        latlon_bbox=fa.get("latlon_bbox", None) if fa is not None else None,
    )
    for fa in focus_areas
]

So you can specify different focus areas for each dataset, do you still prefer the dict approach?

[icedoom888]

Update the configs structure, now you can define focus areas for each dataset in the multi dataset setting, and all will be dispayed for each cllback. This also reduces the number of callbacks needed.

Here an example:

focus_areas:
      datasets:
        data: 
          - latlon_bbox: [30.0, -20.0, 60.0, 40.0]
          - latlon_bbox: [18.0, 73.0, 54.0, 135.0]
          - latlon_bbox: [0.0, -90.0, 180.0, 90.0]

Whic produces the following:

[JPXKQX]

Is this change is intended?

[icedoom888]

This is a new file alltogether, just downloaded and uploaded it