MIXSQPX: A JAX-compatible interface to the mixSQP algorithm

This package provides a jax-compatible interface to the mixSQP algorithm described in,

Kim, Y., Carbonetto, P., Stephens, M., & Anitescu, M. (2020). A Fast Algorithm for Maximum Likelihood Estimation of Mixture Proportions Using Sequential Quadratic Programming. Journal of Computational and Graphical Statistics, 29(2), 261–273.

The implementation of the mixSQP algorithm in this package was ported from the R implementation by the authors of this paper which uses the Rarmadillo package as an interface to the Armadillo C++ library.

This package provides a function mixsolve, which can be transformed with JAX transformations such as jit.

Installation

This branch uses uv as the project frontend and hatchling as the build backend. Building the package requires a working Rust toolchain because the default backend is a bundled Rust shared library.

Local Development

# Install uv if you don't have it already.
curl -LsSf https://astral.sh/uv/install.sh | sh

# Clone the repository.
git clone https://github.com/pdiegert/mixsqpx.git
cd mixsqpx

# Install Rust if you don't have it already.
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

# Sync the development environment and install the package editable.
uv sync --group test --group examples

uv sync builds the Rust library through the Hatch build hook and installs the package in editable mode.

Verification

Run the test suite to verify your installation:

uv run --group test pytest tests

Build a wheel locally with:

uv build

Benchmark the Rust backend explicitly with:

uv run --group test python scripts/benchmark_mixsolve.py

Solver accuracy and diagnostics

The Rust solver certifies the full simplex KKT conditions against the original component matrix, including the zero-weight components. The SVD approximation is used for search curvature; the objective, convergence certificate and returned derivatives use the full regularized matrix. Exact active-face Newton steps are used near convergence within the same SQP solve.

tol_svd is an absolute singular-value cutoff. Setting it to zero disables the low-rank approximation. density_offset is a nonnegative offset after row scaling (default 1e-8); zero is allowed. With row-scaled components A, the simplex objective is -mean(log(A @ probabilities + density_offset)). With observation_weights, the mean uses weights normalized to sum to one. The implementation uses the equivalent component matrix A + density_offset, which preserves the homogeneous SQP reformulation and multiplier-one KKT conditions. Gradients and Hessians are with respect to this component matrix at the actual returned, normalized weights.

result = mixsolve(
    log_component_likelihood,
    tol_svd=1e-10,
    density_offset=1e-9,
    control_sqp=1e-9,
    return_info=True,
)
weights = result.probabilities
info = result.info

The typed diagnostics expose status, iterations, active_residual, inactive_residual, simplex_residual, and svd_rank. Status is 0 for certified convergence, 1 for the iteration limit, 2 for stagnation, and 3 for invalid input. Requesting diagnostics permits inspection of an uncertified candidate. Without return_info=True, failure to meet the requested KKT tolerance raises an error. The default successful return remains (probabilities, gradient, hessian). Invalid likelihood inputs return NaN arrays, with status 3 when diagnostics are requested. Restart the Python process after rebuilding the native library.

Consumers differentiating a profiled objective must apply the same offset to the outer objective and implicit derivatives. A positive offset changes the criterion and its row-scaling derivative must be retained. Its magnitude and any sample-size rule are caller-controlled.

Observation weights

Pass a length-nobs array using the keyword-only observation_weights argument:

result = mixsolve(
    log_component_likelihood,
    observation_weights=observation_weights,
    density_offset=1e-9,
    control_sqp=1e-9,
    return_info=True,
)

Omitting this argument uses uniform observation weights. Supplied weights must be finite and nonnegative, with at least one positive entry. They are normalized inside Rust, so multiplying all weights by a positive constant leaves the criterion, fitted masses, derivatives and KKT tolerances unchanged. Normalization also handles finite weights whose unscaled sum would overflow.

Zero-weight observations are excluded from the solve and SVD, including input validation of their log likelihoods. Negative, nonfinite or all-zero weights raise an error. The gradient and Hessian remain derivatives with respect to the mixture masses, evaluated at the returned masses for the normalized weighted criterion. Integer weights are equivalent to repeating observations when the same density_offset is used. The solver does not choose a sample-size rule for the offset.

Observation weights are dynamic JAX inputs and can change between JIT calls. They require the Rust backend and native ABI 3; rebuild the library and restart Python when updating an editable checkout. Profiling layers implement implicit differentiation through the optimal masses.

Legacy C++ Reference Path

This branch is intended to merge as a Rust-first package. The legacy C++/Armadillo backend is kept only as an archival benchmark/reference path.

• Keep the legacy implementation in a separate checkout or worktree.
• Use a separate environment for the legacy package.
• Do not treat the legacy C++ build as part of the default developer install on this branch.

A recommended setup is:

git worktree add ../mixsqpx-legacy legacy-cpp-final

Additional details for benchmarking and the opt-in exact legacy reference test are documented in LEGACY_CPP_REFERENCE.md.

License

This project is licensed under the MIT License - see the LICENCE file for details.

Acknowledgements

This project includes code adapted from the mixSQP project, which is available under the MIT license. We thank the original authors (Youngseok Kim, Peter Carbonetto, Matthew Stephens and Mihai Anitescu) for their work.