feat: activate minmium plotting for integration tests (#669)

## Description
Testing plots it's not part our current testing suite. Since most of
them are implemented via callbacks a simple way would be to activate a
minimal plotting set up for the integration test to ensure we don't
break them in the future.

This PR proposes those changes:
- Spectrum plot it's quite slow so for now it's deactivated for the SG
model
- Since the test dataset for LAM and SG is a small one we just test
plotting tp
- The plots for the interpolator are deactivated - while the callbacks
run the output produced is not scientifically correct. More work will be
done in the future to address this (@mc4117 - notified about this in
slack channel)
- In the hierarchical model - current plotting of edge features is not
support - this plot is for now deactivated. (Discussed with @icedoom888
)

cc: @VeraChristina - having those tests would also be useful in the
context of multiple datasets to ensure there is no feature regression
for plotting.


## What problem does this change solve?
<!-- Describe if it's a bugfix, new feature, doc update, or breaking
change -->

## What issue or task does this change relate to?
<!-- link to Issue Number -->

##  Additional notes ##
<!-- Include any additional information, caveats, or considerations that
the reviewer should be aware of. -->

***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.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

---------

Co-authored-by: Mariana Clare <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

.github/workflows/integration-tests-hpc.yml

@@ -52,7 +52,7 @@ jobs:
             #SBATCH --qos=ng
             #SBATCH --gpus=1
             #SBATCH --gres=gpu:1
-            #SBATCH --mem=30G
+            #SBATCH --mem=64G
           troika_user: ${{ secrets.HPC_CI_INTEGRATION_USER }}
   benchmark-tests:
     runs-on: hpc

training/src/anemoi/training/diagnostics/callbacks/__init__.py

@@ -182,7 +182,10 @@ def get_callbacks(config: DictConfig) -> list[Callback]:
     trainer_callbacks.extend(instantiate(callback, config) for callback in config.diagnostics.callbacks)
 
     # Plotting callbacks
-    trainer_callbacks.extend(instantiate(callback, config) for callback in config.diagnostics.plot.callbacks)
+    if config["training"]["model_task"] != "anemoi.training.train.tasks.GraphInterpolator":
+        trainer_callbacks.extend(instantiate(callback, config) for callback in config.diagnostics.plot.callbacks)
+    else:
+        LOGGER.info("Plotting callbacks have been temporarily deactivated for the TimeInterpolator")
 
     # Extend with config enabled callbacks
     trainer_callbacks.extend(_get_config_enabled_callbacks(config))

training/src/anemoi/training/diagnostics/callbacks/plot.py

@@ -710,7 +710,13 @@ def _plot(
         else:
             LOGGER.warning("There are no trainable node attributes to plot.")
 
-        if len(edge_trainable_modules := self.get_edge_trainable_modules(model)):
+        from anemoi.models.models import AnemoiModelEncProcDecHierarchical
+
+        if isinstance(model, AnemoiModelEncProcDecHierarchical):
+            LOGGER.warning(
+                "Edge trainable features are not supported for Hierarchical models, skipping plot generation.",
+            )
+        elif len(edge_trainable_modules := self.get_edge_trainable_modules(model)):
             fig = plot_graph_edge_features(model, edge_trainable_modules, q_extreme_limit=self.q_extreme_limit)
 
             self._output_figure(

training/src/anemoi/training/diagnostics/callbacks/plot_ens.py

@@ -286,7 +286,7 @@ def on_validation_batch_end(
             trainer,
             pl_module,
             outputs,
-            batch[0][:, :, 0, :, :],
+            batch[:, :, 0, :, :],
             batch_idx,
         )
 

training/tests/integration/config/test_ensemble_crps.yaml

@@ -14,3 +14,45 @@ dataloader:
     end: 2017-01-08 12:00:00
   validation:
     start: 2017-01-08 18:00:00
+
+diagnostics:
+  plot:
+    callbacks:
+      # Add plot callbacks here.
+      - _target_: anemoi.training.diagnostics.callbacks.plot_ens.PlotEnsSample
+        sample_idx: ${diagnostics.plot.sample_idx}
+        per_sample : 2
+        parameters: ${diagnostics.plot.parameters}
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+        #Defining the accumulation levels for precipitation related fields and the colormap
+        accumulation_levels_plot: [0, 0.05, 0.1, 0.25, 0.5, 1, 1.5, 2, 3, 4, 5, 6, 7, 100] # in mm
+        members: null  # None for all members, list for specific members
+
+      # Deterministic callbacks are also overloaded.
+      - _target_: anemoi.training.diagnostics.callbacks.plot_ens.PlotLoss
+        # group parameters by categories when visualizing contributions to the loss
+        # one-parameter groups are possible to highlight individual parameters
+        parameter_groups:
+          moisture: [tp, cp, tcw]
+          sfc_wind: [10u, 10v]
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+      - _target_: anemoi.training.diagnostics.callbacks.plot_ens.PlotSpectrum
+        sample_idx: ${diagnostics.plot.sample_idx}
+        parameters: ${diagnostics.plot.parameters}
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+      - _target_: anemoi.training.diagnostics.callbacks.plot_ens.PlotHistogram
+        sample_idx: ${diagnostics.plot.sample_idx}
+        parameters: ${diagnostics.plot.parameters}
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+        precip_and_related_fields: ["tp", "cp"] # Optional: specify precip fields for special histogram treatment
+      - _target_: anemoi.training.diagnostics.callbacks.plot_ens.GraphTrainableFeaturesPlot
+        every_n_epochs: ${diagnostics.plot.frequency.epoch}
+
+      # Overloaded PlotSample will return the plots for the first ensemble member
+      # - _target_: anemoi.training.diagnostics.callbacks.plot_ens.PlotSample
+      #   sample_idx: ${diagnostics.plot.sample_idx}
+      #   per_sample: 6
+      #   parameters: ${diagnostics.plot.parameters}
+      #   every_n_batches: ${diagnostics.plot.frequency.batch}
+      #   accumulation_levels_plot: [0, 0.05, 0.1, 0.25, 0.5, 1, 1.5, 2, 3, 4, 5, 6, 7, 100] # in mm
+      #   precip_and_related_fields: ["tp", "cp"] # Optional: specify precip fields

training/tests/integration/config/test_lam.yaml

@@ -50,3 +50,41 @@ data:
     - "cos_local_time"
     - "sin_julian_day"
     - "sin_local_time"
+
+diagnostics:
+  plot:
+    # Parameters to plot
+    parameters:
+    - tp
+    callbacks:
+      # Add plot callbacks here
+      - _target_: anemoi.training.diagnostics.callbacks.plot.GraphTrainableFeaturesPlot
+        every_n_epochs: 5
+      - _target_: anemoi.training.diagnostics.callbacks.plot.PlotLoss
+        # group parameters by categories when visualizing contributions to the loss
+        # one-parameter groups are possible to highlight individual parameters
+        parameter_groups:
+          moisture: [tp]
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+      - _target_: anemoi.training.diagnostics.callbacks.plot.PlotSample
+        sample_idx: ${diagnostics.plot.sample_idx}
+        per_sample : 6
+        parameters: ${diagnostics.plot.parameters}
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+        #Defining the accumulation levels for precipitation related fields and the colormap
+        accumulation_levels_plot: [0, 0.05, 0.1, 0.25, 0.5, 1, 1.5, 2, 3, 4, 5, 6, 7, 100] # in mm
+        precip_and_related_fields: ${diagnostics.plot.precip_and_related_fields}
+        colormaps: ${diagnostics.plot.colormaps}
+      - _target_: anemoi.training.diagnostics.callbacks.plot.PlotSpectrum
+        # every_n_batches: 100 # Override for batch frequency
+        # min_delta: 0.01 # Minimum distance between two consecutive points
+        sample_idx: ${diagnostics.plot.sample_idx}
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+        parameters:
+        - tp
+      - _target_: anemoi.training.diagnostics.callbacks.plot.PlotHistogram
+        sample_idx: ${diagnostics.plot.sample_idx}
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+        precip_and_related_fields: ${diagnostics.plot.precip_and_related_fields}
+        parameters:
+        - tp

training/tests/integration/config/test_stretched.yaml

@@ -60,3 +60,42 @@ training:
   scalers:
     node_weights:
       weight_frac_of_total: 0.25
+
+
+diagnostics:
+  plot:
+    # Parameters to plot
+    parameters:
+    - tp
+    callbacks:
+      # Add plot callbacks here
+      - _target_: anemoi.training.diagnostics.callbacks.plot.GraphTrainableFeaturesPlot
+        every_n_epochs: 5
+      - _target_: anemoi.training.diagnostics.callbacks.plot.PlotLoss
+        # group parameters by categories when visualizing contributions to the loss
+        # one-parameter groups are possible to highlight individual parameters
+        parameter_groups:
+          moisture: [tp]
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+      - _target_: anemoi.training.diagnostics.callbacks.plot.PlotSample
+        sample_idx: ${diagnostics.plot.sample_idx}
+        per_sample : 6
+        parameters: ${diagnostics.plot.parameters}
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+        #Defining the accumulation levels for precipitation related fields and the colormap
+        accumulation_levels_plot: [0, 0.05, 0.1, 0.25, 0.5, 1, 1.5, 2, 3, 4, 5, 6, 7, 100] # in mm
+        precip_and_related_fields: ${diagnostics.plot.precip_and_related_fields}
+        colormaps: ${diagnostics.plot.colormaps}
+      # - _target_: anemoi.training.diagnostics.callbacks.plot.PlotSpectrum
+      #   # every_n_batches: 100 # Override for batch frequency
+      #   # min_delta: 0.01 # Minimum distance between two consecutive points
+      #   sample_idx: ${diagnostics.plot.sample_idx}
+      #   every_n_batches: ${diagnostics.plot.frequency.batch}
+      #   parameters:
+      #   - tp
+      - _target_: anemoi.training.diagnostics.callbacks.plot.PlotHistogram
+        sample_idx: ${diagnostics.plot.sample_idx}
+        every_n_batches: ${diagnostics.plot.frequency.batch}
+        precip_and_related_fields: ${diagnostics.plot.precip_and_related_fields}
+        parameters:
+        - tp

training/tests/integration/config/testing_modifications.yaml

@@ -13,6 +13,7 @@ dataloader:
 diagnostics:
   plot:
     callbacks: []
+    asynchronous: False # Whether to plot asynchronously
   log:
     wandb:
       enabled: False