<commit_before>#!/usr/bin/env python3
import argparse
import math
import os.path
import matplotlib.pyplot as plt
import numpy as np
from scipy import linalg

import artistools as at
import artistools.estimators
import artistools.nonthermal


minionfraction = 1.e-4  # minimum number fraction of the total population to include in SF solution

defaultoutputfile = 'spencerfano_cell{cell:03d}_ts{timestep:02d}_{time_days:.0f}d.pdf'


def lossfunction_ergs(energy, nne):
    H = 6.6260755e-27
    ME = 9.1093897e-28
    EV = 1.6021772e-12
    PI = 3.1415926535987
    QE = 4.80325E-10

    eulergamma = 0.577215664901532

    omegap = math.sqrt(4 * math.pi * nne * pow(QE, 2) / ME)
    zetae = H * omegap / 2 / PI
    v = math.sqrt(2 * energy / ME)
    # print(f'omegap {omegap:.2e} s^-1')
    # print(f'zetae {zetae:.2e} erg = {zetae / EV:.2e} eV')
    # print(f'v {v:.2e}  cm/s')
    if energy > 14 * EV:
        return nne * 2 * PI * pow(QE, 4) / energy * math.log(2 * energy / zetae)
    else:
        return nne * 2 * PI * pow(QE, 4) / energy * math.log(ME * pow(v, 3) / (eulergamma * pow(QE, 2) * omegap))


def lossfunction(energy_ev, nne_cgs):
    nne = nne_cgs * 1e6   # convert from cm^-3 to m^-3
    energy = energy_ev * 1.60218e-19  # convert eV to J
    qe = 1.609e-19  # elementary charge in Coulombs
    me = 9.10938e-31  # kg
    h = 6.62606e-34  # J s
    eps0 = 8.854187817e-12  # F / m

    # h = 4.1356e-15   # Planck's constant in eV * s
    # me = 0.510998e6  # mass of electron in eV/c^2
    # c = 29979245800.  # c in cm/s

    eulergamma = 0.577215664901532

    # zetae = h / (2 * math.pi) * math.sqrt((4 * math.pi * nne * (qe ** 2)) / me)
    # omegap = 2 * math.pi * zetae / h

    omegap = math.sqrt(nne * pow(qe, 2) / me / eps0)
    zetae = h * omegap / 2 / math.pi

    v = math.sqrt(2 * energy / me)  # velocity in m/s
    # print(f'omegap {omegap:.2e} s^-1')
    # print(f'zetae {zetae:.2e} J = {zetae / 1.602e-19:.2e} eV')
    # print(f'v {v:.2e} m/s')
    if energy > 14:
        lossfunc = nne * 2 * math.pi * (qe ** 4 / (4 * math.pi * eps0) ** 2) / energy * math.log(2 * energy / zetae)
    else:
        lossfunc = (nne * 2 * math.pi * (qe ** 4 / (4 * math.pi * eps0) ** 2) / energy * math.log(me * (v ** 3) /
                    (eulergamma * (qe ** 2) * omegap / (4 * math.pi * eps0))))

    # lossfunc is now in J / m
    return lossfunc / 1.60218e-19 / 100  # eV / cm


def Psecondary(epsilon, e_p, I, J):
    e_s = epsilon - I
    return 1 / (J * np.arctan((e_p - I) / (2 * J)) * (1 + ((e_s / J) ** 2)))


def get_J(Z, ionstage, ionpot_ev):
    # returns an energy in eV
    # values from Opal et al. 1971 as applied by Kozma & Fransson 1992
    if (ionstage == 1):
        if (Z == 2):  # He I
            return 15.8
        elif (Z == 10):  # Ne I
            return 24.2
        elif (Z == 18):  # Ar I
            return 10.0

    return 0.6 * ionpot_ev


def sfmatrix_add_ionization_shell(engrid, nnion, row, sfmatrix):
    # this code has been optimised and is now an almost unreadable form of the Spencer-Fano equation
    ar_xs_array = at.nonthermal.get_arxs_array_shell(engrid, row)
    deltaen = engrid[1] - engrid[0]
    ionpot_ev = row.ionpot_ev
    J = get_J(row.Z, row.ionstage, ionpot_ev)
    npts = len(engrid)

    arctanenoverj = np.arctan(engrid / J)
    arctanexpb = np.arctan((engrid - engrid[0] - ionpot_ev) / J)
    arctanexpc = np.arctan((engrid - ionpot_ev) / 2 / J)
    prefactor = deltaen * nnion * ar_xs_array / arctanexpc

    for i, en in enumerate(engrid):

        secondintegralrange = np.argwhere(engrid >= 2 * engrid[i] + ionpot_ev)
        if len(secondintegralrange) > 0:
            startindex = secondintegralrange[0][0]
        else:
            startindex = npts

        sfmatrix[i, i:startindex] += prefactor[i:startindex] * (arctanexpc[i:startindex] - arctanexpb[:startindex - i])

        if startindex < npts:
            sfmatrix[i, startindex:] += prefactor[startindex:] * (arctanenoverj[i] - arctanexpb[startindex - i: npts - i])


def addargs(parser):
    parser.add_argument('-modelpath', default='.',
                        help='Path to ARTIS folder')

    parser.add_argument('-timedays', '-time', '-t',
                        help='Time in days to plot')

    parser.add_argument('-timestep', '-ts', type=int, default=70,
                        help='Timestep number to plot')

    parser.add_argument('-modelgridindex', '-cell', type=int, default=0,
                        help='Modelgridindex to plot')

    parser.add_argument('--print-lines', action='store_true', default=False,
                        help='Output details of matching line details to standard out')<commit_msg>Make code more concise<commit_after>#!/usr/bin/env python3
import argparse
import math
import os.path
import matplotlib.pyplot as plt
import numpy as np
from scipy import linalg

import artistools as at
import artistools.estimators
import artistools.nonthermal


minionfraction = 1.e-4  # minimum number fraction of the total population to include in SF solution

defaultoutputfile = 'spencerfano_cell{cell:03d}_ts{timestep:02d}_{time_days:.0f}d.pdf'


def lossfunction_ergs(energy, nne):
    H = 6.6260755e-27
    ME = 9.1093897e-28
    EV = 1.6021772e-12
    PI = 3.1415926535987
    QE = 4.80325E-10

    eulergamma = 0.577215664901532

    omegap = math.sqrt(4 * math.pi * nne * pow(QE, 2) / ME)
    zetae = H * omegap / 2 / PI
    v = math.sqrt(2 * energy / ME)
    # print(f'omegap {omegap:.2e} s^-1')
    # print(f'zetae {zetae:.2e} erg = {zetae / EV:.2e} eV')
    # print(f'v {v:.2e}  cm/s')
    if energy > 14 * EV:
        return nne * 2 * PI * pow(QE, 4) / energy * math.log(2 * energy / zetae)
    else:
        return nne * 2 * PI * pow(QE, 4) / energy * math.log(ME * pow(v, 3) / (eulergamma * pow(QE, 2) * omegap))


def lossfunction(energy_ev, nne_cgs):
    nne = nne_cgs * 1e6   # convert from cm^-3 to m^-3
    energy = energy_ev * 1.60218e-19  # convert eV to J
    qe = 1.609e-19  # elementary charge in Coulombs
    me = 9.10938e-31  # kg
    h = 6.62606e-34  # J s
    eps0 = 8.854187817e-12  # F / m

    # h = 4.1356e-15   # Planck's constant in eV * s
    # me = 0.510998e6  # mass of electron in eV/c^2
    # c = 29979245800.  # c in cm/s

    eulergamma = 0.577215664901532

    # zetae = h / (2 * math.pi) * math.sqrt((4 * math.pi * nne * (qe ** 2)) / me)
    # omegap = 2 * math.pi * zetae / h

    omegap = math.sqrt(nne * pow(qe, 2) / me / eps0)
    zetae = h * omegap / 2 / math.pi

    v = math.sqrt(2 * energy / me)  # velocity in m/s
    # print(f'omegap {omegap:.2e} s^-1')
    # print(f'zetae {zetae:.2e} J = {zetae / 1.602e-19:.2e} eV')
    # print(f'v {v:.2e} m/s')
    if energy > 14:
        lossfunc = nne * 2 * math.pi * (qe ** 4 / (4 * math.pi * eps0) ** 2) / energy * math.log(2 * energy / zetae)
    else:
        lossfunc = (nne * 2 * math.pi * (qe ** 4 / (4 * math.pi * eps0) ** 2) / energy * math.log(me * (v ** 3) /
                    (eulergamma * (qe ** 2) * omegap / (4 * math.pi * eps0))))

    # lossfunc is now in J / m
    return lossfunc / 1.60218e-19 / 100  # eV / cm


def Psecondary(epsilon, e_p, I, J):
    e_s = epsilon - I
    return 1 / (J * np.arctan((e_p - I) / (2 * J)) * (1 + ((e_s / J) ** 2)))


def get_J(Z, ionstage, ionpot_ev):
    # returns an energy in eV
    # values from Opal et al. 1971 as applied by Kozma & Fransson 1992
    if (ionstage == 1):
        if (Z == 2):  # He I
            return 15.8
        elif (Z == 10):  # Ne I
            return 24.2
        elif (Z == 18):  # Ar I
            return 10.0

    return 0.6 * ionpot_ev


def sfmatrix_add_ionization_shell(engrid, nnion, row, sfmatrix):
    # this code has been optimised and is now an almost unreadable form of the Spencer-Fano equation
    ar_xs_array = at.nonthermal.get_arxs_array_shell(engrid, row)
    deltaen = engrid[1] - engrid[0]
    ionpot_ev = row.ionpot_ev
    J = get_J(row.Z, row.ionstage, ionpot_ev)
    npts = len(engrid)

    arctanenoverj = np.arctan(engrid / J)
    arctanexpb = np.arctan((engrid - engrid[0] - ionpot_ev) / J)
    arctanexpc = np.arctan((engrid - ionpot_ev) / 2 / J)
    prefactor = deltaen * nnion * ar_xs_array / arctanexpc

    for i, en in enumerate(engrid):

        startindex = min(2 * i + math.ceil(ionpot_ev / deltaen), npts)

        sfmatrix[i, i:startindex] += prefactor[i:startindex] * (arctanexpc[i:startindex] - arctanexpb[:startindex - i])

        if startindex < npts:
            sfmatrix[i, startindex:] += prefactor[startindex:] * (arctanenoverj[i] - arctanexpb[startindex - i: npts - i])


def addargs(parser):
    parser.add_argument('-modelpath', default='.',
                        help='Path to ARTIS folder')

    parser.add_argument('-timedays', '-time', '-t',
                        help='Time in days to plot')

    parser.add_argument('-timestep', '-ts', type=int, default=70,
                        help='Timestep number to plot')

    parser.add_argument('-modelgridindex', '-cell', type=int, default=0,
                        help='Modelgridindex to plot')

    parser.add_argument('--print-lines', action='store_true', default=False,
                        help='Output details of matching line details to standard out')
{'fraction_non_alphanumeric': 0.14118539386715787, 'fraction_numerical': 0.058402860548271755, 'mean_word_length': 2.309429903191108, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 9, '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': '1041886', 'n_tokens_mistral': 3983, 'n_tokens_neox': 3528, 'n_words': 927}