pounce_nl/nl_quadratic.rs
1//! Degree-≤2 recognition over an [`Expr`] DAG.
2//!
3//! This is the classifier's "is this row a quadratic, and which one?"
4//! question, answered once and reused. `pounce-cli`'s `dispatch` owns the
5//! *routing* decision (which `ProblemClass`, which solver); what lives here
6//! is only the algebra, so the consumers that are not the CLI —
7//! `NlTnlp`'s constant-structure evaluation, the parse-time recognizer, the
8//! `QcqpProblem` extractor — can reach it without depending on the
9//! command-line driver. It moved out of `pounce-cli/src/dispatch.rs` in
10//! Q3 of the #588 series; see
11//! `dev-notes/quadratic-structure-exploitation.md`.
12//!
13//! ## Two properties this module is built around
14//!
15//! **It is iterative.** The walk carries its own work stack rather than
16//! recursing, because the trees it is handed are not shallow: a `.nl`
17//! writer that emits `o0` (binary `+`) chains for a long sum — Pyomo does —
18//! produces a left-deep `Add` tree one level per term. The recursive
19//! predecessor aborted the process somewhere between 4 000 and 6 000 terms
20//! on a 2 MB thread (which is what a test gets, and where the crash was
21//! first reproduced) and between 16 000 and 24 000 on the CLI's 8 MB main
22//! thread. A stack overflow is an abort, not an error return, so the depth a
23//! *recognizer* survives must not depend on which thread called it.
24//!
25//! It is worth knowing what this does **not** fix: `nl_reader`'s parser
26//! recurses too, with a fatter frame — it gives out at ~6 000 on that same
27//! 8 MB thread — so a deep `.nl` file still fails to load, and it fails
28//! before reaching this module. What is fixed is every path where the tree
29//! is already built (`NlProblem::from_expressions`, a model handed across
30//! threads) and the ceiling this module used to impose on the parser's
31//! successor.
32//!
33//! **It never allocates per monomial.** The predecessor keyed monomials on
34//! `BTreeMap<Vec<usize>, f64>`, so every term cost a heap allocation and
35//! every merge cloned one (`entry(m.clone())`). A degree-≤2 form has only
36//! three shapes of term, so [`Quad2`] stores them in three fields with
37//! inline keys and the allocation disappears. The same change removes the
38//! `O(N²)` accumulation on `Add` chains: the old `add` re-scanned the whole
39//! accumulated map for zeros on *every* merge, which is quadratic down a
40//! left-deep chain. Zeros can only appear where a merge touched, so that is
41//! all this one looks at.
42
43use crate::nl_reader::{BinOp, Expr, UnaryOp};
44use std::collections::{BTreeMap, BTreeSet};
45
46/// The symmetric Hessian of a quadratic form, stored as a sparse upper-
47/// triangular (i ≤ j) map of `(i, j) -> ∂²/∂xᵢ∂xⱼ`. Empty means the
48/// expression is (at most) linear.
49pub type QuadHessian = BTreeMap<(usize, usize), f64>;
50
51/// Full quadratic read-out: `(Hessian, [(var, linear coef), …], constant)`.
52/// The linear and constant parts are the pieces AMPL/Pyomo fold into the
53/// nonlinear objective tree (see [`analyze_quadratic_full`]).
54pub type QuadForm = (QuadHessian, Vec<(usize, f64)>, f64);
55
56/// A polynomial of total degree ≤ 2 in its own shape: a constant, the
57/// linear coefficients keyed by variable, and the quadratic coefficients
58/// keyed by the (i ≤ j) variable pair.
59///
60/// This replaces a general `BTreeMap<Vec<usize>, f64>` polynomial. Degree
61/// is a property of the *type* here rather than of the data, so the
62/// "is it still quadratic?" test is three `is_empty()` calls instead of a
63/// scan, and no monomial key is ever allocated or cloned.
64///
65/// ### Zero coefficients
66///
67/// Stored coefficients are nonzero: [`add`](Quad2::add) and
68/// [`mul`](Quad2::mul) drop any entry they leave at exactly zero, which is
69/// what makes [`degree`](Quad2::degree) and [`as_constant`](Quad2::as_constant)
70/// answerable in `O(1)`. `constant` is the one exception and needs none:
71/// both `0.0` and `-0.0` in that field mean "no constant term", every
72/// consumer guards on `!= 0.0`, and [`analyze_quadratic_full`] normalizes
73/// the sign on the way out.
74///
75/// ### …and when dropping one loses something
76///
77/// Dropping is right for the *storage* question and can be wrong for the
78/// *degree* question, and [`lost_terms`](Quad2::lost_terms) is the
79/// difference (gh #683, sharpened by gh #687). A coefficient that reaches
80/// zero is a coefficient that was **summed**, and it is the arithmetic of
81/// that sum — not the drop — that decides whether anything went missing:
82///
83/// * `x − x` folds `fl(1) + fl(−1)` to `0.0`, and that add is **exact**.
84/// The term really is absent, the body really is degree 0, and the maps
85/// are not a lower bound on anything.
86/// * `2⁵³·x² + x² − 2⁵³·x²` loses the `x²` at `fl(2⁵³ + 1) = 2⁵³`, which is
87/// an **inexact** add, and only then does `2⁵³ − 2⁵³` drop the survivor.
88/// The body is `x²`; the maps say degree 0.
89///
90/// The loss happens at the inexact fold; the drop is only where it becomes
91/// visible. So the flag is set from the *fold*: a form whose arithmetic
92/// never rounded — and never flushed a live coefficient to zero — carries
93/// coefficients that are exactly what real arithmetic on the writer's own
94/// literals would give, and a term missing from its maps is a term that is
95/// genuinely absent. Flagging the drop instead refused both of the above
96/// alike, which was sound and cost reach on the first (gh #687).
97#[derive(Debug, Clone, Default, PartialEq)]
98pub struct Quad2 {
99 constant: f64,
100 linear: BTreeMap<usize, f64>,
101 quadratic: QuadHessian,
102 /// Set when a term may be **missing** from the maps — see the type
103 /// docs and [`lost_terms`](Quad2::lost_terms). Sticky: it survives
104 /// every operation the form takes part in, because a term that went
105 /// missing three additions ago is no less missing now.
106 lost_terms: bool,
107 /// Set when some arithmetic in this form's construction **rounded**:
108 /// an add that could not represent its sum, or a product that could
109 /// not represent itself. Sticky, and for the same reason.
110 ///
111 /// Clear ⇒ every stored coefficient is exactly the real-arithmetic
112 /// value of the terms folded into it, so a coefficient that reached
113 /// zero reached it because the real value is zero. That is the whole
114 /// warrant for `lost_terms` staying clear across a cancellation, and
115 /// it is why an inexact *fold* sets this bit even when the coefficient
116 /// it produced is nonzero and stays stored: what it rounded away can
117 /// be cancelled out of sight later, by an add that is itself exact.
118 inexact: bool,
119}
120
121impl Quad2 {
122 /// The degree-0 term.
123 pub fn constant(&self) -> f64 {
124 self.constant
125 }
126
127 /// The degree-1 terms, ascending by variable index.
128 pub fn linear(&self) -> &BTreeMap<usize, f64> {
129 &self.linear
130 }
131
132 /// The degree-2 terms as *polynomial* coefficients keyed `(i ≤ j)` —
133 /// the coefficient of `xᵢxⱼ`, **not** the Hessian entry (they differ by
134 /// a factor of 2 on the diagonal; [`analyze_quadratic_full`] applies it).
135 pub fn quadratic(&self) -> &QuadHessian {
136 &self.quadratic
137 }
138
139 /// Whether a term may be **missing** from this form's maps: something
140 /// was dropped on the way here *and* the arithmetic that produced it
141 /// had already rounded, or a live coefficient was flushed to zero (or
142 /// to `NaN`) outright.
143 ///
144 /// This is what turns the term maps from an answer about degree into a
145 /// lower bound on it (gh #683). When it is set, an empty
146 /// [`quadratic`](Quad2::quadratic) map is no longer evidence that the
147 /// body is affine, and a consumer asking a *degree* question — as
148 /// opposed to a storage or evaluation question — must say "not
149 /// established" instead of "affine". The two questions used to share
150 /// one predicate, which is how a genuinely quadratic row came to be
151 /// reported as proved affine and had its Jacobian frozen for a whole
152 /// solve.
153 ///
154 /// It was originally set by the **drop** alone, which conflated an
155 /// exact cancellation with a lossy one and refused both (gh #687):
156 /// `x − x` is degree 0 by an exact add, and a consumer that treats it
157 /// as "not established" gives up a proved degree, a matrix evaluation
158 /// and — once the classifier gates on this too — the convex route, for
159 /// a body that lost nothing. The gate is now the fold, so the exact
160 /// case keeps its fast paths and the lossy one is refused for the
161 /// reason it deserves. [`inexact`](Quad2::inexact) is the other half of
162 /// that answer.
163 ///
164 /// It is deliberately one bit for the whole form rather than a record
165 /// per monomial: the consumer's question is about the form, and Q3's
166 /// point was that a `Quad2` allocates nothing per term (gh #588, Q3).
167 /// Two bits, now, and still nothing per term — the cost of a per-key
168 /// provenance record is what keeps this a *conservative* answer: an
169 /// inexact fold on one monomial makes a cancellation on an unrelated
170 /// one count as lossy.
171 ///
172 /// A *linear* or constant term that goes missing sets it too, even
173 /// though losing one cannot understate an affine body's degree by
174 /// itself: it can once the form is **multiplied**, because
175 /// [`mul`](Quad2::mul)'s degree guard reads the same maps —
176 /// `(2⁵³ + 1 − 2⁵³)·x · x²` is degree 3 and folds to an apparent
177 /// degree 0. Distinguishing the two would take a third flag to buy
178 /// back reach that the corpus does not contain.
179 pub fn lost_terms(&self) -> bool {
180 self.lost_terms
181 }
182
183 /// Whether any arithmetic behind this form's coefficients **rounded**.
184 ///
185 /// Clear means every stored coefficient is exact, which is what makes a
186 /// coefficient of zero a proof of absence rather than a lower bound —
187 /// see [`lost_terms`](Quad2::lost_terms), which is this bit read at the
188 /// moment a term is dropped. Public because it is the thing a test
189 /// about the sharpened gate has to be able to see; no consumer routes
190 /// on it.
191 pub fn inexact(&self) -> bool {
192 self.inexact
193 }
194
195 pub(crate) fn of_constant(c: f64) -> Self {
196 Quad2 {
197 // `-0.0` and `0.0` alike mean "no constant term"; normalizing
198 // here keeps `Neg(Const(0.0))` from reporting `-0.0`.
199 constant: if c != 0.0 { c } else { 0.0 },
200 ..Quad2::default()
201 }
202 }
203
204 pub(crate) fn of_var(i: usize) -> Self {
205 let mut q = Quad2::default();
206 q.linear.insert(i, 1.0);
207 q
208 }
209
210 /// Total degree: 0, 1, or 2.
211 pub(crate) fn degree(&self) -> usize {
212 if !self.quadratic.is_empty() {
213 2
214 } else if !self.linear.is_empty() {
215 1
216 } else {
217 0
218 }
219 }
220
221 /// The value, when this form has no variables in it.
222 pub(crate) fn as_constant(&self) -> Option<f64> {
223 (self.degree() == 0).then_some(self.constant)
224 }
225
226 /// Number of stored (nonzero) variable terms.
227 fn width(&self) -> usize {
228 self.linear.len() + self.quadratic.len()
229 }
230
231 /// `a + b`.
232 ///
233 /// Two things here are not incidental. **Only the entries the smaller
234 /// side contributes are re-checked for zero** — the predecessor
235 /// re-scanned the whole accumulated map on every merge, which is `O(k²)`
236 /// down a `k`-term `Add` chain, and an `o0` chain is how a long sum
237 /// reaches this code. And **the smaller side is merged into the larger**,
238 /// whichever operand it is, so the same chain costs `O(k log k)` leaning
239 /// either way rather than only when it leans left. Choosing the direction
240 /// is free of arithmetic consequence: IEEE addition is commutative bit
241 /// for bit, `-0.0` included.
242 pub(crate) fn add(a: Quad2, b: Quad2) -> Quad2 {
243 let (mut acc, small) = if a.width() >= b.width() {
244 (a, b)
245 } else {
246 (b, a)
247 };
248 // Whether anything either side already knew it had rounded. Read
249 // once, before any merge, so the verdict below does not depend on
250 // the order the maps happen to iterate in.
251 let carried_inexact = acc.inexact || small.inexact;
252 let (mut dropped, mut inexact) = (false, false);
253 if small.constant != 0.0 {
254 let was = acc.constant;
255 acc.constant += small.constant;
256 // Same rule as a coefficient, for the same reason: the degree-0
257 // term is exempt from the "no stored zeros" invariant but not
258 // from arithmetic, and `mul` reads `constant == 0.0` as
259 // *annihilates*. A constant that cancelled *inexactly* is not
260 // one that was proven absent (gh #683); one that cancelled
261 // exactly is (gh #687).
262 inexact |= !add_is_exact(was, small.constant, acc.constant);
263 dropped |= was != 0.0 && acc.constant == 0.0;
264 }
265 for (i, c) in &small.linear {
266 let m = merge(&mut acc.linear, *i, *c);
267 dropped |= m.dropped;
268 inexact |= m.inexact;
269 }
270 for (k, c) in &small.quadratic {
271 let m = merge(&mut acc.quadratic, *k, *c);
272 dropped |= m.dropped;
273 inexact |= m.inexact;
274 }
275 // A term is only *lost* if something was dropped and some fold on
276 // the way here could not represent itself. Neither half is enough
277 // alone: an exact cancellation drops a term that was really zero,
278 // and an inexact fold that keeps its coefficient has hidden
279 // nothing — yet.
280 acc.lost_terms |= small.lost_terms || (dropped && (carried_inexact || inexact));
281 acc.inexact = carried_inexact || inexact;
282 acc
283 }
284
285 pub(crate) fn neg(mut self) -> Quad2 {
286 self.constant = -self.constant;
287 for c in self.linear.values_mut() {
288 *c = -*c;
289 }
290 for c in self.quadratic.values_mut() {
291 *c = -*c;
292 }
293 self
294 }
295
296 /// `self · s`, for a scalar `s`.
297 ///
298 /// The `prune` is not decoration, and neither is the flag it sets. A
299 /// nonzero coefficient times a nonzero `s` can still land on zero by
300 /// underflow — `1e-300 x0²` divided by `1e300` is reachable from a real
301 /// `.nl` body, through [`Op::Div`]'s reciprocal — so this is the
302 /// gh #683 shape again, arrived at by scaling rather than by summing.
303 /// Leaving the flushed entry stored would make an arithmetically
304 /// constant form report degree 2 to [`Quad2::degree`] and be refused as
305 /// degree 3 the moment anything multiplied it; dropping it without
306 /// recording the loss would make a genuinely degree-2 body look affine
307 /// to [`Quad2::lost_terms`]'s consumers. `prune` does the first and
308 /// this does the second. (`neg` needs neither: negation cannot reach
309 /// zero from a coefficient that was not already there, and it never
310 /// rounds.)
311 ///
312 /// A coefficient here can only reach zero by **underflow** — both
313 /// factors are nonzero — so unlike a cancelling sum there is no exact
314 /// case to spare (gh #687): `10⁻³⁰⁰ · 10⁻³⁰⁰` is a term the form can no
315 /// longer see, whatever the flags said before.
316 pub(crate) fn scale(mut self, s: f64) -> Quad2 {
317 if s == 0.0 {
318 // An exact zero annihilates every term, so this is a proof of
319 // degree 0 rather than a loss — nothing is being discarded that
320 // survived the multiplication, and `x · 0` is exact. The flags
321 // still ride along: the one route here is division by an
322 // infinity, and `0 · NaN` is not zero.
323 return Quad2 {
324 lost_terms: self.lost_terms,
325 inexact: self.inexact,
326 ..Quad2::default()
327 };
328 }
329 let was = self.constant;
330 self.constant *= s;
331 if was != 0.0 {
332 self.lost_terms |= self.constant == 0.0;
333 self.inexact |= !mul_is_exact(was, s, self.constant);
334 }
335 let mut inexact = false;
336 for c in self.linear.values_mut() {
337 let was = *c;
338 *c *= s;
339 inexact |= !mul_is_exact(was, s, *c);
340 }
341 for c in self.quadratic.values_mut() {
342 let was = *c;
343 *c *= s;
344 inexact |= !mul_is_exact(was, s, *c);
345 }
346 self.inexact |= inexact;
347 // A scale can underflow a live coefficient to zero, which is a lost
348 // term and used to be stored as a zero.
349 self.lost_terms |= self.prune();
350 self
351 }
352
353 /// Restore the "no stored zeros" invariant after an operation wrote
354 /// coefficients in bulk. Returns whether anything was dropped; what
355 /// that *means* is the caller's to say, because it depends on how the
356 /// coefficients got there (see [`scale`](Quad2::scale) and
357 /// [`mul`](Quad2::mul)).
358 fn prune(&mut self) -> bool {
359 let before = self.width();
360 self.linear.retain(|_, c| is_live(*c));
361 self.quadratic.retain(|_, c| is_live(*c));
362 self.width() != before
363 }
364
365 /// `self / d`, for a constant `d` the caller has already checked is
366 /// nonzero.
367 ///
368 /// Scales by the **reciprocal** (not `c / d`) so the arithmetic matches
369 /// what the recursive predecessor produced bit for bit; what is new is
370 /// that `fl(1/d)` is the reciprocal only when `d` is a power of two,
371 /// and `fma(r, d, −1) == 0` says so exactly. Recording it matters
372 /// because a coefficient no real arithmetic produced can still cancel
373 /// *exactly* against another one later — and that cancellation would
374 /// otherwise be read as a proof of absence (gh #687).
375 pub(crate) fn div_by_constant(self, d: f64) -> Quad2 {
376 let r = 1.0 / d;
377 let exact = r.is_normal() && d.is_normal() && r.mul_add(d, -1.0) == 0.0;
378 let mut out = self.scale(r);
379 out.inexact |= !exact;
380 out
381 }
382
383 /// Take on another form's [`lost_terms`](Quad2::lost_terms) and
384 /// [`inexact`](Quad2::inexact) history.
385 ///
386 /// The operand a `Div` or a `Pow` reads as a *constant* is a [`Quad2`]
387 /// like any other, and its arithmetic is part of this form's: dividing
388 /// by `fl(10²⁰⁰ · 10²⁰⁰)` is dividing by an infinity, and dividing by a
389 /// constant that swallowed a variable term is worse than that. The
390 /// value is read out through [`as_constant`](Quad2::as_constant), which
391 /// keeps neither fact, so the caller hands the form itself over here.
392 pub(crate) fn absorb_flags(&mut self, other: &Quad2) {
393 self.lost_terms |= other.lost_terms;
394 self.inexact |= other.inexact;
395 }
396
397 /// `self · other`, or `None` when the product would exceed total
398 /// degree 2 — past that the recognizer gives up and the caller routes
399 /// to the general NLP path.
400 pub(crate) fn mul(&self, other: &Quad2) -> Option<Quad2> {
401 if self.degree() + other.degree() > 2 {
402 return None;
403 }
404 let mut out = Quad2::default();
405 // Either operand's missing term is missing from the product too,
406 // and so is either operand's rounding.
407 let mut lost = self.lost_terms || other.lost_terms;
408 let carried_inexact = self.inexact || other.inexact;
409 let mut inexact = false;
410 // A product of two live coefficients that lands on zero underflowed
411 // — `(10⁻²⁰⁰·x)·(10⁻²⁰⁰·x)` is one monomial whose coefficient is not
412 // representable (gh #683) — and unlike a cancelling *sum* there is
413 // no exact case to spare, so `lost` is set on the spot rather than
414 // left to the cancellation rule below.
415 let product = |a: f64, b: f64, lost: &mut bool, inexact: &mut bool| -> f64 {
416 let t = a * b;
417 *lost |= !is_live(t);
418 *inexact |= !mul_is_exact(a, b, t);
419 t
420 };
421 if self.constant != 0.0 && other.constant != 0.0 {
422 // `10⁻²⁰⁰ · 10⁻²⁰⁰` is not zero, and the branches below read a
423 // zero constant as an annihilating one — which would take a
424 // degree-2 product down to nothing without a trace (gh #683).
425 out.constant = product(self.constant, other.constant, &mut lost, &mut inexact);
426 }
427 let mut dropped = false;
428 // constant × (linear, quadratic), both ways round.
429 for (a, b) in [(self, other), (other, self)] {
430 if a.constant == 0.0 {
431 continue;
432 }
433 for (i, c) in &b.linear {
434 let t = product(a.constant, *c, &mut lost, &mut inexact);
435 let m = accumulate(&mut out.linear, *i, t);
436 dropped |= m.dropped;
437 inexact |= m.inexact;
438 }
439 for (k, c) in &b.quadratic {
440 let t = product(a.constant, *c, &mut lost, &mut inexact);
441 let m = accumulate(&mut out.quadratic, *k, t);
442 dropped |= m.dropped;
443 inexact |= m.inexact;
444 }
445 }
446 // linear × linear. The degree guard above means at most one of the
447 // two operands carries quadratic terms, so this runs only when
448 // neither does and no ordering question arises.
449 for (i, a) in &self.linear {
450 for (j, b) in &other.linear {
451 let key = (*i.min(j), *i.max(j));
452 let t = product(*a, *b, &mut lost, &mut inexact);
453 let m = accumulate(&mut out.quadratic, key, t);
454 dropped |= m.dropped;
455 inexact |= m.inexact;
456 }
457 }
458 // What `prune` drops here that `accumulate` did not already see is a
459 // key whose *first* contribution was a zero — an underflowed product,
460 // and `lost` is set for it above. So the cancellation rule is the
461 // same one `add` applies: a drop counts as a loss only when some fold
462 // behind it rounded.
463 dropped |= out.prune();
464 out.lost_terms = lost || (dropped && (carried_inexact || inexact));
465 out.inexact = carried_inexact || inexact;
466 Some(out)
467 }
468}
469
470/// What folding one coefficient into a map did — the two facts
471/// [`Quad2::lost_terms`] is decided from.
472#[derive(Clone, Copy)]
473struct Merged {
474 /// The key is no longer stored: the fold reached exactly zero, or
475 /// `NaN`.
476 dropped: bool,
477 /// The fold **rounded**: what is stored (or what cancelled) is not what
478 /// exact arithmetic on the same two numbers would have left.
479 inexact: bool,
480}
481
482/// Add `c` to `map[key]`, keeping the "no stored zeros" invariant. Returns
483/// what that did, as [`Merged`].
484///
485/// Only the touched key can have become zero, which is what keeps a merge
486/// proportional to what it merged rather than to what it merged *into*.
487fn merge<K: Ord>(map: &mut BTreeMap<K, f64>, key: K, c: f64) -> Merged {
488 use std::collections::btree_map::Entry;
489 match map.entry(key) {
490 Entry::Occupied(mut e) => {
491 let a = *e.get();
492 let v = a + c;
493 let inexact = !add_is_exact(a, c, v);
494 if is_live(v) {
495 e.insert(v);
496 Merged {
497 dropped: false,
498 inexact,
499 }
500 } else {
501 e.remove();
502 Merged {
503 dropped: true,
504 inexact,
505 }
506 }
507 }
508 // Nothing was added to anything: a live `c` is stored as it stands,
509 // and a dead one is only reachable from a caller that already knows
510 // what its own zero means (`mul`'s underflowed products; `add` only
511 // ever passes stored — hence live — coefficients).
512 Entry::Vacant(e) => {
513 if is_live(c) {
514 e.insert(c);
515 Merged {
516 dropped: false,
517 inexact: false,
518 }
519 } else {
520 Merged {
521 dropped: true,
522 inexact: c.is_nan(),
523 }
524 }
525 }
526 }
527}
528
529/// Accumulate `t` into `map[key]`, the way [`Quad2::mul`] builds a product
530/// up term by term, and report what the fold did.
531///
532/// Unlike [`merge`] this leaves a zero stored — `mul` prunes once at the
533/// end — because a key it lands on twice must see the first contribution as
534/// a plain `0.0 + t`, bit for bit what the `+=` this replaced produced.
535/// Only a key that was already nonzero can be said to have *dropped*
536/// anything.
537fn accumulate<K: Ord>(map: &mut BTreeMap<K, f64>, key: K, t: f64) -> Merged {
538 let slot = map.entry(key).or_insert(0.0);
539 let was = *slot;
540 *slot = was + t;
541 Merged {
542 dropped: was != 0.0 && !is_live(*slot),
543 inexact: !add_is_exact(was, t, *slot),
544 }
545}
546
547/// The exactness predicates this recognizer's fold is decided from live in
548/// `pounce-common` since gh #673, because the `.nl` pipeline grew a second
549/// fold — `Σ 2wₖbₖbₖᵀ` in `QuadraticStructure::push_factored_form` — in a
550/// crate that cannot name this one. See [`pounce_common::exact`] for what
551/// they answer and why it is an exact question rather than a tolerance.
552use pounce_common::exact::{add_is_exact, is_live, mul_is_exact};
553
554/// One entry on the recognizer's explicit work stack.
555enum Step<'a> {
556 /// Lower this subexpression onto the value stack.
557 Visit(&'a Expr),
558 /// Combine values already on the value stack.
559 Apply(Op),
560}
561
562/// A pending combination, popped once its operands have been lowered.
563enum Op {
564 Neg,
565 Add,
566 Sub,
567 Mul,
568 Div,
569 Pow,
570 /// n-ary sum over the top `n` values.
571 Sum(usize),
572 /// Not an operation: the value now on top of the stack is the lowering
573 /// of the `Cse` body with this address, so record it before moving on.
574 CacheCse(*const Expr),
575}
576
577/// Lower an [`Expr`] to a [`Quad2`], or `None` if it contains anything the
578/// recognizer cannot prove is a degree-≤2 polynomial (transcendental ops,
579/// division by a non-constant, `Pow` with an exponent ∉ {0, 1, 2}, products
580/// of degree > 2, external calls, comparisons, `if-then-else`, `min`/`max`,
581/// …). `None` ⇒ treat as general nonlinear.
582///
583/// `Cse` nodes are inlined: a reference is mathematically its body, and
584/// every reference is an independent occurrence. A body is nevertheless
585/// lowered **once**, and its [`Quad2`] reused at every later reference
586/// (keyed on `Arc` identity). That is what makes the walk `Θ(nodes)` on a
587/// shared DAG instead of `Θ(2^depth)`; Q4 had to refuse a re-referenced
588/// body outright to bound the cost, and this is the memoization that
589/// refusal was waiting for (gh #588, Q5). It is bitwise neutral — the
590/// lowering of a body is a function of the body, so the reused value is
591/// the one a second lowering would have produced, bit for bit.
592///
593/// The walk is iterative. See the module docs for why that is a
594/// correctness property and not a style choice.
595pub fn recognize_expr(e: &Expr) -> Option<Quad2> {
596 let mut work: Vec<Step<'_>> = vec![Step::Visit(e)];
597 let mut vals: Vec<Quad2> = Vec::new();
598 // Lowered `Cse` bodies, by address. Only successes land here: a body
599 // that fails aborts the whole walk, so there is no negative result to
600 // remember.
601 let mut cse: std::collections::HashMap<*const Expr, Quad2> = std::collections::HashMap::new();
602
603 while let Some(step) = work.pop() {
604 match step {
605 Step::Visit(e) => match e {
606 Expr::Const(c) => vals.push(Quad2::of_constant(*c)),
607 Expr::Var(i) => vals.push(Quad2::of_var(*i)),
608 Expr::Cse(body) => {
609 let key = std::sync::Arc::as_ptr(body);
610 match cse.get(&key) {
611 Some(q) => vals.push(q.clone()),
612 None => {
613 work.push(Step::Apply(Op::CacheCse(key)));
614 work.push(Step::Visit(body));
615 }
616 }
617 }
618 Expr::Sum(items) => {
619 work.push(Step::Apply(Op::Sum(items.len())));
620 // Pushed forward, so they pop back to front and land on
621 // the value stack with item 0 on top — which puts item
622 // 0 at the *end* of the region `Op::Sum` drains, and
623 // item n-1 at its start. The fold there walks that
624 // region in reverse so the sum still accumulates front
625 // to back, the order the recursive version summed in.
626 for it in items {
627 work.push(Step::Visit(it));
628 }
629 }
630 Expr::Unary(UnaryOp::Neg, a) => {
631 work.push(Step::Apply(Op::Neg));
632 work.push(Step::Visit(a));
633 }
634 // Every other unary op is transcendental.
635 Expr::Unary(..) => return None,
636 Expr::Binary(op, a, b) => {
637 let op = match op {
638 BinOp::Add => Op::Add,
639 BinOp::Sub => Op::Sub,
640 BinOp::Mul => Op::Mul,
641 BinOp::Div => Op::Div,
642 BinOp::Pow => Op::Pow,
643 // atan2 and any other binary opcode.
644 _ => return None,
645 };
646 work.push(Step::Apply(op));
647 // `b` under `a`: `a` pops first and is lowered first.
648 work.push(Step::Visit(b));
649 work.push(Step::Visit(a));
650 }
651 // External calls are opaque; comparisons, logicals,
652 // conditionals and n-ary min/max are the control-flow `.nl`
653 // opcodes. None is provably polynomial ⇒ route to NLP.
654 _ => return None,
655 },
656 Step::Apply(Op::CacheCse(key)) => {
657 // The body's value is already on the stack and stays there;
658 // this only records it for the next reference.
659 cse.insert(key, vals.last()?.clone());
660 }
661 Step::Apply(op) => {
662 let combined = match op {
663 Op::CacheCse(_) => unreachable!("handled above"),
664 Op::Sum(n) => {
665 // The items are the top `n` values, item 0 on top,
666 // so item 0 is the *last* thing `drain` yields and
667 // item n-1 the first. Reverse it: floating-point
668 // addition is not associative, and summing back to
669 // front would disagree with the recursive
670 // predecessor — and with the AD tape — by an ulp
671 // whenever two items share a monomial key.
672 let at = vals.len().checked_sub(n)?;
673 let mut acc = Quad2::default();
674 for p in vals.drain(at..).rev() {
675 acc = Quad2::add(acc, p);
676 }
677 acc
678 }
679 Op::Neg => vals.pop()?.neg(),
680 Op::Add => {
681 let (a, b) = pop2(&mut vals)?;
682 Quad2::add(a, b)
683 }
684 Op::Sub => {
685 let (a, b) = pop2(&mut vals)?;
686 Quad2::add(a, b.neg())
687 }
688 Op::Mul => {
689 let (a, b) = pop2(&mut vals)?;
690 a.mul(&b)?
691 }
692 Op::Div => {
693 // Division is polynomial only by a nonzero constant,
694 // and scales by the reciprocal (not `c / d`) so the
695 // arithmetic matches what the recursive predecessor
696 // produced bit for bit.
697 let (a, b) = pop2(&mut vals)?;
698 let d = b.as_constant()?;
699 if d == 0.0 {
700 return None;
701 }
702 let mut out = a.div_by_constant(d);
703 // The divisor's own history comes along: `as_constant`
704 // reads a number out of a form and leaves the flags
705 // behind.
706 out.absorb_flags(&b);
707 out
708 }
709 Op::Pow => {
710 // Polynomial only for constant exponents in {0, 1, 2}.
711 let (a, b) = pop2(&mut vals)?;
712 let exp = b.as_constant()?;
713 let mut out = if exp == 0.0 {
714 Quad2::of_constant(1.0)
715 } else if exp == 1.0 {
716 a
717 } else if exp == 2.0 {
718 a.mul(&a)?
719 } else {
720 return None;
721 };
722 // Same as `Div`: the exponent is read out of a form,
723 // and one that lost a term is not the exponent it
724 // looks like.
725 out.absorb_flags(&b);
726 out
727 }
728 };
729 vals.push(combined);
730 }
731 }
732 }
733
734 debug_assert_eq!(vals.len(), 1, "one value per lowered expression");
735 vals.pop()
736}
737
738/// Pop a binary operator's two operands, left first.
739fn pop2(vals: &mut Vec<Quad2>) -> Option<(Quad2, Quad2)> {
740 let b = vals.pop()?;
741 let a = vals.pop()?;
742 Some((a, b))
743}
744
745/// Attempt to read an expression as a polynomial of total degree ≤ 2 and
746/// return its Hessian (constant, since the form is quadratic). `None` if
747/// the expression is not provably quadratic ⇒ treat as general nonlinear.
748pub fn analyze_quadratic(e: &Expr) -> Option<QuadHessian> {
749 analyze_quadratic_full(e).map(|(h, _, _)| h)
750}
751
752/// Like [`analyze_quadratic`] but also returns the degree-1 (linear)
753/// coefficients *and* the degree-0 (constant) term of the form:
754/// `(Hessian, [(var, coef), …], constant)`.
755///
756/// AMPL folds the linear part of a nonlinear term into the objective's
757/// nonlinear expression tree (the `−6·x₀` of `(x₀−3)²`, say) rather than
758/// the linear section. Callers building the QP objective vector `c` must
759/// add these in, exactly as the NLP path's `eval_f` sums the linear
760/// section *and* the nonlinear tree — otherwise the linear shift is
761/// silently dropped and the convex solve minimizes the wrong objective.
762///
763/// The **constant** is returned for the same reason: AMPL/Pyomo also fold
764/// the objective's degree-0 term into the nonlinear tree (the `+9` of
765/// `(x₀−3)²`), where it does *not* land in `NlProblem::obj_constant`. It
766/// is irrelevant to the minimizer but is part of the *reported objective
767/// value*; dropping it makes the convex solve report an objective off by
768/// that constant versus the NLP path.
769pub fn analyze_quadratic_full(e: &Expr) -> Option<QuadForm> {
770 Some(quad_form_readout(&recognize_expr(e)?))
771}
772
773/// The `(Hessian, linear, constant)` read-out of an already-recognized
774/// form — the second half of [`analyze_quadratic_full`], split out so a
775/// caller holding a [`Quad2`] the *parser* produced (gh #588, Q5) reaches
776/// the identical numbers by the identical route. There is exactly one
777/// conversion in the crate; a second one is a second thing to keep in step.
778pub fn quad_form_readout(q: &Quad2) -> QuadForm {
779 // ∂²(c·xᵢxⱼ)/∂xᵢ∂xⱼ = c for i≠j; ∂²(c·xᵢ²)/∂xᵢ² = 2c.
780 let mut h: QuadHessian = q
781 .quadratic
782 .iter()
783 .map(|(&(i, j), c)| ((i, j), if i == j { 2.0 * c } else { *c }))
784 .collect();
785 // Drop explicit zeros so `is_empty()` means "linear".
786 h.retain(|_, v| v.abs() > 0.0);
787 let lin: Vec<(usize, f64)> = q.linear.iter().map(|(i, c)| (*i, *c)).collect();
788 // `0.0 +` normalizes `-0.0`, which is how this form spells "absent".
789 (h, lin, 0.0 + q.constant)
790}
791
792/// Is this expression **already** a flat sum of monomials — that is, would
793/// reading it as `½xᵀHx + aᵀx + c` reproduce the same additions the writer
794/// wrote, rather than algebraically expanding something it did not?
795///
796/// This is the gate on the constant-structure evaluator (gh #588, Q4), and
797/// it exists because expanding a quadratic is not an accuracy-neutral
798/// rewrite. `(xᵢ − xⱼ)²` evaluated as written squares a small residual;
799/// evaluated as `xᵢ² − 2xᵢxⱼ + xⱼ²` it cancels two large numbers, and on
800/// `airport.nl` — 84 coordinates around 10³, every row a squared distance —
801/// that difference is enough to take an adaptive-μ solve from stopping at a
802/// tiny step in 16 iterations to grinding out the 300-iteration cap at the
803/// same objective. Which is precisely the gh #544 failure mode: the right
804/// answer, slowly.
805///
806/// So the rule is *exactness*, not magnitude: a form is admitted only when
807/// the recognizer's read-out does no algebra the tape would not have done.
808/// A sum of monomials qualifies — which is exactly how AMPL emits the
809/// `qcqp*` family (`o54` over `o2 n0.5 o2 o2 n<c> v<i> v<j>`), so the target
810/// of the phase is unaffected — and a `Pow` or a `Mul` over a non-atomic
811/// operand does not.
812///
813/// What this predicate is **not** is a verdict on whether a body can be
814/// evaluated from constant structure at all. It is a verdict on one
815/// *representation*. Q4 shipped as though the two were the same, and the
816/// consequence was that a model written in factored form kept its tape and
817/// gained nothing — 41 of `airport.nl`'s 42 rows, and every sum of squared
818/// residuals ever written. gh #673 is that gap, and
819/// [`recognize_factored_quadratic`] closes it by keeping the writer's own
820/// grouping instead of expanding it: same constant structure, no algebra
821/// the tape did not do. A body this predicate refuses is offered to that
822/// one before it falls back to a tape.
823///
824/// ## A re-referenced `Cse` body is *skipped*, not refused
825///
826/// This walk and [`recognize_expr`] behind it both inline a `Cse` body at
827/// **every** reference, so a DAG whose bodies are shared cost `Θ(2^depth)`
828/// rather than `Θ(nodes)` before either was memoized. `nl_reader`'s
829/// `shared_dag_walks_are_memoized_not_exponential` builds exactly that shape
830/// — 30 levels, each a `Cse` referenced twice — and taking it through this
831/// gate in Q4 measured **0.00 s → 172 s** on a model the tape path loads
832/// instantly. At depth 40 it would not have returned at all.
833///
834/// Q4 bounded that by ending the walk in `false` at the second reference,
835/// which cost reach to buy termination and was explicitly left for Q5 to
836/// replace. It is replaced here: a body reached a second time in the same
837/// mode is **skipped**, which is sound for exactly the reason
838/// `nl_reader::validate_expr` documents for the same trick — this walk
839/// aborts on the first violation, so reaching a body a second time proves
840/// the first visit found none. The verdict is unchanged for every DAG that
841/// is a tree, and a shared DAG that *is* an expanded quadratic is now
842/// admitted instead of refused. The mode is part of the key: a body is
843/// legal on the sum spine if it is a sum of monomials, and legal inside a
844/// monomial only if it is itself a monomial, so the two answers are cached
845/// apart.
846///
847/// Iterative for the same reason [`recognize_expr`] is.
848pub fn is_expanded_quadratic(e: &Expr) -> bool {
849 // The sum spine: `Add`/`Sub`/`Neg`/`Sum` may nest freely, and every
850 // leaf of that spine must be a monomial.
851 let mut seen: BTreeSet<(*const Expr, bool)> = BTreeSet::new();
852 let mut spine: Vec<&Expr> = vec![e];
853 while let Some(e) = spine.pop() {
854 match e {
855 Expr::Sum(items) => spine.extend(items.iter()),
856 Expr::Binary(BinOp::Add | BinOp::Sub, a, b) => {
857 spine.push(a);
858 spine.push(b);
859 }
860 Expr::Unary(UnaryOp::Neg, a) => spine.push(a),
861 Expr::Cse(body) => {
862 if seen.insert((std::sync::Arc::as_ptr(body), false)) {
863 spine.push(body);
864 }
865 }
866 other => {
867 if !is_monomial(other, &mut seen) {
868 return false;
869 }
870 }
871 }
872 }
873 true
874}
875
876/// Is this expression a single monomial — the leaf shape
877/// [`is_expanded_quadratic`] admits on its sum spine?
878///
879/// Public because the parse-time recognizer has to answer the same
880/// question about a `V`-segment body it has already parsed, and must
881/// answer it with *this* code rather than a second copy of the rule.
882pub fn is_monomial_expr(e: &Expr) -> bool {
883 let mut seen: BTreeSet<(*const Expr, bool)> = BTreeSet::new();
884 is_monomial(e, &mut seen)
885}
886
887/// A single product term: constants and variables multiplied together, with
888/// no addition anywhere inside it. `xᵢxⱼ`, `0.5·c·xᵢ·xⱼ`, `xᵢ²` and `xᵢ/3`
889/// all qualify; `(xᵢ − xⱼ)²` and `(xᵢ + 1)·xⱼ` do not.
890///
891/// Degree is not checked here — [`recognize_expr`] already refused anything
892/// past 2 by the time this runs, and duplicating the rule would only give
893/// the two a way to disagree.
894fn is_monomial(e: &Expr, seen: &mut BTreeSet<(*const Expr, bool)>) -> bool {
895 let mut work: Vec<&Expr> = vec![e];
896 while let Some(e) = work.pop() {
897 match e {
898 Expr::Const(_) | Expr::Var(_) => {}
899 // Shared with the spine walk, and skipped on revisit for the
900 // same reason — see [`is_expanded_quadratic`]. The `true` in the
901 // key is "seen in monomial mode": a body cleared on the spine
902 // has not been cleared here.
903 Expr::Cse(body) => {
904 if seen.insert((std::sync::Arc::as_ptr(body), true)) {
905 work.push(body);
906 }
907 }
908 Expr::Unary(UnaryOp::Neg, a) => work.push(a),
909 Expr::Binary(BinOp::Mul | BinOp::Div, a, b) => {
910 work.push(a);
911 work.push(b);
912 }
913 // `x^2` is one monomial; `(x - y)^2` is an expansion. The
914 // exponent itself may be any constant expression — the
915 // recognizer has already restricted it to {0, 1, 2}.
916 Expr::Binary(BinOp::Pow, a, b) => {
917 if !matches!(
918 a.as_ref(),
919 Expr::Const(_) | Expr::Var(_) | Expr::Unary(UnaryOp::Neg, _)
920 ) {
921 return false;
922 }
923 work.push(a);
924 work.push(b);
925 }
926 _ => return false,
927 }
928 }
929 true
930}
931
932/// True if the expression is the literal constant zero the `.nl` reader
933/// uses for "no nonlinear part".
934pub fn is_trivially_zero(e: &Expr) -> bool {
935 matches!(e, Expr::Const(c) if *c == 0.0)
936}
937
938// ---------------------------------------------------------------------
939// Factored forms — the writer's own grouping, kept
940// ---------------------------------------------------------------------
941
942/// One `w·(bᵀx + d)²` term of a [`FactoredQuadratic`].
943///
944/// `coefs` is `b`, ascending by variable index and free of stored zeros —
945/// the same convention [`Quad2::linear`] keeps, because that is where it
946/// comes from. It may be **empty**, which is a square of a constant: the
947/// term is then `w·d²`, contributes nothing to the gradient or the
948/// Hessian, and is kept rather than folded so that the value is still
949/// computed the way it was written.
950#[derive(Debug, Clone, PartialEq)]
951pub struct SquaredAffine {
952 /// The constant the square is multiplied by, sign of the sum spine
953 /// already folded in.
954 pub weight: f64,
955 /// `b`, ascending by variable index.
956 pub coefs: Vec<(usize, f64)>,
957 /// `d`.
958 pub constant: f64,
959}
960
961/// A degree-2 form kept as `Σ wₖ(bₖᵀx + dₖ)² + aᵀx + c` — the shape the
962/// `.nl` writer wrote, rather than its expansion about the origin.
963///
964/// This is the representation gh #673 asks for and
965/// [`is_expanded_quadratic`] exists to refuse the alternative to. Reading
966/// `(x − 500000)²` back as `x² − 10⁶x + 2.5·10¹¹` cancels five digits;
967/// reading it back as one squared residual repeats exactly the
968/// multiplication the tape performs, so admitting it costs no accuracy at
969/// all.
970///
971/// The degree-≤1 leftovers (`linear`, `constant`) are the monomials the
972/// writer folded into the same tree — the `+ 3y` of `(x − 1)² + 3y`. They
973/// are summed as coefficients rather than kept term by term, which is the
974/// same reassociation the expanded path already makes and is gated the
975/// same way (see [`Quad2::lost_terms`]).
976#[derive(Debug, Clone, PartialEq)]
977pub struct FactoredQuadratic {
978 /// The squared affine terms, in the order the spine walk met them.
979 pub squares: Vec<SquaredAffine>,
980 /// `a`, ascending by variable index.
981 pub linear: Vec<(usize, f64)>,
982 /// `c`.
983 pub constant: f64,
984}
985
986/// Recognize a degree-2 body as a sum of **squared affine forms** plus
987/// degree-≤1 leftovers, keeping the squares factored (gh #673).
988///
989/// This is the second half of the gate [`is_expanded_quadratic`] opens.
990/// That predicate admits a body whose read-out repeats the additions the
991/// writer wrote; a factored body fails it, and used to keep its AD tape for
992/// good reason — expanding it to the origin cancels. What this returns is
993/// the *third* option: a read-out that is not an expansion either, because
994/// it stores the writer's own linear forms and squares them at evaluation
995/// time exactly as the tape does.
996///
997/// ## What is admitted
998///
999/// The sum spine (`Add`/`Sub`/`Neg`/`Sum`) may nest freely. Every leaf of
1000/// it must be one of:
1001///
1002/// * a **square**: a product of constants with exactly one `Pow(base, 2)`,
1003/// where `base` is itself an [`is_expanded_quadratic`] body of degree
1004/// ≤ 1 — `(xᵢ − xⱼ)²`, `0.5·(2x − y + 3)²`, `x²`;
1005/// * a **degree-2 monomial on the diagonal**, `c·xᵢ²`, which is the same
1006/// term wearing a different opcode and is stored as `c·(xᵢ)²`;
1007/// * a **degree-≤1 monomial**, which folds into `a` and `c`.
1008///
1009/// Anything else — a cross monomial `c·xᵢxⱼ`, a product of two *different*
1010/// affine forms, a transcendental — refuses the whole body, which then
1011/// keeps its tape exactly as it did before this function existed. The
1012/// refusal is deliberate rather than incidental: a mixed form would need
1013/// both an expanded and a factored quadratic part stored side by side, and
1014/// nothing in the corpus asks for one.
1015///
1016/// At least one square with a variable in it is required. Without that the
1017/// body is affine or already expanded, and [`is_expanded_quadratic`] is
1018/// the path that serves it.
1019///
1020/// ## Exactness
1021///
1022/// Each admitted square evaluates as `w·(d + Σbᵢxᵢ)²`. Against the tape
1023/// that is a reassociation of an affine sum and a folding of the writer's
1024/// own constants into `w` — the same latitude [`Quad2::add`] and
1025/// [`Quad2::scale`] already take on the expanded path — and *not* an
1026/// algebraic expansion. The `2.4e-5` disagreement gh #673 records for
1027/// `(x − 500000)²` comes from the expansion; there is none here. Measured
1028/// end to end: with the outer sum left naive, `airport.nl`'s 42 rows
1029/// evaluate bit-identically to the 42 tapes they replace.
1030///
1031/// The *terms* are the tape's; the sum over them is not, because
1032/// `QuadraticStructure::value` compensates it (gh #702). That makes the
1033/// read-out slightly better than the tape rather than equal to it, which is
1034/// gh #702's deliberate trade and the one line the fixture sweep moves.
1035///
1036/// The degree-≤1 leftovers are held to the same [`Quad2::lost_terms`] gate
1037/// the expanded path uses (gh #685), so a leftover that cancelled inexactly
1038/// refuses the body rather than dropping a term out of it.
1039///
1040/// ## A `Cse` on the sum spine refuses the body
1041///
1042/// [`is_expanded_quadratic`] may *skip* a re-referenced body because it
1043/// answers yes/no and the first visit already decided it. This one
1044/// accumulates terms, so skipping would silently drop them and visiting
1045/// every reference is `Θ(2^depth)` on a shared DAG — the shape
1046/// `nl_reader`'s `shared_dag_walks_are_memoized_not_exponential` builds.
1047/// So a `Cse` reached on the spine ends the walk, and the body keeps its
1048/// tape, which is what `ConHybrid` is for. Inside a square's base or a
1049/// monomial there is no such problem: [`recognize_expr`] memoizes on `Arc`
1050/// identity and the two predicates answer yes/no.
1051///
1052/// Iterative, for the reason the module docs give.
1053pub fn recognize_factored_quadratic(e: &Expr) -> Option<FactoredQuadratic> {
1054 let mut seen: BTreeSet<(*const Expr, bool)> = BTreeSet::new();
1055 // The spine, each entry carrying the sign the enclosing `Sub`/`Neg`
1056 // chain gives it. `±1` exactly, so folding it into a weight or negating
1057 // a form is not arithmetic.
1058 let mut spine: Vec<(&Expr, f64)> = vec![(e, 1.0)];
1059 let mut squares: Vec<SquaredAffine> = Vec::new();
1060 let mut rest = Quad2::default();
1061
1062 while let Some((e, sign)) = spine.pop() {
1063 match e {
1064 // Pushed back to front so the leaves *pop* in source order.
1065 // The degree-≤1 leftovers are folded into one `Quad2` by
1066 // floating-point addition, and source order is the order a
1067 // reader can reason about — not, to be exact about what this
1068 // buys, the same association `recognize_expr` uses: that folds
1069 // the *tree*, and `Quad2::add` additionally swaps its operands
1070 // by map width. So the two can still differ in the last ulp on
1071 // a leaning spine. What keeps that from mattering is
1072 // `lost_terms`, which refuses any body where the difference
1073 // could be a dropped term rather than a rounded one.
1074 Expr::Sum(items) => spine.extend(items.iter().rev().map(|it| (it, sign))),
1075 Expr::Binary(BinOp::Add, a, b) => {
1076 spine.push((b, sign));
1077 spine.push((a, sign));
1078 }
1079 Expr::Binary(BinOp::Sub, a, b) => {
1080 spine.push((b, -sign));
1081 spine.push((a, sign));
1082 }
1083 Expr::Unary(UnaryOp::Neg, a) => spine.push((a, -sign)),
1084 // See the docs: accumulating terms is what makes skipping
1085 // unsound here, and visiting is what makes it exponential.
1086 Expr::Cse(_) => return None,
1087 leaf => {
1088 if let Some((weight, base)) = peel_square(leaf) {
1089 squares.push(admit_square(sign * weight, base)?);
1090 continue;
1091 }
1092 if !is_monomial(leaf, &mut seen) {
1093 return None;
1094 }
1095 let q = recognize_expr(leaf)?;
1096 match diagonal_square(&q) {
1097 Some((i, c)) => squares.push(SquaredAffine {
1098 weight: sign * c,
1099 coefs: vec![(i, 1.0)],
1100 constant: 0.0,
1101 }),
1102 // A cross term `c·xᵢxⱼ` lands here and refuses: it is
1103 // not a square, and storing it expanded alongside the
1104 // squares is the mixed representation this refuses.
1105 None if !q.quadratic().is_empty() => return None,
1106 None => rest = Quad2::add(rest, if sign < 0.0 { q.neg() } else { q }),
1107 }
1108 }
1109 }
1110 }
1111
1112 // Nothing factored to keep ⇒ this is not our case. `is_expanded_quadratic`
1113 // serves an expanded body and an affine one needs no form at all.
1114 if !squares.iter().any(|t| !t.coefs.is_empty()) {
1115 return None;
1116 }
1117 // The leftovers are held to the expanded path's gate, for its reasons
1118 // (gh #685): a term that went missing inexactly is a term the form no
1119 // longer evaluates.
1120 if !rest.quadratic().is_empty() || rest.lost_terms() {
1121 return None;
1122 }
1123 Some(FactoredQuadratic {
1124 squares,
1125 linear: rest.linear().iter().map(|(&i, &c)| (i, c)).collect(),
1126 // `0.0 +` normalizes `-0.0`, which is how a `Quad2` spells "absent".
1127 constant: 0.0 + rest.constant(),
1128 })
1129}
1130
1131/// `c·xᵢ²` read off a recognized monomial: the one degree-2 shape that is
1132/// a square without being written as one. `None` for anything else,
1133/// including a cross term and anything carrying a linear or constant part.
1134///
1135/// The coefficient is the **polynomial** one — `Quad2::quadratic` is not
1136/// the Hessian — so `c·xᵢ²` is stored with weight `c` and not `2c`.
1137fn diagonal_square(q: &Quad2) -> Option<(usize, f64)> {
1138 if q.lost_terms() || !q.linear().is_empty() || q.constant() != 0.0 {
1139 return None;
1140 }
1141 match q.quadratic().iter().next() {
1142 Some((&(i, j), &c)) if i == j && q.quadratic().len() == 1 => Some((i, c)),
1143 _ => None,
1144 }
1145}
1146
1147/// Turn `weight · base²` into a stored term, or refuse the body.
1148///
1149/// The base has to clear both of the expanded path's gates — the shape gate
1150/// [`is_expanded_quadratic`] (so reading its coefficients back repeats the
1151/// writer's own additions) and the [`Quad2::lost_terms`] gate (so none of
1152/// them went missing) — and be degree ≤ 1, which is what makes the product
1153/// a square rather than a quartic.
1154fn admit_square(weight: f64, base: &Expr) -> Option<SquaredAffine> {
1155 if !is_expanded_quadratic(base) {
1156 return None;
1157 }
1158 let q = recognize_expr(base)?;
1159 if !q.quadratic().is_empty() || q.lost_terms() {
1160 return None;
1161 }
1162 Some(SquaredAffine {
1163 weight,
1164 coefs: q.linear().iter().map(|(&i, &c)| (i, c)).collect(),
1165 constant: 0.0 + q.constant(),
1166 })
1167}
1168
1169/// Split a leaf into `(constant weight, squared base)`, or `None` when it
1170/// is not a constant multiple of exactly one square.
1171///
1172/// The product tree may nest `Mul`, `Div` and `Neg` freely; every leaf of
1173/// it must be a constant except for the single `Pow(base, 2)`, which may
1174/// not sit under a division. Folding several constants into one `weight`
1175/// reassociates the writer's own constants and nothing else.
1176///
1177/// The order is worth stating precisely, because it is not source order:
1178/// this pops an explicit stack, so `a·(b·(c·s))` folds as `((1·c)·b)·a`.
1179/// One multiply's worth of rounding either way, on constants the writer
1180/// wrote next to each other — the same latitude `Quad2::scale` takes on the
1181/// expanded arm — but a reader reconstructing a last-ulp difference by hand
1182/// needs the real order, not the written one.
1183fn peel_square(e: &Expr) -> Option<(f64, &Expr)> {
1184 let mut work: Vec<(&Expr, bool)> = vec![(e, false)];
1185 let mut weight = 1.0f64;
1186 let mut base: Option<&Expr> = None;
1187 while let Some((e, recip)) = work.pop() {
1188 match e {
1189 Expr::Const(c) => weight = if recip { weight / c } else { weight * c },
1190 Expr::Unary(UnaryOp::Neg, a) => {
1191 weight = -weight;
1192 work.push((a, recip));
1193 }
1194 Expr::Binary(BinOp::Mul, a, b) => {
1195 work.push((a, recip));
1196 work.push((b, recip));
1197 }
1198 Expr::Binary(BinOp::Div, a, b) => {
1199 work.push((a, recip));
1200 work.push((b, !recip));
1201 }
1202 Expr::Binary(BinOp::Pow, a, b) if matches!(b.as_ref(), Expr::Const(c) if *c == 2.0) => {
1203 // A square under a division is `1/(…)²`, which is not one.
1204 // A second square makes the leaf degree 4.
1205 if recip || base.is_some() {
1206 return None;
1207 }
1208 base = Some(a);
1209 }
1210 _ => return None,
1211 }
1212 }
1213 // A weight that is not finite (`x²/0`) is not a form anyone should
1214 // evaluate from stored coefficients, and one that is exactly zero has
1215 // annihilated a term the tape still computes.
1216 let base = base?;
1217 (weight.is_finite() && weight != 0.0).then_some((weight, base))
1218}
1219
1220#[cfg(test)]
1221mod tests {
1222 use super::*;
1223
1224 fn sq(i: usize) -> Expr {
1225 Expr::Binary(
1226 BinOp::Pow,
1227 Box::new(Expr::Var(i)),
1228 Box::new(Expr::Const(2.0)),
1229 )
1230 }
1231
1232 #[test]
1233 fn quadratic_diagonal() {
1234 // (x0 - 1)^2 => x0^2 - 2 x0 + 1
1235 let e = Expr::Binary(
1236 BinOp::Pow,
1237 Box::new(Expr::Binary(
1238 BinOp::Sub,
1239 Box::new(Expr::Var(0)),
1240 Box::new(Expr::Const(1.0)),
1241 )),
1242 Box::new(Expr::Const(2.0)),
1243 );
1244 let (h, lin, c) = analyze_quadratic_full(&e).expect("degree-2 polynomial");
1245 assert_eq!(h.get(&(0, 0)), Some(&2.0));
1246 assert_eq!(lin, vec![(0, -2.0)]);
1247 assert_eq!(c, 1.0);
1248 }
1249
1250 #[test]
1251 fn cross_term_hessian() {
1252 // x0 · x1 => H[0,1] = 1
1253 let e = Expr::Binary(BinOp::Mul, Box::new(Expr::Var(0)), Box::new(Expr::Var(1)));
1254 let h = analyze_quadratic(&e).expect("degree-2");
1255 assert_eq!(h.get(&(0, 1)), Some(&1.0));
1256 }
1257
1258 #[test]
1259 fn rejects_transcendental_and_cubic() {
1260 assert!(analyze_quadratic(&Expr::Unary(UnaryOp::Sin, Box::new(Expr::Var(0)))).is_none());
1261 let cubic = Expr::Binary(
1262 BinOp::Pow,
1263 Box::new(Expr::Var(0)),
1264 Box::new(Expr::Const(3.0)),
1265 );
1266 assert!(analyze_quadratic(&cubic).is_none());
1267 // x0² · x1 — degree 3 by multiplication rather than by exponent.
1268 let deg3 = Expr::Binary(BinOp::Mul, Box::new(sq(0)), Box::new(Expr::Var(1)));
1269 assert!(analyze_quadratic(°3).is_none());
1270 }
1271
1272 #[test]
1273 fn division_by_a_constant_scales_by_the_reciprocal() {
1274 // x0² / 3 — the coefficient must be `2 · (1/3)`, which is what the
1275 // Hessian of `x0²/3` came out as before this module existed, and is
1276 // *not* bitwise `2/3`.
1277 let e = Expr::Binary(BinOp::Div, Box::new(sq(0)), Box::new(Expr::Const(3.0)));
1278 let h = analyze_quadratic(&e).expect("degree-2");
1279 assert_eq!(h.get(&(0, 0)), Some(&(2.0 * (1.0 / 3.0))));
1280 // Division by a variable is not polynomial.
1281 let e = Expr::Binary(BinOp::Div, Box::new(sq(0)), Box::new(Expr::Var(1)));
1282 assert!(analyze_quadratic(&e).is_none());
1283 }
1284
1285 #[test]
1286 fn scaling_a_coefficient_to_zero_drops_it_like_cancellation_does() {
1287 // 1e-300·x0² divided by 1e300. Neither the coefficient nor the
1288 // divisor is zero, so the whole-form `s == 0.0` shortcut does not
1289 // fire — the product underflows instead, and the entry has to go
1290 // the same way a cancelled one does.
1291 let tiny = Expr::Binary(BinOp::Mul, Box::new(Expr::Const(1e-300)), Box::new(sq(0)));
1292 let flushed = Expr::Binary(BinOp::Div, Box::new(tiny), Box::new(Expr::Const(1e300)));
1293 let h = analyze_quadratic(&flushed).expect("degree-2 at worst");
1294 assert!(
1295 h.is_empty(),
1296 "underflowed coefficient was kept as a structural nonzero: {h:?}"
1297 );
1298 // Storage and degree are the two halves of gh #683 and this route
1299 // reaches both: the entry is gone from the map, *and* the form
1300 // says so, so a consumer asking whether the body is affine gets
1301 // "not established" rather than "yes".
1302 let q = recognize_expr(&flushed).expect("degree-2 at worst");
1303 assert!(
1304 q.lost_terms(),
1305 "a coefficient that underflowed in `scale` was dropped silently"
1306 );
1307 // And the degree has to go with it for the *storage* question:
1308 // multiplying by another variable must still be recognized rather
1309 // than refused as degree 3.
1310 let times_x1 = Expr::Binary(BinOp::Mul, Box::new(flushed), Box::new(Expr::Var(1)));
1311 assert!(
1312 analyze_quadratic(×_x1).is_some(),
1313 "a form scaled to nothing was refused as degree 3"
1314 );
1315 }
1316
1317 #[test]
1318 fn cancellation_drops_the_term_and_the_degree_with_it() {
1319 // x0² − x0² is linear (empty Hessian), not a quadratic with a zero
1320 // coefficient — otherwise `x0²−x0²` times `x1` would be refused as
1321 // degree 3.
1322 let zero = Expr::Binary(BinOp::Sub, Box::new(sq(0)), Box::new(sq(0)));
1323 let h = analyze_quadratic(&zero).expect("degree-2 at worst");
1324 assert!(h.is_empty());
1325 let times_x1 = Expr::Binary(BinOp::Mul, Box::new(zero), Box::new(Expr::Var(1)));
1326 assert!(analyze_quadratic(×_x1).is_some());
1327 }
1328
1329 /// The gh #687 distinction, at the level it is made: an **exact**
1330 /// cancellation leaves the form complete, and an absorbed term does
1331 /// not. Both bodies drop a coefficient; only one of them lost anything
1332 /// doing it.
1333 #[test]
1334 fn an_exact_cancellation_is_not_a_lost_term() {
1335 // x0² − x0²: `fl(1) + fl(−1)` is exactly `0`, so the body really is
1336 // degree 0 and the maps say so with nothing held back.
1337 let zero = Expr::Binary(BinOp::Sub, Box::new(sq(0)), Box::new(sq(0)));
1338 let q = recognize_expr(&zero).expect("degree-2 at worst");
1339 assert!(q.quadratic().is_empty());
1340 assert!(!q.inexact(), "an exact fold reported rounding");
1341 assert!(
1342 !q.lost_terms(),
1343 "x0² − x0² was refused a fast path it is entitled to",
1344 );
1345
1346 // 2⁵³·x0² + x0² − 2⁵³·x0², front to back: the `x0²` is absorbed by
1347 // `fl(2⁵³ + 1) = 2⁵³` — an inexact add — and the exact `− 2⁵³` only
1348 // makes the loss visible.
1349 let big = (1u64 << 53) as f64;
1350 let scaled = |c: f64| Expr::Binary(BinOp::Mul, Box::new(Expr::Const(c)), Box::new(sq(0)));
1351 let absorbing = Expr::Sum(vec![scaled(big), sq(0), scaled(-big)]);
1352 let q = recognize_expr(&absorbing).expect("degree-2 at worst");
1353 assert!(q.quadratic().is_empty());
1354 assert!(q.inexact(), "the absorbing add was not seen to round");
1355 assert!(
1356 q.lost_terms(),
1357 "a body whose x0² was absorbed was reported complete",
1358 );
1359 }
1360
1361 /// The absorption is recorded where it happens, which is the point of
1362 /// the sharpened gate: `2⁵³·x0² + x0²` has lost the `x0²` already, and
1363 /// the form says so before anything cancels — while its coefficient is
1364 /// still stored and its degree is still 2.
1365 #[test]
1366 fn the_inexact_fold_is_flagged_before_anything_drops() {
1367 let big = (1u64 << 53) as f64;
1368 let scaled = |c: f64| Expr::Binary(BinOp::Mul, Box::new(Expr::Const(c)), Box::new(sq(0)));
1369 let e = Expr::Binary(BinOp::Add, Box::new(scaled(big)), Box::new(sq(0)));
1370 let q = recognize_expr(&e).expect("degree-2");
1371 assert_eq!(q.quadratic().get(&(0, 0)), Some(&big));
1372 assert!(q.inexact(), "fl(2⁵³ + 1) = 2⁵³ was called exact");
1373 // Nothing is missing from the maps *yet*, so the consumers keep
1374 // their fast paths: the degree is not in question here.
1375 assert!(!q.lost_terms());
1376 }
1377
1378 /// The sticky half of the same rule. The rounding and the drop are in
1379 /// different subexpressions — `(2⁵³·x0 + x0)` absorbs, `− 2⁵³·x0`
1380 /// cancels — and the verdict has to survive the trip between them.
1381 #[test]
1382 fn rounding_carried_from_a_subexpression_makes_a_later_drop_a_loss() {
1383 let big = (1u64 << 53) as f64;
1384 let scaled =
1385 |c: f64| Expr::Binary(BinOp::Mul, Box::new(Expr::Const(c)), Box::new(Expr::Var(0)));
1386 let absorbed = Expr::Binary(BinOp::Add, Box::new(scaled(big)), Box::new(Expr::Var(0)));
1387 let e = Expr::Binary(BinOp::Sub, Box::new(absorbed), Box::new(scaled(big)));
1388 let q = recognize_expr(&e).expect("degree-1 at worst");
1389 assert!(q.linear().is_empty());
1390 assert!(
1391 q.lost_terms(),
1392 "the x0 absorbed two additions ago is no less missing now",
1393 );
1394 }
1395
1396 /// Cancellation *inside a product*: `(x0 + x1)·(x0 − x1)` accumulates
1397 /// `+x0x1` and `−x0x1` onto one key and they cancel exactly. The body
1398 /// really has no cross term, so the form must not be demoted for
1399 /// noticing.
1400 #[test]
1401 fn an_exact_cancellation_in_a_product_is_not_a_lost_term() {
1402 let add = |a: Expr, b: Expr| Expr::Binary(BinOp::Add, Box::new(a), Box::new(b));
1403 let sub = |a: Expr, b: Expr| Expr::Binary(BinOp::Sub, Box::new(a), Box::new(b));
1404 let e = Expr::Binary(
1405 BinOp::Mul,
1406 Box::new(add(Expr::Var(0), Expr::Var(1))),
1407 Box::new(sub(Expr::Var(0), Expr::Var(1))),
1408 );
1409 let q = recognize_expr(&e).expect("degree-2");
1410 assert_eq!(q.quadratic().get(&(0, 0)), Some(&1.0));
1411 assert_eq!(q.quadratic().get(&(1, 1)), Some(&-1.0));
1412 assert_eq!(q.quadratic().get(&(0, 1)), None);
1413 assert!(!q.lost_terms(), "x0² − x1² was reported incomplete");
1414 }
1415
1416 /// An underflowing **multiply** has no exact case to spare: both
1417 /// factors are nonzero and the product is a monomial the form can no
1418 /// longer see. `(10⁻²⁰⁰·x0)·(10⁻²⁰⁰·x0)` stays refused, which is gh
1419 /// #683's second reproduction and gh #687's second acceptance case.
1420 #[test]
1421 fn an_underflowing_product_is_still_a_lost_term() {
1422 let tiny = |i: usize| {
1423 Expr::Binary(
1424 BinOp::Mul,
1425 Box::new(Expr::Const(1e-200)),
1426 Box::new(Expr::Var(i)),
1427 )
1428 };
1429 let e = Expr::Binary(BinOp::Mul, Box::new(tiny(0)), Box::new(tiny(0)));
1430 let q = recognize_expr(&e).expect("degree-2 at worst");
1431 assert!(q.quadratic().is_empty());
1432 assert!(
1433 q.lost_terms(),
1434 "a coefficient that underflowed on the multiply was reported absent",
1435 );
1436 }
1437
1438 /// A constant folded away exactly is degree 0 for real — and `mul`
1439 /// reads a zero constant as annihilating, so this is the one drop whose
1440 /// verdict a *product* depends on.
1441 #[test]
1442 fn an_exactly_cancelled_constant_is_not_a_lost_term() {
1443 let e = Expr::Binary(
1444 BinOp::Sub,
1445 Box::new(Expr::Const(3.0)),
1446 Box::new(Expr::Const(3.0)),
1447 );
1448 let q = recognize_expr(&e).expect("degree 0");
1449 assert_eq!(q.as_constant(), Some(0.0));
1450 assert!(!q.lost_terms());
1451
1452 // The same shape one ulp off: `fl(2⁵³ + 1) − 2⁵³` is `0.0` where the
1453 // real value is `1`.
1454 let big = (1u64 << 53) as f64;
1455 let absorbed = Expr::Binary(
1456 BinOp::Add,
1457 Box::new(Expr::Const(big)),
1458 Box::new(Expr::Const(1.0)),
1459 );
1460 let e = Expr::Binary(BinOp::Sub, Box::new(absorbed), Box::new(Expr::Const(big)));
1461 let q = recognize_expr(&e).expect("degree 0");
1462 assert_eq!(q.as_constant(), Some(0.0));
1463 assert!(q.lost_terms(), "an absorbed constant was reported absent");
1464 }
1465
1466 /// Rounding a coefficient is recorded even when nothing is dropped,
1467 /// because that is what a later exact cancellation would be hiding:
1468 /// `x0²/3` cannot be represented, and the form has to remember it.
1469 #[test]
1470 fn a_rounded_scale_is_inexact_but_loses_nothing() {
1471 let e = Expr::Binary(BinOp::Div, Box::new(sq(0)), Box::new(Expr::Const(3.0)));
1472 let q = recognize_expr(&e).expect("degree-2");
1473 assert_eq!(q.quadratic().get(&(0, 0)), Some(&(1.0 / 3.0)));
1474 assert!(q.inexact(), "1 · (1/3) was called exact");
1475 assert!(
1476 !q.lost_terms(),
1477 "a rounded coefficient is not a missing one"
1478 );
1479
1480 // Halving is exact, so the same shape by a power of two is not.
1481 let e = Expr::Binary(BinOp::Div, Box::new(sq(0)), Box::new(Expr::Const(2.0)));
1482 let q = recognize_expr(&e).expect("degree-2");
1483 assert!(!q.inexact(), "x0²/2 rounds nothing");
1484 }
1485
1486 /// The two-sum and the two-product, against the cases the recognizer
1487 /// actually meets. Exactness is not a tolerance, so these are `==`.
1488 #[test]
1489 fn the_exactness_tests_agree_with_the_arithmetic() {
1490 let big = (1u64 << 53) as f64;
1491 assert!(add_is_exact(big, -big, big + -big));
1492 assert!(add_is_exact(1.0, 2.0, 3.0));
1493 assert!(!add_is_exact(big, 1.0, big + 1.0), "2⁵³ + 1 loses the 1");
1494 assert!(!add_is_exact(0.1, 0.2, 0.1 + 0.2));
1495 assert!(!add_is_exact(f64::INFINITY, 1.0, f64::INFINITY));
1496 assert!(!add_is_exact(f64::NAN, 1.0, f64::NAN + 1.0));
1497
1498 assert!(mul_is_exact(3.0, 1.0, 3.0), "the ±1 shortcut");
1499 assert!(mul_is_exact(3.0, 0.5, 1.5));
1500 assert!(mul_is_exact(0.1, 4.0, 0.4));
1501 assert!(!mul_is_exact(3.0, 1.0 / 3.0, 3.0 * (1.0 / 3.0)));
1502 assert!(!mul_is_exact(1e-200, 1e-200, 1e-200 * 1e-200), "underflow");
1503 assert!(!mul_is_exact(1e300, 1e300, 1e300 * 1e300), "overflow");
1504 }
1505
1506 #[test]
1507 fn cse_bodies_are_inlined_at_every_reference() {
1508 // c = x0; c · c is x0².
1509 let body = std::sync::Arc::new(Expr::Var(0));
1510 let e = Expr::Binary(
1511 BinOp::Mul,
1512 Box::new(Expr::Cse(body.clone())),
1513 Box::new(Expr::Cse(body)),
1514 );
1515 let h = analyze_quadratic(&e).expect("degree-2");
1516 assert_eq!(h.get(&(0, 0)), Some(&2.0));
1517 }
1518
1519 #[test]
1520 fn a_wide_nary_sum_does_not_recurse() {
1521 // Σ xᵢ² over 5000 terms as one `o54` node. Distinct keys, so this
1522 // says nothing about *order* — see the test below for that.
1523 const N: usize = 5000;
1524 let e = Expr::Sum((0..N).map(sq).collect());
1525 let h = analyze_quadratic(&e).expect("sum of squares is a QP");
1526 assert_eq!(h.len(), N);
1527 assert_eq!(h.get(&(N - 1, N - 1)), Some(&2.0));
1528 }
1529
1530 /// Summation order is only observable when two items land on the *same*
1531 /// monomial key with magnitudes far enough apart that the addition
1532 /// rounds. `1e16·x0² + x0² + x0²` folds to `1e16` front to back (each
1533 /// `+1` falls below the ulp) and to `1e16 + 2` back to front. The tape
1534 /// — the reference every non-recognized row is still evaluated with —
1535 /// sums front to back, so that is the answer required here.
1536 ///
1537 /// The 5000-square test above is named for this property but cannot see
1538 /// it: 5000 distinct keys never re-add anything.
1539 #[test]
1540 fn a_repeated_monomial_in_an_nary_sum_accumulates_front_to_back() {
1541 let scaled =
1542 |c: f64, i: usize| Expr::Binary(BinOp::Mul, Box::new(Expr::Const(c)), Box::new(sq(i)));
1543 let e = Expr::Sum(vec![scaled(1.0e16, 0), scaled(1.0, 0), scaled(1.0, 0)]);
1544 let h = analyze_quadratic(&e).expect("sum of squares is a QP");
1545 let got = h[&(0, 0)];
1546 assert_eq!(
1547 got.to_bits(),
1548 (2.0 * 1.0e16_f64).to_bits(),
1549 "expected the front-to-back fold, got {got:e}"
1550 );
1551 assert_ne!(got.to_bits(), (2.0 * (1.0e16_f64 + 2.0)).to_bits());
1552 }
1553
1554 /// The gate that decides whether a form may be evaluated from its
1555 /// coefficients. The two directions cost different things, so both are
1556 /// pinned: admitting a factored form loses digits (and, on
1557 /// `airport.nl`, 284 iterations), refusing an expanded one loses the
1558 /// phase's whole point on the family it was built for.
1559 #[test]
1560 fn only_already_expanded_forms_are_admitted() {
1561 // AMPL's `qcqp*` emission: `o54` over `0.5·((c·xᵢ)·xⱼ)`.
1562 let monomial = |c: f64, i: usize, j: usize| {
1563 Expr::Binary(
1564 BinOp::Mul,
1565 Box::new(Expr::Const(0.5)),
1566 Box::new(Expr::Binary(
1567 BinOp::Mul,
1568 Box::new(Expr::Binary(
1569 BinOp::Mul,
1570 Box::new(Expr::Const(c)),
1571 Box::new(Expr::Var(i)),
1572 )),
1573 Box::new(Expr::Var(j)),
1574 )),
1575 )
1576 };
1577 let row = Expr::Sum(vec![monomial(2.0, 0, 0), monomial(3.0, 0, 1)]);
1578 assert!(is_expanded_quadratic(&row));
1579 // `xᵢ²` is a single monomial, not an expansion.
1580 assert!(is_expanded_quadratic(&Expr::Sum(vec![sq(0), sq(1)])));
1581 // A left-deep `Add` chain is still a sum.
1582 let chain = Expr::Binary(BinOp::Add, Box::new(sq(0)), Box::new(sq(1)));
1583 assert!(is_expanded_quadratic(&chain));
1584 // Division by a constant, and a negated term.
1585 assert!(is_expanded_quadratic(&Expr::Binary(
1586 BinOp::Div,
1587 Box::new(sq(0)),
1588 Box::new(Expr::Const(3.0)),
1589 )));
1590 assert!(is_expanded_quadratic(&Expr::Unary(
1591 UnaryOp::Neg,
1592 Box::new(monomial(1.0, 0, 1))
1593 )));
1594
1595 // `(x₀ − x₁)²` — the `airport.nl` shape. Reading it as
1596 // `x₀² − 2x₀x₁ + x₁²` cancels, so it stays on the tape.
1597 let diff = Expr::Binary(BinOp::Sub, Box::new(Expr::Var(0)), Box::new(Expr::Var(1)));
1598 let factored = Expr::Binary(BinOp::Pow, Box::new(diff), Box::new(Expr::Const(2.0)));
1599 assert!(analyze_quadratic(&factored).is_some(), "it is quadratic");
1600 assert!(
1601 !is_expanded_quadratic(&factored),
1602 "but not already expanded"
1603 );
1604
1605 // `(x₀ + 1)·x₁` — expansion by multiplication rather than by power.
1606 let product = Expr::Binary(
1607 BinOp::Mul,
1608 Box::new(Expr::Binary(
1609 BinOp::Add,
1610 Box::new(Expr::Var(0)),
1611 Box::new(Expr::Const(1.0)),
1612 )),
1613 Box::new(Expr::Var(1)),
1614 );
1615 assert!(!is_expanded_quadratic(&product));
1616
1617 // One factored term anywhere in a long expanded sum disqualifies
1618 // the whole row — the cancellation is in that term, not in the sum.
1619 let mixed = Expr::Sum(vec![monomial(1.0, 0, 0), factored]);
1620 assert!(!is_expanded_quadratic(&mixed));
1621
1622 // A `Cse` referenced once is inlined and judged on its body.
1623 let once = Expr::Binary(
1624 BinOp::Mul,
1625 Box::new(Expr::Cse(std::sync::Arc::new(Expr::Var(0)))),
1626 Box::new(Expr::Var(1)),
1627 );
1628 assert!(is_expanded_quadratic(&once));
1629 }
1630
1631 /// A `Cse` body reached twice is walked once, not `2^depth` times, and
1632 /// is **admitted** rather than refused.
1633 ///
1634 /// Q4 refused it, and refused it for cost rather than algebra — this
1635 /// walk and `recognize_expr` behind it both inline a body per
1636 /// reference, so a shared DAG was `Θ(2^depth)`. Q5's memoization is
1637 /// what makes admitting it affordable, so both halves are asserted
1638 /// here: the verdict, and the fact that the test returns at all. The
1639 /// shape is `nl_reader`'s `shared_dag_walks_are_memoized_not_exponential`
1640 /// scaled up; at depth 60 an exponential walk does not finish inside
1641 /// the lifetime of this test run, and `2^60` sums do not fit in an
1642 /// `f64` count either.
1643 #[test]
1644 fn a_shared_cse_body_is_walked_once_and_admitted() {
1645 let mut e = Expr::Var(0);
1646 for _ in 0..60 {
1647 let shared = std::sync::Arc::new(e);
1648 e = Expr::Binary(
1649 BinOp::Add,
1650 Box::new(Expr::Cse(std::sync::Arc::clone(&shared))),
1651 Box::new(Expr::Cse(shared)),
1652 );
1653 }
1654 assert!(is_expanded_quadratic(&e));
1655 // And the algebra agrees: `x + x` sixty times over is `2^60 · x`.
1656 let q = recognize_expr(&e).expect("a sum of one monomial is degree 1");
1657 assert_eq!(q.linear().get(&0).copied(), Some(2.0_f64.powi(60)));
1658 std::mem::forget(e);
1659 }
1660
1661 /// The same shape inside a *monomial*, which is a different question
1662 /// with a different answer and therefore a separately keyed memo: a
1663 /// body that is a sum of monomials is legal on the spine and illegal
1664 /// under a `*`.
1665 #[test]
1666 fn a_shared_body_is_judged_per_context_not_once_and_for_all() {
1667 let sum = std::sync::Arc::new(Expr::Binary(
1668 BinOp::Add,
1669 Box::new(Expr::Var(0)),
1670 Box::new(Expr::Var(1)),
1671 ));
1672 // On the spine, twice: fine.
1673 let spine = Expr::Binary(
1674 BinOp::Add,
1675 Box::new(Expr::Cse(std::sync::Arc::clone(&sum))),
1676 Box::new(Expr::Cse(std::sync::Arc::clone(&sum))),
1677 );
1678 assert!(is_expanded_quadratic(&spine));
1679 // The same body on the spine and then under a product: the product
1680 // is `(x0 + x1) · x1`, a factored form, and the earlier clean visit
1681 // on the spine must not clear it.
1682 let mixed = Expr::Binary(
1683 BinOp::Add,
1684 Box::new(Expr::Cse(std::sync::Arc::clone(&sum))),
1685 Box::new(Expr::Binary(
1686 BinOp::Mul,
1687 Box::new(Expr::Cse(sum)),
1688 Box::new(Expr::Var(1)),
1689 )),
1690 );
1691 assert!(!is_expanded_quadratic(&mixed));
1692 }
1693
1694 /// Deep, and on a default-sized test thread, for the same reason
1695 /// [`recognize_expr`] is iterative: the gate runs on every row of every
1696 /// model, including the ones that overflow a recursive walk.
1697 #[test]
1698 fn the_expansion_gate_does_not_overflow_the_stack() {
1699 const K: usize = 250_000;
1700 let mut e = sq(0);
1701 for i in 1..K {
1702 e = Expr::Binary(BinOp::Add, Box::new(e), Box::new(sq(i)));
1703 }
1704 assert!(is_expanded_quadratic(&e));
1705 std::mem::forget(e);
1706 }
1707
1708 /// The reason this module is iterative.
1709 ///
1710 /// A `.nl` writer that emits `o0` (binary `+`) chains for a long sum
1711 /// hands the recognizer a left-deep `Add` tree one level deep per term.
1712 /// The recursive predecessor aborted the process — a stack overflow is
1713 /// not a catchable error — somewhere under 8 000 terms on a 2 MB thread
1714 /// and under 24 000 on the CLI's 8 MB main thread. The depth below is
1715 /// far past both, and the test runs on a **default-sized test thread**
1716 /// deliberately: what a recognizer survives must not depend on who
1717 /// called it.
1718 ///
1719 /// The tree is leaked rather than dropped, because `Expr`'s derived
1720 /// `Drop` is still recursive and would overflow tearing this down.
1721 /// That is a real and separate defect (pounce#472 works around it in
1722 /// the Python bindings with a big-stack worker thread); it is not what
1723 /// this test is measuring.
1724 #[test]
1725 fn deep_add_chain_does_not_overflow_the_stack() {
1726 const K: usize = 250_000;
1727 let mut e = sq(0);
1728 for i in 1..K {
1729 e = Expr::Binary(BinOp::Add, Box::new(e), Box::new(sq(i)));
1730 }
1731 let h = analyze_quadratic(&e).expect("a sum of squares is a QP at any depth");
1732 assert_eq!(h.len(), K, "every xᵢ² contributes one diagonal entry");
1733 assert_eq!(h.get(&(K - 1, K - 1)), Some(&2.0));
1734 std::mem::forget(e);
1735 }
1736
1737 /// Same shape, right-deep — the value stack, not the work stack, is
1738 /// what grows here. Both live on the heap.
1739 #[test]
1740 fn deep_right_leaning_chain_does_not_overflow_the_stack() {
1741 const K: usize = 250_000;
1742 let mut e = sq(K - 1);
1743 for i in (0..K - 1).rev() {
1744 e = Expr::Binary(BinOp::Add, Box::new(sq(i)), Box::new(e));
1745 }
1746 let h = analyze_quadratic(&e).expect("a sum of squares is a QP at any depth");
1747 assert_eq!(h.len(), K);
1748 std::mem::forget(e);
1749 }
1750
1751 /// A non-quadratic node deep inside a deep tree must return `None`
1752 /// rather than unwind — the bail-out path drops the work and value
1753 /// stacks, and neither is recursive.
1754 #[test]
1755 fn deep_chain_with_a_transcendental_bails_without_overflowing() {
1756 const K: usize = 250_000;
1757 let mut e = Expr::Unary(UnaryOp::Sin, Box::new(Expr::Var(0)));
1758 for i in 1..K {
1759 e = Expr::Binary(BinOp::Add, Box::new(e), Box::new(sq(i)));
1760 }
1761 assert!(analyze_quadratic(&e).is_none());
1762 std::mem::forget(e);
1763 }
1764
1765 // -----------------------------------------------------------------
1766 // Factored forms (gh #673)
1767 // -----------------------------------------------------------------
1768
1769 /// `(base)^2`.
1770 fn sq_of(base: Expr) -> Expr {
1771 Expr::Binary(BinOp::Pow, Box::new(base), Box::new(Expr::Const(2.0)))
1772 }
1773
1774 fn var_minus(i: usize, c: f64) -> Expr {
1775 Expr::Binary(BinOp::Sub, Box::new(Expr::Var(i)), Box::new(Expr::Const(c)))
1776 }
1777
1778 /// The motivating case. `(x − 500000)²` is exactly what
1779 /// `feasible_x0_extreme_row.nl` writes, and expanding it is the 2.4e-5
1780 /// disagreement gh #673 is named after.
1781 #[test]
1782 fn a_shifted_square_is_kept_factored() {
1783 let e = sq_of(var_minus(0, 500_000.0));
1784 assert!(
1785 !is_expanded_quadratic(&e),
1786 "the expanded gate must refuse it"
1787 );
1788 let f = recognize_factored_quadratic(&e).expect("a square of an affine form");
1789 assert_eq!(f.squares.len(), 1);
1790 assert_eq!(f.squares[0].weight, 1.0);
1791 assert_eq!(f.squares[0].coefs, vec![(0, 1.0)]);
1792 assert_eq!(f.squares[0].constant, -500_000.0);
1793 assert!(f.linear.is_empty());
1794 assert_eq!(f.constant, 0.0);
1795 }
1796
1797 /// The accuracy claim, stated as a number rather than as prose: at
1798 /// `x = 500000 + 1e-4` the expansion loses five digits and the factored
1799 /// read-out loses nothing.
1800 #[test]
1801 fn the_factored_read_out_does_not_cancel_where_the_expansion_does() {
1802 let e = sq_of(var_minus(0, 500_000.0));
1803 let x = 500_000.0 + 1e-4;
1804 // What the tape computes: square the residual as written.
1805 let r = x - 500_000.0;
1806 let taped = r * r;
1807
1808 let (h, lin, c) = analyze_quadratic_full(&e).expect("degree 2");
1809 let expanded = 0.5 * h[&(0, 0)] * x * x + lin[0].1 * x + c;
1810 assert!(
1811 (expanded - taped).abs() / taped > 1e-6,
1812 "the expansion is supposed to cancel here, got {expanded} for {taped}"
1813 );
1814
1815 let f = recognize_factored_quadratic(&e).expect("a square");
1816 let t = &f.squares[0];
1817 let l = t.constant + t.coefs[0].1 * x;
1818 // Bit for bit, not within a tolerance.
1819 assert_eq!(t.weight * l * l, taped);
1820 }
1821
1822 /// `airport.nl`'s shape: a row that is a sum of squared coordinate
1823 /// differences, with the writer's grouping kept term by term.
1824 #[test]
1825 fn a_sum_of_squared_differences_is_admitted() {
1826 let diff = |i: usize, j: usize| {
1827 Expr::Binary(BinOp::Sub, Box::new(Expr::Var(i)), Box::new(Expr::Var(j)))
1828 };
1829 let e = Expr::Binary(
1830 BinOp::Add,
1831 Box::new(sq_of(diff(0, 1))),
1832 Box::new(sq_of(diff(2, 3))),
1833 );
1834 let f = recognize_factored_quadratic(&e).expect("two squares");
1835 assert_eq!(f.squares.len(), 2);
1836 let mut sup: Vec<Vec<(usize, f64)>> = f.squares.iter().map(|t| t.coefs.clone()).collect();
1837 sup.sort_by_key(|c| c[0].0);
1838 assert_eq!(
1839 sup,
1840 vec![vec![(0, 1.0), (1, -1.0)], vec![(2, 1.0), (3, -1.0)]]
1841 );
1842 assert!(
1843 f.squares
1844 .iter()
1845 .all(|t| t.weight == 1.0 && t.constant == 0.0)
1846 );
1847 }
1848
1849 /// The sign of the sum spine reaches the weight, and a constant factor
1850 /// folds into it.
1851 #[test]
1852 fn spine_signs_and_constant_factors_fold_into_the_weight() {
1853 // 3·(x₀ − 1)² − 0.5·(x₁ + 2)²
1854 let a = Expr::Binary(
1855 BinOp::Mul,
1856 Box::new(Expr::Const(3.0)),
1857 Box::new(sq_of(var_minus(0, 1.0))),
1858 );
1859 let b = Expr::Binary(
1860 BinOp::Mul,
1861 Box::new(Expr::Const(0.5)),
1862 Box::new(sq_of(Expr::Binary(
1863 BinOp::Add,
1864 Box::new(Expr::Var(1)),
1865 Box::new(Expr::Const(2.0)),
1866 ))),
1867 );
1868 let f = recognize_factored_quadratic(&Expr::Binary(BinOp::Sub, Box::new(a), Box::new(b)))
1869 .expect("two squares");
1870 let mut got: Vec<(f64, f64)> = f.squares.iter().map(|t| (t.weight, t.constant)).collect();
1871 got.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
1872 assert_eq!(got, vec![(-0.5, 2.0), (3.0, -1.0)]);
1873 }
1874
1875 /// A degree-≤1 leftover in the same tree is not a reason to refuse the
1876 /// row; it folds into `a`/`c` and is evaluated there.
1877 #[test]
1878 fn degree_one_leftovers_fold_into_the_linear_part() {
1879 // (x₀ − 1)² + 3·x₁ + 7
1880 let e = Expr::Sum(vec![
1881 sq_of(var_minus(0, 1.0)),
1882 Expr::Binary(
1883 BinOp::Mul,
1884 Box::new(Expr::Const(3.0)),
1885 Box::new(Expr::Var(1)),
1886 ),
1887 Expr::Const(7.0),
1888 ]);
1889 let f = recognize_factored_quadratic(&e).expect("a square plus leftovers");
1890 assert_eq!(f.squares.len(), 1);
1891 assert_eq!(f.linear, vec![(1, 3.0)]);
1892 assert_eq!(f.constant, 7.0);
1893 }
1894
1895 /// A bare `c·xᵢ²` monomial is the same term wearing a different opcode,
1896 /// and is stored as `c·(xᵢ)²` so that a row mixing the two shapes is
1897 /// still served.
1898 #[test]
1899 fn a_diagonal_monomial_is_stored_as_a_square() {
1900 let e = Expr::Binary(
1901 BinOp::Add,
1902 Box::new(sq_of(var_minus(0, 1.0))),
1903 Box::new(Expr::Binary(
1904 BinOp::Mul,
1905 Box::new(Expr::Const(4.0)),
1906 Box::new(Expr::Binary(
1907 BinOp::Mul,
1908 Box::new(Expr::Var(1)),
1909 Box::new(Expr::Var(1)),
1910 )),
1911 )),
1912 );
1913 let f = recognize_factored_quadratic(&e).expect("square + diagonal monomial");
1914 assert_eq!(f.squares.len(), 2);
1915 let diag = f
1916 .squares
1917 .iter()
1918 .find(|t| t.coefs == vec![(1, 1.0)])
1919 .unwrap();
1920 // The *polynomial* coefficient, so `4x₁²` and not `8x₁²` —
1921 // `push_factored_form` is what doubles it into a Hessian entry.
1922 assert_eq!((diag.weight, diag.constant), (4.0, 0.0));
1923 }
1924
1925 /// A cross monomial has no square to be stored as, and storing it
1926 /// expanded next to the squares is the mixed representation this
1927 /// deliberately does not have. The row keeps its tape.
1928 #[test]
1929 fn a_cross_monomial_alongside_a_square_is_refused() {
1930 let e = Expr::Binary(
1931 BinOp::Add,
1932 Box::new(sq_of(var_minus(0, 1.0))),
1933 Box::new(Expr::Binary(
1934 BinOp::Mul,
1935 Box::new(Expr::Var(1)),
1936 Box::new(Expr::Var(2)),
1937 )),
1938 );
1939 assert!(recognize_factored_quadratic(&e).is_none());
1940 }
1941
1942 /// An already-expanded body is `is_expanded_quadratic`'s to serve, and
1943 /// this must not offer a second, slower answer for it.
1944 #[test]
1945 fn an_expanded_body_is_not_claimed_here() {
1946 let e = Expr::Binary(BinOp::Add, Box::new(sq(0)), Box::new(sq(1)));
1947 assert!(is_expanded_quadratic(&e));
1948 // `x²` *is* a square, so this one is genuinely recognized — the
1949 // caller's ordering is what keeps it on the cheaper path.
1950 assert!(recognize_factored_quadratic(&e).is_some());
1951 // A body with no square at all is refused outright.
1952 let cross = Expr::Binary(BinOp::Mul, Box::new(Expr::Var(0)), Box::new(Expr::Var(1)));
1953 assert!(recognize_factored_quadratic(&cross).is_none());
1954 assert!(recognize_factored_quadratic(&Expr::Var(0)).is_none());
1955 }
1956
1957 /// The product of two *different* affine forms is degree 2 and is not a
1958 /// square. Admitting it would mean storing `l₁·l₂`, which this
1959 /// representation cannot express.
1960 #[test]
1961 fn a_product_of_two_different_affine_forms_is_refused() {
1962 let e = Expr::Binary(
1963 BinOp::Mul,
1964 Box::new(var_minus(0, 1.0)),
1965 Box::new(var_minus(1, 2.0)),
1966 );
1967 assert!(recognize_factored_quadratic(&e).is_none());
1968 }
1969
1970 /// A square of a *quadratic* is degree 4, and a transcendental is not a
1971 /// polynomial at all.
1972 #[test]
1973 fn quartics_and_transcendentals_are_refused() {
1974 assert!(recognize_factored_quadratic(&sq_of(sq(0))).is_none());
1975 let s = Expr::Unary(UnaryOp::Sin, Box::new(Expr::Var(0)));
1976 assert!(recognize_factored_quadratic(&sq_of(s)).is_none());
1977 }
1978
1979 /// A base that is not itself a flat sum of monomials is refused: the
1980 /// coefficients read out of `2·(x + 1)` are not the additions the
1981 /// writer wrote, which is the same rule `is_expanded_quadratic` states.
1982 #[test]
1983 fn a_base_the_expanded_gate_refuses_is_refused_here_too() {
1984 let inner = Expr::Binary(
1985 BinOp::Mul,
1986 Box::new(Expr::Const(2.0)),
1987 Box::new(Expr::Binary(
1988 BinOp::Add,
1989 Box::new(Expr::Var(0)),
1990 Box::new(Expr::Const(1.0)),
1991 )),
1992 );
1993 assert!(!is_expanded_quadratic(&inner));
1994 assert!(recognize_factored_quadratic(&sq_of(inner)).is_none());
1995 }
1996
1997 /// A leftover that lost a term inexactly refuses the body, for the
1998 /// reason gh #685 gives: the form would evaluate a row its own tape
1999 /// does not agree with.
2000 #[test]
2001 fn a_lost_leftover_refuses_the_body() {
2002 let big = 9_007_199_254_740_992.0f64; // 2⁵³
2003 let scaled = |c: f64, i: usize| {
2004 Expr::Binary(BinOp::Mul, Box::new(Expr::Const(c)), Box::new(Expr::Var(i)))
2005 };
2006 // (x₀ − 1)² + 2⁵³·x₁ + x₁ − 2⁵³·x₁ — the `x₁` is folded away.
2007 let e = Expr::Sum(vec![
2008 sq_of(var_minus(0, 1.0)),
2009 scaled(big, 1),
2010 Expr::Var(1),
2011 scaled(-big, 1),
2012 ]);
2013 assert!(recognize_factored_quadratic(&e).is_none());
2014 }
2015
2016 /// A `Cse` on the sum spine ends the walk: skipping a second reference
2017 /// would drop its terms and visiting every reference is exponential.
2018 #[test]
2019 fn a_shared_body_on_the_spine_is_refused() {
2020 let shared = std::sync::Arc::new(sq_of(var_minus(0, 1.0)));
2021 let e = Expr::Binary(
2022 BinOp::Add,
2023 Box::new(Expr::Cse(shared.clone())),
2024 Box::new(Expr::Cse(shared)),
2025 );
2026 assert!(recognize_factored_quadratic(&e).is_none());
2027 }
2028
2029 /// A weight that annihilates or overflows is not evaluated from stored
2030 /// coefficients: `0·(x−1)²` has dropped a term the tape still walks,
2031 /// and `(x−1)²/0` is not a number.
2032 #[test]
2033 fn degenerate_weights_are_refused() {
2034 let zero = Expr::Binary(
2035 BinOp::Mul,
2036 Box::new(Expr::Const(0.0)),
2037 Box::new(sq_of(var_minus(0, 1.0))),
2038 );
2039 assert!(recognize_factored_quadratic(&zero).is_none());
2040 let div0 = Expr::Binary(
2041 BinOp::Div,
2042 Box::new(sq_of(var_minus(0, 1.0))),
2043 Box::new(Expr::Const(0.0)),
2044 );
2045 assert!(recognize_factored_quadratic(&div0).is_none());
2046 }
2047
2048 /// Deep, on a default-sized test thread, for the reason every walk in
2049 /// this module is iterative: a least-squares model is a long `o0` chain
2050 /// of squared residuals, which is exactly the shape that aborts a
2051 /// recursive walk — and now exactly the shape this recognizer is for.
2052 ///
2053 /// Leaked rather than dropped: `Expr`'s derived `Drop` is still
2054 /// recursive. See `deep_add_chain_does_not_overflow_the_stack`.
2055 #[test]
2056 fn a_deep_chain_of_squares_does_not_overflow_the_stack() {
2057 const K: usize = 250_000;
2058 let mut e = sq_of(var_minus(0, 1.0));
2059 for i in 1..K {
2060 e = Expr::Binary(
2061 BinOp::Add,
2062 Box::new(e),
2063 Box::new(sq_of(var_minus(i, i as f64))),
2064 );
2065 }
2066 let f = recognize_factored_quadratic(&e).expect("K squares");
2067 assert_eq!(f.squares.len(), K);
2068 std::mem::forget(e);
2069 }
2070}