neural_lam.vis ============== .. py:module:: neural_lam.vis .. autoapi-nested-parse:: Visualization helpers for analysing Neural-LAM predictions and errors. Functions --------- .. autoapisummary:: neural_lam.vis.plot_error_heatmap neural_lam.vis.plot_error_map neural_lam.vis.plot_on_axis neural_lam.vis.plot_prediction neural_lam.vis.plot_spatial_error Module Contents --------------- .. py:function:: plot_error_heatmap(errors: torch.Tensor, datastore: neural_lam.datastore.base.BaseDatastore, title: str | None = None, normalization: str = 'state_std') -> matplotlib.figure.Figure Plot a heatmap of errors across variables and prediction horizons. :param errors: Shape ``(pred_steps, num_state_vars)``. Per-step, per-variable errors. These values are used for the numeric annotations in each cell. :type errors: torch.Tensor :param datastore: Datastore providing variable names, units, and step length. :type datastore: BaseDatastore :param title: Title for the axes. :type title: str, optional :param normalization: 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. :type normalization: {"state_std", "diff_std"}, default "state_std" :returns: The completed heatmap figure. :rtype: matplotlib.figure.Figure .. rubric:: 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]". .. py:function:: plot_error_map(errors: torch.Tensor, datastore: neural_lam.datastore.base.BaseDatastore, title: str | None = None) -> matplotlib.figure.Figure Deprecated: use :func:`plot_error_heatmap` instead. :param errors: The error values to plot. :type errors: torch.Tensor :param datastore: The datastore containing grid metadata. :type datastore: BaseDatastore :param title: The title for the plot. :type title: str, optional :returns: The completed heatmap figure. :rtype: matplotlib.figure.Figure .. py:function:: 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 Plot weather state on a projection-aware axis using datastore metadata. :param ax: The axis to plot on. Should have a cartopy projection. :type ax: matplotlib.axes.Axes :param da: The data to plot. Should have shape (num_grid_nodes,). :type da: xarray.DataArray :param datastore: The datastore containing metadata about the grid. :type datastore: BaseDatastore :param vmin: Minimum value for color scale. :type vmin: float, optional :param vmax: Maximum value for color scale. :type vmax: float, optional :param ax_title: Title for the axis. :type ax_title: str, optional :param cmap: Colormap to use for plotting. :type cmap: str or matplotlib.colors.Colormap, optional :param boundary_alpha: If provided, overlay boundary mask with given alpha transparency. :type boundary_alpha: float, optional :param crop_to_interior: If True, crop the plot to the interior region. :type crop_to_interior: bool, optional :returns: The mesh object created by pcolormesh. :rtype: matplotlib.collections.QuadMesh .. py:function:: 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 Plot an example prediction alongside the ground truth. :param datastore: Datastore providing grid metadata and projection. :type datastore: BaseDatastore :param da_prediction: Shape ``(num_grid_nodes,)``. Predicted field values. :type da_prediction: xarray.DataArray :param da_target: Shape ``(num_grid_nodes,)``. Ground-truth field values. :type da_target: xarray.DataArray :param title: Overall figure title. :type title: str, optional :param vrange: ``(vmin, vmax)`` for the shared colour scale. Inferred from data if not given. :type vrange: tuple of (float, float), optional :param boundary_alpha: Alpha transparency for the boundary overlay (default 0.7). :type boundary_alpha: float, optional :param crop_to_interior: If True, crop the axes to the interior region (default True). :type crop_to_interior: bool, optional :param colorbar_label: Label for the shared colorbar. :type colorbar_label: str, optional :returns: The completed two-panel prediction figure. :rtype: matplotlib.figure.Figure .. py:function:: 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 Plot a spatially resolved error map on a projection-aware axis. :param error: Error magnitudes on the flattened grid. * **Shape**: ``(num_grid_nodes,)`` :type error: torch.Tensor :param datastore: Datastore providing coordinate metadata and boundary masks. :type datastore: BaseDatastore :param title: Optional figure title. :type title: str or None, optional :param vrange: Explicit value range ``(vmin, vmax)`` for the color scale. :type vrange: tuple[float, float] or None, optional :param boundary_alpha: Alpha transparency for the boundary overlay (default 0.7). :type boundary_alpha: float, optional :param crop_to_interior: If True, crop the axes to the interior region (default True). :type crop_to_interior: bool, optional :param colorbar_label: Label for the colorbar. :type colorbar_label: str, optional :returns: Figure handle containing the plotted map. :rtype: matplotlib.figure.Figure