[BACKPORT] Implementing Ellipsoidal Harmonics Functions (#2891)

.gitignore

@@ -86,6 +86,7 @@ mars/learn/cluster/*.c*
 mars/learn/utils/*.c*
 mars/lib/*.c*
 mars/oscar/**/*.c*
+mars/serialization/*.c*
 
 # web bundle file
 mars/services/web/static

docs/source/reference/tensor/special.rst

@@ -45,6 +45,18 @@ Error function
    mars.tensor.special.erf
 
 
+Ellipsoidal harmonics
+---------------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+
+   mars.tensor.special.ellip_harm
+   mars.tensor.special.ellip_harm_2
+   mars.tensor.special.ellip_normal
+
+
 Gamma and related functions
 ---------------------------
 

mars/lib/sparse/__init__.py

@@ -286,6 +286,10 @@ def reciprocal(x, **kw):
 hyperu = partial(call_sparse, "hyperu")
 hyp0f1 = partial(_call_bin, "hyp0f1")
 
+ellip_harm = partial(call_sparse, "ellip_harm")
+ellip_harm_2 = partial(call_sparse, "ellip_harm_2")
+ellip_normal = partial(call_sparse, "ellip_normal")
+
 
 def equal(a, b, **_):
     try:
@@ -727,6 +731,10 @@ def isreal(x, **kw):
     return _call_unary("isreal", x, **kw)
 
 
+def isfortran(x, **kw):
+    return call_sparse("isfortran", x, **kw)
+
+
 def where(cond, x, y):
     if any([i.ndim not in (0, 2) for i in (cond, x, y)]):
         raise NotImplementedError

mars/tensor/special/__init__.py

@@ -106,5 +106,13 @@
         hyp0f1,
         TensorHYP0F1,
     )
+    from .ellip_harm import (
+        ellip_harm,
+        TensorEllipHarm,
+        ellip_harm_2,
+        TensorEllipHarm2,
+        ellip_normal,
+        TensorEllipNormal,
+    )
 except ImportError:  # pragma: no cover
     pass

mars/tensor/special/core.py

@@ -15,7 +15,15 @@
 
 from ... import opcodes
 from ..arithmetic.core import TensorUnaryOp, TensorBinOp, TensorMultiOp
-from ..array_utils import np, cp, sparse, convert_order, as_same_device, device
+from ..array_utils import (
+    np,
+    cp,
+    issparse,
+    sparse,
+    convert_order,
+    as_same_device,
+    device,
+)
 
 
 _func_name_to_special_cls = {}
@@ -72,7 +80,14 @@ def _execute_cpu(cls, op, xp, *args, **kw):
         kw["order"] = op.order
         if kw.get("out") is not None:
             kw["out"] = np.asarray(kw["out"])
-        return cls._get_func(xp)(*args, **kw)
+        try:
+            return cls._get_func(xp)(*args, **kw)
+        except TypeError:
+            kw.pop("order")
+            r = cls._get_func(xp)(*args, **kw)
+            if issparse(r):
+                return r
+            return convert_order(r, op.outputs[0].order.value)
 
     @classmethod
     def execute(cls, ctx, op):

mars/tensor/special/ellip_harm.py

@@ -0,0 +1,57 @@
+# Copyright 1999-2021 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import scipy.special as spspecial
+
+from ..utils import infer_dtype, implement_scipy
+from .core import _register_special_op, TensorSpecialMultiOp
+
+
+@_register_special_op
+class TensorEllipHarm(TensorSpecialMultiOp):
+    _ARG_COUNT = 5
+    _func_name = "ellip_harm"
+
+
+@implement_scipy(spspecial.ellip_harm)
+@infer_dtype(spspecial.ellip_harm)
+def ellip_harm(h2, k2, n, p, s, signm=1, signn=1, **kwargs):
+    op = TensorEllipHarm(**kwargs)
+    return op(h2, k2, n, p, s, signm, signn)
+
+
+@_register_special_op
+class TensorEllipHarm2(TensorSpecialMultiOp):
+    _ARG_COUNT = 5
+    _func_name = "ellip_harm_2"
+
+
+@implement_scipy(spspecial.ellip_harm_2)
+@infer_dtype(spspecial.ellip_harm_2)
+def ellip_harm_2(h2, k2, n, p, s, **kwargs):
+    op = TensorEllipHarm2(**kwargs)
+    return op(h2, k2, n, p, s)
+
+
+@_register_special_op
+class TensorEllipNormal(TensorSpecialMultiOp):
+    _ARG_COUNT = 4
+    _func_name = "ellip_normal"
+
+
+@implement_scipy(spspecial.ellip_normal)
+@infer_dtype(spspecial.ellip_normal)
+def ellip_normal(h2, k2, n, p, **kwargs):
+    op = TensorEllipNormal(**kwargs)
+    return op(h2, k2, n, p)

mars/tensor/special/tests/test_special_execution.py

@@ -159,7 +159,7 @@ def test_triple_execution(setup, func):
     np.testing.assert_array_equal(result.toarray(), expected)
 
 
-@pytest.mark.parametrize("func", ["hyp2f1"])
+@pytest.mark.parametrize("func", ["hyp2f1", "ellip_normal"])
 def test_quadruple_execution(setup, func):
     sp_func = getattr(spspecial, func)
     mt_func = getattr(mt_special, func)
@@ -196,3 +196,46 @@ def test_quadruple_execution(setup, func):
 
     expected = sp_func(raw1.toarray(), raw2, raw3, raw4)
     np.testing.assert_array_equal(result.toarray(), expected)
+
+
+@pytest.mark.parametrize("func", ["ellip_harm", "ellip_harm_2"])
+def test_quintuple_execution(setup, func):
+    sp_func = getattr(spspecial, func)
+    mt_func = getattr(mt_special, func)
+
+    raw1 = np.random.rand(4, 3, 2)
+    raw2 = np.random.rand(4, 3, 2)
+    raw3 = np.random.rand(4, 3, 2)
+    raw4 = np.random.rand(4, 3, 2)
+    raw5 = np.random.rand(4, 3, 2)
+    a = tensor(raw1, chunk_size=3)
+    b = tensor(raw2, chunk_size=3)
+    c = tensor(raw3, chunk_size=3)
+    d = tensor(raw4, chunk_size=3)
+    e = tensor(raw5, chunk_size=3)
+
+    r = mt_func(a, b, c, d, e)
+
+    result = r.execute().fetch()
+    expected = sp_func(raw1, raw2, raw3, raw4, raw5)
+
+    np.testing.assert_array_equal(result, expected)
+
+    # test sparse
+    raw1 = sps.csr_matrix(np.array([0, 1.0, 1.01, np.nan] * 3).reshape(4, 3))
+    a = tensor(raw1, chunk_size=3)
+    raw2 = np.random.rand(4, 3)
+    b = tensor(raw2, chunk_size=3)
+    raw3 = np.random.rand(4, 3)
+    c = tensor(raw3, chunk_size=3)
+    raw4 = np.random.rand(4, 3)
+    d = tensor(raw4, chunk_size=3)
+    raw5 = np.random.rand(4, 3)
+    e = tensor(raw5, chunk_size=3)
+
+    r = mt_func(a, b, c, d, e)
+
+    result = r.execute().fetch()
+
+    expected = sp_func(raw1.toarray(), raw2, raw3, raw4, raw5)
+    np.testing.assert_array_equal(result.toarray(), expected)

setup.cfg

@@ -67,6 +67,7 @@ dev =
     flake8>=3.8.0
     black
 extra =
+    pillow>=7.0.0
     pyarrow>=0.11.0,!=0.16.*
     lz4>=1.0.0
 kubernetes =