Ruff tardis/plasma (#2845)

* Ruff plasma autofix safe fixes

* unsafe fixes ruff

* Black plasma

tardis/plasma/__init__.py

@@ -1,5 +1,5 @@
 """
-Dealing with plasma properties and calculations related to it. 
+Dealing with plasma properties and calculations related to it.
 """
 
 from tardis.plasma.base import BasePlasma

tardis/plasma/base.py

@@ -43,15 +43,15 @@ def __getattr__(self, item):
         if item in self.outputs_dict:
             return self.get_value(item)
         else:
-            super(BasePlasma, self).__getattribute__(item)
+            super().__getattribute__(item)
 
     def __setattr__(self, key, value):
         if key != "module_dict" and key in self.outputs_dict:
             raise AttributeError(
                 "Plasma inputs can only be updated using " "the 'update' method"
             )
         else:
-            super(BasePlasma, self).__setattr__(key, value)
+            super().__setattr__(key, value)
 
     def __dir__(self):
         attrs = [item for item in self.__dict__ if not item.startswith("_")]
@@ -303,7 +303,7 @@ def write_to_dot(self, fname, args=None, latex_label=True):
         print_graph = self.remove_hidden_properties(print_graph)
 
         for node in print_graph:
-            if latex_label == True:
+            if latex_label is True:
                 if hasattr(self.plasma_properties_dict[node], "latex_formula"):
                     print_graph.nodes[str(node)][
                         "label"
@@ -324,13 +324,13 @@ def write_to_dot(self, fname, args=None, latex_label=True):
         for edge in print_graph.edges:
             label = print_graph.edges[edge]["label"]
             print_graph.edges[edge]["label"] = " "
-            if latex_label == True:
+            if latex_label is True:
                 print_graph.edges[edge]["texlbl"] = label
 
         nx.drawing.nx_agraph.write_dot(print_graph, fname)
 
         for line in fileinput.FileInput(fname, inplace=1):
-            if latex_label == True:
+            if latex_label is True:
                 print(
                     line.replace(
                         r'node [label="\N"]',

tardis/plasma/properties/atomic.py

@@ -3,12 +3,10 @@
 
 import numpy as np
 import pandas as pd
-import radioactivedecay as rd
 from numba import njit
 from scipy.interpolate import PchipInterpolator
 from scipy.special import exp1
 
-from tardis import constants as const
 from tardis.plasma.exceptions import IncompleteAtomicData
 from tardis.plasma.properties.base import (
     BaseAtomicDataProperty,
@@ -377,9 +375,7 @@ def calculate(self, atomic_data, continuum_interaction_species):
         if not mask_selected_species.sum():
             raise IncompleteAtomicData(
                 "two photon transition data for the requested "
-                "continuum_interactions species: {}".format(
-                    continuum_interaction_species.values.tolist()
-                )
+                f"continuum_interactions species: {continuum_interaction_species.values.tolist()}"
             )
         two_photon_data = two_photon_data[mask_selected_species]
         index_lower = two_photon_data.index.droplevel("level_number_upper")

tardis/plasma/properties/base.py

@@ -75,7 +75,7 @@ class ProcessingPlasmaProperty(BasePlasmaProperty, metaclass=ABCMeta):
     """
 
     def __init__(self, plasma_parent):
-        super(ProcessingPlasmaProperty, self).__init__()
+        super().__init__()
         self.plasma_parent = plasma_parent
         self._update_inputs()
         self.frozen = False
@@ -132,7 +132,7 @@ class HiddenPlasmaProperty(ProcessingPlasmaProperty, metaclass=ABCMeta):
     """
 
     def __init__(self, plasma_parent):
-        super(HiddenPlasmaProperty, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
 
 
 class TransitionProbabilitiesProperty(
@@ -160,7 +160,7 @@ class BaseAtomicDataProperty(ProcessingPlasmaProperty, metaclass=ABCMeta):
 
     def __init__(self, plasma_parent):
 
-        super(BaseAtomicDataProperty, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
         self.value = None
 
     @abstractmethod

tardis/plasma/properties/continuum_processes/__init__.py

@@ -3,7 +3,7 @@
 )
 from tardis.plasma.properties.continuum_processes.rates import *
 from tardis.plasma.properties.continuum_processes.recomb_rate_coeff import (
-    StimRecombRateFactor,
     SpontRecombRateCoeff,
     StimRecombRateCoeff,
+    StimRecombRateFactor,
 )

tardis/plasma/properties/continuum_processes/fast_array_util.py

@@ -1,11 +1,10 @@
 # It is currently not possible to use scipy.integrate.cumulative_trapezoid in
 # numba. So here is my own implementation.
-from tardis.transport.montecarlo import njit_dict
-
-
 import numpy as np
 from numba import njit, prange
 
+from tardis.transport.montecarlo import njit_dict
+
 
 @njit(**njit_dict)
 def numba_cumulative_trapezoid(f, x):

tardis/plasma/properties/general.py

@@ -1,7 +1,6 @@
 import logging
 
 import numpy as np
-from astropy import units as u
 
 from tardis import constants as const
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
@@ -30,7 +29,7 @@ class BetaRadiation(ProcessingPlasmaProperty):
     latex_formula = (r"\dfrac{1}{k_{B} T_{\textrm{rad}}}",)
 
     def __init__(self, plasma_parent):
-        super(BetaRadiation, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
         self.k_B_cgs = const.k_B.cgs.value
 
     def calculate(self, t_rad):
@@ -71,7 +70,7 @@ class ThermalGElectron(GElectron):
     )
 
     def calculate(self, beta_electron):
-        return super(ThermalGElectron, self).calculate(beta_electron)
+        return super().calculate(beta_electron)
 
 
 class SelectedAtoms(ProcessingPlasmaProperty):
@@ -115,7 +114,7 @@ class BetaElectron(ProcessingPlasmaProperty):
     latex_formula = (r"\frac{1}{K_{B} T_{\textrm{electron}}}",)
 
     def __init__(self, plasma_parent):
-        super(BetaElectron, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
         self.k_B_cgs = const.k_B.cgs.value
 
     def calculate(self, t_electrons):

tardis/plasma/properties/helium_nlte.py

@@ -50,7 +50,7 @@ def calculate(
         helium_population.loc[0, 0] = 0.0
         # He II excited states
         he_two_population = level_boltzmann_factor.loc[2, 1].mul(
-            (float(g.loc[2, 1, 0]) ** (-1.0))
+            float(g.loc[2, 1, 0]) ** (-1.0)
         )
         helium_population.loc[
             1, helium_population.columns
@@ -127,7 +127,7 @@ class HeliumNumericalNLTE(ProcessingPlasmaProperty):
     outputs = ("helium_population",)
 
     def __init__(self, plasma_parent, heating_rate_data_file):
-        super(HeliumNumericalNLTE, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
         self._g_upper = None
         self._g_lower = None
         self.heating_rate_data = np.loadtxt(heating_rate_data_file, unpack=True)
@@ -232,10 +232,8 @@ def calculate(
         # Reading in populations from files
         helium_population = level_boltzmann_factor.loc[2].copy()
         for zone, _ in enumerate(electron_densities):
-            with open(
-                f"He_NLTE_Files/discradfield{zone}.txt", "r"
-            ) as read_file:
-                for level in range(0, 35):
+            with open(f"He_NLTE_Files/discradfield{zone}.txt") as read_file:
+                for level in range(35):
                     level_population = read_file.readline()
                     level_population = float(level_population)
                     helium_population[zone].loc[0, level] = level_population

tardis/plasma/properties/ion_population.py

@@ -203,7 +203,7 @@ def __init__(
         chi_0_species=(20, 2),
         delta_treatment=None,
     ):
-        super(RadiationFieldCorrection, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
         self.departure_coefficient = departure_coefficient
         self.delta_treatment = delta_treatment
         self.chi_0_species = chi_0_species
@@ -288,7 +288,7 @@ def __init__(
         ion_zero_threshold=ION_ZERO_THRESHOLD,
         electron_densities=None,
     ):
-        super(IonNumberDensity, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
         self.ion_zero_threshold = ion_zero_threshold
         self.block_ids = None
         self._electron_densities = electron_densities
@@ -421,7 +421,7 @@ class IonNumberDensityHeNLTE(ProcessingPlasmaProperty):
     def __init__(
         self, plasma_parent, ion_zero_threshold=1e-20, electron_densities=None
     ):
-        super(IonNumberDensityHeNLTE, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
         self.ion_zero_threshold = ion_zero_threshold
         self.block_ids = None
         self._electron_densities = electron_densities

tardis/plasma/properties/level_population.py

@@ -1,6 +1,7 @@
 import logging
-import pandas as pd
+
 import numpy as np
+import pandas as pd
 
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
 
@@ -25,7 +26,7 @@ class LevelNumberDensity(ProcessingPlasmaProperty):
     latex_formula = (r"N_{i,j}\dfrac{bf_{i,j,k}}{Z_{i,j}}",)
 
     def __init__(self, plasma_parent):
-        super(LevelNumberDensity, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
         self._update_inputs()
         self.initialize_indices = True
 

tardis/plasma/properties/nlte.py

@@ -1,13 +1,10 @@
 import logging
-import os
 
-import numpy as np
 import pandas as pd
 
 from tardis.plasma.properties.base import (
     PreviousIterationProperty,
 )
-from tardis.plasma.properties.ion_population import PhiSahaNebular
 
 __all__ = [
     "PreviousElectronDensities",

tardis/plasma/properties/nlte_excitation_data.py

@@ -5,7 +5,7 @@
 logger = logging.getLogger(__name__)
 
 
-class NLTEExcitationData(object):
+class NLTEExcitationData:
     """Data needed for NLTE excitation treatment."""
 
     def __init__(self, atom_data_lines, nlte_excitation_species):

tardis/plasma/properties/nlte_rate_equation_solver.py

@@ -299,9 +299,7 @@ def calculate(
                     < NLTE_POPULATION_SOLVER_TOLERANCE
                 ):
                     logger.debug(
-                        "NLTE ionization solver converged after {} iterations for shell {}".format(
-                            iteration, shell
-                        )
+                        f"NLTE ionization solver converged after {iteration} iterations for shell {shell}"
                     )
                     break
 

tardis/plasma/properties/partition_function.py

@@ -1,11 +1,11 @@
 import logging
 
 import numpy as np
-from numpy.linalg.linalg import LinAlgError
 import pandas as pd
+from numpy.linalg.linalg import LinAlgError
 
-from tardis.plasma.properties.base import ProcessingPlasmaProperty
 from tardis.plasma.exceptions import PlasmaConfigError
+from tardis.plasma.properties.base import ProcessingPlasmaProperty
 
 logger = logging.getLogger(__name__)
 
@@ -158,7 +158,7 @@ def __init__(self, plasma_parent):
         Selects appropriate 'calculate' function based on NLTE config
         options selected.
         """
-        super(LevelBoltzmannFactorNLTE, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
 
         self._update_inputs()
 
@@ -392,6 +392,4 @@ class ThermalLTEPartitionFunction(PartitionFunction):
     latex_name = (r"Z_{i,j}(T_\mathrm{e}",)
 
     def calculate(self, thermal_lte_level_boltzmann_factor):
-        return super(ThermalLTEPartitionFunction, self).calculate(
-            thermal_lte_level_boltzmann_factor
-        )
+        return super().calculate(thermal_lte_level_boltzmann_factor)

tardis/plasma/properties/plasma_input.py

@@ -1,5 +1,4 @@
 from tardis.plasma.properties.base import (
-    ArrayInput,
     Input,
     ProcessingPlasmaProperty,
 )

tardis/plasma/properties/radiative_properties.py

@@ -5,15 +5,13 @@
 from astropy import units as u
 from numba import jit, prange
 
-
 from tardis import constants as const
 from tardis.opacities.macro_atom.base import TransitionProbabilities
 from tardis.plasma.properties.base import (
     ProcessingPlasmaProperty,
     TransitionProbabilitiesProperty,
 )
 
-
 logger = logging.getLogger(__name__)
 
 __all__ = [
@@ -39,7 +37,7 @@ class StimulatedEmissionFactor(ProcessingPlasmaProperty):
     latex_formula = (r"1-\dfrac{g_{lower}n_{upper}}{g_{upper}n_{lower}}",)
 
     def __init__(self, plasma_parent=None, nlte_species=None):
-        super(StimulatedEmissionFactor, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
         self._g_upper = None
         self._g_lower = None
         self.nlte_species = nlte_species
@@ -175,7 +173,7 @@ class RawRadBoundBoundTransProbs(
     transition_probabilities_outputs = ("p_rad_bb",)
 
     def __init__(self, plasma_parent):
-        super(RawRadBoundBoundTransProbs, self).__init__(plasma_parent)
+        super().__init__(plasma_parent)
         self.normalize = False
 
     def calculate(

tardis/plasma/properties/rate_matrix_index.py

@@ -1,4 +1,5 @@
 import pandas as pd
+
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
 
 __all__ = [

tardis/plasma/tests/test_complete_plasmas.py

@@ -146,12 +146,12 @@ def config(self, request):
         config = Configuration.from_yaml(PLASMA_CONFIG_FPATH)
         hash_string = ""
         for prop, value in request.param.items():
-            hash_string = "_".join((hash_string, prop))
+            hash_string = f"{hash_string}_{prop}"
             if prop == "nlte":
                 for nlte_prop, nlte_value in request.param[prop].items():
                     config.plasma.nlte[nlte_prop] = nlte_value
                     if nlte_prop != "species":
-                        hash_string = "_".join((hash_string, nlte_prop))
+                        hash_string = f"{hash_string}_{nlte_prop}"
             else:
                 config.plasma[prop] = value
                 hash_string = "_".join((hash_string, str(value)))

tardis/plasma/tests/test_hdf_plasma.py

@@ -1,4 +1,3 @@
-import numpy as np
 import numpy.testing as npt
 import pandas.testing as pdt
 import pytest

tardis/plasma/tests/test_nlte_excitation.py

@@ -2,13 +2,12 @@
 import numpy.testing as npt
 import pandas as pd
 import pytest
-from numpy.testing import assert_allclose
 
 from tardis.plasma.properties.nlte_excitation_data import NLTEExcitationData
 from tardis.plasma.properties.nlte_rate_equation_solver import (
-    prepare_r_uls_r_lus,
-    prepare_bound_bound_rate_matrix,
     create_coll_exc_deexc_matrix,
+    prepare_bound_bound_rate_matrix,
+    prepare_r_uls_r_lus,
 )