python/dpt/registration.py

1"""Safeguarded limited-memory optimisation in a fixed, scaled parameter chart.23The evaluator owns the CUDA composition. It returns only a scalar objective and4the small parameter gradient, never a detector image. Parameterisations (pose5anchor, units, constrained nuisance variables) remain fixed throughout a solve.6Starting a new chart requires a new solve and therefore empty curvature history.7"""89from __future__ import annotations1011import math12from collections.abc import Callable13from dataclasses import dataclass14from typing import Literal, Protocol1516from dpt.contracts import ContractError, NumericalError, finite_scalar, integer1718Vector = tuple[float, ...]192021@dataclass(frozen=True, slots=True)22class Evaluation:23    """Objective and derivative with respect to the supplied dimensionless chart."""2425    loss: float26    gradient: Vector272829class Evaluator(Protocol):30    def __call__(self, parameters: Vector, /) -> Evaluation:31        """Evaluate current parameters; invalid trials raise a numerical/domain error."""32        ...333435@dataclass(frozen=True, slots=True)36class RecoveryPolicy:37    max_iterations: int = 20038    max_evaluations: int = 200039    memory: int = 1040    line_search_evaluations: int = 3241    gradient_tolerance: float = 1e-642    relative_loss_tolerance: float = 1e-1243    step_tolerance: float = 1e-1044    armijo: float = 1e-445    curvature: float = 0.946    maximum_step: float = 16.047    minimum_step: float = 1e-144849    def __post_init__(self) -> None:50        for name in ("max_iterations", "max_evaluations", "memory", "line_search_evaluations"):51            integer(getattr(self, name), name, minimum=1)52        for name in ("gradient_tolerance", "relative_loss_tolerance", "step_tolerance"):53            finite_scalar(getattr(self, name), name, minimum=0.0)54        for name in ("armijo", "curvature", "minimum_step", "maximum_step"):55            finite_scalar(getattr(self, name), name, minimum=0.0)56        if not 0.0 < self.armijo < self.curvature < 1.0:57            raise ContractError("strong-Wolfe constants must satisfy 0 < c1 < c2 < 1")58        if not 0.0 < self.minimum_step < self.maximum_step:59            raise ContractError("line-search step interval must be positive and nonempty")606162StopReason = Literal[63    "gradient_tolerance",64    "step_tolerance",65    "loss_stagnation",66    "line_search_failed",67    "evaluation_budget",68    "iteration_budget",69    "cancelled",70]717273@dataclass(frozen=True, slots=True)74class IterationRecord:75    iteration: int76    loss: float77    gradient_infinity_norm: float78    step: float79    evaluations: int80    curvature_pairs: int818283@dataclass(frozen=True, slots=True)84class RecoveryResult:85    parameters: Vector86    evaluation: Evaluation87    reason: StopReason88    evaluations: int89    history: tuple[IterationRecord, ...]9091    @property92    def stationary(self) -> bool:93        """Only the gradient criterion establishes numerical stationarity here.9495        Small loss/step changes indicate stagnation; they do not prove the pose96        was recovered or that the current point is an identifiable minimum.97        """98        return self.reason == "gradient_tolerance"99100101def _dot(a: Vector, b: Vector) -> float:102    try:103        return math.fsum(x * y for x, y in zip(a, b, strict=True))104    except (OverflowError, ValueError):105        # Curvature and line-search callers reject an unrepresentable product;106        # it cannot certify a direction or replace the accepted iterate.107        return math.nan108109110def _norm(a: Vector) -> float:111    return max(abs(x) for x in a)112113114def _checked(evaluation: Evaluation, dimension: int) -> Evaluation:115    if len(evaluation.gradient) != dimension:116        raise ContractError("evaluator gradient dimension differs from its parameter chart")117    if not math.isfinite(evaluation.loss) or not all(map(math.isfinite, evaluation.gradient)):118        raise NumericalError("non-finite objective or parameter gradient")119    return evaluation120121122# region book:registration-lbfgs-direction123def _direction(gradient: Vector, pairs: list[tuple[Vector, Vector, float]]) -> Vector:124    """Two-loop recursion; curvature vectors all belong to the same scaled chart."""125    if not pairs:126        return _steepest_direction(gradient)127    q = gradient128    coefficients: list[float] = []129    for displacement, change, reciprocal in reversed(pairs):130        coefficient = reciprocal * _dot(displacement, q)131        coefficients.append(coefficient)132        q = tuple(a - coefficient * b for a, b in zip(q, change, strict=True))133    scale = 1.0134    if pairs:135        displacement, change, _ = pairs[-1]136        denominator = _dot(change, change)137        if not math.isfinite(denominator) or denominator <= 0.0:138            return _steepest_direction(gradient)139        scale = _dot(displacement, change) / denominator140        if not math.isfinite(scale) or scale <= 0.0:141            return _steepest_direction(gradient)142    result = tuple(scale * value for value in q)143    for (displacement, change, reciprocal), coefficient in zip(144        pairs, reversed(coefficients), strict=True145    ):146        beta = reciprocal * _dot(change, result)147        result = tuple(148            a + (coefficient - beta) * b for a, b in zip(result, displacement, strict=True)149        )150    return tuple(-value for value in result)151152153# endregion book:registration-lbfgs-direction154155156def _steepest_direction(gradient: Vector) -> Vector:157    """Bound the initial chart displacement without squaring a large gradient."""158    scale = max(1.0, _norm(gradient))159    return tuple(-value / scale for value in gradient)160161162class _BudgetExhaustedError(Exception):163    pass164165166def _line_search(167    evaluate: Callable[[Vector], Evaluation],168    x: Vector,169    base: Evaluation,170    direction: Vector,171    policy: RecoveryPolicy,172    *,173    initial_step: float = 1.0,174) -> tuple[float, Evaluation] | None:175    """Seek strong Wolfe, accepting strict Armijo decrease at the step bound.176177    Bisection sacrifices polynomial interpolation speed to keep invalid-domain178    endpoints and nonfinite trials unambiguous. No failed or unchecked trial is179    returned as an accepted iterate. The last accepted state lives in the caller.180    """181    slope = _dot(base.gradient, direction)182    if not math.isfinite(slope) or slope >= 0:183        return None184    attempts = 0185186    def trial(alpha: float) -> Evaluation | None:187        nonlocal attempts188        attempts += 1189        point = tuple(value + alpha * step for value, step in zip(x, direction, strict=True))190        if not all(map(math.isfinite, point)):191            return None192        try:193            return evaluate(point)194        except (NumericalError, OverflowError):195            # Domain invalidity may be signalled as NumericalError by a196            # parameterised evaluator. ContractError remains a programming error.197            return None198199    def armijo(alpha: float, candidate: Evaluation) -> bool:200        return candidate.loss <= base.loss + policy.armijo * alpha * slope201202    def zoom(low: float, high: float, low_value: Evaluation) -> tuple[float, Evaluation] | None:203        while attempts < policy.line_search_evaluations:204            alpha = 0.5 * (low + high)205            if abs(high - low) < policy.minimum_step or alpha == low or alpha == high:206                return None207            candidate = trial(alpha)208            if (209                candidate is None210                or not armijo(alpha, candidate)211                or candidate.loss >= low_value.loss212            ):213                high = alpha214                continue215            derivative = _dot(candidate.gradient, direction)216            if abs(derivative) <= -policy.curvature * slope:217                return alpha, candidate218            if derivative * (high - low) >= 0:219                high = low220            low, low_value = alpha, candidate221        return None222223    previous_alpha, previous_value = 0.0, base224    alpha = min(initial_step, policy.maximum_step)225    while attempts < policy.line_search_evaluations:226        candidate = trial(alpha)227        if (228            candidate is None229            or not armijo(alpha, candidate)230            or (previous_alpha > 0 and candidate.loss >= previous_value.loss)231        ):232            return zoom(previous_alpha, alpha, previous_value)233        derivative = _dot(candidate.gradient, direction)234        if abs(derivative) <= -policy.curvature * slope:235            return alpha, candidate236        if alpha == policy.maximum_step and candidate.loss < base.loss:237            return alpha, candidate238        if derivative >= 0:239            return zoom(alpha, previous_alpha, candidate)240        if alpha == policy.maximum_step:241            return None242        previous_alpha, previous_value = alpha, candidate243        alpha = min(2.0 * alpha, policy.maximum_step)244    return None245246247# region book:registration-safeguarded-loop248def recover_parameters(249    evaluator: Evaluator,250    initial: Vector,251    *,252    policy: RecoveryPolicy | None = None,253    observe: Callable[[IterationRecord], None] | None = None,254    cancelled: Callable[[], bool] | None = None,255) -> RecoveryResult:256    """Minimise a deterministic objective in an immutable, dimensionless chart.257258    The evaluator must use the same observation, forward model and parameter259    coordinates for every call. It may reuse device scratch, but must complete260    numerical-status checks before returning. Rejected trials never replace the261    accepted parameters in the result. Observer exceptions are not swallowed.262    """263    selected = policy or RecoveryPolicy()264    x: Vector = tuple(finite_scalar(value, "initial parameter") for value in initial)265    if not x:266        raise ContractError("recovery requires at least one active parameter")267    calls = 0268269    def evaluate(point: Vector) -> Evaluation:270        nonlocal calls271        if calls >= selected.max_evaluations:272            raise _BudgetExhaustedError273        calls += 1274        return _checked(evaluator(point), len(x))275276    current = evaluate(x)277    history: list[IterationRecord] = []278    pairs: list[tuple[Vector, Vector, float]] = []279    reason: StopReason = "iteration_budget"280    history.append(IterationRecord(0, current.loss, _norm(current.gradient), 0.0, calls, 0))281    for iteration in range(1, selected.max_iterations + 1):282        if cancelled is not None and cancelled():283            reason = "cancelled"284            break285        if _norm(current.gradient) <= selected.gradient_tolerance:286            reason = "gradient_tolerance"287            break288        direction = _direction(current.gradient, pairs)289        slope = _dot(direction, current.gradient)290        if not all(map(math.isfinite, direction)) or not math.isfinite(slope) or slope >= 0:291            pairs.clear()292            direction = _steepest_direction(current.gradient)293        try:294            accepted = _line_search(evaluate, x, current, direction, selected)295            if accepted is None:296                # A bad inverse-Hessian estimate need not terminate a sound297                # forward model: try steepest descent once with empty memory.298                pairs.clear()299                direction = _steepest_direction(current.gradient)300                accepted = _line_search(301                    evaluate,302                    x,303                    current,304                    direction,305                    selected,306                    initial_step=max(307                        selected.minimum_step, 2.0 ** -min(16, selected.line_search_evaluations)308                    ),309                )310        except _BudgetExhaustedError:311            reason = "evaluation_budget"312            break313        if accepted is None:314            reason = "line_search_failed"315            break316        alpha, new_value = accepted317        step: Vector = tuple(alpha * value for value in direction)318        new_x: Vector = tuple(a + b for a, b in zip(x, step, strict=True))319        change = tuple(a - b for a, b in zip(new_value.gradient, current.gradient, strict=True))320        curvature = _dot(step, change)321        threshold = 1e-12 * math.hypot(*step) * math.hypot(*change)322        if (323            math.isfinite(curvature)324            and curvature > max(0.0, threshold)325            and math.isfinite(1.0 / curvature)326        ):327            pairs.append((step, change, 1.0 / curvature))328            if len(pairs) > selected.memory:329                del pairs[0]330        old_loss = current.loss331        x, current = new_x, new_value332        record = IterationRecord(333            iteration, current.loss, _norm(current.gradient), alpha, calls, len(pairs)334        )335        history.append(record)336        if observe is not None:337            observe(record)338        if _norm(current.gradient) <= selected.gradient_tolerance:339            reason = "gradient_tolerance"340            break341        if _norm(step) <= selected.step_tolerance * max(1.0, _norm(x)):342            reason = "step_tolerance"343            break344        if abs(old_loss - current.loss) <= selected.relative_loss_tolerance * max(345            1.0, abs(old_loss), abs(current.loss)346        ):347            reason = "loss_stagnation"348            break349    return RecoveryResult(x, current, reason, calls, tuple(history))350351352# endregion book:registration-safeguarded-loop353