Final numba PR (#1285)

* Add montecarlo_seed value to global module

* Add user-input seed to each seed setting

* Make sure montecarlo_seed for global module is set during configuration

* Include packet_seeds within montecarlo_configuration module

* Initialize packet seeds

* Set user seed before instantiating *any* random process!

* Pass packet seeds to montecarlo_main_loop; use said seeds

* Bring rpacket creation outside of if/else

* Initialize seeds within initialize_packets

* Set mu and nu sampling on a per-seed basis

* Use single_packet_seed to index existing packet_seeds array

* Perform packet-by-packet creation of trajectories, nu in packet_source

* Try initializing xis outside of blackbody function

* Remove print statements

* Do not concatenate arrays in packet_source running; add njit

* Name create_array more explicitly

* Make montecarlo seeds attributes of r_packet

* Add rationale for size of random array

* Clean up documentation of blackbody sampling

* Remove extraneous numba import

* Remove extra kwarg in basepacket create_packets

* Add packet_source docstrings

* Justify and change MAX_SEED_VAL

* Refactor packet_source so that all packet properties are made in the same func

* Attach seed to r_packet

* Clean up create_packets func

* Clean up docstrings for packet_source

* @jitclass BlackBodySimpleSource

* Ensure that the iteration is added to the seed each time montecarlorunner is run

* jitclass uniform packets

* Pass iteration to montecarlorunner

* Remove njit from uniform packet energies

* Provide a better-motivated upper bound for the seed value

* Use random.sample instead of np.random.randint to prevent possible repetitions

* Raise more explicit error if too many packets

* Set seed of random module before calling random.sample

* Address merge error in montecarlo_main_loop

* Clarify ValueError for seeds

* Do not generate packet properties on a seed-by-seed basis

* Remove seeds from create_packets function call

* use new random number generator to make packet seeds

* Bump numpy version to 1.19.0

* Pass rng to packet source functions for randomness

* actually define rng in montecarlo/base

* Delete error raised for too many packets

* Allow packet seeds to overlap

* Remove random module import

* Clarify C comparison

* Include numba comparison plots

* Add Testing TARDIS section to docs

* Make sure custom packet source works

* Add disabling to yml file

* Make changes to mc_config_module when necessary

* Add scatter-disabling functionality to global module

* Set tau_sobolev to 0 in plasma setting

* Add rudimentary last_interaction_type

* Address combined > event for tau_trace

* reference module for sigma_thomson

* Explicitly pass sigma_thomson

* Only reset nu for e-scatter

* Split thomson_scatter and line_scatter entirely

* check doppler factor

* Account for interaction type

* Do not hardcode the array size anymore

* Remove extra commented statments

* Remove reference to flag

docs/development/index.rst

@@ -11,7 +11,6 @@ to the Astropy team for designing it.
 .. toctree::
     :maxdepth: 2
 
-    running_tests
     issues
     debug_numba
 

docs/index.rst

@@ -59,6 +59,14 @@ assistance.
     physics/plasma/index
     physics/montecarlo/index
 
+.. toctree::
+    :maxdepth: 2
+    :caption: Testing TARDIS
+    :hidden:
+
+    testing/automated_tests
+    testing/numba
+
 
 .. toctree::
     :maxdepth: 2

docs/running/custom_source.ipynb

@@ -52,9 +52,10 @@
     "    \n",
     "    def __init__(self, seed, truncation_wavelength):\n",
     "        super().__init__(seed)\n",
+    "        self.rng = np.random.default_rng(seed=seed)\n",
     "        self.truncation_wavelength = truncation_wavelength\n",
     "    \n",
-    "    def create_packets(self, T, no_of_packets,\n",
+    "    def create_packets(self, T, no_of_packets, rng,\n",
     "                       drawing_sample_size=None):\n",
     "        \"\"\"\n",
     "        Packet source that generates a truncated Blackbody source.\n",
@@ -70,11 +71,10 @@
     "            Only wavelengths higher than the truncation wavelength\n",
     "            will be sampled.\n",
     "        \"\"\"\n",
-    "        \n",
-    "        \n",
+    "\n",
     "        # Use mus and energies from normal blackbody source.\n",
-    "        mus = self.create_zero_limb_darkening_packet_mus(no_of_packets)\n",
-    "        energies = self.create_uniform_packet_energies(no_of_packets)\n",
+    "        mus = self.create_zero_limb_darkening_packet_mus(no_of_packets, self.rng)\n",
+    "        energies = self.create_uniform_packet_energies(no_of_packets, self.rng)\n",
     "\n",
     "        # If not specified, draw 2 times as many packets and reject any beyond no_of_packets.\n",
     "        if drawing_sample_size is None:\n",
@@ -86,15 +86,15 @@
     "        \n",
     "        # Draw nus from blackbody distribution and reject based on truncation_frequency.\n",
     "        # If more nus.shape[0] > no_of_packets use only the first no_of_packets.\n",
-    "        nus = self.create_blackbody_packet_nus(T, drawing_sample_size)\n",
+    "        nus = self.create_blackbody_packet_nus(T, drawing_sample_size, self.rng)\n",
     "        nus = nus[nus<truncation_frequency][:no_of_packets]\n",
     "        \n",
     "        \n",
     "        # Only required if the truncation wavelength is too big compared to the maximum \n",
     "        # of the blackbody distribution. Keep sampling until nus.shape[0] > no_of_packets.\n",
     "        while nus.shape[0] < no_of_packets:\n",
     "            additional_nus = self.create_blackbody_packet_nus(\n",
-    "                T, drawing_sample_size\n",
+    "                T, drawing_sample_size, self.rng\n",
     "            )\n",
     "            mask = additional_nus < truncation_frequency\n",
     "            additional_nus = additional_nus[mask][:no_of_packets]\n",
@@ -170,4 +170,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

tardis/io/schemas/plasma.yml

@@ -17,6 +17,11 @@ properties:
     default: false
     description: disable electron scattering process in montecarlo part - non-physical
       only for tests
+  disable_line_scattering:
+    type: boolean
+    default: false
+    description: disable line scattering process in montecarlo part - non-physical
+      only for tests
   ionization:
     type: string
     enum:

tardis/montecarlo/base.py

@@ -28,6 +28,12 @@
 logger = logging.getLogger(__name__)
 TARDIS_PATH = TARDIS_PATH[0]
 
+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
+# seed val to the maximum allowed by numpy.
+
 class MontecarloRunner(HDFWriterMixin):
     """
     This class is designed as an interface between the Python part and the
@@ -61,7 +67,7 @@ class MontecarloRunner(HDFWriterMixin):
                                 (const.h / const.k_B)).cgs.value
 
     def __init__(self, seed, spectrum_frequency, virtual_spectrum_range,
-                 sigma_thomson, enable_reflective_inner_boundary,
+                 sigma_thomson, disable_electron_scattering, enable_reflective_inner_boundary,
                  enable_full_relativity, inner_boundary_albedo,
                  line_interaction_type, integrator_settings,
                  v_packet_settings, spectrum_method,
@@ -74,6 +80,7 @@ def __init__(self, seed, spectrum_frequency, virtual_spectrum_range,
         else:
             self.packet_source = packet_source
         # inject different packets
+        self.disable_electron_scattering = disable_electron_scattering
         self.spectrum_frequency = spectrum_frequency
         self.virtual_spectrum_range = virtual_spectrum_range
         self.sigma_thomson = sigma_thomson
@@ -85,6 +92,7 @@ def __init__(self, seed, spectrum_frequency, virtual_spectrum_range,
         self.integrator_settings = integrator_settings
         self.v_packet_settings = v_packet_settings
         self.spectrum_method = spectrum_method
+        self.seed = seed
         self._integrator = None
         self._spectrum_integrated = None
 
@@ -127,11 +135,23 @@ def _initialize_geometry_arrays(self, model):
         self.r_outer_cgs = model.r_outer.to('cm').value
         self.v_inner_cgs = model.v_inner.to('cm/s').value
 
-    def _initialize_packets(self, T, no_of_packets):
+    def _initialize_packets(self, T, no_of_packets, iteration):
+        # 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
+        # because we call random.sample, which references a different internal
+        # state than in the numpy.random module.
+        seed = self.seed + iteration
+        rng = np.random.default_rng(seed=seed)
+        seeds = rng.choice(MAX_SEED_VAL,
+                           no_of_packets,
+                           replace=True
+                           )
         nus, mus, energies = self.packet_source.create_packets(
                 T,
-                no_of_packets
-                )
+                no_of_packets,
+                rng)
+        mc_config_module.packet_seeds = seeds
         self.input_nu = nus
         self.input_mu = mus
         self.input_energy = energies
@@ -206,7 +226,7 @@ def integrator(self):
 
     def run(self, model, plasma, no_of_packets,
             no_of_virtual_packets=0, nthreads=1,
-            last_run=False):
+            last_run=False, iteration=0):
         """
         Run the montecarlo calculation
 
@@ -236,7 +256,7 @@ def run(self, model, plasma, no_of_packets,
         self._initialize_geometry_arrays(model)
 
         self._initialize_packets(model.t_inner.value,
-                                 no_of_packets)
+                                 no_of_packets, iteration)
 
         montecarlo_configuration = configuration_initialize(self,
                                                             no_of_virtual_packets)
@@ -439,15 +459,18 @@ 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 * (u.cm ** 2)
+            sigma_thomson = 1e-200
+            # mc_config_module.disable_electron_scattering = True
         else:
             logger.debug("Electron scattering switched on")
-            sigma_thomson = const.sigma_T.cgs
+            sigma_thomson = const.sigma_T.to('cm^2').value
+            # mc_config_module.disable_electron_scattering = False
 
         spectrum_frequency = quantity_linspace(
             config.spectrum.stop.to('Hz', u.spectral()),
             config.spectrum.start.to('Hz', u.spectral()),
             num=config.spectrum.num + 1)
+        mc_config_module.disable_line_scattering = config.plasma.disable_line_scattering
 
         return cls(seed=config.montecarlo.seed,
                    spectrum_frequency=spectrum_frequency,
@@ -460,6 +483,7 @@ def from_config(cls, config, packet_source=None):
                    integrator_settings=config.spectrum.integrated,
                    v_packet_settings=config.spectrum.virtual,
                    spectrum_method=config.spectrum.method,
+                   disable_electron_scattering=config.plasma.disable_electron_scattering,
                    packet_source=packet_source,
                    debug_packets=config.montecarlo.debug_packets,
                    logger_buffer=config.montecarlo.logger_buffer,

tardis/montecarlo/montecarlo_configuration.py

@@ -1,5 +1,11 @@
+from tardis import constants as const
+
 full_relativity = True
 single_packet_seed = -1
 temporary_v_packet_bins = None
 number_of_vpackets = 0
-line_interaction_type = None
+montecarlo_seed = 0
+line_interaction_type = None
+packet_seeds = []
+disable_electron_scattering = False
+disable_line_scattering = False