[BACKPORT] Integrate part of scipy elliptic functions and integrals (#3111)

docs/source/reference/tensor/special.rst

@@ -62,6 +62,24 @@ Ellipsoidal harmonics
    mars.tensor.special.ellip_normal
 
 
+Elliptic functions and integrals
+---------------------
+
+.. autosummary::
+   :toctree: generated/
+   :nosignatures:
+
+   mars.tensor.special.ellipk
+   mars.tensor.special.ellipkm1
+   mars.tensor.special.ellipkinc
+   mars.tensor.special.ellipe
+   mars.tensor.special.ellipeinc
+   mars.tensor.special.elliprc
+   mars.tensor.special.elliprf
+   mars.tensor.special.elliprg
+   mars.tensor.special.elliprj
+
+
 Gamma and related functions
 ---------------------------
 

mars/lib/sparse/__init__.py

@@ -295,6 +295,17 @@ def reciprocal(x, **kw):
 ellip_harm_2 = partial(call_sparse, "ellip_harm_2")
 ellip_normal = partial(call_sparse, "ellip_normal")
 
+ellipk = partial(_call_unary, "ellipk")
+ellipkm1 = partial(_call_unary, "ellipkm1")
+ellipkinc = partial(_call_bin, "ellipkinc")
+ellipe = partial(_call_unary, "ellipe")
+ellipeinc = partial(_call_bin, "ellipeinc")
+elliprc = partial(_call_bin, "elliprc")
+elliprd = partial(call_sparse, "elliprd")
+elliprf = partial(call_sparse, "elliprf")
+elliprg = partial(call_sparse, "elliprg")
+elliprj = partial(call_sparse, "elliprj")
+
 
 def equal(a, b, **_):
     try:

mars/lib/sparse/array.py

@@ -738,6 +738,13 @@ def _scipy_binary(self, func_name, other):
     erfcinv = partialmethod(_scipy_unary, "erfcinv")
     entr = partialmethod(_scipy_unary, "entr")
 
+    ellipk = partialmethod(_scipy_unary, "ellipk")
+    ellipkm1 = partialmethod(_scipy_unary, "ellipkm1")
+    ellipkinc = partialmethod(_scipy_binary, "ellipkinc")
+    ellipe = partialmethod(_scipy_unary, "ellipe")
+    ellipeinc = partialmethod(_scipy_binary, "ellipeinc")
+    elliprc = partialmethod(_scipy_binary, "elliprc")
+
     rel_entr = partialmethod(_scipy_binary, "rel_entr")
     kl_div = partialmethod(_scipy_binary, "kl_div")
     xlogy = partialmethod(_scipy_binary, "xlogy")

mars/tensor/special/__init__.py

@@ -127,5 +127,27 @@
         ellip_normal,
         TensorEllipNormal,
     )
+    from .ellip_func_integrals import (
+        ellipk,
+        TensorEllipk,
+        ellipkm1,
+        TensorEllipkm1,
+        ellipkinc,
+        TensorEllipkinc,
+        ellipe,
+        TensorEllipe,
+        ellipeinc,
+        TensorEllipeinc,
+        elliprc,
+        TensorElliprc,
+        elliprd,
+        TensorElliprd,
+        elliprf,
+        TensorElliprf,
+        elliprg,
+        TensorElliprg,
+        elliprj,
+        TensorElliprj,
+    )
 except ImportError:  # pragma: no cover
     pass

mars/tensor/special/ellip_func_integrals.py

@@ -0,0 +1,157 @@
+# 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 ..arithmetic.utils import arithmetic_operand
+from ..utils import infer_dtype, implement_scipy
+from .core import (
+    _register_special_op,
+    TensorSpecialBinOp,
+    TensorSpecialUnaryOp,
+    TensorSpecialMultiOp,
+)
+
+
+@_register_special_op
+@arithmetic_operand(sparse_mode="unary")
+class TensorEllipk(TensorSpecialUnaryOp):
+    _func_name = "ellipk"
+
+
+@_register_special_op
+@arithmetic_operand(sparse_mode="unary")
+class TensorEllipkm1(TensorSpecialUnaryOp):
+    _func_name = "ellipkm1"
+
+
+@_register_special_op
+@arithmetic_operand(sparse_mode="binary_and")
+class TensorEllipkinc(TensorSpecialBinOp):
+    _func_name = "ellipkinc"
+
+
+@_register_special_op
+@arithmetic_operand(sparse_mode="unary")
+class TensorEllipe(TensorSpecialUnaryOp):
+    _func_name = "ellipe"
+
+
+@_register_special_op
+@arithmetic_operand(sparse_mode="binary_and")
+class TensorEllipeinc(TensorSpecialBinOp):
+    _func_name = "ellipeinc"
+
+
+@_register_special_op
+@arithmetic_operand(sparse_mode="binary_and")
+class TensorElliprc(TensorSpecialBinOp):
+    _func_name = "elliprc"
+
+
+@_register_special_op
+class TensorElliprd(TensorSpecialMultiOp):
+    _ARG_COUNT = 3
+    _func_name = "elliprd"
+
+
+@_register_special_op
+class TensorElliprf(TensorSpecialMultiOp):
+    _ARG_COUNT = 3
+    _func_name = "elliprf"
+
+
+@_register_special_op
+class TensorElliprg(TensorSpecialMultiOp):
+    _ARG_COUNT = 3
+    _func_name = "elliprg"
+
+
+@_register_special_op
+class TensorElliprj(TensorSpecialMultiOp):
+    _ARG_COUNT = 4
+    _func_name = "elliprj"
+
+
+@implement_scipy(spspecial.ellipk)
+@infer_dtype(spspecial.ellipk)
+def ellipk(x, **kwargs):
+    op = TensorEllipk(**kwargs)
+    return op(x)
+
+
+@implement_scipy(spspecial.ellipkm1)
+@infer_dtype(spspecial.ellipkm1)
+def ellipkm1(x, **kwargs):
+    op = TensorEllipkm1(**kwargs)
+    return op(x)
+
+
+@implement_scipy(spspecial.ellipkinc)
+@infer_dtype(spspecial.ellipkinc)
+def ellipkinc(phi, m, **kwargs):
+    op = TensorEllipkinc(**kwargs)
+    return op(phi, m)
+
+
+@implement_scipy(spspecial.ellipe)
+@infer_dtype(spspecial.ellipe)
+def ellipe(x, **kwargs):
+    op = TensorEllipe(**kwargs)
+    return op(x)
+
+
+@implement_scipy(spspecial.ellipeinc)
+@infer_dtype(spspecial.ellipeinc)
+def ellipeinc(phi, m, **kwargs):
+    op = TensorEllipeinc(**kwargs)
+    return op(phi, m)
+
+
+try:
+
+    @implement_scipy(spspecial.elliprc)
+    @infer_dtype(spspecial.elliprc)
+    def elliprc(x, y, **kwargs):
+        op = TensorElliprc(**kwargs)
+        return op(x, y)
+
+    @implement_scipy(spspecial.elliprd)
+    @infer_dtype(spspecial.elliprd)
+    def elliprd(x, y, z, **kwargs):
+        op = TensorElliprd(**kwargs)
+        return op(x, y, z)
+
+    @implement_scipy(spspecial.elliprf)
+    @infer_dtype(spspecial.elliprf)
+    def elliprf(x, y, z, **kwargs):
+        op = TensorElliprf(**kwargs)
+        return op(x, y, z)
+
+    @implement_scipy(spspecial.elliprg)
+    @infer_dtype(spspecial.elliprg)
+    def elliprg(x, y, z, **kwargs):
+        op = TensorElliprg(**kwargs)
+        return op(x, y, z)
+
+    @implement_scipy(spspecial.elliprj)
+    @infer_dtype(spspecial.elliprj)
+    def elliprj(x, y, z, p, **kwargs):
+        op = TensorElliprj(**kwargs)
+        return op(x, y, z, p)
+
+except AttributeError:
+    # These functions are not implemented before scipy v1.8 so
+    # spsecial.func may cause AttributeError
+    pass