from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Sequence, cast
import astropy.units as u
import numpy as np
from numba import njit
from scipy.special import exp1
from weno4 import weno4
import lightweaver.constants as Const
from .deprecation import deprecated_names
from .utils import sequence_repr
if TYPE_CHECKING:
from .atmosphere import Atmosphere
from .atomic_model import AtomicModel
from .atomic_set import SpeciesStateTable
[docs]
@deprecated_names
@dataclass
class CollisionalRates:
"""
Base class for all CollisionalRates.
"""
j: int
i: int
atom: 'AtomicModel' = field(init=False)
def __repr__(self):
s = 'CollisionalRates(j=%d, i=%d)' % (self.j, self.i)
return s
def setup(self, atom):
pass
def compute_rates(self, atmos: 'Atmosphere', eq_pops: 'SpeciesStateTable', Cmat: np.ndarray):
raise NotImplementedError
def __eq__(self, other: object) -> bool:
if other is self:
return True
return repr(self) == repr(other)
[docs]
@deprecated_names(attrs=('i_level', 'j_level'))
@dataclass(eq=False)
class TemperatureInterpolationRates(CollisionalRates):
"""
Base class for rates defined by interpolating a coefficient on a
temperature grid.
"""
temperature: Sequence[float]
rates: Sequence[float]
def __repr__(self):
s = '%s(j=%d, i=%d, temperature=%s, rates=%s)' % (
type(self).__name__,
self.j,
self.i,
sequence_repr(self.temperature),
sequence_repr(self.rates),
)
return s
def setup(self, atom):
i, j = self.i, self.j
self.i = min(i, j)
self.j = max(i, j)
self.atom = atom
self.j_level = atom.levels[self.j]
self.i_level = atom.levels[self.i]
self.temperature = np.asarray(self.temperature)
self.rates = np.asarray(self.rates)
[docs]
@deprecated_names
@dataclass(eq=False, repr=False)
class Omega(TemperatureInterpolationRates):
"""
Collisional (de-)excitation of ions by electrons (dimensionless).
Omega as in Seaton's collision strength.
Rate scales as 1/(sqrt(T)) exp(DeltaE).
"""
def setup(self, atom):
super().setup(atom)
self.C0 = (
Const.ERydberg
/ np.sqrt(Const.MElectron)
* np.pi
* Const.RBohr**2
* np.sqrt(8.0 / (np.pi * Const.KBoltzmann))
)
def compute_rates(self, atmos: 'Atmosphere', eq_pops: 'SpeciesStateTable', Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nstar = eq_pops.atomic_pops[self.atom.element].n_star
Cdown = self.C0 * atmos.ne * C / (self.j_level.g * np.sqrt(atmos.temperature))
Cmat[self.i, self.j, :] += Cdown
Cmat[self.j, self.i, :] += Cdown * nstar[self.j] / nstar[self.i]
[docs]
@deprecated_names
@dataclass(eq=False, repr=False)
class CI(TemperatureInterpolationRates):
"""
Collisional ionisation by electrons.
Units: s^-1 K^-1/2 m^3
Rate scales as sqrt(T) exp(DeltaE)
"""
def setup(self, atom):
super().setup(atom)
self.dE = self.j_level.E_SI - self.i_level.E_SI
def compute_rates(self, atmos: 'Atmosphere', eq_pops: 'SpeciesStateTable', Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nstar = eq_pops.atomic_pops[self.atom.element].n_star
Cup = (
C
* atmos.ne
* np.exp(-self.dE / (Const.KBoltzmann * atmos.temperature))
* np.sqrt(atmos.temperature)
)
Cmat[self.j, self.i, :] += Cup
Cmat[self.i, self.j, :] += Cup * nstar[self.i] / nstar[self.j]
[docs]
@deprecated_names
@dataclass(eq=False, repr=False)
class CE(TemperatureInterpolationRates):
"""
Collisional (de-)excitation of neutrals by electrons.
Units: s^-1 K^-1/2 m^3
Rate scales as sqrt(T) exp(DeltaE)
"""
def setup(self, atom):
super().setup(atom)
self.gij = self.i_level.g / self.j_level.g
def compute_rates(self, atmos: 'Atmosphere', eq_pops: 'SpeciesStateTable', Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nstar = eq_pops.atomic_pops[self.atom.element].n_star
Cdown = C * atmos.ne * self.gij * np.sqrt(atmos.temperature)
Cmat[self.i, self.j, :] += Cdown
Cmat[self.j, self.i, :] += Cdown * nstar[self.j] / nstar[self.i]
[docs]
@deprecated_names
@dataclass(eq=False, repr=False)
class CP(TemperatureInterpolationRates):
"""
Collisional (de-)excitation by protons.
Units: s^-1 m^3
"""
def compute_rates(self, atmos: 'Atmosphere', eq_pops: 'SpeciesStateTable', Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nProton = eq_pops['H'][-1, :]
Cdown = C * nProton
nstar = eq_pops.atomic_pops[self.atom.element].n_star
Cmat[self.i, self.j, :] += Cdown
Cmat[self.j, self.i, :] += Cdown * nstar[self.j] / nstar[self.i]
[docs]
@deprecated_names
@dataclass(eq=False, repr=False)
class CH(TemperatureInterpolationRates):
"""
Collisions with neutral hydrogen.
Units: s^-1 m^3
"""
def compute_rates(self, atmos: 'Atmosphere', eq_pops: 'SpeciesStateTable', Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nh0 = eq_pops['H'][0, :]
Cup = C * nh0
nstar = eq_pops.atomic_pops[self.atom.element].n_star
Cmat[self.j, self.i, :] += Cup
Cmat[self.i, self.j, :] += Cup * nstar[self.i] / nstar[self.j]
[docs]
@deprecated_names
@dataclass(eq=False, repr=False)
class ChargeExchangeNeutralH(TemperatureInterpolationRates):
"""
Charge exchange with neutral hydrogen.
Units: s^-1 m^3
Note: downward rate only.
"""
def compute_rates(self, atmos: 'Atmosphere', eq_pops: 'SpeciesStateTable', Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nh0 = eq_pops['H'][0, :]
Cdown = C * nh0
Cmat[self.i, self.j, :] += Cdown
[docs]
@deprecated_names
@dataclass(eq=False, repr=False)
class ChargeExchangeProton(TemperatureInterpolationRates):
"""
Charge exchange with protons.
Units: s^-1 m^3
Note: upward rate only.
"""
def compute_rates(self, atmos: 'Atmosphere', eq_pops: 'SpeciesStateTable', Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nProton = eq_pops['H'][-1, :]
Cup = C * nProton
Cmat[self.j, self.i, :] += Cup
def fone(x):
# return np.where(x <= 50.0, np.exp(x) * exp1(x), 1.0/x)
return np.where(x <= 50.0, np.exp(x) * exp1(x), (1.0 - 1.0 / x + 2.0 / x**2) / x)
@njit(cache=True)
def ftwo(x):
p = np.array(
(
1.0000e00,
2.1658e02,
2.0336e04,
1.0911e06,
3.7114e07,
8.3963e08,
1.2889e10,
1.3449e11,
9.4002e11,
4.2571e12,
1.1743e13,
1.7549e13,
1.0806e13,
4.9776e11,
0.0000,
)
)
q = np.array(
(
1.0000e00,
2.1958e02,
2.0984e04,
1.1517e06,
4.0349e07,
9.4900e08,
1.5345e10,
1.7182e11,
1.3249e12,
6.9071e12,
2.3531e13,
4.9432e13,
5.7760e13,
3.0225e13,
3.3641e12,
)
)
def ftwo_impl(x):
if x > 4.0:
px = p[0]
xFact = 1.0
for i in range(1, 15):
xFact /= x
px += p[i] * xFact
qx = q[0]
xFact = 1.0
for i in range(1, 15):
xFact /= x
qx += q[i] * xFact
return px / (qx * x**2)
else:
gamma = 0.5772156649
f0x = np.pi**2 / 12.0
term = 1.0
count = 0.0
fact = 1.0
xFact = 1.0
while abs(term / f0x) > 1e-8:
count += 1.0
fact *= count
xFact *= -x
term = xFact / (count**2 * fact)
f0x += term
if count > 100.0:
raise ValueError('ftwo too slow to converge')
y = np.exp(x) * ((np.log(x) + gamma) ** 2 * 0.5 + f0x)
return y
y = np.empty_like(x)
for i in range(x.shape[0]):
y[i] = ftwo_impl(x[i])
return y
[docs]
@deprecated_names(attrs=('i_level', 'j_level'))
@dataclass
class Ar85Cdi(CollisionalRates):
"""
Collisional ionisation rates based on Arnaud & Rothenflug (1985, ApJS 60).
Units remain in CGS as per paper.
"""
cdi: Sequence[Sequence[float]]
def __repr__(self):
s = 'Ar85Cdi(j=%d, i=%d, cdi=%s)' % (self.j, self.i, sequence_repr(self.cdi))
return s
def setup(self, atom):
i, j = self.i, self.j
self.i = min(i, j)
self.j = max(i, j)
self.atom = atom
self.i_level = atom.levels[self.i]
self.j_level = atom.levels[self.j]
self.cdi = np.array(self.cdi)
def compute_rates(self, atmos: 'Atmosphere', eq_pops: 'SpeciesStateTable', Cmat: np.ndarray):
nstar = eq_pops.atomic_pops[self.atom.element].n_star
Cup = np.zeros(atmos.Nspace)
cdi = cast(np.ndarray, self.cdi)
for m in range(cdi.shape[0]):
xj = cdi[m, 0] * Const.EV / (Const.KBoltzmann * atmos.temperature)
fac = np.exp(-xj) * np.sqrt(xj)
fxj = (
cdi[m, 1]
+ cdi[m, 2] * (1.0 + xj)
+ (cdi[m, 3] - xj * (cdi[m, 1] + cdi[m, 2] * (2.0 + xj))) * fone(xj)
+ cdi[m, 4] * xj * ftwo(xj)
)
fxj *= fac
fac = 6.69e-7 / cdi[m, 0] ** 1.5
Cup += fac * (fxj << u.Unit('cm3')).to('m3').value
Cup[Cup < 0] = 0.0
Cup *= atmos.ne
Cdown = Cup * nstar[self.i] / nstar[self.j]
Cmat[self.i, self.j, :] += Cdown
Cmat[self.j, self.i, :] += Cup
[docs]
@deprecated_names(attrs=('i_level', 'j_level'))
@dataclass
class Burgess(CollisionalRates):
"""
Collisional ionisation from excited states from Burgess & Chidichimo
(1983, MNRAS 203, 1269).
Fudge parameter is dimensionless.
"""
fudge: float = 1.0
def __repr__(self):
s = 'Burgess(j=%d, i=%d, fudge=%g)' % (self.j, self.i, self.fudge)
return s
def setup(self, atom):
i, j = self.i, self.j
self.i = min(i, j)
self.j = max(i, j)
self.atom = atom
self.i_level = atom.levels[self.i]
self.j_level = atom.levels[self.j]
def compute_rates(self, atmos: 'Atmosphere', eq_pops: 'SpeciesStateTable', Cmat: np.ndarray):
nstar = eq_pops.atomic_pops[self.atom.element].n_star
dE = (self.j_level.E_SI - self.i_level.E_SI) / Const.EV
zz = self.i_level.stage
betaB = 0.25 * (np.sqrt((100.0 * zz + 91.0) / (4.0 * zz + 3.0)) - 5.0)
cbar = 2.3
dEkT = dE * Const.EV / (Const.KBoltzmann * atmos.temperature)
dEkT = np.minimum(dEkT, 500)
invdEkT = 1.0 / dEkT
wlog = np.log(1.0 + invdEkT)
wb = wlog ** (betaB / (1.0 + invdEkT))
Cup = (
2.1715e-8
* cbar
* (13.6 / dE) ** 1.5
* np.sqrt(dEkT)
* exp1(dEkT)
* wb
* (atmos.ne << u.Unit('m-3')).to('cm-3').value
)
Cup *= self.fudge
Cdown = Cup * nstar[self.i, :] / nstar[self.j, :]
Cmat[self.j, self.i, :] += Cup
Cmat[self.i, self.j, :] += Cdown