Skip to content

Commit 24290fe

Browse files
RandomY-2RandomY-2
authored
Integrate part of scipy elliptic functions and integrals (#3111)
1 parent 379c9fb commit 24290fe

File tree

7 files changed

+542
-2
lines changed

7 files changed

+542
-2
lines changed

‎docs/source/reference/tensor/special.rst

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,24 @@ Ellipsoidal harmonics
6262
mars.tensor.special.ellip_normal
6363

6464

65+
Elliptic functions and integrals
66+
---------------------
67+
68+
.. autosummary::
69+
:toctree: generated/
70+
:nosignatures:
71+
72+
mars.tensor.special.ellipk
73+
mars.tensor.special.ellipkm1
74+
mars.tensor.special.ellipkinc
75+
mars.tensor.special.ellipe
76+
mars.tensor.special.ellipeinc
77+
mars.tensor.special.elliprc
78+
mars.tensor.special.elliprf
79+
mars.tensor.special.elliprg
80+
mars.tensor.special.elliprj
81+
82+
6583
Gamma and related functions
6684
---------------------------
6785

‎mars/lib/sparse/__init__.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,17 @@ def reciprocal(x, **kw):
295295
ellip_harm_2 = partial(call_sparse, "ellip_harm_2")
296296
ellip_normal = partial(call_sparse, "ellip_normal")
297297

298+
ellipk = partial(_call_unary, "ellipk")
299+
ellipkm1 = partial(_call_unary, "ellipkm1")
300+
ellipkinc = partial(_call_bin, "ellipkinc")
301+
ellipe = partial(_call_unary, "ellipe")
302+
ellipeinc = partial(_call_bin, "ellipeinc")
303+
elliprc = partial(_call_bin, "elliprc")
304+
elliprd = partial(call_sparse, "elliprd")
305+
elliprf = partial(call_sparse, "elliprf")
306+
elliprg = partial(call_sparse, "elliprg")
307+
elliprj = partial(call_sparse, "elliprj")
308+
298309

299310
def equal(a, b, **_):
300311
try:

‎mars/lib/sparse/array.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -738,6 +738,13 @@ def _scipy_binary(self, func_name, other):
738738
erfcinv = partialmethod(_scipy_unary, "erfcinv")
739739
entr = partialmethod(_scipy_unary, "entr")
740740

741+
ellipk = partialmethod(_scipy_unary, "ellipk")
742+
ellipkm1 = partialmethod(_scipy_unary, "ellipkm1")
743+
ellipkinc = partialmethod(_scipy_binary, "ellipkinc")
744+
ellipe = partialmethod(_scipy_unary, "ellipe")
745+
ellipeinc = partialmethod(_scipy_binary, "ellipeinc")
746+
elliprc = partialmethod(_scipy_binary, "elliprc")
747+
741748
rel_entr = partialmethod(_scipy_binary, "rel_entr")
742749
kl_div = partialmethod(_scipy_binary, "kl_div")
743750
xlogy = partialmethod(_scipy_binary, "xlogy")

‎mars/tensor/special/__init__.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,5 +127,27 @@
127127
ellip_normal,
128128
TensorEllipNormal,
129129
)
130+
from .ellip_func_integrals import (
131+
ellipk,
132+
TensorEllipk,
133+
ellipkm1,
134+
TensorEllipkm1,
135+
ellipkinc,
136+
TensorEllipkinc,
137+
ellipe,
138+
TensorEllipe,
139+
ellipeinc,
140+
TensorEllipeinc,
141+
elliprc,
142+
TensorElliprc,
143+
elliprd,
144+
TensorElliprd,
145+
elliprf,
146+
TensorElliprf,
147+
elliprg,
148+
TensorElliprg,
149+
elliprj,
150+
TensorElliprj,
151+
)
130152
except ImportError: # pragma: no cover
131153
pass

‎mars/tensor/special/ellip_func_integrals.py

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
# Copyright 1999-2021 Alibaba Group Holding Ltd.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import scipy.special as spspecial
16+
17+
from ..arithmetic.utils import arithmetic_operand
18+
from ..utils import infer_dtype, implement_scipy
19+
from .core import (
20+
_register_special_op,
21+
TensorSpecialBinOp,
22+
TensorSpecialUnaryOp,
23+
TensorSpecialMultiOp,
24+
)
25+
26+
27+
@_register_special_op
28+
@arithmetic_operand(sparse_mode="unary")
29+
class TensorEllipk(TensorSpecialUnaryOp):
30+
_func_name = "ellipk"
31+
32+
33+
@_register_special_op
34+
@arithmetic_operand(sparse_mode="unary")
35+
class TensorEllipkm1(TensorSpecialUnaryOp):
36+
_func_name = "ellipkm1"
37+
38+
39+
@_register_special_op
40+
@arithmetic_operand(sparse_mode="binary_and")
41+
class TensorEllipkinc(TensorSpecialBinOp):
42+
_func_name = "ellipkinc"
43+
44+
45+
@_register_special_op
46+
@arithmetic_operand(sparse_mode="unary")
47+
class TensorEllipe(TensorSpecialUnaryOp):
48+
_func_name = "ellipe"
49+
50+
51+
@_register_special_op
52+
@arithmetic_operand(sparse_mode="binary_and")
53+
class TensorEllipeinc(TensorSpecialBinOp):
54+
_func_name = "ellipeinc"
55+
56+
57+
@_register_special_op
58+
@arithmetic_operand(sparse_mode="binary_and")
59+
class TensorElliprc(TensorSpecialBinOp):
60+
_func_name = "elliprc"
61+
62+
63+
@_register_special_op
64+
class TensorElliprd(TensorSpecialMultiOp):
65+
_ARG_COUNT = 3
66+
_func_name = "elliprd"
67+
68+
69+
@_register_special_op
70+
class TensorElliprf(TensorSpecialMultiOp):
71+
_ARG_COUNT = 3
72+
_func_name = "elliprf"
73+
74+
75+
@_register_special_op
76+
class TensorElliprg(TensorSpecialMultiOp):
77+
_ARG_COUNT = 3
78+
_func_name = "elliprg"
79+
80+
81+
@_register_special_op
82+
class TensorElliprj(TensorSpecialMultiOp):
83+
_ARG_COUNT = 4
84+
_func_name = "elliprj"
85+
86+
87+
@implement_scipy(spspecial.ellipk)
88+
@infer_dtype(spspecial.ellipk)
89+
def ellipk(x, **kwargs):
90+
op = TensorEllipk(**kwargs)
91+
return op(x)
92+
93+
94+
@implement_scipy(spspecial.ellipkm1)
95+
@infer_dtype(spspecial.ellipkm1)
96+
def ellipkm1(x, **kwargs):
97+
op = TensorEllipkm1(**kwargs)
98+
return op(x)
99+
100+
101+
@implement_scipy(spspecial.ellipkinc)
102+
@infer_dtype(spspecial.ellipkinc)
103+
def ellipkinc(phi, m, **kwargs):
104+
op = TensorEllipkinc(**kwargs)
105+
return op(phi, m)
106+
107+
108+
@implement_scipy(spspecial.ellipe)
109+
@infer_dtype(spspecial.ellipe)
110+
def ellipe(x, **kwargs):
111+
op = TensorEllipe(**kwargs)
112+
return op(x)
113+
114+
115+
@implement_scipy(spspecial.ellipeinc)
116+
@infer_dtype(spspecial.ellipeinc)
117+
def ellipeinc(phi, m, **kwargs):
118+
op = TensorEllipeinc(**kwargs)
119+
return op(phi, m)
120+
121+
122+
try:
123+
124+
@implement_scipy(spspecial.elliprc)
125+
@infer_dtype(spspecial.elliprc)
126+
def elliprc(x, y, **kwargs):
127+
op = TensorElliprc(**kwargs)
128+
return op(x, y)
129+
130+
@implement_scipy(spspecial.elliprd)
131+
@infer_dtype(spspecial.elliprd)
132+
def elliprd(x, y, z, **kwargs):
133+
op = TensorElliprd(**kwargs)
134+
return op(x, y, z)
135+
136+
@implement_scipy(spspecial.elliprf)
137+
@infer_dtype(spspecial.elliprf)
138+
def elliprf(x, y, z, **kwargs):
139+
op = TensorElliprf(**kwargs)
140+
return op(x, y, z)
141+
142+
@implement_scipy(spspecial.elliprg)
143+
@infer_dtype(spspecial.elliprg)
144+
def elliprg(x, y, z, **kwargs):
145+
op = TensorElliprg(**kwargs)
146+
return op(x, y, z)
147+
148+
@implement_scipy(spspecial.elliprj)
149+
@infer_dtype(spspecial.elliprj)
150+
def elliprj(x, y, z, p, **kwargs):
151+
op = TensorElliprj(**kwargs)
152+
return op(x, y, z, p)
153+
154+
except AttributeError:
155+
# These functions are not implemented before scipy v1.8 so
156+
# spsecial.func may cause AttributeError
157+
pass

0 commit comments

Comments
 (0)