469 lines
16 KiB
Go
469 lines
16 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package signal
|
||
|
||
import (
|
||
"math"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// Spectral solution of the Poisson equation on a rectangle with
|
||
// homogeneous boundary data. The sine basis vanishes on the boundary
|
||
// and diagonalises −Δ for Dirichlet data; the cosine basis has zero
|
||
// slope there and does the same for Neumann. Both solves run one
|
||
// separable transform pair per axis and one division per mode, and
|
||
// both refuse the inputs the mathematics refuses: nonzero-mean
|
||
// sources for Neumann, exactly as the periodic solve does.
|
||
|
||
// poissonRect validates the shared rectangle arguments and returns
|
||
// the grid shape and spacings.
|
||
func poissonRect(name string, f *core.Array, lx, ly float64) (rows, cols int, hx, hy float64, err error) {
|
||
if f.NDim() != 2 {
|
||
return 0, 0, 0, 0, base.Errf("%s: f must be rank 2, got shape %s", name, base.ShapeText(f.Shape()))
|
||
}
|
||
if f.Dtype() != core.Float {
|
||
return 0, 0, 0, 0, base.Errf("%s: f must be a float64 array, got %s", name, f.Dtype())
|
||
}
|
||
rows, cols = f.Shape()[0], f.Shape()[1]
|
||
if rows < 3 || cols < 3 {
|
||
return 0, 0, 0, 0, base.Errf("%s: the grid must be at least 3×3 to hold interior points, got %d×%d", name, rows, cols)
|
||
}
|
||
if lx <= 0 || ly <= 0 {
|
||
return 0, 0, 0, 0, base.Errf("%s: the side lengths must be positive, got %g and %g", name, lx, ly)
|
||
}
|
||
// A non-finite source would drive the compatibility tolerances NaN
|
||
// and slip through their gates, so it is refused up front. A dense
|
||
// float64 payload is scanned slice-wise: the elements are the ones
|
||
// FloatAt returns.
|
||
vals := poissonFloats(f)
|
||
for _, v := range vals {
|
||
if math.IsInf(v, 0) || math.IsNaN(v) {
|
||
return 0, 0, 0, 0, base.Errf("%s: f holds the non-finite value %g", name, v)
|
||
}
|
||
}
|
||
hx = lx / float64(cols-1)
|
||
hy = ly / float64(rows-1)
|
||
return rows, cols, hx, hy, nil
|
||
}
|
||
|
||
// poissonFloats streams the grid's elements in element order: the raw
|
||
// payload when the array is dense, a materialised copy through FloatAt
|
||
// when it is a strided view, whose payload order is not the element
|
||
// order. Treat an aliased result as read-only.
|
||
func poissonFloats(f *core.Array) []float64 {
|
||
if !f.Strided() {
|
||
return f.RawFloats()
|
||
}
|
||
out := make([]float64, f.Len())
|
||
for i := range out {
|
||
out[i] = f.FloatAt(i)
|
||
}
|
||
return out
|
||
}
|
||
|
||
// poissonTrapzSum returns the trapezoidal-weighted sum of f over the
|
||
// whole grid: the boundary rows and columns count half, the weight the
|
||
// mirrored stencil's own left null vector carries. It is the constant
|
||
// mode of the cosine transform pair, the measure the Neumann solve is
|
||
// compatible in.
|
||
func poissonTrapzSum(f *core.Array) float64 {
|
||
rows, cols := f.Shape()[0], f.Shape()[1]
|
||
vals := poissonFloats(f)
|
||
sum := 0.0
|
||
for r := range rows {
|
||
wr := 1.0
|
||
if r == 0 || r == rows-1 {
|
||
wr = 0.5
|
||
}
|
||
for c := range cols {
|
||
wc := 1.0
|
||
if c == 0 || c == cols-1 {
|
||
wc = 0.5
|
||
}
|
||
sum += wr * wc * vals[r*cols+c]
|
||
}
|
||
}
|
||
return sum
|
||
}
|
||
|
||
// SolvePoissonDirichlet solves −Δu = f on the rectangle with u held
|
||
// at zero on the whole boundary. f is sampled on the (rows × cols)
|
||
// grid including the boundary rows and columns, whose entries play no
|
||
// role in the solve; u comes back on the same grid with a zero
|
||
// boundary. Row r of f samples y = r·hy, column c samples x = c·hx,
|
||
// with hx = lx/(cols−1) and hy = ly/(rows−1). The sine transform
|
||
// pair along each axis diagonalises the interior second-difference
|
||
// operator, so the error is the stencil's O(h²), decaying with the
|
||
// grid. Non-float dtypes, grids under 3×3 and non-positive side
|
||
// lengths are errors.
|
||
func SolvePoissonDirichlet(f *core.Array, lx, ly float64) (*core.Array, error) {
|
||
const name = "SolvePoissonDirichlet"
|
||
rows, cols, hx, hy, err := poissonRect(name, f, lx, ly)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
interiorR, interiorC := rows-2, cols-2
|
||
// The interior right-hand side, transformed along both axes with
|
||
// the orthonormal DST-I, whose basis vanishes on the boundary.
|
||
spectrum, err := transformBlock(f, 1, 1, interiorR, interiorC, true)
|
||
if err != nil {
|
||
return nil, base.Errf("%s: %w", name, err)
|
||
}
|
||
// Divide by the stencil eigenvalues: two per axis from the cosine
|
||
// identity on the second difference. Each eigenvalue depends on one
|
||
// index alone, so the per-axis values are built once instead of once
|
||
// per mode; the expressions are the ones the inner loop evaluated.
|
||
eigX := make([]float64, interiorC+1)
|
||
for kx := 1; kx <= interiorC; kx++ {
|
||
eigX[kx] = 2 / (hx * hx) * (1 - math.Cos(math.Pi*float64(kx)/float64(interiorC+1)))
|
||
}
|
||
eigY := make([]float64, interiorR+1)
|
||
for ky := 1; ky <= interiorR; ky++ {
|
||
eigY[ky] = 2 / (hy * hy) * (1 - math.Cos(math.Pi*float64(ky)/float64(interiorR+1)))
|
||
}
|
||
for ky := 1; ky <= interiorR; ky++ {
|
||
ly := eigY[ky]
|
||
row := (ky - 1) * interiorC
|
||
for kx := 1; kx <= interiorC; kx++ {
|
||
spectrum[row+kx-1] /= eigX[kx] + ly
|
||
}
|
||
}
|
||
interior, err := inverseBlock(spectrum, interiorR, interiorC, true)
|
||
if err != nil {
|
||
return nil, base.Errf("%s: %w", name, err)
|
||
}
|
||
out := core.New(core.Float, rows, cols)
|
||
into := out.RawFloats()
|
||
for r := range interiorR {
|
||
for c := range interiorC {
|
||
into[(r+1)*cols+(c+1)] = interior[r*interiorC+c]
|
||
}
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// SolvePoissonNeumann solves −Δu = f with zero normal derivative on
|
||
// the whole boundary: u comes back on the full grid, boundary points
|
||
// included, every sample an unknown. Conventions mirror
|
||
// SolvePoissonDirichlet, with the cosine transform pair along each
|
||
// axis. The stencil is the ghost-mirror one, the neighbour outside the
|
||
// boundary standing in for zero slope, so a boundary row couples its
|
||
// mirrored neighbour with 2/h² where an interior row uses 1/h² (the
|
||
// boundary control volume is half an interior one). Its cosine modes
|
||
// cos(πkx x/lx)·cos(πky y/ly), kx = 0…cols−1 and ky = 0…rows−1, carry
|
||
// the eigenvalues 2/hx²·(1−cos(πkx/(cols−1))) +
|
||
// 2/hy²·(1−cos(πky/(rows−1))) per axis, and the error against the
|
||
// closed form is the stencil's O(h²).
|
||
//
|
||
// The operator annihilates the constants, so a solution exists only
|
||
// for sources whose trapezoidal-weighted sum vanishes; that sum is the
|
||
// operator's own compatibility condition (the trapezoidal measure is
|
||
// its left null vector), not the plain mean, and a nonzero one is an
|
||
// error rather than a quietly shifted problem. The free constant of u
|
||
// is fixed by setting the constant mode to zero, so the result has
|
||
// zero trapezoidal mean. Non-float dtypes, grids under 3×3 and
|
||
// non-positive side lengths are errors.
|
||
func SolvePoissonNeumann(f *core.Array, lx, ly float64) (*core.Array, error) {
|
||
const name = "SolvePoissonNeumann"
|
||
rows, cols, hx, hy, err := poissonRect(name, f, lx, ly)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
// The compatibility gate: the constant mode of the transformed
|
||
// source is the trapezoidal-weighted sum, so anything above the
|
||
// accumulated round-off of that sum breaks solvability. The
|
||
// tolerance follows the periodic solve's constant-mode gate, and it
|
||
// is relative to the source: the sum of a trapezoidal grid scales
|
||
// with the sample magnitude and the element count, so an absolute
|
||
// floor would let a tiny source through with a residual far above
|
||
// its own size (a 1e-9 source used to be accepted with a 1.8e-5
|
||
// relative residual).
|
||
sum := poissonTrapzSum(f)
|
||
tol := 64 * base.EpsF * float64(rows*cols) * maxAbsF(f)
|
||
if math.Abs(sum) > tol {
|
||
return nil, base.Errf("%s: f has the nonzero trapezoidal mean %g, the Neumann problem has no solution",
|
||
name, sum/float64((rows-1)*(cols-1)))
|
||
}
|
||
// The pipeline. The mirrored operator M (rows [2,−2] on the
|
||
// boundary, [−1,2,−1] inside, over h²) is diagonalised by the
|
||
// UNNORMALISED DCT-I, whose boundary input weight of 1/2 is the
|
||
// trapezoidal measure; the library carries the orthonormal DCT-I
|
||
// (boundary weight 1/sqrt2), and W = diag(1/2,1,…,1,1/2) turns one
|
||
// into the other. So the forward trip scales by W^(1/2) before the
|
||
// DCT-I, divides by the eigenvalues, and the return trip scales by
|
||
// W^(-1/2) after the second DCT-I, which is its own inverse.
|
||
weighted := core.New(core.Float, rows, cols)
|
||
src := poissonFloats(f)
|
||
dst := weighted.RawFloats()
|
||
for r := range rows {
|
||
sr := 1.0
|
||
if r == 0 || r == rows-1 {
|
||
sr = math.Sqrt2 / 2
|
||
}
|
||
for c := range cols {
|
||
sc := 1.0
|
||
if c == 0 || c == cols-1 {
|
||
sc = math.Sqrt2 / 2
|
||
}
|
||
dst[r*cols+c] = src[r*cols+c] * sr * sc
|
||
}
|
||
}
|
||
spectrum, err := transformBlock(weighted, 0, 0, rows, cols, false)
|
||
if err != nil {
|
||
return nil, base.Errf("%s: %w", name, err)
|
||
}
|
||
// The x eigenvalues depend on kx alone, so the row is built once and
|
||
// read by every ky; the expression is the one the inner loop
|
||
// evaluated. The y eigenvalue was already hoisted to the outer loop.
|
||
eigX := make([]float64, cols)
|
||
for kx := range cols {
|
||
eigX[kx] = 2 / (hx * hx) * (1 - math.Cos(math.Pi*float64(kx)/float64(cols-1)))
|
||
}
|
||
for ky := range rows {
|
||
ly := 2 / (hy * hy) * (1 - math.Cos(math.Pi*float64(ky)/float64(rows-1)))
|
||
row := ky * cols
|
||
for kx := range cols {
|
||
i := row + kx
|
||
if kx == 0 && ky == 0 {
|
||
// The constant mode is the operator's null space: it
|
||
// stays zero, which is what fixes the solution's free
|
||
// constant, instead of dividing by its zero eigenvalue.
|
||
spectrum[i] = 0
|
||
continue
|
||
}
|
||
spectrum[i] /= eigX[kx] + ly
|
||
}
|
||
}
|
||
interior, err := inverseBlock(spectrum, rows, cols, false)
|
||
if err != nil {
|
||
return nil, base.Errf("%s: %w", name, err)
|
||
}
|
||
// The return trip: undo the W^(1/2) scaling, and strip the
|
||
// trapezoidal mean, the measure the solve works in and the one the
|
||
// operator leaves free.
|
||
out := core.New(core.Float, rows, cols)
|
||
into := out.RawFloats()
|
||
sum, weight := 0.0, 0.0
|
||
for r := range rows {
|
||
sr, wr := 1.0, 1.0
|
||
if r == 0 || r == rows-1 {
|
||
sr, wr = math.Sqrt2, 0.5
|
||
}
|
||
for c := range cols {
|
||
sc, wc := 1.0, 1.0
|
||
if c == 0 || c == cols-1 {
|
||
sc, wc = math.Sqrt2, 0.5
|
||
}
|
||
v := interior[r*cols+c] * sr * sc
|
||
into[r*cols+c] = v
|
||
sum += wr * wc * v
|
||
weight += wr * wc
|
||
}
|
||
}
|
||
mean := sum / weight
|
||
for i := range into {
|
||
into[i] -= mean
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// maxAbsF returns the largest absolute sample of a float array.
|
||
func maxAbsF(f *core.Array) float64 {
|
||
largest := 0.0
|
||
for _, v := range poissonFloats(f) {
|
||
largest = max(largest, math.Abs(v))
|
||
}
|
||
return largest
|
||
}
|
||
|
||
// lineTransformPlan carries the chirp-z pieces one transform length
|
||
// needs for the orthonormal DST-I (sine) or DCT-I (cosine) the Poisson
|
||
// solves apply along a line, plus the buffers the per-line work reuses.
|
||
//
|
||
// Both transforms are the real or imaginary part of a sum of complex
|
||
// exponentials over the half-frequency grid the trigonometric identity
|
||
// (j+1)(k+1) = ((j+1)² + (k+1)² − (j−k)²)/2 exposes: a per-sample
|
||
// phase, a per-bin rotation and an even chirp kernel, so the sum over
|
||
// the bins is a circular convolution of length the smallest power of
|
||
// two ≥ 2n−1. The padded full-length route each line used to take ran
|
||
// a convolution of roughly twice that length to reach the same bins.
|
||
type lineTransformPlan struct {
|
||
n int
|
||
m int
|
||
sine bool
|
||
// pre weights the samples, post rotates the bins, kern is the
|
||
// kernel's forward transform on the convolution grid.
|
||
pre []complex128
|
||
post []complex128
|
||
kern []complex128
|
||
norm float64
|
||
// buf carries one line's spectrum through the convolution; the plan
|
||
// is used from one goroutine and the solves are single-threaded, so
|
||
// the buffer is fully rewritten on every apply.
|
||
buf []complex128
|
||
}
|
||
|
||
// newLineTransformPlan builds the plan for n points. The chirp angles
|
||
// are reduced modulo the kernel's period in exact integer arithmetic
|
||
// before the transcendental call, so no angle grows with n.
|
||
func newLineTransformPlan(n int, sine bool) *lineTransformPlan {
|
||
p := &lineTransformPlan{n: n, sine: sine}
|
||
m := 1
|
||
for m < 2*n-1 {
|
||
m <<= 1
|
||
}
|
||
p.m = m
|
||
p.pre = make([]complex128, n)
|
||
p.post = make([]complex128, n)
|
||
p.kern = make([]complex128, m)
|
||
p.buf = make([]complex128, m)
|
||
if sine {
|
||
// DST-I: out[k] = norm·Im[ post_k · Σ_j x_j·pre_j·e^{−iθ(j−k)²/2} ]
|
||
// with θ = π/(n+1), pre_j = e^{iθ(j + j²/2)},
|
||
// post_k = e^{iθ(k+1 + k²/2)}.
|
||
period := 4 * (n + 1)
|
||
theta := math.Pi / float64(n+1)
|
||
for j := range n {
|
||
t := (2*j + j*j) % period
|
||
p.pre[j] = polar(1, theta*float64(t)/2)
|
||
}
|
||
for k := range n {
|
||
t := (2*(k+1) + k*k) % period
|
||
p.post[k] = polar(1, theta*float64(t)/2)
|
||
}
|
||
for l := range n {
|
||
t := (l * l) % period
|
||
w := polar(1, -theta*float64(t)/2)
|
||
p.kern[l] = w
|
||
if l != 0 {
|
||
p.kern[m-l] = w
|
||
}
|
||
}
|
||
p.norm = math.Sqrt(2 / float64(n+1))
|
||
} else {
|
||
// DCT-I: out[k] = norm·Re[ w_k·e^{iθk²/2} ·
|
||
// Σ_j x_j·w_j·e^{iθj²/2}·e^{−iθ(j−k)²/2} ] with θ = π/(n−1)
|
||
// and w the half weight on the endpoints.
|
||
period := 4 * (n - 1)
|
||
theta := math.Pi / float64(n-1)
|
||
for j := range n {
|
||
t := (j * j) % period
|
||
w := 1.0
|
||
if j == 0 || j == n-1 {
|
||
w = math.Sqrt2 / 2
|
||
}
|
||
p.pre[j] = complex(w, 0) * polar(1, theta*float64(t)/2)
|
||
}
|
||
for k := range n {
|
||
t := (k * k) % period
|
||
w := 1.0
|
||
if k == 0 || k == n-1 {
|
||
w = math.Sqrt2 / 2
|
||
}
|
||
p.post[k] = complex(w, 0) * polar(1, theta*float64(t)/2)
|
||
}
|
||
for l := range n {
|
||
t := (l * l) % period
|
||
w := polar(1, -theta*float64(t)/2)
|
||
p.kern[l] = w
|
||
if l != 0 {
|
||
p.kern[m-l] = w
|
||
}
|
||
}
|
||
p.norm = math.Sqrt(2 / float64(n-1))
|
||
}
|
||
// The kernel's spectrum depends only on the plan; it is the same
|
||
// forward transform every line multiplies against.
|
||
transform(p.kern, -1)
|
||
return p
|
||
}
|
||
|
||
// polar returns the unit complex number of the given angle.
|
||
// math.Sincos returns exactly the pair math.Cos and math.Sin produce
|
||
// (verified bit-for-bit), so the value is unchanged.
|
||
func polar(magnitude, angle float64) complex128 {
|
||
s, c := math.Sincos(angle)
|
||
return complex(magnitude*c, magnitude*s)
|
||
}
|
||
|
||
// apply transforms src into dst; both have length p.n and must not
|
||
// alias. The convolution runs on the plan's reused spectrum buffer,
|
||
// which every call overwrites completely before it is read.
|
||
func (p *lineTransformPlan) apply(dst, src []float64) {
|
||
n, m := p.n, p.m
|
||
buf := p.buf
|
||
if p.sine && n == 1 {
|
||
dst[0] = src[0]
|
||
return
|
||
}
|
||
for j := range n {
|
||
buf[j] = complex(src[j], 0) * p.pre[j]
|
||
}
|
||
clear(buf[n:m])
|
||
transform(buf, -1)
|
||
for i := range m {
|
||
buf[i] *= p.kern[i]
|
||
}
|
||
transform(buf, +1)
|
||
scale := complex(1/float64(m), 0)
|
||
if p.sine {
|
||
for k := range n {
|
||
v := buf[k] * scale * p.post[k]
|
||
dst[k] = p.norm * imag(v)
|
||
}
|
||
return
|
||
}
|
||
for k := range n {
|
||
v := buf[k] * scale * p.post[k]
|
||
dst[k] = p.norm * real(v)
|
||
}
|
||
}
|
||
|
||
// transformBlock applies the named self-inverse orthonormal transform
|
||
// (DST-I when sine, DCT-I when not) to every row and then every column
|
||
// of the rows×cols block of f whose top-left corner sits at (row0,
|
||
// col0), returning the flat spectrum. One plan per axis length serves
|
||
// every line of that axis.
|
||
func transformBlock(f *core.Array, row0, col0, rows, cols int, sine bool) ([]float64, error) {
|
||
srcVals := poissonFloats(f)
|
||
srcCols := f.Shape()[1]
|
||
work := make([]float64, rows*cols)
|
||
longest := max(rows, cols)
|
||
line := make([]float64, longest)
|
||
out := make([]float64, longest)
|
||
// poissonFloats materialises a strided source in element order; a
|
||
// raw payload read would keep the physical order instead.
|
||
rowPlan := newLineTransformPlan(cols, sine)
|
||
for r := range rows {
|
||
for c := range cols {
|
||
line[c] = srcVals[(r+row0)*srcCols+(c+col0)]
|
||
}
|
||
rowPlan.apply(out[:cols], line[:cols])
|
||
copy(work[r*cols:(r+1)*cols], out[:cols])
|
||
}
|
||
colPlan := newLineTransformPlan(rows, sine)
|
||
for c := range cols {
|
||
for r := range rows {
|
||
line[r] = work[r*cols+c]
|
||
}
|
||
colPlan.apply(out[:rows], line[:rows])
|
||
for r := range rows {
|
||
work[r*cols+c] = out[r]
|
||
}
|
||
}
|
||
return work, nil
|
||
}
|
||
|
||
// inverseBlock applies the same self-inverse transform to a flat
|
||
// spectrum, the return trip of transformBlock.
|
||
func inverseBlock(spectrum []float64, rows, cols int, sine bool) ([]float64, error) {
|
||
arr, err := core.FromFloats(spectrum, rows, cols)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return transformBlock(arr, 0, 0, rows, cols, sine)
|
||
}
|