We Keep Coding

I just get a ball going directly upwards here, what can i change to make it work

import numpy as np
import matplotlib.pyplot as plot
from IPython.display import HTML
from matplotlib import animation
#setup fig with axis
fig, ax = plot.subplots(figsize=(8,8))
#set axis limits
ax.set(xlim=(-2,2), ylim=(0,600), xlabel="position, metres", ylabel="height, metres", title="falling apple")
#initial params
T = 100.
m = 3
g = 9.81
v0x = 10
H = 553.
#setting calc interval
dt = 0.1
N = int(T/dt)
#arrays
v = np.zeros((N+1 , 2))
x = np.zeros((N+1 , 2))
f = np.zeros((N+1 , 2))
#array start [x ,y] format
v[0] = np.array([0. , H])
x[0] = np.array([v0x , 0.])
# the only force is gravity
f[:] = np.array([0., m * g])
#running the dynamics sim
for n in range(N):
v[n+1] = v[n] + ((f[n]/m) * dt)
x[n+1] = x[n] + (v[n+1] * dt)
#scatter plot
scat_plt = ax.scatter(x[0,0], x[0,1], marker='o', c='#1f77b4', s=200)
## animating
def animate(i):
scat_plt.set_offsets(x[i])
ani = animation.FuncAnimation(fig, func=animate, frames=N)
ani.save('ball.html', writer=animation.HTMLWriter(fps= 1//dt))
plot.close()
ani.save('ball.mp4', fps= 1//dt)
HTML('ball.html')
The out put is just a circle going straight up where as this is supposed to simulate a ball being thrown horizontally off a tower
It would be highly appreciated if someone could suggest any changes to be made to the logic/physics or the code.
Thank you!!

I think you mixed x with v at some point. Also the force should be negative in y. I tried this and it seems to work:
import numpy as np
import matplotlib.pyplot as plot
from IPython.display import HTML
from matplotlib import animation
#setup fig with axis
fig, ax = plot.subplots(figsize=(8,8))
#set axis limits
ax.set(xlim=(-200,200), ylim=(0,600), xlabel="position, metres", ylabel="height, metres", title="falling apple")
#initial params
T = 100.
m = 3
g = 9.81
v0x = 10
H = 553.
#setting calc interval
dt = 0.1
N = int(T/dt)
#arrays
v = np.zeros((N+1 , 2))
x = np.zeros((N+1 , 2))
f = np.zeros((N+1 , 2))
#array start [x ,y] format
x[0] = np.array([0. , H])
v[0] = np.array([v0x , 0.])
# the only force is gravity
f[:] = np.array([0., -m * g])
#running the dynamics sim
for n in range(N):
v[n+1] = v[n] + ((f[n]/m) * dt)
x[n+1] = x[n] + (v[n+1] * dt)
#scatter plot
scat_plt = ax.scatter(x[0,0], x[0,1], marker='o', c='#1f77b4', s=200)
## animating
def animate(i):
scat_plt.set_offsets(x[i])
ani = animation.FuncAnimation(fig, func=animate, frames=N)
ani.save('ball.html', writer=animation.HTMLWriter(fps= 1//dt))
plot.close()
ani.save('ball.gif', fps= 1//dt)
HTML('ball.html')

Python different function depending on value (Coulomb damping diagram)

I'm trying to plot a diagram of Coulomb damping mass-spring model, see image bellow. mi*N is the damping friction force.
Becuase the damping force switches direction depending on the direction of speed vector, we have 2 different equations of motion x(t), I named them x_a(t) and x_b(t). If you look closely on the diagram above, the use of x(t) depends on the period, for the first period which is from 0 to pi/omega_n x_a(t) is used, for the second period which is from pi/omega_n to 2*pi/omega_n we use x_b(t) and so on.
I was thinking to use a piece of code which would go something like
t = i * (pi/omega_n)
for i = 1, 3, 5, 7,...
x_a(t)
for i = 2, 4, 6, 8,...
x_b(t)
But I have no clue how to implement this.
I managed to define the x_a(t) part of the code and plot it with the undampled model, see bellow.
import numpy as np
import matplotlib.pyplot as plt
#constants
k = 2 #(N/m), spring coef
m = 0.04 #(kg), mass
x0 = -0.1 #(m), preload
mi = 0.3 #(), dry dynamic friction coef. ABS-ABS
N = 0.3 #(N), normal contact force
f_tr = mi * N / k #friction force/pring coef - equivalent distance
omega_0 = np.sqrt(k/m)
#time
t = np.linspace(0,5,100)
#undamped model
x_undamp = x0*np.cos(omega_0*t)
dx_undamp = -omega_0*x0*np.sin(omega_0*t)
#damped model
x_damp = (x0+f_tr)*np.cos(omega_0*t)-f_tr
dx_damp = -omega_0*(x0+f_tr)*np.sin(omega_0*t)
#time to x=0
t0 = np.arccos(f_tr/(x0+f_tr))/omega_0
print t0
#plotting
fig, (ax1, ax2) = plt.subplots(1, 2)
fig.suptitle('position on left, velocity on right')
ax1.plot(t, x_undamp,'r', label='x_undamp')
ax1.plot(t, x_damp, 'b', label = 'x_damp')
ax2.plot(t, dx_undamp, 'r', label = 'dx_undamp')
ax2.plot(t, dx_damp, 'b', label = 'dx_damp')
#grids, titles, legends, axis labels
ax1.grid()
ax2.grid()
ax1.set_title('Position vs time')
ax2.set_title('Velocity vs time')
ax1.legend()
ax2.legend()
ax1.set_xlabel('t(s)')
ax1.set_ylabel('x(m)')
ax2.set_xlabel('t(s)')
ax2.set_ylabel('dx(m/s)')
plt.show()

I'm gonna put plot a diagram of Coulomb damping mass-spring model solution here till someone comes with a better one:
from scipy.integrate import odeint
import numpy as np
import matplotlib.pyplot as plt
m = 0.2
k = 2.0
c = 0.1
mus = 0.3
muk = 0.2
g = 9.8
vf = 0.01
v0 = 0.0
t1 = 10
sign = lambda x: np.tanh(100*x)
def Xi(t):
if t < 1 :
return 0
else:
return 1
vXi = np.vectorize(Xi)
def eq(X, t, Xi):
Ff = k * (Xi(t) - X[0])
if np.abs(X[1]) < vf and np.abs(Ff) < mus * m * g :
Ff = k * (Xi(t) - X[0])
else:
Ff = sign(X[1]) * muk * m * g
d2x = (k * (Xi(t) - X[0]) - Ff) / m
return [X[1], d2x]
t = np.linspace(0, t1, 1000)
X0 = [v0, 0]
sol = odeint(func = eq, y0 = X0, t = t, args = (Xi, ), mxstep = 50000, atol = 1e-5)
plt.plot(t, sol[:, 0], 'r-', label = 'Output (x(t))')
plt.plot(t, vXi(t), 'b-', label = 'Input (xi(t))')
plt.ylabel('values')
plt.xlabel('time')
plt.legend(loc='best')
plt.show()

Is there a way to write a math formula on matplotlib plot dynamically?

I'm doing a template for my lab work in Python. To summarize the purpose of it, it's to plot data points and fit a pre-defined model with scipy curve_fit. Usually I fit polynomials or exponential curves. I managed to print the fitting params dynamically on the plot, but I have to manually type in the relevant equation every time. I'm wondering, is there an elegant way to do this dynamically? I've read about sympy, but for the time being I couldn't make it.
Here's the code:
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from datetime import datetime
#two example functions
def f(x, p0, p1):
return p0 * x + p1
def g(x, p0, p1):
return p0 * np.exp(x * p1)
#example data
xval = np.array([0,1,2,3,4,5,6])
yval = np.array([0, 2,3.95,5.8,8.1, 10.2, 12.4])
#curve fitting
popt, pcov = curve_fit(f, xval, yval)
plt.rcParams.update({'font.size': 12})
plt.figure(figsize=(9,7))
plt.plot(xval, yval,'ko', label = 'Data points', markersize = 7)
plt.title('TITLE', fontsize = 15)
plt.grid()
plt.plot(xval, f(xval, *popt),'r-', label = 'Fit')
#printing the params on plot
for idx in range(len(popt)):
plt.text(0.8,0.05+0.05*(idx+1), 'p'+str(idx)+' = {0:.5f}'.format(popt[idx]), transform=plt.gca().transAxes)
#manually writing the equation, that's what I want to print dynamically
plt.text(0.8, 0.05, '$y = p0 \cdot x + p1 $' , transform=plt.gca().transAxes)
plt.text(0.86, 1.01, datetime.today().strftime('%Y.%m.%d.'), transform=plt.gca().transAxes)
plt.text(0 ,1.01, 'NAME', transform=plt.gca().transAxes)
plt.ylabel('Y axis title')
plt.xlabel('X axis title')
plt.legend()
plt.show()
The expected result is:
if I use a function for fitting - let's say g(x, p0, p1) which returns p0 * np.exp(x * p1) then the returned formula itself should be printed on the plot, just like the other one in the example code :
plt.text(0.8, 0.05, '$y = p0 \cdot x + p1 $' , transform=plt.gca().transAxes)
except it's a manual solution.
I really appreciate any suggestions.

I think that you may use sympy package.
It allows to define custom variables, create, expressions and then evaluate it. I'm not sure what is impact on performance
Here is your code with changes:
import numpy as np
import sympy
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from datetime import datetime
#two example functions
x, p0, p1 = sympy.var("x p0 p1")
f = p0 * x + p1
g = p0 * sympy.exp(x*p1)
def partial_fun(sympy_expr):
def res_fun(X, P0, P1):
return np.array([sympy_expr.evalf(subs={x: x_, p0: P0, p1: P1}) for x_ in X], dtype=np.float)
return res_fun
#example data
xval = np.array([0,1,2,3,4,5,6])
yval = np.array([0, 2,3.95,5.8,8.1, 10.2, 12.4])
#curve fitting
popt, pcov = curve_fit(partial_fun(f), xval, yval)
plt.rcParams.update({'font.size': 12})
plt.figure(figsize=(9,7))
plt.plot(xval, yval,'ko', label = 'Data points', markersize = 7)
plt.title('TITLE', fontsize = 15)
plt.grid()
plt.plot(xval, partial_fun(f)(xval, *popt),'r-', label = 'Fit')
#printing the params on plot
for idx in range(len(popt)):
plt.text(0.8,0.05+0.05*(idx+1), 'p'+str(idx)+' = {0:.5f}'.format(popt[idx]), transform=plt.gca().transAxes)
#manually writing the equation, that's what I want to print dynamically
plt.text(0.8, 0.05, f'$y = {f} $' , transform=plt.gca().transAxes)
plt.text(0.86, 1.01, datetime.today().strftime('%Y.%m.%d.'), transform=plt.gca().transAxes)
plt.text(0 ,1.01, 'NAME', transform=plt.gca().transAxes)
plt.ylabel('Y axis title')
plt.xlabel('X axis title')
plt.legend()
plt.show()

I actually managed to make a solution (without sympy though), and I have to type in manually the formulas, but they are selected automatically. I use dictionary for that.
Here's the code:
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from datetime import datetime
fun_dict = {}
#three example functions
def f(x, p0, p1):
return p0 * x + p1
def g(x, p0, p1):
return p0 * np.exp(x * p1)
def h(x, p0, p1, p2):
return p0 * x ** 2 + p1 * x + p2
f_string = '$y = p0 \cdot x + p1 $'
fun_dict['f'] = f_string
g_string = '$y = p0 \cdot e^{p1 \cdot x} $'
fun_dict['g'] = g_string
h_string = '$y = p0 \cdot x^2 + p1 \cdot x + p2$'
fun_dict['h'] = h_string
#example data
xval = np.array([0,1,2,3,4,5,6])
yval = np.array([0, 2,3.95,5.8,8.1, 10.2, 12.4])
def get_fun(func):
popt, _ = curve_fit(func, xval, yval)
return popt, fun_dict[str(func.__name__)], func
popt, str_name, func = get_fun(h)
plt.rcParams.update({'font.size': 12})
plt.figure(figsize=(9,7))
plt.plot(xval, yval,'ko', label = 'Data points', markersize = 7)
plt.title('TITLE', fontsize = 15)
plt.grid()
plt.plot(xval, func(xval, *popt),'r-', label = 'Fit')
for idx in range(len(popt)):
plt.text(0.8,0.05+0.05*(idx+1), 'p'+str(idx)+' = {0:.5f}'.format(popt[idx]), transform=plt.gca().transAxes)
plt.text(0.7, 0.05, str_name, transform=plt.gca().transAxes)
plt.text(0.86, 1.01, datetime.today().strftime('%Y.%m.%d.'), transform=plt.gca().transAxes)
plt.text(0 ,1.01, 'NAME', transform=plt.gca().transAxes)
plt.ylabel('Y axis title')
plt.xlabel('X axis title')
plt.legend()
plt.show()

Python Matplotlib: Turning orbit program into 3D

I have finished my orbiting program in 2D, using real physics, and it is pretty cool. One thing I want to do now, is add a "z" force, acceleration, and velocity to the equation. I have done this, but now the program is just a vertical line orbiting a vertical line:
Since theres no code for creating a sphere like there is for circle, I needed to use trig:
Here's the code:
import numpy as np
import matplotlib.pyplot as plt
import math
import matplotlib.animation as animation
import pdb
from mpl_toolkits.mplot3d import Axes3D
er = 6378100*10#m
mr = 1737400*10#m
em = 5.97219*10**24#kg
mm = 7.34767309*10**22#kg
G = 6.67384*10**(-11)
mv = -1023#m/s
nts = 10000
ts = 10000
def simData():
tmax = ts*nts
jx = 0.0
t = 0.0
d = 384400000
mx = d
my = 0
mz = 0
vx = 0
vy = -mv
vz = 0
while t < tmax:
Fg = G*(em*mm)/d**2
Fgx = -Fg*mx/d
Fgy = -Fg*my/d
Fgz = -Fg*mz/d
mAx = Fgx/mm
mAy = Fgy/mm
mAz = Fgz/mm
vx = vx + mAx*ts
vy = vy + mAy*ts
vz = vz + mAz*ts
mx = mx + vx*ts
my = my + vy*ts
mz = mz + vz*ts
u, v = np.mgrid[0:2*np.pi:20j, 0:np.pi:10j]
x=np.cos(u)*np.sin(v)+mx
y=np.sin(u)*np.sin(v)+my
z=np.cos(v)+mz
ax.plot_wireframe(x, y, z, color="grey")
d = math.sqrt(mx**2+my**2+mz**2)
t = t + ts
yield jx, t
def simPoints(simData):
jx, t = simData[0], simData[1]
time_text.set_text(time_template%(t))
line.set_data(t, jx)
return line, time_text
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_aspect("equal")
line, = ax.plot([], [], 'bo', ms=10)
eu, ev = np.mgrid[0:2*np.pi:20j, 0:np.pi:10j]
eax=np.cos(eu)*np.sin(ev)
eay=np.sin(eu)*np.sin(ev)
eaz=np.cos(ev)
ax.plot_wireframe(eax, eay, eaz, color="b")
time_template = 'Time = %.1f s' # prints running simulation time
time_text = ax.text(0.05, 0.9, 0.9,'', transform=ax.transAxes)
ani = animation.FuncAnimation(fig, simPoints, simData, blit=False,\
interval=10, repeat=True)
plt.show()
I have tested the trig for making circles without this program and it works. Also, any suggestions to making surfaced spheres instead of wireframe spheres?

How to show the changes of multiple images in one figure?

My code has been modified according to many great suggestions from people in this forum. However, I still have some questions about the code.
My code is:
from pylab import *
from numpy import *
N = 100 #lattice points per axis
dt = 1 #time step
dx = 1 #lattice spacing
t = arange(0, 1000000*dt, dt) #time
a = 1 #cofficient
epsilon = 100 #cofficient
M = 1.0 #cofficient
every = 100 #dump an image every
phi_0 = 0.5 #initial mean value of the order parameter
noise = 0.1 #initial amplitude of thermal fluctuations in the order parameter
th = phi_0*ones((N, N)) + noise*(rand(N, N) - 0.5) #initial condition
x, y = meshgrid(fftfreq(int(th.shape[0]), dx), fftfreq(int(th.shape[1]), dx))
k2 = (x*x + y*y) #k is a victor in the Fourier space, k2=x^2+y^2
g = lambda th, a: 4*a*th*(1-th)*(1-2*th) #function g
def update(th, dt, a, k2):
return ifft2((fft2(th)-dt*M*k2*fft2(g(th,a)))/(1+2*epsilon*M*dt*k2**2))
for i in range(size(t)):
print t[i]
if mod(i, every)==0:
imshow(abs(th), vmin=0.0, vmax=1.0)
colorbar()
show()
#savefig('t'+str(i/every).zfill(3)+'.png', dpi=100)
clf()
th=update(th, dt, a, k2)
When I run it, I have to close the figures one by one to see the changes. But I want to demonstrate the changes of the images in one figure. Any good ideas?

Use the "animation" feature of matplotlib, like in
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
def update_line(num, data, line):
line.set_data(data[...,:num])
return line,
fig1 = plt.figure()
data = np.random.rand(2, 25)
l, = plt.plot([], [], 'r-')
plt.xlim(0, 1)
plt.ylim(0, 1)
plt.xlabel('x')
plt.title('test')
line_ani = animation.FuncAnimation(fig1, update_line, 25, fargs=(data, l),
interval=50, blit=True)
plt.show()
Tutorial at :
http://matplotlib.org/1.3.1/examples/animation/index.html