Bug in jax_md.colab_tools.renderer: The renderer does not plot time varying per particle parameters (e.g. diameter)

[AlmaDalco]

The renderer has a bug: It cannot plot time varying parameters per particle.

Description of problem:
renderer.Disk can take as arguments time varying per particle parameters. See snippet of documentation:

 position: An array of shape `(steps, count, dim)` or `(count, dim)` 
 specifying possibly time varying positions. Here `dim` is the spatial dimension.
 size: An array of shape (steps, count)`, `(count,)`, or `()` specifying
 possibly time-varying / per-disk diameters.
 color: An array of shape `(steps, count, 3)` or `(count,)` specifying
 possibly time-varying / per-disk RGB colors.
 count: The number of disks.

However it seems to not be able to plot time varying parameters. It plots the parameter value at first timepoint and does not change the parameters at later timepoints.

How to reproduce bug: run following script

import numpy as onp
from jax import random
from jax_md.colab_tools import renderer
from jax.ops import index_update,index, index_add

n_timepoints = 4
n_particles = 10
#make a random array of positions with size: (n_timepoints,n_particles,2)
pos = 10* onp.random.rand(n_timepoints,n_particles,2)
#make an array of diamaters with size: (n_timepoints,n_particles)
#the diamaters of all particles could vary in time, in this example is always equal to one
diameter = np.ones((n_timepoints,n_particles))
#set one particle's diamater equal to 2 at first timepoint 
diameter = index_update(diameter,index[0,0],2)

#the renderer should plot the diamaters in time, but actually only plots diamaters at first timepoint 
#and indeed one particle has diameter 2 for all timepoint rather than just for first timepoint
renderer.render(box_size, 
                { 'particles': renderer.Disk(pos, diameter)},
                buffer_size=50)

renderer.render(box_size, 
                { 'particles': renderer.Disk(position, diameter, color)})```

[sschoenholz]

Thanks for identifying this bug. The rendering code is probably some of the lightest tested. I implemented a fix; if you continue to see any problems don't hesitate to let me know!