Matplotlib upgrade

[zombie-einstein]

Hey @sash-a, thought I'd grab a look at #253.

Changes here update Matplotlib version and then makes changes to pass type checking, so it builds and passes tests.

But ...

The new API now requires that the animation update function returns a list of artists that have been updated, these are objects returned from plotting calls like a plot object, or text box (see the table here). The state/data of these objects can be updated without redrawing the whole thing.

I've done this rubiks cube environment:

# Create the first frame of the plot and return a list of drawn image artists
images = self._draw(ax, states[0])

def make_frame(state: State) -> Sequence[Artist]:
    # Update the data in each image from the state
    for i, image in enumerate(images):
        image.set_data(state.cube[i])
    # Return the updated images
    return images

# Create the animation object.
self._animation = matplotlib.animation.FuncAnimation(
    fig,
    make_frame,
    frames=states,
    interval=interval,
)

Doing this for all environments is easier for some than others given how they've been written, it means passing all the relevant artists to the animate function when the plot is created, and then correctly updating them inside the update function.

I think the animations will still work as is, as they still generally update the axes by reference inside the animation function, it's just less efficient then the new API that just updates the relevant artists.

[sash-a]

Wow this is an awesome addition, thanks @zombie-einstein!

Which environments are still remaining?

Also I think the way you've done it with make_frame looks good, but I have no experience with the new API

[zombie-einstein]

Which environments are still remaining?

Apart from rubix cube they all need a look at, with varying degrees of complexity. I think the general pattern needs to be something like a draw function that draws the first frame (and also draws the individual frames via the render method), and returns any actors that will be updated during the animation, then use this to initialise the first frame and update actors inside the update-func.

Some also currently have actors that are created dynamically during plotting I've been trying to work out the best pattern for, in some cases it seems it would be more efficient to initialise all the actors ahead of time, and update their states as needed. E.g. for the sudoko environment initialise all the text actors and update the text during animation (rather than create them on the fly).

So there's a decent chunk of work to refactor all the plots. One option could be merge something like this PR, which will updated the requirement and passes linting/tests, and still create plots (just less efficiently), and then incrementally update the individual viewers in separate PRs to make best use of the new API.

[sash-a]

Hi @zombie-einstein sorry for leaving this hanging, I needed a bit of time off 😂

I'll have a deeper look at it this week and give some a proper opinion 😄

[zombie-einstein]

Hey @sash-a, I think this is all working now, though some of the code definitely needs a work over. Turned out to be a bit more involved:

• I went with checking for a new episode from changes in state, then updating the static elements for a new episode
• I worked out how to add and remove artists inside the update function, so used this in the cases where the number of artists changes, e.g. routes in tsp
• In some cases there was an issue where old elements of the plot were not be in removed after a new episode. I think I've addressed these and the plots now all have the expected behaviour.
• I've not fully updated all the animations, in the case where they seem to work using the old functionality I've just re-used that.

I also grabbed a couple other changes whilst looking at this:

• I believe use of pkg_resources for loading images is deprecated, so updated that
• Pulled out graph layout functionality into a shared module
• Added a script (based on your notebook) to generate animations for all the environments to make testing easier

Wanted to check:

• Is it worth pulling some common functionality into the Viewer interface? There's common functionality for stuff like closing figures that is repeated in the implementations
• Do you feel all the viewers need updating in all cases? There's definitely better improved implementations for the viewers using the old functionality, but could pickup in further PRs if they are currently working otherwise.

[sash-a]

Hi @zombie-einstein thanks for this I should have time to take a look today

[sash-a]

Hi @zombie-einstein sorry for not having a look at this, the last two weeks have been a bit crazy

Is it worth pulling some common functionality into the Viewer interface? There's common functionality for stuff like closing figures that is repeated in the implementations

I think that would be nice, but also happy to do this in a subsequent PR?

Do you feel all the viewers need updating in all cases? There's definitely better improved implementations for the viewers using the old functionality, but could pickup in further PRs if they are currently working otherwise.

The main reason I'd like to do this is so that we don't pin to an old matplotlib version anymore so I do think we should update it in all cases, but if it's a bit slower than the old functionality I wouldn't worry because the viewers don't need to be too fast

[sash-a]

This looks good to me. Just a few very minor changes. Can you please explain why there are so many changes for the graph envs? I get there needs to be changes on reset, but not sure why things changed outside of the if not jnp.array_equal(...) block? (Not saying that it's wrong just curious why)

[zombie-einstein]

This looks good to me. Just a few very minor changes. Can you please explain why there are so many changes for the graph envs? I get there needs to be changes on reset, but not sure why things changed outside of the if not jnp.array_equal(...) block? (Not saying that it's wrong just curious why)

Sure, so a couple things for graphs in general:

• For the graphs the reset generally redraws the full plot, and in some cases relatively heavy computation to layout the nodes (which is REALLY slow if done every step). So the full redraw is only done on a new episode, and means removing all the existing artists and replacing them with the updated ones using some mutable storage for artists outside the animation update function.
• In some cases once the nodes and edges are drawn you can just update their properties inside an episode, e.g. updating their colour without needing to redraw them.
• In the worse cases the plot has features that vary in size every step, e.g. in crvp the route changes every step, so need to update all the artists, and so for convenience I've just been using the existing functionality to remove the existing artists and replace them with newly created ones.

Compared to other environments the graphs are a pain due to all the nodes/edges moving episode, and the potential variable number of artists each step. In comparison for something like Suduko the grid is static, and you know there is a fixed grid of numbers that can just be updated in the animation, you don't need to mess about adding and removing artists inside the update function!

[zombie-einstein]

Hi @zombie-einstein sorry for not having a look at this, the last two weeks have been a bit crazy

Is it worth pulling some common functionality into the Viewer interface? There's common functionality for stuff like closing figures that is repeated in the implementations

I think that would be nice, but also happy to do this in a subsequent PR?

Cool, I might just grab it, as where implementations vary it kind of makes it hared to refactor when trying to work out what different implementations are doing.

Do you feel all the viewers need updating in all cases? There's definitely better improved implementations for the viewers using the old functionality, but could pickup in further PRs if they are currently working otherwise.

The main reason I'd like to do this is so that we don't pin to an old matplotlib version anymore so I do think we should update it in all cases, but if it's a bit slower than the old functionality I wouldn't worry because the viewers don't need to be too fast

Everything is functional now with the updated version (the graphs were really the ones that didn't play nicely with the new API) updates would be mostly about efficient use of new API, e.g. updating artists rather then continually redrawing. My main concern was a the size of the PR, but I can chip away at the other envs if you don't mind it.

[sash-a]

Compared to other environments the graphs are a pain due to all the nodes/edges moving episode, and the potential variable number of artists each step. In comparison for something like Suduko the grid is static, and you know there is a fixed grid of numbers that can just be updated in the animation, you don't need to mess about adding and removing artists inside the update function!

I see, thanks for the explainer 👍

Everything is functional now with the updated version (the graphs were really the ones that didn't play nicely with the new API) updates would be mostly about efficient use of new API, e.g. updating artists rather then continually redrawing. My main concern was a the size of the PR, but I can chip away at the other envs if you don't mind it.

Honestly I'm quite happy with the PR as is and would prefer subsequent changes in a future PR? Could you make an issue with a brief explainer of the speed improvements and another for pulling out common functionality?

If you're happy with this I'm happy to merge the PR as is 😄

[zombie-einstein]

Thanks @sash-a sounds like a plan. I'd made a start on the pulling out common functionality so I'll roll that into this.

[sash-a]

Thanks for all the effort, really nice reduction in complexity for most of the viewers! Just have a couple questions and then I need to run it on my side to check everything still renders correctly 😄

[sash-a]

How does the saving work now?

[zombie-einstein]

Having a save path for the render method was kind of inconsistent across different viewers, since the method in the interface doesn't include the option to save I removed it across the board. If you think it's useful I could update across the viewers?

[zombie-einstein]

Thanks for all the effort, really nice reduction in complexity for most of the viewers! Just have a couple questions and then I need to run it on my side to check everything still renders correctly 😄

Thanks @sash-a, I do need to a final pass over this PR. A few of the layouts of the plots need updating now the axes creation has been consolidated, and remove the commented blocks I was just holding onto whilst I complete the tweaks.

[zombie-einstein]

Hi @sash-a , I think this is sorted now if you wanted to grab a look. It'd be great if you could see the animation and renders make sense.

I try do a final check that things still work Ok in notebooks.

[zombie-einstein]

I think this is now the correct way to update the render in a jupyter notebook (as mentioned in #272). I've tested this locally with what I guess is fairly recent version of jupyter. Only potential issue is if this breaks behaviour of older versions of notebooks?

Without this change I just get the message Javascript Error: IPython is not defined when trying to call the render method in a notebook.

[zombie-einstein]

Setting this displays animations nicely as interactive inline animations in Jupyter notebooks. I'm not totally sure this has no side-effects for in other cases (saving animations as gifs seems fine) so it seemed safest to stick it here where we know the user is using this style of matplotlib viewer.

The other option could be just to leave the rendering of animations in notebooks to the user, but seems more annoying to have to find this information out than just setting it somewhere in the package.