gwdetchar-lasso-correlation

#!/usr/bin/env python
# coding=utf-8
# Copyright (C) LIGO Scientific Collaboration (2015-)
#
# This file is part of the GW DetChar python package.
#
# GW DetChar is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# GW DetChar is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with GW DetChar.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import (division, print_function)

import os
import re
import multiprocessing
import sys
from subprocess import call
import tempfile
import atexit
import shutil

from math import (isnan, isinf, log, log10)
import numpy
from scipy.stats import spearmanr
from scipy.interpolate import UnivariateSpline
import astropy.units as u

from matplotlib import use
use('agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

from sklearn import linear_model
from sklearn.preprocessing import scale

from gwpy.table import Table
from pandas.tools.plotting import table as pandasTab
from pandas import DataFrame

from gwpy.timeseries import (TimeSeries, TimeSeriesDict)
from gwpy.time import (Time, from_gps)
from gwpy.plotter import TimeSeriesPlot
from gwpy.plotter import Plot as gwplot
from gwpy.detector import ChannelList
from gwpy.io import nds2 as io_nds2

from gwdetchar import cli
from gwdetchar.io import html

try:
    from LDAStools import frameCPP
except ImportError:
    io_kw = {}
else:
    io_kw = {'format': 'gwf.framecpp', 'type': 'adc'}


def find_outliers(ts, N):
    ts = ts.value  # strip out Quantity extras
    return numpy.nonzero(abs(ts - numpy.mean(ts)) > N*numpy.std(ts))[0]


def remove_outliers(ts, N):
    outliers = find_outliers(ts, N)
    c = 1
    if outliers.any():
        print("-- Found %d outliers in %s, recursively removing"
              % (len(outliers), ts.name))
        while outliers.any():
            cache = outliers
            mask = numpy.ones(len(ts), dtype=bool)
            mask[outliers] = False
            spline = UnivariateSpline(ts[mask].times.value, ts[mask].value,
                                      s=0, k=3)
            ts[outliers] = spline(ts[outliers].times.value)
            outliers = find_outliers(ts, N)
            print("Completed %d removal cycles" % c)
            if numpy.array_equal(outliers, cache):
                print("Outliers did not change, breaking recursion")
                break
            print("%d outliers remain" % len(outliers))
            c += 1


def configure_mpl():
    mpldir = tempfile.mkdtemp()
    atexit.register(shutil.rmtree, mpldir)
    umask = os.umask(0)
    os.umask(umask)
    os.chmod(mpldir, 0o777 & ~umask)
    os.environ['HOME'] = mpldir
    os.environ['MPLCONFIGDIR'] = mpldir

    import matplotlib
    matplotlib.use('agg')
    import matplotlib.pyplot as plt

    class TexManager(matplotlib.texmanager.TexManager):
        texcache = os.path.join(mpldir, 'tex.cache')

    matplotlib.texmanager.TexManager = TexManager
    matplotlib.rcParams['ps.useafm'] = True
    matplotlib.rcParams['pdf.use14corefonts'] = True
    matplotlib.rcParams['text.usetex'] = True


# -- parse command line -------------------------------------------------------

parser = cli.create_parser(
    description=__doc__,
    formatter_class=cli.argparse.ArgumentDefaultsHelpFormatter)
cli.add_gps_start_stop_arguments(parser)
cli.add_ifo_option(parser)
cli.add_nproc_option(parser, default=1)
parser.add_argument('-J', '--nproc-plot', type=int, default=None,
                    help='number of processes to use for plotting')
parser.add_argument('-o', '--output-dir', default=os.curdir,
                    help='output directory for plots')
parser.add_argument('-f', '--channel-file', type=os.path.abspath,
                    help='path for channel file')
parser.add_argument('-T', '--trend-type', default='minute',
                    choices=['second', 'minute'],
                    help='type of trend for correlation')
parser.add_argument('-p', '--primary-channel',
                    default='{ifo}:DMT-SNSH_EFFECTIVE_RANGE_MPC.mean',
                    help='name of primary channel to use')
parser.add_argument('-P', '--primary-frametype',
                    help='frametype for --primary-channel')
parser.add_argument('-O', '--remove-outliers', type=float, default=None,
                    help='Std. dev. limit for removing outliers')
parser.add_argument('-t', '--threshold', type=float, default=0.0001,
                    help='threshold for making a plot')

psig = parser.add_argument_group('Signal processing options')
psig.add_argument('-b', '--band-pass', type=float, nargs=2, default=None,
                  metavar="FLOW FHIGH",
                  help='lower and upper frequencies for bandpass on h(t)')
psig.add_argument('-x', '--filter-padding', type=float, default=3.,
                  help='amount of time (seconds) to pad data for filtering')

lsig = parser.add_argument_group('LASSO options')
lsig.add_argument('-a', '--alpha', default=None, type=float,
                  help='alpha parameter for LASSO fit')
lsig.add_argument('-C', '--no-cluster', action='store_true', default=False,
                  help='do not generate clustered channel plots')
lsig.add_argument('-c', '--cluster-coefficient', default=.85, type=float,
                  help='correlation coefficient threshold for clustering')

args = parser.parse_args()

start = int(args.gpsstart)
end = int(args.gpsend)
pad = args.filter_padding
auto_xlabel = ('Time [hours] from '
               + re.sub(r'\.0+', '',
                        Time(start, format='gps', scale='utc').iso)
               + ' UTC (%d)' % start)

if args.primary_channel == '{ifo}:GDS-CALIB_STRAIN':
    args.primary_frametype = '%s_HOFT_C00' % args.ifo
primary = args.primary_channel.format(ifo=args.ifo)
range_is_primary = (args.primary_channel
                    == '{ifo}:DMT-SNSH_EFFECTIVE_RANGE_MPC.mean')

if not os.path.isdir(args.output_dir):
    os.makedirs(args.output_dir)
os.chdir(args.output_dir)
nprocplot = args.nproc_plot or args.nproc

if args.band_pass:
    try:
        flower, fupper = args.band_pass
    except TypeError:
        flower, fupper = None

    print("-- Loading primary channel data")
    bandts = TimeSeries.get(primary, start-pad, end+pad, verbose=True,
                            nproc=args.nproc)
    if(flower < 0 or fupper >= float(bandts.sample_rate/(1.*u.Hz))*.5):
        print("\tError with band pass: frequency is"
              " out of range for this channel.")
        print("\tBand (Hz): " + str(args.band_pass)
              + ", Channel sample rate (Hz): " + str(bandts.sample_rate))
        quit()

    # get darm BLRMS
    print("-- Filtering data")
    if args.trend_type == 'minute':
        stride = 60
    else:
        stride = 1
    if flower:
        darmblrms = (
            bandts.highpass(flower/2., fstop=flower/4.,
                            filtfilt=False, ftype='butter')
            .notch(60, filtfilt=False)
            .bandpass(flower, fupper, fstop=[flower/2., fupper*1.5],
                      filtfilt=False, ftype='butter')
            .crop(start, end).rms(stride))
        darmblrms.name = '%s %s-%s Hz BLRMS' % (primary, flower, fupper)
    else:
        darmblrms = bandts.notch(60).crop(start, end).rms(stride)
        darmblrms.name = '%s RMS' % primary

    primaryts = darmblrms

else:
    # load primary channel data
    print("-- Loading primary channel data")
    primaryts = TimeSeries.get(primary, start, end,
                               frametype=args.primary_frametype,
                               verbose=True, nproc=args.nproc)

if args.remove_outliers:
    print("-- Removing outliers above %f sigma" % args.remove_outliers)
    remove_outliers(primaryts, args.remove_outliers)

hour_axis = (numpy.arange(len(primaryts.value)))/60
for i in range(len(hour_axis)):
    hour_axis[i] = round(hour_axis[i], 2)

# get aux data
print("-- Loading auxiliary channel data")
host, port = io_nds2.host_resolution_order(args.ifo)[0]
if args.channel_file is None:
    channels = ChannelList.query_nds2('*.mean', host=host, port=port,
                                      type='m-trend')
else:
    with open(args.channel_file, 'r') as f:
        channels = f.read().rstrip('\n').split('\n')
nchan = len(channels)
print("Identified %d channels" % nchan)

if args.trend_type == 'minute':
    frametype = '%s_M' % args.ifo  # for minute trends
else:
    frametype = '%s_T' % args.ifo  # for second trends

auxdata = TimeSeriesDict.get(
    map(str, channels), start, end, verbose=True,
    frametype=frametype, nproc=args.nproc,
    observatory=args.ifo[0], pad=0, **io_kw)

# -- removes flat data to be re-introdused later
flatdata = dict()
gooddata = dict()

for k, ts in auxdata.items():
    flat = ts.value.min() == ts.value.max()
    if flat:
        flatdata[k] = ts
    else:
        gooddata[k] = ts
auxdata = gooddata
flattab = Table(data=(numpy.asarray(flatdata.keys()),), names=('Channels',))

# -- removes NaN data
nandata = dict()
gooddata = dict()

try:
    data = numpy.array([scale(ts.value) for ts in auxdata.values()]).T
except ValueError:
    print("Nan or Inf found in data, removing bad channels...")
    for k, v in auxdata.items():
        hasnan = 0
        for x in range(len(v.value)):
            if isnan(v.value[x]) or isinf(v.value[x]):
                hasnan += 1
        if hasnan > 0:
            nandata[k] = v
        else:
            gooddata[k] = v
    print("Nan and Inf channels removed.")
    auxdata = gooddata
    data = numpy.array([scale(ts.value) for ts in auxdata.values()]).T


# -- perform LASSO regression -------------------------------------------------

# create model
if args.alpha is None:
    alphas = numpy.logspace(-1, 0, 100, endpoint=True)
    primary_scaled = scale(primaryts.value)
    nchans = numpy.zeros(len(alphas))
    coef_path = numpy.zeros((len(data[0, :]), len(alphas)))
    for i in range(0, len(alphas)):
        model = linear_model.Lasso(alpha=alphas[i])
        model.fit(data, primary_scaled)
        nchans[i] = len(numpy.nonzero(model.coef_)[0])
        coef_path[:, i] = model.coef_
    if 0 in nchans:
        badalphas = list()
        for i in range(0, len(alphas)):
            if nchans[i] == 0:
                badalphas.append(i)
        badalphas = badalphas[::-1]
        for i in range(0, len(badalphas)):
            alphas = numpy.delete(alphas, badalphas[i])
            nchans = numpy.delete(nchans, badalphas[i])
            coef_path = numpy.delete(coef_path, badalphas[i], 1)

    X = data
    y = primary_scaled
    n_samples = X.shape[0]
    K_A = 2
    K_B = log(n_samples)

    R = y[:, numpy.newaxis] - numpy.dot(X, coef_path)  # residuals
    mean_squared_error = numpy.mean(R ** 2, axis=0)
    sigma2 = numpy.var(y)

    df = nchans
    eps64 = numpy.finfo('float64').eps
    criterion_A = (n_samples * mean_squared_error / (sigma2 + eps64)
                   + K_A * df)  # Eqns. 2.15--16 in (Zou et al, 2007)
    criterion_B = (n_samples * mean_squared_error / (sigma2 + eps64)
                   + K_B * df)  # Eqns. 2.15--16 in (Zou et al, 2007)
    n_best = numpy.argmin(criterion_B)
    n_best_A = numpy.argmin(criterion_A)

    alpha_ = alphas[n_best]
    coef_ = coef_path[:, n_best]

    alpha_A = alphas[n_best_A]
    coef_A = coef_path[:, n_best_A]
    model = linear_model.Lasso(alpha_)
    model.fit(data, primary_scaled)
else:
    model = linear_model.Lasso(args.alpha)
    model.fit(data, scale(primaryts.value))

# pulls out alphas
usedalpha = 0
if args.alpha is None:
    usedalpha = model.alpha
else:
    usedalpha = args.alpha

# restructure results for convenience
unsorted_results = []
for n, k in enumerate(auxdata.keys()):
    unsorted_results.append([k, model.coef_[n]])
unsorted_results = sorted(unsorted_results, key=lambda x: abs(x[1]),
                          reverse=True)
sorted_results = ([[x[0] for x in unsorted_results],
                  [y[1] for y in unsorted_results]])
resultstab = Table(data=(sorted_results[0], sorted_results[1]),
                   names=('Channel', 'LASSO Coefficient'))

nonzerodata = dict()
nonzerocoef = dict()
usefulcount = 0
i = 0

while(i < len(sorted_results[1]) and abs(sorted_results[1][i]) > 0):
    nonzerodata[sorted_results[0][i]] = auxdata[sorted_results[0][i]]
    nonzerocoef[sorted_results[0][i]] = sorted_results[1][i]
    if abs(sorted_results[1][i]) >= args.threshold:
        usefulcount += 1
    i += 1

usefultab = Table(data=(resultstab['Channel'][0:usefulcount],
                        resultstab['LASSO Coefficient'][0:usefulcount]),
                  names=('Channel', 'LASSO Coefficient'))

zeroed = resultstab['LASSO Coefficient'] == 0
zeroedtab = Table(data=(resultstab[zeroed]['Channel'],), names=('Channels',))

# print results
print('Found {} channels with |LASSO Coefficient| >= {}'.format(
    len(usefultab), args.threshold))
print(usefultab)

gpsstub = '%d-%d' % (start, end-start)
resultsfile = '%s-LASSO_RESULTS-%s.txt' % (args.ifo, gpsstub)
zerofile = '%s-ZERO_COEFFICIENT_CHANNELS-%s.txt' % (args.ifo, gpsstub)
flatfile = '%s-FLAT_CHANNELS-%s.txt' % (args.ifo, gpsstub)

resultstab.write(resultsfile, format='ascii', overwrite=True)
zeroedtab.write(zerofile, format='ascii', overwrite=True)
flattab.write(flatfile, format='ascii', overwrite=True)

df = usefultab.to_pandas()
df.index += 1

# generate LASSO plots
modelFit = model.predict(data)

re_delim = re.compile('[:_-]')
form = '%%.%dd' % len(str(nchan))
p1 = (.1, .15, .9, .9)  # global plot defaults for plot1, lasso model

fig = plt.figure(figsize=(12, 6))
fig.subplots_adjust(*p1)
ax = fig.add_subplot(1, 1, 1)
ax.plot(hour_axis, scale(primaryts.value), label=primary.replace('_', '\_'))
ax.plot(hour_axis, modelFit, label='LASSO model')
ax.margins(x=0)
ax.set_xlabel(auto_xlabel)
ax.set_ylabel('Scaled arbitrary units')
if range_is_primary:
    ax.set_title('LASSO Model of Range')
else:
    ax.set_title('LASSO Model of Primary Channel')
ax.legend(loc='best')
fig.canvas.draw_idle()

plot1 = '%s-LASSO_MODEL-%s.png' % (args.ifo, gpsstub)
try:
    fig.savefig(plot1)
except (RuntimeError, IOError, IndexError):
    try:
        fig.savefig(plot1)
    except (RuntimeError, IOError, IndexError) as e:
        print("Error trying to save plot1 image, %s: " % plot1)
        print(e)
        plot1 = None
plt.close(fig)

# generate plots for channel contribution to model
colors = []
labels = []
fig = plt.figure(figsize=(8, 5))
fig.subplots_adjust(*p1)
ax = fig.add_subplot(1, 1, 1)
firstchan = sorted_results[0][0]

line1 = ax.plot(
    hour_axis,
    scale(primaryts.value),
    label=primary.replace('_', '\_'))
colors.append(line1[0].get_color())
labels.append(line1[0].get_label())

line2 = ax.plot(
    hour_axis,
    scale(nonzerodata[firstchan].value)*nonzerocoef[firstchan],
    label='Channel 1')
colors.append(line2[0].get_color())
labels.append(line2[0].get_label())

for n in range(1, len(nonzerodata)):
    summation = scale(nonzerodata[firstchan].value)*nonzerocoef[firstchan]
    for m in range(n, 0, -1):
        chan = sorted_results[0][m]
        summation = numpy.add(summation, (scale(nonzerodata[chan].value)
                                          * nonzerocoef[chan]))
    line = ax.plot(hour_axis, summation, label='Channels 1-' + str(n+1))
    colors.append(line[0].get_color())
    labels.append(line[0].get_label())

ax.margins(x=0)
ax.set_xlabel(auto_xlabel)
ax.set_ylabel('Scaled arbitrary units')
ax.set_title('Summations of Channel Contributions to Model')
fig.canvas.draw_idle()

plot2 = '%s-LASSO_CHANNEL_SUMMATION-%s.png' % (args.ifo, gpsstub)
try:
    fig.savefig(plot2)
except (RuntimeError, IOError, IndexError):
    try:
        fig.savefig(plot2)
    except (RuntimeError, IOError, IndexError) as e:
        print("Error trying to save plot2 image, %s: " % plot2)
        print(e)
        plot2 = None
plt.close(fig)

legend_fig = plt.figure()
patches = [
    Patch(color=color, label=label)
    for label, color in zip(labels, colors)]
legend = legend_fig.legend(patches, labels, loc='center',
                           frameon=False, fontsize='small')
legend_fig.canvas.draw_idle()

plot2_legend = ('%s-LASSO_CHANNEL_SUMMATION_LEGEND-%s.png'
                % (args.ifo, gpsstub))
try:
    legend_fig.savefig(
        plot2_legend,
        bbox_inches=(legend.get_window_extent()
                     .transformed(legend_fig.dpi_scale_trans.inverted())))
except (RuntimeError, IOError, IndexError):
    try:
        legend_fig.savefig(
            plot2_legend,
            bbox_inches=(legend.get_window_extent()
                         .transformed(legend_fig.dpi_scale_trans.inverted())))
    except (RuntimeError, IOError, IndexError) as e:
        print("Error trying to save plot2_legend image, %s: " % plot2_legend)
        print(e)
        plot2_legend = None
plt.close(legend_fig)

colors = []
labels = []
fig = plt.figure(figsize=(8, 5))
fig.subplots_adjust(*p1)
ax = fig.add_subplot(1, 1, 1)
line = ax.plot(hour_axis, scale(primaryts.value),
               label=primary.replace('_', '\_'))
colors.append(line[0].get_color())
labels.append(line[0].get_label())

for n in range(0, len(nonzerodata)):
    chan = sorted_results[0][n]
    if nonzerodata[chan] is not None:
        line = ax.plot(hour_axis,
                       (scale(nonzerodata[chan].value)
                           * nonzerocoef[chan]),
                       label=chan.replace('_', '\_'))
        colors.append(line[0].get_color())
        labels.append(line[0].get_label())

ax.margins(x=0)
ax.set_xlabel(auto_xlabel)
ax.set_ylabel('Scaled arbitrary units')
ax.set_title('Individual Channel Contributions to Model')
fig.canvas.draw_idle()

plot3 = '%s-LASSO_CHANNEL_CONTRIBUTIONS-%s.png' % (args.ifo, gpsstub)
try:
    fig.savefig(plot3)
except (RuntimeError, IOError, IndexError):
    try:
        fig.savefig(plot3)
    except (RuntimeError, IOError, IndexError) as e:
        print("Error trying to save plot3 image, %s: " % plot3)
        print(e)
        plot3 = None
plt.close(fig)

legend_fig = plt.figure()
patches = [
    Patch(color=color, label=label)
    for label, color in zip(labels, colors)]
legend = legend_fig.legend(patches, labels, loc='center',
                           frameon=False, fontsize='small')
legend_fig.canvas.draw_idle()

plot3_legend = ('%s-LASSO_CHANNEL_CONTRIBUTIONS_LEGEND-%s.png'
                % (args.ifo, gpsstub))
try:
    legend_fig.savefig(
        plot3_legend,
        bbox_inches=(legend.get_window_extent()
                     .transformed(legend_fig.dpi_scale_trans.inverted())))
except (RuntimeError, IOError, IndexError):
    try:
        legend_fig.savefig(
            plot3_legend,
            bbox_inches=(legend.get_window_extent()
                         .transformed(legend_fig.dpi_scale_trans.inverted())))
    except (RuntimeError, IOError, IndexError) as e:
        print("Error trying to save plot3_legend image, %s: " % plot3_legend)
        print(e)
        plot3_legend = None
plt.close(legend_fig)

# process aux channels, making plots
print("-- Processing channels")
counter = multiprocessing.Value('i', 0)
p4 = (.1, .1, .9, .95)  # global plot defaults for plot4, timeseries subplots


def process_channel(input_,):
    configure_mpl()
    chan = input_[1][0]
    ts = input_[1][1]
    lassocoef = nonzerocoef[chan]
    zeroed = lassocoef == 0

    if zeroed:
        plot4 = None
        plot5 = None
        plot6 = None
        pcorr = None
    else:
        plot4 = None
        plot5 = None
        plot6 = None
        if args.trend_type == 'minute':
            pcorr = numpy.corrcoef(ts.value, primaryts.value)[0, 1]
        else:
            pcorr = 0.0
        if(abs(lassocoef) < args.threshold):
            with counter.get_lock():
                counter.value += 1
                pc = 100 * counter.value / len(nonzerodata)
                print("Completed [%d/%d] %3d%% %-50s"
                      % (counter.value, len(nonzerodata), pc,
                         '(%s)' % str(chan)),
                      end='\r')
                sys.stdout.flush()
            return chan, lassocoef, plot4, plot5, plot6, ts

        # create time series subplots
        fig = TimeSeriesPlot(primaryts, ts, sep=True, sharex=True,
                             figsize=(12, 12))
        fig.subplots_adjust(*p4)
        if range_is_primary:
            fig.axes[0].set_ylabel('Sensitive range [Mpc]')
        else:
            fig.axes[0].set_ylabel('Primary channel units')
        fig.axes[1].set_ylabel('Channel units')
        for ax in fig.axes:
            ax.legend(loc='best')
            ax.set_xlim(start, end)
            ax.set_epoch(start)
        channelstub = re_delim.sub('_', str(chan)).replace('_', '-', 1)

        plot4 = '%s_TRENDS-%s.png' % (channelstub, gpsstub)
        try:
            fig.canvas.draw_idle()
            fig.save(plot4)
        except (RuntimeError, IOError, IndexError):
            try:
                fig.canvas.draw_idle()
                fig.save(plot4)
            except (RuntimeError, IOError, IndexError) as e:
                print("Error trying to save plot4 image, %s: " % plot4)
                print(e)
                plot4 = None
        plt.close(fig)

        # create scaled, sign-corrected, and overlayed timeseries
        tsscaled = scale(ts.value)
        if lassocoef < 0:
            tsscaled = numpy.negative(tsscaled)
        fig = plt.figure(figsize=(12, 6))
        fig.subplots_adjust(*p1)
        ax = fig.add_subplot(1, 1, 1)
        ax.plot(
            hour_axis, scale(primaryts.value),
            label=primary.replace('_', '\_'))
        ax.plot(
            hour_axis, tsscaled,
            label=chan.replace('_', '\_'))
        ax.margins(x=0)
        ax.set_xlabel(auto_xlabel)
        ax.set_ylabel('Scaled amplitude [arbitrary units]')
        ax.legend(loc='best')

        plot5 = '%s_COMPARISON-%s.png' % (channelstub, gpsstub)
        try:
            fig.canvas.draw_idle()
            fig.savefig(plot5)
        except (RuntimeError, IOError, IndexError):
            try:
                fig.canvas.draw_idle()
                fig.savefig(plot5)
            except (RuntimeError, IOError, IndexError) as e:
                print("Error trying to save plot5 image, %s: " % plot5)
                print(e)
                plot5 = None
        plt.close(fig)

        # scatter plot
        primaryColor = 'red'
        plotHeight = 6
        plotWidth = 12

        tsCopy = ts.reshape(-1, 1)
        primarytsCopy = primaryts.reshape(-1, 1)
        primaryReg = linear_model.LinearRegression()
        primaryReg.fit(tsCopy, primarytsCopy)
        primaryFit = primaryReg.predict(tsCopy)

        fig = gwplot()
        fig.subplots_adjust(*p1)
        ax = fig.add_subplot(1, 1, 1)
        ax.set_xlabel(chan.replace('_', '\_') + ' [Channel units]')
        if range_is_primary:
            ax.set_ylabel('Sensitive range [Mpc]')
        else:
            ax.set_ylabel('Primary channel units')
        yrange = abs(max(primaryts.value) - min(primaryts.value))
        yupper = max(primaryts.value) + .1 * yrange
        ylower = min(primaryts.value) - .1 * yrange
        ax.set_ylim(ylower, yupper)
        ax.text(.9, .1, 'r = ' + str('{0:.2}'.format(pcorr)),
                verticalalignment='bottom', horizontalalignment='right',
                transform=ax.transAxes, color='black', size=20,
                bbox=dict(boxstyle='square', facecolor='white', alpha=.75,
                edgecolor='black'))
        fig.add_scatter(ts, primaryts, color=primaryColor)
        fig.add_line(ts, primaryFit, color='black')
        fig.set_figheight(plotHeight)
        fig.set_figwidth(plotWidth)

        plot6 = '%s_SCATTER-%s.png' % (channelstub, gpsstub)
        try:
            fig.canvas.draw_idle()
            fig.save(plot6)
        except (RuntimeError, IOError, IndexError):
            try:
                fig.canvas.draw_idle()
                fig.save(plot6)
            except (RuntimeError, IOError, IndexError) as e:
                print("Error trying to save plot6 image, %s: " % plot6)
                print(e)
                plot6 = None
        plt.close(fig)

    # increment counter and print status
    with counter.get_lock():
        counter.value += 1
        pc = 100 * counter.value / len(nonzerodata)
        print("Completed [%d/%d] %3d%% %-50s"
              % (counter.value, len(nonzerodata), pc,
                 '(%s)' % str(chan)),
              end='\r')
        sys.stdout.flush()
    return chan, lassocoef, plot4, plot5, plot6, ts


# process channels
pool = multiprocessing.Pool(nprocplot)
results = pool.map(process_channel, enumerate(nonzerodata.iteritems()))
results = sorted(results, key=lambda x: abs(x[1]), reverse=True)

#  generate clustered time series plots
counter = multiprocessing.Value('i', 0)
p7 = (.135, .15, .95, .9)  # global plot defaults for plot7, clusters


def generate_cluster(input_,):
    configure_mpl()
    currentchan = input_[1][0]
    currentts = input_[1][5]
    current = input_[0]
    cluster_threshold = args.cluster_coefficient
    plot7 = None
    plot7_legend = None

    if current < len(nonzerodata):
        current_cluster = []
        for other, otheritem in enumerate(auxdata.iteritems()):
            otherchan = otheritem[0]
            otherts = otheritem[1]
            if otherchan != currentchan:
                pcorr = numpy.corrcoef(currentts.value, otherts.value)[0, 1]
                if(abs(pcorr) >= cluster_threshold):
                    otherchannelstub = (re_delim.sub('_', otherchan)
                                        .replace('_', '-', 1))
                    current_cluster.append([other, otherts, pcorr,
                                            otherchan, otherchannelstub])
        if len(current_cluster) == 0:
            with counter.get_lock():
                counter.value += 1
                pc = 100 * counter.value / len(nonzerodata)
                print("Completed [%d/%d] %3d%% %-50s"
                      % (counter.value, len(nonzerodata), pc,
                         '(%s)' % str(currentchan)),
                      end='\r')
                sys.stdout.flush()
            return plot7, plot7_legend
        else:
            colors = []
            labels = []
            fig = plt.figure(figsize=(8, 5))
            fig.subplots_adjust(*p7)
            ax = fig.add_subplot(1, 1, 1)
            line = ax.plot(
                hour_axis,
                scale(currentts.value)*numpy.sign(input_[1][1]),
                label=currentchan.replace('_', '\_'))
            colors.append(line[0].get_color())
            labels.append(line[0].get_label())
            current_cluster = sorted(current_cluster, key=lambda x: abs(x[2]),
                                     reverse=True)

            for i in range(0, len(current_cluster)):
                line = ax.plot(
                    hour_axis,
                    (scale(current_cluster[i][1].value)
                        * numpy.sign(input_[1][1])
                        * numpy.sign(current_cluster[i][2])),
                    label=(current_cluster[i][3].replace('_', '\_')
                           + ', r = '
                           + str('{0:.2}'.format(current_cluster[i][2]))))
                colors.append(line[0].get_color())
                labels.append(line[0].get_label())

            ax.margins(x=0)
            ax.set_xlabel(auto_xlabel)
            ax.set_ylabel('Scaled amplitude [arbitrary units]')
            ax.set_title('Highly Correlated Channels')

            plot7 = '%s_CLUSTER-%s.png' % (
                re_delim.sub('_', str(currentchan))
                    .replace('_', '-', 1),
                gpsstub)
            try:
                fig.canvas.draw_idle()
                fig.savefig(plot7)
            except (RuntimeError, IOError, IndexError):
                try:
                    fig.canvas.draw_idle()
                    fig.savefig(plot7)
                except (RuntimeError, IOError, IndexError) as e:
                    print("Error trying to save plot7 image, %s: " % plot7)
                    print(e)
                    plot7 = None
            plt.close(fig)

            legend_fig = plt.figure()
            patches = [
                Patch(color=color, label=label)
                for label, color in zip(labels, colors)]
            legend = legend_fig.legend(patches, labels, loc='center',
                                       frameon=False, fontsize='small')

            plot7_legend = '%s_CLUSTER_LEGEND-%s.png' % (
                re_delim.sub('_', str(currentchan)).replace('_', '-', 1),
                gpsstub)
            try:
                legend_fig.canvas.draw_idle()
                legend_fig.savefig(
                    plot7_legend,
                    bbox_inches=(legend.get_window_extent()
                                 .transformed(legend_fig.dpi_scale_trans
                                              .inverted())))
            except (RuntimeError, IOError, IndexError):
                try:
                    legend_fig.canvas.draw_idle()
                    legend_fig.savefig(
                        plot7_legend,
                        bbox_inches=(legend.get_window_extent()
                                     .transformed(legend_fig.dpi_scale_trans
                                                  .inverted())))
                except (RuntimeError, IOError, IndexError) as e:
                    print("Error trying to save plot7_legend image, %s: "
                          % plot7_legend)
                    print(e)
                    plot7_legend = None
            plt.close(legend_fig)

    with counter.get_lock():
        counter.value += 1
        pc = 100 * counter.value / len(nonzerodata)
        print("Completed [%d/%d] %3d%% %-50s"
              % (counter.value, len(nonzerodata), pc,
                 '(%s)' % str(currentchan)),
              end='\r')
        sys.stdout.flush()
    return plot7, plot7_legend


if args.no_cluster is False:
    print("-- Generating clusters")
    pool = multiprocessing.Pool(nprocplot)
    clusters = pool.map(generate_cluster, enumerate(results))

channelsfile = '%s-CHANNELS-%s.txt' % (args.ifo, gpsstub)
numpy.savetxt(channelsfile, channels, fmt='%s')

# write html
title = '%s LASSO slow correlations: %d-%d' % (args.ifo, start, end)
page = html.new_bootstrap_page(title=title)
page.div(class_='container')

page.div(class_='page-header')
page.h1(title)
page.div.close()  # page-header


# params
def write_param(param, value):
    page.p()
    page.strong('%s: ' % param)
    page.add(str(value))
    page.p.close()


page.h2('Parameters')
page.p('This analysis used the following parameters:')
write_param('Start time', '%s (%d)' % (from_gps(start), start))
write_param('End time', '%s (%d)' % (from_gps(end), end))
write_param('Primary channel',
            '%s (%s)' % (primary, args.primary_frametype or '-'))
write_param('Channels searched',
            '%d (%s)' % (nchan, "<a href= %s target='_blank'>channel list</a>"
                         % channelsfile))
write_param('Number of flat channels',
            '%d (%s)' % (len(flatdata),
                         "<a href= %s target='_blank'>flat channel list</a>"
                         % (flatfile)))
write_param('LASSO coefficient threshold', '%g' % args.threshold)
write_param('Sigma for outlier removal', '%g' % args.remove_outliers)
write_param('Cluster coefficient threshold: ',
            '%g' % args.cluster_coefficient)

page.h2('Model Information')

page.div(class_='model')
page.div(class_='model-body')

page.div(class_='model-info')
write_param('Model', 'LASSO')
write_param('Non-zero coefficients', '%d' % numpy.count_nonzero(model.coef_))
write_param('Alpha', '%g' % usedalpha)
write_param('Zero coefficients',
            '%d (%s)' % (len(zeroedtab),
                         "<a href= %s target='_blank'>zeroed channel list</a>"
                         % (zerofile)))
page.p('<br /><br />%s' % df.to_html(index=True))
page.div.close()  # model-info

page.p('<br /><br />')

page.div(class_='primary-lasso')
page.a(href=plot1, target='_blank')
page.img(class_='lasso-img', src=plot1)
page.a.close()  # primary lasso plot
page.div.close()  # primary-lasso

page.div(
    class_='channel-summation',
    style_='display: flex;flex-wrap:nowrap;justify-content:space-around;')
page.div(style_='display: block;')
page.a(href=plot2, target='_blank')
page.img(class_='channels-summation-img', src=plot2)
page.a.close()  # channels-summation plot
page.div.close()  # close plot2 div
page.div(class_='scroll-container',
         style_='display: block; padding-top: 40px;')
page.div(style_='display: block; overflow:auto; height:400px;')
page.a(href=plot2_legend, target='_blank')
page.img(class_='channels-contrib-img', src=plot2_legend)
page.a.close()  # legend image
page.div.close()  # overflowed window
page.div.close()  # scroll container
page.div.close()  # channels-summation

page.div(
    class_='channels-and-primary',
    style_='display: flex;flex-wrap:nowrap;justify-content:space-around;')
page.div(style_='display: block;')
page.a(href=plot3, target='_blank')
page.img(class_='channels-contrib-img', src=plot3)
page.a.close()  # channels-contrib plot
page.div.close()  # plot3 div
page.div(class_='scroll-container',
         style_='display: block; padding-top: 40px;')
page.div(style_='display: block; overflow:auto; height:400px;')
page.a(href=plot3_legend, target='_blank')
page.img(class_='channels-contrib-img', src=plot3_legend)
page.a.close()  # legend image
page.div.close()  # overflowed window
page.div.close()  # scroll container
page.div.close()  # channels-and-primary

page.div.close()  # model-body
page.div.close()  # model

# results
page.h2('Top Channels')

page.div(class_='panel-group', id_='results')

# for each auxiliary channel create information container and put plots in it
for i, (ch, lassocoef, plot4, plot5, plot6, ts) in enumerate(results):
    # set container color/context based on lasso coefficient
    if lassocoef == 0:
        break
    elif abs(lassocoef) < args.threshold:
        h = '%s [lasso coefficient = %.4f] (Below threshold)' % (ch, lassocoef)
    else:
        h = '%s [lasso coefficient = %.4f]' % (ch, lassocoef)
    if ((lassocoef is None) or (lassocoef == 0)
            or (abs(lassocoef) < args.threshold)):
        context = 'panel-default'
    elif abs(lassocoef) >= .5:
        context = 'panel-danger'
    elif abs(lassocoef) >= .2:
        context = 'panel-warning'
    else:
        context = 'panel-info'
    page.div(class_='panel %s' % context)

    # heading
    page.div(class_='panel-heading')
    page.a(h, class_='panel-title', href='#channel%d' % i,
           **{'data-toggle': 'collapse', 'data-parent': '#results'})
    page.div.close()  # panel-heading
    # body
    page.div(id_='channel%d' % i, class_='panel-collapse collapse')
    page.div(class_='panel-body')
    if lassocoef is None:
        page.p('The amplitude data for this channel is flat (does not change)'
               ' for the chosen time period.')
    elif abs(lassocoef) < args.threshold:
        page.p('LASSO coefficient below the threshold of %g.'
               % (args.threshold))
    else:
        for p in (plot4, plot5, plot6):
            page.a(href=p, target='_blank')
            page.img(class_='img-responsive', src=p)
            page.a.close()
        if args.no_cluster is False:
            if clusters[i][0] is None:
                page.p("<font size = '3'><br />No channels were highly"
                       " correlated with this channel.</font>")
            else:
                page.div(
                    class_='clusters',
                    style_='display: flex;flex-wrap:nowrap;'
                           'justify-content:space-around;')
                page.div(style_='display: block;')
                page.a(href=clusters[i][0], target='_blank')
                page.img(class_='img-responsive', src=clusters[i][0])
                page.a.close()
                page.div.close()  # close plot7 div
                page.div(class_='scroll-container',
                         style_='display: block; padding-top: 40px;')
                page.div(style_='display: block; overflow:auto;'
                                ' height: 375px;')
                page.a(href=clusters[i][1], target='_blank')
                page.img(class_='img-responsive', src=clusters[i][1])
                page.a.close()  # legend image
                page.div.close()  # overflowed window
                page.div.close()  # scroll container
                page.div.close()  # channels-summation
    page.div.close()  # panel-body
    page.div.close()  # panel-collapse
    page.div.close()  # panel
page.div.close()  # panel-group
page.div.close()  # container
with open('index.html', 'w') as f:
    print(str(page), file=f)

print ("-- Process Completed")