I have been working on some code that produces a live graphic of a graph and blocks moving. It works in theory but does not look like how I want. The animation runs very slow when showing both the graph and the blocks together (a code of each alone works much better). How can I get them to work together smoothly?

I am also open to any other advice regarding this code on how I can improve its readability, speed, and functionality.

Thanks!

```
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import keyboard
import pandas
# =============================================================================
# Parameters
# =============================================================================
time = 10 # maximum time for the simulation
h = 0.05 # step size
steps = int(time/h) # number of steps
order = 4 # two second order equations
ICs = (0, 0, 0, 1, 0) # intial conditions; t0, x1, x1dot, x2, x2dot
m1 = 1.0 # kg; mass 1
m2 = 3*m1 # kg; mass 2
k = 3.0 # N/m; spring constant
F = 2.8 # N; forcing force
us = 0.5 # static friction coefficient
uk = 0.3 #k inetic friction coefficient
g = 9.81 # m/s^2; gravity
side1 = m1 # side length of mass 1
side2 = m2 # side length of mass 2
parameters = {'m1':m1, 'm2':m2, 'k':k, 'F':F, 'us':us, 'uk':uk, 'g':g}
units = {'m1':'kg', 'm2':'kg', 'k':'N/m', 'F':'N', 'us':'', 'uk':'', 'g':'m/s^2'}
# For logging the data to a CSV
All_Values = {'t':(), 'x1':(), 'v1':(), 'x2':(), 'v2':()}
# =============================================================================
# Intializing Arrays
# =============================================================================
vars = np.empty(shape=(order,steps)) # each row is another var, i.e. x1,x2,...
# Set initial conditions for each var
for i in range(order):
vars(i)(0) = ICs(i+1)
K = np.zeros(shape=(4,order)) # Each row is k1, k2, k3, k4 for each var
t = np.empty(steps)
t(0) = ICs(0)
# =============================================================================
# Resets the arrays
# =============================================================================
def setup():
global K
# Initializing vars array
vars = np.empty(shape=(order,steps)) # each row is another var, i.e. x1,x2,...
# Set initial conditions for each var
for i in range(order):
vars(i)(0) = ICs(i+1)
K = np.zeros(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, var, varsprev):
# variables are as defined in the 'parameters' dict
m1 = parameters('m1')
m2 = parameters('m2')
k = parameters('k')
F = parameters('F')
us = parameters('us')
uk = parameters('uk')
g = parameters('g')
# spring force based on the initial and current length of the spring
l_eq = vars(2)(0) - vars(0)(0) # equilibrium length
l_current = var(2) - var(0)
F_s = k*(l_current - l_eq)
# static friction until it can no longer fight the spring force
static = m1*g*us
kinetic = m1*g*uk
if F_s <= static:
F_f = F_s
else:
F_f = kinetic
dx1dt = var(1) # v1 = x1dot
# Adjust the sign of friction
# c = 0, no friction; c = 1, friction left; c = -1, friction right
c = 1 # assume friction left unless if-else below says otherwise
# if the block is moving to the left, friction to the right
if var(1) < 0:
c = -1
# if the block isn't moving, depends on direction of spring force
elif var(1) == 0:
# spring pushing left, friction to the right
if F_s < 0:
c = -1
# no velocity nor spring force, no friction
elif F_s == 0:
c = 0
dv1dt = 1/m1 * (F_s - c*F_f)
dx2dt = var(3) # v2 = x2dot
# chose that mass 2 does _not_ experience friction
dv2dt = 1/m2 * (-F_s + F)
return(np.array((dx1dt, dv1dt, dx2dt, dv2dt)))
# =============================================================================
# Plotting function
# =============================================================================
plt.ion() # set interactive mode on
fig, (graph, block) = plt.subplots(2, 1, figsize=(8,10))
fig.suptitle(f'Double Spring Mass with Partial Friction Using RK4 (stepsize: {h})',
y=0.94)
def Plotting(i):
# GRAPH
graph.cla() # clear what is currently graphed
graph.plot(t(:i), vars(0,:i), label=f'x1')
graph.plot(t(:i), vars(1,:i), label=f'v1')
graph.plot(t(:i), vars(2,:i), label=f'x2')
graph.plot(t(:i), vars(3,:i), label=f'v2')
graph.annotate(f'time = {round(t(i),1)}s', xy=(0.5, 0.98),
xycoords='axes fraction', ha='center', va='top')
graph.set_title('Graph')
graph.set_xlabel('time (s)')
graph.set_ylabel('position (m)')
graph.legend(loc=2) # upper left
# =========================================================================
# BLOCKS
block.cla() # clear what is currently graphed
m1x = vars(0)(i)
m2x = vars(2)(i)
COM = (m1*(m1x-side1) + m2*(m2x+side2))/(m1 + m2)
windowOffset = abs(max((m1x-side1) - COM, (m2x+side2) - COM, key = abs))+side2
block.set_xlim(-(COM+windowOffset),COM+windowOffset)
# plt.ylim(0,1.5*side2)
mass1 = patches.Rectangle((m1x-side1,0), side1, side1, facecolor='#1f77b4')
mass2 = patches.Rectangle((m2x, 0), side2, side2, facecolor='#ff7f0e')
# Spring color based on length from equilibrium
l_eq = ICs(3) - ICs(1) # equilibrium length
l_curr = m2x - m1x # current length of spring
if (l_curr == l_eq):
c = 'k'
elif (l_curr > l_eq):
c = '#d62728'
else:
c = '#2ca02c'
spring = patches.ConnectionPatch(xyA=(m1x,side1/2), coordsA = 'data',
xyB=(m2x, side1/2),coordsB='data',
linewidth = 2,
color = c)
block.add_patch(spring)
block.add_patch(mass1)
block.add_patch(mass2)
block.axis('equal')
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)')
# =============================================================================
# User parameter change
# =============================================================================
def paramUpdate(i):
update = True
# print parameters and values
def printParams():
print("nCurrent parameters and values: ")
for j in parameters:
print(j, '=', parameters(j), units(j))
changed = False # if a parameter is changed, this ensures plot resets
while update==True:
try:
printParams()
# ask for parameter to change
key = input("Which parameter would you like to change? (or type "n" to resume plotting) ")
# if 'n', resume plotting from the current i
if key == '' or key == 'n':
return (-1 if changed else i)
# check to see if the parameter exists (i.e. key exists)
if key not in parameters.keys():
raise KeyError
# if the value is not convertable to float, throws 'ValueError'
value = float(input("What is the new value? "))
parameters(key) = value # change the parameter value
printParams()
changed = True # a change has been made; plot will reset
while True:
resume = input("Would you like to continue changing parameters? (y/n) ")(0).lower()
if resume == '' or not resume in ('y', 'n'):
print("Please answer with 'y' or 'n'. ")
else:
update = True if resume == 'y' else False
break
# catching errors
except KeyError:
print("nPlease enter a valid parameter.")
except ValueError:
print("nPlease enter a valid value.")
return -1
print('Press "z" to pause')
print('Press "c" to resume')
print('Press "x" to save and quit')
print('Press "v" to adjust parameters')
# =============================================================================
# main loop that calculates each var value using RK 4th order method
# =============================================================================
pause = False
i=0
while i<(steps-1):
if i == 0:
setup()
# calculates each k value
K(0) = h * ODE(t(i), vars(:,i), vars(:,i-1))
K(1) = h * ODE(t(i) + h/2, vars(:,i) + K(0)/2, vars(:,i-1))
K(2) = h * ODE(t(i) + h/2, vars(:,i) + K(1)/2, vars(:,i-1))
K(3) = h * ODE(t(i) + h, vars(:,i) + K(2), vars(:,i-1))
# combines k values using RK4 method
vars(:,i+1) = vars(:,i) + 1/6 * (K(0) + 2 * K(1) + 2 * K(2) + K(3))
# updates time
t(i+1) = t(i) + h
# Plotting every fourth calculation to speed up plotting
if (i%4 == 0):
Plotting(i)
fig.canvas.draw()
plt.pause(0.01)
# exit
if keyboard.is_pressed('x'):
plt.pause(0.01)
break
# update parameters
if keyboard.is_pressed('v'):
i = paramUpdate(i)
# pause and resume
if keyboard.is_pressed('z'):
pause = True
while pause:
if keyboard.is_pressed('c'):
pause = False
if keyboard.is_pressed('x'):
plt.pause(0.01)
break
if keyboard.is_pressed('v'):
i = paramUpdate(i)
pause = False
i+=1
print("Done")
plt.show()
# plt.close() # closes the plot at then end
# Saving values to a CSV
All_Values('t') = t
All_Values('x1') = vars(0)
All_Values('v1') = vars(1)
All_Values('x2') = vars(2)
All_Values('v2') = vars(3)
df = pandas.DataFrame(All_Values)
df.to_csv('Value_Output.csv', index = False)
```
```