#![allow(clippy::needless_arbitrary_self_type)]
use thermite::math::policy::{DefaultPolicy, Policy, PrecisionPolicy};
use thermite::math::specialized::{SpecializedCoreMath, SpecializedSpatialMath, SpecializedTranscendentalMath};
use thermite::prelude::*;
use crate::Complex;
use self::specialized::{ComplexVector, SpecializedComplexMath};
use crate::vector::RealFloatVector;
macro_rules! decl_complex_math {
($(
$(#[$trait_meta:meta])*
trait $trait:ident<$element:ident> $(: $($bound:ident)&+ )? { $(
$(#[$meta:meta])*
fn $name:ident [ $($generics:tt)* ][$($generic_names:ident),*]( $($arg_name:ident : $arg_ty:ty),* $(,)?) -> $ret:ty;
)*}
)*) => {paste::paste! {$(
#[doc = "" $trait " math functions with customizable policies."]
$(#[$trait_meta])*
#[doc = ""]
#[doc = "Each function takes a [`Policy`] as its first generic argument. For the"]
#[doc = "default-policy versions (same names, no `_p` suffix), see [`" $trait "Math`]."]
#[doc = ""]
#[doc = "Implemented automatically for every type implementing [`Specialized" $trait "Math`]."]
#[thermite::dispatch(Self)]
pub trait [<$trait MathWithPolicy>] $(: $($bound +)+)? {$(
$(#[$meta])* fn [<$name _p>]<P: Policy, $($generics)*>($($arg_name: $arg_ty),*) -> $ret;
)*}
#[doc = "" $trait " math functions using the default policy."]
$(#[$trait_meta])*
#[doc = ""]
#[doc = "Every method here has a counterpart in [`" $trait "MathWithPolicy`] with a `_p`"]
#[doc = "suffix that takes an explicit [`Policy`]."]
#[doc = ""]
#[doc = "Implementors of [`" $trait "MathWithPolicy`] implement this automatically."]
#[thermite::dispatch(Self)]
pub trait [<$trait Math>]: [<$trait MathWithPolicy>] {$(
$(#[$meta])* #[inline(always)] fn $name<$($generics)*>($($arg_name: $arg_ty),*) -> $ret
{ [<$trait MathWithPolicy>]::[<$name _p>]::<DefaultPolicy, $($generic_names),*>($($arg_name),*) }
)*}
impl<M> [<$trait Math>] for M where M: [<$trait MathWithPolicy>] {}
#[thermite::dispatch(Self)]
impl<E: $element, V: FloatVector<Element = E> + $($($bound +)+)?> [<$trait MathWithPolicy>] for V
where V: [<Specialized $trait Math>]<E>
{$(
$(#[$meta])* #[inline(always)] fn [<$name _p>]<P: Policy, $($generics)*>($($arg_name: $arg_ty),*) -> $ret
{ <V as [<Specialized $trait Math>]<E>>::$name::<P, $($generic_names),*>($($arg_name),*) }
)*})*
}};
}
decl_complex_math! {
trait Complex<FloatElement>: ComplexVector {
fn norm[][](self: Self) -> Self::Real;
fn arg[][](self: Self) -> Self::Real;
fn to_polar[][](self: Self) -> (Self::Real, Self::Real);
fn from_polar[][](r: Self::Real, theta: Self::Real) -> Self;
fn powfr[][](self: Self, e: Self::Real) -> Self;
fn expf[][](self: Self, base: Self::Real) -> Self;
fn logr[][](self: Self, base: Self::Real) -> Self;
fn finv[][](self: Self) -> Self;
fn fdiv[][](self: Self, rhs: Self) -> Self;
}
}
pub mod specialized;
#[cfg(feature = "special")]
pub mod special;
impl<V: RealFloatVector> SpecializedComplexMath<Complex<V::Element>> for Complex<V> {
#[inline(always)]
fn norm<P: Policy>(self) -> V {
self.re.hypot_p::<P>(self.im)
}
#[inline(always)]
fn arg<P: Policy>(self) -> V {
self.im.atan2_p::<P>(self.re)
}
#[inline(always)]
fn from_polar<P: Policy>(r: V, theta: V) -> Self {
let theta = if const { P::POLICY.check_overflow } {
theta.nz(r.cmp_eq(V::ZERO).bitandnot(theta.is_finite()))
} else {
theta
};
let (s, c) = theta.sin_cos_p::<P>();
Self::new(r * c, r * s)
}
#[inline(always)]
fn powfr<P: Policy>(self, e: V) -> Self {
let (r, theta) = self.to_polar_p::<P>();
Self::from_polar_p::<P>(r.powf_p::<P>(e), theta * e)
}
#[inline(always)]
fn expf<P: Policy>(self, base: V) -> Self {
let ln_b = finite_log_term::<P, V>(base.ln_p::<P>(), self.im);
Self::from_polar_p::<P>(base.powf_p::<P>(self.re), self.im * ln_b)
}
#[inline(always)]
fn logr<P: Policy>(self, base: V) -> Self {
let (r, theta) = self.to_polar_p::<P>();
let d = base.ln_p::<P>().reciprocal_p::<P>();
Self::new(r.ln_p::<P>() * d, theta * d)
}
#[inline(always)]
fn finv<P: Policy>(self) -> Self {
let inv = self.norm_p::<P>().reciprocal_p::<P>();
self.conj() * inv * inv
}
}
impl<V: RealFloatVector> SpecializedCoreMath<Complex<V::Element>> for Complex<V> {
#[inline(always)]
fn poly_rational<P: Policy, const N: usize, const D: usize>(
self,
numerator: &[Complex<V::Element>; N],
denominator: &[Complex<V::Element>; D],
) -> Self {
let x = self;
if const { P::POLICY.precision.le(thermite::math::policy::PrecisionPolicy::Average) } {
let n = SpecializedCoreMath::poly::<P, N>(x, numerator);
let d = SpecializedCoreMath::poly::<P, D>(x, denominator);
return n.approx_div_p::<P>(d);
}
let invert = x.norm_sqr().cmp_gt(V::ONE);
let mut n0 = Self::EMPTY;
let mut d0 = Self::EMPTY;
if const { P::POLICY.avoid_branching } || !invert.all() {
n0 = SpecializedCoreMath::poly::<P, N>(x, numerator);
d0 = SpecializedCoreMath::poly::<P, D>(x, denominator);
}
let mut z = Self::EMPTY;
let mut n1 = Self::EMPTY;
let mut d1 = Self::EMPTY;
if const { P::POLICY.avoid_branching } || invert.any() {
z = SpecializedCoreMath::reciprocal::<P>(x);
n1 = SpecializedCoreMath::poly_rev::<P, N>(z, numerator);
d1 = SpecializedCoreMath::poly_rev::<P, D>(z, denominator);
}
let n = invert.select(n1, n0);
let d = invert.select(d1, d0);
let res = n.approx_div_p::<P>(d);
if const { N == D } {
return res;
}
if const { P::POLICY.avoid_branching } || invert.any() {
let (u, e) = if const { N < D } { (z, D - N) } else { (x, N - D) };
return invert.select(res * SpecializedCoreMath::powi::<P>(u, e as i32), res);
}
res
}
#[inline(always)]
fn inverse_sqrt<P: Policy>(self) -> Self {
let s = self.sqrt();
let inv = self.norm_p::<P>().reciprocal_p::<P>();
Complex::new(s.re * inv, -(s.im * inv))
}
}
#[inline(always)]
fn mul_i<V: RealFloatVector>(z: Complex<V>) -> Complex<V> {
Complex::new(-z.im, z.re)
}
#[inline(always)]
fn mul_neg_i<V: RealFloatVector>(z: Complex<V>) -> Complex<V> {
Complex::new(z.im, -z.re)
}
#[inline(always)]
fn saturate<P: Policy, V: RealFloatVector>(res: Complex<V>, denom: V, limit: Complex<V>) -> Complex<V> {
if const { !P::POLICY.check_overflow } {
return res;
}
let lost = denom.is_infinite();
if thermite::unlikely(lost.any()) {
return lost.select(limit, res);
}
res
}
#[inline(always)]
fn log_asinh<P: Policy, V: RealFloatVector>(p: Complex<V>) -> Complex<V> {
let s = p.mul_add(p, Complex::ONE).sqrt();
let w = p + s;
if const { !P::POLICY.precision.ge(PrecisionPolicy::Best) } {
return w.ln_p::<P>();
}
let mut res = (p + p.square() / (s + Complex::ONE)).ln_1p_p::<P>();
let flip = w.norm_sqr().cmp_lt(V::ONE) & p.norm_sqr().cmp_gt(V::ONE);
if thermite::unlikely(flip.any()) {
let l = (s - p).ln_p::<P>();
res = flip.select(Complex::new(-l.re, -l.im), res);
}
res
}
#[inline(always)]
fn ln_reciprocal_pair<P: Policy, V: RealFloatVector>(w: Complex<V>, companion: Complex<V>) -> Complex<V> {
if const { !P::POLICY.precision.ge(PrecisionPolicy::Best) } {
return w.ln_p::<P>();
}
let flip = w.norm_sqr().cmp_lt(V::ONE);
let l = flip.select(companion, w).ln_p::<P>();
Complex::new(l.re.neg_c(flip), l.im.neg_c(flip))
}
#[inline(always)]
fn expm1_from<P: Policy, V: RealFloatVector>(bm1: V, phi: V) -> Complex<V> {
let (s, c) = phi.sin_cos_p::<P>();
let cm1 = phi.cos_m1_p::<P>();
Complex::new(bm1.mul_adde(c, cm1), s.mul_adde(bm1, s))
}
#[inline(always)]
fn finite_log_term<P: Policy, V: RealFloatVector>(ln_r: V, d: V) -> V {
if const { !P::POLICY.check_overflow } {
return ln_r;
}
ln_r.zz(d.cmp_ne(V::ZERO))
}
impl<V: RealFloatVector> SpecializedTranscendentalMath<Complex<V::Element>> for Complex<V> {
#[inline(always)]
fn sin_cos<P: Policy>(self) -> (Self, Self) {
let (s, c) = self.re.sin_cos_p::<P>();
let (sh, ch) = self.im.sinh_cosh_p::<P>();
(Complex::new(s * ch, c * sh), Complex::new(c * ch, -(s * sh)))
}
#[inline(always)]
fn sincos_pi<P: Policy>(self) -> (Self, Self) {
let (s, c) = self.re.sincos_pi_p::<P>();
let (sh, ch) = (self.im * <V as thermite::math::FloatConsts>::PI).sinh_cosh_p::<P>();
(Complex::new(s * ch, c * sh), Complex::new(c * ch, -(s * sh)))
}
#[inline(always)]
fn tan<P: Policy>(self) -> Self {
let (two_re, two_im) = (self.re + self.re, self.im + self.im);
let (s, c) = two_re.sin_cos_p::<P>();
let (sh, ch) = two_im.sinh_cosh_p::<P>();
let denom = c + ch;
let res = Complex::new(s, sh) / denom;
saturate::<P, V>(res, denom, Complex::new(V::ZERO, V::ONE.mul_sign(self.im)))
}
#[inline(always)]
fn sinh_cosh<P: Policy>(self) -> (Self, Self) {
let (s, c) = self.im.sin_cos_p::<P>();
let (sh, ch) = self.re.sinh_cosh_p::<P>();
(Complex::new(sh * c, ch * s), Complex::new(ch * c, sh * s))
}
#[inline(always)]
fn tanh<P: Policy>(self) -> Self {
let (two_re, two_im) = (self.re + self.re, self.im + self.im);
let (s, c) = two_im.sin_cos_p::<P>();
let (sh, ch) = two_re.sinh_cosh_p::<P>();
let denom = ch + c;
let res = Complex::new(sh, s) / denom;
saturate::<P, V>(res, denom, Complex::new(V::ONE.mul_sign(self.re), V::ZERO))
}
#[inline(always)]
fn sinc<P: Policy>(self) -> Self {
let is_zero = self.is_zero();
let q = self.sin_p::<P>() / self;
is_zero.select(Self::ONE, q)
}
#[inline(always)]
fn sinc_pi<P: Policy>(self) -> Self {
let is_zero = self.is_zero();
let q = SpecializedTranscendentalMath::sin_pi::<P>(self) / (self * <V as thermite::math::FloatConsts>::PI);
is_zero.select(Self::ONE, q)
}
#[inline(always)]
fn exp<P: Policy>(self) -> Self {
Self::from_polar_p::<P>(self.re.exp_p::<P>(), self.im)
}
#[inline(always)]
fn exph<P: Policy>(self) -> Self {
Self::from_polar_p::<P>(self.re.exph_p::<P>(), self.im)
}
#[inline(always)]
fn exp2<P: Policy>(self) -> Self {
if const { P::POLICY.precision.le(PrecisionPolicy::Average) } {
return (self * V::LN_2).exp_p::<P>();
}
Self::from_polar_p::<P>(self.re.exp2_p::<P>(), self.im * V::LN_2)
}
#[inline(always)]
fn exp10<P: Policy>(self) -> Self {
if const { P::POLICY.precision.le(PrecisionPolicy::Average) } {
return (self * V::LN_10).exp_p::<P>();
}
Self::from_polar_p::<P>(self.re.exp10_p::<P>(), self.im * V::LN_10)
}
#[inline(always)]
fn exp_m1<P: Policy>(self) -> Self {
expm1_from::<P, V>(self.re.exp_m1_p::<P>(), self.im)
}
#[inline(always)]
fn exp2_m1<P: Policy>(self) -> Self {
if const { P::POLICY.precision.le(PrecisionPolicy::Average) } {
return (self * V::LN_2).exp_m1_p::<P>();
}
expm1_from::<P, V>(self.re.exp2_m1_p::<P>(), self.im * V::LN_2)
}
#[inline(always)]
fn exp10_m1<P: Policy>(self) -> Self {
if const { P::POLICY.precision.le(PrecisionPolicy::Average) } {
return (self * V::LN_10).exp_m1_p::<P>();
}
expm1_from::<P, V>(self.re.exp10_m1_p::<P>(), self.im * V::LN_10)
}
#[inline(always)]
fn powf<P: Policy>(self, e: Self) -> Self {
let (r, theta) = self.to_polar_p::<P>();
let ln_r = r.ln_p::<P>();
let ln_r_angle = finite_log_term::<P, V>(ln_r, e.im);
let ln_r_mod = finite_log_term::<P, V>(ln_r, e.re);
let angle = e.im.mul_adde(ln_r_angle, e.re * theta);
let mut modulus = if const { P::POLICY.precision.le(PrecisionPolicy::Average) } {
e.im.nmul_adde(theta, e.re * ln_r_mod).exp_p::<P>()
} else {
r.powf_p::<P>(e.re) * (-e.im * theta).exp_p::<P>()
};
if const { P::POLICY.check_overflow && !P::POLICY.precision.le(PrecisionPolicy::Average) } {
let lost = !modulus.is_finite();
if thermite::unlikely(lost.any()) {
modulus = lost.select(e.im.nmul_adde(theta, e.re * ln_r_mod).exp_p::<P>(), modulus);
}
}
Self::from_polar_p::<P>(modulus, angle)
}
#[inline(always)]
fn cbrt<P: Policy>(self) -> Self {
let (r, theta) = self.to_polar_p::<P>();
let three: V = thermite::const_splat!(int <V::Element>: 3);
Self::from_polar_p::<P>(r.cbrt_p::<P>(), theta / three)
}
#[inline(always)]
fn nth_root<P: Policy, const N: usize>(self) -> Self {
let (r, theta) = self.to_polar_p::<P>();
let n = V::splat(<V::Element as FloatElement>::from_int(N as thermite::LargeInt));
Self::from_polar_p::<P>(r.powf_p::<P>(n.reciprocal_p::<P>()), theta / n)
}
#[inline(always)]
fn ln<P: Policy>(self) -> Self {
let (r, theta) = self.to_polar_p::<P>();
Complex::new(r.ln_p::<P>(), theta)
}
#[inline(always)]
fn ln_1p<P: Policy>(self) -> Self {
let t = self.re.mul_adde(self.re + V::TWO, self.im * self.im);
Complex::new(
t.ln_1p_p::<P>() * <V as FloatVector>::HALF,
self.im.atan2_p::<P>(self.re + V::ONE),
)
}
#[inline(always)]
fn log2<P: Policy>(self) -> Self {
let (r, theta) = self.to_polar_p::<P>();
Complex::new(r.log2_p::<P>(), theta * V::LOG2_E)
}
#[inline(always)]
fn log10<P: Policy>(self) -> Self {
let (r, theta) = self.to_polar_p::<P>();
Complex::new(r.log10_p::<P>(), theta * V::LOG10_E)
}
#[inline(always)]
fn log_n<P: Policy, const N: usize>(self) -> Self {
let ln_n = V::splat(<V::Element as FloatElement>::from_int(N as thermite::LargeInt)).ln_p::<P>();
self.ln_p::<P>() / ln_n
}
#[inline(always)]
fn ln1m_expnx_ext<P: Policy>(self, _lnx: Self) -> Self {
self.ln1m_expnx_p::<P>()
}
#[inline(always)]
fn asin<P: Policy>(self) -> Self {
mul_neg_i(log_asinh::<P, V>(mul_i(self)))
}
#[inline(always)]
fn acos<P: Policy>(self) -> Self {
let is = mul_i(self.nmul_add(self, Self::ONE).sqrt());
mul_neg_i(ln_reciprocal_pair::<P, V>(self + is, self - is))
}
#[inline(always)]
fn atan<P: Policy>(self) -> Self {
let a = Complex::new(V::ONE - self.im, self.re);
let b = Complex::new(V::ONE + self.im, -self.re);
mul_neg_i(a.ln_p::<P>() - b.ln_p::<P>()) * <V as FloatVector>::HALF
}
#[inline(always)]
fn asinh<P: Policy>(self) -> Self {
log_asinh::<P, V>(self)
}
#[inline(always)]
fn acosh<P: Policy>(self) -> Self {
let h = <V as FloatVector>::HALF;
let half_im = self.im * h;
let a = Complex::new(self.re.mul_adde(h, h), half_im).sqrt();
let b = Complex::new(self.re.mul_sube(h, h), half_im).sqrt();
let half_res = (a + b).ln_p::<P>();
half_res + half_res
}
#[inline(always)]
fn atanh<P: Policy>(self) -> Self {
((Self::ONE + self).ln_p::<P>() - (Self::ONE - self).ln_p::<P>()) * <V as FloatVector>::HALF
}
}
#[inline(always)]
fn moduli<V: RealFloatVector, P: Policy, const N: usize>(values: [Complex<V>; N]) -> [V; N] {
let mut out = [V::ZERO; N];
let mut i = 0;
while i < N {
out[i] = values[i].norm_p::<P>();
i += 1;
}
out
}
impl<V: RealFloatVector> SpecializedSpatialMath<Complex<V::Element>> for Complex<V> {
#[inline(always)]
fn hypot_n<P: Policy, const N: usize>(values: [Self; N]) -> Self {
Self::real(<V as thermite::math::SpatialMathWithPolicy>::hypot_n_p::<P, N>(
moduli::<V, P, N>(values),
))
}
#[inline(always)]
fn inv_hypot_n<P: Policy, const N: usize>(values: [Self; N]) -> Self {
Self::real(<V as thermite::math::SpatialMathWithPolicy>::inv_hypot_n_p::<P, N>(
moduli::<V, P, N>(values),
))
}
#[inline(always)]
fn l1_norm<P: Policy>(self) -> Self {
Self::real(self.norm_l1())
}
#[inline(always)]
fn l2_norm_squared<P: Policy>(self) -> Self {
Self::real(self.norm_sqr())
}
#[inline(always)]
fn l2_norm<P: Policy>(self) -> Self {
Self::real(self.norm_p::<P>())
}
}