ndarray/dimension/
broadcast.rs

1use crate::error::*;
2use crate::{Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
3
4/// Calculate the common shape for a pair of array shapes, that they can be broadcasted
5/// to. Return an error if the shapes are not compatible.
6///
7/// Uses the [NumPy broadcasting rules]
8//  (https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html#general-broadcasting-rules).
9pub(crate) fn co_broadcast<D1, D2, Output>(shape1: &D1, shape2: &D2) -> Result<Output, ShapeError>
10where
11    D1: Dimension,
12    D2: Dimension,
13    Output: Dimension,
14{
15    let (k, overflow) = shape1.ndim().overflowing_sub(shape2.ndim());
16    // Swap the order if d2 is longer.
17    if overflow {
18        return co_broadcast::<D2, D1, Output>(shape2, shape1);
19    }
20    // The output should be the same length as shape1.
21    let mut out = Output::zeros(shape1.ndim());
22    for (out, s) in izip!(out.slice_mut(), shape1.slice()) {
23        *out = *s;
24    }
25    for (out, s2) in izip!(&mut out.slice_mut()[k..], shape2.slice()) {
26        if *out != *s2 {
27            if *out == 1 {
28                *out = *s2
29            } else if *s2 != 1 {
30                return Err(from_kind(ErrorKind::IncompatibleShape));
31            }
32        }
33    }
34    Ok(out)
35}
36
37pub trait DimMax<Other: Dimension>
38{
39    /// The resulting dimension type after broadcasting.
40    type Output: Dimension;
41}
42
43/// Dimensions of the same type remain unchanged when co_broadcast.
44/// So you can directly use `D` as the resulting type.
45/// (Instead of `<D as DimMax<D>>::BroadcastOutput`)
46impl<D: Dimension> DimMax<D> for D
47{
48    type Output = D;
49}
50
51macro_rules! impl_broadcast_distinct_fixed {
52    ($smaller:ty, $larger:ty) => {
53        impl DimMax<$larger> for $smaller {
54            type Output = $larger;
55        }
56
57        impl DimMax<$smaller> for $larger {
58            type Output = $larger;
59        }
60    };
61}
62
63impl_broadcast_distinct_fixed!(Ix0, Ix1);
64impl_broadcast_distinct_fixed!(Ix0, Ix2);
65impl_broadcast_distinct_fixed!(Ix0, Ix3);
66impl_broadcast_distinct_fixed!(Ix0, Ix4);
67impl_broadcast_distinct_fixed!(Ix0, Ix5);
68impl_broadcast_distinct_fixed!(Ix0, Ix6);
69impl_broadcast_distinct_fixed!(Ix1, Ix2);
70impl_broadcast_distinct_fixed!(Ix1, Ix3);
71impl_broadcast_distinct_fixed!(Ix1, Ix4);
72impl_broadcast_distinct_fixed!(Ix1, Ix5);
73impl_broadcast_distinct_fixed!(Ix1, Ix6);
74impl_broadcast_distinct_fixed!(Ix2, Ix3);
75impl_broadcast_distinct_fixed!(Ix2, Ix4);
76impl_broadcast_distinct_fixed!(Ix2, Ix5);
77impl_broadcast_distinct_fixed!(Ix2, Ix6);
78impl_broadcast_distinct_fixed!(Ix3, Ix4);
79impl_broadcast_distinct_fixed!(Ix3, Ix5);
80impl_broadcast_distinct_fixed!(Ix3, Ix6);
81impl_broadcast_distinct_fixed!(Ix4, Ix5);
82impl_broadcast_distinct_fixed!(Ix4, Ix6);
83impl_broadcast_distinct_fixed!(Ix5, Ix6);
84impl_broadcast_distinct_fixed!(Ix0, IxDyn);
85impl_broadcast_distinct_fixed!(Ix1, IxDyn);
86impl_broadcast_distinct_fixed!(Ix2, IxDyn);
87impl_broadcast_distinct_fixed!(Ix3, IxDyn);
88impl_broadcast_distinct_fixed!(Ix4, IxDyn);
89impl_broadcast_distinct_fixed!(Ix5, IxDyn);
90impl_broadcast_distinct_fixed!(Ix6, IxDyn);
91
92#[cfg(test)]
93#[cfg(feature = "std")]
94mod tests
95{
96    use super::co_broadcast;
97    use crate::{Dim, DimMax, Dimension, ErrorKind, Ix0, IxDynImpl, ShapeError};
98
99    #[test]
100    fn test_broadcast_shape()
101    {
102        fn test_co<D1, D2>(d1: &D1, d2: &D2, r: Result<<D1 as DimMax<D2>>::Output, ShapeError>)
103        where
104            D1: Dimension + DimMax<D2>,
105            D2: Dimension,
106        {
107            let d = co_broadcast::<D1, D2, <D1 as DimMax<D2>>::Output>(&d1, d2);
108            assert_eq!(d, r);
109        }
110        test_co(&Dim([2, 3]), &Dim([4, 1, 3]), Ok(Dim([4, 2, 3])));
111        test_co(&Dim([1, 2, 2]), &Dim([1, 3, 4]), Err(ShapeError::from_kind(ErrorKind::IncompatibleShape)));
112        test_co(&Dim([3, 4, 5]), &Ix0(), Ok(Dim([3, 4, 5])));
113        let v = vec![1, 2, 3, 4, 5, 6, 7];
114        test_co(&Dim(vec![1, 1, 3, 1, 5, 1, 7]), &Dim([2, 1, 4, 1, 6, 1]), Ok(Dim(IxDynImpl::from(v.as_slice()))));
115        let d = Dim([1, 2, 1, 3]);
116        test_co(&d, &d, Ok(d));
117        test_co(&Dim([2, 1, 2]).into_dyn(), &Dim(0), Err(ShapeError::from_kind(ErrorKind::IncompatibleShape)));
118        test_co(&Dim([2, 1, 1]), &Dim([0, 0, 1, 3, 4]), Ok(Dim([0, 0, 2, 3, 4])));
119        test_co(&Dim([0]), &Dim([0, 0, 0]), Ok(Dim([0, 0, 0])));
120        test_co(&Dim(1), &Dim([1, 0, 0]), Ok(Dim([1, 0, 0])));
121        test_co(&Dim([1, 3, 0, 1, 1]), &Dim([1, 2, 3, 1]), Err(ShapeError::from_kind(ErrorKind::IncompatibleShape)));
122    }
123}
ndarray/dimension/broadcast.rs

ndarray/dimension/
broadcast.rs
