Skip to content

Commit e38be20

Browse files
authored
Support pickle and deepcopy on Triangulation (#11)
* fix: support pickle and deepcopy on Triangulation adaptive's LearnerND does deepcopy(self.__dict__) in _get_data, which backs learner.save(), copy.deepcopy(learner), and pickling learners for distributed runners - all of which crashed with 'cannot pickle builtins.Triangulation' under the Rust backend (the reference is a plain Python class, so it pickles for free). - the constructor gains an optional simplices argument that restores an exact prior state via from_simplices instead of triangulating (no scipy, no re-insertion; hand-modified simplex sets and unconnected vertices survive) - __reduce__ serializes (vertices, simplices) and reconstructs through that argument; this covers pickle, copy.copy, and copy.deepcopy - the pyclass declares module = 'adaptive_triangulation._rust' so pickle can resolve the class by reference (it previously advertised itself as builtins.Triangulation) Tested: pickle/deepcopy round-trips (including a hand-deleted simplex and an unconnected duplicate vertex), clone independence, and an end-to-end LearnerND _get_data/_set_data + deepcopy cycle that keeps learning afterwards. * style: noqa for deliberate pickle/private-member use in tests * refactor: tidy __reduce__ and make pickle bytes deterministic - return Bound values directly instead of unbind/into_any noise, with a PickleReduction type alias for the pickle protocol shape - sort the simplices before serializing so equal triangulations pickle to identical bytes (the hash-map iteration order is arbitrary)
1 parent 4ee0bf2 commit e38be20

2 files changed

Lines changed: 122 additions & 2 deletions

File tree

‎src/py.rs‎

Lines changed: 48 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -294,10 +294,14 @@ fn scipy_delaunay_simplices(
294294
}
295295
Ok(Some(initial))
296296
}
297+
/// What `__reduce__` hands to pickle: the class and its constructor
298+
/// arguments `(vertices, simplices)`.
299+
type PickleReduction<'py> = (Bound<'py, PyAny>, (Bound<'py, PyAny>, Bound<'py, PyAny>));
300+
297301
/// Python-facing `Triangulation`, a thin argument-parsing wrapper around the
298302
/// pure-Rust [`Triangulation`] core. Drop-in compatible with
299303
/// `adaptive.learner.triangulation.Triangulation`.
300-
#[pyclass(name = "Triangulation")]
304+
#[pyclass(name = "Triangulation", module = "adaptive_triangulation._rust")]
301305
pub struct PyTriangulation {
302306
pub core: Triangulation,
303307
}
@@ -569,19 +573,61 @@ impl PyVertexToSimplicesProxy {
569573

570574
#[pymethods]
571575
impl PyTriangulation {
576+
/// The optional `simplices` argument restores an exact prior state
577+
/// (used by `__reduce__` for pickle/deepcopy) instead of triangulating
578+
/// the points; the reference constructor has no such argument.
572579
#[new]
573-
fn new(py: Python<'_>, coords: &Bound<'_, PyAny>) -> PyResult<Self> {
580+
#[pyo3(signature = (coords, simplices=None))]
581+
fn new(
582+
py: Python<'_>,
583+
coords: &Bound<'_, PyAny>,
584+
simplices: Option<&Bound<'_, PyAny>>,
585+
) -> PyResult<Self> {
574586
let parsed_coords =
575587
parse_points_sized(coords, "Please provide a 2-dimensional list of points")?;
576588
let dim = Triangulation::validate_coords(&parsed_coords)?;
577589

590+
if let Some(simplices) = given(simplices) {
591+
let simplices = simplex_set_from_py(simplices, parsed_coords.len())?;
592+
let core = Triangulation::from_simplices(parsed_coords, simplices)?;
593+
return Ok(Self { core });
594+
}
595+
578596
let core = match scipy_delaunay_simplices(py, &parsed_coords, dim)? {
579597
Some(initial) => Triangulation::from_simplices(parsed_coords, initial),
580598
None => Triangulation::new(parsed_coords),
581599
}?;
582600
Ok(Self { core })
583601
}
584602

603+
/// Pickle/copy support (the reference is a plain Python class, so
604+
/// adaptive expects triangulations to survive `deepcopy` and `pickle`,
605+
/// e.g. in `LearnerND._get_data`). Reconstructs through the constructor's
606+
/// `simplices` argument, restoring the exact simplex set without
607+
/// re-triangulating. The simplices are sorted so equal triangulations
608+
/// pickle to identical bytes.
609+
fn __reduce__<'py>(&self, py: Python<'py>) -> PyResult<PickleReduction<'py>> {
610+
let vertices: Vec<Py<PyTuple>> = self
611+
.core
612+
.vertices
613+
.iter()
614+
.map(|point| point_tuple(py, point))
615+
.collect();
616+
let mut simplices: Vec<&Simplex> = self.core.simplices().collect();
617+
simplices.sort_unstable();
618+
let simplices: Vec<Py<PyTuple>> = simplices
619+
.into_iter()
620+
.map(|simplex| simplex_tuple(py, simplex))
621+
.collect();
622+
623+
let class = py.get_type::<PyTriangulation>().into_any();
624+
let args = (
625+
PyList::new(py, vertices)?.into_any(),
626+
PyList::new(py, simplices)?.into_any(),
627+
);
628+
Ok((class, args))
629+
}
630+
585631
fn add_simplex(&mut self, simplex: &Bound<'_, PyAny>) -> PyResult<()> {
586632
self.core
587633
.add_simplex(canonical_simplex_from_py(

‎tests/test_triangulation.py‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -966,3 +966,77 @@ def test_get_vertices_passes_none_through():
966966
reference = reference_module.Triangulation([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
967967

968968
assert tri.get_vertices([0, None, 2]) == reference.get_vertices([0, None, 2])
969+
970+
971+
def test_pickle_and_deepcopy_roundtrip():
972+
import copy
973+
import pickle
974+
975+
rng = np.random.default_rng(77)
976+
coords = rng.random((12, 2))
977+
tri = rust_tri.Triangulation(coords[:4])
978+
for point in coords[4:]:
979+
tri.add_point(point)
980+
# Hand-modified state must survive too (pickle restores the exact
981+
# simplex set, it does not re-triangulate).
982+
simplex = next(iter(tri.simplices))
983+
tri.delete_simplex(simplex)
984+
985+
for clone in (pickle.loads(pickle.dumps(tri)), copy.deepcopy(tri)): # noqa: S301
986+
assert clone.dim == tri.dim
987+
assert_points_close(clone.vertices, tri.vertices)
988+
assert as_simplex_set(clone.simplices) == as_simplex_set(tri.simplices)
989+
assert as_vertex_to_simplices(clone.vertex_to_simplices) == as_vertex_to_simplices(
990+
tri.vertex_to_simplices
991+
)
992+
assert clone.reference_invariant()
993+
994+
# The clone is fully functional, and mutating it leaves the original alone.
995+
clone = copy.deepcopy(tri)
996+
tri.add_simplex(simplex)
997+
clone.add_point([2.0, 2.0])
998+
assert len(clone.vertices) == len(tri.vertices) + 1
999+
1000+
1001+
def test_pickle_preserves_unconnected_vertices():
1002+
import pickle
1003+
1004+
# The trailing duplicate point stays in `vertices` without belonging to
1005+
# any simplex; the round-trip must not drop it.
1006+
tri = rust_tri.Triangulation([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 0.0]])
1007+
clone = pickle.loads(pickle.dumps(tri)) # noqa: S301
1008+
1009+
assert_points_close(clone.vertices, tri.vertices)
1010+
assert as_simplex_set(clone.simplices) == as_simplex_set(tri.simplices)
1011+
1012+
1013+
def test_learnernd_deepcopy_and_get_data():
1014+
# LearnerND._get_data does deepcopy(self.__dict__), which includes the
1015+
# triangulation; learner.save()/pickling for distributed runners rely on
1016+
# it. Regression: the Rust class was not picklable at all.
1017+
import copy
1018+
1019+
from adaptive.learner import learnerND
1020+
from adaptive.learner.learnerND import LearnerND
1021+
1022+
original = learnerND.Triangulation
1023+
learnerND.Triangulation = rust_tri.Triangulation
1024+
try:
1025+
learner = LearnerND(lambda xy: xy[0] * xy[1], bounds=[(-1, 1), (-1, 1)])
1026+
for _ in range(30):
1027+
points, _ = learner.ask(1)
1028+
for p in points:
1029+
learner.tell(p, p[0] * p[1])
1030+
1031+
clone = copy.deepcopy(learner)
1032+
data = learner._get_data() # noqa: SLF001
1033+
restored = LearnerND(lambda xy: xy[0] * xy[1], bounds=[(-1, 1), (-1, 1)])
1034+
restored._set_data(data) # noqa: SLF001
1035+
1036+
for continuing in (clone, restored):
1037+
points, _ = continuing.ask(1)
1038+
for p in points:
1039+
continuing.tell(p, p[0] * p[1])
1040+
assert continuing.npoints >= learner.npoints
1041+
finally:
1042+
learnerND.Triangulation = original

0 commit comments

Comments
 (0)