From bf4fca1f93ebb2acd981d7269a38c6a1f744fa6d Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 10:44:00 -0400 Subject: [PATCH 01/28] Basic skeleton --- tardis/workflows/__init__.py | 0 .../workflows/standard_simulation_solver.py | 99 +++++++++++++++++++ 2 files changed, 99 insertions(+) create mode 100644 tardis/workflows/__init__.py create mode 100644 tardis/workflows/standard_simulation_solver.py diff --git a/tardis/workflows/__init__.py b/tardis/workflows/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py new file mode 100644 index 00000000000..79927b259dd --- /dev/null +++ b/tardis/workflows/standard_simulation_solver.py @@ -0,0 +1,99 @@ +import logging +from astropy import units as u + +from tardis.simulation.convergence import ConvergenceSolver + +# logging support +logger = logging.getLogger(__name__) + + +class StandardSimulationSolver: + def __init__(self, convergence_strategy): + self.simulation_state = None + self.spectrum_solver = None + self.transport_solver = None + self.luminosity_requested = 0 * u.erg / u.s + + # Convergence + self.convergence_strategy = convergence_strategy + self.converged = False + self.consecutive_converges_count = 0 + + # Convergence solvers + self.t_rad_convergence_solver = ConvergenceSolver( + self.convergence_strategy.t_rad + ) + self.w_convergence_solver = ConvergenceSolver( + self.convergence_strategy.w + ) + self.t_inner_convergence_solver = ConvergenceSolver( + self.convergence_strategy.t_inner + ) + + def _get_convergence_status( + self, t_rad, w, t_inner, estimated_t_rad, estimated_w, estimated_t_inner + ): + t_rad_converged = self.t_rad_convergence_solver.get_convergence_status( + t_rad.value, + estimated_t_rad.value, + self.simulation_state.no_of_shells, + ) + + w_converged = self.w_convergence_solver.get_convergence_status( + w, estimated_w, self.simulation_state.no_of_shells + ) + + t_inner_converged = ( + self.t_inner_convergence_solver.get_convergence_status( + t_inner.value, + estimated_t_inner.value, + 1, + ) + ) + + if np.all([t_rad_converged, w_converged, t_inner_converged]): + hold_iterations = self.convergence_strategy.hold_iterations + self.consecutive_converges_count += 1 + logger.info( + f"Iteration converged {self.consecutive_converges_count:d}/{(hold_iterations + 1):d} consecutive " + f"times." + ) + # If an iteration has converged, require hold_iterations more + # iterations to converge before we conclude that the Simulation + # is converged. + return self.consecutive_converges_count == hold_iterations + 1 + + self.consecutive_converges_count = 0 + return False + + def check_convergence(self, emitted_luminosity): + ( + estimated_t_rad, + estimated_dilution_factor, + ) = self.transport_solver.transport_state.calculate_radiationfield_properties() + + estimated_t_inner = self.estimate_t_inner( + self.simulation_state.t_inner, + self.luminosity_requested, + emitted_luminosity, + t_inner_update_exponent=self.convergence_strategy.t_inner_update_exponent, + ) + + converged = self._get_convergence_status( + self.simulation_state.t_radiative, + self.simulation_state.dilution_factor, + self.simulation_state.t_inner, + estimated_t_rad, + estimated_dilution_factor, + estimated_t_inner, + ) + + return converged + + def solve(self): + while not converged: + solve_plasma() + run_montecarlo() + converged = check_convergence() + + run_montecarlo(final=True) From 73f55a4f4b63d18ed7377bd73d384cfd2091c872 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 11:39:28 -0400 Subject: [PATCH 02/28] Lots of progress Solvers pass info to each other Simplified some parts --- .../workflows/standard_simulation_solver.py | 238 ++++++++++++++++-- 1 file changed, 216 insertions(+), 22 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 79927b259dd..1e778d28b9f 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -1,46 +1,68 @@ import logging + +import numpy as np from astropy import units as u +from tardis.io.atom_data.base import AtomData from tardis.simulation.convergence import ConvergenceSolver +from tardis.spectrum.formal_integral import FormalIntegrator # logging support logger = logging.getLogger(__name__) class StandardSimulationSolver: - def __init__(self, convergence_strategy): + def __init__(self, convergence_strategy, atom_data_path): + self.atom_data_path = atom_data_path self.simulation_state = None + self.atom_data = None self.spectrum_solver = None self.transport_solver = None + self.plasma_solver = None self.luminosity_requested = 0 * u.erg / u.s + self.integrated_spectrum_settings = None # Convergence self.convergence_strategy = convergence_strategy - self.converged = False self.consecutive_converges_count = 0 + self.converged = False + self.total_iterations = 1 + self.completed_iterations = 0 # Convergence solvers - self.t_rad_convergence_solver = ConvergenceSolver( - self.convergence_strategy.t_rad + self.t_radiative_convergence_solver = ConvergenceSolver( + self.convergence_strategy.t_radiative ) - self.w_convergence_solver = ConvergenceSolver( - self.convergence_strategy.w + self.dilution_factor_convergence_solver = ConvergenceSolver( + self.convergence_strategy.dilution_factor ) self.t_inner_convergence_solver = ConvergenceSolver( self.convergence_strategy.t_inner ) def _get_convergence_status( - self, t_rad, w, t_inner, estimated_t_rad, estimated_w, estimated_t_inner + self, + t_radiative, + dilution_factor, + t_inner, + estimated_t_radiative, + estimated_dilution_factor, + estimated_t_inner, ): - t_rad_converged = self.t_rad_convergence_solver.get_convergence_status( - t_rad.value, - estimated_t_rad.value, - self.simulation_state.no_of_shells, + t_radiative_converged = ( + self.t_radiative_convergence_solver.get_convergence_status( + t_radiative.value, + estimated_t_radiative.value, + self.simulation_state.no_of_shells, + ) ) - w_converged = self.w_convergence_solver.get_convergence_status( - w, estimated_w, self.simulation_state.no_of_shells + dilution_factor_converged = ( + self.dilution_factor_convergence_solver.get_convergence_status( + dilution_factor, + estimated_dilution_factor, + self.simulation_state.no_of_shells, + ) ) t_inner_converged = ( @@ -51,7 +73,13 @@ def _get_convergence_status( ) ) - if np.all([t_rad_converged, w_converged, t_inner_converged]): + if np.all( + [ + t_radiative_converged, + dilution_factor_converged, + t_inner_converged, + ] + ): hold_iterations = self.convergence_strategy.hold_iterations self.consecutive_converges_count += 1 logger.info( @@ -66,9 +94,28 @@ def _get_convergence_status( self.consecutive_converges_count = 0 return False - def check_convergence(self, emitted_luminosity): + def _get_atom_data(self): + try: + self.atom_data = AtomData.from_hdf(self.atom_data_path) + except TypeError: + logger.debug("Atom Data Cannot be Read from HDF.") + + def estimate_t_inner( + self, + input_t_inner, + luminosity_requested, + emitted_luminosity, + t_inner_update_exponent=-0.5, + ): + luminosity_ratios = ( + (emitted_luminosity / luminosity_requested).to(1).value + ) + + return input_t_inner * luminosity_ratios**t_inner_update_exponent + + def get_convergence_estimates(self, emitted_luminosity): ( - estimated_t_rad, + estimated_t_radiative, estimated_dilution_factor, ) = self.transport_solver.transport_state.calculate_radiationfield_properties() @@ -78,22 +125,169 @@ def check_convergence(self, emitted_luminosity): emitted_luminosity, t_inner_update_exponent=self.convergence_strategy.t_inner_update_exponent, ) + return ( + estimated_t_radiative, + estimated_dilution_factor, + estimated_t_inner, + ) + def check_convergence( + self, + estimated_t_radiative, + estimated_dilution_factor, + estimated_t_inner, + ): converged = self._get_convergence_status( self.simulation_state.t_radiative, self.simulation_state.dilution_factor, self.simulation_state.t_inner, - estimated_t_rad, + estimated_t_radiative, estimated_dilution_factor, estimated_t_inner, ) return converged + def solve_plasma( + self, + estimated_t_radiative, + estimated_dilution_factor, + estimated_t_inner, + ): + next_t_radiative = self.t_rad_convergence_solver.converge( + self.simulation_state.t_radiative, + estimated_t_radiative, + ) + next_dilution_factor = self.dilution_factor_convergence_solver.converge( + self.simulation_state.dilution_factor, + estimated_dilution_factor, + ) + if ( + self.iterations_executed + 1 + ) % self.convergence_strategy.lock_t_inner_cycles == 0: + next_t_inner = self.t_inner_convergence_solver.converge( + self.simulation_state.t_inner, + estimated_t_inner, + ) + else: + next_t_inner = self.simulation_state.t_inner + + self.simulation_state.t_radiative = next_t_radiative + self.simulation_state.dilution_factor = next_dilution_factor + self.simulation_state.blackbody_packet_source.temperature = next_t_inner + + update_properties = dict( + t_rad=self.simulation_state.t_radiative, + w=self.simulation_state.dilution_factor, + ) + # 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, + 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, + ) + + self.plasma_solver.update(**update_properties) + + def solve_montecarlo(self, no_of_real_packets, no_of_virtual_packets): + transport_state = self.transport_solver.initialize_transport_state( + self.simulation_state, + self.plasma_solver, + no_of_real_packets, + no_of_virtual_packets=no_of_virtual_packets, + iteration=self.completed_iterations, + ) + + virtual_packet_energies = self.transport.run( + transport_state, + time_explosion=self.simulation_state.time_explosion, + iteration=self.completed_iterations, + total_iterations=self.total_iterations, + show_progress_bars=self.show_progress_bars, + ) + + return transport_state, virtual_packet_energies + + def solve_spectrum( + self, + transport_state, + virtual_packet_energies=None, + integrated_spectrum_settings=None, + ): + # Set up spectrum solver + self.spectrum_solver.transport_state = transport_state + if virtual_packet_energies is not None: + self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = ( + virtual_packet_energies + ) + + if integrated_spectrum_settings is not None: + # Set up spectrum solver integrator + self.spectrum_solver.integrator_settings = ( + integrated_spectrum_settings + ) + self.spectrum_solver._integrator = FormalIntegrator( + self.simulation_state, self.plasma, self.transport + ) + + def calculate_emitted_luminosity(self, transport_state): + self.spectrum_solver.transport_state = transport_state + + output_energy = ( + self.transport.transport_state.packet_collection.output_energies + ) + if np.sum(output_energy < 0) == len(output_energy): + logger.critical("No r-packet escaped through the outer boundary.") + + emitted_luminosity = self.spectrum_solver.calculate_emitted_luminosity( + self.luminosity_nu_start, self.luminosity_nu_end + ) + return emitted_luminosity + def solve(self): - while not converged: - solve_plasma() - run_montecarlo() - converged = check_convergence() + converged = False + while self.completed_iterations < self.total_iterations - 1: + transport_state, virtual_packet_energies = self.solve_montecarlo() + + emitted_luminosity = self.calculate_emitted_luminosity( + transport_state + ) + + ( + estimated_t_radiative, + estimated_dilution_factor, + estimated_t_inner, + ) = self.get_convergence_estimates(emitted_luminosity) - run_montecarlo(final=True) + self.solve_plasma( + estimated_t_radiative, + estimated_dilution_factor, + estimated_t_inner, + ) + + converged = self.check_convergence( + estimated_t_radiative, + estimated_dilution_factor, + estimated_t_inner, + ) + self.completed_iterations += 1 + + if converged and self.convergence_strategy.stop_if_converged: + break + + transport_state, virtual_packet_energies = self.solve_montecarlo() + self.solve_spectrum( + transport_state, + virtual_packet_energies, + self.integrated_spectrum_settings, + ) From ffbc45e3b67f81e951cb364892abc52e1b241179 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 12:18:36 -0400 Subject: [PATCH 03/28] Adds setup method, refactors some functions, removes unused items --- .../workflows/standard_simulation_solver.py | 251 +++++++++++------- 1 file changed, 155 insertions(+), 96 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 1e778d28b9f..4424620be75 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -1,11 +1,17 @@ import logging +from pathlib import Path import numpy as np from astropy import units as u +from tardis import constants as const from tardis.io.atom_data.base import AtomData +from tardis.model import SimulationState +from tardis.plasma.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 tardis.transport.montecarlo.base import MonteCarloTransportSolver # logging support logger = logging.getLogger(__name__) @@ -19,7 +25,7 @@ def __init__(self, convergence_strategy, atom_data_path): self.spectrum_solver = None self.transport_solver = None self.plasma_solver = None - self.luminosity_requested = 0 * u.erg / u.s + self.virtual_packet_count = 0 self.integrated_spectrum_settings = None # Convergence @@ -28,6 +34,9 @@ def __init__(self, convergence_strategy, atom_data_path): self.converged = False self.total_iterations = 1 self.completed_iterations = 0 + self.luminosity_nu_start = 0 + self.luminosity_nu_end = 0 + self.luminosity_requested = 0 * u.erg / u.s # Convergence solvers self.t_radiative_convergence_solver = ConvergenceSolver( @@ -40,18 +49,73 @@ def __init__(self, convergence_strategy, atom_data_path): self.convergence_strategy.t_inner ) - def _get_convergence_status( + def _get_atom_data(self, configuration): + if "atom_data" in configuration: + if Path(configuration.atom_data).is_absolute(): + atom_data_fname = Path(configuration.atom_data) + else: + atom_data_fname = ( + Path(configuration.config_dirname) / configuration.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 estimate_t_inner( + self, + input_t_inner, + luminosity_requested, + emitted_luminosity, + t_inner_update_exponent=-0.5, + ): + luminosity_ratios = ( + (emitted_luminosity / luminosity_requested).to(1).value + ) + + return input_t_inner * luminosity_ratios**t_inner_update_exponent + + def get_convergence_estimates(self, emitted_luminosity): + ( + estimated_t_radiative, + estimated_dilution_factor, + ) = self.transport_solver.transport_state.calculate_radiationfield_properties() + + estimated_t_inner = self.estimate_t_inner( + self.simulation_state.t_inner, + self.luminosity_requested, + emitted_luminosity, + t_inner_update_exponent=self.convergence_strategy.t_inner_update_exponent, + ) + return ( + estimated_t_radiative, + estimated_dilution_factor, + estimated_t_inner, + ) + + def check_convergence( self, - t_radiative, - dilution_factor, - t_inner, estimated_t_radiative, estimated_dilution_factor, estimated_t_inner, ): t_radiative_converged = ( self.t_radiative_convergence_solver.get_convergence_status( - t_radiative.value, + self.simulation_state.t_radiative.value, estimated_t_radiative.value, self.simulation_state.no_of_shells, ) @@ -59,7 +123,7 @@ def _get_convergence_status( dilution_factor_converged = ( self.dilution_factor_convergence_solver.get_convergence_status( - dilution_factor, + self.simulation_state.dilution_factor, estimated_dilution_factor, self.simulation_state.no_of_shells, ) @@ -67,7 +131,7 @@ def _get_convergence_status( t_inner_converged = ( self.t_inner_convergence_solver.get_convergence_status( - t_inner.value, + self.simulation_state.t_inner.value, estimated_t_inner.value, 1, ) @@ -94,67 +158,14 @@ def _get_convergence_status( self.consecutive_converges_count = 0 return False - def _get_atom_data(self): - try: - self.atom_data = AtomData.from_hdf(self.atom_data_path) - except TypeError: - logger.debug("Atom Data Cannot be Read from HDF.") - - def estimate_t_inner( - self, - input_t_inner, - luminosity_requested, - emitted_luminosity, - t_inner_update_exponent=-0.5, - ): - luminosity_ratios = ( - (emitted_luminosity / luminosity_requested).to(1).value - ) - - return input_t_inner * luminosity_ratios**t_inner_update_exponent - - def get_convergence_estimates(self, emitted_luminosity): - ( - estimated_t_radiative, - estimated_dilution_factor, - ) = self.transport_solver.transport_state.calculate_radiationfield_properties() - - estimated_t_inner = self.estimate_t_inner( - self.simulation_state.t_inner, - self.luminosity_requested, - emitted_luminosity, - t_inner_update_exponent=self.convergence_strategy.t_inner_update_exponent, - ) - return ( - estimated_t_radiative, - estimated_dilution_factor, - estimated_t_inner, - ) - - def check_convergence( - self, - estimated_t_radiative, - estimated_dilution_factor, - estimated_t_inner, - ): - converged = self._get_convergence_status( - self.simulation_state.t_radiative, - self.simulation_state.dilution_factor, - self.simulation_state.t_inner, - estimated_t_radiative, - estimated_dilution_factor, - estimated_t_inner, - ) - - return converged - def solve_plasma( self, + transport_state, estimated_t_radiative, estimated_dilution_factor, estimated_t_inner, ): - next_t_radiative = self.t_rad_convergence_solver.converge( + next_t_radiative = self.t_radiative_convergence_solver.converge( self.simulation_state.t_radiative, estimated_t_radiative, ) @@ -163,7 +174,7 @@ def solve_plasma( estimated_dilution_factor, ) if ( - self.iterations_executed + 1 + self.completed_iterations + 1 ) % self.convergence_strategy.lock_t_inner_cycles == 0: next_t_inner = self.t_inner_convergence_solver.converge( self.simulation_state.t_inner, @@ -182,24 +193,16 @@ def solve_plasma( ) # 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: + estimators = transport_state.radfield_mc_estimators + if "j_blue_estimator" in self.plasma_solver.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, - 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, - ) self.plasma_solver.update(**update_properties) - def solve_montecarlo(self, no_of_real_packets, no_of_virtual_packets): + def solve_montecarlo(self, no_of_real_packets, no_of_virtual_packets=0): transport_state = self.transport_solver.initialize_transport_state( self.simulation_state, self.plasma_solver, @@ -208,59 +211,113 @@ def solve_montecarlo(self, no_of_real_packets, no_of_virtual_packets): iteration=self.completed_iterations, ) - virtual_packet_energies = self.transport.run( + virtual_packet_energies = self.transport_solver.run( transport_state, time_explosion=self.simulation_state.time_explosion, iteration=self.completed_iterations, total_iterations=self.total_iterations, - show_progress_bars=self.show_progress_bars, + show_progress_bars=False, ) return transport_state, virtual_packet_energies - def solve_spectrum( + def initialize_spectrum_solver( self, transport_state, virtual_packet_energies=None, - integrated_spectrum_settings=None, ): # Set up spectrum solver self.spectrum_solver.transport_state = transport_state + if virtual_packet_energies is not None: self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = ( virtual_packet_energies ) - if integrated_spectrum_settings is not None: + if self.integrated_spectrum_settings is not None: # Set up spectrum solver integrator self.spectrum_solver.integrator_settings = ( - integrated_spectrum_settings + self.integrated_spectrum_settings ) self.spectrum_solver._integrator = FormalIntegrator( self.simulation_state, self.plasma, self.transport ) - def calculate_emitted_luminosity(self, transport_state): - self.spectrum_solver.transport_state = transport_state + def setup_solver(self, configuration): + atom_data = self._get_atom_data(configuration) - output_energy = ( - self.transport.transport_state.packet_collection.output_energies + self.simulation_state = SimulationState.from_config( + configuration, + atom_data=atom_data, ) - if np.sum(output_energy < 0) == len(output_energy): - logger.critical("No r-packet escaped through the outer boundary.") - emitted_luminosity = self.spectrum_solver.calculate_emitted_luminosity( - self.luminosity_nu_start, self.luminosity_nu_end + self.plasma_solver = assemble_plasma( + configuration, + self.simulation_state, + atom_data=atom_data, ) - return emitted_luminosity + + self.transport_solver = MonteCarloTransportSolver.from_config( + configuration, + packet_source=self.simulation_state.packet_source, + enable_virtual_packet_logging=False, + ) + + self.luminosity_nu_start = ( + configuration.supernova.luminosity_wavelength_end.to( + u.Hz, u.spectral() + ) + ) + + if u.isclose( + configuration.supernova.luminosity_wavelength_start, 0 * u.angstrom + ): + self.luminosity_nu_end = np.inf * u.Hz + else: + self.luminosity_nu_end = ( + const.c / configuration.supernova.luminosity_wavelength_start + ).to(u.Hz) + + self.real_packet_count = configuration.montecarlo.no_of_packets + + final_iteration_packet_count = ( + configuration.montecarlo.last_no_of_packets + ) + + if ( + final_iteration_packet_count is None + or final_iteration_packet_count < 0 + ): + final_iteration_packet_count = self.real_packet_count + + self.final_iteration_packet_count = int(final_iteration_packet_count) + + self.virtual_packet_count = int( + configuration.montecarlo.no_of_virtual_packets + ) + + self.integrated_spectrum_settings = configuration.spectrum.integrated + self.spectrum_solver = SpectrumSolver.from_config(configuration) def solve(self): converged = False while self.completed_iterations < self.total_iterations - 1: - transport_state, virtual_packet_energies = self.solve_montecarlo() + transport_state, virtual_packet_energies = self.solve_montecarlo( + self.real_packet_count + ) - emitted_luminosity = self.calculate_emitted_luminosity( - transport_state + output_energy = transport_state.packet_collection.output_energies + if np.sum(output_energy < 0) == len(output_energy): + logger.critical( + "No r-packet escaped through the outer boundary." + ) + + self.spectrum_solver.transport_state = transport_state + + emitted_luminosity = ( + self.spectrum_solver.calculate_emitted_luminosity( + self.luminosity_nu_start, self.luminosity_nu_end + ) ) ( @@ -270,6 +327,7 @@ def solve(self): ) = self.get_convergence_estimates(emitted_luminosity) self.solve_plasma( + transport_state, estimated_t_radiative, estimated_dilution_factor, estimated_t_inner, @@ -285,9 +343,10 @@ def solve(self): if converged and self.convergence_strategy.stop_if_converged: break - transport_state, virtual_packet_energies = self.solve_montecarlo() - self.solve_spectrum( + transport_state, virtual_packet_energies = self.solve_montecarlo( + self.final_iteration_packet_count, self.virtual_packet_count + ) + self.initialize_spectrum_solver( transport_state, virtual_packet_energies, - self.integrated_spectrum_settings, ) From e85a6c936f170a5e663d674c8c959fadb03df41c Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 12:24:23 -0400 Subject: [PATCH 04/28] Move convergence setup to setup method --- .../workflows/standard_simulation_solver.py | 33 ++++++++++--------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 4424620be75..948451d388a 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -18,8 +18,7 @@ class StandardSimulationSolver: - def __init__(self, convergence_strategy, atom_data_path): - self.atom_data_path = atom_data_path + def __init__(self): self.simulation_state = None self.atom_data = None self.spectrum_solver = None @@ -29,7 +28,7 @@ def __init__(self, convergence_strategy, atom_data_path): self.integrated_spectrum_settings = None # Convergence - self.convergence_strategy = convergence_strategy + self.convergence_strategy = None self.consecutive_converges_count = 0 self.converged = False self.total_iterations = 1 @@ -38,17 +37,6 @@ def __init__(self, convergence_strategy, atom_data_path): self.luminosity_nu_end = 0 self.luminosity_requested = 0 * u.erg / u.s - # Convergence solvers - self.t_radiative_convergence_solver = ConvergenceSolver( - self.convergence_strategy.t_radiative - ) - self.dilution_factor_convergence_solver = ConvergenceSolver( - self.convergence_strategy.dilution_factor - ) - self.t_inner_convergence_solver = ConvergenceSolver( - self.convergence_strategy.t_inner - ) - def _get_atom_data(self, configuration): if "atom_data" in configuration: if Path(configuration.atom_data).is_absolute(): @@ -240,7 +228,7 @@ def initialize_spectrum_solver( self.integrated_spectrum_settings ) self.spectrum_solver._integrator = FormalIntegrator( - self.simulation_state, self.plasma, self.transport + self.simulation_state, self.plasma_solver, self.transport_solver ) def setup_solver(self, configuration): @@ -299,6 +287,21 @@ def setup_solver(self, configuration): self.integrated_spectrum_settings = configuration.spectrum.integrated self.spectrum_solver = SpectrumSolver.from_config(configuration) + self.convergence_strategy = ( + configuration.montecarlo.convergence_strategy + ) + + # Convergence solvers + self.t_radiative_convergence_solver = ConvergenceSolver( + self.convergence_strategy.t_radiative + ) + self.dilution_factor_convergence_solver = ConvergenceSolver( + self.convergence_strategy.dilution_factor + ) + self.t_inner_convergence_solver = ConvergenceSolver( + self.convergence_strategy.t_inner + ) + def solve(self): converged = False while self.completed_iterations < self.total_iterations - 1: From 76ae4144ad8e0fa231aa50e0c386afc9a2ca1d92 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 13:07:41 -0400 Subject: [PATCH 05/28] Add docs notebook, fix a couple of errors --- docs/workflows/standard_workflow.ipynb | 119 ++++++++++++++++++ .../workflows/standard_simulation_solver.py | 4 +- 2 files changed, 121 insertions(+), 2 deletions(-) create mode 100644 docs/workflows/standard_workflow.ipynb diff --git a/docs/workflows/standard_workflow.ipynb b/docs/workflows/standard_workflow.ipynb new file mode 100644 index 00000000000..6b0f14595fb --- /dev/null +++ b/docs/workflows/standard_workflow.ipynb @@ -0,0 +1,119 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from tardis.workflows.standard_simulation_solver import StandardSimulationSolver\n", + "from tardis.io.configuration.config_reader import Configuration" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = Configuration.from_yaml('../tardis_example.yml')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "solver = StandardSimulationSolver()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "solver.setup_solver(config)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "solver.solve()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "spectrum = solver.spectrum_solver.spectrum_real_packets\n", + "spectrum_virtual = solver.spectrum_solver.spectrum_virtual_packets\n", + "spectrum_integrated = solver.spectrum_solver.spectrum_integrated" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%matplotlib inline\n", + "plt.figure(figsize=(10, 6.5))\n", + "\n", + "spectrum.plot(label=\"Normal packets\")\n", + "spectrum_virtual.plot(label=\"Virtual packets\")\n", + "spectrum_integrated.plot(label='Formal integral')\n", + "\n", + "plt.xlim(500, 9000)\n", + "plt.title(\"TARDIS example model spectrum\")\n", + "plt.xlabel(\"Wavelength [$\\AA$]\")\n", + "plt.ylabel(\"Luminosity density [erg/s/$\\AA$]\")\n", + "plt.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "tardis", + "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.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 948451d388a..0efff4e717e 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -293,10 +293,10 @@ def setup_solver(self, configuration): # Convergence solvers self.t_radiative_convergence_solver = ConvergenceSolver( - self.convergence_strategy.t_radiative + self.convergence_strategy.t_rad ) self.dilution_factor_convergence_solver = ConvergenceSolver( - self.convergence_strategy.dilution_factor + self.convergence_strategy.w ) self.t_inner_convergence_solver = ConvergenceSolver( self.convergence_strategy.t_inner From 1bb1d47d5666351518a65acd0379d9e8b8b1c48d Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 13:15:13 -0400 Subject: [PATCH 06/28] Move setup to init --- docs/workflows/standard_workflow.ipynb | 11 +- .../workflows/standard_simulation_solver.py | 155 ++++++++---------- 2 files changed, 72 insertions(+), 94 deletions(-) diff --git a/docs/workflows/standard_workflow.ipynb b/docs/workflows/standard_workflow.ipynb index 6b0f14595fb..2a35c2164a1 100644 --- a/docs/workflows/standard_workflow.ipynb +++ b/docs/workflows/standard_workflow.ipynb @@ -25,16 +25,7 @@ "metadata": {}, "outputs": [], "source": [ - "solver = StandardSimulationSolver()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "solver.setup_solver(config)" + "solver = StandardSimulationSolver(config)" ] }, { diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 0efff4e717e..2497ad20a73 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -18,24 +18,82 @@ class StandardSimulationSolver: - def __init__(self): - self.simulation_state = None - self.atom_data = None - self.spectrum_solver = None - self.transport_solver = None - self.plasma_solver = None - self.virtual_packet_count = 0 - self.integrated_spectrum_settings = None - + def __init__(self, configuration): # Convergence - self.convergence_strategy = None self.consecutive_converges_count = 0 self.converged = False self.total_iterations = 1 self.completed_iterations = 0 - self.luminosity_nu_start = 0 - self.luminosity_nu_end = 0 - self.luminosity_requested = 0 * u.erg / u.s + + atom_data = self._get_atom_data(configuration) + + self.simulation_state = SimulationState.from_config( + configuration, + atom_data=atom_data, + ) + + self.plasma_solver = assemble_plasma( + configuration, + self.simulation_state, + atom_data=atom_data, + ) + + self.transport_solver = MonteCarloTransportSolver.from_config( + configuration, + packet_source=self.simulation_state.packet_source, + enable_virtual_packet_logging=False, + ) + + self.luminosity_nu_start = ( + configuration.supernova.luminosity_wavelength_end.to( + u.Hz, u.spectral() + ) + ) + + if u.isclose( + configuration.supernova.luminosity_wavelength_start, 0 * u.angstrom + ): + self.luminosity_nu_end = np.inf * u.Hz + else: + self.luminosity_nu_end = ( + const.c / configuration.supernova.luminosity_wavelength_start + ).to(u.Hz) + + self.real_packet_count = configuration.montecarlo.no_of_packets + + final_iteration_packet_count = ( + configuration.montecarlo.last_no_of_packets + ) + + if ( + final_iteration_packet_count is None + or final_iteration_packet_count < 0 + ): + final_iteration_packet_count = self.real_packet_count + + self.final_iteration_packet_count = int(final_iteration_packet_count) + + self.virtual_packet_count = int( + configuration.montecarlo.no_of_virtual_packets + ) + + self.integrated_spectrum_settings = configuration.spectrum.integrated + self.spectrum_solver = SpectrumSolver.from_config(configuration) + + self.convergence_strategy = ( + configuration.montecarlo.convergence_strategy + ) + + # Convergence solvers + self.t_radiative_convergence_solver = ConvergenceSolver( + self.convergence_strategy.t_rad + ) + self.dilution_factor_convergence_solver = ConvergenceSolver( + self.convergence_strategy.w + ) + self.t_inner_convergence_solver = ConvergenceSolver( + self.convergence_strategy.t_inner + ) def _get_atom_data(self, configuration): if "atom_data" in configuration: @@ -231,77 +289,6 @@ def initialize_spectrum_solver( self.simulation_state, self.plasma_solver, self.transport_solver ) - def setup_solver(self, configuration): - atom_data = self._get_atom_data(configuration) - - self.simulation_state = SimulationState.from_config( - configuration, - atom_data=atom_data, - ) - - self.plasma_solver = assemble_plasma( - configuration, - self.simulation_state, - atom_data=atom_data, - ) - - self.transport_solver = MonteCarloTransportSolver.from_config( - configuration, - packet_source=self.simulation_state.packet_source, - enable_virtual_packet_logging=False, - ) - - self.luminosity_nu_start = ( - configuration.supernova.luminosity_wavelength_end.to( - u.Hz, u.spectral() - ) - ) - - if u.isclose( - configuration.supernova.luminosity_wavelength_start, 0 * u.angstrom - ): - self.luminosity_nu_end = np.inf * u.Hz - else: - self.luminosity_nu_end = ( - const.c / configuration.supernova.luminosity_wavelength_start - ).to(u.Hz) - - self.real_packet_count = configuration.montecarlo.no_of_packets - - final_iteration_packet_count = ( - configuration.montecarlo.last_no_of_packets - ) - - if ( - final_iteration_packet_count is None - or final_iteration_packet_count < 0 - ): - final_iteration_packet_count = self.real_packet_count - - self.final_iteration_packet_count = int(final_iteration_packet_count) - - self.virtual_packet_count = int( - configuration.montecarlo.no_of_virtual_packets - ) - - self.integrated_spectrum_settings = configuration.spectrum.integrated - self.spectrum_solver = SpectrumSolver.from_config(configuration) - - self.convergence_strategy = ( - configuration.montecarlo.convergence_strategy - ) - - # Convergence solvers - self.t_radiative_convergence_solver = ConvergenceSolver( - self.convergence_strategy.t_rad - ) - self.dilution_factor_convergence_solver = ConvergenceSolver( - self.convergence_strategy.w - ) - self.t_inner_convergence_solver = ConvergenceSolver( - self.convergence_strategy.t_inner - ) - def solve(self): converged = False while self.completed_iterations < self.total_iterations - 1: From 516fa65bd06b5ce6610866e33994f6cb6b423b8c Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 13:58:14 -0400 Subject: [PATCH 07/28] Fix iteration and convergence --- tardis/workflows/standard_simulation_solver.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 2497ad20a73..ebe4fdab44c 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -2,16 +2,20 @@ from pathlib import Path import numpy as np +import pandas as pd from astropy import units as u +from IPython.display import display from tardis import constants as const from tardis.io.atom_data.base import AtomData +from tardis.io.logger.logger import logging_state from tardis.model import SimulationState from tardis.plasma.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 tardis.transport.montecarlo.base import MonteCarloTransportSolver +from tardis.util.base import is_notebook # logging support logger = logging.getLogger(__name__) @@ -22,11 +26,14 @@ def __init__(self, configuration): # Convergence self.consecutive_converges_count = 0 self.converged = False - self.total_iterations = 1 self.completed_iterations = 0 + self.luminosity_requested = ( + configuration.supernova.luminosity_requested.cgs + ) atom_data = self._get_atom_data(configuration) + # set up states and solvers self.simulation_state = SimulationState.from_config( configuration, atom_data=atom_data, @@ -59,7 +66,10 @@ def __init__(self, configuration): const.c / configuration.supernova.luminosity_wavelength_start ).to(u.Hz) - self.real_packet_count = configuration.montecarlo.no_of_packets + # montecarlo settings + self.total_iterations = int(configuration.montecarlo.iterations) + + self.real_packet_count = int(configuration.montecarlo.no_of_packets) final_iteration_packet_count = ( configuration.montecarlo.last_no_of_packets @@ -77,6 +87,7 @@ def __init__(self, configuration): configuration.montecarlo.no_of_virtual_packets ) + # spectrum settings self.integrated_spectrum_settings = configuration.spectrum.integrated self.spectrum_solver = SpectrumSolver.from_config(configuration) From 550f4e7fc357b1cc1008291f108a31b44f92fe7f Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 14:20:02 -0400 Subject: [PATCH 08/28] Simplify spectrumsolver init --- tardis/simulation/base.py | 8 ++----- tardis/spectrum/base.py | 2 +- tardis/spectrum/tests/test_spectrum_solver.py | 6 ++--- .../workflows/standard_simulation_solver.py | 23 ++++++++----------- 4 files changed, 16 insertions(+), 23 deletions(-) diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 92a658b80d7..f7e9a325c22 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -291,12 +291,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, diff --git a/tardis/spectrum/base.py b/tardis/spectrum/base.py index 2780e893783..c47d089c3f9 100644 --- a/tardis/spectrum/base.py +++ b/tardis/spectrum/base.py @@ -23,7 +23,7 @@ class SpectrumSolver(HDFWriterMixin): hdf_name = "spectrum" def __init__( - self, transport_state, spectrum_frequency_grid, integrator_settings=None + self, transport_state, spectrum_frequency_grid, integrator_settings ): self.transport_state = transport_state self.spectrum_frequency_grid = spectrum_frequency_grid diff --git a/tardis/spectrum/tests/test_spectrum_solver.py b/tardis/spectrum/tests/test_spectrum_solver.py index bb82e5d822c..9f4f5752ff4 100644 --- a/tardis/spectrum/tests/test_spectrum_solver.py +++ b/tardis/spectrum/tests/test_spectrum_solver.py @@ -43,7 +43,7 @@ def test_initialization(self, simulation): transport_state = simulation.transport.transport_state spectrum_frequency_grid = simulation.transport.spectrum_frequency_grid - solver = SpectrumSolver(transport_state, spectrum_frequency_grid) + solver = SpectrumSolver(transport_state, spectrum_frequency_grid, None) assert solver.transport_state == transport_state assert np.array_equal( solver.spectrum_frequency_grid.value, spectrum_frequency_grid.value @@ -60,7 +60,7 @@ def test_spectrum_real_packets(self, simulation): transport_state = simulation.transport.transport_state spectrum_frequency_grid = simulation.transport.spectrum_frequency_grid - solver = SpectrumSolver(transport_state, spectrum_frequency_grid) + solver = SpectrumSolver(transport_state, spectrum_frequency_grid, None) result = solver.spectrum_real_packets.luminosity key = "simulation/spectrum_solver/spectrum_real_packets/luminosity" expected = self.get_expected_data(key) @@ -76,7 +76,7 @@ def test_spectrum_real_packets_reabsorbed(self, simulation): transport_state = simulation.transport.transport_state spectrum_frequency_grid = simulation.transport.spectrum_frequency_grid - solver = SpectrumSolver(transport_state, spectrum_frequency_grid) + solver = SpectrumSolver(transport_state, spectrum_frequency_grid, None) result = solver.spectrum_real_packets_reabsorbed.luminosity key = "simulation/spectrum_solver/spectrum_real_packets_reabsorbed/luminosity" expected = self.get_expected_data(key) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index ebe4fdab44c..6db092f34ed 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -2,20 +2,16 @@ from pathlib import Path import numpy as np -import pandas as pd from astropy import units as u -from IPython.display import display from tardis import constants as const from tardis.io.atom_data.base import AtomData -from tardis.io.logger.logger import logging_state from tardis.model import SimulationState from tardis.plasma.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 tardis.transport.montecarlo.base import MonteCarloTransportSolver -from tardis.util.base import is_notebook # logging support logger = logging.getLogger(__name__) @@ -91,11 +87,11 @@ def __init__(self, configuration): self.integrated_spectrum_settings = configuration.spectrum.integrated self.spectrum_solver = SpectrumSolver.from_config(configuration) + # Convergence solvers self.convergence_strategy = ( configuration.montecarlo.convergence_strategy ) - # Convergence solvers self.t_radiative_convergence_solver = ConvergenceSolver( self.convergence_strategy.t_rad ) @@ -122,10 +118,9 @@ def _get_atom_data(self, configuration): 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" + except TypeError: + logger.exception( + "TypeError 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.", ) @@ -150,7 +145,9 @@ def get_convergence_estimates(self, emitted_luminosity): ( estimated_t_radiative, estimated_dilution_factor, - ) = self.transport_solver.transport_state.calculate_radiationfield_properties() + ) = ( + self.transport_solver.transport_state.calculate_radiationfield_properties() + ) estimated_t_inner = self.estimate_t_inner( self.simulation_state.t_inner, @@ -287,9 +284,9 @@ def initialize_spectrum_solver( self.spectrum_solver.transport_state = transport_state if virtual_packet_energies is not None: - self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = ( - virtual_packet_energies - ) + self.spectrum_solver._montecarlo_virtual_luminosity.value[ + : + ] = virtual_packet_energies if self.integrated_spectrum_settings is not None: # Set up spectrum solver integrator From afca8068597fb645b00c1aad6ab7829002381baa Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 14:30:34 -0400 Subject: [PATCH 09/28] Separate plasma "solver" and simulation state updates --- .../workflows/standard_simulation_solver.py | 29 ++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 6db092f34ed..2791288db03 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -212,9 +212,8 @@ def check_convergence( self.consecutive_converges_count = 0 return False - def solve_plasma( + def solve_simulation_state( self, - transport_state, estimated_t_radiative, estimated_dilution_factor, estimated_t_inner, @@ -241,17 +240,20 @@ def solve_plasma( self.simulation_state.dilution_factor = next_dilution_factor self.simulation_state.blackbody_packet_source.temperature = next_t_inner + def solve_plasma( + self, + transport_state, + ): update_properties = dict( t_rad=self.simulation_state.t_radiative, w=self.simulation_state.dilution_factor, ) # A check to see if the plasma is set with JBluesDetailed, in which # case it needs some extra kwargs. - estimators = transport_state.radfield_mc_estimators if "j_blue_estimator" in self.plasma_solver.outputs_dict: update_properties.update( - t_inner=next_t_inner, - j_blue_estimator=estimators.j_blue_estimator, + t_inner=self.simulation_state.blackbody_packet_source.temperature, + j_blue_estimator=transport_state.radfield_mc_estimators.j_blue_estimator, ) self.plasma_solver.update(**update_properties) @@ -273,6 +275,10 @@ def solve_montecarlo(self, no_of_real_packets, no_of_virtual_packets=0): show_progress_bars=False, ) + output_energy = transport_state.packet_collection.output_energies + if np.sum(output_energy < 0) == len(output_energy): + logger.critical("No r-packet escaped through the outer boundary.") + return transport_state, virtual_packet_energies def initialize_spectrum_solver( @@ -304,12 +310,6 @@ def solve(self): self.real_packet_count ) - output_energy = transport_state.packet_collection.output_energies - if np.sum(output_energy < 0) == len(output_energy): - logger.critical( - "No r-packet escaped through the outer boundary." - ) - self.spectrum_solver.transport_state = transport_state emitted_luminosity = ( @@ -324,13 +324,16 @@ def solve(self): estimated_t_inner, ) = self.get_convergence_estimates(emitted_luminosity) - self.solve_plasma( - transport_state, + self.solve_simulation_state( estimated_t_radiative, estimated_dilution_factor, estimated_t_inner, ) + self.solve_plasma( + transport_state, + ) + converged = self.check_convergence( estimated_t_radiative, estimated_dilution_factor, From 3a47c9f9dd70fcc44de63c36b5098ce743a2f9ae Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 15:02:31 -0400 Subject: [PATCH 10/28] Simplify convergence solver setup with a dict. --- .../workflows/standard_simulation_solver.py | 123 +++++++----------- 1 file changed, 46 insertions(+), 77 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 2791288db03..a485da532bb 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -92,13 +92,14 @@ def __init__(self, configuration): configuration.montecarlo.convergence_strategy ) - self.t_radiative_convergence_solver = ConvergenceSolver( + self.convergence_solvers = {} + self.convergence_solvers["t_radiative"] = ConvergenceSolver( self.convergence_strategy.t_rad ) - self.dilution_factor_convergence_solver = ConvergenceSolver( + self.convergence_solvers["dilution_factor"] = ConvergenceSolver( self.convergence_strategy.w ) - self.t_inner_convergence_solver = ConvergenceSolver( + self.convergence_solvers["t_inner"] = ConvergenceSolver( self.convergence_strategy.t_inner ) @@ -155,58 +156,37 @@ def get_convergence_estimates(self, emitted_luminosity): emitted_luminosity, t_inner_update_exponent=self.convergence_strategy.t_inner_update_exponent, ) - return ( - estimated_t_radiative, - estimated_dilution_factor, - estimated_t_inner, - ) + return { + "t_radiative": estimated_t_radiative, + "dilution_factor": estimated_dilution_factor, + "t_inner": estimated_t_inner, + } def check_convergence( self, - estimated_t_radiative, - estimated_dilution_factor, - estimated_t_inner, + estimated_values, ): - t_radiative_converged = ( - self.t_radiative_convergence_solver.get_convergence_status( - self.simulation_state.t_radiative.value, - estimated_t_radiative.value, - self.simulation_state.no_of_shells, - ) - ) + convergence_statuses = [] - dilution_factor_converged = ( - self.dilution_factor_convergence_solver.get_convergence_status( - self.simulation_state.dilution_factor, - estimated_dilution_factor, - self.simulation_state.no_of_shells, + for key, solver in self.convergence_solvers.items(): + current_value = getattr(self.simulation_state, key) + estimated_value = estimated_values[key] + no_of_shells = ( + self.simulation_state.no_of_shells if key != "t_inner" else 1 ) - ) - - t_inner_converged = ( - self.t_inner_convergence_solver.get_convergence_status( - self.simulation_state.t_inner.value, - estimated_t_inner.value, - 1, + convergence_statuses.append( + solver.get_convergence_status( + current_value, estimated_value, no_of_shells + ) ) - ) - if np.all( - [ - t_radiative_converged, - dilution_factor_converged, - t_inner_converged, - ] - ): + if np.all(convergence_statuses): hold_iterations = self.convergence_strategy.hold_iterations self.consecutive_converges_count += 1 logger.info( f"Iteration converged {self.consecutive_converges_count:d}/{(hold_iterations + 1):d} consecutive " f"times." ) - # If an iteration has converged, require hold_iterations more - # iterations to converge before we conclude that the Simulation - # is converged. return self.consecutive_converges_count == hold_iterations + 1 self.consecutive_converges_count = 0 @@ -214,31 +194,28 @@ def check_convergence( def solve_simulation_state( self, - estimated_t_radiative, - estimated_dilution_factor, - estimated_t_inner, + estimated_values, ): - next_t_radiative = self.t_radiative_convergence_solver.converge( - self.simulation_state.t_radiative, - estimated_t_radiative, - ) - next_dilution_factor = self.dilution_factor_convergence_solver.converge( - self.simulation_state.dilution_factor, - estimated_dilution_factor, - ) - if ( - self.completed_iterations + 1 - ) % self.convergence_strategy.lock_t_inner_cycles == 0: - next_t_inner = self.t_inner_convergence_solver.converge( - self.simulation_state.t_inner, - estimated_t_inner, - ) - else: - next_t_inner = self.simulation_state.t_inner + next_values = {} + + for key, solver in self.convergence_solvers.items(): + if ( + key == "t_inner" + and (self.completed_iterations + 1) + % self.convergence_strategy.lock_t_inner_cycles + != 0 + ): + next_values[key] = getattr(self.simulation_state, key) + else: + next_values[key] = solver.converge( + getattr(self.simulation_state, key), estimated_values[key] + ) - self.simulation_state.t_radiative = next_t_radiative - self.simulation_state.dilution_factor = next_dilution_factor - self.simulation_state.blackbody_packet_source.temperature = next_t_inner + self.simulation_state.t_radiative = next_values["t_radiative"] + self.simulation_state.dilution_factor = next_values["dilution_factor"] + self.simulation_state.blackbody_packet_source.temperature = next_values[ + "t_inner" + ] def solve_plasma( self, @@ -318,27 +295,19 @@ def solve(self): ) ) - ( - estimated_t_radiative, - estimated_dilution_factor, - estimated_t_inner, - ) = self.get_convergence_estimates(emitted_luminosity) + estimated_values = self.get_convergence_estimates( + emitted_luminosity + ) self.solve_simulation_state( - estimated_t_radiative, - estimated_dilution_factor, - estimated_t_inner, + estimated_values, ) self.solve_plasma( transport_state, ) - converged = self.check_convergence( - estimated_t_radiative, - estimated_dilution_factor, - estimated_t_inner, - ) + converged = self.check_convergence(estimated_values) self.completed_iterations += 1 if converged and self.convergence_strategy.stop_if_converged: From 83dd33a43715e1cf6a14f5908573bcdb1e0fd37c Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 23 Jul 2024 15:42:37 -0400 Subject: [PATCH 11/28] Minor formatting change --- tardis/workflows/standard_simulation_solver.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index a485da532bb..7aba3028ae9 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -299,13 +299,9 @@ def solve(self): emitted_luminosity ) - self.solve_simulation_state( - estimated_values, - ) + self.solve_simulation_state(estimated_values) - self.solve_plasma( - transport_state, - ) + self.solve_plasma(transport_state) converged = self.check_convergence(estimated_values) self.completed_iterations += 1 From bf127de161bff84d553bf98d1f4cf6931fc809db Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 24 Jul 2024 13:01:49 -0400 Subject: [PATCH 12/28] Minor refactoring --- .../workflows/standard_simulation_solver.py | 58 ++++++++----------- 1 file changed, 23 insertions(+), 35 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 7aba3028ae9..25f3e958a7d 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -129,33 +129,31 @@ def _get_atom_data(self, configuration): return atom_data - def estimate_t_inner( - self, - input_t_inner, - luminosity_requested, - emitted_luminosity, - t_inner_update_exponent=-0.5, - ): - luminosity_ratios = ( - (emitted_luminosity / luminosity_requested).to(1).value - ) - - return input_t_inner * luminosity_ratios**t_inner_update_exponent - - def get_convergence_estimates(self, emitted_luminosity): + def get_convergence_estimates(self, transport_state): ( estimated_t_radiative, estimated_dilution_factor, - ) = ( - self.transport_solver.transport_state.calculate_radiationfield_properties() + ) = self.transport_solver.transport_state.calculate_radiationfield_properties() + + self.initialize_spectrum_solver( + transport_state, + None, ) - estimated_t_inner = self.estimate_t_inner( - self.simulation_state.t_inner, - self.luminosity_requested, - emitted_luminosity, - t_inner_update_exponent=self.convergence_strategy.t_inner_update_exponent, + emitted_luminosity = self.spectrum_solver.calculate_emitted_luminosity( + self.luminosity_nu_start, self.luminosity_nu_end ) + + luminosity_ratios = ( + (emitted_luminosity / self.luminosity_requested).to(1).value + ) + + estimated_t_inner = ( + self.simulation_state.t_inner + * luminosity_ratios + ** self.convergence_strategy.t_inner_update_exponent + ) + return { "t_radiative": estimated_t_radiative, "dilution_factor": estimated_dilution_factor, @@ -267,9 +265,9 @@ def initialize_spectrum_solver( self.spectrum_solver.transport_state = transport_state if virtual_packet_energies is not None: - self.spectrum_solver._montecarlo_virtual_luminosity.value[ - : - ] = virtual_packet_energies + self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = ( + virtual_packet_energies + ) if self.integrated_spectrum_settings is not None: # Set up spectrum solver integrator @@ -287,17 +285,7 @@ def solve(self): self.real_packet_count ) - self.spectrum_solver.transport_state = transport_state - - emitted_luminosity = ( - self.spectrum_solver.calculate_emitted_luminosity( - self.luminosity_nu_start, self.luminosity_nu_end - ) - ) - - estimated_values = self.get_convergence_estimates( - emitted_luminosity - ) + estimated_values = self.get_convergence_estimates(transport_state) self.solve_simulation_state(estimated_values) From 6b11f0bd4b2205f43b74e134f0e1c4649cad78e2 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 29 Jul 2024 12:16:26 -0400 Subject: [PATCH 13/28] Fixes plasma update step --- tardis/workflows/standard_simulation_solver.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 25f3e958a7d..70448834051 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -7,6 +7,7 @@ from tardis import constants as const from tardis.io.atom_data.base import AtomData 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.spectrum.base import SpectrumSolver @@ -219,9 +220,12 @@ def solve_plasma( self, transport_state, ): + 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. @@ -244,7 +248,6 @@ def solve_montecarlo(self, no_of_real_packets, no_of_virtual_packets=0): virtual_packet_energies = self.transport_solver.run( transport_state, - time_explosion=self.simulation_state.time_explosion, iteration=self.completed_iterations, total_iterations=self.total_iterations, show_progress_bars=False, From 6f705bf969e1ad19009925048facf4f05b909acc Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 12 Aug 2024 12:04:58 -0400 Subject: [PATCH 14/28] Fixes loggers and progress bars Also adds docstrings to methods. Updates some methods to use new functionality of the plasma. Adds requirements for the convergence plots (still broken) --- docs/workflows/standard_workflow.ipynb | 2 +- .../workflows/standard_simulation_solver.py | 383 ++++++++++++++++-- 2 files changed, 358 insertions(+), 27 deletions(-) diff --git a/docs/workflows/standard_workflow.ipynb b/docs/workflows/standard_workflow.ipynb index 2a35c2164a1..f8347a6c563 100644 --- a/docs/workflows/standard_workflow.ipynb +++ b/docs/workflows/standard_workflow.ipynb @@ -25,7 +25,7 @@ "metadata": {}, "outputs": [], "source": [ - "solver = StandardSimulationSolver(config)" + "solver = StandardSimulationSolver(config, show_convergence_plots=True,show_progress_bars=True)" ] }, { diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 70448834051..efe69becc12 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -2,31 +2,43 @@ from pathlib import Path import numpy as np +import pandas as pd from astropy import units as u +from IPython.display import display from tardis import constants as const from tardis.io.atom_data.base import AtomData +from tardis.io.logger.logger import logging_state 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.spectrum.base import SpectrumSolver from tardis.spectrum.formal_integral import FormalIntegrator +from tardis.spectrum.luminosity import ( + calculate_filtered_luminosity, +) from tardis.transport.montecarlo.base import MonteCarloTransportSolver +from tardis.util.base import is_notebook +from tardis.visualization import ConvergencePlots # logging support logger = logging.getLogger(__name__) class StandardSimulationSolver: - def __init__(self, configuration): - # Convergence - self.consecutive_converges_count = 0 - self.converged = False - self.completed_iterations = 0 - self.luminosity_requested = ( - configuration.supernova.luminosity_requested.cgs - ) + def __init__( + self, + configuration, + log_level=None, + specific_log_level=None, + show_progress_bars=False, + show_convergence_plots=False, + convergence_plots_kwargs={}, + ): + # Logging + logging_state(log_level, configuration, specific_log_level) + self.show_progress_bars = show_progress_bars atom_data = self._get_atom_data(configuration) @@ -88,6 +100,14 @@ def __init__(self, configuration): self.integrated_spectrum_settings = configuration.spectrum.integrated self.spectrum_solver = SpectrumSolver.from_config(configuration) + # Convergence settings + self.consecutive_converges_count = 0 + self.converged = False + self.completed_iterations = 0 + self.luminosity_requested = ( + configuration.supernova.luminosity_requested.cgs + ) + # Convergence solvers self.convergence_strategy = ( configuration.montecarlo.convergence_strategy @@ -104,7 +124,52 @@ def __init__(self, configuration): self.convergence_strategy.t_inner ) + # Convergence plots + if show_convergence_plots: + if not is_notebook(): + raise RuntimeError( + "Convergence Plots cannot be displayed in command-line. Set show_convergence_plots " + "to False." + ) + + self.convergence_plots = ConvergencePlots( + iterations=self.total_iterations, **convergence_plots_kwargs + ) + else: + self.convergence_plots = None + + if "export_convergence_plots" in convergence_plots_kwargs: + if not isinstance( + convergence_plots_kwargs["export_convergence_plots"], + bool, + ): + raise TypeError( + "Expected bool in export_convergence_plots argument" + ) + self.export_convergence_plots = convergence_plots_kwargs[ + "export_convergence_plots" + ] + else: + self.export_convergence_plots = False + def _get_atom_data(self, configuration): + """Process atomic data from the configuration + + Parameters + ---------- + configuration : Configuration + TARDIS configuration object + + Returns + ------- + AtomData + Atomic data object + + Raises + ------ + ValueError + If atom data is missing from the configuration + """ if "atom_data" in configuration: if Path(configuration.atom_data).is_absolute(): atom_data_fname = Path(configuration.atom_data) @@ -131,20 +196,57 @@ def _get_atom_data(self, configuration): return atom_data def get_convergence_estimates(self, transport_state): - ( - estimated_t_radiative, - estimated_dilution_factor, - ) = self.transport_solver.transport_state.calculate_radiationfield_properties() + """Compute convergence estimates from the transport state + + Parameters + ---------- + transport_state : MonteCarloTransportState + Transport state object to compute estimates + + Returns + ------- + dict + Convergence estimates + EstimatedRadiationFieldProperties + Dilute radiation file and j_blues dataclass + """ + estimated_radfield_properties = ( + self.transport_solver.radfield_prop_solver.solve( + transport_state.radfield_mc_estimators, + transport_state.time_explosion, + transport_state.time_of_simulation, + transport_state.geometry_state.volume, + transport_state.opacity_state.line_list_nu, + ) + ) + estimated_t_radiative = estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature + estimated_dilution_factor = estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor self.initialize_spectrum_solver( transport_state, None, ) - emitted_luminosity = self.spectrum_solver.calculate_emitted_luminosity( - self.luminosity_nu_start, self.luminosity_nu_end + emitted_luminosity = calculate_filtered_luminosity( + transport_state.emitted_packet_nu, + transport_state.emitted_packet_luminosity, + self.luminosity_nu_start, + self.luminosity_nu_end, + ) + absorbed_luminosity = calculate_filtered_luminosity( + transport_state.reabsorbed_packet_nu, + transport_state.reabsorbed_packet_luminosity, + self.luminosity_nu_start, + self.luminosity_nu_end, ) + if self.convergence_plots is not None: + self.update_convergence_plots( + emitted_luminosity, absorbed_luminosity + ) + + self.log_iteration_results(emitted_luminosity, absorbed_luminosity) + luminosity_ratios = ( (emitted_luminosity / self.luminosity_requested).to(1).value ) @@ -155,16 +257,37 @@ def get_convergence_estimates(self, transport_state): ** self.convergence_strategy.t_inner_update_exponent ) + self.log_plasma_state( + self.simulation_state.t_radiative, + self.simulation_state.dilution_factor, + self.simulation_state.t_inner, + estimated_t_radiative, + estimated_dilution_factor, + estimated_t_inner, + ) + return { "t_radiative": estimated_t_radiative, "dilution_factor": estimated_dilution_factor, "t_inner": estimated_t_inner, - } + }, estimated_radfield_properties def check_convergence( self, estimated_values, ): + """Check convergence status for a dict of estimated values + + Parameters + ---------- + estimated_values : dict + Estimates to check convergence + + Returns + ------- + bool + If convergence has occurred + """ convergence_statuses = [] for key, solver in self.convergence_solvers.items(): @@ -191,10 +314,151 @@ def check_convergence( self.consecutive_converges_count = 0 return False + def update_convergence_plots(self, emitted_luminosity, absorbed_luminosity): + """Updates convergence plotting data + + Parameters + ---------- + emitted_luminosity : Quantity + Current iteration emitted luminosity + absorbed_luminosity : Quantity + Current iteration absorbed luminosity + """ + self.convergence_plots.fetch_data( + name="t_inner", + value=self.simulation_state.t_inner.value, + item_type="value", + ) + self.convergence_plots.fetch_data( + name="t_rad", + value=self.simulation_state.t_radiative, + item_type="iterable", + ) + self.convergence_plots.fetch_data( + name="w", + value=self.simulation_state.dilution_factor, + item_type="iterable", + ) + self.convergence_plots.fetch_data( + name="velocity", + value=self.simulation_state.velocity, + item_type="iterable", + ) + self.convergence_plots.fetch_data( + name="Emitted", + value=emitted_luminosity.value, + item_type="value", + ) + self.convergence_plots.fetch_data( + name="Absorbed", + value=absorbed_luminosity.value, + item_type="value", + ) + self.convergence_plots.fetch_data( + name="Requested", + value=self.luminosity_requested.value, + item_type="value", + ) + + def log_iteration_results(self, emitted_luminosity, absorbed_luminosity): + """Print current iteration information to log at INFO level + + Parameters + ---------- + emitted_luminosity : Quantity + Current iteration emitted luminosity + absorbed_luminosity : Quantity + Current iteration absorbed luminosity + """ + logger.info( + f"\n\tLuminosity emitted = {emitted_luminosity:.3e}\n" + f"\tLuminosity absorbed = {absorbed_luminosity:.3e}\n" + f"\tLuminosity requested = {self.luminosity_requested:.3e}\n" + ) + + def log_plasma_state( + self, + t_rad, + dilution_factor, + t_inner, + next_t_rad, + next_dilution_factor, + next_t_inner, + log_sampling=5, + ): + """ + Logging the change of the plasma state + + Parameters + ---------- + t_rad : astropy.units.Quanity + current t_rad + dilution_factor : np.ndarray + current dilution_factor + next_t_rad : astropy.units.Quanity + next t_rad + next_dilution_factor : np.ndarray + next dilution_factor + log_sampling : int + the n-th shells to be plotted + + Returns + ------- + """ + plasma_state_log = pd.DataFrame( + index=np.arange(len(t_rad)), + columns=["t_rad", "next_t_rad", "w", "next_w"], + ) + plasma_state_log["t_rad"] = t_rad + plasma_state_log["next_t_rad"] = next_t_rad + plasma_state_log["w"] = dilution_factor + plasma_state_log["next_w"] = next_dilution_factor + plasma_state_log.columns.name = "Shell No." + + if is_notebook(): + logger.info("\n\tPlasma stratification:") + + # Displaying the DataFrame only when the logging level is NOTSET, DEBUG or INFO + if logger.level <= logging.INFO: + if not logger.filters: + display( + plasma_state_log.iloc[::log_sampling].style.format( + "{:.3g}" + ) + ) + elif logger.filters[0].log_level == 20: + display( + plasma_state_log.iloc[::log_sampling].style.format( + "{:.3g}" + ) + ) + else: + output_df = "" + plasma_output = plasma_state_log.iloc[::log_sampling].to_string( + float_format=lambda x: f"{x:.3g}", + justify="center", + ) + for value in plasma_output.split("\n"): + output_df = output_df + f"\t{value}\n" + logger.info("\n\tPlasma stratification:") + logger.info(f"\n{output_df}") + + logger.info( + f"\n\tCurrent t_inner = {t_inner:.3f}\n\tExpected t_inner for next iteration = {next_t_inner:.3f}\n" + ) + def solve_simulation_state( self, estimated_values, ): + """Update the simulation state with new inputs computed from previous + iteration estimates. + + Parameters + ---------- + estimated_values : dict + Estimated from the previous iterations + """ next_values = {} for key, solver in self.convergence_solvers.items(): @@ -216,10 +480,19 @@ def solve_simulation_state( "t_inner" ] - def solve_plasma( - self, - transport_state, - ): + def solve_plasma(self, estimated_radfield_properties): + """Update the plasma solution with the new radiation field estimates + + Parameters + ---------- + estimated_radfield_properties : EstimatedRadiationFieldProperties + The radiation field properties to use for updating the plasma + + Raises + ------ + ValueError + If the plasma solver radiative rates type is unknown + """ radiation_field = DilutePlanckianRadiationField( temperature=self.simulation_state.t_radiative, dilution_factor=self.simulation_state.dilution_factor, @@ -229,15 +502,61 @@ def solve_plasma( ) # A check to see if the plasma is set with JBluesDetailed, in which # case it needs some extra kwargs. - if "j_blue_estimator" in self.plasma_solver.outputs_dict: - update_properties.update( - t_inner=self.simulation_state.blackbody_packet_source.temperature, - j_blue_estimator=transport_state.radfield_mc_estimators.j_blue_estimator, + if ( + self.plasma_solver.plasma_solver_settings.RADIATIVE_RATES_TYPE + == "blackbody" + ): + planckian_radiation_field = ( + radiation_field.to_planckian_radiation_field() + ) + j_blues = planckian_radiation_field.calculate_mean_intensity( + self.plasma_solver.atomic_data.lines.nu.values + ) + update_properties["j_blues"] = pd.DataFrame( + j_blues, index=self.plasma_solver.atomic_data.lines.index + ) + elif ( + self.plasma_solver.plasma_solver_settings.RADIATIVE_RATES_TYPE + == "dilute-blackbody" + ): + j_blues = radiation_field.calculate_mean_intensity( + self.plasma_solver.atomic_data.lines.nu.values + ) + update_properties["j_blues"] = pd.DataFrame( + j_blues, index=self.plasma_solver.atomic_data.lines.index + ) + elif ( + self.plasma_solver.plasma_solver_settings.RADIATIVE_RATES_TYPE + == "detailed" + ): + update_properties["j_blues"] = pd.DataFrame( + estimated_radfield_properties.j_blues, + index=self.plasma_solver.atomic_data.lines.index, + ) + else: + raise ValueError( + f"radiative_rates_type type unknown - {self.plasma.plasma_solver_settings.RADIATIVE_RATES_TYPE}" ) self.plasma_solver.update(**update_properties) def solve_montecarlo(self, no_of_real_packets, no_of_virtual_packets=0): + """Solve the MonteCarlo process + + Parameters + ---------- + no_of_real_packets : int + Number of real packets to simulate + no_of_virtual_packets : int, optional + Number of virtual packets to simulate per interaction, by default 0 + + Returns + ------- + MonteCarloTransportState + The new transport state after simulation + ndarray + Array of unnormalized virtual packet energies in each frequency bin + """ transport_state = self.transport_solver.initialize_transport_state( self.simulation_state, self.plasma_solver, @@ -250,7 +569,7 @@ def solve_montecarlo(self, no_of_real_packets, no_of_virtual_packets=0): transport_state, iteration=self.completed_iterations, total_iterations=self.total_iterations, - show_progress_bars=False, + show_progress_bars=self.show_progress_bars, ) output_energy = transport_state.packet_collection.output_energies @@ -264,6 +583,15 @@ def initialize_spectrum_solver( transport_state, virtual_packet_energies=None, ): + """Set up the spectrum solver + + Parameters + ---------- + transport_state : MonteCarloTransportState + The transport state to init with + virtual_packet_energies : ndarray, optional + Array of virtual packet energies binned by frequency, by default None + """ # Set up spectrum solver self.spectrum_solver.transport_state = transport_state @@ -282,17 +610,20 @@ def initialize_spectrum_solver( ) def solve(self): + """Solve the TARDIS simulation until convergence is reached""" converged = False while self.completed_iterations < self.total_iterations - 1: transport_state, virtual_packet_energies = self.solve_montecarlo( self.real_packet_count ) - estimated_values = self.get_convergence_estimates(transport_state) + estimated_values, estimated_radfield_properties = ( + self.get_convergence_estimates(transport_state) + ) self.solve_simulation_state(estimated_values) - self.solve_plasma(transport_state) + self.solve_plasma(estimated_radfield_properties) converged = self.check_convergence(estimated_values) self.completed_iterations += 1 From 0f369af28cfe5b9277235c49b3439d5e4bc7ffd6 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 12 Aug 2024 12:39:00 -0400 Subject: [PATCH 15/28] Fixes convergence plot rendering --- tardis/workflows/standard_simulation_solver.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index efe69becc12..45685e59f11 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -359,6 +359,7 @@ def update_convergence_plots(self, emitted_luminosity, absorbed_luminosity): value=self.luminosity_requested.value, item_type="value", ) + self.convergence_plots.update() def log_iteration_results(self, emitted_luminosity, absorbed_luminosity): """Print current iteration information to log at INFO level @@ -613,6 +614,9 @@ def solve(self): """Solve the TARDIS simulation until convergence is reached""" converged = False while self.completed_iterations < self.total_iterations - 1: + logger.info( + f"\n\tStarting iteration {(self.completed_iterations + 1):d} of {self.total_iterations:d}" + ) transport_state, virtual_packet_energies = self.solve_montecarlo( self.real_packet_count ) @@ -631,6 +635,7 @@ def solve(self): if converged and self.convergence_strategy.stop_if_converged: break + logger.info(f"\n\tStarting final iteration") transport_state, virtual_packet_energies = self.solve_montecarlo( self.final_iteration_packet_count, self.virtual_packet_count ) From 01e5cc320c69ac1d27c2844fdf00847f392ba987 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 12 Aug 2024 12:54:30 -0400 Subject: [PATCH 16/28] Fixes convergence plots in the final iteration --- .../workflows/standard_simulation_solver.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 45685e59f11..aa3fb8a82b7 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -30,6 +30,7 @@ class StandardSimulationSolver: def __init__( self, configuration, + enable_virtual_packet_logging=False, log_level=None, specific_log_level=None, show_progress_bars=False, @@ -57,7 +58,7 @@ def __init__( self.transport_solver = MonteCarloTransportSolver.from_config( configuration, packet_source=self.simulation_state.packet_source, - enable_virtual_packet_logging=False, + enable_virtual_packet_logging=enable_virtual_packet_logging, ) self.luminosity_nu_start = ( @@ -241,7 +242,7 @@ def get_convergence_estimates(self, transport_state): ) if self.convergence_plots is not None: - self.update_convergence_plots( + self.update_convergence_plot_data( emitted_luminosity, absorbed_luminosity ) @@ -314,7 +315,9 @@ def check_convergence( self.consecutive_converges_count = 0 return False - def update_convergence_plots(self, emitted_luminosity, absorbed_luminosity): + def update_convergence_plot_data( + self, emitted_luminosity, absorbed_luminosity + ): """Updates convergence plotting data Parameters @@ -359,7 +362,6 @@ def update_convergence_plots(self, emitted_luminosity, absorbed_luminosity): value=self.luminosity_requested.value, item_type="value", ) - self.convergence_plots.update() def log_iteration_results(self, emitted_luminosity, absorbed_luminosity): """Print current iteration information to log at INFO level @@ -625,6 +627,9 @@ def solve(self): self.get_convergence_estimates(transport_state) ) + if self.convergence_plots is not None: + self.convergence_plots.update() + self.solve_simulation_state(estimated_values) self.solve_plasma(estimated_radfield_properties) @@ -639,6 +644,12 @@ def solve(self): transport_state, virtual_packet_energies = self.solve_montecarlo( self.final_iteration_packet_count, self.virtual_packet_count ) + if self.convergence_plots is not None: + self.get_convergence_estimates(transport_state) + self.convergence_plots.update( + export_convergence_plots=self.export_convergence_plots, + last=True, + ) self.initialize_spectrum_solver( transport_state, virtual_packet_energies, From 2351c8b0a21992d6f2ee97f6a771e1032df1f715 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 12 Aug 2024 13:06:40 -0400 Subject: [PATCH 17/28] black --- .../workflows/standard_simulation_solver.py | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index aa3fb8a82b7..fae2b83e6fe 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -221,8 +221,12 @@ def get_convergence_estimates(self, transport_state): ) ) - estimated_t_radiative = estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature - estimated_dilution_factor = estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor + estimated_t_radiative = ( + estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature + ) + estimated_dilution_factor = ( + estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor + ) self.initialize_spectrum_solver( transport_state, None, @@ -599,9 +603,9 @@ def initialize_spectrum_solver( self.spectrum_solver.transport_state = transport_state if virtual_packet_energies is not None: - self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = ( - virtual_packet_energies - ) + self.spectrum_solver._montecarlo_virtual_luminosity.value[ + : + ] = virtual_packet_energies if self.integrated_spectrum_settings is not None: # Set up spectrum solver integrator @@ -623,9 +627,10 @@ def solve(self): self.real_packet_count ) - estimated_values, estimated_radfield_properties = ( - self.get_convergence_estimates(transport_state) - ) + ( + estimated_values, + estimated_radfield_properties, + ) = self.get_convergence_estimates(transport_state) if self.convergence_plots is not None: self.convergence_plots.update() From 692e3b99c360bd6ce537e935a0d830430a2dae47 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 12 Aug 2024 13:31:54 -0400 Subject: [PATCH 18/28] Simplify convergence plot updating --- .../workflows/standard_simulation_solver.py | 64 ++++++------------- 1 file changed, 19 insertions(+), 45 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index fae2b83e6fe..114b160cb59 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -246,9 +246,16 @@ def get_convergence_estimates(self, transport_state): ) if self.convergence_plots is not None: - self.update_convergence_plot_data( - emitted_luminosity, absorbed_luminosity - ) + plot_data = { + "t_inner": [self.simulation_state.t_inner.value, "value"], + "t_rad": [self.simulation_state.t_radiative, "iterable"], + "w": [self.simulation_state.dilution_factor, "iterable"], + "velocity": [self.simulation_state.velocity, "iterable"], + "Emitted": [emitted_luminosity.value, "value"], + "Absorbed": [absorbed_luminosity.value, "value"], + "Requested": [self.luminosity_requested.value, "value"], + } + self.update_convergence_plot_data(plot_data) self.log_iteration_results(emitted_luminosity, absorbed_luminosity) @@ -319,53 +326,20 @@ def check_convergence( self.consecutive_converges_count = 0 return False - def update_convergence_plot_data( - self, emitted_luminosity, absorbed_luminosity - ): + def update_convergence_plot_data(self, plot_data_dict): """Updates convergence plotting data Parameters ---------- - emitted_luminosity : Quantity - Current iteration emitted luminosity - absorbed_luminosity : Quantity - Current iteration absorbed luminosity + plot_data_dict : dict + Dictionary of data to update of the form {"name": [value, item_type]} """ - self.convergence_plots.fetch_data( - name="t_inner", - value=self.simulation_state.t_inner.value, - item_type="value", - ) - self.convergence_plots.fetch_data( - name="t_rad", - value=self.simulation_state.t_radiative, - item_type="iterable", - ) - self.convergence_plots.fetch_data( - name="w", - value=self.simulation_state.dilution_factor, - item_type="iterable", - ) - self.convergence_plots.fetch_data( - name="velocity", - value=self.simulation_state.velocity, - item_type="iterable", - ) - self.convergence_plots.fetch_data( - name="Emitted", - value=emitted_luminosity.value, - item_type="value", - ) - self.convergence_plots.fetch_data( - name="Absorbed", - value=absorbed_luminosity.value, - item_type="value", - ) - self.convergence_plots.fetch_data( - name="Requested", - value=self.luminosity_requested.value, - item_type="value", - ) + for name, (value, item_type) in plot_data_dict.items(): + self.convergence_plots.fetch_data( + name=name, + value=value, + item_type=item_type, + ) def log_iteration_results(self, emitted_luminosity, absorbed_luminosity): """Print current iteration information to log at INFO level From dfca04e6ccd09b01910d84a013a5483f1a826112 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 12 Aug 2024 14:01:52 -0400 Subject: [PATCH 19/28] Move logging handling to a separate class --- .../workflows/standard_simulation_solver.py | 109 ++--------------- tardis/workflows/workflow_logging.py | 111 ++++++++++++++++++ 2 files changed, 123 insertions(+), 97 deletions(-) create mode 100644 tardis/workflows/workflow_logging.py diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 114b160cb59..1d2102950c5 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -4,11 +4,9 @@ import numpy as np import pandas as pd from astropy import units as u -from IPython.display import display from tardis import constants as const from tardis.io.atom_data.base import AtomData -from tardis.io.logger.logger import logging_state from tardis.model import SimulationState from tardis.plasma.radiation_field import DilutePlanckianRadiationField from tardis.plasma.standard_plasmas import assemble_plasma @@ -21,12 +19,13 @@ from tardis.transport.montecarlo.base import MonteCarloTransportSolver from tardis.util.base import is_notebook from tardis.visualization import ConvergencePlots +from tardis.workflows.workflow_logging import WorkflowLogging # logging support logger = logging.getLogger(__name__) -class StandardSimulationSolver: +class StandardSimulationSolver(WorkflowLogging): def __init__( self, configuration, @@ -37,8 +36,12 @@ def __init__( show_convergence_plots=False, convergence_plots_kwargs={}, ): - # Logging - logging_state(log_level, configuration, specific_log_level) + # set up logging + super().__init__( + configuration, + log_level=log_level, + specific_log_level=specific_log_level, + ) self.show_progress_bars = show_progress_bars atom_data = self._get_atom_data(configuration) @@ -61,6 +64,7 @@ def __init__( enable_virtual_packet_logging=enable_virtual_packet_logging, ) + # Luminosity filter frequencies self.luminosity_nu_start = ( configuration.supernova.luminosity_wavelength_end.to( u.Hz, u.spectral() @@ -227,10 +231,6 @@ def get_convergence_estimates(self, transport_state): estimated_dilution_factor = ( estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor ) - self.initialize_spectrum_solver( - transport_state, - None, - ) emitted_luminosity = calculate_filtered_luminosity( transport_state.emitted_packet_nu, @@ -257,7 +257,9 @@ def get_convergence_estimates(self, transport_state): } self.update_convergence_plot_data(plot_data) - self.log_iteration_results(emitted_luminosity, absorbed_luminosity) + self.log_iteration_results( + emitted_luminosity, absorbed_luminosity, self.luminosity_requested + ) luminosity_ratios = ( (emitted_luminosity / self.luminosity_requested).to(1).value @@ -341,93 +343,6 @@ def update_convergence_plot_data(self, plot_data_dict): item_type=item_type, ) - def log_iteration_results(self, emitted_luminosity, absorbed_luminosity): - """Print current iteration information to log at INFO level - - Parameters - ---------- - emitted_luminosity : Quantity - Current iteration emitted luminosity - absorbed_luminosity : Quantity - Current iteration absorbed luminosity - """ - logger.info( - f"\n\tLuminosity emitted = {emitted_luminosity:.3e}\n" - f"\tLuminosity absorbed = {absorbed_luminosity:.3e}\n" - f"\tLuminosity requested = {self.luminosity_requested:.3e}\n" - ) - - def log_plasma_state( - self, - t_rad, - dilution_factor, - t_inner, - next_t_rad, - next_dilution_factor, - next_t_inner, - log_sampling=5, - ): - """ - Logging the change of the plasma state - - Parameters - ---------- - t_rad : astropy.units.Quanity - current t_rad - dilution_factor : np.ndarray - current dilution_factor - next_t_rad : astropy.units.Quanity - next t_rad - next_dilution_factor : np.ndarray - next dilution_factor - log_sampling : int - the n-th shells to be plotted - - Returns - ------- - """ - plasma_state_log = pd.DataFrame( - index=np.arange(len(t_rad)), - columns=["t_rad", "next_t_rad", "w", "next_w"], - ) - plasma_state_log["t_rad"] = t_rad - plasma_state_log["next_t_rad"] = next_t_rad - plasma_state_log["w"] = dilution_factor - plasma_state_log["next_w"] = next_dilution_factor - plasma_state_log.columns.name = "Shell No." - - if is_notebook(): - logger.info("\n\tPlasma stratification:") - - # Displaying the DataFrame only when the logging level is NOTSET, DEBUG or INFO - if logger.level <= logging.INFO: - if not logger.filters: - display( - plasma_state_log.iloc[::log_sampling].style.format( - "{:.3g}" - ) - ) - elif logger.filters[0].log_level == 20: - display( - plasma_state_log.iloc[::log_sampling].style.format( - "{:.3g}" - ) - ) - else: - output_df = "" - plasma_output = plasma_state_log.iloc[::log_sampling].to_string( - float_format=lambda x: f"{x:.3g}", - justify="center", - ) - for value in plasma_output.split("\n"): - output_df = output_df + f"\t{value}\n" - logger.info("\n\tPlasma stratification:") - logger.info(f"\n{output_df}") - - logger.info( - f"\n\tCurrent t_inner = {t_inner:.3f}\n\tExpected t_inner for next iteration = {next_t_inner:.3f}\n" - ) - def solve_simulation_state( self, estimated_values, diff --git a/tardis/workflows/workflow_logging.py b/tardis/workflows/workflow_logging.py new file mode 100644 index 00000000000..1a77b4c9406 --- /dev/null +++ b/tardis/workflows/workflow_logging.py @@ -0,0 +1,111 @@ +import logging + +import numpy as np +import pandas as pd +from IPython.display import display + +from tardis.io.logger.logger import logging_state +from tardis.util.base import is_notebook + +logger = logging.getLogger(__name__) + + +class WorkflowLogging: + def __init__( + self, + configuration, + log_level=None, + specific_log_level=None, + ): + logging_state(log_level, configuration, specific_log_level) + + def log_iteration_results( + self, emitted_luminosity, absorbed_luminosity, luminosity_requested + ): + """Print current iteration information to log at INFO level + + Parameters + ---------- + emitted_luminosity : Quantity + Current iteration emitted luminosity + absorbed_luminosity : Quantity + Current iteration absorbed luminosity + luminosity_requested : Quantity + The requested luminosity for the simulation + """ + logger.info( + f"\n\tLuminosity emitted = {emitted_luminosity:.3e}\n" + f"\tLuminosity absorbed = {absorbed_luminosity:.3e}\n" + f"\tLuminosity requested = {luminosity_requested:.3e}\n" + ) + + def log_plasma_state( + self, + t_rad, + dilution_factor, + t_inner, + next_t_rad, + next_dilution_factor, + next_t_inner, + log_sampling=5, + ): + """ + Logging the change of the plasma state + + Parameters + ---------- + t_rad : astropy.units.Quanity + current t_rad + dilution_factor : np.ndarray + current dilution_factor + next_t_rad : astropy.units.Quanity + next t_rad + next_dilution_factor : np.ndarray + next dilution_factor + log_sampling : int + the n-th shells to be plotted + + Returns + ------- + """ + plasma_state_log = pd.DataFrame( + index=np.arange(len(t_rad)), + columns=["t_rad", "next_t_rad", "w", "next_w"], + ) + plasma_state_log["t_rad"] = t_rad + plasma_state_log["next_t_rad"] = next_t_rad + plasma_state_log["w"] = dilution_factor + plasma_state_log["next_w"] = next_dilution_factor + plasma_state_log.columns.name = "Shell No." + + if is_notebook(): + logger.info("\n\tPlasma stratification:") + + # Displaying the DataFrame only when the logging level is NOTSET, DEBUG or INFO + if logger.level <= logging.INFO: + if not logger.filters: + display( + plasma_state_log.iloc[::log_sampling].style.format( + "{:.3g}" + ) + ) + elif logger.filters[0].log_level == 20: + display( + plasma_state_log.iloc[::log_sampling].style.format( + "{:.3g}" + ) + ) + else: + output_df = "" + plasma_output = plasma_state_log.iloc[::log_sampling].to_string( + float_format=lambda x: f"{x:.3g}", + justify="center", + ) + for value in plasma_output.split("\n"): + output_df = output_df + f"\t{value}\n" + logger.info("\n\tPlasma stratification:") + logger.info(f"\n{output_df}") + + logger.info( + f"\n\tCurrent t_inner = {t_inner:.3f}\n\tExpected t_inner for next iteration = {next_t_inner:.3f}\n" + ) From b67f6cb2b7fb320e84d40e90a7b2430dc32c64de Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 12 Aug 2024 14:27:55 -0400 Subject: [PATCH 20/28] Add HDF output capability to solver --- .../workflows/standard_simulation_solver.py | 45 +++++++++++++++++-- 1 file changed, 42 insertions(+), 3 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 1d2102950c5..5ffb6ebda5a 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -7,9 +7,11 @@ from tardis import constants as const from tardis.io.atom_data.base import AtomData +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.base import PlasmaStateStorerMixin from tardis.simulation.convergence import ConvergenceSolver from tardis.spectrum.base import SpectrumSolver from tardis.spectrum.formal_integral import FormalIntegrator @@ -25,7 +27,20 @@ logger = logging.getLogger(__name__) -class StandardSimulationSolver(WorkflowLogging): +class StandardSimulationSolver( + WorkflowLogging, PlasmaStateStorerMixin, HDFWriterMixin +): + hdf_properties = [ + "simulation_state", + "plasma_solver", + "transport_solver", + "iterations_w", + "iterations_t_rad", + "iterations_electron_densities", + "iterations_t_inner", + "spectrum_solver", + ] + def __init__( self, configuration, @@ -37,11 +52,13 @@ def __init__( convergence_plots_kwargs={}, ): # set up logging - super().__init__( - configuration, + WorkflowLogging.__init__( + self, + configuration=configuration, log_level=log_level, specific_log_level=specific_log_level, ) + self.show_progress_bars = show_progress_bars atom_data = self._get_atom_data(configuration) @@ -101,6 +118,13 @@ def __init__( configuration.montecarlo.no_of_virtual_packets ) + # set up plasma storage + PlasmaStateStorerMixin.__init__( + self, + iterations=self.total_iterations, + no_of_shells=self.simulation_state.no_of_shells, + ) + # spectrum settings self.integrated_spectrum_settings = configuration.spectrum.integrated self.spectrum_solver = SpectrumSolver.from_config(configuration) @@ -512,6 +536,13 @@ def solve(self): logger.info( f"\n\tStarting iteration {(self.completed_iterations + 1):d} of {self.total_iterations:d}" ) + self.store_plasma_state( + self.completed_iterations, + self.simulation_state.dilution_factor, + self.simulation_state.t_radiative, + self.plasma_solver.electron_densities, + self.simulation_state.t_inner, + ) transport_state, virtual_packet_energies = self.solve_montecarlo( self.real_packet_count ) @@ -538,6 +569,14 @@ def solve(self): transport_state, virtual_packet_energies = self.solve_montecarlo( self.final_iteration_packet_count, self.virtual_packet_count ) + self.store_plasma_state( + self.completed_iterations, + self.simulation_state.dilution_factor, + self.simulation_state.t_radiative, + self.plasma_solver.electron_densities, + self.simulation_state.t_inner, + ) + self.reshape_plasma_state_store(self.completed_iterations) if self.convergence_plots is not None: self.get_convergence_estimates(transport_state) self.convergence_plots.update( From f9f31d414cf6c8c1da7af555a3ce95177d5f4392 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 12 Aug 2024 15:23:15 -0400 Subject: [PATCH 21/28] Move more basic logging back into workflow --- .../workflows/standard_simulation_solver.py | 6 ++++-- tardis/workflows/workflow_logging.py | 20 ------------------- 2 files changed, 4 insertions(+), 22 deletions(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 5ffb6ebda5a..9101c8f13a9 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -281,8 +281,10 @@ def get_convergence_estimates(self, transport_state): } self.update_convergence_plot_data(plot_data) - self.log_iteration_results( - emitted_luminosity, absorbed_luminosity, self.luminosity_requested + logger.info( + f"\n\tLuminosity emitted = {emitted_luminosity:.3e}\n" + f"\tLuminosity absorbed = {absorbed_luminosity:.3e}\n" + f"\tLuminosity requested = {self.luminosity_requested:.3e}\n" ) luminosity_ratios = ( diff --git a/tardis/workflows/workflow_logging.py b/tardis/workflows/workflow_logging.py index 1a77b4c9406..2dc464fe3b9 100644 --- a/tardis/workflows/workflow_logging.py +++ b/tardis/workflows/workflow_logging.py @@ -19,26 +19,6 @@ def __init__( ): logging_state(log_level, configuration, specific_log_level) - def log_iteration_results( - self, emitted_luminosity, absorbed_luminosity, luminosity_requested - ): - """Print current iteration information to log at INFO level - - Parameters - ---------- - emitted_luminosity : Quantity - Current iteration emitted luminosity - absorbed_luminosity : Quantity - Current iteration absorbed luminosity - luminosity_requested : Quantity - The requested luminosity for the simulation - """ - logger.info( - f"\n\tLuminosity emitted = {emitted_luminosity:.3e}\n" - f"\tLuminosity absorbed = {absorbed_luminosity:.3e}\n" - f"\tLuminosity requested = {luminosity_requested:.3e}\n" - ) - def log_plasma_state( self, t_rad, From 4de0e1f77870a8a472c10f077520274bc3891d73 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 12 Aug 2024 15:25:05 -0400 Subject: [PATCH 22/28] Add not-converged error message --- tardis/workflows/standard_simulation_solver.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/standard_simulation_solver.py index 9101c8f13a9..8af19ce8704 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/standard_simulation_solver.py @@ -567,7 +567,12 @@ def solve(self): if converged and self.convergence_strategy.stop_if_converged: break - logger.info(f"\n\tStarting final iteration") + if converged: + logger.info("\n\tStarting final iteration") + else: + logger.error( + "\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 ) From 94f749b836095a0441e4c22544f74c37d2478fc3 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 12 Aug 2024 15:26:56 -0400 Subject: [PATCH 23/28] Update notebook with export option --- docs/workflows/standard_workflow.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/workflows/standard_workflow.ipynb b/docs/workflows/standard_workflow.ipynb index f8347a6c563..4577f8eb574 100644 --- a/docs/workflows/standard_workflow.ipynb +++ b/docs/workflows/standard_workflow.ipynb @@ -25,7 +25,7 @@ "metadata": {}, "outputs": [], "source": [ - "solver = StandardSimulationSolver(config, show_convergence_plots=True,show_progress_bars=True)" + "solver = StandardSimulationSolver(config, show_convergence_plots=True,show_progress_bars=True,convergence_plots_kwargs={\"export_convergence_plots\":True})" ] }, { From c2a636a0f7f9f6eceede6721c58a60c1271c071c Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 13 Aug 2024 16:00:02 -0400 Subject: [PATCH 24/28] Added simple base workflow and changed some verbiage --- docs/workflows/simple_workflow.ipynb | 110 ++++++++ docs/workflows/standard_workflow.ipynb | 12 +- ...ulation_solver.py => simple_simulation.py} | 173 ++----------- tardis/workflows/standard_simulation.py | 243 ++++++++++++++++++ 4 files changed, 376 insertions(+), 162 deletions(-) create mode 100644 docs/workflows/simple_workflow.ipynb rename tardis/workflows/{standard_simulation_solver.py => simple_simulation.py} (72%) create mode 100644 tardis/workflows/standard_simulation.py diff --git a/docs/workflows/simple_workflow.ipynb b/docs/workflows/simple_workflow.ipynb new file mode 100644 index 00000000000..99ed4c02141 --- /dev/null +++ b/docs/workflows/simple_workflow.ipynb @@ -0,0 +1,110 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from tardis.workflows.simple_simulation import SimpleSimulation\n", + "from tardis.io.configuration.config_reader import Configuration" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = Configuration.from_yaml('../tardis_example.yml')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "workflow = SimpleSimulation(config)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "workflow.run()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "spectrum = workflow.spectrum_solver.spectrum_real_packets\n", + "spectrum_virtual = workflow.spectrum_solver.spectrum_virtual_packets\n", + "spectrum_integrated = workflow.spectrum_solver.spectrum_integrated" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%matplotlib inline\n", + "plt.figure(figsize=(10, 6.5))\n", + "\n", + "spectrum.plot(label=\"Normal packets\")\n", + "spectrum_virtual.plot(label=\"Virtual packets\")\n", + "spectrum_integrated.plot(label='Formal integral')\n", + "\n", + "plt.xlim(500, 9000)\n", + "plt.title(\"TARDIS example model spectrum\")\n", + "plt.xlabel(\"Wavelength [$\\AA$]\")\n", + "plt.ylabel(\"Luminosity density [erg/s/$\\AA$]\")\n", + "plt.legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "tardis", + "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.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/workflows/standard_workflow.ipynb b/docs/workflows/standard_workflow.ipynb index 4577f8eb574..c5e64233cb7 100644 --- a/docs/workflows/standard_workflow.ipynb +++ b/docs/workflows/standard_workflow.ipynb @@ -6,7 +6,7 @@ "metadata": {}, "outputs": [], "source": [ - "from tardis.workflows.standard_simulation_solver import StandardSimulationSolver\n", + "from tardis.workflows.standard_simulation import StandardSimulation\n", "from tardis.io.configuration.config_reader import Configuration" ] }, @@ -25,7 +25,7 @@ "metadata": {}, "outputs": [], "source": [ - "solver = StandardSimulationSolver(config, show_convergence_plots=True,show_progress_bars=True,convergence_plots_kwargs={\"export_convergence_plots\":True})" + "workflow = StandardSimulation(config, show_convergence_plots=True,show_progress_bars=True,convergence_plots_kwargs={\"export_convergence_plots\":True})" ] }, { @@ -34,7 +34,7 @@ "metadata": {}, "outputs": [], "source": [ - "solver.solve()" + "workflow.run()" ] }, { @@ -52,9 +52,9 @@ "metadata": {}, "outputs": [], "source": [ - "spectrum = solver.spectrum_solver.spectrum_real_packets\n", - "spectrum_virtual = solver.spectrum_solver.spectrum_virtual_packets\n", - "spectrum_integrated = solver.spectrum_solver.spectrum_integrated" + "spectrum = workflow.spectrum_solver.spectrum_real_packets\n", + "spectrum_virtual = workflow.spectrum_solver.spectrum_virtual_packets\n", + "spectrum_integrated = workflow.spectrum_solver.spectrum_integrated" ] }, { diff --git a/tardis/workflows/standard_simulation_solver.py b/tardis/workflows/simple_simulation.py similarity index 72% rename from tardis/workflows/standard_simulation_solver.py rename to tardis/workflows/simple_simulation.py index 8af19ce8704..608fb0fbd38 100644 --- a/tardis/workflows/standard_simulation_solver.py +++ b/tardis/workflows/simple_simulation.py @@ -7,11 +7,9 @@ from tardis import constants as const from tardis.io.atom_data.base import AtomData -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.base import PlasmaStateStorerMixin from tardis.simulation.convergence import ConvergenceSolver from tardis.spectrum.base import SpectrumSolver from tardis.spectrum.formal_integral import FormalIntegrator @@ -20,47 +18,20 @@ ) from tardis.transport.montecarlo.base import MonteCarloTransportSolver from tardis.util.base import is_notebook -from tardis.visualization import ConvergencePlots from tardis.workflows.workflow_logging import WorkflowLogging # logging support logger = logging.getLogger(__name__) -class StandardSimulationSolver( - WorkflowLogging, PlasmaStateStorerMixin, HDFWriterMixin -): - hdf_properties = [ - "simulation_state", - "plasma_solver", - "transport_solver", - "iterations_w", - "iterations_t_rad", - "iterations_electron_densities", - "iterations_t_inner", - "spectrum_solver", - ] - - def __init__( - self, - configuration, - enable_virtual_packet_logging=False, - log_level=None, - specific_log_level=None, - show_progress_bars=False, - show_convergence_plots=False, - convergence_plots_kwargs={}, - ): - # set up logging - WorkflowLogging.__init__( - self, - configuration=configuration, - log_level=log_level, - specific_log_level=specific_log_level, - ) - - self.show_progress_bars = show_progress_bars +class SimpleSimulation(WorkflowLogging): + show_progress_bars = is_notebook() + enable_virtual_packet_logging = False + log_level = None + specific_log_level = None + def __init__(self, configuration): + super().__init__(configuration, self.log_level, self.specific_log_level) atom_data = self._get_atom_data(configuration) # set up states and solvers @@ -78,7 +49,7 @@ def __init__( self.transport_solver = MonteCarloTransportSolver.from_config( configuration, packet_source=self.simulation_state.packet_source, - enable_virtual_packet_logging=enable_virtual_packet_logging, + enable_virtual_packet_logging=self.enable_virtual_packet_logging, ) # Luminosity filter frequencies @@ -118,13 +89,6 @@ def __init__( configuration.montecarlo.no_of_virtual_packets ) - # set up plasma storage - PlasmaStateStorerMixin.__init__( - self, - iterations=self.total_iterations, - no_of_shells=self.simulation_state.no_of_shells, - ) - # spectrum settings self.integrated_spectrum_settings = configuration.spectrum.integrated self.spectrum_solver = SpectrumSolver.from_config(configuration) @@ -153,34 +117,6 @@ def __init__( self.convergence_strategy.t_inner ) - # Convergence plots - if show_convergence_plots: - if not is_notebook(): - raise RuntimeError( - "Convergence Plots cannot be displayed in command-line. Set show_convergence_plots " - "to False." - ) - - self.convergence_plots = ConvergencePlots( - iterations=self.total_iterations, **convergence_plots_kwargs - ) - else: - self.convergence_plots = None - - if "export_convergence_plots" in convergence_plots_kwargs: - if not isinstance( - convergence_plots_kwargs["export_convergence_plots"], - bool, - ): - raise TypeError( - "Expected bool in export_convergence_plots argument" - ) - self.export_convergence_plots = convergence_plots_kwargs[ - "export_convergence_plots" - ] - else: - self.export_convergence_plots = False - def _get_atom_data(self, configuration): """Process atomic data from the configuration @@ -249,12 +185,8 @@ def get_convergence_estimates(self, transport_state): ) ) - estimated_t_radiative = ( - estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature - ) - estimated_dilution_factor = ( - estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor - ) + estimated_t_radiative = estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature + estimated_dilution_factor = estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor emitted_luminosity = calculate_filtered_luminosity( transport_state.emitted_packet_nu, @@ -262,30 +194,6 @@ def get_convergence_estimates(self, transport_state): self.luminosity_nu_start, self.luminosity_nu_end, ) - absorbed_luminosity = calculate_filtered_luminosity( - transport_state.reabsorbed_packet_nu, - transport_state.reabsorbed_packet_luminosity, - self.luminosity_nu_start, - self.luminosity_nu_end, - ) - - if self.convergence_plots is not None: - plot_data = { - "t_inner": [self.simulation_state.t_inner.value, "value"], - "t_rad": [self.simulation_state.t_radiative, "iterable"], - "w": [self.simulation_state.dilution_factor, "iterable"], - "velocity": [self.simulation_state.velocity, "iterable"], - "Emitted": [emitted_luminosity.value, "value"], - "Absorbed": [absorbed_luminosity.value, "value"], - "Requested": [self.luminosity_requested.value, "value"], - } - self.update_convergence_plot_data(plot_data) - - logger.info( - f"\n\tLuminosity emitted = {emitted_luminosity:.3e}\n" - f"\tLuminosity absorbed = {absorbed_luminosity:.3e}\n" - f"\tLuminosity requested = {self.luminosity_requested:.3e}\n" - ) luminosity_ratios = ( (emitted_luminosity / self.luminosity_requested).to(1).value @@ -297,15 +205,6 @@ def get_convergence_estimates(self, transport_state): ** self.convergence_strategy.t_inner_update_exponent ) - self.log_plasma_state( - self.simulation_state.t_radiative, - self.simulation_state.dilution_factor, - self.simulation_state.t_inner, - estimated_t_radiative, - estimated_dilution_factor, - estimated_t_inner, - ) - return { "t_radiative": estimated_t_radiative, "dilution_factor": estimated_dilution_factor, @@ -349,26 +248,11 @@ def check_convergence( f"Iteration converged {self.consecutive_converges_count:d}/{(hold_iterations + 1):d} consecutive " f"times." ) - return self.consecutive_converges_count == hold_iterations + 1 + return self.consecutive_converges_count >= hold_iterations + 1 self.consecutive_converges_count = 0 return False - def update_convergence_plot_data(self, plot_data_dict): - """Updates convergence plotting data - - Parameters - ---------- - plot_data_dict : dict - Dictionary of data to update of the form {"name": [value, item_type]} - """ - for name, (value, item_type) in plot_data_dict.items(): - self.convergence_plots.fetch_data( - name=name, - value=value, - item_type=item_type, - ) - def solve_simulation_state( self, estimated_values, @@ -518,9 +402,9 @@ def initialize_spectrum_solver( self.spectrum_solver.transport_state = transport_state if virtual_packet_energies is not None: - self.spectrum_solver._montecarlo_virtual_luminosity.value[ - : - ] = virtual_packet_energies + self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = ( + virtual_packet_energies + ) if self.integrated_spectrum_settings is not None: # Set up spectrum solver integrator @@ -531,20 +415,13 @@ def initialize_spectrum_solver( self.simulation_state, self.plasma_solver, self.transport_solver ) - def solve(self): - """Solve the TARDIS simulation until convergence is reached""" + def run(self): + """Run the TARDIS simulation until convergence is reached""" converged = False while self.completed_iterations < self.total_iterations - 1: logger.info( f"\n\tStarting iteration {(self.completed_iterations + 1):d} of {self.total_iterations:d}" ) - self.store_plasma_state( - self.completed_iterations, - self.simulation_state.dilution_factor, - self.simulation_state.t_radiative, - self.plasma_solver.electron_densities, - self.simulation_state.t_inner, - ) transport_state, virtual_packet_energies = self.solve_montecarlo( self.real_packet_count ) @@ -554,9 +431,6 @@ def solve(self): estimated_radfield_properties, ) = self.get_convergence_estimates(transport_state) - if self.convergence_plots is not None: - self.convergence_plots.update() - self.solve_simulation_state(estimated_values) self.solve_plasma(estimated_radfield_properties) @@ -576,20 +450,7 @@ def solve(self): transport_state, virtual_packet_energies = self.solve_montecarlo( self.final_iteration_packet_count, self.virtual_packet_count ) - self.store_plasma_state( - self.completed_iterations, - self.simulation_state.dilution_factor, - self.simulation_state.t_radiative, - self.plasma_solver.electron_densities, - self.simulation_state.t_inner, - ) - self.reshape_plasma_state_store(self.completed_iterations) - if self.convergence_plots is not None: - self.get_convergence_estimates(transport_state) - self.convergence_plots.update( - export_convergence_plots=self.export_convergence_plots, - last=True, - ) + self.initialize_spectrum_solver( transport_state, virtual_packet_energies, diff --git a/tardis/workflows/standard_simulation.py b/tardis/workflows/standard_simulation.py new file mode 100644 index 00000000000..8114a506b5d --- /dev/null +++ b/tardis/workflows/standard_simulation.py @@ -0,0 +1,243 @@ +import logging + +from tardis.io.util import HDFWriterMixin +from tardis.simulation.base import PlasmaStateStorerMixin +from tardis.spectrum.luminosity import ( + calculate_filtered_luminosity, +) +from tardis.util.base import is_notebook +from tardis.visualization import ConvergencePlots +from tardis.workflows.simple_simulation import SimpleSimulation + +# logging support +logger = logging.getLogger(__name__) + + +class StandardSimulation( + SimpleSimulation, PlasmaStateStorerMixin, HDFWriterMixin +): + hdf_properties = [ + "simulation_state", + "plasma_solver", + "transport_solver", + "iterations_w", + "iterations_t_rad", + "iterations_electron_densities", + "iterations_t_inner", + "spectrum_solver", + ] + + def __init__( + self, + configuration, + enable_virtual_packet_logging=False, + log_level=None, + specific_log_level=None, + show_progress_bars=False, + show_convergence_plots=False, + convergence_plots_kwargs={}, + ): + self.show_progress_bars = show_progress_bars + self.log_level = log_level + self.specific_log_level = specific_log_level + self.enable_virtual_packet_logging = enable_virtual_packet_logging + + SimpleSimulationWorkflow.__init__(self, configuration) + + # set up plasma storage + PlasmaStateStorerMixin.__init__( + self, + iterations=self.total_iterations, + no_of_shells=self.simulation_state.no_of_shells, + ) + + # Convergence plots + if show_convergence_plots: + if not is_notebook(): + raise RuntimeError( + "Convergence Plots cannot be displayed in command-line. Set show_convergence_plots " + "to False." + ) + + self.convergence_plots = ConvergencePlots( + iterations=self.total_iterations, **convergence_plots_kwargs + ) + else: + self.convergence_plots = None + + if "export_convergence_plots" in convergence_plots_kwargs: + if not isinstance( + convergence_plots_kwargs["export_convergence_plots"], + bool, + ): + raise TypeError( + "Expected bool in export_convergence_plots argument" + ) + self.export_convergence_plots = convergence_plots_kwargs[ + "export_convergence_plots" + ] + else: + self.export_convergence_plots = False + + def get_convergence_estimates(self, transport_state): + """Compute convergence estimates from the transport state + + Parameters + ---------- + transport_state : MonteCarloTransportState + Transport state object to compute estimates + + Returns + ------- + dict + Convergence estimates + EstimatedRadiationFieldProperties + Dilute radiation file and j_blues dataclass + """ + estimated_radfield_properties = ( + self.transport_solver.radfield_prop_solver.solve( + transport_state.radfield_mc_estimators, + transport_state.time_explosion, + transport_state.time_of_simulation, + transport_state.geometry_state.volume, + transport_state.opacity_state.line_list_nu, + ) + ) + + estimated_t_radiative = estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature + estimated_dilution_factor = estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor + + emitted_luminosity = calculate_filtered_luminosity( + transport_state.emitted_packet_nu, + transport_state.emitted_packet_luminosity, + self.luminosity_nu_start, + self.luminosity_nu_end, + ) + absorbed_luminosity = calculate_filtered_luminosity( + transport_state.reabsorbed_packet_nu, + transport_state.reabsorbed_packet_luminosity, + self.luminosity_nu_start, + self.luminosity_nu_end, + ) + + if self.convergence_plots is not None: + plot_data = { + "t_inner": [self.simulation_state.t_inner.value, "value"], + "t_rad": [self.simulation_state.t_radiative, "iterable"], + "w": [self.simulation_state.dilution_factor, "iterable"], + "velocity": [self.simulation_state.velocity, "iterable"], + "Emitted": [emitted_luminosity.value, "value"], + "Absorbed": [absorbed_luminosity.value, "value"], + "Requested": [self.luminosity_requested.value, "value"], + } + self.update_convergence_plot_data(plot_data) + + logger.info( + f"\n\tLuminosity emitted = {emitted_luminosity:.3e}\n" + f"\tLuminosity absorbed = {absorbed_luminosity:.3e}\n" + f"\tLuminosity requested = {self.luminosity_requested:.3e}\n" + ) + + luminosity_ratios = ( + (emitted_luminosity / self.luminosity_requested).to(1).value + ) + + estimated_t_inner = ( + self.simulation_state.t_inner + * luminosity_ratios + ** self.convergence_strategy.t_inner_update_exponent + ) + + self.log_plasma_state( + self.simulation_state.t_radiative, + self.simulation_state.dilution_factor, + self.simulation_state.t_inner, + estimated_t_radiative, + estimated_dilution_factor, + estimated_t_inner, + ) + + return { + "t_radiative": estimated_t_radiative, + "dilution_factor": estimated_dilution_factor, + "t_inner": estimated_t_inner, + }, estimated_radfield_properties + + def update_convergence_plot_data(self, plot_data_dict): + """Updates convergence plotting data + + Parameters + ---------- + plot_data_dict : dict + Dictionary of data to update of the form {"name": [value, item_type]} + """ + for name, (value, item_type) in plot_data_dict.items(): + self.convergence_plots.fetch_data( + name=name, + value=value, + item_type=item_type, + ) + + def run(self): + """Run the TARDIS simulation until convergence is reached""" + converged = False + while self.completed_iterations < self.total_iterations - 1: + logger.info( + f"\n\tStarting iteration {(self.completed_iterations + 1):d} of {self.total_iterations:d}" + ) + self.store_plasma_state( + self.completed_iterations, + self.simulation_state.dilution_factor, + self.simulation_state.t_radiative, + self.plasma_solver.electron_densities, + self.simulation_state.t_inner, + ) + transport_state, virtual_packet_energies = self.solve_montecarlo( + self.real_packet_count + ) + + ( + estimated_values, + estimated_radfield_properties, + ) = self.get_convergence_estimates(transport_state) + + if self.convergence_plots is not None: + self.convergence_plots.update() + + self.solve_simulation_state(estimated_values) + + self.solve_plasma(estimated_radfield_properties) + + converged = self.check_convergence(estimated_values) + self.completed_iterations += 1 + + if converged and self.convergence_strategy.stop_if_converged: + break + + if converged: + logger.info("\n\tStarting final iteration") + else: + logger.error( + "\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 + ) + self.store_plasma_state( + self.completed_iterations, + self.simulation_state.dilution_factor, + self.simulation_state.t_radiative, + self.plasma_solver.electron_densities, + self.simulation_state.t_inner, + ) + self.reshape_plasma_state_store(self.completed_iterations) + if self.convergence_plots is not None: + self.get_convergence_estimates(transport_state) + self.convergence_plots.update( + export_convergence_plots=self.export_convergence_plots, + last=True, + ) + self.initialize_spectrum_solver( + transport_state, + virtual_packet_energies, + ) From 011c4153f07c34ab6b976e5007d0c50f73ffe56e Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 13 Aug 2024 16:01:44 -0400 Subject: [PATCH 25/28] Fix typo --- tardis/workflows/standard_simulation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/workflows/standard_simulation.py b/tardis/workflows/standard_simulation.py index 8114a506b5d..429752b543c 100644 --- a/tardis/workflows/standard_simulation.py +++ b/tardis/workflows/standard_simulation.py @@ -42,7 +42,7 @@ def __init__( self.specific_log_level = specific_log_level self.enable_virtual_packet_logging = enable_virtual_packet_logging - SimpleSimulationWorkflow.__init__(self, configuration) + SimpleSimulation.__init__(self, configuration) # set up plasma storage PlasmaStateStorerMixin.__init__( From a68ebaa0c6deca6b4e8908f75f2645c85da50bea Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Thu, 15 Aug 2024 15:05:14 -0400 Subject: [PATCH 26/28] Black format --- tardis/simulation/base.py | 8 ++++++-- tardis/workflows/simple_simulation.py | 14 +++++++++----- tardis/workflows/standard_simulation.py | 8 ++++++-- 3 files changed, 21 insertions(+), 9 deletions(-) diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index f7e9a325c22..92a658b80d7 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -291,8 +291,12 @@ 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, diff --git a/tardis/workflows/simple_simulation.py b/tardis/workflows/simple_simulation.py index 608fb0fbd38..1d4bff8a657 100644 --- a/tardis/workflows/simple_simulation.py +++ b/tardis/workflows/simple_simulation.py @@ -185,8 +185,12 @@ def get_convergence_estimates(self, transport_state): ) ) - estimated_t_radiative = estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature - estimated_dilution_factor = estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor + estimated_t_radiative = ( + estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature + ) + estimated_dilution_factor = ( + estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor + ) emitted_luminosity = calculate_filtered_luminosity( transport_state.emitted_packet_nu, @@ -402,9 +406,9 @@ def initialize_spectrum_solver( self.spectrum_solver.transport_state = transport_state if virtual_packet_energies is not None: - self.spectrum_solver._montecarlo_virtual_luminosity.value[:] = ( - virtual_packet_energies - ) + self.spectrum_solver._montecarlo_virtual_luminosity.value[ + : + ] = virtual_packet_energies if self.integrated_spectrum_settings is not None: # Set up spectrum solver integrator diff --git a/tardis/workflows/standard_simulation.py b/tardis/workflows/standard_simulation.py index 429752b543c..ec6c53b1de2 100644 --- a/tardis/workflows/standard_simulation.py +++ b/tardis/workflows/standard_simulation.py @@ -104,8 +104,12 @@ def get_convergence_estimates(self, transport_state): ) ) - estimated_t_radiative = estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature - estimated_dilution_factor = estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor + estimated_t_radiative = ( + estimated_radfield_properties.dilute_blackbody_radiationfield_state.temperature + ) + estimated_dilution_factor = ( + estimated_radfield_properties.dilute_blackbody_radiationfield_state.dilution_factor + ) emitted_luminosity = calculate_filtered_luminosity( transport_state.emitted_packet_nu, From e36ccd51fe15262e5e596f7d17184c8a6e64a212 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Thu, 15 Aug 2024 15:33:37 -0400 Subject: [PATCH 27/28] Fixes spectrum solver test --- tardis/spectrum/tests/test_spectrum_solver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/spectrum/tests/test_spectrum_solver.py b/tardis/spectrum/tests/test_spectrum_solver.py index 9f4f5752ff4..4c84ade6558 100644 --- a/tardis/spectrum/tests/test_spectrum_solver.py +++ b/tardis/spectrum/tests/test_spectrum_solver.py @@ -92,7 +92,7 @@ def test_solve(self, simulation): transport_state = simulation.transport.transport_state spectrum_frequency_grid = simulation.transport.spectrum_frequency_grid - solver = SpectrumSolver(transport_state, spectrum_frequency_grid) + solver = SpectrumSolver(transport_state, spectrum_frequency_grid, None) result_real, result_virtual, result_integrated = solver.solve( transport_state ) From 79df0e31ad9ecb201ff7e6ff85c0403641a9f631 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Fri, 16 Aug 2024 14:41:41 -0400 Subject: [PATCH 28/28] Some suggested refactoring --- tardis/workflows/standard_simulation.py | 74 +++++++++++++++++-------- tardis/workflows/workflow_logging.py | 66 +++++++++++++--------- 2 files changed, 91 insertions(+), 49 deletions(-) diff --git a/tardis/workflows/standard_simulation.py b/tardis/workflows/standard_simulation.py index ec6c53b1de2..1f3ecae2327 100644 --- a/tardis/workflows/standard_simulation.py +++ b/tardis/workflows/standard_simulation.py @@ -16,6 +16,9 @@ class StandardSimulation( SimpleSimulation, PlasmaStateStorerMixin, HDFWriterMixin ): + convergence_plots = None + export_convergence_plots = False + hdf_properties = [ "simulation_state", "plasma_solver", @@ -41,6 +44,7 @@ def __init__( self.log_level = log_level self.specific_log_level = specific_log_level self.enable_virtual_packet_logging = enable_virtual_packet_logging + self.convergence_plots_kwargs = convergence_plots_kwargs SimpleSimulation.__init__(self, configuration) @@ -53,31 +57,53 @@ def __init__( # Convergence plots if show_convergence_plots: - if not is_notebook(): - raise RuntimeError( - "Convergence Plots cannot be displayed in command-line. Set show_convergence_plots " - "to False." - ) + ( + self.convergence_plots, + self.export_convergence_plots, + ) = self.initialize_convergence_plots() + + def initialize_convergence_plots(self): + """Initialize the convergence plot attributes - self.convergence_plots = ConvergencePlots( - iterations=self.total_iterations, **convergence_plots_kwargs + Returns + ------- + ConvergencePlots + The convergence plot instance + bool + If convergence plots are to be exported + + Raises + ------ + RuntimeError + Raised if run outside a notebook + TypeError + Raised if export_convergence_plots is not a bool + """ + if not is_notebook(): + raise RuntimeError( + "Convergence Plots cannot be displayed in command-line. Set show_convergence_plots " + "to False." ) - else: - self.convergence_plots = None - if "export_convergence_plots" in convergence_plots_kwargs: + convergence_plots = ConvergencePlots( + iterations=self.total_iterations, **self.convergence_plots_kwargs + ) + + if "export_convergence_plots" in self.convergence_plots_kwargs: if not isinstance( - convergence_plots_kwargs["export_convergence_plots"], + self.convergence_plots_kwargs["export_convergence_plots"], bool, ): raise TypeError( "Expected bool in export_convergence_plots argument" ) - self.export_convergence_plots = convergence_plots_kwargs[ + export_convergence_plots = self.convergence_plots_kwargs[ "export_convergence_plots" ] else: - self.export_convergence_plots = False + export_convergence_plots = False + + return convergence_plots, export_convergence_plots def get_convergence_estimates(self, transport_state): """Compute convergence estimates from the transport state @@ -124,6 +150,16 @@ def get_convergence_estimates(self, transport_state): self.luminosity_nu_end, ) + luminosity_ratios = ( + (emitted_luminosity / self.luminosity_requested).to(1).value + ) + + estimated_t_inner = ( + self.simulation_state.t_inner + * luminosity_ratios + ** self.convergence_strategy.t_inner_update_exponent + ) + if self.convergence_plots is not None: plot_data = { "t_inner": [self.simulation_state.t_inner.value, "value"], @@ -134,7 +170,7 @@ def get_convergence_estimates(self, transport_state): "Absorbed": [absorbed_luminosity.value, "value"], "Requested": [self.luminosity_requested.value, "value"], } - self.update_convergence_plot_data(plot_data) + self.update_convergence_plot_data(plot_data) logger.info( f"\n\tLuminosity emitted = {emitted_luminosity:.3e}\n" @@ -142,16 +178,6 @@ def get_convergence_estimates(self, transport_state): f"\tLuminosity requested = {self.luminosity_requested:.3e}\n" ) - luminosity_ratios = ( - (emitted_luminosity / self.luminosity_requested).to(1).value - ) - - estimated_t_inner = ( - self.simulation_state.t_inner - * luminosity_ratios - ** self.convergence_strategy.t_inner_update_exponent - ) - self.log_plasma_state( self.simulation_state.t_radiative, self.simulation_state.dilution_factor, diff --git a/tardis/workflows/workflow_logging.py b/tardis/workflows/workflow_logging.py index 2dc464fe3b9..ddb76920dc4 100644 --- a/tardis/workflows/workflow_logging.py +++ b/tardis/workflows/workflow_logging.py @@ -58,34 +58,50 @@ def log_plasma_state( plasma_state_log["next_w"] = next_dilution_factor plasma_state_log.columns.name = "Shell No." + logger.info("\n\tPlasma stratification:") + if is_notebook(): - logger.info("\n\tPlasma stratification:") - - # Displaying the DataFrame only when the logging level is NOTSET, DEBUG or INFO - if logger.level <= logging.INFO: - if not logger.filters: - display( - plasma_state_log.iloc[::log_sampling].style.format( - "{:.3g}" - ) - ) - elif logger.filters[0].log_level == 20: - display( - plasma_state_log.iloc[::log_sampling].style.format( - "{:.3g}" - ) - ) + self.log_dataframe_notebook(plasma_state_log, log_sampling) else: - output_df = "" - plasma_output = plasma_state_log.iloc[::log_sampling].to_string( - float_format=lambda x: f"{x:.3g}", - justify="center", - ) - for value in plasma_output.split("\n"): - output_df = output_df + f"\t{value}\n" - logger.info("\n\tPlasma stratification:") - logger.info(f"\n{output_df}") + self.log_dataframe_console(plasma_state_log, log_sampling) logger.info( f"\n\tCurrent t_inner = {t_inner:.3f}\n\tExpected t_inner for next iteration = {next_t_inner:.3f}\n" ) + + def log_dataframe_notebook(self, dataframe, step): + """Logs a dataframe in a notebook with a step sample + + Parameters + ---------- + dataframe : pd.DataFrame + Dataframe to display + step : int + Step to use when sampling the dataframe + """ + # Displaying the DataFrame only when the logging level is NOTSET, DEBUG or INFO + if logger.level <= logging.INFO: + if not logger.filters: + display(dataframe.iloc[::step].style.format("{:.3g}")) + elif logger.filters[0].log_level == 20: + display(dataframe.iloc[::step].style.format("{:.3g}")) + + def log_dataframe_console(self, dataframe, step): + """Logs a dataframe to console with a step sample + + Parameters + ---------- + dataframe : pd.DataFrame + Dataframe to display + step : int + Step to use when sampling the dataframe + """ + output_df = "" + output = dataframe.iloc[::step].to_string( + float_format=lambda x: f"{x:.3g}", + justify="center", + ) + for value in output.split("\n"): + output_df = output_df + f"\t{value}\n" + + logger.info(f"\n{output_df}")