424 lines
16 KiB
Go
424 lines
16 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package integrate
|
||
|
||
import (
|
||
"math"
|
||
"sync"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/engine"
|
||
|
||
"sourcedock.dev/petrbalvin/tensor/internal/base"
|
||
"sourcedock.dev/petrbalvin/tensor/internal/core"
|
||
)
|
||
|
||
// Two-dimensional evolution equations on the rectangle, the
|
||
// higher-rank siblings of IntegrateHeat1D and IntegrateWave1D. The
|
||
// grid is the rank-2 shape of the initial state: row r samples
|
||
// y = r·dy, column c samples x = c·dx, the boundary ring is held
|
||
// fixed, and the interior carries the dynamics.
|
||
|
||
// pde2dValidate checks the shared rectangle arguments and returns the
|
||
// grid shape.
|
||
func pde2dValidate(name string, u0 *core.Array, dx, dy, tFinal, dt float64, samples int) (rows, cols int, err error) {
|
||
if u0.NDim() != 2 {
|
||
return 0, 0, base.Errf("%s: the initial state must be rank 2, got shape %s", name, base.ShapeText(u0.Shape()))
|
||
}
|
||
if u0.Dtype() == core.Complex {
|
||
return 0, 0, base.Errf("%s: complex initial states are not supported", name)
|
||
}
|
||
rows, cols = u0.Shape()[0], u0.Shape()[1]
|
||
if rows < 3 || cols < 3 {
|
||
return 0, 0, base.Errf("%s: the grid must be at least 3×3 to hold interior points, got %d×%d", name, rows, cols)
|
||
}
|
||
if !(dx > 0) || !(dy > 0) {
|
||
return 0, 0, base.Errf("%s: the spacings must be positive, got %g and %g", name, dx, dy)
|
||
}
|
||
if !(tFinal > 0) || !(dt > 0) {
|
||
return 0, 0, base.Errf("%s: tFinal and dt must be positive, got %g and %g", name, tFinal, dt)
|
||
}
|
||
// As in the 1-D validation: a dt this far below tFinal cannot be
|
||
// honoured, and the wrapped step count would disable the bound
|
||
// silently.
|
||
if tFinal/dt > 1e12 {
|
||
return 0, 0, base.Errf("%s: dt = %g asks for more than 1e12 steps over %g", name, dt, tFinal)
|
||
}
|
||
if samples < 2 {
|
||
return 0, 0, base.Errf("%s: at least two samples are needed, got %d", name, samples)
|
||
}
|
||
// A non-finite entry flows through the stencils and the solvers'
|
||
// zero-pivot guards compare false against NaN, so it is refused
|
||
// up front.
|
||
for i := range u0.Len() {
|
||
if v := u0.FloatAt(i); math.IsNaN(v) || math.IsInf(v, 0) {
|
||
return 0, 0, base.Errf("%s: the initial state holds the non-finite value %g at %d", name, v, i)
|
||
}
|
||
}
|
||
return rows, cols, nil
|
||
}
|
||
|
||
// fromSlice wraps a float slice of known length as an array.
|
||
func fromSlice(vals []float64, n int) *core.Array {
|
||
out := core.New(core.Float, n)
|
||
copy(out.RawFloats(), vals[:n])
|
||
return out
|
||
}
|
||
|
||
// IntegrateHeat2D evolves u_t = κ·Δu on the rectangle by
|
||
// Peaceman-Rachford alternating direction implicit steps: one
|
||
// half-step implicit in x against an explicit y neighbour sum, one
|
||
// implicit in y against an explicit x neighbour sum, each row and
|
||
// column a tridiagonal solve through SolveTridiagonal. The scheme is
|
||
// second order in space and time and unconditionally stable, so no
|
||
// step-size refusal stands between the caller and a coarse first
|
||
// look. The boundary ring is held at the four constant edge values.
|
||
// The return is a (samples × rows × cols) array: the initial state,
|
||
// then the state after each stored interval, the final state forced
|
||
// into the last sample.
|
||
func IntegrateHeat2D(u0 *core.Array, kappa, dx, dy, tFinal, dt float64, samples int,
|
||
boundBottom, boundTop, boundLeft, boundRight float64) (*core.Array, error) {
|
||
const name = "IntegrateHeat2D"
|
||
rows, cols, err := pde2dValidate(name, u0, dx, dy, tFinal, dt, samples)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if !(kappa > 0) || math.IsInf(kappa, 0) {
|
||
return nil, base.Errf("%s: the diffusivity must be positive, got %g", name, kappa)
|
||
}
|
||
// The four boundary constants enter the explicit neighbour sums
|
||
// every step: a non-finite one would flow through the stencils and
|
||
// the solves and publish an all-NaN history with no error.
|
||
if math.IsNaN(boundBottom) || math.IsInf(boundBottom, 0) ||
|
||
math.IsNaN(boundTop) || math.IsInf(boundTop, 0) ||
|
||
math.IsNaN(boundLeft) || math.IsInf(boundLeft, 0) ||
|
||
math.IsNaN(boundRight) || math.IsInf(boundRight, 0) {
|
||
return nil, base.Errf("%s: the boundary values must be finite, got bottom %g, top %g, left %g, right %g",
|
||
name, boundBottom, boundTop, boundLeft, boundRight)
|
||
}
|
||
steps, h := pdeSchedule(tFinal, dt, samples)
|
||
rx := kappa * h / (2 * dx * dx) // the implicit half-step weight in x
|
||
ry := kappa * h / (2 * dy * dy) // ...and in y
|
||
|
||
u := make([]float64, rows*cols)
|
||
copy(u, denseFloats(u0))
|
||
every := steps / (samples - 1)
|
||
out := core.New(core.Float, samples, rows, cols)
|
||
into := out.RawFloats()
|
||
// Sample 0 is the initial state exactly, before any ring
|
||
// enforcement touches the working state.
|
||
copy(into, u)
|
||
written := 1
|
||
// The working state's horizontal ring rows carry the boundary
|
||
// constants from here on, so the stencils read their up and down
|
||
// neighbours straight out of the payload: a ring slot holds exactly
|
||
// the value the boundary test it replaces would have substituted.
|
||
// The step loop below rewrites both rows after every half-step,
|
||
// so the invariant holds for every step, the first included.
|
||
for c := range cols {
|
||
u[c] = boundBottom
|
||
u[(rows-1)*cols+c] = boundTop
|
||
}
|
||
|
||
star := make([]float64, rows*cols)
|
||
var stepErr error
|
||
var errMu sync.Mutex
|
||
// The implicit diagonals are constants of the scheme, one set per
|
||
// orientation, built once with exactly the values the solves read.
|
||
lowerX := make([]float64, cols-3)
|
||
upperX := make([]float64, cols-3)
|
||
diagX := make([]float64, cols-2)
|
||
for i := range lowerX {
|
||
lowerX[i] = -rx
|
||
upperX[i] = -rx
|
||
}
|
||
for i := range diagX {
|
||
diagX[i] = 1 + 2*rx
|
||
}
|
||
lowerY := make([]float64, rows-3)
|
||
upperY := make([]float64, rows-3)
|
||
diagY := make([]float64, rows-2)
|
||
for i := range lowerY {
|
||
lowerY[i] = -ry
|
||
upperY[i] = -ry
|
||
}
|
||
for i := range diagY {
|
||
diagY[i] = 1 + 2*ry
|
||
}
|
||
// Per-solve line scratch, one set per pdeStepFloor lines of each
|
||
// orientation. engine.ParallelMin invokes a sweep closure once per
|
||
// chunk; in parallel the chunk starts differ by at least
|
||
// pdeStepFloor, so two concurrent chunks never share a set, and a
|
||
// sweep the engine runs inline reuses set zero serially.
|
||
triX := make([]triScratch, (rows-2-1)/pdeStepFloor+1)
|
||
triY := make([]triScratch, (cols-2-1)/pdeStepFloor+1)
|
||
for i := range triX {
|
||
triSized(&triX[i], cols-2, cols)
|
||
}
|
||
for i := range triY {
|
||
triSized(&triY[i], rows-2, rows)
|
||
}
|
||
for s := 1; s <= steps; s++ {
|
||
// Half-step one: implicit in x along every interior row. Rows are
|
||
// independent (a row's solve reads its own and its neighbours'
|
||
// values and writes the star row), so the sweep is partitioned
|
||
// over them; each chunk carries its own right-hand side and
|
||
// solve scratch, and the arithmetic of a row is untouched.
|
||
engine.ParallelMin(rows-2, pdeStepFloor, func(rs, re int) {
|
||
tri := &triX[rs/pdeStepFloor]
|
||
rhs := tri.rhs
|
||
for r := rs + 1; r < re+1; r++ {
|
||
// The row's own values and its two neighbours are slices
|
||
// of the payload, so the stencil reads them by offset
|
||
// instead of rebuilding the flat index per element. The
|
||
// horizontal ring rows hold the boundary constants, so
|
||
// the neighbours need no edge test.
|
||
row := u[r*cols : (r+1)*cols]
|
||
up := u[(r+1)*cols : (r+2)*cols]
|
||
down := u[(r-1)*cols : r*cols]
|
||
for c := range cols {
|
||
rhs[c] = row[c] + ry*(up[c]-2*row[c]+down[c])
|
||
}
|
||
// Boundary neighbours of the implicit solve: the row's
|
||
// own left and right ring values enter the right side.
|
||
rhs[1] += rx * boundLeft
|
||
rhs[cols-2] += rx * boundRight
|
||
// The elimination writes the row's interior straight
|
||
// into the star state; the ring values follow.
|
||
dst := star[r*cols+1 : r*cols+cols-1]
|
||
if serr := base.TriSolve(dst, tri.cp, tri.dp, lowerX, diagX, upperX, rhs[1:cols-1]); serr != nil {
|
||
errMu.Lock()
|
||
if stepErr == nil {
|
||
stepErr = base.Errf("%s: %w", name, serr)
|
||
}
|
||
errMu.Unlock()
|
||
return
|
||
}
|
||
star[r*cols] = boundLeft
|
||
star[r*cols+cols-1] = boundRight
|
||
}
|
||
})
|
||
if stepErr != nil {
|
||
return nil, stepErr
|
||
}
|
||
for c := range cols {
|
||
star[c] = boundBottom
|
||
star[(rows-1)*cols+c] = boundTop
|
||
}
|
||
// Half-step two: implicit in y down every interior column,
|
||
// with the explicit x neighbours from the star state. Columns are
|
||
// independent the same way the rows were, and the walk takes them
|
||
// two at a time: the star state's vertical ring columns hold the
|
||
// boundary constants, so both lanes read their neighbours
|
||
// straight out of the payload, and a pair shares its loads, so
|
||
// each line of the state is fetched once per pair instead of
|
||
// three times across three column passes.
|
||
engine.ParallelMin(cols-2, pdeStepFloor, func(cs, ce int) {
|
||
tri := &triY[cs/pdeStepFloor]
|
||
rhs := tri.rhs
|
||
aux := tri.aux
|
||
dst := tri.dst
|
||
for c := cs + 1; c < ce+1; c += 2 {
|
||
if c+1 < ce+1 {
|
||
// Columns c and c+1: one pass over the rows builds
|
||
// both right sides, the shared neighbour loads in
|
||
// registers.
|
||
for r := range rows {
|
||
base := r * cols
|
||
l := star[base+c-1]
|
||
wm := star[base+c]
|
||
e := star[base+c+1]
|
||
rhs[r] = wm + rx*(e-2*wm+l)
|
||
aux[r] = e + rx*(star[base+c+2]-2*e+wm)
|
||
}
|
||
rhs[1] += ry * boundBottom
|
||
rhs[rows-2] += ry * boundTop
|
||
if serr := base.TriSolve(dst, tri.cp, tri.dp, lowerY, diagY, upperY, rhs[1:rows-1]); serr != nil {
|
||
errMu.Lock()
|
||
if stepErr == nil {
|
||
stepErr = base.Errf("%s: %w", name, serr)
|
||
}
|
||
errMu.Unlock()
|
||
return
|
||
}
|
||
u[c] = boundBottom
|
||
u[(rows-1)*cols+c] = boundTop
|
||
// The column's interior comes straight out of the
|
||
// solve's payload: the values are the ones the
|
||
// accessor read.
|
||
for r := 1; r < rows-1; r++ {
|
||
u[r*cols+c] = dst[r-1]
|
||
}
|
||
aux[1] += ry * boundBottom
|
||
aux[rows-2] += ry * boundTop
|
||
if serr := base.TriSolve(dst, tri.cp, tri.dp, lowerY, diagY, upperY, aux[1:rows-1]); serr != nil {
|
||
errMu.Lock()
|
||
if stepErr == nil {
|
||
stepErr = base.Errf("%s: %w", name, serr)
|
||
}
|
||
errMu.Unlock()
|
||
return
|
||
}
|
||
u[c+1] = boundBottom
|
||
u[(rows-1)*cols+c+1] = boundTop
|
||
for r := 1; r < rows-1; r++ {
|
||
u[r*cols+c+1] = dst[r-1]
|
||
}
|
||
continue
|
||
}
|
||
// The odd tail column, when the chunk ends on one.
|
||
for r := range rows {
|
||
base := r * cols
|
||
wm := star[base+c]
|
||
rhs[r] = wm + rx*(star[base+c+1]-2*wm+star[base+c-1])
|
||
}
|
||
rhs[1] += ry * boundBottom
|
||
rhs[rows-2] += ry * boundTop
|
||
if serr := base.TriSolve(dst, tri.cp, tri.dp, lowerY, diagY, upperY, rhs[1:rows-1]); serr != nil {
|
||
errMu.Lock()
|
||
if stepErr == nil {
|
||
stepErr = base.Errf("%s: %w", name, serr)
|
||
}
|
||
errMu.Unlock()
|
||
return
|
||
}
|
||
u[c] = boundBottom
|
||
u[(rows-1)*cols+c] = boundTop
|
||
for r := 1; r < rows-1; r++ {
|
||
u[r*cols+c] = dst[r-1]
|
||
}
|
||
}
|
||
})
|
||
if stepErr != nil {
|
||
return nil, stepErr
|
||
}
|
||
for r := range rows {
|
||
u[r*cols] = boundLeft
|
||
u[r*cols+cols-1] = boundRight
|
||
}
|
||
if s%every == 0 && written < samples {
|
||
copy(into[written*rows*cols:(written+1)*rows*cols], u)
|
||
written++
|
||
}
|
||
}
|
||
copy(into[(samples-1)*rows*cols:], u)
|
||
return out, nil
|
||
}
|
||
|
||
// IntegrateWave2D evolves u_tt = c²·Δu on the rectangle with the
|
||
// boundary ring held at zero, by the explicit central-difference
|
||
// stencil from the velocity Verlet family the 1-D wave solver uses:
|
||
// second order in space and time, with the CFL budget
|
||
// c·dt·sqrt(1/dx² + 1/dy²) ≤ 1 enforced as an error, because the
|
||
// explicit stencil has no honest answer past it. The return contract
|
||
// mirrors IntegrateHeat2D.
|
||
func IntegrateWave2D(u0, v0 *core.Array, c, dx, dy, tFinal, dt float64, samples int) (*core.Array, error) {
|
||
const name = "IntegrateWave2D"
|
||
rows, cols, err := pde2dValidate(name, u0, dx, dy, tFinal, dt, samples)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if v0.NDim() != 2 || v0.Shape()[0] != rows || v0.Shape()[1] != cols {
|
||
return nil, base.Errf("%s: the velocity must be rank 2 on the same grid, got shape %s", name, base.ShapeText(v0.Shape()))
|
||
}
|
||
if v0.Dtype() == core.Complex {
|
||
return nil, base.Errf("%s: complex velocities are not supported", name)
|
||
}
|
||
// As for u0: a non-finite velocity enters the Taylor start and
|
||
// poisons the three-level stencil without an error.
|
||
for i := range v0.Len() {
|
||
if v := v0.FloatAt(i); math.IsNaN(v) || math.IsInf(v, 0) {
|
||
return nil, base.Errf("%s: the velocity holds the non-finite value %g at %d", name, v, i)
|
||
}
|
||
}
|
||
if !(c > 0) {
|
||
return nil, base.Errf("%s: the wave speed must be positive, got %g", name, c)
|
||
}
|
||
cfl := c * dt * math.Sqrt(1/(dx*dx)+1/(dy*dy))
|
||
if cfl > 1 {
|
||
return nil, base.Errf("%s: the CFL number %g exceeds 1 (c = %g, dt = %g, dx = %g, dy = %g); the explicit stencil is unstable there",
|
||
name, cfl, c, dt, dx, dy)
|
||
}
|
||
steps, h := pdeSchedule(tFinal, dt, samples)
|
||
lapW := (c * h) * (c * h)
|
||
// The stencil's spacings are constants of the grid; squaring them
|
||
// once per step is the same product the element loop evaluated.
|
||
dx2 := dx * dx
|
||
dy2 := dy * dy
|
||
|
||
prev := make([]float64, rows*cols) // u^{s-1}
|
||
cur := make([]float64, rows*cols) // u^{s}
|
||
next := make([]float64, rows*cols) // u^{s+1}, the buffer being written
|
||
copy(cur, denseFloats(u0))
|
||
every := steps / (samples - 1)
|
||
out := core.New(core.Float, samples, rows, cols)
|
||
into := out.RawFloats()
|
||
// Slot 0 is the initial state at t = 0; slot j is the state after
|
||
// j·every steps, at the published time j·tFinal/(samples-1), exactly
|
||
// like IntegrateHeat2D.
|
||
copy(into, cur)
|
||
// From there on the boundary ring is held at zero, so clear it in the
|
||
// working buffer the initial state came in (the other two are born
|
||
// zeroed): otherwise a recycled buffer would carry the caller's ring
|
||
// back into the stencil every third step.
|
||
for r := range rows {
|
||
cur[r*cols], cur[r*cols+cols-1] = 0, 0
|
||
}
|
||
for c := range cols {
|
||
cur[c], cur[(rows-1)*cols+c] = 0, 0
|
||
}
|
||
written := 1
|
||
// The three buffers cycle through (prev, cur, next), so every stencil
|
||
// reads u^{s-1} and u^s while writing u^{s+1} into a third buffer: no
|
||
// point is ever overwritten before its neighbours have read it. The
|
||
// first iteration is the Taylor start
|
||
// u¹ = u⁰ + h·v⁰ + h²c²/2·Δu⁰, second order, so the three-level
|
||
// stencil starts honest, and it reads the untouched u⁰ (the initial
|
||
// velocity and the Laplacian of u⁰), not a half-updated state.
|
||
for s := 1; s <= steps; s++ {
|
||
first := s == 1
|
||
for r := 1; r < rows-1; r++ {
|
||
// The row's neighbours are the slices above and below it, so
|
||
// the stencil reads them by offset instead of rebuilding the
|
||
// flat index per element. The terms keep their order.
|
||
mid := cur[r*cols : (r+1)*cols]
|
||
up := cur[(r+1)*cols : (r+2)*cols]
|
||
down := cur[(r-1)*cols : r*cols]
|
||
if first {
|
||
// The Taylor start reads the untouched u⁰ and its
|
||
// initial velocity, not a half-updated state.
|
||
for cc := 1; cc < cols-1; cc++ {
|
||
i := r*cols + cc
|
||
lap := (down[cc]-2*mid[cc]+up[cc])/dy2 +
|
||
(mid[cc-1]-2*mid[cc]+mid[cc+1])/dx2
|
||
next[i] = mid[cc] + h*v0.FloatAt(i) + 0.5*lapW*lap
|
||
}
|
||
continue
|
||
}
|
||
for cc := 1; cc < cols-1; cc++ {
|
||
i := r*cols + cc
|
||
lap := (down[cc]-2*mid[cc]+up[cc])/dy2 +
|
||
(mid[cc-1]-2*mid[cc]+mid[cc+1])/dx2
|
||
next[i] = 2*mid[cc] - prev[i] + lapW*lap
|
||
}
|
||
}
|
||
prev, cur, next = cur, next, prev
|
||
if s%every == 0 && written < samples {
|
||
copy(into[written*rows*cols:(written+1)*rows*cols], cur)
|
||
written++
|
||
}
|
||
}
|
||
copy(into[(samples-1)*rows*cols:], cur)
|
||
return out, nil
|
||
}
|
||
|
||
// pdeStepFloor is the per-worker floor for a 2-D step sweep. One item is
|
||
// one tridiagonal elimination over a stencil row, measured at roughly
|
||
// half a microsecond on a 32-line grid, so a worker needs about eight
|
||
// of them before the split pays for its spawn. The floor is a constant
|
||
// for that reason, not a length-scaled budget. The line scratch is
|
||
// indexed by start/pdeStepFloor on the same constant: parallel chunk
|
||
// starts differ by at least the floor, so concurrent chunks never share
|
||
// a scratch set.
|
||
const pdeStepFloor = 8
|