import numpy as np
from numba import njit
from tardis.montecarlo.montecarlo_numba import (
njit_dict_no_parallel,
)
[docs]@njit(**njit_dict_no_parallel)
def velocity_dvdr(r_packet, numba_model):
"""
Velocity at radius r and dv/dr of current shell
Parameters
----------
r_packet: RPacket
numba_model: NumbaModel
Returns
-----------
v: float, current velocity
frac: float, dv/dr for current shell
"""
shell_id = r_packet.current_shell_id
v_inner = numba_model.v_inner[shell_id]
v_outer = numba_model.v_outer[shell_id]
r_inner = numba_model.r_inner[shell_id]
r_outer = numba_model.r_outer[shell_id]
r = r_packet.r
frac = (v_outer - v_inner) / (r_outer - r_inner)
return v_inner + frac * (r - r_inner), frac
[docs]@njit(**njit_dict_no_parallel)
def tau_sobolev_factor(r_packet, numba_model):
"""
The angle and velocity dependent Tau Sobolev factor component. Is called when ENABLE_NONHOMOLOGOUS_EXPANSION is set to True.
Note: to get Tau Sobolev, this needs to be multiplied by tau_sobolevs found from plasma
Parameters
----------
r_packet: RPacket
numba_model: NumbaModel
Returns
-----------
factor = 1.0 / ((1 - mu * mu) * v / r + mu * mu * dvdr)
"""
v, dvdr = velocity_dvdr(r_packet, numba_model)
r = r_packet.r
mu = r_packet.mu
factor = 1.0 / ((1 - mu * mu) * v / r + mu * mu * dvdr)
return factor
# @njit(**njit_dict_no_parallel)
[docs]def quartic_roots(a, b, c, d, e, threshold):
"""
Solves ax^4 + bx^3 + cx^2 + dx + e = 0, for the real roots greater than the threshold returns (x - threshold).
Uses: https://en.wikipedia.org/wiki/Quartic_function#General_formula_for_roots
Parameters
-----------
a, b, c, d, e: coefficients of the equations ax^4 + bx^3 + cx^2 + dx + e = 0, float
threshold: lower needed limit on roots, float
Returns
-----------
roots: real positive roots of ax^4 + bx^3 + cx^2 + dx + e = 0
"""
roots = np.roots((a, b, c, d, e))
roots = [root for root in roots if isinstance(root, float)]
roots = [root for root in roots if root > threshold]
return roots