Skip to content

Commit

Permalink
New packet source (extending 2352) (#2366)
Browse files Browse the repository at this point in the history 
* Restructure packet_source

* Add skeleton for EnergyDepositionSource

* Copy direction sampling from BlackBodySimpleSource

* Restructure codebase according to new packet source structure

* Sample packet radii from gamma ray energies

* Create placeholders for frequency and energies

* Use linear normalization instead of softmax to get shell probabilities

* remove gammaray stuff

---------

Co-authored-by: xansh <[email protected]>
  • Loading branch information
@wkerzendorf @xansh
wkerzendorf and xansh authored Jul 21, 2023
1 parent 754351f commit 9831b9e
Show file tree
Hide file tree
Showing 5 changed files with 316 additions and 170 deletions.
120 changes: 70 additions & 50 deletions docs/io/optional/custom_source.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,25 +5,30 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"# Running TARDIS with a Custom Packet Source"
"# Running TARDIS with a Custom Packet Source\n"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"By default, TARDIS generates energy packets using its `BasePacketSource` class, which models the photosphere of the supernova as a perfect blackbody (see [Energy Packet Initialization](../../physics/montecarlo/initialization.ipynb)). However, users may create their own packet source, as will be shown in this notebook. In order to do this, a user must create a class (that can but does not have to inherit from `BasePacketSource`) which contains a method `create_packets` that takes in (in this order):\n",
"- the photospheric temperature\n",
"- the number of packets\n",
"- a random number generator\n",
"- the photospheric radius\n",
"By default, TARDIS generates energy packets using its interface `BasePacketSource` class, which has a derived class `BlackBodySimpleSource`, which models the photosphere of the supernova as a perfect blackbody (see [Energy Packet Initialization](../../physics/montecarlo/initialization.ipynb)). However, users may create their own packet source, as will be shown in this notebook. In order to do this, a user must create a class (that inherits from `BasePacketSource`) and implement the following abstract functions:\n",
"\n",
"and returns arrays of the radii, frequencies, propogation directions, and energies of the ensemble of packets (once again see [Energy Packet Initialization](../../physics/montecarlo/initialization.ipynb) for more information). To use your packet source in a run of TARDIS, you must pass an instance of your class into the `run_tardis` function under the `packet_source` keyword argument.\n",
"- create_packet_radii (returns array of packet radii)\n",
"- create_packet_nus (returns array of packet frequencies)\n",
"- create_packet_mus (returns array of packet directions)\n",
"- create_packet_energies (returns array of packet energies. See [Energy Packet Initialization](../../physics/montecarlo/initialization.ipynb) for more information)\n",
"- create_packets (wrapper which calls the above 4 functions, and is the function used by external code)\n",
"- set_state_from_model (set the state of the source from a model object)\n",
"\n",
".. note:: In both the `BasePacketSource` class and in the example here, all packets are generated at the same radius. This need not be true in general (although the `create_packets` method should only take in a single radius as its argument).\n",
"[Note: In this notebook, we have extended the `BlackBodySimpleSource` class because it already implements some of the above functions]\n",
"\n",
"We show an example of how a custom packet source is used:"
"To use your packet source in a run of TARDIS, you must pass an instance of your class into the `run_tardis` function under the `packet_source` keyword argument.\n",
"\n",
".. note:: In both the `BlackBodySimpleSource` class and in the example here, all packets are generated at the same radius. This need not be true in general (though one call of the `create_packets` method will pick the same radius from the packet source state).\n",
"\n",
"We show an example of how a custom packet source is used:\n"
]
},
{
Expand All @@ -36,7 +41,7 @@
"import numpy as np\n",
"from tardis import constants as const\n",
"from astropy import units as u\n",
"from tardis.montecarlo.packet_source import BasePacketSource\n",
"from tardis.montecarlo.packet_source import BlackBodySimpleSource\n",
"from tardis import run_tardis\n",
"import matplotlib.pyplot as plt\n",
"from tardis.io.atom_data import download_atom_data"
Expand All @@ -49,7 +54,7 @@
"outputs": [],
"source": [
"# Download the atomic data used for a run of TARDIS\n",
"download_atom_data('kurucz_cd23_chianti_H_He')"
"download_atom_data(\"kurucz_cd23_chianti_H_He\")"
]
},
{
Expand All @@ -59,65 +64,81 @@
"outputs": [],
"source": [
"# Create a packet source class that contains a create_packets method\n",
"class TruncBlackbodySource(BasePacketSource):\n",
"class TruncBlackbodySource(BlackBodySimpleSource):\n",
" \"\"\"\n",
" Custom inner boundary source class to replace the Blackbody source\n",
" with a truncated Blackbody source.\n",
" Custom inner boundary source class to replace the Blackbody source\n",
" with a truncated Blackbody source.\n",
"\n",
" Parameters\n",
" ----------\n",
" truncation_wavelength : float\n",
" truncation wavelength in Angstrom.\n",
" Only wavelengths higher than the truncation wavelength\n",
" will be sampled.\n",
" radius : float64\n",
" Initial packet radius\n",
" temperature : float\n",
" Absolute Temperature.\n",
" base_seed : int\n",
" Base Seed for random number generator\n",
" \"\"\"\n",
" \n",
" def __init__(self, base_seed, truncation_wavelength):\n",
" super().__init__(base_seed)\n",
" self.rng = np.random.default_rng(seed=base_seed)\n",
"\n",
" def __init__(self, truncation_wavelength=None, **kwargs):\n",
" self.truncation_wavelength = truncation_wavelength\n",
" \n",
" def create_packets(self, T, no_of_packets, radius,\n",
" drawing_sample_size=None):\n",
" super().__init__(**kwargs)\n",
"\n",
" def create_packets(self, no_of_packets, drawing_sample_size=None):\n",
" \"\"\"\n",
" Packet source that generates a truncated Blackbody source.\n",
" \n",
"\n",
" Parameters\n",
" ----------\n",
" T : float\n",
" Blackbody temperature\n",
" no_of_packets : int\n",
" number of packets to be created\n",
" truncation_wavelength : float\n",
" truncation wavelength in Angstrom. \n",
" Only wavelengths higher than the truncation wavelength\n",
" will be sampled.\n",
"\n",
" Returns\n",
" -------\n",
" array\n",
" Packet radii\n",
" array\n",
" Packet frequencies\n",
" array\n",
" Packet directions\n",
" array\n",
" Packet energies\n",
" \"\"\"\n",
" \n",
" # Makes uniform array of packet radii\n",
" radii = np.ones(no_of_packets) * radius\n",
"\n",
" # Makes uniform array of packet radii from blackbody source\n",
" radii = self.create_packet_radii(no_of_packets)\n",
"\n",
" # Use mus and energies from normal blackbody source.\n",
" mus = self.create_zero_limb_darkening_packet_mus(no_of_packets)\n",
" energies = self.create_uniform_packet_energies(no_of_packets)\n",
" mus = self.create_packet_mus(no_of_packets)\n",
" energies = self.create_packet_energies(no_of_packets)\n",
"\n",
" # If not specified, draw 2 times as many packets and reject any beyond no_of_packets.\n",
" if drawing_sample_size is None:\n",
" drawing_sample_size = 2 * no_of_packets\n",
"\n",
" # Blackbody will be truncated below truncation_wavelength / above truncation_frequency.\n",
" truncation_frequency = u.Quantity(self.truncation_wavelength, u.Angstrom).to(\n",
" u.Hz, equivalencies=u.spectral()).value\n",
" \n",
" truncation_frequency = (\n",
" u.Quantity(self.truncation_wavelength, u.Angstrom)\n",
" .to(u.Hz, equivalencies=u.spectral())\n",
" .value\n",
" )\n",
"\n",
" # Draw nus from blackbody distribution and reject based on truncation_frequency.\n",
" # If more nus.shape[0] > no_of_packets use only the first no_of_packets.\n",
" nus = self.create_blackbody_packet_nus(T, drawing_sample_size)\n",
" nus = nus[nus<truncation_frequency][:no_of_packets]\n",
" \n",
" \n",
" # Only required if the truncation wavelength is too big compared to the maximum \n",
" nus = self.create_packet_nus(drawing_sample_size)\n",
" nus = nus[nus < truncation_frequency][:no_of_packets]\n",
"\n",
" # Only required if the truncation wavelength is too big compared to the maximum\n",
" # of the blackbody distribution. Keep sampling until nus.shape[0] > no_of_packets.\n",
" while nus.shape[0] < no_of_packets:\n",
" additional_nus = self.create_blackbody_packet_nus(\n",
" T, drawing_sample_size\n",
" )\n",
" additional_nus = self.create_packet_nus(drawing_sample_size)\n",
" mask = additional_nus < truncation_frequency\n",
" additional_nus = additional_nus[mask][:no_of_packets]\n",
" nus = np.hstack([nus, additional_nus])[:no_of_packets]\n",
" \n",
"\n",
" return radii, nus, mus, energies"
]
},
Expand All @@ -129,7 +150,7 @@
"source": [
"# Call an instance of the packet source class\n",
"packet_source = TruncBlackbodySource(\n",
" 53253, truncation_wavelength=2000\n",
" truncation_wavelength=2000, base_seed=53253\n",
")"
]
},
Expand All @@ -138,7 +159,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"We now run TARDIS both with and without our custom packet source, and we compare the results:"
"We now run TARDIS both with and without our custom packet source, and we compare the results:\n"
]
},
{
Expand All @@ -147,9 +168,8 @@
"metadata": {},
"outputs": [],
"source": [
"mdl = run_tardis('tardis_example.yml',\n",
" packet_source=packet_source)\n",
"mdl_norm = run_tardis('tardis_example.yml')"
"mdl = run_tardis(\"tardis_example.yml\", packet_source=packet_source)\n",
"mdl_norm = run_tardis(\"tardis_example.yml\")"
]
},
{
Expand Down
Loading

0 comments on commit 9831b9e

Please sign in to comment.