class SimpleCOAtomAbund(ChemicalAbund):
    """Class to hold the raw atomic abundaces of C/O/Si for the CO chemistry"""
    def __init__(self, *sizes):
        atom_ids = ['C', 'O', 'Si']
        masses   = [ 12., 16., 28. ]
        
        super(SimpleCOAtomAbund, self).__init__(atom_ids, masses, *sizes)


    def set_solar_abundances(self, muH=1.28):
        """Solar mass fractions of C, O and Si.

        args:
            muH : mean atomic mass, default = 1.28
        """
        self._data[:] = np.outer(np.array([12*2.7e-4, 16*4.9e-4, 28*3.2e-5]),
                                 np.ones(self.size)) / muH class SimpleCOChemOberg(COChemOberg, StaticChem):
    def __init__(self, **kwargs):
        COChemOberg.__init__(self)
        StaticChem.__init__(self, **kwargs) class EquilibriumCOChemOberg(COChemOberg, EquilibriumChem):
    def __init__(self, fix_ratios=False, fix_grains=True, **kwargs):
        COChemOberg.__init__(self, fix_grains)
        EquilibriumChem.__init__(self, 
                                 fix_ratios=fix_ratios, 
                                 **kwargs) class SimpleCOChemMadhu(COChemMadhu, StaticChem):
    def __init__(self, **kwargs):
        COChemMadhu.__init__(self)
        StaticChem.__init__(self, **kwargs) class EquilibriumCOChemMadhu(COChemMadhu, EquilibriumChem):
    def __init__(self, fix_ratios=False, fix_grains=True, **kwargs):
        COChemMadhu.__init__(self, fix_grains)
        EquilibriumChem.__init__(self,
                                 fix_ratios=fix_ratios, 
                                 **kwargs) class SimpleCNOAtomAbund(ChemicalAbund):
    """Class to hold the raw atomic abundaces of C/O/Si for the CNO chemistry"""

    def __init__(self, *sizes):
        atom_ids = ['C', 'N', 'O', 'Si']
        masses = [12., 14., 16., 28.]

        super(SimpleCNOAtomAbund, self).__init__(atom_ids, masses, *sizes)

    def set_solar_abundances(self, muH=1.28):
        """Solar mass fractions of C, O and Si.

        args:
            muH : mean atomic mass, default = 1.28
        """
        m_abund = np.array([12 * 2.7e-4, 14 * 6.8e-5, 16 * 4.9e-4, 28 * 3.2e-5])
        self._data[:] = np.outer(m_abund, np.ones(self.size)) / muH class SimpleCNOChemOberg(CNOChemOberg, StaticChem):
    def __init__(self, fNH3=None, **kwargs):
        CNOChemOberg.__init__(self, fNH3)
        StaticChem.__init__(self, **kwargs) class EquilibriumCNOChemOberg(CNOChemOberg, EquilibriumChem):
    def __init__(self, fNH3=None, fix_ratios=False, fix_grains=True,
                 fix_N=False, **kwargs):
        CNOChemOberg.__init__(self, fNH3, fix_grains, fix_N)
        EquilibriumChem.__init__(self,
                                 fix_ratios=fix_ratios,
                                 **kwargs) class SimpleCNOChemMadhu(CNOChemMadhu, StaticChem):
    def __init__(self, fNH3=None, **kwargs):
        CNOChemMadhu.__init__(self, fNH3)
        StaticChem.__init__(self, **kwargs) class EquilibriumCNOChemMadhu(CNOChemMadhu, EquilibriumChem):
    def __init__(self, fNH3=None, fix_ratios=False, fix_grains=True,
                 fix_N=False, **kwargs):
        CNOChemMadhu.__init__(self, fNH3, fix_grains, fix_N)
        EquilibriumChem.__init__(self,
                                 fix_ratios=fix_ratios,
                                 **kwargs)
import numpy as np
from DiscEvolution.chemistry import (
    SimpleCOAtomAbund,
    SimpleCOChemMadhu, SimpleCOChemOberg,
    EquilibriumCOChemMadhu, EquilibriumCOChemOberg
)
from DiscEvolution.chemistry import (
    SimpleCNOAtomAbund,
    SimpleCNOChemMadhu, SimpleCNOChemOberg,
    EquilibriumCNOChemMadhu, EquilibriumCNOChemOberg
)
# test_simple_chem.py
#
# Author: R. Booth
# Date: 22 - May - 2018
#
# Check behaviour of simple chemistry functions
###############################################################################

# CO chemistry

# CNO chemistry

def _test_chemical_model(ChemModel, Abundances):
    T = np.logspace(0.5, 3, 6)
    
    Xi = Abundances(len(T))
    Xi.set_solar_abundances()

    Chem = ChemModel()
    
    mol = Chem.equilibrium_chem(T, 1e-10, 0.01, Xi)
        
    T *= 3
    Chem.update(0, T, 1e-10, 0.01, mol)

    atom  = mol.gas.atomic_abundance()
    atom += mol.ice.atomic_abundance()
    for X in atom:
        assert(np.allclose(Xi[X], atom[X], rtol=1e-12))


def test_CO_chemistry():
    for Chem in [ SimpleCOChemMadhu, SimpleCOChemOberg,

                  EquilibriumCOChemMadhu, EquilibriumCOChemOberg,]: