‘unwrap’ can be great to have when working with sin()'s and cos()'s. ‘diff’ can be useful in general. Both ideally should eventually be included in pytorch.
Here is a pytorch port of numpy’s implementation:
def unwrap(phi, dim=-1):
assert dim is -1, ‘unwrap only supports dim=-1 for now’
dphi = diff(phi, same_size=False)
dphi_m = ((dphi+np.pi) % (2 * np.pi)) - np.pi
dphi_m[(dphi_m==-np.pi)&(dphi>0)] = np.pi
phi_adj = dphi_m-dphi
phi_adj[dphi.abs()<np.pi] = 0
return phi + phi_adj.cumsum(dim)
def diff(x, dim=-1, same_size=False):
assert dim is -1, ‘diff only supports dim=-1 for now’
if same_size:
return F.pad(x[…,1:]-x[…,:-1], (1,0))
else:
return x[…,1:]-x[…,:-1]