Skip to content
This repository was archived by the owner on Nov 17, 2023. It is now read-only.

Commit 77fdfb2

Browse files
author
Ying
committed
numpy operator arctan2
* change the test code * add @use_np in test code * only support float16, float32 and float64. * fix format error * remove redundant backslash * change wrapper in symbol * delete gpu test * edit test * change infer type * remove redundant **kwargs * change atol and rtol in test
1 parent 90091b1 commit 77fdfb2

File tree

9 files changed

+422
-3
lines changed

9 files changed

+422
-3
lines changed

‎python/mxnet/ndarray/numpy/_op.py

Lines changed: 88 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,8 @@
3333
'rint', 'radians', 'reciprocal', 'square', 'negative', 'fix', 'ceil', 'floor',
3434
'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'tensordot',
3535
'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack', 'mean',
36-
'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var', 'indices', 'copysign']
36+
'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var', 'indices', 'copysign',
37+
'arctan2']
3738

3839

3940
@set_module('mxnet.ndarray.numpy')
@@ -2483,3 +2484,89 @@ def copysign(x1, x2, out=None):
24832484
array([-1., 0., 1.])
24842485
"""
24852486
return _ufunc_helper(x1, x2, _npi.copysign, _np.copysign, _npi.copysign_scalar, _npi.rcopysign_scalar, out)
2487+
2488+
2489+
@set_module('mxnet.ndarray.numpy')
2490+
def arctan2(x1, x2, out=None):
2491+
r"""
2492+
arctan2(x1, x2, out=None)
2493+
2494+
Element-wise arc tangent of ``x1/x2`` choosing the quadrant correctly.
2495+
2496+
The quadrant (i.e., branch) is chosen so that ``arctan2(x1, x2)`` is
2497+
the signed angle in radians between the ray ending at the origin and
2498+
passing through the point (1,0), and the ray ending at the origin and
2499+
passing through the point (`x2`, `x1`). (Note the role reversal: the
2500+
"`y`-coordinate" is the first function parameter, the "`x`-coordinate"
2501+
is the second.) By IEEE convention, this function is defined for
2502+
`x2` = +/-0 and for either or both of `x1` and `x2` = +/-inf (see
2503+
Notes for specific values).
2504+
2505+
This function is not defined for complex-valued arguments; for the
2506+
so-called argument of complex values, use `angle`.
2507+
2508+
Parameters
2509+
----------
2510+
x1 : ndarray or scalar
2511+
`y`-coordinates.
2512+
x2 : ndarray or scalar
2513+
`x`-coordinates. `x2` must be broadcastable to match the shape of
2514+
`x1` or vice versa.
2515+
out : ndarray or None, optional
2516+
A location into which the result is stored. If provided, it must have
2517+
a shape that the inputs broadcast to. If not provided or `None`,
2518+
a freshly-allocated array is returned.
2519+
2520+
Returns
2521+
-------
2522+
out : ndarray or scalar
2523+
Array of angles in radians, in the range ``[-pi, pi]``. This is a scalar if
2524+
`x1` and `x2` are scalars.
2525+
2526+
Notes
2527+
-----
2528+
*arctan2* is identical to the `atan2` function of the underlying
2529+
C library. The following special values are defined in the C
2530+
standard: [1]_
2531+
2532+
====== ====== ================
2533+
`x1` `x2` `arctan2(x1,x2)`
2534+
====== ====== ================
2535+
+/- 0 +0 +/- 0
2536+
+/- 0 -0 +/- pi
2537+
> 0 +/-inf +0 / +pi
2538+
< 0 +/-inf -0 / -pi
2539+
+/-inf +inf +/- (pi/4)
2540+
+/-inf -inf +/- (3*pi/4)
2541+
====== ====== ================
2542+
2543+
Note that +0 and -0 are distinct floating point numbers, as are +inf
2544+
and -inf.
2545+
2546+
This function differs from the original numpy.arange in the following aspects:
2547+
- Only support float16, float32 and float64.
2548+
2549+
References
2550+
----------
2551+
.. [1] ISO/IEC standard 9899:1999, "Programming language C."
2552+
2553+
Examples
2554+
--------
2555+
Consider four points in different quadrants:
2556+
2557+
>>> x = np.array([-1, +1, +1, -1])
2558+
>>> y = np.array([-1, -1, +1, +1])
2559+
>>> np.arctan2(y, x) * 180 / np.pi
2560+
array([-135., -45., 45., 135.])
2561+
2562+
Note the order of the parameters. `arctan2` is defined also when `x2` = 0
2563+
and at several other special points, obtaining values in
2564+
the range ``[-pi, pi]``:
2565+
2566+
>>> x = np.array([1, -1])
2567+
>>> y = np.array([0, 0])
2568+
>>> np.arctan2(x, y)
2569+
array([ 1.5707964, -1.5707964])
2570+
"""
2571+
return _ufunc_helper(x1, x2, _npi.arctan2, _np.arctan2,
2572+
_npi.arctan2_scalar, _npi.rarctan2_scalar, out=out)

‎python/mxnet/numpy/multiarray.py

Lines changed: 87 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,8 @@
5252
'degrees', 'log2', 'log1p', 'rint', 'radians', 'reciprocal', 'square', 'negative',
5353
'fix', 'ceil', 'floor', 'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh',
5454
'tensordot', 'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate',
55-
'stack', 'mean', 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var', 'indices', 'copysign']
55+
'stack', 'mean', 'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var', 'indices', 'copysign',
56+
'arctan2']
5657

5758
# Return code for dispatching indexing function call
5859
_NDARRAY_UNSUPPORTED_INDEXING = -1
@@ -3986,3 +3987,88 @@ def copysign(x1, x2, out=None):
39863987
array([-1., 0., 1.])
39873988
"""
39883989
return _mx_nd_np.copysign(x1, x2, out=out)
3990+
3991+
3992+
@set_module('mxnet.numpy')
3993+
def arctan2(x1, x2, out=None):
3994+
r"""
3995+
arctan2(x1, x2, out=None)
3996+
3997+
Element-wise arc tangent of ``x1/x2`` choosing the quadrant correctly.
3998+
3999+
The quadrant (i.e., branch) is chosen so that ``arctan2(x1, x2)`` is
4000+
the signed angle in radians between the ray ending at the origin and
4001+
passing through the point (1,0), and the ray ending at the origin and
4002+
passing through the point (`x2`, `x1`). (Note the role reversal: the
4003+
"`y`-coordinate" is the first function parameter, the "`x`-coordinate"
4004+
is the second.) By IEEE convention, this function is defined for
4005+
`x2` = +/-0 and for either or both of `x1` and `x2` = +/-inf (see
4006+
Notes for specific values).
4007+
4008+
This function is not defined for complex-valued arguments; for the
4009+
so-called argument of complex values, use `angle`.
4010+
4011+
Parameters
4012+
----------
4013+
x1 : ndarray or scalar
4014+
`y`-coordinates.
4015+
x2 : ndarray or scalar
4016+
`x`-coordinates. `x2` must be broadcastable to match the shape of
4017+
`x1` or vice versa.
4018+
out : ndarray or None, optional
4019+
A location into which the result is stored. If provided, it must have
4020+
a shape that the inputs broadcast to. If not provided or `None`,
4021+
a freshly-allocated array is returned.
4022+
4023+
Returns
4024+
-------
4025+
out : ndarray or scalar
4026+
Array of angles in radians, in the range ``[-pi, pi]``. This is a scalar if
4027+
`x1` and `x2` are scalars.
4028+
4029+
Notes
4030+
-----
4031+
*arctan2* is identical to the `atan2` function of the underlying
4032+
C library. The following special values are defined in the C
4033+
standard: [1]_
4034+
4035+
====== ====== ================
4036+
`x1` `x2` `arctan2(x1,x2)`
4037+
====== ====== ================
4038+
+/- 0 +0 +/- 0
4039+
+/- 0 -0 +/- pi
4040+
> 0 +/-inf +0 / +pi
4041+
< 0 +/-inf -0 / -pi
4042+
+/-inf +inf +/- (pi/4)
4043+
+/-inf -inf +/- (3*pi/4)
4044+
====== ====== ================
4045+
4046+
Note that +0 and -0 are distinct floating point numbers, as are +inf
4047+
and -inf.
4048+
4049+
This function differs from the original numpy.arange in the following aspects:
4050+
- Only support float16, float32 and float64.
4051+
4052+
References
4053+
----------
4054+
.. [1] ISO/IEC standard 9899:1999, "Programming language C."
4055+
4056+
Examples
4057+
--------
4058+
Consider four points in different quadrants:
4059+
4060+
>>> x = np.array([-1, +1, +1, -1])
4061+
>>> y = np.array([-1, -1, +1, +1])
4062+
>>> np.arctan2(y, x) * 180 / np.pi
4063+
array([-135., -45., 45., 135.])
4064+
4065+
Note the order of the parameters. `arctan2` is defined also when `x2` = 0
4066+
and at several other special points, obtaining values in
4067+
the range ``[-pi, pi]``:
4068+
4069+
>>> x = np.array([1, -1])
4070+
>>> y = np.array([0, 0])
4071+
>>> np.arctan2(x, y)
4072+
array([ 1.5707964, -1.5707964])
4073+
"""
4074+
return _mx_nd_np.arctan2(x1, x2, out=out)

‎python/mxnet/symbol/numpy/_symbol.py

Lines changed: 70 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,8 @@
3535
'rint', 'radians', 'reciprocal', 'square', 'negative', 'fix', 'ceil', 'floor',
3636
'trunc', 'logical_not', 'arcsinh', 'arccosh', 'arctanh', 'tensordot',
3737
'linspace', 'expand_dims', 'tile', 'arange', 'split', 'concatenate', 'stack', 'mean',
38-
'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var', 'indices', 'copysign']
38+
'maximum', 'minimum', 'swapaxes', 'clip', 'argmax', 'std', 'var', 'indices', 'copysign',
39+
'arctan2']
3940

4041

4142
def _num_outputs(sym):
@@ -2778,4 +2779,72 @@ def copysign(x1, x2, out=None):
27782779
return _ufunc_helper(x1, x2, _npi.copysign, _np.copysign, _npi.copysign_scalar, _npi.rcopysign_scalar, out)
27792780

27802781

2782+
@set_module('mxnet.symbol.numpy')
2783+
def arctan2(x1, x2, out=None):
2784+
r"""
2785+
arctan2(x1, x2, out=None)
2786+
2787+
Element-wise arc tangent of ``x1/x2`` choosing the quadrant correctly.
2788+
2789+
The quadrant (i.e., branch) is chosen so that ``arctan2(x1, x2)`` is
2790+
the signed angle in radians between the ray ending at the origin and
2791+
passing through the point (1,0), and the ray ending at the origin and
2792+
passing through the point (`x2`, `x1`). (Note the role reversal: the
2793+
"`y`-coordinate" is the first function parameter, the "`x`-coordinate"
2794+
is the second.) By IEEE convention, this function is defined for
2795+
`x2` = +/-0 and for either or both of `x1` and `x2` = +/-inf (see
2796+
Notes for specific values).
2797+
2798+
This function is not defined for complex-valued arguments; for the
2799+
so-called argument of complex values, use `angle`.
2800+
2801+
Parameters
2802+
----------
2803+
x1 : _Symbol or scalar
2804+
`y`-coordinates.
2805+
x2 : _Symbol or scalar
2806+
`x`-coordinates. `x2` must be broadcastable to match the shape of
2807+
`x1` or vice versa.
2808+
out : _Symbol or None, optional
2809+
A location into which the result is stored. If provided, it must have
2810+
a shape that the inputs broadcast to. If not provided or `None`,
2811+
a freshly-allocated array is returned.
2812+
2813+
Returns
2814+
-------
2815+
out : _Symbol or scalar
2816+
Array of angles in radians, in the range ``[-pi, pi]``. This is a scalar if
2817+
`x1` and `x2` are scalars.
2818+
2819+
Notes
2820+
-----
2821+
*arctan2* is identical to the `atan2` function of the underlying
2822+
C library. The following special values are defined in the C
2823+
standard: [1]_
2824+
2825+
====== ====== ================
2826+
`x1` `x2` `arctan2(x1,x2)`
2827+
====== ====== ================
2828+
+/- 0 +0 +/- 0
2829+
+/- 0 -0 +/- pi
2830+
> 0 +/-inf +0 / +pi
2831+
< 0 +/-inf -0 / -pi
2832+
+/-inf +inf +/- (pi/4)
2833+
+/-inf -inf +/- (3*pi/4)
2834+
====== ====== ================
2835+
2836+
Note that +0 and -0 are distinct floating point numbers, as are +inf
2837+
and -inf.
2838+
2839+
This function differs from the original numpy.arange in the following aspects:
2840+
- Only support float16, float32 and float64.
2841+
2842+
References
2843+
----------
2844+
.. [1] ISO/IEC standard 9899:1999, "Programming language C."
2845+
"""
2846+
return _ufunc_helper(x1, x2, _npi.arctan2, _np.arctan2,
2847+
_npi.arctan2_scalar, _npi.rarctan2_scalar, out=out)
2848+
2849+
27812850
_set_np_symbol_class(_Symbol)

‎src/operator/math_functions-inl.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,8 @@ MXNET_BINARY_MATH_FUNC(hypot)
125125

126126
MXNET_BINARY_MATH_FUNC(pow)
127127

128+
MXNET_BINARY_MATH_FUNC(atan2)
129+
128130
template<typename DType> MSHADOW_XINLINE
129131
float id(DType a) {
130132
return static_cast<float>(a);

‎src/operator/mshadow_op.h

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -322,6 +322,16 @@ MXNET_BINARY_MATH_OP(rpower, math::pow(b, a));
322322

323323
MXNET_BINARY_MATH_OP(rpower_grad, math::id(a) * math::log(b));
324324

325+
MXNET_BINARY_MATH_OP(arctan2, math::atan2(a, b));
326+
327+
MXNET_BINARY_MATH_OP(arctan2_grad, math::id(b) / (math::id(a * a + b * b)));
328+
329+
MXNET_BINARY_MATH_OP(arctan2_rgrad, -math::id(a) / (math::id(a * a + b * b)));
330+
331+
MXNET_BINARY_MATH_OP(rarctan2, math::atan2(b, a));
332+
333+
MXNET_BINARY_MATH_OP(rarctan2_grad, math::id(a) / (math::id(a * a + b * b)));
334+
325335
MXNET_UNARY_MATH_OP_NC(nt, a != DType(0) ? DType(0) : DType(1));
326336

327337
MXNET_BINARY_MATH_OP_NC(ge, a >= b ? DType(1) : DType(0));

‎src/operator/numpy/np_elemwise_broadcast_op.cc

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,5 +144,75 @@ MXNET_OPERATOR_REGISTER_NP_BINARY_SCALAR(_backward_npi_rcopysign_scalar)
144144
.set_attr<FCompute>("FCompute<cpu>",
145145
BinaryScalarOp::Backward<cpu, mshadow_op::rcopysign_grad>);
146146

147+
inline bool IsFloatType(const int dtype) {
148+
return (dtype == mshadow::kFloat16 ||
149+
dtype == mshadow::kFloat32 ||
150+
dtype == mshadow::kFloat64);
151+
}
152+
153+
inline bool Arctan2OpType(const nnvm::NodeAttrs& attrs,
154+
std::vector<int>* in_attrs,
155+
std::vector<int>* out_attrs) {
156+
CHECK_EQ(in_attrs->size(), 2U);
157+
CHECK_EQ(out_attrs->size(), 1U);
158+
159+
TYPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
160+
TYPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(1));
161+
TYPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0));
162+
TYPE_ASSIGN_CHECK(*in_attrs, 1, out_attrs->at(0));
163+
// check if it is float16, float32 or float64. If not, raise error.
164+
CHECK(IsFloatType(in_attrs->at(0))) << "Do not support `int` as input.\n";
165+
return out_attrs->at(0) != -1;
166+
}
167+
168+
NNVM_REGISTER_OP(_npi_arctan2)
169+
.set_num_inputs(2)
170+
.set_num_outputs(1)
171+
.set_attr<nnvm::FListInputNames>("FListInputNames",
172+
[](const NodeAttrs& attrs) {
173+
return std::vector<std::string>{"x1", "x2"};
174+
})
175+
.set_attr<mxnet::FInferShape>("FInferShape", BinaryBroadcastShape)
176+
.set_attr<nnvm::FInferType>("FInferType", Arctan2OpType)
177+
.set_attr<FCompute>("FCompute<cpu>", BinaryBroadcastCompute<cpu, mshadow_op::arctan2>)
178+
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_npi_arctan2"})
179+
.set_attr<nnvm::FInplaceOption>("FInplaceOption",
180+
[](const NodeAttrs& attrs) {
181+
return std::vector<std::pair<int, int> >{{0, 0}};
182+
})
183+
.add_argument("x1", "NDArray-or-Symbol", "The input array")
184+
.add_argument("x2", "NDArray-or-Symbol", "The input array");
185+
186+
NNVM_REGISTER_OP(_backward_npi_arctan2)
187+
.set_num_inputs(3)
188+
.set_num_outputs(2)
189+
.set_attr<nnvm::TIsBackward>("TIsBackward", true)
190+
.set_attr<FResourceRequest>("FResourceRequest",
191+
[](const NodeAttrs& attrs) {
192+
return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
193+
})
194+
.set_attr<FCompute>("FCompute<cpu>", BinaryBroadcastBackwardUseIn<cpu, mshadow_op::arctan2_grad,
195+
mshadow_op::arctan2_rgrad>);
196+
197+
MXNET_OPERATOR_REGISTER_NP_BINARY_SCALAR(_npi_arctan2_scalar)
198+
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu, mshadow_op::arctan2>)
199+
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_npi_arctan2_scalar"});
200+
201+
MXNET_OPERATOR_REGISTER_NP_BINARY_SCALAR(_npi_rarctan2_scalar)
202+
.set_attr<FCompute>("FCompute<cpu>", BinaryScalarOp::Compute<cpu, mshadow_op::rarctan2>)
203+
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseIn{"_backward_npi_rarctan2_scalar"});
204+
205+
MXNET_OPERATOR_REGISTER_BINARY(_backward_npi_arctan2_scalar)
206+
.add_argument("scalar", "float", "scalar value")
207+
.set_attr_parser([](NodeAttrs *attrs) { attrs->parsed = std::stod(attrs->dict["scalar"]); })
208+
.set_attr<FCompute>("FCompute<cpu>",
209+
BinaryScalarOp::Backward<cpu, mshadow_op::arctan2_grad>);
210+
211+
MXNET_OPERATOR_REGISTER_BINARY(_backward_npi_rarctan2_scalar)
212+
.add_argument("scalar", "float", "scalar value")
213+
.set_attr_parser([](NodeAttrs *attrs) { attrs->parsed = std::stod(attrs->dict["scalar"]); })
214+
.set_attr<FCompute>("FCompute<cpu>",
215+
BinaryScalarOp::Backward<cpu, mshadow_op::arctan2_rgrad>);
216+
147217
} // namespace op
148218
} // namespace mxnet

0 commit comments

Comments
 (0)