Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 24 additions & 19 deletions src/underworld3/ckdtree.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -716,14 +716,26 @@ cdef class KDTree:
# must trip this guard, so the comparison is >= (issue #399).
if np.any(closest_n >= self.n):
raise RuntimeError(
"Error in rbf_interpolator_local_from_kdtree - a nearest neighbour wasn't found"
f"Cannot build a {nnn}-point stencil: the kd-tree holds "
f"{self.n} point(s) and a nearest neighbour wasn't found. "
"Reduce nnn, or check that the tree is not empty."
)

# np.bool_, not bool: this module cimports the C++ `bool` from libcpp,
# which shadows the Python builtin and will not compile here.
degenerate = np.zeros(coords_converted.shape[0], dtype=np.bool_)

if nnn == 1:
# The guard lives HERE, not in the callers: this early return
# precedes all `order` handling, so a caller-side check leaves the
# other entry point silently returning a nearest-neighbour stencil
# when a linear-exact one was asked for (issue #443).
if order == 1:
raise ValueError(
"order=1 needs at least dim + 2 neighbours to determine "
"the affine tail; nnn=1 selects the raw nearest-neighbour "
"path, which reproduces constants only."
)
return closest_n, np.ones(closest_n.shape), None, degenerate

# can decompose weighting vecotrs as IDW is a linear relationship
Expand Down Expand Up @@ -771,7 +783,10 @@ cdef class KDTree:
f"{degenerate.size} stencils could not support an affine fit "
"(collinear/coplanar neighbours) even after widening, and fell "
"back to inverse-distance weighting.",
stacklevel=3,
# 2 = whichever public method called _local_stencil. The two
# entry points sit at different depths below that, so this is
# the deepest frame that is correct for both.
stacklevel=2,
)

return closest_n, linear_weights, wide, degenerate
Expand Down Expand Up @@ -921,24 +936,14 @@ cdef class KDTree:
print(f"Mapping values with nnn - {nnn} & p {p} ... start", flush=True)

if nnn == 1:
# only use nearest neighbour raw data
if order == 1:
raise ValueError(
"order=1 needs at least dim + 2 neighbours to determine the "
f"affine tail; nnn=1 selects the raw nearest-neighbour path."
)
closest_n, _ = self.find_closest_n_points(
1, np.ascontiguousarray(coords_converted, dtype=np.float64)
# Only the raw nearest-neighbour value; _local_stencil holds the
# order=1 guard so both entry points share it (issue #443).
closest_n, _, _, _ = self._local_stencil(
coords_converted, nnn, p, order
)
# (n, 1) -> (n,): the nearest-neighbour path returns data rows
# directly, so the stencil axis must not survive into the result.
# query(k=1) used to do this reshape for us.
closest_n = closest_n.reshape(-1)
if np.any(closest_n >= self.n):
raise RuntimeError(
"Error in rbf_interpolator_local_from_kdtree - a nearest neighbour wasn't found"
)
return data[closest_n]
# (n, 1) -> (n,): this path returns data rows directly, so the
# stencil axis must not survive into the result.
return data[closest_n.reshape(-1)]

closest_n, n_weights, wide, degenerate = self._local_stencil(
coords_converted, nnn, p, order
Expand Down
14 changes: 4 additions & 10 deletions src/underworld3/discretisation/discretisation_mesh_variables.py
Original file line number Diff line number Diff line change
Expand Up @@ -886,8 +886,8 @@ def _get_kdtree(self):

return self._kdtree

def rbf_interpolate(self, new_coords, meth=0, p=2, verbose=False, nnn=None,
rubbish=None, order=0, monotone=False):
def rbf_interpolate(self, new_coords, nnn=None, p=2, verbose=False,
order=0, monotone=False):
"""Interpolate variable data to new coordinates using RBF.

Uses inverse distance weighting with k-nearest neighbors to
Expand All @@ -904,18 +904,12 @@ def rbf_interpolate(self, new_coords, meth=0, p=2, verbose=False, nnn=None,
----------
new_coords : numpy.ndarray
Target coordinates of shape ``(n_points, dim)``.
meth : int, optional
Interpolation method (reserved, currently unused).
TODO(BUG): issue #428 — ``meth`` and ``rubbish`` are dead
parameters, and ``tests/test_0505_rbf_swarm_mesh.py`` passes its
``nnn`` into ``meth`` positionally, so that test silently
discards it.
nnn : int, optional
Number of nearest neighbours (default: 4 for 3D, 3 for 2D).
p : float, optional
Power parameter for inverse distance weighting (default: 2).
verbose : bool, optional
Print progress information.
nnn : int, optional
Number of nearest neighbors (default: 4 for 3D, 3 for 2D).
order : int, optional
Polynomial reproduction order, 0 (default, bounded) or 1
(constants and linears exact; requires ``nnn >= dim + 2``).
Expand Down
76 changes: 0 additions & 76 deletions src/underworld3/swarm.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,6 @@ def __init__(
proxy_continuous=True,
_register=True,
_proxy=True,
_nn_proxy=False,
varsymbol=None,
rebuild_on_cycle=True,
units=None,
Expand Down Expand Up @@ -357,7 +356,6 @@ def __init__(
self._vtype = vtype
self._proxy_degree = proxy_degree
self._proxy_continuous = proxy_continuous
self._nn_proxy = _nn_proxy
self._create_proxy_variable()

# Inert: kept for backward compatibility with the removed
Expand Down Expand Up @@ -1205,53 +1203,6 @@ def _rbf_to_meshVar(self, meshVar, nnn=None, verbose=False, order=1,

return

def _rbf_reduce_to_meshVar(self, meshVar, verbose=False):
"""
This method updates a mesh variable for the current
swarm & particle variable state by reducing the swarm to
the nearest point for each particle

Here is how it works:

1) for each particle, create a distance-weighted average on the node data
2) check to see which nodes have zero weight / zero contribution and replace with nearest particle value

Todo: caching the k-d trees etc for the proxy-mesh-variable nodal points
Todo: some form of global fall-back for when there are no particles on a processor

"""

# if not proxied, nothing to do. return.
if not self._meshVar:
return

# 1 - Average particles to nodes with distance weighted average

# Use cached KDTree for interpolation (avoids redundant index construction)
kd = meshVar._get_kdtree()

d, n = kd.query(self.swarm.data, k=1, sqr_dists=False) # need actual distances

node_values = np.zeros((meshVar.coords.shape[0], self.num_components))
w = np.zeros(meshVar.coords.shape[0])

if not self._nn_proxy:
for i in range(self.local_size):
# if b[i]:
node_values[n[i], :] += self.data[i, :] / (1.0e-24 + d[i])
w[n[i]] += 1.0 / (1.0e-24 + d[i])

node_values[np.where(w > 0.0)[0], :] /= w[np.where(w > 0.0)[0]].reshape(-1, 1)

# 2 - set NN vals on mesh var where w == 0.0

p_nnmap = self.swarm._get_map(self)

meshVar.data[...] = node_values[...]
meshVar.data[np.where(w == 0.0), :] = self.data[p_nnmap[np.where(w == 0.0)], :]

return

# # Need to be able to unpack as well
# def pack_raw_data_to_petsc(self, data_array):
# """Convert an array in the correct shape for the underlying variable into something that can be loaded into
Expand Down Expand Up @@ -2312,7 +2263,6 @@ def __init__(
proxy_continuous = (True,)
_register = (True,)
_proxy = (True,)
_nn_proxy = (False,)
varsymbol = (None,)
rebuild_on_cycle = (True,)
"""
Expand Down Expand Up @@ -2883,7 +2833,6 @@ def __init__(self, mesh, recycle_rate=0, verbose=False, clip_to_mesh=True):

self._X0_uninitialised = True
self._index = None
self._nnmapdict = {}
self._migration_disabled = False

# Deterministic (SPMD-consistent) creation index — used to order
Expand Down Expand Up @@ -4256,7 +4205,6 @@ def add_variable(
size=1,
dtype=float,
proxy_degree=2,
_nn_proxy=False,
units=None,
):
"""
Expand All @@ -4276,8 +4224,6 @@ def add_variable(
Data type (float or int)
proxy_degree : int, default 2
Degree for mesh proxy variable interpolation
_nn_proxy : bool, default False
Internal parameter for nearest-neighbor proxy
units : str, optional
Physical units for this variable (e.g., "kg/m^3", "m/s")

Expand Down Expand Up @@ -4322,7 +4268,6 @@ def add_variable(
size,
dtype=dtype,
proxy_degree=proxy_degree,
_nn_proxy=_nn_proxy,
units=units,
)

Expand Down Expand Up @@ -4750,27 +4695,6 @@ def _data_layout(self, i, j=None):
if self.vtype == uw.VarType.MATRIX:
return i + j * self.shape[0]

## Check this - the interface to kdtree has changed, are we picking the correct field ?
@timing.routine_timer_decorator
def _get_map(self, var):
# generate tree if not avaiable
kd = self._get_kdtree()

# get or generate map
meshvar_coords = var._meshVar.coords
# we can't use numpy arrays directly as keys in python dicts, so
# we'll use `xxhash` to generate a hash of array.
# this shouldn't be an issue performance wise but we should test to be
# sufficiently confident of this.
import xxhash

h = xxhash.xxh64()
h.update(meshvar_coords)
digest = h.intdigest()
if digest not in self._nnmapdict:
self._nnmapdict[digest] = kd.query(meshvar_coords, k=1, sqr_dists=False)[1]
return self._nnmapdict[digest]

@timing.routine_timer_decorator
def advection(
self,
Expand Down
37 changes: 37 additions & 0 deletions tests/test_0102_kdtree_linear_rbf.py
Original file line number Diff line number Diff line change
Expand Up @@ -439,6 +439,43 @@ def test_interpolation_matrix_agrees_with_the_value_path(dim, order):
assert np.abs(T @ values - direct).max() < 1.0e-12


@pytest.mark.parametrize("dim", [2, 3])
def test_both_entry_points_reject_nnn_1_with_order_1(dim):
"""The guard must live below both APIs, not in one of them (issue #443).

`_local_stencil` early-returns for nnn == 1 before any `order` handling,
so a guard placed in the value path alone left `interpolation_matrix`
silently returning a nearest-neighbour operator — constants-only — when a
linear-exact one was asked for, with nothing recording the downgrade.
"""
rng = np.random.default_rng(300 + dim)
source = rng.random((200, dim))
target = rng.random((15, dim))
data = _linear(source)[:, None]
kdt = uw.kdtree.KDTree(source)

with pytest.raises(ValueError, match="dim . 2 neighbours"):
kdt.rbf_interpolator_local(target, data, 1, 2, False, order=1)

with pytest.raises(ValueError, match="dim . 2 neighbours"):
kdt.interpolation_matrix(target, nnn=1, order=1)

# nnn=1 at order=0 stays legal on both.
assert kdt.rbf_interpolator_local(target, data, 1, 2, False).shape == (15, 1)
assert kdt.interpolation_matrix(target, nnn=1).shape == (15, 200)


def test_stencil_larger_than_the_cloud_reports_what_went_wrong():
"""The old message named a function the caller never invoked."""
rng = np.random.default_rng(17)
source = rng.random((5, 3))
target = rng.random((4, 3))
kdt = uw.kdtree.KDTree(source)

with pytest.raises(RuntimeError, match="20-point stencil.*holds 5 point"):
kdt.interpolation_matrix(target, nnn=20, order=1)


@pytest.mark.parametrize("dim", [2, 3])
def test_interpolation_matrix_rows_are_never_empty(dim):
"""The raw-weights helper zeroes degenerate rows; the operator must not.
Expand Down
Loading