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.

3 files changed, 121 insertions, 117 deletions

diff --git a/fg21sim/configs/manager.py b/fg21sim/configs/manager.py
index 1cddb75..172811a 100644
--- a/fg21sim/configs/manager.py
+++ b/fg21sim/configs/manager.py
@@ -47,7 +47,8 @@ def _get_configspec():
 
 
 def _flatten_dict(d, sep="/", parent_key=""):
-    """Recursively flatten a nested dictionary with keys compressed.
+    """
+    Recursively flatten a nested dictionary with keys compressed.
 
     The dictionary is recursively flattened into a one-level dictionary,
     i.e., all the leaves are raised to the top level.
@@ -104,7 +105,8 @@ def _flatten_dict(d, sep="/", parent_key=""):
 
 
 class ConfigManager:
-    """Manage the default configurations with specifications, as well as
+    """
+    Manage the default configurations with specifications, as well as
     the user configurations.
 
     Both the default configurations and user configurations are validated
@@ -139,7 +141,8 @@ class ConfigManager:
     userconfig = None
 
     def __init__(self, userconfig=None):
-        """Load the bundled default configurations and specifications.
+        """
+        Load the bundled default configurations and specifications.
         If the ``userconfig`` provided, the user configurations is also
         loaded, validated, and merged.
         """
@@ -156,7 +159,8 @@ class ConfigManager:
             self.read_userconfig(userconfig)
 
     def merge(self, config):
-        """Simply merge the given configurations without validation.
+        """
+        Simply merge the given configurations without validation.
 
         Parameters
         ----------
@@ -172,7 +176,8 @@ class ConfigManager:
         self._config.merge(config)
 
     def read_config(self, config):
-        """Read, validate and merge the input config.
+        """
+        Read, validate and merge the input config.
 
         Parameters
         ----------
@@ -191,7 +196,8 @@ class ConfigManager:
         logger.info("Loaded additional config")
 
     def read_userconfig(self, userconfig):
-        """Read user configuration file, validate, and merge into the
+        """
+        Read user configuration file, validate, and merge into the
         default configurations.
 
         Parameters
@@ -235,7 +241,8 @@ class ConfigManager:
         logger.warning("Reset the configurations to the copy of defaults!")
 
     def _validate(self, config):
-        """Validate the config against the specification using a default
+        """
+        Validate the config against the specification using a default
         validator.  The validated config values are returned if success,
         otherwise, the ``ConfigError`` raised with details.
         """
@@ -265,7 +272,8 @@ class ConfigManager:
         return config
 
     def check_all(self, raise_exception=True):
-        """Further check the whole configurations through a set of custom
+        """
+        Further check the whole configurations through a set of custom
         checker functions, which may check one config option against its
         context if necessary to determine whether it has a valid value.
 
@@ -303,7 +311,8 @@ class ConfigManager:
         return config.get(key, fallback)
 
     def getn(self, key, from_default=False):
-        """Get the value of a config option specified by the input key from
+        """
+        Get the value of a config option specified by the input key from
         from the configurations which is a nested dictionary.
 
         Parameters
@@ -343,7 +352,8 @@ class ConfigManager:
             raise KeyError("%s: invalid key" % "/".join(key))
 
     def setn(self, key, value):
-        """Set the value of config option specified by a list of keys or a
+        """
+        Set the value of config option specified by a list of keys or a
         "/"-separated keys string.
 
         The supplied key-value config pair is first used to create a
@@ -408,7 +418,8 @@ class ConfigManager:
             key="/".join(key), val_new=val_new, val_old=val_old))
 
     def get_path(self, key):
-        """Return the absolute path of the file/directory specified by the
+        """
+        Return the absolute path of the file/directory specified by the
         config keys.
 
         Parameters
@@ -456,7 +467,8 @@ class ConfigManager:
 
     @property
     def foregrounds(self):
-        """Get all available and enabled foreground components.
+        """
+        Get all available and enabled foreground components.
 
         Returns
         -------
@@ -472,7 +484,8 @@ class ConfigManager:
 
     @property
     def frequencies(self):
-        """Get or calculate if ``frequency/type = custom`` the frequencies
+        """
+        Get or calculate if ``frequency/type = custom`` the frequencies
         where to perform the simulations.
 
         Returns
@@ -494,7 +507,8 @@ class ConfigManager:
 
     @property
     def logging(self):
-        """Get and prepare the logging configurations for
+        """
+        Get and prepare the logging configurations for
         ``logging.basicConfig()`` to initialize the logging module.
 
         NOTE
@@ -528,7 +542,8 @@ class ConfigManager:
         return logconf
 
     def dump(self, from_default=False, flatten=False):
-        """Dump the configurations as plain Python dictionary.
+        """
+        Dump the configurations as plain Python dictionary.
 
         Parameters
         ----------
@@ -559,7 +574,8 @@ class ConfigManager:
         return data
 
     def save(self, outfile=None, clobber=False, backup=True):
-        """Save the configurations to file.
+        """
+        Save the configurations to file.
 
         Parameters
         ----------
diff --git a/fg21sim/foregrounds.py b/fg21sim/foregrounds.py
index 85242c1..2fceef6 100644
--- a/fg21sim/foregrounds.py
+++ b/fg21sim/foregrounds.py
@@ -21,7 +21,6 @@ from collections import OrderedDict
 import numpy as np
 import astropy.units as au
 from astropy.io import fits
-import healpy as hp
 
 from .galactic import (Synchrotron as GalacticSynchrotron,
                        FreeFree as GalacticFreeFree,
@@ -29,14 +28,15 @@ from .galactic import (Synchrotron as GalacticSynchrotron,
 from .extragalactic import (GalaxyClusters as EGGalaxyClusters,
                             PointSources as EGPointSources)
 from .products import Products
-from .utils.fits import write_fits_healpix
+from .sky import get_sky
 
 
 logger = logging.getLogger(__name__)
 
 
 class Foregrounds:
-    """Interface to the simulations of supported foreground components.
+    """
+    Interface to the simulations of supported foreground components.
 
     All the enabled components are also combined to make the total foreground
     map, as controlled by the configurations.
@@ -108,14 +108,15 @@ class Foregrounds:
         #
         self.filename_pattern = self.configs.getn("output/filename_pattern")
         self.use_float = self.configs.getn("output/use_float")
+        self.checksum = self.configs.getn("output/checksum")
         self.clobber = self.configs.getn("output/clobber")
         self.combine = self.configs.getn("output/combine")
         self.prefix = self.configs.getn("output/combine_prefix")
         self.output_dir = self.configs.get_path("output/output_dir")
-        self.nside = self.configs.getn("common/nside")
 
     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 = {
@@ -127,12 +128,12 @@ class Foregrounds:
         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()
-        header["COMP"] = ("Combined foreground",
-                          "Emission component")
+        header["COMP"] = ("Combined foreground", "Emission component")
         header.add_comment("COMPONENTS: " + ", ".join(self.components_id))
         header["UNIT"] = ("Kelvin", "Map unit")
         header["CREATOR"] = (__name__, "File creator")
@@ -146,20 +147,17 @@ class Foregrounds:
         self.header = header
         logger.info("Created FITS header")
 
-    def _output(self, hpmap, frequency):
-        """Write the simulated free-free map to disk with proper header
+    def _output(self, skymap, frequency):
+        """
+        Write the simulated free-free 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()
@@ -169,11 +167,12 @@ class Foregrounds:
             "File creation date"
         )
         if self.use_float:
-            hpmap = hpmap.astype(np.float32)
-        write_fits_healpix(filepath, hpmap, header=header,
-                           clobber=self.clobber)
-        logger.info("Write combined foreground to file: {0}".format(filepath))
-        return filepath
+            skymap = skymap.astype(np.float32)
+        sky = get_sky(configs=self.configs)
+        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."""
@@ -183,7 +182,8 @@ class Foregrounds:
             comp_obj.preprocess()
 
     def simulate(self):
-        """Simulate the enabled components, as well as combine all the
+        """
+        Simulate the enabled components, as well as combine all the
         simulated components to make up the total foregrounds.
 
         This is the *main interface* to the foreground simulations.
@@ -195,22 +195,22 @@ class Foregrounds:
         In this way, less memory is required, since the number of components
         are generally much less than the number of frequency bins.
         """
-        npix = hp.nside2npix(self.nside)
+        sky = get_sky(configs=self.configs)
         nfreq = len(self.frequencies)
         for freq_id, freq in enumerate(self.frequencies):
             logger.info("[#{0}/{1}] ".format(freq_id+1, nfreq) +
                         "Simulating components at {freq} {unit} ...".format(
                             freq=freq, unit=self.freq_unit))
             if self.combine:
-                hpmap_comb = np.zeros(npix)
+                skymap_comb = np.zeros(shape=sky.shape)
             for comp_id, comp_obj in self.components.items():
-                hpmap, filepath = comp_obj.simulate_frequency(freq)
+                skymap, filepath = comp_obj.simulate_frequency(freq)
                 if filepath is not None:
                     self.products.add_product(comp_id, freq_id, filepath)
                 if self.combine:
-                    hpmap_comb += hpmap
+                    skymap_comb += skymap
             if self.combine:
-                filepath_comb = self._output(hpmap_comb, freq)
+                filepath_comb = self._output(skymap_comb, freq)
                 self.products.add_product("combined", freq_id, filepath_comb)
 
     def postprocess(self):
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."""