mixsqpx 0.2.0a1__tar.gz

mixsqpx-0.2.0a1/.gitignore

@@ -0,0 +1,164 @@
+# misc other
+# scrap/*
+wheelhouse/
+.codex
+
+# pixi environments
+.pixi
+*.egg-info
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+# lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+Cargo.lock
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+

mixsqpx-0.2.0a1/Cargo.toml

@@ -0,0 +1,3 @@
+[workspace]
+members = ["rust/mixsqpx_rust"]
+resolver = "2"

mixsqpx-0.2.0a1/LEGACY_CPP_REFERENCE.md

@@ -0,0 +1,41 @@
+# Legacy C++ Reference Workflow
+
+This branch is Rust-first. The legacy C++/Armadillo backend is preserved only as
+an archival benchmark and reference target.
+
+## Recommended Archival Setup
+
+Create a sibling worktree from the frozen legacy branch:
+
+```bash
+git worktree add ../mixsqpx-legacy legacy-cpp-final
+```
+
+Set up a separate environment in that worktree and build/install the legacy
+package there. Keep that environment separate from the Rust-first `uv` workflow
+used on this branch.
+
+## Benchmarking
+
+On the Rust-first branch:
+
+```bash
+uv run --group test python scripts/benchmark_mixsolve.py
+```
+
+If you want to compare against the legacy backend, do that from the archival
+checkout using its own build and benchmark tooling. The merged branch should not
+carry the Armadillo/OpenBLAS/CMake dependency stack as part of its default
+install path.
+
+## Exact Legacy Reference Tests
+
+The first-step exact legacy parity test is disabled by default on the Rust-first
+branch. Re-enable it only in an environment that actually has the legacy C++
+backend available:
+
+```bash
+MIXSQPX_BACKEND=cpp \
+MIXSQPX_RUN_LEGACY_CPP_REFERENCE_TESTS=1 \
+uv run --group test pytest tests/test_first_sqp_reference.py
+```

mixsqpx-0.2.0a1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Jackson Bunting, Paul Diegert, and Arnaud Maurel
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mixsqpx-0.2.0a1/PKG-INFO

@@ -0,0 +1,119 @@
+Metadata-Version: 2.4
+Name: mixsqpx
+Version: 0.2.0a1
+Summary: JAX interface to the Rust implementation of the mixSQP algorithm
+Author-email: Paul Diegert <pdiegert@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2025 Jackson Bunting, Paul Diegert, and Arnaud Maurel
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Requires-Python: <3.14,>=3.12
+Requires-Dist: jax<=0.9.2,>=0.9.2
+Requires-Dist: jaxlib<=0.9.2,>=0.9.2
+Description-Content-Type: text/markdown
+
+# 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](https://github.com/stephenslab/mixSQP) 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](https://github.com/jax-ml/jax) transformations such as `jit`.
+
+## Installation
+
+This branch uses [`uv`](https://docs.astral.sh/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
+
+```bash
+# 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:
+
+```bash
+uv run --group test pytest tests
+```
+
+Build a wheel locally with:
+
+```bash
+uv build
+```
+
+Benchmark the Rust backend explicitly with:
+
+```bash
+uv run --group test python scripts/benchmark_mixsolve.py
+```
+
+## 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:
+
+```bash
+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](LEGACY_CPP_REFERENCE.md).
+
+## License
+
+This project is licensed under the MIT License - see the [LICENCE](LICENCE) file for details.
+
+## Acknowledgements
+
+This project includes code adapted from the [mixSQP](https://github.com/stephenslab/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.

mixsqpx-0.2.0a1/README.md

@@ -0,0 +1,87 @@
+# 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](https://github.com/stephenslab/mixSQP) 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](https://github.com/jax-ml/jax) transformations such as `jit`.
+
+## Installation
+
+This branch uses [`uv`](https://docs.astral.sh/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
+
+```bash
+# 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:
+
+```bash
+uv run --group test pytest tests
+```
+
+Build a wheel locally with:
+
+```bash
+uv build
+```
+
+Benchmark the Rust backend explicitly with:
+
+```bash
+uv run --group test python scripts/benchmark_mixsolve.py
+```
+
+## 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:
+
+```bash
+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](LEGACY_CPP_REFERENCE.md).
+
+## License
+
+This project is licensed under the MIT License - see the [LICENCE](LICENCE) file for details.
+
+## Acknowledgements
+
+This project includes code adapted from the [mixSQP](https://github.com/stephenslab/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.

mixsqpx-0.2.0a1/RUST_PORT_PLAN.md

@@ -0,0 +1,256 @@
+# Rust Port Plan
+
+## Goal
+
+Replace the current C++/Armadillo/OpenBLAS/CMake extension with a Rust implementation of the solver and a simpler build path, while keeping the Python/JAX API stable.
+
+The target end state is:
+
+- numeric core implemented in Rust
+- dense linear algebra implemented with a Rust-native stack
+- Python package built without CMake, pybind11, Armadillo, or OpenBLAS
+- JAX custom call registered from Python by loading a Rust-built shared library
+
+## Current Status
+
+### Milestone Status
+
+- Milestone 0: complete
+- Milestone 1: complete
+- Milestone 2: complete
+- Milestone 3: complete
+- Milestone 4: accepted as complete
+- Milestone 5: mostly complete
+- Milestone 6: partially started
+
+### Current Technical Direction
+
+The current implementation differs slightly from the original proposal:
+
+- Linear algebra currently uses `nalgebra` and this is the accepted backend for now
+- Shared library build is Cargo `cdylib`
+- Python integration loads the Rust shared library with `ctypes`
+- JAX FFI registration is done from Python through `jax.ffi.register_ffi_target`
+- Near-term packaging direction is `uv` as the frontend and `hatchling` with a cargo build hook as the backend
+- Long-term packaging direction is to consider a cleaner `maturin`-based layout after a deliberate binding refactor
+
+The next packaging milestone can still replace the legacy C++ build, but the present Rust port is no longer just a scaffold: it is an end-to-end backend. A short `faer` spike was completed and retained on the `faer-spike` branch for reference, but `nalgebra` remains the preferred backend for the current implementation.
+
+## Implemented Work
+
+### Milestone 0: Tooling and Repository Setup
+
+Completed:
+
+- Rust toolchain installed and verified locally
+- dedicated `rust-port` branch and worktree created
+- Rust crate added under `rust/mixsqpx_rust`
+- Cargo build produces a shared library artifact for the JAX FFI path
+
+### Milestone 1: Characterize Current Behavior
+
+Completed:
+
+- characterization fixtures added for representative solver runs
+- first-SQP-step reference fixture added for the tuned Gaussian case
+- benchmark harness added for repeatable timing runs
+- parity tests added around solver outputs instead of relying only on smoke tests
+
+Artifacts:
+
+- `tests/fixtures/characterization/*.npz`
+- `tests/fixtures/reference/gaussian_tuned_first_sqp_step.npz`
+- `tests/test_characterization.py`
+- `tests/test_first_sqp_reference.py`
+- `scripts/benchmark_mixsolve.py`
+
+### Milestone 2: Rust Scaffolding and FFI Path
+
+Completed:
+
+- Rust `cdylib` exposes a JAX-compatible entry point
+- Python can load and register the Rust shared library
+- Rust backend can be selected from Python/JAX without changing the public solver interface
+- notebook and test workflows can exercise the Rust backend directly, while legacy C++ comparison remains archival/optional
+
+Artifacts:
+
+- `rust/mixsqpx_rust/src/lib.rs`
+- `src/mixsqpx/rust_backend.py`
+- `tests/test_rust_backend.py`
+
+### Milestone 3: Port Low-Risk Numeric Pieces
+
+Completed:
+
+- helper numeric routines ported to Rust
+- objective evaluation ported to Rust
+- EM update path ported to Rust
+
+Artifacts:
+
+- `rust/mixsqpx_rust/src/numeric.rs`
+- `rust/mixsqpx_rust/src/objective.rs`
+- `rust/mixsqpx_rust/src/mixem.rs`
+
+### Milestone 4: Port Core SQP Solver
+
+Completed and accepted.
+
+Completed work:
+
+- exact SQP solver path ported to Rust
+- truncated-SVD solver path ported to Rust
+- active-set QP routine ported to Rust
+- backtracking line search ported to Rust
+- existing Python/JAX API controls wired through to the Rust solver options
+
+Artifacts:
+
+- `rust/mixsqpx_rust/src/solver.rs`
+- `rust/mixsqpx_rust/src/lib.rs`
+
+## Milestone 4 Outcome
+
+### What Matched Well
+
+The Rust backend matches the legacy implementation closely in the easier and default regimes:
+
+- default characterization cases pass the current parity thresholds
+- first-SQP reference diagnostics show that the SQP point itself matches extremely closely before the active-set update
+- output plots for the Gaussian example are visually indistinguishable between C++ and Rust
+
+### What Did Not Match Exactly
+
+We were not able to achieve full numerical equality with the C++/Armadillo implementation in the tuned Gaussian characterization case.
+
+What the diagnostics showed:
+
+- EM behavior matched
+- the SQP point derivatives matched closely
+- the first meaningful drift appeared inside the active-set / linear-solve numerics
+- later SQP iterations amplified that drift in the tuned case
+
+The remaining differences appear to come from solver numerics rather than wiring, preprocessing, or objective/gradient construction.
+
+### Acceptance Decision
+
+We are accepting the current Rust implementation for Milestone 4.
+
+Reasoning:
+
+- the Rust backend produces equal or better practical results for the problem configurations we care about right now
+- the outputs are visually indistinguishable in the notebook comparisons for the main Gaussian example
+- the tuned-case differences are numerically real, but they are not currently considered product-blocking
+- further work to force Armadillo-level numerical equality does not appear justified relative to the cost
+
+This means Milestone 4 is considered complete under an acceptance criterion of practical parity rather than bitwise or solver-path equality.
+
+## Testing and Benchmark Results
+
+### Functional Testing
+
+Completed testing includes:
+
+- characterization tests against saved C++ outputs
+- first-SQP diagnostic tests on the tuned Gaussian case
+- Rust backend smoke tests through JAX FFI
+- notebook comparisons for both the default Gaussian example and the tuned Gaussian fixture
+
+### Performance Results
+
+Initial notebook timing on a Rust debug build was misleadingly poor.
+After switching the notebooks to use Rust release builds, the Rust backend performed slightly better than the C++ backend in current local runs.
+
+Observed notebook result for the Gaussian example:
+
+- Rust speedup of roughly `1.2x` to `1.5x` versus the C++ backend, depending on the run
+
+Current conclusion:
+
+- the Rust implementation is at least competitive for the tested `4000 x 101` workload
+- no immediate performance blocker is preventing the packaging switch
+- more systematic benchmarking is still appropriate in Milestone 6
+
+### `faer` Spike Result
+
+A short follow-up spike tested `faer` as a drop-in replacement for the current `nalgebra` touchpoints.
+
+Outcome:
+
+- the spike branch is kept as `faer-spike` for reference
+- the default Gaussian case remained strong and was about `1.5x` faster than the C++ backend in the quick notebook-style comparison
+- the tuned Gaussian case regressed badly, running at about `0.2x` of the C++ backend in the same quick comparison
+- the spike also introduced a small Hessian-tolerance regression in one strict first-SQP test
+
+Decision:
+
+- continue with `nalgebra` for now
+- treat `faer` as an optional future optimization experiment rather than the planned default backend
+
+## Remaining Work
+
+### Milestone 5: Switch Packaging and Archive the C++ Build
+
+Implemented:
+
+- replaced the primary `scikit-build-core`/CMake native build path with `hatchling`
+- switched developer/frontend workflows from `pixi` to `uv`
+- added a cargo-driven Hatch build hook that builds the Rust `cdylib` and bundles it into the wheel
+- updated package metadata and install instructions around the Rust backend
+- removed Armadillo/OpenBLAS logic from the default package build path
+- updated the primary wheel-building workflow to use `uv build` and a Rust toolchain directly
+- created a frozen archival reference branch locally as `legacy-cpp-final`
+- documented the archival benchmark/reference workflow in `LEGACY_CPP_REFERENCE.md`
+- kept the characterization fixtures in-repo as the lightweight always-available reference
+- made the exact first-step legacy parity test opt-in rather than part of the default Rust test path
+
+Remaining tasks:
+
+- verify the updated CI workflow in GitHub rather than only locally
+- finish cleaning up notebook/dev tooling that still assumes older local environment layouts
+- decide whether to keep any additional legacy-comparison helpers in-tree beyond the documented archival path
+
+Design intent:
+
+- the branch that eventually merges should be Rust-only by default
+- the legacy C++ backend should be preserved for benchmarking and archaeology, but not carried as an always-supported mainline dependency stack
+- benchmarking against the C++ implementation should remain possible without forcing Armadillo/OpenBLAS/CMake into the everyday development and install path
+- the immediate packaging cutover should minimize binding churn by keeping the current `ctypes` + shared-library loading model
+- a later refactor can revisit the Python/Rust boundary and move toward `maturin` if the team wants a more idiomatic long-term Rust packaging model
+
+Exit criteria:
+
+- clean install works without system BLAS or Armadillo
+- local package build path is materially simpler than the current one
+- default package/backend path uses Rust rather than the legacy C++ extension
+- a documented archival path exists for reproducing legacy C++ benchmarks when needed
+
+### Milestone 6: Performance Validation and Cleanup
+
+Remaining tasks:
+
+- run more systematic benchmarks across representative matrix sizes
+- profile allocation hotspots in the Rust SQP path
+- keep `nalgebra` as the current backend unless future profiling justifies another spike
+- document the accepted numerical differences from the legacy backend
+- remove dead compatibility code once the packaging transition is complete
+
+Exit criteria:
+
+- clear documented performance story for target workloads
+- accepted numerical behavior documented
+- remaining legacy-only code removed or explicitly retained as transitional support
+
+## Risks and Open Questions
+
+- The main unresolved product question is packaging, not solver feasibility.
+- Larger dense problems may still reveal performance or numeric tradeoffs that do not show up in the `4000 x 101` case.
+- A `faer` spike was completed and archived on its own branch; it did not beat the current `nalgebra` implementation consistently enough to justify switching now.
+
+## Immediate Next Steps
+
+1. Update packaging so the Rust backend becomes the primary build path.
+2. Preserve the current parity fixtures and notebooks as acceptance evidence for the Rust backend.
+3. Expand benchmarks across a few larger matrix sizes before final cleanup.
+4. Implement the `uv` + `hatchling` + cargo-hook packaging path first, and treat a future `maturin` migration as a separate refactor once the Rust-only package path is stable.