rate_new_photo.py

import os
import re
import io
import numpy as np
import matplotlib.pyplot as plt
from collections import Counter

import sys
#sys.path.append(os.path.abspath("/scratch/jh2/ps3459/pynucastro/pynucastro/rates"))
import constants as cons

class ChemSpecie:

    """
    a class like the nucleus class above but for chemical species in the ISM
    :var Z:               atomic number
    :var m:               total mass of specie in g
    :var N:               neutron number
    :var e:               number of electrons
    :var gamma:           adiabatic index
    :var chemsign:            chemical sign
    """
    def __init__(self, name, dummy=False):

        #importing sympy here because it interferes with the rest of pynucastro stuff if imported at the top
        import sympy as sp

        # a dummy chemical specie is one that we can use where a chemical specie is needed
        # but it is not considered to be part of the network
        self.dummy = dummy
        self.raw = name

        self.num = 0
        self.end = 1
        # element symbol and atomic weight
        if name.casefold() == "elec" or name.casefold() == 'e':
            self.Z = 0 #number of protons
            self.m = 9.10938188e-28 #mass in g
            self.N = 0 #number of neutrons
            self.e = 1 #number of electons
            self.gamma = 1.66666667 #adiabatic index
            self.chemsign = "E"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "hp" or name.casefold() == "h+":
            self.Z = 1
            self.m = 1.67262158e-24
            self.N = 0
            self.e = 0
            self.gamma = 1.66666667
            self.chemsign = "H+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "h":
            self.Z = 1
            self.m = 1.67353251819e-24
            self.N = 0
            self.e = 1
            self.gamma = 1.66666667
            self.chemsign = "H"
            self.A = self.Z + self.N
            self.kevap70 = 790490323.1199661
            self.EbindBare = 500.0
        elif name.casefold() == "hm" or name.casefold() == 'h-':
            self.Z = 1
            self.m = 1.67444345638e-24
            self.N = 0
            self.e = 2
            self.gamma = 1.66666667
            self.chemsign = "H-"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "dp" or name.casefold() == "d+":
            self.Z = 1
            self.m = 3.34512158e-24
            self.N = 1
            self.e = 0
            self.gamma = 1.66666667
            self.chemsign = "D+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "d":
            self.Z = 1
            self.m = 3.34603251819e-24
            self.N = 1
            self.e = 1
            self.gamma = 1.66666667
            self.chemsign = "D"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "h2p" or name.casefold() == "h2+":
            self.Z = 2
            self.m = 3.34615409819e-24
            self.N = 0
            self.e = 1
            self.gamma = 1.4
            self.chemsign = "H2+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "dm" or name.casefold() == "d-":
            self.Z = 1
            self.m = 3.34694345638e-24
            self.N = 1
            self.e = 2
            self.gamma = 1.66666667
            self.chemsign = "D-"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "h2":
            self.Z = 2
            self.m = 3.34706503638e-24
            self.N = 0
            self.e = 2
            self.gamma = 1.4
            self.chemsign = "H2"
            self.A = self.Z + self.N
            self.kevap70 = 13763786733.050402
            self.EbindBare = 300.0
        elif name.casefold() == "hdp" or name.casefold() == "hd+":
            self.Z = 2
            self.m = 5.01865409819e-24
            self.N = 1
            self.e = 1
            self.gamma = 1.4
            self.chemsign = "HD+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "hd":
            self.Z = 2
            self.m = 5.01956503638e-24
            self.N = 1
            self.e = 2
            self.gamma = 1.4
            self.chemsign = "HD"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "hepp" or name.casefold() == "he++":
            self.Z = 2
            self.m = 6.69024316e-24
            self.N = 2
            self.e = 0
            self.gamma = 1.66666667
            self.chemsign = "HE++"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "hep" or name.casefold() == "he+":
            self.Z = 2
            self.m = 6.69115409819e-24
            self.N = 2
            self.e = 1
            self.gamma = 1.66666667
            self.chemsign = "HE+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "he":
            self.Z = 2
            self.m = 6.69206503638e-24
            self.N = 2
            self.e = 2
            self.gamma = 1.66666667
            self.chemsign = "He"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "h3p" or name.casefold() == "h3+":
            self.Z = 3
            self.m = 5.0196866163760004e-24
            self.N = 0
            self.e = 2
            self.gamma = 1.66666667
            self.chemsign = "H3+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "cm" or name.casefold() == "c-":
            self.Z = 6
            self.m = 2.0077106047315998e-23
            self.N = 6
            self.e = 7
            self.gamma = 1.66666667
            self.chemsign = "C-"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "c":
            self.Z = 6
            self.m = 2.0076195109128e-23
            self.N = 6
            self.e = 6
            self.gamma = 1.66666667
            self.chemsign = "C"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "cp" or name.casefold() == "c+":
            self.Z = 6
            self.m = 2.007528417094e-23
            self.N = 6
            self.e = 5
            self.gamma = 1.66666667
            self.chemsign = "C+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "ch":
            self.Z = 7
            self.m = 2.1749727627316e-23
            self.N = 6
            self.e = 7
            self.gamma = 1.4
            self.chemsign = "CH"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "chp" or name.casefold() == "ch+":
            self.Z = 7
            self.m = 2.1748816689128e-23
            self.N = 6
            self.e = 6
            self.gamma = 1.4
            self.chemsign = "CH+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "om" or name.casefold() == "o-":
            self.Z = 8
            self.m = 2.6769171083692e-23
            self.N = 8
            self.e = 9
            self.gamma = 1.66666667
            self.chemsign = "O-"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "o":
            self.Z = 8
            self.m = 2.6768260145504e-23
            self.N = 8
            self.e = 8
            self.gamma = 1.66666667
            self.chemsign = "O"
            self.A = self.Z + self.N
            self.kevap70 = 28.369278883298623
            self.EbindBare = 1700.0
        elif name.casefold() == "ch2":
            self.Z = 8
            self.m = 2.3423260145503998e-23
            self.N = 6
            self.e = 8
            self.gamma = 1.4
            self.chemsign = "CH2"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "ch3":
            self.Z = 9
            self.m = 2.5096792663692e-23
            self.N = 6
            self.e = 9
            self.gamma = 1.4
            self.chemsign = "CH3"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "ch4":
            self.Z = 10
            self.m = 2.677032518188e-23
            self.N = 6
            self.e = 10
            self.gamma = 1.4
            self.chemsign = "CH4"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "op" or name.casefold() == "o+":
            self.Z = 8
            self.m = 2.6767349207316e-23
            self.N = 8
            self.e = 7
            self.gamma = 1.66666667
            self.chemsign = "O+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "ch2p" or name.casefold() == "ch2+":
            self.Z = 8
            self.m = 2.3422349207316e-23
            self.N = 6
            self.e = 7
            self.gamma = 1.4
            self.chemsign = "CH2+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "oh":
            self.Z = 9
            self.m = 2.8441792663692003e-23
            self.N = 8
            self.e = 9
            self.gamma = 1.66666667
            self.chemsign = "OH"
            self.A = self.Z + self.N
            self.kevap70 = 3649.8804308076387
            self.EbindBare = 1360.0
        elif name.casefold() == "ch3p" or name.casefold() == "ch3+":
            self.Z = 9
            self.m = 2.5095881725504002e-23
            self.N = 6
            self.e = 8
            self.gamma = 1.4
            self.chemsign = "CH3+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "ohp" or name.casefold() == "oh+":
            self.Z = 9
            self.m = 2.8440881725504e-23
            self.N = 8
            self.e = 8
            self.gamma = 1.4
            self.chemsign = "OH+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "h2o":
            self.Z = 10
            self.m = 3.011532518188e-23
            self.N = 8
            self.e = 10
            self.gamma = 1.4
            self.chemsign = "H2O"
            self.A = self.Z + self.N
            self.kevap70 = 1.6588493815567146e-18
            self.EbindBare = 4800.0
        elif name.casefold() == "h2o_total":
            self.Z = 10
            self.m = 0.0
            self.N = 8
            self.e = 10
            self.gamma = 1.4
            self.chemsign = "H2O_total"
            self.A = self.Z + self.N
            self.kevap70 = 0.0
            self.EbindBare = 0.0
        elif name.casefold() == "h2op" or name.casefold() == "h2o+":
            self.Z = 10
            self.m = 3.0114414243692e-23
            self.N = 8
            self.e = 9
            self.gamma = 1.4
            self.chemsign = "H2O+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "h3op" or name.casefold() == "h3o+":
            self.Z = 11
            self.m = 3.178794676188e-23
            self.N = 8
            self.e = 10
            self.gamma = 1.4
            self.chemsign = "H3O+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "c2":
            self.Z = 12
            self.m = 4.0152390218256e-23
            self.N = 12
            self.e = 12
            self.gamma = 1.4
            self.chemsign = "C2"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "co":
            self.Z = 14
            self.m = 4.6844455254632e-23
            self.N = 14
            self.e = 14
            self.gamma = 1.4
            self.chemsign = "CO"
            self.A = self.Z + self.N
            self.kevap70 = 149751.92964078687
            self.EbindBare = 1100.0
        elif name.casefold() == "co_total":
            self.Z = 14
            self.m = 0.0
            self.N = 14
            self.e = 14
            self.gamma = 1.4
            self.chemsign = "CO_total"
            self.A = self.Z + self.N
            self.kevap70 = 0.0
            self.EbindBare = 0.0
        elif name.casefold() == "co2":
            self.Z = 22
            self.m = 7.3612715400136e-23
            self.N = 22
            self.e = 22
            self.gamma = 1.4
            self.chemsign = "CO2"
            self.A = self.Z + self.N
            self.kevap70 = 0.0053743279721931055
            self.EbindBare = 2300.0
        elif name.casefold() == "cop" or name.casefold() == "co+":
            self.Z = 14
            self.m = 4.6843544316444e-23
            self.N = 14
            self.e = 13
            self.gamma = 1.4
            self.chemsign = "CO+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "hco":
            self.Z = 15
            self.m = 4.8517987772819996e-23
            self.N = 14
            self.e = 15
            self.gamma = 1.4
            self.chemsign = "HCO"
            self.A = self.Z + self.N
            self.kevap70 = 149751.92964078687
            self.EbindBare = 1100.0
        elif name.casefold() == "hocp" or name.casefold() == "hoc+":
            self.Z = 15
            self.m = 4.8517076834631994e-23
            self.N = 14
            self.e = 15
            self.gamma = 1.4
            self.chemsign = "HOC+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "hcop" or name.casefold() == "hco+":
            self.Z = 15
            self.m = 4.8517076834631994e-23
            self.N = 14
            self.e = 14
            self.gamma = 1.4
            self.chemsign = "HCO+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "o2":
            self.Z = 16
            self.m = 5.3536520291008e-23
            self.N = 16
            self.e = 16
            self.gamma = 1.4
            self.chemsign = "O2"
            self.A = self.Z + self.N
            self.kevap70 = 17568.77150646201
            self.EbindBare = 1250.0
        elif name.casefold() == "o2p" or name.casefold() == "o2+":
            self.Z = 16
            self.m = 5.353560935282e-23
            self.N = 16
            self.e = 15
            self.gamma = 1.4
            self.chemsign = "O2+"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        elif name.casefold() == "dummy":
            self.Z = 0
            self.m = 0
            self.N = 0
            self.e = 0
            self.gamma = 0
            self.chemsign = "dummy"
            self.A = self.Z + self.N
            self.kevap70 = 0
            self.EbindBare = 1e99
        else:
            raise UnsupportedChemSpecie()

        self.sym_name = sp.symbols(self.chemsign, real=True)


    def __iter__(self):
        #print('iterrr')
        #yield 
        #yield from {
         #   "Z": self.Z,
         #   "m": self.m,
         #   "N": self.N,
         #   "e": self.e,
         #   "gamma": self.gamma,
         #   "chemsign": self.chemsign
        #}.items()
        return self

    def __next__(self):
        if self.num > self.end:
            raise StopIteration
        else:
            self.num += 1
            return self.num - 1


    def __repr__(self):
        return self.chemsign

    def __hash__(self):
        return hash((self.Z, self.m, self.N, self.e, self.gamma, self.chemsign, self.A, self.sym_name, self.kevap70, self.EbindBare))

    def __eq__(self, other):
        #if isinstance(other, ChemSpecie):
        #    return self.Z == other.Z and self.m == other.m and \
        #       self.N == other.N and self.e == other.e and \
        #       self.gamma == other.gamma and self.chemsign == other.chemsign
        #if isinstance(other, tuple):
        #    return (self.Z, self.A) == other
        #return NotImplemented
        return (self.Z, self.m, self.N, self.e, self.gamma, self.chemsign, self.A, self.sym_name, self.kevap70, self.EbindBare) == \
               (other.Z, other.m, other.N, other.e, other.gamma, other.chemsign, other.A, other.sym_name, other.kevap70, other.EbindBare)

    def __lt__(self, other):
        return self.m < other.m

class ChemComposition:
    """a composition holds the mass fractions of the chemspecies in a network
    -- useful for evaluating the rates
    FOR NOW, I ONLY CHANGED SELF.SYMPY TO TREAT COMPOSITIONS AS MASS FRACTIONS
    REST OF THE CODE TREATS THEM AS NUMBER DENSITIES
    """
    def __init__(self, specie, small=1.e-40):
        """specie is an iterable of the specie (ChemSpecie objects) in the network"""
        if not isinstance(specie[0], ChemSpecie):
            raise ValueError("must supply an iterable of ChemSpecie object")
        self.X = {k: small for k in specie}
        self.Y = {k.sym_name: k.sym_name for k in specie}

    def set_all(self, xval):
        """ set all species to a particular value """
        for k in self.X:
            self.X[k] = xval
        return self.X

    def set_specie_massfrac(self, xval):
        """ set specie name to the mass fraction xval """

        if len(self.X) != len(xval):
            raise ValueError("length of species array does not match length of mass fractions array")

        #need a separate counter for xval tuple
        i = 0
        for k in self.X:
            self.X[k] = xval[i]
            i = i+1
        
        self.normalize()
        return self.X

    def set_specie_numberdens(self, nval):
        """ set specie name to the number density nval """

        if len(self.X) != len(nval):
            raise ValueError("length of species array does not match length of number densities array")

        #need a separate counter for xval tuple
        i = 0
        for k in self.X:
            self.X[k] = nval[i]
            i = i+1
        
        return self.X

    def normalize(self):
        """ normalize the mass fractions to sum to 1 """
        X_sum = sum(self.X[k] for k in self.X)

        for k in self.X:
            self.X[k] /= X_sum
    
        return self.X

    def get_specie_numberdens(self, xval, rho):

        massfracs = self.set_specie_massfrac(xval)
        for k in self.X:
            self.X[k] = massfracs[k] * rho / k.m

        return self.X

    def sympy(self):
        #for k in self.Y:
        #    self.Y[k] = k.sym_name

        return self.Y

    def __str__(self):
        ostr = ""
        for k in self.X:
            ostr += f"  X({k}) : {self.X[k]}\n"
        return ostr




class SympyChemRate:
    def __init__(self, reactants=[], products=[], Tmins=None, Tmaxs=None, rate_expr=None, photo=False):
        import sympy as sp
        self.reactants = reactants

        self.products = products

        self.rate_expr = rate_expr

        self.photo = photo

        def rate_function(Tgas, Te, invTe, invT, lnTe, T32, invT32, invsqrT, composition, user_crate, user_Av, 
                          user_ionH, user_ionH2, user_dissH2, user_ionC, user_ionO, user_dissCO, user_dust2gas_ratio, Tdust, chiFUV):
            if rate_expr is not None:

                #if there is no Tmin
                if Tmins is None:
                    #if there is no Tmax
                    if Tmaxs is None:
                        rate = self.get_small(composition) + rate_expr
                    #if there is Tmax
                    else:
                        rate = self.get_small(composition) + sp.Piecewise((rate_expr, Tgas <= Tmaxs), (0, True))

                #if there is Tmin
                else:
                    #if there is no Tmax
                    if Tmaxs is None:
                        rate = self.get_small(composition) + sp.Piecewise((rate_expr, Tgas > Tmins), (0, True))
                    #if there is Tmin and Tmax
                    else:
                        rate = self.get_small(composition) + sp.Piecewise((rate_expr, (Tgas > Tmins) & (Tgas <= Tmaxs)), (0, True))
            else:
                raise ValueError('Rate not valid for reaction: ', self.reactants, ' --> ', self.products)

            return rate

        self.rate_function = rate_function

    def get_Hnuclei(self, composition):
        import sympy as sp
        #TODO: HARDCODED for now
        nH = composition[ChemSpecie('hp').sym_name] + composition[ChemSpecie('h').sym_name] + composition[ChemSpecie('hm').sym_name] + \
             composition[ChemSpecie('h2').sym_name]*2.0 + composition[ChemSpecie('h2p').sym_name]*2.0 + composition[ChemSpecie('oh').sym_name] + \
             composition[ChemSpecie('ch').sym_name] + composition[ChemSpecie('ch2').sym_name]*2.0 + composition[ChemSpecie('ch3').sym_name]*3.0 + \
             composition[ChemSpecie('ch4').sym_name]*4.0 + \
             composition[ChemSpecie('ohp').sym_name] + composition[ChemSpecie('h2op').sym_name]*2.0 + \
             composition[ChemSpecie('h3op').sym_name]*3.0 + composition[ChemSpecie('h2o_total').sym_name]*2.0
        if ChemSpecie('hd').sym_name in composition:
            nH += composition[ChemSpecie('hd').sym_name] + composition[ChemSpecie('hdp').sym_name]

        if self.photo:
            nH += composition[ChemSpecie('chp').sym_name] + composition[ChemSpecie('ch2p').sym_name]*2.0 + composition[ChemSpecie('h3p').sym_name]*3.0

        return nH

    def get_rho(self, composition):
        import sympy as sp
        mass = [ChemSpecie(str(key).lower()).m for key in composition.keys()]
        cc = [x*y for x,y in zip(composition, mass)]
        #density is just sum(number_density*mass)
        rho = sp.Add(*cc)
        return rho

    def get_small(self, composition):
        import sympy as sp
        #values = list(composition.values())
        #nmax = sp.Max(*values)
        #small = 1e-40/(nmax**3)
        small = 1e-99
        return small

    def fHnOj(self, user_Av):
        import sympy as sp
        #shielding function for H2O+ and H3O+
        #following Glover+2010 MNRAS sect 2.2 eqn.4
        bb = sp.Piecewise((sp.exp(-2.55*user_Av + 0.0165*user_Av**2), user_Av <= 15.0), (sp.exp(-2.8*user_Av), user_Av > 15.0))
        return bb

    def dissH2_Martin96_n_Tgas(self, Tgas, composition):
        #Collisional dissociation rate (cm-3/s) by Martin et al. 1996
        # H2+H->H+H+H
        #NOTE: the use of this rate is suggested
        # for high-density regime and in the presence of UV backgrounds.
        import sympy as sp
        k_CIDm = np.zeros((2, 21))
        k_CIDm[0] = (-178.4239, -68.42243, 43.20243, -4.633167, \
                     69.70086, 40870.38, -23705.70, 128.8953, -53.91334, \
                     5.315517, -19.73427, 16780.95, -25786.11, 14.82123, \
                     -4.890915, 0.4749030, -133.8283, -1.164408, 0.8227443, \
                     0.5864073, -2.056313)

        k_CIDm[1] = (-142.7664, 42.70741, -2.027365, -0.2582097, \
                      21.36094, 27535.31, -21467.79, 60.34928, -27.43096, \
                      2.676150, -11.28215, 14254.55, -23125.20, 9.305564, \
                     -2.464009, 0.1985955, 743.0600, -1.174242, 0.7502286, \
                     0.2358848, 2.937507)

        n_H  = self.get_Hnuclei(composition)
        logT = sp.log(Tgas, 10)
        invT = 1.0/Tgas
        logT2 = logT*logT
        logT3 = logT2*logT
        logTv = np.array([1.0, logT, logT2, logT3])
        k_CID = 0.

        i = 0
        while i < 2:
            logk_h1 = k_CIDm[i,0]*logTv[0] + k_CIDm[i,1]*logTv[1] + \
                      k_CIDm[i,2]*logTv[2] + k_CIDm[i,3]*logTv[3] + \
                      k_CIDm[i,4]*sp.log(1.0+k_CIDm[i,5]*invT, 10)

            logk_h2 = k_CIDm[i,6]*invT

            logk_l1 = k_CIDm[i,7]*logTv[0] + k_CIDm[i,8]*logTv[1] + \
                      k_CIDm[i,9]*logTv[2] + k_CIDm[i,10]*sp.log(1.0+k_CIDm[i,11]*invT, 10)

            logk_l2 = k_CIDm[i,12]*invT

            logn_c1 = k_CIDm[i,13]*logTv[0] + k_CIDm[i,14]*logTv[1] + \
                      k_CIDm[i,15]*logTv[2] + k_CIDm[i,16]*invT

            logn_c2 = k_CIDm[i,17] + logn_c1

            p = k_CIDm[i,18] + k_CIDm[i,19]*sp.exp(-Tgas/1.850e3) + \
                k_CIDm[i,20]*sp.exp(-Tgas/4.40e2)

            n_c1 = 1e1**(logn_c1)
            n_c2 = 1e1**(logn_c2)

            logk_CID = logk_h1 - (logk_h1 - logk_l1) / (1.0 + (n_H/n_c1)**p) + \
                       logk_h2 - (logk_h2 - logk_l2) / (1.0 + (n_H/n_c2)**p)

            k_CID = k_CID + 1.e1**logk_CID

            i += 1

        return self.get_small(composition) + k_CID


    def H_recombination_on_dust(self, Tgas, composition, dust2gas_ratio, chiFUV):
        import sympy as sp
        # The following functions compute the recombination rate
        # on dust for H+, He+, C+, Si+, and O+. See Weingartner&Draine 2001
        # dust2gas_ratio, D/D_sol, default is assumed equal to Z/Z_sol

        GHabing = 1.69*chiFUV
        elec_conc = composition[ChemSpecie('elec').sym_name]

        psi = GHabing * sp.sqrt(Tgas) / elec_conc

        recombination_expr = sp.Piecewise(
            (0.0, (elec_conc < 1e-20) | (GHabing <= 0) | (psi <= 0)),
            (1.225e-13 * dust2gas_ratio / (1e0 + 8.074e-6 * psi**(1.378) * (1.e0 + 5.087e2 * Tgas**(0.01586) * psi**(-0.4723 - 1.102e-5 * sp.log(Tgas)))), True)
        )

        return recombination_expr

    def He_recombination_on_dust(self, Tgas, composition, dust2gas_ratio, chiFUV):
        import sympy as sp

        GHabing = 1.69*chiFUV
        elec_conc = composition[ChemSpecie('elec').sym_name]

        psi = GHabing * sp.sqrt(Tgas) / elec_conc

        recombination_expr = sp.Piecewise(
            (0.0, (elec_conc < 1e-20) | (GHabing <= 0) | (psi <= 0)),
            (5.572e-14*dust2gas_ratio/(1.0+3.185e-7*psi**(1.512)*(1.0+5.115e3*Tgas**(3.903e-7)*psi**(-0.4956-5.494e-7*sp.log(Tgas)))), True)
        )

        return recombination_expr

    def C_recombination_on_dust(self, Tgas, composition, dust2gas_ratio, chiFUV):
        import sympy as sp

        GHabing = 1.69*chiFUV
        elec_conc = composition[ChemSpecie('elec').sym_name]

        psi = GHabing * sp.sqrt(Tgas) / elec_conc

        recombination_expr = sp.Piecewise(
            (0.0, (elec_conc < 1e-20) | (GHabing <= 0) | (psi <= 0)),
            (4.558e-13*dust2gas_ratio/(1.0+6.089e-3*psi**(1.128)*(1.0+4.331e2*Tgas**(0.04845)*psi**(-0.8120-1.333e-4*sp.log(Tgas)))), True)
        )

        return recombination_expr

    def O_recombination_on_dust(self, Tgas, composition, dust2gas_ratio, chiFUV):
        import sympy as sp

        recombination_expr = 0.25e0*self.H_recombination_on_dust(Tgas, composition, dust2gas_ratio, chiFUV)
        return recombination_expr

    def dust_stick(self, Tgas, Tdust):
        #dust sticking coefficient, following Hollenbach and McKee 1979
        import sympy as sp

        Tg100 = Tgas * 1e-2
        Td100 = Tdust * 1e-2
        dust_stick = 1.0/(1.0 + 0.4*sp.sqrt(Tg100+Td100) + 0.2*Tg100 + 0.08*Tg100**2)
        return dust_stick

    def krate_stick(self, composition, idx, Tdust, amin, amax, pexp, rho0, d2g, Tgas):
        #sticking rate (1/s), assuming power-law dust distribution
        # example rate is
        #  @format:idx,R,P,rate
        #  1,CO,CO_ice,krate_stick(n(:),idx_CO,1d-7,1d-5,-3.5,3d0,1d-2)
        # n(:): internal status array (number densities, temeperature, etc...)
        # idx : index of the sticking species, e.g. idx_CO
        # Tdust: dust temperature (assume same for all bins), K
        # amin: min grain size, cm
        # amax: max grain size, cm
        # pexp: power-law exponent, usually -3.5
        # rho0: bulk material density, g/cm3, e.g. 3 g/cm3 for silicates
        # d2g: dust to gass mass ratio, usually 0.01
        import sympy as sp
        #get masses: convert the dict keys to a string that can be used on ChemSpecie object to get ".m" (mass)
        mass = [ChemSpecie(str(key).lower()).m for key in composition.keys()]

        # Get inverse mass squared with a condition to handle cases where mass is 0 (for ices)
        imass = [0 if x == 0 else 1./sp.sqrt(x) for x in mass]

        #derived exponents
        p3 = pexp + 3.
        p4 = pexp + 4.

        #original fortran equation is: sum(n(1:nmols)*mass(1:nmols))
        #we first do the product, then use sp.Add to take the sum. * operator unpacks each element of the list and passes to sp.Add
        cc = [x*y for x,y in zip(composition, mass)]
        n_prod_mass = sp.Add(*cc)

        #total dust density, g/cm3
        rhod = n_prod_mass * d2g
        #print('imass: ', imass[idx])
        #print('rhod: ', rhod)
        #print('rho0: ', rho0)
        #print('stick: ', self.dust_stick(Tgas, Tdust))
        #print('')
        
        #compute rate (1/s) coefficient assuming normalization
        k = cons.pre_kvgas_sqrt*sp.sqrt(Tgas) * imass[idx] * rhod / (4./3.*rho0) * p4 / p3 * (amax**p3-amin**p3) / (amax**p4-amin**p4) * self.dust_stick(Tgas, Tdust)
        return k

    def krate_stickSi(self, composition, idx, Tdust, dust2gas_ratio, Tgas):
        #compact version of krate_stick
        #some default values OK for silicates
        import sympy as sp
        amin = 5e-7 #cm (min dust grain size)
        amax = 2.5e-5 #cm (max dust grain size)
        pexp = -3.5
        rho0 = 3.0 #g/cm3
        d2g = (1.0/162.0)*dust2gas_ratio #Solar dust2gas ratio is 1/162 (e.g., Sharda & Krumholz 2022)

        k = self.krate_stick(composition, idx, Tdust, amin, amax, pexp, rho0, d2g, Tgas)
        #print(k)
        #print('')
        return k

    def krate_nonthermal_evaporation(self, specie, Gnot, Av, crate, yieldd):
        #non-thermal evaporation rate (1/s) following Hollenbach 2009,
        #Gnot is the habing flux (1.78 is Draine)
        #Av is the visual extinction
        #crflux the ionization flux of cosmic rays, 1/s
        #yield is the efficiency of the photons to desorb the given molecule
        import sympy as sp
        crnot=1.3e-17
        Fnot=1e8 #desorbing photons flux, 1/s
        ap2=(3e-8)**2 #sites separation squared, cm2

        f70 = 3.16e-19*crate/crnot

        k = Gnot*Fnot*ap2*yieldd*sp.exp(-1.8*Av)
        k += f70*ChemSpecie(str(specie).lower()).kevap70
        return k

    def krate_evaporation(self, specie, Tdust):
        #evaporation rate, 1/s
        import sympy as sp
        nu0 = 1e12 #1/s
        Ebind = ChemSpecie(str(specie).lower()).EbindBare
        k = nu0 * sp.exp(-Ebind/Tdust)
        return k

    def eval(self, T, composition, user_crate, user_Av, user_ionH, user_ionH2, user_dissH2, user_ionC, user_ionO, user_dissCO, 
             user_dust2gas_ratio, Tdust, chiFUV):
        import sympy as sp
        Tgas = T
        Te = Tgas*8.617343e-5 #CHECK KROME FILES!!!
        invTe = 1.0/Te
        invT = 1.0/T
        invTgas = invT
        lnTe = sp.log(Te)
        T32 = Tgas/3e2
        invT32 = 1.0/T32
        invsqrT=1e0/sp.sqrt(Tgas)

        ntot = sum(composition.values())
        Hnuclei = self.get_Hnuclei(composition)
        nidx_H = composition[ChemSpecie('h').sym_name]
        HnOj = self.get_small(composition) + self.fHnOj(user_Av)

        dissH2_Martin96_n_Tgas = self.get_small(composition) + self.dissH2_Martin96_n_Tgas(T, composition)

        H_recombination_on_dust = self.get_small(composition) + self.H_recombination_on_dust(T, composition, user_dust2gas_ratio, chiFUV)
        He_recombination_on_dust = self.get_small(composition) + self.He_recombination_on_dust(T, composition, user_dust2gas_ratio, chiFUV)
        C_recombination_on_dust = self.get_small(composition) + self.C_recombination_on_dust(T, composition, user_dust2gas_ratio, chiFUV)
        O_recombination_on_dust = self.get_small(composition) + self.O_recombination_on_dust(T, composition, user_dust2gas_ratio, chiFUV)

        #stuff for CO freezing out on dust grains
        index_of_CO = None
        for index, key in enumerate(composition):
            if key == ChemSpecie('co').sym_name:
                index_of_CO = index
                #print('CO found at idx: ', index_of_CO)
                break
        if index_of_CO is None:
            raise ValueError('CO not found in the list of species!')
        #now, get the sticking rate
        krate_stickSi_n_idx_CO_tabTdust = self.get_small(composition) +  self.krate_stickSi(composition, index_of_CO, 1.0, user_dust2gas_ratio, T)
        krate_evaporation_n_idx_CO_tabTdust = self.get_small(composition) + self.krate_evaporation(composition[ChemSpecie('co').sym_name], 1.0)

        #stuff for H2O freezing out on dust grains
        index_of_H2O = None
        for index, key in enumerate(composition):
            if key == ChemSpecie('h2o').sym_name:
                index_of_H2O = index
                #print('H2O found at idx: ', index_of_H2O)
                break
        if index_of_H2O is None:
            raise ValueError('H2O not found in the list of species!')
        #now, get the sticking rate            
        krate_stickSi_n_idx_H2O_tabTdust = self.get_small(composition) + self.krate_stickSi(composition, index_of_H2O, 1.0, user_dust2gas_ratio, T)
        krate_evaporation_n_idx_H2O_tabTdust = self.get_small(composition) + self.krate_evaporation(composition[ChemSpecie('h2o').sym_name], 1.0)

        krate_nonthermal_evaporation = self.get_small(composition) + self.krate_nonthermal_evaporation(composition[ChemSpecie('co').sym_name], 1.78e0, user_Av, user_crate, 1e-3)

        rate = self.rate_function(Tgas, Te, invTe, invT, lnTe, T32, invT32, invsqrT, composition, user_crate, user_Av, 
                                  user_ionH, user_ionH2, user_dissH2, user_ionC, user_ionO, user_dissCO, user_dust2gas_ratio, Tdust, chiFUV)

        # Create a substitution dictionary for all intermediate variables
        subs_dict = {
            'Tgas': Tgas,
            'Te': Te,
            'invTe': invTe,
            'invT': invT,
            'invTgas': invTgas,
            'lnTe': lnTe,
            'T32': T32,
            'invT32': invT32,
            'invsqrT': invsqrT,
            'HnOj': HnOj,
            'ntot': ntot,
            'Hnuclei': Hnuclei,
            'nidx_H': nidx_H,
            'dissH2_Martin96_n_Tgas': dissH2_Martin96_n_Tgas,
            'H_recombination_on_dust_n_Tgas': H_recombination_on_dust,
            'He_recombination_on_dust_n_Tgas': He_recombination_on_dust,
            'C_recombination_on_dust_n_Tgas': C_recombination_on_dust,
            'O_recombination_on_dust_n_Tgas': O_recombination_on_dust,
            'krate_stickSi_n_idx_CO_tabTdust': krate_stickSi_n_idx_CO_tabTdust,
            'krate_evaporation_n_idx_CO_tabTdust': krate_evaporation_n_idx_CO_tabTdust,
            'krate_stickSi_n_idx_H2O_tabTdust': krate_stickSi_n_idx_H2O_tabTdust,
            'krate_evaporation_n_idx_H2O_tabTdust': krate_evaporation_n_idx_H2O_tabTdust
        }
        #    'krate_nonthermal_evaporation_idx_CO_1p78_user_Av_user_crate_1em3': krate_nonthermal_evaporation
        #}
        # Substitute intermediate variables into the rate expression
        rate = rate.subs(subs_dict)

        return rate

    def __repr__(self):
        repstring = str(self.reactants) + ' --> ' + str(self.products)
        return repstring

class UnsupportedChemSpecie(BaseException):
    def __init__(self):
        print('The chemical specie you entered is unsupported ', self)
        return


class UnsupportedChemRate(BaseException):
    def __init__(self):
        print('The chemical rate for the specie(s) you entered is unsupported')
        return

class UnsupportedSympyChemRate(BaseException):
    def __init__(self):
        print('The sympy chemical rate for the specie(s) you entered is unsupported')
        return

class UnphysicalTemperature(BaseException):
    def __init__(self):
        print('The temperature you entered is <= 0 K')
        return