neural_lam.vis

neural_lam.vis#

Visualization helpers for analysing Neural-LAM predictions and errors.

Module Contents#

neural_lam.vis.plot_error_map(errors, datastore: neural_lam.datastore.base.BaseRegularGridDatastore, title=None)#

Plot a heatmap of per-variable errors across prediction horizons.

Parameters:

  • errors (torch.Tensor) –

    Error values for each horizon and feature.

    • Shape: (pred_steps, d_f)

  • datastore (BaseRegularGridDatastore) – Datastore providing metadata for labels and units.

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

Returns:

Figure handle containing the rendered heatmap.

Return type:

matplotlib.figure.Figure

neural_lam.vis.plot_prediction(datastore: neural_lam.datastore.base.BaseRegularGridDatastore, da_prediction: xarray.DataArray, da_target: xarray.DataArray, title=None, vrange=None)#

Plot a prediction alongside the corresponding ground truth field.

Parameters:

  • datastore (BaseRegularGridDatastore) – Datastore providing coordinate metadata and projection details.

  • da_prediction (xarray.DataArray) –

    Predicted field flattened over the grid.

    • Shape: (N_grid,)

  • da_target (xarray.DataArray) –

    Ground-truth field flattened over the grid.

    • Shape: (N_grid,)

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

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

Returns:

Figure handle containing the two subplots.

Return type:

matplotlib.figure.Figure

neural_lam.vis.plot_spatial_error(error, datastore: neural_lam.datastore.base.BaseRegularGridDatastore, title=None, vrange=None)#

Plot spatial error magnitudes on the datastore grid.

Parameters:

  • error (torch.Tensor) –

    Error magnitudes on the flattened grid.

    • Shape: (N_grid,)

  • datastore (BaseRegularGridDatastore) – 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.

Returns:

Figure handle containing the plotted map.

Return type:

matplotlib.figure.Figure

Contents