Simplify plotting function.

icecube · Aug 4, 2020 · 0a6644d · 0a6644d
1 parent 9832272
commit 0a6644d
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 58 deletions.
diff --git a/flarestack/cosmo/icecube_diffuse_flux/__init__.py b/flarestack/cosmo/icecube_diffuse_flux/__init__.py
@@ -107,51 +107,48 @@ def get_diffuse_flux_contour(fit="joint_15", contour_name="68"):
 
     return best_f, upper_f, lower_f, e_range
 
-def plot_diffuse_flux(label, contour_68, contour_95, e_range):
+def plot_diffuse_flux(fit="joint_15"):
+    plt.figure()
 
+    best_fit_flux, best_fit_gamma, contour_68, contour_95, e_range = load_fit(fit)
 
-    plt.figure()
+    for i, contour in enumerate([68., 95.]):
+        all_c = [contour_68, contour_95][i]
 
-    # Plot 68% contour
+        best_f, upper_f, lower_f, e_range = get_diffuse_flux_contour(fit=fit, contour_name=contour)
+
+        plt.plot(e_range,
+                 e_range ** 2 * best_f(e_range),
+                 alpha=1.0, color="k")
 
-    for (index, norm) in contour_68:
         plt.plot(e_range,
-                 e_range ** 2 * flux_f(e_range, norm, index),
-                 alpha=0.6)
-    plt.plot(
-        e_range,
-        e_range ** 2 * upper_contour(e_range, contour_68),
-        color="k", label="68% contour", linestyle="-"
-    )
-    plt.plot(
-        e_range,
-        e_range ** 2 * lower_contour(e_range, contour_68),
-        color="k", linestyle="-",
-    )
-
-    # Plot 95% contour
-
-    for (index, norm) in contour_95:
+                 e_range ** 2 * lower_f(e_range),
+                 color=f"C{i}",
+                 label=f"{contour} %",
+                 )
+
         plt.plot(e_range,
-                 e_range ** 2 * flux_f(e_range, norm, index),
-                 alpha=0.3)
-    plt.plot(
-        e_range,
-        e_range ** 2 * upper_contour(e_range, contour_95),
-        color="k", linestyle="--", label="95% contour"
-    )
-    plt.plot(
-        e_range,
-        e_range ** 2 * lower_contour(e_range, contour_95),
-        color="k", linestyle="--"
-    )
+                 e_range ** 2 * upper_f(e_range),
+                 color=f"C{i}",
+                 )
+
+        for (index, norm) in all_c:
+
+            plt.plot(e_range,
+                     e_range**2 * flux_f(e_range, norm, index),
+                     alpha=0.15
+                     )
 
     plt.yscale("log")
     plt.xscale("log")
     plt.legend()
     plt.title(r"Diffuse Flux Global Best Fit ($\nu_{\mu} + \bar{\nu}_{\mu})$")
     plt.ylabel(r"$E^{2}\frac{dN}{dE}$[GeV cm$^{-2}$ s$^{-1}$ sr$^{-1}$]")
     plt.xlabel(r"$E_{\nu}$ [GeV]")
-    savepath = illustration_dir + "diffuse_flux_global_fit.pdf"
+    plt.tight_layout()
+    savepath = f"{illustration_dir}diffuse_flux_global_fit_{fit}.pdf"
+
+    logger.info(f"Saving to {savepath}")
+
     plt.savefig(savepath)
     plt.close()
diff --git a/tests/test_cosmo.py b/tests/test_cosmo.py
@@ -56,34 +56,34 @@ def test_get_diffuse_flux(self):
             logging.info("Calculated values {0}".format(res_100_tev))
             logging.info("Reference  values {0}".format(default_flux_100TeV[i]))
 
-    # def test_neutrino_cosmology(self):
-    #
-    #     fit = "joint_15"
-    #     bestfit = get_diffuse_flux_at_100TeV(fit)
-    #
-    #     e_pdf_dict = {
-    #      "energy_pdf_name": "power_law",
-    #      'source_energy_erg': 1.e48 * u.erg,
-    #      "gamma": bestfit[1]
-    #     }
-    #
-    #     # Use Supernova rate
-    #
-    #     key = "strolger_15"
-    #
-    #     ccsn_rate = get_rate("ccsn", evolution_name=key, rate_name=key)
-    #
-    #     res = calculate_transient_cosmology(
-    #      e_pdf_dict, ccsn_rate, "test_CCSN", zmax=8.0,
-    #      diffuse_fit=fit
-    #     )
-    #
-    #     self.assertAlmostEqual(res.value/true_cosmology.value, 1.0, delta=0.1)
+    def test_neutrino_cosmology(self):
+
+        fit = "joint_15"
+        bestfit = get_diffuse_flux_at_100TeV(fit)
+
+        e_pdf_dict = {
+         "energy_pdf_name": "power_law",
+         'source_energy_erg': 1.e48 * u.erg,
+         "gamma": bestfit[1]
+        }
+
+        # Use Supernova rate
+
+        key = "strolger_15"
+
+        ccsn_rate = get_rate("ccsn", evolution_name=key, rate_name=key)
+
+        res = calculate_transient_cosmology(
+         e_pdf_dict, ccsn_rate, "test_CCSN", zmax=8.0,
+         diffuse_fit=fit
+        )
+
+        self.assertAlmostEqual(res.value/true_cosmology.value, 1.0, delta=0.1)
 
     def test_plotting(self):
 
-        for label, (_, _, contour_68, contour_95, e_range, _) in contours.items():
-            plot_diffuse_flux(label, contour_68, contour_95, e_range)
+        for label in contours.keys():
+            plot_diffuse_flux(label)
 
 if __name__ == '__main__':
     unittest.main()