neural_lam.vis

neural_lam.vis#

Visualization helpers for analysing Neural-LAM predictions and errors.

Functions#

plot_error_heatmap(→ matplotlib.figure.Figure)

Plot a heatmap of errors across variables and prediction horizons.

plot_error_map(→ matplotlib.figure.Figure)

Deprecated: use plot_error_heatmap() instead.

plot_on_axis(→ matplotlib.collections.QuadMesh)

Plot weather state on a projection-aware axis using datastore metadata.

plot_prediction(→ matplotlib.figure.Figure)

Plot an example prediction alongside the ground truth.

plot_spatial_error(→ matplotlib.figure.Figure)

Plot a spatially resolved error map on a projection-aware axis.

Module Contents#

neural_lam.vis.plot_error_heatmap(errors: torch.Tensor, datastore: neural_lam.datastore.base.BaseDatastore, title: str | None = None, normalization: str = 'state_std') → matplotlib.figure.Figure[source]#

Plot a heatmap of errors across variables and prediction horizons.

Parameters:

  • errors (torch.Tensor) – Shape (pred_steps, num_state_vars). Per-step, per-variable errors. These values are used for the numeric annotations in each cell.

  • datastore (BaseDatastore) – Datastore providing variable names, units, and step length.

  • title (str, optional) – Title for the axes.

  • normalization ({"state_std", "diff_std"}, default "state_std") – Color scaling mode. “state_std” divides by climatological std; “diff_std” divides by the typical one-step change magnitude. Both fall back to per-variable max error when the required stats are missing.

Returns:

The completed heatmap figure.

Return type:

matplotlib.figure.Figure

Notes

Color scaling is controlled by normalization; see _get_heatmap_color_values for the full fallback logic. When stats are unavailable the colorbar label includes “[fallback]”.

neural_lam.vis.plot_error_map(errors: torch.Tensor, datastore: neural_lam.datastore.base.BaseDatastore, title: str | None = None) → matplotlib.figure.Figure[source]#

Deprecated: use plot_error_heatmap() instead.

Parameters:

  • errors (torch.Tensor) – The error values to plot.

  • datastore (BaseDatastore) – The datastore containing grid metadata.

  • title (str, optional) – The title for the plot.

Returns:

The completed heatmap figure.

Return type:

matplotlib.figure.Figure

neural_lam.vis.plot_on_axis(ax: cartopy.mpl.geoaxes.GeoAxes, da: xarray.DataArray, datastore: neural_lam.datastore.base.BaseDatastore, vmin: float | None = None, vmax: float | None = None, ax_title: str | None = None, cmap: str | matplotlib.colors.Colormap = 'plasma', boundary_alpha: float | None = None, crop_to_interior: bool = False) → matplotlib.collections.QuadMesh[source]#

Plot weather state on a projection-aware axis using datastore metadata.

Parameters:

  • ax (matplotlib.axes.Axes) – The axis to plot on. Should have a cartopy projection.

  • da (xarray.DataArray) – The data to plot. Should have shape (num_grid_nodes,).

  • datastore (BaseDatastore) – The datastore containing metadata about the grid.

  • vmin (float, optional) – Minimum value for color scale.

  • vmax (float, optional) – Maximum value for color scale.

  • ax_title (str, optional) – Title for the axis.

  • cmap (str or matplotlib.colors.Colormap, optional) – Colormap to use for plotting.

  • boundary_alpha (float, optional) – If provided, overlay boundary mask with given alpha transparency.

  • crop_to_interior (bool, optional) – If True, crop the plot to the interior region.

Returns:

The mesh object created by pcolormesh.

Return type:

matplotlib.collections.QuadMesh

neural_lam.vis.plot_prediction(datastore: neural_lam.datastore.base.BaseDatastore, da_prediction: xarray.DataArray, da_target: xarray.DataArray, title: str | None = None, vrange: tuple[float, float] | None = None, boundary_alpha: float = 0.7, crop_to_interior: bool = True, colorbar_label: str = '') → matplotlib.figure.Figure[source]#

Plot an example prediction alongside the ground truth.

Parameters:

  • datastore (BaseDatastore) – Datastore providing grid metadata and projection.

  • da_prediction (xarray.DataArray) – Shape (num_grid_nodes,). Predicted field values.

  • da_target (xarray.DataArray) – Shape (num_grid_nodes,). Ground-truth field values.

  • title (str, optional) – Overall figure title.

  • vrange (tuple of (float, float), optional) – (vmin, vmax) for the shared colour scale. Inferred from data if not given.

  • boundary_alpha (float, optional) – Alpha transparency for the boundary overlay (default 0.7).

  • crop_to_interior (bool, optional) – If True, crop the axes to the interior region (default True).

  • colorbar_label (str, optional) – Label for the shared colorbar.

Returns:

The completed two-panel prediction figure.

Return type:

matplotlib.figure.Figure

neural_lam.vis.plot_spatial_error(error: torch.Tensor, datastore: neural_lam.datastore.base.BaseDatastore, title: str | None = None, vrange: tuple[float, float] | None = None, boundary_alpha: float = 0.7, crop_to_interior: bool = True, colorbar_label: str = '') → matplotlib.figure.Figure[source]#

Plot a spatially resolved error map on a projection-aware axis.

Parameters:

  • error (torch.Tensor) – Error magnitudes on the flattened grid. * Shape: (num_grid_nodes,)

  • datastore (BaseDatastore) – Datastore providing coordinate metadata and boundary masks.

  • title (str or None, optional) – Optional figure title.

  • vrange (tuple[float, float] or None, optional) – Explicit value range (vmin, vmax) for the color scale.

  • boundary_alpha (float, optional) – Alpha transparency for the boundary overlay (default 0.7).

  • crop_to_interior (bool, optional) – If True, crop the axes to the interior region (default True).

  • colorbar_label (str, optional) – Label for the colorbar.

Returns:

Figure handle containing the plotted map.

Return type:

matplotlib.figure.Figure