import pickle
from dataclasses import dataclass
from typing import TYPE_CHECKING, List, Tuple, Union
from xdrlib import Unpacker
import numpy as np
import lightweaver.constants as Const
from .utils import get_data_path
if TYPE_CHECKING:
from .atmosphere import Atmosphere
[docs]
class Element:
'''
A simple value comparable description of an element (just proton number
Z), that can be quickly and easily compared, whilst allowing access to
things from the periodic table.
'''
def __init__(self, Z: int):
self.Z = Z
def __hash__(self):
return hash(self.Z)
def __eq__(self, other):
return self.Z == other.Z
def __lt__(self, other):
if isinstance(other, Isotope):
return self.Z <= other.Z
return self.Z < other.Z
def __repr__(self):
return 'Element(Z=%d)' % self.Z
def __str__(self):
return '%s(Z=%d)' % (self.name, self.Z)
@property
def mass(self):
'''
Returns the mass of the element in AMU.
'''
return PeriodicTable.massData[self.Z]
@property
def name(self):
'''
Returns the name of the element as a string.
'''
return PeriodicTable.nameMapping[self.Z]
[docs]
class Isotope(Element):
'''
A simple value comparable isotope description, inheriting from Element.
'''
def __init__(self, N, Z):
super().__init__(Z)
self.N = N
def __hash__(self):
return hash((self.N, self.Z))
def __eq__(self, other):
if type(other) is Isotope:
return self.N == other.N and self.Z == other.Z
return False
def __lt__(self, other):
if isinstance(other, Isotope):
return (self.Z, self.N) < (other.Z, other.N)
elif type(other) is Element:
return self.Z < other.Z
return False
def __repr__(self):
return 'Isotope(N=%d, Z=%d)' % (self.N, self.Z)
def __str__(self):
return '%s(N=%d, Z=%d)' % (self.element_name, self.N, self.Z)
@property
def mass(self):
'''
Returns the mass of the isotope in AMU.
'''
return PeriodicTable.massData[(self.N, self.Z)]
@property
def name(self):
'''
Returns the name of the isotope as a string.
'''
# NOTE(cmo): Handle H isotopes
if self.Z == 1 and self.N != 1:
return PeriodicTable.nameMapping[(self.N, self.Z)]
baseName = self.element_name
return '^%d_%s' % (self.N, baseName)
@property
def element(self):
'''
Returns the underlying Element of which this isotope is a family
member.
'''
return PeriodicTable[self.Z]
@property
def element_mass(self):
'''
Returns the average mass of the element.
'''
return super().mass
@property
def element_name(self):
'''
Returns the name of the Element as a string.
'''
return super().name
[docs]
def load_periodic_table_data():
'''
Internal use function to load data from the AtomicMassesNames.pickle data
file.
'''
path = get_data_path() + 'AtomicMassesNames.pickle'
with open(path, 'rb') as pkl:
massData, nameMapping = pickle.load(pkl)
# NOTE(cmo): Manually change names to D and T for H isotopes
nameMapping[(2, 1)] = 'D'
nameMapping['D'] = (2, 1)
nameMapping[(3, 1)] = 'T'
nameMapping['T'] = (3, 1)
# NOTE(cmo): Compute list of isotopes for each eleemnt.
# This can definitely be done more efficiently, but it probably doesn't
# matter. (~5 ms, done once, on startup, from timeit)
isotopes = {}
elements = {}
for target in massData:
if isinstance(target, tuple):
continue
element = Element(Z=target)
elements[target] = element
isotopes[element] = []
for key in massData:
if isinstance(key, int):
continue
if key[1] == target:
N, Z = key[0], element.Z
iso = Isotope(N=N, Z=Z)
isotopes[element].append(iso)
elements[(N, Z)] = iso
return massData, nameMapping, isotopes, elements
[docs]
def normalise_atom_name(n: str) -> str:
'''
Normalises Element names to be two characters long with an uppercase
first letter, and lower case second, or a space in the case of single
character names.
Parameters
----------
n : str
The name to normalise.
Returns
-------
result : str
The normalised name.
'''
strlen = len(n)
if strlen > 2 or strlen == 0:
raise ValueError('%s does not represent valid Element name' % n)
elif strlen == 1:
return n[0].upper()
else:
return n[0].upper() + n[1].lower()
[docs]
class PeriodicTableData:
'''
Container and accessor for the periodic table data. Not intended to be
instantiated by users, instead use the pre-instantiated PeriodicTable
instance.
'''
massData, nameMapping, isos, elems = load_periodic_table_data()
def __getitem__(self, x: Union[str, int, Tuple[int, int], Element]) -> Element:
'''
Allows access to the associated Element or Isotope via a variety of means.
If input is
- an Element or Isotope, then this function returns it.
- an int, then this function returns the element with associated
proton number.
- a tuple of ints, then this function returns the element with
associated (Z, N)
- a str starting with '^' then the str is parsed as an isotope in
the form ^N_AtomName' and the associated Isotope is returned.
- any other str, then the str is parsed as a one or two character
atom identifier (e.g. H or Ca) and the associated Element is
returned.
'''
if isinstance(x, Element):
return x
if isinstance(x, int):
try:
return self.elems[x]
except KeyError:
raise KeyError('Unable to find Element with Z=%d' % x)
if isinstance(x, tuple) and all(isinstance(y, int) for y in x):
try:
return self.elems[x]
except KeyError:
raise KeyError('Unable to find Isotope with (N=%d, Z=%d)' % x)
if isinstance(x, str):
x = x.strip()
if x.startswith('^'):
terms = x[1:].split('_')
if len(terms) != 2:
raise KeyError('Unable to parse Isotope string %s' % x)
N = int(terms[0])
name = normalise_atom_name(terms[1])
try:
Z = self.nameMapping[name]
return self.elems[(N, Z)]
except KeyError:
raise KeyError('Unable to find Isotope from string %s' % x)
else:
name = normalise_atom_name(x)
try:
Z = self.nameMapping[name]
return self.elems[Z]
except KeyError:
raise KeyError('Unable to find Element with name %s' % name)
raise KeyError('Cannot find element from %s' % repr(x))
[docs]
@classmethod
def get_isotopes(cls, e: Element) -> List[Isotope]:
'''
Get all isotopes associated with a certain Element.
'''
if not isinstance(e, Element):
raise ValueError('Requires Element as first argument, got %s' % repr(e))
if isinstance(e, Isotope):
return cls.isos[e.element]
return cls.isos[e]
@property
def elements(self):
'''
Return a sorted list of Elements by proton number Z.
'''
return sorted([e for _, e in self.elems.items() if type(e) is Element])
@property
def isotopes(self):
'''
Return a sorted list of Isotopes by proton number Z.
'''
return sorted([e for _, e in self.elems.items() if type(e) is Isotope])
@property
def nuclides(self):
'''
Return a list of all nuclides (Elements and Isotopes).
'''
return self.elements + self.isotopes
PeriodicTable = PeriodicTableData()
[docs]
class AtomicAbundance:
'''
Container and accessor for atomic abundance data. This can be
instantiated with a subset of atomic abundances, which will be used in
conjunction with the the default values for non-specified elements.
Parameters
----------
abundanceData : dict, optional
Contains the abundance data to override. For elements this should be
dict[Element] = abundance, and for isotopes dict[Isotope] = isotope
fraction.
abundDex : bool, optional
Whether the supplied abundance is in dex (with Hydrogen abundance of
12.0) or in relative Hydrogen abundance (default: True i.e. in dex).
metallicity : float, optional
Enhance the metallic abundance by a factor of 10**metallicity,
(default: 0.0).
'''
def __init__(self, abundanceData: dict=None, abundDex=True, metallicity: float=0.0):
self.abundance = self.load_default_abundance_data()
# NOTE(cmo): Default abundances always in dex
self.dex_to_decimal(self.abundance)
if abundanceData is not None:
if abundDex:
self.dex_to_decimal(abundanceData)
self.abundance.update(abundanceData)
self.metallicity = metallicity
if metallicity != 0.0:
self.apply_metallicity(self.abundance, metallicity)
self.isotopeProportions = {iso: v for iso, v in self.abundance.items()
if type(iso) is Isotope}
self.convert_isotopes_to_abundances()
self.compute_stats()
[docs]
def convert_isotopes_to_abundances(self):
'''
Converts the isotope fractions to relative Hydrogen abundance.
'''
for e in PeriodicTable.elements:
totalProp = 0.0
isos = PeriodicTable.get_isotopes(e)
for iso in isos:
totalProp += self.abundance[iso]
for iso in isos:
if totalProp != 0.0:
self.abundance[iso] /= totalProp
self.abundance[iso] *= self.abundance[e]
[docs]
def compute_stats(self):
'''
Compute the total abundance (totalAbundance), mass per H atom
(massPerH), and average mass per Element (avgMass).
'''
totalAbund = 0.0
avgMass = 0.0
for e in PeriodicTable.elements:
totalAbund += self.abundance[e]
avgMass += self.abundance[e] * e.mass
self.totalAbundance = totalAbund
self.massPerH = avgMass
self.avgMass = avgMass / totalAbund
def __getitem__(self, x: Union[str, int, Tuple[int, int], Element]) -> float:
'''
Returns the abundance of the requested Element or Isotope. All forms
of describing these are accepted as the PeriodicTable is invoked.
'''
return self.abundance[PeriodicTable[x]]
[docs]
def get_primary_isotope(self, x: Element) -> Isotope:
'''
Returns the Isotope with the highest abundance of a particular
Element.
'''
isos = PeriodicTable.get_isotopes(x)
maxIso = isos[0]
maxAbund = self[maxIso]
for iso in isos[1:]:
if (abund := self[iso]) > maxAbund:
maxAbund = abund
maxIso = iso
return iso
[docs]
@staticmethod
def dex_to_decimal(abunds):
'''
Used to convert from absolute abundance in dex to relative fractional
abundance.
'''
for e, v in abunds.items():
if type(e) is Element:
abunds[e] = 10**(v - 12.0)
[docs]
@staticmethod
def load_default_abundance_data() -> dict:
"""
Load the default abundances and convert to required format for
AtomicAbundance class (i.e. dict where dict[Element] = abundance, and
dict[Isotope] = isotope fraction).
"""
with open(get_data_path() + 'AbundancesAsplund09.pickle', 'rb') as pkl:
abundances = pickle.load(pkl)
lwAbundances = {}
for ele in abundances:
Z = ele['elem']['elem']['Z']
abund = ele['elem']['abundance']
lwAbundances[PeriodicTable[Z]] = abund
for iso in ele['isotopes']:
N = iso['N']
prop = iso['proportion']
lwAbundances[PeriodicTable[(N, Z)]] = prop
for e in PeriodicTable.nuclides:
if e not in lwAbundances:
lwAbundances[e] = 0.0
return lwAbundances
DefaultAtomicAbundance = AtomicAbundance()
[docs]
@dataclass
class KuruczPf:
'''
Storage and functions relating to Bob Kurucz's partition functions.
Based on the data used in RH.
Attributes
----------
element : Element
The element associated with this partition funciton.
abundance : float
The abundance of this element.
Tpf : np.ndarray
The temperature grid on which the partition function is defined.
pf : np.ndarray
The partition function data.
ionPot : np.ndarray
The ionisation potential of each level.
'''
element: Element
abundance: float
Tpf: np.ndarray
pf: np.ndarray
ionPot: np.ndarray
[docs]
def lte_ionisation(self, atmos: 'Atmosphere') -> np.ndarray:
'''
Compute the population of the species in each ionisation
stage in a given atmosphere.
Parameters
----------
atmos : Atmosphere
The atmosphere in which to compute the populations.
Returns
-------
pops : np.ndarray
The LTE ionisation populations [Nstage x Nspace].
'''
Nstage = self.ionPot.shape[0]
Nspace = atmos.Nspace
C1 = (Const.HPlanck / (2.0 * np.pi * Const.MElectron)) * Const.HPlanck / Const.KBoltzmann
CtNe = 2.0 * (C1/atmos.temperature)**(-1.5) / atmos.ne
total = np.ones(Nspace)
pops = np.zeros((Nstage, Nspace))
pops[0, :] = 1.0
Uk = np.interp(atmos.temperature, self.Tpf, self.pf[0, :])
for i in range(1, Nstage):
Ukp1 = np.interp(atmos.temperature, self.Tpf, self.pf[i, :])
pops[i, :] = pops[i-1, :] * CtNe * np.exp(Ukp1 - Uk - self.ionPot[i-1] \
/ (Const.KBoltzmann * atmos.temperature))
total += pops[i]
Ukp1, Uk = Uk, Ukp1
pops[0, :] = self.abundance * atmos.nHTot / total
pops[1:,:] *= pops[0, :]
return pops
[docs]
def fjk(self, atmos: 'Atmosphere', k: int) -> Tuple[np.ndarray, np.ndarray]:
'''
Compute the fractional population of the species in each ionisation
stage and partial derivative wrt n_e at one point in a given
atmosphere.
Parameters
----------
atmos : Atmosphere
The atmosphere in which to compute the populations.
k : int
The spatial index at which to compute the populations
Returns
-------
fj : np.ndarray
The fractional populations [Nstage].
dfj : np.ndarray
The derivatives of the fractional populations [Nstage].
'''
Nspace: int = atmos.Nspace
T: float = atmos.temperature[k]
ne: float = atmos.ne[k]
C1 = ((Const.HPlanck / (2.0 * np.pi * Const.MElectron))
* Const.HPlanck / Const.KBoltzmann)
CtNe = 2.0 * (C1/T)**(-1.5) / ne
Nstage: int = self.ionPot.shape[0]
fjk = np.zeros(Nstage)
fjk[0] = 1.0
dfjk = np.zeros(Nstage)
# fjk: fractional population of stage j, at atmospheric index k
# The first stage starts with a "population" of 1, then via Saha we
# compute the relative populations of the other stages, before dividing
# by the sum across these
Uk: float = np.interp(T, self.Tpf, self.pf[0, :])
for j in range(1, Nstage):
Ukp1: float = np.interp(T, self.Tpf, self.pf[j, :])
fjk[j] = fjk[j-1] * CtNe * np.exp(Ukp1 - Uk - self.ionPot[j-1]
/ (Const.KBoltzmann * T))
dfjk[j] = -j * fjk[j] / ne
Uk = Ukp1
sumF = np.sum(fjk)
sumDf = np.sum(dfjk)
fjk /= sumF
dfjk = (dfjk - fjk * sumDf) / sumF
return fjk, dfjk
[docs]
def fj(self, atmos: 'Atmosphere') -> Tuple[np.ndarray, np.ndarray]:
'''
Compute the fractional population of the species in each ionisation
stage and partial derivative wrt n_e for each location in a given
atmosphere.
Parameters
----------
atmos : Atmosphere
The atmosphere in which to compute the populations.
Returns
-------
fj : np.ndarray
The fractional populations [Nstage x Nspace].
dfj : np.ndarray
The derivatives of the fractional populations [Nstage x Nspace].
'''
Nspace: int = atmos.Nspace
T = atmos.temperature
ne = atmos.ne
C1 = ((Const.HPlanck / (2.0 * np.pi * Const.MElectron))
* Const.HPlanck / Const.KBoltzmann)
CtNe = 2.0 * (C1/T)**(-1.5) / ne
Nstage: int = self.ionPot.shape[0]
fj = np.zeros((Nstage, Nspace))
fj[0, :] = 1.0
dfj = np.zeros((Nstage, Nspace))
# fjk: fractional population of stage j, at atmospheric index k
# The first stage starts with a "population" of 1, then via Saha we
# compute the relative populations of the other stages, before dividing
# by the sum across these
Uk = np.interp(T, self.Tpf, self.pf[0, :])
for j in range(1, Nstage):
Ukp1 = np.interp(T, self.Tpf, self.pf[j, 0])
fj[j] = fj[j-1] * CtNe * np.exp(Ukp1 - Uk - self.ionPot[j-1]
/ (Const.KBoltzmann * T))
dfj[j] = -j * fj[j] / ne
Uk[:] = Ukp1
sumF = np.sum(fj, axis=0)
sumDf = np.sum(dfj, axis=0)
fj /= sumF
dfj = (dfj - fj * sumDf) / sumF
return fj, dfj
[docs]
class KuruczPfTable:
'''
Container for all of the Kurucz partition function data, allowing
different paths and AtomicAbundances to be used. Serves to construct the
KuruczPf objects if used.
Parameters
----------
atomicAbundance : AtomicAbundance, optional
The abundance data to use, if non-standard.
kuruczPfPath : str
The path to the Kurucz parition function data in RH's XDR format, if
non-standard.
'''
def __init__(self, atomicAbundance: AtomicAbundance=None, kuruczPfPath: str=None):
if atomicAbundance is None:
atomicAbundance = DefaultAtomicAbundance
self.atomicAbundance = atomicAbundance
kuruczPfPath = get_data_path() + 'pf_Kurucz.input' if kuruczPfPath is None \
else kuruczPfPath
with open(kuruczPfPath, 'rb') as f:
s = f.read()
u = Unpacker(s)
# NOTE(cmo): Each of these terms is simply in flat lists indexed by Atomic Number Z-1
self.Tpf = np.array(u.unpack_array(u.unpack_double))
stages = []
pf = []
ionpot = []
for i in range(99):
z = u.unpack_int()
stages.append(u.unpack_int())
pf.append(np.array(u.unpack_farray(stages[-1] * self.Tpf.shape[0],
u.unpack_double)).reshape(stages[-1],
self.Tpf.shape[0]))
ionpot.append(np.array(u.unpack_farray(stages[-1], u.unpack_double)))
ionpot = [i * Const.HC / Const.CM_TO_M for i in ionpot]
pf = [np.log(p) for p in pf]
self.pf = pf
self.ionpot = ionpot
def __getitem__(self, x: Element) -> KuruczPf:
'''
Used to construct the partition function object for the requested
element.
'''
if type(x) is Isotope:
raise ValueError('Isotopes not supported by KuruczPf')
zm = x.Z - 1
return KuruczPf(element=x, abundance=self.atomicAbundance[x],
Tpf=self.Tpf, pf=self.pf[zm], ionPot=self.ionpot[zm])