Merge f9c2645 into 888575d

tardis/montecarlo/montecarlo_numba/base.py

@@ -5,8 +5,9 @@
 from tardis.montecarlo.montecarlo_numba.r_packet import (
     RPacket,
     PacketStatus,
-    MonteCarloException,
 )
+from tardis.montecarlo.montecarlo_numba.utils import MonteCarloException
+
 from tardis.montecarlo.montecarlo_numba.numba_interface import (
     PacketCollection,
     VPacketCollection,

tardis/montecarlo/montecarlo_numba/calculate_distances.py

@@ -0,0 +1,150 @@
+import math
+
+from numba import njit
+
+from tardis.montecarlo.montecarlo_numba import (
+    njit_dict_no_parallel,
+)
+
+from tardis.montecarlo.montecarlo_numba.numba_config import (
+    C_SPEED_OF_LIGHT,
+    MISS_DISTANCE,
+    SIGMA_THOMSON,
+    ENABLE_FULL_RELATIVITY,
+)
+
+from tardis.montecarlo.montecarlo_numba.utils import MonteCarloException
+
+
+@njit(**njit_dict_no_parallel)
+def calculate_distance_boundary(r, mu, r_inner, r_outer):
+    """
+    Calculate distance to shell boundary in cm.
+
+    Parameters
+    ----------
+    r : float
+       radial coordinate of the RPacket
+    mu : float
+       cosine of the direction of movement
+    r_inner : float
+       inner radius of current shell
+    r_outer : float
+       outer radius of current shell
+    """
+
+    delta_shell = 0
+    if mu > 0.0:
+        # direction outward
+        distance = math.sqrt(r_outer * r_outer + ((mu * mu - 1.0) * r * r)) - (
+            r * mu
+        )
+        delta_shell = 1
+    else:
+        # going inward
+        check = r_inner * r_inner + (r * r * (mu * mu - 1.0))
+
+        if check >= 0.0:
+            # hit inner boundary
+            distance = -r * mu - math.sqrt(check)
+            delta_shell = -1
+        else:
+            # miss inner boundary
+            distance = math.sqrt(
+                r_outer * r_outer + ((mu * mu - 1.0) * r * r)
+            ) - (r * mu)
+            delta_shell = 1
+
+    return distance, delta_shell
+
+
+# @log_decorator
+#'float64(RPacket, float64, int64, float64, float64)'
+@njit(**njit_dict_no_parallel)
+def calculate_distance_line(
+    r_packet, comov_nu, is_last_line, nu_line, time_explosion
+):
+    """
+    Calculate distance until RPacket is in resonance with the next line
+
+    Parameters
+    ----------
+    r_packet : tardis.montecarlo.montecarlo_numba.r_packet.RPacket
+    comov_nu : float
+        comoving frequency at the CURRENT position of the RPacket
+    is_last_line : bool
+        return MISS_DISTANCE if at the end of the line list
+    nu_line : float
+        line to check the distance to
+    time_explosion : float
+        time since explosion in seconds
+
+    Returns
+    -------
+    """
+
+    nu = r_packet.nu
+
+    if is_last_line:
+        return MISS_DISTANCE
+
+    nu_diff = comov_nu - nu_line
+
+    # for numerical reasons, if line is too close, we set the distance to 0.
+    if r_packet.is_close_line:
+        nu_diff = 0.0
+        r_packet.is_close_line = False
+
+    if nu_diff >= 0:
+        distance = (nu_diff / nu) * C_SPEED_OF_LIGHT * time_explosion
+    else:
+        print("WARNING: nu difference is less than 0.0")
+        raise MonteCarloException(
+            "nu difference is less than 0.0; for more"
+            " information, see print statement beforehand"
+        )
+
+    if ENABLE_FULL_RELATIVITY:
+        return calculate_distance_line_full_relativity(
+            nu_line, nu, time_explosion, r_packet
+        )
+    return distance
+
+
+@njit(**njit_dict_no_parallel)
+def calculate_distance_line_full_relativity(
+    nu_line, nu, time_explosion, r_packet
+):
+    # distance = - mu * r + (ct - nu_r * nu_r * sqrt(ct * ct - (1 + r * r * (1 - mu * mu) * (1 + pow(nu_r, -2))))) / (1 + nu_r * nu_r);
+    nu_r = nu_line / nu
+    ct = C_SPEED_OF_LIGHT * time_explosion
+    distance = -r_packet.mu * r_packet.r + (
+        ct
+        - nu_r
+        * nu_r
+        * math.sqrt(
+            ct * ct
+            - (
+                1
+                + r_packet.r
+                * r_packet.r
+                * (1 - r_packet.mu * r_packet.mu)
+                * (1 + 1.0 / (nu_r * nu_r))
+            )
+        )
+    ) / (1 + nu_r * nu_r)
+    return distance
+
+
+@njit(**njit_dict_no_parallel)
+def calculate_distance_electron(electron_density, tau_event):
+    """
+    Calculate distance to Thomson Scattering
+
+    Parameters
+    ----------
+    electron_density : float
+    tau_event : float
+    """
+    # add full_relativity here
+    return tau_event / (electron_density * SIGMA_THOMSON)

tardis/montecarlo/montecarlo_numba/estimators.py

@@ -0,0 +1,76 @@
+from numba import njit
+
+from tardis.montecarlo.montecarlo_numba import (
+    njit_dict_no_parallel,
+)
+
+from tardis.montecarlo.montecarlo_numba.frame_transformations import (
+    calc_packet_energy,
+    calc_packet_energy_full_relativity,
+)
+
+from tardis.montecarlo.montecarlo_numba.numba_config import (
+    ENABLE_FULL_RELATIVITY,
+)
+
+
+@njit(**njit_dict_no_parallel)
+def set_estimators(r_packet, distance, numba_estimator, comov_nu, comov_energy):
+    """
+    Updating the estimators
+    """
+    numba_estimator.j_estimator[r_packet.current_shell_id] += (
+        comov_energy * distance
+    )
+    numba_estimator.nu_bar_estimator[r_packet.current_shell_id] += (
+        comov_energy * distance * comov_nu
+    )
+
+
+@njit(**njit_dict_no_parallel)
+def set_estimators_full_relativity(
+    r_packet, distance, numba_estimator, comov_nu, comov_energy, doppler_factor
+):
+    numba_estimator.j_estimator[r_packet.current_shell_id] += (
+        comov_energy * distance * doppler_factor
+    )
+    numba_estimator.nu_bar_estimator[r_packet.current_shell_id] += (
+        comov_energy * distance * comov_nu * doppler_factor
+    )
+
+
+@njit(**njit_dict_no_parallel)
+def update_line_estimators(
+    estimators, r_packet, cur_line_id, distance_trace, time_explosion
+):
+    """
+    Function to update the line estimators
+
+    Parameters
+    ----------
+    estimators : tardis.montecarlo.montecarlo_numba.numba_interface.Estimators
+    r_packet : tardis.montecarlo.montecarlo_numba.r_packet.RPacket
+    cur_line_id : int
+    distance_trace : float
+    time_explosion : float
+    """
+
+    """ Actual calculation - simplified below
+    r_interaction = math.sqrt(r_packet.r**2 + distance_trace**2 +
+                            2 * r_packet.r * distance_trace * r_packet.mu)
+    mu_interaction = (r_packet.mu * r_packet.r + distance_trace) / r_interaction
+    doppler_factor = 1.0 - mu_interaction * r_interaction /
+    ( time_explosion * C)
+    """
+
+    if not ENABLE_FULL_RELATIVITY:
+        energy = calc_packet_energy(r_packet, distance_trace, time_explosion)
+    else:
+        energy = calc_packet_energy_full_relativity(r_packet)
+
+    estimators.j_blue_estimator[cur_line_id, r_packet.current_shell_id] += (
+        energy / r_packet.nu
+    )
+    estimators.Edotlu_estimator[
+        cur_line_id, r_packet.current_shell_id
+    ] += energy

tardis/montecarlo/montecarlo_numba/frame_transformations.py

@@ -0,0 +1,103 @@
+import math
+
+from numba import njit
+
+from tardis.montecarlo.montecarlo_numba import (
+    njit_dict_no_parallel,
+)
+
+from tardis.montecarlo.montecarlo_numba.numba_config import (
+    C_SPEED_OF_LIGHT,
+    ENABLE_FULL_RELATIVITY,
+)
+
+
+@njit(**njit_dict_no_parallel)
+def get_doppler_factor(r, mu, time_explosion):
+    inv_c = 1 / C_SPEED_OF_LIGHT
+    inv_t = 1 / time_explosion
+    beta = r * inv_t * inv_c
+    if not ENABLE_FULL_RELATIVITY:
+        return get_doppler_factor_partial_relativity(mu, beta)
+    else:
+        return get_doppler_factor_full_relativity(mu, beta)
+
+
+@njit(**njit_dict_no_parallel)
+def get_doppler_factor_partial_relativity(mu, beta):
+    return 1.0 - mu * beta
+
+
+@njit(**njit_dict_no_parallel)
+def get_doppler_factor_full_relativity(mu, beta):
+    return (1.0 - mu * beta) / math.sqrt(1 - beta * beta)
+
+
+@njit(**njit_dict_no_parallel)
+def get_inverse_doppler_factor(r, mu, time_explosion):
+    """
+    Calculate doppler factor for frame transformation
+
+    Parameters
+    ----------
+    r : float
+    mu : float
+    time_explosion : float
+    """
+    inv_c = 1 / C_SPEED_OF_LIGHT
+    inv_t = 1 / time_explosion
+    beta = r * inv_t * inv_c
+    if not ENABLE_FULL_RELATIVITY:
+        return get_inverse_doppler_factor_partial_relativity(mu, beta)
+    else:
+        return get_inverse_doppler_factor_full_relativity(mu, beta)
+
+
+@njit(**njit_dict_no_parallel)
+def get_inverse_doppler_factor_partial_relativity(mu, beta):
+    return 1.0 / (1.0 - mu * beta)
+
+
+@njit(**njit_dict_no_parallel)
+def get_inverse_doppler_factor_full_relativity(mu, beta):
+    return (1.0 + mu * beta) / math.sqrt(1 - beta * beta)
+
+
+@njit(**njit_dict_no_parallel)
+def calc_packet_energy_full_relativity(r_packet):
+    # accurate to 1 / gamma - according to C. Vogl
+    return r_packet.energy
+
+
+@njit(**njit_dict_no_parallel)
+def calc_packet_energy(r_packet, distance_trace, time_explosion):
+    doppler_factor = 1.0 - (
+        (distance_trace + r_packet.mu * r_packet.r)
+        / (time_explosion * C_SPEED_OF_LIGHT)
+    )
+    energy = r_packet.energy * doppler_factor
+    return energy
+
+
+@njit(**njit_dict_no_parallel)
+def angle_aberration_CMF_to_LF(r_packet, time_explosion, mu):
+    """
+    Converts angle aberration from comoving frame to
+    laboratory frame.
+    """
+    ct = C_SPEED_OF_LIGHT * time_explosion
+    beta = r_packet.r / (ct)
+    return (r_packet.mu + beta) / (1.0 + beta * mu)
+
+
+@njit(**njit_dict_no_parallel)
+def angle_aberration_LF_to_CMF(r_packet, time_explosion, mu):
+    """
+
+    c code:
+    double beta = rpacket_get_r (packet) * storage->inverse_time_explosion * INVERSE_C;
+    return (mu - beta) / (1.0 - beta * mu);
+    """
+    ct = C_SPEED_OF_LIGHT * time_explosion
+    beta = r_packet.r / (ct)
+    return (mu - beta) / (1.0 - beta * mu)

tardis/montecarlo/montecarlo_numba/interaction.py

@@ -1,19 +1,22 @@
 from numba import njit
-from tardis.montecarlo.montecarlo_numba import njit_dict,njit_dict_no_parallel
+from tardis.montecarlo.montecarlo_numba import njit_dict, njit_dict_no_parallel
 from tardis.montecarlo.montecarlo_numba.numba_interface import (
     LineInteractionType,
 )
 
 from tardis.montecarlo import (
     montecarlo_configuration as montecarlo_configuration,
 )
-from tardis.montecarlo.montecarlo_numba.r_packet import (
+from tardis.montecarlo.montecarlo_numba.frame_transformations import (
     get_doppler_factor,
     get_inverse_doppler_factor,
-    get_random_mu,
     angle_aberration_CMF_to_LF,
+)
+from tardis.montecarlo.montecarlo_numba.r_packet import (
     test_for_close_line,
+    InteractionType,
 )
+from tardis.montecarlo.montecarlo_numba.utils import get_random_mu
 from tardis.montecarlo.montecarlo_numba.macro_atom import macro_atom
 
 
@@ -54,7 +57,7 @@ def thomson_scatter(r_packet, time_explosion):
 def line_scatter(r_packet, time_explosion, line_interaction_type, numba_plasma):
     """
     Line scatter function that handles the scattering itself, including new angle drawn, and calculating nu out using macro atom
-    
+
     Parameters
     ----------
     r_packet : tardis.montecarlo.montecarlo_numba.r_packet.RPacket