soar-lite 0.1.0__py3-none-any.whl

soar_lite/__init__.py

@@ -0,0 +1,77 @@
+"""soar_lite: first-order n-path descent with stall restarts.
+
+The public face is deliberate and small. The two plugin contracts
+(`LocalSolver`, `GlobalModel`) and the data types they exchange are re-exported
+from the engine so callers can write their own methods and restart-point models
+against them; the concrete implementations (`DiagonalPreconditioner`,
+`SampleGlobalModel`, `BLRGlobalModel`) are re-exported for convenience.
+`n_path_schedule` is the whole algorithm entry point.
+
+Whether a stalled path is re-drawn is the `restart_stalled` knob on
+`n_path_schedule`, not a plugin: *when* a path has stopped is the local solver's
+own stopping test, and *where* it restarts is the global model's, so the engine
+needs no third plugin seam to join them.
+
+Internal helpers (the box/argmin utilities, the run's private state holders) are
+not re-exported here on purpose; they are `_`-private in their own modules. A
+test or a contributor may still import them, but nothing outside the package is
+meant to.
+"""
+
+from soar_lite.domain import Bounds, Domain, DrawFn
+from soar_lite.engine import (
+    EvalF,
+    EvalFG,
+    GlobalModel,
+    LocalSolver,
+    NPathResult,
+    NPathStatus,
+    PathPoint,
+    Point,
+    ServeTrials,
+    ValuePoint,
+    Work,
+    n_path_schedule,
+    pointwise,
+)
+from soar_lite.global_model import (
+    AcquisitionFn,
+    BLRGlobalModel,
+    IdentityProjection,
+    MinimizeFn,
+    Projection,
+    ProjectionFactory,
+    SampleGlobalModel,
+    SphereProjection,
+    descent_minimize,
+)
+from soar_lite.local_solver import DiagonalPreconditioner
+
+__all__ = [
+    "AcquisitionFn",
+    "BLRGlobalModel",
+    "Bounds",
+    "DiagonalPreconditioner",
+    "Domain",
+    "DrawFn",
+    "EvalF",
+    "EvalFG",
+    "GlobalModel",
+    "IdentityProjection",
+    "LocalSolver",
+    "MinimizeFn",
+    "NPathResult",
+    "NPathStatus",
+    "PathPoint",
+    "Point",
+    "Projection",
+    "ProjectionFactory",
+    "SampleGlobalModel",
+    "ServeTrials",
+    "SphereProjection",
+    "ValuePoint",
+    "Work",
+    "descent_minimize",
+    "n_path_schedule",
+    "pointwise",
+]

soar_lite/domain.py

@@ -0,0 +1,368 @@
+"""The search domain: a declared box, its per-dimension scales and its draws.
+
+This module owns the coordinate space a restart model reads and the distribution
+it proposes restart points from - the box (any side possibly unbounded), each
+dimension's prior width in coordinate units, and the per-dimension draw, either
+the default for that dimension or one the caller injected. It is the object a
+`Projection` is built on: the geometry adds a `features` map, this owns
+everything about *where* the model may look and *how* a start is proposed, so a
+model-free restart sampler needs nothing but this and a key.
+
+The box's own spelling lives here too - `Bounds`, and the validation and
+normalization behind it. This module imports nothing from the package, so it is
+the leaf both `n_path_schedule` and the model modules read a box from; the
+check that holds every point to it stays the engine's.
+
+A side may be unbounded, and there are three cases, not two: a bounded dimension
+draws uniformly, a free one (both sides infinite) draws ``N(0, scale_i^2)``, and
+a half-bounded one (exactly one finite side) draws that same normal *truncated*
+to its declared side - the default never steps over the side the caller declared,
+because a start outside it is not a start in the declared domain at all.
+"""
+
+import math
+from collections.abc import Callable, Sequence
+from typing import TypeAlias
+
+import jax
+import jax.numpy as jnp
+from jax.scipy.special import ndtr, ndtri
+
+Bounds: TypeAlias = Sequence[tuple[float | None, float | None]] | jnp.ndarray
+"""The search box in scipy's spelling, which `Domain` normalizes.
+
+One ``(lb, ub)`` pair per dimension, where either side of the pair - the lower
+or the upper bound - may be ``None`` or ``+-inf``, meaning unbounded there. A
+``(d, 2)`` array is the same data row-wise, for callers that already hold one.
+"""
+
+DrawFn: TypeAlias = Callable[[jnp.ndarray, int], jnp.ndarray]
+"""One dimension's draw, injected per dimension.
+
+``sampler(key, n)`` returns ``n`` coordinates for that one dimension, shape
+``(n,)``. Injected per dimension so a restart model can sample where a bounded
+box does not reach; ``None`` leaves the domain's own default in place - uniform
+inside a bounded side, ``N(0, scale_i**2)`` where the domain is free, and that
+normal restricted to the declared side where it is half-bounded.
+"""
+
+
+class Domain:
+    """The coordinates a restart model may read, and how a start is drawn.
+
+    A domain carries the box, the per-dimension scales and the per-dimension
+    draws, and validates all three once at construction: a NaN side, a lower side
+    above its upper, a scale that is not a ``(d,)`` vector of finite positive
+    values, or a sampler count that is not ``d`` are refused here rather than
+    deep inside a clip or a draw. It is concrete - a caller sampling with no
+    surrogate needs nothing more - and `Projection` wraps one to add the feature
+    map a fitted model reads.
+
+    A side may be unbounded. The box then declares only the dimension and which
+    sides are unbounded; `scale` is what carries a length there, because no width
+    exists to read one off. Not even a bounded side's width is a natural unit:
+    `scale` is in coordinate units, so a caller changing a variable from meters
+    to millimeters changes `scale`, never the box.
+
+    `draw` is the one contract: ``(key, n) -> (n, d)`` starts drawn one dimension
+    at a time.
+
+    Example:
+        Sample starts with no engine and no surrogate, the case a domain alone
+        covers:
+
+            >>> import jax
+            >>> from soar_lite import Domain
+
+            >>> domain = Domain([(0.0, 1.0), (None, None)], scale=[1.0, 2.0])
+            >>> domain.draw(jax.random.PRNGKey(0), 3).shape
+            (3, 2)
+
+        The first column is uniform inside its bounded side, the second draws
+        ``N(0, 2^2)`` because nothing there is declared.
+
+    """
+
+    d: int
+    """The dimension: the number of pairs `bounds` declares.
+
+    Every shape in the engine is read off it - ``x0`` is ``(n_paths, d)``, a
+    proposal is ``(d,)`` - so a caller never states the dimension twice.
+    """
+
+    lo: jnp.ndarray
+    """The lower side per dimension, ``(d,)``; ``-inf`` where that side is free."""
+
+    hi: jnp.ndarray
+    """The upper side per dimension, ``(d,)``; ``+inf`` where that side is free."""
+
+    bounds: Bounds
+    """The box as the caller declared it, unbounded sides included."""
+
+    scale: jnp.ndarray
+    """The prior width per dimension, ``(d,)``, in coordinate units.
+
+    It is the length the default draw is measured in on an unbounded side, and
+    what `SphereProjection` normalizes by; a dimension whose spread is far from
+    a unit scale needs its own entry here rather than a different box.
+    """
+
+    center: jnp.ndarray
+    """The box's midpoint per dimension, ``(d,)``; zero outside `bounded_dims`."""
+
+    width: jnp.ndarray
+    """The box's width per dimension, ``(d,)``; zero outside `bounded_dims`."""
+
+    bounded_dims: jnp.ndarray
+    """The dimensions with both sides finite, ``(d,)``.
+
+    `width` and `center` read off it, so a dimension outside it is zero in
+    both.
+    """
+
+    free_dims: jnp.ndarray
+    """The dimensions with neither side finite, ``(d,)``.
+
+    A dimension in neither mask is the half-bounded one, whose default draw
+    alone has a boundary to respect.
+    """
+
+    def __init__(
+        self,
+        bounds: Bounds,
+        *,
+        scale: jnp.ndarray | Sequence[float] | None = None,
+        samplers: Sequence[DrawFn | None] | None = None,
+    ) -> None:
+        """Validate and carry the declared box, the scales and the samplers.
+
+        Args:
+            bounds: the box, scipy-style pairs or a ``(d, 2)`` array. A side may
+                be ``None`` or ``+-inf`` for an unbounded domain; a NaN side, or
+                a lower side above its upper bound, is refused. There is no box
+                for a domain to default to: the dimension and the unbounded
+                sides are what the box declares.
+            scale: per-dimension prior width in coordinate units, a ``(d,)``
+                vector or ``None`` for every dimension at one. A projection that
+                normalizes compares the entries to each other. The common
+                magnitude is not free either: `scale` is the length the default
+                draw is measured in, so a dimension whose coordinates span far
+                more or far less than a unit scale should carry its own entry.
+            samplers: one optional draw per dimension, ``d`` entries; a ``None``
+                entry keeps the default for that dimension - uniform inside a
+                bounded dimension, ``N(0, scale_i^2)`` inside a free one, and
+                that same normal restricted to the domain inside a half-bounded
+                one (exactly one finite side), where drawing the unrestricted
+                normal would step over the side the caller did declare.
+                Dimensions are drawn independently, so a prior that couples them
+                needs its own `GlobalModel`, not a domain.
+
+        Raises:
+            ValueError: If the box is malformed, or `scale` is not a ``(d,)``
+                vector of finite positive values, or `samplers` is not ``d``
+                entries long.
+
+        """
+        self.lo, self.hi = _check_bounds(bounds)
+        # The dimension is what `bounds` declares: read it off the argument, not
+        # off `lo`/`hi`, which `_check_bounds` just derived from it.
+        self.d = len(bounds)
+        self.bounds = bounds
+        self.scale = _check_scale(scale, self.d)
+        self._samplers = _check_samplers(samplers, self.d)
+        # The domain's shape per dimension. An unbounded side has no width, so
+        # `width` is zero there and the center falls back to the origin - the
+        # only center an unbounded dimension has, and `scale` the only length
+        # a caller can state about it.
+        self.bounded_dims = jnp.isfinite(self.lo) & jnp.isfinite(self.hi)
+        # Neither side finite. The half-bounded dimension - one side finite,
+        # one not - is the remaining case, and the only one whose default draw
+        # has a boundary to respect.
+        self.free_dims = ~jnp.isfinite(self.lo) & ~jnp.isfinite(self.hi)
+        self.width = jnp.where(self.bounded_dims, self.hi - self.lo, 0.0)
+        self.center = jnp.where(self.bounded_dims, self.lo + self.width / 2.0, 0.0)
+
+    def draw(self, key: jnp.ndarray, n_points: int) -> jnp.ndarray:
+        """Return ``(n_points, d)`` restart starts, one dimension at a time.
+
+        A dimension with an injected `DrawFn` takes its column from that, else
+        from that dimension's default: uniform inside a bounded side,
+        ``N(0, scale_i^2)`` inside a free one, and the same normal *truncated* to
+        the domain inside a half-bounded one - the default never steps over the
+        side the caller declared. An injected sampler is the caller's own
+        statement about the dimension's spread, so nothing checks the drawn
+        values against `scale`: an arbitrary distribution publishes no spread to
+        check them against.
+
+        Raises:
+            ValueError: If `n_points` is below one, or a column does not come
+                back ``(n_points,)`` long.
+
+        """
+        if n_points <= 0:
+            msg = f"a draw needs at least one point, got {n_points}"
+            raise ValueError(msg)
+        columns: list[jnp.ndarray] = []
+        for j, sub in enumerate(jax.random.split(key, self.d)):
+            sampler = self._samplers[j]
+            if sampler is not None:
+                column = sampler(sub, n_points)
+            elif bool(self.bounded_dims[j]):
+                column = jax.random.uniform(
+                    sub,
+                    (n_points,),
+                    minval=self.lo[j],
+                    maxval=self.hi[j],
+                )
+            elif bool(self.free_dims[j]):
+                column = self.scale[j] * jax.random.normal(sub, (n_points,))
+            else:
+                column = _truncated_normal(
+                    sub,
+                    n_points,
+                    lower=float(self.lo[j]),
+                    upper=float(self.hi[j]),
+                    scale=float(self.scale[j]),
+                )
+            if jnp.shape(column) != (n_points,):
+                msg = (
+                    f"each drawn column must be ({n_points},) long, got "
+                    f"{jnp.shape(column)} for dimension {j}"
+                )
+                raise ValueError(msg)
+            columns.append(column)
+        return jnp.stack(columns, axis=-1)
+
+
+def _check_scale(scale: jnp.ndarray | Sequence[float] | None, d: int) -> jnp.ndarray:
+    """Normalize the per-dimension prior width to ``(d,)``, refusing a bad one.
+
+    ``None`` is every dimension at one, which puts the coordinate unit itself in
+    charge: the caller has either made each variable's spread O(1) or is about to
+    say what it is. Anything but a ``(d,)`` vector is refused rather than
+    broadcast: a scalar or a length-one vector would silently become a uniform
+    vector, which is a real scale but never what a caller reaching for those
+    shapes means.
+    """
+    if scale is None:
+        return jnp.ones(d)
+    given = jnp.asarray(scale)
+    if given.ndim != 1 or given.shape[0] != d:
+        msg = f"scale must be a ({d},) vector or None, got shape {given.shape}"
+        raise ValueError(msg)
+    if not bool(jnp.all(jnp.isfinite(given))) or not bool(jnp.all(given > 0.0)):
+        msg = f"scale must hold finite positive values, got {scale!r}"
+        raise ValueError(msg)
+    return given
+
+
+def _check_samplers(
+    samplers: Sequence[DrawFn | None] | None,
+    d: int,
+) -> tuple[DrawFn | None, ...]:
+    """Normalize the injected draws to one entry per dimension.
+
+    ``None`` is every dimension on its default draw. The length is checked
+    because zipping a short sequence would leave the tail dimensions silently on
+    the default - a wiring error that would look like a working run.
+    """
+    if samplers is None:
+        return (None,) * d
+    given = tuple(samplers)
+    if len(given) != d:
+        msg = f"samplers needs {d} entries, got {len(given)}"
+        raise ValueError(msg)
+    return given
+
+
+def _check_bounds(bounds: Bounds) -> tuple[jnp.ndarray, jnp.ndarray]:
+    """Reject a malformed box and normalize it to ``(lo, hi)``, in one pass.
+
+    Both spellings are read the same way: one `(lb, ub)` pair per dim, or a
+    `(d, 2)` array. A `None` side becomes ∓inf before anything is looked at, so
+    the check and the box it hands out cannot disagree about what a side means.
+    A NaN side is rejected - it would poison every clip silently instead of
+    failing; +-inf sides are legitimate (unbounded dim).
+    """
+    pairs: list[tuple[float | None, float | None]]
+    if isinstance(bounds, jnp.ndarray):
+        if bounds.ndim != 2 or bounds.shape[1] != 2:
+            msg = f"bounds as an array needs shape (d, 2), got {bounds.shape}"
+            raise ValueError(msg)
+        pairs = list(zip(bounds[:, 0].tolist(), bounds[:, 1].tolist(), strict=True))
+    else:
+        pairs = []
+        for j, pair in enumerate(bounds):
+            if len(pair) != 2:
+                msg = f"bounds[{j}] must be a (lb, ub) pair, got {pair!r}"
+                raise ValueError(msg)
+            pairs.append((pair[0], pair[1]))
+    lo: list[float] = []
+    hi: list[float] = []
+    for j, (low, high) in enumerate(pairs):
+        lo_f = -math.inf if low is None else low
+        hi_f = math.inf if high is None else high
+        if math.isnan(lo_f) or math.isnan(hi_f):
+            msg = f"bounds[{j}] has a NaN side: ({low!r}, {high!r})"
+            raise ValueError(msg)
+        if lo_f > hi_f:
+            msg = f"bounds[{j}] has lo={lo_f} > hi={hi_f}"
+            raise ValueError(msg)
+        lo.append(lo_f)
+        hi.append(hi_f)
+    return jnp.asarray(lo), jnp.asarray(hi)
+
+
+_TRUNCATED_TAIL = 9.0
+"""Clip the standard-normal draw to this many scales either way.
+
+Only the domain's own endpoints can push the CDF inverse past it: `ndtri` is
+undefined at 0 and 1, and a truncation far into the tail rounds both bounds of
+the uniform interval onto the same endpoint. The clip keeps such a draw finite
+and costs nothing else - beyond nine scales the normal's mass is negligible, so
+no reachable draw is touched, while a domain truncated that far out collapses
+onto its own boundary, which is the right limit anyway.
+"""
+
+
+def _truncated_normal(
+    key: jnp.ndarray,
+    n_points: int,
+    *,
+    lower: float,
+    upper: float,
+    scale: float,
+) -> jnp.ndarray:
+    """Draw ``N(0, scale^2)`` restricted to ``[lower, upper]``.
+
+    Both bounds take an infinity, so this one routine serves either orientation
+    of a half-bounded domain. Inverse-CDF because the alternative - draw and
+    reject - costs a variable number of rounds that JAX cannot express without
+    a fixed cap, which would bias whatever the cap cuts off.
+
+    Args:
+        key: the PRNG key to draw from.
+        n_points: how many values to draw.
+        lower: the domain's lower bound, possibly ``-inf``.
+        upper: the domain's upper bound, possibly ``+inf``.
+        scale: the prior's standard deviation, in coordinate units.
+
+    Returns:
+        ``(n_points,)`` finite values, all inside ``[lower, upper]``.
+
+    """
+    v = jax.random.uniform(key, (n_points,))
+    # Sample the *tail* probability, not the CDF value. A half-line is one tail,
+    # and a truncation far out makes that tail small enough that `1 - Phi(a)` and
+    # `Phi(a)` agree to the last float32 - the CDF value loses its precision
+    # exactly where the truncated mass still lives. Multiplying the small tail by
+    # `v` keeps the relative precision, so the draw tracks the true truncated
+    # normal instead of piling its far tail onto one clipped value.
+    if lower == -math.inf:
+        p = ndtr(jnp.asarray(upper) / scale)
+        z = ndtri(p * v)
+    else:
+        p = ndtr(-jnp.asarray(lower) / scale)
+        z = -ndtri(p * (1.0 - v))
+    z = jnp.clip(z, -_TRUNCATED_TAIL, _TRUNCATED_TAIL)
+    return jnp.clip(scale * z, lower, upper)

soar_lite/engine/__init__.py

@@ -0,0 +1,64 @@
+"""The engine package: the contracts a plugin is written against, and the machine.
+
+This package is two modules behind one facade.
+`soar_lite.engine.contracts` is what a plugin author reads: the two plugin
+contracts (`LocalSolver`, `GlobalModel`), the data they exchange (`Point`,
+`ValuePoint`, `PathPoint`, `Work`, `NPathResult`, `NPathStatus`), the
+evaluation-seam aliases, and `pointwise`. `soar_lite.engine.schedule` is the
+machine: the `n_path_schedule` loop and the four private holders a run's state
+lives in.
+This module re-exports the public half, so `from soar_lite.engine import
+LocalSolver` names the contract without a caller knowing which file holds it.
+
+The dependency direction is one-way - plugins import these types, never the
+reverse:
+
+    engine/contracts.py   <- LocalSolver, GlobalModel, Point, ValuePoint, PathPoint
+    engine/schedule.py    <- n_path_schedule and the four holders
+    domain.py             <- Domain, Bounds (the box; a leaf)
+    local_solver.py       <- DiagonalPreconditioner (implements LocalSolver)
+    global_model/         <- SampleGlobalModel, BLRGlobalModel and its projections
+
+The two concrete defaults are bound lazily inside `n_path_schedule`, which keeps
+that graph acyclic. The engine is first-order (value and gradient),
+coordinate-free, and owns the flock, the domain and its box test, the shared horizon
+and the restart cadence; the local solver owns direction, step size, acceptance
+and its step's enumeration, and the global model only where a restart lands.
+
+Restarting is two plugin seams and one flag, not a third plugin: `_Paths` marks a
+path stalled when its method gave up or its point was not usable,
+`restart_stalled` says whether such a path is re-drawn, and `GlobalModel.propose`
+says where.
+"""
+
+from soar_lite.engine.contracts import (
+    EvalF,
+    EvalFG,
+    GlobalModel,
+    LocalSolver,
+    NPathResult,
+    NPathStatus,
+    PathPoint,
+    Point,
+    ServeTrials,
+    ValuePoint,
+    Work,
+    pointwise,
+)
+from soar_lite.engine.schedule import n_path_schedule
+
+__all__ = [
+    "EvalF",
+    "EvalFG",
+    "GlobalModel",
+    "LocalSolver",
+    "NPathResult",
+    "NPathStatus",
+    "PathPoint",
+    "Point",
+    "ServeTrials",
+    "ValuePoint",
+    "Work",
+    "n_path_schedule",
+    "pointwise",
+]