"""Utilities for horizontal DAG construction."""
import warnings
from collections.abc import Generator
from dataclasses import dataclass
from typing import TYPE_CHECKING, Self
import numpy as np
import numpy.typing as npt
import pandas as pd
import xarray as xr
from pycontrails import MetDataArray, MetDataset
from pycontrails.core import airports
from pycontrails.physics import geo
from contrailopt.grid_utils import (
bilinear_interp,
localize_horizontally,
select_flight_times,
to_altitude_ft,
)
from contrailopt.slerp import gc_interp, gc_npts, spherical_fwd
if TYPE_CHECKING:
from cartopy.mpl.geoaxes import GeoAxes
[docs]
@dataclass(kw_only=True, slots=True, frozen=True)
class AirportCoords:
"""Coordinates and elevation of an airport, identified by ICAO code."""
icao_code: str
longitude: float
latitude: float
elevation_ft: float
[docs]
@classmethod
def from_icao(cls, icao_code: str) -> Self:
"""Look up airport coordinates by ICAO code."""
airports_df = airports.global_airport_database()
row = airports_df.query(f"icao_code == '{icao_code}'")
if row.empty:
raise ValueError(f"Could not find airport with ICAO code {icao_code}")
return cls(
icao_code=icao_code,
longitude=row["longitude"].item(),
latitude=row["latitude"].item(),
elevation_ft=row["elevation_ft"].item(),
)
@property
def coords(self) -> tuple[float, float]:
"""Return (longitude, latitude) coordinates as a tuple."""
return self.longitude, self.latitude
def _csr_flat_pos(
adj_ptr: npt.NDArray[np.int64],
nodes: npt.NDArray[np.int64],
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
"""Return flat indices into CSR data arrays for a batch of row nodes.
Returns ``(flat_pos, lengths)``, where flat_pos indexes into adj
and lengths[i] is the number of entries for nodes[i].
"""
starts = adj_ptr[nodes]
lengths = adj_ptr[nodes + 1] - starts
flat_starts = np.repeat(starts, lengths)
offsets = np.arange(lengths.sum()) - np.repeat(np.cumsum(lengths) - lengths, lengths)
flat_pos = flat_starts + offsets
return flat_pos, lengths
def _neighbors_batch(
adj_ptr: npt.NDArray[np.int64],
adj: npt.NDArray[np.int64],
nodes: npt.NDArray[np.int64],
) -> npt.NDArray[np.int64]:
"""Determine neighbors (duplicates included with multiplicity) for a batch of nodes."""
flat_pos, _ = _csr_flat_pos(adj_ptr, nodes)
return adj[flat_pos]
def _reachability(
seed: int,
n_nodes: int,
adj_ptr: npt.NDArray[np.int64],
adj: npt.NDArray[np.int64],
) -> npt.NDArray[np.bool_]:
"""Return boolean array of nodes reachable from seed via adjacency."""
seen = np.zeros(n_nodes, dtype=bool)
frontier = np.array([seed], dtype=np.int64)
seen[frontier] = True
while frontier.size:
nxt = _neighbors_batch(adj_ptr, adj, frontier) # all outgoing neighbors of current frontier
nxt = np.unique(nxt) # dedupe duplicates from shared parents
frontier = nxt[~seen[nxt]] # only newly discovered nodes
seen[frontier] = True
return seen
def _reverse_csr(
adj: npt.NDArray[np.int64],
n_nodes: int,
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
"""Compute reverse CSR pointer and edge permutation order."""
order = np.argsort(adj)
rev_ptr = np.zeros(n_nodes + 1, dtype=np.int64)
np.add.at(rev_ptr[1:], adj, 1)
np.cumsum(rev_ptr, out=rev_ptr)
return rev_ptr, order
[docs]
@dataclass(kw_only=True, slots=True, frozen=True)
class HorizontalDAG:
"""Directed graph on (lon, lat) nodes with CSR adjacency.
All geometry (distances, azimuths, interpolation, polygon exclusion) is
computed on a sphere, not on a planar lon/lat grid.
"""
#: Longitude of each node in degrees ``(n,)``. Assumed to be in the range [-180, 180).
#: (This assumption is used in :meth:`crosses_antimeridian`).
lon: npt.NDArray[np.floating]
#: Latitude of each node in degrees ``(n,)``. Assumed to be in the range [-90, 90].
lat: npt.NDArray[np.floating]
#: CSR row pointers ``(n + 1,)``.
#: Neighbors of node ``i`` are ``adj[adj_ptr[i]: adj_ptr[i+1]]``.
adj_ptr: npt.NDArray[np.int64]
#: Neighbor (destination) indices for each directed edge ``(m,)``.
adj: npt.NDArray[np.int64]
#: Great-circle distance in meters for each directed edge ``(m,)``.
edge_dist: npt.NDArray[np.floating]
#: Index of the distinguished origin node.
h_origin: int
#: Index of the distinguished destination node.
h_dest: int
def __repr__(self) -> str:
name = type(self).__name__
return f"{name}({self.n_nodes} nodes, {self.n_edges} edges)"
def __eq__(self, other: object) -> bool:
if not isinstance(other, HorizontalDAG):
# https://docs.python.org/3/reference/datamodel.html#object.__eq__
return NotImplemented
return (
self.h_origin == other.h_origin
and self.h_dest == other.h_dest
and np.array_equal(self.lon, other.lon)
and np.array_equal(self.lat, other.lat)
and np.array_equal(self.adj_ptr, other.adj_ptr)
and np.array_equal(self.adj, other.adj)
and np.array_equal(self.edge_dist, other.edge_dist)
)
@property
def n_nodes(self) -> int:
"""The number of nodes in the graph."""
return len(self.lon)
@property
def n_edges(self) -> int:
"""The number of directed edges in the graph."""
return len(self.adj)
@property
def out_degree(self) -> npt.NDArray[np.int64]:
"""The out-degree of each node."""
return np.diff(self.adj_ptr)
@property
def edge_src(self) -> npt.NDArray[np.int64]:
"""The source endpoint index of each directed edge."""
return np.repeat(np.arange(self.n_nodes, dtype=np.int64), self.out_degree)
@property
def edges(self) -> npt.NDArray[np.int64]:
"""The directed edges of the graph as (src, dest) index pairs in an ``(m, 2)`` array."""
return np.column_stack([self.edge_src, self.adj])
@property
def crosses_antimeridian(self) -> bool:
"""Determine if the great circle between origin and destination crosses the antimeridian."""
return abs(self.lon[self.h_origin].item() - self.lon[self.h_dest].item()) > 180.0
[docs]
def neighbors(self, i: int) -> npt.NDArray[np.int64]:
"""Return the neighbors of a specified node."""
return self.adj[self.adj_ptr[i] : self.adj_ptr[i + 1]]
[docs]
def edge_index(self, src: int, dst: int) -> int:
"""Return the CSR index of the directed edge (src, dst).
Performs a linear scan over the neighbors of ``src``. This could be
replaced with ``np.searchsorted`` if needed, but that would require
enforcing sorted neighbors during construction.
Raises
------
ValueError
If no edge from ``src`` to ``dst`` exists.
"""
start = self.adj_ptr[src].item()
end = self.adj_ptr[src + 1]
row = self.adj[start:end] # row = self.neighbors(src), but we need start again
pos = np.flatnonzero(row == dst)
if len(pos) == 0:
raise ValueError(f"No edge from {src} to {dst}")
return start + pos.item()
[docs]
def neighbors_batch(self, nodes: npt.NDArray[np.int64]) -> npt.NDArray[np.int64]:
"""Return neighbors (duplicates included with multiplicity) for a batch of nodes."""
return _neighbors_batch(self.adj_ptr, self.adj, nodes)
[docs]
def expand_neighbors(
self, nodes: npt.NDArray[np.int64]
) -> tuple[
npt.NDArray[np.int64],
npt.NDArray[np.floating],
npt.NDArray[np.int64],
npt.NDArray[np.int64],
]:
"""Expand CSR adjacency for a batch of nodes into flat edge arrays.
Returns
-------
flat_nbr : npt.NDArray[np.int64]
``(e,)`` neighbor indices for all edges leaving ``nodes``.
flat_dist : npt.NDArray[np.floating]
``(e,)`` edge distances in meters.
src_idx : npt.NDArray[np.int64]
``(e,)`` index into ``nodes`` for each flat entry, so
``nodes[src_idx[k]]`` is the source node of flat edge ``k``.
flat_edge_idx : npt.NDArray[np.int64]
``(e,)`` index of each edge in the CSR arrays (``adj``, ``edge_dist``).
Here ``e = out_degree[nodes].sum()``, the total number of outgoing
edges from all ``nodes``.
"""
flat_pos, lengths = _csr_flat_pos(self.adj_ptr, nodes)
src_idx = np.repeat(np.arange(len(nodes)), lengths)
return self.adj[flat_pos], self.edge_dist[flat_pos], src_idx, flat_pos
[docs]
def adjacency_matrix(self) -> npt.NDArray[np.bool]:
"""Return dense boolean adjacency matrix A where A[i, j] is True for i->j."""
matrix = np.zeros((self.n_nodes, self.n_nodes), dtype=bool)
matrix[self.edge_src, self.adj] = True
return matrix
[docs]
def distance_matrix(self, missing: float = np.inf) -> npt.NDArray[np.floating]:
"""Return dense distance matrix A where A[i, j] is the distance aong i->j.
Distances are set to infinity by default where edges are missing.
"""
matrix = np.full((self.n_nodes, self.n_nodes), missing, dtype=self.edge_dist.dtype)
matrix[self.edge_src, self.adj] = self.edge_dist
return matrix
[docs]
def reverse(self) -> Self:
"""Return a new DAG with all edge directions flipped and origin/dest swapped."""
src = self.edge_src
rev_ptr, rev_order = _reverse_csr(self.adj, self.n_nodes)
rev_adj = src[rev_order]
rev_edge_dist = self.edge_dist[rev_order]
return type(self)(
lon=self.lon,
lat=self.lat,
adj_ptr=rev_ptr,
adj=rev_adj,
edge_dist=rev_edge_dist,
h_origin=self.h_dest,
h_dest=self.h_origin,
)
[docs]
def prune_unreachable(self) -> Self:
"""Return a new DAG with only nodes reachable from origin that also reach dest."""
src = self.edge_src
fwd = _reachability(self.h_origin, self.n_nodes, self.adj_ptr, self.adj)
rev_ptr, rev_order = _reverse_csr(self.adj, self.n_nodes)
rev_adj = src[rev_order]
bwd = _reachability(self.h_dest, self.n_nodes, rev_ptr, rev_adj)
live = fwd & bwd
if not live[self.h_origin] or not live[self.h_dest]:
raise ValueError("Origin or destination became unreachable after pruning")
# Remap node indices
n = live.sum()
new_idx = np.full(self.n_nodes, -1, dtype=np.int64)
new_idx[live] = np.arange(n, dtype=np.int64)
# Filter edges and distances via CSR ordering
mask = live[src] & live[self.adj]
new_src = new_idx[src[mask]]
new_dst = new_idx[self.adj[mask]]
new_edge_dist = self.edge_dist[mask]
# Build new CSR (already sorted by source from parent CSR)
adj_ptr = np.zeros(n + 1, dtype=np.int64)
np.add.at(adj_ptr[1:], new_src, 1)
np.cumsum(adj_ptr, out=adj_ptr)
return type(self)(
lon=self.lon[live],
lat=self.lat[live],
adj_ptr=adj_ptr,
adj=new_dst,
edge_dist=new_edge_dist,
h_origin=new_idx[self.h_origin].item(),
h_dest=new_idx[self.h_dest].item(),
)
[docs]
def prune_edges(self, degree: int) -> Self:
"""Keep the best ``degree`` outgoing and incoming edges per node.
Each edge is scored by the sum of its azimuth deviations: how far the
edge direction deviates from the azimuth toward the destination (at
the tail) plus how far the reverse deviates from the azimuth toward
the origin (at the head). An edge is kept if it ranks among the best
``degree`` outgoing edges of its source or among the best ``degree``
incoming edges of its destination.
Parameters
----------
degree : int
Number of outgoing and incoming edges to keep per node.
Returns
-------
Self
A new DAG with at most ``degree`` outgoing and incoming edges per node.
"""
# Compute azimuth deviation scores for each edge
src = self.edge_src
az_to_dest = geo.azimuth(self.lon, self.lat, self.lon[self.h_dest], self.lat[self.h_dest])
az_to_origin = geo.azimuth(
self.lon, self.lat, self.lon[self.h_origin], self.lat[self.h_origin]
)
az_edge = geo.azimuth(self.lon[src], self.lat[src], self.lon[self.adj], self.lat[self.adj])
az_rev = geo.azimuth(self.lon[self.adj], self.lat[self.adj], self.lon[src], self.lat[src])
delta_tail = np.abs((az_edge - az_to_dest[src] + 180.0) % 360.0 - 180.0)
delta_head = np.abs((az_rev - az_to_origin[self.adj] + 180.0) % 360.0 - 180.0)
score = delta_tail + delta_head # many other variations also work: max, min, p-weighted
# Rank each edge among its source node's outgoing edges by score
out_order = np.lexsort((score, src))
out_rank = np.empty(self.n_edges, dtype=np.int64)
out_rank[out_order] = np.arange(self.n_edges) - self.adj_ptr[src[out_order]]
# Rank each edge among its dest node's incoming edges by score
dst = self.adj
rev_ptr, _ = _reverse_csr(dst, self.n_nodes)
in_order = np.lexsort((score, dst))
in_rank = np.empty(self.n_edges, dtype=np.int64)
in_rank[in_order] = np.arange(self.n_edges) - rev_ptr[dst[in_order]]
keep = (out_rank < degree) | (in_rank < degree)
# Construct new CSR to return
kept_src = src[keep]
kept_dst = self.adj[keep]
kept_dist = self.edge_dist[keep]
adj_ptr = np.zeros(self.n_nodes + 1, dtype=np.int64)
np.add.at(adj_ptr[1:], kept_src, 1)
np.cumsum(adj_ptr, out=adj_ptr)
return type(self)(
lon=self.lon,
lat=self.lat,
adj_ptr=adj_ptr,
adj=kept_dst,
edge_dist=kept_dist,
h_origin=self.h_origin,
h_dest=self.h_dest,
)
[docs]
def exclude_polygons(self, polygons: list[list[tuple[float, float]]]) -> Self:
"""Return a new DAG with edges crossing any polygon removed.
Uses spherely for geodesic intersection tests on the sphere.
Parameters
----------
polygons : list[list[tuple[float, float]]]
List of polygons, where each polygon is a list of ``(lon, lat)`` vertices.
Returns
-------
HorizontalDAG
A new DAG with offending edges removed and then pruned.
"""
import spherely
src = self.edge_src
# Build a linestring for every edge
linestrings = [
spherely.create_linestring([(self.lon[s], self.lat[s]), (self.lon[d], self.lat[d])])
for s, d in zip(src, self.adj, strict=True)
]
edge_geoms = np.array(linestrings)
# Test each polygon against all edges. We could also take a union of all polygons
# and test once if this becomes a bottleneck
excluded = np.zeros(self.n_edges, dtype=bool)
for coords in polygons:
poly = spherely.create_polygon(coords)
excluded |= spherely.intersects(poly, edge_geoms)
# Build filtered edge list
keep = ~excluded
kept_src = src[keep]
kept_dst = self.adj[keep]
kept_dist = self.edge_dist[keep]
# Rebuild CSR
adj_ptr = np.zeros(self.n_nodes + 1, dtype=np.int64)
np.add.at(adj_ptr[1:], kept_src, 1)
np.cumsum(adj_ptr, out=adj_ptr)
return type(self)(
lon=self.lon,
lat=self.lat,
adj_ptr=adj_ptr,
adj=kept_dst,
edge_dist=kept_dist,
h_origin=self.h_origin,
h_dest=self.h_dest,
).prune_unreachable()
[docs]
def sample_edges(
self, spacing_m: float
) -> tuple[
npt.NDArray[np.floating],
npt.NDArray[np.floating],
npt.NDArray[np.int64],
npt.NDArray[np.int64],
]:
"""Sample points along every edge at most ``spacing_m`` meters apart.
Points are uniformly spaced along each edge, and both edge endpoints
(source and destination nodes) are included as samples.
Returns
-------
sample_lon : npt.NDArray[np.floating]
``(s,)`` longitude of each sample point.
sample_lat : npt.NDArray[np.floating]
``(s,)`` latitude of each sample point.
edge_idx : npt.NDArray[np.int64]
``(s,)`` edge index for each sample point.
edge_ptr : npt.NDArray[np.int64]
``(m + 1,)`` CSR-style pointer so edge ``i``'s samples are at
``sample_lon[edge_ptr[i]: edge_ptr[i+1]]``.
Here ``s = edge_ptr[-1]``, the total number of sample points across all edges.
"""
src = self.edge_src
src_lon = self.lon[src]
src_lat = self.lat[src]
dst_lon = self.lon[self.adj]
dst_lat = self.lat[self.adj]
dist = geo.haversine(src_lon, src_lat, dst_lon, dst_lat)
n_samples = np.maximum(np.ceil(dist / spacing_m).astype(int) + 1, 2)
edge_ptr = np.zeros(self.n_edges + 1, dtype=np.int64)
np.cumsum(n_samples, out=edge_ptr[1:])
total = edge_ptr[-1]
local_idx = np.arange(total) - np.repeat(edge_ptr[:-1], n_samples)
frac = local_idx / np.repeat(n_samples - 1, n_samples)
sample_lon, sample_lat = gc_interp(
np.repeat(src_lon, n_samples),
np.repeat(src_lat, n_samples),
np.repeat(dst_lon, n_samples),
np.repeat(dst_lat, n_samples),
frac.astype(src_lon.dtype),
)
edge_idx = np.repeat(np.arange(self.n_edges), n_samples)
return sample_lon, sample_lat, edge_idx, edge_ptr
[docs]
def plot(
self,
ax: "GeoAxes | None" = None,
linewidth: float = 2.0,
show_edges: bool = True,
) -> "GeoAxes":
"""Plot the DAG on a cartopy map."""
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
data_crs = ccrs.Geodetic()
if ax is None:
central_lon = 180.0 if self.crosses_antimeridian else 0.0
proj = ccrs.PlateCarree(central_longitude=central_lon)
_, ax = plt.subplots(subplot_kw={"projection": proj})
ax.set_extent(
[
self.lon.min() - 2.0,
self.lon.max() + 2.0,
self.lat.min() - 2.0,
self.lat.max() + 2.0,
],
crs=data_crs,
)
ax.add_feature(cfeature.LAND, facecolor="whitesmoke")
ax.add_feature(cfeature.COASTLINE, linewidth=0.5)
ax.add_feature(cfeature.BORDERS, linewidth=0.5, edgecolor="gray")
ax.add_feature(cfeature.STATES, linewidth=0.2, edgecolor="gray")
# Draw edges
if show_edges:
edge_src = self.edge_src
segments = np.stack(
[
np.column_stack([self.lon[edge_src], self.lat[edge_src]]),
np.column_stack([self.lon[self.adj], self.lat[self.adj]]),
],
axis=1,
)
lc = LineCollection(
segments,
colors="steelblue",
linewidths=linewidth,
alpha=0.2,
transform=data_crs,
)
ax.add_collection(lc)
# Draw nodes
ax.scatter(self.lon, self.lat, s=2, color="black", transform=data_crs, zorder=5)
ax.plot(
self.lon[self.h_origin],
self.lat[self.h_origin],
"ro",
markersize=8,
transform=data_crs,
zorder=10,
)
ax.plot(
self.lon[self.h_dest],
self.lat[self.h_dest],
"go",
markersize=8,
transform=data_crs,
zorder=10,
)
ax.set_title(f"{self.n_edges} edges, {self.n_nodes} nodes")
return ax
[docs]
@classmethod
def from_network(
cls,
lon: npt.NDArray[np.floating],
lat: npt.NDArray[np.floating],
tail: npt.NDArray[np.int64],
head: npt.NDArray[np.int64],
origin_idx: int = 0,
dest_idx: int = -1,
max_angle_deg: float = 40.0
) -> Self:
"""Build a DAG from a static network graph using the dual azimuth constraint."""
edges, dists = _dual_az_edges(lon, lat, tail, head, origin_idx, dest_idx, max_angle_deg)
n = len(lon)
order = np.argsort(edges[:, 0])
adj_ptr = np.zeros(n + 1, dtype=np.int64)
np.add.at(adj_ptr[1:], edges[order, 0], 1)
np.cumsum(adj_ptr, out=adj_ptr)
return cls(
lon=lon,
lat=lat,
adj_ptr=adj_ptr,
adj=edges[order, 1],
edge_dist=dists[order],
h_origin=origin_idx if origin_idx >= 0 else n + origin_idx,
h_dest=dest_idx if dest_idx >= 0 else n + dest_idx,
)
[docs]
@classmethod
def from_points(
cls,
lon: npt.NDArray[np.floating],
lat: npt.NDArray[np.floating],
origin_idx: int = 0,
dest_idx: int = -1,
max_angle_deg: float = 40.0,
max_dist_m: float = 500_000.0,
) -> Self:
"""Build a DAG from lon/lat arrays using the dual azimuth constraint."""
tail, head = _neighborhood_edges(lon, lat, max_dist_m)
return cls.from_network(lon, lat, tail, head, origin_idx, dest_idx, max_angle_deg)
[docs]
@classmethod
def from_poisson(
cls,
origin_lon: float,
origin_lat: float,
dest_lon: float,
dest_lat: float,
poisson_spacing_m: float = 80_000.0,
max_cross_track: float | None = None,
max_angle_deg: float = 40.0,
max_dist_m: float = 500_000.0,
dtype: type[np.floating] = np.float64,
rng: np.random.Generator | None = None,
) -> Self:
"""Build a DAG from Poisson-disk sampled points along the OD great circle."""
from scipy.stats.qmc import PoissonDisk
gs_distance = geo.haversine(origin_lon, origin_lat, dest_lon, dest_lat).item()
if max_cross_track is None:
max_cross_track = min(1_000_000.0, gs_distance / 4.0)
nx = int(gs_distance / 100_000.0) + 1
# Great circle spine
gc_lons, gc_lats = gc_npts(origin_lon, origin_lat, dest_lon, dest_lat, nx - 2)
gc_lons = np.concatenate([[origin_lon], gc_lons, [dest_lon]])
gc_lats = np.concatenate([[origin_lat], gc_lats, [dest_lat]])
# Perpendicular azimuths along spine
az_fwd = geo.azimuth(gc_lons[:-1], gc_lats[:-1], gc_lons[1:], gc_lats[1:])
az_perp = np.empty(nx)
az_perp[:-1] = az_fwd + 90.0
az_perp[-1] = az_perp[-2]
# Poisson disk sampling with physically uniform spacing.
# The corridor is gs_distance x 2*max_cross_track in physical space.
# We sample in [0, aspect] x [0, 1] so the radius is isotropic in
# physical units, then normalize the along-track coordinate.
corridor_width = 2.0 * max_cross_track
aspect = gs_distance / corridor_width
unit_radius = poisson_spacing_m / corridor_width
poisson_disk = PoissonDisk(
d=2,
radius=unit_radius,
l_bounds=[0, 0],
u_bounds=[aspect, 1],
rng=rng,
)
pts = poisson_disk.fill_space()
t = pts[:, 0] / aspect
cross = pts[:, 1] * corridor_width - max_cross_track
# Prepend origin and append dest
t = np.concatenate([[0.0], t, [1.0]])
cross = np.concatenate([[0.0], cross, [0.0]], dtype=dtype)
# Project to (lon, lat) dealing with potential antimeridian crossing via unwrapping.
t_gc = np.linspace(0.0, 1.0, nx)
gc_lons_unwrap = np.unwrap(gc_lons, period=360.0)
az_perp_unwrap = np.unwrap(az_perp, period=360.0)
lon_base = np.interp(t, t_gc, gc_lons_unwrap).astype(dtype)
lon_base = (lon_base + 180.0) % 360.0 - 180.0
lat_base = np.interp(t, t_gc, gc_lats).astype(dtype)
az_base = np.interp(t, t_gc, az_perp_unwrap).astype(dtype)
lon, lat = spherical_fwd(lon_base, lat_base, az_base, cross)
lon = (lon + 180.0) % 360.0 - 180.0
return cls.from_points(
lon,
lat,
origin_idx=0,
dest_idx=len(lon) - 1,
max_angle_deg=max_angle_deg,
max_dist_m=max_dist_m,
)
[docs]
def topo_wavefronts(self) -> Generator[npt.NDArray[np.int64], None, None]:
"""Yield topological wavefronts reachable from origin.
Only nodes reachable from ``h_origin`` are emitted. Unreachable nodes
(those with incoming edges from outside the reachable subgraph) are
excluded.
The first wavefront contains only the origin. Wavefront *k* contains
nodes whose reachable in-degree drops to zero after removing wavefronts
0 ... k-1.
"""
reachable = _reachability(self.h_origin, self.n_nodes, self.adj_ptr, self.adj)
live_edges = reachable[self.edge_src]
in_degree = np.zeros(self.n_nodes, dtype=np.int64)
np.add.at(in_degree, self.adj[live_edges], 1)
wave_nodes = np.array([self.h_origin])
while wave_nodes.size:
yield wave_nodes
neighbors = self.neighbors_batch(wave_nodes)
np.subtract.at(in_degree, neighbors, 1)
candidates = np.unique(neighbors)
filt = in_degree[candidates] == 0
wave_nodes = candidates[filt]
[docs]
@dataclass(kw_only=True, slots=True, frozen=True)
class Track:
"""An ordered sequence of timed waypoints along a flown path.
The flight-profile optimizer (:func:`contrailopt.optimize.solve_track`) follows a fixed
lateral path, so it needs only the ordered nodes, their times, and along-track distances.
This interface supplies ``lon``, ``lat``, ``node_time``, ``cum_dist``,
``segment_dist``, ``n_nodes``, ``h_origin``, ``h_dest`` computed straight from
the coordinates, so it can often be used in place of :class:`HorizontalDAG`.
"""
lon: npt.NDArray[np.floating]
lat: npt.NDArray[np.floating]
node_time: npt.NDArray[np.datetime64]
@property
def n_nodes(self) -> int:
"""The number of waypoints."""
return len(self.lon)
@property
def h_origin(self) -> int:
"""Index of the origin node (always the first waypoint)."""
return 0
@property
def h_dest(self) -> int:
"""Index of the destination node (always the last waypoint)."""
return len(self.lon) - 1
@property
def segment_dist(self) -> npt.NDArray[np.floating]:
"""Great-circle distance between consecutive waypoints ``(n - 1,)``."""
return geo.haversine(self.lon[:-1], self.lat[:-1], self.lon[1:], self.lat[1:])
@property
def cum_dist(self) -> npt.NDArray[np.floating]:
"""Cumulative along-track distance at each waypoint ``(n,)``, starting at zero."""
out = np.zeros(self.n_nodes, dtype=self.lon.dtype)
np.cumsum(self.segment_dist, out=out[1:])
return out
@property
def crosses_antimeridian(self) -> bool:
"""Whether the origin-to-destination span wraps the antimeridian."""
return abs(self.lon[0].item() - self.lon[-1].item()) > 180.0
[docs]
def plot(self, ax: "GeoAxes | None" = None, linewidth: float = 2.0) -> "GeoAxes":
"""Plot the track on a cartopy map."""
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
data_crs = ccrs.Geodetic()
if ax is None:
central_lon = 180.0 if self.crosses_antimeridian else 0.0
proj = ccrs.PlateCarree(central_longitude=central_lon)
_, ax = plt.subplots(subplot_kw={"projection": proj})
ax.set_extent(
[
self.lon.min() - 2.0,
self.lon.max() + 2.0,
self.lat.min() - 2.0,
self.lat.max() + 2.0,
],
crs=data_crs,
)
ax.add_feature(cfeature.LAND, facecolor="whitesmoke")
ax.add_feature(cfeature.COASTLINE, linewidth=0.5)
ax.add_feature(cfeature.BORDERS, linewidth=0.5, edgecolor="gray")
ax.add_feature(cfeature.STATES, linewidth=0.2, edgecolor="gray")
ax.plot(self.lon, self.lat, color="steelblue", linewidth=linewidth, transform=data_crs)
ax.plot(self.lon[0], self.lat[0], "ro", markersize=8, transform=data_crs, zorder=10)
ax.plot(self.lon[-1], self.lat[-1], "go", markersize=8, transform=data_crs, zorder=10)
ax.set_title(f"{self.n_nodes} waypoints")
return ax
def _warn_if_ef_present(ds: xr.Dataset) -> None:
"""Warn if the input carries some variant of energy forcing other than 'eef_per_m'."""
for key in ("ef", "eef", "ef_per_m", "energy_forcing", "effective_energy_forcing"):
if key in ds:
warnings.warn(
f"Found '{key}' in the input dataset; it will be ignored. The solver expects "
"effective energy forcing per distance in J / m under the name 'eef_per_m'.",
stacklevel=3,
)
[docs]
def validate_flight_profile(ds: xr.Dataset, n_nodes: int) -> xr.Dataset:
"""Validate and format a ``(waypoint, altitude_ft)`` flight profile for the track solver.
Returns a dataset carrying the met variables under their original names
(``air_temperature``, ``u_wind``, ``v_wind``, and optional ``eef_per_m``),
each cast to float32 and oriented ``(waypoint, altitude_ft)``.
NaN in the core weather variables raises, while NaN in ``eef_per_m`` is zero-filled.
"""
if ds.sizes["waypoint"] != n_nodes:
raise ValueError(f"ds has {ds.sizes['waypoint']} waypoints but dag has {n_nodes} nodes")
_warn_if_ef_present(ds)
for required in ("air_temperature", "u_wind", "v_wind"):
if required not in ds:
raise ValueError(f"flight profile is missing required variable '{required}'")
data_vars = {}
for name in ("air_temperature", "u_wind", "v_wind", "eef_per_m"):
if name not in ds:
continue
col = ds[name].transpose("waypoint", "altitude_ft").values.astype(np.float32, copy=False)
if name == "eef_per_m":
col = np.nan_to_num(col, nan=0.0)
elif np.isnan(col).any():
raise ValueError(f"NaN values found in '{name}'")
data_vars[name] = (("waypoint", "altitude_ft"), col)
return xr.Dataset(data_vars, coords={"altitude_ft": ds["altitude_ft"].values})
[docs]
@dataclass(kw_only=True, slots=True, frozen=True)
class EdgeInterpolation:
"""Met fields interpolated at sample points."""
air_temperature: npt.NDArray[np.floating]
eastward_wind: npt.NDArray[np.floating]
northward_wind: npt.NDArray[np.floating]
eef_per_m: npt.NDArray[np.floating] | None
[docs]
@dataclass(kw_only=True, slots=True, frozen=True)
class EdgeMetLookup:
"""Pre-interpolated met data on edge sample points."""
#: ``xr.Dataset`` with dims ``(sample, altitude_ft, time)`` containing
#: weather variables interpolated onto edge sample coordinates.
ds: xr.Dataset
#: CSR-style pointer array ``(n_edges + 1,)``. Samples for edge ``i``
#: are at indices ``edge_ptr[i]:edge_ptr[i+1]``.
edge_ptr: npt.NDArray[np.int64]
#: Edge index for each sample point ``(n_samples,)``.
edge_idx: npt.NDArray[np.int64]
#: Longitude of each sample point ``(n_samples,)``.
sample_lon: npt.NDArray[np.floating]
#: Latitude of each sample point ``(n_samples,)``.
sample_lat: npt.NDArray[np.floating]
#: Cumulative distance from edge source to each sample point in meters ``(n_samples,)``.
cum_dist: npt.NDArray[np.floating]
#: Distance in meters from this sample to the next ``(n_samples,)``.
#: The last sample of each edge has ``delta_dist = 0``.
#: Equal to ``diff(cum_dist)`` within each edge.
delta_dist: npt.NDArray[np.floating]
#: Azimuth in radians from each sample to the next ``(n_samples,)``.
#: The last sample of each edge copies the previous sample's azimuth.
sample_azimuth: npt.NDArray[np.floating]
def __post_init__(self) -> None:
required = {"air_temperature", "eastward_wind", "northward_wind"}
missing = required - set(self.ds)
if missing:
raise ValueError(f"Met dataset missing required variables: {missing}")
def __repr__(self) -> str:
n_samples = len(self.edge_idx)
n_edges = len(self.edge_ptr) - 1
n_fl = self.ds.sizes["altitude_ft"]
n_time = self.ds.sizes["time"]
name = type(self).__name__
return f"{name}({n_edges} edges, {n_samples} samples, {n_fl} FLs, {n_time} time steps)"
def __call__(
self,
sample_idxs: npt.NDArray[np.int64],
times: npt.NDArray[np.datetime64],
fl_idx: npt.NDArray[np.int64] | None = None,
) -> EdgeInterpolation:
"""Interpolate all variables at given sample indices and times.
Parameters
----------
sample_idxs : npt.NDArray[np.int64]
1D array of sample indices to query.
times : npt.NDArray[np.datetime64]
2D array of time coordinates with shape ``(n_sample, n_fl)``, where
``n_sample = len(sample_idxs)``. Each FL gets its own query time
(e.g. to account for FL-dependent climb duration).
fl_idx : npt.NDArray[np.int64] | None
Flight level indices into the ``altitude_ft`` dimension. If ``None``
(default), all FLs are returned with shape ``(n_sample, n_fl)``.
If an array, outputs are ``(n_sample, len(fl_idx))``.
Returns
-------
EdgeInterpolation
Interpolated met fields at the requested sample and time coordinates.
"""
time_coords = self.ds["time"].values # (n_time,) datetime64[ns]
time_s = (time_coords - time_coords[0]) / np.timedelta64(1, "s")
query_s = (times - time_coords[0]) / np.timedelta64(1, "s")
n_time = len(time_coords)
fp = np.arange(n_time, dtype=np.float64)
t_frac = np.interp(query_s, time_s, fp).astype(np.float32) # np.interp returns float64
if np.any(~np.isfinite(t_frac)): # idiot check
raise RuntimeError("Non-finite t_frac values")
t_lo = np.floor(t_frac).astype(np.int16) # n_time << int16.max, and f32 + int16 = f32
t_hi = np.minimum(t_lo + 1, n_time - 1)
w = t_frac - t_lo
fl_idx = np.arange(self.ds.sizes["altitude_ft"]) if fl_idx is None else fl_idx
def _lerp(name: str) -> npt.NDArray[np.floating]:
data = self.ds[name].values # (n_total_samples, n_fl, n_time)
lo = data[sample_idxs[:, np.newaxis], fl_idx[np.newaxis, :], t_lo]
hi = data[sample_idxs[:, np.newaxis], fl_idx[np.newaxis, :], t_hi]
return lo + w * (hi - lo)
return EdgeInterpolation(
air_temperature=_lerp("air_temperature"),
eastward_wind=_lerp("eastward_wind"),
northward_wind=_lerp("northward_wind"),
eef_per_m=_lerp("eef_per_m") if "eef_per_m" in self.ds else None,
)
[docs]
@classmethod
def from_met(
cls,
met: MetDataset | xr.Dataset,
dag: HorizontalDAG,
altitude_ft: npt.NDArray[np.floating],
takeoff_time: pd.Timestamp,
flight_hours: int,
spacing_m: float,
eef: xr.DataArray | MetDataArray | None = None,
) -> Self:
"""Interpolate met data onto ``dag`` edge sample points.
Parameters
----------
met : MetDataset | xr.Dataset
Gridded met dataset with "air_temperature", "eastward_wind", and "northward_wind".
If "eef_per_m" is present, it will also be included in the output with
NaN values filled to 0.0 (no EEF forecast is treated as zero forcing).
NaN values in weather variables are not allowed and will raise an error.
Either a pycontrails ``MetDataset`` or a raw ``xr.Dataset`` with similar
structure can be passed.
dag : HorizontalDAG
Horizontal DAG whose edges will be sampled.
altitude_ft : npt.NDArray[np.floating]
An array of altitudes in feet to interpolate onto.
takeoff_time : pd.Timestamp
Departure time for the flight, used to select met time steps.
flight_hours : int
Number of hourly time steps to retain starting from takeoff_time.
spacing_m : float
Spacing in meters between sample points along edges. Passed to ``dag.sample_edges``.
eef : xr.DataArray | MetDataArray | None, default None
Optional "eef_per_m" DataArray on its own lon/lat grid.
If provided, EEF is interpolated onto sample points independently from the
weather grid, avoiding the need to pre-merge onto a common grid.
Takes precedence over "eef_per_m" in ``met`` if both are present.
Assumed to adhere to pycontrails ``MetDataArray`` conventions.
Returns
-------
EdgeMetLookup
EdgeMetLookup with weather interpolated onto ``(sample, altitude_ft, time)`` dims.
"""
sample_lon, sample_lat, edge_idx, edge_ptr = dag.sample_edges(spacing_m=spacing_m)
# Compute distance from edge source to each sample
edge_src = dag.edge_src[edge_idx]
src_lon = dag.lon[edge_src]
src_lat = dag.lat[edge_src]
cum_dist = geo.haversine(src_lon, src_lat, sample_lon, sample_lat)
cum_dist[edge_ptr[:-1]] = 0.0 # defensive, not strictly needed
# Compute distance and azimuth from one sample to the next (used for wind calcs)
last = edge_ptr[1:] - 1
delta_dist = np.empty_like(cum_dist)
delta_dist[:-1] = np.diff(cum_dist)
delta_dist[last] = 0.0
sample_azimuth = np.empty_like(cum_dist)
sample_azimuth[:-1] = np.deg2rad(
geo.azimuth(sample_lon[:-1], sample_lat[:-1], sample_lon[1:], sample_lat[1:])
)
sample_azimuth[last] = sample_azimuth[last - 1] # copy previous azimuth for last sample
if isinstance(met, MetDataset):
ds = met.data
metdataset_init = False
else:
ds = met
metdataset_init = True
_warn_if_ef_present(ds)
# Ensure variables
variables = ["air_temperature", "eastward_wind", "northward_wind"]
if "eef_per_m" in ds and eef is None:
variables.append("eef_per_m")
ds = ds[variables]
# Downselect in time
usable = select_flight_times(ds, takeoff_time, flight_hours)
ds = ds.sel(time=usable)
# Now run through MetDataset constructor if needed. Do after .sel(time) for memory-sake
if metdataset_init:
ds = MetDataset(ds).data
# Downselect horizontally to reduce memory consumption
ds = localize_horizontally(ds, sample_lon, sample_lat)
# Convert to altitude_ft coordinates and select vertically on FL choices
ds = to_altitude_ft(ds, altitude_ft)
# Interpolate horizontally onto sample points
# Calling ds.interp chews up too much memory and the pycontrails RGI isn't
# exactly designed for this, so just call custom numpy-based bilinear_interp
ds = bilinear_interp(ds, sample_lon, sample_lat)
# Raise on NaN in core weather - downstream computations would be poisoned.
for var in ("air_temperature", "eastward_wind", "northward_wind"):
if ds[var].isnull().any():
raise ValueError(f"NaN values found in '{var}' after interpolation onto samples")
# NaN-fill eef_per_m with 0.0. If NaNs are kept, downstream computations would be poisoned.
if "eef_per_m" in ds: # if eef is not None, this is skipped
ds["eef_per_m"] = ds["eef_per_m"].fillna(0.0)
# If a separate eef DataArray is provided, interpolate it onto sample points independently
if eef is not None:
if isinstance(eef, MetDataArray):
da_eef = eef.data
metdataset_init = False
else:
da_eef = eef
metdataset_init = True
ds_eef = da_eef.to_dataset(name="eef_per_m")
# Downselect in time
ds_eef = ds_eef.sel(time=usable)
# Now run through MetDataset constructor if needed. Do after .sel(time) for memory-sake
if metdataset_init:
ds_eef = MetDataset(ds_eef).data
# Downselect horizontally to reduce memory consumption
ds_eef = localize_horizontally(ds_eef, sample_lon, sample_lat)
ds_eef = to_altitude_ft(ds_eef, altitude_ft)
da_eef = bilinear_interp(ds_eef, sample_lon, sample_lat)["eef_per_m"].fillna(0.0)
# Bypass xarray coord alignment - eef and met altitude_ft values may differ slightly
# snapping to the same altitude_ft (we use sel(..., method="nearest", tolerance=50.0))
# in some places), so we can't rely on xarray to automatically align
ds["eef_per_m"] = (("sample", "altitude_ft", "time"), da_eef.values)
return cls(
ds=ds,
edge_ptr=edge_ptr,
edge_idx=edge_idx,
sample_lon=sample_lon,
sample_lat=sample_lat,
cum_dist=cum_dist,
delta_dist=delta_dist,
sample_azimuth=sample_azimuth,
)
def _neighborhood_edges(
lon: npt.NDArray[np.floating],
lat: npt.NDArray[np.floating],
max_dist_m: float = 500_000.0
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
"""Build network of directed edges based on distance constraint.
A pair of directed edges tail -> head is returned for each pair
of nodes within ``max_dist_m`` of each other.
Returns
-------
tail : npt.NDArray[np.int64]
``(m,)`` array of tail indices
head : npt.NDArray[np.int64]
``(m,)`` array of head indices
"""
dist = geo.haversine(
lon[:, np.newaxis],
lat[:, np.newaxis],
lon[np.newaxis, :],
lat[np.newaxis, :],
)
return np.nonzero((dist > 0.0) & (dist <= max_dist_m))
def _dual_az_edges(
lon: npt.NDArray[np.floating],
lat: npt.NDArray[np.floating],
tail: npt.NDArray[np.int64],
head: npt.NDArray[np.int64],
origin_idx: int,
dest_idx: int,
max_angle_deg: float = 40.0
) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.floating]]:
"""Filter directed edges using a dual azimuth constraint.
For each pair of nodes within ``max_dist_m``, the directed edge tail -> head
is included iff:
1. The azimuth from tail to head is within ``max_angle_deg`` of the azimuth
from tail to the destination (node ``dest_idx``).
2. The azimuth from head to tail is within ``max_angle_deg`` of the azimuth
from head to the origin (node ``origin_idx``).
These two conditions ensure that each edge roughly points toward the destination
and away from the origin. Together, they constrain edges to lie within a
football-shaped corridor between origin and destination and guarantee that each
edge is forward-pointing.
Returns
-------
edges : npt.NDArray[np.int64]
``(m, 2)`` array of ``[tail, head]`` index pairs.
edge_dist : npt.NDArray[np.floating]
``(m,)`` haversine distances in meters. The dtype matches the input lon/lat dtype.
"""
# Precompute per-node azimuths from each node to dest and origin
az_to_dest = geo.azimuth(lon, lat, lon[dest_idx], lat[dest_idx])
az_to_origin = geo.azimuth(lon, lat, lon[origin_idx], lat[origin_idx])
# Compute azimuths for all candidate edges
az_at_tail = geo.azimuth(lon[tail], lat[tail], lon[head], lat[head])
at_at_head = geo.azimuth(lon[head], lat[head], lon[tail], lat[tail])
# Compute delta angles for azimuth(tail -> head) vs azimuth(tail -> dest)
# and for azimuth(head -> tail) vs azimuth(head -> origin)
delta_tail = np.abs((az_at_tail - az_to_dest[tail] + 180.0) % 360.0 - 180.0)
delta_head = np.abs((at_at_head - az_to_origin[head] + 180.0) % 360.0 - 180.0)
keep = (delta_tail <= max_angle_deg) & (delta_head <= max_angle_deg)
edges = np.column_stack([tail[keep], head[keep]])
edge_dist = geo.haversine(lon[tail[keep]], lat[tail[keep]], lon[head[keep]], lat[head[keep]])
return edges, edge_dist