Reimplementing the code using JAX

[htjb]

The goal is to reimplement maxsmooth in JAX to hopefully speed up the code and modernise the package with type hinting, more standardised docstrings and formatting with Ruff. An example call to the new code is shown below. Currently it tests every possible combination of signs on the derivatives.

"""Example use case for the qp solver."""

import time

import jax
import matplotlib.pyplot as plt
import tqdm
from jax import numpy as jnp

from maxsmooth.derivatives import make_derivative_functions
from maxsmooth.models import normalised_polynomial, normalised_polynomial_basis
from maxsmooth.qp import qp

jax.config.update("jax_enable_x64", True)

function = normalised_polynomial
basis_function = normalised_polynomial_basis

key = jax.random.PRNGKey(0)
x = jnp.linspace(50, 150, 100)
y = 5e6 * x ** (-2.5) + 0.01 * jax.random.normal(key, x.shape)
N = 10
pivot_point = len(x) // 2

qp = jax.jit(
    qp, static_argnames=("N", "pivot_point", "function", "basis_function")
)
start = time.time()
sol = qp(x, y, N, pivot_point, function, basis_function)
end = time.time()
print(f"First Call: QP solved in {end - start:.5f} seconds")

start = time.time()
sol = qp(x, y, N, pivot_point, function, basis_function)
end = time.time()
print(f"Second Call: QP solved in {end - start:.5f} seconds")

vmapped_function = jax.vmap(function, in_axes=(0, None, None, None))

objective_values = []
for i in tqdm.tqdm(range(len(sol["params"]))):
    params = sol["params"][i]
    # plt.plot(x, y, 'o', label='data')
    fit = vmapped_function(x, x[pivot_point], y[pivot_point], params)
    obj_val = jnp.sum((y - fit) ** 2)
    objective_values.append(obj_val)


plt.plot(objective_values, "o-")
plt.xlabel("Solution index")
plt.ylabel("Objective value")
plt.title("Objective values for different constraint sign combinations")
plt.show()

Key Changes:

• Replacing CVXOPT with Jaxopt
• Rewritten models in with Jax numpy and jit compiled
• Using Jax grad and Jax jacobian to get the gradients and prefactors on the model parameters for the gradient e.g. the $G$ in $Ga \leq h$
• Vectorised call to fit different sign spaces
• Added pyproject.toml and uv files.
• Added pre-commit-config.yaml
• Added docstrings, type hinting and linting
• Added a linting workflow and a workflow to check version number
• Added a PR template
• Added Contribution guidelines

To do:

• Update the documentation
• Re-write tests
• Re-write the best_basis class as a function
• Examples need rewriting
• implement some form of the sign exploration/flipping

[htjb]

Jaxopt is deprecated so I'm not sure it is a good idea to rely on it. There is another code called qpax which might be a viable alternative.

https://github.com/kevin-tracy/qpax

[htjb]

MPAX also looks good (https://github.com/MIT-Lu-Lab/MPAX).

After some initial testing, qpax seems to struggle to converge.