Behaviour of grad for arctan2

[JanakyMurthy]

I would like to understand how grad works for jnp.arctan2. The code illustrates the issue I would like to understand.

def f(x):
    # We do this to bring x in the range -pi to pi
    return jnp.arctan2(jnp.sin(x),jnp.cos(x))

df = grad(f)
thetas = jnp.arange(-2*jnp.pi,2*jnp.pi+0.1,0.1)
dfs = [df(t) for t in thetas]
plt.figure()
plt.plot(thetas,dfs,'g*')
plt.show()

Can someone help me understand why this is happening and derive an expression for df(x) that explains this behavior? Thanks!

[Joshuaalbert]

Derivative of arctan(x) is 1/(1+x^2). Now using chain rule d arctan(sin(theta)/cos(theta)) / dtheta = 1/(1+sin(theta)^2/cos(theta)^2) sec(theta)^2 = cos(theta)^2/(cos(theta)^2+sin(theta)^2) sec(theta)^2 = 1. I see you're doing phase-wrapping. It's pretty clear that phase wrapping should have unit derivative except at wraps (where the derivative is not defined).