From ab5f4fdc4c23c4cf78683137cd962cc69067eea3 Mon Sep 17 00:00:00 2001 From: Chris Barnes Date: Mon, 18 Dec 2023 16:52:18 +0000 Subject: [PATCH] Private method for SDF internals --- python/ncollpyde/_ncollpyde.pyi | 3 ++ python/ncollpyde/main.py | 39 ++++++++++++++++----- src/interface.rs | 62 +++++++++++++++++++++++++++++++++ src/utils.rs | 13 +++++-- 4 files changed, 106 insertions(+), 11 deletions(-) diff --git a/python/ncollpyde/_ncollpyde.pyi b/python/ncollpyde/_ncollpyde.pyi index f94a936..1d60b2c 100644 --- a/python/ncollpyde/_ncollpyde.pyi +++ b/python/ncollpyde/_ncollpyde.pyi @@ -30,3 +30,6 @@ class TriMeshWrapper: def intersections_many_threaded( self, src_points: Points, tgt_points: Points ) -> Tuple[List[int], Points, List[bool]]: ... + def sdf_intersections( + self, points: Points, vectors: Points + ) -> Tuple[npt.NDArray[np.float_], npt.NDArray[np.float_]]: ... diff --git a/python/ncollpyde/main.py b/python/ncollpyde/main.py index 0e651ba..bc8a463 100644 --- a/python/ncollpyde/main.py +++ b/python/ncollpyde/main.py @@ -2,7 +2,7 @@ import random import warnings from multiprocessing import cpu_count -from typing import TYPE_CHECKING, Optional, Tuple, Union +from typing import TYPE_CHECKING, Optional, Tuple, Union, List import numpy as np from numpy.typing import ArrayLike, NDArray @@ -218,6 +218,34 @@ def contains( return self._impl.contains(coords, self._interpret_threads(threads)) + def _as_points(self, points: ArrayLike) -> NDArray: + p = np.asarray(points, self.dtype) + if p.shape[1:] != (3,): + raise ValueError("Points must be Nx3 array-like") + return p + + def _validate_points(self, *points: ArrayLike) -> List[NDArray]: + """Ensure that arrays are equal-length sets of points.""" + ndim = None + out = [] + + for p_raw in points: + p = self._as_points(p_raw) + nd = p.shape[1:] + if ndim is None: + ndim = nd + elif ndim != nd: + raise ValueError("Point arrays are not the same shape") + out.append(p) + + return out + + def _sdf_intersections( + self, points: ArrayLike, vectors: ArrayLike + ) -> Tuple[NDArray, NDArray]: + p, v = self._validate_points(points, vectors) + return self._impl.sdf_intersections(p, v) + def intersections( self, src_points: ArrayLike, @@ -257,14 +285,7 @@ def intersections( float array of locations (Nx3), bool array of is_backface (N) """ - src = np.asarray(src_points, self.dtype) - tgt = np.asarray(tgt_points, self.dtype) - - if src.shape != tgt.shape: - raise ValueError("Source and target points arrays must be the same shape") - - if src.shape[1:] != (3,): - raise ValueError("Points must be Nx3 array-like") + src, tgt = self._validate_points(src_points, tgt_points) if self._interpret_threads(threads): idxs, points, bfs = self._impl.intersections_many_threaded( diff --git a/src/interface.rs b/src/interface.rs index f3a5319..7e2eeee 100644 --- a/src/interface.rs +++ b/src/interface.rs @@ -4,6 +4,7 @@ use std::iter::repeat_with; use numpy::ndarray::{Array, Zip}; use numpy::{IntoPyArray, PyArray1, PyArray2, PyReadonlyArray2}; use parry3d_f64::math::{Point, Vector}; +use parry3d_f64::query::{Ray, RayCast}; use parry3d_f64::shape::TriMesh; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; @@ -157,6 +158,67 @@ impl TriMeshWrapper { PyArray2::from_vec2(py, &[point_to_vec(&aabb.mins), point_to_vec(&aabb.maxs)]).unwrap() } + pub fn sdf_intersections<'py>( + &self, + py: Python<'py>, + points: PyReadonlyArray2, + vecs: PyReadonlyArray2, + ) -> (&'py PyArray1, &'py PyArray1) { + let diameter = self.mesh.local_bounding_sphere().radius() * 2.0; + + let (dists, dot_norms): (Vec, Vec) = points + .as_array() + .rows() + .into_iter() + .map(|p| Point::new(p[0], p[1], p[2])) + .zip( + vecs.as_array() + .rows() + .into_iter() + .map(|v| Vector::new(v[0], v[1], v[2]).normalize()), + ) + .map(|(p, v)| { + let ray = Ray::new(p, v); + if let Some(inter) = self.mesh.cast_local_ray_and_get_normal( + &ray, diameter, false, // unused + ) { + (inter.toi, v.dot(&inter.normal)) + } else { + (Precision::INFINITY, Precision::NAN) + } + }) + .unzip(); + ( + PyArray1::from_vec(py, dists), + PyArray1::from_vec(py, dot_norms), + ) + } + + // pub fn sdf_intersections_threaded<'py>( + // &self, + // py: Python<'py>, + // points: PyReadonlyArray2, + // vecs: PyReadonlyArray2, + // ) -> (&'py PyArray1, &'py PyArray1) { + // let diameter = self.mesh.local_bounding_sphere().radius() * 2.0; + + // Zip::from(points.as_array().rows()) + // .and(vecs.as_array().rows()) + // .par_map_collect(|point, vector| { + // let p = Point::new(point[0], point[1], point[2]); + // let v = Vector::new(vector[0], vector[1], vector[2]).normalize(); + + // let ray = Ray::new(p, v); + // if let Some(inter) = self.mesh.cast_local_ray_and_get_normal( + // &ray, diameter, false, // unused + // ) { + // (inter.toi, v.dot(&inter.normal)) + // } else { + // (Precision::INFINITY, Precision::NAN) + // } + // }) + // } + pub fn intersections_many<'py>( &self, py: Python<'py>, diff --git a/src/utils.rs b/src/utils.rs index b1e37a3..2081de6 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -1,6 +1,6 @@ use parry3d_f64::math::{Isometry, Point, Vector}; use parry3d_f64::query::{PointQuery, Ray, RayCast}; -use parry3d_f64::shape::TriMesh; +use parry3d_f64::shape::{FeatureId, TriMesh}; use rand::Rng; pub type Precision = f64; @@ -61,11 +61,20 @@ pub fn points_cross_mesh( src_point: &Point, tgt_point: &Point, ) -> Option<(Point, bool)> { + points_cross_mesh_info(mesh, src_point, tgt_point) + .map(|(inter, _, ft)| (inter, mesh.is_backface(ft))) +} + +pub fn points_cross_mesh_info( + mesh: &TriMesh, + src_point: &Point, + tgt_point: &Point, +) -> Option<(Point, Vector, FeatureId)> { let ray = Ray::new(*src_point, tgt_point - src_point); mesh.cast_local_ray_and_get_normal( &ray, 1.0, false, // unused ) - .map(|i| (ray.point_at(i.toi), mesh.is_backface(i.feature))) + .map(|i| (ray.point_at(i.toi), i.normal, i.feature)) } pub fn dist_from_mesh(mesh: &TriMesh, point: &Point, rays: Option<&[Vector]>) -> f64 {