def test_chianti_ingest_levels_count(memory_session, ch_ingester,
atomic_number, ion_charge, levels_count):
ch_ingester.ingest(levels=True)
ion = Ion.as_unique(memory_session,
atomic_number=atomic_number,
ion_charge=ion_charge)
assert len(ion.levels) == levels_count
def test_gfall_ingester_ingest_lines(memory_session, gfall_ingester,
atomic_number, ion_charge,
level_index_lower, level_index_upper,
exp_wavelength, exp_gf_value):
gfall_ingester.ingest(levels=True, lines=True)
ion = Ion.as_unique(memory_session,
atomic_number=atomic_number,
ion_charge=ion_charge)
data_source = DataSource.as_unique(memory_session, short_name="ku_latest")
lower_level = memory_session.query(Level).\
filter(and_(Level.data_source==data_source,
Level.ion==ion,
Level.level_index==level_index_lower)).one()
upper_level = memory_session.query(Level). \
filter(and_(Level.data_source == data_source,
Level.ion == ion,
Level.level_index == level_index_upper)).one()
line = memory_session.query(Line).\
filter(and_(Line.data_source==data_source,
Line.lower_level==lower_level,
Line.upper_level==upper_level)).one()
wavelength = line.wavelengths[0].quantity
gf_value = line.gf_values[0].quantity
assert_quantity_allclose(wavelength, exp_wavelength)
assert_quantity_allclose(gf_value, exp_gf_value)
def test_line_quantities_query(foo_session, atomic_number, ion_charge, ds_short_name,
lower_level_index, upper_level_index,
expected_wavelength, expected_a_value, expected_gf_value):
data_source = DataSource.as_unique(foo_session, short_name=ds_short_name)
ion = Ion.as_unique(foo_session, atomic_number=atomic_number, ion_charge=ion_charge)
lower_level = foo_session.query(Level).\
filter(and_(Level.data_source==data_source,
Level.ion == ion,
Level.level_index == lower_level_index)).one()
upper_level = foo_session.query(Level). \
filter(and_(Level.data_source == data_source,
Level.ion == ion,
Level.level_index == upper_level_index)).one()
line = foo_session.query(Line).\
filter(and_(Line.data_source == data_source,
Line.lower_level == lower_level,
Line.upper_level == upper_level)).one()
wavelength = line.wavelengths[0].quantity
a_value = line.a_values[0].quantity
gf_value = line.gf_values[0].quantity
assert_quantity_allclose(wavelength, expected_wavelength)
assert_quantity_allclose(a_value, expected_a_value)
assert_quantity_allclose(gf_value, expected_gf_value)
def ingest_lines(self):
logger.info("Ingesting lines from `{}`.".format(
self.data_source.short_name))
for rdr in self.ion_readers:
atomic_number = rdr.ion.Z
ion_charge = rdr.ion.Ion - 1
ion = Ion.as_unique(self.session,
atomic_number=atomic_number,
ion_charge=ion_charge)
try:
bound_lines = rdr.bound_lines
except ChiantiIonReaderError:
logger.info("Lines not found for ion {} {}.".format(
convert_atomic_number2symbol(atomic_number), ion_charge))
continue
logger.info("Ingesting lines for {} {}.".format(
convert_atomic_number2symbol(atomic_number), ion_charge))
lvl_index2id = self.get_lvl_index2id(ion)
for index, row in bound_lines.iterrows():
# index: (lower_level_index, upper_level_index)
lower_level_index, upper_level_index = index
try:
lower_level_id = int(lvl_index2id.loc[lower_level_index])
upper_level_id = int(lvl_index2id.loc[upper_level_index])
except KeyError:
raise IngesterError(
"Levels from this source have not been found."
"You must ingest levels before transitions")
# Create a new line
line = Line(
lower_level_id=lower_level_id,
upper_level_id=upper_level_id,
data_source=self.data_source,
wavelengths=[
LineWavelength(quantity=row["wavelength"] * u.AA,
data_source=self.data_source,
medium=MEDIUM_VACUUM,
method=row["method"])
],
a_values=[
LineAValue(quantity=row["a_value"] * u.Unit("s**-1"),
data_source=self.data_source)
],
gf_values=[
LineGFValue(quantity=row["gf_value"],
data_source=self.data_source)
])
self.session.add(line)
def ingest_levels(self, levels=None):
if levels is None:
levels = self.gfall_reader.levels
# Select ions
if self.ions is not None:
levels = levels.reset_index().\
join(self.ions, how="inner",
on=["atomic_number", "ion_charge"]).\
set_index(["atomic_number", "ion_charge", "level_index"])
print("Ingesting levels from {}".format(self.data_source.short_name))
for ion_index, ion_levels in levels.groupby(level=["atomic_number", "ion_charge"]):
atomic_number, ion_charge = ion_index
ion = Ion.as_unique(self.session, atomic_number=atomic_number, ion_charge=ion_charge)
print("Ingesting levels for {} {}".format(convert_atomic_number2symbol(atomic_number), ion_charge))
for index, row in ion_levels.iterrows():
level_index = index[2] # index: (atomic_number, ion_charge, level_index)
ion.levels.append(
Level(level_index=level_index,
data_source=self.data_source,
J=row["j"],
energies=[
LevelEnergy(quantity=row["energy"]*u.Unit("cm-1"),
method=row["method"],
data_source=self.data_source)
])
)
def ingest_lines(self, lines=None):
if lines is None:
lines = self.gfall_reader.lines
# Select ions
if self.ions is not None:
lines = lines.reset_index(). \
join(self.ions, how="inner",
on=["atomic_number", "ion_charge"]). \
set_index(["atomic_number", "ion_charge", "level_index_lower", "level_index_upper"])
print("Ingesting lines from {}".format(self.data_source.short_name))
for ion_index, ion_lines in lines.groupby(
level=["atomic_number", "ion_charge"]):
atomic_number, ion_charge = ion_index
ion = Ion.as_unique(self.session,
atomic_number=atomic_number,
ion_charge=ion_charge)
print("Ingesting lines for {} {}".format(
convert_atomic_number2symbol(atomic_number), ion_charge))
lvl_index2id = self.get_lvl_index2id(ion)
for index, row in ion_lines.iterrows():
# index: (atomic_number, ion_charge, lower_level_index, upper_level_index)
lower_level_index, upper_level_index = index[2:]
try:
lower_level_id = int(lvl_index2id.loc[lower_level_index])
upper_level_id = int(lvl_index2id.loc[upper_level_index])
except KeyError:
raise IngesterError(
"Levels from this source have not been found."
"You must ingest levels before transitions")
medium = MEDIUM_VACUUM if row[
"wavelength"] <= GFALL_AIR_THRESHOLD else MEDIUM_AIR
# Create a new line
line = Line(lower_level_id=lower_level_id,
upper_level_id=upper_level_id,
data_source=self.data_source,
wavelengths=[
LineWavelength(quantity=row["wavelength"] *
u.nm,
medium=medium,
data_source=self.data_source)
],
gf_values=[
LineGFValue(quantity=row["gf"],
data_source=self.data_source)
])
self.session.add(line)
def test_level_g_hybrid_attribute_expression(foo_session, atomic_number, ion_charge, ds_short_name, level_index, exp_g):
data_source = DataSource.as_unique(foo_session, short_name=ds_short_name)
ion = Ion.as_unique(foo_session, atomic_number=atomic_number, ion_charge=ion_charge)
g = foo_session.query(Level.g). \
filter(and_(Level.data_source == data_source,
Level.ion == ion,
Level.level_index == level_index)).scalar()
assert g == exp_g
def test_chianti_ingest_e_col_count(memory_session, ch_ingester, atomic_number,
ion_charge, e_col_count):
ch_ingester.ingest(levels=True, collisions=True)
ion = Ion.as_unique(memory_session,
atomic_number=atomic_number,
ion_charge=ion_charge)
cnt = memory_session.query(ECollision).join(
ECollision.lower_level).filter(Level.ion == ion).count()
assert cnt == e_col_count
def test_chianti_ingest_lines_count(memory_session, ch_ingester, atomic_number,
ion_charge, lines_count):
ch_ingester.ingest(levels=True, lines=True)
ion = Ion.as_unique(memory_session,
atomic_number=atomic_number,
ion_charge=ion_charge)
cnt = memory_session.query(Line).join(
Line.lower_level).filter(Level.ion == ion).count()
assert cnt == lines_count
def ingest_lines(self, lines=None):
if lines is None:
lines = self.gfall_reader.lines
# Select ions
if self.ions is not None:
lines = lines.reset_index(). \
join(self.ions, how="inner",
on=["atomic_number", "ion_charge"]). \
set_index(["atomic_number", "ion_charge", "level_index_lower", "level_index_upper"])
print("Ingesting lines from {}".format(self.data_source.short_name))
for ion_index, ion_lines in lines.groupby(level=["atomic_number", "ion_charge"]):
atomic_number, ion_charge = ion_index
ion = Ion.as_unique(self.session, atomic_number=atomic_number, ion_charge=ion_charge)
print("Ingesting lines for {} {}".format(convert_atomic_number2symbol(atomic_number), ion_charge))
lvl_index2id = self.get_lvl_index2id(ion)
for index, row in ion_lines.iterrows():
# index: (atomic_number, ion_charge, lower_level_index, upper_level_index)
lower_level_index, upper_level_index = index[2:]
try:
lower_level_id = int(lvl_index2id.loc[lower_level_index])
upper_level_id = int(lvl_index2id.loc[upper_level_index])
except KeyError:
raise IngesterError("Levels from this source have not been found."
"You must ingest levels before transitions")
medium = MEDIUM_VACUUM if row["wavelength"] <= GFALL_AIR_THRESHOLD else MEDIUM_AIR
# Create a new line
line = Line(
lower_level_id=lower_level_id,
upper_level_id=upper_level_id,
data_source=self.data_source,
wavelengths=[
LineWavelength(quantity=row["wavelength"] * u.nm,
medium=medium,
data_source=self.data_source)
],
gf_values=[
LineGFValue(quantity=row["gf"],
data_source=self.data_source)
]
)
self.session.add(line)
def test_gfall_ingester_ingest_levels(memory_session, gfall_ingester, atomic_number, ion_charge, level_index,
exp_energy, exp_j, exp_method):
gfall_ingester.ingest(levels=True, lines=False)
ion = Ion.as_unique(memory_session, atomic_number=atomic_number, ion_charge=ion_charge)
data_source = DataSource.as_unique(memory_session, short_name="ku_latest")
level, energy = memory_session.query(Level, LevelEnergy).\
filter(and_(Level.ion==ion,
Level.level_index==level_index),
Level.data_source==data_source).\
join(Level.energies).one()
assert_almost_equal(level.J, exp_j)
assert_quantity_allclose(energy.quantity, exp_energy.to(u.eV, equivalencies=u.spectral()))
assert energy.method == exp_method
def test_ionization_energies_query(foo_session, atomic_number, ion_charge, ds_short_name,
level_index, method, expected_level_energy_value):
data_source = DataSource.as_unique(foo_session, short_name=ds_short_name)
ion = Ion.as_unique(foo_session, atomic_number=atomic_number, ion_charge=ion_charge)
level, level_energy_value = foo_session.query(Level,
LevelEnergy.quantity.to(u.Unit("cm-1"), equivalencies=u.spectral()).value). \
filter(and_(Level.ion == ion,
Level.data_source == data_source),
Level.level_index == level_index). \
join(Level.energies). \
filter(and_(LevelEnergy.data_source == data_source,
LevelEnergy.method == method)).one()
assert_almost_equal(level_energy_value, expected_level_energy_value)
def test_level_query(foo_session, atomic_number, ion_charge, ds_short_name, level_index,
configuration, term, L, J, spin_multiplicity, parity):
data_source = DataSource.as_unique(foo_session, short_name=ds_short_name)
ion = Ion.as_unique(foo_session, atomic_number=atomic_number, ion_charge=ion_charge)
level = foo_session.query(Level).\
filter(and_(Level.data_source == data_source,
Level.ion == ion,
Level.level_index == level_index)).one()
assert level.configuration == configuration
assert level.term == term
assert level.L == L
assert_almost_equal(level.J, J)
assert level.spin_multiplicity == spin_multiplicity
assert level.parity == parity
def ingest_levels(self):
logger.info("Ingesting levels from `{}`.".format(
self.data_source.short_name))
for rdr in self.ion_readers:
atomic_number = rdr.ion.Z
ion_charge = rdr.ion.Ion - 1
ion = Ion.as_unique(self.session,
atomic_number=atomic_number,
ion_charge=ion_charge)
try:
bound_levels = rdr.bound_levels
except ChiantiIonReaderError:
logger.info("Levels not found for ion {} {}.".format(
convert_atomic_number2symbol(atomic_number), ion_charge))
continue
logger.info("Ingesting levels for {} {}.".format(
convert_atomic_number2symbol(atomic_number), ion_charge))
# ToDo: Determine parity from configuration
for index, row in bound_levels.iterrows():
level = Level(ion=ion,
data_source=self.data_source,
level_index=index,
configuration=row["configuration"],
term=row["term"],
L=row["L"],
J=row["J"],
spin_multiplicity=row["spin_multiplicity"])
level.energies = []
for column, method in [('energy', 'meas'),
('energy_theoretical', 'theor')]:
if row[column] != -1: # check if the value exists
level.energies.append(
LevelEnergy(quantity=row[column] * u.Unit("cm-1"),
data_source=self.data_source,
method=method), )
self.session.add(level)
def test_line_wavelength_medium(foo_session, atomic_number, ion_charge, ds_short_name,
lower_level_index, upper_level_index, expected_medium):
data_source = DataSource.as_unique(foo_session, short_name=ds_short_name)
ion = Ion.as_unique(foo_session, atomic_number=atomic_number, ion_charge=ion_charge)
lower_level = foo_session.query(Level). \
filter(and_(Level.data_source == data_source,
Level.ion == ion,
Level.level_index == lower_level_index)).one()
upper_level = foo_session.query(Level). \
filter(and_(Level.data_source == data_source,
Level.ion == ion,
Level.level_index == upper_level_index)).one()
line = foo_session.query(Line). \
filter(and_(Line.data_source == data_source,
Line.lower_level == lower_level,
Line.upper_level == upper_level)).one()
wavelength = line.wavelengths[0]
assert wavelength.medium == expected_medium
def ingest_ionization_energies(self, ioniz_energies=None):
if ioniz_energies is None:
ioniz_energies = self.parser.prepare_ioniz_energies()
print("Ingesting ionization energies from {}".format(self.data_source.short_name))
for index, row in ioniz_energies.iterrows():
atomic_number, ion_charge = index
# Query for an existing ion; create if doesn't exists
ion = Ion.as_unique(self.session,
atomic_number=atomic_number, ion_charge=ion_charge)
ion.energies = [
IonizationEnergy(ion=ion,
data_source=self.data_source,
quantity=row['ionization_energy_value'] * u.eV,
uncert=row['ionization_energy_uncert'],
method=row['ionization_energy_method'])
]
def ingest_ionization_energies(self, ioniz_energies=None):
if ioniz_energies is None:
ioniz_energies = self.parser.prepare_ioniz_energies()
print("Ingesting ionization energies from {}".format(self.data_source.short_name))
for index, row in ioniz_energies.iterrows():
atomic_number, ion_charge = index
# Query for an existing ion; create if doesn't exists
ion = Ion.as_unique(self.session,
atomic_number=atomic_number, ion_charge=ion_charge)
ion.energies = [
IonizationEnergy(ion=ion,
data_source=self.data_source,
quantity=row['ionization_energy_value'] * u.eV,
uncert=row['ionization_energy_uncert'],
method=row['ionization_energy_method'])
]
def ingest_levels(self, levels=None):
if levels is None:
levels = self.gfall_reader.levels
# Select ions
if self.ions is not None:
levels = levels.reset_index().\
join(self.ions, how="inner",
on=["atomic_number", "ion_charge"]).\
set_index(["atomic_number", "ion_charge", "level_index"])
logger.info("Ingesting levels from `{}`.".format(
self.data_source.short_name))
for ion_index, ion_levels in levels.groupby(
level=["atomic_number", "ion_charge"]):
atomic_number, ion_charge = ion_index
ion = Ion.as_unique(self.session,
atomic_number=atomic_number,
ion_charge=ion_charge)
logger.info("Ingesting levels for {} {}.".format(
convert_atomic_number2symbol(atomic_number), ion_charge))
for index, row in ion_levels.iterrows():
# index: (atomic_number, ion_charge, level_index)
level_index = index[2]
ion.levels.append(
Level(level_index=level_index,
data_source=self.data_source,
J=row["j"],
energies=[
LevelEnergy(quantity=row["energy"] *
u.Unit("cm-1"),
method=row["method"],
data_source=self.data_source)
]))
def test_gfall_ingester_ingest_lines_wavelength_medium(memory_session, gfall_ingester, atomic_number, ion_charge,
level_index_lower, level_index_upper,
exp_wavelength, exp_medium):
gfall_ingester.ingest(levels=True, lines=True)
ion = Ion.as_unique(memory_session, atomic_number=atomic_number, ion_charge=ion_charge)
data_source = DataSource.as_unique(memory_session, short_name="ku_latest")
lower_level = memory_session.query(Level). \
filter(and_(Level.data_source == data_source,
Level.ion == ion,
Level.level_index == level_index_lower)).one()
upper_level = memory_session.query(Level). \
filter(and_(Level.data_source == data_source,
Level.ion == ion,
Level.level_index == level_index_upper)).one()
line = memory_session.query(Line). \
filter(and_(Line.data_source == data_source,
Line.lower_level == lower_level,
Line.upper_level == upper_level)).one()
wavelength = line.wavelengths[0]
assert_quantity_allclose(wavelength.quantity, exp_wavelength)
assert wavelength.medium == exp_medium
def ingest_ground_levels(self, ground_levels=None):
if ground_levels is None:
ground_levels = self.parser.prepare_ground_levels()
print("Ingesting ground levels from {}".format(
self.data_source.short_name))
for index, row in ground_levels.iterrows():
atomic_number, ion_charge = index
# Replace nan with None
row = row.where(pd.notnull(row), None)
ion = Ion.as_unique(self.session,
atomic_number=atomic_number,
ion_charge=ion_charge)
try:
spin_multiplicity = int(row["spin_multiplicity"])
except TypeError: # Raised when the variable is None
spin_multiplicity = None
try:
parity = int(row["parity"])
except TypeError: # Raised when the variable is None
parity = None
ion.levels.append(
Level(data_source=self.data_source,
configuration=row["configuration"],
term=row["term"],
L=row["L"],
spin_multiplicity=spin_multiplicity,
parity=parity,
J=row["J"],
energies=[
LevelEnergy(quantity=0, data_source=self.data_source)
]))
def ingest_ground_levels(self, ground_levels=None):
if ground_levels is None:
ground_levels = self.parser.prepare_ground_levels()
print("Ingesting ground levels from {}".format(self.data_source.short_name))
for index, row in ground_levels.iterrows():
atomic_number, ion_charge = index
# Replace nan with None
row = row.where(pd.notnull(row), None)
ion = Ion.as_unique(self.session, atomic_number=atomic_number, ion_charge=ion_charge)
try:
spin_multiplicity = int(row["spin_multiplicity"])
except TypeError: # Raised when the variable is None
spin_multiplicity = None
try:
parity = int(row["parity"])
except TypeError: # Raised when the variable is None
parity = None
ion.levels.append(
Level(data_source=self.data_source,
configuration=row["configuration"],
term=row["term"],
L=row["L"],
spin_multiplicity=spin_multiplicity,
parity=parity,
J=row["J"],
energies=[
LevelEnergy(quantity=0, data_source=self.data_source)
])
)
def test_e_collision_quantities_query(foo_session, atomic_number, ion_charge, ds_short_name,
lower_level_index, upper_level_index,
expected_energy, expected_temp_strengths):
data_source = DataSource.as_unique(foo_session, short_name=ds_short_name)
ion = Ion.as_unique(foo_session, atomic_number=atomic_number, ion_charge=ion_charge)
lower_level = foo_session.query(Level). \
filter(and_(Level.data_source == data_source,
Level.ion == ion,
Level.level_index == lower_level_index)).one()
upper_level = foo_session.query(Level). \
filter(and_(Level.data_source == data_source,
Level.ion == ion,
Level.level_index == upper_level_index)).one()
e_col = foo_session.query(ECollision). \
filter(and_(ECollision.data_source == data_source,
ECollision.lower_level == lower_level,
ECollision.upper_level == upper_level)).one()
energy = e_col.energies[0].quantity
assert_quantity_allclose(energy, expected_energy)
for temp_strength, expected_temp_strength in zip(e_col.temp_strengths_tuple, expected_temp_strengths):
assert_allclose(temp_strength, expected_temp_strength)
def ingest_collisions(self):
logger.info("Ingesting collisions from `{}`.".format(
self.data_source.short_name))
for rdr in self.ion_readers:
atomic_number = rdr.ion.Z
ion_charge = rdr.ion.Ion - 1
ion = Ion.as_unique(self.session,
atomic_number=atomic_number,
ion_charge=ion_charge)
try:
bound_collisions = rdr.bound_collisions
except ChiantiIonReaderError:
logger.info("Collisions not found for ion {} {}.".format(
convert_atomic_number2symbol(atomic_number), ion_charge))
continue
logger.info("Ingesting collisions for {} {}.".format(
convert_atomic_number2symbol(atomic_number), ion_charge))
lvl_index2id = self.get_lvl_index2id(ion)
for index, row in bound_collisions.iterrows():
# index: (lower_level_index, upper_level_index)
lower_level_index, upper_level_index = index
try:
lower_level_id = int(lvl_index2id.loc[lower_level_index])
upper_level_id = int(lvl_index2id.loc[upper_level_index])
except KeyError:
raise IngesterError(
"Levels from this source have not been found."
"You must ingest levels before transitions")
# Create a new electron collision
e_col = ECollision(
lower_level_id=lower_level_id,
upper_level_id=upper_level_id,
data_source=self.data_source,
bt92_ttype=row["ttype"],
bt92_cups=row["cups"],
energies=[
ECollisionEnergy(quantity=row["energy"] * u.rydberg,
data_source=self.data_source)
],
gf_values=[
ECollisionGFValue(quantity=row["gf_value"],
data_source=self.data_source)
])
e_col.temp_strengths = [
ECollisionTempStrength(temp=temp, strength=strength)
for temp, strength in zip(row["temperatures"],
row["collision_strengths"])
]
self.session.add(e_col)
def test_ion_as_unique(foo_session, atomic_number, ion_charge):
ion = foo_session.query(Ion).get((atomic_number, ion_charge))
ion2 = Ion.as_unique(foo_session, atomic_number=atomic_number, ion_charge=ion_charge)
assert ion is ion2
