Add astro/oskar/{runOSKAR.py,vis2images.py}

1 files changed, 564 insertions, 0 deletions

diff --git a/astro/oskar/runOSKAR.py b/astro/oskar/runOSKAR.py
new file mode 100755
index 0000000..fa0ce2d
--- /dev/null
+++ b/astro/oskar/runOSKAR.py
@@ -0,0 +1,564 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2017 Weitian LI <liweitianux@live.com>
+# MIT license
+#
+# 2017-04-07
+#
+
+"""
+Run OSKAR to simulate the visibilities from the sky model specified
+by a FITS image.
+
+
+Credits
+-------
+[1] GitHub: OxfordSKA/OSKAR
+    https://github.com/OxfordSKA/OSKAR
+[2] GitHub: OxfordSKA/EoR - Emma_files/sim_tidy.py
+    https://github.com/OxfordSKA/EoR/blob/master/Emma_files/sim_tidy.py
+"""
+
+import os
+import sys
+import subprocess
+import configparser
+import argparse
+import logging
+
+import numpy as np
+import astropy.io.fits as fits
+import astropy.units as au
+import astropy.constants as ac
+from astropy.wcs import WCS
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(os.path.basename(sys.argv[0]))
+
+
+class Settings:
+    """
+    OSKAR settings manager.
+    """
+    def __init__(self, infile):
+        self.infile = infile
+        self.config = configparser.ConfigParser(interpolation=None)
+        self.config.read(infile)
+        logger.info("Read in configuration file: %s" % infile)
+        self.init_oskar_settings()
+        self.update_oskar_settings(self.config)
+
+    @property
+    def my_settings(self):
+        return self.config["my"]
+
+    @property
+    def dryrun(self):
+        return self.my_settings.getboolean("dryrun", fallback=False)
+
+    @property
+    def clobber(self):
+        return self.my_settings.getboolean("clobber", fallback=False)
+
+    @property
+    def quiet(self):
+        return self.my_settings.getboolean("quiet", fallback=False)
+
+    @property
+    def oskar_bin(self):
+        oskar = self.my_settings.get("oskar_bin",
+                                     fallback="oskar_sim_interferometer")
+        return os.path.expanduser(oskar)
+
+    @property
+    def output_settings_fn(self):
+        """
+        String format pattern for the output OSKAR settings file.
+        """
+        default = "settings/sim_interferometer_{freq:.2f}.ini"
+        return self.my_settings.get("output_settings_fn", fallback=default)
+
+    @property
+    def output_skymodel_fn(self):
+        """
+        String format pattern for the output OSKAR sky model file.
+        """
+        default = "skymodel/skymodel_{freq:.2f}.txt"
+        return self.my_settings.get("output_skymodel_fn", fallback=default)
+
+    @property
+    def output_skyfits_fn(self):
+        """
+        String format pattern for the output FITS slice of the sky model.
+        """
+        default = "skymodel/skymodel_{freq:.2f}.fits"
+        return self.my_settings.get("output_skyfits_fn", fallback=default)
+
+    @property
+    def output_ms_fn(self):
+        """
+        String format pattern for the output simulated visibility
+        data in MeasurementSet format.
+        """
+        default = "visibility/visibility_{freq:.2f}.ms"
+        return self.my_settings.get("output_ms_fn", fallback=default)
+
+    @property
+    def output_vis_fn(self):
+        """
+        String format pattern for the output simulated visibility
+        data in OSKAR binary format.
+        """
+        default = "visibility/visibility_{freq:.2f}.oskar"
+        return self.my_settings.get("output_vis_fn", fallback=default)
+
+    @property
+    def telescope_model(self):
+        """
+        Telescope model used for visibility simulations.
+        """
+        return self.my_settings["telescope_model"]
+
+    @property
+    def input_cube(self):
+        """
+        Input FITS spectral cube.
+        """
+        return self.my_settings["input_cube"]
+
+    @property
+    def image_size(self):
+        """
+        Width/X and height/Y of the input FITS image (unit: pixel)
+        """
+        size = self.my_settings["image_size"].split(",")
+        return (int(size[0]), int(size[1]))
+
+    @property
+    def image_pixsize(self):
+        """
+        Pixel size of the input FITS image (unit: arcsec)
+        """
+        return self.my_settings.getfloat("image_pixsize")
+
+    @property
+    def frequency(self):
+        """
+        Frequency of the input image. (unit: MHz)
+
+        NOTE: required if the above input FITS file is not a cube, but
+              a 2D image.
+        """
+        return self.my_settings.getfloat("frequency")
+
+    @property
+    def bandwidth(self):
+        """
+        Bandwidth of the input image. (unit: MHz)
+        """
+        return self.my_settings.getfloat("bandwidth")
+
+    @property
+    def ra0(self):
+        """
+        R.A. of the center of the input sky field.
+        unit: deg
+        """
+        return self.my_settings.getfloat("ra0", fallback=0.0)
+
+    @property
+    def dec0(self):
+        """
+        Dec. of the center of the input sky field.
+        unit: deg
+        """
+        return self.my_settings.getfloat("dec0", fallback=-27.0)
+
+    @property
+    def use_gpus(self):
+        """
+        Whether to GPUs
+        """
+        return self.my_settings.getboolean("use_gpus", fallback=False)
+
+    @property
+    def start_time(self):
+        """
+        Start time of the simulating observation
+        """
+        # This default time keeps 'EoR0' region above horizon for 12 hours.
+        # SKA EoR0 region: (ra, dec) = (0, -27) [deg]
+        default = "2000-01-01T03:30:00.000"
+        return self.my_settings.getfloat("start_time", fallback=default)
+
+    @property
+    def obs_length(self):
+        """
+        Observation length of time (unit: s).
+        """
+        default = 12.0 * 3600  # 12 hours
+        return self.my_settings.getfloat("obs_length", fallback=default)
+
+    @property
+    def obs_interval(self):
+        """
+        Observation interval providing the number of time steps in the
+        output data (unit: s).
+        """
+        default = 10.0  # [s]
+        return self.my_settings.getfloat("obs_interval", fallback=default)
+
+    @property
+    def time_average(self):
+        """
+        Correlator time-average duration to simulate time-averaging smearing
+        (unit: s).
+        """
+        default = 10.0  # [s]
+        return self.my_settings.getfloat("time_average", fallback=default)
+
+    def init_oskar_settings(self):
+        """
+        Initialize a `ConfigParser` instance with the default settings
+        for 'oskar_sim_interferometer'.
+        """
+        settings = configparser.ConfigParser()
+        settings.read_dict({
+            "General": {
+                "app": "oskar_sim_interferometer",
+            },
+            "simulator": {
+                "use_gpus": self.use_gpus,
+                "max_sources_per_chunk": 65536,
+                "double_precision": "true",
+                "keep_log_file": "true",
+            },
+            "sky": {
+                "advanced/apply_horizon_clip": "false",
+            },
+            "observation": {
+                "phase_centre_ra_deg": self.ra0,
+                "phase_centre_dec_deg": self.dec0,
+                "start_time_utc": self.start_time,
+                "length": self.obs_length,
+                "num_time_steps":
+                    int(np.ceil(self.obs_length/self.obs_interval)),
+                "num_channels": 1,
+            },
+            "telescope": {
+                "input_directory": self.telescope_model,
+                "pol_mode": "Scalar",
+                "normalise_beams_at_phase_centre": "true",
+                "allow_station_beam_duplication": "true",
+                "aperture_array/array_pattern/enable": "true",
+                "aperture_array/element_pattern/functional_type": "Dipole",
+                "aperture_array/element_pattern/dipole_length": 0.5,
+                "aperture_array/element_pattern/dipole_length_units":
+                    "Wavelengths",
+                "station_type": "Aperture array",
+            },
+            "interferometer": {
+                "channel_bandwidth_hz": self.bandwidth * 1e6,
+                "time_average_sec": self.time_average,
+                "uv_filter_min": "min",
+                "uv_filter_max": "max",
+                "uv_filter_units": "Wavelengths",
+            }
+        })
+        self.oskar_settings = settings
+        logger.info("Initialized 'oskar_settings'")
+
+    def update_oskar_settings(self, config):
+        """
+        Update the OSKAR settings with the loaded user configurations.
+        """
+        for section in self.oskar_settings.sections():
+            if section in config:
+                for key, value in config[section].items():
+                    self.oskar_settings[section][key] = value
+                    logger.info("oskar_settings: [%s]%s = %s" % (
+                        section, key, value))
+        logger.info("Updated 'oskar_settings'")
+
+    def write_oskar_settings(self, outfile, clobber=False):
+        """
+        Write the settings file for 'oskar_sim_interferometer'.
+        """
+        if os.path.exists(outfile) and (not clobber):
+            raise OSError("oskar settings file already exists: " % outfile)
+        with open(outfile, "w") as fp:
+            # NOTE: OSKAR do NOT like space around '='
+            self.oskar_settings.write(fp, space_around_delimiters=False)
+        logger.info("Wrote oskar settings file: %s" % outfile)
+
+
+class SpectralCube:
+    """
+    Manipulate the FITS spectral cube.
+
+    NOTE: The FITS data as `numpy.ndarray` has the opposite index
+          ordering, which likes the Fortran style, i.e., fastest
+          changing axis last: data[frequency, y, x]
+    """
+    def __init__(self, infile):
+        self.infile = infile
+        with fits.open(infile) as hdulist:
+            self.header = hdulist[0].header
+            self.cube = hdulist[0].data
+        self.wcs = WCS(self.header)
+        logger.info("Loaded FITS spectral cube: %s" % infile)
+        logger.info("Spectral cube: width=%d, height=%d" %
+                    (self.width, self.height))
+        if not self.is_cube:
+            logger.warning("NOT a spectral cube!")
+        else:
+            logger.info("Number of frequencies: %d" % self.nfreq)
+
+    @property
+    def naxis(self):
+        return self.header["NAXIS"]
+
+    @property
+    def is_cube(self):
+        return self.naxis == 3
+
+    @property
+    def width(self):
+        """
+        Width of the image, i.e., X axis.
+        """
+        return self.header["NAXIS1"]
+
+    @property
+    def height(self):
+        """
+        Height of the image, i.e., Y axis.
+        """
+        return self.header["NAXIS2"]
+
+    @property
+    def nfreq(self):
+        return self.header["NAXIS3"]
+
+    @property
+    def frequencies(self):
+        """
+        Frequencies of this cube. (unit: MHz)
+        If the input file is not a cube, then return 'None'.
+        """
+        if not self.is_cube:
+            logger.warning("Input FITS file is not a spectral cube: %s" %
+                           self.infile)
+            return None
+
+        nfreq = self.nfreq
+        pix = np.zeros(shape=(nfreq, self.naxis), dtype=np.int)
+        pix[:, -1] = np.arange(nfreq)
+        world = self.wcs.wcs_pix2world(pix, 0)
+        freqMHz = world[:, -1] / 1e6  # Hz -> MHz
+        return freqMHz
+
+    def get_slice(self, nth=0):
+        """
+        Extract the specified nth frequency slice from the cube.
+        """
+        if not self.is_cube:
+            logger.warning("Input FITS file is not a spectral cube: %s" %
+                           self.infile)
+            return self.cube
+        else:
+            return self.cube[nth, :, :]
+
+
+class SkyModel:
+    """
+    OSKAR sky model.
+    """
+    def __init__(self, image, freq, pixsize, ra0, dec0):
+        self.image = image  # K (brightness temperature)
+        self.freq = freq  # MHz
+        self.pixsize = pixsize  # arcsec
+        self.ra0 = ra0  # deg
+        self.dec0 = dec0  # deg
+        logger.info("SkyModel: Loaded image @ %.2f [MHz]" % freq)
+
+    @property
+    def wcs(self):
+        """
+        WCS for the given image assuming the 'SIN' projection.
+        """
+        shape = self.image.shape
+        delta = self.pixsize / 3600.0  # deg
+        wcs_ = WCS(naxis=2)
+        wcs_.wcs.ctype = ["RA---SIN", "DEC--SIN"]
+        wcs_.wcs.crval = np.array([self.ra0, self.dec0])
+        wcs_.wcs.crpix = np.array([shape[1], shape[0]]) / 2.0 + 1
+        wcs_.wcs.cdelt = np.array([delta, delta])
+        return wcs_
+
+    @property
+    def fits_header(self):
+        header = self.wcs.to_header()
+        header["BUNIT"] = ("Jy/pixel", "Brightness unit")
+        header["FREQ"] = (self.freq, "Frequency [MHz]")
+        header["RA0"] = (self.ra0, "Center R.A. [deg]")
+        header["DEC0"] = (self.dec0, "Center Dec. [deg]")
+        return header
+
+    @property
+    def factor_K2JyPixel(self):
+        """
+        Conversion factor to convert brightness unit from 'K' to 'Jy/pixel'
+
+        http://www.iram.fr/IRAMFR/IS/IS2002/html_1/node187.html
+        """
+        pixarea = np.deg2rad(self.pixsize/3600.0) ** 2  # [sr]
+        kB = ac.k_B.si.value  # Boltzmann constant [J/K]
+        c0 = ac.c.si.value  # speed of light in vacuum [m/s]
+        freqHz = self.freq * 1e6  # [Hz]
+        factor = 2*kB * 1.0e26 * pixarea * (freqHz/c0)**2
+        return factor
+
+    @property
+    def ra_dec(self):
+        """
+        Calculate the (ra, dec) of each image pixel using the above WCS.
+
+        NOTE: axis ordering difference between numpy array and FITS
+        """
+        shape = self.image.shape
+        wcs = self.wcs
+        x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
+        pix = np.column_stack([x.flatten(), y.flatten()])
+        world = wcs.wcs_pix2world(pix, 0)
+        ra = world[:, 0].reshape(shape)
+        dec = world[:, 1].reshape(shape)
+        return (ra, dec)
+
+    @property
+    def sky(self):
+        """
+        OSKAR sky model array converted from the input image.
+
+        Columns
+        -------
+        ra : (J2000) right ascension (deg)
+        dec : (J2000) declination (deg)
+        flux : source (Stokes I) flux density (Jy)
+        """
+        ra, dec = self.ra_dec
+        ra = ra.flatten()
+        dec = dec.flatten()
+        flux = self.image.flatten() * self.factor_K2JyPixel
+        mask = flux > 1e-40
+        sky_ = np.column_stack([ra[mask], dec[mask], flux[mask]])
+        return sky_
+
+    def write_sky_model(self, outfile, clobber=False):
+        """
+        Write the converted sky model for simulation.
+        """
+        if os.path.exists(outfile) and (not clobber):
+            raise OSError("oskar sky model file already exists: " % outfile)
+        sky = self.sky
+        header=("Frequency = %.3f [MHz]\n" % self.freq +
+                "Pixel size = %.2f arcsec\n" % self.pixsize +
+                "RA0 = %.4f [deg]\n" % self.ra0 +
+                "Dec0 = %.4f [deg]\n" % self.dec0 +
+                "Number of sources = %d\n\n" % len(sky) +
+                "R.A.[deg]    Dec.[deg]    flux[Jy]")
+        np.savetxt(outfile, sky, fmt='%.10e, %.10e, %.10e', header=header)
+        logger.info("Wrote oskar sky model file: %s" % outfile)
+
+    def write_fits(self, outfile, oldheader=None, clobber=False):
+        if os.path.exists(outfile) and (not clobber):
+            raise OSError("Sky FITS already exists: " % outfile)
+        if oldheader is not None:
+            header = oldheader
+            header.extend(self.fits_header, update=True)
+        else:
+            header = self.fits_header
+        image = self.image * self.factor_K2JyPixel
+        hdu = fits.PrimaryHDU(data=image, header=header)
+        hdu.writeto(outfile, overwrite=True)
+        logger.info("Wrote sky FITS to file: %s" % outfile)
+
+
+class Oskar:
+    """
+    Run OSKAR simulations
+    """
+    def __init__(self, settings):
+        self.settings = settings
+
+    def run(self, settingsfile, dryrun=False):
+        cmd = [self.settings.oskar_bin]
+        if self.settings.quiet:
+            cmd += ["--quiet"]
+        cmd += [settingsfile]
+        logger.info("Running OSKAR simulator: CMD: %s" % " ".join(cmd))
+        if dryrun:
+            logger.info("Dry run!")
+        else:
+            subprocess.check_call(cmd)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Run OSKAR to simulate visibilities")
+    parser.add_argument("config", help="Configuration file")
+    args = parser.parse_args()
+
+    settings = Settings(args.config)
+    clobber = settings.clobber
+    image_cube = SpectralCube(settings.input_cube)
+    frequencies = image_cube.frequencies  # [MHz]
+    if frequencies is None:
+        frequencies = [settings.frequency]
+    logger.info("Number of image slices/frequencies: %d" % len(frequencies))
+
+    for nth, freq in enumerate(frequencies):
+        logger.info(">>> Processing #%d/%d image slice @ %.2f [MHz] <<<" %
+                    (nth+1, len(frequencies), freq))
+        settingsfile = settings.output_settings_fn.format(freq=freq)
+        skymodelfile = settings.output_skymodel_fn.format(freq=freq)
+        skyfitsfile = settings.output_skyfits_fn.format(freq=freq)
+        msfile = settings.output_ms_fn.format(freq=freq)
+        visfile = settings.output_vis_fn.format(freq=freq)
+        for filepath in [settingsfile, skymodelfile, skyfitsfile,
+                         msfile, visfile]:
+            dname = os.path.dirname(filepath)
+            if not os.path.isdir(dname):
+                os.makedirs(dname)
+
+        newconfig = configparser.ConfigParser()
+        newconfig.read_dict({
+            "sky": {
+                "oskar_sky_model/file": skymodelfile,
+            },
+            "observation": {
+                "start_frequency_hz": freq * 1e6,
+            },
+            "interferometer": {
+                "oskar_vis_filename": visfile,
+                "ms_filename": msfile,
+            },
+        })
+        settings.update_oskar_settings(newconfig)
+        settings.write_oskar_settings(outfile=settingsfile, clobber=clobber)
+
+        image_slice = image_cube.get_slice(nth)
+        skymodel = SkyModel(image=image_slice, freq=freq,
+                            pixsize=settings.image_pixsize,
+                            ra0=settings.ra0, dec0=settings.dec0)
+        skymodel.write_sky_model(skymodelfile, clobber=clobber)
+        skymodel.write_fits(skyfitsfile, oldheader=image_cube.header,
+                            clobber=clobber)
+
+        oskar = Oskar(settings)
+        oskar.run(settingsfile, dryrun=settings.dryrun)
+
+
+if __name__ == '__main__':
+    main()