def run_example(with_plots=True):
"""
This is the same example from the Sundials package (idasRoberts_FSA_dns.c)
This simple example problem for IDA, due to Robertson,
is from chemical kinetics, and consists of the following three
equations::
dy1/dt = -p1*y1 + p2*y2*y3
dy2/dt = p1*y1 - p2*y2*y3 - p3*y2**2
0 = y1 + y2 + y3 - 1
"""
def f(t, y, yd, p):
res1 = -p[0]*y[0]+p[1]*y[1]*y[2]-yd[0]
res2 = p[0]*y[0]-p[1]*y[1]*y[2]-p[2]*y[1]**2-yd[1]
res3 = y[0]+y[1]+y[2]-1
return N.array([res1,res2,res3])
#The initial conditons
y0 = [1.0, 0.0, 0.0] #Initial conditions for y
yd0 = [0.1, 0.0, 0.0] #Initial conditions for dy/dt
p0 = [0.040, 1.0e4, 3.0e7] #Initial conditions for parameters
#Create an Assimulo implicit problem
imp_mod = Implicit_Problem(f, y0, yd0,p0=p0)
#Create an Assimulo implicit solver (IDA)
imp_sim = IDA(imp_mod) #Create a IDA solver
#Sets the paramters
imp_sim.atol = N.array([1.0e-8, 1.0e-14, 1.0e-6])
imp_sim.algvar = [1.0,1.0,0.0]
imp_sim.suppress_alg = False #Suppres the algebraic variables on the error test
imp_sim.continuous_output = True #Store data continuous during the simulation
imp_sim.pbar = p0
imp_sim.suppress_sens = False #Dont suppress the sensitivity variables in the error test.
#Let Sundials find consistent initial conditions by use of 'IDA_YA_YDP_INIT'
imp_sim.make_consistent('IDA_YA_YDP_INIT')
#Simulate
t, y, yd = imp_sim.simulate(4,400) #Simulate 4 seconds with 400 communication points
print imp_sim.p_sol[0][-1] , imp_sim.p_sol[1][-1], imp_sim.p_sol[0][-1]
#Basic test
nose.tools.assert_almost_equal(y[-1][0], 9.05518032e-01, 4)
nose.tools.assert_almost_equal(y[-1][1], 2.24046805e-05, 4)
nose.tools.assert_almost_equal(y[-1][2], 9.44595637e-02, 4)
nose.tools.assert_almost_equal(imp_sim.p_sol[0][-1][0], -1.8761, 2) #Values taken from the example in Sundials
nose.tools.assert_almost_equal(imp_sim.p_sol[1][-1][0], 2.9614e-06, 8)
nose.tools.assert_almost_equal(imp_sim.p_sol[2][-1][0], -4.9334e-10, 12)
#Plot
if with_plots:
P.plot(t,y)
P.show()
def run_example(with_plots=True):
#Create an instance of the problem
iter_mod = Extended_Problem() #Create the problem
iter_sim = IDA(iter_mod) #Create the solver
iter_sim.verbosity = 0
iter_sim.continuous_output = True
#Simulate
t, y, yd = iter_sim.simulate(10.0,1000) #Simulate 10 seconds with 1000 communications points
#Basic test
nose.tools.assert_almost_equal(y[-1][0],8.0)
nose.tools.assert_almost_equal(y[-1][1],3.0)
nose.tools.assert_almost_equal(y[-1][2],2.0)
#Plot
if with_plots:
P.plot(t,y)
P.show()
def run_example(with_plots=True):
#Create an instance of the problem
iter_mod = Extended_Problem() #Create the problem
iter_sim = IDA(iter_mod) #Create the solver
iter_sim.verbosity = 0
iter_sim.continuous_output = True
#Simulate
t, y, yd = iter_sim.simulate(
10.0, 1000) #Simulate 10 seconds with 1000 communications points
#Basic test
nose.tools.assert_almost_equal(y[-1][0], 8.0)
nose.tools.assert_almost_equal(y[-1][1], 3.0)
nose.tools.assert_almost_equal(y[-1][2], 2.0)
#Plot
if with_plots:
P.plot(t, y)
P.show()
def simulate(self, tmax, nout=500, maxsteps=5000, rtol=1e-8, verbose=False, test_jac_args=None):
if not hasattr(self, 't0'):
self.assemble(verbose)
if test_jac_args is not None:
self.test_jacobian(*test_jac_args)
sim = IDA(self)
flag, _, _ = sim.make_consistent('IDA_YA_YDP_INIT')
translation = {0:'SUCCESS', 1:'TSTOP_RETURN', 2:'ROOT_RETURN',
99:'WARNING', -1 :'TOO_MUCH_WORK', -2 :'TOO_MUCH_ACC',
-3 :'ERR_FAIL', -4 :'CONV_FAIL', -5 :'LINIT_FAIL',
-6 :'LSETUP_FAIL', -7 :'LSOLVE_FAIL', -8 :'RES_FAIL',
-9 :'REP_RES_ERR', -10:'RTFUNC_FAIL', -11:'CONSTR_FAIL',
-12:'FIRST_RES_FAIL', -13:'LINESEARCH_FAIL', -14:'NO_RECOVERY',
-20:'MEM_NULL', -21:'MEM_FAIL', -22:'ILL_INPUT', -23:'NO_MALLOC',
-24:'BAD_EWT', -25:'BAD_K', -26:'BAD_T', -27:'BAD_DKY'}
if flag < 0:
raise ArithmeticError('make_consistent failed with flag = IDA_%s' % translation[flag])
if flag != 0:
warn('make_consistent returned IDA_%s' % translation[flag])
sim.rtol = rtol
sim.maxsteps = maxsteps
T,Y,Yd = sim.simulate(tmax, nout)
ncnt = len(self.nodes)
nterm = len(self.terminals)
self.T = T
Vs = Y [:, :ncnt]
dVsdt = Yd[:, :ncnt]
Is = Y [:, ncnt:(ncnt+nterm)]
dIsdt = Yd[:, ncnt:(ncnt+nterm)]
Qs = Y [:, (ncnt+nterm):]
dQsdt = Yd[:, (ncnt+nterm):]
for i,node in enumerate(self.nodes):
for term in node:
term.V = Vs[:, i]
term.dVdt = dVsdt[:, i]
for i,term in enumerate(self.terminals):
term.I = Is[:, i]
term.dIdt = dIsdt[:, i]
for i,state in enumerate(self.states):
state.val = Qs[:, i]
state.der = dQsdt[:, i]
def run_example(with_plots=True):
#Defines the residual
def f(t,y,yd):
res_0 = yd[0]-y[2]
res_1 = yd[1]-y[3]
res_2 = yd[2]+y[4]*y[0]
res_3 = yd[3]+y[4]*y[1]+9.82
#res_4 = y[0]**2+y[1]**2-1
res_4 = y[2]**2+y[3]**2-y[4]*(y[0]**2+y[1]**2)-y[1]*9.82
return N.array([res_0,res_1,res_2,res_3,res_4])
#Defines the jacobian
def jac(c,t,y,yd):
jacobian = N.zeros([len(y),len(y)])
#Derivative
jacobian[0,0] = 1*c
jacobian[1,1] = 1*c
jacobian[2,2] = 1*c
jacobian[3,3] = 1*c
#Differentiated
jacobian[0,2] = -1
jacobian[1,3] = -1
jacobian[2,0] = y[4]
jacobian[3,1] = y[4]
jacobian[4,0] = y[0]*2*y[4]*-1
jacobian[4,1] = y[1]*2*y[4]*-1-9.82
jacobian[4,2] = y[2]*2
jacobian[4,3] = y[3]*2
#Algebraic
jacobian[2,4] = y[0]
jacobian[3,4] = y[1]
jacobian[4,4] = -(y[0]**2+y[1]**2)
return jacobian
#The initial conditons
y0 = [1.0,0.0,0.0,0.0,5] #Initial conditions
yd0 = [0.0,0.0,0.0,-9.82,0.0] #Initial conditions
#Create an Assimulo implicit problem
imp_mod = Implicit_Problem(f,y0,yd0)
#Sets the options to the problem
imp_mod.jac = jac #Sets the jacobian
imp_mod.algvar = [1.0,1.0,1.0,1.0,0.0] #Set the algebraic components
imp_mod.name = 'Test Jacobian'
#Create an Assimulo implicit solver (IDA)
imp_sim = IDA(imp_mod) #Create a IDA solver
#Sets the paramters
imp_sim.atol = 1e-6 #Default 1e-6
imp_sim.rtol = 1e-6 #Default 1e-6
imp_sim.suppress_alg = True #Suppres the algebraic variables on the error test
#Let Sundials find consistent initial conditions by use of 'IDA_YA_YDP_INIT'
imp_sim.make_consistent('IDA_YA_YDP_INIT')
#Simulate
t, y, yd = imp_sim.simulate(5,1000) #Simulate 5 seconds with 1000 communication points
#Basic tests
nose.tools.assert_almost_equal(y[-1][0],0.9401995, places=4)
nose.tools.assert_almost_equal(y[-1][1],-0.34095124, places=4)
nose.tools.assert_almost_equal(yd[-1][0], -0.88198927, places=4)
nose.tools.assert_almost_equal(yd[-1][1], -2.43227069, places=4)
#Plot
if with_plots:
P.plot(t,y)
P.show()
def run_example(with_plots=True):
"""
This is the same example from the Sundials package (idasRoberts_FSA_dns.c)
This simple example problem for IDA, due to Robertson,
is from chemical kinetics, and consists of the following three
equations::
dy1/dt = -p1*y1 + p2*y2*y3
dy2/dt = p1*y1 - p2*y2*y3 - p3*y2**2
0 = y1 + y2 + y3 - 1
"""
def f(t, y, yd, p):
res1 = -p[0] * y[0] + p[1] * y[1] * y[2] - yd[0]
res2 = p[0] * y[0] - p[1] * y[1] * y[2] - p[2] * y[1]**2 - yd[1]
res3 = y[0] + y[1] + y[2] - 1
return N.array([res1, res2, res3])
#The initial conditons
y0 = [1.0, 0.0, 0.0] #Initial conditions for y
yd0 = [0.1, 0.0, 0.0] #Initial conditions for dy/dt
p0 = [0.040, 1.0e4, 3.0e7] #Initial conditions for parameters
#Create an Assimulo implicit problem
imp_mod = Implicit_Problem(f, y0, yd0, p0=p0)
#Create an Assimulo implicit solver (IDA)
imp_sim = IDA(imp_mod) #Create a IDA solver
#Sets the paramters
imp_sim.atol = N.array([1.0e-8, 1.0e-14, 1.0e-6])
imp_sim.algvar = [1.0, 1.0, 0.0]
imp_sim.suppress_alg = False #Suppres the algebraic variables on the error test
imp_sim.continuous_output = True #Store data continuous during the simulation
imp_sim.pbar = p0
imp_sim.suppress_sens = False #Dont suppress the sensitivity variables in the error test.
#Let Sundials find consistent initial conditions by use of 'IDA_YA_YDP_INIT'
imp_sim.make_consistent('IDA_YA_YDP_INIT')
#Simulate
t, y, yd = imp_sim.simulate(
4, 400) #Simulate 4 seconds with 400 communication points
print imp_sim.p_sol[0][-1], imp_sim.p_sol[1][-1], imp_sim.p_sol[0][-1]
#Basic test
nose.tools.assert_almost_equal(y[-1][0], 9.05518032e-01, 4)
nose.tools.assert_almost_equal(y[-1][1], 2.24046805e-05, 4)
nose.tools.assert_almost_equal(y[-1][2], 9.44595637e-02, 4)
nose.tools.assert_almost_equal(
imp_sim.p_sol[0][-1][0], -1.8761,
2) #Values taken from the example in Sundials
nose.tools.assert_almost_equal(imp_sim.p_sol[1][-1][0], 2.9614e-06, 8)
nose.tools.assert_almost_equal(imp_sim.p_sol[2][-1][0], -4.9334e-10, 12)
#Plot
if with_plots:
P.plot(t, y)
P.show()
def simulate(model, init_cond, start_time=0., final_time=1., input=(lambda t: []), ncp=500, blt=True,
causalization_options=sp.CausalizationOptions(), expand_to_sx=True, suppress_alg=False,
tol=1e-8, solver="IDA"):
"""
Simulate model from CasADi Interface using CasADi.
init_cond is a dictionary containing initial conditions for all variables.
"""
if blt:
t_0 = timing.time()
blt_model = sp.BLTModel(model, causalization_options)
blt_time = timing.time() - t_0
print("BLT analysis time: %.3f s" % blt_time)
blt_model._model = model
model = blt_model
if causalization_options['closed_form']:
solved_vars = model._solved_vars
solved_expr = model._solved_expr
#~ for (var, expr) in itertools.izip(solved_vars, solved_expr):
#~ print('%s := %s' % (var.getName(), expr))
return model
dh() # This is not a debug statement!
# Extract model variables
model_states = [var for var in model.getVariables(model.DIFFERENTIATED) if not var.isAlias()]
model_derivatives = [var for var in model.getVariables(model.DERIVATIVE) if not var.isAlias()]
model_algs = [var for var in model.getVariables(model.REAL_ALGEBRAIC) if not var.isAlias()]
model_inputs = [var for var in model.getVariables(model.REAL_INPUT) if not var.isAlias()]
states = [var.getVar() for var in model_states]
derivatives = [var.getMyDerivativeVariable().getVar() for var in model_states]
algebraics = [var.getVar() for var in model_algs]
inputs = [var.getVar() for var in model_inputs]
n_x = len(states)
n_y = len(states) + len(algebraics)
n_w = len(algebraics)
n_u = len(inputs)
# Create vectorized model variables
t = model.getTimeVariable()
y = MX.sym("y", n_y)
yd = MX.sym("yd", n_y)
u = MX.sym("u", n_u)
# Extract the residuals and substitute the (x,z) variables for the old variables
scalar_vars = states + algebraics + derivatives + inputs
vector_vars = [y[k] for k in range(n_y)] + [yd[k] for k in range(n_x)] + [u[k] for k in range(n_u)]
[dae] = substitute([model.getDaeResidual()], scalar_vars, vector_vars)
# Fix parameters
if not blt:
# Sort parameters
par_kinds = [model.BOOLEAN_CONSTANT,
model.BOOLEAN_PARAMETER_DEPENDENT,
model.BOOLEAN_PARAMETER_INDEPENDENT,
model.INTEGER_CONSTANT,
model.INTEGER_PARAMETER_DEPENDENT,
model.INTEGER_PARAMETER_INDEPENDENT,
model.REAL_CONSTANT,
model.REAL_PARAMETER_INDEPENDENT,
model.REAL_PARAMETER_DEPENDENT]
pars = reduce(list.__add__, [list(model.getVariables(par_kind)) for
par_kind in par_kinds])
# Get parameter values
model.calculateValuesForDependentParameters()
par_vars = [par.getVar() for par in pars]
par_vals = [model.get_attr(par, "_value") for par in pars]
# Eliminate parameters
[dae] = casadi.substitute([dae], par_vars, par_vals)
# Extract initial conditions
y0 = [init_cond[var.getName()] for var in model_states] + [init_cond[var.getName()] for var in model_algs]
yd0 = [init_cond[var.getName()] for var in model_derivatives] + n_w * [0.]
# Create residual CasADi functions
dae_res = MXFunction([t, y, yd, u], [dae])
dae_res.setOption("name", "complete_dae_residual")
dae_res.init()
###################
#~ import matplotlib.pyplot as plt
#~ h = MX.sym("h")
#~ iter_matrix_expr = dae_res.jac(2)/h + dae_res.jac(1)
#~ iter_matrix = MXFunction([t, y, yd, u, h], [iter_matrix_expr])
#~ iter_matrix.init()
#~ n = 100;
#~ hs = np.logspace(-8, 1, n);
#~ conds = [np.linalg.cond(iter_matrix.call([0, y0, yd0, input(0), hval])[0].toArray()) for hval in hs]
#~ plt.close(1)
#~ plt.figure(1)
#~ plt.loglog(hs, conds, 'b-')
#~ #plt.gca().invert_xaxis()
#~ plt.grid('on')
#~ didx = range(4, 12) + range(30, 33)
#~ aidx = [i for i in range(33) if i not in didx]
#~ didx = range(10)
#~ aidx = []
#~ F = MXFunction([t, y, yd, u], [dae[didx]])
#~ F.init()
#~ G = MXFunction([t, y, yd, u], [dae[aidx]])
#~ G.init()
#~ dFddx = F.jac(2)[:, :n_x]
#~ dFdx = F.jac(1)[:, :n_x]
#~ dFdy = F.jac(1)[:, n_x:]
#~ dGdx = G.jac(1)[:, :n_x]
#~ dGdy = G.jac(1)[:, n_x:]
#~ E_matrix = MXFunction([t, y, yd, u, h], [dFddx])
#~ E_matrix.init()
#~ E_cond = np.linalg.cond(E_matrix.call([0, y0, yd0, input(0), hval])[0].toArray())
#~ iter_matrix_expr = vertcat([horzcat([dFddx + h*dFdx, h*dFdy]), horzcat([dGdx, dGdy])])
#~ iter_matrix = MXFunction([t, y, yd, u, h], [iter_matrix_expr])
#~ iter_matrix.init()
#~ n = 100
#~ hs = np.logspace(-8, 1, n)
#~ conds = [np.linalg.cond(iter_matrix.call([0, y0, yd0, input(0), hval])[0].toArray()) for hval in hs]
#~ plt.loglog(hs, conds, 'b--')
#~ plt.gca().invert_xaxis()
#~ plt.grid('on')
#~ plt.xlabel('$h$')
#~ plt.ylabel('$\kappa$')
#~ plt.show()
#~ dh()
###################
# Expand to SX
if expand_to_sx:
dae_res = SXFunction(dae_res)
dae_res.init()
# Create DAE residual Assimulo function
def dae_residual(t, y, yd):
dae_res.setInput(t, 0)
dae_res.setInput(y, 1)
dae_res.setInput(yd, 2)
dae_res.setInput(input(t), 3)
dae_res.evaluate()
return dae_res.getOutput(0).toArray().reshape(-1)
# Set up simulator
problem = Implicit_Problem(dae_residual, y0, yd0, start_time)
if solver == "IDA":
simulator = IDA(problem)
elif solver == "Radau5DAE":
simulator = Radau5DAE(problem)
else:
raise ValueError("Unknown solver %s" % solver)
simulator.rtol = tol
simulator.atol = 1e-4 * np.array([model.get_attr(var, "nominal") for var in model_states + model_algs])
#~ simulator.atol = tol * np.ones([n_y, 1])
simulator.report_continuously = True
# Log method order
if solver == "IDA":
global order
order = []
def handle_result(solver, t, y, yd):
global order
order.append(solver.get_last_order())
solver.t_sol.extend([t])
solver.y_sol.extend([y])
solver.yd_sol.extend([yd])
problem.handle_result = handle_result
# Suppress algebraic variables
if suppress_alg:
if isinstance(suppress_alg, bool):
simulator.algvar = n_x * [True] + (n_y - n_x) * [False]
else:
simulator.algvar = n_x * [True] + suppress_alg
simulator.suppress_alg = True
# Simulate
t_0 = timing.time()
(t, y, yd) = simulator.simulate(final_time, ncp)
simul_time = timing.time() - t_0
stats = {'time': simul_time, 'steps': simulator.statistics['nsteps']}
if solver == "IDA":
stats['order'] = order
# Generate result for time and inputs
class SimulationResult(dict):
pass
res = SimulationResult()
res.stats = stats
res['time'] = t
if u.numel() > 0:
input_names = [var.getName() for var in model_inputs]
for name in input_names:
res[name] = []
for time in t:
input_val = input(time)
for (name, val) in itertools.izip(input_names, input_val):
res[name].append(val)
# Create results for everything else
if blt:
# Iteration variables
i = 0
for var in model_states:
res[var.getName()] = y[:, i]
res[var.getMyDerivativeVariable().getName()] = yd[:, i]
i += 1
for var in model_algs:
res[var.getName()] = y[:, i]
i += 1
# Create function for computing solved algebraics
for (_, sol_alg) in model._explicit_solved_algebraics:
res[sol_alg.name] = []
alg_sol_f = casadi.MXFunction(model._known_vars + model._explicit_unsolved_vars, model._solved_expr)
alg_sol_f.init()
if expand_to_sx:
alg_sol_f = casadi.SXFunction(alg_sol_f)
alg_sol_f.init()
# Compute solved algebraics
for k in xrange(len(t)):
for (i, var) in enumerate(model._known_vars + model._explicit_unsolved_vars):
alg_sol_f.setInput(res[var.getName()][k], i)
alg_sol_f.evaluate()
for (j, sol_alg) in model._explicit_solved_algebraics:
res[sol_alg.name].append(alg_sol_f.getOutput(j).toScalar())
else:
res_vars = model_states + model_algs
for (i, var) in enumerate(res_vars):
res[var.getName()] = y[:, i]
der_var = var.getMyDerivativeVariable()
if der_var is not None:
res[der_var.getName()] = yd[:, i]
# Add results for all alias variables (only treat time-continuous variables) and convert to array
if blt:
res_model = model._model
else:
res_model = model
for var in res_model.getAllVariables():
if var.getVariability() == var.CONTINUOUS:
res[var.getName()] = np.array(res[var.getModelVariable().getName()])
res["time"] = np.array(res["time"])
res._blt_model = blt_model
return res
yp[2] = -yp[2]
print("DING")
solver.yd = yp
solver.y = y
t0 = 0
tfinal = 1
ncp = 500
y0 = [0., -0.10344, -0.65, 0. ,0. , 0., -0.6911, -0.14161]
yp0 = [0., 0., 0., 0., 0., 1.40059e2, 0., 0.]
switches0 = [False, True, False]
mod = Implicit_Problem(woodpeckertoy, y0, yp0, t0, sw0 = switches0)
mod.state_events = state_events
mod.handle_event = handle_event
sim = IDA(mod)
sim.suppress_alg = True
sim.rtol=1.e-6
sim.atol[3:8] = 1e6
sim.algvar[3:8] = 0.
t, y, yp = sim.simulate(tfinal, ncp)
P.plot(t,y[:,0:3])
P.legend(["z","phiS", "phiB", "z vel", "phiS vel", "phiB vel"])
P.ylabel('y')
P.xlabel('x')
fontlabel_size = 20
tick_size = 14
params = {'lines.markersize' : 0, 'axes.labelsize': fontlabel_size, 'text.fontsize': fontlabel_size, 'legend.fontsize': fontlabel_size, 'xtick.labelsize': tick_size, 'ytick.labelsize': tick_size}
P.rcParams.update(params)
P.show()
def run_example(with_plots=True):
"""
This example show how to use Assimulo and IDA for simulating sensitivities
for initial conditions.::
0 = dy1/dt - -(k01+k21+k31)*y1 - k12*y2 - k13*y3 - b1
0 = dy2/dt - k21*y1 + (k02+k12)*y2
0 = dy3/dt - k31*y1 + k13*y3
y1(0) = p1, y2(0) = p2, y3(0) = p3
p1=p2=p3 = 0
See http://sundials.2283335.n4.nabble.com/Forward-sensitivities-for-initial-conditions-td3239724.html
"""
def f(t, y, yd,p):
y1,y2,y3 = y
yd1,yd2,yd3 = yd
k01 = 0.0211
k02 = 0.0162
k21 = 0.0111
k12 = 0.0124
k31 = 0.0039
k13 = 0.000035
b1 = 49.3
res_0 = -yd1 -(k01+k21+k31)*y1+k12*y2+k13*y3+b1
res_1 = -yd2 + k21*y1-(k02+k12)*y2
res_2 = -yd3 + k31*y1-k13*y3
return N.array([res_0,res_1,res_2])
#The initial conditions
y0 = [0.0,0.0,0.0] #Initial conditions for y
yd0 = [49.3,0.,0.]
p0 = [0.0, 0.0, 0.0] #Initial conditions for parameters
yS0 = N.array([[1,0,0],[0,1,0],[0,0,1.]])
#Create an Assimulo implicit problem
imp_mod = Implicit_Problem(f,y0,yd0,p0=p0)
#Sets the options to the problem
imp_mod.yS0=yS0
#Create an Assimulo explicit solver (IDA)
imp_sim = IDA(imp_mod)
#Sets the paramters
imp_sim.rtol = 1e-7
imp_sim.atol = 1e-6
imp_sim.pbar = [1,1,1] #pbar is used to estimate the tolerances for the parameters
imp_sim.continuous_output = True #Need to be able to store the result using the interpolate methods
imp_sim.sensmethod = 'SIMULTANEOUS' #Defines the sensitvity method used
imp_sim.suppress_sens = False #Dont suppress the sensitivity variables in the error test.
#Simulate
t, y, yd = imp_sim.simulate(400) #Simulate 400 seconds
#Basic test
nose.tools.assert_almost_equal(y[-1][0], 1577.6552477,3)
nose.tools.assert_almost_equal(y[-1][1], 611.9574565, 3)
nose.tools.assert_almost_equal(y[-1][2], 2215.88563217, 3)
nose.tools.assert_almost_equal(imp_sim.p_sol[0][1][0], 1.0)
#Plot
if with_plots:
P.figure(1)
P.subplot(221)
P.plot(t, N.array(imp_sim.p_sol[0])[:,0],
t, N.array(imp_sim.p_sol[0])[:,1],
t, N.array(imp_sim.p_sol[0])[:,2])
P.title("Parameter p1")
P.legend(("p1/dy1","p1/dy2","p1/dy3"))
P.subplot(222)
P.plot(t, N.array(imp_sim.p_sol[1])[:,0],
t, N.array(imp_sim.p_sol[1])[:,1],
t, N.array(imp_sim.p_sol[1])[:,2])
P.title("Parameter p2")
P.legend(("p2/dy1","p2/dy2","p2/dy3"))
P.subplot(223)
P.plot(t, N.array(imp_sim.p_sol[2])[:,0],
t, N.array(imp_sim.p_sol[2])[:,1],
t, N.array(imp_sim.p_sol[2])[:,2])
P.title("Parameter p3")
P.legend(("p3/dy1","p3/dy2","p3/dy3"))
P.subplot(224)
P.plot(t,y)
P.show()
def PostREC(plot_request,Z_initial,Z_final,step_size,number_of_particles, manual_matter_temperature, source_function,chemical_model):
global c,kb,h_p,H0,W0,T_0,kbev,Z,cooling_function,flux,k,T_m, hev ###Global parameters defined for use throughout the function
###Mianly cosmological parameter values, numerical constants etc.
c = 2.99792458*(10**8) #Speed of light (SI units)
kb = 1.38065*(10**-23) #Boltzmann constant (SI units)
h_p = 6.626068*(10**-34) #Planck's constant (SI units)
H0 = 69.7 #100h
W0 = 1.0 #Total matter density (flat universe)
T_0 = 2.725 #Modern day temperature of the CMB
kbev = 8.6173324*(10**-5) #Boltzman constant in eV
hev = 4.135667*(10**-15) #Plancks constant in eV
################################## Program Functions #################################################
##Residual is the main function that the solver, IDA, uses to solve the system of equations. Each "residual"
##(res_1, res_2 etc) corresponds to a certain chemical species. The reaction eqautions array contains the
##reaction equation for that particular species. Values for the rates are substituted in (k) along with the
##matter and radiation temperatures. IDA then trys to make the value of the residual as close to 0 as possible
##by doing so it finds the equilibrium state of the system.
##list of chemical species and the array elements that each corresponds to
#species_list = ["[H]","[H+]","[H-]","[H2]","[H2+]","[e-]","[hv]","[D]","[D-]","[D+]","[HD]","[HD+]","[He]","[He+]","[He++]","[HeH+]",[Li], [Li+], [Li-], [LiH], [LiH+], [H2D+], [D2]]
# 0 1 2 3 4 5 6 9 10 11 12 13 14 15 16 17 18 19 20 21 22 , 23 , 24
def residual(t,y,yd):
res_0 = reaction_equations[0](k,Z,y[8],y[7],y) - yd[0]
res_1 = reaction_equations[1](k,Z,y[8],y[7],y) - yd[1]
res_2 = reaction_equations[2](k,Z,y[8],y[7],y) - yd[2]
res_3 = reaction_equations[3](k,Z,y[8],y[7],y) - yd[3]
res_4 = reaction_equations[4](k,Z,y[8],y[7],y) - yd[4]
res_5 = reaction_equations[5](k,Z,y[8],y[7],y) - yd[5]
photon_density = 0.0 -yd[6]
radiation_temp = 0.0 -yd[7]
matter_temp = 0.0 -yd[8]
res_9 = reaction_equations[7](k,Z,y[8],y[7],y) - yd[9]
res_10 =reaction_equations[8](k,Z,y[8],y[7],y) -yd[10]
res_11 = reaction_equations[9](k,Z,y[8],y[7],y)-yd[11]
res_12 = reaction_equations[10](k,Z,y[8],y[7],y) - yd[12]
res_13 = reaction_equations[11](k,Z,y[8],y[7],y) -yd[13]
res_14 = reaction_equations[12](k,Z,y[8],y[7],y)-yd[14]
res_15 = reaction_equations[13](k,Z,y[8],y[7],y) - yd[15]
res_16 = reaction_equations[14](k,Z,y[8],y[7],y)- yd[16]
res_17 = reaction_equations[15](k,Z,y[8],y[7],y) -yd[17]
res_18 = reaction_equations[16](k,Z,y[8],y[7],y) -yd[18]
res_19 = reaction_equations[17](k,Z,y[8],y[7],y) -yd[19]
res_20 = reaction_equations[18](k,Z,y[8],y[7],y) -yd[20]
res_21 =reaction_equations[19](k,Z,y[8],y[7],y) -yd[21]
res_22 =reaction_equations[20](k,Z,y[8],y[7],y)-yd[22]
res_23 =reaction_equations[21](k,Z,y[8],y[7],y)-yd[23]
res_24 =reaction_equations[22](k,Z,y[8],y[7],y)-yd[24]
return array([res_0,res_1,res_2,res_3,res_4,res_5,photon_density,radiation_temp,matter_temp,res_9,res_10,res_11,res_12,res_13,res_14,res_15,res_16,res_17,res_18,res_19,res_20,res_21,res_22,res_23,res_24])
##deriv provides a jacobian matrix for IDA so that it has apprpriate initial conditions and can solve the equations more efficiently.
def deriv(y,t):
res_0 = reaction_equations[0](k,Z,y[8],y[7],y)
res_1 = reaction_equations[1](k,Z,y[8],y[7],y)
res_2 = reaction_equations[2](k,Z,y[8],y[7],y)
res_3 = reaction_equations[3](k,Z,y[8],y[7],y)
res_4 = reaction_equations[4](k,Z,y[8],y[7],y)
res_5 = reaction_equations[5](k,Z,y[8],y[7],y)
photon_density = 0.0
radiation_temp = 0.0
matter_temp = -cooling_function(y[8],y[7],y[0],y[3],y[12],y[5],y[1],y[4],y[14],y[15],y[16],y[21])
res_9 = reaction_equations[7](k,Z,y[8],y[7],y)
res_10 =reaction_equations[8](k,Z,y[8],y[7],y)
res_11 = reaction_equations[9](k,Z,y[8],y[7],y)
res_12 = reaction_equations[10](k,Z,y[8],y[7],y)
res_13 = reaction_equations[11](k,Z,y[8],y[7],y)
res_14 = reaction_equations[12](k,Z,y[8],y[7],y)
res_15 = reaction_equations[13](k,Z,y[8],y[7],y)
res_16 = reaction_equations[14](k,Z,y[8],y[7],y)
res_17 = reaction_equations[15](k,Z,y[8],y[7],y)
res_18 = reaction_equations[16](k,Z,y[8],y[7],y)
res_19 = reaction_equations[17](k,Z,y[8],y[7],y)
res_20 = reaction_equations[18](k,Z,y[8],y[7],y)
res_21 =reaction_equations[19](k,Z,y[8],y[7],y)
res_22 =reaction_equations[20](k,Z,y[8],y[7],y)
res_23 =reaction_equations[21](k,Z,y[8],y[7],y)
res_24 =reaction_equations[22](k,Z,y[8],y[7],y)
return array([res_0,res_1,res_2,res_3,res_4,res_5,photon_density,radiation_temp,matter_temp,res_9,res_10,res_11,res_12,res_13,res_14,res_15,res_16,res_17,res_18,res_19,res_20,res_21,res_22,res_23,res_24])
##The error_checking function performs an error test on the results. It does this by comparing the final total number of particles to the initial number.
##If there is a large discrepancy, there is a bug somewhere within the code or the solution is unstable. The user is then warned.
def error_checking(error_test,end_line,Z_initial,Z_final):
atoms_lost = error_test*(((1+Z_final)**3)/((1+Z_initial)**3))-(end_line[0]+end_line[1]+end_line[2]+(2*end_line[3])+(2*end_line[4])+end_line[9]+end_line[10]+end_line[11]+(2*end_line[12])+(2*end_line[13])+end_line[15]+end_line[16]+(2*end_line[17])+end_line[18]+end_line[19]+end_line[20]+(2*end_line[21])+(2*end_line[22])+(3*end_line[23])+end_line[14]+(2*end_line[24]))
return(atoms_lost)
##The initial conditions function calculates the initial conditions ready for integration
def initial_conditions(Z_initial,number_of_particles, manual_matter_temperature): #This function returns the initial conditions of the chosen epoch ready to be passed to the numerical integrator
no_H = number_of_particles/((2.68*(10**-5))+0.240+(10**-9.9)+1.0) #Calculate the hydrogen fraction using the fractional abundance of other elements
no_He = 0.240*no_H #Calculate the fractional abundances of elements using nucelosynthesis results
no_Deu = 2.68*(10**-5)*no_H
no_Lithium = (10**-9.9)*no_H
recombination_data = loadtxt('data.dat', delimiter=' ', dtype = (str)) #Loads in an array giving us the recombination fraction and ratio of matter to radiation temperature
with open('data.dat','r') as f:
rows=[map(float,L.strip().split(' ')) for L in f]
arr=array(rows)
free_electron_fraction = arr[8000-Z_initial][1] #Reads in the free electron fraction from the loaded data
matter_radiation_temperature_ratio = arr[8000-Z_initial][2] #Reads the matter radiation temperature ratio from the loaded data
radiation_temperature = (2.725)*(1+Z_initial) #Calculates the initial radiation temperature using the initial redshift
if manual_matter_temperature == 0: #Calculates the matter temperature by first checking if a manual
#value was entered, if not then it is calculated from the ratio.
matter_temperature = matter_radiation_temperature_ratio*radiation_temperature
else:
matter_temperature = manual_matter_temperature
no_e = no_H*free_electron_fraction #Calculate number of electrons from free electron fraction
if free_electron_fraction > 1.0: #Calculate the ionised fraction of elements using ionised fraction.
no_He_plus = (free_electron_fraction - 1.0)*no_He
no_He = no_He - no_He_plus
no_H_plus = no_H
no_H = 0.0
no_Deu_plus = no_Deu
no_Deu = 0.0
else:
no_He_plus = 0.0
no_H_plus = free_electron_fraction*no_H
no_H = no_H - no_H_plus
no_Deu_plus = free_electron_fraction*no_Deu
no_Deu = no_Deu - no_Deu_plus
photon_density = (number_of_particles*((10**10)/4.5)) # Calculates the number density of photons using the photon to baryon ratio
return ([no_H,no_H_plus,0,0,0,no_e,photon_density,radiation_temperature,matter_temperature,no_Deu,0.0,no_Deu_plus,0.0,0.0,no_He,no_He_plus,0.0,0.0,0.0,no_Lithium,0.0,0.0,0.0,0.0,0.0],(no_H+no_H_plus+no_Deu+no_Deu_plus+no_He+no_He_plus+no_Lithium))
##It retuns the abundances of each chemical species, in accordance with the labelling below. It also returns number densities to be used by the error checking function later on.
#species_list = ["[H]","[H+]","[H-]","[H2]","[H2+]","[e-]","[hv]","[D]","[D-]","[D+]","[HD]","[HD+]","[He]","[He+]","[He++]","[HeH+]",[Li], [Li+], [Li-], [LiH], [LiH+], [H2D+], [D2]]
# 0 1 2 3 4 5 6 9 10 11 12 13 14 15 16 17 18 19 20 21 22 , 23 , 24
def photo_ionisation_calculator(source_function): ##This function calculates the effect of photoionising flux on the pre-galacti gas
##refer to reaction_and_rte_writer.py for more info and sources for these rates
flux = [0] * 12 ##Prepare list ready to be filled with values
##Reaction H + hv -> H+ + e-
##Integrate the cross section over the required frequency range through use of the source function, repeat for other reactions
flux[0] = quad(lambda v: ((6.30*(10**-18))*(((3.28846551*(10**15))/v)**4.0)*((exp(4-(4*atan((((v/(3.28846551*(10**15)))-1.0)**0.5))/(((v/(3.28846551*(10**15)))-1)**0.5))))/(1-exp((-2*pi)/(((v/(3.28846551*(10**15)))-1)**0.5)))))*(4.0*pi*((h_p*v)**-1))*eval(source_function), (3.28846551*(10**15)), 10**20)[0]
##Reaction He+ + hv -> He++ + e-
flux[1] = quad(lambda v: ((6.30*(10**-18))*(((3.28846551*(10**15))/v)**4.0)*((exp(4-(4*atan((((v/(3.28846551*(10**15)))-1)**0.5))/(((v/(3.28846551*(10**15)))-1)**0.5))))/(1-exp((-2*pi)/(((v/(3.28846551*(10**15)))-1)**0.5)))))*(4.0*pi*((h_p*v)**-1))*eval(source_function), (3.28846551*(10**15)), 10**20)[0]
##Reaction He + hv -> He+ + e-
flux[2] = quad(lambda v: (7.42*(10**-18)*(((1.66*(v/(24.6/(4.13566*(10**-15))))**-2.05))-(0.66*((v/(24.6/(4.13566*(10**-15))))**-3.05))))*(4.0*pi*((h_p*v)**-1))*eval(source_function),(5.948253*(10**15)), 10**20)[0]
##reaction H- + hv -> H + e-
flux[3] = quad(lambda v: (7.928*(10**5)*((v-(0.755/(4.13566*(10**-15))))**1.5)*(v**-3))*(4.0*pi*((h_p*v)**-1))*eval(source_function),(1.8255*(10**14)), 10**20)[0]
##reaction D + hv -> D+ + e-
flux[4] = quad(lambda v: ((6.30*(10**-18))*(((3.28846551*(10**15))/v)**4.0)*((exp(4-(4*atan((((v/(3.28846551*(10**15)))-1.0)**0.5))/(((v/(3.28846551*(10**15)))-1)**0.5))))/(1-exp((-2*pi)/(((v/(3.28846551*(10**15)))-1)**0.5)))))*(4.0*pi*((h_p*v)**-1))*eval(source_function), (3.28846551*(10**15)), 10**20)[0]
##reaction: H2 + hv -> H2+ + e- (predisscoaition reaction) there are 3 elements for this element, one for
##the effect of each frequency band
flux[5] = quad(lambda v: (6.2*(10**-18)*((4.135667516*(10**-15))*v)-(9.4*(10**-17)))*(4.0*pi*((h_p*v)**-1))*eval(source_function), (3.728539*(10**15)),(3.9896824*(10**15)))[0]
flux[6] = quad(lambda v: (1.4*(10**-18)*((4.135667516*(10**-15))*v)-(1.48*(10**-17)))*(4.0*pi*((h_p*v)**-1))*eval(source_function), (3.9896824*(10**15)), (4.2798411*(10**15)))[0]
flux[7] = quad(lambda v: (2.5*(10**-14)*(((4.135667516*(10**-15))*v)**-2.71))*(4.0*pi*((h_p*v)**-1))*eval(source_function), (4.2798411*(10**15)), 10**20)[0]
##Reaction: H2+ + hv -> H + H+
flux[8] = quad(lambda v: (10**((-1.6547717*(10**6))+(1.866033*(10**5)*log(v))-(7.8986431*(10**3)*(log(v)**2))+(148.73693*(log(v)**3))-(1.0513032*(log(v)**4))))*(4.0*pi*((h_p*v)**-1))*eval(source_function), (6.407671772*(10**14)), 10**20)[0]
##Reaction: H2+ + hv -> 2H+ + e-
flux[9] = quad(lambda v: (10**((-16.926)-((4.528*(10**-2))*((4.135667516*(10**-15))*v))+(2.238*(10**-4)*(((4.135667516*(10**-15))*v)**2))+(4.425*(10**-7)*(((4.135667516*(10**-15))*v)**3))))*(4.0*pi*((h_p*v)**-1))*eval(source_function), (7.253968*(10**15)), (2.176190413*(10**16)))[0]
##reaction: H2 + hv -> H2* -> H + H"
source_function_jw = source_function
v = 12.87/hev
j_w = eval(source_function)
del v
flux[10] = ((1.1*(10**8)*j_w))
##reaction H2 + hv -> H2+ + e- (Direct photodissociation into lyman warner continuum bands)
##different cross sections need to be used fordifferent frequeny bands, these are shown below, along with the limits (in eV).
##These are all taken from tom Abels 1996 paper "Modelling Primordial gas in Numerical Cosmology"
##Total ionisation rate is the sum of the ionisation from each frequency band.
sigmaL01 = lambda v: ((10**-18)*(10**(15.1289-1.05139*hev*v))) ##14.675 < hv < 16.820
sigmaL02 = lambda v:((10**-18)*(10**(-31.41 + (1.8042*(10**-2)*((hev*v)**3)) - (4.2339*(10**-5))*((hev*v)**5)))) ##16.820 < hv < 17.6
sigmaW0 = lambda v:((10**-18)*(10**(13.5311-(0.9182618*hev*v)))) ##14.675 < hv < 17.7
sigmaL1 = lambda v:((10**-18)*(10**(12.0218406-(0.819429*hev*v)))) ## 14.159 < hv < 15.302
sigmaL2 = lambda v:((10**-18)*(10**(16.04644-(1.082438*hev*v)))) ##15.302 < hv < 17.2
sigmaW1 = lambda v:((10**-18)*(10**(12.87367-(0.85088597*hev*v)))) ##14.159 < hv < 17.2
band1 = quad(lambda v: ((0.25*(sigmaL01(v)+sigmaW0(v))+(0.75*sigmaL1(v)+sigmaW1(v)))*(4.0*pi*((h_p*v)**-1))*eval(source_function)), (3.4236*(10**15)),(3.7000*(10**15)))[0]
band2 = quad(lambda v: ((0.25*(sigmaL01(v)+sigmaW0(v))+(0.75*sigmaL2(v)+sigmaW1(v)))*(4.0*pi*((h_p*v)**-1))*eval(source_function)), (3.7000*(10**15)),(4.067*(10**15)))[0]
band3 = quad(lambda v: ((0.25*(sigmaL02(v)+sigmaW0(v))+(0.75*sigmaL2(v)+sigmaW1(v)))*(4.0*pi*((h_p*v)**-1))*eval(source_function)), (4.067*(10**15)), (4.2798*(10**15)))[0]
flux[11] = band1+band2+band3
return flux
##The age_converter gives the age of the Universe as any redshift Z in Gyrs. This
##is converted into seconds using the formula below.
def range_calculator(Z_initial,Z_final):
range = (Z_initial-Z_final)*(10**9)*365.25*24*60*60
return range
def file_unpacking(chemical_model): ##Performs necessary file unpacking for reaction rates and cooling rates
pickle_file0 = open("previous_use.p","rb") ##Checks to see what the previous chemical model in use was. If it is the same as the current use
##then the arrays are loaded straight away - nothing needs to be recalculated
previous_model_used = pickle.load(pickle_file0)
if (previous_model_used == chemical_model):
pickle_file1 =open("reactionratearray.p","rb")
pickle_file4 = open("reaction_equations.p","rb")
else: ##Otherwise, the chemical equations and rates are loaded in accordance with the model chosen.
reaction_and_rate_selector.reaction_writer(chemical_model)
pickle_file1 =open("reactionratearray.p","rb")
pickle_file4 = open("reaction_equations.p","rb")
pickle.dump(chemical_model,open("previous_use.p","wb"))
array_size_lookup = {1:92 , 2:39 , 3 : 20, 4 : 42} ##Dictionary containing the size of the array for each model.
array_size = array_size_lookup[chemical_model] ##Returns the size of the array for the chosen model.
rate_array = pickle.load(pickle_file1) ##Loads the reactions, cooling rates, photoionisation rates and regular rates
reactions = pickle.load(pickle_file4)
return (rate_array,array_size,reactions)
################################### end of program functions ##############################################
global reaction_equations ##Define a variable ready to be filled with the reaction equations
rate_array, array_size,reactions = file_unpacking(chemical_model) ##Unpack the equations
reaction_equations = [0] * len(reactions) ##Turn the unpacked chemical equations into something the program can use,
##In this case each reaction is converted into a lambda function
for i in range(0,len(reactions)):
if (reactions[i] == ''):
reaction_equations[i] = lambda k,Z,T_m,T_r,y : 0.0
else:
reaction_equations[i] = eval('lambda k,Z,T_m,T_r,y: ' + reactions[i])
k = [0] * array_size ##Turn the unpacked reaction rates into something the program can use,
##lambda functions once again.
for i in range(0,array_size):
k[i] = eval(rate_array[i][9])
##Now we need a cooling function for the gas, each term within the function is due to a different cooling/heating process
cooling_function = lambda T_m,T_r,n_H,n_H2,n_HD,n_e,n_Hplus,n_H2plus,n_He,n_Heplus,n_Heplusplus,n_LiH: ((-1.017*(10**-37)*(T_r**4)*(T_m-T_r)*n_e*n_H) + \
(1.27*(10**-21)*(T_m**0.5)*((1+(T_m*10**-5)**0.5)**-1.0)*exp(-157809.1/T_m))*n_e*n_H + \
(9.38*(10**-22)*(T_m**0.5)*((1+(T_m*10**-5)**0.5)**-1.0)*exp(-285335.4/T_m))*n_e*n_He + \
(4.95*(10**-22)*(T_m**0.5)*((1+(T_m*10**-5)**0.5)**-1.0)*exp(-631515.0/T_m))*n_e*n_Heplus + \
(8.70*(10**-27)*(T_m**0.5)*((T_m*10**-3)**-0.2)*((1+(T_m*10**-6)**0.7)**-1.0))*n_e*n_Hplus + \
(1.55*(10**-26)*(T_m**0.3647))*n_e*n_Heplus + \
(3.48*(10**-26)*(T_m**0.5)*((T_m*10**-3)**-0.2)*((1+(T_m*10**-6)**0.7)**-1.0))*n_e*n_Heplusplus + \
(1.24*(10**-13)*(T_m**-1.5)*exp(-47000.0/T_m)*(1+0.3*exp(-94000.0/T_m)))*n_e*n_Heplus + \
(7.50*(10**-19)*((1+(T_m*10**-5)**0.5)**-1)*exp(-118348.0/T_m))*n_e*n_H +\
(5.54*(10**-17)*(T_m**-0.397)*((1+(T_m*10**-5)**0.5)**-1.0)*exp(-473638.0/T_m))*n_e*(n_Heplus) + \
(1.42*(10**-27)*(T_m**0.5)*(1.10+0.34*exp(-((5.50-log10(T_m))**2)/3.0)))*(n_e*(n_Hplus+n_Heplus+4*n_Heplusplus)) + \
(10**(-103.0 + (97.59*log10(T_m))-(48.05*((log10(T_m))**2.0))+(10.80*((log10(T_m))**3))-(0.9032*((log10(T_m))**4))))*n_H2*n_H + \
(1.1*(10**-19)*(T_m**-0.34)*exp(-3025.0/T_m))*n_H2plus*n_e + \
(1.36*(10**-22)*exp(-3152.0/T_m))*n_H2plus*n_H + \
(10**(-36.42+5.95*log10(T_m)-0.526*((log10(T_m))**2)))*n_Hplus*n_H + \
(10**(-42.45906+(21.90083*T_m)-(10.1954*(T_m**2))+(2.19788*(T_m**3))-(0.17286*(T_m**4))))*n_H*n_HD + \
(10**(-31.47+(8.817*(log10(T_m)))-(4.1444*((log10(T_m))**2))+(0.8292*((log10(T_m))**3))-(0.04996*((log10(T_m))**4))))*n_LiH)/(kbev)
##Determine whether or not a source function was entered, if it was, then find the relevant flux for the reactions
if source_function != "0.0":
flux = photo_ionisation_calculator(source_function)
else:
flux = [0] * 12
################################# BEGINNING OF INTEGRATION PROCEDURE ###########################################################
##Begin integration procedure
print "\n"
Z = Z_initial
yinit, error_test_begin = initial_conditions(Z_initial,number_of_particles,manual_matter_temperature) ##Get the initial conditions, and get the initial values ready
##for the error testing function to make use of
yd0init = deriv(yinit,0.0) ##Get the jacobian for IDA to use (faster solving)
Z_initial_s = new_ageconverter.age_converter(Z_initial) ##Convert Z value into age of universe
Z_final_s = new_ageconverter.age_converter(Z_initial-step_size) ##Convert Z value into age of universe
time_range = range_calculator(Z_final_s,Z_initial_s) ##Calculate the actual time in seconds to integrate over by using the step size.
model = Implicit_Problem(residual,yinit,yd0init,0.0) ##Set up the problem
sim = IDA(model)
sim.verbosity = 50 ##Set verbosity to low(not needed)
t,y,yd = sim.simulate(time_range) ##Perform the simulation for one step
eoa = (shape(y)[0])-1
end_line = y[(shape(y)[0])-1,:]
yd0init = yd[eoa,:] ##Provide the new jacobian from the newly calculated results, ready for the next step
plot_values = y[eoa,:]
previous_time_range = time_range
no_steps = (Z_initial-Z_final)/step_size ##calculate the total number of steps needed to complete the process
count = 1 ##Used to keep track of integration progress
pcent = (100/no_steps)*count ##Integration progress as a percentage
for i in arange((Z_initial-step_size),Z_final-step_size,-step_size): ##Z_final+step_size is used since it then subtracts the step size in the iteration
##Recompute the flux for some reactions as Z has changed
Z = i
Z_initial_s = new_ageconverter.age_converter(i) ##Convert Z value into age of universe
Z_final_s = new_ageconverter.age_converter(i-step_size) ##Convert Z value into age of universe
time_range = range_calculator(Z_final_s,Z_initial_s) ##Calculate the actual time in seconds to integrate over
factor = (((1.0+i)/(1.0+i+step_size))**3) ##Upon each step, each chemical species is multiplied by this factor to account for the expanding Universe
pcent = (100/no_steps)*count ##Print the progress percentage and value for the user
count = count + 1
sys.stdout.write("\rCurrent Value, Z = " + str(i) + "\t" + str(pcent) + "%")
sys.stdout.flush()
##re-initialise the values ready for integration during the next step
##each of the number densities is reduced by a factor of step-size**3
##due to the expansion of the universe in this time
##the radiation temperature and matter temperature are set according
##to the redshift, over Z = 200 the two are coupled together, and so
##are equal, after this the matter temperature evolves adiabtically
##The cooling function also comes into effect at the lower Z values
if i > 200:
yinit = [x * factor for x in end_line[0:6]] + \
[end_line[6]*(((1.0+i)/(1.0+i+step_size))**4)] + \
[(T_0*(1.0+i))] + [(T_0*(1.0+i))] + \
[x * factor for x in end_line[9:25]]
else:
yinit = [x * factor for x in end_line[0:6]] + \
[end_line[6]*(((1.0+i)/(1.0+i+step_size))**4)] + \
[(T_0*(1.0+i))] + [((T_0*((1.0+i)**2))/201.0)-cooling_function(end_line[8],end_line[7],end_line[0],end_line[3],end_line[12],end_line[5],end_line[1],end_line[4],end_line[14],end_line[15],end_line[16],end_line[21])] + \
[x * factor for x in end_line[9:25]]
yd0init = yd[eoa,:] ##Provide the new jacobian from the newly calculated results, ready for the next step
model = Implicit_Problem(residual,yinit,yd0init,previous_time_range) ##Simulate again, using new step
sim = IDA(model)
sim.verbosity = 50
t,y,yd = sim.simulate(time_range+previous_time_range)
eoa = (shape(y)[0])-1
end_line = y[eoa,:] ##store plot values, alter time range and get ready for the next step.
previous_time_range = time_range + previous_time_range
plot_values = vstack((plot_values,y[eoa,:]))
if (error_checking(error_test_begin,end_line,Z_initial,i) < 10**-10) :
print("\n\nError test results: PASSED")
else:
print("\nError test results: FAILED")
final_calculated_values = array(plot_values[(shape(plot_values)[0])-1][:]) ##Finds the final calculated abundacnes, ready to return to the user
print("\nFinal calculated abundances are as follows (number densities are in cm^3) :")
print "\n[H]:\n"
print end_line[0]
print "\n[H+]:\n"
print end_line[1]
print "\n[H-]:\n"
print end_line[2]
print "\n[H2]:\n"
print end_line[3]
print "\n[H2+]:\n"
print end_line[4]
print "\n[e-]:\n"
print end_line[5]
print "\n[photon density]:\n"
print end_line[6]
if plot_request == 2: ##If the user wanted to plot the results, this series of commands is run
plotter.plotting_program(Z_initial,Z_final,plot_values)
return (final_calculated_values) ##Return the final calculated values to the user
def run_example(with_plots=True):
"""
This example show how to use Assimulo and IDA for simulating sensitivities
for initial conditions.::
0 = dy1/dt - -(k01+k21+k31)*y1 - k12*y2 - k13*y3 - b1
0 = dy2/dt - k21*y1 + (k02+k12)*y2
0 = dy3/dt - k31*y1 + k13*y3
y1(0) = p1, y2(0) = p2, y3(0) = p3
p1=p2=p3 = 0
See http://sundials.2283335.n4.nabble.com/Forward-sensitivities-for-initial-conditions-td3239724.html
"""
def f(t, y, yd, p):
y1, y2, y3 = y
yd1, yd2, yd3 = yd
k01 = 0.0211
k02 = 0.0162
k21 = 0.0111
k12 = 0.0124
k31 = 0.0039
k13 = 0.000035
b1 = 49.3
res_0 = -yd1 - (k01 + k21 + k31) * y1 + k12 * y2 + k13 * y3 + b1
res_1 = -yd2 + k21 * y1 - (k02 + k12) * y2
res_2 = -yd3 + k31 * y1 - k13 * y3
return N.array([res_0, res_1, res_2])
#The initial conditions
y0 = [0.0, 0.0, 0.0] #Initial conditions for y
yd0 = [49.3, 0., 0.]
p0 = [0.0, 0.0, 0.0] #Initial conditions for parameters
yS0 = N.array([[1, 0, 0], [0, 1, 0], [0, 0, 1.]])
#Create an Assimulo implicit problem
imp_mod = Implicit_Problem(f, y0, yd0, p0=p0)
#Sets the options to the problem
imp_mod.yS0 = yS0
#Create an Assimulo explicit solver (IDA)
imp_sim = IDA(imp_mod)
#Sets the paramters
imp_sim.rtol = 1e-7
imp_sim.atol = 1e-6
imp_sim.pbar = [
1, 1, 1
] #pbar is used to estimate the tolerances for the parameters
imp_sim.continuous_output = True #Need to be able to store the result using the interpolate methods
imp_sim.sensmethod = 'SIMULTANEOUS' #Defines the sensitvity method used
imp_sim.suppress_sens = False #Dont suppress the sensitivity variables in the error test.
#Simulate
t, y, yd = imp_sim.simulate(400) #Simulate 400 seconds
#Basic test
nose.tools.assert_almost_equal(y[-1][0], 1577.6552477, 3)
nose.tools.assert_almost_equal(y[-1][1], 611.9574565, 3)
nose.tools.assert_almost_equal(y[-1][2], 2215.88563217, 3)
nose.tools.assert_almost_equal(imp_sim.p_sol[0][1][0], 1.0)
#Plot
if with_plots:
P.figure(1)
P.subplot(221)
P.plot(t,
N.array(imp_sim.p_sol[0])[:, 0], t,
N.array(imp_sim.p_sol[0])[:, 1], t,
N.array(imp_sim.p_sol[0])[:, 2])
P.title("Parameter p1")
P.legend(("p1/dy1", "p1/dy2", "p1/dy3"))
P.subplot(222)
P.plot(t,
N.array(imp_sim.p_sol[1])[:, 0], t,
N.array(imp_sim.p_sol[1])[:, 1], t,
N.array(imp_sim.p_sol[1])[:, 2])
P.title("Parameter p2")
P.legend(("p2/dy1", "p2/dy2", "p2/dy3"))
P.subplot(223)
P.plot(t,
N.array(imp_sim.p_sol[2])[:, 0], t,
N.array(imp_sim.p_sol[2])[:, 1], t,
N.array(imp_sim.p_sol[2])[:, 2])
P.title("Parameter p3")
P.legend(("p3/dy1", "p3/dy2", "p3/dy3"))
P.subplot(224)
P.plot(t, y)
P.show()
def run_example(with_plots=True):
#Defines the residual
def f(t, y, yd):
res_0 = yd[0] - y[2]
res_1 = yd[1] - y[3]
res_2 = yd[2] + y[4] * y[0]
res_3 = yd[3] + y[4] * y[1] + 9.82
#res_4 = y[0]**2+y[1]**2-1
res_4 = y[2]**2 + y[3]**2 - y[4] * (y[0]**2 + y[1]**2) - y[1] * 9.82
return N.array([res_0, res_1, res_2, res_3, res_4])
#Defines the jacobian
def jac(c, t, y, yd):
jacobian = N.zeros([len(y), len(y)])
#Derivative
jacobian[0, 0] = 1 * c
jacobian[1, 1] = 1 * c
jacobian[2, 2] = 1 * c
jacobian[3, 3] = 1 * c
#Differentiated
jacobian[0, 2] = -1
jacobian[1, 3] = -1
jacobian[2, 0] = y[4]
jacobian[3, 1] = y[4]
jacobian[4, 0] = y[0] * 2 * y[4] * -1
jacobian[4, 1] = y[1] * 2 * y[4] * -1 - 9.82
jacobian[4, 2] = y[2] * 2
jacobian[4, 3] = y[3] * 2
#Algebraic
jacobian[2, 4] = y[0]
jacobian[3, 4] = y[1]
jacobian[4, 4] = -(y[0]**2 + y[1]**2)
return jacobian
#The initial conditons
y0 = [1.0, 0.0, 0.0, 0.0, 5] #Initial conditions
yd0 = [0.0, 0.0, 0.0, -9.82, 0.0] #Initial conditions
#Create an Assimulo implicit problem
imp_mod = Implicit_Problem(f, y0, yd0)
#Sets the options to the problem
imp_mod.jac = jac #Sets the jacobian
imp_mod.algvar = [1.0, 1.0, 1.0, 1.0, 0.0] #Set the algebraic components
imp_mod.name = 'Test Jacobian'
#Create an Assimulo implicit solver (IDA)
imp_sim = IDA(imp_mod) #Create a IDA solver
#Sets the paramters
imp_sim.atol = 1e-6 #Default 1e-6
imp_sim.rtol = 1e-6 #Default 1e-6
imp_sim.suppress_alg = True #Suppres the algebraic variables on the error test
#Let Sundials find consistent initial conditions by use of 'IDA_YA_YDP_INIT'
imp_sim.make_consistent('IDA_YA_YDP_INIT')
#Simulate
t, y, yd = imp_sim.simulate(
5, 1000) #Simulate 5 seconds with 1000 communication points
#Basic tests
nose.tools.assert_almost_equal(y[-1][0], 0.9401995, places=4)
nose.tools.assert_almost_equal(y[-1][1], -0.34095124, places=4)
nose.tools.assert_almost_equal(yd[-1][0], -0.88198927, places=4)
nose.tools.assert_almost_equal(yd[-1][1], -2.43227069, places=4)
#Plot
if with_plots:
P.plot(t, y)
P.show()
