Source code for stilt.footprint.accessor
"""The ``.stilt`` accessor on a footprint."""
from __future__ import annotations
import datetime as dt
import json
import warnings
from functools import cached_property
from pathlib import Path
from typing import TYPE_CHECKING
import numpy as np
import pandas as pd
import xarray as xr
from stilt._paths import atomic_path
from stilt.footprint.config import FootprintConfig
from stilt.identity import read_footprint_settings
from stilt.receptors import Receptor
from stilt.spatial import Grid, horizontal_dims
from .aggregation import aggregate
from .io import _naive_utc, _with_cf_metadata
from .targets import Geometry
if TYPE_CHECKING:
from stilt.visualization import FootprintPlotAccessor
def _spacing(values: np.ndarray) -> float | None:
"""Return the cell size of a regular axis, or ``None`` with fewer than two cells."""
values = np.asarray(values, dtype=float)
return float(np.median(np.abs(np.diff(values)))) if values.size > 1 else None
def _check_flux_grid(
flux: xr.DataArray, foot: xr.DataArray, y_dim: str, x_dim: str
) -> None:
"""
Raise when the flux cells are finer than the footprint's; warn when cells of the same size are offset.
Nearest-cell sampling is right for flux cells as large as the
footprint's or larger. A finer flux would give each footprint cell one
small flux cell, and grids of one size offset by a fraction of a cell
give each footprint cell one of the two flux cells it straddles.
"""
if y_dim not in flux.dims or x_dim not in flux.dims:
return # sample_field says what is wrong with the flux's dimensions
for dim in (x_dim, y_dim):
flux_axis, foot_axis = flux[dim].to_numpy(), foot[dim].to_numpy()
flux_size, foot_size = _spacing(flux_axis), _spacing(foot_axis)
if flux_size is None or foot_size is None or foot_size == 0:
continue
if flux_size < 0.99 * foot_size:
raise ValueError(
f"The flux cells ({flux_size:g} in {dim}) are finer than the "
f"footprint's ({foot_size:g}), so each footprint cell would take "
"one small flux cell in place of the mean over its area. Put the "
"flux on the footprint grid first, as an area-weighted mean "
"(xESMF's conservative regridding), or make the footprints on the "
"flux's grid (sim.calc_footprint(grid=...))."
)
if flux_size <= 1.01 * foot_size:
shift = (float(foot_axis[0]) - float(flux_axis[0])) / flux_size
offset = abs(shift - round(shift))
if offset > 0.1:
warnings.warn(
f"The flux cell centers are {offset:.2f} of a cell from the "
f"footprint's in {dim}, so each footprint cell takes one of "
"the two flux cells it straddles. Put the flux on the "
"footprint grid first for an exact sum.",
stacklevel=3,
)
[docs]
@xr.register_dataarray_accessor("stilt")
class FootprintAccessor:
"""
PYSTILT methods on a footprint, as ``foot.stilt``.
A footprint is an :class:`xarray.DataArray`, so sums, selections, and
plots are plain xarray (``foot.sum("time")``,
``foot.sel(time=slice(a, b))``). The accessor adds what needs the
receptor or the footprint settings, which ride in the array's
attributes. Recent xarray keeps attributes through arithmetic. Older
versions drop them (``foot * 2``), so with those, call these methods
on the footprint as read, or use ``xr.set_options(keep_attrs=True)``.
Examples
--------
>>> foot = sim.footprint # doctest: +SKIP
>>> foot.stilt.receptor # doctest: +SKIP
>>> foot.stilt.enhancement(flux).sum() # doctest: +SKIP
>>> foot.stilt.aggregate(counties, time_bins=bins) # doctest: +SKIP
"""
def __init__(self, foot: xr.DataArray) -> None:
self._foot = foot
def _attr(self, key: str) -> str:
"""Return a PYSTILT attribute, or raise when arithmetic has dropped it."""
value = self._foot.attrs.get(key)
if value is None:
raise ValueError(
f"This array has no {key!r} attribute. Older xarray versions "
"drop attributes in arithmetic. Call this on the footprint as "
"read, or use xr.set_options(keep_attrs=True)."
)
return value
# xarray keeps one accessor per array, so each attribute is parsed once.
[docs]
@cached_property
def receptor(self) -> Receptor:
"""The receptor the footprint belongs to."""
return Receptor.from_json(self._attr("stilt_receptor"))
@cached_property
def _settings(self) -> tuple[FootprintConfig, str | None]:
return read_footprint_settings(
json.loads(self._attr("stilt_footprint")), "footprint"
)
@property
def config(self) -> FootprintConfig:
"""The footprint settings (grid, smoothing, and particle transforms)."""
return self._settings[0]
@property
def geometry_hash(self) -> str | None:
"""Hash of the geometry the grid was derived for, or ``None`` when the footprint has no geometry."""
return self._settings[1]
@property
def grid(self) -> Grid:
"""The grid the footprint is on."""
grid = self.config.grid
if grid is None: # calc_footprint() never makes a footprint without one
raise ValueError("The footprint settings have no grid.")
return grid
@property
def name(self) -> str:
"""The footprint's name, usually the variant name."""
return str(self._foot.attrs.get("stilt_name", ""))
@property
def plot(self) -> FootprintPlotAccessor:
"""Plotting methods, such as ``foot.stilt.plot.map()``."""
from stilt.visualization import FootprintPlotAccessor
return FootprintPlotAccessor(self._foot)
[docs]
def to_netcdf(self, path: str | Path) -> Path:
"""
Write the footprint to a CF-1.8 NetCDF file.
:func:`stilt.read_footprint` reads it back with its receptor and
settings.
Parameters
----------
path : str or Path
File to write.
Returns
-------
Path
The path written to.
"""
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
ds = xr.Dataset({"foot": self._foot})
# NetCDF stores naive times; footprint times are UTC.
ds = ds.assign_coords(time=_naive_utc(ds["time"].values))
ds = _with_cf_metadata(ds, grid=self.grid)
ds.attrs["time_created"] = (
dt.datetime.now(dt.UTC).replace(tzinfo=None).isoformat()
)
with atomic_path(path) as tmp:
ds.to_netcdf(tmp, encoding={"foot": {"zlib": True, "complevel": 4}})
return path
[docs]
def enhancement(self, flux: xr.DataArray) -> xr.DataArray:
"""
Return the modeled enhancement at the receptor, footprint times flux summed over the grid.
The flux is taken at each footprint cell center from the nearest
flux cell (:func:`stilt.sampling.sample_field`). That is right for
flux cells as large as the footprint's or larger. A finer flux is
an error, since each footprint cell would take one small flux cell
in place of the mean over its area: put the flux on the footprint
grid first, as an area-weighted mean (xESMF's conservative
regridding), or make the footprints on the flux's grid
(``sim.calc_footprint(grid=...)``). A flux on cells of the
footprint's size whose centers are offset from the footprint's by
more than a tenth of a cell gets a warning, since each footprint
cell then straddles two flux cells.
Parameters
----------
flux : xarray.DataArray
Surface flux on a ``lat``/``lon`` grid (``y``/``x`` for a
projected footprint), in µmol m⁻² s⁻¹ for an enhancement in
ppm. A flux with a ``time`` dimension is read as steps that start
at their labels, as ``aggregate`` and ``project.jacobian`` read
them: each footprint hour takes the last flux step that starts at
or before the hour's start. A monthly flux labeled on the 1st
gives every hour of July 20 July's flux. An hour before the first
label gets zero flux, and an hour after the last label takes the
last step, since the labels do not say where it ends (add a step
of zeros to end it). A flux labeled at the middle of its steps
must be relabeled first: shift its labels back by half a step,
``flux.assign_coords(time=flux.time - half_step)``. The labels
must strictly increase.
Returns
-------
xarray.DataArray
Enhancement for each footprint time step. ``.sum()`` gives the
total. A footprint already summed over time gives the total
alone.
Raises
------
ValueError
If the flux has a ``time`` dimension and the footprint has none,
since a time-summed footprint cannot say which flux step each
hour meets, if the flux cells are finer than the footprint's, or
if the flux's time labels do not strictly increase (a repeated
label is an error).
"""
from stilt.sampling import sample_field
foot = self._foot
y_dim, x_dim = horizontal_dims(foot)
_check_flux_grid(flux, foot, y_dim, x_dim)
if "time" in flux.dims and "time" not in foot.dims:
raise ValueError(
"The flux varies in time and the footprint has no time "
"dimension. Use the footprint before summing it over time."
)
yy, xx = np.meshgrid(
foot[y_dim].to_numpy(), foot[x_dim].to_numpy(), indexing="ij"
)
if "time" in flux.dims:
layers = [
sample_field(
flux,
xx.ravel(),
yy.ravel(),
times=np.full(xx.size, t),
fill_value=0.0,
step_starts=True,
).reshape(yy.shape)
for t in foot["time"].to_numpy()
]
sampled = xr.DataArray(
np.stack(layers),
dims=["time", y_dim, x_dim],
coords={"time": foot["time"]},
)
else:
values = sample_field(flux, xx.ravel(), yy.ravel(), fill_value=0.0)
sampled = xr.DataArray(values.reshape(yy.shape), dims=[y_dim, x_dim])
# By dimension name, so any dimension order gives the same answer.
return (foot * sampled).sum([y_dim, x_dim]).rename("enhancement")
[docs]
def aggregate(
self,
target: Geometry,
time_bins: pd.IntervalIndex,
) -> pd.DataFrame:
"""
Sum the footprint onto the cells of another grid or set of polygons, per time bin.
A footprint value belongs to its whole cell, so it is summed when
cells are combined. Each footprint cell is split among the target
cells it overlaps, in proportion to the overlapping area. For a
target whose cells are whole blocks of footprint cells, this is a
plain block sum. Footprint values outside the target are dropped.
Time steps are summed within each of ``time_bins``.
Parameters
----------
target : Grid, Mesh, or Zones
Cells to sum onto:
- :class:`~stilt.Grid`, every cell of a regular grid, in
``Grid.index`` order.
- :class:`~stilt.Mesh`, any polygons (a shapefile, H3 hexagons,
nested grids). Rows are indexed by cell id.
- :class:`~stilt.Zones`, groups of the cells of a grid or mesh.
Use it to keep only some cells of a grid.
A target in another coordinate system is projected onto the
footprint grid.
time_bins : pandas.IntervalIndex
Time intervals to sum over, such as the time steps of a flux
inventory. They must be closed on the left
(``closed="left"``): each bin holds the footprint hours that
start in it.
Returns
-------
pandas.DataFrame
One row per target cell and one column per time bin, labeled
by the bin's left edge. Rows are indexed by ``(x, y)`` for grid
targets and by cell id for meshes and zones. Cells and bins the
footprint does not reach are 0.
Raises
------
TypeError
If ``time_bins`` is not a :class:`pandas.IntervalIndex`.
ValueError
If ``time_bins`` is not closed on the left.
"""
return aggregate(self._foot, target, time_bins)