inner, dot, einsum

Author: kngwyu

src/lib.rs

@@ -43,6 +43,7 @@ pub mod npyffi;
 pub mod npyiter;
 mod readonly;
 mod slice_box;
+mod sum_products;
 
 pub use crate::array::{
     get_array_module, PyArray, PyArray1, PyArray2, PyArray3, PyArray4, PyArray5, PyArray6,
@@ -59,7 +60,8 @@ pub use crate::readonly::{
     PyReadonlyArray, PyReadonlyArray1, PyReadonlyArray2, PyReadonlyArray3, PyReadonlyArray4,
     PyReadonlyArray5, PyReadonlyArray6, PyReadonlyArrayDyn,
 };
-pub use ndarray::{Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
+pub use crate::sum_products::{dot, einsum_impl, inner};
+pub use ndarray::{array, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
 
 /// Test readme
 #[doc(hidden)]
@@ -72,3 +74,16 @@ pub mod doc_test {
     }
     doc_comment!(include_str!("../README.md"), readme);
 }
+
+#[macro_export]
+macro_rules! pyarray {
+    ($py: ident, $([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*) => {{
+        $crate::IntoPyArray::into_pyarray($crate::array![$([$([$($x,)*],)*],)*], $py)
+    }};
+    ($py: ident, $([$($x:expr),* $(,)*]),+ $(,)*) => {{
+        $crate::IntoPyArray::into_pyarray($crate::array![$([$($x,)*],)*], $py)
+    }};
+    ($py: ident, $($x:expr),* $(,)*) => {{
+        $crate::IntoPyArray::into_pyarray($crate::array![$($x,)*], $py)
+    }};
+}

src/npyffi/array.rs

@@ -286,7 +286,7 @@ impl PyArrayAPI {
     impl_api![273; PyArray_ResultType(narrs: npy_intp, arr: *mut *mut PyArrayObject, ndtypes: npy_intp, dtypes: *mut *mut PyArray_Descr) -> *mut PyArray_Descr];
     impl_api![274; PyArray_CanCastArrayTo(arr: *mut PyArrayObject, to: *mut PyArray_Descr, casting: NPY_CASTING) -> npy_bool];
     impl_api![275; PyArray_CanCastTypeTo(from: *mut PyArray_Descr, to: *mut PyArray_Descr, casting: NPY_CASTING) -> npy_bool];
-    impl_api![276; PyArray_EinsteinSum(subscripts: *mut c_char, nop: npy_intp, op_in: *mut *mut PyArrayObject, dtype: *mut PyArray_Descr, order: NPY_ORDER, casting: NPY_CASTING, out: *mut PyArrayObject) -> *mut PyArrayObject];
+    impl_api![276; PyArray_EinsteinSum(subscripts: *mut c_char, nop: npy_intp, op_in: *mut *mut PyArrayObject, dtype: *mut PyArray_Descr, order: NPY_ORDER, casting: NPY_CASTING, out: *mut PyArrayObject) -> *mut PyObject];
     impl_api![277; PyArray_NewLikeArray(prototype: *mut PyArrayObject, order: NPY_ORDER, dtype: *mut PyArray_Descr, subok: c_int) -> *mut PyObject];
     impl_api![278; PyArray_GetArrayParamsFromObject(op: *mut PyObject, requested_dtype: *mut PyArray_Descr, writeable: npy_bool, out_dtype: *mut *mut PyArray_Descr, out_ndim: *mut c_int, out_dims: *mut npy_intp, out_arr: *mut *mut PyArrayObject, context: *mut PyObject) -> c_int];
     impl_api![279; PyArray_ConvertClipmodeSequence(object: *mut PyObject, modes: *mut NPY_CLIPMODE, n: c_int) -> c_int];

src/sum_products.rs

@@ -0,0 +1,68 @@
+use crate::npyffi::{PyArrayObject, NPY_CASTING, NPY_ORDER};
+use crate::{Element, PyArray, PY_ARRAY_API};
+use ndarray::Dimension;
+use pyo3::{AsPyPointer, FromPyPointer, PyAny, PyNativeType, PyResult};
+
+/// Return the inner product of two arrays.
+pub fn inner<'py, T, DIN1, DIN2, DOUT>(
+    array1: &'py PyArray<T, DIN1>,
+    array2: &'py PyArray<T, DIN2>,
+) -> PyResult<&'py PyArray<T, DOUT>>
+where
+    DIN1: Dimension,
+    DIN2: Dimension,
+    DOUT: Dimension,
+    T: Element,
+{
+    let result = unsafe { PY_ARRAY_API.PyArray_InnerProduct(array1.as_ptr(), array2.as_ptr()) };
+    let obj = unsafe { PyAny::from_owned_ptr_or_err(array1.py(), result)? };
+    obj.extract()
+}
+
+/// Return the dot product of two arrays.
+pub fn dot<'py, T, DIN1, DIN2, DOUT>(
+    array1: &'py PyArray<T, DIN1>,
+    array2: &'py PyArray<T, DIN2>,
+) -> PyResult<&'py PyArray<T, DOUT>>
+where
+    DIN1: Dimension,
+    DIN2: Dimension,
+    DOUT: Dimension,
+    T: Element,
+{
+    let result = unsafe { PY_ARRAY_API.PyArray_MatrixProduct(array1.as_ptr(), array2.as_ptr()) };
+    let obj = unsafe { PyAny::from_owned_ptr_or_err(array1.py(), result)? };
+    obj.extract()
+}
+
+pub unsafe fn einsum_impl<'py, T, DIN, DOUT>(
+    dummy_array: &'py PyArray<T, DIN>,
+    subscripts: &str,
+    arrays: &[*mut PyArrayObject],
+) -> PyResult<&'py PyArray<T, DOUT>>
+where
+    DIN: Dimension,
+    DOUT: Dimension,
+    T: Element,
+{
+    let subscripts = std::ffi::CStr::from_bytes_with_nul(subscripts.as_bytes()).unwrap();
+    let result = PY_ARRAY_API.PyArray_EinsteinSum(
+        subscripts.as_ptr() as _,
+        arrays.len() as _,
+        arrays.as_ptr() as _,
+        std::ptr::null_mut(),
+        NPY_ORDER::NPY_KEEPORDER,
+        NPY_CASTING::NPY_NO_CASTING,
+        std::ptr::null_mut(),
+    );
+    let obj = PyAny::from_owned_ptr_or_err(dummy_array.py(), result)?;
+    obj.extract()
+}
+
+#[macro_export]
+macro_rules! einsum {
+    ($subscripts: literal, $first_array: ident $(,$array: ident)* $(,)*) => {{
+        let arrays = [$first_array.as_array_ptr(), $($array.as_array_ptr(),)*];
+        unsafe { $crate::einsum_impl($first_array, concat!($subscripts, "\0"), &arrays) }
+    }};
+}

tests/sum_products.rs

@@ -0,0 +1,64 @@
+use numpy::{array, dot, einsum, inner, pyarray, PyArray1, PyArray2};
+
+#[test]
+fn test_dot() {
+    pyo3::Python::with_gil(|py| {
+        let a = pyarray![py, [1, 0], [0, 1]];
+        let b = pyarray![py, [4, 1], [2, 2]];
+        let c = dot(a, b).unwrap();
+        assert_eq!(c.readonly().as_array(), array![[4, 1], [2, 2]]);
+        let a = pyarray![py, 1, 2, 3];
+        let err: pyo3::PyResult<&PyArray2<_>> = dot(a, b);
+        let err = err.unwrap_err();
+        assert!(err.to_string().contains("not aligned"), "{}", err);
+    })
+}
+
+#[test]
+fn test_inner() {
+    pyo3::Python::with_gil(|py| {
+        let a = pyarray![py, 1, 2, 3];
+        let b = pyarray![py, 0, 1, 0];
+        let c = inner(a, b).unwrap();
+        assert_eq!(c.readonly().as_array(), ndarray::arr0(2));
+        let a = pyarray![py, [1, 0], [0, 1]];
+        let b = pyarray![py, [4, 1], [2, 2]];
+        let c = inner(a, b).unwrap();
+        assert_eq!(c.readonly().as_array(), array![[4, 2], [1, 2]]);
+        let a = pyarray![py, 1, 2, 3];
+        let err: pyo3::PyResult<&PyArray2<_>> = inner(a, b);
+        let err = err.unwrap_err();
+        assert!(err.to_string().contains("not aligned"), "{}", err);
+    })
+}
+
+#[test]
+fn test_einsum() {
+    pyo3::Python::with_gil(|py| {
+        let a = PyArray1::<i32>::arange(py, 0, 25, 1)
+            .reshape([5, 5])
+            .unwrap();
+        let b = pyarray![py, 0, 1, 2, 3, 4];
+        let c = pyarray![py, [0, 1, 2], [3, 4, 5]];
+        assert_eq!(
+            einsum!("ii", a).unwrap().readonly().as_array(),
+            ndarray::arr0(60)
+        );
+        assert_eq!(
+            einsum!("ii->i", a).unwrap().readonly().as_array(),
+            array![0, 6, 12, 18, 24],
+        );
+        assert_eq!(
+            einsum!("ij->i", a).unwrap().readonly().as_array(),
+            array![10, 35, 60, 85, 110],
+        );
+        assert_eq!(
+            einsum!("ji", c).unwrap().readonly().as_array(),
+            array![[0, 3], [1, 4], [2, 5]],
+        );
+        assert_eq!(
+            einsum!("ij,j", a, b).unwrap().readonly().as_array(),
+            array![30, 80, 130, 180, 230],
+        );
+    })
+}