dengo / primordial_chemistry /

The default branch has multiple heads

Author: Matthew Turk <>
Affiliation: UCSD
  Copyright (C) 2010 Matthew Turk.  All Rights Reserved.

  This file is part of the primordial_chemistry package.

  This file is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 3 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program.  If not, see <>.

import numpy as na
from chemistry_constants import tevk, tiny, mh

reaction_rates_table = dict()

class ReactionRate(object):
    def __init__(self, name, values):
        self.values = values = name

    def __call__(self, quantities):
        T = quantities['T']
        return na.interp(T, self.T, self.values)

    def init_temperature(cls, T_bounds, n_bins=1024):
        cls.n_bins = 1024
        cls.T = na.logspace(
            na.log10(T_bounds[0]), na.log10(T_bounds[1]), n_bins)
        cls.logT = na.log(cls.T)
        cls.tev = cls.T / tevk
        cls.logtev = na.log(cls.tev)
        cls.T_bounds = T_bounds

class Reaction(object):
    def __init__(self, rate, left_side, right_side): = rate.replace("k","r")
        self.rate = reaction_rates_table[rate]
        self.left_side = left_side
        self.right_side = right_side
        self.considered = set( ( for n, s in left_side + right_side) )

    def down_species(self):
        return [ for n, s in self.left_side]

    def up_species(self):
        return [ for n, s in self.right_side]

    def __call__(self, quantities, up_derivatives, down_derivatives):
        # We just calculate our net derivatives and stick them in the right
        # place
        r = self.rate(quantities)
        for n, s in self.left_side:
            r *= s.number_density(quantities)**n
        for n, s in self.left_side:
            down_derivatives[] += r * n * s.weight
        for n, s in self.right_side:
            up_derivatives[] += r * n * s.weight
        return r

    def __repr__(self):
        a = "%s : " % \
          + " + ".join( ["%s*%s" % (i, for i, s in self.left_side] ) \
          + " => " \
          + " + ".join( ["%s*%s" % (i, for i, s in self.right_side] )
        return a

    def print_c_calculation(self):
        # This function assumes we're inside a loop
        st = ""
        st += "for (i = 0; i < data->nvals ; i++)\n"
        st += "{\n"
        st += "  tv = 0.0\n"
        for n, s in self.left_side:
            #r *= s.number_density(quantities)**n
            st += "  tv -= pow(y[%s], %s);\n" % (, n)
        for n, s in self.right_side:
            #up_derivatives[] += r * n * s.weight
            st += "  tv += y[%s]*%s;\n" % (, n)
        st += "  ydot[i] * data->stride + ri;\n"
        st += "}\n"
        return st

    def print_c_reaction_value(self, species_varnames):
        st = []
        for n, s in self.left_side:
            st += [" * ".join([
                for i in xrange(n)])]
        return " * ".join(st)

class Species(object):
    def __init__(self, name, weight, free_electrons = 0.0, equilibrium = False,
                 computed = False): = name
        self.weight = weight
        self.free_electrons = free_electrons
        self.equilibrium = equilibrium
        self.computed = computed
        if equilibrium and computed: raise RuntimeError

    def number_density(self, quantities):
        return quantities[]/self.weight

    def print_c_convert_number_density(self, input, output):
        return "%s = %s / %s;" % (output, input, self.weight)

    def print_c_number_density(self, input):
        return "(%s / %s)" % (input, self.weight)

    def __repr__(self):
        return "Species: %s" % (

class Constraint(object):

class ChargeConservation(Constraint):
    def __call__(self, quantities, up_derivatives, down_derivatives, dt):
        quantities["de"] = (quantities["HII"]
            + quantities["HeII"] / 4.0
            + quantities["HeIII"] / 2.0
            + quantities["H2II"] / 2.0
            - quantities["HM"])
        quantities["de"] = 0.0
        for q in quantities.species_list:
            quantities["de"] += q.free_electrons * q.number_density(quantities)

class Floor(Constraint):
    def __call__(self, quantities, up_derivatives, down_derivatives, dt):
        for s in quantities.species_list:
            quantities[] = max(quantities[], 1e-30)

class ChemicalHeating(Constraint):
    def __call__(self, quantities, up_derivatives, down_derivatives, dt):
        # Get the total mass
        rho = sum(quantities[i] for i in
        dH2 = (up_derivatives["H2I"] - down_derivatives["H2I"])*dt
        quantities["T"] += dH2 * 51998.0/rho

constraints = [ChemicalHeating(), ChargeConservation(), Floor()]

class QuantitiesTable(object):
    def __init__(self, species_list, initial_values = None):
        self._names = {}
        for i,s in enumerate(species_list):
            self._names[] = i
        self.species_list = species_list
        self.values = na.zeros(len(species_list), dtype='float64')
        if initial_values is not None:
            for s, v in initial_values.items():
                self.values[self._names[s]] = v

    def __getitem__(self, name):
        return self.values[self._names[name]]

    def __setitem__(self, name, value):
        self.values[self._names[name]] = value

    def __iter__(self):
        for i, s in enumerate(self.species_list):
            yield self.values[self._names[]]

    def get_by_name(self, name):
        return self.species_list[self._names[name]]
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.