0

In previous versions of my code I cleared the plot and replotted the data up until the most recently calculated data point. In an attempt to speed up the plotting, I am trying to switch to using set_data so I am not clearing every step. Unfortunately, I am unable to produce a line with this change.

Below you can see my attempt at using set_data for the upper left graphic, but the rest I left using my original method. I hope to use this for all parts of my figure, but I am not sure if this method works for patches.

Please let me know how I can fix this issue.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.gridspec as gridspec
 
# =============================================================================
# Parameters
# =============================================================================
SIM_TIME = 10
STEP_SIZE = 0.05
steps = int(SIM_TIME/STEP_SIZE)
STEPS_PER_FRAME = 4
order = 4  # Two second order equations
 
ICs = [0., 0., 0., 1., 0.]  # Intial conditions; t0, x1, x1dot, x2, x2dot

parameters = {'a':[1.0,  'unit'],
              'b':[2.0,  'unit'],
              'c':[3.0,  'unit']}

# =============================================================================
# Intializing Arrays
# =============================================================================
x_and_v = np.empty(shape=(order, steps))  # Each row is a var, i.e. x1,x2,...
 
# Set initial conditions for each var
for i in range(order):
    x_and_v[i][0] = ICs[i+1]
    
K = np.empty(shape=(4, order)) # Each row is k1, k2, k3, k4 for each var
t = np.empty(steps)
t[0] = ICs[0]

# =============================================================================
# ODE function
# =============================================================================
def ODE(t, curr):
    dx1dt = parameters['a'][0]*t
    dv1dt = parameters['b'][0]*t
    dx2dt = parameters['a'][0]*t
    dv2dt = parameters['c'][0]*t    
    return np.array([dx1dt, dv1dt, dx2dt, dv2dt])

# =============================================================================
# Runge-Kutta (4th Order) Method
# =============================================================================
def RK4(i):
    # calculates each k value
    K[0] = STEP_SIZE * ODE(t[i], x_and_v[:, i])
    K[1] = STEP_SIZE * ODE(t[i] + STEP_SIZE/2, x_and_v[:, i] + K[0]/2)
    K[2] = STEP_SIZE * ODE(t[i] + STEP_SIZE/2, x_and_v[:, i] + K[1]/2)
    K[3] = STEP_SIZE * ODE(t[i] + STEP_SIZE, x_and_v[:, i] + K[2])
    
    return 1/6 * (K[0] + 2*K[1] + 2*K[2] + K[3])

# =============================================================================
# Plotting function
# =============================================================================
plt.close('all')
plt.ion()
fig = plt.figure(figsize=(10, 12))
gs = gridspec.GridSpec(2, 2)
graph_left_x = fig.add_subplot(gs[0, 0])
graph_right_x = fig.add_subplot(gs[0, 1], sharey=graph_left_x)
block = fig.add_subplot(gs[1, :])

fig.suptitle(f'Title (stepsize: {STEP_SIZE})',
             y=0.94)

graph_left_x_line, = graph_left_x.plot(t[0], x_and_v[0][0], 
                                       label='Position', 
                                       color='#1248a1')

def Plotting(i):
    """Plotting x_and_v with time counter in the top middle"""
    graph_left_x_line.set_data(t[i], x_and_v[0][i])
    graph_left_x.annotate(f'time = {round(t[i],1)}s', xy=(0.5, 0.98), 
                           xycoords='axes fraction', ha='center', va='top')
    graph_left_x.set_title('Left Mass')
    graph_left_x.set_ylabel('position [m]')
    graph_left_x.legend(loc=2)
    graph_left_x.relim()
    
    graph_right_x.cla()
    graph_right_x.plot(t[:i], x_and_v[2, :i], 
                       label='Position',
                       color='#ba7000')
    graph_right_x.annotate(f'time = {round(t[i],1)}s', xy=(0.5, 0.98), 
                           xycoords='axes fraction', ha='center', va='top') 
    graph_right_x.set_title('Right Mass')
    graph_right_x.legend(loc=2)
    
    """Animated blocks and spring with time counter in the top middle"""
    block.cla()
        
    m1x = x_and_v[0][i]
    m2x = x_and_v[2][i]
    
    side1 = parameters['a'][0]
    side2 = parameters['b'][0]
    
    block.set_ylim(0, max(side1, side2))
    
    mass1 = patches.Rectangle(
        (m1x - side1, 0),
        side1, side1,
        facecolor='#1f77b4')
    mass2 = patches.Rectangle(
        (m2x, 0),
        side2, side2,
        facecolor='#ff7f0e')
    
    spring = patches.ConnectionPatch(xyA=(m1x, min(side1, side2)/2),
                                     coordsA='data',
                                     xyB=(m2x, min(side1, side2)/2),
                                     coordsB='data',
                                     linewidth=2, color='k')
    
    block.add_patch(spring)
    block.add_patch(mass1)
    block.add_patch(mass2)
    
    block.annotate(f'time = {round(t[i],1)}s', xy=(0.5, 0.98), 
                   xycoords='axes fraction', ha='center', va='top')
    
    block.set_title('Block Animation')
    block.set_xlabel('Position [m]')
    block.axis('equal')
    
    fig.canvas.draw()
    plt.pause(0.01)


# =============================================================================
# Main loop that calculates each x and v value using RK 4th order method
# =============================================================================
i = 0
while i < (steps-1):
    x_and_v[:, i+1] = x_and_v[:, i] + RK4(i)
    t[i+1] = t[i] + STEP_SIZE
    
    if i % STEPS_PER_FRAME == 0:
        Plotting(i)

    i += 1

print('Done')
 
plt.show()
# plt.close()  # closes the plot at then end

Edit: I have checked this post but did not find the solution to work for me.

jared
  • 4,165
  • 1
  • 8
  • 31
  • You go probably not far enough. The matplotlib plot is a scene graph, so instead of constructing the graph anew for each frame, what should work better is to construct it once and then just manipulate the data of the objects, as you intended with `set_data` for the lines. // Why do you not use the `matplotlib.animate` framework? – Lutz Lehmann Jul 08 '20 at 10:25
  • That was what I was hoping to achieve with `set_data` but I cannot seem to get anything to show up. Regarding `matplotlib.animate`, I am confused about how the function works for me to apply it to this project, which is why I haven't used it. – jared Jul 08 '20 at 13:35

1 Answers1

0

My mistake with this problem was actually a pretty basic mistake and misunderstanding of the set_data function. I thought that I needed to pass into it a new data point, but really you need to pass the entire data set, but with updated points.

So, the code ended up looking as follows:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.gridspec as gridspec
 
# =============================================================================
# Parameters
# =============================================================================
SIM_TIME = 10
STEP_SIZE = 0.05
steps = int(SIM_TIME/STEP_SIZE)
STEPS_PER_FRAME = 4
order = 4  # Two second order equations
 
ICs = [0., 0., 0., 1., 0.]  # Intial conditions; t0, x1, x1dot, x2, x2dot

parameters = {'a':[1.0,  'unit'],
              'b':[2.0,  'unit'],
              'c':[3.0,  'unit']}

# =============================================================================
# Intializing Arrays
# =============================================================================
x_and_v = np.empty(shape=(order, steps))  # Each row is a var, i.e. x1,x2,...
 
# Set initial conditions for each var
for i in range(order):
    x_and_v[i][0] = ICs[i+1]
    
K = np.empty(shape=(4, order)) # Each row is k1, k2, k3, k4 for each var
t = np.empty(steps)
t[0] = ICs[0]

# =============================================================================
# ODE function
# =============================================================================
def ODE(t, curr):
    dx1dt = parameters['a'][0]*t
    dv1dt = parameters['b'][0]*t
    dx2dt = parameters['a'][0]*t
    dv2dt = parameters['c'][0]*t    
    return np.array([dx1dt, dv1dt, dx2dt, dv2dt])

# =============================================================================
# Runge-Kutta (4th Order) Method
# =============================================================================
def RK4(i):
    # calculates each k value
    K[0] = STEP_SIZE * ODE(t[i], x_and_v[:, i])
    K[1] = STEP_SIZE * ODE(t[i] + STEP_SIZE/2, x_and_v[:, i] + K[0]/2)
    K[2] = STEP_SIZE * ODE(t[i] + STEP_SIZE/2, x_and_v[:, i] + K[1]/2)
    K[3] = STEP_SIZE * ODE(t[i] + STEP_SIZE, x_and_v[:, i] + K[2])
    
    return 1/6 * (K[0] + 2*K[1] + 2*K[2] + K[3])

# =============================================================================
# Plotting function
# =============================================================================
plt.close('all')
plt.ion()
fig = plt.figure(figsize=(10, 12))

fig.suptitle(f'Title (stepsize: {STEP_SIZE})',
             y=0.94)

gs = gridspec.GridSpec(2, 2)
graph_left = fig.add_subplot(gs[0, 0])
graph_right = fig.add_subplot(gs[0, 1], sharey=graph_left)

graph_left_x_line, = graph_left.plot(np.array([]), np.array([]), 
                                       label='Position', 
                                       color='#1248a1')
graph_left_v_line, = graph_left.plot(np.array([]), np.array([]), 
                                       label='Velocity', 
                                       color='#77a7f7')
graph_left.set_ylabel('position [m]/velocity [m/s]')
graph_left.set_xlabel('time [s]')
graph_left.legend(loc=2)

graph_right_x_line, = graph_right.plot(np.array([]), np.array([]), 
                                       label='Position', 
                                       color='#ba7000')
graph_right_v_line, = graph_right.plot(np.array([]), np.array([]), 
                                       label='Velocity', 
                                       color='#f7c477')

graph_right.set_xlabel('time [s]')
graph_right.legend(loc=2)

block = fig.add_subplot(gs[1, :])
block.set_title('Block Animation')
block.set_xlabel('Position [m]')

def Plotting(i):
    """Plotting x_and_v with time counter in the top middle"""
    graph_left_x_line.set_data(t[:i], x_and_v[0,:i])
    graph_left_v_line.set_data(t[:i], x_and_v[1,:i])
    graph_left.relim()
    graph_left.autoscale_view()

    graph_right_x_line.set_data(t[:i], x_and_v[2, :i])     
    graph_right_v_line.set_data(t[:i], x_and_v[3, :i])
    graph_right.relim()
    graph_right.autoscale_view()
    
    """Animated blocks and spring with time counter in the top middle"""
    block.cla()
        
    m1x = x_and_v[0][i]
    m2x = x_and_v[2][i]
    
    side1 = parameters['a'][0]
    side2 = parameters['b'][0]
    
    block.set_ylim(0, max(side1, side2))
    
    mass1 = patches.Rectangle(
        (m1x - side1, 0),
        side1, side1,
        facecolor='#1f77b4')
    mass2 = patches.Rectangle(
        (m2x, 0),
        side2, side2,
        facecolor='#ff7f0e')
    
    spring = patches.ConnectionPatch(xyA=(m1x, min(side1, side2)/2),
                                     coordsA='data',
                                     xyB=(m2x, min(side1, side2)/2),
                                     coordsB='data',
                                     linewidth=2, color='k')
    
    block.add_patch(spring)
    block.add_patch(mass1)
    block.add_patch(mass2)
    
    block.annotate(f'time = {round(t[i],1)}s', xy=(0.5, 0.98), 
                   xycoords='axes fraction', ha='center', va='top')
    
    block.set_title('Block Animation')
    block.set_xlabel('Position [m]')
    block.axis('equal')
    
    fig.canvas.draw()
    plt.pause(0.01)


# =============================================================================
# Main loop that calculates each x and v value using RK 4th order method
# =============================================================================
i = 0
while i < (steps-1):
    x_and_v[:, i+1] = x_and_v[:, i] + RK4(i)
    t[i+1] = t[i] + STEP_SIZE
    
    if i % STEPS_PER_FRAME == 0:
        Plotting(i)

    i += 1

print('Done')
 
plt.show()
# plt.close()  # closes the plot at then end
jared
  • 4,165
  • 1
  • 8
  • 31