From fd4369f1e8f7c8583301078ea67cc34d8e752727 Mon Sep 17 00:00:00 2001 From: Andrew Date: Mon, 2 Dec 2024 10:21:41 -0500 Subject: [PATCH] Remove tau sobolevs from plasma (#2874) * Create legacy plasma graph * Set up legacy plasma collections and factory * Legacy macro atom transition probs * Modifications to the formal integral, and simple workflow testbed * Partial fix for the formal integral * Drop macro atom data from plasma * Fixes formal integral tests * Resolve macro atom solver comments * Address opacity state and solver comments * Use restructured sobolev calculations for legacy plasma modules * Fixes formal integral test failures * Simplify workflow radiation field setup * Fixes standard workflow * Adds tests for the sobolev calculations * Fixes CUDA formal integral tests * Remove plasma atomic data variable assignment * Clean up macro atom a bit --- tardis/opacities/macro_atom/base.py | 10 -- .../opacities/macro_atom/macroatom_solver.py | 76 +++++++-- tardis/opacities/opacity_solver.py | 79 ++++++--- tardis/opacities/opacity_state.py | 40 ++++- tardis/opacities/tau_sobolev.py | 103 ++++++------ tardis/opacities/tests/test_opacity_solver.py | 3 +- .../tests/test_opacity_state_numba.py | 2 +- tardis/opacities/tests/test_tau_sobolev.py | 36 +++++ tardis/plasma/assembly/__init__.py | 1 + tardis/plasma/assembly/base.py | 93 ++++++----- tardis/plasma/assembly/legacy_assembly.py | 6 +- .../tests/test_collisional_transitions.py | 4 +- .../properties/legacy_property_collections.py | 150 ++++++++++++++++++ .../plasma/properties/property_collections.py | 27 +--- .../plasma/properties/radiative_properties.py | 48 +----- tardis/plasma/standard_plasmas.py | 12 +- tardis/simulation/base.py | 10 +- tardis/spectrum/formal_integral.py | 83 ++++++---- .../tests/test_cuda_formal_integral.py | 4 +- tardis/workflows/simple_tardis_workflow.py | 30 +++- tardis/workflows/standard_tardis_workflow.py | 17 +- 21 files changed, 565 insertions(+), 269 deletions(-) create mode 100644 tardis/opacities/tests/test_tau_sobolev.py create mode 100644 tardis/plasma/properties/legacy_property_collections.py diff --git a/tardis/opacities/macro_atom/base.py b/tardis/opacities/macro_atom/base.py index 2287e456567..8330c045323 100644 --- a/tardis/opacities/macro_atom/base.py +++ b/tardis/opacities/macro_atom/base.py @@ -142,7 +142,6 @@ def calculate_transition_probability( transition_probabilities = np.empty( (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 @@ -287,7 +286,6 @@ def _calculate_transition_probability( 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 @@ -304,14 +302,6 @@ def _calculate_transition_probability( ) 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 ): diff --git a/tardis/opacities/macro_atom/macroatom_solver.py b/tardis/opacities/macro_atom/macroatom_solver.py index 8f92602abf6..db2bf98f102 100644 --- a/tardis/opacities/macro_atom/macroatom_solver.py +++ b/tardis/opacities/macro_atom/macroatom_solver.py @@ -5,8 +5,7 @@ from tardis.opacities.macro_atom.macroatom_state import MacroAtomState -class MacroAtomSolver(object): - +class MacroAtomSolver: initialize: bool = True normalize: bool = True @@ -20,7 +19,6 @@ def __init__(self, initialize=True, normalize=True): normalize: bool Whether or not to normalize the transition probabilities to unity. Default True """ - self.initialize = initialize self.normalize = normalize @@ -32,7 +30,6 @@ def initialize_transition_probabilities(self, atomic_data): atomic_data : tardis.io.atom_data.AtomData Atomic Data """ - coef_and_block_ref = initialize_transition_probabilities(atomic_data) self.transition_probability_coef = coef_and_block_ref[ "transition_probability_coef" @@ -40,7 +37,7 @@ def initialize_transition_probabilities(self, atomic_data): self.block_references = coef_and_block_ref["block_references"] self.initialize = False - def solve_transition_probabilities( + def solve_legacy_transition_probabilities( self, atomic_data, legacy_plasma, @@ -80,12 +77,57 @@ def solve_transition_probabilities( return transition_probabilities + def solve_transition_probabilities( + self, + atomic_data, + mean_intensities, + tau_sobolev, + beta_sobolev, + stimulated_emission_factor, + ): + """Solve the basic transition probabilities for the macroatom + + Parameters + ---------- + atomic_data : tardis.io.atom_data.AtomData + Atomic Data + mean_intensities : pd.DataFrame + Mean intensity of the radiation field for each shell + tau_sobolev : pd.DataFrame + Expansion Optical Depths + beta_sobolev : pd.DataFrame + Modified expansion Optical Depths + stimulated_emission_factor : np.ndarray + + Returns + ------- + pd.DataFrame + Transition Probabilities + """ + if self.initialize: + self.initialize_transition_probabilities(atomic_data) + + transition_probabilities = calculate_transition_probabilities( + atomic_data, + beta_sobolev, + mean_intensities, + stimulated_emission_factor, + tau_sobolev, + self.transition_probability_coef, + self.block_references, + normalize=self.normalize, + ) + + return transition_probabilities + def solve( self, legacy_plasma, atomic_data, tau_sobolev, stimulated_emission_factor, + beta_sobolev=None, + legacy_mode=True, ): """Solved the Macro Atom State @@ -104,13 +146,23 @@ def solve( tardis.opacities.macroatom_state.MacroAtomState State of the macro atom ready to be placed into the OpacityState """ - - transition_probabilities = self.solve_transition_probabilities( - atomic_data, - legacy_plasma, - tau_sobolev, - stimulated_emission_factor, - ) + if legacy_mode: + transition_probabilities = ( + self.solve_legacy_transition_probabilities( + atomic_data, + legacy_plasma, + tau_sobolev, + stimulated_emission_factor, + ) + ) + else: + transition_probabilities = self.solve_transition_probabilities( + atomic_data, + legacy_plasma.j_blues, + tau_sobolev, + beta_sobolev, + stimulated_emission_factor, + ) macro_block_references = atomic_data.macro_atom_references[ "block_references" diff --git a/tardis/opacities/opacity_solver.py b/tardis/opacities/opacity_solver.py index 4561984bcd8..3b17a7c29fe 100644 --- a/tardis/opacities/opacity_solver.py +++ b/tardis/opacities/opacity_solver.py @@ -1,18 +1,23 @@ -from tardis.opacities.tau_sobolev import calculate_sobolev_line_opacity -from tardis.opacities.opacity_state import ( - OpacityState, -) import numpy as np import pandas as pd +from tardis.opacities.opacity_state import ( + OpacityState, +) +from tardis.opacities.tau_sobolev import ( + calculate_beta_sobolev, + calculate_sobolev_line_opacity, +) -class OpacitySolver(object): +class OpacitySolver: line_interaction_type: str = "scatter" disable_line_scattering: bool = False def __init__( - self, line_interaction_type="scatter", disable_line_scattering=False + self, + line_interaction_type="scatter", + disable_line_scattering=False, ): """Solver class for opacities @@ -22,11 +27,10 @@ def __init__( "scatter", "downbranch", or "macroatom" disable_line_scattering: bool """ - self.line_interaction_type = line_interaction_type self.disable_line_scattering = disable_line_scattering - def solve(self, legacy_plasma) -> OpacityState: + def legacy_solve(self, plasma) -> OpacityState: """ Solves the opacity state @@ -39,33 +43,64 @@ def solve(self, legacy_plasma) -> OpacityState: ------- OpacityState """ - atomic_data = legacy_plasma.atomic_data + if self.disable_line_scattering: + tau_sobolev = pd.DataFrame( + np.zeros( + ( + plasma.atomic_data.lines.shape[0], # number of lines + plasma.number_density.shape[1], # number of shells + ), + dtype=np.float64, + ), + index=plasma.atomic_data.lines.index, + ) + else: + tau_sobolev = calculate_sobolev_line_opacity( + plasma.atomic_data.lines, + plasma.level_number_density, + plasma.time_explosion, + plasma.stimulated_emission_factor, + ) + + opacity_state = OpacityState.from_legacy_plasma(plasma, tau_sobolev) + + return opacity_state + + def solve(self, plasma) -> OpacityState: + """ + Solves the opacity state + Parameters + ---------- + plasma : tarids.plasma.BasePlasma + legacy base plasma + + Returns + ------- + OpacityState + """ if self.disable_line_scattering: tau_sobolev = pd.DataFrame( np.zeros( ( - legacy_plasma.atomic_data.lines.shape[ - 0 - ], # number of lines - legacy_plasma.number_density.shape[ - 1 - ], # number of shells + plasma.atomic_data.lines.shape[0], # number of lines + plasma.number_density.shape[1], # number of shells ), dtype=np.float64, ), - index=atomic_data.lines.index, + index=plasma.atomic_data.lines.index, ) else: tau_sobolev = calculate_sobolev_line_opacity( - atomic_data.lines, - legacy_plasma.level_number_density, - legacy_plasma.time_explosion, - legacy_plasma.stimulated_emission_factor, + plasma.atomic_data.lines, + plasma.level_number_density, + plasma.time_explosion, + plasma.stimulated_emission_factor, ) + beta_sobolev = calculate_beta_sobolev(tau_sobolev) - opacity_state = OpacityState.from_legacy_plasma( - legacy_plasma, tau_sobolev + opacity_state = OpacityState.from_plasma( + plasma, tau_sobolev, beta_sobolev ) return opacity_state diff --git a/tardis/opacities/opacity_state.py b/tardis/opacities/opacity_state.py index 3442515419e..031a336481e 100644 --- a/tardis/opacities/opacity_state.py +++ b/tardis/opacities/opacity_state.py @@ -2,10 +2,10 @@ from numba import float64, int64 from numba.experimental import jitclass -from tardis.opacities.tau_sobolev import calculate_sobolev_line_opacity -from tardis.transport.montecarlo.configuration import montecarlo_globals from tardis.opacities.continuum.continuum_state import ContinuumState from tardis.opacities.macro_atom.macroatom_state import MacroAtomState +from tardis.opacities.tau_sobolev import calculate_sobolev_line_opacity +from tardis.transport.montecarlo.configuration import montecarlo_globals class OpacityState: @@ -15,6 +15,7 @@ def __init__( t_electrons, line_list_nu, tau_sobolev, + beta_sobolev, continuum_state, ): """ @@ -26,6 +27,7 @@ def __init__( t_electrons : numpy.ndarray line_list_nu : pd.DataFrame tau_sobolev : pd.DataFrame + beta_sobolev : pd.DataFrame continuum_state: tardis.opacities.continuum.continuum_state.ContinuumState """ self.electron_density = electron_density @@ -34,6 +36,8 @@ def __init__( self.tau_sobolev = tau_sobolev + self.beta_sobolev = beta_sobolev + # Continuum Opacity Data self.continuum_state = continuum_state @@ -53,7 +57,38 @@ def from_legacy_plasma(cls, plasma, tau_sobolev): ------- OpacityStatePython """ + if hasattr(plasma, "photo_ion_cross_sections"): + continuum_state = ContinuumState.from_legacy_plasma(plasma) + else: + continuum_state = None + + return cls( + plasma.electron_densities, + plasma.t_electrons, + plasma.atomic_data.lines.nu, + tau_sobolev, + plasma.beta_sobolev, + continuum_state, + ) + @classmethod + def from_plasma(cls, plasma, tau_sobolev, beta_sobolev): + """ + Generates an OpacityStatePython object from a tardis BasePlasma + + Parameters + ---------- + plasma : tarids.plasma.BasePlasma + legacy base plasma + tau_sobolev : pd.DataFrame + Expansion Optical Depths + beta_sobolev : pd.DataFrame + Modified expansion Optical Depths + + Returns + ------- + OpacityStatePython + """ if hasattr(plasma, "photo_ion_cross_sections"): continuum_state = ContinuumState.from_legacy_plasma(plasma) else: @@ -66,6 +101,7 @@ def from_legacy_plasma(cls, plasma, tau_sobolev): plasma.t_electrons, atomic_data.lines.nu, tau_sobolev, + beta_sobolev, continuum_state, ) diff --git a/tardis/opacities/tau_sobolev.py b/tardis/opacities/tau_sobolev.py index cf667f9f0ff..1206253732f 100644 --- a/tardis/opacities/tau_sobolev.py +++ b/tardis/opacities/tau_sobolev.py @@ -1,6 +1,7 @@ import numpy as np import pandas as pd from astropy import units as u +from numba import jit, prange from tardis import constants as const from tardis.plasma.properties.base import ProcessingPlasmaProperty @@ -72,6 +73,45 @@ def calculate_sobolev_line_opacity( ) +def calculate_beta_sobolev(tau_sobolevs): + """Calculate the beta Sobolev values based on the provided tau_sobolevs. + Values from the previous iteration can be provided. + + Parameters + ---------- + tau_sobolevs : pd.DataFrame + Tau Sobolev opacities. + + Returns + ------- + pd.DataFrame + The latest Beta Sobolev opacities. + """ + + @jit(nopython=True, parallel=True) + def numba_calculate_beta_sobolev(tau_sobolevs, beta_sobolevs): + for i in prange(len(tau_sobolevs)): + if tau_sobolevs[i] > 1e3: + beta_sobolevs[i] = tau_sobolevs[i] ** -1 + elif tau_sobolevs[i] < 1e-4: + beta_sobolevs[i] = 1 - 0.5 * tau_sobolevs[i] + else: + beta_sobolevs[i] = (1 - np.exp(-tau_sobolevs[i])) / ( + tau_sobolevs[i] + ) + return beta_sobolevs + + beta_sobolev = pd.DataFrame( + 0.0, index=tau_sobolevs.index, columns=tau_sobolevs.columns + ) + + numba_calculate_beta_sobolev( + tau_sobolevs.values.ravel(), beta_sobolev.values.ravel() + ) + + return beta_sobolev + + class TauSobolev(ProcessingPlasmaProperty): """ Attributes @@ -88,28 +128,12 @@ class TauSobolev(ProcessingPlasmaProperty): 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, - f_lu, - wavelength_cm, ): """ Calculate Sobolev line opacity. @@ -136,35 +160,24 @@ def calculate( ------ 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( - lines_lower_level_index, axis=0, mode="raise" - ) - tau_sobolevs = ( - self.sobolev_coefficient - * f_lu - * wavelength - * time_explosion - * n_lower - * stimulated_emission_factor + return calculate_sobolev_line_opacity( + lines, + level_number_density, + time_explosion, + 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 + ---------- + beta_sobolev : Numpy Array, dtype float + """ + + outputs = ("beta_sobolev",) + latex_name = (r"\beta_{\textrm{sobolev}}",) + + def calculate(self, tau_sobolevs): + return calculate_beta_sobolev(tau_sobolevs) diff --git a/tardis/opacities/tests/test_opacity_solver.py b/tardis/opacities/tests/test_opacity_solver.py index 3c7ac59c177..ce7ac53f938 100644 --- a/tardis/opacities/tests/test_opacity_solver.py +++ b/tardis/opacities/tests/test_opacity_solver.py @@ -21,14 +21,13 @@ def test_opacity_solver( nb_simulation_verysimple, line_interaction_type, disable_line_scattering ): - legacy_plasma = nb_simulation_verysimple.plasma opacity_solver = OpacitySolver( line_interaction_type=line_interaction_type, disable_line_scattering=disable_line_scattering, ) - actual = opacity_solver.solve(legacy_plasma) + actual = opacity_solver.legacy_solve(legacy_plasma) pdt.assert_series_equal( actual.electron_density, legacy_plasma.electron_densities diff --git a/tardis/opacities/tests/test_opacity_state_numba.py b/tardis/opacities/tests/test_opacity_state_numba.py index c826a638b77..167fd2673b6 100644 --- a/tardis/opacities/tests/test_opacity_state_numba.py +++ b/tardis/opacities/tests/test_opacity_state_numba.py @@ -25,7 +25,7 @@ def test_opacity_state_to_numba( line_interaction_type=line_interaction_type, disable_line_scattering=False, ) - opacity_state = opacity_solver.solve(legacy_plasma) + opacity_state = opacity_solver.legacy_solve(legacy_plasma) if line_interaction_type in ("downbranch", "macroatom"): macro_atom_state = MacroAtomSolver().solve( legacy_plasma, diff --git a/tardis/opacities/tests/test_tau_sobolev.py b/tardis/opacities/tests/test_tau_sobolev.py new file mode 100644 index 00000000000..5d16cef4cfa --- /dev/null +++ b/tardis/opacities/tests/test_tau_sobolev.py @@ -0,0 +1,36 @@ +import numpy.testing as npt +import pandas.testing as pdt + +from tardis.opacities.tau_sobolev import ( + calculate_beta_sobolev, + calculate_sobolev_line_opacity, +) + + +def test_calculate_sobolev_line_opacity( + nb_simulation_verysimple, regression_data +): + legacy_plasma = nb_simulation_verysimple.plasma + + actual = calculate_sobolev_line_opacity( + legacy_plasma.lines, + legacy_plasma.level_number_density, + legacy_plasma.time_explosion, + legacy_plasma.stimulated_emission_factor, + ) + expected = regression_data.sync_dataframe(actual) + pdt.assert_frame_equal(actual, expected) + + +def test_calculate_beta_sobolevs(nb_simulation_verysimple, regression_data): + legacy_plasma = nb_simulation_verysimple.plasma + + tau_sobolevs = calculate_sobolev_line_opacity( + legacy_plasma.lines, + legacy_plasma.level_number_density, + legacy_plasma.time_explosion, + legacy_plasma.stimulated_emission_factor, + ) + actual = calculate_beta_sobolev(tau_sobolevs) + expected = regression_data.sync_ndarray(actual) + npt.assert_allclose(actual, expected) diff --git a/tardis/plasma/assembly/__init__.py b/tardis/plasma/assembly/__init__.py index e69de29bb2d..19081d171f2 100644 --- a/tardis/plasma/assembly/__init__.py +++ b/tardis/plasma/assembly/__init__.py @@ -0,0 +1 @@ +from tardis.plasma.assembly.base import PlasmaSolverFactory diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py index f661a394b68..68ee1c7df02 100644 --- a/tardis/plasma/assembly/base.py +++ b/tardis/plasma/assembly/base.py @@ -1,3 +1,4 @@ +import importlib import logging import numpy as np @@ -22,26 +23,6 @@ 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.transport.montecarlo.estimators.continuum_radfield_properties import ( DiluteBlackBodyContinuumPropertiesSolver, @@ -97,6 +78,9 @@ class PlasmaSolverFactory: plasma_modules: list = [] kwargs: dict = {} property_kwargs: dict = {} + plasma_collection = importlib.import_module( + "tardis.plasma.properties.legacy_property_collections" + ) def __init__( self, @@ -154,15 +138,23 @@ def parse_plasma_config(self, plasma_config): plasma_config.continuum_interaction.enable_two_photon_decay ) - def prepare_factory(self, selected_atomic_numbers, config=None): + def prepare_factory( + self, selected_atomic_numbers, property_collections, config=None + ): """ Set up the plasma factory. Parameters ---------- + selected_atomic_numbers : list of int + Selected atomic numbers in the simulation. + property_collections : str + The property collection module to be used in the plasma assembly. config : object, optional Configuration object containing plasma settings (default: None). """ + self.plasma_collection = importlib.import_module(property_collections) + self.atom_data.prepare_atom_data( selected_atomic_numbers, line_interaction_type=self.line_interaction_type, @@ -172,7 +164,10 @@ def prepare_factory(self, selected_atomic_numbers, config=None): self.check_continuum_interaction_species() - self.plasma_modules = basic_inputs + basic_properties + self.plasma_modules = ( + self.plasma_collection.basic_inputs + + self.plasma_collection.basic_properties + ) self.setup_analytical_approximations() self.property_kwargs[RadiationFieldCorrection] = dict( @@ -181,12 +176,12 @@ def prepare_factory(self, selected_atomic_numbers, config=None): 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 + self.plasma_modules += self.plasma_collection.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.plasma_modules += self.plasma_collection.macro_atom_properties self.setup_helium_treatment() @@ -237,9 +232,11 @@ def setup_helium_treatment(self): # 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 + self.plasma_modules += self.plasma_collection.helium_nlte_properties elif self.helium_treatment == "numerical-nlte": - self.plasma_modules += helium_numerical_nlte_properties + self.plasma_modules += ( + self.plasma_collection.helium_numerical_nlte_properties + ) if self.heating_rate_data_file in ["none", None]: raise PlasmaConfigError("Heating rate data file not specified") self.property_kwargs[HeliumNumericalNLTE] = dict( @@ -255,7 +252,9 @@ def setup_helium_treatment(self): ): self.plasma_modules.append(LevelNumberDensity) else: - self.plasma_modules += helium_lte_properties + self.plasma_modules += ( + self.plasma_collection.helium_lte_properties + ) def setup_legacy_nlte(self, nlte_config): """ @@ -266,7 +265,7 @@ def setup_legacy_nlte(self, nlte_config): nlte_config : dict A dictionary containing the NLTE configuration. """ - self.plasma_modules += nlte_properties + self.plasma_modules += self.plasma_collection.nlte_properties self.plasma_modules.append( LevelBoltzmannFactorNLTE.from_config(nlte_config) ) @@ -283,18 +282,26 @@ def setup_analytical_approximations(self): None """ if self.excitation_analytical_approximation == "lte": - self.plasma_modules += lte_excitation_properties + self.plasma_modules += ( + self.plasma_collection.lte_excitation_properties + ) elif self.excitation_analytical_approximation == "dilute-lte": - self.plasma_modules += dilute_lte_excitation_properties + self.plasma_modules += ( + self.plasma_collection.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": - self.plasma_modules += lte_ionization_properties + self.plasma_modules += ( + self.plasma_collection.lte_ionization_properties + ) elif self.ionization_analytical_approximation == "nebular": - self.plasma_modules += nebular_ionization_properties + self.plasma_modules += ( + self.plasma_collection.nebular_ionization_properties + ) else: raise PlasmaConfigError( f'Invalid excitation analytical approximation. Configured as {self.ionization_analytical_approximation} but needs to be either "lte" or "nebular"' @@ -427,14 +434,20 @@ def setup_continuum_interactions(self): f"macroatom (instead of {self.line_interaction_type})." ) - self.plasma_modules += continuum_interaction_properties - self.plasma_modules += continuum_interaction_inputs + self.plasma_modules += ( + self.plasma_collection.continuum_interaction_properties + ) + self.plasma_modules += ( + self.plasma_collection.continuum_interaction_inputs + ) if self.enable_adiabatic_cooling: - self.plasma_modules += adiabatic_cooling_properties + self.plasma_modules += ( + self.plasma_collection.adiabatic_cooling_properties + ) if self.enable_two_photon_decay: - self.plasma_modules += two_photon_properties + self.plasma_modules += self.plasma_collection.two_photon_properties transition_probabilities_outputs = [ plasma_property.transition_probabilities_outputs @@ -470,12 +483,16 @@ def setup_continuum_interactions(self): "nlte_excitation_species": self.nlte_excitation_species, } if self.nlte_solver == "lu": - self.plasma_modules += nlte_lu_solver_properties + self.plasma_modules += ( + self.plasma_collection.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 + self.plasma_modules += ( + self.plasma_collection.nlte_root_solver_properties + ) else: raise PlasmaConfigError( f"NLTE solver type unknown - {self.nlte_solver}" diff --git a/tardis/plasma/assembly/legacy_assembly.py b/tardis/plasma/assembly/legacy_assembly.py index b250f1cbeda..58234a71fc7 100644 --- a/tardis/plasma/assembly/legacy_assembly.py +++ b/tardis/plasma/assembly/legacy_assembly.py @@ -25,7 +25,11 @@ def assemble_plasma(config, simulation_state, atom_data=None): atom_data, config, ) - plasma_solver_factory.prepare_factory(atomic_numbers, config) + plasma_solver_factory.prepare_factory( + atomic_numbers, + "tardis.plasma.properties.legacy_property_collections", + config, + ) dilute_planckian_radiation_field = DilutePlanckianRadiationField( simulation_state.t_radiative, simulation_state.dilution_factor ) diff --git a/tardis/plasma/detailed_balance/tests/test_collisional_transitions.py b/tardis/plasma/detailed_balance/tests/test_collisional_transitions.py index 9b5d80ea574..ed073729074 100644 --- a/tardis/plasma/detailed_balance/tests/test_collisional_transitions.py +++ b/tardis/plasma/detailed_balance/tests/test_collisional_transitions.py @@ -46,7 +46,9 @@ def legacy_cmfgen_collision_rate_plasma_solver(nlte_atomic_dataset): # plasma_solver_factory.continuum_interaction_species = ["He I"] plasma_solver_factory.line_interaction_type = "macroatom" - plasma_solver_factory.prepare_factory([1]) + plasma_solver_factory.prepare_factory( + [1], "tardis.plasma.properties.legacy_property_collections" + ) plasma_solver_factory.plasma_modules += [ YgData, ContinuumInteractionSpecies, diff --git a/tardis/plasma/properties/legacy_property_collections.py b/tardis/plasma/properties/legacy_property_collections.py new file mode 100644 index 00000000000..573c31c95c3 --- /dev/null +++ b/tardis/plasma/properties/legacy_property_collections.py @@ -0,0 +1,150 @@ +from tardis.opacities.continuum.bound_free import BoundFreeOpacity +from tardis.opacities.macro_atom.base import ( + NonMarkovChainTransitionProbabilities, + TransitionProbabilities, +) +from tardis.opacities.macro_atom.continuum_processes.collisional_ion_trans_prob import ( + RawCollIonTransProbs, +) +from tardis.opacities.tau_sobolev import BetaSobolev, TauSobolev +from tardis.plasma.properties import * + + +class PlasmaPropertyCollection(list): + pass + + +basic_inputs = PlasmaPropertyCollection( + [ + DilutePlanckianRadField, + NumberDensity, + TimeExplosion, + AtomicData, + JBlues, + LinkTRadTElectron, + HeliumTreatment, + ContinuumInteractionSpecies, + NLTEIonizationSpecies, + NLTEExcitationSpecies, + ] +) +basic_properties = PlasmaPropertyCollection( + [ + TRadiative, + DilutionFactor, + BetaRadiation, + Levels, + Lines, + PartitionFunction, + GElectron, + IonizationData, + LinesLowerLevelIndex, + LinesUpperLevelIndex, + TauSobolev, + StimulatedEmissionFactor, + SelectedAtoms, + ElectronTemperature, + ] +) +lte_ionization_properties = PlasmaPropertyCollection([PhiSahaLTE]) +lte_excitation_properties = PlasmaPropertyCollection([LevelBoltzmannFactorLTE]) +macro_atom_properties = PlasmaPropertyCollection( + [BetaSobolev, TransitionProbabilities, MacroAtomData] +) +nebular_ionization_properties = PlasmaPropertyCollection( + [PhiSahaNebular, ZetaData, BetaElectron, RadiationFieldCorrection] +) +dilute_lte_excitation_properties = PlasmaPropertyCollection( + [LevelBoltzmannFactorDiluteLTE] +) +non_nlte_properties = PlasmaPropertyCollection([LevelBoltzmannFactorNoNLTE]) +nlte_properties = PlasmaPropertyCollection( + [ + LevelBoltzmannFactorNLTE, + NLTEData, + PreviousElectronDensities, + PreviousBetaSobolev, + BetaSobolev, + ] +) +nlte_root_solver_properties = PlasmaPropertyCollection( + [NLTEIndexHelper, NLTEPopulationSolverRoot] +) +nlte_lu_solver_properties = PlasmaPropertyCollection( + [NLTEIndexHelper, NLTEPopulationSolverLU] +) +helium_nlte_properties = PlasmaPropertyCollection( + [ + HeliumNLTE, + RadiationFieldCorrection, + ZetaData, + BetaElectron, + LevelNumberDensityHeNLTE, + IonNumberDensityHeNLTE, + ] +) +helium_lte_properties = PlasmaPropertyCollection( + [LevelNumberDensity, IonNumberDensity] +) +helium_numerical_nlte_properties = PlasmaPropertyCollection( + [HeliumNumericalNLTE] +) +continuum_interaction_inputs = PlasmaPropertyCollection( + [ + PhotoIonRateCoeff, + StimRecombRateFactor, + BfHeatingRateCoeffEstimator, + StimRecombCoolingRateCoeffEstimator, + YgData, + ] +) +continuum_interaction_properties = PlasmaPropertyCollection( + [ + StimRecombRateCoeff, + PhotoIonizationData, + SpontRecombRateCoeff, + ThermalLevelBoltzmannFactorLTE, + ThermalLTEPartitionFunction, + BetaElectron, + ThermalGElectron, + ThermalPhiSahaLTE, + SahaFactor, + CorrPhotoIonRateCoeff, + SpontRecombCoolingRateCoeff, + RawRecombTransProbs, + RawPhotoIonTransProbs, + RawRadBoundBoundTransProbs, + MarkovChainTransProbs, + NonContinuumTransProbsMask, + YgInterpolator, + CollExcRateCoeff, + CollDeexcRateCoeff, + RawCollisionTransProbs, + MarkovChainIndex, + MarkovChainTransProbsCollector, + NonMarkovChainTransitionProbabilities, + MonteCarloTransProbs, + FreeFreeCoolingRate, + FreeBoundCoolingRate, + BoundFreeOpacity, + LevelNumberDensityLTE, + PhotoIonBoltzmannFactor, + FreeBoundEmissionCDF, + LevelIdxs2LineIdx, + LevelIdxs2TransitionIdx, + CollIonRateCoeffSeaton, + CollRecombRateCoeff, + RawCollIonTransProbs, + ContinuumInteractionHandler, + BetaSobolev, + ] +) +adiabatic_cooling_properties = PlasmaPropertyCollection([AdiabaticCoolingRate]) +two_photon_properties = PlasmaPropertyCollection( + [ + RawTwoPhotonTransProbs, + TwoPhotonData, + TwoPhotonEmissionCDF, + TwoPhotonFrequencySampler, + ] +) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index 0a71bb984de..940f265448c 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -1,12 +1,3 @@ -from tardis.opacities.continuum.bound_free import BoundFreeOpacity -from tardis.opacities.macro_atom.base import ( - NonMarkovChainTransitionProbabilities, - TransitionProbabilities, -) -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 * @@ -40,7 +31,6 @@ class PlasmaPropertyCollection(list): IonizationData, LinesLowerLevelIndex, LinesUpperLevelIndex, - TauSobolev, StimulatedEmissionFactor, SelectedAtoms, ElectronTemperature, @@ -48,9 +38,7 @@ class PlasmaPropertyCollection(list): ) lte_ionization_properties = PlasmaPropertyCollection([PhiSahaLTE]) lte_excitation_properties = PlasmaPropertyCollection([LevelBoltzmannFactorLTE]) -macro_atom_properties = PlasmaPropertyCollection( - [BetaSobolev, TransitionProbabilities, MacroAtomData] -) +macro_atom_properties = [] nebular_ionization_properties = PlasmaPropertyCollection( [PhiSahaNebular, ZetaData, BetaElectron, RadiationFieldCorrection] ) @@ -63,8 +51,6 @@ class PlasmaPropertyCollection(list): LevelBoltzmannFactorNLTE, NLTEData, PreviousElectronDensities, - PreviousBetaSobolev, - BetaSobolev, ] ) nlte_root_solver_properties = PlasmaPropertyCollection( @@ -111,22 +97,13 @@ class PlasmaPropertyCollection(list): SahaFactor, CorrPhotoIonRateCoeff, SpontRecombCoolingRateCoeff, - RawRecombTransProbs, - RawPhotoIonTransProbs, - RawRadBoundBoundTransProbs, - MarkovChainTransProbs, - NonContinuumTransProbsMask, YgInterpolator, CollExcRateCoeff, CollDeexcRateCoeff, RawCollisionTransProbs, MarkovChainIndex, - MarkovChainTransProbsCollector, - NonMarkovChainTransitionProbabilities, - MonteCarloTransProbs, FreeFreeCoolingRate, FreeBoundCoolingRate, - BoundFreeOpacity, LevelNumberDensityLTE, PhotoIonBoltzmannFactor, FreeBoundEmissionCDF, @@ -134,9 +111,7 @@ class PlasmaPropertyCollection(list): LevelIdxs2TransitionIdx, CollIonRateCoeffSeaton, CollRecombRateCoeff, - RawCollIonTransProbs, ContinuumInteractionHandler, - BetaSobolev, ] ) adiabatic_cooling_properties = PlasmaPropertyCollection([AdiabaticCoolingRate]) diff --git a/tardis/plasma/properties/radiative_properties.py b/tardis/plasma/properties/radiative_properties.py index 859bf85cae4..6802b88d99d 100644 --- a/tardis/plasma/properties/radiative_properties.py +++ b/tardis/plasma/properties/radiative_properties.py @@ -3,7 +3,6 @@ import numpy as np import pandas as pd 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 @@ -16,7 +15,6 @@ __all__ = [ "StimulatedEmissionFactor", - "BetaSobolev", "RawRadBoundBoundTransProbs", ] @@ -96,9 +94,9 @@ def calculate( ) # the following line probably can be removed as well - stimulated_emission_factor[ - np.isneginf(stimulated_emission_factor) - ] = 0.0 + stimulated_emission_factor[np.isneginf(stimulated_emission_factor)] = ( + 0.0 + ) stimulated_emission_factor[ meta_stable_upper & (stimulated_emission_factor < 0) ] = 0.0 @@ -118,46 +116,6 @@ def calculate( return stimulated_emission_factor -class BetaSobolev(ProcessingPlasmaProperty): - """ - Attributes - ---------- - beta_sobolev : Numpy Array, dtype float - """ - - outputs = ("beta_sobolev",) - latex_name = (r"\beta_{\textrm{sobolev}}",) - - def calculate(self, tau_sobolevs): - if getattr(self, "beta_sobolev", None) is None: - initial = 0.0 - else: - initial = self.beta_sobolev - - beta_sobolev = pd.DataFrame( - initial, index=tau_sobolevs.index, columns=tau_sobolevs.columns - ) - - self.calculate_beta_sobolev( - tau_sobolevs.values.ravel(), beta_sobolev.values.ravel() - ) - return beta_sobolev - - @staticmethod - @jit(nopython=True, parallel=True) - def calculate_beta_sobolev(tau_sobolevs, beta_sobolevs): - for i in prange(len(tau_sobolevs)): - if tau_sobolevs[i] > 1e3: - beta_sobolevs[i] = tau_sobolevs[i] ** -1 - elif tau_sobolevs[i] < 1e-4: - beta_sobolevs[i] = 1 - 0.5 * tau_sobolevs[i] - else: - beta_sobolevs[i] = (1 - np.exp(-tau_sobolevs[i])) / ( - tau_sobolevs[i] - ) - return beta_sobolevs - - class RawRadBoundBoundTransProbs( TransitionProbabilities, TransitionProbabilitiesProperty ): diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 7468233be6b..2f9267d80ea 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -17,12 +17,7 @@ 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 ( +from tardis.plasma.properties.legacy_property_collections import ( adiabatic_cooling_properties, basic_inputs, basic_properties, @@ -42,6 +37,11 @@ non_nlte_properties, two_photon_properties, ) +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.plasma.radiation_field import DilutePlanckianRadiationField from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 3e8822c5de0..c4615db9313 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -289,12 +289,8 @@ def advance_state(self, emitted_luminosity): ) ) - 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_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, @@ -449,7 +445,7 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0): f"\n\tStarting iteration {(self.iterations_executed + 1):d} of {self.iterations:d}" ) - opacity_state = self.opacity.solve(self.plasma) + opacity_state = self.opacity.legacy_solve(self.plasma) if self.macro_atom is not None: if montecarlo_globals.CONTINUUM_PROCESSES_ENABLED: macro_atom_state = MacroAtomState.from_legacy_plasma( diff --git a/tardis/spectrum/formal_integral.py b/tardis/spectrum/formal_integral.py index fdbef19ebf8..18dfedb5975 100644 --- a/tardis/spectrum/formal_integral.py +++ b/tardis/spectrum/formal_integral.py @@ -10,6 +10,7 @@ from tardis import constants as const from tardis.opacities.opacity_state import ( + opacity_state_to_numba, opacity_state_initialize, ) from tardis.spectrum.formal_integral_cuda import ( @@ -19,9 +20,6 @@ from tardis.transport.montecarlo import njit_dict, njit_dict_no_parallel from tardis.transport.montecarlo.configuration import montecarlo_globals from tardis.transport.montecarlo.configuration.constants import SIGMA_THOMSON -from tardis.transport.montecarlo.numba_interface import ( - opacity_state_initialize, -) C_INV = 3.33564e-11 PI = np.pi @@ -272,7 +270,15 @@ class FormalIntegrator: points : int64 """ - def __init__(self, simulation_state, plasma, transport, points=1000): + def __init__( + self, + simulation_state, + plasma, + transport, + opacity_state=None, + macro_atom_state=None, + points=1000, + ): self.simulation_state = simulation_state self.transport = transport self.points = points @@ -280,15 +286,26 @@ def __init__(self, simulation_state, plasma, transport, points=1000): self.montecarlo_configuration = ( self.transport.montecarlo_configuration ) - if plasma: - self.plasma = opacity_state_initialize( + if plasma and opacity_state and macro_atom_state: + self.opacity_state = opacity_state_to_numba( + opacity_state, + macro_atom_state, + transport.line_interaction_type, + ) + self.atomic_data = plasma.atomic_data + self.plasma = plasma + self.levels_index = plasma.levels + elif plasma: + self.opacity_state = opacity_state_initialize( plasma, transport.line_interaction_type, self.montecarlo_configuration.DISABLE_LINE_SCATTERING, ) self.atomic_data = plasma.atomic_data - self.original_plasma = plasma + self.plasma = plasma self.levels_index = plasma.levels + else: + self.opacity_state = None def generate_numba_objects(self): """instantiate the numba interface objects @@ -304,11 +321,12 @@ def generate_numba_objects(self): self.transport.r_outer_i / self.simulation_state.time_explosion.to("s").value, ) - self.opacity_state = opacity_state_initialize( - self.original_plasma, - self.transport.line_interaction_type, - self.montecarlo_configuration.DISABLE_LINE_SCATTERING, - ) + if self.opacity_state is None: + self.opacity_state = opacity_state_initialize( + self.plasma, + self.transport.line_interaction_type, + self.montecarlo_configuration.DISABLE_LINE_SCATTERING, + ) if self.transport.use_gpu: self.integrator = CudaFormalIntegrator( self.numba_radial_1d_geometry, @@ -423,7 +441,7 @@ def make_source_function(self): # macro_ref = self.atomic_data.macro_atom_references macro_ref = self.atomic_data.macro_atom_references # macro_data = self.atomic_data.macro_atom_data - macro_data = self.original_plasma.macro_atom_data + macro_data = self.plasma.atomic_data.macro_atom_data no_lvls = len(self.levels_index) no_shells = len(simulation_state.dilution_factor) @@ -431,7 +449,7 @@ def make_source_function(self): if transport.line_interaction_type == "macroatom": internal_jump_mask = (macro_data.transition_type >= 0).values ma_int_data = macro_data[internal_jump_mask] - internal = self.original_plasma.transition_probabilities[ + internal = self.opacity_state.transition_probabilities[ internal_jump_mask ] @@ -442,7 +460,7 @@ def make_source_function(self): montecarlo_transport_state.packet_collection.time_of_simulation * simulation_state.volume ) - exptau = 1 - np.exp(-self.original_plasma.tau_sobolevs) + exptau = 1 - np.exp(-self.opacity_state.tau_sobolev) Edotlu = ( Edotlu_norm_factor * exptau @@ -475,7 +493,7 @@ def make_source_function(self): upper_level_index = self.atomic_data.lines.index.droplevel( "level_number_lower" ) - e_dot_lu = pd.DataFrame(Edotlu.values, index=upper_level_index) + e_dot_lu = pd.DataFrame(Edotlu.value, index=upper_level_index) e_dot_u = e_dot_lu.groupby(level=[0, 1, 2]).sum() e_dot_u_src_idx = macro_ref.loc[e_dot_u.index].references_idx.values @@ -486,7 +504,7 @@ def make_source_function(self): q_indices = (source_level_idx, destination_level_idx) for shell in range(no_shells): Q = sp.coo_matrix( - (internal[shell], q_indices), shape=(no_lvls, no_lvls) + (internal[:, shell], q_indices), shape=(no_lvls, no_lvls) ) inv_N = sp.identity(no_lvls) - Q e_dot_u_vec = np.zeros(no_lvls) @@ -498,16 +516,17 @@ def make_source_function(self): "ion_number", "source_level_number", ] # To make the q_ul e_dot_u product work, could be cleaner - transitions = self.original_plasma.atomic_data.macro_atom_data[ - self.original_plasma.atomic_data.macro_atom_data.transition_type - == -1 + transitions = self.plasma.atomic_data.macro_atom_data[ + self.plasma.atomic_data.macro_atom_data.transition_type == -1 ].copy() transitions_index = transitions.set_index( ["atomic_number", "ion_number", "source_level_number"] ).index.copy() - tmp = self.original_plasma.transition_probabilities[ - (self.atomic_data.macro_atom_data.transition_type == -1).values - ] + tmp = pd.DataFrame( + self.opacity_state.transition_probabilities[ + (self.atomic_data.macro_atom_data.transition_type == -1).values + ] + ) q_ul = tmp.set_index(transitions_index) t = simulation_state.time_explosion.value t = simulation_state.time_explosion.value @@ -526,7 +545,7 @@ def make_source_function(self): # Jredlu should already by in the correct order, i.e. by wavelength of # the transition l->u (similar to Jbluelu) - Jredlu = Jbluelu * np.exp(-self.original_plasma.tau_sobolevs) + att_S_ul + Jredlu = Jbluelu * np.exp(-self.opacity_state.tau_sobolev) + att_S_ul if self.interpolate_shells > 0: ( att_S_ul, @@ -543,11 +562,9 @@ def make_source_function(self): transport.r_outer_i = ( montecarlo_transport_state.geometry_state.r_outer ) - transport.tau_sobolevs_integ = ( - self.original_plasma.tau_sobolevs.values - ) + transport.tau_sobolevs_integ = self.opacity_state.tau_sobolev transport.electron_densities_integ = ( - self.original_plasma.electron_densities.values + self.opacity_state.electron_density ) return att_S_ul, Jredlu, Jbluelu, e_dot_u @@ -557,7 +574,7 @@ def interpolate_integrator_quantities( ): transport = self.transport mct_state = transport.transport_state - plasma = self.original_plasma + plasma = self.plasma nshells = self.interpolate_shells r_middle = ( mct_state.geometry_state.r_inner + mct_state.geometry_state.r_outer @@ -583,15 +600,15 @@ def interpolate_integrator_quantities( # (as in the MC simulation) transport.tau_sobolevs_integ = interp1d( r_middle, - plasma.tau_sobolevs, + self.opacity_state.tau_sobolev, fill_value="extrapolate", kind="nearest", )(r_middle_integ) att_S_ul = interp1d(r_middle, att_S_ul, fill_value="extrapolate")( r_middle_integ ) - Jredlu = pd.DataFrame( - interp1d(r_middle, Jredlu, fill_value="extrapolate")(r_middle_integ) + Jredlu = interp1d(r_middle, Jredlu, fill_value="extrapolate")( + r_middle_integ ) Jbluelu = interp1d(r_middle, Jbluelu, fill_value="extrapolate")( r_middle_integ @@ -616,7 +633,7 @@ def formal_integral(self, nu, N): res = self.make_source_function() att_S_ul = res[0].flatten(order="F") - Jred_lu = res[1].values.flatten(order="F") + Jred_lu = res[1].flatten(order="F") Jblue_lu = res[2].flatten(order="F") self.generate_numba_objects() diff --git a/tardis/spectrum/tests/test_cuda_formal_integral.py b/tardis/spectrum/tests/test_cuda_formal_integral.py index be69c8eef40..e141edad3da 100644 --- a/tardis/spectrum/tests/test_cuda_formal_integral.py +++ b/tardis/spectrum/tests/test_cuda_formal_integral.py @@ -353,12 +353,12 @@ def test_full_formal_integral( res_numba = formal_integrator_numba.make_source_function() att_S_ul_numba = res_numba[0].flatten(order="F") - Jred_lu_numba = res_numba[1].values.flatten(order="F") + Jred_lu_numba = res_numba[1].flatten(order="F") Jblue_lu_numba = res_numba[2].flatten(order="F") res_cuda = formal_integrator_cuda.make_source_function() att_S_ul_cuda = res_cuda[0].flatten(order="F") - Jred_lu_cuda = res_cuda[1].values.flatten(order="F") + Jred_lu_cuda = res_cuda[1].flatten(order="F") Jblue_lu_cuda = res_cuda[2].flatten(order="F") formal_integrator_numba.generate_numba_objects() diff --git a/tardis/workflows/simple_tardis_workflow.py b/tardis/workflows/simple_tardis_workflow.py index 1f5b3cbbf39..18e48ead456 100644 --- a/tardis/workflows/simple_tardis_workflow.py +++ b/tardis/workflows/simple_tardis_workflow.py @@ -9,7 +9,7 @@ from tardis.model import SimulationState from tardis.opacities.macro_atom.macroatom_solver import MacroAtomSolver from tardis.opacities.opacity_solver import OpacitySolver -from tardis.plasma.assembly.legacy_assembly import assemble_plasma +from tardis.plasma.assembly import PlasmaSolverFactory from tardis.plasma.radiation_field import DilutePlanckianRadiationField from tardis.simulation.convergence import ConvergenceSolver from tardis.spectrum.base import SpectrumSolver @@ -41,10 +41,22 @@ def __init__(self, configuration): atom_data=atom_data, ) - self.plasma_solver = assemble_plasma( + plasma_solver_factory = PlasmaSolverFactory( + atom_data, + configuration, + ) + + plasma_solver_factory.prepare_factory( + self.simulation_state.abundance.index, + "tardis.plasma.properties.property_collections", configuration, - self.simulation_state, - atom_data=atom_data, + ) + + self.plasma_solver = plasma_solver_factory.assemble( + self.simulation_state.elemental_number_density, + self.simulation_state.radiation_field_state, + self.simulation_state.time_explosion, + self.simulation_state._electron_densities, ) line_interaction_type = configuration.plasma.line_interaction_type @@ -337,6 +349,8 @@ def solve_opacity(self): self.plasma_solver.atomic_data, opacity_state.tau_sobolev, self.plasma_solver.stimulated_emission_factor, + opacity_state.beta_sobolev, + legacy_mode=False, ) return { @@ -394,6 +408,7 @@ def solve_montecarlo( def initialize_spectrum_solver( self, transport_state, + opacity_states, virtual_packet_energies=None, ): """Set up the spectrum solver @@ -419,7 +434,11 @@ def initialize_spectrum_solver( self.integrated_spectrum_settings ) self.spectrum_solver._integrator = FormalIntegrator( - self.simulation_state, self.plasma_solver, self.transport_solver + self.simulation_state, + self.plasma_solver, + self.transport_solver, + opacity_states["opacity_state"], + opacity_states["macro_atom_state"], ) def run(self): @@ -465,5 +484,6 @@ def run(self): self.initialize_spectrum_solver( transport_state, + opacity_states, virtual_packet_energies, ) diff --git a/tardis/workflows/standard_tardis_workflow.py b/tardis/workflows/standard_tardis_workflow.py index f5842943f32..e4ed2433799 100644 --- a/tardis/workflows/standard_tardis_workflow.py +++ b/tardis/workflows/standard_tardis_workflow.py @@ -219,18 +219,10 @@ def run(self): self.simulation_state.t_inner, ) - self.opacity_state = self.opacity_solver.solve(self.plasma_solver) - - if self.macro_atom_solver is not None: - self.macro_atom_state = self.macro_atom_solver.solve( - self.plasma_solver, - self.plasma_solver.atomic_data, - self.opacity_state.tau_sobolev, - self.plasma_solver.stimulated_emission_factor, - ) + opacity_states = self.solve_opacity() transport_state, virtual_packet_energies = self.solve_montecarlo( - self.real_packet_count + opacity_states, self.real_packet_count ) ( @@ -258,7 +250,9 @@ def run(self): "\n\tITERATIONS HAVE NOT CONVERGED, starting final iteration" ) transport_state, virtual_packet_energies = self.solve_montecarlo( - self.final_iteration_packet_count, self.virtual_packet_count + opacity_states, + self.final_iteration_packet_count, + self.virtual_packet_count, ) self.store_plasma_state( self.completed_iterations, @@ -276,5 +270,6 @@ def run(self): ) self.initialize_spectrum_solver( transport_state, + opacity_states, virtual_packet_energies, )