Skip to main content

thermite_special/specialized/generic/
elliptic.rs

1//! Elliptic integrals.
2//!
3//! The complete integrals K and E use the arithmetic-geometric mean (AGM), which is
4//! branchless and needs only one `sqrt` per iteration - far cheaper for SIMD than the
5//! Carlson duplication algorithm used for the incomplete forms (see [`carlson_rf`] etc.).
6//!
7//! The algorithms are classical (Gauss/Legendre AGM; Carlson's symmetric forms); this is
8//! an original SIMD implementation, not a port of any particular source.
9
10#![allow(clippy::extra_unused_type_parameters)]
11
12use thermite::{
13    math::{TranscendentalMathWithPolicy as _, policy::Policy},
14    prelude::*,
15    register::FloatElement,
16};
17
18use crate::specialized::SpecializedSpecialMath;
19use thermite::math::policy::DenormalBehavior;
20use thermite::{const_element, const_splat};
21
22/// Legendre integral kinds for [`ellint_impl`]'s `KIND` const parameter.
23pub const KIND_F: u8 = 1; // first kind: F(phi, k) / K(k)
24pub const KIND_E: u8 = 2; // second kind: E(phi, k) / E(k)
25pub const KIND_PI: u8 = 3; // third kind: Pi(n, phi, k)
26pub const KIND_D: u8 = 4; // D(phi, k) = (F - E) / k^2
27
28// `v.scale(const_element!(ratio <E>: 1 / 3))` is `v * const_splat!(ratio <E>: 1 / 3)` on CPU
29// and a single `OpVectorTimesScalar` on SPIR-V. `const_splat!` is used where the constant is
30// an FMA or vector operand and `scale` does not fit.
31
32/// Complete elliptic integrals of the first and second kind, `(K(k), E(k))`, evaluated
33/// together from a single AGM pass (they share the iteration).
34///
35/// ```text
36/// K(k) = pi / (2 * AGM(1, k')),     k' = sqrt(1 - k^2)
37/// E(k) = K(k) * (1 - sum_{n>=0} 2^{n-1} c_n^2)
38/// ```
39///
40/// Valid for `|k| <= 1`; `|k| > 1` gives NaN. The endpoint `|k| = 1` is pinned to
41/// `(inf, 1)` under `check_overflow`. See the note at the end of the function.
42#[inline(always)]
43pub fn agm_complete_ke<P, E, V>(k: V) -> (V, V)
44where
45    P: Policy,
46    E: FloatElement,
47    V: FloatVector<Element = E>,
48{
49    // k' = sqrt(1 - k^2); one_minus_sq is cancellation-free as |k| -> 1 (FMA or factored form).
50    let mut a = V::ONE;
51    let b0 = k.one_minus_sq().sqrt();
52    let mut b = b0;
53    let mut c = k;
54
55    // sum starts with the n = 0 term: 2^{-1} c_0^2 = k^2 / 2.
56    let mut sum = (c * c).scale(const_element!(ratio <E>: 1 / 2));
57    let mut pow2 = V::ONE; // 2^{n-1} for the first in-loop term (n = 1) is 2^0 = 1
58
59    // AGM converges quadratically, so this is a handful of iterations; the masked break
60    // stops once every lane's remaining contribution is below the rounding threshold.
61    let thresh = V::SQRT_EPSILON; // c_n^2 ~ eps once |c_n| ~ sqrt(eps)
62    let mut iter = 0;
63    loop {
64        V::_loop_hint();
65
66        let an = (a + b).scale(const_element!(ratio <E>: 1 / 2));
67        let bn = (a * b).sqrt();
68        c = (a - b).scale(const_element!(ratio <E>: 1 / 2));
69        a = an;
70        b = bn;
71
72        sum = pow2.mul_adde(c * c, sum); // sum += pow2 * c^2
73        pow2 = pow2 + pow2;
74
75        iter += 1;
76        if iter >= 24 || c.abs().cmp_le(a * thresh).all() {
77            break;
78        }
79    }
80
81    let k_int = V::FRAC_PI_2 / a;
82    let e_int = k_int.nmul_adde(sum, k_int); // k_int * (1 - sum)
83
84    if const { P::POLICY.check_overflow } {
85        // |k| = 1 makes k' = 0, and AGM(1, 0) = 0, a limit the iteration cannot reach, since
86        // `a` merely halves every pass and stops at the cap (K comes out as (pi/2) * 2^24, and
87        // E as the inf*0 of that, pi/4). K really does diverge there, but E(1) = 1 exactly, so
88        // both are pinned rather than left to the loop. |k| > 1 makes k' NaN, not zero, so the
89        // out-of-domain NaN still propagates.
90        let deg = b0.is_zero();
91        (deg.select(V::INFINITY, k_int), deg.select(V::ONE, e_int))
92    } else {
93        (k_int, e_int)
94    }
95}
96
97/// The arithmetic-geometric mean `$\mathrm{AGM}(a, b)$` of two non-negative arguments.
98///
99/// The same recurrence [`agm_complete_ke`] runs, without that function's `E` accumulator
100/// and without its `a = 1` starting pin:
101///
102/// ```math
103/// a_{n+1} = \frac{a_n + b_n}{2}, \qquad b_{n+1} = \sqrt{a_n b_n}
104/// ```
105///
106/// Both sequences converge to the common limit quadratically. From an extreme starting
107/// ratio the logarithm of that ratio roughly halves each pass until the two arguments are
108/// within a factor of a few, after which the correct digits double per pass, so the
109/// iteration cap covers the whole representable range with room to spare, and the loop
110/// stays uniform across lanes rather than data-dependent.
111///
112/// Kept beside `agm_complete_ke` on purpose: they are one recurrence, and a change to
113/// either is nearly always a change to both.
114///
115/// # Domain
116///
117/// Defined for `a, b >= 0`, and symmetric in its arguments. A negative argument makes the
118/// geometric mean's sign ambiguous after the first pass (`(a + b)/2` can be negative
119/// while `sqrt(ab)` is not), so negatives give NaN under `check_overflow` rather than a
120/// plausible wrong value. `AGM(a, 0) = 0` for every `a`, a limit the iteration approaches
121/// but cannot reach (`b` is pinned at zero while `a` merely halves), so it is pinned too,
122/// as is `AGM(inf, b) = inf`, which otherwise leaves the loop as `inf - inf`. The one
123/// pairing with no limit at all, a zero against an infinity, is NaN.
124///
125/// # Range
126///
127/// `sqrt(ab)` is formed as a single product, so two arguments both above `sqrt(MAX)`
128/// (about 1.3e154 in binary64, 1.8e19 in binary32) overflow to infinity even though the
129/// mean itself is perfectly representable. The AGM is homogeneous,
130/// `$\mathrm{AGM}(ca, cb) = c\,\mathrm{AGM}(a, b)$`, so a caller working up there should
131/// scale both arguments by a common power of two, which is exact.
132#[inline(always)]
133pub fn agm<P, E, V>(a: V, b: V) -> V
134where
135    P: Policy,
136    E: FloatElement,
137    V: FloatVector<Element = E>,
138{
139    let mut x = a;
140    let mut y = b;
141
142    // The limit sits between the two sequences, roughly at their midpoint, so `x` is off
143    // by about half the _current_ gap, linearly, not quadratically. The test is therefore
144    // on the gap going in, exactly as `agm_complete_ke` tests its `c`: a gap of
145    // `sqrt(eps)` before the pass leaves `eps/2` after it, since one pass squares it
146    // (`x' - y' = (sqrt x - sqrt y)^2 / 2`). Testing the gap coming out instead stops a
147    // whole pass early and costs half the mantissa: measured 5.8e-11 at AGM(1, sqrt 2).
148    let thresh = V::SQRT_EPSILON;
149    let mut iter = 0;
150    loop {
151        V::_loop_hint();
152
153        let gap = x - y;
154
155        let xn = (x + y).scale(const_element!(ratio <E>: 1 / 2));
156        let yn = (x * y).sqrt();
157        x = xn;
158        y = yn;
159
160        iter += 1;
161        if iter >= 24 || gap.abs().cmp_le(x * thresh).all() {
162            break;
163        }
164    }
165
166    if const { P::POLICY.check_overflow } {
167        let zero = a.is_zero() | b.is_zero();
168        let inf = a.is_infinite() | b.is_infinite();
169
170        x = zero.select(V::ZERO, x);
171        x = inf.select(V::INFINITY, x);
172
173        // Negative arguments, and the one indeterminate pairing (a zero against an
174        // infinity), where neither pin above is the limit.
175        let nan = a.cmp_lt(V::ZERO) | b.cmp_lt(V::ZERO) | (zero & inf);
176        x = nan.select(V::NAN, x);
177    }
178
179    x
180}
181
182/// True where at least two of three non-negative arguments are zero, the interior
183/// singularity shared by `R_F`, `R_D` and `R_J`, all of which diverge there while the
184/// duplication loop only walks toward the pole until it hits its iteration cap.
185///
186/// Phrased as "the second-smallest argument is zero", which is exact. The cheaper-looking
187/// product test `x*y + x*z + y*z == 0` is not: three denormal arguments underflow every
188/// product to zero and would be misread as the singular case.
189#[inline(always)]
190fn two_or_more_zero<V: FloatVector>(x: V, y: V, z: V) -> V::Mask {
191    x.min(y).max(z.min(x.max(y))).is_zero()
192}
193
194/// Compile-time per-element constants for the elliptic routines, so values that are otherwise a
195/// few runtime ops (e.g. the Carlson convergence threshold, three sequential `sqrt`s) become a
196/// constant splat. Declared for `f32`/`f64`; add more element types as needed.
197pub trait EllipticConsts {
198    /// `(3 * eps)^(1/8)`: the relative-deviation threshold at which the Carlson 7th-order Taylor
199    /// tail (`~deviation^8`) drops below rounding. Equals `sqrt(sqrt(sqrt(eps + eps + eps)))`.
200    const CARLSON_THRESH: Self;
201
202    /// `|t|` below which `carlson_rc` takes its 8-term series in `t = (y - x)/x` instead of
203    /// the closed `atan`/`ln` forms. The tail is `t^8/17`, so this is where that drops below
204    /// `eps`: `1/128` at binary32/64 (`1.5e-17`). A wider element type needs it smaller
205    /// (double-double takes `2^-14`), or every `R_J`, which calls `R_C(1, 1 + E_n)` with a
206    /// small `E_n` on every step, is capped at the series' truncation.
207    const RC_SERIES_THRESH: Self;
208}
209
210impl EllipticConsts for f32 {
211    const CARLSON_THRESH: f32 = 0.15637917816638947;
212    const RC_SERIES_THRESH: f32 = 0.0078125;
213}
214
215impl EllipticConsts for f64 {
216    const CARLSON_THRESH: f64 = 0.012674918778210762;
217    const RC_SERIES_THRESH: f64 = 0.0078125;
218}
219
220/// The Carlson convergence threshold splatted to the vector type `V`.
221#[inline(always)]
222fn carlson_thresh<V: FloatVector<Element: EllipticConsts>>() -> V {
223    V::splat(<V::Element as EllipticConsts>::CARLSON_THRESH)
224}
225
226/// Carlson symmetric integral of the first kind, `R_F(x, y, z)`, via the duplication
227/// algorithm. All Boost special cases are omitted: the duplication converges for any
228/// valid input (e.g. a single zero argument becomes positive after one step), so the
229/// loop runs uniformly across lanes and stops once every lane has converged.
230///
231/// Two or more zero arguments is the one exception. `R_F` diverges there, and the
232/// duplication only walks toward the pole, so it is pinned to infinity under
233/// `check_overflow`.
234#[inline(always)]
235pub fn carlson_rf<P, E, V>(x: V, y: V, z: V) -> V
236where
237    P: Policy,
238    E: FloatElement + EllipticConsts,
239    V: FloatVector<Element = E>,
240{
241    let quarter: V = const_splat!(ratio <E>: 1 / 4);
242    let thresh = carlson_thresh::<V>();
243
244    let mut xn = x;
245    let mut yn = y;
246    let mut zn = z;
247    let mut an = (x + y + z).scale(const_element!(ratio <E>: 1 / 3));
248    let a0 = an;
249    let mut fmn = V::ONE; // 4^-n
250    // Convergence bound. The deviation identity |An - vn| = fmn |A0 - v0| means the current max
251    // deviation is fmn * q0, so `fmn * q0 <= An * thresh` is the stop test. Fold thresh in once
252    // (q = q0/thresh) and the per-iteration test becomes just `fmn * q <= An` - no recomputing
253    // |An - vn| every step.
254    let q = (a0 - x).abs().max((a0 - y).abs()).max((a0 - z).abs()) / thresh;
255
256    let mut iter = 0;
257    loop {
258        V::_loop_hint();
259
260        let rx = xn.sqrt();
261        let ry = yn.sqrt();
262        let rz = zn.sqrt();
263        let lambda = rx.mul_adde(ry + rz, ry * rz); // rx*(ry+rz) + ry*rz; 2-deep vs serial FMA chain
264        // (v + lambda)/4 == v/4 + lambda/4; with true FMA, premultiplying lambda lets each update
265        // be a single fused `v.mul_adde(1/4, lambda/4)`. Without FMA that extra mul is wasted, so
266        // keep the plain add-then-scale form there.
267        if const { matches!(V::HAS_NATIVE_FMA, thermite::tribool::True) } {
268            let lq = lambda * quarter;
269            an = an.mul_add(quarter, lq);
270            xn = xn.mul_add(quarter, lq);
271            yn = yn.mul_add(quarter, lq);
272            zn = zn.mul_add(quarter, lq);
273        } else {
274            an = (an + lambda) * quarter;
275            xn = (xn + lambda) * quarter;
276            yn = (yn + lambda) * quarter;
277            zn = (zn + lambda) * quarter;
278        }
279        fmn *= quarter;
280        iter += 1;
281        if iter >= 30 || (fmn * q).cmp_le(an).all() {
282            break;
283        }
284    }
285
286    // Deviation identity: (An - xn)/An = 4^-n (A0 - x)/An. We use the right-hand form: the
287    // initial deviation (A0 - x) is at full precision, whereas (An - xn) is catastrophic
288    // cancellation once xn -> An (it costs many digits in the converged, e.g. x = 0, case).
289    let scale = fmn / an; // one division, shared by all deviations (vs one div each)
290    let xd = (a0 - x) * scale;
291    let yd = (a0 - y) * scale;
292    let zd = -xd - yd;
293    let e2 = xd.mul_sube(yd, zd * zd); // X*Y - Z*Z
294    let e3 = xd * yd * zd;
295
296    // 7th-order Taylor expansion (Carlson 2015). Like rdj_poly, split by total degree into three
297    // *independent* FMA chains run in parallel and summed - a ~3-deep critical path instead of the
298    // serial `1 + e3*A + e2*B` form (~5 deep, B being a 4-deep chain). The e's are tiny deviations,
299    // so the regrouped sum is as accurate as the serial form (corrections never cancel against 1).
300    let e2_2 = e2 * e2;
301    // Linear:    1 - 1/10 e2 + 1/14 e3
302    let lin = e2.mul_adde(const_splat!(ratio <E>: -1 / 10), V::ONE);
303    let lin = e3.mul_adde(const_splat!(ratio <E>: 1 / 14), lin);
304    // Quadratic: 1/24 e2^2 - 3/44 e2 e3 + 3/104 e3^2
305    let quad = (e2 * e3).mul_adde(
306        const_splat!(ratio <E>: -3 / 44),
307        e2_2.scale(const_element!(ratio <E>: 1 / 24)),
308    );
309    let quad = (e3 * e3).mul_adde(const_splat!(ratio <E>: 3 / 104), quad);
310    // Cubic:     -5/208 e2^3 + 1/16 e2^2 e3
311    let cub = (e2_2 * e3).mul_adde(
312        const_splat!(ratio <E>: 1 / 16),
313        (e2_2 * e2).scale(const_element!(ratio <E>: -5 / 208)),
314    );
315
316    let poly = lin + (quad + cub);
317    let rf = poly / an.sqrt();
318
319    if const { P::POLICY.check_overflow } {
320        two_or_more_zero(x, y, z).select(V::INFINITY, rf)
321    } else {
322        rf
323    }
324}
325
326/// Carlson symmetric integral `R_D(x, y, z) = R_J(x, y, z, z)` (degenerate third kind,
327/// used for the second-kind incomplete integral). Same duplication scheme as [`carlson_rf`],
328/// plus an accumulated sum term.
329#[inline(always)]
330pub fn carlson_rd<P, E, V>(x: V, y: V, z: V) -> V
331where
332    P: Policy,
333    E: FloatElement + EllipticConsts,
334    V: FloatVector<Element = E>,
335{
336    let quarter: V = const_splat!(ratio <E>: 1 / 4);
337    let thresh = carlson_thresh::<V>();
338
339    let mut xn = x;
340    let mut yn = y;
341    let mut zn = z;
342    let mut an = ((x + y) + (z + z + z)).scale(const_element!(ratio <E>: 1 / 5)); // (x + y + 3z) / 5; grouped for ILP
343    let a0 = an;
344    let mut sum = V::ZERO;
345    let mut fac = V::ONE; // 4^-n
346    // Convergence bound q = q0 / thresh; the loop tests `fac * q <= An` (see carlson_rf).
347    let q = (a0 - x).abs().max((a0 - y).abs()).max((a0 - z).abs()) / thresh;
348
349    let mut iter = 0;
350    loop {
351        V::_loop_hint();
352
353        let rx = xn.sqrt();
354        let ry = yn.sqrt();
355        let rz = zn.sqrt();
356        let lambda = rx.mul_adde(ry + rz, ry * rz); // rx*(ry+rz) + ry*rz; 2-deep vs serial FMA chain
357        sum += fac / (rz * (zn + lambda));
358        // (v + lambda)/4 as a fused FMA when available (see carlson_rf).
359        if const { matches!(V::HAS_NATIVE_FMA, thermite::tribool::True) } {
360            let lq = lambda * quarter;
361            an = an.mul_add(quarter, lq);
362            xn = xn.mul_add(quarter, lq);
363            yn = yn.mul_add(quarter, lq);
364            zn = zn.mul_add(quarter, lq);
365        } else {
366            an = (an + lambda) * quarter;
367            xn = (xn + lambda) * quarter;
368            yn = (yn + lambda) * quarter;
369            zn = (zn + lambda) * quarter;
370        }
371        fac *= quarter;
372        iter += 1;
373        if iter >= 30 || (fac * q).cmp_le(an).all() {
374            break;
375        }
376    }
377
378    // Reconstruct from initial deviations (see carlson_rf) to avoid An - xn cancellation.
379    let scale = fac / an; // one division, shared by the deviations
380    let xd = (a0 - x) * scale;
381    let yd = (a0 - y) * scale;
382    let zd = (xd + yd).scale(const_element!(ratio <E>: -1 / 3));
383    let xy = xd * yd;
384    let zz = zd * zd;
385    let xy3 = xy.scale(const_element!(ratio <E>: 3 / 1)); // 3 xy, shared by e3 and e4
386    let e2 = zz.mul_adde(const_splat!(int <E>: -6), xy); // xy - 6 zz
387    let e3 = zz.mul_adde(const_splat!(int <E>: -8), xy3) * zd; // (3 xy - 8 zz) zd
388    let e4 = zz.mul_adde(const_splat!(int <E>: -3), xy3) * zz; // (3 xy - 3 zz) zz = 3 (xy - zz) zz
389    let e5 = xy * (zz * zd);
390
391    let taylor = fac * rdj_poly_n::<E, V>(e2, e3, e4, e5) / (an * an.sqrt()); // fac * An^(-3/2) * poly
392    let rd = sum.mul_adde(const_splat!(int <E>: 3), taylor); // taylor + 3 * sum
393
394    if const { P::POLICY.check_overflow } {
395        // `z == 0` alone also diverges, but that one arrives on its own: the accumulated
396        // `fac / (rz * (zn + lambda))` term divides by zero and carries the infinity out.
397        two_or_more_zero(x, y, z).select(V::INFINITY, rd)
398    } else {
399        rd
400    }
401}
402
403/// Shared 5th-order Taylor tail polynomial for R_D and R_J (Carlson 2015) - they use the
404/// same form in the deviation variables E2..E5.
405#[inline(always)]
406fn rdj_poly_n<E, V>(e2: V, e3: V, e4: V, e5: V) -> V
407where
408    E: FloatElement,
409    V: FloatVector<Element = E>,
410{
411    // e2^2 feeds several higher terms; compute once. Then three *independent* FMA chains (by
412    // total degree) run in parallel and are summed - a ~7-deep critical path instead of one
413    // 12-deep serial Horner chain. The e's are tiny deviations, so the regrouped sum is as
414    // accurate as the serial form (the corrections never cancel against the leading 1).
415    let e2_2 = e2 * e2;
416    // Linear:    1 - 3/14 e2 + 1/6 e3 - 3/22 e4 + 3/26 e5
417    let lin = e2.mul_adde(const_splat!(ratio <E>: -3 / 14), V::ONE);
418    let lin = e3.mul_adde(const_splat!(ratio <E>: 1 / 6), lin);
419    let lin = e4.mul_adde(const_splat!(ratio <E>: -3 / 22), lin);
420    let lin = e5.mul_adde(const_splat!(ratio <E>: 3 / 26), lin);
421    // Quadratic: 9/88 e2^2 - 9/52 e2 e3 + 3/40 e3^2 + 3/20 e2 e4
422    let quad = (e2 * e3).mul_adde(
423        const_splat!(ratio <E>: -9 / 52),
424        e2_2.scale(const_element!(ratio <E>: 9 / 88)),
425    );
426    let quad = (e3 * e3).mul_adde(const_splat!(ratio <E>: 3 / 40), quad);
427    let quad = (e2 * e4).mul_adde(const_splat!(ratio <E>: 3 / 20), quad);
428    // Cubic+:    -1/16 e2^3 + 45/272 e2^2 e3 - 9/68 (e3 e4 + e2 e5)
429    let cub = (e2_2 * e3).mul_adde(
430        const_splat!(ratio <E>: 45 / 272),
431        (e2_2 * e2).scale(const_element!(ratio <E>: -1 / 16)),
432    );
433    let cub = (e3 * e4 + e2 * e5).mul_adde(const_splat!(ratio <E>: -9 / 68), cub);
434
435    lin + (quad + cub)
436}
437
438/// Carlson symmetric integral of the second kind, `R_G(x, y, z)`, as a combination of
439/// [`carlson_rf`] and [`carlson_rd`] (Carlson 2015):
440///
441/// ```text
442/// R_G = (z * R_F(x,y,z) - (x-z)(y-z) * R_D(x,y,z) / 3 + sqrt(x*y/z)) / 2
443/// ```
444///
445/// The arguments are sorted to `hi >= mid >= lo` and substituted as `x = hi, z = mid, y = lo`,
446/// the ordering that keeps `(x-z)(y-z)` from cancelling and puts the middle value (the divisor)
447/// in `z`. That form needs `mid > 0`; two or more zero arguments divide by it, so under
448/// `check_overflow` they take the closed form `R_G(x, 0, 0) = sqrt(x)/2` instead (which also
449/// covers `R_G(0,0,0) = 0`). Unlike `R_F`/`R_D`/`R_J`, `R_G` stays finite there. Not wired into
450/// the public `ellint` surface - no Legendre form needs it; provided for direct use
451/// (e.g. `E(k) = 2 R_G(0, 1-k^2, 1)`).
452#[inline(always)]
453pub fn carlson_rg<P, E, V>(x: V, y: V, z: V) -> V
454where
455    P: Policy,
456    E: FloatElement + EllipticConsts,
457    V: FloatVector<Element = E>,
458{
459    let lo = x.min(y).min(z);
460    let hi = x.max(y).max(z);
461    let mid = (x + y + z) - (lo + hi); // grouped: (x+y+z) and (lo+hi) form in parallel
462    // R_F is fully symmetric; R_D's third argument must be the middle value.
463    let rf = carlson_rf::<P, E, V>(hi, lo, mid);
464    let rd = carlson_rd::<P, E, V>(hi, lo, mid);
465    let root = (hi * lo / mid).sqrt();
466    let prod = (hi - mid) * (lo - mid) * rd;
467    // (mid*rf + sqrt(xy/z) - (x-z)(y-z) rd / 3) / 2
468    let rg = prod
469        .mul_adde(const_splat!(ratio <E>: -1 / 3), mid.mul_adde(rf, root))
470        .scale(const_element!(ratio <E>: 1 / 2));
471
472    if const { P::POLICY.check_overflow } {
473        // Sorted, so `mid == 0` is exactly "two or more arguments are zero". There the general
474        // form is 0/0 and `rf`/`rd` are themselves infinite, but the integral is not: it
475        // collapses to sqrt(hi)/2.
476        mid.is_zero()
477            .select(hi.sqrt().scale(const_element!(ratio <E>: 1 / 2)), rg)
478    } else {
479        rg
480    }
481}
482
483/// Carlson degenerate integral `R_C(x, y) = R_F(x, y, y)`, closed form. Assumes `y > 0`
484/// (the only cases that arise inside R_J and the third-kind reductions); the `y < 0`
485/// Cauchy-principal-value branch is not handled here.
486///
487/// Writing `t = (y - x)/x`, `R_C(x, y) = S(t)/sqrt(x)` where
488/// `S(t) = atan(sqrt(t))/sqrt(t) = 1 - t/3 + t^2/5 - t^3/7 + ...` is smooth and the *same*
489/// series for either sign of `t` (`atanh` for `t < 0`). The closed `atan`/`ln` forms lose
490/// ~`sqrt(eps)` precision as `t -> 0` (e.g. `ln(1 + sqrt|t|)` with tiny `sqrt|t|`), which
491/// matters because R_J calls this with `y -> x` every iteration once `p` nears an argument.
492/// So for small `|t|` we use the series instead; the two agree to full precision at the
493/// crossover. This keeps R_C (and hence R_J near the `p == arg` degeneracy) accurate without
494/// resorting to extended precision.
495#[inline(always)]
496pub fn carlson_rc<P, E, V>(x: V, y: V) -> V
497where
498    P: Policy,
499    E: FloatElement + EllipticConsts,
500    V: SpecializedSpecialMath<E>,
501{
502    let d = y - x;
503    let absd = d.abs();
504
505    // The closed form divides by sad = sqrt(|y-x|) (both branches) and sqrt(y) (neg branch).
506    // With hardware rsqrt (f32) compute each 1/sqrt directly and recover the roots by multiply -
507    // no full sqrt or divide in the hot path. Without it (f64, where rsqrt = rcp(sqrt)) a plain
508    // a/sqrt(b) is a single divide, so sqrt + div stays optimal. `irx = 1/sqrt(x)` also scales
509    // the series (poly * irx) in both paths. The carried `scale` is 1/sad on the rsqrt path
510    // (multiply) and sad on the divide path; the combine below picks the matching op.
511    // The rsqrt arm only pays when the hardware estimate is BOTH present and permitted.
512    // Under `Preserve` the estimate is forbidden (`rsqrtps` treats a subnormal operand as
513    // zero in hardware whatever MXCSR says), so `inverse_sqrt_p` becomes an exact
514    // `1/sqrt` and this arm would spend three of them where the other spends one divide.
515    let (s, neg_arg, irx, scale) =
516        if const { V::HAS_APPROX_RSQRT && !matches!(P::POLICY.denormal_behavior, DenormalBehavior::Preserve) } {
517            let isad = absd.inverse_sqrt_p::<P>(); // 1/sqrt(|y-x|)
518            let irx = x.inverse_sqrt_p::<P>(); // 1/sqrt(x)
519            let iry = y.inverse_sqrt_p::<P>(); // 1/sqrt(y)
520            let sad = absd * isad; // sqrt(|y-x|)
521            let rx = x * irx; // sqrt(x)
522            (sad * irx, (rx + sad) * iry, irx, isad)
523        } else {
524            let sad = absd.sqrt(); // sqrt(|y-x|)
525            let rx = x.sqrt();
526            let irx = rx.approx_reciprocal_p::<P>(); // 1/sqrt(x)
527            (sad * irx, (rx + sad) / y.sqrt(), irx, sad)
528        };
529
530    // x < y: atan(s)  ;  x > y: ln((sqrt(x) + sqrt(x-y))/sqrt(y)) = atanh(s). Both scaled by 1/sad.
531    let num = d.cmp_gt(V::ZERO).select(s.atan_p::<P>(), neg_arg.ln_p::<P>());
532    // Must match the arm chosen above, since `scale` is 1/sad on one path and sad on the other.
533    let closed = if const { V::HAS_APPROX_RSQRT && !matches!(P::POLICY.denormal_behavior, DenormalBehavior::Preserve) }
534    {
535        num * scale
536    } else {
537        num / scale
538    };
539
540    // Series for small |t|, where the closed forms cancel: S(t)/sqrt(x). S is univariate in t,
541    // S(t) = 1 - t/3 + t^2/5 - ... ; evaluated leading-coefficient-first via poly_rev (Estrin + FMA).
542    let t = d / x;
543
544    let mut res = closed;
545
546    // |t| < 1/128 ~ 0.0078 at binary64: series is accurate to <1e-16 with these 8 terms, and
547    // the closed forms are already degrading there. Above it the closed forms are accurate.
548    // t == 0 (x == y) is covered by the series limit S(0) = 1. The threshold is per element
549    // type (`EllipticConsts`) because the tail is the series', not the type's.
550
551    let small = t.abs().cmp_lt(V::splat(E::RC_SERIES_THRESH));
552
553    if const { P::POLICY.avoid_branching } || small.any() {
554        let series = t.poly_rev_n_p::<P, _>(&[
555            <E as FloatElement>::ConstRatio::<-1, 15>::VALUE,
556            <E as FloatElement>::ConstRatio::<1, 13>::VALUE,
557            <E as FloatElement>::ConstRatio::<-1, 11>::VALUE,
558            <E as FloatElement>::ConstRatio::<1, 9>::VALUE,
559            <E as FloatElement>::ConstRatio::<-1, 7>::VALUE,
560            <E as FloatElement>::ConstRatio::<1, 5>::VALUE,
561            <E as FloatElement>::ConstRatio::<-1, 3>::VALUE,
562            <E as FloatElement>::ConstRatio::<1, 1>::VALUE,
563        ]) * irx;
564
565        res = small.select(series, res);
566    }
567
568    res
569}
570
571/// Carlson symmetric integral of the third kind, `R_J(x, y, z, p)`, via duplication.
572/// Handles `p < 0` (a Cauchy principal value) through Carlson's transform to a positive
573/// parameter. Each step accumulates an `R_C` term, so this is the most expensive Carlson
574/// primitive.
575///
576/// Accuracy note: when `p` coincides with one of `x, y, z` (so `(p-x)(p-y)(p-z) -> 0`), the
577/// per-step `R_C(1, b)` term has `b -> 1` every iteration. That used to lose ~7 digits, but
578/// [`carlson_rc`] now switches to its small-argument series there, so this case holds full
579/// precision without the `R_D` special-case or extended precision Boost resorts to.
580#[inline(always)]
581pub fn carlson_rj<P, E, V>(x: V, y: V, z: V, p: V) -> V
582where
583    P: Policy,
584    E: FloatElement + EllipticConsts,
585    V: SpecializedSpecialMath<E>,
586{
587    let quarter: V = const_splat!(ratio <E>: 1 / 4);
588    let thresh = carlson_thresh::<V>();
589
590    // R_J is symmetric in (x, y, z); sort so `hi` is the largest. For p < 0 the integral is
591    // a Cauchy principal value, mapped to a positive parameter p' via Carlson's transform.
592    let lo = x.min(y).min(z);
593    let hi = x.max(y).max(z);
594    let mid = (x + y + z) - (lo + hi); // grouped: (x+y+z) and (lo+hi) form in parallel
595
596    let neg = p.cmp_lt(V::ZERO);
597    let q = -p; // |p| on the p < 0 lanes
598    // p' = (hi(lo + mid + q) - lo*mid) / (hi + q)  (> 0 for sorted lo<=mid<=hi, p < 0)
599    let p_new = hi.mul_sube(lo + mid + q, lo * mid) / (hi + q);
600    let p_eff = neg.select(p_new, p); // R_J parameter: p' where p<0, else p (both > 0)
601
602    let mut xn = lo;
603    let mut yn = mid;
604    let mut zn = hi;
605    let mut pn = p_eff;
606    let mut an = ((lo + mid + hi) + (p_eff + p_eff)).scale(const_element!(ratio <E>: 1 / 5)); // (x + y + z + 2p) / 5; grouped for ILP
607    let a0 = an;
608    let mut rc_sum = V::ZERO;
609    let mut fmn = V::ONE; // 4^-n
610    // Convergence bound qb = q0 / thresh; the loop tests `fmn * qb <= An` (see carlson_rf).
611    let qb = (a0 - lo)
612        .abs()
613        .max((a0 - mid).abs())
614        .max((a0 - hi).abs().max((a0 - p_eff).abs()))
615        / thresh;
616
617    let mut iter = 0;
618    loop {
619        V::_loop_hint();
620
621        let rx = xn.sqrt();
622        let ry = yn.sqrt();
623        let rz = zn.sqrt();
624        let rp = pn.sqrt();
625        let dn = (rp + rx) * (rp + ry) * (rp + rz);
626        // b = 1 + E_n computed stably (avoids the E_n ~ -1 cancellation); R_C(1, 1+E_n).
627        let inner = ry.mul_adde(rz, rx.mul_adde(ry + rz, pn)); // pn + rx(ry+rz) + ry rz
628        // 2 * rp * inner / dn, balanced so numerator (rp*inner) and denominator (dn/2) form in
629        // parallel before the divide, shortening the critical path.
630        let b = (rp * inner) / dn.scale(const_element!(ratio <E>: 1 / 2));
631        rc_sum = (fmn / dn).mul_adde(carlson_rc::<P, E, V>(V::ONE, b), rc_sum); // += (fmn/dn) R_C
632
633        let lambda = rx.mul_adde(ry + rz, ry * rz); // rx*(ry+rz) + ry*rz; 2-deep vs serial FMA chain
634        // (v + lambda)/4 as a fused FMA when available (see carlson_rf).
635        if const { matches!(V::HAS_NATIVE_FMA, thermite::tribool::True) } {
636            let lq = lambda * quarter;
637            an = an.mul_add(quarter, lq);
638            xn = xn.mul_add(quarter, lq);
639            yn = yn.mul_add(quarter, lq);
640            zn = zn.mul_add(quarter, lq);
641            pn = pn.mul_add(quarter, lq);
642        } else {
643            an = (an + lambda) * quarter;
644            xn = (xn + lambda) * quarter;
645            yn = (yn + lambda) * quarter;
646            zn = (zn + lambda) * quarter;
647            pn = (pn + lambda) * quarter;
648        }
649        fmn *= quarter;
650        iter += 1;
651        if iter >= 30 || (fmn * qb).cmp_le(an).all() {
652            break;
653        }
654    }
655
656    // Reconstruct from initial deviations (see carlson_rf) to avoid An - xn cancellation.
657    let scale = fmn / an; // one division, shared by all three deviations
658    let xd = (a0 - lo) * scale;
659    let yd = (a0 - mid) * scale;
660    let zd = (a0 - hi) * scale;
661    let pd = (xd + yd + zd).scale(const_element!(ratio <E>: -1 / 2));
662    let xyz = xd * yd * zd;
663    let pp = pd * pd;
664    let ppd = pp * pd; // pd^3, shared by e3 and e4
665    let sym = yd.mul_adde(zd, xd * (yd + zd)); // xd*yd + xd*zd + yd*zd
666    let e2 = pp.mul_adde(const_splat!(int <E>: -3), sym); // sym - 3 pd^2
667    // e2 is the latest-arriving input (it trails the fmn/an divide through pd/pp/sym). Precompute
668    // the e2-independent parts of e3/e4 - which LLVM can't hoist itself, FP adds don't reassociate
669    // without fast-math - so each e-term is a single FMA past e2 instead of a 2-3 deep chain.
670    // e3 = xyz + 2 e2 pd + 4 pd^3
671    let pre3 = ppd.mul_adde(const_splat!(int <E>: 4), xyz); // 4 pd^3 + xyz
672    let e3 = e2.mul_adde(pd + pd, pre3); // 2 pd e2 + pre3
673    // e4 = (2 xyz + e2 pd + 3 pd^3) pd = e2 pp + (3 pd^3 + 2 xyz) pd
674    let pre4 = ppd.mul_adde(const_splat!(int <E>: 3), xyz + xyz) * pd; // (3 pd^3 + 2 xyz) pd = 3 pp^2 + 2 xyz pd
675    let e4 = e2.mul_adde(pp, pre4); // e2 pp + pre4
676    let e5 = xyz * pp;
677
678    let taylor = fmn * rdj_poly_n::<E, V>(e2, e3, e4, e5) / (an * an.sqrt());
679    let rj = rc_sum.mul_adde(const_splat!(int <E>: 6), taylor); // taylor + 6 * rc_sum
680
681    let out = if const { P::POLICY.avoid_branching } || neg.any() {
682        // Cauchy PV recombination for p < 0 (Carlson):
683        //   R_J = ((p'-z) R_J(x,y,z,p') - 3 R_F(x,y,z) + 3 sqrt(xyz/(xy+p'q)) R_C(xy+p'q, p'q)) / (z+q)
684        let rf = carlson_rf::<P, E, V>(lo, mid, hi);
685        let xy = lo * mid;
686        let xyz = xy * hi;
687        let pq = p_new * q;
688        let rc = carlson_rc::<P, E, V>(xy + pq, pq);
689        // ((p'-z) R_J - 3 R_F + 3 sqrt(xyz/(xy+p'q)) R_C) / (z+q)
690        //   = ((p'-z) R_J + 3 (sqrt(..) R_C - R_F)) / (z+q)
691        let root = (xyz / (xy + pq)).sqrt();
692        let val_neg =
693            (p_new - hi).mul_adde(rj, root.mul_sube(rc, rf).scale(const_element!(ratio <E>: 3 / 1))) / (hi + q);
694        neg.select(val_neg, rj)
695    } else {
696        rj
697    };
698
699    if const { P::POLICY.check_overflow } {
700        // Two singular sets. `p == 0` mostly arrives on its own (the per-step `R_C(1, 0)` is
701        // infinite), but not when an argument is zero too: `dn` is then zero as well and the
702        // `R_C` argument becomes 0/0. Since `mid` is already sorted, both tests are a compare.
703        (mid.is_zero() | p.is_zero()).select(V::INFINITY, out)
704    } else {
705        out
706    }
707}
708
709/// Legendre elliptic integral, const-generic over kind and completeness.
710///
711/// `KIND` is one of [`KIND_F`], [`KIND_E`], [`KIND_D`], [`KIND_PI`]. `COMPLETE` selects
712/// `phi = pi/2` (the AGM path for F/E/D, Carlson R_F+R_J for Pi); otherwise the incomplete
713/// form is evaluated via Carlson at a range-reduced amplitude. Argument is the modulus `k`;
714/// `n` is the characteristic, used only by the third kind ([`KIND_PI`]).
715///
716/// Incomplete amplitudes are reduced into `[-pi/2, pi/2]` using the quasi-period identity
717/// `I(phi + m*pi) = I(phi) + 2m * I_complete` (the integrands have period pi, and one full
718/// period equals twice the complete value). Reduction by `phi - m*pi` loses precision for
719/// very large `|phi|` (argument cancellation); a Cody-Waite split would extend the range.
720#[inline(always)]
721pub fn ellint_impl<P, E, V, const KIND: u8, const COMPLETE: bool>(phi: V, k: V, n: V) -> V
722where
723    P: Policy,
724    E: FloatElement + EllipticConsts,
725    V: SpecializedSpecialMath<E>, // : FloatVector + (via blanket) TranscendentalMathWithPolicy
726{
727    const {
728        assert!(
729            KIND == KIND_F || KIND == KIND_E || KIND == KIND_D || KIND == KIND_PI,
730            "ellint_impl: KIND must be KIND_F, KIND_E, KIND_D, or KIND_PI"
731        );
732    }
733    if const { COMPLETE } {
734        if const { KIND == KIND_PI } {
735            // Pi(n, k) = R_F(0, 1-k^2, 1) + (n/3) R_J(0, 1-k^2, 1, 1-n)  (the AGM does not
736            // cover the third kind, so the complete Pi still goes through Carlson).
737            let w = k.one_minus_sq(); // 1 - k^2
738            let rf = carlson_rf::<P, E, V>(V::ZERO, w, V::ONE);
739            let rj = carlson_rj::<P, E, V>(V::ZERO, w, V::ONE, V::ONE - n);
740            n.scale(const_element!(ratio <E>: 1 / 3)).mul_adde(rj, rf) // (n/3) rj + rf
741        } else {
742            let (kk, ee) = agm_complete_ke::<P, E, V>(k);
743            if const { KIND == KIND_F } {
744                kk
745            } else if const { KIND == KIND_E } {
746                ee
747            } else {
748                // D(k) = (K - E) / k^2
749                let d = (kk - ee) / (k * k);
750                if const { P::POLICY.check_overflow } {
751                    // K(0) == E(0), so k = 0 is 0/0 here. The limit is finite: D(k) -> pi/4.
752                    k.is_zero().select(V::FRAC_PI_4, d)
753                } else {
754                    d
755                }
756            }
757        }
758    } else {
759        // Range-reduce phi into [-pi/2, pi/2]; m counts the half-periods stripped off.
760        //
761        // A single `m*PI` deliberately, rather than a Cody-Waite split of pi. Such a
762        // split exists because a naive reduction leaves an absolute error of about
763        // |phi|*eps in the reduced angle - but here
764        // the result is `F(phi_red) + 2m*K`, which grows with |phi| in the same
765        // proportion, so that error stays at roughly one ULP of the returned value no
766        // matter how large phi gets. Measured at m = 1e9: reduction error ~7e-7 against
767        // a result of ~3.4e9 whose ULP is ~4.8e-7.
768        //
769        // Reduction precision pays off when the output does *not* grow with the input -
770        // sin and cos, whose range is fixed, are where it is worth the extra products.
771        let m = (phi * V::FRAC_1_PI).round();
772        let phi_red = m.nmul_adde(V::PI, phi);
773
774        let (s, cphi) = phi_red.sin_cos_p::<P>();
775        let c2 = cphi * cphi;
776        let k2 = k * k;
777        // w = 1 - k^2 sin^2(phi), rewritten as (1 - k^2) + k^2 cos^2(phi).
778        //
779        // The direct form subtracts two quantities that both approach 1 at the domain
780        // boundary |k sin(phi)| = 1, and `k * sin(phi)` rounds to exactly 1 before the
781        // subtraction ever runs, so w collapses to 0 while the true value is the (tiny,
782        // nonzero) cos^2 term. This form reads that term straight off cos(phi) instead.
783        // For |k| <= 1 both addends are non-negative, so it cannot cancel at all.
784        let w = k2.mul_adde(c2, k.one_minus_sq());
785        let rf = carlson_rf::<P, E, V>(c2, w, V::ONE);
786        let core = if const { KIND == KIND_F } {
787            // F(phi, k) = sin(phi) * R_F(cos^2 phi, 1 - k^2 sin^2 phi, 1)
788            s * rf
789        } else if const { KIND == KIND_PI } {
790            // Pi(n, phi, k) = sin(phi) R_F + (n/3) sin^3(phi) R_J(c2, w, 1, 1 - n sin^2 phi)
791            let pp = n.nmul_adde(s * s, V::ONE);
792            let rj = carlson_rj::<P, E, V>(c2, w, V::ONE, pp);
793            let s3 = s * s * s;
794            (n.scale(const_element!(ratio <E>: 1 / 3)) * s3).mul_adde(rj, s * rf) // (n/3) s^3 rj + s rf
795        } else {
796            let rd = carlson_rd::<P, E, V>(c2, w, V::ONE);
797            let s3 = s * s * s;
798            if const { KIND == KIND_E } {
799                // E(phi, k) = sin(phi) R_F - (k^2 / 3) sin^3(phi) R_D
800                (k2.scale(const_element!(ratio <E>: 1 / 3)) * s3).nmul_adde(rd, s * rf)
801            } else {
802                // D(phi, k) = (1/3) sin^3(phi) R_D
803                s3.scale(const_element!(ratio <E>: 1 / 3)) * rd
804            }
805        };
806
807        // Quasi-period correction: I(phi) = core + 2m * I_complete. Skip the (expensive)
808        // recursive complete evaluation when no lane was reduced, unless the policy forbids
809        // branching. Must use `select`, not a bare `2m * complete` add: for |k| > 1 lanes the
810        // complete value is NaN, but |k| > 1 forces phi < pi/2 so m = 0 there - the select
811        // keeps `core` for those lanes and avoids 0 * NaN poisoning them.
812        if const { P::POLICY.avoid_branching } || !m.is_zero().all() {
813            let complete = ellint_impl::<P, E, V, KIND, true>(phi, k, n);
814            m.is_zero().select(core, (m + m).mul_adde(complete, core))
815        } else {
816            core
817        }
818    }
819}
820
821/// A Carlson symmetric elliptic integral request. The implementors are small structs that carry
822/// the integral's arguments as named fields ([`CarlsonRf`], [`CarlsonRc`], [`CarlsonRd`],
823/// [`CarlsonRj`], [`CarlsonRg`]), so each kind has exactly its own arguments - no dummy slots, and
824/// the special roles (`R_J`'s parameter `p`, `R_D`'s repeated `z`) are named at the call site.
825pub trait CarlsonKind {
826    /// The vector type returned (and the field type of the request struct).
827    type Output;
828    /// Evaluate the integral under precision policy `P`, consuming the request.
829    fn eval<P: Policy>(self) -> Self::Output;
830}
831
832/// The Jacobi zeta function `$Z(\varphi, k)$`.
833///
834/// The oscillating part of the incomplete integral of the second kind (what is left of
835/// `$E(\varphi, k)$` once its linear growth is removed):
836///
837/// ```math
838/// Z(\varphi, k) = E(\varphi, k) - \frac{E(k)}{K(k)} F(\varphi, k)
839/// ```
840///
841/// Odd in `phi`, `pi`-periodic, and exactly zero at every multiple of `pi/2`.
842///
843/// That defining difference is **not** how it is evaluated. Both terms grow with `phi` while
844/// `Z` does not, so the subtraction cancels wherever `Z` is small, which is near the zeros,
845/// i.e. everywhere the function is most delicate. The Carlson form used instead has no
846/// subtraction in it at all:
847///
848/// ```math
849/// Z(\varphi, k) = \frac{k^2 \sin\varphi \cos\varphi \sqrt{1 - k^2\sin^2\varphi}}{3 K(k)}
850///                 R_J(0,\ k'^2,\ 1,\ 1 - k^2\sin^2\varphi)
851/// ```
852///
853/// and `$1 - k^2\sin^2\varphi$` is itself formed as `$k'^2 + k^2\cos^2\varphi$`, a sum of two
854/// non-negative terms, so it cannot cancel either. Measured against the defining difference
855/// at 40 digits, worst relative error 3.1e-15 over `k` to 0.999 and `|phi|` to 4.5.
856///
857/// No sign fixup is needed: `sin` is odd and every other factor is even in `phi`, so the
858/// oddness falls out. `k = 1` is the one modulus with no Carlson form (`$k'^2 = 0$` gives
859/// `R_J` two zero arguments and `K` is infinite) and takes the limit
860/// `$\sin\varphi\,\operatorname{sign}(\cos\varphi)$` instead.
861#[inline(always)]
862pub fn jacobi_zeta<P, E, V>(phi: V, k: V) -> V
863where
864    P: Policy,
865    E: FloatElement + EllipticConsts,
866    V: FloatVector<Element = E> + SpecializedSpecialMath<E>,
867{
868    let k2 = k * k;
869    let kp = k.one_minus_sq();
870    let (sin_phi, cos_phi) = phi.sin_cos_p::<P>();
871    let c2 = cos_phi * cos_phi;
872
873    // 1 - k^2 sin^2(phi), written as a sum of non-negative terms so it never cancels.
874    let one_minus_ks2 = k2.mul_adde(c2, kp);
875
876    let (k_complete, _) = agm_complete_ke::<P, E, V>(k);
877    let rj = carlson_rj::<P, E, V>(V::ZERO, kp, V::ONE, one_minus_ks2);
878
879    let num = (k2 * sin_phi * cos_phi * one_minus_ks2.sqrt() * rj).scale(const_element!(ratio <E>: 1 / 3));
880    let mut z = num.approx_div_p::<P>(k_complete);
881
882    if const { P::POLICY.check_overflow } {
883        // k = 1 leaves R_J with two zero arguments and K infinite, so the quotient is NaN
884        // rather than the limit. Mathematica's simplification of Z(phi, 1) is the signed sine.
885        let unit = k.abs().cmp_eq(V::ONE);
886        if const { P::POLICY.avoid_branching } || thermite::unlikely(unit.any()) {
887            z = unit.select(sin_phi.copysign(cos_phi * sin_phi), z);
888        }
889    }
890
891    z
892}
893
894/// Heuman's lambda function `$\Lambda_0(\varphi, k)$`.
895///
896/// ```math
897/// \Lambda_0(\varphi, k) = \frac{2}{\pi}\Big[E(k) F(\varphi, k') + K(k) E(\varphi, k')
898///                                          - K(k) F(\varphi, k')\Big]
899/// ```
900///
901/// with `$k' = \sqrt{1-k^2}$` the complementary modulus. `$\Lambda_0(0, k) = 0$` and
902/// `$\Lambda_0(\pi/2, k) = 1$`, which is what makes it the natural companion to the complete
903/// integral of the third kind. It is also the standard closed form for the off-axis field of a
904/// circular current loop.
905///
906/// Inside `$|\varphi| \le \pi/2$` a Carlson form avoids the three-term difference above:
907///
908/// ```math
909/// \Lambda_0 = \frac{2}{\pi}\frac{k'^2 \sin\varphi\cos\varphi}{\delta}
910///             \left[R_F(0, k'^2, 1) + \frac{k^2}{3\delta^2} R_J(0, k'^2, 1, p)\right],
911/// \qquad \delta^2 = 1 - k'^2\sin^2\varphi
912/// ```
913///
914/// **The parameter `p` is the delicate part.** Its textbook spelling is `$1 - k^2/\delta^2$`,
915/// which is exactly zero at `$\varphi = \pi/2$` and therefore rounds _negative_ just before
916/// it, and a negative fourth argument sends `R_J` into its Cauchy-principal-value branch,
917/// which is a different function. Since `$\delta^2 = k^2 + k'^2\cos^2\varphi$`, the parameter
918/// is identically `$k'^2\cos^2\varphi/\delta^2$`, a ratio of non-negative quantities that is
919/// correct at the endpoint and cannot go negative. Measured against the defining form at 40
920/// digits: 3.6e-15 worst with that spelling, versus a NaN with the textbook one.
921///
922/// Beyond `$|\varphi| > \pi/2$` the Carlson form no longer applies and the identity
923/// `$\Lambda_0 = F(\varphi,k')/K(k') + \tfrac{2}{\pi}K(k) Z(\varphi, k')$` takes over. That
924/// arm costs three more elliptic evaluations, so it is gated on a lane actually needing it:
925/// the function's usual domain is `$[0, \pi/2]$`.
926#[inline(always)]
927pub fn heuman_lambda<P, E, V>(phi: V, k: V) -> V
928where
929    P: Policy,
930    E: FloatElement + EllipticConsts,
931    V: FloatVector<Element = E> + SpecializedSpecialMath<E>,
932{
933    let k2 = k * k;
934    let kp = k.one_minus_sq();
935    let (sin_phi, cos_phi) = phi.sin_cos_p::<P>();
936    let c2 = cos_phi * cos_phi;
937
938    // delta^2 = 1 - k'^2 sin^2(phi) = k^2 + k'^2 cos^2(phi). The same trick as `jacobi_zeta`'s,
939    // with k^2 and k'^2 swapped: two non-negative terms, so no cancellation.
940    let d2 = kp.mul_adde(c2, k2);
941    let delta = d2.sqrt();
942
943    // p = 1 - k^2/delta^2, in the one form that stays non-negative through phi = pi/2. The
944    // subtraction rounds negative just short of the endpoint, and a negative fourth argument
945    // is a different R_J (the Cauchy principal value).
946    let p = (kp * c2).approx_div_p::<P>(d2);
947
948    let rf = carlson_rf::<P, E, V>(V::ZERO, kp, V::ONE);
949    let rj = carlson_rj::<P, E, V>(V::ZERO, kp, V::ONE, p);
950
951    let bracket = rf + (k2 * rj).scale(const_element!(ratio <E>: 1 / 3)).approx_div_p::<P>(d2);
952    let scale = (kp * sin_phi * cos_phi).approx_div_p::<P>(delta * V::FRAC_PI_2);
953    let mut result = scale * bracket;
954
955    // Outside [-pi/2, pi/2] the Carlson form does not hold and the Legendre identity takes
956    // over. Three more elliptic evaluations, so only when a lane is actually out there.
957    let far = phi.abs().cmp_gt(V::FRAC_PI_2);
958    if const { P::POLICY.avoid_branching } || thermite::unlikely(far.any()) {
959        let k_prime = kp.sqrt();
960        let f_inc = ellint_impl::<P, E, V, KIND_F, false>(phi, k_prime, k);
961        let (k_prime_complete, _) = agm_complete_ke::<P, E, V>(k_prime);
962        let (k_complete, _) = agm_complete_ke::<P, E, V>(k);
963
964        let ratio = f_inc.approx_div_p::<P>(k_prime_complete);
965        let zeta = jacobi_zeta::<P, E, V>(phi, k_prime);
966        result = far.select(ratio + (k_complete * zeta).approx_div_p::<P>(V::FRAC_PI_2), result);
967    }
968
969    result
970}
971
972/// Legendre elliptic integral request, dual to [`CarlsonKind`] (see it for the named-field
973/// rationale). The implementors are [`EllintK`]/[`EllintF`] (1st kind, complete/incomplete),
974/// [`EllintE`]/[`EllintEInc`] (2nd kind), [`EllintD`]/[`EllintDInc`], and
975/// [`EllintPi`]/[`EllintPiInc`] (3rd kind). Completeness is encoded by the *fields*: a complete
976/// integral has no `phi`.
977pub trait EllipticKind {
978    /// The vector (or scalar) type returned.
979    type Output;
980    /// Evaluate the integral under precision policy `P`, consuming the request.
981    fn eval<P: Policy>(self) -> Self::Output;
982}
983
984use thermite::math::scalar::Unwrap;
985
986/// The wider (vector) form of a scalar request struct - the inverse of [`Unwrap`] for these structs.
987/// `Unwrap` only maps vector -> scalar (`type Unwrapped`), and Rust cannot invert an associated type,
988/// so this names the forward (scalar -> vector) direction. It is what lets the `Scalar*` math-trait
989/// layer take a scalar request (`CarlsonRf<f64>`), wrap it into a width-1 vector request
990/// (`CarlsonRf<Vector<f64>>`), evaluate on the (vector-only) backend, then unwrap the scalar result -
991/// without implementing the backend twice. Implemented only for scalar-element request structs.
992pub trait WrapTo {
993    /// The vector request struct whose `Unwrapped` is `Self`.
994    type Wrapped: Unwrap<Unwrapped = Self>;
995}
996
997/// Generates a request struct (`$name<V>` with named fields) plus its `Unwrap` (field-wise
998/// vector<->scalar) and `WrapTo` (scalar -> width-1 vector) impls. The `CarlsonKind` / `EllipticKind`
999/// impl is added by the caller, bounded on `FloatVector` - so the compute backend is vector-only.
1000macro_rules! request_struct {
1001    ($(#[$meta:meta])* $name:ident { $($field:ident),* }) => {
1002        #[derive(Debug, Clone, Copy)]
1003        $(#[$meta])* pub struct $name<V> {
1004            $(pub $field: V,)*
1005        }
1006
1007        impl<V: Unwrap> Unwrap for $name<V> {
1008            type Unwrapped = $name<<V as Unwrap>::Unwrapped>;
1009            #[inline(always)]
1010            fn wrap(value: Self::Unwrapped) -> Self {
1011                $name { $($field: Unwrap::wrap(value.$field),)* }
1012            }
1013            #[inline(always)]
1014            fn unwrap(self) -> Self::Unwrapped {
1015                $name { $($field: self.$field.unwrap(),)* }
1016            }
1017        }
1018
1019        impl<E> WrapTo for $name<E>
1020        where
1021            E: FloatElement + thermite::register::FloatRegister<Storage = E>,
1022            thermite::Vector<E>: Unwrap<Unwrapped = E>,
1023        {
1024            type Wrapped = $name<thermite::Vector<E>>;
1025        }
1026    };
1027}
1028
1029/// Carlson request structs. `$func` is the free function the struct evaluates (all take `<P, E, V>`).
1030macro_rules! decl_carlson {
1031    ($( $(#[$meta:meta])* struct $name:ident { $($field:ident),* } => $func:ident; )*) => {$(
1032        request_struct! { $(#[$meta])* $name { $($field),* } }
1033
1034        impl<E, V> CarlsonKind for $name<V>
1035        where
1036            E: FloatElement + EllipticConsts,
1037            V: FloatVector<Element = E> + SpecializedSpecialMath<E>,
1038        {
1039            type Output = V;
1040            #[inline(always)]
1041            fn eval<P: Policy>(self) -> V {
1042                $func::<P, E, V>($(self.$field),*)
1043            }
1044        }
1045    )*};
1046}
1047
1048decl_carlson! {
1049    /// Carlson `R_F(x, y, z)` - symmetric integral of the first kind. See [`CarlsonKind`].
1050    struct CarlsonRf { x, y, z } => carlson_rf;
1051
1052    /// Carlson `R_C(x, y) = R_F(x, y, y)` - degenerate first kind. See [`CarlsonKind`].
1053    struct CarlsonRc { x, y } => carlson_rc;
1054
1055    /// Carlson `R_D(x, y, z) = R_J(x, y, z, z)` - degenerate third kind; `z` is the repeated argument.
1056    /// See [`CarlsonKind`].
1057    struct CarlsonRd { x, y, z } => carlson_rd;
1058
1059    /// Carlson `R_J(x, y, z, p)` - symmetric integral of the third kind; `p` is the parameter (a
1060    /// Cauchy principal value when `p < 0`). See [`CarlsonKind`].
1061    struct CarlsonRj { x, y, z, p } => carlson_rj;
1062
1063    /// Carlson `R_G(x, y, z)` - symmetric integral of the second kind. See [`CarlsonKind`].
1064    struct CarlsonRg { x, y, z } => carlson_rg;
1065}
1066
1067/// Legendre request structs. Each maps to `ellint_impl::<KIND, COMPLETE>(phi, k, n)`; the three
1068/// trailing field names give the (phi, k, n) arguments - absent slots reuse `k` (ignored: COMPLETE
1069/// drops phi, non-Pi drops n).
1070macro_rules! decl_ellint {
1071    ($( $(#[$meta:meta])* struct $name:ident { $($field:ident),* } = [$kind:expr, $complete:expr]($phi:ident, $k:ident, $n:ident); )*) => {$(
1072        request_struct! { $(#[$meta])* $name { $($field),* } }
1073
1074        impl<E, V> EllipticKind for $name<V>
1075        where
1076            E: FloatElement + EllipticConsts,
1077            V: FloatVector<Element = E> + SpecializedSpecialMath<E>,
1078        {
1079            type Output = V;
1080            #[inline(always)]
1081            fn eval<P: Policy>(self) -> V {
1082                ellint_impl::<P, E, V, { $kind }, { $complete }>(self.$phi, self.$k, self.$n)
1083            }
1084        }
1085    )*};
1086}
1087
1088decl_ellint! {
1089    /// Complete elliptic integral of the first kind, `K(k)`.
1090    struct EllintK { k } = [KIND_F, true](k, k, k);
1091
1092    /// Incomplete elliptic integral of the first kind, `F(phi, k)`.
1093    struct EllintF { phi, k } = [KIND_F, false](phi, k, k);
1094
1095    /// Complete elliptic integral of the second kind, `E(k)`.
1096    struct EllintE { k } = [KIND_E, true](k, k, k);
1097
1098    /// Incomplete elliptic integral of the second kind, `E(phi, k)`.
1099    struct EllintEInc { phi, k } = [KIND_E, false](phi, k, k);
1100
1101    /// Complete `D(k) = (K - E) / k^2`.
1102    struct EllintD { k } = [KIND_D, true](k, k, k);
1103
1104    /// Incomplete `D(phi, k)`.
1105    struct EllintDInc { phi, k } = [KIND_D, false](phi, k, k);
1106
1107    /// Complete elliptic integral of the third kind, `Pi(n, k)`.
1108    struct EllintPi { n, k } = [KIND_PI, true](k, k, n);
1109
1110    /// Incomplete elliptic integral of the third kind, `Pi(n, phi, k)`.
1111    struct EllintPiInc { n, phi, k } = [KIND_PI, false](phi, k, n);
1112}
1113
1114/// The members of the elliptic family that are _not_ Legendre integrals and so do not route
1115/// through [`ellint_impl`]. Same request-struct shape as [`decl_carlson`], evaluating a free
1116/// function of the struct's own fields, but implementing [`EllipticKind`] so they reach callers
1117/// through the same `ellint` entry point as their siblings.
1118macro_rules! decl_ellint_fn {
1119    ($( $(#[$meta:meta])* struct $name:ident { $($field:ident),* } => $func:ident; )*) => {$(
1120        request_struct! { $(#[$meta])* $name { $($field),* } }
1121
1122        impl<E, V> EllipticKind for $name<V>
1123        where
1124            E: FloatElement + EllipticConsts,
1125            V: FloatVector<Element = E> + SpecializedSpecialMath<E>,
1126        {
1127            type Output = V;
1128            #[inline(always)]
1129            fn eval<P: Policy>(self) -> V {
1130                $func::<P, E, V>($(self.$field),*)
1131            }
1132        }
1133    )*};
1134}
1135
1136decl_ellint_fn! {
1137    /// Jacobi zeta `Z(phi, k)`: the oscillating part of `E(phi, k)`. See `jacobi_zeta`.
1138    struct JacobiZeta { phi, k } => jacobi_zeta;
1139
1140    /// Heuman's lambda `Lambda_0(phi, k)`, the complementary-modulus companion to the complete
1141    /// integral of the third kind. See `heuman_lambda`.
1142    struct HeumanLambda { phi, k } => heuman_lambda;
1143}
Last built: 2026-09-08 21:35:55 UTC