feat: initial release

Assisted-by: GLM 5.3 Flash

signal/stft.go

@@ -0,0 +1,158 @@
+// 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"
+	"sourcedock.dev/petrbalvin/tensor/internal/engine"
+)
+
+// The short-time Fourier transform: the signal is cut into
+// overlapping segments, each segment is windowed and transformed, and
+// the frames stack into a time-by-frequency picture. The spectrogram
+// is the magnitude squared of it, scaled per frame exactly as Welch's
+// estimate scales its periodograms, so averaging a spectrogram over
+// its frames reproduces WelchPSD on the same segmentation.
+
+// STFTOptions tunes the transform. Segment is the length of one
+// frame in samples, Overlap the samples two neighbouring frames share
+// (segment/2 is the usual choice) and Window names the taper: "hann",
+// "hamming" or "box", with "hann" implied by the zero value.
+type STFTOptions struct {
+	Segment int
+	Overlap int
+	Window  string
+}
+
+// framesOf computes the frame count of a signal of n samples under
+// the segment and overlap geometry, the count Welch's estimate uses.
+func framesOf(n, segment, overlap int) int {
+	return (n - overlap) / (segment - overlap)
+}
+
+// STFT transforms the vector x frame by frame and returns the
+// complex frames as a (frames × segment) array in row order: row t
+// holds the spectrum of the t-th frame, bin k of frame t is the
+// transform of x[t·hop : t·hop+segment] under the window, with hop =
+// segment − overlap. The Fourier definition treats each windowed
+// segment as one period, which is the standard convention here.
+func STFT(x *core.Array, opts STFTOptions) (*core.Array, error) {
+	const name = "STFT"
+	if x.NDim() != 1 {
+		return nil, base.Errf("%s: the signal must be a vector, got shape %s", name, base.ShapeText(x.Shape()))
+	}
+	if x.Dtype() == core.Complex {
+		return nil, base.Errf("%s: complex arrays are not supported", name)
+	}
+	n := x.Len()
+	if opts.Segment < 2 || opts.Segment > n {
+		return nil, base.Errf("%s: segment must be between 2 and %d, got %d", name, n, opts.Segment)
+	}
+	if opts.Overlap < 0 || opts.Overlap >= opts.Segment {
+		return nil, base.Errf("%s: overlap must be in [0, %d), got %d", name, opts.Segment, opts.Overlap)
+	}
+	window := opts.Window
+	if window == "" {
+		window = "hann"
+	}
+	taper, err := windowTaper(window, opts.Segment)
+	if err != nil {
+		return nil, base.Errf("%s: %w", name, err)
+	}
+	hop := opts.Segment - opts.Overlap
+	frames := framesOf(n, opts.Segment, opts.Overlap)
+	if frames < 1 {
+		return nil, base.Errf("%s: %d samples fill only one %d-sample segment with overlap %d",
+			name, n, opts.Segment, opts.Overlap)
+	}
+	src := widenFloats(x)
+	segment := opts.Segment
+	out, oerr := core.Zeros(core.Complex, frames, segment)
+	if oerr != nil {
+		return nil, oerr
+	}
+	// A frame's windowed samples go straight into its output row and are
+	// transformed there: the row is a private, fully overwritten buffer,
+	// so the copy a separate scratch segment needed is one pass fewer and
+	// the bits land exactly where the copy would have put them.
+	rows := out.RawComplexes()[:out.Len()]
+	engine.ParallelMin(frames, welchMinSegmentsPerWorker, func(fs, fe int) {
+		for f := fs; f < fe; f++ {
+			start := f * hop
+			z := rows[f*segment : (f+1)*segment]
+			for i := range segment {
+				z[i] = complex(src[start+i]*taper[i], 0)
+			}
+			transform(z, -1)
+		}
+	})
+	return out, nil
+}
+
+// Spectrogram returns the one-sided power spectrogram of the vector x
+// sampled at fs hertz as a (frames × segment/2+1) array: every frame
+// carries the periodogram of its windowed segment with the same
+// one-sided doubling and window-power normalisation Welch's estimate
+// applies, in units of x²/Hz. Averaging the frames reproduces
+// WelchPSD on the same geometry.
+func Spectrogram(x *core.Array, fs float64, opts STFTOptions) (*core.Array, error) {
+	const name = "Spectrogram"
+	if !(fs > 0) || math.IsInf(fs, 0) {
+		return nil, base.Errf("%s: fs must be positive and finite, got %g", name, fs)
+	}
+	// The segment geometry is checked before the taper is built: the
+	// taper allocates Segment samples, so a hostile segment must be
+	// refused here rather than after the allocation. STFT repeats the
+	// same checks on its own terms.
+	n := x.Len()
+	if opts.Segment < 2 || opts.Segment > n {
+		return nil, base.Errf("%s: segment must be between 2 and %d, got %d", name, n, opts.Segment)
+	}
+	if opts.Overlap < 0 || opts.Overlap >= opts.Segment {
+		return nil, base.Errf("%s: overlap must be in [0, %d), got %d", name, opts.Segment, opts.Overlap)
+	}
+	window := opts.Window
+	if window == "" {
+		window = "hann"
+	}
+	taper, err := windowTaper(window, opts.Segment)
+	if err != nil {
+		return nil, base.Errf("%s: %w", name, err)
+	}
+	wPower := 0.0
+	for i := range taper {
+		wPower += taper[i] * taper[i]
+	}
+	if wPower == 0 {
+		return nil, base.Errf("%s: the window has zero power", name)
+	}
+	framesSpec, err := STFT(x, opts)
+	if err != nil {
+		return nil, base.Errf("%s: %w", name, err)
+	}
+	frames := framesSpec.Shape()[0]
+	segment := opts.Segment
+	bins := segment/2 + 1
+	scale := 1 / (fs * wPower)
+	out, oerr := core.Zeros(core.Float, frames, bins)
+	if oerr != nil {
+		return nil, oerr
+	}
+	spec := framesSpec.RawComplexes()[:framesSpec.Len()]
+	vals := out.RawFloats()
+	for f := range frames {
+		for k := range bins {
+			mag := base.AbsComplex(spec[f*segment+k])
+			p := mag * mag * scale
+			if k > 0 && k < segment-k {
+				p *= 2
+			}
+			vals[f*bins+k] = p
+		}
+	}
+	return out, nil
+}