Develop Reference

How to draw a set of points using OpenCV Python

I've got a set of points (coordinates X and Y) generated by a mathematical expression and I want to draw the resulting figure in a specific position of the screen (I'd like to determine the position in which to center the drawn figure).
I tried using the following code to test if the formula resulted in the correct figure. But now I need to draw the same contour on a pre-existing image at a specific position.
B = 185
L = 250
W = (L-B)/6
D = (L/2)-L/4
x = np.linspace(-L/2, L/2, 500)
y1 = []
y2 = []
for X in x:
termo1 = sqrt((L**2 - 4*X**2) / (L**2 + 8*W*X + 4*W**2))
termo2 = ((sqrt(5.5*L**2 + 11*L*W + 4*W**2) * (sqrt(3)*B*B - 2*D*sqrt(L**2 + 2*W*L + 4*W**2))
) / (sqrt(3)*B*L*(sqrt(5.5*L**2 + 11*L*W + 4*W**2) - 2*sqrt(L**2 + 2*W*L + 4*W**2))))
termo3 = 1 - sqrt((L*(L**2 + 8*W*X + 4*W**2)) / (2*(L - 2*W)*X**2 +
(L**2 + 8*L*W - 4*W**2)*X + 2*L*W**2 + L**2*W + L**2*W + L**3))
calculo = B/2 * termo1 * (1-termo2 * termo3)
y1.append(calculo)
calculo = -B/2 * termo1 * (1-termo2 * termo3)
y2.append(calculo)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.spines['left'].set_position('center')
ax.spines['bottom'].set_position('zero')
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')
ax.xaxis.set_ticks_position('bottom')
ax.yaxis.set_ticks_position('left')
plt.plot(x, y1, 'r')
plt.plot(x, y2, 'r')
plt.show()

You can do this by creating an onclick event; it will take the mouse cords when click and use them as an offset...I think that's what you are asking for? Though with the current plot x/y limits, it won't show up depending on where you click so I added those in the configuration of the plot.
import numpy as np
from math import sqrt
import matplotlib.pyplot as plt
import os
import matplotlib.image as mpimg
def onclick(event):
global ix, iy
ix, iy = event.xdata, event.ydata
plt.plot(x + ix, y1+ iy, 'r')
plt.plot(x + ix, y2+ iy, 'r')
plt.show()
fig.canvas.mpl_disconnect(cid)
return
B = 185
L = 250
W = (L-B)/6
D = (L/2)-L/4
x = np.linspace(-L/2, L/2, 500)
y1 = []
y2 = []
for X in x:
termo1 = sqrt((L**2 - 4*X**2) / (L**2 + 8*W*X + 4*W**2))
termo2 = ((sqrt(5.5*L**2 + 11*L*W + 4*W**2) * (sqrt(3)*B*B - 2*D*sqrt(L**2 + 2*W*L + 4*W**2))
) / (sqrt(3)*B*L*(sqrt(5.5*L**2 + 11*L*W + 4*W**2) - 2*sqrt(L**2 + 2*W*L + 4*W**2))))
termo3 = 1 - sqrt((L*(L**2 + 8*W*X + 4*W**2)) / (2*(L - 2*W)*X**2 +
(L**2 + 8*L*W - 4*W**2)*X + 2*L*W**2 + L**2*W + L**2*W + L**3))
calculo = B/2 * termo1 * (1-termo2 * termo3)
y1.append(calculo)
calculo = -B/2 * termo1 * (1-termo2 * termo3)
y2.append(calculo)
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_xlim([-500,500])
ax.set_ylim([-500,500])
ax.spines['left'].set_position('center')
ax.spines['bottom'].set_position('zero')
ax.spines['right'].set_color('none')
ax.spines['top'].set_color('none')
ax.xaxis.set_ticks_position('bottom')
ax.yaxis.set_ticks_position('left')
coords = []
cid = fig.canvas.mpl_connect('button_press_event', onclick)
plt.show()

I am not able to get spring in the simulation

I am simulating spring-mass system. This is the code.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
k = 1
m = 1
def f(x, u, t):
return -k/m*x
x_graph = []
t_graph = []
u_graph = []
y_graph = []
def func(x_0, u_0, t_0, h):
for i in range(1, 1000):
m1 = h*u_0
k1 = h*f(x_0, u_0, t_0)
m2 = h*(u_0 + 0.5*k1)
k2 = h*f(x_0+0.5*m1, u_0+0.5*k1, t_0+0.5*h)
m3 = h*(u_0 + 0.5*k2)
k3 = h*f(x_0+0.5*m2, u_0+0.5*k2, t_0+0.5*h)
m4 = h*(u_0 + k3)
k4 = h*f(x_0+m3, u_0+k3, t_0+h)
x_0 += (m1 + 2*m2 + 2*m3 + m4)/6
u_0 += (k1 + 2*k2 + 2*k3 + k4)/6
t_0 += h
x_graph.append(x_0)
t_graph.append(t_0)
u_graph.append(u_0)
y_graph.append(0)
return x_0
print(func(0, 5, np.pi, 0.01))
fig, ax = plt.subplots()
ax.set(xlim=(-5.1, 5.1), ylim=(-0.05,1))
# ax.grid()
def animate(i):
l1.set_data(x_graph[:i],y_graph[:i])
return l1,
l1, = ax.plot([],[], 'o-',markevery=[-1])
ani = animation.FuncAnimation(fig, animate,frames=len(t_graph),interval =5, blit=True)
# plt.show()
ani.save('pandemic.gif', writer='ffmpeg')
I want to have a spring at point (-5,0) which changes its position as the block moves. My current output is this
Edit1:- I have added code which includes data for x_graph, y_graph and t_graph.

Correct answer has already been provided. I am adding a different solution. This solution will suit if the spring is attached on the left to a rigid wall and the ball is moving. The spring stretches and compresses based on the distance between the wall and the ball. This looks like below.
I changed your code very slightly to do this. Please beware that some of the values are hard coded so one has to change them to suit the requirement. It is very easy to change them depending on the x_lim values.
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
k = 1
m = 1
def f(x, u, t):
return -k/m*x
x_graph = []
t_graph = []
u_graph = []
y_graph = []
def func(x_0, u_0, t_0, h):
for i in range(1, 1000):
m1 = h*u_0
k1 = h*f(x_0, u_0, t_0)
m2 = h*(u_0 + 0.5*k1)
k2 = h*f(x_0+0.5*m1, u_0+0.5*k1, t_0+0.5*h)
m3 = h*(u_0 + 0.5*k2)
k3 = h*f(x_0+0.5*m2, u_0+0.5*k2, t_0+0.5*h)
m4 = h*(u_0 + k3)
k4 = h*f(x_0+m3, u_0+k3, t_0+h)
x_0 += (m1 + 2*m2 + 2*m3 + m4)/6
u_0 += (k1 + 2*k2 + 2*k3 + k4)/6
t_0 += h
x_graph.append(x_0)
t_graph.append(t_0)
u_graph.append(u_0)
y_graph.append(0)
return x_0
print(func(0, 5, np.pi, 0.01))
fig, ax = plt.subplots()
ax.set(xlim=(-5.1, 5.1), ylim=(-1,1))
# ax.grid()
def animate(i):
l1.set_data(x_graph[:i],y_graph[:i])
sin_x = []
sin_y = []
for j in np.linspace(-5,x_graph[i],250):
sin_x.append(j)
sin_y.append(0.15*np.sin((j+5)*(2*np.pi)*10/(x_graph[i]+5)))
l2.set_data(sin_x, sin_y)
return l1,l2,
l1, = ax.plot([],[], 'o',markevery=[-1])
l2, = ax.plot([],[], '-')
ani = animation.FuncAnimation(fig, animate, frames=len(t_graph),interval =0.5, blit=True)
# plt.show()
ani.save('pandemic.gif', writer='imagemagic')

I think you are trying to simulate a mass osculating around a point.
This small edit should work
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
k = 1
m = 1
def f(x, u, t):
return -k/m*x
def func(x_0, u_0, t_0, h):
x_graph = []
t_graph = []
u_graph = []
y_graph = []
for i in range(1, 1000):
m1 = h*u_0
k1 = h*f(x_0, u_0, t_0)
m2 = h*(u_0 + 0.5*k1)
k2 = h*f(x_0+0.5*m1, u_0+0.5*k1, t_0+0.5*h)
m3 = h*(u_0 + 0.5*k2)
k3 = h*f(x_0+0.5*m2, u_0+0.5*k2, t_0+0.5*h)
m4 = h*(u_0 + k3)
k4 = h*f(x_0+m3, u_0+k3, t_0+h)
x_0 += (m1 + 2*m2 + 2*m3 + m4)/6
u_0 += (k1 + 2*k2 + 2*k3 + k4)/6
t_0 += h
x_graph.append(x_0)
t_graph.append(t_0)
u_graph.append(u_0)
y_graph.append(0)
return t_graph, x_graph, y_graph, u_graph
t, x, y, u = func(0, 5, np.pi, 0.01)
fig, ax = plt.subplots()
ax.set(xlim=(-5.1, 5.1), ylim=(-0.05,1))
# ax.grid()
def animate(i):
i, j = [0, x[i]], [0, y[i]]
l1.set_data(i, j)
return l1,
l1, = ax.plot([],[], 'o-',markevery=[-1])
l2, = ax.plot([],[], 'o-',markevery=[-1])
l3, = ax.plot([],[], 'o-',markevery=[-1])
ani = animation.FuncAnimation(fig, animate,frames=len(t),interval =5, blit=True)
fig.show()
ani.save('pandemic.mp4', writer='ffmpeg')

I need help figuring out how to get the code to give a triple gaussian distribution

So the purpose of my code is to use inputted data points to give a gaussian plot distribution. I figured out how to make it work with a double gaussian but I'm having a lot of trouble adding a third. Im not quite sure what I'm doing wrong. 1 of the errors I keep getting is an Index Error saying that the list index is out of range. I would appreciate any help with this.
Heres my code:
from pylab import *
import numpy as np
from numpy import loadtxt
from scipy.optimize import leastsq
from scipy.optimize import least_squares
from scipy.stats import iqr
import math
import matplotlib.pyplot as plt
import sys
matplotlib.rcParams['mathtext.default'] = 'regular'
fitfunc_triple = lambda p, x: np.abs(p[0]) * exp(-0.5 * ((x - p[1]) / p[2]) ** 2) + np.abs(p[3]) * exp(
-0.5 * ((x - p[4]) / p[5]) ** 2) + np.abs(p[6]) * exp(-0.5 * ((x - p[7])/ p[8] **2 ))
fitfunc_double = lambda p, x: np.abs(p[0]) * exp(-0.5 * ((x - p[1]) / p[2]) ** 2) + np.abs(p[3]) * exp(
-0.5 * ((x - p[4]) / p[5]) ** 2)
fitfunc_single = lambda p, x: np.abs(p[0]) * exp(-0.5 * ((x - p[1]) / p[2]) ** 2)
errfunc = lambda p, x, y: (y - fitfunc(p, x))
dataR = np.loadtxt("/Users/Safi/Library/Preferences/PyCharmCE2018.1/scratches/rspecial1385.a2261.dat5", skiprows=0)
RA = dataR[:, 0]
DEC = dataR[:, 1]
VELR = dataR[:, 2]
REDSH = dataR[:, 3]
RADD = dataR[:, 4]
sl = 3E5
zbar = np.mean(REDSH)
vc = zbar * sl
VEL = vc + sl * ((REDSH - zbar) / (1 + zbar))
wdith = 200
iters = 10
sig2 = 500
binN = int(math.ceil((np.max(VEL) - np.min(VEL)) / wdith))
sys.stdout = open(str(wdith) + "_1sigma_" + str(iters) + "_sig2_" + str(sig2) + ".txt", "w")
plt.figure(1)
y, x, _ = plt.hist(VEL, binN, alpha=0.5, label='data')
x = (x[1:] + x[:-1]) / 2 # for len(x)==len(y)
data = np.vstack((x, y)).T
xdata = data[:, 0]
ydata = data[:, 1]
yerr = ydata ** 0.5
init = [10, 69500, 1200, 5, 68000, sig2]
bds = ([0, 66000, 800, 0, 66000, sig2], [50, 70000, 1750, 20, 70000, sig2 + 0.01])
def index_outlier(data):
inter_quart = iqr(data) * 1.5
bd2 = np.percentile(data, 75) + inter_quart
bd1 = np.percentile(data, 25) - inter_quart
index = []
for i in [i for i, x in enumerate(data) if x < bd1 or x > bd2]:
index.append(i)
return (index)
#### Bootstrapping Estimation Function ####
def fit_bootstrap(fitfunc, datax, datay, init, bds, sigma, iterations=iters):
errfunc = lambda p, x, y: (y - fitfunc(p, x))
# Fit first time
pfit = least_squares(errfunc, init, bounds=bds, args=(datax, datay), max_nfev=10000)
model = fitfunc(pfit.x, datax)
residuals = pfit.fun
# Random data sets are generated and fitted
ps = []
for i in range(iterations):
randomdataY = []
for k in range(len(sigma)):
randomDelta = np.random.normal(0., sigma[k], 1)
randomdataY.append(datay[k] + randomDelta)
out = np.concatenate(randomdataY)
randomfit = least_squares(errfunc, init, bounds=bds, args=(datax, out))
ps.append(randomfit.x)
# Removing outliers
# Finding outliers and indexing them
master_list = []
indexed = []
for k in range(len(ps[0])): # 0-6
it = []
for i in range(len(ps)): # 0-1000
it.append(ps[i][k])
master_list.append(it)
# indexed.append(index_outlier(master_list[k]))
# # List of outlier indicies
# flat_list=[item for sublist in indexed for item in sublist]
# no_dups= list(set(flat_list))
# # Removing bad fits
# for k in range(len(master_list)):
# for i in sorted(no_dups,reverse=True):
# del master_list[k][i]
pfit_bootstrap = []
perr_bootstrap = []
for i in master_list:
pfit_bootstrap.append(np.median(i))
perr_pos = np.round(np.percentile(i, 84) - np.median(i), 4)
perr_neg = np.round(np.median(i) - np.percentile(i, 16), 4)
perr_bootstrap.append(str('[+') + str(perr_pos) + str(',-') + str(perr_neg) + str(']'))
return (pfit_bootstrap, perr_bootstrap, residuals, pfit.nfev, master_list)
pfit, perr, residuals, nfev, master_list = fit_bootstrap(fitfunc_double, xdata, ydata, init, bds, yerr)
pfit1, perr1, residuals1, nfev1, master_list1 = fit_bootstrap(fitfunc_single, xdata, ydata, init, bds, yerr)
more_data = np.linspace(np.min(xdata), np.max(xdata), 1000)
real_func = fitfunc_double(pfit, more_data)
real_func1 = fitfunc_single(pfit1, more_data)
######## Saving Coefficients #########
A1 = pfit[0]
m1 = pfit[1]
s1 = pfit[2]
A2 = pfit[3]
m2 = pfit[4]
s2 = pfit[5]
A3 = pfit[6]
m3 = pfit[7]
s3 = pfit[8]
pecp = VEL - vc
m1p = m1 - vc
m2p = m2 - vc
m3p = m3 - vc
xdatap = xdata - vc
plt.figure(6)
plt.hist(pecp, binN, alpha=.5, label='data', color='skyblue')
xhmax = np.amax(pecp + 1500)
xhmin = np.amin(pecp - 1500)
xh = np.linspace(xhmin, xhmax, 50)
# yh1=(mlab.normpdf(xh, c[1], c[2]))
yh1 = np.abs(A1) * exp(-0.5 * (((xh - m1p) / (s1)) ** 2))
yh2 = np.abs(A2) * exp(-0.5 * (((xh - m2p) / (s2)) ** 2))
yh3 = np.abs(A3) * exp(-0.5 * (((xh - m3p) / (s3)) ** 2))
plt.plot(xh, yh1, color='b', linewidth=2)
plt.plot(xh, yh2, color='r', linewidth=2)
plt.plot(xh, yh3, color='g', linewidth=2)
plt.plot(xh, yh1 + yh2 + yh3, color='purple', linewidth=3)
# plt.errorbar(xdatap,y,xerr=wdith/2,ls='none', yerr=yerr,color='k',linewidth=2)
# plt.plot(xdatap, ydata,'.',color='k')
plt.ylim(0, np.max(ydata) + 2)
plt.xlabel('Peculiar Velocity (km/s)')
plt.ylabel('N$_{gal}$')
plt.text(-4800, 15, '$\mu_{2}$-$\mu_{1}$ = ' + str(int(m2 - m1)) + ' km/s')
plt.savefig(str(wdith) + "_1sigma_" + str(iters) + "_sig2_" + str(sig2) + "_hist.ps")
divi = -1800
memlow = np.array([[0 for x in range(2)] for y in range(1)])
memhigh = np.array([[0 for x in range(2)] for y in range(1)])
j = 0
k = 0
plt.show()

split step Fourier method in Python

Sorry that this post is long but I am trying to simulate two dimensional Schrodinger equation in python using split-step method.
One dimensional problem of this equation has been explained in this post:
https://jakevdp.github.io/blog/2012/09/05/quantum-python/
I tried adding another dimension and modified the operators following the above post but I am confused as how to plot the psi function now that I added another dimension, here is my modified class equation by adding and extra dimension, I also modified the helper functions now that I added extra dimension:
import numpy as np
from scipy.fftpack import fft, ifft
class Schrodinger(object):
def __init__(self, x,y,psi_x0,psi_y0,V_x,V_y,k0=None, hbar=1, m=1, t0=0.0):
self.x,self.y,psi_x0,psi_y0,self.V_x,V_y=map(np.asarray(x,y,psi_x0,psi_y0,V_x,V_y))
N=self.x.size
assert self.x.shape==(N,)
assert self.y.shape==(N,)
assert psi_x0.shape==(N,)
assert self.V_x.shape==(N,)
assert self.V_y.shape==(N,)
self.hbar = hbar
self.m = m
self.t = t0
self.dt_ = None
self.N = len(x)
self.dx=self.x[1]-self.x[0]
self.dy=self.y[1]-self.y[0]
self.dk = 2 * np.pi / (self.N * self.dx)
if k0 == None:
self.k0 = -0.5 * self.N * self.dk
else:
self.k0 = k0
self.k = self.k0 + self.dk * np.arange(self.N)
self.psi_x = psi_x0
self.psi_y=psi_y0
self.compute_k_from_x()
self.compute_k_from_y()
self.x_evolve_half = None
self.x_evolve = None
self.y_evolve_half = None
self.y_evolve = None
self.k_evolve = None
self.psi_x_line = None
self.psi_y_line = None
self.psi_k_line = None
self.V_x_line = None
self.V_y_line = None
def _set_psi_x(self,psi_x):
self.psi_mode_x=(psi_x*np.exp(-1j*self.k[0]*self.x)*self.dx/np.sqrt(2*np.pi))
def _get_psi_x(self):
return (self.psi_mode_x*np.exp(1j*self.k[0])*np.sqrt(2*np.pi)/self.dx)
def _set_psi_y(self,psi_y):
self.psi_mode_y=(psi_y*np.exp(-1j*self.k[0]*self.y)*self.dy/np.sqrt(2*np.pi))
def _get_psi_y(self):
return (self.psi_mode_y*np.exp(1j*self.k[0])*np.sqrt(2*np.pi)/self.dy)
def _set_psi_k_x(self,psi_k_x):
self.psi_mode_k_x=psi_k_x*np.exp(1j*self.x[0]*self.dk*np.arange(self.N))
def _get_psi_k_x(self):
return self.psi_mode_k_x*np.exp(-1j*self.x[0]*self.dk*np.arange(self.N))
def _set_psi_k_y(self,psi_k_y):
self.psi_mode_k_y=psi_k_y*np.exp(1j*self.y[0]*self.dk*np.arange(self.N))
def _get_psi_k_y(self):
return self.psi_mode_k_y*np.exp(-1j*self.y[0]*self.dk*np.arange(self.N))
def _get_dt(self):
return self.dt_
def _set_dt(self,dt):
if dt!=self.dt_:
self.dt_=dt
self.x_evolve_half=np.exp(-0.5*1j*self.V_x/self.hbar*dt)
self.y_evolve_half=np.exp(-0.5*1j*self.V_y/self.hbar*dt)
self.x_evolve=self.x_evolve_half*self.x_evolve_half
self.y_evolve=self.y_evolve_half*self.y_evolve_half
self.k_evolve=np.exp(-0.5*1j*self.hbar/self.m*(self.k*self.k)*dt)
psi_x=property(_get_psi_x,_set_psi_x)
psi_y=property(_get_psi_y,_set_psi_y)
psi_k_x=property(_get_psi_k_x,_set_psi_k_x)
psi_k_y=property(_get_psi_k_y,_set_psi_k_y)
dz=property(_get_dt,_set_dt)
def compute_k_from_x(self):
self.psi_mode_k_x=fft(self.psi_mode_x)
def compute_k_from_y(self):
self.psi_mode_k_y=fft(self.psi_mode_y)
def compute_x_from_k(self):
self.psi_mode_x=ifft(self.psi_mode_k_x)
def compute_y_from_k(self):
self.psi_mode_y=ifft(self.psi_mode_k_y)
def time_step(self,dt,Nsteps=1):
self.dt=dt
if Nsteps>0:
self.psi_mode_x*=self.x_evolve_half
self.psi_mode_y*=self.y_evolve_half
for i in range(Nsteps-1):
self.compute_k_from_x()
self.compute_k_from_y()
self.psi_mode_k_x*=self.k_evolve
self.psi_mode_k_y*=self.k_evolve
self.compute_x_from_k()
self.compute_y_from_k()
self.psi_mode_x*=self.x_evolve
self.psi_mode_y*=self.y_evolve
self.compute_k_from_x()
self.compute_k_from_y()
self.psi_mode_k_x*=self.k_evolve
self.psi_mode_k_y*=self.k_evolve
self.compute_x_from_k()
self.psi_mode_x*=self.x_evolve_half
self.compute_y_from_k()
self.psi_mode_y*=self.y_evolve_half
self.compute_k_from_x()
self.compute_k_from_y()
self.t+=dt*Nsteps
def gauss_x(x, a, x0, k0):
return ((a*np.sqrt(np.pi))**(-0.5)* np.exp(-0.5*((x-x0)* 1./a)**2 +1j*x*k0))
def gauss_y(y, a, y0, k0):
return ((a*np.sqrt(np.pi))**(-0.5)* np.exp(-0.5*((y-y0)* 1./a)**2 +1j*y*k0))
def gauss_k_x(k,a,x0,k0):
return ((a/np.sqrt(np.pi))**0.5* np.exp(-0.5*(a*(k- k0))** 2- 1j*(k- k0)*x0))
def gauss_k_y(k,a,y0,k0):
return ((a/np.sqrt(np.pi))**0.5* np.exp(-0.5*(a*(k- k0))** 2- 1j*(k- k0)*y0))
def theta(x):
x = np.asarray(x)
y = np.zeros(x.shape)
y[x > 0] = 1.0
return y
def square_barrier(x, width, height):
return height * (theta(x) - theta(x - width))
dt = 0.01
N_steps = 50
t_max = 120
frames = int(t_max / float(N_steps * dt))
hbar = 1.0
m = 1.9
N = 2 ** 11
dx = 0.1
dy=0.1
x = dx * (np.arange(N) - 0.5 * N)
y = dy * (np.arange(N) - 0.5 * N)
V0 = 1.5
L = hbar / np.sqrt(2 * m * V0)
a = 3 * L
x0 = -60 * L
y0 = -60 * L
V_x = square_barrier(x, a, V0)
V_y = square_barrier(y, a, V0)
V_x[x < -98] = 1E6
V_x[x > 98] = 1E6
V_y[y < -98] = 1E6
V_y[y > 98] = 1E6
p0 = np.sqrt(2 * m * 0.2 * V0)
dp2 = p0 * p0 * 1./80
d = hbar / np.sqrt(2 * dp2)
k0 = p0 / hbar
v0 = p0 / m
psi_x0 = gauss_x(x, d, x0, k0)
psi_y0 = gauss_y(y, d, y0, k0)
S =Schrodinger(x=x,y=y,psi_x0=psi_x0,psi_y0=psi_y0,V_x=V,V_y=V_y,hbar=hbar,m=m,k0=-28)
the code for one-dimensional plot and the animation is as below:
from matplotlib import pyplot as pl
from matplotlib import animation
fig = pl.figure()
xlim = (-100, 100)
klim = (-5, 5)
ymin = 0
ymax = V0
ax1 = fig.add_subplot(211, xlim=xlim,
ylim=(ymin - 0.2 * (ymax - ymin),
ymax + 0.2 * (ymax - ymin)))
psi_x_line, = ax1.plot([], [], c='r', label=r'$|\psi(x)|$')
V_x_line, = ax1.plot([], [], c='k', label=r'$V(x)$')
center_line = ax1.axvline(0, c='k', ls=':',
label = r"$x_0 + v_0t$")
title = ax1.set_title("")
ax1.legend(prop=dict(size=12))
ax1.set_xlabel('$x$')
ax1.set_ylabel(r'$|\psi(x)|$')
ymin = abs(S.psi_k).min()
ymax = abs(S.psi_k).max()
ax2 = fig.add_subplot(212, xlim=klim,
ylim=(ymin - 0.2 * (ymax - ymin),
ymax + 0.2 * (ymax - ymin)))
psi_k_line, = ax2.plot([], [], c='r', label=r'$|\psi(k)|$')
p0_line1 = ax2.axvline(-p0 / hbar, c='k', ls=':', label=r'$\pm p_0$')
p0_line2 = ax2.axvline(p0 / hbar, c='k', ls=':')
mV_line = ax2.axvline(np.sqrt(2 * V0) / hbar, c='k', ls='--',
label=r'$\sqrt{2mV_0}$')
ax2.legend(prop=dict(size=12))
ax2.set_xlabel('$k$')
ax2.set_ylabel(r'$|\psi(k)|$')
V_x_line.set_data(S.x, S.V_x)
# Animate plot
def init():
psi_x_line.set_data([], [])
V_x_line.set_data([], [])
center_line.set_data([], [])
psi_k_line.set_data([], [])
title.set_text("")
return (psi_x_line, V_x_line, center_line, psi_k_line, title)
def animate(i):
S.time_step(dt, N_steps)
psi_x_line.set_data(S.x, 4 * abs(S.psi_x))
V_x_line.set_data(S.x, S.V_x)
center_line.set_data(2 * [x0 + S.t * p0 / m], [0, 1])
psi_k_line.set_data(S.k, abs(S.psi_k))
title.set_text("t = %.2f" % S.t)
return (psi_x_line, V_x_line, center_line, psi_k_line, title)
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=frames, interval=30, blit=True)
pl.show()

Changing variables in a function to create multiple plots on same axes

I am currently trying to plot a function which describes linear perturbation growth in cosmology for different world models. I would like to be able to have all the curves on the same set of axes, but am struggling with setting it up.
My aim is to plot this function D, with respect to z, but have multiple plots with varying density parameters ($\Omega$).
I have managed two solutions but both aren't working perfectly, the first is very inefficient (adding new functions for each set of parameters):
z = np.arange(0.0,10,0.1)
#density parameters
MOm = 1.0
MOv = 0.0
COm = 0.3
COv = 0.7
H0 = 70
def Mf(z):
A = (5/2)*MOm*(H0**2)
H = H0 * np.sqrt( MOm*((1+z)**3) + MOv )
return A * ((1+z)/(H**3))
def MF(z):
res = np.zeros_like(z)
for i,val in enumerate(z):
y,err = integrate.quad(Mf,val,np.inf)
res[i] = y
return res
def MD(z):
return (H0 * np.sqrt( MOm*((1+z)**3) + MOv )) * MF(z)
def Cf(z):
A = (5/2)*COm*(H0**2)
H = H0 * np.sqrt( COm*((1+z)**3) + COv )
return A * ((1+z)/(H**3))
def CF(z):
res = np.zeros_like(z)
for i,val in enumerate(z):
y,err = integrate.quad(Cf,val,np.inf)
res[i] = y
return res
def CD(z):
return (H0 * np.sqrt( COm*((1+z)**3) + COv )) * CF(z)
plt.plot(z,MD(z),label="Matter Dominated")
plt.plot(z,CD(z),label="Current Epoch")
So I tried to make it simpler with a for loop, but have been unable to work out how to add labels to each plot inside the loop:
Om = (1.0,0.3)
Ov = (0.0,0.7)
for param1,param2 in zip(Om,Ov):
def f(z):
A = (5/2)*param1*(H0**2)
H = H0 * np.sqrt( param1*((1+z)**3) + param2 )
return A * ((1+z)/(H**3))
def F(z):
res = np.zeros_like(z)
for i,val in enumerate(z):
y,err = integrate.quad(f,val,np.inf)
res[i] = y
return res
def D(z):
return (H0 * np.sqrt( param1*((1+z)**3) + param2 )) * F(z)
plt.plot(z,D(z))
Could someone please help explain an efficient method of doing so? Or how to add labels to plots on the fly with a for loop. Any help would be greatly appreciated.

So I tried to make it simpler with a for loop, but have been unable to work out how to add labels to each plot inside the loop
from scipy import integrate
from matplotlib import pyplot as plt
MOm = 1.0
MOv = 0.0
COm = 0.3
COv = 0.7
z = np.arange(0.0,10,0.1)
H0 = 70
Om = (1.0,0.3)
Ov = (0.0,0.7)
fig = plt.figure(1)
for param1,param2 in zip(Om,Ov):
def f(z):
A = (5/2)*param1*(H0**2)
H = H0 * np.sqrt( param1*((1+z)**3) + param2 )
return A * ((1+z)/(H**3))
def F(z):
res = np.zeros_like(z)
for i,val in enumerate(z):
y,err = integrate.quad(f,val,np.inf)
res[i] = y
return res
def D(z):
return (H0 * np.sqrt( param1*((1+z)**3) + param2 )) * F(z)
## Now define labels as you need and labels as follows:
plt.plot(z,D(z),label = 'Om: {}, Ov: {}'.format(param1,param2))
plt.legend()

You can create a label in the loop based on the two parameters using a predefined string which you format with the respective values.
label="Om {}, Ov {}".format(param1, param2)
In total, that would give:
import numpy as np
import scipy.integrate as integrate
import matplotlib.pyplot as plt
z = np.arange(0.0,10,0.1)
MOm = 1.0
MOv = 0.0
COm = 0.3
COv = 0.7
H0 = 70
Om = (1.0,0.3)
Ov = (0.0,0.7)
plt.figure(figsize=(3.8,2.4))
for param1,param2 in zip(Om,Ov):
def f(z):
A = (5/2)*param1*(H0**2)
H = H0 * np.sqrt( param1*((1+z)**3) + param2 )
return A * ((1+z)/(H**3))
def F(z):
res = np.zeros_like(z)
for i,val in enumerate(z):
y,err = integrate.quad(f,val,np.inf)
res[i] = y
return res
def D(z):
return (H0 * np.sqrt( param1*((1+z)**3) + param2 )) * F(z)
plt.plot(z,D(z), label="Om {}, Ov {}".format(param1, param2))
plt.legend()
plt.show()