"""Boundaries sampled from a fitted Fourier descriptor model.
The model is a joint PCA Gaussian over the Fourier descriptors of canonicalized
polygons plus a shifted log-normal over their vertex counts. A model fitted to
EMODnet offshore wind-farm polygons ships with the package. Fitting is outside
the package, and models in the older per-frequency format remain loadable.
"""
from __future__ import annotations
import json
from dataclasses import dataclass
from functools import cache
from pathlib import Path
import numpy as np
from shapely.geometry import Polygon
from shapely.geometry.base import BaseGeometry
from .._random import as_rng
from .._validation import _validate_finite, _validate_integer, _validate_nonnegative
from .base import (
Boundary,
Point2D,
_optional_area,
_oriented,
_resolve_center,
)
__all__ = ["FourierPolygon"]
_DEFAULT_MODEL_PATH = Path(__file__).parents[1] / "data" / "shape_model.npz"
@dataclass(repr=False)
class _ShapeModel:
"""Fitted generative model consumed by :class:`FourierPolygon`.
Holds a joint PCA Gaussian over Fourier descriptors (``Z_mean``,
``Z_components``, ``score_chol``) and a shifted log-normal over vertex
counts (the ``v_*`` scalars). Older models store per-frequency factors in
``Z_chol`` instead. ``K`` equals ``Z_mean.shape[0]`` and is not stored.
"""
Z_mean: np.ndarray # (K,) complex per-bin mean of Fourier coefficients
Z_chol: np.ndarray | None # legacy per-bin (Re, Im) covariance factor
K: int # Fourier truncation length (equals Z_mean.shape[0])
v_shift: float # rigid shift, log-normal is fit on log(v - v_shift)
v_log_mu: float # mean of log(v - v_shift) over the trimmed population
v_log_sigma: float # std of log(v - v_shift) over the trimmed population
v_min: int # inclusive lower bound on sampled vertex counts
v_max: int # inclusive upper bound on sampled vertex counts
Z_components: np.ndarray | None = None # (P, K) complex joint PCA loadings
score_chol: np.ndarray | None = None # (P, P) PCA-score covariance factor
metadata: dict | None = None
def __repr__(self) -> str:
# The arrays run to thousands of numbers, too many to print.
return "<_ShapeModel>"
def save(self, path: str | Path) -> Path:
"""Write the model to ``path`` as an ``.npz`` archive and return its path."""
path = Path(path).with_suffix(".npz")
arrays: dict[str, np.ndarray] = {
"schema_version": np.asarray(2 if self.Z_components is not None else 1),
"Z_mean": self.Z_mean,
"v_shift": np.asarray(self.v_shift),
"v_log_mu": np.asarray(self.v_log_mu),
"v_log_sigma": np.asarray(self.v_log_sigma),
"v_min": np.asarray(self.v_min),
"v_max": np.asarray(self.v_max),
"metadata_json": np.asarray(
json.dumps(self.metadata or {}, sort_keys=True)
),
}
if self.Z_components is not None:
assert self.score_chol is not None
arrays["Z_components"] = self.Z_components
arrays["score_chol"] = self.score_chol
else:
assert self.Z_chol is not None
arrays["Z_chol"] = self.Z_chol
# NumPy typing workaround. All archive values are arrays.
np.savez(path, **arrays) # ty: ignore[invalid-argument-type]
return path
@classmethod
def load(cls, path: str | Path) -> "_ShapeModel":
"""Load a model written by :meth:`save`."""
with np.load(Path(path).with_suffix(".npz")) as f:
Z_mean = f["Z_mean"]
Z_chol = f["Z_chol"] if "Z_chol" in f else None
Z_components = f["Z_components"] if "Z_components" in f else None
score_chol = f["score_chol"] if "score_chol" in f else None
v_shift = float(f["v_shift"])
v_log_mu = float(f["v_log_mu"])
v_log_sigma = float(f["v_log_sigma"])
v_min = int(f["v_min"])
v_max = int(f["v_max"])
metadata = {}
if "metadata_json" in f:
metadata = json.loads(str(f["metadata_json"].item()))
return cls(
Z_mean=Z_mean,
Z_chol=Z_chol,
K=int(Z_mean.shape[0]),
v_shift=v_shift,
v_log_mu=v_log_mu,
v_log_sigma=v_log_sigma,
v_min=v_min,
v_max=v_max,
Z_components=Z_components,
score_chol=score_chol,
metadata=metadata,
)
@cache
def _load_default_model() -> _ShapeModel:
"""Load the packaged model once. Its arrays are shared, so read-only."""
model = _ShapeModel.load(_DEFAULT_MODEL_PATH)
for array in (model.Z_mean, model.Z_chol, model.Z_components, model.score_chol):
if array is not None:
array.flags.writeable = False
return model
def _simplify_to_v(
boundary: np.ndarray, target_v: int, max_iterations: int = 28
) -> Polygon | None:
"""Douglas-Peucker-simplify ``boundary`` to about ``target_v`` vertices.
Bisects the tolerance. The result may miss ``target_v`` by one. Returns the
closest valid polygon, or ``None`` if every step was degenerate.
"""
poly = Polygon(boundary)
if not poly.is_valid or poly.area <= 0:
return None
diam = float(np.linalg.norm(np.ptp(boundary, axis=0)))
lo, hi = 0.0, diam
best: Polygon | None = None
best_diff = float("inf")
for _ in range(max_iterations):
mid = 0.5 * (lo + hi)
s = poly.simplify(mid, preserve_topology=True)
if not isinstance(s, Polygon) or not s.is_valid or s.area <= 0:
hi = mid
continue
v = len(s.exterior.coords) - 1
diff = abs(v - target_v)
if diff < best_diff:
best, best_diff = s, diff
if v == target_v:
return s
if v > target_v:
lo = mid
else:
hi = mid
return best
def _sample_vertex_count(model: _ShapeModel, rng: np.random.Generator) -> int:
"""Draw a vertex count from the fitted shifted log-normal.
Rejection-samples in batches of 64 until a count lies within
``[v_min, v_max]``.
"""
while True:
y = np.exp(rng.normal(model.v_log_mu, model.v_log_sigma, size=64))
v = np.rint(y + model.v_shift).astype(int)
ok = v[(v >= model.v_min) & (v <= model.v_max)]
if len(ok):
return int(ok[0])
[docs]
class FourierPolygon(Boundary):
"""Polygon sampled from a fitted Fourier descriptor model.
Each attempt draws PCA scores scaled by ``temperature``, inverse-FFTs the
coefficients to a smooth curve, and simplifies it to a vertex count drawn
from the model. Invalid results are retried up to ``max_attempts`` times,
then a ``RuntimeError`` is raised. ``model`` is a model or ``.npz`` path and
defaults to the packaged model. The shape is scaled to ``area`` and moved
to ``center``. A heading is drawn per sample unless ``rotation`` pins it in
degrees.
"""
# The model as passed, so the packaged default prints as None.
_repr_attributes = {**Boundary._repr_attributes, "model": "_model_argument"}
def __init__(
self,
model: _ShapeModel | str | Path | None = None,
*,
temperature: float = 2.0,
max_attempts: int = 128,
area: float | None = None,
center: Point2D | None = None,
rotation: float | None = None,
seed: int | np.random.Generator | None = None,
):
self.max_attempts = _validate_integer("max_attempts", max_attempts)
self._model_argument = model
if model is None:
model = _load_default_model()
if isinstance(model, (str, Path)):
model = _ShapeModel.load(model)
self.model = model
self.temperature = _validate_nonnegative("temperature", temperature)
self.target_area = _optional_area(area)
self.center, self._positioned = _resolve_center(center)
self.rotation = (
None if rotation is None else _validate_finite("rotation", rotation)
)
self.seed = self._resolve_own_seed(seed)
def _build_geometry(self) -> BaseGeometry:
rng = as_rng(self.seed)
model = self.model
for _ in range(self.max_attempts):
if model.Z_components is not None:
n_scores = model.Z_components.shape[0]
u = rng.standard_normal(n_scores) * self.temperature
scores = model.score_chol @ u
# A complex vector-matrix product through BLAS spins every BLAS
# thread for a few microseconds of work, so einsum forms it.
Z = model.Z_mean + np.einsum("i,ij->j", scores, model.Z_components)
else:
assert model.Z_chol is not None
u = rng.standard_normal((model.K, 2)) * self.temperature
delta = np.einsum("kij,kj->ki", model.Z_chol, u)
Z = model.Z_mean + delta[:, 0] + 1j * delta[:, 1]
Z[0] = 0 # canonicalization centers each polygon, so the DC bin is 0
z = np.fft.ifft(Z)
boundary = np.column_stack([z.real, z.imag])
v = _sample_vertex_count(model, rng)
result = _simplify_to_v(boundary, v)
if result is not None and result.is_valid and result.area > 1e-6:
return _oriented(result, self.rotation, rng)
raise RuntimeError(
f"could not generate a valid shape in {self.max_attempts} attempts"
)
