I am using pseudo-spectral method to solve for KdV equation u_t + u*u_x + u_xxx = 0. After simplification by Fourier Transform, I got two equations with two variables:
uhat = vhat * exp(-i*k^3*t)
d(vhat)/dt =-0.5i * k * exp(-i*k^3*t)*F[u^2]
where F represents Fourier Transform, uhat = F[u], vhat = F[v]
I'd like to use RK4 to solve for u by è ifft(uhat)è eventually. Now I got two variables uhat and vhat, I have 2 ideas for solving uhat and vhat:
To treat them as a system of ODEs to implement RK4.
Treat the equation 2 above as the to-be solved ODE to solve chat by
RK4, and calculate uhat = vhat * exp(-i*k^2*delta_t) after each
time step delta_t.
I got problem to implement both ideas. Here are the codes for the 2nd idea above.
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)
y[i+1,:] = y[i+1,:] * np.exp(delta_t * 1j * kx**3)
return y
#====================================================================
#----- Parameters for PDE -----
L = 100
n = 512
delta_t = 0.1
tmax = 20
c1 = 1.5 # amplitude of 1st wave
c2 = 0.5 # amplitude of 2nd wave
#----- Constructing the grid and kernel functions
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.* np.pi/L)*np.concatenate((kx1,kx2)); kx[0] = 1e-6
#----- Initial Condition -----
z1 = np.sqrt(c1)/2. * (x-0.1*L)
z2 = np.sqrt(c2)/2. * (x-0.4*L)
soliton = 6*(0.5 * c1 / (np.cosh(z1))**2 + 0.5 * c2/ (np.cosh(z2))**2)
uhat0 = np.fft.fft(soliton)
vhat0 = uhat0
#----- Define ODE -----
def wprime(t,vhat):
g = -0.5 * 1j* kx * np.exp(-1j * kx**3 * t)
return g * np.fft.fft(np.real(np.fft.ifft(np.exp(1j*kx**3*t)*vhat)**2))
#====================================================================
#----- Compute the numerical solution -----
TimeStart = 0.
TimeEnd = tmax+delta_t
TimeSpan = np.arange(TimeStart, TimeEnd, delta_t)
w_sol = RK4(wprime,TimeSpan, vhat0)
#----- Animate the numerical solution -----
fig = plt.figure()
ims = []
for i in TimeSpan:
w = np.real(np.fft.ifft(w_sol[i,:]))
im = plt.plot(x,w)
ims.append([im])
ani = animation.ArtistAnimation(fig, ims, interval=100, blit=False)
plt.show()
The RK4 part was from #LutzL.
In general it looks good. Better use a FuncAnimation, reduces the memory footprint and works:
#----- Animate the numerical solution -----
fig = plt.figure()
ax = plt.axes(ylim=(-5,5))
line, = ax.plot(x, soliton, '-')
NT = TimeSpan.size;
def animate(i):
i = i % (NT + 10)
if i<NT:
w = np.real(np.fft.ifft(w_sol[i,:]))
line.set_ydata(w)
return line,
anim = animation.FuncAnimation(fig, animate, interval=50, blit=True)
plt.show()
However, the first graph remains standing, I do not know why.
Related
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 have a problem with time and path optimization. I couldn't define two objectives for both time and path simultaneously. Python reads the last objective and gives result according to that way.
Could you please help me to solve this optimization problem? Thanks..
import matplotlib.pyplot as plt
from gekko import GEKKO
# Gekko model
m = GEKKO(remote=False)
# Time points
nt = 501 # nt=101
tm = np.linspace(0, 1, nt) # tm = np.linspace(0, 100, nt)
m.time = tm
# Variables
g = m.Const(value=9.80665)
V = m.Const(value=200) # velocity
Xi = m.Var(value=0, lb=-2*np.pi, ub=2*np.pi) # Heading angle value=-np.pi dene
x = m.Var(value=0, lb=-100000, ub=100000) # x position
y = m.Var(value=0, lb=-100000, ub=100000) # y position
pathx = m.Const(value=70000) # intended distance in x direction
pathy = m.Const(value=20000) # intended distance in y direction
p = np.zeros(nt) # final time=1
p[-1] = 1.0
final = m.Param(value=p)
m.options.MAX_ITER = 1000000 # iteration number
# Optimize Final Time
tf = m.FV(value=1.0, lb=0.0001, ub=1000.0)
tf.STATUS = 1
# Controlled parameters
Mu = m.MV(value=0, lb=-1, ub=1) # solver controls bank angle
Mu.STATUS = 1
Mu.DCOST = 1e-3
# Equations
m.Equation(x.dt() == tf * (V * (m.cos(Xi))))
m.Equation(y.dt() == tf * (V * (m.sin(Xi))))
m.Equation(Xi.dt() == tf * (g * m.tan(Mu)) / V )
# Objective Function
w = 1e4
m.Minimize(w * (x * final - pathx) ** 2) # 1D part (x)
m.Minimize(w * (pathy - y * final) ** 2) # 2D part (y)
m.Obj(tf)
'''
Multiple objectives are combined into a single objective with summation in Gekko. If there are multiple separate objectives that should be considered separately then a Pareto front is a possible approach to evaluate the tradeoff (see Section 6.5 of the Optimization book). However, this isn't implemented in Gekko.
One correct is that m.Equation(y * final <= pathy ) should be m.Equation((y-pathy) * final <= 0) to avoid a potential infeasible solution if pathy<0.
import matplotlib.pyplot as plt
from gekko import GEKKO
import numpy as np
# Gekko model
m = GEKKO(remote=False)
# Time points
nt = 51 # nt=101
tm = np.linspace(0, 1, nt) # tm = np.linspace(0, 100, nt)
m.time = tm
# Variables
g = m.Const(value=9.80665)
V = m.Const(value=200) # velocity
Xi = m.Var(value=0, lb=-2*np.pi, ub=2*np.pi) # Heading angle value=-np.pi dene
x = m.Var(value=0, lb=-100000, ub=100000) # x position
y = m.Var(value=0, lb=-100000, ub=100000) # y position
pathx = m.Const(value=70000) # intended distance in x direction
pathy = m.Const(value=20000) # intended distance in y direction
p = np.zeros(nt) # final time=1
p[-1] = 1.0
final = m.Param(value=p)
m.options.MAX_ITER = 1000000 # iteration number
# Optimize Final Time
tf = m.FV(value=1.0, lb=0.0001, ub=1000.0)
tf.STATUS = 1
# Controlled parameters
Mu = m.MV(value=0, lb=-1, ub=1) # solver controls bank angle
Mu.STATUS = 1
Mu.DCOST = 1e-3
# Equations
m.Equation(x.dt() == tf * (V * (m.cos(Xi))))
m.Equation(y.dt() == tf * (V * (m.sin(Xi))))
m.Equation(Xi.dt() == tf * (g * m.tan(Mu)) / V )
m.Equation((x-pathx) * final <= 0)
m.Equation((y-pathy) * final <= 0)
# Objective Function
w = 1e4
m.Minimize(w * (x * final - pathx) ** 2) # 1D part (x)
m.Minimize(w * (y * final - pathy) ** 2) # 2D part (y)
#in here python reads the last objective how can i run this two objectives simultaneously.
#m.Obj(Xi * final)
m.Minimize(tf)
#m.Maximize(0.2 * mass * final) # objective function
m.options.IMODE = 6
#m.options.NODES = 2
#m.options.MV_TYPE = 1
#m.options.SOLVER = 3
# m.open_folder() # to search for infeasibilities
m.solve(disp=False)
print('Final Time: ' + str(tf.value[0]))
tm = np.linspace(0,tf.value[0],nt)
#tm = tm * tf.value[0]
fig, axs = plt.subplots(3)
fig.suptitle('Results')
axs[0].plot(tm, Mu.value, 'b-', lw=2, label=r'$Bank$')
axs[0].legend(loc='best')
axs[1].plot(tm, x.value, 'r--', lw=2, label=r'$x$')
axs[1].legend(loc='best')
axs[2].plot(tm, y.value, 'p-', lw=2, label=r'$y$')
axs[2].legend(loc='best')
plt.xlabel('Time')
#plt.ylabel('Value')
plt.show()
plt.figure()
plt.plot(x.value, y.value)
plt.show()
Can anyone tell me what is wrong with this code? It is from https://jakevdp.github.io/blog/2012/09/05/quantum-python/ .
Everything in it worked out except the title of the plot.I can't figure it out.
It should look like this
but when the code is run, it polts this
Here is the code given:-
"""
General Numerical Solver for the 1D Time-Dependent Schrodinger's equation.
author: Jake Vanderplas
email: vanderplas#astro.washington.edu
website: http://jakevdp.github.com
license: BSD
Please feel free to use and modify this, but keep the above information. Thanks!
"""
import numpy as np
from matplotlib import pyplot as pl
from matplotlib import animation
from scipy.fftpack import fft,ifft
class Schrodinger(object):
"""
Class which implements a numerical solution of the time-dependent
Schrodinger equation for an arbitrary potential
"""
def __init__(self, x, psi_x0, V_x,
k0 = None, hbar=1, m=1, t0=0.0):
"""
Parameters
----------
x : array_like, float
length-N array of evenly spaced spatial coordinates
psi_x0 : array_like, complex
length-N array of the initial wave function at time t0
V_x : array_like, float
length-N array giving the potential at each x
k0 : float
the minimum value of k. Note that, because of the workings of the
fast fourier transform, the momentum wave-number will be defined
in the range
k0 < k < 2*pi / dx
where dx = x[1]-x[0]. If you expect nonzero momentum outside this
range, you must modify the inputs accordingly. If not specified,
k0 will be calculated such that the range is [-k0,k0]
hbar : float
value of planck's constant (default = 1)
m : float
particle mass (default = 1)
t0 : float
initial tile (default = 0)
"""
# Validation of array inputs
self.x, psi_x0, self.V_x = map(np.asarray, (x, psi_x0, V_x))
N = self.x.size
assert self.x.shape == (N,)
assert psi_x0.shape == (N,)
assert self.V_x.shape == (N,)
# Set internal parameters
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.dk = 2 * np.pi / (self.N * self.dx)
# set momentum scale
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.compute_k_from_x()
# variables which hold steps in evolution of the
self.x_evolve_half = None
self.x_evolve = None
self.k_evolve = None
# attributes used for dynamic plotting
self.psi_x_line = None
self.psi_k_line = None
self.V_x_line = None
def _set_psi_x(self, psi_x):
self.psi_mod_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_mod_x * np.exp(1j * self.k[0] * self.x)
* np.sqrt(2 * np.pi) / self.dx)
def _set_psi_k(self, psi_k):
self.psi_mod_k = psi_k * np.exp(1j * self.x[0]
* self.dk * np.arange(self.N))
def _get_psi_k(self):
return self.psi_mod_k * np.exp(-1j * self.x[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.x_evolve = self.x_evolve_half * self.x_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_k = property(_get_psi_k, _set_psi_k)
dt = property(_get_dt, _set_dt)
def compute_k_from_x(self):
self.psi_mod_k = fft(self.psi_mod_x)
def compute_x_from_k(self):
self.psi_mod_x = ifft(self.psi_mod_k)
def time_step(self, dt, Nsteps = 1):
"""
Perform a series of time-steps via the time-dependent
Schrodinger Equation.
Parameters
----------
dt : float
the small time interval over which to integrate
Nsteps : float, optional
the number of intervals to compute. The total change
in time at the end of this method will be dt * Nsteps.
default is N = 1
"""
self.dt = dt
if Nsteps > 0:
self.psi_mod_x *= self.x_evolve_half
for i in xrange(Nsteps - 1):
self.compute_k_from_x()
self.psi_mod_k *= self.k_evolve
self.compute_x_from_k()
self.psi_mod_x *= self.x_evolve
self.compute_k_from_x()
self.psi_mod_k *= self.k_evolve
self.compute_x_from_k()
self.psi_mod_x *= self.x_evolve_half
self.compute_k_from_x()
self.t += dt * Nsteps
######################################################################
# Helper functions for gaussian wave-packets
def gauss_x(x, a, x0, k0):
"""
a gaussian wave packet of width a, centered at x0, with momentum k0
"""
return ((a * np.sqrt(np.pi)) ** (-0.5)
* np.exp(-0.5 * ((x - x0) * 1. / a) ** 2 + 1j * x * k0))
def gauss_k(k,a,x0,k0):
"""
analytical fourier transform of gauss_x(x), above
"""
return ((a / np.sqrt(np.pi))**0.5
* np.exp(-0.5 * (a * (k - k0)) ** 2 - 1j * (k - k0) * x0))
######################################################################
# Utility functions for running the animation
def theta(x):
"""
theta function :
returns 0 if x<=0, and 1 if x>0
"""
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))
######################################################################
# Create the animation
# specify time steps and duration
dt = 0.01
N_steps = 50
t_max = 120
frames = int(t_max / float(N_steps * dt))
# specify constants
hbar = 1.0 # planck's constant
m = 1.9 # particle mass
# specify range in x coordinate
N = 2 ** 11
dx = 0.1
x = dx * (np.arange(N) - 0.5 * N)
# specify potential
V0 = 1.5
L = hbar / np.sqrt(2 * m * V0)
a = 3 * L
x0 = -60 * L
V_x = square_barrier(x, a, V0)
V_x[x < -98] = 1E6
V_x[x > 98] = 1E6
# specify initial momentum and quantities derived from it
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)
# define the Schrodinger object which performs the calculations
S = Schrodinger(x=x,
psi_x0=psi_x0,
V_x=V_x,
hbar=hbar,
m=m,
k0=-28)
######################################################################
# Set up plot
fig = pl.figure()
# plotting limits
xlim = (-100, 100)
klim = (-5, 5)
# top axes show the x-space data
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)|$')
# bottom axes show the k-space data
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)
# call the animator. blit=True means only re-draw the parts that have changed.
anim = animation.FuncAnimation(fig, animate, init_func=init,
frames=frames, interval=30, blit=True)
# uncomment the following line to save the video in mp4 format. This
# requires either mencoder or ffmpeg to be installed on your system
#anim.save('schrodinger_barrier.mp4', fps=15, extra_args=['-vcodec', 'libx264'])
pl.show()
For your ease I mention the lines of this code which I do suspect for the error are line number 238,247,255,286,287,296,297
Thanks in advance
The problem is resolved when blit=False, though it may slow down your animation.
Just quoting from a previous answer:
"Possible solutions are:
Put the title inside the axes.
Don't use blitting"
See: How to update plot title with matplotlib using animation?
You also need ffmpeg installed. There are other answers on stackoverflow that help you through that installation. But for this script, here are my recommended new lines you need to add, assuming you're using Windows:
pl.rcParams['animation.ffmpeg_path'] = r"PUT_YOUR_FULL_PATH_TO_FFMPEG_HERE\ffmpeg.exe"
Writer = animation.writers['ffmpeg']
Then adjust the anim.save() line to:
anim.save('schrodinger_barrier.mp4', writer=Writer(fps=15, codec='libx264'))
I have following python code, and would like to:
Plot the same function in 1 (only one) figure with many different (lets say 4) 'v0' and 'theta' values, each trajectory in a different color.
Make 4 plots in 4 different figures, so that it looks like a square with 4 plots of 4 different 'v0' and 'theta' values
Make a widget to vary the v0 and theta values as the user wants with the mouse.
import numpy as np
import scipy.integrate as integrate
import matplotlib.pyplot as plt
%matplotlib inline
theta = 45.
theta = theta * np.pi/180.
v0 = 20.0
g = 9.81
R = 0.035
m = 0.057
rho = 1.2041
C = 0.5
k = (0.5*np.pi*R**2*C*rho)/m
x0=0
y0=10
vx0 = v0*np.sin(theta)
vy0 =
v0*np.cos(theta)
print(vx0)
print(vy0)
def f_func(X_vek,time):
f = np.zeros(4)
f[0] = X_vek[2]
f[1] = X_vek[3]
f[2] = - k*(f[0]**2 + f[1]**2)**(0.5)*f[0]
f[3] = -g - k*(f[0]**2 + f[1]**2)**(0.5)*f[1]
return f
X0 = [ x0, y0, vx0, vy0]
t0 = 0. tf = 10
tau = 0.05
t = np.arange(t0,tf,tau)
X = integrate.odeint(f_func,X0,t)
x = X[:,0]
y = X[:,1]
vx = X[:,2]
vy = X[:,3]
mask = y >= 0
plt.scatter(x[mask],y[mask])
plt.scatter(x[mask],y[mask])
plt.xlabel('x') plt.ylabel('y') plt.show()
I could do point 1 and 2 of my question with changing the values after plotting, then calculate vx0 and vy0 again and then call the integrate function and finally plot again, but that's kinda weird and not clean. Is there any better way to do that? like an array of different v0 and theta values or something?
Thanks!
Make your code as a function:
def func(theta=45, v0=20):
theta = theta * np.pi/180.
g = 9.81
R = 0.035
m = 0.057
rho = 1.2041
C = 0.5
k = (0.5*np.pi*R**2*C*rho)/m
x0=0
y0=10
vx0 = v0*np.sin(theta)
vy0 = v0*np.cos(theta)
def f_func(X_vek,time):
f0, f1 = X_vek[2:4].tolist()
f2 = - k*(f0**2 + f1**2)**(0.5)*f0
f3 = -g - k*(f0**2 + f1**2)**(0.5)*f1
return [f0, f1, f2, f3]
X0 = [ x0, y0, vx0, vy0]
t0 = 0.
tf = 10
tau = 0.05
t = np.arange(t0,tf,tau)
X = integrate.odeint(f_func,X0,t)
x = X[:,0]
y = X[:,1]
vx = X[:,2]
vy = X[:,3]
mask = y >= 0
return x[mask], y[mask]
then you can plot it with different parameters:
plt.plot(*func())
plt.plot(*func(theta=30))
plt.xlabel('x')
plt.ylabel('y')
plt.show()
I suggest you use Holoviews to make dynamic graph:
import holoviews as hv
hv.extension("bokeh")
hv.DynamicMap(
lambda theta, v0:hv.Curve(func(theta, v0)).redim.range(x=(0, 50), y=(0, 50)),
kdims=[hv.Dimension("theta", range=(0, 80), default=40),
hv.Dimension("v0", range=(1, 40), default=20)])
Here is the result:
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()