411 lines
13 KiB
Go
411 lines
13 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package core
|
||
|
||
import (
|
||
"math"
|
||
"math/cmplx"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/engine"
|
||
)
|
||
|
||
// Special functions for physics, cosmology and chemistry: the gamma
|
||
// and beta families, the error function family, Legendre polynomials,
|
||
// spherical harmonics and spherical Bessel functions. Element-wise
|
||
// functions run on real arrays and promote ints and float32 per the
|
||
// ladder; complex inputs are refused. Values outside a function's
|
||
// domain follow the IEEE behaviour of the underlying implementation
|
||
// (math.Gamma, for example, returns +Inf at non-positive integers and
|
||
// NaN where undefined).
|
||
|
||
// Gamma returns the gamma function Γ(x) of each element.
|
||
func Gamma(a *Array) (*Array, error) { return a.realFunc("Gamma", math.Gamma) }
|
||
|
||
// LnGamma returns the natural logarithm of |Γ(x)| of each element. The
|
||
// sign of Γ for negative arguments is dropped; callers that need it
|
||
// evaluate Gamma directly.
|
||
func LnGamma(a *Array) (*Array, error) {
|
||
return a.realFunc("LnGamma", func(x float64) float64 {
|
||
lg, _ := math.Lgamma(x)
|
||
return lg
|
||
})
|
||
}
|
||
|
||
// Beta returns the Euler beta function B(x, y) = Γ(x)Γ(y)/Γ(x+y)
|
||
// element-wise. Both arrays must have the same shape.
|
||
func Beta(x, y *Array) (*Array, error) {
|
||
if !sameShape(x.shape, y.shape) {
|
||
return nil, errf("Beta: shape mismatch %s vs %s", shapeText(x.shape), shapeText(y.shape))
|
||
}
|
||
if x.dt == Complex || y.dt == Complex {
|
||
return nil, errf("Beta: 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, yv := x.floatAt(i), y.floatAt(i)
|
||
// In log space: the Γ product overflows to Inf well before
|
||
// B itself leaves the float64 range (B(100, 80) ≈ 3e-67),
|
||
// and Inf/Inf would turn the answer into NaN.
|
||
lx, sx := math.Lgamma(xv)
|
||
ly, sy := math.Lgamma(yv)
|
||
lxy, sxy := math.Lgamma(xv + yv)
|
||
sign := sx * sy * sxy
|
||
out.floats[i] = float64(sign) * math.Exp(lx+ly-lxy)
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// Erf and Erfc return the error function and the complementary error
|
||
// function of each element.
|
||
func Erf(a *Array) (*Array, error) { return a.realFunc("Erf", math.Erf) }
|
||
func Erfc(a *Array) (*Array, error) { return a.realFunc("Erfc", math.Erfc) }
|
||
|
||
// Sinc returns the normalised sinc sin(πx)/(πx) of each element, with
|
||
// Sinc(0) = 1. The π-normalised convention is the one signal
|
||
// processing and interpolation use.
|
||
func Sinc(a *Array) (*Array, error) {
|
||
return a.realFunc("Sinc", func(x float64) float64 {
|
||
if x == 0 {
|
||
return 1
|
||
}
|
||
return math.Sin(math.Pi*x) / (math.Pi * x)
|
||
})
|
||
}
|
||
|
||
// Cosm1 returns cos(x) − 1 of each element, kept accurate where the
|
||
// direct subtraction loses: as a float64 subtraction cos(x) − 1 has
|
||
// no correct significant bit under about |x| ≈ 1e-8, while the factored
|
||
// even series in x² holds full precision down to the smallest
|
||
// representable argument, where the answer is exactly −x²/2 as far as
|
||
// the format can see it. The standard library carries log1p and expm1
|
||
// for the logarithmic and exponential side of the same problem and no
|
||
// cosinusoidal counterpart, which is the gap this fills.
|
||
func Cosm1(a *Array) (*Array, error) {
|
||
return a.realFunc("Cosm1", cosm1)
|
||
}
|
||
|
||
// cosm1Crossover splits the series and the direct subtraction: below
|
||
// it the factored series ends well past the rounding floor, above it
|
||
// the cancellation the subtraction suffers still leaves every bit the
|
||
// answer has.
|
||
const cosm1Crossover = math.Pi / 4
|
||
|
||
// cosm1 evaluates cos(x) − 1 at one point. With t = x² the series
|
||
// cos x − 1 = −(t/2)·(1 − t/12·(1 − t/30·(1 − t/56·(1 − t/90·
|
||
// (1 − t/132·(1 − t/182·(1 − t/240))))))) truncates around 2e-18
|
||
// relative at the crossover and falls quadratically below it; the
|
||
// leading −t/2 factor carries the sign, so negative arguments need no
|
||
// branch.
|
||
func cosm1(x float64) float64 {
|
||
t := x * x
|
||
if math.Abs(x) >= cosm1Crossover {
|
||
return math.Cos(x) - 1
|
||
}
|
||
return -0.5 * t * (1 - t/12*(1-t/30*(1-t/56*(1-t/90*(1-t/132*(1-t/182*(1-t/240)))))))
|
||
}
|
||
|
||
// Legendre returns the Legendre polynomial P_l(x) of each element,
|
||
// computed by the Bonnet recurrence. The degree l must be ≥ 0; the
|
||
// convention is P_0 = 1, P_1 = x, no extra scaling.
|
||
func Legendre(l int, x *Array) (*Array, error) {
|
||
if l < 0 {
|
||
return nil, errf("Legendre: degree must be ≥ 0, got %d", l)
|
||
}
|
||
if x.dt == Complex {
|
||
return nil, errf("Legendre: 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)
|
||
p0, p1 := 1.0, xv
|
||
for k := 1; k < l; k++ {
|
||
p0, p1 = p1, ((2*float64(k)+1)*xv*p1-float64(k)*p0)/float64(k+1)
|
||
}
|
||
if l == 0 {
|
||
p1 = 1
|
||
}
|
||
out.floats[i] = p1
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// LegendreAssociated returns the associated Legendre function
|
||
// P_l^m(x) of each element, with the Condon-Shortley phase (−1)^m
|
||
// folded in and no extra normalisation. Requires |m| ≤ l; a negative m
|
||
// uses the standard relation P_l^(−m) = (−1)^m (l−m)!/(l+m)! P_l^m.
|
||
func LegendreAssociated(l, m int, x *Array) (*Array, error) {
|
||
if l < 0 {
|
||
return nil, errf("LegendreAssociated: degree must be ≥ 0, got %d", l)
|
||
}
|
||
if m < -l || m > l {
|
||
return nil, errf("LegendreAssociated: order %d out of range for degree %d", m, l)
|
||
}
|
||
if x.dt == Complex {
|
||
return nil, errf("LegendreAssociated: complex arrays are not supported")
|
||
}
|
||
am, neg := m, false
|
||
if am < 0 {
|
||
am, neg = -m, true
|
||
}
|
||
mm := float64(am)
|
||
// Seed sign from the Condon-Shortley phase (−1)^m.
|
||
seedSign := 1.0
|
||
if am%2 == 1 {
|
||
seedSign = -1
|
||
}
|
||
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)
|
||
// P_m^m = (−1)^m (2m−1)!! (1−x²)^(m/2).
|
||
fact := 1.0
|
||
for k := 1; k <= am; k++ {
|
||
fact *= float64(2*k - 1)
|
||
}
|
||
pmm := seedSign * fact * math.Pow(1-xv*xv, mm/2)
|
||
if l == am {
|
||
out.floats[i] = pmm
|
||
continue
|
||
}
|
||
// P_{m+1}^m = x (2m+1) P_m^m, then the fixed-order
|
||
// recurrence up the degree.
|
||
pmp1 := xv * (2*mm + 1) * pmm
|
||
pl := pmp1
|
||
for k := am + 2; k <= l; k++ {
|
||
pl = ((2*float64(k)-1)*xv*pmp1 - float64(k+am-1)*pmm) / float64(k-am)
|
||
pmm, pmp1 = pmp1, pl
|
||
}
|
||
out.floats[i] = pl
|
||
}
|
||
})
|
||
if neg {
|
||
// P_l^(−m) = (−1)^m (l−m)!/(l+m)! P_l^m.
|
||
ratio := math.Exp(LnFactorial(l-am) - LnFactorial(l+am))
|
||
engine.Parallel(out.Len(), func(s, e int) {
|
||
for i := s; i < e; i++ {
|
||
out.floats[i] *= seedSign * ratio
|
||
}
|
||
})
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// LnFactorial returns the natural logarithm of n!.
|
||
func LnFactorial(n int) float64 {
|
||
if n < 0 {
|
||
return math.NaN()
|
||
}
|
||
lg, _ := math.Lgamma(float64(n + 1))
|
||
return lg
|
||
}
|
||
|
||
// SphericalHarmonic returns the complex spherical harmonic
|
||
// Y_l^m(θ, φ) element-wise over the polar angle θ and azimuth φ,
|
||
// using the Condon-Shortley phase and the orthonormal convention
|
||
// N·P_l^m(cos θ)·e^{imφ} with N = √((2l+1)/(4π)·(l−m)!/(l+m)!).
|
||
// theta and phi must have the same shape.
|
||
func SphericalHarmonic(l, m int, theta, phi *Array) (*Array, error) {
|
||
if l < 0 {
|
||
return nil, errf("SphericalHarmonic: degree must be ≥ 0, got %d", l)
|
||
}
|
||
if m < -l || m > l {
|
||
return nil, errf("SphericalHarmonic: order %d out of range for degree %d", m, l)
|
||
}
|
||
if !sameShape(theta.shape, phi.shape) {
|
||
return nil, errf("SphericalHarmonic: shape mismatch %s vs %s",
|
||
shapeText(theta.shape), shapeText(phi.shape))
|
||
}
|
||
if theta.dt == Complex || phi.dt == Complex {
|
||
return nil, errf("SphericalHarmonic: complex arrays are not supported")
|
||
}
|
||
am, conj := m, false
|
||
if am < 0 {
|
||
am, conj = -m, true
|
||
}
|
||
// P_l^m takes the cosine of the polar angle, not the angle.
|
||
cosTheta, err := Cos(theta)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
p, err := LegendreAssociated(l, am, cosTheta)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
norm := math.Sqrt((2*float64(l) + 1) / (4 * math.Pi) *
|
||
math.Exp(LnFactorial(l-am)-LnFactorial(l+am)))
|
||
// The Condon-Shortley phase (−1)^m is already folded into the
|
||
// associated Legendre function; the normalisation carries no
|
||
// extra sign.
|
||
out := &Array{shape: append([]int{}, theta.shape...), dt: Complex}
|
||
out.alloc(theta.Len())
|
||
engine.Parallel(theta.Len(), func(s, e int) {
|
||
for i := s; i < e; i++ {
|
||
y := complex(norm*p.floatAt(i), 0) *
|
||
cmplx.Exp(complex(0, float64(am)*phi.floatAt(i)))
|
||
if conj {
|
||
// Y_{l,−m} = (−1)^m conj(Y_{l,m}).
|
||
phase := 1.0
|
||
if am%2 == 1 {
|
||
phase = -1
|
||
}
|
||
y = complex(phase, 0) * cmplx.Conj(y)
|
||
}
|
||
out.complexes[i] = y
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// SphericalBesselJ returns the spherical Bessel function of the first
|
||
// kind j_l(x) of each element. Orders at or below |x| climb the upward
|
||
// recurrence from the exact closed forms j₀ = sin x/x and
|
||
// j₁ = sin x/x² − cos x/x, the stable direction there at O(l) cost;
|
||
// higher orders run the downward Miller recurrence, which is what the
|
||
// small roots need and whose start sits above the turning point at
|
||
// order l + |x|.
|
||
func SphericalBesselJ(l int, x *Array) (*Array, error) {
|
||
if l < 0 {
|
||
return nil, errf("SphericalBesselJ: degree must be ≥ 0, got %d", l)
|
||
}
|
||
if x.dt == Complex {
|
||
return nil, errf("SphericalBesselJ: 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++ {
|
||
out.floats[i] = sphericalBesselJ(l, x.floatAt(i))
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// SphericalBesselY returns the spherical Bessel function of the second
|
||
// kind y_l(x) of each element, evaluated by the upward recurrence from
|
||
// y_0 = −cos(x)/x, the stable direction. Every y_l diverges to −Inf at
|
||
// x = 0.
|
||
func SphericalBesselY(l int, x *Array) (*Array, error) {
|
||
if l < 0 {
|
||
return nil, errf("SphericalBesselY: degree must be ≥ 0, got %d", l)
|
||
}
|
||
if x.dt == Complex {
|
||
return nil, errf("SphericalBesselY: 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)
|
||
if xv == 0 {
|
||
// Divergence at the origin: fill this element and carry
|
||
// on, the rest of the chunk still needs its values.
|
||
out.floats[i] = math.Inf(-1)
|
||
continue
|
||
}
|
||
y0 := -math.Cos(xv) / xv
|
||
if l == 0 {
|
||
// Assign and carry on: a return here would abandon the
|
||
// rest of the worker's chunk at zero.
|
||
out.floats[i] = y0
|
||
continue
|
||
}
|
||
y1 := -math.Cos(xv)/(xv*xv) - math.Sin(xv)/xv
|
||
for k := 1; k < l; k++ {
|
||
y0, y1 = y1, (2*float64(k)+1)/xv*y1-y0
|
||
}
|
||
out.floats[i] = y1
|
||
}
|
||
})
|
||
return out, nil
|
||
}
|
||
|
||
// sphericalBesselJ evaluates j_l at one point. Below the argument the
|
||
// upward recurrence climbs from the exact j₀ and j₁ at O(l) cost; at
|
||
// and above it the downward Miller recurrence, started well above the
|
||
// turning point, is the stable direction.
|
||
func sphericalBesselJ(l int, x float64) float64 {
|
||
if x == 0 {
|
||
if l == 0 {
|
||
return 1
|
||
}
|
||
return 0
|
||
}
|
||
if math.Abs(x) < 1e-4 {
|
||
// Small-x power series. Below ~1e-6 the unscaled Miller seed
|
||
// overflows on the way down and the renormalisation turns into
|
||
// Inf/Inf = NaN; the two-term series is relative ~1e-8 exact
|
||
// over the whole branch and underflows to the true limit for
|
||
// large l.
|
||
num := math.Pow(x, float64(l))
|
||
dbl := 1.0 // (2l+1)!!
|
||
for k := 3; k <= 2*l+1; k += 2 {
|
||
dbl *= float64(k)
|
||
}
|
||
return num / dbl * (1 - x*x/(2*float64(2*l+3)))
|
||
}
|
||
if float64(l) <= math.Abs(x) {
|
||
return sphericalJUpward(l, x)
|
||
}
|
||
// |x| keeps the start above l for negative arguments too: a start
|
||
// at or below l would never pass the order on the way down and
|
||
// renormalise against a garbage seed.
|
||
ceiling := l + int(math.Abs(x)) + 40
|
||
jp1, j := 0.0, 1.0 // j_{k+1}, j_k, seeded at k = ceiling
|
||
jl := 0.0
|
||
for k := ceiling; k >= 1; k-- {
|
||
if k == l {
|
||
jl = j
|
||
}
|
||
jp1, j = j, (2*float64(k)+1)/x*j-jp1
|
||
if aj := math.Abs(j); aj > 1e200 {
|
||
// The unscaled walk overflows for degree-argument
|
||
// combinations whose true value is representable; the
|
||
// renormalisation cancels any common factor.
|
||
jp1 /= aj
|
||
j /= aj
|
||
jl /= aj
|
||
}
|
||
}
|
||
// j now holds the unscaled j_0, which for l = 0 is the answer.
|
||
if l == 0 {
|
||
jl = j
|
||
}
|
||
return sphericalJ0(x) * jl / j
|
||
}
|
||
|
||
// sphericalJUpward climbs the spherical Bessel functions of the first
|
||
// kind from the exact closed forms j₀ = sin x/x and
|
||
// j₁ = sin x/x² − cos x/x by the upward recurrence
|
||
// jₖ₊₁ = (2k+1)/x·jₖ − jₖ₋₁, the stable direction while the order
|
||
// stays at or below the argument. Both seeds carry the absolute scale,
|
||
// so no renormalisation is needed, and they carry the parity for free:
|
||
// the recurrence reproduces j_l(−x) = (−1)^l j_l(x). The caller
|
||
// guarantees x ≠ 0 and 0 ≤ l ≤ |x|.
|
||
func sphericalJUpward(l int, x float64) float64 {
|
||
j0 := sphericalJ0(x)
|
||
j1 := math.Sin(x)/(x*x) - math.Cos(x)/x
|
||
if l == 0 {
|
||
return j0
|
||
}
|
||
for k := 1; k < l; k++ {
|
||
j0, j1 = j1, (2*float64(k)+1)/x*j1-j0
|
||
}
|
||
return j1
|
||
}
|
||
|
||
// sphericalJ0 is the exact j_0 = sin(x)/x with the x = 0 limit.
|
||
func sphericalJ0(x float64) float64 {
|
||
if x == 0 {
|
||
return 1
|
||
}
|
||
return math.Sin(x) / x
|
||
}
|