Metadata-Version: 2.5
Name: nlls-gram
Version: 2.12.1
Summary: Metric-aware Levenberg-Marquardt nonlinear least-squares for JAX
Project-URL: Homepage, https://github.com/HighDimensionalEconLab/nlls_gram
Project-URL: Repository, https://github.com/HighDimensionalEconLab/nlls_gram
Project-URL: Documentation, https://highdimensionaleconlab.github.io/nlls_gram/
Project-URL: Issues, https://github.com/HighDimensionalEconLab/nlls_gram/issues
Author-email: Jesse Perla <jesseperla@gmail.com>
License-Expression: MIT
License-File: LICENSE
Keywords: jax,levenberg-marquardt,metric,nonlinear-least-squares,optimization
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Mathematics
Requires-Python: >=3.12
Requires-Dist: jax>=0.11.2
Description-Content-Type: text/markdown

# nlls_gram

[![CI](https://github.com/HighDimensionalEconLab/nlls_gram/actions/workflows/ci.yml/badge.svg)](https://github.com/HighDimensionalEconLab/nlls_gram/actions/workflows/ci.yml)
[![Docs](https://github.com/HighDimensionalEconLab/nlls_gram/actions/workflows/docs.yml/badge.svg)](https://highdimensionaleconlab.github.io/nlls_gram/)
[![PyPI](https://img.shields.io/pypi/v/nlls-gram.svg)](https://pypi.org/project/nlls-gram/)
[![Python versions](https://img.shields.io/pypi/pyversions/nlls-gram.svg)](https://pypi.org/project/nlls-gram/)
[![License: MIT](https://img.shields.io/github/license/HighDimensionalEconLab/nlls_gram)](https://github.com/HighDimensionalEconLab/nlls_gram/blob/main/LICENSE)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)

Levenberg-Marquardt nonlinear least squares for JAX pytrees, aimed at solving
systems of equations. The core use case is **underdetermined** systems — more
parameters than residuals — where a zero-residual root is not unique and
something must select *which* interpolant is returned. Two solvers differ in
where that selection lives:

- **`RidgeLevenbergMarquardt`** puts it in the **objective**, minimizing
  \(\|r(x)\|^2 + \lambda\,\|x_m\|_W^2\) for a positive-definite metric
  \(W\) on the metric block of \(x = [x_m; x_f]\) (the free block stays
  unpenalized). Annealing \(\lambda\) toward zero converges to the
  minimum-seminorm — e.g. minimum-RKHS-norm — interpolant, by classical
  nonlinear Tikhonov regularization. Every inner problem is a well-posed NLLS.
- **`LevenbergMarquardt`** puts it in the **damping geometry**: standard
  damped LM whose trust region is measured in \(W\), so the small-damping
  Gauss-Newton limit is the minimum-\(W\)-norm correction.

Both take a residual over `(x)`, `(x, args)`, or `(x, args, p)`, flatten any
pytree `x`, expose per-step `update(...)` and an internally jitted `solve(...)`
loop with callbacks and multi-start, and differentiate `solve(...).x` with
respect to `p` through a custom implicit rule — no unrolling.

## Install

```bash
uv add nlls-gram
```

For GPU use, install the JAX accelerator build that matches your hardware:

```bash
uv add nlls-gram "jax[cuda13]"
```

## Minimum-RKHS-norm interpolation

```python
import jax.numpy as jnp
from nlls_gram import AnnealRidge, RidgeLevenbergMarquardt, RepeatedFactorMetric

# W = blockdiag(K, K): the RKHS seminorm over two coefficient blocks. The
# constructor takes the FACTOR; shift a semidefinite K by epsilon*I first.
metric = RepeatedFactorMetric(jnp.linalg.cholesky(K, upper=True), repeats=2)

solver = RidgeLevenbergMarquardt(collocation_residual, metric=metric, ridge=1e-4)

# Anneal the ridge toward the interpolating limit on stationarity.
anneal = AnnealRidge(ridge_floor=1e-10)
result = solver.solve(
    x0, callback=anneal, user_state=anneal.init_state(), gtol=1e-8, atol=1e-8
)
```

## General nonlinear least squares

```python
import jax, jax.numpy as jnp
from nlls_gram import LevenbergMarquardt


def residual(x, args, p):
    return args["design"] @ x - p["target"]


solver = LevenbergMarquardt(residual)
result = solver.solve(
    jnp.zeros(8), {"design": design}, p={"target": y}, max_steps=200, atol=1e-8
)

# The solution is differentiable in p, at a cost independent of max_steps.
sensitivity = jax.grad(
    lambda p: jnp.sum(
        solver.solve(jnp.zeros(8), {"design": design}, p=p, max_steps=200, atol=1e-8).x
        ** 2
    )
)({"target": y})
```

Pass `metric=` to weight the damping geometry.

## Linear solvers

The same typed configs serve both solvers, forward and in the implicit-AD
role: `Cholesky()` (the default, auto-selecting the smaller of the dual and
normal systems), `QR()` (damping-row QR — stable at tiny damping and
rank-safe), `CG(precond)` and `GramCG(precond)` (matrix-free in parameter and
residual space), and `SVD()` for rank-deficient tangents. A knob that exists
for only one method is a field on that method, so it cannot be passed with
another.

## Row-wise residuals

For an expensive per-state residual vmapped over a simulated path, replace
that `jax.vmap` with `vmap_rows`: the same values, with each row's gradient
computed once per linearization and reused across the Jacobian assembly
instead of one pass of the row function per row and Jacobian column (or
cotangent row).

```python
from nlls_gram import vmap_rows


@vmap_rows(in_axes=(0, None, None))
def euler_residual(state, policy, p): ...
```

## Docs

<https://highdimensionaleconlab.github.io/nlls_gram/>

## License

MIT.
