Instantly share code, notes, and snippets.

nikola-j/atan2.py

Last active April 27, 2023 13:36
Atan2 pytorch onnx
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
 def my_atan2(y, x): pi = torch.from_numpy(np.array([np.pi])).to(y.device, y.dtype) ans = torch.atan(y / (x + 1e-6)) ans += ((y > 0) & (x < 0)) * pi ans -= ((y < 0) & (x < 0)) * pi ans *= (1 - ((y > 0) & (x == 0)) * 1.0) ans += ((y > 0) & (x == 0)) * (pi / 2) ans *= (1 - ((y < 0) & (x == 0)) * 1.0) ans += ((y < 0) & (x == 0)) * (-pi / 2) return ans

Atan2 PyTorch implementation that can be exported to onnx

mrhe13 commented Mar 29, 2022

Thanks a lot!!!!!!atan2 is not supported in onnx and i've been trapped for a long time....

evi-Genius commented Apr 13, 2022

Thanks! It works.

gzchenjiajun commented Jun 1, 2022

How does this write into the project?

kometa-triatlon commented Sep 23, 2022

Better add an epsilon when dividing by x:

``````ans = torch.atan(y / (x + 1e-6))
``````

nikola-j commented Sep 23, 2022

@kometa-triatlon no need, torch atan handles inf, eg:

```a = torch.from_numpy(np.array([1.0]))
b = torch.from_numpy(np.array([0.0]))
torch.atan(a/b)```

Result

`tensor([1.5708], dtype=torch.float64)`

kometa-triatlon commented Sep 23, 2022 • edited

There can be a combination of parameters that would surprise you ;)

``````zero = torch.tensor([0.])
one = torch.tensor([1.])

print(my_atan2(zero, one), torch.atan2(zero, one))
print(my_atan2(one, zero), torch.atan2(one, zero))
print(my_atan2(one, one), torch.atan2(one, one))
print(my_atan2(zero, zero), torch.atan2(zero, zero))
``````

Result:

``````tensor([0.]) tensor([0.])
tensor([1.5708]) tensor([1.5708])
tensor([0.7854]) tensor([0.7854])
tensor([nan]) tensor([0.])
``````

With adding epsilon (ans = torch.atan(y/ (x + 1e-10))):

``````tensor([0.]) tensor([0.])
tensor([1.5708]) tensor([1.5708])
tensor([0.7854]) tensor([0.7854])
tensor([0.]) tensor([0.])
``````

nikola-j commented Sep 24, 2022 • edited

But atan2 of 0,0 is undefined, not 0. I think it's better to return nan, that way you would know that something is wrong.
Edit: okay I see most libraries handle 0,0 as 0, I'll add the epsilon, thanks for the suggestion

nnbtam99 commented Oct 6, 2022 • edited

Hi @nikola-j , thank you for sharing. Have you tried this implementation with complex Tensors? If possible, could you share how you derive the aboved algorithm?

Edited: I found it here: https://en.wikipedia.org/wiki/Atan2. Thank you!!!

candlewill commented Apr 27, 2023

This optimized version includes the following improvements:

Used torch.tensor to create the pi tensor directly instead of using torch.from_numpy.
Defined eps as a separate variable, making it easier to adjust if needed.
These improvements make the code more readable while maintaining performance optimizations.

```def onnx_atan2(y, x):
# Create a pi tensor with the same device and data type as y
pi = torch.tensor(np.pi, device=y.device, dtype=y.dtype)
half_pi = pi / 2
eps = 1e-6

# Compute the arctangent of y/x
ans = torch.atan(y / (x + eps))

# Create boolean tensors representing positive, negative, and zero values of y and x
y_positive = y > 0
y_negative = y < 0
x_negative = x < 0
x_zero = x == 0

# Adjust ans based on the positive, negative, and zero values of y and x
ans += torch.where(y_positive & x_negative, pi, torch.zeros_like(ans))  # Quadrants I and II
ans -= torch.where(y_negative & x_negative, pi, torch.zeros_like(ans))  # Quadrants III and IV
ans = torch.where(y_positive & x_zero, half_pi, ans)  # Positive y-axis
ans = torch.where(y_negative & x_zero, -half_pi, ans)  # Negative y-axis

return ans```