<commit_before>import logging

import numpy as np
import pandas as pd
from scipy import interpolate
from astropy import constants, units
import math
from scipy.integrate import trapz

from tardis.plasma.properties.base import ProcessingPlasmaProperty
from tardis.plasma.exceptions import PlasmaIonizationError

logger = logging.getLogger(__name__)

__all__ = ['PhiSahaNebular', 'PhiSahaLTE', 'PhiSahaNLTE', 'PhiSahaNoNLTE',
           'RadiationFieldCorrection', 'IonNumberDensity']

class PhiSahaLTE(ProcessingPlasmaProperty):
    """
    Outputs:
    general_phi : Pandas DataFrame
        Used as basis for PhiSahaLTE and PhiSahaNebular. Identical output to
        PhiSahaLTE. Separate property required as PhiSahaNebular is based
        on PhiSahaLTE, but the code cannot deal with the inclusion of two
        properties that generate a property called 'phi'.
    """
    outputs = ('general_phi',)
    latex_name = ('\\Phi',)
    latex_formula = ('\\dfrac{2Z_{i,j+1}}{Z_{i,j}}\\Big(\
                     \\dfrac{2\\pi m_{e}/\\beta_{\\textrm{rad}}}{h^2}\
                     \\Big)^{3/2}e^{\\dfrac{-\\chi_{i,j}}{kT_{\
                     \\textrm{rad}}}}',)

    @staticmethod
    def calculate(g_electron, beta_rad, partition_function, ionization_data):
        def calculate_phis(group):
            return group[1:] / group[:-1].values
        phis = partition_function.groupby(level='atomic_number').apply(
            calculate_phis)
        phis = pd.DataFrame(phis.values, index=phis.index.droplevel(0))
        phi_coefficient = (2 * g_electron * np.exp(np.outer(
            ionization_data.ionization_energy.ix[phis.index].values,
            -beta_rad)))
        return phis * phi_coefficient

class PhiSahaNebular(ProcessingPlasmaProperty):
    """
    Outputs:
    phi_saha_nebular: Pandas DataFrame
        The ionization equilibrium as calculated using a modified version of
        the Saha equation that accounts for dilution of the radiation field.
    """
    outputs = ('general_phi',)
    latex_name = ('\\Phi',)
    latex_formula = ('W(\\delta\\zeta_{i,j}+W(1-\\zeta_{i,j}))\\left(\
                     \\dfrac{T_{\\textrm{electron}}}{T_{\\textrm{rad}}}\
                     \\right)^{1/2}',)

    @staticmethod
    def calculate(t_rad, w, zeta_data, t_electrons, delta,
            g_electron, beta_rad, partition_function, ionization_data):
        phi_lte = PhiSahaLTE.calculate(g_electron, beta_rad,
            partition_function, ionization_data)
        zeta = PhiSahaNebular.get_zeta_values(zeta_data, phi_lte, t_rad)
        phis = phi_lte * w * ((zeta * delta) + w * (1 - zeta)) * \
               (t_electrons/t_rad) ** .5
        return phis

    @staticmethod
    def get_zeta_values(zeta_data, general_phi, t_rad):
        try:
            zeta = interpolate.interp1d(zeta_data.columns.values, zeta_data.ix[
                general_phi.index].values)(t_rad)
            zeta = zeta.astype(float)
        except ValueError:
            raise ValueError('t_rads outside of zeta factor interpolation'
                             ' zeta_min={0:.2f} zeta_max={1:.2f} '
                             '- requested {2}'.format(
                zeta_data.columns.values.min(), zeta_data.columns.values.max(),
                t_rad))
        return zeta

class PhiSahaNLTE(ProcessingPlasmaProperty):
    outputs = ('phi',)
    def calculate(self, general_phi, g_electron, beta_rad, partition_function,
        ionization_data, nlte_ionization_species, beta_electron, t_rad,
        previous_electron_densities, w):
        for species in nlte_ionization_species:
            number_of_ions = species+1
            partition_function_species, ionization_data_species = \
                self.filter_phi_input_data(species, partition_function,
                    ionization_data)
            phi_lte = PhiSahaLTE.calculate(g_electron, beta_rad,
                partition_function_species, ionization_data_species)
            r_lu_matrix = np.zeros((species+1, species+1, len(beta_rad)))
            r_lu_matrix_reshaped = r_lu_matrix.reshape(((species+1)**2,
                len(beta_rad)))
            r_ul_matrix = np.zeros((species+1, species+1, len(beta_rad)))
            r_ul_matrix_reshaped = r_ul_matrix.reshape(((species+1)**2,
                len(beta_rad)))
            for ion in range(0, species):
                ionization_energy = ionization_data.ix[species].ix[
                    ion+1].ionization_energy * units.erg
                ionization_nu = ionization_energy / constants.h.cgs
                alpha_st, alpha_sp = self.calculate_rate_upper_lower(phi_lte,
                    ionization_nu, species, ion, beta_electron, t_rad, w)
                gamma = alpha_sp / phi_lte.ix[species].ix[ion+1]
                lower_index = (number_of_ions + 1) * ion
                upper_index = lower_index + number_of_ions
                r_lu_matrix_reshaped[lower_index] = -gamma
                r_lu_matrix_reshaped[upper_index] = gamma
                r_ul_matrix_reshaped[lower_index+1] = (alpha_st + alpha_sp) * \
                    previous_electron_densities
                r_ul_matrix_reshaped[upper_index+1] = -(alpha_st + alpha_sp) *\
                    previous_electron_densities
            rates_matrix = r_lu_matrix + r_ul_matrix
            rates_matrix[0, :, :] = 1.0
            x = np.zeros(rates_matrix.shape[0])
            x[0] = 1.0
            for i in xrange(len(beta_rad)):
                phi = \
                    np.linalg.solve(rates_matrix[:, :, i], x)
                general_phi[i].ix[species] = \
                    (phi[1:] / phi[:-1])
        return general_phi

    @staticmethod
    def calculate_rate_upper_lower(phi_lte, ionization_nu, species, ion,
        beta_electron, t_rad, w):
#Assuming for now that a_{ik} = (v_{i}/v)^3
        sp_coefficient = 8 * np.pi * phi_lte.ix[species].ix[ion+1] * (
            constants.c.cgs.value)**(-2.0) * (ionization_nu)**(3.0)
        x_interval = 9*ionization_nu.value
        x_values = np.arange(ionization_nu.value, 10*ionization_nu.value, (
            x_interval/1000.0))
        exponential_coefficient = -constants.h.cgs.value * beta_electron
        alpha_sp_dataframe = pd.DataFrame(1.0, index=range(len(
            exponential_coefficient)),
            columns=range(len(x_values)))
        y_values_sp = (np.exp(alpha_sp_dataframe.mul(exponential_coefficient,
            axis=0).mul(x_values, axis=1))).mul(1/x_values, axis=1)
        alpha_sp = trapz(y_values_sp, x_values) * sp_coefficient
        alpha_st_dataframe = pd.DataFrame(1.0, index=range(len(
            exponential_coefficient)),
            columns=range(len(x_values)))
        st_coefficient = 4 * np.pi * phi_lte.ix[species].ix[ion+1] * (
            1 / constants.h.cgs.value) * (ionization_nu)**(3.0)
        J_nu_dataframe = pd.DataFrame(1.0, index=range(len(
            exponential_coefficient)), columns=range(len(x_values)))
        J_nu = (J_nu_dataframe * 2.0 * constants.h.cgs.value * \
            constants.c.cgs.value**(-2.0)).mul(x_values**(3.0), axis=1).mul((
            w / np.exp(exponential_coefficient * t_rad) - 1.0), axis=0)
        y_values_st = (np.exp(alpha_st_dataframe.mul(exponential_coefficient,
            axis=0).mul(x_values, axis=1))).mul(x_values**-4.0, axis=1).mul(
            J_nu, axis=1)
        alpha_st = trapz(y_values_st, x_values) * st_coefficient
        return alpha_st, alpha_sp

    @staticmethod
    def filter_phi_input_data(atomic_number, partition_function,
        ionization_data):
        index_tuples = []
        for ion in range(atomic_number+1):
            index_tuples.append((atomic_number, ion))
        partition_function_index = pd.MultiIndex.from_tuples(index_tuples,
            names=['atomic_number', 'ion_number'])
        ionization_data_index = pd.MultiIndex.from_tuples(index_tuples[1:],
            names=['atomic_number', 'ion_number'])
        partition_function = pd.DataFrame(
            partition_function.ix[atomic_number].values,
            index=partition_function_index, columns=partition_function.columns)
        ionization_data = pd.DataFrame(
            ionization_data.ix[atomic_number].values,
            index=ionization_data_index, columns=['ionization_energy'])
        return partition_function, ionization_data

class PhiSahaNoNLTE(ProcessingPlasmaProperty):
    outputs = ('phi',)
    def calculate(self, general_phi):
        return general_phi

class RadiationFieldCorrection(ProcessingPlasmaProperty):
    """<commit_msg>Change to J_nu calculation.<commit_after>import logging

import numpy as np
import pandas as pd
from scipy import interpolate
from astropy import constants, units
import math
from scipy.integrate import trapz

from tardis.plasma.properties.base import ProcessingPlasmaProperty
from tardis.plasma.exceptions import PlasmaIonizationError

logger = logging.getLogger(__name__)

__all__ = ['PhiSahaNebular', 'PhiSahaLTE', 'PhiSahaNLTE', 'PhiSahaNoNLTE',
           'RadiationFieldCorrection', 'IonNumberDensity']

class PhiSahaLTE(ProcessingPlasmaProperty):
    """
    Outputs:
    general_phi : Pandas DataFrame
        Used as basis for PhiSahaLTE and PhiSahaNebular. Identical output to
        PhiSahaLTE. Separate property required as PhiSahaNebular is based
        on PhiSahaLTE, but the code cannot deal with the inclusion of two
        properties that generate a property called 'phi'.
    """
    outputs = ('general_phi',)
    latex_name = ('\\Phi',)
    latex_formula = ('\\dfrac{2Z_{i,j+1}}{Z_{i,j}}\\Big(\
                     \\dfrac{2\\pi m_{e}/\\beta_{\\textrm{rad}}}{h^2}\
                     \\Big)^{3/2}e^{\\dfrac{-\\chi_{i,j}}{kT_{\
                     \\textrm{rad}}}}',)

    @staticmethod
    def calculate(g_electron, beta_rad, partition_function, ionization_data):
        def calculate_phis(group):
            return group[1:] / group[:-1].values
        phis = partition_function.groupby(level='atomic_number').apply(
            calculate_phis)
        phis = pd.DataFrame(phis.values, index=phis.index.droplevel(0))
        phi_coefficient = (2 * g_electron * np.exp(np.outer(
            ionization_data.ionization_energy.ix[phis.index].values,
            -beta_rad)))
        return phis * phi_coefficient

class PhiSahaNebular(ProcessingPlasmaProperty):
    """
    Outputs:
    phi_saha_nebular: Pandas DataFrame
        The ionization equilibrium as calculated using a modified version of
        the Saha equation that accounts for dilution of the radiation field.
    """
    outputs = ('general_phi',)
    latex_name = ('\\Phi',)
    latex_formula = ('W(\\delta\\zeta_{i,j}+W(1-\\zeta_{i,j}))\\left(\
                     \\dfrac{T_{\\textrm{electron}}}{T_{\\textrm{rad}}}\
                     \\right)^{1/2}',)

    @staticmethod
    def calculate(t_rad, w, zeta_data, t_electrons, delta,
            g_electron, beta_rad, partition_function, ionization_data):
        phi_lte = PhiSahaLTE.calculate(g_electron, beta_rad,
            partition_function, ionization_data)
        zeta = PhiSahaNebular.get_zeta_values(zeta_data, phi_lte, t_rad)
        phis = phi_lte * w * ((zeta * delta) + w * (1 - zeta)) * \
               (t_electrons/t_rad) ** .5
        return phis

    @staticmethod
    def get_zeta_values(zeta_data, general_phi, t_rad):
        try:
            zeta = interpolate.interp1d(zeta_data.columns.values, zeta_data.ix[
                general_phi.index].values)(t_rad)
            zeta = zeta.astype(float)
        except ValueError:
            raise ValueError('t_rads outside of zeta factor interpolation'
                             ' zeta_min={0:.2f} zeta_max={1:.2f} '
                             '- requested {2}'.format(
                zeta_data.columns.values.min(), zeta_data.columns.values.max(),
                t_rad))
        return zeta

class PhiSahaNLTE(ProcessingPlasmaProperty):
    outputs = ('phi',)
    def calculate(self, general_phi, g_electron, beta_rad, partition_function,
        ionization_data, nlte_ionization_species, beta_electron, t_rad,
        previous_electron_densities, w):
        for species in nlte_ionization_species:
            number_of_ions = species+1
            partition_function_species, ionization_data_species = \
                self.filter_phi_input_data(species, partition_function,
                    ionization_data)
            phi_lte = PhiSahaLTE.calculate(g_electron, beta_rad,
                partition_function_species, ionization_data_species)
            r_lu_matrix = np.zeros((species+1, species+1, len(beta_rad)))
            r_lu_matrix_reshaped = r_lu_matrix.reshape(((species+1)**2,
                len(beta_rad)))
            r_ul_matrix = np.zeros((species+1, species+1, len(beta_rad)))
            r_ul_matrix_reshaped = r_ul_matrix.reshape(((species+1)**2,
                len(beta_rad)))
            for ion in range(0, species):
                ionization_energy = ionization_data.ix[species].ix[
                    ion+1].ionization_energy * units.erg
                ionization_nu = ionization_energy / constants.h.cgs
                alpha_st, alpha_sp = self.calculate_rate_upper_lower(phi_lte,
                    ionization_nu, species, ion, beta_electron, t_rad, w,
                    beta_rad)
                gamma = alpha_sp / phi_lte.ix[species].ix[ion+1]
                lower_index = (number_of_ions + 1) * ion
                upper_index = lower_index + number_of_ions
                r_lu_matrix_reshaped[lower_index] = -gamma
                r_lu_matrix_reshaped[upper_index] = gamma
                r_ul_matrix_reshaped[lower_index+1] = (alpha_st + alpha_sp) * \
                    previous_electron_densities
                r_ul_matrix_reshaped[upper_index+1] = -(alpha_st + alpha_sp) *\
                    previous_electron_densities
            rates_matrix = r_lu_matrix + r_ul_matrix
            rates_matrix[0, :, :] = 1.0
            x = np.zeros(rates_matrix.shape[0])
            x[0] = 1.0
            for i in xrange(len(beta_rad)):
                phi = \
                    np.linalg.solve(rates_matrix[:, :, i], x)
                general_phi[i].ix[species] = \
                    (phi[1:] / phi[:-1])
        return general_phi

    @staticmethod
    def calculate_rate_upper_lower(phi_lte, ionization_nu, species, ion,
        beta_electron, t_rad, w, beta_rad):
#Assuming for now that a_{ik} = (v_{i}/v)^3
        sp_coefficient = 8 * np.pi * phi_lte.ix[species].ix[ion+1] * (
            constants.c.cgs.value)**(-2.0) * (ionization_nu)**(3.0)
        x_interval = 9*ionization_nu.value
        x_values = np.arange(ionization_nu.value, 10*ionization_nu.value, (
            x_interval/1000.0))
        exponential_coefficient = -constants.h.cgs.value * beta_electron
        alpha_sp_dataframe = pd.DataFrame(1.0, index=range(len(
            exponential_coefficient)),
            columns=range(len(x_values)))
        y_values_sp = (np.exp(alpha_sp_dataframe.mul(exponential_coefficient,
            axis=0).mul(x_values, axis=1))).mul(1/x_values, axis=1)
        alpha_sp = trapz(y_values_sp, x_values) * sp_coefficient
        alpha_st_dataframe = pd.DataFrame(1.0, index=range(len(
            exponential_coefficient)),
            columns=range(len(x_values)))
        st_coefficient = 4 * np.pi * phi_lte.ix[species].ix[ion+1] * (
            1 / constants.h.cgs.value) * (ionization_nu)**(3.0)
        J_nu = PhiSahaNLTE.get_J_nu_values(x_values, beta_rad, w)
        y_values_st = (np.exp(alpha_st_dataframe.mul(exponential_coefficient,
            axis=0).mul(x_values, axis=1))).mul(x_values**-4.0, axis=1).mul(
            J_nu)
        alpha_st = trapz(y_values_st, x_values) * st_coefficient
        return alpha_st, alpha_sp

    @staticmethod
    def get_J_nu_values(x_values, beta_rad, w):
        exponential_coefficient_j_nu = -constants.h.cgs.value * beta_rad
        J_nu = pd.DataFrame(1.0, index=range(len(
            exponential_coefficient_j_nu)), columns=range(len(x_values)))
        for nu_index in range(len(x_values)):
            nu = x_values[nu_index]
            nu_interval = 9*nu
            nu_values = np.arange(nu, 10*nu, (nu_interval/1000.0))
            J_nu_dataframe = pd.DataFrame(1.0, index=range(len(
                exponential_coefficient_j_nu)), columns=range(len(nu_values)))
            J_nu_exponential = J_nu_dataframe.mul(nu_values, axis=1).mul(
                exponential_coefficient_j_nu, axis=0)
            J_nu_exp_term = 1 / (J_nu_exponential - 1.0)
            J_nu_values = (J_nu_exp_term * 2.0 * constants.h.cgs.value *
                constants.c.cgs.value**(-2.0)).mul(nu_values**(3.0),
                axis=1).mul(w, axis=0)
            J_nu[nu_index] = trapz(J_nu_values, nu_values)
        return J_nu

    @staticmethod
    def filter_phi_input_data(atomic_number, partition_function,
        ionization_data):
        index_tuples = []
        for ion in range(atomic_number+1):
            index_tuples.append((atomic_number, ion))
        partition_function_index = pd.MultiIndex.from_tuples(index_tuples,
            names=['atomic_number', 'ion_number'])
        ionization_data_index = pd.MultiIndex.from_tuples(index_tuples[1:],
            names=['atomic_number', 'ion_number'])
        partition_function = pd.DataFrame(
            partition_function.ix[atomic_number].values,
            index=partition_function_index, columns=partition_function.columns)
        ionization_data = pd.DataFrame(
            ionization_data.ix[atomic_number].values,
            index=ionization_data_index, columns=['ionization_energy'])
        return partition_function, ionization_data

class PhiSahaNoNLTE(ProcessingPlasmaProperty):
    outputs = ('phi',)
    def calculate(self, general_phi):
        return general_phi

class RadiationFieldCorrection(ProcessingPlasmaProperty):
    """
{'fraction_non_alphanumeric': 0.11143400022353861, 'fraction_numerical': 0.011065161506650274, 'mean_word_length': 2.5074480595844766, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1073118', 'n_tokens_mistral': 5915, 'n_tokens_neox': 5608, 'n_words': 963}