I have implemented a simple randomized, population-based optimization method - Grey Wolf optimizer. I am having some trouble with properly capturing the Matplotlib plots at each iteration using the camera package.
I am running GWO for the objective function f(x,y) = x^2 + y^2. I can only see the candidate solutions converging to the minima, but the contour plot doesn't show up.
Do you have any suggestions, how can I display the contour plot in the background?
GWO Algorithm implementation
%matplotlib notebook
import matplotlib.pyplot as plt
import numpy as np
from celluloid import Camera
import ffmpeg
import pillow
# X : Position vector of the initial population
# n : Initial population size
def gwo(f,max_iterations,LB,UB):
fig = plt.figure()
camera = Camera(fig)
def random_population_uniform(m,a,b):
dims = len(a)
x = [list(a + np.multiply(np.random.rand(dims),b - a)) for i in range(m)]
return np.array(x)
def search_agent_fitness(fitness):
alpha = 0
if fitness[1] < fitness[alpha]:
alpha, beta = 1, alpha
else:
beta = 1
if fitness[2] > fitness[alpha] and fitness[2] < fitness[beta]:
beta, delta = 2, beta
elif fitness[2] < fitness[alpha]:
alpha,beta,delta = 2,alpha,beta
else:
delta = 2
for i in range(3,len(fitness)):
if fitness[i] <= fitness[alpha]:
alpha, beta,delta = i, alpha, beta
elif fitness[i] > fitness[alpha] and fitness[i]<= fitness[beta]:
beta,delta = i,beta
elif fitness[i] > fitness[beta] and fitness[i]<= fitness[delta]:
delta = i
return alpha, beta, delta
def plot_search_agent_positions(f,X,alpha,beta,delta,a,b):
# Plot the positions of search agents
x = X[:,0]
y = X[:,1]
s = plt.scatter(x,y,c='gray',zorder=1)
s = plt.scatter(x[alpha],y[alpha],c='red',zorder=1)
s = plt.scatter(x[beta],y[beta],c='blue',zorder=1)
s = plt.scatter(x[delta],y[delta],c='green',zorder=1)
camera.snap()
# Initialize the position of the search agents
X = random_population_uniform(50,np.array(LB),np.array(UB))
n = len(X)
l = 1
# Plot the first image on screen
x = np.linspace(LB[0],LB[1],1000)
y = np.linspace(LB[0],UB[1],1000)
X1,X2 = np.meshgrid(x,y)
Z = f(X1,X2)
cont = plt.contour(X1,X2,Z,20,linewidths=0.75)
while (l < max_iterations):
# Take the x,y coordinates of the initial population
x = X[:,0]
y = X[:,1]
# Calculate the objective function for each search agent
fitness = list(map(f,x,y))
# Update alpha, beta and delta
alpha,beta,delta = search_agent_fitness(fitness)
# Plot search agent positions
plot_search_agent_positions(f,X,alpha,beta,delta,LB,UB)
# a decreases linearly from 2 to 0
a = 2 - l *(2 / max_iterations)
# Update the position of search agents including the Omegas
for i in range(n):
x_prey = X[alpha]
r1 = np.random.rand(2) #r1 is a random vector in [0,1] x [0,1]
r2 = np.random.rand(2) #r2 is a random vector in [0,1] x [0,1]
A1 = 2*a*r1 - a
C1 = 2*r2
D_alpha = np.abs(C1 * x_prey - X[i])
X_1 = x_prey - A1*D_alpha
x_prey = X[beta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A2 = 2*a*r1 - a
C2 = 2*r2
D_beta = np.abs(C2 * x_prey - X[i])
X_2 = x_prey - A2*D_beta
x_prey = X[delta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A3 = 2*a*r1 - a
C3 = 2*r2
D_delta = np.abs(C3 * x_prey - X[i])
X_3 = x_prey - A3*D_delta
X[i] = (X_1 + X_2 + X_3)/3
l = l + 1
return X[alpha],camera
Function call
# define the objective function
def f(x,y):
return x**2 + y**2
minimizer,camera = gwo(f,7,[-10,-10],[10,10])
animation = camera.animate(interval = 1000, repeat = True,
repeat_delay = 500)
Is it possible that the line x = np.linspace(LB[0],LB[1],1000) should be x = np.linspace(LB[0],UB[1],1000) instead? With your current definition of x, x is an array only filled with the value -10 which means that you are unlikely to find a contour.
Another thing that you might want to do is to move the cont = plt.contour(X1,X2,Z,20,linewidths=0.75) line inside of your plot_search_agent_positions function to ensure that the contour is plotted at each iteration of the animation.
Once you make those changes, the code looks like that:
import matplotlib.pyplot as plt
import numpy as np
from celluloid import Camera
import ffmpeg
import PIL
from matplotlib import animation, rc
from IPython.display import HTML, Image # For GIF
from scipy.interpolate import griddata
rc('animation', html='html5')
# X : Position vector of the initial population
# n : Initial population size
def gwo(f,max_iterations,LB,UB):
fig = plt.figure()
fig.gca(aspect='equal')
camera = Camera(fig)
def random_population_uniform(m,a,b):
dims = len(a)
x = [list(a + np.multiply(np.random.rand(dims),b - a)) for i in range(m)]
return np.array(x)
def search_agent_fitness(fitness):
alpha = 0
if fitness[1] < fitness[alpha]:
alpha, beta = 1, alpha
else:
beta = 1
if fitness[2] > fitness[alpha] and fitness[2] < fitness[beta]:
beta, delta = 2, beta
elif fitness[2] < fitness[alpha]:
alpha,beta,delta = 2,alpha,beta
else:
delta = 2
for i in range(3,len(fitness)):
if fitness[i] <= fitness[alpha]:
alpha, beta,delta = i, alpha, beta
elif fitness[i] > fitness[alpha] and fitness[i]<= fitness[beta]:
beta,delta = i,beta
elif fitness[i] > fitness[beta] and fitness[i]<= fitness[delta]:
delta = i
return alpha, beta, delta
def plot_search_agent_positions(f,X,alpha,beta,delta,a,b,X1,X2,Z):
# Plot the positions of search agents
x = X[:,0]
y = X[:,1]
s = plt.scatter(x,y,c='gray',zorder=1)
s = plt.scatter(x[alpha],y[alpha],c='red',zorder=1)
s = plt.scatter(x[beta],y[beta],c='blue',zorder=1)
s = plt.scatter(x[delta],y[delta],c='green',zorder=1)
Z=f(X1,X2)
cont=plt.contour(X1,X2,Z,levels=20,colors='k',norm=True)
plt.clabel(cont, cont.levels, inline=True, fontsize=10)
camera.snap()
# Initialize the position of the search agents
X = random_population_uniform(50,np.array(LB),np.array(UB))
n = len(X)
l = 1
# Plot the first image on screen
x = np.linspace(LB[0],UB[1],1000)
y = np.linspace(LB[0],UB[1],1000)
X1,X2 = np.meshgrid(x,y)
Z=f(X1,X2)
while (l < max_iterations):
# Take the x,y coordinates of the initial population
x = X[:,0]
y = X[:,1]
# Calculate the objective function for each search agent
fitness = list(map(f,x,y))
# Update alpha, beta and delta
alpha,beta,delta = search_agent_fitness(fitness)
# Plot search agent positions
plot_search_agent_positions(f,X,alpha,beta,delta,LB,UB,X1,X2,Z)
# a decreases linearly from 2 to 0
a = 2 - l *(2 / max_iterations)
# Update the position of search agents including the Omegas
for i in range(n):
x_prey = X[alpha]
r1 = np.random.rand(2) #r1 is a random vector in [0,1] x [0,1]
r2 = np.random.rand(2) #r2 is a random vector in [0,1] x [0,1]
A1 = 2*a*r1 - a
C1 = 2*r2
D_alpha = np.abs(C1 * x_prey - X[i])
X_1 = x_prey - A1*D_alpha
x_prey = X[beta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A2 = 2*a*r1 - a
C2 = 2*r2
D_beta = np.abs(C2 * x_prey - X[i])
X_2 = x_prey - A2*D_beta
x_prey = X[delta]
r1 = np.random.rand(2)
r2 = np.random.rand(2)
A3 = 2*a*r1 - a
C3 = 2*r2
D_delta = np.abs(C3 * x_prey - X[i])
X_3 = x_prey - A3*D_delta
X[i] = (X_1 + X_2 + X_3)/3
l = l + 1
return X[alpha],camera
# define the objective function
def f(x,y):
return x**2 + y**2
minimizer,camera = gwo(f,7,[-10,-10],[10,10])
animation = camera.animate(interval = 1000, repeat = True,repeat_delay = 500)
And the output gives:
Related
I want to simulate a data generating process via tensor methods. In the end, the data will be exported to a csv file such that each row corresponds to a time period and each column corresponds to a unit. The following code
import numpy as np
import pandas as pd
import random
import tensorly as tl
from itertools import product
import matplotlib.pyplot as plt
import tensorly.decomposition
np.random.seed(1812)
# Data Generation
L = 0.05
H = 0.05
dx = 0.0025
dy = 0.0025
tmax = 60
dt = 0.01
epsilon = 0.0001
alpha = 0.5e-5+np.random.random()*1e-5
SimulateData = []
SimulateDataNoNoise = []
r_x = alpha*dt/dx**2
r_y = alpha*dt/dy**2
fo = r_x + r_y
if fo > 0.5:
msg = f'Current Fo = {fo}, which is numerically unstable (>0.5)'
raise ValueError(msg)
# x, y meshgrid based on dx, dy
nx = int(L/dx + 1)
ny = int(H/dy + 1)
X, Y = np.meshgrid(np.linspace(0, L, nx), np.linspace(0, H, ny))
# center point of the domain
ic = int((nx-1)/2)
jc = int((ny-1)/2)
# initial and boundary conditions
S = np.zeros((ny, nx))
def enforceBdy(S):
''' Enforces the boundary conditions on S, the temperature values on the domain's grid points'''
S[:, 0] = 1
S[:, -1] = 1
S[0, :] = 1
S[-1, :] = 1
return S
S = enforceBdy(S)
def Laplace(T):
'''Computes the Laplacian operator, del-squared on the data'''
tmp_x, tmp_y = np.gradient(T, dx, dy)
tmp_x, _ = np.gradient(tmp_x, dx)
_, tmp_y = np.gradient(tmp_y, dy)
return tmp_x+tmp_y
# iteration
nmax = int(tmax/dt)
for n in range(nmax):
dSdt = alpha*Laplace(S)
S = S + dSdt*dt
S = enforceBdy(S)
if n % 100 == 0:
noise = np.random.normal(size=S.shape)*.1
SimulateData.append(S.copy()+noise)
SimulateDataNoNoise.append(S.copy())
# check for convergence
err = np.abs(dSdt*dt).max()
if err <= epsilon:
break
#
# Creates Tensor
X = np.stack(SimulateData, 2)
nx,ny,nt = X.shape
# CP Decomposition
err = []
for i in range(1,11):
CP_Heat = tl.decomposition.parafac(X,i)
reconstructed = tl.kruskal_to_tensor(CP_Heat)
err.append(((X-reconstructed)**2).sum())
AIC1 = [2*e + 2*(i+1) for i,e in enumerate(err)]
AIC2 = [2*e + (i+1)*nx+(i+1)*ny+(i+1)*nt for i,e in enumerate(err)]
AIC = AIC2
idxmin = np.argmin(AIC)
R = idxmin+1
min_AIC = AIC[idxmin]
Y = np.zeros((21,40))
beta = np.random.randint(low=-0,high=15,size=21).reshape(-1,1)
for i in range(40):
RHS = 15 + X[:,:,i]#beta + np.random.normal(size=21).reshape(-1,1)
Y[:,i] = RHS.ravel()
Y
np.savetxt("Sim1.csv", Y, delimiter=",")
Returns a CSV file of 21 rows and 40 columns. Suppose, however, I wanted 40 or 70 rows in the final file with 40 columns. How would I do this? When I try with the number 22
import numpy as np
import pandas as pd
import random
import tensorly as tl
from itertools import product
import matplotlib.pyplot as plt
import tensorly.decomposition
np.random.seed(1812)
# Data Generation
L = 0.05
H = 0.05
dx = 0.0025
dy = 0.0025
tmax = 60
dt = 0.01
epsilon = 0.0001
alpha = 0.5e-5+np.random.random()*1e-5
SimulateData = []
SimulateDataNoNoise = []
r_x = alpha*dt/dx**2
r_y = alpha*dt/dy**2
fo = r_x + r_y
if fo > 0.5:
msg = f'Current Fo = {fo}, which is numerically unstable (>0.5)'
raise ValueError(msg)
# x, y meshgrid based on dx, dy
nx = int(L/dx + 1)
ny = int(H/dy + 1)
X, Y = np.meshgrid(np.linspace(0, L, nx), np.linspace(0, H, ny))
# center point of the domain
ic = int((nx-1)/2)
jc = int((ny-1)/2)
# initial and boundary conditions
S = np.zeros((ny, nx))
def enforceBdy(S):
''' Enforces the boundary conditions on S, the temperature values on the domain's grid points'''
S[:, 0] = 1
S[:, -1] = 1
S[0, :] = 1
S[-1, :] = 1
return S
S = enforceBdy(S)
def Laplace(T):
'''Computes the Laplacian operator, del-squared on the data'''
tmp_x, tmp_y = np.gradient(T, dx, dy)
tmp_x, _ = np.gradient(tmp_x, dx)
_, tmp_y = np.gradient(tmp_y, dy)
return tmp_x+tmp_y
# iteration
nmax = int(tmax/dt)
for n in range(nmax):
dSdt = alpha*Laplace(S)
S = S + dSdt*dt
S = enforceBdy(S)
if n % 100 == 0:
noise = np.random.normal(size=S.shape)*.1
SimulateData.append(S.copy()+noise)
SimulateDataNoNoise.append(S.copy())
# check for convergence
err = np.abs(dSdt*dt).max()
if err <= epsilon:
break
#
# Creates Tensor
X = np.stack(SimulateData, 2)
nx,ny,nt = X.shape
# CP Decomposition
err = []
for i in range(1,11):
CP_Heat = tl.decomposition.parafac(X,i)
reconstructed = tl.kruskal_to_tensor(CP_Heat)
err.append(((X-reconstructed)**2).sum())
AIC1 = [2*e + 2*(i+1) for i,e in enumerate(err)]
AIC2 = [2*e + (i+1)*nx+(i+1)*ny+(i+1)*nt for i,e in enumerate(err)]
AIC = AIC2
idxmin = np.argmin(AIC)
R = idxmin+1
min_AIC = AIC[idxmin]
Y = np.zeros((22,40))
beta = np.random.randint(low=-0,high=15,size=22).reshape(-1,1)
for i in range(40):
RHS = 15 + X[:,:,i]#beta + np.random.normal(size=22).reshape(-1,1)
Y[:,i] = RHS.ravel()
Y
np.savetxt("Sim1.csv", Y, delimiter=",")
Python throws an exception saying "(size 22 is different from 21)", but I'm unclear on where the 21 comes from when I do not specify the number 21 anywhere in my code.
Here is my code:
# Libraries
import numpy as np
from scipy.integrate import odeint
from scipy.integrate import quad
# Constant parameters
m_Rb = 1.443*10**-25 #mass of rubidium 87
k_b = 1.38*10**-23
hbar = 1.05*10**-34
L = 38.116*10**6 #spontaneous decay rate
epsilon_0 = 8.85418782*10**-12 #permittivity of free space
# Changable paramaters
lmbda = 700*10**-9 #wavelength of laser light
k = (2*np.pi)/lmbda #wavevector of laser light
V = 1.25*10**-4 #volume of MOT space
length = 5*10**-2 #length of MOT
Bohr = 9.274*10**-24 #Bohr magneton value
B = 5*10**-4 #magnetic field strength
# Scattering force
I_sat = 1 #intensity
n0 = 1 #refraction constant for medium
E_0 = ((2*I_sat)/(3*10**8*n0*epsilon_0))
Rabi = (E_0/hbar) * (4.2275*1.6*10**-19*0.529*10**-10)
n = 20
delta_omega = np.array([-20*10**6, -15*10**6, -10*10**6, -5*10**6])
def F(v):
i = 0
R_i = []
while i<len(delta_omega):
delta = delta_omega[i] + (k*v)
R_scat = L/2 * (Rabi**2/2)/(delta**2+(Rabi**2/2)+(L**2/4)) # scattering rate
R_i.append(R_scat)
i = i+1
R = np.sum(R_i)
R_total = (L*R)/(2*R + L)
def dXdt(t, X):
G, E = X
dGdt = E*(L+R_total) - G*R_total
dEdt = G*R_total - E*(L+R_total)
return [dGdt, dEdt]
t = np.linspace(0, 5, n)
solve = odeint(dXdt, [1, 0], t, tfirst=True)
G = solve.T[0]
E = solve.T[1]
F = hbar*k*(G-E)*R_total
return F
#stopping distance
ds = []
dt = 0.00001
x = np.zeros(n)
time = np.zeros(n)
#v = np.linspace(-700, 700, n)
v = np.zeros(n)
time[0] = 0
x[0] = 0
v[0] = 800
a = 0
print(n)
for a in np.arange(1, n):
time[a]=dt*a
solve = F(v=a)
F_int=solve[1]
v[a] = v[a-1] - dt*(F_int/m_Rb)*x[a-1]
x[a] = x[a-1] + dt*v[a-1]
for s in x:
if abs(s)<=length and abs(s)>0:
ds.append(abs(s))
# Capture velocity
dss = np.max(ds)
I want to find the value of velocity for which ds is at its maximum value (i.e. what's the maximum velocity that a particle can have and not exceed 0.05) but I don't know how to do this. I have calculated the distances travelled by a number of particles, and then made an array (ds) containing the particles which stop within the value of length. dss gives the maximum distance a particle has travelled, but it is the corresponding velocity for this value that I want to find.
I was about to plot a Poincare section of the following DE, which is quite meaningful to have a periodic potential function V(x) = - cos(x) in this equation.
After calculating the solution using RK4 with time interval dt = 0.001, the one that python drew was as the following plot.
But according to the textbook(referred to 2E by J.M.T. Thompson and H.B. Stewart), the section would look like as
:
it has so much difference. For my personal opinion, since Poincare section does not appear as what writers draw, there must be some error in my code. However, I actually done for other forced oscillation DE, including Duffing's equation, and obtained the identical one as those in the textbook. So, I was wodering if there are some typos in the equation given by the textbook, or somewhere else. I posted my code, but might be quite messy to understand. So appreicate dealing with it.
import numpy as np
import matplotlib.pylab as plt
import matplotlib as mpl
import sys
import time
state = [1]
def print_percent_done(index, total, state, title='Please wait'):
percent_done2 = (index+1)/total*100
percent_done = round(percent_done2, 1)
print(f'\t⏳{title}: {percent_done}% done', end='\r')
if percent_done2 > 99.9 and state[0]:
print('\t✅'); state = [0]
####
no = 1
####
def multiple(n, q):
m = n; i = 0
while m >= 0:
m -= q
i += 1
return min(abs(n - (i - 1)*q), abs(i*q - n))
# system(2)
#Basic info.
filename = 'sinPotentialWell'
# a = 1
# alpha = 0.01
# w = 4
w0 = .5
n = 1000000
h = .01
t_0 = 0
x_0 = 0.1
y_0 = 0
A = [(t_0, x_0, y_0)]
def f(t, x, y):
return y
def g(t, x, y):
return -0.5*y - np.sin(x) + 1.1*np.sin(0.5*t)
for i in range(n):
t0 = A[i][0]; x0 = A[i][1]; y0 = A[i][2]
k1 = f(t0, x0, y0)
u1 = g(t0, x0, y0)
k2 = f(t0 + h/2, x0 + h*k1/2, y0 + h*u1/2)
u2 = g(t0 + h/2, x0 + h*k1/2, y0 + h*u1/2)
k3 = f(t0 + h/2, x0 + h*k2/2, y0 + h*u2/2)
u3 = g(t0 + h/2, x0 + h*k2/2, y0 + h*u2/2)
k4 = f(t0 + h, x0 + h*k3, y0 + h*u3)
u4 = g(t0 + h, x0 + h*k3, y0 + h*u3)
t = t0 + h
x = x0 + (k1 + 2*k2 + 2*k3 + k4)*h/6
y = y0 + (u1 + 2*u2 + 2*u3 + u4)*h/6
A.append([t, x, y])
if i%1000 == 0: print_percent_done(i, n, state, 'Solving given DE')
#phase diagram
print('showing 3d_(x, y, phi) graph')
PHI=[[]]; X=[[]]; Y=[[]]
PHI_period1 = []; X_period1 = []; Y_period1 = []
for i in range(n):
if w0*A[i][0]%(2*np.pi) < 1 and w0*A[i-1][0]%(2*np.pi) > 6:
PHI.append([]); X.append([]); Y.append([])
PHI_period1.append((w0*A[i][0])%(2*np.pi)); X_period1.append(A[i][1]); Y_period1.append(A[i][2])
phi_period1 = np.array(PHI_period1); x_period1 = np.array(X_period1); y_period1 = np.array(Y_period1)
print('showing Poincare Section at phi=0')
plt.plot(x_period1, y_period1, 'gs', markersize = 2)
plt.plot()
plt.title('phi=0 Poincare Section')
plt.xlabel('x'); plt.ylabel('y')
plt.show()
If you factor out some of the computation blocks, you can make the code more flexible and computations more direct. No need to reconstruct something if you can construct it in the first place. You want to catch the points where w0*t is a multiple of 2*pi, so just construct the time loops so you integrate in chunks of 2*pi/w0 and only remember the interesting points.
num_plot_points = 2000
h = .01
t,x,y = t_0,x_0,y_0
x_section,y_section = [],[]
T = 2*np.pi/w0
for k in range(num_plot_points):
t = 0;
while t < T-1.2*h:
x,y = RK4step(t,x,y,h)
t += h
x,y = RK4step(t,x,y,T-t)
if k%100 == 0: print_percent_done(k, num_plot_points, state, 'Solving given DE')
x_section.append(x); y_section.append(y)
with RK4step just containing the code of the RK4 step.
This will not solve the mystery. The veil gets lifted if you consider that x is the angle theta (of a forced pendulum with friction) on a circle. Thus to get points with the same spacial location it needs to be reduced by multiples of 2*pi. Doing that,
plt.plot([x%(2*np.pi) for x in x_section], y_section, 'gs', markersize = 2)
results in the expected plot
I don't know that how make the code the three graph in damping harmonic oscillation model,
[X - t(time)], [V(velocity) - t(time)], [a(acceleration) - t(time)] graph
i can make the [X - t(time)] graph
but i don`t know how to make another graphs..
import numpy as np
from matplotlib import pyplot as plt
# mx'' = - bx' - kx
x_0 = 3
v_0 = 0
y_0 = np.array([x_0,v_0]) # first array
def Euler_Method(f,a,b,y0,step):
t = np.linspace(a,b,step)
h = t[1] - t[0]
Y = [y0]
N = len(t)
n = 0
y = y0
for n in range(0,N-1) :
y = y + h*f(y,t[n])
Y.append(y)
n = n+1
Y = np.array(Y)
return Y, t
def harmonic(y,t) :
k = 50
m = 200
b = 20 # drag coefficient
a = (-1*k/m)*y[0] - (b/m)*y[1] # x'' = a, y[0] : first position
v = y[1] # v = first velocity : y[1]
f = np.array([v,a])
return f
a = Euler_Method(harmonic, 0, 100, y_0, 100000)
X = a[0][:,0]
t = a[1]
plt.plot(t,X)
plt.show()
Why can't you just take the derivative of X to get V and A?
V = np.diff(X)
A = np.diff(V)
fig, (ax1, ax2, ax3) = plt.subplots(3)
fig.suptitle('Vertically stacked subplots')
ax1.plot(t, X)
ax2.plot(t[1:], V)
ax3.plot(t[2:], A)
plt.show()
Gives,
I have a complexed valued system from a PDE problem, the odeint() in Python cannot deal with it. I wrote a RK4 module to solve my system. It seems to work, however, the computed values are obvious incorrect. At the second time step, whole computed values are zero. Here are my code:
import matplotlib.pyplot as plt
import numpy as np
import drawnow
import time
import math
### Parameters ###
L = 20
n = 64
delta_t = 1.
tmax = 10
miu = 1e-6
x2 = np.linspace(-L/2,L/2, n+1)
x = x2[:n] # periodic B.C. #0 = #n
kx1 = np.linspace(0,n/2-1,n/2)
kx2 = np.linspace(1,n/2, n/2)
kx2 = -1*kx2[::-1]
kx = (2.*math.pi/L)*np.concatenate((kx1,kx2)); kx[0] = 1e-6
ky = kx; y = x
X, Y = np.meshgrid(x, y)
KX,KY = np.meshgrid(kx,ky)
K = KX**2 + KY**2
K2 = np.reshape(K, n**2,1)
### Initial Condition ###
vorticity = np.exp(-0.25*X**2 - 2.*Y**2)
wt = np.fft.fft2(vorticity)
wt2 = np.reshape(wt, n**2, 1) # wt2 is initial condition
### Define ODE ###
def yprime(t,rhs):
global miu, K, K2,n,KX, KY, wt2, wt
psit = -wt/ K
psix = np.real(np.fft.ifft2(1j*KX*psit))
psiy = np.real(np.fft.ifft2(1j*KY*psit))
wx = np.real(np.fft.ifft2(1j*KX*wt))
wy = np.real(np.fft.ifft2(1j*KY*wt))
rhs = -miu * K2 * wt2 + np.reshape(np.fft.fft2(wx*psiy - wy*psix), n**2,1)
return rhs
def RK4(domain,wt2,tmax):
w = np.empty((tmax+1,n**2))
w = w + 0j
t = np.empty(tmax+1) # length
w[0,:] = wt2 # enter initial conditions in y
t[0] = domain[0]
for i in range(1,tmax):
t[i+1] = t[i]+delta_t
w[i+1,:] = RK4Step(t[i], w[i,:],delta_t)
return w
def RK4Step(t,w,delta_t):
k1 = yprime(t,w)
k2 = yprime(t+0.5*delta_t, w+0.5*k1*delta_t)
k3 = yprime(t+0.5*delta_t, w+0.5*k2*delta_t)
k4 = yprime(t+delta_t, w+k3*delta_t)
return w + (k1+2*k2+2*k3+k4)*delta_t/6.
### Prediction ###
TimeStart = 0.
TimeEnd = tmax+1
TimeSpan = np.arange(TimeStart, TimeEnd, delta_t)
wt2_sol = RK4(TimeSpan, wt2, tmax)
for i in TimeSpan:
w = np.real(np.fft.ifft2(np.reshape(wt2_sol[i,:], (n, n))))
plt.pcolor(X,Y,w,shading = 'interp',cmap='jet')
drawnow
time.sleep(0.2)
plt.show()
Any idea why it doesn't work? In addition, I like to make a short video based on the solution. the function 'drawnow' and 'time.sleep() do not seem to work here.
Thank you!
My cleaned up version. Changing the number of inner steps does not change the quality of the result.
Make the Runge-Kutta solver (more) universal, input time array with (times[0],y0) being the initial point of the IVP
replace def yprime(t,rhs): with def yprime(t,wt):, since wt is your state variable, and rhs is the result. So wt is a local variable in yprime. Eliminate rhs by direct assembly in the return statement.
remove all reshape operation, act on the vector space of 2D arrays, numpy is good in treating matrices as some other kind of vector
add matplotlib.animate for the pre-generated image sequence. The tutorial for that seemed easier than the function-based animation
Played with arange to replace linspace in generation of kx. Better option is probably to use fftshift to swap the halves of the frequency array
.
import numpy as np
import math
from matplotlib import pyplot as plt
from matplotlib import animation
#----- Numerical integration of ODE via fixed-step classical Runge-Kutta -----
def RK4Step(yprime,t,y,dt):
k1 = yprime(t , y )
k2 = yprime(t+0.5*dt, y+0.5*k1*dt)
k3 = yprime(t+0.5*dt, y+0.5*k2*dt)
k4 = yprime(t+ dt, y+ k3*dt)
return y + (k1+2*k2+2*k3+k4)*dt/6.
def RK4(yprime,times,y0):
y = np.empty(times.shape+y0.shape,dtype=y0.dtype)
y[0,:] = y0 # enter initial conditions in y
steps = 4
for i in range(times.size-1):
dt = (times[i+1]-times[i])/steps
y[i+1,:] = y[i,:]
for k in range(steps):
y[i+1,:] = RK4Step(yprime, times[i]+k*dt, y[i+1,:], dt)
return y
#====================================================================
#----- Parameters for PDE -----
L = 20
n = 64
delta_t = 1.
tmax = 10
miu = 1e-6
#----- Constructing the grid and kernel functions
x2 = np.linspace(-L/2,L/2, n+1)
x = x2[:n] # periodic B.C. #0 = #n
y = x
kx = np.linspace( -n/2 , n/2-1, n)
kx = (2.*math.pi/L)*np.concatenate((np.arange(0,n/2),np.arange(-n/2,0)));
kx[0] = 1e-6
ky = kx;
X, Y = np.meshgrid(x, y)
KX,KY = np.meshgrid(kx,ky)
K = KX**2 + KY**2
#----- Initial Condition -----
vorticity = np.exp(-0.25*X**2 - 2.*Y**2)
wt0 = np.fft.fft2(vorticity)
#----- Define ODE -----
def wprime(t,wt):
global miu, K, K2,n,KX, KY
psit = -wt / K
psix = np.real(np.fft.ifft2(1j*KX*psit))
psiy = np.real(np.fft.ifft2(1j*KY*psit))
wx = np.real(np.fft.ifft2(1j*KX*wt))
wy = np.real(np.fft.ifft2(1j*KY*wt))
return -miu * K * wt + np.fft.fft2(wx*psiy - wy*psix)
#====================================================================
#----- Compute the numerical solution -----
TimeStart = 0.
TimeEnd = tmax+delta_t
TimeSpan = np.arange(TimeStart, TimeEnd, delta_t)
wt_sol = RK4(wprime,TimeSpan, wt0)
#----- Animate the numerical solution -----
fig = plt.figure()
ims = []
for i in TimeSpan:
w = np.real(np.fft.ifft2(wt_sol[i,:]))
im = plt.pcolor(X,Y,w,edgecolors='none',cmap='jet')
ims.append([im])
ani = animation.ArtistAnimation(fig, ims, interval=50, blit=True,
repeat_delay=1000)
#ani.save('PDE-animation.mp4')
plt.show()