Skip to main content

thermite_complex/math/
mod.rs

1// The `fn name[..][..](self: Self, ..)` shape in the `decl_complex_math!` invocation
2// below is the macro DSL's, as in `thermite::math`.
3#![allow(clippy::needless_arbitrary_self_type)]
4
5//! Math kernels for [`Complex`].
6//!
7//! Implements the `Specialized*Math` traits. Complex vectors thereby get
8//! [`CoreMath`](thermite::math::CoreMath),
9//! [`TranscendentalMath`](thermite::math::TranscendentalMath) and
10//! [`SpatialMath`](thermite::math::SpatialMath) (with their `_p` policy forms)
11//! from the blanket impls in `thermite::math`, plus
12//! [`SpecializedComplexMath`] for the operations
13//! over the real part.
14//!
15//! Each kernel evaluates its complex function through the *inner* vector's policy
16//! math (`sin_cos_p`, `exp_p`, `atan2_p`, ...), carrying the caller's policy all the
17//! way down. Anything expressible by composition (`sin`, `cos`, `tan_pi`,
18//! `powi`, `sqrt1pm1`, `compound`, ...) comes from the trait defaults, which are
19//! already correct over C.
20//!
21//! [`RealMath`](thermite::math::RealMath) is not implemented; see the note at the
22//! bottom of this file.
23
24use thermite::math::policy::{DefaultPolicy, Policy, PrecisionPolicy};
25use thermite::math::specialized::{SpecializedCoreMath, SpecializedSpatialMath, SpecializedTranscendentalMath};
26use thermite::prelude::*;
27
28use crate::Complex;
29use self::specialized::{ComplexVector, SpecializedComplexMath};
30use crate::vector::RealFloatVector;
31
32// A copy of thermite::math's (private) decl_math!, dropping the ScalarMath
33// aggregate, which only makes sense for bare f32/f64. The rest is unchanged, so
34// ComplexMath is generated as TranscendentalMath is, #[dispatch] trampolines and
35// all.
36macro_rules! decl_complex_math {
37    ($(
38        $(#[$trait_meta:meta])*
39        trait $trait:ident<$element:ident> $(: $($bound:ident)&+ )? { $(
40            $(#[$meta:meta])*
41            fn $name:ident [ $($generics:tt)* ][$($generic_names:ident),*]( $($arg_name:ident : $arg_ty:ty),* $(,)?) -> $ret:ty;
42        )*}
43    )*) => {paste::paste! {$(
44        #[doc = "" $trait " math functions with customizable policies."]
45        $(#[$trait_meta])*
46        #[doc = ""]
47        #[doc = "Each function takes a [`Policy`] as its first generic argument. For the"]
48        #[doc = "default-policy versions (same names, no `_p` suffix), see [`" $trait "Math`]."]
49        #[doc = ""]
50        #[doc = "Implemented automatically for every type implementing [`Specialized" $trait "Math`]."]
51        #[thermite::dispatch(Self)]
52        pub trait [<$trait MathWithPolicy>] $(: $($bound +)+)? {$(
53            $(#[$meta])* fn [<$name _p>]<P: Policy, $($generics)*>($($arg_name: $arg_ty),*) -> $ret;
54        )*}
55
56        #[doc = "" $trait " math functions using the default policy."]
57        $(#[$trait_meta])*
58        #[doc = ""]
59        #[doc = "Every method here has a counterpart in [`" $trait "MathWithPolicy`] with a `_p`"]
60        #[doc = "suffix that takes an explicit [`Policy`]."]
61        #[doc = ""]
62        #[doc = "Implementors of [`" $trait "MathWithPolicy`] implement this automatically."]
63        #[thermite::dispatch(Self)]
64        pub trait [<$trait Math>]: [<$trait MathWithPolicy>] {$(
65            $(#[$meta])* #[inline(always)] fn $name<$($generics)*>($($arg_name: $arg_ty),*) -> $ret
66            { [<$trait MathWithPolicy>]::[<$name _p>]::<DefaultPolicy, $($generic_names),*>($($arg_name),*) }
67        )*}
68
69        impl<M> [<$trait Math>] for M where M: [<$trait MathWithPolicy>] {}
70
71        // The FloatVector<Element = E> bound is what ties E down, as in core.
72        #[thermite::dispatch(Self)]
73        impl<E: $element, V: FloatVector<Element = E> + $($($bound +)+)?> [<$trait MathWithPolicy>] for V
74            where V: [<Specialized $trait Math>]<E>
75        {$(
76            $(#[$meta])* #[inline(always)] fn [<$name _p>]<P: Policy, $($generics)*>($($arg_name: $arg_ty),*) -> $ret
77            { <V as [<Specialized $trait Math>]<E>>::$name::<P, $($generic_names),*>($($arg_name),*) }
78        )*})*
79    }};
80}
81
82decl_complex_math! {
83    /// Operations whose result is real (modulus, argument, polar form) or whose
84    /// argument is (a real power, base, or logarithm base), which the `Self -> Self`
85    /// core families cannot express.
86    ///
87    /// The purely complex operations (`exp`, `ln`, `sin`, `sqrt`, `powf`, ...) fit
88    /// the core families and come from
89    /// [`TranscendentalMath`](thermite::math::TranscendentalMath) as they do for
90    /// any other vector.
91    trait Complex<FloatElement>: ComplexVector {
92        /// The modulus `$|z|$`, as a real value.
93        ///
94        /// Uses `hypot`, so it does not overflow for large components the way
95        /// `sqrt(norm_sqr())` would.
96        fn norm[][](self: Self) -> Self::Real;
97
98        /// The principal argument `arg(z)`, in `(-pi, pi]`, as a real value.
99        fn arg[][](self: Self) -> Self::Real;
100
101        /// Converts to polar form `(r, theta)`, such that `self == r * exp(i*theta)`.
102        fn to_polar[][](self: Self) -> (Self::Real, Self::Real);
103
104        /// Builds a complex number from a polar representation `r * exp(i*theta)`.
105        fn from_polar[][](r: Self::Real, theta: Self::Real) -> Self;
106
107        /// Raises `self` to a real power.
108        ///
109        /// The complex-exponent form is [`powf`](thermite::math::TranscendentalMath::powf).
110        fn powfr[][](self: Self, e: Self::Real) -> Self;
111
112        /// Raises a real base to the complex power `self`.
113        fn expf[][](self: Self, base: Self::Real) -> Self;
114
115        /// The logarithm of `self` in an arbitrary real base.
116        ///
117        /// The complex-base form is [`log`](thermite::math::TranscendentalMath::log).
118        fn logr[][](self: Self, base: Self::Real) -> Self;
119
120        /// `1/self`, scaling by the modulus and not its square.
121        ///
122        /// Survives the magnitudes where [`inv`](ComplexVector::inv) would have
123        /// `norm_sqr()` overflow to infinity or underflow to zero.
124        fn finv[][](self: Self) -> Self;
125
126        /// `self/rhs`, scaling by the modulus and not its square.
127        ///
128        /// Survives the magnitudes where `/` would have `rhs.norm_sqr()` overflow
129        /// to infinity or underflow to zero.
130        fn fdiv[][](self: Self, rhs: Self) -> Self;
131    }
132}
133
134pub mod specialized;
135
136#[cfg(feature = "special")]
137pub mod special;
138
139// --- SpecializedComplexMath: the real-valued and real-argument operations ---
140
141impl<V: RealFloatVector> SpecializedComplexMath<Complex<V::Element>> for Complex<V> {
142    #[inline(always)]
143    fn norm<P: Policy>(self) -> V {
144        self.re.hypot_p::<P>(self.im)
145    }
146
147    #[inline(always)]
148    fn arg<P: Policy>(self) -> V {
149        self.im.atan2_p::<P>(self.re)
150    }
151
152    #[inline(always)]
153    fn from_polar<P: Policy>(r: V, theta: V) -> Self {
154        // C99 takes `e^(-inf + iy)` to +-0 for every non-finite `y`, and the same rule
155        // is what every caller here wants: once the modulus has collapsed to zero the
156        // result is the origin whatever direction it was approached from, but
157        // `r * cos(theta)` still reads `0 * NaN`. `z^(1 + i)` at z = 0 has a genuine
158        // -inf angle - the spiral never settles - and comes out NaN without this.
159        //
160        // The `is_finite` test is what keeps it honest: a finite angle is left exactly
161        // as it was, signed zeros included, so nothing well-defined moves.
162        let theta = if const { P::POLICY.check_overflow } {
163            theta.nz(r.cmp_eq(V::ZERO).bitandnot(theta.is_finite()))
164        } else {
165            theta
166        };
167
168        let (s, c) = theta.sin_cos_p::<P>();
169
170        Self::new(r * c, r * s)
171    }
172
173    #[inline(always)]
174    fn powfr<P: Policy>(self, e: V) -> Self {
175        // z^y = (r e^(i t))^y = r^y e^(i t y)
176        let (r, theta) = self.to_polar_p::<P>();
177
178        Self::from_polar_p::<P>(r.powf_p::<P>(e), theta * e)
179    }
180
181    #[inline(always)]
182    fn expf<P: Policy>(self, base: V) -> Self {
183        // b^(a + ci) = b^a * e^(i c ln b)
184        //
185        // `ln b` only ever reaches the angle here, so it can be masked outright.
186        let ln_b = finite_log_term::<P, V>(base.ln_p::<P>(), self.im);
187
188        Self::from_polar_p::<P>(base.powf_p::<P>(self.re), self.im * ln_b)
189    }
190
191    #[inline(always)]
192    fn logr<P: Policy>(self, base: V) -> Self {
193        // log_b(z) = ln(z) / ln(b); one real reciprocal, then scale both components.
194        let (r, theta) = self.to_polar_p::<P>();
195        let d = base.ln_p::<P>().reciprocal_p::<P>();
196
197        Self::new(r.ln_p::<P>() * d, theta * d)
198    }
199
200    #[inline(always)]
201    fn finv<P: Policy>(self) -> Self {
202        // Scaling twice by 1/|z| keeps every intermediate in range, where
203        // conj(z) / |z|^2 would overflow the square.
204        //
205        // Taking the reciprocal explicitly, rather than writing `(conj/n)/n`:
206        // `Div<V>` is already a reciprocal-and-scale, so the quotient form spent
207        // two divisions computing the same 1/|z| twice. `digamma` calls this in a
208        // loop.
209        let inv = self.norm_p::<P>().reciprocal_p::<P>();
210
211        self.conj() * inv * inv
212    }
213}
214
215// --- SpecializedCoreMath ---
216
217impl<V: RealFloatVector> SpecializedCoreMath<Complex<V::Element>> for Complex<V> {
218    /// `P(z)/Q(z)`, evaluated directly or through `1/z` depending on which is better
219    /// conditioned.
220    ///
221    /// Overridden only to change *which quantity* that decision is made on. The
222    /// generic default tests `x.cmp_gt(ONE)`, which over C is the lexicographic order
223    /// on `(re, im)` - so it keys off the real part alone and will happily evaluate
224    /// the direct form at `z = 10^150 i`, overflowing, while reporting that `z` is
225    /// "not greater than one". The condition that actually matters is `|z| > 1`.
226    ///
227    /// Both forms are the same rational function (`P_rev(1/z)/Q_rev(1/z)` differs from
228    /// `P(z)/Q(z)` only by `z^(D-N)`, corrected below), so this is a conditioning fix,
229    /// not a correctness one - except where the wrong choice overflows outright.
230    #[inline(always)]
231    fn poly_rational<P: Policy, const N: usize, const D: usize>(
232        self,
233        numerator: &[Complex<V::Element>; N],
234        denominator: &[Complex<V::Element>; D],
235    ) -> Self {
236        let x = self;
237
238        if const { P::POLICY.precision.le(thermite::math::policy::PrecisionPolicy::Average) } {
239            let n = SpecializedCoreMath::poly::<P, N>(x, numerator);
240            let d = SpecializedCoreMath::poly::<P, D>(x, denominator);
241
242            return n.approx_div_p::<P>(d);
243        }
244
245        let invert = x.norm_sqr().cmp_gt(V::ONE);
246
247        let mut n0 = Self::EMPTY;
248        let mut d0 = Self::EMPTY;
249
250        if const { P::POLICY.avoid_branching } || !invert.all() {
251            n0 = SpecializedCoreMath::poly::<P, N>(x, numerator);
252            d0 = SpecializedCoreMath::poly::<P, D>(x, denominator);
253        }
254
255        let mut z = Self::EMPTY;
256        let mut n1 = Self::EMPTY;
257        let mut d1 = Self::EMPTY;
258
259        if const { P::POLICY.avoid_branching } || invert.any() {
260            z = SpecializedCoreMath::reciprocal::<P>(x);
261            n1 = SpecializedCoreMath::poly_rev::<P, N>(z, numerator);
262            d1 = SpecializedCoreMath::poly_rev::<P, D>(z, denominator);
263        }
264
265        let n = invert.select(n1, n0);
266        let d = invert.select(d1, d0);
267
268        let res = n.approx_div_p::<P>(d);
269
270        // Same degree: the reversed form is the same value, nothing to undo.
271        if const { N == D } {
272            return res;
273        }
274
275        if const { P::POLICY.avoid_branching } || invert.any() {
276            let (u, e) = if const { N < D } { (z, D - N) } else { (x, N - D) };
277
278            return invert.select(res * SpecializedCoreMath::powi::<P>(u, e as i32), res);
279        }
280
281        res
282    }
283
284    // 1/sqrt(z) = conj(sqrt(z)) / |z|, since |sqrt(z)|^2 == |z|.
285    #[inline(always)]
286    fn inverse_sqrt<P: Policy>(self) -> Self {
287        let s = self.sqrt();
288        let inv = self.norm_p::<P>().reciprocal_p::<P>();
289
290        Complex::new(s.re * inv, -(s.im * inv))
291    }
292}
293
294// --- SpecializedTranscendentalMath ---
295
296/// `$iz = -\operatorname{Im} z + i\operatorname{Re} z$`: a component swap and a sign flip.
297///
298/// Worth a helper because the natural spelling is not free. `z.mul_add(Self::I, w)`
299/// is a complex FMA - four inner FMAs - over a constant of zeros and ones, and IEEE
300/// forbids folding `a*0.0 + b` to `b` (`a` may be infinite, and the zero has a sign),
301/// so all four survive into the assembly.
302#[inline(always)]
303fn mul_i<V: RealFloatVector>(z: Complex<V>) -> Complex<V> {
304    Complex::new(-z.im, z.re)
305}
306
307/// `$-iz = \operatorname{Im} z - i\operatorname{Re} z$`. See [`mul_i`].
308#[inline(always)]
309fn mul_neg_i<V: RealFloatVector>(z: Complex<V>) -> Complex<V> {
310    Complex::new(z.im, -z.re)
311}
312
313/// Replaces the lanes where a hyperbolic quotient has overflowed with its limit.
314///
315/// `tan`/`tanh` divide by `cosh` plus a bounded term. Past `$|2x| \approx 710$` in
316/// binary64 that denominator is infinite while the numerator's `sinh` is too, so the
317/// quotient is `inf/inf` - `NaN` in the component that should have saturated, and a
318/// signed zero in the other. The function itself is perfectly well behaved there and
319/// tends to a unit along one axis.
320///
321/// Keyed on the denominator rather than a magnitude threshold, so it fires exactly on
322/// the lanes that lost the value and needs no per-element constant. `is_infinite`, not
323/// `!is_finite`: a `NaN` argument must still produce `NaN`.
324///
325/// A no-op unless the policy sets
326/// [`check_overflow`](thermite::math::policy::PolicyParameters::check_overflow).
327#[inline(always)]
328fn saturate<P: Policy, V: RealFloatVector>(res: Complex<V>, denom: V, limit: Complex<V>) -> Complex<V> {
329    if const { !P::POLICY.check_overflow } {
330        return res;
331    }
332
333    let lost = denom.is_infinite();
334
335    if thermite::unlikely(lost.any()) {
336        return lost.select(limit, res);
337    }
338
339    res
340}
341
342/// `$\operatorname{asinh}(p) = \ln(p + \sqrt{1 + p^2})$`, conditioned at both ends.
343///
344/// Serves `asin` too: `$(iz)^2 = -z^2$`, so `$\sqrt{1 - z^2}$` is this function's
345/// `$\sqrt{1 + p^2}$` at `$p = iz$`, and `$\operatorname{asin}(z) = -i\operatorname{asinh}(iz)$`
346/// exactly. One kernel, two functions.
347///
348/// The naive `$\ln(p + s)$` loses everything at both extremes:
349///
350/// - **Large `$|p|$`**: `$s \to \mp p$` and the sum cancels. At `$z = 10^8 i$` the
351///   `asin` sum is `$-10^8 + 10^8$` - *exactly* zero - so `$\ln 0$` returned an
352///   infinity where the true value is `19.11i`. The companion `$s - p$` is the other
353///   root, and `$(p + s)(s - p) = s^2 - p^2 = 1$`, so it is both exact and the
354///   well-conditioned one. `$|w| < 1$` tests which cancelled, and `$\ln w = -\ln w'$`
355///   holds outright rather than up to `$2\pi i$`, the companion lying in the right
356///   half-plane exactly when it is selected.
357/// - **Small `$|p|$`**: `$w = 1 + O(p)$`, and forming that sum rounds away the very
358///   `$p$` the answer consists of - `asinh(1e-8)` kept 8 of its 16 digits. Feeding
359///   `$w - 1$` to `ln_1p` instead fixes it, provided `$w - 1$` is *not* formed by
360///   subtracting: `$s - 1 = p^2/(s + 1)$` has no cancellation of its own.
361///
362/// # Policy
363///
364/// Both corrections are gated at [`Best`](PrecisionPolicy::Best) and above, together
365/// costing one complex division and, only where a lane actually cancels, a second
366/// logarithm. Below that the plain `$\ln(p + s)$` stands, as before.
367#[inline(always)]
368fn log_asinh<P: Policy, V: RealFloatVector>(p: Complex<V>) -> Complex<V> {
369    let s = p.mul_add(p, Complex::ONE).sqrt();
370    let w = p + s;
371
372    if const { !P::POLICY.precision.ge(PrecisionPolicy::Best) } {
373        return w.ln_p::<P>();
374    }
375
376    // w - 1, without ever forming the difference.
377    let mut res = (p + p.square() / (s + Complex::ONE)).ln_1p_p::<P>();
378
379    // `|w| < 1` alone is too eager. For tiny `p` with a negative real part `w = 1 + p`
380    // sits just under one without anything having cancelled, and the log branch would
381    // undo the `ln_1p` correction above. Genuine cancellation needs the two terms to be
382    // large and nearly opposite, so `|p| > 1` as well - which is also where `u` itself
383    // stops being trustworthy, `u = w - 1 ~ -1` then being a difference of two terms of
384    // size `|p|`.
385    let flip = w.norm_sqr().cmp_lt(V::ONE) & p.norm_sqr().cmp_gt(V::ONE);
386
387    if thermite::unlikely(flip.any()) {
388        let l = (s - p).ln_p::<P>();
389
390        res = flip.select(Complex::new(-l.re, -l.im), res);
391    }
392
393    res
394}
395
396/// `$\ln w$`, taking the reciprocal companion where `w` has cancelled.
397///
398/// The inverse trigonometric and hyperbolic functions are each `$\ln$` of a sum of two
399/// terms whose *difference* is the algebraically conjugate root, and in every case the
400/// product of the two is exactly one:
401///
402/// ```text
403/// asin :  (iz + s)(s - iz) = s^2 + z^2 = 1      s = sqrt(1 - z^2)
404/// acos :  (z + is)(z - is) = z^2 + s^2 = 1      s = sqrt(1 - z^2)
405/// asinh:  (z + s)(s - z)   = s^2 - z^2 = 1      s = sqrt(1 + z^2)
406/// ```
407///
408/// So exactly one of the pair is well conditioned: for large `$|z|$` the two terms are
409/// nearly equal in magnitude, one sum cancels to nothing and the other doubles. At
410/// `$z = 10^8 i$` the `asin` sum is `$-10^8 + 10^8$`, *exactly zero*, and `$\ln 0$`
411/// returned an infinity for a true value of `19.11i`.
412///
413/// Since `$ww' = 1$`, `$|w| < 1$` is an exact test for which one cancelled, and
414/// `$\ln w = -\ln w'$` - unambiguously, not merely up to `$2\pi i$`, because the
415/// companion is in the right half-plane exactly when it is the one being selected.
416///
417/// One `ln` either way: the argument is blended *before* the logarithm, so the cost
418/// over the naive form is a complex add, a `norm_sqr` and two selects.
419///
420/// # Policy
421///
422/// Gated at [`Best`](PrecisionPolicy::Best) and above. Below it the cancelling form is
423/// used unconditionally, as it was before - the failure needs `$|z| \gg 1$`, and the
424/// lower tiers do not promise the digits that are lost.
425#[inline(always)]
426fn ln_reciprocal_pair<P: Policy, V: RealFloatVector>(w: Complex<V>, companion: Complex<V>) -> Complex<V> {
427    if const { !P::POLICY.precision.ge(PrecisionPolicy::Best) } {
428        return w.ln_p::<P>();
429    }
430
431    let flip = w.norm_sqr().cmp_lt(V::ONE);
432
433    let l = flip.select(companion, w).ln_p::<P>();
434
435    Complex::new(l.re.neg_c(flip), l.im.neg_c(flip))
436}
437
438/// `$b^{a+ci} - 1$` from the real `$b^a - 1$` and the angle `$\phi = c\ln b$`.
439///
440/// ```text
441/// b^(a+ci) - 1 = (b^a cos(phi) - 1) + i b^a sin(phi)
442///              = (bm1 cos(phi) + (cos(phi) - 1)) + i (bm1 + 1) sin(phi)
443/// ```
444///
445/// `exp_m1` and `cos_m1` are the cancellation-free primitives, keeping the relative
446/// accuracy near `z = 0` that the naive form loses entirely. Shared by the `e`, `2`
447/// and `10` bases, which differ only in which real `*_m1` supplies `bm1` and in the
448/// scale of `phi`.
449#[inline(always)]
450fn expm1_from<P: Policy, V: RealFloatVector>(bm1: V, phi: V) -> Complex<V> {
451    let (s, c) = phi.sin_cos_p::<P>();
452    let cm1 = phi.cos_m1_p::<P>();
453
454    Complex::new(bm1.mul_adde(c, cm1), s.mul_adde(bm1, s))
455}
456
457/// Masks `+-inf` out of a `d * ln r` angle term.
458///
459/// `ln r` is -inf at r = 0 and +inf at r = inf, so a real exponent - `d == 0`, the
460/// overwhelmingly common case - forms `0 * inf` where the limit is plainly 0, and one
461/// NaN there takes the whole result with it: `(0+0i)^2` was NaN on the strength of it.
462/// Killing the infinity before the product exists is a compare and an AND, no branch,
463/// and it leaves the term identical wherever `d` is not zero.
464///
465/// This is what reduces `powf` to the `powfr` formula, which never forms `ln r` at all.
466#[inline(always)]
467fn finite_log_term<P: Policy, V: RealFloatVector>(ln_r: V, d: V) -> V {
468    if const { !P::POLICY.check_overflow } {
469        return ln_r;
470    }
471
472    ln_r.zz(d.cmp_ne(V::ZERO))
473}
474
475impl<V: RealFloatVector> SpecializedTranscendentalMath<Complex<V::Element>> for Complex<V> {
476    /// `sin(a + bi) = sin(a)cosh(b) + i*cos(a)sinh(b)`,
477    /// `cos(a + bi) = cos(a)cosh(b) - i*sin(a)sinh(b)`.
478    #[inline(always)]
479    fn sin_cos<P: Policy>(self) -> (Self, Self) {
480        let (s, c) = self.re.sin_cos_p::<P>();
481        let (sh, ch) = self.im.sinh_cosh_p::<P>();
482
483        (Complex::new(s * ch, c * sh), Complex::new(c * ch, -(s * sh)))
484    }
485
486    /// `$\sin(\pi z)$` and `$\cos(\pi z)$`, from the *real* `sincos_pi`.
487    ///
488    /// Must be overridden rather than left to the default. That default is
489    /// `sin_cos(z * pi)`, which rounds `pi * Re z` before doing any reduction and so
490    /// throws away the exact argument reduction real `sincos_pi` performs near the
491    /// integers - precisely where the Gamma reflection formulas put their poles, and
492    /// where `sin(pi z)` passes through zero. It is also no more work: one real
493    /// `sincos_pi` and one real `sinh_cosh`, the same two calls the default makes.
494    #[inline(always)]
495    fn sincos_pi<P: Policy>(self) -> (Self, Self) {
496        let (s, c) = self.re.sincos_pi_p::<P>();
497        let (sh, ch) = (self.im * <V as thermite::math::FloatConsts>::PI).sinh_cosh_p::<P>();
498
499        (Complex::new(s * ch, c * sh), Complex::new(c * ch, -(s * sh)))
500    }
501
502    /// `$\tan(a + bi) = \frac{\sin 2a + i\sinh 2b}{\cos 2a + \cosh 2b}$`
503    ///
504    /// The doubled-angle form takes one real division, where the default
505    /// (`sin_cos` then a complex divide) takes a complex one.
506    ///
507    /// Under [`check_overflow`](thermite::math::policy::PolicyParameters::check_overflow)
508    /// the saturation is handled: `$\tan(z) \to i\,\mathrm{sign}(b)$` as `$|b|$` grows,
509    /// but `$\sinh$` and `$\cosh$` both overflow past `$|2b| \approx 710$` and the
510    /// quotient becomes `inf/inf`. `tan(1 + 400i)` was `NaN`.
511    #[inline(always)]
512    fn tan<P: Policy>(self) -> Self {
513        let (two_re, two_im) = (self.re + self.re, self.im + self.im);
514
515        let (s, c) = two_re.sin_cos_p::<P>();
516        let (sh, ch) = two_im.sinh_cosh_p::<P>();
517
518        let denom = c + ch;
519        let res = Complex::new(s, sh) / denom;
520
521        saturate::<P, V>(res, denom, Complex::new(V::ZERO, V::ONE.mul_sign(self.im)))
522    }
523
524    /// `sinh(a + bi) = sinh(a)cos(b) + i*cosh(a)sin(b)`,
525    /// `cosh(a + bi) = cosh(a)cos(b) + i*sinh(a)sin(b)`.
526    #[inline(always)]
527    fn sinh_cosh<P: Policy>(self) -> (Self, Self) {
528        let (s, c) = self.im.sin_cos_p::<P>();
529        let (sh, ch) = self.re.sinh_cosh_p::<P>();
530
531        (Complex::new(sh * c, ch * s), Complex::new(ch * c, sh * s))
532    }
533
534    /// `tanh(a + bi) = (sinh(2a) + i*sin(2b)) / (cosh(2a) + cos(2b))`.
535    ///
536    /// Saturates to `$\mathrm{sign}(a)$` for large `$|a|$` under `check_overflow`; see
537    /// [`tan`](Self::tan), of which this is the transpose.
538    #[inline(always)]
539    fn tanh<P: Policy>(self) -> Self {
540        let (two_re, two_im) = (self.re + self.re, self.im + self.im);
541
542        let (s, c) = two_im.sin_cos_p::<P>();
543        let (sh, ch) = two_re.sinh_cosh_p::<P>();
544
545        let denom = ch + c;
546        let res = Complex::new(sh, s) / denom;
547
548        saturate::<P, V>(res, denom, Complex::new(V::ONE.mul_sign(self.re), V::ZERO))
549    }
550
551    /// `$\mathrm{sinc}(z) = \sin(z)/z$`, with the removable singularity filled in.
552    #[inline(always)]
553    fn sinc<P: Policy>(self) -> Self {
554        let is_zero = self.is_zero();
555
556        // 0/0 = NaN at the origin, so the guard has to be a select; the quotient
557        // cannot be patched up after the fact.
558        let q = self.sin_p::<P>() / self;
559
560        is_zero.select(Self::ONE, q)
561    }
562
563    /// `$\mathrm{sinc}_\pi(z) = \frac{\sin(\pi z)}{\pi z}$`, singularity filled in.
564    ///
565    /// Overridden so the zeros are *exact*. The default is `sinc(z * pi)`, which
566    /// rounds `pi * Re z` before reducing; the subsequent division by `pi z` cancels
567    /// most of that error, so the default is accurate to about an ulp - but at a
568    /// non-zero integer it returns ~1e-16 rather than zero. Going through the real
569    /// `sin_pi`, which is exactly zero there, makes this exactly zero too.
570    ///
571    /// That is the property that makes `sinc_pi` an *interpolating* kernel: Lanczos
572    /// and sinc resampling reproduce their samples only if the kernel vanishes at
573    /// every non-zero integer.
574    #[inline(always)]
575    fn sinc_pi<P: Policy>(self) -> Self {
576        let is_zero = self.is_zero();
577
578        // As in `sinc`: 0/0 is NaN at the origin, so the guard must be a select.
579        // Trait-qualified: the blanket `TranscendentalMath::sin_pi` is equally in scope.
580        let q = SpecializedTranscendentalMath::sin_pi::<P>(self) / (self * <V as thermite::math::FloatConsts>::PI);
581
582        is_zero.select(Self::ONE, q)
583    }
584
585    /// `e^(a + bi) = e^a * (cos(b) + i*sin(b))`.
586    #[inline(always)]
587    fn exp<P: Policy>(self) -> Self {
588        Self::from_polar_p::<P>(self.re.exp_p::<P>(), self.im)
589    }
590
591    /// `exph(z) = e^z / 2`.
592    #[inline(always)]
593    fn exph<P: Policy>(self) -> Self {
594        Self::from_polar_p::<P>(self.re.exph_p::<P>(), self.im)
595    }
596
597    /// `$2^z = 2^a e^{ib\ln 2}$`.
598    ///
599    /// # Policy
600    ///
601    /// Above [`Average`](PrecisionPolicy::Average) the real part goes through the real
602    /// `exp2`. Rescaling it as `exp(a ln 2)` instead rounds `a ln 2` first, and `exp`
603    /// then amplifies that rounding by the argument - `exp2(1000)` is wrong in its
604    /// tenth digit (~300 ulp) that way. The imaginary part can afford the multiply
605    /// either way, feeding a `sincos` that reduces its own argument.
606    ///
607    /// At or below `Average` the rescaled form is used, `exp2` being the more
608    /// expensive kernel and 300 ulp being well inside that tier's budget.
609    #[inline(always)]
610    fn exp2<P: Policy>(self) -> Self {
611        if const { P::POLICY.precision.le(PrecisionPolicy::Average) } {
612            return (self * V::LN_2).exp_p::<P>();
613        }
614
615        Self::from_polar_p::<P>(self.re.exp2_p::<P>(), self.im * V::LN_2)
616    }
617
618    /// `$10^z = 10^a e^{ib\ln 10}$`. See [`exp2`](Self::exp2), including the policy split.
619    #[inline(always)]
620    fn exp10<P: Policy>(self) -> Self {
621        if const { P::POLICY.precision.le(PrecisionPolicy::Average) } {
622            return (self * V::LN_10).exp_p::<P>();
623        }
624
625        Self::from_polar_p::<P>(self.re.exp10_p::<P>(), self.im * V::LN_10)
626    }
627
628    /// `$e^z - 1$`, without the cancellation of forming `$e^z$` and subtracting one.
629    #[inline(always)]
630    fn exp_m1<P: Policy>(self) -> Self {
631        expm1_from::<P, V>(self.re.exp_m1_p::<P>(), self.im)
632    }
633
634    /// `$2^z - 1$`. Through the real `exp2_m1` above `Average`, as [`exp2`](Self::exp2) is.
635    #[inline(always)]
636    fn exp2_m1<P: Policy>(self) -> Self {
637        if const { P::POLICY.precision.le(PrecisionPolicy::Average) } {
638            return (self * V::LN_2).exp_m1_p::<P>();
639        }
640
641        expm1_from::<P, V>(self.re.exp2_m1_p::<P>(), self.im * V::LN_2)
642    }
643
644    /// `$10^z - 1$`. Through the real `exp10_m1` above `Average`, as [`exp10`](Self::exp10) is.
645    #[inline(always)]
646    fn exp10_m1<P: Policy>(self) -> Self {
647        if const { P::POLICY.precision.le(PrecisionPolicy::Average) } {
648            return (self * V::LN_10).exp_m1_p::<P>();
649        }
650
651        expm1_from::<P, V>(self.re.exp10_m1_p::<P>(), self.im * V::LN_10)
652    }
653
654    /// `z^w`, the principal value.
655    #[inline(always)]
656    fn powf<P: Policy>(self, e: Self) -> Self {
657        // z^w = (r e^(i t))^(c + di)
658        //     = r^c e^(-d t) * (cos(c t + d ln r) + i sin(c t + d ln r))
659        //     = from_polar(e^(c ln r - d t), c t + d ln r)
660        //
661        // The angle needs `ln r` at every policy, so it is computed once up front.
662        let (r, theta) = self.to_polar_p::<P>();
663        let ln_r = r.ln_p::<P>();
664
665        // Both products against `ln r` are `0 * inf` at the degenerate points, each for
666        // its own exponent part: `d ln r` in the angle, `c ln r` in the fused exponent.
667        // `0^0` is 1 and `0^2` is 0 only once both are masked.
668        let ln_r_angle = finite_log_term::<P, V>(ln_r, e.im);
669        let ln_r_mod = finite_log_term::<P, V>(ln_r, e.re);
670
671        let angle = e.im.mul_adde(ln_r_angle, e.re * theta);
672
673        let mut modulus = if const { P::POLICY.precision.le(PrecisionPolicy::Average) } {
674            // Fused exponent: one `exp` for the whole modulus. `ln r` is already in
675            // hand and `powf` is `exp(c ln r)` underneath, so the split form costs a
676            // second `exp` and a second `ln` for an extended-precision `c ln r` that
677            // this tier is not paying for.
678            e.im.nmul_adde(theta, e.re * ln_r_mod).exp_p::<P>()
679        } else {
680            // `powf` carries more precision through `c ln r` than the fused exponent
681            // can, which is what keeps `|z^w|` near an ulp once `|c ln r|` is large.
682            r.powf_p::<P>(e.re) * (-e.im * theta).exp_p::<P>()
683        };
684
685        if const { P::POLICY.check_overflow && !P::POLICY.precision.le(PrecisionPolicy::Average) } {
686            // `r^c` leaves the range on its own where `e^{-dt}` would have brought the
687            // product back - `(-1e200)^(2 + 300i)` is about 1e-9 and the split form
688            // gives `inf`, or `NaN` from the mirror-image `0 * inf`. The fused exponent
689            // has no such intermediate, so it covers those lanes.
690            let lost = !modulus.is_finite();
691
692            if thermite::unlikely(lost.any()) {
693                modulus = lost.select(e.im.nmul_adde(theta, e.re * ln_r_mod).exp_p::<P>(), modulus);
694            }
695        }
696
697
698        Self::from_polar_p::<P>(modulus, angle)
699    }
700
701    /// The principal cube root.
702    ///
703    /// This does not agree with the real cube root of a negative real: the real
704    /// cube root of -8 is -2, the principal complex one `$1 + i\sqrt{3}$`.
705    #[inline(always)]
706    fn cbrt<P: Policy>(self) -> Self {
707        let (r, theta) = self.to_polar_p::<P>();
708
709        // 1/3 is not representable, so divide; multiplying by a rounded reciprocal
710        // loses a bit at the higher policies.
711        let three: V = thermite::const_splat!(int <V::Element>: 3);
712
713        Self::from_polar_p::<P>(r.cbrt_p::<P>(), theta / three)
714    }
715
716    /// The principal `N`th root, `$z^{1/N} = |z|^{1/N} e^{i\arg(z)/N}$`.
717    ///
718    /// The trait default is real-only: for odd `N` it takes `abs()` and restores the
719    /// sign afterwards, which over C collapses `z` to its modulus and returns a real
720    /// root.
721    #[inline(always)]
722    fn nth_root<P: Policy, const N: usize>(self) -> Self {
723        let (r, theta) = self.to_polar_p::<P>();
724
725        // Not `const_splat!`: `N` is a generic parameter, which cannot appear in the
726        // const operation that arm expands to.
727        let n = V::splat(<V::Element as FloatElement>::from_int(N as thermite::LargeInt));
728
729        Self::from_polar_p::<P>(r.powf_p::<P>(n.reciprocal_p::<P>()), theta / n)
730    }
731
732    /// The principal natural logarithm: `ln(z) = ln|z| + i*arg(z)`.
733    ///
734    /// Branch cut on `(-inf, 0]`, continuous from above; `-pi <= Im(ln z) <= pi`.
735    #[inline(always)]
736    fn ln<P: Policy>(self) -> Self {
737        let (r, theta) = self.to_polar_p::<P>();
738
739        Complex::new(r.ln_p::<P>(), theta)
740    }
741
742    /// `$\ln(1 + z)$`, without the cancellation of forming `1 + z` first.
743    #[inline(always)]
744    fn ln_1p<P: Policy>(self) -> Self {
745        // |1 + z|^2 - 1 = 2a + a^2 + b^2 = a(a + 2) + b^2, so
746        //   Re = 0.5 * ln_1p(a(a + 2) + b^2),  Im = atan2(b, 1 + a)
747        // The real part goes through ln_1p to keep the accuracy near z = 0 that
748        // ln(1 + z) would lose.
749        let t = self.re.mul_adde(self.re + V::TWO, self.im * self.im);
750
751        Complex::new(
752            t.ln_1p_p::<P>() * <V as FloatVector>::HALF,
753            self.im.atan2_p::<P>(self.re + V::ONE),
754        )
755    }
756
757    /// `$\log_2 z = \log_2|z| + i\arg(z)\log_2 e$`.
758    ///
759    /// Through the real `log2` rather than `ln(z) * log2(e)`, which is the same work
760    /// (one multiply fewer, in fact - the argument is scaled but `ln|z|` is not) and
761    /// picks up whatever the element's own `log2` does. Measured identical to the
762    /// rescaled form on f64, where thermite's `log2` *is* `ln * LOG2_E`; f32 has a
763    /// dedicated kernel, so no policy gate is warranted either way.
764    #[inline(always)]
765    fn log2<P: Policy>(self) -> Self {
766        let (r, theta) = self.to_polar_p::<P>();
767
768        Complex::new(r.log2_p::<P>(), theta * V::LOG2_E)
769    }
770
771    /// `$\log_{10} z = \log_{10}|z| + i\arg(z)\log_{10} e$`. See [`log2`](Self::log2).
772    #[inline(always)]
773    fn log10<P: Policy>(self) -> Self {
774        let (r, theta) = self.to_polar_p::<P>();
775
776        Complex::new(r.log10_p::<P>(), theta * V::LOG10_E)
777    }
778
779    /// `log_N(z) = ln(z) / ln(N)` for a compile-time integer base.
780    #[inline(always)]
781    fn log_n<P: Policy, const N: usize>(self) -> Self {
782        // See `nth_root`: a generic `N` rules `const_splat!` out here.
783        let ln_n = V::splat(<V::Element as FloatElement>::from_int(N as thermite::LargeInt)).ln_p::<P>();
784
785        self.ln_p::<P>() / ln_n
786    }
787
788    /// `$\ln(1 - e^{-z})$`.
789    ///
790    /// The `_ext` form lets a real vector reuse an already-computed `ln(x)`. There
791    /// is no such shortcut over C, so it forwards to the plain form.
792    #[inline(always)]
793    fn ln1m_expnx_ext<P: Policy>(self, _lnx: Self) -> Self {
794        self.ln1m_expnx_p::<P>()
795    }
796
797    // --- inverse trigonometric / hyperbolic functions ---
798
799    /// `asin(z) = -i ln(iz + sqrt(1 - z^2))`.
800    ///
801    /// Branch cuts on `(-inf, -1)` (continuous from above) and `(1, inf)`
802    /// (continuous from below); `-pi/2 <= Re(asin z) <= pi/2`.
803    #[inline(always)]
804    fn asin<P: Policy>(self) -> Self {
805        // asin(z) = -i*asinh(iz); see [`log_asinh`] for the conditioning.
806        mul_neg_i(log_asinh::<P, V>(mul_i(self)))
807    }
808
809    /// `acos(z) = -i ln(z + i sqrt(1 - z^2))`.
810    ///
811    /// Branch cuts on `(-inf, -1)` and `(1, inf)`; `0 <= Re(acos z) <= pi`.
812    #[inline(always)]
813    fn acos<P: Policy>(self) -> Self {
814        let is = mul_i(self.nmul_add(self, Self::ONE).sqrt());
815
816        mul_neg_i(ln_reciprocal_pair::<P, V>(self + is, self - is))
817    }
818
819    /// `atan(z) = (ln(1 + iz) - ln(1 - iz)) / (2i)`.
820    ///
821    /// Branch cuts on `(-inf*i, -i]` and `[i, inf*i)`; `-pi/2 <= Re(atan z) <= pi/2`.
822    #[inline(always)]
823    fn atan<P: Policy>(self) -> Self {
824        // 1 +- iz = (1 -+ Im z) +- i*Re z: two adds, no multiply by the unit.
825        let a = Complex::new(V::ONE - self.im, self.re);
826        let b = Complex::new(V::ONE + self.im, -self.re);
827
828        // z / (2i) == -0.5i * z, and the -i is the free swap-and-negate.
829        mul_neg_i(a.ln_p::<P>() - b.ln_p::<P>()) * <V as FloatVector>::HALF
830    }
831
832    /// `asinh(z) = ln(z + sqrt(1 + z^2))`.
833    #[inline(always)]
834    fn asinh<P: Policy>(self) -> Self {
835        log_asinh::<P, V>(self)
836    }
837
838    /// `acosh(z) = 2 ln(sqrt((z+1)/2) + sqrt((z-1)/2))`.
839    ///
840    /// Branch cut on `(-inf, 1)`, continuous from above.
841    #[inline(always)]
842    fn acosh<P: Policy>(self) -> Self {
843        // (z +- 1)/2 scales by a *real* half, so the complex FMA (four inner FMAs,
844        // half of them against a zero imaginary part) is one FMA and one multiply.
845        let h = <V as FloatVector>::HALF;
846        let half_im = self.im * h;
847
848        let a = Complex::new(self.re.mul_adde(h, h), half_im).sqrt();
849        let b = Complex::new(self.re.mul_sube(h, h), half_im).sqrt();
850
851        let half_res = (a + b).ln_p::<P>();
852
853        half_res + half_res
854    }
855
856    /// `atanh(z) = (ln(1 + z) - ln(1 - z)) / 2`.
857    ///
858    /// Branch cuts on `(-inf, -1]` and `[1, inf)`.
859    #[inline(always)]
860    fn atanh<P: Policy>(self) -> Self {
861        ((Self::ONE + self).ln_p::<P>() - (Self::ONE - self).ln_p::<P>()) * <V as FloatVector>::HALF
862    }
863}
864
865// --- SpecializedSpatialMath: the norms, as real-valued complex numbers ---
866
867/// The modulus of each term, for the `hypot` family.
868#[inline(always)]
869fn moduli<V: RealFloatVector, P: Policy, const N: usize>(values: [Complex<V>; N]) -> [V; N] {
870    let mut out = [V::ZERO; N];
871
872    let mut i = 0;
873    while i < N {
874        out[i] = values[i].norm_p::<P>();
875        i += 1;
876    }
877
878    out
879}
880
881impl<V: RealFloatVector> SpecializedSpatialMath<Complex<V::Element>> for Complex<V> {
882    /// `$\sqrt{\sum_i |z_i|^2}$`, as a real complex number.
883    ///
884    /// Must be overridden, and not only for tuning. The generic `hypot_n` changes
885    /// *meaning* over C depending on the precision policy: its high-precision path
886    /// opens with `abs()`, which here is the modulus, so everything after it is real
887    /// and the result is the norm - but the `PrecisionPolicy::Worst` path skips that
888    /// and squares directly, giving the analytic continuation `sqrt(sum z_i^2)`
889    /// instead. Two different functions behind one name, chosen by a policy.
890    ///
891    /// This pins the norm, matching [`l2_norm`](Self::l2_norm) and the crate docs.
892    /// Taking the modulus of each term first costs `N` extra square roots and buys
893    /// the same overflow safety the real `hypot_n` has.
894    #[inline(always)]
895    fn hypot_n<P: Policy, const N: usize>(values: [Self; N]) -> Self {
896        Self::real(<V as thermite::math::SpatialMathWithPolicy>::hypot_n_p::<P, N>(
897            moduli::<V, P, N>(values),
898        ))
899    }
900
901    /// `$1/\sqrt{\sum_i |z_i|^2}$`, as a real complex number. See [`hypot_n`](Self::hypot_n).
902    #[inline(always)]
903    fn inv_hypot_n<P: Policy, const N: usize>(values: [Self; N]) -> Self {
904        Self::real(<V as thermite::math::SpatialMathWithPolicy>::inv_hypot_n_p::<P, N>(
905            moduli::<V, P, N>(values),
906        ))
907    }
908
909    /// `|re| + |im|`, as a real complex number.
910    #[inline(always)]
911    fn l1_norm<P: Policy>(self) -> Self {
912        Self::real(self.norm_l1())
913    }
914
915    /// `$|z|^2 = z\bar{z}$`, as a real complex number.
916    #[inline(always)]
917    fn l2_norm_squared<P: Policy>(self) -> Self {
918        Self::real(self.norm_sqr())
919    }
920
921    /// The modulus `$|z|$`, as a real complex number.
922    ///
923    /// The default `sqrt(l2_norm_squared())` squares the range and overflows for
924    /// large components; `hypot` does not.
925    #[inline(always)]
926    fn l2_norm<P: Policy>(self) -> Self {
927        Self::real(self.norm_p::<P>())
928    }
929}
930
931// --- No SpecializedRealMath ---
932//
933// RealMath is the part of the math library that assumes an ordered field: atan2 (a
934// quadrant of the real plane), wrap_angle, to_degrees/to_radians, step, smoothstep
935// and its inverse, logaddexp, rescale. Their defaults are written over max, abs and
936// clamp as orderings. Over C they would return a plausible-looking result with no
937// meaning.
938//
939// Each public math trait blankets off its own specialized trait, and leaving this
940// one unimplemented costs nothing: CoreMath, TranscendentalMath and SpatialMath are
941// unaffected. The argument of a complex number is ComplexMath::arg, returning a real
942// value where atan2 would have to return a complex one.