Make mean_axis return Option, like mean does

Author: LukeMathWalker

examples/column_standardize.rs

@@ -23,9 +23,9 @@ fn main() {
                           [ 2.,  2.,  2.]];
 
     println!("{:8.4}", data);
-    println!("{:8.4} (Mean axis=0)", data.mean_axis(Axis(0)));
+    println!("{:8.4} (Mean axis=0)", data.mean_axis(Axis(0)).unwrap());
 
-    data -= &data.mean_axis(Axis(0));
+    data -= &data.mean_axis(Axis(0)).unwrap();
     println!("{:8.4}", data);
 
     data /= &std(&data, Axis(0));

src/numeric/impl_numeric.rs

@@ -150,8 +150,9 @@ impl<A, S, D> ArrayBase<S, D>
 
     /// Return mean along `axis`.
     ///
-    /// **Panics** if `axis` is out of bounds, if the length of the axis is
-    /// zero and division by zero panics for type `A`, or if `A::from_usize()`
+    /// Return `None` if the length of the axis is zero.
+    ///
+    /// **Panics** if `axis` is out of bounds or if `A::from_usize()`
     /// fails for the axis length.
     ///
     /// ```
@@ -160,19 +161,25 @@ impl<A, S, D> ArrayBase<S, D>
     /// let a = arr2(&[[1., 2., 3.],
     ///                [4., 5., 6.]]);
     /// assert!(
-    ///     a.mean_axis(Axis(0)) == aview1(&[2.5, 3.5, 4.5]) &&
-    ///     a.mean_axis(Axis(1)) == aview1(&[2., 5.]) &&
+    ///     a.mean_axis(Axis(0)).unwrap() == aview1(&[2.5, 3.5, 4.5]) &&
+    ///     a.mean_axis(Axis(1)).unwrap() == aview1(&[2., 5.]) &&
     ///
-    ///     a.mean_axis(Axis(0)).mean_axis(Axis(0)) == aview0(&3.5)
+    ///     a.mean_axis(Axis(0)).unwrap().mean_axis(Axis(0)).unwrap() == aview0(&3.5)
     /// );
     /// ```
-    pub fn mean_axis(&self, axis: Axis) -> Array<A, D::Smaller>
+    pub fn mean_axis(&self, axis: Axis) -> Option<Array<A, D::Smaller>>
         where A: Clone + Zero + FromPrimitive + Add<Output=A> + Div<Output=A>,
               D: RemoveAxis,
     {
-        let n = A::from_usize(self.len_of(axis)).expect("Converting axis length to `A` must not fail.");
-        let sum = self.sum_axis(axis);
-        sum / &aview0(&n)
+        let axis_length = self.len_of(axis);
+        if axis_length == 0 {
+            None
+        } else {
+            let axis_length = A::from_usize(axis_length)
+                .expect("Converting axis length to `A` must not fail.");
+            let sum = self.sum_axis(axis);
+            Some(sum / &aview0(&axis_length))
+        }
     }
 
     /// Return variance along `axis`.

tests/array.rs

@@ -931,10 +931,10 @@ fn sum_mean()
     let a = arr2(&[[1., 2.], [3., 4.]]);
     assert_eq!(a.sum_axis(Axis(0)), arr1(&[4., 6.]));
     assert_eq!(a.sum_axis(Axis(1)), arr1(&[3., 7.]));
-    assert_eq!(a.mean_axis(Axis(0)), arr1(&[2., 3.]));
-    assert_eq!(a.mean_axis(Axis(1)), arr1(&[1.5, 3.5]));
+    assert_eq!(a.mean_axis(Axis(0)), Some(arr1(&[2., 3.])));
+    assert_eq!(a.mean_axis(Axis(1)), Some(arr1(&[1.5, 3.5])));
     assert_eq!(a.sum_axis(Axis(1)).sum_axis(Axis(0)), arr0(10.));
-    assert_eq!(a.view().mean_axis(Axis(1)), aview1(&[1.5, 3.5]));
+    assert_eq!(a.view().mean_axis(Axis(1)).unwrap(), aview1(&[1.5, 3.5]));
     assert_eq!(a.sum(), 10.);
 }
 
@@ -947,11 +947,9 @@ fn sum_mean_empty() {
         Array::zeros((2, 3)),
     );
     let a = Array1::<f32>::ones(0).mean_axis(Axis(0));
-    assert_eq!(a.shape(), &[]);
-    assert!(a[()].is_nan());
+    assert_eq!(a, None);
     let a = Array3::<f32>::ones((2, 0, 3)).mean_axis(Axis(1));
-    assert_eq!(a.shape(), &[2, 3]);
-    a.mapv(|x| assert!(x.is_nan()));
+    assert_eq!(a, None);
 }
 
 #[test]

tests/complex.rs

@@ -22,5 +22,5 @@ fn complex_mat_mul()
     let r = a.dot(&e);
     println!("{}", a);
     assert_eq!(r, a);
-    assert_eq!(a.mean_axis(Axis(0)), arr1(&[c(1.5, 1.), c(2.5, 0.)]));
+    assert_eq!(a.mean_axis(Axis(0)).unwrap(), arr1(&[c(1.5, 1.), c(2.5, 0.)]));
 }