From c9921d9d6af4459e1aeffbc9279cc5733901d780 Mon Sep 17 00:00:00 2001 From: Ayushi Daksh Date: Tue, 18 Jul 2023 08:43:58 +0530 Subject: [PATCH] Make all setters atomic operations --- tardis/analysis.py | 57 +++++++++++++++++------- tardis/visualization/widgets/grotrian.py | 25 ++++++++--- 2 files changed, 61 insertions(+), 21 deletions(-) diff --git a/tardis/analysis.py b/tardis/analysis.py index 7e2c3ce65b7..6ead15f9a43 100644 --- a/tardis/analysis.py +++ b/tardis/analysis.py @@ -10,6 +10,8 @@ import numpy as np import pandas as pd +INVALID_ION_ERROR_MSG = "Atomic number, ion_number pair not present in model" + class LastLineInteraction(object): @classmethod @@ -86,8 +88,13 @@ def atomic_number(self): @atomic_number.setter def atomic_number(self, value): - self._atomic_number = value - self.update_last_interaction_filter() + old_atomic_number = self._atomic_number + try: + self._atomic_number = value + self.update_last_interaction_filter() + except: + self._atomic_number = old_atomic_number + raise ValueError(INVALID_ION_ERROR_MSG) @property def ion_number(self): @@ -95,13 +102,25 @@ def ion_number(self): @ion_number.setter def ion_number(self, value): - self._ion_number = value - self.update_last_interaction_filter() + old_ion_number = self._ion_number + try: + self._ion_number = value + self.update_last_interaction_filter() + except: + self._ion_number = old_ion_number + raise ValueError(INVALID_ION_ERROR_MSG) def set_ion(self, atomic_number, ion_number): - self._atomic_number = atomic_number - self._ion_number = ion_number - self.update_last_interaction_filter() + old_atomic_number = self._atomic_number + old_ion_number = self._ion_number + try: + self._atomic_number = atomic_number + self._ion_number = ion_number + self.update_last_interaction_filter() + except: + self._atomic_number = old_atomic_number + self._ion_number = old_ion_number + raise ValueError(INVALID_ION_ERROR_MSG) @property def shell(self): @@ -109,8 +128,13 @@ def shell(self): @shell.setter def shell(self, value): - self._shell = value - self.update_last_interaction_filter() + old_shell = self._shell + try: + self._shell = value + self.update_last_interaction_filter() + except: + self._shell = old_shell + raise ValueError("Invalid shell number") def update_last_interaction_filter(self): if self.packet_filter_mode == "packet_out_nu": @@ -142,29 +166,32 @@ def update_last_interaction_filter(self): self.last_line_interaction_shell_id == self.shell ) - self.last_line_in = self.lines.iloc[ + last_line_in = self.lines.iloc[ self.last_line_interaction_in_id[packet_filter] ] - self.last_line_out = self.lines.iloc[ + last_line_out = self.lines.iloc[ self.last_line_interaction_out_id[packet_filter] ] if self.atomic_number is not None: - self.last_line_in = self.last_line_in.xs( + last_line_in = last_line_in.xs( self.atomic_number, level="atomic_number", drop_level=False ) - self.last_line_out = self.last_line_out.xs( + last_line_out = last_line_out.xs( self.atomic_number, level="atomic_number", drop_level=False ) if self.ion_number is not None: - self.last_line_in = self.last_line_in.xs( + last_line_in = last_line_in.xs( self.ion_number, level="ion_number", drop_level=False ) - self.last_line_out = self.last_line_out.xs( + last_line_out = last_line_out.xs( self.ion_number, level="ion_number", drop_level=False ) + self.last_line_in = last_line_in + self.last_line_out = last_line_out + last_line_in_count = self.last_line_in.line_id.value_counts() last_line_out_count = self.last_line_out.line_id.value_counts() diff --git a/tardis/visualization/widgets/grotrian.py b/tardis/visualization/widgets/grotrian.py index d50611a53b1..9ee196fb346 100644 --- a/tardis/visualization/widgets/grotrian.py +++ b/tardis/visualization/widgets/grotrian.py @@ -296,7 +296,15 @@ def filter_mode(self): @filter_mode.setter def filter_mode(self, value): assert value in self.FILTER_MODES + + # Set the atomic_number and ion_number in the appropriate analysis object + self._line_interaction_analysis[value].set_ion( + self.atomic_number, self.ion_number + ) + self._line_interaction_analysis[value].shell = self.shell + self._filter_mode = value + self._compute_transitions() @property @@ -310,6 +318,16 @@ def set_ion(self, atomic_number, ion_number): Sets the atomic number and ion number """ assert type(atomic_number) is int and type(ion_number) is int + if (atomic_number, ion_number) not in self._level_energy_data.index or ( + atomic_number, + ion_number, + ) not in self._level_population_data.index: + raise ValueError( + "The (atomic_number, ion_number) pair doesn't exist in model" + ) + self._line_interaction_analysis[self.filter_mode].set_ion( + atomic_number, ion_number + ) self._atomic_number = atomic_number self._ion_number = ion_number @@ -342,6 +360,7 @@ def shell(self): @shell.setter def shell(self, value): assert value is None or type(value) is int + self._line_interaction_analysis[self.filter_mode].shell = value self._shell = value self._compute_level_data() self._compute_transitions() @@ -360,12 +379,6 @@ def _compute_transitions(self): """ Computes the excitation/de-excitation line transition data for the arrows in the widget """ - # Get relevant lines for current simulation - self._line_interaction_analysis[self.filter_mode].set_ion( - self.atomic_number, self.ion_number - ) - self._line_interaction_analysis[self.filter_mode].shell = self.shell - ### Get the excitation/de-excitation transitions from LastLineInteraction object excite_lines = ( self._line_interaction_analysis[self.filter_mode]