From d5835104869e45fcb438a8597d6ae0419faabc3d Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 30 Apr 2024 17:38:31 -0400 Subject: [PATCH 01/89] change numberdensity to input --- tardis/plasma/properties/general.py | 18 ----- tardis/plasma/properties/plasma_input.py | 18 ++++- .../plasma/properties/property_collections.py | 2 +- tardis/plasma/standard_plasmas.py | 81 +++++++++---------- tardis/plasma/tests/test_complete_plasmas.py | 8 +- 5 files changed, 61 insertions(+), 66 deletions(-) diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index 9988f656411..e1a41306077 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -12,7 +12,6 @@ __all__ = [ "BetaRadiation", "GElectron", - "NumberDensity", "IsotopeNumberDensity", "SelectedAtoms", "ElectronTemperature", @@ -79,23 +78,6 @@ def calculate(self, beta_electron): return super(ThermalGElectron, self).calculate(beta_electron) -class NumberDensity(ProcessingPlasmaProperty): - """ - Attributes - ---------- - number_density : Pandas DataFrame, dtype float - Indexed by atomic number, columns corresponding to zones - """ - - outputs = ("number_density",) - latex_name = ("N_{i}",) - - @staticmethod - def calculate(atomic_mass, abundance, density): - number_densities = abundance * density - return number_densities.div(atomic_mass.loc[abundance.index], axis=0) - - class IsotopeNumberDensity(ProcessingPlasmaProperty): """ Calculate the atom number density based on isotope mass. diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index 2d7b67ca4d6..288e44042e7 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -1,10 +1,14 @@ -from tardis.plasma.properties.base import Input, ArrayInput, DataFrameInput +from tardis.plasma.properties.base import ( + ArrayInput, + Input, +) __all__ = [ "TRadiative", "DilutionFactor", "AtomicData", "Abundance", + "NumberDensity", "IsotopeAbundance", "Density", "TimeExplosion", @@ -160,3 +164,15 @@ class NLTEIonizationSpecies(Input): class NLTEExcitationSpecies(Input): outputs = ("nlte_excitation_species",) + + +class NumberDensity(Input): + """ + Attributes + ---------- + number_density : Pandas DataFrame, dtype float + Indexed by atomic number, columns corresponding to zones + """ + + outputs = ("number_density",) + latex_name = ("N_{i}",) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index deb2405347b..8565d740e04 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -10,6 +10,7 @@ class PlasmaPropertyCollection(list): TRadiative, Abundance, Density, + NumberDensity, TimeExplosion, AtomicData, DilutionFactor, @@ -29,7 +30,6 @@ class PlasmaPropertyCollection(list): PartitionFunction, GElectron, IonizationData, - NumberDensity, LinesLowerLevelIndex, LinesUpperLevelIndex, TauSobolev, diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 42d81af1e1b..db8ca8876ee 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -1,57 +1,53 @@ -from astropy import units as u -import os import logging import numpy as np import pandas as pd +from astropy import units as u -from tardis.io.atom_data import AtomData +from tardis.plasma import BasePlasma +from tardis.plasma.exceptions import PlasmaConfigError +from tardis.plasma.properties import ( + HeliumNumericalNLTE, + IonNumberDensity, + IonNumberDensityHeNLTE, + JBluesBlackBody, + JBluesDetailed, + JBluesDiluteBlackBody, + LevelBoltzmannFactorNLTE, + MarkovChainTransProbsCollector, + RadiationFieldCorrection, + StimulatedEmissionFactor, +) +from tardis.plasma.properties.base import TransitionProbabilitiesProperty from tardis.plasma.properties.level_population import LevelNumberDensity from tardis.plasma.properties.nlte_rate_equation_solver import ( NLTEPopulationSolverLU, NLTEPopulationSolverRoot, ) -from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper -from tardis.util.base import species_string_to_tuple -from tardis.plasma import BasePlasma -from tardis.plasma.properties.base import TransitionProbabilitiesProperty from tardis.plasma.properties.property_collections import ( + adiabatic_cooling_properties, basic_inputs, basic_properties, + continuum_interaction_inputs, + continuum_interaction_properties, + detailed_j_blues_inputs, + detailed_j_blues_properties, + dilute_lte_excitation_properties, + helium_lte_properties, + helium_nlte_properties, + helium_numerical_nlte_properties, lte_excitation_properties, lte_ionization_properties, macro_atom_properties, - dilute_lte_excitation_properties, nebular_ionization_properties, - non_nlte_properties, - nlte_properties, - helium_nlte_properties, - helium_numerical_nlte_properties, - helium_lte_properties, - detailed_j_blues_properties, - detailed_j_blues_inputs, - continuum_interaction_properties, - continuum_interaction_inputs, - adiabatic_cooling_properties, - two_photon_properties, - isotope_properties, nlte_lu_solver_properties, + nlte_properties, nlte_root_solver_properties, + non_nlte_properties, + two_photon_properties, ) -from tardis.plasma.exceptions import PlasmaConfigError - -from tardis.plasma.properties import ( - LevelBoltzmannFactorNLTE, - JBluesBlackBody, - JBluesDiluteBlackBody, - JBluesDetailed, - RadiationFieldCorrection, - StimulatedEmissionFactor, - HeliumNumericalNLTE, - IonNumberDensity, - IonNumberDensityHeNLTE, - MarkovChainTransProbsCollector, -) +from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper +from tardis.util.base import species_string_to_tuple logger = logging.getLogger(__name__) @@ -122,6 +118,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): t_rad=simulation_state.t_radiative, abundance=simulation_state.abundance, density=simulation_state.density, + number_density=simulation_state.composition.elemental_number_density, atomic_data=atom_data, time_explosion=simulation_state.time_explosion, w=simulation_state.dilution_factor, @@ -138,7 +135,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): if line_interaction_type != "macroatom": raise PlasmaConfigError( "Continuum interactions require line_interaction_type " - "macroatom (instead of {}).".format(line_interaction_type) + f"macroatom (instead of {line_interaction_type})." ) plasma_modules += continuum_interaction_properties @@ -171,8 +168,9 @@ def assemble_plasma(config, simulation_state, atom_data=None): if config.plasma.nlte_ionization_species: nlte_ionization_species = config.plasma.nlte_ionization_species for species in nlte_ionization_species: - if not ( - species in config.plasma.continuum_interaction.species + if ( + species + not in config.plasma.continuum_interaction.species ): raise PlasmaConfigError( f"NLTE ionization species {species} not in continuum species." @@ -180,8 +178,9 @@ def assemble_plasma(config, simulation_state, atom_data=None): if config.plasma.nlte_excitation_species: nlte_excitation_species = config.plasma.nlte_excitation_species for species in nlte_excitation_species: - if not ( - species in config.plasma.continuum_interaction.species + if ( + species + not in config.plasma.continuum_interaction.species ): raise PlasmaConfigError( f"NLTE excitation species {species} not in continuum species." @@ -199,9 +198,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): plasma_modules += nlte_root_solver_properties else: raise PlasmaConfigError( - "NLTE solver type unknown - {}".format( - config.plasma.nlte_solver - ) + f"NLTE solver type unknown - {config.plasma.nlte_solver}" ) kwargs.update( diff --git a/tardis/plasma/tests/test_complete_plasmas.py b/tardis/plasma/tests/test_complete_plasmas.py index ed7d4428852..a28aeba63e8 100644 --- a/tardis/plasma/tests/test_complete_plasmas.py +++ b/tardis/plasma/tests/test_complete_plasmas.py @@ -193,7 +193,7 @@ def test_plasma_properties(self, plasma, attr): def test_levels(self, plasma): actual = pd.DataFrame(plasma.levels) - key = f"plasma/levels" + key = "plasma/levels" expected = pd.read_hdf(self.regression_data.fpath, key) pdt.assert_frame_equal(actual, expected) @@ -202,13 +202,13 @@ def test_scalars_properties(self, plasma, attr): actual = getattr(plasma, attr) if hasattr(actual, "cgs"): actual = actual.cgs.value - key = f"plasma/scalars" + key = "plasma/scalars" expected = pd.read_hdf(self.regression_data.fpath, key)[attr] npt.assert_equal(actual, expected) def test_helium_treatment(self, plasma): actual = plasma.helium_treatment - key = f"plasma/scalars" + key = "plasma/scalars" expected = pd.read_hdf(self.regression_data.fpath, key)[ "helium_treatment" ] @@ -217,6 +217,6 @@ def test_helium_treatment(self, plasma): def test_zeta_data(self, plasma): if hasattr(plasma, "zeta_data"): actual = plasma.zeta_data - key = f"plasma/zeta_data" + key = "plasma/zeta_data" expected = pd.read_hdf(self.regression_data.fpath, key) npt.assert_allclose(actual, expected.values) From 7aaba5d73904e5f1d0bd97cadb1437eda1f9e20f Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 4 May 2024 16:18:47 -0400 Subject: [PATCH 02/89] fixed number density --- tardis/model/base.py | 19 ++++++++ tardis/plasma/properties/general.py | 43 ------------------- .../plasma/properties/property_collections.py | 3 -- tardis/plasma/standard_plasmas.py | 2 +- 4 files changed, 20 insertions(+), 47 deletions(-) diff --git a/tardis/model/base.py b/tardis/model/base.py index cf4e2da7548..72d4c157a20 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -161,6 +161,25 @@ def t_radiative(self, new_t_radiative): "Trying to set t_radiative for different number of shells." ) + @property + def elemental_number_density(self): + elemental_number_density = ( + ( + self.composition.elemental_mass_fraction + * self.composition.density + ) + .divide(self.composition.element_masses, axis=0) + .dropna() + ) + elemental_number_density = elemental_number_density.iloc[ + :, + self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index, + ] + elemental_number_density.columns = range( + len(elemental_number_density.columns) + ) + return elemental_number_density + @property def radius(self): return self.time_explosion * self.velocity diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index e1a41306077..55823b0fe14 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -12,7 +12,6 @@ __all__ = [ "BetaRadiation", "GElectron", - "IsotopeNumberDensity", "SelectedAtoms", "ElectronTemperature", "BetaElectron", @@ -77,48 +76,6 @@ class ThermalGElectron(GElectron): def calculate(self, beta_electron): return super(ThermalGElectron, self).calculate(beta_electron) - -class IsotopeNumberDensity(ProcessingPlasmaProperty): - """ - Calculate the atom number density based on isotope mass. - - Attributes - ---------- - isotope_number_density : Pandas DataFrame, dtype float - Indexed by atomic number, columns corresponding to zones - """ - - outputs = ("isotope_number_density",) - latex_name = ("N_{i}",) - - @staticmethod - def calculate(isotope_mass, isotope_abundance, density): - """ - Calculate the atom number density based on isotope mass. - - Parameters - ---------- - isotope_mass : pandas.DataFrame - Masses of isotopes. - isotope_abundance : pandas.DataFrame - Fractional abundance of isotopes. - density : pandas.DataFrame - Density of each shell. - - Returns - ------- - pandas.DataFrame - Indexed by atomic number, columns corresponding to zones. - """ - number_densities = isotope_abundance * density - isotope_number_density_array = ( - number_densities.to_numpy() / isotope_mass.to_numpy() - ) - return pd.DataFrame( - isotope_number_density_array, index=isotope_abundance.index - ) - - class SelectedAtoms(ProcessingPlasmaProperty): """ Attributes diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index 8565d740e04..eb6c15da97e 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -153,6 +153,3 @@ class PlasmaPropertyCollection(list): TwoPhotonFrequencySampler, ] ) -isotope_properties = PlasmaPropertyCollection( - [IsotopeAbundance, IsotopeMass, IsotopeNumberDensity] -) diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index db8ca8876ee..40b11aca649 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -118,7 +118,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): t_rad=simulation_state.t_radiative, abundance=simulation_state.abundance, density=simulation_state.density, - number_density=simulation_state.composition.elemental_number_density, + number_density=simulation_state.elemental_number_density, atomic_data=atom_data, time_explosion=simulation_state.time_explosion, w=simulation_state.dilution_factor, From 32b22ea4a84e2f53bc52b568727759a4be57e675 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sun, 5 May 2024 19:55:06 -0400 Subject: [PATCH 03/89] some fixes --- tardis/plasma/properties/atomic.py | 2 -- tardis/tests/test_tardis_full_formal_integral.py | 5 +++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/tardis/plasma/properties/atomic.py b/tardis/plasma/properties/atomic.py index f206264d3e4..f5153baf697 100644 --- a/tardis/plasma/properties/atomic.py +++ b/tardis/plasma/properties/atomic.py @@ -76,8 +76,6 @@ def _filter_atomic_property(self, levels, selected_atoms): return levels[levels.index.isin(selected_atoms, level="atomic_number")] def _set_index(self, levels): - # levels = levels.set_index(['atomic_number', 'ion_number', - # 'level_number']) return ( levels.index, levels["energy"], diff --git a/tardis/tests/test_tardis_full_formal_integral.py b/tardis/tests/test_tardis_full_formal_integral.py index f28203045f1..d3a0eeaecc5 100644 --- a/tardis/tests/test_tardis_full_formal_integral.py +++ b/tardis/tests/test_tardis_full_formal_integral.py @@ -1,11 +1,12 @@ from pathlib import Path -import pytest + import numpy.testing as npt +import pytest from astropy import units as u from astropy.tests.helper import assert_quantity_allclose -from tardis.simulation.base import Simulation from tardis.io.configuration.config_reader import Configuration +from tardis.simulation.base import Simulation config_line_modes = ["downbranch", "macroatom"] interpolate_shells = [-1, 30] From 261bbad532567b949a66d2db760e824a0b14d538 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sun, 5 May 2024 20:14:04 -0400 Subject: [PATCH 04/89] removing density --- tardis/plasma/properties/property_collections.py | 1 - tardis/plasma/standard_plasmas.py | 1 - tardis/plasma/tests/test_hdf_plasma.py | 2 +- 3 files changed, 1 insertion(+), 3 deletions(-) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index eb6c15da97e..a266c1618db 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -9,7 +9,6 @@ class PlasmaPropertyCollection(list): [ TRadiative, Abundance, - Density, NumberDensity, TimeExplosion, AtomicData, diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 40b11aca649..da772508140 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -117,7 +117,6 @@ def assemble_plasma(config, simulation_state, atom_data=None): kwargs = dict( t_rad=simulation_state.t_radiative, abundance=simulation_state.abundance, - density=simulation_state.density, number_density=simulation_state.elemental_number_density, atomic_data=atom_data, time_explosion=simulation_state.time_explosion, diff --git a/tardis/plasma/tests/test_hdf_plasma.py b/tardis/plasma/tests/test_hdf_plasma.py index d0c4f4c6433..13169bdfa4e 100644 --- a/tardis/plasma/tests/test_hdf_plasma.py +++ b/tardis/plasma/tests/test_hdf_plasma.py @@ -86,7 +86,7 @@ def test_atomic_data_uuid(simulation_verysimple, regression_data): assert actual == expected -COLLECTION_PROPERTIES = ["t_rad", "w", "density"] +COLLECTION_PROPERTIES = ["t_rad", "w"] @pytest.mark.parametrize("attr", COLLECTION_PROPERTIES) From df3c406c27fafa5c6eeb13cbe7c4b7f067251153 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sun, 5 May 2024 20:28:09 -0400 Subject: [PATCH 05/89] remove atomic and isotope mass --- tardis/plasma/properties/atomic.py | 66 ------------------- .../plasma/properties/property_collections.py | 1 - 2 files changed, 67 deletions(-) diff --git a/tardis/plasma/properties/atomic.py b/tardis/plasma/properties/atomic.py index f5153baf697..c917c08c421 100644 --- a/tardis/plasma/properties/atomic.py +++ b/tardis/plasma/properties/atomic.py @@ -31,8 +31,6 @@ "Lines", "LinesLowerLevelIndex", "LinesUpperLevelIndex", - "AtomicMass", - "IsotopeMass", "IonizationData", "ZetaData", "NLTEData", @@ -538,70 +536,6 @@ def calculate(self, level_idxs2line_idx, level_idxs2continuum_idx): return level_idxs2transition_idx - -class AtomicMass(ProcessingPlasmaProperty): - """ - Attributes - ---------- - atomic_mass : pandas.Series - Atomic masses of the elements used. Indexed by atomic number. - """ - - outputs = ("atomic_mass",) - - def calculate(self, atomic_data, selected_atoms): - if getattr(self, self.outputs[0]) is not None: - return (getattr(self, self.outputs[0]),) - else: - return atomic_data.atom_data.loc[selected_atoms].mass - - -class IsotopeMass(ProcessingPlasmaProperty): - """ - Attributes - ---------- - isotope_mass : pandas.Series - Masses of the isotopes used. Indexed by isotope name e.g. 'Ni56'. - """ - - outputs = ("isotope_mass",) - - def calculate(self, isotope_abundance): - """ - Determine mass of each isotope. - - Parameters - ---------- - isotope_abundance : pandas.DataFrame - Fractional abundance of isotopes. - - Returns - ------- - pandas.DataFrame - Masses of the isotopes used. Indexed by isotope name e.g. 'Ni56'. - """ - if getattr(self, self.outputs[0]) is not None: - return (getattr(self, self.outputs[0]),) - else: - if isotope_abundance.empty: - return None - isotope_mass_dict = {} - for Z, A in isotope_abundance.index: - element_name = rd.utils.Z_to_elem(Z) - isotope_name = element_name + str(A) - - isotope_mass_dict[(Z, A)] = rd.Nuclide(isotope_name).atomic_mass - - isotope_mass_df = pd.DataFrame.from_dict( - isotope_mass_dict, orient="index", columns=["mass"] - ) - isotope_mass_df.index = pd.MultiIndex.from_tuples( - isotope_mass_df.index - ) - isotope_mass_df.index.names = ["atomic_number", "mass_number"] - return isotope_mass_df / const.N_A - - class IonizationData(BaseAtomicDataProperty): """ Attributes diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index a266c1618db..d93b6016442 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -25,7 +25,6 @@ class PlasmaPropertyCollection(list): BetaRadiation, Levels, Lines, - AtomicMass, PartitionFunction, GElectron, IonizationData, From d05e819353400c01f93d31c79aa41495fa97661a Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 6 May 2024 14:22:13 -0400 Subject: [PATCH 06/89] add isotopic_number_density --- tardis/model/base.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tardis/model/base.py b/tardis/model/base.py index 72d4c157a20..fbc15f3f64b 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -180,6 +180,26 @@ def elemental_number_density(self): ) return elemental_number_density + @property + def isotopic_number_density(self): + isotopic_number_density = ( + self.composition.isotopic_mass_fraction * self.composition.density + ).divide( + self.composition.isotope_masses.loc[ + self.composition.isotopic_mass_fraction.index + ] + * u.u.to(u.g), + axis=0, + ) + isotopic_number_density = isotopic_number_density.iloc[ + :, + self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index, + ] + isotopic_number_density.columns = range( + len(isotopic_number_density.columns) + ) + return isotopic_number_density + @property def radius(self): return self.time_explosion * self.velocity From e8049426d4dc46d6f13e09b456debfb013598607 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 6 May 2024 16:11:43 -0400 Subject: [PATCH 07/89] add opacities package --- benchmarks/transport_montecarlo_opacities.py | 2 +- benchmarks/transport_montecarlo_packet.py | 7 +- tardis/energy_input/gamma_packet_loop.py | 35 +++--- tardis/energy_input/gamma_ray_estimators.py | 19 ++-- tardis/energy_input/gamma_ray_interactions.py | 26 ++--- tardis/energy_input/gamma_ray_transport.py | 7 +- tardis/energy_input/tests/test_util.py | 2 +- tardis/energy_input/util.py | 13 ++- tardis/opacities/__init__.py | 0 .../montecarlo => opacities}/opacities.py | 0 tardis/opacities/tau_sobolev.py | 104 ++++++++++++++++++ .../plasma/properties/property_collections.py | 1 + .../plasma/properties/radiative_properties.py | 70 ------------ .../montecarlo/single_packet_loop.py | 16 +-- .../montecarlo/tests/test_opacities.py | 2 +- .../transport/montecarlo/tests/test_packet.py | 8 +- tardis/transport/montecarlo/vpacket.py | 39 +++---- tardis/transport/r_packet_transport.py | 14 +-- 18 files changed, 196 insertions(+), 169 deletions(-) create mode 100644 tardis/opacities/__init__.py rename tardis/{transport/montecarlo => opacities}/opacities.py (100%) create mode 100644 tardis/opacities/tau_sobolev.py diff --git a/benchmarks/transport_montecarlo_opacities.py b/benchmarks/transport_montecarlo_opacities.py index b56fed5a2e6..c5fe0196f47 100644 --- a/benchmarks/transport_montecarlo_opacities.py +++ b/benchmarks/transport_montecarlo_opacities.py @@ -4,7 +4,7 @@ from asv_runner.benchmarks.mark import parameterize -import tardis.transport.montecarlo.opacities as calculate_opacity +import tardis.opacities.opacities as calculate_opacity from benchmarks.benchmark_base import BenchmarkBase diff --git a/benchmarks/transport_montecarlo_packet.py b/benchmarks/transport_montecarlo_packet.py index 1c2f79bdf66..ae1cd9b00e7 100644 --- a/benchmarks/transport_montecarlo_packet.py +++ b/benchmarks/transport_montecarlo_packet.py @@ -5,14 +5,13 @@ import numpy as np from asv_runner.benchmarks.mark import parameterize, skip_benchmark -import tardis.transport.montecarlo.estimators.radfield_mc_estimators +import tardis.opacities.opacities as opacities +import tardis.transport.frame_transformations as frame_transformations +import tardis.transport.geometry.calculate_distances as calculate_distances import tardis.transport.montecarlo.estimators.radfield_mc_estimators import tardis.transport.montecarlo.numba_interface as numba_interface -import tardis.transport.montecarlo.opacities as opacities import tardis.transport.montecarlo.r_packet as r_packet import tardis.transport.montecarlo.utils as utils -import tardis.transport.frame_transformations as frame_transformations -import tardis.transport.geometry.calculate_distances as calculate_distances import tardis.transport.r_packet_transport as r_packet_transport from benchmarks.benchmark_base import BenchmarkBase from tardis.model.geometry.radial1d import NumbaRadial1DGeometry diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py index 66d83ac131a..385febda679 100644 --- a/tardis/energy_input/gamma_packet_loop.py +++ b/tardis/energy_input/gamma_packet_loop.py @@ -1,34 +1,33 @@ import numpy as np from numba import njit -from tardis.transport.montecarlo import njit_dict_no_parallel -from tardis.transport.montecarlo.opacities import ( - compton_opacity_calculation, - photoabsorption_opacity_calculation, - pair_creation_opacity_calculation, - photoabsorption_opacity_calculation_kasen, - kappa_calculation, - pair_creation_opacity_artis, - SIGMA_T, -) +from tardis.energy_input.gamma_ray_estimators import deposition_estimator_kasen from tardis.energy_input.gamma_ray_grid import ( distance_trace, move_packet, ) +from tardis.energy_input.gamma_ray_interactions import ( + compton_scatter, + get_compton_fraction_artis, + pair_creation_packet, + scatter_type, +) +from tardis.energy_input.GXPacket import GXPacketStatus from tardis.energy_input.util import ( - doppler_factor_3d, C_CGS, H_CGS_KEV, + doppler_factor_3d, get_index, ) -from tardis.energy_input.GXPacket import GXPacketStatus -from tardis.energy_input.gamma_ray_interactions import ( - get_compton_fraction_artis, - scatter_type, - compton_scatter, - pair_creation_packet, +from tardis.opacities.opacities import ( + SIGMA_T, + compton_opacity_calculation, + kappa_calculation, + pair_creation_opacity_artis, + pair_creation_opacity_calculation, + photoabsorption_opacity_calculation, ) -from tardis.energy_input.gamma_ray_estimators import deposition_estimator_kasen +from tardis.transport.montecarlo import njit_dict_no_parallel @njit(**njit_dict_no_parallel) diff --git a/tardis/energy_input/gamma_ray_estimators.py b/tardis/energy_input/gamma_ray_estimators.py index 87a313775b2..e28ba5029ff 100644 --- a/tardis/energy_input/gamma_ray_estimators.py +++ b/tardis/energy_input/gamma_ray_estimators.py @@ -1,19 +1,19 @@ import numpy as np from numba import njit -from tardis.transport.montecarlo import njit_dict_no_parallel -from tardis.transport.montecarlo.opacities import ( - compton_opacity_calculation, - SIGMA_T, - photoabsorption_opacity_calculation, - kappa_calculation, -) from tardis.energy_input.util import ( + ELECTRON_MASS_ENERGY_KEV, + H_CGS_KEV, angle_aberration_gamma, doppler_factor_3d, - H_CGS_KEV, - ELECTRON_MASS_ENERGY_KEV, ) +from tardis.opacities.opacities import ( + SIGMA_T, + compton_opacity_calculation, + kappa_calculation, + photoabsorption_opacity_calculation, +) +from tardis.transport.montecarlo import njit_dict_no_parallel def compton_emissivity_estimator(packet, distance): @@ -32,7 +32,6 @@ def compton_emissivity_estimator(packet, distance): float64, int Unnormalized emissivity estimator, line index """ - cmf_direction = angle_aberration_gamma( packet.get_direction_vector(), packet.location_r ) diff --git a/tardis/energy_input/gamma_ray_interactions.py b/tardis/energy_input/gamma_ray_interactions.py index c70ec6babd2..649a526af72 100644 --- a/tardis/energy_input/gamma_ray_interactions.py +++ b/tardis/energy_input/gamma_ray_interactions.py @@ -1,22 +1,22 @@ import numpy as np from numba import njit -from tardis.transport.montecarlo import njit_dict_no_parallel -from tardis.transport.montecarlo.opacities import ( - compton_opacity_partial, - kappa_calculation, -) +from tardis.energy_input.GXPacket import GXPacketStatus from tardis.energy_input.util import ( - get_random_unit_vector, - euler_rodrigues, - compton_theta_distribution, - get_perpendicular_vector, - angle_aberration_gamma, - doppler_factor_3d, ELECTRON_MASS_ENERGY_KEV, H_CGS_KEV, + angle_aberration_gamma, + compton_theta_distribution, + doppler_factor_3d, + euler_rodrigues, + get_perpendicular_vector, + get_random_unit_vector, ) -from tardis.energy_input.GXPacket import GXPacketStatus +from tardis.opacities.opacities import ( + compton_opacity_partial, + kappa_calculation, +) +from tardis.transport.montecarlo import njit_dict_no_parallel @njit(**njit_dict_no_parallel) @@ -195,7 +195,6 @@ def compton_scatter(photon, compton_angle): float64 Photon phi direction """ - # get comoving frame direction comov_direction = angle_aberration_gamma( photon.direction, photon.location, photon.time_current @@ -254,7 +253,6 @@ def pair_creation_packet(packet): GXPacket outgoing packet """ - probability_gamma = ( 2 * ELECTRON_MASS_ENERGY_KEV / (H_CGS_KEV * packet.nu_cmf) ) diff --git a/tardis/energy_input/gamma_ray_transport.py b/tardis/energy_input/gamma_ray_transport.py index 867414ee1bd..bd769191195 100644 --- a/tardis/energy_input/gamma_ray_transport.py +++ b/tardis/energy_input/gamma_ray_transport.py @@ -1,14 +1,15 @@ import logging + +import astropy.units as u import numpy as np import pandas as pd -import astropy.units as u import radioactivedecay as rd from tardis.energy_input.energy_source import ( get_all_isotopes, setup_input_energy, ) -from tardis.transport.montecarlo.opacities import M_P +from tardis.opacities.opacities import M_P # Energy: keV, exported as eV for SF solver # distance: cm @@ -317,10 +318,12 @@ def decay_chain_energies( def fractional_decay_energy(decay_energy): """Function to calculate fractional decay energy + Parameters ---------- decay_energy : Dict dictionary of decay chain energies for each isotope in each shell + Returns ------- fractional_decay_energy : Dict diff --git a/tardis/energy_input/tests/test_util.py b/tardis/energy_input/tests/test_util.py index 6d0268dd8b4..052a68efc20 100644 --- a/tardis/energy_input/tests/test_util.py +++ b/tardis/energy_input/tests/test_util.py @@ -8,7 +8,7 @@ klein_nishina, spherical_to_cartesian, ) -from tardis.transport.montecarlo.opacities import ( +from tardis.opacities.opacities import ( kappa_calculation, ) diff --git a/tardis/energy_input/util.py b/tardis/energy_input/util.py index 5f47c5bbda4..061017158b7 100644 --- a/tardis/energy_input/util.py +++ b/tardis/energy_input/util.py @@ -1,10 +1,10 @@ import astropy.units as u -import tardis.constants as const import numpy as np from numba import njit +import tardis.constants as const +from tardis.opacities.opacities import kappa_calculation from tardis.transport.montecarlo import njit_dict_no_parallel -from tardis.transport.montecarlo.opacities import kappa_calculation R_ELECTRON_SQUARED = (const.a0.cgs.value * const.alpha.cgs.value**2.0) ** 2.0 ELECTRON_MASS_ENERGY_KEV = (const.m_e * const.c**2.0).to("keV").value @@ -42,7 +42,8 @@ def spherical_to_cartesian(r, theta, phi): def get_random_unit_vector(): """Generate a random unit vector - Returns: + Returns + ------- array: random unit vector """ theta = get_random_theta_photon() @@ -218,9 +219,9 @@ def klein_nishina(energy, theta_C): https://en.wikipedia.org/wiki/Klein%E2%80%93Nishina_formula .. math:: - \frac{r_e}{2} [1 + \kappa (1 - \cos\theta_C)]^{-2} \left( 1 + \cos^2\theta_C + \frac{\kappa^2 (1 - \cos\theta_C)^2}{1 + \kappa(1 - \cos\theta_C)}\right) + \frac{r_e}{2} [1 + \\kappa (1 - \\cos\theta_C)]^{-2} \\left( 1 + \\cos^2\theta_C + \frac{\\kappa^2 (1 - \\cos\theta_C)^2}{1 + \\kappa(1 - \\cos\theta_C)}\right) - where :math:`\kappa = E / (m_e c^2)` + where :math:`\\kappa = E / (m_e c^2)` Parameters ---------- @@ -228,6 +229,7 @@ def klein_nishina(energy, theta_C): Packet energy theta_C : float Compton angle + Returns ------- float @@ -275,6 +277,7 @@ def compton_theta_distribution(energy, sample_resolution=100): @njit(**njit_dict_no_parallel) def get_random_theta_photon(): """Get a random theta direction between 0 and pi + Returns ------- float diff --git a/tardis/opacities/__init__.py b/tardis/opacities/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tardis/transport/montecarlo/opacities.py b/tardis/opacities/opacities.py similarity index 100% rename from tardis/transport/montecarlo/opacities.py rename to tardis/opacities/opacities.py diff --git a/tardis/opacities/tau_sobolev.py b/tardis/opacities/tau_sobolev.py new file mode 100644 index 00000000000..95997683ea3 --- /dev/null +++ b/tardis/opacities/tau_sobolev.py @@ -0,0 +1,104 @@ +import numpy as np +import pandas as pd +from astropy import units as u + +from tardis import constants as const +from tardis.plasma.properties.base import ProcessingPlasmaProperty + +SOBOLEV_COEFFICIENT = ( + ( + ((np.pi * const.e.gauss**2) / (const.m_e.cgs * const.c.cgs)) + * u.cm + * u.s + / u.cm**3 + ) + .to(1) + .value +) + + +def calculate_sobolev_line_opacity( + lines, + level_number_density, + time_explosion, + stimulated_emission_factor, +): + tau_sobolevs = (lines.wavelength_cm * lines.f_lu).values[np.newaxis].T + tau_sobolevs *= ( + SOBOLEV_COEFFICIENT + * time_explosion.to(u.s).value + * stimulated_emission_factor + ) + tau_sobolevs *= level_number_density.reindex( + lines.droplevel(-1).index + ).values + return tau_sobolevs + + +class TauSobolev(ProcessingPlasmaProperty): + """ + Attributes + ---------- + tau_sobolev : Pandas DataFrame, dtype float + Sobolev optical depth for each line. Indexed by line. + Columns as zones. + """ + + outputs = ("tau_sobolevs",) + latex_name = (r"\tau_{\textrm{sobolev}}",) + latex_formula = ( + r"\dfrac{\pi e^{2}}{m_{e} c}f_{lu}\lambda t_{exp}\ + n_{lower} \Big(1-\dfrac{g_{lower}n_{upper}}{g_{upper}n_{lower}}\Big)", + ) + + def __init__(self, plasma_parent): + super(TauSobolev, self).__init__(plasma_parent) + self.sobolev_coefficient = ( + ( + ((np.pi * const.e.gauss**2) / (const.m_e.cgs * const.c.cgs)) + * u.cm + * u.s + / u.cm**3 + ) + .to(1) + .value + ) + + def calculate( + self, + lines, + level_number_density, + lines_lower_level_index, + time_explosion, + stimulated_emission_factor, + j_blues, + f_lu, + wavelength_cm, + ): + f_lu = f_lu.values[np.newaxis].T + wavelength = wavelength_cm.values[np.newaxis].T + n_lower = level_number_density.values.take( + lines_lower_level_index, axis=0, mode="raise" + ) + tau_sobolevs = ( + self.sobolev_coefficient + * f_lu + * wavelength + * time_explosion + * n_lower + * stimulated_emission_factor + ) + + if np.any(np.isnan(tau_sobolevs)) or np.any( + np.isinf(np.abs(tau_sobolevs)) + ): + raise ValueError( + "Some tau_sobolevs are nan, inf, -inf in tau_sobolevs." + " Something went wrong!" + ) + + return pd.DataFrame( + tau_sobolevs, + index=lines.index, + columns=np.array(level_number_density.columns), + ) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index d93b6016442..26d08db5d64 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -1,4 +1,5 @@ from tardis.plasma.properties import * +from tardis.opacities.tau_sobolev import TauSobolev class PlasmaPropertyCollection(list): diff --git a/tardis/plasma/properties/radiative_properties.py b/tardis/plasma/properties/radiative_properties.py index f3b5671de46..99d34e91924 100644 --- a/tardis/plasma/properties/radiative_properties.py +++ b/tardis/plasma/properties/radiative_properties.py @@ -16,7 +16,6 @@ __all__ = [ "StimulatedEmissionFactor", - "TauSobolev", "BetaSobolev", "TransitionProbabilities", "RawRadBoundBoundTransProbs", @@ -115,75 +114,6 @@ def calculate( return stimulated_emission_factor -class TauSobolev(ProcessingPlasmaProperty): - """ - Attributes - ---------- - tau_sobolev : Pandas DataFrame, dtype float - Sobolev optical depth for each line. Indexed by line. - Columns as zones. - """ - - outputs = ("tau_sobolevs",) - latex_name = (r"\tau_{\textrm{sobolev}}",) - latex_formula = ( - r"\dfrac{\pi e^{2}}{m_{e} c}f_{lu}\lambda t_{exp}\ - n_{lower} \Big(1-\dfrac{g_{lower}n_{upper}}{g_{upper}n_{lower}}\Big)", - ) - - def __init__(self, plasma_parent): - super(TauSobolev, self).__init__(plasma_parent) - self.sobolev_coefficient = ( - ( - ((np.pi * const.e.gauss**2) / (const.m_e.cgs * const.c.cgs)) - * u.cm - * u.s - / u.cm**3 - ) - .to(1) - .value - ) - - def calculate( - self, - lines, - level_number_density, - lines_lower_level_index, - time_explosion, - stimulated_emission_factor, - j_blues, - f_lu, - wavelength_cm, - ): - f_lu = f_lu.values[np.newaxis].T - wavelength = wavelength_cm.values[np.newaxis].T - n_lower = level_number_density.values.take( - lines_lower_level_index, axis=0, mode="raise" - ) - tau_sobolevs = ( - self.sobolev_coefficient - * f_lu - * wavelength - * time_explosion - * n_lower - * stimulated_emission_factor - ) - - if np.any(np.isnan(tau_sobolevs)) or np.any( - np.isinf(np.abs(tau_sobolevs)) - ): - raise ValueError( - "Some tau_sobolevs are nan, inf, -inf in tau_sobolevs." - " Something went wrong!" - ) - - return pd.DataFrame( - tau_sobolevs, - index=lines.index, - columns=np.array(level_number_density.columns), - ) - - class BetaSobolev(ProcessingPlasmaProperty): """ Attributes diff --git a/tardis/transport/montecarlo/single_packet_loop.py b/tardis/transport/montecarlo/single_packet_loop.py index 693cd75d6ea..2560c814410 100644 --- a/tardis/transport/montecarlo/single_packet_loop.py +++ b/tardis/transport/montecarlo/single_packet_loop.py @@ -1,6 +1,14 @@ from numba import njit from tardis import constants as const +from tardis.opacities.opacities import ( + chi_continuum_calculator, + chi_electron_calculator, +) +from tardis.transport.frame_transformations import ( + get_doppler_factor, + get_inverse_doppler_factor, +) from tardis.transport.montecarlo.estimators.radfield_estimator_calcs import ( update_bound_free_estimators, ) @@ -9,19 +17,11 @@ line_scatter, thomson_scatter, ) -from tardis.transport.montecarlo.opacities import ( - chi_continuum_calculator, - chi_electron_calculator, -) from tardis.transport.montecarlo.r_packet import ( InteractionType, PacketStatus, ) from tardis.transport.montecarlo.vpacket import trace_vpacket_volley -from tardis.transport.frame_transformations import ( - get_doppler_factor, - get_inverse_doppler_factor, -) from tardis.transport.r_packet_transport import ( move_packet_across_shell_boundary, move_r_packet, diff --git a/tardis/transport/montecarlo/tests/test_opacities.py b/tardis/transport/montecarlo/tests/test_opacities.py index 53933960579..07096d6e0c6 100644 --- a/tardis/transport/montecarlo/tests/test_opacities.py +++ b/tardis/transport/montecarlo/tests/test_opacities.py @@ -1,7 +1,7 @@ import numpy.testing as npt import pytest -from tardis.transport.montecarlo.opacities import ( +from tardis.opacities.opacities import ( compton_opacity_calculation, kappa_calculation, pair_creation_opacity_calculation, diff --git a/tardis/transport/montecarlo/tests/test_packet.py b/tardis/transport/montecarlo/tests/test_packet.py index b0c602b6b08..c0e27c128e2 100644 --- a/tardis/transport/montecarlo/tests/test_packet.py +++ b/tardis/transport/montecarlo/tests/test_packet.py @@ -1,14 +1,14 @@ import numpy as np import pytest -import tardis.transport.montecarlo.montecarlo_configuration as numba_config +import tardis.opacities.opacities as opacities +import tardis.transport.frame_transformations as frame_transformations +import tardis.transport.geometry.calculate_distances as calculate_distances import tardis.transport.montecarlo.estimators.radfield_mc_estimators +import tardis.transport.montecarlo.montecarlo_configuration as numba_config import tardis.transport.montecarlo.numba_interface as numba_interface -import tardis.transport.montecarlo.opacities as opacities import tardis.transport.montecarlo.r_packet as r_packet import tardis.transport.montecarlo.utils as utils -import tardis.transport.frame_transformations as frame_transformations -import tardis.transport.geometry.calculate_distances as calculate_distances import tardis.transport.r_packet_transport as r_packet_transport from tardis import constants as const from tardis.model.geometry.radial1d import NumbaRadial1DGeometry diff --git a/tardis/transport/montecarlo/vpacket.py b/tardis/transport/montecarlo/vpacket.py index ee084a4d580..3aacd8a5be9 100644 --- a/tardis/transport/montecarlo/vpacket.py +++ b/tardis/transport/montecarlo/vpacket.py @@ -1,35 +1,31 @@ import math import numpy as np -from tardis.transport.montecarlo.opacities import ( - chi_continuum_calculator, -) -from numba import float64, int64 -from numba import njit +from numba import float64, int64, njit from numba.experimental import jitclass -from tardis.transport.montecarlo import njit_dict_no_parallel -from tardis.transport.montecarlo.r_packet import ( - PacketStatus, +from tardis.opacities.opacities import ( + chi_continuum_calculator, ) -from tardis.transport.r_packet_transport import ( - move_packet_across_shell_boundary, +from tardis.transport.frame_transformations import ( + angle_aberration_CMF_to_LF, + angle_aberration_LF_to_CMF, + get_doppler_factor, ) - from tardis.transport.geometry.calculate_distances import ( calculate_distance_boundary, calculate_distance_line, ) - -from tardis.transport.frame_transformations import ( - get_doppler_factor, - angle_aberration_LF_to_CMF, - angle_aberration_CMF_to_LF, -) - +from tardis.transport.montecarlo import njit_dict_no_parallel from tardis.transport.montecarlo.numba_config import ( - SIGMA_THOMSON, C_SPEED_OF_LIGHT, + SIGMA_THOMSON, +) +from tardis.transport.montecarlo.r_packet import ( + PacketStatus, +) +from tardis.transport.r_packet_transport import ( + move_packet_across_shell_boundary, ) vpacket_spec = [ @@ -45,7 +41,7 @@ @jitclass(vpacket_spec) -class VPacket(object): +class VPacket: def __init__( self, r, @@ -175,6 +171,7 @@ def trace_vpacket( ): """ Trace single vpacket. + Parameters ---------- v_packet @@ -185,7 +182,6 @@ def trace_vpacket( ------- """ - tau_trace_combined = 0.0 while True: ( @@ -262,7 +258,6 @@ def trace_vpacket_volley( opacity_state : [type] [description] """ - if (r_packet.nu < vpacket_collection.v_packet_spawn_start_frequency) or ( r_packet.nu > vpacket_collection.v_packet_spawn_end_frequency ): diff --git a/tardis/transport/r_packet_transport.py b/tardis/transport/r_packet_transport.py index 598e60adfdc..97976a2de43 100644 --- a/tardis/transport/r_packet_transport.py +++ b/tardis/transport/r_packet_transport.py @@ -1,20 +1,18 @@ import numpy as np from numba import njit -from tardis.transport.montecarlo import njit_dict_no_parallel +from tardis.transport.frame_transformations import ( + get_doppler_factor, +) from tardis.transport.geometry.calculate_distances import ( calculate_distance_boundary, - calculate_distance_electron, calculate_distance_line, ) +from tardis.transport.montecarlo import njit_dict_no_parallel from tardis.transport.montecarlo.estimators.radfield_estimator_calcs import ( - update_line_estimators, update_base_estimators, + update_line_estimators, ) -from tardis.transport.frame_transformations import ( - get_doppler_factor, -) -from tardis.transport.montecarlo.opacities import calculate_tau_electron from tardis.transport.montecarlo.r_packet import ( InteractionType, PacketStatus, @@ -48,7 +46,6 @@ def trace_packet( Returns ------- """ - r_inner = numba_radial_1d_geometry.r_inner[r_packet.current_shell_id] r_outer = numba_radial_1d_geometry.r_outer[r_packet.current_shell_id] @@ -199,7 +196,6 @@ def move_r_packet( distance : float distance in cm """ - doppler_factor = get_doppler_factor( r_packet.r, r_packet.mu, time_explosion, enable_full_relativity ) From e573f16eca2e02369573ed3ce8a5c66e0efed658 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 6 May 2024 21:26:35 -0400 Subject: [PATCH 08/89] Update imports in property_collections.py, base.py, test_numba_interface.py, transport_montecarlo_numba_interface.py, conftest.py, formal_integral.py, base.py, and macro_atom.py --- .../transport_montecarlo_numba_interface.py | 2 +- tardis/opacities/tau_sobolev.py | 58 ++--- .../plasma/properties/property_collections.py | 4 + .../plasma/properties/radiative_properties.py | 142 +---------- tardis/plasma/properties/util/macro_atom.py | 61 ----- tardis/radiation_field/base.py | 2 +- tardis/transport/montecarlo/base.py | 2 +- .../transport/montecarlo/formal_integral.py | 6 +- .../transport/montecarlo/numba_interface.py | 240 ------------------ tardis/transport/montecarlo/tests/conftest.py | 2 +- .../montecarlo/tests/test_numba_interface.py | 2 +- 11 files changed, 36 insertions(+), 485 deletions(-) delete mode 100644 tardis/plasma/properties/util/macro_atom.py diff --git a/benchmarks/transport_montecarlo_numba_interface.py b/benchmarks/transport_montecarlo_numba_interface.py index 9b999c7c848..3d8ead2ed04 100644 --- a/benchmarks/transport_montecarlo_numba_interface.py +++ b/benchmarks/transport_montecarlo_numba_interface.py @@ -5,7 +5,7 @@ import numpy as np from asv_runner.benchmarks.mark import parameterize -import tardis.transport.montecarlo.numba_interface as numba_interface +import tardis.opacities.opacity_state as numba_interface from benchmarks.benchmark_base import BenchmarkBase diff --git a/tardis/opacities/tau_sobolev.py b/tardis/opacities/tau_sobolev.py index 95997683ea3..302e00a6f64 100644 --- a/tardis/opacities/tau_sobolev.py +++ b/tardis/opacities/tau_sobolev.py @@ -23,16 +23,25 @@ def calculate_sobolev_line_opacity( time_explosion, stimulated_emission_factor, ): - tau_sobolevs = (lines.wavelength_cm * lines.f_lu).values[np.newaxis].T - tau_sobolevs *= ( - SOBOLEV_COEFFICIENT + tau_sobolevs = ( + (lines.wavelength_cm * lines.f_lu).values[np.newaxis].T + * SOBOLEV_COEFFICIENT * time_explosion.to(u.s).value * stimulated_emission_factor + * level_number_density.reindex(lines.droplevel(-1).index).values + ) + + if np.any(np.isnan(tau_sobolevs)) or np.any(np.isinf(np.abs(tau_sobolevs))): + raise ValueError( + "Some tau_sobolevs are nan, inf, -inf in tau_sobolevs." + " Something went wrong!" + ) + + return pd.DataFrame( + tau_sobolevs, + index=lines.index, + columns=np.array(level_number_density.columns), ) - tau_sobolevs *= level_number_density.reindex( - lines.droplevel(-1).index - ).values - return tau_sobolevs class TauSobolev(ProcessingPlasmaProperty): @@ -68,37 +77,12 @@ def calculate( self, lines, level_number_density, - lines_lower_level_index, time_explosion, stimulated_emission_factor, - j_blues, - f_lu, - wavelength_cm, ): - f_lu = f_lu.values[np.newaxis].T - wavelength = wavelength_cm.values[np.newaxis].T - n_lower = level_number_density.values.take( - lines_lower_level_index, axis=0, mode="raise" - ) - tau_sobolevs = ( - self.sobolev_coefficient - * f_lu - * wavelength - * time_explosion - * n_lower - * stimulated_emission_factor - ) - - if np.any(np.isnan(tau_sobolevs)) or np.any( - np.isinf(np.abs(tau_sobolevs)) - ): - raise ValueError( - "Some tau_sobolevs are nan, inf, -inf in tau_sobolevs." - " Something went wrong!" - ) - - return pd.DataFrame( - tau_sobolevs, - index=lines.index, - columns=np.array(level_number_density.columns), + return calculate_sobolev_line_opacity( + lines, + level_number_density, + time_explosion, + stimulated_emission_factor, ) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index 26d08db5d64..3b1e8ebc5cd 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -1,3 +1,7 @@ +from tardis.opacities.macro_atom.base import ( + NonMarkovChainTransitionProbabilities, + TransitionProbabilities, +) from tardis.plasma.properties import * from tardis.opacities.tau_sobolev import TauSobolev diff --git a/tardis/plasma/properties/radiative_properties.py b/tardis/plasma/properties/radiative_properties.py index 99d34e91924..6facf1e8d27 100644 --- a/tardis/plasma/properties/radiative_properties.py +++ b/tardis/plasma/properties/radiative_properties.py @@ -3,21 +3,20 @@ import numpy as np import pandas as pd from astropy import units as u -from tardis import constants as const from numba import jit, prange +from tardis import constants as const +from tardis.opacities.macro_atom.base import TransitionProbabilities from tardis.plasma.properties.base import ( ProcessingPlasmaProperty, TransitionProbabilitiesProperty, ) -from tardis.plasma.properties.util import macro_atom logger = logging.getLogger(__name__) __all__ = [ "StimulatedEmissionFactor", "BetaSobolev", - "TransitionProbabilities", "RawRadBoundBoundTransProbs", "NonMarkovChainTransitionProbabilities", ] @@ -154,143 +153,6 @@ def calculate_beta_sobolev(tau_sobolevs, beta_sobolevs): return beta_sobolevs -class TransitionProbabilities(ProcessingPlasmaProperty): - """ - Attributes - ---------- - transition_probabilities : Pandas DataFrame, dtype float - """ - - outputs = ("transition_probabilities",) - - def __init__(self, plasma_parent): - super(TransitionProbabilities, self).__init__(plasma_parent) - self.initialize = True - self.normalize = True - - def calculate( - self, - atomic_data, - beta_sobolev, - j_blues, - stimulated_emission_factor, - tau_sobolevs, - ): - # I wonder why? - # Not sure who wrote this but the answer is that when the plasma is - # first initialised (before the first iteration, without temperature - # values etc.) there are no j_blues values so this just prevents - # an error. Aoife. - if len(j_blues) == 0: - return None - macro_atom_data = self._get_macro_atom_data(atomic_data) - if self.initialize: - self.initialize_macro_atom_transition_type_filters( - atomic_data, macro_atom_data - ) - self.transition_probability_coef = ( - self._get_transition_probability_coefs(macro_atom_data) - ) - self.initialize = False - transition_probabilities = self._calculate_transition_probability( - macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor - ) - transition_probabilities = pd.DataFrame( - transition_probabilities, - index=macro_atom_data.transition_line_id, - columns=tau_sobolevs.columns, - ) - return transition_probabilities - - def _calculate_transition_probability( - self, macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor - ): - transition_probabilities = np.empty( - (self.transition_probability_coef.shape[0], beta_sobolev.shape[1]) - ) - # trans_old = self.calculate_transition_probabilities(macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor) - transition_type = macro_atom_data.transition_type.values - lines_idx = macro_atom_data.lines_idx.values - tpos = macro_atom_data.transition_probability.values - macro_atom.calculate_transition_probabilities( - tpos, - beta_sobolev.values, - j_blues.values, - stimulated_emission_factor, - transition_type, - lines_idx, - self.block_references, - transition_probabilities, - self.normalize, - ) - return transition_probabilities - - def calculate_transition_probabilities( - self, macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor - ): - transition_probabilities = self.prepare_transition_probabilities( - macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor - ) - return transition_probabilities - - def initialize_macro_atom_transition_type_filters( - self, atomic_data, macro_atom_data - ): - self.transition_up_filter = macro_atom_data.transition_type.values == 1 - self.transition_up_line_filter = macro_atom_data.lines_idx.values[ - self.transition_up_filter - ] - self.block_references = np.hstack( - ( - atomic_data.macro_atom_references.block_references, - len(macro_atom_data), - ) - ) - - @staticmethod - def _get_transition_probability_coefs(macro_atom_data): - return macro_atom_data.transition_probability.values[np.newaxis].T - - def prepare_transition_probabilities( - self, macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor - ): - current_beta_sobolev = beta_sobolev.values.take( - macro_atom_data.lines_idx.values, axis=0, mode="raise" - ) - transition_probabilities = ( - self.transition_probability_coef * current_beta_sobolev - ) - j_blues = j_blues.take( - self.transition_up_line_filter, axis=0, mode="raise" - ) - macro_stimulated_emission = stimulated_emission_factor.take( - self.transition_up_line_filter, axis=0, mode="raise" - ) - transition_probabilities[self.transition_up_filter] *= ( - j_blues * macro_stimulated_emission - ) - return transition_probabilities - - def _normalize_transition_probabilities(self, transition_probabilities): - macro_atom.normalize_transition_probabilities( - transition_probabilities, self.block_references - ) - - @staticmethod - def _get_macro_atom_data(atomic_data): - try: - return atomic_data.macro_atom_data - except: - logger.debug( - "Macro Atom Data was not found. Instead returning All Macro Atom Data" - ) - return atomic_data.macro_atom_data_all - - -class NonMarkovChainTransitionProbabilities(TransitionProbabilities): - outputs = ("non_markov_transition_probabilities",) - - class RawRadBoundBoundTransProbs( TransitionProbabilities, TransitionProbabilitiesProperty ): diff --git a/tardis/plasma/properties/util/macro_atom.py b/tardis/plasma/properties/util/macro_atom.py deleted file mode 100644 index c72781f4b41..00000000000 --- a/tardis/plasma/properties/util/macro_atom.py +++ /dev/null @@ -1,61 +0,0 @@ -import numpy as np -from numba import njit - -from tardis import constants as const -from tardis.transport.montecarlo import njit_dict - -h_cgs = const.h.cgs.value -c = const.c.to("cm/s").value -kb = const.k_B.cgs.value -inv_c2 = 1 / (c**2) - - -@njit(**njit_dict) -def calculate_transition_probabilities( - transition_probability_coef, - beta_sobolev, - j_blues, - stimulated_emission_factor, - transition_type, - lines_idx, - block_references, - transition_probabilities, - normalize, -): - """ - Calculates transition probabilities for macro_atom interactions - - transition_probability_coef must be a 1D array - transition_type, lines_idx, and block_references must be int-type arrays - beta_sobolev, j_blues,stimulated_emission_factor, and transition_probabilities must be 2D array - """ - norm_factor = np.zeros(transition_probabilities.shape[1]) - - for i in range(transition_probabilities.shape[0]): - line_idx = lines_idx[i] - for j in range(transition_probabilities.shape[1]): - transition_probabilities[i, j] = ( - transition_probability_coef[i] * beta_sobolev[line_idx, j] - ) - if transition_type[i] == 1: - for j in range(transition_probabilities.shape[1]): - transition_probabilities[i, j] *= ( - stimulated_emission_factor[line_idx, j] - * j_blues[line_idx, j] - ) - - if normalize: - for i in range(block_references.shape[0] - 1): - for k in range(transition_probabilities.shape[1]): - norm_factor[k] = 0.0 - for j in range(block_references[i], block_references[i + 1]): - for k in range(transition_probabilities.shape[1]): - norm_factor[k] += transition_probabilities[j, k] - for k in range(transition_probabilities.shape[1]): - if norm_factor[k] != 0.0: - norm_factor[k] = 1 / norm_factor[k] - else: - norm_factor[k] = 1.0 - for j in range(block_references[i], block_references[i + 1]): - for k in range(transition_probabilities.shape[1]): - transition_probabilities[j, k] *= norm_factor[k] diff --git a/tardis/radiation_field/base.py b/tardis/radiation_field/base.py index 9614c5f4621..4fe0658c20c 100644 --- a/tardis/radiation_field/base.py +++ b/tardis/radiation_field/base.py @@ -2,7 +2,7 @@ from astropy import units as u from tardis.transport.montecarlo.packet_source import BasePacketSource -from tardis.transport.montecarlo.numba_interface import OpacityState +from tardis.opacities.opacity_state import OpacityState class MonteCarloRadiationFieldState: diff --git a/tardis/transport/montecarlo/base.py b/tardis/transport/montecarlo/base.py index d607616ed71..c97b6d6a1b9 100644 --- a/tardis/transport/montecarlo/base.py +++ b/tardis/transport/montecarlo/base.py @@ -7,6 +7,7 @@ from tardis import constants as const from tardis.io.logger import montecarlo_tracking as mc_tracker from tardis.io.util import HDFWriterMixin +from tardis.opacities.opacity_state import opacity_state_initialize from tardis.transport.montecarlo.estimators.radfield_mc_estimators import ( initialize_estimator_statistics, ) @@ -20,7 +21,6 @@ from tardis.transport.montecarlo.formal_integral import FormalIntegrator from tardis.transport.montecarlo.numba_interface import ( NumbaModel, - opacity_state_initialize, ) from tardis.transport.montecarlo.r_packet import ( rpacket_trackers_to_dataframe, diff --git a/tardis/transport/montecarlo/formal_integral.py b/tardis/transport/montecarlo/formal_integral.py index e27df091055..6b3154d01b4 100644 --- a/tardis/transport/montecarlo/formal_integral.py +++ b/tardis/transport/montecarlo/formal_integral.py @@ -10,12 +10,14 @@ from numba.experimental import jitclass +from tardis.opacities.opacity_state import ( + OpacityState, + opacity_state_initialize, +) from tardis.transport.montecarlo.numba_config import SIGMA_THOMSON from tardis.transport.montecarlo import njit_dict, njit_dict_no_parallel from tardis.transport.montecarlo.numba_interface import ( - opacity_state_initialize, NumbaModel, - OpacityState, ) from tardis.transport.montecarlo.formal_integral_cuda import ( CudaFormalIntegrator, diff --git a/tardis/transport/montecarlo/numba_interface.py b/tardis/transport/montecarlo/numba_interface.py index f6c9198444d..1962315ef4e 100644 --- a/tardis/transport/montecarlo/numba_interface.py +++ b/tardis/transport/montecarlo/numba_interface.py @@ -28,246 +28,6 @@ def __init__(self, time_explosion): self.time_explosion = time_explosion -opacity_state_spec = [ - ("electron_density", float64[:]), - ("t_electrons", float64[:]), - ("line_list_nu", float64[:]), - ("tau_sobolev", float64[:, :]), - ("transition_probabilities", float64[:, :]), - ("line2macro_level_upper", int64[:]), - ("macro_block_references", int64[:]), - ("transition_type", int64[:]), - ("destination_level_id", int64[:]), - ("transition_line_id", int64[:]), - ("bf_threshold_list_nu", float64[:]), - ("p_fb_deactivation", float64[:, :]), - ("photo_ion_nu_threshold_mins", float64[:]), - ("photo_ion_nu_threshold_maxs", float64[:]), - ("photo_ion_block_references", int64[:]), - ("chi_bf", float64[:, :]), - ("x_sect", float64[:]), - ("phot_nus", float64[:]), - ("ff_opacity_factor", float64[:]), - ("emissivities", float64[:, :]), - ("photo_ion_activation_idx", int64[:]), - ("k_packet_idx", int64), -] - - -@jitclass(opacity_state_spec) -class OpacityState(object): - def __init__( - self, - electron_density, - t_electrons, - line_list_nu, - tau_sobolev, - transition_probabilities, - line2macro_level_upper, - macro_block_references, - transition_type, - destination_level_id, - transition_line_id, - bf_threshold_list_nu, - p_fb_deactivation, - photo_ion_nu_threshold_mins, - photo_ion_nu_threshold_maxs, - photo_ion_block_references, - chi_bf, - x_sect, - phot_nus, - ff_opacity_factor, - emissivities, - photo_ion_activation_idx, - k_packet_idx, - ): - """ - Plasma for the Numba code - - Parameters - ---------- - electron_density : numpy.ndarray - t_electrons : numpy.ndarray - line_list_nu : numpy.ndarray - tau_sobolev : numpy.ndarray - transition_probabilities : numpy.ndarray - line2macro_level_upper : numpy.ndarray - macro_block_references : numpy.ndarray - transition_type : numpy.ndarray - destination_level_id : numpy.ndarray - transition_line_id : numpy.ndarray - bf_threshold_list_nu : numpy.ndarray - """ - - self.electron_density = electron_density - self.t_electrons = t_electrons - self.line_list_nu = line_list_nu - self.tau_sobolev = tau_sobolev - self.bf_threshold_list_nu = bf_threshold_list_nu - - #### Macro Atom transition probabilities - self.transition_probabilities = transition_probabilities - self.line2macro_level_upper = line2macro_level_upper - - self.macro_block_references = macro_block_references - self.transition_type = transition_type - - # Destination level is not needed and/or generated for downbranch - self.destination_level_id = destination_level_id - self.transition_line_id = transition_line_id - self.p_fb_deactivation = p_fb_deactivation - - # Continuum Opacity Data - self.photo_ion_nu_threshold_mins = photo_ion_nu_threshold_mins - self.photo_ion_nu_threshold_maxs = photo_ion_nu_threshold_maxs - - self.photo_ion_block_references = photo_ion_block_references - self.chi_bf = chi_bf - self.x_sect = x_sect - self.phot_nus = phot_nus - self.ff_opacity_factor = ff_opacity_factor - self.emissivities = emissivities - self.photo_ion_activation_idx = photo_ion_activation_idx - self.k_packet_idx = k_packet_idx - - -def opacity_state_initialize( - plasma, - line_interaction_type, - disable_line_scattering, - continuum_processes_enabled, -): - """ - Initialize the OpacityState object and copy over the data over from TARDIS Plasma - - Parameters - ---------- - plasma : tardis.plasma.BasePlasma - line_interaction_type : enum - """ - - electron_densities = plasma.electron_densities.values - t_electrons = plasma.t_electrons - line_list_nu = plasma.atomic_data.lines.nu.values - tau_sobolev = np.ascontiguousarray( - plasma.tau_sobolevs.values.copy(), dtype=np.float64 - ) - if disable_line_scattering: - tau_sobolev *= 0 - - if line_interaction_type == "scatter": - # to adhere to data types, we must have an array of minimum size 1 - array_size = 1 - transition_probabilities = np.zeros( - (array_size, array_size), dtype=np.float64 - ) # to adhere to data types - line2macro_level_upper = np.zeros(array_size, dtype=np.int64) - macro_block_references = np.zeros(array_size, dtype=np.int64) - transition_type = np.zeros(array_size, dtype=np.int64) - destination_level_id = np.zeros(array_size, dtype=np.int64) - transition_line_id = np.zeros(array_size, dtype=np.int64) - else: - transition_probabilities = np.ascontiguousarray( - plasma.transition_probabilities.values.copy(), dtype=np.float64 - ) - line2macro_level_upper = ( - plasma.atomic_data.lines_upper2macro_reference_idx - ) - # TODO: Fix setting of block references for non-continuum mode - - if continuum_processes_enabled: - macro_block_references = plasma.macro_block_references - else: - macro_block_references = plasma.atomic_data.macro_atom_references[ - "block_references" - ].values - transition_type = plasma.macro_atom_data["transition_type"].values - - # Destination level is not needed and/or generated for downbranch - destination_level_id = plasma.macro_atom_data[ - "destination_level_idx" - ].values - transition_line_id = plasma.macro_atom_data["lines_idx"].values - if continuum_processes_enabled: - bf_threshold_list_nu = plasma.nu_i.loc[ - plasma.level2continuum_idx.index - ].values - p_fb_deactivation = np.ascontiguousarray( - plasma.p_fb_deactivation.values.copy(), dtype=np.float64 - ) - - phot_nus = plasma.photo_ion_cross_sections.nu.loc[ - plasma.level2continuum_idx.index - ] - photo_ion_block_references = np.pad( - phot_nus.groupby(level=[0, 1, 2], sort=False) - .count() - .values.cumsum(), - [1, 0], - ) - photo_ion_nu_threshold_mins = ( - phot_nus.groupby(level=[0, 1, 2], sort=False).first().values - ) - photo_ion_nu_threshold_maxs = ( - phot_nus.groupby(level=[0, 1, 2], sort=False).last().values - ) - - chi_bf = plasma.chi_bf.loc[plasma.level2continuum_idx.index].values - x_sect = plasma.photo_ion_cross_sections.x_sect.loc[ - plasma.level2continuum_idx.index - ].values - - phot_nus = phot_nus.values - ff_opacity_factor = ( - plasma.ff_cooling_factor / np.sqrt(t_electrons) - ).astype(np.float64) - emissivities = plasma.fb_emission_cdf.loc[ - plasma.level2continuum_idx.index - ].values - photo_ion_activation_idx = plasma.photo_ion_idx.loc[ - plasma.level2continuum_idx.index, "destination_level_idx" - ].values - k_packet_idx = np.int64(plasma.k_packet_idx) - else: - bf_threshold_list_nu = np.zeros(0, dtype=np.float64) - p_fb_deactivation = np.zeros((0, 0), dtype=np.float64) - photo_ion_nu_threshold_mins = np.zeros(0, dtype=np.float64) - photo_ion_nu_threshold_maxs = np.zeros(0, dtype=np.float64) - photo_ion_block_references = np.zeros(0, dtype=np.int64) - chi_bf = np.zeros((0, 0), dtype=np.float64) - x_sect = np.zeros(0, dtype=np.float64) - phot_nus = np.zeros(0, dtype=np.float64) - ff_opacity_factor = np.zeros(0, dtype=np.float64) - emissivities = np.zeros((0, 0), dtype=np.float64) - photo_ion_activation_idx = np.zeros(0, dtype=np.int64) - k_packet_idx = np.int64(-1) - - return OpacityState( - electron_densities, - t_electrons, - line_list_nu, - tau_sobolev, - transition_probabilities, - line2macro_level_upper, - macro_block_references, - transition_type, - destination_level_id, - transition_line_id, - bf_threshold_list_nu, - p_fb_deactivation, - photo_ion_nu_threshold_mins, - photo_ion_nu_threshold_maxs, - photo_ion_block_references, - chi_bf, - x_sect, - phot_nus, - ff_opacity_factor, - emissivities, - photo_ion_activation_idx, - k_packet_idx, - ) - - rpacket_tracker_spec = [ ("length", int64), ("seed", int64), diff --git a/tardis/transport/montecarlo/tests/conftest.py b/tardis/transport/montecarlo/tests/conftest.py index 659f2037be9..eae3827f577 100644 --- a/tardis/transport/montecarlo/tests/conftest.py +++ b/tardis/transport/montecarlo/tests/conftest.py @@ -3,6 +3,7 @@ import pytest import numpy as np from numba import njit +from tardis.opacities.opacity_state import opacity_state_initialize from tardis.transport.montecarlo.packet_collections import ( VPacketCollection, ) @@ -15,7 +16,6 @@ from tardis.transport.montecarlo.numba_interface import ( - opacity_state_initialize, NumbaModel, ) diff --git a/tardis/transport/montecarlo/tests/test_numba_interface.py b/tardis/transport/montecarlo/tests/test_numba_interface.py index 5d779ce36d2..0b41c863eef 100644 --- a/tardis/transport/montecarlo/tests/test_numba_interface.py +++ b/tardis/transport/montecarlo/tests/test_numba_interface.py @@ -1,5 +1,5 @@ import pytest -import tardis.transport.montecarlo.numba_interface as numba_interface +import tardis.opacities.opacity_state as numba_interface import numpy.testing as npt import numpy as np From 73f2ddcc015e7fb977225fee4059f809fd9b6bf3 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 6 May 2024 21:27:13 -0400 Subject: [PATCH 09/89] Add calculate_transition_probabilities function to util.py in macro_atom package --- tardis/opacities/macro_atom/base.py | 144 ++++++++++++++++ tardis/opacities/macro_atom/util.py | 61 +++++++ tardis/opacities/opacity_state.py | 249 ++++++++++++++++++++++++++++ 3 files changed, 454 insertions(+) create mode 100644 tardis/opacities/macro_atom/base.py create mode 100644 tardis/opacities/macro_atom/util.py create mode 100644 tardis/opacities/opacity_state.py diff --git a/tardis/opacities/macro_atom/base.py b/tardis/opacities/macro_atom/base.py new file mode 100644 index 00000000000..251ddeaf792 --- /dev/null +++ b/tardis/opacities/macro_atom/base.py @@ -0,0 +1,144 @@ +from tardis.plasma.properties.base import ProcessingPlasmaProperty +from tardis.plasma.properties.radiative_properties import logger +from tardis.opacities.macro_atom import util + + +import numpy as np +import pandas as pd + + +class TransitionProbabilities(ProcessingPlasmaProperty): + """ + Attributes + ---------- + transition_probabilities : Pandas DataFrame, dtype float + """ + + outputs = ("transition_probabilities",) + + def __init__(self, plasma_parent): + super(TransitionProbabilities, self).__init__(plasma_parent) + self.initialize = True + self.normalize = True + + def calculate( + self, + atomic_data, + beta_sobolev, + j_blues, + stimulated_emission_factor, + tau_sobolevs, + ): + # I wonder why? + # Not sure who wrote this but the answer is that when the plasma is + # first initialised (before the first iteration, without temperature + # values etc.) there are no j_blues values so this just prevents + # an error. Aoife. + if len(j_blues) == 0: + return None + macro_atom_data = self._get_macro_atom_data(atomic_data) + if self.initialize: + self.initialize_macro_atom_transition_type_filters( + atomic_data, macro_atom_data + ) + self.transition_probability_coef = ( + self._get_transition_probability_coefs(macro_atom_data) + ) + self.initialize = False + transition_probabilities = self._calculate_transition_probability( + macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor + ) + transition_probabilities = pd.DataFrame( + transition_probabilities, + index=macro_atom_data.transition_line_id, + columns=tau_sobolevs.columns, + ) + return transition_probabilities + + def _calculate_transition_probability( + self, macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor + ): + transition_probabilities = np.empty( + (self.transition_probability_coef.shape[0], beta_sobolev.shape[1]) + ) + # trans_old = self.calculate_transition_probabilities(macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor) + transition_type = macro_atom_data.transition_type.values + lines_idx = macro_atom_data.lines_idx.values + tpos = macro_atom_data.transition_probability.values + util.calculate_transition_probabilities( + tpos, + beta_sobolev.values, + j_blues.values, + stimulated_emission_factor, + transition_type, + lines_idx, + self.block_references, + transition_probabilities, + self.normalize, + ) + return transition_probabilities + + def calculate_transition_probabilities( + self, macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor + ): + transition_probabilities = self.prepare_transition_probabilities( + macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor + ) + return transition_probabilities + + def initialize_macro_atom_transition_type_filters( + self, atomic_data, macro_atom_data + ): + self.transition_up_filter = macro_atom_data.transition_type.values == 1 + self.transition_up_line_filter = macro_atom_data.lines_idx.values[ + self.transition_up_filter + ] + self.block_references = np.hstack( + ( + atomic_data.macro_atom_references.block_references, + len(macro_atom_data), + ) + ) + + @staticmethod + def _get_transition_probability_coefs(macro_atom_data): + return macro_atom_data.transition_probability.values[np.newaxis].T + + def prepare_transition_probabilities( + self, macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor + ): + current_beta_sobolev = beta_sobolev.values.take( + macro_atom_data.lines_idx.values, axis=0, mode="raise" + ) + transition_probabilities = ( + self.transition_probability_coef * current_beta_sobolev + ) + j_blues = j_blues.take( + self.transition_up_line_filter, axis=0, mode="raise" + ) + macro_stimulated_emission = stimulated_emission_factor.take( + self.transition_up_line_filter, axis=0, mode="raise" + ) + transition_probabilities[self.transition_up_filter] *= ( + j_blues * macro_stimulated_emission + ) + return transition_probabilities + + def _normalize_transition_probabilities(self, transition_probabilities): + util.normalize_transition_probabilities( + transition_probabilities, self.block_references + ) + + @staticmethod + def _get_macro_atom_data(atomic_data): + try: + return atomic_data.macro_atom_data + except: + logger.debug( + "Macro Atom Data was not found. Instead returning All Macro Atom Data" + ) + return atomic_data.macro_atom_data_all + + +class NonMarkovChainTransitionProbabilities(TransitionProbabilities): + outputs = ("non_markov_transition_probabilities",) diff --git a/tardis/opacities/macro_atom/util.py b/tardis/opacities/macro_atom/util.py new file mode 100644 index 00000000000..c72781f4b41 --- /dev/null +++ b/tardis/opacities/macro_atom/util.py @@ -0,0 +1,61 @@ +import numpy as np +from numba import njit + +from tardis import constants as const +from tardis.transport.montecarlo import njit_dict + +h_cgs = const.h.cgs.value +c = const.c.to("cm/s").value +kb = const.k_B.cgs.value +inv_c2 = 1 / (c**2) + + +@njit(**njit_dict) +def calculate_transition_probabilities( + transition_probability_coef, + beta_sobolev, + j_blues, + stimulated_emission_factor, + transition_type, + lines_idx, + block_references, + transition_probabilities, + normalize, +): + """ + Calculates transition probabilities for macro_atom interactions + + transition_probability_coef must be a 1D array + transition_type, lines_idx, and block_references must be int-type arrays + beta_sobolev, j_blues,stimulated_emission_factor, and transition_probabilities must be 2D array + """ + norm_factor = np.zeros(transition_probabilities.shape[1]) + + for i in range(transition_probabilities.shape[0]): + line_idx = lines_idx[i] + for j in range(transition_probabilities.shape[1]): + transition_probabilities[i, j] = ( + transition_probability_coef[i] * beta_sobolev[line_idx, j] + ) + if transition_type[i] == 1: + for j in range(transition_probabilities.shape[1]): + transition_probabilities[i, j] *= ( + stimulated_emission_factor[line_idx, j] + * j_blues[line_idx, j] + ) + + if normalize: + for i in range(block_references.shape[0] - 1): + for k in range(transition_probabilities.shape[1]): + norm_factor[k] = 0.0 + for j in range(block_references[i], block_references[i + 1]): + for k in range(transition_probabilities.shape[1]): + norm_factor[k] += transition_probabilities[j, k] + for k in range(transition_probabilities.shape[1]): + if norm_factor[k] != 0.0: + norm_factor[k] = 1 / norm_factor[k] + else: + norm_factor[k] = 1.0 + for j in range(block_references[i], block_references[i + 1]): + for k in range(transition_probabilities.shape[1]): + transition_probabilities[j, k] *= norm_factor[k] diff --git a/tardis/opacities/opacity_state.py b/tardis/opacities/opacity_state.py new file mode 100644 index 00000000000..6bfa5e308c0 --- /dev/null +++ b/tardis/opacities/opacity_state.py @@ -0,0 +1,249 @@ +import numpy as np +from numba import float64, int64 +from numba.experimental import jitclass + +from tardis.opacities.tau_sobolev import calculate_sobolev_line_opacity + +opacity_state_spec = [ + ("electron_density", float64[:]), + ("t_electrons", float64[:]), + ("line_list_nu", float64[:]), + ("tau_sobolev", float64[:, :]), + ("transition_probabilities", float64[:, :]), + ("line2macro_level_upper", int64[:]), + ("macro_block_references", int64[:]), + ("transition_type", int64[:]), + ("destination_level_id", int64[:]), + ("transition_line_id", int64[:]), + ("bf_threshold_list_nu", float64[:]), + ("p_fb_deactivation", float64[:, :]), + ("photo_ion_nu_threshold_mins", float64[:]), + ("photo_ion_nu_threshold_maxs", float64[:]), + ("photo_ion_block_references", int64[:]), + ("chi_bf", float64[:, :]), + ("x_sect", float64[:]), + ("phot_nus", float64[:]), + ("ff_opacity_factor", float64[:]), + ("emissivities", float64[:, :]), + ("photo_ion_activation_idx", int64[:]), + ("k_packet_idx", int64), +] + + +@jitclass(opacity_state_spec) +class OpacityState: + def __init__( + self, + electron_density, + t_electrons, + line_list_nu, + tau_sobolev, + transition_probabilities, + line2macro_level_upper, + macro_block_references, + transition_type, + destination_level_id, + transition_line_id, + bf_threshold_list_nu, + p_fb_deactivation, + photo_ion_nu_threshold_mins, + photo_ion_nu_threshold_maxs, + photo_ion_block_references, + chi_bf, + x_sect, + phot_nus, + ff_opacity_factor, + emissivities, + photo_ion_activation_idx, + k_packet_idx, + ): + """ + Plasma for the Numba code + + Parameters + ---------- + electron_density : numpy.ndarray + t_electrons : numpy.ndarray + line_list_nu : numpy.ndarray + tau_sobolev : numpy.ndarray + transition_probabilities : numpy.ndarray + line2macro_level_upper : numpy.ndarray + macro_block_references : numpy.ndarray + transition_type : numpy.ndarray + destination_level_id : numpy.ndarray + transition_line_id : numpy.ndarray + bf_threshold_list_nu : numpy.ndarray + """ + self.electron_density = electron_density + self.t_electrons = t_electrons + self.line_list_nu = line_list_nu + self.tau_sobolev = tau_sobolev + self.bf_threshold_list_nu = bf_threshold_list_nu + + #### Macro Atom transition probabilities + self.transition_probabilities = transition_probabilities + self.line2macro_level_upper = line2macro_level_upper + + self.macro_block_references = macro_block_references + self.transition_type = transition_type + + # Destination level is not needed and/or generated for downbranch + self.destination_level_id = destination_level_id + self.transition_line_id = transition_line_id + self.p_fb_deactivation = p_fb_deactivation + + # Continuum Opacity Data + self.photo_ion_nu_threshold_mins = photo_ion_nu_threshold_mins + self.photo_ion_nu_threshold_maxs = photo_ion_nu_threshold_maxs + + self.photo_ion_block_references = photo_ion_block_references + self.chi_bf = chi_bf + self.x_sect = x_sect + self.phot_nus = phot_nus + self.ff_opacity_factor = ff_opacity_factor + self.emissivities = emissivities + self.photo_ion_activation_idx = photo_ion_activation_idx + self.k_packet_idx = k_packet_idx + + +def opacity_state_initialize( + plasma, + line_interaction_type, + disable_line_scattering, + continuum_processes_enabled, +): + """ + Initialize the OpacityState object and copy over the data over from TARDIS Plasma + + Parameters + ---------- + plasma : tardis.plasma.BasePlasma + line_interaction_type : enum + """ + electron_densities = plasma.electron_densities.values + t_electrons = plasma.t_electrons + line_list_nu = plasma.atomic_data.lines.nu.values + + tau_sobolev_df = calculate_sobolev_line_opacity( + plasma.atomic_data.lines, + plasma.level_number_density, + plasma.time_explosion, + plasma.stimulated_emission_factor, + ) + + tau_sobolev = np.ascontiguousarray(tau_sobolev_df, dtype=np.float64) + + if disable_line_scattering: + tau_sobolev *= 0 + + if line_interaction_type == "scatter": + # to adhere to data types, we must have an array of minimum size 1 + array_size = 1 + transition_probabilities = np.zeros( + (array_size, array_size), dtype=np.float64 + ) # to adhere to data types + line2macro_level_upper = np.zeros(array_size, dtype=np.int64) + macro_block_references = np.zeros(array_size, dtype=np.int64) + transition_type = np.zeros(array_size, dtype=np.int64) + destination_level_id = np.zeros(array_size, dtype=np.int64) + transition_line_id = np.zeros(array_size, dtype=np.int64) + else: + transition_probabilities = np.ascontiguousarray( + plasma.transition_probabilities.values.copy(), dtype=np.float64 + ) + line2macro_level_upper = ( + plasma.atomic_data.lines_upper2macro_reference_idx + ) + # TODO: Fix setting of block references for non-continuum mode + + if continuum_processes_enabled: + macro_block_references = plasma.macro_block_references + else: + macro_block_references = plasma.atomic_data.macro_atom_references[ + "block_references" + ].values + transition_type = plasma.macro_atom_data["transition_type"].values + + # Destination level is not needed and/or generated for downbranch + destination_level_id = plasma.macro_atom_data[ + "destination_level_idx" + ].values + transition_line_id = plasma.macro_atom_data["lines_idx"].values + if continuum_processes_enabled: + bf_threshold_list_nu = plasma.nu_i.loc[ + plasma.level2continuum_idx.index + ].values + p_fb_deactivation = np.ascontiguousarray( + plasma.p_fb_deactivation.values.copy(), dtype=np.float64 + ) + + phot_nus = plasma.photo_ion_cross_sections.nu.loc[ + plasma.level2continuum_idx.index + ] + photo_ion_block_references = np.pad( + phot_nus.groupby(level=[0, 1, 2], sort=False) + .count() + .values.cumsum(), + [1, 0], + ) + photo_ion_nu_threshold_mins = ( + phot_nus.groupby(level=[0, 1, 2], sort=False).first().values + ) + photo_ion_nu_threshold_maxs = ( + phot_nus.groupby(level=[0, 1, 2], sort=False).last().values + ) + + chi_bf = plasma.chi_bf.loc[plasma.level2continuum_idx.index].values + x_sect = plasma.photo_ion_cross_sections.x_sect.loc[ + plasma.level2continuum_idx.index + ].values + + phot_nus = phot_nus.values + ff_opacity_factor = ( + plasma.ff_cooling_factor / np.sqrt(t_electrons) + ).astype(np.float64) + emissivities = plasma.fb_emission_cdf.loc[ + plasma.level2continuum_idx.index + ].values + photo_ion_activation_idx = plasma.photo_ion_idx.loc[ + plasma.level2continuum_idx.index, "destination_level_idx" + ].values + k_packet_idx = np.int64(plasma.k_packet_idx) + else: + bf_threshold_list_nu = np.zeros(0, dtype=np.float64) + p_fb_deactivation = np.zeros((0, 0), dtype=np.float64) + photo_ion_nu_threshold_mins = np.zeros(0, dtype=np.float64) + photo_ion_nu_threshold_maxs = np.zeros(0, dtype=np.float64) + photo_ion_block_references = np.zeros(0, dtype=np.int64) + chi_bf = np.zeros((0, 0), dtype=np.float64) + x_sect = np.zeros(0, dtype=np.float64) + phot_nus = np.zeros(0, dtype=np.float64) + ff_opacity_factor = np.zeros(0, dtype=np.float64) + emissivities = np.zeros((0, 0), dtype=np.float64) + photo_ion_activation_idx = np.zeros(0, dtype=np.int64) + k_packet_idx = np.int64(-1) + + return OpacityState( + electron_densities, + t_electrons, + line_list_nu, + tau_sobolev, + transition_probabilities, + line2macro_level_upper, + macro_block_references, + transition_type, + destination_level_id, + transition_line_id, + bf_threshold_list_nu, + p_fb_deactivation, + photo_ion_nu_threshold_mins, + photo_ion_nu_threshold_maxs, + photo_ion_block_references, + chi_bf, + x_sect, + phot_nus, + ff_opacity_factor, + emissivities, + photo_ion_activation_idx, + k_packet_idx, + ) From 52d5523d84e210c5b336493f1ddfa023a88bcdf6 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Thu, 9 May 2024 17:01:57 -0400 Subject: [PATCH 10/89] Add calculate_transition_probabilities function to util.py in macro_atom package --- .../macro_atom/transition_probabilities.py | 425 ++++++++++++++++++ 1 file changed, 425 insertions(+) create mode 100644 tardis/opacities/macro_atom/transition_probabilities.py diff --git a/tardis/opacities/macro_atom/transition_probabilities.py b/tardis/opacities/macro_atom/transition_probabilities.py new file mode 100644 index 00000000000..75f841b1eb2 --- /dev/null +++ b/tardis/opacities/macro_atom/transition_probabilities.py @@ -0,0 +1,425 @@ +import logging + +import numpy as np +import pandas as pd + +from scipy import sparse as sp + +from tardis.plasma.properties.base import ProcessingPlasmaProperty +from tardis.plasma.properties.continuum_processes import ( + get_ground_state_multi_index, +) +from tardis.transport.montecarlo.macro_atom import ( + MacroAtomTransitionType, +) + +__all__ = [ + "MarkovChainTransProbs", + "MarkovChainIndex", + "MarkovChainTransProbsCollector", + "NonContinuumTransProbsMask", + "MonteCarloTransProbs", +] + +logger = logging.getLogger(__name__) + + +def normalize_trans_probs(p): + """ + Normalize a set of transition probabilities. + + Parameters + ---------- + p : pandas.DataFrame, dtype float + Unnormalized transition probabilities. Indexed by + source_level_idx, destination_level_idx. + + Returns + ------- + pandas.DataFrame, dtype float + Normalized transition probabilities: the sum of + all probabilites with the same source_level_idx sum to one. + Indexed by source_level_idx, destination_level_idx. + """ + p = p.astype(np.float64) + p_summed = p.groupby(level=0).sum() + p_summed[p_summed == 0] = 1 + index = p.index.get_level_values("source_level_idx") + p_norm = p / p_summed.loc[index].values + assert np.all(np.isfinite(p_norm)) + return p_norm + + +class SpMatrixSeriesConverterMixin(object): + @staticmethod + def series2matrix(series, idx2reduced_idx): + """ + Convert a Pandas Series to a sparse matrix and re-index it. + + Parameters + ---------- + series : pandas.Series, dtype float + Rates or transition probabilities. Indexed by + source_level_idx, destination_level_idx. + idx2reduced_idx : pandas.Series + Values of (compact) matrix index. Indexed by references_idx. + Maps the references_idx of a level to the index + used in the sparse matrix. + + Returns + ------- + scipy.sparse.coo.coo_matrix + Sparse matrix of rates or transition probabilites. + """ + q_indices = ( + series.index.get_level_values(0), + series.index.get_level_values(1), + ) + q_indices = ( + idx2reduced_idx.loc[q_indices[0]].values, + idx2reduced_idx.loc[q_indices[1]].values, + ) + max_idx = idx2reduced_idx.max() + 1 + matrix = sp.coo_matrix( + (series.astype(np.float64), q_indices), shape=(max_idx, max_idx) + ) + return matrix + + @staticmethod + def matrix2series(matrix, idx2reduced_idx, names=None): + """ + Convert a sparse matrix to a Pandas Series and index it. + + Parameters + ---------- + matrix : scipy.sparse.coo.coo_matrix + Sparse matrix of rates or transition probabilites. + idx2reduced_idx : pandas.Series + Values of (compact) matrix index. Indexed by references_idx. + Maps the references_idx of a level to the index + used in the sparse matrix. + names : array-like, optional + Names of levels in MultiIndex of returned Series. + + Returns + ------- + pandas.Series + Rates or transition probabilities. Indexed by + source_level_idx, destination_level_idx. + """ + reduced_idx2idx = pd.Series( + idx2reduced_idx.index, index=idx2reduced_idx + ) + matrix = matrix.tocoo() + index = pd.MultiIndex.from_arrays( + [reduced_idx2idx.loc[matrix.row], reduced_idx2idx.loc[matrix.col]] + ) + series = pd.Series(matrix.data, index=index) + if names: + series.index.names = names + return series + + +class MarkovChainIndex(ProcessingPlasmaProperty): + """ + Attributes + ---------- + idx2mkv_idx : pandas.Series, dtype int + k_packet_idx : int + Macro atom level idx corresponding to a k-packet. + idx2deactivation_idx : pandas.Series, dtype int + """ + + outputs = ("idx2mkv_idx", "k_packet_idx", "idx2deactivation_idx") + + def calculate(self, atomic_data, continuum_interaction_species): + ma_ref = atomic_data.macro_atom_references + mask = ma_ref.index.droplevel("source_level_number").isin( + continuum_interaction_species + ) + mask2 = ma_ref.index.isin( + get_ground_state_multi_index(continuum_interaction_species) + ) + mask = np.logical_or(mask, mask2) + idx = ma_ref[mask].references_idx.values + idx2mkv_idx = pd.Series(np.arange(len(idx)), index=idx) + idx2mkv_idx.loc["k"] = idx2mkv_idx.max() + 1 + + k_packet_idx = ma_ref.references_idx.max() + 1 + + idx2deactivation_idx = idx2mkv_idx + k_packet_idx + 1 + return idx2mkv_idx, k_packet_idx, idx2deactivation_idx + + +class NonContinuumTransProbsMask(ProcessingPlasmaProperty): + """ + Attributes + ---------- + non_continuum_trans_probs_mask : numpy.ndarray, dtype bool + """ + + outputs = ("non_continuum_trans_probs_mask",) + + def calculate(self, atomic_data, continuum_interaction_species): + # I don't have to remove the ground states of + # the next higher ionization states of the continuum species + # since they only contain zero probabilities. + continuum_trans_probs_mask = atomic_data.macro_atom_data.set_index( + ["atomic_number", "ion_number"] + ).index.isin(continuum_interaction_species) + non_continuum_trans_probs_mask = np.logical_not( + continuum_trans_probs_mask + ) + return non_continuum_trans_probs_mask + + +class MarkovChainTransProbsCollector(ProcessingPlasmaProperty): + """ + Attributes + ---------- + p_combined : pandas.DataFrame, dtype float + Combined and normalized transition probabilities. + Indexed by source_level_idx, destination_level_idx. + """ + + outputs = ("p_combined",) + + def __init__(self, plasma_parent, inputs): + super().__init__(plasma_parent) + self.inputs = inputs + + def calculate(self, *args): + p = pd.concat(args) + p = p.groupby(level=[0, 1, 2]).sum() + p = normalize_trans_probs(p) + return p + + +class MarkovChainTransProbs( + ProcessingPlasmaProperty, SpMatrixSeriesConverterMixin +): + outputs = ("N", "R", "B", "p_deactivation") + latex_name = ("N", "R", "B", r"p_\textrm{deactivation}") + """ + Attributes + ---------- + N : pandas.DataFrame, dtype float + Fundamental matrix of the Markov-chain macro atom. + Indexed by source_level_idx, destination_level_idx. + Expected number of visits to destination_level_idx starting + from souce_level_idx (before being absorbed). + R : pandas.DataFrame, dtype float + Deactivation probabilities of the Markov-chain macro atom. + Indexed by source_level_idx. + Probability of deactivation/absorption in source_level_idx. + B : pandas.DataFrame, dtype float + Absorbing probabilities of the Markov-chain macro atom. + Indexed by source_level_idx, destination_level_idx. + Probability of being absorbed in destination_level_idx when + starting from source_level_idx. + p_deactivation : pandas.DataFrame, dtype float + Redistribution probabilities after deactivation of the Markov-chain + macro atom. Indexed by source_level_idx, destination_level_idx. + Probability of an r-packet being emitted in the transition + (source_level_idx --> destination_level_idx) after deactivation + in source_level_idx. + """ + + def calculate(self, p_combined, idx2mkv_idx): + p = p_combined + p_internal = p.xs(0, level="transition_type") + p_deactivation = normalize_trans_probs( + p.xs(-1, level="transition_type") + ) + + N = pd.DataFrame(columns=p_internal.columns) + B = pd.DataFrame(columns=p_internal.columns) + R = pd.DataFrame(columns=p_internal.columns, index=idx2mkv_idx.index) + R.index.name = "source_level_idx" + for column in p_internal: + Q = self.series2matrix(p_internal[column], idx2mkv_idx) + inv_N = sp.identity(Q.shape[0]) - Q + N1 = sp.linalg.inv(inv_N.tocsc()) + R1 = (1 - np.asarray(Q.sum(axis=1))).flatten() + B1 = N1.multiply(R1) + N1 = self.matrix2series( + N1, idx2mkv_idx, names=p_internal.index.names + ) + B1 = self.matrix2series( + B1, idx2mkv_idx, names=p_internal.index.names + ) + N[column] = N1 + B[column] = B1 + R[column] = R1 + N = N.sort_index() + B = B.sort_index() + return N, R, B, p_deactivation + + +class MonteCarloTransProbs(ProcessingPlasmaProperty): + outputs = ( + "non_continuum_trans_probs", + "level_absorption_probs", + "deactivation_channel_probs", + "transition_probabilities", + "macro_block_references", + "macro_atom_data", + ) + """ + Attributes + ---------- + non_continuum_trans_probs + level_absorption_probs + deactivation_channel_probs + transition_probabilities + macro_block_references + macro_atom_data + """ + + def calculate( + self, + non_markov_transition_probabilities, + atomic_data, + non_continuum_trans_probs_mask, + k_packet_idx, + idx2deactivation_idx, + level_idxs2transition_idx, + p_deactivation, + cool_rate_fb, + cool_rate_fb_tot, + level2continuum_idx, + B, + ): + # Prepare the transition probabilities for the non continuum species + macro_atom_data = atomic_data.macro_atom_data + transition_info = macro_atom_data[ + ["lines_idx", "transition_type"] + ].set_index(non_markov_transition_probabilities.index) + non_continuum_trans_probs = pd.concat( + [transition_info, non_markov_transition_probabilities], axis=1 + ) + index = macro_atom_data.set_index( + ["source_level_idx", "destination_level_idx"] + ).index + non_continuum_trans_probs = non_continuum_trans_probs.set_index(index) + non_continuum_trans_probs = non_continuum_trans_probs[ + non_continuum_trans_probs_mask + ] + + # Prepare the level absorption probabilities for the continuum species + level_absorption_probs = B.copy() + level_absorption_probs.insert(0, "lines_idx", -1) + level_absorption_probs.insert(0, "transition_type", 3) + destination_level_idx = level_absorption_probs.index.get_level_values( + "destination_level_idx" + ) + source_level_idx = level_absorption_probs.rename( + index={"k": k_packet_idx} + ).index.get_level_values("source_level_idx") + destination_level_idx = idx2deactivation_idx.loc[destination_level_idx] + absorption_index = pd.MultiIndex.from_arrays( + [source_level_idx, destination_level_idx], names=B.index.names + ) + level_absorption_probs.index = absorption_index + + # Prepare the free-bound cooling probabilities + fb_cooling_probs = ( + cool_rate_fb + / cool_rate_fb_tot.values + * p_deactivation.loc[("k"), ("bf")] + ) + continuum_idx = level2continuum_idx.loc[fb_cooling_probs.index].values + fb_cooling_probs.index = pd.MultiIndex.from_product( + [["k"], np.ones(len(fb_cooling_probs), dtype=int) * -1], + names=p_deactivation.index.names, + ) + fb_cooling_probs.insert(0, "lines_idx", continuum_idx) + fb_cooling_probs.insert( + 0, + "transition_type", + level_idxs2transition_idx.at[("k", "bf"), "transition_type"], + ) + + # Check if there are two-photon decays + if "two-photon" in p_deactivation.index.get_level_values(1): + two_photon_index = p_deactivation[ + p_deactivation.index.get_level_values(1) == "two-photon" + ].index + level_idxs2transition_idx_two_photon = pd.DataFrame( + [[-1, MacroAtomTransitionType.TWO_PHOTON.value]], + index=two_photon_index, + columns=level_idxs2transition_idx.columns, + ) + level_idxs2transition_idx = pd.concat( + [ + level_idxs2transition_idx_two_photon, + level_idxs2transition_idx, + ] + ) + + # Prepare the deactivation channel probabilities for the continuum species + deactivation_channel_probs = p_deactivation.copy() + deactivation_channel_probs = pd.concat( + [ + level_idxs2transition_idx.reindex( + deactivation_channel_probs.index + ), + deactivation_channel_probs, + ], + axis=1, + ).reindex(deactivation_channel_probs.index) + + deactivation_channel_probs = deactivation_channel_probs.drop( + ("k", "bf") + ) + deactivation_channel_probs = pd.concat( + [deactivation_channel_probs, fb_cooling_probs], sort=False + ) + + source_level_idx = deactivation_channel_probs.index.get_level_values( + "source_level_idx" + ) + + source_level_idx = idx2deactivation_idx.loc[source_level_idx] + destination_level_idx = np.ones_like(source_level_idx) * -1 + deactivation_index = pd.MultiIndex.from_arrays( + [source_level_idx, destination_level_idx], names=B.index.names + ) + + deactivation_channel_probs.index = deactivation_index + + # Combine everything + combined_trans_probs = pd.concat( + [ + level_absorption_probs, + deactivation_channel_probs, + non_continuum_trans_probs, + ] + ) + combined_trans_probs = combined_trans_probs.sort_index() + + block_references = ( + combined_trans_probs[0].groupby("source_level_idx").count().cumsum() + ) + continous_index = np.arange(block_references.index.max() + 1) + block_references = ( + block_references.reindex(continous_index).ffill().astype(int) + ) # This is needed because some macro atom levels have no transitions + block_references = np.pad(block_references, (1, 0), constant_values=0.0) + + macro_atom_info = combined_trans_probs[ + ["transition_type", "lines_idx"] + ].reset_index() + + combined_trans_probs = combined_trans_probs.drop( + ["lines_idx", "transition_type"], axis="columns" + ) + + return ( + non_continuum_trans_probs, + level_absorption_probs, + deactivation_channel_probs, + combined_trans_probs, + block_references, + macro_atom_info, + ) From 9530fee2ba5db7ea5ab38dda3270f8288ee2fcd8 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Thu, 9 May 2024 17:02:20 -0400 Subject: [PATCH 11/89] Remove unused imports and update plasma properties --- tardis/opacities/macro_atom/base.py | 9 +- tardis/plasma/properties/__init__.py | 12 +- tardis/plasma/properties/general.py | 3 +- tardis/plasma/properties/j_blues.py | 4 +- .../plasma/properties/radiative_properties.py | 1 - .../properties/transition_probabilities.py | 425 ------------------ tardis/plasma/standard_plasmas.py | 2 +- 7 files changed, 15 insertions(+), 441 deletions(-) delete mode 100644 tardis/plasma/properties/transition_probabilities.py diff --git a/tardis/opacities/macro_atom/base.py b/tardis/opacities/macro_atom/base.py index 251ddeaf792..6d649b8b12c 100644 --- a/tardis/opacities/macro_atom/base.py +++ b/tardis/opacities/macro_atom/base.py @@ -1,11 +1,12 @@ -from tardis.plasma.properties.base import ProcessingPlasmaProperty -from tardis.plasma.properties.radiative_properties import logger -from tardis.opacities.macro_atom import util - +import logging import numpy as np import pandas as pd +from tardis.opacities.macro_atom import util +from tardis.plasma.properties.base import ProcessingPlasmaProperty + +logger = logging.getLogger(__name__) class TransitionProbabilities(ProcessingPlasmaProperty): """ diff --git a/tardis/plasma/properties/__init__.py b/tardis/plasma/properties/__init__.py index 8ec73b81bcb..63d473f01f9 100644 --- a/tardis/plasma/properties/__init__.py +++ b/tardis/plasma/properties/__init__.py @@ -5,17 +5,17 @@ Every property has a calculate function that returns the values of its outputs. """ +from tardis.opacities.macro_atom.transition_probabilities import * from tardis.plasma.properties.atomic import * +from tardis.plasma.properties.continuum_processes import * from tardis.plasma.properties.general import * +from tardis.plasma.properties.helium_nlte import * from tardis.plasma.properties.ion_population import * +from tardis.plasma.properties.j_blues import * from tardis.plasma.properties.level_population import * +from tardis.plasma.properties.nlte import * +from tardis.plasma.properties.nlte_rate_equation_solver import * from tardis.plasma.properties.partition_function import * from tardis.plasma.properties.plasma_input import * from tardis.plasma.properties.radiative_properties import * -from tardis.plasma.properties.nlte import * -from tardis.plasma.properties.j_blues import * -from tardis.plasma.properties.continuum_processes import * -from tardis.plasma.properties.transition_probabilities import * -from tardis.plasma.properties.helium_nlte import * from tardis.plasma.properties.rate_matrix_index import * -from tardis.plasma.properties.nlte_rate_equation_solver import * diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index 55823b0fe14..6b2cb7b68b5 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -1,10 +1,9 @@ import logging import numpy as np -import pandas as pd from astropy import units as u -from tardis import constants as const +from tardis import constants as const from tardis.plasma.properties.base import ProcessingPlasmaProperty logger = logging.getLogger(__name__) diff --git a/tardis/plasma/properties/j_blues.py b/tardis/plasma/properties/j_blues.py index 9ffbe3c1b5c..1f34f0a276b 100644 --- a/tardis/plasma/properties/j_blues.py +++ b/tardis/plasma/properties/j_blues.py @@ -1,10 +1,10 @@ import numpy as np import pandas as pd -from tardis import constants as const +from tardis import constants as const from tardis.plasma.properties.base import ( - ProcessingPlasmaProperty, DataFrameInput, + ProcessingPlasmaProperty, ) from tardis.util.base import intensity_black_body diff --git a/tardis/plasma/properties/radiative_properties.py b/tardis/plasma/properties/radiative_properties.py index 6facf1e8d27..37170b5c9e6 100644 --- a/tardis/plasma/properties/radiative_properties.py +++ b/tardis/plasma/properties/radiative_properties.py @@ -18,7 +18,6 @@ "StimulatedEmissionFactor", "BetaSobolev", "RawRadBoundBoundTransProbs", - "NonMarkovChainTransitionProbabilities", ] C_EINSTEIN = ( diff --git a/tardis/plasma/properties/transition_probabilities.py b/tardis/plasma/properties/transition_probabilities.py deleted file mode 100644 index 75f841b1eb2..00000000000 --- a/tardis/plasma/properties/transition_probabilities.py +++ /dev/null @@ -1,425 +0,0 @@ -import logging - -import numpy as np -import pandas as pd - -from scipy import sparse as sp - -from tardis.plasma.properties.base import ProcessingPlasmaProperty -from tardis.plasma.properties.continuum_processes import ( - get_ground_state_multi_index, -) -from tardis.transport.montecarlo.macro_atom import ( - MacroAtomTransitionType, -) - -__all__ = [ - "MarkovChainTransProbs", - "MarkovChainIndex", - "MarkovChainTransProbsCollector", - "NonContinuumTransProbsMask", - "MonteCarloTransProbs", -] - -logger = logging.getLogger(__name__) - - -def normalize_trans_probs(p): - """ - Normalize a set of transition probabilities. - - Parameters - ---------- - p : pandas.DataFrame, dtype float - Unnormalized transition probabilities. Indexed by - source_level_idx, destination_level_idx. - - Returns - ------- - pandas.DataFrame, dtype float - Normalized transition probabilities: the sum of - all probabilites with the same source_level_idx sum to one. - Indexed by source_level_idx, destination_level_idx. - """ - p = p.astype(np.float64) - p_summed = p.groupby(level=0).sum() - p_summed[p_summed == 0] = 1 - index = p.index.get_level_values("source_level_idx") - p_norm = p / p_summed.loc[index].values - assert np.all(np.isfinite(p_norm)) - return p_norm - - -class SpMatrixSeriesConverterMixin(object): - @staticmethod - def series2matrix(series, idx2reduced_idx): - """ - Convert a Pandas Series to a sparse matrix and re-index it. - - Parameters - ---------- - series : pandas.Series, dtype float - Rates or transition probabilities. Indexed by - source_level_idx, destination_level_idx. - idx2reduced_idx : pandas.Series - Values of (compact) matrix index. Indexed by references_idx. - Maps the references_idx of a level to the index - used in the sparse matrix. - - Returns - ------- - scipy.sparse.coo.coo_matrix - Sparse matrix of rates or transition probabilites. - """ - q_indices = ( - series.index.get_level_values(0), - series.index.get_level_values(1), - ) - q_indices = ( - idx2reduced_idx.loc[q_indices[0]].values, - idx2reduced_idx.loc[q_indices[1]].values, - ) - max_idx = idx2reduced_idx.max() + 1 - matrix = sp.coo_matrix( - (series.astype(np.float64), q_indices), shape=(max_idx, max_idx) - ) - return matrix - - @staticmethod - def matrix2series(matrix, idx2reduced_idx, names=None): - """ - Convert a sparse matrix to a Pandas Series and index it. - - Parameters - ---------- - matrix : scipy.sparse.coo.coo_matrix - Sparse matrix of rates or transition probabilites. - idx2reduced_idx : pandas.Series - Values of (compact) matrix index. Indexed by references_idx. - Maps the references_idx of a level to the index - used in the sparse matrix. - names : array-like, optional - Names of levels in MultiIndex of returned Series. - - Returns - ------- - pandas.Series - Rates or transition probabilities. Indexed by - source_level_idx, destination_level_idx. - """ - reduced_idx2idx = pd.Series( - idx2reduced_idx.index, index=idx2reduced_idx - ) - matrix = matrix.tocoo() - index = pd.MultiIndex.from_arrays( - [reduced_idx2idx.loc[matrix.row], reduced_idx2idx.loc[matrix.col]] - ) - series = pd.Series(matrix.data, index=index) - if names: - series.index.names = names - return series - - -class MarkovChainIndex(ProcessingPlasmaProperty): - """ - Attributes - ---------- - idx2mkv_idx : pandas.Series, dtype int - k_packet_idx : int - Macro atom level idx corresponding to a k-packet. - idx2deactivation_idx : pandas.Series, dtype int - """ - - outputs = ("idx2mkv_idx", "k_packet_idx", "idx2deactivation_idx") - - def calculate(self, atomic_data, continuum_interaction_species): - ma_ref = atomic_data.macro_atom_references - mask = ma_ref.index.droplevel("source_level_number").isin( - continuum_interaction_species - ) - mask2 = ma_ref.index.isin( - get_ground_state_multi_index(continuum_interaction_species) - ) - mask = np.logical_or(mask, mask2) - idx = ma_ref[mask].references_idx.values - idx2mkv_idx = pd.Series(np.arange(len(idx)), index=idx) - idx2mkv_idx.loc["k"] = idx2mkv_idx.max() + 1 - - k_packet_idx = ma_ref.references_idx.max() + 1 - - idx2deactivation_idx = idx2mkv_idx + k_packet_idx + 1 - return idx2mkv_idx, k_packet_idx, idx2deactivation_idx - - -class NonContinuumTransProbsMask(ProcessingPlasmaProperty): - """ - Attributes - ---------- - non_continuum_trans_probs_mask : numpy.ndarray, dtype bool - """ - - outputs = ("non_continuum_trans_probs_mask",) - - def calculate(self, atomic_data, continuum_interaction_species): - # I don't have to remove the ground states of - # the next higher ionization states of the continuum species - # since they only contain zero probabilities. - continuum_trans_probs_mask = atomic_data.macro_atom_data.set_index( - ["atomic_number", "ion_number"] - ).index.isin(continuum_interaction_species) - non_continuum_trans_probs_mask = np.logical_not( - continuum_trans_probs_mask - ) - return non_continuum_trans_probs_mask - - -class MarkovChainTransProbsCollector(ProcessingPlasmaProperty): - """ - Attributes - ---------- - p_combined : pandas.DataFrame, dtype float - Combined and normalized transition probabilities. - Indexed by source_level_idx, destination_level_idx. - """ - - outputs = ("p_combined",) - - def __init__(self, plasma_parent, inputs): - super().__init__(plasma_parent) - self.inputs = inputs - - def calculate(self, *args): - p = pd.concat(args) - p = p.groupby(level=[0, 1, 2]).sum() - p = normalize_trans_probs(p) - return p - - -class MarkovChainTransProbs( - ProcessingPlasmaProperty, SpMatrixSeriesConverterMixin -): - outputs = ("N", "R", "B", "p_deactivation") - latex_name = ("N", "R", "B", r"p_\textrm{deactivation}") - """ - Attributes - ---------- - N : pandas.DataFrame, dtype float - Fundamental matrix of the Markov-chain macro atom. - Indexed by source_level_idx, destination_level_idx. - Expected number of visits to destination_level_idx starting - from souce_level_idx (before being absorbed). - R : pandas.DataFrame, dtype float - Deactivation probabilities of the Markov-chain macro atom. - Indexed by source_level_idx. - Probability of deactivation/absorption in source_level_idx. - B : pandas.DataFrame, dtype float - Absorbing probabilities of the Markov-chain macro atom. - Indexed by source_level_idx, destination_level_idx. - Probability of being absorbed in destination_level_idx when - starting from source_level_idx. - p_deactivation : pandas.DataFrame, dtype float - Redistribution probabilities after deactivation of the Markov-chain - macro atom. Indexed by source_level_idx, destination_level_idx. - Probability of an r-packet being emitted in the transition - (source_level_idx --> destination_level_idx) after deactivation - in source_level_idx. - """ - - def calculate(self, p_combined, idx2mkv_idx): - p = p_combined - p_internal = p.xs(0, level="transition_type") - p_deactivation = normalize_trans_probs( - p.xs(-1, level="transition_type") - ) - - N = pd.DataFrame(columns=p_internal.columns) - B = pd.DataFrame(columns=p_internal.columns) - R = pd.DataFrame(columns=p_internal.columns, index=idx2mkv_idx.index) - R.index.name = "source_level_idx" - for column in p_internal: - Q = self.series2matrix(p_internal[column], idx2mkv_idx) - inv_N = sp.identity(Q.shape[0]) - Q - N1 = sp.linalg.inv(inv_N.tocsc()) - R1 = (1 - np.asarray(Q.sum(axis=1))).flatten() - B1 = N1.multiply(R1) - N1 = self.matrix2series( - N1, idx2mkv_idx, names=p_internal.index.names - ) - B1 = self.matrix2series( - B1, idx2mkv_idx, names=p_internal.index.names - ) - N[column] = N1 - B[column] = B1 - R[column] = R1 - N = N.sort_index() - B = B.sort_index() - return N, R, B, p_deactivation - - -class MonteCarloTransProbs(ProcessingPlasmaProperty): - outputs = ( - "non_continuum_trans_probs", - "level_absorption_probs", - "deactivation_channel_probs", - "transition_probabilities", - "macro_block_references", - "macro_atom_data", - ) - """ - Attributes - ---------- - non_continuum_trans_probs - level_absorption_probs - deactivation_channel_probs - transition_probabilities - macro_block_references - macro_atom_data - """ - - def calculate( - self, - non_markov_transition_probabilities, - atomic_data, - non_continuum_trans_probs_mask, - k_packet_idx, - idx2deactivation_idx, - level_idxs2transition_idx, - p_deactivation, - cool_rate_fb, - cool_rate_fb_tot, - level2continuum_idx, - B, - ): - # Prepare the transition probabilities for the non continuum species - macro_atom_data = atomic_data.macro_atom_data - transition_info = macro_atom_data[ - ["lines_idx", "transition_type"] - ].set_index(non_markov_transition_probabilities.index) - non_continuum_trans_probs = pd.concat( - [transition_info, non_markov_transition_probabilities], axis=1 - ) - index = macro_atom_data.set_index( - ["source_level_idx", "destination_level_idx"] - ).index - non_continuum_trans_probs = non_continuum_trans_probs.set_index(index) - non_continuum_trans_probs = non_continuum_trans_probs[ - non_continuum_trans_probs_mask - ] - - # Prepare the level absorption probabilities for the continuum species - level_absorption_probs = B.copy() - level_absorption_probs.insert(0, "lines_idx", -1) - level_absorption_probs.insert(0, "transition_type", 3) - destination_level_idx = level_absorption_probs.index.get_level_values( - "destination_level_idx" - ) - source_level_idx = level_absorption_probs.rename( - index={"k": k_packet_idx} - ).index.get_level_values("source_level_idx") - destination_level_idx = idx2deactivation_idx.loc[destination_level_idx] - absorption_index = pd.MultiIndex.from_arrays( - [source_level_idx, destination_level_idx], names=B.index.names - ) - level_absorption_probs.index = absorption_index - - # Prepare the free-bound cooling probabilities - fb_cooling_probs = ( - cool_rate_fb - / cool_rate_fb_tot.values - * p_deactivation.loc[("k"), ("bf")] - ) - continuum_idx = level2continuum_idx.loc[fb_cooling_probs.index].values - fb_cooling_probs.index = pd.MultiIndex.from_product( - [["k"], np.ones(len(fb_cooling_probs), dtype=int) * -1], - names=p_deactivation.index.names, - ) - fb_cooling_probs.insert(0, "lines_idx", continuum_idx) - fb_cooling_probs.insert( - 0, - "transition_type", - level_idxs2transition_idx.at[("k", "bf"), "transition_type"], - ) - - # Check if there are two-photon decays - if "two-photon" in p_deactivation.index.get_level_values(1): - two_photon_index = p_deactivation[ - p_deactivation.index.get_level_values(1) == "two-photon" - ].index - level_idxs2transition_idx_two_photon = pd.DataFrame( - [[-1, MacroAtomTransitionType.TWO_PHOTON.value]], - index=two_photon_index, - columns=level_idxs2transition_idx.columns, - ) - level_idxs2transition_idx = pd.concat( - [ - level_idxs2transition_idx_two_photon, - level_idxs2transition_idx, - ] - ) - - # Prepare the deactivation channel probabilities for the continuum species - deactivation_channel_probs = p_deactivation.copy() - deactivation_channel_probs = pd.concat( - [ - level_idxs2transition_idx.reindex( - deactivation_channel_probs.index - ), - deactivation_channel_probs, - ], - axis=1, - ).reindex(deactivation_channel_probs.index) - - deactivation_channel_probs = deactivation_channel_probs.drop( - ("k", "bf") - ) - deactivation_channel_probs = pd.concat( - [deactivation_channel_probs, fb_cooling_probs], sort=False - ) - - source_level_idx = deactivation_channel_probs.index.get_level_values( - "source_level_idx" - ) - - source_level_idx = idx2deactivation_idx.loc[source_level_idx] - destination_level_idx = np.ones_like(source_level_idx) * -1 - deactivation_index = pd.MultiIndex.from_arrays( - [source_level_idx, destination_level_idx], names=B.index.names - ) - - deactivation_channel_probs.index = deactivation_index - - # Combine everything - combined_trans_probs = pd.concat( - [ - level_absorption_probs, - deactivation_channel_probs, - non_continuum_trans_probs, - ] - ) - combined_trans_probs = combined_trans_probs.sort_index() - - block_references = ( - combined_trans_probs[0].groupby("source_level_idx").count().cumsum() - ) - continous_index = np.arange(block_references.index.max() + 1) - block_references = ( - block_references.reindex(continous_index).ffill().astype(int) - ) # This is needed because some macro atom levels have no transitions - block_references = np.pad(block_references, (1, 0), constant_values=0.0) - - macro_atom_info = combined_trans_probs[ - ["transition_type", "lines_idx"] - ].reset_index() - - combined_trans_probs = combined_trans_probs.drop( - ["lines_idx", "transition_type"], axis="columns" - ) - - return ( - non_continuum_trans_probs, - level_absorption_probs, - deactivation_channel_probs, - combined_trans_probs, - block_references, - macro_atom_info, - ) diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index da772508140..2a72c5eb6fa 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -60,7 +60,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): Parameters ---------- config : io.config_reader.Configuration - model : model.SimulationState + simulation_state : model.SimulationState atom_data : atomic.AtomData If None, an attempt will be made to read the atomic data from config. From 2fb7f92763349da97dc489f6de9e8aee55d989d0 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Thu, 9 May 2024 17:43:01 -0400 Subject: [PATCH 12/89] add __init__ to macroatom --- tardis/opacities/macro_atom/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 tardis/opacities/macro_atom/__init__.py diff --git a/tardis/opacities/macro_atom/__init__.py b/tardis/opacities/macro_atom/__init__.py new file mode 100644 index 00000000000..e69de29bb2d From 6ac3e883652fe694aa0ba18c74e0ec5325f9170f Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 10 May 2024 14:18:51 -0400 Subject: [PATCH 13/89] blackify tardis --- tardis/opacities/macro_atom/base.py | 1 + tardis/plasma/properties/atomic.py | 1 + tardis/plasma/properties/general.py | 1 + 3 files changed, 3 insertions(+) diff --git a/tardis/opacities/macro_atom/base.py b/tardis/opacities/macro_atom/base.py index 6d649b8b12c..5bfc7e1bb52 100644 --- a/tardis/opacities/macro_atom/base.py +++ b/tardis/opacities/macro_atom/base.py @@ -8,6 +8,7 @@ logger = logging.getLogger(__name__) + class TransitionProbabilities(ProcessingPlasmaProperty): """ Attributes diff --git a/tardis/plasma/properties/atomic.py b/tardis/plasma/properties/atomic.py index c917c08c421..72fd4b6b472 100644 --- a/tardis/plasma/properties/atomic.py +++ b/tardis/plasma/properties/atomic.py @@ -536,6 +536,7 @@ def calculate(self, level_idxs2line_idx, level_idxs2continuum_idx): return level_idxs2transition_idx + class IonizationData(BaseAtomicDataProperty): """ Attributes diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index 6b2cb7b68b5..cfc56fba66c 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -75,6 +75,7 @@ class ThermalGElectron(GElectron): def calculate(self, beta_electron): return super(ThermalGElectron, self).calculate(beta_electron) + class SelectedAtoms(ProcessingPlasmaProperty): """ Attributes From 801c92b1055d8633b076b3d221d95430c7988b5a Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 10 May 2024 14:56:22 -0400 Subject: [PATCH 14/89] blackified --- tardis/plasma/properties/atomic.py | 1 + tardis/plasma/properties/general.py | 1 + 2 files changed, 2 insertions(+) diff --git a/tardis/plasma/properties/atomic.py b/tardis/plasma/properties/atomic.py index c917c08c421..72fd4b6b472 100644 --- a/tardis/plasma/properties/atomic.py +++ b/tardis/plasma/properties/atomic.py @@ -536,6 +536,7 @@ def calculate(self, level_idxs2line_idx, level_idxs2continuum_idx): return level_idxs2transition_idx + class IonizationData(BaseAtomicDataProperty): """ Attributes diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index 55823b0fe14..660991c1d32 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -76,6 +76,7 @@ class ThermalGElectron(GElectron): def calculate(self, beta_electron): return super(ThermalGElectron, self).calculate(beta_electron) + class SelectedAtoms(ProcessingPlasmaProperty): """ Attributes From ae0650dcab4fad00e8b7cbee6d70f2700cedc634 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 10 May 2024 17:40:02 -0400 Subject: [PATCH 15/89] chore: Update imports and remove unused code --- tardis/model/base.py | 2 +- tardis/model/parse_input.py | 2 +- tardis/model/radiation_field_state.py | 70 ------------------- tardis/radiation_field/base.py | 34 --------- tardis/radiation_field/opacities/__init__.py | 1 - .../continuum_radfield_properties.py | 2 +- .../estimators/dilute_blackbody_properties.py | 2 +- 7 files changed, 4 insertions(+), 109 deletions(-) delete mode 100644 tardis/model/radiation_field_state.py delete mode 100644 tardis/radiation_field/base.py delete mode 100644 tardis/radiation_field/opacities/__init__.py diff --git a/tardis/model/base.py b/tardis/model/base.py index adb86b99bd6..a55361c0b86 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -23,7 +23,7 @@ parse_packet_source, ) from tardis.transport.montecarlo.packet_source import BlackBodySimpleSource -from tardis.model.radiation_field_state import ( +from tardis.radiation_field.planck_rad_field import ( DiluteBlackBodyRadiationFieldState, ) from tardis.util.base import is_valid_nuclide_or_elem diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py index 935e1994be1..66f12ad49ca 100644 --- a/tardis/model/parse_input.py +++ b/tardis/model/parse_input.py @@ -17,7 +17,7 @@ from tardis.model.geometry.radial1d import HomologousRadial1DGeometry from tardis.model.matter.composition import Composition from tardis.model.matter.decay import IsotopicMassFraction -from tardis.model.radiation_field_state import ( +from tardis.radiation_field.planck_rad_field import ( DiluteBlackBodyRadiationFieldState, ) from tardis.transport.montecarlo.packet_source import ( diff --git a/tardis/model/radiation_field_state.py b/tardis/model/radiation_field_state.py deleted file mode 100644 index dff0bd65bbc..00000000000 --- a/tardis/model/radiation_field_state.py +++ /dev/null @@ -1,70 +0,0 @@ -import numpy as np -from astropy import units as u - -from tardis.util.base import intensity_black_body - -from typing import Union - - -class DiluteBlackBodyRadiationFieldState: - """ - Represents the state of a dilute thermal radiation field. - - - Parameters - ---------- - t_radiative : u.Quantity - Radiative temperature in each shell - dilution_factor : numpy.ndarray - Dilution Factors in each shell - geometry: tardis.model.Radial1DModel - The geometry of the model that uses to constrains the active shells - """ - - def __init__( - self, - t_radiative: u.Quantity, - dilution_factor: np.ndarray, - geometry=None, - ): - # ensuring that the radiation_field has both - # dilution_factor and t_radiative equal length - assert len(t_radiative) == len(dilution_factor) - if ( - geometry is not None - ): # check the active shells only (this is used when setting up the radiation_field_state) - assert np.all( - t_radiative[ - geometry.v_inner_boundary_index : geometry.v_outer_boundary_index - ] - > 0 * u.K - ) - assert np.all( - dilution_factor[ - geometry.v_inner_boundary_index : geometry.v_outer_boundary_index - ] - > 0 - ) - else: - assert np.all(t_radiative > 0 * u.K) - assert np.all(dilution_factor > 0) - self.t_radiative = t_radiative - self.dilution_factor = dilution_factor - - def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]): - """ - Calculate the intensity of the radiation field at a given frequency. - - Parameters - ---------- - nu : u.Quantity - Frequency at which the intensity is to be calculated - - Returns - ------- - intensity : u.Quantity - Intensity of the radiation field at the given frequency - """ - return self.dilution_factor * intensity_black_body( - nu[np.newaxis].T, self.t_radiative - ) diff --git a/tardis/radiation_field/base.py b/tardis/radiation_field/base.py deleted file mode 100644 index 4fe0658c20c..00000000000 --- a/tardis/radiation_field/base.py +++ /dev/null @@ -1,34 +0,0 @@ -import numpy as np -from astropy import units as u - -from tardis.transport.montecarlo.packet_source import BasePacketSource -from tardis.opacities.opacity_state import OpacityState - - -class MonteCarloRadiationFieldState: - """_summary_ - - Parameters - ---------- - t_radiative : u.Quantity - Radiative temperature in each shell - dilution_factor : numpy.ndarray - Dilution Factors in each shell - opacities : OpacityState - Opacity container object - packet_source : SourceFunction - Source function for radiative transfer, for example a packet_source - """ - - def __init__( - self, - t_radiative: u.Quantity, - dilution_factor: np.ndarray, - opacities: OpacityState, - packet_source: BasePacketSource, - ): - self.t_radiative = t_radiative - self.dilution_factor = dilution_factor - self.t_rad = self.t_radiative - self.opacities = opacities - self.packet_source = packet_source diff --git a/tardis/radiation_field/opacities/__init__.py b/tardis/radiation_field/opacities/__init__.py deleted file mode 100644 index 47f2c79e60d..00000000000 --- a/tardis/radiation_field/opacities/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from tardis.radiation_field.opacities.base import * diff --git a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py index dcd3b9e1560..64b0064e3e7 100644 --- a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py +++ b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py @@ -6,7 +6,7 @@ import tardis.constants as const from tardis.io.atom_data import AtomData -from tardis.model.radiation_field_state import ( +from tardis.radiation_field.planck_rad_field import ( DiluteBlackBodyRadiationFieldState, ) from tardis.transport.montecarlo.estimators.radfield_mc_estimators import ( diff --git a/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py b/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py index 8e7c98c7ed7..b097356a6aa 100644 --- a/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py +++ b/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py @@ -3,7 +3,7 @@ from scipy.special import zeta from tardis import constants as const -from tardis.model.radiation_field_state import ( +from tardis.radiation_field.planck_rad_field import ( DiluteBlackBodyRadiationFieldState, ) From 69dc1b5dd3c67d139b025623aed06f88ef0703de Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 10 May 2024 17:40:10 -0400 Subject: [PATCH 16/89] chore: Add PlanckRadiationField and DilutePlanckRadiationField classes --- tardis/radiation_field/planck_rad_field.py | 88 ++++++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 tardis/radiation_field/planck_rad_field.py diff --git a/tardis/radiation_field/planck_rad_field.py b/tardis/radiation_field/planck_rad_field.py new file mode 100644 index 00000000000..7e4926b65f1 --- /dev/null +++ b/tardis/radiation_field/planck_rad_field.py @@ -0,0 +1,88 @@ +from astropy import units as u +import numpy as np +from tardis.util.base import intensity_black_body + + +class PlanckRadiationField: + def __init__(self, temperature) -> None: + self.temperature = u.Quantity(temperature, u.K) + + def calculate_mean_intensity(self, nu): + return intensity_black_body( + nu.values[np.newaxis].T, self.temperature.value + ) + + +class DilutePlanckRadiationField: + def __init__(self, temperature, dilution_factor) -> None: + self.temperature = u.Quantity(temperature, u.K) + self.dilution_factor = dilution_factor + + def calculate_mean_intensity(self, nu): + return self.dilution_factor * intensity_black_body( + nu.values[np.newaxis].T, self.temperature.value + ) + + +class DiluteBlackBodyRadiationFieldState: + """ + Represents the state of a dilute thermal radiation field. + + + Parameters + ---------- + t_radiative : u.Quantity + Radiative temperature in each shell + dilution_factor : numpy.ndarray + Dilution Factors in each shell + geometry: tardis.model.Radial1DModel + The geometry of the model that uses to constrains the active shells + """ + + def __init__( + self, + t_radiative: u.Quantity, + dilution_factor: np.ndarray, + geometry=None, + ): + # ensuring that the radiation_field has both + # dilution_factor and t_radiative equal length + assert len(t_radiative) == len(dilution_factor) + if ( + geometry is not None + ): # check the active shells only (this is used when setting up the radiation_field_state) + assert np.all( + t_radiative[ + geometry.v_inner_boundary_index : geometry.v_outer_boundary_index + ] + > 0 * u.K + ) + assert np.all( + dilution_factor[ + geometry.v_inner_boundary_index : geometry.v_outer_boundary_index + ] + > 0 + ) + else: + assert np.all(t_radiative > 0 * u.K) + assert np.all(dilution_factor > 0) + self.t_radiative = t_radiative + self.dilution_factor = dilution_factor + + def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]): + """ + Calculate the intensity of the radiation field at a given frequency. + + Parameters + ---------- + nu : u.Quantity + Frequency at which the intensity is to be calculated + + Returns + ------- + intensity : u.Quantity + Intensity of the radiation field at the given frequency + """ + return self.dilution_factor * intensity_black_body( + nu[np.newaxis].T, self.t_radiative + ) From 448d5c481fe6a1d57c34b0e7ca58c7efe2e01df9 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 10 May 2024 17:44:27 -0400 Subject: [PATCH 17/89] chore: Update imports and remove unused code --- tardis/model/base.py | 7 +------ tardis/radiation_field/planck_rad_field.py | 5 ++++- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/tardis/model/base.py b/tardis/model/base.py index a55361c0b86..41cf7ca3cc6 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -5,7 +5,6 @@ import numpy as np from astropy import units as u -from tardis import constants from tardis.io.configuration.config_reader import Configuration from tardis.io.configuration.config_validator import validate_dict from tardis.io.model.readers.csvy import ( @@ -18,13 +17,9 @@ parse_csvy_composition, parse_csvy_geometry, parse_csvy_radiation_field_state, + parse_packet_source, parse_radiation_field_state, parse_structure_config, - parse_packet_source, -) -from tardis.transport.montecarlo.packet_source import BlackBodySimpleSource -from tardis.radiation_field.planck_rad_field import ( - DiluteBlackBodyRadiationFieldState, ) from tardis.util.base import is_valid_nuclide_or_elem diff --git a/tardis/radiation_field/planck_rad_field.py b/tardis/radiation_field/planck_rad_field.py index 7e4926b65f1..c643440ca7c 100644 --- a/tardis/radiation_field/planck_rad_field.py +++ b/tardis/radiation_field/planck_rad_field.py @@ -1,5 +1,8 @@ -from astropy import units as u +from typing import Union + import numpy as np +from astropy import units as u + from tardis.util.base import intensity_black_body From c1dc52daf979332d190309e49223e55ae9e11ea3 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 13 May 2024 10:31:55 -0400 Subject: [PATCH 18/89] removed density --- tardis/plasma/properties/plasma_input.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index 288e44042e7..ae8b6ab4adc 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -10,7 +10,6 @@ "Abundance", "NumberDensity", "IsotopeAbundance", - "Density", "TimeExplosion", "JBlueEstimator", "LinkTRadTElectron", @@ -80,18 +79,6 @@ class IsotopeAbundance(Input): outputs = ("isotope_abundance",) -class Density(ArrayInput): - """ - Attributes - ---------- - density : Numpy array, dtype float - Total density values - """ - - outputs = ("density",) - latex_name = (r"\rho",) - - class TimeExplosion(Input): """ Attributes From 5e020545d10e43cca181a73d863dac4820178081 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Thu, 6 Jun 2024 15:51:11 -0400 Subject: [PATCH 19/89] ruff output --- tardis/opacities/macro_atom/transition_probabilities.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tardis/opacities/macro_atom/transition_probabilities.py b/tardis/opacities/macro_atom/transition_probabilities.py index 75f841b1eb2..6facd5ee70a 100644 --- a/tardis/opacities/macro_atom/transition_probabilities.py +++ b/tardis/opacities/macro_atom/transition_probabilities.py @@ -2,7 +2,6 @@ import numpy as np import pandas as pd - from scipy import sparse as sp from tardis.plasma.properties.base import ProcessingPlasmaProperty @@ -50,7 +49,7 @@ def normalize_trans_probs(p): return p_norm -class SpMatrixSeriesConverterMixin(object): +class SpMatrixSeriesConverterMixin: @staticmethod def series2matrix(series, idx2reduced_idx): """ From 28c43485b86b081196c2e402f2cc37e519f747c8 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Thu, 6 Jun 2024 16:59:03 -0400 Subject: [PATCH 20/89] cleanup and adding object mode --- tardis/plasma/properties/base.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tardis/plasma/properties/base.py b/tardis/plasma/properties/base.py index 50de6154d08..4c790b596d3 100644 --- a/tardis/plasma/properties/base.py +++ b/tardis/plasma/properties/base.py @@ -1,6 +1,6 @@ import logging -from abc import ABCMeta, abstractmethod, abstractproperty +from abc import ABCMeta, abstractmethod import numpy as np import pandas as pd @@ -31,7 +31,8 @@ class BasePlasmaProperty(object, metaclass=ABCMeta): Used to label nodes when plotting graphs """ - @abstractproperty + @property + @abstractmethod def outputs(self): pass @@ -144,7 +145,8 @@ class TransitionProbabilitiesProperty( track all transition probabilities and to later combine them. """ - @abstractproperty + @property + @abstractmethod def transition_probabilities_outputs(self): pass @@ -202,6 +204,12 @@ def _set_output_value(self, output, value): setattr(self, output, np.array(value, copy=False)) +class ObjectInput(Input): + def set_value(self, value): + for output in self.outputs: + self._set_output_value(self, output, getattr(value, output)) + + class DataFrameInput(Input): def _set_output_value(self, output, value): setattr(self, output, np.array(pd.DataFrame(value), copy=False)) From cfffe5bf9f37f926e71950040b98f396b674ba5f Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 7 Jun 2024 15:22:53 -0400 Subject: [PATCH 21/89] starting to make radiation_field a thing --- tardis/model/parse_input.py | 10 ++--- tardis/plasma/properties/base.py | 15 ++++--- tardis/plasma/properties/plasma_input.py | 10 +++-- tardis/plasma/standard_plasmas.py | 7 ++++ tardis/radiation_field/planck_rad_field.py | 41 ++++++------------- .../continuum_radfield_properties.py | 8 ++-- .../estimators/dilute_blackbody_properties.py | 4 +- .../montecarlo/montecarlo_transport_state.py | 2 +- 8 files changed, 45 insertions(+), 52 deletions(-) diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py index 66f12ad49ca..c3cf8cafcf4 100644 --- a/tardis/model/parse_input.py +++ b/tardis/model/parse_input.py @@ -18,7 +18,7 @@ from tardis.model.matter.composition import Composition from tardis.model.matter.decay import IsotopicMassFraction from tardis.radiation_field.planck_rad_field import ( - DiluteBlackBodyRadiationFieldState, + DilutePlanckianRadiationField, ) from tardis.transport.montecarlo.packet_source import ( BlackBodySimpleSource, @@ -581,9 +581,7 @@ def parse_radiation_field_state( assert len(dilution_factor) == geometry.no_of_shells - return DiluteBlackBodyRadiationFieldState( - t_radiative, dilution_factor, geometry - ) + return DilutePlanckianRadiationField(t_radiative, dilution_factor, geometry) def initialize_packet_source( @@ -714,9 +712,7 @@ def parse_csvy_radiation_field_state( 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): diff --git a/tardis/plasma/properties/base.py b/tardis/plasma/properties/base.py index 4c790b596d3..1bac9e7658e 100644 --- a/tardis/plasma/properties/base.py +++ b/tardis/plasma/properties/base.py @@ -1,10 +1,9 @@ import logging - from abc import ABCMeta, abstractmethod + 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 ---------- @@ -115,7 +114,7 @@ def update(self): 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): @@ -205,9 +204,15 @@ def _set_output_value(self, output, value): class ObjectInput(Input): + + input_object_map = {} # mapping output names from input object attributes def set_value(self, value): for output in self.outputs: - self._set_output_value(self, output, getattr(value, output)) + if output in self.input_object_map: + object_attr = self.input_object_map[output] + self._set_output_value(output, getattr(value, object_attr)) + else: + self._set_output_value(output, getattr(value, output)) class DataFrameInput(Input): diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index 288e44042e7..44ab2622527 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -1,7 +1,4 @@ -from tardis.plasma.properties.base import ( - ArrayInput, - Input, -) +from tardis.plasma.properties.base import ArrayInput, Input, ObjectInput __all__ = [ "TRadiative", @@ -176,3 +173,8 @@ class NumberDensity(Input): outputs = ("number_density",) latex_name = ("N_{i}",) + + +class DilutePlanckianRadFieldInput(ObjectInput): + input_object_map = {"t_rad": "temperature_kelvin", "w": "dilution_factor"} + outputs = ("t_rad", "w") diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 2a72c5eb6fa..92322a38a3f 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -70,6 +70,13 @@ def assemble_plasma(config, simulation_state, atom_data=None): : plasma.BasePlasma """ + + if (config.plasma.ionization == "nebular") or ( + config.plasma.excitation == "dilute-lte" + ): + radiation_field = "dilute_planckian_radiation_field" + else: + pass # Convert the nlte species list to a proper format. nlte_species = [ species_string_to_tuple(s) for s in config.plasma.nlte.species diff --git a/tardis/radiation_field/planck_rad_field.py b/tardis/radiation_field/planck_rad_field.py index c643440ca7c..be58a0de665 100644 --- a/tardis/radiation_field/planck_rad_field.py +++ b/tardis/radiation_field/planck_rad_field.py @@ -6,35 +6,14 @@ from tardis.util.base import intensity_black_body -class PlanckRadiationField: - def __init__(self, temperature) -> None: - self.temperature = u.Quantity(temperature, u.K) - - def calculate_mean_intensity(self, nu): - return intensity_black_body( - nu.values[np.newaxis].T, self.temperature.value - ) - - -class DilutePlanckRadiationField: - def __init__(self, temperature, dilution_factor) -> None: - self.temperature = u.Quantity(temperature, u.K) - self.dilution_factor = dilution_factor - - def calculate_mean_intensity(self, nu): - return self.dilution_factor * intensity_black_body( - nu.values[np.newaxis].T, self.temperature.value - ) - - -class DiluteBlackBodyRadiationFieldState: +class DilutePlanckianRadiationField: """ Represents the state of a dilute thermal radiation field. Parameters ---------- - t_radiative : u.Quantity + temperature : u.Quantity Radiative temperature in each shell dilution_factor : numpy.ndarray Dilution Factors in each shell @@ -44,18 +23,18 @@ class DiluteBlackBodyRadiationFieldState: def __init__( self, - t_radiative: u.Quantity, + temperature: u.Quantity, dilution_factor: np.ndarray, geometry=None, ): # ensuring that the radiation_field has both # dilution_factor and t_radiative equal length - assert len(t_radiative) == len(dilution_factor) + assert len(temperature) == len(dilution_factor) if ( geometry is not None ): # check the active shells only (this is used when setting up the radiation_field_state) assert np.all( - t_radiative[ + temperature[ geometry.v_inner_boundary_index : geometry.v_outer_boundary_index ] > 0 * u.K @@ -67,11 +46,15 @@ def __init__( > 0 ) else: - assert np.all(t_radiative > 0 * u.K) + assert np.all(temperature > 0 * u.K) assert np.all(dilution_factor > 0) - self.t_radiative = t_radiative + self.temperature = temperature self.dilution_factor = dilution_factor + @property + def temperature_kelvin(self): + return self.temperature.to(u.K).value + def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]): """ Calculate the intensity of the radiation field at a given frequency. @@ -87,5 +70,5 @@ def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]): Intensity of the radiation field at the given frequency """ return self.dilution_factor * intensity_black_body( - nu[np.newaxis].T, self.t_radiative + nu[np.newaxis].T, self.temperature ) diff --git a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py index 64b0064e3e7..64db0f8f6e0 100644 --- a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py +++ b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py @@ -7,7 +7,7 @@ import tardis.constants as const from tardis.io.atom_data import AtomData from tardis.radiation_field.planck_rad_field import ( - DiluteBlackBodyRadiationFieldState, + DilutePlanckianRadiationField, ) from tardis.transport.montecarlo.estimators.radfield_mc_estimators import ( RadiationFieldMCEstimators, @@ -76,7 +76,7 @@ def __init__(self, atom_data: AtomData) -> None: def solve( self, - dilute_blackbody_radiationfield_state: DiluteBlackBodyRadiationFieldState, + dilute_blackbody_radiationfield_state: DilutePlanckianRadiationField, t_electrons: u.Quantity, ): """ @@ -202,7 +202,7 @@ def calculate_stimulated_recomb_rate_factor( def calculate_mean_intensity_photo_ion_table( self, - dilute_blackbody_radiationfield_state: DiluteBlackBodyRadiationFieldState, + dilute_blackbody_radiationfield_state: DilutePlanckianRadiationField, ): mean_intensity = ( dilute_blackbody_radiationfield_state.calculate_mean_intensity( @@ -213,7 +213,7 @@ def calculate_mean_intensity_photo_ion_table( mean_intensity, index=self.atom_data.photoionization_data.index, columns=np.arange( - len(dilute_blackbody_radiationfield_state.t_radiative) + len(dilute_blackbody_radiationfield_state.temperature) ), ) return mean_intensity_df diff --git a/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py b/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py index b097356a6aa..b8a33125ecf 100644 --- a/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py +++ b/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py @@ -4,7 +4,7 @@ from tardis import constants as const from tardis.radiation_field.planck_rad_field import ( - DiluteBlackBodyRadiationFieldState, + DilutePlanckianRadiationField, ) DILUTION_FACTOR_ESTIMATOR_CONSTANT = ( @@ -52,6 +52,6 @@ def solve(self, radfield_mc_estimators, time_of_simulation, volume): * volume ) - return DiluteBlackBodyRadiationFieldState( + return DilutePlanckianRadiationField( temperature_radiative, dilution_factor ) diff --git a/tardis/transport/montecarlo/montecarlo_transport_state.py b/tardis/transport/montecarlo/montecarlo_transport_state.py index 30bceccc3a8..f4802101773 100644 --- a/tardis/transport/montecarlo/montecarlo_transport_state.py +++ b/tardis/transport/montecarlo/montecarlo_transport_state.py @@ -106,7 +106,7 @@ def calculate_radiationfield_properties(self): ) return ( - dilute_bb_radfield.t_radiative, + dilute_bb_radfield.temperature, dilute_bb_radfield.dilution_factor, ) From 41779e2702834b46284858c729fccc0e2da4507d Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sun, 9 Jun 2024 09:51:04 -0400 Subject: [PATCH 22/89] switched over to old tau_sobolev calculation --- tardis/opacities/tau_sobolev.py | 34 ++++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/tardis/opacities/tau_sobolev.py b/tardis/opacities/tau_sobolev.py index 302e00a6f64..7431c014b2b 100644 --- a/tardis/opacities/tau_sobolev.py +++ b/tardis/opacities/tau_sobolev.py @@ -77,12 +77,36 @@ def calculate( self, lines, level_number_density, + lines_lower_level_index, time_explosion, stimulated_emission_factor, + f_lu, + wavelength_cm, ): - return calculate_sobolev_line_opacity( - lines, - level_number_density, - time_explosion, - stimulated_emission_factor, + f_lu = f_lu.values[np.newaxis].T + wavelength = wavelength_cm.values[np.newaxis].T + n_lower = level_number_density.values.take( + lines_lower_level_index, axis=0, mode="raise" + ) + tau_sobolevs = ( + self.sobolev_coefficient + * f_lu + * wavelength + * time_explosion + * n_lower + * stimulated_emission_factor + ) + + if np.any(np.isnan(tau_sobolevs)) or np.any( + np.isinf(np.abs(tau_sobolevs)) + ): + raise ValueError( + "Some tau_sobolevs are nan, inf, -inf in tau_sobolevs." + " Something went wrong!" + ) + + return pd.DataFrame( + tau_sobolevs, + index=lines.index, + columns=np.array(level_number_density.columns), ) From 0d1bc4db43303ae2d25199c37359bd24df1aad43 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sun, 9 Jun 2024 16:49:21 -0400 Subject: [PATCH 23/89] renamed function to indicate numba use --- tardis/opacities/macro_atom/base.py | 2 +- tardis/opacities/macro_atom/util.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tardis/opacities/macro_atom/base.py b/tardis/opacities/macro_atom/base.py index 5bfc7e1bb52..32b71187a09 100644 --- a/tardis/opacities/macro_atom/base.py +++ b/tardis/opacities/macro_atom/base.py @@ -67,7 +67,7 @@ def _calculate_transition_probability( transition_type = macro_atom_data.transition_type.values lines_idx = macro_atom_data.lines_idx.values tpos = macro_atom_data.transition_probability.values - util.calculate_transition_probabilities( + util.fast_calculate_transition_probabilities( tpos, beta_sobolev.values, j_blues.values, diff --git a/tardis/opacities/macro_atom/util.py b/tardis/opacities/macro_atom/util.py index c72781f4b41..8cb5eed473a 100644 --- a/tardis/opacities/macro_atom/util.py +++ b/tardis/opacities/macro_atom/util.py @@ -4,14 +4,14 @@ from tardis import constants as const from tardis.transport.montecarlo import njit_dict -h_cgs = const.h.cgs.value -c = const.c.to("cm/s").value -kb = const.k_B.cgs.value -inv_c2 = 1 / (c**2) +H_CGS = const.h.cgs.value +C = const.c.to("cm/s").value +K_BOLTZMANN = const.k_B.cgs.value +INV_C2 = 1 / (C**2) @njit(**njit_dict) -def calculate_transition_probabilities( +def fast_calculate_transition_probabilities( transition_probability_coef, beta_sobolev, j_blues, From c92474988308a4920c85c3e1261debf347ce2023 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 10 Jun 2024 19:40:44 -0400 Subject: [PATCH 24/89] address comments --- tardis/opacities/macro_atom/base.py | 5 --- tardis/opacities/macro_atom/util.py | 5 --- tardis/opacities/tau_sobolev.py | 58 +++++++++++++++++++++++++++++ 3 files changed, 58 insertions(+), 10 deletions(-) diff --git a/tardis/opacities/macro_atom/base.py b/tardis/opacities/macro_atom/base.py index 32b71187a09..d23e5473198 100644 --- a/tardis/opacities/macro_atom/base.py +++ b/tardis/opacities/macro_atom/base.py @@ -126,11 +126,6 @@ def prepare_transition_probabilities( ) return transition_probabilities - def _normalize_transition_probabilities(self, transition_probabilities): - util.normalize_transition_probabilities( - transition_probabilities, self.block_references - ) - @staticmethod def _get_macro_atom_data(atomic_data): try: diff --git a/tardis/opacities/macro_atom/util.py b/tardis/opacities/macro_atom/util.py index 8cb5eed473a..97eb98f1a64 100644 --- a/tardis/opacities/macro_atom/util.py +++ b/tardis/opacities/macro_atom/util.py @@ -4,11 +4,6 @@ from tardis import constants as const from tardis.transport.montecarlo import njit_dict -H_CGS = const.h.cgs.value -C = const.c.to("cm/s").value -K_BOLTZMANN = const.k_B.cgs.value -INV_C2 = 1 / (C**2) - @njit(**njit_dict) def fast_calculate_transition_probabilities( diff --git a/tardis/opacities/tau_sobolev.py b/tardis/opacities/tau_sobolev.py index 7431c014b2b..cf667f9f0ff 100644 --- a/tardis/opacities/tau_sobolev.py +++ b/tardis/opacities/tau_sobolev.py @@ -23,6 +23,34 @@ def calculate_sobolev_line_opacity( time_explosion, stimulated_emission_factor, ): + """ + Calculates the Sobolev line opacity based on the provided parameters. + + Parameters + ---------- + lines : pandas.DataFrame + DataFrame containing information about spectral lines. + level_number_density : pandas.DataFrame + DataFrame with level number densities. + time_explosion : astropy.units.Quantity + Time since explosion. + stimulated_emission_factor : float + Factor for stimulated emission. + + Returns + ------- + pandas.DataFrame + Calculated Sobolev line opacity values. + + Raises + ------ + ValueError + If any calculated tau_sobolevs are nan or inf. + + Examples + -------- + >>> calculate_sobolev_line_opacity(lines_data, level_density_data, time_exp, stim_factor) + """ tau_sobolevs = ( (lines.wavelength_cm * lines.f_lu).values[np.newaxis].T * SOBOLEV_COEFFICIENT @@ -83,6 +111,36 @@ def calculate( f_lu, wavelength_cm, ): + """ + Calculate Sobolev line opacity. + + Calculates the Sobolev line opacity based on the provided parameters. + + Parameters + ---------- + lines : pandas.DataFrame + DataFrame containing information about spectral lines. + level_number_density : pandas.DataFrame + DataFrame with level number densities. + time_explosion : astropy.units.Quantity + Time since explosion. + stimulated_emission_factor : float + Factor for stimulated emission. + + Returns + ------- + pandas.DataFrame + Calculated Sobolev line opacity values. + + Raises + ------ + ValueError + If any calculated tau_sobolevs are nan or inf. + + Examples + -------- + >>> calculate_sobolev_line_opacity(lines_data, level_density_data, time_exp, stim_factor) + """ f_lu = f_lu.values[np.newaxis].T wavelength = wavelength_cm.values[np.newaxis].T n_lower = level_number_density.values.take( From af7e39e6398451e78380abcb8c7a68de1cd920fa Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 10 Jun 2024 22:04:43 -0400 Subject: [PATCH 25/89] added dilute planckian radiation field --- tardis/model/base.py | 4 ++-- tardis/plasma/properties/plasma_input.py | 21 +++++++++++++------ .../plasma/properties/property_collections.py | 3 ++- tardis/plasma/standard_plasmas.py | 11 ++++++---- tardis/radiation_field/__init__.py | 3 +++ 5 files changed, 29 insertions(+), 13 deletions(-) diff --git a/tardis/model/base.py b/tardis/model/base.py index 41cf7ca3cc6..4bfd68ca6e1 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -141,14 +141,14 @@ def dilution_factor(self, new_dilution_factor): @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: diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index defa0f4543f..f9d70216dcc 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -1,4 +1,8 @@ -from tardis.plasma.properties.base import ArrayInput, Input, ObjectInput +from tardis.plasma.properties.base import ( + ArrayInput, + Input, + ProcessingPlasmaProperty, +) __all__ = [ "TRadiative", @@ -17,10 +21,11 @@ "ContinuumInteractionSpecies", "NLTEIonizationSpecies", "NLTEExcitationSpecies", + "DilutePlanckianRadField", ] -class TRadiative(ArrayInput): +class TRadiative(ProcessingPlasmaProperty): """ Attributes ---------- @@ -30,8 +35,11 @@ class TRadiative(ArrayInput): outputs = ("t_rad",) latex_name = (r"T_{\textrm{rad}}",) + def calculate(self, dilute_planckian_radiation_field): + return dilute_planckian_radiation_field.temperature.cgs.value -class DilutionFactor(ArrayInput): + +class DilutionFactor(ProcessingPlasmaProperty): """ Attributes ---------- @@ -42,6 +50,8 @@ class DilutionFactor(ArrayInput): outputs = ("w",) latex_name = ("W",) + def calculate(self, dilute_planckian_radiation_field): + return dilute_planckian_radiation_field.dilution_factor class AtomicData(Input): @@ -162,6 +172,5 @@ class NumberDensity(Input): latex_name = ("N_{i}",) -class DilutePlanckianRadFieldInput(ObjectInput): - input_object_map = {"t_rad": "temperature_kelvin", "w": "dilution_factor"} - outputs = ("t_rad", "w") +class DilutePlanckianRadField(Input): + outputs = ("dilute_planckian_radiation_field",) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index 3b1e8ebc5cd..dc72e8b217f 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -12,11 +12,12 @@ class PlasmaPropertyCollection(list): basic_inputs = PlasmaPropertyCollection( [ - TRadiative, + DilutePlanckianRadField, Abundance, NumberDensity, TimeExplosion, AtomicData, + TRadiative, DilutionFactor, LinkTRadTElectron, HeliumTreatment, diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 92322a38a3f..05391d7e399 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -47,6 +47,7 @@ two_photon_properties, ) from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper +from tardis.radiation_field import DilutePlanckianRadiationField from tardis.util.base import species_string_to_tuple logger = logging.getLogger(__name__) @@ -70,13 +71,13 @@ def assemble_plasma(config, simulation_state, atom_data=None): : plasma.BasePlasma """ - if (config.plasma.ionization == "nebular") or ( config.plasma.excitation == "dilute-lte" ): radiation_field = "dilute_planckian_radiation_field" else: - pass + radiation_field = "dilute_planckian_radiation_field" + # Convert the nlte species list to a proper format. nlte_species = [ species_string_to_tuple(s) for s in config.plasma.nlte.species @@ -121,13 +122,15 @@ def assemble_plasma(config, simulation_state, atom_data=None): for s in config.plasma.nlte_excitation_species ] + dilute_planckian_radiation_field = DilutePlanckianRadiationField( + simulation_state.t_radiative, simulation_state.dilution_factor + ) kwargs = dict( - t_rad=simulation_state.t_radiative, + dilute_planckian_radiation_field=dilute_planckian_radiation_field, abundance=simulation_state.abundance, number_density=simulation_state.elemental_number_density, atomic_data=atom_data, time_explosion=simulation_state.time_explosion, - w=simulation_state.dilution_factor, link_t_rad_t_electron=config.plasma.link_t_rad_t_electron, continuum_interaction_species=continuum_interaction_species, nlte_ionization_species=nlte_ionization_species, diff --git a/tardis/radiation_field/__init__.py b/tardis/radiation_field/__init__.py index e69de29bb2d..ba732d148bf 100644 --- a/tardis/radiation_field/__init__.py +++ b/tardis/radiation_field/__init__.py @@ -0,0 +1,3 @@ +from tardis.radiation_field.planck_rad_field import ( + DilutePlanckianRadiationField, +) From 36fa4b60ecb771688d67fbfdc8699897f588f18d Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 12 Jun 2024 21:38:46 -0400 Subject: [PATCH 26/89] refactor: Convert species lists to proper format in assemble_plasma function --- astropy_helpers | 1 + docs/physics/nlte/nlte_rates.ipynb | 960 ++++++++++++++++++ .../plasma/construction_simple_plasma.ipynb | 115 +++ tardis/plasma/nlte/__init__.py | 0 tardis/plasma/nlte/nlte_data.py | 123 +++ tardis/plasma/standard_plasmas.py | 22 +- test.txt | 2 + 7 files changed, 1209 insertions(+), 14 deletions(-) create mode 160000 astropy_helpers create mode 100644 docs/physics/nlte/nlte_rates.ipynb create mode 100644 docs/physics/plasma/construction_simple_plasma.ipynb create mode 100644 tardis/plasma/nlte/__init__.py create mode 100644 tardis/plasma/nlte/nlte_data.py create mode 100644 test.txt diff --git a/astropy_helpers b/astropy_helpers new file mode 160000 index 00000000000..9f82aac6c21 --- /dev/null +++ b/astropy_helpers @@ -0,0 +1 @@ +Subproject commit 9f82aac6c2141b425e2d639560f7260189d90b54 diff --git a/docs/physics/nlte/nlte_rates.ipynb b/docs/physics/nlte/nlte_rates.ipynb new file mode 100644 index 00000000000..3f0d08a38c7 --- /dev/null +++ b/docs/physics/nlte/nlte_rates.ipynb @@ -0,0 +1,960 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 25, + "metadata": {}, + "outputs": [], + "source": [ + "from tardis.plasma.nlte import nlte_data\n", + "from tardis.io.atom_data.util import download_atom_data\n", + "from tardis.io.atom_data import AtomData\n", + "\n", + "from scipy import sparse\n", + "import pandas as pd\n", + "import numpy as np" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Atomic Data kurucz_cd23_chianti_H_He already exists in /Users/wkerzend/projects/tardis/tardis-data/kurucz_cd23_chianti_H_He.h5. Will not download - override with force_download=True.\n" + ] + } + ], + "source": [ + "download_atom_data('kurucz_cd23_chianti_H_He')" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "NLTE_SPECIES = [(1,0)]" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "atom_data = AtomData.from_hdf('kurucz_cd23_chianti_H_He.h5')" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/var/folders/fg/nwmb1mss6kq3hwhj10dt0qh00000gn/T/ipykernel_25979/2147926970.py:1: PerformanceWarning: indexing past lexsort depth may impact performance.\n", + " test = atom_data.lines.loc[NLTE_SPECIES[0]]\n" + ] + } + ], + "source": [ + "test = atom_data.lines.loc[NLTE_SPECIES[0]]" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "325 µs ± 11.2 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n" + ] + } + ], + "source": [ + "%%timeit\n", + "\n", + "test.A_ul.unstack(fill_value=0)" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [], + "source": [ + "full_level_multi_index = pd.MultiIndex.from_product([atom_data.levels.loc[NLTE_SPECIES[0]].index, atom_data.levels.loc[NLTE_SPECIES[0]].index])" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
level_number0123456789...15161718192021222324
level_number
0NaN6.273665e+08NaN6.262443e+081.670556e+08NaN1.664219e+08NaNNaN6.805429e+07...NaN3.438596e+07NaN3.438596e+07NaNNaNNaNNaNNaNNaN
1NaNNaNNaNNaNNaN2.105076e+06NaN5.378844e+07NaNNaN...NaNNaN429939.240273NaN7.855694e+06NaNNaNNaNNaNNaN
2NaNNaNNaNNaN2.244369e+07NaN2.244401e+07NaNNaN9.661716e+06...NaN4.954019e+06NaN4.945181e+06NaNNaNNaNNaNNaNNaN
3NaNNaNNaNNaNNaN4.209949e+06NaN1.079587e+076.474973e+07NaN...NaNNaN859850.753800NaN1.574627e+069.435974e+06NaNNaNNaNNaN
4NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaN302288.604192NaN2.830168e+06NaNNaNNaNNaNNaN
5NaNNaNNaNNaNNaNNaNNaNNaNNaN3.062191e+06...NaN1.637221e+06NaN1.633173e+06NaNNaNNaNNaNNaNNaN
6NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaN602531.413900NaN5.650034e+053.394761e+06NaNNaNNaNNaN
7NaNNaNNaNNaNNaNNaNNaNNaNNaN3.479209e+05...NaN1.495171e+05NaN1.491122e+04NaNNaN4.240070e+06NaNNaNNaN
8NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...1.379376e+07NaNNaN1.349097e+05NaNNaN3.029557e+054.544335e+06NaNNaN
9NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaN215292.953168NaN1.238973e+06NaNNaNNaNNaNNaN
10NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaN7.372768e+05NaN7.372907e+05NaNNaNNaNNaNNaNNaN
11NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaN430569.564168NaN2.477853e+051.482659e+06NaNNaNNaNNaN
12NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaN1.884757e+05NaN1.888855e+04NaNNaN2.410167e+06NaNNaNNaN
13NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaN1.695886e+05NaNNaN1.726362e+052.584466e+06NaNNaN
14NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaN5.047039e+042.403367e+03NaNNaN3.828462e+06NaN
15NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaNNaN4.806703e+04NaNNaN2.736774e+054.386991e+06
16NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN
17NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN
18NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN
19NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN
20NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN
21NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN
22NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN
23NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN
24NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN...NaNNaNNaNNaNNaNNaNNaNNaNNaNNaN
\n", + "

25 rows × 25 columns

\n", + "
" + ], + "text/plain": [ + "level_number 0 1 2 3 4 5 \\\n", + "level_number \n", + "0 NaN 6.273665e+08 NaN 6.262443e+08 1.670556e+08 NaN \n", + "1 NaN NaN NaN NaN NaN 2.105076e+06 \n", + "2 NaN NaN NaN NaN 2.244369e+07 NaN \n", + "3 NaN NaN NaN NaN NaN 4.209949e+06 \n", + "4 NaN NaN NaN NaN NaN NaN \n", + "5 NaN NaN NaN NaN NaN NaN \n", + "6 NaN NaN NaN NaN NaN NaN \n", + "7 NaN NaN NaN NaN NaN NaN \n", + "8 NaN NaN NaN NaN NaN NaN \n", + "9 NaN NaN NaN NaN NaN NaN \n", + "10 NaN NaN NaN NaN NaN NaN \n", + "11 NaN NaN NaN NaN NaN NaN \n", + "12 NaN NaN NaN NaN NaN NaN \n", + "13 NaN NaN NaN NaN NaN NaN \n", + "14 NaN NaN NaN NaN NaN NaN \n", + "15 NaN NaN NaN NaN NaN NaN \n", + "16 NaN NaN NaN NaN NaN NaN \n", + "17 NaN NaN NaN NaN NaN NaN \n", + "18 NaN NaN NaN NaN NaN NaN \n", + "19 NaN NaN NaN NaN NaN NaN \n", + "20 NaN NaN NaN NaN NaN NaN \n", + "21 NaN NaN NaN NaN NaN NaN \n", + "22 NaN NaN NaN NaN NaN NaN \n", + "23 NaN NaN NaN NaN NaN NaN \n", + "24 NaN NaN NaN NaN NaN NaN \n", + "\n", + "level_number 6 7 8 9 ... \\\n", + "level_number ... \n", + "0 1.664219e+08 NaN NaN 6.805429e+07 ... \n", + "1 NaN 5.378844e+07 NaN NaN ... \n", + "2 2.244401e+07 NaN NaN 9.661716e+06 ... \n", + "3 NaN 1.079587e+07 6.474973e+07 NaN ... \n", + "4 NaN NaN NaN NaN ... \n", + "5 NaN NaN NaN 3.062191e+06 ... \n", + "6 NaN NaN NaN NaN ... \n", + "7 NaN NaN NaN 3.479209e+05 ... \n", + "8 NaN NaN NaN NaN ... \n", + "9 NaN NaN NaN NaN ... \n", + "10 NaN NaN NaN NaN ... \n", + "11 NaN NaN NaN NaN ... \n", + "12 NaN NaN NaN NaN ... \n", + "13 NaN NaN NaN NaN ... \n", + "14 NaN NaN NaN NaN ... \n", + "15 NaN NaN NaN NaN ... \n", + "16 NaN NaN NaN NaN ... \n", + "17 NaN NaN NaN NaN ... \n", + "18 NaN NaN NaN NaN ... \n", + "19 NaN NaN NaN NaN ... \n", + "20 NaN NaN NaN NaN ... \n", + "21 NaN NaN NaN NaN ... \n", + "22 NaN NaN NaN NaN ... \n", + "23 NaN NaN NaN NaN ... \n", + "24 NaN NaN NaN NaN ... \n", + "\n", + "level_number 15 16 17 18 \\\n", + "level_number \n", + "0 NaN 3.438596e+07 NaN 3.438596e+07 \n", + "1 NaN NaN 429939.240273 NaN \n", + "2 NaN 4.954019e+06 NaN 4.945181e+06 \n", + "3 NaN NaN 859850.753800 NaN \n", + "4 NaN NaN 302288.604192 NaN \n", + "5 NaN 1.637221e+06 NaN 1.633173e+06 \n", + "6 NaN NaN 602531.413900 NaN \n", + "7 NaN 1.495171e+05 NaN 1.491122e+04 \n", + "8 1.379376e+07 NaN NaN 1.349097e+05 \n", + "9 NaN NaN 215292.953168 NaN \n", + "10 NaN 7.372768e+05 NaN 7.372907e+05 \n", + "11 NaN NaN 430569.564168 NaN \n", + "12 NaN 1.884757e+05 NaN 1.888855e+04 \n", + "13 NaN NaN NaN 1.695886e+05 \n", + "14 NaN NaN NaN NaN \n", + "15 NaN NaN NaN NaN \n", + "16 NaN NaN NaN NaN \n", + "17 NaN NaN NaN NaN \n", + "18 NaN NaN NaN NaN \n", + "19 NaN NaN NaN NaN \n", + "20 NaN NaN NaN NaN \n", + "21 NaN NaN NaN NaN \n", + "22 NaN NaN NaN NaN \n", + "23 NaN NaN NaN NaN \n", + "24 NaN NaN NaN NaN \n", + "\n", + "level_number 19 20 21 22 \\\n", + "level_number \n", + "0 NaN NaN NaN NaN \n", + "1 7.855694e+06 NaN NaN NaN \n", + "2 NaN NaN NaN NaN \n", + "3 1.574627e+06 9.435974e+06 NaN NaN \n", + "4 2.830168e+06 NaN NaN NaN \n", + "5 NaN NaN NaN NaN \n", + "6 5.650034e+05 3.394761e+06 NaN NaN \n", + "7 NaN NaN 4.240070e+06 NaN \n", + "8 NaN NaN 3.029557e+05 4.544335e+06 \n", + "9 1.238973e+06 NaN NaN NaN \n", + "10 NaN NaN NaN NaN \n", + "11 2.477853e+05 1.482659e+06 NaN NaN \n", + "12 NaN NaN 2.410167e+06 NaN \n", + "13 NaN NaN 1.726362e+05 2.584466e+06 \n", + "14 5.047039e+04 2.403367e+03 NaN NaN \n", + "15 NaN 4.806703e+04 NaN NaN \n", + "16 NaN NaN NaN NaN \n", + "17 NaN NaN NaN NaN \n", + "18 NaN NaN NaN NaN \n", + "19 NaN NaN NaN NaN \n", + "20 NaN NaN NaN NaN \n", + "21 NaN NaN NaN NaN \n", + "22 NaN NaN NaN NaN \n", + "23 NaN NaN NaN NaN \n", + "24 NaN NaN NaN NaN \n", + "\n", + "level_number 23 24 \n", + "level_number \n", + "0 NaN NaN \n", + "1 NaN NaN \n", + "2 NaN NaN \n", + "3 NaN NaN \n", + "4 NaN NaN \n", + "5 NaN NaN \n", + "6 NaN NaN \n", + "7 NaN NaN \n", + "8 NaN NaN \n", + "9 NaN NaN \n", + "10 NaN NaN \n", + "11 NaN NaN \n", + "12 NaN NaN \n", + "13 NaN NaN \n", + "14 3.828462e+06 NaN \n", + "15 2.736774e+05 4.386991e+06 \n", + "16 NaN NaN \n", + "17 NaN NaN \n", + "18 NaN NaN \n", + "19 NaN NaN \n", + "20 NaN NaN \n", + "21 NaN NaN \n", + "22 NaN NaN \n", + "23 NaN NaN \n", + "24 NaN NaN \n", + "\n", + "[25 rows x 25 columns]" + ] + }, + "execution_count": 28, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "test.reindex(full_level_multi_index, fill_value=np.nan).A_ul.unstack(fill_value=0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "tardis-devel", + "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.1.0" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/physics/plasma/construction_simple_plasma.ipynb b/docs/physics/plasma/construction_simple_plasma.ipynb new file mode 100644 index 00000000000..44849be37d3 --- /dev/null +++ b/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 Date: Fri, 14 Jun 2024 12:47:44 -0400 Subject: [PATCH 27/89] moved radiation field into plasma. Resulting in some renames --- tardis/plasma/radiation_field/__init__.py | 3 + .../radiation_field/planck_rad_field.py | 74 +++++++++++++++++++ 2 files changed, 77 insertions(+) create mode 100644 tardis/plasma/radiation_field/__init__.py create mode 100644 tardis/plasma/radiation_field/planck_rad_field.py diff --git a/tardis/plasma/radiation_field/__init__.py b/tardis/plasma/radiation_field/__init__.py new file mode 100644 index 00000000000..307ea6046c5 --- /dev/null +++ b/tardis/plasma/radiation_field/__init__.py @@ -0,0 +1,3 @@ +from tardis.plasma.radiation_field.planck_rad_field import ( + DilutePlanckianRadiationField, +) diff --git a/tardis/plasma/radiation_field/planck_rad_field.py b/tardis/plasma/radiation_field/planck_rad_field.py new file mode 100644 index 00000000000..be58a0de665 --- /dev/null +++ b/tardis/plasma/radiation_field/planck_rad_field.py @@ -0,0 +1,74 @@ +from typing import Union + +import numpy as np +from astropy import units as u + +from tardis.util.base import intensity_black_body + + +class DilutePlanckianRadiationField: + """ + Represents the state of a dilute thermal radiation field. + + + Parameters + ---------- + temperature : u.Quantity + Radiative temperature in each shell + dilution_factor : numpy.ndarray + Dilution Factors in each shell + geometry: tardis.model.Radial1DModel + The geometry of the model that uses to constrains the active shells + """ + + def __init__( + self, + temperature: u.Quantity, + dilution_factor: np.ndarray, + geometry=None, + ): + # ensuring that the radiation_field has both + # dilution_factor and t_radiative equal length + assert len(temperature) == len(dilution_factor) + if ( + geometry is not None + ): # check the active shells only (this is used when setting up the radiation_field_state) + assert np.all( + temperature[ + geometry.v_inner_boundary_index : geometry.v_outer_boundary_index + ] + > 0 * u.K + ) + assert np.all( + dilution_factor[ + geometry.v_inner_boundary_index : geometry.v_outer_boundary_index + ] + > 0 + ) + else: + assert np.all(temperature > 0 * u.K) + assert np.all(dilution_factor > 0) + self.temperature = temperature + self.dilution_factor = dilution_factor + + @property + def temperature_kelvin(self): + return self.temperature.to(u.K).value + + def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]): + """ + Calculate the intensity of the radiation field at a given frequency. + + Parameters + ---------- + nu : u.Quantity + Frequency at which the intensity is to be calculated + + Returns + ------- + intensity : u.Quantity + Intensity of the radiation field at the given frequency + """ + return self.dilution_factor * intensity_black_body( + nu[np.newaxis].T, self.temperature + ) From 2f042f25adcf9b24a4763adf8f2cc051337a6076 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 19 Jun 2024 16:18:26 -0400 Subject: [PATCH 28/89] some fixes --- tardis/model/parse_input.py | 2 +- .../radiation_field/planck_rad_field.py | 2 +- tardis/plasma/standard_plasmas.py | 2 +- tardis/radiation_field/__init__.py | 3 - tardis/radiation_field/planck_rad_field.py | 74 ------------------- tardis/simulation/base.py | 9 ++- .../continuum_radfield_properties.py | 2 +- .../estimators/dilute_blackbody_properties.py | 2 +- 8 files changed, 11 insertions(+), 85 deletions(-) delete mode 100644 tardis/radiation_field/__init__.py delete mode 100644 tardis/radiation_field/planck_rad_field.py diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py index 8d01e022bf8..8da8080ec11 100644 --- a/tardis/model/parse_input.py +++ b/tardis/model/parse_input.py @@ -17,7 +17,7 @@ from tardis.model.geometry.radial1d import HomologousRadial1DGeometry from tardis.model.matter.composition import Composition from tardis.model.matter.decay import IsotopicMassFraction -from tardis.radiation_field.planck_rad_field import ( +from tardis.plasma.radiation_field.planck_rad_field import ( DilutePlanckianRadiationField, ) from tardis.transport.montecarlo.packet_source import ( diff --git a/tardis/plasma/radiation_field/planck_rad_field.py b/tardis/plasma/radiation_field/planck_rad_field.py index be58a0de665..bcd443081e2 100644 --- a/tardis/plasma/radiation_field/planck_rad_field.py +++ b/tardis/plasma/radiation_field/planck_rad_field.py @@ -47,7 +47,7 @@ def __init__( ) else: assert np.all(temperature > 0 * u.K) - assert np.all(dilution_factor > 0) + assert np.all(dilution_factor >= 0) self.temperature = temperature self.dilution_factor = dilution_factor diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 072093344b4..a048a59c254 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -47,7 +47,7 @@ two_photon_properties, ) from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper -from tardis.radiation_field import DilutePlanckianRadiationField +from tardis.plasma.radiation_field import DilutePlanckianRadiationField from tardis.util.base import species_string_to_tuple logger = logging.getLogger(__name__) diff --git a/tardis/radiation_field/__init__.py b/tardis/radiation_field/__init__.py deleted file mode 100644 index ba732d148bf..00000000000 --- a/tardis/radiation_field/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from tardis.radiation_field.planck_rad_field import ( - DilutePlanckianRadiationField, -) diff --git a/tardis/radiation_field/planck_rad_field.py b/tardis/radiation_field/planck_rad_field.py deleted file mode 100644 index be58a0de665..00000000000 --- a/tardis/radiation_field/planck_rad_field.py +++ /dev/null @@ -1,74 +0,0 @@ -from typing import Union - -import numpy as np -from astropy import units as u - -from tardis.util.base import intensity_black_body - - -class DilutePlanckianRadiationField: - """ - Represents the state of a dilute thermal radiation field. - - - Parameters - ---------- - temperature : u.Quantity - Radiative temperature in each shell - dilution_factor : numpy.ndarray - Dilution Factors in each shell - geometry: tardis.model.Radial1DModel - The geometry of the model that uses to constrains the active shells - """ - - def __init__( - self, - temperature: u.Quantity, - dilution_factor: np.ndarray, - geometry=None, - ): - # ensuring that the radiation_field has both - # dilution_factor and t_radiative equal length - assert len(temperature) == len(dilution_factor) - if ( - geometry is not None - ): # check the active shells only (this is used when setting up the radiation_field_state) - assert np.all( - temperature[ - geometry.v_inner_boundary_index : geometry.v_outer_boundary_index - ] - > 0 * u.K - ) - assert np.all( - dilution_factor[ - geometry.v_inner_boundary_index : geometry.v_outer_boundary_index - ] - > 0 - ) - else: - assert np.all(temperature > 0 * u.K) - assert np.all(dilution_factor > 0) - self.temperature = temperature - self.dilution_factor = dilution_factor - - @property - def temperature_kelvin(self): - return self.temperature.to(u.K).value - - def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]): - """ - Calculate the intensity of the radiation field at a given frequency. - - Parameters - ---------- - nu : u.Quantity - Frequency at which the intensity is to be calculated - - Returns - ------- - intensity : u.Quantity - Intensity of the radiation field at the given frequency - """ - return self.dilution_factor * intensity_black_body( - nu[np.newaxis].T, self.temperature - ) diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 219e20f363f..0470beb656d 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -17,6 +17,7 @@ from tardis.model.parse_input import initialize_packet_source from tardis.transport.montecarlo.base import MonteCarloTransportSolver from tardis.plasma.standard_plasmas import assemble_plasma +from tardis.plasma.radiation_field import DilutePlanckianRadiationField from tardis.simulation.convergence import ConvergenceSolver from tardis.util.base import is_notebook from tardis.visualization import ConvergencePlots @@ -339,10 +340,12 @@ def advance_state(self): # Bad test to see if this is a nlte run if "nlte_data" in self.plasma.outputs_dict: self.plasma.store_previous_properties() - + radiation_field = DilutePlanckianRadiationField( + temperature=self.simulation_state.t_radiative, + dilution_factor=self.simulation_state.dilution_factor, + ) update_properties = dict( - t_rad=self.simulation_state.t_radiative, - w=self.simulation_state.dilution_factor, + dilute_planckian_radiation_field=radiation_field ) # A check to see if the plasma is set with JBluesDetailed, in which # case it needs some extra kwargs. diff --git a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py index 64db0f8f6e0..da742cc1ee3 100644 --- a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py +++ b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py @@ -6,7 +6,7 @@ import tardis.constants as const from tardis.io.atom_data import AtomData -from tardis.radiation_field.planck_rad_field import ( +from tardis.plasma.radiation_field.planck_rad_field import ( DilutePlanckianRadiationField, ) from tardis.transport.montecarlo.estimators.radfield_mc_estimators import ( diff --git a/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py b/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py index b8a33125ecf..c82612584ea 100644 --- a/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py +++ b/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py @@ -3,7 +3,7 @@ from scipy.special import zeta from tardis import constants as const -from tardis.radiation_field.planck_rad_field import ( +from tardis.plasma.radiation_field.planck_rad_field import ( DilutePlanckianRadiationField, ) From d849f859a8108adac65ae7ef82f8e1f599f42590 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 19 Jun 2024 16:18:55 -0400 Subject: [PATCH 29/89] black montecarlo --- tardis/transport/montecarlo/base.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tardis/transport/montecarlo/base.py b/tardis/transport/montecarlo/base.py index 8d4b4599e5a..ce62695eb4d 100644 --- a/tardis/transport/montecarlo/base.py +++ b/tardis/transport/montecarlo/base.py @@ -187,9 +187,9 @@ def run( total_iterations=total_iterations, ) - transport_state._montecarlo_virtual_luminosity.value[:] = ( - v_packets_energy_hist - ) + transport_state._montecarlo_virtual_luminosity.value[ + : + ] = v_packets_energy_hist transport_state.last_interaction_type = last_interaction_tracker.types transport_state.last_interaction_in_nu = last_interaction_tracker.in_nus transport_state.last_line_interaction_in_id = ( From 17864364cb18351fa5e0f22ccab999c5fa94d635 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 19 Jun 2024 17:00:34 -0400 Subject: [PATCH 30/89] chore: Initialize atom data and simulation state in `initialization.py` --- tardis/simulation/base.py | 94 +++++++------------------ tardis/simulation/initialization.py | 105 ++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 70 deletions(-) create mode 100644 tardis/simulation/initialization.py diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 0470beb656d..4926d4ad984 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -1,7 +1,6 @@ import logging import time from collections import OrderedDict -from pathlib import Path import numpy as np import pandas as pd @@ -10,19 +9,19 @@ import tardis from tardis import constants as const -from tardis.io.atom_data.base import AtomData from tardis.io.configuration.config_reader import ConfigurationError from tardis.io.util import HDFWriterMixin -from tardis.model import SimulationState -from tardis.model.parse_input import initialize_packet_source -from tardis.transport.montecarlo.base import MonteCarloTransportSolver -from tardis.plasma.standard_plasmas import assemble_plasma from tardis.plasma.radiation_field import DilutePlanckianRadiationField +from tardis.plasma.standard_plasmas import assemble_plasma from tardis.simulation.convergence import ConvergenceSolver +from tardis.simulation.initialization import ( + initialize_atom_data, + initialize_simulation_state, +) +from tardis.transport.montecarlo.base import MonteCarloTransportSolver from tardis.util.base import is_notebook from tardis.visualization import ConvergencePlots -# Adding logging support logger = logging.getLogger(__name__) @@ -496,12 +495,12 @@ def log_plasma_state( ---------- t_rad : astropy.units.Quanity current t_rad - w : astropy.units.Quanity - current w + dilution_factor : np.ndarray + current dilution_factor next_t_rad : astropy.units.Quanity next t_rad - next_w : astropy.units.Quanity - next_w + next_dilution_factor : np.ndarray + next dilution_factor log_sampling : int the n-th shells to be plotted @@ -617,6 +616,9 @@ def from_config( show_convergence_plots=False, show_progress_bars=True, legacy_mode_enabled=False, + atom_data=None, + plasma=None, + transport=None, **kwargs, ): """ @@ -638,70 +640,22 @@ def from_config( # Allow overriding some config structures. This is useful in some # unit tests, and could be extended in all the from_config classmethods. - atom_data = kwargs.get("atom_data", None) - if atom_data is None: - if "atom_data" in config: - if Path(config.atom_data).is_absolute(): - atom_data_fname = Path(config.atom_data) - else: - atom_data_fname = ( - Path(config.config_dirname) / config.atom_data - ) - - else: - raise ValueError( - "No atom_data option found in the configuration." - ) - - logger.info(f"\n\tReading Atomic Data from {atom_data_fname}") - - try: - atom_data = AtomData.from_hdf(atom_data_fname) - except TypeError as e: - print( - e, - "Error might be from the use of an old-format of the atomic database, \n" - "please see https://github.com/tardis-sn/tardis-refdata/tree/master/atom_data" - " for the most recent version.", - ) - raise - if "model" in kwargs: - simulation_state = kwargs["model"] - else: - if hasattr(config, "csvy_model"): - simulation_state = SimulationState.from_csvy( - config, - atom_data=atom_data, - legacy_mode_enabled=legacy_mode_enabled, - ) - else: - simulation_state = SimulationState.from_config( - config, - atom_data=atom_data, - legacy_mode_enabled=legacy_mode_enabled, - ) - if packet_source is not None: - simulation_state.packet_source = initialize_packet_source( - config, - simulation_state.geometry, - packet_source, - legacy_mode_enabled, - ) - if "plasma" in kwargs: - plasma = kwargs["plasma"] - else: + atom_data = initialize_atom_data(config, atom_data=atom_data) + simulation_state = initialize_simulation_state( + config, packet_source, legacy_mode_enabled, kwargs, atom_data + ) + if plasma is None: plasma = assemble_plasma( config, simulation_state, atom_data=atom_data, ) - if "transport" in kwargs: - if packet_source is not None: - raise ConfigurationError( - "Cannot specify packet_source and transport at the same time." - ) - transport = kwargs["transport"] - else: + + if (transport is not None) and (packet_source is not None): + raise ConfigurationError( + "Cannot specify packet_source and transport at the same time." + ) + if transport is None: transport = MonteCarloTransportSolver.from_config( config, packet_source=simulation_state.packet_source, diff --git a/tardis/simulation/initialization.py b/tardis/simulation/initialization.py new file mode 100644 index 00000000000..e95ab558e69 --- /dev/null +++ b/tardis/simulation/initialization.py @@ -0,0 +1,105 @@ +import logging +from pathlib import Path + +from tardis.io.atom_data.base import AtomData +from tardis.model import SimulationState +from tardis.model.parse_input import initialize_packet_source + +logger = logging.getLogger(__name__) + + +def initialize_atom_data(config, atom_data=None): + """ + Initialize atom data for the simulation. + + Parameters + ---------- + config : object + The configuration object containing information about the atom data. + atom_data : object, optional + Existing atom data to be used, if provided. + + Returns + ------- + object + The initialized atom data. + + Raises + ------ + ValueError + If no atom_data option is found in the configuration. + """ + if atom_data is None: + if "atom_data" in config: + if Path(config.atom_data).is_absolute(): + atom_data_fname = Path(config.atom_data) + else: + atom_data_fname = Path(config.config_dirname) / config.atom_data + + else: + raise ValueError("No atom_data option found in the configuration.") + + logger.info(f"\n\tReading Atomic Data from {atom_data_fname}") + + try: + atom_data = AtomData.from_hdf(atom_data_fname) + except TypeError as e: + print( + e, + "Error might be from the use of an old-format of the atomic database, \n" + "please see https://github.com/tardis-sn/tardis-refdata/tree/master/atom_data" + " for the most recent version.", + ) + raise + + return atom_data + + +def initialize_simulation_state( + config, packet_source, legacy_mode_enabled, kwargs, atom_data +): + """ + Initialize the simulation state. + + Parameters + ---------- + config : object + The configuration object for the simulation. + packet_source : object + The packet source for the simulation. + legacy_mode_enabled : bool + Flag indicating if legacy mode is enabled. + kwargs : dict + Additional keyword arguments. + atom_data : object + The atom data for the simulation. + + Returns + ------- + object + The initialized simulation state. + """ + if "model" in kwargs: + simulation_state = kwargs["model"] + else: + if hasattr(config, "csvy_model"): + simulation_state = SimulationState.from_csvy( + config, + atom_data=atom_data, + legacy_mode_enabled=legacy_mode_enabled, + ) + else: + simulation_state = SimulationState.from_config( + config, + atom_data=atom_data, + legacy_mode_enabled=legacy_mode_enabled, + ) + if packet_source is not None: + simulation_state.packet_source = initialize_packet_source( + config, + simulation_state.geometry, + packet_source, + legacy_mode_enabled, + ) + + return simulation_state From 5740708e4c602ed9655143eff0268fe533a8ea15 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 19 Jun 2024 17:16:06 -0400 Subject: [PATCH 31/89] updating the documentation --- tardis/simulation/base.py | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 4926d4ad984..cb814fcf9d1 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -622,20 +622,35 @@ def from_config( **kwargs, ): """ - Create a new Simulation instance from a Configuration object. + Create a simulation instance from the provided configuration. Parameters ---------- - config : tardis.io.config_reader.Configuration - + config : object + The configuration object for the simulation. + packet_source : object, optional + The packet source for the simulation. + virtual_packet_logging : bool, optional + Flag indicating virtual packet logging. + show_convergence_plots : bool, optional + Flag indicating whether to show convergence plots. + show_progress_bars : bool, optional + Flag indicating whether to show progress bars. + legacy_mode_enabled : bool, optional + Flag indicating if legacy mode is enabled. + atom_data : object, optional + The atom data for the simulation. + plasma : object, optional + The plasma object for the simulation. + transport : object, optional + The transport solver for the simulation. **kwargs - Allow overriding some structures, such as model, plasma, atomic data - and the transport, instead of creating them from the configuration - object. + Additional keyword arguments. Returns ------- - Simulation + object + The created simulation instance. """ # Allow overriding some config structures. This is useful in some # unit tests, and could be extended in all the from_config classmethods. From 92f4e90cacc92e34c5f91c445ff7b1badec773f7 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 19 Jun 2024 20:18:18 -0400 Subject: [PATCH 32/89] feat: Add EstimatedRadiationFieldProperties class for Monte Carlo estimators --- tardis/transport/montecarlo/estimators/base.py | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 tardis/transport/montecarlo/estimators/base.py diff --git a/tardis/transport/montecarlo/estimators/base.py b/tardis/transport/montecarlo/estimators/base.py new file mode 100644 index 00000000000..641fa045028 --- /dev/null +++ b/tardis/transport/montecarlo/estimators/base.py @@ -0,0 +1,11 @@ +from dataclasses import dataclass + +import numpy as np + +from tardis.plasma.radiation_field import DilutePlanckianRadiationField + + +@dataclass +class EstimatedRadiationFieldProperties: + dilute_blackbody_radiationfield_state: DilutePlanckianRadiationField + j_blues: np.ndarray From e3f1819ce89f77db2448ef87ca6838d9471b44dd Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 21 Jun 2024 18:21:43 -0400 Subject: [PATCH 33/89] trying to get rid of j_blues in plasma with MC restructure --- tardis/plasma/properties/general.py | 18 ------ tardis/plasma/properties/plasma_input.py | 11 ---- tardis/simulation/base.py | 37 ++++++++--- tardis/transport/montecarlo/base.py | 14 +++- .../montecarlo/estimators/__init__.py | 3 + .../estimators/dilute_blackbody_properties.py | 64 +++++++++++++++++-- .../montecarlo/montecarlo_transport_state.py | 33 +--------- 7 files changed, 104 insertions(+), 76 deletions(-) diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index cfc56fba66c..b63acb66559 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -122,21 +122,3 @@ def __init__(self, plasma_parent): def calculate(self, t_electrons): return 1 / (self.k_B_cgs * t_electrons) - - -class LuminosityInner(ProcessingPlasmaProperty): - outputs = ("luminosity_inner",) - - @staticmethod - def calculate(r_inner, t_inner): - return ( - 4 * np.pi * const.sigma_sb.cgs * r_inner[0] ** 2 * t_inner**4 - ).to("erg/s") - - -class TimeSimulation(ProcessingPlasmaProperty): - outputs = ("time_simulation",) - - @staticmethod - def calculate(luminosity_inner): - return 1.0 * u.erg / luminosity_inner diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index f9d70216dcc..31f37f1ad48 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -15,8 +15,6 @@ "JBlueEstimator", "LinkTRadTElectron", "HeliumTreatment", - "RInner", - "TInner", "Volume", "ContinuumInteractionSpecies", "NLTEIonizationSpecies", @@ -125,15 +123,6 @@ class LinkTRadTElectron(Input): class HeliumTreatment(Input): outputs = ("helium_treatment",) - -class RInner(Input): - outputs = ("r_inner",) - - -class TInner(Input): - outputs = ("t_inner",) - - class Volume(Input): outputs = ("volume",) diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index cb814fcf9d1..9f56307787c 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -262,10 +262,23 @@ def advance_state(self): ------- converged : bool """ - ( - estimated_t_rad, - estimated_dilution_factor, - ) = self.transport.transport_state.calculate_radiationfield_properties() + estimated_radfield_properties = ( + self.transport.radfield_prop_solver.solve( + self.transport.transport_state.radfield_mc_estimators, + self.transport.transport_state.time_explosion, + self.transport.transport_state.time_of_simulation, + self.transport.transport_state.geometry_state.volume, + self.transport.transport_state.opacity_state.line_list_nu, + ) + ) + + estimated_t_rad = ( + estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature + ) + estimated_dilution_factor = ( + estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor + ) + estimated_t_inner = self.estimate_t_inner( self.simulation_state.t_inner, self.luminosity_requested, @@ -334,11 +347,6 @@ def advance_state(self): self.simulation_state.dilution_factor = next_dilution_factor self.simulation_state.blackbody_packet_source.temperature = next_t_inner - # model.calculate_j_blues() equivalent - # model.update_plasmas() equivalent - # Bad test to see if this is a nlte run - if "nlte_data" in self.plasma.outputs_dict: - self.plasma.store_previous_properties() radiation_field = DilutePlanckianRadiationField( temperature=self.simulation_state.t_radiative, dilution_factor=self.simulation_state.dilution_factor, @@ -346,6 +354,16 @@ def advance_state(self): update_properties = dict( dilute_planckian_radiation_field=radiation_field ) + + # model.calculate_j_blues() equivalent + # model.update_plasmas() equivalent + # Bad test to see if this is a nlte run + if "nlte_data" in self.plasma.outputs_dict: + self.plasma.store_previous_properties() + update_properties = dict( + j_blues=estimated_radfield_properties.j_blues + ) + # A check to see if the plasma is set with JBluesDetailed, in which # case it needs some extra kwargs. @@ -382,7 +400,6 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0): self.transport.run( transport_state, - time_explosion=self.simulation_state.time_explosion, iteration=self.iterations_executed, total_iterations=self.iterations, show_progress_bars=self.show_progress_bars, diff --git a/tardis/transport/montecarlo/base.py b/tardis/transport/montecarlo/base.py index ce62695eb4d..9f206852f72 100644 --- a/tardis/transport/montecarlo/base.py +++ b/tardis/transport/montecarlo/base.py @@ -10,6 +10,9 @@ montecarlo_main_loop, numba_config, ) +from tardis.transport.montecarlo.estimators.dilute_blackbody_properties import ( + MCRadiationFieldPropertiesSolver, +) from tardis.transport.montecarlo.estimators.radfield_mc_estimators import ( initialize_estimator_statistics, ) @@ -49,6 +52,7 @@ class MonteCarloTransportSolver(HDFWriterMixin): def __init__( self, + radfield_prop_solver, spectrum_frequency, virtual_spectrum_spawn_range, enable_full_relativity, @@ -64,6 +68,7 @@ def __init__( use_gpu=False, montecarlo_configuration=None, ): + self.radfield_prop_solver = radfield_prop_solver # inject different packets self.spectrum_frequency = spectrum_frequency self.virtual_spectrum_spawn_range = virtual_spectrum_spawn_range @@ -124,6 +129,7 @@ def initialize_transport_state( spectrum_frequency=self.spectrum_frequency, geometry_state=geometry_state, opacity_state=opacity_state, + time_explosion=simulation_state.time_explosion, ) transport_state.enable_full_relativity = ( @@ -142,7 +148,6 @@ def initialize_transport_state( def run( self, transport_state, - time_explosion, iteration=0, total_iterations=0, show_progress_bars=True, @@ -176,7 +181,7 @@ def run( ) = montecarlo_main_loop( transport_state.packet_collection, transport_state.geometry_state, - time_explosion.cgs.value, + transport_state.time_explosion.cgs.value, transport_state.opacity_state, self.montecarlo_configuration, transport_state.radfield_mc_estimators, @@ -289,7 +294,12 @@ def from_config( config.montecarlo.tracking.initial_array_length ) + radfield_prop_solver = MCRadiationFieldPropertiesSolver( + config.plasma.w_epsilon + ) + return cls( + radfield_prop_solver=radfield_prop_solver, spectrum_frequency=spectrum_frequency, virtual_spectrum_spawn_range=config.montecarlo.virtual_spectrum_spawn_range, enable_full_relativity=config.montecarlo.enable_full_relativity, diff --git a/tardis/transport/montecarlo/estimators/__init__.py b/tardis/transport/montecarlo/estimators/__init__.py index e69de29bb2d..f0b7f1abf73 100644 --- a/tardis/transport/montecarlo/estimators/__init__.py +++ b/tardis/transport/montecarlo/estimators/__init__.py @@ -0,0 +1,3 @@ +from tardis.transport.montecarlo.estimators.base import ( + EstimatedRadiationFieldProperties, +) diff --git a/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py b/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py index c82612584ea..8682ec7a156 100644 --- a/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py +++ b/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py @@ -6,6 +6,9 @@ from tardis.plasma.radiation_field.planck_rad_field import ( DilutePlanckianRadiationField, ) +from tardis.transport.montecarlo.estimators.base import ( + EstimatedRadiationFieldProperties, +) DILUTION_FACTOR_ESTIMATOR_CONSTANT = ( (const.c**2 / (2 * const.h)) @@ -18,11 +21,20 @@ ).cgs.value -class MCDiluteBlackBodyRadFieldSolver: - def __init__(self) -> None: - pass +class MCRadiationFieldPropertiesSolver: + w_epsilon = 1e-10 + + def __init__(self, w_epsilon=1e-10) -> None: + self.w_epsilon = w_epsilon - def solve(self, radfield_mc_estimators, time_of_simulation, volume): + def solve( + self, + radfield_mc_estimators, + time_explosion, + time_of_simulation, + volume, + line_list_nu, + ): """ Calculate an updated radiation field from the :math: `\\bar{nu}_\\textrm{estimator}` and :math:`\\J_\\textrm{estimator}` @@ -39,6 +51,26 @@ def solve(self, radfield_mc_estimators, time_of_simulation, volume): t_radiative : astropy.units.Quantity (float) dilution_factor : numpy.ndarray (float) """ + dilute_planck_rad_field = self.estimate_dilute_planck_radiation_field( + radfield_mc_estimators, time_of_simulation, volume + ) + j_blues = self.estimate_jblues( + radfield_mc_estimators.j_blue_estimator, + dilute_planck_rad_field, + time_explosion, + time_of_simulation, + volume, + line_list_nu, + ) + + return EstimatedRadiationFieldProperties( + dilute_blackbody_radiationfield_state=dilute_planck_rad_field, + j_blues=j_blues, + ) + + def estimate_dilute_planck_radiation_field( + self, radfield_mc_estimators, time_of_simulation, volume + ): temperature_radiative = ( T_RADIATIVE_ESTIMATOR_CONSTANT * radfield_mc_estimators.nu_bar_estimator @@ -51,7 +83,29 @@ def solve(self, radfield_mc_estimators, time_of_simulation, volume): * time_of_simulation.value * volume ) - return DilutePlanckianRadiationField( temperature_radiative, dilution_factor ) + + def estimate_jblues( + self, + j_blue_estimator, + estimated_radfield_state, + time_explosion, + time_of_simulation, + volume, + line_list_nu, + ): + j_blues_norm_factor = ( + const.c.cgs + * time_explosion + / (4 * np.pi * time_of_simulation * volume) + ) + j_blues = j_blue_estimator * j_blues_norm_factor.cgs.value + planck_j_blues = estimated_radfield_state.calculate_mean_intensity( + line_list_nu + ) + zero_j_blues = j_blues == 0.0 + j_blues[zero_j_blues] = self.w_epsilon * planck_j_blues[zero_j_blues] + + return j_blues diff --git a/tardis/transport/montecarlo/montecarlo_transport_state.py b/tardis/transport/montecarlo/montecarlo_transport_state.py index 14d678a5b9c..582127a56bb 100644 --- a/tardis/transport/montecarlo/montecarlo_transport_state.py +++ b/tardis/transport/montecarlo/montecarlo_transport_state.py @@ -5,7 +5,7 @@ from tardis.io.util import HDFWriterMixin from tardis.transport.montecarlo.estimators.dilute_blackbody_properties import ( - MCDiluteBlackBodyRadFieldSolver, + MCRadiationFieldPropertiesSolver, ) from tardis.transport.montecarlo.formal_integral import IntegrationError from tardis.spectrum import TARDISSpectrum @@ -62,9 +62,11 @@ def __init__( spectrum_frequency, geometry_state, opacity_state, + time_explosion, rpacket_tracker=None, vpacket_tracker=None, ): + self.time_explosion = time_explosion self.packet_collection = packet_collection self.radfield_mc_estimators = radfield_mc_estimators self.spectrum_frequency = spectrum_frequency @@ -81,35 +83,6 @@ def __init__( self.rpacket_tracker = rpacket_tracker self.vpacket_tracker = vpacket_tracker - def calculate_radiationfield_properties(self): - """ - Calculate an updated radiation field from the :math: - `\\bar{nu}_\\textrm{estimator}` and :math:`\\J_\\textrm{estimator}` - calculated in the montecarlo simulation. - The details of the calculation can be found in the documentation. - - Parameters - ---------- - nubar_estimator : np.ndarray (float) - j_estimator : np.ndarray (float) - - Returns - ------- - t_radiative : astropy.units.Quantity (float) - dilution_factor : numpy.ndarray (float) - """ - dilute_bb_solver = MCDiluteBlackBodyRadFieldSolver() - dilute_bb_radfield = dilute_bb_solver.solve( - self.radfield_mc_estimators, - self.time_of_simulation, - self.geometry_state.volume, - ) - - return ( - dilute_bb_radfield.temperature, - dilute_bb_radfield.dilution_factor, - ) - @property def output_nu(self): return self.packet_collection.output_nus * u.Hz From a99fbb87b50cb6a0b544162b4ccc06da021fcf44 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 22 Jun 2024 18:08:34 -0400 Subject: [PATCH 34/89] completely restructure j_blues. Estimators and all. --- tardis/plasma/base.py | 45 +++++---- tardis/plasma/properties/__init__.py | 1 - .../plasma/properties/continuum_processes.py | 37 +++----- tardis/plasma/properties/general.py | 2 - tardis/plasma/properties/j_blues.py | 94 ------------------- tardis/plasma/properties/plasma_input.py | 8 +- .../plasma/properties/property_collections.py | 15 +-- .../radiation_field/planck_rad_field.py | 58 ++++++++++++ tardis/plasma/standard_plasmas.py | 48 ++++++---- tardis/simulation/base.py | 46 +++++++-- 10 files changed, 172 insertions(+), 182 deletions(-) delete mode 100644 tardis/plasma/properties/j_blues.py diff --git a/tardis/plasma/base.py b/tardis/plasma/base.py index ef785486673..b5189f3798a 100644 --- a/tardis/plasma/base.py +++ b/tardis/plasma/base.py @@ -1,24 +1,36 @@ -import os -import re +import fileinput import logging +import re import tempfile -import fileinput +import dataclasses import networkx as nx -from tardis.plasma.exceptions import PlasmaMissingModule, NotInitializedModule -from tardis.plasma.properties.base import * from tardis.io.util import PlasmaWriterMixin +from tardis.plasma.exceptions import NotInitializedModule, PlasmaMissingModule +from tardis.plasma.properties.base import * logger = logging.getLogger(__name__) +@dataclasses.dataclass(frozen=True) +class PlasmaSolverSettings: + RADIATIVE_RATES_TYPE: str = "blackbody" + + class BasePlasma(PlasmaWriterMixin): outputs_dict = {} hdf_name = "plasma" - def __init__(self, plasma_properties, property_kwargs=None, **kwargs): + def __init__( + self, + plasma_properties, + plasma_solver_settings, + property_kwargs=None, + **kwargs, + ): + self.plasma_solver_settings = plasma_solver_settings self.outputs_dict = {} self.input_properties = [] self.plasma_properties = self._init_properties( @@ -63,7 +75,6 @@ def _build_graph(self): :param plasma_modules: :return: """ - self.graph = nx.DiGraph() # Adding all nodes self.graph.add_nodes_from( @@ -200,8 +211,7 @@ def freeze(self, *args): for key in args: if key not in self.outputs_dict: raise PlasmaMissingModule( - "Trying to freeze property {0}" - " that is unavailable".format(key) + f"Trying to freeze property {key}" " that is unavailable" ) self.outputs_dict[key].frozen = True @@ -224,8 +234,7 @@ def thaw(self, *args): for key in args: if key not in self.outputs_dict: raise PlasmaMissingModule( - "Trying to thaw property {0}" - " that is unavailable".format(key) + f"Trying to thaw property {key}" " that is unavailable" ) self.outputs_dict[key].frozen = False @@ -249,7 +258,6 @@ def _resolve_update_list(self, changed_properties): : list all affected modules. """ - descendants_ob = [] for plasma_property in changed_properties: @@ -284,11 +292,10 @@ def write_to_dot(self, fname, args=None, latex_label=True): enables/disables writing LaTeX equations and edge labels into the file. """ - try: - import pygraphviz + pass except: - logger.warn( + logger.warning( "pygraphviz missing. Plasma graph will not be " "generated." ) return @@ -303,7 +310,7 @@ def write_to_dot(self, fname, args=None, latex_label=True): ] = f"\\\\textrm{{{node}: }}" node_list = self.plasma_properties_dict[node] formulae = node_list.latex_formula - for output in range(0, len(formulae)): + for output in range(len(formulae)): formula = formulae[output] label = formula.replace("\\", "\\\\") print_graph.nodes[str(node)]["label"] += label @@ -341,7 +348,7 @@ def write_to_dot(self, fname, args=None, latex_label=True): ) if args is not None: - with open(fname, "r") as file: + with open(fname) as file: lines = file.readlines() for newline in args: @@ -374,7 +381,7 @@ def write_to_tex(self, fname_graph, scale=0.5, args=None, latex_label=True): try: import dot2tex except: - logger.warn( + logger.warning( "dot2tex missing. Plasma graph will not be " "generated." ) return @@ -383,7 +390,7 @@ def write_to_tex(self, fname_graph, scale=0.5, args=None, latex_label=True): self.write_to_dot(temp_fname, args=args, latex_label=latex_label) - with open(temp_fname, "r") as file: + with open(temp_fname) as file: dot_string = file.read().replace("\\\\", "\\") texcode = dot2tex.dot2tex( diff --git a/tardis/plasma/properties/__init__.py b/tardis/plasma/properties/__init__.py index 63d473f01f9..3363cbd9bcd 100644 --- a/tardis/plasma/properties/__init__.py +++ b/tardis/plasma/properties/__init__.py @@ -11,7 +11,6 @@ from tardis.plasma.properties.general import * from tardis.plasma.properties.helium_nlte import * from tardis.plasma.properties.ion_population import * -from tardis.plasma.properties.j_blues import * from tardis.plasma.properties.level_population import * from tardis.plasma.properties.nlte import * from tardis.plasma.properties.nlte_rate_equation_solver import * diff --git a/tardis/plasma/properties/continuum_processes.py b/tardis/plasma/properties/continuum_processes.py index 06121cc4e20..a95c7f2dea3 100644 --- a/tardis/plasma/properties/continuum_processes.py +++ b/tardis/plasma/properties/continuum_processes.py @@ -5,18 +5,17 @@ from numba import njit, prange from tardis import constants as const -from tardis.transport.montecarlo.estimators.util import ( - bound_free_estimator_array2frame, - integrate_array_by_blocks, -) -from tardis.transport.montecarlo import njit_dict from tardis.plasma.exceptions import PlasmaException from tardis.plasma.properties.base import ( Input, ProcessingPlasmaProperty, TransitionProbabilitiesProperty, ) -from tardis.plasma.properties.j_blues import JBluesDiluteBlackBody +from tardis.transport.montecarlo import njit_dict +from tardis.transport.montecarlo.estimators.util import ( + bound_free_estimator_array2frame, + integrate_array_by_blocks, +) __all__ = [ "SpontRecombRateCoeff", @@ -356,8 +355,7 @@ def calculate( photo_ion_norm_factor, photo_ion_block_references, photo_ion_index, - t_rad, - w, + dilute_planckian_radiation_field, level2continuum_idx, ): # Used for initialization @@ -366,8 +364,7 @@ def calculate( photo_ion_cross_sections, photo_ion_block_references, photo_ion_index, - t_rad, - w, + dilute_planckian_radiation_field, ) else: gamma_estimator = bound_free_estimator_array2frame( @@ -382,13 +379,12 @@ def calculate_from_dilute_bb( photo_ion_cross_sections, photo_ion_block_references, photo_ion_index, - t_rad, - w, + dilute_planckian_radiation_field, ): nu = photo_ion_cross_sections["nu"] x_sect = photo_ion_cross_sections["x_sect"] - j_nus = JBluesDiluteBlackBody.calculate( - photo_ion_cross_sections, nu, t_rad, w + j_nus = dilute_planckian_radiation_field.calculate_mean_intensity( + nu, ) gamma = j_nus.multiply(4.0 * np.pi * x_sect / nu / H, axis=0) gamma = integrate_array_by_blocks( @@ -416,8 +412,7 @@ def calculate( photo_ion_norm_factor, photo_ion_block_references, photo_ion_index, - t_rad, - w, + dilute_planckian_radiation_field, phi_ik, t_electrons, boltzmann_factor_photo_ion, @@ -429,8 +424,7 @@ def calculate( photo_ion_cross_sections, photo_ion_block_references, photo_ion_index, - t_rad, - w, + dilute_planckian_radiation_field, t_electrons, boltzmann_factor_photo_ion, ) @@ -447,16 +441,13 @@ def calculate_from_dilute_bb( photo_ion_cross_sections, photo_ion_block_references, photo_ion_index, - t_rad, - w, + dilute_planckian_radiation_field, t_electrons, boltzmann_factor_photo_ion, ): nu = photo_ion_cross_sections["nu"] x_sect = photo_ion_cross_sections["x_sect"] - j_nus = JBluesDiluteBlackBody.calculate( - photo_ion_cross_sections, nu, t_rad, w - ) + j_nus = dilute_planckian_radiation_field.calculate_mean_intensity(nu) j_nus *= boltzmann_factor_photo_ion alpha_stim = j_nus.multiply(4.0 * np.pi * x_sect / nu / H, axis=0) alpha_stim = integrate_array_by_blocks( diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index b63acb66559..cb6e46ef650 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -14,8 +14,6 @@ "SelectedAtoms", "ElectronTemperature", "BetaElectron", - "LuminosityInner", - "TimeSimulation", "ThermalGElectron", ] diff --git a/tardis/plasma/properties/j_blues.py b/tardis/plasma/properties/j_blues.py deleted file mode 100644 index 1f34f0a276b..00000000000 --- a/tardis/plasma/properties/j_blues.py +++ /dev/null @@ -1,94 +0,0 @@ -import numpy as np -import pandas as pd - -from tardis import constants as const -from tardis.plasma.properties.base import ( - DataFrameInput, - ProcessingPlasmaProperty, -) -from tardis.util.base import intensity_black_body - - -class JBluesBlackBody(ProcessingPlasmaProperty): - """ - Attributes - ---------- - lte_j_blues : Pandas DataFrame, dtype float - J_blue values as calculated in LTE. - """ - - outputs = ("j_blues",) - latex_name = "J^{b}_{lu(LTE)}" - - @staticmethod - def calculate(lines, nu, t_rad): - j_blues = intensity_black_body(nu.values[np.newaxis].T, t_rad) - j_blues = pd.DataFrame( - j_blues, index=lines.index, columns=np.arange(len(t_rad)) - ) - return j_blues - - -class JBluesDiluteBlackBody(ProcessingPlasmaProperty): - outputs = ("j_blues",) - latex_name = r"J_{\textrm{blue}}" - - @staticmethod - def calculate(lines, nu, t_rad, w): - j_blues = w * intensity_black_body(nu.values[np.newaxis].T, t_rad) - j_blues = pd.DataFrame( - j_blues, index=lines.index, columns=np.arange(len(t_rad)) - ) - return j_blues - - -class JBluesDetailed(ProcessingPlasmaProperty): - outputs = ("j_blues",) - latex_name = "J_{\\textrm{blue}}" - - def __init__(self, plasma_parent, w_epsilon): - super(JBluesDetailed, self).__init__(plasma_parent) - self.w_epsilon = w_epsilon - - def calculate( - self, lines, nu, t_rad, w, j_blues_norm_factor, j_blue_estimator - ): - # Used for initialization - if len(j_blue_estimator) == 0: - return JBluesDiluteBlackBody.calculate(lines, nu, t_rad, w) - else: - j_blues = pd.DataFrame( - j_blue_estimator * j_blues_norm_factor.value, - index=lines.index, - columns=np.arange(len(t_rad)), - ) - - for i in range(len(t_rad)): - zero_j_blues = j_blues[i] == 0.0 - j_blues[i][ - zero_j_blues - ] = self.w_epsilon * intensity_black_body( - nu[zero_j_blues].values, t_rad[i] - ) - return j_blues - - -class JBluesNormFactor(ProcessingPlasmaProperty): - - outputs = ("j_blues_norm_factor",) - latex = ( - r"\frac{c time_\textrm{simulation}}}{4\pi" - r"time_\textrm{simulation} volume}" - ) - - @staticmethod - def calculate(time_explosion, time_simulation, volume): - return ( - const.c.cgs - * time_explosion - / (4 * np.pi * time_simulation * volume) - ) - - -class JBluesEstimator(DataFrameInput): - outputs = ("j_blue_estimator",) diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index 31f37f1ad48..917c1ec4193 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -12,7 +12,7 @@ "NumberDensity", "IsotopeAbundance", "TimeExplosion", - "JBlueEstimator", + "JBlues", "LinkTRadTElectron", "HeliumTreatment", "Volume", @@ -96,15 +96,15 @@ class TimeExplosion(Input): latex_name = (r"t_{\textrm{exp}}",) -class JBlueEstimator(ArrayInput): +class JBlues(Input): """ Attributes ---------- j_blue_estimators : Numpy array """ - outputs = ("j_blue_estimators",) - latex_name = (r"J_{\textrm{blue-estimator}}",) + outputs = ("j_blues",) + latex_name = (r"J_{\textrm{blue}}",) class LinkTRadTElectron(Input): diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index dc72e8b217f..119423568d7 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -17,8 +17,7 @@ class PlasmaPropertyCollection(list): NumberDensity, TimeExplosion, AtomicData, - TRadiative, - DilutionFactor, + JBlues, LinkTRadTElectron, HeliumTreatment, ContinuumInteractionSpecies, @@ -28,6 +27,8 @@ class PlasmaPropertyCollection(list): ) basic_properties = PlasmaPropertyCollection( [ + TRadiative, + DilutionFactor, BetaRadiation, Levels, Lines, @@ -85,18 +86,10 @@ class PlasmaPropertyCollection(list): helium_numerical_nlte_properties = PlasmaPropertyCollection( [HeliumNumericalNLTE] ) -detailed_j_blues_inputs = PlasmaPropertyCollection( - [JBluesEstimator, RInner, TInner, Volume] -) -detailed_j_blues_properties = PlasmaPropertyCollection( - [JBluesDetailed, JBluesNormFactor, LuminosityInner, TimeSimulation] -) continuum_interaction_inputs = PlasmaPropertyCollection( [ StimRecombRateCoeffEstimator, PhotoIonRateCoeffEstimator, - RInner, - TInner, Volume, BfHeatingRateCoeffEstimator, StimRecombCoolingRateCoeffEstimator, @@ -114,9 +107,7 @@ class PlasmaPropertyCollection(list): ThermalGElectron, ThermalPhiSahaLTE, SahaFactor, - TimeSimulation, PhotoIonEstimatorsNormFactor, - LuminosityInner, StimRecombRateCoeff, CorrPhotoIonRateCoeff, SpontRecombCoolingRateCoeff, diff --git a/tardis/plasma/radiation_field/planck_rad_field.py b/tardis/plasma/radiation_field/planck_rad_field.py index bcd443081e2..a411062737c 100644 --- a/tardis/plasma/radiation_field/planck_rad_field.py +++ b/tardis/plasma/radiation_field/planck_rad_field.py @@ -72,3 +72,61 @@ def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]): return self.dilution_factor * intensity_black_body( nu[np.newaxis].T, self.temperature ) + + def to_planckian_radiation_field(self): + return PlanckianRadiationField(self.temperature) + + +class PlanckianRadiationField: + """ + Represents the state of a dilute thermal radiation field. + + + Parameters + ---------- + temperature : u.Quantity + Radiative temperature in each shell + dilution_factor : numpy.ndarray + Dilution Factors in each shell + geometry: tardis.model.Radial1DModel + The geometry of the model that uses to constrains the active shells + """ + + def __init__( + self, + temperature: u.Quantity, + geometry=None, + ): + if ( + geometry is not None + ): # check the active shells only (this is used when setting up the radiation_field_state) + assert np.all( + temperature[ + geometry.v_inner_boundary_index : geometry.v_outer_boundary_index + ] + > 0 * u.K + ) + else: + assert np.all(temperature > 0 * u.K) + assert np.all(dilution_factor >= 0) + self.temperature = temperature + + @property + def temperature_kelvin(self): + return self.temperature.to(u.K).value + + def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]): + """ + Calculate the intensity of the radiation field at a given frequency. + + Parameters + ---------- + nu : u.Quantity + Frequency at which the intensity is to be calculated + + Returns + ------- + intensity : u.Quantity + Intensity of the radiation field at the given frequency + """ + return intensity_black_body(nu[np.newaxis].T, self.temperature) diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index a048a59c254..caca28b4b70 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -5,14 +5,12 @@ from astropy import units as u from tardis.plasma import BasePlasma +from tardis.plasma.base import PlasmaSolverSettings from tardis.plasma.exceptions import PlasmaConfigError from tardis.plasma.properties import ( HeliumNumericalNLTE, IonNumberDensity, IonNumberDensityHeNLTE, - JBluesBlackBody, - JBluesDetailed, - JBluesDiluteBlackBody, LevelBoltzmannFactorNLTE, MarkovChainTransProbsCollector, RadiationFieldCorrection, @@ -30,8 +28,6 @@ basic_properties, continuum_interaction_inputs, continuum_interaction_properties, - detailed_j_blues_inputs, - detailed_j_blues_properties, dilute_lte_excitation_properties, helium_lte_properties, helium_nlte_properties, @@ -213,22 +209,37 @@ def assemble_plasma(config, simulation_state, atom_data=None): r_inner=simulation_state.r_inner.to(u.cm), t_inner=simulation_state.t_inner, ) - if config.plasma.radiative_rates_type == "blackbody": - plasma_modules.append(JBluesBlackBody) - elif config.plasma.radiative_rates_type == "dilute-blackbody": - plasma_modules.append(JBluesDiluteBlackBody) - elif config.plasma.radiative_rates_type == "detailed": - plasma_modules += detailed_j_blues_properties + detailed_j_blues_inputs - kwargs.update( - r_inner=simulation_state.r_inner.to(u.cm), - t_inner=simulation_state.t_inner, - volume=simulation_state.volume, - j_blue_estimator=None, + + ##### RADIATIVE RATES SETUP + + plasma_solver_settings = PlasmaSolverSettings( + RADIATIVE_RATES_TYPE=config.plasma.radiative_rates_type + ) + + if (plasma_solver_settings.RADIATIVE_RATES_TYPE == "dilute-blackbody") or ( + plasma_solver_settings.RADIATIVE_RATES_TYPE == "detailed" + ): + kwargs["j_blues"] = pd.DataFrame( + dilute_planckian_radiation_field.calculate_mean_intensity( + atom_data.lines["nu"].values + ), + index=atom_data.lines.index, + ) + + elif plasma_solver_settings.RADIATIVE_RATES_TYPE == "blackbody": + planckian_rad_field = ( + dilute_planckian_radiation_field.to_planckian_radiation_field() + ) + kwargs["j_blues"] = pd.DataFrame( + planckian_rad_field.calculate_mean_intensity( + atom_data.lines["nu"].values + ), + index=atom_data.lines.index, ) - property_kwargs[JBluesDetailed] = {"w_epsilon": config.plasma.w_epsilon} + else: raise ValueError( - f"radiative_rates_type type unknown - {config.plasma.radiative_rates_type}" + f"radiative_rates_type type unknown - {plasma_solver_settings.RADIATIVE_RATES_TYPE}" ) if config.plasma.excitation == "lte": @@ -335,6 +346,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): plasma = BasePlasma( plasma_properties=plasma_modules, property_kwargs=property_kwargs, + plasma_solver_settings=plasma_solver_settings, **kwargs, ) diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 9f56307787c..fff263f76aa 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -358,21 +358,49 @@ def advance_state(self): # model.calculate_j_blues() equivalent # model.update_plasmas() equivalent # Bad test to see if this is a nlte run - if "nlte_data" in self.plasma.outputs_dict: - self.plasma.store_previous_properties() - update_properties = dict( - j_blues=estimated_radfield_properties.j_blues + if ( + self.plasma.plasma_solver_settings.RADIATIVE_RATES_TYPE + == "blackbody" + ): + j_blues = radiation_field.calculate_mean_intensity( + self.plasma.atomic_data.lines.nu.values + ) + update_properties["j_blues"] = pd.DataFrame( + j_blues, index=self.plasma.atomic_data.lines.index + ) + raise NotImplementedError( + "This is not right yet - calculate dilute mean intensity" ) + elif ( + self.plasma.plasma_solver_settings.RADIATIVE_RATES_TYPE + == "dilute-blackbody" + ): + j_blues = radiation_field.calculate_mean_intensity( + self.plasma.atomic_data.lines.nu.values + ) + update_properties["j_blues"] = pd.DataFrame( + j_blues, index=self.plasma.atomic_data.lines.index + ) + elif ( + self.plasma.plasma_solver_settings.RADIATIVE_RATES_TYPE + == "detailed" + ): + update_properties["j_blues"] = pd.DataFrame( + estimated_radfield_properties.j_blues, + index=self.plasma.atomic_data.lines.index, + ) + else: + raise ValueError( + f"radiative_rates_type type unknown - {self.plasma.plasma_solver_settings.RADIATIVE_RATES_TYPE}" + ) + + self.plasma.store_previous_properties() # A check to see if the plasma is set with JBluesDetailed, in which # case it needs some extra kwargs. estimators = self.transport.transport_state.radfield_mc_estimators - if "j_blue_estimator" in self.plasma.outputs_dict: - update_properties.update( - t_inner=next_t_inner, - j_blue_estimator=estimators.j_blue_estimator, - ) + if "gamma_estimator" in self.plasma.outputs_dict: update_properties.update( gamma_estimator=estimators.photo_ion_estimator, From 2d7e1c55de96fce027af53737956c8af03ccc6fb Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 22 Jun 2024 18:36:26 -0400 Subject: [PATCH 35/89] cleanup for the restructure --- tardis/plasma/base.py | 2 +- tardis/plasma/radiation_field/planck_rad_field.py | 1 - tardis/simulation/base.py | 8 ++++---- 3 files changed, 5 insertions(+), 6 deletions(-) diff --git a/tardis/plasma/base.py b/tardis/plasma/base.py index b5189f3798a..8a42c32dfdb 100644 --- a/tardis/plasma/base.py +++ b/tardis/plasma/base.py @@ -1,8 +1,8 @@ +import dataclasses import fileinput import logging import re import tempfile -import dataclasses import networkx as nx diff --git a/tardis/plasma/radiation_field/planck_rad_field.py b/tardis/plasma/radiation_field/planck_rad_field.py index a411062737c..69d656bd0d6 100644 --- a/tardis/plasma/radiation_field/planck_rad_field.py +++ b/tardis/plasma/radiation_field/planck_rad_field.py @@ -108,7 +108,6 @@ def __init__( ) else: assert np.all(temperature > 0 * u.K) - assert np.all(dilution_factor >= 0) self.temperature = temperature @property diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index fff263f76aa..141879e2302 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -362,15 +362,15 @@ def advance_state(self): self.plasma.plasma_solver_settings.RADIATIVE_RATES_TYPE == "blackbody" ): - j_blues = radiation_field.calculate_mean_intensity( + planckian_radiation_field = ( + radiation_field.to_planckian_radiation_field() + ) + j_blues = planckian_radiation_field.calculate_mean_intensity( self.plasma.atomic_data.lines.nu.values ) update_properties["j_blues"] = pd.DataFrame( j_blues, index=self.plasma.atomic_data.lines.index ) - raise NotImplementedError( - "This is not right yet - calculate dilute mean intensity" - ) elif ( self.plasma.plasma_solver_settings.RADIATIVE_RATES_TYPE == "dilute-blackbody" From 818c6d438089b94fdc14d62fb6314351c4d20f29 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 23 Jul 2024 16:38:54 -0400 Subject: [PATCH 36/89] remove parse_input.py --- tardis/model/parse_input.py | 718 ------------------------------------ 1 file changed, 718 deletions(-) delete mode 100644 tardis/model/parse_input.py diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py deleted file mode 100644 index 8da8080ec11..00000000000 --- a/tardis/model/parse_input.py +++ /dev/null @@ -1,718 +0,0 @@ -import logging -import os - -import numpy as np -import pandas as pd -from astropy import units as u - -from tardis import constants as const -from tardis.io.model.parse_density_configuration import ( - calculate_density_after_time, - parse_config_v1_density, - parse_csvy_density, -) -from tardis.io.model.readers.base import read_abundances_file, read_density_file -from tardis.io.model.readers.csvy import parse_csv_abundances -from tardis.io.model.readers.generic_readers import read_uniform_abundances -from tardis.model.geometry.radial1d import HomologousRadial1DGeometry -from tardis.model.matter.composition import Composition -from tardis.model.matter.decay import IsotopicMassFraction -from tardis.plasma.radiation_field.planck_rad_field import ( - DilutePlanckianRadiationField, -) -from tardis.transport.montecarlo.packet_source import ( - BlackBodySimpleSource, - BlackBodySimpleSourceRelativistic, -) -from tardis.util.base import quantity_linspace - -logger = logging.getLogger(__name__) - - -def parse_structure_config(config, time_explosion, enable_homology=True): - """ - Parse the structure configuration data. - - Parameters - ---------- - config : object - The configuration data. - time_explosion : float - The time of the explosion. - enable_homology : bool, optional - Whether to enable homology (default is True). - - Returns - ------- - electron_densities : object - The parsed electron densities. - temperature : object - The parsed temperature. - geometry : object - The parsed geometry. - density : object - The parsed density. - - Raises - ------ - NotImplementedError - If the structure configuration type is not supported. - - Notes - ----- - This function parses the structure configuration data and returns the parsed electron - densities, temperature, geometry, and density. The structure configuration can be of - type 'specific' or 'file'. If it is of type 'specific', the velocity and density are - parsed from the configuration. If it is of type 'file', the velocity and density are - read from a file. The parsed data is used to create a homologous radial 1D geometry object. - """ - electron_densities = None - temperature = None - structure_config = config.model.structure - if structure_config.type == "specific": - velocity = quantity_linspace( - structure_config.velocity.start, - structure_config.velocity.stop, - structure_config.velocity.num + 1, - ).cgs - density = parse_config_v1_density(config) - - elif structure_config.type == "file": - if os.path.isabs(structure_config.filename): - structure_config_fname = structure_config.filename - else: - structure_config_fname = os.path.join( - config.config_dirname, structure_config.filename - ) - - ( - time_0, - velocity, - density_0, - electron_densities, - temperature, - ) = read_density_file(structure_config_fname, structure_config.filetype) - density_0 = density_0.insert(0, 0) - - density = calculate_density_after_time( - density_0, time_0, time_explosion - ) - - else: - raise NotImplementedError - - # Note: This is the number of shells *without* taking in mind the - # v boundaries. - if len(density) == len(velocity): - logger.warning( - "Number of density points larger than number of shells. Assuming inner point irrelevant" - ) - density = density[1:] - geometry = HomologousRadial1DGeometry( - velocity[:-1], # v_inner - velocity[1:], # v_outer - v_inner_boundary=structure_config.get("v_inner_boundary", None), - v_outer_boundary=structure_config.get("v_outer_boundary", None), - time_explosion=time_explosion, - ) - return electron_densities, temperature, geometry, density - - -def parse_csvy_geometry( - config, csvy_model_config, csvy_model_data, time_explosion -): - """ - Parse the geometry data from a CSVY model. - - Parameters - ---------- - config : object - The configuration data. - csvy_model_config : object - The configuration data of the CSVY model. - csvy_model_data : object - The data of the CSVY model. - time_explosion : float - The time of the explosion. - - Returns - ------- - geometry : object - The parsed geometry. - - Raises - ------ - None. - - Notes - ----- - This function parses the geometry data from a CSVY model. It extracts the velocity - information from the CSVY model configuration or data. The parsed velocity data is - used to create a homologous radial 1D geometry object, which is returned. - """ - if hasattr(config, "model"): - if hasattr(config.model, "v_inner_boundary"): - v_boundary_inner = config.model.v_inner_boundary - else: - v_boundary_inner = None - - if hasattr(config.model, "v_outer_boundary"): - v_boundary_outer = config.model.v_outer_boundary - else: - v_boundary_outer = None - else: - v_boundary_inner = None - v_boundary_outer = None - - if hasattr(csvy_model_config, "velocity"): - velocity = quantity_linspace( - csvy_model_config.velocity.start, - csvy_model_config.velocity.stop, - csvy_model_config.velocity.num + 1, - ).cgs - else: - velocity_field_index = [ - field["name"] for field in csvy_model_config.datatype.fields - ].index("velocity") - velocity_unit = u.Unit( - csvy_model_config.datatype.fields[velocity_field_index]["unit"] - ) - velocity = csvy_model_data["velocity"].values * velocity_unit - velocity = velocity.to("cm/s") - - geometry = HomologousRadial1DGeometry( - velocity[:-1], # v_inner - velocity[1:], # v_outer - v_inner_boundary=v_boundary_inner, - v_outer_boundary=v_boundary_outer, - time_explosion=time_explosion, - ) - return geometry - - -def parse_abundance_config(config, geometry, time_explosion): - """ - Parse the abundance configuration data. - - Parameters - ---------- - config : object - The configuration data. - geometry : object - The geometry of the model. - time_explosion : float - The time of the explosion. - - Returns - ------- - nuclide_mass_fraction : object - The parsed nuclide mass fraction. - - raw_isotope_abundance : object - The parsed raw isotope abundance. This is the isotope abundance data before decay. - - Raises - ------ - None. - - Notes - ----- - This function parses the abundance configuration data and returns the parsed nuclide - mass fraction. The abundance configuration can be of type 'uniform' or 'file'. If it - is of type 'uniform', the abundance and isotope abundance are read using the - 'read_uniform_abundances' function. If it is of type 'file', the abundance and - isotope abundance are read from a file using the 'read_abundances_file' function. - The parsed data is then processed to replace NaN values with 0.0, remove rows with - zero sum, and normalize the data if necessary. The resulting nuclide mass fraction - is returned. - """ - abundances_section = config.model.abundances - isotope_abundance = pd.DataFrame() - - if abundances_section.type == "uniform": - abundance, isotope_abundance = read_uniform_abundances( - abundances_section, geometry.no_of_shells - ) - - elif abundances_section.type == "file": - if os.path.isabs(abundances_section.filename): - abundances_fname = abundances_section.filename - else: - abundances_fname = os.path.join( - config.config_dirname, abundances_section.filename - ) - - index, abundance, isotope_abundance = read_abundances_file( - abundances_fname, abundances_section.filetype - ) - - abundance = abundance.replace(np.nan, 0.0) - abundance = abundance[abundance.sum(axis=1) > 0] - - norm_factor = abundance.sum(axis=0) + isotope_abundance.sum(axis=0) - - if np.any(np.abs(norm_factor - 1) > 1e-12): - logger.warning( - "Abundances have not been normalized to 1. - normalizing" - ) - abundance /= norm_factor - isotope_abundance /= norm_factor - # The next line is if the abundances are given via dict - # and not gone through the schema validator - raw_isotope_abundance = isotope_abundance - model_isotope_time_0 = config.model.abundances.get( - "model_isotope_time_0", 0.0 * u.day - ) - isotope_abundance = IsotopicMassFraction( - isotope_abundance, time_0=model_isotope_time_0 - ).decay(time_explosion) - - nuclide_mass_fraction = convert_to_nuclide_mass_fraction( - isotope_abundance, abundance - ) - return nuclide_mass_fraction, raw_isotope_abundance - - -def convert_to_nuclide_mass_fraction(isotopic_mass_fraction, mass_fraction): - """ - Convert the abundance and isotope abundance data to nuclide mass fraction. - - Parameters - ---------- - isotope_abundance : pandas.DataFrame - The isotope abundance data. - abundance : pandas.DataFrame - The abundance data. - - Returns - ------- - nuclide_mass_fraction : pandas.DataFrame - The converted nuclide mass fraction. - - Raises - ------ - None. - - Notes - ----- - This function converts the abundance and isotope abundance data to nuclide mass fraction. - If the abundance data is not None, it is converted to nuclide mass fraction by mapping - the abundance index to nuclide indices using the 'convert_element2nuclide_index' function. - The resulting abundance data is then concatenated with the isotope abundance data to - obtain the final nuclide mass fraction. - """ - nuclide_mass_fraction = pd.DataFrame() - if mass_fraction is not None: - mass_fraction.index = Composition.convert_element2nuclide_index( - mass_fraction.index - ) - nuclide_mass_fraction = mass_fraction - else: - nuclide_mass_fraction = pd.DataFrame() - - if isotopic_mass_fraction is not None: - nuclide_mass_fraction = pd.concat( - [nuclide_mass_fraction, isotopic_mass_fraction] - ) - return nuclide_mass_fraction - - -def parse_csvy_composition( - atom_data, csvy_model_config, csvy_model_data, time_explosion, geometry -): - """ - Parse the composition data from a CSVY model. - - Parameters - ---------- - atom_data : object - The atom data used for parsing. - csvy_model_config : object - The configuration data of the CSVY model. - csvy_model_data : object - The data of the CSVY model. - time_explosion : float - The time of the explosion. - geometry : object - The geometry of the model. - - Returns - ------- - density : object - The parsed density data. - abundance : object - The parsed abundance data. - isotope_abundance : object - The parsed isotope abundance data. - elemental_mass : object - The elemental mass data. - - Raises - ------ - None. - - Notes - ----- - This function parses the composition data from a CSVY model. It calls the 'parse_density_csvy' - function to parse the density data, and the 'parse_abundance_csvy' function to parse the abundance - and isotope abundance data. The parsed data is returned as density, abundance, isotope_abundance, - and elemental_mass. - """ - density = parse_density_csvy( - csvy_model_config, csvy_model_data, time_explosion - ) - - nuclide_mass_fraction, raw_isotope_mass_fraction = parse_abundance_csvy( - csvy_model_config, csvy_model_data, geometry, time_explosion - ) - return Composition( - density, - nuclide_mass_fraction, - raw_isotope_mass_fraction, - atom_data.atom_data.mass.copy(), - ) - - -def parse_abundance_csvy( - csvy_model_config, csvy_model_data, geometry, time_explosion -): - """ - Parse the abundance data from a CSVY model. - - Parameters - ---------- - csvy_model_config : object - The configuration data of the CSVY model. - csvy_model_data : object - The data of the CSVY model. - geometry : object - The geometry of the model. - - Returns - ------- - abundance : pd.DataFrame - The parsed abundance data. - isotope_abundance : pandas.DataFrame - The parsed isotope abundance data. - - Raises - ------ - None. - - Notes - ----- - This function parses the abundance data from a CSVY model. If the CSVY model - configuration contains an 'abundance' attribute, it uses the 'read_uniform_abundances' - function to parse the abundance and isotope abundance data. Otherwise, it uses the - 'parse_csv_abundances' function to parse the data. The parsed data is then processed - to replace NaN values with 0.0, remove rows with zero sum, and normalize the data - if necessary. The resulting abundance and isotope abundance arrays are returned. - """ - if hasattr(csvy_model_config, "abundance"): - abundances_section = csvy_model_config.abundance - mass_fraction, isotope_mass_fraction = read_uniform_abundances( - abundances_section, geometry.no_of_shells - ) - else: - _, mass_fraction, isotope_mass_fraction = parse_csv_abundances( - csvy_model_data - ) - mass_fraction = mass_fraction.loc[:, 1:] - mass_fraction.columns = np.arange(mass_fraction.shape[1]) - isotope_mass_fraction = isotope_mass_fraction.loc[:, 1:] - isotope_mass_fraction.columns = np.arange( - isotope_mass_fraction.shape[1] - ) - - mass_fraction = mass_fraction.replace(np.nan, 0.0) - mass_fraction = mass_fraction[mass_fraction.sum(axis=1) > 0] - isotope_mass_fraction = isotope_mass_fraction.replace(np.nan, 0.0) - isotope_mass_fraction = isotope_mass_fraction[ - isotope_mass_fraction.sum(axis=1) > 0 - ] - norm_factor = mass_fraction.sum(axis=0) + isotope_mass_fraction.sum(axis=0) - - if np.any(np.abs(norm_factor - 1) > 1e-12): - logger.warning( - "Abundances have not been normalized to 1. - normalizing" - ) - mass_fraction /= norm_factor - isotope_mass_fraction /= norm_factor - - raw_isotope_mass_fraction = isotope_mass_fraction - isotope_mass_fraction = IsotopicMassFraction( - isotope_mass_fraction, time_0=csvy_model_config.model_isotope_time_0 - ).decay(time_explosion) - return ( - convert_to_nuclide_mass_fraction(isotope_mass_fraction, mass_fraction), - raw_isotope_mass_fraction, - ) - - -def parse_density_csvy(csvy_model_config, csvy_model_data, time_explosion): - """ - Parse the density data from a CSVY model. - - Parameters - ---------- - csvy_model_config : object - The configuration data of the CSVY model. - csvy_model_data : object - The data of the CSVY model. - time_explosion : float - The time of the explosion. - - Returns - ------- - density : object - The parsed density data. - - Raises - ------ - None. - - Notes - ----- - This function parses the density data from a CSVY model. If the CSVY model configuration - contains a 'density' attribute, it uses the 'parse_csvy_density' function to parse the - density data. Otherwise, it calculates the density data using the 'calculate_density_after_time' - function. The parsed density data is returned. - """ - if hasattr(csvy_model_config, "density"): - density = parse_csvy_density(csvy_model_config, time_explosion) - else: - time_0 = csvy_model_config.model_density_time_0 - density_field_index = [ - field["name"] for field in csvy_model_config.datatype.fields - ].index("density") - density_unit = u.Unit( - csvy_model_config.datatype.fields[density_field_index]["unit"] - ) - density_0 = csvy_model_data["density"].values * density_unit - # Removing as thee new architecture removes the 0th shell already - # density_0 = density_0.to("g/cm^3")[1:] - # density_0 = density_0.insert(0, 0) - density = calculate_density_after_time( - density_0, time_0, time_explosion - ) - - return density - - -def parse_radiation_field_state( - config, t_radiative, geometry, dilution_factor=None, packet_source=None -): - """ - Parses the radiation field state based on the provided configuration, radiative temperature, geometry, dilution factor, and packet source. - - Parameters - ---------- - config : Config - The configuration object. - t_radiative : {None, Quantity}, optional - The radiative temperature. If None, it is calculated based on the initial_t_rad value in the plasma configuration. - geometry : Geometry - The geometry object. - dilution_factor : {None, ndarray}, optional - The dilution factor. If None, it is calculated based on the geometry. - packet_source : {None, PacketSource}, optional - The packet source object. - - Returns - ------- - DiluteThermalRadiationFieldState - The parsed radiation field state. - - Raises - ------ - AssertionError - If the length of t_radiative or dilution_factor is not compatible with the geometry. - """ - if t_radiative is None: - if 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 - ) - - assert len(t_radiative) == 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 DilutePlanckianRadiationField(t_radiative, dilution_factor, geometry) - - -def initialize_packet_source( - config, geometry, packet_source, legacy_mode_enabled -): - """ - Initialize the packet source based on config and geometry - - Parameters - ---------- - config : Config - The configuration object containing the supernova and plasma settings. - geometry : Geometry - The geometry object containing the inner radius information. - packet_source : BasePacketSource - The packet source object based on the configuration and geometry. - - Returns - ------- - packet_source : BasePacketSource - The packet source object based on the configuration and geometry. - - Raises - ------ - ValueError - If both t_inner and luminosity_requested are None. - """ - if config.montecarlo.enable_full_relativity: - packet_source = BlackBodySimpleSourceRelativistic( - base_seed=config.montecarlo.seed, - time_explosion=config.supernova.time_explosion, - legacy_mode_enabled=legacy_mode_enabled, - ) - else: - packet_source = BlackBodySimpleSource( - base_seed=config.montecarlo.seed, - legacy_mode_enabled=legacy_mode_enabled, - ) - - luminosity_requested = config.supernova.luminosity_requested - if config.plasma.initial_t_inner > 0.0 * u.K: - packet_source.radius = geometry.r_inner_active[0] - packet_source.temperature = config.plasma.initial_t_inner - - elif (config.plasma.initial_t_inner < 0.0 * u.K) and ( - luminosity_requested is not None - ): - packet_source.radius = geometry.r_inner_active[0] - packet_source.set_temperature_from_luminosity(luminosity_requested) - else: - raise ValueError( - "Both t_inner and luminosity_requested cannot be None." - ) - - return packet_source - - -def parse_packet_source(config, geometry, legacy_mode_enabled): - """ - Parse the packet source based on the given configuration and geometry. - - Parameters - ---------- - config : Config - The configuration object containing the supernova and plasma settings. - geometry : Geometry - The geometry object containing the inner radius information. - - Returns - ------- - packet_source : BlackBodySimpleSource - The packet source object based on the configuration and geometry. - """ - if config.montecarlo.enable_full_relativity: - packet_source = BlackBodySimpleSourceRelativistic( - base_seed=config.montecarlo.seed, - time_explosion=config.supernova.time_explosion, - legacy_mode_enabled=legacy_mode_enabled, - ) - else: - packet_source = BlackBodySimpleSource( - base_seed=config.montecarlo.seed, - legacy_mode_enabled=legacy_mode_enabled, - ) - - return initialize_packet_source( - config, geometry, packet_source, legacy_mode_enabled - ) - - -def parse_csvy_radiation_field_state( - config, csvy_model_config, csvy_model_data, geometry, packet_source -): - 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 DilutePlanckianRadiationField(t_radiative, dilution_factor, geometry) - - -def calculate_t_radiative_from_t_inner(geometry, packet_source): - """ - Calculates the radiative temperature based on the inner temperature and the geometry of the system. - - Parameters - ---------- - geometry : Geometry - The geometry object. - packet_source : PacketSource - The packet source object. - - Returns - ------- - 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): - return 0.5 * ( - 1 - - np.sqrt( - 1 - - ( - geometry.r_inner[geometry.v_inner_boundary_index] ** 2 - / geometry.r_middle**2 - ) - .to(1) - .value - ) - ) From 3ec64f180acc69557559be9544ad73a47f613a81 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 23 Jul 2024 16:54:44 -0400 Subject: [PATCH 37/89] chore: Refactor radiation field configuration parsing and state creation Refactor the `parse_radiation_field_configuration.py` module to improve code organization and readability. Update the import statements to reflect the changes in the module structure. Replace the deprecated `DiluteBlackBodyRadiationFieldState` class with the new `DilutePlanckianRadiationField` class from the `tardis.plasma.radiation_field` module. This change ensures consistency and compatibility with the latest codebase. Also, update the `tardis/simulation/base.py` module to import the `DilutePlanckianRadiationField` class from the `tardis.plasma.radiation_field` module. This change ensures that the correct class is used for creating the radiation field in the simulation. --- .../io/model/parse_radiation_field_configuration.py | 13 +++---------- tardis/simulation/base.py | 5 +++-- 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/tardis/io/model/parse_radiation_field_configuration.py b/tardis/io/model/parse_radiation_field_configuration.py index e12c319feb2..d311ad94923 100644 --- a/tardis/io/model/parse_radiation_field_configuration.py +++ b/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__) @@ -67,9 +64,7 @@ def parse_radiation_field_state_from_config( 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( @@ -132,9 +127,7 @@ def parse_radiation_field_state_from_csvy( 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): diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 48a401fe04b..e6933a4651e 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -17,14 +17,15 @@ ) from tardis.io.util import HDFWriterMixin from tardis.model import SimulationState +from tardis.plasma.radiation_field import DilutePlanckianRadiationField from tardis.plasma.standard_plasmas import assemble_plasma -from tardis.spectrum.formal_integral import FormalIntegrator from tardis.simulation.convergence import ConvergenceSolver +from tardis.spectrum.base import SpectrumSolver +from tardis.spectrum.formal_integral import FormalIntegrator from tardis.transport.montecarlo.base import MonteCarloTransportSolver from tardis.transport.montecarlo.configuration import montecarlo_globals from tardis.util.base import is_notebook from tardis.visualization import ConvergencePlots -from tardis.spectrum.base import SpectrumSolver # Adding logging support logger = logging.getLogger(__name__) From f1f0cddead4ae286c71c71d042e6d82e481b446e Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 23 Jul 2024 16:56:03 -0400 Subject: [PATCH 38/89] revert astropy_helpers --- astropy_helpers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/astropy_helpers b/astropy_helpers index 9f82aac6c21..3b45ed3191c 160000 --- a/astropy_helpers +++ b/astropy_helpers @@ -1 +1 @@ -Subproject commit 9f82aac6c2141b425e2d639560f7260189d90b54 +Subproject commit 3b45ed3191ceb45c574db304ec0f33282d2e4a98 From 2a7d460e7c7d892cb5cf224895b888dd15152257 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 23 Jul 2024 16:56:52 -0400 Subject: [PATCH 39/89] remove astropy_helpers --- astropy_helpers | 1 - 1 file changed, 1 deletion(-) delete mode 160000 astropy_helpers diff --git a/astropy_helpers b/astropy_helpers deleted file mode 160000 index 3b45ed3191c..00000000000 --- a/astropy_helpers +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 3b45ed3191ceb45c574db304ec0f33282d2e4a98 From dbc34acc25b3cf68ec79103b25692f16d5ca9b1f Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 23 Jul 2024 17:16:25 -0400 Subject: [PATCH 40/89] removed test.txt --- test.txt | 2 -- 1 file changed, 2 deletions(-) delete mode 100644 test.txt diff --git a/test.txt b/test.txt deleted file mode 100644 index 1f2a7fffe05..00000000000 --- a/test.txt +++ /dev/null @@ -1,2 +0,0 @@ -conda-linux-64.lock:https://conda.anaconda.org/conda-forge/linux-64/pandas-2.2.2-py312h1d6d2e6_1.conda#ae00b61f3000d2284d1f2584d4dfafa8 -conda-osx-64.lock:https://conda.anaconda.org/conda-forge/osx-64/pandas-2.2.2-py312h1171441_1.conda#240737937f1f046b0e03ecc11ac4ec98 From 251d1cbb32cf2bb1a2b9ebb884aa2de30eb0bca4 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 23 Jul 2024 17:38:26 -0400 Subject: [PATCH 41/89] blackified code --- tardis/plasma/properties/base.py | 1 + tardis/plasma/properties/plasma_input.py | 1 + tardis/simulation/base.py | 6 +++--- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tardis/plasma/properties/base.py b/tardis/plasma/properties/base.py index 1bac9e7658e..351b42a4d02 100644 --- a/tardis/plasma/properties/base.py +++ b/tardis/plasma/properties/base.py @@ -206,6 +206,7 @@ def _set_output_value(self, output, value): class ObjectInput(Input): input_object_map = {} # mapping output names from input object attributes + def set_value(self, value): for output in self.outputs: if output in self.input_object_map: diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index f9d70216dcc..06ff2a23c94 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -50,6 +50,7 @@ class DilutionFactor(ProcessingPlasmaProperty): outputs = ("w",) latex_name = ("W",) + def calculate(self, dilute_planckian_radiation_field): return dilute_planckian_radiation_field.dilution_factor diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index e6933a4651e..09fd0cf942d 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -400,9 +400,9 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0): # Set up spectrum solver self.spectrum_solver.transport_state = transport_state - self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = ( - v_packets_energy_hist - ) + self.spectrum_solver._montecarlo_virtual_luminosity.value[ + : + ] = v_packets_energy_hist output_energy = ( self.transport.transport_state.packet_collection.output_energies From d49a49c4c24bceaeba6781aa4050c3740ce2e275 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 23 Jul 2024 18:01:42 -0400 Subject: [PATCH 42/89] cleanup simulation from merges --- tardis/simulation/base.py | 8 ++++---- tardis/simulation/initialization.py | 4 +++- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index c587956849e..cc67a409b51 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -10,14 +10,14 @@ import tardis from tardis import constants as const from tardis.io.configuration.config_reader import ConfigurationError -from tardis.io.model.parse_packet_source_configuration import ( - initialize_packet_source, -) from tardis.io.util import HDFWriterMixin -from tardis.model import SimulationState from tardis.plasma.radiation_field import DilutePlanckianRadiationField from tardis.plasma.standard_plasmas import assemble_plasma from tardis.simulation.convergence import ConvergenceSolver +from tardis.simulation.initialization import ( + initialize_atom_data, + initialize_simulation_state, +) from tardis.spectrum.base import SpectrumSolver from tardis.spectrum.formal_integral import FormalIntegrator from tardis.transport.montecarlo.base import MonteCarloTransportSolver diff --git a/tardis/simulation/initialization.py b/tardis/simulation/initialization.py index e95ab558e69..e140d61b647 100644 --- a/tardis/simulation/initialization.py +++ b/tardis/simulation/initialization.py @@ -2,8 +2,10 @@ from pathlib import Path from tardis.io.atom_data.base import AtomData +from tardis.io.model.parse_packet_source_configuration import ( + initialize_packet_source, +) from tardis.model import SimulationState -from tardis.model.parse_input import initialize_packet_source logger = logging.getLogger(__name__) From adbaf7fa59f111f64bb2a26d324e955aa5cd8636 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 12:13:53 -0400 Subject: [PATCH 43/89] fix --- tardis/io/model/readers/csvy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/io/model/readers/csvy.py b/tardis/io/model/readers/csvy.py index 2c03272b5ec..7e6c1f1c73b 100644 --- a/tardis/io/model/readers/csvy.py +++ b/tardis/io/model/readers/csvy.py @@ -23,7 +23,7 @@ def load_csvy(fname): ------- 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: From bb19fdb276844e5da5881f0afce009a4d8d3485c Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 13:10:33 -0400 Subject: [PATCH 44/89] remove unused Input --- tardis/plasma/properties/base.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/tardis/plasma/properties/base.py b/tardis/plasma/properties/base.py index 351b42a4d02..c81038143fc 100644 --- a/tardis/plasma/properties/base.py +++ b/tardis/plasma/properties/base.py @@ -202,20 +202,6 @@ class ArrayInput(Input): def _set_output_value(self, output, value): setattr(self, output, np.array(value, copy=False)) - -class ObjectInput(Input): - - input_object_map = {} # mapping output names from input object attributes - - def set_value(self, value): - for output in self.outputs: - if output in self.input_object_map: - object_attr = self.input_object_map[output] - self._set_output_value(output, getattr(value, object_attr)) - else: - self._set_output_value(output, getattr(value, output)) - - class DataFrameInput(Input): def _set_output_value(self, output, value): setattr(self, output, np.array(pd.DataFrame(value), copy=False)) From aca109977e9154d1192583694fa4959c73896a5f Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 13:46:20 -0400 Subject: [PATCH 45/89] added description --- tardis/plasma/properties/plasma_input.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index 06ff2a23c94..ac77507f978 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -27,6 +27,8 @@ class TRadiative(ProcessingPlasmaProperty): """ + Radiative temperature property. + Attributes ---------- t_rad : Numpy Array, dtype float @@ -41,6 +43,8 @@ def calculate(self, dilute_planckian_radiation_field): class DilutionFactor(ProcessingPlasmaProperty): """ + Dilution factor of the radiation field. + Attributes ---------- w : Numpy Array, dtype float between 0 and 1 From 54ffe6854c62d0233e57e0f7a5c6f0ddd1b036f6 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 14:03:42 -0400 Subject: [PATCH 46/89] blackiefied codebase --- tardis/plasma/properties/base.py | 1 + tardis/simulation/base.py | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tardis/plasma/properties/base.py b/tardis/plasma/properties/base.py index c81038143fc..9ed4f2f5229 100644 --- a/tardis/plasma/properties/base.py +++ b/tardis/plasma/properties/base.py @@ -202,6 +202,7 @@ class ArrayInput(Input): def _set_output_value(self, output, value): setattr(self, output, np.array(value, copy=False)) + class DataFrameInput(Input): def _set_output_value(self, output, value): setattr(self, output, np.array(pd.DataFrame(value), copy=False)) diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index cc67a409b51..906d7e0a99e 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -397,9 +397,9 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0): # Set up spectrum solver self.spectrum_solver.transport_state = transport_state - self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = ( - v_packets_energy_hist - ) + self.spectrum_solver._montecarlo_virtual_luminosity.value[ + : + ] = v_packets_energy_hist output_energy = ( self.transport.transport_state.packet_collection.output_energies From 2ea3122abc4ad0a291c92a69b15f52d30ddb3df4 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 16:26:39 -0400 Subject: [PATCH 47/89] cleanup of branch --- tardis/transport/montecarlo/base.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tardis/transport/montecarlo/base.py b/tardis/transport/montecarlo/base.py index cd5cf9bfe59..efb6ced9e9d 100644 --- a/tardis/transport/montecarlo/base.py +++ b/tardis/transport/montecarlo/base.py @@ -7,9 +7,9 @@ from tardis import constants as const from tardis.io.logger import montecarlo_tracking as mc_tracker from tardis.io.util import HDFWriterMixin -from tardis.transport.montecarlo.montecarlo_main_loop import ( - montecarlo_main_loop, - numba_config, +from tardis.transport.montecarlo.configuration.base import ( + MonteCarloConfiguration, + configuration_initialize, ) from tardis.transport.montecarlo.estimators.dilute_blackbody_properties import ( MCRadiationFieldPropertiesSolver, @@ -17,6 +17,9 @@ from tardis.transport.montecarlo.estimators.radfield_mc_estimators import ( initialize_estimator_statistics, ) +from tardis.transport.montecarlo.montecarlo_main_loop import ( + montecarlo_main_loop, +) from tardis.transport.montecarlo.montecarlo_transport_state import ( MonteCarloTransportState, ) From 9800fc821ee65b2fde77a09441c245c6b19e91e8 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 16:36:46 -0400 Subject: [PATCH 48/89] blackify code --- tardis/plasma/properties/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tardis/plasma/properties/base.py b/tardis/plasma/properties/base.py index c81038143fc..9ed4f2f5229 100644 --- a/tardis/plasma/properties/base.py +++ b/tardis/plasma/properties/base.py @@ -202,6 +202,7 @@ class ArrayInput(Input): def _set_output_value(self, output, value): setattr(self, output, np.array(value, copy=False)) + class DataFrameInput(Input): def _set_output_value(self, output, value): setattr(self, output, np.array(pd.DataFrame(value), copy=False)) From d19e423e7e4eb46b96f199c636168a1055d8433b Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 16:43:15 -0400 Subject: [PATCH 49/89] chore: Refactor atom data parsing and simulation state initialization --- tardis/io/model/parse_atom_data.py | 52 ++++++++++++++++++ .../model/parse_simulation_state.py} | 55 +------------------ tardis/simulation/base.py | 10 ++-- 3 files changed, 58 insertions(+), 59 deletions(-) create mode 100644 tardis/io/model/parse_atom_data.py rename tardis/{simulation/initialization.py => io/model/parse_simulation_state.py} (50%) diff --git a/tardis/io/model/parse_atom_data.py b/tardis/io/model/parse_atom_data.py new file mode 100644 index 00000000000..ccd25805df5 --- /dev/null +++ b/tardis/io/model/parse_atom_data.py @@ -0,0 +1,52 @@ +from tardis.io.atom_data.base import AtomData +from tardis.simulation.initialization import logger + + +from pathlib import Path + + +def parse_atom_data(config, atom_data=None): + """ + Parse atom data for the simulation. + + Parameters + ---------- + config : object + The configuration object containing information about the atom data. + atom_data : object, optional + Existing atom data to be used, if provided. + + Returns + ------- + object + The initialized atom data. + + Raises + ------ + ValueError + If no atom_data option is found in the configuration. + """ + if atom_data is None: + if "atom_data" in config: + if Path(config.atom_data).is_absolute(): + atom_data_fname = Path(config.atom_data) + else: + atom_data_fname = Path(config.config_dirname) / config.atom_data + + else: + raise ValueError("No atom_data option found in the configuration.") + + logger.info(f"\n\tReading Atomic Data from {atom_data_fname}") + + try: + atom_data = AtomData.from_hdf(atom_data_fname) + except TypeError as e: + print( + e, + "Error might be from the use of an old-format of the atomic database, \n" + "please see https://github.com/tardis-sn/tardis-refdata/tree/master/atom_data" + " for the most recent version.", + ) + raise + + return atom_data diff --git a/tardis/simulation/initialization.py b/tardis/io/model/parse_simulation_state.py similarity index 50% rename from tardis/simulation/initialization.py rename to tardis/io/model/parse_simulation_state.py index e140d61b647..ff9773bd126 100644 --- a/tardis/simulation/initialization.py +++ b/tardis/io/model/parse_simulation_state.py @@ -1,63 +1,10 @@ -import logging -from pathlib import Path - -from tardis.io.atom_data.base import AtomData from tardis.io.model.parse_packet_source_configuration import ( initialize_packet_source, ) from tardis.model import SimulationState -logger = logging.getLogger(__name__) - - -def initialize_atom_data(config, atom_data=None): - """ - Initialize atom data for the simulation. - - Parameters - ---------- - config : object - The configuration object containing information about the atom data. - atom_data : object, optional - Existing atom data to be used, if provided. - - Returns - ------- - object - The initialized atom data. - - Raises - ------ - ValueError - If no atom_data option is found in the configuration. - """ - if atom_data is None: - if "atom_data" in config: - if Path(config.atom_data).is_absolute(): - atom_data_fname = Path(config.atom_data) - else: - atom_data_fname = Path(config.config_dirname) / config.atom_data - - else: - raise ValueError("No atom_data option found in the configuration.") - - logger.info(f"\n\tReading Atomic Data from {atom_data_fname}") - - try: - atom_data = AtomData.from_hdf(atom_data_fname) - except TypeError as e: - print( - e, - "Error might be from the use of an old-format of the atomic database, \n" - "please see https://github.com/tardis-sn/tardis-refdata/tree/master/atom_data" - " for the most recent version.", - ) - raise - - return atom_data - -def initialize_simulation_state( +def parse_simulation_state( config, packet_source, legacy_mode_enabled, kwargs, atom_data ): """ diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 906d7e0a99e..8ce3047f2af 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -14,10 +14,10 @@ from tardis.plasma.radiation_field import DilutePlanckianRadiationField from tardis.plasma.standard_plasmas import assemble_plasma from tardis.simulation.convergence import ConvergenceSolver -from tardis.simulation.initialization import ( - initialize_atom_data, - initialize_simulation_state, +from tardis.io.model.parse_simulation_state import ( + parse_simulation_state, ) +from tardis.io.model.parse_atom_data import parse_atom_data from tardis.spectrum.base import SpectrumSolver from tardis.spectrum.formal_integral import FormalIntegrator from tardis.transport.montecarlo.base import MonteCarloTransportSolver @@ -673,8 +673,8 @@ def from_config( # Allow overriding some config structures. This is useful in some # unit tests, and could be extended in all the from_config classmethods. - atom_data = initialize_atom_data(config, atom_data=atom_data) - simulation_state = initialize_simulation_state( + atom_data = parse_atom_data(config, atom_data=atom_data) + simulation_state = parse_simulation_state( config, packet_source, legacy_mode_enabled, kwargs, atom_data ) if plasma is None: From b3bf356d241506573c8f7e4a3c10f733c5972fe3 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 16:54:57 -0400 Subject: [PATCH 50/89] restructure logger --- tardis/io/model/parse_atom_data.py | 8 ++++---- tardis/io/model/parse_packet_source_configuration.py | 6 +++--- tardis/io/model/parse_simulation_state.py | 8 ++++---- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/tardis/io/model/parse_atom_data.py b/tardis/io/model/parse_atom_data.py index ccd25805df5..5cfe30197d7 100644 --- a/tardis/io/model/parse_atom_data.py +++ b/tardis/io/model/parse_atom_data.py @@ -1,9 +1,9 @@ -from tardis.io.atom_data.base import AtomData -from tardis.simulation.initialization import logger - - +import logging from pathlib import Path +from tardis.io.atom_data.base import AtomData + +logger = logging.getLogger(__name__) def parse_atom_data(config, atom_data=None): """ diff --git a/tardis/io/model/parse_packet_source_configuration.py b/tardis/io/model/parse_packet_source_configuration.py index 90aab02e556..2e711001b76 100644 --- a/tardis/io/model/parse_packet_source_configuration.py +++ b/tardis/io/model/parse_packet_source_configuration.py @@ -46,7 +46,7 @@ def initialize_packet_source(packet_source, config, geometry): return packet_source -def parse_packet_source_from_config(config, geometry, legacy_mode_enabled): +def parse_packet_source_from_config(config, geometry, enable_legacy_mode): """ Parse the packet source based on the given configuration and geometry. @@ -66,12 +66,12 @@ def parse_packet_source_from_config(config, geometry, legacy_mode_enabled): packet_source = BlackBodySimpleSourceRelativistic( base_seed=config.montecarlo.seed, time_explosion=config.supernova.time_explosion, - legacy_mode_enabled=legacy_mode_enabled, + legacy_mode_enabled=enable_legacy_mode, ) else: packet_source = BlackBodySimpleSource( base_seed=config.montecarlo.seed, - legacy_mode_enabled=legacy_mode_enabled, + legacy_mode_enabled=enable_legacy_mode, ) return initialize_packet_source(packet_source, config, geometry) diff --git a/tardis/io/model/parse_simulation_state.py b/tardis/io/model/parse_simulation_state.py index ff9773bd126..ffd5266a017 100644 --- a/tardis/io/model/parse_simulation_state.py +++ b/tardis/io/model/parse_simulation_state.py @@ -5,7 +5,7 @@ def parse_simulation_state( - config, packet_source, legacy_mode_enabled, kwargs, atom_data + config, packet_source, enable_legacy_mode, kwargs, atom_data ): """ Initialize the simulation state. @@ -35,20 +35,20 @@ def parse_simulation_state( simulation_state = SimulationState.from_csvy( config, atom_data=atom_data, - legacy_mode_enabled=legacy_mode_enabled, + legacy_mode_enabled=enable_legacy_mode, ) else: simulation_state = SimulationState.from_config( config, atom_data=atom_data, - legacy_mode_enabled=legacy_mode_enabled, + legacy_mode_enabled=enable_legacy_mode, ) if packet_source is not None: simulation_state.packet_source = initialize_packet_source( config, simulation_state.geometry, packet_source, - legacy_mode_enabled, + enable_legacy_mode, ) return simulation_state From 734ffc80e395e566c012bcbaf596b089e61e0c92 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 17:11:50 -0400 Subject: [PATCH 51/89] working on continuum radfield properties --- .../montecarlo/estimators/continuum_radfield_properties.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py index da742cc1ee3..4112723dade 100644 --- a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py +++ b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py @@ -6,6 +6,7 @@ import tardis.constants as const from tardis.io.atom_data import AtomData +from tardis.plasma.properties.continuum_processes import PhotoIonBoltzmannFactor from tardis.plasma.radiation_field.planck_rad_field import ( DilutePlanckianRadiationField, ) @@ -16,7 +17,6 @@ bound_free_estimator_array2frame, integrate_array_by_blocks, ) -from tardis.plasma.properties.continuum_processes import PhotoIonBoltzmannFactor H = const.h.cgs.value From d69533a4806389a057f8f91112b74b705d677a53 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 17:58:18 -0400 Subject: [PATCH 52/89] Refactor continuum processes module structure --- tardis/io/atom_data/base.py | 2 +- tardis/opacities/continuum/__init__.py | 0 tardis/opacities/continuum/bound_free.py | 36 ++ .../continuum_processes/__init__.py | 0 .../collisional_ion_trans_prob.py | 117 ++++++ .../macro_atom/transition_probabilities.py | 2 +- tardis/opacities/opacities.py | 6 +- tardis/plasma/properties/__init__.py | 2 +- tardis/plasma/properties/atomic.py | 2 +- .../continuum_processes/__init__.py | 0 .../continuum_processes/fast_array_util.py | 66 ++++ .../photo_ion_rate_coeff.py | 80 ++++ .../rates.py} | 364 +----------------- .../stim_recomb_rate_coeff.py | 85 ++++ tardis/plasma/properties/ion_population.py | 4 +- tardis/plasma/properties/util/__init__.py | 3 - .../util/integrate_array_by_blocks.py | 35 -- .../continuum_radfield_properties.py | 19 +- .../tests/test_continuum_property_solver.py | 8 +- 19 files changed, 416 insertions(+), 415 deletions(-) create mode 100644 tardis/opacities/continuum/__init__.py create mode 100644 tardis/opacities/continuum/bound_free.py create mode 100644 tardis/opacities/macro_atom/continuum_processes/__init__.py create mode 100644 tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py create mode 100644 tardis/plasma/properties/continuum_processes/__init__.py create mode 100644 tardis/plasma/properties/continuum_processes/fast_array_util.py create mode 100644 tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py rename tardis/plasma/properties/{continuum_processes.py => continuum_processes/rates.py} (69%) create mode 100644 tardis/plasma/properties/continuum_processes/stim_recomb_rate_coeff.py delete mode 100644 tardis/plasma/properties/util/__init__.py delete mode 100644 tardis/plasma/properties/util/integrate_array_by_blocks.py diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 90628b0ec72..7011f995841 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -9,7 +9,7 @@ from tardis import constants as const from tardis.io.atom_data.util import resolve_atom_data_fname -from tardis.plasma.properties.continuum_processes import ( +from tardis.plasma.properties.continuum_proc.continuum_processes import ( get_ground_state_multi_index, ) diff --git a/tardis/opacities/continuum/__init__.py b/tardis/opacities/continuum/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tardis/opacities/continuum/bound_free.py b/tardis/opacities/continuum/bound_free.py new file mode 100644 index 00000000000..e103f942ade --- /dev/null +++ b/tardis/opacities/continuum/bound_free.py @@ -0,0 +1,36 @@ +from tardis.plasma.exceptions import PlasmaException +from tardis.plasma.properties.base import ProcessingPlasmaProperty + + +class BoundFreeOpacity(ProcessingPlasmaProperty): + """ + Attributes + ---------- + chi_bf : pandas.DataFrame, dtype float + Bound-free opacity corrected for stimulated emission. + """ + + outputs = ("chi_bf",) + latex_name = (r"\chi^{\textrm{bf}}",) + + def calculate( + self, + photo_ion_cross_sections, + t_electrons, + phi_ik, + level_number_density, + lte_level_number_density, + boltzmann_factor_photo_ion, + ): + cross_section = photo_ion_cross_sections["x_sect"].values + + n_i = level_number_density.loc[photo_ion_cross_sections.index] + lte_n_i = lte_level_number_density.loc[photo_ion_cross_sections.index] + chi_bf = (n_i - lte_n_i * boltzmann_factor_photo_ion).multiply( + cross_section, axis=0 + ) + + num_neg_elements = (chi_bf < 0).sum().sum() + if num_neg_elements: + raise PlasmaException("Negative values in bound-free opacity.") + return chi_bf diff --git a/tardis/opacities/macro_atom/continuum_processes/__init__.py b/tardis/opacities/macro_atom/continuum_processes/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py b/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py new file mode 100644 index 00000000000..9747faa988b --- /dev/null +++ b/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py @@ -0,0 +1,117 @@ +from tardis.plasma.properties.base import ( + ProcessingPlasmaProperty, + TransitionProbabilitiesProperty, +) +from tardis.plasma.properties.continuum_processes.rates import ( + H, + IndexSetterMixin, +) + + +import pandas as pd + + +class RawCollIonTransProbs(TransitionProbabilitiesProperty, IndexSetterMixin): + """ + Attributes + ---------- + p_coll_ion : pandas.DataFrame, dtype float + The unnormalized transition probabilities for + collisional ionization. + p_coll_recomb : pandas.DataFrame, dtype float + The unnormalized transition probabilities for + collisional recombination. + cool_rate_coll_ion : pandas.DataFrame, dtype float + The collisional ionization cooling rates of the electron gas. + """ + + outputs = ("p_coll_ion", "p_coll_recomb", "cool_rate_coll_ion") + transition_probabilities_outputs = ( + "p_coll_ion", + "p_coll_recomb", + "cool_rate_coll_ion", + ) + latex_name = ( + r"p^{\textrm{coll ion}}", + r"p^{\textrm{coll recomb}}", + r"C^{\textrm{ion}}", + ) + + def calculate( + self, + coll_ion_coeff, + coll_recomb_coeff, + nu_i, + photo_ion_idx, + electron_densities, + energy_i, + level_number_density, + ): + p_coll_ion = coll_ion_coeff.multiply(energy_i, axis=0) + p_coll_ion = p_coll_ion.multiply(electron_densities, axis=1) + p_coll_ion = self.set_index(p_coll_ion, photo_ion_idx, reverse=False) + + coll_recomb_rate = coll_recomb_coeff.multiply( + electron_densities, axis=1 + ) # The full rate is obtained from this by multiplying by the + # electron density and ion number density. + p_recomb_deactivation = coll_recomb_rate.multiply(nu_i, axis=0) * H + p_recomb_deactivation = self.set_index( + p_recomb_deactivation, photo_ion_idx, transition_type=-1 + ) + p_recomb_deactivation = p_recomb_deactivation.groupby(level=[0]).sum() + index_dd = pd.MultiIndex.from_product( + [p_recomb_deactivation.index.values, ["k"], [0]], + names=list(photo_ion_idx.columns) + ["transition_type"], + ) + p_recomb_deactivation = p_recomb_deactivation.set_index(index_dd) + + p_recomb_internal = coll_recomb_rate.multiply(energy_i, axis=0) + p_recomb_internal = self.set_index( + p_recomb_internal, photo_ion_idx, transition_type=0 + ) + p_coll_recomb = pd.concat([p_recomb_deactivation, p_recomb_internal]) + + cool_rate_coll_ion = (coll_ion_coeff * electron_densities).multiply( + nu_i * H, axis=0 + ) + level_lower_index = coll_ion_coeff.index + cool_rate_coll_ion = ( + cool_rate_coll_ion + * level_number_density.loc[level_lower_index].values + ) + cool_rate_coll_ion = self.set_index( + cool_rate_coll_ion, photo_ion_idx, reverse=False + ) + cool_rate_coll_ion = cool_rate_coll_ion.groupby( + level="destination_level_idx" + ).sum() + ion_cool_index = pd.MultiIndex.from_product( + [["k"], cool_rate_coll_ion.index.values, [0]], + names=list(photo_ion_idx.columns) + ["transition_type"], + ) + cool_rate_coll_ion = cool_rate_coll_ion.set_index(ion_cool_index) + return p_coll_ion, p_coll_recomb, cool_rate_coll_ion + + +class CollRecombRateCoeff(ProcessingPlasmaProperty): + """ + Attributes + ---------- + coll_recomb_coeff : pandas.DataFrame, dtype float + The rate coefficient for collisional recombination. + Multiply with the electron density squared and the ion number density + to obtain the total rate. + + Notes + ----- + The collisional recombination rate coefficient is calculated from the + collisional ionization rate coefficient based on the requirement of detailed + balance. + """ + + outputs = ("coll_recomb_coeff",) + latex_name = (r"c_{\kappa\textrm{i,}}",) + + def calculate(self, phi_ik, coll_ion_coeff): + return coll_ion_coeff.multiply(phi_ik.loc[coll_ion_coeff.index]) diff --git a/tardis/opacities/macro_atom/transition_probabilities.py b/tardis/opacities/macro_atom/transition_probabilities.py index 6facd5ee70a..c22d337a6fb 100644 --- a/tardis/opacities/macro_atom/transition_probabilities.py +++ b/tardis/opacities/macro_atom/transition_probabilities.py @@ -5,7 +5,7 @@ from scipy import sparse as sp from tardis.plasma.properties.base import ProcessingPlasmaProperty -from tardis.plasma.properties.continuum_processes import ( +from tardis.plasma.properties.continuum_proc.continuum_processes import ( get_ground_state_multi_index, ) from tardis.transport.montecarlo.macro_atom import ( diff --git a/tardis/opacities/opacities.py b/tardis/opacities/opacities.py index dd7aeae83d4..71f193b9389 100644 --- a/tardis/opacities/opacities.py +++ b/tardis/opacities/opacities.py @@ -278,7 +278,8 @@ def compton_opacity_partial(energy, fraction): @njit(**njit_dict_no_parallel) def compton_opacity_calculation(energy, electron_density): - """Calculate the Compton scattering opacity for a given energy + """ + Calculate the Compton scattering opacity for a given energy (Rybicki & Lightman, 1979) $ @@ -364,7 +365,8 @@ def photoabsorption_opacity_calculation( def photoabsorption_opacity_calculation_kasen( energy, number_density, proton_count ): - """Calculates photoabsorption opacity for a given energy + """ + Calculates photoabsorption opacity for a given energy Approximate treatment from Kasen et al. (2006) Parameters diff --git a/tardis/plasma/properties/__init__.py b/tardis/plasma/properties/__init__.py index 3363cbd9bcd..d15c03afc74 100644 --- a/tardis/plasma/properties/__init__.py +++ b/tardis/plasma/properties/__init__.py @@ -7,7 +7,7 @@ from tardis.opacities.macro_atom.transition_probabilities import * from tardis.plasma.properties.atomic import * -from tardis.plasma.properties.continuum_processes import * +from tardis.plasma.properties.continuum_proc.continuum_processes import * from tardis.plasma.properties.general import * from tardis.plasma.properties.helium_nlte import * from tardis.plasma.properties.ion_population import * diff --git a/tardis/plasma/properties/atomic.py b/tardis/plasma/properties/atomic.py index 72fd4b6b472..5dcc206c784 100644 --- a/tardis/plasma/properties/atomic.py +++ b/tardis/plasma/properties/atomic.py @@ -15,7 +15,7 @@ HiddenPlasmaProperty, ProcessingPlasmaProperty, ) -from tardis.plasma.properties.continuum_processes import ( +from tardis.plasma.properties.continuum_proc.continuum_processes import ( A0, BETA_COLL, K_B, diff --git a/tardis/plasma/properties/continuum_processes/__init__.py b/tardis/plasma/properties/continuum_processes/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tardis/plasma/properties/continuum_processes/fast_array_util.py b/tardis/plasma/properties/continuum_processes/fast_array_util.py new file mode 100644 index 00000000000..8e48a0dda80 --- /dev/null +++ b/tardis/plasma/properties/continuum_processes/fast_array_util.py @@ -0,0 +1,66 @@ +# It is currently not possible to use scipy.integrate.cumulative_trapezoid in +# numba. So here is my own implementation. +from tardis.transport.montecarlo import njit_dict + + +import numpy as np +from numba import njit, prange + + +@njit(**njit_dict) +def numba_cumulative_trapezoid(f, x): + """ + Cumulatively integrate f(x) using the composite trapezoidal rule. + + Parameters + ---------- + f : numpy.ndarray, dtype float + Input array to integrate. + x : numpy.ndarray, dtype float + The coordinate to integrate along. + + Returns + ------- + numpy.ndarray, dtype float + The result of cumulative integration of f along x + """ + integ = (np.diff(x) * (f[1:] + f[:-1]) / 2.0).cumsum() + return integ / integ[-1] + + +@njit(**njit_dict) +def cumulative_integrate_array_by_blocks(f, x, block_references): + """ + Cumulatively integrate a function over blocks. + + This function cumulatively integrates a function `f` defined at + locations `x` over blocks given in `block_references`. + + Parameters + ---------- + f : numpy.ndarray, dtype float + Input array to integrate. Shape is (N_freq, N_shells), where + N_freq is the number of frequency values and N_shells is the number + of computational shells. + x : numpy.ndarray, dtype float + The sample points corresponding to the `f` values. Shape is (N_freq,). + block_references : numpy.ndarray, dtype int + The start indices of the blocks to be integrated. Shape is (N_blocks,). + + Returns + ------- + numpy.ndarray, dtype float + Array with cumulatively integrated values. Shape is (N_freq, N_shells) + same as f. + """ + n_rows = len(block_references) - 1 + integrated = np.zeros_like(f) + for i in prange(f.shape[1]): # columns + # TODO: Avoid this loop through vectorization of cumulative_trapezoid + for j in prange(n_rows): # rows + start = block_references[j] + stop = block_references[j + 1] + integrated[start + 1 : stop, i] = numba_cumulative_trapezoid( + f[start:stop, i], x[start:stop] + ) + return integrated diff --git a/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py b/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py new file mode 100644 index 00000000000..1f4c719cb39 --- /dev/null +++ b/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py @@ -0,0 +1,80 @@ +import numpy as np +import pandas as pd + +from tardis.plasma.properties.base import Input +from tardis.plasma.properties.continuum_processes.rates import H +from tardis.transport.montecarlo.estimators.util import ( + ProcessingPlasmaProperty, + bound_free_estimator_array2frame, + integrate_array_by_blocks, +) + + +class PhotoIonRateCoeffEstimator(Input): + """ + Attributes + ---------- + gamma_estimator : pandas.DataFrame, dtype float + Unnormalized MC estimator for the rate coefficient for radiative + ionization. + """ + + outputs = ("gamma_estimator",) + latex_name = (r"\gamma_\textrm{estim}",) + + +class PhotoIonRateCoeff(ProcessingPlasmaProperty): + """ + Attributes + ---------- + gamma : pandas.DataFrame, dtype float + The rate coefficient for radiative ionization. + """ + + outputs = ("gamma",) + latex_name = (r"\gamma",) + + def calculate( + self, + photo_ion_cross_sections, + gamma_estimator, + photo_ion_norm_factor, + photo_ion_block_references, + photo_ion_index, + dilute_planckian_radiation_field, + level2continuum_idx, + ): + # Used for initialization + if gamma_estimator is None: + gamma = self.calculate_from_dilute_bb( + photo_ion_cross_sections, + photo_ion_block_references, + photo_ion_index, + dilute_planckian_radiation_field, + ) + else: + gamma_estimator = bound_free_estimator_array2frame( + gamma_estimator, level2continuum_idx + ) + gamma = gamma_estimator * photo_ion_norm_factor.value + + return gamma + + @staticmethod + def calculate_from_dilute_bb( + photo_ion_cross_sections, + photo_ion_block_references, + photo_ion_index, + dilute_planckian_radiation_field, + ): + nu = photo_ion_cross_sections["nu"] + x_sect = photo_ion_cross_sections["x_sect"] + j_nus = dilute_planckian_radiation_field.calculate_mean_intensity( + nu, + ) + gamma = j_nus.multiply(4.0 * np.pi * x_sect / nu / H, axis=0) + gamma = integrate_array_by_blocks( + gamma.values, nu.values, photo_ion_block_references + ) + gamma = pd.DataFrame(gamma, index=photo_ion_index) + return gamma diff --git a/tardis/plasma/properties/continuum_processes.py b/tardis/plasma/properties/continuum_processes/rates.py similarity index 69% rename from tardis/plasma/properties/continuum_processes.py rename to tardis/plasma/properties/continuum_processes/rates.py index a95c7f2dea3..9c9c6ef96ba 100644 --- a/tardis/plasma/properties/continuum_processes.py +++ b/tardis/plasma/properties/continuum_processes/rates.py @@ -2,7 +2,7 @@ import numpy as np import pandas as pd -from numba import njit, prange +from numba import njit from tardis import constants as const from tardis.plasma.exceptions import PlasmaException @@ -11,19 +11,15 @@ ProcessingPlasmaProperty, TransitionProbabilitiesProperty, ) -from tardis.transport.montecarlo import njit_dict +from tardis.plasma.properties.continuum_processes.fast_array_util import ( + cumulative_integrate_array_by_blocks, + numba_cumulative_trapezoid, +) from tardis.transport.montecarlo.estimators.util import ( - bound_free_estimator_array2frame, integrate_array_by_blocks, ) __all__ = [ - "SpontRecombRateCoeff", - "StimRecombRateCoeff", - "PhotoIonRateCoeff", - "PhotoIonEstimatorsNormFactor", - "PhotoIonRateCoeffEstimator", - "StimRecombRateCoeffEstimator", "CorrPhotoIonRateCoeff", "BfHeatingRateCoeffEstimator", "StimRecombCoolingRateCoeffEstimator", @@ -36,7 +32,6 @@ "AdiabaticCoolingRate", "FreeFreeCoolingRate", "FreeBoundCoolingRate", - "BoundFreeOpacity", "LevelNumberDensityLTE", "PhotoIonBoltzmannFactor", "FreeBoundEmissionCDF", @@ -44,8 +39,6 @@ "TwoPhotonEmissionCDF", "TwoPhotonFrequencySampler", "CollIonRateCoeffSeaton", - "CollRecombRateCoeff", - "RawCollIonTransProbs", ] N_A = const.N_A.cgs.value @@ -69,67 +62,6 @@ logger = logging.getLogger(__name__) -# It is currently not possible to use scipy.integrate.cumulative_trapezoid in -# numba. So here is my own implementation. -@njit(**njit_dict) -def numba_cumulative_trapezoid(f, x): - """ - Cumulatively integrate f(x) using the composite trapezoidal rule. - - Parameters - ---------- - f : numpy.ndarray, dtype float - Input array to integrate. - x : numpy.ndarray, dtype float - The coordinate to integrate along. - - Returns - ------- - numpy.ndarray, dtype float - The result of cumulative integration of f along x - """ - integ = (np.diff(x) * (f[1:] + f[:-1]) / 2.0).cumsum() - return integ / integ[-1] - - -@njit(**njit_dict) -def cumulative_integrate_array_by_blocks(f, x, block_references): - """ - Cumulatively integrate a function over blocks. - - This function cumulatively integrates a function `f` defined at - locations `x` over blocks given in `block_references`. - - Parameters - ---------- - f : numpy.ndarray, dtype float - Input array to integrate. Shape is (N_freq, N_shells), where - N_freq is the number of frequency values and N_shells is the number - of computational shells. - x : numpy.ndarray, dtype float - The sample points corresponding to the `f` values. Shape is (N_freq,). - block_references : numpy.ndarray, dtype int - The start indices of the blocks to be integrated. Shape is (N_blocks,). - - Returns - ------- - numpy.ndarray, dtype float - Array with cumulatively integrated values. Shape is (N_freq, N_shells) - same as f. - """ - n_rows = len(block_references) - 1 - integrated = np.zeros_like(f) - for i in prange(f.shape[1]): # columns - # TODO: Avoid this loop through vectorization of cumulative_trapezoid - for j in prange(n_rows): # rows - start = block_references[j] - stop = block_references[j + 1] - integrated[start + 1 : stop, i] = numba_cumulative_trapezoid( - f[start:stop, i], x[start:stop] - ) - return integrated - - def get_ion_multi_index(multi_index_full, next_higher=True): """ Calculate the corresponding ion MultiIndex for a level MultiIndex. @@ -337,126 +269,6 @@ def calculate( return fb_emission_cdf -class PhotoIonRateCoeff(ProcessingPlasmaProperty): - """ - Attributes - ---------- - gamma : pandas.DataFrame, dtype float - The rate coefficient for radiative ionization. - """ - - outputs = ("gamma",) - latex_name = (r"\gamma",) - - def calculate( - self, - photo_ion_cross_sections, - gamma_estimator, - photo_ion_norm_factor, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - level2continuum_idx, - ): - # Used for initialization - if gamma_estimator is None: - gamma = self.calculate_from_dilute_bb( - photo_ion_cross_sections, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - ) - else: - gamma_estimator = bound_free_estimator_array2frame( - gamma_estimator, level2continuum_idx - ) - gamma = gamma_estimator * photo_ion_norm_factor.value - - return gamma - - @staticmethod - def calculate_from_dilute_bb( - photo_ion_cross_sections, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - ): - nu = photo_ion_cross_sections["nu"] - x_sect = photo_ion_cross_sections["x_sect"] - j_nus = dilute_planckian_radiation_field.calculate_mean_intensity( - nu, - ) - gamma = j_nus.multiply(4.0 * np.pi * x_sect / nu / H, axis=0) - gamma = integrate_array_by_blocks( - gamma.values, nu.values, photo_ion_block_references - ) - gamma = pd.DataFrame(gamma, index=photo_ion_index) - return gamma - - -class StimRecombRateCoeff(ProcessingPlasmaProperty): - """ - Attributes - ---------- - alpha_stim : pandas.DataFrame, dtype float - The rate coefficient for stimulated recombination. - """ - - outputs = ("alpha_stim",) - latex_name = (r"\alpha^{\textrm{stim}}",) - - def calculate( - self, - photo_ion_cross_sections, - alpha_stim_estimator, - photo_ion_norm_factor, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - phi_ik, - t_electrons, - boltzmann_factor_photo_ion, - level2continuum_idx, - ): - # Used for initialization - if alpha_stim_estimator is None: - alpha_stim = self.calculate_from_dilute_bb( - photo_ion_cross_sections, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - t_electrons, - boltzmann_factor_photo_ion, - ) - else: - alpha_stim_estimator = bound_free_estimator_array2frame( - alpha_stim_estimator, level2continuum_idx - ) - alpha_stim = alpha_stim_estimator * photo_ion_norm_factor - alpha_stim *= phi_ik.loc[alpha_stim.index] - return alpha_stim - - @staticmethod - def calculate_from_dilute_bb( - photo_ion_cross_sections, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - t_electrons, - boltzmann_factor_photo_ion, - ): - nu = photo_ion_cross_sections["nu"] - x_sect = photo_ion_cross_sections["x_sect"] - j_nus = dilute_planckian_radiation_field.calculate_mean_intensity(nu) - j_nus *= boltzmann_factor_photo_ion - alpha_stim = j_nus.multiply(4.0 * np.pi * x_sect / nu / H, axis=0) - alpha_stim = integrate_array_by_blocks( - alpha_stim.values, nu.values, photo_ion_block_references - ) - alpha_stim = pd.DataFrame(alpha_stim, index=photo_ion_index) - return alpha_stim - - class RawRecombTransProbs(TransitionProbabilitiesProperty, IndexSetterMixin): """ Attributes @@ -546,32 +358,6 @@ def calculate(time_simulation, volume): return (time_simulation * volume * H) ** -1 -class PhotoIonRateCoeffEstimator(Input): - """ - Attributes - ---------- - gamma_estimator : pandas.DataFrame, dtype float - Unnormalized MC estimator for the rate coefficient for radiative - ionization. - """ - - outputs = ("gamma_estimator",) - latex_name = (r"\gamma_\textrm{estim}",) - - -class StimRecombRateCoeffEstimator(Input): - """ - Attributes - ---------- - alpha_stim_estimator : pandas.DataFrame, dtype float - Unnormalized MC estimator for the rate coefficient for stimulated - recombination. - """ - - outputs = ("alpha_stim_estimator",) - latex_name = (r"\alpha^{\textrm{stim}}_\textrm{estim}",) - - class StimRecombCoolingRateCoeffEstimator(Input): """ Attributes @@ -943,40 +729,6 @@ def calculate( return cool_rate_fb_tot, cool_rate_fb, p_fb_deactivation -class BoundFreeOpacity(ProcessingPlasmaProperty): - """ - Attributes - ---------- - chi_bf : pandas.DataFrame, dtype float - Bound-free opacity corrected for stimulated emission. - """ - - outputs = ("chi_bf",) - latex_name = (r"\chi^{\textrm{bf}}",) - - def calculate( - self, - photo_ion_cross_sections, - t_electrons, - phi_ik, - level_number_density, - lte_level_number_density, - boltzmann_factor_photo_ion, - ): - x_sect = photo_ion_cross_sections["x_sect"].values - - n_i = level_number_density.loc[photo_ion_cross_sections.index] - lte_n_i = lte_level_number_density.loc[photo_ion_cross_sections.index] - chi_bf = (n_i - lte_n_i * boltzmann_factor_photo_ion).multiply( - x_sect, axis=0 - ) - - num_neg_elements = (chi_bf < 0).sum().sum() - if num_neg_elements: - raise PlasmaException("Negative values in bound-free opacity.") - return chi_bf - - class LevelNumberDensityLTE(ProcessingPlasmaProperty): """ Attributes @@ -1064,109 +816,3 @@ def _calculate_factor(self, nu_i, t_electrons): def _calculate_u0s(nu, t_electrons): u0s = nu[np.newaxis].T / t_electrons * (H / K_B) return u0s - - -class CollRecombRateCoeff(ProcessingPlasmaProperty): - """ - Attributes - ---------- - coll_recomb_coeff : pandas.DataFrame, dtype float - The rate coefficient for collisional recombination. - Multiply with the electron density squared and the ion number density - to obtain the total rate. - - Notes - ----- - The collisional recombination rate coefficient is calculated from the - collisional ionization rate coefficient based on the requirement of detailed - balance. - """ - - outputs = ("coll_recomb_coeff",) - latex_name = (r"c_{\kappa\textrm{i,}}",) - - def calculate(self, phi_ik, coll_ion_coeff): - return coll_ion_coeff.multiply(phi_ik.loc[coll_ion_coeff.index]) - - -class RawCollIonTransProbs(TransitionProbabilitiesProperty, IndexSetterMixin): - """ - Attributes - ---------- - p_coll_ion : pandas.DataFrame, dtype float - The unnormalized transition probabilities for - collisional ionization. - p_coll_recomb : pandas.DataFrame, dtype float - The unnormalized transition probabilities for - collisional recombination. - cool_rate_coll_ion : pandas.DataFrame, dtype float - The collisional ionization cooling rates of the electron gas. - """ - - outputs = ("p_coll_ion", "p_coll_recomb", "cool_rate_coll_ion") - transition_probabilities_outputs = ( - "p_coll_ion", - "p_coll_recomb", - "cool_rate_coll_ion", - ) - latex_name = ( - r"p^{\textrm{coll ion}}", - r"p^{\textrm{coll recomb}}", - r"C^{\textrm{ion}}", - ) - - def calculate( - self, - coll_ion_coeff, - coll_recomb_coeff, - nu_i, - photo_ion_idx, - electron_densities, - energy_i, - level_number_density, - ): - p_coll_ion = coll_ion_coeff.multiply(energy_i, axis=0) - p_coll_ion = p_coll_ion.multiply(electron_densities, axis=1) - p_coll_ion = self.set_index(p_coll_ion, photo_ion_idx, reverse=False) - - coll_recomb_rate = coll_recomb_coeff.multiply( - electron_densities, axis=1 - ) # The full rate is obtained from this by multiplying by the - # electron density and ion number density. - p_recomb_deactivation = coll_recomb_rate.multiply(nu_i, axis=0) * H - p_recomb_deactivation = self.set_index( - p_recomb_deactivation, photo_ion_idx, transition_type=-1 - ) - p_recomb_deactivation = p_recomb_deactivation.groupby(level=[0]).sum() - index_dd = pd.MultiIndex.from_product( - [p_recomb_deactivation.index.values, ["k"], [0]], - names=list(photo_ion_idx.columns) + ["transition_type"], - ) - p_recomb_deactivation = p_recomb_deactivation.set_index(index_dd) - - p_recomb_internal = coll_recomb_rate.multiply(energy_i, axis=0) - p_recomb_internal = self.set_index( - p_recomb_internal, photo_ion_idx, transition_type=0 - ) - p_coll_recomb = pd.concat([p_recomb_deactivation, p_recomb_internal]) - - cool_rate_coll_ion = (coll_ion_coeff * electron_densities).multiply( - nu_i * H, axis=0 - ) - level_lower_index = coll_ion_coeff.index - cool_rate_coll_ion = ( - cool_rate_coll_ion - * level_number_density.loc[level_lower_index].values - ) - cool_rate_coll_ion = self.set_index( - cool_rate_coll_ion, photo_ion_idx, reverse=False - ) - cool_rate_coll_ion = cool_rate_coll_ion.groupby( - level="destination_level_idx" - ).sum() - ion_cool_index = pd.MultiIndex.from_product( - [["k"], cool_rate_coll_ion.index.values, [0]], - names=list(photo_ion_idx.columns) + ["transition_type"], - ) - cool_rate_coll_ion = cool_rate_coll_ion.set_index(ion_cool_index) - return p_coll_ion, p_coll_recomb, cool_rate_coll_ion diff --git a/tardis/plasma/properties/continuum_processes/stim_recomb_rate_coeff.py b/tardis/plasma/properties/continuum_processes/stim_recomb_rate_coeff.py new file mode 100644 index 00000000000..9e376442279 --- /dev/null +++ b/tardis/plasma/properties/continuum_processes/stim_recomb_rate_coeff.py @@ -0,0 +1,85 @@ +import numpy as np +import pandas as pd +from tardis.plasma.properties.base import Input +from tardis.plasma.properties.continuum_processes.rates import H +from tardis.transport.montecarlo.estimators.util import ( + bound_free_estimator_array2frame, + integrate_array_by_blocks, + ProcessingPlasmaProperty, +) + + +class StimRecombRateCoeffEstimator(Input): + """ + Attributes + ---------- + alpha_stim_estimator : pandas.DataFrame, dtype float + Unnormalized MC estimator for the rate coefficient for stimulated + recombination. + """ + + outputs = ("alpha_stim_estimator",) + latex_name = (r"\alpha^{\textrm{stim}}_\textrm{estim}",) + + +class StimRecombRateCoeff(ProcessingPlasmaProperty): + """ + Attributes + ---------- + alpha_stim : pandas.DataFrame, dtype float + The rate coefficient for stimulated recombination. + """ + + outputs = ("alpha_stim",) + latex_name = (r"\alpha^{\textrm{stim}}",) + + def calculate( + self, + photo_ion_cross_sections, + alpha_stim_estimator, + photo_ion_norm_factor, + photo_ion_block_references, + photo_ion_index, + dilute_planckian_radiation_field, + phi_ik, + t_electrons, + boltzmann_factor_photo_ion, + level2continuum_idx, + ): + # Used for initialization + if alpha_stim_estimator is None: + alpha_stim = self.calculate_from_dilute_bb( + photo_ion_cross_sections, + photo_ion_block_references, + photo_ion_index, + dilute_planckian_radiation_field, + t_electrons, + boltzmann_factor_photo_ion, + ) + else: + alpha_stim_estimator = bound_free_estimator_array2frame( + alpha_stim_estimator, level2continuum_idx + ) + alpha_stim = alpha_stim_estimator * photo_ion_norm_factor + alpha_stim *= phi_ik.loc[alpha_stim.index] + return alpha_stim + + @staticmethod + def calculate_from_dilute_bb( + photo_ion_cross_sections, + photo_ion_block_references, + photo_ion_index, + dilute_planckian_radiation_field, + t_electrons, + boltzmann_factor_photo_ion, + ): + nu = photo_ion_cross_sections["nu"] + x_sect = photo_ion_cross_sections["x_sect"] + j_nus = dilute_planckian_radiation_field.calculate_mean_intensity(nu) + j_nus *= boltzmann_factor_photo_ion + alpha_stim = j_nus.multiply(4.0 * np.pi * x_sect / nu / H, axis=0) + alpha_stim = integrate_array_by_blocks( + alpha_stim.values, nu.values, photo_ion_block_references + ) + alpha_stim = pd.DataFrame(alpha_stim, index=photo_ion_index) + return alpha_stim diff --git a/tardis/plasma/properties/ion_population.py b/tardis/plasma/properties/ion_population.py index 794ecde510d..1c42028fbfd 100644 --- a/tardis/plasma/properties/ion_population.py +++ b/tardis/plasma/properties/ion_population.py @@ -8,7 +8,9 @@ from tardis.plasma.exceptions import PlasmaIonizationError from tardis.plasma.properties.base import ProcessingPlasmaProperty -from tardis.plasma.properties.continuum_processes import get_ion_multi_index +from tardis.plasma.properties.continuum_proc.continuum_processes import ( + get_ion_multi_index, +) logger = logging.getLogger(__name__) diff --git a/tardis/plasma/properties/util/__init__.py b/tardis/plasma/properties/util/__init__.py deleted file mode 100644 index 9aec5b1c60a..00000000000 --- a/tardis/plasma/properties/util/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -""" -Calculating transition probabilities for macro_atom interactions. -""" diff --git a/tardis/plasma/properties/util/integrate_array_by_blocks.py b/tardis/plasma/properties/util/integrate_array_by_blocks.py deleted file mode 100644 index 3573b0f6447..00000000000 --- a/tardis/plasma/properties/util/integrate_array_by_blocks.py +++ /dev/null @@ -1,35 +0,0 @@ -import numpy as np -from numba import njit, prange - -from tardis.transport.montecarlo import njit_dict - - -@njit(**njit_dict) -def integrate_array_by_blocks(f, x, block_references): - """ - Integrate a function over blocks. - - This function integrates a function `f` defined at locations `x` - over blocks given in `block_references`. - - Parameters - ---------- - f : numpy.ndarray, dtype float - 2D input array to integrate. - x : numpy.ndarray, dtype float - 1D array with the sample points corresponding to the `f` values. - block_references : numpy.ndarray, dtype int - 1D array with the start indices of the blocks to be integrated. - - Returns - ------- - numpy.ndarray, dtype float - 2D array with integrated values. - """ - integrated = np.zeros((len(block_references) - 1, f.shape[1])) - for i in prange(f.shape[1]): # columns - for j in prange(len(integrated)): # rows - start = block_references[j] - stop = block_references[j + 1] - integrated[j, i] = np.trapz(f[start:stop, i], x[start:stop]) - return integrated diff --git a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py index 4112723dade..448575dd483 100644 --- a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py +++ b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py @@ -6,7 +6,9 @@ import tardis.constants as const from tardis.io.atom_data import AtomData -from tardis.plasma.properties.continuum_processes import PhotoIonBoltzmannFactor +from tardis.plasma.properties.continuum_proc.continuum_processes import ( + PhotoIonBoltzmannFactor, +) from tardis.plasma.radiation_field.planck_rad_field import ( DilutePlanckianRadiationField, ) @@ -139,20 +141,23 @@ def calculate_photo_ionization_rate_coefficient( ----- Equation 16 in Lucy 2003. """ - gamma = mean_intensity_photo_ion_df.multiply( + photo_ion_rate_coefficient = mean_intensity_photo_ion_df.multiply( 4.0 * np.pi * self.atom_data.photoionization_data.x_sect / (self.atom_data.photoionization_data.nu * H), axis=0, ) - gamma = integrate_array_by_blocks( - gamma.values, + photo_ion_rate_coefficient = integrate_array_by_blocks( + photo_ion_rate_coefficient.values, self.atom_data.photoionization_data.nu.values, self.atom_data.photo_ion_block_references, ) - gamma = pd.DataFrame(gamma, index=self.atom_data.photo_ion_unique_index) - return gamma + photo_ion_rate_coefficient = pd.DataFrame( + photo_ion_rate_coefficient, + index=self.atom_data.photo_ion_unique_index, + ) + return photo_ion_rate_coefficient def calculate_stimulated_recomb_rate_factor( self, @@ -221,5 +226,5 @@ def calculate_mean_intensity_photo_ion_table( @dataclass class ContinuumProperties: - stimulated_recomb_rate_factor: pd.DataFrame + stimulated_recombination_rate_coefficient: pd.DataFrame photo_ionization_rate_coefficient: pd.DataFrame diff --git a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py index b3e322c2aa2..aa50fe96990 100644 --- a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py +++ b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py @@ -43,9 +43,9 @@ def test_continuum_estimators( continuum_simulation.plasma.gamma, ) stimulated_recomb_rate_coeff = ( - continuum_properties_dilute_bb.stimulated_recomb_rate_factor + continuum_properties_dilute_bb.stimulated_recombination_rate_coefficient * continuum_plasma.phi_ik.loc[ - continuum_properties_dilute_bb.stimulated_recomb_rate_factor.index + continuum_properties_dilute_bb.stimulated_recombination_rate_coefficient.index ] ) pdt.assert_frame_equal( @@ -76,9 +76,9 @@ def test_continuum_estimators( continuum_simulation.plasma.gamma, ) stimulated_recomb_rate_coeff = ( - continuum_properties_mc.stimulated_recomb_rate_factor + continuum_properties_mc.stimulated_recombination_rate_coefficient * continuum_plasma.phi_ik.loc[ - continuum_properties_dilute_bb.stimulated_recomb_rate_factor.index + continuum_properties_dilute_bb.stimulated_recombination_rate_coefficient.index ] ) pdt.assert_frame_equal( From 98f83401a54c98e51ad060e2ce926e77f17f5e88 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Wed, 24 Jul 2024 18:09:22 -0400 Subject: [PATCH 53/89] Refactor opacities module structure --- benchmarks/transport_montecarlo_opacities.py | 3 +- tardis/energy_input/gamma_packet_loop.py | 4 +- tardis/energy_input/gamma_ray_estimators.py | 4 +- .../collisional_ion_trans_prob.py | 8 ++-- .../macro_atom/transition_probabilities.py | 2 +- .../properties/continuum_processes/rates.py | 33 ----------------- ...omb_rate_coeff.py => recomb_rate_coeff.py} | 37 ++++++++++++++++++- .../montecarlo/tests/test_opacities.py | 4 +- 8 files changed, 51 insertions(+), 44 deletions(-) rename tardis/plasma/properties/continuum_processes/{stim_recomb_rate_coeff.py => recomb_rate_coeff.py} (69%) diff --git a/benchmarks/transport_montecarlo_opacities.py b/benchmarks/transport_montecarlo_opacities.py index 52ff45ccd41..e536e5c49a5 100644 --- a/benchmarks/transport_montecarlo_opacities.py +++ b/benchmarks/transport_montecarlo_opacities.py @@ -4,6 +4,7 @@ from asv_runner.benchmarks.mark import parameterize +import tardis.opacities.compton_opacity_calculation import tardis.opacities.opacities as calculate_opacity from benchmarks.benchmark_base import BenchmarkBase @@ -28,7 +29,7 @@ class BenchmarkMontecarloMontecarloNumbaOpacities(BenchmarkBase): } ) def time_compton_opacity_calculation(self, electron_number_density, energy): - calculate_opacity.compton_opacity_calculation( + tardis.opacities.compton_opacity_calculation.compton_opacity_calculation( energy, electron_number_density ) diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py index 385febda679..faac2164eef 100644 --- a/tardis/energy_input/gamma_packet_loop.py +++ b/tardis/energy_input/gamma_packet_loop.py @@ -19,9 +19,11 @@ doppler_factor_3d, get_index, ) +from tardis.opacities.compton_opacity_calculation import ( + compton_opacity_calculation, +) from tardis.opacities.opacities import ( SIGMA_T, - compton_opacity_calculation, kappa_calculation, pair_creation_opacity_artis, pair_creation_opacity_calculation, diff --git a/tardis/energy_input/gamma_ray_estimators.py b/tardis/energy_input/gamma_ray_estimators.py index e28ba5029ff..13ebc9e84c1 100644 --- a/tardis/energy_input/gamma_ray_estimators.py +++ b/tardis/energy_input/gamma_ray_estimators.py @@ -7,9 +7,11 @@ angle_aberration_gamma, doppler_factor_3d, ) +from tardis.opacities.compton_opacity_calculation import ( + compton_opacity_calculation, +) from tardis.opacities.opacities import ( SIGMA_T, - compton_opacity_calculation, kappa_calculation, photoabsorption_opacity_calculation, ) diff --git a/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py b/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py index 9747faa988b..944fc16c172 100644 --- a/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py +++ b/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py @@ -1,15 +1,15 @@ +import pandas as pd + +from tardis import constants as const from tardis.plasma.properties.base import ( ProcessingPlasmaProperty, TransitionProbabilitiesProperty, ) from tardis.plasma.properties.continuum_processes.rates import ( - H, IndexSetterMixin, ) - -import pandas as pd - +H = const.h.cgs.value class RawCollIonTransProbs(TransitionProbabilitiesProperty, IndexSetterMixin): """ diff --git a/tardis/opacities/macro_atom/transition_probabilities.py b/tardis/opacities/macro_atom/transition_probabilities.py index c22d337a6fb..0ae783d15a9 100644 --- a/tardis/opacities/macro_atom/transition_probabilities.py +++ b/tardis/opacities/macro_atom/transition_probabilities.py @@ -5,7 +5,7 @@ from scipy import sparse as sp from tardis.plasma.properties.base import ProcessingPlasmaProperty -from tardis.plasma.properties.continuum_proc.continuum_processes import ( +from tardis.plasma.properties.continuum_processes.rates import ( get_ground_state_multi_index, ) from tardis.transport.montecarlo.macro_atom import ( diff --git a/tardis/plasma/properties/continuum_processes/rates.py b/tardis/plasma/properties/continuum_processes/rates.py index 9c9c6ef96ba..42610b46e52 100644 --- a/tardis/plasma/properties/continuum_processes/rates.py +++ b/tardis/plasma/properties/continuum_processes/rates.py @@ -161,39 +161,6 @@ def set_index(p, photo_ion_idx, transition_type=0, reverse=True): return p -class SpontRecombRateCoeff(ProcessingPlasmaProperty): - """ - Attributes - ---------- - alpha_sp : pandas.DataFrame, dtype float - The rate coefficient for spontaneous recombination. - """ - - outputs = ("alpha_sp",) - latex_name = (r"\alpha^{\textrm{sp}}",) - - def calculate( - self, - photo_ion_cross_sections, - t_electrons, - photo_ion_block_references, - photo_ion_index, - phi_ik, - boltzmann_factor_photo_ion, - ): - x_sect = photo_ion_cross_sections["x_sect"].values - nu = photo_ion_cross_sections["nu"].values - - alpha_sp = 8 * np.pi * x_sect * nu**2 / C**2 - alpha_sp = alpha_sp[:, np.newaxis] - alpha_sp = alpha_sp * boltzmann_factor_photo_ion - alpha_sp = integrate_array_by_blocks( - alpha_sp, nu, photo_ion_block_references - ) - alpha_sp = pd.DataFrame(alpha_sp, index=photo_ion_index) - return alpha_sp * phi_ik.loc[alpha_sp.index] - - class SpontRecombCoolingRateCoeff(ProcessingPlasmaProperty): """ Attributes diff --git a/tardis/plasma/properties/continuum_processes/stim_recomb_rate_coeff.py b/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py similarity index 69% rename from tardis/plasma/properties/continuum_processes/stim_recomb_rate_coeff.py rename to tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py index 9e376442279..1fe4132f709 100644 --- a/tardis/plasma/properties/continuum_processes/stim_recomb_rate_coeff.py +++ b/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py @@ -1,7 +1,7 @@ import numpy as np import pandas as pd -from tardis.plasma.properties.base import Input -from tardis.plasma.properties.continuum_processes.rates import H +from tardis.plasma.properties.base import Input, ProcessingPlasmaProperty +from tardis.plasma.properties.continuum_processes.rates import C, H from tardis.transport.montecarlo.estimators.util import ( bound_free_estimator_array2frame, integrate_array_by_blocks, @@ -83,3 +83,36 @@ def calculate_from_dilute_bb( ) alpha_stim = pd.DataFrame(alpha_stim, index=photo_ion_index) return alpha_stim + + +class SpontRecombRateCoeff(ProcessingPlasmaProperty): + """ + Attributes + ---------- + alpha_sp : pandas.DataFrame, dtype float + The rate coefficient for spontaneous recombination. + """ + + outputs = ("alpha_sp",) + latex_name = (r"\alpha^{\textrm{sp}}",) + + def calculate( + self, + photo_ion_cross_sections, + t_electrons, + photo_ion_block_references, + photo_ion_index, + phi_ik, + boltzmann_factor_photo_ion, + ): + cross_section = photo_ion_cross_sections["x_sect"].values + nu = photo_ion_cross_sections["nu"].values + + alpha_sp = 8 * np.pi * cross_section * nu**2 / C**2 + alpha_sp = alpha_sp[:, np.newaxis] + alpha_sp = alpha_sp * boltzmann_factor_photo_ion + alpha_sp = integrate_array_by_blocks( + alpha_sp, nu, photo_ion_block_references + ) + alpha_sp = pd.DataFrame(alpha_sp, index=photo_ion_index) + return alpha_sp * phi_ik.loc[alpha_sp.index] diff --git a/tardis/transport/montecarlo/tests/test_opacities.py b/tardis/transport/montecarlo/tests/test_opacities.py index 07096d6e0c6..c5461112aed 100644 --- a/tardis/transport/montecarlo/tests/test_opacities.py +++ b/tardis/transport/montecarlo/tests/test_opacities.py @@ -1,8 +1,10 @@ import numpy.testing as npt import pytest -from tardis.opacities.opacities import ( +from tardis.opacities.compton_opacity_calculation import ( compton_opacity_calculation, +) +from tardis.opacities.opacities import ( kappa_calculation, pair_creation_opacity_calculation, photoabsorption_opacity_calculation, From 84671ca3b54bc743d7a53504de9851f38aeefe1f Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Thu, 25 Jul 2024 11:00:18 -0400 Subject: [PATCH 54/89] cleanup from restructure --- tardis/io/atom_data/base.py | 2 +- tardis/plasma/properties/__init__.py | 2 +- tardis/plasma/properties/atomic.py | 2 +- .../photo_ion_rate_coeff.py | 67 +------------------ tardis/plasma/properties/ion_population.py | 2 +- .../plasma/properties/property_collections.py | 8 +-- tardis/simulation/base.py | 33 ++++++--- .../continuum_radfield_properties.py | 2 +- 8 files changed, 30 insertions(+), 88 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 7011f995841..ad870ae60fa 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -9,7 +9,7 @@ from tardis import constants as const from tardis.io.atom_data.util import resolve_atom_data_fname -from tardis.plasma.properties.continuum_proc.continuum_processes import ( +from tardis.plasma.properties.continuum_processes.rates import ( get_ground_state_multi_index, ) diff --git a/tardis/plasma/properties/__init__.py b/tardis/plasma/properties/__init__.py index d15c03afc74..671d19065f6 100644 --- a/tardis/plasma/properties/__init__.py +++ b/tardis/plasma/properties/__init__.py @@ -7,7 +7,7 @@ from tardis.opacities.macro_atom.transition_probabilities import * from tardis.plasma.properties.atomic import * -from tardis.plasma.properties.continuum_proc.continuum_processes import * +from tardis.plasma.properties.continuum_processes.rates import * from tardis.plasma.properties.general import * from tardis.plasma.properties.helium_nlte import * from tardis.plasma.properties.ion_population import * diff --git a/tardis/plasma/properties/atomic.py b/tardis/plasma/properties/atomic.py index 5dcc206c784..7c212c1933e 100644 --- a/tardis/plasma/properties/atomic.py +++ b/tardis/plasma/properties/atomic.py @@ -15,7 +15,7 @@ HiddenPlasmaProperty, ProcessingPlasmaProperty, ) -from tardis.plasma.properties.continuum_proc.continuum_processes import ( +from tardis.plasma.properties.continuum_processes.rates import ( A0, BETA_COLL, K_B, diff --git a/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py b/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py index 1f4c719cb39..c3bf1fff744 100644 --- a/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py +++ b/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py @@ -1,16 +1,8 @@ -import numpy as np -import pandas as pd - from tardis.plasma.properties.base import Input from tardis.plasma.properties.continuum_processes.rates import H -from tardis.transport.montecarlo.estimators.util import ( - ProcessingPlasmaProperty, - bound_free_estimator_array2frame, - integrate_array_by_blocks, -) -class PhotoIonRateCoeffEstimator(Input): +class PhotoIonRateCoeff(Input): """ Attributes ---------- @@ -19,62 +11,5 @@ class PhotoIonRateCoeffEstimator(Input): ionization. """ - outputs = ("gamma_estimator",) - latex_name = (r"\gamma_\textrm{estim}",) - - -class PhotoIonRateCoeff(ProcessingPlasmaProperty): - """ - Attributes - ---------- - gamma : pandas.DataFrame, dtype float - The rate coefficient for radiative ionization. - """ - outputs = ("gamma",) latex_name = (r"\gamma",) - - def calculate( - self, - photo_ion_cross_sections, - gamma_estimator, - photo_ion_norm_factor, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - level2continuum_idx, - ): - # Used for initialization - if gamma_estimator is None: - gamma = self.calculate_from_dilute_bb( - photo_ion_cross_sections, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - ) - else: - gamma_estimator = bound_free_estimator_array2frame( - gamma_estimator, level2continuum_idx - ) - gamma = gamma_estimator * photo_ion_norm_factor.value - - return gamma - - @staticmethod - def calculate_from_dilute_bb( - photo_ion_cross_sections, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - ): - nu = photo_ion_cross_sections["nu"] - x_sect = photo_ion_cross_sections["x_sect"] - j_nus = dilute_planckian_radiation_field.calculate_mean_intensity( - nu, - ) - gamma = j_nus.multiply(4.0 * np.pi * x_sect / nu / H, axis=0) - gamma = integrate_array_by_blocks( - gamma.values, nu.values, photo_ion_block_references - ) - gamma = pd.DataFrame(gamma, index=photo_ion_index) - return gamma diff --git a/tardis/plasma/properties/ion_population.py b/tardis/plasma/properties/ion_population.py index 1c42028fbfd..10f440b95c7 100644 --- a/tardis/plasma/properties/ion_population.py +++ b/tardis/plasma/properties/ion_population.py @@ -8,7 +8,7 @@ from tardis.plasma.exceptions import PlasmaIonizationError from tardis.plasma.properties.base import ProcessingPlasmaProperty -from tardis.plasma.properties.continuum_proc.continuum_processes import ( +from tardis.plasma.properties.continuum_processes.rates import ( get_ion_multi_index, ) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index 119423568d7..92fee62af76 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -2,8 +2,8 @@ NonMarkovChainTransitionProbabilities, TransitionProbabilities, ) -from tardis.plasma.properties import * from tardis.opacities.tau_sobolev import TauSobolev +from tardis.plasma.properties import * class PlasmaPropertyCollection(list): @@ -88,9 +88,6 @@ class PlasmaPropertyCollection(list): ) continuum_interaction_inputs = PlasmaPropertyCollection( [ - StimRecombRateCoeffEstimator, - PhotoIonRateCoeffEstimator, - Volume, BfHeatingRateCoeffEstimator, StimRecombCoolingRateCoeffEstimator, YgData, @@ -99,15 +96,12 @@ class PlasmaPropertyCollection(list): continuum_interaction_properties = PlasmaPropertyCollection( [ PhotoIonizationData, - SpontRecombRateCoeff, - PhotoIonRateCoeff, ThermalLevelBoltzmannFactorLTE, ThermalLTEPartitionFunction, BetaElectron, ThermalGElectron, ThermalPhiSahaLTE, SahaFactor, - PhotoIonEstimatorsNormFactor, StimRecombRateCoeff, CorrPhotoIonRateCoeff, SpontRecombCoolingRateCoeff, diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 13e893d85d2..5adb4bec945 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -10,18 +10,21 @@ import tardis from tardis import constants as const from tardis.io.configuration.config_reader import ConfigurationError +from tardis.io.model.parse_atom_data import parse_atom_data +from tardis.io.model.parse_simulation_state import ( + parse_simulation_state, +) from tardis.io.util import HDFWriterMixin from tardis.plasma.radiation_field import DilutePlanckianRadiationField from tardis.plasma.standard_plasmas import assemble_plasma from tardis.simulation.convergence import ConvergenceSolver -from tardis.io.model.parse_simulation_state import ( - parse_simulation_state, -) -from tardis.io.model.parse_atom_data import parse_atom_data from tardis.spectrum.base import SpectrumSolver from tardis.spectrum.formal_integral import FormalIntegrator from tardis.transport.montecarlo.base import MonteCarloTransportSolver from tardis.transport.montecarlo.configuration import montecarlo_globals +from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( + MCContinuumPropertiesSolver, +) from tardis.util.base import is_notebook from tardis.visualization import ConvergencePlots @@ -406,14 +409,24 @@ def advance_state(self, emitted_luminosity): # A check to see if the plasma is set with JBluesDetailed, in which # case it needs some extra kwargs. - estimators = self.transport.transport_state.radfield_mc_estimators + radfield_mc_estimators = ( + self.transport.transport_state.radfield_mc_estimators + ) - if "gamma_estimator" in self.plasma.outputs_dict: + if "gamma" in self.plasma.outputs_dict: + continuum_property_solver = MCContinuumPropertiesSolver( + self.atom_data + ) + estimated_continuum_properties = continuum_property_solver.solve( + radfield_mc_estimators, + self.transport.transport_state.time_of_simulation, + self.transport.transport_state.geometry_state.volume, + ) update_properties.update( - gamma_estimator=estimators.photo_ion_estimator, - alpha_stim_estimator=estimators.stim_recomb_estimator, - bf_heating_coeff_estimator=estimators.bf_heating_estimator, - stim_recomb_cooling_coeff_estimator=estimators.stim_recomb_cooling_estimator, + gamma=estimated_continuum_properties.photo_ion_coeff, + alpha_stim_coeff=estimated_continuum_properties.stim_recomb_estimator, + bf_heating_coeff_estimator=radfield_mc_estimators.bf_heating_estimator, + stim_recomb_cooling_coeff_estimator=radfield_mc_estimators.stim_recomb_cooling_estimator, ) self.plasma.update(**update_properties) diff --git a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py index 448575dd483..c5b5a1177b5 100644 --- a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py +++ b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py @@ -6,7 +6,7 @@ import tardis.constants as const from tardis.io.atom_data import AtomData -from tardis.plasma.properties.continuum_proc.continuum_processes import ( +from tardis.plasma.properties.continuum_processes.rates import ( PhotoIonBoltzmannFactor, ) from tardis.plasma.radiation_field.planck_rad_field import ( From 7782d07de8e63a43a7ca4bfdcdc2dbbb2ae7cb7c Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Thu, 25 Jul 2024 16:18:37 -0400 Subject: [PATCH 55/89] cleanup --- .../collisional_ion_trans_prob.py | 24 ------------------- tardis/plasma/properties/__init__.py | 2 +- .../continuum_processes/__init__.py | 7 ++++++ .../properties/continuum_processes/rates.py | 23 ++++++++++++++++++ .../continuum_processes/recomb_rate_coeff.py | 8 +++---- .../plasma/properties/property_collections.py | 1 + tardis/simulation/base.py | 7 ++++-- 7 files changed, 41 insertions(+), 31 deletions(-) diff --git a/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py b/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py index 944fc16c172..c3bde599279 100644 --- a/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py +++ b/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py @@ -2,7 +2,6 @@ from tardis import constants as const from tardis.plasma.properties.base import ( - ProcessingPlasmaProperty, TransitionProbabilitiesProperty, ) from tardis.plasma.properties.continuum_processes.rates import ( @@ -92,26 +91,3 @@ def calculate( ) cool_rate_coll_ion = cool_rate_coll_ion.set_index(ion_cool_index) return p_coll_ion, p_coll_recomb, cool_rate_coll_ion - - -class CollRecombRateCoeff(ProcessingPlasmaProperty): - """ - Attributes - ---------- - coll_recomb_coeff : pandas.DataFrame, dtype float - The rate coefficient for collisional recombination. - Multiply with the electron density squared and the ion number density - to obtain the total rate. - - Notes - ----- - The collisional recombination rate coefficient is calculated from the - collisional ionization rate coefficient based on the requirement of detailed - balance. - """ - - outputs = ("coll_recomb_coeff",) - latex_name = (r"c_{\kappa\textrm{i,}}",) - - def calculate(self, phi_ik, coll_ion_coeff): - return coll_ion_coeff.multiply(phi_ik.loc[coll_ion_coeff.index]) diff --git a/tardis/plasma/properties/__init__.py b/tardis/plasma/properties/__init__.py index 671d19065f6..3363cbd9bcd 100644 --- a/tardis/plasma/properties/__init__.py +++ b/tardis/plasma/properties/__init__.py @@ -7,7 +7,7 @@ from tardis.opacities.macro_atom.transition_probabilities import * from tardis.plasma.properties.atomic import * -from tardis.plasma.properties.continuum_processes.rates import * +from tardis.plasma.properties.continuum_processes import * from tardis.plasma.properties.general import * from tardis.plasma.properties.helium_nlte import * from tardis.plasma.properties.ion_population import * diff --git a/tardis/plasma/properties/continuum_processes/__init__.py b/tardis/plasma/properties/continuum_processes/__init__.py index e69de29bb2d..c94ed467ffd 100644 --- a/tardis/plasma/properties/continuum_processes/__init__.py +++ b/tardis/plasma/properties/continuum_processes/__init__.py @@ -0,0 +1,7 @@ +from tardis.plasma.properties.continuum_processes.photo_ion_rate_coeff import ( + PhotoIonRateCoeff, +) +from tardis.plasma.properties.continuum_processes.rates import * +from tardis.plasma.properties.continuum_processes.recomb_rate_coeff import ( + StimRecombRateCoeff, +) diff --git a/tardis/plasma/properties/continuum_processes/rates.py b/tardis/plasma/properties/continuum_processes/rates.py index 42610b46e52..b7f9a055bd3 100644 --- a/tardis/plasma/properties/continuum_processes/rates.py +++ b/tardis/plasma/properties/continuum_processes/rates.py @@ -783,3 +783,26 @@ def _calculate_factor(self, nu_i, t_electrons): def _calculate_u0s(nu, t_electrons): u0s = nu[np.newaxis].T / t_electrons * (H / K_B) return u0s + + +class CollRecombRateCoeff(ProcessingPlasmaProperty): + """ + Attributes + ---------- + coll_recomb_coeff : pandas.DataFrame, dtype float + The rate coefficient for collisional recombination. + Multiply with the electron density squared and the ion number density + to obtain the total rate. + + Notes + ----- + The collisional recombination rate coefficient is calculated from the + collisional ionization rate coefficient based on the requirement of detailed + balance. + """ + + outputs = ("coll_recomb_coeff",) + latex_name = (r"c_{\kappa\textrm{i,}}",) + + def calculate(self, phi_ik, coll_ion_coeff): + return coll_ion_coeff.multiply(phi_ik.loc[coll_ion_coeff.index]) diff --git a/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py b/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py index 1fe4132f709..35ab06246f6 100644 --- a/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py +++ b/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py @@ -1,15 +1,15 @@ import numpy as np import pandas as pd + from tardis.plasma.properties.base import Input, ProcessingPlasmaProperty from tardis.plasma.properties.continuum_processes.rates import C, H from tardis.transport.montecarlo.estimators.util import ( bound_free_estimator_array2frame, integrate_array_by_blocks, - ProcessingPlasmaProperty, ) -class StimRecombRateCoeffEstimator(Input): +class StimRecombRateCoeff(Input): """ Attributes ---------- @@ -18,11 +18,11 @@ class StimRecombRateCoeffEstimator(Input): recombination. """ - outputs = ("alpha_stim_estimator",) + outputs = ("alpha_stim",) latex_name = (r"\alpha^{\textrm{stim}}_\textrm{estim}",) -class StimRecombRateCoeff(ProcessingPlasmaProperty): +class StimRecombRateCoeffOLD(ProcessingPlasmaProperty): """ Attributes ---------- diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index f4c87f9f7d5..c30fe78d417 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -3,6 +3,7 @@ TransitionProbabilities, ) from tardis.opacities.tau_sobolev import TauSobolev +from tardis.opacities.continuum.bound_free import BoundFreeOpacity from tardis.plasma.properties import * diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index ca6f128b836..0334bfd9393 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -368,6 +368,11 @@ def advance_state(self, emitted_luminosity): # model.calculate_j_blues() equivalent # model.update_plasmas() equivalent # Bad test to see if this is a nlte run + + if "nlte_data" in self.plasma.outputs_dict: + self.plasma.store_previous_properties() + + # JBlues solver if ( self.plasma.plasma_solver_settings.RADIATIVE_RATES_TYPE == "blackbody" @@ -404,8 +409,6 @@ def advance_state(self, emitted_luminosity): f"radiative_rates_type type unknown - {self.plasma.plasma_solver_settings.RADIATIVE_RATES_TYPE}" ) - self.plasma.store_previous_properties() - # A check to see if the plasma is set with JBluesDetailed, in which # case it needs some extra kwargs. From cfc5078e0e0ad5eb7930069d639ccd964f1da29e Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Thu, 25 Jul 2024 17:01:30 -0400 Subject: [PATCH 56/89] clean up --- tardis/plasma/properties/property_collections.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index c30fe78d417..61e0d0881f3 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -4,6 +4,9 @@ ) from tardis.opacities.tau_sobolev import TauSobolev from tardis.opacities.continuum.bound_free import BoundFreeOpacity +from tardis.opacities.macro_atom.continuum_processes.collisional_ion_trans_prob import ( + RawCollIonTransProbs, +) from tardis.plasma.properties import * From ee864bbeee9a2aa2a84c882e75cff9116a1262f5 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Thu, 25 Jul 2024 17:35:24 -0400 Subject: [PATCH 57/89] more cleanup --- .../continuum_processes/__init__.py | 1 + .../properties/continuum_processes/rates.py | 11 +------- tardis/plasma/properties/plasma_input.py | 5 ---- .../plasma/properties/property_collections.py | 6 +++-- tardis/plasma/standard_plasmas.py | 26 ++++++++++++++----- .../montecarlo/tests/test_opacities.py | 4 +-- 6 files changed, 27 insertions(+), 26 deletions(-) diff --git a/tardis/plasma/properties/continuum_processes/__init__.py b/tardis/plasma/properties/continuum_processes/__init__.py index c94ed467ffd..29d4d5503ef 100644 --- a/tardis/plasma/properties/continuum_processes/__init__.py +++ b/tardis/plasma/properties/continuum_processes/__init__.py @@ -4,4 +4,5 @@ from tardis.plasma.properties.continuum_processes.rates import * from tardis.plasma.properties.continuum_processes.recomb_rate_coeff import ( StimRecombRateCoeff, + SpontRecombRateCoeff, ) diff --git a/tardis/plasma/properties/continuum_processes/rates.py b/tardis/plasma/properties/continuum_processes/rates.py index b7f9a055bd3..4dd697cfd64 100644 --- a/tardis/plasma/properties/continuum_processes/rates.py +++ b/tardis/plasma/properties/continuum_processes/rates.py @@ -28,6 +28,7 @@ "RawPhotoIonTransProbs", "CollDeexcRateCoeff", "CollExcRateCoeff", + "CollRecombRateCoeff", "RawCollisionTransProbs", "AdiabaticCoolingRate", "FreeFreeCoolingRate", @@ -315,16 +316,6 @@ def calculate( ) return gamma_corr - -class PhotoIonEstimatorsNormFactor(ProcessingPlasmaProperty): - outputs = ("photo_ion_norm_factor",) - latex_name = (r"\frac{1}{t_\textrm{simulation volume h}}",) - - @staticmethod - def calculate(time_simulation, volume): - return (time_simulation * volume * H) ** -1 - - class StimRecombCoolingRateCoeffEstimator(Input): """ Attributes diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index 76aa5f48a01..aac02e58308 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -15,7 +15,6 @@ "JBlues", "LinkTRadTElectron", "HeliumTreatment", - "Volume", "ContinuumInteractionSpecies", "NLTEIonizationSpecies", "NLTEExcitationSpecies", @@ -128,10 +127,6 @@ class LinkTRadTElectron(Input): class HeliumTreatment(Input): outputs = ("helium_treatment",) -class Volume(Input): - outputs = ("volume",) - - class ContinuumInteractionSpecies(Input): """ Attributes diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index 61e0d0881f3..b1fa21c29ff 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -1,12 +1,12 @@ +from tardis.opacities.continuum.bound_free import BoundFreeOpacity from tardis.opacities.macro_atom.base import ( NonMarkovChainTransitionProbabilities, TransitionProbabilities, ) -from tardis.opacities.tau_sobolev import TauSobolev -from tardis.opacities.continuum.bound_free import BoundFreeOpacity from tardis.opacities.macro_atom.continuum_processes.collisional_ion_trans_prob import ( RawCollIonTransProbs, ) +from tardis.opacities.tau_sobolev import TauSobolev from tardis.plasma.properties import * @@ -93,6 +93,7 @@ class PlasmaPropertyCollection(list): ) continuum_interaction_inputs = PlasmaPropertyCollection( [ + PhotoIonRateCoeff, BfHeatingRateCoeffEstimator, StimRecombCoolingRateCoeffEstimator, YgData, @@ -101,6 +102,7 @@ class PlasmaPropertyCollection(list): continuum_interaction_properties = PlasmaPropertyCollection( [ PhotoIonizationData, + SpontRecombRateCoeff, ThermalLevelBoltzmannFactorLTE, ThermalLTEPartitionFunction, BetaElectron, diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index caca28b4b70..1846beeea14 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -2,7 +2,6 @@ import numpy as np import pandas as pd -from astropy import units as u from tardis.plasma import BasePlasma from tardis.plasma.base import PlasmaSolverSettings @@ -44,6 +43,9 @@ ) from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper from tardis.plasma.radiation_field import DilutePlanckianRadiationField +from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( + DiluteBlackBodyContinuumPropertiesSolver, +) from tardis.util.base import species_string_to_tuple logger = logging.getLogger(__name__) @@ -129,6 +131,9 @@ def assemble_plasma(config, simulation_state, atom_data=None): plasma_modules = basic_inputs + basic_properties property_kwargs = {} + + ########### SETTING UP CONTINUUM INTERACTIONS + if len(config.plasma.continuum_interaction.species) > 0: line_interaction_type = config.plasma.line_interaction_type if line_interaction_type != "macroatom": @@ -200,14 +205,23 @@ def assemble_plasma(config, simulation_state, atom_data=None): f"NLTE solver type unknown - {config.plasma.nlte_solver}" ) + # initializing rates + t_electron = ( + config.plasma.link_t_rad_t_electron + * dilute_planckian_radiation_field.temperature.to(u.K).value + ) + initial_continuum_solver = DiluteBlackBodyContinuumPropertiesSolver( + atom_data + ) + initial_continuum_properties = initial_continuum_solver.solve( + dilute_planckian_radiation_field, t_electron + ) + kwargs.update( - gamma_estimator=None, + gamma=initial_continuum_properties.photo_ionization_rate_coefficient, bf_heating_coeff_estimator=None, stim_recomb_cooling_coeff_estimator=None, - alpha_stim_estimator=None, - volume=simulation_state.volume, - r_inner=simulation_state.r_inner.to(u.cm), - t_inner=simulation_state.t_inner, + alpha_stim=config.plasma.link_t_rad_t_electron, ) ##### RADIATIVE RATES SETUP diff --git a/tardis/transport/montecarlo/tests/test_opacities.py b/tardis/transport/montecarlo/tests/test_opacities.py index c5461112aed..5cda6b97f82 100644 --- a/tardis/transport/montecarlo/tests/test_opacities.py +++ b/tardis/transport/montecarlo/tests/test_opacities.py @@ -1,9 +1,7 @@ import numpy.testing as npt import pytest +from tardis.opacities.opacities import compton_opacity_calculation -from tardis.opacities.compton_opacity_calculation import ( - compton_opacity_calculation, -) from tardis.opacities.opacities import ( kappa_calculation, pair_creation_opacity_calculation, From 0ebcef8497e8b3105072bcd0239c7540cdc33628 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Thu, 25 Jul 2024 17:41:07 -0400 Subject: [PATCH 58/89] cleanup standard_plasmas.py --- tardis/plasma/standard_plasmas.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 1846beeea14..e78a54a481c 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -206,7 +206,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): ) # initializing rates - t_electron = ( + t_electrons = ( config.plasma.link_t_rad_t_electron * dilute_planckian_radiation_field.temperature.to(u.K).value ) @@ -214,14 +214,14 @@ def assemble_plasma(config, simulation_state, atom_data=None): atom_data ) initial_continuum_properties = initial_continuum_solver.solve( - dilute_planckian_radiation_field, t_electron + dilute_planckian_radiation_field, t_electrons ) kwargs.update( gamma=initial_continuum_properties.photo_ionization_rate_coefficient, bf_heating_coeff_estimator=None, stim_recomb_cooling_coeff_estimator=None, - alpha_stim=config.plasma.link_t_rad_t_electron, + alpha_stim=initial_continuum_properties.stimulated_recombination_rate_coefficient, ) ##### RADIATIVE RATES SETUP From cf1396a5fd044daa25bc6ed7e612182fee5390fd Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 26 Jul 2024 11:35:36 -0400 Subject: [PATCH 59/89] working nlte ionizations --- .../continuum_processes/__init__.py | 3 +- .../continuum_processes/recomb_rate_coeff.py | 55 ++----------------- .../plasma/properties/property_collections.py | 3 +- tardis/plasma/standard_plasmas.py | 3 +- .../continuum_radfield_properties.py | 3 +- .../tests/test_continuum_property_solver.py | 8 +-- 6 files changed, 17 insertions(+), 58 deletions(-) diff --git a/tardis/plasma/properties/continuum_processes/__init__.py b/tardis/plasma/properties/continuum_processes/__init__.py index 29d4d5503ef..572d3172c4e 100644 --- a/tardis/plasma/properties/continuum_processes/__init__.py +++ b/tardis/plasma/properties/continuum_processes/__init__.py @@ -3,6 +3,7 @@ ) from tardis.plasma.properties.continuum_processes.rates import * from tardis.plasma.properties.continuum_processes.recomb_rate_coeff import ( - StimRecombRateCoeff, + StimRecombRateFactor, SpontRecombRateCoeff, + StimRecombRateCoeff, ) diff --git a/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py b/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py index 35ab06246f6..a6ea2879d0b 100644 --- a/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py +++ b/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py @@ -9,7 +9,7 @@ ) -class StimRecombRateCoeff(Input): +class StimRecombRateFactor(Input): """ Attributes ---------- @@ -18,11 +18,11 @@ class StimRecombRateCoeff(Input): recombination. """ - outputs = ("alpha_stim",) + outputs = ("alpha_stim_factor",) latex_name = (r"\alpha^{\textrm{stim}}_\textrm{estim}",) -class StimRecombRateCoeffOLD(ProcessingPlasmaProperty): +class StimRecombRateCoeff(ProcessingPlasmaProperty): """ Attributes ---------- @@ -35,54 +35,10 @@ class StimRecombRateCoeffOLD(ProcessingPlasmaProperty): def calculate( self, - photo_ion_cross_sections, - alpha_stim_estimator, - photo_ion_norm_factor, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, + alpha_stim_factor, phi_ik, - t_electrons, - boltzmann_factor_photo_ion, - level2continuum_idx, - ): - # Used for initialization - if alpha_stim_estimator is None: - alpha_stim = self.calculate_from_dilute_bb( - photo_ion_cross_sections, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - t_electrons, - boltzmann_factor_photo_ion, - ) - else: - alpha_stim_estimator = bound_free_estimator_array2frame( - alpha_stim_estimator, level2continuum_idx - ) - alpha_stim = alpha_stim_estimator * photo_ion_norm_factor - alpha_stim *= phi_ik.loc[alpha_stim.index] - return alpha_stim - - @staticmethod - def calculate_from_dilute_bb( - photo_ion_cross_sections, - photo_ion_block_references, - photo_ion_index, - dilute_planckian_radiation_field, - t_electrons, - boltzmann_factor_photo_ion, ): - nu = photo_ion_cross_sections["nu"] - x_sect = photo_ion_cross_sections["x_sect"] - j_nus = dilute_planckian_radiation_field.calculate_mean_intensity(nu) - j_nus *= boltzmann_factor_photo_ion - alpha_stim = j_nus.multiply(4.0 * np.pi * x_sect / nu / H, axis=0) - alpha_stim = integrate_array_by_blocks( - alpha_stim.values, nu.values, photo_ion_block_references - ) - alpha_stim = pd.DataFrame(alpha_stim, index=photo_ion_index) - return alpha_stim + return alpha_stim_factor * phi_ik.loc[alpha_stim_factor.index] class SpontRecombRateCoeff(ProcessingPlasmaProperty): @@ -99,7 +55,6 @@ class SpontRecombRateCoeff(ProcessingPlasmaProperty): def calculate( self, photo_ion_cross_sections, - t_electrons, photo_ion_block_references, photo_ion_index, phi_ik, diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index b1fa21c29ff..4cf62169ddb 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -94,6 +94,7 @@ class PlasmaPropertyCollection(list): continuum_interaction_inputs = PlasmaPropertyCollection( [ PhotoIonRateCoeff, + StimRecombRateFactor, BfHeatingRateCoeffEstimator, StimRecombCoolingRateCoeffEstimator, YgData, @@ -101,6 +102,7 @@ class PlasmaPropertyCollection(list): ) continuum_interaction_properties = PlasmaPropertyCollection( [ + StimRecombRateCoeff, PhotoIonizationData, SpontRecombRateCoeff, ThermalLevelBoltzmannFactorLTE, @@ -109,7 +111,6 @@ class PlasmaPropertyCollection(list): ThermalGElectron, ThermalPhiSahaLTE, SahaFactor, - StimRecombRateCoeff, CorrPhotoIonRateCoeff, SpontRecombCoolingRateCoeff, RawRecombTransProbs, diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index e78a54a481c..7468233be6b 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -2,6 +2,7 @@ import numpy as np import pandas as pd +from astropy import units as u from tardis.plasma import BasePlasma from tardis.plasma.base import PlasmaSolverSettings @@ -221,7 +222,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): gamma=initial_continuum_properties.photo_ionization_rate_coefficient, bf_heating_coeff_estimator=None, stim_recomb_cooling_coeff_estimator=None, - alpha_stim=initial_continuum_properties.stimulated_recombination_rate_coefficient, + alpha_stim_factor=initial_continuum_properties.stimulated_recombination_rate_factor, ) ##### RADIATIVE RATES SETUP diff --git a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py index c5b5a1177b5..3f49a6629b1 100644 --- a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py +++ b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py @@ -226,5 +226,6 @@ def calculate_mean_intensity_photo_ion_table( @dataclass class ContinuumProperties: - stimulated_recombination_rate_coefficient: pd.DataFrame + # this is not the rate coefficient but misses Phi I_K + stimulated_recombination_rate_factor: pd.DataFrame photo_ionization_rate_coefficient: pd.DataFrame diff --git a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py index aa50fe96990..5fb75a0228c 100644 --- a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py +++ b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py @@ -43,9 +43,9 @@ def test_continuum_estimators( continuum_simulation.plasma.gamma, ) stimulated_recomb_rate_coeff = ( - continuum_properties_dilute_bb.stimulated_recombination_rate_coefficient + continuum_properties_dilute_bb.stimulated_recombination_rate_factor * continuum_plasma.phi_ik.loc[ - continuum_properties_dilute_bb.stimulated_recombination_rate_coefficient.index + continuum_properties_dilute_bb.stimulated_recombination_rate_factor.index ] ) pdt.assert_frame_equal( @@ -76,9 +76,9 @@ def test_continuum_estimators( continuum_simulation.plasma.gamma, ) stimulated_recomb_rate_coeff = ( - continuum_properties_mc.stimulated_recombination_rate_coefficient + continuum_properties_mc.stimulated_recombination_rate_factor * continuum_plasma.phi_ik.loc[ - continuum_properties_dilute_bb.stimulated_recombination_rate_coefficient.index + continuum_properties_dilute_bb.stimulated_recombination_rate_factor.index ] ) pdt.assert_frame_equal( From fcf2563ea6e8f40ac30ece6c7b01e0d4f3af6437 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 26 Jul 2024 11:51:00 -0400 Subject: [PATCH 60/89] start of assemble plasma cleanup --- .../configuration/tests/test_config_reader.py | 2 +- tardis/plasma/standard_plasmas.py | 368 ------------------ tardis/plasma/tests/test_nlte_solver.py | 2 +- .../tests/test_tardis_model_density_config.py | 2 +- tardis/simulation/base.py | 2 +- 5 files changed, 4 insertions(+), 372 deletions(-) delete mode 100644 tardis/plasma/standard_plasmas.py diff --git a/tardis/io/configuration/tests/test_config_reader.py b/tardis/io/configuration/tests/test_config_reader.py index f525557ca63..2fe68ea061d 100644 --- a/tardis/io/configuration/tests/test_config_reader.py +++ b/tardis/io/configuration/tests/test_config_reader.py @@ -10,7 +10,7 @@ from astropy.units import Quantity from tardis.io.configuration.config_reader import Configuration from tardis.plasma.exceptions import PlasmaConfigError -from tardis.plasma.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.standard_plasmas import assemble_plasma def test_convergence_section_parser(): diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py deleted file mode 100644 index 7468233be6b..00000000000 --- a/tardis/plasma/standard_plasmas.py +++ /dev/null @@ -1,368 +0,0 @@ -import logging - -import numpy as np -import pandas as pd -from astropy import units as u - -from tardis.plasma import BasePlasma -from tardis.plasma.base import PlasmaSolverSettings -from tardis.plasma.exceptions import PlasmaConfigError -from tardis.plasma.properties import ( - HeliumNumericalNLTE, - IonNumberDensity, - IonNumberDensityHeNLTE, - LevelBoltzmannFactorNLTE, - MarkovChainTransProbsCollector, - RadiationFieldCorrection, - StimulatedEmissionFactor, -) -from tardis.plasma.properties.base import TransitionProbabilitiesProperty -from tardis.plasma.properties.level_population import LevelNumberDensity -from tardis.plasma.properties.nlte_rate_equation_solver import ( - NLTEPopulationSolverLU, - NLTEPopulationSolverRoot, -) -from tardis.plasma.properties.property_collections import ( - adiabatic_cooling_properties, - basic_inputs, - basic_properties, - continuum_interaction_inputs, - continuum_interaction_properties, - dilute_lte_excitation_properties, - helium_lte_properties, - helium_nlte_properties, - helium_numerical_nlte_properties, - lte_excitation_properties, - lte_ionization_properties, - macro_atom_properties, - nebular_ionization_properties, - nlte_lu_solver_properties, - nlte_properties, - nlte_root_solver_properties, - non_nlte_properties, - two_photon_properties, -) -from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper -from tardis.plasma.radiation_field import DilutePlanckianRadiationField -from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( - DiluteBlackBodyContinuumPropertiesSolver, -) -from tardis.util.base import species_string_to_tuple - -logger = logging.getLogger(__name__) - - -def assemble_plasma(config, simulation_state, atom_data=None): - """ - Create a BasePlasma instance from a Configuration object - and a SimulationState. - - Parameters - ---------- - config : io.config_reader.Configuration - simulation_state : model.SimulationState - atom_data : atomic.AtomData - If None, an attempt will be made to read the atomic data - from config. - - Returns - ------- - : plasma.BasePlasma - - """ - # Convert the nlte species list to a proper format. - nlte_species = [ - species_string_to_tuple(species) - for species in config.plasma.nlte.species - ] - - # Convert the continuum interaction species list to a proper format. - continuum_interaction_species = [ - species_string_to_tuple(species) - for species in config.plasma.continuum_interaction.species - ] - continuum_interaction_species = pd.MultiIndex.from_tuples( - continuum_interaction_species, names=["atomic_number", "ion_number"] - ) - - atom_data.prepare_atom_data( - simulation_state.abundance.index, - line_interaction_type=config.plasma.line_interaction_type, - continuum_interaction_species=continuum_interaction_species, - nlte_species=nlte_species, - ) - - # Check if continuum interaction species are in selected_atoms - continuum_atoms = continuum_interaction_species.get_level_values( - "atomic_number" - ) - continuum_atoms_in_selected_atoms = np.all( - continuum_atoms.isin(atom_data.selected_atomic_numbers) - ) - if not continuum_atoms_in_selected_atoms: - raise PlasmaConfigError( - "Not all continuum interaction species " - "belong to atoms that have been specified " - "in the configuration." - ) - - nlte_ionization_species = [ - species_string_to_tuple(species) - for species in config.plasma.nlte_ionization_species - ] - nlte_excitation_species = [ - species_string_to_tuple(species) - for species in config.plasma.nlte_excitation_species - ] - - dilute_planckian_radiation_field = DilutePlanckianRadiationField( - simulation_state.t_radiative, simulation_state.dilution_factor - ) - kwargs = dict( - dilute_planckian_radiation_field=dilute_planckian_radiation_field, - abundance=simulation_state.abundance, - number_density=simulation_state.elemental_number_density, - atomic_data=atom_data, - time_explosion=simulation_state.time_explosion, - link_t_rad_t_electron=config.plasma.link_t_rad_t_electron, - continuum_interaction_species=continuum_interaction_species, - nlte_ionization_species=nlte_ionization_species, - nlte_excitation_species=nlte_excitation_species, - ) - - plasma_modules = basic_inputs + basic_properties - property_kwargs = {} - - ########### SETTING UP CONTINUUM INTERACTIONS - - if len(config.plasma.continuum_interaction.species) > 0: - line_interaction_type = config.plasma.line_interaction_type - if line_interaction_type != "macroatom": - raise PlasmaConfigError( - "Continuum interactions require line_interaction_type " - f"macroatom (instead of {line_interaction_type})." - ) - - plasma_modules += continuum_interaction_properties - plasma_modules += continuum_interaction_inputs - - if config.plasma.continuum_interaction.enable_adiabatic_cooling: - plasma_modules += adiabatic_cooling_properties - - if config.plasma.continuum_interaction.enable_two_photon_decay: - plasma_modules += two_photon_properties - - transition_probabilities_outputs = [ - plasma_property.transition_probabilities_outputs - for plasma_property in plasma_modules - if issubclass(plasma_property, TransitionProbabilitiesProperty) - ] - transition_probabilities_outputs = [ - item - for sublist in transition_probabilities_outputs - for item in sublist - ] - - property_kwargs[MarkovChainTransProbsCollector] = { - "inputs": transition_probabilities_outputs - } - if ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ): - if config.plasma.nlte_ionization_species: - nlte_ionization_species = config.plasma.nlte_ionization_species - for species in nlte_ionization_species: - if ( - species - not in config.plasma.continuum_interaction.species - ): - raise PlasmaConfigError( - f"NLTE ionization species {species} not in continuum species." - ) - if config.plasma.nlte_excitation_species: - nlte_excitation_species = config.plasma.nlte_excitation_species - for species in nlte_excitation_species: - if ( - species - not in config.plasma.continuum_interaction.species - ): - raise PlasmaConfigError( - f"NLTE excitation species {species} not in continuum species." - ) - property_kwargs[NLTEIndexHelper] = { - "nlte_ionization_species": config.plasma.nlte_ionization_species, - "nlte_excitation_species": config.plasma.nlte_excitation_species, - } - if config.plasma.nlte_solver == "lu": - plasma_modules += nlte_lu_solver_properties - logger.warning( - "LU solver will be inaccurate for NLTE excitation, proceed with caution." - ) - elif config.plasma.nlte_solver == "root": - plasma_modules += nlte_root_solver_properties - else: - raise PlasmaConfigError( - f"NLTE solver type unknown - {config.plasma.nlte_solver}" - ) - - # initializing rates - t_electrons = ( - config.plasma.link_t_rad_t_electron - * dilute_planckian_radiation_field.temperature.to(u.K).value - ) - initial_continuum_solver = DiluteBlackBodyContinuumPropertiesSolver( - atom_data - ) - initial_continuum_properties = initial_continuum_solver.solve( - dilute_planckian_radiation_field, t_electrons - ) - - kwargs.update( - gamma=initial_continuum_properties.photo_ionization_rate_coefficient, - bf_heating_coeff_estimator=None, - stim_recomb_cooling_coeff_estimator=None, - alpha_stim_factor=initial_continuum_properties.stimulated_recombination_rate_factor, - ) - - ##### RADIATIVE RATES SETUP - - plasma_solver_settings = PlasmaSolverSettings( - RADIATIVE_RATES_TYPE=config.plasma.radiative_rates_type - ) - - if (plasma_solver_settings.RADIATIVE_RATES_TYPE == "dilute-blackbody") or ( - plasma_solver_settings.RADIATIVE_RATES_TYPE == "detailed" - ): - kwargs["j_blues"] = pd.DataFrame( - dilute_planckian_radiation_field.calculate_mean_intensity( - atom_data.lines["nu"].values - ), - index=atom_data.lines.index, - ) - - elif plasma_solver_settings.RADIATIVE_RATES_TYPE == "blackbody": - planckian_rad_field = ( - dilute_planckian_radiation_field.to_planckian_radiation_field() - ) - kwargs["j_blues"] = pd.DataFrame( - planckian_rad_field.calculate_mean_intensity( - atom_data.lines["nu"].values - ), - index=atom_data.lines.index, - ) - - else: - raise ValueError( - f"radiative_rates_type type unknown - {plasma_solver_settings.RADIATIVE_RATES_TYPE}" - ) - - if config.plasma.excitation == "lte": - plasma_modules += lte_excitation_properties - elif config.plasma.excitation == "dilute-lte": - plasma_modules += dilute_lte_excitation_properties - - if config.plasma.ionization == "lte": - plasma_modules += lte_ionization_properties - elif config.plasma.ionization == "nebular": - plasma_modules += nebular_ionization_properties - - if nlte_species: - plasma_modules += nlte_properties - nlte_conf = config.plasma.nlte - plasma_modules.append(LevelBoltzmannFactorNLTE.from_config(nlte_conf)) - property_kwargs[StimulatedEmissionFactor] = dict( - nlte_species=nlte_species - ) - else: - plasma_modules += non_nlte_properties - - if config.plasma.line_interaction_type in ("downbranch", "macroatom"): - if not config.plasma.continuum_interaction.species: - plasma_modules += macro_atom_properties - - if "delta_treatment" in config.plasma: - property_kwargs[RadiationFieldCorrection] = dict( - delta_treatment=config.plasma.delta_treatment - ) - - if ( - config.plasma.helium_treatment == "recomb-nlte" - or config.plasma.helium_treatment == "numerical-nlte" - ) and ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ): - # Prevent the user from using helium NLTE treatment with - # NLTE ionization and excitation treatment. This is because - # the helium_nlte_properties could overwrite the NLTE ionization - # and excitation ion number and electron densities. - # helium_numerical_nlte_properties is also included here because - # it is currently in the same if else block, and thus may block - # the addition of the components from the else block. - raise PlasmaConfigError( - "Helium NLTE treatment is incompatible with the NLTE eonization and excitation treatment." - ) - - # TODO: Disentangle these if else block such that compatible components - # can be added independently. - if config.plasma.helium_treatment == "recomb-nlte": - plasma_modules += helium_nlte_properties - elif config.plasma.helium_treatment == "numerical-nlte": - plasma_modules += helium_numerical_nlte_properties - # TODO: See issue #633 - if config.plasma.heating_rate_data_file in ["none", None]: - raise PlasmaConfigError("Heating rate data file not specified") - else: - property_kwargs[HeliumNumericalNLTE] = dict( - heating_rate_data_file=config.plasma.heating_rate_data_file - ) - else: - # If nlte ionization species are present, we don't want to add the - # IonNumberDensity from helium_lte_properties, since we want - # to use the IonNumberDensity provided by the NLTE solver. - if ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ): - plasma_modules += [LevelNumberDensity] - else: - plasma_modules += helium_lte_properties - - if simulation_state._electron_densities is not None: - electron_densities = pd.Series( - simulation_state._electron_densities.cgs.value - ) - if config.plasma.helium_treatment == "numerical-nlte": - property_kwargs[IonNumberDensityHeNLTE] = dict( - electron_densities=electron_densities - ) - elif ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ) and config.plasma.nlte_solver == "root": - property_kwargs[NLTEPopulationSolverRoot] = dict( - electron_densities=electron_densities - ) - elif ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ) and config.plasma.nlte_solver == "lu": - property_kwargs[NLTEPopulationSolverLU] = dict( - electron_densities=electron_densities - ) - else: - property_kwargs[IonNumberDensity] = dict( - electron_densities=electron_densities - ) - - kwargs["helium_treatment"] = config.plasma.helium_treatment - - plasma = BasePlasma( - plasma_properties=plasma_modules, - property_kwargs=property_kwargs, - plasma_solver_settings=plasma_solver_settings, - **kwargs, - ) - - return plasma diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py index 48045d3cf09..35b9b572d3f 100644 --- a/tardis/plasma/tests/test_nlte_solver.py +++ b/tardis/plasma/tests/test_nlte_solver.py @@ -14,7 +14,7 @@ calculate_jacobian_matrix, calculate_rate_matrix, ) -from tardis.plasma.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.standard_plasmas import assemble_plasma @pytest.fixture diff --git a/tardis/plasma/tests/test_tardis_model_density_config.py b/tardis/plasma/tests/test_tardis_model_density_config.py index 20c9ec46cda..7b5be056ddc 100644 --- a/tardis/plasma/tests/test_tardis_model_density_config.py +++ b/tardis/plasma/tests/test_tardis_model_density_config.py @@ -4,7 +4,7 @@ from tardis.io.configuration.config_reader import Configuration from tardis.model import SimulationState -from tardis.plasma.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.standard_plasmas import assemble_plasma @pytest.fixture diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 0334bfd9393..b5be5b46b60 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -16,7 +16,7 @@ ) from tardis.io.util import HDFWriterMixin from tardis.plasma.radiation_field import DilutePlanckianRadiationField -from tardis.plasma.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.standard_plasmas import assemble_plasma from tardis.simulation.convergence import ConvergenceSolver from tardis.spectrum.base import SpectrumSolver from tardis.spectrum.formal_integral import FormalIntegrator From 83240744b37de1e2127eaa2d219a3d308bbada7e Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 26 Jul 2024 11:51:33 -0400 Subject: [PATCH 61/89] cleanup standard_plasmas.py --- tardis/plasma/assembly/__init__.py | 0 tardis/plasma/assembly/base.py | 379 +++++++++++++++++++++++++++++ 2 files changed, 379 insertions(+) create mode 100644 tardis/plasma/assembly/__init__.py create mode 100644 tardis/plasma/assembly/base.py diff --git a/tardis/plasma/assembly/__init__.py b/tardis/plasma/assembly/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py new file mode 100644 index 00000000000..2a3c6724568 --- /dev/null +++ b/tardis/plasma/assembly/base.py @@ -0,0 +1,379 @@ +import logging + +import numpy as np +import pandas as pd +from astropy import units as u + +from tardis.plasma import BasePlasma +from tardis.plasma.base import PlasmaSolverSettings +from tardis.plasma.exceptions import PlasmaConfigError +from tardis.plasma.properties import ( + HeliumNumericalNLTE, + IonNumberDensity, + IonNumberDensityHeNLTE, + LevelBoltzmannFactorNLTE, + MarkovChainTransProbsCollector, + RadiationFieldCorrection, + StimulatedEmissionFactor, +) +from tardis.plasma.properties.base import TransitionProbabilitiesProperty +from tardis.plasma.properties.level_population import LevelNumberDensity +from tardis.plasma.properties.nlte_rate_equation_solver import ( + NLTEPopulationSolverLU, + NLTEPopulationSolverRoot, +) +from tardis.plasma.properties.property_collections import ( + adiabatic_cooling_properties, + basic_inputs, + basic_properties, + continuum_interaction_inputs, + continuum_interaction_properties, + dilute_lte_excitation_properties, + helium_lte_properties, + helium_nlte_properties, + helium_numerical_nlte_properties, + lte_excitation_properties, + lte_ionization_properties, + macro_atom_properties, + nebular_ionization_properties, + nlte_lu_solver_properties, + nlte_properties, + nlte_root_solver_properties, + non_nlte_properties, + two_photon_properties, +) +from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper +from tardis.plasma.radiation_field import DilutePlanckianRadiationField +from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( + DiluteBlackBodyContinuumPropertiesSolver, +) +from tardis.util.base import species_string_to_tuple + +logger = logging.getLogger(__name__) + + +def assemble_plasma(config, simulation_state, atom_data=None): + """ + Create a BasePlasma instance from a Configuration object + and a SimulationState. + + Parameters + ---------- + config : io.config_reader.Configuration + simulation_state : model.SimulationState + atom_data : atomic.AtomData + If None, an attempt will be made to read the atomic data + from config. + + Returns + ------- + : plasma.BasePlasma + + """ + # Convert the nlte species list to a proper format. + nlte_species = [ + species_string_to_tuple(species) + for species in config.plasma.nlte.species + ] + + # Convert the continuum interaction species list to a proper format. + continuum_interaction_species = [ + species_string_to_tuple(species) + for species in config.plasma.continuum_interaction.species + ] + continuum_interaction_species = pd.MultiIndex.from_tuples( + continuum_interaction_species, names=["atomic_number", "ion_number"] + ) + + atom_data.prepare_atom_data( + simulation_state.abundance.index, + line_interaction_type=config.plasma.line_interaction_type, + continuum_interaction_species=continuum_interaction_species, + nlte_species=nlte_species, + ) + + # Check if continuum interaction species are in selected_atoms + continuum_atoms = continuum_interaction_species.get_level_values( + "atomic_number" + ) + + continuum_atoms_in_selected_atoms = np.all( + continuum_atoms.isin(atom_data.selected_atomic_numbers) + ) + if not continuum_atoms_in_selected_atoms: + raise PlasmaConfigError( + "Not all continuum interaction species " + "belong to atoms that have been specified " + "in the configuration." + ) + + nlte_ionization_species = [ + species_string_to_tuple(species) + for species in config.plasma.nlte_ionization_species + ] + nlte_excitation_species = [ + species_string_to_tuple(species) + for species in config.plasma.nlte_excitation_species + ] + + dilute_planckian_radiation_field = DilutePlanckianRadiationField( + simulation_state.t_radiative, simulation_state.dilution_factor + ) + kwargs = dict( + dilute_planckian_radiation_field=dilute_planckian_radiation_field, + abundance=simulation_state.abundance, + number_density=simulation_state.elemental_number_density, + atomic_data=atom_data, + time_explosion=simulation_state.time_explosion, + link_t_rad_t_electron=config.plasma.link_t_rad_t_electron, + continuum_interaction_species=continuum_interaction_species, + nlte_ionization_species=nlte_ionization_species, + nlte_excitation_species=nlte_excitation_species, + ) + + plasma_modules = basic_inputs + basic_properties + property_kwargs = {} + + ########### SETTING UP CONTINUUM INTERACTIONS + + if len(config.plasma.continuum_interaction.species) > 0: + setup_continuum_interactions( + config, + atom_data, + dilute_planckian_radiation_field, + kwargs, + property_kwargs, + ) + + ##### RADIATIVE RATES SETUP + + plasma_solver_settings = PlasmaSolverSettings( + RADIATIVE_RATES_TYPE=config.plasma.radiative_rates_type + ) + + if (plasma_solver_settings.RADIATIVE_RATES_TYPE == "dilute-blackbody") or ( + plasma_solver_settings.RADIATIVE_RATES_TYPE == "detailed" + ): + kwargs["j_blues"] = pd.DataFrame( + dilute_planckian_radiation_field.calculate_mean_intensity( + atom_data.lines["nu"].values + ), + index=atom_data.lines.index, + ) + + elif plasma_solver_settings.RADIATIVE_RATES_TYPE == "blackbody": + planckian_rad_field = ( + dilute_planckian_radiation_field.to_planckian_radiation_field() + ) + kwargs["j_blues"] = pd.DataFrame( + planckian_rad_field.calculate_mean_intensity( + atom_data.lines["nu"].values + ), + index=atom_data.lines.index, + ) + + else: + raise ValueError( + f"radiative_rates_type type unknown - {plasma_solver_settings.RADIATIVE_RATES_TYPE}" + ) + + if config.plasma.excitation == "lte": + plasma_modules += lte_excitation_properties + elif config.plasma.excitation == "dilute-lte": + plasma_modules += dilute_lte_excitation_properties + + if config.plasma.ionization == "lte": + plasma_modules += lte_ionization_properties + elif config.plasma.ionization == "nebular": + plasma_modules += nebular_ionization_properties + + if nlte_species: + plasma_modules += nlte_properties + nlte_conf = config.plasma.nlte + plasma_modules.append(LevelBoltzmannFactorNLTE.from_config(nlte_conf)) + property_kwargs[StimulatedEmissionFactor] = dict( + nlte_species=nlte_species + ) + else: + plasma_modules += non_nlte_properties + + if config.plasma.line_interaction_type in ("downbranch", "macroatom"): + if not config.plasma.continuum_interaction.species: + plasma_modules += macro_atom_properties + + if "delta_treatment" in config.plasma: + property_kwargs[RadiationFieldCorrection] = dict( + delta_treatment=config.plasma.delta_treatment + ) + + if ( + config.plasma.helium_treatment == "recomb-nlte" + or config.plasma.helium_treatment == "numerical-nlte" + ) and ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ): + # Prevent the user from using helium NLTE treatment with + # NLTE ionization and excitation treatment. This is because + # the helium_nlte_properties could overwrite the NLTE ionization + # and excitation ion number and electron densities. + # helium_numerical_nlte_properties is also included here because + # it is currently in the same if else block, and thus may block + # the addition of the components from the else block. + raise PlasmaConfigError( + "Helium NLTE treatment is incompatible with the NLTE eonization and excitation treatment." + ) + + # TODO: Disentangle these if else block such that compatible components + # can be added independently. + if config.plasma.helium_treatment == "recomb-nlte": + plasma_modules += helium_nlte_properties + elif config.plasma.helium_treatment == "numerical-nlte": + plasma_modules += helium_numerical_nlte_properties + # TODO: See issue #633 + if config.plasma.heating_rate_data_file in ["none", None]: + raise PlasmaConfigError("Heating rate data file not specified") + else: + property_kwargs[HeliumNumericalNLTE] = dict( + heating_rate_data_file=config.plasma.heating_rate_data_file + ) + else: + # If nlte ionization species are present, we don't want to add the + # IonNumberDensity from helium_lte_properties, since we want + # to use the IonNumberDensity provided by the NLTE solver. + if ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ): + plasma_modules += [LevelNumberDensity] + else: + plasma_modules += helium_lte_properties + + if simulation_state._electron_densities is not None: + electron_densities = pd.Series( + simulation_state._electron_densities.cgs.value + ) + if config.plasma.helium_treatment == "numerical-nlte": + property_kwargs[IonNumberDensityHeNLTE] = dict( + electron_densities=electron_densities + ) + elif ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ) and config.plasma.nlte_solver == "root": + property_kwargs[NLTEPopulationSolverRoot] = dict( + electron_densities=electron_densities + ) + elif ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ) and config.plasma.nlte_solver == "lu": + property_kwargs[NLTEPopulationSolverLU] = dict( + electron_densities=electron_densities + ) + else: + property_kwargs[IonNumberDensity] = dict( + electron_densities=electron_densities + ) + + kwargs["helium_treatment"] = config.plasma.helium_treatment + + plasma = BasePlasma( + plasma_properties=plasma_modules, + property_kwargs=property_kwargs, + plasma_solver_settings=plasma_solver_settings, + **kwargs, + ) + + return plasma + + +def setup_continuum_interactions( + config, + atom_data, + dilute_planckian_radiation_field, + kwargs, + property_kwargs, +): + continuum_plasma_modules = [] + line_interaction_type = config.plasma.line_interaction_type + if line_interaction_type != "macroatom": + raise PlasmaConfigError( + "Continuum interactions require line_interaction_type " + f"macroatom (instead of {line_interaction_type})." + ) + + continuum_plasma_modules += continuum_interaction_properties + continuum_plasma_modules += continuum_interaction_inputs + + if config.plasma.continuum_interaction.enable_adiabatic_cooling: + continuum_plasma_modules += adiabatic_cooling_properties + + if config.plasma.continuum_interaction.enable_two_photon_decay: + continuum_plasma_modules += two_photon_properties + + transition_probabilities_outputs = [ + plasma_property.transition_probabilities_outputs + for plasma_property in continuum_plasma_modules + if issubclass(plasma_property, TransitionProbabilitiesProperty) + ] + transition_probabilities_outputs = [ + item for sublist in transition_probabilities_outputs for item in sublist + ] + + property_kwargs[MarkovChainTransProbsCollector] = { + "inputs": transition_probabilities_outputs + } + if ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ): + if config.plasma.nlte_ionization_species: + nlte_ionization_species = config.plasma.nlte_ionization_species + for species in nlte_ionization_species: + if species not in config.plasma.continuum_interaction.species: + raise PlasmaConfigError( + f"NLTE ionization species {species} not in continuum species." + ) + if config.plasma.nlte_excitation_species: + nlte_excitation_species = config.plasma.nlte_excitation_species + for species in nlte_excitation_species: + if species not in config.plasma.continuum_interaction.species: + raise PlasmaConfigError( + f"NLTE excitation species {species} not in continuum species." + ) + property_kwargs[NLTEIndexHelper] = { + "nlte_ionization_species": config.plasma.nlte_ionization_species, + "nlte_excitation_species": config.plasma.nlte_excitation_species, + } + if config.plasma.nlte_solver == "lu": + continuum_plasma_modules += nlte_lu_solver_properties + logger.warning( + "LU solver will be inaccurate for NLTE excitation, proceed with caution." + ) + elif config.plasma.nlte_solver == "root": + continuum_plasma_modules += nlte_root_solver_properties + else: + raise PlasmaConfigError( + f"NLTE solver type unknown - {config.plasma.nlte_solver}" + ) + + # initializing rates + t_electrons = ( + config.plasma.link_t_rad_t_electron + * dilute_planckian_radiation_field.temperature.to(u.K).value + ) + initial_continuum_solver = DiluteBlackBodyContinuumPropertiesSolver( + atom_data + ) + initial_continuum_properties = initial_continuum_solver.solve( + dilute_planckian_radiation_field, t_electrons + ) + + kwargs.update( + gamma=initial_continuum_properties.photo_ionization_rate_coefficient, + bf_heating_coeff_estimator=None, + stim_recomb_cooling_coeff_estimator=None, + alpha_stim_factor=initial_continuum_properties.stimulated_recombination_rate_factor, + ) + return continuum_plasma_modules From 60239f8045c51b581dacaf382bd04b6203fc60f1 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 26 Jul 2024 11:59:37 -0400 Subject: [PATCH 62/89] updated tests --- tardis/io/configuration/tests/test_config_reader.py | 2 +- tardis/plasma/tests/test_nlte_solver.py | 2 +- tardis/plasma/tests/test_tardis_model_density_config.py | 2 +- tardis/simulation/base.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tardis/io/configuration/tests/test_config_reader.py b/tardis/io/configuration/tests/test_config_reader.py index 2fe68ea061d..e224af75fe0 100644 --- a/tardis/io/configuration/tests/test_config_reader.py +++ b/tardis/io/configuration/tests/test_config_reader.py @@ -10,7 +10,7 @@ from astropy.units import Quantity from tardis.io.configuration.config_reader import Configuration from tardis.plasma.exceptions import PlasmaConfigError -from tardis.plasma.assembly.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.base import assemble_plasma def test_convergence_section_parser(): diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py index 35b9b572d3f..2936a0d7339 100644 --- a/tardis/plasma/tests/test_nlte_solver.py +++ b/tardis/plasma/tests/test_nlte_solver.py @@ -14,7 +14,7 @@ calculate_jacobian_matrix, calculate_rate_matrix, ) -from tardis.plasma.assembly.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.base import assemble_plasma @pytest.fixture diff --git a/tardis/plasma/tests/test_tardis_model_density_config.py b/tardis/plasma/tests/test_tardis_model_density_config.py index 7b5be056ddc..df6edac0c8d 100644 --- a/tardis/plasma/tests/test_tardis_model_density_config.py +++ b/tardis/plasma/tests/test_tardis_model_density_config.py @@ -4,7 +4,7 @@ from tardis.io.configuration.config_reader import Configuration from tardis.model import SimulationState -from tardis.plasma.assembly.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.base import assemble_plasma @pytest.fixture diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index b5be5b46b60..3d62b615cc5 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -16,7 +16,7 @@ ) from tardis.io.util import HDFWriterMixin from tardis.plasma.radiation_field import DilutePlanckianRadiationField -from tardis.plasma.assembly.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.base import assemble_plasma from tardis.simulation.convergence import ConvergenceSolver from tardis.spectrum.base import SpectrumSolver from tardis.spectrum.formal_integral import FormalIntegrator From 5e4902a9199e4db97f66c15795c3d14fb8f415e0 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 26 Jul 2024 13:46:17 -0400 Subject: [PATCH 63/89] some more cleanup --- .../estimators/continuum_radfield_properties.py | 7 +++---- .../estimators/tests/test_continuum_property_solver.py | 8 ++++---- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py index 3f49a6629b1..9ee55da09ba 100644 --- a/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py +++ b/tardis/transport/montecarlo/estimators/continuum_radfield_properties.py @@ -109,7 +109,7 @@ def solve( photo_ion_rate_coeff = self.calculate_photo_ionization_rate_coefficient( mean_intensity_photo_ion_df ) - stimulated_recomb_rate_coeff = ( + stimulated_recomb_rate_factor = ( self.calculate_stimulated_recomb_rate_factor( mean_intensity_photo_ion_df, photo_ion_boltzmann_factor, @@ -117,7 +117,7 @@ def solve( ) return ContinuumProperties( - stimulated_recomb_rate_coeff, photo_ion_rate_coeff + stimulated_recomb_rate_factor, photo_ion_rate_coeff ) def calculate_photo_ionization_rate_coefficient( @@ -214,14 +214,13 @@ def calculate_mean_intensity_photo_ion_table( self.atom_data.photoionization_data.nu.values ) ) - mean_intensity_df = pd.DataFrame( + return pd.DataFrame( mean_intensity, index=self.atom_data.photoionization_data.index, columns=np.arange( len(dilute_blackbody_radiationfield_state.temperature) ), ) - return mean_intensity_df @dataclass diff --git a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py index 5fb75a0228c..2751861736b 100644 --- a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py +++ b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py @@ -65,10 +65,10 @@ def test_continuum_estimators( ) continuum_plasma.update( - gamma_estimator=transport_state.radfield_mc_estimators.photo_ion_estimator, - alpha_stim_estimator=transport_state.radfield_mc_estimators.stim_recomb_estimator, - bf_heating_coeff_estimator=transport_state.radfield_mc_estimators.bf_heating_estimator, - stim_recomb_cooling_coeff_estimator=transport_state.radfield_mc_estimators.stim_recomb_cooling_estimator, + gamma=continuum_properties_mc.photo_ionization_rate_coefficient, + alpha_stim_factor=continuum_properties_mc.stimulated_recombination_rate_factor, + bf_heating_coeff_estimator=None, + stim_recomb_cooling_coeff_estimator=None, ) pdt.assert_frame_equal( From 9f92256a5e0b39b80b9bcb0f343560e5a634114d Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 26 Jul 2024 13:48:03 -0400 Subject: [PATCH 64/89] fix benchmarks --- benchmarks/opacities_opacity.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/benchmarks/opacities_opacity.py b/benchmarks/opacities_opacity.py index e536e5c49a5..de571cc2184 100644 --- a/benchmarks/opacities_opacity.py +++ b/benchmarks/opacities_opacity.py @@ -4,9 +4,9 @@ from asv_runner.benchmarks.mark import parameterize -import tardis.opacities.compton_opacity_calculation import tardis.opacities.opacities as calculate_opacity from benchmarks.benchmark_base import BenchmarkBase +from tardis.opacities.opacities import compton_opacity_calculation class BenchmarkMontecarloMontecarloNumbaOpacities(BenchmarkBase): @@ -29,9 +29,7 @@ class BenchmarkMontecarloMontecarloNumbaOpacities(BenchmarkBase): } ) def time_compton_opacity_calculation(self, electron_number_density, energy): - tardis.opacities.compton_opacity_calculation.compton_opacity_calculation( - energy, electron_number_density - ) + compton_opacity_calculation(energy, electron_number_density) @parameterize( { From 3c7ee67695f60ad5654b5deff0c9e419797e59f9 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 26 Jul 2024 14:10:19 -0400 Subject: [PATCH 65/89] cleanup assembly --- tardis/plasma/assembly/legacy_assembly.py | 379 ++++++++++++++++++++++ 1 file changed, 379 insertions(+) create mode 100644 tardis/plasma/assembly/legacy_assembly.py diff --git a/tardis/plasma/assembly/legacy_assembly.py b/tardis/plasma/assembly/legacy_assembly.py new file mode 100644 index 00000000000..2a3c6724568 --- /dev/null +++ b/tardis/plasma/assembly/legacy_assembly.py @@ -0,0 +1,379 @@ +import logging + +import numpy as np +import pandas as pd +from astropy import units as u + +from tardis.plasma import BasePlasma +from tardis.plasma.base import PlasmaSolverSettings +from tardis.plasma.exceptions import PlasmaConfigError +from tardis.plasma.properties import ( + HeliumNumericalNLTE, + IonNumberDensity, + IonNumberDensityHeNLTE, + LevelBoltzmannFactorNLTE, + MarkovChainTransProbsCollector, + RadiationFieldCorrection, + StimulatedEmissionFactor, +) +from tardis.plasma.properties.base import TransitionProbabilitiesProperty +from tardis.plasma.properties.level_population import LevelNumberDensity +from tardis.plasma.properties.nlte_rate_equation_solver import ( + NLTEPopulationSolverLU, + NLTEPopulationSolverRoot, +) +from tardis.plasma.properties.property_collections import ( + adiabatic_cooling_properties, + basic_inputs, + basic_properties, + continuum_interaction_inputs, + continuum_interaction_properties, + dilute_lte_excitation_properties, + helium_lte_properties, + helium_nlte_properties, + helium_numerical_nlte_properties, + lte_excitation_properties, + lte_ionization_properties, + macro_atom_properties, + nebular_ionization_properties, + nlte_lu_solver_properties, + nlte_properties, + nlte_root_solver_properties, + non_nlte_properties, + two_photon_properties, +) +from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper +from tardis.plasma.radiation_field import DilutePlanckianRadiationField +from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( + DiluteBlackBodyContinuumPropertiesSolver, +) +from tardis.util.base import species_string_to_tuple + +logger = logging.getLogger(__name__) + + +def assemble_plasma(config, simulation_state, atom_data=None): + """ + Create a BasePlasma instance from a Configuration object + and a SimulationState. + + Parameters + ---------- + config : io.config_reader.Configuration + simulation_state : model.SimulationState + atom_data : atomic.AtomData + If None, an attempt will be made to read the atomic data + from config. + + Returns + ------- + : plasma.BasePlasma + + """ + # Convert the nlte species list to a proper format. + nlte_species = [ + species_string_to_tuple(species) + for species in config.plasma.nlte.species + ] + + # Convert the continuum interaction species list to a proper format. + continuum_interaction_species = [ + species_string_to_tuple(species) + for species in config.plasma.continuum_interaction.species + ] + continuum_interaction_species = pd.MultiIndex.from_tuples( + continuum_interaction_species, names=["atomic_number", "ion_number"] + ) + + atom_data.prepare_atom_data( + simulation_state.abundance.index, + line_interaction_type=config.plasma.line_interaction_type, + continuum_interaction_species=continuum_interaction_species, + nlte_species=nlte_species, + ) + + # Check if continuum interaction species are in selected_atoms + continuum_atoms = continuum_interaction_species.get_level_values( + "atomic_number" + ) + + continuum_atoms_in_selected_atoms = np.all( + continuum_atoms.isin(atom_data.selected_atomic_numbers) + ) + if not continuum_atoms_in_selected_atoms: + raise PlasmaConfigError( + "Not all continuum interaction species " + "belong to atoms that have been specified " + "in the configuration." + ) + + nlte_ionization_species = [ + species_string_to_tuple(species) + for species in config.plasma.nlte_ionization_species + ] + nlte_excitation_species = [ + species_string_to_tuple(species) + for species in config.plasma.nlte_excitation_species + ] + + dilute_planckian_radiation_field = DilutePlanckianRadiationField( + simulation_state.t_radiative, simulation_state.dilution_factor + ) + kwargs = dict( + dilute_planckian_radiation_field=dilute_planckian_radiation_field, + abundance=simulation_state.abundance, + number_density=simulation_state.elemental_number_density, + atomic_data=atom_data, + time_explosion=simulation_state.time_explosion, + link_t_rad_t_electron=config.plasma.link_t_rad_t_electron, + continuum_interaction_species=continuum_interaction_species, + nlte_ionization_species=nlte_ionization_species, + nlte_excitation_species=nlte_excitation_species, + ) + + plasma_modules = basic_inputs + basic_properties + property_kwargs = {} + + ########### SETTING UP CONTINUUM INTERACTIONS + + if len(config.plasma.continuum_interaction.species) > 0: + setup_continuum_interactions( + config, + atom_data, + dilute_planckian_radiation_field, + kwargs, + property_kwargs, + ) + + ##### RADIATIVE RATES SETUP + + plasma_solver_settings = PlasmaSolverSettings( + RADIATIVE_RATES_TYPE=config.plasma.radiative_rates_type + ) + + if (plasma_solver_settings.RADIATIVE_RATES_TYPE == "dilute-blackbody") or ( + plasma_solver_settings.RADIATIVE_RATES_TYPE == "detailed" + ): + kwargs["j_blues"] = pd.DataFrame( + dilute_planckian_radiation_field.calculate_mean_intensity( + atom_data.lines["nu"].values + ), + index=atom_data.lines.index, + ) + + elif plasma_solver_settings.RADIATIVE_RATES_TYPE == "blackbody": + planckian_rad_field = ( + dilute_planckian_radiation_field.to_planckian_radiation_field() + ) + kwargs["j_blues"] = pd.DataFrame( + planckian_rad_field.calculate_mean_intensity( + atom_data.lines["nu"].values + ), + index=atom_data.lines.index, + ) + + else: + raise ValueError( + f"radiative_rates_type type unknown - {plasma_solver_settings.RADIATIVE_RATES_TYPE}" + ) + + if config.plasma.excitation == "lte": + plasma_modules += lte_excitation_properties + elif config.plasma.excitation == "dilute-lte": + plasma_modules += dilute_lte_excitation_properties + + if config.plasma.ionization == "lte": + plasma_modules += lte_ionization_properties + elif config.plasma.ionization == "nebular": + plasma_modules += nebular_ionization_properties + + if nlte_species: + plasma_modules += nlte_properties + nlte_conf = config.plasma.nlte + plasma_modules.append(LevelBoltzmannFactorNLTE.from_config(nlte_conf)) + property_kwargs[StimulatedEmissionFactor] = dict( + nlte_species=nlte_species + ) + else: + plasma_modules += non_nlte_properties + + if config.plasma.line_interaction_type in ("downbranch", "macroatom"): + if not config.plasma.continuum_interaction.species: + plasma_modules += macro_atom_properties + + if "delta_treatment" in config.plasma: + property_kwargs[RadiationFieldCorrection] = dict( + delta_treatment=config.plasma.delta_treatment + ) + + if ( + config.plasma.helium_treatment == "recomb-nlte" + or config.plasma.helium_treatment == "numerical-nlte" + ) and ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ): + # Prevent the user from using helium NLTE treatment with + # NLTE ionization and excitation treatment. This is because + # the helium_nlte_properties could overwrite the NLTE ionization + # and excitation ion number and electron densities. + # helium_numerical_nlte_properties is also included here because + # it is currently in the same if else block, and thus may block + # the addition of the components from the else block. + raise PlasmaConfigError( + "Helium NLTE treatment is incompatible with the NLTE eonization and excitation treatment." + ) + + # TODO: Disentangle these if else block such that compatible components + # can be added independently. + if config.plasma.helium_treatment == "recomb-nlte": + plasma_modules += helium_nlte_properties + elif config.plasma.helium_treatment == "numerical-nlte": + plasma_modules += helium_numerical_nlte_properties + # TODO: See issue #633 + if config.plasma.heating_rate_data_file in ["none", None]: + raise PlasmaConfigError("Heating rate data file not specified") + else: + property_kwargs[HeliumNumericalNLTE] = dict( + heating_rate_data_file=config.plasma.heating_rate_data_file + ) + else: + # If nlte ionization species are present, we don't want to add the + # IonNumberDensity from helium_lte_properties, since we want + # to use the IonNumberDensity provided by the NLTE solver. + if ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ): + plasma_modules += [LevelNumberDensity] + else: + plasma_modules += helium_lte_properties + + if simulation_state._electron_densities is not None: + electron_densities = pd.Series( + simulation_state._electron_densities.cgs.value + ) + if config.plasma.helium_treatment == "numerical-nlte": + property_kwargs[IonNumberDensityHeNLTE] = dict( + electron_densities=electron_densities + ) + elif ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ) and config.plasma.nlte_solver == "root": + property_kwargs[NLTEPopulationSolverRoot] = dict( + electron_densities=electron_densities + ) + elif ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ) and config.plasma.nlte_solver == "lu": + property_kwargs[NLTEPopulationSolverLU] = dict( + electron_densities=electron_densities + ) + else: + property_kwargs[IonNumberDensity] = dict( + electron_densities=electron_densities + ) + + kwargs["helium_treatment"] = config.plasma.helium_treatment + + plasma = BasePlasma( + plasma_properties=plasma_modules, + property_kwargs=property_kwargs, + plasma_solver_settings=plasma_solver_settings, + **kwargs, + ) + + return plasma + + +def setup_continuum_interactions( + config, + atom_data, + dilute_planckian_radiation_field, + kwargs, + property_kwargs, +): + continuum_plasma_modules = [] + line_interaction_type = config.plasma.line_interaction_type + if line_interaction_type != "macroatom": + raise PlasmaConfigError( + "Continuum interactions require line_interaction_type " + f"macroatom (instead of {line_interaction_type})." + ) + + continuum_plasma_modules += continuum_interaction_properties + continuum_plasma_modules += continuum_interaction_inputs + + if config.plasma.continuum_interaction.enable_adiabatic_cooling: + continuum_plasma_modules += adiabatic_cooling_properties + + if config.plasma.continuum_interaction.enable_two_photon_decay: + continuum_plasma_modules += two_photon_properties + + transition_probabilities_outputs = [ + plasma_property.transition_probabilities_outputs + for plasma_property in continuum_plasma_modules + if issubclass(plasma_property, TransitionProbabilitiesProperty) + ] + transition_probabilities_outputs = [ + item for sublist in transition_probabilities_outputs for item in sublist + ] + + property_kwargs[MarkovChainTransProbsCollector] = { + "inputs": transition_probabilities_outputs + } + if ( + config.plasma.nlte_ionization_species + or config.plasma.nlte_excitation_species + ): + if config.plasma.nlte_ionization_species: + nlte_ionization_species = config.plasma.nlte_ionization_species + for species in nlte_ionization_species: + if species not in config.plasma.continuum_interaction.species: + raise PlasmaConfigError( + f"NLTE ionization species {species} not in continuum species." + ) + if config.plasma.nlte_excitation_species: + nlte_excitation_species = config.plasma.nlte_excitation_species + for species in nlte_excitation_species: + if species not in config.plasma.continuum_interaction.species: + raise PlasmaConfigError( + f"NLTE excitation species {species} not in continuum species." + ) + property_kwargs[NLTEIndexHelper] = { + "nlte_ionization_species": config.plasma.nlte_ionization_species, + "nlte_excitation_species": config.plasma.nlte_excitation_species, + } + if config.plasma.nlte_solver == "lu": + continuum_plasma_modules += nlte_lu_solver_properties + logger.warning( + "LU solver will be inaccurate for NLTE excitation, proceed with caution." + ) + elif config.plasma.nlte_solver == "root": + continuum_plasma_modules += nlte_root_solver_properties + else: + raise PlasmaConfigError( + f"NLTE solver type unknown - {config.plasma.nlte_solver}" + ) + + # initializing rates + t_electrons = ( + config.plasma.link_t_rad_t_electron + * dilute_planckian_radiation_field.temperature.to(u.K).value + ) + initial_continuum_solver = DiluteBlackBodyContinuumPropertiesSolver( + atom_data + ) + initial_continuum_properties = initial_continuum_solver.solve( + dilute_planckian_radiation_field, t_electrons + ) + + kwargs.update( + gamma=initial_continuum_properties.photo_ionization_rate_coefficient, + bf_heating_coeff_estimator=None, + stim_recomb_cooling_coeff_estimator=None, + alpha_stim_factor=initial_continuum_properties.stimulated_recombination_rate_factor, + ) + return continuum_plasma_modules From 0890747ac77b314fe34789beaf526be7e1bdb93b Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 26 Jul 2024 14:10:37 -0400 Subject: [PATCH 66/89] cleanup assembly --- .../configuration/tests/test_config_reader.py | 2 +- tardis/plasma/assembly/base.py | 404 ++---------------- tardis/plasma/tests/test_nlte_solver.py | 2 +- .../tests/test_tardis_model_density_config.py | 2 +- tardis/simulation/base.py | 2 +- .../tests/test_continuum_property_solver.py | 3 +- 6 files changed, 40 insertions(+), 375 deletions(-) diff --git a/tardis/io/configuration/tests/test_config_reader.py b/tardis/io/configuration/tests/test_config_reader.py index e224af75fe0..784797a2f03 100644 --- a/tardis/io/configuration/tests/test_config_reader.py +++ b/tardis/io/configuration/tests/test_config_reader.py @@ -10,7 +10,7 @@ from astropy.units import Quantity from tardis.io.configuration.config_reader import Configuration from tardis.plasma.exceptions import PlasmaConfigError -from tardis.plasma.assembly.base import assemble_plasma +from tardis.plasma.assembly.legacy_assembly import assemble_plasma def test_convergence_section_parser(): diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index 2a3c6724568..b8b384d6419 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -1,379 +1,45 @@ -import logging - -import numpy as np import pandas as pd -from astropy import units as u - -from tardis.plasma import BasePlasma -from tardis.plasma.base import PlasmaSolverSettings -from tardis.plasma.exceptions import PlasmaConfigError -from tardis.plasma.properties import ( - HeliumNumericalNLTE, - IonNumberDensity, - IonNumberDensityHeNLTE, - LevelBoltzmannFactorNLTE, - MarkovChainTransProbsCollector, - RadiationFieldCorrection, - StimulatedEmissionFactor, -) -from tardis.plasma.properties.base import TransitionProbabilitiesProperty -from tardis.plasma.properties.level_population import LevelNumberDensity -from tardis.plasma.properties.nlte_rate_equation_solver import ( - NLTEPopulationSolverLU, - NLTEPopulationSolverRoot, -) -from tardis.plasma.properties.property_collections import ( - adiabatic_cooling_properties, - basic_inputs, - basic_properties, - continuum_interaction_inputs, - continuum_interaction_properties, - dilute_lte_excitation_properties, - helium_lte_properties, - helium_nlte_properties, - helium_numerical_nlte_properties, - lte_excitation_properties, - lte_ionization_properties, - macro_atom_properties, - nebular_ionization_properties, - nlte_lu_solver_properties, - nlte_properties, - nlte_root_solver_properties, - non_nlte_properties, - two_photon_properties, -) -from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper -from tardis.plasma.radiation_field import DilutePlanckianRadiationField -from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( - DiluteBlackBodyContinuumPropertiesSolver, -) from tardis.util.base import species_string_to_tuple -logger = logging.getLogger(__name__) - - -def assemble_plasma(config, simulation_state, atom_data=None): - """ - Create a BasePlasma instance from a Configuration object - and a SimulationState. - - Parameters - ---------- - config : io.config_reader.Configuration - simulation_state : model.SimulationState - atom_data : atomic.AtomData - If None, an attempt will be made to read the atomic data - from config. - - Returns - ------- - : plasma.BasePlasma - - """ - # Convert the nlte species list to a proper format. - nlte_species = [ - species_string_to_tuple(species) - for species in config.plasma.nlte.species - ] - - # Convert the continuum interaction species list to a proper format. - continuum_interaction_species = [ - species_string_to_tuple(species) - for species in config.plasma.continuum_interaction.species - ] - continuum_interaction_species = pd.MultiIndex.from_tuples( - continuum_interaction_species, names=["atomic_number", "ion_number"] - ) - - atom_data.prepare_atom_data( - simulation_state.abundance.index, - line_interaction_type=config.plasma.line_interaction_type, - continuum_interaction_species=continuum_interaction_species, - nlte_species=nlte_species, - ) - - # Check if continuum interaction species are in selected_atoms - continuum_atoms = continuum_interaction_species.get_level_values( - "atomic_number" - ) - - continuum_atoms_in_selected_atoms = np.all( - continuum_atoms.isin(atom_data.selected_atomic_numbers) - ) - if not continuum_atoms_in_selected_atoms: - raise PlasmaConfigError( - "Not all continuum interaction species " - "belong to atoms that have been specified " - "in the configuration." - ) - - nlte_ionization_species = [ - species_string_to_tuple(species) - for species in config.plasma.nlte_ionization_species - ] - nlte_excitation_species = [ - species_string_to_tuple(species) - for species in config.plasma.nlte_excitation_species - ] - - dilute_planckian_radiation_field = DilutePlanckianRadiationField( - simulation_state.t_radiative, simulation_state.dilution_factor - ) - kwargs = dict( - dilute_planckian_radiation_field=dilute_planckian_radiation_field, - abundance=simulation_state.abundance, - number_density=simulation_state.elemental_number_density, - atomic_data=atom_data, - time_explosion=simulation_state.time_explosion, - link_t_rad_t_electron=config.plasma.link_t_rad_t_electron, - continuum_interaction_species=continuum_interaction_species, - nlte_ionization_species=nlte_ionization_species, - nlte_excitation_species=nlte_excitation_species, - ) - - plasma_modules = basic_inputs + basic_properties - property_kwargs = {} - - ########### SETTING UP CONTINUUM INTERACTIONS - - if len(config.plasma.continuum_interaction.species) > 0: - setup_continuum_interactions( - config, - atom_data, - dilute_planckian_radiation_field, - kwargs, - property_kwargs, - ) - - ##### RADIATIVE RATES SETUP - - plasma_solver_settings = PlasmaSolverSettings( - RADIATIVE_RATES_TYPE=config.plasma.radiative_rates_type - ) - - if (plasma_solver_settings.RADIATIVE_RATES_TYPE == "dilute-blackbody") or ( - plasma_solver_settings.RADIATIVE_RATES_TYPE == "detailed" - ): - kwargs["j_blues"] = pd.DataFrame( - dilute_planckian_radiation_field.calculate_mean_intensity( - atom_data.lines["nu"].values - ), - index=atom_data.lines.index, - ) - - elif plasma_solver_settings.RADIATIVE_RATES_TYPE == "blackbody": - planckian_rad_field = ( - dilute_planckian_radiation_field.to_planckian_radiation_field() - ) - kwargs["j_blues"] = pd.DataFrame( - planckian_rad_field.calculate_mean_intensity( - atom_data.lines["nu"].values - ), - index=atom_data.lines.index, - ) - - else: - raise ValueError( - f"radiative_rates_type type unknown - {plasma_solver_settings.RADIATIVE_RATES_TYPE}" - ) - if config.plasma.excitation == "lte": - plasma_modules += lte_excitation_properties - elif config.plasma.excitation == "dilute-lte": - plasma_modules += dilute_lte_excitation_properties - - if config.plasma.ionization == "lte": - plasma_modules += lte_ionization_properties - elif config.plasma.ionization == "nebular": - plasma_modules += nebular_ionization_properties - - if nlte_species: - plasma_modules += nlte_properties - nlte_conf = config.plasma.nlte - plasma_modules.append(LevelBoltzmannFactorNLTE.from_config(nlte_conf)) - property_kwargs[StimulatedEmissionFactor] = dict( - nlte_species=nlte_species - ) - else: - plasma_modules += non_nlte_properties - - if config.plasma.line_interaction_type in ("downbranch", "macroatom"): - if not config.plasma.continuum_interaction.species: - plasma_modules += macro_atom_properties - - if "delta_treatment" in config.plasma: - property_kwargs[RadiationFieldCorrection] = dict( - delta_treatment=config.plasma.delta_treatment - ) - - if ( - config.plasma.helium_treatment == "recomb-nlte" - or config.plasma.helium_treatment == "numerical-nlte" - ) and ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ): - # Prevent the user from using helium NLTE treatment with - # NLTE ionization and excitation treatment. This is because - # the helium_nlte_properties could overwrite the NLTE ionization - # and excitation ion number and electron densities. - # helium_numerical_nlte_properties is also included here because - # it is currently in the same if else block, and thus may block - # the addition of the components from the else block. - raise PlasmaConfigError( - "Helium NLTE treatment is incompatible with the NLTE eonization and excitation treatment." - ) +class PlasmaSolverFactory: - # TODO: Disentangle these if else block such that compatible components - # can be added independently. - if config.plasma.helium_treatment == "recomb-nlte": - plasma_modules += helium_nlte_properties - elif config.plasma.helium_treatment == "numerical-nlte": - plasma_modules += helium_numerical_nlte_properties - # TODO: See issue #633 - if config.plasma.heating_rate_data_file in ["none", None]: - raise PlasmaConfigError("Heating rate data file not specified") - else: - property_kwargs[HeliumNumericalNLTE] = dict( - heating_rate_data_file=config.plasma.heating_rate_data_file - ) - else: - # If nlte ionization species are present, we don't want to add the - # IonNumberDensity from helium_lte_properties, since we want - # to use the IonNumberDensity provided by the NLTE solver. - if ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ): - plasma_modules += [LevelNumberDensity] - else: - plasma_modules += helium_lte_properties + nlte_species: list + continuum_interaction_species: pd.MultiIndex - if simulation_state._electron_densities is not None: - electron_densities = pd.Series( - simulation_state._electron_densities.cgs.value + def __init__(self, config) -> None: + self.set_nlte_species_from_string(config.plasma.nlte.species) + self.set_continuum_interaction_species_from_string( + config.plasma.continuum_interaction.species ) - if config.plasma.helium_treatment == "numerical-nlte": - property_kwargs[IonNumberDensityHeNLTE] = dict( - electron_densities=electron_densities - ) - elif ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ) and config.plasma.nlte_solver == "root": - property_kwargs[NLTEPopulationSolverRoot] = dict( - electron_densities=electron_densities - ) - elif ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ) and config.plasma.nlte_solver == "lu": - property_kwargs[NLTEPopulationSolverLU] = dict( - electron_densities=electron_densities - ) - else: - property_kwargs[IonNumberDensity] = dict( - electron_densities=electron_densities - ) - kwargs["helium_treatment"] = config.plasma.helium_treatment - - plasma = BasePlasma( - plasma_properties=plasma_modules, - property_kwargs=property_kwargs, - plasma_solver_settings=plasma_solver_settings, - **kwargs, - ) - - return plasma - - -def setup_continuum_interactions( - config, - atom_data, - dilute_planckian_radiation_field, - kwargs, - property_kwargs, -): - continuum_plasma_modules = [] - line_interaction_type = config.plasma.line_interaction_type - if line_interaction_type != "macroatom": - raise PlasmaConfigError( - "Continuum interactions require line_interaction_type " - f"macroatom (instead of {line_interaction_type})." - ) - - continuum_plasma_modules += continuum_interaction_properties - continuum_plasma_modules += continuum_interaction_inputs - - if config.plasma.continuum_interaction.enable_adiabatic_cooling: - continuum_plasma_modules += adiabatic_cooling_properties - - if config.plasma.continuum_interaction.enable_two_photon_decay: - continuum_plasma_modules += two_photon_properties - - transition_probabilities_outputs = [ - plasma_property.transition_probabilities_outputs - for plasma_property in continuum_plasma_modules - if issubclass(plasma_property, TransitionProbabilitiesProperty) - ] - transition_probabilities_outputs = [ - item for sublist in transition_probabilities_outputs for item in sublist - ] + # Convert the continuum interaction species list to a proper format. - property_kwargs[MarkovChainTransProbsCollector] = { - "inputs": transition_probabilities_outputs - } - if ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species + def set_continuum_interaction_species_from_string( + self, continuum_interaction_species ): - if config.plasma.nlte_ionization_species: - nlte_ionization_species = config.plasma.nlte_ionization_species - for species in nlte_ionization_species: - if species not in config.plasma.continuum_interaction.species: - raise PlasmaConfigError( - f"NLTE ionization species {species} not in continuum species." - ) - if config.plasma.nlte_excitation_species: - nlte_excitation_species = config.plasma.nlte_excitation_species - for species in nlte_excitation_species: - if species not in config.plasma.continuum_interaction.species: - raise PlasmaConfigError( - f"NLTE excitation species {species} not in continuum species." - ) - property_kwargs[NLTEIndexHelper] = { - "nlte_ionization_species": config.plasma.nlte_ionization_species, - "nlte_excitation_species": config.plasma.nlte_excitation_species, - } - if config.plasma.nlte_solver == "lu": - continuum_plasma_modules += nlte_lu_solver_properties - logger.warning( - "LU solver will be inaccurate for NLTE excitation, proceed with caution." - ) - elif config.plasma.nlte_solver == "root": - continuum_plasma_modules += nlte_root_solver_properties - else: - raise PlasmaConfigError( - f"NLTE solver type unknown - {config.plasma.nlte_solver}" - ) - - # initializing rates - t_electrons = ( - config.plasma.link_t_rad_t_electron - * dilute_planckian_radiation_field.temperature.to(u.K).value - ) - initial_continuum_solver = DiluteBlackBodyContinuumPropertiesSolver( - atom_data - ) - initial_continuum_properties = initial_continuum_solver.solve( - dilute_planckian_radiation_field, t_electrons - ) - - kwargs.update( - gamma=initial_continuum_properties.photo_ionization_rate_coefficient, - bf_heating_coeff_estimator=None, - stim_recomb_cooling_coeff_estimator=None, - alpha_stim_factor=initial_continuum_properties.stimulated_recombination_rate_factor, - ) - return continuum_plasma_modules + continuum_interaction_species = [ + species_string_to_tuple(species) + for species in continuum_interaction_species + ] + self.continuum_interaction_species = pd.MultiIndex.from_tuples( + continuum_interaction_species, names=["atomic_number", "ion_number"] + ) + + def set_nlte_species_from_string(self, nlte_species): + """ + Sets the non-LTE species from a string representation. + + Parameters + ---------- + nlte_species : str + A string representation of the non-LTE species. + + Returns + ------- + None + This method does not return anything. + """ + self.nlte_species = [ + species_string_to_tuple(species) for species in nlte_species + ] diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py index 2936a0d7339..27d5bafc617 100644 --- a/tardis/plasma/tests/test_nlte_solver.py +++ b/tardis/plasma/tests/test_nlte_solver.py @@ -14,7 +14,7 @@ calculate_jacobian_matrix, calculate_rate_matrix, ) -from tardis.plasma.assembly.base import assemble_plasma +from tardis.plasma.assembly.legacy_assembly import assemble_plasma @pytest.fixture diff --git a/tardis/plasma/tests/test_tardis_model_density_config.py b/tardis/plasma/tests/test_tardis_model_density_config.py index df6edac0c8d..5551e7c4eb4 100644 --- a/tardis/plasma/tests/test_tardis_model_density_config.py +++ b/tardis/plasma/tests/test_tardis_model_density_config.py @@ -4,7 +4,7 @@ from tardis.io.configuration.config_reader import Configuration from tardis.model import SimulationState -from tardis.plasma.assembly.base import assemble_plasma +from tardis.plasma.assembly.legacy_assembly import assemble_plasma @pytest.fixture diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 3d62b615cc5..1626efbc4b4 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -16,7 +16,7 @@ ) from tardis.io.util import HDFWriterMixin from tardis.plasma.radiation_field import DilutePlanckianRadiationField -from tardis.plasma.assembly.base import assemble_plasma +from tardis.plasma.assembly.legacy_assembly import assemble_plasma from tardis.simulation.convergence import ConvergenceSolver from tardis.spectrum.base import SpectrumSolver from tardis.spectrum.formal_integral import FormalIntegrator diff --git a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py index 5fb75a0228c..eb831f26e57 100644 --- a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py +++ b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py @@ -1,14 +1,13 @@ from copy import deepcopy -import numpy.testing as npt import pandas.testing as pdt import pytest +from tardis.simulation import Simulation from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( DiluteBlackBodyContinuumPropertiesSolver, MCContinuumPropertiesSolver, ) -from tardis.simulation import Simulation @pytest.mark.continuum From 7540e55f24b8c5d76924c939eeec36c4b4653630 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Fri, 26 Jul 2024 17:30:35 -0400 Subject: [PATCH 67/89] working on the restructure --- tardis/plasma/assembly/base.py | 209 ++++++++++++++++++++- tardis/plasma/assembly/legacy_assembly.py | 2 +- tardis/plasma/base.py | 2 - tardis/plasma/properties/ion_population.py | 4 +- 4 files changed, 207 insertions(+), 10 deletions(-) diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index b8b384d6419..d331144b77b 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -1,27 +1,208 @@ +import numpy as np import pandas as pd + +from tardis.plasma.exceptions import PlasmaConfigError +from tardis.plasma.properties.property_collections import ( + basic_inputs, + basic_properties, + dilute_lte_excitation_properties, + lte_excitation_properties, + lte_ionization_properties, + nebular_ionization_properties, + nlte_properties, + non_nlte_properties, +) from tardis.util.base import species_string_to_tuple +from tardis.plasma.properties import ( + HeliumNumericalNLTE, + IonNumberDensity, + IonNumberDensityHeNLTE, + LevelBoltzmannFactorNLTE, + MarkovChainTransProbsCollector, + RadiationFieldCorrection, + StimulatedEmissionFactor, +) + + +def map_species_from_string(species): + return [species_string_to_tuple(spec) for spec in species] + + class PlasmaSolverFactory: - nlte_species: list continuum_interaction_species: pd.MultiIndex + line_interaction_type: str = "scatter" + + legacy_nlte_species: list = [] + nlte_exciation_species: list = [] + nlte_ionization_species: list = [] + plasma_modules: list = [] + kwargs: dict = {} + property_kwargs: dict = {} + + excitation_analytical_approximation: str = "lte" + ionization_analytical_approximation: str = "lte" + + radiative_rates_type: str = "dilute-blackbody" + + nebular_ionization_delta_treatment: tuple # species to use for the delta_treatment in nebular ionization ML93 + + def __init__(self, config, atom_data, atomic_numbers) -> None: - def __init__(self, config) -> None: self.set_nlte_species_from_string(config.plasma.nlte.species) self.set_continuum_interaction_species_from_string( config.plasma.continuum_interaction.species ) + self.line_interaction_type = config.plasma.line_interaction_type - # Convert the continuum interaction species list to a proper format. + self.atom_data = atom_data + self.atom_data.prepare_atom_data( + atomic_numbers, + line_interaction_type=config.plasma.line_interaction_type, + continuum_interaction_species=self.continuum_interaction_species, + nlte_species=self.legacy_nlte_species, + ) + + #### THIS IS VERY BAD BUT FOR NOW IS CHICKEN/EGG + # Check if continuum interaction species are in selected_atoms + continuum_atoms = self.continuum_interaction_species.get_level_values( + "atomic_number" + ) + + continuum_atoms_in_selected_atoms = np.all( + continuum_atoms.isin(atom_data.selected_atomic_numbers) + ) + if not continuum_atoms_in_selected_atoms: + raise PlasmaConfigError( + "Not all continuum interaction species " + "belong to atoms that have been specified " + "in the configuration." + ) + ##### ---------------------------- + + self.plasma_modules = basic_inputs + basic_properties + + self.kwargs = dict( + link_t_rad_t_electron=config.plasma.link_t_rad_t_electron, + continuum_interaction_species=continuum_interaction_species, + nlte_ionization_species=nlte_ionization_species, + nlte_excitation_species=nlte_excitation_species, + ) + + self.setup_analytical_approximations(config) + + self.setup_legacy_nlte(config.plasma.nlte) + if self.line_interaction_type in ("downbranch", "macroatom") and ( + len(self.continuum_interaction_species) == 0 + ): + self.setup_legacy_macro_atom(config) + + def setup_legacy_macro_atom(self, macro_atom_config=None): + self.plasma_modules += macro_atom_properties + + if macro_atom_config is not None: + self.plasma_modules.append( + MarkovChainTransProbsCollector.from_config(macro_atom_config) + ) + + def setup_legacy_nlte(self, nlte_config): + if len(self.legacy_nlte_species) > 0: + self.plasma_modules += nlte_properties + self.plasma_modules.append( + LevelBoltzmannFactorNLTE.from_config(nlte_config) + ) + self.property_kwargs[StimulatedEmissionFactor] = dict( + nlte_species=self.legacy_nlte_species + ) + else: + self.plasma_modules += non_nlte_properties + + def setup_analytical_approximations(self, plasma_config): + """ + Setup the analytical approximations for excitation and ionization. + + Returns + ------- + None + """ + self.excitation_analytical_approximation = plasma_config.excitation + self.ionization_analytical_approximation = plasma_config.ionization + plasma_modules = [] + if self.excitation_analytical_approximation == "lte": + plasma_modules += lte_excitation_properties + elif self.excitation_analytical_approximation == "dilute-lte": + plasma_modules += dilute_lte_excitation_properties + else: + raise PlasmaConfigError( + f'Invalid excitation analytical approximation. Configured as {self.excitation_analytical_approximation} but needs to be either "lte" or "dilute-lte"' + ) + + if self.ionization_analytical_approximation == "lte": + plasma_modules += lte_ionization_properties + elif self.ionization_analytical_approximation == "nebular": + plasma_modules += nebular_ionization_properties + if "delta_treatment" in plasma_config: + self.property_kwargs[RadiationFieldCorrection] = dict( + delta_treatment=plasma_config.delta_treatment + ) + else: + raise PlasmaConfigError( + f'Invalid excitation analytical approximation. Configured as {self.ionization_analytical_approximation} but needs to be either "lte" or "nebular"' + ) + + def setup_radiative_rates(self, radiative_rates_config): + ##### RADIATIVE RATES SETUP + + if (self.radiative_rates_type == "dilute-blackbody") or ( + self.radiative_rates_type == "detailed" + ): + self.kwargs["j_blues"] = pd.DataFrame( + dilute_planckian_radiation_field.calculate_mean_intensity( + atom_data.lines["nu"].values + ), + index=atom_data.lines.index, + ) + + elif plasma_solver_settings.RADIATIVE_RATES_TYPE == "blackbody": + planckian_rad_field = ( + dilute_planckian_radiation_field.to_planckian_radiation_field() + ) + kwargs["j_blues"] = pd.DataFrame( + planckian_rad_field.calculate_mean_intensity( + atom_data.lines["nu"].values + ), + index=atom_data.lines.index, + ) + + else: + raise ValueError( + f"radiative_rates_type type unknown - {plasma_solver_settings.RADIATIVE_RATES_TYPE}" + ) + + self.radiative_rates_type = radiative_rates_config.type def set_continuum_interaction_species_from_string( self, continuum_interaction_species ): + """ + Set the continuum interaction species from a list of species strings. + + Parameters + ---------- + continuum_interaction_species : list of str + List of species strings representing the continuum interaction species. + + Returns + ------- + None + """ continuum_interaction_species = [ species_string_to_tuple(species) for species in continuum_interaction_species ] + self.continuum_interaction_species = pd.MultiIndex.from_tuples( continuum_interaction_species, names=["atomic_number", "ion_number"] ) @@ -40,6 +221,22 @@ def set_nlte_species_from_string(self, nlte_species): None This method does not return anything. """ - self.nlte_species = [ - species_string_to_tuple(species) for species in nlte_species - ] + self.legacy_nlte_species = map_species_from_string(nlte_species) + + def assemble( + self, dilute_planckian_radiation_field, simulation_state, atom_data + ): + + kwargs = dict( + time_explosion=simulation_state.time_explosion, + dilute_planckian_radiation_field=dilute_planckian_radiation_field, + abundance=simulation_state.abundance, + number_density=simulation_state.elemental_number_density, + atomic_data=atom_data, + ) + + return BasePlasma( + plasma_properties=self.plasma_modules, + property_kwargs=self.property_kwargs, + **kwargs, + ) diff --git a/tardis/plasma/assembly/legacy_assembly.py b/tardis/plasma/assembly/legacy_assembly.py index 2a3c6724568..bd163497ce7 100644 --- a/tardis/plasma/assembly/legacy_assembly.py +++ b/tardis/plasma/assembly/legacy_assembly.py @@ -221,7 +221,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): # it is currently in the same if else block, and thus may block # the addition of the components from the else block. raise PlasmaConfigError( - "Helium NLTE treatment is incompatible with the NLTE eonization and excitation treatment." + "Helium NLTE treatment is incompatible with the NLTE ionization and excitation treatment." ) # TODO: Disentangle these if else block such that compatible components diff --git a/tardis/plasma/base.py b/tardis/plasma/base.py index 8a42c32dfdb..b6b781a691d 100644 --- a/tardis/plasma/base.py +++ b/tardis/plasma/base.py @@ -26,11 +26,9 @@ class BasePlasma(PlasmaWriterMixin): def __init__( self, plasma_properties, - plasma_solver_settings, property_kwargs=None, **kwargs, ): - self.plasma_solver_settings = plasma_solver_settings self.outputs_dict = {} self.input_properties = [] self.plasma_properties = self._init_properties( diff --git a/tardis/plasma/properties/ion_population.py b/tardis/plasma/properties/ion_population.py index 10f440b95c7..1d342163e21 100644 --- a/tardis/plasma/properties/ion_population.py +++ b/tardis/plasma/properties/ion_population.py @@ -221,7 +221,9 @@ def calculate( self._set_chi_0(ionization_data) if self.delta_treatment is None: if self.departure_coefficient is None: - departure_coefficient = 1.0 / w + departure_coefficient = ( + 1.0 / w + ) # see Equation 13 and explanations on page 451 lower right in ML 93 else: departure_coefficient = self.departure_coefficient radiation_field_correction = -np.ones( From e29ae2212817900c4734d923151ef96248b7eec8 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 14:54:25 -0400 Subject: [PATCH 68/89] slowly fixing the assembly module --- tardis/plasma/assembly/base.py | 352 +++++++++++++++++---- tardis/plasma/assembly/legacy_assembly.py | 357 +--------------------- 2 files changed, 302 insertions(+), 407 deletions(-) diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index d331144b77b..2b990027c36 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -1,30 +1,53 @@ +import logging + import numpy as np import pandas as pd +from astropy import units as u +from tardis.plasma import BasePlasma from tardis.plasma.exceptions import PlasmaConfigError +from tardis.plasma.properties import ( + HeliumNumericalNLTE, + IonNumberDensity, + IonNumberDensityHeNLTE, + LevelBoltzmannFactorNLTE, + MarkovChainTransProbsCollector, + RadiationFieldCorrection, + StimulatedEmissionFactor, +) +from tardis.plasma.properties.base import TransitionProbabilitiesProperty +from tardis.plasma.properties.level_population import LevelNumberDensity +from tardis.plasma.properties.nlte_rate_equation_solver import ( + NLTEPopulationSolverLU, + NLTEPopulationSolverRoot, +) from tardis.plasma.properties.property_collections import ( + adiabatic_cooling_properties, basic_inputs, basic_properties, + continuum_interaction_inputs, + continuum_interaction_properties, dilute_lte_excitation_properties, + helium_lte_properties, + helium_nlte_properties, + helium_numerical_nlte_properties, lte_excitation_properties, lte_ionization_properties, + macro_atom_properties, nebular_ionization_properties, + nlte_lu_solver_properties, nlte_properties, + nlte_root_solver_properties, non_nlte_properties, + two_photon_properties, ) -from tardis.util.base import species_string_to_tuple - - -from tardis.plasma.properties import ( - HeliumNumericalNLTE, - IonNumberDensity, - IonNumberDensityHeNLTE, - LevelBoltzmannFactorNLTE, - MarkovChainTransProbsCollector, - RadiationFieldCorrection, - StimulatedEmissionFactor, +from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper +from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( + DiluteBlackBodyContinuumPropertiesSolver, ) +from tardis.util.base import species_string_to_tuple +logger = logging.getLogger(__name__) def map_species_from_string(species): return [species_string_to_tuple(spec) for spec in species] @@ -32,22 +55,36 @@ def map_species_from_string(species): class PlasmaSolverFactory: - continuum_interaction_species: pd.MultiIndex - line_interaction_type: str = "scatter" + ## Analytical Approximations + excitation_analytical_approximation: str = "lte" + ionization_analytical_approximation: str = "lte" + nebular_ionization_delta_treatment: tuple # species to use for the delta_treatment in nebular ionization ML93 + link_t_rad_t_electron: float = 1.0 + + radiative_rates_type: str = "dilute-blackbody" + + ## Statistical Balance Solver legacy_nlte_species: list = [] - nlte_exciation_species: list = [] + + nlte_excitation_species: list = [] nlte_ionization_species: list = [] - plasma_modules: list = [] - kwargs: dict = {} - property_kwargs: dict = {} + nlte_solver: str = "lu" - excitation_analytical_approximation: str = "lte" - ionization_analytical_approximation: str = "lte" + ## Helium Treatment options + helium_treatment: str = "none" + heating_rate_data_file: str = "none" - radiative_rates_type: str = "dilute-blackbody" + ## Continuum Interaction + continuum_interaction_species: pd.MultiIndex - nebular_ionization_delta_treatment: tuple # species to use for the delta_treatment in nebular ionization ML93 + ## Opacities + line_interaction_type: str = "scatter" + + ## Assembly properties + plasma_modules: list = [] + kwargs: dict = {} + property_kwargs: dict = {} def __init__(self, config, atom_data, atomic_numbers) -> None: @@ -83,31 +120,105 @@ def __init__(self, config, atom_data, atomic_numbers) -> None: ##### ---------------------------- self.plasma_modules = basic_inputs + basic_properties + self.link_t_rad_t_electron = config.plasma.link_t_rad_t_electron - self.kwargs = dict( - link_t_rad_t_electron=config.plasma.link_t_rad_t_electron, - continuum_interaction_species=continuum_interaction_species, - nlte_ionization_species=nlte_ionization_species, - nlte_excitation_species=nlte_excitation_species, - ) - - self.setup_analytical_approximations(config) + self.setup_analytical_approximations(config.plasma) self.setup_legacy_nlte(config.plasma.nlte) if self.line_interaction_type in ("downbranch", "macroatom") and ( len(self.continuum_interaction_species) == 0 ): - self.setup_legacy_macro_atom(config) + self.setup_legacy_macro_atom() + + self.helium_treatment = config.plasma.helium_treatment + self.heating_rate_data_file = config.plasma.heating_rate_data_file + self.setup_helium_treatment() + if len(config.plasma.continuum_interaction.species) > 0: + self.setup_continuum_interactions( + config.plasma.continuum_interaction + ) - def setup_legacy_macro_atom(self, macro_atom_config=None): - self.plasma_modules += macro_atom_properties + def setup_helium_treatment(self): + """ + Set up the helium treatment for the plasma assembly. - if macro_atom_config is not None: - self.plasma_modules.append( - MarkovChainTransProbsCollector.from_config(macro_atom_config) + Parameters + ---------- + helium_treatment : str + The type of helium treatment to be used. Possible values are: + - "recomb-nlte": Use recombination NLTE treatment for helium. + - "numerical-nlte": Use numerical NLTE treatment for helium. + + heating_rate_data_file : str or None + The path to the heating rate data file. Required when using + "numerical-nlte" helium treatment. + + Raises + ------ + PlasmaConfigError + If the helium NLTE treatment is incompatible with the NLTE ionization + and excitation treatment. + + If the heating rate data file is not specified when using + "numerical-nlte" helium treatment. + """ + if ( + self.helium_treatment == "recomb-nlte" + or self.helium_treatment == "numerical-nlte" + ) and ( + len(self.nlte_ionization_species + self.nlte_excitation_species) > 0 + ): + # Prevent the user from using helium NLTE treatment with + # NLTE ionization and excitation treatment. This is because + # the helium_nlte_properties could overwrite the NLTE ionization + # and excitation ion number and electron densities. + # helium_numerical_nlte_properties is also included here because + # it is currently in the same if else block, and thus may block + # the addition of the components from the else block. + raise PlasmaConfigError( + "Helium NLTE treatment is incompatible with the NLTE ionization and excitation treatment." ) + # TODO: Disentangle these if else block such that compatible components + # can be added independently. + if self.helium_treatment == "recomb-nlte": + self.plasma_modules += helium_nlte_properties + elif self.helium_treatment == "numerical-nlte": + self.plasma_modules += helium_numerical_nlte_properties + if heating_rate_data_file in ["none", None]: + raise PlasmaConfigError("Heating rate data file not specified") + self.property_kwargs[HeliumNumericalNLTE] = dict( + heating_rate_data_file=heating_rate_data_file + ) + else: + # If nlte ionization species are present, we don't want to add the + # IonNumberDensity from helium_lte_properties, since we want + # to use the IonNumberDensity provided by the NLTE solver. + if ( + len(self.nlte_ionization_species + self.nlte_excitation_species) + > 0 + ): + self.plasma_modules.append(LevelNumberDensity) + else: + self.plasma_modules += helium_lte_properties + + def setup_legacy_macro_atom(self): + self.plasma_modules += macro_atom_properties + def setup_legacy_nlte(self, nlte_config): + """ + Set up the non-LTE (NLTE) properties for the legacy species. + + Parameters: + ----------- + nlte_config : dict + A dictionary containing the NLTE configuration. + + Notes: + ------ + This method adds the NLTE properties for the legacy species to the plasma modules. + If there are no legacy NLTE species, it adds the non-NLTE properties instead. + """ if len(self.legacy_nlte_species) > 0: self.plasma_modules += nlte_properties self.plasma_modules.append( @@ -129,20 +240,20 @@ def setup_analytical_approximations(self, plasma_config): """ self.excitation_analytical_approximation = plasma_config.excitation self.ionization_analytical_approximation = plasma_config.ionization - plasma_modules = [] + if self.excitation_analytical_approximation == "lte": - plasma_modules += lte_excitation_properties + self.plasma_modules += lte_excitation_properties elif self.excitation_analytical_approximation == "dilute-lte": - plasma_modules += dilute_lte_excitation_properties + self.plasma_modules += dilute_lte_excitation_properties else: raise PlasmaConfigError( f'Invalid excitation analytical approximation. Configured as {self.excitation_analytical_approximation} but needs to be either "lte" or "dilute-lte"' ) if self.ionization_analytical_approximation == "lte": - plasma_modules += lte_ionization_properties + self.plasma_modules += lte_ionization_properties elif self.ionization_analytical_approximation == "nebular": - plasma_modules += nebular_ionization_properties + self.plasma_modules += nebular_ionization_properties if "delta_treatment" in plasma_config: self.property_kwargs[RadiationFieldCorrection] = dict( delta_treatment=plasma_config.delta_treatment @@ -152,36 +263,34 @@ def setup_analytical_approximations(self, plasma_config): f'Invalid excitation analytical approximation. Configured as {self.ionization_analytical_approximation} but needs to be either "lte" or "nebular"' ) - def setup_radiative_rates(self, radiative_rates_config): - ##### RADIATIVE RATES SETUP - + def initialize_j_blues(self, dilute_planckian_radiation_field, lines_df): if (self.radiative_rates_type == "dilute-blackbody") or ( self.radiative_rates_type == "detailed" ): - self.kwargs["j_blues"] = pd.DataFrame( + j_blues = pd.DataFrame( dilute_planckian_radiation_field.calculate_mean_intensity( - atom_data.lines["nu"].values + lines_df.nu.values ), - index=atom_data.lines.index, + index=lines_df.index, ) - elif plasma_solver_settings.RADIATIVE_RATES_TYPE == "blackbody": + elif self.radiative_rates_type == "blackbody": planckian_rad_field = ( dilute_planckian_radiation_field.to_planckian_radiation_field() ) - kwargs["j_blues"] = pd.DataFrame( + j_blues = pd.DataFrame( planckian_rad_field.calculate_mean_intensity( - atom_data.lines["nu"].values + lines_df.nu.values ), - index=atom_data.lines.index, + index=lines_df.index, ) else: raise ValueError( - f"radiative_rates_type type unknown - {plasma_solver_settings.RADIATIVE_RATES_TYPE}" + f"radiative_rates_type type unknown - {self.radiative_rates_type}" ) - self.radiative_rates_type = radiative_rates_config.type + return j_blues def set_continuum_interaction_species_from_string( self, continuum_interaction_species @@ -223,18 +332,149 @@ def set_nlte_species_from_string(self, nlte_species): """ self.legacy_nlte_species = map_species_from_string(nlte_species) - def assemble( - self, dilute_planckian_radiation_field, simulation_state, atom_data - ): + def setup_continuum_interactions(self, config_continuum_interaction): + line_interaction_type = self.line_interaction_type + if line_interaction_type != "macroatom": + raise PlasmaConfigError( + "Continuum interactions require line_interaction_type " + f"macroatom (instead of {line_interaction_type})." + ) + + self.plasma_modules += continuum_interaction_properties + self.plasma_modules += continuum_interaction_inputs + + if config_continuum_interaction.enable_adiabatic_cooling: + self.plasma_modules += adiabatic_cooling_properties + + if config_continuum_interaction.enable_two_photon_decay: + self.plasma_modules += two_photon_properties + + transition_probabilities_outputs = [ + plasma_property.transition_probabilities_outputs + for plasma_property in self.plasma_modules + if issubclass(plasma_property, TransitionProbabilitiesProperty) + ] + transition_probabilities_outputs = [ + item + for sublist in transition_probabilities_outputs + for item in sublist + ] + + self.property_kwargs[MarkovChainTransProbsCollector] = { + "inputs": transition_probabilities_outputs + } + if len(self.nlte_ionization_species + self.nlte_excitation_species) > 0: + if len(self.nlte_ionization_species) > 0: + nlte_ionization_species = self.nlte_ionization_species + for species in nlte_ionization_species: + if species not in self.continuum_interaction_species: + raise PlasmaConfigError( + f"NLTE ionization species {species} not in continuum species." + ) + if len(self.nlte_excitation_species) > 0: + nlte_excitation_species = self.nlte_excitation_species + for species in nlte_excitation_species: + if species not in self.continuum_interaction.species: + raise PlasmaConfigError( + f"NLTE excitation species {species} not in continuum species." + ) + self.property_kwargs[NLTEIndexHelper] = { + "nlte_ionization_species": self.nlte_ionization_species, + "nlte_excitation_species": self.nlte_excitation_species, + } + if self.nlte_solver == "lu": + self.plasma_modules += nlte_lu_solver_properties + logger.warning( + "LU solver will be inaccurate for NLTE excitation, proceed with caution." + ) + elif self.nlte_solver == "root": + self.plasma_modules += nlte_root_solver_properties + else: + raise PlasmaConfigError( + f"NLTE solver type unknown - {self.nlte_solver}" + ) + + def setup_electron_densities(self, electron_densities): + if self.helium_treatment == "numerical-nlte": + self.property_kwargs[IonNumberDensityHeNLTE] = dict( + electron_densities=electron_densities + ) + elif ( + len(self.nlte_ionization_species + self.nlte_excitation_species) > 0 + ) and self.nlte_solver == "root": + self.property_kwargs[NLTEPopulationSolverRoot] = dict( + electron_densities=electron_densities + ) + elif ( + len(self.nlte_ionization_species + self.nlte_excitation_species) > 0 + ) and self.nlte_solver == "lu": + self.property_kwargs[NLTEPopulationSolverLU] = dict( + electron_densities=electron_densities + ) + else: + self.property_kwargs[IonNumberDensity] = dict( + electron_densities=electron_densities + ) + + def initialize_continuum_properties(self, dilute_planckian_radiation_field): + """ + Initialize the continuum properties of the plasma. + + Parameters + ---------- + dilute_planckian_radiation_field : DilutePlanckianRadiationField + The dilute Planckian radiation field. + + Returns + ------- + initial_continuum_properties : `~tardis.plasma.properties.ContinuumProperties` + The initial continuum properties of the plasma. + """ + t_electrons = dilute_planckian_radiation_field.temperature.to(u.K).value + + initial_continuum_solver = DiluteBlackBodyContinuumPropertiesSolver( + self.atom_data + ) + initial_continuum_properties = initial_continuum_solver.solve( + dilute_planckian_radiation_field, t_electrons + ) + return initial_continuum_properties + + def assemble(self, dilute_planckian_radiation_field, simulation_state): + j_blues = self.initialize_j_blues( + dilute_planckian_radiation_field, self.atom_data.lines + ) kwargs = dict( time_explosion=simulation_state.time_explosion, dilute_planckian_radiation_field=dilute_planckian_radiation_field, abundance=simulation_state.abundance, number_density=simulation_state.elemental_number_density, - atomic_data=atom_data, + link_t_rad_t_electron=self.link_t_rad_t_electron, + atomic_data=self.atom_data, + j_blues=j_blues, + continuum_interaction_species=self.continuum_interaction_species, + nlte_ionization_species=self.nlte_ionization_species, + nlte_excitation_species=self.nlte_excitation_species, ) + if len(self.continuum_interaction_species) > 0: + initial_continuum_properties = self.initialize_continuum_properties( + dilute_planckian_radiation_field + ) + kwargs.update( + gamma=initial_continuum_properties.photo_ionization_rate_coefficient, + bf_heating_coeff_estimator=None, + stim_recomb_cooling_coeff_estimator=None, + alpha_stim_factor=initial_continuum_properties.stimulated_recombination_rate_factor, + ) + + if simulation_state._electron_densities is not None: + electron_densities = pd.Series( + simulation_state._electron_densities.cgs.value + ) + self.setup_electron_densities(electron_densities) + kwargs["helium_treatment"] = self.helium_treatment return BasePlasma( plasma_properties=self.plasma_modules, property_kwargs=self.property_kwargs, diff --git a/tardis/plasma/assembly/legacy_assembly.py b/tardis/plasma/assembly/legacy_assembly.py index bd163497ce7..7e45266dffc 100644 --- a/tardis/plasma/assembly/legacy_assembly.py +++ b/tardis/plasma/assembly/legacy_assembly.py @@ -1,55 +1,5 @@ -import logging - -import numpy as np -import pandas as pd -from astropy import units as u - -from tardis.plasma import BasePlasma -from tardis.plasma.base import PlasmaSolverSettings -from tardis.plasma.exceptions import PlasmaConfigError -from tardis.plasma.properties import ( - HeliumNumericalNLTE, - IonNumberDensity, - IonNumberDensityHeNLTE, - LevelBoltzmannFactorNLTE, - MarkovChainTransProbsCollector, - RadiationFieldCorrection, - StimulatedEmissionFactor, -) -from tardis.plasma.properties.base import TransitionProbabilitiesProperty -from tardis.plasma.properties.level_population import LevelNumberDensity -from tardis.plasma.properties.nlte_rate_equation_solver import ( - NLTEPopulationSolverLU, - NLTEPopulationSolverRoot, -) -from tardis.plasma.properties.property_collections import ( - adiabatic_cooling_properties, - basic_inputs, - basic_properties, - continuum_interaction_inputs, - continuum_interaction_properties, - dilute_lte_excitation_properties, - helium_lte_properties, - helium_nlte_properties, - helium_numerical_nlte_properties, - lte_excitation_properties, - lte_ionization_properties, - macro_atom_properties, - nebular_ionization_properties, - nlte_lu_solver_properties, - nlte_properties, - nlte_root_solver_properties, - non_nlte_properties, - two_photon_properties, -) -from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper +from tardis.plasma.assembly.base import PlasmaSolverFactory from tardis.plasma.radiation_field import DilutePlanckianRadiationField -from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( - DiluteBlackBodyContinuumPropertiesSolver, -) -from tardis.util.base import species_string_to_tuple - -logger = logging.getLogger(__name__) def assemble_plasma(config, simulation_state, atom_data=None): @@ -70,310 +20,15 @@ def assemble_plasma(config, simulation_state, atom_data=None): : plasma.BasePlasma """ - # Convert the nlte species list to a proper format. - nlte_species = [ - species_string_to_tuple(species) - for species in config.plasma.nlte.species - ] - - # Convert the continuum interaction species list to a proper format. - continuum_interaction_species = [ - species_string_to_tuple(species) - for species in config.plasma.continuum_interaction.species - ] - continuum_interaction_species = pd.MultiIndex.from_tuples( - continuum_interaction_species, names=["atomic_number", "ion_number"] - ) - - atom_data.prepare_atom_data( - simulation_state.abundance.index, - line_interaction_type=config.plasma.line_interaction_type, - continuum_interaction_species=continuum_interaction_species, - nlte_species=nlte_species, + atomic_numbers = simulation_state.abundance.index + plasma_solver_factory = PlasmaSolverFactory( + config, atom_data, atomic_numbers ) - # Check if continuum interaction species are in selected_atoms - continuum_atoms = continuum_interaction_species.get_level_values( - "atomic_number" - ) - - continuum_atoms_in_selected_atoms = np.all( - continuum_atoms.isin(atom_data.selected_atomic_numbers) - ) - if not continuum_atoms_in_selected_atoms: - raise PlasmaConfigError( - "Not all continuum interaction species " - "belong to atoms that have been specified " - "in the configuration." - ) - - nlte_ionization_species = [ - species_string_to_tuple(species) - for species in config.plasma.nlte_ionization_species - ] - nlte_excitation_species = [ - species_string_to_tuple(species) - for species in config.plasma.nlte_excitation_species - ] - dilute_planckian_radiation_field = DilutePlanckianRadiationField( simulation_state.t_radiative, simulation_state.dilution_factor ) - kwargs = dict( - dilute_planckian_radiation_field=dilute_planckian_radiation_field, - abundance=simulation_state.abundance, - number_density=simulation_state.elemental_number_density, - atomic_data=atom_data, - time_explosion=simulation_state.time_explosion, - link_t_rad_t_electron=config.plasma.link_t_rad_t_electron, - continuum_interaction_species=continuum_interaction_species, - nlte_ionization_species=nlte_ionization_species, - nlte_excitation_species=nlte_excitation_species, - ) - - plasma_modules = basic_inputs + basic_properties - property_kwargs = {} - - ########### SETTING UP CONTINUUM INTERACTIONS - - if len(config.plasma.continuum_interaction.species) > 0: - setup_continuum_interactions( - config, - atom_data, - dilute_planckian_radiation_field, - kwargs, - property_kwargs, - ) - - ##### RADIATIVE RATES SETUP - - plasma_solver_settings = PlasmaSolverSettings( - RADIATIVE_RATES_TYPE=config.plasma.radiative_rates_type - ) - - if (plasma_solver_settings.RADIATIVE_RATES_TYPE == "dilute-blackbody") or ( - plasma_solver_settings.RADIATIVE_RATES_TYPE == "detailed" - ): - kwargs["j_blues"] = pd.DataFrame( - dilute_planckian_radiation_field.calculate_mean_intensity( - atom_data.lines["nu"].values - ), - index=atom_data.lines.index, - ) - - elif plasma_solver_settings.RADIATIVE_RATES_TYPE == "blackbody": - planckian_rad_field = ( - dilute_planckian_radiation_field.to_planckian_radiation_field() - ) - kwargs["j_blues"] = pd.DataFrame( - planckian_rad_field.calculate_mean_intensity( - atom_data.lines["nu"].values - ), - index=atom_data.lines.index, - ) - - else: - raise ValueError( - f"radiative_rates_type type unknown - {plasma_solver_settings.RADIATIVE_RATES_TYPE}" - ) - - if config.plasma.excitation == "lte": - plasma_modules += lte_excitation_properties - elif config.plasma.excitation == "dilute-lte": - plasma_modules += dilute_lte_excitation_properties - - if config.plasma.ionization == "lte": - plasma_modules += lte_ionization_properties - elif config.plasma.ionization == "nebular": - plasma_modules += nebular_ionization_properties - - if nlte_species: - plasma_modules += nlte_properties - nlte_conf = config.plasma.nlte - plasma_modules.append(LevelBoltzmannFactorNLTE.from_config(nlte_conf)) - property_kwargs[StimulatedEmissionFactor] = dict( - nlte_species=nlte_species - ) - else: - plasma_modules += non_nlte_properties - - if config.plasma.line_interaction_type in ("downbranch", "macroatom"): - if not config.plasma.continuum_interaction.species: - plasma_modules += macro_atom_properties - - if "delta_treatment" in config.plasma: - property_kwargs[RadiationFieldCorrection] = dict( - delta_treatment=config.plasma.delta_treatment - ) - - if ( - config.plasma.helium_treatment == "recomb-nlte" - or config.plasma.helium_treatment == "numerical-nlte" - ) and ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ): - # Prevent the user from using helium NLTE treatment with - # NLTE ionization and excitation treatment. This is because - # the helium_nlte_properties could overwrite the NLTE ionization - # and excitation ion number and electron densities. - # helium_numerical_nlte_properties is also included here because - # it is currently in the same if else block, and thus may block - # the addition of the components from the else block. - raise PlasmaConfigError( - "Helium NLTE treatment is incompatible with the NLTE ionization and excitation treatment." - ) - - # TODO: Disentangle these if else block such that compatible components - # can be added independently. - if config.plasma.helium_treatment == "recomb-nlte": - plasma_modules += helium_nlte_properties - elif config.plasma.helium_treatment == "numerical-nlte": - plasma_modules += helium_numerical_nlte_properties - # TODO: See issue #633 - if config.plasma.heating_rate_data_file in ["none", None]: - raise PlasmaConfigError("Heating rate data file not specified") - else: - property_kwargs[HeliumNumericalNLTE] = dict( - heating_rate_data_file=config.plasma.heating_rate_data_file - ) - else: - # If nlte ionization species are present, we don't want to add the - # IonNumberDensity from helium_lte_properties, since we want - # to use the IonNumberDensity provided by the NLTE solver. - if ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ): - plasma_modules += [LevelNumberDensity] - else: - plasma_modules += helium_lte_properties - - if simulation_state._electron_densities is not None: - electron_densities = pd.Series( - simulation_state._electron_densities.cgs.value - ) - if config.plasma.helium_treatment == "numerical-nlte": - property_kwargs[IonNumberDensityHeNLTE] = dict( - electron_densities=electron_densities - ) - elif ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ) and config.plasma.nlte_solver == "root": - property_kwargs[NLTEPopulationSolverRoot] = dict( - electron_densities=electron_densities - ) - elif ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ) and config.plasma.nlte_solver == "lu": - property_kwargs[NLTEPopulationSolverLU] = dict( - electron_densities=electron_densities - ) - else: - property_kwargs[IonNumberDensity] = dict( - electron_densities=electron_densities - ) - - kwargs["helium_treatment"] = config.plasma.helium_treatment - - plasma = BasePlasma( - plasma_properties=plasma_modules, - property_kwargs=property_kwargs, - plasma_solver_settings=plasma_solver_settings, - **kwargs, - ) - - return plasma - - -def setup_continuum_interactions( - config, - atom_data, - dilute_planckian_radiation_field, - kwargs, - property_kwargs, -): - continuum_plasma_modules = [] - line_interaction_type = config.plasma.line_interaction_type - if line_interaction_type != "macroatom": - raise PlasmaConfigError( - "Continuum interactions require line_interaction_type " - f"macroatom (instead of {line_interaction_type})." - ) - - continuum_plasma_modules += continuum_interaction_properties - continuum_plasma_modules += continuum_interaction_inputs - - if config.plasma.continuum_interaction.enable_adiabatic_cooling: - continuum_plasma_modules += adiabatic_cooling_properties - - if config.plasma.continuum_interaction.enable_two_photon_decay: - continuum_plasma_modules += two_photon_properties - - transition_probabilities_outputs = [ - plasma_property.transition_probabilities_outputs - for plasma_property in continuum_plasma_modules - if issubclass(plasma_property, TransitionProbabilitiesProperty) - ] - transition_probabilities_outputs = [ - item for sublist in transition_probabilities_outputs for item in sublist - ] - - property_kwargs[MarkovChainTransProbsCollector] = { - "inputs": transition_probabilities_outputs - } - if ( - config.plasma.nlte_ionization_species - or config.plasma.nlte_excitation_species - ): - if config.plasma.nlte_ionization_species: - nlte_ionization_species = config.plasma.nlte_ionization_species - for species in nlte_ionization_species: - if species not in config.plasma.continuum_interaction.species: - raise PlasmaConfigError( - f"NLTE ionization species {species} not in continuum species." - ) - if config.plasma.nlte_excitation_species: - nlte_excitation_species = config.plasma.nlte_excitation_species - for species in nlte_excitation_species: - if species not in config.plasma.continuum_interaction.species: - raise PlasmaConfigError( - f"NLTE excitation species {species} not in continuum species." - ) - property_kwargs[NLTEIndexHelper] = { - "nlte_ionization_species": config.plasma.nlte_ionization_species, - "nlte_excitation_species": config.plasma.nlte_excitation_species, - } - if config.plasma.nlte_solver == "lu": - continuum_plasma_modules += nlte_lu_solver_properties - logger.warning( - "LU solver will be inaccurate for NLTE excitation, proceed with caution." - ) - elif config.plasma.nlte_solver == "root": - continuum_plasma_modules += nlte_root_solver_properties - else: - raise PlasmaConfigError( - f"NLTE solver type unknown - {config.plasma.nlte_solver}" - ) - - # initializing rates - t_electrons = ( - config.plasma.link_t_rad_t_electron - * dilute_planckian_radiation_field.temperature.to(u.K).value - ) - initial_continuum_solver = DiluteBlackBodyContinuumPropertiesSolver( - atom_data - ) - initial_continuum_properties = initial_continuum_solver.solve( - dilute_planckian_radiation_field, t_electrons - ) - kwargs.update( - gamma=initial_continuum_properties.photo_ionization_rate_coefficient, - bf_heating_coeff_estimator=None, - stim_recomb_cooling_coeff_estimator=None, - alpha_stim_factor=initial_continuum_properties.stimulated_recombination_rate_factor, + return plasma_solver_factory.assemble( + dilute_planckian_radiation_field, simulation_state ) - return continuum_plasma_modules From 930462fe0284b24c67e89a5014c40b432e6bac68 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 14:55:20 -0400 Subject: [PATCH 69/89] blackify --- .../continuum_processes/collisional_ion_trans_prob.py | 1 + tardis/plasma/properties/continuum_processes/rates.py | 1 + tardis/plasma/properties/plasma_input.py | 1 + tardis/simulation/base.py | 6 +++--- 4 files changed, 6 insertions(+), 3 deletions(-) diff --git a/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py b/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py index c3bde599279..6fd362303f1 100644 --- a/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py +++ b/tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py @@ -10,6 +10,7 @@ H = const.h.cgs.value + class RawCollIonTransProbs(TransitionProbabilitiesProperty, IndexSetterMixin): """ Attributes diff --git a/tardis/plasma/properties/continuum_processes/rates.py b/tardis/plasma/properties/continuum_processes/rates.py index 4dd697cfd64..f627debcd75 100644 --- a/tardis/plasma/properties/continuum_processes/rates.py +++ b/tardis/plasma/properties/continuum_processes/rates.py @@ -316,6 +316,7 @@ def calculate( ) return gamma_corr + class StimRecombCoolingRateCoeffEstimator(Input): """ Attributes diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index aac02e58308..5ed71b85f49 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -127,6 +127,7 @@ class LinkTRadTElectron(Input): class HeliumTreatment(Input): outputs = ("helium_treatment",) + class ContinuumInteractionSpecies(Input): """ Attributes diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index a805bafdf9e..599fe2b0caf 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -456,9 +456,9 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0): # Set up spectrum solver self.spectrum_solver.transport_state = transport_state - self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = ( - v_packets_energy_hist - ) + self.spectrum_solver._montecarlo_virtual_luminosity.value[ + : + ] = v_packets_energy_hist output_energy = ( self.transport.transport_state.packet_collection.output_energies From 9e47b81a79dfc47446ef15889be780b23c4ae02e Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 14:58:41 -0400 Subject: [PATCH 70/89] reverse the import pygraphviz --- tardis/plasma/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/plasma/base.py b/tardis/plasma/base.py index 8a42c32dfdb..87af2cbd479 100644 --- a/tardis/plasma/base.py +++ b/tardis/plasma/base.py @@ -293,7 +293,7 @@ def write_to_dot(self, fname, args=None, latex_label=True): edge labels into the file. """ try: - pass + import pygraphviz except: logger.warning( "pygraphviz missing. Plasma graph will not be " "generated." From 8d67db8f8f2bf274751c68235db540107df89515 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 14:59:54 -0400 Subject: [PATCH 71/89] fix docstrings --- tardis/plasma/radiation_field/planck_rad_field.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tardis/plasma/radiation_field/planck_rad_field.py b/tardis/plasma/radiation_field/planck_rad_field.py index 69d656bd0d6..ab3a4b741db 100644 --- a/tardis/plasma/radiation_field/planck_rad_field.py +++ b/tardis/plasma/radiation_field/planck_rad_field.py @@ -8,7 +8,7 @@ class DilutePlanckianRadiationField: """ - Represents the state of a dilute thermal radiation field. + Represents the state of a dilute planckian radiation field. Parameters @@ -79,7 +79,7 @@ def to_planckian_radiation_field(self): class PlanckianRadiationField: """ - Represents the state of a dilute thermal radiation field. + Represents the state of a planckian radiation field. Parameters From f6f3160efc1fa5692a3703073d0be6f492dbee23 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 15:54:29 -0400 Subject: [PATCH 72/89] fixed all plasma --- tardis/plasma/assembly/base.py | 35 ++++++++++++++++++++++------------ 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index 2b990027c36..a7994a8f676 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -64,6 +64,8 @@ class PlasmaSolverFactory: radiative_rates_type: str = "dilute-blackbody" + delta_treatment = None + ## Statistical Balance Solver legacy_nlte_species: list = [] @@ -123,21 +125,34 @@ def __init__(self, config, atom_data, atomic_numbers) -> None: self.link_t_rad_t_electron = config.plasma.link_t_rad_t_electron self.setup_analytical_approximations(config.plasma) + if "delta_treatment" in config.plasma: + self.property_kwargs[RadiationFieldCorrection] = dict( + delta_treatment=config.plasma.delta_treatment + ) self.setup_legacy_nlte(config.plasma.nlte) if self.line_interaction_type in ("downbranch", "macroatom") and ( len(self.continuum_interaction_species) == 0 ): self.setup_legacy_macro_atom() + self.delta_treatment = config.plasma.get("delta_treatment", None) + if self.delta_treatment is not None: + self.property_kwargs[RadiationFieldCorrection] = dict( + delta_treatment=config.plasma.delta_treatment + ) self.helium_treatment = config.plasma.helium_treatment self.heating_rate_data_file = config.plasma.heating_rate_data_file self.setup_helium_treatment() + + self.nlte_solver = config.plasma.nlte_solver + self.nlte_ionization_species = config.plasma.nlte_ionization_species + self.nlte_excitation_species = config.plasma.nlte_excitation_species if len(config.plasma.continuum_interaction.species) > 0: self.setup_continuum_interactions( config.plasma.continuum_interaction ) - + self.radiative_rates_type = config.plasma.radiative_rates_type def setup_helium_treatment(self): """ Set up the helium treatment for the plasma assembly. @@ -254,10 +269,6 @@ def setup_analytical_approximations(self, plasma_config): self.plasma_modules += lte_ionization_properties elif self.ionization_analytical_approximation == "nebular": self.plasma_modules += nebular_ionization_properties - if "delta_treatment" in plasma_config: - self.property_kwargs[RadiationFieldCorrection] = dict( - delta_treatment=plasma_config.delta_treatment - ) else: raise PlasmaConfigError( f'Invalid excitation analytical approximation. Configured as {self.ionization_analytical_approximation} but needs to be either "lte" or "nebular"' @@ -333,11 +344,11 @@ def set_nlte_species_from_string(self, nlte_species): self.legacy_nlte_species = map_species_from_string(nlte_species) def setup_continuum_interactions(self, config_continuum_interaction): - line_interaction_type = self.line_interaction_type - if line_interaction_type != "macroatom": + + if self.line_interaction_type != "macroatom": raise PlasmaConfigError( "Continuum interactions require line_interaction_type " - f"macroatom (instead of {line_interaction_type})." + f"macroatom (instead of {self.line_interaction_type})." ) self.plasma_modules += continuum_interaction_properties @@ -364,17 +375,17 @@ def setup_continuum_interactions(self, config_continuum_interaction): "inputs": transition_probabilities_outputs } if len(self.nlte_ionization_species + self.nlte_excitation_species) > 0: - if len(self.nlte_ionization_species) > 0: + if self.nlte_ionization_species: nlte_ionization_species = self.nlte_ionization_species for species in nlte_ionization_species: - if species not in self.continuum_interaction_species: + if species not in config_continuum_interaction.species: raise PlasmaConfigError( f"NLTE ionization species {species} not in continuum species." ) - if len(self.nlte_excitation_species) > 0: + if self.nlte_excitation_species: nlte_excitation_species = self.nlte_excitation_species for species in nlte_excitation_species: - if species not in self.continuum_interaction.species: + if species not in config_continuum_interaction.species: raise PlasmaConfigError( f"NLTE excitation species {species} not in continuum species." ) From 7f1c30107af90f4d048155ac80cc3ae3c1696348 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 16:07:42 -0400 Subject: [PATCH 73/89] slow fixes --- tardis/plasma/assembly/base.py | 13 +++++++++---- tardis/plasma/base.py | 2 ++ 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index a7994a8f676..b0e20cea9ae 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -5,6 +5,7 @@ from astropy import units as u from tardis.plasma import BasePlasma +from tardis.plasma.base import PlasmaSolverSettings from tardis.plasma.exceptions import PlasmaConfigError from tardis.plasma.properties import ( HeliumNumericalNLTE, @@ -224,13 +225,13 @@ def setup_legacy_nlte(self, nlte_config): """ Set up the non-LTE (NLTE) properties for the legacy species. - Parameters: - ----------- + Parameters + ---------- nlte_config : dict A dictionary containing the NLTE configuration. - Notes: - ------ + Notes + ----- This method adds the NLTE properties for the legacy species to the plasma modules. If there are no legacy NLTE species, it adds the non-NLTE properties instead. """ @@ -455,6 +456,9 @@ def assemble(self, dilute_planckian_radiation_field, simulation_state): j_blues = self.initialize_j_blues( dilute_planckian_radiation_field, self.atom_data.lines ) + plasma_solver_settings = PlasmaSolverSettings( + RADIATIVE_RATES_TYPE=self.radiative_rates_type + ) kwargs = dict( time_explosion=simulation_state.time_explosion, @@ -489,5 +493,6 @@ def assemble(self, dilute_planckian_radiation_field, simulation_state): return BasePlasma( plasma_properties=self.plasma_modules, property_kwargs=self.property_kwargs, + plasma_solver_settings=plasma_solver_settings, **kwargs, ) diff --git a/tardis/plasma/base.py b/tardis/plasma/base.py index b6b781a691d..fefc1ffce24 100644 --- a/tardis/plasma/base.py +++ b/tardis/plasma/base.py @@ -27,6 +27,7 @@ def __init__( self, plasma_properties, property_kwargs=None, + plasma_solver_settings=None, **kwargs, ): self.outputs_dict = {} @@ -34,6 +35,7 @@ def __init__( self.plasma_properties = self._init_properties( plasma_properties, property_kwargs, **kwargs ) + self.plasma_solver_settings = plasma_solver_settings self._build_graph() self.update(**kwargs) From e05aeca41bb55343868370a254c07b7403c61dd4 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 17:02:15 -0400 Subject: [PATCH 74/89] Refactor recomb_rate_coeff.py and test_continuum_property_solver.py --- .../properties/continuum_processes/recomb_rate_coeff.py | 1 - .../estimators/tests/test_continuum_property_solver.py | 8 ++++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py b/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py index a6ea2879d0b..86d6e3b68c3 100644 --- a/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py +++ b/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py @@ -4,7 +4,6 @@ from tardis.plasma.properties.base import Input, ProcessingPlasmaProperty from tardis.plasma.properties.continuum_processes.rates import C, H from tardis.transport.montecarlo.estimators.util import ( - bound_free_estimator_array2frame, integrate_array_by_blocks, ) diff --git a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py index eb831f26e57..ef561073a21 100644 --- a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py +++ b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py @@ -64,10 +64,10 @@ def test_continuum_estimators( ) continuum_plasma.update( - gamma_estimator=transport_state.radfield_mc_estimators.photo_ion_estimator, - alpha_stim_estimator=transport_state.radfield_mc_estimators.stim_recomb_estimator, - bf_heating_coeff_estimator=transport_state.radfield_mc_estimators.bf_heating_estimator, - stim_recomb_cooling_coeff_estimator=transport_state.radfield_mc_estimators.stim_recomb_cooling_estimator, + gamma=continuum_properties_mc.photo_ionization_rate_coefficient, + alpha_stim_factor=continuum_properties_mc.stimulated_recombination_rate_factor, + bf_heating_coeff_estimator=None, + stim_recomb_cooling_coeff_estimator=None, ) pdt.assert_frame_equal( From efdb5b1b6ad073d05cf0fcfe9adc4f27cd0cefca Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 17:40:59 -0400 Subject: [PATCH 75/89] fixing assembly --- tardis/plasma/assembly/base.py | 106 ++++++++++++++++----------------- 1 file changed, 53 insertions(+), 53 deletions(-) diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index b0e20cea9ae..7ac60dfcbfd 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -89,71 +89,60 @@ class PlasmaSolverFactory: kwargs: dict = {} property_kwargs: dict = {} - def __init__(self, config, atom_data, atomic_numbers) -> None: - - self.set_nlte_species_from_string(config.plasma.nlte.species) - self.set_continuum_interaction_species_from_string( - config.plasma.continuum_interaction.species - ) - self.line_interaction_type = config.plasma.line_interaction_type - + def __init__(self, config, atom_data, selected_atomic_numbers) -> None: + self.parse_plasma_config(config.plasma) self.atom_data = atom_data self.atom_data.prepare_atom_data( - atomic_numbers, + selected_atomic_numbers, line_interaction_type=config.plasma.line_interaction_type, continuum_interaction_species=self.continuum_interaction_species, nlte_species=self.legacy_nlte_species, ) - - #### THIS IS VERY BAD BUT FOR NOW IS CHICKEN/EGG - # Check if continuum interaction species are in selected_atoms - continuum_atoms = self.continuum_interaction_species.get_level_values( - "atomic_number" - ) - - continuum_atoms_in_selected_atoms = np.all( - continuum_atoms.isin(atom_data.selected_atomic_numbers) - ) - if not continuum_atoms_in_selected_atoms: - raise PlasmaConfigError( - "Not all continuum interaction species " - "belong to atoms that have been specified " - "in the configuration." - ) - ##### ---------------------------- + self.check_continuum_interaction_species() self.plasma_modules = basic_inputs + basic_properties - self.link_t_rad_t_electron = config.plasma.link_t_rad_t_electron - self.setup_analytical_approximations(config.plasma) - if "delta_treatment" in config.plasma: - self.property_kwargs[RadiationFieldCorrection] = dict( - delta_treatment=config.plasma.delta_treatment - ) + self.setup_analytical_approximations() + self.property_kwargs[RadiationFieldCorrection] = dict( + delta_treatment=self.delta_treatment + ) self.setup_legacy_nlte(config.plasma.nlte) + if self.line_interaction_type in ("downbranch", "macroatom") and ( len(self.continuum_interaction_species) == 0 ): - self.setup_legacy_macro_atom() - self.delta_treatment = config.plasma.get("delta_treatment", None) - if self.delta_treatment is not None: - self.property_kwargs[RadiationFieldCorrection] = dict( - delta_treatment=config.plasma.delta_treatment - ) + self.plasma_modules += macro_atom_properties - self.helium_treatment = config.plasma.helium_treatment - self.heating_rate_data_file = config.plasma.heating_rate_data_file self.setup_helium_treatment() - self.nlte_solver = config.plasma.nlte_solver - self.nlte_ionization_species = config.plasma.nlte_ionization_species - self.nlte_excitation_species = config.plasma.nlte_excitation_species if len(config.plasma.continuum_interaction.species) > 0: self.setup_continuum_interactions( config.plasma.continuum_interaction ) - self.radiative_rates_type = config.plasma.radiative_rates_type + + def parse_plasma_config(self, plasma_config): + self.set_continuum_interaction_species_from_string( + plasma_config.continuum_interaction.species + ) + self.set_nlte_species_from_string(plasma_config.nlte.species) + self.line_interaction_type = plasma_config.line_interaction_type + self.link_t_rad_t_electron = plasma_config.link_t_rad_t_electron + + self.excitation_analytical_approximation = plasma_config.excitation + self.ionization_analytical_approximation = plasma_config.ionization + self.delta_treatment = plasma_config.get("delta_treatment", None) + + self.helium_treatment = plasma_config.helium_treatment + self.heating_rate_data_file = plasma_config.heating_rate_data_file + + self.nlte_ionization_species = plasma_config.nlte_ionization_species + self.nlte_excitation_species = plasma_config.nlte_excitation_species + + self.nlte_solver = plasma_config.nlte_solver + + self.radiative_rates_type = plasma_config.radiative_rates_type + def setup_helium_treatment(self): """ Set up the helium treatment for the plasma assembly. @@ -201,10 +190,10 @@ def setup_helium_treatment(self): self.plasma_modules += helium_nlte_properties elif self.helium_treatment == "numerical-nlte": self.plasma_modules += helium_numerical_nlte_properties - if heating_rate_data_file in ["none", None]: + if self.heating_rate_data_file in ["none", None]: raise PlasmaConfigError("Heating rate data file not specified") self.property_kwargs[HeliumNumericalNLTE] = dict( - heating_rate_data_file=heating_rate_data_file + heating_rate_data_file=self.heating_rate_data_file ) else: # If nlte ionization species are present, we don't want to add the @@ -218,9 +207,6 @@ def setup_helium_treatment(self): else: self.plasma_modules += helium_lte_properties - def setup_legacy_macro_atom(self): - self.plasma_modules += macro_atom_properties - def setup_legacy_nlte(self, nlte_config): """ Set up the non-LTE (NLTE) properties for the legacy species. @@ -246,7 +232,7 @@ def setup_legacy_nlte(self, nlte_config): else: self.plasma_modules += non_nlte_properties - def setup_analytical_approximations(self, plasma_config): + def setup_analytical_approximations(self): """ Setup the analytical approximations for excitation and ionization. @@ -254,9 +240,6 @@ def setup_analytical_approximations(self, plasma_config): ------- None """ - self.excitation_analytical_approximation = plasma_config.excitation - self.ionization_analytical_approximation = plasma_config.ionization - if self.excitation_analytical_approximation == "lte": self.plasma_modules += lte_excitation_properties elif self.excitation_analytical_approximation == "dilute-lte": @@ -328,6 +311,23 @@ def set_continuum_interaction_species_from_string( continuum_interaction_species, names=["atomic_number", "ion_number"] ) + def check_continuum_interaction_species(self): + + continuum_atoms = self.continuum_interaction_species.get_level_values( + "atomic_number" + ) + + continuum_atoms_in_selected_atoms = np.all( + continuum_atoms.isin(self.atom_data.selected_atomic_numbers) + ) + + if not continuum_atoms_in_selected_atoms: + raise PlasmaConfigError( + "Not all continuum interaction species " + "belong to atoms that have been specified " + "in the configuration." + ) + def set_nlte_species_from_string(self, nlte_species): """ Sets the non-LTE species from a string representation. From 90783072088872f151b5ca44e6ba9e64d286ccb8 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 18:00:40 -0400 Subject: [PATCH 76/89] restructure the read in --- .../plasma/construction_simple_plasma.ipynb | 80 ++++++++----------- tardis/plasma/assembly/base.py | 7 +- tardis/plasma/assembly/legacy_assembly.py | 4 +- 3 files changed, 40 insertions(+), 51 deletions(-) diff --git a/docs/physics/plasma/construction_simple_plasma.ipynb b/docs/physics/plasma/construction_simple_plasma.ipynb index 44849be37d3..35ab4bdaec4 100644 --- a/docs/physics/plasma/construction_simple_plasma.ipynb +++ b/docs/physics/plasma/construction_simple_plasma.ipynb @@ -9,56 +9,31 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": 4, "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 1\u001b[0m plasma_solver_factory \u001b[38;5;241m=\u001b[39m \u001b[43mPlasmaSolverFactory\u001b[49m\u001b[43m(\u001b[49m\u001b[43matom_data\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mselected_atomic_numbers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n", + "\u001b[0;31mTypeError\u001b[0m: PlasmaSolverFactory.__init__() missing 1 required positional argument: 'atom_data'" + ] + } + ], "source": [ - "test = DilutePlanckianRadFieldInput()\n", - "test.set_value(d_radfield)" + "plasma_solver_factory = PlasmaSolverFactory(atom_data, selected_atomic_numbers=[1, 2])" ] }, { @@ -107,7 +93,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.5" + "version": "3.12.4" } }, "nbformat": 4, diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index 7ac60dfcbfd..d19503ce100 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -89,12 +89,13 @@ class PlasmaSolverFactory: kwargs: dict = {} property_kwargs: dict = {} - def __init__(self, config, atom_data, selected_atomic_numbers) -> None: - self.parse_plasma_config(config.plasma) + def __init__(self, atom_data, selected_atomic_numbers, config=None) -> None: + if config is not None: + self.parse_plasma_config(config.plasma) self.atom_data = atom_data self.atom_data.prepare_atom_data( selected_atomic_numbers, - line_interaction_type=config.plasma.line_interaction_type, + line_interaction_type=self.line_interaction_type, continuum_interaction_species=self.continuum_interaction_species, nlte_species=self.legacy_nlte_species, ) diff --git a/tardis/plasma/assembly/legacy_assembly.py b/tardis/plasma/assembly/legacy_assembly.py index 7e45266dffc..665cd4d6758 100644 --- a/tardis/plasma/assembly/legacy_assembly.py +++ b/tardis/plasma/assembly/legacy_assembly.py @@ -22,7 +22,9 @@ def assemble_plasma(config, simulation_state, atom_data=None): """ atomic_numbers = simulation_state.abundance.index plasma_solver_factory = PlasmaSolverFactory( - config, atom_data, atomic_numbers + atom_data, + atomic_numbers, + config, ) dilute_planckian_radiation_field = DilutePlanckianRadiationField( From 683e8118f022166f9f483b6d3ea8fd329b4d23c1 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 18:21:24 -0400 Subject: [PATCH 77/89] fixup plasma assemble to be clean --- .../plasma/construction_simple_plasma.ipynb | 72 ++++++++++++++----- tardis/plasma/assembly/base.py | 52 ++++++++------ tardis/plasma/assembly/legacy_assembly.py | 4 +- tardis/plasma/properties/general.py | 4 +- tardis/plasma/properties/nlte.py | 4 +- .../plasma/properties/property_collections.py | 1 - 6 files changed, 92 insertions(+), 45 deletions(-) diff --git a/docs/physics/plasma/construction_simple_plasma.ipynb b/docs/physics/plasma/construction_simple_plasma.ipynb index 35ab4bdaec4..00e94706acf 100644 --- a/docs/physics/plasma/construction_simple_plasma.ipynb +++ b/docs/physics/plasma/construction_simple_plasma.ipynb @@ -9,9 +9,46 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 1, "metadata": {}, - "outputs": [], + "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": "a4d91ae0803f463c96fe296874d0bacd", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Iterations: 0/? [00:00 1\u001b[0m plasma_solver_factory \u001b[38;5;241m=\u001b[39m \u001b[43mPlasmaSolverFactory\u001b[49m\u001b[43m(\u001b[49m\u001b[43matom_data\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mselected_atomic_numbers\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n", - "\u001b[0;31mTypeError\u001b[0m: PlasmaSolverFactory.__init__() missing 1 required positional argument: 'atom_data'" - ] - } - ], + "outputs": [], "source": [ "plasma_solver_factory = PlasmaSolverFactory(atom_data, selected_atomic_numbers=[1, 2])" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "plasma_solver_factory.assemble(d_radfield)" + ] + }, { "cell_type": "markdown", "metadata": {}, diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index d19503ce100..39b29e4ceb9 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -79,7 +79,11 @@ class PlasmaSolverFactory: heating_rate_data_file: str = "none" ## Continuum Interaction - continuum_interaction_species: pd.MultiIndex + continuum_interaction_species: pd.MultiIndex = pd.MultiIndex.from_tuples( + [], names=["atomic_number", "ion_number"] + ) + enable_adiabatic_cooling: bool = False + enable_two_photon_decay: bool = False ## Opacities line_interaction_type: str = "scatter" @@ -107,8 +111,8 @@ def __init__(self, atom_data, selected_atomic_numbers, config=None) -> None: self.property_kwargs[RadiationFieldCorrection] = dict( delta_treatment=self.delta_treatment ) - - self.setup_legacy_nlte(config.plasma.nlte) + if config is not None: + self.setup_legacy_nlte(config.plasma.nlte) if self.line_interaction_type in ("downbranch", "macroatom") and ( len(self.continuum_interaction_species) == 0 @@ -117,10 +121,8 @@ def __init__(self, atom_data, selected_atomic_numbers, config=None) -> None: self.setup_helium_treatment() - if len(config.plasma.continuum_interaction.species) > 0: - self.setup_continuum_interactions( - config.plasma.continuum_interaction - ) + if len(self.continuum_interaction_species) > 0: + self.setup_continuum_interactions() def parse_plasma_config(self, plasma_config): self.set_continuum_interaction_species_from_string( @@ -144,6 +146,13 @@ def parse_plasma_config(self, plasma_config): self.radiative_rates_type = plasma_config.radiative_rates_type + self.enable_adiabatic_cooling = ( + plasma_config.continuum_interaction.enable_adiabatic_cooling + ) + self.enable_two_photon_decay = ( + plasma_config.continuum_interaction.enable_two_photon_decay + ) + def setup_helium_treatment(self): """ Set up the helium treatment for the plasma assembly. @@ -345,7 +354,7 @@ def set_nlte_species_from_string(self, nlte_species): """ self.legacy_nlte_species = map_species_from_string(nlte_species) - def setup_continuum_interactions(self, config_continuum_interaction): + def setup_continuum_interactions(self): if self.line_interaction_type != "macroatom": raise PlasmaConfigError( @@ -356,10 +365,10 @@ def setup_continuum_interactions(self, config_continuum_interaction): self.plasma_modules += continuum_interaction_properties self.plasma_modules += continuum_interaction_inputs - if config_continuum_interaction.enable_adiabatic_cooling: + if self.enable_adiabatic_cooling: self.plasma_modules += adiabatic_cooling_properties - if config_continuum_interaction.enable_two_photon_decay: + if self.enable_two_photon_decay: self.plasma_modules += two_photon_properties transition_probabilities_outputs = [ @@ -380,14 +389,14 @@ def setup_continuum_interactions(self, config_continuum_interaction): if self.nlte_ionization_species: nlte_ionization_species = self.nlte_ionization_species for species in nlte_ionization_species: - if species not in config_continuum_interaction.species: + if species not in self.continuum_interaction_species: raise PlasmaConfigError( f"NLTE ionization species {species} not in continuum species." ) if self.nlte_excitation_species: nlte_excitation_species = self.nlte_excitation_species for species in nlte_excitation_species: - if species not in config_continuum_interaction.species: + if species not in self.continuum_interaction_species: raise PlasmaConfigError( f"NLTE excitation species {species} not in continuum species." ) @@ -453,7 +462,13 @@ def initialize_continuum_properties(self, dilute_planckian_radiation_field): ) return initial_continuum_properties - def assemble(self, dilute_planckian_radiation_field, simulation_state): + def assemble( + self, + number_densities, + dilute_planckian_radiation_field, + time_explosion, + electron_densities=None, + ): j_blues = self.initialize_j_blues( dilute_planckian_radiation_field, self.atom_data.lines ) @@ -462,10 +477,9 @@ def assemble(self, dilute_planckian_radiation_field, simulation_state): ) kwargs = dict( - time_explosion=simulation_state.time_explosion, + time_explosion=time_explosion, dilute_planckian_radiation_field=dilute_planckian_radiation_field, - abundance=simulation_state.abundance, - number_density=simulation_state.elemental_number_density, + number_density=number_densities, link_t_rad_t_electron=self.link_t_rad_t_electron, atomic_data=self.atom_data, j_blues=j_blues, @@ -485,10 +499,8 @@ def assemble(self, dilute_planckian_radiation_field, simulation_state): alpha_stim_factor=initial_continuum_properties.stimulated_recombination_rate_factor, ) - if simulation_state._electron_densities is not None: - electron_densities = pd.Series( - simulation_state._electron_densities.cgs.value - ) + if electron_densities is not None: + electron_densities = pd.Series(electron_densities.cgs.value) self.setup_electron_densities(electron_densities) kwargs["helium_treatment"] = self.helium_treatment return BasePlasma( diff --git a/tardis/plasma/assembly/legacy_assembly.py b/tardis/plasma/assembly/legacy_assembly.py index 665cd4d6758..a87c682586e 100644 --- a/tardis/plasma/assembly/legacy_assembly.py +++ b/tardis/plasma/assembly/legacy_assembly.py @@ -32,5 +32,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): ) return plasma_solver_factory.assemble( - dilute_planckian_radiation_field, simulation_state + simulation_state.elemental_number_density, + dilute_planckian_radiation_field, + simulation_state.time_explosion, ) diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index cb6e46ef650..7a66d16cc05 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -84,8 +84,8 @@ class SelectedAtoms(ProcessingPlasmaProperty): outputs = ("selected_atoms",) - def calculate(self, abundance): - return abundance.index + def calculate(self, number_density): + return number_density.index class ElectronTemperature(ProcessingPlasmaProperty): diff --git a/tardis/plasma/properties/nlte.py b/tardis/plasma/properties/nlte.py index f5a9a424046..a6e4d7ab230 100644 --- a/tardis/plasma/properties/nlte.py +++ b/tardis/plasma/properties/nlte.py @@ -29,7 +29,7 @@ class PreviousElectronDensities(PreviousIterationProperty): def set_initial_value(self, kwargs): initial_value = pd.Series( 1000000.0, - index=kwargs["abundance"].columns, + index=kwargs["number_density"].columns, ) self._set_initial_value(initial_value) @@ -47,6 +47,6 @@ def set_initial_value(self, kwargs): initial_value = pd.DataFrame( 1.0, index=kwargs["atomic_data"].lines.index, - columns=kwargs["abundance"].columns, + columns=kwargs["number_density"].columns, ) self._set_initial_value(initial_value) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index 4cf62169ddb..aef95cce730 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -18,7 +18,6 @@ class PlasmaPropertyCollection(list): [ DilutePlanckianRadField, DilutePlanckianRadField, - Abundance, NumberDensity, TimeExplosion, AtomicData, From d381f02450ec0aa5214857677bf71d05f7294d0a Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sat, 27 Jul 2024 18:56:21 -0400 Subject: [PATCH 78/89] fix shell info widget --- .../plasma/construction_simple_plasma.ipynb | 122 ++++++++++++------ tardis/visualization/widgets/shell_info.py | 2 +- .../widgets/tests/test_shell_info.py | 8 +- 3 files changed, 87 insertions(+), 45 deletions(-) diff --git a/docs/physics/plasma/construction_simple_plasma.ipynb b/docs/physics/plasma/construction_simple_plasma.ipynb index 00e94706acf..e0315f3f4c7 100644 --- a/docs/physics/plasma/construction_simple_plasma.ipynb +++ b/docs/physics/plasma/construction_simple_plasma.ipynb @@ -9,48 +9,12 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": 9, "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": "a4d91ae0803f463c96fe296874d0bacd", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Iterations: 0/? [00:00 1\u001b[0m \u001b[43mplasma_solver_factory\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43massemble\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnumber_densities\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43md_radfield\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m5\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mu\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mday\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/python/tardis/tardis/plasma/assembly/base.py:506\u001b[0m, in \u001b[0;36mPlasmaSolverFactory.assemble\u001b[0;34m(self, number_densities, dilute_planckian_radiation_field, time_explosion, electron_densities)\u001b[0m\n\u001b[1;32m 504\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msetup_electron_densities(electron_densities)\n\u001b[1;32m 505\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhelium_treatment\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhelium_treatment\n\u001b[0;32m--> 506\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mBasePlasma\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 507\u001b[0m \u001b[43m \u001b[49m\u001b[43mplasma_properties\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mplasma_modules\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 508\u001b[0m \u001b[43m \u001b[49m\u001b[43mproperty_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mproperty_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 509\u001b[0m \u001b[43m \u001b[49m\u001b[43mplasma_solver_settings\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mplasma_solver_settings\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 510\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 511\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/python/tardis/tardis/io/util.py:195\u001b[0m, in \u001b[0;36mHDFWriterMixin.__new__\u001b[0;34m(cls, *args, **kwargs)\u001b[0m\n\u001b[1;32m 193\u001b[0m instance \u001b[38;5;241m=\u001b[39m \u001b[38;5;28msuper\u001b[39m(HDFWriterMixin, \u001b[38;5;28mcls\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__new__\u001b[39m(\u001b[38;5;28mcls\u001b[39m)\n\u001b[1;32m 194\u001b[0m instance\u001b[38;5;241m.\u001b[39moptional_hdf_properties \u001b[38;5;241m=\u001b[39m []\n\u001b[0;32m--> 195\u001b[0m \u001b[43minstance\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 196\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m instance\n", + "File \u001b[0;32m~/python/tardis/tardis/plasma/base.py:39\u001b[0m, in \u001b[0;36mBasePlasma.__init__\u001b[0;34m(self, plasma_properties, property_kwargs, plasma_solver_settings, **kwargs)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplasma_properties \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_init_properties(\n\u001b[1;32m 36\u001b[0m plasma_properties, property_kwargs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m 37\u001b[0m )\n\u001b[1;32m 38\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplasma_solver_settings \u001b[38;5;241m=\u001b[39m plasma_solver_settings\n\u001b[0;32m---> 39\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_build_graph\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 40\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mupdate(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n", + "File \u001b[0;32m~/python/tardis/tardis/plasma/base.py:101\u001b[0m, in \u001b[0;36mBasePlasma._build_graph\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 99\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m \u001b[38;5;28minput\u001b[39m \u001b[38;5;129;01min\u001b[39;00m plasma_property\u001b[38;5;241m.\u001b[39minputs:\n\u001b[1;32m 100\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28minput\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutputs_dict:\n\u001b[0;32m--> 101\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m PlasmaMissingModule(\n\u001b[1;32m 102\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mModule \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mplasma_property\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m requires input \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 103\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28minput\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m which has not been added\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 104\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m to this plasma\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 105\u001b[0m )\n\u001b[1;32m 106\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 107\u001b[0m position \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutputs_dict[\u001b[38;5;28minput\u001b[39m]\u001b[38;5;241m.\u001b[39moutputs\u001b[38;5;241m.\u001b[39mindex(\u001b[38;5;28minput\u001b[39m)\n", + "\u001b[0;31mPlasmaMissingModule\u001b[0m: Module PartitionFunction requires input level_boltzmann_factor which has not been added to this plasma" + ] + } + ], + "source": [ + "plasma_solver_factory.assemble(number_densities, d_radfield, 5 * u.day)" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "[tardis.plasma.properties.plasma_input.DilutePlanckianRadField,\n", + " tardis.plasma.properties.plasma_input.DilutePlanckianRadField,\n", + " tardis.plasma.properties.plasma_input.NumberDensity,\n", + " tardis.plasma.properties.plasma_input.TimeExplosion,\n", + " tardis.plasma.properties.plasma_input.AtomicData,\n", + " tardis.plasma.properties.plasma_input.JBlues,\n", + " tardis.plasma.properties.plasma_input.LinkTRadTElectron,\n", + " tardis.plasma.properties.plasma_input.HeliumTreatment,\n", + " tardis.plasma.properties.plasma_input.ContinuumInteractionSpecies,\n", + " tardis.plasma.properties.plasma_input.NLTEIonizationSpecies,\n", + " tardis.plasma.properties.plasma_input.NLTEExcitationSpecies,\n", + " tardis.plasma.properties.plasma_input.TRadiative,\n", + " tardis.plasma.properties.plasma_input.DilutionFactor,\n", + " tardis.plasma.properties.general.BetaRadiation,\n", + " tardis.plasma.properties.atomic.Levels,\n", + " tardis.plasma.properties.atomic.Lines,\n", + " tardis.plasma.properties.partition_function.PartitionFunction,\n", + " tardis.plasma.properties.general.GElectron,\n", + " tardis.plasma.properties.atomic.IonizationData,\n", + " tardis.plasma.properties.atomic.LinesLowerLevelIndex,\n", + " tardis.plasma.properties.atomic.LinesUpperLevelIndex,\n", + " tardis.opacities.tau_sobolev.TauSobolev,\n", + " tardis.plasma.properties.radiative_properties.StimulatedEmissionFactor,\n", + " tardis.plasma.properties.general.SelectedAtoms,\n", + " tardis.plasma.properties.general.ElectronTemperature,\n", + " tardis.plasma.properties.partition_function.LevelBoltzmannFactorLTE,\n", + " tardis.plasma.properties.ion_population.PhiSahaLTE,\n", + " tardis.plasma.properties.level_population.LevelNumberDensity,\n", + " tardis.plasma.properties.ion_population.IonNumberDensity]" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "plasma_solver_factory.plasma_modules" ] }, { diff --git a/tardis/visualization/widgets/shell_info.py b/tardis/visualization/widgets/shell_info.py index 811386be4ae..2cb93c91dd6 100644 --- a/tardis/visualization/widgets/shell_info.py +++ b/tardis/visualization/widgets/shell_info.py @@ -190,7 +190,7 @@ def __init__(self, sim_model): super().__init__( sim_model.simulation_state.t_radiative, sim_model.simulation_state.dilution_factor, - sim_model.plasma.abundance, + sim_model.simulation_state.abundance, sim_model.plasma.number_density, sim_model.plasma.ion_number_density, sim_model.plasma.level_number_density, diff --git a/tardis/visualization/widgets/tests/test_shell_info.py b/tardis/visualization/widgets/tests/test_shell_info.py index 7a99374dbc5..df7f4f168d9 100644 --- a/tardis/visualization/widgets/tests/test_shell_info.py +++ b/tardis/visualization/widgets/tests/test_shell_info.py @@ -16,7 +16,7 @@ def base_shell_info(simulation_verysimple): return BaseShellInfo( simulation_verysimple.simulation_state.t_radiative, simulation_verysimple.simulation_state.dilution_factor, - simulation_verysimple.plasma.abundance, + simulation_verysimple.simulation_state.abundance, simulation_verysimple.plasma.number_density, simulation_verysimple.plasma.ion_number_density, simulation_verysimple.plasma.level_number_density, @@ -58,12 +58,14 @@ def test_element_count_data( ): element_count_data = base_shell_info.element_count(1) assert element_count_data.shape == ( - len(simulation_verysimple.plasma.abundance[shell_num - 1]), + len( + simulation_verysimple.simulation_state.abundance[shell_num - 1] + ), 2, ) assert np.allclose( element_count_data.iloc[:, -1].map(np.float64), - simulation_verysimple.plasma.abundance[shell_num - 1], + simulation_verysimple.simulation_state.abundance[shell_num - 1], ) @pytest.mark.parametrize(("atomic_num", "shell_num"), [(12, 1), (20, 20)]) From fcce9e0aeb70b0819e9e0eeca4fb51f6581ae957 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sun, 28 Jul 2024 00:30:08 -0400 Subject: [PATCH 79/89] fix the widgets --- .../plasma/construction_simple_plasma.ipynb | 228 +++++++++++++----- tardis/model/base.py | 1 + tardis/plasma/assembly/base.py | 70 +++--- tardis/plasma/assembly/legacy_assembly.py | 2 +- .../plasma/properties/radiative_properties.py | 2 + tardis/visualization/widgets/shell_info.py | 2 +- 6 files changed, 208 insertions(+), 97 deletions(-) diff --git a/docs/physics/plasma/construction_simple_plasma.ipynb b/docs/physics/plasma/construction_simple_plasma.ipynb index e0315f3f4c7..58023347add 100644 --- a/docs/physics/plasma/construction_simple_plasma.ipynb +++ b/docs/physics/plasma/construction_simple_plasma.ipynb @@ -9,18 +9,73 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 1, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "LOADING\n", + "Scalene extension successfully loaded. Note: Scalene currently only\n", + "supports CPU+GPU profiling inside Jupyter notebooks. For full Scalene\n", + "profiling, use the command line version. To profile in line mode, use\n", + "`%scrun [options] statement`. To profile in cell mode, use `%%scalene\n", + "[options]` followed by your code.\n", + "\n", + "NOTE: in Jupyter notebook on MacOS, Scalene cannot profile child\n", + "processes. Do not run to try Scalene with multiprocessing in Jupyter\n", + "Notebook.\n" + ] + }, + { + "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": "72ded6233c124eeba35259b3c035ab59", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Iterations: 0/? [00:00 1\u001b[0m \u001b[43mplasma_solver_factory\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43massemble\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnumber_densities\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43md_radfield\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m5\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mu\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mday\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/python/tardis/tardis/plasma/assembly/base.py:506\u001b[0m, in \u001b[0;36mPlasmaSolverFactory.assemble\u001b[0;34m(self, number_densities, dilute_planckian_radiation_field, time_explosion, electron_densities)\u001b[0m\n\u001b[1;32m 504\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msetup_electron_densities(electron_densities)\n\u001b[1;32m 505\u001b[0m kwargs[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhelium_treatment\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhelium_treatment\n\u001b[0;32m--> 506\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mBasePlasma\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 507\u001b[0m \u001b[43m \u001b[49m\u001b[43mplasma_properties\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mplasma_modules\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 508\u001b[0m \u001b[43m \u001b[49m\u001b[43mproperty_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mproperty_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 509\u001b[0m \u001b[43m \u001b[49m\u001b[43mplasma_solver_settings\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mplasma_solver_settings\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 510\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 511\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/python/tardis/tardis/io/util.py:195\u001b[0m, in \u001b[0;36mHDFWriterMixin.__new__\u001b[0;34m(cls, *args, **kwargs)\u001b[0m\n\u001b[1;32m 193\u001b[0m instance \u001b[38;5;241m=\u001b[39m \u001b[38;5;28msuper\u001b[39m(HDFWriterMixin, \u001b[38;5;28mcls\u001b[39m)\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__new__\u001b[39m(\u001b[38;5;28mcls\u001b[39m)\n\u001b[1;32m 194\u001b[0m instance\u001b[38;5;241m.\u001b[39moptional_hdf_properties \u001b[38;5;241m=\u001b[39m []\n\u001b[0;32m--> 195\u001b[0m \u001b[43minstance\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__init__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 196\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m instance\n", - "File \u001b[0;32m~/python/tardis/tardis/plasma/base.py:39\u001b[0m, in \u001b[0;36mBasePlasma.__init__\u001b[0;34m(self, plasma_properties, property_kwargs, plasma_solver_settings, **kwargs)\u001b[0m\n\u001b[1;32m 35\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplasma_properties \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_init_properties(\n\u001b[1;32m 36\u001b[0m plasma_properties, property_kwargs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs\n\u001b[1;32m 37\u001b[0m )\n\u001b[1;32m 38\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplasma_solver_settings \u001b[38;5;241m=\u001b[39m plasma_solver_settings\n\u001b[0;32m---> 39\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_build_graph\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 40\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mupdate(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n", - "File \u001b[0;32m~/python/tardis/tardis/plasma/base.py:101\u001b[0m, in \u001b[0;36mBasePlasma._build_graph\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 99\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m \u001b[38;5;28minput\u001b[39m \u001b[38;5;129;01min\u001b[39;00m plasma_property\u001b[38;5;241m.\u001b[39minputs:\n\u001b[1;32m 100\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28minput\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutputs_dict:\n\u001b[0;32m--> 101\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m PlasmaMissingModule(\n\u001b[1;32m 102\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mModule \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mplasma_property\u001b[38;5;241m.\u001b[39mname\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m requires input \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 103\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28minput\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m which has not been added\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 104\u001b[0m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m to this plasma\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m 105\u001b[0m )\n\u001b[1;32m 106\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 107\u001b[0m position \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutputs_dict[\u001b[38;5;28minput\u001b[39m]\u001b[38;5;241m.\u001b[39moutputs\u001b[38;5;241m.\u001b[39mindex(\u001b[38;5;28minput\u001b[39m)\n", - "\u001b[0;31mPlasmaMissingModule\u001b[0m: Module PartitionFunction requires input level_boltzmann_factor which has not been added to this plasma" + "name": "stdout", + "output_type": "stream", + "text": [ + "SCRUN MAGIC\n" ] + }, + { + "data": { + "text/html": [ + "\n", + " \n", + " " + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" } ], "source": [ - "plasma_solver_factory.assemble(number_densities, d_radfield, 5 * u.day)" + "%scrun plasma_solver = plasma_solver_factory.assemble(number_densities, d_radfield, 5 * u.day)" ] }, { "cell_type": "code", - "execution_count": 16, + "execution_count": 36, "metadata": {}, "outputs": [ { "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
0
atomic_numberion_number
102.245918e-04
13.011070e+08
201.405189e-15
12.206015e-03
23.011070e+08
\n", + "
" + ], "text/plain": [ - "[tardis.plasma.properties.plasma_input.DilutePlanckianRadField,\n", - " tardis.plasma.properties.plasma_input.DilutePlanckianRadField,\n", - " tardis.plasma.properties.plasma_input.NumberDensity,\n", - " tardis.plasma.properties.plasma_input.TimeExplosion,\n", - " tardis.plasma.properties.plasma_input.AtomicData,\n", - " tardis.plasma.properties.plasma_input.JBlues,\n", - " tardis.plasma.properties.plasma_input.LinkTRadTElectron,\n", - " tardis.plasma.properties.plasma_input.HeliumTreatment,\n", - " tardis.plasma.properties.plasma_input.ContinuumInteractionSpecies,\n", - " tardis.plasma.properties.plasma_input.NLTEIonizationSpecies,\n", - " tardis.plasma.properties.plasma_input.NLTEExcitationSpecies,\n", - " tardis.plasma.properties.plasma_input.TRadiative,\n", - " tardis.plasma.properties.plasma_input.DilutionFactor,\n", - " tardis.plasma.properties.general.BetaRadiation,\n", - " tardis.plasma.properties.atomic.Levels,\n", - " tardis.plasma.properties.atomic.Lines,\n", - " tardis.plasma.properties.partition_function.PartitionFunction,\n", - " tardis.plasma.properties.general.GElectron,\n", - " tardis.plasma.properties.atomic.IonizationData,\n", - " tardis.plasma.properties.atomic.LinesLowerLevelIndex,\n", - " tardis.plasma.properties.atomic.LinesUpperLevelIndex,\n", - " tardis.opacities.tau_sobolev.TauSobolev,\n", - " tardis.plasma.properties.radiative_properties.StimulatedEmissionFactor,\n", - " tardis.plasma.properties.general.SelectedAtoms,\n", - " tardis.plasma.properties.general.ElectronTemperature,\n", - " tardis.plasma.properties.partition_function.LevelBoltzmannFactorLTE,\n", - " tardis.plasma.properties.ion_population.PhiSahaLTE,\n", - " tardis.plasma.properties.level_population.LevelNumberDensity,\n", - " tardis.plasma.properties.ion_population.IonNumberDensity]" + " 0\n", + "atomic_number ion_number \n", + "1 0 2.245918e-04\n", + " 1 3.011070e+08\n", + "2 0 1.405189e-15\n", + " 1 2.206015e-03\n", + " 2 3.011070e+08" ] }, - "execution_count": 16, + "execution_count": 36, "metadata": {}, "output_type": "execute_result" } ], "source": [ - "plasma_solver_factory.plasma_modules" + "plasma_solver.ion_number_density" ] }, { diff --git a/tardis/model/base.py b/tardis/model/base.py index d3e1ce6cee4..27e11d2f20d 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -95,6 +95,7 @@ class SimulationState(HDFWriterMixin): "density", "r_inner", "time_explosion", + "abundance", ] hdf_name = "simulation_state" diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index 39b29e4ceb9..90b366f3a4d 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -103,28 +103,21 @@ def __init__(self, atom_data, selected_atomic_numbers, config=None) -> None: continuum_interaction_species=self.continuum_interaction_species, nlte_species=self.legacy_nlte_species, ) - self.check_continuum_interaction_species() - - self.plasma_modules = basic_inputs + basic_properties - - self.setup_analytical_approximations() - self.property_kwargs[RadiationFieldCorrection] = dict( - delta_treatment=self.delta_treatment - ) - if config is not None: - self.setup_legacy_nlte(config.plasma.nlte) - if self.line_interaction_type in ("downbranch", "macroatom") and ( - len(self.continuum_interaction_species) == 0 - ): - self.plasma_modules += macro_atom_properties + def parse_plasma_config(self, plasma_config): + """ + Parse the plasma configuration. - self.setup_helium_treatment() + Parameters + ---------- + plasma_config : PlasmaConfig + The plasma configuration object containing the plasma parameters. - if len(self.continuum_interaction_species) > 0: - self.setup_continuum_interactions() + Returns + ------- + None - def parse_plasma_config(self, plasma_config): + """ self.set_continuum_interaction_species_from_string( plasma_config.continuum_interaction.species ) @@ -153,6 +146,30 @@ def parse_plasma_config(self, plasma_config): plasma_config.continuum_interaction.enable_two_photon_decay ) + def setup_factory(self, config=None): + self.check_continuum_interaction_species() + + self.plasma_modules = basic_inputs + basic_properties + + self.setup_analytical_approximations() + self.property_kwargs[RadiationFieldCorrection] = dict( + delta_treatment=self.delta_treatment + ) + if (config is not None) and len(self.legacy_nlte_species) > 0: + self.setup_legacy_nlte(config.plasma.nlte) + else: + self.plasma_modules += non_nlte_properties + + if self.line_interaction_type in ("downbranch", "macroatom") and ( + len(self.continuum_interaction_species) == 0 + ): + self.plasma_modules += macro_atom_properties + + self.setup_helium_treatment() + + if len(self.continuum_interaction_species) > 0: + self.setup_continuum_interactions() + def setup_helium_treatment(self): """ Set up the helium treatment for the plasma assembly. @@ -231,16 +248,13 @@ def setup_legacy_nlte(self, nlte_config): This method adds the NLTE properties for the legacy species to the plasma modules. If there are no legacy NLTE species, it adds the non-NLTE properties instead. """ - if len(self.legacy_nlte_species) > 0: - self.plasma_modules += nlte_properties - self.plasma_modules.append( - LevelBoltzmannFactorNLTE.from_config(nlte_config) - ) - self.property_kwargs[StimulatedEmissionFactor] = dict( - nlte_species=self.legacy_nlte_species - ) - else: - self.plasma_modules += non_nlte_properties + self.plasma_modules += nlte_properties + self.plasma_modules.append( + LevelBoltzmannFactorNLTE.from_config(nlte_config) + ) + self.property_kwargs[StimulatedEmissionFactor] = dict( + nlte_species=self.legacy_nlte_species + ) def setup_analytical_approximations(self): """ diff --git a/tardis/plasma/assembly/legacy_assembly.py b/tardis/plasma/assembly/legacy_assembly.py index a87c682586e..916ebccaab7 100644 --- a/tardis/plasma/assembly/legacy_assembly.py +++ b/tardis/plasma/assembly/legacy_assembly.py @@ -26,7 +26,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): atomic_numbers, config, ) - + plasma_solver_factory.setup_factory(config) dilute_planckian_radiation_field = DilutePlanckianRadiationField( simulation_state.t_radiative, simulation_state.dilution_factor ) diff --git a/tardis/plasma/properties/radiative_properties.py b/tardis/plasma/properties/radiative_properties.py index a1b8d0a5d8e..2490a53683f 100644 --- a/tardis/plasma/properties/radiative_properties.py +++ b/tardis/plasma/properties/radiative_properties.py @@ -5,6 +5,7 @@ from astropy import units as u from numba import jit, prange + from tardis import constants as const from tardis.opacities.macro_atom.base import TransitionProbabilities from tardis.plasma.properties.base import ( @@ -12,6 +13,7 @@ TransitionProbabilitiesProperty, ) + logger = logging.getLogger(__name__) __all__ = [ diff --git a/tardis/visualization/widgets/shell_info.py b/tardis/visualization/widgets/shell_info.py index 2cb93c91dd6..8fb011609d6 100644 --- a/tardis/visualization/widgets/shell_info.py +++ b/tardis/visualization/widgets/shell_info.py @@ -216,7 +216,7 @@ def __init__(self, hdf_fpath): super().__init__( sim_data["/simulation/simulation_state/t_radiative"], sim_data["/simulation/simulation_state/dilution_factor"], - sim_data["/simulation/plasma/abundance"], + sim_data["/simulation/simulation_state/abundance"], sim_data["/simulation/plasma/number_density"], sim_data["/simulation/plasma/ion_number_density"], sim_data["/simulation/plasma/level_number_density"], From ac141ec3085f1734b2aec277fcf02fc8cf0e0aad Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sun, 28 Jul 2024 10:04:22 -0400 Subject: [PATCH 80/89] fixing hopefully last bugs --- tardis/plasma/assembly/base.py | 29 +++++++++++++---------- tardis/plasma/assembly/legacy_assembly.py | 1 + 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index 90b366f3a4d..13354ea2f61 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -79,9 +79,7 @@ class PlasmaSolverFactory: heating_rate_data_file: str = "none" ## Continuum Interaction - continuum_interaction_species: pd.MultiIndex = pd.MultiIndex.from_tuples( - [], names=["atomic_number", "ion_number"] - ) + continuum_interaction_species: list = [] enable_adiabatic_cooling: bool = False enable_two_photon_decay: bool = False @@ -100,10 +98,17 @@ def __init__(self, atom_data, selected_atomic_numbers, config=None) -> None: self.atom_data.prepare_atom_data( selected_atomic_numbers, line_interaction_type=self.line_interaction_type, - continuum_interaction_species=self.continuum_interaction_species, + continuum_interaction_species=self.continuum_interaction_species_multi_index, nlte_species=self.legacy_nlte_species, ) + @property + def continuum_interaction_species_multi_index(self): + return pd.MultiIndex.from_tuples( + map_species_from_string(self.continuum_interaction_species), + names=["atomic_number", "ion_number"], + ) + def parse_plasma_config(self, plasma_config): """ Parse the plasma configuration. @@ -118,7 +123,7 @@ def parse_plasma_config(self, plasma_config): None """ - self.set_continuum_interaction_species_from_string( + self.continuum_interaction_species = ( plasma_config.continuum_interaction.species ) self.set_nlte_species_from_string(plasma_config.nlte.species) @@ -326,19 +331,17 @@ def set_continuum_interaction_species_from_string( ------- None """ - continuum_interaction_species = [ + self.continuum_interaction_species = [ species_string_to_tuple(species) for species in continuum_interaction_species ] - self.continuum_interaction_species = pd.MultiIndex.from_tuples( - continuum_interaction_species, names=["atomic_number", "ion_number"] - ) - def check_continuum_interaction_species(self): - continuum_atoms = self.continuum_interaction_species.get_level_values( - "atomic_number" + continuum_atoms = ( + self.continuum_interaction_species_multi_index.get_level_values( + "atomic_number" + ) ) continuum_atoms_in_selected_atoms = np.all( @@ -497,7 +500,7 @@ def assemble( link_t_rad_t_electron=self.link_t_rad_t_electron, atomic_data=self.atom_data, j_blues=j_blues, - continuum_interaction_species=self.continuum_interaction_species, + continuum_interaction_species=self.continuum_interaction_species_multi_index, nlte_ionization_species=self.nlte_ionization_species, nlte_excitation_species=self.nlte_excitation_species, ) diff --git a/tardis/plasma/assembly/legacy_assembly.py b/tardis/plasma/assembly/legacy_assembly.py index 916ebccaab7..6f12fde618b 100644 --- a/tardis/plasma/assembly/legacy_assembly.py +++ b/tardis/plasma/assembly/legacy_assembly.py @@ -35,4 +35,5 @@ def assemble_plasma(config, simulation_state, atom_data=None): simulation_state.elemental_number_density, dilute_planckian_radiation_field, simulation_state.time_explosion, + simulation_state._electron_densities, ) From 01b4df7711cbb752cfe564aed4329f2ec05e0e4f Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Sun, 28 Jul 2024 10:52:13 -0400 Subject: [PATCH 81/89] slowly fix up assembly --- tardis/plasma/assembly/base.py | 39 +++++++++++++++++++++++++++++----- 1 file changed, 34 insertions(+), 5 deletions(-) diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index 13354ea2f61..7afe3215b9f 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -152,6 +152,25 @@ def parse_plasma_config(self, plasma_config): ) def setup_factory(self, config=None): + """ + Set up the plasma factory. + + Parameters + ---------- + config : object, optional + Configuration object containing plasma settings (default: None). + + Notes + ----- + This method performs the necessary setup steps for the plasma factory. + It checks the continuum interaction species, sets up the plasma modules, + sets up analytical approximations, sets up radiation field correction, + sets up legacy NLTE species if present, sets up macro atom properties if + line interaction type is 'downbranch' or 'macroatom' and there are no + continuum interaction species, and sets up helium treatment. Finally, + it sets up continuum interactions if there are any. + + """ self.check_continuum_interaction_species() self.plasma_modules = basic_inputs + basic_properties @@ -247,11 +266,6 @@ def setup_legacy_nlte(self, nlte_config): ---------- nlte_config : dict A dictionary containing the NLTE configuration. - - Notes - ----- - This method adds the NLTE properties for the legacy species to the plasma modules. - If there are no legacy NLTE species, it adds the non-NLTE properties instead. """ self.plasma_modules += nlte_properties self.plasma_modules.append( @@ -337,7 +351,13 @@ def set_continuum_interaction_species_from_string( ] def check_continuum_interaction_species(self): + """ + Check if all continuum interaction species belong to atoms that have been specified in the configuration. + Raises + ------ + PlasmaConfigError: If not all continuum interaction species belong to specified atoms. + """ continuum_atoms = ( self.continuum_interaction_species_multi_index.get_level_values( "atomic_number" @@ -372,7 +392,16 @@ def set_nlte_species_from_string(self, nlte_species): self.legacy_nlte_species = map_species_from_string(nlte_species) def setup_continuum_interactions(self): + """ + Set up continuum interactions for the plasma assembly. + Raises + ------ + PlasmaConfigError: If the line_interaction_type is not "macroatom". + PlasmaConfigError: If an NLTE ionization species is not in the continuum species. + PlasmaConfigError: If an NLTE excitation species is not in the continuum species. + PlasmaConfigError: If the NLTE solver type is unknown. + """ if self.line_interaction_type != "macroatom": raise PlasmaConfigError( "Continuum interactions require line_interaction_type " From 4ba4afb2bff7058a77285b6f583c9a28e721209d Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 29 Jul 2024 13:31:52 -0400 Subject: [PATCH 82/89] working on getting the notebook running --- .../plasma/construction_simple_plasma.ipynb | 39 +++++++------------ tardis/plasma/assembly/base.py | 12 ------ 2 files changed, 13 insertions(+), 38 deletions(-) diff --git a/docs/physics/plasma/construction_simple_plasma.ipynb b/docs/physics/plasma/construction_simple_plasma.ipynb index 58023347add..a8719bb41cc 100644 --- a/docs/physics/plasma/construction_simple_plasma.ipynb +++ b/docs/physics/plasma/construction_simple_plasma.ipynb @@ -4,7 +4,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Constructing a simple plasma" + "# Constructing a simple thermal plasma" ] }, { @@ -12,22 +12,6 @@ "execution_count": 1, "metadata": {}, "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "LOADING\n", - "Scalene extension successfully loaded. Note: Scalene currently only\n", - "supports CPU+GPU profiling inside Jupyter notebooks. For full Scalene\n", - "profiling, use the command line version. To profile in line mode, use\n", - "`%scrun [options] statement`. To profile in cell mode, use `%%scalene\n", - "[options]` followed by your code.\n", - "\n", - "NOTE: in Jupyter notebook on MacOS, Scalene cannot profile child\n", - "processes. Do not run to try Scalene with multiprocessing in Jupyter\n", - "Notebook.\n" - ] - }, { "name": "stderr", "output_type": "stream", @@ -39,7 +23,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "72ded6233c124eeba35259b3c035ab59", + "model_id": "adc8142e04de4ac4ad9fb6d9eba83286", "version_major": 2, "version_minor": 0 }, @@ -53,7 +37,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "32dafcfb559b481e902613c28442602a", + "model_id": "36b88ed7168e46c8864e17b0107f8f71", "version_major": 2, "version_minor": 0 }, @@ -69,13 +53,18 @@ "import numpy as np\n", "import pandas as pd\n", "from astropy import units as u\n", - "%load_ext scalene\n", + "\n", "from tardis.io.atom_data import AtomData\n", "from tardis.plasma.assembly.base import PlasmaSolverFactory\n", "from tardis.plasma.radiation_field import DilutePlanckianRadiationField\n", - "#from tardis.plasma.properties.plasma_input import T\n", - "\n", - "%load_ext pyinstrument\n" + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the default configuration the plasma solver assumes the conditions for a Local Thermodynamic Equilibrium meaning a planckian radiation field and a maxwellian electron velocity distribution at the same templerature" ] }, { @@ -154,9 +143,7 @@ "output_type": "display_data" } ], - "source": [ - "%scrun plasma_solver = plasma_solver_factory.assemble(number_densities, d_radfield, 5 * u.day)" - ] + "source": [] }, { "cell_type": "code", diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index 7afe3215b9f..163193e6b51 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -121,7 +121,6 @@ def parse_plasma_config(self, plasma_config): Returns ------- None - """ self.continuum_interaction_species = ( plasma_config.continuum_interaction.species @@ -159,17 +158,6 @@ def setup_factory(self, config=None): ---------- config : object, optional Configuration object containing plasma settings (default: None). - - Notes - ----- - This method performs the necessary setup steps for the plasma factory. - It checks the continuum interaction species, sets up the plasma modules, - sets up analytical approximations, sets up radiation field correction, - sets up legacy NLTE species if present, sets up macro atom properties if - line interaction type is 'downbranch' or 'macroatom' and there are no - continuum interaction species, and sets up helium treatment. Finally, - it sets up continuum interactions if there are any. - """ self.check_continuum_interaction_species() From 7d8e6e04d8830e6bee8d5cd3b659f9ce9d3bc6b8 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 29 Jul 2024 13:43:35 -0400 Subject: [PATCH 83/89] Refactor code to address comments --- tardis/energy_input/gamma_packet_loop.py | 4 +--- tardis/energy_input/gamma_ray_estimators.py | 4 +--- .../properties/continuum_processes/photo_ion_rate_coeff.py | 1 - .../properties/continuum_processes/recomb_rate_coeff.py | 5 +++-- tardis/plasma/properties/property_collections.py | 1 - tardis/transport/montecarlo/base.py | 2 +- ...lute_blackbody_properties.py => mc_rad_field_solver.py} | 0 tardis/transport/montecarlo/montecarlo_transport_state.py | 7 ------- 8 files changed, 6 insertions(+), 18 deletions(-) rename tardis/transport/montecarlo/estimators/{dilute_blackbody_properties.py => mc_rad_field_solver.py} (100%) diff --git a/tardis/energy_input/gamma_packet_loop.py b/tardis/energy_input/gamma_packet_loop.py index faac2164eef..385febda679 100644 --- a/tardis/energy_input/gamma_packet_loop.py +++ b/tardis/energy_input/gamma_packet_loop.py @@ -19,11 +19,9 @@ doppler_factor_3d, get_index, ) -from tardis.opacities.compton_opacity_calculation import ( - compton_opacity_calculation, -) from tardis.opacities.opacities import ( SIGMA_T, + compton_opacity_calculation, kappa_calculation, pair_creation_opacity_artis, pair_creation_opacity_calculation, diff --git a/tardis/energy_input/gamma_ray_estimators.py b/tardis/energy_input/gamma_ray_estimators.py index 13ebc9e84c1..e28ba5029ff 100644 --- a/tardis/energy_input/gamma_ray_estimators.py +++ b/tardis/energy_input/gamma_ray_estimators.py @@ -7,11 +7,9 @@ angle_aberration_gamma, doppler_factor_3d, ) -from tardis.opacities.compton_opacity_calculation import ( - compton_opacity_calculation, -) from tardis.opacities.opacities import ( SIGMA_T, + compton_opacity_calculation, kappa_calculation, photoabsorption_opacity_calculation, ) diff --git a/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py b/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py index c3bf1fff744..96e7d5096f6 100644 --- a/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py +++ b/tardis/plasma/properties/continuum_processes/photo_ion_rate_coeff.py @@ -1,5 +1,4 @@ from tardis.plasma.properties.base import Input -from tardis.plasma.properties.continuum_processes.rates import H class PhotoIonRateCoeff(Input): diff --git a/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py b/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py index a6ea2879d0b..16ddf8a7caf 100644 --- a/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py +++ b/tardis/plasma/properties/continuum_processes/recomb_rate_coeff.py @@ -1,13 +1,14 @@ import numpy as np import pandas as pd +import tardis.constants as const from tardis.plasma.properties.base import Input, ProcessingPlasmaProperty -from tardis.plasma.properties.continuum_processes.rates import C, H from tardis.transport.montecarlo.estimators.util import ( - bound_free_estimator_array2frame, integrate_array_by_blocks, ) +C = const.c.cgs.value + class StimRecombRateFactor(Input): """ diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index 4cf62169ddb..ea4c8ec37fb 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -16,7 +16,6 @@ class PlasmaPropertyCollection(list): basic_inputs = PlasmaPropertyCollection( [ - DilutePlanckianRadField, DilutePlanckianRadField, Abundance, NumberDensity, diff --git a/tardis/transport/montecarlo/base.py b/tardis/transport/montecarlo/base.py index efb6ced9e9d..242c5d3123b 100644 --- a/tardis/transport/montecarlo/base.py +++ b/tardis/transport/montecarlo/base.py @@ -11,7 +11,7 @@ MonteCarloConfiguration, configuration_initialize, ) -from tardis.transport.montecarlo.estimators.dilute_blackbody_properties import ( +from tardis.transport.montecarlo.estimators.mc_rad_field_solver import ( MCRadiationFieldPropertiesSolver, ) from tardis.transport.montecarlo.estimators.radfield_mc_estimators import ( diff --git a/tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py b/tardis/transport/montecarlo/estimators/mc_rad_field_solver.py similarity index 100% rename from tardis/transport/montecarlo/estimators/dilute_blackbody_properties.py rename to tardis/transport/montecarlo/estimators/mc_rad_field_solver.py diff --git a/tardis/transport/montecarlo/montecarlo_transport_state.py b/tardis/transport/montecarlo/montecarlo_transport_state.py index 324be096534..cc49707cc13 100644 --- a/tardis/transport/montecarlo/montecarlo_transport_state.py +++ b/tardis/transport/montecarlo/montecarlo_transport_state.py @@ -1,14 +1,8 @@ import warnings -import numpy as np from astropy import units as u from tardis.io.util import HDFWriterMixin -from tardis.transport.montecarlo.estimators.dilute_blackbody_properties import ( - MCRadiationFieldPropertiesSolver, -) -from tardis.spectrum.formal_integral import IntegrationError -from tardis.spectrum.spectrum import TARDISSpectrum class MonteCarloTransportState(HDFWriterMixin): @@ -63,7 +57,6 @@ def __init__( rpacket_tracker=None, vpacket_tracker=None, ): - self.time_explosion = time_explosion self.packet_collection = packet_collection self.radfield_mc_estimators = radfield_mc_estimators self.enable_full_relativity = False From 9f93d847ed8ee5f57f41a70d2d9f4e10f577be5c Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 30 Jul 2024 10:40:43 -0400 Subject: [PATCH 84/89] remove abundance --- tardis/plasma/properties/property_collections.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index ea4c8ec37fb..0a71bb984de 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -17,7 +17,6 @@ class PlasmaPropertyCollection(list): basic_inputs = PlasmaPropertyCollection( [ DilutePlanckianRadField, - Abundance, NumberDensity, TimeExplosion, AtomicData, From 31bfec695ac39fb54cffe5ede9f467e349e0c118 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 6 Aug 2024 14:14:23 +0200 Subject: [PATCH 85/89] added pydantic --- tardis/plasma/assembly/base.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index 163193e6b51..c3dc7f3d87c 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -3,7 +3,9 @@ import numpy as np import pandas as pd from astropy import units as u +from pydantic import BaseModel +from tardis.io.atom_data import AtomData from tardis.plasma import BasePlasma from tardis.plasma.base import PlasmaSolverSettings from tardis.plasma.exceptions import PlasmaConfigError @@ -54,18 +56,23 @@ def map_species_from_string(species): return [species_string_to_tuple(spec) for spec in species] -class PlasmaSolverFactory: +class PlasmaSolverFactory(BaseModel): + class Config: + arbitrary_types_allowed = True + extra = "allow" ## Analytical Approximations excitation_analytical_approximation: str = "lte" ionization_analytical_approximation: str = "lte" - nebular_ionization_delta_treatment: tuple # species to use for the delta_treatment in nebular ionization ML93 + nebular_ionization_delta_treatment: ( + tuple + ) = () # species to use for the delta_treatment in nebular ionization ML93 link_t_rad_t_electron: float = 1.0 radiative_rates_type: str = "dilute-blackbody" - delta_treatment = None + delta_treatment: float | None = None ## Statistical Balance Solver legacy_nlte_species: list = [] @@ -92,6 +99,7 @@ class PlasmaSolverFactory: property_kwargs: dict = {} def __init__(self, atom_data, selected_atomic_numbers, config=None) -> None: + super().__init__() if config is not None: self.parse_plasma_config(config.plasma) self.atom_data = atom_data From 88aff491bf6147dadaf7318e9df9e67a33b1ed96 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 6 Aug 2024 19:22:51 -0400 Subject: [PATCH 86/89] possible fix for test problem Co-authored-by: Atharva Arya --- tardis/plasma/assembly/base.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index c3dc7f3d87c..96d0328855d 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -3,9 +3,7 @@ import numpy as np import pandas as pd from astropy import units as u -from pydantic import BaseModel -from tardis.io.atom_data import AtomData from tardis.plasma import BasePlasma from tardis.plasma.base import PlasmaSolverSettings from tardis.plasma.exceptions import PlasmaConfigError @@ -56,10 +54,7 @@ def map_species_from_string(species): return [species_string_to_tuple(spec) for spec in species] -class PlasmaSolverFactory(BaseModel): - class Config: - arbitrary_types_allowed = True - extra = "allow" +class PlasmaSolverFactory: ## Analytical Approximations excitation_analytical_approximation: str = "lte" @@ -99,7 +94,10 @@ class Config: property_kwargs: dict = {} def __init__(self, atom_data, selected_atomic_numbers, config=None) -> None: - super().__init__() + self.plasma_modules = [] + self.kwargs = {} + self.property_kwargs = {} + if config is not None: self.parse_plasma_config(config.plasma) self.atom_data = atom_data From dcd157fa9bded34347af742417a4a06fba0960e0 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Tue, 6 Aug 2024 20:42:55 -0400 Subject: [PATCH 87/89] fix problem with opacity solver --- tardis/opacities/opacity_solver.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tardis/opacities/opacity_solver.py b/tardis/opacities/opacity_solver.py index 9b3b63f48ad..d5e821b99a7 100644 --- a/tardis/opacities/opacity_solver.py +++ b/tardis/opacities/opacity_solver.py @@ -46,7 +46,9 @@ def solve(self, legacy_plasma) -> OpacityState: legacy_plasma.atomic_data.lines.shape[ 0 ], # number of lines - legacy_plasma.abundance.shape[1], # number of shells + legacy_plasma.number_density.shape[ + 1 + ], # number of shells ), dtype=np.float64, ), From ed750c49711aed98c15fc962a685f7e5c09e353a Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 12 Aug 2024 09:24:34 -0400 Subject: [PATCH 88/89] chore: Refactor plasma assembly base and notebook Refactor the `PlasmaSolverFactory` class in `tardis/plasma/assembly/base.py` to improve code organization and readability. Remove unused imports and reformat the code using the black linter. Also, remove the unnecessary code block in the Jupyter notebook `docs/physics/plasma/construction_simple_plasma.ipynb`. These changes aim to enhance the maintainability and clarity of the codebase. --- .../plasma/construction_simple_plasma.ipynb | 36 ----------- tardis/plasma/assembly/base.py | 61 +++++++++++++++++++ 2 files changed, 61 insertions(+), 36 deletions(-) diff --git a/docs/physics/plasma/construction_simple_plasma.ipynb b/docs/physics/plasma/construction_simple_plasma.ipynb index a8719bb41cc..10b53802885 100644 --- a/docs/physics/plasma/construction_simple_plasma.ipynb +++ b/docs/physics/plasma/construction_simple_plasma.ipynb @@ -109,42 +109,6 @@ "plasma_solver_factory.setup_factory()\n" ] }, - { - "cell_type": "code", - "execution_count": 7, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "SCRUN MAGIC\n" - ] - }, - { - "data": { - "text/html": [ - "\n", - " \n", - " " - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [] - }, { "cell_type": "code", "execution_count": 36, diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index 96d0328855d..1e82a731e29 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -55,6 +55,67 @@ def map_species_from_string(species): class PlasmaSolverFactory: + """Factory class for creating plasma solvers. + + atom_data : object + Object containing atomic data. + selected_atomic_numbers : list + List of selected atomic numbers. + + Attributes + ---------- + excitation_analytical_approximation : str + Analytical approximation for excitation (default: "lte"). + ionization_analytical_approximation : str + Analytical approximation for ionization (default: "lte"). + nebular_ionization_delta_treatment : tuple + Species to use for the delta_treatment in nebular ionization ML93 (default: ()). + link_t_rad_t_electron : float + Link between t_rad and t_electron (default: 1.0). + radiative_rates_type : str + Type of radiative rates (default: "dilute-blackbody"). + delta_treatment : float or None + Delta treatment (default: None). + legacy_nlte_species : list + List of legacy non-LTE species (default: []). + nlte_excitation_species : list + List of non-LTE excitation species (default: []). + nlte_ionization_species : list + List of non-LTE ionization species (default: []). + nlte_solver : str + Non-LTE solver (default: "lu"). + Helium treatment options (default: "none"). + heating_rate_data_file : str + Heating rate data file (default: "none"). + continuum_interaction_species : list + List of continuum interaction species (default: []). + enable_adiabatic_cooling : bool + Flag for enabling adiabatic cooling (default: False). + enable_two_photon_decay : bool + Flag for enabling two-photon decay (default: False). + line_interaction_type : str + Type of line interaction (default: "scatter"). + plasma_modules : list + List of plasma modules (default: []). + kwargs : dict + Additional keyword arguments (default: {}). + property_kwargs : dict + Additional keyword arguments for properties (default: {}). + + Methods + ------- + parse_plasma_config(plasma_config) + continuum_interaction_species_multi_index() + Get the continuum interaction species as a multi-index. + setup_factory(config) + setup_helium_treatment() + setup_legacy_nlte(nlte_config) + Set up the non-LTE properties for the legacy species. + setup_analytical_approximations() + Set up the analytical approximations for excitation and ionization. + initialize_j_blues(dilute_planckian_radiation_field, lines_df) + Initialize j_blues. + """ ## Analytical Approximations excitation_analytical_approximation: str = "lte" From 554fa7e8ec18c3834b46fabefaa46e5d790c4675 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 12 Aug 2024 09:37:20 -0400 Subject: [PATCH 89/89] added docstrings --- tardis/plasma/assembly/base.py | 44 ++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index 1e82a731e29..7470ee08ea9 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -357,6 +357,26 @@ def setup_analytical_approximations(self): ) def initialize_j_blues(self, dilute_planckian_radiation_field, lines_df): + """ + Initialize the j_blues DataFrame based on the radiative_rates_type and the dilute_planckian_radiation_field. + + Parameters + ---------- + dilute_planckian_radiation_field : object + The dilute Planckian radiation field object. + lines_df : pandas.DataFrame + The DataFrame containing lines information. + + Returns + ------- + pandas.DataFrame + The DataFrame with calculated mean intensity values. + + Raises + ------ + ValueError + If the radiative_rates_type is unknown. + """ if (self.radiative_rates_type == "dilute-blackbody") or ( self.radiative_rates_type == "detailed" ): @@ -570,6 +590,30 @@ def assemble( time_explosion, electron_densities=None, ): + """ + Assemble the plasma based on the provided parameters and settings. + + Parameters + ---------- + number_densities : dict + Dictionary of number densities for different species. + dilute_planckian_radiation_field : object + The dilute Planckian radiation field object. + time_explosion : float + The time of explosion. + electron_densities : array-like, optional + Optional electron densities. + + Returns + ------- + BasePlasma + The assembled plasma object. + + Raises + ------ + ValueError + If an error occurs during assembly. + """ j_blues = self.initialize_j_blues( dilute_planckian_radiation_field, self.atom_data.lines )