jaddle 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- jaddle-0.1.0/LICENSE.txt +21 -0
- jaddle-0.1.0/PKG-INFO +464 -0
- jaddle-0.1.0/README.md +419 -0
- jaddle-0.1.0/jaddle/__init__.py +1 -0
- jaddle-0.1.0/jaddle/glop_helpers.py +112 -0
- jaddle-0.1.0/jaddle/highs_helpers.py +90 -0
- jaddle-0.1.0/jaddle/jaddle_basic_types.py +272 -0
- jaddle-0.1.0/jaddle/jaddle_convex.py +1340 -0
- jaddle-0.1.0/jaddle/jaddle_linear.py +3224 -0
- jaddle-0.1.0/jaddle/jaddle_optimisers.py +223 -0
- jaddle-0.1.0/jaddle/presolve.py +242 -0
- jaddle-0.1.0/jaddle/sciopt_helpers.py +176 -0
- jaddle-0.1.0/jaddle.egg-info/PKG-INFO +464 -0
- jaddle-0.1.0/jaddle.egg-info/SOURCES.txt +21 -0
- jaddle-0.1.0/jaddle.egg-info/dependency_links.txt +1 -0
- jaddle-0.1.0/jaddle.egg-info/requires.txt +29 -0
- jaddle-0.1.0/jaddle.egg-info/top_level.txt +1 -0
- jaddle-0.1.0/pyproject.toml +61 -0
- jaddle-0.1.0/setup.cfg +4 -0
- jaddle-0.1.0/tests/test_convex.py +103 -0
- jaddle-0.1.0/tests/test_examples.py +44 -0
- jaddle-0.1.0/tests/test_linear.py +161 -0
- jaddle-0.1.0/tests/test_presolve.py +102 -0
jaddle-0.1.0/LICENSE.txt
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Brendan van Rooyen
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
jaddle-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,464 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: jaddle
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: First-order primal-dual solvers for large-scale linear and convex programs, written in JAX
|
|
5
|
+
Author: Brendan van Rooyen
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/BrendanVR/Jaddle
|
|
8
|
+
Project-URL: Source, https://github.com/BrendanVR/Jaddle
|
|
9
|
+
Project-URL: Issues, https://github.com/BrendanVR/Jaddle/issues
|
|
10
|
+
Keywords: optimization,linear programming,convex optimization,pdhg,pdlp,saddle point,jax,gpu
|
|
11
|
+
Classifier: Development Status :: 3 - Alpha
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
18
|
+
Classifier: Topic :: Scientific/Engineering :: Mathematics
|
|
19
|
+
Requires-Python: >=3.11
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
License-File: LICENSE.txt
|
|
22
|
+
Requires-Dist: jax>=0.10
|
|
23
|
+
Requires-Dist: optax>=0.2.8
|
|
24
|
+
Requires-Dist: numpy>=2.0
|
|
25
|
+
Requires-Dist: scipy>=1.13
|
|
26
|
+
Requires-Dist: highspy>=1.14
|
|
27
|
+
Provides-Extra: cuda12
|
|
28
|
+
Requires-Dist: jax[cuda12]>=0.10; extra == "cuda12"
|
|
29
|
+
Provides-Extra: cuda13
|
|
30
|
+
Requires-Dist: jax[cuda13]>=0.10; extra == "cuda13"
|
|
31
|
+
Provides-Extra: scip
|
|
32
|
+
Requires-Dist: pyscipopt>=6.0; extra == "scip"
|
|
33
|
+
Provides-Extra: glop
|
|
34
|
+
Requires-Dist: ortools>=9.15; extra == "glop"
|
|
35
|
+
Provides-Extra: examples
|
|
36
|
+
Requires-Dist: matplotlib>=3.8; extra == "examples"
|
|
37
|
+
Requires-Dist: scikit-learn>=1.5; extra == "examples"
|
|
38
|
+
Provides-Extra: test
|
|
39
|
+
Requires-Dist: pytest>=8; extra == "test"
|
|
40
|
+
Provides-Extra: dev
|
|
41
|
+
Requires-Dist: jaddle[examples,test]; extra == "dev"
|
|
42
|
+
Requires-Dist: build; extra == "dev"
|
|
43
|
+
Requires-Dist: twine; extra == "dev"
|
|
44
|
+
Dynamic: license-file
|
|
45
|
+
|
|
46
|
+
<img width="1536" height="1024" alt="jaddle_logo_full" src="https://github.com/user-attachments/assets/d4e61518-3557-4273-b4d0-0d89510e0f3f" />
|
|
47
|
+
|
|
48
|
+
# Jaddle: The JAX Saddle Solver
|
|
49
|
+
|
|
50
|
+
*First-order primal–dual solvers for large-scale linear and convex programs, written entirely in JAX.*
|
|
51
|
+
|
|
52
|
+
Jaddle solves constrained optimisation problems by casting them as saddle-point
|
|
53
|
+
problems on the Lagrangian and running primal–dual iterations to convergence.
|
|
54
|
+
It provides two solvers:
|
|
55
|
+
|
|
56
|
+
- **`jaddle.jaddle_linear`**: a restarted, adaptively step-sized PDHG-family
|
|
57
|
+
solver for sparse linear programs, in the tradition of PDLP and cuPDLP.
|
|
58
|
+
- **`jaddle.jaddle_convex`**: a general solver for smooth convex programs. The
|
|
59
|
+
user supplies the objective and constraints as ordinary JAX functions, and
|
|
60
|
+
Jaddle derives everything else by automatic differentiation.
|
|
61
|
+
|
|
62
|
+
Both solvers run unchanged on CPU, GPU and TPU. Jaddle is not intended to
|
|
63
|
+
replace a commercial or production-grade solver. It aims to be a compact,
|
|
64
|
+
readable and hackable platform for first-order optimisation that still holds
|
|
65
|
+
up on demanding benchmarks.
|
|
66
|
+
|
|
67
|
+
## Contents
|
|
68
|
+
|
|
69
|
+
- [Why JAX](#why-jax)
|
|
70
|
+
- [Installation](#installation)
|
|
71
|
+
- [Quick start](#quick-start)
|
|
72
|
+
- [The linear programming solver](#the-linear-programming-solver)
|
|
73
|
+
- [The convex programming solver](#the-convex-programming-solver)
|
|
74
|
+
- [Numerical precision](#numerical-precision)
|
|
75
|
+
- [Benchmarks](#benchmarks)
|
|
76
|
+
- [Repository layout](#repository-layout)
|
|
77
|
+
|
|
78
|
+
## Why JAX
|
|
79
|
+
|
|
80
|
+
Jaddle has no hand-written kernels. Every solver in the library is expressed
|
|
81
|
+
in high-level JAX, and JAX's compiler and transformations supply the
|
|
82
|
+
performance and portability. This brings several concrete benefits.
|
|
83
|
+
|
|
84
|
+
**Compiled performance from Python.** Each epoch of iterations is traced once
|
|
85
|
+
and compiled by XLA into a single fused program. Primal and dual updates, step
|
|
86
|
+
size line searches, projections and residual computations run on the
|
|
87
|
+
accelerator without returning control to Python between iterations. The inner
|
|
88
|
+
loop of the LP solver is limited by the speed of sparse matrix–vector products,
|
|
89
|
+
which is the limit any first-order LP method faces, rather than by interpreter
|
|
90
|
+
overhead.
|
|
91
|
+
|
|
92
|
+
**Hardware portability.** The same source runs on a laptop CPU, a consumer GPU
|
|
93
|
+
or a TPU pod. There are no device-specific code paths to maintain, and moving
|
|
94
|
+
between hardware targets requires no code changes. Switching from CPU to GPU is
|
|
95
|
+
a matter of installing the appropriate `jax` wheel.
|
|
96
|
+
|
|
97
|
+
**Automatic differentiation.** The convex solver needs only the objective and
|
|
98
|
+
the constraint functions. Gradients of the Lagrangian with respect to the
|
|
99
|
+
primal and dual variables come from `jax.grad`, so users never derive,
|
|
100
|
+
implement or debug a gradient by hand. Any function JAX can trace can be used
|
|
101
|
+
as an objective or constraint, from a least-squares loss to a neural network.
|
|
102
|
+
|
|
103
|
+
**Composable optimisation primitives.** The convex solver builds on
|
|
104
|
+
[Optax](https://github.com/google-deepmind/optax). It uses Optax's projection
|
|
105
|
+
operators, and in `alternating` mode it accepts any Optax gradient
|
|
106
|
+
transformation as the update rule. A new primal–dual variant is often only a
|
|
107
|
+
few lines of code: for example, giving the primal player plain gradient descent
|
|
108
|
+
and the dual player Nesterov momentum.
|
|
109
|
+
|
|
110
|
+
**Precision as a configuration choice.** A single call switches the whole
|
|
111
|
+
library between float64, float32 and float16 arithmetic. Nothing is
|
|
112
|
+
duplicated per precision, so the trade-off between accuracy and throughput can
|
|
113
|
+
be explored without code changes.
|
|
114
|
+
|
|
115
|
+
**Functional state.** Solver state is an immutable pytree (`SaddleState`, plus
|
|
116
|
+
the step size state). Warm starting, checkpointing and resuming a solve
|
|
117
|
+
therefore amount to passing a value back in. The solver keeps no hidden mutable
|
|
118
|
+
state.
|
|
119
|
+
|
|
120
|
+
**A small, readable codebase.** Because the compiler handles fusion, memory
|
|
121
|
+
layout and device placement, the full library is roughly 5,700 lines of
|
|
122
|
+
Python. Algorithms remain close to their mathematical description, which makes
|
|
123
|
+
Jaddle well suited to research and experimentation.
|
|
124
|
+
|
|
125
|
+
## Installation
|
|
126
|
+
|
|
127
|
+
Jaddle requires Python 3.11 or later. Clone the repository and install it:
|
|
128
|
+
|
|
129
|
+
```bash
|
|
130
|
+
git clone https://github.com/BrendanVR/Jaddle.git
|
|
131
|
+
cd Jaddle
|
|
132
|
+
pip install -e .
|
|
133
|
+
```
|
|
134
|
+
|
|
135
|
+
This installs the core dependencies: `jax`, `optax`, `numpy`, `scipy` and
|
|
136
|
+
`highspy`, and runs on CPU. Optional features are available as extras, which
|
|
137
|
+
can be combined, for example `pip install -e ".[cuda12,examples]"`:
|
|
138
|
+
|
|
139
|
+
| Extra | Installs | Used for |
|
|
140
|
+
|---|---|---|
|
|
141
|
+
| `cuda12`, `cuda13` | CUDA-enabled JAX | GPU execution. Choose the extra that matches your CUDA driver; see the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html) for TPU and other platforms. |
|
|
142
|
+
| `scip` | `pyscipopt` | Reading and presolving MPS files with SCIP (`jaddle.sciopt_helpers`) |
|
|
143
|
+
| `glop` | `ortools` | Presolving with OR-Tools glop (`jaddle.glop_helpers`). Also requires building the helper with `tools/glop_presolve/build.sh`. |
|
|
144
|
+
| `examples` | `matplotlib`, `scikit-learn` | The convex examples |
|
|
145
|
+
| `test` | `pytest` | Running the test suite |
|
|
146
|
+
| `dev` | `examples`, `test`, `build`, `twine` | Development and packaging |
|
|
147
|
+
|
|
148
|
+
Jaddle is developed on Linux, including WSL 2. Continuous integration tests
|
|
149
|
+
it on Python 3.11, 3.12 and 3.13.
|
|
150
|
+
|
|
151
|
+
### Running the tests
|
|
152
|
+
|
|
153
|
+
```bash
|
|
154
|
+
pip install -e ".[examples,test]"
|
|
155
|
+
pytest
|
|
156
|
+
```
|
|
157
|
+
|
|
158
|
+
The suite runs on CPU in double precision and takes well under a minute. It
|
|
159
|
+
checks both solvers against exact or reference solutions (SciPy's HiGHS
|
|
160
|
+
interface for LPs, closed-form solutions for the convex problems), the
|
|
161
|
+
defined-variable presolve, and the self-contained examples. GitHub Actions runs
|
|
162
|
+
it on every push and pull request.
|
|
163
|
+
|
|
164
|
+
## Quick start
|
|
165
|
+
|
|
166
|
+
### A linear program
|
|
167
|
+
|
|
168
|
+
```python
|
|
169
|
+
import numpy as np
|
|
170
|
+
import scipy.sparse as sp
|
|
171
|
+
import jaddle.jaddle_linear as jl
|
|
172
|
+
|
|
173
|
+
# minimise 3 x1 + 2 x2 subject to x1 + x2 >= 4, x >= 0
|
|
174
|
+
lp = jl.LP(
|
|
175
|
+
c=np.array([3.0, 2.0]),
|
|
176
|
+
A_eq=sp.csc_matrix((0, 2)), b_eq=np.zeros(0),
|
|
177
|
+
A_ineq=sp.csc_matrix([[-1.0, -1.0]]), b_ineq=np.array([-4.0]),
|
|
178
|
+
lower_bounds=np.zeros(2), upper_bounds=np.full(2, np.inf),
|
|
179
|
+
)
|
|
180
|
+
|
|
181
|
+
result = jl.solve(lp, verbose=True)
|
|
182
|
+
x = result["solution"].primal
|
|
183
|
+
print(result["stop_reason"], lp.objective(x))
|
|
184
|
+
```
|
|
185
|
+
|
|
186
|
+
### A convex program
|
|
187
|
+
|
|
188
|
+
```python
|
|
189
|
+
import jax.numpy as jnp
|
|
190
|
+
import jaddle.jaddle_convex as jc
|
|
191
|
+
|
|
192
|
+
y = jnp.linspace(-1.0, 1.0, 100) ** 3 # data to fit
|
|
193
|
+
|
|
194
|
+
cp = jc.JaddleCP(
|
|
195
|
+
num_variables=100,
|
|
196
|
+
objective=lambda p: jnp.sum((p - y) ** 2),
|
|
197
|
+
constraints_eq=lambda p: jnp.zeros(0),
|
|
198
|
+
constraints_ineq=lambda p: p[:-1] - p[1:], # p must be non-decreasing
|
|
199
|
+
lower_bounds=-jnp.ones(100),
|
|
200
|
+
upper_bounds=jnp.ones(100),
|
|
201
|
+
)
|
|
202
|
+
|
|
203
|
+
result = jc.solve(cp)
|
|
204
|
+
p = result["solution"].primal
|
|
205
|
+
```
|
|
206
|
+
|
|
207
|
+
The [`examples/`](https://github.com/BrendanVR/Jaddle/blob/main/examples/README.md) directory contains complete, annotated
|
|
208
|
+
versions of both, together with examples that load and solve MIPLIB instances.
|
|
209
|
+
|
|
210
|
+
## The linear programming solver
|
|
211
|
+
|
|
212
|
+
### Problem form
|
|
213
|
+
|
|
214
|
+
`jaddle_linear.solve` accepts linear programs in the standard form
|
|
215
|
+
|
|
216
|
+
```
|
|
217
|
+
minimise cᵀx
|
|
218
|
+
subject to A_eq x = b_eq
|
|
219
|
+
A_ineq x <= b_ineq
|
|
220
|
+
l <= x <= u
|
|
221
|
+
```
|
|
222
|
+
|
|
223
|
+
where bounds may be infinite. Problems can be built directly from NumPy and
|
|
224
|
+
SciPy sparse arrays with `jl.LP`. They can also be converted from MPS files
|
|
225
|
+
through HiGHS (`highs_helpers.highs_to_standard_form_sparse`), SCIP
|
|
226
|
+
(`sciopt_helpers.scip_to_standard_form_sparse`) or OR-Tools glop
|
|
227
|
+
(`glop_helpers.glop_presolve`). Internally, constraint matrices are stored as
|
|
228
|
+
JAX `BCOO` sparse arrays.
|
|
229
|
+
|
|
230
|
+
The solver seeks a saddle point of the Lagrangian
|
|
231
|
+
|
|
232
|
+
```
|
|
233
|
+
L(x, y) = cᵀx + y_eqᵀ(A_eq x − b_eq) + y_ineqᵀ(A_ineq x − b_ineq), y_ineq >= 0,
|
|
234
|
+
```
|
|
235
|
+
|
|
236
|
+
minimising over `x` within its bounds and maximising over the duals `y`.
|
|
237
|
+
|
|
238
|
+
### Algorithm
|
|
239
|
+
|
|
240
|
+
Each iteration takes a projected primal descent step and a projected dual
|
|
241
|
+
ascent step. The key components follow the design of PDLP and cuPDLP, with
|
|
242
|
+
some refinements of our own.
|
|
243
|
+
|
|
244
|
+
**Update schemes** (`update_mode`)
|
|
245
|
+
|
|
246
|
+
| Mode | Description |
|
|
247
|
+
|---|---|
|
|
248
|
+
| `"alternating"` (default) | Gauss–Seidel primal–dual iteration: a primal step, then a dual step evaluated at the new primal iterate. |
|
|
249
|
+
| `"pdhg"` | Chambolle–Pock PDHG: the dual step is taken at the extrapolated point `2x⁽ᵏ⁺¹⁾ − x⁽ᵏ⁾`. |
|
|
250
|
+
| `"halpern"` | Restarted Halpern-anchored PDHG. Each step blends the PDHG update back towards an anchor point, which accelerates convergence of the last iterate. Best used with restarts. |
|
|
251
|
+
|
|
252
|
+
**Adaptive step size.** All modes use a cuPDLP-style adaptive step. At every
|
|
253
|
+
iteration the step `η` is checked against a local estimate of the interaction
|
|
254
|
+
term `(Δy)ᵀ A (Δx)`. A step that overshoots is rejected and shrunk. An accepted
|
|
255
|
+
step is advanced conservatively. The step size learned during a solve is
|
|
256
|
+
carried across restarts.
|
|
257
|
+
|
|
258
|
+
**Primal weight.** The primal and dual step sizes are `η / k` and `η · k`. The
|
|
259
|
+
primal weight `k` is initialised from the ratio `‖c‖ / ‖b‖`. At each restart it
|
|
260
|
+
is rebalanced from the relative movement of the primal and dual iterates, using
|
|
261
|
+
a log-space blend controlled by `k_theta` and clamped to
|
|
262
|
+
`[1/k_scale, k_scale]`.
|
|
263
|
+
|
|
264
|
+
**Adaptive restarts.** A restart resets the averaging (and, for `halpern`, the
|
|
265
|
+
anchor) and keeps the current iterate as a warm start. This stops the
|
|
266
|
+
iteration from settling into slow rotational orbits. Restarts fire on any of
|
|
267
|
+
three conditions, all judged by a normalised KKT merit:
|
|
268
|
+
|
|
269
|
+
- *sufficient progress*: the merit falls below `restart_decay` times its value
|
|
270
|
+
at the last restart;
|
|
271
|
+
- *stalling*: the merit falls below `necessary_decay` times its value at the
|
|
272
|
+
start of the cycle and then starts to rise;
|
|
273
|
+
- *no progress*: the cycle reaches its length cap.
|
|
274
|
+
|
|
275
|
+
The merit can be checked within an epoch (`restart_check_every`), so a restart
|
|
276
|
+
does not have to wait for the epoch boundary.
|
|
277
|
+
|
|
278
|
+
**Diagonal preconditioning.** Before solving, the problem is equilibrated by
|
|
279
|
+
Ruiz (L∞) scaling followed by Pock–Chambolle (L1) scaling. By default the
|
|
280
|
+
scaling is applied to the augmented matrix `[[A, b], [cᵀ, 0]]`, so cost and
|
|
281
|
+
right-hand-side magnitudes also inform the scaling. The objective and
|
|
282
|
+
right-hand side are then normalised. All reported quantities are mapped back
|
|
283
|
+
to the original problem's units.
|
|
284
|
+
|
|
285
|
+
**Termination.** The solver stops when the standard LP optimality certificate
|
|
286
|
+
holds to a relative tolerance, following the PDLP and HiGHS conventions:
|
|
287
|
+
|
|
288
|
+
| Test | Normalised by | Tolerance argument |
|
|
289
|
+
|---|---|---|
|
|
290
|
+
| Primal feasibility | `1 + ‖b‖` | `primal_feasibility_tolerance` |
|
|
291
|
+
| Dual feasibility | `1 + ‖c‖` | `dual_feasibility_tolerance` |
|
|
292
|
+
| Duality gap | `1 + \|primal obj\| + \|dual obj\|` | `dual_gap_tolerance` |
|
|
293
|
+
|
|
294
|
+
All three tolerances default to `1e-3`. The norms are L2 by default
|
|
295
|
+
(`termination_norm`). Reduced costs are split between the dual residual and
|
|
296
|
+
the dual objective as in PDLP, so far-away finite bounds do not swamp the gap.
|
|
297
|
+
|
|
298
|
+
### Principal options
|
|
299
|
+
|
|
300
|
+
| Argument | Default | Purpose |
|
|
301
|
+
|---|---|---|
|
|
302
|
+
| `update_mode` | `"alternating"` | Iteration scheme (see above). |
|
|
303
|
+
| `iterations_per_epoch` | `256` | Iterations per compiled epoch. Metrics, logging and termination checks run between epochs. |
|
|
304
|
+
| `max_epochs`, `max_seconds` | `None` | Iteration and wall-clock budgets. `max_seconds` includes setup and compilation. |
|
|
305
|
+
| `restarts`, `epochs_per_restart` | `True`, `10` | Adaptive restarts and the length cap of a restart cycle. |
|
|
306
|
+
| `restart_decay`, `necessary_decay` | `0.2`, `0.8` | Thresholds for the sufficient-progress and stalling restarts. |
|
|
307
|
+
| `k_scale`, `k_theta`, `k_init` | `1e8`, `0.5`, `None` | Clamp band, smoothing and initial value for the primal weight. |
|
|
308
|
+
| `ruiz_iterations`, `pc_iterations` | `10`, `1` | Number of preconditioning sweeps. |
|
|
309
|
+
| `initial_solution`, `initial_opt_state` | `None` | Warm start from a previous solve's `solution` and `opt_state`. |
|
|
310
|
+
| `vertex_bias` | `0.0` | Small cost perturbation that steers the solver to a vertex of the optimal face rather than its interior. Useful before crossover. |
|
|
311
|
+
| `verbose`, `log_every` | `False`, `1` | Per-epoch progress logging. |
|
|
312
|
+
|
|
313
|
+
### Results
|
|
314
|
+
|
|
315
|
+
`solve` returns a dictionary with the following keys:
|
|
316
|
+
|
|
317
|
+
| Key | Contents |
|
|
318
|
+
|---|---|
|
|
319
|
+
| `solution` | A `SaddleState(primal, dual_ineq, dual_eq)` in the original problem's units. |
|
|
320
|
+
| `converged` | Whether a termination criterion was met. |
|
|
321
|
+
| `stop_reason` | `"certificate"`, `"primal_stall"` (only when the opt-in `primal_stop` heuristic is enabled), `"max_epochs"`, `"time_limit"` or `"interrupted"`. |
|
|
322
|
+
| `opt_state` | Final step size and primal weight, for warm starting. |
|
|
323
|
+
| `solve_seconds` | Time spent in the iteration loop, including the first-epoch compile. |
|
|
324
|
+
| `corrected_seconds` | Iteration-loop time with the one-off compile amortised out. |
|
|
325
|
+
| `epochs` | Number of epochs run. |
|
|
326
|
+
|
|
327
|
+
### Supporting tools
|
|
328
|
+
|
|
329
|
+
- **Presolve.** `jaddle.presolve.eliminate_defined_variables` substitutes out
|
|
330
|
+
variables defined by dense equality rows (`z = Σ aⱼxⱼ`). Rows of this kind
|
|
331
|
+
hide objective structure from the scaling and can stall first-order methods.
|
|
332
|
+
External presolve is available through HiGHS, SCIP and glop.
|
|
333
|
+
- **Certification.** `jl.evaluate_lp_certificate` reports the relative primal,
|
|
334
|
+
dual and gap residuals of any primal–dual pair independently of a solve.
|
|
335
|
+
- **Polishing.** `jl.solve_with_polishing`, `jl.primal_polish` and
|
|
336
|
+
`jl.dual_polish` refine a near-optimal point towards a tighter certificate.
|
|
337
|
+
- **Diagnostics.** `jl.lp_summary_statistics` prints the problem dimensions and
|
|
338
|
+
coefficient ranges, which is the first thing to check when an instance
|
|
339
|
+
misbehaves.
|
|
340
|
+
|
|
341
|
+
## The convex programming solver
|
|
342
|
+
|
|
343
|
+
### Problem form
|
|
344
|
+
|
|
345
|
+
`jaddle_convex.solve` accepts a `JaddleCP`:
|
|
346
|
+
|
|
347
|
+
```python
|
|
348
|
+
jc.JaddleCP(
|
|
349
|
+
num_variables, # n
|
|
350
|
+
objective, # f(x) -> scalar, convex and differentiable
|
|
351
|
+
constraints_eq, # h(x) -> vector, constrained to h(x) = 0
|
|
352
|
+
constraints_ineq, # g(x) -> vector, constrained to g(x) <= 0
|
|
353
|
+
lower_bounds, # l, may contain -inf
|
|
354
|
+
upper_bounds, # u, may contain +inf
|
|
355
|
+
)
|
|
356
|
+
```
|
|
357
|
+
|
|
358
|
+
All functions must be traceable by JAX. The solver forms the Lagrangian
|
|
359
|
+
`f(x) + λᵀg(x) + μᵀh(x)` with `λ >= 0`, and obtains every gradient it needs by
|
|
360
|
+
automatic differentiation. No derivatives are required from the user.
|
|
361
|
+
|
|
362
|
+
### Algorithm
|
|
363
|
+
|
|
364
|
+
**Update schemes** (`update_mode`)
|
|
365
|
+
|
|
366
|
+
| Mode | Gradient evaluations per iteration | Description |
|
|
367
|
+
|---|---|---|
|
|
368
|
+
| `"extragradient"` (default) | 2 | Korpelevich extragradient: a look-ahead step followed by a corrector step. Contractive for monotone problems. |
|
|
369
|
+
| `"forward_reflected"` | 1 | Malitsky–Tam forward-reflected-backward: one gradient per iteration plus a reflection term built from the previous gradient. Requires the adaptive step. |
|
|
370
|
+
| `"alternating"` | 1 | Primal step, then a dual step at the new primal point. The update rule is a user-supplied Optax optimiser. |
|
|
371
|
+
|
|
372
|
+
**Adaptive step size.** With `adaptive_eta="auto"` (the default), the
|
|
373
|
+
extragradient and forward-reflected schemes use a Malitsky–Tam line search
|
|
374
|
+
based on a local Lipschitz estimate. The two gradient evaluations that
|
|
375
|
+
extragradient already performs give an estimate of the local Lipschitz
|
|
376
|
+
constant at no extra cost. The step is accepted while it remains within the
|
|
377
|
+
admissible bound, and is rejected and shrunk otherwise. The learned step is
|
|
378
|
+
carried across restarts.
|
|
379
|
+
|
|
380
|
+
**Primal weight.** As in the LP solver, a primal weight `k` balances the primal
|
|
381
|
+
and dual step sizes. It is initialised from the gradient of the objective and
|
|
382
|
+
the constraint values at the origin, rebalanced at each restart, and clamped
|
|
383
|
+
to `[1/k_scale, k_scale]` (default `k_scale = 10`).
|
|
384
|
+
|
|
385
|
+
**Custom optimisers.** In `alternating` mode, `optimiser` accepts any Optax
|
|
386
|
+
`GradientTransformation`. It is applied separately to the primal and dual
|
|
387
|
+
players through `jaddle_optimisers.create_saddle_optimiser`. Ready-made
|
|
388
|
+
choices include `jo.gd`, `jo.optimistic_gd`, `jo.optimistic_adadelta` and
|
|
389
|
+
`jo.gd_dual_momentum`.
|
|
390
|
+
|
|
391
|
+
**Restarts and averaging.** Adaptive restarts (`restarts=True`) and iterate
|
|
392
|
+
averaging (`average=True`, with an optional `weight_function`) are available
|
|
393
|
+
but off by default.
|
|
394
|
+
|
|
395
|
+
**Termination.** The solver stops when all three of the following fall below
|
|
396
|
+
their tolerances:
|
|
397
|
+
|
|
398
|
+
| Test | Tolerance argument | Default |
|
|
399
|
+
|---|---|---|
|
|
400
|
+
| Stationarity (projected primal gradient norm of the Lagrangian) | `primal_grad_norm_tolerance` | `1e-2` |
|
|
401
|
+
| Primal feasibility | `primal_feasibility_tolerance` | `1e-3` |
|
|
402
|
+
| Complementary slackness | `complementarity_slack_tolerance` | `1e-3` |
|
|
403
|
+
|
|
404
|
+
### Principal options
|
|
405
|
+
|
|
406
|
+
| Argument | Default | Purpose |
|
|
407
|
+
|---|---|---|
|
|
408
|
+
| `update_mode` | `"extragradient"` | Iteration scheme (see above). |
|
|
409
|
+
| `adaptive_eta` | `"auto"` | Adaptive line-searched step. Pass a float to set the seed, or `None` to use the optimiser's fixed learning rate. |
|
|
410
|
+
| `optimiser` | `None` | Optax optimiser for `alternating` mode. When `None`, gradient descent with learning rate 0.5 is used. |
|
|
411
|
+
| `iterations_per_epoch` | `1000` | Iterations per compiled epoch. |
|
|
412
|
+
| `max_epochs`, `max_seconds` | `None` | Iteration and wall-clock budgets. |
|
|
413
|
+
| `k_scale`, `k_theta`, `k_init` | `10.0`, `0.5`, `None` | Primal weight controls. `k_scale=None` disables the primal weight. |
|
|
414
|
+
| `restarts`, `epochs_per_restart` | `False`, `10` | Adaptive restarts. |
|
|
415
|
+
| `initial_solution`, `initial_opt_state` | `None` | Warm start. |
|
|
416
|
+
|
|
417
|
+
`solve` returns a dictionary containing `solution` (a `SaddleState`),
|
|
418
|
+
`converged`, `stop_reason`, `opt_state` and `solve_seconds`.
|
|
419
|
+
|
|
420
|
+
## Numerical precision
|
|
421
|
+
|
|
422
|
+
Call `jaddle.jaddle_optimisers.configure_jax(profile)` before creating any
|
|
423
|
+
arrays:
|
|
424
|
+
|
|
425
|
+
| Profile | Arithmetic | Recommended for |
|
|
426
|
+
|---|---|---|
|
|
427
|
+
| `"float64"` (default) | Double precision, full-precision matrix products | Real-world LPs and any problem needing tight tolerances. Wide-dynamic-range instances can stall on rounding error in single precision. |
|
|
428
|
+
| `"float32"` | Single precision, TF32 matrix products | Throughput on well-conditioned problems. |
|
|
429
|
+
| `"float16"` | Half precision, accumulating in float32 | Memory- and bandwidth-limited experiments. |
|
|
430
|
+
|
|
431
|
+
The profile can also be set through the `JADDLE_JAX_PROFILE` environment
|
|
432
|
+
variable. `configure_jax` also enables a persistent XLA compilation cache in
|
|
433
|
+
`~/.cache/jaddle_jax`, so repeated solves of problems with the same shape skip
|
|
434
|
+
recompilation.
|
|
435
|
+
|
|
436
|
+
## Benchmarks
|
|
437
|
+
|
|
438
|
+
Jaddle's LP solver is evaluated on the LP relaxations of the 383 MIPLIB 2017
|
|
439
|
+
instances used in the PDLP papers (100k–10M nonzeros). With default solver
|
|
440
|
+
settings, HiGHS presolve, a relative tolerance of `1e-4` and a 600-second
|
|
441
|
+
limit per instance, Jaddle certifies optimality on **374 of the 383**
|
|
442
|
+
instances.
|
|
443
|
+
|
|
444
|
+
The benchmark harnesses, instance diagnostics and instructions for reproducing
|
|
445
|
+
these results are described in [`benchmarks/README.md`](https://github.com/BrendanVR/Jaddle/blob/main/benchmarks/README.md).
|
|
446
|
+
|
|
447
|
+
## Repository layout
|
|
448
|
+
|
|
449
|
+
| Path | Contents |
|
|
450
|
+
|---|---|
|
|
451
|
+
| [`jaddle/jaddle_linear.py`](https://github.com/BrendanVR/Jaddle/blob/main/jaddle/jaddle_linear.py) | LP solver, scaling, certification and polishing |
|
|
452
|
+
| [`jaddle/jaddle_convex.py`](https://github.com/BrendanVR/Jaddle/blob/main/jaddle/jaddle_convex.py) | Convex solver |
|
|
453
|
+
| [`jaddle/jaddle_basic_types.py`](https://github.com/BrendanVR/Jaddle/blob/main/jaddle/jaddle_basic_types.py) | `SaddleState`, `JaddleCP`, `LP` and `JaddleLP` |
|
|
454
|
+
| [`jaddle/jaddle_optimisers.py`](https://github.com/BrendanVR/Jaddle/blob/main/jaddle/jaddle_optimisers.py) | Precision profiles and Optax saddle-point optimisers |
|
|
455
|
+
| [`jaddle/presolve.py`](https://github.com/BrendanVR/Jaddle/blob/main/jaddle/presolve.py) | Defined-variable elimination and postsolve |
|
|
456
|
+
| [`jaddle/highs_helpers.py`](https://github.com/BrendanVR/Jaddle/blob/main/jaddle/highs_helpers.py), [`sciopt_helpers.py`](https://github.com/BrendanVR/Jaddle/blob/main/jaddle/sciopt_helpers.py), [`glop_helpers.py`](https://github.com/BrendanVR/Jaddle/blob/main/jaddle/glop_helpers.py) | Converters and presolve interfaces for HiGHS, SCIP and OR-Tools glop |
|
|
457
|
+
| [`examples/`](https://github.com/BrendanVR/Jaddle/blob/main/examples/README.md) | Worked LP and convex examples |
|
|
458
|
+
| [`benchmarks/`](https://github.com/BrendanVR/Jaddle/blob/main/benchmarks/README.md) | MIPLIB and convex benchmark harnesses |
|
|
459
|
+
| [`tools/glop_presolve/`](https://github.com/BrendanVR/Jaddle/tree/main/tools/glop_presolve) | C++ helper exposing glop's presolver |
|
|
460
|
+
| [`tests/`](https://github.com/BrendanVR/Jaddle/tree/main/tests) | Test suite (`pytest`) |
|
|
461
|
+
|
|
462
|
+
## License
|
|
463
|
+
|
|
464
|
+
Jaddle is released under the MIT License. See [LICENSE.txt](https://github.com/BrendanVR/Jaddle/blob/main/LICENSE.txt).
|