313 lines
10 KiB
Go
313 lines
10 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package core
|
||
|
||
import (
|
||
"math"
|
||
"math/big"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/engine"
|
||
)
|
||
|
||
// Classical orthogonal polynomials on the real line, each element-wise
|
||
// over real arrays with integer degree or order. The recurrences are
|
||
// the stable ones for every family, and the normalisations follow the
|
||
// standard physics conventions: Hermite physicists' Hₙ with leading
|
||
// coefficient 2ⁿ, Laguerre Lₙ (α = 0) and generalised Lₙ^α, Chebyshev
|
||
// of the first and second kind Tₙ and Uₙ on [−1, 1].
|
||
|
||
// Hermite returns the physicists' Hermite polynomial Hₙ(x) of each
|
||
// element, by the recurrence Hₙ₊₁ = 2x·Hₙ − 2n·Hₙ₋₁.
|
||
func Hermite(n int, x *Array) (*Array, error) {
|
||
if n < 0 {
|
||
return nil, errf("Hermite: degree must be ≥ 0, got %d", n)
|
||
}
|
||
if x.dt == Complex {
|
||
return nil, errf("Hermite: complex arrays are not supported")
|
||
}
|
||
out := &Array{shape: append([]int{}, x.shape...), dt: Float}
|
||
out.alloc(x.Len())
|
||
engine.Parallel(x.Len(), func(s, e int) {
|
||
for i := s; i < e; i++ {
|
||
xv := x.floatAt(i)
|
||
h0, h1 := 1.0, 2*xv
|
||
for k := 1; k < n; k++ {
|
||
h0, h1 = h1, 2*xv*h1-2*float64(k)*h0
|
||
}
|
||
if n == 0 {
|
||
h1 = 1
|
||
}
|
||
out.floats[i] = h1
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// Laguerre returns the generalised Laguerre polynomial Lₙ^α(x) of each
|
||
// element, by the recurrence (n+1)Lₙ₊₁ = (2n+1+α−x)Lₙ − (n+α)Lₙ₋₁.
|
||
// The degree n must be ≥ 0.
|
||
func Laguerre(n int, alpha float64, x *Array) (*Array, error) {
|
||
if n < 0 {
|
||
return nil, errf("Laguerre: degree must be ≥ 0, got %d", n)
|
||
}
|
||
if x.dt == Complex {
|
||
return nil, errf("Laguerre: complex arrays are not supported")
|
||
}
|
||
if n == 0 {
|
||
out := &Array{shape: append([]int{}, x.shape...), dt: Float}
|
||
out.alloc(x.Len())
|
||
for i := range x.Len() {
|
||
out.floats[i] = 1
|
||
}
|
||
return out, nil
|
||
}
|
||
out := &Array{shape: append([]int{}, x.shape...), dt: Float}
|
||
out.alloc(x.Len())
|
||
engine.Parallel(x.Len(), func(s, e int) {
|
||
for i := s; i < e; i++ {
|
||
xv := x.floatAt(i)
|
||
l0, l1 := 1.0, 1.0+float64(alpha)-xv
|
||
for k := 1; k < n; k++ {
|
||
kf := float64(k)
|
||
l0, l1 = l1, ((2*kf+1+alpha-xv)*l1-(kf+alpha)*l0)/(kf+1)
|
||
}
|
||
out.floats[i] = l1
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// ChebyshevT returns the Chebyshev polynomial of the first kind
|
||
// Tₙ(x) = cos(n·arccos x) of each element, by the recurrence
|
||
// Tₙ₊₁ = 2x·Tₙ − Tₙ₋₁.
|
||
func ChebyshevT(n int, x *Array) (*Array, error) {
|
||
if n < 0 {
|
||
return nil, errf("ChebyshevT: degree must be ≥ 0, got %d", n)
|
||
}
|
||
if x.dt == Complex {
|
||
return nil, errf("ChebyshevT: complex arrays are not supported")
|
||
}
|
||
out := &Array{shape: append([]int{}, x.shape...), dt: Float}
|
||
out.alloc(x.Len())
|
||
engine.Parallel(x.Len(), func(s, e int) {
|
||
for i := s; i < e; i++ {
|
||
xv := x.floatAt(i)
|
||
t0, t1 := 1.0, xv
|
||
for k := 1; k < n; k++ {
|
||
t0, t1 = t1, 2*xv*t1-t0
|
||
}
|
||
if n == 0 {
|
||
t1 = 1
|
||
}
|
||
out.floats[i] = t1
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// ChebyshevU returns the Chebyshev polynomial of the second kind
|
||
// Uₙ(x) of each element, by the recurrence Uₙ₊₁ = 2x·Uₙ − Uₙ₋₁ with
|
||
// U₀ = 1, U₁ = 2x.
|
||
func ChebyshevU(n int, x *Array) (*Array, error) {
|
||
if n < 0 {
|
||
return nil, errf("ChebyshevU: degree must be ≥ 0, got %d", n)
|
||
}
|
||
if x.dt == Complex {
|
||
return nil, errf("ChebyshevU: complex arrays are not supported")
|
||
}
|
||
out := &Array{shape: append([]int{}, x.shape...), dt: Float}
|
||
out.alloc(x.Len())
|
||
engine.Parallel(x.Len(), func(s, e int) {
|
||
for i := s; i < e; i++ {
|
||
xv := x.floatAt(i)
|
||
u0, u1 := 1.0, 2*xv
|
||
for k := 1; k < n; k++ {
|
||
u0, u1 = u1, 2*xv*u1-u0
|
||
}
|
||
if n == 0 {
|
||
u1 = 1
|
||
}
|
||
out.floats[i] = u1
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// SphericalHarmonicReal returns the real spherical harmonic
|
||
// Y_l^m(θ, φ) built from the complex one, both branches under the
|
||
// standard convention Y_real = √2·(−1)^m·Re(Y_l^m) for m > 0 and
|
||
// Y_real = √2·(−1)^m·Im(Y_l^|m|) for m < 0, and for m = 0 the plain
|
||
// Y_l⁰. Values are real and element-wise over theta and phi, which
|
||
// must share a shape.
|
||
func SphericalHarmonicReal(l, m int, theta, phi *Array) (*Array, error) {
|
||
// m = 0 is purely real in the complex form, but the contract here
|
||
// is a Float array, so the value is extracted rather than handing
|
||
// the complex Y_l⁰ through unchanged.
|
||
yc, err := SphericalHarmonic(l, absInt(m), theta, phi)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
out := &Array{shape: append([]int{}, theta.shape...), dt: Float}
|
||
out.alloc(theta.Len())
|
||
root2 := math.Sqrt2
|
||
// The standard real forms carry the (−1)^m phase on both sides of
|
||
// zero, mirroring the Condon-Shortley factor the complex harmonics
|
||
// carry.
|
||
sign := 1.0
|
||
if m%2 != 0 {
|
||
sign = -1
|
||
}
|
||
for i := range theta.Len() {
|
||
z := yc.complexAt(i)
|
||
switch {
|
||
case m == 0:
|
||
out.floats[i] = real(z)
|
||
case m > 0:
|
||
out.floats[i] = sign * root2 * real(z)
|
||
default:
|
||
out.floats[i] = sign * root2 * imag(z)
|
||
}
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// Airy returns the Airy functions of the first and second kind at
|
||
// each element, Ai in the first return slot and Bi in the second.
|
||
// The power series about zero is used, which converges for every
|
||
// real x but is numerically dependable only within |x| ≤ 8; outside
|
||
// that window the functions return NaN rather than a silently wrong
|
||
// value. Positive arguments above 2 sum the series in extended
|
||
// precision, where the float64 reconstruction of Ai would cancel the
|
||
// growing solution and lose digits. The scaling convention is the
|
||
// standard Ai, Bi with Ai(0) = 1/(3^{2/3}Γ(2/3)),
|
||
// Bi(0) = 1/(3^{1/6}Γ(1/3)).
|
||
func Airy(a *Array) (ai, bi *Array, err error) {
|
||
if a.dt == Complex {
|
||
return nil, nil, errf("Airy: complex arrays are not supported")
|
||
}
|
||
ai = &Array{shape: append([]int{}, a.shape...), dt: Float}
|
||
ai.alloc(a.Len())
|
||
bi = &Array{shape: append([]int{}, a.shape...), dt: Float}
|
||
bi.alloc(a.Len())
|
||
engine.Parallel(a.Len(), func(s, e int) {
|
||
for i := s; i < e; i++ {
|
||
x := a.floatAt(i)
|
||
if math.Abs(x) > 8 {
|
||
ai.floats[i] = math.NaN()
|
||
bi.floats[i] = math.NaN()
|
||
continue
|
||
}
|
||
aiv, biv := airySeries(x)
|
||
ai.floats[i] = aiv
|
||
bi.floats[i] = biv
|
||
}
|
||
})
|
||
return ai, bi, nil
|
||
}
|
||
|
||
// airySeries evaluates the Maclaurin pair
|
||
// Ai(x) = c1·f(x) − c2·g(x), Bi(x) = √3·(c1·f(x) + c2·g(x))
|
||
// with f = Σ 3ᵏ(1/3)ₖx^{3k}/(3k)!, g = Σ 3ᵏ(2/3)ₖx^{3k+1}/(3k+1)!,
|
||
// summed by the term-ratio recurrence to machine precision. For
|
||
// positive arguments above airyBigCut the Ai reconstruction runs in
|
||
// extended precision instead: Bi is a sum of like-signed terms and
|
||
// keeps the float64 series.
|
||
func airySeries(x float64) (aiv, biv float64) {
|
||
const (
|
||
c1 = 0.35502805388781723926
|
||
c2 = 0.25881940379280679840
|
||
)
|
||
tf, tg := 1.0, x
|
||
sumF, sumG := 1.0, x
|
||
for k := range 400 {
|
||
tf *= 3 * x * x * x * (float64(k) + 1.0/3) / float64((3*k+1)*(3*k+2)*(3*k+3))
|
||
tg *= 3 * x * x * x * (float64(k) + 2.0/3) / float64((3*k+2)*(3*k+3)*(3*k+4))
|
||
sumF += tf
|
||
sumG += tg
|
||
if math.Abs(tf) < 1e-18*math.Abs(sumF) && math.Abs(tg) < 1e-18*math.Abs(sumG) {
|
||
break
|
||
}
|
||
}
|
||
biv = math.Sqrt(3) * (c1*sumF + c2*sumG)
|
||
if x > airyBigCut {
|
||
aiv = airyBig(x)
|
||
return aiv, biv
|
||
}
|
||
aiv = c1*sumF - c2*sumG
|
||
return aiv, biv
|
||
}
|
||
|
||
// airyBigCut is the positive argument above which Ai is reconstructed
|
||
// in extended precision. At the cut the float64 difference c1·f − c2·g
|
||
// has lost 2ζ/ln2 ≈ 5 bits to the cancellation, which caps its
|
||
// relative accuracy at about 1e-14; below the cut the float64 series
|
||
// still holds fourteen digits.
|
||
const airyBigCut = 2.0
|
||
|
||
// The Maclaurin constants of the extended-precision branch, Ai(0) and
|
||
// −Ai'(0) to 50 significant digits: at the working sizes the
|
||
// cancellation needs, the float64 roundings of airySeries would
|
||
// reintroduce the very error the branch removes. The regression test
|
||
// pins their float64 roundings against the float64 constants above.
|
||
const (
|
||
airyC1Digits = "0.35502805388781723926006318600418317639797917419918"
|
||
airyC2Digits = "0.25881940379280679840518356018920396347909113835493"
|
||
)
|
||
|
||
// airyBig evaluates Ai(x) for x > 0 in extended precision. The
|
||
// reconstruction c1·f − c2·g cancels e^{2ζ}, ζ = 2x^{3/2}/3: both
|
||
// products carry the growing solution e^{ζ} while Ai itself is the
|
||
// decaying e^{−ζ} one, so the float64 difference loses 2ζ/ln2 ≈
|
||
// 1.92·x^{3/2} bits (4.5e-3 relative at x = 8). The same two series are
|
||
// therefore summed in math/big with a working size scaled to that loss,
|
||
// following the fresnelBig pattern; the term ratios here are the
|
||
// simplified forms x³/((3k+2)(3k+3)) and x³/((3k+3)(3k+4)) the float64
|
||
// recurrence above evaluates with the (3k+1) and (3k+2) factors still
|
||
// in place.
|
||
func airyBig(x float64) float64 {
|
||
bits := uint(64 + int(1.9235*x*math.Sqrt(x)) + 96)
|
||
xf := new(big.Float).SetPrec(bits).SetFloat64(x)
|
||
x3 := new(big.Float).SetPrec(bits).Mul(xf, xf)
|
||
x3.Mul(x3, xf)
|
||
c1, _, err := big.ParseFloat(airyC1Digits, 10, bits, big.ToNearestEven)
|
||
if err != nil {
|
||
return math.NaN()
|
||
}
|
||
c2, _, err := big.ParseFloat(airyC2Digits, 10, bits, big.ToNearestEven)
|
||
if err != nil {
|
||
return math.NaN()
|
||
}
|
||
tf := new(big.Float).SetPrec(bits).SetInt64(1)
|
||
tg := new(big.Float).SetPrec(bits).Set(xf)
|
||
sumF := new(big.Float).SetPrec(bits).Set(tf)
|
||
sumG := new(big.Float).SetPrec(bits).Set(tg)
|
||
// The terms are all positive for x > 0, so a cut in absolute size
|
||
// relative to a sum that stays above 1 is a relative cut.
|
||
tiny := new(big.Float).SetPrec(bits).SetMantExp(big.NewFloat(1), -int(bits)+20)
|
||
for k := range 400 {
|
||
r1 := new(big.Float).SetPrec(bits).Quo(x3, big.NewFloat(float64((3*k+2)*(3*k+3))))
|
||
r2 := new(big.Float).SetPrec(bits).Quo(x3, big.NewFloat(float64((3*k+3)*(3*k+4))))
|
||
tf.Mul(tf, r1)
|
||
tg.Mul(tg, r2)
|
||
sumF.Add(sumF, tf)
|
||
sumG.Add(sumG, tg)
|
||
if tf.Cmp(tiny) < 0 && tg.Cmp(tiny) < 0 {
|
||
break
|
||
}
|
||
}
|
||
term := new(big.Float).SetPrec(bits).Mul(c2, sumG)
|
||
aiv := new(big.Float).SetPrec(bits).Mul(c1, sumF)
|
||
aiv.Sub(aiv, term)
|
||
out, _ := aiv.Float64()
|
||
return out
|
||
}
|
||
|
||
// absInt returns |n|.
|
||
func absInt(n int) int {
|
||
if n < 0 {
|
||
return -n
|
||
}
|
||
return n
|
||
}
|