Skip to content

Commit

Permalink
array now returns a reference in the original vector instead of a copy
Browse files Browse the repository at this point in the history
  • Loading branch information
dirkvdb committed Feb 17, 2025
1 parent b1fe169 commit 700d93c
Show file tree
Hide file tree
Showing 5 changed files with 119 additions and 25 deletions.
2 changes: 1 addition & 1 deletion infra-rs
3 changes: 2 additions & 1 deletion python/geodynamix/geodynamix.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,8 @@ class raster:
@property
def array(self) -> np.ndarray:
"""
Returns the raster as a numpy array. Data is copied so changes to the array will not affect the raster.
Returns the raster as a numpy array. The data in the array is a view on the raster data so changes to the array will affect the raster.
This also means that the array cannot be changed in size, only in content.
"""
...
@property
Expand Down
76 changes: 74 additions & 2 deletions src/raster.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
use numpy::PyArrayDescr;
use pyo3::exceptions::PyBufferError;
use pyo3::ffi;
use pyo3::prelude::*;
use std::ffi::c_void;
use std::ffi::CString;
use std::ops::Add;
use std::ops::Div;
use std::ops::Mul;
Expand Down Expand Up @@ -71,8 +75,9 @@ impl Raster {
}

#[getter]
pub fn array<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
utils::raster_array(py, &self.raster)
pub fn array<'py>(slf: Bound<'py, Self>, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
//utils::raster_array(py, &slf.borrow().raster)
utils::raster_buffer_array(py, slf)
}

#[getter]
Expand Down Expand Up @@ -109,6 +114,26 @@ impl Raster {
pub fn __truediv__(&self, rhs_object: Bound<'_, PyAny>) -> PyResult<Raster> {
impl_raster_op!(div, self.raster, rhs_object)
}
/// # Safety
/// This function is unsafe because it exposes a raw pointer to the Python buffer protocol.
pub unsafe fn __getbuffer__(
slf: Bound<'_, Self>,
view: *mut ffi::Py_buffer,
flags: std::ffi::c_int,
) -> PyResult<()> {
fill_view_from_data(
view,
flags,
slf.borrow().raster.raw_data_u8_slice(),
slf.into_any(),
)
}

/// # Safety
/// This function is unsafe because it exposes a raw pointer to the Python buffer protocol.
pub unsafe fn __releasebuffer__(&self, _buffer: *mut ffi::Py_buffer) {
// No need to release, the raster still owns the data
}
}

impl From<PythonDenseArray> for Raster {
Expand Down Expand Up @@ -144,3 +169,50 @@ pub fn raster_equal(

Ok(array1 == array2)
}

/// # Safety
///
/// `view` must be a valid pointer to `ffi::Py_buffer`, or null
/// `data` must outlive the Python lifetime of `owner` (i.e. data must be owned by owner, or data
/// must be static data)
unsafe fn fill_view_from_data(
view: *mut ffi::Py_buffer,
flags: std::ffi::c_int,
data: &[u8],
owner: Bound<'_, PyAny>,
) -> PyResult<()> {
if view.is_null() {
return Err(PyBufferError::new_err("View is null"));
}

(*view).obj = owner.into_ptr();
(*view).buf = data.as_ptr() as *mut c_void;
(*view).len = data.len() as isize;
(*view).readonly = 0;
(*view).itemsize = 1;

(*view).format = if (flags & ffi::PyBUF_FORMAT) == ffi::PyBUF_FORMAT {
let msg = CString::new("B").unwrap();
msg.into_raw()
} else {
std::ptr::null_mut()
};

(*view).ndim = 1;
(*view).shape = if (flags & ffi::PyBUF_ND) == ffi::PyBUF_ND {
&mut (*view).len
} else {
std::ptr::null_mut()
};

(*view).strides = if (flags & ffi::PyBUF_STRIDES) == ffi::PyBUF_STRIDES {
&mut (*view).itemsize
} else {
std::ptr::null_mut()
};

(*view).suboffsets = std::ptr::null_mut();
(*view).internal = std::ptr::null_mut();

Ok(())
}
23 changes: 22 additions & 1 deletion src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use pyo3::{exceptions::PyValueError, prelude::*, types::IntoPyDict};

use geo::{raster::algo, Array, ArrayDataType, ArrayMetadata, ArrayNum, DenseArray};

use crate::{PythonDenseArray, RasterMetadata};
use crate::{PythonDenseArray, Raster, RasterMetadata};

fn convert_array<'py, T: ArrayNum + numpy::Element>(
py: Python<'py>,
Expand Down Expand Up @@ -47,6 +47,27 @@ pub fn raster_mask<'py>(py: Python<'py>, raster: &PythonDenseArray) -> PyResult<
}
}

pub fn raster_buffer_array<'py>(
py: Python<'py>,
raster: Bound<'py, Raster>,
) -> PyResult<Bound<'py, PyAny>> {
let np = PyModule::import(py, "numpy")?;

let shape = [
raster.borrow().raster.rows().count(),
raster.borrow().raster.columns().count(),
];
let dtype =
array_type_to_numpy_dtype(py, raster.borrow().raster.data_type()).into_pyobject(py)?;
let kwargs = vec![
("dtype", dtype.into_pyobject(py)?.into_any()),
("buffer", raster.into_pyobject(py)?.into_any()),
];

np.getattr("ndarray")?
.call((shape.into_pyobject(py)?,), Some(&kwargs.into_py_dict(py)?))
}

pub fn raster_masked_array<'py>(
py: Python<'py>,
raster: &PythonDenseArray,
Expand Down
40 changes: 20 additions & 20 deletions tests/testgdx.py
Original file line number Diff line number Diff line change
Expand Up @@ -182,31 +182,31 @@ def test_write_map(self):
os.remove("writtenraster.asc")
self.assertTrue(np.allclose(raster, written_raster.array, equal_nan=True))

# def test_modify_raster_using_numpy(self):
# # modify the raster data using the ndarry accessor
# # verify that the internal raster data has changed by writing it to disk
# # and comparing the result
def test_modify_raster_using_numpy(self):
# modify the raster data using the ndarry accessor
# verify that the internal raster data has changed by writing it to disk
# and comparing the result

# create_test_file("raster.asc", test_raster)
# raster = gdx.read("raster.asc")
create_test_file("raster.asc", test_raster)
raster = gdx.read("raster.asc")

# raster.array.fill(44)
raster.array.fill(44)

# gdx.write(raster, "writtenraster.asc")
# written_raster = gdx.read("writtenraster.asc")
# os.remove("writtenraster.asc")
gdx.write(raster, "writtenraster.asc")
written_raster = gdx.read("writtenraster.asc")
os.remove("writtenraster.asc")

# expected = np.array(
# [
# [44, 44, 44, 44, 44],
# [44, 44, 44, 44, 44],
# [44, 44, 44, 44, 44],
# [44, 44, 44, 44, 44],
# ],
# dtype="f",
# )
expected = np.array(
[
[44, 44, 44, 44, 44],
[44, 44, 44, 44, 44],
[44, 44, 44, 44, 44],
[44, 44, 44, 44, 44],
],
dtype="f",
)

# self.assertTrue(np.allclose(expected, written_raster.array, equal_nan=True))
self.assertTrue(np.allclose(expected, written_raster.array, equal_nan=True))

def test_replace_value(self):
expected = np.array(
Expand Down

0 comments on commit 700d93c

Please sign in to comment.