Skip to content

Commit

Permalink
Integrate remaining error functions and fresnel integrals except ```f…
Browse files Browse the repository at this point in the history
…resnel_zeros``` (#3172)
  • Loading branch information
RandomY-2 committed Jul 4, 2022
1 parent 60e19c3 commit 3484e28
Show file tree
Hide file tree
Showing 7 changed files with 239 additions and 435 deletions.
4 changes: 4 additions & 0 deletions docs/source/reference/tensor/special.rst
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,11 @@ Error functions and fresnel integrals
mars.tensor.special.erfi
mars.tensor.special.erfinv
mars.tensor.special.erfcinv
mars.tensor.special.wofz
mars.tensor.special.dawsn
mars.tensor.special.fresnel
mars.tensor.special.modfresnelp
mars.tensor.special.modfresnelm


Ellipsoidal harmonics
Expand Down
3 changes: 3 additions & 0 deletions mars/lib/sparse/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,9 @@ def reciprocal(x, **kw):
erfi = partial(_call_unary, "erfi")
erfinv = partial(_call_unary, "erfinv")
erfcinv = partial(_call_unary, "erfcinv")
wofz = partial(_call_unary, "wofz")
dawsn = partial(_call_unary, "dawsn")
voigt_profile = partial(call_sparse, "voigt_profile")

jv = partial(_call_bin, "jv")
jve = partial(_call_bin, "jve")
Expand Down
2 changes: 2 additions & 0 deletions mars/lib/sparse/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -736,6 +736,8 @@ def _scipy_binary(self, func_name, other):
erfi = partialmethod(_scipy_unary, "erfi")
erfinv = partialmethod(_scipy_unary, "erfinv")
erfcinv = partialmethod(_scipy_unary, "erfcinv")
wofz = partialmethod(_scipy_unary, "wofz")
dawsn = partialmethod(_scipy_unary, "dawsn")
entr = partialmethod(_scipy_unary, "entr")

ellipk = partialmethod(_scipy_unary, "ellipk")
Expand Down
10 changes: 10 additions & 0 deletions mars/tensor/special/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,18 @@
TensorErfinv,
erfcinv,
TensorErfcinv,
wofz,
TensorWofz,
dawsn,
TensorDawsn,
fresnel,
TensorFresnel,
modfresnelp,
TensorModFresnelP,
modfresnelm,
TensorModFresnelM,
voigt_profile,
TensorVoigtProfile,
)
from .gamma_funcs import (
gamma,
Expand Down
66 changes: 66 additions & 0 deletions mars/tensor/special/err_fresnel.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from ..utils import infer_dtype, implement_scipy
from .core import (
TensorSpecialUnaryOp,
TensorSpecialMultiOp,
TensorTupleOp,
_register_special_op,
)
Expand Down Expand Up @@ -59,12 +60,42 @@ class TensorErfcinv(TensorSpecialUnaryOp):
_func_name = "erfcinv"


@_register_special_op
@arithmetic_operand(sparse_mode="unary")
class TensorWofz(TensorSpecialUnaryOp):
_func_name = "wofz"


@_register_special_op
@arithmetic_operand(sparse_mode="unary")
class TensorDawsn(TensorSpecialUnaryOp):
_func_name = "dawsn"


@_register_special_op
class TensorFresnel(TensorTupleOp):
_func_name = "fresnel"
_n_outputs = 2


@_register_special_op
class TensorModFresnelP(TensorTupleOp):
_func_name = "modfresnelp"
_n_outputs = 2


@_register_special_op
class TensorModFresnelM(TensorTupleOp):
_func_name = "modfresnelm"
_n_outputs = 2


@_register_special_op
class TensorVoigtProfile(TensorSpecialMultiOp):
_ARG_COUNT = 3
_func_name = "voigt_profile"


@implement_scipy(spspecial.erf)
@infer_dtype(spspecial.erf)
def erf(x, out=None, where=None, **kwargs):
Expand Down Expand Up @@ -152,8 +183,43 @@ def erfcinv(x, out=None, where=None, **kwargs):
return op(x, out=out, where=where)


@implement_scipy(spspecial.wofz)
@infer_dtype(spspecial.wofz)
def wofz(x, out=None, where=None, **kwargs):
op = TensorWofz(**kwargs)
return op(x, out=out, where=where)


@implement_scipy(spspecial.dawsn)
@infer_dtype(spspecial.dawsn)
def dawsn(x, out=None, where=None, **kwargs):
op = TensorDawsn(**kwargs)
return op(x, out=out, where=where)


@implement_scipy(spspecial.fresnel)
@infer_dtype(spspecial.fresnel, multi_outputs=True)
def fresnel(x, out=None, **kwargs):
op = TensorFresnel(**kwargs)
return op(x, out=out)


@implement_scipy(spspecial.modfresnelp)
@infer_dtype(spspecial.modfresnelp, multi_outputs=True)
def modfresnelp(x, out=None, **kwargs):
op = TensorModFresnelP(**kwargs)
return op(x, out=out)


@implement_scipy(spspecial.modfresnelm)
@infer_dtype(spspecial.modfresnelm, multi_outputs=True)
def modfresnelm(x, out=None, **kwargs):
op = TensorModFresnelM(**kwargs)
return op(x, out=out)


@implement_scipy(spspecial.voigt_profile)
@infer_dtype(spspecial.voigt_profile)
def voigt_profile(x, sigma, gamma, **kwargs):
op = TensorVoigtProfile(**kwargs)
return op(x, sigma, gamma)

0 comments on commit 3484e28

Please sign in to comment.