animate 1D data with LineConnection

Question

The example code below creates 2 plots that attempt to convey what the desired animation is seeking to do. Note that the plot on the right is similar to the plot on the left, but represents the position of a wave after the passage of time t. The goal is to animate the downward movement of waves, with the waves depicted using color. The reason for adopting an approach that uses LineConnections is because the spacing is not necessarily regular (notice that the first two values in z don't follow the same pattern as the rest of the data in z. Is there a way to modify the small reproducible example below to animate the downward movement of the color-filled waves? I've not found examples of the animation feature that uses LineConnections. [Also, I'm open to other approaches that would accomodate irregularly spaced data (i.e., z position may be irregularly spaced)]

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

# x position of 1D profile
x = [0.5 for i in np.arange(20)]
# z position of 1D profile
z = [59.95, 59.70, 59.25, 58.75, 59.25,
     58.75, 58.25, 57.75, 57.25, 56.75,
     56.25, 55.75, 55.25, 54.75, 54.25,
     53.75, 53.25, 52.75, 52.25, 51.75]

points = np.array([x, z]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

norm = plt.Normalize(
    0.0, 2.0
)

# data at time t1
t = 1
t1 = np.sin(2 * np.pi * (np.linspace(0,2,len(z)) - 0.01 * t)) + 1
# data as time t20
t = 20
t20 = np.sin(2 * np.pi * (np.linspace(0,2,len(z)) - 0.01 * t)) + 1

cmap = 'viridis'
lc1 = LineCollection(segments, cmap=cmap, norm=norm)
lc1.set_array(t1)
lc1.set_linewidth(50)

lc2 = LineCollection(segments, cmap=cmap, norm=norm)
lc2.set_array(t20)
lc2.set_linewidth(50)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(3,6), tight_layout=True)
ax1.add_collection(lc1)
line2 = ax2.add_collection(lc2)
plt.colorbar(
    line2, 
    shrink=1.0, 
    ax=ax, 
    label="Temperature", 
    location='right', 
    pad=0.05, 
    fraction=0.05, 
    aspect=7
)

ax1.set_xlim(0.4, 0.6)
ax1.set_ylim(51, 60.5)
ax2.set_xlim(0.4, 0.6)
ax2.set_ylim(51, 60.5)

plt.show()

Answer 1

You can use the set_array method that you already found within the animate function that is passed to FuncAnimation:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib import animation

# x position of 1D profile
x = [0.5 for i in np.arange(20)]
# z position of 1D profile
z = [59.95, 59.70, 59.25, 58.75, 59.25,
     58.75, 58.25, 57.75, 57.25, 56.75,
     56.25, 55.75, 55.25, 54.75, 54.25,
     53.75, 53.25, 52.75, 52.25, 51.75]

points = np.array([x, z]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)

norm = plt.Normalize(
    0.0, 2.0
)

# data at time t1
t = 1
data = np.sin(2 * np.pi * (np.linspace(0,2,len(z)) - 0.01 * t)) + 1

cmap = 'viridis'
lc1 = LineCollection(segments, cmap=cmap, norm=norm, array=data, linewidth=50)

fig, ax = plt.subplots(figsize=(3,6), tight_layout=True)
ax.add_collection(lc1)
plt.colorbar(
    lc1, 
    shrink=1.0, 
    label="Temperature", 
    location='right', 
    pad=0.05, 
    fraction=0.05, 
    aspect=7
)

ax.set_xlim(0.4, 0.6)
ax.set_ylim(51, 60.5)

def animate(t):
    # data at time t
    data = np.sin(2 * np.pi * (np.linspace(0,2,len(z)) - 0.01 * t)) + 1    
    lc1.set_array(data)
    
    return lc1,

ani = animation.FuncAnimation(
    fig, animate, interval=20, blit=True, save_count=100)

ani.save('waves.gif')
plt.show()