Fit with optimistix BFGS and Levenberg-Marquardt

optimistix solves minimization and least-squares problems on JAX pytrees. BFGS minimizes potential from flatppl.Module.logdensity() in unconstrained coordinates. Levenberg-Marquardt needs residuals rather than a density, so a FlatPPL query supplies them. The model is a Normal likelihood with unknown mean and precision under a conjugate Normal-Gamma prior; every fitted point has a closed form.

Runtime budget: 120 s.

from math import lgamma, log, pi
from typing import Any

import jax
import jax.numpy as jnp
import numpy as np
import optimistix as optx

from flatppl import flatppl

The model and its closed forms

The likelihood L peaks at the sample mean and tau = n / S with S = sum((y - mean)^2). The posterior peaks at m_n and (a_n - 1/2) / b_n.

data = np.array([1.2, 0.4, 2.1, 1.7, 0.9, 1.5, 0.3, 1.1])
model = flatppl(
    """
    tau ~ Gamma(shape = 3.0, rate = 2.0)
    mu ~ Normal(mu = 0.0, sigma = 1.0 / sqrt(2.0 * tau))
    y ~ iid(Normal(mu = mu, sigma = 1.0 / sqrt(tau)), 8)
    observed = external(cartpow(reals, 8))
    prior = lawof(record(mu = mu, tau = tau))
    L = likelihoodof(kernelof(record(y = y), mu = mu, tau = tau), record(y = observed))
    posterior = bayesupdate(L, prior)
    """
).set(observed=data)

a0, b0, k0, m0, n = 3.0, 2.0, 2.0, 0.0, data.size
mean = data.mean()
squares = ((data - mean) ** 2).sum()
kn, an = k0 + n, a0 + n / 2
mn = (k0 * m0 + n * mean) / kn
bn = b0 + 0.5 * squares + k0 * n * (mean - m0) ** 2 / (2 * kn)
mle = {"mu": mean, "tau": n / squares}
mode = {"mu": mn, "tau": (an - 0.5) / bn}

BFGS in unconstrained coordinates

jacobian=False keeps each density itself in the free coordinates (mu, log tau), so its minimizer maps through constrain to the mode in model coordinates. optimistix passes args to the function; here it is unused.

Float64 is used here because BFGS stops when the function stops changing, which in float32 happens about 1e-4 short of the optimum. FlatPPL itself does not require it. jax.enable_x64 scopes it to these fits.

def potential_of(fit: Any) -> Any:
    """Wrap ``fit.potential`` in optimistix's ``fn(y, args)`` form."""

    def potential(position: Any, args: None) -> jax.Array:
        return fit.potential(position)

    return potential


solver = optx.BFGS(rtol=1e-9, atol=1e-9)
for name, measure, expected in [("MAP", "posterior", mode), ("MLE", "L", mle)]:
    with jax.enable_x64(True):
        fit = model.logdensity(measure, jacobian=False, dtype="float64")
        solution = optx.minimise(potential_of(fit), solver, fit.init)
        point = {
            key: float(value) for key, value in fit.constrain(solution.value).items()
        }
    print(name, {key: round(value, 8) for key, value in point.items()})
    for key, value in expected.items():
        assert abs(point[key] - value) < 1e-8 * abs(value), (name, key)

MAP {'mu': 0.92, 'tau': 1.47794452}
MLE {'mu': 1.15, 'tau': 2.98507463}

Levenberg-Marquardt on residuals

A least-squares solver wants the residual vector, not its summed square. This query returns observed - mu for a given mu; its least-squares solution is the sample mean, the maximum likelihood mu for any tau. It runs in float64 for the same stopping reason.

with jax.enable_x64(True):
    residuals = (
        flatppl(
            """
            observed = external(cartpow(reals, 8))
            mu = elementof(reals)
            inputs = mu
            outputs = observed .- mu
            """
        )
        .set(observed=data)
        .compile(dtype="float64")
    )
    least_squares = optx.least_squares(
        lambda mu, args: residuals(mu),
        optx.LevenbergMarquardt(rtol=1e-9, atol=1e-9),
        jnp.float64(0.0),
    )
    fitted_mean = float(least_squares.value)
print("Levenberg-Marquardt mu:", round(fitted_mean, 8))
assert abs(fitted_mean - mean) < 1e-8 * mean

Levenberg-Marquardt mu: 1.15

The posterior’s normalization

The posterior in free coordinates, with its Jacobian, integrates to the closed-form Normal-Gamma evidence.

log_evidence = (
    lgamma(an) - lgamma(a0) + a0 * log(b0) - an * log(bn)
    + 0.5 * log(k0 / kn) - n / 2 * log(2 * pi)
)  # fmt: skip
density = model.logdensity("posterior")
mu_axis = jnp.linspace(-3.0, 5.0, 401)
log_tau_axis = jnp.linspace(-4.0, 3.0, 401)
grid = jax.vmap(
    lambda m: jax.vmap(lambda t: density({"mu": m, "tau": t}))(log_tau_axis)
)(mu_axis)
total = np.trapezoid(
    np.trapezoid(np.exp(np.asarray(grid)), np.asarray(log_tau_axis), axis=1),
    np.asarray(mu_axis),
)
print(f"log evidence: grid {np.log(total):.6f}, closed form {log_evidence:.6f}")
assert abs(np.log(total) - log_evidence) < 1e-6

log evidence: grid -10.558731, closed form -10.558731