From df439d895f6ee3a25c5df3cb8d63d6439a756943 Mon Sep 17 00:00:00 2001 From: Wolfgang Kerzendorf Date: Mon, 11 Dec 2023 11:06:01 -0500 Subject: [PATCH] Restructure radiation field (#2474) * restructure of geometry * add radial1d boundary logic * black format * several fixes * fix epsilon * add testing of boundaries * change the r_inner_active * first integration with `from_config` working * hunting down density indexing bug * all model tests (without csvy) pass * more fixes * fix of model to simulation_state * fix inner boundary packet error * fix some leftovers * final fix for csvy * blackify * restructure to readers and remove some leftover code * further cleanup * first start of the restructure * add comment about removing quantitiness * add velocity check * add new abundance functions * remove default units * add new matter module * several restructures. move decay from io to tardis/model/matter * mid restructure * slow progress on abundance refactor * include effective_element_masses * further updates * fix of csvy readers * last fixes to custom abundance widget * restructuring the radiation field * some cleanup * cleanup * further cleanup * clearnup * changes * remove matter base.py * removing matter * fixed last tests * add several parsers * add composition * remove isotopemassfractions * some small fixes * fixed w, t_rad to dilution_factor, t_radiative * add setters for w, t_rad * several fixes for instantiating the right t_rads. also ensuring that they are positive and len(t_rad) == dilution factor * add the from_config atom_data * fixing model.ipynb * fix some more tests * fix tests * fix some of the issues wdigets * assert mass fractions positive. Change IsotopeAbundances to Massfractions * add radiation_field_state * Fix isotope mass fraction typo and t_radiatve duplicate definition * Document remove functionality for effective element masses * Fix variable name in convert_to_nuclide_mass_fraction function * Refactor test names and variable names in test_base.py * Refactor gamma ray simulation setup and test functions * fix gamma ray tests * Add model_isotope_time_0 property to abundances schema * Add test for montecarlo main loop with vpacket tracking * Refactor config_reader and config_validator modules * Fix model_isotope_time_0 property in model_definitions.yml * Refactor simulation module to use pathlib * remove bad comment and refactor with ruff * Fix variable name in test_simulation_state_mass() function * Fix unit conversion in calculate_cell_masses function * Fix logger debug message in CSV reader and update method names in Composition class * black formatting * Refactor code to use composition variable in test_gamma_ray_transport.py * Refactor imports in atom_web_download.py and update return statement in config_internal.py * fix for documentation not builiding * add fix for model_isotope_time_0 * fix grid test * restructure the grid * final change * add stuff * add vpacket_log * Refactor model_reader.py and shell_info.py changing w to dilution_factor * Update damping constants and variable names in doc notebook * fix w in several places in tardis * slowly fixing radiation_field changes * last fix * change hdf from model to simulation_state * fixing model->simulation-state * Fix model_reader.py imports and formatting issues * Refactor variable names for clarity and consistency and being correct * Update variable names in model.ipynb * Update variable name for average temperature calculation * Refactor variable names in model reader and simulation base * Refactor store_model_to_hdf to store_simulation_state_to_hdf * Refactor HDF file handling and model reader imports * add docstr --- docs/io/optional/callback_example.ipynb | 2 +- docs/physics/setup/model.ipynb | 38 ++-- .../update_and_conv/update_and_conv.ipynb | 44 ++-- tardis/io/model/hdf.py | 27 +-- tardis/io/model/model_reader.py | 110 ++-------- tardis/io/model_reader.py | 103 +-------- tardis/io/tests/test_model_reader.py | 144 +++++++----- tardis/model/base.py | 204 +++++------------ tardis/model/geometry/radial1d.py | 18 ++ tardis/model/matter/composition.py | 12 + tardis/model/parse_input.py | 205 +++++++++++++++++- tardis/model/radiation_field_state.py | 32 +++ tardis/model/tests/test_base.py | 6 +- tardis/model/tests/test_density.py | 4 +- .../montecarlo_numba/formal_integral.py | 21 +- .../montecarlo_numba/tests/test_base.py | 97 +++++++-- tardis/simulation/base.py | 49 +++-- tardis/visualization/widgets/shell_info.py | 15 +- .../widgets/tests/test_shell_info.py | 4 +- 19 files changed, 633 insertions(+), 502 deletions(-) create mode 100644 tardis/model/radiation_field_state.py diff --git a/docs/io/optional/callback_example.ipynb b/docs/io/optional/callback_example.ipynb index 488aa27392f..0b13d5bb8cd 100644 --- a/docs/io/optional/callback_example.ipynb +++ b/docs/io/optional/callback_example.ipynb @@ -25,7 +25,7 @@ "outputs": [], "source": [ "def average_temp(sim):\n", - " t_rads = sim.simulation_state.t_rad\n", + " t_rads = sim.simulation_state.t_radiative\n", " volumes = sim.simulation_state.volume\n", " avg = sum(t_rads*volumes) / sum(volumes)\n", " print(f\"Average temperature for iteration {sim.iterations_executed}: {avg}\")" diff --git a/docs/physics/setup/model.ipynb b/docs/physics/setup/model.ipynb index 26e36fca7b7..4e1df5b2780 100644 --- a/docs/physics/setup/model.ipynb +++ b/docs/physics/setup/model.ipynb @@ -461,9 +461,9 @@ " 'Ni57':.1,\n", " 'Cr51':.4}\n", "\n", - "abund_isotopes_model = SimulationState.from_config(abund_isotopes_config, atom_data=atom_data)\n", + "abund_isotopes_simulation_state = SimulationState.from_config(abund_isotopes_config, atom_data=atom_data)\n", "\n", - "abund_isotopes_model.abundance" + "abund_isotopes_simulation_state.abundance" ] }, { @@ -529,15 +529,15 @@ "t_rad_config.model.structure.velocity.stop = 2000 * u.km/u.s\n", "t_rad_config.model.structure.velocity.num = 20\n", "\n", - "t_rad_model = SimulationState.from_config(t_rad_config, atom_data=atom_data)\n", + "t_radiative_simulation_state = SimulationState.from_config(t_rad_config, atom_data=atom_data)\n", "\n", - "print('t_inner:\\n', t_rad_model.t_inner)\n", - "print('t_rad:\\n', t_rad_model.t_rad)\n", + "print('t_inner:\\n', t_radiative_simulation_state.t_inner)\n", + "print('t_rad:\\n', t_radiative_simulation_state.t_radiative)\n", "\n", - "plt.plot(t_rad_model.r_middle, t_rad_model.t_rad)\n", - "plt.scatter(t_rad_model.radius[0], t_rad_model.t_inner)\n", - "plt.xlabel(f'Radius ({t_rad_model.r_middle.unit})')\n", - "plt.ylabel(f'Radiative Temperature ({t_rad_model.t_rad.unit})');" + "plt.plot(t_radiative_simulation_state.r_middle, t_radiative_simulation_state.t_radiative)\n", + "plt.scatter(t_radiative_simulation_state.radius[0], t_radiative_simulation_state.t_inner)\n", + "plt.xlabel(f'Radius ({t_radiative_simulation_state.r_middle.unit})')\n", + "plt.ylabel(f'Radiative Temperature ({t_radiative_simulation_state.t_radiative.unit})');" ] }, { @@ -580,24 +580,24 @@ }, "outputs": [], "source": [ - "w_config = copy.deepcopy(base_config)\n", + "dilution_factor_config = copy.deepcopy(base_config)\n", "\n", - "w_config.supernova.time_explosion = 10 * u.day\n", + "dilution_factor_config.supernova.time_explosion = 10 * u.day\n", "\n", "# This line is necessary to indicate we are using a built-in shell structure\n", "# and density. Do not change it.\n", - "w_config.model.structure.type = 'specific'\n", + "dilution_factor_config.model.structure.type = 'specific'\n", "\n", - "w_config.model.structure.velocity.start = 1000 * u.km/u.s\n", - "w_config.model.structure.velocity.stop = 2000 * u.km/u.s\n", - "w_config.model.structure.velocity.num = 20\n", + "dilution_factor_config.model.structure.velocity.start = 1000 * u.km/u.s\n", + "dilution_factor_config.model.structure.velocity.stop = 2000 * u.km/u.s\n", + "dilution_factor_config.model.structure.velocity.num = 20\n", "\n", - "w_model = SimulationState.from_config(w_config, atom_data=atom_data)\n", + "dilution_factor_simulation_state = SimulationState.from_config(dilution_factor_config, atom_data=atom_data)\n", "\n", - "print('w:\\n', w_model.w)\n", + "print(f'dilution_factor: {dilution_factor_simulation_state.dilution_factor}')\n", "\n", - "plt.plot(w_model.r_middle, w_model.w)\n", - "plt.xlabel(f'Radius ({w_model.r_middle.unit})')\n", + "plt.plot(dilution_factor_simulation_state.r_middle, dilution_factor_simulation_state.dilution_factor)\n", + "plt.xlabel(f'Radius ({dilution_factor_simulation_state.r_middle.unit})')\n", "plt.ylabel(f'Dilution Factor');" ] } diff --git a/docs/physics/update_and_conv/update_and_conv.ipynb b/docs/physics/update_and_conv/update_and_conv.ipynb index b92b8f134b0..971249e6a75 100644 --- a/docs/physics/update_and_conv/update_and_conv.ipynb +++ b/docs/physics/update_and_conv/update_and_conv.ipynb @@ -191,15 +191,15 @@ "\n", "# We manually put in the damping constants and t_inner_update_exponent for\n", "# illustrative purposes:\n", - "d_t_rad = 0.5\n", - "d_w = 0.3\n", - "d_t_inner = 0.7\n", + "damping_t_radiative = 0.5\n", + "damping_dilution_factor = 0.3\n", + "damping_t_inner = 0.7\n", "t_inner_update_exponent = -0.5\n", "\n", "# We set the above values into the configuration:\n", - "tardis_config.montecarlo.convergence_strategy.t_rad.damping_constant = d_t_rad\n", - "tardis_config.montecarlo.convergence_strategy.w.damping_constant = d_w\n", - "tardis_config.montecarlo.convergence_strategy.t_inner.damping_constant = d_t_inner\n", + "tardis_config.montecarlo.convergence_strategy.t_rad.damping_constant = damping_t_radiative\n", + "tardis_config.montecarlo.convergence_strategy.w.damping_constant = damping_dilution_factor\n", + "tardis_config.montecarlo.convergence_strategy.t_inner.damping_constant = damping_t_inner\n", "tardis_config.montecarlo.convergence_strategy.t_inner_update_exponent = t_inner_update_exponent" ] }, @@ -212,7 +212,7 @@ "source": [ "sim = Simulation.from_config(tardis_config)\n", "\n", - "model = sim.simulation_state\n", + "simulation_state = sim.simulation_state\n", "plasma = sim.plasma\n", "transport = sim.transport" ] @@ -233,7 +233,7 @@ "metadata": {}, "outputs": [], "source": [ - "model.t_rad" + "simulation_state.t_radiative" ] }, { @@ -243,7 +243,7 @@ "metadata": {}, "outputs": [], "source": [ - "model.w" + "simulation_state.dilution_factor" ] }, { @@ -253,7 +253,7 @@ "metadata": {}, "outputs": [], "source": [ - "model.t_inner" + "simulation_state.t_inner" ] }, { @@ -344,8 +344,8 @@ "metadata": {}, "outputs": [], "source": [ - "V = model.volume\n", - "Delta_t = transport.calculate_time_of_simulation(model)\n", + "V = simulation_state.volume\n", + "Delta_t = transport.calculate_time_of_simulation(simulation_state)\n", "prefactor = 1 / (4 * np.pi * V * Delta_t)\n", "J = prefactor * j_estimator\n", "J" @@ -410,7 +410,7 @@ "metadata": {}, "outputs": [], "source": [ - "t_rad_updated = model.t_rad + d_t_rad * (t_rad_estimated - model.t_rad)\n", + "t_rad_updated = simulation_state.t_radiative + damping_t_radiative * (t_rad_estimated - simulation_state.t_radiative)\n", "t_rad_updated" ] }, @@ -421,8 +421,8 @@ "metadata": {}, "outputs": [], "source": [ - "w_estimated = ( j_estimator / (4 * const.sigma_sb.cgs * t_rad_estimated**4 * V * Delta_t) ).decompose()\n", - "w_estimated" + "dilution_factor_estimated = ( j_estimator / (4 * const.sigma_sb.cgs * t_rad_estimated**4 * V * Delta_t) ).decompose()\n", + "dilution_factor_estimated" ] }, { @@ -432,8 +432,8 @@ "metadata": {}, "outputs": [], "source": [ - "w_updated = model.w + d_w * (w_estimated - model.w)\n", - "w_updated" + "dilution_factor_updated = simulation_state.dilution_factor + damping_dilution_factor * (dilution_factor_estimated - simulation_state.dilution_factor)\n", + "dilution_factor_updated" ] }, { @@ -484,7 +484,7 @@ "metadata": {}, "outputs": [], "source": [ - "t_inner_estimated = model.t_inner * (L_output / L_requested)**t_inner_update_exponent\n", + "t_inner_estimated = simulation_state.t_inner * (L_output / L_requested)**t_inner_update_exponent\n", "t_inner_estimated" ] }, @@ -495,7 +495,7 @@ "metadata": {}, "outputs": [], "source": [ - "t_inner_updated = model.t_inner + d_t_inner * (t_inner_estimated - model.t_inner)\n", + "t_inner_updated = simulation_state.t_inner + damping_t_inner * (t_inner_estimated - simulation_state.t_inner)\n", "t_inner_updated" ] }, @@ -536,7 +536,7 @@ "metadata": {}, "outputs": [], "source": [ - "model.t_rad" + "simulation_state.t_radiative" ] }, { @@ -548,7 +548,7 @@ }, "outputs": [], "source": [ - "model.w" + "simulation_state.dilution_factor" ] }, { @@ -558,7 +558,7 @@ "metadata": {}, "outputs": [], "source": [ - "model.t_inner" + "simulation_state.t_inner" ] }, { diff --git a/tardis/io/model/hdf.py b/tardis/io/model/hdf.py index 5d066f1e793..1621b1134ce 100644 --- a/tardis/io/model/hdf.py +++ b/tardis/io/model/hdf.py @@ -1,27 +1,28 @@ -from tardis.io.model.model_reader import simulation_state_to_dict - - import h5py +from tardis.io.model.model_reader import simulation_state_to_dict + -def store_model_to_hdf(model, fname): +def store_simulation_state_to_hdf(simulation_state, fname): """ Stores data from SimulationState object into a hdf file. Parameters ---------- - model : tardis.model.SimulationState - filename : str + simulation_state : tardis.model.SimulationState + fname : str """ with h5py.File(fname, "a") as f: - model_group = f.require_group("model") - model_group.clear() + simulation_state_group = f.require_group("simulation_state") + simulation_state_group.clear() - model_dict = simulation_state_to_dict(model) + simulation_state_dict = simulation_state_to_dict(simulation_state) - for key, value in model_dict.items(): + for key, value in simulation_state_dict.items(): if key.endswith("_cgs"): - model_group.create_dataset(key, data=value[0]) - model_group.create_dataset(key + "_unit", data=value[1]) + simulation_state_group.create_dataset(key, data=value[0]) + simulation_state_group.create_dataset( + key + "_unit", data=value[1] + ) else: - model_group.create_dataset(key, data=value) + simulation_state_group.create_dataset(key, data=value) diff --git a/tardis/io/model/model_reader.py b/tardis/io/model/model_reader.py index 91c78ee5be0..49be1cf5c1f 100644 --- a/tardis/io/model/model_reader.py +++ b/tardis/io/model/model_reader.py @@ -1,4 +1,7 @@ -# reading different model files +import logging + +import h5py +from astropy import units as u from tardis.io.configuration.config_reader import ConfigurationNameSpace from tardis.montecarlo.base import MontecarloTransport @@ -6,11 +9,6 @@ BlackBodySimpleSource, BlackBodySimpleSourceRelativistic, ) -from tardis.io.model.density import calculate_density_after_time -from astropy import units as u -import h5py - -import logging # Adding logging support logger = logging.getLogger(__name__) @@ -31,7 +29,6 @@ def transport_to_dict(transport): v_packet_settings : dict virtual_spectrum_spawn_range : dict """ - transport_dict = { "Edotlu_estimator": transport.Edotlu_estimator, "bf_heating_estimator": transport.bf_heating_estimator, @@ -104,7 +101,6 @@ def store_transport_to_hdf(transport, fname): transport : tardis.montecarlo.MontecarloTransport filename : str """ - with h5py.File(fname, "a") as f: transport_group = f.require_group("transport") transport_group.clear() @@ -155,7 +151,6 @@ def transport_from_hdf(fname): ------- new_transport : tardis.montecarlo.MontecarloTransport """ - d = {} # Loading data from hdf file @@ -282,7 +277,7 @@ def transport_from_hdf(fname): return new_transport -def simulation_state_to_dict(model): +def simulation_state_to_dict(simulation_state): """ Retrieves all the data from a SimulationState object and returns a dictionary. @@ -296,18 +291,17 @@ def simulation_state_to_dict(model): isotope_abundance : dict """ simulation_state_dict = { - "velocity_cgs": model.velocity.cgs, - "abundance": model.abundance, - "time_explosion_cgs": model.time_explosion.cgs, - "t_inner_cgs": model.t_inner.cgs, - "t_radiative_cgs": model.t_radiative.cgs, - "dilution_factor": model.dilution_factor, - "v_boundary_inner_cgs": model.v_boundary_inner.cgs, - "v_boundary_outer_cgs": model.v_boundary_outer.cgs, - "w": model.w, - "t_rad_cgs": model.t_rad.cgs, - "r_inner_cgs": model.r_inner.cgs, - "density_cgs": model.density.cgs, + "velocity_cgs": simulation_state.velocity.cgs, + "abundance": simulation_state.abundance, + "time_explosion_cgs": simulation_state.time_explosion.cgs, + "t_inner_cgs": simulation_state.t_inner.cgs, + "t_radiative_cgs": simulation_state.t_radiative.cgs, + "dilution_factor": simulation_state.dilution_factor, + "v_boundary_inner_cgs": simulation_state.v_boundary_inner.cgs, + "v_boundary_outer_cgs": simulation_state.v_boundary_outer.cgs, + "dilution_factor": simulation_state.dilution_factor, + "r_inner_cgs": simulation_state.r_inner.cgs, + "density_cgs": simulation_state.density.cgs, } for key, value in simulation_state_dict.items(): @@ -318,75 +312,3 @@ def simulation_state_to_dict(model): ] return simulation_state_dict - - -def model_from_hdf(fname): - """ - Creates a SimulationState object using data stored in a hdf file. - - Parameters - ---------- - fname : str - - Returns - ------- - new_model : tardis.model.SimulationState - """ - - from tardis.model import SimulationState - - d = {} - - # Loading data from hdf file - with h5py.File(fname, "r") as f: - model_group = f["model"] - for key, value in model_group.items(): - if not key.endswith("_unit"): - if type(value) == h5py._hl.dataset.Dataset: - if isinstance(value[()], bytes): - d[key] = value[()].decode("utf-8") - else: - d[key] = value[()] - else: - data_inner = {} - for key_inner, value_inner in value.items(): - if isinstance(value_inner[()], bytes): - data_inner[key] = value_inner[()].decode("utf-8") - else: - data_inner[key] = value_inner[()] - data_inner[key_inner] = value_inner[()] - d[key] = data_inner - - for key, value in model_group.items(): - if key.endswith("_unit"): - d[key[:-5]] = [d[key[:-5]], value[()]] - - # Converting cgs data to astropy quantities - for key, value in d.items(): - if key.endswith("_cgs"): - d[key] = u.Quantity(value[0], unit=u.Unit(value[1].decode("utf-8"))) - - homologous_density = calculate_density_after_time( - d["homologous_density"]["density_0"], - d["homologous_density"]["time_0"], - d["time_explosion_cgs"], - ) - - new_model = SimulationState( - velocity=d["velocity_cgs"], - density=homologous_density, - abundance=d["abundance"], - isotope_abundance=None, - time_explosion=d["time_explosion_cgs"], - elemental_mass=None, - t_inner=d["t_inner_cgs"], - t_radiative=d["t_radiative_cgs"], - dilution_factor=d["dilution_factor"], - v_boundary_inner=d["v_boundary_inner_cgs"], - v_boundary_outer=d["v_boundary_outer_cgs"], - ) - - new_model.t_rad = d["t_rad_cgs"] - new_model.w = d["w"] - - return new_model diff --git a/tardis/io/model_reader.py b/tardis/io/model_reader.py index 2ca6236fa3b..6856481b71e 100644 --- a/tardis/io/model_reader.py +++ b/tardis/io/model_reader.py @@ -1,23 +1,23 @@ # reading different model files -from tardis.io.configuration.config_reader import ConfigurationNameSpace -from tardis.montecarlo.base import MontecarloTransport -from tardis.montecarlo.packet_source import ( - BlackBodySimpleSource, - BlackBodySimpleSourceRelativistic, -) -from tardis.util.base import parse_quantity, is_valid_nuclide_or_elem - +import logging import warnings + +import h5py import numpy as np -from numpy import recfromtxt, genfromtxt import pandas as pd from astropy import units as u +from numpy import recfromtxt from radioactivedecay import Nuclide from radioactivedecay.utils import Z_DICT, elem_to_Z -import h5py -import logging +from tardis.io.configuration.config_reader import ConfigurationNameSpace +from tardis.montecarlo.base import MontecarloTransport +from tardis.montecarlo.packet_source import ( + BlackBodySimpleSource, + BlackBodySimpleSourceRelativistic, +) +from tardis.util.base import is_valid_nuclide_or_elem, parse_quantity # Adding logging support logger = logging.getLogger(__name__) @@ -75,7 +75,7 @@ def read_density_file(filename, filetype): if invalid_volume_mask.sum() > 0: message = "\n".join( [ - f"cell {i:d}: v_inner {v_inner_i:s}, v_outer " f"{v_outer_i:s}" + f"cell {i:d}: v_inner {v_inner_i:s}, v_outer {v_outer_i:s}" for i, v_inner_i, v_outer_i in zip( np.arange(len(v_outer))[invalid_volume_mask], v_inner[invalid_volume_mask], @@ -116,7 +116,6 @@ def read_abundances_file( outer_boundary_index : int index of the outer shell, default None """ - file_parsers = { "simple_ascii": read_simple_ascii_abundances, "artis": read_simple_ascii_abundances, @@ -215,7 +214,6 @@ def read_simple_ascii_density(fname): data : pandas.DataFrame data frame containing index, velocity (in km/s) and density """ - with open(fname) as fh: time_of_model_string = fh.readline().strip() time_of_model = parse_quantity(time_of_model_string) @@ -253,7 +251,6 @@ def read_artis_density(fname): data : pandas.DataFrame data frame containing index, velocity (in km/s) and density """ - with open(fname) as fh: for i, line in enumerate(open(fname)): if i == 0: @@ -389,7 +386,6 @@ def read_cmfgen_composition(fname, delimiter=r"\s+"): fname : str filename of the csv file """ - warnings.warn( "The current CMFGEN model parser is deprecated", DeprecationWarning ) @@ -414,7 +410,6 @@ def read_csv_composition(fname, delimiter=r"\s+"): fname : str filename of the csv file """ - return read_csv_isotope_abundances( fname, delimiter=delimiter, skip_columns=0, skip_rows=[1] ) @@ -458,7 +453,6 @@ def read_csv_isotope_abundances( abundances : pandas.DataFrame isotope_abundance : pandas.MultiIndex """ - df = pd.read_csv( fname, comment="#", sep=delimiter, skiprows=skip_rows, index_col=0 ) @@ -506,7 +500,6 @@ def parse_csv_abundances(csvy_data): abundances : pandas.DataFrame isotope_abundance : pandas.MultiIndex """ - abundance_col_names = [ name for name in csvy_data.columns if is_valid_nuclide_or_elem(name) ] @@ -557,7 +550,6 @@ def transport_to_dict(transport): v_packet_settings : dict virtual_spectrum_spawn_range : dict """ - transport_dict = { "Edotlu_estimator": transport.Edotlu_estimator, "bf_heating_estimator": transport.bf_heating_estimator, @@ -630,7 +622,6 @@ def store_transport_to_hdf(transport, fname): transport : tardis.montecarlo.MontecarloTransport filename : str """ - with h5py.File(fname, "a") as f: transport_group = f.require_group("transport") transport_group.clear() @@ -681,7 +672,6 @@ def transport_from_hdf(fname): ------- new_transport : tardis.montecarlo.MontecarloTransport """ - d = {} # Loading data from hdf file @@ -866,72 +856,3 @@ def store_model_to_hdf(model, fname): model_group.create_dataset(key + "_unit", data=value[1]) else: model_group.create_dataset(key, data=value) - - -def model_from_hdf(fname): - """ - Creates a SimulationState object using data stored in a hdf file. - - Parameters - ---------- - fname : str - - Returns - ------- - new_model : tardis.model.SimulationState - """ - - from tardis.model import SimulationState - - d = {} - - # Loading data from hdf file - with h5py.File(fname, "r") as f: - model_group = f["model"] - for key, value in model_group.items(): - if not key.endswith("_unit"): - if type(value) == h5py._hl.dataset.Dataset: - if isinstance(value[()], bytes): - d[key] = value[()].decode("utf-8") - else: - d[key] = value[()] - else: - data_inner = {} - for key_inner, value_inner in value.items(): - if isinstance(value_inner[()], bytes): - data_inner[key] = value_inner[()].decode("utf-8") - else: - data_inner[key] = value_inner[()] - data_inner[key_inner] = value_inner[()] - d[key] = data_inner - - for key, value in model_group.items(): - if key.endswith("_unit"): - d[key[:-5]] = [d[key[:-5]], value[()]] - - # Converting cgs data to astropy quantities - for key, value in d.items(): - if key.endswith("_cgs"): - d[key] = u.Quantity(value[0], unit=u.Unit(value[1].decode("utf-8"))) - - homologous_density = HomologousDensity( - d["homologous_density"]["density_0"], d["homologous_density"]["time_0"] - ) - - new_model = SimulationState( - velocity=d["velocity_cgs"], - homologous_density=homologous_density, - abundance=d["abundance"], - isotope_abundance=None, - time_explosion=d["time_explosion_cgs"], - t_inner=d["t_inner_cgs"], - t_radiative=d["t_radiative_cgs"], - dilution_factor=d["dilution_factor"], - v_boundary_inner=d["v_boundary_inner_cgs"], - v_boundary_outer=d["v_boundary_outer_cgs"], - ) - - new_model.t_rad = d["t_rad_cgs"] - new_model.w = d["w"] - - return new_model diff --git a/tardis/io/tests/test_model_reader.py b/tardis/io/tests/test_model_reader.py index f690115dffc..7cbd41f4c99 100644 --- a/tardis/io/tests/test_model_reader.py +++ b/tardis/io/tests/test_model_reader.py @@ -1,17 +1,16 @@ -import os -from pathlib import Path - -from astropy import units as u +import h5py import numpy as np import pytest -import h5py +from astropy import units as u from tardis.io.configuration.config_reader import Configuration +from tardis.io.model.hdf import store_simulation_state_to_hdf from tardis.io.model.model_reader import ( simulation_state_to_dict, - transport_to_dict, store_transport_to_hdf, + transport_to_dict, ) +from tardis.io.model.readers.artis import read_artis_density from tardis.io.model.readers.cmfgen import ( read_cmfgen_composition, read_cmfgen_density, @@ -21,8 +20,6 @@ read_simple_ascii_abundances, read_uniform_abundances, ) -from tardis.io.model.hdf import store_model_to_hdf -from tardis.io.model.readers.artis import read_artis_density @pytest.fixture @@ -133,85 +130,134 @@ def test_simple_read_cmfgen_density(cmfgen_fname): def test_model_to_dict(simulation_verysimple): - model = simulation_verysimple.simulation_state + simulation_state = simulation_verysimple.simulation_state - model_dict = simulation_state_to_dict(model) + simulation_state_dict = simulation_state_to_dict(simulation_state) # Check model dictionary assert np.array_equal( - model_dict["velocity_cgs"][0], model.velocity.cgs.value + simulation_state_dict["velocity_cgs"][0], + simulation_state.velocity.cgs.value, + ) + assert ( + simulation_state_dict["velocity_cgs"][1] + == simulation_state.velocity.cgs.unit.to_string() + ) + assert np.array_equal( + simulation_state_dict["abundance"], simulation_state.abundance + ) + assert np.array_equal( + simulation_state_dict["time_explosion_cgs"][0], + simulation_state.time_explosion.value, + ) + assert ( + simulation_state_dict["time_explosion_cgs"][1] + == simulation_state.time_explosion.unit.to_string() + ) + assert np.array_equal( + simulation_state_dict["t_inner_cgs"][0], + simulation_state.t_inner.cgs.value, + ) + assert ( + simulation_state_dict["t_inner_cgs"][1] + == simulation_state.t_inner.unit.to_string() + ) + assert np.array_equal( + simulation_state_dict["t_radiative_cgs"][0], + simulation_state.t_radiative.cgs.value, + ) + assert ( + simulation_state_dict["t_radiative_cgs"][1] + == simulation_state.t_radiative.unit.to_string() + ) + assert np.array_equal( + simulation_state_dict["dilution_factor"], + simulation_state.dilution_factor, ) - assert model_dict["velocity_cgs"][1] == model.velocity.cgs.unit.to_string() - assert np.array_equal(model_dict["abundance"], model.abundance) assert np.array_equal( - model_dict["time_explosion_cgs"][0], model.time_explosion.value + simulation_state_dict["v_boundary_inner_cgs"][0], + simulation_state.v_boundary_inner.cgs.value, ) assert ( - model_dict["time_explosion_cgs"][1] - == model.time_explosion.unit.to_string() + simulation_state_dict["v_boundary_inner_cgs"][1] + == simulation_state.v_boundary_inner.cgs.unit.to_string() ) - assert np.array_equal(model_dict["t_inner_cgs"][0], model.t_inner.cgs.value) - assert model_dict["t_inner_cgs"][1] == model.t_inner.unit.to_string() assert np.array_equal( - model_dict["t_radiative_cgs"][0], model.t_radiative.cgs.value + simulation_state_dict["v_boundary_outer_cgs"][0], + simulation_state.v_boundary_outer.cgs.value, ) assert ( - model_dict["t_radiative_cgs"][1] == model.t_radiative.unit.to_string() + simulation_state_dict["v_boundary_outer_cgs"][1] + == simulation_state.v_boundary_outer.cgs.unit.to_string() ) - assert np.array_equal(model_dict["dilution_factor"], model.dilution_factor) + assert np.array_equal( - model_dict["v_boundary_inner_cgs"][0], model.v_boundary_inner.cgs.value + simulation_state_dict["r_inner_cgs"][0], + simulation_state.r_inner.cgs.value, ) assert ( - model_dict["v_boundary_inner_cgs"][1] - == model.v_boundary_inner.cgs.unit.to_string() + simulation_state_dict["r_inner_cgs"][1] + == simulation_state.r_inner.cgs.unit.to_string() ) assert np.array_equal( - model_dict["v_boundary_outer_cgs"][0], model.v_boundary_outer.cgs.value + simulation_state_dict["density_cgs"][0], + simulation_state.density.cgs.value, ) assert ( - model_dict["v_boundary_outer_cgs"][1] - == model.v_boundary_outer.cgs.unit.to_string() + simulation_state_dict["density_cgs"][1] + == simulation_state.density.cgs.unit.to_string() ) - assert np.array_equal(model_dict["w"], model.w) - assert np.array_equal(model_dict["t_rad_cgs"][0], model.t_rad.cgs.value) - assert model_dict["t_rad_cgs"][1] == model.t_rad.cgs.unit.to_string() - assert np.array_equal(model_dict["r_inner_cgs"][0], model.r_inner.cgs.value) - assert model_dict["r_inner_cgs"][1] == model.r_inner.cgs.unit.to_string() - assert np.array_equal(model_dict["density_cgs"][0], model.density.cgs.value) - assert model_dict["density_cgs"][1] == model.density.cgs.unit.to_string() def test_store_model_to_hdf(simulation_verysimple, tmp_path): - model = simulation_verysimple.simulation_state + simulation_state = simulation_verysimple.simulation_state - fname = tmp_path / "model.h5" + fname = tmp_path / "simulation_state.h5" # Store model object - store_model_to_hdf(model, fname) + store_simulation_state_to_hdf(simulation_state, fname) # Check file contents with h5py.File(fname) as f: - assert np.array_equal(f["model/velocity_cgs"], model.velocity.cgs.value) - assert np.array_equal(f["model/abundance"], model.abundance) assert np.array_equal( - f["model/time_explosion_cgs"], model.time_explosion.cgs.value + f["simulation_state/velocity_cgs"], + simulation_state.velocity.cgs.value, + ) + assert np.array_equal( + f["simulation_state/abundance"], simulation_state.abundance + ) + assert np.array_equal( + f["simulation_state/time_explosion_cgs"], + simulation_state.time_explosion.cgs.value, + ) + assert np.array_equal( + f["simulation_state/t_inner_cgs"], + simulation_state.t_inner.cgs.value, + ) + assert np.array_equal( + f["simulation_state/t_radiative_cgs"], + simulation_state.t_radiative.cgs.value, + ) + assert np.array_equal( + f["simulation_state/dilution_factor"], + simulation_state.dilution_factor, + ) + assert np.array_equal( + f["simulation_state/v_boundary_inner_cgs"], + simulation_state.v_boundary_inner.cgs.value, ) - assert np.array_equal(f["model/t_inner_cgs"], model.t_inner.cgs.value) assert np.array_equal( - f["model/t_radiative_cgs"], model.t_radiative.cgs.value + f["simulation_state/v_boundary_outer_cgs"], + simulation_state.v_boundary_outer.cgs.value, ) - assert np.array_equal(f["model/dilution_factor"], model.dilution_factor) assert np.array_equal( - f["model/v_boundary_inner_cgs"], model.v_boundary_inner.cgs.value + f["simulation_state/r_inner_cgs"], + simulation_state.r_inner.cgs.value, ) assert np.array_equal( - f["model/v_boundary_outer_cgs"], model.v_boundary_outer.cgs.value + f["simulation_state/density_cgs"], + simulation_state.density.cgs.value, ) - assert np.array_equal(f["model/w"], model.w) - assert np.array_equal(f["model/t_rad_cgs"], model.t_rad.cgs.value) - assert np.array_equal(f["model/r_inner_cgs"], model.r_inner.cgs.value) - assert np.array_equal(f["model/density_cgs"], model.density.cgs.value) def test_transport_to_dict(simulation_verysimple): diff --git a/tardis/model/base.py b/tardis/model/base.py index ac2ef5ea0b0..1820670b7ba 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -15,12 +15,15 @@ from tardis.model.matter.composition import Composition from tardis.model.parse_input import ( parse_abundance_config, - parse_composition_csvy, + parse_csvy_composition, parse_csvy_geometry, + parse_csvy_radiation_field_state, + parse_radiation_field_state, parse_structure_config, + parse_packet_source, ) from tardis.montecarlo.packet_source import BlackBodySimpleSource -from tardis.radiation_field.base import MonteCarloRadiationFieldState +from tardis.model.radiation_field_state import DiluteThermalRadiationFieldState from tardis.util.base import is_valid_nuclide_or_elem logger = logging.getLogger(__name__) @@ -83,7 +86,7 @@ class SimulationState(HDFWriterMixin): hdf_properties = [ "t_inner", - "w", + "dilution_factor", "t_radiative", "v_inner", "v_outer", @@ -91,111 +94,46 @@ class SimulationState(HDFWriterMixin): "r_inner", "time_explosion", ] - hdf_name = "model" + hdf_name = "simulation_state" def __init__( self, geometry, composition, + radiation_field_state, time_explosion, - t_inner, - luminosity_requested=None, - t_radiative=None, - dilution_factor=None, + packet_source, electron_densities=None, ): self.geometry = geometry self.composition = composition - + self.radiation_field_state = radiation_field_state self.time_explosion = time_explosion self._electron_densities = electron_densities - - self.blackbody_packet_source = BlackBodySimpleSource( - self.r_inner[0], t_inner - ) - if t_inner is None: - if luminosity_requested is not None: - self.blackbody_packet_source.set_temperature_from_luminosity( - luminosity_requested - ) - else: - raise ValueError( - "Both t_inner and luminosity_requested cannot be None." - ) - else: - self.blackbody_packet_source.temperature = t_inner - - if t_radiative is None: - lambda_wien_inner = ( - constants.b_wien / self.blackbody_packet_source.temperature - ) - t_radiative = constants.b_wien / ( - lambda_wien_inner - * ( - 1 - + (self.v_middle - self.geometry.v_inner_boundary) - / constants.c - ) - ) - - elif len(t_radiative) == self.no_of_shells + 1: - t_radiative = t_radiative[ - self.geometry.v_inner_boundary_index - + 1 : self.geometry.v_outer_boundary_index - + 1 - ] - else: - assert len(t_radiative) == self.no_of_shells - - if dilution_factor is None: - dilution_factor = 0.5 * ( - 1 - - np.sqrt( - 1 - (self.r_inner[0] ** 2 / self.r_middle**2).to(1).value - ) - ) - elif len(dilution_factor) != self.no_of_shells: - dilution_factor = dilution_factor[ - self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index - ] - assert len(dilution_factor) == self.no_of_shells - - self.radiation_field = MonteCarloRadiationFieldState( - t_radiative, dilution_factor, None, None - ) - - @property - def w(self): - return self.dilution_factor - - @w.setter - def w(self, value): - self.dilution_factor = value + self.packet_source = packet_source @property - def t_rad(self): - return self.t_radiative - - @t_rad.setter - def t_rad(self, value): - self.t_radiative = value + def blackbody_packet_source(self): + return self.packet_source @property def t_inner(self): - return self.blackbody_packet_source.temperature + return self.packet_source.temperature @t_inner.setter def t_inner(self, value): - self.blackbody_packet_source.temperature = value + self.packet_source.temperature = value @property def dilution_factor(self): - return self.radiation_field.dilution_factor + return self.radiation_field_state.dilution_factor[ + self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index + ] @dilution_factor.setter - def dilution_factor(self, value): - if len(value) == self.no_of_shells: - self.radiation_field.dilution_factor = value + def dilution_factor(self, new_dilution_factor): + if len(new_dilution_factor) == self.no_of_shells: + self.radiation_field_state.dilution_factor = new_dilution_factor else: raise ValueError( "Trying to set dilution_factor for unmatching number" @@ -204,15 +142,17 @@ def dilution_factor(self, value): @property def t_radiative(self): - return self.radiation_field.t_radiative + return self.radiation_field_state.t_radiative[ + self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index + ] @t_radiative.setter - def t_radiative(self, value): - if len(value) == self.no_of_shells: - self.radiation_field.t_radiative = value + def t_radiative(self, new_t_radiative): + if len(new_t_radiative) == self.no_of_shells: + self.radiation_field_state.t_radiative = new_t_radiative else: raise ValueError( - "Trying to set t_radiative for unmatching number of shells." + "Trying to set t_radiative for different number of shells." ) @property @@ -306,7 +246,7 @@ def from_config(cls, config, atom_data): ( electron_densities, - temperature, + t_radiative, geometry, density, ) = parse_structure_config(config, time_explosion) @@ -320,32 +260,30 @@ def from_config(cls, config, atom_data): density, nuclide_mass_fraction, atom_data.atom_data.mass.copy() ) - if temperature is not None: - t_radiative = temperature - elif config.plasma.initial_t_rad > 0 * u.K: - t_radiative = ( - np.ones(geometry.no_of_shells + 1) * config.plasma.initial_t_rad - ) - else: - t_radiative = None + nuclide_mass_fraction = parse_abundance_config( + config, geometry, time_explosion + ) - #### Here starts the packetsource section + # using atom_data.mass.copy() to ensure that the original atom_data is not modified + composition = Composition( + density, nuclide_mass_fraction, atom_data.atom_data.mass.copy() + ) - if config.plasma.initial_t_inner < 0.0 * u.K: - luminosity_requested = config.supernova.luminosity_requested - t_inner = None - else: - luminosity_requested = None - t_inner = config.plasma.initial_t_inner + packet_source = parse_packet_source(config, geometry) + radiation_field_state = parse_radiation_field_state( + config, + t_radiative, + geometry, + dilution_factor=None, + packet_source=packet_source, + ) return cls( geometry=geometry, composition=composition, + radiation_field_state=radiation_field_state, time_explosion=time_explosion, - t_radiative=t_radiative, - t_inner=t_inner, - luminosity_requested=luminosity_requested, - dilution_factor=None, + packet_source=packet_source, electron_densities=electron_densities, ) @@ -415,13 +353,12 @@ def from_csvy(cls, config, atom_data=None): time_explosion = config.supernova.time_explosion.cgs electron_densities = None - temperature = None geometry = parse_csvy_geometry( config, csvy_model_config, csvy_model_data, time_explosion ) - composition = parse_composition_csvy( + composition = parse_csvy_composition( atom_data, csvy_model_config, csvy_model_data, @@ -429,52 +366,17 @@ def from_csvy(cls, config, atom_data=None): geometry, ) - # TODO -- implement t_radiative - # t_radiative = None - if temperature: - t_radiative = temperature - elif hasattr(csvy_model_data, "columns"): - if "t_rad" in csvy_model_data.columns: - t_rad_field_index = [ - field["name"] for field in csvy_model_config.datatype.fields - ].index("t_rad") - t_rad_unit = u.Unit( - csvy_model_config.datatype.fields[t_rad_field_index]["unit"] - ) - t_radiative = ( - csvy_model_data["t_rad"].iloc[0:].values * t_rad_unit - ) - else: - t_radiative = None - - dilution_factor = None - if hasattr(csvy_model_data, "columns"): - if "dilution_factor" in csvy_model_data.columns: - dilution_factor = ( - csvy_model_data["dilution_factor"].iloc[0:].to_numpy() - ) - - elif config.plasma.initial_t_rad > 0 * u.K: - t_radiative = ( - np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad - ) - else: - t_radiative = None + packet_source = parse_packet_source(config, geometry) - if config.plasma.initial_t_inner < 0.0 * u.K: - luminosity_requested = config.supernova.luminosity_requested - t_inner = None - else: - luminosity_requested = None - t_inner = config.plasma.initial_t_inner + radiation_field_state = parse_csvy_radiation_field_state( + config, csvy_model_config, csvy_model_data, geometry, packet_source + ) return cls( geometry=geometry, composition=composition, time_explosion=time_explosion, - t_radiative=t_radiative, - t_inner=t_inner, - luminosity_requested=luminosity_requested, - dilution_factor=dilution_factor, + radiation_field_state=radiation_field_state, + packet_source=packet_source, electron_densities=electron_densities, ) diff --git a/tardis/model/geometry/radial1d.py b/tardis/model/geometry/radial1d.py index c870b817882..a130a80cb65 100644 --- a/tardis/model/geometry/radial1d.py +++ b/tardis/model/geometry/radial1d.py @@ -76,6 +76,16 @@ def __init__( "Requesting inner boundary below inner shell. Extrapolating the inner cell" ) + @property + def v_middle(self): + return (self.v_inner + self.v_outer) / 2.0 + + @property + def v_middle_active(self): + return self.v_middle[ + self.v_inner_boundary_index : self.v_outer_boundary_index + ] + @property def v_inner_boundary_index(self): return np.clip( @@ -126,6 +136,14 @@ def r_outer(self): def r_outer_active(self): return (self.v_outer_active * self.time_explosion).cgs + @property + def r_middle(self): + return (self.v_middle * self.time_explosion).cgs + + @property + def r_middle_active(self): + return (self.v_middle_active * self.time_explosion).cgs + @property def volume(self): """Volume in shell computed from r_outer and r_inner""" diff --git a/tardis/model/matter/composition.py b/tardis/model/matter/composition.py index eb3ebd8efb6..e48e7a57682 100644 --- a/tardis/model/matter/composition.py +++ b/tardis/model/matter/composition.py @@ -8,6 +8,18 @@ def compile_rd_isotope_masses(): + """ + Compiles the masses of isotopes from the default data in RD_DEFAULT_DATA. + + Parameters + ---------- + None + + Returns + ------- + pandas.Series + A series containing the masses of isotopes, indexed by atomic number and mass number. + """ atomic_numbers = [] mass_numbers = [] nuclide_masses = [] diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py index 401457a2dfd..2dc8e76da99 100644 --- a/tardis/model/parse_input.py +++ b/tardis/model/parse_input.py @@ -5,6 +5,7 @@ import pandas as pd from astropy import units as u +from tardis import constants as const from tardis.io.model.parse_density_configuration import ( calculate_density_after_time, parse_config_v1_density, @@ -16,6 +17,8 @@ from tardis.model.geometry.radial1d import HomologousRadial1DGeometry from tardis.model.matter.composition import Composition from tardis.model.matter.decay import IsotopicMassFraction +from tardis.model.radiation_field_state import DiluteThermalRadiationFieldState +from tardis.montecarlo.packet_source import BlackBodySimpleSource from tardis.util.base import quantity_linspace logger = logging.getLogger(__name__) @@ -25,6 +28,42 @@ def parse_structure_config(config, time_explosion, enable_homology=True): """ Parse the structure configuration data. + Parameters + ---------- + config : object + The configuration data. + time_explosion : float + The time of the explosion. + enable_homology : bool, optional + Whether to enable homology (default is True). + + Returns + ------- + electron_densities : object + The parsed electron densities. + temperature : object + The parsed temperature. + geometry : object + The parsed geometry. + density : object + The parsed density. + + Raises + ------ + NotImplementedError + If the structure configuration type is not supported. + + Notes + ----- + This function parses the structure configuration data and returns the parsed electron + densities, temperature, geometry, and density. The structure configuration can be of + type 'specific' or 'file'. If it is of type 'specific', the velocity and density are + parsed from the configuration. If it is of type 'file', the velocity and density are + read from a file. The parsed data is used to create a homologous radial 1D geometry object. + """ + """ + Parse the structure configuration data. + Parameters ---------- config : object @@ -305,7 +344,7 @@ def convert_to_nuclide_mass_fraction(isotopic_mass_fraction, mass_fraction): return nuclide_mass_fraction -def parse_composition_csvy( +def parse_csvy_composition( atom_data, csvy_model_config, csvy_model_data, time_explosion, geometry ): """ @@ -479,3 +518,167 @@ def parse_density_csvy(csvy_model_config, csvy_model_data, time_explosion): ) return density + + +def parse_radiation_field_state( + config, t_radiative, geometry, dilution_factor=None, packet_source=None +): + """ + Parses the radiation field state based on the provided configuration, radiative temperature, geometry, dilution factor, and packet source. + + Parameters + ---------- + config : Config + The configuration object. + t_radiative : {None, Quantity}, optional + The radiative temperature. If None, it is calculated based on the initial_t_rad value in the plasma configuration. + geometry : Geometry + The geometry object. + dilution_factor : {None, ndarray}, optional + The dilution factor. If None, it is calculated based on the geometry. + packet_source : {None, PacketSource}, optional + The packet source object. + + Returns + ------- + DiluteThermalRadiationFieldState + The parsed radiation field state. + + Raises + ------ + AssertionError + If the length of t_radiative or dilution_factor is not compatible with the geometry. + """ + if t_radiative is None: + if config.plasma.initial_t_rad > 0 * u.K: + t_radiative = ( + np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad + ) + else: + t_radiative = calculate_t_radiative_from_t_inner( + geometry, packet_source + ) + + assert len(t_radiative) == geometry.no_of_shells + + if dilution_factor is None: + dilution_factor = calculate_geometric_dilution_factor(geometry) + elif len(dilution_factor) != geometry.no_of_shells: + dilution_factor = dilution_factor[ + geometry.v_inner_boundary_index : geometry.v_outer_boundary_index + ] + assert len(dilution_factor) == geometry.no_of_shells + + return DiluteThermalRadiationFieldState(t_radiative, dilution_factor) + + +def parse_packet_source(config, geometry): + """ + Parse the packet source based on the given configuration and geometry. + + Parameters + ---------- + config : Config + The configuration object containing the supernova and plasma settings. + geometry : Geometry + The geometry object containing the inner radius information. + + Returns + ------- + packet_source : BlackBodySimpleSource + The packet source object based on the configuration and geometry. + + Raises + ------ + ValueError + If both t_inner and luminosity_requested are None. + + """ + luminosity_requested = config.supernova.luminosity_requested + if config.plasma.initial_t_inner > 0.0 * u.K: + packet_source = BlackBodySimpleSource( + radius=geometry.r_inner[0], + temperature=config.plasma.initial_t_inner, + ) + elif (config.plasma.initial_t_inner < 0.0 * u.K) and ( + luminosity_requested is not None + ): + packet_source = BlackBodySimpleSource(radius=geometry.r_inner[0]) + packet_source.set_temperature_from_luminosity(luminosity_requested) + else: + raise ValueError( + "Both t_inner and luminosity_requested cannot be None." + ) + return packet_source + + +def parse_csvy_radiation_field_state( + config, csvy_model_config, csvy_model_data, geometry, packet_source +): + t_radiative = None + dilution_factor = None + + if hasattr(csvy_model_data, "columns") and ( + "t_rad" in csvy_model_data.columns + ): + t_rad_field_index = [ + field["name"] for field in csvy_model_config.datatype.fields + ].index("t_rad") + t_rad_unit = u.Unit( + csvy_model_config.datatype.fields[t_rad_field_index]["unit"] + ) + t_radiative = csvy_model_data["t_rad"].iloc[1:].values * t_rad_unit + + elif config.plasma.initial_t_rad > 0 * u.K: + t_radiative = ( + np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad + ) + t_radiative = ( + np.ones(geometry.no_of_shells) * config.plasma.initial_t_rad + ) + else: + t_radiative = calculate_t_radiative_from_t_inner( + geometry, packet_source + ) + + if hasattr(csvy_model_data, "columns") and ( + "dilution_factor" in csvy_model_data.columns + ): + dilution_factor = csvy_model_data["dilution_factor"].iloc[1:].values + else: + dilution_factor = calculate_geometric_dilution_factor(geometry) + + return DiluteThermalRadiationFieldState(t_radiative, dilution_factor) + + +def calculate_t_radiative_from_t_inner(geometry, packet_source): + """ + Calculates the radiative temperature based on the inner temperature and the geometry of the system. + + Parameters + ---------- + geometry : Geometry + The geometry object. + packet_source : PacketSource + The packet source object. + + Returns + ------- + Quantity + The calculated radiative temperature. + """ + lambda_wien_inner = const.b_wien / packet_source.temperature + t_radiative = const.b_wien / ( + lambda_wien_inner + * (1 + (geometry.v_middle - geometry.v_inner_boundary) / const.c) + ) + return t_radiative + + +def calculate_geometric_dilution_factor(geometry): + return 0.5 * ( + 1 + - np.sqrt( + 1 - (geometry.r_inner[0] ** 2 / geometry.r_middle**2).to(1).value + ) + ) diff --git a/tardis/model/radiation_field_state.py b/tardis/model/radiation_field_state.py new file mode 100644 index 00000000000..f6b1356f76b --- /dev/null +++ b/tardis/model/radiation_field_state.py @@ -0,0 +1,32 @@ +from tardis.montecarlo.montecarlo_numba.numba_interface import OpacityState + + +import numpy as np +from astropy import units as u + + +class DiluteThermalRadiationFieldState: + """ + Represents the state of a dilute thermal radiation field. + + + Parameters + ---------- + t_radiative : u.Quantity + Radiative temperature in each shell + dilution_factor : numpy.ndarray + Dilution Factors in each shell + """ + + def __init__( + self, + t_radiative: u.Quantity, + dilution_factor: np.ndarray, + ): + # ensuring that the radiation_field has both + # dilution_factor and t_radiative equal length + assert len(t_radiative) == len(dilution_factor) + assert np.all(t_radiative > 0 * u.K) + assert np.all(dilution_factor > 0) + self.t_radiative = t_radiative + self.dilution_factor = dilution_factor diff --git a/tardis/model/tests/test_base.py b/tardis/model/tests/test_base.py index 773d2302e62..c9a8a2f8819 100644 --- a/tardis/model/tests/test_base.py +++ b/tardis/model/tests/test_base.py @@ -471,7 +471,7 @@ def to_hdf_buffer(hdf_file_path, simulation_verysimple): def test_hdf_simulation_state_scalars( hdf_file_path, simulation_verysimple, attr ): - path = "model/scalars" + path = "simulation_state/scalars" expected = pd.read_hdf(hdf_file_path, path)[attr] actual = getattr(simulation_verysimple.simulation_state, attr) if hasattr(actual, "cgs"): @@ -479,14 +479,14 @@ def test_hdf_simulation_state_scalars( assert_almost_equal(actual, expected) -simulation_state_nparray_attrs = ["w", "v_inner", "v_outer"] +simulation_state_nparray_attrs = ["dilution_factor", "v_inner", "v_outer"] @pytest.mark.parametrize("attr", simulation_state_nparray_attrs) def test_hdf_simulation_state_nparray( hdf_file_path, simulation_verysimple, attr ): - path = f"model/{attr}" + path = f"simulation_state/{attr}" expected = pd.read_hdf(hdf_file_path, path) actual = getattr(simulation_verysimple.simulation_state, attr) if hasattr(actual, "cgs"): diff --git a/tardis/model/tests/test_density.py b/tardis/model/tests/test_density.py index fd7035ec9f7..907e2ef3e24 100644 --- a/tardis/model/tests/test_density.py +++ b/tardis/model/tests/test_density.py @@ -17,7 +17,7 @@ def test_hdf_density_0(hdf_file_path, simulation_verysimple): actual = simulation_verysimple.simulation_state.density if hasattr(actual, "cgs"): actual = actual.cgs.value - path = os.path.join("model", "density") + path = "simulation_state/density" expected = pd.read_hdf(hdf_file_path, path) assert_almost_equal(actual, expected.values) @@ -26,6 +26,6 @@ def test_hdf_time_0(hdf_file_path, simulation_verysimple): actual = simulation_verysimple.simulation_state.time_explosion if hasattr(actual, "cgs"): actual = actual.cgs.value - path = os.path.join("model", "scalars") + path = "simulation_state/scalars" expected = pd.read_hdf(hdf_file_path, path)["time_explosion"] assert_almost_equal(actual, expected) diff --git a/tardis/montecarlo/montecarlo_numba/formal_integral.py b/tardis/montecarlo/montecarlo_numba/formal_integral.py index c97c017360d..bdf170a4065 100644 --- a/tardis/montecarlo/montecarlo_numba/formal_integral.py +++ b/tardis/montecarlo/montecarlo_numba/formal_integral.py @@ -145,7 +145,6 @@ def numba_formal_integral( nu_end = nu_ends[i] nu_end_idx = nu_ends_idxs[i] for _ in range(max(nu_end_idx - pline, 0)): - # calculate e-scattering optical depth to next resonance point zend = ( model.time_explosion @@ -281,7 +280,6 @@ class FormalIntegrator(object): """ def __init__(self, model, plasma, transport, points=1000): - self.model = model self.transport = transport self.points = points @@ -418,7 +416,7 @@ def make_source_function(self): Numpy array containing ( 1 - exp(-tau_ul) ) S_ul ordered by wavelength of the transition u -> l """ - model = self.model + simulation_state = self.model transport = self.transport # macro_ref = self.atomic_data.macro_atom_references @@ -427,7 +425,7 @@ def make_source_function(self): macro_data = self.original_plasma.macro_atom_data no_lvls = len(self.levels_index) - no_shells = len(model.w) + no_shells = len(simulation_state.dilution_factor) if transport.line_interaction_type == "macroatom": internal_jump_mask = (macro_data.transition_type >= 0).values @@ -439,7 +437,9 @@ def make_source_function(self): source_level_idx = ma_int_data.source_level_idx.values destination_level_idx = ma_int_data.destination_level_idx.values - Edotlu_norm_factor = 1 / (transport.time_of_simulation * model.volume) + Edotlu_norm_factor = 1 / ( + transport.time_of_simulation * simulation_state.volume + ) exptau = 1 - np.exp(-self.original_plasma.tau_sobolevs) Edotlu = Edotlu_norm_factor * exptau * transport.Edotlu_estimator @@ -448,8 +448,13 @@ def make_source_function(self): Jbluelu_norm_factor = ( ( const.c.cgs - * model.time_explosion - / (4 * np.pi * transport.time_of_simulation * model.volume) + * simulation_state.time_explosion + / ( + 4 + * np.pi + * transport.time_of_simulation + * simulation_state.volume + ) ) .to("1/(cm^2 s)") .value @@ -495,7 +500,7 @@ def make_source_function(self): (self.atomic_data.macro_atom_data.transition_type == -1).values ] q_ul = tmp.set_index(transitions_index) - t = model.time_explosion.value + t = simulation_state.time_explosion.value lines = self.atomic_data.lines.set_index("line_id") wave = lines.wavelength_cm.loc[ transitions.transition_line_id diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_base.py b/tardis/montecarlo/montecarlo_numba/tests/test_base.py index cf10aa19bd2..320826cf68d 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_base.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_base.py @@ -17,36 +17,96 @@ def test_montecarlo_radial1d(): assert False -def test_montecarlo_main_loop( +@pytest.fixture(scope="function") +def montecarlo_main_loop_config( config_montecarlo_1e5_verysimple, - atomic_dataset, - tardis_ref_path, - tmpdir, - set_seed_fixture, - random_call_fixture, - request, ): montecarlo_configuration.LEGACY_MODE_ENABLED = True # Setup model config from verysimple - atomic_data = deepcopy(atomic_dataset) + config_montecarlo_1e5_verysimple.montecarlo.last_no_of_packets = 1e5 config_montecarlo_1e5_verysimple.montecarlo.no_of_virtual_packets = 0 config_montecarlo_1e5_verysimple.montecarlo.iterations = 1 config_montecarlo_1e5_verysimple.plasma.line_interaction_type = "macroatom" del config_montecarlo_1e5_verysimple["config_dirname"] + return config_montecarlo_1e5_verysimple + + +def test_montecarlo_main_loop( + montecarlo_main_loop_config, + tardis_ref_path, + request, + atomic_dataset, +): + atomic_dataset = deepcopy(atomic_dataset) + montecarlo_main_loop_simulation = Simulation.from_config( + montecarlo_main_loop_config, + atom_data=atomic_dataset, + virtual_packet_logging=False, + ) + montecarlo_main_loop_simulation.run_convergence() + montecarlo_main_loop_simulation.run_final() + + compare_fname = os.path.join( + tardis_ref_path, "montecarlo_1e5_compare_data.h5" + ) + if request.config.getoption("--generate-reference"): + montecarlo_main_loop_simulation.to_hdf(compare_fname, overwrite=True) + + # Load compare data from refdata + expected_nu = pd.read_hdf( + compare_fname, key="/simulation/transport/output_nu" + ).values + expected_energy = pd.read_hdf( + compare_fname, key="/simulation/transport/output_energy" + ).values + expected_nu_bar_estimator = pd.read_hdf( + compare_fname, key="/simulation/transport/nu_bar_estimator" + ).values + expected_j_estimator = pd.read_hdf( + compare_fname, key="/simulation/transport/j_estimator" + ).values + + actual_energy = montecarlo_main_loop_simulation.transport.output_energy + actual_nu = montecarlo_main_loop_simulation.transport.output_nu + actual_nu_bar_estimator = ( + montecarlo_main_loop_simulation.transport.nu_bar_estimator + ) + actual_j_estimator = montecarlo_main_loop_simulation.transport.j_estimator + + # Compare + npt.assert_allclose( + actual_nu_bar_estimator, expected_nu_bar_estimator, rtol=1e-13 + ) + npt.assert_allclose(actual_j_estimator, expected_j_estimator, rtol=1e-13) + npt.assert_allclose(actual_energy.value, expected_energy, rtol=1e-13) + npt.assert_allclose(actual_nu.value, expected_nu, rtol=1e-13) + + +@pytest.mark.xfail(reason="need to store virtual packet data in hdf5") +def test_montecarlo_main_loop_vpacket_log( + montecarlo_main_loop_config, + tardis_ref_path, + request, + atomic_dataset, +): + atomic_dataset = deepcopy(atomic_dataset) + montecarlo_main_loop_config.montecarlo.no_of_virtual_packets = 5 - sim = Simulation.from_config( - config_montecarlo_1e5_verysimple, atom_data=atomic_data + montecarlo_main_loop_simulation = Simulation.from_config( + montecarlo_main_loop_config, + atom_data=atomic_dataset, + virtual_packet_logging=True, ) - sim.run_convergence() - sim.run_final() + montecarlo_main_loop_simulation.run_convergence() + montecarlo_main_loop_simulation.run_final() compare_fname = os.path.join( tardis_ref_path, "montecarlo_1e5_compare_data.h5" ) if request.config.getoption("--generate-reference"): - sim.to_hdf(compare_fname, overwrite=True) + montecarlo_main_loop_simulation.to_hdf(compare_fname, overwrite=True) # Load compare data from refdata expected_nu = pd.read_hdf( @@ -62,10 +122,13 @@ def test_montecarlo_main_loop( compare_fname, key="/simulation/transport/j_estimator" ).values - actual_energy = sim.transport.output_energy - actual_nu = sim.transport.output_nu - actual_nu_bar_estimator = sim.transport.nu_bar_estimator - actual_j_estimator = sim.transport.j_estimator + transport = montecarlo_main_loop_simulation.transport + actual_energy = transport.output_energy + actual_nu = transport.output_nu + actual_nu_bar_estimator = transport.nu_bar_estimator + actual_j_estimator = montecarlo_main_loop_simulation.transport.j_estimator + actual_vpacket_log_nus = transport.virt_packet_nus + actual_vpacket_log_energies = transport.virt_packet_energies # Compare npt.assert_allclose( diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 2f3b68c83bf..88089363bf8 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -289,8 +289,8 @@ def advance_state(self): ) converged = self._get_convergence_status( - self.simulation_state.t_rad, - self.simulation_state.w, + self.simulation_state.t_radiative, + self.simulation_state.dilution_factor, self.simulation_state.t_inner, estimated_t_rad, estimated_w, @@ -299,13 +299,13 @@ def advance_state(self): # calculate_next_plasma_state equivalent # FIXME: Should convergence strategy have its own class? - next_t_rad = self.damped_converge( - self.simulation_state.t_rad, + next_t_radiative = self.damped_converge( + self.simulation_state.t_radiative, estimated_t_rad, self.convergence_strategy.t_rad.damping_constant, ) - next_w = self.damped_converge( - self.simulation_state.w, + next_dilution_factor = self.damped_converge( + self.simulation_state.dilution_factor, estimated_w, self.convergence_strategy.w.damping_constant, ) @@ -328,11 +328,13 @@ def advance_state(self): ) self.convergence_plots.fetch_data( name="t_rad", - value=self.simulation_state.t_rad, + value=self.simulation_state.t_radiative, item_type="iterable", ) self.convergence_plots.fetch_data( - name="w", value=self.simulation_state.w, item_type="iterable" + name="w", + value=self.simulation_state.dilution_factor, + item_type="iterable", ) self.convergence_plots.fetch_data( name="velocity", @@ -341,15 +343,15 @@ def advance_state(self): ) self.log_plasma_state( - self.simulation_state.t_rad, - self.simulation_state.w, + self.simulation_state.t_radiative, + self.simulation_state.dilution_factor, self.simulation_state.t_inner, - next_t_rad, - next_w, + next_t_radiative, + next_dilution_factor, next_t_inner, ) - self.simulation_state.t_rad = next_t_rad - self.simulation_state.w = next_w + 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 # model.calculate_j_blues() equivalent @@ -359,7 +361,8 @@ def advance_state(self): self.plasma.store_previous_properties() update_properties = dict( - t_rad=self.simulation_state.t_rad, w=self.simulation_state.w + 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. @@ -431,8 +434,8 @@ def run_convergence(self): while self.iterations_executed < self.iterations - 1: self.store_plasma_state( self.iterations_executed, - self.simulation_state.w, - self.simulation_state.t_rad, + self.simulation_state.dilution_factor, + self.simulation_state.t_radiative, self.plasma.electron_densities, self.simulation_state.t_inner, ) @@ -456,8 +459,8 @@ def run_final(self): """ self.store_plasma_state( self.iterations_executed, - self.simulation_state.w, - self.simulation_state.t_rad, + self.simulation_state.dilution_factor, + self.simulation_state.t_radiative, self.plasma.electron_densities, self.simulation_state.t_inner, ) @@ -485,10 +488,10 @@ def run_final(self): def log_plasma_state( self, t_rad, - w, + dilution_factor, t_inner, next_t_rad, - next_w, + next_dilution_factor, next_t_inner, log_sampling=5, ): @@ -517,8 +520,8 @@ def log_plasma_state( ) plasma_state_log["t_rad"] = t_rad plasma_state_log["next_t_rad"] = next_t_rad - plasma_state_log["w"] = w - plasma_state_log["next_w"] = next_w + plasma_state_log["w"] = dilution_factor + plasma_state_log["next_w"] = next_dilution_factor plasma_state_log.columns.name = "Shell No." if is_notebook(): diff --git a/tardis/visualization/widgets/shell_info.py b/tardis/visualization/widgets/shell_info.py index 0ac4f104cf3..205d727d037 100644 --- a/tardis/visualization/widgets/shell_info.py +++ b/tardis/visualization/widgets/shell_info.py @@ -18,7 +18,7 @@ class BaseShellInfo: def __init__( self, t_radiative, - w, + dilution_factor, abundance, number_density, ion_number_density, @@ -30,7 +30,7 @@ def __init__( ---------- t_radiative : array_like Radiative Temperature of each shell of simulation - w : array_like + dilution_factor : array_like Dilution Factor (W) of each shell of simulation model abundance : pandas.DataFrame Fractional abundance of elements where row labels are atomic number @@ -46,7 +46,7 @@ def __init__( number, ion number, level number) and column labels are shell number """ self.t_radiative = t_radiative - self.w = w + self.dilution_factor = dilution_factor self.abundance = abundance self.number_density = number_density self.ion_number_density = ion_number_density @@ -62,7 +62,10 @@ def shells_data(self): simulation model """ shells_temp_w = pd.DataFrame( - {"Rad. Temp.": self.t_radiative, "W": self.w} + { + "Rad. Temp.": self.t_radiative, + "Dilution Factor": self.dilution_factor, + } ) shells_temp_w.index = range( 1, len(self.t_radiative) + 1 @@ -185,7 +188,7 @@ def __init__(self, sim_model): """ super().__init__( sim_model.simulation_state.t_radiative, - sim_model.simulation_state.w, + sim_model.simulation_state.dilution_factor, sim_model.plasma.abundance, sim_model.plasma.number_density, sim_model.plasma.ion_number_density, @@ -211,7 +214,7 @@ def __init__(self, hdf_fpath): with pd.HDFStore(hdf_fpath, "r") as sim_data: super().__init__( sim_data["/simulation/simulation_state/t_radiative"], - sim_data["/simulation/simulation_state/w"], + sim_data["/simulation/simulation_state/dilution_factor"], sim_data["/simulation/plasma/abundance"], sim_data["/simulation/plasma/number_density"], sim_data["/simulation/plasma/ion_number_density"], diff --git a/tardis/visualization/widgets/tests/test_shell_info.py b/tardis/visualization/widgets/tests/test_shell_info.py index 4c08233d8d2..0afefc5e854 100644 --- a/tardis/visualization/widgets/tests/test_shell_info.py +++ b/tardis/visualization/widgets/tests/test_shell_info.py @@ -14,7 +14,7 @@ def base_shell_info(simulation_verysimple): return BaseShellInfo( simulation_verysimple.simulation_state.t_radiative, - simulation_verysimple.simulation_state.w, + simulation_verysimple.simulation_state.dilution_factor, simulation_verysimple.plasma.abundance, simulation_verysimple.plasma.number_density, simulation_verysimple.plasma.ion_number_density, @@ -48,7 +48,7 @@ def test_shells_data(self, base_shell_info, simulation_verysimple): ) assert np.allclose( shells_data.iloc[:, 1].map(np.float64), - simulation_verysimple.simulation_state.w, + simulation_verysimple.simulation_state.dilution_factor, ) @pytest.mark.parametrize("shell_num", [1, 20])