Merge pull request #388 from jcapriot/is_scalar_bugfix

improve scalar test to handle arbitrary dimensional ndarrays
simpeg · Dec 17, 2024 · edba2be · edba2be
2 parents 8d780d2 + 06e057c
commit edba2be
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git a/discretize/utils/code_utils.py b/discretize/utils/code_utils.py
@@ -25,7 +25,11 @@ def is_scalar(f):
     """
     if isinstance(f, SCALARTYPES):
         return True
-    elif isinstance(f, np.ndarray) and f.size == 1 and isinstance(f[0], SCALARTYPES):
+    elif (
+        isinstance(f, np.ndarray)
+        and f.size == 1
+        and isinstance(f.reshape(-1)[0], SCALARTYPES)
+    ):
         return True
     return False
 

diff --git a/tests/base/test_utils.py b/tests/base/test_utils.py
@@ -238,9 +238,15 @@ def test_is_scalar(self):
         self.assertTrue(is_scalar(1.0))
         self.assertTrue(is_scalar(1))
         self.assertTrue(is_scalar(1j))
-        self.assertTrue(is_scalar(np.r_[1.0]))
-        self.assertTrue(is_scalar(np.r_[1]))
-        self.assertTrue(is_scalar(np.r_[1j]))
+        self.assertTrue(is_scalar(np.array(1.0)))
+        self.assertTrue(is_scalar(np.array(1)))
+        self.assertTrue(is_scalar(np.array(1j)))
+        self.assertTrue(is_scalar(np.array([1.0])))
+        self.assertTrue(is_scalar(np.array([1])))
+        self.assertTrue(is_scalar(np.array([1j])))
+        self.assertTrue(is_scalar(np.array([[1.0]])))
+        self.assertTrue(is_scalar(np.array([[1]])))
+        self.assertTrue(is_scalar(np.array([[1j]])))
 
     def test_as_array_n_by_dim(self):
         true = np.array([[1, 2, 3]])