lipschitz_octree and unique_sparse_voxel_corners

[alecjacobson]

and bumps on a few template issues along the way

[alecjacobson]

igl.lipschitz_octree is the first time I'm trying to expose taking a python function nb::callable as input. In this case it's to call the implicit function on a batch of points (all corners of an octree's current depth). Since the python version is quite a bit slower than the C++ version, I'm worried the overhead is high. It might also just be that the C++ version is parallelized and the simple numpy function I was testing was not.

[alecjacobson]

I have no idea if it will work with a function operating on a pytorch array. That'd be cool.

[alecjacobson]

Cool. It does kind of seem like it works with pytorch. I'm not sure the best practice, but since the udf function will get called with a double type Eigen::Matrix, what worked for me was to make sure my pytorch-based udf function on the python side checks the input and possible converts it to a pytorch tensor.

Sure seems like a lot of copying, but at least the functionality seems to be there.

Unsurprisingly, when the dense grid fits into device memory and the implicit function isn't too complicated, then this O(n²) method doesn't necessarily beat the naive dense O(n³) evaluation. This is especially the case if the device is mps or gpu.

The O(n³) also applies to memory so it's pretty abrupt when the dense grid falls out of memory and becomes intensely slow.

I'm probably not going to optimize this anymore any time soon, but I leave this test example here for posterity:

import torch
import igl
import time
import polyscope as ps
import numpy as np

## On Alec's M4 laptop mps produces 
device = torch.device("mps")
#device = torch.device("cpu")
if device.type == "mps":
    device_dtype = torch.float32
else:
    device_dtype = torch.float64

def sdSphere(Q,s):
    return torch.linalg.norm(Q,axis=1) - s

def sdBox(Q,b,unused):
    X = Q[:,0]
    Y = Q[:,1]
    Z = Q[:,2]
    b = torch.as_tensor(b, dtype=Q.dtype, device=Q.device)
    q = torch.abs(Q) - b
    zero = torch.zeros(1, dtype=q.dtype, device=q.device)
    return torch.linalg.norm(torch.maximum(q, zero), axis=1) + \
       torch.minimum(torch.maximum(q[:,0], torch.maximum(q[:,1], q[:,2])), zero)

def opUnion(d1,d2,unused):
    return torch.minimum(d1,d2)

def mix(a, b, t):
    return a + t * (b - a)

def opSmoothUnion(d1,d2,k):
    h = torch.clip(0.5 + 0.5*(d2-d1)/k, 0.0, 1.0)
    return mix(d2, d1, h) - k * h * (1.0 - h)

def sdf(Q):
    # ensure that Q is a torch tensor
    if not isinstance(Q, torch.Tensor):
        Q = torch.tensor(Q, dtype=device_dtype, device=device)
    sphere_translation = np.array([0.0, 0.0, 0.2], dtype=np.float64)
    sphere_translation = torch.tensor(sphere_translation, dtype=Q.dtype, device=Q.device)
    box_translation = np.array([0.0,0.0,-0.4]);
    box_translation = torch.tensor(box_translation, dtype=Q.dtype, device=Q.device)
    return opSmoothUnion(
        sdSphere(Q-sphere_translation,0.4),
        sdBox(
                Q-box_translation,
                np.array([7+5/8,3+5/8,2+1/4],dtype=np.float64)/8,
                0.1),
        0.22)

def udf(Q):
    # ensure that Q is a torch tensor
    if not isinstance(Q, torch.Tensor):
        Q = torch.tensor(Q, dtype=device_dtype, device=device)
    U = torch.abs(sdf(Q))
    return U.cpu().numpy().astype(np.float64)


origin = np.array([-1,-1,-1],dtype=np.float64)
h0 = 2.0
max_depth = 9

# time before
start_time = time.time()

ijk = igl.lipschitz_octree(origin,h0,max_depth,udf)
print("lipschitz_octree:", time.time() - start_time)

start_time = time.time()
h = h0 / (2**max_depth)
unique_ijk, J, unique_corners = igl.unique_sparse_voxel_corners(origin,h0,max_depth,ijk)
print("unique_sparse_voxel_corners:", time.time() - start_time)

start_time = time.time()
unique_S = sdf(unique_corners)
unique_S = unique_S.cpu().numpy().astype(np.float64)
print("sdf:", time.time() - start_time)

start_time = time.time()
V,F = igl.marching_cubes(unique_S,unique_corners,J,0.0)
print("marching_cubes:", time.time() - start_time)

res = np.array([2**max_depth+1, 2**max_depth+1, 2**max_depth+1])
start_time = time.time()
GV = igl.grid(res)
print("grid:", time.time() - start_time)
GV = GV*h0 + origin

# move to mps device
GV = torch.tensor(GV, dtype=torch.float32, device=device)
start_time = time.time()
GS = sdf(GV)
print("sdf grid:", time.time() - start_time)
start_time = time.time()
GV = GV.to(torch.float32).cpu().numpy().astype(np.float64)
GS = GS.to(torch.float32).cpu().numpy().astype(np.float64)
gV,gF,E2V = igl.marching_cubes(GS,GV,res[0],res[1],res[2])
print("marching_cubes grid:", time.time() - start_time)


ps.init()
ps.register_point_cloud("unique_corners", unique_corners, radius=h/10.0)
ps.register_surface_mesh("mc", V, F)
ps.register_surface_mesh("mc_grid", gV, gF)
ps.set_give_focus_on_show(True)
# z axis up
ps.set_up_dir("z_up")
ps.show()