Merge pull request #593 from helmholtz-analytics/features/546-arctan2

Feature arctan2

CHANGELOG.md

@@ -4,6 +4,7 @@
 - [#577](https://github.com/helmholtz-analytics/heat/pull/577) Add ndim property in dndarray
 - [#578](https://github.com/helmholtz-analytics/heat/pull/578) Bugfix: Bad variable in reshape
 - [#580](https://github.com/helmholtz-analytics/heat/pull/580) New feature: fliplr()
+- [#593](https://github.com/helmholtz-analytics/heat/pull/593) New feature arctan2()
 - [#594](https://github.com/helmholtz-analytics/heat/pull/594) New feature: Advanced indexing
 - [#594](https://github.com/helmholtz-analytics/heat/pull/594) Bugfix: getitem and setitem memory consumption heavily reduced
 

heat/core/tests/test_trigonometrics.py

@@ -165,6 +165,48 @@ def test_arctan(self):
         with self.assertRaises(TypeError):
             ht.arctan("hello world")
 
+    def test_arctan2(self):
+        float32_y = torch.randn(30, device=device)
+        float32_x = torch.randn(30, device=device)
+
+        float32_comparison = torch.atan2(float32_y, float32_x)
+        float32_arctan2 = ht.arctan2(ht.array(float32_y), ht.array(float32_x))
+
+        self.assertIsInstance(float32_arctan2, ht.DNDarray)
+        self.assertEqual(float32_arctan2.dtype, ht.float32)
+        self.assertTrue(torch.allclose(float32_arctan2._DNDarray__array, float32_comparison))
+
+        float64_y = torch.randn(30, dtype=torch.float64, device=device)
+        float64_x = torch.randn(30, dtype=torch.float64, device=device)
+
+        float64_comparison = torch.atan2(float64_y, float64_x)
+        float64_arctan2 = ht.arctan2(ht.array(float64_y), ht.array(float64_x))
+
+        self.assertIsInstance(float64_arctan2, ht.DNDarray)
+        self.assertEqual(float64_arctan2.dtype, ht.float64)
+        self.assertTrue(torch.allclose(float64_arctan2._DNDarray__array, float64_comparison))
+
+        # Rare Special Case with integers
+        int32_x = ht.array([-1, +1, +1, -1])
+        int32_y = ht.array([-1, -1, +1, +1])
+
+        int32_comparison = ht.array([-135.0, -45.0, 45.0, 135.0], dtype=ht.float64)
+        int32_arctan2 = ht.arctan2(int32_y, int32_x) * 180 / ht.pi
+
+        self.assertIsInstance(int32_arctan2, ht.DNDarray)
+        self.assertEqual(int32_arctan2.dtype, ht.float64)
+        self.assertTrue(ht.allclose(int32_arctan2, int32_comparison))
+
+        int16_x = ht.array([-1, +1, +1, -1], dtype=ht.int16)
+        int16_y = ht.array([-1, -1, +1, +1], dtype=ht.int16)
+
+        int16_comparison = ht.array([-135.0, -45.0, 45.0, 135.0], dtype=ht.float32)
+        int16_arctan2 = ht.arctan2(int16_y, int16_x) * 180 / ht.pi
+
+        self.assertIsInstance(int16_arctan2, ht.DNDarray)
+        self.assertEqual(int16_arctan2.dtype, ht.float32)
+        self.assertTrue(ht.allclose(int16_arctan2, int16_comparison))
+
     def test_arcsin(self):
         # base elements
         elements = [-1.0, -0.83, -0.12, 0.0, 0.24, 0.67, 1.0]

heat/core/trigonometrics.py

@@ -1,12 +1,15 @@
 import torch
 from .constants import pi
 from .operations import __local_op as local_op
+from .operations import __binary_op as binary_op
+from . import types
 
 
 __all__ = [
     "arccos",
     "arcsin",
     "arctan",
+    "arctan2",
     "cos",
     "cosh",
     "deg2rad",
@@ -99,6 +102,36 @@ def arctan(x, out=None):
     return local_op(torch.atan, x, out)
 
 
+def arctan2(x1, x2):
+    """
+    Element-wise arc tangent of ``x1/x2`` choosing the quadrant correctly.
+    Returns a new ``DNDarray`` with the signed angles in radians between vector (``x2``,``x1``) and vector (1,0)
+
+    Parameters
+    ----------
+    x1 : DNDarray
+         y-coordinates
+    x2 : DNDarray
+         x-coordinates. If ``x1.shape!=x2.shape``, they must be broadcastable to a common shape (which becomes the shape of the output).
+
+    Returns
+    -------
+    DNDarray
+
+    Examples
+    --------
+    >>> x = ht.array([-1, +1, +1, -1])
+    >>> y = ht.array([-1, -1, +1, +1])
+    >>> ht.arctan2(y, x) * 180 / ht.pi
+    tensor([-135.0000,  -45.0000,   45.0000,  135.0000], dtype=torch.float64)
+    """
+    # Cast integer to float because torch.atan2() only supports integer types on PyTorch 1.5.0.
+    x1 = x1.astype(types.promote_types(x1.dtype, types.float))
+    x2 = x2.astype(types.promote_types(x2.dtype, types.float))
+
+    return binary_op(torch.atan2, x1, x2)
+
+
 def cos(x, out=None):
     """
     Return the trigonometric cosine, element-wise.