Files
tensor/internal/core/diff.go
T

104 lines
3.3 KiB
Go
Raw Normal View History

2026-09-03 10:00:00 +02:00
// 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]
}
}
}
}