Skip to content

Commit 4b44e56

Browse files
authored
Complete the drop-in API for neighbor-aware losses; numpy fast paths in the PyO3 layer (#7)
* feat: complete the drop-in surface for adaptive's neighbor-aware losses LearnerND with a neighbor-aware loss (triangle_loss / curvature_loss, nth_neighbors >= 1) calls five Triangulation methods that were missing, plus two behaviors only that path exercises. Add them: - get_vertex, get_neighbors_from_vertices, get_simplices_attached_to_points, get_opposing_vertices, get_face_sharing_neighbors - same semantics and exception types as the reference; get_opposing_vertices and get_simplices_attached_to_points are answered natively from the facet/vertex indexes instead of the reference's set scans - get_vertices passes None entries through as None (adaptive feeds the result of get_opposing_vertices, which contains None on the hull, straight back in) - SimplicesProxy supports the binary set operators in both operand orders (adaptive computes `... - tri.simplices`), delegating to a snapshot set Pinned by cross-validation tests against the reference on random 2D/3D triangulations, exception-type parity for unknown simplices, and an end-to-end LearnerND run with curvature_loss_function() driving 100 points through the previously uncovered code path. * perf: numpy fast paths in the PyO3 layer - parse_point / parse_points bulk-copy f64 numpy arrays instead of iterating them as Python objects (other dtypes and nested sequences keep the generic path) - VerticesProxy.__array__ builds a PyArray2 directly instead of round-tripping through a list of Python tuples and numpy.array; dtype is honored via astype, and copy needs no handling because the snapshot is always freshly allocated Output types are unchanged everywhere (tuples/lists/sets), so the drop-in contract is untouched; all 112 tests pass. A/B microbenchmark vs main (best of 5): np.asarray(tri.vertices) snapshots 50x faster (0.189s -> 0.004s per 2000 calls on a 500-vertex mesh), constructor from a numpy array ~14% faster, add_point unchanged. * style: avoid camelcase-as-lowercase import alias in LearnerND test (ruff N813)
1 parent 810ad01 commit 4b44e56

3 files changed

Lines changed: 358 additions & 32 deletions

File tree

‎src/py.rs‎

Lines changed: 190 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,13 @@
66
//! Python exception types raised by `adaptive.learner.triangulation` in the
77
//! same situations.
88
9-
use numpy::PyArray2;
9+
use numpy::{PyArray2, PyReadonlyArray1, PyReadonlyArray2};
1010
use pyo3::exceptions::{
1111
PyAssertionError, PyIndexError, PyNotImplementedError, PyRuntimeError, PyTypeError,
1212
PyValueError, PyZeroDivisionError,
1313
};
1414
use pyo3::prelude::*;
15-
use pyo3::types::{PyAny, PyDict, PyList, PyModule, PySet, PyTuple};
15+
use pyo3::types::{PyAny, PyList, PyModule, PySet, PyTuple};
1616
use rustc_hash::FxHashSet;
1717

1818
use crate::geometry::GeometryError;
@@ -34,6 +34,10 @@ impl TriangulationError {
3434
}
3535

3636
pub(crate) fn parse_point(obj: &Bound<'_, PyAny>) -> PyResult<Vec<f64>> {
37+
// Bulk-copy f64 numpy arrays instead of iterating Python objects.
38+
if let Ok(array) = obj.extract::<PyReadonlyArray1<'_, f64>>() {
39+
return Ok(array.as_array().to_vec());
40+
}
3741
let Ok(iter) = obj.try_iter() else {
3842
return Err(PyTypeError::new_err("Expected an iterable of floats"));
3943
};
@@ -62,6 +66,15 @@ fn parse_points_impl(
6266
type_error_message: &str,
6367
require_sized: bool,
6468
) -> PyResult<Vec<Vec<f64>>> {
69+
// Bulk-copy f64 numpy arrays instead of iterating Python objects row by
70+
// row; other dtypes and nested sequences take the generic path below.
71+
if let Ok(array) = obj.extract::<PyReadonlyArray2<'_, f64>>() {
72+
return Ok(array
73+
.as_array()
74+
.outer_iter()
75+
.map(|row| row.to_vec())
76+
.collect());
77+
}
6578
if require_sized {
6679
ensure_sized(obj, type_error_message)?;
6780
}
@@ -217,14 +230,6 @@ pub(crate) fn simplex_set_py<'a>(
217230
Ok(PySet::new(py, &tuples)?.into())
218231
}
219232

220-
pub(crate) fn point_list_py(py: Python<'_>, points: &[Vec<f64>]) -> PyResult<Py<PyAny>> {
221-
let tuples: Vec<Py<PyAny>> = points
222-
.iter()
223-
.map(|point| point_tuple(py, point).into())
224-
.collect();
225-
Ok(PyList::new(py, tuples)?.into())
226-
}
227-
228233
pub(crate) fn index_list_py(py: Python<'_>, indices: &[usize]) -> PyResult<Py<PyAny>> {
229234
Ok(PyList::new(py, indices.iter().copied())?.into())
230235
}
@@ -315,6 +320,16 @@ impl PySimplicesProxy {
315320
vertex: Some(vertex),
316321
}
317322
}
323+
324+
/// The current simplices as a real Python set, for the set-operator
325+
/// protocol below.
326+
fn snapshot_set(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
327+
let triangulation = self.triangulation.bind(py).borrow();
328+
match self.vertex {
329+
Some(vertex) => simplex_set_py(py, triangulation.core.simplices_of(vertex)),
330+
None => simplex_set_py(py, triangulation.core.simplices()),
331+
}
332+
}
318333
}
319334

320335
/// Lazy, sequence-like view of the vertex coordinates (supports `__array__`
@@ -406,6 +421,66 @@ impl PySimplicesProxy {
406421
None => triangulation.core.num_simplices(),
407422
}
408423
}
424+
425+
// The reference exposes `simplices` as a real set, so callers combine it
426+
// with sets freely (adaptive does `... - tri.simplices`). Delegate the
427+
// binary set operators to a snapshot set, in both operand orders.
428+
429+
fn __sub__(&self, py: Python<'_>, other: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
430+
Ok(self
431+
.snapshot_set(py)?
432+
.bind(py)
433+
.call_method1("__sub__", (other,))?
434+
.unbind())
435+
}
436+
437+
fn __rsub__(&self, py: Python<'_>, other: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
438+
Ok(other
439+
.call_method1("__sub__", (self.snapshot_set(py)?,))?
440+
.unbind())
441+
}
442+
443+
fn __and__(&self, py: Python<'_>, other: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
444+
Ok(self
445+
.snapshot_set(py)?
446+
.bind(py)
447+
.call_method1("__and__", (other,))?
448+
.unbind())
449+
}
450+
451+
fn __rand__(&self, py: Python<'_>, other: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
452+
Ok(other
453+
.call_method1("__and__", (self.snapshot_set(py)?,))?
454+
.unbind())
455+
}
456+
457+
fn __or__(&self, py: Python<'_>, other: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
458+
Ok(self
459+
.snapshot_set(py)?
460+
.bind(py)
461+
.call_method1("__or__", (other,))?
462+
.unbind())
463+
}
464+
465+
fn __ror__(&self, py: Python<'_>, other: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
466+
Ok(other
467+
.call_method1("__or__", (self.snapshot_set(py)?,))?
468+
.unbind())
469+
}
470+
471+
fn __xor__(&self, py: Python<'_>, other: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
472+
Ok(self
473+
.snapshot_set(py)?
474+
.bind(py)
475+
.call_method1("__xor__", (other,))?
476+
.unbind())
477+
}
478+
479+
fn __rxor__(&self, py: Python<'_>, other: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
480+
Ok(other
481+
.call_method1("__xor__", (self.snapshot_set(py)?,))?
482+
.unbind())
483+
}
409484
}
410485

411486
#[pymethods]
@@ -466,22 +541,16 @@ impl PyVerticesProxy {
466541
dtype: Option<&Bound<'_, PyAny>>,
467542
copy: Option<bool>,
468543
) -> PyResult<Py<PyAny>> {
544+
// The snapshot is always a freshly allocated array, so the `copy`
545+
// argument needs no special handling. Build it directly instead of
546+
// round-tripping through a list of Python tuples.
547+
let _ = copy;
469548
let triangulation = self.triangulation.bind(py).borrow();
470-
let vertices = point_list_py(py, &triangulation.core.vertices)?;
471-
let numpy = PyModule::import(py, "numpy")?;
472-
let kwargs = PyDict::new(py);
473-
if let Some(dtype) = dtype {
474-
kwargs.set_item("dtype", dtype)?;
549+
let array = PyArray2::from_vec2(py, &triangulation.core.vertices)?;
550+
match dtype {
551+
Some(dtype) if !dtype.is_none() => Ok(array.call_method1("astype", (dtype,))?.unbind()),
552+
_ => Ok(array.into_any().unbind()),
475553
}
476-
let array = if copy == Some(false) {
477-
numpy.call_method("asarray", (vertices,), Some(&kwargs))?
478-
} else {
479-
if let Some(copy) = copy {
480-
kwargs.set_item("copy", copy)?;
481-
}
482-
numpy.call_method("array", (vertices,), Some(&kwargs))?
483-
};
484-
Ok(array.into())
485554
}
486555
}
487556

@@ -609,20 +678,32 @@ impl PyTriangulation {
609678
self.core.dim
610679
}
611680

681+
/// `None` entries pass through as `None`, like the reference (which maps
682+
/// every index through `get_vertex`); adaptive relies on this when it
683+
/// feeds the result of `get_opposing_vertices` straight back in.
612684
#[pyo3(name = "get_vertices")]
613685
fn get_vertices_method(
614686
&self,
615687
py: Python<'_>,
616688
indices: &Bound<'_, PyAny>,
617689
) -> PyResult<Py<PyAny>> {
618-
let indices = ordered_indices_from_py(indices, self.core.vertices.len())?;
619-
point_list_py(
620-
py,
621-
&self
622-
.core
623-
.get_vertices(&indices)
624-
.map_err(TriangulationError::into_pyerr)?,
625-
)
690+
let Ok(iter) = indices.try_iter() else {
691+
return Err(PyTypeError::new_err(
692+
"Expected an iterable of vertex indices",
693+
));
694+
};
695+
let mut items: Vec<Py<PyAny>> = Vec::new();
696+
for item in iter {
697+
let item = item?;
698+
if item.is_none() {
699+
items.push(py.None());
700+
continue;
701+
}
702+
let index = normalize_index(item.extract::<isize>()?, self.core.vertices.len())
703+
.map_err(TriangulationError::into_pyerr)?;
704+
items.push(point_tuple(py, &self.core.vertices[index]).into());
705+
}
706+
Ok(PyList::new(py, items)?.into())
626707
}
627708

628709
fn locate_point(&self, py: Python<'_>, point: &Bound<'_, PyAny>) -> PyResult<Py<PyAny>> {
@@ -824,6 +905,83 @@ impl PyTriangulation {
824905
Ok((simplex_set_py(py, &deleted)?, simplex_set_py(py, &added)?))
825906
}
826907

908+
#[pyo3(signature = (index))]
909+
fn get_vertex(&self, py: Python<'_>, index: Option<isize>) -> PyResult<Py<PyAny>> {
910+
match index {
911+
None => Ok(py.None()),
912+
Some(index) => {
913+
let index = normalize_index(index, self.core.vertices.len())
914+
.map_err(TriangulationError::into_pyerr)?;
915+
Ok(point_tuple(py, &self.core.vertices[index]).into())
916+
}
917+
}
918+
}
919+
920+
fn get_neighbors_from_vertices(
921+
&self,
922+
py: Python<'_>,
923+
simplex: &Bound<'_, PyAny>,
924+
) -> PyResult<Py<PyAny>> {
925+
let indices = ordered_indices_from_py(simplex, self.core.vertices.len())?;
926+
let neighbors = self
927+
.core
928+
.neighbors_from_vertices(&indices)
929+
.map_err(TriangulationError::into_pyerr)?;
930+
simplex_set_py(py, &neighbors)
931+
}
932+
933+
fn get_simplices_attached_to_points(
934+
&self,
935+
py: Python<'_>,
936+
indices: &Bound<'_, PyAny>,
937+
) -> PyResult<Py<PyAny>> {
938+
let indices = ordered_indices_from_py(indices, self.core.vertices.len())?;
939+
let attached = self
940+
.core
941+
.simplices_attached_to_points(&indices)
942+
.map_err(TriangulationError::into_pyerr)?;
943+
simplex_set_py(py, &attached)
944+
}
945+
946+
fn get_opposing_vertices(
947+
&self,
948+
py: Python<'_>,
949+
simplex: &Bound<'_, PyAny>,
950+
) -> PyResult<Py<PyTuple>> {
951+
let simplex = ordered_indices_from_py(simplex, self.core.vertices.len())?;
952+
let opposing = self
953+
.core
954+
.opposing_vertices(&simplex)
955+
.map_err(TriangulationError::into_pyerr)?;
956+
Ok(PyTuple::new(py, opposing)?.into())
957+
}
958+
959+
/// Keep only the simplices sharing a whole face with `simplex`. A pure
960+
/// set filter on the arguments, like the reference implementation; the
961+
/// simplices are not required to be part of the triangulation.
962+
fn get_face_sharing_neighbors(
963+
&self,
964+
py: Python<'_>,
965+
neighbors: &Bound<'_, PyAny>,
966+
simplex: &Bound<'_, PyAny>,
967+
) -> PyResult<Py<PyAny>> {
968+
let neighbors = simplex_set_from_py(neighbors, self.core.vertices.len())?;
969+
let simplex = ordered_indices_from_py(simplex, self.core.vertices.len())?;
970+
let simplex_set: FxHashSet<usize> = simplex.iter().copied().collect();
971+
let sharing: FxHashSet<Simplex> = neighbors
972+
.into_iter()
973+
.filter(|neighbor| {
974+
let shared: FxHashSet<usize> = neighbor
975+
.iter()
976+
.copied()
977+
.filter(|vertex| simplex_set.contains(vertex))
978+
.collect();
979+
shared.len() == self.core.dim
980+
})
981+
.collect();
982+
simplex_set_py(py, &sharing)
983+
}
984+
827985
fn vertex_invariant(&self, _vertex: usize) -> PyResult<bool> {
828986
Err(PyNotImplementedError::new_err("vertex_invariant"))
829987
}

‎src/triangulation.rs‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -822,6 +822,89 @@ impl Triangulation {
822822
.collect())
823823
}
824824

825+
/// All simplices that share at least one vertex with the given vertex
826+
/// set (the union of the vertices' simplex sets).
827+
pub fn neighbors_from_vertices(
828+
&self,
829+
indices: &[usize],
830+
) -> Result<FxHashSet<Simplex>, TriangulationError> {
831+
self.validate_simplex_indices(indices)?;
832+
let mut neighbors = FxHashSet::default();
833+
for &vertex in indices {
834+
neighbors.extend(self.simplices_of(vertex).cloned());
835+
}
836+
Ok(neighbors)
837+
}
838+
839+
/// All simplices sharing exactly `dim` vertices (a whole facet) with the
840+
/// given vertex set. For a full simplex these are its facet neighbours;
841+
/// the simplex itself (sharing dim+1 vertices) is excluded.
842+
pub fn simplices_attached_to_points(
843+
&self,
844+
indices: &[usize],
845+
) -> Result<FxHashSet<Simplex>, TriangulationError> {
846+
self.validate_simplex_indices(indices)?;
847+
let index_set: FxHashSet<usize> = indices.iter().copied().collect();
848+
let mut seen: FxHashSet<SimplexId> = FxHashSet::default();
849+
let mut attached = FxHashSet::default();
850+
for &vertex in indices {
851+
for &id in &self.vertex_to_ids[vertex] {
852+
if !seen.insert(id) {
853+
continue;
854+
}
855+
let simplex = self.simplex_by_id(id);
856+
let shared = simplex
857+
.iter()
858+
.filter(|vertex| index_set.contains(vertex))
859+
.count();
860+
if shared == self.dim {
861+
attached.insert(simplex.clone());
862+
}
863+
}
864+
}
865+
Ok(attached)
866+
}
867+
868+
/// For each vertex of `simplex` (in the given order), the vertex of the
869+
/// facet neighbour on the other side of the opposite facet, or `None`
870+
/// when that facet lies on the hull. Errors when the simplex is not part
871+
/// of the triangulation.
872+
pub fn opposing_vertices(
873+
&self,
874+
simplex: &[usize],
875+
) -> Result<Vec<Option<PointIndex>>, TriangulationError> {
876+
let mut sorted = simplex.to_vec();
877+
sorted.sort_unstable();
878+
self.validate_simplex_indices(&sorted)?;
879+
let Some(&own_id) = self.ids.get(&sorted) else {
880+
return Err(TriangulationError::Value(
881+
"Provided simplex is not part of the triangulation".to_string(),
882+
));
883+
};
884+
885+
simplex
886+
.iter()
887+
.map(|&vertex| {
888+
let position = sorted
889+
.iter()
890+
.position(|&other| other == vertex)
891+
.expect("vertex comes from the simplex itself");
892+
let facet = facet_excluding(&sorted, position);
893+
let neighbour = self
894+
.facets
895+
.get(&facet)
896+
.and_then(|incident| incident.iter().copied().find(|&id| id != own_id));
897+
Ok(neighbour.map(|id| {
898+
self.simplex_by_id(id)
899+
.iter()
900+
.copied()
901+
.find(|other| !sorted.contains(other))
902+
.expect("facet neighbour has exactly one vertex outside the simplex")
903+
}))
904+
})
905+
.collect()
906+
}
907+
825908
fn simplex_points(
826909
&self,
827910
simplex: &[usize],

0 commit comments

Comments
 (0)