proxax 0.0.1__py3-none-any.whl

proxax/__init__.py

@@ -0,0 +1,79 @@
+"""Certified first-order proximal point methods for composite optimization."""
+
+from __future__ import annotations
+
+from proxax import certificates
+from proxax import constraints
+from proxax import optax_compat
+from proxax import oracle
+from proxax import problem as _problem
+from proxax import prox
+from proxax import solve as _solve
+from proxax import solvers
+from proxax import tree
+
+__version__ = "0.0.1"
+
+ACACG = solvers.ACACG
+AG = solvers.AG
+AIPP = solvers.AIPP
+APD = solvers.APD
+AdapAPG = solvers.AdapAPG
+AdapFISTA = solvers.AdapFISTA
+CompositeProblem = _problem.CompositeProblem
+DAICG = solvers.DAICG
+ECG = solvers.ECG
+IAICG = solvers.IAICG
+IAIPAL = solvers.IAIPAL
+IALM = solvers.IALM
+KKTResiduals = _problem.KKTResiduals
+NCFISTA = solvers.NCFISTA
+QPAIPP = solvers.QPAIPP
+QPAPD = solvers.QPAPD
+SProxALM = solvers.SProxALM
+Solution = solvers.Solution
+SolveStats = solvers.SolveStats
+SolveStatus = solvers.SolveStatus
+Solver = solvers.Solver
+SpectralOracle = oracle.SpectralOracle
+SpectralProblem = oracle.SpectralProblem
+StoppingConfig = solvers.StoppingConfig
+UPFAG = solvers.UPFAG
+solve = _solve.solve
+
+__all__ = [
+    "ACACG",
+    "AG",
+    "AIPP",
+    "APD",
+    "AdapAPG",
+    "AdapFISTA",
+    "CompositeProblem",
+    "DAICG",
+    "ECG",
+    "IAICG",
+    "IAIPAL",
+    "IALM",
+    "KKTResiduals",
+    "NCFISTA",
+    "QPAIPP",
+    "QPAPD",
+    "SProxALM",
+    "Solution",
+    "SolveStats",
+    "SolveStatus",
+    "Solver",
+    "SpectralOracle",
+    "SpectralProblem",
+    "StoppingConfig",
+    "UPFAG",
+    "__version__",
+    "certificates",
+    "constraints",
+    "optax_compat",
+    "oracle",
+    "prox",
+    "solve",
+    "solvers",
+    "tree",
+]

proxax/_loop.py

@@ -0,0 +1,67 @@
+"""Dual-mode loop execution supporting both JAX while_loop and eager Python."""
+
+from __future__ import annotations
+
+import time
+import typing
+from collections import abc
+
+import jax
+
+from proxax import _typing
+
+StateT = typing.TypeVar("StateT")
+
+
+def while_loop(
+    cond_fun: abc.Callable[[StateT], jax.Array],
+    body_fun: abc.Callable[[StateT], StateT],
+    init_val: StateT,
+    *,
+    jit: bool = True,
+    callback: _typing.CallbackFn[StateT] | None = None,
+    time_limit: float | None = None,
+    on_time_limit: abc.Callable[[StateT], StateT] | None = None,
+) -> StateT:
+  """Executes a stateful loop via `jax.lax.while_loop` or eager Python.
+
+  Args:
+    cond_fun: Continuation predicate returning a boolean JAX scalar.
+    body_fun: Pure state transition function `state -> next_state`.
+    init_val: Initial loop carry PyTree.
+    jit: If True, executes `jax.lax.while_loop` (XLA-compiled and
+      `vmap`-compatible). If False, runs a Python `while` loop supporting
+      Python callbacks and wall-clock time limits.
+    callback: Optional Python callback run after each iteration when
+      `jit=False`. Returning True stops the loop early.
+    time_limit: Optional wall-clock limit in seconds (supported when
+      `jit=False`).
+    on_time_limit: Optional state updater applied when `time_limit` triggers.
+
+  Returns:
+    Final loop state `StateT`.
+
+  Raises:
+    ValueError: If `callback` or `time_limit` is supplied when `jit=True`.
+  """
+  if jit:
+    if callback is not None:
+      raise ValueError("Python `callback` requires `jit=False`.")
+    if time_limit is not None:
+      raise ValueError("Wall-clock `time_limit` requires `jit=False`.")
+    return jax.lax.while_loop(cond_fun, body_fun, init_val)
+
+  state = init_val
+  start_time = time.perf_counter()
+  while bool(cond_fun(state)):
+    if (
+        time_limit is not None
+        and (time.perf_counter() - start_time) >= time_limit
+    ):
+      if on_time_limit is not None:
+        state = on_time_limit(state)
+      break
+    state = body_fun(state)
+    if callback is not None and callback(state):
+      break
+  return state

proxax/_typing.py

@@ -0,0 +1,39 @@
+"""Core type aliases and structural protocols for proxax."""
+
+from __future__ import annotations
+
+import typing
+
+import jax
+
+Array: typing.TypeAlias = jax.Array
+Scalar: typing.TypeAlias = jax.Array | float
+PyTree: typing.TypeAlias = typing.Any
+StateT = typing.TypeVar("StateT")
+StateContraT = typing.TypeVar("StateContraT", contravariant=True)
+
+
+@typing.runtime_checkable
+class SmoothFn(typing.Protocol):
+  """Callable evaluating a scalar smooth objective on a PyTree."""
+
+  def __call__(self, x: PyTree, /) -> jax.Array:
+    """Evaluates the smooth function value at `x`."""
+    ...
+
+
+@typing.runtime_checkable
+class ValueAndGradFn(typing.Protocol):
+  """Callable returning both scalar value and PyTree gradient at `x`."""
+
+  def __call__(self, x: PyTree, /) -> tuple[jax.Array, PyTree]:
+    """Evaluates `(f(x), grad_f(x))`."""
+    ...
+
+
+class CallbackFn(typing.Protocol[StateContraT]):
+  """Optional eager-mode callback invoked after each outer iteration."""
+
+  def __call__(self, state: StateContraT, /) -> bool | None:
+    """Inspects solver state; returning True requests early termination."""
+    ...

proxax/certificates/__init__.py

@@ -0,0 +1,14 @@
+"""Stationarity and KKT certificate refinement routines."""
+
+from __future__ import annotations
+
+from proxax.certificates import stationarity
+
+RefinementResult = stationarity.RefinementResult
+refine = stationarity.refine
+
+__all__ = [
+    "RefinementResult",
+    "refine",
+    "stationarity",
+]

proxax/certificates/stationarity.py

@@ -0,0 +1,168 @@
+"""Stationarity refinement mapping approximate iterates to certified pairs."""
+
+from __future__ import annotations
+
+import dataclasses
+
+import jax
+import jax.numpy as jnp
+
+from proxax import _loop
+from proxax import _typing
+from proxax import tree
+from proxax.oracle import composite
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class RefinementResult:
+  """Certified approximate stationary pair `(x_hat, v_hat)`.
+
+  Guarantees exact subdifferential membership
+  `v_hat in grad_f(x_hat) + dh(x_hat)` with monotone composite descent
+  `phi(x_hat) <= phi(y)`.
+  """
+
+  x_hat: _typing.PyTree
+  v_hat: _typing.PyTree
+  obj_val: jax.Array
+  lipschitz_used: jax.Array
+  smooth_evals: jax.Array
+  grad_evals: jax.Array
+  prox_evals: jax.Array
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class _RefineCarry:
+  """Internal backtracking loop state for `refine`."""
+
+  lipschitz: jax.Array
+  x_trial: _typing.PyTree
+  f_trial: jax.Array
+  accepted: jax.Array
+  backtracks: jax.Array
+
+
+def refine(
+    oracle: composite.CompositeOracleProtocol,
+    y: _typing.PyTree,
+    lipschitz_init: _typing.Scalar,
+    *,
+    backtrack_factor: float = 2.0,
+    max_backtracks: int = 40,
+    jit: bool = True,
+) -> RefinementResult:
+  """Computes a certified pair `(x_hat, v_hat)` from an approximate iterate `y`.
+
+  Performs one adaptive composite gradient step from `y` with smoothness `L`:
+
+  - `x_hat = prox_{(1 / L) h}(y - (1 / L) grad_f(y))`
+  - `v_hat = grad_f(x_hat) - grad_f(y) + L * (y - x_hat)`
+
+  Because `L * (y - x_hat) - grad_f(y) in dh(x_hat)` holds exactly by the
+  optimality condition of `prox`, `v_hat in grad_f(x_hat) + dh(x_hat)` is
+  guaranteed without subgradient selection heuristics.
+
+  Args:
+    oracle: Target `CompositeOracleProtocol`.
+    y: Approximate domain point `y in dom(h)`.
+    lipschitz_init: Initial trial smoothness constant `L_0 > 0`.
+    backtrack_factor: Multiplicative step-up factor `> 1` during line search.
+    max_backtracks: Maximum number of backtracking doublings.
+    jit: Whether to run line search via `jax.lax.while_loop` (`True`) or eager
+      Python (`False`).
+
+  Returns:
+    `RefinementResult` containing `(x_hat, v_hat)`, objective value, and counts.
+  """
+  dtype = tree.infer_dtype(y)
+  f_y, grad_y = oracle.smooth_value_and_grad(y)
+  num_slack = 10.0 * tree.eps(dtype) * (1.0 + jnp.abs(f_y))
+
+  if oracle.curvature_upper is not None:
+    L_fixed = jnp.maximum(
+        jnp.asarray(oracle.curvature_upper, dtype=dtype),
+        jnp.asarray(1e-12, dtype=dtype),
+    )
+    inv_L = 1.0 / L_fixed
+    x_hat = oracle.prox(tree.add_scaled(y, -inv_L, grad_y), inv_L)
+    f_x_hat, grad_x_hat = oracle.smooth_value_and_grad(x_hat)
+    v_hat = tree.add(
+        tree.sub(grad_x_hat, grad_y),
+        tree.scale(L_fixed, tree.sub(y, x_hat)),
+    )
+    obj_val = f_x_hat + oracle.proximal_value(x_hat)
+    return RefinementResult(
+        x_hat=x_hat,
+        v_hat=v_hat,
+        obj_val=obj_val,
+        lipschitz_used=L_fixed,
+        smooth_evals=jnp.asarray(2, dtype=jnp.int32),
+        grad_evals=jnp.asarray(2, dtype=jnp.int32),
+        prox_evals=jnp.asarray(1, dtype=jnp.int32),
+    )
+
+  L0 = jnp.maximum(
+      jnp.asarray(lipschitz_init, dtype=dtype),
+      jnp.asarray(1e-12, dtype=dtype),
+  )
+
+  def _eval_trial(
+      L_val: jax.Array,
+  ) -> tuple[_typing.PyTree, jax.Array, jax.Array]:
+    step = 1.0 / L_val
+    x_t = oracle.prox(tree.add_scaled(y, -step, grad_y), step)
+    f_t = oracle.smooth_value(x_t)
+    diff = tree.sub(x_t, y)
+    quad_upper = (
+        f_y + tree.vdot(grad_y, diff) + 0.5 * L_val * tree.norm_sq(diff)
+    )
+    ok = f_t <= quad_upper + num_slack
+    return x_t, f_t, ok
+
+  x0_trial, f0_trial, ok0 = _eval_trial(L0)
+  init_carry = _RefineCarry(
+      lipschitz=L0,
+      x_trial=x0_trial,
+      f_trial=f0_trial,
+      accepted=ok0,
+      backtracks=jnp.asarray(0, dtype=jnp.int32),
+  )
+
+  def _cond(carry: _RefineCarry) -> jax.Array:
+    return jnp.logical_and(
+        jnp.logical_not(carry.accepted),
+        carry.backtracks < max_backtracks,
+    )
+
+  def _body(carry: _RefineCarry) -> _RefineCarry:
+    next_L = carry.lipschitz * backtrack_factor
+    x_t, f_t, ok = _eval_trial(next_L)
+    return _RefineCarry(
+        lipschitz=next_L,
+        x_trial=x_t,
+        f_trial=f_t,
+        accepted=ok,
+        backtracks=carry.backtracks + 1,
+    )
+
+  final_carry = _loop.while_loop(_cond, _body, init_carry, jit=jit)
+  x_hat = final_carry.x_trial
+  L_used = final_carry.lipschitz
+  f_x_hat, grad_x_hat = oracle.smooth_value_and_grad(x_hat)
+  v_hat = tree.add(
+      tree.sub(grad_x_hat, grad_y),
+      tree.scale(L_used, tree.sub(y, x_hat)),
+  )
+  obj_val = f_x_hat + oracle.proximal_value(x_hat)
+  trials = 1 + final_carry.backtracks
+  return RefinementResult(
+      x_hat=x_hat,
+      v_hat=v_hat,
+      obj_val=obj_val,
+      lipschitz_used=L_used,
+      smooth_evals=trials + 2,
+      grad_evals=jnp.asarray(2, dtype=jnp.int32),
+      prox_evals=trials,
+  )

proxax/constraints/__init__.py

@@ -0,0 +1,62 @@
+"""Linear operators, convex sets, and set-inclusion constraints."""
+
+from __future__ import annotations
+
+from proxax.constraints import linear_op
+from proxax.constraints import sets
+
+BoundCallableLinearOperator = linear_op.BoundCallableLinearOperator
+BoundConstraint = linear_op.BoundConstraint
+BoundLinearOperator = linear_op.BoundLinearOperator
+BoundProductConstraint = linear_op.BoundProductConstraint
+BoundSetConstraint = linear_op.BoundSetConstraint
+BoxSet = sets.BoxSet
+CallableLinearOperator = linear_op.CallableLinearOperator
+Constraint = linear_op.Constraint
+ConvexSet = sets.ConvexSet
+LinearOperator = linear_op.LinearOperator
+LowerBound = sets.LowerBound
+MatrixLinearOperator = linear_op.MatrixLinearOperator
+ProductConstraint = linear_op.ProductConstraint
+ProductSet = sets.ProductSet
+SecondOrderCone = sets.SecondOrderCone
+SetConstraint = linear_op.SetConstraint
+Singleton = sets.Singleton
+UpperBound = sets.UpperBound
+ZeroSet = sets.ZeroSet
+as_linear_operator = linear_op.as_linear_operator
+box = linear_op.box
+equality = linear_op.equality
+greater_equal = linear_op.greater_equal
+in_set = linear_op.in_set
+less_equal = linear_op.less_equal
+normalize_constraints = linear_op.normalize_constraints
+
+__all__ = [
+    "BoundCallableLinearOperator",
+    "BoundConstraint",
+    "BoundLinearOperator",
+    "BoundProductConstraint",
+    "BoundSetConstraint",
+    "BoxSet",
+    "CallableLinearOperator",
+    "Constraint",
+    "ConvexSet",
+    "LinearOperator",
+    "LowerBound",
+    "MatrixLinearOperator",
+    "ProductConstraint",
+    "ProductSet",
+    "SecondOrderCone",
+    "SetConstraint",
+    "Singleton",
+    "UpperBound",
+    "ZeroSet",
+    "as_linear_operator",
+    "box",
+    "equality",
+    "greater_equal",
+    "in_set",
+    "less_equal",
+    "normalize_constraints",
+]