From 12a8f6ec3d227b6c19916001f8ef5fa673ab0588 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 11 Nov 2024 12:10:09 -0500 Subject: [PATCH] Modifications to the formal integral, and simple workflow testbed --- tardis/spectrum/formal_integral.py | 43 +++++++++++----------- tardis/workflows/simple_tardis_workflow.py | 35 +++++++++++++++--- 2 files changed, 52 insertions(+), 26 deletions(-) diff --git a/tardis/spectrum/formal_integral.py b/tardis/spectrum/formal_integral.py index dfe7d755f17..080c855990e 100644 --- a/tardis/spectrum/formal_integral.py +++ b/tardis/spectrum/formal_integral.py @@ -10,7 +10,7 @@ from tardis import constants as const from tardis.opacities.opacity_state import ( - opacity_state_initialize, + opacity_state_to_numba, ) from tardis.spectrum.formal_integral_cuda import ( CudaFormalIntegrator, @@ -19,9 +19,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 M_PI = np.arccos(-1) @@ -272,7 +269,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 @@ -281,10 +286,10 @@ def __init__(self, simulation_state, plasma, transport, points=1000): self.transport.montecarlo_configuration ) if plasma: - self.plasma = opacity_state_initialize( - plasma, + self.plasma = opacity_state_to_numba( + opacity_state, + macro_atom_state, transport.line_interaction_type, - self.montecarlo_configuration.DISABLE_LINE_SCATTERING, ) self.atomic_data = plasma.atomic_data self.original_plasma = plasma @@ -304,7 +309,7 @@ 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.opacity_state = opacity_state_to_numba( self.original_plasma, self.transport.line_interaction_type, self.montecarlo_configuration.DISABLE_LINE_SCATTERING, @@ -431,9 +436,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_jump_mask - ] + internal = self.plasma.transition_probabilities[internal_jump_mask] source_level_idx = ma_int_data.source_level_idx.values destination_level_idx = ma_int_data.destination_level_idx.values @@ -442,7 +445,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.plasma.tau_sobolev) Edotlu = ( Edotlu_norm_factor * exptau @@ -475,7 +478,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 @@ -505,7 +508,7 @@ def make_source_function(self): transitions_index = transitions.set_index( ["atomic_number", "ion_number", "source_level_number"] ).index.copy() - tmp = self.original_plasma.transition_probabilities[ + tmp = self.plasma.transition_probabilities[ (self.atomic_data.macro_atom_data.transition_type == -1).values ] q_ul = tmp.set_index(transitions_index) @@ -526,7 +529,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.plasma.tau_sobolev) + att_S_ul if self.interpolate_shells > 0: ( att_S_ul, @@ -543,11 +546,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.plasma.tau_sobolev.values transport.electron_densities_integ = ( - self.original_plasma.electron_densities.values + self.plasma.electron_densities.values ) return att_S_ul, Jredlu, Jbluelu, e_dot_u @@ -583,7 +584,7 @@ def interpolate_integrator_quantities( # (as in the MC simulation) transport.tau_sobolevs_integ = interp1d( r_middle, - plasma.tau_sobolevs, + plasma.tau_sobolev, fill_value="extrapolate", kind="nearest", )(r_middle_integ) diff --git a/tardis/workflows/simple_tardis_workflow.py b/tardis/workflows/simple_tardis_workflow.py index 1f5b3cbbf39..b07caa1cef9 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,27 @@ 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, + ) + + dilute_planckian_radiation_field = DilutePlanckianRadiationField( + self.simulation_state.t_radiative, + self.simulation_state.dilution_factor, + ) + + self.plasma_solver = plasma_solver_factory.assemble( + self.simulation_state.elemental_number_density, + dilute_planckian_radiation_field, + self.simulation_state.time_explosion, + self.simulation_state._electron_densities, ) line_interaction_type = configuration.plasma.line_interaction_type @@ -337,6 +354,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 +413,7 @@ def solve_montecarlo( def initialize_spectrum_solver( self, transport_state, + opacity_states, virtual_packet_energies=None, ): """Set up the spectrum solver @@ -419,7 +439,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 +489,6 @@ def run(self): self.initialize_spectrum_solver( transport_state, + opacity_states, virtual_packet_energies, )