feat(sdf): S12-C PyTorch integration + Python bindings
Add batch evaluation/gradient functions for SDF trees (sdf_torch.h/cpp). Add sdf_gradient.h with finite-difference gradient utilities. Add Python modules vde.sdf (high-level SDF API) and vde.torch_sdf (PyTorch autograd). Add test_sdf_torch_bridge.cpp with batch eval/gradient tests. Integrate into vde_sdf library and test suite.
This commit is contained in:
茂之钳
@@ -0,0 +1,132 @@
|
||||
"""
|
||||
vde.sdf — Implicit Modeling & Differentiable Geometry
|
||||
|
||||
Usage:
|
||||
from vde.sdf import Sphere, Box, Union, evaluate, to_mesh
|
||||
|
||||
s = Sphere(1.0)
|
||||
b = Box([0.5, 0.5, 0.5])
|
||||
u = Union(s, b)
|
||||
|
||||
d = evaluate(u, [0.5, 0, 0])
|
||||
mesh = to_mesh(u, resolution=64)
|
||||
"""
|
||||
import numpy as np
|
||||
|
||||
# Re-export C++ bindings
|
||||
from vde._vde import (
|
||||
SdfNode, SdfOp, SdfParams, MCMesh, SdfBBox,
|
||||
evaluate, sdf_to_mesh, sdf_to_mesh_lambda,
|
||||
estimate_bounds, estimate_bbox,
|
||||
)
|
||||
|
||||
# ── Convenience constructors ──────────────────────
|
||||
|
||||
def sphere(radius=1.0):
|
||||
"""Create a sphere SDF node."""
|
||||
return SdfNode.sphere(radius)
|
||||
|
||||
def box(half_extents):
|
||||
"""Create a box SDF node.
|
||||
Args:
|
||||
half_extents: [hx, hy, hz] half-size in each dimension
|
||||
"""
|
||||
return SdfNode.box(half_extents)
|
||||
|
||||
def cylinder(radius=1.0, height=2.0):
|
||||
"""Create a cylinder SDF node."""
|
||||
return SdfNode.cylinder(radius, height)
|
||||
|
||||
def torus(major_radius=1.0, minor_radius=0.3):
|
||||
"""Create a torus SDF node."""
|
||||
return SdfNode.torus(major_radius, minor_radius)
|
||||
|
||||
def capsule(a, b, radius=0.5):
|
||||
"""Create a capsule SDF node between points a and b."""
|
||||
return SdfNode.capsule(a, b, radius)
|
||||
|
||||
def cone(angle=0.5, height=2.0):
|
||||
"""Create a cone SDF node."""
|
||||
return SdfNode.cone(angle, height)
|
||||
|
||||
def plane(normal, offset=0.0):
|
||||
"""Create a plane SDF node."""
|
||||
return SdfNode.plane(normal, offset)
|
||||
|
||||
def ellipsoid(radii):
|
||||
"""Create an ellipsoid SDF node."""
|
||||
return SdfNode.ellipsoid(radii)
|
||||
|
||||
# ── CSG Operations ────────────────────────────────
|
||||
|
||||
def union(a, b):
|
||||
"""Union of two SDF nodes (min)."""
|
||||
return SdfNode.op_union(a, b)
|
||||
|
||||
def intersection(a, b):
|
||||
"""Intersection of two SDF nodes (max)."""
|
||||
return SdfNode.op_intersection(a, b)
|
||||
|
||||
def difference(a, b):
|
||||
"""Difference: subtract b from a."""
|
||||
return SdfNode.op_difference(a, b)
|
||||
|
||||
def smooth_union(a, b, k=0.1):
|
||||
"""Smooth union with blend radius k."""
|
||||
return SdfNode.smooth_union(a, b, k)
|
||||
|
||||
def smooth_intersection(a, b, k=0.1):
|
||||
"""Smooth intersection with blend radius k."""
|
||||
return SdfNode.smooth_intersection(a, b, k)
|
||||
|
||||
def smooth_difference(a, b, k=0.1):
|
||||
"""Smooth difference with blend radius k."""
|
||||
return SdfNode.smooth_difference(a, b, k)
|
||||
|
||||
# ── Modifiers ─────────────────────────────────────
|
||||
|
||||
def round_shape(node, r=0.1):
|
||||
"""Round (fatten) a shape."""
|
||||
return SdfNode.round(node, r)
|
||||
|
||||
def onion(node, thickness=0.1):
|
||||
"""Create a thin shell."""
|
||||
return SdfNode.onion(node, thickness)
|
||||
|
||||
# ── Domain Transforms ─────────────────────────────
|
||||
|
||||
def translate(node, offset):
|
||||
"""Translate a shape."""
|
||||
return SdfNode.translate(node, offset)
|
||||
|
||||
def rotate(node, angle):
|
||||
"""Rotate around Y axis (radians)."""
|
||||
return SdfNode.rotate(node, angle)
|
||||
|
||||
def scale(node, factors):
|
||||
"""Non-uniform scale."""
|
||||
return SdfNode.scale(node, factors)
|
||||
|
||||
def repeat(node, cell):
|
||||
"""Infinite repeat."""
|
||||
return SdfNode.repeat(node, cell)
|
||||
|
||||
def twist(node, amount):
|
||||
"""Twist around Y axis."""
|
||||
return SdfNode.twist(node, amount)
|
||||
|
||||
def bend(node, amount):
|
||||
"""Bend deformation."""
|
||||
return SdfNode.bend(node, amount)
|
||||
|
||||
# ── Mesh Conversion ───────────────────────────────
|
||||
|
||||
def to_mesh(node, resolution=64, level=0.0):
|
||||
"""Convert SDF tree to triangle mesh.
|
||||
Returns: MCMesh with .vertices and .triangles (both numpy arrays)
|
||||
"""
|
||||
return sdf_to_mesh(node, resolution, level)
|
||||
|
||||
def to_mesh_lambda(fn, bmin, bmax, resolution=64, level=0.0):
|
||||
"""Convert lambda SDF to mesh."""
|
||||
return sdf_to_mesh_lambda(fn, bmin, bmax, resolution, level)
|
||||
@@ -0,0 +1,185 @@
|
||||
"""
|
||||
PyTorch integration for differentiable SDF operations.
|
||||
|
||||
Provides torch.autograd.Function wrappers that allow SDF
|
||||
evaluation within PyTorch computation graphs.
|
||||
|
||||
Usage:
|
||||
import torch
|
||||
from vde.torch_sdf import SdfEvaluator
|
||||
|
||||
evaluator = SdfEvaluator(node)
|
||||
x = torch.randn(100, 3, requires_grad=True)
|
||||
d = evaluator(x) # SDF values
|
||||
loss = d.abs().mean()
|
||||
loss.backward() # x.grad contains df/dx
|
||||
"""
|
||||
import torch
|
||||
|
||||
class SdfEvaluator(torch.autograd.Function):
|
||||
"""Differentiable SDF evaluation for PyTorch.
|
||||
|
||||
Forward: evaluate SDF at input points.
|
||||
Backward: compute gradient of SDF w.r.t. input points.
|
||||
|
||||
Usage:
|
||||
evaluator = SdfEvaluator.apply
|
||||
d = evaluator(node, points) # where node is an SdfNode
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, node, points):
|
||||
"""Forward pass: evaluate SDF.
|
||||
|
||||
Args:
|
||||
node: SdfNode (from C++ binding)
|
||||
points: torch.Tensor of shape (N, 3)
|
||||
Returns:
|
||||
distances: torch.Tensor of shape (N,)
|
||||
"""
|
||||
import numpy as np
|
||||
from vde._vde import evaluate
|
||||
|
||||
# Convert to numpy for C++ evaluation
|
||||
pts_np = points.detach().cpu().numpy().astype(np.float64)
|
||||
n = pts_np.shape[0]
|
||||
|
||||
distances = np.zeros(n, dtype=np.float64)
|
||||
for i in range(n):
|
||||
p = pts_np[i]
|
||||
distances[i] = evaluate(node, p)
|
||||
|
||||
ctx.save_for_backward(points)
|
||||
ctx.node = node # Store node reference for backward
|
||||
|
||||
return torch.from_numpy(distances).to(points.device).to(points.dtype)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
"""Backward pass: gradient of SDF w.r.t. input points.
|
||||
|
||||
df/dx = gradient of SDF at each input point
|
||||
Chain rule: dL/dx = dL/d(sdf) * d(sdf)/dx
|
||||
"""
|
||||
import numpy as np
|
||||
from vde._vde import evaluate
|
||||
|
||||
points = ctx.saved_tensors[0]
|
||||
node = ctx.node
|
||||
|
||||
pts_np = points.detach().cpu().numpy().astype(np.float64)
|
||||
n = pts_np.shape[0]
|
||||
|
||||
gradients = np.zeros((n, 3), dtype=np.float64)
|
||||
h = 1e-6
|
||||
|
||||
for i in range(n):
|
||||
p = pts_np[i]
|
||||
# Central finite differences
|
||||
for j in range(3):
|
||||
pp = p.copy()
|
||||
pm = p.copy()
|
||||
pp[j] += h
|
||||
pm[j] -= h
|
||||
gradients[i, j] = (evaluate(node, pp) - evaluate(node, pm)) / (2.0 * h)
|
||||
|
||||
grad_input = torch.from_numpy(gradients).to(points.device).to(points.dtype)
|
||||
grad_node = None # Don't backprop to SDF parameters (for now)
|
||||
|
||||
return grad_node, grad_input * grad_output.unsqueeze(-1)
|
||||
|
||||
|
||||
class SdfShapeFitter:
|
||||
"""Fit SDF shapes to target point clouds with gradient descent.
|
||||
|
||||
Usage:
|
||||
fitter = SdfShapeFitter(sphere(1.0))
|
||||
optimized = fitter.fit(target_points, lr=0.01, epochs=100)
|
||||
"""
|
||||
|
||||
def __init__(self, node):
|
||||
self.node = node
|
||||
|
||||
def fit(self, target_points, lr=0.01, epochs=100, verbose=False):
|
||||
"""Fit SDF shape to target point cloud.
|
||||
|
||||
Args:
|
||||
target_points: torch.Tensor of shape (N, 3)
|
||||
lr: learning rate
|
||||
epochs: number of optimization steps
|
||||
verbose: print progress
|
||||
|
||||
Returns:
|
||||
optimized node (SdfNode)
|
||||
"""
|
||||
import numpy as np
|
||||
from vde._vde import evaluate
|
||||
|
||||
pts = target_points.detach().cpu().numpy().astype(np.float64)
|
||||
n = pts.shape[0]
|
||||
h = 1e-6
|
||||
|
||||
# Get current parameters
|
||||
params = self.node.params
|
||||
param_list = self._params_to_list(params)
|
||||
|
||||
for epoch in range(epochs):
|
||||
# Forward: evaluate SDF at all points
|
||||
loss = 0.0
|
||||
grads = {k: 0.0 for k in self._param_names()}
|
||||
|
||||
for i in range(n):
|
||||
p = pts[i]
|
||||
d = evaluate(self.node, p)
|
||||
loss += abs(d)
|
||||
|
||||
# Compute numerical gradient w.r.t. parameters
|
||||
for name, idx in self._param_indices().items():
|
||||
orig_val = param_list[idx]
|
||||
param_list[idx] = orig_val + h
|
||||
self._list_to_params(param_list, self.node.params)
|
||||
dp = evaluate(self.node, p)
|
||||
|
||||
param_list[idx] = orig_val - h
|
||||
self._list_to_params(param_list, self.node.params)
|
||||
dm = evaluate(self.node, p)
|
||||
|
||||
param_list[idx] = orig_val
|
||||
grads[name] += (abs(dp) - abs(dm)) / (2.0 * h)
|
||||
|
||||
self._list_to_params(param_list, self.node.params)
|
||||
loss /= n
|
||||
for name in grads:
|
||||
grads[name] /= n
|
||||
|
||||
# Gradient descent step
|
||||
for name, idx in self._param_indices().items():
|
||||
param_list[idx] -= lr * grads[name]
|
||||
|
||||
self._list_to_params(param_list, self.node.params)
|
||||
|
||||
if verbose and epoch % 10 == 0:
|
||||
print(f"Epoch {epoch}: loss = {loss:.6f}")
|
||||
|
||||
return self.node
|
||||
|
||||
def _param_names(self):
|
||||
return ['radius', 'extent_x', 'extent_y', 'extent_z', 'height']
|
||||
|
||||
def _param_indices(self):
|
||||
return {
|
||||
'radius': 0,
|
||||
'extent_x': 1, 'extent_y': 2, 'extent_z': 3,
|
||||
'height': 4,
|
||||
}
|
||||
|
||||
def _params_to_list(self, params):
|
||||
return [params.radius, params.extents.x, params.extents.y,
|
||||
params.extents.z, params.height]
|
||||
|
||||
def _list_to_params(self, lst, params):
|
||||
params.radius = lst[0]
|
||||
params.extents.x = lst[1]
|
||||
params.extents.y = lst[2]
|
||||
params.extents.z = lst[3]
|
||||
params.height = lst[4]
|
||||
Reference in New Issue
Block a user