diff --git a/include/vde/sdf/sdf_gradient.h b/include/vde/sdf/sdf_gradient.h new file mode 100644 index 0000000..8faa3b7 --- /dev/null +++ b/include/vde/sdf/sdf_gradient.h @@ -0,0 +1,38 @@ +#pragma once +#include "vde/core/point.h" +#include +#include + +namespace vde::sdf { + +using core::Point3D; +using core::Vector3D; + +/// Compute numerical gradient of a scalar field f at point p +/// Uses central finite differences with step size h +template +[[nodiscard]] Vector3D gradient(Fn&& f, const Point3D& p, double h = 1e-6) { + double fx = f(p); + const double eps_x = h * (1.0 + std::abs(p.x())); + const double eps_y = h * (1.0 + std::abs(p.y())); + const double eps_z = h * (1.0 + std::abs(p.z())); + + double dfdx = (f(Point3D(p.x() + eps_x, p.y(), p.z())) - + f(Point3D(p.x() - eps_x, p.y(), p.z()))) / (2.0 * eps_x); + double dfdy = (f(Point3D(p.x(), p.y() + eps_y, p.z())) - + f(Point3D(p.x(), p.y() - eps_y, p.z()))) / (2.0 * eps_y); + double dfdz = (f(Point3D(p.x(), p.y(), p.z() + eps_z)) - + f(Point3D(p.x(), p.y(), p.z() - eps_z))) / (2.0 * eps_z); + + return Vector3D(dfdx, dfdy, dfdz); +} + +/// Evaluate scalar field and its gradient at point p (joint pass) +/// Returns (value, gradient) +template +[[nodiscard]] std::tuple evaluate_with_gradient( + Fn&& f, const Point3D& p, double h = 1e-6) { + return {f(p), gradient(std::forward(f), p, h)}; +} + +} // namespace vde::sdf diff --git a/include/vde/sdf/sdf_torch.h b/include/vde/sdf/sdf_torch.h new file mode 100644 index 0000000..e596712 --- /dev/null +++ b/include/vde/sdf/sdf_torch.h @@ -0,0 +1,35 @@ +#pragma once +#include "vde/sdf/sdf_gradient.h" +#include "vde/sdf/sdf_tree.h" +#include "vde/core/point.h" +#include + +// This header is optional — only compiled if Torch is available + +namespace vde::sdf { + +/// Evaluate SDF on a batch of points (for batched computation) +/// @param root SDF tree +/// @param points Flat array of x,y,z triplets (size = 3*n) +/// @param n Number of points +/// @param distances Output array (size = n) +void evaluate_batch(const SdfNodePtr& root, + const double* points, int n, + double* distances); + +/// Compute gradients for a batch of points +/// @param root SDF tree +/// @param points Flat array of x,y,z triplets +/// @param n Number of points +/// @param gradients Output array (size = 3*n, interleaved x,y,z) +void gradient_batch(const SdfNodePtr& root, + const double* points, int n, + double* gradients); + +/// Compute both values and gradients in one pass +void evaluate_with_gradient_batch(const SdfNodePtr& root, + const double* points, int n, + double* distances, + double* gradients); + +} // namespace vde::sdf diff --git a/python/vde/sdf.py b/python/vde/sdf.py new file mode 100644 index 0000000..538546d --- /dev/null +++ b/python/vde/sdf.py @@ -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) diff --git a/python/vde/torch_sdf.py b/python/vde/torch_sdf.py new file mode 100644 index 0000000..83bd5e1 --- /dev/null +++ b/python/vde/torch_sdf.py @@ -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] diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 0a8d63e..3083ac6 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -171,6 +171,7 @@ add_library(vde_sdf STATIC sdf/sdf_primitives.cpp sdf/sdf_tree.cpp sdf/sdf_to_mesh.cpp + sdf/sdf_torch.cpp ) target_include_directories(vde_sdf PUBLIC ${CMAKE_SOURCE_DIR}/include diff --git a/src/sdf/sdf_torch.cpp b/src/sdf/sdf_torch.cpp new file mode 100644 index 0000000..79385a3 --- /dev/null +++ b/src/sdf/sdf_torch.cpp @@ -0,0 +1,44 @@ +#include "vde/sdf/sdf_torch.h" +#include "vde/sdf/sdf_gradient.h" + +namespace vde::sdf { + +void evaluate_batch(const SdfNodePtr& root, + const double* points, int n, + double* distances) { + auto fn = [&](const Point3D& p) -> double { return evaluate(root, p); }; + for (int i = 0; i < n; ++i) { + Point3D p(points[i*3], points[i*3+1], points[i*3+2]); + distances[i] = fn(p); + } +} + +void gradient_batch(const SdfNodePtr& root, + const double* points, int n, + double* gradients) { + auto fn = [&](const Point3D& p) -> double { return evaluate(root, p); }; + for (int i = 0; i < n; ++i) { + Point3D p(points[i*3], points[i*3+1], points[i*3+2]); + Vector3D g = gradient(fn, p); + gradients[i*3] = g.x(); + gradients[i*3+1] = g.y(); + gradients[i*3+2] = g.z(); + } +} + +void evaluate_with_gradient_batch(const SdfNodePtr& root, + const double* points, int n, + double* distances, + double* gradients) { + auto fn = [&](const Point3D& p) -> double { return evaluate(root, p); }; + for (int i = 0; i < n; ++i) { + Point3D p(points[i*3], points[i*3+1], points[i*3+2]); + auto [val, grad] = evaluate_with_gradient(fn, p); + distances[i] = val; + gradients[i*3] = grad.x(); + gradients[i*3+1] = grad.y(); + gradients[i*3+2] = grad.z(); + } +} + +} // namespace vde::sdf diff --git a/tests/sdf/CMakeLists.txt b/tests/sdf/CMakeLists.txt index a53f637..69fb88d 100644 --- a/tests/sdf/CMakeLists.txt +++ b/tests/sdf/CMakeLists.txt @@ -2,3 +2,4 @@ add_vde_test(test_sdf_primitives) add_vde_test(test_sdf_operations) add_vde_test(test_sdf_tree) add_vde_test(test_sdf_to_mesh) +add_vde_test(test_sdf_torch_bridge) diff --git a/tests/sdf/test_sdf_torch_bridge.cpp b/tests/sdf/test_sdf_torch_bridge.cpp new file mode 100644 index 0000000..0ef344a --- /dev/null +++ b/tests/sdf/test_sdf_torch_bridge.cpp @@ -0,0 +1,168 @@ +#include +#include "vde/sdf/sdf_torch.h" +#include "vde/sdf/sdf_gradient.h" +#include "vde/sdf/sdf_tree.h" +#include "vde/sdf/sdf_primitives.h" +#include "vde/core/point.h" +#include +#include +#include + +using namespace vde::sdf; +using vde::core::Point3D; + +constexpr double EPS = 1e-5; + +// ═══════════════════════════════════════════════════ +// Helper: generate random points +// ═══════════════════════════════════════════════════ + +std::vector random_points(int n) { + std::mt19937 rng(42); + std::uniform_real_distribution dist(-3.0, 3.0); + std::vector pts(3 * n); + for (int i = 0; i < 3 * n; ++i) { + pts[i] = dist(rng); + } + return pts; +} + +// ═══════════════════════════════════════════════════ +// evaluate_batch +// ═══════════════════════════════════════════════════ + +TEST(SdfTorchBridge, EvaluateBatch_MatchesSequential) { + auto root = SdfNode::sphere(2.0); + const int n = 100; + auto pts_arr = random_points(n); + + std::vector batch_dist(n); + evaluate_batch(root, pts_arr.data(), n, batch_dist.data()); + + for (int i = 0; i < n; ++i) { + Point3D p(pts_arr[i*3], pts_arr[i*3+1], pts_arr[i*3+2]); + double expected = evaluate(root, p); + EXPECT_NEAR(batch_dist[i], expected, EPS); + } +} + +TEST(SdfTorchBridge, EvaluateBatch_CSGTree) { + auto a = SdfNode::sphere(1.5); + auto b = SdfNode::box(Point3D(0.8, 0.8, 0.8)); + auto root = SdfNode::op_union(a, b); + const int n = 50; + auto pts_arr = random_points(n); + + std::vector batch_dist(n); + evaluate_batch(root, pts_arr.data(), n, batch_dist.data()); + + for (int i = 0; i < n; ++i) { + Point3D p(pts_arr[i*3], pts_arr[i*3+1], pts_arr[i*3+2]); + double expected = evaluate(root, p); + EXPECT_NEAR(batch_dist[i], expected, EPS); + } +} + +// ═══════════════════════════════════════════════════ +// gradient_batch +// ═══════════════════════════════════════════════════ + +TEST(SdfTorchBridge, GradientBatch_MatchesSequential) { + auto root = SdfNode::sphere(2.0); + const int n = 100; + auto pts_arr = random_points(n); + + std::vector batch_grad(3 * n); + gradient_batch(root, pts_arr.data(), n, batch_grad.data()); + + auto fn = [&](const Point3D& p) { return evaluate(root, p); }; + double h = 1e-6; + + for (int i = 0; i < n; ++i) { + Point3D p(pts_arr[i*3], pts_arr[i*3+1], pts_arr[i*3+2]); + + // Sequential gradient via central FD + double dfdx = (fn(Point3D(p.x()+h, p.y(), p.z())) - + fn(Point3D(p.x()-h, p.y(), p.z()))) / (2.0*h); + double dfdy = (fn(Point3D(p.x(), p.y()+h, p.z())) - + fn(Point3D(p.x(), p.y()-h, p.z()))) / (2.0*h); + double dfdz = (fn(Point3D(p.x(), p.y(), p.z()+h)) - + fn(Point3D(p.x(), p.y(), p.z()-h))) / (2.0*h); + + EXPECT_NEAR(batch_grad[i*3], dfdx, 1e-4); + EXPECT_NEAR(batch_grad[i*3+1], dfdy, 1e-4); + EXPECT_NEAR(batch_grad[i*3+2], dfdz, 1e-4); + } +} + +TEST(SdfTorchBridge, GradientBatch_Box) { + auto root = SdfNode::box(Point3D(2.0, 1.0, 3.0)); + const int n = 50; + auto pts_arr = random_points(n); + + std::vector batch_grad(3 * n); + gradient_batch(root, pts_arr.data(), n, batch_grad.data()); + + // Check that gradient at origin (inside box) is zero or close + Point3D origin(0, 0, 0); + Vector3D g = gradient([&](const Point3D& p) { return evaluate(root, p); }, origin); + // Gradient near origin should point outward (or be near-internal max) + (void)g; // just verifying no crash + + for (int i = 0; i < n; ++i) { + // All gradients should be finite + EXPECT_TRUE(std::isfinite(batch_grad[i*3])); + EXPECT_TRUE(std::isfinite(batch_grad[i*3+1])); + EXPECT_TRUE(std::isfinite(batch_grad[i*3+2])); + } +} + +// ═══════════════════════════════════════════════════ +// evaluate_with_gradient_batch +// ═══════════════════════════════════════════════════ + +TEST(SdfTorchBridge, EvaluateWithGradientBatch_BothCorrect) { + auto root = SdfNode::sphere(1.5); + const int n = 50; + auto pts_arr = random_points(n); + + std::vector batch_dist(n); + std::vector batch_grad(3 * n); + evaluate_with_gradient_batch(root, pts_arr.data(), n, + batch_dist.data(), batch_grad.data()); + + // Compare distances against sequential + std::vector seq_dist(n); + evaluate_batch(root, pts_arr.data(), n, seq_dist.data()); + for (int i = 0; i < n; ++i) { + EXPECT_NEAR(batch_dist[i], seq_dist[i], EPS); + } + + // Compare gradients against sequential + std::vector seq_grad(3 * n); + gradient_batch(root, pts_arr.data(), n, seq_grad.data()); + for (int i = 0; i < 3 * n; ++i) { + EXPECT_NEAR(batch_grad[i], seq_grad[i], 1e-10); + } +} + +TEST(SdfTorchBridge, EvaluateWithGradientBatch_CSGTree) { + // Smooth union of sphere and box + auto sphere = SdfNode::sphere(1.0); + auto box_node = SdfNode::box(Point3D(0.8, 0.8, 0.8)); + auto root = SdfNode::smooth_union(sphere, box_node, 0.2); + const int n = 30; + auto pts_arr = random_points(n); + + std::vector batch_dist(n); + std::vector batch_grad(3 * n); + evaluate_with_gradient_batch(root, pts_arr.data(), n, + batch_dist.data(), batch_grad.data()); + + for (int i = 0; i < n; ++i) { + EXPECT_TRUE(std::isfinite(batch_dist[i])); + EXPECT_TRUE(std::isfinite(batch_grad[i*3])); + EXPECT_TRUE(std::isfinite(batch_grad[i*3+1])); + EXPECT_TRUE(std::isfinite(batch_grad[i*3+2])); + } +}