-
-
Notifications
You must be signed in to change notification settings - Fork 407
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
andrewfullard
wants to merge
24
commits into
tardis-sn:master
Choose a base branch
from
andrewfullard:rates-matrix-solver
base: master
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+868
−72
Open
Rates matrix solver #2865
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 7bfe4e2
Add ion number source and destination as index columns
andrewfullard 5f8b512
Rates matrix solver
andrewfullard b37d575
Normalized level population solver (not working)
andrewfullard 83f1ff1
Correct shape of assignment
andrewfullard d5e8f4f
Correctly updating the dataframe (still chained assignment)
andrewfullard a37abb0
Adds docstrings
andrewfullard e8b167f
Fixes level population indexing and add TODO for the rate matrix solver
andrewfullard a9240d2
Refactor of the rate matrix solver to adapt to different inputs
andrewfullard 5d44378
Add consistent
andrewfullard 09cc117
Adds support for no approximation of collisional rates
andrewfullard 10cae7e
Fix transformation of rate matrix
andrewfullard e33fe28
Make electron_distribution a dataclass
andrewfullard 30d6f2f
Better name for rate matrix variable
andrewfullard 1dd7e79
Adds simple tests for level pop solver
andrewfullard 3363872
Correct index in col strength test
andrewfullard a1a2352
Resolve test indexing issue
andrewfullard 4e6e10d
Specify CGS for the electron distribution
andrewfullard 2f18918
Rename electron_distribution to electron_energy_distribution
andrewfullard 670fae4
Adds basic regression data test to the level population solver
andrewfullard d9444cb
Adds proper rate matrix and level population solver tests
andrewfullard 1b3659f
Rename electron energy dist. input
andrewfullard b0f91ef
Adds example notebooks and comment on tests
andrewfullard 4bd4206
Remove outdated notebook
andrewfullard File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
File renamed without changes.
File renamed without changes.
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
from tardis.plasma.electron_distribution.base import ElectronDistribution |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,17 @@ | ||
from dataclasses import dataclass | ||
|
||
from astropy import units as u | ||
|
||
|
||
@dataclass | ||
class ElectronDistribution: | ||
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 | ||
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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( | ||
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 | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | ||
|
||
|
||
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, | ||
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 = { | ||
"radiative": radiation_field, | ||
"electron": electron_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 * 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) | ||
|
||
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 | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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, | ||
) |
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
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