From 0bac3e2f16e77bf88be45f1f02ef3e9c2848e2e0 Mon Sep 17 00:00:00 2001 From: fncnt Date: Mon, 15 Dec 2025 13:24:51 +0100 Subject: [PATCH 1/3] Python: minor tweaks to psssh.py example --- python/examples/psssh.py | 1 + python/examples/requirements.txt | 1 + 2 files changed, 2 insertions(+) diff --git a/python/examples/psssh.py b/python/examples/psssh.py index 0cc5472..3dc3d59 100644 --- a/python/examples/psssh.py +++ b/python/examples/psssh.py @@ -121,6 +121,7 @@ def harmonic_potential(x, r, k): compiled = nutpie.compile_pymc_model(model) # zelll's Python bindings don't benefit from nutpie's multithreading # unless a free-threaded Python build is used + # however, num_dual currently has no pre-built free-threaded wheels trace = nutpie.sample( compiled, save_warmup=False, diff --git a/python/examples/requirements.txt b/python/examples/requirements.txt index 47cb536..ca3f0ff 100644 --- a/python/examples/requirements.txt +++ b/python/examples/requirements.txt @@ -4,3 +4,4 @@ numpy==2.2.6 nutpie==0.15.2 pymc==5.25.1 pytensor==2.31.7 +numba==0.63.1 From 41679de81bf53f5ea06509f7f8229d074a2b3349 Mon Sep 17 00:00:00 2001 From: fncnt Date: Mon, 15 Dec 2025 13:25:19 +0100 Subject: [PATCH 2/3] Python: update dep:pyo3 --- python/Cargo.toml | 4 ++-- python/src/lib.rs | 35 +++++++++++++++++++++++------------ 2 files changed, 25 insertions(+), 14 deletions(-) diff --git a/python/Cargo.toml b/python/Cargo.toml index f43a32a..7f599dd 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -17,9 +17,9 @@ crate-type = ["cdylib"] [dependencies] bincode = { version = "2.0", features = ["serde"] } -pyo3 = { version = "0.26", features = ["experimental-inspect"] } +pyo3 = { version = "0.27", features = ["experimental-inspect"] } zelll = { path = "..", features = ["serde"] } serde = { version = "1.0" } [dev-dependencies] -pyo3-introspection = "0.26.0" \ No newline at end of file +pyo3-introspection = "0.27.0" \ No newline at end of file diff --git a/python/src/lib.rs b/python/src/lib.rs index a0df621..82b2838 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -12,11 +12,11 @@ use pyo3::types::PyIterator; // TODO: where we accept an arbitrary `ParticlesIterable` as we currently do // TODO: but also perhaps clone into some buffer so we don't have to hold the GIL for too long #[derive(Clone)] -struct ParticlesIterable<'py> { - inner: Bound<'py, PyAny>, +struct ParticlesIterable<'a, 'py> { + inner: Borrowed<'a, 'py, PyAny>, } -impl<'py> IntoIterator for ParticlesIterable<'py> { +impl<'a, 'py> IntoIterator for ParticlesIterable<'a, 'py> { type Item = [f64; 3]; type IntoIter = ParticlesIterator<'py>; @@ -45,7 +45,7 @@ impl<'py> Iterator for ParticlesIterator<'py> { // otherwise, retry with next element loop { match self.iter.next().transpose() { - Ok(Some(p)) => match <[f64; 3] as FromPyObject>::extract_bound(&p) { + Ok(Some(p)) => match <[f64; 3] as FromPyObject>::extract(p.as_borrowed()) { Ok(p) => break Some(p), Err(_) => (), }, @@ -109,12 +109,14 @@ impl PyCellGrid { /// see `CellGrid`. #[new] #[pyo3(signature = (particles: "typing.Iterable[typing.Sequence[float]] | None" = None, /, cutoff=1.0) -> "zelll.CellGrid")] - fn new<'py>(py: Python<'py>, particles: Option<&Bound<'py, PyAny>>, cutoff: f64) -> Self { + fn new<'py>(py: Python<'py>, particles: Option>, cutoff: f64) -> Self { let inner = match particles { Some(p) => { // TODO: see if we can simplify ParticlesIterable, it wraps Bound<>, so holds the GIL // TODO: would like to call ::new() detached from the GIL if that's possible? - let particles = ParticlesIterable { inner: p.clone() }; + let particles = ParticlesIterable { + inner: p.as_borrowed(), + }; CellGrid::new(particles, cutoff) } // nutpie+multithreading needs to serialize/deserialize PyCellGrid @@ -152,11 +154,11 @@ impl PyCellGrid { #[pyo3(signature = (particles: "typing.Iterable[typing.Sequence[float]]", /, cutoff=None))] fn rebuild<'py>( mut slf: PyRefMut<'_, Self>, - particles: &Bound<'py, PyAny>, + particles: Bound<'py, PyAny>, cutoff: Option, ) -> () { let particles = ParticlesIterable { - inner: particles.clone(), + inner: particles.as_borrowed(), }; slf.inner.rebuild_mut(particles, cutoff); @@ -201,9 +203,9 @@ impl PyCellGrid { #[pyo3(signature = (coordinates: "typing.Sequence[float]") -> "zelll.CellQueryIter | None")] fn query_neighbors( slf: PyRef<'_, Self>, - coordinates: &Bound<'_, PyAny>, + coordinates: Bound<'_, PyAny>, ) -> Option { - PyCellQueryIter::new(slf, coordinates) + PyCellQueryIter::new(slf, coordinates.as_borrowed()) } /// Given 3D `coordinates`, this returns a list over all particles in the neighborhood @@ -324,6 +326,12 @@ impl PyCellGridIter { // } // } +// impl Drop for PyCellQueryIter { +// fn drop(&mut self) { +// eprintln!("Dropping PyCellQueryIter"); +// } +// } + #[pymethods] impl PyCellGridIter { fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { @@ -352,8 +360,11 @@ pub struct PyCellQueryIter { } impl PyCellQueryIter { - fn new(py_cellgrid: PyRef<'_, PyCellGrid>, coordinates: &Bound<'_, PyAny>) -> Option { - let coordinates = <[f64; 3] as FromPyObject>::extract_bound(coordinates).ok()?; + fn new( + py_cellgrid: PyRef<'_, PyCellGrid>, + coordinates: Borrowed<'_, '_, PyAny>, + ) -> Option { + let coordinates = <[f64; 3] as FromPyObject>::extract(coordinates).ok()?; let iter = Box::new((&py_cellgrid).inner.query_neighbors(coordinates)?); // SAFETY: see PyCellGridIter let iter = unsafe { From 19ef3392e920ed2937cacdd84881fbf1ae5e55de Mon Sep 17 00:00:00 2001 From: fncnt Date: Mon, 15 Dec 2025 16:30:01 +0100 Subject: [PATCH 3/3] Python: adjust internal comments re: using iterators from multiple threads after some more testing --- python/src/lib.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/python/src/lib.rs b/python/src/lib.rs index 82b2838..2120c62 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -272,9 +272,7 @@ impl PyCellGrid { #[pyclass(name = "CellGridIter", module = "zelll", unsendable)] pub struct PyCellGridIter { // TODO: it looks like we probably don't need `_owner` - // TODO: also haven't yet figured out how to migrate this from IntoPy to IntoPyObject - // TODO: likely need more unsafe using Py::from_*_pointer() - // _owner: PyObject, + // _owner: Py, // TODO: `_keep_borrow` is enough to maintain correct drop order *and* prevents `PyCellGrid` // TODO: from being mutated while `PyCellGridIter` is still alive _keep_borrow: PyRef<'static, PyCellGrid>, @@ -284,9 +282,11 @@ pub struct PyCellGridIter { impl PyCellGridIter { fn new(py_cellgrid: PyRef<'_, PyCellGrid>) -> Self { // let py = py_cellgrid.py(); - // let _owner = (&py_cellgrid).into_py(py); + // let _owner = (&py_cellgrid) + // .into_py_any(py) + // .expect("could not store owner internally"); let iter = Box::new((&py_cellgrid).inner.particle_pairs()); - // SAFETY: but the idea is that `_keep_borrow` makes sure that `iter`s lifetime can be extended + // SAFETY: the idea is that `_keep_borrow` makes sure that `iter`s lifetime can be extended // SAFETY: replicating some ideas from // SAFETY: https://github.com/PyO3/pyo3/issues/1085 and // SAFETY: https://github.com/PyO3/pyo3/issues/1089