Custom source packets

docs/running/custom_source.ipynb

@@ -0,0 +1,173 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Running TARDIS with a custom packet source"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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 import run_tardis\n",
+    "import matplotlib.pyplot as plt\n",
+    "from tardis.io.atom_data import download_atom_data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "download_atom_data('kurucz_cd23_chianti_H_He')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Custom packet source class that is derived from BasePacketSource. The method create_packets (which returns ```nus, mus, energies```) has to be defined."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class TruncBlackbodySource(BasePacketSource):\n",
+    "    \"\"\"\n",
+    "        Custom inner boundary source class to replace the Blackbody source\n",
+    "        with a truncated Blackbody source.\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    def __init__(self, seed, truncation_wavelength):\n",
+    "        super().__init__(seed)\n",
+    "        self.truncation_wavelength = truncation_wavelength\n",
+    "    \n",
+    "    def create_packets(self, T, no_of_packets,\n",
+    "                       drawing_sample_size=None):\n",
+    "        \"\"\"\n",
+    "        Packet source that generates a truncated Blackbody source.\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",
+    "        \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",
+    "\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",
+    "        # 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",
+    "        # 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",
+    "            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",
+    "        return nus, mus, energies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "packet_source = TruncBlackbodySource(\n",
+    "    53253, truncation_wavelength=2000\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mdl = run_tardis('tardis_example.yml',\n",
+    "                 packet_source=packet_source)\n",
+    "mdl_norm = run_tardis('tardis_example.yml')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline\n",
+    "plt.plot(mdl.runner.spectrum_virtual.wavelength,\n",
+    "         mdl.runner.spectrum_virtual.luminosity_density_lambda,\n",
+    "         color='red', label='truncated blackbody')\n",
+    "plt.plot(mdl_norm.runner.spectrum_virtual.wavelength,\n",
+    "         mdl_norm.runner.spectrum_virtual.luminosity_density_lambda,\n",
+    "         color='blue', label='normal blackbody')\n",
+    "plt.xlabel('$\\lambda [\\AA]$')\n",
+    "plt.ylabel('$L_\\lambda$ [erg/s/$\\AA$]')\n",
+    "plt.xlim(500, 10000)\n",
+    "plt.legend()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "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.6.7"
+  },
+  "nbsphinx": {
+   "execute": "always",
+   "timeout": 600
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

docs/running/index.rst

@@ -24,6 +24,7 @@ A quick start guide to run TARDIS can be found
     commandline
     gui
     interaction/index
+    custom_source
 
 
 .. toctree::

tardis/base.py

@@ -1,6 +1,7 @@
 # functions that are important for the general usage of TARDIS
 
-def run_tardis(config, atom_data=None, simulation_callbacks=[]):
+def run_tardis(config, atom_data=None, packet_source=None,
+               simulation_callbacks=[]):
     """
     This function is one of the core functions to run TARDIS from a given
     config object.
@@ -34,7 +35,9 @@ def run_tardis(config, atom_data=None, simulation_callbacks=[]):
     except TypeError:
         tardis_config = Configuration.from_config_dict(config)
 
-    simulation = Simulation.from_config(tardis_config, atom_data=atom_data)
+    simulation = Simulation.from_config(tardis_config,
+                                        packet_source=packet_source,
+                                        atom_data=atom_data)
     for cb in simulation_callbacks:
         simulation.add_callback(cb)
 

tardis/montecarlo/base.py

@@ -13,7 +13,7 @@
 
 from tardis.util.base import quantity_linspace
 from tardis.io.util import HDFWriterMixin
-from tardis.montecarlo import montecarlo, packet_source
+from tardis.montecarlo import montecarlo, packet_source as source
 from tardis.montecarlo.formal_integral import FormalIntegrator
 
 import numpy as np
@@ -57,10 +57,15 @@ def __init__(self, seed, spectrum_frequency, virtual_spectrum_range,
                  sigma_thomson, enable_reflective_inner_boundary,
                  enable_full_relativity, inner_boundary_albedo,
                  line_interaction_type, integrator_settings,
-                 v_packet_settings, spectrum_method):
+                 v_packet_settings, spectrum_method,
+                 packet_source=None):
 
         self.seed = seed
-        self.packet_source = packet_source.BlackBodySimpleSource(seed)
+        if packet_source is None:
+            self.packet_source = source.BlackBodySimpleSource(seed)
+        else:
+            self.packet_source = packet_source
+        # inject different packets
         self.spectrum_frequency = spectrum_frequency
         self.virtual_spectrum_range = virtual_spectrum_range
         self.sigma_thomson = sigma_thomson
@@ -406,7 +411,7 @@ def calculate_f_lambda(self, wavelength):
         pass
 
     @classmethod
-    def from_config(cls, config):
+    def from_config(cls, config, packet_source=None):
         """
         Create a new MontecarloRunner instance from a Configuration object.
 
@@ -441,4 +446,5 @@ def from_config(cls, config):
                    line_interaction_type=config.plasma.line_interaction_type,
                    integrator_settings=config.spectrum.integrated,
                    v_packet_settings=config.spectrum.virtual,
-                   spectrum_method=config.spectrum.method)
+                   spectrum_method=config.spectrum.method,
+                   packet_source=packet_source)

tardis/montecarlo/formal_integral.py

@@ -13,7 +13,6 @@
 class IntegrationError(Exception):
     pass
 
-
 class FormalIntegrator(object):
 
     def __init__(self, model, plasma, runner, points=1000):