Is there a way to numerically solve the following PDE in Python?
The second term on the RHS has a derivative with respect to time as well as space.
I tried using Py-PDE package in Python, it solves only the form dy(x,t)/dt = f(y(x,t)) so I tried to use a root finding algorithm similar to scipy fzero to get the solution to dy(x,t)/dt - f(y(x,t),dy(x,t)/dt) = 0 (solving for dy(x,t)/dt).
class nonlinearPDE(pde.PDEBase):
def __init__(self, bc={"derivative":0}):
self.bc = bc #boundary conditions for operators
def _make_pde_rhs_numba(self, state):
"""numba-compiled implementation of the PDE"""
laplace = state.grid.make_operator("laplace", bc=self.bc)
def findroot(f, df, x0, nmax):
"""Newton–Raphson method"""
for i in range(nmax):
x0 = x0 - f(x0)/df(x0)
return x0
#jit
def pde_rhs(y, t):
func = lambda dydt : dydt - a*laplace(y) - b*laplace(dydt)
dydt = findroot(func, lambda x : 1, 0, 1)
return dydt
return pde_rhs
However, when the program tries to solve the PDE it throws an error:
File "...\anaconda3\lib\site-packages\pde\solvers\controller.py", line 191, in run
t = stepper(state, t, t_break)
File "...\anaconda3\lib\site-packages\pde\solvers\scipy.py", line 82, in stepper
sol = integrate.solve_ivp(
File "...\anaconda3\lib\site-packages\scipy\integrate\_ivp\ivp.py", line 542, in solve_ivp
solver = method(fun, t0, y0, tf, vectorized=vectorized, **options)
File "...\anaconda3\lib\site-packages\scipy\integrate\_ivp\rk.py", line 94, in __init__
self.f = self.fun(self.t, self.y)
File "...\anaconda3\lib\site-packages\scipy\integrate\_ivp\base.py", line 138, in fun
return self.fun_single(t, y)
File "...\anaconda3\lib\site-packages\scipy\integrate\_ivp\base.py", line 20, in fun_wrapped
return np.asarray(fun(t, y), dtype=dtype)
File "...\anaconda3\lib\site-packages\pde\solvers\scipy.py", line 74, in rhs_helper
return rhs(state_flat.reshape(shape), t).flat # type: ignore
File "...\anaconda3\lib\site-packages\numba\core\dispatcher.py", line 420, in _compile_for_args
error_rewrite(e, 'typing')
File "...\anaconda3\lib\site-packages\numba\core\dispatcher.py", line 361, in error_rewrite
raise e.with_traceback(None)
TypingError: Cannot capture the non-constant value associated with variable 'y' in a function that will escape.
Since no one has posted an answer yet, I managed to get a minimal working example by using scipy odeint with a method of lines to solve the PDE, that is, by discretizing the Laplace operator, and then wrapping the differential equation inside fsolve to get dydt:
import numpy as np
from scipy.integrate import odeint
from scipy.optimize import fsolve
a=1
b=1
t = np.linspace(0,1,100)
y0 = np.asarray([1,0.5,0]) #starting values
N=len(y0)
def laplace(y):
""" Laplace operator in cartesian coordinate system """
d2y_0 = [2*(y[1] - y[0])] #at i=0
d2y_N = [2*(y[-2] - y[-1])] #at i=N
d2y_i = [y[i+2] - 2*y[i+1] + y[i] for i in range(N-2)] #elsewhere
return np.asarray(d2y_0 + d2y_i + d2y_N)
def rhs(y, dydt, t):
""" RHS of PDE including dydt term """
return a*laplace(y) + b*laplace(dydt)
def pde(y, t):
""" Function that solves for dydt using root finding """
return fsolve(lambda dydt : dydt - rhs(y, dydt, t),np.zeros(N))
#calculate
sol = odeint(pde, y0, t)
#plot odeint with fsolve
import matplotlib.pyplot as plt
plt.figure()
plt.plot(t,sol)
plt.grid(ls='--')
plt.xlabel('$t$')
plt.ylabel('$y_i$')
plt.title('odeint with fsolve')
plt.legend(["$i=$"+str(i) for i in range(N)])
#plot theoretical
u10=y0[0]
u20=y0[1]
u30=y0[2]
E1=np.exp(-2*a*t/(1+2*b)) #1st exponential term
E2=np.exp(-4*a*t/(1+4*b)) #2nd exponential term
u1 = 1/4*((1+2*E1+E2)*u10 - 2*(E2-1)*u20 + (1-2*E1+E2)*u30)
u2 = 1/4*((1-E2)*u10 + 2*(E2+1)*u20 + (1-E2)*u30)
u3 = 1/4*((1-2*E1+E2)*u10 - 2*(E2-1)*u20 + (1+2*E1+E2)*u30)
plt.figure()
plt.plot(t,u1,t,u2,t,u3)
plt.grid(ls='--')
plt.xlabel('$t$')
plt.ylabel('$y_i$')
plt.title('Theoretical')
plt.legend(["$i=$"+str(i) for i in range(N)])
Note that the same Laplace discretization method allows us to solve a system of ODEs to get the exact analytical solution with which we verify our numerical calculation (for a N=3 size system):
>>> np.allclose(sol[:,0],u1)
True
>>> np.allclose(sol[:,1],u2)
True
>>> np.allclose(sol[:,2],u3)
True
Seems it works as intended.
Related
I am trying to solve a set of nonlinear equations using Sympy. This is my code with a few numbers changed and only one input value (the real code runs through 170,000 lines of data):
import sympy as sp
K0 = 2.00*10**-4
x2, y2 = sp.symbols('x2, y2')
x0 = 500
y0 = 500
R1 = ((x2 - x0)**2 + (y2 - y0)**2)**0.5
R2 = K0 * R1
H2 = sp.atan(R2 * (x2 - x0)/R1)
V2 = sp.atan(R2 * (y2 - y0)/R1)
x, y = sp.symbols('x, y')
x0 = 1.0
y0 = 2.0
x = R1 * H2
y = R1 * V2
dat = sp.nsolve([x - x0, y - y0], [x2, y2], [512, 512]) # This line is the problem
print "dat = %f, %f" % (dat[0], dat[1])
Purely with Python, the code runs fine and produces a good output - but it's very slow. To speed up the process, I used Cython to compile a .pyx file with the exact same code (in a definition, def test()), yet during run-time I'm met with:
File "test.py", line 3, in <module>
demo.test()
File "demo.pyx", line 17, in demo.test
dat = sp.nsolve([x - x0, y - y0], [x2, y2], [512, 512])
File "C:\...\site-packages\sympy\utilities\decorator.py", line 91, in func_wrapper
return func(*args, **kwargs)
File "C:\...\site-packages\sympy\solvers\solvers.py", line 2847, in nsolve
x = findroot(f, x0, J=J, **kwargs)
File "C:\...\site-packages\mpmath\calculus\optimization.py", line 960, in findroot
for x, error in iterations:
File "C:\...\site-packages\mpmath\calculus\optimization.py", line 658, in __iter__
s = self.ctx.lu_solve(Jx, fxn)
File "C:\...\site-packages\mpmath\matrices\linalg.py", line 227, in lu_solve
A, p = ctx.LU_decomp(A)
File "C:\...\site-packages\mpmath\matrices\linalg.py", line 137, in LU_decomp
raise ZeroDivisionError('matrix is numerically singular')
ZeroDivisionError: matrix is numerically singular
I have narrowed down the problem to the x - x0 and y - y0 part. For some reason, the compiled code cannot handle finding the roots when they're not equal to 0. Can nsolve simply not be converted to C using Cython? Is there something I'm missing that I have to do with Cython?
You can use sympy.lambdify together with e.g. SciPy's solvers. If that is not fast enough you could use symengine.Lambdify.
Getting the function signatures right, and deriving the Jacobian, requires you to jump through hoops. If you want to use a library for that I have written pyneqsys:
>>> from pyneqsys.symbolic import SymbolicSys
>>> neqsys = SymbolicSys([x2, y2], [x - x0, y - y0])
>>> neqsys.solve([512, 512])
Out[4]:
(array([ 547.28609349, 594.58064617]),
fjac: array([[ 0.91320338, 0.4075041 ],
[-0.4075041 , 0.91320338]])
fun: array([ -1.37667655e-13, 1.52011737e-12])
message: 'The solution converged.'
nfev: 17
njev: 2
qtf: array([ 1.55620322e-10, 4.63225371e-10])
r: array([ 0.02751454, 0.023682 , 0.03261983])
status: 1
success: True
x: array([ 547.28609349, 594.58064617]))
If those 170 000 solves involves gradually changing parameters pyneqsys can exploit that (by propagating the solution as guess between solves). It can also use symengine.Lambdify automatically by setting the environment variable SYM_BACKEND=sympysymengine
I'm trying to solve the equation Iy'' + b|y'|y' + ky = 0 and fit the coefficients to data.
This is the code I have so far (ready to run):
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
import pandas as pd
from scipy.optimize import curve_fit
# Define derivatives of function
def f(y, t, I, b, k):
theta, omega = y
derivs = [omega, -b / I * omega * abs(omega) - k / I * theta]
return derivs
# integrate the function
def yint(t, I, b, k, y0, y1):
y = odeint(f, [y0, y1], t, args=(I, b, k))
return y.ravel()
# define t and y to fit
y0 = [.5, 0]
t = np.arange(0, 3, .01)
y = np.cos(t)*np.e**(-.01*t)
# fit
vals, cov = curve_fit(yint, t, y, p0=[0.002245, 1e-5, 0.2492, y0[0], y[1]])
However, when I run the function, I get the error:
Traceback (most recent call last):
File "---", line 24, in <module>
vals, cov = curve_fit(yint, t, y, p0=[0.002245, 1e-5, 0.2492, y0[0], y[1]])
File "---.py", line 578, in curve_fit
res = leastsq(func, p0, args=args, full_output=1, **kw)
File "---.py", line 371, in leastsq
shape, dtype = _check_func('leastsq', 'func', func, x0, args, n)
File "---.py", line 20, in _check_func
res = atleast_1d(thefunc(*((x0[:numinputs],) + args)))
File "---.py", line 447, in _general_function
return function(xdata, *params) - ydata
ValueError: operands could not be broadcast together with shapes (600,) (300,)
Any thoughts on how to fix this?
The problem is that function yint returns an array of shape (600,) for the argument of shape (300,). Think again about yint: it solves a second-order differential equation by representing it as a system of two first-order equations. So the result of y = odeint(...) has two columns, one for the solution itself, the second for its derivative. Its shape is (300,2). Mashing the solution and its derivative together with ravel does not make sense. Instead, you should only take the actual solution, that's the thing you are fitting.
So, replace
return y.ravel()
with
return y[:, 0]
I'm trying to calculate the roots for a function using the scipy function fsolve, but an error keeps flagging:
TypeError: 'numpy.array' object is not callable
I assume it's probably easier to define the equation as a function but I've tried that a few times to no avail.
Code:
import scipy
import numpy as np
import matplotlib.pyplot as plt
from scipy import optimize
# Constants
wavelength = 0.6328
ncore = 1.462420
nclad = 1.457420
a = 8.335
# Mode Order
l = 0
# Mode parameters
V = (2 * np.pi * a / wavelength) * np.sqrt(ncore**2 - nclad**2)
U = np.arange(0, V, 0.01)
W = np.sqrt(V**2-U**2)
func = U * scipy.special.jv(l+1, U) / scipy.special.jv(l, U) - W * scipy.special.kv(l+1, W) / scipy.special.kv(l, W)
from scipy.optimize import fsolve
x = fsolve(func,0)
print x
StackTrace:
Traceback (most recent call last):
File "<ipython-input-52-081a9cc9c0ea>", line 1, in <module>
runfile('/home/luke/Documents/PythonPrograms/ModeSolver_StepIndex/ModeSolver_StepIndex.py', wdir='/home/luke/Documents/PythonPrograms/ModeSolver_StepIndex')
File "/usr/lib/python2.7/site-packages/spyderlib/widgets/externalshell/sitecustomize.py", line 580, in runfile
execfile(filename, namespace)
File "/home/luke/Documents/PythonPrograms/ModeSolver_StepIndex/ModeSolver_StepIndex.py", line 52, in <module>
x = fsolve(func,0)
File "/usr/lib64/python2.7/site-packages/scipy/optimize/minpack.py", line 140, in fsolve
res = _root_hybr(func, x0, args, jac=fprime, **options)
File "/usr/lib64/python2.7/site-packages/scipy/optimize/minpack.py", line 197, in _root_hybr
shape, dtype = _check_func('fsolve', 'func', func, x0, args, n, (n,))
File "/usr/lib64/python2.7/site-packages/scipy/optimize/minpack.py", line 20, in _check_func
res = atleast_1d(thefunc(*((x0[:numinputs],) + args)))
TypeError: 'numpy.ndarray' object is not callable
That is because fsolve takes a function as argument.
Try this, Note you still will encounter some runtime error , you will have to check if your return from func is properly constructed, I will leave that for you to figure out.
import scipy
import numpy as np
import matplotlib.pyplot as plt
from scipy import optimize
# Constants
wavelength = 0.6328
ncore = 1.462420
nclad = 1.457420
a = 8.335
# Mode Order
# l = 0
# Mode parameters
V = (2 * np.pi * a / wavelength) * np.sqrt(ncore**2 - nclad**2)
U = np.arange(0, V, 0.01)
W = np.sqrt(V**2-U**2)
def func(l):
return U * scipy.special.jv(l+1, U) / scipy.special.jv(l, U) - W * scipy.special.kv(l+1, W) / scipy.special.kv(l, W)
from scipy.optimize import fsolve
x = fsolve(func,0)
print x
You need to pass a function to fsolve not an array.
If I just print your func:
func
array([ -1.04882076e+01, -1.04881526e+01, -1.04879876e+01,
-1.04877125e+01, -1.04873274e+01, -1.04868321e+01,
-1.04862266e+01, -1.04855109e+01, -1.04846847e+01,
-1.04837481e+01, -1.04827008e+01, -1.04815428e+01,
-1.04802738e+01, -1.04788938e+01, -1.04774024e+01,
-1.04757996e+01, -1.04740850e+01, -1.04722585e+01,
-1.04703198e+01, -1.04682686e+01, -1.04661046e+01,
-1.04638275e+01, -1.04614371e+01, -1.04589330e+01,
-1.04563147e+01, -1.04535820e+01, -1.04507345e+01,
-1.04477718e+01, -1.04446934e+01, -1.04414988e+01,
... ]
that's an array but you want a function. Something like this works:
def linear(x):
return 2*x+4
fsolve(linear, 0)
don't know how one could define your function, though.
I am trying to fit resistivity vs temperature data to Bloch-Gruneisen formula for resistivity in metals:
function
as you can see there is an integral function with a parametric limit. I don't know how to implement an algorithm to run a least squares fit.
I came up with:
import matplotlib.pyplot as plt
import numpy as np
import pylab as pl
import scipy as sp
from scipy.optimize import leastsq
#retrieve data from file
data = pl.loadtxt('salita.txt')
Temp = data[:, 0]
Res = data[:, 2]
def debye_func(p, T, r):
rho0, AD, TD = p
coeff = AD*np.power(T, 5)/np.power(TD, 4)
f = np.power(x^5)/np.power(np.sinh(x), 2) #function to integrate
err_debye = r - rho0 - coeff * #integral???
return err_debye
p0 = sp.array([0.0001 , 0.00001, 50])
plsq = leastsq(debye_func, p0, args=(Temp, Res))
print plsq
Ideas on how could I write it?
EDIT: my code has become:
import matplotlib.pyplot as plt
import numpy as np
import pylab as pl
import scipy as sp
from scipy.optimize import leastsq
from scipy.integrate import quad
#retrieve data from file
data = pl.loadtxt('salita.txt')
Temp = data[:, 0]
Res = data[:, 2]
def debye_integrand(x):
return np.power(x, 5)/np.power(np.sinh(x), 2)
def debye_func(p, T, r):
rho0, AD, TD = p
coeff = AD*np.power(T, 5)/np.power(TD, 4)
err_debye = r - rho0 - coeff * quad(debye_integrand, 0, TD/(2*T))
return err_debye
p0 = sp.array([0.0001 , 0.00001, 50])
plsq = leastsq(debye_func, p0, args=(Temp, Res))
print plsq
Now I get a ValueError:
Traceback (most recent call last):
File "debye.py", line 24, in <module>
plsq = leastsq(debye_func, p0, args=(Temp, Res))
File "/System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python/scipy/optimize/minpack.py", line 348, in leastsq
m = _check_func('leastsq', 'func', func, x0, args, n)[0]
File "/System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python/scipy/optimize/minpack.py", line 14, in _check_func
res = atleast_1d(thefunc(*((x0[:numinputs],) + args)))
File "debye.py", line 19, in debye_func
err_debye = r - rho0 - coeff * quad(debye_integrand, 0, TD/(2*T))
File "/System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python/scipy/integrate/quadpack.py", line 247, in quad
retval = _quad(func,a,b,args,full_output,epsabs,epsrel,limit,points)
File "/System/Library/Frameworks/Python.framework/Versions/2.7/Extras/lib/python/scipy/integrate/quadpack.py", line 296, in _quad
if (b != Inf and a != -Inf):
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
I think that means I'm providing leastsq an argument it can't take, but I don't know how to modify my code.
EDIT2: I solved my equation analytically with Maxima and I got
import matplotlib.pyplot as plt
import numpy as np
import pylab as pl
import scipy as sp
from scipy.optimize import leastsq
from scipy.integrate import quad
from scipy.special import zetac
from mpmath import polylog
#retrieve data from file
data = pl.loadtxt('salita.txt')
Temp = data[:, 0]
Res = data[:, 2]
def debye_integrand(x):
return np.power(x, 5)/np.power(np.sinh(x), 2)
def debye_func(p, T, r, integral):
rho0, AD, TD = p
coeff = AD*np.power(T, 5)/np.power(TD, 4)
den = np.exp(TD/T) -1
m1 = 5*((TD/(2*T))**4)*np.log(np.exp(TD/(2*T)+1)*(np.exp(TD/T)-1)+120*polylog(5, np.exp(TD/(T))*(1-np.exp(TD/(2*T)))
m2 = 120*(TD/(2*T))*polylog(4, np.exp(TD/(2*T)))*(np.exp(np.exp(TD/T))-1)+60*((TD/(2*T))**2)*polylog(3, np.exp(TD/(2*T))*(1-np.exp((TD/(2*T)))
m3 = 20*((TD/(2*T))**3)*polylog(2, np.exp(TD/(2*T))*(np.exp(TD/T)-1)+120**polylog(5, -np.exp(TD/(2*T)))*(1-np.exp(TD/T))
m4 = 120*(TD/(2*T))*polylog(4, -np.exp(TD/(2*T)))*(np.exp(TD/T)-1)+60*((TD/(2*T))**2)*polylog(3, -np.exp(TD/(2*T)))*(1-np.exp(TD/T))
m5 = 20*((TD/(2*T))**3)*polylog(2, -np.exp(TD/(2*T)))*(np.exp(TD/T)-1) -2*((TD/(2*T))**5)*np.exp(TD/T)
m6 = 5*((TD/(2*T))**4)*np.log(1-np.exp(TD/(2*T))*(np.exp(TD/T)-1)
zeta = 15.0*zetac(5)/2
integral = (m1+m2+m3+m4+m5+m6)/den +zeta
err_debye = r - rho0 - coeff * integral
return err_debye
#initizalied with Einstein model fit
p0 = sp.array([0.00001 , 0.0000001, 70.0])
plsq = leastsq(debye_func, p0, args=(Temp, Res))
print plsq
It says SyntaxError: invalid syntax in m2. I tried to do it with loops in the numerical way, but I didn't succeed.
My .txt file is here, if you want to try. First column is temperature, third one is resistivity.
You could, for instance, separately define the integrand function,
def debye_integrand(x, n):
return x**n/((np.exp(x) - 1)*(1 - np.exp(-x)))
and then use scipy.integrate.quad to do this integration numerically,
from scipy.integrate import quad
def debye_func(p, T, r):
# [...] the rest of your code from above here
err_debye = r - rho0 - coeff * quad(debye_integrand, 0, T/TD, args=(n,))
return np.sum(err_debye**2)
That's the general idea, and this might need to be adapted further to your code. An ideal solution would be to find an analytical solution to that integral, or rewrite it with classical integral functions from scipy.special, but it might not be straightforward (see below).
Also you should use the more general scipy.opitimize.minimize function instead of the least-square fit, since it provides algorithms that are more efficient and robust for non-smooth optimizations. The default optimization method BFGS is a good start.
Edit: actually, there is an analytical solution to this integral (for n=5), that you can obtain, for instance, with Maxima,
>> integrate(x**5/((exp(x) - 1)*(1 - exp(-x))), x, 0, a)
where a is the integration limit, li_k the Polylogarithm function of order k (see mpmath.polylog) and ζ is the Riemann Zeta function (see scipy.special.zetac).
Although, depending on your needs, it might be faster to just go with a numerical integration (or pre-calculated table lookup) rather than puyting all of this together, and converting it to python.
Edit 2: Here the final solution with analytical calculation of the integral,
import numpy as np
import mpmath as mp
from scipy.optimize import minimize
from scipy.integrate import quad
import matplotlib.pyplot as plt
def debye_integral_sym_scalar(x):
"""
Calculate the Debye integral for a scalar using multi precision math,
as otherwise it overflows with 64bit floats
"""
exp_x = mp.exp(x)
m1 = -120*mp.polylog(5, exp_x)
m2 = 120*x*mp.polylog(4, exp_x)
m3 = -60*x**2*mp.polylog(3, exp_x)
m4 = 20*x**3*mp.polylog(2, exp_x)
m5 = 5*x**4*mp.log(1 - exp_x)
m6 = - x**5*exp_x
return m1 + m2 + m3 + m4 + m5 + m6/(exp_x - 1) + 120*mp.zeta(5)
# this is the actual function that we can use
def debye_integral_sym(x):
f = np.vectorize(debye_integral_sym_scalar, otypes=[np.complex])
return f(x).real
def debye_integrand(x, n):
return x**n/((np.exp(x) - 1)*(1 - np.exp(-x)))
# test that debye_integral_sym returns the same result as quad
a = 10.0
res0 = quad(debye_integrand, 0, a, args=(5,))[0]
res1 = debye_integral_sym(a)
np.testing.assert_allclose(res0, res1)
def resistivity_fit(p, T):
rho0, AD, TD = p
coeff = AD*np.power(T, 5)/np.power(TD, 4)
return rho0 + coeff * debye_integral_sym(TD/(2*T))
def debye_err_func(p, T, r):
return np.sum((r - resistivity_fit(p, T))**2)
# wget "http://pastebin.com/raw.php?i=tvzcdxYA" -O salita.txt
data = np.loadtxt('salita.txt')
temp_exp = data[:, 0]
res_exp = data[:, 2]
p0 = np.array([0.0001 , 0.00001, 50])
p_opt = minimize(debye_err_func, p0, args=(temp_exp, res_exp))
print p_opt
temp = np.linspace(temp_exp.min(), temp_exp.max(), 100)
plt.plot(temp_exp, res_exp, '.', label='Experimental data')
plt.plot(temp, resistivity_fit(p_opt.x, temp), 'r', label='Bloch-Gruneisen fit')
plt.legend(loc='best')
plt.xlabel('Temperature [K]')
plt.ylabel('Resistivity')
plt.show()
With the output of the optimization function,
status: 0
success: True
njev: 5
nfev: 25
hess_inv: array([[ 7.32764243e-01, -4.89555962e-01, -1.93879729e-08],
[ -4.89555962e-01, 3.27690582e-01, -2.09510086e-08],
[ -1.93879729e-08, -2.09510086e-08, 1.00000000e+00]])
fun: 1.784420370873494e-11
x: array([ 9.96468440e-06, 7.40349389e-06, 5.00000000e+01])
message: 'Optimization terminated successfully.'
jac: array([ -1.11880569e-06, 1.28115957e-06, 2.31303410e-12])
and the resulting plot,
I am a python beginner, currently using scipy's odeint to compute a coupled ODE system, however, when I run, python shell always tell me that
>>>
Excess work done on this call (perhaps wrong Dfun type).
Run with full_output = 1 to get quantitative information.
>>>
So, I have to change my time step and final time, in order to make it integratable. To do this, I need to try a different combinations, which is quite a pain. Could anyone tell me how can I ask odeint to automatically vary the time step and final time to successfully integrate this ode system?
and here is part of the code which has called odeint:
def main(t, init_pop_a, init_pop_b, *args, **kwargs):
"""
solve the obe for a given set of parameters
"""
# construct initial condition
# initially, rho_ee = 0
rho_init = zeros((16,16))*1j ########
rho_init[1,1] = init_pop_a
rho_init[2,2] = init_pop_b
rho_init[0,0] = 1 - (init_pop_a + init_pop_b)########
rho_init_ravel, params = to_1d(rho_init)
# perform the integration
result = odeint(wrapped_bloch3, rho_init_ravel, t, args=args)
# BUG: need to pass kwargs
# rewrap the result
return from_1d(result, params, prepend=(len(t),))
things = [2*pi, 20*pi, 0,0, 0,0, 0.1,100]
Omega_a, Omega_b, Delta_a, Delta_b, \
init_pop_a, init_pop_b, tstep, tfinal = things
args = ( Delta_a, Delta_b, Omega_a, Omega_b )
t = arange(0, tfinal + tstep, tstep)
data = main(t, init_pop_a, init_pop_b, *args)
plt.plot(t,abs(data[:,4,4]))
where wrapped_bloch3 is the function compute dy/dt.
EDIT: I note you already got an answer here: complex ODE systems in scipy
odeint does not work with complex-valued equations. I get
from scipy.integrate import odeint
import numpy as np
def func(t, y):
return 1 + 1j
t = np.linspace(0, 1, 200)
y = odeint(func, 0, t)
# -> This outputs:
#
# TypeError: can't convert complex to float
# odepack.error: Result from function call is not a proper array of floats.
You can solve your equation by the other ode solver:
from scipy.integrate import ode
import numpy as np
def myodeint(func, y0, t):
y0 = np.array(y0, complex)
func2 = lambda t, y: func(y, t) # odeint has these the other way :/
sol = ode(func2).set_integrator('zvode').set_initial_value(y0, t=t[0])
y = [sol.integrate(tp) for tp in t[1:]]
y.insert(0, y0)
return np.array(y)
def func(y, t, alpha):
return 1j*alpha*y
alpha = 3.3
t = np.linspace(0, 1, 200)
y = myodeint(lambda y, t: func(y, t, alpha), [1, 0, 0], t)