Source code for maltpynt.plot

# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Quicklook plots."""
from __future__ import (absolute_import, unicode_literals, division,
                        print_function)

try:
    import matplotlib.pyplot as plt
except:
    # Permit to import the module anyway if matplotlib is not present
    pass
from .io import load_data, get_file_type, load_pds
from .io import is_string
from .base import create_gti_mask, _assign_value_if_none
from .base import detection_level
from .fspec import rms_normalize_pds
import logging
import numpy as np
from .io import MP_FILE_EXTENSION


def _next_color(ax):
    try:
        return next(ax._get_lines.color_cycle)
    except:
        return next(ax._get_lines.prop_cycler)['color']


def _baseline_fun(x, a):
    """A constant function."""
    return a


def _value_or_none(dict_like, key):
    try:
        return dict_like[key]
    except:
        return None


[docs]def plot_generic(fnames, vars, errs=None, figname=None, xlog=None, ylog=None): """Generic plotting function.""" if is_string(fnames): fnames = [fnames] figname = _assign_value_if_none(figname, '{0} vs {1}'.format(vars[1], vars[0])) plt.figure(figname) ax = plt.gca() if xlog: ax.set_xscale('log', nonposx='clip') if ylog: ax.set_yscale('log', nonposy='clip') xlabel, ylabel = vars xlabel_err, ylabel_err = None, None if errs is not None: xlabel_err, ylabel_err = errs for i, fname in enumerate(fnames): data = load_data(fname) color = _next_color(ax) xdata = data[xlabel] ydata = data[ylabel] xdata_err = _value_or_none(data, xlabel_err) ydata_err = _value_or_none(data, ylabel_err) plt.errorbar(xdata, ydata, yerr=ydata_err, xerr=xdata_err, fmt='-', drawstyle='steps-mid', color=color, label=fname) plt.xlabel(vars[0]) plt.ylabel(vars[1]) plt.legend()
[docs]def plot_pds(fnames, figname=None, xlog=None, ylog=None): """Plot a list of PDSs, or a single one.""" from scipy.optimize import curve_fit import collections if is_string(fnames): fnames = [fnames] figlabel = fnames[0] for i, fname in enumerate(fnames): pdsdata = load_pds(fname) try: freq = pdsdata['freq'] xlog = _assign_value_if_none(xlog, False) ylog = _assign_value_if_none(ylog, False) except: flo = pdsdata['flo'] fhi = pdsdata['fhi'] freq = (fhi + flo) / 2 xlog = _assign_value_if_none(xlog, True) ylog = _assign_value_if_none(ylog, True) pds = pdsdata['pds'] epds = pdsdata['epds'] npds = pdsdata['npds'] norm = pdsdata['norm'] rebin = pdsdata['rebin'] ctrate = pdsdata['ctrate'] back_ctrate = pdsdata['back_ctrate'] nbin = len(pds[1:]) lev = detection_level(nbin, n_summed_spectra=npds, n_rebin=rebin) if norm == "rms": lev, _ = rms_normalize_pds(lev, 0, source_ctrate=ctrate, back_ctrate=back_ctrate) p, pcov = curve_fit(_baseline_fun, freq, pds, p0=[2], sigma=epds) logging.info('White noise level is {0}'.format(p[0])) pds -= p[0] if xlog and ylog: plt.figure('PDS - Loglog ' + figlabel) else: plt.figure('PDS ' + figlabel) ax = plt.gca() color = _next_color(ax) if xlog: ax.set_xscale('log', nonposx='clip') if ylog: ax.set_yscale('log', nonposy='clip') level = lev - p[0] if norm == 'Leahy' or (norm == 'rms' and (not xlog or not ylog)): plt.errorbar(freq[1:], pds[1:], yerr=epds[1:], fmt='-', drawstyle='steps-mid', color=color, label=fname) elif norm == 'rms' and xlog and ylog: plt.errorbar(freq[1:], pds[1:] * freq[1:], yerr=epds[1:] * freq[1:], fmt='-', drawstyle='steps-mid', color=color, label=fname) level *= freq if np.any(level < 0): continue if isinstance(level, collections.Iterable): plt.plot(freq, level, ls='--', color=color) else: plt.axhline(level, ls='--', color=color) plt.xlabel('Frequency') if norm == 'rms': plt.ylabel('(rms/mean)^2') elif norm == 'Leahy': plt.ylabel('Leahy power') plt.legend() if figname is not None: plt.savefig(figname)
[docs]def plot_cospectrum(fnames, figname=None, xlog=None, ylog=None): """Plot the cospectra from a list of CPDSs, or a single one.""" if is_string(fnames): fnames = [fnames] figlabel = fnames[0] for fname in fnames: pdsdata = load_pds(fname) try: freq = pdsdata['freq'] xlog = _assign_value_if_none(xlog, False) ylog = _assign_value_if_none(ylog, False) except: flo = pdsdata['flo'] fhi = pdsdata['fhi'] freq = (fhi + flo) / 2 xlog = _assign_value_if_none(xlog, True) ylog = _assign_value_if_none(ylog, True) cpds = pdsdata['cpds'] cospectrum = cpds.real if xlog and ylog: plt.figure('Cospectrum - Loglog ' + figlabel) else: plt.figure('Cospectrum ' + figlabel) ax = plt.gca() if xlog: ax.set_xscale('log', nonposx='clip') if ylog: ax.set_yscale('log', nonposy='clip') plt.xlabel('Frequency') if xlog and ylog: plt.plot(freq[1:], freq[1:] * cospectrum[1:], drawstyle='steps-mid', label=fname) plt.ylabel('Cospectrum * Frequency') else: plt.plot(freq[1:], cospectrum[1:], drawstyle='steps-mid', label=fname) plt.ylabel('Cospectrum') plt.legend() if figname is not None: plt.savefig(figname)
[docs]def plot_lc(lcfiles, figname=None, fromstart=False, xlog=None, ylog=None): """Plot a list of light curve files, or a single one.""" if is_string(lcfiles): lcfiles = [lcfiles] figlabel = lcfiles[0] plt.figure('LC ' + figlabel) for lcfile in lcfiles: logging.info('Loading %s...' % lcfile) lcdata = load_data(lcfile) time = lcdata['time'] lc = lcdata['lc'] gti = lcdata['GTI'] instr = lcdata['Instr'] if fromstart: time -= lcdata['Tstart'] gti -= lcdata['Tstart'] if instr == 'PCA': # If RXTE, plot per PCU count rate npcus = lcdata['nPCUs'] lc /= npcus for g in gti: plt.axvline(g[0], ls='-', color='red') plt.axvline(g[1], ls='--', color='red') good = create_gti_mask(time, gti) plt.plot(time, lc, drawstyle='steps-mid', color='grey') plt.plot(time[good], lc[good], drawstyle='steps-mid', label=lcfile) plt.xlabel('Time (s)') if instr == 'PCA': plt.ylabel('light curve (Ct/bin/PCU)') else: plt.ylabel('light curve (Ct/bin)') plt.legend() if figname is not None: plt.savefig(figname)
[docs]def main(args=None): """Main function called by the `MPplot` command line script.""" import argparse description = \ 'Plot the content of MaLTPyNT light curves and frequency spectra' parser = argparse.ArgumentParser(description=description) parser.add_argument("files", help="List of files", nargs='+') parser.add_argument("--noplot", help="Only create images, do not plot", default=False, action='store_true') parser.add_argument("--figname", help="Figure name", default=None, type=str) parser.add_argument("--xlog", help="Use logarithmic X axis", default=None, action='store_true') parser.add_argument("--ylog", help="Use logarithmic Y axis", default=None, action='store_true') parser.add_argument("--xlin", help="Use linear X axis", default=None, action='store_true') parser.add_argument("--ylin", help="Use linear Y axis", default=None, action='store_true') parser.add_argument("--fromstart", help="Times are measured from the start of the " "observation (only relevant for light curves)", default=False, action='store_true') parser.add_argument("--axes", nargs=2, type=str, help="Plot two variables contained in the file", default=None) args = parser.parse_args(args) if args.noplot and args.figname is None: args.figname = args.files[0].replace(MP_FILE_EXTENSION, '.png') if args.xlin is not None: args.xlog = False if args.ylin is not None: args.ylog = False for fname in args.files: ftype, contents = get_file_type(fname) if args.axes is not None: plot_generic(fname, args.axes, xlog=args.xlog, ylog=args.ylog, figname=args.figname) continue if ftype == 'lc': plot_lc(fname, fromstart=args.fromstart, xlog=args.xlog, ylog=args.ylog, figname=args.figname) elif ftype[-4:] == 'cpds': plot_cospectrum(fname, xlog=args.xlog, ylog=args.ylog, figname=args.figname) elif ftype[-3:] == 'pds': plot_pds(fname, xlog=args.xlog, ylog=args.ylog, figname=args.figname) if not args.noplot: plt.show()