Skip to content

Commit

Permalink
Continuum branch [attempted merge with master] (tardis-sn#1886)
Browse files Browse the repository at this point in the history
* Add calculation of bound-free opacity for numba MC

* Fix bound-free frequency sorting error

* Add determination of macro atom activation level after bound-free absorption

skip-checks: true

* Fix sorting error in bound-free opacity interpolator

* Calculate factor used in free-free cooling from fundamental constants

* Add determination of macro atom activation level after continuum absorption

* Fix renaming bug

* Add calculation of free-free opacity

* Add note about free-free Gaunt factor

* Add function for combined bf and ff opacity calculation

* Make it easier to retrieve the transition_idx of a MA deactivation channel

* Swap level order in transition_idx references to deactivation

* Reformat some lines

* Add transition_types and idxs for MA deactivation by cooling

* Add mask for retrieving transition probabilities of non continuum species

* Combine the Markov chain and normal transition probabilities

* Reformat some lines

* Replace dummy k-packet idx in interaction handling

* Change output name of combined continuum transition probabilities

* Add block references for combined continuum transition probabilities

* Activate macro atom by level idx instead of line idx

* Make macro atom data a plasma property to enable modularity

* Split transition probabilities and transition info

* Fix macro atom test

* Adapt transition types in plasma to match the ones in the numba macro atom

* Make the macro atom return the transition type

* Add individual fb cooling probabilities to MC transition probabilities

* Make bound-free cross sections available in the MC calculation

* Initialize the arrays for the continuum estimators

* Convert continuum estimator arrays to DataFrames

* Make t_electrons available in MC calculation

* Propagate continuum estimators into plasma

* Add continuum estimators to Numba estimators

* Reformat some lines

* Made path selection a little bit faster and simplifiedw

* Add sorted list of threshold frequencies for photoionization to Numba MC

* Added selection fucntion for continuum process, WIP

* Moved continuum selection inside of tracepacket, still need a docstring

* added plasma properties that perform the sampling of the free-free and free-bound emission processes

* restrcited sampling range between continuum ids

* Record bound-free estimators during Numba MC

* Began work on Continuum jitclass

* Implemented the Continuum as a jitclass which caches the state at a given shell_id and fequency

* Added the continuum samplers to the  attribute of the continuum processes file.  Fixed implementation after merge

* Added temporary change to make all continuum processes thomson scatter so that the simulation finished in a finite amount of time for testing during development

* fixed the continuum samplers so they are now properly processed.  Also fixed some bugs in the samplers

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

* started addition of free-free emission handler

* Add calculation of bound-free opacity for numba MC

* Fix bound-free frequency sorting error

* Add determination of macro atom activation level after bound-free absorption

skip-checks: true

* Fix sorting error in bound-free opacity interpolator

* Calculate factor used in free-free cooling from fundamental constants

* Add determination of macro atom activation level after continuum absorption

* Fix renaming bug

* Add calculation of free-free opacity

* Add note about free-free Gaunt factor

* Add function for combined bf and ff opacity calculation

* Make it easier to retrieve the transition_idx of a MA deactivation channel

* Swap level order in transition_idx references to deactivation

* Reformat some lines

* Add transition_types and idxs for MA deactivation by cooling

* Add mask for retrieving transition probabilities of non continuum species

* Combine the Markov chain and normal transition probabilities

* Reformat some lines

* Replace dummy k-packet idx in interaction handling

* Change output name of combined continuum transition probabilities

* Add block references for combined continuum transition probabilities

* Activate macro atom by level idx instead of line idx

* Make macro atom data a plasma property to enable modularity

* Split transition probabilities and transition info

* Fix macro atom test

* Adapt transition types in plasma to match the ones in the numba macro atom

* Make the macro atom return the transition type

* Add individual fb cooling probabilities to MC transition probabilities

* Make bound-free cross sections available in the MC calculation

* Initialize the arrays for the continuum estimators

* Convert continuum estimator arrays to DataFrames

* Make t_electrons available in MC calculation

* Propagate continuum estimators into plasma

* Add continuum estimators to Numba estimators

* Reformat some lines

* Add sorted list of threshold frequencies for photoionization to Numba MC

* Record bound-free estimators during Numba MC

* Add estimator for stimulated recombination cooling rate

* Add column name to estimator DataFrames

* 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

* 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

* Fix bug in calculation of corrected photoionization rate coefficient

* Format some lines

* Update tardis/montecarlo/montecarlo_numba/interaction.py

Co-authored-by: Christian Vogl <[email protected]>

* No longer recalculate the local continuum opacities before interacting

* Resolved comment from Christian, do the same process for BF cooling as BF emission

* Minor formatting, currently we still don't actually run through any continuum processes because the new macroatom is not in

* Resolved a couple of comments. (#5)

* 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

* No longer recalculate the local continuum opacities before interacting

* Resolved comment from Christian, do the same process for BF cooling as BF emission

* Minor formatting, currently we still don't actually run through any continuum processes because the new macroatom is not in

* Switch to macro atom data and probabilities with continuum processes

* Raise an exception if more than one two-photon decays are requested

* Add transition info for two-photon decays to numba macro atom data

* Add frequency sampler for two-photon emission

* Updated the bound-free frequency sampler.  Fixed an issue arising from the indexing

* Now allows for continuum processes to occur

* Added proper interaction type to r_packet for continuum processes

* Added bound-free cooling and adiabatic cooling, debugging in progress.  Current issue with macroatom running out of the block

* fixed updatin gof distance_continuum in r_packet, set the chi_bf to zero when there's no cuirrent continuum

* removed print statements for packet index

* Removed excess print statements, added some comments for clarity

* removed a couple commented-out print statements

* added docstrings

* fixed error from typo

* Added conditional compilation of the Continuum jitclass such that normal methods are replaces with dummy methods when we are not worrying about the continuum opacities.  Updated the conditional initialization of the plasma as well accordingly

* updated docstring for montecarlo main loop, should be noted that passing a CPUDispatcher used as a constructor to the main loop function is slow according to numba docs, in the future a factory function should generate the main loop

* added some fixed for running without continuum species

* Added proper flag for adiabatic cooling

* removed extra print statement to check for continuum processes in testing

* removed commented out print statements

* Added configuration parameter for inclusion of continuum processes and had it set during the initialization of the simulation object.  Will be useful for future numba optimization as a compile-time constant

* switched check for continuum processes to rely on config parameter

* Added conditional paths for continuum interactions in the macroatom events based upon compile time constant

* Fixed vpackets throwing a close line error

* Implimented partial fixed for formal integral, development ongoing

* Fixed merge issues

* Fix treatment at end of line list

* Updated continuum interactions to appropriately use comoving frequencies when exciting a macroatom

* added copy method to continuum class

* Added continuum interactions to vpacket

* added continuum interactions to vpacket

* Fixed arg ordering so args come before kwargs, removed conversion of lists to numpy arrays before concatenations

* Vpackets now all share one continuum object since they go sequentially

* decided to go with copies of continuum

* Defer updating continuum if parameters have not changed since last call

* removed check for updating continuum since vpackets always cross a shell boundary before updating

* went back to single instance of continuum to share among vpackets since deffering updating of the continuum is unneccesary

* realtivity check

* Added a step to recompile the main_loop to make sure things get properly set on new runs

* Removed reindexing of photoionization data to be compatable with carsus

* Added fix from chvogl for handling duplicate lines

* Removed the recompile since it caused issues with the lifetimes of some objects

* Fixed bug introduced by relativity fix.  Now recompiling the main loops works and the relativity flag has a namespace

* fixed problems
Co-authored-by: Andreas Flörs <[email protected]>
Co-authored-by: Christian Vogl <[email protected]>

* fixed tracker

* Added

Co-authored-by: Jack O'Brien <[email protected]>

* trying to fix the tests

Co-authored-by: Jack O'Brien <[email protected]>

* remove recompile step due to some numba screwup related tardis-sn#1908

Co-authored-by Jack O'Brien <[email protected]>

Co-authored-by: Christian Vogl <[email protected]>
Co-authored-by: rodot- <[email protected]>
Co-authored-by: Jack O'Brien <[email protected]>
Co-authored-by: Sona Chitchyan <[email protected]>
  • Loading branch information
5 people authored Feb 24, 2022
1 parent 69be60c commit e86ade3
Show file tree
Hide file tree
Showing 29 changed files with 1,533 additions and 145 deletions.
2 changes: 2 additions & 0 deletions tardis/io/logger/tests/test_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@
from tardis.simulation import Simulation
from tardis.io.logger.logger import LOGGING_LEVELS
from tardis import run_tardis
import pytest

pytestmark = pytest.mark.skip(reason='logging testing slow and disabled for now')

def test_logging_simulation(atomic_data_fname, caplog):
"""
Expand Down
31 changes: 31 additions & 0 deletions tardis/montecarlo/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,30 @@ def _initialize_estimator_arrays(self, tau_sobolev_shape):
self.Edotlu_estimator = np.zeros(tau_sobolev_shape)
# TODO: this is the wrong attribute naming style.

def _initialize_continuum_estimator_arrays(self, gamma_shape):
"""
Initialize the arrays for the MC estimators for continuum processes.
Parameters
----------
gamma_shape : tuple
Shape of the array with the photoionization rate coefficients.
"""
self.photo_ion_estimator = np.zeros(gamma_shape, dtype=np.float64)
self.stim_recomb_estimator = np.zeros(gamma_shape, dtype=np.float64)
self.stim_recomb_cooling_estimator = np.zeros(
gamma_shape, dtype=np.float64
)
self.bf_heating_estimator = np.zeros(gamma_shape, dtype=np.float64)

self.stim_recomb_cooling_estimator = np.zeros(
gamma_shape, dtype=np.float64
)

self.photo_ion_estimator_statistics = np.zeros(
gamma_shape, dtype=np.int64
)

def _initialize_geometry_arrays(self, model):
"""
Generate the cgs like geometry arrays for the montecarlo part
Expand Down Expand Up @@ -304,6 +328,13 @@ def run(
# Initializing estimator array
self._initialize_estimator_arrays(plasma.tau_sobolevs.shape)

if not plasma.continuum_interaction_species.empty:
gamma_shape = plasma.gamma.shape
else:
gamma_shape = (0, 0)

self._initialize_continuum_estimator_arrays(gamma_shape)

self._initialize_geometry_arrays(model)

self._initialize_packets(
Expand Down
1 change: 1 addition & 0 deletions tardis/montecarlo/montecarlo_configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,3 +16,4 @@
LEGACY_MODE_ENABLED = False
VPACKET_LOGGING = False
RPACKET_TRACKING = False
CONTINUUM_PROCESSES_ENABLED = False
58 changes: 35 additions & 23 deletions tardis/montecarlo/montecarlo_numba/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
numba_plasma_initialize,
Estimators,
configuration_initialize,
create_continuum_class
)

from tardis.montecarlo import (
Expand All @@ -29,6 +30,7 @@
from numba.typed import List
from tardis.util.base import update_iterations_pbar, update_packet_pbar

#ContinuumObject = create_continuum_class(plasma)

def montecarlo_radial1d(
model,
Expand All @@ -39,6 +41,8 @@ def montecarlo_radial1d(
show_progress_bars,
runner,
):
#montecarlo_main_loop.recompile()
#global ContinuumObject
packet_collection = PacketCollection(
runner.input_r,
runner.input_nu,
Expand All @@ -59,11 +63,18 @@ def montecarlo_radial1d(
runner.nu_bar_estimator,
runner.j_blue_estimator,
runner.Edotlu_estimator,
runner.photo_ion_estimator,
runner.stim_recomb_estimator,
runner.bf_heating_estimator,
runner.stim_recomb_cooling_estimator,
runner.photo_ion_estimator_statistics,
)
packet_seeds = montecarlo_configuration.packet_seeds

number_of_vpackets = montecarlo_configuration.number_of_vpackets
if iteration == 0: montecarlo_main_loop.recompile() # Make sure we update
ContinuumObject = create_continuum_class(plasma)

#if iteration == 0: montecarlo_main_loop.recompile() # Make sure we update
(
v_packets_energy_hist,
last_interaction_type,
Expand All @@ -87,6 +98,7 @@ def montecarlo_radial1d(
runner.spectrum_frequency.value,
number_of_vpackets,
packet_seeds,
ContinuumObject,
montecarlo_configuration.VPACKET_LOGGING,
iteration=iteration,
show_progress_bars=show_progress_bars,
Expand All @@ -102,29 +114,21 @@ def montecarlo_radial1d(

if montecarlo_configuration.VPACKET_LOGGING and number_of_vpackets > 0:
runner.virt_packet_nus = np.concatenate(
np.array(virt_packet_nus)
).ravel()
virt_packet_nus).ravel()
runner.virt_packet_energies = np.concatenate(
np.array(virt_packet_energies)
).ravel()
virt_packet_energies).ravel()
runner.virt_packet_initial_mus = np.concatenate(
np.array(virt_packet_initial_mus)
).ravel()
virt_packet_initial_mus).ravel()
runner.virt_packet_initial_rs = np.concatenate(
np.array(virt_packet_initial_rs)
).ravel()
virt_packet_initial_rs).ravel()
runner.virt_packet_last_interaction_in_nu = np.concatenate(
np.array(virt_packet_last_interaction_in_nu)
).ravel()
virt_packet_last_interaction_in_nu).ravel()
runner.virt_packet_last_interaction_type = np.concatenate(
np.array(virt_packet_last_interaction_type)
).ravel()
virt_packet_last_interaction_type).ravel()
runner.virt_packet_last_line_interaction_in_id = np.concatenate(
np.array(virt_packet_last_line_interaction_in_id)
).ravel()
virt_packet_last_line_interaction_in_id).ravel()
runner.virt_packet_last_line_interaction_out_id = np.concatenate(
np.array(virt_packet_last_line_interaction_out_id)
).ravel()
virt_packet_last_line_interaction_out_id).ravel()
update_iterations_pbar(1)

# Condition for Checking if RPacket Tracking is enabled
Expand All @@ -141,6 +145,7 @@ def montecarlo_main_loop(
spectrum_frequency,
number_of_vpackets,
packet_seeds,
ContinuumObject,
virtual_packet_logging,
iteration,
show_progress_bars,
Expand All @@ -155,16 +160,20 @@ def montecarlo_main_loop(
----------
packet_collection : PacketCollection
numba_model : NumbaModel
numba_plasma : NumbaPlasma
estimators : NumbaEstimators
spectrum_frequency : astropy.units.Quantity
frequency bins
frequency binspas
number_of_vpackets : int
VPackets released per interaction
packet_seeds : numpy.array
ContinuumObject : numba.experimental.jitclass.boxing.Continuum
Constructor method for local continuum jitclass
virtual_packet_logging : bool
Option to enable virtual packet logging.
"""
output_nus = np.empty_like(packet_collection.packets_output_nu)

output_nus = np.empty_like(packet_collection.packets_input_nu)
last_interaction_types = (
np.ones_like(packet_collection.packets_output_nu, dtype=np.int64) * -1
)
Expand Down Expand Up @@ -221,6 +230,7 @@ def montecarlo_main_loop(
total_iterations=total_iterations,
)


if montecarlo_configuration.single_packet_seed != -1:
seed = packet_seeds[montecarlo_configuration.single_packet_seed]
np.random.seed(seed)
Expand All @@ -235,16 +245,18 @@ def montecarlo_main_loop(
seed,
i,
)
continuum = ContinuumObject()
vpacket_collection = vpacket_collections[i]
tracked_rpacket = rpacket_trackers[i]
rpacket_tracker = rpacket_trackers[i]

single_packet_loop(
loop = single_packet_loop(
r_packet,
continuum,
numba_model,
numba_plasma,
estimators,
vpacket_collection,
tracked_rpacket,
rpacket_tracker
)

output_nus[i] = r_packet.nu
Expand Down Expand Up @@ -336,7 +348,7 @@ def montecarlo_main_loop(

packet_collection.packets_output_energy[:] = output_energies[:]
packet_collection.packets_output_nu[:] = output_nus[:]

#print("Finished Main Loop")
return (
v_packets_energy_hist,
last_interaction_types,
Expand Down
6 changes: 2 additions & 4 deletions tardis/montecarlo/montecarlo_numba/calculate_distances.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,11 @@
njit_dict_no_parallel,
)

import tardis.montecarlo.montecarlo_numba.numba_config as nc
from tardis.montecarlo.montecarlo_numba.numba_config import (
C_SPEED_OF_LIGHT,
MISS_DISTANCE,
SIGMA_THOMSON,
ENABLE_FULL_RELATIVITY,
CLOSE_LINE_THRESHOLD,
)

Expand Down Expand Up @@ -59,8 +59,6 @@ def calculate_distance_boundary(r, mu, r_inner, r_outer):
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
Expand Down Expand Up @@ -105,7 +103,7 @@ def calculate_distance_line(
" information, see print statement beforehand"
)

if ENABLE_FULL_RELATIVITY:
if nc.ENABLE_FULL_RELATIVITY:
return calculate_distance_line_full_relativity(
nu_line, nu, time_explosion, r_packet
)
Expand Down
71 changes: 65 additions & 6 deletions tardis/montecarlo/montecarlo_numba/estimators.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
from math import exp
from numba import njit

from tardis.montecarlo.montecarlo_numba import numba_config as nc
from tardis.montecarlo.montecarlo_numba.numba_config import H, KB

from tardis.montecarlo.montecarlo_numba import (
njit_dict_no_parallel,
)
Expand All @@ -9,11 +13,6 @@
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):
"""
Expand All @@ -39,6 +38,66 @@ def set_estimators_full_relativity(
)


@njit(**njit_dict_no_parallel)
def update_bound_free_estimators(
comov_nu,
comov_energy,
shell_id,
distance,
numba_estimator,
t_electron,
x_sect_bfs,
current_continua,
bf_threshold_list_nu,
):
"""
Update the estimators for bound-free processes.
Parameters
----------
comov_nu : float
comov_energy : float
shell_id : int
distance : float
numba_estimator : tardis.montecarlo.montecarlo_numba.numba_interface.Estimators
t_electron : float
Electron temperature in the current cell.
x_sect_bfs : numpy.ndarray, dtype float
Photoionization cross-sections of all bound-free continua for
which absorption is possible for frequency `comov_nu`.
current_continua : numpy.ndarray, dtype int
Continuum ids for which absorption is possible for frequency `comov_nu`.
bf_threshold_list_nu : numpy.ndarray, dtype float
Threshold frequencies for photoionization sorted by decreasing frequency.
"""
# TODO: Add full relativity mode
boltzmann_factor = exp(-(H * comov_nu) / (KB * t_electron))
for i, current_continuum in enumerate(current_continua):
photo_ion_rate_estimator_increment = (
comov_energy * distance * x_sect_bfs[i] / comov_nu
)
numba_estimator.photo_ion_estimator[
current_continuum, shell_id
] += photo_ion_rate_estimator_increment
numba_estimator.stim_recomb_estimator[current_continuum, shell_id] += (
photo_ion_rate_estimator_increment * boltzmann_factor
)
numba_estimator.photo_ion_estimator_statistics[
current_continuum, shell_id
] += 1

nu_th = bf_threshold_list_nu[current_continuum]
bf_heating_estimator_increment = (
comov_energy * distance * x_sect_bfs[i] * (1 - nu_th / comov_nu)
)
numba_estimator.bf_heating_estimator[
current_continuum, shell_id
] += bf_heating_estimator_increment
numba_estimator.stim_recomb_cooling_estimator[
current_continuum, shell_id
] += (bf_heating_estimator_increment * boltzmann_factor)


@njit(**njit_dict_no_parallel)
def update_line_estimators(
estimators, r_packet, cur_line_id, distance_trace, time_explosion
Expand All @@ -63,7 +122,7 @@ def update_line_estimators(
( time_explosion * C)
"""

if not ENABLE_FULL_RELATIVITY:
if not nc.ENABLE_FULL_RELATIVITY:
energy = calc_packet_energy(r_packet, distance_trace, time_explosion)
else:
energy = calc_packet_energy_full_relativity(r_packet)
Expand Down
16 changes: 9 additions & 7 deletions tardis/montecarlo/montecarlo_numba/formal_integral.py
Original file line number Diff line number Diff line change
Expand Up @@ -195,14 +195,14 @@ def numba_formal_integral(
return L


integrator_spec = [
("model", NumbaModel.class_type.instance_type),
("plasma", NumbaPlasma.class_type.instance_type),
("points", int64),
]
#integrator_spec = [
# ("model", NumbaModel.class_type.instance_type),
# ("plasma", NumbaPlasma.class_type.instance_type),
# ("points", int64),
#]


@jitclass(integrator_spec)
#@jitclass(integrator_spec)
class NumbaFormalIntegrator(object):
"""
Helper class for performing the formal integral
Expand Down Expand Up @@ -383,8 +383,10 @@ def make_source_function(self):
model = self.model
runner = self.runner

#macro_ref = self.atomic_data.macro_atom_references
macro_ref = self.atomic_data.macro_atom_references
macro_data = self.atomic_data.macro_atom_data
#macro_data = self.atomic_data.macro_atom_data
macro_data = self.original_plasma.macro_atom_data

no_lvls = len(self.atomic_data.levels)
no_shells = len(model.w)
Expand Down
Loading

0 comments on commit e86ade3

Please sign in to comment.