Merge pull request #160 from karllark/enhance_model_plot

Enhance model plot

Author: karllark

measure_extinction/model.py

@@ -1,6 +1,7 @@
 import copy
 import numpy as np
 import matplotlib.pyplot as plt
+from matplotlib.ticker import ScalarFormatter
 import astropy.units as u
 from astropy.table import QTable
 import scipy.optimize as op
@@ -417,7 +418,7 @@ def stellar_sed(self, moddata):
             # check for any zero distance cases, requested parameters are directly on a model
             # in this case set the distance to 0.01 of the min distance so this model dominates
             tvals = dist2[gsindxs] == 0.0
-            if (sum(tvals) > 0):
+            if sum(tvals) > 0:
                 dist2[gsindxs[tvals]] = 0.01 * np.min(dist2[gsindxs[~tvals]])
 
             weights = 1.0 / np.sqrt(dist2[gsindxs])
@@ -964,7 +965,14 @@ def fit_sampler(
 
         return (outmod, flat_samples, sampler)
 
-    def plot(self, obsdata, modinfo, resid_range=10.0, lyaplot=False):
+    def plot(
+        self,
+        obsdata,
+        modinfo,
+        resid_range=10.0,
+        lyaplot=False,
+        xticks=[0.1, 0.2, 0.3, 0.5, 0.7, 1.0, 2.0],
+    ):
         """
         Standard plot showing the data and best fit.
 
@@ -978,6 +986,12 @@ def plot(self, obsdata, modinfo, resid_range=10.0, lyaplot=False):
 
         resid_range : float
             percentage value for the +/- range for the residual plot
+
+        lyaplot : boolean
+            set to add two panels giving the Ly-alpha fit and residuals
+
+        xticks : vector
+            set to a vector of floats giving the values for the xticks
         """
         # plotting setup for easier to read plots
         fontsize = 16
@@ -1035,7 +1049,7 @@ def plot(self, obsdata, modinfo, resid_range=10.0, lyaplot=False):
         for cspec in obsdata.data.keys():
             if cspec == "BAND":
                 ptype = "o"
-                rcolor = "g"
+                rcolor = "k"
             else:
                 ptype = "-"
                 rcolor = "k"
@@ -1052,7 +1066,13 @@ def plot(self, obsdata, modinfo, resid_range=10.0, lyaplot=False):
             multval = multlam * nvals
 
             if first_pass:
-                plabs = ["Obs", "Star", "w/ Foreground", "w/ Dust Ext", "w/ Dust+Gas Ext"]
+                plabs = [
+                    "Obs",
+                    "Star",
+                    "w/ Foreground",
+                    "w/ Dust Ext",
+                    "w/ Dust+Gas Ext",
+                ]
             else:
                 plabs = [None, None, None, None, None]
 
@@ -1061,15 +1081,28 @@ def plot(self, obsdata, modinfo, resid_range=10.0, lyaplot=False):
                     cax.plot(
                         cwaves, modsed_nofore[cspec] * multlam, "c" + ptype, alpha=0.2
                     )
-                    cax.plot(cwaves, modsed_nofore[cspec] * multval, "c" + ptype, label=plabs[1])
-                    cax.plot(cwaves, modsed[cspec] * multval, "b" + ptype, label=plabs[2])
+                    cax.plot(
+                        cwaves,
+                        modsed_nofore[cspec] * multval,
+                        "c" + ptype,
+                        label=plabs[1],
+                    )
+                    cax.plot(
+                        cwaves, modsed[cspec] * multval, "b" + ptype, label=plabs[2]
+                    )
                 else:
-                    cax.plot(cwaves, modsed[cspec] * multval, "b" + ptype, label=plabs[1])
+                    cax.plot(
+                        cwaves, modsed[cspec] * multval, "b" + ptype, label=plabs[1]
+                    )
                 cax.plot(cwaves, modsed[cspec] * multlam, "b" + ptype, alpha=0.2)
                 cax.plot(cwaves, ext_modsed[cspec] * multlam, "g" + ptype, alpha=0.2)
-                cax.plot(cwaves, ext_modsed[cspec] * multval, "g" + ptype, label=plabs[3])
+                cax.plot(
+                    cwaves, ext_modsed[cspec] * multval, "g" + ptype, label=plabs[3]
+                )
                 cax.plot(cwaves, hi_ext_modsed[cspec] * multlam, "r" + ptype, alpha=0.2)
-                cax.plot(cwaves, hi_ext_modsed[cspec] * multval, "r" + ptype, label=plabs[4])
+                cax.plot(
+                    cwaves, hi_ext_modsed[cspec] * multval, "r" + ptype, label=plabs[4]
+                )
 
                 gvals = obsdata.data[cspec].fluxes > 0.0
                 cax.plot(
@@ -1151,6 +1184,13 @@ def plot(self, obsdata, modinfo, resid_range=10.0, lyaplot=False):
         axes[1].set_xscale("log")
         ax.set_yscale("log")
 
+        if xticks is not None:
+            for tax in [ax, axes[1]]:
+                tax.xaxis.set_major_formatter(ScalarFormatter())
+                tax.xaxis.set_minor_formatter(ScalarFormatter())
+                tax.set_xticks(xticks, minor=True)
+                tax.tick_params(axis="x", which="minor", labelsize=fontsize * 0.8)
+
         ydelt = yrange[1] - yrange[0]
         yrange[0] = 10 ** (yrange[0] - 0.1 * ydelt)
         yrange[1] = 10 ** (yrange[1] + 0.1 * ydelt)

measure_extinction/plotting/plot_model.py

@@ -18,7 +18,9 @@ def main():
         choices=["chains", "corner", "bestfit"],
         default="bestfit",
     )
-    parser.add_argument("--residrange", help="residual range in percentage", default=50.0, type=float)
+    parser.add_argument(
+        "--residrange", help="residual range in percentage", default=50.0, type=float
+    )
     parser.add_argument("--burnfrac", help="burn fraction", default=0.5, type=float)
     parser.add_argument(
         "--obspath",
@@ -117,7 +119,9 @@ def main():
         fig = memod.plot_sampler_corner(flat_samples)
         save_str = "_mefit_corner"
     else:
-        fig = memod.plot(reddened_star, modinfo, lyaplot=True, resid_range=args.residrange)
+        fig = memod.plot(
+            reddened_star, modinfo, lyaplot=True, resid_range=args.residrange
+        )
         save_str = "_mefit_mcmc"
 
     # plot or save to a file
@@ -128,5 +132,6 @@ def main():
     else:
         plt.show()
 
+
 if __name__ == "__main__":
     main()