1 回答

TA貢獻1906條經驗 獲得超10個贊
我在這個問題上的錯誤實際上是一個非常基本的錯誤和對set_data功能的誤解。我以為我需要傳遞一個新的數據點,但實際上你需要傳遞整個數據集,但要有更新的點。
所以,代碼最終看起來如下:
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
添加回答
舉報