blue-sampler 0.1.0__tar.gz

blue_sampler-0.1.0/PKG-INFO

@@ -0,0 +1,71 @@
+Metadata-Version: 2.4
+Name: blue_sampler
+Version: 0.1.0
+Summary: Stealthy point-pattern sampling on the unit torus
+Project-URL: Repository, https://github.com/For-a-few-DPPs-more/hyperuniform-samplers
+License: MIT
+Keywords: blue noise,jax,point process,sampling,stealthy
+Requires-Python: >=3.10
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: matplotlib
+Requires-Dist: numpy
+Requires-Dist: requests
+Requires-Dist: squarenet
+Description-Content-Type: text/markdown
+
+# blue-sampler
+
+Generate **stealthy point patterns** — low-discrepancy, spectrally isotropic
+samples on the unit torus [0, 1)^D.
+
+Stealthy patterns suppress long-range density fluctuations while remaining
+aperiodic.  They are useful in rendering, quadrature, and computational
+physics wherever quasi-random, isotropic spatial coverage is needed.
+
+## Installation
+
+```bash
+pip install blue_sampler
+```
+
+## Quick start
+
+```python
+import blue_sampler as blue
+
+# 10 000 points in 2-D
+x = blue.sample(10_000)
+blue.plot(x)
+blue.plot_structure_factor(x)
+
+# 3-D
+x = blue.sample(5_000, D=3)
+```
+## Supported dimensions
+
+| D | Notes |
+|---|-------|
+| 2 | Fast, recommended for exploration |
+| 3 | ~3× slower than 2-D |
+| 4 | Requires more iterations (set automatically by `Config.auto`) |
+| 5 | Experimental |
+
+## Algorithm overview
+
+The pipeline alternates between:
+
+1. **Spatial gradient** — short-range Gaussian repulsion via
+   neighbour convolution on the torus.
+2. **Spectral gradient** — minimises the structure factor S(k) for k below
+   a chosen cut-off, using a set of all the wave vectors within an integer
+   half-ball.
+3. **Grid assignment** (SquareNet) — periodic re-assignment to a regular
+   grid for efficient sparse local operations.
+
+For N ≤ 3 000 a direct O(N²) bootstrap is used.  For larger N a
+hierarchical strategy clones and refines a coarser solution.
+
+## License
+
+MIT

blue_sampler-0.1.0/README.MD

@@ -0,0 +1,55 @@
+# blue-sampler
+
+Generate **stealthy point patterns** — low-discrepancy, spectrally isotropic
+samples on the unit torus [0, 1)^D.
+
+Stealthy patterns suppress long-range density fluctuations while remaining
+aperiodic.  They are useful in rendering, quadrature, and computational
+physics wherever quasi-random, isotropic spatial coverage is needed.
+
+## Installation
+
+```bash
+pip install blue_sampler
+```
+
+## Quick start
+
+```python
+import blue_sampler as blue
+
+# 10 000 points in 2-D
+x = blue.sample(10_000)
+blue.plot(x)
+blue.plot_structure_factor(x)
+
+# 3-D
+x = blue.sample(5_000, D=3)
+```
+## Supported dimensions
+
+| D | Notes |
+|---|-------|
+| 2 | Fast, recommended for exploration |
+| 3 | ~3× slower than 2-D |
+| 4 | Requires more iterations (set automatically by `Config.auto`) |
+| 5 | Experimental |
+
+## Algorithm overview
+
+The pipeline alternates between:
+
+1. **Spatial gradient** — short-range Gaussian repulsion via
+   neighbour convolution on the torus.
+2. **Spectral gradient** — minimises the structure factor S(k) for k below
+   a chosen cut-off, using a set of all the wave vectors within an integer
+   half-ball.
+3. **Grid assignment** (SquareNet) — periodic re-assignment to a regular
+   grid for efficient sparse local operations.
+
+For N ≤ 3 000 a direct O(N²) bootstrap is used.  For larger N a
+hierarchical strategy clones and refines a coarser solution.
+
+## License
+
+MIT

blue_sampler-0.1.0/README.md

@@ -0,0 +1,55 @@
+# blue-sampler
+
+Generate **stealthy point patterns** — low-discrepancy, spectrally isotropic
+samples on the unit torus [0, 1)^D.
+
+Stealthy patterns suppress long-range density fluctuations while remaining
+aperiodic.  They are useful in rendering, quadrature, and computational
+physics wherever quasi-random, isotropic spatial coverage is needed.
+
+## Installation
+
+```bash
+pip install blue_sampler
+```
+
+## Quick start
+
+```python
+import blue_sampler as blue
+
+# 10 000 points in 2-D
+x = blue.sample(10_000)
+blue.plot(x)
+blue.plot_structure_factor(x)
+
+# 3-D
+x = blue.sample(5_000, D=3)
+```
+## Supported dimensions
+
+| D | Notes |
+|---|-------|
+| 2 | Fast, recommended for exploration |
+| 3 | ~3× slower than 2-D |
+| 4 | Requires more iterations (set automatically by `Config.auto`) |
+| 5 | Experimental |
+
+## Algorithm overview
+
+The pipeline alternates between:
+
+1. **Spatial gradient** — short-range Gaussian repulsion via
+   neighbour convolution on the torus.
+2. **Spectral gradient** — minimises the structure factor S(k) for k below
+   a chosen cut-off, using a set of all the wave vectors within an integer
+   half-ball.
+3. **Grid assignment** (SquareNet) — periodic re-assignment to a regular
+   grid for efficient sparse local operations.
+
+For N ≤ 3 000 a direct O(N²) bootstrap is used.  For larger N a
+hierarchical strategy clones and refines a coarser solution.
+
+## License
+
+MIT

blue_sampler-0.1.0/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ["hatchling >= 1.26"]
+build-backend = "hatchling.build"
+
+[project]
+name        = "blue_sampler"
+version     = "0.1.0"
+description = "Stealthy point-pattern sampling on the unit torus"
+readme      = "README.md"
+license     = { text = "MIT" }
+requires-python = ">=3.10"
+
+keywords = [
+    "blue noise",
+    "stealthy",
+    "point process",
+    "sampling",
+    "jax",
+]
+
+dependencies = [
+    "numpy",
+    "jax",
+    "jaxlib",
+    "matplotlib",
+    "requests",
+    "squarenet",
+]
+
+[project.urls]
+Repository = "https://github.com/For-a-few-DPPs-more/hyperuniform-samplers"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/blue_sampler"]
+
+[tool.ruff]
+line-length = 100

blue_sampler-0.1.0/src/blue_sampler/__init__.py

@@ -0,0 +1,24 @@
+"""
+blue_sampler
+============
+
+Generate stealthy point patterns — low-discrepancy, spectrally isotropic
+samples on the unit torus [0, 1)^D.
+
+Quick start
+-----------
+>>> import blue_sampler as blue
+>>> x = blue.sample(N=10_000, D=2)          # (10000, 2) array
+>>> blue.plot(x)
+>>> blue.plot_structure_factor(x)
+"""
+from .run import sample
+from .viz import plot, plot_structure_factor
+from .math_utils import structure_factor
+
+__all__ = [
+    "sample",
+    "plot",
+    "plot_structure_factor",
+    "structure_factor",
+]

blue_sampler-0.1.0/src/blue_sampler/kernels.py

@@ -0,0 +1,100 @@
+"""
+Kernels for energy gradient
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+from .math_utils import torus_delta, clean_grad
+
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Kernel functions  (JAX)
+# ──────────────────────────────────────────────────────────────────────────────
+
+def gauss_kernel(
+    x: jnp.ndarray,
+    y: jnp.ndarray,
+    sigma2: float,
+) -> jnp.ndarray:
+    """
+    Isotropic Gaussian repulsion kernel on the torus.
+
+    Returns the *gradient* contribution  (y − x) * exp(−‖y−x‖² / σ²).
+    """
+    delta = torus_delta(y - x)
+    dist2 = jnp.sum(delta ** 2, axis=-1, keepdims=True)
+    return clean_grad(delta * jnp.exp(-dist2 / sigma2))
+
+
+def gauss_sin_kernel(
+    x: jnp.ndarray,
+    y: jnp.ndarray,
+    a: float,
+    b: float,
+    c: float,
+) -> jnp.ndarray:
+    """
+    Trigonometric Gaussian kernel — more stable in contexts were
+    sigma2 is not << 1 e.g. high dimension or low number of points
+    -> using discontinuous torus_delta would become problematic.
+
+    Parameters
+    ----------
+    a, b, c : pre-computed scale factors (derived from sigma²).
+    """
+    delta    = a * (y - x)
+    cos_term = b * (1.0 - jnp.cos(delta))
+    sin_term = c * jnp.sin(delta)
+    dist2 = jnp.sum(cos_term, axis=-1, keepdims=True)
+    return clean_grad(sin_term * jnp.exp(-dist2))
+
+def spectral_kernel(x, k, k_):
+    """
+    spectral kernel directly target spectral energy. Only usable
+    for small subsets of preselected wavevectors
+    """
+    phase   = jnp.sum(k * x, axis=-1, keepdims=True)
+    ek      = clean_grad(jnp.exp(phase))
+    Sk      = jnp.sum(ek, axis=0, keepdims=True)
+    return jnp.real(Sk * k_ * jnp.conjugate(ek))
+
+def reduce_kernel(kernel, x, params, init=None):
+    """
+    Generic reduction over a kernel.
+
+    Parameters
+    ----------
+    kernel : callable
+        Function of the form kernel(x, param) -> contribution.
+    x : array
+        State.
+    params : PyTree
+        Collection of parameters passed to kernel.
+    init : array, optional
+        Initial accumulator. Defaults to zeros_like(x).
+
+    Returns
+    -------
+    array
+        Sum of all kernel contributions.
+    """
+    if init is None:
+        init = jnp.zeros_like(x)
+
+    def body(acc, param):
+        return acc + kernel(x, param), None
+
+    out, _ = jax.lax.scan(body, init, params)
+    return out
+
+def micro_shift_kernel(x, shift, kernel, Axes):
+    contrib = kernel(x, jnp.roll(x, shift, axis=Axes))
+    return contrib - jnp.roll(contrib, -shift, axis=Axes)
+
+def micro_grad(x_val, SHIFTS, LR_spatial, S):
+    out = reduce_kernel(micro_shift_kernel, x_val, SHIFTS)
+    return (LR_spatial / S) * out

blue_sampler-0.1.0/src/blue_sampler/math_utils.py

@@ -0,0 +1,254 @@
+"""
+Low-level mathematical helpers
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Lattice helpers
+# ──────────────────────────────────────────────────────────────────────────────
+
+def drop_symmetric(directions: np.ndarray) -> np.ndarray:
+    """
+    Keep only one representative from each direction pair {v, -v}.
+
+    The canonical representative is the one whose *first non-zero component*
+    is positive.
+
+    Parameters
+    ----------
+    directions : (M, D) int array
+
+    Returns
+    -------
+    (K, D) int array  with K ≤ M // 2 + 1
+    """
+    first_nz_idx = (directions != 0).argmax(axis=1)
+    first_nz_val = directions[np.arange(len(directions)), first_nz_idx]
+    return directions[first_nz_val > 0]
+
+
+def integers_in_half_ball(radius: float, D: int) -> np.ndarray:
+    """
+    Return all non-zero integer lattice vectors inside a sphere of *radius*,
+    keeping only one vector per direction pair.
+
+    Parameters
+    ----------
+    radius : float
+    D : int
+
+    Returns
+    -------
+    (M, D) int32 array
+    """
+    if radius <= 0.9:
+        return np.zeros((0, D), dtype=np.int32)
+    if radius <= 1.9:
+        return np.eye(D, dtype=np.int32)
+
+    r   = np.arange(-radius, radius + 1)
+    pts = np.stack(np.meshgrid(*(r,) * D, indexing="ij"), axis=-1).reshape(-1, D)
+    d2  = np.sum(pts ** 2, axis=-1)
+    return drop_symmetric(pts[(d2 > 0) & (d2 <= radius ** 2)])
+
+
+def simplex(D: int) -> np.ndarray:
+    """
+    Vertices of a regular simplex centred at the origin in R^D.
+
+    Returns
+    -------
+    (D+1, D) float64 array
+    """
+    if D == 1:
+        return np.array([-1.0, 1.0])[:, None]
+    null = np.zeros((D, 1))
+    tip  = np.zeros((1, D))
+    tip[0, -1] = 1.0
+    base = np.hstack((simplex(D - 1), null))
+    return np.vstack((np.sqrt(1.0 - (1.0 / D) ** 2) * base - tip / D, tip))
+
+
+def grid_shape(N: int, D: int) -> tuple[tuple[int, ...], int, tuple[int, ...]]:
+    """
+    Smallest D-hypercube grid that contains at least *N* points.
+
+    Returns
+    -------
+    IJK   : shape tuple  e.g. (32, 32) for D=2
+    total : total number of grid slots  (I^D)
+    axes  : tuple(range(D))
+    """
+    I    = int(np.ceil(N ** (1.0 / D)))
+    IJK  = (I,) * D
+    return IJK, I ** D, tuple(range(D))
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Torus geometry  (JAX)
+# ──────────────────────────────────────────────────────────────────────────────
+
+def torus_wrap(x: jnp.ndarray) -> jnp.ndarray:
+    """Wrap coordinates into [0, 1)^D."""
+    return x - jnp.floor(x)
+
+
+def torus_delta(delta: jnp.ndarray) -> jnp.ndarray:
+    """Shortest signed displacement on the unit torus."""
+    return delta - jnp.round(delta)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Gradient / status helpers  (JAX)
+# ──────────────────────────────────────────────────────────────────────────────
+
+def clean_grad(x: jnp.ndarray) -> jnp.ndarray:
+    """Replace NaN gradient contributions (fictive points) with 0."""
+    return jnp.nan_to_num(x, nan=0.0)
+
+
+def clean_points(x: jnp.ndarray) -> jnp.ndarray:
+    """
+    Preserve the NaN status flag of empty grid slots after a torus wrap.
+
+    The last coordinate of each grid slot encodes whether the slot is real
+    (0.0) or fictive (NaN).  ``torus_wrap`` can corrupt this flag, so we
+    re-round it here.
+    """
+    return x.at[..., -1:].set(jnp.round(x[..., -1:]))
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Wave-vector preparation
+# ──────────────────────────────────────────────────────────────────────────────
+
+def prepare_wave_vectors(
+    Ks: np.ndarray,
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    """
+    Build JAX arrays for the spectral gradient.
+
+    Parameters
+    ----------
+    Ks : (M, D) integer wave-vector matrix
+
+    Returns
+    -------
+    K_w : complex array of shape (M, 1, D+1)  — phase multipliers
+    K_  : complex array of shape (M, 1, D+1)  — normalised duals
+    """
+    K  = 2.0 * jnp.pi * Ks * 1j
+    K  = jnp.concatenate((K, np.zeros((len(K), 1))), axis=1)[:, None, :]
+    Kn = (jnp.abs(K) ** 2).sum(axis=-1, keepdims=True)
+    return K, -K / Kn
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Grid initialisation helpers
+# ──────────────────────────────────────────────────────────────────────────────
+
+def prepare_points(
+    x: np.ndarray | None,
+    N_asked: int,
+    IJK: tuple[int, ...],
+    D: int,
+) -> jnp.ndarray:
+    """
+    Pad *N_asked* real points to fill the I^D grid.
+
+    Fictive slots receive a NaN status coordinate so gradients ignore them.
+
+    Parameters
+    ----------
+    x       : (N_asked, D) array or *None* (random initialisation).
+    N_asked : number of real points.
+    IJK     : grid shape tuple.
+    D       : spatial dimension.
+
+    Returns
+    -------
+    jnp.ndarray of shape (*IJK, D+1)
+    """
+    if x is None:
+        x = np.random.rand(N_asked, D)
+    else:
+        x = np.asarray(x).reshape(N_asked, D)
+
+    total             = int(np.prod(IJK))
+    xfull             = np.random.rand(total, D + 1)
+    xfull[:, -1]      = 0.0        # status = 0   → real
+    xfull[:N_asked, :D] = x
+    xfull[N_asked:, D]  = np.nan   # status = NaN → fictive
+    return jnp.array(xfull.reshape(*IJK, D + 1))
+
+def random_rotations(x, batch_size, Dout, Din):
+    Q, _      = np.linalg.qr(np.random.randn(batch_size, Dout, Din))
+    offsets   = np.einsum(
+        "nij,kj->nki", Q, x
+    )
+    return offsets
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Structure factor
+# ──────────────────────────────────────────────────────────────────────────────
+
+def structure_factor(
+    points: np.ndarray,
+    nbins: int = 100,
+    resolution: float = 30.0,
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Estimate the radial structure factor S(k) via scattering intensity.
+
+    Parameters
+    ----------
+    points     : (N, D) array of point coordinates in [0, 1)^D.
+    nbins      : number of radial bins.
+    resolution : how many random wave-vectors to sample per bin.
+
+    Returns
+    -------
+    k : (M,) float array — bin centres.
+    S : (M,) float array — mean S(k) per non-empty bin.
+    """
+    pts = np.asarray(points)
+    N, D = pts.shape
+
+    kmed = int(1_000 ** (1.0 / D))
+    kmax = int(2 * N ** (1.0 / D))
+    bins = np.linspace(0, kmax, nbins)
+
+    # Random + deterministic wave-vector sampling
+    nvecs = np.random.randint(-kmax, kmax + 1, size=(int(resolution * nbins), D))
+    nvecs = np.concatenate([nvecs, integers_in_half_ball(kmed, D)], axis=0)
+    nvecs = nvecs[np.any(nvecs != 0, axis=1)]
+
+    knorm   = np.linalg.norm(nvecs, axis=1)
+    bin_idx = np.searchsorted(bins, knorm) - 1
+    valid   = (bin_idx >= 0) & (bin_idx < len(bins) - 1)
+    nvecs, bin_idx = nvecs[valid], bin_idx[valid]
+
+    kvecs = jnp.array(2.0 * np.pi * nvecs)
+    pts_j = jnp.array(pts)
+
+    def Sk_one(k: jnp.ndarray) -> jnp.ndarray:
+        rho = jnp.sum(jnp.exp(1j * (pts_j @ k)), axis=0)
+        return jnp.abs(rho) ** 2 / N
+
+    Sk = np.asarray(jax.lax.map(Sk_one, kvecs))
+
+    n_bins = len(bins) - 1
+    S_sum  = np.bincount(bin_idx, weights=Sk, minlength=n_bins)
+    counts = np.bincount(bin_idx,             minlength=n_bins)
+
+    S      = np.zeros_like(S_sum, dtype=float)
+    nz     = counts > 0
+    S[nz]  = S_sum[nz] / counts[nz]
+
+    centres = 0.5 * (bins[:-1] + bins[1:])
+    return centres[nz], S[nz]

blue_sampler-0.1.0/src/blue_sampler/progress.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import sys
+import time
+
+
+class ProgressLogger:
+    """Hierarchical \r-based progress display for nested pipeline levels."""
+
+    def __init__(self, verbose: int):
+        self.verbose = verbose
+        self.level   = -1
+
+    def enter_level(self, N: int, D: int, N_ITER: int) -> _LevelCtx:
+        """Push a new recursion level and return its context."""
+        self.level += 1
+        return _LevelCtx(self, N, D, N_ITER)
+
+    def exit_level(self) -> None:
+        self.level -= 1
+
+    def _prefix(self) -> str:
+        return f"[L{self.level}] "
+
+    def write(self, msg: str, newline: bool = False) -> None:
+        if self.verbose < 1:
+            return
+        sys.stdout.write(f"\r{self._prefix()}{msg}    ")
+        if newline:
+            sys.stdout.write("\n")
+        sys.stdout.flush()
+
+
+class _LevelCtx:
+    """Tracks timing and tick state for a single pipeline level."""
+
+    def __init__(self, logger: ProgressLogger, N: int, D: int, N_ITER: int):
+        self._log    = logger
+        self.N       = N
+        self.D       = D
+        self.N_ITER  = N_ITER
+        self._tick   = 0
+        self._t0: float | None     = None
+        self._t_iter: float | None = None
+
+    def on_compile(self) -> None:
+        self._log.write("compiling JAX kernel…")
+
+    def on_bruteforce_start(self) -> None:
+        self._log.write(f"bruteforce  N={self.N}  D={self.D} …")
+
+    def on_bruteforce_done(self) -> None:
+        self._log.write("bruteforce done ✓", newline=True)
+
+    def tick(self) -> None:
+        """Called once per gridification callback (= one full iteration). Drives the ETA display."""
+        now = time.perf_counter()
+        self._tick += 1
+
+        if self._tick == 1:
+            self._t0 = now
+            self._log.write(f"{self._bar()} — calibrating…")
+            return
+
+        if self._tick == 2:
+            self._t_iter = now - self._t0  # type: ignore[operator]
+
+        self._log.write(f"{self._bar()} — {self._eta(now)} remaining")
+
+    def done(self) -> None:
+        self._log.write(f"{self._bar(done=True)} — done ✓", newline=True)
+
+    def _bar(self, done: bool = False) -> str:
+        filled = self.N_ITER if done else max(0, self._tick - 1)
+        W      = 20
+        n_fill = int(W * filled / self.N_ITER)
+        bar    = "▓" * n_fill + "░" * (W - n_fill)
+        return f"{filled}/{self.N_ITER} [{bar}]"
+
+    def _eta(self, now: float) -> str:
+        if self._t_iter is None:
+            return "?"
+        remaining = (self.N_ITER - (self._tick - 1)) * self._t_iter
+        return f"~{remaining:.0f}s" if remaining < 60 else f"~{remaining / 60:.1f}min"

blue_sampler-0.1.0/src/blue_sampler/run.py

@@ -0,0 +1,271 @@
+from __future__ import annotations
+
+import numpy as np
+import jax
+import jax.numpy as jnp
+from squarenet import SquareNet
+
+from .math_utils import (
+    integers_in_half_ball,
+    simplex,
+    grid_shape,
+    torus_wrap,
+    clean_points,
+    prepare_wave_vectors,
+    prepare_points,
+    random_rotations,
+)
+from .kernels import (
+    gauss_kernel,
+    gauss_sin_kernel,
+    spectral_kernel,
+)
+from .progress import ProgressLogger, _LevelCtx
+
+
+_PRESETS: dict[int, dict] = {
+    2: dict(spatial_radius=7, spectral_radius=7, LR_spatial=0.1,  LR_spectral=0.1, expension_factor=0.3, S=1.0),
+    3: dict(spatial_radius=5, spectral_radius=5, LR_spatial=0.1,  LR_spectral=0.1, expension_factor=0.3, S=1.0),
+    4: dict(spatial_radius=3, spectral_radius=3, LR_spatial=0.01, LR_spectral=0.1, expension_factor=1.0, S=0.5),
+}
+
+
+# ── Bruteforce (small N) ──────────────────────────────────────────────────────
+
+def _build_bruteforce(N: int, D: int, ctx: _LevelCtx):
+    """AOT-compile a gradient-descent sampler for N ≤ ~3 000 points."""
+    DX     = 1.0 / N ** (1.0 / D)
+    S      = 1.0
+    sigma2 = S * 2.0 * DX ** 2
+    high_D = sigma2 >= 0.03
+
+    lr_table = {2: 0.4, 3: 0.1, 4: 0.05, 5: 0.01}
+    lr    = lr_table.get(D, 0.01)
+    Niter = 1_000 if high_D else 3_000
+
+    if high_D:
+        a = 2.0 * jnp.pi
+        b = 2.0 / (sigma2 * a ** 2)
+        c = 1.0 / (2.0 * S * jnp.pi)
+        kernel = lambda x, y: gauss_sin_kernel(x, y, a, b, c)
+    else:
+        kernel = lambda x, y: gauss_kernel(x, y, sigma2)
+
+    def grad(x):
+        return jax.vmap(lambda xi: kernel(xi[None], x).sum(axis=0))(x)
+
+    @jax.jit
+    def _run(x):
+        def step(_, x):
+            return torus_wrap(x - lr * grad(x))
+        return jax.lax.fori_loop(0, Niter, step, x)
+
+    ctx.on_compile()
+    compiled = _run.lower(jax.ShapeDtypeStruct((N, D), jnp.float32)).compile()
+
+    def sample_fn(init: np.ndarray | None = None) -> jnp.ndarray:
+        ctx.on_bruteforce_start()
+        if init is None:
+            init = np.random.rand(N, D)
+        out = compiled(jnp.asarray(init))
+        out.block_until_ready()
+        ctx.on_bruteforce_done()
+        return out
+
+    return sample_fn
+
+
+# ── Core pipeline ─────────────────────────────────────────────────────────────
+
+def _run_pipeline(
+    N: int,
+    D: int,
+    N_ITER: int,
+    logger: ProgressLogger,
+    *,
+    x: np.ndarray | None = None,
+    S: float,
+    expension_factor: float,
+    LR_spatial: float,
+    LR_spectral: float,
+    spatial_radius: int,
+    spectral_radius: int,
+    N_PER_STEP: int,
+    _is_root: bool = False,
+    _is_leaf: bool = True,
+) -> np.ndarray:
+    """Recursive stealthy-sampling pipeline. Spawns child pipelines when N is large."""
+    ctx = logger.enter_level(N, D, N_ITER)
+
+    try:
+        Dsimp      = min(D, 3)
+        IJK, _, Axes = grid_shape(N, D)
+        Nsqrt      = N ** 0.5
+        Ncbrt      = N ** (1.0 / D)
+        is_root    = _is_root or (N <= 2_000) or (x is not None)
+        sigma2     = S * 2.0 * (1.0 / Ncbrt) ** 2
+        high_D     = sigma2 >= 0.03
+
+        SHIFTS        = integers_in_half_ball(spatial_radius, D)
+        Ks            = integers_in_half_ball(spectral_radius, D)
+        K_w, K_       = prepare_wave_vectors(Ks)
+        Clone_simplex = jnp.array(simplex(Dsimp))
+
+        if D == 4:
+            N_ITER *= 2
+        if D >= 5:
+            N_ITER *= 6
+
+        if high_D:
+            a = 2.0 * jnp.pi
+            b = 2.0 / (sigma2 * a ** 2)
+            c = 1.0 / (2.0 * S * jnp.pi)
+            micro_kernel = lambda x_val, y_val: gauss_sin_kernel(x_val, y_val, a, b, c)
+        else:
+            micro_kernel = lambda x_val, y_val: gauss_kernel(x_val, y_val, sigma2)
+
+        def micro_grad(x_val):
+            def body(acc, shift):
+                contrib = micro_kernel(x_val, jnp.roll(x_val, shift, axis=Axes))
+                return acc + contrib - jnp.roll(contrib, -shift, axis=Axes), None
+            out, _ = jax.lax.scan(body, jnp.zeros_like(x_val), SHIFTS)
+            return out
+
+        def macro_grad(x_val):
+            x_flat = x_val.reshape(-1, D + 1)
+            def body(acc, args):
+                k, k_ = args
+                return acc + spectral_kernel(x_flat, k, k_), None
+            out, _ = jax.lax.scan(body, jnp.zeros_like(x_flat), (K_w[:, 0], K_[:, 0]))
+            return out.reshape(*IJK, D + 1)
+
+        sn = SquareNet(gridshape=IJK, max_iter=50, verbose=0)
+
+        def _gridify_numpy(x_val: np.ndarray) -> np.ndarray:
+            ctx.tick()
+            flat = torus_wrap(np.random.permutation(x_val.reshape(-1, D + 1)) - 0.5)
+            sn.fit(flat[:, :D], method="ultimate")
+            return sn.map(flat)
+
+        def gridify(x_val: jnp.ndarray) -> jnp.ndarray:
+            return jax.pure_callback(
+                _gridify_numpy,
+                jax.ShapeDtypeStruct(x_val.shape, x_val.dtype),
+                x_val,
+            )
+
+        def clone(x_val: np.ndarray) -> np.ndarray:
+            """Expand N//(Dsimp+1) parents into N children via simplex offsets."""
+            x_val     = x_val.reshape(-1, D + 1)
+            x_val     = x_val[np.isfinite(x_val[:, -1]), :D]
+            x_val     = np.random.permutation(x_val)
+            N_parents = N // (Dsimp + 1)
+            N_keep    = N - (Dsimp + 1) * N_parents
+            offsets   = random_rotations(Clone_simplex, N_parents, D, Dsimp) * (expension_factor / Ncbrt)
+            children  = (x_val[:N_parents, None, :] + offsets).reshape(-1, D)
+            if N_keep > 0:
+                children = np.concatenate([x_val[N_parents:], children], axis=0)
+            return np.asarray(torus_wrap(jnp.array(children)))
+
+        @jax.jit
+        def run_iters(x_val: jnp.ndarray) -> jnp.ndarray:
+            def step(i, x_val):
+                x_val = jax.lax.cond(
+                    i % N_PER_STEP == 0,
+                    gridify,
+                    lambda val: val,
+                    x_val,
+                )
+                return clean_points(torus_wrap(
+                    x_val
+                    - (LR_spatial / S) * micro_grad(x_val)
+                    - (LR_spectral / (Nsqrt * Ncbrt)) * macro_grad(x_val)
+                ))
+            return jax.lax.fori_loop(0, N_ITER * N_PER_STEP, step, x_val)
+
+        if is_root:
+            xparent = _build_bruteforce(N, D, ctx)(x)
+            x_pts   = prepare_points(np.asarray(xparent), N, IJK, D)
+        else:
+            N_child = N // (Dsimp + 1) + N % (Dsimp + 1)
+            xparent = clone(
+                _run_pipeline(
+                    N=N_child,
+                    D=D,
+                    logger=logger,
+                    S=S,
+                    expension_factor=expension_factor,
+                    LR_spatial=LR_spatial,
+                    LR_spectral=LR_spectral,
+                    spatial_radius=spatial_radius,
+                    spectral_radius=spectral_radius,
+                    N_ITER=N_ITER,
+                    N_PER_STEP=N_PER_STEP,
+                    _is_root=False,
+                    _is_leaf=False,
+                )
+            )
+            x_pts = prepare_points(xparent, N, IJK, D)
+            x_pts = run_iters(x_pts)
+            ctx.done()
+
+        if _is_leaf:
+            x_pts = np.array(x_pts.reshape(-1, D + 1))
+            x_pts = x_pts[np.isfinite(x_pts[:, -1]), :D]
+
+        return x_pts
+
+    finally:
+        logger.exit_level()
+
+
+# ── Public entry point ────────────────────────────────────────────────────────
+
+def sample(
+    N: int,
+    D: int,
+    bruteforce: bool = False,
+    N_ITER: int  = 6,
+    verbose: int = 1,
+) -> np.ndarray:
+    """
+    Generate N stealthy points in [0, 1)^D.
+
+    Parameters
+    ----------
+    N        : number of output points.
+    D        : spatial dimension (2–4, or ≥5 falls back to bruteforce).
+    N_ITER   : pipeline iterations — more is better but slower.
+    verbose  : 0 = silent, 1 = live progress.
+    """
+    logger = ProgressLogger(verbose)
+
+    if bruteforce or N <= 3_000 or D >= 5:
+        reason = (
+            f"D={D} ≥ 5" if D >= 5
+            else f"N={N} ≤ 3 000" if N <= 3_000
+            else "bruteforce flag"
+        )
+        ctx  = logger.enter_level(N, D, 0, reason)
+        blue = _build_bruteforce(N, D, ctx)
+        out  = np.array(blue())
+        logger.exit_level()
+        return out
+
+    preset = _PRESETS[D]
+    return _run_pipeline(
+        N=N,
+        D=D,
+        N_ITER=N_ITER,
+        logger=logger,
+        x=None,
+        S=preset["S"],
+        expension_factor=preset["expension_factor"],
+        LR_spatial=preset["LR_spatial"],
+        LR_spectral=preset["LR_spectral"],
+        spatial_radius=preset["spatial_radius"],
+        spectral_radius=preset["spectral_radius"],
+        N_PER_STEP=10,
+        _is_root=False,
+        _is_leaf=True,
+    )

blue_sampler-0.1.0/src/blue_sampler/viz.py

@@ -0,0 +1,118 @@
+"""
+Visualisation helpers 
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+from .math_utils import structure_factor as _structure_factor
+
+
+def plot(
+    points: np.ndarray,
+    max_scatter: int = 30_000,
+    ax: plt.Axes | None = None,
+    **scatter_kw,
+) -> plt.Figure:
+    """
+    Scatter plot of a 2-D or 3-D point set.
+
+    For large point sets the view is automatically zoomed so that at most
+    *max_scatter* points are displayed.
+
+    Parameters
+    ----------
+    points : array-like
+        Point coordinates, shape (N, D) with D ∈ {2, 3}.
+        Higher-dimensional arrays are silently projected onto the first 3 axes.
+    max_scatter : int
+        Maximum number of points to draw.  Excess points are cropped by
+        zooming into the lower-left corner of the domain.
+    ax : matplotlib Axes | None
+        Existing axes to draw into.  When *None* a new figure is created.
+    **scatter_kw
+        Extra keyword arguments forwarded to ``ax.scatter``.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    """
+    pts = np.asarray(points).reshape(-1, np.asarray(points).shape[-1])
+    D   = min(pts.shape[-1], 3)
+    pts = pts[:, :D]
+
+    if len(pts) > max_scatter:
+        zoom = (max_scatter / len(pts)) ** (1.0 / D)
+        pts  = pts[(pts <= zoom).all(axis=1)]
+
+    kw = dict(s=0.4, color="black")
+    kw.update(scatter_kw)
+
+    if ax is None:
+        fig = plt.figure(figsize=(8, 8))
+        if D == 2:
+            ax = fig.add_subplot(111)
+        else:
+            ax = fig.add_subplot(111, projection="3d")
+    else:
+        fig = ax.get_figure()
+
+    if D == 2:
+        ax.scatter(pts[:, 0], pts[:, 1], **kw)
+    else:
+        ax.scatter(pts[:, 0], pts[:, 1], pts[:, 2], **kw)
+
+    ax.set_axis_off()
+    plt.tight_layout()
+    plt.show()
+    return fig
+
+
+def plot_structure_factor(
+    points: np.ndarray,
+    bins: int = 100,
+    resolution: float = 30.0,
+    ax: plt.Axes | None = None,
+    **plot_kw,
+) -> plt.Figure:
+    """
+    Log-log plot of the radial structure factor S(k).
+
+    Parameters
+    ----------
+    points : (N, D) array
+        Point coordinates in [0, 1)^D.
+    bins : int
+        Number of radial bins for the structure-factor estimate.
+    resolution : float
+        Random wave-vector density (vectors per bin) for the estimate.
+    ax : matplotlib Axes | None
+        Existing axes to draw into.  When *None* a new figure is created.
+    **plot_kw
+        Extra keyword arguments forwarded to ``ax.loglog``.
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+    """
+    pts = np.asarray(points).reshape(-1, np.asarray(points).shape[-1])
+    k, S = _structure_factor(pts, nbins=bins, resolution=resolution)
+
+    kw = dict(marker="o", markersize=2, linewidth=1)
+    kw.update(plot_kw)
+
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(7, 5))
+    else:
+        fig = ax.get_figure()
+
+    ax.loglog(k, S, **kw)
+    ax.set_xlabel("k")
+    ax.set_ylabel("S(k)")
+    ax.set_title("Structure factor  (log-log)")
+    ax.grid(True, which="both", alpha=0.4)
+    plt.tight_layout()
+    plt.show()
+    return fig