invexapi 0.1.0__py3-none-any.whl

invexapi/__init__.py

@@ -0,0 +1,45 @@
+from . import metadata, optim, penalties
+from .metadata import DesignDecision, Invariant, Provenance
+from .optim import LinearizedADMM
+from .penalties import (
+    Certificate,
+    FiniteDifference2D,
+    Identity,
+    L1Penalty,
+    LinearOperator,
+    Loss,
+    LogInvexPenalty,
+    ManualVerifier,
+    Penalty,
+    QuasinormInvexPenalty,
+    Reference,
+    Sum,
+    TikhonovPenalty,
+    Verifier,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "optim",
+    "penalties",
+    "metadata",
+    "Loss",
+    "Penalty",
+    "Certificate",
+    "Reference",
+    "Verifier",
+    "ManualVerifier",
+    "Sum",
+    "QuasinormInvexPenalty",
+    "LogInvexPenalty",
+    "TikhonovPenalty",
+    "L1Penalty",
+    "LinearOperator",
+    "Identity",
+    "FiniteDifference2D",
+    "LinearizedADMM",
+    "Provenance",
+    "DesignDecision",
+    "Invariant",
+]

invexapi/metadata.py

@@ -0,0 +1,139 @@
+"""Rich prose docstrings elsewhere in this library tend to tangle three different
+kinds of information together: which source file/kernel an implementation was
+transcribed from, why one design was chosen over a rejected alternative, and what
+invariant must be preserved across future edits. This module gives each of those
+its own small dataclass, plus a ``@documented`` decorator that attaches them to a
+class/function as introspectable attributes and registers it so a tool can retrieve 
+every documented item as plain JSON via ``dump_all_json()`` without parsing docstrings 
+or importing this library at all (if it just wants a static snapshot generated once 
+and shipped elsewhere).
+
+Deliberately plain stdlib dataclasses not something being validated as external input, 
+and invexapi's only required dependency is PyTorch.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from typing import Any, Callable, Optional, Sequence
+
+__all__ = [
+    "Reference",
+    "Provenance",
+    "DesignDecision",
+    "Invariant",
+    "documented",
+    "dump_all",
+    "dump_all_json",
+]
+
+
+@dataclass(frozen=True)
+class Reference:
+    """A citation backing a mathematical or design claim."""
+
+    authors: str
+    title: str
+    venue: str
+    year: int
+    locator: str  # e.g. "Lemma 1, item 5 (Eq. 10)", "Section 3.1.4"
+
+
+@dataclass(frozen=True)
+class Provenance:
+    """Which source file(s)/kernel(s) an implementation was transcribed from."""
+
+    files: tuple[str, ...]
+    description: str
+
+
+@dataclass(frozen=True)
+class DesignDecision:
+    """A choice made, an alternative rejected, and why.
+
+    ``reference`` is set when the decision is backed by a specific paper result
+    (e.g. "this prox formula because Theorem 3 proves it's the global optimizer"),
+    left ``None`` for plain engineering judgement calls.
+    """
+
+    choice: str
+    rejected: str
+    rationale: str
+    reference: Optional[Reference] = None
+
+
+@dataclass(frozen=True)
+class Invariant:
+    """A constraint that must be preserved across future edits."""
+
+    statement: str
+
+
+_REGISTRY: list[type | Callable[..., Any]] = []
+
+
+def documented(
+    *,
+    provenance: Sequence[Provenance] = (),
+    decisions: Sequence[DesignDecision] = (),
+    invariants: Sequence[Invariant] = (),
+    example: Optional[Callable[[], Any]] = None,
+):
+    """Class/function decorator attaching structured documentation.
+
+    ``provenance``/``decisions``/``invariants`` describe the decorated class or
+    function itself, not any particular instance, so they're stored as plain
+    attributes rather than requiring instantiation (unlike
+    :class:`invexapi.penalties.Certificate`, which genuinely can vary per instance).
+
+    ``example``, if given, is a zero-argument callable building a representative
+    instance purely so :func:`dump_all` can read its ``Certificate``s — only
+    meaningful for classes (e.g. concrete ``Penalty``s) whose certificates don't
+    depend on constructor arguments.
+    """
+
+    def wrap(obj):
+        obj.provenance = tuple(provenance)
+        obj.design_decisions = tuple(decisions)
+        obj.invariants = tuple(invariants)
+        obj._metadata_example = example
+        _REGISTRY.append(obj)
+        return obj
+
+    return wrap
+
+
+def _certificates_for(obj) -> dict[str, Any]:
+    example = getattr(obj, "_metadata_example", None)
+    if example is None:
+        return {}
+    instance = example()
+    certs = {}
+    for name in ("convex", "invex", "quasi_convex", "quasi_invex"):
+        cert = getattr(instance, name, None)
+        if cert is not None:
+            certs[name] = asdict(cert)
+    return certs
+
+
+def dump_all() -> list[dict[str, Any]]:
+    """Every ``@documented`` class/function as plain, JSON-serializable dicts."""
+    entries = []
+    for obj in _REGISTRY:
+        entries.append(
+            {
+                "name": obj.__qualname__,
+                "module": obj.__module__,
+                "provenance": [asdict(p) for p in obj.provenance],
+                "design_decisions": [asdict(d) for d in obj.design_decisions],
+                "invariants": [asdict(i) for i in obj.invariants],
+                "certificates": _certificates_for(obj),
+            }
+        )
+    return entries
+
+
+def dump_all_json(**kwargs: Any) -> str:
+    """:func:`dump_all`, serialized to a JSON string."""
+    return json.dumps(dump_all(), **kwargs)

invexapi/optim/__init__.py

@@ -0,0 +1,7 @@
+from .admm import LinearizedADMM
+from .base import Solver
+from .conjugate_gradient import NonlinearCG
+from .fista import FISTA
+from .gradient_descent import GradientDescent
+
+__all__ = ["Solver", "GradientDescent", "FISTA", "NonlinearCG", "LinearizedADMM"]

invexapi/optim/_certification.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import warnings
+
+from ..metadata import DesignDecision, documented
+
+
+@documented(
+    decisions=[
+        DesignDecision(
+            choice="duck-type obj via getattr(obj, ..., None), don't require isinstance(obj, Loss)",
+            rejected="requiring obj to be a Loss instance before checking certificates",
+            rationale=(
+                "optimizers stay duck-typed throughout this library; an object with "
+                "no invex/convex attributes at all is treated identically to a Loss "
+                "with those attributes set to None — no certificate either way"
+            ),
+        )
+    ],
+)
+def warn_if_unproven(obj, label: str) -> None:
+    """Warn if ``obj`` carries no certificate backing a global-optimum claim."""
+    if getattr(obj, "invex", None) is None and getattr(obj, "convex", None) is None:
+        warnings.warn(
+            f"{label} carries no convex/invex certificate — convergence to a "
+            "global optimum is not guaranteed.",
+            stacklevel=3,
+        )

invexapi/optim/_linesearch.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import torch
+
+
+def backtracking(
+    value_fn,
+    x: torch.Tensor,
+    direction: torch.Tensor,
+    grad: torch.Tensor,
+    step0: float = 1.0,
+    shrink: float = 0.5,
+    c1: float = 1e-4,
+    max_backtracks: int = 50,
+) -> float:
+    """Armijo backtracking line search along ``direction`` from ``x``.
+
+    Returns a step size ``t`` such that ``value_fn(x + t*direction)`` gives
+    sufficient decrease, or the smallest tried step if none satisfies it.
+    """
+    f0 = value_fn(x)
+    slope = torch.sum(grad * direction)
+    t = step0
+    for _ in range(max_backtracks):
+        if value_fn(x + t * direction) <= f0 + c1 * t * slope:
+            return t
+        t *= shrink
+    return t

invexapi/optim/admm.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+from typing import Callable, Optional
+
+import torch
+
+from ..metadata import DesignDecision, Invariant, Provenance, documented
+from ..penalties.operators import Identity, LinearOperator
+from .base import Solver
+
+__all__ = ["LinearizedADMM"]
+
+
+@documented(
+    provenance=[
+        Provenance(
+            files=(
+                "code from conference paper Global Optimality for Nonlinear "
+                "Constrained Restoration Problems via Invexity",
+            ),
+            description=(
+                "TV-regularized denoising via ADMM with a linearized/relaxed "
+                "x-update (fixed alpha=0.5 extrapolation, fixed rho, no adaptive "
+                "scheme); z-update and dual update are standard exact ADMM"
+            ),
+        )
+    ],
+    decisions=[
+        DesignDecision(
+            choice="name the class LinearizedADMM, not ADMM",
+            rejected="calling it ADMM (unqualified)",
+            rationale=(
+                "the x-update is a single relaxed/extrapolated gradient step, not "
+                "an exact ADMM subproblem solve — Linearized ADMM is a real, "
+                "published variant name; calling it plain ADMM would imply a "
+                "stronger per-iteration guarantee than this implementation gives"
+            ),
+        ),
+        DesignDecision(
+            choice="warn on smooth and penalty separately, not via a combined Sum",
+            rejected="reusing invexapi.penalties.Sum like FISTA does",
+            rationale=(
+                "Sum assumes smooth and penalty share x's domain; here penalty "
+                "operates on D@x, generally a different shape (e.g. TV's D stacks "
+                "a vertical+horizontal gradient field), so no single combined "
+                "object can be constructed to certify"
+            ),
+        ),
+    ],
+    invariants=[
+        Invariant(
+            "the primal residual used for the stopping criterion and the dual "
+            "update must be computed from the NEWLY updated z, not the z from "
+            "before this iteration's prox step — using the old z is trivially "
+            "zero whenever x0 already minimizes smooth alone (e.g. x0=y for "
+            "0.5||x-y||^2), causing a spurious immediate 'convergence'."
+        )
+    ],
+)
+class LinearizedADMM(Solver):
+    """Linearized ADMM for ``min_x smooth(x) + penalty(D@x)``.
+
+    Named ``LinearizedADMM``, not ``ADMM``, because its x-update is a single
+    relaxed/extrapolated gradient step (transcribed exactly from the source ADMM
+    scripts), not an exact ADMM subproblem solve.
+
+    ``smooth`` needs ``.grad(x)`` (and optionally ``.value(x)``); ``penalty`` needs
+    ``.prox(z, step)`` (and optionally ``.value(z)``), operating on ``D@x``'s
+    domain, which may differ in shape from ``x`` (e.g. total-variation's ``D``
+    stacks a vertical and horizontal gradient field). ``D`` defaults to
+    :class:`~invexapi.penalties.operators.Identity`. ``project``, if given, is
+    applied to ``x`` after every update (e.g. ``lambda x: x.clamp(min=0)`` for
+    natural-image pixel constraints, matching the source scripts) — defaulting to
+    no projection, since forcing non-negativity isn't appropriate for a general
+    solver.
+    """
+
+    def __init__(
+        self,
+        smooth,
+        penalty,
+        D: Optional[LinearOperator] = None,
+        rho: float = 1.0,
+        alpha: float = 0.5,
+        project: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+        max_iter: int = 100,
+        tol: float = 1e-6,
+    ):
+        super().__init__(max_iter=max_iter, tol=tol)
+        self.smooth = smooth
+        self.penalty = penalty
+        self.D = D if D is not None else Identity()
+        self.rho = rho
+        self.alpha = alpha
+        self.project = project
+
+    def run(self, x0: torch.Tensor):
+        # Sum (used by FISTA) assumes smooth and penalty share x's domain; here
+        # penalty operates on D@x, generally a different shape, so there is no
+        # single combined object to check — warn on each half separately instead.
+        self._warn_if_unproven(self.smooth, "smooth")
+        self._warn_if_unproven(self.penalty, "penalty")
+
+        x = x0.clone()
+        x_relaxed = x0.clone()
+        z = self.D.apply(x0)
+        d = torch.zeros_like(z)
+        history = []
+
+        for _ in range(self.max_iter):
+            x1 = (1.0 - self.alpha) * x_relaxed + self.alpha * x
+            grad_coupling = self.rho * self.D.adjoint(self.D.apply(x1) - z + d)
+            grad_smooth = self.smooth.grad(x1)
+            x_new = x - (1.0 / (2.0 * self.rho)) * (grad_coupling + grad_smooth)
+            if self.project is not None:
+                x_new = self.project(x_new)
+
+            x_relaxed = (1.0 - self.alpha) * x_relaxed + self.alpha * x_new
+
+            Dx = self.D.apply(x_new)
+            z = self.penalty.prox(Dx + d / self.rho, 1.0 / self.rho)
+            primal_residual = Dx - z
+            d = d + self.rho * primal_residual
+
+            if hasattr(self.smooth, "value") and hasattr(self.penalty, "value"):
+                history.append((self.smooth.value(x_new) + self.penalty.value(Dx)).item())
+
+            x = x_new
+
+            # Standard ADMM stopping criterion: the primal residual Dx - z (not
+            # x's own movement, which can be zero for a degenerate iteration or
+            # two before the z/dual feedback has caught up, e.g. when x0 already
+            # minimizes smooth alone) must vanish.
+            if primal_residual.norm() < self.tol * max(Dx.norm(), 1.0):
+                break
+
+        return x_relaxed, history

invexapi/optim/base.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import List, Tuple
+
+import torch
+
+from ..metadata import DesignDecision, documented
+from ._certification import warn_if_unproven
+
+__all__ = ["Solver"]
+
+
+@documented(
+    decisions=[
+        DesignDecision(
+            choice="standardize only max_iter/tol and the run(x0)->(x,history) contract",
+            rejected="a fully unified constructor signature across all solvers",
+            rationale=(
+                "GradientDescent/NonlinearCG take a single duck-typed objective "
+                "while FISTA takes smooth+penalty — a real shape difference, not "
+                "incidental duplication to force into one signature"
+            ),
+        )
+    ],
+)
+class Solver(ABC):
+    """Common shape for GradientDescent, FISTA, and NonlinearCG."""
+
+    def __init__(self, max_iter: int = 100, tol: float = 1e-6):
+        self.max_iter = max_iter
+        self.tol = tol
+
+    @abstractmethod
+    def run(self, x0: torch.Tensor) -> Tuple[torch.Tensor, List[float]]:
+        """Run the solver from ``x0``, returning ``(x_final, objective_history)``."""
+
+    def _warn_if_unproven(self, obj, label: str) -> None:
+        warn_if_unproven(obj, label)
+
+    def _record(self, history: List[float], obj, x: torch.Tensor) -> None:
+        """Append ``obj.value(x)`` to ``history`` if ``obj`` exposes ``value``."""
+        if hasattr(obj, "value"):
+            history.append(obj.value(x).item())

invexapi/optim/conjugate_gradient.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ..metadata import DesignDecision, Invariant, documented
+from ._linesearch import backtracking
+from .base import Solver
+
+__all__ = ["NonlinearCG"]
+
+
+@documented(
+    decisions=[
+        DesignDecision(
+            choice="Polak-Ribière+ (PR beta clamped to >= 0, i.e. automatic restart)",
+            rejected="plain Polak-Ribière or Fletcher-Reeves without restart",
+            rationale=(
+                "clamping the PR coefficient to be non-negative automatically "
+                "restarts the direction to steepest descent whenever the plain PR "
+                "formula would go negative, which is known to improve robustness "
+                "on non-quadratic (e.g. invex) objectives over unrestarted variants"
+            ),
+        )
+    ],
+    invariants=[
+        Invariant("beta_pr is always clamped to >= 0 before use as the restart coefficient.")
+    ],
+)
+class NonlinearCG(Solver):
+    """Nonlinear conjugate gradient (Polak-Ribière+, with restart) on any
+    objective exposing ``.grad(x)`` (and, for the default backtracking line
+    search, ``.value(x)``).
+
+    Suited to general (invex or otherwise) differentiable objectives, unlike
+    linear CG which only solves quadratic ``Ax=b`` systems.
+    """
+
+    def __init__(
+        self,
+        objective,
+        step: Optional[float] = None,
+        max_iter: int = 100,
+        tol: float = 1e-6,
+    ):
+        super().__init__(max_iter=max_iter, tol=tol)
+        if step is None and not hasattr(objective, "value"):
+            raise ValueError("objective must implement value(x) when step is None")
+        self.objective = objective
+        self.step = step
+
+    def run(self, x0: torch.Tensor):
+        self._warn_if_unproven(self.objective, "objective")
+        x = x0.clone()
+        history = []
+
+        grad = self.objective.grad(x)
+        direction = -grad
+
+        for _ in range(self.max_iter):
+            self._record(history, self.objective, x)
+            if grad.norm() < self.tol:
+                break
+
+            if self.step is not None:
+                t = self.step
+            else:
+                t = backtracking(self.objective.value, x, direction, grad)
+            x = x + t * direction
+
+            grad_new = self.objective.grad(x)
+            beta_pr = torch.sum(grad_new * (grad_new - grad)) / torch.sum(grad * grad).clamp_min(1e-12)
+            beta = torch.clamp(beta_pr, min=0.0)  # Polak-Ribière+ restart
+
+            direction = -grad_new + beta * direction
+            grad = grad_new
+
+        return x, history

invexapi/optim/fista.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import math
+
+import torch
+
+from ..metadata import DesignDecision, documented
+from ..penalties import Sum
+from .base import Solver
+
+__all__ = ["FISTA"]
+
+
+@documented(
+    decisions=[
+        DesignDecision(
+            choice="warn/record off self.combined = Sum(smooth, penalty), not smooth/penalty individually",
+            rejected="checking smooth's and penalty's certificates separately",
+            rationale=(
+                "a global-optimum claim is about the COMBINED objective smooth+penalty "
+                "minimizes, and certificates don't compose from the parts (see "
+                "invexapi.penalties.Sum) — checking the parts separately could miss "
+                "that the combination itself is uncertified even when both halves are"
+            ),
+        )
+    ],
+)
+class FISTA(Solver):
+    """Accelerated proximal gradient (FISTA) for ``min_x smooth(x) + penalty(x)``.
+
+    ``smooth`` must implement ``.grad(x)`` (and optionally ``.value(x)`` for the
+    tracked history); ``penalty`` must implement ``.prox(x, step)`` (and optionally
+    ``.value(x)``). ``step`` should be <= 1/L where L is the Lipschitz constant of
+    ``smooth.grad``.
+    """
+
+    def __init__(self, smooth, penalty, step: float, max_iter: int = 100, tol: float = 1e-6):
+        super().__init__(max_iter=max_iter, tol=tol)
+        self.smooth = smooth
+        self.penalty = penalty
+        self.combined = Sum(smooth, penalty)
+        self.step = step
+
+    def run(self, x0: torch.Tensor):
+        self._warn_if_unproven(self.combined, "smooth+penalty")
+        x_prev = x0.clone()
+        y = x0.clone()
+        t_prev = 1.0
+        history = []
+
+        for _ in range(self.max_iter):
+            grad = self.smooth.grad(y)
+            x = self.penalty.prox(y - self.step * grad, self.step)
+
+            self._record(history, self.combined, x)
+
+            if (x - x_prev).norm() < self.tol * max(x_prev.norm(), 1.0):
+                x_prev = x
+                break
+
+            t = (1.0 + math.sqrt(1.0 + 4.0 * t_prev ** 2)) / 2.0
+            y = x + ((t_prev - 1.0) / t) * (x - x_prev)
+
+            x_prev = x
+            t_prev = t
+
+        return x_prev, history

invexapi/optim/gradient_descent.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._linesearch import backtracking
+from .base import Solver
+
+__all__ = ["GradientDescent"]
+
+
+class GradientDescent(Solver):
+    """Gradient descent on any objective exposing ``.grad(x)`` (and, for the
+    default backtracking line search, ``.value(x)``).
+
+    If ``step`` is given, a fixed step size is used; otherwise each iteration
+    picks its step via Armijo backtracking.
+    """
+
+    def __init__(
+        self,
+        objective,
+        step: Optional[float] = None,
+        max_iter: int = 100,
+        tol: float = 1e-6,
+    ):
+        super().__init__(max_iter=max_iter, tol=tol)
+        if step is None and not hasattr(objective, "value"):
+            raise ValueError("objective must implement value(x) when step is None")
+        self.objective = objective
+        self.step = step
+
+    def run(self, x0: torch.Tensor):
+        self._warn_if_unproven(self.objective, "objective")
+        x = x0.clone()
+        history = []
+        for _ in range(self.max_iter):
+            grad = self.objective.grad(x)
+            self._record(history, self.objective, x)
+            if grad.norm() < self.tol:
+                break
+            if self.step is not None:
+                t = self.step
+            else:
+                t = backtracking(self.objective.value, x, -grad, grad)
+            x = x - t * grad
+        return x, history

invexapi/penalties/__init__.py

@@ -0,0 +1,23 @@
+from .base import Certificate, Loss, ManualVerifier, Penalty, Reference, Sum, Verifier
+from .convex import TikhonovPenalty
+from .l1 import L1Penalty
+from .log import LogInvexPenalty
+from .operators import FiniteDifference2D, Identity, LinearOperator
+from .quasinorm import QuasinormInvexPenalty
+
+__all__ = [
+    "Loss",
+    "Penalty",
+    "Certificate",
+    "Reference",
+    "Verifier",
+    "ManualVerifier",
+    "Sum",
+    "QuasinormInvexPenalty",
+    "LogInvexPenalty",
+    "TikhonovPenalty",
+    "L1Penalty",
+    "LinearOperator",
+    "Identity",
+    "FiniteDifference2D",
+]