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

@@ -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(

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,

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)

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)

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