text:
"""Sequence-to-sequence model for human motion prediction."""

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

class Seq2SeqModel(nn.Module):
    """Sequence-to-sequence model for human motion prediction"""

    def __init__(
        self,
        architecture,
        source_seq_len,
        target_seq_len,
        rnn_size,  # hidden recurrent layer size
        num_layers,
        max_gradient_norm,
        batch_size,
        learning_rate,
        learning_rate_decay_factor,
        loss_to_use,
        number_of_actions,
        one_hot=True,
        residual_velocities=False,
        output_dim=67,
        dropout=0.0,
        dtype=torch.float32,
        num_traj=0,
    ):
        """Create the model.

    Args:
      architecture: [basic, tied] whether to tie the decoder and decoder.
      source_seq_len: lenght of the input sequence.
      target_seq_len: lenght of the target sequence.
      rnn_size: number of units in the rnn.
      num_layers: number of rnns to stack.
      max_gradient_norm: gradients will be clipped to maximally this norm.
      batch_size: the size of the batches used during training;
        the model construction is independent of batch_size, so it can be
        changed after initialization if this is convenient, e.g., for decoding.
      learning_rate: learning rate to start with.
      learning_rate_decay_factor: decay learning rate by this much when needed.
      loss_to_use: [supervised, sampling_based]. Whether to use ground truth in
        each timestep to compute the loss after decoding, or to feed back the
        prediction from the previous time-step.
      number_of_actions: number of classes we have.
      one_hot: whether to use one_hot encoding during train/test (sup models).
      residual_velocities: whether to use a residual connection that models velocities.
      dtype: the data type to use to store internal variables.
    """
        super(Seq2SeqModel, self).__init__()

        self.HUMAN_SIZE = output_dim
        self.input_size = self.HUMAN_SIZE + num_traj + (number_of_actions if one_hot else 0)

        print("One hot is ", one_hot)
        print("Input size is %d" % self.input_size)

        # Summary writers for train and test runs

        self.source_seq_len = source_seq_len
        self.target_seq_len = target_seq_len
        self.rnn_size = rnn_size
        self.batch_size = batch_size
        self.dropout = dropout
        self.num_layers = num_layers
        # === Create the RNN that will keep the state ===
        print("rnn_size = {0}".format(rnn_size))
        self.cell = torch.nn.GRUCell(self.input_size, self.rnn_size)
        if num_layers > 1:
            self.cell2 = torch.nn.GRUCell(self.rnn_size, self.rnn_size)

        self.fc1 = nn.Linear(self.rnn_size * num_layers, self.HUMAN_SIZE)

    def forward(self, encoder_inputs, decoder_inputs, use_cuda):
        # This appears to be coded in a slightly odd way, but is retained from the
        # pytorch port of Martinez et al. => some confidence it's correct.
        def loop_function(prev, i):
            return prev

        batchsize = encoder_inputs.shape[0]
        encoder_inputs = torch.transpose(encoder_inputs, 0, 1)
        decoder_inputs = torch.transpose(decoder_inputs, 0, 1)

        state = torch.zeros(batchsize, self.rnn_size)
        if self.num_layers > 1:
            state2 = torch.zeros(batchsize, self.rnn_size)
        if use_cuda:
            state = state.cuda()
            if self.num_layers > 1:
                state2 = state2.cuda()

        for i in range(self.source_seq_len - 1):
            state = self.cell(encoder_inputs[i], state)
            if self.num_layers > 1:
                state2 = self.cell2(state, state2)
                state2 = F.dropout(state2, self.dropout, training=self.training)

            state = F.dropout(state, self.dropout, training=self.training)

        outputs = []
        prev = None
        for i, inp in enumerate(decoder_inputs):
            if loop_function is not None and prev is not None:
                inp = torch.cat((loop_function(prev, i), inp[:, self.HUMAN_SIZE :]), 1)

            inp = inp.detach()

            state = self.cell(inp, state)
            if self.num_layers > 1:
                state2 = self.cell2(state, state2)
                state_out = torch.cat((state, state2), dim=1)
            else:
                state_out = state

            output = inp[:, 0 : self.HUMAN_SIZE] + self.fc1(F.dropout(state_out, self.dropout, training=self.training))

            outputs.append(output.view([1, batchsize, self.HUMAN_SIZE]))
            if loop_function is not None:
                prev = output

        #    return outputs, state

        outputs = torch.cat(outputs, 0)
        return torch.transpose(outputs, 0, 1)

    def get_batch(self, data_Y, data_U, actions, stratify=False):
        """Get a random batch of data from the specified bucket, prepare for step.

    Args
      data: a list of sequences of size n-by-d to fit the model to.
      actions: a list of the actions we are using
    Returns
      The tuple (encoder_inputs, decoder_inputs, decoder_outputs);
      the constructed batches have the proper format to call step(...) later.
    """

        # Select entries at random
        probs = np.array([y.shape[0] for y in data_Y])
        probs = probs / probs.sum()
        bsz = self.batch_size if not stratify else len(data_Y)

        if not stratify:
            chosen_keys = np.random.choice(len(data_Y), self.batch_size, p=probs)
            bsz = self.batch_size
        else:
            bsz = len(data_Y)
            chosen_keys = list(range(bsz))

        # How many frames in total do we need?
        total_frames = self.source_seq_len + self.target_seq_len
        traj_size = data_U[1].shape[1]

        encoder_inputs = np.zeros((bsz, self.source_seq_len - 1, self.input_size), dtype=float)
        decoder_inputs = np.zeros((bsz, self.target_seq_len, self.input_size), dtype=float)
        decoder_outputs = np.zeros((bsz, self.target_seq_len, self.HUMAN_SIZE), dtype=float)

        for i in range(bsz):

            the_key = chosen_keys[i]

            # Get the number of frames
            n = data_Y[the_key].shape[0]
            if n < total_frames:
                how_did_we_get_here = 1
            assert n >= total_frames, "n of file {:d} too small.".format(the_key)

            # Sample somewherein the middle
            idx = np.random.randint(0, n - total_frames) if n > total_frames else 0

            # Select the data around the sampled points
            data_Y_sel = data_Y[the_key][idx : idx + total_frames, :]
            data_U_sel = data_U[the_key][idx : idx + total_frames, :]

            # Add the data
            encoder_inputs[i, :, 0 : self.HUMAN_SIZE] = data_Y_sel[0 : self.source_seq_len - 1, :]
            encoder_inputs[i, :, self.HUMAN_SIZE :] = data_U_sel[0 : self.source_seq_len - 1, :]  # <= done
            decoder_inputs[i, :, 0 : self.HUMAN_SIZE] = data_Y_sel[
                self.source_seq_len - 1 : self.source_seq_len + self.target_seq_len - 1, :
            ]
            decoder_inputs[i, :, self.HUMAN_SIZE :] = data_U_sel[
                self.source_seq_len - 1 : self.source_seq_len + self.target_seq_len - 1, :
            ]
            decoder_outputs[i, :, 0 : self.HUMAN_SIZE] = data_Y_sel[
                self.source_seq_len : self.source_seq_len + self.target_seq_len, :
            ]  # <= done

        return encoder_inputs, decoder_inputs, decoder_outputs

    def get_test_batch(self, data_Y, data_U, batch_size):
        """Get a random batch of data from the specified bucket, prepare for step.

    Args
      data: a list of sequences of size n-by-d to fit the model to.
      actions: a list of the actions we are using
    Returns
      The tuple (encoder_inputs, decoder_inputs, decoder_outputs);
      the constructed batches have the proper format to call step(...) later.
    """

        k_ahead = 6

        if batch_size <= 0:
            batch_size = len(data_Y) - k_ahead
            chosen_keys = range(batch_size)
        else:
            chosen_keys = np.random.choice(
                len(data_Y) - k_ahead, batch_size, replace=batch_size > len(data_Y) - k_ahead
            )

        # How many frames in total do we need?
        source_len = 64
        target_len = 64 * k_ahead
        total_frames = self.source_seq_len * (k_ahead + 1)

        encoder_inputs = np.zeros((batch_size, source_len - 1, self.input_size), dtype=float)
        decoder_inputs = np.zeros((batch_size, target_len, self.input_size), dtype=float)
        decoder_outputs = np.zeros((batch_size, target_len, self.HUMAN_SIZE), dtype=float)

        for i in chosen_keys:
            # Add the data
            encoder_inputs[i, :, 0 : self.HUMAN_SIZE] = data_Y[i].T[0 : source_len - 1, :]
            encoder_inputs[i, :, self.HUMAN_SIZE :] = data_U[i].T[0 : source_len - 1, :]

            decoder_inputs[i, 0, 0 : self.HUMAN_SIZE] = data_Y[i].T[source_len - 1, :]
            decoder_inputs[i, 0, self.HUMAN_SIZE :] = data_U[i].T[source_len - 1, :]
            for k in range(k_ahead - 1):
                decoder_inputs[i, 64 * k + 1 : 64 * (k + 1) + 1, 0 : self.HUMAN_SIZE] = data_Y[i + 1 + k].T
                decoder_inputs[i, 64 * k + 1 : 64 * (k + 1) + 1, self.HUMAN_SIZE :] = data_U[i + 1 + k].T
                decoder_outputs[i, 64 * k : 64 * (k + 1), :] = data_Y[i + 1 + k].T
            decoder_inputs[i, 64 * (k_ahead - 1) + 1 : 64 * k_ahead + 1, 0 : self.HUMAN_SIZE] = data_Y[i + k_ahead].T[
                0:63, :
            ]
            decoder_inputs[i, 64 * (k_ahead - 1) + 1 : 64 * k_ahead + 1, self.HUMAN_SIZE :] = data_U[i + k_ahead].T[
                0:63, :
            ]
            decoder_outputs[i, 64 * (k_ahead - 1) : 64 * k_ahead, :] = data_Y[i + k_ahead].T

        return encoder_inputs, decoder_inputs, decoder_outputs

    def find_indices_srnn(self, data, action):
        """
    Find the same action indices as in SRNN.
    See https://github.com/asheshjain399/RNNexp/blob/master/structural_rnn/CRFProblems/H3.6m/processdata.py#L325
    """

        # Used a fixed dummy seed, following
        # https://github.com/asheshjain399/RNNexp/blob/srnn/structural_rnn/forecastTrajectories.py#L29
        SEED = 1234567890
        rng = np.random.RandomState(SEED)

        subject = 5
        subaction1 = 1
        subaction2 = 2

        T1 = data[(subject, action, subaction1, "even")].shape[0]
        T2 = data[(subject, action, subaction2, "even")].shape[0]
        prefix, suffix = 50, 100

        idx = []
        idx.append(rng.randint(16, T1 - prefix - suffix))
        idx.append(rng.randint(16, T2 - prefix - suffix))
        idx.append(rng.randint(16, T1 - prefix - suffix))
        idx.append(rng.randint(16, T2 - prefix - suffix))
        idx.append(rng.randint(16, T1 - prefix - suffix))
        idx.append(rng.randint(16, T2 - prefix - suffix))
        idx.append(rng.randint(16, T1 - prefix - suffix))
        idx.append(rng.randint(16, T2 - prefix - suffix))
        return idx

    def get_batch_srnn(self, data, action):
        """
    Get a random batch of data from the specified bucket, prepare for step.

    Args
      data: dictionary with k:v, k=((subject, action, subsequence, 'even')),
        v=nxd matrix with a sequence of poses
      action: the action to load data from
    Returns
      The tuple (encoder_inputs, decoder_inputs, decoder_outputs);
      the constructed batches have the proper format to call step(...) later.
    """

        actions = [
            "directions",
            "discussion",
            "eating",
            "greeting",
            "phoning",
            "posing",
            "purchases",
            "sitting",
            "sittingdown",
            "smoking",
            "takingphoto",
            "waiting",
            "walking",
            "walkingdog",
            "walkingtogether",
        ]

        if not action in actions:
            raise ValueError("Unrecognized action {0}".format(action))

        frames = {}
        frames[action] = self.find_indices_srnn(data, action)

        batch_size = 8  # we always evaluate 8 seeds
        subject = 5  # we always evaluate on subject 5
        source_seq_len = self.source_seq_len
        target_seq_len = self.target_seq_len

        seeds = [(action, (i % 2) + 1, frames[action][i]) for i in range(batch_size)]

        encoder_inputs = np.zeros((batch_size, source_seq_len - 1, self.input_size), dtype=float)
        decoder_inputs = np.zeros((batch_size, target_seq_len, self.input_size), dtype=float)
        decoder_outputs = np.zeros((batch_size, target_seq_len, self.input_size), dtype=float)

        # Compute the number of frames needed
        total_frames = source_seq_len + target_seq_len

        # Reproducing SRNN's sequence subsequence selection as done in
        # https://github.com/asheshjain399/RNNexp/blob/master/structural_rnn/CRFProblems/H3.6m/processdata.py#L343
        for i in range(batch_size):

            _, subsequence, idx = seeds[i]
            idx = idx + 50

            data_sel = data[(subject, action, subsequence, "even")]

            data_sel = data_sel[(idx - source_seq_len) : (idx + target_seq_len), :]

            encoder_inputs[i, :, :] = data_sel[0 : source_seq_len - 1, :]
            decoder_inputs[i, :, :] = data_sel[source_seq_len - 1 : (source_seq_len + target_seq_len - 1), :]
            decoder_outputs[i, :, :] = data_sel[source_seq_len:, :]

        return encoder_inputs, decoder_inputs, decoder_outputs

max_stars_repo_name: ornithos/pytorch-mtds-mocap

max_issues_repo_name: ornithos/pytorch-mtds-mocap

max_forks_repo_name: ornithos/pytorch-mtds-mocap


text:
# -*- coding: utf-8 -*-
"""
Created on Thu Sep 18 18:34:12 2014

%reset -f
%pylab
%clear
%load_ext autoreload
%autoreload 2

@author: isaac
"""

from __future__ import division

import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import UnivariateSpline, interp1d

# setup better plots
import plots
reload(plots)
from plots import bmap, rcj, tl, tickout

import squirrel
reload(squirrel)

import eqns
reload(eqns)

# %% Load in the raw data

dph = 20
fname_cl = 'Data/Heers2011/cl_alpha_{}.csv'
fname_cd = 'Data/Heers2011/cd_alpha_{}.csv'

# load in raw data (digized plot has greater resolution than experiment)
cl_raw = np.genfromtxt(fname_cl.format(dph), delimiter=',')
cd_raw = np.genfromtxt(fname_cd.format(dph), delimiter=',')

# first column of digitized data in radians
cl_raw[:, 0] = np.deg2rad(cl_raw[:, 0])
cd_raw[:, 0] = np.deg2rad(cd_raw[:, 0])

# actual data sampled at -10 to 80 deg aoa in 10 deg increments
aoa_deg = np.arange(-10, 81, 10)
aoa_rad = np.deg2rad(aoa_deg)

# determine the actual experimental values
kind = 'linear'
Cl_raw = interp1d(cl_raw[:, 0], cl_raw[:, 1], kind=kind, bounds_error=False)
Cd_raw = interp1d(cd_raw[:, 0], cd_raw[:, 1], kind=kind, bounds_error=False)
Cl = Cl_raw(aoa_rad)
Cd = Cd_raw(aoa_rad)

# linear fit between points
Cl_fun = UnivariateSpline(aoa_rad, Cl, s=0, k=1)
Cd_fun = UnivariateSpline(aoa_rad, Cd, s=0, k=1)
Clprime_fun = Cl_fun.derivative()
Cdprime_fun = Cd_fun.derivative()

# fit a smoothing spline to experimental data
cl_fun = UnivariateSpline(aoa_rad, Cl, s=.0001, k=3)
cd_fun = UnivariateSpline(aoa_rad, Cd, s=.0001, k=3)
clprime_fun = cl_fun.derivative()
cdprime_fun = cd_fun.derivative()

# angles to evaluate spline over
angle_min, angle_max = aoa_rad.min(), aoa_rad.max()
alpha, dalpha = np.linspace(angle_min, angle_max, 1001, retstep=True)

# evaluate the spline
cl = cl_fun(alpha)
cd = cd_fun(alpha)
clprime = clprime_fun(alpha)
cdprime = cdprime_fun(alpha)

# evaluate the linear fit
Clf = Cl_fun(alpha)
Cdf = Cd_fun(alpha)
Clprime = Clprime_fun(alpha)
Cdprime = Cdprime_fun(alpha)

# %% Cl, Cd, and ClCd curves for paper (updated)

fig, ax = plt.subplots()
ax.axvline(0, color='gray', lw=1)
ax.axhline(0, color='gray', lw=1)
ax.plot(np.rad2deg(aoa_rad), Cl, 'o', ms=6, label=r'$C_L$')
ax.plot(np.rad2deg(aoa_rad), Cd, 's', ms=6, label=r'$C_D$')
ax.plot(np.rad2deg(aoa_rad), Cl / Cd, '^', ms=6, label=r'$C_L/C_D$')
ax.plot(np.rad2deg(alpha), cl, color=bmap[0], lw=1.5)
ax.plot(np.rad2deg(alpha), cd, color=bmap[1], lw=1.5)
ax.plot(np.rad2deg(alpha), cl / cd, color=bmap[2], lw=1.5)

ax.set_xlim(-15, 65)
ax.set_ylim(-2, 3)
ax.legend(loc='lower right', frameon=False, fontsize=18)
ax.set_xlabel(r'$\alpha$', fontsize=18)
ax.set_ylabel('force coefficients', fontsize=18)

plt.draw()
# add degree symbol to angles
ticks = ax.get_xticklabels()
newticks = []
for tick in ticks:
    text = tick.get_text()
    newticks.append(text + u'\u00B0')
ax.set_xticklabels(newticks)

ax.text(2, 2.75, 'chukar 20 d.p.h.', {'fontsize': 18})

[ttl.set_size(18) for ttl in ax.get_xticklabels()]
[ttl.set_size(18) for ttl in ax.get_yticklabels()]

rcj(ax)
tl(fig)

fig.savefig('Figures/figure4c_chukar.pdf', transparent=True)

# %% Find equilibrium points

pitches = np.deg2rad(np.linspace(-25, 25, 4000))
gammas = np.deg2rad(np.linspace(10, 90, 1000))
angle_rng = (angle_min, angle_max)

ch_equil_exp = eqns.pitch_bifurcation(pitches, gammas, Cl_fun, Cd_fun)
ch_equil_spl = eqns.pitch_bifurcation(pitches, gammas, cl_fun, cd_fun,
                                      angle_rng=angle_rng)

# %% Classify the stability of fixed points

ch_td_exp, ch_ev_exp = eqns.tau_delta(ch_equil_exp, Cl_fun, Cd_fun,
                                      Clprime_fun, Cdprime_fun)
ch_td_spl, ch_ev_spl = eqns.tau_delta(ch_equil_spl, cl_fun, cd_fun,
                                      clprime_fun, cdprime_fun,
                                      angle_rng=angle_rng)

# %% Stability analysis for individual squirrels

ch_nuni_exp, ch_uni_exp, ch_class_exp = eqns.classify_fp(ch_td_exp)
ch_nuni_spl, ch_uni_spl, ch_class_spl = eqns.classify_fp(ch_td_spl)

possible_class = ['saddle point', 'unstable focus', 'unstable node',
                  'stable focus', 'stable node']
bfbmap = [bmap[0], bmap[4], bmap[2], bmap[3], bmap[1]]

# %% Spline bifurcation plot for paper

rd = np.rad2deg

gam_high = angle_rng[0] - pitches  # closer to 0
gam_low = angle_rng[1] - pitches  # closer to 90

fig, ax = plt.subplots()

ax.fill_between(rd(pitches), rd(gam_high), 0, color='gray', alpha=.1, lw=0)
ax.fill_between(rd(pitches), rd(gam_low), 60, color='gray', alpha=.1, lw=0)

ax.axvline(0, color='gray')
ax.axvline(13, color='gray')

for ii, fp_kind in enumerate(possible_class):
    idx = np.where(ch_class_spl == fp_kind)[0]
    if len(idx) > 0:
        ax.plot(rd(ch_equil_spl[idx, 0]), rd(ch_equil_spl[idx, 1]), 'o',
                c=bfbmap[ii], ms=2.5, label=fp_kind)

_leg = ax.legend(loc='lower left', markerscale=3, fancybox=True, framealpha=.75,
                 frameon=True, fontsize=16)
_leg.get_frame().set_color('w')
ax.set_xlim(-15, 15)
ax.set_ylim(60, 0)
#ax.set_ylabel(r'$\gamma^*$, equilibrium glide angle', fontsize=18)
#ax.set_xlabel(r'$\theta$, pitch angle', fontsize=18)
ax.set_ylabel(r'$\gamma^*$    ', fontsize=18, rotation=0)
ax.set_xlabel(r'$\theta$', fontsize=18)

ax.text(-13, 5, 'chukar 20 d.p.h.', {'fontsize': 18})

plt.draw()

# add degree symbol to angles
ticks = ax.get_xticklabels()
newticks = []
for tick in ticks:
    text = tick.get_text()
    newticks.append(text + u'\u00B0')
ax.set_xticklabels(newticks)

ticks = ax.get_yticklabels()
newticks = []
for tick in ticks:
    text = tick.get_text()
    newticks.append(text + u'\u00B0')
ax.set_yticklabels(newticks)

[ttl.set_size(18) for ttl in ax.get_xticklabels()]
[ttl.set_size(18) for ttl in ax.get_yticklabels()]

rcj(ax)
tl(fig)

fig.savefig('Figures/figure6c_bifurcation_chukar.pdf',
            transparent=True)

# %% Plot the equilibrium glide velocity on a polar plot

fig, ax = plt.subplots()

for ii, fp_kind in enumerate(possible_class):

    idx = np.where(ch_class_spl == fp_kind)[0]
    if len(idx) > 0:
        geq = ch_equil_spl[idx, 1]
        veq = eqns.v_equil(geq, cl_fun, cd_fun)
        vxeq = veq * np.cos(geq)
        vzeq = -veq * np.sin(geq)

        ax.plot(vxeq, vzeq, 'o', c=bfbmap[ii], ms=2, label=fp_kind,
                mec=bfbmap[ii])

ax.axis('equal', adjustable='box')
ax.set_xlim(0, 1.5)
ax.set_ylim(-1.5, 0)
ax.set_xlabel(r"$\hat{v}_x$", fontsize=18)
ax.set_ylabel(r"$\hat{v}_z$", fontsize=18)
_leg = ax.legend(loc='best', markerscale=4, frameon=False, framealpha=.75)

ax.xaxis.set_label_position('top')
ax.xaxis.tick_top()
rcj(ax, ['bottom', 'right'])
tl(fig)

fig.savefig('Figures/appendix_bifurcation_chukar.pdf',
            transparent=True)

# %% Velocity polar diagram, pitch = 0

pitch = 0

afdict = dict(cli=cl_fun, cdi=cd_fun, clip=clprime_fun, cdip=cdprime_fun)

lims = (vxlim, vzlim) = np.r_[0, 1.25], np.r_[0, -1.25]
tvec = np.linspace(0, 30, 351)

reload(plots)
from plots import phase_plotter as ppr

fig, ax = ppr(afdict, pitch, lims, angle_rng, tvec, ngrid=201,
              nseed=501, nseed_skip=25, quiver=False, skip=10, seed=False,
              timer=True, gamtest=gammas, extrap=None,
              traj=plots.ps_traj_dp5, fig=None, ax=None)

lab = 'chukar 20 d.p.h., ' + r'$\theta=$0' + u'\u00B0'
ax.text(.05, -1, lab, fontsize=16)

fig.savefig('Figures/figure6ci_vpd0_chukar.pdf', transparent=True)

# %% Velocity polar diagram, pitch = 0 with Z nullclines

pitch = 0

afdict = dict(cli=cl_fun, cdi=cd_fun, clip=clprime_fun, cdip=cdprime_fun)

lims = (vxlim, vzlim) = np.r_[0, 1.25], np.r_[0, -1.25]
tvec = np.linspace(0, 30, 351)

reload(plots)
from plots import phase_plotter as ppr

fig, ax = ppr(afdict, pitch, lims, angle_rng, tvec, ngrid=201,
              nseed=501, nseed_skip=25, quiver=False, skip=10, seed=False,
              timer=True, gamtest=gammas, extrap=None,
              traj=plots.ps_traj_dp5,
              nullcline_x=False, nullcline_z=True,
              fig=None, ax=None)

lab = 'chukar 20 d.p.h., ' + r'$\theta=$0' + u'\u00B0'
ax.text(.05, -1, lab, fontsize=16)

fig.savefig('Figures/figure6ci_vpd0_nullcline_chukar.pdf', transparent=True)

# %% Velocity polar diagram, pitch = 13

pitch = np.deg2rad(13)

fig, ax = ppr(afdict, pitch, lims, angle_rng, tvec, ngrid=201,
              nseed=501, nseed_skip=25, quiver=False, skip=10, seed=False,
              timer=True, gamtest=gammas, extrap=None,
              traj=plots.ps_traj_dp5, fig=None, ax=None)

lab = 'chukar 20 d.p.h., ' + r'$\theta=$13' + u'\u00B0'
ax.text(.05, -1, lab, fontsize=16)

fig.savefig('Figures/figure6cii_vpd13_chukar.pdf', transparent=True)

# %% Velocity polar diagram, pitch = 13 with Z nullcline

pitch = np.deg2rad(13)

fig, ax = ppr(afdict, pitch, lims, angle_rng, tvec, ngrid=201,
              nseed=501, nseed_skip=25, quiver=False, skip=10, seed=False,
              timer=True, gamtest=gammas, extrap=None,
              traj=plots.ps_traj_dp5, nullcline_z=True, fig=None, ax=None)

lab = 'chukar 20 d.p.h., ' + r'$\theta=$13' + u'\u00B0'
ax.text(.05, -1, lab, fontsize=16)

fig.savefig('Figures/figure6cii_vpd13_nullcline_chukar.pdf', transparent=True)

ext: py

reason: import numpy,from scipy


text:
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import random

import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn import metrics
from sklearn import tree
from sklearn.metrics import confusion_matrix
from sklearn.metrics import accuracy_score
from sklearn.ensemble import RandomForestClassifier

from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import cross_val_score

data = pd.read_csv(r"C:\Users\patel\AppData\Local\Programs\Python\Python38\Programs\Assignment Programs\Day 6\Data Set\pacific.csv")

print(data.head(6))

data.Status = pd.Categorical(data.Status)
data['Status'] = data.Status.cat.codes

print(data.head())

random.seed(2)
pred_columns = data[:]
pred_columns.drop(['Status'],axis=1,inplace=True)
pred_columns.drop(['Event'],axis=1,inplace=True)
pred_columns.drop(['Latitude'],axis=1,inplace=True)
pred_columns.drop(['Longitude'],axis=1,inplace=True)
pred_columns.drop(['ID'],axis=1,inplace=True)
pred_columns.drop(['Name'],axis=1,inplace=True)
pred_columns.drop(['Date'],axis=1,inplace=True)
pred_columns.drop(['Time'],axis=1,inplace=True)
prediction_var = pred_columns.columns

print(list(prediction_var))

#...............................................   PREDICTION AND TESTING OF DATA VALUES   .........................................................................

train, test = train_test_split(data, test_size = 0.3,random_state=8)

print("\n \n The shape of the Train Data is    :      ",train.shape)
print("\n \n The shape of the Test Data is    :      ",test.shape)

train_X = train[prediction_var]
train_y=train['Status']

test_X= test[prediction_var]
test_y =test['Status']

#############################################################################################################

model=RandomForestClassifier(n_estimators=100)

model.fit(train_X,train_y)

prediction=model.predict(test_X)

print("\n \n Accuracy of RANDOM FOREST ALGORITHM    is       :          ",metrics.accuracy_score(prediction,test_y))

##############################################################################################################

'''
model = tree.DecisionTreeClassifier()
model.fit(train_X,train_y)

pred_y=model.predict(test_X)

#df=pd.DataFrame(pred_y,test_y)

#print(df)

print("\n \n Accuracy of DECISION TREE ALGORITHM    is       :          ",metrics.accuracy_score(pred_y,test_y))

'''

max_stars_repo_path: Logistic Reg_Random Forest Algo.py

max_stars_repo_name: DevisPatel/Python-Programs

max_issues_repo_path: Logistic Reg_Random Forest Algo.py

max_issues_repo_name: DevisPatel/Python-Programs

max_forks_repo_path: Logistic Reg_Random Forest Algo.py

max_forks_repo_name: DevisPatel/Python-Programs

reason: import numpy


text:
from collections import defaultdict
import functools
import itertools
import math
from math import ceil
import tempfile
from typing import Dict, Iterator, List, NamedTuple, Tuple

from flask import make_response, Markup, render_template
import matplotlib
import matplotlib.pyplot as plt
import numba
import numpy as np
import pathos

from steady_cell_phenotype._util import (
    batcher,
    BinCounter,
    complete_search_generator,
    get_phased_trajectory,
    HashableNdArray,
    process_model_text,
    random_search_generator,
    )

matplotlib.use("agg")

class TrajectoryStatistics(NamedTuple):
    trajectory_length_counts: Dict[HashableNdArray, BinCounter]
    data_counts: Dict[HashableNdArray, np.ndarray]
    means: Dict[HashableNdArray, np.ndarray]
    variances: Dict[HashableNdArray, np.ndarray]

    @classmethod
    def make(cls, num_variables: int):
        return TrajectoryStatistics(
                trajectory_length_counts=defaultdict(lambda: BinCounter()),
                data_counts=defaultdict(lambda: np.zeros(shape=0, dtype=int)),
                means=defaultdict(lambda: np.zeros(shape=(0, num_variables))),
                variances=defaultdict(lambda: np.zeros(shape=(0, num_variables))),
                )

def batch_trajectory_process(batch, update_fn) -> TrajectoryStatistics:
    num_samples = batch.shape[0]
    num_variables = batch.shape[1]

    trajectory_statistics = TrajectoryStatistics.make(num_variables)
    # unpack
    trajectory_length_counts = trajectory_statistics.trajectory_length_counts
    data_counts = trajectory_statistics.data_counts
    means = trajectory_statistics.means
    scaled_variances = trajectory_statistics.variances

    for batch_idx in range(num_samples):
        trajectory, phase_state = get_phased_trajectory(batch[batch_idx, :], update_fn)
        trajectory_length_counts[phase_state].add(len(trajectory))

        # extract arrays, to reduce dict lookup again
        data_count = data_counts[phase_state]
        mean = means[phase_state]
        scaled_variance = scaled_variances[phase_state]

        old_len: int = data_count.shape[0]
        trajectory_len: int = trajectory.shape[0]

        if trajectory_len == 0:
            continue

        # resize, if necessary
        if old_len < trajectory_len:
            # extend data_count
            new_data_count: np.ndarray = np.zeros(shape=trajectory_len, dtype=int)
            new_data_count[:old_len] = data_count
            data_count: np.ndarray = new_data_count
            data_counts[phase_state] = new_data_count
            # extend means
            new_mean: np.ndarray = np.zeros(shape=(trajectory_len, num_variables))
            new_mean[:old_len, :] = mean
            mean = new_mean
            means[phase_state] = new_mean
            # extend scaled variance
            new_scaled_variance: np.ndarray = np.zeros(
                    shape=(trajectory_len, num_variables)
                    )
            new_scaled_variance[:old_len] = scaled_variance
            scaled_variance = new_scaled_variance
            scaled_variances[phase_state] = new_scaled_variance

        # do it once
        reversed_trajectory = trajectory[::-1, :]

        # Welford's online algorithm for mean and variance calculation. See Knuth Vol 2, pg 232
        data_count[:trajectory_len] += 1
        old_mean = np.array(mean[:trajectory_len])  # copy
        mean[:trajectory_len] += \
            (reversed_trajectory - old_mean) / np.expand_dims(data_count[:trajectory_len], axis=-1)
        scaled_variance[:trajectory_len] += (reversed_trajectory - old_mean) * \
                                            (reversed_trajectory - mean[:trajectory_len])

    # need to create a new tuple, as means and scaled variances can get clobbered.
    return TrajectoryStatistics(
            trajectory_length_counts=trajectory_length_counts,
            data_counts=data_counts,
            means=means,
            variances=scaled_variances,
            )

def reducer(init_stats: TrajectoryStatistics,
            new_stats: TrajectoryStatistics) -> TrajectoryStatistics:
    # unpack
    trajectory_length_counts: Dict[HashableNdArray, BinCounter]
    data_counts: Dict[HashableNdArray, np.ndarray]
    means: Dict[HashableNdArray, np.ndarray]
    variances: Dict[HashableNdArray, np.ndarray]
    trajectory_length_counts, data_counts, means, variances = init_stats

    # unpack
    next_trajectory_length_counts: Dict[HashableNdArray, BinCounter]
    next_data_counts: Dict[HashableNdArray, np.ndarray]
    next_means: Dict[HashableNdArray, np.ndarray]
    next_scaled_variances: Dict[HashableNdArray, np.ndarray]
    (
        next_trajectory_length_counts,
        next_data_counts,
        next_means,
        next_scaled_variances,
        ) = new_stats

    for phase_point in next_trajectory_length_counts.keys():
        num_variables = next_means[phase_point].shape[1]
        if phase_point not in trajectory_length_counts:
            # just copy it
            trajectory_length_counts[phase_point] = next_trajectory_length_counts[phase_point]
            data_counts[phase_point] = next_data_counts[phase_point]
            means[phase_point] = next_means[phase_point]
            # remember to scale!
            variances[phase_point] = next_scaled_variances[
                                         phase_point
                                     ] / np.expand_dims(data_counts[phase_point], axis=-1)
        else:
            # merge
            trajectory_length_counts[phase_point].add(
                    next_trajectory_length_counts[phase_point]
                    )
            new_max_len = trajectory_length_counts[phase_point].max

            # data counts
            old_data_count: np.ndarray = data_counts[phase_point]
            next_data_count: np.ndarray = next_data_counts[phase_point]
            new_data_counts = np.zeros(shape=new_max_len, dtype=int)
            new_data_counts[: old_data_count.shape[0]] = old_data_count
            new_data_counts[: next_data_count.shape[0]] += next_data_count

            data_counts[phase_point] = new_data_counts

            # means
            old_mean: np.ndarray = means[phase_point]
            next_mean: np.ndarray = next_means[phase_point]
            new_mean: np.ndarray = np.zeros(
                    shape=(new_max_len, num_variables), dtype=np.float64
                    )
            max_overlap = min(old_mean.shape[0], next_mean.shape[0])
            # weighted average on the overlap, copy over on the tail
            new_mean[:max_overlap] = (
                                             np.expand_dims(old_data_count[:max_overlap], axis=-1)
                                             * old_mean[:max_overlap]
                                             + np.expand_dims(next_data_count[:max_overlap], axis=-1)
                                             * next_mean[:max_overlap]
                                     ) / np.expand_dims(new_data_counts[:max_overlap], axis=-1)
            if old_mean.shape[0] > next_mean.shape[0]:
                new_mean[max_overlap:, :] = old_mean[max_overlap:, :]
            else:
                new_mean[max_overlap:, :] = next_mean[max_overlap:, :]

            means[phase_point] = new_mean

            # scaled variances
            old_variance: np.ndarray = variances[phase_point]
            next_scaled_variance: np.ndarray = next_scaled_variances[phase_point]
            new_variance: np.ndarray = np.zeros(
                    shape=(new_max_len, num_variables), dtype=np.float64
                    )
            # "weighted average" on the overlap, copy over on the tail
            old_proportion = np.expand_dims(
                    old_data_count[:max_overlap] / new_data_counts[:max_overlap], axis=-1
                    )
            next_proportion = np.expand_dims(
                    next_data_count[:max_overlap] / new_data_counts[:max_overlap], axis=-1
                    )
            new_variance[:max_overlap] = (
                    old_proportion * old_variance[:max_overlap]
                    + next_proportion
                    * (next_scaled_variance[:max_overlap] / np.expand_dims(next_data_count[:max_overlap], axis=-1))
                    + ((old_mean[:max_overlap] - next_mean[:max_overlap]) ** 2)
                    * old_proportion
                    * next_proportion)
            if old_variance.shape[0] > next_scaled_variance.shape[0]:
                new_variance[max_overlap:, :] = old_variance[max_overlap:, :]
            else:
                new_variance[max_overlap:, :] = next_scaled_variance[
                                                max_overlap:, :
                                                ] / np.expand_dims(next_data_count[max_overlap:], axis=-1)

            variances[phase_point] = new_variance

    return TrajectoryStatistics(trajectory_length_counts, data_counts, means, variances)

def compute_cycles(*,
                   model_text: str,
                   knockouts: Dict[str, int],
                   continuous: Dict[str, bool],
                   num_iterations: int,
                   visualize_variables: Dict[str, bool]):
    """

    Parameters
    ----------
    model_text
    knockouts
    continuous
    num_iterations
    visualize_variables

    Returns
    -------

    """
    # turn boolean-valued dictionary to list of variables to visualize
    visualized_variables: Tuple[str] = tuple(
            var for var in visualize_variables if visualize_variables[var]
            )

    # create an update function and equation system
    variables, update_fn, equation_system = process_model_text(
            model_text, knockouts, continuous
            )

    # compile and warm up the jitter
    update_fn = numba.jit(update_fn)
    update_fn(np.zeros(len(variables), dtype=np.int64))

    # associate variable names with their index in vectors
    variable_idx: Dict[str, int] = dict(zip(variables, range(len(variables))))

    constants: List[str] = equation_system.constant_variables()
    constants_vals: Dict[str, int] = {
        const: int(equation_system[const]) for const in constants
        }

    # decide if we will perform a complete state space search or not
    state_space_size = 3 ** (len(variables) - len(constants))
    perform_complete_search = state_space_size <= num_iterations

    state_generator: Iterator[np.ndarray]
    if perform_complete_search:
        state_generator = complete_search_generator(
                variables=variables, constants_vals=constants_vals
                )
        num_iterations = state_space_size
    else:
        state_generator = random_search_generator(
                num_iterations=num_iterations,
                variables=variables,
                constants_vals=constants_vals,
                )

    max_threads = max(1, min(pathos.multiprocessing.cpu_count() - 1, math.floor(state_space_size / 1000)))

    batch_generator = batcher(
            state_generator, variables, batch_size=min(1000, num_iterations // max_threads)
            )

    trajectory_length_counts: Dict[HashableNdArray, BinCounter] = dict()
    data_counts: Dict[HashableNdArray, np.ndarray] = dict()
    means: Dict[HashableNdArray, np.ndarray] = dict()
    variances: Dict[HashableNdArray, np.ndarray] = dict()
    with pathos.multiprocessing.ProcessPool(nodes=max_threads) as pool:
        trajectory_length_counts, data_counts, means, variances = functools.reduce(
                reducer,
                pool.uimap(
                        batch_trajectory_process, batch_generator, itertools.repeat(update_fn)
                        ),
                TrajectoryStatistics(
                        trajectory_length_counts, data_counts, means, variances
                        ),
                )

    phase_points = trajectory_length_counts.keys()

    # recreate limit cycles from phase point keys
    limit_cycles: Dict[HashableNdArray, List] = dict()  # record limit sets with their ordering. i.e. as cycles
    for phase_point in phase_points:
        cycle = list()
        t_state = phase_point
        state = phase_point.array
        # do-while
        cycle.append(t_state)
        state = update_fn(state)
        t_state = HashableNdArray(state)
        while t_state != phase_point:
            cycle.append(t_state)
            state = update_fn(state)
            t_state = HashableNdArray(state)
        # record cycle
        limit_cycles[phase_point] = cycle

    # create images for each variable
    limit_set_stats_images = {phase_point: dict() for phase_point in phase_points}
    with tempfile.TemporaryDirectory() as tmp_dir_name:
        plt.rcParams["svg.fonttype"] = "none"
        for phase_point, var in itertools.product(phase_points, visualized_variables):
            cycle = limit_cycles[phase_point]
            var_idx = variable_idx[var]

            var_means = np.flipud(means[phase_point][:, var_idx])
            var_stdevs = np.flipud(np.sqrt(variances[phase_point][:, var_idx]))

            plt.figure(figsize=(6, 4))

            # means, since phased will fix on single value at end
            plt.plot(list(var_means) + [cycle[0][var_idx]], color="#1f3d87")
            plt.fill_between(
                    range(len(var_means) + 1),
                    list(var_means - var_stdevs) + [cycle[0][var_idx]],
                    list(var_means + var_stdevs) + [cycle[0][var_idx]],
                    color="grey",
                    alpha=0.25,
                    )
            # dashed cycle in `converging` region
            plt.plot(
                    range(len(var_means) - len(cycle), len(var_means) + 1),
                    [cycle[idx][var_idx] for idx in range(len(cycle))]
                    + [cycle[0][var_idx]],
                    color="#F24F00",
                    linestyle="dashed",
                    )
            # solid cycle in `converged` region
            plt.plot(
                    range(len(var_means), len(var_means) + len(cycle) + 1),
                    [cycle[idx][var_idx] for idx in range(len(cycle))]
                    + [cycle[0][var_idx]],
                    color="#F24F00",
                    )

            plt.title(f"Phased Envelope of trajectories for {var}")

            plt.xlabel("Distance to phase-fixing point on limit cycle")
            max_tick_goal = 40
            tick_gap = max(1, ceil((len(var_means) + len(cycle) + 1) / max_tick_goal))
            x_tick_locations = [
                x
                for x in range(0, len(var_means) + len(cycle) + 1)
                if (x - len(var_means)) % tick_gap == 0
                ]
            x_tick_labels = list(map(lambda x: len(var_means) - x, x_tick_locations))
            plt.xticks(x_tick_locations, x_tick_labels, rotation=90)
            plt.xlim([0, len(var_means) + len(cycle)])
            plt.axvline(len(var_means), linestyle="dotted", color="grey")
            plt.axvline(len(var_means) - len(cycle), linestyle="dotted", color="grey")

            plt.ylabel("State")
            plt.yticks([0, 1, 2])
            for y_val in range(3):
                plt.axhline(y=y_val, linestyle="dotted", color="grey", alpha=0.5)
            plt.ylim([0 - 0.1, 2 + 0.1])

            plt.tight_layout()

            image_filename = "dist-" + str(hash(phase_point)) + "-" + var + ".svg"
            plt.savefig(
                    tmp_dir_name + "/" + image_filename, transparent=True, pad_inches=0.0
                    )
            plt.close()
            with open(tmp_dir_name + "/" + image_filename, "r") as image:
                limit_set_stats_images[phase_point][var] = Markup(image.read())

    # print('var images complete')

    # cycle list sorted by frequency
    cycle_list = list(limit_cycles.values())
    cycle_list.sort(
            key=lambda cyc: trajectory_length_counts[cyc[0]].total(), reverse=True
            )

    # tally of number of cycles of each length, sorted by cycle length
    cycle_len_counts = defaultdict(lambda: 0)
    for cycle in cycle_list:
        cycle_len_counts[len(cycle)] += 1
    cycle_len_counts = tuple(sorted(cycle_len_counts.items()))

    length_distribution_images = dict()
    with tempfile.TemporaryDirectory() as tmp_dir_name:
        # create length distribution plots
        plt.rcParams["svg.fonttype"] = "none"
        for phase_point in phase_points:
            length_dist = trajectory_length_counts[phase_point].trimmed_bins()
            plt.figure(figsize=(4, 3))
            plt.bar(x=range(len(length_dist)), height=length_dist, color="#002868")
            plt.title("Distribution of lengths of paths")
            plt.xlabel("Length of path")
            plt.ylabel("Number of states")
            plt.tight_layout()
            image_filename = "dist-" + str(hash(phase_point)) + ".svg"
            plt.savefig(
                    tmp_dir_name + "/" + image_filename, transparent=True, pad_inches=0.0
                    )
            plt.close()
            with open(tmp_dir_name + "/" + image_filename, "r") as image:
                length_distribution_images[phase_point] = Markup(image.read())

    def get_key(cyc):
        return cyc[0]

    def get_count(cyc):
        return trajectory_length_counts[get_key(cyc)].total()

    cycles = [
        {
            "states": cycle,
            "len": len(cycle),
            "count": get_count(cycle),
            "percent": 100 * get_count(cycle) / num_iterations,
            "len-dist-image": length_distribution_images[get_key(cycle)]
            if get_key(cycle) in length_distribution_images
            else "",
            "limit-set-stats-images": limit_set_stats_images[get_key(cycle)]
            if get_key(cycle) in limit_set_stats_images
            else dict(),
            }
        for cycle in cycle_list]

    # respond with the results-of-computation page
    return make_response(
            render_template(
                    "compute-cycles.html",
                    cycles=cycles,
                    variables=variables,
                    cycle_len_counts=cycle_len_counts,
                    complete_results=perform_complete_search,
                    )
            )

max_stars_repo_name: knappa/steadycellphenotype

max_issues_repo_name: knappa/steadycellphenotype

max_forks_repo_name: knappa/steadycellphenotype

reason: import numpy,import numba


text:
import matplotlib.pyplot as plt
import numpy as np

euler = False
v_verlet = True

## Standard form of the damped oscillator equation
# accel + a*vel + b*vel^2 + c*pos = d

##Coefficients
a, b, c, d = 0.5, 0.05, 1, 1

## Initial conditions
init_pos = 0
init_vel = 10

dt = 0.1
time = np.arange(0,50+dt,dt)

pos = [init_pos]
vel = [init_vel]

##Euler forward integration is unstable, first order and not symplectic
if euler == True:
    for i in range(len(time)-1):
        new_vel = vel[i] + (-a*vel[i]-b*vel[i]**2-c*pos[i]+d)*dt
        new_pos = pos[i] + vel[i]*dt

        vel.append(new_vel)
        pos.append(new_pos)

## Velocity Verlet is stable for oscillatory motion and is symplectic
## Velocity Verlet is second order accurate
elif v_verlet == True:

    for i in range(len(time)-1):

        #Calculate intermediate velocity (velocity Verlet)
        int_vel = vel[i] + dt/2*(-a*vel[i]-b*vel[i]**2-c*pos[i]+d)
        new_pos = pos[i] + dt*int_vel
        new_vel = int_vel + dt/2*(-a*int_vel-b*int_vel**2-c*new_pos+d)

        vel.append(new_vel)
        pos.append(new_pos)

plt.plot(time,pos,label="Position")
plt.plot(time,vel,label="Velocity")
plt.xlabel("Time")
plt.legend()
plt.savefig("Damped Oscillator a="+str(a)+" b="+str(b)+" c="+str(c)+" d="+str(d)+".pdf")
plt.show()


text:
# -*- coding: utf-8 -*-
"""
Created on Fri Jun 16 04:51:14 2017

@author: gamma
"""

import numpy as np
import pickle

matrixShape = ##Shape##
matrixDType = ##Dtype##
workingDir = ##workingDir##
log = ""

LogFile = open(workingDir+"logfile.txt", 'w')
MatrixDumps = open(workingDir+"Matrix.txt", 'w')

try:
    matrix = np.zeros(matrixShape,dtype = matrixDType)
except Exception as e:
    log+= str(e)+"\n"
try:
    LogFile.write(str(matrix))
except Exception as e:
    log+= str(e)+"\n"

MatrixDumps.write(np.ndarray.dumps(matrix))

max_stars_repo_name: MahmoudAbdelRahman/GhPlotLib

max_issues_repo_name: MahmoudAbdelRahman/GhPlotLib

max_forks_repo_name: MahmoudAbdelRahman/GhPlotLib


text:
import atexit
import contextlib
import glob
import logging
import os
import sys
import tarfile
from collections import defaultdict
from datetime import datetime
import time

import jinja2
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scipy
import scipy.ndimage
import seaborn as sns
import tabulate
from matplotlib.mlab import griddata
from matplotlib.ticker import FormatStrFormatter, LogLocator, MultipleLocator
from mpl_toolkits.axes_grid1 import ImageGrid
from root_numpy import root2array
import ROOT
from scipy.stats import chi2

from NPFit.NPFit import kde
from NPFit.NPFit.makeflow import (fluctuate, max_likelihood_fit, multi_signal,
                                  multidim_grid, multidim_np)
from NPFit.NPFit.nll import fit_nll
from NPFit.NPFit.parameters import conversion, label, nlo
from NPFit.NPFit.scaling import load_fitted_scan
from NPFitProduction.NPFitProduction.cross_sections import CrossSectionScan
from NPFitProduction.NPFitProduction.utils import (cartesian_product,
                                                   sorted_combos)

tweaks = {
    "lines.markeredgewidth": 0.0,
    "lines.linewidth": 8,
    "lines.markersize": 23,
    "patch.edgecolor": "black",
    "legend.facecolor": "white",
    "legend.frameon": True,
    "legend.edgecolor": "white",
    "legend.fontsize": "medium",
    "legend.handletextpad": 0.5,
    "mathtext.fontset": "custom",
    "mathtext.rm": "Bitstream Vera Sans",
    "mathtext.it": "Bitstream Vera Sans:italic",
    "mathtext.bf": "Bitstream Vera Sans:bold",
    "axes.labelsize": "x-large",
    "axes.titlesize": "x-large",
    "xtick.labelsize": "x-large",
    "xtick.major.size": 5,
    "ytick.major.size": 5,
    "xtick.direction": "in",
    "ytick.direction": "in",
    "ytick.labelsize": "x-large",
}
sns.set(context="poster", style="white", font_scale=1.5, rc=tweaks)

x_min, x_max, y_min, y_max = np.array([0.200, 1.200, 0.550, 2.250])

def round(num, sig_figs):
    return str(float('{0:.{1}e}'.format(num, sig_figs - 1)))

def get_masked_colormap(bottom_map, top_map, norm, width, masked_value):
    low = masked_value - width / 2.
    high = masked_value + width / 2.
    if low > norm.vmin:
        colors = zip(np.linspace(0., norm(low), 100), bottom_map(np.linspace(0.1, 1., 100)))
        colors += [(norm(low), 'gray')]
    else:
        colors = [(0., 'gray')]
    if high < norm.vmax:
        colors += [(norm(high), 'gray')]
        colors += zip(np.linspace(norm(high), 1., 100), top_map(np.linspace(0.1, 1., 100)))
    else:
        colors += [(1., 'gray')]

    cmap = matplotlib.colors.LinearSegmentedColormap.from_list('masked_map', colors)

    return cmap

def get_stacked_colormaps(cmaps, interfaces, norm):
    colors = []
    low = 0.
    for cmap, interface in zip(cmaps, interfaces):
        colors += zip(np.linspace(low, norm(interface), 100), cmap(np.linspace(0, 1., 100)))
        low = norm(interface)

    colors += zip(np.linspace(low, 1., 100), cmaps[-1](np.linspace(0, 1., 100)))

    cmap = matplotlib.colors.LinearSegmentedColormap.from_list('masked_map', colors)

    return cmap

class Plotter(object):

    def __init__(self, config):
        atexit.register(self.webify)
        self.config = config

    def webify(self):
        path = os.path.join(os.environ['LOCALRT'], 'src', 'NPFit', 'NPFit', 'data')
        env = jinja2.Environment(loader=jinja2.FileSystemLoader(path))
        env.filters["datetime"] = lambda d: datetime.fromtimestamp(d).strftime('%a, %d %b %Y, %H:%M')
        env.filters["basename"] = lambda d: os.path.basename(d)
        env.tests["sum"] = lambda s: s == "Total"
        template = env.get_template('template.html')

        tfilename = os.path.join(self.config['outdir'], 'plots', 'plots.tar.gz')
        tfile = tarfile.open(tfilename, 'w:gz')
        for root, dirs, files in os.walk(os.path.abspath(self.config['outdir'])):
            for path in [os.path.join(root, d) for d in dirs] + [root]:
                pngs = glob.glob(os.path.join(path, '*.png'))
                pdfs = glob.glob(os.path.join(path, '*.pdf'))
                plots = [os.path.splitext(os.path.basename(name))[0] for name in pngs]
                if plots:
                    plots.sort()
                    [tfile.add(f, arcname=os.path.basename(name)) for f in [os.path.join(path, name) for name in pdfs]]

                dirs = ['..']
                files = []
                names = [x for x in os.listdir(path) if not x.startswith('.') and '#' not in x and x != 'index.html']
                for name in names:
                    if '.png' not in name and '.pdf' not in name:
                        fullname = os.path.join(path, name)
                        if os.path.isdir(fullname):
                            if name not in ['javascripts', 'stylesheets']:
                                dirs.append(name)
                        else:
                            files.append({'name': name, 'timestamp': os.path.getmtime(fullname)})

                with open(os.path.join(path, 'index.html'), 'w') as f:
                    logging.info('updating index {}'.format(os.path.join(path, 'index.html')))
                    f.write(template.render(
                        dirs=sorted(dirs),
                        files=sorted(files),
                        plots=plots,
                    ).encode('utf-8'))

        tfile.close()

    @contextlib.contextmanager
    def saved_figure(self, x_label, y_label, name, header=False, figsize=(11, 11), dpi=None):
        fig, ax = plt.subplots(figsize=figsize)
        lumi = str(self.config['luminosity']) + ' fb$^{-1}$ (13 TeV)'
        if header:
            plt.title(lumi, loc='right', fontweight='normal')
            plt.title(r'CMS', loc='left', fontweight='bold')
            if header == 'preliminary':
                plt.text(0.155, 1.009, r'Preliminary', style='italic', transform=ax.transAxes)

        try:
            yield ax

        finally:
            logging.info('saving {}'.format(name))
            plt.xlabel(x_label, horizontalalignment='right', x=1.0)
            plt.ylabel(y_label, horizontalalignment='right', y=1.0)
            if dpi is None:
                plt.savefig(os.path.join(self.config['outdir'], 'plots', '{}.pdf'.format(name)), bbox_inches='tight')
            else:
                plt.savefig(os.path.join(self.config['outdir'], 'plots', '{}.pdf'.format(name)), bbox_inches='tight',
                        dpi=dpi)
            plt.savefig(os.path.join(self.config['outdir'], 'plots', '{}.png'.format(name)), bbox_inches='tight',
                    dpi=dpi)
            plt.close()

class Plot(object):

    def __init__(self, subdir):
        self.subdir = subdir

    def specify(self):
        """Make all inputs

        This method should add all of the commands producing input files which
        are needed by the Plot to the MakeflowSpecification. This command will
        be run for each plot in the analysis config file. By specifying this
        per-plot, commenting out plots in the config will remove their inputs
        from the Makefile, so that only the needed inputs are produced.

        """
        pass

    def write(self, config):
        """Write the plots

        This command should actually produce and save the plots.
        """
        try:
            os.makedirs(os.path.join(config['outdir'], 'plots', self.subdir))
        except OSError:
            pass  # the directory has already been made

def get_errs(scan, dimension, processes, config, clip=None):
    fits = None
    mgs = None
    for coefficients in sorted_combos(config['coefficients'], dimension):
        for process in processes:
            if process in scan.points[coefficients]:
                fit = scan.evaluate(coefficients, scan.points[coefficients][process], process)
                mg, _ = scan.scales(coefficients, process)

                if fits is None:
                    fits = fit
                    mgs = mg
                else:
                    fits = np.concatenate([fits, fit])
                    mgs = np.concatenate([mgs, mg])
    percent_errs = (mgs - fits) / mgs * 100
    np.random.shuffle(percent_errs)
    percent_errs = percent_errs[:clip]
    avg_abs_percent_errs = sum(np.abs(percent_errs)) / len(mgs[:clip])

    return percent_errs, avg_abs_percent_errs

class FitQualityByPoints(Plot):

    def __init__(self, dimensions=[1], processes=['ttZ', 'ttH', 'ttW'], points=None, subdir='fit'):
        self.dimensions = dimensions
        self.processes = processes
        self.subdir = subdir
        if points is None:
            self.fitpoints = np.array(range(10, 1000, 100))
        else:
            self.fitpoints = np.array(points)

    def specify(self, config, spec, index):
        spec.add(['cross_sections.npz'], [], ['run', 'plot', '--index', str(index), config['fn']])

    def write(self, config, plotter, args):
        super(FitQualityByPoints, self).write(config)

        name = os.path.join(self.subdir, 'fit_quality_by_points')
        y_label = r'$\frac{100}{n}\sum^n_{i=1} |\frac{\mu_{\mathrm{MG}} - \mu_{\mathrm{fit}}}{\mu_{\mathrm{MG}}}|_i$'
        loc = MultipleLocator(5)
        with plotter.saved_figure('fit points', y_label, name,
                figsize=(17, 11)) as ax:
            labels = []
            table = []
            scan = load_fitted_scan(config, 'cross_sections.npz')
            testpoints = dict((dimension, 0) for dimension, _, _ in self.dimensions)
            for dimension, _, _ in self.dimensions:
                for k in scan.points.keys():
                    if len(k) == dimension:
                        testpoints[dimension] += sum([len(v) for v in scan.points[k].values()])
            avg_abs_percent_errs = np.zeros(len(self.fitpoints))

            for dimension, marker, color in self.dimensions:
                for index, points in enumerate(self.fitpoints):
                    scan = load_fitted_scan(config, 'cross_sections.npz', maxpoints=points, dimension=dimension)
                    errs, avg_abs_percent_err = get_errs(scan, dimension, self.processes, config,
                            clip=min(testpoints.values()))

                    table.append([points, avg_abs_percent_err, len(errs), dimension])
                    avg_abs_percent_errs[index] = avg_abs_percent_err

                ax.plot(self.fitpoints, avg_abs_percent_errs, marker=marker, markersize=10, linewidth=1, linestyle='none',
                        label='fit dimension {}'.format(dimension), color=color)
            ax.set_yscale('log', subsy=range(10))
            ax.xaxis.set_minor_locator(loc)
            ax.tick_params(axis='x', length=5, which='minor')
            plt.legend(fontsize='large')

        headers = ['fit points', 'average absolute percent error', 'total test points', 'dimension']
        with open(os.path.join(config['outdir'], 'fit_quality_by_points.txt'), 'w') as f:
            f.write(tabulate.tabulate(table, headers=headers) + '\n')

class FitQualityByDim(Plot):

    def __init__(self, fit_dimensions=[1], eval_dimensions=None, processes=['ttZ', 'ttH', 'ttW'], fit_to_test_ratio=1,
            subdir='fit', xmin=-100, xmax=50):
        self.fit_dimensions = fit_dimensions
        self.eval_dimensions = eval_dimensions
        self.processes = processes
        self.fit_to_test_ratio = float(fit_to_test_ratio)
        self.subdir = subdir
        self.xmin = xmin
        self.xmax = xmax

    def specify(self, config, spec, index):
        spec.add(['cross_sections.npz'], [], ['run', 'plot', '--index', str(index), config['fn']])

    def write(self, config, plotter, args):
        super(FitQualityByDim, self).write(config)

        name = os.path.join(self.subdir, 'fit_quality_by_dim')
        x_label = r'$100 \frac{\mu_{\mathrm{MG}} - \mu_{\mathrm{fit}}}{\mu_{\mathrm{MG}}}$'
        scan = CrossSectionScan(os.path.join(config['outdir'], 'cross_sections.npz'))

        testpoints = dict(('_'.join([str(d) for d in dimensions]), 0) for dimensions in self.fit_dimensions)
        for dimensions in self.fit_dimensions:
            tag = '_'.join([str(d) for d in dimensions])
            for k in scan.points.keys():
                if len(k) in dimensions:
                    testpoints[tag] += sum([len(v) for v in scan.points[k].values()])
        with plotter.saved_figure(x_label, 'counts', name) as ax:
            labels = []
            table = []

            for fit_dims in self.fit_dimensions:
                scan.fit(dimensions=fit_dims)
                for eval_dim in (self.eval_dimensions if self.eval_dimensions is not None else fit_dims):
                    errs, _ = get_errs(scan, eval_dim, self.processes, config, clip=min(testpoints.values()))
                    errs[errs < self.xmin] = self.xmin
                    errs[errs > self.xmax] = self.xmax
                    if self.eval_dimensions is None:
                        label = '{fdim} fit $\sigma$={s:0.2f}'.format(
                            fdim=' and '.join(('{}d'.format(dim) for dim in fit_dims)),
                            s=np.std(errs),
                        )
                    else:
                        label = '{fdim} fit, {edim}d evaluation $\sigma_{{{sfdim},{edim}}}$={s:0.2f}'.format(
                            fdim=' and '.join(('{}d'.format(dim) for dim in fit_dims)),
                            edim=eval_dim,
                            sfdim='+'.join((str(dim) for dim in fit_dims)),
                            s=np.std(errs),
                        )
                    ax.hist(errs, 70, histtype='step', fill=False, label=label, linewidth=4)
                    bad = errs[(errs < -5) | (errs > 5)]
                ax.set_yscale('log', subsy=range(10))
                if self.eval_dimensions is not None:
                    plt.ylim(ymin=0, ymax=2e7)
                    plt.xlim(xmin=-100)
                ax.yaxis.set_tick_params('minor', size=5)
                plt.legend(title='{:,} test points'.format(len(errs)), loc='upper left')

class NewPhysicsScaling2D(Plot):

    def __init__(
            self,
            processes=['ttZ', 'ttH', 'ttW'],
            subdir='scaling2d',
            dimensionless=False,
            dimensions=None,
            maxnll=None,
            match_zwindows=False,
            madgraph=False,
            numvalues=100,
            points=10000,
            dpi=500,
            profile=True):
        self.subdir = subdir
        self.processes = processes
        self.dimensionless = dimensionless
        self.dimensions = dimensions if dimensions is not None else [2]
        self.maxnll = maxnll
        self.match_zwindows = match_zwindows
        self.madgraph = madgraph
        self.numvalues = numvalues
        self.points = points
        self.dpi = dpi
        self.profile = profile

    def specify(self, config, spec, index):
        inputs = []
        if self.maxnll is not None:
            inputs += multidim_np(config, spec, 2, points=self.points)

        for coefficients in sorted_combos(config['coefficients'], 2):
            cmd = 'run plot --coefficients {coefficients} --index {index} {fn}'
            base = os.path.join(config['outdir'], 'plots', self.subdir, '_'.join(coefficients))
            outputs = [base + ext for ext in ['.pdf', '.png']]
            spec.add(inputs + ['cross_sections.npz'], outputs, cmd.format(coefficients=' '.join(coefficients), index=index, fn=config['fn']))
            if self.profile:
                workspace = os.path.join(config['outdir'], 'workspaces', '{}.root'.format('_'.join(config['coefficients'])))
                cmd = ['run', 'combine', '--snapshot'] + list(coefficients) + [config['fn']]
                snapshot = os.path.join(config['outdir'], 'snapshots', '{}.root'.format('_'.join(coefficients)))
                spec.add([workspace], [snapshot], cmd)

    def write(self, config, plotter, args):
        super(NewPhysicsScaling2D, self).write(config)
        scan = CrossSectionScan(os.path.join(config['outdir'], 'cross_sections.npz'))
        scan.fit(dimensions=self.dimensions)

        if self.match_zwindows:
            zmin = None
            zmax = None
            for coefficients in sorted_combos(args.coefficients, 2):
                madgraph = scan.dataframe(coefficients)
                if zmin is None:
                    zmin = min(madgraph[self.processes].min())
                    zmax = max(madgraph[self.processes].max())
                else:
                    zmin = min(min(madgraph[self.processes].min()), zmin)
                    zmax = max(max(madgraph[self.processes].max()), zmax)

        for coefficients in sorted_combos(args.coefficients, 2):
            tag = '_'.join(coefficients)
            name = os.path.join(self.subdir, tag)
            x = coefficients[0]
            y = coefficients[1]

            x_label = label[x] + ('' if self.dimensionless else r' $/\Lambda^2\, / \,\mathrm{TeV}^{-2}$')
            y_label = label[y] + ('' if self.dimensionless else r' $/\Lambda^2\, / \,\mathrm{TeV}^{-2}$')
            x_conv = 1. if self.dimensionless else conversion[x]
            y_conv = 1. if self.dimensionless else conversion[y]

            if coefficients in scan.points.keys():
                madgraph = scan.dataframe(coefficients)
            else:
                madgraph = scan.dataframe(sorted(config['coefficients']))
            xmin = madgraph[x].min()
            xmax = madgraph[x].max()
            ymin = madgraph[y].min()
            ymax = madgraph[y].max()

            if self.maxnll:
                try:
                    data = root2array(os.path.join(config['outdir'], 'scans', '{}.total.root'.format(tag)))
                except IOError as e:
                    print 'input data missing, will not match nll for {}'.format(tag)
                    continue

                zi = 2 * data['deltaNLL']
                xi = data[x]
                yi = data[y]
                xmin = xi[zi < self.maxnll].min()
                ymin = yi[zi < self.maxnll].min()
                xmax = xi[zi < self.maxnll].max()
                ymax = yi[zi < self.maxnll].max()
                window = (madgraph[x] > xmin) & (madgraph[x] < xmax) & (madgraph[y] > ymin) & (madgraph[y] < ymax)
                madgraph = madgraph[window]

            if not self.match_zwindows:
                zmin = min(madgraph[self.processes].min())
                zmax = max(madgraph[self.processes].max())

            norm = matplotlib.colors.LogNorm(vmin=min(zmin, 0.9), vmax=zmax)
            masked_map = get_masked_colormap(
                    sns.light_palette("navy", as_cmap=True),
                    sns.light_palette((210, 90, 60), input="husl", as_cmap=True),
                    norm,
                    0.1,
                    1.0
            )

            fig = plt.figure(figsize=(30, 9))
            grid = ImageGrid(
                fig, 111,
                nrows_ncols=(1, len(self.processes)),
                axes_pad=0.45,
                share_all=True,
                cbar_location="right",
                cbar_mode="single",
                cbar_size="7%",
                cbar_pad=0.15,
                aspect=False
            )

            offplane_coefficients = []
            offplane_values = []
            if self.madgraph:
                df = madgraph
                for c in sorted(config['coefficients']):
                    if c not in coefficients:
                        offplane_coefficients += [c]
                        offplane_values += [0.0]
            else:
                if self.profile:
                    f = ROOT.TFile(os.path.join(config['outdir'], 'snapshots',
                        '{}.root'.format('_'.join(coefficients))))
                    w = f.Get('w')
                    w.loadSnapshot('MultiDimFit')
                values = []
                columns = sorted(config['coefficients'])
                for column in columns:
                    if column in coefficients:
                        if column == x:
                            values += [np.linspace(xmin, xmax, self.numvalues)]
                        if column == y:
                            values += [np.linspace(ymin, ymax, self.numvalues)]
                    else:
                        offplane_coefficients += [column]
                        if self.profile:
                            var = w.var(column)
                            values += [np.array([var.getVal()])]
                            offplane_values += [var.getVal() * (1. if self.dimensionless else conversion[column])]
                        else:
                            values += [np.zeros(1)]
                            offplane_values += [0.]
                start = time.time()
                df = scan.dataframe(columns, evaluate_points=cartesian_product(*values))
                print('took {:.1f} seconds to get df '.format(time.time() - start))

            for ax, process in zip(grid, self.processes):
                columns = list(coefficients) + [process]
                data = df[columns]
                xx = data[x].values * x_conv
                yy = data[y].values * y_conv
                X, Y = np.meshgrid(np.linspace(xmin * x_conv, xmax * x_conv, self.numvalues * 10), np.linspace(ymin *
                    y_conv, ymax * y_conv, self.numvalues * 10))
                Z = scipy.interpolate.griddata((xx, yy), data[process].values, (X, Y), method='nearest')
                msize = 200 if self.madgraph else 25
                marker = 'o' if self.madgraph else 's'
                start = time.time()
                if self.madgraph:
                    scatter = ax.scatter(
                            xx,
                            yy,
                            c=data[process],
                            norm=norm,
                            s=msize,
                            marker=marker,
                            cmap=masked_map,
                            edgecolors='face',
                            rasterized=True
                    )
                else:
                    scatter = ax.pcolormesh(
                            X,
                            Y,
                            Z,
                            rasterized=True,
                            norm=norm,
                            cmap=masked_map)
                if not self.profile:
                    ax.scatter(
                            [0.0],
                            [0.0],
                            c='red',
                            s=600,
                            linewidths=0,
                            marker='*',
                            label='SM'
                    )
                ax.set_ylabel(y_label, horizontalalignment='right', y=1.0)
                ax.set_xlim([xx.min(), xx.max()])
                ax.set_ylim([yy.min(), yy.max()])
                ax.annotate(
                    label[process],
                    xy=(0.5, 0.85),
                    xycoords='axes fraction',
                    horizontalalignment='center',
                    fontsize='x-large',
                    bbox=dict(boxstyle="round,pad=.5", fc="white", ec="none")
                )

            bar = fig.colorbar(
                    scatter,
                    cax=ax.cax,
                    label='$\sigma_{NP+SM} / \sigma_{SM}$',
                    ticks=LogLocator(subs=range(10)),
                )
            # there is bug in this version of matplotlib ignores zorder, so redraw ticklines
            for t in ax.cax.yaxis.get_ticklines():
                ax.cax.add_artist(t)

            ax.legend(fancybox=True)
            if self.dimensionless:
                point = ', '.join([label[c] for c in offplane_coefficients])
            else:
                point = ', '.join(['$\\frac{{{}}}{{\Lambda^2}}$'.format(label[c].replace('$', '')) for c in offplane_coefficients])
            title = fig.suptitle(
                    '({}) = ({}){}'.format(
                        point,
                        ', '.join([round(v, 2) for v in offplane_values]),
                        ' TeV$^{-2}$' if not self.dimensionless else ''),
                    fontsize="x-large"
            )
            title.set_position([.5, 1.05])

            logging.info('saving {}'.format(name))
            ax.set_xlabel(x_label, horizontalalignment='right', x=1.0)
            start = time.time()
            plt.savefig(os.path.join(config['outdir'], 'plots', '{}.pdf'.format(name)), bbox_inches='tight', dpi=self.dpi)
            plt.savefig(os.path.join(config['outdir'], 'plots', '{}.png'.format(name)), bbox_inches='tight',
                    dpi=self.dpi)
            plt.close()
            print('took {:.1f} seconds to save plot '.format(time.time() - start))

class NewPhysicsScaling(Plot):

    def __init__(
            self,
            processes=[('ttZ', '+', '#2fd164')],
            subdir='scaling',
            overlay_result=False,
            dimensionless=False,
            match_nll_window=True,
            points=300):
        self.subdir = subdir
        self.processes = processes
        self.overlay_result = overlay_result
        self.dimensionless = dimensionless
        self.match_nll_window = match_nll_window
        self.points = points

    def specify(self, config, spec, index):
        inputs = ['cross_sections.npz']
        if self.match_nll_window:
            inputs = multidim_np(config, spec, 1, points=self.points)

        for coefficient in config['coefficients']:
            spec.add(inputs, [], ['run', 'plot', '--coefficients', coefficient, '--index', index, config['fn']])

    def write(self, config, plotter, args):
        super(NewPhysicsScaling, self).write(config)
        # FIXME does this work letting WQ transfer?
        scan = CrossSectionScan(os.path.join(config['outdir'], 'cross_sections.npz'))
        if self.match_nll_window:
            nll = fit_nll(config, transform=False, dimensionless=self.dimensionless)

        for coefficient in args.coefficients:
            conv = 1. if self.dimensionless else conversion[coefficient]
            if not np.any([p in scan.points[coefficient] for p, _, _ in self.processes]):
                continue
            with plotter.saved_figure(
                    label[coefficient] + ('' if self.dimensionless else r' $/\Lambda^2$'),
                    '$\sigma_{NP+SM} / \sigma_{SM}$',
                    os.path.join(self.subdir, coefficient)) as ax:

                for process, marker, c in self.processes:
                    x = scan.points[coefficient][process]
                    y, errs = scan.scales(coefficient, process)
                    if self.match_nll_window:
                        xmin = nll[coefficient]['x'][nll[coefficient]['y'] < 13].min()
                        xmax = nll[coefficient]['x'][nll[coefficient]['y'] < 13].max()
                    else:
                        xmin = min(x * conv)
                        xmax = max(x * conv)

                    xi = np.linspace(xmin, xmax, 10000).reshape(10000, 1)
                    yi = scan.evaluate(coefficient, xi, process)
                    ax.plot(xi * conv, yi, color='#C6C6C6')
                    ax.plot(x * conv, y, marker, mfc='none', markeredgewidth=2, markersize=15, label=label[process],
                            color=c)

                if self.overlay_result:
                    colors = ['black', 'gray']
                    for (x, _), color in zip(nll[coefficient]['best fit'], colors):
                        plt.axvline(
                            x=x,
                            ymax=0.5,
                            linestyle='-',
                            color=color,
                            label='Best fit\n {}$={:.2f}$'.format(label[coefficient], x)
                        )
                    for (low, high), color in zip(nll[coefficient]['one sigma'], colors):
                        plt.axvline(
                            x=low,
                            ymax=0.5,
                            linestyle='--',
                            color=color,
                            label='$1 \sigma [{:03.2f}, {:03.2f}]$'.format(low, high)
                        )
                        plt.axvline(
                            x=high,
                            ymax=0.5,
                            linestyle='--',
                            color=color
                        )
                    for (low, high), color in zip(nll[coefficient]['two sigma'], colors):
                        plt.axvline(
                            x=low,
                            ymax=0.5,
                            linestyle=':',
                            color=color,
                            label='$2 \sigma [{:03.2f}, {:03.2f}]$'.format(low, high)
                        )
                        plt.axvline(
                            x=high,
                            ymax=0.5,
                            linestyle=':',
                            color=color
                        )

                plt.xlim(xmin=xmin, xmax=xmax)
                plt.ylim(ymin=0, ymax=3)
                plt.title(r'CMS Simulation', loc='left', fontweight='bold')
                plt.title(r'MG5_aMC@NLO LO', loc='right', size=27)
                ax.legend(loc='upper center')
                if self.match_nll_window:
                    ax.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))

class NLL2D(Plot):

    def __init__(self, subdir='nll2d', dimensionless=False, draw='mesh', maxnll=12, vmin=0.05, points=2000, dpi=400,
            freeze=False):
        self.subdir = subdir
        self.dimensionless = dimensionless
        if draw not in ['mesh', 'scatter', None]:
            raise NotImplementedError('can only draw mesh, scatter, or None')
        self.draw = draw
        self.maxnll = maxnll
        self.vmin = vmin
        self.points = points
        self.dpi = dpi
        self.freeze = freeze

    def specify(self, config, spec, index):
        multidim_np(config, spec, 2, points=self.points, freeze=self.freeze)

        for coefficients in sorted_combos(config['coefficients'], 2):
            label = '{}{}'.format('_'.join(coefficients), '_frozen' if self.freeze else '')
            inputs = [
                config['fn'],
                os.path.join(config['outdir'], 'scans', '{}.total.root'.format(label))
            ]
            cmd = 'run plot --coefficients {coefficients} --index {index} {fn}'
            outputs = [os.path.join(config['outdir'], 'plots', self.subdir, '{}.pdf'.format(label))]
            spec.add(inputs, outputs, cmd.format(coefficients=' '.join(coefficients), index=index, fn=config['fn']))

    def write(self, config, plotter, args):
        super(NLL2D, self).write(config)

        levels = sorted(chi2.isf([0.05, 0.32], 2))
        labels = ['68% CL', '95% CL']
        for coefficients in sorted_combos(args.coefficients, 2):
            tag = '{}{}'.format('_'.join(coefficients), '_frozen' if self.freeze else '')
            try:
                data = root2array(os.path.join(config['outdir'], 'scans', '{}.total.root'.format(tag)))
            except IOError as e:
                print 'input data missing, skipping {}'.format(tag)
                continue

            x = coefficients[0]
            y = coefficients[1]
            zi = 2 * data['deltaNLL']
            zi = zi - zi.min()
            xi = data[x]
            yi = data[y]
            x_label = label[x] + ('' if self.dimensionless else r' $/\Lambda^2$')
            y_label = label[y] + ('' if self.dimensionless else r' $/\Lambda^2$')
            if not self.dimensionless:
                xi *= conversion[x]
                yi *= conversion[y]
                x_label += '$\ [\mathrm{TeV}^{-2}]$'
                y_label += '$\ [\mathrm{TeV}^{-2}]$'

            xmin = xi[zi < self.maxnll].min()
            xmax = xi[zi < self.maxnll].max()
            ymin = yi[zi < self.maxnll].min()
            ymax = yi[zi < self.maxnll].max()
            window = (xi > xmin) & (xi < xmax) & (yi > ymin) & (yi < ymax)

            gridsize = 50
            X, Y = np.meshgrid(
                    np.linspace(xmin, xmax, gridsize),
                    np.linspace(ymin, ymax, gridsize)
            )
            zmin = 0.1
            xi = xi[window]
            yi = yi[window]
            zi = zi[window]
            bf_x = xi[zi.argmin()]
            bf_y = yi[zi.argmin()]

            Z = scipy.interpolate.griddata((xi, yi), zi, (X, Y), method='nearest')
            Z = scipy.ndimage.filters.gaussian_filter(Z, 1.6)
            with plotter.saved_figure(
                    x_label,
                    y_label,
                    os.path.join(self.subdir, tag),
                    header=config['header'],
                    figsize=(15, 11),
                    dpi=self.dpi) as ax:

                draw = self.draw if (x, y) != ('c2G', 'c3G') else 'scatter'
                if draw == 'mesh':
                    contour = plt.tricontour(
                            X.ravel(),
                            Y.ravel(),
                            Z.ravel(),
                            levels,
                            colors=['black', 'black'],
                            linestyles=['--', '-']
                    )
                else:
                    contour = plt.tricontour(
                            xi,
                            yi,
                            zi,
                            levels,
                            colors=['black', 'black'],
                            linestyles=['--', '-']
                    )

                for i, l in enumerate(labels):
                    contour.collections[i].set_label(l)

                plt.plot(
                    bf_x,
                    bf_y,
                    mew=3,
                    marker="x",
                    linestyle='None',
                    color='black',
                    label='best fit',
                    zorder=10
                )

                if self.draw in ['scatter', 'mesh']:
                    np.clip(Z, self.vmin, zi.max(), Z)
                    np.clip(zi, self.vmin, zi.max(), zi)

                    if draw is 'mesh':
                        points = ax.pcolormesh(
                                X,
                                Y,
                                Z,
                            norm=matplotlib.colors.LogNorm(),
                            cmap=get_stacked_colormaps(
                                [sns.light_palette("red", reverse=True, as_cmap=True), sns.light_palette("blue", as_cmap=True)],
                                interfaces=levels[:-1],
                                norm=matplotlib.colors.LogNorm(vmin=self.vmin, vmax=Z.max())
                            ),
                            rasterized=True
                        )
                    elif draw is 'scatter':
                        points = ax.scatter(
                            xi,
                            yi,
                            c=zi,
                            norm=matplotlib.colors.LogNorm(),
                            s=600,
                            marker='s',
                            linewidths=0,
                            cmap=get_stacked_colormaps(
                                [sns.light_palette("red", reverse=True, as_cmap=True), sns.light_palette("blue", as_cmap=True)],
                                interfaces=levels[:-1],
                                norm=matplotlib.colors.LogNorm(vmin=self.vmin, vmax=zi.max())
                            ),
                            rasterized=True
                        )
                    bar = plt.colorbar(
                        points,
                        label='$-2\ \Delta\ \mathrm{ln}\ \mathrm{L}$' + (' (asimov data)' if config['asimov data'] else ''),
                        ticks=LogLocator(subs=range(10))
                    )
                    for t in bar.ax.get_yticklines():
                        bar.ax.add_artist(t)
                ax.legend(fancybox=True, ncol=3)
                plt.ylim(ymin=ymin, ymax=ymax)
                plt.xlim(xmin=xmin, xmax=xmax)
                # ax.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))

class NLL(Plot):

    def __init__(self, subdir='nll', transform=True, dimensionless=False, points=300):
        self.subdir = subdir
        self.transform = transform
        self.dimensionless = dimensionless
        self.points = points

    def specify(self, config, spec, index):
        inputs = multidim_np(config, spec, 1, points=self.points)

        for coefficient in config['coefficients']:
            spec.add(inputs, [], ['run', 'plot', '--coefficient', coefficient, '--index', index, config['fn']])

    def write(self, config, plotter, args):
        super(NLL, self).write(config)
        data = fit_nll(config, self.transform, self.dimensionless)
        scan = CrossSectionScan(os.path.join(config['outdir'], 'cross_sections.npz'))

        for coefficient in args.coefficients:
            info = data[coefficient]
            for p in config['processes']:
                s0, s1, s2 = scan.construct(p, [coefficient])
                if not ((s1 > 1e-5) or (s2 > 1e-5)):
                    continue  # coefficient has no effect on any of the scaled processes
            x_label = '{} {}'.format(info['label'].replace('\ \mathrm{TeV}^{-2}', ''), info['units'])

            with plotter.saved_figure(
                    x_label,
                    '$-2\ \Delta\ \mathrm{ln}\ \mathrm{L}$' + (' (asimov data)' if config['asimov data'] else ''),
                    os.path.join(self.subdir, coefficient),
                    header=config['header']) as ax:
                ax.plot(info['x'], info['y'], color='black')

                for i, (x, y) in enumerate(info['best fit']):
                    if i == 0:
                        plt.axvline(
                            x=x,
                            ymax=0.5,
                            linestyle=':',
                            color='black',
                            label='Best fit',
                        )
                    else:
                        plt.axvline(
                            x=x,
                            ymax=0.5,
                            linestyle=':',
                            color='black',
                        )
                for i, (low, high) in enumerate(info['one sigma']):
                    ax.plot([low, high], [1.0, 1.0], '--', label=r'68% CL' if (i == 0) else '', color='blue')
                for i, (low, high) in enumerate(info['two sigma']):
                    ax.plot([low, high], [3.84, 3.84], '-.', label=r'95% CL' if (i == 0) else '', color='#ff321a')

                ax.legend(loc='upper center')
                ax.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))
                plt.ylim(ymin=0, ymax=12)
                plt.xlim(xmin=info['x'][info['y'] < 13].min(), xmax=info['x'][info['y'] < 13].max())
                if info['transformed']:
                    plt.xlim(xmin=0)

def two_signal_best_fit(config, ax, signals, theory_errors, tag, contours):
    # TODO switch this to mixin with to signal and contour method
    limits = root2array(os.path.join(config['outdir'], 'best-fit-{}.root'.format(tag)))

    x = limits['r_{}'.format(signals[0])] * nlo[signals[0]]
    y = limits['r_{}'.format(signals[1])] * nlo[signals[1]]
    z = 2 * limits['deltaNLL']

    if contours:
        levels = {
            2.30: '  1 $\sigma$',
            5.99: '  2 $\sigma$',
            # 11.83: ' 3 $\sigma$',
            # 19.33: ' 4 $\sigma$',
            # 28.74: ' 5 $\sigma$'
        }

        xi = np.linspace(x_min, x_max, 1000)
        yi = np.linspace(y_min, y_max, 1000)
        zi = griddata(x, y, z, xi, yi, interp='linear')

        cs = plt.contour(xi, yi, zi, sorted(levels.keys()), colors='black', linewidths=2)
        plt.clabel(cs, fmt=levels)

    handles = []
    labels = []

    bf, = plt.plot(
        x[z.argmin()],
        y[z.argmin()],
        color='black',
        mew=3,
        markersize=17,
        marker="*",
        linestyle='None'
    )
    handles.append(bf)
    labels.append('2D best fit')

    if theory_errors:
        x_process = signals[0]
        y_process = signals[1]
        xerr_low, xerr_high = np.array(theory_errors[y_process]) * nlo[y_process]
        yerr_low, yerr_high = np.array(theory_errors[x_process]) * nlo[x_process]
        theory = plt.errorbar(
            nlo[x_process], nlo[y_process],
            yerr=[[xerr_low], [xerr_high]],
            xerr=[[yerr_low], [yerr_high]],
            capsize=5,
            mew=2,
            color='black',
            ls='',
            marker='o',
            markersize=10,
            linewidth=3
        )
        handles.append(theory)
        labels.append('{} theory\n[1610.07922]'.format(label['ttV']))

    ax.set_xlim([x_min, x_max])
    ax.set_autoscalex_on(False)

    return handles, labels

class TwoProcessCrossSectionSM(Plot):

    def __init__(self, subdir='.', signals=['ttW', 'ttZ'], theory_errors=None, tag=None, numpoints=500, chunksize=500, contours=True):
        self.subdir = subdir
        self.signals = signals
        self.theory_errors = theory_errors
        if tag:
            self.tag = tag
        else:
            self.tag = '-'.join(signals)
        self.numpoints = numpoints
        self.chunksize = chunksize
        self.contours = contours

    def specify(self, config, spec, index):
        inputs = multi_signal(self.signals, self.tag, spec, config)
        for signal in self.signals:
            inputs += max_likelihood_fit(signal, spec, config)
        if self.contours:
            inputs += multidim_grid(config, self.tag, self.numpoints, self.chunksize, spec)

        spec.add(inputs, [],  ['run', 'plot', '--index', index, config['fn']])

    def write(self, config, plotter, args):
        x = self.signals[0]
        y = self.signals[1]
        with plotter.saved_figure(
                label['sigma {}'.format(x)],
                label['sigma {}'.format(y)],
                self.tag,
                header=config['header']) as ax:
            handles, labels = two_signal_best_fit(config, ax, self.signals, self.theory_errors, self.tag, self.contours)

            data = root2array(os.path.join(config['outdir'], 'best-fit-{}.root'.format(x)))

            x_cross_section = plt.axvline(x=data['limit'][0] * nlo[x], color='black')
            x_error = ax.axvspan(
                data['limit'][1] * nlo[x],
                data['limit'][2] * nlo[x],
                alpha=0.5,
                color='#FA6900',
                linewidth=0.0
            )
            handles.append((x_cross_section, x_error))
            labels.append('{} 1D $\pm$ $1\sigma$'.format(label[x]))

            data = root2array(os.path.join(config['outdir'], 'best-fit-{}.root'.format(y)))

            y_cross_section = plt.axhline(y=data['limit'][0] * nlo[y], color='black')
            y_error = ax.axhspan(
                data['limit'][1] * nlo[y],
                data['limit'][2] * nlo[y],
                color='#69D2E7',
                alpha=0.5,
                linewidth=0.0
            )
            handles.append((y_cross_section, y_error))
            labels.append('{} 1D $\pm$ $1\sigma$'.format(label[y]))

            plt.legend(handles, labels)

class TwoProcessCrossSectionSMAndNP(Plot):

    def __init__(self,
             subdir='.',
             signals=['ttW', 'ttZ'],
             theory_errors=None,
             tag=None,
             transform=True,
             dimensionless=False,
             points=300):
        self.subdir = subdir
        self.signals = signals
        self.theory_errors = theory_errors
        if tag:
            self.tag = tag
        else:
            self.tag = '-'.join(signals)
        self.transform = transform
        self.dimensionless = dimensionless
        self.points = points

    def specify(self, config, spec, index):
        inputs = multi_signal(self.signals, self.tag, spec, config)
        inputs += multidim_np(config, spec, 1, points=self.points)
        inputs += fluctuate(config, spec)

        spec.add(inputs, [],  ['run', 'plot', '--index', index, config['fn']])

    def write(self, config, plotter, args):
        x_proc = self.signals[0]
        y_proc = self.signals[1]
        nll = fit_nll(config, self.transform, self.dimensionless)

        table = []
        scales = ['r_{}'.format(x) for x in config['processes']]
        for coefficient in args.coefficients:
            data = np.load(os.path.join(config['outdir'], 'fluctuations-{}.npy'.format(coefficient)))[()]
            if np.isnan(data['x_sec_{}'.format(x_proc)]).any() or np.isnan(data['x_sec_{}'.format(y_proc)]).any():
                print 'skipping coefficient {} with nan fluctuations'.format(coefficient)
                continue

            with plotter.saved_figure(
                    label['sigma {}'.format(x_proc)],
                    label['sigma {}'.format(y_proc)],
                    os.path.join(self.subdir, '{}_{}'.format(self.tag, coefficient)),
                    header=config['header']) as ax:
                handles, labels = two_signal_best_fit(config, ax, self.signals, self.theory_errors, self.tag, False)

                x = data['x_sec_{}'.format(x_proc)]
                y = data['x_sec_{}'.format(y_proc)]

                try:
                    kdehist = kde.kdehist2(x, y, [70, 70])
                    clevels = sorted(kde.confmap(kdehist[0], [.6827, .9545]))
                    contour = ax.contour(kdehist[1], kdehist[2], kdehist[0], clevels, colors=['#ff321a', 'blue'], linestyles=['-.', '--'])
                    for handle, l in zip(contour.collections[::-1], ['68% CL', '95% CL']):
                        handles.append(handle)
                        labels.append(l)
                except Exception as e:
                    print 'problem making contour for {}: {}'.format(coefficient, e)

                colors = ['black', 'gray']
                for (bf, _), color in zip(nll[coefficient]['best fit'], colors):
                    table.append([coefficient, '{:.2f}'.format(bf), '{:.2f}'.format(data[0][coefficient])] + ['{:.2f}'.format(data[0][i]) for i in scales])
                    point, = plt.plot(
                        data[0]['x_sec_{}'.format(x_proc)],
                        data[0]['x_sec_{}'.format(y_proc)],
                        color=color,
                        markeredgecolor=color,
                        mew=3,
                        markersize=17,
                        marker="x",
                        linestyle='None'
                    )
                    handles.append(point)

                    labels.append("Best fit\n{}".format(nll[coefficient]['label']))
                plt.legend(handles, labels, loc='upper right', fontsize=27)
                plt.ylim(ymin=y_min, ymax=y_max)
                plt.xlim(xmin=x_min, xmax=x_max)
                ax.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))
                ax.xaxis.set_major_formatter(FormatStrFormatter('%.1f'))

        print tabulate.tabulate(table, headers=['coefficient', 'bf', 'coefficient value'] + config['processes'])

def plot(args, config):

    plotter = Plotter(config)

    if not args.coefficients:
        args.coefficients = config['coefficients']

    if args.header:
        config['header'] = args.header

    if args.index is not None:
        config['plots'][args.index].write(config, plotter, args)
    else:
        for p in config['plots']:
            p.write(config, plotter, args)

max_stars_repo_name: annawoodard/EffectiveTTV

max_issues_repo_name: annawoodard/NPFit

max_forks_repo_name: annawoodard/NPFit

reason: import numpy,import scipy,from scipy


text:
import pandas as pd
import numpy as np
import cPickle as pickle

from sklearn import preprocessing
from sklearn.feature_extraction import DictVectorizer

from sklearn_pandas import DataFrameMapper

from param import config
# Gather some features
def build_features(features, data):
    # remove NaNs
    data.fillna(0, inplace=True)
    data.loc[data.Open.isnull(), 'Open'] = 1
    # Use some properties directly
    features.extend(['Store', 'CompetitionDistance', 'CompetitionOpenSinceMonth',
                     'CompetitionOpenSinceYear', 'Promo', 'Promo2', 'Promo2SinceWeek', 'Promo2SinceYear'])

    # add some more with a bit of preprocessing
    features.append('SchoolHoliday')
    data['SchoolHoliday'] = data['SchoolHoliday'].astype(float)
    #
    #features.append('StateHoliday')
    #data.loc[data['StateHoliday'] == 'a', 'StateHoliday'] = '1'
    #data.loc[data['StateHoliday'] == 'b', 'StateHoliday'] = '2'
    #data.loc[data['StateHoliday'] == 'c', 'StateHoliday'] = '3'
    #data['StateHoliday'] = data['StateHoliday'].astype(float)

    features.append('DayOfWeek')
    features.append('month')
    features.append('day')
    features.append('year')
    data['year'] = data.Date.apply(lambda x: x.split('-')[0])
    data['year'] = data['year'].astype(float)
    data['month'] = data.Date.apply(lambda x: x.split('-')[1])
    data['month'] = data['month'].astype(float)
    data['day'] = data.Date.apply(lambda x: x.split('-')[2])
    data['day'] = data['day'].astype(float)

    features.append('StoreType')
    data.loc[data['StoreType'] == 'a', 'StoreType'] = '1'
    data.loc[data['StoreType'] == 'b', 'StoreType'] = '2'
    data.loc[data['StoreType'] == 'c', 'StoreType'] = '3'
    data.loc[data['StoreType'] == 'd', 'StoreType'] = '4'
    data['StoreType'] = data['StoreType'].astype(float)

    features.append('Assortment')
    data.loc[data['Assortment'] == 'a', 'Assortment'] = '1'
    data.loc[data['Assortment'] == 'b', 'Assortment'] = '2'
    data.loc[data['Assortment'] == 'c', 'Assortment'] = '3'
    data['Assortment'] = data['Assortment'].astype(float)

def extract_feature(path, train, test, type, feat_names):
    y_train = train[config.target].values
    train.drop(config.target, axis=1, inplace=True)
    y_test = [0] * len(test.index)
    if type == "valid":
        y_test = test[config.target].values
        test.drop(config.target, axis=1, inplace=True)

    if feat_names.count("stand") > 0:
        feat = "stand"
        train_s = train.copy()
        test_s = test.copy()

        with open("%s/train.%s.feat.pkl" %(path, feat), "wb") as f:
            pickle.dump([train_s, y_train], f, -1)
        with open("%s/%s.%s.feat.pkl" %(path, type, feat), "wb") as f:
            pickle.dump([test_s, y_test], f, -1)

if __name__ == "__main__":
    # load data
    train = pd.read_csv(config.origin_train_path, index_col=0)
    test = pd.read_csv(config.origin_test_path, index_col=0)

    features = []
    build_features(features, train)
    build_features([], test)
    features.append('Sales')
    train = train[features]
    test = test[features]

    ## load kfold object
    with open("%s/fold.pkl" % config.data_folder, 'rb') as i_f:
        skf = pickle.load(i_f)

    # extract features
    feat_names = config.feat_names

    # for cross validation
    print "Extract feature for cross validation"
    for iter in range(config.kiter):
        for fold, (trainInd, validInd) in enumerate(skf[iter]):
            path = "%s/iter%d/fold%d" %(config.data_folder, iter, fold)
            sub_train = train.iloc[trainInd].copy()
            sub_val = train.iloc[validInd].copy()
            # extract feature
            extract_feature(path, sub_train, sub_val, "valid", feat_names)
    print "Done"

    # for train/test
    print "Extract feature for train/test"
    path = "%s/all" % config.data_folder
    extract_feature(path, train, test, "test", feat_names)
    print "Done"

max_stars_repo_name: Keesiu/meta-kaggle

max_issues_repo_name: Keesiu/meta-kaggle

max_forks_repo_name: Keesiu/meta-kaggle