Rates matrix solver (#2865)

* Rename detailed_balance to equilibrium

* Add ion number source and destination as index columns

* Rates matrix solver

* Normalized level population solver (not working)

* Correct shape of assignment

* Correctly updating the dataframe (still chained assignment)

* Adds docstrings

* Fixes level population indexing and add TODO for the rate matrix solver

* Refactor of the rate matrix solver to adapt to different inputs

Also include electron distribution class to hold relevant info.

* Add consistent

* Adds support for no approximation of collisional rates

* Fix transformation of rate matrix

* Make electron_distribution a dataclass

* Better name for rate matrix variable

* Adds simple tests for level pop solver

* Correct index in col strength test

* Resolve test indexing issue

* Specify CGS for the electron distribution

* Rename electron_distribution to electron_energy_distribution

* Adds basic regression data test to the level population solver

* Adds proper rate matrix and level population solver tests

* Rename electron energy dist. input

* Adds example notebooks and comment on tests

* Remove outdated notebook

* Address comment

* Pre-run tardis solver workflow

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
+}

docs/physics/plasma/detailed_balance/comparison.ipynb → 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",
@@ -128,7 +128,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -144,7 +144,7 @@
     "\n",
     "chianti_sim_state = SimulationState.from_config(config, atom_data=chianti_atom_data)\n",
     "\n",
-    "reference_coeff = pd.HDFStore(\"collisional_rate_coefficients.hdf\")\n",
+    "reference_coeff = pd.HDFStore(\"collisional_rate_coefficients.h5\")\n",
     "reference_rate_coeff_df = pd.concat([reference_coeff[\"coll_exc_coeff\"], reference_coeff[\"coll_deexc_coeff\"].reorder_levels([0,1,3,2])])\n",
     "\n",
     "temperature = reference_coeff[\"t_electrons\"].values * u.K\n",