From 45140361c023a6bd9afc5d773834fb1908771ceb Mon Sep 17 00:00:00 2001 From: Jack O'Brien Date: Sat, 27 Jul 2024 14:57:31 -0400 Subject: [PATCH] Weighted Packet Sampler (#2718) * Added a class for a weighted packet sampler which samples packets uniformly in frequency space then assigns energy to each packet based on their black body amplitude * Changed the hdf_name for the weighted source * Added fixture for the weighted packet source and a integration test against the regression data for the montecarlo_main_loop with the default packet source * Testing readding a fixture in the weighted sampler tests * moved the montecarlo config fixture to conftest.py * removed legacy mode from the simple weighted packet source fixture * updated path the the montecarlo main loop test data * updated path the the montecarlo main loop test data part 2 * fixed typo * fixed typo 2 * ran ruff * ran black * Added the unit test for the blackbodyweightedsource * Fixed a typo in one of the fixtures * Updated the documentation and cleaned up the implimentation of the weighted packet source * fixed weighted packet source base clas * Use the rng for sampling rather than np.random as otherwise there are issues with the tests * Fixed issues with the rng and made the unit tests avoid legacy mode * ran black --- tardis/transport/montecarlo/tests/conftest.py | 37 ++++++-- .../tests/test_montecarlo_main_loop.py | 15 ---- .../tests/test_weighted_packet_source.py | 60 +++++++++++++ ...test_weighted_packet_source_integration.py | 66 ++++++++++++++ .../montecarlo/weighted_packet_source.py | 89 +++++++++++++++++++ 5 files changed, 245 insertions(+), 22 deletions(-) create mode 100644 tardis/transport/montecarlo/tests/test_weighted_packet_source.py create mode 100644 tardis/transport/montecarlo/tests/test_weighted_packet_source_integration.py create mode 100644 tardis/transport/montecarlo/weighted_packet_source.py diff --git a/tardis/transport/montecarlo/tests/conftest.py b/tardis/transport/montecarlo/tests/conftest.py index 201221e9bf0..61f80e4aac0 100644 --- a/tardis/transport/montecarlo/tests/conftest.py +++ b/tardis/transport/montecarlo/tests/conftest.py @@ -1,23 +1,39 @@ from copy import deepcopy -import pytest import numpy as np +import pytest from numba import njit -from tardis.opacities.opacity_state import opacity_state_initialize -from tardis.transport.montecarlo.packet_collections import ( - VPacketCollection, -) +from tardis.opacities.opacity_state import opacity_state_initialize from tardis.simulation import Simulation from tardis.transport.montecarlo import RPacket from tardis.transport.montecarlo.estimators.radfield_mc_estimators import ( RadiationFieldMCEstimators, ) - - from tardis.transport.montecarlo.numba_interface import ( opacity_state_initialize, ) +from tardis.transport.montecarlo.packet_collections import ( + VPacketCollection, +) +from tardis.transport.montecarlo.weighted_packet_source import ( + BlackBodyWeightedSource, +) + + +@pytest.fixture(scope="function") +def montecarlo_main_loop_config( + config_montecarlo_1e5_verysimple, +): + # Setup model config from verysimple + + config_montecarlo_1e5_verysimple.montecarlo.last_no_of_packets = 1e5 + config_montecarlo_1e5_verysimple.montecarlo.no_of_virtual_packets = 0 + config_montecarlo_1e5_verysimple.montecarlo.iterations = 1 + config_montecarlo_1e5_verysimple.plasma.line_interaction_type = "macroatom" + + del config_montecarlo_1e5_verysimple["config_dirname"] + return config_montecarlo_1e5_verysimple @pytest.fixture(scope="package") @@ -28,6 +44,13 @@ def nb_simulation_verysimple(config_verysimple, atomic_dataset): return sim +@pytest.fixture(scope="package") +def simple_weighted_packet_source(): + packet_source = BlackBodyWeightedSource(base_seed=1963) + + return packet_source + + @pytest.fixture(scope="package") def verysimple_opacity_state(nb_simulation_verysimple): return opacity_state_initialize( diff --git a/tardis/transport/montecarlo/tests/test_montecarlo_main_loop.py b/tardis/transport/montecarlo/tests/test_montecarlo_main_loop.py index 295cf8b703a..bcbfd4f7ea2 100644 --- a/tardis/transport/montecarlo/tests/test_montecarlo_main_loop.py +++ b/tardis/transport/montecarlo/tests/test_montecarlo_main_loop.py @@ -11,21 +11,6 @@ def test_montecarlo_radial1d(): assert False -@pytest.fixture(scope="function") -def montecarlo_main_loop_config( - config_montecarlo_1e5_verysimple, -): - # Setup model config from verysimple - - config_montecarlo_1e5_verysimple.montecarlo.last_no_of_packets = 1e5 - config_montecarlo_1e5_verysimple.montecarlo.no_of_virtual_packets = 0 - config_montecarlo_1e5_verysimple.montecarlo.iterations = 1 - config_montecarlo_1e5_verysimple.plasma.line_interaction_type = "macroatom" - - del config_montecarlo_1e5_verysimple["config_dirname"] - return config_montecarlo_1e5_verysimple - - def test_montecarlo_main_loop( montecarlo_main_loop_config, regression_data, diff --git a/tardis/transport/montecarlo/tests/test_weighted_packet_source.py b/tardis/transport/montecarlo/tests/test_weighted_packet_source.py new file mode 100644 index 00000000000..34b8f4599b7 --- /dev/null +++ b/tardis/transport/montecarlo/tests/test_weighted_packet_source.py @@ -0,0 +1,60 @@ +import os + +from astropy import units as u +import numpy as np +import pandas as pd +import pytest +from numpy.testing import assert_allclose + +from tardis.transport.montecarlo.weighted_packet_source import ( + BlackBodyWeightedSource, +) +from tardis.tests.fixtures.regression_data import RegressionData + + +class TestBlackBodyWeightedSource: + @pytest.fixture(scope="class") + def blackbodyweightedsource(self, request): + """ + Create blackbodyweightedsource instance. + + Yields + ------- + tardis.transport.montecarlo.packet_source.blackbodyweightedsource + """ + cls = type(self) + bb = BlackBodyWeightedSource( + radius=123, + temperature=10000 * u.K, + base_seed=1963, + legacy_mode_enabled=False, + ) + yield bb + + def test_bb_nus(self, regression_data, blackbodyweightedsource): + actual_nus = blackbodyweightedsource.create_packet_nus(100).value + expected_nus = regression_data.sync_ndarray(actual_nus) + assert_allclose(actual_nus, expected_nus) + + def test_bb_mus(self, regression_data, blackbodyweightedsource): + actual_mus = blackbodyweightedsource.create_packet_mus(100) + expected_mus = regression_data.sync_ndarray(actual_mus) + assert_allclose(actual_mus, expected_mus) + + def test_bb_energies(self, regression_data, blackbodyweightedsource): + actual_unif_energies = blackbodyweightedsource.create_packet_energies( + 100 + ).value + expected_unif_energies = regression_data.sync_ndarray( + actual_unif_energies + ) + assert_allclose(actual_unif_energies, expected_unif_energies) + + def test_bb_attributes(self, regression_data, blackbodyweightedsource): + actual_bb = blackbodyweightedsource + expected_bb = regression_data.sync_hdf_store(actual_bb)[ + "/black_body_weighted_source/scalars" + ] + assert_allclose(expected_bb.base_seed, actual_bb.base_seed) + assert_allclose(expected_bb.temperature, actual_bb.temperature.value) + assert_allclose(expected_bb.radius, actual_bb.radius) diff --git a/tardis/transport/montecarlo/tests/test_weighted_packet_source_integration.py b/tardis/transport/montecarlo/tests/test_weighted_packet_source_integration.py new file mode 100644 index 00000000000..a59904a2fde --- /dev/null +++ b/tardis/transport/montecarlo/tests/test_weighted_packet_source_integration.py @@ -0,0 +1,66 @@ +from copy import deepcopy + +import numpy.testing as npt +import pandas as pd + +from tardis.simulation import Simulation + + +def test_montecarlo_main_loop_weighted( + montecarlo_main_loop_config, + regression_data, + atomic_dataset, + simple_weighted_packet_source, +): + atomic_dataset = deepcopy(atomic_dataset) + montecarlo_main_loop_simulation_weighted = Simulation.from_config( + montecarlo_main_loop_config, + atom_data=atomic_dataset, + virtual_packet_logging=False, + legacy_mode_enabled=True, + ) + montecarlo_main_loop_simulation_weighted.packet_source = ( + simple_weighted_packet_source + ) + montecarlo_main_loop_simulation_weighted.run_convergence() + montecarlo_main_loop_simulation_weighted.run_final() + + # Get the montecarlo simple regression data + regression_data_dir = ( + regression_data.absolute_regression_data_dir.absolute().parents[0] + / "test_montecarlo_main_loop/test_montecarlo_main_loop.h5" + ) + expected_hdf_store = pd.HDFStore(regression_data_dir, mode="r") + + # Load compare data from refdata + + expected_nu = expected_hdf_store[ + "/simulation/transport/transport_state/output_nu" + ] + expected_energy = expected_hdf_store[ + "/simulation/transport/transport_state/output_energy" + ] + expected_nu_bar_estimator = expected_hdf_store[ + "/simulation/transport/transport_state/nu_bar_estimator" + ] + expected_j_estimator = expected_hdf_store[ + "/simulation/transport/transport_state/j_estimator" + ] + expected_hdf_store.close() + transport_state = ( + montecarlo_main_loop_simulation_weighted.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.radfield_mc_estimators.nu_bar_estimator + ) + actual_j_estimator = transport_state.radfield_mc_estimators.j_estimator + + # Compare + npt.assert_allclose( + actual_nu_bar_estimator, expected_nu_bar_estimator, rtol=1e-2 + ) + npt.assert_allclose(actual_j_estimator, expected_j_estimator, rtol=1e-2) + npt.assert_allclose(actual_energy, expected_energy, rtol=1e-2) + npt.assert_allclose(actual_nu, expected_nu, rtol=1e-2) diff --git a/tardis/transport/montecarlo/weighted_packet_source.py b/tardis/transport/montecarlo/weighted_packet_source.py new file mode 100644 index 00000000000..dbdbb85efd3 --- /dev/null +++ b/tardis/transport/montecarlo/weighted_packet_source.py @@ -0,0 +1,89 @@ +import numexpr as ne +import numpy as np +from astropy import constants as const +from astropy import units as u + +from tardis.transport.montecarlo.packet_source import ( + BasePacketSource, + HDFWriterMixin, + BlackBodySimpleSource, +) +from tardis.util.base import intensity_black_body + + +class BlackBodyWeightedSource(BlackBodySimpleSource): + """ + Simple packet source that generates Blackbody packets for the Montecarlo + part. + + Parameters + ---------- + radius : astropy.units.Quantity + Initial packet radius + temperature : astropy.units.Quantity + Absolute Temperature. + base_seed : int + Base Seed for random number generator + legacy_secondary_seed : int + Secondary seed for global numpy rng (Deprecated: Legacy reasons only) + """ + + hdf_properties = ["radius", "temperature", "base_seed"] + hdf_name = "black_body_weighted_source" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._reseed( + self.base_seed + ) # Needed because base_source doesn't seed by default + + def create_packet_nus(self, no_of_packets, l_samples=1000): + """ + Create packet :math:`\\nu` distributed uniformly over + bounds taken from the BlackBodySimpleSource distribution + + Parameters + ---------- + no_of_packets : int + l_samples : int + number of l_samples needed for sampling from BlackBodySimpleSource + + Returns + ------- + array of frequencies + numpy.ndarray + """ + nus = super().create_packet_nus(no_of_packets, l_samples) + + nu_min = nus.min() + nu_max = nus.max() + + self.nus = ( + self.rng.uniform(nu_min.cgs.value, nu_max.cgs.value, no_of_packets) + * nus.unit + ) + + return self.nus + + def create_packet_energies(self, no_of_packets): + """ + Set energy weight for each packet from the relative contribution to + the Planck Distribution + + Parameters + ---------- + no_of_packets : int + number of packets + + Returns + ------- + energies for packets + numpy.ndarray + """ + try: + self.nus + except AttributeError: + self.nus = self.create_packet_nus(no_of_packets) + + intensity = intensity_black_body(self.nus.cgs.value, self.temperature) + return intensity / intensity.sum() * u.erg