219 lines
6.8 KiB
Go
219 lines
6.8 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package signal
|
||
|
||
import (
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
import "slices"
|
||
|
||
import "sourcedock.dev/petrbalvin/tensor/internal/engine"
|
||
|
||
// Numerically robust scalar reduction and finite-difference stencils
|
||
// for grid data. SumKahan applies Kahan's compensated summation to the
|
||
// whole array: the same sum `Sum` computes, but with the running
|
||
// rounding error carried alongside, so long chains of similar
|
||
// magnitudes stop bleeding digits. The finite-difference stencils
|
||
// produce the central-difference gradient and Laplacian of grid data,
|
||
// the raw material of explicit PDE steppers.
|
||
|
||
// SumKahan returns the sum of all elements by Kahan's compensated
|
||
// summation. For arrays with mixed magnitudes the result carries more
|
||
// of the small terms than a plain left-to-right accumulation.
|
||
// core.Complex arrays are refused.
|
||
func SumKahan(a *core.Array) (float64, error) {
|
||
if a.Dtype() == core.Complex {
|
||
return 0, base.Errf("SumKahan: complex arrays are not supported")
|
||
}
|
||
src := widenFloats(a)
|
||
var sum, comp float64
|
||
// Bounded by the element count: a rebased view's aliased payload
|
||
// runs past its extent.
|
||
for i := range a.Len() {
|
||
v := src[i]
|
||
v -= comp
|
||
t := sum + v
|
||
comp = (t - sum) - v
|
||
sum = t
|
||
}
|
||
return sum, nil
|
||
}
|
||
|
||
// Gradient1D returns the central-difference derivative of a 1-D
|
||
// uniform grid signal y with spacing dx. Interior points use the
|
||
// second-order central stencil; the endpoints fall back to the
|
||
// first-order one-sided stencil, so the result has the same length as
|
||
// the input.
|
||
func Gradient1D(y *core.Array, dx float64) (*core.Array, error) {
|
||
if y.NDim() != 1 {
|
||
return nil, base.Errf("Gradient1D: needs a rank-1 signal, got shape %s", base.ShapeText(y.Shape()))
|
||
}
|
||
if y.Dtype() == core.Complex {
|
||
return nil, base.Errf("Gradient1D: complex arrays are not supported")
|
||
}
|
||
n := y.Len()
|
||
if n < 2 {
|
||
return nil, base.Errf("Gradient1D: needs at least 2 points, got %d", n)
|
||
}
|
||
if dx == 0 {
|
||
return nil, base.Errf("Gradient1D: spacing must not be zero")
|
||
}
|
||
src := widenFloats(y)
|
||
out := core.New(core.Float, n)
|
||
dst := out.RawFloats()
|
||
dst[0] = (src[1] - src[0]) / dx
|
||
for i := 1; i < n-1; i++ {
|
||
dst[i] = (src[i+1] - src[i-1]) / (2 * dx)
|
||
}
|
||
dst[n-1] = (src[n-1] - src[n-2]) / dx
|
||
return out, nil
|
||
}
|
||
|
||
// Laplacian returns the second-derivative Laplacian of a grid signal.
|
||
// 1-D input takes one spacing, 2-D input two (dx, dy for a row-major
|
||
// (rows, cols) grid, 5-point stencil), 3-D input three (7-point
|
||
// stencil). Boundary points copy their nearest interior value; the
|
||
// caller owns the boundary condition. A rank outside 1..3, a zero
|
||
// spacing, a missing spacing, or a degenerate extent is an error:
|
||
// every axis needs at least 2 points, and every axis of a 2-D or 3-D
|
||
// grid at least 3, so the interior stencil and the boundary copies
|
||
// have somewhere to stand.
|
||
func Laplacian(a *core.Array, spacings ...float64) (*core.Array, error) {
|
||
const name = "Laplacian"
|
||
if a.Dtype() == core.Complex {
|
||
return nil, base.Errf("%s: complex arrays are not supported", name)
|
||
}
|
||
if len(spacings) < a.NDim() {
|
||
return nil, base.Errf("%s: needs %d spacings for rank %d, got %d",
|
||
name, a.NDim(), a.NDim(), len(spacings))
|
||
}
|
||
if slices.Contains(spacings, 0) {
|
||
return nil, base.Errf("%s: spacing must not be zero", name)
|
||
}
|
||
switch a.NDim() {
|
||
case 1:
|
||
if a.Len() < 2 {
|
||
return nil, base.Errf("%s: needs at least 2 points, got %d", name, a.Len())
|
||
}
|
||
case 2, 3:
|
||
for i, d := range a.Shape() {
|
||
if d < 3 {
|
||
return nil, base.Errf("%s: axis %d has extent %d, the stencil needs at least 3 points",
|
||
name, i, d)
|
||
}
|
||
}
|
||
default:
|
||
return nil, base.Errf("%s: needs a 1-D, 2-D or 3-D array, got rank %d", name, a.NDim())
|
||
}
|
||
switch a.NDim() {
|
||
case 1:
|
||
return laplacian1D(a, spacings[0]), nil
|
||
case 2:
|
||
return laplacian2D(a, spacings[0], spacings[1]), nil
|
||
default:
|
||
return laplacian3D(a, spacings[0], spacings[1], spacings[2]), nil
|
||
}
|
||
}
|
||
|
||
// laplacian1D applies the 3-point second-difference stencil along the
|
||
// single axis. Interior: (y[i+1] − 2y[i] + y[i−1]) / dx². For n = 2
|
||
// there is no interior point, and the edge copy leaves both entries
|
||
// zero.
|
||
func laplacian1D(y *core.Array, dx float64) *core.Array {
|
||
n := y.Len()
|
||
out := core.New(core.Float, n)
|
||
inv := 1 / (dx * dx)
|
||
src := widenFloats(y)
|
||
dst := out.RawFloats()
|
||
for i := 1; i < n-1; i++ {
|
||
dst[i] = (src[i+1] - 2*src[i] + src[i-1]) * inv
|
||
}
|
||
dst[0] = dst[1]
|
||
dst[n-1] = dst[n-2]
|
||
return out
|
||
}
|
||
|
||
// laplacian2D applies the 5-point stencil on a row-major grid of
|
||
// shape (ny, nx) with spacings dx, dy.
|
||
func laplacian2D(y *core.Array, dx, dy float64) *core.Array {
|
||
ny, nx := y.Shape()[0], y.Shape()[1]
|
||
out := core.New(core.Float, ny, nx)
|
||
invx := 1 / (dx * dx)
|
||
invy := 1 / (dy * dy)
|
||
src := widenFloats(y)
|
||
dst := out.RawFloats()
|
||
engine.ParallelMin(ny, workFloorFor(nx*6), func(s, e int) {
|
||
for j := s; j < e; j++ {
|
||
for i := 1; i < nx-1; i++ {
|
||
idx := j*nx + i
|
||
c := src[idx]
|
||
acc := (src[idx+1] + src[idx-1] - 2*c) * invx
|
||
if j > 0 && j < ny-1 {
|
||
acc += (src[(j-1)*nx+i] + src[(j+1)*nx+i] - 2*c) * invy
|
||
}
|
||
dst[idx] = acc
|
||
}
|
||
dst[j*nx] = dst[j*nx+1]
|
||
dst[j*nx+nx-1] = dst[j*nx+nx-2]
|
||
}
|
||
})
|
||
for i := range nx {
|
||
dst[i] = dst[nx+i]
|
||
dst[(ny-1)*nx+i] = dst[(ny-2)*nx+i]
|
||
}
|
||
return out
|
||
}
|
||
|
||
// laplacian3D applies the 7-point stencil on a row-major grid of
|
||
// shape (nz, ny, nx) with spacings dx, dy, dz. Boundary planes copy
|
||
// their nearest interior plane, the caller owns the boundary
|
||
// condition.
|
||
func laplacian3D(y *core.Array, dx, dy, dz float64) *core.Array {
|
||
nz, ny, nx := y.Shape()[0], y.Shape()[1], y.Shape()[2]
|
||
out := core.New(core.Float, nz, ny, nx)
|
||
invx := 1 / (dx * dx)
|
||
invy := 1 / (dy * dy)
|
||
invz := 1 / (dz * dz)
|
||
src := widenFloats(y)
|
||
dst := out.RawFloats()
|
||
engine.ParallelMin(nz, workFloorFor(ny*nx*9), func(zs, ze int) {
|
||
for k := zs; k < ze; k++ {
|
||
for j := range ny {
|
||
for i := 1; i < nx-1; i++ {
|
||
idx := (k*ny+j)*nx + i
|
||
c := src[idx]
|
||
acc := (src[idx+1] + src[idx-1] - 2*c) * invx
|
||
if j > 0 && j < ny-1 {
|
||
acc += (src[(k*ny+j-1)*nx+i] + src[(k*ny+j+1)*nx+i] - 2*c) * invy
|
||
}
|
||
if k > 0 && k < nz-1 {
|
||
acc += (src[((k-1)*ny+j)*nx+i] + src[((k+1)*ny+j)*nx+i] - 2*c) * invz
|
||
}
|
||
dst[idx] = acc
|
||
}
|
||
// The x boundaries copy their nearest interior value.
|
||
dst[(k*ny+j)*nx] = dst[(k*ny+j)*nx+1]
|
||
dst[(k*ny+j)*nx+nx-1] = dst[(k*ny+j)*nx+nx-2]
|
||
}
|
||
}
|
||
})
|
||
// The y boundaries copy their nearest interior row.
|
||
for k := range nz {
|
||
for i := range nx {
|
||
dst[k*ny*nx+i] = dst[(k*ny+1)*nx+i]
|
||
dst[((k+1)*ny-1)*nx+i] = dst[((k+1)*ny-2)*nx+i]
|
||
}
|
||
}
|
||
// The z boundaries copy their nearest interior plane.
|
||
for j := range ny {
|
||
for i := range nx {
|
||
dst[j*nx+i] = dst[(ny+j)*nx+i]
|
||
dst[((nz-1)*ny+j)*nx+i] = dst[((nz-2)*ny+j)*nx+i]
|
||
}
|
||
}
|
||
return out
|
||
}
|