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
37/// A trait to specify when one dimension is strictly larger than another.
38///
39/// Broadcasting two arrays together frequently requires typing the resultant
40/// array has having a dimensionality equal to the maximum of the two input arrays.
41/// This trait is what determines that typing.
42///
43/// For example, `Ix1: DimMax<Ix0>`, but not vice-versa.
44pub trait DimMax<Other: Dimension>
45{
46    /// The resulting dimension type after broadcasting.
47    type Output: Dimension;
48}
49
50/// Dimensions of the same type remain unchanged when co_broadcast.
51/// So you can directly use `D` as the resulting type.
52/// (Instead of `<D as DimMax<D>>::BroadcastOutput`)
53impl<D: Dimension> DimMax<D> for D
54{
55    type Output = D;
56}
57
58macro_rules! impl_broadcast_distinct_fixed {
59    ($smaller:ty, $larger:ty) => {
60        impl DimMax<$larger> for $smaller {
61            type Output = $larger;
62        }
63
64        impl DimMax<$smaller> for $larger {
65            type Output = $larger;
66        }
67    };
68}
69
70impl_broadcast_distinct_fixed!(Ix0, Ix1);
71impl_broadcast_distinct_fixed!(Ix0, Ix2);
72impl_broadcast_distinct_fixed!(Ix0, Ix3);
73impl_broadcast_distinct_fixed!(Ix0, Ix4);
74impl_broadcast_distinct_fixed!(Ix0, Ix5);
75impl_broadcast_distinct_fixed!(Ix0, Ix6);
76impl_broadcast_distinct_fixed!(Ix1, Ix2);
77impl_broadcast_distinct_fixed!(Ix1, Ix3);
78impl_broadcast_distinct_fixed!(Ix1, Ix4);
79impl_broadcast_distinct_fixed!(Ix1, Ix5);
80impl_broadcast_distinct_fixed!(Ix1, Ix6);
81impl_broadcast_distinct_fixed!(Ix2, Ix3);
82impl_broadcast_distinct_fixed!(Ix2, Ix4);
83impl_broadcast_distinct_fixed!(Ix2, Ix5);
84impl_broadcast_distinct_fixed!(Ix2, Ix6);
85impl_broadcast_distinct_fixed!(Ix3, Ix4);
86impl_broadcast_distinct_fixed!(Ix3, Ix5);
87impl_broadcast_distinct_fixed!(Ix3, Ix6);
88impl_broadcast_distinct_fixed!(Ix4, Ix5);
89impl_broadcast_distinct_fixed!(Ix4, Ix6);
90impl_broadcast_distinct_fixed!(Ix5, Ix6);
91impl_broadcast_distinct_fixed!(Ix0, IxDyn);
92impl_broadcast_distinct_fixed!(Ix1, IxDyn);
93impl_broadcast_distinct_fixed!(Ix2, IxDyn);
94impl_broadcast_distinct_fixed!(Ix3, IxDyn);
95impl_broadcast_distinct_fixed!(Ix4, IxDyn);
96impl_broadcast_distinct_fixed!(Ix5, IxDyn);
97impl_broadcast_distinct_fixed!(Ix6, IxDyn);
98
99#[cfg(test)]
100#[cfg(feature = "std")]
101mod tests
102{
103    use super::co_broadcast;
104    use crate::{Dim, DimMax, Dimension, ErrorKind, Ix0, IxDynImpl, ShapeError};
105
106    #[test]
107    fn test_broadcast_shape()
108    {
109        fn test_co<D1, D2>(d1: &D1, d2: &D2, r: Result<<D1 as DimMax<D2>>::Output, ShapeError>)
110        where
111            D1: Dimension + DimMax<D2>,
112            D2: Dimension,
113        {
114            let d = co_broadcast::<D1, D2, <D1 as DimMax<D2>>::Output>(&d1, d2);
115            assert_eq!(d, r);
116        }
117        test_co(&Dim([2, 3]), &Dim([4, 1, 3]), Ok(Dim([4, 2, 3])));
118        test_co(&Dim([1, 2, 2]), &Dim([1, 3, 4]), Err(ShapeError::from_kind(ErrorKind::IncompatibleShape)));
119        test_co(&Dim([3, 4, 5]), &Ix0(), Ok(Dim([3, 4, 5])));
120        let v = vec![1, 2, 3, 4, 5, 6, 7];
121        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()))));
122        let d = Dim([1, 2, 1, 3]);
123        test_co(&d, &d, Ok(d));
124        test_co(&Dim([2, 1, 2]).into_dyn(), &Dim(0), Err(ShapeError::from_kind(ErrorKind::IncompatibleShape)));
125        test_co(&Dim([2, 1, 1]), &Dim([0, 0, 1, 3, 4]), Ok(Dim([0, 0, 2, 3, 4])));
126        test_co(&Dim([0]), &Dim([0, 0, 0]), Ok(Dim([0, 0, 0])));
127        test_co(&Dim(1), &Dim([1, 0, 0]), Ok(Dim([1, 0, 0])));
128        test_co(&Dim([1, 3, 0, 1, 1]), &Dim([1, 2, 3, 1]), Err(ShapeError::from_kind(ErrorKind::IncompatibleShape)));
129    }
130}
ndarray/dimension/broadcast.rs

ndarray/dimension/
broadcast.rs
