proxax 0.0.1__tar.gz

proxax-0.0.1/.github/workflows/publish.yml

@@ -0,0 +1,54 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    name: Build and test distributions
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v7
+      - name: Run lint, type-check, docs, and unit tests
+        env:
+          XLA_PYTHON_CLIENT_PREALLOCATE: "false"
+        run: |
+          uv sync --python 3.12 --extra dev --extra docs
+          uv run ruff check .
+          uv run mypy
+          uv run sphinx-build -W -b html docs site
+          uv run pytest -v
+      - name: Build sdist and wheel
+        run: uv build
+      - name: Upload distribution artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+
+  publish:
+    name: Publish to PyPI
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/proxax
+    permissions:
+      id-token: write
+      contents: read
+    steps:
+      - uses: astral-sh/setup-uv@v7
+      - name: Download distribution artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+      - name: Publish package to PyPI
+        run: uv publish --trusted-publishing always

proxax-0.0.1/.gitignore

@@ -0,0 +1,32 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Distribution / packaging
+build/
+dist/
+*.egg-info/
+.eggs/
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# Testing, type-checking, and linting caches
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+.pyright/
+.coverage
+htmlcov/
+
+# Documentation build output
+site/
+
+# IDE & OS artifacts
+.vscode/
+.idea/
+.DS_Store
+Thumbs.db

proxax-0.0.1/AGENTS.md

@@ -0,0 +1,56 @@
+# `proxax` — Workspace Rules & Agent Context
+
+`proxax` (https://github.com/wwkong/proxax) is a Python library of **certified first-order proximal point methods** for unconstrained and set/cone-constrained nonconvex composite optimization in **JAX** and **Optax**, inspired by Weiwei (William) Kong's MATLAB framework **[`wwkong/nc_opt`](https://github.com/wwkong/nc_opt)** and papers (including **`AIPP`**, **`QP-AIPP`**, **`PF.APD` / `APD`** at https://wwkong.github.io/files/publications/apd.pdf, and **`NL-IAPIAL` / `AIDAL`**).
+
+Read `DESIGN.md` for mathematical formulations and `docs/nc_opt_reference.md` for the detailed mapping to `wwkong/nc_opt`, the complete **`PF.APD` (`apd.pdf`)** specification, and the Phase 2+ roadmap.
+
+---
+
+## 1. Environment & Verification Commands (WSL2 Primary Workspace)
+
+- **Canonical Repo Path:** `/home/wkong/Projects/proxax` (`origin`: `https://github.com/wwkong/proxax.git`, branch `main`).
+- **Package Manager & Runtime:** `uv` (`pyproject.toml` + `uv.lock`) with Python 3.14, `jax[cuda13]`, and `optax`.
+- **Runtime Dependency Rule:** `src/proxax/` must depend **strictly on `jax` and `optax` only** (zero SciPy / CVXPY / Equinox / Flax runtime dependencies).
+- **GPU Memory Rule (12 GB RTX 3060 shared with Windows host):**
+  - Always ensure `XLA_PYTHON_CLIENT_PREALLOCATE=false` when running JAX on GPU (`tests/conftest.py` and `docs/conf.py` set this automatically).
+  - **Keep unit tests tiny, fast (< 10s total), and low-memory** (2D–4D problems, tight iteration caps, `jax.lax.while_loop` over long Python dispatch loops). Never introduce large matrices or long-running benchmarks into `tests/`.
+
+Run all four verification gates via `uv` before committing:
+
+```bash
+uv run ruff check .
+uv run mypy
+XLA_PYTHON_CLIENT_PREALLOCATE=false uv run pytest -v
+XLA_PYTHON_CLIENT_PREALLOCATE=false uv run sphinx-build -W -b html docs site
+```
+
+---
+
+## 2. Strict Coding & Style Conventions (Google Python Style Guide)
+
+1. **80-Column Line Limit:** Enforced by `ruff` (`line-length = 80`, `indent-width = 2` or standard 2/4-space blocks; match existing module style and pass `uv run ruff check .`).
+2. **Module-Only Imports (`TID253` / Google Style):**
+   - Always import modules, never classes/functions directly (except `typing` / `collections.abc` when module-qualified):
+     - Good: `import dataclasses`, `import typing`, `from collections import abc`, `from proxax import tree`, `from proxax.prox import base as prox_base`.
+     - Bad: `from dataclasses import dataclass`, `from proxax.prox.base import ProxOperator`.
+   - Avoid `B008` function calls in default parameter values (bind module-level finals like `_ZERO_PROX: typing.Final = prox_base.Zero()`).
+3. **Immutable PyTree Dataclasses & Strict Typing (`mypy --strict`):**
+   - Use `@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)` for solvers, configs, problems, states, and results.
+   - When registering dataclasses with `meta_fields` / `data_fields`, call `jax.tree_util.register_dataclass(Cls, data_fields=[...], meta_fields=[...])` after the class definition so `mypy --strict` resolves types cleanly against JAX stubs.
+   - Use `tree.eps(dtype)` (`float(np.finfo(dtype).eps)`) instead of untyped `jnp.finfo(dtype).eps`.
+4. **Sphinx / Google Docstrings:**
+   - Avoid trailing `'s` attached directly to single-backtick inline literals (write `` `Constraint` objects `` instead of `` `Constraint`s ``) and keep blank lines around bullet/formula blocks so `sphinx-build -W` passes with zero warnings.
+
+---
+
+## 3. Core Mathematical & Numerical Invariants (Do Not Break)
+
+1. **Exact Stationarity Certificates (`src/proxax/certificates/stationarity.py`):**
+   - Every solver returns a certified refined point `x_hat = prox_{(1/L) h}(y - (1/L) grad_f(y))` and exact composite subgradient `v_hat = grad_f(x_hat) - grad_f(y) + L * (y - x_hat) in grad_f(x_hat) + dh(x_hat)` (plus `A^*(p_hat)` in constrained problems), never raw unrefined inner iterates.
+2. **Strongly-Convex 1-Prox `ACG` (`src/proxax/solvers/inner/acg.py`):**
+   - Maintains the exact $\mu$-strongly convex estimating sequence with $c_\mu = 1 + \mu A_t$, $b_\mu = 1 + 2\mu A_t$, $d_\mu = 1 + \mu A_{t+1}$, convex combination weights $(A_t d_\mu, a_t c_\mu)$, and exact $\varepsilon$-subgradient pair $(u_{t+1}, \eta_{t+1}) \in \partial_{\eta_{t+1}} \psi(y_{t+1})$.
+3. **IEEE `float32` Adaptive Smoothness Noise-Floor Guard:**
+   - Only relax `L_next = L_used * ls_decrease` when `0.5 * L_used * ||y_next - x_tilde||^2 > 100 * eps * (1 + |f|)`. Without this guard, `L` decays toward `0` when step norms fall below $\sqrt{\epsilon_{\text{mach}}} \approx 3.5 \times 10^{-4}$ in `float32`.
+   - When `lower_curvature == 0.0` (convex `f`), calibrate default `lam0 = 10.0 / max(L0, 1e-6)` rather than huge values like `1e3` so `1 / lam` stays well above the `float32` precision floor in the $\sigma$-criterion.
+4. **Curvature Preservation Across Penalty / ALM Stages (`src/proxax/oracle/composite.py`):**
+   - `PenalizedOracle` adds the convex Moreau envelope $\psi_c(x; \bar{p}) = \frac{c}{2}\|w(x) - \Pi_{\mathcal{C}}(w(x))\|_2^2 - \frac{\|\bar{p}\|_2^2}{2c}$ where $w(x) = \mathcal{A}(x) + \bar{p}/c$ and $\bar{p} = (1 - \theta) p$. Because $\psi_c$ is convex, `curvature_lower` $m$ of $f$ is unchanged as $c_k$ grows, so proximal steps $\lambda_k$ must **never** shrink as $O(1/c_k)$.

proxax-0.0.1/DESIGN.md

@@ -0,0 +1,57 @@
+# Architecture & MVP Implementation Plan: `proxax`
+
+## 1. Core Architectural Pillars & Standards
+
+**`proxax`** is a JAX/Optax-native Python library for certified first-order proximal point methods on unconstrained and linearly/set-constrained nonconvex composite optimization problems.
+
+### Engineering & Style Standards
+- **Dependencies:** Strictly **`jax`** and **`optax`** at runtime (leveraging `optax.projections` and `optax.tree` directly; zero third-party framework lock-in).
+- **Code Style:** **Google Python Style Guide**, strictly enforced via `ruff` (80-column line limit, `convention = "google"`, module-only imports via `flake8-tidy-imports` / `isort`) and `mypy --strict`:
+  - Module-only imports everywhere (`import dataclasses`, `import typing`, `from collections import abc`, `from proxax import prox`).
+  - `@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)` for all configs, solvers, problems, states, and results (paired with `jax.tree_util.register_dataclass` for PyTree states/oracles). Unary math primitives (e.g., `prox.L1(0.1)`) keep their single primary value positional and optional metadata `kw_only`.
+  - Structural subtyping via **`typing.Protocol`** (`ProxOperator`, `ConvexSet`, `LinearOperator`, `CompositeOracleProtocol`).
+- **Documentation:** **Sphinx + MyST-Parser + Furo + MathJax** (`uv run sphinx-build -W -b html docs site`).
+- **Environment & Packaging:** **`uv`** (`pyproject.toml` + `uv.lock`) in WSL2 (`/home/wkong/Projects/proxax`), pushed to `https://github.com/wwkong/proxax`.
+
+---
+
+## 2. Unified Mathematical Formulations
+
+### 2.1. Unconstrained Composite Problem
+$$
+\min_{x \in \mathcal{X}} \; \phi(x) := f(x) + h(x)
+$$
+- $\mathcal{X}$: Arbitrary PyTree of JAX arrays (inner product $\langle u, v \rangle = \text{optax.tree.vdot}(u, v)$).
+- $f: \mathcal{X} \to \mathbb{R}$ (`smooth`): Differentiable, $(m, L^+)$-weakly convex/smooth ($-\frac{m}{2}\|x-y\|^2 \le f(x) - f(y) - \langle \nabla f(y), x-y \rangle \le \frac{L^+}{2}\|x-y\|^2$).
+- $h: \mathcal{X} \to (-\infty, +\infty]$ (`proximal`): Proper, closed, $\mu$-strongly convex ($\mu \ge 0$) with cheap $\operatorname{prox}_{\gamma h}$ and domain projection $\Pi_{\operatorname{dom}(h)}$.
+- **Certified Output:** Refined pair $(\hat{x}, \hat{v})$ satisfying $\hat{v} \in \nabla f(\hat{x}) + \partial h(\hat{x})$ and $\|\hat{v}\| \le \rho$.
+
+### 2.2. Set / Cone-Constrained Composite Problem
+$$
+\min_{x \in \mathcal{X}} \; f(x) + h(x) \quad \text{s.t.} \quad \mathcal{A}(x) \in \mathcal{C}
+$$
+- $\mathcal{A}: \mathcal{X} \to \mathcal{Y}$: Linear map between PyTrees (dense/sparse matrix, `(matvec, rmatvec)` pair, or forward-only callable bound lazily to `x0` via `jax.linear_transpose`).
+- $\mathcal{C} \subseteq \mathcal{Y}$: **Any closed convex set** with Euclidean projection $\Pi_{\mathcal{C}}$ (singletons $\{b\}$ for $\mathcal{A}(x) = b$, halfspaces $(-\infty, b]$, boxes $[l, u]$, SOC, PSD cone, or named PyTree dicts via `ProductConvexSet`).
+- **Unified Penalty / AL Smooth Term** (for multiplier $\bar{p}$ and penalty $c > 0$, with $\bar{p} = 0$ for pure quadratic penalty in `QPAIPP` and $\bar{p} = (1-\theta)p$ for ALM variants):
+  $$
+  \psi_c(x; \bar{p}) = \frac{c}{2}\left\|w(x) - \Pi_{\mathcal{C}}(w(x))\right\|^2 - \frac{\|\bar{p}\|^2}{2c}, \qquad w(x) := \mathcal{A}(x) + \frac{\bar{p}}{c}
+  $$
+  with smooth gradient $\nabla_x \psi_c(x; \bar{p}) = \mathcal{A}^* \hat{p}(x)$, where $\hat{p}(x) := c\big(w(x) - \Pi_{\mathcal{C}}(w(x))\big) \in \mathcal{N}_{\mathcal{C}}\big(\Pi_{\mathcal{C}}(w(x))\big)$.
+- **Certified KKT Output:** $(\hat{x}, \hat{p}, \hat{v})$ with $\hat{v} \in \nabla f(\hat{x}) + \partial h(\hat{x}) + \mathcal{A}^* \hat{p}$, $\|\hat{v}\| \le \rho$, and $\operatorname{dist}_{\mathcal{C}}(\mathcal{A}(\hat{x})) = \|\mathcal{A}(\hat{x}) - \Pi_{\mathcal{C}}(\mathcal{A}(\hat{x}))\| \le \eta$.
+
+---
+
+## 3. Key Architectural Safeguards
+
+1. **Zero Re-Compilation Oracles:** `ProxRegularizedOracle` and `PenalizedOracle` are `@jax.tree_util.register_dataclass` structs with static callables in `meta_fields` and dynamic JAX arrays (`center`, `smooth_reg`, `prox_reg`, `multiplier`, `penalty`) in `data_fields`.
+2. **Safe Curvature Splitting:** `ProxRegularizedOracle` splits $\frac{1}{\lambda} = \alpha_{\text{smooth}} + \mu_{\text{prox}}$ explicitly—never claiming $\mu_{\text{prox}} > 0$ unless `curvature_lower` $m$ is known or validated by adaptive descent/convexity checks.
+3. **`vmap`-Fast Nested Loops vs. Optax Adapter:** `proxax.solve(..., jit=True)` uses nested `jax.lax.while_loop`s so `refine` only runs once per outer step even under `jax.vmap`. `proxax.optax_compat` provides the flattened per-inner-step `GradientTransformationExtraArgs` state machine while keeping `params` equal to the feasible primal iterate $y_t$.
+4. **Safe Initialization & Flexible Stopping:** `solver.init` projects `x0` onto $\operatorname{dom}(h)$ via `proximal.domain_project(x0)` before any gradient evaluation, and `StoppingConfig` supports both $\ell_2$ and $\ell_\infty$ norms plus optional relative tolerances.
+
+---
+
+## 4. MVP Execution Milestones (Atomic Git Commits)
+
+1. **Milestone 1 — Tooling & Core Math Primitives:** `_typing.py`, `_loop.py`, `tree.py`, `prox/`, `constraints/`, `oracle/`, `certificates/stationarity.py`, `problem.py` + unit tests.
+2. **Milestone 2 — `ACG` & Unconstrained `AIPP`:** `solvers/base.py`, `solvers/inner/acg.py`, `solvers/unconstrained/aipp.py`, `proxax.solve()` + nonconvex composite tests (`jit=True`, `jit=False`, `vmap`).
+3. **Milestone 3 — Constrained `QPAIPP` & Optax State Machine:** `solvers/constrained/qp_aipp.py`, `optax_compat/transforms.py` + constrained nonconvex QP and PyTree Optax training-loop tests.

proxax-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 William Kong
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

proxax-0.0.1/PKG-INFO

@@ -0,0 +1,95 @@
+Metadata-Version: 2.5
+Name: proxax
+Version: 0.0.1
+Summary: Certified first-order proximal point methods for composite optimization in JAX and Optax.
+Project-URL: Homepage, https://github.com/wwkong/proxax
+Project-URL: Repository, https://github.com/wwkong/proxax
+Project-URL: Issues, https://github.com/wwkong/proxax/issues
+Author-email: William Kong <wwkong92@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: augmented-lagrangian,composite-optimization,jax,nonconvex-optimization,optax,optimization,proximal
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.4.30
+Requires-Dist: optax>=0.2.2
+Provides-Extra: cuda12
+Requires-Dist: jax[cuda12]>=0.4.30; extra == 'cuda12'
+Provides-Extra: cuda13
+Requires-Dist: jax[cuda13]>=0.5.0; extra == 'cuda13'
+Provides-Extra: dev
+Requires-Dist: mypy>=1.11; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.6; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: furo>=2024.8; extra == 'docs'
+Requires-Dist: myst-parser>=4.0; extra == 'docs'
+Requires-Dist: sphinx-autodoc-typehints>=2.4; extra == 'docs'
+Requires-Dist: sphinx>=8.0; extra == 'docs'
+Description-Content-Type: text/markdown
+
+# `proxax`
+
+Certified first-order proximal point methods for unconstrained and constrained
+nonconvex composite optimization in **JAX** and **Optax**.
+
+## Installation
+
+```bash
+pip install proxax
+```
+
+Or with `uv`:
+
+```bash
+uv add proxax
+```
+
+## Problem Classes
+
+1. **Unconstrained Nonconvex Composite Optimization:**
+   $$\min_{x \in \mathcal{X}} \; f(x) + h(x)$$
+   where $f$ (`smooth`) is differentiable and $(m, L^+)$-weakly convex/smooth on
+   an arbitrary JAX PyTree $\mathcal{X}$, and $h$ (`proximal`) is proper,
+   closed, convex, and proximable.
+
+2. **Set / Cone-Constrained Nonconvex Composite Optimization:**
+   $$\min_{x \in \mathcal{X}} \; f(x) + h(x) \quad \text{s.t.} \quad \mathcal{A}(x) \in \mathcal{C}$$
+   where $\mathcal{A}: \mathcal{X} \to \mathcal{Y}$ is a linear map between
+   PyTrees and $\mathcal{C} \subseteq \mathcal{Y}$ is a closed convex set with
+   Euclidean projection $\Pi_{\mathcal{C}}$.
+
+## Quick Example
+
+```python
+import jax.numpy as jnp
+import proxax
+from proxax import constraints
+from proxax import prox
+from proxax import solvers
+
+problem = proxax.CompositeProblem(
+    smooth=lambda x: 0.5 * jnp.dot(x, q_mat @ x) + jnp.dot(c_vec, x),
+    proximal=prox.Box(lower=-1.0, upper=1.0),
+    constraints=constraints.equality(lambda x: a_mat @ x, rhs=b_vec),
+)
+
+sol = proxax.solve(problem, x0, solver=solvers.QPAIPP(opt_tol=1e-5, feas_tol=1e-5))
+```
+
+## Available Certified Solvers (`wwkong/nc_opt` Suite)
+
+- **Unconstrained Composite (`CompositeProblem`):** `AIPP`, `APD` (`PF.APD`), `ECG`, `AG`, `UPFAG`, `NCFISTA`, `AdapFISTA`, `ACACG`, `AdapAPG`
+- **Unconstrained Spectral (`SpectralProblem` / `NuclearNorm`):** `IAICG`, `DAICG`
+- **Set / Cone-Constrained (`CompositeProblem` with `constraints`):** `QPAIPP`, `QPAPD`, `IALM` (`NL-IAPIAL` / `AIDAL`), `SProxALM`, `IAIPAL`
+- **Optax Single-Step Transforms (`proxax.optax_compat`):** `proximal`, `aipp`, `apd`, `qp_aipp`, `qp_apd`, `ialm`

proxax-0.0.1/README.md

@@ -0,0 +1,55 @@
+# `proxax`
+
+Certified first-order proximal point methods for unconstrained and constrained
+nonconvex composite optimization in **JAX** and **Optax**.
+
+## Installation
+
+```bash
+pip install proxax
+```
+
+Or with `uv`:
+
+```bash
+uv add proxax
+```
+
+## Problem Classes
+
+1. **Unconstrained Nonconvex Composite Optimization:**
+   $$\min_{x \in \mathcal{X}} \; f(x) + h(x)$$
+   where $f$ (`smooth`) is differentiable and $(m, L^+)$-weakly convex/smooth on
+   an arbitrary JAX PyTree $\mathcal{X}$, and $h$ (`proximal`) is proper,
+   closed, convex, and proximable.
+
+2. **Set / Cone-Constrained Nonconvex Composite Optimization:**
+   $$\min_{x \in \mathcal{X}} \; f(x) + h(x) \quad \text{s.t.} \quad \mathcal{A}(x) \in \mathcal{C}$$
+   where $\mathcal{A}: \mathcal{X} \to \mathcal{Y}$ is a linear map between
+   PyTrees and $\mathcal{C} \subseteq \mathcal{Y}$ is a closed convex set with
+   Euclidean projection $\Pi_{\mathcal{C}}$.
+
+## Quick Example
+
+```python
+import jax.numpy as jnp
+import proxax
+from proxax import constraints
+from proxax import prox
+from proxax import solvers
+
+problem = proxax.CompositeProblem(
+    smooth=lambda x: 0.5 * jnp.dot(x, q_mat @ x) + jnp.dot(c_vec, x),
+    proximal=prox.Box(lower=-1.0, upper=1.0),
+    constraints=constraints.equality(lambda x: a_mat @ x, rhs=b_vec),
+)
+
+sol = proxax.solve(problem, x0, solver=solvers.QPAIPP(opt_tol=1e-5, feas_tol=1e-5))
+```
+
+## Available Certified Solvers (`wwkong/nc_opt` Suite)
+
+- **Unconstrained Composite (`CompositeProblem`):** `AIPP`, `APD` (`PF.APD`), `ECG`, `AG`, `UPFAG`, `NCFISTA`, `AdapFISTA`, `ACACG`, `AdapAPG`
+- **Unconstrained Spectral (`SpectralProblem` / `NuclearNorm`):** `IAICG`, `DAICG`
+- **Set / Cone-Constrained (`CompositeProblem` with `constraints`):** `QPAIPP`, `QPAPD`, `IALM` (`NL-IAPIAL` / `AIDAL`), `SProxALM`, `IAIPAL`
+- **Optax Single-Step Transforms (`proxax.optax_compat`):** `proximal`, `aipp`, `apd`, `qp_aipp`, `qp_apd`, `ialm`

proxax-0.0.1/docs/api.md

@@ -0,0 +1,121 @@
+# API Reference
+
+## Top-Level Problem & Dispatcher
+
+```{eval-rst}
+.. automodule:: proxax.problem
+   :members:
+
+.. automodule:: proxax.solve
+   :members:
+```
+
+## Proximal Operators (`proxax.prox`)
+
+```{eval-rst}
+.. automodule:: proxax.prox.base
+   :members:
+
+.. automodule:: proxax.prox.norms
+   :members:
+
+.. automodule:: proxax.prox.indicators
+   :members:
+
+.. automodule:: proxax.prox.combinators
+   :members:
+```
+
+## Convex Sets & Constraints (`proxax.constraints`)
+
+```{eval-rst}
+.. automodule:: proxax.constraints.sets
+   :members:
+
+.. automodule:: proxax.constraints.linear_op
+   :members:
+```
+
+## Oracles & Stationarity Refinement
+
+```{eval-rst}
+.. automodule:: proxax.oracle.terms
+   :members:
+
+.. automodule:: proxax.oracle.composite
+   :members:
+
+.. automodule:: proxax.oracle.spectral
+   :members:
+
+.. automodule:: proxax.certificates.stationarity
+   :members:
+```
+
+## Solvers (`proxax.solvers`)
+
+```{eval-rst}
+.. automodule:: proxax.solvers.base
+   :members:
+
+.. automodule:: proxax.solvers.inner.acg
+   :members:
+
+.. automodule:: proxax.solvers.inner.pf_acg
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.aipp
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.apd
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.ecg
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.ag
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.upfag
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.nc_fista
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.adap_fista
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.ac_acg
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.adap_apg
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.ia_icg
+   :members:
+
+.. automodule:: proxax.solvers.unconstrained.da_icg
+   :members:
+
+.. automodule:: proxax.solvers.constrained.qp_aipp
+   :members:
+
+.. automodule:: proxax.solvers.constrained.qp_apd
+   :members:
+
+.. automodule:: proxax.solvers.constrained.ialm
+   :members:
+
+.. automodule:: proxax.solvers.constrained.s_prox_alm
+   :members:
+
+.. automodule:: proxax.solvers.constrained.iaipal
+   :members:
+```
+
+## Optax Compatibility (`proxax.optax_compat`)
+
+```{eval-rst}
+.. automodule:: proxax.optax_compat.transforms
+   :members:
+```

proxax-0.0.1/docs/conf.py

@@ -0,0 +1,21 @@
+"""Sphinx configuration for proxax documentation."""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
+
+project = "proxax"
+author = "William Kong"
+release = "0.0.1"
+extensions = [
+    "myst_parser",
+    "sphinx.ext.autodoc",
+    "sphinx.ext.mathjax",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.viewcode",
+]
+myst_enable_extensions = ["dollarmath", "amsmath"]
+html_theme = "furo"
+autodoc_typehints = "description"

proxax-0.0.1/docs/index.md

@@ -0,0 +1,22 @@
+# `proxax` Documentation
+
+Certified first-order proximal point methods for unconstrained and constrained
+nonconvex composite optimization in **JAX** and **Optax**.
+
+## Mathematical Scope
+
+1. **Unconstrained Nonconvex Composite & Spectral Optimization (`AIPP`, `APD`, `ECG`, `AG`, `UPFAG`, `NCFISTA`, `AdapFISTA`, `ACACG`, `AdapAPG`, `IAICG`, `DAICG`):**
+
+   $$\min_{x \in \mathcal{X}} \; f(x) + h(x)$$
+
+2. **Set / Cone-Constrained Nonconvex Composite Optimization (`QPAIPP`, `QPAPD`, `IALM`, `SProxALM`, `IAIPAL`):**
+
+   $$\min_{x \in \mathcal{X}} \; f(x) + h(x) \quad \text{s.t.} \quad \mathcal{A}(x) \in \mathcal{C}$$
+
+```{toctree}
+:maxdepth: 2
+:caption: Contents
+
+nc_opt_reference
+api
+```

proxax-0.0.1/docs/nc_opt_reference.md

@@ -0,0 +1,74 @@
+# Reference & Roadmap from `wwkong/nc_opt` and Kong's `APD`
+
+`proxax` brings the certified nonconvex composite proximal-point methodology of **[`wwkong/nc_opt`](https://github.com/wwkong/nc_opt)** (MATLAB) into an idiomatic, `jit`/`vmap`-compatible **JAX + Optax** architecture.
+
+---
+
+## 1. Pointers to `wwkong/nc_opt` & Primary Papers
+
+- **Upstream MATLAB Repository:** https://github.com/wwkong/nc_opt
+  - `src/alg_framework/` — Core outer proximal-point templates, adaptive curvature / line-search routines, inner `ACG` subroutines, and refinement maps.
+  - `src/solvers/` — Unconstrained and linearly/cone-constrained composite solvers (`AIPP`, `R-AIPP`, `APD` / `PF.APD`, `QP-AIPP`, `NL-IAPIAL`, `AIDAL`, `AG`).
+  - `tests/papers/apd/` — Reference MATLAB implementation and experiments for `PF.APD` (`APD`): https://github.com/wwkong/nc_opt/tree/master/tests/papers/apd
+- **Author Publications Page:** https://wwkong.github.io/publication/
+
+### Key Papers by Weiwei (William) Kong et al.
+
+1. **Parameter-Free Accelerated Proximal Descent (`PF.APD` / `APD`):**
+   - W. Kong, *"Complexity-Optimal and Parameter-Free First-Order Methods for Finding Stationary Points of Composite Optimization Problems,"* SIAM Journal on Optimization, 34(3):3005–3032, 2024 (DOI: `10.1137/22M1498826`, PDF: https://wwkong.github.io/files/publications/apd.pdf).
+2. **Unconstrained `AIPP` & Constrained `QP-AIPP`:**
+   - W. Kong, J. G. Melo, and R. D. C. Monteiro, *"Complexity of a quadratic penalty accelerated inexact proximal point method for solving linearly constrained nonconvex composite programs,"* SIAM Journal on Optimization, 29(4):2566–2593, 2019.
+   - W. Kong, J. G. Melo, and R. D. C. Monteiro, *"An efficient adaptive accelerated inexact proximal point method for solving linearly constrained nonconvex composite problems,"* Comput. Math. Appl., 76:305–346, 2020 (`R-QP-AIPP`).
+3. **Proximal & Dampened Augmented Lagrangian (`NL-IAPIAL` / `AIDAL`):**
+   - W. Kong, J. G. Melo, and R. D. C. Monteiro, *"Iteration-complexity of a proximal augmented Lagrangian method for solving nonconvex composite optimization problems with nonlinear convex constraints,"* Math. Oper. Res., 48(2):1066–1094, 2023.
+   - W. Kong, J. G. Melo, and R. D. C. Monteiro, *"Iteration-complexity of an inner accelerated inexact proximal augmented Lagrangian method based on the classical Lagrangian function,"* SIAM Journal on Optimization, 33(1):181–210, 2023.
+
+---
+
+## 2. Deep Dive: Kong's `APD` (`PF.APD` + `PF.ACG`, SIAM J. Optim. 2024)
+
+Unlike "Accelerated Primal-Dual" methods in saddle-point literature, Kong's **`APD`** stands for **Parameter-Free Accelerated Proximal Descent (`PF.APD`)** (https://wwkong.github.io/files/publications/apd.pdf, Algorithms 3.1–3.4). It is a premier target for `proxax`'s next unconstrained and penalty-subproblem solver:
+
+- **Truly Parameter-Free & Optimal in Both Regimes:** Requires **no knowledge** of $(m, M)$ or boundedness of $\operatorname{dom} h$, achieving $\tilde{\mathcal{O}}(\sqrt{m M}\Delta_0 / \varepsilon^2)$ resolvent evaluations when $f$ is $m$-weakly convex and automatically improving to $\tilde{\mathcal{O}}(\sqrt{M}\min\{\sqrt{\Delta_0}/\varepsilon, d_0/\sqrt{\varepsilon}\})$ when $f$ is convex.
+- **Rescaled Subproblem (`Algorithm 3.4`, Eq. (3.12)):** At outer step $k$ with lower-curvature estimate $\hat{m} > 0$, `PF.ACG` (`Algorithm 3.2`) is called with $\mu = 1/2$ on:
+
+  $$\psi^s(z) = \frac{f(z)}{2\hat{m}} + \frac{1}{2}\|z - z_k\|^2, \qquad \psi^n(z) = \frac{h(z)}{2\hat{m}}.$$
+
+- **Exact Certificate at Every Inner Iteration (No Extra `refine` Needed):**
+  In `Algorithm 3.1`, $y_{j+1} = \operatorname{prox}_{\psi^n / (L_{j+1} + \mu)}\big(\tilde{x}_j - \nabla \psi^s(\tilde{x}_j) / (L_{j+1} + \mu)\big)$ with line search $\psi^s(y_{j+1}) - \ell_{\psi^s}(y_{j+1}; \tilde{x}_j) \le \frac{L_{j+1}}{2}\|y_{j+1} - \tilde{x}_j\|^2$, and `Algorithm 3.2` forms:
+
+  $$r_{j+1} = \nabla \psi^s(y_{j+1}) - \nabla \psi^s(\tilde{x}_j) + (L_{j+1} + \mu)(\tilde{x}_j - y_{j+1}) \in \nabla \psi^s(y_{j+1}) + \partial \psi^n(y_{j+1}).$$
+
+  Rescaling gives $v = 2\hat{m}(r_{j+1} + z_k - y_{j+1}) \in \nabla f(y_{j+1}) + \partial h(y_{j+1})$ directly!
+- **Dynamic "Bad" vs. "Good" Inner Stopping Checks (Eqs. (3.5)–(3.6)):**
+  - **Bad termination (`Step 2`, Eq. (3.5)):** If $\mu A_{j+1}\|y_{j+1} - \tilde{x}_j\|^2 \le \|y_{j+1} - y_0\|^2$ or $\psi(y_0) \ge \psi(y_{j+1}) + \langle r_{j+1}, y_0 - y_{j+1}\rangle$ fails, local $\mu$-strong convexity is violated; `Algorithm 3.3` multiplies $\hat{m} \leftarrow \alpha \hat{m}$ and retries.
+  - **Good termination (`Step 3`, Eq. (3.6)):** Stops when $\|r_{j+1}\|^2 \le \sigma^2 \|y_{j+1} - y_0\|^2$ and $\|r_{j+1} + y_0 - y_{j+1}\|^2 \le \theta\big[\psi(y_0) - \psi(y_{j+1}) + \frac{1}{2}\|y_{j+1} - y_0\|^2\big]$ ($\theta > 2$, $\sigma = \rho \in (0, 1)$), certifying proximal descent (Eqs. (2.2)–(2.4)).
+- **Practical Curvature Decay Heuristic (Eq. (6.1)):**
+  In `nc_opt/tests/papers/apd`, $\alpha = \beta = 2$, $\rho = 1/\sqrt{\alpha}$, with $L_0 \leftarrow L_0 / (1 + \beta/2)$ before `PF.ACG` and $m_{k+1} \leftarrow \max\{m_0, m_{k+1} / (1 + \alpha/2)\}$ across outer steps (guarding against the `float32` noise floor in JAX).
+- **Constrained Penalty / ALM Integration (Sections 5 & 7 of `apd.pdf`):**
+  Using `APD` as the inner stage solver inside `QPAIPP` (`QP-APD`, with $m_0 = \varepsilon$) or proximal `iALM` yields a **completely parameter-free** constrained solver with $\tilde{\mathcal{O}}(\varepsilon^{-3})$ nonconvex complexity that automatically improves to $\tilde{\mathcal{O}}(\varepsilon^{-1.5})$ when $f$ is convex!
+
+---
+
+## 3. Mapping Current `proxax` Modules to `nc_opt`
+
+| `nc_opt` Concept / Component | `proxax` Module & Symbols | Notes |
+| :--- | :--- | :--- |
+| Oracle `(f, grad_f, h, prox_h)` + curvature `(m, M)` | `proxax.oracle.terms.SmoothTerm`, `proxax.prox.base.ProxOperator`, `proxax.oracle.composite.CompositeOracle` | Arbitrary JAX PyTrees + automatic `domain_project` onto $\operatorname{cl}(\operatorname{dom}(h))$. |
+| Spectral oracle `SpectralOracle.m` + `refine_ICG.m` | `proxax.oracle.spectral` (`SpectralOracle`, `SpectralProblem`, `refine_icg`), `proxax.prox.norms.NuclearNorm` | Curvature redistribution $\xi$, thin-SVD 1D vector reduction, and exact matrix subgradient refinement. |
+| Refinement $(y) \mapsto (\hat{x}, \hat{v})$ | `proxax.certificates.stationarity.refine` | Exact $\hat{v} \in \nabla f(\hat{x}) + \partial h(\hat{x})$ via one adaptive composite gradient step. |
+| Proximal subproblem $\phi(x) + \frac{1}{2\lambda_k}\|x - z_k\|^2$ | `proxax.oracle.composite.ProxRegularizedOracle`, `RescaledProxOracle` | Splits $1/\lambda_k = \alpha_{\text{smooth}} + \mu_{\text{prox}}$ (`AIPP`) or rescales $(\phi + \hat{m}\|\cdot - z_k\|^2)/(2\hat{m})$ (`APD`). |
+| Inner `ACG`, `PF.ACG`, & `AdapAPG` | `proxax.solvers.inner.acg`, `proxax.solvers.inner.pf_acg`, `proxax.solvers.unconstrained.adap_apg.AdapAPG` | Tracks $(u_t, \eta_t)$ (`ACG`), exact refinement-free $r_{j+1}$ (`PF.ACG`), and Lin & Xiao (2014) adaptive APG (`AdapAPG`). |
+| Unconstrained `AIPP` & `APD` (`PF.APD`) | `proxax.solvers.unconstrained` (`AIPP`, `APD`) | Known-$m$ / adaptive-$\lambda_k$ `AIPP` and parameter-free `PF.APD` + Eq. (6.1) decay. |
+| Direct composite & FISTA-type solvers (`ECG`, `AG`, `UPFAG`, `NC_FISTA`, `ADAP_FISTA`, `AC_ACG`) | `proxax.solvers.unconstrained` (`ECG`, `AG`, `UPFAG`, `NCFISTA`, `AdapFISTA`, `ACACG`) | Exact composite gradient, Ghadimi–Lan `AG` & `UPFAG`, Liang–Monteiro–Sim `NC-FISTA` / `ADAP-NC-FISTA`, and Liang–Monteiro `AC-ACG`. |
+| Spectral ICG solvers (`IA_ICG.m`, `DA_ICG.m`) | `proxax.solvers.unconstrained` (`IAICG`, `DAICG`) | Inner-accelerated and doubly-accelerated spectral composite solvers (Kong & Monteiro 2020). |
+| Set/cone constraint $\mathcal{A}(x) \in \mathcal{C}$ & Moreau envelope | `proxax.constraints.sets`, `proxax.constraints.linear_op`, `proxax.oracle.composite.PenalizedOracle` | Matrix or PyTree callable $\mathcal{A}$ (`jax.linear_transpose` adjoint) + any `ConvexSet` projection $\Pi_{\mathcal{C}}$. |
+| Constrained penalty & AL frameworks (`penalty.m`, `iALM.m`, `AIDAL.m`, `sProxALM.m`, `IAIPAL.m`) | `proxax.solvers.constrained` (`QPAIPP`, `QPAPD`, `IALM`, `SProxALM`, `IAIPAL`) | Geometric penalty (`QPAIPP`, `QPAPD`), dampened proximal AL (`IALM`), smoothed proximal ALM (`SProxALM`), and single-`ACG`-step proximal AL (`IAIPAL`). |
+| Optax step-by-step transforms | `proxax.optax_compat.transforms` (`proximal`, `aipp`, `apd`, `qp_aipp`, `qp_apd`, `ialm`) | Single-inner-step PyTree state machines compatible with `jax.jit` and `jax.lax.while_loop`. |
+
+---
+
+## 4. Completed `wwkong/nc_opt` Coverage & Extensions
+
+1. **Complete `nc_opt` Solver & Framework Coverage:** Every solver and outer framework in `wwkong/nc_opt` (`ACG`, `PF.ACG`, `AIPP`, `APD`, `ECG`, `AG`, `UPFAG`, `NCFISTA`, `AdapFISTA`, `ACACG`, `AdapAPG`, `SpectralOracle` + `IAICG` + `DAICG`, `QPAIPP`, `QPAPD`, `IALM` / `NL-IAPIAL` / `AIDAL`, `SProxALM`, and `IAIPAL`) is implemented with exact KKT/subdifferential stationarity certificates and `jax.jit` compatibility.
+2. **Future Extensions:** Add `sets.PSDCone` (`jnp.linalg.eigh`) and a standalone `benchmarks/` suite reproducing the paper tables from `wwkong/nc_opt/tests/papers/`.