[Restructure/plasma] Refactor the simulation state initialization

docs/physics/plasma/construction_simple_plasma.ipynb

@@ -0,0 +1,115 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Constructing a simple plasma"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/wkerzend/python/tardis/tardis/__init__.py:20: UserWarning: Astropy is already imported externally. Astropy should be imported after TARDIS.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bbbd27367e48465696aa4e40f25b8496",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Iterations:          0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6830a2c86f754397999e6801ccca1001",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Packets:             0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import numpy as np\n",
+    "from astropy import units as u\n",
+    "from tardis.radiation_field.planck_rad_field import DilutePlanckianRadiationField\n",
+    "#from tardis.plasma.properties.plasma_input import T"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "temperatures_rad = np.ones(10) * 10000 * u.K\n",
+    "dilution_factor = np.ones(10) * 0.1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "d_radfield = DilutePlanckianRadiationField(temperature=temperatures_rad, dilution_factor=dilution_factor)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test = DilutePlanckianRadFieldInput()\n",
+    "test.set_value(d_radfield)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tardis/io/model/parse_radiation_field_configuration.py

@@ -1,5 +1,4 @@
 import logging
-import os
 
 import numpy as np
 from astropy import units as u
@@ -8,9 +7,7 @@
 from tardis.io.model.parse_geometry_configuration import (
     parse_structure_from_config,
 )
-from tardis.model.radiation_field_state import (
-    DiluteBlackBodyRadiationFieldState,
-)
+from tardis.plasma.radiation_field import DilutePlanckianRadiationField
 
 logger = logging.getLogger(__name__)
 
@@ -42,7 +39,7 @@
     AssertionError
         If the length of t_radiative or dilution_factor is not compatible with the geometry.
     """
     (
         density_time,
         velocity,
         density,
@@ -50,26 +47,24 @@
         temperature,
     ) = parse_structure_from_config(config)
 
     if temperature is None:
         if config.plasma.initial_t_rad > 0 * u.K:
             temperature = (
                 np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad
             )
         else:
             temperature = calculate_t_radiative_from_t_inner(
                 geometry, packet_source
             )
 
     assert len(temperature) == geometry.no_of_shells
 
     if dilution_factor is None:
         dilution_factor = calculate_geometric_dilution_factor(geometry)
 
     assert len(dilution_factor) == geometry.no_of_shells
 
-    return DiluteBlackBodyRadiationFieldState(
-        temperature, dilution_factor, geometry
-    )
+    return DilutePlanckianRadiationField(temperature, dilution_factor, geometry)
 
 
 def parse_radiation_field_state_from_csvy(
@@ -95,46 +90,44 @@
     DiluteThermalRadiationFieldState
         The parsed radiation field state.
     """
     t_radiative = None
     dilution_factor = None
 
     if hasattr(csvy_model_data, "columns") and (
         "t_rad" in csvy_model_data.columns
     ):
         t_rad_field_index = [
             field["name"] for field in csvy_model_config.datatype.fields
         ].index("t_rad")
         t_rad_unit = u.Unit(
             csvy_model_config.datatype.fields[t_rad_field_index]["unit"]
         )
         t_radiative = csvy_model_data["t_rad"].iloc[1:].values * t_rad_unit
 
     elif config.plasma.initial_t_rad > 0 * u.K:
         t_radiative = (
             np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad
         )
     else:
         t_radiative = calculate_t_radiative_from_t_inner(
             geometry, packet_source
         )
 
     if np.any(t_radiative < 1000 * u.K):
         logging.critical(
             "Radiative temperature is too low in some of the shells, temperatures below 1000K "
             f"(e.g., T_rad = {t_radiative[np.argmin(t_radiative)]} in shell {np.argmin(t_radiative)} in your model) "
             "are not accurately handled by TARDIS.",
         )
 
     if hasattr(csvy_model_data, "columns") and (
         "dilution_factor" in csvy_model_data.columns
     ):
         dilution_factor = csvy_model_data["dilution_factor"].iloc[1:].values
     else:
         dilution_factor = calculate_geometric_dilution_factor(geometry)
 
-    return DiluteBlackBodyRadiationFieldState(
-        t_radiative, dilution_factor, geometry
-    )
+    return DilutePlanckianRadiationField(t_radiative, dilution_factor, geometry)
 
 
 def calculate_t_radiative_from_t_inner(geometry, packet_source):
@@ -153,12 +146,12 @@
     Quantity
         The calculated radiative temperature.
     """
     lambda_wien_inner = const.b_wien / packet_source.temperature
     t_radiative = const.b_wien / (
         lambda_wien_inner
         * (1 + (geometry.v_middle - geometry.v_inner_boundary) / const.c)
     )
     return t_radiative
 
 
 def calculate_geometric_dilution_factor(geometry):
@@ -174,7 +167,7 @@
     np.array
         The dilution factors for the inout geometry
     """
     return 0.5 * (
         1
         - np.sqrt(
             1

tardis/io/model/readers/csvy.py

@@ -23,7 +23,7 @@
     -------
     yaml_dict : dictionary
                 YAML part of the csvy file
-    data : pandas.dataframe
+    data : pandas.DataFrame
             csv data from csvy file
     """
     with open(fname) as fh:
@@ -111,14 +111,14 @@
     isotope_mass_fraction : pandas.MultiIndex
     """
 
     mass_fraction_col_names = [
         name for name in csvy_data.columns if is_valid_nuclide_or_elem(name)
     ]
     df = csvy_data.loc[:, mass_fraction_col_names]
 
     df = df.transpose()
 
     mass_fractions = pd.DataFrame(
         columns=np.arange(df.shape[1]),
         index=pd.Index([], name="atomic_number"),
         dtype=np.float64,
@@ -127,20 +127,20 @@
     isotope_index = pd.MultiIndex(
         [[]] * 2, [[]] * 2, names=["atomic_number", "mass_number"]
     )
     isotope_mass_fractions = pd.DataFrame(
         columns=np.arange(df.shape[1]), index=isotope_index, dtype=np.float64
     )
 
     for element_symbol_string in df.index[0:]:
         if element_symbol_string in Z_DICT.values():
             z = elem_to_Z(element_symbol_string)
             mass_fractions.loc[z, :] = df.loc[element_symbol_string].tolist()
         else:
             nuc = Nuclide(element_symbol_string)
             z = nuc.Z
             mass_no = nuc.A
             isotope_mass_fractions.loc[(z, mass_no), :] = df.loc[
                 element_symbol_string
             ].tolist()
 
     return mass_fractions.index, mass_fractions, isotope_mass_fractions

tardis/model/base.py

@@ -146,14 +146,14 @@
 
     @property
     def t_radiative(self):
-        return self.radiation_field_state.t_radiative[
+        return self.radiation_field_state.temperature[
             self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index
         ]
 
     @t_radiative.setter
     def t_radiative(self, new_t_radiative):
         if len(new_t_radiative) == self.no_of_shells:
-            self.radiation_field_state.t_radiative[
+            self.radiation_field_state.temperature[
                 self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index
             ] = new_t_radiative
         else:
@@ -289,17 +289,17 @@
         """
         time_explosion = config.supernova.time_explosion.cgs
 
         geometry = parse_geometry_from_config(config, time_explosion)
 
         composition, electron_densities = parse_composition_from_config(
             atom_data, config, time_explosion, geometry
         )
 
         packet_source = parse_packet_source_from_config(
             config, geometry, legacy_mode_enabled
         )
 
         radiation_field_state = parse_radiation_field_state_from_config(
             config,
             geometry,
             dilution_factor=None,
@@ -382,11 +382,11 @@
 
         electron_densities = None
 
         geometry = parse_geometry_from_csvy(
             config, csvy_model_config, csvy_model_data, time_explosion
         )
 
         composition = parse_composition_from_csvy(
             atom_data,
             csvy_model_config,
             csvy_model_data,
@@ -394,11 +394,11 @@
             geometry,
         )
 
         packet_source = parse_packet_source_from_config(
             config, geometry, legacy_mode_enabled
         )
 
         radiation_field_state = parse_radiation_field_state_from_csvy(
             config, csvy_model_config, csvy_model_data, geometry, packet_source
         )
 

tardis/plasma/nlte/nlte_data.py

@@ -0,0 +1,123 @@
+import logging
+
+import numpy as np
+from scipy import interpolate
+
+logger = logging.getLogger(__name__)
+
+
+class NLTEData:
+    def __init__(self, atom_data, nlte_species):
+        self.atom_data = atom_data
+        self.lines = atom_data.lines.reset_index()
+        self.nlte_species = nlte_species
+
+        if nlte_species:
+            logger.info("Preparing the NLTE data")
+            self._init_indices()
+            if atom_data.collision_data is not None:
+                self._create_collision_coefficient_matrix()
+
+    def _init_indices(self):
+        self.lines_idx = {}
+        self.lines_level_number_lower = {}
+        self.lines_level_number_upper = {}
+        self.A_uls = {}
+        self.B_uls = {}
+        self.B_lus = {}
+
+        for species in self.nlte_species:
+            lines_idx = np.where(
+                (self.lines.atomic_number == species[0])
+                & (self.lines.ion_number == species[1])
+            )
+            self.lines_idx[species] = lines_idx
+            self.lines_level_number_lower[
+                species
+            ] = self.lines.level_number_lower.values[lines_idx].astype(int)
+            self.lines_level_number_upper[
+                species
+            ] = self.lines.level_number_upper.values[lines_idx].astype(int)
+
+            self.A_uls[species] = self.atom_data.lines.A_ul.values[lines_idx]
+            self.B_uls[species] = self.atom_data.lines.B_ul.values[lines_idx]
+            self.B_lus[species] = self.atom_data.lines.B_lu.values[lines_idx]
+
+    def _create_collision_coefficient_matrix(self):
+        self.C_ul_interpolator = {}
+        self.delta_E_matrices = {}
+        self.g_ratio_matrices = {}
+        collision_group = self.atom_data.collision_data.groupby(
+            level=["atomic_number", "ion_number"]
+        )
+        for species in self.nlte_species:
+            no_of_levels = self.atom_data.levels.loc[species].energy.count()
+            C_ul_matrix = np.zeros(
+                (
+                    no_of_levels,
+                    no_of_levels,
+                    len(self.atom_data.collision_data_temperatures),
+                )
+            )
+            delta_E_matrix = np.zeros((no_of_levels, no_of_levels))
+            g_ratio_matrix = np.zeros((no_of_levels, no_of_levels))
+
+            for (
+                (
+                    atomic_number,
+                    ion_number,
+                    level_number_lower,
+                    level_number_upper,
+                ),
+                line,
+            ) in collision_group.get_group(species).iterrows():
+                # line.columns : delta_e, g_ratio, temperatures ...
+                C_ul_matrix[
+                    level_number_lower, level_number_upper, :
+                ] = line.values[2:]
+                delta_E_matrix[level_number_lower, level_number_upper] = line[
+                    "delta_e"
+                ]
+                # TODO TARDISATOMIC fix change the g_ratio to be the otherway round - I flip them now here.
+                g_ratio_matrix[level_number_lower, level_number_upper] = (
+                    1 / line["g_ratio"]
+                )
+            self.C_ul_interpolator[species] = interpolate.interp1d(
+                self.atom_data.collision_data_temperatures, C_ul_matrix
+            )
+            self.delta_E_matrices[species] = delta_E_matrix
+
+            self.g_ratio_matrices[species] = g_ratio_matrix
+
+    def get_collision_matrix(self, species, t_electrons):
+        """
+        Creat collision matrix by interpolating the C_ul values for
+        the desired temperatures.
+        """
+        c_ul_matrix = self.C_ul_interpolator[species](t_electrons)
+        no_of_levels = c_ul_matrix.shape[0]
+        c_ul_matrix[np.isnan(c_ul_matrix)] = 0.0
+
+        # TODO in tardisatomic the g_ratio is the other way round - here I'll flip it in prepare_collision matrix
+
+        c_lu_matrix = (
+            c_ul_matrix
+            * np.exp(
+                -self.delta_E_matrices[species].reshape(
+                    (no_of_levels, no_of_levels, 1)
+                )
+                / t_electrons.reshape((1, 1, t_electrons.shape[0]))
+            )
+            * self.g_ratio_matrices[species].reshape(
+                (no_of_levels, no_of_levels, 1)
+            )
+        )
+        return c_ul_matrix + c_lu_matrix.transpose(1, 0, 2)
+
+
+class NLTEMatrixFactory:
+    def __init__(self, nlte_species):
+        pass
+
+    def generate_indices(self):
+        pass

tardis/plasma/properties/base.py

@@ -1,10 +1,9 @@
 import logging
+from abc import ABCMeta, abstractmethod
 
-from abc import ABCMeta, abstractmethod, abstractproperty
 import numpy as np
 import pandas as pd
 
-
 __all__ = [
     "BasePlasmaProperty",
     "BaseAtomicDataProperty",
@@ -19,7 +18,7 @@
 logger = logging.getLogger(__name__)
 
 
-class BasePlasmaProperty(object, metaclass=ABCMeta):
+class BasePlasmaProperty(metaclass=ABCMeta):
     """
     Attributes
     ----------
@@ -31,7 +30,8 @@
                  Used to label nodes when plotting graphs
     """
 
-    @abstractproperty
+    @property
+    @abstractmethod
     def outputs(self):
         pass
 
@@ -114,7 +114,7 @@
                 for i, output in enumerate(self.outputs):
                     setattr(self, output, new_values[i])
         else:
-            logger.info("{} has been frozen!".format(self.name))
+            logger.info(f"{self.name} has been frozen!")
 
     @abstractmethod
     def calculate(self, *args, **kwargs):
@@ -144,7 +144,8 @@
     track all transition probabilities and to later combine them.
     """
 
-    @abstractproperty
+    @property
+    @abstractmethod
     def transition_probabilities_outputs(self):
         pass