Repository: rust-ndarray/ndarray (rust)

Issue:
Implementation of numpy's meshgrid function for ndarray
I have made an attempt to try to replicate the core functionality of numpy's  [meshgrid](https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html) function, using ndarray.
My attempt, which seems to be yielding identical results to the numpy version is hosted [here](https://github.com/jreniel/meshgridrs).

My hope is that this function, or an equivalent, makes it to ndarray, since I feel like `meshgrid` is one of those core numpy functions that we tend to use repeatedly in my field, and it would be an invaluable addition to the ndarray crate.

To be clear, I'm asking to not necessarily use the same function that I've posted here, I'll be happy too if the ndarray devs implement a better solution. All I would like to see is to have this function be a part of the ndarray collection of functions.

Pull request:
Meshgrid implementation
Closes #1355, but depends first on #1440. Could have its dependence on #1440 removed if needed faster.

Patch:
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 6ebdc843..13fb9d0d 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -1,5 +1,7 @@
 on:
   pull_request:
+    paths-ignore:
+      - '.github/workflows/latest-deps.yaml'
   merge_group:
   push:
     branches:
diff --git a/src/free_functions.rs b/src/free_functions.rs
index c1889cec..a2ad6137 100644
--- a/src/free_functions.rs
+++ b/src/free_functions.rs
@@ -9,6 +9,7 @@
 use alloc::vec;
 #[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
+use meshgrid_impl::Meshgrid;
 #[allow(unused_imports)]
 use std::compile_error;
 use std::mem::{forget, size_of};
@@ -45,6 +46,8 @@ use crate::{imp_prelude::*, LayoutRef};
 ///
 /// This macro uses `vec![]`, and has the same ownership semantics;
 /// elements are moved into the resulting `Array`.
+/// If running with `no_std`, this may require that you `use alloc::vec`
+/// before being able to use the `array!` macro.
 ///
 /// Use `array![...].into_shared()` to create an `ArcArray`.
 ///
@@ -336,3 +339,408 @@ pub fn rcarr3<A: Clone, const N: usize, const M: usize>(xs: &[[[A; M]; N]]) -> A
 {
     arr3(xs).into_shared()
 }
+
+/// The indexing order for [`meshgrid`]; see there for more details.
+///
+/// Controls whether the first argument to `meshgrid` will fill the rows or columns of the outputs.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum MeshIndex
+{
+    /// Cartesian indexing.
+    ///
+    /// The first argument of `meshgrid` will repeat over the columns of the output.
+    ///
+    /// Note: this is the default in `numpy`.
+    XY,
+    /// Matrix indexing.
+    ///
+    /// The first argument of `meshgrid` will repeat over the rows of the output.
+    IJ,
+}
+
+mod meshgrid_impl
+{
+    use super::MeshIndex;
+    use crate::extension::nonnull::nonnull_debug_checked_from_ptr;
+    use crate::{
+        ArrayBase,
+        ArrayRef1,
+        ArrayView,
+        ArrayView2,
+        ArrayView3,
+        ArrayView4,
+        ArrayView5,
+        ArrayView6,
+        Axis,
+        Data,
+        Dim,
+        IntoDimension,
+        Ix1,
+        LayoutRef1,
+    };
+
+    /// Construct the correct strides for the `idx`-th entry into meshgrid
+    fn construct_strides<A, const N: usize>(
+        arr: &LayoutRef1<A>, idx: usize, indexing: MeshIndex,
+    ) -> <[usize; N] as IntoDimension>::Dim
+    where [usize; N]: IntoDimension
+    {
+        let mut ret = [0; N];
+        if idx < 2 && indexing == MeshIndex::XY {
+            ret[1 - idx] = arr.stride_of(Axis(0)) as usize;
+        } else {
+            ret[idx] = arr.stride_of(Axis(0)) as usize;
+        }
+        Dim(ret)
+    }
+
+    /// Construct the correct shape for the `idx`-th entry into meshgrid
+    fn construct_shape<A, const N: usize>(
+        arrays: [&LayoutRef1<A>; N], indexing: MeshIndex,
+    ) -> <[usize; N] as IntoDimension>::Dim
+    where [usize; N]: IntoDimension
+    {
+        let mut ret = arrays.map(|a| a.len());
+        if indexing == MeshIndex::XY {
+            ret.swap(0, 1);
+        }
+        Dim(ret)
+    }
+
+    /// A trait to encapsulate static dispatch for [`meshgrid`](super::meshgrid); see there for more details.
+    ///
+    /// The inputs should always be some sort of 1D array.
+    /// The outputs should always be ND arrays where N is the number of inputs.
+    ///
+    /// Where possible, this trait tries to return array views rather than allocating additional memory.
+    pub trait Meshgrid
+    {
+        type Output;
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output;
+    }
+
+    macro_rules! meshgrid_body {
+        ($count:literal, $indexing:expr, $(($arr:expr, $idx:literal)),+) => {
+        {
+            let shape = construct_shape([$($arr),+], $indexing);
+            (
+                $({
+                    let strides = construct_strides::<_, $count>($arr, $idx, $indexing);
+                    unsafe { ArrayView::new(nonnull_debug_checked_from_ptr($arr.as_ptr() as *mut A), shape, strides) }
+                }),+
+            )
+        }
+    };
+    }
+
+    impl<'a, 'b, A> Meshgrid for (&'a ArrayRef1<A>, &'b ArrayRef1<A>)
+    {
+        type Output = (ArrayView2<'a, A>, ArrayView2<'b, A>);
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
+        {
+            meshgrid_body!(2, indexing, (arrays.0, 0), (arrays.1, 1))
+        }
+    }
+
+    impl<'a, 'b, S1, S2, A: 'b + 'a> Meshgrid for (&'a ArrayBase<S1, Ix1>, &'b ArrayBase<S2, Ix1>)
+    where
+        S1: Data<Elem = A>,
+        S2: Data<Elem = A>,
+    {
+        type Output = (ArrayView2<'a, A>, ArrayView2<'b, A>);
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
+        {
+            Meshgrid::meshgrid((&**arrays.0, &**arrays.1), indexing)
+        }
+    }
+
+    impl<'a, 'b, 'c, A> Meshgrid for (&'a ArrayRef1<A>, &'b ArrayRef1<A>, &'c ArrayRef1<A>)
+    {
+        type Output = (ArrayView3<'a, A>, ArrayView3<'b, A>, ArrayView3<'c, A>);
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
+        {
+            meshgrid_body!(3, indexing, (arrays.0, 0), (arrays.1, 1), (arrays.2, 2))
+        }
+    }
+
+    impl<'a, 'b, 'c, S1, S2, S3, A: 'b + 'a + 'c> Meshgrid
+        for (&'a ArrayBase<S1, Ix1>, &'b ArrayBase<S2, Ix1>, &'c ArrayBase<S3, Ix1>)
+    where
+        S1: Data<Elem = A>,
+        S2: Data<Elem = A>,
+        S3: Data<Elem = A>,
+    {
+        type Output = (ArrayView3<'a, A>, ArrayView3<'b, A>, ArrayView3<'c, A>);
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
+        {
+            Meshgrid::meshgrid((&**arrays.0, &**arrays.1, &**arrays.2), indexing)
+        }
+    }
+
+    impl<'a, 'b, 'c, 'd, A> Meshgrid for (&'a ArrayRef1<A>, &'b ArrayRef1<A>, &'c ArrayRef1<A>, &'d ArrayRef1<A>)
+    {
+        type Output = (ArrayView4<'a, A>, ArrayView4<'b, A>, ArrayView4<'c, A>, ArrayView4<'d, A>);
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
+        {
+            meshgrid_body!(4, indexing, (arrays.0, 0), (arrays.1, 1), (arrays.2, 2), (arrays.3, 3))
+        }
+    }
+
+    impl<'a, 'b, 'c, 'd, S1, S2, S3, S4, A: 'a + 'b + 'c + 'd> Meshgrid
+        for (&'a ArrayBase<S1, Ix1>, &'b ArrayBase<S2, Ix1>, &'c ArrayBase<S3, Ix1>, &'d ArrayBase<S4, Ix1>)
+    where
+        S1: Data<Elem = A>,
+        S2: Data<Elem = A>,
+        S3: Data<Elem = A>,
+        S4: Data<Elem = A>,
+    {
+        type Output = (ArrayView4<'a, A>, ArrayView4<'b, A>, ArrayView4<'c, A>, ArrayView4<'d, A>);
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
+        {
+            Meshgrid::meshgrid((&**arrays.0, &**arrays.1, &**arrays.2, &**arrays.3), indexing)
+        }
+    }
+
+    impl<'a, 'b, 'c, 'd, 'e, A> Meshgrid
+        for (&'a ArrayRef1<A>, &'b ArrayRef1<A>, &'c ArrayRef1<A>, &'d ArrayRef1<A>, &'e ArrayRef1<A>)
+    {
+        type Output = (ArrayView5<'a, A>, ArrayView5<'b, A>, ArrayView5<'c, A>, ArrayView5<'d, A>, ArrayView5<'e, A>);
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
+        {
+            meshgrid_body!(5, indexing, (arrays.0, 0), (arrays.1, 1), (arrays.2, 2), (arrays.3, 3), (arrays.4, 4))
+        }
+    }
+
+    impl<'a, 'b, 'c, 'd, 'e, S1, S2, S3, S4, S5, A: 'a + 'b + 'c + 'd + 'e> Meshgrid
+        for (
+            &'a ArrayBase<S1, Ix1>,
+            &'b ArrayBase<S2, Ix1>,
+            &'c ArrayBase<S3, Ix1>,
+            &'d ArrayBase<S4, Ix1>,
+            &'e ArrayBase<S5, Ix1>,
+        )
+    where
+        S1: Data<Elem = A>,
+        S2: Data<Elem = A>,
+        S3: Data<Elem = A>,
+        S4: Data<Elem = A>,
+        S5: Data<Elem = A>,
+    {
+        type Output = (ArrayView5<'a, A>, ArrayView5<'b, A>, ArrayView5<'c, A>, ArrayView5<'d, A>, ArrayView5<'e, A>);
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
+        {
+            Meshgrid::meshgrid((&**arrays.0, &**arrays.1, &**arrays.2, &**arrays.3, &**arrays.4), indexing)
+        }
+    }
+
+    impl<'a, 'b, 'c, 'd, 'e, 'f, A> Meshgrid
+        for (
+            &'a ArrayRef1<A>,
+            &'b ArrayRef1<A>,
+            &'c ArrayRef1<A>,
+            &'d ArrayRef1<A>,
+            &'e ArrayRef1<A>,
+            &'f ArrayRef1<A>,
+        )
+    {
+        type Output = (
+            ArrayView6<'a, A>,
+            ArrayView6<'b, A>,
+            ArrayView6<'c, A>,
+            ArrayView6<'d, A>,
+            ArrayView6<'e, A>,
+            ArrayView6<'f, A>,
+        );
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
+        {
+            meshgrid_body!(6, indexing, (arrays.0, 0), (arrays.1, 1), (arrays.2, 2), (arrays.3, 3), (arrays.4, 4), (arrays.5, 5))
+        }
+    }
+
+    impl<'a, 'b, 'c, 'd, 'e, 'f, S1, S2, S3, S4, S5, S6, A: 'a + 'b + 'c + 'd + 'e + 'f> Meshgrid
+        for (
+            &'a ArrayBase<S1, Ix1>,
+            &'b ArrayBase<S2, Ix1>,
+            &'c ArrayBase<S3, Ix1>,
+            &'d ArrayBase<S4, Ix1>,
+            &'e ArrayBase<S5, Ix1>,
+            &'f ArrayBase<S6, Ix1>,
+        )
+    where
+        S1: Data<Elem = A>,
+        S2: Data<Elem = A>,
+        S3: Data<Elem = A>,
+        S4: Data<Elem = A>,
+        S5: Data<Elem = A>,
+        S6: Data<Elem = A>,
+    {
+        type Output = (
+            ArrayView6<'a, A>,
+            ArrayView6<'b, A>,
+            ArrayView6<'c, A>,
+            ArrayView6<'d, A>,
+            ArrayView6<'e, A>,
+            ArrayView6<'f, A>,
+        );
+
+        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
+        {
+            Meshgrid::meshgrid((&**arrays.0, &**arrays.1, &**arrays.2, &**arrays.3, &**arrays.4, &**arrays.5), indexing)
+        }
+    }
+}
+
+/// Create coordinate matrices from coordinate vectors.
+///
+/// Given an N-tuple of 1D coordinate vectors, return an N-tuple of ND coordinate arrays.
+/// This is particularly useful for computing the outputs of functions with N arguments over
+/// regularly spaced grids.
+///
+/// The `indexing` argument can be controlled by [`MeshIndex`] to support both Cartesian and
+/// matrix indexing. In the two-dimensional case, inputs of length `N` and `M` will create
+/// output arrays of size `(M, N)` when using [`MeshIndex::XY`] and size `(N, M)` when using
+/// [`MeshIndex::IJ`].
+///
+/// # Example
+/// ```
+/// use ndarray::{array, meshgrid, MeshIndex};
+///
+/// let arr1 = array![1, 2];
+/// let arr2 = array![3, 4];
+/// let arr3 = array![5, 6];
+///
+/// // Cartesian indexing
+/// let (res1, res2) = meshgrid((&arr1, &arr2), MeshIndex::XY);
+/// assert_eq!(res1, array![
+///     [1, 2],
+///     [1, 2],
+/// ]);
+/// assert_eq!(res2, array![
+///     [3, 3],
+///     [4, 4],
+/// ]);
+///
+/// // Matrix indexing
+/// let (res1, res2) = meshgrid((&arr1, &arr2), MeshIndex::IJ);
+/// assert_eq!(res1, array![
+///     [1, 1],
+///     [2, 2],
+/// ]);
+/// assert_eq!(res2, array![
+///     [3, 4],
+///     [3, 4],
+/// ]);
+///
+/// let (_, _, res3) = meshgrid((&arr1, &arr2, &arr3), MeshIndex::XY);
+/// assert_eq!(res3, array![
+///     [[5, 6],
+///      [5, 6]],
+///     [[5, 6],
+///      [5, 6]],
+/// ]);
+/// ```
+pub fn meshgrid<T: Meshgrid>(arrays: T, indexing: MeshIndex) -> T::Output
+{
+    Meshgrid::meshgrid(arrays, indexing)
+}
+
+#[cfg(test)]
+mod tests
+{
+    use super::s;
+    use crate::{meshgrid, Axis, MeshIndex};
+    #[cfg(not(feature = "std"))]
+    use alloc::vec;
+
+    #[test]
+    fn test_meshgrid2()
+    {
+        let x = array![1, 2, 3];
+        let y = array![4, 5, 6, 7];
+        let (xx, yy) = meshgrid((&x, &y), MeshIndex::XY);
+        assert_eq!(xx, array![[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]);
+        assert_eq!(yy, array![[4, 4, 4], [5, 5, 5], [6, 6, 6], [7, 7, 7]]);
+
+        let (xx, yy) = meshgrid((&x, &y), MeshIndex::IJ);
+        assert_eq!(xx, array![[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]]);
+        assert_eq!(yy, array![[4, 5, 6, 7], [4, 5, 6, 7], [4, 5, 6, 7]]);
+    }
+
+    #[test]
+    fn test_meshgrid3()
+    {
+        let x = array![1, 2, 3];
+        let y = array![4, 5, 6, 7];
+        let z = array![-1, -2];
+        let (xx, yy, zz) = meshgrid((&x, &y, &z), MeshIndex::XY);
+        assert_eq!(xx, array![
+            [[1, 1], [2, 2], [3, 3]],
+            [[1, 1], [2, 2], [3, 3]],
+            [[1, 1], [2, 2], [3, 3]],
+            [[1, 1], [2, 2], [3, 3]],
+        ]);
+        assert_eq!(yy, array![
+            [[4, 4], [4, 4], [4, 4]],
+            [[5, 5], [5, 5], [5, 5]],
+            [[6, 6], [6, 6], [6, 6]],
+            [[7, 7], [7, 7], [7, 7]],
+        ]);
+        assert_eq!(zz, array![
+            [[-1, -2], [-1, -2], [-1, -2]],
+            [[-1, -2], [-1, -2], [-1, -2]],
+            [[-1, -2], [-1, -2], [-1, -2]],
+            [[-1, -2], [-1, -2], [-1, -2]],
+        ]);
+
+        let (xx, yy, zz) = meshgrid((&x, &y, &z), MeshIndex::IJ);
+        assert_eq!(xx, array![
+            [[1, 1], [1, 1], [1, 1], [1, 1]],
+            [[2, 2], [2, 2], [2, 2], [2, 2]],
+            [[3, 3], [3, 3], [3, 3], [3, 3]],
+        ]);
+        assert_eq!(yy, array![
+            [[4, 4], [5, 5], [6, 6], [7, 7]],
+            [[4, 4], [5, 5], [6, 6], [7, 7]],
+            [[4, 4], [5, 5], [6, 6], [7, 7]],
+        ]);
+        assert_eq!(zz, array![
+            [[-1, -2], [-1, -2], [-1, -2], [-1, -2]],
+            [[-1, -2], [-1, -2], [-1, -2], [-1, -2]],
+            [[-1, -2], [-1, -2], [-1, -2], [-1, -2]],
+        ]);
+    }
+
+    #[test]
+    fn test_meshgrid_from_offset()
+    {
+        let x = array![1, 2, 3];
+        let x = x.slice(s![1..]);
+        let y = array![4, 5, 6];
+        let y = y.slice(s![1..]);
+        let (xx, yy) = meshgrid((&x, &y), MeshIndex::XY);
+        assert_eq!(xx, array![[2, 3], [2, 3]]);
+        assert_eq!(yy, array![[5, 5], [6, 6]]);
+    }
+
+    #[test]
+    fn test_meshgrid_neg_stride()
+    {
+        let x = array![1, 2, 3];
+        let x = x.slice(s![..;-1]);
+        assert!(x.stride_of(Axis(0)) < 0); // Setup for test
+        let y = array![4, 5, 6];
+        let (xx, yy) = meshgrid((&x, &y), MeshIndex::XY);
+        assert_eq!(xx, array![[3, 2, 1], [3, 2, 1], [3, 2, 1]]);
+        assert_eq!(yy, array![[4, 4, 4], [5, 5, 5], [6, 6, 6]]);
+    }
+}


Test patch:
diff --git a/.github/workflows/latest-deps.yaml b/.github/workflows/latest-deps.yaml
index f2f3d848..3b28169e 100644
--- a/.github/workflows/latest-deps.yaml
+++ b/.github/workflows/latest-deps.yaml
@@ -8,6 +8,9 @@ on:
     # Sorry if this ruins your weekend, future maintainer...
     - cron: '0 12 * * FRI'
   workflow_dispatch: # For running manually
+  pull_request:
+    paths:
+      - '.github/workflows/latest-deps.yaml'
 
 env:
   CARGO_TERM_COLOR: always
@@ -41,7 +44,7 @@ jobs:
 
   latest_deps_msrv:
     runs-on: ubuntu-latest
-    name: Check Latest Dependencies on MSRV (${{ env.MSRV }})
+    name: Check Latest Dependencies on MSRV
     steps:
       - name: Check Out Repo
         uses: actions/checkout@v4