Gradient plot is ok?

[jecampagne]

Hi, here is a snippet on the usage of jax.vmap and jax.grad to get a gradient flow as well as contour plots on top of each other

import numpy as np


import jax
import jaxopt
import jax.numpy as jnp
jax.config.update("jax_enable_x64", True)

import matplotlib as mpl
import matplotlib.pyplot as plt
mpl.rc('image', cmap='jet')
mpl.rcParams['font.size'] = 18
mpl.rcParams["font.family"] = "Times New Roman"

sample_size = 5_000
sigma_e = 3.0         # true value of parameter error sigma
random_num_generator = np.random.RandomState(0)
x_range=20
xi = x_range * random_num_generator.rand(sample_size)
e = random_num_generator.normal(0, sigma_e, sample_size)
yi = 2*xi*np.exp(-np.sin(2*np.pi*xi/3.)) +e

def nl_lik(params, X, y):
    residuals = params[0]*X*jnp.exp(-jnp.sin(X*params[1])) - y
    return jnp.mean(residuals ** 2)

fig = plt.figure(figsize=(10,10))
grid = x,y = np.mgrid[0.0:3.0:101j,2.0:2.5:101j]

points = np.swapaxes(grid,0,-1).reshape(-1,2)
v = jax.vmap(lambda p,x,y: nl_lik(p,x,y), in_axes = (0,None,None))(points,xi,yi)
v = np.swapaxes(v.reshape(101,101),0,-1)
plt.contourf(x,y,v, levels = 100)
plt.contour(x,y,v, levels = 20, colors = 'w')

grid = x,y = np.mgrid[0.0:3.0:50j,2.0:2.5:10j]
points = np.swapaxes(grid,0,-1).reshape(-1,2)
gradients = jax.vmap(
    jax.grad(
        lambda p,x,y: nl_lik(p,x,y)
    ), in_axes = (0,None,None)
)(points,xi,yi)

plt.quiver(
    points[:,0],
    points[:,1],
    gradients[:,0],
    gradients[:,1],
    angles = 'xy',
    scale = 3.5e4
)

The result looks strange to me

as I would have expected the arrow to be orthogonal to contour levels. I'm wrong somewhere ?
Thanks

[jakevdp]

I don't think your function is returning the gradients that you think it is: grad(f)(p, x, y) will vary the components of p simultaneously. To compute the components of a vector gradient, you'll need to compute partial derivatives with respect to one direction at a time, for example using the argnums argument of jax.grad. Here's a simple example:

import jax.numpy as jnp
import jax
import matplotlib.pyplot as plt

f = lambda x, y: -(x ** 2 + y ** 2)
x, y = jnp.mgrid[-5:5:20j,-5:5:20j]

z = f(x, y)
grad_x = jax.vmap(jax.vmap(jax.grad(f, argnums=0)))(x, y)
grad_y = jax.vmap(jax.vmap(jax.grad(f, argnums=1)))(x, y)

fig, ax = plt.subplots(figsize=(8, 6))
ax.contourf(x, y, z, levels=100)
ax.quiver(x, y, grad_x, grad_y)

[jakevdp]

Ah, I see: in the notebook, the first input contains pairs of (x, y) inputs. So the equivalent example would be something like this:

import jax.numpy as jnp
import jax
import matplotlib.pyplot as plt

x, y = jnp.mgrid[-5:5:20j,-5:5:20j]
points = jnp.vstack([x.ravel(), y.ravel()]).T
f = lambda p: -(p[0] ** 2 + p[1] ** 2)
z, gradient = jax.vmap(jax.value_and_grad(f))(points)

fig, ax = plt.subplots(figsize=(8, 6))
ax.contourf(x, y, z.reshape(x.shape), levels=100)
ax.quiver(points[:, 0], points[:, 1], gradient[:, 0], gradient[:, 1])

I would compare this logic to the logic in your own example to see where the quiver gradients are going wrong.

[jakevdp]

Actually, looking again I don't think there's anything wrong with your gradients: the issue is the aspect ratio of your charts.

The data axis ratio of your plot is about 6:1, which can make it so that vector gradient representations don't visually line up with the gradient shown by the contours. You should try plotting your function with an axis ratio of 1:

plt.gca().set_aspect(1)

When I do this for your chart, the gradients appear consistent with the contours:

[jecampagne]

humm...

def jax_gradhess(model,p,xdata,ydata):
    print("pt:",p)
    print("jax grad:",jax.jacfwd(model)(p,xdata,ydata))
    print("jax hess:\n",jax.hessian(model)(p,xdata,ydata))

def fdiff_gradhess(model,p,xdata,ydata,h=0.01):
    print("pt:",p)
    f11 = lambda h1,h2:model(p+h1*np.array([1,0])+h2*np.array([0,1]),xdata,ydata)
    f10 = lambda h: f11(h,0.)
    f01 = lambda h: f11(0.,h)
    gxx = (f10(h)-f10(-h))/(2*h)
    gyy = (f01(h)-f01(-h))/(2*h)
    print("df. grad:", np.array([gxx,gyy]))
    hxx = (-f10(2*h)+16*f10(h)-30*f10(0.)+16*f10(-h)-f10(-2*h))/(12*h**2)
    hyy = (-f01(2*h)+16*f01(h)-30*f01(0.)+16*f01(-h)-f01(-2*h))/(12*h**2)
    hxy = hyx = (f11(h,h)-f11(h,-h)-f11(-h,h)+f11(-h,-h))/(4*h**2)
    print("df. hess:\n",np.array([[hxx,hxy],[hyx,hyy]]))
    
pt = np.array(nllik_model.x)
jax_gradhess(nl_lik,pt,xi,yi)
fdiff_gradhess(nl_lik,pt,xi,yi,h=0.00001)

I got a correct agreement of jax gradient/hessian compare to Finite Differenc central approx

pt: [1.99867887 2.09412972]
jax grad: [ 3.36652751e-06 -1.44656690e-03]
jax hess:
 [[   559.60721694   1089.19599011]
 [  1089.19599011 199169.16859078]]
pt: [1.99867887 2.09412972]
df. grad: [ 3.36655148e-06 -1.43256509e-03]
df. hess:
 [[   559.60726571   1089.19599118]
 [  1089.19599118 199169.16873234]]

so the problem is localised is this part ?

grid = np.mgrid[0.0:3.0:50j,2.0:2.5:10j]
points = np.swapaxes(grid,0,-1).reshape(-1,2)
gradients = jax.vmap(
    jax.grad(
        lambda p,x,y: nl_lik(p,x,y)
    ), in_axes = (0,None,None)
)(points,xi,yi)

[jakevdp]

I think our comments crossed - I believe your gradients are correct, but the chart aspect ratio results in a misleading visualization. See my previous comment.

[jecampagne]

oh ! Yes I see, I have made some test and you certainly put the finger in the right spot. Thanks. Now when I have very quite assymetric range the plot is not so nice so I do not now if one has a fix... Anyhow, Thanks for your help.