Fix bug in relativistic packet source

tardis/montecarlo/packet_source.py

@@ -313,7 +313,7 @@ def create_packets(self, no_of_packets):
         self.beta = ((self.radius / self.time_explosion) / const.c).to("")
         return super().create_packets(no_of_packets)
 
-    def create_packet_nus(self, no_of_packets):
+    def create_packet_mus(self, no_of_packets):
         """
         Create zero-limb-darkening packet :math:`\mu^\prime` distributed
         according to :math:`\\mu^\\prime=2 \\frac{\\mu^\\prime + \\beta}{2 \\beta + 1}`.
@@ -327,7 +327,7 @@ def create_packet_nus(self, no_of_packets):
 
         Returns
         -------
-        array of frequencies
+        Directions for packets
             numpy.ndarray
         """
         z = self.rng.random(no_of_packets)

tardis/montecarlo/tests/test_packet_source.py

@@ -3,8 +3,12 @@
 import numpy as np
 import pandas as pd
 import pytest
+from numpy.testing import assert_allclose
 
-from tardis.montecarlo.packet_source import BlackBodySimpleSource
+from tardis.montecarlo.packet_source import (
+    BlackBodySimpleSource,
+    BlackBodySimpleSourceRelativistic,
+)
 from tardis.montecarlo import (
     montecarlo_configuration as montecarlo_configuration,
 )
@@ -43,6 +47,22 @@ def blackbodysimplesource(self, request):
         yield cls.bb
         montecarlo_configuration.LEGACY_MODE_ENABLED = False
 
+    @pytest.fixture(scope="class")
+    def blackbody_simplesource_relativistic(self, request):
+        """
+        Create BlackBodySimpleSourceRelativistic instance.
+
+        Yields
+        -------
+        tardis.montecarlo.packet_source.BlackBodySimpleSourceRelativistic
+        """
+        montecarlo_configuration.LEGACY_MODE_ENABLED = True
+        bb_rel = BlackBodySimpleSourceRelativistic(
+            base_seed=1963, legacy_second_seed=2508
+        )
+        yield bb_rel
+        montecarlo_configuration.LEGACY_MODE_ENABLED = False
+
     def test_bb_packet_sampling(
         self,
         request,
@@ -51,8 +71,6 @@ def test_bb_packet_sampling(
         blackbodysimplesource,
     ):
         """
-        Test generate_plot_mpl method.
-
         Parameters
         ----------
         request : _pytest.fixtures.SubRequest
@@ -75,3 +93,47 @@ def test_bb_packet_sampling(
         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.montecarlo.packet_source.BlackBodySimpleSourceRelativistic
+        """
+        blackbody_simplesource_relativistic.temperature = 10000
+        blackbody_simplesource_relativistic.beta = 0.25
+
+        nus = blackbody_simplesource_relativistic.create_packet_nus(100)
+        unif_energies = (
+            blackbody_simplesource_relativistic.create_packet_energies(100)
+        )
+        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,
+            ]
+        )
+
+        assert_allclose(nus, expected_nus)
+        assert_allclose(unif_energies, expected_unif_energies)
+        assert_allclose(mus, expected_mus, rtol=1e-6)