adding units

docs/physics/montecarlo/initialization.ipynb

@@ -60,6 +60,7 @@
    "source": [
     "import numpy as np\n",
     "from tardis.montecarlo.packet_source import BlackBodySimpleSource\n",
+    "from astropy import units as u\n",
     "from tardis import constants as const\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
@@ -86,7 +87,7 @@
     "seed = 1\n",
     "\n",
     "# Radius of the supernova's photosphere in cm\n",
-    "r_phot = 1e15\n",
+    "r_phot = 1e15 * u.cm\n",
     "\n",
     "# Number of packets generated\n",
     "n_packets = 40000"
@@ -108,10 +109,14 @@
    "outputs": [],
    "source": [
     "# Temperature in K\n",
-    "temperature = 5000\n",
+    "temperature = 10000 * u.K\n",
     "\n",
-    "luminosity = 4 * np.pi * (r_phot**2) * const.sigma_sb.cgs.value * (temperature**4)\n",
-    "print('Luminosity:', luminosity, 'erg/s')"
+    "luminosity = 4 * np.pi * (r_phot**2) * const.sigma_sb * (temperature**4)\n",
+    "\n",
+    "# Makes sure the luminosity is given in erg/s\n",
+    "luminosity = luminosity.to('erg/s')\n",
+    "\n",
+    "print('Luminosity:', luminosity)"
    ]
   },
   {
@@ -121,11 +126,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Luminocity in erg/s\n",
-    "luminosity = 1e40\n",
+    "# Luminosity in erg/s\n",
+    "luminosity = 7e42 * u.erg / u.s\n",
+    "\n",
+    "temperature = (luminosity / (4 * np.pi * const.sigma_sb))**0.25 / np.sqrt(r_phot)\n",
     "\n",
-    "temperature = (luminosity / (4 * np.pi * const.sigma_sb.cgs.value))**0.25 / np.sqrt(r_phot)\n",
-    "print('Temperature:', temperature, 'K')"
+    "# Makes sure the termperature is given in K\n",
+    "temperature = temperature.to('K')\n",
+    "\n",
+    "print('Temperature:', temperature)"
    ]
   },
   {
@@ -146,7 +155,17 @@
     "# We define our packet source\n",
     "packet_source = BlackBodySimpleSource(seed)\n",
     "\n",
-    "radii, nus, mus, energies = packet_source.create_packets(temperature, n_packets, rng, r_phot)\n",
+    "radii, nus, mus, energies = packet_source.create_packets(\n",
+    "    temperature.value, \n",
+    "    n_packets, \n",
+    "    rng, \n",
+    "    r_phot)\n",
+    "\n",
+    "# Sets the energies in units of ergs\n",
+    "energies *= u.erg\n",
+    "\n",
+    "# Sets the frequencies in units of Hz\n",
+    "nus *= u.Hz\n",
     "\n",
     "print('Energies:', energies)\n",
     "print('Radii:', radii)"
@@ -168,12 +187,12 @@
    "outputs": [],
    "source": [
     "# Time of simulation\n",
-    "t_simulation = 1 / luminosity\n",
-    "print('Time of simulation:', t_simulation, 'seconds')\n",
+    "t_simulation = 1 * u.erg / luminosity\n",
+    "print('Time of simulation:', t_simulation)\n",
     "\n",
     "# Array of luminosity contribution by each packet\n",
     "lumin_per_packet = energies / t_simulation\n",
-    "print('Luminosity per packet:', lumin_per_packet, 'erg/s')"
+    "print('Luminosity per packet:', lumin_per_packet)"
    ]
   },
   {
@@ -195,9 +214,9 @@
    },
    "outputs": [],
    "source": [
-    "h = const.h.cgs.value\n",
-    "c2 = const.c.cgs.value**2\n",
-    "kB = const.k_B.cgs.value\n",
+    "h = const.h.cgs\n",
+    "c2 = const.c.cgs**2\n",
+    "kB = const.k_B.cgs\n",
     "\n",
     "def planck_function(nu):\n",
     "    return 8 * np.pi**2 * r_phot**2 * h * nu**3 / (c2 * (np.exp(h * nu / (kB * temperature)) - 1))"
@@ -225,7 +244,9 @@
     "\n",
     "# In the histogram plot below, the weights argument is used \n",
     "# to make sure our plotted spectrum has the correct y-axis scale\n",
-    "plt.hist(nus, bins=bins, weights=lumin_per_packet/bin_width)\n",
+    "plt.hist(nus.value,\n",
+    "         bins=bins,\n",
+    "         weights=lumin_per_packet/bin_width)\n",
     "\n",
     "# We plot the planck function for comparison\n",
     "plt.plot(nus_planck, planck_function(nus_planck))\n",