Blackify TARDIS again (#1971)

* Blackify TARDIS again * reformat montecarlo
tardis-sn · Apr 25, 2022 · 01544d1 · 01544d1
1 parent c4f560d
commit 01544d1
Show file tree

Hide file tree

Showing 47 changed files with 404 additions and 345 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -22,8 +22,6 @@ exclude = '''
     | buck-out
     | build
     | dist
-    | model         # temporary - to remove later
-    | montecarlo    # temporary
   )/
   | setup.py
   | docs/conf.py

diff --git a/tardis/__init__.py b/tardis/__init__.py
@@ -17,8 +17,8 @@
 
 from astropy import physical_constants, astronomical_constants
 
-physical_constants.set('codata2014')
-astronomical_constants.set('iau2012')
+physical_constants.set("codata2014")
+astronomical_constants.set("iau2012")
 
 # ----------------------------------------------------------------------------
 

diff --git a/tardis/analysis/opacities.py b/tardis/analysis/opacities.py
@@ -342,15 +342,15 @@ def _calc_thomson_scattering_opacity(self):
             sigma_T = const.sigma_T
         except AttributeError:
             logger.warning("using astropy < 1.1.1: setting sigma_T manually")
-            sigma_T = 6.65245873e-29 * units.m ** 2
+            sigma_T = 6.65245873e-29 * units.m**2
 
         edens = self.mdl.plasma.electron_densities.values
 
         try:
             edens.to("1/cm^3")
         except AttributeError:
             logger.info("setting electron density units by hand (cm^-3)")
-            edens = edens / units.cm ** 3
+            edens = edens / units.cm**3
 
         kappa_thom = edens * sigma_T
 
@@ -376,9 +376,7 @@ def _calc_planck_mean_opacity(self):
             bb_nu = Blackbody(T)
 
             tmp = (
-                bb_nu(self.nu_bins[:-1])                
-                * delta_nu
-                * self.kappa_tot[:, 0]
+                bb_nu(self.nu_bins[:-1]) * delta_nu * self.kappa_tot[:, 0]
             ).sum()
             tmp /= (bb_nu(self.nu_bins[:-1], T) * delta_nu).sum()
 

diff --git a/tardis/conftest.py b/tardis/conftest.py
@@ -171,6 +171,7 @@ def config_montecarlo_1e5_verysimple():
     config = Configuration.from_yaml(path)
     return config
 
+
 @pytest.fixture(scope="session")
 def simulation_verysimple(config_verysimple, atomic_dataset):
     atomic_data = deepcopy(atomic_dataset)

diff --git a/tardis/gui/widgets.py b/tardis/gui/widgets.py
@@ -162,8 +162,8 @@ def shell_picker(self, shell, mouseevent):
         """Enable picking shells in the shell plot."""
         if mouseevent.xdata is None:
             return False, dict()
-        mouse_r2 = mouseevent.xdata ** 2 + mouseevent.ydata ** 2
-        if shell.r_inner ** 2 < mouse_r2 < shell.r_outer ** 2:
+        mouse_r2 = mouseevent.xdata**2 + mouseevent.ydata**2
+        if shell.r_inner**2 < mouse_r2 < shell.r_outer**2:
             return True, dict()
         return False, dict()
 

diff --git a/tardis/io/logger/tests/test_logging.py b/tardis/io/logger/tests/test_logging.py
@@ -7,7 +7,10 @@
 from tardis import run_tardis
 import pytest
 
-pytestmark = pytest.mark.skip(reason='logging testing slow and disabled for now')
+pytestmark = pytest.mark.skip(
+    reason="logging testing slow and disabled for now"
+)
+
 
 def test_logging_simulation(atomic_data_fname, caplog):
     """

diff --git a/tardis/io/tests/test_model_reader.py b/tardis/io/tests/test_model_reader.py
@@ -55,8 +55,8 @@ def test_simple_read_artis_density(artis_density_fname):
     assert np.isclose(0.00114661 * u.day, time_of_model, atol=1e-7 * u.day)
     assert np.isclose(
         mean_density[23],
-        0.2250048 * u.g / u.cm ** 3,
-        atol=1.0e-6 * u.g / u.cm ** 3,
+        0.2250048 * u.g / u.cm**3,
+        atol=1.0e-6 * u.g / u.cm**3,
     )
     assert len(mean_density) == 69
     assert len(velocity) == len(mean_density) + 1
@@ -115,11 +115,11 @@ def test_simple_read_cmfgen_density(cmfgen_fname):
     assert np.isclose(0.976 * u.day, time_of_model, atol=1e-7 * u.day)
     assert np.isclose(
         mean_density[4],
-        4.2539537e-09 * u.g / u.cm ** 3,
-        atol=1.0e-6 * u.g / u.cm ** 3,
+        4.2539537e-09 * u.g / u.cm**3,
+        atol=1.0e-6 * u.g / u.cm**3,
     )
     assert np.isclose(
-        electron_densities[5], 2.6e14 * u.cm ** -3, atol=1.0e-6 * u.cm ** -3
+        electron_densities[5], 2.6e14 * u.cm**-3, atol=1.0e-6 * u.cm**-3
     )
     assert len(mean_density) == 9
     assert len(velocity) == len(mean_density) + 1
diff --git a/tardis/model/tests/test_base.py b/tardis/model/tests/test_base.py
@@ -173,31 +173,31 @@ def test_abundances(self):
     def test_densities(self):
         assert_almost_equal(
             self.model.density[0].to(u.Unit("g/cm3")).value,
-            9.7656229e-11 / 13.0 ** 3,
+            9.7656229e-11 / 13.0**3,
         )
         assert_almost_equal(
             self.model.density[1].to(u.Unit("g/cm3")).value,
-            4.8170911e-11 / 13.0 ** 3,
+            4.8170911e-11 / 13.0**3,
         )
         assert_almost_equal(
             self.model.density[2].to(u.Unit("g/cm3")).value,
-            2.5600000e-11 / 13.0 ** 3,
+            2.5600000e-11 / 13.0**3,
         )
         assert_almost_equal(
             self.model.density[3].to(u.Unit("g/cm3")).value,
-            1.4450533e-11 / 13.0 ** 3,
+            1.4450533e-11 / 13.0**3,
         )
         assert_almost_equal(
             self.model.density[4].to(u.Unit("g/cm3")).value,
-            8.5733893e-11 / 13.0 ** 3,
+            8.5733893e-11 / 13.0**3,
         )
         assert_almost_equal(
             self.model.density[5].to(u.Unit("g/cm3")).value,
-            5.3037103e-11 / 13.0 ** 3,
+            5.3037103e-11 / 13.0**3,
         )
         assert_almost_equal(
             self.model.density[6].to(u.Unit("g/cm3")).value,
-            3.3999447e-11 / 13.0 ** 3,
+            3.3999447e-11 / 13.0**3,
         )
 
 
@@ -313,6 +313,7 @@ def test_model_decay(simple_isotope_abundance):
 
 class TestModelState:
     """Test the ModelState class."""
+
     def setup(self):
         """Initialize config and model."""
         filename = "tardis_configv1_verysimple.yml"
@@ -387,6 +388,7 @@ def test_density(self):
 def to_hdf_buffer(hdf_file_path, simulation_verysimple):
     simulation_verysimple.model.to_hdf(hdf_file_path, overwrite=True)
 
+
 model_scalar_attrs = ["t_inner"]
 
 

diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py
@@ -29,7 +29,7 @@
 logger = logging.getLogger(__name__)
 
 
-MAX_SEED_VAL = 2 ** 32 - 1
+MAX_SEED_VAL = 2**32 - 1
 
 # MAX_SEED_VAL must be multiple orders of magnitude larger than no_of_packets;
 # otherwise, each packet would not have its own seed. Here, we set the max
@@ -73,14 +73,14 @@ class MontecarloRunner(HDFWriterMixin):
 
     hdf_name = "runner"
     w_estimator_constant = (
-        (const.c ** 2 / (2 * const.h))
-        * (15 / np.pi ** 4)
+        (const.c**2 / (2 * const.h))
+        * (15 / np.pi**4)
         * (const.h / const.k_B) ** 4
         / (4 * np.pi)
     ).cgs.value
 
     t_rad_estimator_constant = (
-        (np.pi ** 4 / (15 * 24 * zeta(5, 1))) * (const.h / const.k_B)
+        (np.pi**4 / (15 * 24 * zeta(5, 1))) * (const.h / const.k_B)
     ).cgs.value
 
     def __init__(
@@ -126,7 +126,7 @@ def __init__(
         self.seed = seed
         self._integrator = None
         self._spectrum_integrated = None
-        self.use_gpu=use_gpu
+        self.use_gpu = use_gpu
 
         self.virt_logging = virtual_packet_logging
         self.virt_packet_last_interaction_type = np.ones(2) * -1
@@ -269,10 +269,11 @@ def spectrum_virtual(self):
     @property
     def spectrum_integrated(self):
         if self._spectrum_integrated is None:
-            #This was changed from unpacking to specific attributes as compute
-            #is not used in calculate_spectrum
+            # This was changed from unpacking to specific attributes as compute
+            # is not used in calculate_spectrum
             self._spectrum_integrated = self.integrator.calculate_spectrum(
-                self.spectrum_frequency[:-1], points=self.integrator_settings.points, 
+                self.spectrum_frequency[:-1],
+                points=self.integrator_settings.points,
                 interpolate_shells=self.integrator_settings.interpolate_shells,
             )
         return self._spectrum_integrated
@@ -560,7 +561,7 @@ def calculate_radiationfield_properties(self):
         w = self.j_estimator / (
             4
             * const.sigma_sb.cgs.value
-            * t_rad ** 4
+            * t_rad**4
             * self.time_of_simulation.value
             * self.volume.value
         )
@@ -584,7 +585,7 @@ def calculate_luminosity_inner(self, model):
             * np.pi
             * const.sigma_sb.cgs
             * model.r_inner[0] ** 2
-            * model.t_inner ** 4
+            * model.t_inner**4
         ).to("erg/s")
 
     def calculate_time_of_simulation(self, model):
@@ -641,8 +642,8 @@ def from_config(
             config.spectrum.start.to("Hz", u.spectral()),
             num=config.spectrum.num + 1,
         )
-        running_mode=config.spectrum.integrated.compute.upper()
-        
+        running_mode = config.spectrum.integrated.compute.upper()
+
         if running_mode == "GPU":
             if cuda.is_available():
                 use_gpu = True
@@ -663,7 +664,7 @@ def from_config(
                 """An invalid option for compute was passed. The three
                 valid values are 'GPU', 'CPU', and 'Automatic'."""
             )
-        
+
         mc_config_module.disable_line_scattering = (
             config.plasma.disable_line_scattering
         )

diff --git a/tardis/montecarlo/montecarlo_numba/base.py b/tardis/montecarlo/montecarlo_numba/base.py
@@ -26,6 +26,7 @@
 from numba.typed import List
 from tardis.util.base import update_iterations_pbar, update_packet_pbar
 
+
 def montecarlo_radial1d(
     model,
     plasma,
@@ -100,22 +101,28 @@ def montecarlo_radial1d(
     runner.last_line_interaction_out_id = last_line_interaction_out_id
 
     if montecarlo_configuration.VPACKET_LOGGING and number_of_vpackets > 0:
-        runner.virt_packet_nus = np.concatenate(
-            virt_packet_nus).ravel()
+        runner.virt_packet_nus = np.concatenate(virt_packet_nus).ravel()
         runner.virt_packet_energies = np.concatenate(
-            virt_packet_energies).ravel()
+            virt_packet_energies
+        ).ravel()
         runner.virt_packet_initial_mus = np.concatenate(
-            virt_packet_initial_mus).ravel()
+            virt_packet_initial_mus
+        ).ravel()
         runner.virt_packet_initial_rs = np.concatenate(
-            virt_packet_initial_rs).ravel()
+            virt_packet_initial_rs
+        ).ravel()
         runner.virt_packet_last_interaction_in_nu = np.concatenate(
-            virt_packet_last_interaction_in_nu).ravel()
+            virt_packet_last_interaction_in_nu
+        ).ravel()
         runner.virt_packet_last_interaction_type = np.concatenate(
-            virt_packet_last_interaction_type).ravel()
+            virt_packet_last_interaction_type
+        ).ravel()
         runner.virt_packet_last_line_interaction_in_id = np.concatenate(
-            virt_packet_last_line_interaction_in_id).ravel()
+            virt_packet_last_line_interaction_in_id
+        ).ravel()
         runner.virt_packet_last_line_interaction_out_id = np.concatenate(
-            virt_packet_last_line_interaction_out_id).ravel()
+            virt_packet_last_line_interaction_out_id
+        ).ravel()
     update_iterations_pbar(1)
     # Condition for Checking if RPacket Tracking is enabled
     if montecarlo_configuration.RPACKET_TRACKING:
@@ -145,7 +152,7 @@ def montecarlo_main_loop(
     ----------
     packet_collection : PacketCollection
     numba_model : NumbaModel
-	numba_plasma : NumbaPlasma
+        numba_plasma : NumbaPlasma
     estimators : NumbaEstimators
     spectrum_frequency : astropy.units.Quantity
         frequency binspas
@@ -188,7 +195,7 @@ def montecarlo_main_loop(
             )
         )
         rpacket_trackers.append(RPacketTracker())
-        
+
     # Arrays for vpacket logging
     virt_packet_nus = []
     virt_packet_energies = []
@@ -233,7 +240,7 @@ def montecarlo_main_loop(
             numba_plasma,
             estimators,
             vpacket_collection,
-            rpacket_tracker
+            rpacket_tracker,
         )
 
         output_nus[i] = r_packet.nu

diff --git a/tardis/montecarlo/montecarlo_numba/calculate_distances.py b/tardis/montecarlo/montecarlo_numba/calculate_distances.py
@@ -15,7 +15,10 @@
 )
 
 from tardis.montecarlo.montecarlo_numba.utils import MonteCarloException
-from tardis.montecarlo.montecarlo_numba.r_packet import print_r_packet_properties
+from tardis.montecarlo.montecarlo_numba.r_packet import (
+    print_r_packet_properties,
+)
+
 
 @njit(**njit_dict_no_parallel)
 def calculate_distance_boundary(r, mu, r_inner, r_outer):
@@ -96,9 +99,7 @@ def calculate_distance_line(
     if nu_diff >= 0:
         distance = (nu_diff / nu) * C_SPEED_OF_LIGHT * time_explosion
     else:
-        raise MonteCarloException(
-            "nu difference is less than 0.0"
-        )
+        raise MonteCarloException("nu difference is less than 0.0")
 
     if nc.ENABLE_FULL_RELATIVITY:
         return calculate_distance_line_full_relativity(