I need to generate an animating Heatmap but not working?

I have an RL system I want to deploy on streamlit. However in this system, I need to update a heatmap to position my agent and end goal. The plot is deployed fine but the plot does not animate, I have no idea what to do.
Please guide me on this.
Here is the update code.

fig,ax = plt.subplots(1,1,figsize = (15,15))
ax.imshow(State)
ax.set_xticks(np.arange(0,n_cols))
ax.set_xticks(np.arange(-.5, n_cols, 1), minor=True)
ax.set_yticks(np.arange(-.5, n_rows, 1), minor=True)
ax.grid(which = 'minor',color='w', linestyle='-', linewidth=2)
figure = st.pyplot(fig,clear_figure=True)
#To create figure above code

        
def plot(State,s,fig,ax,windy,n_rows,n_cols,figure):
  State[s[0],s[1]] = 1
  ax.imshow(State)
  ax.set_xticks(np.arange(0,n_cols))
  ax.set_xticks(np.arange(-.5, n_cols, 1), minor=True)
  ax.set_yticks(np.arange(-.5, n_rows, 1), minor=True)
  ax.grid(which = 'minor',color='w', linestyle='-', linewidth=2)
  figure.pyplot(fig=fig, clear_figure=False)
  if s[1] not in windy:
    State[s[0],s[1]] = 0
  else:
    State[s[0],s[1]] = 0.4

this is how I run my code in a loop 
time.sleep(1)
plot(State, s, fig, ax, windy_columns,n_rows,n_cols,figure)
#All in a loop

Please guide me what I am doing wrong as it is like a similar solution I saw for animating matplotlib figure

You are in the right direction generating the frames that will compose the animation. However, the trick of sequentially throwing single pictures to the screen won’t do the trick (:melting_face: Matlab flashbacks), you’ll need to throw them to a video encoder. Here is an example using the matplotlib.animation module:

matplotlibanimate

Code:

import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np
import streamlit as st
from itertools import cycle

## Generate some data
randomdata = np.random.random(size=(5,3,3))
harvest_cycle = cycle([i for i in randomdata])
harvest = next(harvest_cycle)

## Initialize figure
fig,ax = plt.subplots(figsize=[4,4])
ax.set_aspect('equal')
heatmap = ax.imshow(harvest)

st.markdown("## :camera: Initial condition")
st.pyplot(fig)

## Animate 
def update(frame_number): 
    ax.cla() ## Clears the data passed to the axes
    ##  This is important so the figure does not store data
    ##  of every single frame and slows everything down
    
    ax.imshow(next(harvest_cycle))
    ax.set_title(f"Frame: {frame_number}",fontsize=24)

animation = FuncAnimation(fig,update,frames=50)

## Saved mp4 video
animation.save("movie.mp4")

## Keep jshmtl representation
animjs = animation.to_jshtml()

st.markdown("## :movie_camera: Enconded MP4 movie")
st.video("movie.mp4")

st.markdown("## :control_knobs: Using `jshtml` from matplotlib")
st.components.v1.html(animjs,height=600)
3 Likes