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} = \frac{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}