thermite_special/specialized/
bessel.rs

1//! The Bessel family's glue: what sits between the kernels in
2//! [`generic::bessel`](super::generic::bessel) and the per-element `SpecialMath` impls.
3//!
4//! Three things live here. The **reflection helpers** for negative whole orders, which
5//! every entry point applies after evaluating at `|N|`. The **entry-point stamping
6//! macros** for the spherical and Airy families, whose bodies differ per element only in
7//! the tables they name. And [`BesselDetails`], the per-arithmetic decisions of the
8//! real-order `I`/`K` kernel, which is what lets `thermite-complex` run that kernel in
9//! complex arithmetic (together with the doc-hidden [`kernels`] re-export it reaches the
10//! bodies through).
11
12use thermite::{
13    math::{PrimalProjection, specialized::SpecializedPrimalMath},
14    register::FloatElement,
15};
16
17/// Whether a negative Bessel order owes its result a sign flip.
18///
19/// The four families reflect differently at negative **integer** order:
20/// `$J_{-n} = (-1)^n J_n$` and `$Y_{-n} = (-1)^n Y_n$`, while `$I_{-n} = I_n$` and
21/// `$K_{-n} = K_n$` outright. Every kernel evaluates at `|N|` (the absolute value cannot live
22/// in a const-generic argument without `generic_const_exprs`), so this is the sign the entry
23/// point still owes, and is `true` only for `J`/`Y` at odd negative orders.
24///
25/// A flip is exact, so it costs no accuracy and the `Reference` tier stays bit-identical.
26#[inline(always)]
27pub(crate) const fn bessel_reflect_negates(n: i32) -> bool {
28    n < 0 && n % 2 != 0
29}
30
31/// The per-lane twin of [`bessel_reflect_negates`], for the runtime-order `J`/`Y` entries.
32///
33/// Returns `(|nu|, flip)`: the magnitude to evaluate at, and the lanes owed a negation. The
34/// recurrences only walk upward from order 0, so a negative order has to arrive as its
35/// magnitude and be reflected afterwards.
36///
37/// The parity test stays in the float domain deliberately. `nu` is an exact integer here, so
38/// `nu/2` is either integral or exactly half-integral, and comparing against its own floor
39/// answers "odd?" without casting an integer mask across to the float mask type.
40#[inline(always)]
41pub(crate) fn bessel_reflect_v<V: thermite::vector::FloatVector>(nu: V) -> (V, V::Mask) {
42    let na = nu.abs();
43    let half = na * V::HALF;
44
45    (na, nu.cmp_lt(V::ZERO) & half.cmp_ne(half.floor()))
46}
47
48/// Stamps the spherical Bessel entry points for one element type.
49///
50/// The four families share two kernels and one derivative formula. What differs per entry is
51/// which slot of the returned pair is the value, whether the neighbour enters negated, and
52/// whether the exponential is factored out.
53macro_rules! impl_sph_bessel_entries {
54    ($e:ty, $bi0:expr) => {
55        #[inline(always)]
56        fn sph_bessel_j_n<P: Policy, const N: usize>(self) -> Self {
57            generic::bessel::spherical::sph_jy_impl_n::<P, $e, _, N>(self).1
58        }
59
60        #[inline(always)]
61        fn sph_bessel_y_n<P: Policy, const N: usize>(self) -> Self {
62            generic::bessel::spherical::sph_jy_impl_n::<P, $e, _, N>(self).3
63        }
64
65        #[inline(always)]
66        fn sph_bessel_i_n<P: Policy, const N: usize>(self) -> Self {
67            generic::bessel::spherical::sph_ik_impl_n::<P, $e, _, N,false>(self, $bi0.far_threshold).1
68        }
69
70        #[inline(always)]
71        fn sph_bessel_i_scaled_n<P: Policy, const N: usize>(self) -> Self {
72            generic::bessel::spherical::sph_ik_impl_n::<P, $e, _, N,true>(self, $bi0.far_threshold).1
73        }
74
75        #[inline(always)]
76        fn sph_bessel_k_n<P: Policy, const N: usize>(self) -> Self {
77            generic::bessel::spherical::sph_ik_impl_n::<P, $e, _, N,false>(self, $bi0.far_threshold).3
78        }
79
80        #[inline(always)]
81        fn sph_bessel_k_scaled_n<P: Policy, const N: usize>(self) -> Self {
82            generic::bessel::spherical::sph_ik_impl_n::<P, $e, _, N,true>(self, $bi0.far_threshold).3
83        }
84
85        #[inline(always)]
86        fn sph_bessel_j_with_deriv_n<P: Policy, const N: usize>(self) -> (Self, Self) {
87            let (prev, v, _, _) = generic::bessel::spherical::sph_jy_impl_n::<P, $e, _, N>(self);
88            (v, generic::bessel::spherical::sph_deriv_n::<$e, _, N,false>(self, prev, v))
89        }
90
91        #[inline(always)]
92        fn sph_bessel_y_with_deriv_n<P: Policy, const N: usize>(self) -> (Self, Self) {
93            let (_, _, prev, v) = generic::bessel::spherical::sph_jy_impl_n::<P, $e, _, N>(self);
94            (v, generic::bessel::spherical::sph_deriv_n::<$e, _, N,false>(self, prev, v))
95        }
96
97        #[inline(always)]
98        fn sph_bessel_i_with_deriv_n<P: Policy, const N: usize, const SCALED: bool>(self) -> (Self, Self) {
99            let (prev, v, _, _) = generic::bessel::spherical::sph_ik_impl_n::<P, $e, _, N,SCALED>(self, $bi0.far_threshold);
100            let d = generic::bessel::spherical::sph_deriv_n::<$e, _, N,false>(self, prev, v);
101            // `d/dx (e^{-x} i_n) = e^{-x}(i_n' - i_n)`: the scaling's own derivative, which
102            // does not cancel.
103            (v, if const { SCALED } { d - v } else { d })
104        }
105
106        #[inline(always)]
107        fn sph_bessel_k_with_deriv_n<P: Policy, const N: usize, const SCALED: bool>(self) -> (Self, Self) {
108            let (_, _, prev, v) = generic::bessel::spherical::sph_ik_impl_n::<P, $e, _, N,SCALED>(self, $bi0.far_threshold);
109            // `k`'s neighbour enters negated: it is the decaying solution.
110            let d = generic::bessel::spherical::sph_deriv_n::<$e, _, N,true>(self, prev, v);
111            // At the origin `d + v` is `-inf + inf`, but the derivative of `e^{x} k_n` is still
112            // `-inf` there, which `d` already holds.
113            (v, if const { SCALED } { self.is_zero().select(d, d + v) } else { d })
114        }
115
116        // ---- runtime-order twins, one line each -------------------------------------------
117
118        #[inline(always)]
119        fn sph_bessel_j<P: Policy>(self, n: u32) -> Self {
120            generic::bessel::spherical::sph_jy_impl::<P, $e, _>(self, n).1
121        }
122
123        #[inline(always)]
124        fn sph_bessel_y<P: Policy>(self, n: u32) -> Self {
125            generic::bessel::spherical::sph_jy_impl::<P, $e, _>(self, n).3
126        }
127
128        #[inline(always)]
129        fn sph_bessel_i<P: Policy>(self, n: u32) -> Self {
130            generic::bessel::spherical::sph_ik_impl::<P, $e, _, false>(self, n, $bi0.far_threshold).1
131        }
132
133        #[inline(always)]
134        fn sph_bessel_i_scaled<P: Policy>(self, n: u32) -> Self {
135            generic::bessel::spherical::sph_ik_impl::<P, $e, _, true>(self, n, $bi0.far_threshold).1
136        }
137
138        #[inline(always)]
139        fn sph_bessel_k<P: Policy>(self, n: u32) -> Self {
140            generic::bessel::spherical::sph_ik_impl::<P, $e, _, false>(self, n, $bi0.far_threshold).3
141        }
142
143        #[inline(always)]
144        fn sph_bessel_k_scaled<P: Policy>(self, n: u32) -> Self {
145            generic::bessel::spherical::sph_ik_impl::<P, $e, _, true>(self, n, $bi0.far_threshold).3
146        }
147
148        #[inline(always)]
149        fn sph_bessel_j_with_deriv<P: Policy>(self, n: u32) -> (Self, Self) {
150            let (prev, v, _, _) = generic::bessel::spherical::sph_jy_impl::<P, $e, _>(self, n);
151            (v, generic::bessel::spherical::sph_deriv::<$e, _, false>(self, n, prev, v))
152        }
153
154        #[inline(always)]
155        fn sph_bessel_y_with_deriv<P: Policy>(self, n: u32) -> (Self, Self) {
156            let (_, _, prev, v) = generic::bessel::spherical::sph_jy_impl::<P, $e, _>(self, n);
157            (v, generic::bessel::spherical::sph_deriv::<$e, _, false>(self, n, prev, v))
158        }
159
160        #[inline(always)]
161        fn sph_bessel_i_with_deriv<P: Policy, const SCALED: bool>(self, n: u32) -> (Self, Self) {
162            let (prev, v, _, _) = generic::bessel::spherical::sph_ik_impl::<P, $e, _, SCALED>(self, n, $bi0.far_threshold);
163            let d = generic::bessel::spherical::sph_deriv::<$e, _, false>(self, n, prev, v);
164            (v, if const { SCALED } { d - v } else { d })
165        }
166
167        #[inline(always)]
168        fn sph_bessel_k_with_deriv<P: Policy, const SCALED: bool>(self, n: u32) -> (Self, Self) {
169            let (_, _, prev, v) = generic::bessel::spherical::sph_ik_impl::<P, $e, _, SCALED>(self, n, $bi0.far_threshold);
170            let d = generic::bessel::spherical::sph_deriv::<$e, _, true>(self, n, prev, v);
171            (v, if const { SCALED } { self.is_zero().select(d, d + v) } else { d })
172        }
173    };
174}
175
176/// Stamps the ten Airy entry points for one element type.
177///
178/// They differ only in which of the four outputs each asks the kernel for, and whether the
179/// exponential is factored out: four `const bool`s and one more. Writing them out twice, once
180/// per element, would be eighty lines of near-identical text in which a single transposed flag
181/// would be invisible. Here the flags line up in a column.
182///
183/// The flag order is `(Ai, Ai', Bi, Bi')`, matching the return tuple and SciPy's `airy`.
184macro_rules! impl_airy_entries {
185    ($e:ty, $nh:literal, $ne:literal, $no:literal, $fnum:literal, $fden:literal, $lgamma:expr, $zero:expr, $bi0:expr) => {
186        /// `(Ai, Ai', Bi, Bi')`, all four.
187        #[inline(always)]
188        fn airy_tuple<P: Policy>(self) -> (Self, Self, Self, Self) {
189            generic::bessel::airy::airy_impl::<P, $e, _, $nh, $ne, $no, $fnum, $fden, false, true, true, true, true>(
190                self, $lgamma, $zero, $bi0.far_threshold,
191            )
192        }
193
194        /// `(Ai, Ai', Bi, Bi')`, all four, with the exponential factored out on the positive axis.
195        #[inline(always)]
196        fn airy_tuple_scaled<P: Policy>(self) -> (Self, Self, Self, Self) {
197            generic::bessel::airy::airy_impl::<P, $e, _, $nh, $ne, $no, $fnum, $fden, true, true, true, true, true>(
198                self, $lgamma, $zero, $bi0.far_threshold,
199            )
200        }
201
202        /// `Ai` alone: one Bessel pass, and `K` only within it.
203        #[inline(always)]
204        fn airy_ai<P: Policy>(self) -> Self {
205            generic::bessel::airy::airy_impl::<P, $e, _, $nh, $ne, $no, $fnum, $fden, false, true, false, false, false>(
206                self, $lgamma, $zero, $bi0.far_threshold,
207            )
208            .0
209        }
210
211        /// `e^zeta Ai` on the positive axis.
212        #[inline(always)]
213        fn airy_ai_scaled<P: Policy>(self) -> Self {
214            generic::bessel::airy::airy_impl::<P, $e, _, $nh, $ne, $no, $fnum, $fden, true, true, false, false, false>(
215                self, $lgamma, $zero, $bi0.far_threshold,
216            )
217            .0
218        }
219
220        /// `Bi` alone: one Bessel pass, `I` and `K` within it.
221        #[inline(always)]
222        fn airy_bi<P: Policy>(self) -> Self {
223            generic::bessel::airy::airy_impl::<P, $e, _, $nh, $ne, $no, $fnum, $fden, false, false, false, true, false>(
224                self, $lgamma, $zero, $bi0.far_threshold,
225            )
226            .2
227        }
228
229        /// `e^-zeta Bi` on the positive axis.
230        #[inline(always)]
231        fn airy_bi_scaled<P: Policy>(self) -> Self {
232            generic::bessel::airy::airy_impl::<P, $e, _, $nh, $ne, $no, $fnum, $fden, true, false, false, true, false>(
233                self, $lgamma, $zero, $bi0.far_threshold,
234            )
235            .2
236        }
237
238        /// `Ai'` alone: the order-2/3 pass, `K` only.
239        #[inline(always)]
240        fn airy_ai_prime<P: Policy>(self) -> Self {
241            generic::bessel::airy::airy_impl::<P, $e, _, $nh, $ne, $no, $fnum, $fden, false, false, true, false, false>(
242                self, $lgamma, $zero, $bi0.far_threshold,
243            )
244            .1
245        }
246
247        /// `e^zeta Ai'` on the positive axis.
248        #[inline(always)]
249        fn airy_ai_prime_scaled<P: Policy>(self) -> Self {
250            generic::bessel::airy::airy_impl::<P, $e, _, $nh, $ne, $no, $fnum, $fden, true, false, true, false, false>(
251                self, $lgamma, $zero, $bi0.far_threshold,
252            )
253            .1
254        }
255
256        /// `Bi'` alone: the order-2/3 pass, `I` and `K`.
257        #[inline(always)]
258        fn airy_bi_prime<P: Policy>(self) -> Self {
259            generic::bessel::airy::airy_impl::<P, $e, _, $nh, $ne, $no, $fnum, $fden, false, false, false, false, true>(
260                self, $lgamma, $zero, $bi0.far_threshold,
261            )
262            .3
263        }
264
265        /// `e^-zeta Bi'` on the positive axis.
266        #[inline(always)]
267        fn airy_bi_prime_scaled<P: Policy>(self) -> Self {
268            generic::bessel::airy::airy_impl::<P, $e, _, $nh, $ne, $no, $fnum, $fden, true, false, false, false, true>(
269                self, $lgamma, $zero, $bi0.far_threshold,
270            )
271            .3
272        }
273    };
274}
275
276/// Per-arithmetic decisions of the real-order modified Bessel kernel,
277/// [`bessel_ik_real`](kernels::bessel_ik_real).
278///
279/// The kernel's arithmetic (Temme's series, two continued fractions, the Wronskian, the
280/// asymptotic series) is the same over R and over C. What changes is every place it
281/// _compares_ the argument: region selects by magnitude, the domain test, the overflow
282/// corner of the exponential. On a real vector those are plain comparisons. On `Complex`
283/// `cmp_lt` is a lexicographic sort order, not a modulus, and would route far-off-axis
284/// points into the wrong arm. Same split as [`ExpIntDetails`](super::ExpIntDetails).
285///
286/// The **order** is always real (`V::Primal`), which is why every threshold here is a
287/// primal and why the kernel takes the order as a separate real vector.
288///
289/// The blanket impl below covers every primal type (real `f32`/`f64` vectors,
290/// `Compensated`) with the real-line defaults. `Complex` overrides all of it.
291pub trait BesselDetails<V: thermite::vector::FloatVector + PrimalProjection> {
292    /// Lanes in the small-argument region, `|z| <= 2`, where `K` comes from Temme's series.
293    #[inline(always)]
294    fn near(z: V) -> V::Mask {
295        z.cmp_le(V::TWO)
296    }
297
298    /// Lanes with `|z| >= threshold`, per lane: where `I` takes the asymptotic series.
299    #[inline(always)]
300    fn beyond(z: V, threshold: V::Primal) -> V::Mask {
301        z.cmp_ge(V::from_primal(threshold))
302    }
303
304    /// Lanes inside the kernel's domain: the open positive axis, or the closed right
305    /// half-plane less the origin. The origin itself is selected to its limits by the
306    /// kernel. Everything else outside this mask is NaN.
307    #[inline(always)]
308    fn valid(z: V) -> V::Mask {
309        z.cmp_gt(V::ZERO)
310    }
311
312    /// Lanes where a single `e^z` overflows before `e^z * a` does, so the exponential is
313    /// halved and applied twice: `Re z >= threshold`.
314    #[inline(always)]
315    fn exp_far(z: V, threshold: V::Primal) -> V::Mask {
316        z.cmp_ge(V::from_primal(threshold))
317    }
318
319    /// Whether the large-argument expansion of `I` carries its second exponential.
320    ///
321    /// `$I_\nu(z) \sim \frac{e^{z}}{\sqrt{2\pi z}}\sum(-1)^k a_k z^{-k} + \frac{e^{-z \pm
322    /// (\nu+1/2)\pi i}}{\sqrt{2\pi z}}\sum a_k z^{-k}$` (DLMF 10.40.5). On the real line the
323    /// second term is `$e^{-2x}$` relative and below epsilon wherever the arm runs, so the
324    /// default drops it. Off the axis its modulus is `$e^{-2\,\mathrm{Re}\,z}$`, which on
325    /// the imaginary axis is **one**, so `Complex` keeps it.
326    const ASYM_TWO_TERMS: bool = false;
327
328    /// The exponent of that second term in the _scaled_ domain, `$-2z \pm (\nu+1/2)\pi i$`
329    /// with the sign of `Im z`. Only read when [`ASYM_TWO_TERMS`](Self::ASYM_TWO_TERMS).
330    #[inline(always)]
331    fn asym_second_exponent(z: V, _nu: V::Primal) -> V {
332        z
333    }
334}
335
336impl<E: FloatElement, V: thermite::vector::FloatVector<Element = E> + SpecializedPrimalMath<E>> BesselDetails<V> for V {}
337
338/// The shared Bessel kernels a downstream arithmetic instantiates for itself.
339///
340/// `thermite-complex` runs `bessel_ik_real` in complex arithmetic with its own
341/// [`BesselDetails`], and nothing here is API for anyone else.
342#[doc(hidden)]
343pub mod kernels {
344    pub use crate::specialized::generic::bessel::ik::{asymptotic_series_g, unscale_i_pair_masked};
345    pub use crate::specialized::generic::bessel::ik_real::bessel_ik_real;
346}
thermite_special/specialized/bessel.rs

thermite_special/specialized/
bessel.rs
