from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Sequence, cast
import numpy as np
from numba import njit
from scipy.special import exp1
from weno4 import weno4
import lightweaver.constants as Const
from .utils import sequence_repr
if TYPE_CHECKING:
from .atmosphere import Atmosphere
from .atomic_model import AtomicModel
from .atomic_set import SpeciesStateTable
[docs]
@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', eqPops: 'SpeciesStateTable',
Cmat: np.ndarray):
raise NotImplementedError
def __eq__(self, other: object) -> bool:
if other is self:
return True
return repr(self) == repr(other)
[docs]
@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.jLevel = atom.levels[self.j]
self.iLevel = atom.levels[self.i]
self.temperature = np.asarray(self.temperature)
self.rates = np.asarray(self.rates)
[docs]
@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', eqPops: 'SpeciesStateTable',
Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nstar = eqPops.atomicPops[self.atom.element].nStar
Cdown = self.C0 * atmos.ne * C / (self.jLevel.g * np.sqrt(atmos.temperature))
Cmat[self.i, self.j, :] += Cdown
Cmat[self.j, self.i, :] += Cdown * nstar[self.j] / nstar[self.i]
[docs]
@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.jLevel.E_SI - self.iLevel.E_SI
def compute_rates(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable',
Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nstar = eqPops.atomicPops[self.atom.element].nStar
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]
@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.iLevel.g / self.jLevel.g
def compute_rates(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable',
Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nstar = eqPops.atomicPops[self.atom.element].nStar
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]
@dataclass(eq=False, repr=False)
class CP(TemperatureInterpolationRates):
'''
Collisional (de-)excitation by protons.
Units: s^-1 m^3
'''
def compute_rates(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable',
Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nProton = eqPops['H'][-1, :]
Cdown = C * nProton
nstar = eqPops.atomicPops[self.atom.element].nStar
Cmat[self.i, self.j, :] += Cdown
Cmat[self.j, self.i, :] += Cdown * nstar[self.j] / nstar[self.i]
[docs]
@dataclass(eq=False, repr=False)
class CH(TemperatureInterpolationRates):
'''
Collisions with neutral hydrogen.
Units: s^-1 m^3
'''
def compute_rates(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable',
Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nh0 = eqPops['H'][0, :]
Cup = C * nh0
nstar = eqPops.atomicPops[self.atom.element].nStar
Cmat[self.j, self.i, :] += Cup
Cmat[self.i, self.j, :] += Cup * nstar[self.i] / nstar[self.j]
[docs]
@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', eqPops: 'SpeciesStateTable',
Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nh0 = eqPops['H'][0, :]
Cdown = C * nh0
nstar = eqPops.atomicPops[self.atom.element].nStar
Cmat[self.i, self.j, :] += Cdown
[docs]
@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', eqPops: 'SpeciesStateTable',
Cmat: np.ndarray):
C = weno4(atmos.temperature, self.temperature, self.rates)
C[C < 0.0] = 0.0
nProton = eqPops['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.0000e+00, 2.1658e+02, 2.0336e+04, 1.0911e+06, 3.7114e+07,
8.3963e+08, 1.2889e+10, 1.3449e+11, 9.4002e+11, 4.2571e+12,
1.1743e+13, 1.7549e+13, 1.0806e+13, 4.9776e+11, 0.0000))
q = np.array((1.0000e+00, 2.1958e+02, 2.0984e+04, 1.1517e+06, 4.0349e+07,
9.4900e+08, 1.5345e+10, 1.7182e+11, 1.3249e+12, 6.9071e+12,
2.3531e+13, 4.9432e+13, 5.7760e+13, 3.0225e+13, 3.3641e+12))
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]
@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.iLevel = atom.levels[self.i]
self.jLevel = atom.levels[self.j]
self.cdi = np.array(self.cdi)
def compute_rates(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable',
Cmat: np.ndarray):
nstar = eqPops.atomicPops[self.atom.element].nStar
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 * Const.CM_TO_M**3
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]
@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.iLevel = atom.levels[self.i]
self.jLevel = atom.levels[self.j]
def compute_rates(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable',
Cmat: np.ndarray):
nstar = eqPops.atomicPops[self.atom.element].nStar
dE = (self.jLevel.E_SI - self.iLevel.E_SI) / Const.EV
zz = self.iLevel.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 * Const.CM_TO_M**3)
Cup *= self.fudge
Cdown = Cup * nstar[self.i, :] / nstar[self.j, :]
Cmat[self.j, self.i, :] += Cup
Cmat[self.i, self.j, :] += Cdown