0.1.0 release

.github/workflows/run_tests.yml

@@ -7,7 +7,7 @@ jobs:
   run-test:
     strategy:
       matrix:
-        python-version: [ "3.10", "3.12" ]
+        python-version: [ "3.11", "3.13" ]
         os: [ ubuntu-latest ]
       fail-fast: false
     runs-on: ${{ matrix.os }}

.gitignore

@@ -12,3 +12,5 @@ examples/data
 .DS_Store
 .devcontainer
 .benchmarks
+uv.lock
+.python-version

.pre-commit-config.yaml

@@ -33,4 +33,4 @@ repos:
     rev: v1.1.405
     hooks:
     - id: pyright
-      additional_dependencies: ["equinox", "jax", "lineax", "pytest", "optax", "diffrax"]
+      additional_dependencies: ["equinox", "jax", "lineax", "pytest", "optax", "diffrax", "sif2jax", "fire", "matplotlib"]

CONTRIBUTING.md

@@ -4,7 +4,7 @@ Contributions (pull requests) are very welcome! Here's how to get started.
 
 ---
 
-**Getting started**
+## Getting started
 
 First fork the library on GitHub.
 
@@ -19,7 +19,7 @@ pre-commit install  # `pre-commit` is installed by `pip` on the previous line
 
 ---
 
-**If you're making changes to the code:**
+### If you're making changes to the code:**
 
 Now make your changes. Make sure to include additional tests if necessary.
 
@@ -30,6 +30,13 @@ pip install -e '.[tests]'
 pytest  # `pytest` is installed by `pip` on the previous line.
 ```
 
+If your changes could affect solver (or compilation) performance, please run the benchmark tests with 
+
+```bash
+pytest benchmarks/ --benchmark-only
+```
+
+You can run benchmarks before or after your change, and also save more extensive results for analysis. For more on this, skip to the "Benchmarking" section below. 
 Then push your changes back to your fork of the repository:
 
 ```bash
@@ -40,7 +47,7 @@ Finally, open a pull request on GitHub!
 
 ---
 
-**If you're making changes to the documentation:**
+### If you're making changes to the documentation:
 
 Make your changes. You can then build the documentation by doing
 
@@ -50,3 +57,44 @@ mkdocs serve
 ```
 
 You can then see your local copy of the documentation by navigating to `localhost:8000` in a web browser.
+
+
+## Benchmarking
+
+If you're interested in more extensive benchmarking - for instance when contributing a new solver - this section is for you. (Note that benchmarks are not run by default, and `--benchmark-only` is required to override this.)
+
+You can save benchmark results with
+
+```
+pytest benchmarks/ --benchmark-save=<benchmark_name> --benchmark-only
+```
+
+and compare against previous runs with `pytest --benchmark-compare`, which will automatically pull in the last saved commit, but also takes run iDs as arguments. See the `pytest-benchmark` [documentation](https://pytest-benchmark.readthedocs.io/en/latest/usage.html#commandline-options) for more command line options. 
+The `benchmark-autosave` option will specify the commit iD, instead of a user-defined name.
+Make sure that you are running benchmarks with a clean working tree, so you can trace how changes affect performance!
+
+For convenience, we support some custom flags: 
+
+- `--min-dimension=<int>, --max-dimension=<int>` benchmarks can be run on a subset of problems based on problem size.
+- `--scipy` benchmarks of our solvers are run against the corresponding Python implementation. You might want to limit problem dimension here, they can be quite slow.
+
+pytest's `-k` flags also work in this setting to enable selective execution of benchmarking functions.
+
+**Analysing benchmark results**
+
+You can find a script to analyse benchmark results in `benchmarks/profile.py`. Run it with
+
+```bash
+python benchmarks/profile.py <platform> <python_version> <precision> <iD> <kind> *solver_names
+```
+
+Where platform refers to the platform on which the benchmarks were run (e.g. Darwin), precision is the numerical precision, e.g. 32bit, and iD is the benchmark run, a four-digit integer.
+These are necessary to identify the saved results for the specific run. `kind` specifies if `runtime` or `compilation` benchmarks are to be compared, and solver names should be given as strings. These are defined in `benchmarks/test_benchmarks.py` for every benchmarked solver, e.g. `optx.BFGS`. Putting this together, an example call would be
+
+```bash
+python benchmarks/profile.py Darwin 3.13 64bit 0001 runtime optx.BFGS optx.LBFGS
+```
+
+**If you are contributing a solver**
+
+In this case, you're probably reasonably familiar with the alternatives out there - if implementations we could compare to exist, please add them to the listed solvers in `benchmarks/test_benchmarks.py`, including hyperparameters such as solver tolerances to get as fair of a comparison as is feasible.

README.md

@@ -17,7 +17,7 @@ Features include:
 pip install optimistix
 ```
 
-Requires Python 3.10+ and JAX 0.4.38+ and [Equinox](https://github.com/patrick-kidger/equinox) 0.11.11+.
+Requires Python 3.11+.
 
 ## Documentation
 

benchmarks/conftest.py

@@ -0,0 +1,59 @@
+import os  # noqa: I001
+
+os.environ["EQX_ON_ERROR"] = "nan"  # Make sure this is set before importing equinox
+import equinox.internal as eqxi
+import jax
+import pytest
+
+
+jax.config.update("jax_enable_x64", True)
+
+
+@pytest.fixture
+def getkey():
+    return eqxi.GetKey()
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--scipy",
+        action="store_true",
+        dest="scipy",
+        default=False,
+        help="Benchmark against scipy solvers.",
+    )
+    parser.addoption(
+        "--max-dimension",
+        action="store",
+        type=int,
+        default=None,
+        help=(
+            "Maximum dimension for optimization variables. "
+            "Tests with higher dimensions will be skipped."
+        ),
+    )
+    parser.addoption(
+        "--min-dimension",
+        action="store",
+        type=int,
+        default=None,
+        help=(
+            "Minimum dimension for optimization variables. "
+            "Tests with higher dimensions will be skipped."
+        ),
+    )
+
+
+def pytest_configure(config):
+    global _max_dimension
+    _max_dimension = config.getoption("--max-dimension")
+    global _min_dimension
+    _min_dimension = config.getoption("--min-dimension")
+
+
+def get_max_dimension():
+    return _max_dimension
+
+
+def get_min_dimension():
+    return _min_dimension

benchmarks/levenberg-marquardt.py

@@ -137,11 +137,8 @@ def transpose(self, state, options):
     def conj(self, state, options):
         assert False
 
-    def allow_dependent_columns(self, operator):
-        assert False
-
-    def allow_dependent_rows(self, operator):
-        assert False
+    def assume_full_rank(self) -> bool:
+        return True
 
 
 # Default option for Optimistix. (QR linear solver.)