fixed deprecation for nd calls to math operations

Signed-off-by: Nick Papior <[email protected]>

src/sisl/_core/geometry.py

@@ -3057,7 +3057,7 @@ def a2o(self, atoms: AtomsIndex, all: bool = False) -> ndarray:
 
         # Create ranges
         if is_integral:
-            return _a.arangei(ob, oe)
+            return _a.arangei(ob[0], oe[0])
 
         return _a.array_arange(ob, oe)
 

src/sisl/_core/orbital.py

@@ -711,8 +711,8 @@ def psi(self, r, m: int = 0):
         r = _a.asarray(r)
         s = r.shape[:-1]
         # Convert to spherical coordinates
-        n, idx, r, theta, phi = cart2spher(r, theta=m != 0, cos_phi=True, maxR=self.R)
-        p = _a.zerosd(n)
+        idx, r, theta, phi = cart2spher(r, theta=m != 0, cos_phi=True, maxR=self.R)
+        p = _a.zerosd(s)
         if len(idx) > 0:
             p[idx] = self.psi_spher(r, theta, phi, m, cos_phi=True)
             # Reduce memory immediately
@@ -1570,10 +1570,8 @@ def psi(self, r):
         r = _a.asarray(r)
         s = r.shape[:-1]
         # Convert to spherical coordinates
-        n, idx, r, theta, phi = cart2spher(
-            r, theta=self.m != 0, cos_phi=True, maxR=self.R
-        )
-        p = _a.zerosd(n)
+        idx, r, theta, phi = cart2spher(r, theta=self.m != 0, cos_phi=True, maxR=self.R)
+        p = _a.zerosd(s)
         if len(idx) > 0:
             p[idx] = self.psi_spher(r, theta, phi, cos_phi=True)
             # Reduce memory immediately

src/sisl/utils/mathematics.py

@@ -161,8 +161,6 @@ def cart2spher(r, theta: bool = True, cos_phi: bool = False, maxR=None):
 
     Returns
     -------
-    n : int
-       number of total points, only for `maxR` different from ``None``
     idx : numpy.ndarray
        indices of points with ``r <= maxR``
     r : numpy.ndarray
@@ -174,36 +172,40 @@ def cart2spher(r, theta: bool = True, cos_phi: bool = False, maxR=None):
        If `cos_phi` is ``True`` this is :math:`\cos(\phi)`, otherwise
        :math:`\phi` is returned (the polar angle from the :math:`z` axis)
     """
-    r = _a.asarray(r).reshape(-1, 3)
-    n = r.shape[0]
+    r = _a.asarray(r)
     if maxR is None:
         rr = sqrt(square(r).sum(-1))
         if theta:
-            theta = arctan2(r[:, 1], r[:, 0])
+            theta = arctan2(r[..., 1], r[..., 0])
         else:
             theta = None
         phi = zeros_like(rr)
         idx = rr != 0.0
-        divide(r[:, 2], rr, out=phi, where=idx)
+        divide(r[..., 2], rr, out=phi, where=idx)
         if not cos_phi:
             arccos(phi, out=phi, where=idx)
         return rr, theta, phi
 
+    if r.ndim != 2:
+        raise NotImplementedError(
+            "cart2spher(..., maxR=1) not allowed for !=2D arrays."
+        )
+
     rr = square(r).sum(-1)
     idx = indices_le(rr, maxR**2)
     r = take(r, idx, 0)
     rr = sqrt(take(rr, idx))
     if theta:
-        theta = arctan2(r[:, 1], r[:, 0])
+        theta = arctan2(r[..., 1], r[..., 0])
     else:
         theta = None
 
     phi = zeros_like(rr)
     idx0 = rr != 0.0
-    divide(r[:, 2], rr, out=phi, where=idx0)
+    divide(r[..., 2], rr, out=phi, where=idx0)
     if not cos_phi:
         arccos(phi, out=phi, where=idx0)
-    return n, idx, rr, theta, phi
+    return idx, rr, theta, phi
 
 
 def spherical_harm(m, l, theta, phi):

src/sisl/utils/tests/test_mathematics.py

@@ -64,3 +64,21 @@ def test_curl_4d_not_3():
     a = np.random.rand(4, 3, 4, 3)
     with pytest.raises(ValueError):
         curl(a, axis=1, axisv=2)
+
+
+@pytest.mark.parametrize("nd", [0, 1, 2, 3])
+def test_cart2spher_nd(nd):
+    r = np.random.rand(*([3] * (nd + 1)))
+    R, theta, phi = cart2spher(r)
+    assert R.ndim == nd
+    assert theta.ndim == nd
+    assert phi.ndim == nd
+
+
+def test_cart2spher_nd_maxr():
+    r = np.random.rand(3, 3)
+    idx, R, theta, phi = cart2spher(r, maxR=0.5)
+    assert idx.ndim == 1
+    assert R.ndim == 1
+    assert theta.ndim == 1
+    assert phi.ndim == 1