Merge pull request #173 from PyO3/einsum

Add bindings to inner, dot and einsum

Author: kngwyu

src/array.rs

@@ -64,10 +64,9 @@ use crate::slice_box::SliceBox;
 /// # Example
 /// ```
 /// # #[macro_use] extern crate ndarray;
-/// use pyo3::{GILGuard, Python};
 /// use numpy::PyArray;
 /// use ndarray::Array;
-/// Python::with_gil(|py| {
+/// pyo3::Python::with_gil(|py| {
 ///     let pyarray = PyArray::arange(py, 0., 4., 1.).reshape([2, 2]).unwrap();
 ///     let array = array![[3., 4.], [5., 6.]];
 ///     assert_eq!(
@@ -78,6 +77,8 @@ use crate::slice_box::SliceBox;
 /// ```
 pub struct PyArray<T, D>(PyAny, PhantomData<T>, PhantomData<D>);
 
+/// Zero-dimensional array.
+pub type PyArray0<T> = PyArray<T, Ix0>;
 /// One-dimensional array.
 pub type PyArray1<T> = PyArray<T, Ix1>;
 /// Two-dimensional array.
@@ -218,10 +219,9 @@ impl<T, D> PyArray<T, D> {
     ///
     /// # Example
     /// ```
-    /// use pyo3::{GILGuard, Python, Py};
     /// use numpy::PyArray1;
-    /// fn return_py_array() -> Py<PyArray1<i32>> {
-    ///    Python::with_gil(|py| PyArray1::zeros(py, [5], false).to_owned())
+    /// fn return_py_array() -> pyo3::Py<PyArray1<i32>> {
+    ///    pyo3::Python::with_gil(|py| PyArray1::zeros(py, [5], false).to_owned())
     /// }
     /// let array = return_py_array();
     /// pyo3::Python::with_gil(|py| {
@@ -594,8 +594,6 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
     }
 
     /// Get dynamic dimensioned array from fixed dimension array.
-    ///
-    /// See [get](#method.get) for usage.
     pub fn to_dyn(&self) -> &PyArray<T, IxDyn> {
         let python = self.py();
         unsafe { PyArray::from_borrowed_ptr(python, self.as_ptr()) }
@@ -710,6 +708,15 @@ impl<T: Element, D: Dimension> PyArray<T, D> {
     }
 }
 
+impl<T: Copy + Element> PyArray<T, Ix0> {
+    /// Get the element of zero-dimensional PyArray.
+    ///
+    /// See [inner](../fn.inner.html) for example.
+    pub fn item(&self) -> T {
+        unsafe { *self.data() }
+    }
+}
+
 impl<T: Element> PyArray<T, Ix1> {
     /// Construct one-dimension PyArray from slice.
     ///

src/lib.rs

@@ -43,10 +43,11 @@ 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,
-    PyArrayDyn,
+    get_array_module, PyArray, PyArray0, PyArray1, PyArray2, PyArray3, PyArray4, PyArray5,
+    PyArray6, PyArrayDyn,
 };
 pub use crate::convert::{IntoPyArray, NpyIndex, ToNpyDims, ToPyArray};
 pub use crate::dtype::{c32, c64, DataType, Element, PyArrayDescr};
@@ -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,29 @@ pub mod doc_test {
     }
     doc_comment!(include_str!("../README.md"), readme);
 }
+
+/// Create a [PyArray](./array/struct.PyArray.html) with one, two or three dimensions.
+/// This macro is backed by
+/// [`ndarray::array`](https://docs.rs/ndarray/latest/ndarray/macro.array.html).
+///
+/// # Example
+/// ```
+/// pyo3::Python::with_gil(|py| {
+///     let array = numpy::pyarray![py, [1, 2], [3, 4]];
+///     assert_eq!(
+///         array.readonly().as_array(),
+///         ndarray::array![[1, 2], [3, 4]]
+///     );
+/// });
+#[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,121 @@
+use crate::npyffi::{NPY_CASTING, NPY_ORDER};
+use crate::{Element, PyArray, PY_ARRAY_API};
+use ndarray::{Dimension, IxDyn};
+use pyo3::{AsPyPointer, FromPyPointer, PyAny, PyNativeType, PyResult};
+use std::ffi::CStr;
+
+/// Return the inner product of two arrays.
+///
+/// # Example
+/// ```
+/// pyo3::Python::with_gil(|py| {
+///     let array = numpy::pyarray![py, 1, 2, 3];
+///     let inner: &numpy::PyArray0::<_> = numpy::inner(array, array).unwrap();
+///     assert_eq!(inner.item(), 14);
+/// });
+/// ```
+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 obj = unsafe {
+        let result = PY_ARRAY_API.PyArray_InnerProduct(array1.as_ptr(), array2.as_ptr());
+        PyAny::from_owned_ptr_or_err(array1.py(), result)?
+    };
+    obj.extract()
+}
+
+/// Return the dot product of two arrays.
+///
+/// # Example
+/// ```
+/// pyo3::Python::with_gil(|py| {
+///     let a = numpy::pyarray![py, [1, 0], [0, 1]];
+///     let b = numpy::pyarray![py, [4, 1], [2, 2]];
+///     let dot: &numpy::PyArray2::<_> = numpy::dot(a, b).unwrap();
+///     assert_eq!(
+///         dot.readonly().as_array(),
+///         ndarray::array![[4, 1], [2, 2]]
+///     );
+/// });
+/// ```
+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 obj = unsafe {
+        let result = PY_ARRAY_API.PyArray_MatrixProduct(array1.as_ptr(), array2.as_ptr());
+        PyAny::from_owned_ptr_or_err(array1.py(), result)?
+    };
+    obj.extract()
+}
+
+/// Return the Einstein summation convention of given tensors.
+///
+/// We also provide the [einsum macro](./macro.einsum.html).
+pub fn einsum_impl<'py, T, DOUT>(
+    subscripts: &str,
+    arrays: &[&'py PyArray<T, IxDyn>],
+) -> PyResult<&'py PyArray<T, DOUT>>
+where
+    DOUT: Dimension,
+    T: Element,
+{
+    let subscripts: std::borrow::Cow<CStr> = if subscripts.ends_with("\0") {
+        CStr::from_bytes_with_nul(subscripts.as_bytes())
+            .unwrap()
+            .into()
+    } else {
+        std::ffi::CString::new(subscripts).unwrap().into()
+    };
+    let obj = unsafe {
+        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(),
+        );
+        PyAny::from_owned_ptr_or_err(arrays[0].py(), result)?
+    };
+    obj.extract()
+}
+
+/// Return the Einstein summation convention of given tensors.
+///
+/// For more about the Einstein summation convention, you may reffer to
+/// [the numpy document](https://numpy.org/doc/stable/reference/generated/numpy.einsum.html).
+///
+/// # Example
+/// ```
+/// pyo3::Python::with_gil(|py| {
+///     let a = numpy::PyArray::arange(py, 0, 2 * 3 * 4, 1).reshape([2, 3, 4]).unwrap();
+///     let b = numpy::pyarray![py, [20, 30], [40, 50], [60, 70]];
+///     let einsum = numpy::einsum!("ijk,ji->ik", a, b).unwrap();
+///     assert_eq!(
+///         einsum.readonly().as_array(),
+///         ndarray::array![[640,  760,  880, 1000], [2560, 2710, 2860, 3010]]
+///     );
+/// });
+/// ```
+#[macro_export]
+macro_rules! einsum {
+    ($subscripts: literal $(,$array: ident)+ $(,)*) => {{
+        let arrays = [$($array.to_dyn(),)+];
+        unsafe { $crate::einsum_impl(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],
+        );
+    })
+}