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`
