helper_functions.py 11.2 KB
Newer Older
Gregory Ashton's avatar
Gregory Ashton committed
1 2 3 4 5 6
"""
Provides helpful functions to facilitate ease-of-use of pyfstat
"""

import os
import sys
7
import subprocess
Gregory Ashton's avatar
Gregory Ashton committed
8 9 10
import argparse
import logging
import inspect
11
import peakutils
Gregory Ashton's avatar
Gregory Ashton committed
12
from functools import wraps
13
from scipy.stats.distributions import ncx2
14
import numpy as np
15
import lal
16
import lalpulsar
Gregory Ashton's avatar
Gregory Ashton committed
17

18 19 20 21 22 23 24 25 26
# workaround for matplotlib on X-less remote logins
if 'DISPLAY' in os.environ:
    import matplotlib.pyplot as plt
else:
    logging.info('No $DISPLAY environment variable found, so importing \
                  matplotlib.pyplot with non-interactive "Agg" backend.')
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
Gregory Ashton's avatar
Gregory Ashton committed
27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44

def set_up_optional_tqdm():
    try:
        from tqdm import tqdm
    except ImportError:
        def tqdm(x, *args, **kwargs):
            return x
    return tqdm


def set_up_matplotlib_defaults():
    plt.switch_backend('Agg')
    plt.rcParams['text.usetex'] = True
    plt.rcParams['axes.formatter.useoffset'] = False


def set_up_command_line_arguments():
    parser = argparse.ArgumentParser()
45 46 47
    parser.add_argument("-v", "--verbose", action="store_true",
                        help="Increase output verbosity [logging.DEBUG]")
    parser.add_argument("-q", "--quite", action="store_true",
48
                        help="Decrease output verbosity [logging.WARNING]")
Gregory Ashton's avatar
Gregory Ashton committed
49 50
    parser.add_argument("--no-interactive", help="Don't use interactive",
                        action="store_true")
51 52 53 54 55 56 57 58 59
    parser.add_argument("-c", "--clean", action="store_true",
                        help="Force clean data, never use cached data")
    fu_parser = parser.add_argument_group(
        'follow-up options', 'Options related to MCMCFollowUpSearch')
    fu_parser.add_argument('-s', "--setup-only", action="store_true",
                           help="Only generate the setup file, don't run")
    fu_parser.add_argument(
        "--no-template-counting", action="store_true",
        help="No counting of templates, useful if the setup is predefined")
60 61 62
    parser.add_argument(
        '-N', type=int, default=3, metavar='N',
        help="Number of threads to use when running in parallel")
Gregory Ashton's avatar
Gregory Ashton committed
63 64 65
    parser.add_argument('unittest_args', nargs='*')
    args, unknown = parser.parse_known_args()
    sys.argv[1:] = args.unittest_args
66

Gregory Ashton's avatar
Gregory Ashton committed
67 68 69
    if args.quite or args.no_interactive:
        def tqdm(x, *args, **kwargs):
            return x
70 71
    else:
        tqdm = set_up_optional_tqdm()
72

Gregory Ashton's avatar
Gregory Ashton committed
73 74
    logger = logging.getLogger()
    stream_handler = logging.StreamHandler()
75 76 77
    stream_handler.setFormatter(logging.Formatter(
        '%(asctime)s %(levelname)-8s: %(message)s', datefmt='%H:%M'))

Gregory Ashton's avatar
Gregory Ashton committed
78
    if args.quite:
79
        logger.setLevel(logging.WARNING)
Gregory Ashton's avatar
Gregory Ashton committed
80
        stream_handler.setLevel(logging.WARNING)
81
    elif args.verbose:
82
        logger.setLevel(logging.DEBUG)
Gregory Ashton's avatar
Gregory Ashton committed
83
        stream_handler.setLevel(logging.DEBUG)
84
    else:
85
        logger.setLevel(logging.INFO)
86
        stream_handler.setLevel(logging.INFO)
87

Gregory Ashton's avatar
Gregory Ashton committed
88
    logger.addHandler(stream_handler)
89
    return args, tqdm
Gregory Ashton's avatar
Gregory Ashton committed
90 91


92
def get_ephemeris_files():
93
    """ Returns the earth_ephem and sun_ephem """
Gregory Ashton's avatar
Gregory Ashton committed
94
    config_file = os.path.expanduser('~')+'/.pyfstat.conf'
95 96
    env_var = 'LALPULSAR_DATADIR'
    please = 'Please provide the ephemerides paths when initialising searches.'
Gregory Ashton's avatar
Gregory Ashton committed
97 98 99 100 101 102 103 104 105
    if os.path.isfile(config_file):
        d = {}
        with open(config_file, 'r') as f:
            for line in f:
                k, v = line.split('=')
                k = k.replace(' ', '')
                for item in [' ', "'", '"', '\n']:
                    v = v.replace(item, '')
                d[k] = v
106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
        try:
            earth_ephem = d['earth_ephem']
            sun_ephem = d['sun_ephem']
        except:
            logging.warning('No [earth/sun]_ephem found in '+config_file+'. '+please)
            earth_ephem = None
            sun_ephem = None
    elif env_var in os.environ.keys():
        earth_ephem = os.path.join(os.environ[env_var],'earth00-40-DE421.dat.gz')
        sun_ephem = os.path.join(os.environ[env_var],'sun00-40-DE421.dat.gz')
        if not ( os.path.isfile(earth_ephem) and os.path.isfile(sun_ephem) ):
            earth_ephem = os.path.join(os.environ[env_var],'earth00-19-DE421.dat.gz')
            sun_ephem = os.path.join(os.environ[env_var],'sun00-19-DE421.dat.gz')
            if not ( os.path.isfile(earth_ephem) and os.path.isfile(sun_ephem) ):
                logging.warning('No [earth/sun]00-[19/40]-DE421 ephemerides '
                                'found in the '+os.environ[env_var]+' directory. '+please)
                earth_ephem = None
                sun_ephem = None
Gregory Ashton's avatar
Gregory Ashton committed
124
    else:
125 126
        logging.warning('No '+config_file+' file or $'+env_var+' environment '
                        'variable found. '+please)
Gregory Ashton's avatar
Gregory Ashton committed
127 128 129 130 131 132 133 134 135 136 137 138 139 140
        earth_ephem = None
        sun_ephem = None
    return earth_ephem, sun_ephem


def round_to_n(x, n):
    if not x:
        return 0
    power = -int(np.floor(np.log10(abs(x)))) + (n - 1)
    factor = (10 ** power)
    return round(x * factor) / factor


def texify_float(x, d=2):
Gregory Ashton's avatar
Gregory Ashton committed
141 142
    if x == 0:
        return 0
Gregory Ashton's avatar
Gregory Ashton committed
143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172
    if type(x) == str:
        return x
    x = round_to_n(x, d)
    if 0.01 < abs(x) < 100:
        return str(x)
    else:
        power = int(np.floor(np.log10(abs(x))))
        stem = np.round(x / 10**power, d)
        if d == 1:
            stem = int(stem)
        return r'${}{{\times}}10^{{{}}}$'.format(stem, power)


def initializer(func):
    """ Decorator function to automatically assign the parameters to self """
    names, varargs, keywords, defaults = inspect.getargspec(func)

    @wraps(func)
    def wrapper(self, *args, **kargs):
        for name, arg in list(zip(names[1:], args)) + list(kargs.items()):
            setattr(self, name, arg)

        for name, default in zip(reversed(names), reversed(defaults)):
            if not hasattr(self, name):
                setattr(self, name, default)

        func(self, *args, **kargs)

    return wrapper

173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195

def get_peak_values(frequencies, twoF, threshold_2F, F0=None, F0range=None):
    if F0:
        cut_idxs = np.abs(frequencies - F0) < F0range
        frequencies = frequencies[cut_idxs]
        twoF = twoF[cut_idxs]
    idxs = peakutils.indexes(twoF, thres=1.*threshold_2F/np.max(twoF))
    F0maxs = frequencies[idxs]
    twoFmaxs = twoF[idxs]
    freq_err = frequencies[1] - frequencies[0]
    return F0maxs, twoFmaxs, freq_err*np.ones(len(idxs))


def get_comb_values(F0, frequencies, twoF, period, N=4):
    if period == 'sidereal':
        period = 23*60*60 + 56*60 + 4.0616
    elif period == 'terrestrial':
        period = 86400
    freq_err = frequencies[1] - frequencies[0]
    comb_frequencies = [n*1/period for n in range(-N, N+1)]
    comb_idxs = [np.argmin(np.abs(frequencies-F0-F)) for F in comb_frequencies]
    return comb_frequencies, twoF[comb_idxs], freq_err*np.ones(len(comb_idxs))

196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226

def compute_P_twoFstarcheck(twoFstarcheck, twoFcheck, M0, plot=False):
    """ Returns the unnormalised pdf of twoFstarcheck given twoFcheck """
    upper = 4+twoFstarcheck + 0.5*(2*(4*M0+2*twoFcheck))
    rho2starcheck = np.linspace(1e-1, upper, 500)
    integrand = (ncx2.pdf(twoFstarcheck, 4*M0, rho2starcheck)
                 * ncx2.pdf(twoFcheck, 4, rho2starcheck))
    if plot:
        fig, ax = plt.subplots()
        ax.plot(rho2starcheck, integrand)
        fig.savefig('test')
    return np.trapz(integrand, rho2starcheck)


def compute_pstar(twoFcheck_obs, twoFstarcheck_obs, m0, plot=False):
    M0 = 2*m0 + 1
    upper = 4+twoFcheck_obs + (2*(4*M0+2*twoFcheck_obs))
    twoFstarcheck_vals = np.linspace(1e-1, upper, 500)
    P_twoFstarcheck = np.array(
        [compute_P_twoFstarcheck(twoFstarcheck, twoFcheck_obs, M0)
         for twoFstarcheck in twoFstarcheck_vals])
    C = np.trapz(P_twoFstarcheck, twoFstarcheck_vals)
    idx = np.argmin(np.abs(twoFstarcheck_vals - twoFstarcheck_obs))
    if plot:
        fig, ax = plt.subplots()
        ax.plot(twoFstarcheck_vals, P_twoFstarcheck)
        ax.fill_between(twoFstarcheck_vals[:idx+1], 0, P_twoFstarcheck[:idx+1])
        ax.axvline(twoFstarcheck_vals[idx])
        fig.savefig('test')
    pstar_l = np.trapz(P_twoFstarcheck[:idx+1]/C, twoFstarcheck_vals[:idx+1])
    return 2*np.min([pstar_l, 1-pstar_l])
227 228


229
def run_commandline(cl, log_level=20, raise_error=True, return_output=True):
230
    """Run a string cmd as a subprocess, check for errors and return output.
231

232 233 234 235 236 237 238 239 240 241 242
    Parameters
    ----------
    cl: str
        Command to run
    log_level: int
        See https://docs.python.org/2/library/logging.html#logging-levels,
        default is '20' (INFO)

    """

    logging.log(log_level, 'Now executing: ' + cl)
243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261
    if return_output:
        try:
            out = subprocess.check_output(cl,                       # what to run
                                          stderr=subprocess.STDOUT, # catch errors
                                          shell=True,               # proper environment etc
                                          universal_newlines=True,  # properly display linebreaks in error/output printing
                                         )
        except subprocess.CalledProcessError as e:
            logging.log(log_level, 'Execution failed: {}'.format(e.output))
            if raise_error:
                raise
            else:
                out = 0
        os.system('\n')
        return(out)
    else:
        process = subprocess.Popen(cl, shell=True)
        process.communicate()

262

263

264
def convert_array_to_gsl_matrix(array):
265
    gsl_matrix = lal.gsl_matrix(*array.shape)
266 267
    gsl_matrix.data = array
    return gsl_matrix
268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285


def get_sft_array(sftfilepattern, data_duration, F0, dF0):
    """ Return the raw data from a set of sfts """

    SFTCatalog = lalpulsar.SFTdataFind(
        sftfilepattern, lalpulsar.SFTConstraints())
    MultiSFTs = lalpulsar.LoadMultiSFTs(SFTCatalog, F0-dF0, F0+dF0)
    SFTs = MultiSFTs.data[0]
    data = []
    for sft in SFTs.data:
        data.append(np.abs(sft.data.data))
    data = np.array(data).T
    n, nsfts = data.shape
    freqs = np.linspace(sft.f0, sft.f0+n*sft.deltaF, n)
    times = np.linspace(0, data_duration, nsfts)

    return times, freqs, data
286 287


288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309
def get_covering_band(tref, tstart, tend, F0, F1, F2):
    """ Get the covering band using XLALCWSignalCoveringBand

    Parameters
    ----------
    tref, tstart, tend: int
        The reference, start, and end times of interest
    F0, F1, F1:
        Frequency and spin-down of the signal

    Note: this is similar to the function
    `injection_helper_functions.get_frequency_range_of_signal`, however this
    does not use the sky position and calculates an estimate for a full year
    search over any sky position. In this sense, it is much more conservative.

    Returns
    -------
    F0min, F0max: float
        Estimates of the minimum and maximum frequencies of the signal during
        the search

    """
310 311 312 313
    tref = lal.LIGOTimeGPS(tref)
    tstart = lal.LIGOTimeGPS(tstart)
    tend = lal.LIGOTimeGPS(tend)
    psr = lalpulsar.PulsarSpinRange()
314 315 316
    psr.fkdot[0] = F0
    psr.fkdot[1] = F1
    psr.fkdot[2] = F2
317 318 319 320
    psr.refTime = tref
    return lalpulsar.CWSignalCoveringBand(tstart, tend, psr, 0, 0, 0)


Gregory Ashton's avatar
Gregory Ashton committed
321 322 323 324 325 326 327
def twoFDMoffThreshold(twoFon, knee=400, twoFDMoffthreshold_below_threshold=62,
                       prefactor=0.9, offset=0.5):
    """ Calculation of the 2F_DMoff threshold, see Eq 2 of arXiv:1707.5286 """
    if twoFon <= knee:
        return twoFDMoffthreshold_below_threshold
    else:
        return 10**(prefactor*np.log10(twoFon-offset))