PacketSource should contain .to_hdf() method

tardis/io/util.py

@@ -242,7 +242,7 @@ def to_hdf_util(
                     path_or_buf, complevel=complevel, complib=complib
                 )
             except TypeError as e:
-                if e.message == "Expected bytes, got HDFStore":
+                if str(e) == "Expected bytes, got HDFStore":
                     # when path_or_buf is an HDFStore buffer instead
                     logger.debug(
                         "Expected bytes, got HDFStore. Changing path to HDF buffer"

tardis/transport/montecarlo/packet_source.py

@@ -6,7 +6,7 @@
 from tardis.transport.montecarlo.packet_collections import (
     PacketCollection,
 )
-
+from tardis.io.util import HDFWriterMixin
 from astropy import units as u
 
 
@@ -126,7 +126,7 @@ def calculate_radfield_luminosity(self):
         ).to("erg/s")
 
 
-class BlackBodySimpleSource(BasePacketSource):
+class BlackBodySimpleSource(BasePacketSource, HDFWriterMixin):
     """
     Simple packet source that generates Blackbody packets for the Montecarlo
     part.
@@ -143,6 +143,9 @@ class BlackBodySimpleSource(BasePacketSource):
         Secondary seed for global numpy rng (Deprecated: Legacy reasons only)
     """
 
+    hdf_properties = ["radius", "temperature", "base_seed"]
+    hdf_name = "black_body_simple_source"
+
     @classmethod
     def from_simulation_state(cls, simulation_state, *args, **kwargs):
         return cls(
@@ -276,7 +279,7 @@ def set_temperature_from_luminosity(self, luminosity: u.Quantity):
         ).to("K")
 
 
-class BlackBodySimpleSourceRelativistic(BlackBodySimpleSource):
+class BlackBodySimpleSourceRelativistic(BlackBodySimpleSource, HDFWriterMixin):
     """
     Simple packet source that generates Blackbody packets for the Montecarlo
     part.
@@ -295,6 +298,8 @@ class BlackBodySimpleSourceRelativistic(BlackBodySimpleSource):
         Secondary seed for global numpy rng (Deprecated: Legacy reasons only)
     """
 
+    hdf_properties = ["time_explosion", "radius", "temperature", "base_seed"]
+
     @classmethod
     def from_simulation_state(cls, simulation_state, *args, **kwargs):
         return cls(

tardis/transport/montecarlo/tests/test_packet_source.py

@@ -13,26 +13,10 @@
 from tardis.transport.montecarlo import (
     montecarlo_configuration as montecarlo_configuration,
 )
+from tardis.tests.fixtures.regression_data import RegressionData
 
 
-class TestPacketSource:
-    @pytest.fixture(scope="class")
-    def packet_unit_test_fpath(self, tardis_ref_path):
-        """
-        Path to `packet_unittest.h5`.
-
-        Parameters
-        ----------
-        tardis_ref_path : pd.HDFStore
-
-        Returns
-        -------
-        os.path
-        """
-        return os.path.abspath(
-            os.path.join(tardis_ref_path, "packet_unittest.h5")
-        )
-
+class TestBlackBodySimpleSource:
     @pytest.fixture(scope="class")
     def blackbodysimplesource(self, request):
         """
@@ -43,97 +27,98 @@ def blackbodysimplesource(self, request):
         tardis.transport.montecarlo.packet_source.BlackBodySimpleSource
         """
         cls = type(self)
-        cls.bb = BlackBodySimpleSource(
-            base_seed=1963, legacy_mode_enabled=True, legacy_second_seed=2508
+        bb = BlackBodySimpleSource(
+            radius=123,
+            temperature=10000 * u.K,
+            base_seed=1963,
+            legacy_second_seed=2508,
+            legacy_mode_enabled=True,
+        )
+        yield bb
+
+    def test_bb_nus(self, regression_data, blackbodysimplesource):
+        actual_nus = blackbodysimplesource.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, blackbodysimplesource):
+        actual_mus = blackbodysimplesource.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, blackbodysimplesource):
+        actual_unif_energies = blackbodysimplesource.create_packet_energies(
+            100
+        ).value
+        expected_unif_energies = regression_data.sync_ndarray(
+            actual_unif_energies
         )
-        yield cls.bb
+        assert_allclose(actual_unif_energies, expected_unif_energies)
 
+    def test_bb_attributes(self, regression_data, blackbodysimplesource):
+        actual_bb = blackbodysimplesource
+        expected_bb = regression_data.sync_hdf_store(actual_bb)[
+            "/black_body_simple_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)
+
+
+class TestBlackBodySimpleSourceRel:
     @pytest.fixture(scope="class")
-    def blackbody_simplesource_relativistic(self, request):
+    def blackbody_simplesource_relativistic(self):
         """
         Create BlackBodySimpleSourceRelativistic instance.
 
         Yields
         -------
-        tardis.transport.montecarlo.packet_source.BlackBodySimpleSourceRelativistic
+        tardis.montecarlo.packet_source.BlackBodySimpleSourceRelativistic
         """
         bb_rel = BlackBodySimpleSourceRelativistic(
-            base_seed=1963, legacy_mode_enabled=True, legacy_second_seed=2508
+            time_explosion=1123187,
+            base_seed=1963,
+            legacy_second_seed=2508,
+            legacy_mode_enabled=True,
         )
+        bb_rel.temperature = 10000 * u.K
+        bb_rel.beta = 0.25
         yield bb_rel
 
-    def test_bb_packet_sampling(
-        self,
-        request,
-        tardis_ref_data,
-        packet_unit_test_fpath,
-        blackbodysimplesource,
-    ):
-        """
-        Parameters
-        ----------
-        request : _pytest.fixtures.SubRequest
-        tardis_ref_data: pd.HDFStore
-        packet_unit_test_fpath: os.path
-        """
-        if request.config.getoption("--generate-reference"):
-            ref_bb = pd.read_hdf(packet_unit_test_fpath, key="/blackbody")
-            ref_bb.to_hdf(
-                tardis_ref_data, key="/packet_unittest/blackbody", mode="a"
-            )
-            pytest.skip("Reference data was generated during this run.")
-
-        ref_df = tardis_ref_data["/packet_unittest/blackbody"]
-        self.bb.temperature = 10000 * u.K
-        nus = self.bb.create_packet_nus(100).value
-        mus = self.bb.create_packet_mus(100)
-        unif_energies = self.bb.create_packet_energies(100).value
-        assert np.all(np.isclose(nus, ref_df["nus"]))
-        assert np.all(np.isclose(mus, ref_df["mus"]))
-        assert np.all(np.isclose(unif_energies, ref_df["energies"]))
-
-    def test_bb_packet_sampling_relativistic(
-        self,
-        tardis_ref_data,
-        blackbody_simplesource_relativistic,
-    ):
-        """
-        Parameters
-        ----------
-        tardis_ref_data : pd.HDFStore
-        blackbody_simplesource_relativistic : tardis.transport.montecarlo.packet_source.BlackBodySimpleSourceRelativistic
-        """
-        blackbody_simplesource_relativistic.temperature = 10000 * u.K
-        blackbody_simplesource_relativistic.beta = 0.25
+    def test_bb_nus(self, regression_data, blackbody_simplesource_relativistic):
+        actual_nus = blackbody_simplesource_relativistic.create_packet_nus(
+            100
+        ).value
+        expected_nus = regression_data.sync_ndarray(actual_nus)
+        assert_allclose(actual_nus, expected_nus)
 
-        nus = blackbody_simplesource_relativistic.create_packet_nus(100).value
-        unif_energies = (
+    def test_bb_energies(
+        self, regression_data, blackbody_simplesource_relativistic
+    ):
+        actual_unif_energies = (
             blackbody_simplesource_relativistic.create_packet_energies(
                 100
             ).value
         )
-        blackbody_simplesource_relativistic._reseed(2508)
-        mus = blackbody_simplesource_relativistic.create_packet_mus(10)
-
         gamma = np.sqrt(1 - blackbody_simplesource_relativistic.beta**2) ** -1
-        ref_df = tardis_ref_data["/packet_unittest/blackbody"]
-        expected_nus = ref_df["nus"]
-        expected_unif_energies = ref_df["energies"] * 1.6 / gamma
-        expected_mus = np.array(
-            [
-                0.60420546,
-                0.49899691,
-                0.69583288,
-                0.96812652,
-                0.01544154,
-                0.93562304,
-                0.44306545,
-                0.77010037,
-                0.896973,
-                0.67876489,
-            ]
+        expected_unif_energies = regression_data.sync_ndarray(
+            actual_unif_energies
         )
+        assert_allclose(actual_unif_energies, expected_unif_energies)
 
-        assert_allclose(nus, expected_nus)
-        assert_allclose(unif_energies, expected_unif_energies)
-        assert_allclose(mus, expected_mus, rtol=1e-6)
+    def test_bb_mus(self, regression_data, blackbody_simplesource_relativistic):
+        blackbody_simplesource_relativistic._reseed(2508)
+        actual_mus = blackbody_simplesource_relativistic.create_packet_mus(10)
+        expected_mus = regression_data.sync_ndarray(actual_mus)
+        assert_allclose(actual_mus, expected_mus)
+
+    def test_bb_attributes(
+        self, regression_data, blackbody_simplesource_relativistic
+    ):
+        actual_bb = blackbody_simplesource_relativistic
+        expected_bb = regression_data.sync_hdf_store(actual_bb)[
+            "/black_body_simple_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.time_explosion, actual_bb.time_explosion)