243 lines
8.8 KiB
Go
243 lines
8.8 KiB
Go
// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
package core
|
|
|
|
import (
|
|
"math/bits"
|
|
|
|
"sourcedock.dev/petrbalvin/tensor/internal/engine"
|
|
)
|
|
|
|
// Quasi-random sequences: Halton points, the low-
|
|
// discrepancy workhorse of Monte Carlo that needs no tables. Each
|
|
// coordinate runs the radical inverse in its own prime base, which
|
|
// stratifies every b^k block of points evenly through the unit
|
|
// hypercube: the property random sampling only has in expectation.
|
|
|
|
// HaltonPoints returns the first n Halton points of the given
|
|
// dimension as an (n, dim) float64 array, skipping the leading skip
|
|
// points (the early Halton coordinates correlate visibly in high
|
|
// dimensions; the standard cure is to drop them). dim must be at
|
|
// least 1 and at most 32: beyond that the available small primes
|
|
// run out and the stratification degrades; a larger request is an
|
|
// error, not a silently worse sequence. n+skip must stay below 2^32,
|
|
// the index budget both quasi-random constructors enforce; past it
|
|
// the point indices wrap the int arithmetic and every point collapses
|
|
// back to the origin instead of advancing.
|
|
func HaltonPoints(n, dim, skip int) (*Array, error) {
|
|
const name = "HaltonPoints"
|
|
if n < 0 {
|
|
return nil, errf("%s: n must be non-negative, got %d", name, n)
|
|
}
|
|
if dim < 1 || dim > 32 {
|
|
return nil, errf("%s: dim must lie in [1, 32], got %d", name, dim)
|
|
}
|
|
if skip < 0 {
|
|
return nil, errf("%s: skip must be non-negative, got %d", name, skip)
|
|
}
|
|
if uint64(skip)+uint64(n) >= 1<<32 {
|
|
return nil, errf("%s: n + skip must stay below 2^32, the index budget, got %d + %d",
|
|
name, n, skip)
|
|
}
|
|
bases := firstPrimes(dim)
|
|
out := &Array{shape: []int{n, dim}, dt: Float}
|
|
out.alloc(n * dim)
|
|
// The points are independent, so the walk splits over disjoint
|
|
// point ranges: every point writes its own slots and no value is
|
|
// accumulated, so the split cannot move a single coordinate.
|
|
engine.ParallelMin(n, copyMinPerWorker, func(s, e int) {
|
|
for p := s; p < e; p++ {
|
|
idx := p + skip + 1 // the point indexed 0 would be the origin
|
|
for d := range dim {
|
|
out.floats[p*dim+d] = radicalInverse(idx, bases[d])
|
|
}
|
|
}
|
|
})
|
|
return out, nil
|
|
}
|
|
|
|
// radicalInverse reflects the base-b digits of i through the radix
|
|
// point: the van der Corput core of every Halton coordinate.
|
|
func radicalInverse(i int, b int) float64 {
|
|
f := 1.0
|
|
r := 0.0
|
|
for i > 0 {
|
|
f /= float64(b)
|
|
r += f * float64(i%b)
|
|
i /= b
|
|
}
|
|
return r
|
|
}
|
|
|
|
// firstPrimes returns the first count primes by trial division.
|
|
func firstPrimes(count int) []int {
|
|
primes := make([]int, 0, count)
|
|
candidate := 2
|
|
for len(primes) < count {
|
|
isPrime := true
|
|
for _, p := range primes {
|
|
if p*p > candidate {
|
|
break
|
|
}
|
|
if candidate%p == 0 {
|
|
isPrime = false
|
|
break
|
|
}
|
|
}
|
|
if isPrime {
|
|
primes = append(primes, candidate)
|
|
}
|
|
candidate++
|
|
}
|
|
return primes
|
|
}
|
|
|
|
// Sobol sequences: the digital low-discrepancy companion to
|
|
// Halton. Each coordinate runs its own linear recurrence over GF(2)
|
|
// driven by direction numbers derived from a primitive polynomial, so
|
|
// unlike Halton every dimension shares the same base-2 lattice and the
|
|
// first 2^k points are exactly stratified through every coordinate,
|
|
// not just on average. The parameters below are the Joe and Kuo (2008)
|
|
// initialisation table, the current standard, for the first 40
|
|
// dimensions; dimension 1 is the plain Gray-coded van der Corput
|
|
// sequence and carries no polynomial.
|
|
|
|
// sobolParams are the initialisation parameters of one dimension: the
|
|
// polynomial degree s, the primitive polynomial coefficient a whose
|
|
// bits select the recurrence taps, and the s odd initialisation
|
|
// integers m_i with 1 <= m_i < 2^i.
|
|
type sobolParams struct {
|
|
s uint32
|
|
a uint32
|
|
m []uint32
|
|
}
|
|
|
|
var sobolTable = [...]sobolParams{
|
|
{0, 0, nil}, // 1: plain van der Corput
|
|
{1, 0, []uint32{1}}, // 2
|
|
{2, 1, []uint32{1, 3}}, // 3
|
|
{3, 1, []uint32{1, 3, 1}}, // 4
|
|
{3, 2, []uint32{1, 1, 1}}, // 5
|
|
{4, 1, []uint32{1, 1, 3, 3}}, // 6
|
|
{4, 4, []uint32{1, 3, 5, 13}}, // 7
|
|
{5, 2, []uint32{1, 1, 5, 5, 17}}, // 8
|
|
{5, 4, []uint32{1, 1, 5, 5, 5}}, // 9
|
|
{5, 7, []uint32{1, 1, 7, 11, 19}}, // 10
|
|
{5, 11, []uint32{1, 1, 5, 1, 1}}, // 11
|
|
{5, 13, []uint32{1, 1, 1, 3, 11}}, // 12
|
|
{5, 14, []uint32{1, 3, 5, 5, 31}}, // 13
|
|
{6, 1, []uint32{1, 3, 3, 9, 7, 49}}, // 14
|
|
{6, 13, []uint32{1, 1, 1, 15, 21, 21}}, // 15
|
|
{6, 16, []uint32{1, 3, 1, 13, 27, 49}}, // 16
|
|
{6, 19, []uint32{1, 1, 1, 15, 7, 5}}, // 17
|
|
{6, 22, []uint32{1, 3, 1, 15, 13, 25}}, // 18
|
|
{6, 25, []uint32{1, 1, 5, 5, 19, 61}}, // 19
|
|
{7, 1, []uint32{1, 3, 7, 11, 23, 15, 103}}, // 20
|
|
{7, 4, []uint32{1, 3, 7, 13, 13, 15, 69}}, // 21
|
|
{7, 7, []uint32{1, 1, 3, 13, 7, 35, 63}}, // 22
|
|
{7, 8, []uint32{1, 3, 5, 9, 1, 25, 53}}, // 23
|
|
{7, 14, []uint32{1, 3, 1, 13, 9, 35, 107}}, // 24
|
|
{7, 19, []uint32{1, 3, 1, 5, 27, 61, 31}}, // 25
|
|
{7, 21, []uint32{1, 1, 5, 11, 19, 41, 61}}, // 26
|
|
{7, 28, []uint32{1, 3, 5, 3, 3, 13, 69}}, // 27
|
|
{7, 31, []uint32{1, 1, 7, 13, 1, 19, 1}}, // 28
|
|
{7, 32, []uint32{1, 3, 7, 5, 13, 19, 59}}, // 29
|
|
{7, 37, []uint32{1, 1, 3, 9, 25, 29, 41}}, // 30
|
|
{7, 41, []uint32{1, 3, 5, 13, 23, 1, 55}}, // 31
|
|
{7, 42, []uint32{1, 3, 7, 3, 13, 59, 17}}, // 32
|
|
{7, 50, []uint32{1, 3, 1, 3, 5, 53, 69}}, // 33
|
|
{7, 55, []uint32{1, 1, 5, 5, 23, 33, 13}}, // 34
|
|
{7, 56, []uint32{1, 1, 7, 7, 1, 61, 123}}, // 35
|
|
{7, 59, []uint32{1, 1, 7, 9, 13, 61, 49}}, // 36
|
|
{7, 62, []uint32{1, 3, 3, 5, 3, 55, 33}}, // 37
|
|
{8, 14, []uint32{1, 3, 1, 15, 31, 13, 49, 245}}, // 38
|
|
{8, 21, []uint32{1, 3, 5, 15, 31, 59, 63, 97}}, // 39
|
|
{8, 22, []uint32{1, 3, 1, 11, 11, 11, 77, 249}}, // 40
|
|
}
|
|
|
|
// sobolDirections fills v with the 32-bit direction numbers of the
|
|
// given dimension. The first s come straight from the initialisation
|
|
// integers scaled into their binary place; the rest follow the
|
|
// recurrence v_i = v_{i-s} ^ (v_{i-s} >> s) ^ Σ a_k·v_{i-k} over the
|
|
// polynomial taps.
|
|
func sobolDirections(p sobolParams, v []uint32) {
|
|
s := int(p.s)
|
|
if s == 0 {
|
|
// Dimension 1: the direction numbers are the binary places
|
|
// themselves, which makes the sequence Gray-coded van der Corput.
|
|
for i := range v {
|
|
v[i] = 1 << (31 - i)
|
|
}
|
|
return
|
|
}
|
|
for i := range s {
|
|
v[i] = p.m[i] << (31 - i)
|
|
}
|
|
for i := s; i < len(v); i++ {
|
|
v[i] = v[i-s] ^ (v[i-s] >> uint(s))
|
|
for k := 1; k < s; k++ {
|
|
if p.a>>(uint(s-1-k))&1 != 0 {
|
|
v[i] ^= v[i-k]
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// SobolPoints returns the first n Sobol points of the given dimension
|
|
// as an (n, dim) float64 array, skipping the leading skip points (the
|
|
// sequence starts at the origin index, which carries no information
|
|
// and is dropped, exactly as HaltonPoints drops it). dim must be at
|
|
// least 1 and at most 40: that is the width of the initialisation
|
|
// table, and a larger request is an error, not a silently worse
|
|
// sequence. n+skip must stay below 2^32, the period the 32-bit
|
|
// direction numbers give; beyond it the index arithmetic would wrap
|
|
// back to the origin instead of advancing.
|
|
func SobolPoints(n, dim, skip int) (*Array, error) {
|
|
const name = "SobolPoints"
|
|
if n < 0 {
|
|
return nil, errf("%s: n must be non-negative, got %d", name, n)
|
|
}
|
|
if dim < 1 || dim > len(sobolTable) {
|
|
return nil, errf("%s: dim must lie in [1, %d], got %d", name, len(sobolTable), dim)
|
|
}
|
|
if skip < 0 {
|
|
return nil, errf("%s: skip must be non-negative, got %d", name, skip)
|
|
}
|
|
if uint64(skip)+uint64(n) >= 1<<32 {
|
|
return nil, errf("%s: n + skip must stay below 2^32, the sequence period, got %d + %d",
|
|
name, n, skip)
|
|
}
|
|
const width = 32
|
|
out := &Array{shape: []int{n, dim}, dt: Float}
|
|
out.alloc(n * dim)
|
|
// Each coordinate walks its own points. The worker seeds its range's
|
|
// Gray code from scratch, then advances one flip at a time: the Gray
|
|
// codes of consecutive indices differ in the lowest set bit of the
|
|
// newer index, so x picks up exactly the direction numbers the
|
|
// from-scratch walk would XOR, and the exclusive or combines them
|
|
// exactly whatever the order.
|
|
engine.ParallelMin(n, copyMinPerWorker, func(s, e int) {
|
|
for d := range dim {
|
|
var v [width]uint32
|
|
sobolDirections(sobolTable[d], v[:])
|
|
idx := s + skip + 1 // the point indexed 0 would be the origin
|
|
gray := uint32(idx) ^ uint32(idx>>1)
|
|
var x uint32
|
|
for b := range width {
|
|
if gray&(1<<uint(b)) != 0 {
|
|
x ^= v[b]
|
|
}
|
|
}
|
|
for p := s; p < e; p++ {
|
|
out.floats[p*dim+d] = float64(x) / (1 << 32)
|
|
if p+1 < e {
|
|
idx++
|
|
x ^= v[bits.TrailingZeros32(uint32(idx))]
|
|
}
|
|
}
|
|
}
|
|
})
|
|
return out, nil
|
|
}
|