<commit_before>import numpy as np


class CoronalEquilibrium(object):
    def __init__(self, atomic_data):
        self.S = atomic_data.ionisation_coeff
        self.alpha = atomic_data.recombination_coeff
        self.element = atomic_data.element
        self.nuclear_charge = atomic_data.nuclear_charge

    def ionisation_stage_distribution(self, temperature, density):
        y = np.zeros((self.nuclear_charge + 1, len(temperature)))
        y[0] = np.ones_like(temperature)
        for k in xrange(self.nuclear_charge):
            S = self.S(k, temperature, density)
            alpha = self.alpha(k, temperature, density)
            y[k+1] = y[k] * S / alpha

        y /= y.sum(0) # fractional abundance
        return FractionalAbundance(y, temperature, density)


class FractionalAbundance(object):
    def __init__(self, y, temperature, density):
        self.y = y
        self.temperature = temperature
        self.density = density

    def plot_vs_temperature(self):
        from matplotlib.pyplot import gca
        ax = gca()

        ax.loglog(self.temperature, self.y.T)
        ax.set_xlabel('$T_\mathrm{e}\ [\mathrm{eV}]$')
        ax.set_ylim(0.05, 1.3)
        self._annotate_ionisation_stages(ax)

    def _annotate_ionisation_stages(self, ax):
        max_pos = self.y.argmax(axis=-1)
        for i in xrange(self.y.shape[0]):
            index = max_pos[i]
            xy = self.temperature[index], self.y[i, index]
            s = '$%d^+$' % (i,)
            ax.annotate(s, xy, ha='center')


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    from atomic_data import AtomicData

    ad = AtomicData.from_element('Ar')
    coronal = CoronalEquilibrium(ad)

    temperature = np.logspace(0, 5, 300)
    y = coronal.ionisation_stage_distribution(temperature, density=1e19)

    plt.figure(1); plt.clf()
    y.plot_vs_temperature()

    plt.draw()
    plt.show()
<commit_msg>CoronalEquilibrium: pass *kwargs to plot function.<commit_after>import numpy as np


class CoronalEquilibrium(object):
    def __init__(self, atomic_data):
        self.S = atomic_data.ionisation_coeff
        self.alpha = atomic_data.recombination_coeff
        self.element = atomic_data.element
        self.nuclear_charge = atomic_data.nuclear_charge

    def ionisation_stage_distribution(self, temperature, density):
        y = np.zeros((self.nuclear_charge + 1, len(temperature)))
        y[0] = np.ones_like(temperature)
        for k in xrange(self.nuclear_charge):
            S = self.S(k, temperature, density)
            alpha = self.alpha(k, temperature, density)
            y[k+1] = y[k] * S / alpha

        y /= y.sum(0) # fractional abundance
        return FractionalAbundance(y, temperature, density)


class FractionalAbundance(object):
    def __init__(self, y, temperature, density):
        self.y = y
        self.temperature = temperature
        self.density = density

    def plot_vs_temperature(self, **kwargs):
        from matplotlib.pyplot import gca
        ax = gca()

        ax.loglog(self.temperature, self.y.T, **kwargs)
        ax.set_xlabel('$T_\mathrm{e}\ [\mathrm{eV}]$')
        ax.set_ylim(0.05, 1.3)
        self._annotate_ionisation_stages(ax)

    def _annotate_ionisation_stages(self, ax):
        max_pos = self.y.argmax(axis=-1)
        for i in xrange(self.y.shape[0]):
            index = max_pos[i]
            xy = self.temperature[index], self.y[i, index]
            s = '$%d^+$' % (i,)
            ax.annotate(s, xy, ha='center')


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    from atomic_data import AtomicData

    ad = AtomicData.from_element('Ar')
    coronal = CoronalEquilibrium(ad)

    temperature = np.logspace(0, 5, 300)
    y = coronal.ionisation_stage_distribution(temperature, density=1e19)

    plt.figure(1); plt.clf()
    y.plot_vs_temperature()

    plt.draw()
    plt.show()

{'fraction_non_alphanumeric': 0.1178433889602054, 'fraction_numerical': 0.010269576379974325, 'mean_word_length': 2.685903500473037, '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': '162696', 'n_tokens_mistral': 1345, 'n_tokens_neox': 1252, 'n_words': 261}