diff --git a/docs/physics/plasma/equilibrium/chianti_solver.ipynb b/docs/physics/plasma/equilibrium/chianti_solver.ipynb new file mode 100644 index 00000000000..d4f77ccfb5b --- /dev/null +++ b/docs/physics/plasma/equilibrium/chianti_solver.ipynb @@ -0,0 +1,127 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "os.environ['XUVTOP'] = \"~/chianti\"\n", + "\n", + "try:\n", + " import ChiantiPy.core as ch\n", + " from ChiantiPy.tools.io import versionRead\n", + "\n", + "except ImportError:\n", + " # Shamefully copied from their GitHub source:\n", + " import chianti.core as ch\n", + " def versionRead():\n", + " \"\"\"\n", + " Read the version number of the CHIANTI database\n", + " \"\"\"\n", + " xuvtop = os.environ['XUVTOP']\n", + " vFileName = os.path.join(xuvtop, 'VERSION')\n", + " vFile = open(vFileName)\n", + " versionStr = vFile.readline()\n", + " vFile.close()\n", + " return versionStr.strip()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "compare_shell = 0\n", + "species = (14, 1)\n", + "species_string = 'si_2'" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "t_electron = 9967.4884" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "electron_density = 1e6" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "chianti_species = ch.ion(species_string, temperature=t_electron, eDensity=electron_density)\n", + "chianti_species.populate()\n", + "chianti_population = chianti_species.Population['population']" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([[9.99999989e-01, 1.06967533e-08, 8.19021735e-17, 1.63935872e-16,\n", + " 3.82270728e-16, 8.00932946e-18, 1.60440977e-17, 1.62648689e-17,\n", + " 2.44571411e-17, 9.10400264e-17, 3.26639365e-18, 6.54275441e-18,\n", + " 7.80689247e-18, 1.17291171e-17, 7.51999263e-18, 1.01449887e-17,\n", + " 6.56279954e-17, 2.63927140e-18, 5.28562769e-18, 6.84869862e-18,\n", + " 1.02728743e-17, 6.36689557e-18, 8.48045348e-18, 5.33290944e-18,\n", + " 6.25920357e-18]])" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "chianti_population" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "carsus", + "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.12.5" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/physics/plasma/detailed_balance/collisional_rate_coefficients.hdf b/docs/physics/plasma/equilibrium/collisional_rate_coefficients.hdf similarity index 100% rename from docs/physics/plasma/detailed_balance/collisional_rate_coefficients.hdf rename to docs/physics/plasma/equilibrium/collisional_rate_coefficients.hdf diff --git a/docs/physics/plasma/detailed_balance/comparison.ipynb b/docs/physics/plasma/equilibrium/comparison.ipynb similarity index 99% rename from docs/physics/plasma/detailed_balance/comparison.ipynb rename to docs/physics/plasma/equilibrium/comparison.ipynb index f15ceb6145c..2d1b5f4ea9a 100644 --- a/docs/physics/plasma/detailed_balance/comparison.ipynb +++ b/docs/physics/plasma/equilibrium/comparison.ipynb @@ -38,7 +38,7 @@ "from astropy import units as u\n", "\n", "from tardis.io.atom_data import AtomData\n", - "from tardis.plasma.detailed_balance.rates import (\n", + "from tardis.plasma.equilibrium.rates import (\n", " RadiativeRatesSolver,\n", " ThermalCollisionalRateSolver,\n", " UpsilonRegemorterSolver,\n", diff --git a/docs/physics/plasma/detailed_balance/rates.ipynb b/docs/physics/plasma/equilibrium/rates.ipynb similarity index 100% rename from docs/physics/plasma/detailed_balance/rates.ipynb rename to docs/physics/plasma/equilibrium/rates.ipynb diff --git a/docs/physics/plasma/equilibrium/tardis_solver.ipynb b/docs/physics/plasma/equilibrium/tardis_solver.ipynb new file mode 100644 index 00000000000..74279400dba --- /dev/null +++ b/docs/physics/plasma/equilibrium/tardis_solver.ipynb @@ -0,0 +1,227 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Comparison between TARDIS and ChiantiPy" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "import astropy.units as u\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "from tardis.io.atom_data import AtomData\n", + "from tardis.io.configuration.config_reader import Configuration\n", + "from tardis.model.base import SimulationState\n", + "from tardis.plasma.equilibrium.rates import (\n", + " RadiativeRatesSolver,\n", + " ThermalCollisionalRateSolver,\n", + ")\n", + "from tardis.plasma.radiation_field import (\n", + " DilutePlanckianRadiationField,\n", + ")\n", + "\n", + "home = str(Path('~').expanduser())\n", + "\n", + "config = Configuration.from_yaml(home+\"/tardis/tardis/plasma/tests/data/plasma_base_test_config.yml\")\n", + "\n", + "ion_slice = (2, 0, slice(None), slice(None))\n", + "\n", + "\n", + "def get_radiative_rate_solver(radiative_transitions):\n", + " rad_rate_solver = RadiativeRatesSolver(radiative_transitions)\n", + " return rad_rate_solver\n", + "\n", + "def get_chianti_collisional_rate_solver(atom_data, radiative_transitions,):\n", + " col_strength_temperatures = atom_data.collision_data_temperatures\n", + " col_strengths = atom_data.collision_data.loc[ion_slice, :]\n", + " collisional_rate_solver = ThermalCollisionalRateSolver(atom_data.levels, radiative_transitions, col_strength_temperatures, col_strengths, 'chianti', \"none\")\n", + " return collisional_rate_solver\n", + "\n", + "chianti_atom_data = AtomData.from_hdf(home+'/carsus/docs/kurucz_cd23_chianti_H_Si.h5')\n", + "\n", + "chianti_radiative_transitions = chianti_atom_data.lines.loc[ion_slice, :]\n", + "\n", + "chianti_sim_state = SimulationState.from_config(config, atom_data=chianti_atom_data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from tardis.plasma.electron_energy_distribution import (\n", + " ThermalElectronEnergyDistribution,\n", + ")\n", + "\n", + "chianti_radiative_rate_solver = get_radiative_rate_solver(chianti_radiative_transitions)\n", + "\n", + "chianti_collisional_rate_solver = get_chianti_collisional_rate_solver(chianti_atom_data, chianti_radiative_transitions)\n", + "\n", + "temperature = chianti_sim_state.t_radiative\n", + "\n", + "# need a better way to make a custom radiation field where the intensity is zero\n", + "# in specific locations as desired\n", + "rad_field = DilutePlanckianRadiationField(chianti_sim_state.t_radiative, dilution_factor=np.zeros_like(chianti_sim_state.t_radiative))\n", + "electron_dist = ThermalElectronEnergyDistribution(0, chianti_sim_state.t_radiative, 1e6 * u.g/u.cm**3)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rate_solvers = [(chianti_radiative_rate_solver, \"radiative\"), (chianti_collisional_rate_solver, \"electron\")]\n", + "\n", + "lte_rate_solvers = [(chianti_collisional_rate_solver, \"electron\")]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rad_rates = chianti_radiative_rate_solver.solve(rad_field)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "col_rates = chianti_collisional_rate_solver.solve(electron_dist.temperature)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from tardis.plasma.equilibrium.level_populations import LevelPopulationSolver\n", + "from tardis.plasma.equilibrium.rate_matrix import RateMatrix\n", + "\n", + "rate_matrix_solver = RateMatrix(rate_solvers, chianti_atom_data.levels)\n", + "\n", + "rate_matrix = rate_matrix_solver.solve(rad_field, electron_dist)\n", + "\n", + "lte_rate_matrix = RateMatrix(lte_rate_solvers, chianti_atom_data.levels).solve(rad_field, electron_dist)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "solver = LevelPopulationSolver(rate_matrix, chianti_atom_data.levels)\n", + "\n", + "level_pops = solver.solve()\n", + "\n", + "lte_level_pops = LevelPopulationSolver(lte_rate_matrix, chianti_atom_data.levels).solve()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "chianti_levels = [3.42881304e-01, 6.57118589e-01, 4.18696636e-09, 5.67602616e-08,\n", + " 4.63312038e-08, 1.20482595e-11, 1.97010580e-11, 3.48183541e-15,\n", + " 6.88699562e-16, 2.78695447e-16, 4.20834423e-16, 1.25473722e-15,\n", + " 2.49671361e-15, 8.92884539e-17, 1.65423100e-16, 2.11333508e-17,\n", + " 1.11695201e-17, 1.69708563e-17, 2.70670610e-17, 5.36487601e-17,\n", + " 8.72033865e-18, 1.17118808e-17, 3.44693669e-16, 5.11955673e-16,\n", + " 2.81010802e-18, 3.50318771e-18, 5.25951995e-18, 3.28532989e-18,\n", + " 6.57049588e-18]\n", + "\n", + "chianti_levels_he_1 = [1.00000000e+00, 1.75073626e-11, 1.8167685e-8, 1.79805022e-19,\n", + " 1.07614011e-19, 3.58865323e-20, 8.18076407e-23, 1.76624918e-21,\n", + " 3.95823826e-22, 1.30154195e-21, 7.81338357e-22, 2.60269199e-22,\n", + " 1.83743361e-22, 1.21199070e-22, 7.93684710e-23, 7.69690651e-23,\n", + " 5.55204217e-24, 4.28655703e-23, 8.91723472e-23, 3.08213125e-22,\n", + " 1.84912140e-22, 6.16147989e-23, 5.28143670e-23, 3.77505779e-23,\n", + " 2.26343284e-23, 3.19064683e-23, 8.89209221e-24, 9.39180485e-23,\n", + " 5.19053997e-23, 2.53784273e-23, 2.27283788e-24, 1.11837286e-22,\n", + " 3.85509048e-23, 1.07725425e-22, 6.46134793e-23, 2.15124099e-23,\n", + " 2.22314242e-23, 1.59215321e-23, 9.54500244e-24, 1.48014892e-23,\n", + " 3.80078829e-23, 4.89438878e-23, 2.71814793e-23, 1.79147356e-23,\n", + " 9.80689599e-23, 7.25122090e-23, 8.04471105e-23, 1.27910690e-23,\n", + " 1.51551660e-24]\n", + "\n", + "chianti_levels_o_2 = [9.50427666e-01, 2.93020096e-02, 1.96811489e-02, 3.73612176e-04,\n", + " 2.15563005e-04, 1.19850233e-18, 7.62882704e-19, 3.73304507e-19,\n", + " 6.86402347e-22, 4.51096625e-22, 1.86694313e-23, 3.93108508e-23,\n", + " 5.19716753e-23, 5.48509200e-24, 1.05025627e-23, 4.23498947e-25,\n", + " 5.73694095e-23, 2.27177963e-24, 4.53847634e-24, 6.68361893e-24,\n", + " 8.38430369e-24, 9.60077484e-25, 6.19972361e-25, 1.04127656e-24,\n", + " 1.66000813e-24, 2.48443321e-24, 6.16103282e-24, 9.08726063e-24,\n", + " 3.02895913e-24, 1.64793516e-25, 7.92929866e-26, 2.66317598e-23,\n", + " 6.44175783e-23, 5.28146898e-25, 7.36488325e-25]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Current issues:\n", + "\n", + "- Rates matrix layout is different between Chianti and TARDIS despite identical data\n", + "- Still some large % differences for high-energy lines\n", + "- He I is different for the second and third levels\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import astropy.constants as const\n", + "\n", + "plt.scatter(chianti_atom_data.levels.loc[2,0].energy * u.erg.to('eV'), chianti_levels_he_1, marker='o', label='CHIANTI')\n", + "plt.scatter(chianti_atom_data.levels.loc[2,0].energy * u.erg.to('eV'), level_pops.loc[2,0,:][0], marker='x', label='TARDIS')\n", + "plt.scatter(chianti_atom_data.levels.loc[2,0].energy * u.erg.to('eV'), lte_level_pops.loc[2,0,:][0], marker='x', label='TARDIS col only')\n", + "plt.xlabel(\"Energy (eV)\")\n", + "plt.ylabel(\"Population\")\n", + "plt.semilogy()\n", + "plt.legend()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "tardis", + "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.12.4" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/physics/plasma/detailed_balance/test_continuum_template_wkerzen_rate_coeffs.yml b/docs/physics/plasma/equilibrium/test_continuum_template_wkerzen_rate_coeffs.yml similarity index 100% rename from docs/physics/plasma/detailed_balance/test_continuum_template_wkerzen_rate_coeffs.yml rename to docs/physics/plasma/equilibrium/test_continuum_template_wkerzen_rate_coeffs.yml diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 33e8e4bbca2..40acbe0b233 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -171,7 +171,6 @@ class AtomData: # Either all or none of the related dataframes must be given related_groups = [ ("macro_atom_data_all", "macro_atom_references_all"), - ("collision_data", "collision_data_temperatures"), ] @classmethod @@ -203,35 +202,32 @@ def from_hdf(cls, fname=None): store["metadata"].loc[("format", "version")].value ) carsus_version = tuple(map(int, carsus_version_str.split("."))) - if carsus_version == (1, 0): - # Checks for various collisional data from Carsus files - if "collisions_data" in store: - try: + # Checks for various collisional data from Carsus files + if "collisions_data" in store: + try: + if carsus_version == (1, 0): dataframes["collision_data_temperatures"] = store[ "collisions_metadata" ].temperatures - if "cmfgen" in store["collisions_metadata"].dataset: - dataframes["yg_data"] = store["collisions_data"] - dataframes["collision_data"] = "dummy value" - elif ( - "chianti" - in store["collisions_metadata"].dataset - ): - dataframes["collision_data"] = store[ - "collisions_data" - ] - else: - raise KeyError( - "Atomic Data Collisions Not a Valid Chanti or CMFGEN Carsus Data File" - ) - except KeyError as e: - logger.warning( - "Atomic Data is not a Valid Carsus Atomic Data File" + if "cmfgen" in store["collisions_metadata"].dataset: + dataframes["yg_data"] = store["collisions_data"] + dataframes["collision_data"] = "dummy value" + elif "chianti" in store["collisions_metadata"].dataset: + dataframes["collision_data"] = store[ + "collisions_data" + ] + else: + raise KeyError( + "Atomic Data Collisions Not a Valid Chanti or CMFGEN Carsus Data File" ) - raise - dataframes["levels"] = store["levels_data"] - dataframes["lines"] = store["lines_data"] - else: + except KeyError as e: + logger.warning( + "Atomic Data is not a Valid Carsus Atomic Data File" + ) + raise + dataframes["levels"] = store["levels_data"] + dataframes["lines"] = store["lines_data"] + if carsus_version != (1, 0) and carsus_version != (2, 0): raise ValueError( f"Current carsus version, {carsus_version}, is not supported." ) @@ -528,9 +524,9 @@ def prepare_continuum_interaction_data(self, continuum_interaction_species): ) level_idxs2continuum_idx = photo_ion_levels_idx.copy() - level_idxs2continuum_idx[ - "continuum_idx" - ] = self.level2continuum_edge_idx + level_idxs2continuum_idx["continuum_idx"] = ( + self.level2continuum_edge_idx + ) self.level_idxs2continuum_idx = level_idxs2continuum_idx.set_index( ["source_level_idx", "destination_level_idx"] ) @@ -564,31 +560,33 @@ def prepare_macro_atom_data( self.macro_atom_references = self.macro_atom_references.loc[ self.macro_atom_references["count_down"] > 0 ] - self.macro_atom_references.loc[ - :, "count_total" - ] = self.macro_atom_references["count_down"] - self.macro_atom_references.loc[ - :, "block_references" - ] = np.hstack( - ( - 0, - np.cumsum( - self.macro_atom_references["count_down"].values[:-1] - ), + self.macro_atom_references.loc[:, "count_total"] = ( + self.macro_atom_references["count_down"] + ) + self.macro_atom_references.loc[:, "block_references"] = ( + np.hstack( + ( + 0, + np.cumsum( + self.macro_atom_references["count_down"].values[ + :-1 + ] + ), + ) ) ) elif line_interaction_type == "macroatom": - self.macro_atom_references.loc[ - :, "block_references" - ] = np.hstack( - ( - 0, - np.cumsum( - self.macro_atom_references["count_total"].values[ - :-1 - ] - ), + self.macro_atom_references.loc[:, "block_references"] = ( + np.hstack( + ( + 0, + np.cumsum( + self.macro_atom_references[ + "count_total" + ].values[:-1] + ), + ) ) ) diff --git a/tardis/plasma/detailed_balance/rates/__init__.py b/tardis/plasma/detailed_balance/rates/__init__.py deleted file mode 100644 index b4101b7d28c..00000000000 --- a/tardis/plasma/detailed_balance/rates/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -from tardis.plasma.detailed_balance.rates.collision_strengths import ( - UpsilonCMFGENSolver, - UpsilonRegemorterSolver, -) -from tardis.plasma.detailed_balance.rates.collisional_rates import ( - ThermalCollisionalRateSolver, -) -from tardis.plasma.detailed_balance.rates.radiative_rates import ( - RadiativeRatesSolver, -) diff --git a/tardis/plasma/electron_energy_distribution/__init__.py b/tardis/plasma/electron_energy_distribution/__init__.py new file mode 100644 index 00000000000..0aa2df215cf --- /dev/null +++ b/tardis/plasma/electron_energy_distribution/__init__.py @@ -0,0 +1,4 @@ +from tardis.plasma.electron_energy_distribution.base import ( + ElectronEnergyDistribution, + ThermalElectronEnergyDistribution, +) diff --git a/tardis/plasma/electron_energy_distribution/base.py b/tardis/plasma/electron_energy_distribution/base.py new file mode 100644 index 00000000000..bbd9cc48542 --- /dev/null +++ b/tardis/plasma/electron_energy_distribution/base.py @@ -0,0 +1,28 @@ +from dataclasses import dataclass + +from astropy import units as u + + +@dataclass +class ElectronEnergyDistribution: + """Electron energy distribution. + + energy : Quantity + Electron energies in ergs. + """ + + energy: u.Quantity + + +@dataclass +class ThermalElectronEnergyDistribution(ElectronEnergyDistribution): + """Thermal electron temperature and density distribution. + + temperature : Quantity + Electron temperatures in Kelvin. + number_density : Quantity + Electron number densities in g/cm^3. + """ + + temperature: u.Quantity + number_density: u.Quantity diff --git a/tardis/plasma/detailed_balance/__init__.py b/tardis/plasma/equilibrium/__init__.py similarity index 100% rename from tardis/plasma/detailed_balance/__init__.py rename to tardis/plasma/equilibrium/__init__.py diff --git a/tardis/plasma/equilibrium/level_populations.py b/tardis/plasma/equilibrium/level_populations.py new file mode 100644 index 00000000000..0907c9de64b --- /dev/null +++ b/tardis/plasma/equilibrium/level_populations.py @@ -0,0 +1,69 @@ +import numpy as np +import pandas as pd + + +class LevelPopulationSolver: + def __init__(self, rates_matrices: pd.DataFrame, levels: pd.DataFrame): + """Solve the normalized level population values from the rate matrices. + + Parameters + ---------- + rates_matrices : pd.DataFrame + DataFrame of rate matrices indexed by atomic number and ion number, + with each column being a cell. + levels : pd.DataFrame + DataFrame of energy levels. + """ + self.rates_matrices = rates_matrices + self.levels = levels + + def __calculate_level_population(self, rates_matrix: np.ndarray): + """Helper function to calculate the normalized, per-level boltzmann factor. + + Parameters + ---------- + rates_matrix : np.ndarray + The rate matrix for a given species and cell. + + Returns + ------- + np.ndarray + The normalized, per-level population. + """ + normalized_ion_population = np.zeros(rates_matrix.shape[0]) + normalized_ion_population[0] = 1.0 + normalized_level_population = np.linalg.solve( + rates_matrix[:, :], normalized_ion_population + ) + return normalized_level_population + + def solve(self): + """Solves the normalized level population values from the rate matrices. + + Returns + ------- + pd.DataFrame + Normalized level population values indexed by atomic number, ion + number and level number. Columns are cells. + """ + normalized_level_populations = pd.DataFrame( + index=self.levels.index, + columns=self.rates_matrices.columns, + dtype=np.float64, + ) + + # try converting the set of vectors into a single 2D array and then applying index + for species_id in self.rates_matrices.index: + # TODO: resolve the chained assignment here. Maybe an intermediate df + # is needed + + solved_matrices = self.rates_matrices.loc[species_id].apply( + lambda rates_matrix: self.__calculate_level_population( + rates_matrix + ) + ) + normalized_level_populations.loc[species_id, :].update( + np.vstack(solved_matrices.values).T + ) + + return normalized_level_populations diff --git a/tardis/plasma/equilibrium/rate_matrix.py b/tardis/plasma/equilibrium/rate_matrix.py new file mode 100644 index 00000000000..a61fb7045e2 --- /dev/null +++ b/tardis/plasma/equilibrium/rate_matrix.py @@ -0,0 +1,106 @@ +import numpy as np +import pandas as pd +from scipy.sparse import coo_matrix + + +class RateMatrix: + def __init__( + self, + rate_solvers: list, + levels: pd.DataFrame, + ): + """Constructs the rate matrix from an arbitrary number of rate solvers. + + Parameters + ---------- + rate_solvers : list + List of rate solver objects. + levels : pd.DataFrame + DataFrame of energy levels. + """ + self.rate_solvers = rate_solvers + self.levels = levels + + def solve( + self, + radiation_field, + thermal_electron_energy_distribution, + ): + """Construct the compiled rate matrix dataframe. + + Parameters + ---------- + radiation_field : RadiationField + Radiation field containing radiative temperature. + thermal_electron_energy_distribution : ThermalElectronEnergyDistribution + Distribution of electrons in the plasma, containing electron energies, + temperatures and number densities. + + Returns + ------- + pd.DataFrame + A DataFrame of rate matrices indexed by atomic number and ion number, + with each column being a cell. + """ + required_arg = { + "radiative": radiation_field, + "electron": thermal_electron_energy_distribution.temperature, + } + + rates_df_list = [ + solver.solve(required_arg[arg]) for solver, arg in self.rate_solvers + ] + # Extract all indexes + all_indexes = set() + for df in rates_df_list: + all_indexes.update(df.index) + + # Create a union of all indexes + all_indexes = sorted(all_indexes) + + # Reindex each dataframe to ensure consistent indices + rates_df_list = [ + df.reindex(all_indexes, fill_value=0) for df in rates_df_list + ] + + # Multiply rates by electron number density where appropriate + rates_df_list = [ + rates_df * thermal_electron_energy_distribution.number_density + if solver_arg_tuple[1] == "electron" + else rates_df + for solver_arg_tuple, rates_df in zip( + self.rate_solvers, rates_df_list + ) + ] + + rates_df = sum(rates_df_list) + + grouped_rates_df = rates_df.groupby( + level=("atomic_number", "ion_number") + ) + + rate_matrices = pd.DataFrame( + index=grouped_rates_df.groups.keys(), columns=rates_df.columns + ) + + for species_id, rates in grouped_rates_df: + number_of_levels = self.levels.energy.loc[species_id].count() + for shell in range(len(rates.columns)): + matrix = coo_matrix( + ( + rates[shell], + ( + rates.index.get_level_values( + "level_number_destination" + ), + rates.index.get_level_values("level_number_source"), + ), + ), + shape=(number_of_levels, number_of_levels), + ) + matrix_array = matrix.toarray() + np.fill_diagonal(matrix_array, -np.sum(matrix_array, axis=0)) + matrix_array[0, :] = 1 + rate_matrices.loc[species_id, shell] = matrix_array + + return rate_matrices diff --git a/tardis/plasma/equilibrium/rates/__init__.py b/tardis/plasma/equilibrium/rates/__init__.py new file mode 100644 index 00000000000..2d7e4b79bc2 --- /dev/null +++ b/tardis/plasma/equilibrium/rates/__init__.py @@ -0,0 +1,10 @@ +from tardis.plasma.equilibrium.rates.collision_strengths import ( + UpsilonCMFGENSolver, + UpsilonRegemorterSolver, +) +from tardis.plasma.equilibrium.rates.collisional_rates import ( + ThermalCollisionalRateSolver, +) +from tardis.plasma.equilibrium.rates.radiative_rates import ( + RadiativeRatesSolver, +) diff --git a/tardis/plasma/detailed_balance/rates/collision_strengths.py b/tardis/plasma/equilibrium/rates/collision_strengths.py similarity index 100% rename from tardis/plasma/detailed_balance/rates/collision_strengths.py rename to tardis/plasma/equilibrium/rates/collision_strengths.py diff --git a/tardis/plasma/detailed_balance/rates/collisional_rates.py b/tardis/plasma/equilibrium/rates/collisional_rates.py similarity index 74% rename from tardis/plasma/detailed_balance/rates/collisional_rates.py rename to tardis/plasma/equilibrium/rates/collisional_rates.py index 32ae38dbfae..8143ec9ca05 100644 --- a/tardis/plasma/detailed_balance/rates/collisional_rates.py +++ b/tardis/plasma/equilibrium/rates/collisional_rates.py @@ -3,7 +3,7 @@ from astropy import units as u from tardis import constants as const -from tardis.plasma.detailed_balance.rates.collision_strengths import ( +from tardis.plasma.equilibrium.rates.collision_strengths import ( UpsilonChiantiSolver, UpsilonCMFGENSolver, UpsilonRegemorterSolver, @@ -41,13 +41,13 @@ def __init__( ) self.radiative_transitions = radiative_transitions # find the transitions that have radiative rate data but no collisional data - missing_collision_strengths_index = ( + self.missing_collision_strengths_index = ( radiative_transitions.index.difference( thermal_collisional_strengths.index ) ) self.all_collisional_strengths_index = ( - missing_collision_strengths_index.append( + self.missing_collision_strengths_index.append( thermal_collisional_strengths.index ).sort_values() ) @@ -75,9 +75,13 @@ def __init__( if collisional_strength_approximation == "regemorter": self.thermal_collision_strength_approximator = ( UpsilonRegemorterSolver( - radiative_transitions.loc[missing_collision_strengths_index] + radiative_transitions.loc[ + self.missing_collision_strengths_index + ] ) ) + else: + self.thermal_collision_strength_approximator = None def solve(self, temperatures_electron): thermal_all_collision_strengths = self.calculate_collision_strengths( @@ -116,6 +120,27 @@ def solve(self, temperatures_electron): "level_number_source", "level_number_destination", ] + + collision_rates_coeff_df = collision_rates_coeff_df.reset_index() + + # Add the new columns by duplicating the ion_number column + collision_rates_coeff_df["ion_number_source"] = ( + collision_rates_coeff_df["ion_number"] + ) + collision_rates_coeff_df["ion_number_destination"] = ( + collision_rates_coeff_df["ion_number"] + ) + + collision_rates_coeff_df = collision_rates_coeff_df.set_index( + [ + "atomic_number", + "ion_number", + "ion_number_source", + "ion_number_destination", + "level_number_source", + "level_number_destination", + ] + ) return collision_rates_coeff_df def calculate_collision_strengths(self, temperatures_electron): @@ -135,11 +160,20 @@ def calculate_collision_strengths(self, temperatures_electron): thermal_collision_strengths = ( self.thermal_collision_strength_solver.solve(temperatures_electron) ) - thermal_collision_strength_approximated = ( - self.thermal_collision_strength_approximator.solve( - temperatures_electron + + if self.thermal_collision_strength_approximator is None: + thermal_collision_strength_approximated = pd.DataFrame( + 0, + index=self.missing_collision_strengths_index, + columns=np.arange(len(temperatures_electron)), + dtype=np.float64, + ) + else: + thermal_collision_strength_approximated = ( + self.thermal_collision_strength_approximator.solve( + temperatures_electron + ) ) - ) return pd.concat( [ diff --git a/tardis/plasma/detailed_balance/rates/radiative_rates.py b/tardis/plasma/equilibrium/rates/radiative_rates.py similarity index 76% rename from tardis/plasma/detailed_balance/rates/radiative_rates.py rename to tardis/plasma/equilibrium/rates/radiative_rates.py index ab6fb22b444..248383bb505 100644 --- a/tardis/plasma/detailed_balance/rates/radiative_rates.py +++ b/tardis/plasma/equilibrium/rates/radiative_rates.py @@ -54,4 +54,22 @@ def solve(self, radiation_field): "level_number_source", "level_number_destination", ] + + rates_df = rates_df.reset_index() + + # Add the new columns by duplicating the ion_number column + rates_df["ion_number_source"] = rates_df["ion_number"] + rates_df["ion_number_destination"] = rates_df["ion_number"] + + rates_df = rates_df.set_index( + [ + "atomic_number", + "ion_number", + "ion_number_source", + "ion_number_destination", + "level_number_source", + "level_number_destination", + ] + ) + return rates_df diff --git a/tardis/plasma/detailed_balance/tests/__init__.py b/tardis/plasma/equilibrium/tests/__init__.py similarity index 100% rename from tardis/plasma/detailed_balance/tests/__init__.py rename to tardis/plasma/equilibrium/tests/__init__.py diff --git a/tardis/plasma/equilibrium/tests/conftest.py b/tardis/plasma/equilibrium/tests/conftest.py new file mode 100644 index 00000000000..c915e6a0784 --- /dev/null +++ b/tardis/plasma/equilibrium/tests/conftest.py @@ -0,0 +1,70 @@ +from pathlib import Path + +import pytest + +from tardis.io.atom_data import AtomData +from tardis.io.configuration.config_reader import Configuration +from tardis.model.base import SimulationState +from tardis.plasma.equilibrium.rates import ( + RadiativeRatesSolver, + ThermalCollisionalRateSolver, +) + + +@pytest.fixture +def new_chianti_atomic_dataset_si(tardis_regression_path): + atomic_data_fname = ( + tardis_regression_path / "atom_data" / "kurucz_cd23_chianti_Si.h5" + ) + return AtomData.from_hdf(atomic_data_fname) + + +@pytest.fixture(params=[(14, 1, slice(None), slice(None))]) +def radiative_transitions(new_chianti_atomic_dataset_si, request): + return new_chianti_atomic_dataset_si.lines.loc[request.param, :] + + +@pytest.fixture +def radiative_rate_solver(radiative_transitions): + return RadiativeRatesSolver(radiative_transitions) + + +@pytest.fixture(params=[(14, 1, slice(None), slice(None))]) +def collisional_rate_solver( + new_chianti_atomic_dataset_si, radiative_transitions, request +): + col_strength_temperatures = ( + new_chianti_atomic_dataset_si.collision_data_temperatures + ) + col_strengths = new_chianti_atomic_dataset_si.collision_data.loc[ + request.param, : + ] + return ThermalCollisionalRateSolver( + new_chianti_atomic_dataset_si.levels, + radiative_transitions, + col_strength_temperatures, + col_strengths, + "chianti", + ) + + +@pytest.fixture +def rate_solver_list(radiative_rate_solver, collisional_rate_solver): + return [ + (radiative_rate_solver, "radiative"), + (collisional_rate_solver, "electron"), + ] + + +@pytest.fixture +def collisional_simulation_state(new_chianti_atomic_dataset_si): + config = Configuration.from_yaml( + Path("tardis") + / "plasma" + / "tests" + / "data" + / "plasma_base_test_config.yml" + ) + return SimulationState.from_config( + config, atom_data=new_chianti_atomic_dataset_si + ) diff --git a/tardis/plasma/detailed_balance/tests/test_collisional_transitions.py b/tardis/plasma/equilibrium/tests/test_collisional_transitions.py similarity index 94% rename from tardis/plasma/detailed_balance/tests/test_collisional_transitions.py rename to tardis/plasma/equilibrium/tests/test_collisional_transitions.py index 9b5d80ea574..8cf0d371999 100644 --- a/tardis/plasma/detailed_balance/tests/test_collisional_transitions.py +++ b/tardis/plasma/equilibrium/tests/test_collisional_transitions.py @@ -12,7 +12,7 @@ PlasmaSolverFactory, convert_species_to_multi_index, ) -from tardis.plasma.detailed_balance.rates import ( +from tardis.plasma.equilibrium.rates import ( # UpsilonCMFGENSolver, ThermalCollisionalRateSolver, # RadiativeRatesSolver, @@ -141,18 +141,24 @@ def test_thermal_collision_rates( collision_strengths_type="cmfgen", collisional_strength_approximation="regemorter", ) - coll_rates_coeff = therm_coll_rate_solver.solve([10000, 20000] * u.K) + coll_rates_coeff = therm_coll_rate_solver.solve([10000, 20000] * u.K).loc[ + :, :, 0, 0, :, : + ] pdt.assert_frame_equal( coll_rates_coeff.iloc[:435], legacy_cmfgen_collision_rate_plasma_solver.coll_exc_coeff, + # to allow comparison between old and new versions check_names=False, + check_column_type=False, ) pdt.assert_frame_equal( coll_rates_coeff.iloc[435:], legacy_cmfgen_collision_rate_plasma_solver.coll_deexc_coeff.swaplevel( "level_number_lower", "level_number_upper" ), + # to allow comparison between old and new versions check_names=False, + check_column_type=False, ) @@ -184,7 +190,7 @@ def test_legacy_chianti_collisional_strengths( npt.assert_allclose( collision_strengths[0, 1, :], - chianti_collisional_rates.loc[1, 0, 1, 0], + chianti_collisional_rates.loc[1, 0, 0, 0, 1, 0], rtol=1e-4, atol=1e-13, ) diff --git a/tardis/plasma/equilibrium/tests/test_level_populations.py b/tardis/plasma/equilibrium/tests/test_level_populations.py new file mode 100644 index 00000000000..203f6bc2a43 --- /dev/null +++ b/tardis/plasma/equilibrium/tests/test_level_populations.py @@ -0,0 +1,73 @@ +import astropy.units as u +import numpy as np +import pandas as pd +import pytest + +from tardis.plasma.electron_energy_distribution import ( + ThermalElectronEnergyDistribution, +) +from tardis.plasma.equilibrium.level_populations import LevelPopulationSolver +from tardis.plasma.equilibrium.rate_matrix import RateMatrix +from tardis.plasma.radiation_field import ( + DilutePlanckianRadiationField, +) + + +class TestLevelPopulationSolver: + @pytest.fixture(autouse=True) + def setup( + self, + rate_solver_list, + new_chianti_atomic_dataset_si, + collisional_simulation_state, + ): + rate_matrix_solver = RateMatrix( + rate_solver_list, new_chianti_atomic_dataset_si.levels + ) + + rad_field = DilutePlanckianRadiationField( + collisional_simulation_state.t_radiative, + dilution_factor=np.zeros_like( + collisional_simulation_state.t_radiative + ), + ) + electron_dist = ThermalElectronEnergyDistribution( + 0, collisional_simulation_state.t_radiative, 1e6 * u.g / u.cm**3 + ) + + rates_matrices = rate_matrix_solver.solve(rad_field, electron_dist) + self.solver = LevelPopulationSolver( + rates_matrices, new_chianti_atomic_dataset_si.levels + ) + + def test_calculate_level_population_simple(self): + """Test solving a 2-level ion.""" + rates_matrix = np.array([[1, 1], [2, -2]]) + expected_population = np.array([0.5, 0.5]) + result = self.solver._LevelPopulationSolver__calculate_level_population( + rates_matrix + ) + np.testing.assert_array_almost_equal(result, expected_population) + + def test_calculate_level_population_empty(self): + """Test empty rate matrix.""" + rates_matrix = np.array([[]]) + with pytest.raises(np.linalg.LinAlgError): + self.solver._LevelPopulationSolver__calculate_level_population( + rates_matrix + ) + + def test_calculate_level_population_zeros(self): + """Test zero rate matrix.""" + rates_matrix = np.array([[0, 0], [0, 0]]) + with pytest.raises(np.linalg.LinAlgError): + self.solver._LevelPopulationSolver__calculate_level_population( + rates_matrix + ) + + def test_solve(self, regression_data): + """Test the solve method.""" + expected_populations = False + result = self.solver.solve() + expected_populations = regression_data.sync_dataframe(result) + pd.testing.assert_frame_equal(result, expected_populations) diff --git a/tardis/plasma/equilibrium/tests/test_rate_matrix.py b/tardis/plasma/equilibrium/tests/test_rate_matrix.py new file mode 100644 index 00000000000..d2897f1eea8 --- /dev/null +++ b/tardis/plasma/equilibrium/tests/test_rate_matrix.py @@ -0,0 +1,36 @@ +import astropy.units as u +import numpy as np +import pandas.testing as pdt + +from tardis.plasma.electron_energy_distribution import ( + ThermalElectronEnergyDistribution, +) +from tardis.plasma.equilibrium.rate_matrix import RateMatrix +from tardis.plasma.radiation_field import ( + DilutePlanckianRadiationField, +) + + +def test_rate_matrix_solver( + new_chianti_atomic_dataset_si, + rate_solver_list, + collisional_simulation_state, + regression_data, +): + rate_matrix_solver = RateMatrix( + rate_solver_list, new_chianti_atomic_dataset_si.levels + ) + + rad_field = DilutePlanckianRadiationField( + collisional_simulation_state.t_radiative, + dilution_factor=np.zeros_like(collisional_simulation_state.t_radiative), + ) + electron_dist = ThermalElectronEnergyDistribution( + 0, collisional_simulation_state.t_radiative, 1e6 * u.g / u.cm**3 + ) + + actual = rate_matrix_solver.solve(rad_field, electron_dist) + + expected = regression_data.sync_dataframe(actual) + + pdt.assert_frame_equal(actual, expected)