Move montecarlo to transport and flatten monetcarlo numba

tardis/base.py

@@ -34,7 +34,7 @@ def run_tardis(
         the atomic data. Atomic data to use for this TARDIS simulation. If set to None (i.e. default),
         the atomic data will be loaded according to keywords set in the configuration
     packet_source : class, optional
-        A custom packet source class or a child class of `tardis.montecarlo.packet_source`
+        A custom packet source class or a child class of `tardis.transport.montecarlo.packet_source`
         used to override the TARDIS `BasePacketSource` class.
     simulation_callbacks : list of lists, default: `[]`, optional
         Set of callbacks to call at the end of every iteration of the Simulation.

tardis/energy_input/gamma_packet_loop.py

@@ -1,8 +1,8 @@
 import numpy as np
 from numba import njit
 
-from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel
-from tardis.montecarlo.montecarlo_numba.opacities import (
+from tardis.transport.montecarlo import njit_dict_no_parallel
+from tardis.transport.montecarlo.opacities import (
     compton_opacity_calculation,
     photoabsorption_opacity_calculation,
     pair_creation_opacity_calculation,

tardis/energy_input/gamma_ray_estimators.py

@@ -1,8 +1,8 @@
 import numpy as np
 from numba import njit
 
-from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel
-from tardis.montecarlo.montecarlo_numba.opacities import (
+from tardis.transport.montecarlo import njit_dict_no_parallel
+from tardis.transport.montecarlo.opacities import (
     compton_opacity_calculation,
     SIGMA_T,
     photoabsorption_opacity_calculation,

tardis/energy_input/gamma_ray_grid.py

@@ -1,7 +1,7 @@
 import numpy as np
 from numba import njit
 
-from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel
+from tardis.transport.montecarlo import njit_dict_no_parallel
 from tardis.energy_input.util import (
     doppler_factor_3d,
     solve_quadratic_equation,

tardis/energy_input/gamma_ray_interactions.py

@@ -1,8 +1,8 @@
 import numpy as np
 from numba import njit
 
-from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel
-from tardis.montecarlo.montecarlo_numba.opacities import (
+from tardis.transport.montecarlo import njit_dict_no_parallel
+from tardis.transport.montecarlo.opacities import (
     compton_opacity_partial,
     kappa_calculation,
 )

tardis/energy_input/gamma_ray_packet_source.py

@@ -14,7 +14,7 @@
     get_index,
     get_random_unit_vector,
 )
-from tardis.montecarlo.packet_source import BasePacketSource
+from tardis.transport.montecarlo.packet_source import BasePacketSource
 
 
 class RadioactivePacketSource(BasePacketSource):

tardis/energy_input/gamma_ray_transport.py

@@ -8,7 +8,7 @@
     get_all_isotopes,
     setup_input_energy,
 )
-from tardis.montecarlo.montecarlo_numba.opacities import M_P
+from tardis.transport.montecarlo.opacities import M_P
 
 # Energy: keV, exported as eV for SF solver
 # distance: cm

tardis/energy_input/samplers.py

@@ -1,7 +1,7 @@
 import numpy as np
 from numba import njit
 
-from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel
+from tardis.transport.montecarlo import njit_dict_no_parallel
 
 
 @njit(**njit_dict_no_parallel)

tardis/energy_input/tests/test_util.py

@@ -8,7 +8,7 @@
     klein_nishina,
     spherical_to_cartesian,
 )
-from tardis.montecarlo.montecarlo_numba.opacities import (
+from tardis.transport.montecarlo.opacities import (
     kappa_calculation,
 )
 

tardis/energy_input/util.py

@@ -3,8 +3,8 @@
 import numpy as np
 from numba import njit
 
-from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel
-from tardis.montecarlo.montecarlo_numba.opacities import kappa_calculation
+from tardis.transport.montecarlo import njit_dict_no_parallel
+from tardis.transport.montecarlo.opacities import kappa_calculation
 
 R_ELECTRON_SQUARED = (const.a0.cgs.value * const.alpha.cgs.value**2.0) ** 2.0
 ELECTRON_MASS_ENERGY_KEV = (const.m_e * const.c**2.0).to("keV").value

tardis/io/model_reader.py

@@ -12,8 +12,8 @@
 from radioactivedecay.utils import Z_DICT, elem_to_Z
 
 from tardis.io.configuration.config_reader import ConfigurationNameSpace
-from tardis.montecarlo.base import MonteCarloTransportSolver
-from tardis.montecarlo.packet_source import (
+from tardis.transport.montecarlo.base import MonteCarloTransportSolver
+from tardis.transport.montecarlo.packet_source import (
     BlackBodySimpleSource,
     BlackBodySimpleSourceRelativistic,
 )
@@ -541,7 +541,7 @@ def transport_to_dict(transport):
 
     Parameters
     ----------
-    transport : tardis.montecarlo.MontecarloTransport
+    transport : tardis.transport.montecarlo.MontecarloTransport
 
     Returns
     -------
@@ -619,7 +619,7 @@ def store_transport_to_hdf(transport, fname):
 
     Parameters
     ----------
-    transport : tardis.montecarlo.MontecarloTransport
+    transport : tardis.transport.montecarlo.MontecarloTransport
     filename : str
     """
     with h5py.File(fname, "a") as f:
@@ -670,7 +670,7 @@ def transport_from_hdf(fname):
 
     Returns
     -------
-    new_transport : tardis.montecarlo.MontecarloTransport
+    new_transport : tardis.transport.montecarlo.MontecarloTransport
     """
     d = {}
 

tardis/model/base.py

@@ -22,7 +22,7 @@
     parse_structure_config,
     parse_packet_source,
 )
-from tardis.montecarlo.packet_source import BlackBodySimpleSource
+from tardis.transport.montecarlo.packet_source import BlackBodySimpleSource
 from tardis.model.radiation_field_state import (
     DiluteBlackBodyRadiationFieldState,
 )

tardis/model/parse_input.py

@@ -20,7 +20,7 @@
 from tardis.model.radiation_field_state import (
     DiluteBlackBodyRadiationFieldState,
 )
-from tardis.montecarlo.packet_source import (
+from tardis.transport.montecarlo.packet_source import (
     BlackBodySimpleSource,
     BlackBodySimpleSourceRelativistic,
 )

tardis/plasma/properties/continuum_processes.py

@@ -5,11 +5,11 @@
 from numba import njit, prange
 
 from tardis import constants as const
-from tardis.montecarlo.estimators.util import (
+from tardis.transport.montecarlo.estimators.util import (
     bound_free_estimator_array2frame,
     integrate_array_by_blocks,
 )
-from tardis.montecarlo.montecarlo_numba import njit_dict
+from tardis.transport.montecarlo import njit_dict
 from tardis.plasma.exceptions import PlasmaException
 from tardis.plasma.properties.base import (
     Input,

tardis/plasma/properties/transition_probabilities.py

@@ -9,7 +9,7 @@
 from tardis.plasma.properties.continuum_processes import (
     get_ground_state_multi_index,
 )
-from tardis.montecarlo.montecarlo_numba.macro_atom import (
+from tardis.transport.montecarlo.macro_atom import (
     MacroAtomTransitionType,
 )
 

tardis/plasma/properties/util/integrate_array_by_blocks.py

@@ -1,7 +1,7 @@
 import numpy as np
 from numba import njit, prange
 
-from tardis.montecarlo.montecarlo_numba import njit_dict
+from tardis.transport.montecarlo import njit_dict
 
 
 @njit(**njit_dict)

tardis/plasma/properties/util/macro_atom.py

@@ -2,7 +2,7 @@
 from numba import njit
 
 from tardis import constants as const
-from tardis.montecarlo.montecarlo_numba import njit_dict
+from tardis.transport.montecarlo import njit_dict
 
 h_cgs = const.h.cgs.value
 c = const.c.to("cm/s").value

tardis/radiation_field/base.py

@@ -1,8 +1,8 @@
 import numpy as np
 from astropy import units as u
 
-from tardis.montecarlo.packet_source import BasePacketSource
-from tardis.montecarlo.montecarlo_numba.numba_interface import OpacityState
+from tardis.transport.montecarlo.packet_source import BasePacketSource
+from tardis.transport.montecarlo.numba_interface import OpacityState
 
 
 class MonteCarloRadiationFieldState:

tardis/simulation/base.py

@@ -15,7 +15,7 @@
 from tardis.io.util import HDFWriterMixin
 from tardis.model import SimulationState
 from tardis.model.parse_input import initialize_packet_source
-from tardis.montecarlo.base import MonteCarloTransportSolver
+from tardis.transport.montecarlo.base import MonteCarloTransportSolver
 from tardis.plasma.standard_plasmas import assemble_plasma
 from tardis.util.base import is_notebook
 from tardis.visualization import ConvergencePlots
@@ -95,7 +95,7 @@ class Simulation(PlasmaStateStorerMixin, HDFWriterMixin):
     iterations : int
     model : tardis.model.SimulationState
     plasma : tardis.plasma.BasePlasma
-    transport : tardis.montecarlo.MontecarloTransport
+    transport : tardis.transport.montecarlo.MontecarloTransport
     no_of_packets : int
     last_no_of_packets : int
     no_of_virtual_packets : int

tardis/transport/frame_transformations.py

@@ -2,11 +2,11 @@
 
 from numba import njit
 
-from tardis.montecarlo.montecarlo_numba import (
+from tardis.transport.montecarlo import (
     njit_dict_no_parallel,
 )
 
-from tardis.montecarlo.montecarlo_numba.numba_config import C_SPEED_OF_LIGHT
+from tardis.transport.montecarlo.numba_config import C_SPEED_OF_LIGHT
 
 
 @njit(**njit_dict_no_parallel)

tardis/transport/geometry/calculate_distances.py

@@ -2,19 +2,19 @@
 
 from numba import njit
 
-from tardis.montecarlo.montecarlo_numba import (
+from tardis.transport.montecarlo import (
     njit_dict_no_parallel,
 )
 
-from tardis.montecarlo.montecarlo_numba.numba_config import (
+from tardis.transport.montecarlo.numba_config import (
     C_SPEED_OF_LIGHT,
     MISS_DISTANCE,
     SIGMA_THOMSON,
     CLOSE_LINE_THRESHOLD,
 )
 
-from tardis.montecarlo.montecarlo_numba.utils import MonteCarloException
-from tardis.montecarlo.montecarlo_numba.r_packet import (
+from tardis.transport.montecarlo.utils import MonteCarloException
+from tardis.transport.montecarlo.r_packet import (
     print_r_packet_properties,
 )
 
@@ -75,7 +75,7 @@ def calculate_distance_line(
 
     Parameters
     ----------
-    r_packet : tardis.montecarlo.montecarlo_numba.r_packet.RPacket
+    r_packet : tardis.transport.montecarlo.r_packet.RPacket
     comov_nu : float
         comoving frequency at the CURRENT position of the RPacket
     is_last_line : bool

tardis/transport/montecarlo/__init__.py

@@ -0,0 +1,28 @@
+"""
+Faciliating the MonteCarlo iterations. 
+
+During a simulation run, a number of MonteCarlo iterations specified
+in the configuration are run using the numba compiler.
+Most of the iterations are used to calculate the steady-state plasma 
+properties and with the last iteration, the spectrum is determined.
+"""
+
+"""
+Implements the main loop of the MonteCarlo routine.
+"""
+
+from llvmlite import binding
+
+binding.set_option("tmp", "-non-global-value-max-name-size=2048")
+njit_dict = {"fastmath": True, "error_model": "numpy", "parallel": True}
+njit_dict_no_parallel = {
+    "fastmath": True,
+    "error_model": "numpy",
+    "parallel": False,
+}
+
+from tardis.transport.montecarlo.r_packet import RPacket
+from tardis.transport.montecarlo.montecarlo_main_loop import montecarlo_main_loop
+from tardis.transport.montecarlo.packet_collections import (
+    PacketCollection,
+)

tardis/montecarlo/base.py → tardis/transport/montecarlo/base.py

@@ -7,28 +7,28 @@
 from tardis import constants as const
 from tardis.io.logger import montecarlo_tracking as mc_tracker
 from tardis.io.util import HDFWriterMixin
-from tardis.montecarlo.estimators.radfield_mc_estimators import (
+from tardis.transport.montecarlo.estimators.radfield_mc_estimators import (
     initialize_estimator_statistics,
 )
-from tardis.montecarlo.montecarlo_configuration import (
+from tardis.transport.montecarlo.montecarlo_configuration import (
     MonteCarloConfiguration,
     configuration_initialize,
 )
-from tardis.montecarlo.montecarlo_numba import (
+from tardis.transport.montecarlo import (
     montecarlo_main_loop,
-    numba_config,
 )
-from tardis.montecarlo.montecarlo_numba.formal_integral import FormalIntegrator
-from tardis.montecarlo.montecarlo_numba.numba_interface import (
+from tardis.transport.montecarlo.formal_integral import FormalIntegrator
+from tardis.transport.montecarlo.numba_interface import (
     NumbaModel,
     opacity_state_initialize,
 )
-from tardis.montecarlo.montecarlo_numba.r_packet import (
+from tardis.transport.montecarlo.r_packet import (
     rpacket_trackers_to_dataframe,
 )
-from tardis.montecarlo.montecarlo_transport_state import (
+from tardis.transport.montecarlo.montecarlo_transport_state import (
     MonteCarloTransportState,
 )
+from tardis.transport.montecarlo import numba_config
 from tardis.util.base import (
     quantity_linspace,
     refresh_packet_pbar,

tardis/montecarlo/estimators/continuum_radfield_properties.py → tardis/transport/montecarlo/estimators/continuum_radfield_properties.py

@@ -9,10 +9,10 @@
 from tardis.model.radiation_field_state import (
     DiluteBlackBodyRadiationFieldState,
 )
-from tardis.montecarlo.estimators.radfield_mc_estimators import (
+from tardis.transport.montecarlo.estimators.radfield_mc_estimators import (
     RadiationFieldMCEstimators,
 )
-from tardis.montecarlo.estimators.util import (
+from tardis.transport.montecarlo.estimators.util import (
     bound_free_estimator_array2frame,
     integrate_array_by_blocks,
 )