Merge pull request tardis-sn#1318 from wkerzendorf/numba_montecarlo_p…

…rofile1

Numba montecarlo profile1

tardis/montecarlo/base.py

@@ -19,6 +19,7 @@
 from tardis.montecarlo.montecarlo_numba import montecarlo_logger as mc_logger
 from tardis.montecarlo.montecarlo_numba.numba_interface import (
     configuration_initialize)
+from tardis.montecarlo.montecarlo_numba import numba_config
 
 import numpy as np
 
@@ -61,7 +62,7 @@ class MontecarloRunner(HDFWriterMixin):
                                 (const.h / const.k_B)).cgs.value
 
     def __init__(self, seed, spectrum_frequency, virtual_spectrum_spawn_range,
-                 sigma_thomson, disable_electron_scattering, enable_reflective_inner_boundary,
+                 disable_electron_scattering, enable_reflective_inner_boundary,
                  enable_full_relativity, inner_boundary_albedo,
                  line_interaction_type, integrator_settings,
                  v_packet_settings, spectrum_method,
@@ -77,10 +78,10 @@ def __init__(self, seed, spectrum_frequency, virtual_spectrum_spawn_range,
         self.disable_electron_scattering = disable_electron_scattering
         self.spectrum_frequency = spectrum_frequency
         self.virtual_spectrum_spawn_range = virtual_spectrum_spawn_range
-        self.sigma_thomson = sigma_thomson
         self.enable_reflective_inner_boundary = enable_reflective_inner_boundary
         self.inner_boundary_albedo = inner_boundary_albedo
         self.enable_full_relativity = enable_full_relativity
+        numba_config.ENABLE_FULL_RELATIVITY = enable_full_relativity
         self.line_interaction_type = line_interaction_type
         self.single_packet_seed = single_packet_seed
         self.integrator_settings = integrator_settings
@@ -450,12 +451,14 @@ def from_config(cls, config, packet_source=None):
 
         """
         if config.plasma.disable_electron_scattering:
-            logger.warn('Disabling electron scattering - this is not physical')
-            sigma_thomson = 1e-200
+            logger.warn('Disabling electron scattering - this is not physical.'
+                        'Likely bug in formal integral - '
+                        'will not give same results.')
+            numba_config.SIGMA_THOMSON = 1e-200
             # mc_config_module.disable_electron_scattering = True
         else:
             logger.debug("Electron scattering switched on")
-            sigma_thomson = const.sigma_T.to('cm^2').value
+            numba_config.SIGMA_THOMSON = const.sigma_T.to('cm^2').value
             # mc_config_module.disable_electron_scattering = False
 
         spectrum_frequency = quantity_linspace(
@@ -467,7 +470,6 @@ def from_config(cls, config, packet_source=None):
         return cls(seed=config.montecarlo.seed,
                    spectrum_frequency=spectrum_frequency,
                    virtual_spectrum_spawn_range=config.montecarlo.virtual_spectrum_spawn_range,
-                   sigma_thomson=sigma_thomson,
                    enable_reflective_inner_boundary=config.montecarlo.enable_reflective_inner_boundary,
                    inner_boundary_albedo=config.montecarlo.inner_boundary_albedo,
                    enable_full_relativity=config.montecarlo.enable_full_relativity,

tardis/montecarlo/montecarlo_numba/base.py

@@ -32,17 +32,15 @@ def montecarlo_radial1d(model, plasma, runner):
 
     v_packets_energy_hist, last_interaction_type = montecarlo_main_loop(
         packet_collection, numba_model, numba_plasma, estimators,
-        runner.spectrum_frequency.value, number_of_vpackets, packet_seeds,
-        runner.sigma_thomson)
+        runner.spectrum_frequency.value, number_of_vpackets, packet_seeds)
     runner._montecarlo_virtual_luminosity.value[:] = v_packets_energy_hist
     runner.last_interaction_type = last_interaction_type
 
 
 @njit(**njit_dict, nogil=True)
 def montecarlo_main_loop(packet_collection, numba_model, numba_plasma,
                          estimators, spectrum_frequency,
-                         number_of_vpackets, packet_seeds,
-                         sigma_thomson):
+                         number_of_vpackets, packet_seeds):
     """
     This is the main loop of the MonteCarlo routine that generates packets 
     and sends them through the ejecta. 
@@ -80,7 +78,7 @@ def montecarlo_main_loop(packet_collection, numba_model, numba_plasma,
             number_of_vpackets,
             montecarlo_configuration.temporary_v_packet_bins)
         loop = single_packet_loop(r_packet, numba_model, numba_plasma, estimators,
-                           vpacket_collection, sigma_thomson)
+                           vpacket_collection)
         # if loop and 'stop' in loop:
         #     raise MonteCarloException
 

tardis/montecarlo/montecarlo_numba/interaction.py

@@ -9,8 +9,6 @@
     angle_aberration_CMF_to_LF, test_for_close_line)
 from tardis.montecarlo.montecarlo_numba.macro_atom import macro_atom
 
-CLOSE_LINE_THRESHOLD = 1e-7
-
 @njit(**njit_dict)
 def thomson_scatter(r_packet, time_explosion):
     """
@@ -46,32 +44,6 @@ def thomson_scatter(r_packet, time_explosion):
         )
 
 
-"""
-void
-montecarlo_thomson_scatter (rpacket_t * packet, storage_model_t * storage,
-                            double distance, rk_state *mt_state)
-{
-  move_packet (packet, storage, distance);
-  double doppler_factor = rpacket_doppler_factor (packet, storage);
-  double comov_nu = rpacket_get_nu (packet) * doppler_factor;
-  double comov_energy = rpacket_get_energy (packet) * doppler_factor;
-  rpacket_set_mu (packet, 2.0 * rk_double (mt_state) - 1.0);
-  double inverse_doppler_factor = rpacket_inverse_doppler_factor (packet, storage);
-  rpacket_set_nu (packet, comov_nu * inverse_doppler_factor);
-  rpacket_set_energy (packet, comov_energy * inverse_doppler_factor);
-  rpacket_reset_tau_event (packet, mt_state);
-  storage->last_interaction_type[rpacket_get_id (packet)] = 1;
-
-  angle_aberration_CMF_to_LF (packet, storage);
-
-  if (rpacket_get_virtual_packet_flag (packet) > 0)
-    {
-      create_vpacket (storage, packet, mt_state);
-    }
-}
-
-"""
-
 @njit(**njit_dict)
 def line_scatter(r_packet, time_explosion, line_interaction_type, numba_plasma):
     #increment_j_blue_estimator(packet, storage, distance, line2d_idx);

tardis/montecarlo/montecarlo_numba/numba_config.py

@@ -0,0 +1,8 @@
+from tardis import constants as const
+
+SIGMA_THOMSON = const.sigma_T.to('cm^2').value
+CLOSE_LINE_THRESHOLD = 1e-7
+C_SPEED_OF_LIGHT = const.c.to('cm/s').value
+MISS_DISTANCE = 1e99
+
+ENABLE_FULL_RELATIVITY = False

tardis/montecarlo/montecarlo_numba/numba_interface.py

@@ -1,6 +1,7 @@
 from enum import IntEnum
 
-from numba import float64, int64, jitclass, boolean
+from numba import float64, int64, boolean
+from numba.experimental import jitclass
 import numpy as np
 
 from astropy import units as u