From 2efb2f8334f780d69dff6a5ca05a5b2af6f049d9 Mon Sep 17 00:00:00 2001 From: Nicholas Ehsan Roy Date: Wed, 7 Oct 2026 04:38:33 +0200 Subject: [PATCH 1/4] wip(tests): numerical coupling reference, nonlinear search cells, interface-reading twins Work in progress: the float64 pass-map reference, the nonlinear cells of the coupling search and the twins whose interface reading is the edge-mapped graph's. Tests only. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013UkCde7g23gTziUjYAvnKD [skip ci] --- tests/property/coupling_reference.py | 539 ++++++++++++++ tests/property/geometry_graphs.py | 192 +++++ .../test_coupling_nonlinear_search.py | 685 ++++++++++++++++++ .../property/test_coupling_targeted_search.py | 77 +- .../property/test_interface_reading_twins.py | 305 ++++++++ 5 files changed, 1765 insertions(+), 33 deletions(-) create mode 100644 tests/property/coupling_reference.py create mode 100644 tests/property/test_coupling_nonlinear_search.py create mode 100644 tests/property/test_interface_reading_twins.py diff --git a/tests/property/coupling_reference.py b/tests/property/coupling_reference.py new file mode 100644 index 00000000..86b9a80b --- /dev/null +++ b/tests/property/coupling_reference.py @@ -0,0 +1,539 @@ +"""A numerical float64 reference for one coupling group, with no closed form. + +``test_coupling_targeted_search.py`` scores ``coupling_diagnostics()`` +against :class:`~tests.property.coupled_topologies.LinearModel`: the fixed +point and the pass's Jacobian of a group of *linear* relays, in closed +form. A group with a nonlinear node, a dtype-widening transform or a +mapping that reads a moving geometry has no closed form. This module +computes the same answers numerically, for any small group (up to a few +dozen scalars) built from any nodes and edges. + +**How the pass map is obtained** (:class:`PassReference`). The caller +builds an *x64 twin* of the graph under test -- the same nodes, edges, +schedule and sub-cycling, every float in float64 -- whose group runs +``max_iterations=1`` with no acceleration and ``predictor="linear"`` +(:func:`twin_knobs`). One step of that twin is one coupling pass and is +differentiated straight through (there is no fixed point to apply the +implicit function theorem at). The pass starts from the predictor's +extrapolation ``2 pred_0 - pred_1`` of the last two results, which the +step keeps in three ``_meta`` slots of the state; writing the iterate +``x`` into both and ``2`` into the count makes the step compute + + ``P(x; pre) = one pass of the group from the iterate x, each member + integrating from the pre-step state pre`` + +with *pre* the members' fields of the state handed in. The map is the +graph's own compiled step (``gm._raw_step_fn``, the callable ``step()`` +runs) called on a state: nothing of the estimator is called, and the only +thing written that ``set_node_state`` could not write is those three +slots. Their layout is not assumed either: :meth:`PassReference.of` +steps the twin once and finds the order of nodes whose flattened floating +fields equal the stored ``pred_0`` bit for bit. + +From ``P``, in float64 (``jax.jacfwd`` for every derivative, so every +dependence the pass has on the iterate is in it: transforms, mappings, +and a geometry field where a member holds one): + +* :meth:`PassReference.fixed_point`: Picard passes, then Newton on + ``P(x) - x`` with the dense Jacobian, to a residual of a few float64 + ``eps`` of the fields; the residual history is returned as the evidence; +* :meth:`PassReference.jacobian`: the dense ``dP/dx`` at any iterate, and + :func:`radius` of it; +* :meth:`PassReference.distance`: the distance from a returned iterate + to the fixed point in the norm ``spectral_error_bound`` documents (each + field of the norm's reading over its own largest magnitude at the + returned state; a root mean square under the relative norms); +* :meth:`PassReference.gradient_error`: the relative error of the + implicit derivative ``(I - dP/dx)^{-1} dP/dc`` taken at the returned + iterate against the same dense solve at the fixed point, for every + scalar constant ``c`` (the form ``gradient_relative_error_bound`` is + read in) or for a scalar loss ``w . x``; +* :meth:`PassReference.residual`: the exact residual ``P(x) - x`` of a + returned state in the group's norm; +* :meth:`PassReference.nonlinearity`: ``||(I - J(x))^{-1} (J_mean - + J(x))||`` with ``J_mean`` the mean Jacobian over the segment from the + fixed point to ``x`` -- the exact factor by which a linear bound taken + at ``x`` can miss on a nonlinear map (see the method). + +**What it catches.** Anything by which the estimator's numbers disagree +with the map the solve iterates: a term of the pass missing from the +estimator's Jacobian (a mapping, a transform, a geometry dependence), a +reading taken at another time level than the pass takes it, a Jacobian +taken at another state than the returned one, a bound in another norm +than the one documented, a float32 analysis that lost the quantity. + +**What it does not catch.** + +* A defect in the pass itself (the solve iterating a wrong map): the + reference inherits it. The closed-form and time-level references are + what hold the pass. +* A difference between the single-pass branch (``max_iterations=1``) and + the pass an iterating group runs. :func:`passes_compose` checks that on + a twin run for several passes; the search's validation runs it. +* A difference between the float32 graph's pass and its float64 twin's + beyond rounding (a node that branches on its dtype). +* A graph with more than one coupling group (the others would run their + own single pass), or a group whose members hold non-float state the + pass map depends on. + +**Cost.** Per twin: the twin's own build, and two compiles (the pass +with its Jacobian, one program; the constants' sensitivities, only where +a gradient is scored), a second or two each for a group of a dozen +scalars. Per example: a few dozen jitted calls, a few milliseconds. +""" + +from __future__ import annotations + +import contextlib +import dataclasses +import itertools +import math +from typing import Callable, Optional, Sequence + +import jax +import jax.numpy as jnp +import numpy as np +from jax.flatten_util import ravel_pytree + +#: The float64 rounding of one evaluation. +EPS64 = float(np.finfo(np.float64).eps) +#: A fixed point is accepted where the pass moves no entry by more than +#: this many float64 ``eps`` of its field's magnitude (one pass of a few +#: dozen operations rounds at a handful; a resolvent of 1e3 leaves the +#: Newton iterate that far from a stationary float). +FIXED_POINT_ULPS = 2.0 ** 12 + + +@contextlib.contextmanager +def x64(): + """Run the body under ``jax_enable_x64``, restoring the setting.""" + prior = jax.config.read("jax_enable_x64") + jax.config.update("jax_enable_x64", True) + try: + yield + finally: + jax.config.update("jax_enable_x64", prior) + + +def twin_knobs(knobs: dict) -> dict: + """The x64 twin's group configuration for a group under *knobs*: the + schedule and the sub-cycling of the group under test (what decides the + pass), one pass, and the predictor the iterate is handed in through.""" + keep = ("iteration_mode", "subcycling", "boundary_interpolation") + return {**{k: knobs[k] for k in keep if k in knobs}, + "max_iterations": 1, "predictor": "linear"} + + +def radius(J: np.ndarray) -> float: + """The spectral radius of *J*.""" + J = np.asarray(J, np.float64) + return float(np.max(np.abs(np.linalg.eigvals(J)))) if J.size else 0.0 + + +@dataclasses.dataclass(frozen=True) +class Norm: + """The norm a report states its bound in. + + ``fields(x)`` returns the norm's reading of a flat iterate as a list + of arrays (JAX or NumPy): a member's field under ``"l2"`` and + ``"mixed"``, what each internal edge delivers under ``"interface"``. + Each is weighted by ``1 / (rtol max|field|)`` at the returned state + (``rtol`` 1 under ``"l2"``); ``rms`` divides the sum of squares by the + number of entries. + """ + + fields: Callable + rtol: float = 1.0 + rms: bool = False + + def weights(self, x_returned, also=None) -> list: + """One weight per field; *also* is a second iterate whose magnitude + counts too (the residual's own weights read both).""" + out = [] + second = self.fields(also) if also is not None else None + for k, f in enumerate(self.fields(x_returned)): + ref = float(np.max(np.abs(np.asarray(f)))) if np.size(f) else 0.0 + if second is not None and np.size(second[k]): + ref = max(ref, float(np.max(np.abs(np.asarray(second[k]))))) + out.append(1.0 / (self.rtol * ref) if ref > 0 else 0.0) + return out + + def of_difference(self, a, b, weights) -> float: + """``||w (fields(a) - fields(b))||``, a root mean square under ``rms``.""" + total, count = 0.0, 0 + for w, fa, fb in zip(weights, self.fields(a), self.fields(b)): + d = np.asarray(fa, np.float64) - np.asarray(fb, np.float64) + total += float(np.sum((w * d) ** 2)) + count += d.size + return math.sqrt(total / count if self.rms and count else total) + + +@dataclasses.dataclass(frozen=True) +class FixedPoint: + """A fixed point and the evidence it is one.""" + + x: np.ndarray + #: ``max |P(x) - x|`` over each field's magnitude, in float64 ``eps``. + ulps: float + converged: bool + #: The same measure after every Picard pass and Newton step taken. + history: tuple + picard: int + newton: int + + +class PassReference: + """The one-pass map of an x64 twin's coupling group and what follows from it. + + Build with :meth:`of`; :meth:`at` binds a pre-step state and the + constants (both ride as arguments of the jitted map, so one twin + serves every draw of its cell). + """ + + def __init__(self, gm, key, layout, constants_of): + self.gm = gm + self.key = key + #: ``((node, field, shape, start, stop), ...)`` of the flat iterate. + self.layout = layout + self.size = layout[-1][4] if layout else 0 + self._constants_of = constants_of + self._slots = tuple(f"coupling_{key}_pred_{s}" for s in ("0", "1", "count")) + self._ext = gm._default_external_inputs() # noqa: SLF001 + self._step = gm._raw_step_fn # noqa: SLF001 + self._jit_both = jax.jit(lambda x, pre, params: ( + self._pass(x, pre, params), jax.jacfwd(self._pass)(x, pre, params))) + self._jit_sens = None + self._pre = None + self._params = None + + # -- construction ------------------------------------------------------- + + @classmethod + def of(cls, gm, key: Optional[str] = None, *, params: Optional[dict] = None, + constants: Optional[Callable] = None) -> "PassReference": + """The reference of compiled x64 twin *gm* (built under + :func:`twin_knobs`) for its group *key* (its only one by default). + + The layout of the predictor's slots is found by stepping the twin + once from the state it holds, with *params* (``gm.params`` by + default): give it a state and constants under which no two members + end the step with equal fields (a drawn example; not the zeros a + fresh graph holds), or the stored vector does not say which member + is where and this raises. + + *constants* maps ``gm.params`` to the sub-tree of constants the + gradient is taken with respect to; by default every floating leaf + of the members' parameters and of the mappings. + """ + with x64(): + groups = {"+".join(sorted(g.nodes)): g for g in gm._coupling_groups} # noqa: SLF001 + assert len(groups) == 1, ( + f"the reference handles a graph with one coupling group, not {sorted(groups)}") + key = next(iter(groups)) if key is None else key + group = groups[key] + assert group.max_iterations == 1 and group.predictor == "linear", ( + "the twin's group runs one pass under the linear predictor (twin_knobs)") + state = gm._state # noqa: SLF001 + slots = tuple(f"coupling_{key}_pred_{s}" for s in ("0", "1", "count")) + assert all(s in state["_meta"] for s in slots), sorted(state["_meta"]) + ext = gm._default_external_inputs() # noqa: SLF001 + after = gm._raw_step_fn(state, ext, gm.params if params is None else params) # noqa: SLF001 + stored = np.asarray(after["_meta"][slots[0]]) + members = sorted(group.nodes) + + def fields_of(name): + return [f for f in sorted(state[name]) + if jnp.issubdtype(jnp.asarray(state[name][f]).dtype, jnp.floating)] + + orders = [tuple(members), tuple(n for n in gm.node_names if n in group.nodes), + tuple(group.nodes)] + if len(members) <= 6: + orders += list(itertools.permutations(members)) + matches = [] + for order in dict.fromkeys(orders): + flat = np.concatenate([np.ravel(np.asarray(after[n][f])) + for n in order for f in fields_of(n)]) + if flat.shape == stored.shape and np.array_equal(flat, stored): + matches.append(order) + assert matches, ( + f"no order of {members} flattens the stepped state to the stored pred_0") + assert len(matches) == 1, ( + f"the stored pred_0 does not determine the layout (it is the stepped state " + f"under each of {matches[:4]}): give the twin a state and constants under " + f"which the members' fields differ after a step") + layout, at = [], 0 + for n in matches[0]: + for f in fields_of(n): + shape = tuple(np.shape(state[n][f])) + size = int(np.prod(shape, dtype=np.int64)) + assert jnp.asarray(state[n][f]).dtype == jnp.float64, ( + f"{n}.{f} is {jnp.asarray(state[n][f]).dtype} in the twin, not float64") + layout.append((n, f, shape, at, at + size)) + at += size + + def default_constants(params): + def floats(tree): + return {k: v for k, v in tree.items() + if jnp.issubdtype(jnp.asarray(v).dtype, jnp.floating)} + return {"nodes": {n: floats(params["nodes"].get(n, {})) for n in members}, + "mappings": {k: floats(v) + for k, v in params.get("mappings", {}).items()}} + + return cls(gm, key, tuple(layout), constants or default_constants) + + def at(self, pre_state: dict, params: dict) -> "PassReference": + """Bind the pre-step state (``gm._state`` after the caller set it: + the members' fields are *pre*; outside nodes are read as the step + reads them) and the constants. Returns ``self``.""" + self._pre, self._params = pre_state, params + return self + + # -- the flat iterate --------------------------------------------------- + + def flat(self, state: dict) -> np.ndarray: + """``{node: {field: array}}`` (any float dtype) as the flat float64 iterate.""" + return np.concatenate([np.ravel(np.asarray(state[n][f], np.float64)) + for n, f, _shape, _a, _b in self.layout]) + + def field(self, x, node: str, field: str): + """One field of flat iterate *x* (NumPy or JAX), in its shape.""" + for n, f, shape, a, b in self.layout: + if (n, f) == (node, field): + return x[a:b].reshape(shape) + raise KeyError((node, field)) + + def member_fields(self, x) -> list: + """Every floating field of every member: the reading of ``"l2"`` + and ``"mixed"``.""" + return [x[a:b] for _n, _f, _shape, a, b in self.layout] + + def _flat_of(self, state): + return jnp.concatenate([jnp.ravel(state[n][f]) for n, f, _s, _a, _b in self.layout]) + + def _pass(self, x, pre, params): + p0, p1, count = self._slots + meta = {**pre["_meta"], p0: x, p1: x, count: jnp.asarray(2, pre["_meta"][count].dtype)} + return self._flat_of(self._step({**pre, "_meta": meta}, self._ext, params)) + + # -- the map and its derivatives --------------------------------------- + + def apply(self, x) -> np.ndarray: + """``P(x; pre)`` in float64.""" + return self._both(x)[0] + + def jacobian(self, x) -> np.ndarray: + """The dense ``dP/dx`` at iterate *x*.""" + return self._both(x)[1] + + def _both(self, x): + with x64(): + P, J = self._jit_both(jnp.asarray(x, jnp.float64), self._pre, self._params) + return np.asarray(P), np.asarray(J) + + def sensitivities(self, x) -> np.ndarray: + """``dP/dc`` at iterate *x*: one column per scalar constant + (:meth:`constant_names` names them).""" + with x64(): + theta, restore = ravel_pytree(self._constants_of(self._params)) + if self._jit_sens is None: + def moved(x_, theta_, pre, params): + c = restore(theta_) + nodes = {n: {**params["nodes"].get(n, {}), **c["nodes"].get(n, {})} + for n in params["nodes"]} + mappings = {k: {**v, **c.get("mappings", {}).get(k, {})} + for k, v in params.get("mappings", {}).items()} + return self._pass(x_, pre, {**params, "nodes": nodes, "mappings": mappings}) + self._jit_sens = jax.jit(jax.jacfwd(moved, argnums=1)) + return np.asarray(self._jit_sens(jnp.asarray(x, jnp.float64), theta, self._pre, + self._params)) + + def constant_names(self) -> list: + """``node.leaf[i]`` / ``mapping:key.leaf[i]`` per column of :meth:`sensitivities`.""" + with x64(): + tree = self._constants_of(self._params) + _flat, restore = ravel_pytree(tree) + index = restore(jnp.arange(_flat.size, dtype=jnp.float64)) + names = [None] * int(_flat.size) + for kind in ("nodes", "mappings"): + for owner, leaves in index.get(kind, {}).items(): + for leaf, where in leaves.items(): + for local, k in enumerate(np.ravel(np.asarray(where)).astype(int)): + prefix = "" if kind == "nodes" else "mapping:" + names[k] = f"{prefix}{owner}.{leaf}[{local}]" + return names + + # -- the fixed point ---------------------------------------------------- + + def _ulps(self, x, P) -> float: + if not (np.all(np.isfinite(x)) and np.all(np.isfinite(P))): + return math.inf + worst = 0.0 + for _n, _f, _shape, a, b in self.layout: + scale = max(float(np.max(np.abs(x[a:b]))), float(np.max(np.abs(P[a:b]))), + float(np.finfo(np.float64).tiny)) + with np.errstate(over="ignore"): + worst = max(worst, float(np.max(np.abs(P[a:b] - x[a:b])) / (EPS64 * scale))) + return worst + + def fixed_point(self, start, *, picard: int = 12, newton: int = 40) -> FixedPoint: + """The fixed point of ``P(.; pre)`` the iteration from *start* (the + returned iterate) leads to: *picard* passes of the map itself, + then Newton on ``P(x) - x`` with the dense Jacobian until the pass + moves no entry by more than :data:`FIXED_POINT_ULPS` float64 ``eps`` + of its field, and then while it still improves.""" + x = np.asarray(start, np.float64) + history = [] + for _ in range(picard): + P = self.apply(x) + history.append(self._ulps(x, P)) + if not math.isfinite(history[-1]): + break # the map left float range: Newton from the last finite iterate + x = P + if history[-1] <= 1.0: + break + n_picard = len(history) + steps = 0 + best, best_ulps = x, math.inf + eye = np.eye(self.size) + for _ in range(newton): + P, J = self._both(x) + if not (np.all(np.isfinite(P)) and np.all(np.isfinite(J))): + break + now = self._ulps(x, P) + history.append(now) + if now < best_ulps: + best, best_ulps = x, now + elif best_ulps <= FIXED_POINT_ULPS: + break # no longer improving, and already a fixed point + if now <= 1.0: + break + try: + x = x + np.linalg.solve(eye - J, P - x) + except np.linalg.LinAlgError: + break + steps += 1 + return FixedPoint(best, best_ulps, best_ulps <= FIXED_POINT_ULPS, tuple(history), + n_picard, steps) + + # -- what a report is scored against ----------------------------------- + + def norm(self, kind: str, rtol: float = 1e-6, fields: Optional[Callable] = None) -> Norm: + """The group's norm: ``"l2"`` (weights ``1 / max|field|``, a plain + 2-norm), ``"mixed"`` or ``"interface"`` (``1 / (rtol max|field|)``, + a root mean square). *fields* is the reading; the members' + floating fields by default, which ``"interface"`` is not.""" + assert kind in ("l2", "mixed", "interface"), kind + assert fields is not None or kind != "interface", ( + "the interface norm reads what the internal edges deliver: pass fields=") + return Norm(fields or self.member_fields, 1.0 if kind == "l2" else float(rtol), + kind != "l2") + + def distance(self, returned, fixed: FixedPoint, norm: Norm) -> float: + """The distance from *returned* to the fixed point in *norm* at the + returned state's weights.""" + x = np.asarray(returned, np.float64) + return norm.of_difference(x, fixed.x, norm.weights(x)) + + def residual(self, returned, norm: Norm) -> float: + """``||P(x) - x||`` of the returned state in *norm*, each field over + the larger of its magnitude before and after the pass (the weights + the step's own residual uses).""" + x = np.asarray(returned, np.float64) + P = self.apply(x) + return norm.of_difference(P, x, norm.weights(x, also=P)) + + def implicit_derivatives(self, x) -> np.ndarray: + """``(I - dP/dx)^{-1} dP/dc`` at iterate *x*: the implicit-function + derivative of the fixed point with respect to every scalar + constant, as taken at *x* (one column per constant).""" + x = np.asarray(x, np.float64) + return np.linalg.solve(np.eye(self.size) - self.jacobian(x), self.sensitivities(x)) + + def gradient_error(self, returned, fixed: FixedPoint, norm: Norm, *, + loss: Optional[Sequence[float]] = None): + """``(worst, column)``: the worst relative error, over the scalar + constants, of the implicit derivative taken at the returned iterate + against the one at the fixed point. + + With *loss* (a cotangent ``w``: the loss is ``w . x``) the error is + of the scalar ``d loss / d c``, ``|w . (g_k - g*)| / |w . g_k|``. + Without, of the vector ``d x / d c`` in *norm* (linear fields: + selections of the iterate) at the returned state's weights, the + form ``gradient_relative_error_bound`` documents. + """ + x = np.asarray(returned, np.float64) + g_k, g_star = self.implicit_derivatives(x), self.implicit_derivatives(fixed.x) + weights = norm.weights(x) + worst, column = 0.0, None + zero = np.zeros(self.size) + for c in range(g_k.shape[1]): + if loss is not None: + w = np.asarray(loss, np.float64) + size, miss = abs(float(w @ g_k[:, c])), abs(float(w @ (g_k[:, c] - g_star[:, c]))) + else: + size = norm.of_difference(g_k[:, c], zero, weights) + miss = norm.of_difference(g_k[:, c], g_star[:, c], weights) + if size > 0 and miss / size > worst: + worst, column = miss / size, c + return worst, column + + def nonlinearity(self, returned, fixed: FixedPoint, norm: Norm, *, nodes: int = 8) -> float: + """``h = ||T (I - J(x))^{-1} (J_mean - J(x)) T^+||`` with ``T`` the + weighted (linear) reading of *norm* and ``J_mean`` the mean of + ``dP/dx`` over the segment from the fixed point to ``x``. + + ``x - x* = (I - J_mean)^{-1} (x - P(x))`` exactly (the mean value + theorem in integral form), so the error a *linear* analysis at + ``x`` computes, ``e = (I - J(x))^{-1} (x - P(x))``, misses the true + one by ``(I - E)^{-1}`` with ``E = (I - J(x))^{-1} (J_mean - + J(x))``: the true distance is at most ``1 / (1 - h)`` times the + linear one where ``h < 1``. Zero on an affine map. *nodes*: the + Gauss-Legendre points the mean is taken with. + """ + x = np.asarray(returned, np.float64) + J = self.jacobian(x) + t, w = np.polynomial.legendre.leggauss(nodes) + mean = sum(wk * self.jacobian(fixed.x + 0.5 * (tk + 1.0) * (x - fixed.x)) + for tk, wk in zip(t, w)) / 2.0 + T = self.reading_matrix(norm, x) + E = np.linalg.solve(np.eye(self.size) - J, mean - J) + return float(np.linalg.norm(T @ E @ np.linalg.pinv(T), 2)) + + def reading_matrix(self, norm: Norm, returned) -> np.ndarray: + """The weighted reading of a *linear* norm as a matrix: its rows + applied to a flat iterate are ``w fields(x)``.""" + x = np.asarray(returned, np.float64) + weights = norm.weights(x) + zero = [np.asarray(f, np.float64) for f in norm.fields(np.zeros(self.size))] + columns = [] + for k in range(self.size): + unit = np.zeros(self.size) + unit[k] = 1.0 + columns.append(np.concatenate([ + w * (np.ravel(np.asarray(f, np.float64)) - np.ravel(z)) + for w, f, z in zip(weights, norm.fields(unit), zero)])) + return np.stack(columns, axis=1) if columns else np.zeros((0, 0)) + + +def passes_compose(reference: PassReference, several, passes: int, start) -> float: + """How far *passes* passes of an iterating twin are from *passes* + compositions of the reference's single pass, in float64 ``eps`` of + each field. + + *several* is a second x64 twin of the same graph whose group runs + ``max_iterations=passes`` with no acceleration, no predictor and a + tolerance it cannot meet, holding the same pre-step state and + constants the reference is bound to; *start* is that state's flat + iterate. The single-pass branch is another branch of the step than + the iteration (``CouplingGroup.max_iterations``): this is the check + that both run the same pass. + """ + with x64(): + after = several._raw_step_fn(several._state, several._default_external_inputs(), # noqa: SLF001 + reference._params) # noqa: SLF001 + got = reference.flat({n: {f: np.asarray(v) for f, v in after[n].items()} + for n, _f, _s, _a, _b in reference.layout}) + x = np.asarray(start, np.float64) + for _ in range(passes): + x = reference.apply(x) + return reference._ulps(got, x) # noqa: SLF001 diff --git a/tests/property/geometry_graphs.py b/tests/property/geometry_graphs.py index 0d45dd7e..fc9410fd 100644 --- a/tests/property/geometry_graphs.py +++ b/tests/property/geometry_graphs.py @@ -425,6 +425,173 @@ def inline_geometry(graph: GGraph) -> GGraph: return GGraph(nodes, edges, list(graph.groups)) +# --------------------------------------------------------------------------- +# The twin with an identical interface reading +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class _Relayed: + sf: str # the source field the edge read + out: str # the state field that holds the mapped value + mapping: Any + gfield: Optional[str] # the source's own geometry field, or None (a static mapping) + leaves: tuple # the mapping's weight leaves, held as node parameters + + +class RelayedSourceNode(SimulationNode): + """A source node with the mappings of its outgoing edges inside it: the + relay of :func:`relay_twin`. + + Same name, timestep and evaluation count as the node it wraps, the + same state and one more field per relayed edge: ``out = + mapping.apply(new[sf], weights, new[gfield])``, computed by ``update`` + from the fields it has just produced (and by ``initial_state`` from + the initial ones), so in every state the graph holds -- and in every + iterate of a group, which is a state some pass returned -- ``out`` is + the mapped value of the ``sf`` beside it. The mapping's weights are + parameters of this node (``"_"``), so they stay constants + of the step that a gradient or a per-step override can reach. + """ + + def __init__(self, inner: SimulationNode, relayed): + params = dict(inner.params) + for r in relayed: + weights = r.mapping.params_pytree() + for leaf in r.leaves: + params[f"{r.out}_{leaf}"] = jnp.asarray(weights[leaf]) + super().__init__(inner.name, inner.delta_t, **params) + self._wrapped = inner + self._relayed = tuple(relayed) + self._outs = frozenset(r.out for r in relayed) + + def _mapped(self, r: _Relayed, fields, p): + weights = {leaf: p[f"{r.out}_{leaf}"] for leaf in r.leaves} or None + if r.gfield is None: + return r.mapping.apply(fields[r.sf], weights) + return r.mapping.apply(fields[r.sf], weights, fields[r.gfield]) + + def initial_state(self): + state = dict(self._wrapped.initial_state()) + for r in self._relayed: + assert r.sf in state, f"{self.name}.{r.sf} is not a state field (a flux is not relayed)" + state[r.out] = self._mapped(r, state, self.params) + return state + + def param_specs(self): + return self._wrapped.param_specs() + + def update_evaluations(self): + return self._wrapped.update_evaluations() + + def boundary_input_spec(self): + return self._wrapped.boundary_input_spec() + + def update(self, state, boundary_inputs, dt, *, params=None): + p = self.params if params is None else {**self.params, **params} + inner = {k: v for k, v in state.items() if k not in self._outs} + new = dict(self._wrapped.update(inner, boundary_inputs, dt, params=params)) + for r in self._relayed: + new[r.out] = self._mapped(r, new, p) + return new + + +def relay_twin(graph: GGraph) -> GGraph: + """*graph* with every mapped edge's mapping moved into a relay on its + source, so that the twin's **interface reading is the edge-mapped + graph's**. + + ``convergence_norm="interface"`` reads what each internal edge + delivers: for ``S.sf -> T.tf`` through mapping ``m`` and transform + ``t``, ``t(m(S.sf))``. The node-inlined twin of + :func:`inline_geometry` moves ``m`` into the *target*, so its edge + delivers the raw ``S.sf`` and its interface norm is another norm: the + two graphs' diagnostics agree only as far as two norms of one solve + do. Here the mapped value is a state field of the source, + ``S."_to__" = m(S.sf)``, and the edge is the plain edge + ``S."_to__" -> T.tf`` with the same transform and additive + flag in the same place of the edge order: it delivers ``t(m(S.sf))``, + the same reading entry for entry, with the mapping (and, for a source + anchor, its geometry) inside the relay. + + The relay is fused with the source (:class:`RelayedSourceNode`) and is + not a node of its own: a separate relay in the group would be fed by + an internal edge ``S.sf -> relay`` that the interface norm reads too, + and outside the group it would sit in the group's cycle. + + What the twin holds equal by construction, and what it does not: + + * equal: every value a node reads, so the fixed point and every + iterate of an unaccelerated solve; the interface norm's reading, so + the residual, the pass count, the verdict, the spectrum taken on the + reading (``rho_spectral``, ``spectral_error_bound``) and the flags; + * not equal: anything that reads the members' *state*, which has the + relay's field in it here -- the ``"l2"`` and ``"mixed"`` norms, an + accelerator's secant vectors, the gradient bound's norm (the raw + source fields: ``S.sf`` there, the relay's field here) -- and the + float floor, which counts an inner product for a mapping the norm + evaluates and none for a field read as stored. + + A geometry-dependent mapping is relayed where its geometry is its + source's (``anchor="source"``: the relay reads the geometry field + beside the value, both as the update has just produced them). A + target-anchored geometry is refused: the relay would need the + target's geometry over one more internal edge, which the norm would + read. A flux source field is refused too (it is not state). + """ + edges: list = [] + relayed: dict = {} + for e in graph.edges: + if e.mapping is None: + edges.append(e) + continue + gfield = None + if e.geometry is not None: + anchor, gfield = e.geometry + if anchor != "source": + raise NotImplementedError( + f"{e}: a target-anchored geometry has no relay twin (the relay would read " + f"the target's geometry over an internal edge of its own)") + out = f"{e.sf}_to_{e.dst}_{e.tf}" + leaves = tuple(sorted(e.mapping.params_pytree())) + relayed.setdefault(e.src, []).append(_Relayed(e.sf, out, e.mapping, gfield, leaves)) + edges.append(GEdge(e.src, e.dst, out, e.tf, transform=e.transform, additive=e.additive)) + nodes = [RelayedSourceNode(nd, relayed[nd.name]) if nd.name in relayed else nd + for nd in graph.nodes] + return GGraph(nodes, edges, list(graph.groups)) + + +def transform_twin(graph: GGraph) -> GGraph: + """*graph* with every *static* mapping written as its edge's transform: + the same nodes, the same state, plain edges, and on each formerly + mapped edge the transform ``v -> t(m(v))`` (the mapping with its own + weights, then the edge's transform ``t``). + + The interface norm reads an edge through its mapping and then its + transform, so the twin's reading is the edge-mapped graph's and -- + unlike :func:`relay_twin` -- so is its state. A mapping that reads a + geometry cannot be written this way (a transform sees the value + only), which is what :func:`relay_twin` is for. The weights are + constants of the transform here and parameters of the step there, so + the gradient bound, which probes every parameter, is the one number + the two graphs do not share. + """ + edges = [] + for e in graph.edges: + if e.mapping is None: + edges.append(e) + continue + assert e.geometry is None, f"{e}: a geometry-dependent mapping is not a transform" + then = resolve_transform(e.transform) + + def through(v, _m=e.mapping, _t=then): + v = _m.apply(v, None) + return v if _t is None else _t(v) + + edges.append(GEdge(e.src, e.dst, e.sf, e.tf, transform=through, additive=e.additive)) + return GGraph(list(graph.nodes), edges, list(graph.groups)) + + def build(graph: GGraph, *, compile: bool = True) -> GraphManager: """A ``GraphManager`` of *graph*, compiled. @@ -471,6 +638,13 @@ def build(graph: GGraph, *, compile: bool = True) -> GraphManager: # deleted: set ``DIAGNOSTICS_READ_GEOMETRY = True`` when a later phase makes # the diagnostics read the geometry, and the original comparisons and cases # run again. +# +# Waiting for that phase too: ``RELAY_INTERFACE_CASES`` (defined after +# ``case``, below), the interface norm over source-anchored geometry edges, +# whose relay twin (``relay_twin``) has the edge-mapped graph's interface +# reading. ``tests/property/test_interface_reading_twins.py`` asserts each +# refused today and its relay twin reporting; with the switch set it +# compares the two reports as it compares a static mapping's today. # ============================================================================= #: Whether coupling diagnostics account for a moving geometry. @@ -613,6 +787,24 @@ def case(label: str, *, group: Optional[dict] = None, **kw) -> Case: return Case(label, group=None if group is None else tuple(sorted(group.items())), **kw) +#: PHASE 1 (the block of that name above): the cases whose interface +#: diagnostics are held to their relay twin's once diagnostics read a +#: geometry. Source anchors on both edges (a target-anchored geometry has +#: no relay twin), the geometry moving with the iterate, Gauss-Seidel (the +#: relay doubles the scalars a Jacobi pass carries, past what the spectrum +#: resolves). +RELAY_INTERFACE_CASES = [ + case("geometry edges, interface norm, three passes", adv=0.3, dtype="float64", + down="source", up="source", + group=dict(max_iterations=3, convergence_norm="interface", rtol=1e-6, + diagnostics=True)), + case("multilinear geometry edges, interface norm, three passes", kind="multilinear", + adv=0.3, down="source", up="source", + group=dict(max_iterations=3, convergence_norm="interface", rtol=1e-6, + diagnostics=True)), +] + + def _grid_of(c: Case): return GRIDS[c.d] diff --git a/tests/property/test_coupling_nonlinear_search.py b/tests/property/test_coupling_nonlinear_search.py new file mode 100644 index 00000000..8c0ca7ba --- /dev/null +++ b/tests/property/test_coupling_nonlinear_search.py @@ -0,0 +1,685 @@ +"""The coupling search on nonlinear groups, scored against a numerical float64 reference. + +``test_coupling_targeted_search.py`` scores ``coupling_diagnostics()`` on +groups of *linear* relays, whose fixed point and Jacobian are closed +forms. Its own docstring names what that leaves out: a nonlinear node +(the Newton-Kantorovich factor of the gradient bound is exactly one on an +affine map, and the Jacobian is the same at every iterate). This module +adds nonlinear cells to that search, scored by the same four scores +against :mod:`tests.property.coupling_reference` -- the pass map of an x64 +twin of the graph, its fixed point by Newton and its dense Jacobian by +``jax.jacfwd``. + +**The reference is validated first** (the first tests below): on the +linear cells of the search it reproduces the closed form's fixed point, +Jacobian, spectral radius, distance, residual and gradient error to +float64 rounding, and several passes of an iterating twin are the same +number of compositions of its single pass. + +**The nonlinear cells** (:data:`CELLS`) are structures of the linear +search with every member a :class:`NRelay`: the linear relay with each +port's input ``u`` passed through a nonlinearity ``phi`` centred on a +parameter ``c`` before its gain, + +* ``"saturating"``: ``c + tanh(s (u - c)) / s`` (a saturating gain); +* ``"quadratic"``: ``u + s (u - c)^2`` (a quadratic term); +* ``"product"``: ``u_j + s (u_j - c_j) (u_j' - c_j')`` with ``j'`` the + node's next port (the product of two fields' deviations; a node with + one port squares its own). + +``phi(c) = c`` and ``phi'(c) = 1`` for each, so with ``c`` the input at +the linear group's fixed point the nonlinear group has *the same fixed +point and the same Jacobian there* as the linear one: the numbers of a +case are drawn exactly as the linear search draws them (loop gain, +non-normal gains, a small field, a change of units, the start's offset: +:func:`tests.property.test_coupling_targeted_search.values_of`), and the +linear closed form is a second, independent check of the reference's +fixed point on every nonlinear example. One more number is drawn, the +*curve* ``s max|c|``: how nonlinear the map is over a field's own size. +Away from the fixed point -- where a capped or early-stopped solve +returns -- the Jacobian is another matrix, and only the numerical +reference knows it. + +**The scores** are those of the linear search, restated for a map that +is not affine: + +1. *bound* (CPL-088, which claims the bound for a linear ``F`` and calls + it asymptotic for a nonlinear one): the true distance over + ``spectral_error_bound`` times ``1 / (1 - h)``, with ``h`` the + reference's measure of how far the Jacobian moves between the returned + iterate and the fixed point + (:meth:`~tests.property.coupling_reference.PassReference.nonlinearity`). + ``x - x* = (I - J_mean)^{-1} (x - F(x))`` exactly, so a bound that is + right for the linearisation at the returned iterate can be short by + that factor and no more: a score over one is a wrong number, not a + nonlinear map. Not scored where ``h >= 1`` (nothing taken at the + returned iterate bounds the distance; the fraction is reported). +2. *radius* / *radius_strict* (CPL-087: "the spectral radius of dF/dx at + the returned state"): as the linear search scores them, against the + reference's Jacobian **at the returned iterate**. +3. *gradient* (CPL-093, CPL-095): the true relative error of the implicit + derivative taken at the returned iterate, the worst over every scalar + constant (the nonlinearity's ``c`` and ``s`` included), over + ``gradient_relative_error_bound`` where ``gradient_bound_usable`` -- + the flag that certifies the Newton-Kantorovich check, which on these + cells is not trivially passed. +4. *floor* (CPL-097, CPL-100): the exact residual of the returned state + above the reported one, over the reported floor. + +Per push: the validation on one linear cell and each score on +:data:`PER_PUSH_CELLS` at the house floor of examples. Slow: the +validation on every per-push draw of the linear search, and the hunt over +every nonlinear cell. What goes over a threshold is pinned at the foot +of the module, as in the linear search. +""" + +from __future__ import annotations + +import dataclasses +import functools +import math +import os +import warnings + +os.environ.setdefault("JAX_PLATFORMS", "cpu") + +import jax.numpy as jnp +import numpy as np +import pytest +from hypothesis import strategies as st + +from maddening.core.coupling.mapping import matrix_mapping +from maddening.core.graph_manager import GraphManager +from maddening.core.node import BoundaryInputSpec, SimulationNode +from tests.property import coupled_graphs as cg +from tests.property import coupled_topologies as ct +from tests.property import coupling_reference as cr +from tests.property import test_coupling_targeted_search as linear +from tests.property.sysid_transform_grid import precision +from tests.property.targeted_search import PER_PUSH, SLOW, targeted_search + +EPS64 = cr.EPS64 + + +# --------------------------------------------------------------------------- +# What both halves share: the norm's reading on a topology of relays +# --------------------------------------------------------------------------- + + +def edge_fields(topo: ct.Topology, values: dict, ref: cr.PassReference, *, raw: bool = False): + """The interface norm's reading of a flat iterate: what each internal + edge delivers (its source's ``x`` through the edge's ``H`` and its + transform's factor), or with *raw* the source field it reads -- one + field per internal edge either way, as that norm counts them.""" + edges = [(e, np.asarray(values["H"][i], np.float64) if e.mapped else None) + for i, e in enumerate(topo.edges) if topo.internal(e)] + + def fields(x): + out = [] + for e, H in edges: + src = np.asarray(ref.field(x, e.src, "x"), np.float64) + if raw: + out.append(src) + else: + out.append(ct.TRANSFORM_FACTORS[e.transform] * (src if H is None else H @ src)) + return out + + return fields + + +def norms_of(cfg: dict, topo: ct.Topology, values: dict, ref: cr.PassReference) -> tuple: + """``(norm, raw norm)`` of a group under *cfg*: the norm its error + bound is stated in, and the one its gradient bound is (the raw source + fields under ``"interface"``; the same norm otherwise).""" + kind, rtol = cfg["convergence_norm"], cfg["rtol"] + if kind != "interface": + norm = ref.norm(kind, rtol) + return norm, norm + return (ref.norm(kind, rtol, edge_fields(topo, values, ref)), + ref.norm(kind, rtol, edge_fields(topo, values, ref, raw=True))) + + +def state_weights(ref: cr.PassReference, x: np.ndarray) -> np.ndarray: + """``1 / max|field|`` per entry of flat iterate *x*.""" + w = np.ones(ref.size) + for _n, _f, _shape, a, b in ref.layout: + top = float(np.max(np.abs(x[a:b]))) + w[a:b] = 1.0 / top if top > 0 else 1.0 + return w + + +def bound_reference(ref: cr.PassReference, values: dict, built_twin: ct.Built) -> cr.PassReference: + """*ref* bound to the start and the constants of *values* on its twin.""" + with cr.x64(): + ct.set_initial(built_twin, values) + params = params_for(built_twin, values) + return ref.at(built_twin.gm._state, params) # noqa: SLF001 + + +# --------------------------------------------------------------------------- +# The reference against the closed form, on the linear cells +# --------------------------------------------------------------------------- + + +#: The example a twin's layout probe steps (see ``PassReference.of``). +_PROBE = linear.Case(0, 1, 0.6, False, 1.0, 0.0, 1.0, 0, 0.5, 0) + + +@functools.lru_cache(maxsize=8) +def _linear_twin(index: int) -> tuple: + cell = linear.CELLS[index] + with cr.x64(): + built = ct.build(cell.topo, cr.twin_knobs(cell.knobs), dtype="float64", + mapping_kind=cell.mapping_kind) + # A drawn example, so that the members' fields differ after a step + # and the predictor slot's layout is determined. + values = linear.values_of(dataclasses.replace(_PROBE, cell=index)) + ct.set_initial(built, values) + return built, cr.PassReference.of(built.gm, params=ct.params_for(built, values)) + + +def against_the_closed_form(case: linear.Case) -> dict: + """How far each answer of the numerical reference is from the closed + form's, on one drawn linear case (the step the search itself scores).""" + cell = linear.CELLS[case.cell] + topo = cell.topo + values = linear.values_of(case) + with precision(cell.dtype == "float64"): + (step,) = ct.run(linear._built(case.cell), values, 1) # noqa: SLF001 + out = dict(finite=all(np.all(np.isfinite(s["x"])) for s in step.state.values())) + if not out["finite"]: + return out + model = ct.LinearModel(topo, values, dtype=cell.dtype, group_cfgs=cell.cfgs) + built_twin, ref = _linear_twin(case.cell) + ref = bound_reference(ref, values, built_twin) + x = ref.flat(step.state) + fixed = ref.fixed_point(x) + members, off, _k = model._group_layout(0) # noqa: SLF001 + # The model's stacking order, in the reference's. + perm = np.concatenate([np.arange(off[n], off[n] + topo.node(n).n) + for n, _f, _s, _a, _b in ref.layout]) + exact = model.group_fixed_point(0, step.pre, step.state) + x_star = np.concatenate([np.asarray(exact[m], np.float64) for m in members])[perm] + J_exact = linear._pass_jacobian(model)[np.ix_(perm, perm)] # noqa: SLF001 + J = ref.jacobian(fixed.x) + norm, raw = norms_of(cell.cfgs[0], topo, values, ref) + resolvent = float(np.linalg.norm(np.linalg.inv(np.eye(ref.size) - J_exact), 2)) + dist, dist_exact = ref.distance(x, fixed, norm), model.returned_weight_distance( + 0, step.pre, step.state) + _dn, _b, detail = model.group_report_consistency(0, step.pre, step.state, 0.0) + grad, grad_exact = ref.gradient_error(x, fixed, raw)[0], linear._gradient_error( # noqa: SLF001 + model, step.pre, step.state) + scale = max(float(np.max(np.abs(x_star))), np.finfo(np.float64).tiny) + # A relative norm divides a field by ``rtol`` times its size: one float64 + # rounding of a field is this much of the norm's unit. + unit = EPS64 / norm.rtol + out.update( + converged=fixed.converged, ulps=fixed.ulps, resolvent=resolvent, + fixed_point=float(np.max(np.abs(fixed.x - x_star))) / (EPS64 * scale * resolvent), + jacobian=float(np.max(np.abs(J - J_exact))) / ( + EPS64 * max(float(np.max(np.abs(J_exact))), np.finfo(np.float64).tiny)), + radius=abs(cr.radius(J) - linear._radius(J_exact)), # noqa: SLF001 + distance=abs(dist - dist_exact), distance_exact=dist_exact, norm_unit=unit, + residual=abs(ref.residual(x, norm) - detail["residual_true"]), + residual_exact=detail["residual_true"], + gradient=abs(grad - grad_exact), gradient_exact=grad_exact) + return out + + +#: What "to float64 rounding" allows each difference, as measured over +#: every per-push draw of the linear search (see +#: ``test_the_reference_reproduces_the_closed_form_on_every_per_push_draw``): +#: the fixed point in float64 ``eps`` of the field times the resolvent's +#: norm; the Jacobian in ``eps`` of its largest entry; the radius absolute; +#: the distance and the residual in the norm's own rounding unit ``eps / +#: rtol`` times the resolvent's norm (a relative norm divides a field by +#: ``rtol`` times its size); the gradient error absolute beside one. +ALLOWED = dict(fixed_point=2.0 ** 10, jacobian=2.0 ** 6, radius=2.0 ** 10 * EPS64, + distance=2.0 ** 12, residual=2.0 ** 12, gradient=1e-9) + + +def closed_form_misses(seen: dict) -> dict: + """The differences of :func:`against_the_closed_form` over what + :data:`ALLOWED` gives each (at most 1 where the reference agrees).""" + if not seen["finite"]: + return {} + unit = seen["norm_unit"] * max(seen["resolvent"], 1.0) + return dict( + converged=0.0 if seen["converged"] else math.inf, + fixed_point=seen["fixed_point"] / ALLOWED["fixed_point"], + jacobian=seen["jacobian"] / ALLOWED["jacobian"], + radius=seen["radius"] / ALLOWED["radius"], + distance=seen["distance"] / (ALLOWED["distance"] * ( + unit + EPS64 * seen["distance_exact"])), + residual=seen["residual"] / (ALLOWED["residual"] * ( + unit + EPS64 * seen["residual_exact"])), + gradient=seen["gradient"] / (ALLOWED["gradient"] * max(seen["gradient_exact"], 1.0))) + + +# --------------------------------------------------------------------------- +# The nonlinear relay and its cells +# --------------------------------------------------------------------------- + +KINDS = ("saturating", "quadratic", "product") + + +def phi(kind: str, u: list, c: list, s: list, xp=np) -> list: + """Each port's input through the nonlinearity (module docstring); + ``xp`` is ``numpy`` (the restatement the scores use) or ``jax.numpy``.""" + d = [uj - cj for uj, cj in zip(u, c)] + k = len(u) + if kind == "saturating": + return [c[j] + xp.tanh(s[j] * d[j]) / s[j] for j in range(k)] + if kind == "quadratic": + return [u[j] + s[j] * d[j] * d[j] for j in range(k)] + assert kind == "product", kind + return [u[j] + s[j] * d[j] * d[(j + 1) % k] for j in range(k)] + + +class NRelay(SimulationNode): + """``x <- alpha x_pre + b + sum_j G_j phi(u_j; c_j, s_j)``, at any float + dtype: :class:`~tests.property.coupled_topologies.TRelay` with a + nonlinearity on every port. Every constant is a parameter of the + step, so a compiled graph serves every draw.""" + + def __init__(self, name, timestep, *, kind, n, ports, alpha=0.0, dtype="float32"): + dt_ = jnp.dtype(dtype) + params = {"b": jnp.zeros(n, dt_)} + for j in range(ports): + params[f"G{j}"] = jnp.zeros((n, n), dt_) + params[f"c{j}"] = jnp.zeros(n, dt_) + params[f"s{j}"] = jnp.ones((), dt_) + super().__init__(name, timestep, **params) + self._kind, self._n, self._k = kind, int(n), int(ports) + self._alpha, self._dtype = float(alpha), dt_ + + def initial_state(self): + return {"x": jnp.zeros(self._n, self._dtype)} + + def boundary_input_spec(self): + return {f"u{j}": BoundaryInputSpec(shape=(self._n,), dtype=self._dtype, + default=jnp.zeros(self._n, self._dtype)) + for j in range(self._k)} + + def update(self, state, boundary_inputs, dt, *, params=None): + p = self.params if params is None else {**self.params, **params} + dtype = self._dtype + f = jnp.asarray(self._alpha, dtype) * state["x"] + p["b"] + u = [boundary_inputs.get(f"u{j}", jnp.zeros(self._n, dtype)) for j in range(self._k)] + v = phi(self._kind, u, [p[f"c{j}"] for j in range(self._k)], + [p[f"s{j}"] for j in range(self._k)], xp=jnp) + for j in range(self._k): + f = f + p[f"G{j}"] @ v[j] + return {"x": f.astype(dtype)} + + def update_evaluations(self): + return 1 + + +def _tri() -> ct.Topology: + """Three members of two entries: ``a`` reads ``c``; ``b`` reads ``a``; + ``c`` reads ``a`` and ``b`` on two ports (the two fields a product + multiplies).""" + b = ct.TopologyBuilder() + b.node("a", 2, alpha=0.5) + b.node("b", 2, alpha=0.0) + b.node("c", 2, alpha=-0.25) + b.edge("c", "a") + b.edge("a", "b") + b.edge("a", "c") + b.edge("b", "c") + b.group("a", "b", "c") + return b.build("tri") + + +#: The single-rate structures of the linear search (the centre ``c`` of a +#: nonlinearity is the input at the fixed point, which a sub-cycled +#: member reads interpolated), and ``tri``. +STRUCTURES = {**{name: linear.STRUCTURES[name] + for name in ("ring-2", "ring-3", "ring-5", "pair-3", "hub", "mapped")}, + "tri": _tri()} + + +@dataclasses.dataclass(frozen=True) +class Cell: + """What one compiled nonlinear graph bakes in: a structure, a dtype, + a configuration and a cap as the linear search's cells hold them, and + the nonlinearity: one of :data:`KINDS` on every member, or ``"each"`` + -- the three in turn, by the members' order.""" + + structure: str + dtype: str + knob: int + cap: int + kind: str + mapping_kind: str = "matrix" + + @property + def topo(self) -> ct.Topology: + return STRUCTURES[self.structure] + + @property + def knobs(self) -> dict: + return linear.Cell("ring-3", self.dtype, self.knob, self.cap).knobs + + @property + def cfgs(self) -> list: + return ct.group_cfgs_of([self.knobs]) + + def kind_of(self, name: str) -> str: + return (KINDS[self.topo.names.index(name) % len(KINDS)] if self.kind == "each" + else self.kind) + + +def _cells() -> tuple: + """Every structure with every nonlinearity at both dtypes, the + configurations and the caps rotated as the linear search rotates them.""" + out = [] + for s, name in enumerate(STRUCTURES): + for q, kind in enumerate(KINDS): + for t, dtype in enumerate(("float32", "float64")): + knob = (s + 3 * t + 2 * q) % len(linear.KNOBS) + out.append(Cell(name, dtype, knob, linear.CAPS[(s + t + q) % 2], kind)) + return tuple(out) + + +#: The per-push cell: ``tri`` with a saturating gain on ``a``, a quadratic +#: term on ``b`` and the product of two fields on ``c``, in float32, +#: stopped after five Jacobi passes (a returned iterate away from the +#: fixed point, where the Jacobian is another matrix). One cell, because +#: a cell's cost is the compile of the graph under test with its +#: diagnostics (seconds); the three nonlinearities apart, at both dtypes +#: and under every configuration, are the slow hunt's. +_FIRST = (Cell("tri", "float32", 1, 5, "each"),) +CELLS = _FIRST + tuple(c for c in _cells() if c not in _FIRST) +PER_PUSH_CELLS = tuple(range(len(_FIRST))) +ALL_CELLS = tuple(range(len(CELLS))) +#: The slow hunt's blocks (each cell compiles a graph and a twin). +BLOCKS = tuple(ALL_CELLS[k::6] for k in range(6)) + + +def build(cell: Cell, knobs: dict, dtype: str) -> ct.Built: + """A compiled graph of *cell*'s structure with every member an + :class:`NRelay` (``ct.build`` for these nodes: dense mappings, the + group last).""" + topo = cell.topo + assert all(nd.timestep == topo.nodes[0].timestep and not nd.leaves and not nd.flux + and not nd.beta for nd in topo.nodes), "single-rate plain relays only" + assert len(topo.groups) == 1 and set(topo.groups[0]) == set(topo.names) + gm = GraphManager() + for nd in topo.nodes: + gm.add_node(NRelay(nd.name, nd.timestep, kind=cell.kind_of(nd.name), n=nd.n, ports=nd.ports, + alpha=nd.alpha, dtype=dtype)) + keys = {} + for i, e in enumerate(topo.edges): + mapping = (matrix_mapping(np.zeros((topo.node(e.dst).n, topo.node(e.src).n), + jnp.dtype(dtype))) if e.mapped else None) + gm.add_edge(e.src, e.dst, e.field, f"u{e.port}", transform=e.transform, + additive=e.additive, mapping=mapping) + if e.mapped: + keys[i] = gm._edges[i].key # noqa: SLF001 + gm.add_coupling_group(list(topo.groups[0]), **cg.live_knobs(knobs)) + with warnings.catch_warnings(record=True): + warnings.simplefilter("always") + gm.compile() + return ct.Built(gm, topo, str(dtype), keys, [], topo.names, + tuple(range(len(topo.edges)))) + + +def params_for(built: ct.Built, values: dict) -> dict: + """``ct.params_for`` and, where *values* holds them, every node's + nonlinearity constants.""" + params = ct.params_for(built, values) + for name, v in values.get("nonlinear", {}).items(): + node = params["nodes"][name] + for j, (c, s) in enumerate(zip(v["c"], v["s"])): + node[f"c{j}"] = jnp.asarray(c, node[f"c{j}"].dtype) + node[f"s{j}"] = jnp.asarray(s, node[f"s{j}"].dtype) + return params + + +@functools.lru_cache(maxsize=max(len(b) for b in BLOCKS) + 1) +def _built(index: int) -> tuple: + """``(the graph under test, its x64 twin, the twin's reference)``.""" + cell = CELLS[index] + with precision(cell.dtype == "float64"): + built = build(cell, cell.knobs, cell.dtype) + with cr.x64(): + twin = build(cell, cr.twin_knobs(cell.knobs), "float64") + values = values_of(Case(index, _PROBE, 1.0)) + ct.set_initial(twin, values) + return built, twin, cr.PassReference.of(twin.gm, params=params_for(twin, values)) + + +@dataclasses.dataclass(frozen=True) +class Case: + """One drawn problem on one nonlinear cell: the numbers of the linear + search's case *base* (its ``cell`` is not read) and the curve.""" + + cell: int + base: linear.Case + #: ``s max|c|``: the nonlinearity over a field's own size. + curve: float + + @property + def eps(self) -> float: + return float(np.finfo(CELLS[self.cell].dtype).eps) + + +def delivered(topo: ct.Topology, values: dict, x: dict) -> dict: + """``{node: [u_0, ...]}``: what every port reads when each source + holds ``x[source]`` (float64).""" + out = {nd.name: [np.zeros(nd.n) for _ in range(nd.ports)] for nd in topo.nodes} + for i, e in enumerate(topo.edges): + src = np.asarray(x[e.src], np.float64) + v = np.asarray(values["H"][i], np.float64) @ src if e.mapped else src + out[e.dst][e.port] = out[e.dst][e.port] + ct.TRANSFORM_FACTORS[e.transform] * v + return out + + +def values_of(case: Case) -> dict: + """The linear search's values for ``case.base`` on this cell, and each + port's nonlinearity centred on its input at the linear fixed point.""" + cell = CELLS[case.cell] + topo = cell.topo + values = linear.values_of(case.base, cell=cell) + model = ct.LinearModel(topo, values, dtype=cell.dtype, group_cfgs=cell.cfgs) + pre = {m: {"x": np.asarray(values["nodes"][m]["x0"])} for m in topo.names} + exact = model.group_fixed_point(0, pre, pre) + values["fixed_point"] = {m: np.asarray(exact[m], np.float64) for m in topo.names} + centres = delivered(topo, values, values["fixed_point"]) + dt = np.dtype(cell.dtype) + values["nonlinear"] = {} + for m in topo.names: + c = [np.asarray(cj, dt) for cj in centres[m]] + s = [np.asarray(case.curve / max(float(np.max(np.abs(cj))), np.finfo(dt).tiny), dt) + for cj in c] + values["nonlinear"][m] = {"c": c, "s": s} + return values + + +def _cancellation(cell: Cell, values: dict, pre: dict, state: dict) -> float: + """How far the worst member cancels inside itself (the linear search's + ``_cancellation``, for an :class:`NRelay`): the magnitudes its update + sums over its field's largest entry.""" + topo = cell.topo + reads = delivered(topo, values, {m: state[m]["x"] for m in topo.names}) + worst = 1.0 + for nd in topo.nodes: + v, nl = values["nodes"][nd.name], values["nonlinear"][nd.name] + terms = (abs(nd.alpha) * np.abs(np.asarray(pre[nd.name]["x"], np.float64)) + + np.abs(np.asarray(v["b"], np.float64))) + through = phi(cell.kind_of(nd.name), reads[nd.name], [np.asarray(c, np.float64) for c in nl["c"]], + [float(s) for s in nl["s"]]) + for j in range(nd.ports): + terms = terms + np.abs(np.asarray(v["G"][j], np.float64)) @ np.abs(through[j]) + size = float(np.max(np.abs(np.asarray(state[nd.name]["x"], np.float64)))) + if size > 0: + worst = max(worst, float(np.max(terms)) / size) + return worst + + +def run_once(built: ct.Built, values: dict) -> ct.Step: + """One step from the drawn start (``ct.run`` with this module's constants).""" + ct.set_initial(built, values) + params = params_for(built, values) + gm = built.gm + pre = ct._snapshot(gm, {}) # noqa: SLF001 + gm.step(params=params) + key = built.topo.group_key(0) + return ct.Step(pre, ct._snapshot(gm, {}), {0: dict(gm.coupling_diagnostics()[key])}, # noqa: SLF001 + {0: cg.group_meta(gm, key)}) + + +def leaves_float_range(cell: Cell, values: dict, ref: cr.PassReference) -> bool: + """Whether plain passes from the drawn start leave *cell*'s float + range before its cap, on a cell with a quasi-Newton acceleration. + + Such an example is not stepped: under ``"iqn-ils"`` and ``"iqn-imvj"`` + a group whose iterate becomes non-finite before the cap never returns + from ``step()`` (the finding pinned by + ``test_a_group_that_leaves_float_range_under_iqn_still_returns``), and + a search cannot score a call that does not return. Every other + acceleration returns a non-finite state, which is scored as unusable. + *ref* is bound to *values*. + """ + if not str(cell.knobs.get("acceleration", "none")).startswith("iqn"): + return False + x = ref.flat({m: {"x": values["nodes"][m]["x0"]} for m in cell.topo.names}) + limit = math.sqrt(float(np.finfo(cell.dtype).max)) + for _ in range(min(cell.cap, 60)): + with np.errstate(all="ignore"): + x = ref.apply(x) + if not np.all(np.isfinite(x)) or float(np.max(np.abs(x))) > limit: + return True + return False + + +@functools.lru_cache(maxsize=4096) +def observe(case: Case) -> dict: + """One step of *case* and the scores of what it reported.""" + cell = CELLS[case.cell] + topo = cell.topo + values = values_of(case) + built, twin, ref = _built(case.cell) + ref = bound_reference(ref, values, twin) + if leaves_float_range(cell, values, ref): + # Not stepped (see ``leaves_float_range``): nothing is scored. + return dict(bound=0.0, radius=0.0, radius_strict=0.0, gradient=0.0, floor=0.0, + spectral_usable=False, gradient_usable=False, floor_reported=False, + referenced=False, near=False, stepped=False, report={}) + with precision(cell.dtype == "float64"): + step = run_once(built, values) + d = dict(step.reports[0]) + floor = linear._reported_floor(built.gm, topo.group_key(0), step.metas[0], d) # noqa: SLF001 + out = dict(bound=0.0, radius=0.0, radius_strict=0.0, gradient=0.0, floor=0.0, + spectral_usable=bool(d["spectral_usable"]), + gradient_usable=bool(d["gradient_bound_usable"]), + floor_reported=math.isfinite(floor), referenced=False, near=False, stepped=True, + report={k: d[k] for k in ("iterations", "converged", "residual", "rho_spectral", + "spectral_error_bound", "spectral_usable", + "gradient_relative_error_bound", + "gradient_bound_usable", "precision_limited")}) + finite = all(np.all(np.isfinite(s["x"])) for s in step.state.values()) + if not finite or not math.isfinite(d["residual"]): + return out + x = ref.flat(step.state) + fixed = ref.fixed_point(x) + out["report"].update(reference_ulps=fixed.ulps, reference_steps=(fixed.picard, fixed.newton)) + if not fixed.converged: + return out # no reference: nothing is scored (the fraction is held to a floor) + out["referenced"] = True + exact = ref.flat({m: {"x": values["fixed_point"][m]} for m in topo.names}) + out["report"]["fixed_point_vs_linear"] = float( + np.max(np.abs(fixed.x - exact)) / max(float(np.max(np.abs(exact))), 1e-300)) + norm, raw = norms_of(cell.cfgs[0], topo, values, ref) + residual = float(d["residual"]) + cancels = _cancellation(cell, values, step.pre, step.state) + allowed = ((residual + cancels * floor) / (residual + floor) + if out["floor_reported"] and residual + floor > 0 else 1.0) + out["report"].update(floor=floor, cancellation=cancels) + + if out["floor_reported"]: + true = ref.residual(x, norm) + above = true - residual * (1.0 + 2.0 ** 8 * case.eps) + out["floor"] = max(0.0, above) / max(cancels * floor, 1e-300) + out["report"]["residual_true"] = true + + # CPL-087: the Jacobian at the returned state. + J = ref.jacobian(x) + out["report"].update(rho_true=cr.radius(J), rho_at_fixed_point=cr.radius( + ref.jacobian(fixed.x))) + linear.radius_scores(out, J, state_weights(ref, x), float(d["rho_spectral"]), case.eps, + case.base.seed) + + if out["spectral_usable"]: + dist = ref.distance(x, fixed, norm) + h = ref.nonlinearity(x, fixed, norm) + bound = float(d["spectral_error_bound"]) * allowed + out["report"].update(distance=dist, nonlinearity=h, + distance_over_bound=dist / bound if bound > 0 else math.inf) + out["near"] = h < 1.0 + if out["near"]: + # The bound of the linearisation at the returned iterate is + # short of the distance by at most ``1 / (1 - h)``. + reach = bound / (1.0 - h) + out["bound"] = (math.inf if math.isnan(bound) else + dist / reach if reach > 0 else (math.inf if dist > 0 else 0.0)) + + if out["gradient_usable"]: + true, column = ref.gradient_error(x, fixed, raw) + bound = float(d["gradient_relative_error_bound"]) * allowed + 64.0 * case.eps + out["gradient"] = math.inf if math.isnan(bound) else true / bound + out["report"].update(gradient_error=true, gradient_constant=( + None if column is None else ref.constant_names()[column])) + return out + + +# --------------------------------------------------------------------------- +# The strategy and the searches +# --------------------------------------------------------------------------- + +#: The curves drawn, in decades: a hundredth (all but linear over a +#: field's own size) to a hundred (a returned iterate a thousandth of its +#: field from the fixed point still sees the Jacobian move by a tenth). +CURVES = (-2.0, 2.0) + +THRESHOLD = linear.THRESHOLD +FLAG = linear.FLAG +SEARCHES = linear.SEARCHES +#: The least fraction of a hunt's examples with the flag set, and with a +#: reference (a fixed point Newton reached). +USABLE_FLOOR = 0.25 +REFERENCED_FLOOR = 0.9 + + +def cases(cells=ALL_CELLS, domain: linear.Domain = linear.CLAIMED, curves=CURVES): + """Draw a :class:`Case` on one of *cells*: the linear search's numbers + within *domain* and a curve.""" + return st.builds(Case, cell=st.sampled_from(tuple(cells)), base=linear.cases((0,), domain), + curve=st.floats(*curves).map(lambda x: 10.0 ** x)) + + +def search(name: str, *, cells=PER_PUSH_CELLS, domain: linear.Domain = linear.CLAIMED, + profile=None, fail: bool = True): + """Run search *name*; returns ``(report, {fraction name: value})``. + + The default profile is the per-push one, the same draws for every + score (an example is a step and a Newton solve; the four searches + share what the random phase draws).""" + profile = PER_PUSH if profile is None else profile + drawn = [] + + def score(case: Case): + drawn.append(case) + seen = observe(case) + return seen[name], seen["report"] + + report = targeted_search(cases(cells, domain), score, THRESHOLD[name], profile=profile, + label=name, fail=fail) + seen = [observe(c) for c in drawn] + count = max(len(seen), 1) + return report, dict(usable=sum(s[FLAG[name]] for s in seen) / count, + referenced=sum(s["referenced"] for s in seen) / count, + near=sum(s["near"] for s in seen) / count, + stepped=sum(s["stepped"] for s in seen) / count) diff --git a/tests/property/test_coupling_targeted_search.py b/tests/property/test_coupling_targeted_search.py index e524456f..c31e5069 100644 --- a/tests/property/test_coupling_targeted_search.py +++ b/tests/property/test_coupling_targeted_search.py @@ -237,9 +237,13 @@ def eps(self) -> float: return float(np.finfo(CELLS[self.cell].dtype).eps) -def values_of(case: Case) -> dict: - """The gains, biases, mappings and start of *case*, rounded to its dtype.""" - cell = CELLS[case.cell] +def values_of(case: Case, cell: "Cell | None" = None) -> dict: + """The gains, biases, mappings and start of *case*, rounded to its dtype. + + *cell*: the cell to draw on in place of ``CELLS[case.cell]`` (the + nonlinear cells of ``test_coupling_nonlinear_search.py`` draw the same + numbers on cells of their own).""" + cell = CELLS[case.cell] if cell is None else cell topo, cfgs = cell.topo, cell.cfgs members = topo.groups[0] rng = np.random.default_rng(case.seed) @@ -467,6 +471,41 @@ def moved(node=None, port=None, edge=None, entry=None): return worst +def radius_scores(out: dict, J: np.ndarray, wv: np.ndarray, rho_reported: float, eps: float, + seed: int) -> None: + """Write the two radius scores of a reported ``rho_spectral`` into + *out* (``"radius"`` and ``"radius_strict"``; the module docstring says + what each is), given the pass's Jacobian *J* at the state the report + describes and that state's weights *wv*. *eps*: the analysis dtype's + (the group's, at least float32).""" + A = (wv[:, None] * J) / wv[None, :] + rho = _radius(J) + if math.isfinite(rho_reported): + allowance = max(_radius_allowance(A, rho, eps, seed), 1e-300) + norm_A = float(np.linalg.norm(A, 2)) + off = abs(rho_reported - rho) if out["spectral_usable"] else 0.0 + resolved = np.linalg.matrix_rank(J) <= SPECTRAL_KRYLOV_STEPS + if not resolved: + # "An estimate otherwise, which spectral_usable reports": + # past eight independent scalars nothing is called exact. A + # settled estimate is read as good to the margin the flag + # tests it by, 5% of ``1 - rho``. + allowance += SPECTRAL_SETTLED_FRACTION * max(1.0 - rho, 0.0) + out["report"]["rank"] = int(np.linalg.matrix_rank(J)) + if norm_A > 0 and np.linalg.norm(A @ A.T - A.T @ A, 2) <= 1e-9 * norm_A ** 2: + # Normal in the norm's weights: an estimate "from below", + # whatever the flag. + off = max(off, rho_reported - rho) + # The statement a user can check: within the flag's own margin of + # the radius, wherever the flag is set and the group has no more + # scalars than the Krylov steps resolve. + if out["spectral_usable"] and resolved: + out["radius_strict"] = abs(rho_reported - rho) / max( + SPECTRAL_SETTLED_FRACTION * (1.0 - rho_reported), 1e-300) + out["radius"] = off / allowance + out["report"]["jacobian_norm"] = norm_A + + @functools.lru_cache(maxsize=4096) def observe(case: Case) -> dict: """One step of *case* and the four scores of what it reported.""" @@ -505,36 +544,8 @@ def observe(case: Case) -> dict: J = _pass_jacobian(model) wv = _state_weights(model, step.state) - A = (wv[:, None] * J) / wv[None, :] - rho = _radius(J) - rho_reported = float(d["rho_spectral"]) - out["report"]["rho_true"] = rho - # The analysis runs in the group's dtype (at least float32). - eps_analysis = case.eps - if math.isfinite(rho_reported): - allowance = max(_radius_allowance(A, rho, eps_analysis, case.seed), 1e-300) - norm_A = float(np.linalg.norm(A, 2)) - off = abs(rho_reported - rho) if out["spectral_usable"] else 0.0 - resolved = np.linalg.matrix_rank(J) <= SPECTRAL_KRYLOV_STEPS - if not resolved: - # "An estimate otherwise, which spectral_usable reports": - # past eight independent scalars nothing is called exact. A - # settled estimate is read as good to the margin the flag - # tests it by, 5% of ``1 - rho``. - allowance += SPECTRAL_SETTLED_FRACTION * max(1.0 - rho, 0.0) - out["report"]["rank"] = int(np.linalg.matrix_rank(J)) - if norm_A > 0 and np.linalg.norm(A @ A.T - A.T @ A, 2) <= 1e-9 * norm_A ** 2: - # Normal in the norm's weights: an estimate "from below", - # whatever the flag. - off = max(off, rho_reported - rho) - # The statement a user can check: within the flag's own margin of - # the radius, wherever the flag is set and the group has no more - # scalars than the Krylov steps resolve. - if out["spectral_usable"] and resolved: - out["radius_strict"] = abs(rho_reported - rho) / max( - SPECTRAL_SETTLED_FRACTION * (1.0 - rho_reported), 1e-300) - out["radius"] = off / allowance - out["report"]["jacobian_norm"] = norm_A + out["report"]["rho_true"] = _radius(J) + radius_scores(out, J, wv, float(d["rho_spectral"]), case.eps, case.seed) if out["spectral_usable"]: dist = model.returned_weight_distance(0, step.pre, step.state) diff --git a/tests/property/test_interface_reading_twins.py b/tests/property/test_interface_reading_twins.py new file mode 100644 index 00000000..217c2335 --- /dev/null +++ b/tests/property/test_interface_reading_twins.py @@ -0,0 +1,305 @@ +"""Twins of an edge-mapped graph whose interface reading is its own. + +``convergence_norm="interface"`` reads what each internal edge delivers: +the source field through the edge's mapping and then its transform. The +node-inlined twin of ``tests/property/test_differential_geometry_edges.py`` +moves a mapping into the *target* node, so its edges deliver the raw +source field and its interface norm is another norm: reports of the two +graphs under that norm can agree only loosely. The two twins here keep +the reading (``tests/property/geometry_graphs.py``): + +* :func:`~tests.property.geometry_graphs.transform_twin` writes a static + mapping as its edge's transform. Same state, same reading: **every + number of the report but the gradient bound is the edge-mapped graph's** + (measured bit for bit on jaxlib 0.11.0, under both schedules, Aitken + and IQN-ILS, in float32 and float64; held here to :data:`ULPS` of the + dtype's ``eps``, since they are two programs). The gradient bound + probes every parameter of the step, and the mapping's weights are + parameters there and constants of a transform here. +* :func:`~tests.property.geometry_graphs.relay_twin` moves the mapping + into a relay on the *source*, whose state field **is** the mapped + value: a plain edge delivers it, so the reading is the same entry for + entry, and a geometry the mapping reads can sit inside the relay. The + state has the relay's field in it, so what is equal is what the reading + decides -- every iterate of an unaccelerated solve, the residual, the + pass count, the verdict, ``rho_spectral`` where the spectrum is + resolved, the flags -- and ``spectral_error_bound`` is equal only to + :data:`RELAY_BOUND_GAP`: its factor is the norm of a resolvent + compressed onto a Krylov basis of the *state*, which is another basis + with the relay's field in it (measured 1.1% to 3.4% apart under + Gauss-Seidel). The float floor is not compared (13% apart: the norm + counts an inner product for a mapping it evaluates and none for a field + it reads as stored), nor the gradient bound (its norm is the raw source + fields, which are other fields here). Under Jacobi the relay doubles + the scalars the pass carries (18 for 9 on the two-body graph), past + what eight Krylov steps resolve, so the relay comparison is made under + Gauss-Seidel. + +**What this proves on today's tree** (no geometry): the library reads a +static mapped internal edge as the step delivers it, and both twins say +so; a fault seeded in the library's reading (the edge read without its +mapping) breaks both equalities (the PR that added this file measured +it; ``docs/developer_guide/testing_standards.md``). And that the relay +twin is the edge-mapped graph where it has a source-anchored *moving* +geometry: on the solve path, which 0.4.0 supports, the two step bit for +bit. + +**What waits** (:data:`~tests.property.geometry_graphs.RELAY_INTERFACE_CASES`, +in the harness's phase-1 block): the interface norm over a geometry edge +is refused at compile in 0.4.0. Each such case asserts the refusal and +that its relay twin compiles and reports today; when +``DIAGNOSTICS_READ_GEOMETRY`` is set the same test compares the two +reports as the static cases are compared. +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("JAX_PLATFORMS", "cpu") + +import numpy as np +import pytest + +from tests.property import coupled_graphs as cg +from tests.property import geometry_graphs as gg +from tests.property import test_coupling_targeted_search as linear + +KEY = "F+P" +#: Float ``eps`` of the group's dtype two programs of the same arithmetic +#: may put between a reported number and its twin's. Zero was measured; +#: an Arnoldi process turns one rounding of a product into a few of the +#: radius, and 2**10 of them is still five decades under what a reading +#: without its mapping moves (the seeded fault: 0.6% to 60%). +ULPS = 2.0 ** 10 +#: How far the relay twin's ``spectral_error_bound`` factor may be from +#: the edge-mapped graph's (module docstring): three times the largest gap +#: measured. +RELAY_BOUND_GAP = 0.10 +#: Every number of a report but the gradient bound's two keys. +NUMBERS = ("residual", "amplification", "error_estimate", "gradient_error_estimate", + "rho_spectral", "spectral_error_bound") +FLAGS = ("converged", "ratio_usable", "spectral_usable", "precision_limited") + +_P_HOLDS = dict(down="target", up="source") +_SOURCES = dict(down="source", up="source") + + +def _interface(dtype: str, **knobs) -> gg.Case: + return gg.case(f"{dtype} {knobs}", dtype=dtype, adv=0.3, **_P_HOLDS, + group=dict(convergence_norm="interface", rtol=1e-6, diagnostics=True, + **knobs)) + + +#: A solve stopped at a cap of three passes (a residual far above its +#: float floor, so it is a number and not a rounding), per push; the +#: accelerated ones in the slow lane. +STATIC_PER_PUSH = [ + _interface("float32", max_iterations=3), + _interface("float64", max_iterations=3, iteration_mode="jacobi"), +] +STATIC_SLOW = [ + _interface("float64", max_iterations=3), + _interface("float32", max_iterations=3, iteration_mode="jacobi"), + _interface("float64", max_iterations=6, acceleration="aitken"), + _interface("float32", max_iterations=6, acceleration="iqn-ils"), +] +RELAY_PER_PUSH = [_interface("float64", max_iterations=3)] +RELAY_SLOW = [_interface("float32", max_iterations=3)] + + +def _report(gm) -> dict: + """``coupling_diagnostics()`` of the two-body group, with the floor it used.""" + report = dict(gm.coupling_diagnostics()[KEY]) + report["floor"] = linear._reported_floor(gm, KEY, cg.group_meta(gm, KEY), report) # noqa: SLF001 + return report + + +def _close(a: float, b: float, eps: float, what) -> None: + a, b = float(a), float(b) + assert np.isfinite(a) and np.isfinite(b), (what, a, b) + assert abs(a - b) <= ULPS * eps * max(abs(a), abs(b)), (what, a, b) + + +def _same_states(c: gg.Case, a: dict, b: dict, step: int) -> None: + """Every field the edge-mapped graph holds is the twin's, within rounding.""" + for name in a: + for field, x in a[name].items(): + y = b[name][field] + assert x.dtype == y.dtype and x.shape == y.shape, (name, field) + scale = float(max(np.max(np.abs(x)), np.max(np.abs(y)))) + assert float(np.max(np.abs(x.astype(np.float64) - y.astype(np.float64)))) <= ( + 8 * float(np.finfo(x.dtype).eps) * scale), (c.label, step, name, field) + + +def compare(c: gg.Case, edge, twin, *, relay: bool, steps: int = 3) -> None: + """The reports of *edge* and *twin* after each of *steps* steps.""" + eps = float(np.finfo(np.dtype(c.dtype)).eps) + for step in range(1, steps + 1): + edge.step() + twin.step() + _same_states(c, gg.snapshot(edge), gg.snapshot(twin), step) + ra, rb = _report(edge), _report(twin) + where = (c.label, step) + assert "not_usable_reason" not in ra and "not_usable_reason" not in rb, where + assert int(ra["iterations"]) == int(rb["iterations"]), (where, ra, rb) + for flag in FLAGS: + assert bool(ra[flag]) is bool(rb[flag]), (where, flag, ra, rb) + assert bool(ra["spectral_usable"]), ("premise: a usable report", where, ra) + if relay: + for name in ("residual", "amplification", "error_estimate", "rho_spectral"): + _close(ra[name], rb[name], eps, (where, name)) + fa = float(ra["spectral_error_bound"]) / (float(ra["residual"]) + ra["floor"]) + fb = float(rb["spectral_error_bound"]) / (float(rb["residual"]) + rb["floor"]) + assert abs(fa - fb) <= RELAY_BOUND_GAP * max(fa, fb), (where, fa, fb) + continue + for name in NUMBERS + ("floor",): + _close(ra[name], rb[name], eps, (where, name)) + + +def assert_reports_as_its_transform_twin(c: gg.Case) -> None: + with gg.x64(c.needs_x64): + static = gg.static_twin(c) + assert all(e.geometry is None for e in static.edges) and sum( + e.mapping is not None for e in static.edges) == 2, "premise: two static mapped edges" + compare(c, gg.build(static), gg.build(gg.transform_twin(gg.static_twin(c))), relay=False) + + +def assert_reports_as_its_relay_twin(c: gg.Case) -> None: + with gg.x64(c.needs_x64): + compare(c, gg.build(gg.static_twin(c)), gg.build(gg.relay_twin(gg.static_twin(c))), + relay=True) + + +@pytest.mark.parametrize("c", STATIC_PER_PUSH, ids=repr) +def test_a_static_mapped_edge_under_the_interface_norm_reports_as_its_transform_twin(c): + """Per push; slow sibling + :func:`test_every_static_mapped_edge_reports_as_its_transform_twin`.""" + assert_reports_as_its_transform_twin(c) + + +# Slow: two graphs with diagnostics compiled per case. +# Per push: tests/property/test_interface_reading_twins.py::test_a_static_mapped_edge_under_the_interface_norm_reports_as_its_transform_twin +@pytest.mark.slow +@pytest.mark.parametrize("c", STATIC_SLOW, ids=repr) +def test_every_static_mapped_edge_reports_as_its_transform_twin(c): + assert_reports_as_its_transform_twin(c) + + +@pytest.mark.parametrize("c", RELAY_PER_PUSH, ids=repr) +def test_a_static_mapped_edge_under_the_interface_norm_reports_as_its_relay_twin(c): + """Per push; slow sibling + :func:`test_every_static_mapped_edge_reports_as_its_relay_twin`.""" + assert_reports_as_its_relay_twin(c) + + +# Slow: two graphs with diagnostics compiled per case. +# Per push: tests/property/test_interface_reading_twins.py::test_a_static_mapped_edge_under_the_interface_norm_reports_as_its_relay_twin +@pytest.mark.slow +@pytest.mark.parametrize("c", RELAY_SLOW, ids=repr) +def test_every_static_mapped_edge_reports_as_its_relay_twin(c): + assert_reports_as_its_relay_twin(c) + + +def test_the_relay_twin_delivers_the_mapped_value_over_plain_edges_only(): + """The twin's structure: the same nodes and groups, no mapping on any + edge, and each formerly mapped edge reading the relay's field, whose + initial value is the mapping of the source field beside it.""" + c = STATIC_PER_PUSH[0] + with gg.x64(c.needs_x64): + static, twin = gg.static_twin(c), gg.relay_twin(gg.static_twin(c)) + assert [nd.name for nd in twin.nodes] == [nd.name for nd in static.nodes] + assert twin.groups == static.groups and len(twin.edges) == len(static.edges) + for e, t in zip(static.edges, twin.edges): + assert t.mapping is None and t.geometry is None, t + assert (t.src, t.dst, t.tf, t.transform, t.additive) == ( + e.src, e.dst, e.tf, e.transform, e.additive) + assert (t.sf == e.sf) is (e.mapping is None), (e, t) + if e.mapping is not None: + state = twin.node(e.src).initial_state() + np.testing.assert_array_equal( + np.asarray(state[t.sf]), np.asarray(e.mapping.apply(state[e.sf], None))) + assert all(e.mapping is None and e.transform is not None + for e in gg.transform_twin(gg.static_twin(c)).edges) + + +def test_a_target_anchored_geometry_has_no_relay_twin(): + """The relay would read the target's geometry over an internal edge of + its own, which the interface norm would read: refused, not built.""" + c = gg.case("target anchors", adv=0.3, down="target", up="target") + with pytest.raises(NotImplementedError, match="target-anchored geometry has no relay twin"): + gg.relay_twin(gg.two_body(c)) + + +# --------------------------------------------------------------------------- +# The relay twin of a moving geometry, on the solve path +# --------------------------------------------------------------------------- + +#: Source-anchored geometry edges whose geometry moves with the input +#: (inside a group, with the iterate), under the norms 0.4.0 supports. +MOVING = [ + gg.case("plain step", adv=0.3, **_SOURCES), + gg.case("group, Gauss-Seidel", adv=0.3, **_SOURCES, + group=dict(max_iterations=200, tolerance=1e-5)), + pytest.param(gg.case("plain step, multilinear", kind="multilinear", adv=0.3, **_SOURCES), + marks=pytest.mark.slow), + pytest.param(gg.case("group, Jacobi, three passes", adv=0.3, **_SOURCES, + group=dict(max_iterations=3, iteration_mode="jacobi")), + marks=pytest.mark.slow), + pytest.param(gg.case("group, multilinear, float64", kind="multilinear", dtype="float64", + adv=0.3, **_SOURCES, group=dict(max_iterations=3)), + marks=pytest.mark.slow), + pytest.param(gg.case("group, P before F", order=("P", "F"), adv=0.3, **_SOURCES, + group=dict(max_iterations=3)), marks=pytest.mark.slow), +] + + +# Slow (the marked cases): two graphs compiled per case. +# Per push: tests/property/test_interface_reading_twins.py::test_the_relay_twin_of_a_moving_source_anchored_geometry_steps_as_the_edge_mapped_graph +@pytest.mark.parametrize("c", MOVING, ids=repr) +def test_the_relay_twin_of_a_moving_source_anchored_geometry_steps_as_the_edge_mapped_graph(c): + """The relay reads the geometry field beside the value, both as its + update has just produced them -- the time level a source-anchored + geometry edge reads ("from the dict the value is read from"). Every + field the edge-mapped graph holds, after every step, within rounding + (zero was measured), with the geometry moved and the pass count the + same.""" + with gg.x64(c.needs_x64): + edge, twin = gg.build(gg.two_body(c)), gg.build(gg.relay_twin(gg.two_body(c))) + before = gg.snapshot(edge) + for step in range(1, c.steps + 1): + edge.step() + twin.step() + _same_states(c, gg.snapshot(edge), gg.snapshot(twin), step) + if c.group is not None: + assert int(edge.coupling_diagnostics()[KEY]["iterations"]) == int( + twin.coupling_diagnostics()[KEY]["iterations"]), (c.label, step) + after = gg.snapshot(edge) + field = "pos" if c.kind == "multilinear" else "A" + assert np.any(after["F"][field] != before["F"][field]), "premise: the geometry moved" + + +# --------------------------------------------------------------------------- +# What waits for diagnostics that read a geometry +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("c", gg.RELAY_INTERFACE_CASES, ids=repr) +def test_the_interface_norm_over_a_geometry_edge_reports_as_its_relay_twin(c): + """PHASE 1 (see ``geometry_graphs``): the interface norm over a + geometry edge is refused at compile, and the relay twin -- plain edges + only -- compiles and reports a usable spectrum today: the report the + edge-mapped graph is held to once ``DIAGNOSTICS_READ_GEOMETRY`` is + set, by the comparison of the static cases.""" + with gg.x64(c.needs_x64): + twin = gg.build(gg.relay_twin(gg.two_body(c))) + if not gg.DIAGNOSTICS_READ_GEOMETRY: + keys = [e.key for e in gg.build(gg.two_body(c), compile=False).edges + if e.geometry is not None] + gg.assert_interface_norm_refused(lambda: gg.build(gg.two_body(c)), keys) + twin.step() + report = twin.coupling_diagnostics()[KEY] + assert bool(report["spectral_usable"]) and np.isfinite( + float(report["spectral_error_bound"])), report + return + compare(c, gg.build(gg.two_body(c)), twin, relay=True) From 826463ab45462e9b5bf5e0e61a314bb6e8b15696 Mon Sep 17 00:00:00 2001 From: Nicholas Ehsan Roy Date: Wed, 7 Oct 2026 05:21:17 +0200 Subject: [PATCH 2/4] wip(tests): nonlinear search pins, twins restructured, seeded-fault table Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013UkCde7g23gTziUjYAvnKD [skip ci] --- tests/duration_allowlist.txt | 4 + tests/property/coupling_reference.py | 105 ++++- tests/property/geometry_graphs.py | 50 +++ .../test_coupling_nonlinear_search.py | 411 +++++++++++++++++- .../property/test_interface_reading_twins.py | 120 +++-- 5 files changed, 616 insertions(+), 74 deletions(-) diff --git a/tests/duration_allowlist.txt b/tests/duration_allowlist.txt index a75e4369..1877ec08 100644 --- a/tests/duration_allowlist.txt +++ b/tests/duration_allowlist.txt @@ -49,3 +49,7 @@ tests/property/test_sysid_targeted_search.py::test_no_wrong_fit_among_the_fixed_ tests/compliance/test_claims_inventories.py::test_every_cited_test_exists_and_runs_on_every_push[coupling_claims.yaml] # kept: the guard that every test a coupling claim cites exists and is collected per push; it must run on every push, and its time is a subprocess collection that grows with the number of cited tests (20.3 s and 20.8 s on CI) tests/property/test_coupling_targeted_search.py::test_a_usable_error_bound_is_never_below_the_distance_per_push # kept: the per-push witness of the slow hunt over coupling_diagnostics()'s four numbers; its seconds are the four compiled cells (solver="ift" with diagnostics) that the other three searches, the seeds and the pins then share tests/core/test_coupling_claims_in_every_domain.py::test_the_claim_holds_for_each_member_of_a_vmapped_step[CPL-006] # kept: the per-push witness of the vmap cell of CPL-006; its time is tracing a batched coupled step with diagnostics on, and it measured 20.2 s and 21.9 s on CI against the 20 s line (2026-10-06), the second with the ninth Jacobian-vector product of the spectral self-check +tests/property/test_coupling_nonlinear_search.py::test_the_reference_reproduces_the_closed_form_on_a_linear_cell # kept: the per-push check that the numerical coupling reference is the closed form (the reference is what the nonlinear cells and the later geometry cells are scored against); its seconds are one cell of the linear search compiled (which that search then does not pay) and its twin +tests/property/test_coupling_nonlinear_search.py::test_the_nonlinear_fixed_point_is_the_linear_one_and_the_jacobian_moves # kept: carries the compile of the one per-push nonlinear cell (solver="ift" with diagnostics) and of its float64 twin, which the four per-push nonlinear searches then share; the coupling search has no other nonlinear coverage per push +tests/property/test_interface_reading_twins.py::test_a_static_mapped_edge_under_the_interface_norm_reports_a_usable_spectrum[float32 {'max_iterations': 3}] # kept: carries the compile of the edge-mapped graph the per-push interface-reading comparison is made against (a group with diagnostics under the interface norm) +tests/property/test_interface_reading_twins.py::test_a_static_mapped_edge_under_the_interface_norm_reports_as_its_transform_twin[float32 {'max_iterations': 3}] # kept: the only per-push comparison of coupling_diagnostics() under the interface norm with a twin that has the same reading (a reading without its mapping moves the residual by 8.8% here and nothing else per push sees it); its seconds are the twin's compile diff --git a/tests/property/coupling_reference.py b/tests/property/coupling_reference.py index 86b9a80b..e42d0495 100644 --- a/tests/property/coupling_reference.py +++ b/tests/property/coupling_reference.py @@ -50,6 +50,8 @@ read in) or for a scalar loss ``w . x``; * :meth:`PassReference.residual`: the exact residual ``P(x) - x`` of a returned state in the group's norm; +* :meth:`PassReference.finite_difference_gap`: a Jacobian-vector product + (the pass's, or a reading's) against a central difference; * :meth:`PassReference.nonlinearity`: ``||(I - J(x))^{-1} (J_mean - J(x))||`` with ``J_mean`` the mean Jacobian over the segment from the fixed point to ``x`` -- the exact factor by which a linear bound taken @@ -449,34 +451,58 @@ def implicit_derivatives(self, x) -> np.ndarray: x = np.asarray(x, np.float64) return np.linalg.solve(np.eye(self.size) - self.jacobian(x), self.sensitivities(x)) - def gradient_error(self, returned, fixed: FixedPoint, norm: Norm, *, - loss: Optional[Sequence[float]] = None): - """``(worst, column)``: the worst relative error, over the scalar - constants, of the implicit derivative taken at the returned iterate - against the one at the fixed point. - - With *loss* (a cotangent ``w``: the loss is ``w . x``) the error is - of the scalar ``d loss / d c``, ``|w . (g_k - g*)| / |w . g_k|``. - Without, of the vector ``d x / d c`` in *norm* (linear fields: - selections of the iterate) at the returned state's weights, the - form ``gradient_relative_error_bound`` documents. + def gradient_errors(self, returned, fixed: FixedPoint, norm: Norm, *, + loss: Optional[Sequence[float]] = None) -> tuple: + """``(miss, at_iterate, at_fixed_point)``, one entry per scalar + constant: the size of the implicit derivative taken at the returned + iterate, of the one at the fixed point, and of their difference. + + With *loss* (a cotangent ``w``: the loss is ``w . x``) the sizes + are of the scalar ``d loss / d c``. Without, of the vector ``d x / + d c`` in *norm* (linear fields: selections of the iterate) at the + returned state's weights, the form + ``gradient_relative_error_bound`` documents. """ x = np.asarray(returned, np.float64) g_k, g_star = self.implicit_derivatives(x), self.implicit_derivatives(fixed.x) + if loss is not None: + w = np.asarray(loss, np.float64) + return np.abs(w @ (g_k - g_star)), np.abs(w @ g_k), np.abs(w @ g_star) weights = norm.weights(x) - worst, column = 0.0, None zero = np.zeros(self.size) - for c in range(g_k.shape[1]): - if loss is not None: - w = np.asarray(loss, np.float64) - size, miss = abs(float(w @ g_k[:, c])), abs(float(w @ (g_k[:, c] - g_star[:, c]))) - else: - size = norm.of_difference(g_k[:, c], zero, weights) - miss = norm.of_difference(g_k[:, c], g_star[:, c], weights) - if size > 0 and miss / size > worst: - worst, column = miss / size, c + columns = range(g_k.shape[1]) + return (np.asarray([norm.of_difference(g_k[:, c], g_star[:, c], weights) for c in columns]), + np.asarray([norm.of_difference(g_k[:, c], zero, weights) for c in columns]), + np.asarray([norm.of_difference(g_star[:, c], zero, weights) for c in columns])) + + def gradient_error(self, returned, fixed: FixedPoint, norm: Norm, *, + loss: Optional[Sequence[float]] = None, + columns: Optional[Sequence[bool]] = None): + """``(worst, column)``: the worst relative error ``|g_k - g*| / + |g_k|`` of :meth:`gradient_errors` over the scalar constants + (those *columns* marks, by the order of :meth:`constant_names`; + all of them by default).""" + miss, at_iterate, _at_fixed = self.gradient_errors(returned, fixed, norm, loss=loss) + worst, column = 0.0, None + for c in range(len(miss)): + if at_iterate[c] <= 0 or (columns is not None and not columns[c]): + continue + if miss[c] / at_iterate[c] > worst: + worst, column = float(miss[c] / at_iterate[c]), c return worst, column + def gradient_resolution(self, returned, norm: Norm) -> float: + """What this reference cannot resolve of a relative gradient error: + 2**10 float64 ``eps`` times the resolvent's norm at the returned + iterate and the spread of *norm*'s weights (a derivative is a + dense solve, and a field a thousandth of another weighs a + thousand times as much).""" + x = np.asarray(returned, np.float64) + weights = [w for w in norm.weights(x) if w > 0] + spread = max(weights) / min(weights) if weights else 1.0 + resolvent = float(np.linalg.norm(np.linalg.inv(np.eye(self.size) - self.jacobian(x)), 2)) + return 2.0 ** 10 * EPS64 * spread * max(resolvent, 1.0) + def nonlinearity(self, returned, fixed: FixedPoint, norm: Norm, *, nodes: int = 8) -> float: """``h = ||T (I - J(x))^{-1} (J_mean - J(x)) T^+||`` with ``T`` the weighted (linear) reading of *norm* and ``J_mean`` the mean of @@ -499,6 +525,43 @@ def nonlinearity(self, returned, fixed: FixedPoint, norm: Norm, *, nodes: int = E = np.linalg.solve(np.eye(self.size) - J, mean - J) return float(np.linalg.norm(T @ E @ np.linalg.pinv(T), 2)) + def finite_difference_gap(self, x, reading: Optional[Callable] = None, *, + directions: int = 4, step: float = 1e-6, seed: int = 0) -> float: + """How far a Jacobian-vector product is from a central difference, + relative to the product: the worst over *directions* seeded + directions of ``|J v - (f(x + h v) - f(x - h v)) / 2h| / |J v|``. + + ``f`` is the pass ``P`` and ``J`` its :meth:`jacobian`; with + *reading* (a JAX function of a flat iterate, differentiated here by + ``jax.jvp``) it is that reading and its own product -- the check a + run-time self-test of a reading's Jacobian would make. ``h`` is + *step* times the iterate's size, so the difference is good to + about ``step ** 2`` on a smooth map. + """ + x = np.asarray(x, np.float64) + rng = np.random.default_rng(seed) + h = step * max(float(np.max(np.abs(x))), 1.0) + J = None if reading is not None else self.jacobian(x) + worst = 0.0 + for _ in range(directions): + v = rng.normal(size=x.shape) + if reading is None: + product = J @ v + difference = (self.apply(x + h * v) - self.apply(x - h * v)) / (2.0 * h) + else: + with x64(): + product = np.asarray(jax.jvp( + lambda z: jnp.concatenate([jnp.ravel(f) for f in reading(z)]), + (jnp.asarray(x),), (jnp.asarray(v),))[1]) + difference = (np.concatenate([np.ravel(np.asarray(f)) for f in reading( + jnp.asarray(x + h * v))]) - np.concatenate([ + np.ravel(np.asarray(f)) for f in reading(jnp.asarray(x - h * v))]) + ) / (2.0 * h) + size = float(np.linalg.norm(product)) + if size > 0: + worst = max(worst, float(np.linalg.norm(product - difference)) / size) + return worst + def reading_matrix(self, norm: Norm, returned) -> np.ndarray: """The weighted reading of a *linear* norm as a matrix: its rows applied to a flat iterate are ``w fields(x)``.""" diff --git a/tests/property/geometry_graphs.py b/tests/property/geometry_graphs.py index fc9410fd..e6531f59 100644 --- a/tests/property/geometry_graphs.py +++ b/tests/property/geometry_graphs.py @@ -44,6 +44,56 @@ :func:`two_body` builds the graph the differential and gradient tests draw from: a grid-side body ``F`` and a point-side body ``P`` joined in both directions, in or out of a coupling group, at one rate or two. + +**The twins with the edge-mapped graph's interface reading** +(:func:`transform_twin`, :func:`relay_twin`; their docstrings say what +each holds equal) and **the faults later work is held to.** Diagnostics +that read a geometry (the interface norm as a criterion with a geometry +edge; the bounds in that reading with the geometry term) do not exist in +0.4.0. The table lists the faults that work must be caught on, the +instrument expected to catch each, and -- where the fault can be seeded +on today's tree in an analogous static-mapping or solve-path form -- the +signal measured when it was (scratch copies of ``src/``, jaxlib 0.11.0, +CPU; the instruments are ``tests/property/test_interface_reading_twins.py`` +and, for the numerical reference, ``tests/property/coupling_reference.py`` +and ``tests/property/test_coupling_nonlinear_search.py``): + +== ========================================= ================================== ========================================== +# fault instrument analogue today, and the measured signal +== ========================================= ================================== ========================================== +1 the geometry term dropped from the the numerical reference: its none (a static mapping has no geometry + reading's Jacobian Jacobian is ``jacfwd`` of the term). The reference's own sensitivity: + pass with the geometry field in a radius reported 10% low fails the + the iterate (radius, bound and per-push radius search; an error bound + gradient scores); the relay halved, or a gradient bound a twentieth, + twin's ``rho_spectral`` fails a per-push seed (mutants R1 to R3). +2 the geometry read at the wrong time the relay twin (residual, pass solve path: a group's source-anchored + level inside the reading count, ``rho_spectral``); the geometry read from the pre-step state + time-level reference of in place of the dict its value is read + ``test_geometry_time_levels.py`` from. The relay twin's step-for-step + for the solve it must agree with equality fails after two passes: a + field 2.6e-4 apart, 5.4e-7 allowed, zero + unmutated (mutant F2-F3). A converged + solve does not see it: at a fixed point + the iterate and the pass agree. +3 the geometry taken from the iterate the relay twin, whose relay the same seeded fault, which is the + instead of the pre-step state (or the computes the mapped value from reverse (the pre-step state where the + reverse) the fields of one update; the iterate was due): F2-F3. + numerical reference +4 the geometry's own rounding left out of the numerical reference's floor not seeded. The relay twin does not + the precision floor score (the exact residual of the judge a floor: its own is 13% from the + returned state above the reported edge-mapped graph's by construction + one, over the reported floor) on (a field read as stored against a + cells with a float32 geometry mapping the norm evaluates). +5 the mapping dropped from the interface both twins seeded (mutant F5): the transform twin's + reading and the relay twin's equalities both + fail on the per-push case (residuals + 8.8% apart, 1.2e-4 allowed). +6 the reading's Jacobian-vector product ``PassReference a reading with a term its tangent does + disagreeing with a finite difference .finite_difference_gap`` (the not see: gap 1e-3 or more where the + (a later run-time self-check) check itself, for the pass or honest reading's is under 1e-7 + for any reading written in JAX) (``test_the_reference_s_jacobian_is_...``). +== ========================================= ================================== ========================================== """ from __future__ import annotations diff --git a/tests/property/test_coupling_nonlinear_search.py b/tests/property/test_coupling_nonlinear_search.py index 8c0ca7ba..b05ddfe3 100644 --- a/tests/property/test_coupling_nonlinear_search.py +++ b/tests/property/test_coupling_nonlinear_search.py @@ -58,11 +58,19 @@ the returned state"): as the linear search scores them, against the reference's Jacobian **at the returned iterate**. 3. *gradient* (CPL-093, CPL-095): the true relative error of the implicit - derivative taken at the returned iterate, the worst over every scalar - constant (the nonlinearity's ``c`` and ``s`` included), over - ``gradient_relative_error_bound`` where ``gradient_bound_usable`` -- - the flag that certifies the Newton-Kantorovich check, which on these - cells is not trivially passed. + derivative taken at the returned iterate, the worst over the scalar + constants, over ``gradient_relative_error_bound`` where + ``gradient_bound_usable`` -- the flag that certifies the + Newton-Kantorovich check, which on these cells is not trivially + passed. Two scores, by the constant: ``"gradient"`` over the gains, + biases and mapping weights, and ``"gradient_vanishing"`` over the + centre ``c`` and the curve ``s`` of every nonlinearity, which the + fixed point does not respond to (``phi`` depends on neither at ``u = + c``) and every other iterate does, so that ``|g_k - g*|`` is about + ``|g_k|``: the relative error is of order one however close the + iterate is, and a usable bound below that is short for that constant. + Both allow what the float64 reference itself cannot resolve + (:meth:`~tests.property.coupling_reference.PassReference.gradient_resolution`). 4. *floor* (CPL-097, CPL-100): the exact residual of the returned state above the reported one, over the reported floor. @@ -83,6 +91,7 @@ os.environ.setdefault("JAX_PLATFORMS", "cpu") +import jax import jax.numpy as jnp import numpy as np import pytest @@ -531,6 +540,15 @@ def run_once(built: ct.Built, values: dict) -> ct.Step: {0: cg.group_meta(gm, key)}) +def does_not_move_the_fixed_point(constant: str) -> bool: + """Whether *constant* (``node.leaf[i]``, as the reference names it) is + a nonlinearity's centre ``c`` or curve ``s``: at the fixed point + ``u = c`` and ``phi`` depends on neither, so the fixed point's + derivative with respect to it is zero (to the rounding of ``c``).""" + leaf = constant.split(".", 1)[1] + return leaf[0] in "cs" and leaf[1].isdigit() + + def leaves_float_range(cell: Cell, values: dict, ref: cr.PassReference) -> bool: """Whether plain passes from the drawn start leave *cell*'s float range before its cap, on a cell with a quasi-Newton acceleration. @@ -565,15 +583,16 @@ def observe(case: Case) -> dict: ref = bound_reference(ref, values, twin) if leaves_float_range(cell, values, ref): # Not stepped (see ``leaves_float_range``): nothing is scored. - return dict(bound=0.0, radius=0.0, radius_strict=0.0, gradient=0.0, floor=0.0, + return dict(bound=0.0, radius=0.0, radius_strict=0.0, gradient=0.0, gradient_vanishing=0.0, + floor=0.0, spectral_usable=False, gradient_usable=False, floor_reported=False, referenced=False, near=False, stepped=False, report={}) with precision(cell.dtype == "float64"): step = run_once(built, values) d = dict(step.reports[0]) floor = linear._reported_floor(built.gm, topo.group_key(0), step.metas[0], d) # noqa: SLF001 - out = dict(bound=0.0, radius=0.0, radius_strict=0.0, gradient=0.0, floor=0.0, - spectral_usable=bool(d["spectral_usable"]), + out = dict(bound=0.0, radius=0.0, radius_strict=0.0, gradient=0.0, gradient_vanishing=0.0, + floor=0.0, spectral_usable=bool(d["spectral_usable"]), gradient_usable=bool(d["gradient_bound_usable"]), floor_reported=math.isfinite(floor), referenced=False, near=False, stepped=True, report={k: d[k] for k in ("iterations", "converged", "residual", "rho_spectral", @@ -627,11 +646,15 @@ def observe(case: Case) -> dict: dist / reach if reach > 0 else (math.inf if dist > 0 else 0.0)) if out["gradient_usable"]: - true, column = ref.gradient_error(x, fixed, raw) - bound = float(d["gradient_relative_error_bound"]) * allowed + 64.0 * case.eps - out["gradient"] = math.inf if math.isnan(bound) else true / bound - out["report"].update(gradient_error=true, gradient_constant=( - None if column is None else ref.constant_names()[column])) + bound = (float(d["gradient_relative_error_bound"]) * allowed + 64.0 * case.eps + + ref.gradient_resolution(x, raw)) + names = ref.constant_names() + for score, vanishing in (("gradient", False), ("gradient_vanishing", True)): + true, column = ref.gradient_error(x, fixed, raw, columns=[ + does_not_move_the_fixed_point(n) is vanishing for n in names]) + out[score] = math.inf if math.isnan(bound) else true / bound + out["report"].update({f"{score}_error": true, f"{score}_constant": ( + None if column is None else ref.constant_names()[column])}) return out @@ -644,8 +667,8 @@ def observe(case: Case) -> dict: #: field from the fixed point still sees the Jacobian move by a tenth). CURVES = (-2.0, 2.0) -THRESHOLD = linear.THRESHOLD -FLAG = linear.FLAG +THRESHOLD = {**linear.THRESHOLD, "gradient_vanishing": linear.THRESHOLD["gradient"]} +FLAG = {**linear.FLAG, "gradient_vanishing": "gradient_usable"} SEARCHES = linear.SEARCHES #: The least fraction of a hunt's examples with the flag set, and with a #: reference (a fixed point Newton reached). @@ -683,3 +706,361 @@ def score(case: Case): referenced=sum(s["referenced"] for s in seen) / count, near=sum(s["near"] for s in seen) / count, stepped=sum(s["stepped"] for s in seen) / count) + + +# --------------------------------------------------------------------------- +# The reference, validated +# --------------------------------------------------------------------------- + + +def _per_push_draws(cells) -> list: + """The cases the linear search's four per-push searches draw, on *cells*.""" + drawn = [] + for name in linear.SEARCHES: + def score(case, name=name): + drawn.append(case) + return linear.scorer(name)(case) + targeted_search(linear.cases(cells, linear.CLAIMED), score, math.inf, + profile=linear.EVERY_PUSH.seeded(sorted(linear.THRESHOLD).index(name)), + label=name) + return list(dict.fromkeys(drawn)) + + +def assert_the_reference_reproduces_the_closed_form(drawn) -> dict: + """Every answer of the reference within :data:`ALLOWED` of the closed + form on each of *drawn*; the worst miss per answer.""" + worst, compared = {}, 0 + for case in drawn: + seen = against_the_closed_form(case) + misses = closed_form_misses(seen) + compared += bool(misses) + over = {k: v for k, v in misses.items() if v > 1.0} + assert not over, f"{case}: {over} ({seen})" + for k, v in misses.items(): + worst[k] = max(worst.get(k, 0.0), v) + assert compared >= 0.9 * len(drawn), (compared, len(drawn)) + return worst + + +def test_the_reference_reproduces_the_closed_form_on_a_linear_cell(): + """Per push: the ring of three float64 scalars under Jacobi and the + seed shapes on it. Slow sibling: + :func:`test_the_reference_reproduces_the_closed_form_on_every_per_push_draw`.""" + drawn = [] + # Forty derandomised draws of the linear search's generator (not + # scored here: the slow sibling takes the search's own draws). + targeted_search(linear.cases((0,), linear.CLAIMED), lambda c: (drawn.append(c) or 0.0, None), + math.inf, profile=dataclasses.replace(PER_PUSH, max_examples=40)) + drawn += [c for c in linear.SEEDS.values() if c.cell == 0] + assert len(set(drawn)) >= 20, len(set(drawn)) + assert_the_reference_reproduces_the_closed_form(drawn) + + +# Slow: a twin compiled for each per-push cell of the linear search and for +# its multi-rate cell, and every one of its per-push draws (about 600). +# Per push: tests/property/test_coupling_nonlinear_search.py::test_the_reference_reproduces_the_closed_form_on_a_linear_cell +@pytest.mark.slow +def test_the_reference_reproduces_the_closed_form_on_every_per_push_draw(): + """Every per-push draw of the linear search (the five per-push cells: + both dtypes, both schedules, the l2 and interface norms, a mapped + ring, twelve scalars), the seed shapes, and the same number of draws + on the multi-rate cell (a sub-cycled member, linear interpolation).""" + drawn = _per_push_draws(linear.PER_PUSH_CELLS) + list(linear.SEEDS.values()) + multirate = len(linear.PER_PUSH_CELLS) + assert linear.CELLS[multirate].structure.endswith("multirate") + drawn += _per_push_draws((multirate,))[:150] + worst = assert_the_reference_reproduces_the_closed_form(drawn) + print(f"{len(drawn)} draws; the worst miss over what is allowed: " + + ", ".join(f"{k} {v:.3g}" for k, v in sorted(worst.items()))) + + +@pytest.mark.parametrize("index", [0, pytest.param(2, marks=pytest.mark.slow), + pytest.param(len(linear.PER_PUSH_CELLS), + marks=pytest.mark.slow)]) +def test_several_passes_of_an_iterating_twin_are_compositions_of_the_single_pass(index): + """The single-pass branch of the step (``max_iterations=1``) runs the + pass an iterating group iterates: three passes of a twin that iterates + are three compositions of the reference's map, to float64 rounding. + (The ring under Jacobi per push; the mapped ring under Gauss-Seidel + and the multi-rate ring in the slow lane.)""" + cell = linear.CELLS[index] + values = linear.values_of(dataclasses.replace(_PROBE, cell=index)) + built_twin, ref = _linear_twin(index) + ref = bound_reference(ref, values, built_twin) + knobs = {**cr.twin_knobs(cell.knobs), "max_iterations": 3, "predictor": "none", + "tolerance": 0.0} + with cr.x64(): + several = ct.build(cell.topo, knobs, dtype="float64", mapping_kind=cell.mapping_kind) + ct.set_initial(several, values) + start = ref.flat({m: {"x": values["nodes"][m]["x0"]} for m in cell.topo.groups[0]}) + assert cr.passes_compose(ref, several.gm, 3, start) <= 2.0 ** 6 + # The premise: the pass depends on the iterate, apart from the pre-step state. + assert np.any(ref.apply(start) != ref.apply(start + 1.0)) + + +def test_the_reference_refuses_a_twin_it_cannot_read(): + """One group, one pass, the linear predictor: anything else is refused + when the reference is built, not answered wrongly.""" + cell = linear.CELLS[0] + with cr.x64(): + iterating = ct.build(cell.topo, {"max_iterations": 3}, dtype="float64") + with pytest.raises(AssertionError, match="one pass under the linear predictor"): + cr.PassReference.of(iterating.gm) + fresh = ct.build(cell.topo, cr.twin_knobs(cell.knobs), dtype="float64") + # A fresh graph: every field zero after its step, so the slot does + # not say which member is where. + with pytest.raises(AssertionError, match="does not determine the layout"): + cr.PassReference.of(fresh.gm) + + +# --------------------------------------------------------------------------- +# The searches on the nonlinear cells +# --------------------------------------------------------------------------- + + +def test_the_nonlinear_fixed_point_is_the_linear_one_and_the_jacobian_moves(): + """The premise of the cells, on the per-push one: the reference's fixed + point is the linear closed form's (the nonlinearity is centred on + it), and away from it the Jacobian is another matrix -- the part no + closed form here knows.""" + case = Case(0, linear.Case(0, 7, 0.6, True, 1.0, 0.0, 1.0, 0, 1e-1, 0), 30.0) + seen = observe(case) + assert seen["stepped"] and seen["referenced"], seen + report = seen["report"] + assert report["fixed_point_vs_linear"] <= 2.0 ** 10 * case.eps ** 2 + 2.0 ** 10 * EPS64, report + assert abs(report["rho_at_fixed_point"] - 0.6) <= 1e-5, report + assert abs(report["rho_true"] - report["rho_at_fixed_point"]) >= 1e-3, report + assert report["nonlinearity"] >= 1e-3, report + + +#: Examples on the per-push cell on which a number is close to its limit +#: (found by a 500-example search of that cell), so that the per-push lane +#: sees a bound that became too small without drawing for it: +#: ``{name: (case, score, the least the score must still be)}``. +SEEDS = { + # Five Jacobi passes from a whole field away: the distance is 0.88 of + # ``spectral_error_bound``. + "the-error-bound-an-eighth-above-the-distance": ( + Case(0, linear.Case(0, 65535, 0.05, True, 1.0, 0.0, 1.0, -3, 1.0, 2), 1.0), + "bound", 0.75), + # A curve of 21 over a field: the radius is 0.399 at the returned + # iterate and 0.333 at the fixed point (``h`` = 0.14), and the + # derivative with respect to a gain is 24% off, a fourteenth of + # ``gradient_relative_error_bound``. + "the-jacobian-a-fifth-from-the-fixed-point-s": ( + Case(0, linear.Case(0, 1, 1.0 / 3.0, False, 1.0, 0.0, 1.0, -1, 1.0, 0), + 21.544346900318832), "gradient", 0.05), +} + + +@pytest.mark.parametrize("seed", sorted(SEEDS)) +def test_every_score_holds_on_the_nonlinear_seed_shapes(seed): + case, score, least = SEEDS[seed] + seen = observe(case) + assert seen["stepped"] and seen["referenced"] and seen[FLAG[score]], seen + over = {name: seen[name] for name in SEARCHES + ("radius_strict",) + if seen[name] > THRESHOLD[name]} + assert not over, f"{seed}: {over} ({seen['report']})" + assert seen[score] >= least, ( + f"premise: {score} is {seen[score]!r} on {seed}, no longer near its limit " + f"({seen['report']})") + + +def _held(name: str, fractions: dict) -> None: + assert fractions["referenced"] >= REFERENCED_FLOOR * fractions["stepped"], (name, fractions) + assert fractions["usable"] > 0, f"{name}: no example had its flag set ({fractions})" + + +def test_a_usable_error_bound_reaches_the_distance_on_the_nonlinear_cell_per_push(): + _report, fractions = search("bound") + _held("bound", fractions) + assert fractions["near"] > 0, fractions + + +def test_a_settled_spectral_radius_is_the_radius_at_the_returned_iterate_per_push(): + _report, fractions = search("radius") + _held("radius", fractions) + + +def test_a_usable_gradient_bound_is_never_below_the_error_on_the_nonlinear_cell_per_push(): + _report, fractions = search("gradient") + _held("gradient", fractions) + + +def test_the_floor_covers_what_the_reported_residual_misses_on_the_nonlinear_cell_per_push(): + _report, fractions = search("floor") + _held("floor", fractions) + + +# Slow: 115 random examples a search on each of six blocks of seven cells, +# each cell a compile of the graph with its diagnostics and of its twin; +# the blocks outermost, so the searches share a block's compiled graphs. +# Per push: tests/property/test_coupling_nonlinear_search.py::test_a_usable_error_bound_reaches_the_distance_on_the_nonlinear_cell_per_push +@pytest.mark.slow +@pytest.mark.parametrize("block,name", [(b, n) for b in range(len(BLOCKS)) + for n in SEARCHES + ("radius_strict",)]) +def test_the_hunt_finds_no_number_on_the_wrong_side_of_a_nonlinear_group(block, name): + """Not shrunk: the example that fails is reported as drawn (a shrink + here is thousands of steps).""" + profile = dataclasses.replace(SLOW, max_examples=115).seeded(1000 + block, shrink=False) + report, fractions = search(name, cells=BLOCKS[block], profile=profile) + print(f"{name}, block {block}: worst {report}; {fractions}") + assert fractions["referenced"] >= REFERENCED_FLOOR * fractions["stepped"], fractions + assert fractions["usable"] >= USABLE_FLOOR, ( + f"{name}, block {block}: only {fractions['usable']:.2f} of the examples had the flag " + f"set (floor {USABLE_FLOOR})") + + +def test_the_reference_s_jacobian_is_the_central_difference_of_its_pass(): + """The dense Jacobian against a central difference of the pass map, on + the per-push nonlinear cell away from its fixed point, and the same + check on a reading (the interface norm's, written in JAX): the + self-check a reading's Jacobian-vector product can be held to.""" + case = Case(0, linear.Case(0, 7, 0.6, True, 1.0, 0.0, 1.0, 0, 1e-1, 0), 30.0) + cell = CELLS[0] + values = values_of(case) + _built_graph, twin, ref = _built(0) + ref = bound_reference(ref, values, twin) + x = ref.flat({m: {"x": values["nodes"][m]["x0"]} for m in cell.topo.names}) + assert ref.finite_difference_gap(x) <= 1e-7 + H = {i: jnp.asarray(values["H"][i], jnp.float64) for i in values["H"]} + + def reading(z): + out = [] + for i, e in enumerate(cell.topo.edges): + src = ref.field(z, e.src, "x") + out.append(ct.TRANSFORM_FACTORS[e.transform] * (H[i] @ src if e.mapped else src)) + return out + + assert ref.finite_difference_gap(x, reading) <= 1e-7 + # A reading whose product drops a term (one the tangent does not see) + # is caught by the same check. + + def dropped(z): + return [f + jnp.sum(jax.lax.stop_gradient(z) ** 2) for f in reading(z)] + + assert ref.finite_difference_gap(x, dropped) >= 1e-3 + + +# --------------------------------------------------------------------------- +# What the search found +# --------------------------------------------------------------------------- + +#: The findings of the hunt on this tree, each an example as drawn (the +#: hunt does not shrink). Strict: a fix turns each green. ``EDGE``: at +#: the edge of what the claim is made on, as the reason says; a ``CORE`` +#: finding would be inside it (the hunt found none). +KNOWN = { + # A two-member float32 ring with saturating gains, started on its + # fixed point: one pass, residual 0, precision_limited. The curve + # ``s`` of a nonlinearity moves the fixed point by nothing (the + # derivative with respect to it is zero there) and the returned + # iterate, a float32 rounding away, by a little: relative error 2.7, + # ``gradient_relative_error_bound`` 1.4e-6 with the flag set. + "CPL-093-EDGE-a-constant-the-fixed-point-does-not-respond-to-at-a-stalled-start": ( + Case(1, linear.Case(0, 0, 0.05, False, 1.0, 0.0, 1.0, 0, 0.0, 0), 1.0), + "gradient_vanishing", + "FINDING (EDGE, CPL-093): for a constant whose derivative is zero at the fixed point " + "the relative gradient error is of order one at any other iterate; at a start stalled " + "on its float floor the usable bound reads 1.4e-6"), + # The same constant on a float64 hub converged in four passes from a + # start a whole field away: the bound reads 1.09 for a relative error + # of 1.134 (4% short). + "CPL-093-EDGE-a-constant-the-fixed-point-does-not-respond-to-four-percent-short": ( + Case(26, linear.Case(0, 2619, 0.05, False, 1.0, 0.0, 1.0, 3, 1.0, 3), 10.0), + "gradient_vanishing", + "FINDING (EDGE, CPL-093): for a constant whose derivative is zero at the fixed point " + "the usable bound reads 1.09 for a relative error of 1.13"), +} + + +# Slow: each pin compiles its cell (a graph with diagnostics and a twin). +# Per push: tests/property/test_coupling_nonlinear_search.py::test_a_usable_gradient_bound_is_never_below_the_error_on_the_nonlinear_cell_per_push +@pytest.mark.slow +@pytest.mark.parametrize("case,score", [ + pytest.param(case, score, marks=pytest.mark.xfail(strict=True, reason=reason)) + for case, score, reason in KNOWN.values()], ids=list(KNOWN)) +def test_a_known_defect_the_nonlinear_search_reached_is_fixed(case, score): + seen = observe(case) + assert seen["stepped"] and seen["referenced"], seen + assert seen[FLAG[score]], f"the flag is no longer set: {seen['report']}" + assert seen[score] <= THRESHOLD[score], ( + f"{score} is {seen[score]!r}, over {THRESHOLD[score]!r}: {seen['report']}") + + +#: The example the hunt stopped on: the fan-out hub with products of two +#: fields in float64, Jacobi under IQN-ILS, cap 120, non-normal gains, +#: started a whole field from its fixed point. Plain passes from that +#: start leave float64 range on the eighth. +_NEVER_RETURNS = Case(30, linear.Case(0, 3397, 0.6568215254870227, True, 1.0, 0.0, 1.0, 6, 1.0, 6), + 1.0) +_STEP_IN_A_SUBPROCESS = """ +import sys +from tests.property import test_coupling_nonlinear_search as nl +from tests.property.sysid_transform_grid import precision +case = nl._NEVER_RETURNS +cell = nl.CELLS[case.cell] +from tests.property import coupled_graphs as cg +with precision(True): + built = nl.build(cell, cg.live_knobs({**cell.knobs, "acceleration": sys.argv[1]}), cell.dtype) + values = nl.values_of(case) + nl.run_once(built, nl.values_of(nl.Case(case.cell, nl._PROBE, 1.0))) # compiled, and returns + print("COMPILED", flush=True) + step = nl.run_once(built, values) + print("RETURNED", int(step.reports[0]["iterations"]), bool(step.reports[0]["converged"]), + flush=True) +""" + + +def _steps_in_a_subprocess(acceleration: str, seconds: float) -> str: + """The output of one step of :data:`_NEVER_RETURNS` under + *acceleration*, in a process of its own given *seconds* after its + compile (a call that does not return cannot be scored in this one).""" + import pathlib # noqa: PLC0415 + import subprocess # noqa: PLC0415 + import sys # noqa: PLC0415 + + root = pathlib.Path(__file__).resolve().parents[2] + proc = subprocess.Popen([sys.executable, "-c", _STEP_IN_A_SUBPROCESS, acceleration], + stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, + cwd=root, env={**os.environ, "PYTHONPATH": os.pathsep.join( + [str(root / "src"), str(root), + os.environ.get("PYTHONPATH", "")])}) + try: + assert proc.stdout is not None + assert proc.stdout.readline().strip() == "COMPILED", "the benign step did not return" + try: + out, _err = proc.communicate(timeout=seconds) + except subprocess.TimeoutExpired: + return "TIMEOUT" + return out.strip() + finally: + proc.kill() + proc.wait() + + +# Slow: a subprocess that compiles the cell and is then given 30 s for a +# step that takes 0.05 s where it returns. +# Per push: tests/property/test_coupling_nonlinear_search.py::test_a_usable_error_bound_reaches_the_distance_on_the_nonlinear_cell_per_push +@pytest.mark.slow +def test_a_group_that_leaves_float_range_without_a_quasi_newton_acceleration_returns(): + """The control of the pin below: the same graph and start with no + acceleration returns from its step at the cap, not converged.""" + assert _steps_in_a_subprocess("none", 30.0) == "RETURNED 120 False" + + +# Slow: as the control above. +# Per push: tests/property/test_coupling_nonlinear_search.py::test_a_usable_error_bound_reaches_the_distance_on_the_nonlinear_cell_per_push +@pytest.mark.slow +@pytest.mark.parametrize("acceleration", ["iqn-ils", "iqn-imvj"]) +@pytest.mark.xfail(strict=True, reason=( + "FINDING (EDGE, CPL-053/CPL-084): a group under acceleration='iqn-ils' or 'iqn-imvj' whose " + "iterate becomes non-finite before max_iterations never returns from step() (measured: no " + "return in 240 s; 0.05 s for a step of the same compiled graph that converges)")) +def test_a_group_that_leaves_float_range_under_iqn_still_returns(acceleration): + """``converged`` False "means ... the group hit max_iterations" + (CPL-053) and ``residual`` is "inf on a non-finite state" (CPL-084): + both say the step returns. With diagnostics on or off, under the l2 + norm as under the mixed one; at a cap of 12, before the iterate is + non-finite, it returns.""" + assert _steps_in_a_subprocess(acceleration, 30.0).startswith("RETURNED") diff --git a/tests/property/test_interface_reading_twins.py b/tests/property/test_interface_reading_twins.py index 217c2335..7e91ddc9 100644 --- a/tests/property/test_interface_reading_twins.py +++ b/tests/property/test_interface_reading_twins.py @@ -38,8 +38,8 @@ **What this proves on today's tree** (no geometry): the library reads a static mapped internal edge as the step delivers it, and both twins say so; a fault seeded in the library's reading (the edge read without its -mapping) breaks both equalities (the PR that added this file measured -it; ``docs/developer_guide/testing_standards.md``). And that the relay +mapping) breaks both equalities (measured when this file was added: the +two residuals 8.8% apart; ``docs/developer_guide/testing_standards.md``). And that the relay twin is the edge-mapped graph where it has a source-anchored *moving* geometry: on the solve path, which 0.4.0 supports, the two step bit for bit. @@ -48,12 +48,13 @@ in the harness's phase-1 block): the interface norm over a geometry edge is refused at compile in 0.4.0. Each such case asserts the refusal and that its relay twin compiles and reports today; when -``DIAGNOSTICS_READ_GEOMETRY`` is set the same test compares the two +``DIAGNOSTICS_READ_GEOMETRY`` is set the slow test compares the two reports as the static cases are compared. """ from __future__ import annotations +import functools import os os.environ.setdefault("JAX_PLATFORMS", "cpu") @@ -69,8 +70,9 @@ #: Float ``eps`` of the group's dtype two programs of the same arithmetic #: may put between a reported number and its twin's. Zero was measured; #: an Arnoldi process turns one rounding of a product into a few of the -#: radius, and 2**10 of them is still five decades under what a reading -#: without its mapping moves (the seeded fault: 0.6% to 60%). +#: radius, and 2**10 of them (1.2e-4 in float32) is still far under what +#: a reading without its mapping moves (the seeded fault: the residual +#: 8.8% apart on the per-push case). ULPS = 2.0 ** 10 #: How far the relay twin's ``spectral_error_bound`` factor may be from #: the edge-mapped graph's (module docstring): three times the largest gap @@ -92,20 +94,33 @@ def _interface(dtype: str, **knobs) -> gg.Case: #: A solve stopped at a cap of three passes (a residual far above its -#: float floor, so it is a number and not a rounding), per push; the -#: accelerated ones in the slow lane. -STATIC_PER_PUSH = [ - _interface("float32", max_iterations=3), - _interface("float64", max_iterations=3, iteration_mode="jacobi"), -] +#: float floor, so it is a number and not a rounding). Per push: one +#: case, whose edge-mapped graph both twins are compared with (a compile +#: of a group with diagnostics is seconds; :func:`_edge` keeps it). The +#: other dtype, the other schedule and the accelerations: the slow lane. +STATIC_PER_PUSH = [_interface("float32", max_iterations=3)] STATIC_SLOW = [ _interface("float64", max_iterations=3), + _interface("float64", max_iterations=3, iteration_mode="jacobi"), _interface("float32", max_iterations=3, iteration_mode="jacobi"), _interface("float64", max_iterations=6, acceleration="aitken"), _interface("float32", max_iterations=6, acceleration="iqn-ils"), ] -RELAY_PER_PUSH = [_interface("float64", max_iterations=3)] -RELAY_SLOW = [_interface("float32", max_iterations=3)] +#: The relay twin's comparison is the slow lane's (a third compile); per +#: push the relay twin is held to the edge-mapped graph on the solve path +#: (:data:`MOVING`) and the seeded reading fault to the transform twin. +RELAY_SLOW = [_interface("float32", max_iterations=3), _interface("float64", max_iterations=3)] + + +@functools.lru_cache(maxsize=2) +def _edge(c: gg.Case): + """The compiled edge-mapped graph of static case *c* (two static + mapped internal edges).""" + with gg.x64(c.needs_x64): + static = gg.static_twin(c) + assert all(e.geometry is None for e in static.edges) and sum( + e.mapping is not None for e in static.edges) == 2, "premise: two static mapped edges" + return gg.build(static) def _report(gm) -> dict: @@ -135,6 +150,8 @@ def _same_states(c: gg.Case, a: dict, b: dict, step: int) -> None: def compare(c: gg.Case, edge, twin, *, relay: bool, steps: int = 3) -> None: """The reports of *edge* and *twin* after each of *steps* steps.""" eps = float(np.finfo(np.dtype(c.dtype)).eps) + edge.reset_state() + twin.reset_state() for step in range(1, steps + 1): edge.step() twin.step() @@ -159,16 +176,26 @@ def compare(c: gg.Case, edge, twin, *, relay: bool, steps: int = 3) -> None: def assert_reports_as_its_transform_twin(c: gg.Case) -> None: with gg.x64(c.needs_x64): - static = gg.static_twin(c) - assert all(e.geometry is None for e in static.edges) and sum( - e.mapping is not None for e in static.edges) == 2, "premise: two static mapped edges" - compare(c, gg.build(static), gg.build(gg.transform_twin(gg.static_twin(c))), relay=False) + compare(c, _edge(c), gg.build(gg.transform_twin(gg.static_twin(c))), relay=False) def assert_reports_as_its_relay_twin(c: gg.Case) -> None: with gg.x64(c.needs_x64): - compare(c, gg.build(gg.static_twin(c)), gg.build(gg.relay_twin(gg.static_twin(c))), - relay=True) + compare(c, _edge(c), gg.build(gg.relay_twin(gg.static_twin(c))), relay=True) + + +@pytest.mark.parametrize("c", STATIC_PER_PUSH, ids=repr) +def test_a_static_mapped_edge_under_the_interface_norm_reports_a_usable_spectrum(c): + """The premise of the comparisons, and the compile of the edge-mapped + graph they share (a test's seconds here are a compile each).""" + with gg.x64(c.needs_x64): + edge = _edge(c) + edge.reset_state() + edge.step() + report = _report(edge) + assert bool(report["spectral_usable"]) and not bool(report["converged"]), report + assert float(report["residual"]) > 2.0 ** 10 * report["floor"], ( + "premise: a residual far above its float floor", report) @pytest.mark.parametrize("c", STATIC_PER_PUSH, ids=repr) @@ -186,18 +213,11 @@ def test_every_static_mapped_edge_reports_as_its_transform_twin(c): assert_reports_as_its_transform_twin(c) -@pytest.mark.parametrize("c", RELAY_PER_PUSH, ids=repr) -def test_a_static_mapped_edge_under_the_interface_norm_reports_as_its_relay_twin(c): - """Per push; slow sibling - :func:`test_every_static_mapped_edge_reports_as_its_relay_twin`.""" - assert_reports_as_its_relay_twin(c) - - # Slow: two graphs with diagnostics compiled per case. -# Per push: tests/property/test_interface_reading_twins.py::test_a_static_mapped_edge_under_the_interface_norm_reports_as_its_relay_twin +# Per push: tests/property/test_interface_reading_twins.py::test_a_static_mapped_edge_under_the_interface_norm_reports_as_its_transform_twin @pytest.mark.slow @pytest.mark.parametrize("c", RELAY_SLOW, ids=repr) -def test_every_static_mapped_edge_reports_as_its_relay_twin(c): +def test_a_static_mapped_edge_under_the_interface_norm_reports_as_its_relay_twin(c): assert_reports_as_its_relay_twin(c) @@ -239,8 +259,14 @@ def test_a_target_anchored_geometry_has_no_relay_twin(): #: (inside a group, with the iterate), under the norms 0.4.0 supports. MOVING = [ gg.case("plain step", adv=0.3, **_SOURCES), - gg.case("group, Gauss-Seidel", adv=0.3, **_SOURCES, - group=dict(max_iterations=200, tolerance=1e-5)), + # Two passes, not a converged solve: at a fixed point the iterate and + # the pass agree, and a geometry read from the wrong one of them is the + # same number. + gg.case("group, Gauss-Seidel, two passes", adv=0.3, **_SOURCES, + group=dict(max_iterations=2)), + pytest.param(gg.case("group, Gauss-Seidel, converged", adv=0.3, **_SOURCES, + group=dict(max_iterations=200, tolerance=1e-5)), + marks=pytest.mark.slow), pytest.param(gg.case("plain step, multilinear", kind="multilinear", adv=0.3, **_SOURCES), marks=pytest.mark.slow), pytest.param(gg.case("group, Jacobi, three passes", adv=0.3, **_SOURCES, @@ -284,19 +310,37 @@ def test_the_relay_twin_of_a_moving_source_anchored_geometry_steps_as_the_edge_m # --------------------------------------------------------------------------- +@pytest.mark.parametrize("c", gg.RELAY_INTERFACE_CASES, ids=repr) +def test_the_interface_norm_over_a_geometry_edge_is_refused_and_its_relay_twin_builds(c): + """PHASE 1 (see ``geometry_graphs``), per push: the interface norm over + a geometry edge is refused at compile, and the relay twin -- plain + edges only -- is accepted. Slow sibling: + :func:`test_the_interface_norm_over_a_geometry_edge_reports_as_its_relay_twin`.""" + with gg.x64(c.needs_x64): + keys = [e.key for e in gg.build(gg.two_body(c), compile=False).edges + if e.geometry is not None] + assert len(keys) == 2, keys + if gg.DIAGNOSTICS_READ_GEOMETRY: + gg.build(gg.two_body(c)) # accepted, once diagnostics read a geometry + else: + gg.assert_interface_norm_refused(lambda: gg.build(gg.two_body(c)), keys) + twin = gg.build(gg.relay_twin(gg.two_body(c))) + assert all(e.geometry is None and e.mapping is None for e in twin.edges) + + +# Slow: a relay twin with diagnostics compiled per case (and, once +# diagnostics read a geometry, the edge-mapped graph beside it). +# Per push: tests/property/test_interface_reading_twins.py::test_the_interface_norm_over_a_geometry_edge_is_refused_and_its_relay_twin_builds +@pytest.mark.slow @pytest.mark.parametrize("c", gg.RELAY_INTERFACE_CASES, ids=repr) def test_the_interface_norm_over_a_geometry_edge_reports_as_its_relay_twin(c): - """PHASE 1 (see ``geometry_graphs``): the interface norm over a - geometry edge is refused at compile, and the relay twin -- plain edges - only -- compiles and reports a usable spectrum today: the report the - edge-mapped graph is held to once ``DIAGNOSTICS_READ_GEOMETRY`` is - set, by the comparison of the static cases.""" + """PHASE 1 (see ``geometry_graphs``): the relay twin of a refused case + reports a usable spectrum today -- the report the edge-mapped graph is + held to once ``DIAGNOSTICS_READ_GEOMETRY`` is set, by the comparison of + the static cases.""" with gg.x64(c.needs_x64): twin = gg.build(gg.relay_twin(gg.two_body(c))) if not gg.DIAGNOSTICS_READ_GEOMETRY: - keys = [e.key for e in gg.build(gg.two_body(c), compile=False).edges - if e.geometry is not None] - gg.assert_interface_norm_refused(lambda: gg.build(gg.two_body(c)), keys) twin.step() report = twin.coupling_diagnostics()[KEY] assert bool(report["spectral_usable"]) and np.isfinite( From ebe02a29d2264db5b03c2ec0c479b20fc634fe6c Mon Sep 17 00:00:00 2001 From: Nicholas Ehsan Roy Date: Wed, 7 Oct 2026 05:37:40 +0200 Subject: [PATCH 3/4] test(coupling): a numerical reference, nonlinear search cells and interface-reading twins Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013UkCde7g23gTziUjYAvnKD --- docs/developer_guide/testing_standards.md | 29 +++++++++++++++++++ .../test_coupling_nonlinear_search.py | 23 ++++++++++----- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/docs/developer_guide/testing_standards.md b/docs/developer_guide/testing_standards.md index c7938c48..ff75400d 100644 --- a/docs/developer_guide/testing_standards.md +++ b/docs/developer_guide/testing_standards.md @@ -407,6 +407,35 @@ A score is 0 where its flag is False: an unusable number is never a failure. A s **To add a shape:** if it is a number (a magnitude, a ratio, a start), add a field to `Case`, apply it in `values_of`, draw it in `cases` behind a `Domain` field, and add an example to `SEEDS`. If it is structure (an edge, a node size, a knob), add a topology to `search_topologies()` or a row to `KNOBS`; `_cells()` picks it up for the slow profile, and it joins `_FIRST` only if it should cost every push a compile (about 3 s). Then run the hunt over the new domain on three seeds (`search(name, cells=..., domain=..., profile=SLOW.seeded(k, shrink=False), fail=False)` gives the examples-to-find). An over-threshold example that is a defect goes into `KNOWN` with the claim it breaks; one that the claim does not promise goes into the score as an allowance with its reason. Prove a new score can fail with seeded faults in a scratch copy of `src/` (a halved factor, a dropped term, a flag forced True), each of which the per-push profile must find. +### A numerical reference, nonlinear cells, and twins with the same interface reading + +The coupling search above scores against a closed form, which exists only for linear relays. Three instruments extend it to groups that have none: a nonlinear node today, and a mapping that reads a moving geometry once the diagnostics read one. + +**The numerical reference** (`tests/property/coupling_reference.py`, `PassReference`) computes, in float64, what the closed form gave: the group's fixed point, the dense Jacobian of one pass at any iterate, and from them the spectral radius, the distance of a returned iterate in the norm the report states its bound in, the exact residual, and the error of the implicit gradient for every scalar constant. + +- **How the pass map is obtained.** The caller builds an *x64 twin* of the graph under test: the same nodes, edges, schedule and sub-cycling in float64, its group at `max_iterations=1` with no acceleration and `predictor="linear"` (`twin_knobs`). One step of the twin is one pass, differentiated straight through. The pass starts from the predictor's extrapolation `2 pred_0 - pred_1`, which the step keeps in three `_meta` slots of the state; writing the iterate `x` into both gives `P(x; pre)`: one pass from the iterate `x`, each member integrating from the pre-step state. The map is the graph's own compiled step called on a state, so nothing of the estimator is called, and `jax.jacfwd` of it holds every dependence the pass has on the iterate (transforms, mappings, a geometry field a member holds). The slots' layout is found, not assumed: the twin is stepped once on a drawn example and the order of members whose flattened fields equal the stored `pred_0` is taken. +- **What it catches:** anything by which the estimator's numbers disagree with the map the solve iterates (a term of the pass missing from the estimator's Jacobian, a reading at another time level than the pass's, a Jacobian taken at another state than the returned one, a bound in another norm). **What it does not:** a defect in the pass itself, which the reference inherits (the closed-form and time-level references hold the pass); a difference between the single-pass branch and the pass an iterating group runs (checked separately: three passes of an iterating twin are three compositions of the reference's map, to float64 rounding); a node that branches on its dtype; a graph with more than one coupling group. +- **Validated against the closed form** on every per-push draw of the linear search (600 draws on its five per-push cells, and 150 on the multi-rate cell, in the slow lane; forty draws on one cell per push). Worst differences measured: the radius 1.0e-15; the Jacobian 1.6 float64 `eps` of its largest entry; the fixed point 5.7 `eps` of the field times the resolvent's norm; the distance 2.4e-12 of itself; the gradient error 1.9e-12. `ALLOWED` in `tests/property/test_coupling_nonlinear_search.py` holds each to at most 2**12 roundings. +- **Cost:** per twin, its build and two compiles (the pass with its Jacobian; the constants' sensitivities), together under two seconds for a dozen scalars; per example, a few dozen jitted calls (about 20 ms). The seconds of a cell are the compile of the graph under test with its diagnostics. + +**The nonlinear cells** (`tests/property/test_coupling_nonlinear_search.py`) are structures of the linear search whose members pass each port's input through a nonlinearity centred on a parameter `c` before its gain: a saturating gain `c + tanh(s (u - c)) / s`, a quadratic term `u + s (u - c)^2`, or the product of two fields' deviations. Each has `phi(c) = c` and `phi'(c) = 1`, so with `c` the input at the linear group's fixed point the nonlinear group has the same fixed point and the same Jacobian *there*: the numbers of a case are drawn exactly as the linear search draws them, plus the curve `s max|c|` (0.01 to 100), and the linear closed form checks the reference's fixed point on every example. Away from the fixed point, where a capped solve returns, the Jacobian is another matrix, and only the reference knows it. + +| Score | What changes for a map that is not affine | +|---|---| +| error bound (CPL-088) | the claim is for a linear `F` and "asymptotic" otherwise, so the distance is divided by the bound times `1 / (1 - h)`, with `h = ‖(I − J(x))⁻¹ (J_mean − J(x))‖` measured by the reference (`J_mean` the mean Jacobian between the fixed point and the returned iterate). `x − x* = (I − J_mean)⁻¹ (x − F(x))` exactly, so a bound that is right for the linearisation at the returned iterate is short by at most that factor: a score over one is a wrong number, not a nonlinear map. Unscored where `h ≥ 1` | +| spectral radius (CPL-087) | scored against the reference's Jacobian **at the returned iterate**, which is what the claim says and which differs from the fixed point's here | +| gradient bound (CPL-093, CPL-095) | two scores: over the gains, biases and mapping weights (`gradient`), and over each nonlinearity's `c` and `s` (`gradient_vanishing`), which the fixed point does not respond to, so that the relative error is of order one at any other iterate | +| floor (CPL-097, CPL-100) | the exact residual is the reference's | + +Per push: one cell (`tri`, three members with one nonlinearity each, float32, five Jacobi passes), twenty derandomised examples a score, and two seed examples on which a number is within an eighth of its limit. Slow: 115 random examples a score on each of six blocks of seven cells (every structure, each nonlinearity, both dtypes, every configuration), not shrunk (a shrink is thousands of steps). An example that plain passes carry out of float range is not stepped on a cell with a quasi-Newton acceleration (see the pinned finding in the module). To add a nonlinearity, add a branch to `phi` with `phi(c) = c` and `phi'(c) = 1`; to add a structure, add it to `STRUCTURES` (single-rate: a sub-cycled member reads its input interpolated, and `c` would not be its input at the fixed point). + +**The twins with the same interface reading** (`tests/property/geometry_graphs.py`, `tests/property/test_interface_reading_twins.py`). `convergence_norm="interface"` reads what each internal edge delivers. The node-inlined twin of a geometry edge moves the mapping into the *target*, so its edges deliver the raw source field and its interface norm is another norm. Two twins keep the reading: + +- `transform_twin` writes a static mapping as its edge's transform: the same state and the same reading, so every number of the report but the gradient bound is the edge-mapped graph's (bit for bit as measured; held to 2**10 `eps`). The gradient bound probes the step's parameters, and the weights are constants of a transform. +- `relay_twin` moves the mapping into a relay fused with the *source* node, whose state field is the mapped value; a plain edge delivers it. The reading is the same entry for entry, and a geometry the mapping reads sits inside the relay (source anchors only: a target-anchored geometry would reach the relay over an internal edge the norm would read too, and so would a relay that is a node of its own). The state has the relay's field in it, so what is equal is what the reading decides: every iterate of an unaccelerated solve, the residual, the pass count, the verdict, `rho_spectral` where the spectrum is resolved, the flags. `spectral_error_bound`'s factor is 1.1% to 3.4% apart (it is the norm of a resolvent compressed onto a Krylov basis of the state), the float floor 13% apart, and the gradient bound is in another norm; under Jacobi the relay doubles the scalars a pass carries, so the comparison is made under Gauss-Seidel. + +On a static mapped edge both twins report as the edge-mapped graph does, and a fault seeded in the library's reading (the edge read without its mapping) puts the two residuals 8.8% apart. On the solve path the relay twin of a *moving* source-anchored geometry steps as the edge-mapped graph bit for bit, and a geometry read from the pre-step state in place of the iterate breaks that after two passes (a converged solve does not see it: at a fixed point the two agree). The interface norm over a geometry edge is refused in 0.4.0; `RELAY_INTERFACE_CASES`, in the harness's phase-1 block, are the cases whose reports are compared with their relay twins' once `DIAGNOSTICS_READ_GEOMETRY` is set. The harness module's docstring lists the six faults that later work is held to, the instrument for each and the signal measured where the fault can be seeded today. + ## Test time budget Every push runs the default lane (everything not marked diff --git a/tests/property/test_coupling_nonlinear_search.py b/tests/property/test_coupling_nonlinear_search.py index b05ddfe3..8dd9758c 100644 --- a/tests/property/test_coupling_nonlinear_search.py +++ b/tests/property/test_coupling_nonlinear_search.py @@ -586,7 +586,7 @@ def observe(case: Case) -> dict: return dict(bound=0.0, radius=0.0, radius_strict=0.0, gradient=0.0, gradient_vanishing=0.0, floor=0.0, spectral_usable=False, gradient_usable=False, floor_reported=False, - referenced=False, near=False, stepped=False, report={}) + referenced=False, near=False, stepped=False, finite=False, report={}) with precision(cell.dtype == "float64"): step = run_once(built, values) d = dict(step.reports[0]) @@ -595,6 +595,7 @@ def observe(case: Case) -> dict: floor=0.0, spectral_usable=bool(d["spectral_usable"]), gradient_usable=bool(d["gradient_bound_usable"]), floor_reported=math.isfinite(floor), referenced=False, near=False, stepped=True, + finite=False, report={k: d[k] for k in ("iterations", "converged", "residual", "rho_spectral", "spectral_error_bound", "spectral_usable", "gradient_relative_error_bound", @@ -602,6 +603,7 @@ def observe(case: Case) -> dict: finite = all(np.all(np.isfinite(s["x"])) for s in step.state.values()) if not finite or not math.isfinite(d["residual"]): return out + out["finite"] = True x = ref.flat(step.state) fixed = ref.fixed_point(x) out["report"].update(reference_ulps=fixed.ulps, reference_steps=(fixed.picard, fixed.newton)) @@ -670,10 +672,14 @@ def observe(case: Case) -> dict: THRESHOLD = {**linear.THRESHOLD, "gradient_vanishing": linear.THRESHOLD["gradient"]} FLAG = {**linear.FLAG, "gradient_vanishing": "gradient_usable"} SEARCHES = linear.SEARCHES -#: The least fraction of a hunt's examples with the flag set, and with a -#: reference (a fixed point Newton reached). +#: The least fraction of a hunt's examples with the flag set (measured: +#: 0.64 to 1.00 by block and score; the linear search holds a half, and a +#: hunt here climbs towards starts that diverge), and the least fraction of +#: the examples that returned a finite state for which the reference found +#: a fixed point (measured: 0.93 to 1.00; a capped solve from a start a +#: whole field away can return where Newton reaches none). USABLE_FLOOR = 0.25 -REFERENCED_FLOOR = 0.9 +REFERENCED_FLOOR = 0.75 def cases(cells=ALL_CELLS, domain: linear.Domain = linear.CLAIMED, curves=CURVES): @@ -705,7 +711,8 @@ def score(case: Case): return report, dict(usable=sum(s[FLAG[name]] for s in seen) / count, referenced=sum(s["referenced"] for s in seen) / count, near=sum(s["near"] for s in seen) / count, - stepped=sum(s["stepped"] for s in seen) / count) + stepped=sum(s["stepped"] for s in seen) / count, + finite=sum(s["finite"] for s in seen) / count) # --------------------------------------------------------------------------- @@ -774,6 +781,8 @@ def test_the_reference_reproduces_the_closed_form_on_every_per_push_draw(): + ", ".join(f"{k} {v:.3g}" for k, v in sorted(worst.items()))) +# Slow (the marked cells): a twin and an iterating twin compiled per cell. +# Per push: tests/property/test_coupling_nonlinear_search.py::test_several_passes_of_an_iterating_twin_are_compositions_of_the_single_pass @pytest.mark.parametrize("index", [0, pytest.param(2, marks=pytest.mark.slow), pytest.param(len(linear.PER_PUSH_CELLS), marks=pytest.mark.slow)]) @@ -867,7 +876,7 @@ def test_every_score_holds_on_the_nonlinear_seed_shapes(seed): def _held(name: str, fractions: dict) -> None: - assert fractions["referenced"] >= REFERENCED_FLOOR * fractions["stepped"], (name, fractions) + assert fractions["referenced"] >= REFERENCED_FLOOR * fractions["finite"], (name, fractions) assert fractions["usable"] > 0, f"{name}: no example had its flag set ({fractions})" @@ -905,7 +914,7 @@ def test_the_hunt_finds_no_number_on_the_wrong_side_of_a_nonlinear_group(block, profile = dataclasses.replace(SLOW, max_examples=115).seeded(1000 + block, shrink=False) report, fractions = search(name, cells=BLOCKS[block], profile=profile) print(f"{name}, block {block}: worst {report}; {fractions}") - assert fractions["referenced"] >= REFERENCED_FLOOR * fractions["stepped"], fractions + assert fractions["referenced"] >= REFERENCED_FLOOR * fractions["finite"], fractions assert fractions["usable"] >= USABLE_FLOOR, ( f"{name}, block {block}: only {fractions['usable']:.2f} of the examples had the flag " f"set (floor {USABLE_FLOOR})") From 75748b6239a771811b3e0762ce13aa29573fa733 Mon Sep 17 00:00:00 2001 From: Nicholas Ehsan Roy Date: Wed, 7 Oct 2026 05:46:22 +0200 Subject: [PATCH 4/4] test(coupling): hold the referenced fraction to the examples that returned a finite state Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_013UkCde7g23gTziUjYAvnKD --- docs/developer_guide/testing_standards.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/developer_guide/testing_standards.md b/docs/developer_guide/testing_standards.md index ff75400d..d9d579d6 100644 --- a/docs/developer_guide/testing_standards.md +++ b/docs/developer_guide/testing_standards.md @@ -415,7 +415,7 @@ The coupling search above scores against a closed form, which exists only for li - **How the pass map is obtained.** The caller builds an *x64 twin* of the graph under test: the same nodes, edges, schedule and sub-cycling in float64, its group at `max_iterations=1` with no acceleration and `predictor="linear"` (`twin_knobs`). One step of the twin is one pass, differentiated straight through. The pass starts from the predictor's extrapolation `2 pred_0 - pred_1`, which the step keeps in three `_meta` slots of the state; writing the iterate `x` into both gives `P(x; pre)`: one pass from the iterate `x`, each member integrating from the pre-step state. The map is the graph's own compiled step called on a state, so nothing of the estimator is called, and `jax.jacfwd` of it holds every dependence the pass has on the iterate (transforms, mappings, a geometry field a member holds). The slots' layout is found, not assumed: the twin is stepped once on a drawn example and the order of members whose flattened fields equal the stored `pred_0` is taken. - **What it catches:** anything by which the estimator's numbers disagree with the map the solve iterates (a term of the pass missing from the estimator's Jacobian, a reading at another time level than the pass's, a Jacobian taken at another state than the returned one, a bound in another norm). **What it does not:** a defect in the pass itself, which the reference inherits (the closed-form and time-level references hold the pass); a difference between the single-pass branch and the pass an iterating group runs (checked separately: three passes of an iterating twin are three compositions of the reference's map, to float64 rounding); a node that branches on its dtype; a graph with more than one coupling group. -- **Validated against the closed form** on every per-push draw of the linear search (600 draws on its five per-push cells, and 150 on the multi-rate cell, in the slow lane; forty draws on one cell per push). Worst differences measured: the radius 1.0e-15; the Jacobian 1.6 float64 `eps` of its largest entry; the fixed point 5.7 `eps` of the field times the resolvent's norm; the distance 2.4e-12 of itself; the gradient error 1.9e-12. `ALLOWED` in `tests/property/test_coupling_nonlinear_search.py` holds each to at most 2**12 roundings. +- **Validated against the closed form** on every per-push draw of the linear search (its 556 distinct examples on five cells, its seed shapes, and 150 draws on the multi-rate cell, in the slow lane; forty draws on one cell per push). Worst differences measured: the radius 1.0e-15; the Jacobian 1.6 float64 `eps` of its largest entry; the fixed point 5.7 `eps` of the field times the resolvent's norm; the distance 2.4e-12 of itself; the gradient error 1.9e-12. `ALLOWED` in `tests/property/test_coupling_nonlinear_search.py` holds each to at most 2**12 roundings. - **Cost:** per twin, its build and two compiles (the pass with its Jacobian; the constants' sensitivities), together under two seconds for a dozen scalars; per example, a few dozen jitted calls (about 20 ms). The seconds of a cell are the compile of the graph under test with its diagnostics. **The nonlinear cells** (`tests/property/test_coupling_nonlinear_search.py`) are structures of the linear search whose members pass each port's input through a nonlinearity centred on a parameter `c` before its gain: a saturating gain `c + tanh(s (u - c)) / s`, a quadratic term `u + s (u - c)^2`, or the product of two fields' deviations. Each has `phi(c) = c` and `phi'(c) = 1`, so with `c` the input at the linear group's fixed point the nonlinear group has the same fixed point and the same Jacobian *there*: the numbers of a case are drawn exactly as the linear search draws them, plus the curve `s max|c|` (0.01 to 100), and the linear closed form checks the reference's fixed point on every example. Away from the fixed point, where a capped solve returns, the Jacobian is another matrix, and only the reference knows it.