Fix relativistic packet initialization

tardis/montecarlo/base.py

@@ -106,7 +106,12 @@ def __init__(
 
         self.seed = seed
         if packet_source is None:
-            self.packet_source = source.BlackBodySimpleSource(seed)
+            if not enable_full_relativity:
+                self.packet_source = source.BlackBodySimpleSource(seed)
+            else:
+                self.packet_source = source.BlackBodySimpleSourceRelativistic(
+                    seed
+                )
         else:
             self.packet_source = packet_source
         # inject different packets
@@ -202,7 +207,9 @@ def _initialize_geometry_arrays(self, model):
         self.v_inner_cgs = model.v_inner.to("cm/s").value
         self.v_outer_cgs = model.v_outer.to("cm/s").value
 
-    def _initialize_packets(self, T, no_of_packets, iteration, radius):
+    def _initialize_packets(
+        self, T, no_of_packets, iteration, radius, time_explosion
+    ):
         # the iteration 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
@@ -211,9 +218,15 @@ def _initialize_packets(self, T, no_of_packets, iteration, radius):
         seed = self.seed + iteration
         rng = np.random.default_rng(seed=seed)
         seeds = rng.choice(MAX_SEED_VAL, no_of_packets, replace=True)
-        radii, nus, mus, energies = self.packet_source.create_packets(
-            T, no_of_packets, rng, radius
-        )
+        if not self.enable_full_relativity:
+            radii, nus, mus, energies = self.packet_source.create_packets(
+                T, no_of_packets, rng, radius
+            )
+        else:
+            radii, nus, mus, energies = self.packet_source.create_packets(
+                T, no_of_packets, rng, radius, time_explosion
+            )
+
         mc_config_module.packet_seeds = seeds
         self.input_r = radii
         self.input_nu = nus
@@ -344,7 +357,11 @@ def run(
         self._initialize_geometry_arrays(model)
 
         self._initialize_packets(
-            model.t_inner.value, no_of_packets, iteration, model.r_inner[0]
+            model.t_inner.value,
+            no_of_packets,
+            iteration,
+            model.r_inner[0],
+            model.time_explosion,
         )
 
         configuration_initialize(self, no_of_virtual_packets)

tardis/montecarlo/montecarlo_numba/vpacket.py

@@ -28,7 +28,10 @@
     angle_aberration_CMF_to_LF,
 )
 
-from tardis.montecarlo.montecarlo_numba.numba_config import SIGMA_THOMSON
+from tardis.montecarlo.montecarlo_numba.numba_config import (
+    SIGMA_THOMSON,
+    C_SPEED_OF_LIGHT,
+)
 
 vpacket_spec = [
     ("r", float64),
@@ -244,6 +247,11 @@ def trace_vpacket_volley(
         v_packet_on_inner_boundary = True
         mu_min = 0.0
 
+        if montecarlo_configuration.full_relativity:
+            inv_c = 1 / C_SPEED_OF_LIGHT
+            inv_t = 1 / numba_model.time_explosion
+            beta_inner = numba_radial_1d_geometry.r_inner[0] * inv_t * inv_c
+
     mu_bin = (1.0 - mu_min) / no_of_vpackets
     r_packet_doppler_factor = get_doppler_factor(
         r_packet.r, r_packet.mu, numba_model.time_explosion
@@ -252,7 +260,16 @@ def trace_vpacket_volley(
         v_packet_mu = mu_min + i * mu_bin + np.random.random() * mu_bin
 
         if v_packet_on_inner_boundary:  # The weights are described in K&S 2014
-            weight = 2 * v_packet_mu / no_of_vpackets
+            if not montecarlo_configuration.full_relativity:
+                weight = 2 * v_packet_mu / no_of_vpackets
+            else:
+                weight = (
+                    2
+                    * (v_packet_mu + beta_inner)
+                    / (2 * beta_inner + 1)
+                    / no_of_vpackets
+                )
+
         else:
             weight = (1 - mu_min) / (2 * no_of_vpackets)
 

tardis/montecarlo/packet_source.py

@@ -139,3 +139,76 @@ def create_packets(self, T, no_of_packets, rng, radius):
         energies = self.create_uniform_packet_energies(no_of_packets, rng)
 
         return radii, nus, mus, energies
+
+
+class BlackBodySimpleSourceRelativistic(BlackBodySimpleSource):
+    def create_packets(self, T, no_of_packets, rng, radius, time_explosion):
+        """Generate relativistic black-body packet properties as arrays
+
+        Parameters
+        ----------
+        T : float64
+            Temperature
+        no_of_packets : int
+            Number of packets
+        rng : numpy random number generator
+        radius : float64
+            Initial packet radius
+        time_explosion: float64
+            Time elapsed since explosion
+
+        Returns
+        -------
+        array
+            Packet radii
+        array
+            Packet frequencies
+        array
+            Packet directions
+        array
+            Packet energies
+        """
+        self.beta = ((radius / time_explosion) / const.c).to("")
+        return super().create_packets(T, no_of_packets, rng, radius)
+
+    def create_zero_limb_darkening_packet_mus(self, no_of_packets, rng):
+        """
+        Create zero-limb-darkening packet :math:`\mu^\prime` distributed
+        according to :math:`\\mu^\\prime=2 \\frac{\\mu^\\prime + \\beta}{2 \\beta + 1}`.
+        The complicated distribution is due to the fact that the inner boundary
+        on which the packets are initialized is not comoving with the material.
+
+        Parameters
+        ----------
+        no_of_packets : int
+            number of packets to be created
+        """
+        z = rng.random(no_of_packets)
+        beta = self.beta
+        return -beta + np.sqrt(beta**2 + 2 * beta * z + z)
+
+    def create_uniform_packet_energies(self, no_of_packets, rng):
+        """
+        Uniformly distribute energy in arbitrary units where the ensemble of
+        packets has energy of 1 multiplied by relativistic correction factors.
+
+        Parameters
+        ----------
+        no_of_packets : int
+            number of packets
+
+        Returns
+        -------
+            energies for packets : numpy.ndarray
+        """
+        beta = self.beta
+        gamma = 1.0 / np.sqrt(1 - beta**2)
+        static_inner_boundary2cmf_factor = (2 * beta + 1) / (1 - beta**2)
+        energies = np.ones(no_of_packets) / no_of_packets
+        # In principle, the factor gamma should be applied to the time of
+        # simulation to account for time dilation between the lab and comoving
+        # frame. However, all relevant quantities (luminosities, estimators, ...)
+        # are calculated as ratios of packet energies and the time of simulation.
+        # Thus, we can absorb the factor gamma in the packet energies, which is
+        # more convenient.
+        return energies * static_inner_boundary2cmf_factor / gamma

tardis/transport/r_packet_transport.py

@@ -15,9 +15,7 @@
 from tardis.transport.frame_transformations import (
     get_doppler_factor,
 )
-from tardis.montecarlo.montecarlo_numba.numba_config import (
-    ENABLE_FULL_RELATIVITY,
-)
+import tardis.montecarlo.montecarlo_numba.numba_config as nc
 from tardis.montecarlo.montecarlo_numba.opacities import calculate_tau_electron
 from tardis.montecarlo.montecarlo_numba.r_packet import (
     InteractionType,
@@ -211,8 +209,8 @@ def move_r_packet(r_packet, distance, time_explosion, numba_estimator):
         comov_nu = r_packet.nu * doppler_factor
         comov_energy = r_packet.energy * doppler_factor
 
-        # Account for length contraction and angle aberration
-        if ENABLE_FULL_RELATIVITY:
+        # Account for length contraction
+        if nc.ENABLE_FULL_RELATIVITY:
             distance *= doppler_factor
 
         set_estimators(