# -*- coding: utf-8 -*-
""" Handles reading and interpolation of atmopshere (grid) data """
import logging
import numpy as np
from flex.extensions.bindata import MultipleDataExtension
from ..abund import Abund
from ..data_structure import (
Collection,
CollectionFactory,
absolute,
array,
asfloat,
asstr,
lowercase,
oneof,
this,
uppercase,
)
logger = logging.getLogger(__name__)
[docs]class AtmosphereError(RuntimeError):
"""Something went wrong with the atmosphere interpolation"""
[docs]@CollectionFactory
class Atmosphere(Collection):
"""
Atmosphere structure
contains all information to describe the solar atmosphere
i.e. temperature etc in the different layers
as well as stellar parameters and abundances
"""
# fmt: off
_fields = [
("teff", 5770, asfloat, this, "float: effective temperature in Kelvin"),
("logg", 4.0, asfloat, this, "float: surface gravity in log10(cgs)"),
("abund", Abund.solar(), this, this, "Abund: elemental abundances"),
("vturb", 0, absolute, this, "float: turbulence velocity in km/s"),
("lonh", 0, asfloat, this, "float: ?"),
("source", "marcs2012.sav", this, this, "str: datafile name of this data, or atmosphere grid/file"),
("method", "grid", lowercase(oneof("grid", "embedded")), this,
"str: whether the data source is a grid or a fixed atmosphere"),
("geom", "PP", uppercase(oneof("PP", "SPH", None)), this,
"str: the geometry of the atmopshere model"),
("radius", 0, asfloat, this, "float: radius of the spherical model"),
("height", None, array(None, "f8"), this, "array: height of the spherical model"),
("opflag", [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], array(20, int), this,
"array of size (20,): opacity flags"),
("wlstd", 5000, asfloat, this, "float: wavelength standard deviation"),
("depth", None, uppercase(oneof(None, "RHOX", "TAU")), this,
"str: flag that determines whether to use RHOX or TAU for calculations"),
("interp", None, uppercase(oneof(None, "RHOX", "TAU")), this,
"str: flag that determines whether to use RHOX or TAU for interpolation"),
("rhox", None, array(None, "f8"), this,
"array: mass column density"),
("tau", None, array(None, "f8"), this,
"array: continuum optical depth"),
("temp", None, array(None, "f8"), this,
"array: temperature profile in Kelvin"),
("rho", None, array(None, "f8"), this,
"array: density profile"),
("xna", None, array(None, "f8"), this,
"array: number density of atoms in 1/cm**3"),
("xne", None, array(None, "f8"), this,
"array: number density of electrons in 1/cm**3"),
("citation_info", "", asstr, this, "str: citation for this atmosphere"),
]
# fmt: on
# TODO: pick the right geometry for the grid, based on whether it has height/radius values or not
def __init__(self, *args, **kwargs):
monh = kwargs.pop("monh", kwargs.pop("feh", 0))
abund = kwargs.pop("abund", "empty")
abund_format = kwargs.pop("abund_format", "sme")
if "geom" in kwargs.keys() and kwargs["geom"] == "":
kwargs["geom"] = None
super().__init__(**kwargs)
self.abund = Abund(monh=monh, pattern=abund, type=abund_format)
@property
def monh(self):
"""float: metallicity"""
return self.abund.monh
@monh.setter
def monh(self, value):
self.abund.monh = value
@property
def names(self):
return self._names + ["monh"]
@property
def dtype(self):
obj = lambda: None
obj.names = [n.lower() for n in self.names]
return obj
@property
def ndep(self):
scale = self.temp
if scale is not None:
return scale.shape[0]
return None
@ndep.setter
def ndep(self, value):
pass
def _save(self):
data = {}
ext2 = self.abund._save()
header = ext2.header
data["abund"] = ext2.data
header["abund_format"] = header["type"]
del header["type"]
for name in self._names:
value = self[name]
if isinstance(value, np.ndarray):
data[name] = value
elif value is None or isinstance(value, Abund):
pass
else:
# if isinstance(value, (np.floating, np.integer, np.str, float, int, str)):
header[name] = value
ext = MultipleDataExtension(header, data)
return ext
@classmethod
def _load(cls, ext):
header = ext.header
header.update(ext.data)
obj = cls(**header)
return obj
[docs]@CollectionFactory
class AtmosphereGrid(np.recarray):
"""
A grid of atmospheres, used for the interpolation
of model atmospheres. Each entry represents one
atmosphere model.
"""
# fmt: off
_fields = [
("source", None, asstr, this, "str: datafile name of this data"),
("method", "grid", lowercase(oneof("grid", "embedded")), this,
"str: whether the data source is a grid or a fixed atmosphere"),
("geom", None, uppercase(oneof(None, "PP", "SPH")), this,
"str: the geometry of the atmopshere model"),
("depth", None, uppercase(oneof(None, "RHOX", "TAU")), this,
"str: flag that determines whether to use RHOX or TAU for calculations"),
("interp", None, uppercase(oneof(None, "RHOX", "TAU")), this,
"str: flag that determines whether to use RHOX or TAU for interpolation"),
("abund_format", "sme", oneof(*Abund._formats), this,
"str: format of the Abundance field, as defined by the Abund class"),
("citation_info", "", asstr, this, "str: Citation text to cite in your papers"),
]
# fmt: on
def __new__(cls, natmo, npoints, **kwargs):
dtype = [
("teff", "f4"),
("logg", "f4"),
("monh", "f4"),
("vturb", "f4"),
("lonh", "f4"),
("radius", "f4"),
("height", f"({npoints},)f4"),
("wlstd", "f4"),
("opflag", "(20,)i4"),
("abund", "(99,)f4"),
("temp", f"({npoints},)f4"),
("rhox", f"({npoints},)f4"),
("tau", f"({npoints},)f4"),
("rho", f"({npoints},)f4"),
("xna", f"({npoints},)f4"),
("xne", f"({npoints},)f4"),
]
names = [s[0].lower() for s in dtype]
titles = [s[0].upper() for s in dtype]
data = super().__new__(cls, (natmo,), dtype=dtype, names=names, titles=titles)
data.interp = "TAU"
data.depth = "RHOX"
data.method = "grid"
data.geom = "PP"
data.source = ""
data.citation_info = ""
data.abund_format = "sme"
data.info = ""
return data
def __array_finalize__(self, obj):
if obj is None:
return
self.interp = getattr(obj, "interp", "TAU")
self.depth = getattr(obj, "depth", "RHOX")
self.method = getattr(obj, "method", "grid")
self.geom = getattr(obj, "geom", "PP")
self.source = getattr(obj, "source", "")
self.citation_info = getattr(obj, "citation_info", "")
self.abund_format = getattr(obj, "abund_format", "sme")
self.info = getattr(obj, "info", "")
def __reduce__(self):
# Get the parent's __reduce__ tuple
pickled_state = super(AtmosphereGrid, self).__reduce__()
# Create our own tuple to pass to __setstate__
new_state = pickled_state[2] + (
self.interp,
self.depth,
self.source,
self.geom,
self.citation_info,
self.method,
self.abund_format,
self.info,
)
# Return a tuple that replaces the parent's __setstate__ tuple with our own
return (pickled_state[0], pickled_state[1], new_state)
def __setstate__(self, state):
self.interp = state[-8]
self.depth = state[-7]
self.source = state[-6]
self.geom = state[-5]
self.citation_info = state[-4]
self.method = state[-3]
self.abund_format = state[-2]
self.info = state[-1]
# Call the parent's __setstate__ with the other tuple elements.
super(AtmosphereGrid, self).__setstate__(state[0:-8])
def __getitem__(self, key):
"""Overwrite the getitem routine, so we keep additional
properties and/or return an atmosphere object, when only
one record is returned"""
cls = self.__class__
value = super().__getitem__(key)
if isinstance(value, cls) and value.size == 1:
return value[0]
if isinstance(value, np.record):
kwargs = {s: value[s] for s in value.dtype.names}
value = Atmosphere(**kwargs)
if isinstance(value, (Atmosphere, cls)):
for name in self._names:
setattr(value, name, getattr(self, name))
return value
def __str__(self):
return self.source
def __repr__(self):
return str(self)
[docs] def get(self, teff, logg, monh):
mask = self.teff == teff
mask &= self.logg == logg
mask &= self.monh == monh
return self[mask]
@property
def ndep(self):
return self.shape[1]
[docs] def save(self, filename):
"""Save the Atmopshere grid to a file using a numpy save format"""
header = {
"interp": self.interp,
"depth": self.depth,
"source": self.source,
"geom": self.geom,
"citation_info": self.citation_info,
"method": self.method,
"abund_format": self.abund_format,
"info": self.info,
}
names = list(header.keys())
values = [header[n] for n in names]
header = np.rec.array(values, names=names)
np.savez(filename, data=self, header=header)
return
[docs] @classmethod
def load(cls, filename):
"""Load the atmosphere grid data from disk"""
data = np.load(filename)
self = data["data"].view(cls)
header = data["header"]
for k in header.dtype.names:
v = header[k][()]
setattr(self, k, v)
return self