Add tests for ensemblestats (#66)

* Add check to dimensions of input array
* Add tests to ensemblestats.mean
* Add checks to excprob
* Add tests for excprob

pysteps/postprocessing/ensemblestats.py

@@ -13,6 +13,7 @@
 
 import numpy as np
 
+
 def mean(X, ignore_nan=False, X_thr=None):
     """Compute the mean value from a forecast ensemble field.
 
@@ -32,6 +33,16 @@ def mean(X, ignore_nan=False, X_thr=None):
     out : ndarray
         Array of shape (m,n) containing the ensemble mean.
     """
+
+    X = np.asanyarray(X)
+    X_ndim = X.ndim
+
+    if X_ndim > 3 or X_ndim <= 1:
+        raise Exception('Number of dimensions of X should be 2 or 3.' +
+                        'It was: {}'.format(X_ndim))
+    elif X.ndim == 2:
+        X = X[None, ...]
+
     if ignore_nan or X_thr is not None:
         if X_thr is not None:
             X = X.copy()
@@ -41,6 +52,7 @@ def mean(X, ignore_nan=False, X_thr=None):
     else:
         return np.mean(X, axis=0)
 
+
 def excprob(X, X_thr, ignore_nan=False):
     """For a given forecast ensemble field, compute exceedance probabilities
     for the given intensity thresholds.
@@ -59,10 +71,18 @@ def excprob(X, X_thr, ignore_nan=False):
     Returns
     -------
     out : ndarray
-        Array of shape (len(X_thr),m,n) containing the exceedance probabilities 
-        for the given intensity thresholds. If len(X_thr)=1, the first dimension 
-        is dropped.
+        Array of shape (len(X_thr),m,n) containing the exceedance probabilities
+        for the given intensity thresholds.
+        If len(X_thr)=1, the first dimension is dropped.
     """
+    #  Checks
+    X = np.asanyarray(X)
+    X_ndim = X.ndim
+
+    if X_ndim < 3:
+        raise Exception('Number of dimensions of X should be 3 or more.' +
+                        ' It was: {}'.format(X_ndim))
+
     P = []
 
     if np.isscalar(X_thr):
@@ -75,7 +95,7 @@ def excprob(X, X_thr, ignore_nan=False):
         X_ = X.copy()
 
         X_[X >= x] = 1.0
-        X_[X <  x] = 0.0
+        X_[X < x] = 0.0
 
         if ignore_nan:
             P.append(np.nanmean(X_, axis=0))

pysteps/tests/test_ensemblestats.py

@@ -0,0 +1,78 @@
+# -*- coding: utf-8 -*-
+
+import pytest
+import numpy as np
+from pysteps.postprocessing.ensemblestats import mean
+from pysteps.postprocessing.ensemblestats import excprob
+from numpy.testing import assert_array_almost_equal
+
+# CREATE DATASETS TO TEST
+
+a = np.arange(9, dtype=float).reshape(3, 3)
+b = np.tile(a, (4, 1, 1))
+b1 = b.copy()
+b1[3] = np.nan
+a1 = a.copy()
+a1[:] = np.nan
+a2 = a.copy()
+a2[0, :] = np.nan
+
+#  test data
+test_data = [
+    (a, False, None, a),
+    (b, False, None, a),
+    (b1, True, None, a),
+    (b1, False, None, a1),
+    (b, False, 0.0, a),
+    (b, False, 3.0, a2),
+    (b, True, 3.0, a2),
+    (b1, True, 3.0, a2),
+    ]
+
+
+@pytest.mark.parametrize("X, ignore_nan, X_thr, expected", test_data)
+def test_ensemblestats_mean(X, ignore_nan, X_thr, expected):
+    """Test ensemblestats mean."""
+    assert_array_almost_equal(mean(X, ignore_nan, X_thr), expected)
+
+
+#  test exceptions
+test_exceptions = [(0), (None), (a[0, :]),
+                   (np.tile(a, (4, 1, 1, 1))),
+                   ]
+
+
+@pytest.mark.parametrize("X", test_exceptions)
+def test_exceptions_mean(X):
+    with pytest.raises(Exception):
+        mean(X)
+
+
+#  test data
+b2 = b.copy()
+b2[2, 2, 2] = np.nan
+
+test_data = [
+    (b, 2.0, False, np.array([[0., 0., 1.], [1., 1., 1.], [1., 1., 1.]])),
+    (b2, 2.0, False, np.array([[0., 0., 1.], [1., 1., 1.], [1., 1., np.nan]])),
+    (b2, 2.0, True, np.array([[0., 0., 1.], [1., 1., 1.], [1., 1., 1.]])),
+    ]
+
+
+@pytest.mark.parametrize("X, X_thr, ignore_nan, expected", test_data)
+def test_ensemblestats_excprob(X, X_thr, ignore_nan, expected):
+    """Test ensemblestats excprob."""
+    assert_array_almost_equal(excprob(X, X_thr, ignore_nan), expected)
+
+
+#  test exceptions
+test_exceptions = [(0), (None),
+                   (a[0, :]),
+                   (a),
+                   ]
+
+
+@pytest.mark.parametrize("X", test_exceptions)
+def test_exceptions_excprob(X):
+    with pytest.raises(Exception):
+        excprob(X, 2.0)