def thomson_scatter(r_packet, time_explosion):
"""
Thomson scattering — no longer line scattering
2) get the doppler factor at that position with the old angle
3) convert the current energy and nu into the comoving
frame with the old mu
4) Scatter and draw new mu - update mu
5) Transform the comoving energy and nu back using the new mu
Parameters
----------
r_packet : RPacket
time_explosion: float
time since explosion in seconds
"""
old_doppler_factor = get_doppler_factor(r_packet.r, r_packet.mu,
time_explosion)
comov_nu = r_packet.nu * old_doppler_factor
comov_energy = r_packet.energy * old_doppler_factor
r_packet.mu = get_random_mu()
inverse_new_doppler_factor = get_inverse_doppler_factor(
r_packet.r, r_packet.mu, time_explosion)
r_packet.nu = comov_nu * inverse_new_doppler_factor
r_packet.energy = comov_energy * inverse_new_doppler_factor
if montecarlo_configuration.full_relativity:
r_packet.mu = angle_aberration_CMF_to_LF(r_packet, time_explosion,
r_packet.mu)
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
r_packet: RPacket
time_explosion: float
line_interaction_type: enum
numba_plasma: NumbaPlasma
"""
old_doppler_factor = get_doppler_factor(r_packet.r, r_packet.mu,
time_explosion)
r_packet.mu = get_random_mu()
inverse_new_doppler_factor = get_inverse_doppler_factor(
r_packet.r, r_packet.mu, time_explosion)
comov_energy = r_packet.energy * old_doppler_factor
r_packet.energy = comov_energy * inverse_new_doppler_factor
if line_interaction_type == LineInteractionType.SCATTER:
line_emission(r_packet, r_packet.next_line_id, time_explosion,
numba_plasma)
else: # includes both macro atom and downbranch - encoded in the transition probabilities
emission_line_id = macro_atom(r_packet, numba_plasma)
line_emission(r_packet, emission_line_id, time_explosion, numba_plasma)
def test_unphysical_doppler_factor(mu, r, inv_t_exp, packet, model):
# Set the params from test cases here
# TODO: add relativity tests
time_explosion = 1 / inv_t_exp
# Perform any other setups just before this, they can be additional calls
# to other methods or introduction of some temporary variables
with pytest.raises(r_packet.SuperluminalError):
obtained = r_packet.get_doppler_factor(r, mu, time_explosion)
def v_packet_initialize_line_id(v_packet, numba_plasma, numba_model):
inverse_line_list_nu = numba_plasma.line_list_nu[::-1]
doppler_factor = r_packet.get_doppler_factor(
v_packet.r, v_packet.mu, numba_model.time_explosion
)
comov_nu = v_packet.nu * doppler_factor
next_line_id = len(numba_plasma.line_list_nu) - np.searchsorted(
inverse_line_list_nu, comov_nu
)
v_packet.next_line_id = next_line_id
def test_get_doppler_factor(mu, r, inv_t_exp, expected):
# Set the params from test cases here
# TODO: add relativity tests
time_explosion = 1 / inv_t_exp
# Perform any other setups just before this, they can be additional calls
# to other methods or introduction of some temporary variables
obtained = r_packet.get_doppler_factor(r, mu, time_explosion)
# Perform required assertions
assert_almost_equal(obtained, expected)
def test_frame_transformations(packet, mu, r, inv_t_exp, full_relativity):
packet = r_packet.RPacket(r=r, mu=mu, energy=packet.energy, nu=packet.nu)
mc.full_relativity = bool(full_relativity)
mc.full_relativity = full_relativity
inverse_doppler_factor = r_packet.get_inverse_doppler_factor(
r, mu, 1 / inv_t_exp)
r_packet.angle_aberration_CMF_to_LF(packet, 1 / inv_t_exp, packet.mu)
doppler_factor = r_packet.get_doppler_factor(r, mu, 1 / inv_t_exp)
mc.full_relativity = False
assert_almost_equal(doppler_factor * inverse_doppler_factor, 1.0)
def test_compute_distance2line_relativistic(mu, r, t_exp, nu, nu_line,
full_relativity, packet, runner):
packet = r_packet.RPacket(r=r, nu=nu, mu=mu, energy=packet.energy)
# packet.nu_line = nu_line
numba_estimator = Estimators(
runner.j_estimator,
runner.nu_bar_estimator,
runner.j_blue_estimator,
runner.Edotlu_estimator,
)
mc.full_relativity = bool(full_relativity)
doppler_factor = r_packet.get_doppler_factor(r, mu, t_exp)
comov_nu = packet.nu * doppler_factor
distance = r_packet.calculate_distance_line(packet, comov_nu, nu_line,
t_exp)
r_packet.move_r_packet(packet, distance, t_exp, numba_estimator)
doppler_factor = r_packet.get_doppler_factor(r, mu, t_exp)
comov_nu = packet.nu * doppler_factor
mc.full_relativity = False
assert_allclose(comov_nu, nu_line, rtol=1e-14)
def test_move_r_packet(
packet_params,
expected_params,
packet,
model,
estimators,
ENABLE_FULL_RELATIVITY,
):
distance = 1.0e13
packet.nu = packet_params["nu"]
packet.mu = packet_params["mu"]
packet.energy = packet_params["energy"]
packet.r = packet_params["r"]
numba_config.ENABLE_FULL_RELATIVITY = ENABLE_FULL_RELATIVITY
r_packet.move_r_packet.recompile(
) # This must be done as move_r_packet was jitted with ENABLE_FULL_RELATIVITY
doppler_factor = r_packet.get_doppler_factor(packet.r, packet.mu,
model.time_explosion)
r_packet.move_r_packet(packet, distance, model.time_explosion, estimators)
assert_almost_equal(packet.mu, expected_params["mu"])
assert_almost_equal(packet.r, expected_params["r"])
expected_j = expected_params["j"]
expected_nubar = expected_params["nubar"]
if ENABLE_FULL_RELATIVITY:
expected_j *= doppler_factor
expected_nubar *= doppler_factor
numba_config.ENABLE_FULL_RELATIVITY = False
assert_allclose(estimators.j_estimator[packet.current_shell_id],
expected_j,
rtol=5e-7)
assert_allclose(
estimators.nu_bar_estimator[packet.current_shell_id],
expected_nubar,
rtol=5e-7,
)
def test_compute_distance2line(packet_params, expected_params, packet, model):
packet = r_packet.RPacket(packet.r, packet.mu, packet.nu, packet.energy)
nu_line = packet_params["nu_line"]
# packet.next_line_id = packet_params['next_line_id']
# packet.last_line = packet_params['last_line']
time_explosion = model.time_explosion
doppler_factor = r_packet.get_doppler_factor(packet.r, packet.mu,
time_explosion)
comov_nu = packet.nu * doppler_factor
d_line = 0
obtained_tardis_error = None
try:
d_line = r_packet.calculate_distance_line(packet, comov_nu, nu_line,
time_explosion)
except r_packet.MonteCarloException:
obtained_tardis_error = r_packet.MonteCarloException
assert_almost_equal(d_line, expected_params["d_line"])
assert obtained_tardis_error == expected_params["tardis_error"]
def test_move_packet(packet_params, expected_params, packet, model,
full_relativity):
distance = 1e13
packet.nu = packet_params["nu"]
packet.mu = packet_params["mu"]
packet.energy = packet_params["energy"]
packet.r = packet_params["r"]
# model.full_relativity = full_relativity
mc.full_relativity = full_relativity
doppler_factor = r_packet.get_doppler_factor(packet.r, packet.mu,
model.time_explosion)
numba_estimator = Estimators(packet_params["j"], packet_params["nu_bar"],
0, 0)
r_packet.move_r_packet(packet, distance, model.time_explosion,
numba_estimator)
assert_almost_equal(packet.mu, expected_params["mu"])
assert_almost_equal(packet.r, expected_params["r"])
expected_j = expected_params["j"]
expected_nubar = expected_params["nubar"]
if full_relativity:
expected_j *= doppler_factor
expected_nubar *= doppler_factor
mc.full_relativity = False
assert_allclose(
numba_estimator.j_estimator[packet.current_shell_id],
expected_j,
rtol=5e-7,
)
assert_allclose(
numba_estimator.nu_bar_estimator[packet.current_shell_id],
expected_nubar,
rtol=5e-7,
)
def test_calculate_distance_line(packet_params, expected_params, static_packet,
model):
nu_line = packet_params["nu_line"]
is_last_line = packet_params["is_last_line"]
time_explosion = model.time_explosion
doppler_factor = r_packet.get_doppler_factor(static_packet.r,
static_packet.mu,
time_explosion)
comov_nu = static_packet.nu * doppler_factor
d_line = 0
obtained_tardis_error = None
try:
d_line = r_packet.calculate_distance_line(static_packet, comov_nu,
is_last_line, nu_line,
time_explosion)
except r_packet.MonteCarloException:
obtained_tardis_error = r_packet.MonteCarloException
assert_almost_equal(d_line, expected_params["d_line"])
assert obtained_tardis_error == expected_params["tardis_error"]
def trace_vpacket_within_shell(v_packet, numba_model, numba_plasma):
"""
Trace VPacket within one shell (relatively simple operation)
"""
r_inner = numba_model.r_inner[v_packet.current_shell_id]
r_outer = numba_model.r_outer[v_packet.current_shell_id]
distance_boundary, delta_shell = calculate_distance_boundary(
v_packet.r, v_packet.mu, r_inner, r_outer
)
# defining start for line interaction
start_line_id = v_packet.next_line_id
# e scattering initialization
cur_electron_density = numba_plasma.electron_density[
v_packet.current_shell_id
]
tau_electron = calculate_tau_electron(
cur_electron_density, distance_boundary
)
tau_trace_combined = tau_electron
# Calculating doppler factor
doppler_factor = get_doppler_factor(
v_packet.r, v_packet.mu, numba_model.time_explosion
)
comov_nu = v_packet.nu * doppler_factor
cur_line_id = start_line_id
for cur_line_id in range(start_line_id, len(numba_plasma.line_list_nu)):
# if tau_trace_combined > 10: ### FIXME ?????
# break
nu_line = numba_plasma.line_list_nu[cur_line_id]
# TODO: Check if this is what the C code does
tau_trace_line = numba_plasma.tau_sobolev[
cur_line_id, v_packet.current_shell_id
]
if cur_line_id == len(numba_plasma.line_list_nu) - 1:
is_last_line = True
else:
is_last_line = False
distance_trace_line = calculate_distance_line(
v_packet,
comov_nu,
is_last_line,
nu_line,
numba_model.time_explosion,
)
if cur_line_id != (len(numba_plasma.line_list_nu) - 1):
test_for_close_line(
v_packet,
cur_line_id,
numba_plasma.line_list_nu[cur_line_id - 1],
numba_plasma,
)
if distance_boundary <= distance_trace_line:
break
tau_trace_combined += tau_trace_line
else:
if cur_line_id == (len(numba_plasma.line_list_nu) - 1):
cur_line_id += 1
v_packet.next_line_id = cur_line_id
return tau_trace_combined, distance_boundary, delta_shell
def trace_vpacket_volley(
r_packet, vpacket_collection, numba_model, numba_plasma
):
"""
Shoot a volley of vpackets (the vpacket collection specifies how many)
from the current position of the rpacket.
Parameters
----------
r_packet : [type]
[description]
vpacket_collection : [type]
[description]
numba_model : [type]
[description]
numba_plasma : [type]
[description]
"""
if (r_packet.nu < vpacket_collection.v_packet_spawn_start_frequency) or (
r_packet.nu > vpacket_collection.v_packet_spawn_end_frequency
):
return
no_of_vpackets = vpacket_collection.number_of_vpackets
if no_of_vpackets == 0:
return
### TODO theoretical check for r_packet nu within vpackets bins - is done somewhere else I think
if r_packet.r > numba_model.r_inner[0]: # not on inner_boundary
r_inner_over_r = numba_model.r_inner[0] / r_packet.r
mu_min = -math.sqrt(1 - r_inner_over_r * r_inner_over_r)
v_packet_on_inner_boundary = False
if montecarlo_configuration.full_relativity:
mu_min = angle_aberration_LF_to_CMF(
r_packet, numba_model.time_explosion, mu_min
)
else:
v_packet_on_inner_boundary = True
mu_min = 0.0
mu_bin = (1.0 - mu_min) / no_of_vpackets
r_packet_doppler_factor = get_doppler_factor(
r_packet.r, r_packet.mu, numba_model.time_explosion
)
for i in range(no_of_vpackets):
v_packet_mu = mu_min + i * mu_bin + np.random.random() * mu_bin
if v_packet_on_inner_boundary: # The weights are described in K&S 2014
weight = 2 * v_packet_mu / no_of_vpackets
else:
weight = (1 - mu_min) / (2 * no_of_vpackets)
# C code: next line, angle_aberration_CMF_to_LF( & virt_packet, storage);
if montecarlo_configuration.full_relativity:
v_packet_mu = angle_aberration_CMF_to_LF(
r_packet, numba_model.time_explosion, v_packet_mu
)
v_packet_doppler_factor = get_doppler_factor(
r_packet.r, v_packet_mu, numba_model.time_explosion
)
# transform between r_packet mu and v_packet_mu
doppler_factor_ratio = r_packet_doppler_factor / v_packet_doppler_factor
v_packet_nu = r_packet.nu * doppler_factor_ratio
v_packet_energy = r_packet.energy * weight * doppler_factor_ratio
v_packet = VPacket(
r_packet.r,
v_packet_mu,
v_packet_nu,
v_packet_energy,
r_packet.current_shell_id,
r_packet.next_line_id,
i,
r_packet.is_close_line,
)
if r_packet.next_line_id <= (len(numba_plasma.line_list_nu) - 1):
test_for_close_line(
v_packet,
r_packet.next_line_id,
numba_plasma.line_list_nu[r_packet.next_line_id - 1],
numba_plasma,
)
tau_vpacket = trace_vpacket(v_packet, numba_model, numba_plasma)
v_packet.energy *= math.exp(-tau_vpacket)
vpacket_collection.set_properties(
v_packet.nu,
v_packet.energy,
r_packet.last_interaction_in_nu,
r_packet.last_interaction_type,
r_packet.last_line_interaction_in_id,
r_packet.last_line_interaction_out_id,
)
