galactic/synchrotron: Update to support sky.py

* Also update foregrounds.py to use sky.py
* Minor fixes to configs/manager.py

TODO: update synchrotron/add_smallscales() to also work with sky patch.

1 files changed, 61 insertions, 73 deletions

diff --git a/fg21sim/galactic/synchrotron.py b/fg21sim/galactic/synchrotron.py
index af59539..7af9088 100644
--- a/fg21sim/galactic/synchrotron.py
+++ b/fg21sim/galactic/synchrotron.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2016 Weitian LI <liweitianux@live.com>
+# Copyright (c) 2016-2017 Weitian LI <weitian@aaronly.me>
 # MIT license
 
 """
@@ -14,7 +14,7 @@ from astropy.io import fits
 import astropy.units as au
 import healpy as hp
 
-from ..utils.fits import read_fits_healpix, write_fits_healpix
+from ..sky import get_sky
 
 
 logger = logging.getLogger(__name__)
@@ -54,7 +54,8 @@ class Synchrotron:
         self.template_unit = au.Unit(
             self.configs.getn(comp+"/template_unit"))
         self.indexmap_path = self.configs.get_path(comp+"/indexmap")
-        self.smallscales = self.configs.getn(comp+"/add_smallscales")
+        self.add_smallscales = self.configs.getn(comp+"/add_smallscales")
+        self.smallscales_added = False
         self.lmin = self.configs.getn(comp+"/lmin")
         self.lmax = self.configs.getn(comp+"/lmax")
         self.prefix = self.configs.getn(comp+"/prefix")
@@ -65,43 +66,24 @@ class Synchrotron:
         self.use_float = self.configs.getn("output/use_float")
         self.checksum = self.configs.getn("output/checksum")
         self.clobber = self.configs.getn("output/clobber")
-        self.nside = self.configs.getn("common/nside")
         self.freq_unit = au.Unit(self.configs.getn("frequency/unit"))
         #
         logger.info("Loaded and setup configurations")
 
-    def _load_template(self):
-        """Load the template map, and upgrade/downgrade the resolution
-        to match the output Nside.
-        """
-        self.template, self.template_header = read_fits_healpix(
-            self.template_path)
-        template_nside = self.template_header["NSIDE"]
-        logger.info("Loaded template map from {0} (Nside={1})".format(
-            self.template_path, template_nside))
-        # Upgrade/downgrade resolution
-        if template_nside != self.nside:
-            self.template = hp.ud_grade(self.template, nside_out=self.nside)
-            logger.info("Upgrade/downgrade template map from Nside "
-                        "{0} to {1}".format(template_nside, self.nside))
+    def _load_maps(self):
+        """Load the template map and spectral index map."""
+        sky = get_sky(self.configs)
+        self.template = sky.load(self.template_path)
+        self.indexmap = sky.load(self.indexmap_path)
 
-    def _load_indexmap(self):
-        """Load the spectral index map, and upgrade/downgrade the resolution
-        to match the output Nside.
+    def _add_smallscales(self):
         """
-        self.indexmap, self.indexmap_header = read_fits_healpix(
-            self.indexmap_path)
-        indexmap_nside = self.indexmap_header["NSIDE"]
-        logger.info("Loaded spectral index map from {0} (Nside={1})".format(
-            self.indexmap_path, indexmap_nside))
-        # Upgrade/downgrade resolution
-        if indexmap_nside != self.nside:
-            self.indexmap = hp.ud_grade(self.indexmap, nside_out=self.nside)
-            logger.info("Upgrade/downgrade spectral index map from Nside "
-                        "{0} to {1}".format(indexmap_nside, self.nside))
+        Add fluctuations on small scales to the template map.
 
-    def _add_smallscales(self):
-        """Add fluctuations on small scales to the template map.
+        NOTE:
+        Only when the input template is a HEALPix map, this function
+        will be applied to add the small-scale fluctuations, which assuming
+        a angular power spectrum model.
 
         XXX/TODO:
         * Support using different models.
@@ -114,7 +96,11 @@ class Synchrotron:
             "An improved source-subtracted and destriped 408-MHz all-sky map"
             Sec. 4.2: Small-scale fluctuations
         """
-        if (not self.smallscales) or (hasattr(self, "hpmap_smallscales")):
+        if (not self.add_smallscales) or (self.smallscales_added):
+            return
+        if self.template.type_ != "healpix":
+            logger.warning("Input template map is NOT a HEALPix map; " +
+                           "skip adding small-scale fluctuations!")
             return
         # To add small scale fluctuations
         # model: Remazeilles15
@@ -130,15 +116,16 @@ class Synchrotron:
                        1.0 - np.exp(-ell[ell_idx]**2 * sigma_tp**2))
         cl[ell < self.lmin] = cl[self.lmin]
         # generate a realization of the Gaussian random field
-        gss = hp.synfast(cls=cl, nside=self.nside, new=True)
+        gss = hp.synfast(cls=cl, nside=self.template.nside, new=True)
         # whiten the Gaussian random field
         gss = (gss - gss.mean()) / gss.std()
-        self.hpmap_smallscales = alpha * gss * self.template**beta
-        self.template += self.hpmap_smallscales
-        logger.info("Added small-scale fluctuations")
+        hpmap_smallscales = alpha * gss * self.template.data**beta
+        self.template.data += hpmap_smallscales
+        logger.info("Added small-scale fluctuations to template map")
 
     def _make_filepath(self, **kwargs):
-        """Make the path of output file according to the filename pattern
+        """
+        Make the path of output file according to the filename pattern
         and output directory loaded from configurations.
         """
         data = {
@@ -150,7 +137,8 @@ class Synchrotron:
         return filepath
 
     def _make_header(self):
-        """Make the header with detail information (e.g., parameters and
+        """
+        Make the header with detail information (e.g., parameters and
         history) for the simulated products.
         """
         header = fits.Header()
@@ -168,20 +156,17 @@ class Synchrotron:
         self.header = header
         logger.info("Created FITS header")
 
-    def output(self, hpmap, frequency):
-        """Write the simulated synchrotron map to disk with proper
+    def output(self, skymap, frequency):
+        """
+        Write the simulated synchrotron map to disk with proper
         header keywords and history.
 
         Returns
         -------
-        filepath : str
-            The (absolute) path to the output HEALPix map file.
+        outfile : str
+            The (absolute) path to the output sky map file.
         """
-        if not os.path.exists(self.output_dir):
-            os.mkdir(self.output_dir)
-            logger.info("Created output dir: {0}".format(self.output_dir))
-        #
-        filepath = self._make_filepath(frequency=frequency)
+        outfile = self._make_filepath(frequency=frequency)
         if not hasattr(self, "header"):
             self._make_header()
         header = self.header.copy()
@@ -191,14 +176,16 @@ class Synchrotron:
             "File creation date"
         )
         if self.use_float:
-            hpmap = hpmap.astype(np.float32)
-        write_fits_healpix(filepath, hpmap, header=header,
-                           clobber=self.clobber, checksum=self.checksum)
-        logger.info("Write simulated map to file: {0}".format(filepath))
-        return filepath
+            skymap = skymap.astype(np.float32)
+        sky = self.template.copy()
+        sky.data = skymap
+        sky.header = header
+        sky.write(outfile, clobber=self.clobber, checksum=self.checksum)
+        return outfile
 
     def preprocess(self):
-        """Perform the preparation procedures for the final simulations.
+        """
+        Perform the preparation procedures for the final simulations.
 
         Attributes
         ----------
@@ -211,54 +198,55 @@ class Synchrotron:
             return
         #
         logger.info("{name}: preprocessing ...".format(name=self.name))
-        self._load_template()
-        self._load_indexmap()
+        self._load_maps()
         self._add_smallscales()
         #
         self._preprocessed = True
 
     def simulate_frequency(self, frequency):
-        """Transform the template map to the requested frequency,
+        """
+        Transform the template map to the requested frequency,
         according to the spectral model and using an spectral index map.
 
         Returns
         -------
-        hpmap_f : 1D `~numpy.ndarray`
-            The HEALPix map (RING ordering) at the input frequency.
+        skymap_f : 1D `~numpy.ndarray`
+            The sky map at the input frequency.
         filepath : str
-            The (absolute) path to the output HEALPix file if saved,
+            The (absolute) path to the output file if saved,
             otherwise ``None``.
         """
         self.preprocess()
         #
         logger.info("Simulating {name} map at {freq} ({unit}) ...".format(
             name=self.name, freq=frequency, unit=self.freq_unit))
-        hpmap_f = (self.template *
-                   (frequency / self.template_freq) ** self.indexmap)
+        skymap_f = (self.template.data *
+                    (frequency / self.template_freq) ** self.indexmap.data)
         #
         if self.save:
-            filepath = self.output(hpmap_f, frequency)
+            filepath = self.output(skymap_f, frequency)
         else:
             filepath = None
-        return (hpmap_f, filepath)
+        return (skymap_f, filepath)
 
     def simulate(self, frequencies):
-        """Simulate the synchrotron map at the specified frequencies.
+        """
+        Simulate the synchrotron map at the specified frequencies.
 
         Returns
         -------
-        hpmaps : list[1D `~numpy.ndarray`]
-            List of HEALPix maps (in RING ordering) at each frequency.
+        skymaps : list[1D `~numpy.ndarray`]
+            List of sky maps at each frequency.
         paths : list[str]
-            List of (absolute) path to the output HEALPix maps.
+            List of (absolute) path to the output sky maps.
         """
-        hpmaps = []
+        skymaps = []
         paths = []
         for f in np.array(frequencies, ndmin=1):
-            hpmap_f, filepath = self.simulate_frequency(f)
-            hpmaps.append(hpmap_f)
+            skymap_f, filepath = self.simulate_frequency(f)
+            skymaps.append(skymap_f)
             paths.append(filepath)
-        return (hpmaps, paths)
+        return (skymaps, paths)
 
     def postprocess(self):
         """Perform the post-simulation operations before the end."""