lightweaver 0.15.0__cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

lightweaver/iteration_update.py

@@ -0,0 +1,134 @@
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List
+
+if TYPE_CHECKING:
+    from . import Context
+
+@dataclass
+class IterationUpdate:
+    '''
+    Stores the results of an iteration of one of the backend functions, and
+    determines how to format this for printing. All changes refer to relative
+    change.
+
+    Attributes
+    ----------
+    ctx : Context
+        The context with which this update is associated.
+    crsw :  float
+        The current value of the collisional radiative switching parameter.
+    updatedJ : bool
+        Whether the iteration affected the global J grid.
+    dJMax : float
+        The maximum change in J.
+    dJMaxIdx : int
+        The index of the maximum change of J in a flattened array of J.
+    updatedPops : bool
+        Whether the active atomic populations were modified by the iteration.
+    dPops : List[float]
+        The maximum change in each active population.
+    dPopsMaxIdx : List[int]
+        The location of the maximum change in each population in the flattened
+        population array.
+    ngAccelerated : List[bool]
+        Whether the atomic populations were modified by Ng Acceleration (per species, due to thresholding).
+    updatedNe : bool
+        Whether the electron density in the atmosphere was affected by the iteration.
+    dNeMax : float
+        The maximum change in the electron density.
+    dNeMaxIdx : int
+        The location of the maximum change in the electron density array.
+    updatedRho : bool
+        Whether the iteration affected the value of rhoPrd on PRD lines.
+    NprdSubIter : int
+        The number of PRD sub-iterations taken (if multiple),
+    dRho : List[float]
+        The maximum change in rho for each spectral line treated with PRD, in
+        the order of the lines on each activeAtom. These values are repeated for
+        each sub-iteration < NprdSubIter.
+    dRhoMaxIdx : List[int]
+        The location of the maximum change in rho for each PRD line.
+    updatedJPrd : bool
+        Whether the PRD iteration affected J.
+    dJPrdMax : float
+        The maximum change in J during each PRD sub-iteration.
+    dJPrdMaxIdx : int
+        The location of the maximum change in J for each PRD sub-iteration.
+    dPopsMax : float
+        The maximum population change (including ne) over the iteration
+        (read-only property).
+    dRhoMax : float
+        The maximum change in the PRD rho value for any line in the final
+        subiteration (read-only property).
+    '''
+    ctx: 'Context'
+    crsw: float = 1.0
+    updatedJ: bool = False
+    dJMax: float = 0.0
+    dJMaxIdx: int = 0
+
+    updatedPops: bool = False
+    dPops: List[float] = field(default_factory=list)
+    dPopsMaxIdx: List[int] = field(default_factory=list)
+    ngAccelerated: List[bool] = field(default_factory=list)
+
+    updatedNe: bool = False
+    dNeMax: float = 0.0
+    dNeMaxIdx: int = 0
+
+    updatedRho: bool = False
+    NprdSubIter: int = 0
+    dRho: List[float] = field(default_factory=list)
+    dRhoMaxIdx: List[int] = field(default_factory=list)
+    updatedJPrd: bool = False
+    dJPrdMax: List[float] = field(default_factory=list)
+    dJPrdMaxIdx: List[int] = field(default_factory=list)
+
+    @property
+    def dPopsMax(self) -> float:
+        if len(self.dPops) == 0:
+            if self.updatedNe:
+                return self.dNeMax
+            else:
+                return 0.0
+
+        result = max(self.dPops)
+        if self.updatedNe:
+            result = max(result, self.dNeMax)
+        return result
+
+    @property
+    def dRhoMax(self) -> float:
+        if self.NprdSubIter == 0:
+            return 0.0
+        finalSubIterStart = (self.NprdSubIter - 1) * self.ctx.kwargs['spect'].NprdTrans
+        return max(self.dRho[finalSubIterStart:])
+
+    def compact_representation(self):
+        '''
+        Produce a compact string representation of the object (similar to
+        Lightweaver < v0.8).
+        '''
+        chunks = []
+        if self.crsw != 1.0:
+            chunks.append(f'CRSW: {self.crsw:.2e}')
+
+        if self.updatedJ:
+            chunks.append(f'dJ = {self.dJMax:.2e}')
+
+        if self.updatedPops:
+            for idx, delta in enumerate(self.dPops):
+                atomName = self.ctx.activeAtoms[idx].atomicModel.element.name
+                accel = ' (accelerated)' if self.ngAccelerated[idx] else ''
+                chunks.append(f'    {atomName} delta = {delta:6.4e}{accel}')
+
+        if self.updatedNe:
+            delta = self.dNeMax
+            chunks.append(f'    ne delta = {delta:6.4e}')
+
+        if self.updatedRho:
+            iterCount = self.NprdSubIter
+            dRhoMax = self.dRhoMax
+            chunks.append(f'    PRD dRho = {dRhoMax:.2e}, (sub-iterations: {iterCount})')
+
+        return '\n'.join(chunks)

lightweaver/molecule.py

@@ -0,0 +1,225 @@
+from collections import OrderedDict
+from typing import List, Optional, Tuple, Union
+
+import astropy.units as u
+import numpy as np
+from numba import njit
+from parse import parse
+
+import lightweaver.constants as Const
+from .atomic_table import Element, PeriodicTable
+
+
+# TODO(cmo): This should really be done with a generator/coroutine
+def get_next_line(data: List[str]) -> Optional[str]:
+    if len(data) == 0:
+        return None
+    for i, d in enumerate(data):
+        if d.strip().startswith('#') or d.strip() == '':
+            continue
+        break
+    d = data[i]
+    if i == len(data) - 1:
+        data[:] = []
+        return d.strip()
+    data[:] = data[i+1:]
+    return d.strip()
+
+def get_constituent(name: str) -> Tuple[int, str]:
+    res = parse('{:d}{!s}', name)
+    if res is None:
+        constituent = (1, name)
+    else:
+        constituent = (res[0], res[1])
+    return constituent
+
+def equilibrium_constant_kurucz_70(tempRange, mk, Ediss, eqc):
+    minTemp = tempRange[0]
+    maxTemp = tempRange[1]
+    kB = Const.KBoltzmann
+    CM_TO_M = u.Unit('cm').to('m')
+
+    @njit('float64(float64)')
+    def kurucz_70(T):
+        if T < minTemp or T > maxTemp:
+            return 0.0
+
+        kT = kB * T
+        eq = eqc[0]
+        for i in range(1, eqc.shape[0]):
+            eq = eq * T + eqc[i]
+        arg = Ediss / kT + eq - 1.5 * mk * np.log(T)
+        eq = np.exp(arg)
+        return eq * (CM_TO_M**3)**mk
+    return kurucz_70
+
+def equilibrium_constant_kurucz_85(tempRange, mk, Ediss, eqc):
+    minTemp = tempRange[0]
+    maxTemp = tempRange[1]
+    kB = Const.KBoltzmann
+    CM_TO_M = u.Unit('cm').to('m')
+
+    @njit('float64(float64)')
+    def kurucz_85(T):
+        if T < minTemp or T > maxTemp:
+            return 0.0
+
+        t = T * 1e-4
+        kT = kB * T
+        eq = eqc[0]
+        for i in range(1, eqc.shape[0]):
+            eq = eq * t + eqc[i]
+        eq = np.exp(Ediss / kT + eq - 1.5 * mk * np.log(T))
+        return eq * (CM_TO_M**3)**mk
+    return kurucz_85
+
+def equilibrium_constant_sauval_tatum(tempRange, Ediss, eqc):
+    minTemp = tempRange[0]
+    maxTemp = tempRange[1]
+    kB = Const.KBoltzmann
+    THETA0 = Const.Theta0
+    Ediss = Ediss / Const.EV
+
+    @njit('float64(float64)')
+    def sauval_tatum(T):
+        if T < minTemp or T > maxTemp:
+            return 0.0
+
+        theta = THETA0 / T
+        t = np.log10(theta)
+        kT = kB * T
+
+        eq = eqc[0]
+        for i in range(1, eqc.shape[0]):
+            eq = eq * t + eqc[i]
+        eq = 10**(Ediss * theta - eq) * kT
+        return eq
+    return sauval_tatum
+
+
+class Molecule:
+    '''
+    Simple class for working with RH molecule definitions.
+
+    Parameters
+    ----------
+    filePath : str
+        Path from which to load molecular data. Use
+        `get_default_molecule_path` for the path to the default RH
+        distribution of molecule files (C2, CH, CN, CO, CaH, H2+, H2, H2O,
+        HF, LiH, MgH, N2, NH, NO, O2, OH, SiO, TiO) all of which have the
+        '.molecule' extension.
+    '''
+    def __init__(self, filePath: str):
+        with open(filePath, 'r') as f:
+            lines = f.readlines()
+
+        l = get_next_line(lines)
+        self.name = l
+        l = get_next_line(lines)
+        self.charge = int(l)
+        if self.charge < 0 or self.charge > 1:
+            raise ValueError(('Only neutral or singly charged positive molecules '
+                              'are allowed (%s)') % self.name)
+
+        structure = get_next_line(lines)
+        constituents = [get_constituent(s.strip()) for s in structure.split(',')]
+        self.elements = [PeriodicTable[c[1]] for c in constituents]
+        self.elementCount = [c[0] for c in constituents]
+        self.Nnuclei = sum(self.elementCount)
+
+        l = get_next_line(lines)
+        self.Ediss = float(l) * Const.EV
+
+        fitStr = get_next_line(lines)
+        self.formationTempRange = [float(f) for f in get_next_line(lines).split()]
+        if len(self.formationTempRange) != 2:
+            raise ValueError("Expected two entries for formation temperature range (%s)"
+                              % self.name)
+
+        pfCoeffs = get_next_line(lines).split()
+        Npf = int(pfCoeffs[0].strip())
+        if len(pfCoeffs) != Npf+1:
+            raise ValueError("Unexpected number of partition function fit parameters (%s)"
+                             % self.name)
+        self.pfCoeffs = np.array([float(f.strip()) for f in pfCoeffs[1:]][::-1])
+
+        eqcCoeffs = get_next_line(lines).split()
+        Neqc = int(eqcCoeffs[0].strip())
+        if len(eqcCoeffs) != Neqc+1:
+            raise ValueError(("Unexpected number of equilibrium coefficient "
+                              "fit parameters (%s)") % self.name)
+        self.eqcCoeffs = np.array([float(f.strip()) for f in eqcCoeffs[1:]][::-1])
+
+        self.weight = 0.0
+        for count, ele in zip(self.elementCount, self.elements):
+            self.weight += count * ele.mass
+
+        if fitStr == 'KURUCZ_70':
+            self.equilibrium_constant = equilibrium_constant_kurucz_70(
+                                                self.formationTempRange,
+                                                self.Nnuclei - 1 - self.charge,
+                                                self.Ediss, self.eqcCoeffs)
+        elif fitStr == 'KURUCZ_85':
+            self.equilibrium_constant = equilibrium_constant_kurucz_85(
+                                                self.formationTempRange,
+                                                self.Nnuclei - 1 - self.charge,
+                                                self.Ediss, self.eqcCoeffs)
+        elif fitStr == 'SAUVAL_TATUM_84':
+            self.equilibrium_constant = equilibrium_constant_sauval_tatum(
+                                                self.formationTempRange,
+                                                self.Ediss, self.eqcCoeffs)
+        else:
+            raise ValueError(('Unknown molecular equilibrium constant fit '
+                              'method %s in molecule %s') % (fitStr, self.name))
+
+class MolecularTable:
+    '''
+    Stores a set of molecular models, can be indexed by model name to return
+    the associated model (as a string).
+    '''
+    def __init__(self, paths: Optional[List[str]]=None):
+
+        self.molecules: List[Molecule] = []
+        if paths is None:
+            self.indices = OrderedDict()
+            return
+
+        for path in paths:
+            self.molecules.append(Molecule(path))
+
+        self.indices = OrderedDict(zip([m.name for m in self.molecules],
+                                   list(range(len(self.molecules)))))
+
+    def __getitem__(self, name: str) -> Molecule:
+        name = name.upper()
+        return self.molecules[self.indices[name]]
+
+    def __contains__(self, name: Union[int, Tuple[int, int], str, Element]) -> bool:
+        if type(name) is not str:
+            return False
+
+        name = name.upper() # type: ignore
+        return name in self.indices.keys()
+
+    def __len__(self) -> int:
+        return len(self.molecules)
+
+    def __iter__(self) -> 'MolecularTableIterator':
+        return MolecularTableIterator(self)
+
+class MolecularTableIterator():
+    def __init__(self, table: MolecularTable):
+        self.table = table
+        self.index = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.index < len(self.table):
+            mol = self.table.molecules[self.index]
+            self.index += 1
+            return mol
+
+        raise StopIteration

lightweaver/multi.py

@@ -0,0 +1,113 @@
+import re
+from dataclasses import dataclass
+from typing import Tuple
+
+import astropy.units as u
+import numpy as np
+
+import lightweaver.constants as C
+from .atmosphere import Atmosphere, ScaleType
+
+
+@dataclass
+class MultiMetadata:
+    '''
+    Metadata that is stored in a MULTI atmosphere, but doesn't really belong
+    in a Lightweaver atmosphere.
+    '''
+    name: str
+    logG: float
+
+def read_multi_atmos(filename: str) -> Tuple[MultiMetadata, Atmosphere]:
+    '''
+    Load a MULTI atmosphere definition from a file for use in Lightweaver.
+
+    Parameters
+    ----------
+    filename : str
+        The path to load the file from.
+
+    Returns
+    -------
+    meta : MultiMetadata
+        Additional metadata from the MULTI metadata that doesn't appear in a
+        Lightweaver atmosphere (name, gravitational acceleration used).
+    atmos : Atmosphere
+        The atmosphere loaded into the standard Lightweaver format.
+
+    Raises
+    ------
+    ValueError
+        if file isn't found, or cannot be parsed correctly.
+    '''
+    try:
+        with open(filename, 'r') as f:
+            lines = f.readlines()
+    except FileNotFoundError:
+        raise ValueError('Atmosphere file not found (%s)' % filename)
+
+    def get_line(commentPattern=r'^\s*\*'):
+        while len(lines) > 0:
+            line = lines.pop(0)
+            if not re.match(commentPattern, line):
+                return line.strip()
+        return None
+
+    atmosName = get_line()
+
+    scaleStr = get_line()
+    logG = float(get_line()) - 2 # For conversion to log[m.s^-2]
+    Nspace = int(get_line())
+
+    dscale = np.zeros(Nspace)
+    temp = np.zeros(Nspace)
+    ne = np.zeros(Nspace)
+    vlos = np.zeros(Nspace)
+    vturb = np.zeros(Nspace)
+    for k in range(Nspace):
+        vals = get_line().split()
+        vals = [float(v) for v in vals]
+        dscale[k] = vals[0]
+        temp[k] = vals[1]
+        ne[k] = vals[2]
+        vlos[k] = vals[3]
+        vturb[k] = vals[4]
+
+    scaleMode = scaleStr[0].upper()
+    if scaleMode == 'M':
+        scaleType = ScaleType.ColumnMass
+        dscale = ((10**dscale) << u.Unit('g cm-2')).to('kg m-2').value
+    elif scaleMode == 'T':
+        scaleType = ScaleType.Tau500
+        dscale = 10**dscale
+    elif scaleMode == 'H':
+        scaleType = ScaleType.Geometric
+        dscale = (dscale << u.Unit('km')).to('m').value
+    else:
+        raise ValueError('Unknown scale type: %s (expected M, T, or H)' % scaleStr)
+
+    vlos = (vlos << u.Unit('km s-1')).to('m s-1').value
+    vturb = (vturb << u.Unit('km s-1')).to('m s-1').value
+    ne = (ne << u.Unit('cm-3')).to('m-3').value
+
+    if len(lines) <= Nspace:
+        raise ValueError('Hydrogen populations not supplied!')
+
+    hPops = np.zeros((6, Nspace))
+    for k in range(Nspace):
+        vals = get_line().split()
+        vals = [float(v) for v in vals]
+        hPops[:, k] = vals
+
+    hPops = (hPops << u.Unit('cm-3')).to('m-3').value
+
+    meta = MultiMetadata(atmosName, logG)
+    atmos = Atmosphere.make_1d(scale=scaleType,
+                               depthScale=dscale,
+                               temperature=temp,
+                               vlos=vlos,
+                               vturb=vturb,
+                               ne=ne,
+                               hydrogenPops=hPops)
+
+    return (meta, atmos)

lightweaver/nr_update.py

@@ -0,0 +1,106 @@
+import numpy as np
+
+from .atomic_set import lte_pops
+from .atomic_table import PeriodicTable
+
+
+def nr_post_update(self, fdCollisionRates=True, hOnly=False,
+                   timeDependentData=None, chunkSize=5,
+                   ngUpdate=None, printUpdate=None, extraParams=None):
+    '''
+    Compute the Newton-Raphson terms for updating the electron density
+    through charge conservation. Is attached to the Context object.
+    Parameters
+    ----------
+    fdCollisionRates : bool, optional
+        Whether to use a finite difference approximation to the collisional
+        rates to find the population gradient WRT ne (default: True i.e. use
+        finite-difference, if False, collisional rates are ignored for this
+        process.)
+    hOnly : bool, optional
+        Ignore atoms other than Hydrogen (the primary electron contributor)
+        (default: False)
+    timeDependentData : dict, optional
+        The presence of this argument indicates that the time-dependent
+        formalism should be used. Should contain the keys 'dt' with a
+        floating point timestep, and 'nPrev' with a list of population
+        vectors from the start of the time integration step, in the order of
+        the active atoms. This latter term can be obtained from the previous
+        state provied by `Context.time_dep_update`.
+    chunkSize : int, optional
+        Not currently used.
+    ngUpdate : bool, optional
+        Whether to apply Ng Acceleration (default: None, to apply automatic
+        behaviour), will only accelerate if the counter on the Ng accelerator
+        has seen enough steps since the previous acceleration (set in Context
+        initialisation).
+    printUpdate : bool, optional
+        Whether to print information on the size of the update (default:
+        None, to apply automatic behaviour).
+    extraParams : dict, optional
+        Dict of extra parameters to be converted through the
+        `dict2ExtraParams` function and passed onto the C++ core.
+    Returns
+    -------
+    dPops : float
+        The maximum relative change of any of the NLTE populations in the
+        atmosphere.
+    '''
+    if self.activeAtoms[0].element != PeriodicTable[1]:
+        raise ValueError('Calling nr_post_update without Hydrogen active.')
+
+    if ngUpdate is None:
+        ngUpdate = self.conserveCharge
+
+    if printUpdate is None:
+        printUpdate = ngUpdate
+
+    atoms = self.activeAtoms[:1] if hOnly else self.activeAtoms
+    crswVal = self.crswCallback.val
+
+    if hOnly:
+        backgroundAtoms = [model for model in self.kwargs['spect'].radSet if model.element != PeriodicTable[1]]
+    else:
+        backgroundAtoms = self.kwargs['spect'].radSet.detailedAtoms + self.kwargs['spect'].radSet.passiveAtoms
+
+    backgroundNe = np.zeros_like(self.atmos.ne)
+    for atomModel in backgroundAtoms:
+        lteStages = np.array([l.stage for l in atomModel.levels])
+        atom = self.kwargs['eqPops'].atomicPops[atomModel.element]
+        backgroundNe += (lteStages[:, None] * atom.n[:, :]).sum(axis=0)
+
+    neStart = np.copy(self.atmos.ne)
+
+    dC = []
+    if fdCollisionRates:
+        for atom in atoms:
+            atom.compute_collisions(fillDiagonal=True)
+            Cprev = np.copy(atom.C)
+            pertSize = 1e-4
+            pert = neStart * pertSize
+            self.atmos.ne[:] += pert
+            nStarPrev = np.copy(atom.nStar)
+            atom.nStar[:] = lte_pops(atom.atomicModel, self.atmos.temperature,
+                                     self.atmos.ne, atom.nTotal)
+            atom.compute_collisions(fillDiagonal=True)
+            self.atmos.ne[:] = neStart
+            atom.nStar[:] = nStarPrev
+            dC.append(crswVal * (atom.C - Cprev) / pert)
+            atom.C[:] = Cprev
+
+    self._nr_post_update_impl(atoms, dC, backgroundNe,
+                              timeDependentData=timeDependentData, chunkSize=chunkSize, extraParams=extraParams)
+    self.eqPops.update_lte_atoms_Hmin_pops(self.atmos.pyAtmos, conserveCharge=False, quiet=True)
+
+    if ngUpdate:
+        update = self.rel_diff_ng_accelerate(printUpdate=printUpdate)
+    else:
+        update = self.rel_diff_pops(printUpdate=printUpdate)
+    neDiff = ((np.asarray(self.atmos.ne) - neStart)
+                / np.asarray(self.atmos.ne))
+    neDiffMaxIdx = neDiff.argmax()
+    neDiffMax = neDiff[neDiffMaxIdx]
+    update.updatedNe = True
+    update.dNeMax = neDiffMax
+    update.dNeMaxIdx = neDiffMaxIdx
+    return update