Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rates matrix solver #2865

Open
wants to merge 24 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 17 commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
956d172
Rename detailed_balance to equilibrium
andrewfullard Oct 14, 2024
7bfe4e2
Add ion number source and destination as index columns
andrewfullard Oct 14, 2024
5f8b512
Rates matrix solver
andrewfullard Oct 22, 2024
b37d575
Normalized level population solver (not working)
andrewfullard Oct 22, 2024
83f1ff1
Correct shape of assignment
andrewfullard Oct 22, 2024
d5e8f4f
Correctly updating the dataframe (still chained assignment)
andrewfullard Oct 22, 2024
a37abb0
Adds docstrings
andrewfullard Oct 22, 2024
e8b167f
Fixes level population indexing and add TODO for the rate matrix solver
andrewfullard Oct 28, 2024
a9240d2
Refactor of the rate matrix solver to adapt to different inputs
andrewfullard Oct 28, 2024
5d44378
Add consistent
andrewfullard Oct 28, 2024
09cc117
Adds support for no approximation of collisional rates
andrewfullard Oct 29, 2024
10cae7e
Fix transformation of rate matrix
andrewfullard Oct 29, 2024
e33fe28
Make electron_distribution a dataclass
andrewfullard Nov 5, 2024
30d6f2f
Better name for rate matrix variable
andrewfullard Nov 5, 2024
1dd7e79
Adds simple tests for level pop solver
andrewfullard Nov 5, 2024
3363872
Correct index in col strength test
andrewfullard Nov 5, 2024
a1a2352
Resolve test indexing issue
andrewfullard Nov 5, 2024
4e6e10d
Specify CGS for the electron distribution
andrewfullard Nov 5, 2024
2f18918
Rename electron_distribution to electron_energy_distribution
andrewfullard Nov 12, 2024
670fae4
Adds basic regression data test to the level population solver
andrewfullard Nov 12, 2024
d9444cb
Adds proper rate matrix and level population solver tests
andrewfullard Nov 12, 2024
1b3659f
Rename electron energy dist. input
andrewfullard Nov 18, 2024
b0f91ef
Adds example notebooks and comment on tests
andrewfullard Nov 18, 2024
4bd4206
Remove outdated notebook
andrewfullard Nov 18, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this should be .h5

File renamed without changes.
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
10 changes: 0 additions & 10 deletions tardis/plasma/detailed_balance/rates/__init__.py

This file was deleted.

1 change: 1 addition & 0 deletions tardis/plasma/electron_distribution/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from tardis.plasma.electron_distribution.base import ElectronDistribution
17 changes: 17 additions & 0 deletions tardis/plasma/electron_distribution/base.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
from dataclasses import dataclass

Check warning on line 1 in tardis/plasma/electron_distribution/base.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/electron_distribution/base.py#L1

Added line #L1 was not covered by tests

from astropy import units as u

Check warning on line 3 in tardis/plasma/electron_distribution/base.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/electron_distribution/base.py#L3

Added line #L3 was not covered by tests


@dataclass
class ElectronDistribution:

Check warning on line 7 in tardis/plasma/electron_distribution/base.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/electron_distribution/base.py#L6-L7

Added lines #L6 - L7 were not covered by tests
andrewfullard marked this conversation as resolved.
Show resolved Hide resolved
"""Electron temperature and density distribution.

temperature : Quantity
Electron temperatures.
number_density : Quantity
andrewfullard marked this conversation as resolved.
Show resolved Hide resolved
Electron number densities.
"""

temperature: u.Quantity
number_density: u.Quantity

Check warning on line 17 in tardis/plasma/electron_distribution/base.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/electron_distribution/base.py#L16-L17

Added lines #L16 - L17 were not covered by tests
69 changes: 69 additions & 0 deletions tardis/plasma/equilibrium/level_populations.py
Original file line number Diff line number Diff line change
@@ -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(

Check warning on line 49 in tardis/plasma/equilibrium/level_populations.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/level_populations.py#L49

Added line #L49 was not covered by tests
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:

Check warning on line 56 in tardis/plasma/equilibrium/level_populations.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/level_populations.py#L56

Added line #L56 was not covered by tests
# TODO: resolve the chained assignment here. Maybe an intermediate df
# is needed

solved_matrices = self.rates_matrices.loc[species_id].apply(

Check warning on line 60 in tardis/plasma/equilibrium/level_populations.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/level_populations.py#L60

Added line #L60 was not covered by tests
lambda rates_matrix: self.__calculate_level_population(
rates_matrix
)
)
normalized_level_populations.loc[species_id, :].update(

Check warning on line 65 in tardis/plasma/equilibrium/level_populations.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/level_populations.py#L65

Added line #L65 was not covered by tests
np.vstack(solved_matrices.values).T
)

return normalized_level_populations

Check warning on line 69 in tardis/plasma/equilibrium/level_populations.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/level_populations.py#L69

Added line #L69 was not covered by tests
106 changes: 106 additions & 0 deletions tardis/plasma/equilibrium/rate_matrix.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
import numpy as np
import pandas as pd
from scipy.sparse import coo_matrix

Check warning on line 3 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L1-L3

Added lines #L1 - L3 were not covered by tests


class RateMatrix:
def __init__(

Check warning on line 7 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L6-L7

Added lines #L6 - L7 were not covered by tests
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

Check warning on line 22 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L21-L22

Added lines #L21 - L22 were not covered by tests

def solve(

Check warning on line 24 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L24

Added line #L24 was not covered by tests
self,
radiation_field,
electron_distribution,
andrewfullard marked this conversation as resolved.
Show resolved Hide resolved
):
"""Construct the compiled rate matrix dataframe.

Parameters
----------
radiation_field : RadiationField
Radiation field containing radiative temperature.
electron_distribution : ElectronDistribution
Distribution of electrons in the plasma, containing electron
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 = {

Check warning on line 45 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L45

Added line #L45 was not covered by tests
"radiative": radiation_field,
"electron": electron_distribution.temperature,
}

rates_df_list = [

Check warning on line 50 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L50

Added line #L50 was not covered by tests
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)

Check warning on line 56 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L54-L56

Added lines #L54 - L56 were not covered by tests

# Create a union of all indexes
all_indexes = sorted(all_indexes)

Check warning on line 59 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L59

Added line #L59 was not covered by tests

# Reindex each dataframe to ensure consistent indices
rates_df_list = [

Check warning on line 62 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L62

Added line #L62 was not covered by tests
df.reindex(all_indexes, fill_value=0) for df in rates_df_list
]

# Multiply rates by electron number density where appropriate
rates_df_list = [

Check warning on line 67 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L67

Added line #L67 was not covered by tests
rates_df * electron_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)

Check warning on line 76 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L76

Added line #L76 was not covered by tests

grouped_rates_df = rates_df.groupby(

Check warning on line 78 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L78

Added line #L78 was not covered by tests
level=("atomic_number", "ion_number")
)

rate_matrices = pd.DataFrame(

Check warning on line 82 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L82

Added line #L82 was not covered by tests
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(

Check warning on line 89 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L86-L89

Added lines #L86 - L89 were not covered by tests
(
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

Check warning on line 104 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L101-L104

Added lines #L101 - L104 were not covered by tests

return rate_matrices

Check warning on line 106 in tardis/plasma/equilibrium/rate_matrix.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rate_matrix.py#L106

Added line #L106 was not covered by tests
10 changes: 10 additions & 0 deletions tardis/plasma/equilibrium/rates/__init__.py
Original file line number Diff line number Diff line change
@@ -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,
)
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -41,13 +41,13 @@
)
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()
)
Expand Down Expand Up @@ -75,9 +75,13 @@
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

Check warning on line 84 in tardis/plasma/equilibrium/rates/collisional_rates.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rates/collisional_rates.py#L84

Added line #L84 was not covered by tests

def solve(self, temperatures_electron):
thermal_all_collision_strengths = self.calculate_collision_strengths(
Expand Down Expand Up @@ -116,6 +120,27 @@
"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):
Expand All @@ -135,11 +160,20 @@
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(

Check warning on line 165 in tardis/plasma/equilibrium/rates/collisional_rates.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rates/collisional_rates.py#L165

Added line #L165 was not covered by tests
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(
[
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,4 +54,22 @@
"level_number_source",
"level_number_destination",
]

rates_df = rates_df.reset_index()

Check warning on line 58 in tardis/plasma/equilibrium/rates/radiative_rates.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rates/radiative_rates.py#L58

Added line #L58 was not covered by tests

# 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"]

Check warning on line 62 in tardis/plasma/equilibrium/rates/radiative_rates.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rates/radiative_rates.py#L61-L62

Added lines #L61 - L62 were not covered by tests

rates_df = rates_df.set_index(

Check warning on line 64 in tardis/plasma/equilibrium/rates/radiative_rates.py

View check run for this annotation

Codecov / codecov/patch

tardis/plasma/equilibrium/rates/radiative_rates.py#L64

Added line #L64 was not covered by tests
[
"atomic_number",
"ion_number",
"ion_number_source",
"ion_number_destination",
"level_number_source",
"level_number_destination",
]
)

return rates_df
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -141,18 +141,22 @@ 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,
check_names=False,
check_column_type=False,
andrewfullard marked this conversation as resolved.
Show resolved Hide resolved
)
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"
),
check_names=False,
check_column_type=False,
)


Expand Down Expand Up @@ -184,7 +188,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,
)
Loading
Loading