104 lines
3.3 KiB
Go
104 lines
3.3 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
||
// SPDX-License-Identifier: MIT
|
||
|
||
package core
|
||
|
||
// Discrete differences: the workhorse behind finite-difference
|
||
// derivatives and signal detrending, along any single axis.
|
||
|
||
// Diff takes the successive differences along one axis, order times:
|
||
// order 1 is out[i] = a[i+1] − a[i] along the axis, order 2 applies
|
||
// it again, and so on. The axis shrinks by order; the axis must
|
||
// therefore hold more elements than the order, and axis must name one
|
||
// of the array's axes. Complex arrays are fine: differences carry no
|
||
// ordering assumption. Int arrays keep their dtype, because the
|
||
// difference of two int64 is the int64 difference; everything else
|
||
// produces float.
|
||
func Diff(a *Array, order, axis int) (*Array, error) {
|
||
if order < 1 {
|
||
return nil, errf("Diff: the order must be at least 1, got %d", order)
|
||
}
|
||
if axis < 0 || axis >= a.NDim() {
|
||
return nil, errf("Diff: axis %d is outside the %d axes of shape %s", axis, a.NDim(), shapeText(a.Shape()))
|
||
}
|
||
if a.Shape()[axis] <= order {
|
||
return nil, errf("Diff: axis %d holds %d elements, more than the order %d is needed",
|
||
axis, a.Shape()[axis], order)
|
||
}
|
||
cur := a
|
||
for range order {
|
||
next, err := diffOnce(cur, axis)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
cur = next
|
||
}
|
||
return cur, nil
|
||
}
|
||
|
||
// diffOnce applies one round of differences along the axis. Int
|
||
// differences stay int64 and complex ones stay complex; every other
|
||
// dtype produces float, the only route that used to widen the int side
|
||
// through float64 and round neighbours above 2^53 together.
|
||
//
|
||
// The output is walked run by run: for one position along the trailing
|
||
// dimensions the two neighbours sit a fixed stride apart, so a run of
|
||
// the output is a plain elementwise subtraction with the dtype
|
||
// dispatched once, not a per-element coordinate fold.
|
||
func diffOnce(a *Array, axis int) (*Array, error) {
|
||
if !a.isContiguous() {
|
||
// A strided view's payload window is not the run the walk needs,
|
||
// so reduce it to a dense copy first; the elements, and with
|
||
// them the differences, are the ones the accessors returned.
|
||
a = a.materialise()
|
||
}
|
||
shape := a.Shape()
|
||
outShape := append([]int{}, shape...)
|
||
outShape[axis]--
|
||
dt := Float
|
||
switch a.dt {
|
||
case Complex:
|
||
dt = Complex
|
||
case Int:
|
||
dt = Int
|
||
}
|
||
out := &Array{shape: outShape, dt: dt}
|
||
out.alloc(out.Len())
|
||
tail := 1
|
||
for d := axis + 1; d < len(shape); d++ {
|
||
tail *= shape[d]
|
||
}
|
||
head := 1
|
||
for d := range axis {
|
||
head *= shape[d]
|
||
}
|
||
n := shape[axis]
|
||
switch dt {
|
||
case Int:
|
||
diffRuns(out.ints, a.ints[:a.Len()], head, n, tail)
|
||
case Complex:
|
||
diffRuns(out.complexes, a.complexes[:a.Len()], head, n, tail)
|
||
default:
|
||
// The source widens exactly as FloatAt widens it; a float64
|
||
// source is read in place.
|
||
diffRuns(out.floats, floatPayload(a), head, n, tail)
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
// diffRuns fills dst with the successive differences of src along an
|
||
// axis of n elements that steps by tail elements, for each of head
|
||
// outer positions.
|
||
func diffRuns[T int64 | float64 | complex128](dst, src []T, head, n, tail int) {
|
||
for h := range head {
|
||
base := h * n * tail
|
||
dstBase := h * (n - 1) * tail
|
||
for i := range tail {
|
||
s, d := base+i, dstBase+i
|
||
for k := range n - 1 {
|
||
dst[d+k*tail] = src[s+(k+1)*tail] - src[s+k*tail]
|
||
}
|
||
}
|
||
}
|
||
}
|