treams_core/numerics/
mod.rs

1//! Numerical support without physics: complex helpers, forward-mode jets, broadcasting,
2//! fallible allocation and the parallel policy. treams-rs extension.
3
4pub(crate) mod broadcast;
5pub(crate) mod jet;
6mod memory;
7pub(crate) mod parallel;
8
9pub(crate) use jet::Jet;
10pub(crate) use memory::{filled, reserve, zeros};
11
12use crate::Complex;
13
14/// The Euler-Mascheroni constant.
15pub(crate) const EULER: f64 = std::f64::consts::EULER_GAMMA;
16
17/// The principal square root of `z`, equal to `z.sqrt()` up to rounding.
18///
19/// The algebraic form avoids the `atan2`, `sin` and `cos` of num-complex's polar form.
20/// Inputs whose larger part lies outside `[1e-150, 1e150]` are scaled first, so the
21/// squares neither overflow nor underflow. The sign of a zero imaginary part picks the
22/// side of the branch cut on the negative real axis: `complex_sqrt(-1 - 0i) = -i`.
23#[inline]
24pub(crate) fn complex_sqrt(z: Complex) -> Complex {
25    if z.im == 0.0 || !finite(z) {
26        return z.sqrt();
27    }
28    let scale = z.re.abs().max(z.im.abs());
29    if !(1e-150..=1e150).contains(&scale) {
30        return complex_sqrt(z / scale) * scale.sqrt();
31    }
32    let norm = z.re.mul_add(z.re, z.im * z.im).sqrt();
33    let dominant = norm.midpoint(z.re.abs()).sqrt();
34    if z.re >= 0.0 {
35        Complex::new(dominant, z.im / (2.0 * dominant))
36    } else {
37        Complex::new(z.im.abs() / (2.0 * dominant), dominant.copysign(z.im))
38    }
39}
40
41/// `(-1)^n` as a float, for the sign factors of degrees and orders.
42#[inline]
43pub(crate) fn parity(n: i32) -> f64 {
44    if n % 2 == 0 { 1.0 } else { -1.0 }
45}
46
47/// The bits of a real mode label, such as an axial wavenumber, as an exact hash key.
48/// Adding zero turns -0.0 into +0.0, so the two zeros, which compare equal, share a key.
49pub(crate) fn label_bits(x: f64) -> u64 {
50    (x + 0.0).to_bits()
51}
52
53/// Whether both parts of `z` are finite.
54pub(crate) fn finite(z: Complex) -> bool {
55    z.re.is_finite() && z.im.is_finite()
56}
57
58/// `numerator / denominator` without overflow or underflow in the denominator's square.
59///
60/// num-complex divides by `|denominator|^2`, which overflows from about `1e154` and
61/// underflows below about `1e-154`. Dividing both operands by the larger part of the
62/// denominator first keeps that square near one.
63pub(crate) fn ratio(numerator: Complex, denominator: Complex) -> Complex {
64    let scale = denominator.re.abs().max(denominator.im.abs());
65    (numerator / scale) / (denominator / scale)
66}
67
68/// Extreme scales and branch sides of the complex helpers.
69#[cfg(test)]
70mod tests {
71    use super::{complex_sqrt, ratio};
72    use crate::Complex;
73
74    #[test]
75    fn principal_square_root_preserves_scales_and_branch_sides() {
76        for exponent in [-300, -155, -149, 0, 149, 155, 300] {
77            for re in [-0.7, 0.0, 0.7] {
78                for im in [-0.4, -0.0, 0.0, 0.4] {
79                    let z = Complex::new(re, im) * 10.0_f64.powi(exponent);
80                    let actual = complex_sqrt(z);
81                    let expected = z.sqrt();
82                    assert!((actual - expected).norm() <= 3e-14 * expected.norm());
83                    assert_eq!(actual.im.is_sign_negative(), expected.im.is_sign_negative());
84                    assert!((actual * actual - z).norm() <= 3e-14 * z.norm());
85                }
86            }
87        }
88        for im in [-1e-300, 1e-300] {
89            let z = Complex::new(-1.0, im);
90            let root = complex_sqrt(z);
91            assert!((root * root - z).im.abs() < 1e-313);
92            assert!(root.re > 0.0);
93        }
94    }
95    #[test]
96    fn complex_division_preserves_extreme_scales() {
97        let scales = [-300, -155, -149, -1, 0, 1, 149, 155, 300];
98        let a = Complex::new(0.3, -0.2);
99        let b = Complex::new(0.4, 0.1);
100        for n in scales {
101            for d in scales {
102                let difference: i32 = n - d;
103                if difference.abs() > 300 {
104                    continue;
105                }
106                let expected = (a / b) * 10.0_f64.powi(difference);
107                let actual = ratio(a * 10.0_f64.powi(n), b * 10.0_f64.powi(d));
108                assert!(
109                    (actual - expected).norm() <= 3e-14 * expected.norm(),
110                    "{n}/{d}: {actual} != {expected}"
111                );
112            }
113        }
114    }
115}
treams_core/numerics/mod.rs

treams_core/numerics/
mod.rs
