elastiqp 0.2.0__tar.gz → 0.2.1__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.
- {elastiqp-0.2.0 → elastiqp-0.2.1}/.github/workflows/wheels.yml +1 -1
- {elastiqp-0.2.0 → elastiqp-0.2.1}/PKG-INFO +1 -1
- {elastiqp-0.2.0 → elastiqp-0.2.1}/bindings/bindings.cc +1 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/bindings/jax_ffi.cc +25 -68
- {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/das.hpp +36 -8
- {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/elastiqp_c.h +4 -1
- {elastiqp-0.2.0 → elastiqp-0.2.1}/package.xml +1 -1
- {elastiqp-0.2.0 → elastiqp-0.2.1}/python/elastiqp/jax.py +47 -135
- {elastiqp-0.2.0 → elastiqp-0.2.1}/src/elastiqp_c.cc +6 -2
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_c_api.c +1 -1
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_c_parity.cc +1 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_das.cc +57 -3
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_jax_ffi.py +99 -29
- {elastiqp-0.2.0 → elastiqp-0.2.1}/.clang-format +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/.gitignore +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/CMakeLists.txt +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/LICENSE +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/NOTICE +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/README.md +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/cmake/elastiqpConfig.cmake.in +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/cmake/elastiqp_install.cmake +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/examples/basic_usage.py +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/examples/constraint_conflict_demo.py +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/examples/jax_usage.py +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/examples/torch_usage.py +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/common.hpp +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/elastiqp.hpp +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/ipm.hpp +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/kkt_vjp.hpp +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/include/elastiqp/pdal.hpp +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/pyproject.toml +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/python/elastiqp/__init__.py +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/python/elastiqp/torch.py +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/data/warmstart_dual_overshoot.npz +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/data/warmstart_stale_duals.npz +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/support/drift_traj.hpp +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/support/problem_gen.hpp +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/support/test_util.hpp +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_bcl_creep.cc +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_bindings.py +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_das_alloc.cc +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_gap_creep.cc +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_ipm.cc +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_lp.cc +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_pdal.cc +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_ruiz.cc +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_solvers.cc +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/tests/test_torch.py +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/third_party/xla_ffi/README.md +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/third_party/xla_ffi/xla/ffi/api/api.h +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/third_party/xla_ffi/xla/ffi/api/c_api.h +0 -0
- {elastiqp-0.2.0 → elastiqp-0.2.1}/third_party/xla_ffi/xla/ffi/api/ffi.h +0 -0
|
@@ -285,6 +285,7 @@ NB_MODULE(_core, m) {
|
|
|
285
285
|
.def_rw("eps_rel", &S::eps_rel)
|
|
286
286
|
.def_rw("sing_tol", &S::sing_tol)
|
|
287
287
|
.def_rw("zero_tol", &S::zero_tol)
|
|
288
|
+
.def_rw("cond_iters", &S::cond_iters)
|
|
288
289
|
.def_rw("eps_prox", &S::eps_prox)
|
|
289
290
|
.def_rw("eta_prox", &S::eta_prox)
|
|
290
291
|
.def_rw("prox_relaxation", &S::prox_relaxation)
|
|
@@ -1,16 +1,16 @@
|
|
|
1
1
|
// XLA FFI handlers for ElastiQP. Every call builds a fresh solver (XLA
|
|
2
2
|
// assumes purity), so warm starting is an explicit argument.
|
|
3
3
|
//
|
|
4
|
-
// ElastiqpSolve:
|
|
5
|
-
//
|
|
6
|
-
// ElastiqpSolveWarm: (Q, q, A, b, G, h, penalty, x0, y0, z0)
|
|
7
|
-
// -> (x, t, y, z_t, z, info) (das and pdal only)
|
|
4
|
+
// ElastiqpSolve: (Q, q, A, b, G, h, penalty, x0, y0, z0)
|
|
5
|
+
// -> (x, t, y, z_t, z, xr, tr, yr, z_t_r, z_r, info)
|
|
8
6
|
//
|
|
9
7
|
// method: 0 = das, 1 = pdal, 2 = ipm. max_iter is the backend's outer budget.
|
|
8
|
+
// warm != 0 seeds the solve from (x0, y0, z0) (das and pdal only); otherwise
|
|
9
|
+
// they are ignored.
|
|
10
10
|
// The *r block is the kappa-relaxed point the Python wrappers differentiate
|
|
11
11
|
// at; it copies the tight solution when target_kappa <= 0. info = [converged,
|
|
12
|
-
// iters, relax_converged]
|
|
13
|
-
//
|
|
12
|
+
// iters, relax_converged]. A has 0 rows and b 0 entries when there are no
|
|
13
|
+
// equalities.
|
|
14
14
|
|
|
15
15
|
#include <algorithm>
|
|
16
16
|
#include <cstdint>
|
|
@@ -134,9 +134,11 @@ void write_solution(const elastiqp::Solution& sol, int64_t n, int64_t m,
|
|
|
134
134
|
|
|
135
135
|
ffi::Error ElastiqpSolveImpl(
|
|
136
136
|
double eps_abs, int64_t max_iter, int64_t ruiz, double target_kappa,
|
|
137
|
-
int64_t
|
|
138
|
-
ffi::Buffer<ffi::F64>
|
|
139
|
-
ffi::Buffer<ffi::F64>
|
|
137
|
+
int64_t warm, int64_t method, ffi::Buffer<ffi::F64> Q,
|
|
138
|
+
ffi::Buffer<ffi::F64> q, ffi::Buffer<ffi::F64> A, ffi::Buffer<ffi::F64> b,
|
|
139
|
+
ffi::Buffer<ffi::F64> G, ffi::Buffer<ffi::F64> h,
|
|
140
|
+
ffi::Buffer<ffi::F64> penalty, ffi::Buffer<ffi::F64> x0,
|
|
141
|
+
ffi::Buffer<ffi::F64> y0, ffi::Buffer<ffi::F64> z0,
|
|
140
142
|
ffi::ResultBuffer<ffi::F64> x, ffi::ResultBuffer<ffi::F64> t,
|
|
141
143
|
ffi::ResultBuffer<ffi::F64> y, ffi::ResultBuffer<ffi::F64> z_t,
|
|
142
144
|
ffi::ResultBuffer<ffi::F64> z, ffi::ResultBuffer<ffi::F64> xr,
|
|
@@ -146,37 +148,9 @@ ffi::Error ElastiqpSolveImpl(
|
|
|
146
148
|
const int64_t n = q.dimensions()[0];
|
|
147
149
|
const int64_t m = b.dimensions()[0];
|
|
148
150
|
const int64_t p = h.dimensions()[0];
|
|
149
|
-
|
|
150
|
-
|
|
151
|
-
|
|
152
|
-
dispatch(method, eps_abs, max_iter, ruiz, target_kappa, n, m, p,
|
|
153
|
-
Q.typed_data(), q.typed_data(), A.typed_data(), b.typed_data(),
|
|
154
|
-
G.typed_data(), h.typed_data(), penalty.typed_data(), nullptr,
|
|
155
|
-
nullptr, nullptr, sol, rsol);
|
|
156
|
-
if (err.failure()) return err;
|
|
157
|
-
|
|
158
|
-
write_solution(sol, n, m, p, x, t, y, z_t, z);
|
|
159
|
-
write_solution(rsol, n, m, p, xr, tr, yr, z_t_r, z_r);
|
|
160
|
-
info->typed_data()[0] = static_cast<double>(sol.converged);
|
|
161
|
-
info->typed_data()[1] = static_cast<double>(sol.iters);
|
|
162
|
-
info->typed_data()[2] = static_cast<double>(rsol.converged);
|
|
163
|
-
return ffi::Error::Success();
|
|
164
|
-
}
|
|
165
|
-
|
|
166
|
-
ffi::Error ElastiqpSolveWarmImpl(
|
|
167
|
-
double eps_abs, int64_t max_iter, int64_t ruiz, int64_t method,
|
|
168
|
-
ffi::Buffer<ffi::F64> Q, ffi::Buffer<ffi::F64> q, ffi::Buffer<ffi::F64> A,
|
|
169
|
-
ffi::Buffer<ffi::F64> b, ffi::Buffer<ffi::F64> G, ffi::Buffer<ffi::F64> h,
|
|
170
|
-
ffi::Buffer<ffi::F64> penalty, ffi::Buffer<ffi::F64> x0,
|
|
171
|
-
ffi::Buffer<ffi::F64> y0, ffi::Buffer<ffi::F64> z0,
|
|
172
|
-
ffi::ResultBuffer<ffi::F64> x, ffi::ResultBuffer<ffi::F64> t,
|
|
173
|
-
ffi::ResultBuffer<ffi::F64> y, ffi::ResultBuffer<ffi::F64> z_t,
|
|
174
|
-
ffi::ResultBuffer<ffi::F64> z, ffi::ResultBuffer<ffi::F64> info) {
|
|
175
|
-
const int64_t n = q.dimensions()[0];
|
|
176
|
-
const int64_t m = b.dimensions()[0];
|
|
177
|
-
const int64_t p = h.dimensions()[0];
|
|
178
|
-
if (x0.dimensions()[0] != n || y0.dimensions()[0] != m ||
|
|
179
|
-
z0.dimensions()[0] != p) {
|
|
151
|
+
const bool w = warm != 0;
|
|
152
|
+
if (w && (x0.dimensions()[0] != n || y0.dimensions()[0] != m ||
|
|
153
|
+
z0.dimensions()[0] != p)) {
|
|
180
154
|
return ffi::Error(ffi::ErrorCode::kInvalidArgument,
|
|
181
155
|
"warm start (x0, y0, z0) must have shapes (n,), (m,), "
|
|
182
156
|
"(p,) matching (q, b, h)");
|
|
@@ -184,15 +158,18 @@ ffi::Error ElastiqpSolveWarmImpl(
|
|
|
184
158
|
|
|
185
159
|
elastiqp::Solution sol, rsol;
|
|
186
160
|
ffi::Error err =
|
|
187
|
-
dispatch(method, eps_abs, max_iter, ruiz,
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
161
|
+
dispatch(method, eps_abs, max_iter, ruiz, target_kappa, n, m, p,
|
|
162
|
+
Q.typed_data(), q.typed_data(), A.typed_data(), b.typed_data(),
|
|
163
|
+
G.typed_data(), h.typed_data(), penalty.typed_data(),
|
|
164
|
+
w ? x0.typed_data() : nullptr, w ? y0.typed_data() : nullptr,
|
|
165
|
+
w ? z0.typed_data() : nullptr, sol, rsol);
|
|
191
166
|
if (err.failure()) return err;
|
|
192
167
|
|
|
193
168
|
write_solution(sol, n, m, p, x, t, y, z_t, z);
|
|
169
|
+
write_solution(rsol, n, m, p, xr, tr, yr, z_t_r, z_r);
|
|
194
170
|
info->typed_data()[0] = static_cast<double>(sol.converged);
|
|
195
171
|
info->typed_data()[1] = static_cast<double>(sol.iters);
|
|
172
|
+
info->typed_data()[2] = static_cast<double>(rsol.converged);
|
|
196
173
|
return ffi::Error::Success();
|
|
197
174
|
}
|
|
198
175
|
|
|
@@ -204,6 +181,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(ElastiqpSolve, ElastiqpSolveImpl,
|
|
|
204
181
|
.Attr<int64_t>("max_iter")
|
|
205
182
|
.Attr<int64_t>("ruiz")
|
|
206
183
|
.Attr<double>("target_kappa")
|
|
184
|
+
.Attr<int64_t>("warm")
|
|
207
185
|
.Attr<int64_t>("method")
|
|
208
186
|
.Arg<ffi::Buffer<ffi::F64>>() // Q
|
|
209
187
|
.Arg<ffi::Buffer<ffi::F64>>() // q
|
|
@@ -212,6 +190,9 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(ElastiqpSolve, ElastiqpSolveImpl,
|
|
|
212
190
|
.Arg<ffi::Buffer<ffi::F64>>() // G
|
|
213
191
|
.Arg<ffi::Buffer<ffi::F64>>() // h
|
|
214
192
|
.Arg<ffi::Buffer<ffi::F64>>() // penalty
|
|
193
|
+
.Arg<ffi::Buffer<ffi::F64>>() // x0
|
|
194
|
+
.Arg<ffi::Buffer<ffi::F64>>() // y0
|
|
195
|
+
.Arg<ffi::Buffer<ffi::F64>>() // z0
|
|
215
196
|
.Ret<ffi::Buffer<ffi::F64>>() // x
|
|
216
197
|
.Ret<ffi::Buffer<ffi::F64>>() // t
|
|
217
198
|
.Ret<ffi::Buffer<ffi::F64>>() // y
|
|
@@ -224,27 +205,3 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(ElastiqpSolve, ElastiqpSolveImpl,
|
|
|
224
205
|
.Ret<ffi::Buffer<ffi::F64>>() // z_r
|
|
225
206
|
.Ret<ffi::Buffer<ffi::F64>>() // info
|
|
226
207
|
);
|
|
227
|
-
|
|
228
|
-
XLA_FFI_DEFINE_HANDLER_SYMBOL(ElastiqpSolveWarm, ElastiqpSolveWarmImpl,
|
|
229
|
-
ffi::Ffi::Bind()
|
|
230
|
-
.Attr<double>("eps_abs")
|
|
231
|
-
.Attr<int64_t>("max_iter")
|
|
232
|
-
.Attr<int64_t>("ruiz")
|
|
233
|
-
.Attr<int64_t>("method")
|
|
234
|
-
.Arg<ffi::Buffer<ffi::F64>>() // Q
|
|
235
|
-
.Arg<ffi::Buffer<ffi::F64>>() // q
|
|
236
|
-
.Arg<ffi::Buffer<ffi::F64>>() // A
|
|
237
|
-
.Arg<ffi::Buffer<ffi::F64>>() // b
|
|
238
|
-
.Arg<ffi::Buffer<ffi::F64>>() // G
|
|
239
|
-
.Arg<ffi::Buffer<ffi::F64>>() // h
|
|
240
|
-
.Arg<ffi::Buffer<ffi::F64>>() // penalty
|
|
241
|
-
.Arg<ffi::Buffer<ffi::F64>>() // x0
|
|
242
|
-
.Arg<ffi::Buffer<ffi::F64>>() // y0
|
|
243
|
-
.Arg<ffi::Buffer<ffi::F64>>() // z0
|
|
244
|
-
.Ret<ffi::Buffer<ffi::F64>>() // x
|
|
245
|
-
.Ret<ffi::Buffer<ffi::F64>>() // t
|
|
246
|
-
.Ret<ffi::Buffer<ffi::F64>>() // y
|
|
247
|
-
.Ret<ffi::Buffer<ffi::F64>>() // z_t
|
|
248
|
-
.Ret<ffi::Buffer<ffi::F64>>() // z
|
|
249
|
-
.Ret<ffi::Buffer<ffi::F64>>() // info
|
|
250
|
-
);
|
|
@@ -21,8 +21,12 @@ struct Settings {
|
|
|
21
21
|
double eps_rel = 0.0;
|
|
22
22
|
// Working-set LDL^T pivot below this = dependent row.
|
|
23
23
|
double sing_tol = 3.7e-11;
|
|
24
|
-
// Min/max
|
|
24
|
+
// Min/max eigenvalue ratio accepted for Q_s before adding a prox shift,
|
|
25
|
+
// judged from the Cholesky pivots (see cond_iters).
|
|
25
26
|
double zero_tol = 1e-11;
|
|
27
|
+
// Inverse-iteration steps confirming the smallest eigenvalue of Q_s against
|
|
28
|
+
// zero_tol when the pivot ratio is below sqrt(zero_tol); 0 = pivots only.
|
|
29
|
+
int cond_iters = 3;
|
|
26
30
|
// Initial prox shift (x max|diag Q_s|) when Q_s is not PD; 0 disables.
|
|
27
31
|
double eps_prox = 1e-6;
|
|
28
32
|
// Outer-loop tolerance on eps*|x - xc|_inf; 0 -> eps_abs.
|
|
@@ -109,6 +113,7 @@ class Solver {
|
|
|
109
113
|
fx_.resize(n_);
|
|
110
114
|
fr_.resize(mp_);
|
|
111
115
|
xu_.resize(n_);
|
|
116
|
+
eig_work_.resize(n_);
|
|
112
117
|
changed_.reserve(static_cast<size_t>(mp_));
|
|
113
118
|
changed_cols_.resize(n_, mp_);
|
|
114
119
|
warm_x_.setZero(n_);
|
|
@@ -305,16 +310,23 @@ class Solver {
|
|
|
305
310
|
int inner = 0;
|
|
306
311
|
st = ldp(inner);
|
|
307
312
|
sol_.iters += inner;
|
|
308
|
-
if (st == Status::kNumerics &&
|
|
313
|
+
if (st == Status::kNumerics && settings.eps_prox > 0 &&
|
|
309
314
|
escalations < settings.prox_escalations) {
|
|
310
315
|
// Inner loop broke down: recenter at the current x, grow the prox shift
|
|
311
|
-
// 100x, restart.
|
|
316
|
+
// 100x (or turn it on), restart.
|
|
312
317
|
escalations++;
|
|
313
318
|
prox_escalated_ = true;
|
|
314
319
|
for (int i = 0; i < static_cast<int>(W_.size()); ++i)
|
|
315
320
|
lam_full_[W_[static_cast<size_t>(i)]] = lam_[i];
|
|
316
|
-
|
|
317
|
-
if (
|
|
321
|
+
double eps_next = 100.0 * eps_;
|
|
322
|
+
if (eps_ > 0) {
|
|
323
|
+
xc_ = llt_.matrixU().solve(u_ - v_);
|
|
324
|
+
} else {
|
|
325
|
+
xc_ = x_;
|
|
326
|
+
eps_next = settings.eps_prox *
|
|
327
|
+
std::max(1.0, Qs_.diagonal().cwiseAbs().maxCoeff());
|
|
328
|
+
}
|
|
329
|
+
if (!factor(eps_next)) return finish(Status::kNumerics);
|
|
318
330
|
for (int i = 0; i < mp_; ++i)
|
|
319
331
|
tol_[i] = (settings.eps_abs + settings.eps_rel * std::abs(rhs_[i])) *
|
|
320
332
|
scale_[i] * dr_[i];
|
|
@@ -653,6 +665,9 @@ class Solver {
|
|
|
653
665
|
pmax = std::max(pmax, piv * piv);
|
|
654
666
|
}
|
|
655
667
|
ok = pmin > settings.zero_tol * pmax && pmin > 0;
|
|
668
|
+
if (ok && settings.cond_iters > 0 &&
|
|
669
|
+
pmin < std::sqrt(settings.zero_tol) * pmax)
|
|
670
|
+
ok = min_eig_estimate() > settings.zero_tol * (scale + eps_);
|
|
656
671
|
}
|
|
657
672
|
if (ok) break;
|
|
658
673
|
if (settings.eps_prox <= 0) return false;
|
|
@@ -667,6 +682,18 @@ class Solver {
|
|
|
667
682
|
rhs_dirty_ = true;
|
|
668
683
|
return true;
|
|
669
684
|
}
|
|
685
|
+
// Upper estimate of the smallest eigenvalue of the factored matrix
|
|
686
|
+
double min_eig_estimate() {
|
|
687
|
+
for (int i = 0; i < n_; ++i)
|
|
688
|
+
eig_work_[i] = 1.0 + static_cast<double>(i) / n_;
|
|
689
|
+
double est = std::numeric_limits<double>::infinity();
|
|
690
|
+
for (int it = 0; it < settings.cond_iters; ++it) {
|
|
691
|
+
eig_work_ /= eig_work_.norm();
|
|
692
|
+
llt_.solveInPlace(eig_work_);
|
|
693
|
+
est = 1.0 / eig_work_.norm();
|
|
694
|
+
}
|
|
695
|
+
return est;
|
|
696
|
+
}
|
|
670
697
|
void normalize_row(int i) {
|
|
671
698
|
const double nrm = Mt_.col(i).norm();
|
|
672
699
|
scale_[i] = nrm > 1e-300 ? 1.0 / nrm : 1.0;
|
|
@@ -1322,9 +1349,10 @@ class Solver {
|
|
|
1322
1349
|
double eps_ = 0.0;
|
|
1323
1350
|
MatrixXd Mt_;
|
|
1324
1351
|
VectorXd scale_, hi_, d_, v_, u_, uS_, mu_, x_, xc_, lam_full_;
|
|
1325
|
-
// Scratch: Ruiz factors, a user-frame / unscaled n-vector, the
|
|
1326
|
-
// update_rows re-solves, and the rows
|
|
1327
|
-
|
|
1352
|
+
// Scratch: Ruiz factors, a user-frame / unscaled n-vector, the inverse
|
|
1353
|
+
// iteration vector, the columns update_rows re-solves, and the rows
|
|
1354
|
+
// refactor_working_set re-adds.
|
|
1355
|
+
VectorXd fx_, fr_, xu_, eig_work_;
|
|
1328
1356
|
MatrixXd changed_cols_;
|
|
1329
1357
|
std::vector<int> changed_, rows_;
|
|
1330
1358
|
std::vector<RowState> state_;
|
|
@@ -74,7 +74,8 @@ typedef struct {
|
|
|
74
74
|
* elastiqp_option_info() gives their ranges. Out-of-range values are rejected.
|
|
75
75
|
*
|
|
76
76
|
* Tolerances and iteration limits apply on the next solve. Scaling and
|
|
77
|
-
* factorization options (ruiz*, zero_tol, eps_prox,
|
|
77
|
+
* factorization options (ruiz*, zero_tol, cond_iters, eps_prox,
|
|
78
|
+
* reuse_factorization)
|
|
78
79
|
* apply on refactorization; updating Q forces it. Disabling ruiz retains
|
|
79
80
|
* existing scaling. Explicit warm starts apply even with warm_start = 0.
|
|
80
81
|
*/
|
|
@@ -105,6 +106,8 @@ typedef struct {
|
|
|
105
106
|
|
|
106
107
|
double relax_reg;
|
|
107
108
|
int relax_factor_retries;
|
|
109
|
+
|
|
110
|
+
int cond_iters; /* appended to keep the struct layout of earlier fields */
|
|
108
111
|
} ElastiQPSettings;
|
|
109
112
|
|
|
110
113
|
/*
|
|
@@ -2,7 +2,7 @@
|
|
|
2
2
|
<?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
|
|
3
3
|
<package format="3">
|
|
4
4
|
<name>elastiqp</name>
|
|
5
|
-
<version>0.2.
|
|
5
|
+
<version>0.2.1</version>
|
|
6
6
|
<description>ElastiQP: An elastic QP solver for robot control</description>
|
|
7
7
|
<maintainer email="danielpmorton@gmail.com">Daniel Morton</maintainer>
|
|
8
8
|
<license>Apache 2.0</license>
|
|
@@ -8,7 +8,7 @@ Supported:
|
|
|
8
8
|
Notes:
|
|
9
9
|
- vmap is sequential mode
|
|
10
10
|
- Requires float64
|
|
11
|
-
- Warm-starting is
|
|
11
|
+
- Warm-starting is differentiable (the warm start itself gets a zero gradient)
|
|
12
12
|
- Warm-starting requires explicit passing of the previous result
|
|
13
13
|
- Recommended default kappa for differentiability: 1e-3
|
|
14
14
|
"""
|
|
@@ -66,11 +66,6 @@ _lib = ctypes.cdll.LoadLibrary(str(_find_library()))
|
|
|
66
66
|
jax.ffi.register_ffi_target(
|
|
67
67
|
"elastiqp_solve", jax.ffi.pycapsule(_lib.ElastiqpSolve), platform="cpu"
|
|
68
68
|
)
|
|
69
|
-
jax.ffi.register_ffi_target(
|
|
70
|
-
"elastiqp_solve_warm",
|
|
71
|
-
jax.ffi.pycapsule(_lib.ElastiqpSolveWarm),
|
|
72
|
-
platform="cpu",
|
|
73
|
-
)
|
|
74
69
|
|
|
75
70
|
|
|
76
71
|
class Result(NamedTuple):
|
|
@@ -91,10 +86,14 @@ def _ffi_solve(
|
|
|
91
86
|
G,
|
|
92
87
|
h,
|
|
93
88
|
penalty,
|
|
89
|
+
x0,
|
|
90
|
+
y0,
|
|
91
|
+
z0,
|
|
94
92
|
eps_abs,
|
|
95
93
|
max_iter,
|
|
96
94
|
ruiz,
|
|
97
95
|
target_kappa,
|
|
96
|
+
warm,
|
|
98
97
|
method,
|
|
99
98
|
):
|
|
100
99
|
n = Q.shape[-1]
|
|
@@ -105,60 +104,9 @@ def _ffi_solve(
|
|
|
105
104
|
# (x, t, y, z_t, z) tight solution, then the kappa-relaxed central
|
|
106
105
|
# point (identical to the tight block when target_kappa <= 0), then
|
|
107
106
|
# info = [converged, iters, relax_converged].
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
vec(p),
|
|
111
|
-
vec(m),
|
|
112
|
-
vec(p),
|
|
113
|
-
vec(p),
|
|
114
|
-
vec(n),
|
|
115
|
-
vec(p),
|
|
116
|
-
vec(m),
|
|
117
|
-
vec(p),
|
|
118
|
-
vec(p),
|
|
119
|
-
vec(3),
|
|
120
|
-
]
|
|
107
|
+
sol = [vec(n), vec(p), vec(m), vec(p), vec(p)]
|
|
108
|
+
out_types = sol + sol + [vec(3)]
|
|
121
109
|
call = jax.ffi.ffi_call("elastiqp_solve", out_types, vmap_method="sequential")
|
|
122
|
-
return call(
|
|
123
|
-
Q,
|
|
124
|
-
q,
|
|
125
|
-
A,
|
|
126
|
-
b,
|
|
127
|
-
G,
|
|
128
|
-
h,
|
|
129
|
-
penalty,
|
|
130
|
-
eps_abs=np.float64(eps_abs),
|
|
131
|
-
max_iter=np.int64(max_iter),
|
|
132
|
-
ruiz=np.int64(bool(ruiz)),
|
|
133
|
-
target_kappa=np.float64(target_kappa),
|
|
134
|
-
method=np.int64(_METHOD_ID[method]),
|
|
135
|
-
)
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
def _ffi_solve_warm(
|
|
139
|
-
Q,
|
|
140
|
-
q,
|
|
141
|
-
A,
|
|
142
|
-
b,
|
|
143
|
-
G,
|
|
144
|
-
h,
|
|
145
|
-
penalty,
|
|
146
|
-
x0,
|
|
147
|
-
y0,
|
|
148
|
-
z0,
|
|
149
|
-
eps_abs,
|
|
150
|
-
max_iter,
|
|
151
|
-
ruiz,
|
|
152
|
-
method,
|
|
153
|
-
):
|
|
154
|
-
n = Q.shape[-1]
|
|
155
|
-
m = b.shape[-1]
|
|
156
|
-
p = h.shape[-1]
|
|
157
|
-
batch = Q.shape[:-2]
|
|
158
|
-
vec = lambda d: jax.ShapeDtypeStruct(batch + (d,), jnp.float64)
|
|
159
|
-
# (x, t, y, z_t, z), then info = [converged, iters].
|
|
160
|
-
out_types = [vec(n), vec(p), vec(m), vec(p), vec(p), vec(2)]
|
|
161
|
-
call = jax.ffi.ffi_call("elastiqp_solve_warm", out_types, vmap_method="sequential")
|
|
162
110
|
return call(
|
|
163
111
|
Q,
|
|
164
112
|
q,
|
|
@@ -173,56 +121,12 @@ def _ffi_solve_warm(
|
|
|
173
121
|
eps_abs=np.float64(eps_abs),
|
|
174
122
|
max_iter=np.int64(max_iter),
|
|
175
123
|
ruiz=np.int64(bool(ruiz)),
|
|
124
|
+
target_kappa=np.float64(target_kappa),
|
|
125
|
+
warm=np.int64(bool(warm)),
|
|
176
126
|
method=np.int64(_METHOD_ID[method]),
|
|
177
127
|
)
|
|
178
128
|
|
|
179
129
|
|
|
180
|
-
@partial(jax.custom_vjp, nondiff_argnums=(10, 11, 12, 13))
|
|
181
|
-
def _solve_warm(
|
|
182
|
-
Q,
|
|
183
|
-
q,
|
|
184
|
-
A,
|
|
185
|
-
b,
|
|
186
|
-
G,
|
|
187
|
-
h,
|
|
188
|
-
penalty,
|
|
189
|
-
x0,
|
|
190
|
-
y0,
|
|
191
|
-
z0,
|
|
192
|
-
eps_abs,
|
|
193
|
-
max_iter,
|
|
194
|
-
ruiz,
|
|
195
|
-
method,
|
|
196
|
-
):
|
|
197
|
-
return _ffi_solve_warm(
|
|
198
|
-
Q,
|
|
199
|
-
q,
|
|
200
|
-
A,
|
|
201
|
-
b,
|
|
202
|
-
G,
|
|
203
|
-
h,
|
|
204
|
-
penalty,
|
|
205
|
-
x0,
|
|
206
|
-
y0,
|
|
207
|
-
z0,
|
|
208
|
-
eps_abs,
|
|
209
|
-
max_iter,
|
|
210
|
-
ruiz,
|
|
211
|
-
method,
|
|
212
|
-
)
|
|
213
|
-
|
|
214
|
-
|
|
215
|
-
def _solve_warm_fwd(*args):
|
|
216
|
-
raise TypeError("elastiqp.jax.solve is not differentiable with warm_start")
|
|
217
|
-
|
|
218
|
-
|
|
219
|
-
def _solve_warm_bwd(*args):
|
|
220
|
-
raise TypeError("elastiqp.jax.solve is not differentiable with warm_start")
|
|
221
|
-
|
|
222
|
-
|
|
223
|
-
_solve_warm.defvjp(_solve_warm_fwd, _solve_warm_bwd)
|
|
224
|
-
|
|
225
|
-
|
|
226
130
|
def _outer(a, b):
|
|
227
131
|
return a[..., :, None] * b[..., None, :]
|
|
228
132
|
|
|
@@ -275,7 +179,9 @@ def _kkt_bwd(res, ct):
|
|
|
275
179
|
return Qb, qb, Ab, bb, Gb, hb, penalty_b
|
|
276
180
|
|
|
277
181
|
|
|
278
|
-
|
|
182
|
+
# target_kappa is only used by _solve_fwd; custom_vjp needs the primal and
|
|
183
|
+
# fwd signatures to match. (x0, y0, z0) are ignored unless warm.
|
|
184
|
+
@partial(jax.custom_vjp, nondiff_argnums=(10, 11, 12, 13, 14, 15))
|
|
279
185
|
def _solve(
|
|
280
186
|
Q,
|
|
281
187
|
q,
|
|
@@ -284,10 +190,14 @@ def _solve(
|
|
|
284
190
|
G,
|
|
285
191
|
h,
|
|
286
192
|
penalty,
|
|
193
|
+
x0,
|
|
194
|
+
y0,
|
|
195
|
+
z0,
|
|
287
196
|
eps_abs,
|
|
288
197
|
max_iter,
|
|
289
198
|
ruiz,
|
|
290
199
|
target_kappa,
|
|
200
|
+
warm,
|
|
291
201
|
method,
|
|
292
202
|
):
|
|
293
203
|
# Tight solution when not differentiating
|
|
@@ -299,10 +209,14 @@ def _solve(
|
|
|
299
209
|
G,
|
|
300
210
|
h,
|
|
301
211
|
penalty,
|
|
212
|
+
x0,
|
|
213
|
+
y0,
|
|
214
|
+
z0,
|
|
302
215
|
eps_abs,
|
|
303
216
|
max_iter,
|
|
304
217
|
ruiz,
|
|
305
218
|
0.0,
|
|
219
|
+
warm,
|
|
306
220
|
method,
|
|
307
221
|
)
|
|
308
222
|
return tuple(out[:5]) + (out[10],) # solution + info
|
|
@@ -316,10 +230,14 @@ def _solve_fwd(
|
|
|
316
230
|
G,
|
|
317
231
|
h,
|
|
318
232
|
penalty,
|
|
233
|
+
x0,
|
|
234
|
+
y0,
|
|
235
|
+
z0,
|
|
319
236
|
eps_abs,
|
|
320
237
|
max_iter,
|
|
321
238
|
ruiz,
|
|
322
239
|
target_kappa,
|
|
240
|
+
warm,
|
|
323
241
|
method,
|
|
324
242
|
):
|
|
325
243
|
if not target_kappa > 0:
|
|
@@ -335,18 +253,27 @@ def _solve_fwd(
|
|
|
335
253
|
G,
|
|
336
254
|
h,
|
|
337
255
|
penalty,
|
|
256
|
+
x0,
|
|
257
|
+
y0,
|
|
258
|
+
z0,
|
|
338
259
|
eps_abs,
|
|
339
260
|
max_iter,
|
|
340
261
|
ruiz,
|
|
341
262
|
target_kappa,
|
|
263
|
+
warm,
|
|
342
264
|
method,
|
|
343
265
|
)
|
|
344
266
|
x, t, y, z_t, z, xr, tr, yr, z_t_r, z_r, info = out
|
|
345
267
|
return (x, t, y, z_t, z, info), (Q, A, G, h, xr, tr, yr, z_t_r, z_r)
|
|
346
268
|
|
|
347
269
|
|
|
348
|
-
def _solve_bwd(eps_abs, max_iter, ruiz, target_kappa, method, res, ct):
|
|
349
|
-
|
|
270
|
+
def _solve_bwd(eps_abs, max_iter, ruiz, target_kappa, warm, method, res, ct):
|
|
271
|
+
# Zero cotangent for the warm start: a converged solve of a QP with a
|
|
272
|
+
# unique solution does not depend on where the solver starts. The
|
|
273
|
+
# (x, y, z) cotangents share the (x0, y0, z0) shapes.
|
|
274
|
+
xb, _, yb, _, zb, _ = ct
|
|
275
|
+
zeros = (jnp.zeros_like(xb), jnp.zeros_like(yb), jnp.zeros_like(zb))
|
|
276
|
+
return _kkt_bwd(res, ct) + zeros
|
|
350
277
|
|
|
351
278
|
|
|
352
279
|
_solve.defvjp(_solve_fwd, _solve_bwd)
|
|
@@ -396,7 +323,9 @@ def solve(
|
|
|
396
323
|
derivatives. Defaults to 1e-3.
|
|
397
324
|
warm_start (tuple | Result, optional): Explicit warm start, either a
|
|
398
325
|
previous Result or a (x, y, z) tuple, (das and pdal only).
|
|
399
|
-
|
|
326
|
+
Differentiable: the gradients are the cold solve's, and the warm
|
|
327
|
+
start itself gets a zero gradient (exact when the QP has a unique
|
|
328
|
+
solution and the solve converges). Defaults to None.
|
|
400
329
|
|
|
401
330
|
Returns:
|
|
402
331
|
Result: Solution to the elastic QP
|
|
@@ -447,32 +376,11 @@ def solve(
|
|
|
447
376
|
f"{q.shape}, {b.shape}, {h.shape}; got "
|
|
448
377
|
f"{x0.shape}, {y0.shape}, {z0.shape}"
|
|
449
378
|
)
|
|
450
|
-
|
|
451
|
-
|
|
452
|
-
|
|
453
|
-
|
|
454
|
-
|
|
455
|
-
G,
|
|
456
|
-
h,
|
|
457
|
-
penalty,
|
|
458
|
-
x0,
|
|
459
|
-
y0,
|
|
460
|
-
z0,
|
|
461
|
-
float(eps_abs),
|
|
462
|
-
int(max_iter),
|
|
463
|
-
bool(ruiz),
|
|
464
|
-
method,
|
|
465
|
-
)
|
|
466
|
-
return Result(
|
|
467
|
-
x=x,
|
|
468
|
-
t=t,
|
|
469
|
-
y=y,
|
|
470
|
-
z_t=z_t,
|
|
471
|
-
z=z,
|
|
472
|
-
converged=info[..., 0].astype(jnp.int32),
|
|
473
|
-
iters=info[..., 1].astype(jnp.int32),
|
|
474
|
-
)
|
|
475
|
-
|
|
379
|
+
else:
|
|
380
|
+
# Ignored by the solver (warm=False)
|
|
381
|
+
x0 = jnp.zeros(q.shape, jnp.float64)
|
|
382
|
+
y0 = jnp.zeros(b.shape, jnp.float64)
|
|
383
|
+
z0 = jnp.zeros(h.shape, jnp.float64)
|
|
476
384
|
x, t, y, z_t, z, info = _solve(
|
|
477
385
|
Q,
|
|
478
386
|
q,
|
|
@@ -481,10 +389,14 @@ def solve(
|
|
|
481
389
|
G,
|
|
482
390
|
h,
|
|
483
391
|
penalty,
|
|
392
|
+
x0,
|
|
393
|
+
y0,
|
|
394
|
+
z0,
|
|
484
395
|
float(eps_abs),
|
|
485
396
|
int(max_iter),
|
|
486
397
|
bool(ruiz),
|
|
487
398
|
float(target_kappa),
|
|
399
|
+
warm_start is not None,
|
|
488
400
|
method,
|
|
489
401
|
)
|
|
490
402
|
# Recall: info = [converged, iters, relax_converged]. The converged
|
|
@@ -120,8 +120,8 @@ constexpr Option kOptions[] = {
|
|
|
120
120
|
ELASTIQP_D(sing_tol, 0.0,
|
|
121
121
|
"Working-set LDL' pivot below this = dependent row."),
|
|
122
122
|
ELASTIQP_D(zero_tol, 0.0,
|
|
123
|
-
"Min/max
|
|
124
|
-
"a prox shift."),
|
|
123
|
+
"Min/max eigenvalue ratio of Q_s (from its Cholesky pivots, "
|
|
124
|
+
"see cond_iters) accepted before adding a prox shift."),
|
|
125
125
|
ELASTIQP_D(eps_prox, 0.0,
|
|
126
126
|
"Initial prox shift (x max|diag Q_s|) when Q_s is not PD; 0 "
|
|
127
127
|
"disables."),
|
|
@@ -159,6 +159,10 @@ constexpr Option kOptions[] = {
|
|
|
159
159
|
ELASTIQP_I(relax_factor_retries,
|
|
160
160
|
"relax(): x100 regularization retries on a failed "
|
|
161
161
|
"factorization."),
|
|
162
|
+
ELASTIQP_I(cond_iters,
|
|
163
|
+
"Inverse-iteration steps confirming the smallest eigenvalue "
|
|
164
|
+
"of Q_s against zero_tol when the pivots look small; 0: "
|
|
165
|
+
"pivots only."),
|
|
162
166
|
};
|
|
163
167
|
|
|
164
168
|
#undef ELASTIQP_D
|
|
@@ -286,7 +286,7 @@ static void options(void) {
|
|
|
286
286
|
int i, j, type = -1, ok = 1, unique = 1, n_opt = elastiqp_num_options();
|
|
287
287
|
|
|
288
288
|
elastiqp_setup(&work, &qp, NULL);
|
|
289
|
-
check("options:
|
|
289
|
+
check("options: 23 listed", n_opt == 23);
|
|
290
290
|
for (i = 0; i < n_opt; ++i) {
|
|
291
291
|
ok &= elastiqp_option_info(i, &name, &type, &def, &lo, &hi, &desc) == 0;
|
|
292
292
|
ok &= name != NULL && desc != NULL && desc[0] != '\0';
|
|
@@ -14,9 +14,10 @@
|
|
|
14
14
|
// 3. Robustness: a badly scaled parametrization with and without Ruiz;
|
|
15
15
|
// dependent / inconsistent equalities (certificate before iterating);
|
|
16
16
|
// every inequality row duplicated (dependent working sets, the cycle
|
|
17
|
-
// guard's regime, refactors());
|
|
18
|
-
//
|
|
19
|
-
//
|
|
17
|
+
// guard's regime, refactors()); a near-singular Q the Cholesky pivots
|
|
18
|
+
// hide (the eigenvalue check, and the prox fallback without it); the
|
|
19
|
+
// drift-gated Ruiz refresh through a row exploding past the
|
|
20
|
+
// limit_scaling clamp and a row decaying to the noise floor.
|
|
20
21
|
// 4. Factorization reuse: constant Q, a fifth of G changing per tick,
|
|
21
22
|
// same answers as a full refactorization and only those rows re-solved.
|
|
22
23
|
// 5. The pinch/release ticks of test_gap_creep: saturated -> active ->
|
|
@@ -294,6 +295,58 @@ void ProxLoop(std::mt19937& rng) {
|
|
|
294
295
|
"status");
|
|
295
296
|
}
|
|
296
297
|
|
|
298
|
+
// Near-singular Q whose null direction is small on the trailing coordinates
|
|
299
|
+
// (as in diff IK, J'J + 1e-12 I with 6 task rows on 7 DOF): the Cholesky
|
|
300
|
+
// pivots overstate its smallest eigenvalue, so only the inverse-iteration
|
|
301
|
+
// check turns the prox shift on up front. Without it (cond_iters = 0), the
|
|
302
|
+
// inner loop can break down and must recover through the prox fallback.
|
|
303
|
+
void NearSingularQ(std::mt19937& rng) {
|
|
304
|
+
const int n = 7, p = 14, cases = 20;
|
|
305
|
+
int solved = 0, prox = 0, escalated = 0;
|
|
306
|
+
int solved_off = 0, escalated_off = 0;
|
|
307
|
+
double worst = 0.0;
|
|
308
|
+
for (int k = 0; k < cases; ++k) {
|
|
309
|
+
VectorXd v = problem_gen::Randn(rng, n, 1);
|
|
310
|
+
v.tail(2) *= 0.05;
|
|
311
|
+
v.normalize();
|
|
312
|
+
const MatrixXd J = problem_gen::Randn(rng, n - 1, n) *
|
|
313
|
+
(MatrixXd::Identity(n, n) - v * v.transpose());
|
|
314
|
+
const MatrixXd Q = J.transpose() * J + 1e-12 * MatrixXd::Identity(n, n);
|
|
315
|
+
const VectorXd q = -J.transpose() * problem_gen::Randn(rng, n - 1, 1);
|
|
316
|
+
MatrixXd G(p, n);
|
|
317
|
+
G << MatrixXd::Identity(n, n), -MatrixXd::Identity(n, n);
|
|
318
|
+
const VectorXd h = VectorXd::Constant(p, 0.5) +
|
|
319
|
+
0.1 * problem_gen::Randn(rng, p, 1);
|
|
320
|
+
const VectorXd w = VectorXd::Constant(p, 1e3);
|
|
321
|
+
das::Solver s;
|
|
322
|
+
s.setup(Q, q, G, h, w);
|
|
323
|
+
const elastiqp::Solution sol = s.solve();
|
|
324
|
+
solved += sol.converged;
|
|
325
|
+
prox += s.proximal();
|
|
326
|
+
escalated += s.prox_escalated();
|
|
327
|
+
QPData qp;
|
|
328
|
+
qp.Q = Q;
|
|
329
|
+
qp.q = q;
|
|
330
|
+
qp.A = MatrixXd(0, n);
|
|
331
|
+
qp.b = VectorXd(0);
|
|
332
|
+
qp.G = G;
|
|
333
|
+
qp.h = h;
|
|
334
|
+
worst = std::max(worst, Kkt(qp, w, sol));
|
|
335
|
+
das::Solver off;
|
|
336
|
+
off.settings.cond_iters = 0;
|
|
337
|
+
off.setup(Q, q, G, h, w);
|
|
338
|
+
solved_off += off.solve().converged;
|
|
339
|
+
escalated_off += off.prox_escalated();
|
|
340
|
+
}
|
|
341
|
+
Check("near-singular Q: solved, prox on up front",
|
|
342
|
+
solved == cases && prox == cases && escalated == 0,
|
|
343
|
+
static_cast<double>(solved), "solved");
|
|
344
|
+
Check("near-singular Q: elastic KKT residual", worst <= 1e-6, worst, "kkt");
|
|
345
|
+
Check("near-singular Q, cond_iters = 0: solved via the prox fallback",
|
|
346
|
+
solved_off == cases && escalated_off > 0,
|
|
347
|
+
static_cast<double>(escalated_off), "escalated");
|
|
348
|
+
}
|
|
349
|
+
|
|
297
350
|
// Drift-gated Ruiz refresh through the update path: a row exploding past
|
|
298
351
|
// the 1e4 clamp in one tick (rescaled() fires, drift back to O(1)) and a
|
|
299
352
|
// row decaying to the noise floor (left alone, solved as if absent).
|
|
@@ -470,6 +523,7 @@ void RandomSuite() {
|
|
|
470
523
|
EqualityCases(rng);
|
|
471
524
|
DegenerateCase(rng);
|
|
472
525
|
ProxLoop(rng);
|
|
526
|
+
NearSingularQ(rng);
|
|
473
527
|
RuizRefresh(rng);
|
|
474
528
|
ReuseBench();
|
|
475
529
|
}
|
|
@@ -14,6 +14,29 @@ import numpy as np
|
|
|
14
14
|
RESULTS = []
|
|
15
15
|
|
|
16
16
|
|
|
17
|
+
def ffi_solve_cold(
|
|
18
|
+
Q, q, A, b, G, h, penalty, eps_abs, max_iter, ruiz, target_kappa, method
|
|
19
|
+
):
|
|
20
|
+
# Raw FFI call without a warm start (x0, y0, z0 are ignored placeholders)
|
|
21
|
+
zeros = (jnp.zeros_like(q), jnp.zeros_like(b), jnp.zeros_like(h))
|
|
22
|
+
return elastiqp.jax._ffi_solve(
|
|
23
|
+
Q,
|
|
24
|
+
q,
|
|
25
|
+
A,
|
|
26
|
+
b,
|
|
27
|
+
G,
|
|
28
|
+
h,
|
|
29
|
+
penalty,
|
|
30
|
+
*zeros,
|
|
31
|
+
eps_abs,
|
|
32
|
+
max_iter,
|
|
33
|
+
ruiz,
|
|
34
|
+
target_kappa,
|
|
35
|
+
False,
|
|
36
|
+
method,
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
|
|
17
40
|
def check(name, ok, detail=""):
|
|
18
41
|
RESULTS.append(ok)
|
|
19
42
|
print(f" {name:<52s} {detail:<24s} {'OK' if ok else 'FAIL'}")
|
|
@@ -198,7 +221,7 @@ def main():
|
|
|
198
221
|
return w_loss @ s.x + w_t @ s.t
|
|
199
222
|
|
|
200
223
|
def loss_relaxed(Q_, q_, A_, b_, G_, h_, penalty_, kap=kappa, method="das"):
|
|
201
|
-
out =
|
|
224
|
+
out = ffi_solve_cold(
|
|
202
225
|
Q_,
|
|
203
226
|
q_,
|
|
204
227
|
A_,
|
|
@@ -317,7 +340,7 @@ def main():
|
|
|
317
340
|
# relaxation would otherwise mean silently wrong gradients). s_t = t and
|
|
318
341
|
# s_in = h + t - Gx are the slacks of t >= 0 and Gx - t <= h; s_in is
|
|
319
342
|
# reconstructed from x, so it carries the O(tol) primal residual.
|
|
320
|
-
out =
|
|
343
|
+
out = ffi_solve_cold(*args, 1e-11, 300, False, kappa, "das")
|
|
321
344
|
xr, tr, z_t_r, z_r, info = out[5], out[6], out[8], out[9], out[10]
|
|
322
345
|
s_in_r = args[5] + tr - args[4] @ xr
|
|
323
346
|
comp = max(
|
|
@@ -334,7 +357,7 @@ def main():
|
|
|
334
357
|
# silently: an absurd kappa converges the solve (info[0]) but stalls the
|
|
335
358
|
# relaxation (info[2]), and solve() folds that into converged on the
|
|
336
359
|
# differentiated path.
|
|
337
|
-
out =
|
|
360
|
+
out = ffi_solve_cold(*args, 1e-11, 300, False, 1e8, "das")
|
|
338
361
|
tight_ok, relax_bad = float(out[10][0]) == 1.0, float(out[10][2]) == 0.0
|
|
339
362
|
|
|
340
363
|
def conv_smooth(q_):
|
|
@@ -468,9 +491,7 @@ def main():
|
|
|
468
491
|
try:
|
|
469
492
|
jax.grad(
|
|
470
493
|
lambda q_: jnp.sum(
|
|
471
|
-
elastiqp.jax.solve(
|
|
472
|
-
Qx, q_, Gx, hx, 10.0, A=Ax, b=bx, target_kappa=0.0
|
|
473
|
-
).x
|
|
494
|
+
elastiqp.jax.solve(Qx, q_, Gx, hx, 10.0, A=Ax, b=bx, target_kappa=0.0).x
|
|
474
495
|
)
|
|
475
496
|
)(qx)
|
|
476
497
|
msg = None
|
|
@@ -517,7 +538,7 @@ def main():
|
|
|
517
538
|
for m in elastiqp.jax.METHODS
|
|
518
539
|
}
|
|
519
540
|
relaxed = {
|
|
520
|
-
m:
|
|
541
|
+
m: ffi_solve_cold(*args, 1e-11, 300, False, kappa, m)
|
|
521
542
|
for m in elastiqp.jax.METHODS
|
|
522
543
|
}
|
|
523
544
|
for m in ("pdal", "ipm"):
|
|
@@ -566,9 +587,9 @@ def main():
|
|
|
566
587
|
int(a_sol.converged) == 1 and da < FFI_NB_TOL,
|
|
567
588
|
f"|dx|={da:.1e}",
|
|
568
589
|
)
|
|
569
|
-
a_jit = jax.jit(
|
|
570
|
-
|
|
571
|
-
)
|
|
590
|
+
a_jit = jax.jit(lambda q_: elastiqp.jax.solve(Qx, q_, Gx, hx, 10.0, A=Ax, b=bx).x)(
|
|
591
|
+
qx
|
|
592
|
+
)
|
|
572
593
|
check("das under jit", float(jnp.abs(a_jit - a_sol.x).max()) == 0.0, "")
|
|
573
594
|
# Hard rows (penalty = inf, das only): finite gradients, none w.r.t. the
|
|
574
595
|
# hard row's penalty, close to a stiff elastic penalty in its place.
|
|
@@ -671,8 +692,15 @@ def main():
|
|
|
671
692
|
check("warm sequence has working-set changes", ws_changed > 0, f"{ws_changed}")
|
|
672
693
|
try:
|
|
673
694
|
elastiqp.jax.solve(
|
|
674
|
-
Qw,
|
|
675
|
-
|
|
695
|
+
Qw,
|
|
696
|
+
drift[0][0],
|
|
697
|
+
Gw,
|
|
698
|
+
drift[0][1],
|
|
699
|
+
10.0,
|
|
700
|
+
A=Aw,
|
|
701
|
+
b=drift[0][2],
|
|
702
|
+
method="ipm",
|
|
703
|
+
warm_start=prev,
|
|
676
704
|
)
|
|
677
705
|
ipm_raised = False
|
|
678
706
|
except ValueError:
|
|
@@ -755,22 +783,57 @@ def main():
|
|
|
755
783
|
f"|dx|={dxv:.1e} max iters={int(v_warm.iters.max())}",
|
|
756
784
|
)
|
|
757
785
|
|
|
758
|
-
|
|
759
|
-
|
|
760
|
-
|
|
761
|
-
|
|
762
|
-
|
|
763
|
-
|
|
786
|
+
# A warm-started solve has the cold solve's gradients (the solution does
|
|
787
|
+
# not depend on the start), and a zero gradient for the warm start itself
|
|
788
|
+
def loss_w(Q_, q_, G_, h_, pen_, warm, method):
|
|
789
|
+
s = elastiqp.jax.solve(
|
|
790
|
+
Q_, q_, G_, h_, pen_, A=Aw, b=b_k, method=method, warm_start=warm
|
|
791
|
+
)
|
|
792
|
+
return jnp.sum(s.x**2) + jnp.sum(jnp.sin(s.x))
|
|
793
|
+
|
|
794
|
+
pen_k = jnp.full(h_k.shape, 10.0)
|
|
795
|
+
warm = (p_sol.x, p_sol.y, p_sol.z)
|
|
796
|
+
for method in ("das", "pdal"):
|
|
797
|
+
g_cold = jax.grad(loss_w, argnums=(0, 1, 2, 3, 4))(
|
|
798
|
+
Qw, q_k, Gw, h_k, pen_k, None, method
|
|
799
|
+
)
|
|
800
|
+
*g_warm, g_start = jax.grad(loss_w, argnums=(0, 1, 2, 3, 4, 5))(
|
|
801
|
+
Qw, q_k, Gw, h_k, pen_k, warm, method
|
|
802
|
+
)
|
|
803
|
+
dg = max(
|
|
804
|
+
float(jnp.abs(a - b).max()) for a, b in zip(g_cold, g_warm, strict=True)
|
|
805
|
+
)
|
|
806
|
+
g0 = max(float(jnp.abs(g).max()) for g in g_start)
|
|
807
|
+
check(
|
|
808
|
+
f"grad with warm_start ({method}) == cold",
|
|
809
|
+
dg < 1e-8 and g0 == 0.0,
|
|
810
|
+
f"|dgrad|={dg:.1e} |grad warm start|={g0:.1e}",
|
|
811
|
+
)
|
|
812
|
+
|
|
813
|
+
# Warm-started lax.scan works under jit + grad and matches the cold scan.
|
|
814
|
+
# The carry's cotangent is zero by construction, so this is a smoke test
|
|
815
|
+
# of the warm path, not an independent check of the gradient through it
|
|
816
|
+
def scan_loss(q_, warm_carry):
|
|
817
|
+
def body(carry, h_):
|
|
818
|
+
s = elastiqp.jax.solve(
|
|
819
|
+
Qw,
|
|
820
|
+
q_,
|
|
821
|
+
Gw,
|
|
822
|
+
h_,
|
|
823
|
+
10.0,
|
|
824
|
+
A=Aw,
|
|
825
|
+
b=b_k,
|
|
826
|
+
warm_start=carry if warm_carry else None,
|
|
764
827
|
)
|
|
765
|
-
|
|
766
|
-
|
|
767
|
-
|
|
768
|
-
|
|
769
|
-
|
|
770
|
-
|
|
771
|
-
|
|
772
|
-
|
|
773
|
-
)
|
|
828
|
+
return (s.x, s.y, s.z), jnp.sum(s.x**2)
|
|
829
|
+
|
|
830
|
+
_, losses = jax.lax.scan(body, warm, hs)
|
|
831
|
+
return jnp.sum(losses)
|
|
832
|
+
|
|
833
|
+
g_cold = jax.jit(jax.grad(scan_loss), static_argnums=1)(q_k, False)
|
|
834
|
+
g_warm = jax.jit(jax.grad(scan_loss), static_argnums=1)(q_k, True)
|
|
835
|
+
dgs = float(jnp.abs(g_cold - g_warm).max())
|
|
836
|
+
check("grad through a warm-started lax.scan", dgs < 1e-8, f"|dgrad|={dgs:.1e}")
|
|
774
837
|
try:
|
|
775
838
|
elastiqp.jax.solve(
|
|
776
839
|
Qw,
|
|
@@ -825,8 +888,15 @@ def main():
|
|
|
825
888
|
|
|
826
889
|
def x0_of_h0(h0, method):
|
|
827
890
|
return elastiqp.jax.solve(
|
|
828
|
-
Qb,
|
|
829
|
-
|
|
891
|
+
Qb,
|
|
892
|
+
qb,
|
|
893
|
+
Gb,
|
|
894
|
+
hb.at[0].set(h0),
|
|
895
|
+
10.0,
|
|
896
|
+
A=Ab,
|
|
897
|
+
b=bb,
|
|
898
|
+
method=method,
|
|
899
|
+
target_kappa=1e-9,
|
|
830
900
|
).x[0]
|
|
831
901
|
|
|
832
902
|
for method in ["das", "pdal", "ipm"]:
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|