diff --git a/docs/index.rst b/docs/index.rst index e22c73e86ca..57e2feae9f9 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -78,6 +78,7 @@ Mission Statement physics/montecarlo/index physics/est_and_conv/index physics/spectrum/index + physics/energy_input/index .. toctree:: diff --git a/docs/physics/energy_input/gammaray_deposition.ipynb b/docs/physics/energy_input/gammaray_deposition.ipynb new file mode 100644 index 00000000000..d813eefb906 --- /dev/null +++ b/docs/physics/energy_input/gammaray_deposition.ipynb @@ -0,0 +1,306 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Gamma-ray energy deposition\n", + "\n", + "This notebook provides the initial implementation of Gamma-ray energy deposition into an arbitrary ejecta.\n", + "It is a WORK IN PROGRESS and should NOT be used for any science work until further notice." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Main loop\n", + "\n", + "Generates a simple 1D ejecta and a list of gamma-ray objects.\n", + "\n", + "Runs packets of gamma-rays through the ejecta. Handles interactions by calling the appropriate function. \n", + "\n", + "Adds deposited energy and output energy to 2 different dataframes." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "scrolled": false, + "tags": [] + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "import astropy as apy\n", + "\n", + "from tardis.energy_input.indivisible_packets import main_gamma_ray_loop\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from tardis.model import Radial1DModel\n", + "from tardis.io.config_reader import Configuration\n", + "\n", + "# Set up packet count\n", + "num_packets = 50000\n", + "\n", + "# Lock seed\n", + "np.random.seed(1)\n", + "\n", + "# Adjust model\n", + "config = Configuration.from_yaml(\"../../tardis/io/tests/data/tardis_configv1_density_exponential_nebular.yml\")\n", + "config.model.structure.velocity.start = 1 * apy.units.km / apy.units.s\n", + "config.model.structure.density.rho_0 = 5e2 * apy.units.g / (apy.units.cm ** 3)\n", + "config.supernova.time_explosion = 2.0 * apy.units.d\n", + "\n", + "config.atom_data = \"kurucz_cd23_chianti_H_He.h5\"\n", + "\n", + "# Create model\n", + "model = Radial1DModel.from_config(config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Generate plasma" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "from tardis.plasma.properties import Density, Abundance, IsotopeAbundance, IsotopeNumberDensity, AtomicData, AtomicMass, IsotopeMass, NumberDensity, SelectedAtoms\n", + "from tardis.plasma.base import BasePlasma\n", + "from tardis.io.atom_data import AtomData\n", + "\n", + "input = [Density, Abundance, IsotopeAbundance, AtomicData, AtomicMass, IsotopeNumberDensity, NumberDensity, SelectedAtoms, IsotopeMass]\n", + "\n", + "plasma = BasePlasma(\n", + " plasma_properties=input,\n", + " density = model.density,\n", + " abundance = model.abundance,\n", + " isotope_abundance = model.raw_isotope_abundance,\n", + " atomic_data = AtomData.from_hdf(config.atom_data)\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "# Compute energy deposition rate\n", + "# ejecta_energy_df is the deposited energy\n", + "# ejecta_plot_energy_df is information for plotting\n", + "# escape_energy is the escaping energy\n", + "# decayed_packet_count is the number of packets created per shell\n", + "# energy_plot_positrons is the deposited energy from positrons\n", + "# estimated_deposition is the deposited energy from the Kasen (2006) estimator (currently not functional)\n", + "(\n", + " energy_df,\n", + " energy_plot_df,\n", + " escape_energy,\n", + " decayed_packet_count,\n", + " energy_plot_positrons,\n", + " estimated_deposition\n", + ") = main_gamma_ray_loop(\n", + " num_packets,\n", + " model,\n", + " plasma,\n", + " time_steps=50,\n", + " time_end=50.0,\n", + ")\n", + "\n", + "ejecta_energy = energy_plot_df[\"energy_input\"]\n", + "ejecta_energy_r = energy_plot_df[\"energy_input_r\"]\n", + "energy_input_time = energy_plot_df[\"energy_input_time\"]\n", + "energy_input_type = energy_plot_df[\"energy_input_type\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Plotting results\n", + "\n", + "Energy deposited at a given radius" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig = plt.figure(dpi=150, facecolor='w')\n", + "ax = fig.add_subplot(111)\n", + "\n", + "scatter = ax.scatter(np.array(ejecta_energy_r)/np.max(model.v_outer.value), np.array(ejecta_energy), s=1, alpha=0.1)\n", + "ax.set_xlabel(\"r/R\")\n", + "ax.set_ylabel(\"E (keV)\")\n", + "ax.semilogy();" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Interactions binned by radius" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig = plt.figure(dpi=150, facecolor='w')\n", + "ax = fig.add_subplot(111)\n", + "ax.hist(np.array(ejecta_energy_r), bins=200)\n", + "#ax.set_xlim(0, 1)\n", + "ax.set_xlabel(\"r (cm)\")\n", + "ax.set_ylabel(\"Interaction count\");" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Density Profile" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig = plt.figure(dpi=150, facecolor='w')\n", + "plt.semilogy(model.r_middle/np.max(model.r_outer), model.density)\n", + "plt.plot(0,0)\n", + "plt.ylabel(\"Density g cm$^{-3}$\")\n", + "plt.xlabel(\"r/R\");" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Fraction of energy escaping from the ejecta" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "np.sum(np.sum(escape_energy) / (np.sum(escape_energy) + np.sum(energy_df)))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Spectrum of escape energy at the final time step" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from tardis.energy_input.energy_source import read_artis_lines\n", + "\n", + "ni56_lines = read_artis_lines(\"ni56\")\n", + "co56_lines = read_artis_lines(\"co56\")\n", + "\n", + "plt.figure(figsize=(12, 6), dpi=150)\n", + "plt.step(escape_energy.index, escape_energy.iloc[:,49], label=\"$\\gamma$-spectrum\", where=\"post\")\n", + "plt.xlabel(\"Energy (keV)\")\n", + "plt.ylabel(\"keV/keV/s\");\n", + "plt.loglog()\n", + "plt.xlim(100, 4000)\n", + "#plt.ylim(0.01, 100)\n", + "plt.vlines(ni56_lines.energy*1000, 0.01, 0.02, color=\"k\", label=\"Ni56\")\n", + "plt.vlines(co56_lines.energy*1000, 0.01, 0.02, color=\"r\", label=\"Co56\")\n", + "plt.legend()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Energy deposition rate\n", + "\n", + "Dataframe index is the radial grid location. Columns are time steps in seconds" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "energy_df" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Energy deposition rate versus radius at time_explosion" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig = plt.figure(dpi=150, facecolor='w')\n", + "plt.semilogy(model.v_middle/np.max(model.v_outer), energy_df.iloc[:, 0])\n", + "plt.xlabel(\"r/R\")\n", + "plt.ylabel(\"Energy deposition rate [eV/s$]\");" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "interpreter": { + "hash": "1d9f0af959d2c458edd906b0042305fe807dd53effdb40eb4325690d1852d0c6" + }, + "kernelspec": { + "display_name": "Python 3.8.13 ('tardis_new')", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.13" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/physics/energy_input/index.rst b/docs/physics/energy_input/index.rst new file mode 100644 index 00000000000..5aecf0aa908 --- /dev/null +++ b/docs/physics/energy_input/index.rst @@ -0,0 +1,6 @@ +*************************** +Gamma Ray Energy Deposition +*************************** + +.. toctree:: + gammaray_deposition \ No newline at end of file diff --git a/tardis/energy_input/GXPacket.py b/tardis/energy_input/GXPacket.py new file mode 100644 index 00000000000..282a0d11b57 --- /dev/null +++ b/tardis/energy_input/GXPacket.py @@ -0,0 +1,70 @@ +from enum import IntEnum +import numpy as np +from numba import int64, float64 +from numba.experimental import jitclass + + +class GXPacketStatus(IntEnum): + BETA_DECAY = -1 + COMPTON_SCATTER = 0 + PHOTOABSORPTION = 1 + PAIR_CREATION = 2 + IN_PROCESS = 3 + END = 4 + + +gxpacket_spec = [ + ("location", float64[:]), + ("direction", float64[:]), + ("energy_rf", float64), + ("energy_cmf", float64), + ("nu_rf", float64), + ("nu_cmf", float64), + ("status", int64), + ("shell", int64), + ("time_current", float64), + ("tau", float64), +] + + +@jitclass(gxpacket_spec) +class GXPacket(object): + """ + Indivisible gamma-ray packet + """ + + def __init__( + self, + location, + direction, + energy_rf, + energy_cmf, + nu_rf, + nu_cmf, + status, + shell, + time_current, + ): + self.location = location + self.direction = direction + self.energy_rf = energy_rf + self.energy_cmf = energy_cmf + self.nu_rf = nu_rf + self.nu_cmf = nu_cmf + self.status = status + self.shell = shell + self.time_current = time_current + # TODO: rename to tau_event + self.tau = -np.log(np.random.random()) + + def get_location_r(self): + """Calculate radius of the packet + + Returns: + float: packet radius + """ + return np.sqrt( + self.location[0] ** 2.0 + + self.location[1] ** 2.0 + + self.location[2] ** 2.0 + ) diff --git a/tardis/energy_input/calculate_opacity.py b/tardis/energy_input/calculate_opacity.py new file mode 100644 index 00000000000..9aeae20686e --- /dev/null +++ b/tardis/energy_input/calculate_opacity.py @@ -0,0 +1,249 @@ +import numpy as np +import astropy.units as u +import radioactivedecay as rd +from numba import njit + +from tardis import constants as const +from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel +from tardis.energy_input.util import kappa_calculation + +# Force cgs +M_P = const.m_p.to(u.g).value +MASS_SI = rd.Nuclide("Si-28").atomic_mass * M_P +MASS_FE = rd.Nuclide("Fe-56").atomic_mass * M_P +SIGMA_T = const.sigma_T.cgs.value +FINE_STRUCTURE = const.alpha.value + + +@njit(**njit_dict_no_parallel) +def compton_opacity_partial(energy, fraction): + """Partial Compton scattering opacity, from artis file + gamma.cc + + Parameters + ---------- + energy : float + packet energy in terms of electron rest energy + fraction : float + 1 + 2 * packet energy + + Returns + ------- + np.float64 + Compton scattering opacity + """ + term_1 = ( + ((energy**2.0) - (2.0 * energy) - 2.0) + * np.log(fraction) + / energy + / energy + ) + term_2 = (((fraction**2.0) - 1.0) / (fraction**2.0)) / 2.0 + term_3 = ((fraction - 1.0) / energy) * ( + (1 / energy) + (2.0 / fraction) + (1.0 / (energy * fraction)) + ) + return 3.0 * SIGMA_T * (term_1 + term_2 + term_3) / (8 * energy) + + +@njit(**njit_dict_no_parallel) +def compton_opacity_calculation(energy, electron_density): + """Calculate the Compton scattering opacity for a given energy + (Rybicki & Lightman, 1979) + + $ + \\rho / 2 p_m \\times 3/4 \\sigma_T ((1 + \kappa) / \kappa^3 + ((2\kappa(1 + \kappa)) / (1 + 2\kappa) - \ln(1 + 2\kappa) + 1/(2\kappa) \ln(1 + 2\kappa) + - (1 + 3\kappa) / (1 + 2\kappa)^2) + $ + + Parameters + ---------- + energy : float + The energy of the photon + ejecta_density : float + The density of the ejecta + + Returns + ------- + float + The Compton scattering opacity + """ + kappa = kappa_calculation(energy) + + a = 1.0 + 2.0 * kappa + + sigma_KN = ( + 3.0 + / 4.0 + * SIGMA_T + * ( + (1.0 + kappa) + / kappa**3.0 + * ((2.0 * kappa * (1.0 + kappa)) / a - np.log(a)) + + 1.0 / (2.0 * kappa) * np.log(a) + - (1.0 + 3 * kappa) / a**2.0 + ) + ) + + return electron_density * sigma_KN + + +@njit(**njit_dict_no_parallel) +def photoabsorption_opacity_calculation( + energy, ejecta_density, iron_group_fraction +): + """Calculates photoabsorption opacity for a given energy + Approximate treatment from Ambwani & Sutherland (1988) + Magic numbers are from the approximate treatment + + Parameters + ---------- + energy : float + Photon energy + ejecta_density : float + The density of the ejecta + iron_group_fraction : float + Fraction of iron group elements in the shell + + Returns + ------- + float + Photoabsorption opacity + """ + si_opacity = ( + 1.16e-24 + * (energy / 100.0) ** -3.13 + * ejecta_density + / MASS_SI + * (1.0 - iron_group_fraction) + ) + + fe_opacity = ( + 25.7e-24 + * (energy / 100.0) ** -3.0 + * ejecta_density + / MASS_FE + * iron_group_fraction + ) + + return si_opacity + fe_opacity + + +@njit(**njit_dict_no_parallel) +def photoabsorption_opacity_calculation_kasen( + energy, number_density, proton_count +): + """Calculates photoabsorption opacity for a given energy + Approximate treatment from Kasen et al. (2006) + + Parameters + ---------- + energy : float + Photon energy + number_density : float + The number density of the ejecta for each atom + proton_count : float + Number of protons for each atom in the ejecta + + Returns + ------- + float + Photoabsorption opacity + """ + kappa = kappa_calculation(energy) + + opacity = (FINE_STRUCTURE**4.0) * 8.0 * np.sqrt(2) * (kappa**-3.5) + # Note- this should actually be atom_number_density * (atom_proton_number ** 5) + return ( + SIGMA_T + * opacity + * np.sum((number_density / proton_count) * proton_count**5) + ) + + +@njit(**njit_dict_no_parallel) +def pair_creation_opacity_calculation( + energy, ejecta_density, iron_group_fraction +): + """Calculates pair creation opacity for a given energy + Approximate treatment from Ambwani & Sutherland (1988) + + Parameters + ---------- + energy : float + Photon energy + ejecta_density : float + The density of the ejecta + iron_group_fraction : float + Fraction of iron group elements in the shell + + Returns + ------- + float + Pair creation opacity + """ + z_si = 14 + z_fe = 26 + + si_proton_ratio = z_si**2.0 / MASS_SI + fe_proton_ratio = z_fe**2.0 / MASS_FE + + multiplier = ejecta_density * ( + si_proton_ratio * (1.0 - iron_group_fraction) + + fe_proton_ratio * iron_group_fraction + ) + + # Conditions prevent divide by zero + # Ambwani & Sutherland (1988) + if energy > 1022 and energy < 1500: + opacity = multiplier * 1.0063 * (energy / 1000 - 1.022) * 1.0e-27 + elif energy >= 1500: + opacity = ( + multiplier * (0.0481 + 0.301 * (energy / 1000 - 1.5)) * 1.0e-27 + ) + else: + opacity = 0 + + return opacity + + +@njit(**njit_dict_no_parallel) +def pair_creation_opacity_artis(energy, ejecta_density, iron_group_fraction): + """Calculates pair creation opacity for a given energy + Approximate treatment from Ambwani & Sutherland (1988) + as implemented in ARTIS + + Parameters + ---------- + energy : float + Photon energy + ejecta_density : float + The density of the ejecta + iron_group_fraction : float + Fraction of iron group elements in the shell + + Returns + ------- + float + Pair creation opacity + """ + # Conditions prevent divide by zero + # Ambwani & Sutherland (1988) + if energy > 1022: + if energy > 1500: + opacity_si = (0.0481 + (0.301 * (energy - 1500))) * 196.0e-27 + opacity_fe = (0.0481 + (0.301 * (energy - 1500))) * 784.0e-27 + else: + opacity_si = 1.0063 * (energy - 1022) * 196.0e-27 + opacity_fe = 1.0063 * (energy - 1022) * 784.0e-27 + + opacity_si *= ejecta_density / M_P / 28 + opacity_fe *= ejecta_density / M_P / 56 + + opacity = (opacity_fe * iron_group_fraction) + ( + opacity_si * (1.0 - iron_group_fraction) + ) + else: + opacity = 0 + + return opacity diff --git a/tardis/energy_input/energy_source.py b/tardis/energy_input/energy_source.py new file mode 100644 index 00000000000..db715e2118e --- /dev/null +++ b/tardis/energy_input/energy_source.py @@ -0,0 +1,352 @@ +import numpy as np +import pandas as pd +from numba import njit +import radioactivedecay as rd + +from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel +from tardis.util.base import ( + atomic_number2element_symbol, +) +from tardis.energy_input.util import ( + convert_half_life_to_astropy_units, +) + + +@njit(**njit_dict_no_parallel) +def sample_mass(masses, inner_radius, outer_radius): + """Samples location weighted by mass + + Parameters + ---------- + masses : array + Shell masses + inner_radius : array + Inner radii + outer_radius : array + Outer radii + + Returns + ------- + float + Sampled radius + int + Sampled shell index + """ + norm_mass = masses / np.sum(masses) + cdf = np.cumsum(norm_mass) + shell = np.searchsorted(cdf, np.random.random()) + + z = np.random.random() + radius = ( + z * inner_radius[shell] ** 3.0 + (1.0 - z) * outer_radius[shell] ** 3.0 + ) ** (1.0 / 3.0) + + return radius, shell + + +@njit(**njit_dict_no_parallel) +def create_energy_cdf(energy, intensity): + """Creates a CDF of given intensities + + Parameters + ---------- + energy : One-dimensional Numpy Array, dtype float + Array of energies + intensity : One-dimensional Numpy Array, dtype float + Array of intensities + + Returns + ------- + One-dimensional Numpy Array, dtype float + Sorted energy array + One-dimensional Numpy Array, dtype float + CDF where each index corresponds to the energy in + the sorted array + """ + energy.sort() + sorted_indices = np.argsort(energy) + sorted_intensity = intensity[sorted_indices] + norm_intensity = sorted_intensity / np.sum(sorted_intensity) + cdf = np.cumsum(norm_intensity) + + return energy, cdf + + +@njit(**njit_dict_no_parallel) +def sample_energy_distribution(energy_sorted, cdf): + """Randomly samples a CDF of energies + + Parameters + ---------- + energy_sorted : One-dimensional Numpy Array, dtype float + Sorted energy array + cdf : One-dimensional Numpy Array, dtype float + CDF + + Returns + ------- + float + Sampled energy + """ + index = np.searchsorted(cdf, np.random.random()) + + return energy_sorted[index] + + +def setup_input_energy(nuclear_data, source): + """Sets up energy distribution and CDF for a + source of decay radiation. + + Parameters + ---------- + nuclear_data : Pandas dataframe + Dataframe of nuclear decay properties + source : str + Type of decay radiation + + Returns + ------- + One-dimensional Numpy Array, dtype float + Sorted energy array + One-dimensional Numpy Array, dtype float + CDF where each index corresponds to the energy in + the sorted array + """ + intensity = nuclear_data.query("type==" + source)["intensity"].values + energy = nuclear_data.query("type==" + source)["energy"].values + energy_sorted, cdf = create_energy_cdf(energy, intensity) + + return energy_sorted, cdf + + +def intensity_ratio(nuclear_data, source_1, source_2): + """Determined the ratio of intensities between two + sources of decay radiation + + Parameters + ---------- + nuclear_data : pandas.Dataframe + Dataframe of nuclear decay properties + source_1 : str + Type of decay radiation to compare + source_2 : str + Type of decay radiation to compare + + Returns + ------- + float + Fractional intensity of source_1 + float + Fractional intensity of source_2 + float + Number of decay products per decay + """ + intensity_1 = nuclear_data.query("type==" + source_1)["intensity"].values + intensity_2 = nuclear_data.query("type==" + source_2)["intensity"].values + total_intensity = np.sum(intensity_1) + np.sum(intensity_2) + # Factor of 100 is because intensities are in percent + scale_factor = total_intensity / 100 + return ( + np.sum(intensity_1) / total_intensity, + np.sum(intensity_2) / total_intensity, + scale_factor, + ) + + +def ni56_chain_energy( + taus, time_start, time_end, number_ni56, ni56_lines, co56_lines +): + """Calculate the energy from the Ni56 chain + + Parameters + ---------- + taus : array float64 + Mean half-life for each isotope + time_start : float + Start time in days + time_end : float + End time in days + number_ni56 : int + Number of Ni56 atoms at time_start + ni56_lines : DataFrame + Ni56 lines and intensities + co56_lines : DataFrame + Co56 lines and intensities + + Returns + ------- + float + Total energy from Ni56 decay + """ + total_ni56 = -taus["Ni56"] * ( + np.exp(-time_end / taus["Ni56"]) - np.exp(-time_start / taus["Ni56"]) + ) + total_co56 = -taus["Co56"] * ( + np.exp(-time_end / taus["Co56"]) - np.exp(-time_start / taus["Co56"]) + ) + + total_energy = pd.DataFrame() + + total_energy["Ni56"] = number_ni56 * ( + (ni56_lines.energy * 1000 * ni56_lines.intensity).sum() + / taus["Ni56"] + * total_ni56 + ) + + total_energy["Co56"] = number_ni56 * ( + (co56_lines.energy * 1000 * co56_lines.intensity).sum() + / (taus["Ni56"] - taus["Co56"]) + * (total_ni56 - total_co56) + ) + + return total_energy + + +def ni56_chain_energy_choice( + taus, time_start, time_end, number_ni56, ni56_lines, co56_lines, isotope +): + """Calculate the energy from the Ni56 or Co56 chain + + Parameters + ---------- + taus : array float64 + Mean half-life for each isotope + time_start : float + Start time in days + time_end : float + End time in days + number_ni56 : int + Number of Ni56 atoms at time_start + ni56_lines : DataFrame + Ni56 lines and intensities + co56_lines : DataFrame + Co56 lines and intensities + isotope : string + Isotope chain to calculate energy for + + Returns + ------- + float + Total energy from decay chain + """ + total_ni56 = -taus["Ni56"] * ( + np.exp(-time_end / taus["Ni56"]) - np.exp(-time_start / taus["Ni56"]) + ) + total_co56 = -taus["Co56"] * ( + np.exp(-time_end / taus["Co56"]) - np.exp(-time_start / taus["Co56"]) + ) + + if isotope == "Ni56": + total_energy = number_ni56 * ( + (ni56_lines.energy * 1000 * ni56_lines.intensity).sum() + / taus["Ni56"] + * total_ni56 + ) + else: + total_energy = number_ni56 * ( + (co56_lines.energy * 1000 * co56_lines.intensity).sum() + / (taus["Ni56"] - taus["Co56"]) + * (total_ni56 - total_co56) + ) + + return total_energy + + +def get_all_isotopes(abundances): + """Get the possible isotopes present over time + for a given starting abundance + + Parameters + ---------- + abundances : DataFrame + Current isotope abundances + + Returns + ------- + list + List of isotope names + """ + progenitors = [ + f"{rd.utils.Z_DICT[i[0]]}-{i[1]}" for i in abundances.T.columns + ] + + isotopes = set(progenitors) + check = True + + while check == True: + progeny = set(isotopes) + + for i in isotopes: + for p in rd.Nuclide(i).progeny(): + if p != "SF": + progeny.add(p) + + if progeny == isotopes: + check = False + else: + isotopes |= progeny + + isotopes = [i.replace("-", "") for i in isotopes] + return isotopes + + +def get_tau(meta, isotope_string): + """Calculate the mean lifetime of an isotope + + Parameters + ---------- + meta : DataFrame + Isotope metadata + isotope_string : str + Isotope of interest + + Returns + ------- + float + Mean lifetime of isotope + """ + isotope_meta = meta.loc[isotope_string] + half_life = isotope_meta.loc[isotope_meta["key"] == "Parent T1/2 value"][ + "value" + ].values[0] + half_life = convert_half_life_to_astropy_units(half_life) + return half_life / np.log(2) + + +def get_isotope_string(atom_number, atom_mass): + """Get the isotope string in the format e.g. Ni56 + + Parameters + ---------- + atom_number : int + Atomic number + atom_mass : int + Atomic mass + + Returns + ------- + str + Isotope string in the format e.g. Ni56 + """ + return atomic_number2element_symbol(atom_number) + str(atom_mass) + + +def read_artis_lines(isotope, path_to_data): + """Reads lines of ARTIS format + + Parameters + ---------- + isotope : string + Isotope to read e.g. Ni56 + + Returns + ------- + pd.DataFrame + Energies and intensities of the isotope lines + """ + return pd.read_csv( + path_to_data + isotope + ".txt", + names=["energy", "intensity"], + sep=" ", + index_col=False, + ) diff --git a/tardis/energy_input/gamma_ray_estimators.py b/tardis/energy_input/gamma_ray_estimators.py new file mode 100644 index 00000000000..cb484886b72 --- /dev/null +++ b/tardis/energy_input/gamma_ray_estimators.py @@ -0,0 +1,139 @@ +import numpy as np +from numba import njit + +from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel +from tardis.energy_input.util import ( + angle_aberration_gamma, + doppler_gamma, + H_CGS_KEV, + ELECTRON_MASS_ENERGY_KEV, + kappa_calculation, +) +from tardis.energy_input.calculate_opacity import ( + compton_opacity_calculation, + SIGMA_T, + photoabsorption_opacity_calculation, +) + + +def compton_emissivity_estimator(packet, distance): + """Compton scattering emissivity estimator for integral + calculations + + Parameters + ---------- + packet : GXPacket + Packet that needs its emissivity calculated + distance : float64 + Distance packet has travelled + + Returns + ------- + float64, int + Unnormalized emissivity estimator, line index + """ + + cmf_direction = angle_aberration_gamma( + packet.get_direction_vector(), packet.location_r + ) + + cmf_angle = np.dot(cmf_direction, [1, 0, 0]) + + frequency_factor = ( + 1 + + H_CGS_KEV * packet.nu_cmf / ELECTRON_MASS_ENERGY_KEV + + (1.0 - cmf_angle) + ) + + line_index = GET_NEAREST_LINE_REDWARD_FUNCTION( + packet.nu_cmf / frequency_factor + ) + + partial_cross_section = ( + 3.0 + / 16.0 + / np.pi + * SIGMA_T + / frequency_factor + / frequency_factor + * (frequency_factor + (1.0 / frequency_factor) + cmf_angle**2.0 - 1.0) + ) + + doppler_factor = doppler_gamma( + packet.get_direction_vector(), packet.location_r + ) + + emissivity = ( + packet.energy_rf + * partial_cross_section + * distance + * doppler_factor**2.0 + / frequency_factor + ) + + return emissivity, line_index + + +def pair_creation_estimator(packet, pair_creation_opacity, distance): + """Calculates the emissivity for pair creation gamma-rays + + Parameters + ---------- + packet : GXPacket + Packet that needs its emissivity calculated + pair_creation_opacity : float64 + Opacity of the pair creation process + distance : float64 + Distance packet has travelled + + Returns + ------- + float64 + Emissivity estimator + """ + normalized_energy = ( + 2 * ELECTRON_MASS_ENERGY_KEV / (H_CGS_KEV * packet.nu_cmf) + ) + + emissivity = ( + pair_creation_opacity * normalized_energy * packet.energy_rf * distance + ) + + return emissivity + + +@njit(**njit_dict_no_parallel) +def get_average_compton_fraction(energy): + def f(x, mu): + return 1.0 / (1.0 + x * (1.0 - mu)) + + def cross_section(x, mu): + return ( + (3.0 * SIGMA_T) + / (16.0 * np.pi) + * f(x, mu) ** 2.0 + * (f(x, mu) + 1.0 / f(x, mu) - (1.0 - mu**2)) + ) + + x = kappa_calculation(energy) + mus = np.linspace(-1, 1, 100) + dmu = mus[1] - mus[0] + sum = 0 + norm = 0 + + for mu in mus: + sum += cross_section(x, mu) * f(x, mu) * dmu + norm += cross_section(x, mu) * dmu + + integral = 1.0 - sum / norm + + return 1 - integral + + +@njit(**njit_dict_no_parallel) +def deposition_estimator_kasen(energy, ejecta_density, iron_group_fraction): + return get_average_compton_fraction(energy) * compton_opacity_calculation( + energy, ejecta_density + ) + photoabsorption_opacity_calculation( + energy, ejecta_density, iron_group_fraction + ) diff --git a/tardis/energy_input/gamma_ray_grid.py b/tardis/energy_input/gamma_ray_grid.py new file mode 100644 index 00000000000..5c8a09b2246 --- /dev/null +++ b/tardis/energy_input/gamma_ray_grid.py @@ -0,0 +1,144 @@ +import numpy as np +from numba import njit + +from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel +from tardis.energy_input.util import ( + doppler_gamma, + solve_quadratic_equation, + C_CGS, +) + + +@njit(**njit_dict_no_parallel) +def calculate_distance_radial(photon, r_inner, r_outer): + """ + Calculates 3D distance to shell from gamma ray position + + Parameters + ---------- + photon : GXPhoton object + r_inner : float + r_outer : float + + Returns + ------- + distance : float + + """ + + # solve the quadratic distance equation for the inner and + # outer shell boundaries + inner_1, inner_2 = solve_quadratic_equation( + photon.location, photon.direction, r_inner + ) + outer_1, outer_2 = solve_quadratic_equation( + photon.location, photon.direction, r_outer + ) + + final_position_inner_1 = photon.location + photon.direction * inner_1 + final_position_inner_2 = photon.location + photon.direction * inner_2 + final_position_outer_1 = photon.location + photon.direction * outer_1 + final_position_outer_2 = photon.location + photon.direction * outer_2 + + if np.dot(final_position_inner_1, photon.direction) > 0: + inner_1 = -1 + if np.dot(final_position_inner_2, photon.direction) > 0: + inner_2 = -1 + if np.dot(final_position_outer_1, photon.direction) < 0: + outer_1 = -1 + if np.dot(final_position_outer_2, photon.direction) < 0: + outer_2 = -1 + + distances = np.array([inner_1, inner_2, outer_1, outer_2]) + + # the correct distance is the shortest positive distance + distance_list = [i for i in distances if i > 0] + + if not distance_list: + print(photon.get_location_r() - r_inner) + print(photon.get_location_r() - r_outer) + print(photon.get_location_r()) + print(photon.location, photon.direction, r_inner, r_outer) + print(distances) + print(photon.shell) + raise ValueError("No root found for distance calculation!") + + shortest = min(distance_list) + shell_change = 1 + + if shortest == (inner_1 or inner_2): + shell_change = -1 + + return shortest, shell_change + + +@njit(**njit_dict_no_parallel) +def distance_trace( + photon, + inner_velocity, + outer_velocity, + total_opacity, + current_time, + next_time, +): + """ + Traces distance traveled by gamma ray and finds distance to + next interaction and boundary + + Parameters + ---------- + photon : GXPhoton object + inner_velocity : One dimensional Numpy array, dtype float + outer_velocity : One dimensional Numpy array, dtype float + total_opacity : float + current_time : float + next_time : float + + Returns + ------- + distance_interaction : float + distance_boundary : float + distance_time : float + shell_change : int + """ + distance_boundary, shell_change = calculate_distance_radial( + photon, + inner_velocity[photon.shell] * current_time, + outer_velocity[photon.shell] * current_time, + ) + + distance_interaction = photon.tau / total_opacity + distance_time = (next_time - photon.time_current) * C_CGS + return distance_interaction, distance_boundary, distance_time, shell_change + + +@njit(**njit_dict_no_parallel) +def move_packet(packet, distance): + """ + Moves packet a distance along its direction vector + + Parameters + ---------- + packet : GXPacket object + distance : float + + Returns + ------- + packet : GXPacket object + + """ + location_old = packet.location + direction = packet.direction + + location_new = location_old + distance * direction + + packet.location = location_new + + doppler_factor = doppler_gamma( + packet.direction, packet.location, packet.time_current + ) + + packet.nu_cmf = packet.nu_rf * doppler_factor + packet.energy_cmf = packet.energy_rf * doppler_factor + + return packet diff --git a/tardis/energy_input/gamma_ray_interactions.py b/tardis/energy_input/gamma_ray_interactions.py new file mode 100644 index 00000000000..d93d28bdbba --- /dev/null +++ b/tardis/energy_input/gamma_ray_interactions.py @@ -0,0 +1,311 @@ +import numpy as np +from numba import njit + +from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel +from tardis.energy_input.util import ( + get_random_unit_vector, + kappa_calculation, + euler_rodrigues, + compton_theta_distribution, + get_perpendicular_vector, + angle_aberration_gamma, + doppler_gamma, + ELECTRON_MASS_ENERGY_KEV, + H_CGS_KEV, +) +from tardis.energy_input.GXPacket import GXPacketStatus +from tardis.energy_input.calculate_opacity import compton_opacity_partial + + +@njit(**njit_dict_no_parallel) +def get_compton_angle(energy): + """ + Computes the compton angle from the Klein-Nishina equation. + Computes the lost energy due to this angle + + Parameters + ---------- + energy : float + Photon energy + + Returns + ------- + compton_angle : float + Compton scattering angle + lost_energy : float + Energy lost based on angle + new_energy : float + Photon energy + """ + theta_angles, theta_distribution = compton_theta_distribution(energy) + + z = np.random.random() + + # get Compton scattering angle + compton_angle = np.interp(z, theta_distribution, theta_angles) + # Energy calculations + new_energy = energy / ( + 1.0 + kappa_calculation(energy) * (1.0 - np.cos(compton_angle)) + ) + lost_energy = energy - new_energy + + return compton_angle, lost_energy, new_energy + + +@njit(**njit_dict_no_parallel) +def get_compton_fraction(energy): + """ + Computes the compton angle from the Klein-Nishina equation. + Determines the probability of absorption from this angle. + + Parameters + ---------- + energy : float + Photon energy + + Returns + ------- + compton_angle : float + Compton scattering angle + compton_fraction : float + Fraction of energy lost + """ + theta_angles, theta_distribution = compton_theta_distribution(energy) + + z = np.random.random() + + # get Compton scattering angle + compton_angle = np.interp(z, theta_distribution, theta_angles) + # Energy calculations + fraction = 1 / ( + 1.0 + kappa_calculation(energy) * (1.0 - np.cos(compton_angle)) + ) + + return compton_angle, fraction + + +@njit(**njit_dict_no_parallel) +def get_compton_fraction_artis(energy): + """Gets the Compton scattering/absorption fraction + and angle following the scheme in ARTIS + + Parameters + ---------- + energy : float + Energy of the gamma-ray + + Returns + ------- + float + Scattering angle + float + Compton scattering fraction + """ + energy_norm = kappa_calculation(energy) + + fraction_max = 1.0 + 2.0 * energy_norm + fraction_min = 1.0 + + normalization = np.random.random() * compton_opacity_partial( + energy_norm, fraction_max + ) + + epsilon = 1.0e20 + count = 0 + + while epsilon > 1.0e-4: + fraction_try = (fraction_max + fraction_min) / 2.0 + sigma_try = compton_opacity_partial(energy_norm, fraction_try) + + if sigma_try > normalization: + fraction_max = fraction_try + epsilon = (sigma_try - normalization) / normalization + else: + fraction_min = fraction_try + epsilon = (normalization - sigma_try) / normalization + + count += 1 + if count > 1000: + print("Error, failure to get a Compton fraction") + break + + angle = np.arccos(1.0 - ((fraction_try - 1) / energy_norm)) + + return angle, fraction_try + + +@njit(**njit_dict_no_parallel) +def get_compton_fraction_urilight(energy): + """Gets the Compton scattering/absorption fraction + and angle following the scheme in Urilight + + Parameters + ---------- + energy : float + Energy of the gamma-ray + + Returns + ------- + float + Scattering angle + float + Compton scattering fraction + """ + E0 = kappa_calculation(energy) + + x0 = 1.0 / (1.0 + 2.0 * E0) + + accept = False + while not accept: + + z = np.random.random(3) + alpha1 = np.log(1.0 / x0) + alpha2 = (1.0 - x0**2.0) / 2.0 + if z[1] < alpha1 / (alpha1 + alpha2): + x = x0 ** z[2] + else: + x = np.sqrt(x0**2.0 + (1.0 - x0**2.0) * z[2]) + + f = (1.0 - x) / x / E0 + sin2t = f * (2.0 - f) + ge = 1.0 - x / (1 + x**2.0) * sin2t + if ge > z[3]: + accept = True + + cost = 1.0 - f + + return np.arccos(cost), x + + +@njit(**njit_dict_no_parallel) +def compton_scatter(photon, compton_angle): + """ + Changes the direction of the gamma-ray by the Compton scattering angle + + Parameters + ---------- + photon : GXPhoton object + compton_angle : float + + Returns + ------- + float64 + Photon theta direction + float64 + Photon phi direction + """ + + # get comoving frame direction + comov_direction = angle_aberration_gamma( + photon.direction, photon.location, photon.time_current + ) + + # compute an arbitrary perpendicular vector to the comoving direction + orthogonal_vector = get_perpendicular_vector(comov_direction) + # determine a random vector with compton_angle to the comoving direction + new_vector = np.dot( + euler_rodrigues(compton_angle, orthogonal_vector), + comov_direction, + ) + + # draw a random angle from [0,2pi] + phi = 2.0 * np.pi * np.random.random() + # rotate the vector with compton_angle around the comoving direction + final_compton_scattered_vector = np.dot( + euler_rodrigues(phi, comov_direction), new_vector + ) + + norm_phi = np.dot( + final_compton_scattered_vector, final_compton_scattered_vector + ) + + norm_theta = np.dot(final_compton_scattered_vector, comov_direction) + + assert ( + np.abs(norm_phi - 1) < 1e-8 + ), "Error, norm of Compton scatter vector is not 1!" + + assert ( + np.abs(norm_theta - np.cos(compton_angle)) < 1e-8 + ), "Error, difference between new vector angle and Compton angle is more than 0!" + + # Calculate the angle aberration of the final direction + final_direction = angle_aberration_gamma( + final_compton_scattered_vector, photon.location, photon.time_current + ) + + return final_direction + + +@njit(**njit_dict_no_parallel) +def pair_creation_packet(packet): + """ + Pair creation randomly scatters the packet + or destroys it, based on the frequency + + Parameters + ---------- + packet : GXPacket + incoming packet + + Returns + ------- + GXPacket + outgoing packet + """ + + probability_gamma = ( + 2 * ELECTRON_MASS_ENERGY_KEV / (H_CGS_KEV * packet.nu_cmf) + ) + + if np.random.random() > probability_gamma: + packet.status = GXPacketStatus.PHOTOABSORPTION + return packet + + new_direction = get_random_unit_vector() + + # Calculate aberration of the random angle for the rest frame + final_direction = angle_aberration_gamma( + new_direction, packet.location, -1 * packet.time_current + ) + + packet.direction = final_direction + + doppler_factor = doppler_gamma( + packet.direction, packet.location, packet.time_current + ) + + packet.nu_cmf = ELECTRON_MASS_ENERGY_KEV / H_CGS_KEV + packet.nu_rf = packet.nu_cmf / doppler_factor + packet.energy_rf = packet.energy_cmf / doppler_factor + + return packet + + +@njit(**njit_dict_no_parallel) +def scatter_type(compton_opacity, photoabsorption_opacity, total_opacity): + """ + Determines the scattering type based on process opacities + + Parameters + ---------- + compton_opacity : float + photoabsorption_opacity : float + total_opacity : float + + Returns + ------- + status : GXPacketStatus + Scattering process the photon encounters + + """ + z = np.random.random() + + if z <= (compton_opacity / total_opacity): + status = GXPacketStatus.COMPTON_SCATTER + elif z <= (compton_opacity + photoabsorption_opacity) / total_opacity: + status = GXPacketStatus.PHOTOABSORPTION + else: + status = GXPacketStatus.PAIR_CREATION + + return status diff --git a/tardis/energy_input/indivisible_packets.py b/tardis/energy_input/indivisible_packets.py new file mode 100644 index 00000000000..34b7b192f93 --- /dev/null +++ b/tardis/energy_input/indivisible_packets.py @@ -0,0 +1,920 @@ +import numpy as np +from tqdm.auto import tqdm +import pandas as pd +import astropy.units as u +from numba import njit +from numba.typed import List + +from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel +from tardis.energy_input.GXPacket import GXPacket, GXPacketStatus +from tardis.energy_input.gamma_ray_grid import ( + distance_trace, + move_packet, +) +from tardis.energy_input.energy_source import ( + ni56_chain_energy, + read_artis_lines, +) +from tardis.energy_input.calculate_opacity import ( + compton_opacity_calculation, + photoabsorption_opacity_calculation, + pair_creation_opacity_calculation, + photoabsorption_opacity_calculation_kasen, + pair_creation_opacity_artis, + SIGMA_T, +) +from tardis.energy_input.gamma_ray_interactions import ( + get_compton_fraction_artis, + scatter_type, + compton_scatter, + pair_creation_packet, +) +from tardis.energy_input.util import ( + doppler_gamma, + C_CGS, + H_CGS_KEV, + kappa_calculation, + get_index, + get_random_unit_vector, +) +from tardis import constants as const +from tardis.energy_input.gamma_ray_estimators import deposition_estimator_kasen + +# Energy: keV, exported as eV for SF solver +# distance: cm +# mass: g +# time: s + + +@njit(**njit_dict_no_parallel) +def sample_energy(energy, intensity): + """Samples energy from energy and intensity + + Parameters + ---------- + energy : One-dimensional Numpy Array, dtype float + Array of energies + intensity : One-dimensional Numpy Array, dtype float + Array of intensities + + Returns + ------- + float + Energy + """ + z = np.random.random() + + average = (energy * intensity).sum() + total = 0 + for (e, i) in zip(energy, intensity): + total += e * i / average + if z <= total: + return e + + return False + + +@njit(**njit_dict_no_parallel) +def sample_decay_time(isotope, ni56_tau, co56_tau): + """Samples the decay time from the mean half-life + of the isotopes (needs restructuring for more isotopes) + + Parameters + ---------- + isotope : string + Isotope as e.g. Ni56 + taus : dict + Mean half-life for each isotope + + Returns + ------- + float + Decay time in seconds + """ + if isotope == "Ni56": + decay_time = -ni56_tau * np.log(np.random.random()) + else: + decay_time = -ni56_tau * np.log(np.random.random()) - co56_tau * np.log( + np.random.random() + ) + return decay_time + + +@njit(**njit_dict_no_parallel) +def initialize_packets( + decays_per_shell, + input_energy, + ni56_lines, + co56_lines, + inner_velocities, + outer_velocities, + inv_volume_time, + times, + energy_df_rows, + effective_times, + ni56_tau, + co56_tau, +): + """Initialize a list of GXPacket objects for the simulation + to operate on. + + Parameters + ---------- + decays_per_shell : array int64 + Number of decays per simulation shell + input_energy : float64 + Total input energy from decay + ni56_lines : array float64 + Lines and intensities for Ni56 + co56_lines : array float64 + Lines and intensities for Co56 + inner_velocities : array float64 + Inner velocities of the shells + outer_velocities : array float64 + Outer velocities of the shells + inv_volume_time : array float64 + Inverse volume with time + times : array float64 + Simulation time steps + energy_df_rows : list + Setup list for energy DataFrame output + effective_times : array float64 + Middle time of the time step + ni56_tau, co56_tau : float64 + Mean half-life for each isotope + + Returns + ------- + list + List of GXPacket objects + array + Array of main output dataframe rows + array + Array of plotting output dataframe rows + array + Array of positron output dataframe rows + """ + + packets = List() + + number_of_packets = np.sum(decays_per_shell) + + print("Total packets:", number_of_packets) + + packet_energy = input_energy / number_of_packets + + print("Energy per packet", packet_energy) + + ni56_energy = (ni56_lines[:, 0] * ni56_lines[:, 1]).sum() + co56_energy = (co56_lines[:, 0] * co56_lines[:, 1]).sum() + + ni56_fraction = ni56_energy / (ni56_energy + co56_energy) + + energy_plot_df_rows = np.zeros((number_of_packets, 8)) + energy_plot_positron_rows = np.zeros((number_of_packets, 4)) + + j = 0 + for k, shell in enumerate(decays_per_shell): + z = np.random.random(shell) + + initial_radii = ( + z * inner_velocities[k] ** 3.0 + + (1.0 - z) * outer_velocities[k] ** 3.0 + ) ** (1.0 / 3.0) + + for i in range(shell): + decay_time = np.inf + while decay_time > times[-1]: + ni_or_co_random = np.random.random() + + if ni_or_co_random > ni56_fraction: + decay_type = "Co56" + energy = co56_lines[:, 0] * 1000 + intensity = co56_lines[:, 1] + # positron energy scaled by intensity + positron_energy = 0.63 * 1000 * 0.19 + positron_fraction = positron_energy / np.sum( + energy * intensity + ) + else: + decay_type = "Ni56" + energy = ni56_lines[:, 0] * 1000 + intensity = ni56_lines[:, 1] + positron_fraction = 0 + + decay_time = sample_decay_time(decay_type, ni56_tau, co56_tau) + + cmf_energy = sample_energy(energy, intensity) + + energy_factor = 1.0 + if decay_time < times[0]: + energy_factor = decay_time / times[0] + decay_time = times[0] + + decay_time_index = get_index(decay_time, times) + + energy_df_rows[k, decay_time_index] += ( + positron_fraction * packet_energy * 1000 + ) + + scaled_r = initial_radii[i] * effective_times[decay_time_index] + + initial_location = scaled_r * get_random_unit_vector() + initial_direction = get_random_unit_vector() + initial_energy = packet_energy * energy_factor + + # draw a random gamma-ray in shell + packet = GXPacket( + initial_location, + initial_direction, + 1.0, + initial_energy, + 0.0, + 0.0, + GXPacketStatus.IN_PROCESS, + k, + decay_time, + ) + + packet.energy_rf = packet.energy_cmf / doppler_gamma( + packet.direction, + packet.location, + packet.time_current, + ) + + energy_plot_df_rows[j] = np.array( + [ + i, + packet.energy_rf, + packet.get_location_r(), + packet.time_current, + int(packet.status), + 0, + 0, + 0, + ] + ) + + energy_plot_positron_rows[j] = [ + j, + positron_fraction * packet_energy * 1000, + # * inv_volume_time[packet.shell, decay_time_index], + packet.get_location_r(), + packet.time_current, + ] + + packet.nu_cmf = cmf_energy / H_CGS_KEV + + packet.nu_rf = packet.nu_cmf / doppler_gamma( + packet.direction, + packet.location, + packet.time_current, + ) + + packets.append(packet) + + j += 1 + + return ( + packets, + energy_df_rows, + energy_plot_df_rows, + energy_plot_positron_rows, + ) + + +def main_gamma_ray_loop( + num_decays, + model, + plasma, + time_steps=10, + time_end=80.0, + grey_opacity=-1, + spectrum_bins=500, + time_space="log", + photoabsorption_opacity="tardis", + pair_creation_opacity="tardis", + seed=1, + path_to_artis_lines="", +): + """Main loop that determines the gamma ray propagation + + Parameters + ---------- + num_decays : int + Number of decays requested + model : tardis.Radial1DModel + The tardis model to calculate gamma ray propagation through + plasma : tardis.plasma.BasePlasma + The tardis plasma with calculated atomic number density + time_steps : int + Number of time steps requested + time_end : float + End time of simulation in days + grey_opacity : float + Grey photoabsorption opacity for gamma-rays in cm^2 g^-1, set to -1 to turn off + spectrum_bins : int + Number of energy bins for the gamma-ray spectrum + time_space : str + `'log'` for log-space time steps, otherwise time steps will be linear + photoabsorption_opacity : str + Set the photoabsorption opacity calculation. + Defaults to Ambwani & Sutherland (1988) approximation. + `'kasen'` uses the Kasen et al. 2006 method. + pair_creation_opacity : str + Set the pair creation opacity calculation. Defaults to Ambwani & Sutherland (1988) approximation. + `'artis'` uses the ARTIS implementation of the Ambwani & Sutherland (1988) approximation. + + Returns + ------- + pandas.DataFrame + Energy per shell per time in units of eV/s/cm^-3 + pandas.DataFrame + Columns: + packet index, + Energy input per packet, + radius of deposition, + time of deposition, + compton opacity, + photoabsorption opacity, + pair creation opacity + pandas.DataFrame + Energy of escaping packets + numpy.ndarray + Packets emitted per shell + pandas.DataFrame + Energy from positrons + pandas.DataFrame + Estimated energy deposition in units of keV/s/cm^-3 + """ + # Note: not best numpy practice, but works better in numba than the alternatives + np.random.seed(seed) + + # Enforce cgs + outer_velocities = model.v_outer.to("cm/s").value + inner_velocities = model.v_inner.to("cm/s").value + ejecta_density = model.density.to("g/cm^3").value + ejecta_volume = model.volume.to("cm^3").value + ejecta_velocity_volume = ( + 4 * np.pi / 3 * (outer_velocities**3.0 - inner_velocities**3.0) + ) + time_explosion = model.time_explosion.to("s").value + number_of_shells = model.no_of_shells + raw_isotope_abundance = model.raw_isotope_abundance + + shell_masses = ejecta_volume * ejecta_density + + time_start = time_explosion + time_end *= u.d.to(u.s) + + assert ( + time_start < time_end + ), "Error, simulation start time greater than end time!" + + if time_space == "log": + times = np.zeros(time_steps + 1) + + # log time steps + for i in range(time_steps + 1): + times[i] = ( + np.log(time_start) + + (np.log(time_end) - np.log(time_start)) / time_steps * i + ) + times[i] = np.exp(times[i]) + else: + times = np.linspace(time_start, time_end, time_steps + 1) + + dt_array = np.diff(times) + effective_time_array = np.array( + [np.sqrt(times[i] * times[i + 1]) for i in range(time_steps)] + ) + + # Use isotopic number density + for atom_number in plasma.isotope_number_density.index.get_level_values(0): + plasma.number_density.loc[ + atom_number + ] = plasma.isotope_number_density.loc[atom_number].values + + # Calculate electron number density + electron_number_density = ( + plasma.number_density.mul(plasma.number_density.index, axis=0) + ).sum() + electron_number_density_time = np.zeros( + (len(ejecta_velocity_volume), len(effective_time_array)) + ) + mass_density_time = np.zeros( + (len(ejecta_velocity_volume), len(effective_time_array)) + ) + electron_number = (electron_number_density * ejecta_volume).to_numpy() + inv_volume_time = np.zeros( + (len(ejecta_velocity_volume), len(effective_time_array)) + ) + + # Pre-calculate quantities as they change with time + for i, t in enumerate(effective_time_array): + inv_volume_time[:, i] = (1.0 / ejecta_velocity_volume) / (t**3.0) + mass_density_time[:, i] = shell_masses * inv_volume_time[:, i] + electron_number_density_time[:, i] = ( + electron_number * inv_volume_time[:, i] + ) + + energy_df_rows = np.zeros((number_of_shells, time_steps)) + + # Calculate number of packets per shell based on the mass of Ni56 + mass_ni56 = raw_isotope_abundance.loc[(28, 56)] * shell_masses + number_ni56 = mass_ni56 / 56 * const.N_A + total_number_ni56 = number_ni56.sum() + + decayed_packet_count = np.zeros(len(number_ni56), dtype=np.int64) + + for i, shell in enumerate(number_ni56): + decayed_packet_count[i] = round(num_decays * shell / total_number_ni56) + + # hardcoded Ni and Co 56 mean half-lives + taus = { + "Ni56": 6.075 * u.d.to("s") / np.log(2), + "Co56": 77.233 * u.d.to("s") / np.log(2), + } + + # This will use the decay radiation database and be a more complex network eventually + ni56_lines = read_artis_lines("ni56", path_to_artis_lines) + co56_lines = read_artis_lines("co56", path_to_artis_lines) + + # urilight chooses to have 0 as the baseline for this calculation + # but time_start should also be valid + total_energy = ni56_chain_energy( + taus, 0, time_end, number_ni56, ni56_lines, co56_lines + ) + + print("Total gamma-ray energy") + print(total_energy.sum().sum() * u.keV.to("erg")) + + print("Total positron energy") + print(total_energy["Co56"].sum(axis=0) * 0.0337 * u.keV.to("erg")) + + # Taking iron group to be elements 21-30 + # Used as part of the approximations for photoabsorption and pair creation + # Dependent on atomic data + iron_group_fraction_per_shell = model.abundance.loc[(21):(30)].sum(axis=0) + + # Need to update volume for positron deposition to be time-dependent + print("Initializing packets") + ( + packets, + energy_df_rows, + energy_plot_df_rows, + energy_plot_positron_rows, + ) = initialize_packets( + decayed_packet_count, + total_energy.sum().sum(), + ni56_lines.to_numpy(), + co56_lines.to_numpy(), + inner_velocities, + outer_velocities, + inv_volume_time, + times, + energy_df_rows, + effective_time_array, + taus["Ni56"], + taus["Co56"], + ) + + print("Total positron energy from packets") + print((energy_df_rows).sum().sum() * u.eV.to("erg")) + + total_cmf_energy = 0 + total_rf_energy = 0 + + for p in packets: + total_cmf_energy += p.energy_cmf + total_rf_energy += p.energy_rf + + print("Total CMF energy") + print(total_cmf_energy) + + # Below is the Artis compensation for their method of packet rejection + """ + energy_ratio = total_energy.sum().sum() / total_cmf_energy + + print("Energy ratio") + print(energy_ratio) + + for p in packets: + p.energy_cmf *= energy_ratio + p.energy_rf *= energy_ratio + + for e in energy_df_rows: + e *= energy_ratio + + for row in energy_plot_df_rows: + row[1] *= energy_ratio + """ + print("Total RF energy") + print(total_rf_energy) + + energy_bins = np.logspace(2, 4, spectrum_bins) + energy_out = np.zeros((len(energy_bins - 1), time_steps)) + + # Process packets + ( + energy_df_rows, + energy_plot_df_rows, + energy_out, + deposition_estimator, + ) = gamma_packet_loop( + packets, + grey_opacity, + photoabsorption_opacity, + pair_creation_opacity, + electron_number_density_time, + mass_density_time, + inv_volume_time, + iron_group_fraction_per_shell.to_numpy(), + inner_velocities, + outer_velocities, + times, + dt_array, + effective_time_array, + energy_bins, + energy_df_rows, + energy_plot_df_rows, + energy_out, + ) + + # DataFrame of energy information + energy_plot_df = pd.DataFrame( + data=energy_plot_df_rows, + columns=[ + "packet_index", + "energy_input", + "energy_input_r", + "energy_input_time", + "energy_input_type", + "compton_opacity", + "photoabsorption_opacity", + "total_opacity", + ], + ) + + # DataFrame of positron energies + energy_plot_positrons = pd.DataFrame( + data=energy_plot_positron_rows, + columns=[ + "packet_index", + "energy_input", + "energy_input_r", + "energy_input_time", + ], + ) + + # DataFrame of estimated deposition + # Multiply dataframes by inv_volume_time array + # if per unit volume is needed + energy_estimated_deposition = ( + pd.DataFrame(data=deposition_estimator, columns=times[:-1]) + ) / dt_array + + # Energy is eV/s + energy_df = pd.DataFrame(data=energy_df_rows, columns=times[:-1]) / dt_array + + final_energy = 0 + for p in packets: + final_energy += p.energy_rf + + print("Final energy to test for conservation") + print(final_energy) + + escape_energy = pd.DataFrame( + data=energy_out, columns=times[:-1], index=energy_bins + ) + + return ( + energy_df, + energy_plot_df, + escape_energy, + decayed_packet_count, + energy_plot_positrons, + energy_estimated_deposition, + ) + + +@njit(**njit_dict_no_parallel) +def process_packet_path(packet): + """Move the packet through interactions + + Parameters + ---------- + packet : GXPacket + Packet for processing + + Returns + ------- + GXPacket + Packet after processing + float + Energy injected into the ejecta + """ + if packet.status == GXPacketStatus.COMPTON_SCATTER: + comoving_freq_energy = packet.nu_cmf * H_CGS_KEV + + compton_angle, compton_fraction = get_compton_fraction_artis( + comoving_freq_energy + ) + + # Packet is no longer a gamma-ray, destroy it + if np.random.random() < 1 / compton_fraction: + packet.nu_cmf = packet.nu_cmf / compton_fraction + + packet.direction = compton_scatter(packet, compton_angle) + + # Calculate rest frame frequency after scaling by the fraction that remains + doppler_factor = doppler_gamma( + packet.direction, + packet.location, + packet.time_current, + ) + + packet.nu_rf = packet.nu_cmf / doppler_factor + packet.energy_rf = packet.energy_cmf / doppler_factor + + ejecta_energy_gained = 0.0 + else: + packet.status = GXPacketStatus.PHOTOABSORPTION + + if packet.status == GXPacketStatus.PAIR_CREATION: + packet = pair_creation_packet(packet) + ejecta_energy_gained = 0.0 + + if packet.status == GXPacketStatus.PHOTOABSORPTION: + # Ejecta gains comoving energy + ejecta_energy_gained = packet.energy_cmf + + return packet, ejecta_energy_gained + + +@njit(**njit_dict_no_parallel) +def gamma_packet_loop( + packets, + grey_opacity, + photoabsorption_opacity_type, + pair_creation_opacity_type, + electron_number_density_time, + mass_density_time, + inv_volume_time, + iron_group_fraction_per_shell, + inner_velocities, + outer_velocities, + times, + dt_array, + effective_time_array, + energy_bins, + energy_df_rows, + energy_plot_df_rows, + energy_out, +): + """Propagates packets through the simulation + + Parameters + ---------- + packets : list + List of GXPacket objects + grey_opacity : float + Grey opacity value in cm^2/g + electron_number_density_time : array float64 + Electron number densities with time + mass_density_time : array float64 + Mass densities with time + inv_volume_time : array float64 + Inverse volumes with time + iron_group_fraction_per_shell : array float64 + Iron group fraction per shell + inner_velocities : array float64 + Inner velocities of the shells + outer_velocities : array float64 + Inner velocities of the shells + times : array float64 + Simulation time steps + dt_array : array float64 + Simulation delta-time steps + effective_time_array : array float64 + Simulation middle time steps + energy_bins : array float64 + Bins for escaping gamma-rays + energy_df_rows : array float64 + Energy output + energy_plot_df_rows : array float64 + Energy output for plotting + energy_out : array float64 + Escaped energy array + + Returns + ------- + array float64 + Energy output + array float64 + Energy output for plotting + array float64 + Escaped energy array + + Raises + ------ + ValueError + Packet time index less than zero + """ + escaped_packets = 0 + scattered_packets = 0 + packet_count = len(packets) + print("Entering gamma ray loop for " + str(packet_count) + " packets") + + deposition_estimator = np.zeros_like(energy_df_rows) + + for i in range(packet_count): + packet = packets[i] + time_index = get_index(packet.time_current, times) + + if time_index < 0: + print(packet.time_current, time_index) + raise ValueError("Packet time index less than 0!") + + scattered = False + + initial_energy = packet.energy_cmf + + while packet.status == GXPacketStatus.IN_PROCESS: + # Get delta-time value for this step + dt = dt_array[time_index] + # Calculate packet comoving energy for opacities + comoving_energy = H_CGS_KEV * packet.nu_cmf + + if grey_opacity < 0: + doppler_factor = doppler_gamma( + packet.direction, + packet.location, + effective_time_array[time_index], + ) + + kappa = kappa_calculation(comoving_energy) + + # artis threshold for Thomson scattering + if kappa < 1e-2: + compton_opacity = ( + SIGMA_T + * electron_number_density_time[packet.shell, time_index] + ) + else: + compton_opacity = compton_opacity_calculation( + comoving_energy, + electron_number_density_time[packet.shell, time_index], + ) + + if photoabsorption_opacity_type == "kasen": + # currently not functional, requires proton count and + # electron count per isotope + photoabsorption_opacity = 0 + # photoabsorption_opacity_calculation_kasen() + else: + photoabsorption_opacity = ( + photoabsorption_opacity_calculation( + comoving_energy, + mass_density_time[packet.shell, time_index], + iron_group_fraction_per_shell[packet.shell], + ) + ) + + if pair_creation_opacity_type == "artis": + pair_creation_opacity = pair_creation_opacity_artis( + comoving_energy, + mass_density_time[packet.shell, time_index], + iron_group_fraction_per_shell[packet.shell], + ) + else: + pair_creation_opacity = pair_creation_opacity_calculation( + comoving_energy, + mass_density_time[packet.shell, time_index], + iron_group_fraction_per_shell[packet.shell], + ) + + else: + compton_opacity = 0.0 + pair_creation_opacity = 0.0 + photoabsorption_opacity = ( + grey_opacity * mass_density_time[packet.shell, time_index] + ) + + # convert opacities to rest frame + total_opacity = ( + compton_opacity + + photoabsorption_opacity + + pair_creation_opacity + ) * doppler_factor + + packet.tau = -np.log(np.random.random()) + + ( + distance_interaction, + distance_boundary, + distance_time, + shell_change, + ) = distance_trace( + packet, + inner_velocities, + outer_velocities, + total_opacity, + effective_time_array[time_index], + times[time_index + 1], + ) + + distance = min( + distance_interaction, distance_boundary, distance_time + ) + + packet.time_current += distance / C_CGS + + packet = move_packet(packet, distance) + + deposition_estimator[packet.shell, time_index] += ( + (initial_energy * 1000) + * distance + * (packet.energy_cmf / initial_energy) + * deposition_estimator_kasen( + comoving_energy, + mass_density_time[packet.shell, time_index], + iron_group_fraction_per_shell[packet.shell], + ) + ) + + if distance == distance_time: + time_index += 1 + + if time_index > len(effective_time_array) - 1: + # Packet ran out of time + packet.status = GXPacketStatus.END + else: + packet.shell = get_index( + packet.get_location_r(), + inner_velocities * effective_time_array[time_index], + ) + + elif distance == distance_interaction: + + packet.status = scatter_type( + compton_opacity, + photoabsorption_opacity, + total_opacity, + ) + + packet, ejecta_energy_gained = process_packet_path(packet) + + # Save packets to dataframe rows + # convert KeV to eV / s / cm^3 + energy_df_rows[packet.shell, time_index] += ( + ejecta_energy_gained * 1000 + ) + + energy_plot_df_rows[i] = np.array( + [ + i, + ejecta_energy_gained * 1000 + # * inv_volume_time[packet.shell, time_index] + / dt, + packet.get_location_r(), + packet.time_current, + packet.shell, + compton_opacity, + photoabsorption_opacity, + pair_creation_opacity, + ] + ) + + if packet.status == GXPacketStatus.PHOTOABSORPTION: + # Packet destroyed, go to the next packet + break + else: + packet.status = GXPacketStatus.IN_PROCESS + scattered = True + + else: + packet.shell += shell_change + + if packet.shell > len(mass_density_time[:, 0]) - 1: + rest_energy = packet.nu_rf * H_CGS_KEV + bin_index = get_index(rest_energy, energy_bins) + bin_width = ( + energy_bins[bin_index + 1] - energy_bins[bin_index] + ) + energy_out[bin_index, time_index] += rest_energy / ( + bin_width * dt + ) + packet.status = GXPacketStatus.END + escaped_packets += 1 + if scattered: + scattered_packets += 1 + elif packet.shell < 0: + packet.energy_rf = 0.0 + packet.energy_cmf = 0.0 + packet.status = GXPacketStatus.END + + print("Escaped packets:", escaped_packets) + print("Scattered packets:", scattered_packets) + + return energy_df_rows, energy_plot_df_rows, energy_out, deposition_estimator diff --git a/tardis/energy_input/nuclear_energy_source.py b/tardis/energy_input/nuclear_energy_source.py new file mode 100644 index 00000000000..ca4c3f83cb6 --- /dev/null +++ b/tardis/energy_input/nuclear_energy_source.py @@ -0,0 +1,53 @@ +from nuclear.ejecta import Ejecta +from nuclear.io.nndc import get_decay_radiation_database, store_decay_radiation + + +def decay_nuclides(shell_mass, initial_composition, epoch): + """Decay model + + Parameters + ---------- + shell_mass : float + Mass of the shell in grams + initial_composition : DataFrame + Initial ejecta composition + epoch : float + Time in days + + Returns + ------- + DataFrame + New composition at time epoch + """ + decay_model = Ejecta(shell_mass, initial_composition) + + new_fractions = decay_model.decay(epoch) + return new_fractions + + +def get_decay_database( + isotope_abundance, +): + """Gets the decay radiation database for a set + of isotopes + + Parameters + ---------- + isotope_abundance : DataFrame + DataFrame of simulation isotope masses per shell + + Returns + ------- + DataFrame + Decay radiation database + DataFrame + Metadata for the decay radiation database + """ + for column in isotope_abundance: + if column == "Fe56": + continue + store_decay_radiation(column, force_update=False) + + decay_rad_db, meta = get_decay_radiation_database() + + return decay_rad_db, meta diff --git a/tardis/energy_input/tests/conftest.py b/tardis/energy_input/tests/conftest.py new file mode 100644 index 00000000000..b36623ef660 --- /dev/null +++ b/tardis/energy_input/tests/conftest.py @@ -0,0 +1,26 @@ +import pytest +import numpy as np + +from tardis.energy_input.GXPacket import GXPacket, GXPacketStatus +from tardis.energy_input.util import H_CGS_KEV + + +@pytest.fixture(scope="function") +def basic_gamma_ray(): + """basic gamma ray fixture + + Returns + ------- + GXPacket + """ + return GXPacket( + location=np.array([1.36375693e13, 4.10589818e14, 9.11718168e14]), + direction=np.array([-0.97113853, 0.23134328, -0.05805379]), + energy_rf=1e52, + energy_cmf=1e52, + nu_rf=1000.0e3 / H_CGS_KEV, + nu_cmf=1000.0e3 / H_CGS_KEV, + status=GXPacketStatus.IN_PROCESS, + shell=1, + time_current=1000, + ) diff --git a/tardis/energy_input/tests/test_calculate_opacity.py b/tardis/energy_input/tests/test_calculate_opacity.py new file mode 100644 index 00000000000..41da4301b9f --- /dev/null +++ b/tardis/energy_input/tests/test_calculate_opacity.py @@ -0,0 +1,78 @@ +import pytest +import numpy.testing as npt + +import tardis.energy_input.calculate_opacity as calculate_opacity + + +@pytest.mark.parametrize( + ["electron_number_density", "energy", "expected"], + [ + (1.0e11, 511.0, 2.865396624016367e-14), + (1e15, 255.5, 3.743906253489761e-10), + (1e5, 511.0e7, 4.318577913631238e-26), + ], +) +def test_compton_opacity_calculation(energy, electron_number_density, expected): + """ + Parameters + ---------- + energy : float + electron_number_density : float + """ + opacity = calculate_opacity.compton_opacity_calculation( + energy, electron_number_density + ) + + npt.assert_almost_equal(opacity, expected) + + +@pytest.mark.parametrize( + ["ejecta_density", "energy", "iron_group_fraction", "expected"], + [ + (1.0, 511.0, 0.0, 0.00015028056615643418), + (1e-2, 255.5, 0.5, 8.903267700390038e-05), + (1e-2, 255.5, 0.25, 5.1069068712110425e-05), + (1e5, 511.0e7, 1.0, 0.0), + ], +) +def test_photoabsorption_opacity_calculation( + energy, ejecta_density, iron_group_fraction, expected +): + """ + Parameters + ---------- + energy : float + ejecta_density : float + iron_group_fraction : float + """ + opacity = calculate_opacity.photoabsorption_opacity_calculation( + energy, ejecta_density, iron_group_fraction + ) + + npt.assert_almost_equal(opacity, expected) + + +@pytest.mark.parametrize( + ["ejecta_density", "energy", "iron_group_fraction", "expected"], + [ + (1.0, 511.0, 0.0, 0.0), + (1e-2, 1500, 0.5, 2.743980356831218e-06), + (1e-2, 1200, 0.25, 8.846018943383742e-06), + (1e5, 511.0e7, 1.0, 1111355719.7411418), + ], +) +def test_pair_creation_opacity_calculation( + energy, ejecta_density, iron_group_fraction, expected +): + """ + Parameters + ---------- + energy : float + ejecta_density : float + iron_group_fraction : float + """ + opacity = calculate_opacity.pair_creation_opacity_calculation( + energy, ejecta_density, iron_group_fraction + ) + + npt.assert_almost_equal(opacity, expected) diff --git a/tardis/energy_input/tests/test_energy_source.py b/tardis/energy_input/tests/test_energy_source.py new file mode 100644 index 00000000000..42b42ceefdd --- /dev/null +++ b/tardis/energy_input/tests/test_energy_source.py @@ -0,0 +1,47 @@ +import pytest +import numpy as np +import numpy.testing as npt + +from tardis.energy_input.energy_source import ( + create_energy_cdf, +) + + +@pytest.mark.parametrize( + ["energy", "intensity", "expected_cdf"], + [ + (np.array([100.0, 50.0]), np.array([1.0, 1.0]), np.array([0.5, 1.0])), + (np.array([50.0, 100.0]), np.array([0.0, 1.0]), np.array([0.0, 1.0])), + ], +) +def test_create_energy_cdf(energy, intensity, expected_cdf): + """ + Parameters + ---------- + energy : One-dimensional Numpy Array, dtype float + intensity : One-dimensional Numpy Array, dtype float + expected_cdf : One-dimensional Numpy Array, dtype float + """ + actual_energy, actual_cdf = create_energy_cdf(energy, intensity) + expected_energy = np.sort(energy) + + npt.assert_array_almost_equal_nulp(actual_cdf, expected_cdf) + npt.assert_array_almost_equal_nulp(actual_energy, expected_energy) + + +@pytest.mark.xfail(reason="To be implemented") +def test_sample_energy_distribution(): + """To test the energy distribution sample""" + assert False + + +@pytest.mark.xfail(reason="To be implemented") +def test_setup_input_energy(): + """To test setting up the input energy""" + assert False + + +@pytest.mark.xfail(reason="To be implemented") +def test_intensity_ratio(): + """To test the intensity ratio""" + assert False diff --git a/tardis/energy_input/tests/test_gamma_ray_grid.py b/tardis/energy_input/tests/test_gamma_ray_grid.py new file mode 100644 index 00000000000..11a6f394ae5 --- /dev/null +++ b/tardis/energy_input/tests/test_gamma_ray_grid.py @@ -0,0 +1,38 @@ +import pytest +import numpy.testing as npt +import numpy as np + +from tardis.energy_input.gamma_ray_grid import ( + calculate_distance_radial, + distance_trace, + move_packet, +) + + +@pytest.mark.xfail(reason="To be implemented") +def test_calculate_distance_radial(): + """Test the radial distance calculation""" + assert False + + +@pytest.mark.xfail(reason="To be implemented") +def test_distance_trace(): + """Test the distance trace""" + assert False + + +def test_move_packet(basic_gamma_ray): + """ + + Parameters + ---------- + basic_gamma_ray : GammaRay object + """ + packet = basic_gamma_ray + distance = 1.0e15 + + new = packet.location + distance * packet.direction + + actual = move_packet(packet, distance) + + npt.assert_almost_equal(actual.location, new) diff --git a/tardis/energy_input/tests/test_gamma_ray_interactions.py b/tardis/energy_input/tests/test_gamma_ray_interactions.py new file mode 100644 index 00000000000..fedc16d9619 --- /dev/null +++ b/tardis/energy_input/tests/test_gamma_ray_interactions.py @@ -0,0 +1,72 @@ +import pytest +import numpy as np +import numpy.testing as npt + +from tardis.energy_input.gamma_ray_interactions import ( + compton_scatter, + pair_creation_packet, + scatter_type, +) +from tardis.energy_input.util import ELECTRON_MASS_ENERGY_KEV, H_CGS_KEV + +from tardis.energy_input.GXPacket import GXPacketStatus + + +@pytest.mark.xfail(reason="To be implemented") +def test_get_compton_angle(): + """Test the Compton angle calculation""" + assert False + + +@pytest.mark.xfail(reason="To be implemented") +def test_compton_scatter(): + """Test Compton scattering""" + assert False + + +def test_pair_creation(basic_gamma_ray): + """ + + Parameters + ---------- + basic_gamma_ray : GammaRay object + """ + np.random.seed(2) + + initial_direction = basic_gamma_ray.direction + basic_gamma_ray.nu_cmf = 2 * ELECTRON_MASS_ENERGY_KEV / H_CGS_KEV + + pair_creation_packet(basic_gamma_ray) + + npt.assert_almost_equal( + basic_gamma_ray.nu_cmf, ELECTRON_MASS_ENERGY_KEV / H_CGS_KEV + ) + for i, j in zip(initial_direction, basic_gamma_ray.direction): + assert i != j + + +@pytest.mark.parametrize( + ["compton_opacity", "photoabsorption_opacity", "total_opacity", "expected"], + [ + (1, 0, 1, GXPacketStatus.COMPTON_SCATTER), + (0, 1, 1, GXPacketStatus.PHOTOABSORPTION), + (0, 0, 1, GXPacketStatus.PAIR_CREATION), + ], +) +def test_scatter_type( + compton_opacity, photoabsorption_opacity, total_opacity, expected +): + """Test the scattering type + + Parameters + ---------- + compton_opacity : float + photoabsorption_opacity : float + total_opacity : float + expected : list + Expected parameters + """ + actual = scatter_type( + compton_opacity, photoabsorption_opacity, total_opacity + ) + assert actual == expected diff --git a/tardis/energy_input/tests/test_util.py b/tardis/energy_input/tests/test_util.py new file mode 100644 index 00000000000..6ea70d8f022 --- /dev/null +++ b/tardis/energy_input/tests/test_util.py @@ -0,0 +1,124 @@ +from random import random +import pytest +import astropy.units as u +import numpy.testing as npt +import numpy as np + +import tardis.energy_input.util as util +from tardis.energy_input.util import ( + R_ELECTRON_SQUARED, + get_perpendicular_vector, +) +from tardis import constants as const + + +@pytest.mark.parametrize( + ["r", "theta", "phi", "expected_x", "expected_y", "expected_z"], + [ + (1, 0, 0, 0, 0, 1), + (1, np.pi, 0, 0, 0, -1), + (1, np.pi / 2, 0, 1, 0, 0), + (1, np.pi / 2, np.pi, -1, 0, 0), + (1, np.pi / 2, np.pi / 2, 0, 1, 0), + ], +) +def test_spherical_to_cartesian( + r, theta, phi, expected_x, expected_y, expected_z +): + actual_x, actual_y, actual_z = util.spherical_to_cartesian(r, theta, phi) + npt.assert_almost_equal(actual_x, expected_x) + npt.assert_almost_equal(actual_y, expected_y) + npt.assert_almost_equal(actual_z, expected_z) + + +@pytest.mark.xfail(reason="To be removed") +def test_doppler_gamma(): + """Test the doppler shift""" + assert False + + +@pytest.mark.xfail(reason="To be removed") +def test_angle_aberration_gamma(): + """Test the angle aberration equation""" + assert False + + +@pytest.mark.parametrize( + ["energy", "expected"], + [ + (511.0, 1.0000021334560507), + (255.5, 0.5000010667280254), + (0.0, 0.0), + (511.0e7, 10000021.334560508), + ], +) +def test_kappa_calculation(energy, expected): + """ + + Parameters + ---------- + energy : float + expected : float + """ + kappa = util.kappa_calculation(energy) + npt.assert_almost_equal(kappa, expected) + + +@pytest.mark.xfail(reason="To be removed") +def test_euler_rodrigues(): + """Test Euler-Rodrigues rotation""" + assert False + + +@pytest.mark.xfail(reason="To be implemented") +def test_solve_quadratic_equation(): + """Test the quadratic solver""" + assert False + + +@pytest.mark.parametrize( + ["energy", "theta_C"], + [ + (511.0e3, 1.0), + (255.5e3, np.pi), + (0.0, 2.0 * np.pi), + (511.0e10, np.pi / 2.0), + ], +) +def test_klein_nishina(energy, theta_C): + """ + + Parameters + ---------- + energy : float + theta_C : float + In radians + """ + actual = util.klein_nishina(energy, theta_C) + + kappa = util.kappa_calculation(energy) + + expected = ( + R_ELECTRON_SQUARED + / 2 + * (1.0 + kappa * (1.0 - np.cos(theta_C))) ** -2.0 + * ( + 1.0 + + np.cos(theta_C) ** 2.0 + + (kappa**2.0 * (1.0 - np.cos(theta_C)) ** 2.0) + / (1.0 + kappa * (1.0 - np.cos(theta_C))) + ) + ) + + npt.assert_almost_equal(actual, expected) + + +def test_get_perpendicular_vector(): + """Test the perpendicular vector calculation""" + input_vector = np.array([0.3, 0.4, 0.5]) + + random_perpendicular_vector = get_perpendicular_vector(input_vector) + + dot_product = np.dot(input_vector, random_perpendicular_vector) + + npt.assert_almost_equal(dot_product, 0.0) diff --git a/tardis/energy_input/util.py b/tardis/energy_input/util.py new file mode 100644 index 00000000000..9c1633046b1 --- /dev/null +++ b/tardis/energy_input/util.py @@ -0,0 +1,405 @@ +import astropy.units as u +import tardis.constants as const +import numpy as np +from numba import njit + +from tardis.montecarlo.montecarlo_numba import njit_dict_no_parallel + +R_ELECTRON_SQUARED = (const.a0.cgs.value * const.alpha.cgs.value**2.0) ** 2.0 +ELECTRON_MASS_ENERGY_KEV = (const.m_e * const.c**2.0).to("keV").value +BOUNDARY_THRESHOLD = 1e-7 +KEV2ERG = (1000 * u.eV).to("erg").value +C_CGS = const.c.cgs.value +H_CGS_KEV = const.h.to("keV s").value + + +@njit(**njit_dict_no_parallel) +def spherical_to_cartesian(r, theta, phi): + """Converts spherical coordinates to Cartesian + + Parameters + ---------- + r : float64 + Radius + theta : float64 + Theta angle in radians + phi : float64 + Phi angle in radians + + Returns + ------- + array + x, y, z coordinates + """ + x = r * np.cos(phi) * np.sin(theta) + y = r * np.sin(phi) * np.sin(theta) + z = r * np.cos(theta) + return np.array([x, y, z]) + + +@njit(**njit_dict_no_parallel) +def get_random_unit_vector(): + """Generate a random unit vector + + Returns: + array: random unit vector + """ + theta = get_random_theta_photon() + phi = get_random_phi_photon() + + vector = spherical_to_cartesian(1, theta, phi) + + return normalize_vector(vector) + + +@njit(**njit_dict_no_parallel) +def doppler_gamma(direction_vector, position_vector, time): + """Doppler shift for photons in 3D + + Parameters + ---------- + direction_vector : array + position_vector : array + time : float + + Returns + ------- + float + Doppler factor + """ + velocity_vector = position_vector / time + return 1 - (np.dot(direction_vector, velocity_vector) / C_CGS) + + +@njit(**njit_dict_no_parallel) +def angle_aberration_gamma(direction_vector, position_vector, time): + """Angle aberration formula for photons in 3D + + Parameters + ---------- + direction_vector : array + position_vector : array + time : float + + Returns + ------- + array + New direction after aberration + """ + velocity_vector = position_vector / time + direction_dot_velocity = np.dot(direction_vector, velocity_vector) + + gamma = 1.0 / np.sqrt( + 1.0 - (np.dot(velocity_vector, velocity_vector) / (C_CGS**2.0)) + ) + + factor_a = gamma * (1.0 - direction_dot_velocity / C_CGS) + + factor_b = ( + gamma - (gamma**2.0 * direction_dot_velocity / (gamma + 1.0) / C_CGS) + ) / C_CGS + + output_vector = (direction_vector - (velocity_vector * factor_b)) / factor_a + + return output_vector + + +@njit(**njit_dict_no_parallel) +def kappa_calculation(energy): + """ + Calculates kappa for various other calculations + i.e. energy normalized to electron rest energy + 511.0 KeV + + Parameters + ---------- + energy : float + + Returns + ------- + kappa : float + + """ + return energy / ELECTRON_MASS_ENERGY_KEV + + +@njit(**njit_dict_no_parallel) +def euler_rodrigues(theta, direction): + """ + Calculates the Euler-Rodrigues rotation matrix + + Parameters + ---------- + theta : float + angle of rotation in radians + direction : One dimensional Numpy array, dtype float + x, y, z direction vector + + Returns + ------- + rotation matrix : Two dimensional Numpy array, dtype float + + """ + a = np.cos(theta / 2) + b = direction[0] * np.sin(theta / 2) + c = direction[1] * np.sin(theta / 2) + d = direction[2] * np.sin(theta / 2) + + er11 = a**2.0 + b**2.0 - c**2.0 - d**2.0 + er12 = 2.0 * (b * c - a * d) + er13 = 2.0 * (b * d + a * c) + + er21 = 2.0 * (b * c + a * d) + er22 = a**2.0 + c**2.0 - b**2.0 - d**2.0 + er23 = 2.0 * (c * d - a * b) + + er31 = 2.0 * (b * d - a * c) + er32 = 2.0 * (c * d + a * b) + er33 = a**2.0 + d**2.0 - b**2.0 - c**2.0 + + return np.array( + [[er11, er12, er13], [er21, er22, er23], [er31, er32, er33]] + ) + + +@njit(**njit_dict_no_parallel) +def solve_quadratic_equation(position, direction, radius): + """ + Solves the quadratic equation for the distance to the shell boundary + + Parameters + ---------- + position : array + direction : array + radius : float + + Returns + ------- + solution_1 : float + solution_2 : float + + """ + a = np.sum(direction**2) + b = 2.0 * np.sum(position * direction) + c = -(radius**2) + np.sum(position**2) + root = b**2 - 4 * a * c + solution_1 = -np.inf + solution_2 = -np.inf + if root > 0.0: + solution_1 = (-b + np.sqrt(root)) / (2 * a) + solution_2 = (-b - np.sqrt(root)) / (2 * a) + elif root == 0: + solution_1 = -b / (2 * a) + + return solution_1, solution_2 + + +@njit(**njit_dict_no_parallel) +def solve_quadratic_equation_expanding(position, direction, time, radius): + """ + Solves the quadratic equation for the distance to an expanding shell boundary + + Parameters + ---------- + position : array + direction : array + time : float + radius : float + + Returns + ------- + solution_1 : float + solution_2 : float + + """ + light_distance = time * C_CGS + a = np.dot(direction, direction) - (radius / light_distance) ** 2.0 + b = 2.0 * (np.dot(position, direction) - radius**2.0 / light_distance) + c = np.dot(position, position) - radius**2.0 + root = b**2.0 - 4.0 * a * c + solution_1 = -np.inf + solution_2 = -np.inf + if root > 0.0: + solution_1 = (-b + np.sqrt(root)) / (2.0 * a) + solution_2 = (-b - np.sqrt(root)) / (2.0 * a) + elif root == 0: + solution_1 = -b / (2.0 * a) + + return solution_1, solution_2 + + +@njit(**njit_dict_no_parallel) +def klein_nishina(energy, theta_C): + """ + Calculates the Klein-Nishina equation + + https://en.wikipedia.org/wiki/Klein%E2%80%93Nishina_formula + + .. math:: + \frac{r_e}{2} [1 + \kappa (1 - \cos\theta_C)]^{-2} \left( 1 + \cos^2\theta_C + \frac{\kappa^2 (1 - \cos\theta_C)^2}{1 + \kappa(1 - \cos\theta_C)}\right) + + where :math:`\kappa = E / (m_e c^2)` + + Parameters + ---------- + energy : float + Packet energy + theta_C : float + Compton angle + Returns + ------- + float + Klein-Nishina solution + """ + kappa = kappa_calculation(energy) + return ( + R_ELECTRON_SQUARED + / 2 + * (1.0 + kappa * (1.0 - np.cos(theta_C))) ** -2.0 + * ( + 1.0 + + np.cos(theta_C) ** 2.0 + + (kappa**2.0 * (1.0 - np.cos(theta_C)) ** 2.0) + / (1.0 + kappa * (1.0 - np.cos(theta_C))) + ) + ) + + +@njit(**njit_dict_no_parallel) +def compton_theta_distribution(energy, sample_resolution=100): + """ + Calculates the cumulative distribution function of theta angles + for Compton Scattering + + Parameters + ---------- + energy : float + sample_resolution : int + + Returns + ------- + theta_angles : One dimensional Numpy array, dtype float + norm_theta_distribution : One dimensional Numpy array, dtype float + + """ + theta_angles = np.linspace(0, np.pi, sample_resolution) + + theta_distribution = np.cumsum(klein_nishina(energy, theta_angles)) + norm_theta_distribution = theta_distribution / np.max(theta_distribution) + + return theta_angles, norm_theta_distribution + + +@njit(**njit_dict_no_parallel) +def get_random_theta_photon(): + """Get a random theta direction between 0 and pi + Returns + ------- + float + Random theta direction + """ + return np.arccos(1.0 - 2.0 * np.random.random()) + + +@njit(**njit_dict_no_parallel) +def get_random_phi_photon(): + """Get a random phi direction between 0 and 2 * pi + + Returns + ------- + float + Random phi direction + """ + return 2.0 * np.pi * np.random.random() + + +def convert_half_life_to_astropy_units(half_life_string): + """Converts input half-life to use astropy units + + Parameters + ---------- + half_life_string : str + Half-life as a string + + Returns + ------- + astropy Quantity + Half-life in seconds + """ + value, unit = half_life_string.split(" ") + try: + nominal_value, magnitude = value.split("×") + base, exponent = magnitude.split("+") + nominal_value = float(nominal_value) * float(base) ** float(exponent) + except ValueError: + nominal_value = float(value) + if unit == "y": + unit = "yr" + half_life_with_unit = nominal_value * u.Unit(unit) + return half_life_with_unit.to(u.s) + + +@njit(**njit_dict_no_parallel) +def normalize_vector(vector): + """ + Normalizes a vector in Cartesian coordinates + + Parameters + ---------- + vector : One-dimensional Numpy Array, dtype float + Input vector + + Returns + ------- + One-dimensional Numpy Array, dtype float + Normalized vector + """ + return vector / np.linalg.norm(vector) + + +@njit(**njit_dict_no_parallel) +def get_perpendicular_vector(original_direction): + """ + Computes a random vector which is perpendicular to the input vector + + Parameters + ---------- + original_direction : array + + Returns + ------- + numpy.ndarray + Perpendicular vector to the input + """ + random_vector = get_random_unit_vector() + perpendicular_vector = np.cross(original_direction, random_vector) + return normalize_vector(perpendicular_vector) + + +@njit(**njit_dict_no_parallel) +def get_index(value, array): + """Get the index that places a value + at array[i] < array <= vec[i+1] + + Parameters + ---------- + value : float + Value to locate + array : array + Array to search + + Returns + ------- + int + Index + """ + if value <= array[0]: + return 0 + elif value > array[-1]: + return len(array) - 1 + + i = 0 + while value > array[i + 1]: + i += 1 + + return i diff --git a/tardis/io/tests/data/tardis_configv1_density_exponential_nebular.yml b/tardis/io/tests/data/tardis_configv1_density_exponential_nebular.yml new file mode 100644 index 00000000000..64bacc62ffa --- /dev/null +++ b/tardis/io/tests/data/tardis_configv1_density_exponential_nebular.yml @@ -0,0 +1,50 @@ +tardis_config_version: v1.0 + +supernova: + luminosity_requested: 9.44 log_lsun + time_explosion: 150 day + +atom_data: kurucz_atom_pure.h5 + +model: + structure: + type: specific + + velocity: + start: 1.1e4 km/s + stop: 20000 km/s + num: 20 + + density: + type: exponential + time_0: 20. second + rho_0: 3.e2 g/cm^3 + v_0: 3000. km/s + + abundances: + type: uniform + O: 0.19 + Mg: 0.03 + Si: 0.12 + S: 0.09 + Ar: 0.04 + Ca: 0.03 + Ni56: 0.5 + +plasma: + ionization: nebular + excitation: dilute-lte + radiative_rates_type: dilute-blackbody + line_interaction_type: scatter + +montecarlo: + seed: 23111963 + no_of_packets: 2.0e+5 + iterations: 30 + last_no_of_packets: 5.0e+5 + no_of_virtual_packets: 5 + +spectrum: + start: 500 angstrom + stop: 20000 angstrom + num: 10000