import pickle
from dataclasses import dataclass
from typing import TYPE_CHECKING, List, Tuple, Union
from .deprecation import deprecated_names
try:
from xdrlib import Unpacker
except ImportError:
from mda_xdrlib.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.mass_data[self.Z]
@property
def name(self):
"""
Returns the name of the element as a string.
"""
return PeriodicTable.name_mapping[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.mass_data[(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.name_mapping[(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:
mass_data, name_mapping = pickle.load(pkl)
# NOTE(cmo): Manually change names to D and T for H isotopes
name_mapping[(2, 1)] = 'D'
name_mapping['D'] = (2, 1)
name_mapping[(3, 1)] = 'T'
name_mapping['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 mass_data:
if isinstance(target, tuple):
continue
element = Element(Z=target)
elements[target] = element
isotopes[element] = []
for key in mass_data:
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 mass_data, name_mapping, 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]
@deprecated_names(attrs=('mass_data', 'name_mapping'))
class PeriodicTableData:
"""
Container and accessor for the periodic table data. Not intended to be
instantiated by users, instead use the pre-instantiated PeriodicTable
instance.
"""
mass_data, name_mapping, 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.name_mapping[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.name_mapping[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]
@deprecated_names(attrs=('avg_mass', 'isotope_proportions', 'mass_per_h', 'total_abundance'))
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
----------
abundance_data : dict, optional
Contains the abundance data to override. For elements this should be
dict[Element] = abundance, and for isotopes dict[Isotope] = isotope
fraction.
abund_dex : 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, abundance_data: dict = None, abund_dex=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 abundance_data is not None:
if abund_dex:
self.dex_to_decimal(abundance_data)
self.abundance.update(abundance_data)
self.metallicity = metallicity
if metallicity != 0.0:
self.apply_metallicity(self.abundance, metallicity)
self.isotope_proportions = {
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 (total_abundance), mass per H atom
(mass_per_h), and average mass per Element (avg_mass).
"""
totalAbund = 0.0
avg_mass = 0.0
for e in PeriodicTable.elements:
totalAbund += self.abundance[e]
avg_mass += self.abundance[e] * e.mass
self.total_abundance = totalAbund
self.mass_per_h = avg_mass
self.avg_mass = avg_mass / 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 maxIso
[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]
@deprecated_names
@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.
ion_pot : np.ndarray
The ionisation potential of each level.
"""
element: Element
abundance: float
Tpf: np.ndarray
pf: np.ndarray
ion_pot: 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.ion_pot.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.ion_pot[i - 1] / (Const.KBoltzmann * atmos.temperature))
)
total += pops[i]
Ukp1, Uk = Uk, Ukp1
pops[0, :] = self.abundance * atmos.nh_tot / 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].
"""
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.ion_pot.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.ion_pot[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.ion_pot.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, :])
fj[j] = (
fj[j - 1] * CtNe * np.exp(Ukp1 - Uk - self.ion_pot[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]
@deprecated_names(attrs=('atomic_abundance',))
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
----------
atomic_abundance : AtomicAbundance, optional
The abundance data to use, if non-standard.
kurucz_pf_path : str
The path to the Kurucz parition function data in RH's XDR format, if
non-standard.
"""
def __init__(self, atomic_abundance: AtomicAbundance = None, kurucz_pf_path: str = None):
if atomic_abundance is None:
atomic_abundance = DefaultAtomicAbundance
self.atomic_abundance = atomic_abundance
kurucz_pf_path = (
get_data_path() + 'pf_Kurucz.input' if kurucz_pf_path is None else kurucz_pf_path
)
with open(kurucz_pf_path, '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):
u.unpack_int() # Atomic number, implied by position.
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_CM 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.atomic_abundance[x],
Tpf=self.Tpf,
pf=self.pf[zm],
ion_pot=self.ionpot[zm],
)