"""
Calculating a footprint from particles, as STILT-R's ``calc_footprint`` does.
Footprints match STILT-R at ``rtol=1e-7`` per cell, except on a projected
grid, where the smoothing (``_compute_kernel_bandwidths``) and the test for
the early-time interpolation (``_interpolate_early_timesteps``) differ on
purpose, and in the smoothing of a cloud across 180 degrees on a grid that
goes all the way around (``_unwrap_around_cloud``,
``_fold_padding_around_the_globe``). Run the ``fidelity`` tests before
merging any change to this module.
"""
from __future__ import annotations
from collections.abc import Sequence
from pathlib import Path
from typing import Any, cast
import numpy as np
import pandas as pd
import xarray as xr
from stilt.exceptions import EmptyFootprint
from stilt.footprint.config import FootprintConfig
from stilt.particles import FOOTPRINT_COLUMNS, check_particles
from stilt.receptors import Receptor
from stilt.spatial import Grid, _grid_cell_starts, is_longlat
from stilt.transforms import apply_transforms
from .io import _footprint_array
def _make_gauss_kernel(rs: tuple[float, float], sigma: float) -> np.ndarray:
"""2D Gaussian smoothing kernel on a physical-unit grid, normalized to sum=1."""
if sigma == 0:
return np.array([[1.0]])
d = 3 * sigma
nx = 1 + 2 * int(np.floor(d / rs[0]))
ny = 1 + 2 * int(np.floor(d / rs[1]))
x = (np.arange(nx) - nx // 2) * rs[0]
y = (np.arange(ny) - ny // 2) * rs[1]
xx, yy = np.meshgrid(x, y, indexing="ij")
m = np.exp(-(xx**2 + yy**2) / (2 * sigma**2))
w = m / m.sum()
return np.where(np.isnan(w), 1.0, w)
#: The particle columns a footprint is made from. STILT-R's ``simulation_step``
#: passes ``calc_footprint`` only these (``footprint_varsiwant``), so the rows the
#: early-time interpolation inserts have no missing values and are kept.
_FOOTPRINT_INPUTS = ("particle", "age", "lon", "lat", "foot")
def _r_seq(start: float, stop: float, by: float) -> np.ndarray:
"""Return R's ``seq(start, stop, by = by)``: ``start + (0:n) * by``, capped at ``stop``."""
n = int((stop - start) / by + 1e-10)
return np.minimum(start + np.arange(n + 1) * by, stop)
def _interpolation_times(time_sign: int) -> np.ndarray:
"""
Return STILT-R's interpolation times for the first 100 minutes, in minutes.
They are built as R builds them, so they match STILT-R's to the last bit:
R's ``seq(0, 10, by = 0.1)`` holds ``2.9000000000000004``, not ``2.9``.
A position interpolated at a time one bit off can fall in the next grid
cell, which moves that step's ``foot`` with it.
"""
times = np.concatenate(
[_r_seq(0, 10, 0.1), _r_seq(10.2, 20, 0.2), _r_seq(20.5, 100, 0.5)]
)
return times * time_sign
def _wrap_longitudes(lon: Any, *, west: Any = 0.0) -> Any:
"""Return *lon* shifted by whole turns into ``[west, west + 360)``; both are numbers, arrays, or Series."""
return (((lon - west) % 360) + 360) % 360 + west
def _covers_the_globe(xmin: float, xmax: float) -> bool:
"""Return whether longitudes *xmin* to *xmax* span a whole number of turns (STILT-R's ``xdist == 0``)."""
return ((180 - xmin) - (-180 - xmax)) % 360 == 0
def _wrap_antimeridian_longitudes(
p: pd.DataFrame, *, xmin: float, xmax: float
) -> tuple[pd.DataFrame, float, float, bool]:
"""
Shift longitudes to 0 to 360 when the grid crosses the antimeridian.
Only for longitude/latitude grids. A grid that covers the globe is
anchored to -180 to 180 and the longitudes are left as they are. Returns
``(p, xmin, xmax, wrapped)``.
"""
if _covers_the_globe(xmin, xmax):
return p, -180.0, 180.0, False
if (xmax < xmin) or (xmax > 180):
p = p.copy()
p["lon"] = _wrap_longitudes(p["lon"])
xmin = _wrap_longitudes(xmin)
xmax = _wrap_longitudes(xmax)
return p, xmin, xmax, True
return p, xmin, xmax, False
def _interpolate_early_timesteps(
p: pd.DataFrame,
*,
xres: float,
yres: float,
time_sign: int,
crs: str = "+proj=longlat",
) -> pd.DataFrame:
"""
Add interpolated positions in the first 100 minutes when particles jump grid cells.
Near the receptor, particles can move more than a grid cell per output
step. When the median step does, along either axis of the grid,
positions and ``foot`` are linearly interpolated onto finer times (every
0.1 min to 10 min, 0.2 min to 20 min, and 0.5 min to 100 min). ``foot``
is then rescaled so each of those three windows keeps its total, as
STILT-R does.
*p* holds longitudes and latitudes in degrees. On a projected *crs* the
steps are measured on the grid, from the particles' positions projected
to it, so they compare with *xres* and *yres* in its units. STILT-R
compares the steps in degrees with the cell size in the grid's units,
which on a grid in meters never interpolates. The interpolation itself
is done in degrees on any grid.
"""
early = cast(pd.DataFrame, p[np.abs(p["age"]) < 100])
if early.empty:
return p
# Match R: per-particle median(abs(diff(lon/lat))), then median across particles.
# should_interpolate = median(dx) > xres OR median(dy) > yres
sorted_early = early.sort_values(by=["particle", "age"])
if not is_longlat(crs):
# The steps along the grid's axes, in its units.
from pyproj import Transformer
tr = Transformer.from_crs("EPSG:4326", crs, always_xy=True)
x, y = tr.transform(
sorted_early["lon"].to_numpy(), sorted_early["lat"].to_numpy()
)
sorted_early = sorted_early.assign(lon=x, lat=y)
diffs = cast(
pd.DataFrame,
sorted_early.groupby("particle", sort=False)[["lon", "lat"]].diff().abs(),
)
per_particle_med = cast(
pd.DataFrame, diffs.groupby(sorted_early["particle"]).median()
)
dx_values = per_particle_med["lon"].dropna().to_numpy()
dy_values = per_particle_med["lat"].dropna().to_numpy()
dx_med = float(np.median(dx_values)) if dx_values.size else np.nan
dy_med = float(np.median(dy_values)) if dy_values.size else np.nan
if (np.isnan(dx_med) or dx_med <= xres) and (np.isnan(dy_med) or dy_med <= yres):
return p
t_new = _interpolation_times(time_sign)
# Store pre-interpolation foot sums per window for rescaling later.
atime = np.abs(p["age"])
foot_sums = [
p.loc[atime <= 10, "foot"].sum(),
p.loc[(atime > 10) & (atime <= 20), "foot"].sum(),
p.loc[(atime > 20) & (atime <= 100), "foot"].sum(),
]
p = _interpolate_particle_tracks(p, t_new=t_new)
p["age"] = p["age"].round(1)
# Rescale foot so total influence per window matches original.
atime = np.abs(p["age"])
masks = [
atime <= 10,
(atime > 10) & (atime <= 20),
(atime > 20) & (atime <= 100),
]
for mask, total in zip(masks, foot_sums, strict=False):
s = p.loc[mask, "foot"].sum()
if s > 0:
p.loc[mask, "foot"] *= total / s
return p
def _interpolate_particle_tracks(p: pd.DataFrame, *, t_new: np.ndarray) -> pd.DataFrame:
"""Interpolate lon/lat/foot for each particle on a dense time grid."""
frames: list[pd.DataFrame] = []
original_columns = list(p.columns)
for number, group in p.groupby("particle", sort=False):
source_time = group["age"].to_numpy(dtype=float)
target_time = np.unique(np.concatenate([source_time, t_new]))
order = np.argsort(source_time)
sorted_time = source_time[order]
unique_time, unique_idx = np.unique(sorted_time, return_index=True)
# Mirror STILT-R's full_join(expand.grid(...), by = c("particle", "time")):
# original rows keep every column, while inserted rows only receive the
# interpolated lon/lat/foot values below. A later dropna() therefore
# matches R's na.omit() across all columns.
expanded = pd.DataFrame(
{
"particle": np.full(len(target_time), number),
"age": target_time,
}
)
join_group = group.copy(deep=False)
join_group["age"] = join_group["age"].astype(float)
frame = expanded.merge(
join_group, on=["particle", "age"], how="left", sort=False
)
for col in ["lon", "lat", "foot"]:
values = group[col].to_numpy(dtype=float)[order][unique_idx]
if len(unique_time) >= 2:
# left/right=nan matches R's na_interp: no extrapolation
# beyond the observed [min_time, max_time] range.
frame[col] = np.interp(
target_time, unique_time, values, left=np.nan, right=np.nan
)
elif len(unique_time) == 1:
frame[col] = np.full(len(target_time), values[0])
else:
frame[col] = np.full(len(target_time), np.nan)
frames.append(frame.loc[:, original_columns])
if not frames:
return p.iloc[0:0].copy()
return (
pd.concat(frames, ignore_index=True)
.dropna()
.sort_values(["particle", "age"], ascending=[True, False], kind="stable")
.reset_index(drop=True)
)
def _project_particles_to_crs(
p: pd.DataFrame,
*,
crs: str,
xmin: float,
xmax: float,
ymin: float,
ymax: float,
) -> tuple[pd.DataFrame, float, float, float, float]:
"""
Project particle positions and the grid bounds to the grid's CRS.
Grid bounds are given in degrees and projected along with the
particles.
"""
from pyproj import Transformer
tr = Transformer.from_crs("EPSG:4326", crs, always_xy=True)
p = p.copy()
p["lon"], p["lat"] = tr.transform(p["lon"].values, p["lat"].values)
corners_x, corners_y = tr.transform([xmin, xmax], [ymin, ymax])
return (
p,
float(np.min(corners_x)),
float(np.max(corners_x)),
float(np.min(corners_y)),
float(np.max(corners_y)),
)
def _compute_kernel_bandwidths(
p: pd.DataFrame, *, smooth_factor: float, crs: str = "+proj=longlat"
) -> tuple[pd.DataFrame, np.ndarray]:
"""
Return ``(kernel_df, w)``, the Gaussian kernel width for each ``rtime``, in *crs* units.
*p* holds the particles' longitudes and latitudes in degrees, before any
projection. The width grows with the spread of the particles and the
time since release, as in STILT-R, and is divided by ``cos(lat)`` so it
covers the same distance at high latitudes::
w = smooth_factor * 0.06 * varsum**0.25 * (|rtime| / 1440)**0.5 / cos(lat)
This is in degrees of longitude. On a projected grid it is converted to
the grid's units by the length of one degree of longitude at the
particles' mean position (``_grid_units_per_degree_longitude``), so the
kernel is as wide from east to west, on the ground, as on a
longitude/latitude grid. STILT-R applies the formula to the projected
coordinates instead, which on a grid in meters gives a kernel a few
meters wide.
On any grid, each ``rtime``'s longitudes are first shifted by whole
turns to within 180 degrees of that step's circular mean
(``_unwrap_around_cloud``), so a cloud has the spread and mean position
it has on the ground wherever on the globe it is, across 180 degrees
included, and wherever it was released. STILT-R measures the spread of
the longitudes as written, so on a global grid (where it does not shift
them) a cloud across 180 degrees gets a kernel many times too wide.
"""
p = p.assign(lon=_unwrap_around_cloud(p))
kernel_df = (
p.groupby("rtime")
.agg(
long_var=("lon", "var"),
lati_var=("lat", "var"),
lat_mean=("lat", "mean"),
lon_mean=("lon", "mean"),
)
.reset_index()
.dropna()
)
if kernel_df.empty:
# Single-particle / no-variance case: zero-width kernel per rtime.
rtime_vals = np.sort(np.asarray(p["rtime"].dropna().unique()))
kernel_df = pd.DataFrame(
{
"rtime": rtime_vals,
"lat_mean": [float(p["lat"].to_numpy().mean())] * len(rtime_vals),
}
)
return kernel_df, np.zeros(len(kernel_df), dtype=float)
kernel_df["varsum"] = kernel_df["long_var"] + kernel_df["lati_var"]
di = kernel_df["varsum"].to_numpy() ** 0.25
ti = (np.abs(kernel_df["rtime"].to_numpy()) / 1440) ** 0.5
grid_conv = np.cos(kernel_df["lat_mean"].to_numpy() * np.pi / 180)
w = smooth_factor * 0.06 * di * ti / grid_conv
if not is_longlat(crs):
# The mean of unwrapped longitudes can pass 180; only those go back.
lon_mean = kernel_df["lon_mean"].to_numpy()
outside = (lon_mean < -180.0) | (lon_mean >= 180.0)
lon_mean = np.where(outside, _wrap_longitudes(lon_mean, west=-180.0), lon_mean)
w = w * _grid_units_per_degree_longitude(
crs, lon=lon_mean, lat=kernel_df["lat_mean"].to_numpy()
)
# A step whose cloud is centered where the projection has no value
# (UTM near the equator 81 to 99 degrees from its meridian, a
# geostationary grid off the disk) has no width. Its particles cannot
# reach the grid, so the step is dropped, as STILT-R's na.omit drops a
# step whose particles project to infinity.
# If no step has a width, the particles are not smoothed, as when no
# step has a spread (above).
finite = np.isfinite(w)
if not finite.any():
w = np.zeros(len(w), dtype=float)
elif not finite.all():
kernel_df, w = kernel_df[finite].reset_index(drop=True), w[finite]
return kernel_df, w
def _unwrap_around_cloud(p: pd.DataFrame) -> pd.Series:
"""
Return the longitudes of *p*, each ``rtime``'s within 180 degrees of that step's circular mean.
The circular mean is the direction of the mean of the unit vectors, so
it is the cloud's center wherever the cloud is, and the window it opens
has its seam on the far side of the globe from the cloud. The center is
expressed in the frame the longitudes are written in (-180 to 180, or
0 to 360 on a grid that crosses 180 degrees), so a longitude is shifted,
by whole turns (``_wrap_longitudes``), only when it lies more than 180
degrees from the center. The others are returned as they are, to the
last bit.
"""
lon = p["lon"]
rad = np.radians(lon.to_numpy(dtype=float))
steps = p["rtime"]
mean_sin = pd.Series(np.sin(rad), index=p.index).groupby(steps).transform("mean")
mean_cos = pd.Series(np.cos(rad), index=p.index).groupby(steps).transform("mean")
circular = np.degrees(np.arctan2(mean_sin, mean_cos))
# The same center, a whole number of turns away, nearest the plain mean.
plain = lon.groupby(steps).transform("mean")
center = plain + ((circular - plain + 180.0) % 360.0 - 180.0)
west = center - 180.0
outside = (lon < west) | (lon >= west + 360.0)
return lon.where(~outside, _wrap_longitudes(lon, west=west))
def _grid_units_per_degree_longitude(
crs: str, *, lon: np.ndarray, lat: np.ndarray
) -> np.ndarray:
"""
Return the length of one degree of longitude at each ``(lon, lat)``, in *crs* units.
It is measured with the projection itself: the distance between the
points a hundredth of a degree east and west, divided by the
difference in degrees. So it includes the projection's scale at that
point and its units (meters, feet, kilometers), and does not depend on
how the projection's axes are turned. It is NaN where the projection
has no value. Next to a projection's cut (a
Lambert conformal or Mercator grid is cut at ``lon_0 + 180``) one of
the two points is across it, so the distances from the point to each
side differ by more than a factor of two; there the shorter one is
used.
"""
from pyproj import Transformer
tr = Transformer.from_crs("EPSG:4326", crs, always_xy=True)
step = 0.01
x, y = (np.asarray(v) for v in tr.transform(lon, lat))
x_west, y_west = (np.asarray(v) for v in tr.transform(lon - step, lat))
x_east, y_east = (np.asarray(v) for v in tr.transform(lon + step, lat))
# Where the projection has no value it gives inf, and the length is NaN.
with np.errstate(invalid="ignore"):
across = np.hypot(x_east - x_west, y_east - y_west) / (2 * step)
west = np.hypot(x - x_west, y - y_west) / step
east = np.hypot(x_east - x, y_east - y) / step
cut = np.maximum(west, east) > 2 * np.minimum(west, east)
return np.where(cut, np.minimum(west, east), across)
def _padded_axes(
*,
xmin: float,
ymin: float,
xres: float,
yres: float,
n_lon: int,
n_lat: int,
xbuf: int,
ybuf: int,
) -> tuple[np.ndarray, np.ndarray]:
"""
Return the lower cell edges of the output grid padded by *xbuf* and *ybuf* cells per side.
The padding, the size of the largest kernel, lets particles just outside
the grid be smoothed into it.
"""
glong_buf = xmin - xbuf * xres + np.arange(n_lon + 2 * xbuf) * xres
glati_buf = ymin - ybuf * yres + np.arange(n_lat + 2 * ybuf) * yres
return glong_buf, glati_buf
def _filter_and_rasterize_particles(
p: pd.DataFrame,
*,
glong_buf: np.ndarray,
glati_buf: np.ndarray,
xbufh: int,
ybufh: int,
xmin: float,
xmax: float,
ymin: float,
ymax: float,
xres: float,
yres: float,
time_integrate: bool,
) -> tuple[pd.DataFrame, np.ndarray]:
"""
Sum ``foot`` by padded grid cell and time step for particles on the padded grid.
Returns ``(p, layers)``. ``p`` has columns ``loi, lai, time, rtime,
foot, layer``, and ``layers`` holds the sorted hour indices, or ``[0]``
when ``time_integrate`` is set. Both are empty when no particle is on
the padded grid.
"""
filtered = cast(
pd.DataFrame,
p[
(p["foot"] > 0)
& (p["lon"] >= xmin - xbufh * xres)
& (p["lon"] < xmax + xbufh * xres)
& (p["lat"] >= ymin - ybufh * yres)
& (p["lat"] < ymax + ybufh * yres)
].copy(),
)
if filtered.empty:
return filtered, np.zeros(0, dtype=int)
filtered["loi"] = (
np.searchsorted(
glong_buf,
filtered["lon"].to_numpy(dtype=float),
side="right",
)
- 1
)
filtered["lai"] = (
np.searchsorted(
glati_buf,
filtered["lat"].to_numpy(dtype=float),
side="right",
)
- 1
)
# Sum foot for particles in the same cell at the same time step.
p = cast(
pd.DataFrame,
filtered.groupby(["loi", "lai", "age", "rtime"], as_index=False)["foot"].sum(),
)
# time_integrate=True collapses all steps into a single layer; otherwise
# bin into hourly layers for time-resolved output.
p["layer"] = 0 if time_integrate else np.floor(p["age"] / 60).astype(int)
layers = np.sort(np.asarray(p["layer"].unique(), dtype=int))
return p, layers
# Most (cell, kernel cell) pairs ``_stamp_kernel`` holds at once, about 16 to
# 50 bytes each. It keeps memory bounded when late-time kernels are wide.
_MAX_STAMP_PAIRS = 2**20
def _stamp_kernel(
out: np.ndarray,
rows: np.ndarray,
cols: np.ndarray,
values: np.ndarray,
kernel: np.ndarray,
) -> None:
"""
Add ``values[j] * kernel``, centered on cell ``(rows[j], cols[j])``, to *out* in place.
This is STILT-R's ``permute.f90``: the kernel is added at each occupied
cell, so the cost is cells x kernel cells. Kernel mass that falls outside
*out* is dropped, as a convolution with zeros beyond the edge drops it.
Cells are stamped in chunks of at most ``_MAX_STAMP_PAIRS`` (cell, kernel
cell) pairs, so memory stays bounded however wide the kernel is.
"""
n_rows, n_cols = out.shape
hx, hy = kernel.shape[0] // 2, kernel.shape[1] // 2
dx, dy = np.meshgrid(np.arange(-hx, hx + 1), np.arange(-hy, hy + 1), indexing="ij")
dx, dy, weights = dx.ravel(), dy.ravel(), kernel.ravel()
chunk = max(1, _MAX_STAMP_PAIRS // weights.size)
for start in range(0, len(values), chunk):
r = rows[start : start + chunk]
c = cols[start : start + chunk]
v = values[start : start + chunk]
# The box these kernels cover, cut to the edges of out.
r0, r1 = r.min() - hx, r.max() + hx + 1
c0, c1 = c.min() - hy, c.max() + hy + 1
crosses_edge = r0 < 0 or c0 < 0 or r1 > n_rows or c1 > n_cols
r0, r1 = max(r0, 0), min(r1, n_rows)
c0, c1 = max(c0, 0), min(c1, n_cols)
width = c1 - c0
# Flat index in the box of each (cell, kernel cell) pair, and its mass.
idx = ((r - r0) * width + (c - c0))[:, None] + (dx * width + dy)
mass = v[:, None] * weights
if crosses_edge:
rr = r[:, None] + dx
cc = c[:, None] + dy
inside = (rr >= 0) & (rr < n_rows) & (cc >= 0) & (cc < n_cols)
idx, mass = idx[inside], mass[inside]
out[r0:r1, c0:c1] += np.bincount(
idx.ravel(), weights=mass.ravel(), minlength=(r1 - r0) * width
).reshape(r1 - r0, width)
def _fold_padding_around_the_globe(
foot_arr: np.ndarray, *, pad: int, n_lon: int
) -> np.ndarray:
"""
Add the padding of a padded array to the cells it lies beside, around the globe.
*foot_arr* has *pad* cells of padding on each side of its first axis and
*n_lon* cells between them. On a grid that goes all the way around, the
cells past the east edge are the first cells again, and the cells past
the west edge are the last, so the smoothing that spilled into the
padding is returned to them instead of being dropped.
"""
folded = np.zeros((n_lon, *foot_arr.shape[1:]), dtype=foot_arr.dtype)
start = 0
while start < foot_arr.shape[0]:
cell = (start - pad) % n_lon
stop = min(start + n_lon - cell, foot_arr.shape[0])
folded[cell : cell + stop - start] += foot_arr[start:stop]
start = stop
return folded
def _accumulate_smoothed_footprint(
p: pd.DataFrame,
*,
layers: np.ndarray,
n_lon_buf: int,
n_lat_buf: int,
rtimes: np.ndarray,
w: np.ndarray,
rs: tuple[float, float],
) -> np.ndarray:
"""
Grid each time step's ``foot`` on the padded grid and smooth it with a Gaussian kernel.
Returns an array of shape ``(n_lon_buf, n_lat_buf, n_layers)``. Each
step's particles are summed by cell, then the step's kernel is added at
each occupied cell (``_stamp_kernel``), as STILT-R's ``permute.f90``
does. This equals a convolution of the gridded step with the kernel, but
costs cells x kernel cells instead of grid cells x kernel cells.
"""
foot_arr = np.zeros((n_lon_buf, n_lat_buf, len(layers)), dtype=float)
kernel_cache: dict[float, np.ndarray] = {}
layer_index = {int(layer): i for i, layer in enumerate(layers)}
for key, step in p.groupby(["layer", "rtime"], sort=False):
layer, rtime_val = cast(tuple[Any, Any], key)
i = layer_index[int(layer)]
# Nearest-neighbor kernel bandwidth for this rtime.
step_w_idx = int(np.argmin(np.abs(rtimes - rtime_val)))
step_w = float(w[step_w_idx])
if step_w not in kernel_cache:
kernel_cache[step_w] = _make_gauss_kernel(rs, step_w)
k = kernel_cache[step_w]
loi_arr = step["loi"].to_numpy().astype(int)
lai_arr = step["lai"].to_numpy().astype(int)
foot_vals = step["foot"].to_numpy(dtype=float)
valid = (
(loi_arr >= 0)
& (loi_arr < n_lon_buf)
& (lai_arr >= 0)
& (lai_arr < n_lat_buf)
)
# Sum the particles in each occupied cell, then stamp the kernel once per cell.
cells, inverse = np.unique(
loi_arr[valid] * n_lat_buf + lai_arr[valid], return_inverse=True
)
if len(cells) == 0:
continue
cell_foot = np.bincount(inverse, weights=foot_vals[valid])
_stamp_kernel(
foot_arr[:, :, i], cells // n_lat_buf, cells % n_lat_buf, cell_foot, k
)
return foot_arr