From 52d582ab452a1f8aca61bcd903e91eb61ce66e2d Mon Sep 17 00:00:00 2001 From: Aaron LI Date: Sat, 13 May 2017 15:36:11 +0800 Subject: Add astro/oskar/{runOSKAR.py,vis2images.py} --- astro/oskar/runOSKAR.py | 564 ++++++++++++++++++++++++++++++++++++++++++++++ astro/oskar/vis2images.py | 176 +++++++++++++++ 2 files changed, 740 insertions(+) create mode 100755 astro/oskar/runOSKAR.py create mode 100755 astro/oskar/vis2images.py (limited to 'astro/oskar') 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 +# 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() diff --git a/astro/oskar/vis2images.py b/astro/oskar/vis2images.py new file mode 100755 index 0000000..40fe01b --- /dev/null +++ b/astro/oskar/vis2images.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +# +# Copyright (c) 2017 Weitian LI +# MIT license +# +# 2017-04-08 +# + +""" +Make images from simulated visibilities using CASA 'clean' task. + + +Credits +------- +[1] CASA: Common Astronomy Software Applications + https://casa.nrao.edu/ +[2] CASA: clean + https://casa.nrao.edu/docs/TaskRef/clean-task.html +[2] GitHub: OxfordSKA/EoR - EoR_pipeline/make_images.py + https://github.com/OxfordSKA/EoR/blob/master/EoR_pipeline/make_images.py +""" + +import os +import sys +import argparse +import subprocess + +try: + from configparser import ConfigParser +except ImportError: + # CASA (4.7) ships Python 2.7 + from ConfigParser import ConfigParser + + +class Settings: + """ + Manage settings for imaging. + """ + def __init__(self, infile): + self.infile = infile + # Python 2.7's ConfigParser doesn't have parameter 'interpolation' + config = ConfigParser() + config.read(infile) + + DEFAULTS = { + "casa_bin": "casa", + "output_ms_fn": "visibility/visibility_{freq:.2f}.ms", + "output_image_rootname": "image/image_{freq:.2f}", + "clean_niter": '500', # int + "clean_gridmode": "widefield", + "clean_wprojplanes": '256', # int + "clean_weighting": "natural", + "clean_uvrange": "", + } + + casa_bin = config.get("my", "casa_bin", vars=DEFAULTS) + self.casa_bin = os.path.expanduser(casa_bin) + + # Width/X and height/Y of the input FITS image (unit: pixel) + size = config.get("my", "image_size").split(",") + self.image_width = int(size[0]) + self.image_height = int(size[1]) + self.image_size = (self.image_width, self.image_height) + + # Pixel size of the input FITS image (unit: arcsec) + self.image_pixsize = config.getfloat("my", "image_pixsize") + + # String format pattern for the output simulated visibility + # data in MeasurementSet format. + self.output_ms_fn = config.get("my", "output_ms_fn", vars=DEFAULTS) + + # String format pattern for the output image rootname created + # from visibility using CASA 'clean' task. + self.output_image_rootname = config.get( + "my", "output_image_rootname", vars=DEFAULTS) + + # Number of iteration over which to clean (i.e., deconvolve the + # dirty image) + # NOTE: Python 2.7's .getint() not support 'vars' parameter + self.clean_niter = int(config.get("my", "clean_niter", vars=DEFAULTS)) + + # Apply corrections for non-coplanar effects during imaging + # using the W-Projection algorithm + self.clean_gridmode = config.get("my", "clean_gridmode", + vars=DEFAULTS) + + # Number of pre-computed w-planes used for the W-Projection + # algorithm + self.clean_wprojplanes = int(config.get("my", "clean_wprojplanes", + vars=DEFAULTS)) + + # Decides how much weight is given to uv grid points to allow + # for different sampling densities + self.clean_weighting = config.get("my", "clean_weighting", + vars=DEFAULTS) + + # Range of baselines to include when generating the image + self.clean_uvrange = config.get("my", "clean_uvrange", vars=DEFAULTS) + + +class Imager: + """ + Imager using CASA 'clean' task to create image from visibility. + """ + def __init__(self, ms, rootname): + self.ms = ms + self.rootname = rootname + + def make_image(self, settings): + """ + Make image from visibility using 'clean' task. + """ + default(clean) + ret = clean( + vis=self.ms, + imagename=self.rootname, + niter=settings.clean_niter, + gridmode=settings.clean_gridmode, + wprojplanes=settings.clean_wprojplanes, + uvrange=settings.clean_uvrange, + weighting=settings.clean_weighting, + imsize=[settings.image_width, settings.image_height], + cell=[settings.image_pixsize, settings.image_pixsize] + ) + return ret + + def export_fits(self): + """ + Export create image & psf into FITS. + """ + for imgtype in ["image", "psf"]: + imgfile = "%s.%s" % (self.rootname, imgtype) + fitsfile = imgfile + ".fits" + exportfits(imagename=imgfile, fitsimage=fitsfile) + + +def main_casa(): + imgroot = sys.argv[-1] + msfile = sys.argv[-2] + configfile = sys.argv[-3] + + settings = Settings(configfile) + imager = Imager(msfile, imgroot) + imager.make_image(settings) + imager.export_fits() + + +def main(): + parser = argparse.ArgumentParser( + description="Make images from visibilities using CASA") + parser.add_argument("config", help="Configuration file") + parser.add_argument("frequency", type=float, + help="frequency slice to imaging [MHz]") + args = parser.parse_args() + settings = Settings(args.config) + + msfile = settings.output_ms_fn.format(freq=args.frequency) + imgroot = settings.output_image_rootname.format(freq=args.frequency) + dname = os.path.dirname(imgroot) + if not os.path.isdir(dname): + os.makedirs(dname) + + cmd = [ + settings.casa_bin, "--nogui", "--nologger", "--log2term", + "-c", __file__, args.config, msfile, imgroot + ] + print("CMD: %s" % " ".join(cmd)) + subprocess.check_call(cmd) + + +if __name__ == "__main__": + progname = os.path.basename(sys.argv[0]) + if progname in ["casa", "casapy", "casapy.py"]: + main_casa() + else: + main() -- cgit v1.2.2