neural-implicit-queries
neural-implicit-queries copied to clipboard
`closest_point` can return a point further than `eps` from the surface
Hi, I've been digging through the kd tree implementation for closest_point
and noticed a mismatch between the grid cell size and the sample points used for intersection checking. It looks like the samples are taken from the faces of a box which is 2x the actual current cell size (see here). And also eps_cube_width
is only half the width of the largest cube fitting in the sphere of radius eps - which should be 2 * eps / sqrt(d)
.
One consequence of this is that it's possible that closest_point
returns a point that is further than eps
from the true surface.
Here's a minimal example with a 2D planar surface before (b5081cd) and after (409d4bc) fixing the mismatch (code for plots below):
Would be great to have your input on this! Luckily from my limited experiments it doesn't seem to happen very often in practice, maybe because the problematic grid cells are ruled out by range analysis in most cases. I didn't check other kd-tree based methods, but I guess same would apply there.
Code for plots (requires running on one of the commits above for the debugging info):
import jax.numpy as jnp
from jax import random as jrnd
from jax import vmap
from matplotlib import pyplot as plt
from functools import partial
import mlp, affine, kd_tree
def get_grid(min, max, n=100):
x, y = jnp.linspace(min[0], max[0], n), jnp.linspace(min[1], max[1], n)
x, y = jnp.meshgrid(x, y)
return jnp.stack([x, y], axis=-1)
def draw_box(ax, upper, lower, *args, **kwargs):
ax.plot(
[upper[0], upper[0], lower[0], lower[0], upper[0]],
[upper[1], lower[1], lower[1], upper[1], upper[1]],
*args, **kwargs,
)
eps = 0.1
class NarrowSurface(affine.AffineImplicitFunction):
def __call__(self, params, x):
return x[1] - eps # SDF of plane with normal [0, 1] going through [0, eps]
def classify_box(self, *args, **kwargs):
return kd_tree.SIGN_UNKNOWN # simulate case where range analysis can't classify any regions
f_affine = NarrowSurface(
mlp.func_from_spec(mode='affine'),
affine.AffineContext('affine_fixed')
)
lower, upper = jnp.array([-3., -3.]), jnp.array([3., 3.])
inputs = get_grid(lower, upper, n=1000)
fig, ax = plt.subplots(1, 1, figsize=(10, 8))
sdf_value = vmap(vmap(partial(f_affine, None)))(inputs)
query_points = jnp.array([-0.2, 2.])[None]
gt_closest_points = query_points.at[0, 1].set(eps)
dist, closest_points, iters = kd_tree.closest_point(
f_affine, None, lower, upper, query_points, eps=eps, return_iters=True
)
vlim = eps ** 2
img = ax.imshow(
sdf_value,
vmin=-vlim,
vmax=vlim,
cmap='seismic',
extent=[-3, 3, -3, 3],
origin='lower',
)
for i in range(len(query_points)):
ax.add_patch(plt.Circle(closest_points[i], eps, color='r', linestyle='--', linewidth=2, fill=False))
ax.plot([], [], 'r--', label=f'radius {eps=}')
ax.set_xlim(-1, 1)
ax.set_ylim(-2 * eps, 2.5)
i = -1
ax.add_patch(plt.Circle(query_points[0], iters[i].min_dist[0], color='w', linestyle='--', linewidth=2, fill=False))
ax.plot([], [], 'w--', label='final query_min_dist')
draw_box(ax, iters[i].min_upper[0], iters[i].min_lower[0], '--k', label='final kd-tree grid cell')
ax.plot(*closest_points.T, 'b*', label='estimated closest point')
ax.plot(*query_points.T, 'k*', label='query point')
ax.plot(*gt_closest_points.T, '*g', label='true closest point')
ax.legend()
plt.show()
fig.savefig('out.png', bbox_inches='tight', pad_inches=0)