<commit_before>import numpy as np
from scipy.integrate import odeint

from abundance import FractionalAbundance

class RateEquations(object):
    def __init__(self, atomic_data):
        self.atomic_data = atomic_data
        self.nuclear_charge = atomic_data.nuclear_charge

    def _set_temperature_and_density_grid(self, temperature, density):
        self.temperature = temperature
        self.density = density

    def _set_initial_conditions(self):
        self.y_shape = (self.nuclear_charge + 1, len(self.temperature))
        self._init_y()
        self._init_coeffs()

    def _init_y(self):
        y = np.zeros(self.y_shape)
        y[0] = np.ones_like(self.temperature)
        self.y = y.ravel()
        self.dydt = np.zeros(self.y_shape)

    def _init_coeffs(self):
        S = np.zeros(self.y_shape)
        alpha = np.zeros(self.y_shape)

        recombination_coeff = self.atomic_data.coeffs['recombination']
        ionisation_coeff = self.atomic_data.coeffs['ionisation']
        for k in xrange(self.nuclear_charge):
            S[k] = ionisation_coeff(k, self.temperature, self.density)
            alpha[k] = recombination_coeff(k, self.temperature, self.density)

        self.S = S
        self.alpha = alpha

    def derivs(self, y_, t0):
        """right hand side of the rate equations"""

        dydt = self.dydt
        S = self.S
        alpha = self.alpha
        ne = self.density

        y = y_.reshape(self.y_shape)

        for k in xrange(self.nuclear_charge + 1): #FIXME: number of charge states
            if k == 0:
                gain = y[k+1]*alpha[k]
                loss = y[k] * S[k]
            elif k == self.nuclear_charge:
                gain = y[k-1]*S[k-1]
                loss = y[k] * alpha[k-1]
            else:
                # eq 2 with alpha[k+1] => alpha[k] substitiution
                gain = y[k-1]*S[k-1] + y[k+1]*alpha[k]
                loss = y[k] * (S[k] + alpha[k-1])

            dydt[k] = ne * (gain - loss)

        return dydt.ravel()

    def solve(self, time, temperature, density):
        """
        Integrate the rate equations.

        Parameters
        ----------
        time : array
            A sequence of time points for which to solve.
        temperature : array
            Electron temperature grid to solve on [eV].
        density : array
            Electron density grid to solve on [m-3].
        """

        self._set_temperature_and_density_grid(temperature, density)
        self._set_initial_conditions()
        solution  = odeint(self.derivs, self.y, time)

        abundances = []
        for s in solution.reshape(time.shape + self.y_shape):
            abundances.append(FractionalAbundance(self.atomic_data, s, self.temperature,
                self.density))

        return RateEquationsSolution(time, abundances)


class RateEquationsSolution(object):
    def __init__(self, times, abundances):
        self.times = times<commit_msg>Optimised version of derivs() using vector slicing.<commit_after>import numpy as np
from scipy.integrate import odeint

from abundance import FractionalAbundance

class RateEquations(object):
    def __init__(self, atomic_data):
        self.atomic_data = atomic_data
        self.nuclear_charge = atomic_data.nuclear_charge

    def _set_temperature_and_density_grid(self, temperature, density):
        self.temperature = temperature
        self.density = density

    def _set_initial_conditions(self):
        self.y_shape = (self.nuclear_charge + 1, len(self.temperature))
        self._init_y()
        self._init_coeffs()

    def _init_y(self):
        y = np.zeros(self.y_shape)
        y[0] = np.ones_like(self.temperature)
        self.y = y.ravel()
        self.dydt = np.zeros(self.y_shape)

    def _init_coeffs(self):
        S = np.zeros(self.y_shape)
        alpha = np.zeros(self.y_shape)

        recombination_coeff = self.atomic_data.coeffs['recombination']
        ionisation_coeff = self.atomic_data.coeffs['ionisation']
        for k in xrange(self.nuclear_charge):
            S[k] = ionisation_coeff(k, self.temperature, self.density)
            alpha[k] = recombination_coeff(k, self.temperature, self.density)

        self.S = S
        self.alpha = alpha

    def derivs(self, y_, t0):
        """right hand side of the rate equations"""

        dydt = self.dydt
        S = self.S
        alpha = self.alpha
        ne = self.density

        y = y_.reshape(self.y_shape)

        for k in xrange(self.nuclear_charge + 1): #FIXME: number of charge states
            if k == 0:
                gain = y[k+1]*alpha[k]
                loss = y[k] * S[k]
            elif k == self.nuclear_charge:
                gain = y[k-1]*S[k-1]
                loss = y[k] * alpha[k-1]
            else:
                # eq 2 with alpha[k+1] => alpha[k] substitiution
                gain = y[k-1]*S[k-1] + y[k+1]*alpha[k]
                loss = y[k] * (S[k] + alpha[k-1])

            dydt[k] = ne * (gain - loss)

        return dydt.ravel()

    def derivs_optimized(self, y_, t0):
        """
        Optimised version of derivs using array slicing.  It should give the
        same as derivs().
        """

        dydt = self.dydt
        S = self.S
        alpha_to = self.alpha
        ne = self.density

        y = y_.reshape(self.y_shape)
        current = slice(1, -1)
        upper = slice(2, None)
        lower = slice(None, -2)
        dydt[current]  = y[lower] * S[lower]
        dydt[current] += y[upper] * alpha_to[current]
        dydt[current] -= y[current] * S[current]
        dydt[current] -= y[current] * alpha_to[lower]

        current, upper = 0, 1 # neutral and single ionised state
        dydt[current] = y[upper] * alpha_to[current] - y[current] * S[current]

        current, lower = -1, -2 # fully stripped and 1 electron state
        dydt[current] = y[lower] * S[lower] - y[current] * alpha_to[lower]
        dydt *= ne

        return dydt.ravel()

    def solve(self, time, temperature, density):
        """
        Integrate the rate equations.

        Parameters
        ----------
        time : array
            A sequence of time points for which to solve.
        temperature : array
            Electron temperature grid to solve on [eV].
        density : array
            Electron density grid to solve on [m-3].
        """

        self._set_temperature_and_density_grid(temperature, density)
        self._set_initial_conditions()
        solution  = odeint(self.derivs_optimized, self.y, time)

        abundances = []
        for s in solution.reshape(time.shape + self.y_shape):
            abundances.append(FractionalAbundance(self.atomic_data, s, self.temperature,
                self.density))

        return RateEquationsSolution(time, abundances)


class RateEquationsSolution(object):
    def __init__(self, times, abundances):
        self.times = times
{'fraction_non_alphanumeric': 0.1097737819025522, 'fraction_numerical': 0.006090487238979118, 'mean_word_length': 2.0517699115044246, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, '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': '4901911', 'n_tokens_mistral': 2262, 'n_tokens_neox': 2114, 'n_words': 501}