diff --git a/docs/io/optional/callback_example.ipynb b/docs/io/optional/callback_example.ipynb index 0b13d5bb8cd..04fa294f37e 100644 --- a/docs/io/optional/callback_example.ipynb +++ b/docs/io/optional/callback_example.ipynb @@ -78,7 +78,7 @@ "source": [ "def append_num_emitted_to_list(sim, lst):\n", " if sim.iterations_executed < sim.iterations:\n", - " num_emitted_packets = len(sim.transport.emitted_packet_nu)\n", + " num_emitted_packets = len(sim.transport.transport_state.emitted_packet_nu)\n", " lst.append(num_emitted_packets)" ] }, diff --git a/docs/io/optional/custom_source.ipynb b/docs/io/optional/custom_source.ipynb index 101df3df447..767421a792d 100644 --- a/docs/io/optional/custom_source.ipynb +++ b/docs/io/optional/custom_source.ipynb @@ -42,6 +42,9 @@ "from tardis import constants as const\n", "from astropy import units as u\n", "from tardis.montecarlo.packet_source import BlackBodySimpleSource\n", + "from tardis.montecarlo.montecarlo_numba.packet_collections import (\n", + " PacketCollection,\n", + ")\n", "from tardis import run_tardis\n", "import matplotlib.pyplot as plt\n", "from tardis.io.atom_data import download_atom_data" @@ -87,7 +90,7 @@ " self.truncation_wavelength = truncation_wavelength\n", " super().__init__(**kwargs)\n", "\n", - " def create_packets(self, no_of_packets, drawing_sample_size=None):\n", + " def create_packets(self, no_of_packets, drawing_sample_size=None, seed_offset=0, *args, **kwargs):\n", " \"\"\"\n", " Packet source that generates a truncated Blackbody source.\n", "\n", @@ -108,12 +111,17 @@ " Packet energies\n", " \"\"\"\n", "\n", + " self._reseed(self.base_seed + seed_offset)\n", + " packet_seeds = self.rng.choice(\n", + " self.MAX_SEED_VAL, no_of_packets, replace=True\n", + " )\n", + "\n", " # Makes uniform array of packet radii from blackbody source\n", - " radii = self.create_packet_radii(no_of_packets)\n", + " radii = self.create_packet_radii(no_of_packets, *args, **kwargs)\n", "\n", " # Use mus and energies from normal blackbody source.\n", - " mus = self.create_packet_mus(no_of_packets)\n", - " energies = self.create_packet_energies(no_of_packets)\n", + " mus = self.create_packet_mus(no_of_packets, *args, **kwargs)\n", + " energies = self.create_packet_energies(no_of_packets, *args, **kwargs)\n", "\n", " # If not specified, draw 2 times as many packets and reject any beyond no_of_packets.\n", " if drawing_sample_size is None:\n", @@ -128,18 +136,22 @@ "\n", " # Draw nus from blackbody distribution and reject based on truncation_frequency.\n", " # If more nus.shape[0] > no_of_packets use only the first no_of_packets.\n", - " nus = self.create_packet_nus(drawing_sample_size)\n", + " nus = self.create_packet_nus(drawing_sample_size, *args, **kwargs)\n", " nus = nus[nus < truncation_frequency][:no_of_packets]\n", "\n", " # Only required if the truncation wavelength is too big compared to the maximum\n", " # of the blackbody distribution. Keep sampling until nus.shape[0] > no_of_packets.\n", " while nus.shape[0] < no_of_packets:\n", - " additional_nus = self.create_packet_nus(drawing_sample_size)\n", + " additional_nus = self.create_packet_nus(drawing_sample_size, *args, **kwargs)\n", " mask = additional_nus < truncation_frequency\n", " additional_nus = additional_nus[mask][:no_of_packets]\n", " nus = np.hstack([nus, additional_nus])[:no_of_packets]\n", "\n", - " return radii, nus, mus, energies" + " radiation_field_luminosity = (\n", + " self.calculate_radfield_luminosity().to(u.erg / u.s).value\n", + " )\n", + "\n", + " return PacketCollection(radii, nus, mus, energies, packet_seeds, radiation_field_luminosity)" ] }, { @@ -179,11 +191,11 @@ "outputs": [], "source": [ "%matplotlib inline\n", - "plt.plot(mdl.transport.spectrum_virtual.wavelength,\n", - " mdl.transport.spectrum_virtual.luminosity_density_lambda,\n", + "plt.plot(mdl.transport.transport_state.spectrum_virtual.wavelength,\n", + " mdl.transport.transport_state.spectrum_virtual.luminosity_density_lambda,\n", " color='red', label='truncated blackbody (custom packet source)')\n", - "plt.plot(mdl_norm.transport.spectrum_virtual.wavelength,\n", - " mdl_norm.transport.spectrum_virtual.luminosity_density_lambda,\n", + "plt.plot(mdl_norm.transport.transport_state.spectrum_virtual.wavelength,\n", + " mdl_norm.transport.transport_state.spectrum_virtual.luminosity_density_lambda,\n", " color='blue', label='normal blackbody (default packet source)')\n", "plt.xlabel('$\\lambda [\\AA]$')\n", "plt.ylabel('$L_\\lambda$ [erg/s/$\\AA$]')\n", @@ -208,7 +220,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.7.10" + "version": "3.11.5" } }, "nbformat": 4, diff --git a/docs/physics/montecarlo/initialization.ipynb b/docs/physics/montecarlo/initialization.ipynb index f53662dee62..2ab262e6da1 100644 --- a/docs/physics/montecarlo/initialization.ipynb +++ b/docs/physics/montecarlo/initialization.ipynb @@ -79,6 +79,9 @@ "source": [ "import numpy as np\n", "from tardis.montecarlo.packet_source import BlackBodySimpleSource\n", + "from tardis.montecarlo.montecarlo_numba.packet_collections import (\n", + " PacketCollection,\n", + ")\n", "from astropy import units as u\n", "from tardis import constants as const\n", "import matplotlib.pyplot as plt" @@ -164,24 +167,21 @@ "# We define our packet source\n", "packet_source = BlackBodySimpleSource(base_seed=base_seed)\n", "\n", - "# Create separate packet seeds to sample each packet's montecarlo trajectory\n", - "packet_seeds = packet_source.create_packet_seeds(n_packets, seed_offset)\n", - "\n", "# Set radii and temperature from model\n", "packet_source.radius = r_boundary_inner\n", - "packet_source.temperature = temperature_inner.value\n", + "packet_source.temperature = temperature_inner\n", "\n", "# Create packets\n", - "radii, nus, mus, energies = packet_source.create_packets(n_packets)\n", + "packet_collection = packet_source.create_packets(n_packets)\n", "\n", "# Sets the energies in units of ergs\n", - "energies *= u.erg\n", + "packet_collection.initial_energies *= u.erg\n", "\n", "# Sets the frequencies in units of Hz\n", - "nus *= u.Hz\n", + "packet_collection.initial_nus *= u.Hz\n", "\n", - "print(\"Energies:\", energies)\n", - "print(\"Radii:\", radii)" + "print(\"Energies:\", packet_collection.initial_energies)\n", + "print(\"Radii:\", packet_collection.initial_radii)" ] }, { @@ -205,7 +205,7 @@ "print(\"Time of simulation:\", t_simulation)\n", "\n", "# Array of luminosity contribution by each packet\n", - "lumin_per_packet = energies / t_simulation\n", + "lumin_per_packet = packet_collection.initial_energies / t_simulation\n", "print(\"Luminosity per packet:\", lumin_per_packet)" ] }, @@ -263,15 +263,15 @@ "source": [ "# We set important quantites for making our histogram\n", "bins = 200\n", - "nus_planck = np.linspace(min(nus), max(nus), bins)\n", + "nus_planck = np.linspace(min(packet_collection.initial_nus), max(packet_collection.initial_nus), bins).value\n", "bin_width = nus_planck[1] - nus_planck[0]\n", "\n", "# In the histogram plot below, the weights argument is used\n", "# to make sure our plotted spectrum has the correct y-axis scale\n", - "plt.hist(nus.value, bins=bins, weights=lumin_per_packet / bin_width)\n", + "plt.hist(packet_collection.initial_nus.value, bins=bins, weights=lumin_per_packet / bin_width)\n", "\n", "# We plot the planck function for comparison\n", - "plt.plot(nus_planck, planck_function(nus_planck))\n", + "plt.plot(nus_planck * u.Hz, planck_function(nus_planck * u.Hz))\n", "\n", "plt.xlabel(\"Frequency (Hz)\")\n", "plt.ylabel(\"Luminosity density w.r.t. frequency (erg/s/Hz)\")\n", @@ -296,7 +296,7 @@ "source": [ "x = np.linspace(0, 1, 1000)\n", "\n", - "plt.hist(mus, bins=bins, density=True)\n", + "plt.hist(packet_collection.initial_mus, bins=bins, density=True)\n", "plt.plot(x, 2 * x)\n", "plt.xlabel(\"Propagation direction\")\n", "plt.ylabel(\"Probability density\")\n", @@ -312,7 +312,7 @@ "source": [ "thetas = np.linspace(0, np.pi / 2, 1000)\n", "\n", - "plt.hist(np.arccos(mus), bins=bins, density=True)\n", + "plt.hist(np.arccos(packet_collection.initial_mus), bins=bins, density=True)\n", "plt.plot(thetas, np.sin(2 * thetas))\n", "plt.xlabel(\"Angle with normal (rad)\")\n", "plt.ylabel(\"Probability density\")\n", @@ -348,7 +348,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.13" + "version": "3.11.5" } }, "nbformat": 4, diff --git a/docs/physics/spectrum/basic.ipynb b/docs/physics/spectrum/basic.ipynb index d0d61571ea7..2ac6c24799f 100644 --- a/docs/physics/spectrum/basic.ipynb +++ b/docs/physics/spectrum/basic.ipynb @@ -47,7 +47,7 @@ "source": [ "from tardis.io.configuration.config_reader import Configuration\n", "from tardis.simulation import Simulation\n", - "from tardis.montecarlo import TARDISSpectrum\n", + "from tardis.montecarlo.spectrum import TARDISSpectrum\n", "from tardis.io.atom_data.util import download_atom_data\n", "from astropy import units as u\n", "import numpy as np\n", @@ -118,7 +118,7 @@ "metadata": {}, "outputs": [], "source": [ - "nus = sim.transport.output_nu\n", + "nus = sim.transport.transport_state.output_nu\n", "nus" ] }, @@ -129,7 +129,7 @@ "metadata": {}, "outputs": [], "source": [ - "energies = sim.transport.output_energy\n", + "energies = sim.transport.transport_state.output_energy\n", "energies" ] }, @@ -183,7 +183,7 @@ "metadata": {}, "outputs": [], "source": [ - "luminosities = energies / sim.transport.time_of_simulation\n", + "luminosities = energies / sim.transport.transport_state.time_of_simulation\n", "luminosities" ] }, @@ -203,7 +203,7 @@ "metadata": {}, "outputs": [], "source": [ - "emitted_mask = sim.transport.emitted_packet_mask\n", + "emitted_mask = sim.transport.transport_state.emitted_packet_mask\n", "emitted_mask" ] }, @@ -460,7 +460,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.7.12" + "version": "3.11.5" } }, "nbformat": 4, diff --git a/docs/physics/update_and_conv/update_and_conv.ipynb b/docs/physics/update_and_conv/update_and_conv.ipynb index 971249e6a75..b0c8da0d24c 100644 --- a/docs/physics/update_and_conv/update_and_conv.ipynb +++ b/docs/physics/update_and_conv/update_and_conv.ipynb @@ -311,7 +311,7 @@ }, "outputs": [], "source": [ - "j_estimator = transport.j_estimator * (u.erg * u.cm) \n", + "j_estimator = transport.transport_state.estimators.j_estimator * (u.erg * u.cm) \n", "j_estimator" ] }, @@ -324,7 +324,7 @@ }, "outputs": [], "source": [ - "nu_bar_estimator = transport.nu_bar_estimator * (u.erg * u.cm * u.Hz)\n", + "nu_bar_estimator = transport.transport_state.estimators.nu_bar_estimator * (u.erg * u.cm * u.Hz)\n", "nu_bar_estimator" ] }, @@ -345,7 +345,7 @@ "outputs": [], "source": [ "V = simulation_state.volume\n", - "Delta_t = transport.calculate_time_of_simulation(simulation_state)\n", + "Delta_t = transport.transport_state.time_of_simulation\n", "prefactor = 1 / (4 * np.pi * V * Delta_t)\n", "J = prefactor * j_estimator\n", "J" @@ -462,7 +462,7 @@ "#nu_lower = tardis_config.supernova.luminosity_wavelength_end.to(u.Hz, u.spectral)\n", "#nu_upper = tardis_config.supernova.luminosity_wavelength_start.to(u.Hz, u.spectral)\n", "\n", - "L_output = transport.calculate_emitted_luminosity(0,np.inf)\n", + "L_output = transport.transport_state.calculate_emitted_luminosity(0,np.inf)\n", "L_output" ] }, @@ -570,6 +570,14 @@ "source": [ "plasma.electron_densities" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bd488259", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { @@ -588,7 +596,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.13" + "version": "3.11.5" }, "vscode": { "interpreter": { diff --git a/docs/quickstart.ipynb b/docs/quickstart.ipynb index b694d9dcca1..33e69d0d773 100644 --- a/docs/quickstart.ipynb +++ b/docs/quickstart.ipynb @@ -145,9 +145,9 @@ "metadata": {}, "outputs": [], "source": [ - "spectrum = sim.transport.spectrum\n", - "spectrum_virtual = sim.transport.spectrum_virtual\n", - "spectrum_integrated = sim.transport.spectrum_integrated" + "spectrum = sim.transport.transport_state.spectrum\n", + "spectrum_virtual = sim.transport.transport_state.spectrum_virtual\n", + "spectrum_integrated = sim.transport.transport_state.spectrum_integrated" ] }, { @@ -191,7 +191,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.17" + "version": "3.11.5" } }, "nbformat": 4, diff --git a/tardis/analysis.py b/tardis/analysis.py index 6ead15f9a43..56c4eb1bbd5 100644 --- a/tardis/analysis.py +++ b/tardis/analysis.py @@ -15,14 +15,14 @@ class LastLineInteraction(object): @classmethod - def from_model(cls, model, packet_filter_mode="packet_out_nu"): + def from_simulation(cls, simulation, packet_filter_mode="packet_out_nu"): return cls( - model.transport.last_line_interaction_in_id, - model.transport.last_line_interaction_out_id, - model.transport.last_line_interaction_shell_id, - model.transport.output_nu, - model.transport.last_interaction_in_nu, - model.plasma.atomic_data.lines, + simulation.transport.last_line_interaction_in_id, + simulation.transport.last_line_interaction_out_id, + simulation.transport.last_line_interaction_shell_id, + simulation.transport.transport_state.packet_collection.output_nus, + simulation.transport.last_interaction_in_nu, + simulation.plasma.atomic_data.lines, packet_filter_mode, ) diff --git a/tardis/gui/widgets.py b/tardis/gui/widgets.py index d5e11922857..7397575da5a 100644 --- a/tardis/gui/widgets.py +++ b/tardis/gui/widgets.py @@ -406,7 +406,6 @@ def match_dicts(self, dict1, dict2): # dict1<=dict2 elif isinstance(dict2[key], list): if isinstance(dict2[key][1], list): - # options = dict2[key][1] #This is passed by reference. # So copy the list manually. options = [ @@ -937,8 +936,8 @@ def __init__(self, parent, wavelength_start, wavelength_end, tablecreator): f"Line Interaction: {wavelength_start:.2f} - {wavelength_end:.2f} (A) " ) self.layout = QtWidgets.QVBoxLayout() - packet_nu_line_interaction = analysis.LastLineInteraction.from_model( - self.parent.model + packet_nu_line_interaction = ( + analysis.LastLineInteraction.from_simulation(self.parent.model) ) packet_nu_line_interaction.packet_filter_mode = "packet_nu" packet_nu_line_interaction.wavelength_start = ( @@ -946,8 +945,8 @@ def __init__(self, parent, wavelength_start, wavelength_end, tablecreator): ) packet_nu_line_interaction.wavelength_end = wavelength_end * u.angstrom - line_in_nu_line_interaction = analysis.LastLineInteraction.from_model( - self.parent.model + line_in_nu_line_interaction = ( + analysis.LastLineInteraction.from_simulation(self.parent.model) ) line_in_nu_line_interaction.packet_filter_mode = "line_in_nu" line_in_nu_line_interaction.wavelength_start = ( diff --git a/tardis/io/model/hdf.py b/tardis/io/model/hdf.py deleted file mode 100644 index 1621b1134ce..00000000000 --- a/tardis/io/model/hdf.py +++ /dev/null @@ -1,28 +0,0 @@ -import h5py - -from tardis.io.model.model_reader import simulation_state_to_dict - - -def store_simulation_state_to_hdf(simulation_state, fname): - """ - Stores data from SimulationState object into a hdf file. - - Parameters - ---------- - simulation_state : tardis.model.SimulationState - fname : str - """ - with h5py.File(fname, "a") as f: - simulation_state_group = f.require_group("simulation_state") - simulation_state_group.clear() - - simulation_state_dict = simulation_state_to_dict(simulation_state) - - for key, value in simulation_state_dict.items(): - if key.endswith("_cgs"): - simulation_state_group.create_dataset(key, data=value[0]) - simulation_state_group.create_dataset( - key + "_unit", data=value[1] - ) - else: - simulation_state_group.create_dataset(key, data=value) diff --git a/tardis/io/model/model_reader.py b/tardis/io/model/model_reader.py deleted file mode 100644 index 49be1cf5c1f..00000000000 --- a/tardis/io/model/model_reader.py +++ /dev/null @@ -1,314 +0,0 @@ -import logging - -import h5py -from astropy import units as u - -from tardis.io.configuration.config_reader import ConfigurationNameSpace -from tardis.montecarlo.base import MontecarloTransport -from tardis.montecarlo.packet_source import ( - BlackBodySimpleSource, - BlackBodySimpleSourceRelativistic, -) - -# Adding logging support -logger = logging.getLogger(__name__) - - -def transport_to_dict(transport): - """ - Retrieves all the data from a transport object and returns a dictionary. - - Parameters - ---------- - transport : tardis.montecarlo.MontecarloTransport - - Returns - ------- - transport_dict : dict - integrator_settings : dict - v_packet_settings : dict - virtual_spectrum_spawn_range : dict - """ - transport_dict = { - "Edotlu_estimator": transport.Edotlu_estimator, - "bf_heating_estimator": transport.bf_heating_estimator, - "disable_electron_scattering": transport.disable_electron_scattering, - "enable_full_relativity": transport.enable_full_relativity, - "enable_reflective_inner_boundary": transport.enable_reflective_inner_boundary, - "inner_boundary_albedo": transport.inner_boundary_albedo, - "input_energy": transport.input_energy, - "input_mu": transport.input_mu, - "input_nu": transport.input_nu, - "input_r": transport.input_r, - "j_blue_estimator": transport.j_blue_estimator, - "j_estimator": transport.j_estimator, - "last_interaction_in_nu": transport.last_interaction_in_nu, - "last_interaction_type": transport.last_interaction_type, - "last_line_interaction_in_id": transport.last_line_interaction_in_id, - "last_line_interaction_out_id": transport.last_line_interaction_out_id, - "last_line_interaction_shell_id": transport.last_line_interaction_shell_id, - "line_interaction_type": transport.line_interaction_type, - "nu_bar_estimator": transport.nu_bar_estimator, - "photo_ion_estimator": transport.photo_ion_estimator, - "photo_ion_estimator_statistics": transport.photo_ion_estimator_statistics, - "r_inner": transport.r_inner_cgs, - "r_outer": transport.r_outer_cgs, - "packet_source_base_seed": transport.packet_source.base_seed, - "spectrum_frequency_cgs": transport.spectrum_frequency, - "spectrum_method": transport.spectrum_method, - "stim_recomb_cooling_estimator": transport.stim_recomb_cooling_estimator, - "stim_recomb_estimator": transport.stim_recomb_estimator, - "time_of_simulation_cgs": transport.time_of_simulation, - "use_gpu": transport.use_gpu, - "v_inner": transport.v_inner_cgs, - "v_outer": transport.v_outer_cgs, - "nthreads": transport.nthreads, - "virt_logging": transport.virt_logging, - "virt_packet_energies": transport.virt_packet_energies, - "virt_packet_initial_mus": transport.virt_packet_initial_mus, - "virt_packet_initial_rs": transport.virt_packet_initial_rs, - "virt_packet_last_interaction_in_nu": transport.virt_packet_last_interaction_in_nu, - "virt_packet_last_interaction_type": transport.virt_packet_last_interaction_type, - "virt_packet_last_line_interaction_in_id": transport.virt_packet_last_line_interaction_in_id, - "virt_packet_last_line_interaction_out_id": transport.virt_packet_last_line_interaction_out_id, - "virt_packet_last_line_interaction_shell_id": transport.virt_packet_last_line_interaction_shell_id, - "virt_packet_nus": transport.virt_packet_nus, - "volume_cgs": transport.volume, - } - - for key, value in transport_dict.items(): - if key.endswith("_cgs"): - transport_dict[key] = [value.cgs.value, value.unit.to_string()] - - integrator_settings = transport.integrator_settings - v_packet_settings = transport.v_packet_settings - virtual_spectrum_spawn_range = transport.virtual_spectrum_spawn_range - - return ( - transport_dict, - integrator_settings, - v_packet_settings, - virtual_spectrum_spawn_range, - ) - - -def store_transport_to_hdf(transport, fname): - """ - Stores data from MontecarloTransport object into a hdf file. - - Parameters - ---------- - transport : tardis.montecarlo.MontecarloTransport - filename : str - """ - with h5py.File(fname, "a") as f: - transport_group = f.require_group("transport") - transport_group.clear() - - ( - transport_data, - integrator_settings, - v_packet_settings, - virtual_spectrum_spawn_range, - ) = transport_to_dict(transport) - - for key, value in transport_data.items(): - if key.endswith("_cgs"): - transport_group.create_dataset(key, data=value[0]) - transport_group.create_dataset(key + "_unit", data=value[1]) - else: - if value is not None: - transport_group.create_dataset(key, data=value) - - integrator_settings_group = transport_group.create_group( - "integrator_settings" - ) - for key, value in integrator_settings.items(): - integrator_settings_group.create_dataset(key, data=value) - - v_packet_settings_group = transport_group.create_group( - "v_packet_settings" - ) - for key, value in v_packet_settings.items(): - v_packet_settings_group.create_dataset(key, data=value) - - virtual_spectrum_spawn_range_group = transport_group.create_group( - "virtual_spectrum_spawn_range" - ) - for key, value in virtual_spectrum_spawn_range.items(): - virtual_spectrum_spawn_range_group.create_dataset(key, data=value) - - -def transport_from_hdf(fname): - """ - Creates a MontecarloTransport object using data stored in a hdf file. - - Parameters - ---------- - fname : str - - Returns - ------- - new_transport : tardis.montecarlo.MontecarloTransport - """ - d = {} - - # Loading data from hdf file - with h5py.File(fname, "r") as f: - transport_group = f["transport"] - for key, value in transport_group.items(): - if not key.endswith("_unit"): - if isinstance(value, h5py.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 transport_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"))) - - # Using packet source seed to packet source - if not d["enable_full_relativity"]: - d["packet_source"] = BlackBodySimpleSource(d["packet_source_base_seed"]) - else: - d["packet_source"] = BlackBodySimpleSourceRelativistic( - d["packet_source_base_seed"] - ) - - # Converting virtual spectrum spawn range values to astropy quantities - vssr = d["virtual_spectrum_spawn_range"] - d["virtual_spectrum_spawn_range"] = { - "start": u.Quantity(vssr["start"], unit=u.Unit("Angstrom")), - "end": u.Quantity(vssr["end"], unit=u.Unit("Angstrom")), - } - - # Converting dictionaries to ConfigurationNameSpace - d["integrator_settings"] = ConfigurationNameSpace(d["integrator_settings"]) - d["v_packet_settings"] = ConfigurationNameSpace(d["v_packet_settings"]) - d["virtual_spectrum_spawn_range"] = ConfigurationNameSpace( - d["virtual_spectrum_spawn_range"] - ) - - # Creating a transport object and storing data - new_transport = MontecarloTransport( - spectrum_frequency=d["spectrum_frequency_cgs"], - virtual_spectrum_spawn_range=d["virtual_spectrum_spawn_range"], - disable_electron_scattering=d["disable_electron_scattering"], - enable_reflective_inner_boundary=d["enable_reflective_inner_boundary"], - enable_full_relativity=d["enable_full_relativity"], - inner_boundary_albedo=d["inner_boundary_albedo"], - line_interaction_type=d["line_interaction_type"], - integrator_settings=d["integrator_settings"], - v_packet_settings=d["v_packet_settings"], - spectrum_method=d["spectrum_method"], - packet_source=d["packet_source"], - nthreads=d["nthreads"], - virtual_packet_logging=d["virt_logging"], - use_gpu=d["use_gpu"], - ) - - new_transport.Edotlu_estimator = d["Edotlu_estimator"] - new_transport.bf_heating_estimator = d["bf_heating_estimator"] - new_transport.input_energy = d["input_energy"] - new_transport.input_mu = d["input_mu"] - new_transport.input_nu = d["input_nu"] - new_transport.input_r = d["input_r_cgs"] - new_transport.j_blue_estimator = d["j_blue_estimator"] - new_transport.j_estimator = d["j_estimator"] - new_transport.last_interaction_in_nu = d["last_interaction_in_nu"] - new_transport.last_interaction_type = d["last_interaction_type"] - new_transport.last_line_interaction_in_id = d["last_line_interaction_in_id"] - new_transport.last_line_interaction_out_id = d[ - "last_line_interaction_out_id" - ] - new_transport.last_line_interaction_shell_id = d[ - "last_line_interaction_shell_id" - ] - new_transport.nu_bar_estimator = d["nu_bar_estimator"] - new_transport.photo_ion_estimator = d["photo_ion_estimator"] - new_transport.photo_ion_estimator_statistics = d[ - "photo_ion_estimator_statistics" - ] - new_transport.r_inner_cgs = d["r_inner"] - new_transport.r_outer_cgs = d["r_outer"] - new_transport.stim_recomb_cooling_estimator = d[ - "stim_recomb_cooling_estimator" - ] - new_transport.stim_recomb_estimator = d["stim_recomb_estimator"] - new_transport.time_of_simulation = d["time_of_simulation_cgs"] - new_transport.v_inner_cgs = d["v_inner"] - new_transport.v_outer_cgs = d["v_outer"] - new_transport.virt_packet_energies = d["virt_packet_energies"] - new_transport.virt_packet_initial_mus = d["virt_packet_initial_mus"] - new_transport.virt_packet_initial_rs = d["virt_packet_initial_rs"] - new_transport.virt_packet_last_interaction_in_nu = d[ - "virt_packet_last_interaction_in_nu" - ] - new_transport.virt_packet_last_interaction_type = d[ - "virt_packet_last_interaction_type" - ] - new_transport.virt_packet_last_line_interaction_in_id = d[ - "virt_packet_last_line_interaction_in_id" - ] - new_transport.virt_packet_last_line_interaction_out_id = d[ - "virt_packet_last_line_interaction_out_id" - ] - new_transport.virt_packet_last_line_interaction_shell_id = d[ - "virt_packet_last_line_interaction_shell_id" - ] - new_transport.virt_packet_nus = d["virt_packet_nus"] - new_transport.volume = d["volume_cgs"] - - return new_transport - - -def simulation_state_to_dict(simulation_state): - """ - Retrieves all the data from a SimulationState object and returns a dictionary. - - Parameters - ---------- - transport : tardis.model.SimulationState - - Returns - ------- - model_dict : dict - isotope_abundance : dict - """ - simulation_state_dict = { - "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(): - if hasattr(value, "unit"): - simulation_state_dict[key] = [ - value.cgs.value, - value.cgs.unit.to_string(), - ] - - return simulation_state_dict diff --git a/tardis/io/model_reader.py b/tardis/io/model_reader.py index 6856481b71e..67beb541ab2 100644 --- a/tardis/io/model_reader.py +++ b/tardis/io/model_reader.py @@ -12,7 +12,7 @@ from radioactivedecay.utils import Z_DICT, elem_to_Z from tardis.io.configuration.config_reader import ConfigurationNameSpace -from tardis.montecarlo.base import MontecarloTransport +from tardis.montecarlo.base import MonteCarloTransportSolver from tardis.montecarlo.packet_source import ( BlackBodySimpleSource, BlackBodySimpleSourceRelativistic, @@ -726,7 +726,7 @@ def transport_from_hdf(fname): ) # Creating a transport object and storing data - new_transport = MontecarloTransport( + new_transport = MonteCarloTransportSolver( spectrum_frequency=d["spectrum_frequency_cgs"], virtual_spectrum_spawn_range=d["virtual_spectrum_spawn_range"], disable_electron_scattering=d["disable_electron_scattering"], @@ -739,7 +739,7 @@ def transport_from_hdf(fname): spectrum_method=d["spectrum_method"], packet_source=d["packet_source"], nthreads=d["nthreads"], - virtual_packet_logging=d["virt_logging"], + enable_virtual_packet_logging=d["virt_logging"], use_gpu=d["use_gpu"], ) diff --git a/tardis/io/tests/test_model_reader.py b/tardis/io/tests/test_model_reader.py index 7cbd41f4c99..97a6ac7fe63 100644 --- a/tardis/io/tests/test_model_reader.py +++ b/tardis/io/tests/test_model_reader.py @@ -4,12 +4,6 @@ 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, - 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, @@ -127,344 +121,3 @@ def test_simple_read_cmfgen_density(cmfgen_fname): ) assert len(mean_density) == 9 assert len(velocity) == len(mean_density) + 1 - - -def test_model_to_dict(simulation_verysimple): - simulation_state = simulation_verysimple.simulation_state - - simulation_state_dict = simulation_state_to_dict(simulation_state) - - # Check model dictionary - assert np.array_equal( - 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 np.array_equal( - simulation_state_dict["v_boundary_inner_cgs"][0], - simulation_state.v_boundary_inner.cgs.value, - ) - assert ( - simulation_state_dict["v_boundary_inner_cgs"][1] - == simulation_state.v_boundary_inner.cgs.unit.to_string() - ) - assert np.array_equal( - simulation_state_dict["v_boundary_outer_cgs"][0], - simulation_state.v_boundary_outer.cgs.value, - ) - assert ( - simulation_state_dict["v_boundary_outer_cgs"][1] - == simulation_state.v_boundary_outer.cgs.unit.to_string() - ) - - assert np.array_equal( - simulation_state_dict["r_inner_cgs"][0], - simulation_state.r_inner.cgs.value, - ) - assert ( - simulation_state_dict["r_inner_cgs"][1] - == simulation_state.r_inner.cgs.unit.to_string() - ) - assert np.array_equal( - simulation_state_dict["density_cgs"][0], - simulation_state.density.cgs.value, - ) - assert ( - simulation_state_dict["density_cgs"][1] - == simulation_state.density.cgs.unit.to_string() - ) - - -def test_store_model_to_hdf(simulation_verysimple, tmp_path): - simulation_state = simulation_verysimple.simulation_state - - fname = tmp_path / "simulation_state.h5" - - # Store model object - store_simulation_state_to_hdf(simulation_state, fname) - - # Check file contents - with h5py.File(fname) as f: - assert np.array_equal( - 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["simulation_state/v_boundary_outer_cgs"], - simulation_state.v_boundary_outer.cgs.value, - ) - assert np.array_equal( - f["simulation_state/r_inner_cgs"], - simulation_state.r_inner.cgs.value, - ) - assert np.array_equal( - f["simulation_state/density_cgs"], - simulation_state.density.cgs.value, - ) - - -def test_transport_to_dict(simulation_verysimple): - transport = simulation_verysimple.transport - transport_data = transport.__dict__ - - ( - transport_dict, - integrator_settings, - v_packet_settings, - virtual_spectrum_spawn_range, - ) = transport_to_dict(transport) - - # Check transport dictionary - for key, value in transport_dict.items(): - if isinstance(value, np.ndarray): - if key + "_cgs" in transport_data.keys(): - assert np.array_equal(value, transport_data[key + "_cgs"]) - else: - assert np.array_equal(value, transport_data[key]) - elif isinstance(value, list): - assert np.array_equal(value[0], transport_data[key[:-4]].value) - assert value[1] == transport_data[key[:-4]].unit.to_string() - elif key == "packet_source_base_seed": # Check packet source base seed - assert value == transport_data["packet_source"].base_seed - else: - assert value == transport_data[key] - - # Check integrator settings - for key, value in integrator_settings.items(): - assert value == transport.integrator_settings[key] - - # Check v_packet settings - for key, value in v_packet_settings.items(): - assert value == transport.v_packet_settings[key] - - # Check virtual spectrum spawn range - for key, value in virtual_spectrum_spawn_range.items(): - assert value.value == transport.virtual_spectrum_spawn_range[key].value - - -def test_store_transport_to_hdf(simulation_verysimple, tmp_path): - transport = simulation_verysimple.transport - transport_data = transport.__dict__ - fname = tmp_path / "transport.h5" - # s - # Store transport object - store_transport_to_hdf(transport, fname) - # - # Check file contents - with h5py.File(fname) as f: - assert np.array_equal( - f["transport/Edotlu_estimator"], transport_data["Edotlu_estimator"] - ) - assert np.array_equal( - f["transport/bf_heating_estimator"], - transport_data["bf_heating_estimator"], - ) - assert ( - f["transport/enable_full_relativity"][()] - == transport_data["enable_full_relativity"] - ) - assert ( - f["transport/enable_reflective_inner_boundary"][()] - == transport_data["enable_reflective_inner_boundary"] - ) - assert ( - f["transport/inner_boundary_albedo"][()] - == transport_data["inner_boundary_albedo"] - ) - assert np.array_equal( - f["transport/input_energy"], transport_data["input_energy"] - ) - assert np.array_equal( - f["transport/input_mu"], transport_data["input_mu"] - ) - assert np.array_equal( - f["transport/input_nu"], transport_data["input_nu"] - ) - assert np.array_equal(f["transport/input_r"], transport_data["input_r"]) - assert np.array_equal( - f["transport/j_blue_estimator"], transport_data["j_blue_estimator"] - ) - assert np.array_equal( - f["transport/j_estimator"], transport_data["j_estimator"] - ) - assert np.array_equal( - f["transport/last_interaction_in_nu"], - transport_data["last_interaction_in_nu"], - ) - assert np.array_equal( - f["transport/last_interaction_type"], - transport_data["last_interaction_type"], - ) - assert np.array_equal( - f["transport/last_line_interaction_in_id"], - transport_data["last_line_interaction_in_id"], - ) - assert np.array_equal( - f["transport/last_line_interaction_out_id"], - transport_data["last_line_interaction_out_id"], - ) - assert np.array_equal( - f["transport/last_line_interaction_shell_id"], - transport_data["last_line_interaction_shell_id"], - ) - if hasattr(f["transport/line_interaction_type"][()], "decode"): - assert ( - f["transport/line_interaction_type"][()].decode("utf-8") - == transport_data["line_interaction_type"] - ) - else: - assert np.array_equal( - f["transport/line_interaction_type"][()], - transport_data["line_interaction_type"], - ) - assert np.array_equal( - f["transport/nu_bar_estimator"], transport_data["nu_bar_estimator"] - ) - assert np.array_equal( - f["transport/photo_ion_estimator"], - transport_data["photo_ion_estimator"], - ) - assert np.array_equal( - f["transport/photo_ion_estimator_statistics"], - transport_data["photo_ion_estimator_statistics"], - ) - assert np.array_equal( - f["transport/r_inner"], transport_data["r_inner_cgs"] - ) - assert np.array_equal( - f["transport/r_outer"], transport_data["r_outer_cgs"] - ) - assert ( - f["transport/packet_source_base_seed"][()] - == transport_data["packet_source"].base_seed - ) - assert np.array_equal( - f["transport/spectrum_frequency_cgs"], - transport_data["spectrum_frequency"].value, - ) - if hasattr(f["transport/spectrum_method"][()], "decode"): - assert ( - f["transport/spectrum_method"][()].decode("utf-8") - == transport_data["spectrum_method"] - ) - else: - assert np.array_equal( - f["transport/spectrum_method"][()], - transport_data["spectrum_method"], - ) - assert np.array_equal( - f["transport/stim_recomb_cooling_estimator"], - transport_data["stim_recomb_cooling_estimator"], - ) - assert np.array_equal( - f["transport/stim_recomb_estimator"], - transport_data["stim_recomb_estimator"], - ) - assert np.array_equal( - f["transport/time_of_simulation_cgs"], - transport_data["time_of_simulation"].value, - ) - assert f["transport/use_gpu"][()] == transport_data["use_gpu"] - assert np.array_equal( - f["transport/v_inner"], transport_data["v_inner_cgs"] - ) - assert np.array_equal( - f["transport/v_outer"], transport_data["v_outer_cgs"] - ) - assert f["transport/nthreads"][()] == transport_data["nthreads"] - assert f["transport/virt_logging"][()] == transport_data["virt_logging"] - assert np.array_equal( - f["transport/virt_packet_energies"], - transport_data["virt_packet_energies"], - ) - assert np.array_equal( - f["transport/virt_packet_initial_mus"], - transport_data["virt_packet_initial_mus"], - ) - assert np.array_equal( - f["transport/virt_packet_initial_rs"], - transport_data["virt_packet_initial_rs"], - ) - assert np.array_equal( - f["transport/virt_packet_last_interaction_in_nu"], - transport_data["virt_packet_last_interaction_in_nu"], - ) - assert np.array_equal( - f["transport/virt_packet_last_interaction_type"], - transport_data["virt_packet_last_interaction_type"], - ) - assert np.array_equal( - f["transport/virt_packet_last_line_interaction_in_id"], - transport_data["virt_packet_last_line_interaction_in_id"], - ) - assert np.array_equal( - f["transport/virt_packet_last_line_interaction_out_id"], - transport_data["virt_packet_last_line_interaction_out_id"], - ) - assert np.array_equal( - f["transport/virt_packet_last_line_interaction_shell_id"], - transport_data["virt_packet_last_line_interaction_shell_id"], - ) - assert np.array_equal( - f["transport/virt_packet_nus"], transport_data["virt_packet_nus"] - ) - assert np.array_equal( - f["transport/volume_cgs"], transport_data["volume"].value - ) diff --git a/tardis/model/geometry/radial1d.py b/tardis/model/geometry/radial1d.py index a130a80cb65..53f21fbeed3 100644 --- a/tardis/model/geometry/radial1d.py +++ b/tardis/model/geometry/radial1d.py @@ -188,6 +188,7 @@ def to_numba(self): ("r_outer", float64[:]), ("v_inner", float64[:]), ("v_outer", float64[:]), + ("volume", float64[:]), ] @@ -203,8 +204,10 @@ def __init__(self, r_inner, r_outer, v_inner, v_outer): r_outer : numpy.ndarray v_inner : numpy.ndarray v_outer : numpy.ndarray + volume : numpy.ndarray """ self.r_inner = r_inner self.r_outer = r_outer self.v_inner = v_inner self.v_outer = v_outer + self.volume = (4 / 3) * np.pi * (self.r_outer**3 - self.r_inner**3) diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py index 2dc8e76da99..061aecb3633 100644 --- a/tardis/model/parse_input.py +++ b/tardis/model/parse_input.py @@ -18,7 +18,10 @@ 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.montecarlo.packet_source import ( + BlackBodySimpleSource, + BlackBodySimpleSourceRelativistic, +) from tardis.util.base import quantity_linspace logger = logging.getLogger(__name__) @@ -572,9 +575,9 @@ def parse_radiation_field_state( return DiluteThermalRadiationFieldState(t_radiative, dilution_factor) -def parse_packet_source(config, geometry): +def initialize_packet_source(config, geometry, packet_source): """ - Parse the packet source based on the given configuration and geometry. + Initialize the packet source based on config and geometry Parameters ---------- @@ -582,36 +585,64 @@ def parse_packet_source(config, geometry): The configuration object containing the supernova and plasma settings. geometry : Geometry The geometry object containing the inner radius information. + packet_source : BasePacketSource + The packet source object based on the configuration and geometry. Returns ------- - packet_source : BlackBodySimpleSource + packet_source : BasePacketSource 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, - ) + packet_source.radius = geometry.r_inner[0] + packet_source.temperature = config.plasma.initial_t_inner + elif (config.plasma.initial_t_inner < 0.0 * u.K) and ( luminosity_requested is not None ): - packet_source = BlackBodySimpleSource(radius=geometry.r_inner[0]) + packet_source.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_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. + """ + if config.montecarlo.enable_full_relativity: + packet_source = BlackBodySimpleSourceRelativistic( + base_seed=config.montecarlo.seed, + time_explosion=config.supernova.time_explosion, + ) + else: + packet_source = BlackBodySimpleSource(base_seed=config.montecarlo.seed) + + return initialize_packet_source(config, geometry, packet_source) + + def parse_csvy_radiation_field_state( config, csvy_model_config, csvy_model_data, geometry, packet_source ): diff --git a/tardis/model/tests/test_base.py b/tardis/model/tests/test_base.py index c9a8a2f8819..3aedf574420 100644 --- a/tardis/model/tests/test_base.py +++ b/tardis/model/tests/test_base.py @@ -176,7 +176,9 @@ def setup(self, example_configuration_dir, atomic_dataset): ) def test_initial_temperature(self): - assert_almost_equal(self.simulation_state.t_inner.value, 2508) + assert_almost_equal( + self.simulation_state.packet_source.temperature.value, 2508 + ) class TestModelFromArtisDensityAbundancesAllAscii: diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py index 32fdc6fa88d..e2df4666d67 100644 --- a/tardis/montecarlo/base.py +++ b/tardis/montecarlo/base.py @@ -1,4 +1,3 @@ -import os import logging import warnings @@ -7,16 +6,14 @@ from numba import set_num_threads from numba import cuda -from scipy.special import zeta -import tardis -from tardis.montecarlo.spectrum import TARDISSpectrum from tardis.util.base import quantity_linspace from tardis.io.util import HDFWriterMixin from tardis.montecarlo import packet_source as source from tardis.montecarlo.montecarlo_numba.formal_integral import FormalIntegrator +from tardis.montecarlo.montecarlo_numba.estimators import initialize_estimators from tardis.montecarlo import montecarlo_configuration as mc_config_module - +from tardis.montecarlo.montecarlo_state import MonteCarloTransportState from tardis.montecarlo.montecarlo_numba import montecarlo_radial1d from tardis.montecarlo.montecarlo_numba.numba_interface import ( @@ -31,29 +28,19 @@ # TODO: refactor this into more parts -class MontecarloTransport(HDFWriterMixin): +class MonteCarloTransportSolver(HDFWriterMixin): """ This class is designed as an interface between the Python part and the montecarlo C-part """ hdf_properties = [ - "output_nu", - "output_energy", - "nu_bar_estimator", - "j_estimator", - "montecarlo_virtual_luminosity", + "transport_state", "last_interaction_in_nu", "last_interaction_type", "last_line_interaction_in_id", "last_line_interaction_out_id", "last_line_interaction_shell_id", - "packet_luminosity", - "spectrum", - "spectrum_virtual", - "spectrum_reabsorbed", - "time_of_simulation", - "emitted_packet_mask", ] vpacket_hdf_properties = [ @@ -69,16 +56,6 @@ class MontecarloTransport(HDFWriterMixin): ] hdf_name = "transport" - w_estimator_constant = ( - (const.c**2 / (2 * const.h)) - * (15 / np.pi**4) - * (const.h / const.k_B) ** 4 - / (4 * np.pi) - ).cgs.value - - t_rad_estimator_constant = ( - (np.pi**4 / (15 * 24 * zeta(5, 1))) * (const.h / const.k_B) - ).cgs.value def __init__( self, @@ -93,7 +70,7 @@ def __init__( v_packet_settings, spectrum_method, packet_source, - virtual_packet_logging, + enable_virtual_packet_logging, nthreads=1, debug_packets=False, logger_buffer=1, @@ -113,10 +90,12 @@ def __init__( self.v_packet_settings = v_packet_settings self.spectrum_method = spectrum_method self._integrator = None - self._spectrum_integrated = None + self.use_gpu = use_gpu - self.virt_logging = virtual_packet_logging + self.virt_logging = enable_virtual_packet_logging + + # Length 2 for initialization - will be removed in next PR self.virt_packet_last_interaction_type = np.ones(2) * -1 self.virt_packet_last_interaction_in_nu = np.ones(2) * -1.0 self.virt_packet_last_line_interaction_in_id = np.ones(2) * -1 @@ -144,84 +123,16 @@ def __init__( if self.spectrum_method == "integrated": self.optional_hdf_properties.append("spectrum_integrated") - def _initialize_estimator_arrays(self, tau_sobolev_shape): - """ - Initialize the output arrays of the montecarlo simulation. - - Parameters - ---------- - tau_sobolev_shape : tuple - tuple for the tau_sobolev_shape - """ - - # Estimators - self.j_estimator = np.zeros(tau_sobolev_shape[1], dtype=np.float64) - self.nu_bar_estimator = np.zeros(tau_sobolev_shape[1], dtype=np.float64) - self.j_blue_estimator = np.zeros(tau_sobolev_shape) - self.Edotlu_estimator = np.zeros(tau_sobolev_shape) - # TODO: this is the wrong attribute naming style. - - def _initialize_continuum_estimator_arrays(self, gamma_shape): - """ - Initialize the arrays for the MC estimators for continuum processes. - - Parameters - ---------- - gamma_shape : tuple - Shape of the array with the photoionization rate coefficients. - """ - self.photo_ion_estimator = np.zeros(gamma_shape, dtype=np.float64) - self.stim_recomb_estimator = np.zeros(gamma_shape, dtype=np.float64) - self.stim_recomb_cooling_estimator = np.zeros( - gamma_shape, dtype=np.float64 - ) - self.bf_heating_estimator = np.zeros(gamma_shape, dtype=np.float64) - - self.stim_recomb_cooling_estimator = np.zeros( - gamma_shape, dtype=np.float64 - ) - - self.photo_ion_estimator_statistics = np.zeros( - gamma_shape, dtype=np.int64 - ) - - def _initialize_geometry_arrays(self, simulation_state): - """ - Generate the cgs like geometry arrays for the montecarlo part - - Parameters - ---------- - model : model.SimulationState - """ - self.r_inner_cgs = simulation_state.r_inner.to("cm").value - self.r_outer_cgs = simulation_state.r_outer.to("cm").value - self.v_inner_cgs = simulation_state.v_inner.to("cm/s").value - self.v_outer_cgs = simulation_state.v_outer.to("cm/s").value - - def _initialize_packets(self, model, no_of_packets, iteration): + def _initialize_packets(self, no_of_packets, iteration): # the iteration (passed as seed_offset) is added each time to preserve randomness # across different simulations with the same temperature, # for example. - mc_config_module.packet_seeds = self.packet_source.create_packet_seeds( - no_of_packets, iteration - ) - - # Set latest state of packet source - self.packet_source.set_state_from_model(model) # Create packets - radii, nus, mus, energies = self.packet_source.create_packets( - no_of_packets + self.packet_collection = self.packet_source.create_packets( + no_of_packets, seed_offset=iteration ) - self.input_r = radii.to(u.cm).value - self.input_nu = nus - self.input_mu = mus - self.input_energy = energies - - self._output_nu = np.ones(no_of_packets, dtype=np.float64) * -99.0 - self._output_energy = np.ones(no_of_packets, dtype=np.float64) * -99.0 - self.last_line_interaction_in_id = -1 * np.ones( no_of_packets, dtype=np.int64 ) @@ -234,66 +145,6 @@ def _initialize_packets(self, model, no_of_packets, iteration): self.last_interaction_type = -1 * np.ones(no_of_packets, dtype=np.int64) self.last_interaction_in_nu = np.zeros(no_of_packets, dtype=np.float64) - self._montecarlo_virtual_luminosity = u.Quantity( - np.zeros_like(self.spectrum_frequency.value), "erg / s" - ) - - @property - def spectrum(self): - return TARDISSpectrum( - self.spectrum_frequency, self.montecarlo_emitted_luminosity - ) - - @property - def spectrum_reabsorbed(self): - return TARDISSpectrum( - self.spectrum_frequency, self.montecarlo_reabsorbed_luminosity - ) - - @property - def spectrum_virtual(self): - if np.all(self.montecarlo_virtual_luminosity == 0): - warnings.warn( - "MontecarloTransport.spectrum_virtual" - "is zero. Please run the montecarlo simulation with" - "no_of_virtual_packets > 0", - UserWarning, - ) - - return TARDISSpectrum( - self.spectrum_frequency, self.montecarlo_virtual_luminosity - ) - - @property - def spectrum_integrated(self): - if self._spectrum_integrated is None: - # This was changed from unpacking to specific attributes as compute - # is not used in calculate_spectrum - self._spectrum_integrated = self.integrator.calculate_spectrum( - self.spectrum_frequency[:-1], - points=self.integrator_settings.points, - interpolate_shells=self.integrator_settings.interpolate_shells, - ) - return self._spectrum_integrated - - @property - def integrator(self): - if self._integrator is None: - warnings.warn( - "MontecarloTransport.integrator: " - "The FormalIntegrator is not yet available." - "Please run the montecarlo simulation at least once.", - UserWarning, - ) - if self.enable_full_relativity: - raise NotImplementedError( - "The FormalIntegrator is not yet implemented for the full " - "relativity mode. " - "Please run with config option enable_full_relativity: " - "False." - ) - return self._integrator - def run( self, simulation_state, @@ -323,27 +174,31 @@ def run( set_num_threads(self.nthreads) - self.time_of_simulation = self.calculate_time_of_simulation( - simulation_state - ) - self.volume = simulation_state.volume - - # Initializing estimator array - self._initialize_estimator_arrays(plasma.tau_sobolevs.shape) - if not plasma.continuum_interaction_species.empty: gamma_shape = plasma.gamma.shape else: gamma_shape = (0, 0) - self._initialize_continuum_estimator_arrays(gamma_shape) + # Initializing estimator array + estimators = initialize_estimators( + plasma.tau_sobolevs.shape, gamma_shape + ) - self._initialize_geometry_arrays(simulation_state) + self._initialize_packets(no_of_packets, iteration) - self._initialize_packets( - simulation_state, - no_of_packets, - iteration, + self.transport_state = MonteCarloTransportState( + self.packet_collection, + estimators, + simulation_state.volume.cgs.copy(), + spectrum_frequency=self.spectrum_frequency, + geometry_state=simulation_state.geometry.to_numba(), + ) + self.transport_state.enable_full_relativity = ( + self.enable_full_relativity + ) + self.transport_state.integrator_settings = self.integrator_settings + self.transport_state._integrator = FormalIntegrator( + simulation_state, plasma, self ) configuration_initialize(self, no_of_virtual_packets) @@ -351,19 +206,19 @@ def run( simulation_state, plasma, iteration, - no_of_packets, + self.packet_collection, + self.transport_state.estimators, total_iterations, show_progress_bars, self, ) - self._integrator = FormalIntegrator(simulation_state, plasma, self) def legacy_return(self): return ( - self.output_nu, - self.output_energy, - self.j_estimator, - self.nu_bar_estimator, + self.transport_state.packet_collection.output_nus, + self.transport_state.packet_collection.output_energies, + self.transport_state.estimators.j_estimator, + self.transport_state.estimators.nu_bar_estimator, self.last_line_interaction_in_id, self.last_line_interaction_out_id, self.last_interaction_type, @@ -375,235 +230,9 @@ def get_line_interaction_id(self, line_interaction_type): line_interaction_type ) - @property - def output_nu(self): - return u.Quantity(self._output_nu, u.Hz) - - @property - def output_energy(self): - return u.Quantity(self._output_energy, u.erg) - - @property - def virtual_packet_nu(self): - try: - return u.Quantity(self.virt_packet_nus, u.Hz) - except AttributeError: - warnings.warn( - "MontecarloTransport.virtual_packet_nu:" - "Set 'virtual_packet_logging: True' in the configuration file" - "to access this property" - "It should be added under 'virtual' property of 'spectrum' property", - UserWarning, - ) - return None - - @property - def virtual_packet_energy(self): - try: - return u.Quantity(self.virt_packet_energies, u.erg) - except AttributeError: - warnings.warn( - "MontecarloTransport.virtual_packet_energy:" - "Set 'virtual_packet_logging: True' in the configuration file" - "to access this property" - "It should be added under 'virtual' property of 'spectrum' property", - UserWarning, - ) - return None - - @property - def virtual_packet_luminosity(self): - try: - return self.virtual_packet_energy / self.time_of_simulation - except TypeError: - warnings.warn( - "MontecarloTransport.virtual_packet_luminosity:" - "Set 'virtual_packet_logging: True' in the configuration file" - "to access this property" - "It should be added under 'virtual' property of 'spectrum' property", - UserWarning, - ) - return None - - @property - def packet_luminosity(self): - return self.output_energy / self.time_of_simulation - - @property - def emitted_packet_mask(self): - return self.output_energy >= 0 - - @property - def emitted_packet_nu(self): - return self.output_nu[self.emitted_packet_mask] - - @property - def reabsorbed_packet_nu(self): - return self.output_nu[~self.emitted_packet_mask] - - @property - def emitted_packet_luminosity(self): - return self.packet_luminosity[self.emitted_packet_mask] - - @property - def reabsorbed_packet_luminosity(self): - return -self.packet_luminosity[~self.emitted_packet_mask] - - @property - def montecarlo_reabsorbed_luminosity(self): - return u.Quantity( - np.histogram( - self.reabsorbed_packet_nu, - weights=self.reabsorbed_packet_luminosity, - bins=self.spectrum_frequency, - )[0], - "erg / s", - ) - - @property - def montecarlo_emitted_luminosity(self): - return u.Quantity( - np.histogram( - self.emitted_packet_nu, - weights=self.emitted_packet_luminosity, - bins=self.spectrum_frequency, - )[0], - "erg / s", - ) - - @property - def montecarlo_virtual_luminosity(self): - return ( - self._montecarlo_virtual_luminosity[:-1] - / self.time_of_simulation.value - ) - - def calculate_emitted_luminosity( - self, luminosity_nu_start, luminosity_nu_end - ): - """ - Calculate emitted luminosity. - - Parameters - ---------- - luminosity_nu_start : astropy.units.Quantity - luminosity_nu_end : astropy.units.Quantity - - Returns - ------- - astropy.units.Quantity - """ - luminosity_wavelength_filter = ( - self.emitted_packet_nu > luminosity_nu_start - ) & (self.emitted_packet_nu < luminosity_nu_end) - - emitted_luminosity = self.emitted_packet_luminosity[ - luminosity_wavelength_filter - ].sum() - - return emitted_luminosity - - def calculate_reabsorbed_luminosity( - self, luminosity_nu_start, luminosity_nu_end - ): - """ - Calculate reabsorbed luminosity. - - Parameters - ---------- - luminosity_nu_start : astropy.units.Quantity - luminosity_nu_end : astropy.units.Quantity - - Returns - ------- - astropy.units.Quantity - """ - luminosity_wavelength_filter = ( - self.reabsorbed_packet_nu > luminosity_nu_start - ) & (self.reabsorbed_packet_nu < luminosity_nu_end) - - reabsorbed_luminosity = self.reabsorbed_packet_luminosity[ - luminosity_wavelength_filter - ].sum() - - return reabsorbed_luminosity - - def calculate_radiationfield_properties(self): - """ - Calculate an updated radiation field from the :math: - `\\bar{nu}_\\textrm{estimator}` and :math:`\\J_\\textrm{estimator}` - calculated in the montecarlo simulation. - The details of the calculation can be found in the documentation. - - Parameters - ---------- - nubar_estimator : np.ndarray (float) - j_estimator : np.ndarray (float) - - Returns - ------- - t_rad : astropy.units.Quantity (float) - w : numpy.ndarray (float) - """ - - t_rad = ( - self.t_rad_estimator_constant - * self.nu_bar_estimator - / self.j_estimator - ) - w = self.j_estimator / ( - 4 - * const.sigma_sb.cgs.value - * t_rad**4 - * self.time_of_simulation.value - * self.volume.value - ) - - return t_rad * u.K, w - - def calculate_luminosity_inner(self, model): - """ - Calculate inner luminosity. - - Parameters - ---------- - model : model.SimulationState - - Returns - ------- - astropy.units.Quantity - """ - return ( - 4 - * np.pi - * const.sigma_sb.cgs - * model.r_inner[0] ** 2 - * model.t_inner**4 - ).to("erg/s") - - def calculate_time_of_simulation(self, model): - """ - Calculate time of montecarlo simulation. - - Parameters - ---------- - model : model.SimulationState - - Returns - ------- - float - """ - return 1.0 * u.erg / self.calculate_luminosity_inner(model) - - def calculate_f_nu(self, frequency): - pass - - def calculate_f_lambda(self, wavelength): - pass - @classmethod def from_config( - cls, config, packet_source=None, virtual_packet_logging=False + cls, config, packet_source, enable_virtual_packet_logging=False ): """ Create a new MontecarloTransport instance from a Configuration object. @@ -646,10 +275,7 @@ def from_config( but no CUDA GPU is available.""" ) elif running_mode == "AUTOMATIC": - if cuda.is_available(): - use_gpu = True - else: - use_gpu = False + use_gpu = bool(cuda.is_available()) elif running_mode == "CPU": use_gpu = False else: @@ -666,16 +292,6 @@ def from_config( config.montecarlo.tracking.initial_array_length ) - if packet_source is None: - if not config.montecarlo.enable_full_relativity: - packet_source = source.BlackBodySimpleSource( - base_seed=config.montecarlo.seed - ) - else: - packet_source = source.BlackBodySimpleSourceRelativistic( - base_seed=config.montecarlo.seed - ) - return cls( spectrum_frequency=spectrum_frequency, virtual_spectrum_spawn_range=config.montecarlo.virtual_spectrum_spawn_range, @@ -690,9 +306,9 @@ def from_config( packet_source=packet_source, debug_packets=config.montecarlo.debug_packets, logger_buffer=config.montecarlo.logger_buffer, - virtual_packet_logging=( + enable_virtual_packet_logging=( config.spectrum.virtual.virtual_packet_logging - | virtual_packet_logging + | enable_virtual_packet_logging ), nthreads=config.montecarlo.nthreads, tracking_rpacket=config.montecarlo.tracking.track_rpacket, diff --git a/tardis/montecarlo/montecarlo_numba/__init__.py b/tardis/montecarlo/montecarlo_numba/__init__.py index a4eb6dc3c3a..b7f1bb07f19 100644 --- a/tardis/montecarlo/montecarlo_numba/__init__.py +++ b/tardis/montecarlo/montecarlo_numba/__init__.py @@ -14,4 +14,6 @@ from tardis.montecarlo.montecarlo_numba.r_packet import RPacket from tardis.montecarlo.montecarlo_numba.base import montecarlo_radial1d -from tardis.montecarlo.montecarlo_numba.numba_interface import PacketCollection +from tardis.montecarlo.montecarlo_numba.packet_collections import ( + PacketCollection, +) diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py index 13f32f32d96..f69f3546654 100644 --- a/tardis/montecarlo/montecarlo_numba/base.py +++ b/tardis/montecarlo/montecarlo_numba/base.py @@ -2,6 +2,11 @@ from numba.np.ufunc.parallel import get_thread_id, get_num_threads import numpy as np +from tardis.montecarlo.montecarlo_numba.estimators import Estimators +from tardis.montecarlo.montecarlo_numba.packet_collections import ( + VPacketCollection, +) + from tardis.montecarlo.montecarlo_numba.r_packet import ( RPacket, @@ -9,12 +14,9 @@ ) from tardis.montecarlo.montecarlo_numba.numba_interface import ( - PacketCollection, - VPacketCollection, RPacketTracker, NumbaModel, opacity_state_initialize, - Estimators, ) from tardis.montecarlo import ( @@ -37,38 +39,19 @@ def montecarlo_radial1d( simulation_state, plasma, iteration, - no_of_packets, + packet_collection, + estimators, total_iterations, show_progress_bars, transport, ): - packet_collection = PacketCollection( - transport.input_r, - transport.input_nu, - transport.input_mu, - transport.input_energy, - transport._output_nu, - transport._output_energy, - ) - numba_radial_1d_geometry = simulation_state.geometry.to_numba() + numba_radial_1d_geometry = transport.transport_state.geometry_state numba_model = NumbaModel( simulation_state.time_explosion.to("s").value, ) opacity_state = opacity_state_initialize( plasma, transport.line_interaction_type ) - estimators = Estimators( - transport.j_estimator, - transport.nu_bar_estimator, - transport.j_blue_estimator, - transport.Edotlu_estimator, - transport.photo_ion_estimator, - transport.stim_recomb_estimator, - transport.bf_heating_estimator, - transport.stim_recomb_cooling_estimator, - transport.photo_ion_estimator_statistics, - ) - packet_seeds = montecarlo_configuration.packet_seeds number_of_vpackets = montecarlo_configuration.number_of_vpackets ( @@ -96,14 +79,15 @@ def montecarlo_radial1d( estimators, transport.spectrum_frequency.value, number_of_vpackets, - packet_seeds, montecarlo_configuration.VPACKET_LOGGING, iteration=iteration, show_progress_bars=show_progress_bars, - no_of_packets=no_of_packets, total_iterations=total_iterations, ) - transport._montecarlo_virtual_luminosity.value[:] = v_packets_energy_hist + + transport.transport_state._montecarlo_virtual_luminosity.value[ + : + ] = v_packets_energy_hist transport.last_interaction_type = last_interaction_type transport.last_interaction_in_nu = last_interaction_in_nu transport.last_line_interaction_in_id = last_line_interaction_in_id @@ -152,11 +136,9 @@ def montecarlo_main_loop( estimators, spectrum_frequency, number_of_vpackets, - packet_seeds, virtual_packet_logging, iteration, show_progress_bars, - no_of_packets, total_iterations, ): """ @@ -178,22 +160,15 @@ def montecarlo_main_loop( virtual_packet_logging : bool Option to enable virtual packet logging. """ - output_nus = np.empty_like(packet_collection.packets_input_nu) - last_interaction_types = ( - np.ones_like(packet_collection.packets_output_nu, dtype=np.int64) * -1 - ) - output_energies = np.empty_like(packet_collection.packets_output_nu) + no_of_packets = len(packet_collection.initial_nus) + output_nus = np.empty(no_of_packets, dtype=np.float64) + output_energies = np.empty(no_of_packets, dtype=np.float64) - last_interaction_in_nus = np.empty_like(packet_collection.packets_output_nu) - last_line_interaction_in_ids = ( - np.ones_like(packet_collection.packets_output_nu, dtype=np.int64) * -1 - ) - last_line_interaction_out_ids = ( - np.ones_like(packet_collection.packets_output_nu, dtype=np.int64) * -1 - ) - last_line_interaction_shell_ids = ( - np.ones_like(packet_collection.packets_output_nu, dtype=np.int64) * -1 - ) + last_interaction_in_nus = np.empty(no_of_packets, dtype=np.float64) + last_interaction_types = -1 * np.ones(no_of_packets, dtype=np.int64) + last_line_interaction_in_ids = -np.ones(no_of_packets, dtype=np.int64) + last_line_interaction_out_ids = -np.ones(no_of_packets, dtype=np.int64) + last_line_interaction_shell_ids = -np.ones(no_of_packets, dtype=np.int64) v_packets_energy_hist = np.zeros_like(spectrum_frequency) delta_nu = spectrum_frequency[1] - spectrum_frequency[0] @@ -259,16 +234,15 @@ def montecarlo_main_loop( total_iterations=total_iterations, ) - seed = packet_seeds[i] - np.random.seed(seed) r_packet = RPacket( - packet_collection.packets_input_radius[i], - packet_collection.packets_input_mu[i], - packet_collection.packets_input_nu[i], - packet_collection.packets_input_energy[i], - seed, + packet_collection.initial_radii[i], + packet_collection.initial_mus[i], + packet_collection.initial_nus[i], + packet_collection.initial_energies[i], + packet_collection.packet_seeds[i], i, ) + np.random.seed(r_packet.seed) local_estimators = estimator_list[tid] vpacket_collection = vpacket_collections[i] rpacket_tracker = rpacket_trackers[i] @@ -381,8 +355,8 @@ def montecarlo_main_loop( for rpacket_tracker in rpacket_trackers: rpacket_tracker.finalize_array() - packet_collection.packets_output_energy[:] = output_energies[:] - packet_collection.packets_output_nu[:] = output_nus[:] + packet_collection.output_energies[:] = output_energies[:] + packet_collection.output_nus[:] = output_nus[:] return ( v_packets_energy_hist, last_interaction_types, diff --git a/tardis/montecarlo/montecarlo_numba/estimators.py b/tardis/montecarlo/montecarlo_numba/estimators.py index feb06107658..3bae2ec640a 100644 --- a/tardis/montecarlo/montecarlo_numba/estimators.py +++ b/tardis/montecarlo/montecarlo_numba/estimators.py @@ -1,5 +1,9 @@ from math import exp -from numba import njit + +import numpy as np + +from numba import njit, float64, int64 +from numba.experimental import jitclass from tardis.montecarlo.montecarlo_numba import numba_config as nc from tardis.montecarlo.montecarlo_numba.numba_config import H, KB @@ -14,15 +18,123 @@ ) +def initialize_estimators(tau_sobolev_shape, gamma_shape): + """ + Initializes the estimators used in the Monte Carlo simulation. + + Parameters + ---------- + tau_sobolev_shape : tuple + Shape of the array with the Sobolev optical depth. + gamma_shape : tuple + Shape of the array with the photoionization rate coefficients. + + Returns + ------- + Estimators + The initialized estimators. + + Examples + -------- + >>> tau_sobolev_shape = (10, 20) + >>> gamma_shape = (5, 5) + >>> initialize_estimators(tau_sobolev_shape, gamma_shape) + + """ + + j_estimator = np.zeros(tau_sobolev_shape[1], dtype=np.float64) + nu_bar_estimator = np.zeros(tau_sobolev_shape[1], dtype=np.float64) + j_blue_estimator = np.zeros(tau_sobolev_shape) + Edotlu_estimator = np.zeros(tau_sobolev_shape) + + photo_ion_estimator = np.zeros(gamma_shape, dtype=np.float64) + stim_recomb_estimator = np.zeros(gamma_shape, dtype=np.float64) + stim_recomb_cooling_estimator = np.zeros(gamma_shape, dtype=np.float64) + bf_heating_estimator = np.zeros(gamma_shape, dtype=np.float64) + + stim_recomb_cooling_estimator = np.zeros(gamma_shape, dtype=np.float64) + + photo_ion_estimator_statistics = np.zeros(gamma_shape, dtype=np.int64) + return Estimators( + j_estimator, + nu_bar_estimator, + j_blue_estimator, + Edotlu_estimator, + photo_ion_estimator, + stim_recomb_estimator, + bf_heating_estimator, + stim_recomb_cooling_estimator, + photo_ion_estimator_statistics, + ) + + +base_estimators_spec = [ + ("j_estimator", float64[:]), + ("nu_bar_estimator", float64[:]), + ("j_blue_estimator", float64[:, :]), + ("Edotlu_estimator", float64[:, :]), +] + +continuum_estimators_spec = [ + ("photo_ion_estimator", float64[:, :]), + ("stim_recomb_estimator", float64[:, :]), + ("bf_heating_estimator", float64[:, :]), + ("stim_recomb_cooling_estimator", float64[:, :]), + ("photo_ion_estimator_statistics", int64[:, :]), +] + + +@jitclass(base_estimators_spec + continuum_estimators_spec) +class Estimators(object): + def __init__( + self, + j_estimator, + nu_bar_estimator, + j_blue_estimator, + Edotlu_estimator, + photo_ion_estimator, + stim_recomb_estimator, + bf_heating_estimator, + stim_recomb_cooling_estimator, + photo_ion_estimator_statistics, + ): + self.j_estimator = j_estimator + self.nu_bar_estimator = nu_bar_estimator + self.j_blue_estimator = j_blue_estimator + self.Edotlu_estimator = Edotlu_estimator + self.photo_ion_estimator = photo_ion_estimator + self.stim_recomb_estimator = stim_recomb_estimator + self.bf_heating_estimator = bf_heating_estimator + self.stim_recomb_cooling_estimator = stim_recomb_cooling_estimator + self.photo_ion_estimator_statistics = photo_ion_estimator_statistics + + def increment(self, other): + self.j_estimator += other.j_estimator + self.nu_bar_estimator += other.nu_bar_estimator + self.j_blue_estimator += other.j_blue_estimator + self.Edotlu_estimator += other.Edotlu_estimator + self.photo_ion_estimator += other.photo_ion_estimator + self.stim_recomb_estimator += other.stim_recomb_estimator + self.bf_heating_estimator += other.bf_heating_estimator + self.stim_recomb_cooling_estimator += ( + other.stim_recomb_cooling_estimator + ) + self.photo_ion_estimator_statistics += ( + other.photo_ion_estimator_statistics + ) + + @njit(**njit_dict_no_parallel) -def set_estimators(r_packet, distance, numba_estimator, comov_nu, comov_energy): +def update_base_estimators( + r_packet, distance, estimator_state, comov_nu, comov_energy +): """ Updating the estimators """ - numba_estimator.j_estimator[r_packet.current_shell_id] += ( + estimator_state.j_estimator[r_packet.current_shell_id] += ( comov_energy * distance ) - numba_estimator.nu_bar_estimator[r_packet.current_shell_id] += ( + estimator_state.nu_bar_estimator[r_packet.current_shell_id] += ( comov_energy * distance * comov_nu ) @@ -33,7 +145,7 @@ def update_bound_free_estimators( comov_energy, shell_id, distance, - numba_estimator, + estimator_state, t_electron, x_sect_bfs, current_continua, @@ -65,13 +177,13 @@ def update_bound_free_estimators( photo_ion_rate_estimator_increment = ( comov_energy * distance * x_sect_bfs[i] / comov_nu ) - numba_estimator.photo_ion_estimator[ + estimator_state.photo_ion_estimator[ current_continuum, shell_id ] += photo_ion_rate_estimator_increment - numba_estimator.stim_recomb_estimator[current_continuum, shell_id] += ( + estimator_state.stim_recomb_estimator[current_continuum, shell_id] += ( photo_ion_rate_estimator_increment * boltzmann_factor ) - numba_estimator.photo_ion_estimator_statistics[ + estimator_state.photo_ion_estimator_statistics[ current_continuum, shell_id ] += 1 @@ -79,10 +191,10 @@ def update_bound_free_estimators( bf_heating_estimator_increment = ( comov_energy * distance * x_sect_bfs[i] * (1 - nu_th / comov_nu) ) - numba_estimator.bf_heating_estimator[ + estimator_state.bf_heating_estimator[ current_continuum, shell_id ] += bf_heating_estimator_increment - numba_estimator.stim_recomb_cooling_estimator[ + estimator_state.stim_recomb_cooling_estimator[ current_continuum, shell_id ] += (bf_heating_estimator_increment * boltzmann_factor) diff --git a/tardis/montecarlo/montecarlo_numba/formal_integral.py b/tardis/montecarlo/montecarlo_numba/formal_integral.py index bdf170a4065..c1339989865 100644 --- a/tardis/montecarlo/montecarlo_numba/formal_integral.py +++ b/tardis/montecarlo/montecarlo_numba/formal_integral.py @@ -1,4 +1,3 @@ -import sys import warnings import numpy as np import pandas as pd @@ -9,7 +8,7 @@ from tardis import constants as const from numba import njit, char, float64, int64, typeof, byte, prange from numba.experimental import jitclass -import pdb + from tardis.montecarlo.montecarlo_numba.numba_config import SIGMA_THOMSON from tardis.montecarlo import montecarlo_configuration as mc_config_module @@ -279,8 +278,8 @@ class FormalIntegrator(object): points : int64 """ - def __init__(self, model, plasma, transport, points=1000): - self.model = model + def __init__(self, simulation_state, plasma, transport, points=1000): + self.simulation_state = simulation_state self.transport = transport self.points = points if plasma: @@ -299,11 +298,13 @@ def generate_numba_objects(self): self.numba_radial_1d_geometry = NumbaRadial1DGeometry( self.transport.r_inner_i, self.transport.r_outer_i, - self.transport.r_inner_i / self.model.time_explosion.to("s").value, - self.transport.r_outer_i / self.model.time_explosion.to("s").value, + self.transport.r_inner_i + / self.simulation_state.time_explosion.to("s").value, + self.transport.r_outer_i + / self.simulation_state.time_explosion.to("s").value, ) self.numba_model = NumbaModel( - self.model.time_explosion.cgs.value, + self.simulation_state.time_explosion.cgs.value, ) self.opacity_state = opacity_state_initialize( self.original_plasma, self.transport.line_interaction_type @@ -340,7 +341,7 @@ def raise_or_return(message): warnings.warn(message) return False - for obj in (self.model, self.plasma, self.transport): + for obj in (self.simulation_state, self.plasma, self.transport): if obj is None: return raise_or_return( "The integrator is missing either model, plasma or " @@ -375,7 +376,7 @@ def calculate_spectrum( self.check(raises) N = points or self.points if interpolate_shells == 0: # Default Value - interpolate_shells = max(2 * self.model.no_of_shells, 80) + interpolate_shells = max(2 * self.simulation_state.no_of_shells, 80) warnings.warn( "The number of interpolate_shells was not " f"specified. The value was set to {interpolate_shells}." @@ -416,8 +417,9 @@ def make_source_function(self): Numpy array containing ( 1 - exp(-tau_ul) ) S_ul ordered by wavelength of the transition u -> l """ - simulation_state = self.model + simulation_state = self.simulation_state transport = self.transport + mct_state = transport.transport_state # macro_ref = self.atomic_data.macro_atom_references macro_ref = self.atomic_data.macro_atom_references @@ -438,10 +440,13 @@ def make_source_function(self): destination_level_idx = ma_int_data.destination_level_idx.values Edotlu_norm_factor = 1 / ( - transport.time_of_simulation * simulation_state.volume + mct_state.packet_collection.time_of_simulation + * simulation_state.volume ) exptau = 1 - np.exp(-self.original_plasma.tau_sobolevs) - Edotlu = Edotlu_norm_factor * exptau * transport.Edotlu_estimator + Edotlu = ( + Edotlu_norm_factor * exptau * mct_state.estimators.Edotlu_estimator + ) # The following may be achieved by calling the appropriate plasma # functions @@ -452,7 +457,7 @@ def make_source_function(self): / ( 4 * np.pi - * transport.time_of_simulation + * mct_state.time_of_simulation * simulation_state.volume ) ) @@ -461,7 +466,10 @@ def make_source_function(self): ) # Jbluelu should already by in the correct order, i.e. by wavelength of # the transition l->u - Jbluelu = transport.j_blue_estimator * Jbluelu_norm_factor + Jbluelu = ( + transport.transport_state.estimators.j_blue_estimator + * Jbluelu_norm_factor + ) upper_level_index = self.atomic_data.lines.index.droplevel( "level_number_lower" @@ -501,6 +509,7 @@ def make_source_function(self): ] q_ul = tmp.set_index(transitions_index) t = simulation_state.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 @@ -527,8 +536,8 @@ def make_source_function(self): att_S_ul, Jredlu, Jbluelu, e_dot_u ) else: - transport.r_inner_i = transport.r_inner_cgs - transport.r_outer_i = transport.r_outer_cgs + transport.r_inner_i = mct_state.geometry_state.r_inner + transport.r_outer_i = mct_state.geometry_state.r_outer transport.tau_sobolevs_integ = ( self.original_plasma.tau_sobolevs.values ) @@ -542,12 +551,17 @@ def interpolate_integrator_quantities( self, att_S_ul, Jredlu, Jbluelu, e_dot_u ): transport = self.transport + mct_state = transport.transport_state plasma = self.original_plasma nshells = self.interpolate_shells - r_middle = (transport.r_inner_cgs + transport.r_outer_cgs) / 2.0 + r_middle = ( + mct_state.geometry_state.r_inner + mct_state.geometry_state.r_outer + ) / 2.0 r_integ = np.linspace( - transport.r_inner_cgs[0], transport.r_outer_cgs[-1], nshells + mct_state.geometry_state.r_inner[0], + mct_state.geometry_state.r_outer[-1], + nshells, ) transport.r_inner_i = r_integ[:-1] transport.r_outer_i = r_integ[1:] @@ -601,7 +615,7 @@ def formal_integral(self, nu, N): self.generate_numba_objects() L, I_nu_p = self.integrator.formal_integral( - self.model.t_inner, + self.simulation_state.t_inner, nu, nu.shape[0], att_S_ul, diff --git a/tardis/montecarlo/montecarlo_numba/numba_interface.py b/tardis/montecarlo/montecarlo_numba/numba_interface.py index 50be58a6d8f..9b8e850bfa7 100644 --- a/tardis/montecarlo/montecarlo_numba/numba_interface.py +++ b/tardis/montecarlo/montecarlo_numba/numba_interface.py @@ -267,164 +267,6 @@ def opacity_state_initialize(plasma, line_interaction_type): ) -packet_collection_spec = [ - ("packets_input_radius", float64[:]), - ("packets_input_nu", float64[:]), - ("packets_input_mu", float64[:]), - ("packets_input_energy", float64[:]), - ("packets_output_nu", float64[:]), - ("packets_output_energy", float64[:]), -] - - -@jitclass(packet_collection_spec) -class PacketCollection(object): - def __init__( - self, - packets_input_radius, - packets_input_nu, - packets_input_mu, - packets_input_energy, - packets_output_nu, - packets_output_energy, - ): - self.packets_input_radius = packets_input_radius - self.packets_input_nu = packets_input_nu - self.packets_input_mu = packets_input_mu - self.packets_input_energy = packets_input_energy - self.packets_output_nu = packets_output_nu - self.packets_output_energy = packets_output_energy - - -vpacket_collection_spec = [ - ("rpacket_index", int64), - ("spectrum_frequency", float64[:]), - ("v_packet_spawn_start_frequency", float64), - ("v_packet_spawn_end_frequency", float64), - ("nus", float64[:]), - ("energies", float64[:]), - ("initial_mus", float64[:]), - ("initial_rs", float64[:]), - ("idx", int64), - ("number_of_vpackets", int64), - ("length", int64), - ("last_interaction_in_nu", float64[:]), - ("last_interaction_type", int64[:]), - ("last_interaction_in_id", int64[:]), - ("last_interaction_out_id", int64[:]), - ("last_interaction_shell_id", int64[:]), -] - - -@jitclass(vpacket_collection_spec) -class VPacketCollection(object): - def __init__( - self, - rpacket_index, - spectrum_frequency, - v_packet_spawn_start_frequency, - v_packet_spawn_end_frequency, - number_of_vpackets, - temporary_v_packet_bins, - ): - self.spectrum_frequency = spectrum_frequency - self.v_packet_spawn_start_frequency = v_packet_spawn_start_frequency - self.v_packet_spawn_end_frequency = v_packet_spawn_end_frequency - self.nus = np.empty(temporary_v_packet_bins, dtype=np.float64) - self.energies = np.empty(temporary_v_packet_bins, dtype=np.float64) - self.initial_mus = np.empty(temporary_v_packet_bins, dtype=np.float64) - self.initial_rs = np.empty(temporary_v_packet_bins, dtype=np.float64) - self.number_of_vpackets = number_of_vpackets - self.last_interaction_in_nu = np.zeros( - temporary_v_packet_bins, dtype=np.float64 - ) - self.last_interaction_type = -1 * np.ones( - temporary_v_packet_bins, dtype=np.int64 - ) - self.last_interaction_in_id = -1 * np.ones( - temporary_v_packet_bins, dtype=np.int64 - ) - self.last_interaction_out_id = -1 * np.ones( - temporary_v_packet_bins, dtype=np.int64 - ) - self.last_interaction_shell_id = -1 * np.ones( - temporary_v_packet_bins, dtype=np.int64 - ) - self.idx = 0 - self.rpacket_index = rpacket_index - self.length = temporary_v_packet_bins - - def set_properties( - self, - nu, - energy, - initial_mu, - initial_r, - last_interaction_in_nu, - last_interaction_type, - last_interaction_in_id, - last_interaction_out_id, - last_interaction_shell_id, - ): - if self.idx >= self.length: - temp_length = self.length * 2 + self.number_of_vpackets - temp_nus = np.empty(temp_length, dtype=np.float64) - temp_energies = np.empty(temp_length, dtype=np.float64) - temp_initial_mus = np.empty(temp_length, dtype=np.float64) - temp_initial_rs = np.empty(temp_length, dtype=np.float64) - temp_last_interaction_in_nu = np.empty( - temp_length, dtype=np.float64 - ) - temp_last_interaction_type = np.empty(temp_length, dtype=np.int64) - temp_last_interaction_in_id = np.empty(temp_length, dtype=np.int64) - temp_last_interaction_out_id = np.empty(temp_length, dtype=np.int64) - temp_last_interaction_shell_id = np.empty( - temp_length, dtype=np.int64 - ) - - temp_nus[: self.length] = self.nus - temp_energies[: self.length] = self.energies - temp_initial_mus[: self.length] = self.initial_mus - temp_initial_rs[: self.length] = self.initial_rs - temp_last_interaction_in_nu[ - : self.length - ] = self.last_interaction_in_nu - temp_last_interaction_type[ - : self.length - ] = self.last_interaction_type - temp_last_interaction_in_id[ - : self.length - ] = self.last_interaction_in_id - temp_last_interaction_out_id[ - : self.length - ] = self.last_interaction_out_id - temp_last_interaction_shell_id[ - : self.length - ] = self.last_interaction_shell_id - - self.nus = temp_nus - self.energies = temp_energies - self.initial_mus = temp_initial_mus - self.initial_rs = temp_initial_rs - self.last_interaction_in_nu = temp_last_interaction_in_nu - self.last_interaction_type = temp_last_interaction_type - self.last_interaction_in_id = temp_last_interaction_in_id - self.last_interaction_out_id = temp_last_interaction_out_id - self.last_interaction_shell_id = temp_last_interaction_shell_id - self.length = temp_length - - self.nus[self.idx] = nu - self.energies[self.idx] = energy - self.initial_mus[self.idx] = initial_mu - self.initial_rs[self.idx] = initial_r - self.last_interaction_in_nu[self.idx] = last_interaction_in_nu - self.last_interaction_type[self.idx] = last_interaction_type - self.last_interaction_in_id[self.idx] = last_interaction_in_id - self.last_interaction_out_id[self.idx] = last_interaction_out_id - self.last_interaction_shell_id[self.idx] = last_interaction_shell_id - self.idx += 1 - - rpacket_tracker_spec = [ ("length", int64), ("seed", int64), @@ -535,62 +377,6 @@ def finalize_array(self): self.interaction_type = self.interaction_type[: self.num_interactions] -base_estimators_spec = [ - ("j_estimator", float64[:]), - ("nu_bar_estimator", float64[:]), - ("j_blue_estimator", float64[:, :]), - ("Edotlu_estimator", float64[:, :]), -] - -continuum_estimators_spec = [ - ("photo_ion_estimator", float64[:, :]), - ("stim_recomb_estimator", float64[:, :]), - ("bf_heating_estimator", float64[:, :]), - ("stim_recomb_cooling_estimator", float64[:, :]), - ("photo_ion_estimator_statistics", int64[:, :]), -] - - -@jitclass(base_estimators_spec + continuum_estimators_spec) -class Estimators(object): - def __init__( - self, - j_estimator, - nu_bar_estimator, - j_blue_estimator, - Edotlu_estimator, - photo_ion_estimator, - stim_recomb_estimator, - bf_heating_estimator, - stim_recomb_cooling_estimator, - photo_ion_estimator_statistics, - ): - self.j_estimator = j_estimator - self.nu_bar_estimator = nu_bar_estimator - self.j_blue_estimator = j_blue_estimator - self.Edotlu_estimator = Edotlu_estimator - self.photo_ion_estimator = photo_ion_estimator - self.stim_recomb_estimator = stim_recomb_estimator - self.bf_heating_estimator = bf_heating_estimator - self.stim_recomb_cooling_estimator = stim_recomb_cooling_estimator - self.photo_ion_estimator_statistics = photo_ion_estimator_statistics - - def increment(self, other): - self.j_estimator += other.j_estimator - self.nu_bar_estimator += other.nu_bar_estimator - self.j_blue_estimator += other.j_blue_estimator - self.Edotlu_estimator += other.Edotlu_estimator - self.photo_ion_estimator += other.photo_ion_estimator - self.stim_recomb_estimator += other.stim_recomb_estimator - self.bf_heating_estimator += other.bf_heating_estimator - self.stim_recomb_cooling_estimator += ( - other.stim_recomb_cooling_estimator - ) - self.photo_ion_estimator_statistics += ( - other.photo_ion_estimator_statistics - ) - - def configuration_initialize(transport, number_of_vpackets): if transport.line_interaction_type == "macroatom": montecarlo_configuration.line_interaction_type = ( diff --git a/tardis/montecarlo/montecarlo_numba/packet_collections.py b/tardis/montecarlo/montecarlo_numba/packet_collections.py new file mode 100644 index 00000000000..244125686ed --- /dev/null +++ b/tardis/montecarlo/montecarlo_numba/packet_collections.py @@ -0,0 +1,170 @@ +import numpy as np +from numba.experimental import jitclass +from numba import float64, int64 + +packet_collection_spec = [ + ("initial_radii", float64[:]), + ("initial_nus", float64[:]), + ("initial_mus", float64[:]), + ("initial_energies", float64[:]), + ("packet_seeds", int64[:]), + ("time_of_simulation", float64), + ("radiation_field_luminosity", float64), # + ("output_nus", float64[:]), + ("output_energies", float64[:]), +] + + +@jitclass(packet_collection_spec) +class PacketCollection: + def __init__( + self, + initial_radii, + initial_nus, + initial_mus, + initial_energies, + packet_seeds, + radiation_field_luminosity, + ): + self.initial_radii = initial_radii + self.initial_nus = initial_nus + self.initial_mus = initial_mus + self.initial_energies = initial_energies + self.packet_seeds = packet_seeds + self.radiation_field_luminosity = radiation_field_luminosity + self.time_of_simulation = ( + 1 / radiation_field_luminosity + ) # 1 erg / luminosity + self.output_nus = np.ones_like(initial_radii, dtype=np.float64) * -99.0 + self.output_energies = ( + np.ones_like(initial_radii, dtype=np.float64) * -99.0 + ) + + +vpacket_collection_spec = [ + ("rpacket_index", int64), + ("spectrum_frequency", float64[:]), + ("v_packet_spawn_start_frequency", float64), + ("v_packet_spawn_end_frequency", float64), + ("nus", float64[:]), + ("energies", float64[:]), + ("initial_mus", float64[:]), + ("initial_rs", float64[:]), + ("idx", int64), + ("number_of_vpackets", int64), + ("length", int64), + ("last_interaction_in_nu", float64[:]), + ("last_interaction_type", int64[:]), + ("last_interaction_in_id", int64[:]), + ("last_interaction_out_id", int64[:]), + ("last_interaction_shell_id", int64[:]), +] + + +@jitclass(vpacket_collection_spec) +class VPacketCollection(object): + def __init__( + self, + rpacket_index, + spectrum_frequency, + v_packet_spawn_start_frequency, + v_packet_spawn_end_frequency, + number_of_vpackets, + temporary_v_packet_bins, + ): + self.spectrum_frequency = spectrum_frequency + self.v_packet_spawn_start_frequency = v_packet_spawn_start_frequency + self.v_packet_spawn_end_frequency = v_packet_spawn_end_frequency + self.nus = np.empty(temporary_v_packet_bins, dtype=np.float64) + self.energies = np.empty(temporary_v_packet_bins, dtype=np.float64) + self.initial_mus = np.empty(temporary_v_packet_bins, dtype=np.float64) + self.initial_rs = np.empty(temporary_v_packet_bins, dtype=np.float64) + self.number_of_vpackets = number_of_vpackets + self.last_interaction_in_nu = np.zeros( + temporary_v_packet_bins, dtype=np.float64 + ) + self.last_interaction_type = -1 * np.ones( + temporary_v_packet_bins, dtype=np.int64 + ) + self.last_interaction_in_id = -1 * np.ones( + temporary_v_packet_bins, dtype=np.int64 + ) + self.last_interaction_out_id = -1 * np.ones( + temporary_v_packet_bins, dtype=np.int64 + ) + self.last_interaction_shell_id = -1 * np.ones( + temporary_v_packet_bins, dtype=np.int64 + ) + self.idx = 0 + self.rpacket_index = rpacket_index + self.length = temporary_v_packet_bins + + def set_properties( + self, + nu, + energy, + initial_mu, + initial_r, + last_interaction_in_nu, + last_interaction_type, + last_interaction_in_id, + last_interaction_out_id, + last_interaction_shell_id, + ): + if self.idx >= self.length: + temp_length = self.length * 2 + self.number_of_vpackets + temp_nus = np.empty(temp_length, dtype=np.float64) + temp_energies = np.empty(temp_length, dtype=np.float64) + temp_initial_mus = np.empty(temp_length, dtype=np.float64) + temp_initial_rs = np.empty(temp_length, dtype=np.float64) + temp_last_interaction_in_nu = np.empty( + temp_length, dtype=np.float64 + ) + temp_last_interaction_type = np.empty(temp_length, dtype=np.int64) + temp_last_interaction_in_id = np.empty(temp_length, dtype=np.int64) + temp_last_interaction_out_id = np.empty(temp_length, dtype=np.int64) + temp_last_interaction_shell_id = np.empty( + temp_length, dtype=np.int64 + ) + + temp_nus[: self.length] = self.nus + temp_energies[: self.length] = self.energies + temp_initial_mus[: self.length] = self.initial_mus + temp_initial_rs[: self.length] = self.initial_rs + temp_last_interaction_in_nu[ + : self.length + ] = self.last_interaction_in_nu + temp_last_interaction_type[ + : self.length + ] = self.last_interaction_type + temp_last_interaction_in_id[ + : self.length + ] = self.last_interaction_in_id + temp_last_interaction_out_id[ + : self.length + ] = self.last_interaction_out_id + temp_last_interaction_shell_id[ + : self.length + ] = self.last_interaction_shell_id + + self.nus = temp_nus + self.energies = temp_energies + self.initial_mus = temp_initial_mus + self.initial_rs = temp_initial_rs + self.last_interaction_in_nu = temp_last_interaction_in_nu + self.last_interaction_type = temp_last_interaction_type + self.last_interaction_in_id = temp_last_interaction_in_id + self.last_interaction_out_id = temp_last_interaction_out_id + self.last_interaction_shell_id = temp_last_interaction_shell_id + self.length = temp_length + + self.nus[self.idx] = nu + self.energies[self.idx] = energy + self.initial_mus[self.idx] = initial_mu + self.initial_rs[self.idx] = initial_r + self.last_interaction_in_nu[self.idx] = last_interaction_in_nu + self.last_interaction_type[self.idx] = last_interaction_type + self.last_interaction_in_id[self.idx] = last_interaction_in_id + self.last_interaction_out_id[self.idx] = last_interaction_out_id + self.last_interaction_shell_id[self.idx] = last_interaction_shell_id + self.idx += 1 diff --git a/tardis/montecarlo/montecarlo_numba/tests/conftest.py b/tardis/montecarlo/montecarlo_numba/tests/conftest.py index c1ded1e9055..996f3b0d1d4 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/conftest.py +++ b/tardis/montecarlo/montecarlo_numba/tests/conftest.py @@ -3,17 +3,19 @@ import pytest import numpy as np from numba import njit +from tardis.montecarlo.montecarlo_numba.estimators import Estimators +from tardis.montecarlo.montecarlo_numba.packet_collections import ( + VPacketCollection, +) from tardis.simulation import Simulation -from tardis.montecarlo.montecarlo_numba import RPacket, PacketCollection -from tardis.montecarlo.montecarlo_numba.numba_interface import Estimators +from tardis.montecarlo.montecarlo_numba import RPacket +from tardis.montecarlo.montecarlo_numba.estimators import Estimators from tardis.montecarlo.montecarlo_numba.numba_interface import ( opacity_state_initialize, NumbaModel, - Estimators, - VPacketCollection, ) @@ -94,24 +96,16 @@ def verysimple_3vpacket_collection(nb_simulation_verysimple): @pytest.fixture(scope="package") def verysimple_packet_collection(nb_simulation_verysimple): - transport = nb_simulation_verysimple.transport - return PacketCollection( - transport.input_r, - transport.input_nu, - transport.input_mu, - transport.input_energy, - transport._output_nu, - transport._output_energy, - ) + return nb_simulation_verysimple.transport.packet_collection @pytest.fixture(scope="function") def packet(verysimple_packet_collection): return RPacket( r=7.5e14, - nu=verysimple_packet_collection.packets_input_nu[0], - mu=verysimple_packet_collection.packets_input_mu[0], - energy=verysimple_packet_collection.packets_input_energy[0], + nu=verysimple_packet_collection.initial_nus[0], + mu=verysimple_packet_collection.initial_mus[0], + energy=verysimple_packet_collection.initial_energies[0], seed=1963, index=0, ) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_base.py b/tardis/montecarlo/montecarlo_numba/tests/test_base.py index 4c73232aa40..1523821f213 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_base.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_base.py @@ -59,20 +59,19 @@ def test_montecarlo_main_loop( "/simulation/transport/j_estimator" ] expected_hdf_store.close() - 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 + transport_state = montecarlo_main_loop_simulation.transport.transport_state + actual_energy = transport_state.packet_collection.output_energies + actual_nu = transport_state.packet_collection.output_nus + actual_nu_bar_estimator = transport_state.estimators.nu_bar_estimator + actual_j_estimator = transport_state.estimators.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) + npt.assert_allclose(actual_energy, expected_energy, rtol=1e-13) + npt.assert_allclose(actual_nu, expected_nu, rtol=1e-13) def test_montecarlo_main_loop_vpacket_log( @@ -111,12 +110,15 @@ def test_montecarlo_main_loop_vpacket_log( ] 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 + transport_state = transport.transport_state + + actual_energy = transport_state.packet_collection.output_energies + actual_nu = transport_state.packet_collection.output_nus + actual_nu_bar_estimator = transport_state.estimators.nu_bar_estimator + actual_j_estimator = transport_state.estimators.j_estimator actual_vpacket_log_nus = transport.virt_packet_nus actual_vpacket_log_energies = transport.virt_packet_energies + expected_hdf_store.close() # Compare npt.assert_allclose( @@ -128,10 +130,8 @@ def test_montecarlo_main_loop_vpacket_log( npt.assert_allclose( actual_j_estimator, expected_j_estimator, rtol=1e-12, atol=1e-15 ) - npt.assert_allclose( - actual_energy.value, expected_energy, rtol=1e-12, atol=1e-15 - ) - npt.assert_allclose(actual_nu.value, expected_nu, rtol=1e-12, atol=1e-15) + npt.assert_allclose(actual_energy, expected_energy, rtol=1e-12, atol=1e-15) + npt.assert_allclose(actual_nu, expected_nu, rtol=1e-12, atol=1e-15) npt.assert_allclose( actual_vpacket_log_nus, expected_vpacket_log_nus, rtol=1e-12, atol=1e-15 ) diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_cuda_formal_integral.py b/tardis/montecarlo/montecarlo_numba/tests/test_cuda_formal_integral.py index 3eab004525a..679b25491c5 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_cuda_formal_integral.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_cuda_formal_integral.py @@ -1,28 +1,17 @@ -import pytest import numpy as np -from tardis import constants as c -from astropy import units as u - -from copy import deepcopy import numpy.testing as ntest +import pytest from numba import cuda -from numba import njit -from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel - -import tardis.montecarlo.montecarlo_numba.formal_integral_cuda as formal_integral_cuda import tardis.montecarlo.montecarlo_numba.formal_integral as formal_integral_numba +import tardis.montecarlo.montecarlo_numba.formal_integral_cuda as formal_integral_cuda +from tardis import constants as c from tardis.model.geometry.radial1d import NumbaRadial1DGeometry -from tardis.montecarlo.montecarlo_numba.numba_interface import NumbaModel - - -from tardis.montecarlo.montecarlo_numba.formal_integral import FormalIntegrator from tardis.montecarlo.montecarlo_numba.formal_integral import ( + FormalIntegrator, NumbaFormalIntegrator, ) - -from tardis.montecarlo.base import MontecarloTransport - +from tardis.montecarlo.montecarlo_numba.numba_interface import NumbaModel # Test cases must also take into account use of a GPU to run. If there is no GPU then the test cases will fail. GPUs_available = cuda.is_available() @@ -359,11 +348,11 @@ def test_full_formal_integral( # The function calculate_spectrum sets this property, but in order to test the CUDA. # version it is done manually, as well as to speed up the test. formal_integrator_numba.interpolate_shells = max( - 2 * formal_integrator_numba.model.no_of_shells, 80 + 2 * formal_integrator_numba.simulation_state.no_of_shells, 80 ) formal_integrator_cuda.interpolate_shells = max( - 2 * formal_integrator_cuda.model.no_of_shells, 80 + 2 * formal_integrator_cuda.simulation_state.no_of_shells, 80 ) res_numba = formal_integrator_numba.make_source_function() @@ -389,7 +378,7 @@ def test_full_formal_integral( formal_integrator_cuda.generate_numba_objects() L_cuda = formal_integrator_cuda.integrator.formal_integral( - formal_integrator_cuda.model.t_inner, + formal_integrator_cuda.simulation_state.t_inner, sim.transport.spectrum.frequency, sim.transport.spectrum.frequency.shape[0], att_S_ul_cuda, @@ -401,7 +390,7 @@ def test_full_formal_integral( )[0] L_numba = formal_integrator_numba.integrator.formal_integral( - formal_integrator_numba.model.t_inner, + formal_integrator_numba.simulation_state.t_inner, sim.transport.spectrum.frequency, sim.transport.spectrum.frequency.shape[0], att_S_ul_numba, diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py index 891a213928d..bdf02a5dcbe 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_packet.py +++ b/tardis/montecarlo/montecarlo_numba/tests/test_packet.py @@ -1,5 +1,6 @@ import pytest import numpy as np +import tardis.montecarlo.montecarlo_numba.estimators import tardis.montecarlo.montecarlo_numba.r_packet as r_packet import tardis.transport.geometry.calculate_distances as calculate_distances @@ -44,7 +45,7 @@ def model(): @pytest.fixture(scope="function") def estimators(): - return numba_interface.Estimators( + return tardis.montecarlo.montecarlo_numba.estimators.Estimators( j_estimator=np.array([0.0, 0.0], dtype=np.float64), nu_bar_estimator=np.array([0.0, 0.0], dtype=np.float64), j_blue_estimator=np.array( @@ -226,7 +227,6 @@ def test_trace_packet( verysimple_estimators, set_seed_fixture, ): - set_seed_fixture(1963) packet.initialize_line_id(verysimple_opacity_state, verysimple_numba_model) distance, interaction_type, delta_shell = r_packet_transport.trace_packet( diff --git a/tardis/montecarlo/montecarlo_state.py b/tardis/montecarlo/montecarlo_state.py new file mode 100644 index 00000000000..82f7c47d4ca --- /dev/null +++ b/tardis/montecarlo/montecarlo_state.py @@ -0,0 +1,370 @@ +import warnings + +import numpy as np +from astropy import units as u +from scipy.special import zeta + +from tardis import constants as const +from tardis.io.util import HDFWriterMixin +from tardis.montecarlo.spectrum import TARDISSpectrum + +DILUTION_FACTOR_ESTIMATOR_CONSTANT = ( + (const.c**2 / (2 * const.h)) + * (15 / np.pi**4) + * (const.h / const.k_B) ** 4 + / (4 * np.pi) +).cgs.value + +T_RADIATIVE_ESTIMATOR_CONSTANT = ( + (np.pi**4 / (15 * 24 * zeta(5, 1))) * (const.h / const.k_B) +).cgs.value + + +class MonteCarloTransportState(HDFWriterMixin): + hdf_properties = [ + "output_nu", + "output_energy", + "nu_bar_estimator", + "j_estimator", + "montecarlo_virtual_luminosity", + "packet_luminosity", + "spectrum", + "spectrum_virtual", + "spectrum_reabsorbed", + "time_of_simulation", + "emitted_packet_mask", + ] + + hdf_name = "transport_state" + + def __init__( + self, + packet_collection, + estimators, + volume, + spectrum_frequency, + geometry_state, + ): + self.packet_collection = packet_collection + self.estimators = estimators + self.volume = volume + self.spectrum_frequency = spectrum_frequency + self._montecarlo_virtual_luminosity = u.Quantity( + np.zeros_like(self.spectrum_frequency.value), "erg / s" + ) + self._integrator = None + self.integrator_settings = None + self._spectrum_integrated = None + self.enable_full_relativity = False + self.geometry_state = geometry_state + + def calculate_radiationfield_properties(self): + """ + Calculate an updated radiation field from the :math: + `\\bar{nu}_\\textrm{estimator}` and :math:`\\J_\\textrm{estimator}` + calculated in the montecarlo simulation. + The details of the calculation can be found in the documentation. + + Parameters + ---------- + nubar_estimator : np.ndarray (float) + j_estimator : np.ndarray (float) + + Returns + ------- + t_radiative : astropy.units.Quantity (float) + dilution_factor : numpy.ndarray (float) + """ + estimated_t_radiative = ( + T_RADIATIVE_ESTIMATOR_CONSTANT + * self.estimators.nu_bar_estimator + / self.estimators.j_estimator + ) * u.K + dilution_factor = self.estimators.j_estimator / ( + 4 + * const.sigma_sb.cgs.value + * estimated_t_radiative.value**4 + * (self.packet_collection.time_of_simulation) + * self.geometry_state.volume + ) + + return estimated_t_radiative, dilution_factor + + @property + def output_nu(self): + return self.packet_collection.output_nus * u.Hz + + @property + def output_energy(self): + return self.packet_collection.output_energies * u.erg + + @property + def nu_bar_estimator(self): + return self.estimators.nu_bar_estimator + + @property + def j_estimator(self): + return self.estimators.j_estimator + + @property + def time_of_simulation(self): + return self.packet_collection.time_of_simulation * u.s + + @property + def packet_luminosity(self): + return ( + self.packet_collection.output_energies + * u.erg + / (self.packet_collection.time_of_simulation * u.s) + ) + + @property + def emitted_packet_mask(self): + return self.packet_collection.output_energies >= 0 + + @property + def emitted_packet_nu(self): + return ( + self.packet_collection.output_nus[self.emitted_packet_mask] * u.Hz + ) + + @property + def reabsorbed_packet_nu(self): + return ( + self.packet_collection.output_nus[~self.emitted_packet_mask] * u.Hz + ) + + @property + def emitted_packet_luminosity(self): + return self.packet_luminosity[self.emitted_packet_mask] + + @property + def reabsorbed_packet_luminosity(self): + return -self.packet_luminosity[~self.emitted_packet_mask] + + @property + def montecarlo_reabsorbed_luminosity(self): + return u.Quantity( + np.histogram( + self.reabsorbed_packet_nu, + weights=self.reabsorbed_packet_luminosity, + bins=self.spectrum_frequency, + )[0], + "erg / s", + ) + + @property + def montecarlo_emitted_luminosity(self): + return u.Quantity( + np.histogram( + self.emitted_packet_nu, + weights=self.emitted_packet_luminosity, + bins=self.spectrum_frequency, + )[0], + "erg / s", + ) + + @property + def montecarlo_virtual_luminosity(self): + return ( + self._montecarlo_virtual_luminosity[:-1] + / self.packet_collection.time_of_simulation + ) + + @property + def spectrum(self): + return TARDISSpectrum( + self.spectrum_frequency, self.montecarlo_emitted_luminosity + ) + + @property + def spectrum_reabsorbed(self): + return TARDISSpectrum( + self.spectrum_frequency, self.montecarlo_reabsorbed_luminosity + ) + + @property + def spectrum_virtual(self): + if np.all(self.montecarlo_virtual_luminosity == 0): + warnings.warn( + "MontecarloTransport.spectrum_virtual" + "is zero. Please run the montecarlo simulation with" + "no_of_virtual_packets > 0", + UserWarning, + ) + + return TARDISSpectrum( + self.spectrum_frequency, self.montecarlo_virtual_luminosity + ) + + @property + def spectrum_integrated(self): + if self._spectrum_integrated is None: + # This was changed from unpacking to specific attributes as compute + # is not used in calculate_spectrum + self._spectrum_integrated = self.integrator.calculate_spectrum( + self.spectrum_frequency[:-1], + points=self.integrator_settings.points, + interpolate_shells=self.integrator_settings.interpolate_shells, + ) + return self._spectrum_integrated + + @property + def integrator(self): + if self._integrator is None: + warnings.warn( + "MontecarloTransport.integrator: " + "The FormalIntegrator is not yet available." + "Please run the montecarlo simulation at least once.", + UserWarning, + ) + if self.enable_full_relativity: + raise NotImplementedError( + "The FormalIntegrator is not yet implemented for the full " + "relativity mode. " + "Please run with config option enable_full_relativity: " + "False." + ) + return self._integrator + + def calculate_emitted_luminosity( + self, luminosity_nu_start, luminosity_nu_end + ): + """ + Calculate emitted luminosity. + + Parameters + ---------- + luminosity_nu_start : astropy.units.Quantity + luminosity_nu_end : astropy.units.Quantity + + Returns + ------- + astropy.units.Quantity + """ + luminosity_wavelength_filter = ( + self.emitted_packet_nu > luminosity_nu_start + ) & (self.emitted_packet_nu < luminosity_nu_end) + + return self.emitted_packet_luminosity[ + luminosity_wavelength_filter + ].sum() + + def calculate_reabsorbed_luminosity( + self, luminosity_nu_start, luminosity_nu_end + ): + """ + Calculate reabsorbed luminosity. + + Parameters + ---------- + luminosity_nu_start : astropy.units.Quantity + luminosity_nu_end : astropy.units.Quantity + + Returns + ------- + astropy.units.Quantity + """ + luminosity_wavelength_filter = ( + self.reabsorbed_packet_nu > luminosity_nu_start + ) & (self.reabsorbed_packet_nu < luminosity_nu_end) + + return self.reabsorbed_packet_luminosity[ + luminosity_wavelength_filter + ].sum() + + @property + def virtual_packet_nu(self): + try: + return u.Quantity(self.virt_packet_nus, u.Hz) + except AttributeError: + warnings.warn( + "MontecarloTransport.virtual_packet_nu:" + "Set 'virtual_packet_logging: True' in the configuration file" + "to access this property" + "It should be added under 'virtual' property of 'spectrum' property", + UserWarning, + ) + return None + + @property + def virtual_packet_energy(self): + try: + return u.Quantity(self.virt_packet_energies, u.erg) + except AttributeError: + warnings.warn( + "MontecarloTransport.virtual_packet_energy:" + "Set 'virtual_packet_logging: True' in the configuration file" + "to access this property" + "It should be added under 'virtual' property of 'spectrum' property", + UserWarning, + ) + return None + + @property + def virtual_packet_luminosity(self): + try: + return self.virtual_packet_energy / ( + self.packet_collection.time_of_simulation * u.s + ) + except TypeError: + warnings.warn( + "MontecarloTransport.virtual_packet_luminosity:" + "Set 'virtual_packet_logging: True' in the configuration file" + "to access this property" + "It should be added under 'virtual' property of 'spectrum' property", + UserWarning, + ) + return None + + @property + def montecarlo_virtual_luminosity(self): + return ( + self._montecarlo_virtual_luminosity[:-1] + / self.packet_collection.time_of_simulation + ) + + @property + def virtual_packet_nu(self): + try: + return u.Quantity(self.virt_packet_nus, u.Hz) + except AttributeError: + warnings.warn( + "MontecarloTransport.virtual_packet_nu:" + "Set 'virtual_packet_logging: True' in the configuration file" + "to access this property" + "It should be added under 'virtual' property of 'spectrum' property", + UserWarning, + ) + return None + + @property + def virtual_packet_energy(self): + try: + return u.Quantity(self.virt_packet_energies, u.erg) + except AttributeError: + warnings.warn( + "MontecarloTransport.virtual_packet_energy:" + "Set 'virtual_packet_logging: True' in the configuration file" + "to access this property" + "It should be added under 'virtual' property of 'spectrum' property", + UserWarning, + ) + return None + + @property + def virtual_packet_luminosity(self): + try: + return ( + self.virtual_packet_energy + / self.packet_collection.time_of_simulation + ) + except TypeError: + warnings.warn( + "MontecarloTransport.virtual_packet_luminosity:" + "Set 'virtual_packet_logging: True' in the configuration file" + "to access this property" + "It should be added under 'virtual' property of 'spectrum' property", + UserWarning, + ) + return None diff --git a/tardis/montecarlo/packet_source.py b/tardis/montecarlo/packet_source.py index 0bd63ea1f90..3350f9ab7a1 100644 --- a/tardis/montecarlo/packet_source.py +++ b/tardis/montecarlo/packet_source.py @@ -6,6 +6,9 @@ from tardis.montecarlo import ( montecarlo_configuration as montecarlo_configuration, ) +from tardis.montecarlo.montecarlo_numba.packet_collections import ( + PacketCollection, +) from astropy import units as u @@ -40,20 +43,6 @@ def __init__(self, base_seed=None, legacy_second_seed=None): def _reseed(self, seed): self.rng = np.random.default_rng(seed=seed) - def create_packet_seeds(self, no_of_packets, seed_offset): - # the iteration (passed as seed_offset) is added each time to preserve randomness - # across different simulations with the same temperature, - # for example. We seed the random module instead of the numpy module - # because we call random.sample, which references a different internal - # state than in the numpy.random module. - self._reseed(self.base_seed + seed_offset) - seeds = self.rng.choice(self.MAX_SEED_VAL, no_of_packets, replace=True) - return seeds - - @abc.abstractmethod - def set_state_from_model(self, model): - pass - @abc.abstractmethod def create_packet_radii(self, no_of_packets, *args, **kwargs): pass @@ -70,7 +59,7 @@ def create_packet_mus(self, no_of_packets, *args, **kwargs): def create_packet_energies(self, no_of_packets, *args, **kwargs): pass - def create_packets(self, no_of_packets, *args, **kwargs): + def create_packets(self, no_of_packets, seed_offset=0, *args, **kwargs): """Generate packet properties as arrays Parameters @@ -89,12 +78,50 @@ def create_packets(self, no_of_packets, *args, **kwargs): array Packet energies """ + # the iteration (passed as seed_offset) is added each time to preserve randomness + # across different simulations with the same temperature, + # for example. We seed the random module instead of the numpy module + # because we call random.sample, which references a different internal + # state than in the numpy.random module. + self._reseed(self.base_seed + seed_offset) + packet_seeds = self.rng.choice( + self.MAX_SEED_VAL, no_of_packets, replace=True + ) + radii = self.create_packet_radii(no_of_packets, *args, **kwargs) nus = self.create_packet_nus(no_of_packets, *args, **kwargs) mus = self.create_packet_mus(no_of_packets, *args, **kwargs) energies = self.create_packet_energies(no_of_packets, *args, **kwargs) + # Check if all arrays have the same length + assert ( + len(radii) == len(nus) == len(mus) == len(energies) == no_of_packets + ) + radiation_field_luminosity = ( + self.calculate_radfield_luminosity().to(u.erg / u.s).value + ) + return PacketCollection( + radii, nus, mus, energies, packet_seeds, radiation_field_luminosity + ) - return radii, nus, mus, energies + def calculate_radfield_luminosity(self): + """ + Calculate inner luminosity. + + Parameters + ---------- + model : model.SimulationState + + Returns + ------- + astropy.units.Quantity + """ + return ( + 4 + * np.pi + * const.sigma_sb.cgs + * self.radius**2 + * self.temperature**4 + ).to("erg/s") class BlackBodySimpleSource(BasePacketSource): @@ -115,21 +142,19 @@ class BlackBodySimpleSource(BasePacketSource): """ @classmethod - def from_model(cls, model, *args, **kwargs): - return cls(model.r_inner[0], model.t_inner.value, *args, **kwargs) + def from_simulation_state(cls, simulation_state, *args, **kwargs): + return cls( + simulation_state.r_inner[0], + simulation_state.t_inner.value, + *args, + **kwargs, + ) def __init__(self, radius=None, temperature=None, **kwargs): self.radius = radius self.temperature = temperature super().__init__(**kwargs) - def set_state_from_model(self, simulation_state): - """ - Set state of packet source (correct state should be ensured before creating packets) - """ - self.radius = simulation_state.r_inner[0] - self.temperature = simulation_state.t_inner.value - def create_packets(self, no_of_packets, *args, **kwargs): if self.radius is None or self.temperature is None: raise ValueError("Black body Radius or Temperature isn't set") @@ -192,7 +217,12 @@ def create_packet_nus(self, no_of_packets, l_samples=1000): xis_prod = np.prod(xis[1:], 0) x = ne.evaluate("-log(xis_prod)/l") - return x * (const.k_B.cgs.value * self.temperature) / const.h.cgs.value + if isinstance(self.temperature, u.Quantity): + temperature = self.temperature.value + else: + temperature = self.temperature + + return x * (const.k_B.cgs.value * temperature) / const.h.cgs.value def create_packet_mus(self, no_of_packets): """ @@ -269,11 +299,11 @@ class BlackBodySimpleSourceRelativistic(BlackBodySimpleSource): """ @classmethod - def from_model(cls, model, *args, **kwargs): + def from_simulation_state(cls, simulation_state, *args, **kwargs): return cls( - model.time_explosion, - model.r_inner[0], - model.t_inner.value, + simulation_state.time_explosion, + simulation_state.r_inner[0], + simulation_state.t_inner.value, *args, **kwargs, ) @@ -282,13 +312,6 @@ def __init__(self, time_explosion=None, **kwargs): self.time_explosion = time_explosion super().__init__(**kwargs) - def set_state_from_model(self, model): - """ - Set state of packet source (correct state should be ensured before creating packets) - """ - self.time_explosion = model.time_explosion - super().set_state_from_model(model) - def create_packets(self, no_of_packets): """Generate relativistic black-body packet properties as arrays diff --git a/tardis/montecarlo/tests/test_base.py b/tardis/montecarlo/tests/test_base.py index 131f245223d..d7cab27c2e8 100644 --- a/tardis/montecarlo/tests/test_base.py +++ b/tardis/montecarlo/tests/test_base.py @@ -4,6 +4,7 @@ import pytest from astropy import units as u from numpy.testing import assert_almost_equal +from pathlib import Path ### # Save and Load @@ -15,20 +16,17 @@ def to_hdf_buffer(hdf_file_path, simulation_verysimple): simulation_verysimple.transport.to_hdf( hdf_file_path, name="transport", overwrite=True ) + simulation_verysimple.transport.transport_state.to_hdf( + hdf_file_path, name="transport_state", overwrite=True + ) transport_properties = [ - "output_nu", - "output_energy", - "nu_bar_estimator", - "j_estimator", - "montecarlo_virtual_luminosity", "last_interaction_in_nu", "last_interaction_type", "last_line_interaction_in_id", "last_line_interaction_out_id", "last_line_interaction_shell_id", - "packet_luminosity", ] @@ -37,6 +35,26 @@ def test_hdf_transport(hdf_file_path, simulation_verysimple, attr): actual = getattr(simulation_verysimple.transport, attr) if hasattr(actual, "cgs"): actual = actual.cgs.value - path = os.path.join("transport", attr) + path = f"transport/{attr}" + expected = pd.read_hdf(hdf_file_path, path) + assert_almost_equal(actual, expected.values) + + +transport_state_properties = [ + "output_nu", + "output_energy", + "nu_bar_estimator", + "j_estimator", + "montecarlo_virtual_luminosity", + "packet_luminosity", +] + + +@pytest.mark.parametrize("attr", transport_state_properties) +def test_hdf_transport_state(hdf_file_path, simulation_verysimple, attr): + actual = getattr(simulation_verysimple.transport.transport_state, attr) + if hasattr(actual, "cgs"): + actual = actual.cgs.value + path = f"transport_state/{attr}" expected = pd.read_hdf(hdf_file_path, path) assert_almost_equal(actual, expected.values) diff --git a/tardis/montecarlo/tests/test_montecarlo.py b/tardis/montecarlo/tests/test_montecarlo.py index 90ea731b290..c3cc904b6b3 100644 --- a/tardis/montecarlo/tests/test_montecarlo.py +++ b/tardis/montecarlo/tests/test_montecarlo.py @@ -8,7 +8,7 @@ import tardis.montecarlo.montecarlo_numba.utils as utils import tardis.montecarlo.montecarlo_configuration as mc from tardis import constants as const -from tardis.montecarlo.montecarlo_numba.numba_interface import Estimators +from tardis.montecarlo.montecarlo_numba.estimators import Estimators from tardis.montecarlo.montecarlo_numba.numba_interface import RPacketTracker diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index 88089363bf8..2cc940874de 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -14,8 +14,9 @@ from tardis.io.configuration.config_reader import ConfigurationError from tardis.io.util import HDFWriterMixin from tardis.model import SimulationState +from tardis.model.parse_input import initialize_packet_source from tardis.montecarlo import montecarlo_configuration as mc_config_module -from tardis.montecarlo.base import MontecarloTransport +from tardis.montecarlo.base import MonteCarloTransportSolver from tardis.montecarlo.montecarlo_numba.r_packet import ( rpacket_trackers_to_dataframe, ) @@ -199,8 +200,10 @@ def __init__( def estimate_t_inner( self, input_t_inner, luminosity_requested, t_inner_update_exponent=-0.5 ): - emitted_luminosity = self.transport.calculate_emitted_luminosity( - self.luminosity_nu_start, self.luminosity_nu_end + emitted_luminosity = ( + self.transport.transport_state.calculate_emitted_luminosity( + self.luminosity_nu_start, self.luminosity_nu_end + ) ) luminosity_ratios = ( @@ -281,7 +284,7 @@ def advance_state(self): ( estimated_t_rad, estimated_w, - ) = self.transport.calculate_radiationfield_properties() + ) = self.transport.transport_state.calculate_radiationfield_properties() estimated_t_inner = self.estimate_t_inner( self.simulation_state.t_inner, self.luminosity_requested, @@ -366,17 +369,19 @@ def advance_state(self): ) # A check to see if the plasma is set with JBluesDetailed, in which # case it needs some extra kwargs. + + estimators = self.transport.transport_state.estimators if "j_blue_estimator" in self.plasma.outputs_dict: update_properties.update( t_inner=next_t_inner, - j_blue_estimator=self.transport.j_blue_estimator, + j_blue_estimator=estimators.j_blue_estimator, ) if "gamma_estimator" in self.plasma.outputs_dict: update_properties.update( - gamma_estimator=self.transport.photo_ion_estimator, - alpha_stim_estimator=self.transport.stim_recomb_estimator, - bf_heating_coeff_estimator=self.transport.bf_heating_estimator, - stim_recomb_cooling_coeff_estimator=self.transport.stim_recomb_cooling_estimator, + 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.update(**update_properties) @@ -396,15 +401,21 @@ def iterate(self, no_of_packets, no_of_virtual_packets=0): total_iterations=self.iterations, show_progress_bars=self.show_progress_bars, ) - output_energy = self.transport.output_energy + 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.transport.calculate_emitted_luminosity( - self.luminosity_nu_start, self.luminosity_nu_end + emitted_luminosity = ( + self.transport.transport_state.calculate_emitted_luminosity( + self.luminosity_nu_start, self.luminosity_nu_end + ) ) - reabsorbed_luminosity = self.transport.calculate_reabsorbed_luminosity( - self.luminosity_nu_start, self.luminosity_nu_end + reabsorbed_luminosity = ( + self.transport.transport_state.calculate_reabsorbed_luminosity( + self.luminosity_nu_start, self.luminosity_nu_end + ) ) if hasattr(self, "convergence_plots"): self.convergence_plots.fetch_data( @@ -681,6 +692,10 @@ def from_config( simulation_state = SimulationState.from_config( config, atom_data=atom_data ) + if packet_source is not None: + simulation_state.packet_source = initialize_packet_source( + config, simulation_state.geometry, packet_source + ) if "plasma" in kwargs: plasma = kwargs["plasma"] else: @@ -696,10 +711,10 @@ def from_config( ) transport = kwargs["transport"] else: - transport = MontecarloTransport.from_config( + transport = MonteCarloTransportSolver.from_config( config, - packet_source=packet_source, - virtual_packet_logging=virtual_packet_logging, + packet_source=simulation_state.packet_source, + enable_virtual_packet_logging=virtual_packet_logging, ) convergence_plots_config_options = [ diff --git a/tardis/simulation/tests/test_simulation.py b/tardis/simulation/tests/test_simulation.py index 258787c612f..c18c7b1ebb3 100644 --- a/tardis/simulation/tests/test_simulation.py +++ b/tardis/simulation/tests/test_simulation.py @@ -88,11 +88,27 @@ def simulation_one_loop( ], ) def test_plasma_estimates(simulation_one_loop, refdata, name): - try: - actual = getattr(simulation_one_loop.transport, name) - except AttributeError: + if name in ["nu_bar_estimator", "j_estimator"]: + actual = getattr( + simulation_one_loop.transport.transport_state.estimators, name + ) + elif name in ["t_radiative", "dilution_factor"]: actual = getattr(simulation_one_loop.simulation_state, name) - if name in ["t_radiative", "output_nu", "output_energy"]: + elif name in ["output_nu", "output_energy"]: + OLD_TO_NEW_DICT = { + "output_nu": "output_nus", + "output_energy": "output_energies", + } + actual = getattr( + simulation_one_loop.transport.transport_state.packet_collection, + OLD_TO_NEW_DICT[name], + ) + else: + try: + actual = getattr(simulation_one_loop.transport, name) + except AttributeError: + actual = getattr(simulation_one_loop.simulation_state, name) + if name in ["t_radiative"]: # removing the quantitiness of the data actual = actual.value actual = pd.Series(actual) diff --git a/tardis/tests/test_tardis_full.py b/tardis/tests/test_tardis_full.py index 84278933668..55a57370707 100644 --- a/tardis/tests/test_tardis_full.py +++ b/tardis/tests/test_tardis_full.py @@ -77,12 +77,17 @@ def get_ref_data(key): def test_j_blue_estimators(self, transport, refdata): j_blue_estimator = refdata("j_blue_estimator").values - npt.assert_allclose(transport.j_blue_estimator, j_blue_estimator) + npt.assert_allclose( + transport.transport_state.estimators.j_blue_estimator, + j_blue_estimator, + ) def test_spectrum(self, transport, refdata): luminosity = u.Quantity(refdata("spectrum/luminosity"), "erg /s") - assert_quantity_allclose(transport.spectrum.luminosity, luminosity) + assert_quantity_allclose( + transport.transport_state.spectrum.luminosity, luminosity + ) def test_virtual_spectrum(self, transport, refdata): luminosity = u.Quantity( @@ -90,7 +95,7 @@ def test_virtual_spectrum(self, transport, refdata): ) assert_quantity_allclose( - transport.spectrum_virtual.luminosity, luminosity + transport.transport_state.spectrum_virtual.luminosity, luminosity ) def test_transport_properties(self, transport): diff --git a/tardis/tests/test_tardis_full_formal_integral.py b/tardis/tests/test_tardis_full_formal_integral.py index 62a4abb0a21..318ce1dcaa9 100644 --- a/tardis/tests/test_tardis_full_formal_integral.py +++ b/tardis/tests/test_tardis_full_formal_integral.py @@ -78,12 +78,17 @@ def get_ref_data(key): def test_j_blue_estimators(self, transport, refdata): j_blue_estimator = refdata("j_blue_estimator").values - npt.assert_allclose(transport.j_blue_estimator, j_blue_estimator) + npt.assert_allclose( + transport.transport_state.estimators.j_blue_estimator, + j_blue_estimator, + ) def test_spectrum(self, transport, refdata): luminosity = u.Quantity(refdata("spectrum/luminosity"), "erg /s") - assert_quantity_allclose(transport.spectrum.luminosity, luminosity) + assert_quantity_allclose( + transport.transport_state.spectrum.luminosity, luminosity + ) def test_spectrum_integrated(self, transport, refdata): luminosity = u.Quantity( @@ -91,5 +96,5 @@ def test_spectrum_integrated(self, transport, refdata): ) assert_quantity_allclose( - transport.spectrum_integrated.luminosity, luminosity + transport.transport_state.spectrum_integrated.luminosity, luminosity ) diff --git a/tardis/transport/r_packet_transport.py b/tardis/transport/r_packet_transport.py index 579e387866d..19763c67edb 100644 --- a/tardis/transport/r_packet_transport.py +++ b/tardis/transport/r_packet_transport.py @@ -10,7 +10,7 @@ ) from tardis.montecarlo.montecarlo_numba.estimators import ( update_line_estimators, - set_estimators, + update_base_estimators, ) from tardis.transport.frame_transformations import ( get_doppler_factor, @@ -212,7 +212,7 @@ def move_r_packet(r_packet, distance, time_explosion, numba_estimator): if nc.ENABLE_FULL_RELATIVITY: distance *= doppler_factor - set_estimators( + update_base_estimators( r_packet, distance, numba_estimator, comov_nu, comov_energy ) diff --git a/tardis/visualization/tools/sdec_plot.py b/tardis/visualization/tools/sdec_plot.py index 7509e95a5a1..c16c13a63a4 100644 --- a/tardis/visualization/tools/sdec_plot.py +++ b/tardis/visualization/tools/sdec_plot.py @@ -88,6 +88,7 @@ def __init__( Time of simulation, having unit of s (second) """ # Save packets properties in a dataframe for easier data manipulation + packet_nus = u.Quantity(packet_nus, u.Hz) self.packets_df = pd.DataFrame( { "nus": packet_nus, @@ -163,9 +164,12 @@ def from_simulation(cls, sim, packets_mode): lines_df = sim.plasma.atomic_data.lines.reset_index().set_index( "line_id" ) - r_inner = sim.simulation_state.r_inner - t_inner = sim.simulation_state.t_inner - time_of_simulation = sim.transport.time_of_simulation + transport_state = sim.transport.transport_state + r_inner = sim.simulation_state.geometry.r_inner + t_inner = sim.simulation_state.packet_source.temperature + time_of_simulation = ( + transport_state.packet_collection.time_of_simulation * u.s + ) if packets_mode == "virtual": return cls( @@ -179,10 +183,10 @@ def from_simulation(cls, sim, packets_mode): sim.transport.virt_packet_energies, "erg" ), r_inner=r_inner, - spectrum_delta_frequency=sim.transport.spectrum_virtual.delta_frequency, - spectrum_frequency_bins=sim.transport.spectrum_virtual._frequency, - spectrum_luminosity_density_lambda=sim.transport.spectrum_virtual.luminosity_density_lambda, - spectrum_wavelength=sim.transport.spectrum_virtual.wavelength, + spectrum_delta_frequency=transport_state.spectrum_virtual.delta_frequency, + spectrum_frequency_bins=transport_state.spectrum_virtual._frequency, + spectrum_luminosity_density_lambda=transport_state.spectrum_virtual.luminosity_density_lambda, + spectrum_wavelength=transport_state.spectrum_virtual.wavelength, t_inner=t_inner, time_of_simulation=time_of_simulation, ) @@ -192,29 +196,29 @@ def from_simulation(cls, sim, packets_mode): # which got emitted return cls( last_interaction_type=sim.transport.last_interaction_type[ - sim.transport.emitted_packet_mask + transport_state.emitted_packet_mask ], last_line_interaction_in_id=sim.transport.last_line_interaction_in_id[ - sim.transport.emitted_packet_mask + transport_state.emitted_packet_mask ], last_line_interaction_out_id=sim.transport.last_line_interaction_out_id[ - sim.transport.emitted_packet_mask + transport_state.emitted_packet_mask ], last_line_interaction_in_nu=sim.transport.last_interaction_in_nu[ - sim.transport.emitted_packet_mask + transport_state.emitted_packet_mask ], lines_df=lines_df, - packet_nus=sim.transport.output_nu[ - sim.transport.emitted_packet_mask + packet_nus=transport_state.packet_collection.output_nus[ + transport_state.emitted_packet_mask ], - packet_energies=sim.transport.output_energy[ - sim.transport.emitted_packet_mask + packet_energies=transport_state.packet_collection.output_energies[ + transport_state.emitted_packet_mask ], r_inner=r_inner, - spectrum_delta_frequency=sim.transport.spectrum.delta_frequency, - spectrum_frequency_bins=sim.transport.spectrum._frequency, - spectrum_luminosity_density_lambda=sim.transport.spectrum.luminosity_density_lambda, - spectrum_wavelength=sim.transport.spectrum.wavelength, + spectrum_delta_frequency=transport_state.spectrum.delta_frequency, + spectrum_frequency_bins=transport_state.spectrum._frequency, + spectrum_luminosity_density_lambda=transport_state.spectrum.luminosity_density_lambda, + spectrum_wavelength=transport_state.spectrum.wavelength, t_inner=t_inner, time_of_simulation=time_of_simulation, ) diff --git a/tardis/visualization/widgets/grotrian.py b/tardis/visualization/widgets/grotrian.py index 9ee196fb346..f15e6132698 100644 --- a/tardis/visualization/widgets/grotrian.py +++ b/tardis/visualization/widgets/grotrian.py @@ -178,7 +178,7 @@ def from_simulation(cls, sim, **kwargs): ) level_population_data = sim.plasma.level_number_density line_interaction_analysis = { - filter_mode: LastLineInteraction.from_model(sim, filter_mode) + filter_mode: LastLineInteraction.from_simulation(sim, filter_mode) for filter_mode in cls.FILTER_MODES } return cls( diff --git a/tardis/visualization/widgets/line_info.py b/tardis/visualization/widgets/line_info.py index c21993a5fec..7ae2d483c0f 100644 --- a/tardis/visualization/widgets/line_info.py +++ b/tardis/visualization/widgets/line_info.py @@ -124,18 +124,21 @@ def from_simulation(cls, sim): ------- LineInfoWidget object """ + transport_state = sim.transport.transport_state return cls( lines_data=sim.plasma.lines.reset_index().set_index("line_id"), line_interaction_analysis={ - filter_mode: LastLineInteraction.from_model(sim, filter_mode) + filter_mode: LastLineInteraction.from_simulation( + sim, filter_mode + ) for filter_mode in cls.FILTER_MODES }, - spectrum_wavelength=sim.transport.spectrum.wavelength, - spectrum_luminosity_density_lambda=sim.transport.spectrum.luminosity_density_lambda.to( + spectrum_wavelength=transport_state.spectrum.wavelength, + spectrum_luminosity_density_lambda=transport_state.spectrum.luminosity_density_lambda.to( "erg/(s AA)" ), - virt_spectrum_wavelength=sim.transport.spectrum_virtual.wavelength, - virt_spectrum_luminosity_density_lambda=sim.transport.spectrum_virtual.luminosity_density_lambda.to( + virt_spectrum_wavelength=transport_state.spectrum_virtual.wavelength, + virt_spectrum_luminosity_density_lambda=transport_state.spectrum_virtual.luminosity_density_lambda.to( "erg/(s AA)" ), )