a Jax convolv2d for real 2D arrays

[jecampagne]

Well,
I was quite surprise to not find an equivalent of scipy fftconvolve in the current JAX version. I probably miss something.
So I have setup a personal crude implementation nb in Colab

I was wandering is my convolve2r function is compliant with jax.vmap & jax.jacrev. Of course any suggestion to speed up and make better implementation is welcome. One thing for instance is that the scipy fft next_fast_len function is not implemented in JAX. This function compute the size in power of small numbers (2,3,..) greater to the original size, to get FFT working well.

def convolve2r(in1,in2):

  # suppose
  # in1 = jnp.asarray(x1)
  # in2 = jnp.asarray(x2)
  #  "assert" in1.ndim == in2.ndim == 2
  #  "assert" in1.size == in2.size
  #  "assert" in1.dtype.kind == 'f' and in2.dtype.kind == "f"

  axes = range(in1.ndim)

  s1 = in1.shape
  s2 = in2.shape
  shape = [s1[i] + s2[i] -1 for i in range(in1.ndim)]
  
  # crude next_fast_len equivalent with just power of 2
  fshape = 2**jnp.ceil(jnp.log2(jnp.array(shape))).astype('uint16')

  fslice = tuple([slice(sz) for sz in shape])

  sp1 = jnp.fft.fftn(in1, fshape, axes=axes)
  sp2 = jnp.fft.fftn(in2, fshape, axes=axes)
  conv = jnp.fft.irfftn(sp1 * sp2, fshape, axes=axes)
  conv = conv[tuple(map(slice, shape))]   #conv = jax.lax.slice(conv,(0,0),(shape))

  def _centered(arr, newshape):
    # Return the center newshape portion of the array.
    newshape = jnp.asarray(newshape)
    currshape = jnp.array(arr.shape)
    startind = (currshape - newshape) // 2
    endind = startind + newshape
    myslice = [slice(startind[k], endind[k]) for k in range(len(endind))]
    return arr[tuple(myslice)]

  return _centered(conv, s1)

Thanks

[jakevdp]

Thanks for the question and the demo code!

As written, I don't think this would be compatible with JAX transforms like jit and vmap because your slice start and end indices are dynamic. I think you could probably fix that by using np.asarray in place of jnp.asarray within the _centered helper function.

The best way to test it is to try it out: put @jax.jit on the line before your function and see if it works 😁

I think we'd be open to adding this to the jax.scipy package if you want to clean it up and open a PR.

[jecampagne]

Making an update of the notebook with a simple exercise where I only consider as Theta parameters the flux and sigma of the galaxy and fix the psf map (and I do not consider distortion)

def model(par, psf):
  flux, sig = par
  gal = f(xx,yy,s=sig)
  gal = flux * gal/sum(gal)
  gal = convolve2r(gal,psf)
  return gal

def loss(par, img_target, psf):
  img_pred = model(par, psf)
  return jnp.mean(jnp.square(img_pred - img_target))  

Then

par_init = jnp.array([1e5,1.])
loss(par_init, gal_mes,psf) # Array(37.1882973, dtype=float64)
jax.grad(loss)(par_init, gal_mes,psf) # Array([ 1.24994667e-03, -1.00068581e+02], dtype=float64)

Is ok

And using jaxopt

gd = jaxopt.GradientDescent(loss, maxiter=1000)
init_params = jnp.array([1.e5,1.])
res = gd.run(init_params,img_target=gal_mes, psf=psf)

one gets the minimizer converging nicely.

So, I was wandering how the gradient process does not see the bottleneck of the convolve2dr computation. I would appreciate a nice explanation.

[jakevdp]

Performing gradient descent does not involve storing the full Jacobean matrix, so I wouldn’t expect you to run into the same memory issues as when you call jacrev directly - that’s why I asked if your goal is really to compute the full Jacobean matrix.

[jecampagne]

So the grad wrt the 2 shape parameters is cleverly traceback to not having to compute the full jacobian as intermediate step, isn't it?

[jecampagne]

Ok, the last thing is the vmap: considering several galaxy images that all should be convolved either with the same PSF or with 1 PSF per galaxy.

[jecampagne]

Well, here is a vmap test

def raw_generator(par,xx,yy):
  flux, sig = par
  gal = f(xx,yy,s=sig)
  gal = flux * gal/jnp.sum(gal)
  return gal

#make a catalog of raw galaxy images (here 2)
catalog_raw = jnp.array([raw_generator((flux,sig), xx,yy) for flux,sig in zip((0.5e5,1.e5),(1.5,2.))])
#apply pSf on each galaxy image thanks to vmap
catalog=jax.vmap(convolve2r, in_axes=(0,None))(catalog_raw, psf)

Seems to give some relevant results. So, I can make a PR