feat: Reconstruction plots callback

[icedoom888]

Description

Introducing focus area to plot only specific regions.
Introducing reconstruction plots to compare input and output and their difference.

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--782.org.readthedocs.build/en/782/

---

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

---

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

[anaprietonem]

some of this seems to be could be better abstracted and integrated into training/src/anemoi/training/diagnostics/maps.py

[icedoom888]

Not sure which parts you refer too here?

[anaprietonem]

I replied on the other PR, let me know if it's unclear

[anaprietonem]

Is it fair to say this is a special case for the plot_predicted_multilevel_flat_sample where we don't have a target but rather the input and the pred? so it would like the plotting a diagnostic variable but where target is input?

Thinking if would dramatically reduce the amount of code if we could generalise the plot_predicted_multilevel_flat_sample to work with that case

[icedoom888]

True, but the function was already quite complex and had many different columns.
As it is a different layout, I preferred to just create a new one rather than extending this...

[anaprietonem]

The issue that I see at the moment is that projection and transform are specific and will just work if one has cartopy installed. So I would like to find a way we can define this and reuse it for the other plots so that we reuse the functionality we currently have but with the option of customising the projections and the transform.

[anaprietonem]

If we don't want to make it completely configurable from the config I think a clearer design would be something like

@dataclass(frozen=True)
class MapContext:
    """Encapsulates map projection, data CRS, and extent handling."""
    projection: object | None
    data_crs: object | None
    pad: float = 0.1

    @classmethod
    def from_latlons(
        cls,
        latlons: np.ndarray,
        *,
        datashader: bool = False,
        pad: float = 0.1,
    ) -> MapContext:
        try:
            import cartopy.crs as ccrs
        except ModuleNotFoundError:
            return cls(projection=None, data_crs=None, pad=pad)

        if datashader:
            return cls(projection=None, data_crs=None, pad=pad)

        projection = lambert_conformal_from_latlon_points(latlons)
        data_crs = ccrs.PlateCarree()

        return cls(projection=projection, data_crs=data_crs, pad=pad)

    def set_extent(
        self,
        ax,
        lon: np.ndarray,
        lat: np.ndarray,
    ) -> None:
        """Set axis extent consistently for projected or flat axes."""
        if self.data_crs is not None and hasattr(ax, "set_extent"):
            ax.set_extent(
                [
                    lon.min() - self.pad,
                    lon.max() + self.pad,
                    lat.min() - self.pad,
                    lat.max() + self.pad,
                ],
                crs=self.data_crs,
            )
        else:
            xmin, xmax = max(lon.min(), -np.pi), min(lon.max(), np.pi)
            ymin, ymax = max(lat.min(), -np.pi / 2), min(lat.max(), np.pi / 2)

            ax.set_xlim(xmin - self.pad, xmax + self.pad)
            ax.set_ylim(ymin - self.pad, ymax + self.pad)

and we then could pass this to the plots like

from map_context import MapContext

map_ctx = MapContext.from_latlons(
    latlons,
    datashader=datashader,
)

fig, ax = plt.subplots(
    n_vars,
    n_plots_per_sample,
    figsize=(n_plots_per_sample * 4, n_vars * 3),
    layout="constrained",
    subplot_kw={"projection": map_ctx.projection},
)

so the MapContext will handle all the logic and code related to cartopy

[anaprietonem]

Also some of those settings could be defined at the level of callback ie if you have a specific AutoencoderPlotSample then there we can define those and try to reuse the plot_predicted_multilevel_flat_sample

[anaprietonem]

On the topic of creating a new function vs using the existing one, from what I see it's mostly the definition of the projection and the fact that for the autoencoder you just one 3 subplots, vs the forecasting one includes 6 subplots right? So I think this could be configured to define this at the callback level and then reuse the functions that we have.

I understand having the functions separated might make it easier to understand but bugs and features in the future would then need to be copied across both of them, making dev workflow more cumbersome so I would prefer we homogenise it.

[anaprietonem]

These needs to be repeated for all callbacks - we should have a function to clip the data as part of the basePlot and then reuse it

[icedoom888]

What do you mean? various callbacks can have different focus_areas, you won't know until you reach every callback plot, no?

[anaprietonem]

Hey Alberto, I had a look I think the focus_area parameter is a nice feature!

I left some comments for clarification and mostly I think the reconstruction plot callback and how the LAEA projection is handled can be abstracted a bit more so the amount of code changes can be reduced. Let me know if anything is not clear

[icedoom888]

Hey @anaprietonem thanks for reviewing, addressed majority of the requested changes, could you clarify some things in the above comments?

[JPXKQX]

Hi Alberto, first of all, thanks for the PR. I think this focus_area functionality can be very useful for many users. Below are a few thoughts and suggestions.

• It seems that this PR contains two distinct contributions: a) the introduction of the focus_area argument, b) the reconstruction maps functionality. As the latter depends on the autoencoder PR being merged first, it might be worth splitting this into two separate PRs, so that the focus_area feature can be merged and made available as soon as possible.

• The need for a focus_area has already been raised by users, so I agree it would be a very valuable. Regarding implementation, one possible improvement would be to move this functionality into a dedicated class in a new file. In this setup, the focus_area could be created in the init method (possibly via a factory method) and then applied in forward(). This could help keep the logic more modular and easier to test.

• A few open questions around naming, where consistency with existing code might help:

  • focus_area. In anemoi-graphs, we usually refer to an area/region of interest. It might be worth considering whether to keep a similar name for coherence (e.g. area_of_interest) or to deliberately use a different term to distinguish concepts. Not sure what is preferred.
  • spatial_mask naming. In the existing configs, we typically refer to this kind of object as node_attribute_name or mask_attr_name
  • latlon_bounds. We already have a similar argument in graphs called area. While that name may be a bit generic, common alternatives in other libraries include: extent (used in matplotlib), bbox (bounding box), clip, ... Something more explicit like latlon_bbox, which specifies that it has to be lat/lon, could also work well.

• In anemoi-graphs and anemoid-dataset, the area (latlon_bounds) is usually stored as a flat tuple of four values rather than a tuple of tuples. This also seems consistent with many other python packages

These comments are mainly intended to raise discussion rather than to suggest strict changes. Happy to hear your thoughts and suggestions. Thanks again for your contribution!

[icedoom888]

Hey @JPXKQX , thanks for taking a look¦

• Yes, i can split this into two separate PRs.
• Not sure what you mean by a separate class and factory. I thought of this as simply another plot argument.
• I would go for: focus_area to separate, mask_attr_name to keep consistency with graphs and latlon_bbox expressed as a flat tuple.

Will ask your review in the spawened PRs ;)