ndarray/
free_functions.rs

1// Copyright 2014-2016 bluss and ndarray developers.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9use alloc::vec;
10#[cfg(not(feature = "std"))]
11use alloc::vec::Vec;
12use meshgrid_impl::Meshgrid;
13#[allow(unused_imports)]
14use std::compile_error;
15use std::mem::{forget, size_of};
16use std::ptr::NonNull;
17
18use crate::{dimension, ArcArray1, ArcArray2};
19use crate::{imp_prelude::*, ArrayPartsSized};
20
21/// Create an **[`Array`]** with one, two, three, four, five, or six dimensions.
22///
23/// ```
24/// use ndarray::array;
25/// let a1 = array![1, 2, 3, 4];
26///
27/// let a2 = array![[1, 2],
28///                 [3, 4]];
29///
30/// let a3 = array![[[1, 2], [3, 4]],
31///                 [[5, 6], [7, 8]]];
32///
33/// let a4 = array![[[[1, 2, 3, 4]]]];
34///
35/// let a5 = array![[[[[1, 2, 3, 4, 5]]]]];
36///
37/// let a6 = array![[[[[[1, 2, 3, 4, 5, 6]]]]]];
38///
39/// assert_eq!(a1.shape(), &[4]);
40/// assert_eq!(a2.shape(), &[2, 2]);
41/// assert_eq!(a3.shape(), &[2, 2, 2]);
42/// assert_eq!(a4.shape(), &[1, 1, 1, 4]);
43/// assert_eq!(a5.shape(), &[1, 1, 1, 1, 5]);
44/// assert_eq!(a6.shape(), &[1, 1, 1, 1, 1, 6]);
45/// ```
46///
47/// This macro uses `vec![]`, and has the same ownership semantics;
48/// elements are moved into the resulting `Array`.
49/// If running with `no_std`, this may require that you `use alloc::vec`
50/// before being able to use the `array!` macro.
51///
52/// Use `array![...].into_shared()` to create an `ArcArray`.
53///
54/// Attempts to crate 7D+ arrays with this macro will lead to
55/// a compiler error, since the difference between a 7D array
56/// of i32 and a 6D array of `[i32; 3]` is ambiguous. Higher-dim
57/// arrays can be created with [`ArrayD`].
58///
59/// ```compile_fail
60/// use ndarray::array;
61/// let a7 = array![[[[[[[1, 2, 3]]]]]]];
62/// // error: Arrays of 7 dimensions or more (or ndarrays of Rust arrays) cannot be constructed with the array! macro.
63/// ```
64#[macro_export]
65macro_rules! array {
66    ($([$([$([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*) => {{
67        compile_error!("Arrays of 7 dimensions or more (or ndarrays of Rust arrays) cannot be constructed with the array! macro.");
68    }};
69    ($([$([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*) => {{
70        $crate::Array6::from(vec![$([$([$([$([$([$($x,)*],)*],)*],)*],)*],)*])
71    }};
72    ($([$([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*) => {{
73        $crate::Array5::from(vec![$([$([$([$([$($x,)*],)*],)*],)*],)*])
74    }};
75    ($([$([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*]),+ $(,)*) => {{
76        $crate::Array4::from(vec![$([$([$([$($x,)*],)*],)*],)*])
77    }};
78    ($([$([$($x:expr),* $(,)*]),+ $(,)*]),+ $(,)*) => {{
79        $crate::Array3::from(vec![$([$([$($x,)*],)*],)*])
80    }};
81    ($([$($x:expr),* $(,)*]),+ $(,)*) => {{
82        $crate::Array2::from(vec![$([$($x,)*],)*])
83    }};
84    ($($x:expr),* $(,)*) => {{
85        $crate::Array::from(vec![$($x,)*])
86    }};
87}
88
89/// Create a zero-dimensional array with the element `x`.
90pub fn arr0<A>(x: A) -> Array0<A>
91{
92    unsafe { ArrayBase::from_shape_vec_unchecked((), vec![x]) }
93}
94
95/// Create a one-dimensional array with elements from `xs`.
96pub fn arr1<A: Clone>(xs: &[A]) -> Array1<A>
97{
98    ArrayBase::from(xs.to_vec())
99}
100
101/// Create a one-dimensional array with elements from `xs`.
102pub fn rcarr1<A: Clone>(xs: &[A]) -> ArcArray1<A>
103{
104    arr1(xs).into_shared()
105}
106
107/// Create a zero-dimensional array view borrowing `x`.
108pub const fn aview0<A>(x: &A) -> ArrayView0<'_, A>
109{
110    ArrayBase {
111        data: ViewRepr::new(),
112        parts: ArrayPartsSized::new(
113            // Safe because references are always non-null.
114            unsafe { NonNull::new_unchecked(x as *const A as *mut A) },
115            Ix0(),
116            Ix0(),
117        ),
118    }
119}
120
121/// Create a one-dimensional array view with elements borrowing `xs`.
122///
123/// **Panics** if the length of the slice overflows `isize`. (This can only
124/// occur if `A` is zero-sized, because slices cannot contain more than
125/// `isize::MAX` number of bytes.)
126///
127/// ```
128/// use ndarray::{aview1, ArrayView1};
129///
130/// let data = [1.0; 1024];
131///
132/// // Create a 2D array view from borrowed data
133/// let a2d = aview1(&data).into_shape_with_order((32, 32)).unwrap();
134///
135/// assert_eq!(a2d.sum(), 1024.0);
136///
137/// // Create a const 1D array view
138/// const C: ArrayView1<'static, f64> = aview1(&[1., 2., 3.]);
139///
140/// assert_eq!(C.sum(), 6.);
141/// ```
142pub const fn aview1<A>(xs: &[A]) -> ArrayView1<'_, A>
143{
144    if size_of::<A>() == 0 {
145        assert!(
146            xs.len() <= isize::MAX as usize,
147            "Slice length must fit in `isize`.",
148        );
149    }
150    ArrayBase {
151        data: ViewRepr::new(),
152        parts: ArrayPartsSized::new(
153            // Safe because references are always non-null.
154            unsafe { NonNull::new_unchecked(xs.as_ptr() as *mut A) },
155            Ix1(xs.len()),
156            Ix1(1),
157        ),
158    }
159}
160
161/// Create a two-dimensional array view with elements borrowing `xs`.
162///
163/// **Panics** if the product of non-zero axis lengths overflows `isize` (This
164/// can only occur if A is zero-sized or if `N` is zero, because slices cannot
165/// contain more than `isize::MAX` number of bytes).
166///
167/// ```
168/// use ndarray::{aview2, ArrayView2};
169///
170/// let data = vec![[1., 2., 3.], [4., 5., 6.]];
171///
172/// let view = aview2(&data);
173/// assert_eq!(view.sum(), 21.);
174///
175/// // Create a const 2D array view
176/// const C: ArrayView2<'static, f64> = aview2(&[[1., 2., 3.], [4., 5., 6.]]);
177/// assert_eq!(C.sum(), 21.);
178/// ```
179pub const fn aview2<A, const N: usize>(xs: &[[A; N]]) -> ArrayView2<'_, A>
180{
181    let cols = N;
182    let rows = xs.len();
183    if size_of::<A>() == 0 {
184        if let Some(n_elems) = rows.checked_mul(cols) {
185            assert!(
186                rows <= isize::MAX as usize
187                    && cols <= isize::MAX as usize
188                    && n_elems <= isize::MAX as usize,
189                "Product of non-zero axis lengths must not overflow isize.",
190            );
191        } else {
192            panic!("Overflow in number of elements.");
193        }
194    } else if N == 0 {
195        assert!(
196            rows <= isize::MAX as usize,
197            "Product of non-zero axis lengths must not overflow isize.",
198        );
199    }
200    // Safe because references are always non-null.
201    let ptr = unsafe { NonNull::new_unchecked(xs.as_ptr() as *mut A) };
202    let dim = Ix2(rows, cols);
203    let strides = if rows == 0 || cols == 0 {
204        Ix2(0, 0)
205    } else {
206        Ix2(cols, 1)
207    };
208    ArrayBase {
209        data: ViewRepr::new(),
210        parts: ArrayPartsSized::new(ptr, dim, strides),
211    }
212}
213
214/// Create a one-dimensional read-write array view with elements borrowing `xs`.
215///
216/// ```
217/// use ndarray::{aview_mut1, s};
218/// // Create an array view over some data, then slice it and modify it.
219/// let mut data = [0; 1024];
220/// {
221///     let mut a = aview_mut1(&mut data).into_shape_with_order((32, 32)).unwrap();
222///     a.slice_mut(s![.., ..;3]).fill(5);
223/// }
224/// assert_eq!(&data[..10], [5, 0, 0, 5, 0, 0, 5, 0, 0, 5]);
225/// ```
226pub fn aview_mut1<A>(xs: &mut [A]) -> ArrayViewMut1<'_, A>
227{
228    ArrayViewMut::from(xs)
229}
230
231/// Create a two-dimensional read-write array view with elements borrowing `xs`.
232///
233/// **Panics** if the product of non-zero axis lengths overflows `isize` (This can only occur if A
234/// is zero-sized because slices cannot contain more than `isize::MAX` number of bytes).
235///
236/// # Example
237///
238/// ```
239/// use ndarray::aview_mut2;
240///
241/// // The inner (nested) and outer arrays can be of any length.
242/// let mut data = [[0.; 2]; 128];
243/// {
244///     // Make a 128 x 2 mut array view then turn it into 2 x 128
245///     let mut a = aview_mut2(&mut data).reversed_axes();
246///     // Make the first row ones and second row minus ones.
247///     a.row_mut(0).fill(1.);
248///     a.row_mut(1).fill(-1.);
249/// }
250/// // look at the start of the result
251/// assert_eq!(&data[..3], [[1., -1.], [1., -1.], [1., -1.]]);
252/// ```
253pub fn aview_mut2<A, const N: usize>(xs: &mut [[A; N]]) -> ArrayViewMut2<'_, A>
254{
255    ArrayViewMut2::from(xs)
256}
257
258/// Create a two-dimensional array with elements from `xs`.
259///
260/// ```
261/// use ndarray::arr2;
262///
263/// let a = arr2(&[[1, 2, 3],
264///                [4, 5, 6]]);
265/// assert!(
266///     a.shape() == [2, 3]
267/// );
268/// ```
269pub fn arr2<A: Clone, const N: usize>(xs: &[[A; N]]) -> Array2<A>
270{
271    Array2::from(xs.to_vec())
272}
273
274macro_rules! impl_from_nested_vec {
275    ($arr_type:ty, $ix_type:tt, $($n:ident),+) => {
276        impl<A, $(const $n: usize),+> From<Vec<$arr_type>> for Array<A, $ix_type>
277        {
278            fn from(mut xs: Vec<$arr_type>) -> Self
279            {
280                let dim = $ix_type(xs.len(), $($n),+);
281                let ptr = xs.as_mut_ptr();
282                let cap = xs.capacity();
283                let expand_len = dimension::size_of_shape_checked(&dim)
284                    .expect("Product of non-zero axis lengths must not overflow isize.");
285                forget(xs);
286                unsafe {
287                    let v = if size_of::<A>() == 0 {
288                        Vec::from_raw_parts(ptr as *mut A, expand_len, expand_len)
289                    } else if $($n == 0 ||)+ false {
290                        Vec::new()
291                    } else {
292                        let expand_cap = cap $(* $n)+;
293                        Vec::from_raw_parts(ptr as *mut A, expand_len, expand_cap)
294                    };
295                    ArrayBase::from_shape_vec_unchecked(dim, v)
296                }
297            }
298        }
299    };
300}
301
302impl_from_nested_vec!([A; N], Ix2, N);
303impl_from_nested_vec!([[A; M]; N], Ix3, N, M);
304impl_from_nested_vec!([[[A; L]; M]; N], Ix4, N, M, L);
305impl_from_nested_vec!([[[[A; K]; L]; M]; N], Ix5, N, M, L, K);
306impl_from_nested_vec!([[[[[A; J]; K]; L]; M]; N], Ix6, N, M, L, K, J);
307
308/// Create a two-dimensional array with elements from `xs`.
309///
310pub fn rcarr2<A: Clone, const N: usize>(xs: &[[A; N]]) -> ArcArray2<A>
311{
312    arr2(xs).into_shared()
313}
314
315/// Create a three-dimensional array with elements from `xs`.
316///
317/// **Panics** if the slices are not all of the same length.
318///
319/// ```
320/// use ndarray::arr3;
321///
322/// let a = arr3(&[[[1, 2],
323///                 [3, 4]],
324///                [[5, 6],
325///                 [7, 8]],
326///                [[9, 0],
327///                 [1, 2]]]);
328/// assert!(
329///     a.shape() == [3, 2, 2]
330/// );
331/// ```
332pub fn arr3<A: Clone, const N: usize, const M: usize>(xs: &[[[A; M]; N]]) -> Array3<A>
333{
334    Array3::from(xs.to_vec())
335}
336
337/// Create a three-dimensional array with elements from `xs`.
338pub fn rcarr3<A: Clone, const N: usize, const M: usize>(xs: &[[[A; M]; N]]) -> ArcArray<A, Ix3>
339{
340    arr3(xs).into_shared()
341}
342
343/// The indexing order for [`meshgrid`]; see there for more details.
344///
345/// Controls whether the first argument to `meshgrid` will fill the rows or columns of the outputs.
346#[derive(Debug, Clone, Copy, PartialEq, Eq)]
347pub enum MeshIndex
348{
349    /// Cartesian indexing.
350    ///
351    /// The first argument of `meshgrid` will repeat over the columns of the output.
352    ///
353    /// Note: this is the default in `numpy`.
354    XY,
355    /// Matrix indexing.
356    ///
357    /// The first argument of `meshgrid` will repeat over the rows of the output.
358    IJ,
359}
360
361mod meshgrid_impl
362{
363    use super::MeshIndex;
364    use crate::extension::nonnull::nonnull_debug_checked_from_ptr;
365    use crate::{
366        ArrayBase,
367        ArrayRef1,
368        ArrayView,
369        ArrayView2,
370        ArrayView3,
371        ArrayView4,
372        ArrayView5,
373        ArrayView6,
374        Axis,
375        Data,
376        Dim,
377        IntoDimension,
378        Ix1,
379        LayoutRef1,
380    };
381
382    /// Construct the correct strides for the `idx`-th entry into meshgrid
383    fn construct_strides<A, const N: usize>(
384        arr: &LayoutRef1<A>, idx: usize, indexing: MeshIndex,
385    ) -> <[usize; N] as IntoDimension>::Dim
386    where [usize; N]: IntoDimension
387    {
388        let mut ret = [0; N];
389        if idx < 2 && indexing == MeshIndex::XY {
390            ret[1 - idx] = arr.stride_of(Axis(0)) as usize;
391        } else {
392            ret[idx] = arr.stride_of(Axis(0)) as usize;
393        }
394        Dim(ret)
395    }
396
397    /// Construct the correct shape for the `idx`-th entry into meshgrid
398    fn construct_shape<A, const N: usize>(
399        arrays: [&LayoutRef1<A>; N], indexing: MeshIndex,
400    ) -> <[usize; N] as IntoDimension>::Dim
401    where [usize; N]: IntoDimension
402    {
403        let mut ret = arrays.map(|a| a.len());
404        if indexing == MeshIndex::XY {
405            ret.swap(0, 1);
406        }
407        Dim(ret)
408    }
409
410    /// A trait to encapsulate static dispatch for [`meshgrid`](super::meshgrid); see there for more details.
411    ///
412    /// The inputs should always be some sort of 1D array.
413    /// The outputs should always be ND arrays where N is the number of inputs.
414    ///
415    /// Where possible, this trait tries to return array views rather than allocating additional memory.
416    pub trait Meshgrid
417    {
418        type Output;
419
420        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output;
421    }
422
423    macro_rules! meshgrid_body {
424        ($count:literal, $indexing:expr, $(($arr:expr, $idx:literal)),+) => {
425        {
426            let shape = construct_shape([$($arr),+], $indexing);
427            (
428                $({
429                    let strides = construct_strides::<_, $count>($arr, $idx, $indexing);
430                    unsafe { ArrayView::new(nonnull_debug_checked_from_ptr($arr.as_ptr() as *mut A), shape, strides) }
431                }),+
432            )
433        }
434    };
435    }
436
437    impl<'a, 'b, A> Meshgrid for (&'a ArrayRef1<A>, &'b ArrayRef1<A>)
438    {
439        type Output = (ArrayView2<'a, A>, ArrayView2<'b, A>);
440
441        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
442        {
443            meshgrid_body!(2, indexing, (arrays.0, 0), (arrays.1, 1))
444        }
445    }
446
447    impl<'a, 'b, S1, S2, A: 'b + 'a> Meshgrid for (&'a ArrayBase<S1, Ix1>, &'b ArrayBase<S2, Ix1>)
448    where
449        S1: Data<Elem = A>,
450        S2: Data<Elem = A>,
451    {
452        type Output = (ArrayView2<'a, A>, ArrayView2<'b, A>);
453
454        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
455        {
456            Meshgrid::meshgrid((&**arrays.0, &**arrays.1), indexing)
457        }
458    }
459
460    impl<'a, 'b, 'c, A> Meshgrid for (&'a ArrayRef1<A>, &'b ArrayRef1<A>, &'c ArrayRef1<A>)
461    {
462        type Output = (ArrayView3<'a, A>, ArrayView3<'b, A>, ArrayView3<'c, A>);
463
464        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
465        {
466            meshgrid_body!(3, indexing, (arrays.0, 0), (arrays.1, 1), (arrays.2, 2))
467        }
468    }
469
470    impl<'a, 'b, 'c, S1, S2, S3, A: 'b + 'a + 'c> Meshgrid
471        for (&'a ArrayBase<S1, Ix1>, &'b ArrayBase<S2, Ix1>, &'c ArrayBase<S3, Ix1>)
472    where
473        S1: Data<Elem = A>,
474        S2: Data<Elem = A>,
475        S3: Data<Elem = A>,
476    {
477        type Output = (ArrayView3<'a, A>, ArrayView3<'b, A>, ArrayView3<'c, A>);
478
479        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
480        {
481            Meshgrid::meshgrid((&**arrays.0, &**arrays.1, &**arrays.2), indexing)
482        }
483    }
484
485    impl<'a, 'b, 'c, 'd, A> Meshgrid for (&'a ArrayRef1<A>, &'b ArrayRef1<A>, &'c ArrayRef1<A>, &'d ArrayRef1<A>)
486    {
487        type Output = (ArrayView4<'a, A>, ArrayView4<'b, A>, ArrayView4<'c, A>, ArrayView4<'d, A>);
488
489        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
490        {
491            meshgrid_body!(4, indexing, (arrays.0, 0), (arrays.1, 1), (arrays.2, 2), (arrays.3, 3))
492        }
493    }
494
495    impl<'a, 'b, 'c, 'd, S1, S2, S3, S4, A: 'a + 'b + 'c + 'd> Meshgrid
496        for (&'a ArrayBase<S1, Ix1>, &'b ArrayBase<S2, Ix1>, &'c ArrayBase<S3, Ix1>, &'d ArrayBase<S4, Ix1>)
497    where
498        S1: Data<Elem = A>,
499        S2: Data<Elem = A>,
500        S3: Data<Elem = A>,
501        S4: Data<Elem = A>,
502    {
503        type Output = (ArrayView4<'a, A>, ArrayView4<'b, A>, ArrayView4<'c, A>, ArrayView4<'d, A>);
504
505        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
506        {
507            Meshgrid::meshgrid((&**arrays.0, &**arrays.1, &**arrays.2, &**arrays.3), indexing)
508        }
509    }
510
511    impl<'a, 'b, 'c, 'd, 'e, A> Meshgrid
512        for (&'a ArrayRef1<A>, &'b ArrayRef1<A>, &'c ArrayRef1<A>, &'d ArrayRef1<A>, &'e ArrayRef1<A>)
513    {
514        type Output = (ArrayView5<'a, A>, ArrayView5<'b, A>, ArrayView5<'c, A>, ArrayView5<'d, A>, ArrayView5<'e, A>);
515
516        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
517        {
518            meshgrid_body!(5, indexing, (arrays.0, 0), (arrays.1, 1), (arrays.2, 2), (arrays.3, 3), (arrays.4, 4))
519        }
520    }
521
522    impl<'a, 'b, 'c, 'd, 'e, S1, S2, S3, S4, S5, A: 'a + 'b + 'c + 'd + 'e> Meshgrid
523        for (
524            &'a ArrayBase<S1, Ix1>,
525            &'b ArrayBase<S2, Ix1>,
526            &'c ArrayBase<S3, Ix1>,
527            &'d ArrayBase<S4, Ix1>,
528            &'e ArrayBase<S5, Ix1>,
529        )
530    where
531        S1: Data<Elem = A>,
532        S2: Data<Elem = A>,
533        S3: Data<Elem = A>,
534        S4: Data<Elem = A>,
535        S5: Data<Elem = A>,
536    {
537        type Output = (ArrayView5<'a, A>, ArrayView5<'b, A>, ArrayView5<'c, A>, ArrayView5<'d, A>, ArrayView5<'e, A>);
538
539        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
540        {
541            Meshgrid::meshgrid((&**arrays.0, &**arrays.1, &**arrays.2, &**arrays.3, &**arrays.4), indexing)
542        }
543    }
544
545    impl<'a, 'b, 'c, 'd, 'e, 'f, A> Meshgrid
546        for (
547            &'a ArrayRef1<A>,
548            &'b ArrayRef1<A>,
549            &'c ArrayRef1<A>,
550            &'d ArrayRef1<A>,
551            &'e ArrayRef1<A>,
552            &'f ArrayRef1<A>,
553        )
554    {
555        type Output = (
556            ArrayView6<'a, A>,
557            ArrayView6<'b, A>,
558            ArrayView6<'c, A>,
559            ArrayView6<'d, A>,
560            ArrayView6<'e, A>,
561            ArrayView6<'f, A>,
562        );
563
564        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
565        {
566            meshgrid_body!(6, indexing, (arrays.0, 0), (arrays.1, 1), (arrays.2, 2), (arrays.3, 3), (arrays.4, 4), (arrays.5, 5))
567        }
568    }
569
570    impl<'a, 'b, 'c, 'd, 'e, 'f, S1, S2, S3, S4, S5, S6, A: 'a + 'b + 'c + 'd + 'e + 'f> Meshgrid
571        for (
572            &'a ArrayBase<S1, Ix1>,
573            &'b ArrayBase<S2, Ix1>,
574            &'c ArrayBase<S3, Ix1>,
575            &'d ArrayBase<S4, Ix1>,
576            &'e ArrayBase<S5, Ix1>,
577            &'f ArrayBase<S6, Ix1>,
578        )
579    where
580        S1: Data<Elem = A>,
581        S2: Data<Elem = A>,
582        S3: Data<Elem = A>,
583        S4: Data<Elem = A>,
584        S5: Data<Elem = A>,
585        S6: Data<Elem = A>,
586    {
587        type Output = (
588            ArrayView6<'a, A>,
589            ArrayView6<'b, A>,
590            ArrayView6<'c, A>,
591            ArrayView6<'d, A>,
592            ArrayView6<'e, A>,
593            ArrayView6<'f, A>,
594        );
595
596        fn meshgrid(arrays: Self, indexing: MeshIndex) -> Self::Output
597        {
598            Meshgrid::meshgrid((&**arrays.0, &**arrays.1, &**arrays.2, &**arrays.3, &**arrays.4, &**arrays.5), indexing)
599        }
600    }
601}
602
603/// Create coordinate matrices from coordinate vectors.
604///
605/// Given an N-tuple of 1D coordinate vectors, return an N-tuple of ND coordinate arrays.
606/// This is particularly useful for computing the outputs of functions with N arguments over
607/// regularly spaced grids.
608///
609/// The `indexing` argument can be controlled by [`MeshIndex`] to support both Cartesian and
610/// matrix indexing. In the two-dimensional case, inputs of length `N` and `M` will create
611/// output arrays of size `(M, N)` when using [`MeshIndex::XY`] and size `(N, M)` when using
612/// [`MeshIndex::IJ`].
613///
614/// # Example
615/// ```
616/// use ndarray::{array, meshgrid, MeshIndex};
617///
618/// let arr1 = array![1, 2];
619/// let arr2 = array![3, 4];
620/// let arr3 = array![5, 6];
621///
622/// // Cartesian indexing
623/// let (res1, res2) = meshgrid((&arr1, &arr2), MeshIndex::XY);
624/// assert_eq!(res1, array![
625///     [1, 2],
626///     [1, 2],
627/// ]);
628/// assert_eq!(res2, array![
629///     [3, 3],
630///     [4, 4],
631/// ]);
632///
633/// // Matrix indexing
634/// let (res1, res2) = meshgrid((&arr1, &arr2), MeshIndex::IJ);
635/// assert_eq!(res1, array![
636///     [1, 1],
637///     [2, 2],
638/// ]);
639/// assert_eq!(res2, array![
640///     [3, 4],
641///     [3, 4],
642/// ]);
643///
644/// let (_, _, res3) = meshgrid((&arr1, &arr2, &arr3), MeshIndex::XY);
645/// assert_eq!(res3, array![
646///     [[5, 6],
647///      [5, 6]],
648///     [[5, 6],
649///      [5, 6]],
650/// ]);
651/// ```
652pub fn meshgrid<T: Meshgrid>(arrays: T, indexing: MeshIndex) -> T::Output
653{
654    Meshgrid::meshgrid(arrays, indexing)
655}
656
657#[cfg(test)]
658mod tests
659{
660    use super::s;
661    use crate::{meshgrid, Axis, MeshIndex};
662    #[cfg(not(feature = "std"))]
663    use alloc::vec;
664
665    #[test]
666    fn test_meshgrid2()
667    {
668        let x = array![1, 2, 3];
669        let y = array![4, 5, 6, 7];
670        let (xx, yy) = meshgrid((&x, &y), MeshIndex::XY);
671        assert_eq!(xx, array![[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]);
672        assert_eq!(yy, array![[4, 4, 4], [5, 5, 5], [6, 6, 6], [7, 7, 7]]);
673
674        let (xx, yy) = meshgrid((&x, &y), MeshIndex::IJ);
675        assert_eq!(xx, array![[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]]);
676        assert_eq!(yy, array![[4, 5, 6, 7], [4, 5, 6, 7], [4, 5, 6, 7]]);
677    }
678
679    #[test]
680    fn test_meshgrid3()
681    {
682        let x = array![1, 2, 3];
683        let y = array![4, 5, 6, 7];
684        let z = array![-1, -2];
685        let (xx, yy, zz) = meshgrid((&x, &y, &z), MeshIndex::XY);
686        assert_eq!(xx, array![
687            [[1, 1], [2, 2], [3, 3]],
688            [[1, 1], [2, 2], [3, 3]],
689            [[1, 1], [2, 2], [3, 3]],
690            [[1, 1], [2, 2], [3, 3]],
691        ]);
692        assert_eq!(yy, array![
693            [[4, 4], [4, 4], [4, 4]],
694            [[5, 5], [5, 5], [5, 5]],
695            [[6, 6], [6, 6], [6, 6]],
696            [[7, 7], [7, 7], [7, 7]],
697        ]);
698        assert_eq!(zz, array![
699            [[-1, -2], [-1, -2], [-1, -2]],
700            [[-1, -2], [-1, -2], [-1, -2]],
701            [[-1, -2], [-1, -2], [-1, -2]],
702            [[-1, -2], [-1, -2], [-1, -2]],
703        ]);
704
705        let (xx, yy, zz) = meshgrid((&x, &y, &z), MeshIndex::IJ);
706        assert_eq!(xx, array![
707            [[1, 1], [1, 1], [1, 1], [1, 1]],
708            [[2, 2], [2, 2], [2, 2], [2, 2]],
709            [[3, 3], [3, 3], [3, 3], [3, 3]],
710        ]);
711        assert_eq!(yy, array![
712            [[4, 4], [5, 5], [6, 6], [7, 7]],
713            [[4, 4], [5, 5], [6, 6], [7, 7]],
714            [[4, 4], [5, 5], [6, 6], [7, 7]],
715        ]);
716        assert_eq!(zz, array![
717            [[-1, -2], [-1, -2], [-1, -2], [-1, -2]],
718            [[-1, -2], [-1, -2], [-1, -2], [-1, -2]],
719            [[-1, -2], [-1, -2], [-1, -2], [-1, -2]],
720        ]);
721    }
722
723    #[test]
724    fn test_meshgrid_from_offset()
725    {
726        let x = array![1, 2, 3];
727        let x = x.slice(s![1..]);
728        let y = array![4, 5, 6];
729        let y = y.slice(s![1..]);
730        let (xx, yy) = meshgrid((&x, &y), MeshIndex::XY);
731        assert_eq!(xx, array![[2, 3], [2, 3]]);
732        assert_eq!(yy, array![[5, 5], [6, 6]]);
733    }
734
735    #[test]
736    fn test_meshgrid_neg_stride()
737    {
738        let x = array![1, 2, 3];
739        let x = x.slice(s![..;-1]);
740        assert!(x.stride_of(Axis(0)) < 0); // Setup for test
741        let y = array![4, 5, 6];
742        let (xx, yy) = meshgrid((&x, &y), MeshIndex::XY);
743        assert_eq!(xx, array![[3, 2, 1], [3, 2, 1], [3, 2, 1]]);
744        assert_eq!(yy, array![[4, 4, 4], [5, 5, 5], [6, 6, 6]]);
745    }
746}
ndarray/free_functions.rs

ndarray/
free_functions.rs
