move cmap checking to utils

mtlam · Jun 4, 2020 · afe99da · afe99da
1 parent b376d6b
commit afe99da
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 15 deletions.
diff --git a/pypulse/archive.py b/pypulse/archive.py
@@ -938,7 +938,7 @@ def getPulsarCalibrator(self):
         freqs, caldata, calerrs = self.getLevels(differences=True)
         pol_type = self.subintheader['POL_TYPE']
         fd_poln = self.header['FD_POLN']
-        return Calibrator(freqs, caldata, calerrs, pol_type=pol_type, fd_poln=fd_poln, verbose=self.verbose)
+        return Calibrator(freqs, caldata, calerrs, pol_type=pol_type, fd_poln=fd_poln, Funit=self.getFrequencyUnit(), Sunit=self.getDataUnit(), verbose=self.verbose)
 
 
     def calibrate(self, psrcalar, fluxcalonar=None, fluxcaloffar=None):
@@ -1363,19 +1363,6 @@ def imshow(self, ax=None, cbar=False, mask=None, show=True,
                 fig = plt.figure()
                 ax = fig.add_subplot(111)
 
-            # If cmap is a string, try to get the colormap,
-            # otherwise use default
-            if isinstance(cmap, str):
-                if cmap in plt.colormaps():
-                    cmap = plt.get_cmap(cmap)
-                else:
-                    cmap = None
-            # Try to use cividis as default colormap, otherwise viridis
-            if cmap is None: 
-                if "cividis" in plt.colormaps():
-                    cmap = plt.cm.cividis
-                else:
-                    cmap = plt.cm.viridis
             #cmap.set_bad(color='k', alpha=1.0)
 
             if mask is not None:

diff --git a/pypulse/utils.py b/pypulse/utils.py
@@ -211,7 +211,21 @@ def decimate(x, window_len, error=False):#, mean=True, remainder=False):
             retval += x[i:length:window_len]
         return retval/window_len
 
-def imshow(x, ax=None, origin='lower', interpolation='nearest', aspect='auto', **kwargs):
+def imshow(x, ax=None, origin='lower', interpolation='nearest', aspect='auto', cmap=None, **kwargs):
+    # If cmap is a string, try to get the colormap,
+    # otherwise use default
+    if isinstance(cmap, str):
+        if cmap in plt.colormaps():
+            cmap = plt.get_cmap(cmap)
+        else:
+            cmap = None
+    # Try to use cividis as default colormap, otherwise viridis
+    if cmap is None: 
+        if "cividis" in plt.colormaps():
+            cmap = plt.cm.cividis
+        else:
+            cmap = plt.cm.viridis
+
     if ax is not None:
         im = ax.imshow(x, origin=origin, interpolation=interpolation, aspect=aspect, **kwargs)
     else: