dengo / dengo /

The default branch has multiple heads

Author: Matthew Turk <>
Affiliation: Columbia University
  Copyright (C) 2012 Matthew Turk.  All Rights Reserved.

  This file is part of the dengo 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
from .reaction_classes import reaction_registry, cooling_registry, \
    count_m, index_i, species_registry
import types
import sympy
from sympy.printing import ccode

class ChemicalNetwork(object):

    energy_term = None

    def __init__(self):
        self.reactions = {}
        self.cooling_actions = {}
        self.required_species = set([])

    def add_reaction(self, reaction):
        if isinstance(reaction, types.StringTypes):
            reaction = reaction_registry[reaction]
        if in self.reactions:
            raise RuntimeError
        self.reactions[] = reaction
        for n, s in reaction.left_side:
        for n, s in reaction.right_side:
        print "Adding reaction: %s" % reaction

    def add_cooling(self, cooling_term):
        if self.energy_term is None:
            self.energy_term = species_registry["ge"]
        if in self.cooling_actions:
            raise RuntimeError
        self.cooling_actions[] = cooling_term
    def init_temperature(self, T_bounds = (1, 1e8), n_bins=1024):
        self.n_bins = 1024
        self.T = na.logspace(na.log(T_bounds[0]),
                             n_bins, base = na.e)
        self.logT = na.log(self.T)
        self.tev = self.T / tevk
        self.logtev = na.log(self.tev)
        self.T_bounds = T_bounds

    def species_total(self, species):
        eq = sympy.sympify("0")
        for rn, rxn in sorted(self.reactions.items()):
            eq += rxn.species_equation(species)
        return eq

    def species_reactions(self, species):
        tr = []
        for rn, rxn in sorted(self.reactions.items()):
            if species in rxn: tr.append(rxn)
        return tr

    def species_list(self):
        species_list = []
        for s in sorted(self.required_species):
        return species_list

    def __iter__(self):
        for rname, rxn in sorted(self.reactions.items()):
            yield rxn

    def print_ccode(self, species, assign_to = None):
        #assign_to = sympy.IndexedBase("d_%s" %, (count_m,))
        if assign_to is None: assign_to = sympy.Symbol("d_%s[i]" %
        if species == self.energy_term:
            return self.print_cooling(assign_to)
        return ccode(species.weight * self.species_total(species),
                     assign_to = assign_to)

    def print_cooling(self, assign_to):
        eq = 0
        for term in self.cooling_actions:
            eq += self.cooling_actions[term].equation
        return ccode(eq, assign_to = assign_to)

    def print_jacobian_component(self, s1, s2, assign_to = None):
        if s1 == self.energy_term:
            st = sum(self.cooling_actions[ca].equation
                     for ca in sorted(self.cooling_actions))
            st = self.species_total(s1)
        if assign_to is None:
            assign_to = sympy.Symbol("d_%s_%s" % (,
        if isinstance(st, (list, tuple)):
            codes = []
            for temp_name, temp_eq in st[0]:
                teq = sympy.sympify(temp_eq)
                codes.append(ccode(teq, assign_to = temp_name))
            codes.append(ccode(st[1], assign_to = assign_to))
            return "\n".join(codes)
        return ccode(sympy.diff(st, s2.symbol), assign_to = assign_to)

    def print_number_density(self):
        eq = 0
        for s in sorted(self.required_species):
            if != 'ge': 
                eq += (1.0/s.weight * s.symbol)
        return ccode(eq)

    def print_mass_density(self):
        eq = 0
        for s in sorted(self.required_species):
            if != 'ge': 
                eq += s.symbol
        return ccode(eq)

    def species_gamma(self, species):
        if == 'H2I' or == 'H2II':
            gamma = sympy.Symbol('gammaH2')
            gamma = sympy.Symbol('gamma')
        return gamma

    def gamma_factor(self):
        eq = sympy.sympify("0")
        for s in sorted(self.required_species):
            if != 'ge':
                eq += (1.0/s.weight * sympy.sympify( / \
                    (self.species_gamma(s) - 1.0)
        return eq

    def temperature_calculation(self, derivative=False):
        # If derivative=True, returns the derivative of
        # temperature with respect to ge.  Otherwise,
        # returns just the temperature function
        ge = sympy.Symbol('ge')
        function_eq = (sympy.Symbol('density') * ge * sympy.Symbol('mh')) / \
            (sympy.Symbol('kb') * (self.gamma_factor()))
        if derivative == True:
            deriv_eq = sympy.diff(function_eq, ge)
            return ccode(deriv_eq)
            return ccode(function_eq)
        return ccode(eq)