pyrun.py

#run the primordial chem network one zone test in python
from mynet import rates, sym_rates, massfracs
import constants as cons

import rate
import rate_collection as ratecol

import numpy as np
from scipy.integrate import solve_ivp
import functools
from pylab import genfromtxt
import sympy as sp
from operator import mul, add


method = int(input('enter 1 for pure python calculation, 2 for sympy calculation '))

if method == 1:

    object = ratecol.ChemRateCollection(rates=list(rates.values()), tdot_switch=0)
    print('Using dT/dt')

    def main_rhs(t,y):
        return object.rhs(t,y)

    def main_fft(y):
        return object.fft(y)

elif method == 2:

    if massfracs == 1:
        raise ValueError('Sympy calculation can only work with number densities!')

    object = ratecol.SympyChemRateCollection(rates=list(sym_rates.values()), tdot_switch=0, ydots_lambdified=True)

    def main_rhs(t,y):
        return object.rhs_numdens(t,y)

    def main_fft(y):
        return object.fft_numdens(y)

    def main_jac(t,y):
        return object.jac_numdens(t,y)

else:
    raise ValueError('Invalid choice!')


def normalize(val):
    aal = [max(0,q) for q in val]
    bal = [q/sum(aal) for q in aal]
    return bal

masses = [cons.mass_elec, cons.mass_hp, cons.mass_h, cons.mass_hm, cons.mass_dp, cons.mass_d, cons.mass_h2p, \
          cons.mass_dm, cons.mass_h2, cons.mass_hdp, cons.mass_hd, cons.mass_hepp, cons.mass_hep, cons.mass_he]

krome = genfromtxt('datafiles/fort22_wD.dat').T

ntot = 1.0 #number density
Tgas = [1e2]
Eint = [10254504084.93553]

y0_labels = ['elec', 'H+', 'H', 'H-', 'D+', 'D', 'H2+', 'D-', 'H2', 'HD+', 'HD', 'He++', 'He+', 'He', 'Tgas']

#number densities of each specie
y0 = [1e-4, 1e-4, 1e0, 1e-40, 1e-40, 1e-40, 1e-40, 1e-40, 1e-6, 1e-40, 1e-40, 1e-40, 1e-40, 0.0775e0]


y0 = [x*ntot for x in y0]
dd = ntot

abd_elec = []
abd_hp = []
abd_h = []
abd_hm = []
abd_dp = []
abd_d = []
abd_h2p = []
abd_dm = []
abd_h2 = []
abd_hdp = []
abd_hd = []
abd_hepp = []
abd_hep = []
abd_he = []

temp = []
time = []
dens = []

rstep = 1
rstep_max = int(input('Input max number of steps '))
tff_fac = float(input('Input fac, where dt = fac * tff '))

while rstep < rstep_max:

    dd1 = dd

    #at each step, y0 just consists of the species, nont the temperature (see below)
    tff = main_fft(y0)

    dtH = tff_fac * tff     #define time-step
    deldd = (dd/tff) * dtH #density increase
    dd = dd + deldd        #update density

    #add temperature separately here
    if rstep == 1:
         y0 = [x*dd/dd1 for x in y0] + Tgas
    else:
         y0 = [x*dd/dd1 for x in y0] + [sol.y[14][0]]

    dt = dtH

    if dd > 1e18:
        break

    if dt < 10:
        break

    #charge conservation  
    y0[0] = -y0[3] - y0[7] + y0[1] + y0[12] + y0[6] + y0[4] + y0[9] + 2.0*y0[11]

    #integrate from 0 till dt, and put sol.y at dt as the updated y0

    if method == 1:
        sol = solve_ivp(main_rhs, t_span=[0,dt], t_eval=[dt], y0=y0, method='BDF', atol=1e-4, rtol=1e-4)
    elif method == 2:
        sol = solve_ivp(fun=lambda t, y: main_rhs(t,y), t_span=[0,dt], t_eval=[dt], y0=y0, method='BDF', atol=1e-4, rtol=1e-4)
        #mass conservation?
        #ns = sol.y.flatten()[:-1]
        #xs = [x*y for x,y in zip(ns, masses)] #ns*mu*mh
        #rho = sum(xs)
        #xs = [x/rho for x in xs] #ns*mu*mh/rho = xs
        #xs = normalize(xs)
        #a1 = [x/y for x,y in zip(xs, masses)] #xs/(mu*mh)
        #y0 = [rho*x for x in a1] #xs*rho/(mu*mh) = ns

    if not sol.success:
        breakpoint()

    #charge conservation  
    y0[0] = -y0[3] - y0[7] + y0[1] + y0[12] + y0[6] + y0[4] + y0[9] + 2.0*y0[11]

    #recreate y0 but do not include temperature
    idx = 0
    y0 = [sol.y[0][idx], sol.y[1][idx], sol.y[2][idx], sol.y[3][idx], sol.y[4][idx], sol.y[5][idx], sol.y[6][idx], sol.y[7][idx], \
          sol.y[8][idx], sol.y[9][idx], sol.y[10][idx], sol.y[11][idx], sol.y[12][idx], sol.y[13][idx]]

    temp.append(sol.y[14][idx])

    time.append(dt)
    dens.append(dd)

    abd_elec.append(y0[0])
    abd_hp.append(y0[1])
    abd_h.append(y0[2])
    abd_hm.append(y0[3])
    abd_dp.append(y0[4])
    abd_d.append(y0[5])
    abd_h2p.append(y0[6])
    abd_dm.append(y0[7])
    abd_h2.append(y0[8])
    abd_hdp.append(y0[9])
    abd_hd.append(y0[10])
    abd_hepp.append(y0[11])
    abd_hep.append(y0[12])
    abd_he.append(y0[13])

    rhotot = sum([a*b for a,b in zip(y0,masses)])

    print('%0.10e %0.10e %0.10e %d' %(dd, rhotot, sol.y[14][idx], rstep))

    rstep += 1