feat: initial release

Assisted-by: GLM 5.3 Flash

stats/histogram2d.go

@@ -0,0 +1,214 @@
+// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
+// SPDX-License-Identifier: MIT
+
+package stats
+
+import (
+	"math"
+	"sync"
+
+	"sourcedock.dev/petrbalvin/tensor/internal/base"
+	"sourcedock.dev/petrbalvin/tensor/internal/core"
+	"sourcedock.dev/petrbalvin/tensor/internal/engine"
+)
+
+// Histogram2D bins the paired samples (x[i], y[i]) into a
+// xBins × yBins count matrix over equal-width bins spanning each
+// axis's own data range, mirroring Histogram's conventions: bins are
+// closed on the left and the last bin absorbs the maximum. The count
+// matrix comes back as a (xBins × yBins) array, the edges as the
+// bin boundaries on each axis. Real inputs only, and the bin counts
+// are bounded by maxHistBins in total, exactly as Histogram's are.
+func Histogram2D(x, y *core.Array, xBins, yBins int) (*core.Array, []float64, []float64, error) {
+	const name = "Histogram2D"
+	if x.Len() != y.Len() {
+		return nil, nil, nil, base.Errf("%s: x and y must share their length, got %d and %d", name, x.Len(), y.Len())
+	}
+	if x.Len() == 0 {
+		return nil, nil, nil, base.Errf("%s: empty samples have no histogram", name)
+	}
+	if x.Dtype() == core.Complex || y.Dtype() == core.Complex {
+		return nil, nil, nil, base.Errf("%s: complex inputs are not supported", name)
+	}
+	if xBins < 1 || yBins < 1 {
+		return nil, nil, nil, base.Errf("%s: needs at least one bin per axis, got %d and %d", name, xBins, yBins)
+	}
+	// The feasibility check runs before any allocation: the count matrix
+	// costs xBins*yBins cells, and refused up front it can neither
+	// overflow the product nor ask the allocator for an unbounded
+	// buffer.
+	if xBins > maxHistBins || yBins > maxHistBins || xBins > maxHistBins/yBins {
+		return nil, nil, nil, base.Errf("%s: the %d x %d bin request exceeds the %d-bin limit",
+			name, xBins, yBins, maxHistBins)
+	}
+	xEdges, err := edgesFor(x, xBins, name, "x")
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	yEdges, err := edgesFor(y, yBins, name, "y")
+	if err != nil {
+		return nil, nil, nil, err
+	}
+	counts := core.New(core.Int, xBins, yBins)
+	ints := counts.RawInts()
+	// The paired samples are read once: a dense float64 payload is walked
+	// in place and any other layout is widened into a private slice, so
+	// the sweep below reads no accessor and every value is the one
+	// FloatAt returns.
+	xVals, yVals := histogramAxis(x), histogramAxis(y)
+	// Counting a cell is exact integer arithmetic and every sample
+	// carries exactly one increment, so the sweep splits over disjoint
+	// slices of the sample and the private counters merge in any order
+	// into the totals a single pass produces.
+	bins := xBins * yBins
+	minPerWorker := max(hist2DSerialSamples, hist2DBinsPerWorker*bins)
+	// The first non-finite pair is the one the serial sweep would name:
+	// the chunks report their own first, and the lowest index wins.
+	bad := -1
+	var mu sync.Mutex
+	note := func(i int) {
+		if i < 0 {
+			return
+		}
+		mu.Lock()
+		if bad < 0 || i < bad {
+			bad = i
+		}
+		mu.Unlock()
+	}
+	engine.ParallelMin(x.Len(), minPerWorker, func(start, end int) {
+		if start == 0 && end == x.Len() {
+			// The whole sample runs inline, the worker policy having
+			// found the split uneconomical: count straight into the
+			// result.
+			note(histogram2DCount(xVals, yVals, 0, xEdges, yEdges, xBins, yBins, ints))
+			return
+		}
+		local := make([]int64, bins)
+		first := histogram2DCount(xVals[start:end], yVals[start:end], start, xEdges, yEdges, xBins, yBins, local)
+		mu.Lock()
+		for i, c := range local {
+			ints[i] += c
+		}
+		mu.Unlock()
+		note(first)
+	})
+	if bad >= 0 {
+		return nil, nil, nil, base.Errf("%s: sample %d (%g, %g) is not finite", name, bad, xVals[bad], yVals[bad])
+	}
+	return counts, xEdges, yEdges, nil
+}
+
+// hist2DBinsPerWorker is the paired-sample count one counting worker
+// must carry per cell before the binning splits across goroutines, and
+// hist2DSerialSamples the count it must carry whatever the cell count:
+// a worker counts into a private matrix of xBins·yBins cells, so the
+// split only pays where the samples drained into it outweigh the cells
+// it costs.
+const (
+	hist2DSerialSamples = 1 << 12
+	hist2DBinsPerWorker = 8
+)
+
+// histogramAxis returns the axis's values as a plain slice, walking a
+// dense float64 payload in place. Element i sits at payload index i,
+// and a rebased view's payload may run past its own count, so the view
+// is cut to the visible elements.
+func histogramAxis(a *core.Array) []float64 {
+	if fs := rawFloats(a); fs != nil {
+		return fs[:a.Len()]
+	}
+	vals := make([]float64, a.Len())
+	for i := range vals {
+		vals[i] = a.FloatAt(i)
+	}
+	return vals
+}
+
+// histogram2DCount adds one slice of the paired sample to counts and
+// returns the index of the first non-finite pair it met, or -1. The bin
+// arithmetic is binOf's own, so a split of the sample changes no count.
+func histogram2DCount(xv, yv []float64, base int, xEdges, yEdges []float64, xBins, yBins int, counts []int64) int {
+	xBase, xStep := xEdges[0], xEdges[1]-xEdges[0]
+	yBase, yStep := yEdges[0], yEdges[1]-yEdges[0]
+	for i := range xv {
+		xi := binOf(xv[i], xBase, xStep, xBins)
+		yi := binOf(yv[i], yBase, yStep, yBins)
+		if xi < 0 || yi < 0 {
+			return base + i
+		}
+		counts[xi*yBins+yi]++
+	}
+	return -1
+}
+
+// edgesFor builds one axis's bin edges over the data's own range.
+func edgesFor(a *core.Array, bins int, name, axis string) ([]float64, error) {
+	lo, hi := math.Inf(1), math.Inf(-1)
+	// The finiteness scan and the range share one pass over a dense
+	// float64 payload, the values being the ones FloatAt returns; every
+	// other layout keeps the accessor scan and the range pass after it.
+	// The payload walk is bounded by Len: a rebased view's payload may
+	// run past its own count.
+	if fs := rawFloats(a); fs != nil {
+		for i, v := range fs[:a.Len()] {
+			if math.IsNaN(v) || math.IsInf(v, 0) {
+				return nil, base.Errf("%s: %s sample %d is not finite (%g)", name, axis, i, v)
+			}
+			if v < lo {
+				lo = v
+			}
+			if v > hi {
+				hi = v
+			}
+		}
+	} else {
+		for i := range a.Len() {
+			if v := a.FloatAt(i); math.IsNaN(v) || math.IsInf(v, 0) {
+				return nil, base.Errf("%s: %s sample %d is not finite (%g)", name, axis, i, v)
+			}
+		}
+		for i := range a.Len() {
+			v := a.FloatAt(i)
+			if v < lo {
+				lo = v
+			}
+			if v > hi {
+				hi = v
+			}
+		}
+	}
+	if lo == hi {
+		lo -= 0.5
+		hi += 0.5
+	}
+	width := (hi - lo) / float64(bins)
+	if math.IsInf(width, 0) || math.IsNaN(width) {
+		// Both float extremes in one axis span more than the float64
+		// range: no finite edges exist, and the bin arithmetic would
+		// clamp every sample into bin 0 over ±Inf edges.
+		return nil, base.Errf("%s: the %s samples span more than the float64 range (%g to %g)", name, axis, lo, hi)
+	}
+	edges := make([]float64, bins+1)
+	for i := range bins + 1 {
+		edges[i] = lo + float64(i)*width
+	}
+	return edges, nil
+}
+
+// binOf locates one sample's bin from the axis's base edge and the step
+// between its first two edges, closing the last bin on the right; a
+// non-finite sample returns −1.
+func binOf(v, base, step float64, bins int) int {
+	if math.IsNaN(v) || math.IsInf(v, 0) {
+		return -1
+	}
+	i := int((v - base) / step)
+	if i < 0 {
+		return 0 // cannot happen over the data's own range, kept guarded
+	}
+	if i >= bins {
+		i = bins - 1
+	}
+	return i
+}