utils.py

"""
Utility module for the various scripts needed to correlate LWA and VLA data.

"""

import os
import re
import math
import time
import ephem
import errno
import fcntl
import numpy as np
import shutil
import tempfile
import subprocess
from datetime import datetime

from lsl import astro
from lsl.common import stations
from lsl.reader import base, drx, drx8, vdif
from lsl.common.dp import fS
from lsl.common.mcs import datetime_to_mjdmpm, delay_to_mcsd, mcsd_to_delay
from lsl.common.metabundle import get_command_script
from lsl.common.metabundleADP import get_command_script as get_command_scriptADP
from lsl.misc.beamformer import calc_delay


__version__ = '1.2'
__all__ = ['get_numa_node_count', 'get_numa_support', 'get_gpu_count',
           'get_gpu_support', 'InterProcessLock', 'EnhancedFixedBody',
           'EnhancedSun', 'EnhancedJupiter', 'multi_column_print',
           'best_freq_units', 'parse_time_string', 'nsround',
           'read_correlator_configuration', 'get_better_time', 'PolyCos']


# List of bright radio sources and pulsars in PyEphem format
_srcs = ["ForA,f|J,03:22:41.70,-37:12:30.0,1",
         "TauA,f|J,05:34:32.00,+22:00:52.0,1", 
         "VirA,f|J,12:30:49.40,+12:23:28.0,1",
         "HerA,f|J,16:51:08.15,+04:59:33.3,1", 
         "SgrA,f|J,17:45:40.00,-29:00:28.0,1", 
         "CygA,f|J,19:59:28.30,+40:44:02.0,1", 
         "CasA,f|J,23:23:27.94,+58:48:42.4,1",
         "3C196,f|J,08:13:36.06,+48:13:02.6,1",
         "3C286,f|J,13:31:08.29,+30:30:33.0,1",
         "3C295,f|J,14:11:20.47,+52:12:09.5,1", ]


def get_numa_node_count():
    # Query lscpu
    lscpu = subprocess.Popen(['lscpu',], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    output, _ = lscpu.communicate()
    output = output.decode()
    output = output.split('\n')
    
    # Look for the number of NUMA nodes
    nn = 1
    for line in output:
        if line.find('NUMA node(s)') != -1:
            nn = int(line.rsplit(':', 1)[1], 10)
    return nn


def get_numa_support():
    # Check the number of NUMA nodes
    nn = get_numa_node_count()
            
    # Check for the numactl utility
    numactl = subprocess.call(['which', 'numactl'], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    
    # If we have more than one NUMA node and numactl we are good to go
    status = False
    if numactl == 0 and nn > 1:
        status = True
        
    return status


def get_gpu_count():
    # Query nvidia-smi
    smi = subprocess.Popen(['nvidia-smi', '-q'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    output, _ = smi.communicate()
    output = output.decode()
    output = output.split('\n')
    
    # Look for the number of NUMA nodes
    gpus = 0
    for line in output:
        if line.find('Attached GPUs') != -1:
            gpus = int(line.rsplit(':', 1)[1], 10)
    return gpus


def get_gpu_support():
    # Check the number of GPUs
    gpus = get_gpu_count()
    
    # Check for the cupy module
    has_cupy = 1
    try:
        import cupy
    except ImportError:
        has_cupy = 0
        
    # If we have at least one GPU and cupy is installed we are ready
    status = False
    if has_cupy == 1 and gpus >= 1:
        status = True
        
    return status


class InterProcessLock(object):
    def __init__(self, name):
        self.name = name
        self.fh = open(f"{self.name}.lock", 'w+')
        self.locked = False
        
    def __del__(self):
        self.unlock()
        self.fh.close()
        
    def __enter__(self):
        self.lock()
        
    def __exit__(self, type, value, tb):
        self.unlock()
        
    def lock(self, block=True):	
        while not self.locked:
            try:
                fcntl.flock(self.fh, fcntl.LOCK_EX|fcntl.LOCK_NB)
                self.locked = True
            except IOError as e:
                if e.errno != errno.EAGAIN:
                    raise
                if not block:
                    break
                time.sleep(0.01)
                
        return self.locked
        
    def unlock(self):
        if not self.locked:
            return False
            
        fcntl.flock(self.fh, fcntl.LOCK_UN)
        self.locked = False
        return True


class EnhancedFixedBody(ephem.FixedBody):
    """
    Sub-class of ephem.FixedBody that allows for pulsar phase and frequency 
    calculation.  This is done through a new '_polycos' attribute that is set
    after initilization to a PolyCos instance.  On calling the 'compute' 
    method, the 'phase' and 'frequency' attributes are set.
    """
    
    def __init__(self, body=None):
        ephem.FixedBody.__init__(self)
        
        if type(body) is ephem.FixedBody:
            for attr in ('name', '_ra', '_dec', '_epoch', '_pa', '_pmra', '_pmdec'):
                value = getattr(body, attr, None)
                if value is not None:
                    setattr(self, attr, value)
                    
    def __setattr__(self, name, value):
        # Validate that the _polycos attribute is set to a PolyCos instance
        if name == '_polycos':
            if type(value) != PolyCos:
                raise ValueError("Must set _polycos with a PolyCos instance")
                
        # Set the attribute if everything is ok
        ephem.FixedBody.__setattr__(self, name, value)
        
    def __getattr__(self, name):
        if name in ('phase', 'frequency', 'period'):
            ## Is this even valid?
            if getattr(self, '_polycos', None) is None:
                raise ValueError("Pulsar parameters cannot be determined for a non-pulsar body")
                
        # Get the attribute if every is ok
        ephem.FixedBody.__getattr__(self, name)
        
    def compute(self, when=None, epoch=ephem.J2000):
        # Basic validation
        if when is None:
            when = ephem.now()
            
        # Compute the basic source parameters via PyEphem
        if type(when) is ephem.Observer:
            ephem.FixedBody.compute(self, when)
        else:
            ephem.FixedBody.compute(self, when, epoch=epoch)
            
        # Compute the pulsar parameters - if applicable
        self.compute_pulsar(when)
        
    def compute_pulsar(self, mjd, mjdf=0.0):
        """
        Compute the pulsar paramaters (if avaliable) with higher precision.
        """
        
        if getattr(self, '_polycos', None) is not None:
            if type(mjd) is ephem.Observer:
                mjd = mjd.date + (astro.DJD_OFFSET - astro.MJD_OFFSET)
                mjdf = 0.0
            elif type(mjd) is ephem.Date:
                mjd = mjd + (astro.DJD_OFFSET - astro.MJD_OFFSET)
                mjdf = 0.0
                
            self.dm = self._polycos.getDM(mjd, mjdf)
            self.phase = self._polycos.getPhase(mjd, mjdf)
            self.frequency = self._polycos.getFrequency(mjd, mjdf)
            self.doppler = self._polycos.getDoppler(mjd, mjdf)
            self.period = 1.0/self.frequency


class EnhancedSun(ephem.Sun):
    """
    Minimal sub-class of ephem.Sun to allow the 'duration' attribute to 
    be set.
    """
    
    def __getattr__(self, name):
        # Catch the _ra, _dec, _epoch, etc. attributes since they don't exist
        # for ephem.Planet sub-classes and we don't want to break things
        if name in ('_ra', '_dec', '_epoch', '_pa', '_pmra', '_pmdec'):
            return 'moving'
            
        # Get the attribute if every is ok
        ephem.Sun.__getattr__(self, name)
        
    def __setattr__(self, name, value):
        # Catch the _ra, _dec, _epoch, etc. attributes since they don't exist
        # for ephem.Planet sub-classes and we don't want to break things
        if name in ('_ra', '_dec', '_epoch', '_pa', '_pmra', '_pmdec'):
            raise AttributeError("Cannot set '%s' on this object" % value)
            
        # Set the attribute if everything is ok
        ephem.Sun.__setattr__(self, name, value)


class EnhancedJupiter(ephem.Jupiter):
    """
    Minimal sub-class of ephem.Jupiter to allow the 'duration' attribute to 
    be set.
    """
    
    def __getattr__(self, name):
        # Catch the _ra, _dec, _epoch, etc. attributes since they don't exist
        # for ephem.Planet sub-classes and we don't want to break things
        if name in ('_ra', '_dec', '_epoch', '_pa', '_pmra', '_pmdec'):
            return 'moving'
            
        # Get the attribute if every is ok
        ephem.Jupiter.__getattr__(self, name)
        
    def __setattr__(self, name, value):
        # Catch the _ra, _dec, _epoch, etc. attributes since they don't exist
        # for ephem.Planet sub-classes and we don't want to break things
        if name in ('_ra', '_dec', '_epoch', '_pa', '_pmra', '_pmdec'):
            raise AttributeError("Cannot set '%s' on this object" % value)
            
        # Set the attribute if everything is ok
        ephem.Jupiter.__setattr__(self, name, value)


def multi_column_print(items, sep=';  ', width=86):
    """
    Multi-column print statement for lists.
    """
    
    # Find the longest item in the list so that we can make aligned columns.
    maxLen = 0
    for i in items:
        l = len(str(i))
        if l > maxLen:
            maxLen = l
    formatter = "%%%is" % maxLen
    
    # Find out how many columns to make using the width the print over, the 
    # maximum item size, and the size of the separator.  When doing this also
    # make sure that we have at least one column
    nCol = width // (maxLen+len(sep))
    if nCol < 1:
        nCol = 1
        
    # Figure out how many rows to use.  This needs to take into acount partial
    # rows with len(items) % nCol != 0.
    nRow = len(items) // nCol + (0 if (len(items) % nCol) == 0 else 1)
    
    # Print
    for r in range(nRow):
        ## Build up the line
        out = sep.join([formatter % str(i) for i in items[r*nCol:(r+1)*nCol]])
        ## Add the separator at the end if this isn't the last line
        if r != nRow-1:
            out += sep
        ## Print
        print(out)


def best_freq_units(freq):
    """Given a numpy array of frequencies in Hz, return a new array with the
    frequencies in the best units possible (kHz, MHz, etc.)."""
    
    # Figure out how large the data are
    try:
        scale = int(math.log10(max(freq)))
    except TypeError:
        scale = int(math.log10(freq))
    if scale >= 9:
        divis = 1e9
        units = 'GHz'
    elif scale >= 6:
        divis = 1e6
        units = 'MHz'
    elif scale >= 3:
        divis = 1e3
        units = 'kHz'
    elif scale >= 0:
        divis = 1
        units = 'Hz'
    elif scale >= -3:
        divis = 1e-3
        units = 'mHz'
    else:
        divis = 1e-6
        units = 'uHz'
        
    # Convert the frequency
    newFreq = freq / divis
    
    # Return units and freq
    return (newFreq, units)


_timeRE = re.compile('^[ \t]*(?P<value>[+-]?\d*\.?\d*([Ee][+-]?\d*)?)[ \t]*(?P<unit>(([kmun]?s)|h|m))?[ \t]*$')
def parse_time_string(value):
    """
    Given a time in the format of "decimal_value unit", convert the string 
    to a floating point time in seconds.  Valid units are:
      * h  - hours
      * ks - kiloseconds
      * m  - minutes
      * s  - seconds
      * ms - milliseconds
      * us - microseconds
      * ns - nanoseconds
    
    If no units are provided, the value is assumed to be in seconds.
    """
    
    try:
        value = float(value)
    except ValueError:
        mtch = _timeRE.match(value)
        if mtch is None:
            raise ValueError(f"Invalid literal for parse_time_string(): {value}")
        value = float(mtch.group('value'))
        unit = mtch.group('unit')
        if unit is not None:
            if unit == 'h':
                value *= 3600.0
            elif unit == 'ks':
                value *= 1e3
            elif unit == 'm':
                value *= 60.0
            elif unit == 'ms':
                value *= 1e-3
            elif unit == 'us':
                value *= 1e-6
            elif unit == 'ns':
                value *= 1e-9
                
    return value


def nsround(value):
    """
    Round a time in seconds to the nearest ns.
    """
    
    return round(value*1e9)/1e9


def _read_correlator_configuration(filename):
    """
    Backend function for read_correlator_configuration.
    """
    
    context = None
    config = {}
    sources = []
    blocks = []
    
    fh = open(filename, 'r')
    for line in fh:
        if line[0] == '#':
            continue
        if len(line) < 3:
            continue
            
        line = line.strip().rstrip()
        
        if line == 'Context':
            temp_context = {'observer':'Unknown', 'project':'Unknown', 'session':None, 'vlaref':None}
        elif line[:8] == 'Observer':
            temp_context['observer'] = line.split(None, 1)[1]
        elif line[:7] == 'Project':
            temp_context['project'] = line.split(None, 1)[1]
        elif line[:7] == 'Session':
            temp_context['session'] = line.split(None, 1)[1]
        elif line[:6] == 'VLARef':
            temp_context['vlaref'] = line.split(None, 1)[1]
        elif line == 'EndContext':
            context = temp_context
            
        elif line == 'Configuration':
            temp_config = {'inttime':None, 'channels':None, 'basis':None}
        elif line[:8] == 'Channels':
            temp_config['channels'] = int(line.split(None, 1)[1], 10)
        elif line[:7] == 'IntTime':
            temp_config['inttime'] = float(line.split(None, 1)[1])
        elif line[:8] == 'PolBasis':
            temp_config['basis'] = line.split(None, 1)[1]
        elif line == 'EndConfiguration':
            config = temp_config
            
        elif line == 'Source':
            source = {'intent':'target', 'duration':0.0}
        elif line[:4] == 'Name':
            source['name'] = line.split(None, 1)[1]
        elif line[:6] == 'Intent':
            source['intent'] = line.split(None, 1)[1].lower()
        elif line[:6] == 'RA2000':
            source['ra'] = line.split(None, 1)[1]
        elif line[:7] == 'Dec2000':
            source['dec'] = line.split(None, 1)[1]
        elif line[:6] == 'Polyco':
            source['polyco'] = line.split(None, 1)[1]
        elif line[:8] == 'Duration':
            source['duration'] = float(line.split(None, 1)[1])
        elif line == 'SourceDone':
            sources.append( source )
            
        elif line == 'Input':
            block = {'fileOffset':0.0}
        elif line[:4] == 'File' and line[4] != 'O':
            block['filename'] = line.split(None, 1)[1]
        elif line[:8] == 'MetaData':
            ## Optional
            block['metadata'] = line.split(None, 1)[1]
        elif line[:4] == 'Type':
            block['type'] = line.split(None, 1)[1].lower()
        elif line[:7] == 'Antenna':
            block['antenna'] = line.split(None, 1)[1]
        elif line[:4] == 'Pols':
            block['pols'] = [v.strip().rstrip() for v in line.split(None, 1)[1].split(',')]
        elif line[:8] == 'Location':
            block['location'] = [float(v) for v in line.split(None, 1)[1].split(',')]
        elif line[:16]  == 'ApparentLocation':
            block['appLocation'] = [float(v) for v in line.split(None, 1)[1].split(',')]
        elif line[:11] == 'ClockOffset':
            block['clockOffset'] = [parse_time_string(v) for v in line.split(None, 1)[1].split(',')]
        elif line[:10] == 'FileOffset':
            block['fileOffset'] = parse_time_string(line.split(None, 1)[1])
        elif line == 'InputDone':
            ## Make sure we have a metaData key since it is optional
            if 'metadata' not in block:
                block['metadata'] = None
            blocks.append( block )
    fh.close()
    
    # Set the context
    config['context'] = context
    
    # Find the reference source
    if 'ra' in sources[0].keys() and 'dec' in sources[0].keys():
        refSource = EnhancedFixedBody()
        refSource.name = sources[0]['name']
        refSource._ra = sources[0]['ra']
        refSource._dec = sources[0]['dec']
        refSource._epoch = ephem.J2000
    else:
        srcs = [EnhancedSun(), EnhancedJupiter()]
        for line in _srcs:
            srcs.append( EnhancedFixedBody(ephem.readdb(line)) )
            
        refSource = None
        for i in range(len(srcs)):
            if srcs[i].name == sources[0]['name']:
                refSource = srcs[i]
                break
        if refSource is None:
            raise ValueError(f"Unknown source '{sources[0]['name']}'")
    refSource.intent = sources[0]['intent']
    refSource.duration = sources[0]['duration']
    try:
        if not os.path.exists(sources[0]['polyco']):
            # Maybe it is relative to the configuration file's path?
            sources[0]['polyco'] = os.path.join(os.path.dirname(filename), sources[0]['polyco'])
        refSource._polycos = PolyCos(sources[0]['polyco'], psrname=refSource.name.replace('PSR', '').replace('_', '').replace(' ',''))
    except KeyError:
        pass
        
    # Sort everything out so that the VDIF files come first
    order = sorted(range(len(blocks)), key=lambda x: blocks[x]['type'][-1])
    blocks = [blocks[o] for o in order]

    # Build up a list of filenames	
    filenames = [block['filename'] for block in blocks]
    
    # Build up a list of metadata filenames
    metanames = [block['metadata'] for block in blocks]
    
    # Build up a list of file offsets
    offsets = [block['fileOffset'] for block in blocks]
    
    # Build up a list of readers
    readers = []
    for block in blocks:
        if block['type'] == 'vdif':
            readers.append( vdif )
        elif block['type'] == 'drx':
            readers.append( drx )
        elif block['type'] == 'drx8':
            readers.append( drx8 )
        else:
            readers.append( None )
            
    # Build up a list of antennas
    antennas = []
    i = 1
    for block in blocks:
        aid = None
        name = block['antenna']
        if name.lower() in ('lwa1', 'lwa-1'):
            aid = 51
        elif name.lower() in ('lwasv', 'lwa-sv'):
            aid = 52
        elif name.lower() in ('lwana', 'lwa-na'):
            aid = 53
        elif name.lower() in ('ovrolwa', 'ovro-lwa'):
            aid = 54
        else:
            for j in range(len(name)):
                try:
                    aid = int(name[j:], 10)
                    if name[:j].lower() == 'lwa':
                        aid += 50
                    break
                except ValueError:
                    pass
        pols = block['pols']
        location = block['location']
        try:
            app_location = block['appLocation']
        except KeyError:
            app_location = None
        clock_offsets = block['clockOffset']
        
        if aid is None:
            raise RuntimeError(f"Cannot convert antenna name '{name}' to a number")
            
        stand = stations.Stand(aid, *location)
        try:
            apparent_stand = stations.Stand(aid, *app_location)
        except TypeError:
            apparent_stand = None
        for pol,offset in zip(pols, clock_offsets):
            cable = stations.Cable('%s-%s' % (name, pol), 0.0, vf=1.0, dd=0.0)
            cable.clock_offset = offset
            
            if pol.lower() == 'x':
                antenna = stations.Antenna(i, stand=stand, cable=cable, pol=0)
            else:
                antenna = stations.Antenna(i, stand=stand, cable=cable, pol=1)
            antenna.apparent_stand = apparent_stand
            antenna.config_name = name.upper().replace('-', '')
            
            antennas.append( antenna )
            i += 1
            
    # Done
    return config, refSource, filenames, metanames, offsets, readers, antennas


def read_correlator_configuration(filename_or_npz):
    """
    Parse a correlator configuration file generated by createConfigFile.py and
    return a seven-element tuple of:
      * a correlator configuration dictionary or None if no configuraiton is found
      * the reference source as a ephem.FixedBody-compatible instance
      * a list of filenames, 
      * a list of metadata tarball names, 
      * a list of file offsets in seconds, 
      * a list of lsl.reader modules to use, and
      * a list of lsl.common.stations.Antenna instances for each file
    """
    
    # Sort out what to do depending on what we were given
    if isinstance(filename_or_npz, np.lib.npyio.NpzFile):
        ## An open .npz file, just work with it
        dataDict = filename_or_npz
        to_close = False
    elif os.path.splitext(filename_or_npz)[1] == '.npz':
        ## A .npz file, open and and then work with it
        dataDict = np.load(filename_or_npz)
        to_close = True
    else:
        ## Something else, just try some stuff
        try:
            ### Could it be a file that has already been opened?
            filename_or_npz = filename_or_npz.filename
        except AttributeError:
            pass
        return _read_correlator_configuration(filename_or_npz)
        
    # Make a temporary directory
    tempdir = tempfile.mkdtemp(prefix='config-')
    
    try:
        # Configuration file - oh, and check for a 'Polyco' definition
        cConfig = dataDict['config']
        tempConfig = os.path.join(tempdir, 'config.txt')
        fh = open(tempConfig, 'w')
        polycos = None
        for line in cConfig:
            try:
                line = line.decode()
            except AttributeError:
                pass
            fh.write(line)
            if line.find('Polyco') != -1:
                polycos = line.strip().rstrip()
        fh.close()
        
        # Polycos (optional)
        if polycos is not None:
            tempPolycos = os.path.basename(polycos.split(None, 1)[1])
            tempPolycos = os.path.join(tempdir, tempPolycos)
            try:
                cPolycos = dataDict['polycos']
                fh = open(tempPolycos, 'wb')
                for line in cPolycos:
                    fh.write(line)
                fh.close()
            except KeyError:
                pass
                
        # Read
        full_config = _read_correlator_configuration(tempConfig)
        
    finally:
        # Cleanup
        shutil.rmtree(tempdir)
        if to_close:
            dataDict.close()
            
    # Done
    return full_config


def get_better_time(frame):
    """
    Given a lsl.reader.vdif.Frame, lsl.reader.drx.Frame instance, or
    lsl.reader.drx8.Frame, return a more accurate time for the frame.  Unlike the
    Frame.get_time() functions, this function returns a two-element tuple of:
      * integer seconds since the UNIX epoch and
      * fractional second
    """
    
    if isinstance(frame, (drx.Frame, drx8.Frame)):
        # HACK - What should T_NOM really be at LWA1???
        tt = frame.payload.timetag
        to = 6660 if frame.header.time_offset else 0
        return list(base.FrameTimestamp.from_dp_timetag(tt, to))
    else:
        return list(frame.time)


class PolyCos(object):
    """
    Class for working with pulsar PolyCos files.
    """
    
    def __init__(self, filename, psrname=None):
        if psrname is None:
            psrname = os.path.basename(filename)
            psrname = psrname.split('_', 1)[0]
            
        from mini_presto.polycos import polycos
        
        self.filename = filename
        self._polycos_base = polycos(psrname, filename)
        
    def getDM(self, mjd, mjdf=0.0):
        """
        Given a MJD value, return the dispersion measure of the pulsar in 
        pc cm^{-3}.
        """
        
        goodpoly = self._polycos_base.select_polyco(mjd, mjdf)
        return self._polycos_base.polycos[goodpoly].DM
        
    def getPhase(self, mjd, mjdf=0.0):
        """
        Given a MJD value, compute the phase of the pulsar.
        """
        
        return self._polycos_base.get_phase(mjd, mjdf)
        
    def getFrequency(self, mjd, mjdf=0.0):
        """
        Given a MJD value, compute the frequency of the pulsar in Hz.
        """
        
        return self._polycos_base.get_freq(mjd, mjdf)
        
    def getDoppler(self, mjd, mjdf=0.0):
        """
        Given a MJD value, return the approximate topocentric Doppler shift of the 
        pulsar (~1 + vObs/c).
        """
        
        return 1.0 + self._polycos_base.get_voverc(mjd, mjdf)