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 |
Thanks a lot!!!!!!atan2 is not supported in onnx and i've been trapped for a long time....
Thanks! It works.
How does this write into the project?
Better add an epsilon when dividing by x:
ans = torch.atan(y / (x + 1e-6))
@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)
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.])
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
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?
Thank you in advance!
Edited: I found it here: https://en.wikipedia.org/wiki/Atan2. Thank you!!!
This optimized version includes the following improvements:
Added comments in English to explain each step in the code.
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
Atan2 PyTorch implementation that can be exported to onnx