Skip to content

Commit 6e47b97

Browse files
committed
Optimize copying nalgebra matrices into NumPy arrays.
1 parent 7c07bff commit 6e47b97

File tree

2 files changed

+18
-4
lines changed

2 files changed

+18
-4
lines changed

src/convert.rs

+8-4
Original file line numberDiff line numberDiff line change
@@ -183,10 +183,14 @@ where
183183

184184
fn to_pyarray<'py>(&self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
185185
unsafe {
186-
let array = PyArray::new(py, (self.nrows(), self.ncols()), false);
187-
for r in 0..self.nrows() {
188-
for c in 0..self.ncols() {
189-
*array.uget_mut((r, c)) = self.get_unchecked((r, c)).clone();
186+
let array = PyArray::<N, _>::new(py, (self.nrows(), self.ncols()), true);
187+
let mut data_ptr = array.data();
188+
if self.data.is_contiguous() {
189+
ptr::copy_nonoverlapping(self.data.ptr(), data_ptr, self.len());
190+
} else {
191+
for item in self.iter() {
192+
data_ptr.write(item.clone());
193+
data_ptr = data_ptr.add(1);
190194
}
191195
}
192196
array

tests/to_py.rs

+10
Original file line numberDiff line numberDiff line change
@@ -303,6 +303,7 @@ fn slice_container_type_confusion() {
303303
#[test]
304304
fn matrix_to_numpy() {
305305
let matrix = nalgebra::Matrix3::<i32>::new(0, 1, 2, 3, 4, 5, 6, 7, 8);
306+
assert!(nalgebra::storage::RawStorage::is_contiguous(&matrix.data));
306307

307308
Python::with_gil(|py| {
308309
let array = matrix.to_pyarray(py);
@@ -312,4 +313,13 @@ fn matrix_to_numpy() {
312313
array![[0, 1, 2], [3, 4, 5], [6, 7, 8]],
313314
);
314315
});
316+
317+
let matrix = matrix.row(0);
318+
assert!(!nalgebra::storage::RawStorage::is_contiguous(&matrix.data));
319+
320+
Python::with_gil(|py| {
321+
let array = matrix.to_pyarray(py);
322+
323+
assert_eq!(array.readonly().as_array(), array![[0, 1, 2]]);
324+
});
315325
}

0 commit comments

Comments
 (0)