Merge pull request #5 from chvogl/monte-carlo-continuum (#4)

* Reworked the continuum process selection function a little bit to work with the new continuum jitclass

* started addition of free-free emission handler

* updated interactions to work with new macroatom

* Added continuum process frquency samplers to the continuum jitclass constructor

* Fixed typo from merge conflict

* Fixed spacing

* small formatting updates to get the changes to track

* propagated continuum through single packet loop

* Cleaned up file, bound-free and free-free emission pathways should be in the correct place now

* added fixture for a new continuum object

* Added placeholder test for calculating the continuum opacities.  I'm weary to put in exact values until we're done

* reworked the contruction of the continuum class so that only a plasma object needs to be passed

* Updated interaction to now get the macro activation level from the continuum object

* refer to previous commit

* Added the continuum_process into the single packet loop so packets can now actually interact with the contniuum

* various small bugfixes to get the numba funcs to compile properly

* Removed extra functions

tardis/montecarlo/montecarlo_numba/base.py

@@ -59,8 +59,7 @@ def montecarlo_radial1d(model, plasma, runner):
     packet_seeds = montecarlo_configuration.packet_seeds
 
     number_of_vpackets = montecarlo_configuration.number_of_vpackets
-    ContinuumObject = create_continuum_class(plasma.chi_continuum_calculator)
-
+    ContinuumObject = create_continuum_class(plasma)
     (
         v_packets_energy_hist,
         last_interaction_type,
@@ -205,15 +204,11 @@ def montecarlo_main_loop(
         vpacket_collection = vpacket_collections[i]
 
         loop = single_packet_loop(
-
             r_packet,
             numba_model,
             numba_plasma,
             estimators,
             vpacket_collection,
-            continuum,
-
-            r_packet, numba_model, numba_plasma, estimators, vpacket_collection,
             continuum
         )
         # if loop and 'stop' in loop:

tardis/montecarlo/montecarlo_numba/interaction.py

@@ -1,4 +1,5 @@
 from numba import njit
+import numpy as np
 from tardis.montecarlo.montecarlo_numba import njit_dict, njit_dict_no_parallel
 from tardis.montecarlo.montecarlo_numba.numba_interface import (
     LineInteractionType,
@@ -16,8 +17,13 @@
     InteractionType,
 )
 from tardis.montecarlo.montecarlo_numba.utils import get_random_mu
-from tardis.montecarlo.montecarlo_numba.macro_atom import macro_atom
+from tardis.montecarlo.montecarlo_numba.macro_atom import (
+        macro_atom, MacroAtomTransitionType
+)
+
 
+
+@njit(**njit_dict_no_parallel)
 def scatter(r_packet, time_explosion):
 
     old_doppler_factor = get_doppler_factor(
@@ -34,7 +40,9 @@ def scatter(r_packet, time_explosion):
 
     return comov_nu, inverse_new_doppler_factor
 
-def continuum_event(r_packet, time_explosion, plasma, chi_continuum_calculator):
+# Maybe make the continuum selection a method of the continuum?
+@njit(**njit_dict_no_parallel)
+def continuum_event(r_packet, time_explosion, continuum, numba_plasma):
 
     comov_nu, inverse_new_doppler_factor = scatter(r_packet, time_explosion)
 
@@ -43,28 +51,83 @@ def continuum_event(r_packet, time_explosion, plasma, chi_continuum_calculator):
 
     zrand = np.random.random()
 
-    (
-        chi_bf,
-        chi_bf_contributions,
-        current_continua,
-        chi_ff,
-    ) = chi_continuum_calculator(comov_nu, r_packet.current_shell_id)
+    # Does this need to be re-calculated?
+    # continuum.calculate(comov_nu, r_packet.current_shell_id)
+
+    # Need to determine if a collisional process or not.  If not:
+
+    # This somehow determines the transition type and continuum id needed
+    # but does not sample new frequencies
+    # This reruns continuum.calculate
+    destination_level_idx = continuum.determine_macro_activation_idx(
+            comov_nu, r_packet.current_shell_id)
+
+    transition_id, transition_type = macro_atom(
+            destination_level_idx, 
+            r_packet.current_shell_id, 
+            numba_plasma
+            )
+
+    # Then use this to get a transition id from the macroatom 
+    if transition_type == MacroAtomTransitionType.FF_EMISSION: 
+        free_free_emission(r_packet, time_explosion, numba_plasma, continuum)
+    elif transition_type == MacroAtomTransitionType.BF_EMISSION:
+        bound_free_emission(r_packet, 
+                time_explosion, 
+                numba_plasma, 
+                continuum, 
+                transition_id)
+
+@njit(**njit_dict_no_parallel)
+def get_current_line_id(nu, numba_plasma):
+    '''
+    Get the next line id corresponding to a packet at frequency nu
+    '''
+
+    reverse_line_list = numba_plasma.line_list_nu[::-1]
+    number_of_lines = len(numba_plasma.line_list_nu)
+    line_id = number_of_lines - np.searchsorted(reverse_line_list, nu)
+    return line_id
+
+
+@njit(**njit_dict_no_parallel)
+def free_free_emission(r_packet, time_explosion, numba_plasma, continuum):
+
+    inverse_doppler_factor = get_inverse_doppler_factor(
+        r_packet.r, r_packet.mu, time_explosion
+    )
+    # Need to get the sampler into numba somehow
+    # maybe I'll update the numba_plasma?
+    comov_nu = continuum.sample_nu_free_free(r_packet.current_shell_id)
+    r_packet.nu = comov_nu * inverse_doppler_factor
+    current_line_id = get_current_line_id(r_packet.nu, numba_plasma) 
+    r_packet.next_line_id = current_line_id
+
+    if montecarlo_configuration.full_relativity:
+        r_packet.mu = angle_aberration_CMF_to_LF(
+            r_packet, time_explosion, r_packet.mu
+        )
 
-    cur_electron_density = numba_plasma.electron_density[
-        r_packet.current_shell_id
-    ]
+@njit(**njit_dict_no_parallel)
+def bound_free_emission(r_packet, time_explosion, numba_plasma, continuum, continuum_id):
 
-    chi_e = cur_electron_density * SIGMA_THOMSON
+    inverse_doppler_factor = get_inverse_doppler_factor(
+        r_packet.r, r_packet.mu, time_explosion
+    )
 
-    chi_nu = chi_bf + chi_ff + chi_e
+    comov_nu = continuum.sample_nu_free_bound(
+            r_packet.current_shell_id, 
+            continuum_id)
+
+    r_packet.nu = comov_nu * inverse_doppler_factor
+    current_line_id = get_current_line_id(r_packet.nu, numba_plasma)
+    r_packet.next_line_id = current_line_id
 
+    if montecarlo_configuration.full_relativity:
+        r_packet.mu = angle_aberration_CMF_to_LF(
+            r_packet, time_explosion, r_packet.mu
+        )
 
-    if zrand < SIGMA_THOMSON / chi_nu:
-        thomson_scatter(r_packet, time_explosion)
-    elif zrand < (SIGMA_THOMSON + chi_bf) / chi_nu:
-        bound_free_absorption(r_packet, time_explosion, plasma)
-    else:
-        free_free_absorption(r_packet, time_explosion)
 
 
 @njit(**njit_dict_no_parallel)

tardis/montecarlo/montecarlo_numba/numba_interface.py

@@ -318,8 +318,13 @@ def set_properties(
 
 
 
-def create_continuum_class(chi_continuum_calculator):
+def create_continuum_class(plasma):
     """called before mainloop"""
+
+    chi_continuum_calculator = plasma.chi_continuum_calculator
+    nu_fb_sampler = plasma.nu_fb_sampler
+    nu_ff_sampler = plasma.nu_ff_sampler
+    get_macro_activation_idx = plasma.determine_continuum_macro_activation_idx
     continuum_spec = [
             ("chi_bf_tot", float64),
             ("chi_bf_contributions", float64[:]),
@@ -348,6 +353,24 @@ def calculate(self, nu, shell):
             self.chi_ff,
             ) = chi_continuum_calculator(nu, shell)
 
+        def sample_nu_free_bound(self, shell, continuum_id):
+
+            return nu_fb_sampler(shell, continuum_id)
+
+        def sample_nu_free_free(self, shell):
+
+            return nu_ff_sampler(shell)
+
+        def determine_macro_activation_idx(self, nu, shell):
+
+            self.calculate(nu, shell) # maybe don't do this here
+            idx = get_macro_activation_idx(
+                    nu, self.chi_bf_tot, self.chi_ff, 
+                    self.chi_bf_contributions, self.current_continua
+                    )
+            return idx
+
+
     @njit(**njit_dict_no_parallel)
     def continuum_constructor():
         return Continuum()

tardis/montecarlo/montecarlo_numba/r_packet.py

@@ -161,17 +161,6 @@ def trace_packet(
     chi_continuum = chi_e + chi_bf + chi_ff
     distance_continuum = tau_event / chi_continuum
 
-    (
-        chi_bf,
-        chi_bf_contributions,
-        current_continua,
-        x_sect_bfs,
-        chi_ff,
-    ) = chi_continuum_calculator(comov_nu, r_packet.current_shell_id)
-
-    chi_continuum = chi_e + chi_bf + chi_ff
-    distance_continuum = tau_event / chi_continuum
-
     cur_line_id = start_line_id  # initializing varibale for Numba
     # - do not remove
     last_line_id = len(numba_plasma.line_list_nu) - 1
@@ -210,6 +199,7 @@ def trace_packet(
         distance = min(distance_trace, distance_boundary, distance_continuum)
 
         if distance_trace != 0:
+
             if distance == distance_boundary:
                 interaction_type = InteractionType.BOUNDARY  # BOUNDARY
                 r_packet.next_line_id = cur_line_id

tardis/montecarlo/montecarlo_numba/single_packet_loop.py

@@ -17,6 +17,7 @@
     InteractionType,
     thomson_scatter,
     line_scatter,
+    continuum_event,
 )
 from tardis.montecarlo.montecarlo_numba.numba_interface import (
     LineInteractionType,
@@ -38,7 +39,7 @@
 @njit
 def single_packet_loop(
     r_packet, numba_model, numba_plasma, estimators, vpacket_collection,
-    chi_bf_interpolator
+    continuum
 ):
     """
     Parameters
@@ -78,7 +79,7 @@ def single_packet_loop(
     while r_packet.status == PacketStatus.IN_PROCESS:
         distance, interaction_type, delta_shell = trace_packet(
             r_packet, numba_model, numba_plasma,
-            estimators, chi_bf_interpolator
+            estimators, continuum
         )
 
         if interaction_type == InteractionType.BOUNDARY:
@@ -116,6 +117,19 @@ def single_packet_loop(
             trace_vpacket_volley(
                 r_packet, vpacket_collection, numba_model, numba_plasma
             )
+        elif interaction_type == InteractionType.CONTINUUM_PROCESS:
+            r_packet.last_interaction_type = InteractionType.CONTINUUM_PROCESS
+            move_r_packet(
+                r_packet, distance, numba_model.time_explosion, estimators
+            )
+            continuum_event(r_packet, numba_model.time_explosion,
+                    continuum, numba_plasma)
+
+            trace_vpacket_volley(
+                r_packet, vpacket_collection, numba_model, numba_plasma
+            )
+
+
         if mc_tracker.DEBUG_MODE:
             r_packet_track_nu.append(r_packet.nu)
             r_packet_track_mu.append(r_packet.mu)

tardis/montecarlo/montecarlo_numba/tests/conftest.py

@@ -24,6 +24,17 @@ def nb_simulation_verysimple(config_verysimple, atomic_dataset):
     return sim
 
 
+@pytest.fixture(scope="package")
+def verysimple_continuum(nu_simulation_verysimple):
+    plasma = nu_simulation_verysimple
+    numba_plasma = numba_interface.numba_plasma_initialize(plasma, "macroatom")
+
+    ContinuumClass = create_continuum_class(plasma)
+
+    continuum = ContinuumClass()
+    return continuum
+
+
 @pytest.fixture()
 def verysimple_collection(nb_simulation_verysimple):
     runner = nb_simulation_verysimple.runner

tardis/montecarlo/montecarlo_numba/tests/test_numba_interface.py

@@ -4,6 +4,47 @@
 import numpy as np
 
 
+@pytest.mark.parametrize(
+       ["current_shell_id", "nu"],
+       [(0, 0.6), (1, 0.4)]
+       ) 
+def test_coninuum_opacities(
+    verysimple_continuum, current_shell_id, nu):
+
+    continuum = verysimple_continuum
+    continuum.calculate(nu, current_shell_id)
+    print(current_shell_id, nu)
+    print(continuum.chi_bf_tot)
+    print(continuum.chi_bf_contributions)
+    print(continuum.current_continuua)
+    print(continuum.x_sect_bfs)
+    print(continuum.chi_ff)
+
+
+@pytest.mark.parametrize(
+       "current_shell_id",
+       [0, 1, 2]
+       ) 
+def test_continuum_free_free_sampler(
+        verysimple_continuum, current_shell_id):
+
+    continuum = verysimple_continuum
+    nu = continuum.sample_nu_free_free(current_shell_id)
+    print(nu)
+
+@pytest.mark.parametrize(
+       ["current_shell_id", "continuum_id"],
+       [(0, 0), (1, 0)]
+       ) 
+def test_coninuum_opacities(
+    verysimple_continuum, current_shell_id, continuum_id):
+
+    continuum = verysimple_continuum
+    nu = continuum.sample_nu_free_bound(current_shell_id, continuum_id)
+    print(nu)
+
+
+
 @pytest.mark.parametrize("input_params", ["scatter", "macroatom", "downbranch"])
 def test_numba_plasma_initialize(nb_simulation_verysimple, input_params):
     line_interaction_type = input_params

tardis/plasma/properties/continuum_processes.py

@@ -1029,7 +1029,7 @@ def nu_fb(shell, continuum_id):
             zrand = np.random.random()
             idx = np.searchsorted(em[start:end], zrand, side="right")
 
-            return phot_nus[idx] - (em[idx] - zrand) / (em[idx] - em[idx-1]) * (phot_nus[idx] - phot_nus[idx-1])
+            return phot_nus[idx] - (em[idx] - zrand) / (em[idx] - em[idx - 1]) * (phot_nus[idx] - phot_nus[idx - 1])
 
         return nu_fb
 

tardis/plasma/properties/property_collections.py

@@ -93,7 +93,7 @@ class PlasmaPropertyCollection(list):
 )
 continuum_interaction_properties = PlasmaPropertyCollection(
     [
-        MacroAtomData,  # TODO: Remove this once everything works
+        MacroAtomData, # TODO: Remove this once everything works
         PhotoIonizationData,
         SpontRecombRateCoeff,
         PhotoIonRateCoeff,