--- a/src/python.rs Fri Nov 28 14:46:35 2025 -0500 +++ b/src/python.rs Thu Dec 04 14:36:22 2025 -0500 @@ -50,7 +50,20 @@ }*/ fn __iter__(slf: Bound<'_, Self>) -> PyResult<Py<$iter>> { - Py::new(slf.py(), $iter { measure_ref: slf.unbind(), next: 0 }) + Py::new(slf.py(), $iter { + measure_ref: slf.unbind(), + next: 0, + pad: false, + }) + } + + /// Same as iterating the object, but pads (or cuts) location vectors to 3 dimensions. + fn iter_padded(slf: Bound<'_, Self>) -> PyResult<Py<$iter>> { + Py::new(slf.py(), $iter { + measure_ref: slf.unbind(), + next: 0, + pad: true, + }) } } @@ -61,6 +74,7 @@ pub struct $iter { measure_ref: Py<$name>, next: usize, + pad: bool, } #[pymethods] @@ -77,11 +91,22 @@ let py = slf.py(); let meas_: PyRef<'_, $name> = slf.measure_ref.extract(py)?; let meas = &(meas_.inner); + let pad = slf.pad; let next = &mut slf.next; Ok((*next < meas.len()).then(|| { let δ = meas[*next]; *next += 1; - (PyArray::from_slice(py, &δ.x.0), δ.α) + if pad { + ( + PyArray::from_iter( + py, + δ.x.iter().copied().chain(std::iter::repeat(0.0)).take(3), + ), + δ.α, + ) + } else { + (PyArray::from_slice(py, &δ.x.0), δ.α) + } })) } }