Skip to content

Commit

Permalink
Refactor xml methods. Fix transition methods. Add possibility of user…
Browse files Browse the repository at this point in the history 
… rate functions.
  • Loading branch information
@mbradle
mbradle committed Aug 28, 2023
1 parent bef5609 commit bbf40d7
Show file tree
Hide file tree
Showing 4 changed files with 159 additions and 114 deletions.
24 changes: 22 additions & 2 deletions .github/workflows/lvlspy_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,15 +43,17 @@ def test_probability():
p = s.compute_equilibrium_probabilities(T)
assert p[0] == 1.0


def test_transitions():
coll = get_collection()
s = coll.get()["al26"]
levels = s.get_levels()
upwards = s.get_upward_transitions_from_level(levels[0])
assert len(upwards) > 0
downwards = s.get_downward_transitions_from_level(levels[len(levels)-1])
downwards = s.get_downward_transitions_from_level(levels[len(levels) - 1])
assert len(downwards) > 0


def test_einstein():
coll = get_collection()
s = coll.get()["al26"]
Expand All @@ -73,13 +75,31 @@ def test_einstein():
assert s.get_level_to_level_transition(levs[0], levs[1]) == None


def test_rates():
def user_rate(temperature):
return 2. * temperature

coll = get_collection()
s = coll.get()["al26"]
levs = s.get_levels()
trans = s.get_level_to_level_transition(levs[1], levs[0])
assert trans.compute_lower_to_upper_rate(
1.0e9
) < trans.compute_upper_to_lower_rate(1.0e9)

assert trans.compute_lower_to_upper_rate(
1.0e9, user_func=user_rate
) == trans.compute_upper_to_lower_rate(1.0e9, user_func=user_rate)


def test_frequency():
coll = get_collection()
s = coll.get()["al26"]
levs = s.get_levels()
trans = s.get_level_to_level_transition(levs[1], levs[0])
assert trans.get_frequency() == 5.520390824072181e19


def test_write_xml():
coll = get_collection()
assert coll.write_to_xml('out.xml') == None
assert coll.write_to_xml("out.xml") == None
213 changes: 107 additions & 106 deletions lvlspy/spcoll.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,8 @@ def __init__(self, species=None):
self.properties = {}
self.spcoll = {}
if species:
for sp in species:
self.spcoll[sp.get_name()] = sp
for my_species in species:
self.spcoll[my_species.get_name()] = my_species

def add_species(self, species):
"""Method to add a species to a collection.
Expand Down Expand Up @@ -86,79 +86,72 @@ def write_to_xml(self, file, pretty_print=True, units="keV"):

self._add_optional_properties(root, self)

my_coll = self.get()
for species_name, species in self.get().items():
my_species = etree.SubElement(root, "species", name=species_name)

for sp in my_coll:
self._add_optional_properties(my_species, species)

tu_dict = {}
xml_levels = etree.SubElement(my_species, "levels")

for transition in my_coll[sp].get_transitions():
e_upper = (transition.get_upper_level()).get_energy()
for level in species.get_levels():
self._add_level_to_xml(xml_levels, species, level, units)

if e_upper not in tu_dict:
tu_dict[e_upper] = []
xml.write(file, pretty_print=pretty_print)

tu_dict[e_upper].append(transition)
def _add_level_to_xml(self, xml_levels, species, level, units):
result = etree.SubElement(xml_levels, "level")
self._add_optional_properties(result, level)
result_props = etree.SubElement(result, "properties")
if units != "keV":
my_energy = etree.SubElement(
result_props, "energy", units=units
)
else:
my_energy = etree.SubElement(result_props, "energy")
my_energy.text = self._get_energy_text(
level.get_energy(), units
)
my_multiplicity = etree.SubElement(
result_props, "multiplicity"
)
my_multiplicity.text = str(level.get_multiplicity())

my_species = etree.SubElement(root, "species", name=sp)
self._add_transitions_to_xml(result, species, level, units)

self._add_optional_properties(my_species, my_coll[sp])
return result

my_levels = etree.SubElement(my_species, "levels")
def _add_transitions_to_xml(self, xml_level, species, level, units):

levels = my_coll[sp].get_levels()
transitions = species.get_downward_transitions_from_level(level)

for level in levels:
my_level = etree.SubElement(my_levels, "level")
self._add_optional_properties(my_level, level)
my_level_props = etree.SubElement(my_level, "properties")
if units != "keV":
my_energy = etree.SubElement(
my_level_props, "energy", units=units
)
else:
my_energy = etree.SubElement(my_level_props, "energy")
my_energy.text = self._get_energy_text(
level.get_energy(), units
)
my_multiplicity = etree.SubElement(
my_level_props, "multiplicity"
)
my_multiplicity.text = str(level.get_multiplicity())

if level.get_energy() in tu_dict:
if len(transitions) == 0:
return

if len(tu_dict[level.get_energy()]) > 0:
my_transitions = etree.SubElement(
my_level, "transitions"
)
for transition in tu_dict[level.get_energy()]:
my_trans = etree.SubElement(
my_transitions, "transition"
)
self._add_optional_properties(my_trans, transition)
lower_level = transition.get_lower_level()
if units != "keV":
my_to_energy = etree.SubElement(
my_trans, "to_energy", units=units
)
else:
my_to_energy = etree.SubElement(
my_trans, "to_energy"
)
my_to_energy.text = self._get_energy_text(
lower_level.get_energy(), units
)
my_to_multiplicity = etree.SubElement(
my_trans, "to_multiplicity"
)
my_to_multiplicity.text = str(
lower_level.get_multiplicity()
)
my_a = etree.SubElement(my_trans, "a")
my_a.text = str(transition.get_einstein_a())
xml_transitions = etree.SubElement(xml_level, "transitions")

xml.write(file, pretty_print=pretty_print)
for transition in transitions:
xml_trans = etree.SubElement(xml_transitions, "transition")
self._add_optional_properties(xml_trans, transition)
lower_level = transition.get_lower_level()
if units != "keV":
xml_to_energy = etree.SubElement(
xml_trans, "to_energy", units=units
)
else:
xml_to_energy = etree.SubElement(
xml_trans, "to_energy"
)
xml_to_energy.text = self._get_energy_text(
lower_level.get_energy(), units
)
xml_to_multiplicity = etree.SubElement(
xml_trans, "to_multiplicity"
)
xml_to_multiplicity.text = str(
lower_level.get_multiplicity()
)
xml_a = etree.SubElement(xml_trans, "a")
xml_a.text = str(transition.get_einstein_a())

def _get_energy_text(self, energy, units):
return str(energy * lv.units_dict[units])
Expand Down Expand Up @@ -265,50 +258,58 @@ def update_from_xml(self, file, xpath=""):

self._update_optional_properties(spcoll, self)

el_species = spcoll.xpath("//species" + xpath)

for xml_species in el_species:
level_dict = {}
my_species = ls.Species(xml_species.attrib["name"])
self._update_optional_properties(xml_species, my_species)
el_level = xml_species.xpath(".//level")
for lev in el_level:
props = lev.xpath(".//properties")
energy = props[0].xpath(".//energy")
multiplicity = props[0].xpath(".//multiplicity")
attributes = energy[0].attrib
if "units" in attributes:
my_level = lv.Level(
float(energy[0].text),
int(multiplicity[0].text),
units=attributes["units"],
)
else:
my_level = lv.Level(
float(energy[0].text), int(multiplicity[0].text)
)
self._update_optional_properties(lev, my_level)

level_dict[my_level.get_energy()] = my_level

my_species.add_level(my_level)

el_trans = lev.xpath(".//transition")
for trans in el_trans:
to_energy = trans.xpath(".//to_energy")
to_a = trans.xpath(".//a")

f_to_energy = self._convert_to_kev(to_energy)
if f_to_energy in level_dict:
my_trans = lt.Transition(
my_level,
level_dict[f_to_energy],
float(to_a[0].text),
)
self._update_optional_properties(trans, my_trans)
my_species.add_transition(my_trans)
for xml_species in spcoll.xpath("//species" + xpath):
self.add_species(self._get_species_from_xml(xml_species))

def _get_species_from_xml(self, xml_species):
level_dict = {}
result = ls.Species(xml_species.attrib["name"])
self._update_optional_properties(xml_species, result)
for xml_level in xml_species.xpath(".//level"):
new_level = self._get_level_from_xml(xml_level)
result.add_level(new_level)
level_dict[new_level.get_energy()] = new_level

for xml_trans in xml_level.xpath(".//transition"):
trans = self._get_transition_from_xml(
xml_trans, new_level, level_dict
)
if trans:
result.add_transition(trans)

return result

def _get_level_from_xml(self, xml_level):
props = xml_level.xpath(".//properties")
energy = props[0].xpath(".//energy")
multiplicity = props[0].xpath(".//multiplicity")
attributes = energy[0].attrib
if "units" in attributes:
result = lv.Level(
float(energy[0].text),
int(multiplicity[0].text),
units=attributes["units"],
)
else:
result = lv.Level(float(energy[0].text), int(multiplicity[0].text))
self._update_optional_properties(xml_level, result)

return result

self.add_species(my_species)
def _get_transition_from_xml(self, xml_trans, upper_level, level_dict):
to_energy = xml_trans.xpath(".//to_energy")
to_a = xml_trans.xpath(".//a")

f_to_energy = self._convert_to_kev(to_energy)
if f_to_energy in level_dict:
result = lt.Transition(
upper_level,
level_dict[f_to_energy],
float(to_a[0].text),
)
self._update_optional_properties(xml_trans, result)
return result
return None

def _convert_to_kev(self, energy):
attributes = energy[0].attrib
Expand Down
4 changes: 2 additions & 2 deletions lvlspy/species.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ def get_upward_transitions_from_level(self, level):
result = []
for transition in self.get_transitions():
if transition.get_lower_level() == level:
result.append(level)
result.append(transition)

return result

Expand All @@ -134,7 +134,7 @@ def get_downward_transitions_from_level(self, level):
result = []
for transition in self.get_transitions():
if transition.get_upper_level() == level:
result.append(level)
result.append(transition)

return result

Expand Down
32 changes: 28 additions & 4 deletions lvlspy/transition.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,34 +118,58 @@ def get_einstein_b_lower_to_upper(self):
/ self.lower_level.get_multiplicity()
)

def compute_lower_to_upper_rate(self, temperature):
def compute_lower_to_upper_rate(self, temperature, user_func=None):
"""Method to compute the total rate for transition from the lower
level to upper level.
Args:
``temperature`` (:obj:`float`) The temperature in K at which to
compute the rate.
``user_func`` (optional): A `function
<https://docs.python.org/3/library/stdtypes.html#functions>`_
that computes the lower level to upper level transition rate.
If supplied, the routine will use this function in place of the
default one, which computes the rate from the appropriate
Einstein coefficient and the blackbody spectrum.
The function must take one :obj:`float` argument giving the
temperature. Other data can be bound to the function.
Returns:
:obj:`float`: The rate (per second).
"""

if user_func:
return user_func(temperature)

return self.get_einstein_b_lower_to_upper() * self._bb(temperature)

def compute_upper_to_lower_rate(self, temperature):
"""Method to compute the total rate for transition from the upper level to
to lower level.
def compute_upper_to_lower_rate(self, temperature, user_func=None):
"""Method to compute the total rate for transition from the upper
level to to lower level.
Args:
``temperature`` (:obj:`float`) The temperature in K at which to
compute the rate.
``user_func`` (optional): A `function
<https://docs.python.org/3/library/stdtypes.html#functions>`_
that computes the upper level to lower level transition rate.
If supplied, the routine will use this function in place of the
default one, which computes the rate from the appropriate
Einstein coefficients and the blackbody spectrum.
The function must take one :obj:`float` argument giving the
temperature. Other data can be bound to the function.
Returns:
:obj:`float`: The rate (per second).
"""

if user_func:
return user_func(temperature)

return (
self.get_einstein_a()
+ self.get_einstein_b_upper_to_lower() * self._bb(temperature)
Expand Down

0 comments on commit bbf40d7

Please sign in to comment.