[Python] A water animation with velocity heatmap and changing surface level in python

I am trying to make an animation of water from data I generated from a simulation.

I have two datasets (Hraw and Uraw) with identical dimensions with x-coordinates and time coordinates. Uraw also contains horizontal velocities taken at equidistant points between 0 and the surface height of the water and the other dataset (Hraw) contains the y-coordinates at which the velocities are measured.

Since the water height and the positions of where the velocities are measured vary over time, I am having some trouble making the heatmap dynamic. I managed to simulate just the surface but the heats are not appearing:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# uraw and hraw have shape (t,x,y) say (174, 200, 50)

# Create grid coordinates
x_coords = np.arange(200)
X = np.tile(x_coords[:, np.newaxis], (1, 50))

# Calculate visualization parameters
max_h = np.max(hraw[:, :, -1]) # Maximum surface height
vmin, vmax = np.min(uraw), np.max(uraw) # Velocity color limits

# Set up figure
fig, ax = plt.subplots(figsize=(12, 7))
ax.set_xlabel('X Position')
ax.set_ylabel('Elevation')
ax.set_ylim(0, max_h * 1.05) # Add 5% padding
ax.set_title('Surface Elevation and Velocity Profile')

# Initial plot
pc = ax.pcolormesh(X, hraw[0], uraw[0],
shading='auto', cmap='viridis',
vmin=vmin, vmax=vmax)
line, = ax.plot(x_coords, hraw[0, :, -1], 'r-', lw=2, label='Surface')
cbar = fig.colorbar(pc, ax=ax)
cbar.set_label('Velocity Magnitude')

def animate(t):
"""Update plot for each timestep"""
# Update velocity heatmap
pc.set_array(uraw[t].ravel())
pc.set_offsets(np.column_stack([X.ravel(), hraw[t].ravel()]))

# Update surface line
line.set_ydata(hraw[t, :, -1])

ax.set_title(f'Time Step: {t}/{hraw.shape[0]-1}')
return pc, line

ani = FuncAnimation(fig, animate, frames=hraw.shape[0], interval=100, blit=False)

plt.show()


Attempt 1

I also did another one where I managed to get the heatmap to work but the heights are kept constant:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

timesteps, nx, ny = hraw.shape
x = np.linspace(0, 1, nx) # arbitrary horizontal scale
y_base = np.linspace(0, 1, ny) # normalized vertical positions
X, Y_base = np.meshgrid(x, y_base)

# Determine the global max absolute velocity for symmetric color scaling
abs_max_vel = np.max(np.abs(uraw))

# Initial setup
fig, ax = plt.subplots(figsize=(10, 6))

U0 = uraw[0].T # shape (50, 200)
H0 = hraw[0, :, -1] # surface heights at timestep 0

# Create the velocity heatmap with symmetric colormap and limits
velocity_mesh = ax.pcolormesh(
x, y_base, U0,
shading='auto',
cmap='seismic',
vmin=-abs_max_vel,
vmax=abs_max_vel
)

surface_line, = ax.plot(x, H0, color='white', lw=2, label='Surface')
colorbar = fig.colorbar(velocity_mesh, ax=ax, label='Velocity (u)')

# Set axis limits
ax.set_xlim(x[0], x[-1])
ax.set_ylim(0, np.max(hraw))
ax.set_xlabel('x')
ax.set_ylabel('Height')
ax.set_title('Water surface and horizontal velocity')
ax.legend()

def update(t):
h = hraw[t] # shape (200, 50)
u = uraw[t] # shape (200, 50)

# Construct dynamic Y grid
Y_dynamic = np.array([np.linspace(0, h[i, -1], ny) for i in range(nx)]).T

# Update pcolormesh values and offsets
velocity_mesh.set_array(u.T.ravel())
velocity_mesh.set_offsets(np.c_[X.ravel(), Y_dynamic.ravel()])

# Update surface line
surface_line.set_ydata(h[:, -1])
ax.set_title(f"Timestep {t}")

return velocity_mesh, surface_line

ani = animation.FuncAnimation(fig, update, frames=timesteps, interval=100, blit=False)
plt.show()


Attempt 2

Basically I want it to have a dynamic surface and heatmap. I managed to get them separately but not together.

Continue reading...