How do I stop python from drawing a new colorbar every iteration

Viewed 22

I made this code as a CFD of sorts for fun, and I want to add a color bar to show the velocity of the fluid in different places. Unfortunately, every time it plots a new frame it also plots a new colorbar rather than refreshing the old one. I'd like to get it to refresh rather than draw a new one entirely. Any help would be appreciated. Plotting Begins on line 70

import numpy as np
from matplotlib import pyplot

plot_every = 100

def distance(x1,y1,x2,y2):
    return np.sqrt((x2-x1)**2 + (y2-y1)**2)

def main():
    Nx = 400 #Cells Across x direction
    Ny = 100 #Cells Across y direction
    #CELL <> NODE
    tau = .53 #kinimatic viscosity
    tymestep = tau
    Nt = 30000 #total iterations

    #Lattice Speeds and Velcoties
    NL = 9 #There are 9 differnct velocites, (up, down, left, right, up-left diag, up-right diag, down-left diag, down-right diag, and zero)
    #NL would be 27 in 3D flow
    cxs = np.array([0,0,1,1,1,0,-1,-1,-1]) #I don't know what this is
    cys = np.array([0,1,1,0,-1,-1,-1,0,1]) #I don't know what this is
    weights = np.array([4/9,1/9,1/36,1/9,1/36,1/9,1/36,1/9,1/36])
    #COMPLETELY DIFFERNT WEIGTS FOR 2D AND 3D FLOW

    #Initial Conditions
    F = np.ones((Ny,Nx,NL)) + 0.01*np.random.randn(Ny,Nx,NL)
    F[:,:,3] = 2.3 #Assigning an inital speed in x direction with right as posative

    #Drawing Our cylinder
    cylinder = np.full((Ny,Nx), False)
    radius = 13
    for y in range(0,Ny):
        for x in range(0,Nx):
            if (distance(Nx//4,Ny//2,x,y) < radius):
                cylinder[y][x] = True


    #main loop
    for it in range(Nt):
        #print(it)

        F[:,-1, [6,7,8]] = F[:,-2, [6,7,8]] #without this, fluid will bounce off of outside walls (you may want this to happen)
        F[:,0, [2,3,4]] = F[:,1, [2,3,4]] #without this, fluid will bounce off of outside walls (you may want this to happen)

        for i, cx, cy in zip(range(NL),cxs, cys): #this line is sligtly differnt than his because I think he made a typo
            F[:,:,i] = np.roll(F[:,:,i], cx, axis = 1)
            F[:,:,i] = np.roll(F[:,:,i], cy, axis = 0)

        bndryF = F[cylinder,:]
        bndryF = bndryF[:, [0,5,6,7,8,1,2,3,4]] #defines what happens in a colsion (reverse the velocity). This works by setting the up vel to down vel etc

        #Fluid Variables
        rho = np.sum(F,2) #density
        ux = np.sum(F * cxs, 2)/rho #x velocity (momentum/mass)
        uy = np.sum(F * cys, 2)/rho #y velocity

        F[cylinder,: ] = bndryF
        ux[cylinder] = 0 #set all velocities in cylinder = 0
        uy[cylinder] = 0 #set all velocities in cylinder = 0

        #collisions
        Feq = np.zeros(F.shape)
        for i, cx, cy, w in zip(range(NL), cxs, cys, weights):
            Feq[:, :, i] = rho * w * (
                1 + 3*(cx*ux + cy*uy) + 9*(cx*ux + cy*uy)**2/2 - 3*(ux**2 + uy**2)/2
            )

        F += -1/tau * (F-Feq)

        if(it%plot_every == 0):
            dfydx = ux[2:, 1:-1] - ux[0:-2, 1: -1]
            dfxdy = uy[1: -1, 2:] - uy[1: -1, 0: -2]
            curl = dfydx - dfxdy

            pyplot.imshow(np.sqrt(ux**2+uy**2),cmap = "bwr")
            #pyplot.imshow(curl, cmap = "bwr")

            pyplot.colorbar(label="Velocity", orientation="horizontal")
            pyplot.pause(0.01)
            pyplot.cla()


if __name__ == "__main__":
    main()
1 Answers

In your code you are adding a new colorbar at every iteration.

As far as I know, it is impossible to update a colorbar. The workaround is to delete the colorbar of the previous time step, and replace it with a new one. This is achieved by the update_colorbar function in the code below.

import numpy as np
from matplotlib import pyplot
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize

plot_every = 100

def distance(x1,y1,x2,y2):
    return np.sqrt((x2-x1)**2 + (y2-y1)**2)

def update_colorbar(fig, cmap, param, norm=None):
    """The name is misleading: here we create a new colorbar which will be
    placed on the same colorbar axis as the original.
    """
    # colorbar axes
    cax = None
    if len(fig.axes) > 1:
        cax = fig.axes[-1]
    
    # remove the previous colorbar, if present
    if cax is not None:
        cax.clear()
    if norm is None:
        norm = Normalize(vmin=np.amin(param), vmax=np.amax(param))
        
    mappable = ScalarMappable(cmap=cmap, norm=norm)
    fig.colorbar(mappable, orientation="horizontal", label="Velocity", cax=cax)
        
def main():
    Nx = 400 #Cells Across x direction
    Ny = 100 #Cells Across y direction
    #CELL <> NODE
    tau = .53 #kinimatic viscosity
    tymestep = tau
    Nt = 30000 #total iterations

    #Lattice Speeds and Velcoties
    NL = 9 #There are 9 differnct velocites, (up, down, left, right, up-left diag, up-right diag, down-left diag, down-right diag, and zero)
    #NL would be 27 in 3D flow
    cxs = np.array([0,0,1,1,1,0,-1,-1,-1]) #I don't know what this is
    cys = np.array([0,1,1,0,-1,-1,-1,0,1]) #I don't know what this is
    weights = np.array([4/9,1/9,1/36,1/9,1/36,1/9,1/36,1/9,1/36])
    #COMPLETELY DIFFERNT WEIGTS FOR 2D AND 3D FLOW

    #Initial Conditions
    F = np.ones((Ny,Nx,NL)) + 0.01*np.random.randn(Ny,Nx,NL)
    F[:,:,3] = 2.3 #Assigning an inital speed in x direction with right as posative

    #Drawing Our cylinder
    cylinder = np.full((Ny,Nx), False)
    radius = 13
    for y in range(0,Ny):
        for x in range(0,Nx):
            if (distance(Nx//4,Ny//2,x,y) < radius):
                cylinder[y][x] = True

    fig, ax = pyplot.subplots()
    cmap = "bwr"
    
    #main loop
    for it in range(Nt):
        # clear previous images
        ax.images.clear()
        #print(it)

        F[:,-1, [6,7,8]] = F[:,-2, [6,7,8]] #without this, fluid will bounce off of outside walls (you may want this to happen)
        F[:,0, [2,3,4]] = F[:,1, [2,3,4]] #without this, fluid will bounce off of outside walls (you may want this to happen)

        for i, cx, cy in zip(range(NL),cxs, cys): #this line is sligtly differnt than his because I think he made a typo
            F[:,:,i] = np.roll(F[:,:,i], cx, axis = 1)
            F[:,:,i] = np.roll(F[:,:,i], cy, axis = 0)

        bndryF = F[cylinder,:]
        bndryF = bndryF[:, [0,5,6,7,8,1,2,3,4]] #defines what happens in a colsion (reverse the velocity). This works by setting the up vel to down vel etc

        #Fluid Variables
        rho = np.sum(F,2) #density
        ux = np.sum(F * cxs, 2)/rho #x velocity (momentum/mass)
        uy = np.sum(F * cys, 2)/rho #y velocity

        F[cylinder,: ] = bndryF
        ux[cylinder] = 0 #set all velocities in cylinder = 0
        uy[cylinder] = 0 #set all velocities in cylinder = 0

        #collisions
        Feq = np.zeros(F.shape)
        for i, cx, cy, w in zip(range(NL), cxs, cys, weights):
            Feq[:, :, i] = rho * w * (
                1 + 3*(cx*ux + cy*uy) + 9*(cx*ux + cy*uy)**2/2 - 3*(ux**2 + uy**2)/2
            )

        F += -1/tau * (F-Feq)

        if(it%plot_every == 0):
            dfydx = ux[2:, 1:-1] - ux[0:-2, 1: -1]
            dfxdy = uy[1: -1, 2:] - uy[1: -1, 0: -2]
            curl = dfydx - dfxdy

            img = np.sqrt(ux**2+uy**2)
            ax.imshow(img ,cmap = cmap)
            #pyplot.imshow(curl, cmap = "bwr")
            
            update_colorbar(fig, cmap, param=img)
                
            pyplot.pause(0.01)

if __name__ == "__main__":
    main()

One thing you can definitely improve is the following line of code, which defines the values visible in the colorbar:

norm = Normalize(vmin=np.amin(param), vmax=np.amax(param))

Specifically, you'd have to choose a wise (conservative) value for vmax=. Currently, vmax=np.amax(param), but the maximum is going to change at every iteration. If I were you, I would chose a value big enough such that np.amax(param) < your_value, in order to ensure consistent colors for each time step.

Related