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.