NF+BF - fix nans as scalar, test nan2zero error

Scalar nan gave an error for fancy indexing.

Author: matthew-brett

nibabel/casting.py

@@ -41,7 +41,7 @@ def _cached_int_clippers(flt_type, int_type):
     return FLT_INT_CLIPS[(flt_type, int_type)]
 
 
-class RoundingError(Exception):
+class CastingError(Exception):
     pass
 
 
@@ -56,7 +56,7 @@ def nice_round(arr, int_type, nan2zero=True, infmax=False):
         Numpy integer type
     nan2zero : {True, False}
         Whether to convert NaN value to zero.  Default is True.  If False, and
-        NaNs are present, raise RoundingError
+        NaNs are present, raise CastingError
     infmax : {False, True}
         If True, set np.inf values in `arr` to be `int_type` integer maximum
         value, -np.inf as `int_type` integer minimum.  If False, merely set infs
@@ -77,23 +77,21 @@ def nice_round(arr, int_type, nan2zero=True, infmax=False):
     arr = np.asarray(arr)
     flt_type = arr.dtype.type
     int_type = np.dtype(int_type).type
+    # Deal with scalar as input; fancy indexing needs 1D
+    shape = arr.shape
+    arr = np.atleast_1d(arr)
     mn, mx = _cached_int_clippers(flt_type, int_type)
     nans = np.isnan(arr)
     have_nans = np.any(nans)
     if not nan2zero and have_nans:
-        raise RoundingError('NaNs in array, nan2zero not True')
+        raise CastingError('NaNs in array, nan2zero not True')
     iarr = np.clip(np.rint(arr), mn, mx).astype(int_type)
     if have_nans:
         iarr[nans] = 0
     if not infmax:
-        return iarr
-    # Deal with scalar as input
-    shape = iarr.shape
-    iarr = np.atleast_1d(iarr)
-    arr = np.atleast_1d(arr)
+        return iarr.reshape(shape)
     ii = np.iinfo(int_type)
     iarr[arr == np.inf] = ii.max
     if ii.min != int(mn):
         iarr[arr == -np.inf] = ii.min
-    iarr.shape = shape
-    return iarr
+    return iarr.reshape(shape)

nibabel/tests/test_casting.py

@@ -3,10 +3,9 @@
 
 import numpy as np
 
-from ..casting import nice_round, int_clippers
+from ..casting import nice_round, int_clippers, CastingError
 
-from numpy.testing import (assert_array_almost_equal,
-                           assert_array_equal)
+from numpy.testing import (assert_array_almost_equal, assert_array_equal)
 
 from nose.tools import (assert_true, assert_equal, assert_raises)
 
@@ -64,3 +63,7 @@ def test_casting():
             assert_array_equal(farr, np.array(arr, dtype=ft))
     # Test scalars work and return scalars
     assert_array_equal(nice_round(np.float32(0), np.int16), [0])
+    # Test scalar nan OK
+    assert_array_equal(nice_round(np.nan, np.int16), [0])
+    # Test nans give error if not nan2zero
+    assert_raises(CastingError, nice_round, np.nan, np.int16, False)