Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Flexible packet initial radius #1592

Merged
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions tardis/montecarlo/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ def _initialize_geometry_arrays(self, model):
self.r_outer_cgs = model.r_outer.to("cm").value
self.v_inner_cgs = model.v_inner.to("cm/s").value

def _initialize_packets(self, T, no_of_packets, iteration):
def _initialize_packets(self, T, no_of_packets, iteration, radius):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not really related to this PR but maybe we should rename T (e.g. temperature).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed. For the restructuring to come!

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we make this an issues

# the iteration is added each time to preserve randomness
# across different simulations with the same temperature,
# for example. We seed the random module instead of the numpy module
Expand All @@ -175,10 +175,11 @@ def _initialize_packets(self, T, no_of_packets, iteration):
seed = self.seed + iteration
rng = np.random.default_rng(seed=seed)
seeds = rng.choice(MAX_SEED_VAL, no_of_packets, replace=True)
nus, mus, energies = self.packet_source.create_packets(
T, no_of_packets, rng
radii, nus, mus, energies = self.packet_source.create_packets(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am wondering if the radius/radii should be an attribute of the packet_source. In the case of the BlackBodySimpleSource, it would be set to r_inner by default. This way we might avoid passing this information through multiple layers.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Possibly, but we also expect the radius to be different for different packets in the nebular phase. So it needs to be an array of radii corresponding to the energy deposition radius.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think there are all sorts of potential ways to configure this - but for now I think it's fine how it is.

T, no_of_packets, rng, radius
)
mc_config_module.packet_seeds = seeds
self.input_r = radii
self.input_nu = nus
self.input_mu = mus
self.input_energy = energies
Expand Down Expand Up @@ -291,7 +292,9 @@ def run(

self._initialize_geometry_arrays(model)

self._initialize_packets(model.t_inner.value, no_of_packets, iteration)
self._initialize_packets(
model.t_inner.value, no_of_packets, iteration, model.r_inner[0]
)

configuration_initialize(self, no_of_virtual_packets)
montecarlo_radial1d(model, plasma, self)
Expand Down
1 change: 1 addition & 0 deletions tardis/montecarlo/montecarlo_numba/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@

def montecarlo_radial1d(model, plasma, runner):
packet_collection = PacketCollection(
runner.input_r,
runner.input_nu,
runner.input_mu,
runner.input_energy,
Expand Down
43 changes: 32 additions & 11 deletions tardis/montecarlo/montecarlo_numba/numba_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ def __init__(
):
"""
Plasma for the Numba code

Parameters
----------
electron_density : numpy.ndarray
Expand Down Expand Up @@ -161,6 +161,7 @@ def numba_plasma_initialize(plasma, line_interaction_type):


packet_collection_spec = [
("packets_input_radius", float64[:]),
("packets_input_nu", float64[:]),
("packets_input_mu", float64[:]),
("packets_input_energy", float64[:]),
Expand All @@ -173,12 +174,14 @@ def numba_plasma_initialize(plasma, line_interaction_type):
class PacketCollection(object):
def __init__(
self,
packets_input_radius,
packets_input_nu,
packets_input_mu,
packets_input_energy,
packets_output_nu,
packets_output_energy,
):
self.packets_input_radius = packets_input_radius
self.packets_input_nu = packets_input_nu
self.packets_input_mu = packets_input_mu
self.packets_input_energy = packets_input_energy
Expand Down Expand Up @@ -220,10 +223,18 @@ def __init__(
self.nus = np.empty(temporary_v_packet_bins, dtype=np.float64)
self.energies = np.empty(temporary_v_packet_bins, dtype=np.float64)
self.number_of_vpackets = number_of_vpackets
self.last_interaction_in_nu = np.zeros(temporary_v_packet_bins, dtype=np.float64)
self.last_interaction_type = -1 * np.ones(temporary_v_packet_bins, dtype=np.int64)
self.last_interaction_in_id = -1 * np.ones(temporary_v_packet_bins, dtype=np.int64)
self.last_interaction_out_id = -1 * np.ones(temporary_v_packet_bins, dtype=np.int64)
self.last_interaction_in_nu = np.zeros(
temporary_v_packet_bins, dtype=np.float64
)
self.last_interaction_type = -1 * np.ones(
temporary_v_packet_bins, dtype=np.int64
)
self.last_interaction_in_id = -1 * np.ones(
temporary_v_packet_bins, dtype=np.int64
)
self.last_interaction_out_id = -1 * np.ones(
temporary_v_packet_bins, dtype=np.int64
)
self.idx = 0
self.rpacket_index = rpacket_index
self.length = temporary_v_packet_bins
Expand All @@ -241,21 +252,31 @@ def set_properties(
temp_length = self.length * 2 + self.number_of_vpackets
temp_nus = np.empty(temp_length, dtype=np.float64)
temp_energies = np.empty(temp_length, dtype=np.float64)
temp_last_interaction_in_nu = np.empty(temp_length, dtype=np.float64)
temp_last_interaction_in_nu = np.empty(
temp_length, dtype=np.float64
)
temp_last_interaction_type = np.empty(temp_length, dtype=np.int64)
temp_last_interaction_in_id = np.empty(temp_length, dtype=np.int64)
temp_last_interaction_out_id = np.empty(temp_length, dtype=np.int64)
temp_nus[: self.length] = self.nus
temp_energies[: self.length] = self.energies
temp_last_interaction_in_nu[: self.length] = self.last_interaction_in_nu
temp_last_interaction_type[: self.length] = self.last_interaction_type
temp_last_interaction_in_id[: self.length] = self.last_interaction_in_id
temp_last_interaction_out_id[: self.length] = self.last_interaction_out_id
temp_last_interaction_in_nu[
: self.length
] = self.last_interaction_in_nu
temp_last_interaction_type[
: self.length
] = self.last_interaction_type
temp_last_interaction_in_id[
: self.length
] = self.last_interaction_in_id
temp_last_interaction_out_id[
: self.length
] = self.last_interaction_out_id

self.nus = temp_nus
self.energies = temp_energies
self.last_interaction_in_nu = temp_last_interaction_in_nu
self.last_interaction_type = temp_last_interaction_type
self.last_interaction_type = temp_last_interaction_type
self.last_interaction_in_id = temp_last_interaction_in_id
self.last_interaction_out_id = temp_last_interaction_out_id
self.length = temp_length
Expand Down
2 changes: 2 additions & 0 deletions tardis/montecarlo/montecarlo_numba/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ def nb_simulation_verysimple(config_verysimple, atomic_dataset):
def verysimple_collection(nb_simulation_verysimple):
runner = nb_simulation_verysimple.runner
return PacketCollection(
runner.input_r,
runner.input_nu,
runner.input_mu,
runner.input_energy,
Expand Down Expand Up @@ -99,6 +100,7 @@ def verysimple_3vpacket_collection(nb_simulation_verysimple):
def verysimple_packet_collection(nb_simulation_verysimple):
runner = nb_simulation_verysimple.runner
return PacketCollection(
runner.input_r,
runner.input_nu,
runner.input_mu,
runner.input_energy,
Expand Down
7 changes: 5 additions & 2 deletions tardis/montecarlo/montecarlo_numba/tests/test_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,9 @@ def test_montecarlo_main_loop(

runner._initialize_geometry_arrays(model)
runner._initialize_estimator_arrays(numba_plasma.tau_sobolev.shape)
runner._initialize_packets(model.t_inner.value, 100000, 0)
runner._initialize_packets(
model.t_inner.value, 100000, 0, model.r_inner[0].value
)

# Init parameters
montecarlo_configuration.v_packet_spawn_start_frequency = (
Expand All @@ -93,6 +95,7 @@ def test_montecarlo_main_loop(

# Init packet collection from runner
packet_collection = PacketCollection(
runner.input_r,
runner.input_nu,
runner.input_mu,
runner.input_energy,
Expand Down Expand Up @@ -130,7 +133,7 @@ def test_montecarlo_main_loop(
for i in range(len(packet_collection.packets_input_nu)):
# Generate packet
packet = r_packet.RPacket(
numba_model.r_inner[0],
packet_collection.packets_input_radius[i],
packet_collection.packets_input_mu[i],
packet_collection.packets_input_nu[i],
packet_collection.packets_input_energy[i],
Expand Down
15 changes: 8 additions & 7 deletions tardis/montecarlo/packet_source.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,9 @@ def create_zero_limb_darkening_packet_mus(no_of_packets, rng):
@staticmethod
def create_uniform_packet_energies(no_of_packets, rng):
"""
Uniformly distribute energy in arbitrary units where the ensemble of
packets has energy of 1.
Uniformly distribute energy in arbitrary units where the ensemble of
packets has energy of 1.

Parameters
----------
no_of_packets : int
Expand Down Expand Up @@ -73,15 +73,15 @@ def create_blackbody_packet_nus(T, no_of_packets, rng, l_samples=1000):
.. math::
x = -\\ln{(\\xi_1\\xi_2\\xi_3\\xi_4)}/l_{\\rm min}\\; .
where :math:`x=h\\nu/kT`

Parameters
----------
T : float
temperature
no_of_packets : int
l_samples : int
number of l_samples needed in the algorithm

Returns
-------
: numpy.ndarray
Expand Down Expand Up @@ -110,9 +110,10 @@ class BlackBodySimpleSource(BasePacketSource):
part.
"""

def create_packets(self, T, no_of_packets, rng):
def create_packets(self, T, no_of_packets, rng, radius):
andrewfullard marked this conversation as resolved.
Show resolved Hide resolved
radii = np.ones(no_of_packets) * radius
nus = self.create_blackbody_packet_nus(T, no_of_packets, rng)
mus = self.create_zero_limb_darkening_packet_mus(no_of_packets, rng)
energies = self.create_uniform_packet_energies(no_of_packets, rng)

return nus, mus, energies
return radii, nus, mus, energies