Consistent units in packet source (#2542)

* Consistent units in packet source

* Format

* Address comments

tardis/montecarlo/packet_source.py

@@ -88,17 +88,17 @@ def create_packets(self, no_of_packets, seed_offset=0, *args, **kwargs):
             self.MAX_SEED_VAL, no_of_packets, replace=True
         )
 
-        radii = self.create_packet_radii(no_of_packets, *args, **kwargs)
-        nus = self.create_packet_nus(no_of_packets, *args, **kwargs)
+        radii = self.create_packet_radii(no_of_packets, *args, **kwargs).value
+        nus = self.create_packet_nus(no_of_packets, *args, **kwargs).value
         mus = self.create_packet_mus(no_of_packets, *args, **kwargs)
-        energies = self.create_packet_energies(no_of_packets, *args, **kwargs)
+        energies = self.create_packet_energies(
+            no_of_packets, *args, **kwargs
+        ).value
         # Check if all arrays have the same length
         assert (
             len(radii) == len(nus) == len(mus) == len(energies) == no_of_packets
         )
-        radiation_field_luminosity = (
-            self.calculate_radfield_luminosity().to(u.erg / u.s).value
-        )
+        radiation_field_luminosity = self.calculate_radfield_luminosity().value
         return PacketCollection(
             radii,
             nus,
@@ -123,7 +123,7 @@ def calculate_radfield_luminosity(self):
         return (
             4
             * np.pi
-            * const.sigma_sb.cgs
+            * const.sigma_sb
             * self.radius**2
             * self.temperature**4
         ).to("erg/s")
@@ -136,9 +136,9 @@ class BlackBodySimpleSource(BasePacketSource):
 
     Parameters
     ----------
-    radius : float64
+    radius : astropy.units.Quantity
         Initial packet radius
-    temperature : float
+    temperature : astropy.units.Quantity
         Absolute Temperature.
     base_seed : int
         Base Seed for random number generator
@@ -150,7 +150,7 @@ class BlackBodySimpleSource(BasePacketSource):
     def from_simulation_state(cls, simulation_state, *args, **kwargs):
         return cls(
             simulation_state.r_inner[0],
-            simulation_state.t_inner.value,
+            simulation_state.t_inner,
             *args,
             **kwargs,
         )
@@ -179,7 +179,7 @@ def create_packet_radii(self, no_of_packets):
         Radii for packets
             numpy.ndarray
         """
-        return np.ones(no_of_packets) * self.radius
+        return np.ones(no_of_packets) * self.radius.cgs
 
     def create_packet_nus(self, no_of_packets, l_samples=1000):
         """
@@ -222,12 +222,7 @@ def create_packet_nus(self, no_of_packets, l_samples=1000):
         xis_prod = np.prod(xis[1:], 0)
         x = ne.evaluate("-log(xis_prod)/l")
 
-        if isinstance(self.temperature, u.Quantity):
-            temperature = self.temperature.value
-        else:
-            temperature = self.temperature
-
-        return x * (const.k_B.cgs.value * temperature) / const.h.cgs.value
+        return (x * (const.k_B * self.temperature) / const.h).cgs
 
     def create_packet_mus(self, no_of_packets):
         """
@@ -266,7 +261,7 @@ def create_packet_energies(self, no_of_packets):
         energies for packets
             numpy.ndarray
         """
-        return np.ones(no_of_packets) / no_of_packets
+        return np.ones(no_of_packets) / no_of_packets * u.erg
 
     def set_temperature_from_luminosity(self, luminosity: u.Quantity):
         """
@@ -291,11 +286,11 @@ class BlackBodySimpleSourceRelativistic(BlackBodySimpleSource):
 
     Parameters
     ----------
-    time_explosion : float 64
+    time_explosion : astropy.units.Quantity
         Time elapsed since explosion
-    radius : float64
+    radius : astropy.units.Quantity
         Initial packet radius
-    temperature : float
+    temperature : astropy.units.Quantity
         Absolute Temperature.
     base_seed : int
         Base Seed for random number generator
@@ -308,7 +303,7 @@ def from_simulation_state(cls, simulation_state, *args, **kwargs):
         return cls(
             simulation_state.time_explosion,
             simulation_state.r_inner[0],
-            simulation_state.t_inner.value,
+            simulation_state.t_inner,
             *args,
             **kwargs,
         )
@@ -338,7 +333,7 @@ def create_packets(self, no_of_packets):
         """
         if self.radius is None or self.time_explosion is None:
             raise ValueError("Black body Radius or Time of Explosion isn't set")
-        self.beta = ((self.radius / self.time_explosion) / const.c).to("")
+        self.beta = (self.radius / self.time_explosion) / const.c
         return super().create_packets(no_of_packets)
 
     def create_packet_mus(self, no_of_packets):
@@ -387,4 +382,4 @@ def create_packet_energies(self, no_of_packets):
         # 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
+        return energies * static_inner_boundary2cmf_factor / gamma * u.erg

tardis/montecarlo/tests/test_packet_source.py

@@ -1,5 +1,6 @@
 import os
 
+from astropy import units as u
 import numpy as np
 import pandas as pd
 import pytest
@@ -76,7 +77,6 @@ def test_bb_packet_sampling(
         request : _pytest.fixtures.SubRequest
         tardis_ref_data: pd.HDFStore
         packet_unit_test_fpath: os.path
-        blackbodysimplesource: tardis.montecarlo.packet_source.BlackBodySimpleSource
         """
         if request.config.getoption("--generate-reference"):
             ref_bb = pd.read_hdf(packet_unit_test_fpath, key="/blackbody")
@@ -86,10 +86,10 @@ def test_bb_packet_sampling(
             pytest.skip("Reference data was generated during this run.")
 
         ref_df = tardis_ref_data["/packet_unittest/blackbody"]
-        self.bb.temperature = 10000
-        nus = self.bb.create_packet_nus(100)
+        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)
+        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"]))
@@ -105,12 +105,14 @@ def test_bb_packet_sampling_relativistic(
         tardis_ref_data : pd.HDFStore
         blackbody_simplesource_relativistic : tardis.montecarlo.packet_source.BlackBodySimpleSourceRelativistic
         """
-        blackbody_simplesource_relativistic.temperature = 10000
+        blackbody_simplesource_relativistic.temperature = 10000 * u.K
         blackbody_simplesource_relativistic.beta = 0.25
 
-        nus = blackbody_simplesource_relativistic.create_packet_nus(100)
+        nus = blackbody_simplesource_relativistic.create_packet_nus(100).value
         unif_energies = (
-            blackbody_simplesource_relativistic.create_packet_energies(100)
+            blackbody_simplesource_relativistic.create_packet_energies(
+                100
+            ).value
         )
         blackbody_simplesource_relativistic._reseed(2508)
         mus = blackbody_simplesource_relativistic.create_packet_mus(10)