# © 2026. Triad National Security, LLC. All rights reserved.
# This program was produced under U.S. Government contract 89233218CNA000001 for Los Alamos
# National Laboratory (LANL), which is operated by Triad National Security, LLC for the U.S.
# Department of Energy/National Nuclear Security Administration. All rights in the program are
# reserved by Triad National Security, LLC, and the U.S. Department of Energy/National Nuclear
# Security Administration. The Government is granted for itself and others acting on its behalf
# a nonexclusive, paid-up, irrevocable worldwide license in this material to reproduce, prepare
# derivative works, distribute copies to the public, perform publicly and display publicly, and
# to permit others to do so.
"""Helper functions for faceting plots by sub-gridcell dimensions."""
from __future__ import annotations
import warnings
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from elm_diagnostics.config.schema import PlotStyleConfig
from elm_diagnostics.io.subgrid import SubgridLevel
# Warn if creating more than this many facets
_MAX_FACETS_NO_WARNING = 16
[docs]
def calculate_facet_layout(n_units: int) -> tuple[int, int]:
"""Calculate optimal (nrows, ncols) layout for n subgrid units.
The layout aims for a roughly square grid, with preference for
wider-than-tall layouts for better use of screen space.
Parameters
----------
n_units : int
Number of sub-gridcell units to plot
Returns
-------
tuple[int, int]
(nrows, ncols) for subplot layout
Examples
--------
>>> calculate_facet_layout(1)
(1, 1)
>>> calculate_facet_layout(3)
(1, 3)
>>> calculate_facet_layout(4)
(2, 2)
>>> calculate_facet_layout(6)
(2, 3)
>>> calculate_facet_layout(9)
(3, 3)
"""
if n_units <= 0:
raise ValueError(f"n_units must be positive, got {n_units}")
if n_units == 1:
return (1, 1)
elif n_units == 2:
return (1, 2)
elif n_units == 3:
return (1, 3)
elif n_units <= 6:
# 4, 5, 6 → 2 rows
return (2, (n_units + 1) // 2)
elif n_units <= 12:
# 7-12 → 3 rows
return (3, (n_units + 2) // 3)
else:
# For large numbers, aim for roughly square
ncols = int(np.ceil(np.sqrt(n_units)))
nrows = int(np.ceil(n_units / ncols))
return (nrows, ncols)
[docs]
def validate_variable_for_subgrid(
da: xr.DataArray,
by: SubgridLevel,
varname: str,
) -> None:
"""Validate that a variable has the requested sub-gridcell dimension.
Parameters
----------
da : xr.DataArray
Variable data array
by : {"column", "pft", "landunit"}
Requested sub-gridcell dimension
varname : str
Variable name (for error message)
Raises
------
ValueError
If the variable does not have the requested dimension, or if
the dimension has size ≤ 1 (making faceting meaningless).
"""
if by not in da.dims:
available = [d for d in da.dims if d in ("column", "pft", "landunit")]
if available:
raise ValueError(
f"Variable '{varname}' does not have dimension '{by}'. "
f"Available sub-gridcell dimensions: {available}. "
f"Use by='{available[0]}' or select a different variable."
)
else:
raise ValueError(
f"Variable '{varname}' does not have dimension '{by}'. "
f"This variable has no sub-gridcell dimensions (column, pft, landunit). "
f"It may be gridcell-averaged. Remove the 'by' parameter or select "
f"a variable with sub-gridcell output."
)
# Check dimension size
size = da.sizes[by]
if size <= 1:
raise ValueError(
f"Variable '{varname}' has dimension '{by}' but size is {size}. "
f"Faceting requires multiple units. Use by=None for single-unit data."
)
[docs]
def get_subgrid_units(da: xr.DataArray, by: SubgridLevel) -> list[int]:
"""Extract list of sub-gridcell unit indices from a DataArray.
Parameters
----------
da : xr.DataArray
Data array with sub-gridcell dimension
by : {"column", "pft", "landunit"}
Sub-gridcell dimension name
Returns
-------
list[int]
Sorted list of unit indices along the specified dimension
"""
coords = da.coords[by].values
return sorted(coords.tolist())