lightweaver 0.16.1__cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

lightweaver/broadening.py

@@ -0,0 +1,605 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence
+
+import astropy.units as u
+import numpy as np
+
+import lightweaver.constants as Const
+from .atomic_table import PeriodicTable
+from .barklem import Barklem
+
+if TYPE_CHECKING:
+    from .atmosphere import Atmosphere
+    from .atomic_model import AtomicLine
+    from .atomic_set import SpeciesStateTable
+
+@dataclass
+class LineBroadeningResult:
+    '''
+    Result expected from instances of `LineBroadening.broaden`.
+    '''
+    natural: np.ndarray
+    Qelast: np.ndarray
+    other: Optional[List] = None
+
+
+@dataclass
+class LineBroadener:
+    '''
+    Base class for broadening implementations. To be used if your broadener
+    does something special and can't just return an array.
+    '''
+    def __repr__(self):
+        raise NotImplementedError
+
+    def setup(self, line: 'AtomicLine'):
+        pass
+
+    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> Any:
+        raise NotImplementedError
+
+@dataclass
+class StandardLineBroadener(LineBroadener):
+    '''
+    Standard base class for broadening implementations. Unless you need to do
+    something weird, inherit from this one.
+    '''
+    def __repr__(self):
+        raise NotImplementedError
+
+    def setup(self, line: 'AtomicLine'):
+        pass
+
+    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
+        raise NotImplementedError
+
+
+@dataclass
+class LineBroadening:
+    '''
+    Standard component of AtomicLine to compute the broadening parameters in
+    a flexible way.
+
+    For most Voigt-like situations, this should be usable without
+    modifications, but this class can be inherited from to make substantial
+    modifications.
+
+    Parameters
+    ----------
+    natural : list of StandardLineBroadener
+        List of broadening terms that are not elastic collisions (separated for PRD).
+    elastic : list of StandardLineBroadener
+        List of elastic broadening terms.
+    other : list of LineBroadener, optional
+        List of other broadening terms, not used by the VoigtLine by default,
+        but existing to provide _options_ (default: None)
+    '''
+    natural: List[StandardLineBroadener]
+    elastic: List[StandardLineBroadener]
+    other: Optional[List[LineBroadener]] = None
+
+    def __repr__(self):
+        otherStr = '' if self.other is None else ', other=%s' % repr(self.other)
+        s = 'LineBroadening(natural=%s, elastic=%s%s)' % (repr(self.natural),
+                                                          repr(self.elastic), otherStr)
+        return s
+
+    def __post_init__(self):
+        if len(self.natural) == 0 and len(self.elastic) == 0:
+            raise ValueError('No standard broadening terms provided to LineBroadening')
+
+    def setup(self, line: 'AtomicLine'):
+        b: LineBroadener
+        for b in self.natural:
+            b.setup(line)
+
+        for b in self.elastic:
+            b.setup(line)
+
+        if self.other is not None:
+            for b in self.other:
+                b.setup(line)
+
+    @staticmethod
+    def sum_broadening_list(broadeners: List[StandardLineBroadener], atmos: 'Atmosphere',
+                            eqPops: 'SpeciesStateTable') -> Optional[np.ndarray]:
+        '''
+        Sums a list of StandardLineBroadeners.
+        '''
+        if len(broadeners) == 0:
+            return None
+
+        result = broadeners[0].broaden(atmos, eqPops)
+        for b in broadeners[1:]:
+            result += b.broaden(atmos, eqPops)
+        return result
+
+    @staticmethod
+    def compute_other_broadening(broadeners: Optional[List[LineBroadener]],
+                                 atmos: 'Atmosphere',
+                                 eqPops: 'SpeciesStateTable') -> Optional[List]:
+        '''
+        Returns a list of the computed broadening terms.
+        '''
+
+        if broadeners is None:
+            return None
+        if len(broadeners) == 0:
+            return None
+
+        result = [b.broaden(atmos, eqPops) for b in broadeners]
+        return result
+
+    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> LineBroadeningResult:
+        '''
+        Computes the broadening, this function is called by the AtomicLine object.
+        '''
+        natural = self.sum_broadening_list(self.natural, atmos, eqPops)
+        Qelast = self.sum_broadening_list(self.elastic, atmos, eqPops)
+
+        others = self.compute_other_broadening(self.other, atmos, eqPops)
+
+        if natural is None:
+            if Qelast is None:
+                raise ValueError(f'Insufficient information provided to {self}')
+            natural = np.zeros_like(Qelast) # type: ignore
+        elif Qelast is None:
+            Qelast = np.zeros_like(natural)
+
+        return LineBroadeningResult(natural=natural,
+                                    Qelast=Qelast, other=others)
+
+
+@dataclass(eq=False)
+class VdwApprox(StandardLineBroadener):
+    '''
+    Base class for van der Waals approximation using a list of coefficients
+    (vals).
+    '''
+    vals: Sequence[float]
+    line: 'AtomicLine' = field(init=False)
+
+    def setup(self, line: 'AtomicLine'):
+        self.line = line
+
+    def __repr__(self):
+        s = '%s(vals=%s)' % (type(self).__name__, repr(self.vals))
+        return s
+
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+
+        if self.vals != other.vals:
+            return False
+
+        try:
+            if self.line != other.line:
+                return False
+        except AttributeError:
+            pass
+
+        return True
+
+
+@dataclass(eq=False, repr=False)
+class VdwUnsold(VdwApprox):
+    '''
+    Implementation of the Unsold method for van der Waals broadening.
+    '''
+    def setup(self, line: 'AtomicLine'):
+        self.line = line
+        if len(self.vals) != 2:
+            raise ValueError('VdwUnsold expects 2 coefficients (%s)' % repr(line))
+
+        Z = line.jLevel.stage + 1
+        cont = line.overlyingContinuumLevel
+
+        deltaR = (Const.ERydberg / (cont.E_SI - line.jLevel.E_SI))**2 \
+                 - (Const.ERydberg / (cont.E_SI - line.iLevel.E_SI))**2
+        fourPiEps0 = 4.0 * np.pi * Const.Epsilon0
+        self.C625 = (2.5 * Const.QElectron**2 / fourPiEps0 * Const.ABarH / fourPiEps0 \
+                     * 2 * np.pi * (Z * Const.RBohr)**2 / Const.HPlanck * deltaR)**0.4
+
+        element = line.atom.element
+
+        self.vRel35He = (8.0 * Const.KBoltzmann / (np.pi*Const.Amu * element.mass)\
+                         * (1.0 + element.mass / PeriodicTable[2].mass))**0.3
+        self.vRel35H = (8.0 * Const.KBoltzmann / (np.pi*Const.Amu * element.mass)\
+                         * (1.0 + element.mass / PeriodicTable[1].mass))**0.3
+
+
+    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
+        '''
+        The function that is called by LineBroadening.
+
+        Parameters
+        ----------
+        atmos : Atmosphere
+            The atmosphere in which to compute the broadening.
+        eqPops : SpeciesStateTable
+            The populations to use for computing the broadening.
+
+        Returns
+        -------
+        broad : np.ndarray
+            An array detailing the broadening at each location in the
+            atmosphere [Nspace].
+        '''
+        heAbund = eqPops.abundance[PeriodicTable[2]]
+        cross = 8.08 * (self.vals[0] * self.vRel35H \
+                             + self.vals[1] * heAbund * self.vRel35He) * self.C625
+        nHGround = eqPops['H'][0, :]
+        broad = cross * atmos.temperature**0.3 * nHGround
+        return broad
+
+
+@dataclass(eq=False, repr=False)
+class VdwBarklem(VdwApprox):
+    '''
+    Implementation of the Barklem method for van der Waals broadening.
+    '''
+    def setup(self, line: 'AtomicLine'):
+        self.line = line
+        if len(self.vals) != 2:
+            raise ValueError('VdwBarklem expects 2 coefficients (%s)' % (repr(line)))
+        newVals = Barklem.get_active_cross_section(line.atom, line, self.vals)
+
+        self.barklemVals = newVals
+
+        Z = line.jLevel.stage + 1
+        cont = line.overlyingContinuumLevel
+
+        deltaR = (Const.ERydberg / (cont.E_SI - line.jLevel.E_SI))**2 \
+                 - (Const.ERydberg / (cont.E_SI - line.iLevel.E_SI))**2
+        fourPiEps0 = 4.0 * np.pi * Const.Epsilon0
+        self.C625 = (2.5 * Const.QElectron**2 / fourPiEps0 * Const.ABarH / fourPiEps0 \
+                     * 2 * np.pi * (Z * Const.RBohr)**2 / Const.HPlanck * deltaR)**0.4
+
+        element = line.atom.element
+
+        self.vRel35He = (8.0 * Const.KBoltzmann / (np.pi*Const.Amu * element.mass)\
+                         * (1.0 + element.mass / PeriodicTable[2].mass))**0.3
+
+
+    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
+        '''
+        The function that is called by LineBroadening.
+
+        Parameters
+        ----------
+        atmos : Atmosphere
+            The atmosphere in which to compute the broadening.
+        eqPops : SpeciesStateTable
+            The populations to use for computing the broadening.
+
+        Returns
+        -------
+        broad : np.ndarray
+            An array detailing the broadening at each location in the
+            atmosphere [Nspace].
+        '''
+        heAbund = eqPops.abundance[PeriodicTable[2]]
+        nHGround = eqPops['H'][0, :]
+        cross = 8.08 * self.barklemVals[2] * heAbund * self.vRel35He * self.C625
+
+        broad = self.barklemVals[0] * atmos.temperature**(0.5*(1.0-self.barklemVals[1])) \
+                 + cross * atmos.temperature**0.3
+        broad *= nHGround
+        return broad
+
+
+@dataclass(eq=False)
+class RadiativeBroadening(StandardLineBroadener):
+    '''
+    Simple constant radiative broadening with coefficient gamma.
+    '''
+    gamma: float
+    line: 'AtomicLine' = field(init=False)
+
+    def setup(self, line: 'AtomicLine'):
+        self.line = line
+
+    def __repr__(self):
+        s = '%s(gamma=%g)' % (type(self).__name__, self.gamma)
+        return s
+
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+
+        if self.gamma != other.gamma:
+            return False
+
+        try:
+            if self.line != other.line:
+                return False
+        except AttributeError:
+            pass
+
+        return True
+
+    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
+        '''
+        The function that is called by LineBroadening.
+
+        Parameters
+        ----------
+        atmos : Atmosphere
+            The atmosphere in which to compute the broadening.
+        eqPops : SpeciesStateTable
+            The populations to use for computing the broadening.
+
+        Returns
+        -------
+        broad : np.ndarray
+            An array detailing the broadening at each location in the
+            atmosphere [Nspace].
+        '''
+        return np.ones_like(atmos.temperature) * self.gamma
+
+@dataclass
+class QuadraticStarkBroadening(StandardLineBroadener):
+    '''
+    Lindholm theory result for Quadratic Stark broadening by electrons and
+    singly ionised particles.
+    Follows HM2014 pp. 238-239, uses C4 from Traving 1960 via RH.
+    '''
+    coeff: float
+    line: 'AtomicLine' = field(init=False)
+
+    def __repr__(self):
+        s = '%s(coeff=%g)' % (type(self).__name__, self.coeff)
+        return s
+
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+
+        if self.coeff != other.coeff:
+            return False
+
+        try:
+            if self.line != other.line:
+                return False
+        except AttributeError:
+            pass
+
+        return True
+
+    def setup(self, line: 'AtomicLine'):
+        self.line = line
+        weight = line.atom.element.mass
+        C = 8.0 * Const.KBoltzmann / (np.pi * Const.Amu * weight)
+        Cm = (1.0 + weight / (Const.MElectron / Const.Amu))**(1.0/6.0)
+        # NOTE(cmo): 28.0 is average atomic weight
+        Cm += (1.0 + weight / (28.0))**(1.0/6.0)
+        self.C = C
+        self.Cm = Cm
+
+        Z = line.iLevel.stage + 1
+        cont = line.overlyingContinuumLevel
+
+        E_Ryd = Const.ERydberg / (1.0 + Const.MElectron / (weight * Const.Amu))
+        neff_l = Z * np.sqrt(E_Ryd / (cont.E_SI - line.iLevel.E_SI))
+        neff_u = Z * np.sqrt(E_Ryd / (cont.E_SI - line.jLevel.E_SI))
+
+        C4 = Const.QElectron**2 / (4.0 * np.pi * Const.Epsilon0) \
+            * Const.RBohr \
+            * (2.0 * np.pi * Const.RBohr**2 / Const.HPlanck) / (18.0 * Z**4) \
+            * ((neff_u * (5.0 * neff_u**2 + 1.0))**2 \
+                - (neff_l * (5.0 * neff_l**2 + 1.0))**2)
+        self.cStark23 = 11.37 * (self.coeff * C4)**(2.0/3.0)
+
+    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
+        '''
+        The function that is called by LineBroadening.
+
+        Parameters
+        ----------
+        atmos : Atmosphere
+            The atmosphere in which to compute the broadening.
+        eqPops : SpeciesStateTable
+            The populations to use for computing the broadening.
+
+        Returns
+        -------
+        broad : np.ndarray
+            An array detailing the broadening at each location in the
+            atmosphere [Nspace].
+        '''
+        vRel = (self.C * atmos.temperature)**(1.0/6.0) * self.Cm
+        stark = self.cStark23 * vRel * atmos.ne
+        return stark
+
+@dataclass
+class MultiplicativeStarkBroadening(StandardLineBroadener):
+    '''
+    Simple expression for multiplicative Stark broadening, assumes that this
+    can be expresed as a constant * ne.
+    '''
+    coeff: float
+
+    def __repr__(self):
+        s = '%s(coeff=%g)' % (type(self).__name__, self.coeff)
+        return s
+
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+
+        if self.coeff != other.coeff:
+            return False
+
+        return True
+
+    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
+        '''
+        The function that is called by LineBroadening.
+
+        Parameters
+        ----------
+        atmos : Atmosphere
+            The atmosphere in which to compute the broadening.
+        eqPops : SpeciesStateTable
+            The populations to use for computing the broadening.
+
+        Returns
+        -------
+        broad : np.ndarray
+            An array detailing the broadening at each location in the
+            atmosphere [Nspace].
+        '''
+        return self.coeff * atmos.ne # type: ignore
+
+@dataclass
+class HydrogenLinearStarkBroadening(StandardLineBroadener):
+    '''
+    Linear Stark broadening for the case of Hydrogen from Sutton 1978 (like
+    RH).
+
+    `reproduceOldRHBug` allows for reproducing an issue in older versions of RH
+    where the broadening effects was a factor os ~2pi too small as it had not
+    been converted to angular frequency.
+    '''
+    line: 'AtomicLine' = field(init=False)
+    reproduceOldRHBug: bool = False
+
+    def __repr__(self):
+        s = f'{type(self).__name__}(reproduceOldRHBug={self.reproduceOldRHBug})'
+        return s
+
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+
+        try:
+            if self.line != other.line:
+                return False
+        except AttributeError:
+            pass
+
+        return True
+
+    def setup(self, line: 'AtomicLine'):
+        self.line = line
+
+        if line.atom.element.Z != 1:
+            raise ValueError('HydrogenicLinearStarkBroadening applied to non-Hydrogen line')
+
+    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
+        '''
+        The function that is called by LineBroadening.
+
+        Parameters
+        ----------
+        atmos : Atmosphere
+            The atmosphere in which to compute the broadening.
+        eqPops : SpeciesStateTable
+            The populations to use for computing the broadening.
+
+        Returns
+        -------
+        broad : np.ndarray
+            An array detailing the broadening at each location in the
+            atmosphere [Nspace].
+        '''
+        nUpper = int(np.round(np.sqrt(0.5*self.line.jLevel.g)))
+        nLower = int(np.round(np.sqrt(0.5*self.line.iLevel.g)))
+
+        a1 = 0.642 if nUpper - nLower == 1 else 1.0
+        C = a1 * 0.6 * (nUpper**2 - nLower**2)
+        if not self.reproduceOldRHBug:
+            C *= 4.0 * np.pi * 0.425
+        GStark = C * u.Unit('m-2').to('cm-2') * atmos.ne**(2.0/3.0)
+        return GStark
+
+@dataclass(eq=False)
+class ScaledExponentBroadening(StandardLineBroadener):
+    '''
+    Broadening implementation following the CRTAF ScaledExponents recipe.
+        scaling * T**a * n_H(0)**b * n_e**c
+
+    scaling : float
+        Scalar multiplication term
+    temperatureExp : float
+        Temperature exponent
+    hydrogenExp : float
+        Neutral (ground-state) hydrogen density exponent
+    electronExp : float
+        Electron density exponent
+    '''
+    scaling: float
+    temperatureExp: float
+    hydrogenExp: float
+    electronExp: float
+    line: 'AtomicLine' = field(init=False)
+
+    def setup(self, line: 'AtomicLine'):
+        self.line = line
+
+    def __repr__(self):
+        s = '%s(scaling=%g, temperatureExp=%g, hydrogenExp=%g, electronExp=%g)' % (type(self).__name__,
+                self.scaling, self.temperatureExp, self.hydrogenExp, self.electronExp)
+        return s
+
+    def __eq__(self, other):
+        if type(self) is not type(other):
+            return False
+
+        if self.scaling != other.scaling:
+            return False
+        if self.temperatureExp != other.temperatureExp:
+            return False
+        if self.hydrogenExp != other.hydrogenExp:
+            return False
+        if self.electronExp != other.electronExp:
+            return False
+
+        try:
+            if self.line != other.line:
+                return False
+        except AttributeError:
+            pass
+
+        return True
+
+    def broaden(self, atmos: 'Atmosphere', eqPops: 'SpeciesStateTable') -> np.ndarray:
+        '''
+        The function that is called by LineBroadening.
+
+        Parameters
+        ----------
+        atmos : Atmosphere
+            The atmosphere in which to compute the broadening.
+        eqPops : SpeciesStateTable
+            The populations to use for computing the broadening.
+
+        Returns
+        -------
+        broad : np.ndarray
+            An array detailing the broadening at each location in the
+            atmosphere [Nspace].
+        '''
+        result = np.ones_like(atmos.temperature)
+
+        if self.temperatureExp > 0.0:
+            if self.temperatureExp == 1.0:
+                result *= atmos.temperature
+            else:
+                result *= atmos.temperature**self.temperatureExp
+
+        if self.hydrogenExp > 0.0:
+            nHGround = eqPops['H'][0, :]
+            if self.hydrogenExp == 1.0:
+                result *= nHGround
+            else:
+                result *= nHGround**self.hydrogenExp
+
+        if self.electronExp > 0.0:
+            if self.electronExp == 1.0:
+                result *= atmos.ne
+            else:
+                result *= atmos.ne**self.electronExp
+
+        return self.scaling * result