jack
62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
"""RBF (thin-plate spline) heatmap field.
|
|
|
|
Uses CuPy for GPU acceleration when available, falls back to NumPy + SciPy.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
from scipy.interpolate import RBFInterpolator
|
|
|
|
try:
|
|
import cupy as _cupy # type: ignore
|
|
BACKEND = "cupy"
|
|
except Exception:
|
|
_cupy = None
|
|
BACKEND = "numpy"
|
|
|
|
_KERNEL = "thin_plate_spline"
|
|
_SMOOTHING = 0.0
|
|
|
|
|
|
def interpolate_field(
|
|
xs: np.ndarray,
|
|
ys: np.ndarray,
|
|
vs: np.ndarray,
|
|
*,
|
|
width: int,
|
|
height: int,
|
|
extent: tuple[float, float, float, float], # (xmin, xmax, ymin, ymax) in mm
|
|
round_clip: bool = False,
|
|
) -> np.ndarray:
|
|
"""Return a (height, width) float64 array of interpolated values.
|
|
|
|
Args:
|
|
xs, ys: sensor positions in mm (1-D arrays, length N)
|
|
vs: sensor values (length N)
|
|
width/height: output grid dimensions in pixels
|
|
extent: (xmin, xmax, ymin, ymax) in the same mm space as xs/ys
|
|
round_clip: if True, pixels outside the inscribed ellipse become NaN
|
|
"""
|
|
coords = np.column_stack([xs, ys])
|
|
rbf = RBFInterpolator(coords, vs, kernel=_KERNEL, smoothing=_SMOOTHING)
|
|
|
|
xmin, xmax, ymin, ymax = extent
|
|
gx = np.linspace(xmin, xmax, width)
|
|
gy = np.linspace(ymin, ymax, height)
|
|
grid_x, grid_y = np.meshgrid(gx, gy)
|
|
flat = np.column_stack([grid_x.ravel(), grid_y.ravel()])
|
|
|
|
# RBFInterpolator always runs on CPU; CuPy only accelerates other ops if added later
|
|
field = rbf(flat).reshape(height, width)
|
|
|
|
if round_clip:
|
|
cx = (xmin + xmax) / 2
|
|
cy = (ymin + ymax) / 2
|
|
rx = (xmax - xmin) / 2
|
|
ry = (ymax - ymin) / 2
|
|
dist = ((grid_x - cx) / rx) ** 2 + ((grid_y - cy) / ry) ** 2
|
|
field = np.where(dist <= 1.0, field, np.nan)
|
|
|
|
return field.astype(np.float64)
|
|
|