"""Token sets for field surrogates.
Field transformers take a set of rows. These helpers turn prophys structures
into such rows. A raster becomes one row per cell with its world coordinates
and values, a point list one row per point, geometry one row per segment or
edge, and `sample_set` turns the Monte Carlo sample of any upstream quantity
into a set, so a surrogate can map a whole distribution to an output.
"""
from __future__ import annotations
from typing import Any
import jax
import jax.numpy as jnp
from ..diagnostics import raise_error
from ..symbolic.expr import Expr, _collect, op, wrap
__all__ = ["field_tokens", "point_tokens", "geometry_tokens", "sample_set", "SampleSet"]
[docs]
def field_tokens(field, channels: int = 1, coords: bool = True) -> Expr:
"""Token rows of a raster field: cell centre coordinates followed by the cell values.
Args:
field: A raster structure with ``values`` ``[..., ny, nx]`` (or
``[..., ny, nx, channels]``), ``origin`` and ``cell`` size.
channels: Number of value channels per cell.
coords: Whether to prepend the world coordinates ``x, y`` of each
cell centre.
Returns:
An expression of shape ``[..., ny * nx, 2 + channels]`` (or
``[..., ny * nx, channels]`` without coordinates). Leading batch
axes of the values are kept.
Raises:
~prophys.diagnostics.ProphysError: ``NN123`` when the field's
channel count differs from `channels`.
"""
def _tok(v):
v = jnp.asarray(v)
if channels == 1:
ny, nx = v.shape[-2:]
vals = v.reshape(v.shape[:-2] + (ny * nx, 1))
else:
ny, nx, c = v.shape[-3:]
if c != channels:
raise_error("NN123", f"field has {c} channels, {channels} declared")
vals = v.reshape(v.shape[:-3] + (ny * nx, c))
if not coords:
return vals
yy, xx = jnp.meshgrid(jnp.arange(ny), jnp.arange(nx), indexing="ij")
x = field.origin[0] + (xx.reshape(-1) + 0.5) * field.cell[1]
y = field.origin[1] + (yy.reshape(-1) + 0.5) * field.cell[0]
xy = jnp.stack([x, y], axis=-1).astype(vals.dtype)
xy = jnp.broadcast_to(xy, vals.shape[:-1] + (2,))
return jnp.concatenate([xy, vals], axis=-1)
return op("field_tokens", _tok, field.values)
[docs]
def point_tokens(points, features: Any = None) -> Expr:
"""Token rows of a point set: coordinates followed by optional per point features.
Args:
points: A point structure with ``coords`` ``[n, 2]``, or an
expression or array of coordinates.
features: Optional features ``[n, f]`` or ``[n]``.
Returns:
An expression of shape ``[n, 2 + f]`` (the coordinates alone when
no features are given).
"""
coords = points.coords if hasattr(points, "coords") else wrap(points)
if features is None:
return coords
def _tok(c, f):
f = jnp.asarray(f)
if f.ndim == c.ndim - 1:
f = f[..., None]
return jnp.concatenate([c, jnp.broadcast_to(f, c.shape[:-1] + f.shape[-1:]).astype(c.dtype)], axis=-1)
return op("point_tokens", _tok, coords, features)
[docs]
def geometry_tokens(geometry) -> Expr:
"""Token rows of a geometry, one row ``[x0, y0, x1, y1]`` per segment or edge.
Works for a line list, a polyline, a polygon (edges including the
closing one) or a polygon list.
Returns:
An expression of shape ``[n_segments, 4]``.
Raises:
~prophys.diagnostics.ProphysError: ``NN123`` for any other structure.
"""
if hasattr(geometry, "segments"):
src, closed = geometry.segments, None
elif hasattr(geometry, "vertices"):
src = geometry.vertices
closed = type(geometry).__name__.startswith("Polygon")
else:
raise_error(
"NN123",
f"cannot make tokens from a {type(geometry).__name__}",
hint="pass a LineList, Polyline, Polygon or PolygonList",
)
def _tok(v):
v = jnp.asarray(v)
if closed is None:
return v.reshape(v.shape[:-2] + (4,)) if v.shape[-2:] == (2, 2) else v
nxt = jnp.roll(v, -1, axis=-2) if closed else v[..., 1:, :]
cur = v if closed else v[..., :-1, :]
seg = jnp.concatenate([cur, nxt], axis=-1)
return seg.reshape((-1, 4)) if seg.ndim == 3 else seg
return op("geometry_tokens", _tok, src)
[docs]
class SampleSet(Expr):
"""Expression node holding a fresh Monte Carlo sample ``[n, k]`` of an expression.
The sample is drawn with fixed random numbers (common random numbers,
set by `seed`), so it is smooth in the model parameters. Every random
variable `expr` depends on is drawn anew, independently of the outer
Monte Carlo draw of the model. Create it with :func:`sample_set`.
Args:
expr: The expression to sample.
n: Sample size (rows).
seed: Seed of the fixed random numbers.
"""
def __init__(self, expr: Any, n: int, seed: int = 0):
from ..domain.random_variable import RandomVariable
self.expr = wrap(expr)
self.n = int(n)
self.seed = int(seed)
found: dict = {}
_collect(self.expr, RandomVariable, found)
self.rvs = found
def _children(self):
# only the parameters are visible, the random variables are drawn here
# and must not join the outer Monte Carlo draw of the model
return tuple(self.expr.params().values())
def _evaluate(self, env):
key = jax.random.PRNGKey(self.seed)
keys = jax.random.split(key, max(len(self.rvs), 1))
env2 = {k: v for k, v in env.items() if k not in self.rvs}
draws = {
name: rv.distribution.sample(k, (self.n,)).evaluate(env) for k, (name, rv) in zip(keys, self.rvs.items())
}
env2["__rv_samples__"] = {**dict(env.get("__rv_samples__", {})), **draws}
v = jnp.asarray(self.expr._evaluate(env2))
v = jnp.broadcast_to(v, (self.n,) + v.shape[1:]) if v.ndim else jnp.broadcast_to(v, (self.n,))
return v.reshape(self.n, -1)
def __repr__(self):
return f"SampleSet({self.n})"
[docs]
def sample_set(expr: Any, n: int, seed: int = 0) -> SampleSet:
"""Turn the Monte Carlo sample of an upstream quantity into a token set.
Returns a :class:`SampleSet` whose value is ``[n, k]``: `n` draws of
`expr` (flattened to `k` columns), drawn independently of the model's
own Monte Carlo draw with fixed random numbers so the set is smooth in
the parameters. Feeding it to a :class:`~prophys.nn.FieldSurrogate`
makes the surrogate a map from a distribution to an output.
Args:
expr: The expression to sample, typically depending on random variables.
n: Sample size.
seed: Seed of the fixed random numbers.
Example:
>>> tokens = nn.sample_set(load, 256) # [256, 1]
>>> damage = nn.FieldSurrogate("damage", rows=256, cols=1)(tokens)
"""
return SampleSet(expr, n, seed)
