neural-implicit-queries icon indicating copy to clipboard operation
neural-implicit-queries copied to clipboard

`closest_point` can return a point further than `eps` from the surface

Open Linusnie opened this issue 1 year ago • 0 comments

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): before after

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)

Linusnie avatar Jul 31 '23 16:07 Linusnie