GeoStack operators: geoops and xrops on one kernel model

pipekit-array        tier 0: adapters, KernelOp, dim-level functions
     │               runs any kernel on NumPy, JAX and xarray
     │   [xarray] extra ──▶ xarray
     ▼
geoops               tier 1: every kernel operator, written once;
     │               geometry on carriers; no xarray in its base
     ▼
xrops                tier 2: native-xarray work (sel by value, groupby,
                     many files, Dataset/DataTree), xrops.lift(geoops op),
                     state physics; imports geoops and xarray

# monthly NDVI climatology over a lazy cube: a geoops kernel, then a by-value reduction
# clim(m, y, x) = (1 / |Tₘ|) ∑_{t ∈ Tₘ} NDVI(t, y, x), Tₘ the dates in month m
ndvi: Operator = geoops.NDVI()
climatology: Operator = xrops.Climatology(op=ndvi, freq="month")
clim: DataArray["month y x"] = climatology(cube)

# the same as a pipekit pipeline
monthly: Operator = xrops.Climatology(freq="month")
pipe: pipekit.Pipeline = ndvi | monthly
clim_2: DataArray["month y x"] = pipe(cube)

# finitevolx/_src/array.py — arrays in, arrays out, no carrier, no CRS
class CGrid(eqx.Module):
    dx_u: Float[Array, "H W"]  # zonal spacing at u points, m
    dy_v: Float[Array, "H W"]  # meridional spacing at v points, m
    area_t: Float[Array, "H W"]  # cell area at T points, m²
    mask_t: Bool[Array, "H W"]  # ocean = True


def divergence(
    u: Float[Array, "... H W"],  # zonal velocity on u points
    v: Float[Array, "... H W"],  # meridional velocity on v points
    grid: CGrid,
) -> Float[Array, "... H W"]:
    # ∇·𝐮 ≈ [δₓ(u Δy) + δᵧ(v Δx)] / A, zero on land
    ...

# somax/geo/operators.py — behind somax[geo]; GeoModels may import GeoStack and GeoNumerics
class Divergence(Operator):
    def __call__(
        self, u: GeoArray["... y x"], v: GeoArray["... y x"]
    ) -> GeoArray["... y x"]:
        # derived geometry, once per call: transform + CRS → spacings and areas (pyproj lives here)
        grid: CGrid = cgrid_from_transform(u.transform, u.crs, mask=u.valid)
        # the L1 call on raw arrays
        div: Float[Array, "... H W"] = divergence(u.data, v.data, grid)
        # rewrapped on the T-point grid
        return u.rewrap(div, dims=u.dims)

# somax/geo/simulate.py — linear shallow water, η the SSH anomaly
def simulate(
    eta0: GeoMap,  # GeoArray["y x"], initial state
    times: Int[np.ndarray, "T"],  # output times tₖ, s after eta0's time; concrete
    depth: float,  # H, m
    dt: float,  # Δt, s
    mask: Bool[Array, "H W"] | None = None,  # ocean cells; None means eta0.valid
) -> GeoImageSeries:  # GeoArray["time y x"], one frame per requested time
    # entry, before tracing: output k is reached after nₖ = tₖ / Δt steps; n = 0 is η₀
    n_out: Int[np.ndarray, "T"] = np.rint(times / dt).astype(int)
    off_grid: bool = bool(np.any(n_out * dt != times))
    gaps: Int[np.ndarray, "T-1"] = np.diff(n_out)
    unsorted: bool = bool(np.any(gaps < 0))
    if off_grid or unsorted or n_out[0] < 0:
        raise ValueError("times must be sorted, non-negative multiples of dt")
    n_total: int = int(n_out[-1])
    # entry: carrier → arrays and derived geometry, once; the domain is the mask the
    # caller passes, so finite placeholders on land never turn land into ocean
    ocean: Bool[Array, "H W"] = eta0.valid if mask is None else mask
    grid: CGrid = cgrid_from_transform(eta0.transform, eta0.crs, mask=ocean)
    eta_init: Float[Array, "H W"] = eta0.data
    zeros: Float[Array, "H W"] = jnp.zeros_like(eta_init)
    # (η, u, v), each (H, W)
    state0: tuple[Float[Array, "H W"], Float[Array, "H W"], Float[Array, "H W"]] = (
        eta_init,
        zeros,
        zeros,
    )

    def step(
        state: tuple[Float[Array, "H W"], ...], _: None
    ) -> tuple[tuple[Float[Array, "H W"], ...], Float[Array, "H W"]]:
        eta, u, v = state
        # ∂η/∂t = −H ∇·𝐮
        div: Float[Array, "H W"] = divergence(u, v, grid)
        eta_next: Float[Array, "H W"] = eta - dt * depth * div
        # ∇η = (∂η/∂y, ∂η/∂x)
        gy: Float[Array, "H W"]
        gx: Float[Array, "H W"]
        gy, gx = gradient(eta_next, grid)
        # ∂𝐮/∂t = −g ∇η, g = 9.81 m s⁻²
        u_next: Float[Array, "H W"] = u - dt * G * gx
        v_next: Float[Array, "H W"] = v - dt * G * gy
        return (eta_next, u_next, v_next), eta_next

    # one trace, no carrier inside the loop: η after each of the S = n_total steps, (S, H, W)
    etas: Float[Array, "S H W"]
    _, etas = jax.lax.scan(step, state0, None, length=n_total)
    # index n is the state after n steps, so η₀ goes first: (S + 1, H, W)
    first: Float[Array, "1 H W"] = eta_init[None]
    traj: Float[Array, "S+1 H W"] = jnp.concatenate([first, etas])
    # the frames at the requested times: (T, H, W)
    frames: Float[Array, "T H W"] = traj[n_out]
    # exit: arrays → a series on eta0's grid, stamped with the requested times;
    # TimeAxis stamps are int64 nanoseconds, so offsets in s become ns first
    offsets_ns: Int[np.ndarray, "T"] = times.astype(np.int64) * 1_000_000_000
    stamps: Int[np.ndarray, "T"] = eta0.time + offsets_ns
    return GeoImageSeries.from_frames(frames, like=eta0, time=stamps)


# vardax/geo/assimilate.py — 4D-Var over a window of along-track SLA.
# vardax is GeoML and may not import GeoModels, so the forward model M is an argument.
# M: an initial map, output times and the ocean mask → the trajectory at those times
ForwardModel = Callable[
    [GeoMap, Int[np.ndarray, "T"], Bool[Array, "H W"]], GeoImageSeries
]


def assimilate(
    background: GeoMap,  # xᵇ, GeoArray["y x"]
    obs: GeoTracks,  # y, GeoPoints["obs"] padded to N, with a valid mask
    times: Int[np.ndarray, "T"],  # the window's output times, s after xᵇ's time
    forward: ForwardModel,  # M, from the caller
    sigma_b: float,  # background error std, m
    sigma_o: float,  # observation error std, m
) -> GeoImageSeries:

    # m_b = 1 where xᵇ is valid (ocean); land and nodata hold 0 and never enter J
    valid_b: Bool[Array, "H W"] = background.valid
    xb: Float[Array, "H W"] = jnp.where(valid_b, background.data, 0.0)

    def cost(x0: Float[Array, "H W"]) -> Float[Array, ""]:
        # m_b ⊙ x₀: land cells are held at 0, so their gradient is 0 and they never move
        x0_ocean: Float[Array, "H W"] = jnp.where(valid_b, x0, 0.0)
        # M(x₀) on the background's domain: the mask is passed, not re-derived from data
        x0_map: GeoMap = background.with_data(x0_ocean)
        traj: GeoImageSeries = forward(x0_map, times, valid_b)
        # H(M(x₀)): sampled at the observations in space and time
        predicted: GeoTracks = geoarray.sample(traj, obs)
        # d = m ⊙ (y − H(M(x₀))); padded slots give 0, never NaN
        resid: Float[Array, "N"] = obs.data - predicted.data
        d: Float[Array, "N"] = jnp.where(obs.valid, resid, 0.0)
        # m_b ⊙ (x₀ − xᵇ): finite everywhere, zero off the valid cells
        dx_raw: Float[Array, "H W"] = x0_ocean - xb
        dx: Float[Array, "H W"] = jnp.where(valid_b, dx_raw, 0.0)
        # J = ½‖m_b ⊙ (x₀ − xᵇ)‖²/σ_b² + ½‖d‖²/σ_o²
        jb: Float[Array, ""] = 0.5 * einx.dot("h w, h w ->", dx, dx) / sigma_b**2
        jo: Float[Array, ""] = 0.5 * einx.dot("n, n ->", d, d) / sigma_o**2
        return jb + jo

    # xₐ = argmin J, by gradient descent on ∇J
    cost_jit: Callable = jax.jit(cost)
    x_a: Float[Array, "H W"] = minimize(cost_jit, xb)
    # the analysis trajectory, as a readable series
    analysis_map: GeoMap = background.with_data(x_a)
    traj_a: GeoImageSeries = forward(analysis_map, times, valid_b)

    # land back to NaN in the output (broadcast over time), so its validity matches xᵇ
    def restore_land(frames: Float[Array, "T H W"]) -> Float[Array, "T H W"]:
        return jnp.where(valid_b, frames, jnp.nan)

    return traj_a.map(restore_land)


# the caller binds the model, in GeoModels or a script: M(x₀) = simulate(x₀, times; H, Δt)
forward: ForwardModel = functools.partial(somax.geo.simulate, depth=4_000.0, dt=600.0)
# daily output over a 7-day window, t = 0 included: (8,)
times: Int[np.ndarray, "T"] = np.arange(0, 8 * 86_400, 86_400)
series: GeoImageSeries = vardax.geo.assimilate(
    background, obs, times, forward, sigma_b=0.05, sigma_o=0.03
)

pipekit
   │
   ▼
pipekit-array ── [xarray] adapter ──▶ xarray
   │                                  ▲
   ▼                                  │
geoarray                              │
   │                                  │
   ▼                                  │
geoops   (never imports xarray        │
   │      or xrops)                   │
   ▼                                  │
xrops ─────────── imports ─────────────┘
   reuses geoops directly and through xrops.lift

# a per-item operator, written once for geoarray carriers: GeoArray["band y x"] → GeoArray["y x"]
detect: Operator = geoops.plume.DetectPlumes(threshold=tau)

# lifted over time by xrops: each slice becomes a GeoImage, runs, and is wrapped back
per_date: Operator = xrops.lift(detect, dim="time")

# lazy, one task per dask chunk: cube is DataArray["time band y x"]
masks: DataArray["time y x"] = per_date(cube)

# L1 kernel in geoops/indices/_src/array.py: "(),()->()"
def normalized_difference(
    a: Float[Array, "*s"], b: Float[Array, "*s"], eps: float
) -> Float[Array, "*s"]:
    # ND = (a − b) / (a + b + ε)
    return (a - b) / (a + b + eps)


# L2 operator: two labels of one dim, then the kernel
class NDVI(Operator):
    backends: tuple[str, ...] = ("numpy", "jax", "xarray")

    def __init__(
        self, nir: str = "B8", red: str = "B4", dim: str = "band", eps: float = 1e-10
    ) -> None:
        self.nir: str = nir
        self.red: str = red
        self.dim: str = dim
        self.eps: float = eps

    def _apply(self, scene: GeoArray["band y x"]) -> GeoArray["y x"]:
        # exact label lookups by name, any carrier
        nir_sel: dict[str, str] = {self.dim: self.nir}
        red_sel: dict[str, str] = {self.dim: self.red}
        nir: GeoArray["y x"] = pipekit_array.sel(scene, **nir_sel)
        red: GeoArray["y x"] = pipekit_array.sel(scene, **red_sel)
        # the kernel, element-wise, through KernelOp
        nd: KernelOp = KernelOp(
            normalized_difference, "(),()->()", core={}, eps=self.eps
        )
        # NDVI = (ρ_NIR − ρ_red) / (ρ_NIR + ρ_red + ε)
        return nd(nir, red)

# L2 geometry operator, array math: pyproj in the operator, a primitive for the resampling
class Reproject(Operator):
    backends: tuple[str, ...] = ("numpy", "jax")

    def __init__(self, crs: str, resolution: float) -> None:
        self.crs: str = crs
        self.resolution: float = resolution

    def _apply(self, scene: GeoArray["band y x"]) -> GeoArray["band y x"]:
        # the target grid: transform and shape in the new CRS
        target: GridSpec = geoarray.target_grid(
            scene, crs=self.crs, resolution=self.resolution
        )
        # derived geometry: source pixel coordinates (r, c) of every target pixel, by pyproj
        rows: Float[Array, "H W"]
        cols: Float[Array, "H W"]
        rows, cols = geoarray.source_coords(scene, target)
        # L1: q(y′, x′) = ∑ wᵢ q(rᵢ, cᵢ), bilinear weights wᵢ, per band
        warped: Float[Array, "C H W"] = map_coordinates(scene.data, rows, cols, order=1)
        # rewrap on the target grid, same dims
        return GeoArray.from_grid(warped, target, dims=scene.dims)

# L1 kernel in xrops/ocn/_src/array.py: consumes the level axis, "(l),(l)->()"
def mixed_layer_depth(
    sigma0: Float[Array, "... L"],  # σ₀, kg m⁻³, L levels last
    depth: Float[Array, "... L"],  # z, m
    threshold: float,  # kg m⁻³
) -> Float[Array, "..."]:
    # numpy or jax.numpy, chosen by the input
    xp: ModuleType = array_namespace(sigma0)
    # σ₀(z_ref), z_ref the shallowest level
    ref: Float[Array, "... 1"] = sigma0[..., :1]
    # Δσ(z) = σ₀(z) − σ₀(z_ref)
    delta: Float[Array, "... L"] = sigma0 - ref
    # 𝟙[Δσ(z) > threshold]
    below: Bool[Array, "... L"] = delta > threshold
    # the first level where Δσ exceeds the threshold
    first: Int[Array, "..."] = xp.argmax(below, axis=-1)
    first_l: Int[Array, "... 1"] = first[..., None]
    # MLD = z at that level
    picked: Float[Array, "... 1"] = xp.take_along_axis(depth, first_l, axis=-1)
    mld: Float[Array, "..."] = picked[..., 0]
    # argmax of an all-False row is 0: no crossing means no MLD, so NaN
    crossed: Bool[Array, "..."] = einx.any("... [l]", below)
    return xp.where(crossed, mld, xp.nan)


# L2: the level dim named, the threshold bound
mld_op: KernelOp = KernelOp(
    mixed_layer_depth, "(l),(l)->()", core={"l": "level"}, threshold=0.03
)

# a GLORYS-like cube, lazy: (time, level, y, x) → (time, y, x)
mld_cube: DataArray["time y x"] = mld_op(sigma0_cube, depth_cube)

# one gridded profile on a geoarray carrier: GeoVolume (level, y, x) → GeoMap (y, x)
mld_map: GeoArray["y x"] = mld_op(sigma0_volume, depth_volume)

# geoops/geom/_src/array.py — L1, spacings as arrays
def central_gradient(
    f: Float[Array, "... H W"],
    dy: Float[Array, "H W"],
    dx: Float[Array, "H W"],
) -> Float[Array, "2 ... H W"]:
    # numpy or jax.numpy, chosen by the input
    xp: ModuleType = array_namespace(f)
    # neighbours along y and x
    f_n: Float[Array, "... H W"] = xp.roll(f, -1, -2)
    f_s: Float[Array, "... H W"] = xp.roll(f, 1, -2)
    f_e: Float[Array, "... H W"] = xp.roll(f, -1, -1)
    f_w: Float[Array, "... H W"] = xp.roll(f, 1, -1)
    # ∂f/∂y ≈ (fᵢ₊₁ − fᵢ₋₁) / 2Δy
    fy: Float[Array, "... H W"] = (f_n - f_s) / (2 * dy)
    # ∂f/∂x ≈ (fⱼ₊₁ − fⱼ₋₁) / 2Δx
    fx: Float[Array, "... H W"] = (f_e - f_w) / (2 * dx)
    # (∂f/∂y, ∂f/∂x) on a new leading axis
    return einx.rearrange("... h w, ... h w -> (1 + 1) ... h w", fy, fx)


# geoops: spacing from the transform (a, e), or geodesic distances for a geographic CRS
grad_op: Operator = geoops.geom.Gradient()
grad_map: GeoArray["component y x"] = grad_op(ssh_map)  # ssh_map: GeoArray["y x"]

# xrops: spacing from coordinate values, spherical metrics on lat/lon; the same geoops kernel
grad_xr: Operator = xrops.Gradient(dims=("lat", "lon"))
grad_cube: DataArray["component time lat lon"] = grad_xr(
    ssh
)  # ssh: DataArray["time lat lon"]

1 pipekit, pipekit-array: adapter, Signature, functions, KernelOp
        │
        ▼
2 geoarray: adapter (geoarray page)
        │
        ▼
3 geoops: slim base install (blocks 5)
  axis → dim, KernelOps, backends,
  gradient kernel, no [xarray] extra
        │
        ▼
4 dependency graph: xrtoolz to tier 2
        │
        ▼
5 xrops (package layout page): rename, fold xrcore,
  Signature re-export, import geoops kernels,
  xrops.lift, xrops.Gradient, state-physics KernelOps