# -*- coding: utf-8 -*-
"""
Created on Sun Jan 23 00:27:00 2022

@author: marco
"""
import pandas as pd
import numpy as np
import os
from scipy.linalg import pinv as inv
import matplotlib.pyplot as plt
os.chdir('C://Users//marco//Desktop//Projects//Bayesian_OLS')
cwd = os.getcwd()
print("Current working directory: {0}".format(cwd))
import warnings                                  # `do not disturbe` mode
warnings.filterwarnings('ignore')

#y = Puller.Banxico(serie="SR16734", name="IGAE", plot=False)
#p = Puller.Banxico(serie="SP1", name="Inflation", plot=False)
#r = Puller.Banxico(serie="SF3270", name="Interest_rate", plot=False)
#m = Puller.Banxico(serie="SF1", name="Money", plot=False)
#df = pd.concat([y, p, r, m], axis=1).dropna()
df = pd.read_excel('inflation.xls', index_col=0)

Y=df.iloc[:,0:1].to_numpy()
X=df.iloc[:,1:].to_numpy()

def plot_dist(df, column, title):
    mu = df[column].mean()
    sigma = df[column].std()
    n, bins, patches = plt.hist(x=df[column], bins='auto',
                                color='#0504aa',
                                alpha=0.7, rwidth=0.85)
    plt.grid(axis='y', alpha=0.75)
    plt.ylabel(None)
    plt.title(title)
    maxfreq = n.max()
    plt.ylim(ymax=np.ceil(maxfreq / 10) * 10 if (maxfreq % 10 > 0) else maxfreq + 10)
    plt.show() #plot

def ar_companion_matrix(beta):
    # dont include constant
    k = beta.shape[0]-1
    FF = np.zeros((k, k))
    #insert identity matrix
    FF[1:k, 0:(k-1)] = np.eye(N=k-1, M=k-1)
    temp = (beta[1:k+1, :]).T
    #state space companion form
    #Insert coeffcients along top row
    FF[0:1,0:k+1] = temp
    return(FF)

def gibbs(X,Y,reps,burn,t0,d0,plot):
    reps = reps # number of Gibbs iterations
    burn = burn # percent of burn-in iterations
    out  = np.zeros((reps, X.shape[1]+1))
    t1  = Y.shape[0]  #number of observations
    b0 =  np.zeros((X.shape[1],1))#Priors
    sigma0 = np.eye((X.shape[1])) # variance matrix
    # priors for sigma2
    t0 = t0
    d0 = d0
    # Starting values
    B = b0
    sigma2 = 1
    for i in range(0,reps):
        M = inv(inv(sigma0) + (1/sigma2) * X.T @ X) @ (inv(sigma0) @ b0 + (1/sigma2)* X.T @ Y)
        V = inv(inv(sigma0) + (1/sigma2) * X.T @ X)
        chck = -1
        while (chck < 1):
            B = M + (np.random.normal(0,1,X.shape[1]) @ np.linalg.cholesky(V)).T.reshape(-1,1)
            b = ar_companion_matrix(B)
            ee = np.max(np.abs(np.linalg.eig(b)[1]))
            if (ee <= 1):
                chck = 1
        # compute residuals
        resids = Y - X @ B
        T2 = t0 + t1
        D1 = d0 + resids.T @ resids
        # keep samples after burn period
        out[i,] = np.append(B.T,sigma2)
        #draw from Inverse Gamma
        z0 = np.random.normal(1,1,t1)
        z0z0 = z0.T @ z0
        sigma2 = D1/z0z0

    out = pd.DataFrame(out[burn:reps,:])

    if plot==True:
        for i in range(0,out.shape[1]):
            if i != out.shape[1]-1:
                plot_dist(df=out, column=[i], title='Estimator distribution of beta ' + str(i))
            else:
                plot_dist(df=out, column=[i], title='Estimator distribution of the variance')

    return(out)

gibbs(X=X,Y=Y,reps=5000,burn=4000,t0 = 1,d0 = 0.1, plot=True)

'''
gibbs_sampler <- function(X,Y,B0,sigma0,sigma2,theta0,D0,reps,out,out1){
  for (i in 1:reps){

    M = solve(solve(sigma0)
              +as.numeric(1/sigma2)*t(X)%*%X)%*%(solve(sigma0)%*%b0
              +as.numeric(1/sigma2)*t(X)%*%Y)
    V = solve(solve(sigma0)+as.numeric(1/sigma2)*t(X)%*%X)
    chck=-1
    while (chck < 1){ #check for stability

      B= M+t(rnorm(ncol(X))%*%chol(V))
      b = ar_companion_matrix(B)
      ee <- max(sapply(eigen(b)$values,abs))
      if(ee<=1){
      }
    }
    # compute residuals
    resids <- Y- X%*%B
    T2 = t0 + t1
    D1 = d0 + t(resids) %*% resids

    # keeps samples after burn period
    out[i,] <- t(matrix(c(t(B),sigma2)))

    #draw from Inverse Gamma
    z0 = rnorm(t1,1)
    z0z0 = t(z0) %*% z0
    sigma2 = D1/z0z0

    # compute 2 year forecasts
    yhat = rep(0,h)
    end = as.numeric(length(Y))
    #yhat[1:2] = Y[(end-1):end,]
    cfactor = sqrt(sigma2)
    X_mat = c(1,rep(0,ncol(X)-1))

    for(m in ncol(X):h){
      for (lag in 1:(ncol(X)-1)){
        #create X matrix with p lags
        X_mat[(lag+1)] = yhat[m-lag]
      }
      # Use X matrix to forecast yhat
      yhat[m] = X_mat %*% B + rnorm(1) * cfactor
    }

    out1[i,] <- yhat
  }
  return = list(out,out1)
}

# Set the parameters
reps = 5000 # number of Gibbs iterations
burn = 4000 # percent of burn-in iterations

# Forecast horizon
h = 24

# Matrix to store posterior coefficients and forecasts
out  = matrix(0, nrow = reps, ncol = ncol(X_BO) + 1)
out1 = matrix(0, nrow = reps, ncol = h)

t1 <- nrow(Y_CA) #number of observations
b0 = matrix(0,ncol(X_CA),1) #Priors
sigma0 <- diag(ncol(X_CA)) # variance matrix
# priors for sigma2
t0= 1
d0=0.1

# Starting values
B = b0
sigma2 = 1
df1 = gibbs_sampler(X_CA,Y_CA, B,sigma0,sigma2,t0,d0,reps,out,out1)

coef <- results[[1]][(burn+1):reps,]
forecasts <- results[[2]][(burn+1):reps,]

const <- mean(coef[,1])
beta1 <- mean(coef[,2])
beta2 <- mean(coef[,3])
beta3 <- mean(coef[,4])
sigma <- mean(coef[,5])


<gh_stars>1-10
import numpy as np
import infomap
import pandas as pd
import scipy
from tqdm import tqdm
import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from scipy.signal import lfilter
import warnings

from tqdm import tqdm_notebook
from functools import reduce
warnings.filterwarnings('ignore')

def cluster_counts(labels,features):
    cluster_indexes=[]
    for c in np.unique(labels):
        cluster_indexes.append(np.where(np.array(labels)==c))
    clusters_features=[]
    for cluster in cluster_indexes:
      # print(cluster)
      cluster_features=[features[j] for j in cluster[0] ]
      clusters_features.append(cluster_features)
    return cluster_indexes,clusters_features

def distances_calculations(features_vectors):
    distances=[]
    n=15
    b = [1.0 / n] * n
    a = 1
    for vector1 in tqdm(features_vectors):
        distance=[]
        for vector2 in features_vectors:
            distance.append(scipy.linalg.norm(vector1-vector2,2))
        distances.append(lfilter(b,a,distance))
    return distances
def weights_init(distances_total):
  weights=[]
  for dist in tqdm(distances_total):
      w=[]
      for d in (dist):
          w.append(1/(0.01+d))
      weights.append(w)
  return weights

import pickle
def save_pkl(variable, name):
    name = name + '.pkl'
    output = open(name, 'wb')
    pickle.dump(variable, output)
    output.close()
def findCommunities(G):
    """
    Partition network with the Infomap algorithm.
    Annotates nodes with 'community' id.
    """

    im = infomap.Infomap("--two-level")

    print("Building Infomap network from a NetworkX graph...")
    for source, target in G.edges:
        im.add_link(source, target)

    print("Find communities with Infomap...")
    im.run()

    print(f"Found {im.num_top_modules} modules with codelength: {im.codelength}")

    communities = im.get_modules()
    nx.set_node_attributes(G, communities, 'community')

#create function for  clearence inside the cluster, by euclidian distance
def new_weights(weights):
  weights_new=[]
  for weight in weights:
    w=weight.copy()
    w=np.array(w)
    w[w<(np.median(w)+np.std(w))]=0
    weights_new.append(w)
  return weights_new

def cluster_clearence(cluster_feature_vector:list):
    print(len(cluster_feature_vector))
    distance=[]
    for vector_1 in tqdm_notebook(cluster_feature_vector):
        dist=[]
        for vector_2 in cluster_feature_vector:
            dist.append(np.linalg.norm(vector_1-vector_2,2))
        dist=np.mean(dist)
        distance.append(dist)
    return np.mean(distance)

def cluster_creation(distances:list,features_all:list,features:list,inter:int,treshold:float):
  final_clusters=[]

  weights=weights_init(distances)
  weights_new=new_weights(weights)

  weights_matrix=np.matrix(weights_new)
  Graph=nx.DiGraph(weights_matrix)
  findCommunities(Graph)
  communities = [v for k,v in nx.get_node_attributes(Graph, 'community').items()]
 # print('number of clusters:', np.unique(communities))
  clusters,cluster_features=cluster_counts(communities,features_all)

  # clusters=[cluster[0] for cluster in clusters ]
  # vectors=[distances for i in range(len(clusters))]
  distance=list(map(cluster_clearence,cluster_features))
  if np.isnan(np.array(distance)).any() == True:
    print('DETECTED NAN')
    a=np.argwhere(np.isnan(distance))
   # print(a)
    distance.pop(a[0][0])
    cluster_features.pop(a[0][0])

  print('clearence of clusters',distance)

  bad_clusters=[]
  for i in range(len(distance)):

    if distance[i]<=treshold:
      final_clusters.append(cluster_features[i])
    else:
      bad_clusters.append(cluster_features[i])
  cluster_length=len(bad_clusters)

  i=0
  # for i in range(len(bad_clusters)):
  final_bad_clusters=[]
  while i<cluster_length:
    #print('ITERATION '+str(i),'Total quantaty of clusters:',len(bad_clusters))
    #print("___"*20)
    small_features=bad_clusters[i]
    small_vectors=distances_calculations(small_features)

    weights=weights_init(small_vectors)
    weights_new=new_weights(weights)

    weights_matrix=np.matrix(weights_new)

    Graph=nx.DiGraph(weights_matrix)
    findCommunities(Graph)
    communities = [v for k,v in nx.get_node_attributes(Graph, 'community').items()]
    #print('number of mini clusters in cluster '+str(i), np.unique(communities))

    clusters,cluster_features=cluster_counts(communities,small_features)

    # clusters=[cluster[0] for cluster in clusters ]
    # vectors=[small_vectors for vec in range(len(clusters))]
    distance=list(map(cluster_clearence,cluster_features))
    if np.isnan(np.array(distance)).any() == True:
      print('DETECTED NAN')
      a=np.argwhere(np.isnan(distance))
      distance.pop(a[0][0])
      cluster_features.pop(a[0][0])
    print('clearence of mini clusters',distance)

    for m in range(len(distance)):
      if distance[m]<=treshold:
        final_clusters.append(cluster_features[m])
      elif i<=inter and distance[m]>treshold:
        bad_clusters.append(cluster_features[m])

      elif i>inter and distance[m]>treshold:
        final_bad_clusters.append(cluster_features[m])

    cluster_length=len(bad_clusters)
    print(cluster_length)
    i=i+1

  bad_clusters_vectors=np.concatenate(final_bad_clusters)
  bad_vectors=distances_calculations(bad_clusters_vectors)

  weights=weights_init(bad_vectors)
  weights_new=new_weights(weights)

  weights_matrix=np.matrix(weights_new)
  Graph=nx.DiGraph(weights_matrix)
  findCommunities(Graph)
  communities = [v for k,v in nx.get_node_attributes(Graph, 'community').items()]

 # print('number of mini clusters in cluster '+str(i), np.unique(communities))

  clusters,cluster_features=cluster_counts(communities,bad_clusters_vectors)

  bad_clusters=[]
  for i in range(len(distance)):

    if distance[i]<=treshold:
      final_clusters.append(cluster_features[i])
    else:
      bad_clusters.append(cluster_features[i])
  cluster_length=len(bad_clusters)
  i=0
 # print('CLUSTERS from bad')
  final_bad_clusters=[]
  while i<cluster_length:
   # print('ITERATION '+str(i),'Total quantaty of clusters from bad cluster:',len(bad_clusters))
    #print("___"*20)
    small_features=bad_clusters[i]
    small_vectors=distances_calculations(small_features)

    weights=weights_init(small_vectors)
    weights_new=new_weights(weights)

    weights_matrix=np.matrix(weights_new)
    Graph=nx.DiGraph(weights_matrix)
    findCommunities(Graph)
    communities = [v for k,v in nx.get_node_attributes(Graph, 'community').items()]
    #print('number of mini clusters in cluster '+str(i), np.unique(communities))

    clusters,cluster_features=cluster_counts(communities,small_features)

    # clusters=[cluster[0] for cluster in clusters ]
    # vectors=[small_vectors for vec in range(len(clusters))]
    distance=list(map(cluster_clearence,cluster_features))
    if np.isnan(np.array(distance)).any() == True:
      #print('DETECTED NAN')
      a=np.argwhere(np.isnan(distance))
      distance.pop(a[0][0])
      cluster_features.pop(a[0][0])
    #print('clearence of mini clusters from bad clusters',distance)

    for m in range(len(distance)):
      if distance[m]<=treshold:
        final_clusters.append(cluster_features[m])
      elif i<=inter and distance[m]>treshold:
        bad_clusters.append(cluster_features[m])

      elif i>inter and distance[m]>treshold:
        final_bad_clusters.append(cluster_features[m])

    cluster_length=len(bad_clusters)
    # print(cluster_length)
    i=i+1

  return final_clusters,final_bad_clusters

def final_cluster(features,cluster_feature):
  final_clusters=[]
  for feature_vector in cluster_feature:
    cluster=[]
    for vector in feature_vector:
      bc=np.bincount(np.where(features==vector)[0])
      cluster.append(bc.argmax())
    final_clusters.append(cluster)
  return final_clusters


"""
Filename: oligopoly.py
Authors: <NAME>, <NAME>, <NAME>
This is an example for the lecture dyn_stack.rst from the QuantEcon
series of lectures by <NAME> and <NAME>.
We deal with a large monopolistic firm who faces costs:
C_t = e Q_t + .5 g Q_t^2 + .5 c (Q_{t+1} - Q_t)^2
where the fringe firms face:
sigma_t = d q_t + .5 h q_t^2 + .5 c (q_{t+1} - q_t)^2
Additionally, there is a linear inverse demand curve of the form:
p_t = A_0 - A_1 (Q_t + \bar{q_t}) + \eta_t,
where:
.. math
    \eta_{t+1} = \rho \eta_t + C_{\varepsilon} \varepsilon_{t+1};
    \varepsilon_{t+1} \sim N(0, 1)
For more details, see the lecture.
"""
import numpy as np
import scipy.linalg as la
from quantecon import LQ
from quantecon.matrix_eqn import solve_discrete_lyapunov
from scipy.optimize import root

def setup_matrices(params):
    """
    This function sets up the A, B, R, Q for the oligopoly problem
    described in the lecture.
    Parameters
    ----------
    params : Array(Float, ndim=1)
        Contains the parameters that describe the problem in the order
        [a0, a1, rho, c_eps, c, d, e, g, h, beta]
    Returns
    -------
    (A, B, Q, R) : Array(Float, ndim=2)
        These matrices describe the oligopoly problem.
    """

    # Left hand side of (37)
    Alhs = np.eye(5)
    Alhs[4, :] = np.array([a0-d, 1., -a1, -a1-h, c])
    Alhsinv = la.inv(Alhs)

    # Right hand side of (37)
    Brhs = np.array([[0., 0., 1., 0., 0.]]).T
    Arhs = np.eye(5)
    Arhs[1, 1] = rho
    Arhs[3, 4] = 1.
    Arhs[4, 4] = c / beta

    # R from equation (40)
    R = np.array([[0., 0., (a0-e)/2., 0., 0.],
                  [0., 0., 1./2., 0., 0.],
                  [(a0-e)/2., 1./2, -a1 - .5*g, -a1/2, 0.],
                  [0., 0., -a1/2, 0., 0.],
                  [0., 0., 0., 0., 0.]])

    Rf = np.array([[0., 0., 0., 0., 0., (a0-d)/2.],
                  [0., 0., 0., 0., 0., 1./2.],
                  [0., 0., 0., 0., 0., -a1/2.],
                  [0., 0., 0., 0., 0., -a1/2.],
                  [0., 0., 0., 0., 0., 0.],
                  [(a0-d)/2., 1./2., -a1/2., -a1/2., 0., -h/2.]])

    Q = np.array([[c/2]])

    A = Alhsinv.dot(Arhs)
    B = Alhsinv.dot(Brhs)

    return A, B, Q, R, Rf

def find_PFd(A, B, Q, R, Rf, beta=.95):
    """
    Taking the parameters A, B, Q, R as found in the `setup_matrices`,
    we find the value function of the optimal linear regulator problem.
    This is steps 2 and 3 in the lecture notes.
    Parameters
    ----------
    (A, B, Q, R) : Array(Float, ndim=2)
        The matrices that describe the oligopoly problem
    Returns
    -------
    (P, F, d) : Array(Float, ndim=2)
        The matrix that describes the value function of the optimal
        linear regulator problem.
    """

    lq = LQ(Q, -R, A, B, beta=beta)
    P, F, d = lq.stationary_values()

    Af = np.vstack((np.hstack([A-np.dot(B,F), np.array([[0., 0., 0., 0., 0.]]).T]),np.array([[0., 0., 0., 0., 0., 1.]])))
    Bf = np.array([[0., 0., 0., 0., 0., 1.]]).T

    lqf = LQ(Q, -Rf, Af, Bf, beta=beta)
    Pf, Ff, df = lqf.stationary_values()

    return P, F, d, Pf, Ff, df

def solve_for_opt_policy(params, eta0=0., Q0=0., q0=0.):
    """
    Taking the parameters as given, solve for the optimal decision rules
    for the firm.
    Parameters
    ----------
    params : Array(Float, ndim=1)
        This holds all of the model parameters in an array
    Returns
    -------
    out :
    """
    # Step 1/2: Formulate/Solve the optimal linear regulator
    (A, B, Q, R, Rf) = setup_matrices(params)
    (P, F, d, Pf, Ff, df) = find_PFd(A, B, Q, R, Rf, beta=beta)

    # Step 3: Convert implementation into state variables (Find coeffs)
    P22 = P[-1, -1]
    P21 = P[-1, :-1]
    P22inv = P22**(-1)

    # Step 4: Find optimal x_0 and \mu_{x, 0}
    z0 = np.array([1., eta0, Q0, q0])
    x0 = -P22inv*np.dot(P21, z0)
    D0 = -np.dot(P22inv, P21)

    # Return -F and -Ff because we use u_t = -F y_t
    return P, -F, D0, Pf, -Ff

# Parameter values
a0 = 100.
a1 = 1.
rho = .8
c_eps = .2
c = 1.
d = 20.
e = 20.
g = .2
h = .2
beta = .95
params = np.array([a0, a1, rho, c_eps, c, d, e, g, h, beta])

P, F, D0, Pf,Ff  = solve_for_opt_policy(params)

# Checking time-inconsistency:
A, B, Q, R, Rf = setup_matrices(params)
# arbitrary initial z_0
y0 = np.array([[1, 1, 1, 1]]).T
# optimal x_0 = i_0
i0 = np.dot(D0,y0)
# iterate one period using the closed-loop system
y1 = np.dot( A + np.dot(B,F) , np.vstack([y0, i0]) )
# the last element of y_1 is x_1 = i_1
i1_0 = y1[-1,0]

# compare this to the case when the leader solves a Stackelberg problem
# in period 1. if in period 1 the leader could choose i1 given
# (1, v_1, Q_1, \bar{q}_1)
i1_1 = np.dot(D0, y1[0:-1,0])

print("P = {}".format(P))
print("-F = {}".format(F))
print("D0 = {}".format(D0))
print("Pf = {}".format(Pf))
print("-Ff = {}".format(Ff))
print("i1_0 = {}".format(i1_0))
print("i1_1 = {}".format(i1_1))


from scipy import optimize
import scipy
import numpy
import matplotlib.pyplot as plt

def func(x, a, b):
        return a + b * numpy.log2(x)

def create_graph():
    with open('./results_n.txt') as f:
        y_pts_n = f.read().splitlines()
    with open('./results_n_2.txt') as f:
        y_pts_n_2 = f.read().splitlines()
    with open('./results_log_p.txt') as f:
        y_pts_log_p = f.read().splitlines()

    line_n, _ = scipy.optimize.curve_fit(lambda t,a,b: a+b*numpy.log2(t), range(2, 25, 2),  y_pts_n)
    line_n_2, _ = scipy.optimize.curve_fit(lambda t,a,b: a+b*numpy.log2(t), range(2, 49, 2),  y_pts_n_2)
    line_log_p, _ = scipy.optimize.curve_fit(lambda t,a,b: a+b*numpy.log2(t), range(2, 111, 2),  y_pts_log_p)

    fig, ax = plt.subplots()
    # ax.plot(range(2, 25, 2), y_pts_n, linestyle='None', marker='o', color='b', label=r"$p(n) = n$")
    ax.plot(range(2, 111, 2), func(range(2, 111, 2), *line_n),
            linewidth=2.0, linestyle='-', color='b', label=r"$Fitted Curve: n$"
    )
    # ax.plot(range(2, 49, 2), y_pts_n_2,
    #         linestyle='None', marker='o', color='r',  label=r"$p(n) = \frac{n}{2}$"
    # )
    ax.plot(range(2, 111, 2), func(range(2, 111, 2), *line_n_2),
            linewidth=2.0, linestyle='-', color='r', label=r"$Fitted Curve: \frac{n}{2}$"
    )
    # ax.plot(range(2, 111, 2), y_pts_log_p,
    #         linestyle='None', marker='o', color='g', label=r"$p(n) = \frac{n}{p} \geq \log \, p$"
    # )
    ax.plot(range(2, 111, 2), func(range(2, 111, 2), *line_log_p),
        linewidth=2.0, linestyle='-', color='g', label=r"$Fitted Curve: \frac{n}{p} \geq \log \, p$"
    )
    ax.set(xlabel='n - points count', ylabel=r'time $(\mu s)$',
           title='Line-of-Sight')
    ax.grid()
    ax.set_ylim(ymin=0)
    ax.legend(loc="upper left")
    fig.savefig("common_graph.pdf", format="pdf")
    plt.show()

if __name__ == '__main__':
    create_graph()


import sympy
import argparse
import numpy as np

import equations
import data
# from dsr_utils import run_dsr
from gp_utils import run_gp
from interpolate import num_diff, num_diff_gp
import pickle
import os
import time

# # set up ODE config
# ode_param = None
# x_id = 0
#
# # data generation config
# freq = 10
# n_sample = 100
# noise_sigma = 0.0
#
# # set up algorithm config
# alg = 'gp'

def run(ode_name, ode_param, x_id, freq, n_sample, noise_ratio, alg, seed, n_seed):
    np.random.seed(999)
    print(freq)

    ode = equations.get_ode(ode_name, ode_param)
    T = ode.T
    init_low = ode.init_low
    init_high = ode.init_high
    has_coef = ode.has_coef

    noise_sigma = ode.std_base * noise_ratio

    dg = data.DataGenerator(ode, T, freq, n_sample, noise_sigma, init_low, init_high)
    yt = dg.generate_data()

    if noise_sigma == 0:
        dxdt_hat = (yt[1:, :, :] - yt[:-1, :, :]) / (dg.solver.t[1:] - dg.solver.t[:-1])[:, None, None]
    elif alg != 'gp':
        dxdt_hat = num_diff(yt, dg, alg)
    else:
        dxdt_hat, xt_hat = num_diff_gp(yt, dg, ode)

    print('Numerical differentiation: Done.')

    # if alg != 'gp':
    X_train = yt[:-1, :, :]
    # else:
    #     X_train = xt_hat[:-1, :, :]
    X_train = X_train.reshape(X_train.shape[0] * X_train.shape[1], X_train.shape[2])

    y_train = dxdt_hat[:, :, x_id].flatten()
    assert X_train.shape[0] == y_train.shape[0]

    if alg == 'tv':
        path_base = 'results/{}/noise-{}/sample-{}/freq-{}/'.format(ode_name, noise_ratio, n_sample, freq)
    elif alg == 'gp':
        path_base = 'results_gp/{}/noise-{}/sample-{}/freq-{}/'.format(ode_name, noise_ratio, n_sample, freq)
    else:
        path_base = 'results_spline/{}/noise-{}/sample-{}/freq-{}/'.format(ode_name, noise_ratio, n_sample, freq)

    if not os.path.isdir(path_base):
        os.makedirs(path_base)

    for s in range(seed, seed+n_seed):
        print(' ')
        print('Running with seed {}'.format(s))
        if x_id == 0:
            path = path_base + 'grad_seed_{}.pkl'.format(s)
        else:
            path = path_base + 'grad_x_{}_seed_{}.pkl'.format(x_id, s)

        if os.path.isfile(path):
            print('Skipping seed {}'.format(s))
            continue
        start = time.time()
        f_hat, est_gp = run_gp(X_train, y_train, ode, x_id, s)
        print(f_hat)
        f_true = ode.get_expression()[x_id]
        if not isinstance(f_true, tuple):
            correct = sympy.simplify(f_hat - f_true) == 0
        else:
            correct_list = [sympy.simplify(f_hat - f) == 0 for f in f_true]
            correct = max(correct_list) == 1

        end = time.time()

        with open(path, 'wb') as f:
            pickle.dump({
                'model': est_gp._program,
                'X_train': X_train,
                'y_train': y_train,
                'seed': s,
                'correct': correct,
                'f_hat': f_hat,
                'ode': ode,
                'noise_ratio': noise_ratio,
                'noise_sigma': noise_sigma,
                'dg': dg,
                'time': end-start,
            }, f)

        print(f_hat)
        print(correct)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--ode_name", help="name of the ode", type=str)
    parser.add_argument("--ode_param", help="parameters of the ode (default: None)", type=str, default=None)
    parser.add_argument("--x_id", help="ID of the equation to be learned", type=int, default=0)
    parser.add_argument("--freq", help="sampling frequency", type=float, default=10)
    parser.add_argument("--n_sample", help="number of trajectories", type=int, default=100)
    parser.add_argument("--noise_sigma", help="noise level (default 0)", type=float, default=0.)
    parser.add_argument("--alg", help="name of the benchmark", type=str, default='tv', choices=['tv', 'spline', 'gp'])
    parser.add_argument("--seed", help="random seed", type=int, default=0)
    parser.add_argument("--n_seed", help="random seed", type=int, default=10)

    args = parser.parse_args()
    print('Running with: ', args)

    if args.ode_param is not None:
        param = [float(x) for x in args.ode_param.split(',')]
    else:
        param = None
    if args.freq >= 1:
        freq = int(args.freq)
    else:
        freq = args.freq
    run(args.ode_name, param, args.x_id, freq, args.n_sample, args.noise_sigma, args.alg, seed=args.seed, n_seed=args.n_seed)


<reponame>bigpeng2012/CarND-Vehicle-Detection<filename>Vehicle_Detection_Tracking.py
#!/usr/bin/env python
# coding: utf-8

# # P5 Vehicle Detection

# ## The goals / steps of this project are the following:
#
# 1.Perform a Histogram of Oriented Gradients (HOG) feature extraction on a labeled training set of images and train a classifier Linear SVM classifier
#
# 2.Optionally, you can also apply a color transform and append binned color features, as well as histograms of color, to your HOG feature vector.
#
# 3.Note: for those first two steps don't forget to normalize your features and randomize a selection for training and testing.
#
# 4.Implement a sliding-window technique and use your trained classifier to search for vehicles in images.
#
# 5.Run your pipeline on a video stream (start with the test_video.mp4 and later implement on full project_video.mp4) and create a heat map of recurring detections frame by frame to reject outliers and follow detected vehicles.
#
# 6.Estimate a bounding box for vehicles detected.

# In[1]:

import cv2
import glob
import numpy as np

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import time

from random import shuffle

from skimage.feature import hog
from sklearn.svm import LinearSVC
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from scipy.ndimage.measurements import label

from sklearn.preprocessing import StandardScaler
from tqdm import tqdm

from moviepy.editor import VideoFileClip
from IPython.display import HTML

get_ipython().run_line_magic('matplotlib', 'inline')

# ### Load Data

# In[2]:

car_images = glob.glob('vehicles/**/*.png')
noncar_images = glob.glob('non-vehicles/**/*.png')
print(len(car_images), len(noncar_images))

# ### Visualize Data Example

# In[3]:

fig, axs = plt.subplots(2,8, figsize=(8, 3))
#fig.subplots_adjust(hspace = .2, wspace=.001)
axs = axs.ravel()

for i in np.arange(8):
    img = cv2.imread(car_images[np.random.randint(0,len(car_images))])
    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    axs[i].axis('off')
    axs[i].set_title('car', fontsize=15)
    axs[i].imshow(img)
for i in np.arange(8,16):
    img = cv2.imread(noncar_images[np.random.randint(0,len(noncar_images))])
    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
    axs[i].axis('off')
    axs[i].set_title('noncar', fontsize=15)
    axs[i].imshow(img)

# ## Feature Extraction

# ### Histogram of Oriented Gradient (HOG) Features

# In[4]:

def get_hog_features(img, orient, pix_per_cell, cell_per_block,
                        vis=False, feature_vec=True):
    # Call with two outputs if vis==True
    if vis == True:
        features, hog_image = hog(img, orientations=orient,
                                  pixels_per_cell=(pix_per_cell, pix_per_cell),
                                  cells_per_block=(cell_per_block, cell_per_block),
                                  transform_sqrt=False,
                                  visualise=True, feature_vector=False)
        return features, hog_image
    # Otherwise call with one output
    else:
        features = hog(img, orientations=orient,
                       pixels_per_cell=(pix_per_cell, pix_per_cell),
                       cells_per_block=(cell_per_block, cell_per_block),
                       transform_sqrt=False,
                       visualise=False, feature_vector=feature_vec)
        return features

# ### Example of HOG Features

# In[5]:

# Generate a random index to look at a car image
ind = np.random.randint(0, len(car_images))
# Read in the image
image = mpimg.imread(car_images[ind])
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
# Define HOG parameters
orient = 10
pix_per_cell = 8
cell_per_block = 2
# Call our function with vis=True to see an image output
features, hog_image = get_hog_features(gray, orient,
                        pix_per_cell, cell_per_block,
                        vis=True, feature_vec=False)

# Plot the examples
fig = plt.figure(figsize = (15,20))
plt.subplot(121)
plt.imshow(image, cmap='gray')
plt.title('Example Car Image')
plt.subplot(122)
plt.imshow(hog_image, cmap='gray')
plt.title('HOG Visualization')
plt.savefig('output_images/hog_visualization.png',
            bbox_inches="tight")

# ### Histogram of Color

# In[6]:

def color_hist(img, nbins=32, bins_range=(0, 256), vis=False):

    bins_range=(0, 256)
    # Compute the histogram of the color channels separately
    channel1_hist = np.histogram(img[:,:,0], bins=nbins, range=bins_range)
    channel2_hist = np.histogram(img[:,:,1], bins=nbins, range=bins_range)
    channel3_hist = np.histogram(img[:,:,2], bins=nbins, range=bins_range)

    # Generating bin centers
    bin_edges = channel1_hist[1] #all three bins are the same size
    bin_centers = (bin_edges[1:] + bin_edges[0:len(bin_edges) - 1])/2

    hist_features = np.concatenate((channel1_hist[0],
                                    channel2_hist[0],
                                    channel3_hist[0]))
    if vis == True:
        return channel1_hist, channel2_hist, channel3_hist, bin_centers
    else:
        return hist_features

# ### Spatial Binning of Color

# In[7]:

def bin_spatial(img, color_space='RGB', size=(32, 32)):
    # Convert image to new color space (if specified)
    if color_space != 'RGB':
        if color_space == 'HSV':
            feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
        elif color_space == 'LUV':
            feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2LUV)
        elif color_space == 'HLS':
            feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2HLS)
        elif color_space == 'YUV':
            feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2YUV)
        elif color_space == 'YCrCb':
            feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2YCrCb)
    else: feature_image = np.copy(img)
    # Use cv2.resize().ravel() to create the feature vector
    features = cv2.resize(feature_image, size).ravel()
    # Return the feature vector
    return features

# ### Extract Features from a list of Imgaes

# In[8]:

# Define a function to extract features from a list of images
# Have this function call bin_spatial() and color_hist()
def extract_features(images, color_space='RGB', spatial_size=(32, 32),
                     hist_bins=32, orient=9, pix_per_cell=8,
                     cell_per_block=2, hog_channel=0,
                     spatial_feat=True, hist_feat=True, hog_feat=True):
    # Create a list to append feature vectors to
    features = []
    # Iterate through the list of images
    for file in tqdm(images):
        file_features = []
        img = mpimg.imread(file) # idea for the progress came from "nhiddink"
        # apply color conversion if other than 'RGB'
        if color_space != 'RGB':
            if color_space == 'HSV':
                feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
            elif color_space == 'LUV':
                feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2LUV)
            elif color_space == 'HLS':
                feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2HLS)
            elif color_space == 'YUV':
                feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2YUV)
            elif color_space == 'YCrCb':
                feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2YCrCb)
        else: feature_image = np.copy(img)

        if spatial_feat == True:
            # Apply bin_spatial() to get spatial color features
            spatial_features = bin_spatial(feature_image, size=spatial_size)
            file_features.append(spatial_features)
        if hist_feat == True:
            # Apply color_hist() also with a color space option now
            hist_features = color_hist(feature_image, nbins=hist_bins)
            file_features.append(hist_features)
        if hog_feat == True:
            # Call get_hog_features() with vis=False, feature_vec=True
            if hog_channel == 'ALL':
                hog_features = []
                for channel in range(feature_image.shape[2]):
                    hog_features.append(get_hog_features(feature_image[:,:,channel],
                                        orient, pix_per_cell, cell_per_block,
                                        vis=False, feature_vec=True))
                hog_features = np.ravel(hog_features)
            else:
                hog_features = get_hog_features(feature_image[:,:,hog_channel], orient,
                                                pix_per_cell, cell_per_block, vis=False,
                                                feature_vec=True)
            file_features.append(hog_features)
        # Append the new feature vector to the features list
        features.append(np.concatenate(file_features))
    # Return list of feature vectors
    return features

# ### Slide Window Search

# In[9]:

def slide_window(img, x_start_stop=[None, None], y_start_stop=[None, None],
                 xy_window=(128, 128), #(64, 64), (96, 96)
                 xy_overlap=(0.5, 0.5)):
    # If x and/or y start/stop positions not defined, set to image size
    if x_start_stop[0] == None:
        x_start_stop[0] = 0
    if x_start_stop[1] == None:
        x_start_stop[1] = img.shape[1]
    if y_start_stop[0] == None:
        y_start_stop[0] = 0
    if y_start_stop[1] == None:
        y_start_stop[1] = img.shape[0]
    # Compute the span of the region to be searched
    xspan = x_start_stop[1] - x_start_stop[0]
    yspan = y_start_stop[1] - y_start_stop[0]
    # Compute the number of pixels per step in x/y
    nx_pix_per_step = np.int(xy_window[0]*(1 - xy_overlap[0]))
    ny_pix_per_step = np.int(xy_window[1]*(1 - xy_overlap[1]))
    # Compute the number of windows in x/y
    nx_buffer = np.int(xy_window[0]*(xy_overlap[0]))
    ny_buffer = np.int(xy_window[1]*(xy_overlap[1]))
    nx_windows = np.int((xspan-nx_buffer)/nx_pix_per_step)
    ny_windows = np.int((yspan-ny_buffer)/ny_pix_per_step)
    # Initialize a list to append window positions to
    window_list = []

    for ys in range(ny_windows):
        for xs in range(nx_windows):
            # Calculate window position
            startx = xs*nx_pix_per_step + x_start_stop[0]
            endx = startx + xy_window[0]
            starty = ys*ny_pix_per_step + y_start_stop[0]
            endy = starty + xy_window[1]
            # Append window position to list
            window_list.append(((startx, starty), (endx, endy)))
    # Return the list of windows
    return window_list

# Define a function to draw bounding boxes

# In[10]:

# Define a function to draw bounding boxes
def draw_boxes(img, bboxes, color=(255, 0, 0), thick=6):
    # Make a copy of the image
    imcopy = np.copy(img)
    # Iterate through the bounding boxes
    for bbox in bboxes:
        # Draw a rectangle given bbox coordinates
        cv2.rectangle(imcopy, bbox[0], bbox[1], color, thick)
    # Return the image copy with boxes drawn
    return imcopy

# ### Sliding Window Example

# In[11]:

test_img = mpimg.imread('test_images/test4.jpg')

windows =  slide_window(test_img,
                        x_start_stop=[None, None],
                        y_start_stop=[400, 656], #tune the parameters
                        xy_window=(64, 64),
                        xy_overlap=(0.5, 0.5))

window_img = draw_boxes(test_img, windows, color=(255,0,0),thick = 6)
plt.imshow(window_img);
matplotlib.rc('xtick', labelsize=15)
matplotlib.rc('ytick', labelsize=15)
plt.title('Sliding Windows Technique:', fontsize=15);
#plt.savefig('output_images/sliding_windows.png', bbox_inches="tight")

# ### Extract Features from a Single Image

# In[12]:

# This function is very similar to extract_features()
# just for a single image rather than list of images
def single_img_features(image, color_space='RGB', spatial_size=(32, 32),
                        hist_bins=32, orient=9,
                        pix_per_cell=8, cell_per_block=2, hog_channel=0,
                        spatial_feat=True, hist_feat=True, hog_feat=True):
    #1) Define an empty list to receive features
    img_features = []
    #2) Apply color conversion if other than 'RGB'
    if color_space != 'RGB':
        if color_space == 'HSV':
            feature_image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
        elif color_space == 'LUV':
            feature_image = cv2.cvtColor(image, cv2.COLOR_RGB2LUV)
        elif color_space == 'HLS':
            feature_image = cv2.cvtColor(image, cv2.COLOR_RGB2HLS)
        elif color_space == 'YUV':
            feature_image = cv2.cvtColor(image, cv2.COLOR_RGB2YUV)
        elif color_space == 'YCrCb':
            feature_image = cv2.cvtColor(image, cv2.COLOR_RGB2YCrCb)
    else: feature_image = np.copy(image)
    #3) Compute spatial features if flag is set
    if spatial_feat == True:
        # Apply bin_spatial() to get spatial color features
        spatial_features = bin_spatial(feature_image, size=spatial_size)
        #4) Append features to list
        img_features.append(spatial_features)
    #5) Compute histogram features if flag is set
    if hist_feat == True:
        # Apply color_hist() also with a color space option now
        hist_features = color_hist(feature_image, nbins=hist_bins)
        #6) Append features to list
        img_features.append(hist_features)
    #7) Compute HOG features if flag is set
    if hog_feat == True:
        if hog_channel == 'ALL':
            hog_features = []
            for channel in range(feature_image.shape[2]):
                hog_features.append(get_hog_features(feature_image[:,:,channel],
                                    orient, pix_per_cell, cell_per_block,
                                    vis=False, feature_vec=True))
            hog_features = np.ravel(hog_features)
        else:
            hog_features = get_hog_features(feature_image[:,:,hog_channel], orient,
                                            pix_per_cell, cell_per_block,
                                            vis=False, feature_vec=True)
        #8) Append features to list
        img_features.append(hog_features)
    #9) Return concatenated array of features
    return np.concatenate(img_features)

# In[13]:

# Define a function you will pass an image
# and the list of windows to be searched (output of slide_windows())
def search_windows(img, windows, clf, scaler, color_space='RGB',
                   spatial_size=(32, 32), hist_bins=32, orient=9,
                   pix_per_cell=8, cell_per_block=2, hog_channel=0,
                   spatial_feat=True, hist_feat=True, hog_feat=True):

    #1) Create an empty list to receive positive detection windows
    on_windows = []
    #2) Iterate over all windows in the list
    for window in windows:
        #3) Extract the test window from original image
        test_img = cv2.resize(img[window[0][1]:window[1][1], window[0][0]:window[1][0]], (64, 64))
        #4) Extract features for that window using single_img_features()
        features = single_img_features(test_img, color_space=color_space,
                            spatial_size=spatial_size, hist_bins=hist_bins,
                            orient=orient, pix_per_cell=pix_per_cell,
                            cell_per_block=cell_per_block,
                            hog_channel=hog_channel, spatial_feat=spatial_feat,
                            hist_feat=hist_feat, hog_feat=hog_feat)
        #5) Scale extracted features to be fed to classifier
        test_features = scaler.transform(np.array(features).reshape(1, -1))
        #6) Predict using your classifier
        prediction = clf.predict(test_features)
        #7) If positive (prediction == 1) then save the window
        if prediction == 1:
            on_windows.append(window)
    #8) Return windows for positive detections
    return on_windows

# ## Train and Test the Classifier

# In[14]:

### TODO: Tweak these parameters and see how the results change.
color_space = 'YCrCb' # Can be RGB, HSV, LUV, HLS, YUV, YCrCb
#color_space = 'RGB'
orient = 10  # HOG orientations   #9
pix_per_cell = 8 # HOG pixels per cell
cell_per_block = 2 # HOG cells per block
hog_channel = 'ALL' # Can be 0, 1, 2, or "ALL"
spatial_size = (32, 32) # Spatial binning dimensions
hist_bins = 64   # Number of histogram bins
spatial_feat = True # Spatial features on or off
hist_feat = True # Histogram features on or off
hog_feat = True # HOG features on or off

car_features = extract_features(car_images, color_space=color_space,
                            spatial_size=spatial_size, hist_bins=hist_bins,
                            orient=orient, pix_per_cell=pix_per_cell,
                            cell_per_block=cell_per_block,
                            hog_channel=hog_channel,
                            spatial_feat=spatial_feat, hist_feat=hist_feat,
                            hog_feat=hog_feat)
notcar_features = extract_features(noncar_images, color_space=color_space,
                            spatial_size=spatial_size, hist_bins=hist_bins,
                            orient=orient, pix_per_cell=pix_per_cell,
                            cell_per_block=cell_per_block,
                            hog_channel=hog_channel,
                            spatial_feat=spatial_feat, hist_feat=hist_feat,
                            hog_feat=hog_feat)

# Create an array stack of feature vectors
X = np.vstack((car_features, notcar_features)).astype(np.float64)
# Fit a per-column scaler
X_scaler = StandardScaler().fit(X)
# Apply the scaler to X
scaled_X = X_scaler.transform(X)

# Define the labels vector
y = np.hstack((np.ones(len(car_features)), np.zeros(len(notcar_features))))

# Split up data into randomized training and test sets
rand_state = np.random.randint(0, 100)
X_train, X_test, y_train, y_test = train_test_split(scaled_X, y,
                                                    test_size=0.2,
                                                    random_state=rand_state)
# Use a linear SVC
svc = LinearSVC()
# Check the training time for the SVC
t=time.time()
svc.fit(X_train, y_train)
t2 = time.time()
print(round(t2-t, 2), 'seconds to train SVC.')
print('Using:', orient, 'orientations', pix_per_cell,
    'pixels per cell and', cell_per_block, 'cells per block')
print('Feature vector length:', len(X_train[0]))

print('Test Accuracy of SVC: {0:.2f}%'.format(round(svc.score(X_test, y_test)*100, 4)))
print()
t = time.time()
print('  Predictions:', svc.predict(X_test[0:10]))
print('       Labels:', y_test[0:10])
t2 = time.time()
print()
print(round(t2-t, 5), 'seconds to predict 10 labels with SVC.')

# ### Test on test_images

# In[15]:

for i in range(1,7):

    fname = 'test_images/test{}.jpg'.format(i)
    image = mpimg.imread(fname)
    draw_image = np.copy(image)

    image = image.astype(np.float32)/255

    windows =  slide_window(test_img,
                            x_start_stop=[600, None],
                            y_start_stop=[400, 656], #tune the parameters
                            xy_window=(128,128),
                            xy_overlap=(.7,.7))

    hot_windows = search_windows(image, windows, svc, X_scaler, color_space=color_space,
                            spatial_size=spatial_size, hist_bins=hist_bins,
                            orient=orient, pix_per_cell=pix_per_cell,
                            cell_per_block=cell_per_block,
                            hog_channel=hog_channel, spatial_feat=spatial_feat,
                            hist_feat=hist_feat, hog_feat=hog_feat)

    window_img = draw_boxes(draw_image, hot_windows)

    f, (ax1, ax2) = plt.subplots(1, 2, figsize=(18,9))
    plt.tight_layout()
    ax1.imshow(draw_image)
    ax1.set_title('Original Image', fontsize=30)
    ax2.imshow(window_img)
    ax2.set_title('Detect Vehicles', fontsize=30)
    #plt.savefig('output_images/windows.png', bbox_inches="tight")

# ### Build Heat Maps to Fix Multiple Detections & False Positives

# In[16]:

def add_heat(heatmap, bbox_list):
    # Iterate through list of bboxes
    for box in bbox_list:
        # Add += 1 for all pixels inside each bbox
        # Assuming each "box" takes the form ((x1, y1), (x2, y2))
        heatmap[box[0][1]:box[1][1], box[0][0]:box[1][0]] += 1

    # Return updated heatmap
    return heatmap# Iterate through list of bboxes

def apply_threshold(heatmap, threshold):
    # Zero out pixels below the threshold
    heatmap[heatmap <= threshold] = 0
    # Return thresholded map
    return heatmap

def draw_labeled_bboxes(img, labels):
    # Iterate through all detected cars
    for car_number in range(1, labels[1]+1):
        # Find pixels with each car_number label value
        nonzero = (labels[0] == car_number).nonzero()
        # Identify x and y values of those pixels
        nonzeroy = np.array(nonzero[0])
        nonzerox = np.array(nonzero[1])
        # Define a bounding box based on min/max x and y
        bbox = ((np.min(nonzerox), np.min(nonzeroy)), (np.max(nonzerox), np.max(nonzeroy)))
        # Draw the box on the image
        cv2.rectangle(img, bbox[0], bbox[1], (255,0,0), 8)
    # Return the image
    return img

def convert_color(img, conv='RGB2YCrCb'):
    if conv == 'RGB2YCrCb':
        return cv2.cvtColor(img, cv2.COLOR_RGB2YCrCb)
    if conv == 'BGR2YCrCb':
        return cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
    if conv == 'RGB2LUV':
        return cv2.cvtColor(img, cv2.COLOR_RGB2LUV)

# In[17]:

def find_cars(img, ystart, ystop, scale, svc, X_scaler,
              orient, pix_per_cell, cell_per_block,
              spatial_size, hist_bins):

    draw_img = np.copy(img)
    img = img.astype(np.float32)/255

    img_tosearch = img[ystart:ystop,:,:]
    ctrans_tosearch = convert_color(img_tosearch, conv='RGB2YCrCb')
    if scale != 1:
        imshape = ctrans_tosearch.shape
        ctrans_tosearch = cv2.resize(ctrans_tosearch,
                                     (np.int(imshape[1]/scale),
                                      np.int(imshape[0]/scale)))

    ch1 = ctrans_tosearch[:,:,0]
    ch2 = ctrans_tosearch[:,:,1]
    ch3 = ctrans_tosearch[:,:,2]

    # Define blocks and steps as above
    nxblocks = (ch1.shape[1] // pix_per_cell) - cell_per_block + 1
    nyblocks = (ch1.shape[0] // pix_per_cell) - cell_per_block + 1
    nfeat_per_block = orient*cell_per_block**2

    # 64 was the orginal sampling rate, with 8 cells and 8 pix per cell
    window = 64
    nblocks_per_window = (window // pix_per_cell) - cell_per_block + 1
    cells_per_step = 2  # Instead of overlap, define how many cells to step
    nxsteps = (nxblocks - nblocks_per_window) // cells_per_step
    nysteps = (nyblocks - nblocks_per_window) // cells_per_step

    # Compute individual channel HOG features for the entire image
    hog1 = get_hog_features(ch1, orient, pix_per_cell, cell_per_block, feature_vec=False)
    hog2 = get_hog_features(ch2, orient, pix_per_cell, cell_per_block, feature_vec=False)
    hog3 = get_hog_features(ch3, orient, pix_per_cell, cell_per_block, feature_vec=False)

    bbox_list=[]
    for xb in range(nxsteps):
        for yb in range(nysteps):
            ypos = yb*cells_per_step
            xpos = xb*cells_per_step
            # Extract HOG for this patch
            hog_feat1 = hog1[ypos:ypos+nblocks_per_window, xpos:xpos+nblocks_per_window].ravel()
            hog_feat2 = hog2[ypos:ypos+nblocks_per_window, xpos:xpos+nblocks_per_window].ravel()
            hog_feat3 = hog3[ypos:ypos+nblocks_per_window, xpos:xpos+nblocks_per_window].ravel()
            hog_features = np.hstack((hog_feat1, hog_feat2, hog_feat3))

            xleft = xpos*pix_per_cell
            ytop = ypos*pix_per_cell

            # Extract the image patch
            subimg = cv2.resize(ctrans_tosearch[ytop:ytop+window, xleft:xleft+window], (64,64))

            # Get color features
            spatial_features = bin_spatial(subimg, size=spatial_size)
            hist_features = color_hist(subimg, nbins=hist_bins)

            # Scale features and make a prediction
            test_features = X_scaler.transform(np.hstack((spatial_features,
                                                          hist_features,
                                                          hog_features)).reshape(1, -1))

            # Scale features and make a prediction
            test_prediction = svc.predict(test_features)

            if test_prediction == 1:
                xbox_left = np.int(xleft*scale)
                ytop_draw = np.int(ytop*scale)
                win_draw = np.int(window*scale)
                cv2.rectangle(draw_img,(xbox_left, ytop_draw+ystart),
                              (xbox_left+win_draw,ytop_draw+win_draw+ystart),
                              (255,0,0),8)
                bbox_list.append(((xbox_left, ytop_draw+ystart),
                                  (xbox_left+win_draw,ytop_draw+win_draw+ystart)))

    return bbox_list

# ### Test on test_images

# In[18]:

for i in range(1,7):

    fname = 'test_images/test{}.jpg'.format(i)
    img = mpimg.imread(fname)

    orient=10
    pix_per_cell=8
    cell_per_block=2
    spatial_size=(32, 32)
    hist_bins=64

    bbox_list = []

    ystart = 380
    ystop = 550
    scale = 1.0
    bbox_list.append(find_cars(img, ystart, ystop, scale, svc, X_scaler,
                                    orient, pix_per_cell, cell_per_block,
                                    spatial_size, hist_bins))

    ystart = 400
    ystop = 600
    scale = 1.5
    bbox_list.append(find_cars(img, ystart, ystop, scale, svc, X_scaler,
                                    orient, pix_per_cell, cell_per_block,
                                    spatial_size, hist_bins))

    ystart = 400
    ystop = 656
    scale = 2.0
    bbox_list.append(find_cars(img, ystart, ystop, scale, svc, X_scaler,
                                    orient, pix_per_cell, cell_per_block,
                                    spatial_size, hist_bins))

    bbox_list = [item for sublist in bbox_list for item in sublist]

    #out_img = draw_boxes(img, bbox_list, random_color=True)
    out_img = draw_boxes(img, bbox_list)

    heat = np.zeros_like(img[:,:,0]).astype(np.float)
    heat = add_heat(heat, bbox_list)
    heat = apply_threshold(heat, 2)

    # Find final boxes from heatmap using label function
    labels = label(heat)
    new_img = draw_labeled_bboxes(np.copy(img), labels)

    f, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(18,9))
    plt.tight_layout()
    ax1.imshow(out_img)
    ax1.set_title('Search Boxes', fontsize=30)
    ax2.imshow(heat, cmap='hot')
    ax2.set_title('Heat Map', fontsize=30)
    ax3.imshow(new_img)
    ax3.set_title('Bounding Boxes', fontsize=30)
    #plt.savefig('output_images/heat_map1.png', bbox_inches="tight")

# ## Final Pipeline

# In[19]:

def detect_vehicle(img):

    orient=10
    pix_per_cell=8
    cell_per_block=2
    spatial_size=(32, 32)
    hist_bins=64

    bbox_list = []

    ystart = 380
    ystop = 550
    scale = 1.0
    bbox_list.append(find_cars(img, ystart, ystop, scale, svc, X_scaler,
                                    orient, pix_per_cell, cell_per_block,
                                    spatial_size, hist_bins))

    ystart = 400
    ystop = 600
    scale = 1.5
    bbox_list.append(find_cars(img, ystart, ystop, scale, svc, X_scaler,
                                    orient, pix_per_cell, cell_per_block,
                                    spatial_size, hist_bins))

    ystart = 400
    ystop = 656
    scale = 2.0
    bbox_list.append(find_cars(img, ystart, ystop, scale, svc, X_scaler,
                                    orient, pix_per_cell, cell_per_block,
                                    spatial_size, hist_bins))
    bbox_list = [item for sublist in bbox_list for item in sublist]

    heat = np.zeros_like(img[:,:,0]).astype(np.float)
    heat = add_heat(heat, bbox_list)
    heat = apply_threshold(heat, 2)

    # Find final boxes from heatmap using label function
    labels = label(heat)
    new_img = draw_labeled_bboxes(np.copy(img), labels)

    return new_img

# ### Test on Video Stream

# In[20]:

clip1 = VideoFileClip("test_video.mp4")
test_output  = "test_output.mp4"
test_clip = clip1.fl_image(detect_vehicle)
get_ipython().run_line_magic('time', 'test_clip.write_videofile(test_output, audio=False)')

# In[21]:

clip2 = VideoFileClip("project_video.mp4")
output = "project_output.mp4"
clip = clip2.fl_image(detect_vehicle)
get_ipython().run_line_magic('time', 'clip.write_videofile(output, audio=False)')

# In[ ]:


<filename>tests/test_series.py
from sympy import (
    symbols, cos, sin, log, sqrt,
    Tuple, pi, Plane, S, I, im,
    Circle, Point,
    Piecewise, And, Eq, Interval, Abs, lambdify
)
from spb.series import (
    LineOver1DRangeSeries, Parametric2DLineSeries, Parametric3DLineSeries,
    SurfaceOver2DRangeSeries, ContourSeries, ParametricSurfaceSeries,
    InteractiveSeries,
    ImplicitSeries,
    Vector2DSeries, Vector3DSeries, SliceVector3DSeries,
    ComplexSeries, ComplexInteractiveSeries, ComplexPointSeries,
    ComplexPointInteractiveSeries,
    GeometrySeries, GeometryInteractiveSeries,
    PlaneSeries, PlaneInteractiveSeries,
    List2DSeries
)
from spb.functions import _process_piecewise
import numpy as np
from pytest import warns, raises

def test_lin_log_scale():
    # Verify that data series create the correct spacing in the data.
    x, y, z = symbols("x, y, z")

    s = LineOver1DRangeSeries(x, (x, 1, 10), adaptive=False, n=50, xscale="linear")
    xx, _ = s.get_data()
    assert np.isclose(xx[1] - xx[0], xx[-1] - xx[-2])

    s = LineOver1DRangeSeries(x, (x, 1, 10), adaptive=False, n=50, xscale="log")
    xx, _ = s.get_data()
    assert not np.isclose(xx[1] - xx[0], xx[-1] - xx[-2])

    s = Parametric2DLineSeries(
        cos(x), sin(x), (x, pi / 2, 1.5 * pi), adaptive=False, n=50, xscale="linear"
    )
    _, _, param = s.get_data()
    assert np.isclose(param[1] - param[0], param[-1] - param[-2])

    s = Parametric2DLineSeries(
        cos(x), sin(x), (x, pi / 2, 1.5 * pi), adaptive=False, n=50, xscale="log"
    )
    _, _, param = s.get_data()
    assert not np.isclose(param[1] - param[0], param[-1] - param[-2])

    s = Parametric3DLineSeries(
        cos(x), sin(x), x, (x, pi / 2, 1.5 * pi), adaptive=False, n=50, xscale="linear"
    )
    _, _, _, param = s.get_data()
    assert np.isclose(param[1] - param[0], param[-1] - param[-2])

    s = Parametric3DLineSeries(
        cos(x), sin(x), x, (x, pi / 2, 1.5 * pi), adaptive=False, n=50, xscale="log"
    )
    _, _, _, param = s.get_data()
    assert not np.isclose(param[1] - param[0], param[-1] - param[-2])

    s = SurfaceOver2DRangeSeries(
        cos(x ** 2 + y ** 2),
        (x, 1, 5),
        (y, 1, 5),
        n=10,
        xscale="linear",
        yscale="linear",
    )
    xx, yy, _ = s.get_data()
    assert np.isclose(xx[0, 1] - xx[0, 0], xx[0, -1] - xx[0, -2])
    assert np.isclose(yy[1, 0] - yy[0, 0], yy[-1, 0] - yy[-2, 0])

    s = SurfaceOver2DRangeSeries(
        cos(x ** 2 + y ** 2), (x, 1, 5), (y, 1, 5), n=10, xscale="log", yscale="log"
    )
    xx, yy, _ = s.get_data()
    assert not np.isclose(xx[0, 1] - xx[0, 0], xx[0, -1] - xx[0, -2])
    assert not np.isclose(yy[1, 0] - yy[0, 0], yy[-1, 0] - yy[-2, 0])

    s = ImplicitSeries(
        cos(x ** 2 + y ** 2) > 0,
        (x, 1, 5),
        (y, 1, 5),
        n=10,
        xscale="linear",
        yscale="linear",
        adaptive=False,
    )
    xx, yy, _, _, _ = s.get_data()
    assert np.isclose(xx[1] - xx[0], xx[-1] - xx[-2])
    assert np.isclose(yy[1] - yy[0], yy[-1] - yy[-2])

    s = ImplicitSeries(
        cos(x ** 2 + y ** 2) > 0,
        (x, 1, 5),
        (y, 1, 5),
        n=10,
        xscale="log",
        yscale="log",
        adaptive=False,
    )
    xx, yy, _, _, _ = s.get_data()
    assert not np.isclose(xx[1] - xx[0], xx[-1] - xx[-2])
    assert not np.isclose(yy[1] - yy[0], yy[-1] - yy[-2])

    s = InteractiveSeries([log(x)], [(x, 1e-05, 1e05)], n=10, xscale="linear")
    xx, yy = s.get_data()
    assert np.isclose(xx[1] - xx[0], xx[-1] - xx[-2])

    s = InteractiveSeries([log(x)], [(x, 1e-05, 1e05)], n=10, xscale="log")
    xx, yy = s.get_data()
    assert not np.isclose(xx[1] - xx[0], xx[-1] - xx[-2])

    s = ComplexSeries(
        cos(x),
        (x, 1e-05, 1e05),
        n=10,
        xscale="linear",
        adaptive=False,
    )
    xx, yy, _ = s.get_data()
    assert np.isclose(xx[0, 1] - xx[0, 0], xx[0, -1] - xx[0, -2])

    s = ComplexSeries(
        cos(x),
        (x, 1e-05, 1e05),
        n=10,
        xscale="log",
        adaptive=False,
    )
    xx, yy, _ = s.get_data()
    assert not np.isclose(xx[0, 1] - xx[0, 0], xx[0, -1] - xx[0, -2])

    s = Vector3DSeries(
        x,
        y,
        z,
        (x, 1, 1e05),
        (y, 1, 1e05),
        (z, 1, 1e05),
        xscale="linear",
        yscale="linear",
        zscale="linear",
    )
    xx, yy, zz, _, _, _ = s.get_data()
    assert np.isclose(
        xx[0, :, 0][1] - xx[0, :, 0][0], xx[0, :, 0][-1] - xx[0, :, 0][-2]
    )
    assert np.isclose(
        yy[:, 0, 0][1] - yy[:, 0, 0][0], yy[:, 0, 0][-1] - yy[:, 0, 0][-2]
    )
    assert np.isclose(
        zz[0, 0, :][1] - zz[0, 0, :][0], zz[0, 0, :][-1] - zz[0, 0, :][-2]
    )

    s = Vector3DSeries(
        x,
        y,
        z,
        (x, 1, 1e05),
        (y, 1, 1e05),
        (z, 1, 1e05),
        xscale="log",
        yscale="log",
        zscale="log",
    )
    xx, yy, zz, _, _, _ = s.get_data()
    assert not np.isclose(
        xx[0, :, 0][1] - xx[0, :, 0][0], xx[0, :, 0][-1] - xx[0, :, 0][-2]
    )
    assert not np.isclose(
        yy[:, 0, 0][1] - yy[:, 0, 0][0], yy[:, 0, 0][-1] - yy[:, 0, 0][-2]
    )
    assert not np.isclose(
        zz[0, 0, :][1] - zz[0, 0, :][0], zz[0, 0, :][-1] - zz[0, 0, :][-2]
    )

def test_data_shape():
    # Verify that the series produces the correct data shape when the input
    # expression is a number.
    u, x, y, z = symbols("u, x:z")

    # scalar expression: it should return a numpy ones array
    s = LineOver1DRangeSeries(1, (x, -5, 5))
    xx, yy = s.get_data()
    assert len(xx) == len(yy)
    assert np.all(yy == 1)

    s = LineOver1DRangeSeries(1, (x, -5, 5), adaptive=False)
    xx, yy = s.get_data()
    assert len(xx) == len(yy)
    assert np.all(yy == 1)

    s = Parametric2DLineSeries(sin(x), 1, (x, 0, pi))
    xx, yy, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(param))
    assert np.all(yy == 1)

    s = Parametric2DLineSeries(1, sin(x), (x, 0, pi))
    xx, yy, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(param))
    assert np.all(xx == 1)

    s = Parametric2DLineSeries(sin(x), 1, (x, 0, pi), adaptive=False)
    xx, yy, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(param))
    assert np.all(yy == 1)

    s = Parametric2DLineSeries(1, sin(x), (x, 0, pi), adaptive=False)
    xx, yy, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(param))
    assert np.all(xx == 1)

    s = Parametric3DLineSeries(cos(x), sin(x), 1, (x, 0, 2 * pi))
    xx, yy, zz, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(zz)) and (len(xx) == len(param))
    assert np.all(zz == 1)

    s = Parametric3DLineSeries(cos(x), 1, x, (x, 0, 2 * pi))
    xx, yy, zz, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(zz)) and (len(xx) == len(param))
    assert np.all(yy == 1)

    s = Parametric3DLineSeries(1, sin(x), x, (x, 0, 2 * pi))
    xx, yy, zz, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(zz)) and (len(xx) == len(param))
    assert np.all(xx == 1)

    s = SurfaceOver2DRangeSeries(1, (x, -2, 2), (y, -3, 3))
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)
    assert np.all(zz == 1)

    s = ParametricSurfaceSeries(1, x, y, (x, 0, 1), (y, 0, 1))
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)
    assert np.all(xx == 1)

    s = ParametricSurfaceSeries(1, 1, y, (x, 0, 1), (y, 0, 1))
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)
    assert np.all(yy == 1)

    s = ParametricSurfaceSeries(x, 1, 1, (x, 0, 1), (y, 0, 1))
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)
    assert np.all(zz == 1)

    s = ComplexSeries(1, (x, -5, 5), modules=None)
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)
    assert np.all(zz == 1)

    s = ComplexSeries(1, (x, -5, 5), modules="mpmath")
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)
    assert np.all(zz == 1)

    s = ComplexSeries(1, (x, -5 - 2 * I, 5 + 2 * I), domain_coloring=True,
            modules=None)
    rr, ii, mag, arg, colors, _ = s.get_data()
    assert (rr.shape == ii.shape) and (rr.shape[:2] == colors.shape[:2])
    assert (rr.shape == mag.shape) and (rr.shape == arg.shape)

    s = ComplexSeries(1, (x, -5 - 2 * I, 5 + 2 * I), domain_coloring=True,
            modules="mpmath")
    rr, ii, mag, arg, colors, _ = s.get_data()
    assert (rr.shape == ii.shape) and (rr.shape[:2] == colors.shape[:2])
    assert (rr.shape == mag.shape) and (rr.shape == arg.shape)

    # Corresponds to LineOver1DRangeSeries
    s = InteractiveSeries([S.One], [Tuple(x, -5, 5)])
    s.update_data(dict())
    xx, yy = s.get_data()
    assert len(xx) == len(yy)
    assert np.all(yy == 1)

    # Corresponds to Parametric2DLineSeries
    s = InteractiveSeries([S.One, sin(x)], [Tuple(x, 0, pi)])
    s.update_data(dict())
    xx, yy, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(param))
    assert np.all(xx == 1)

    s = InteractiveSeries([sin(x), S.One], [Tuple(x, 0, pi)])
    s.update_data(dict())
    xx, yy, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(param))
    assert np.all(yy == 1)

    # Corresponds to Parametric3DLineSeries
    s = InteractiveSeries([cos(x), sin(x), S.One], [(x, 0, 2 * pi)])
    s.update_data(dict())
    xx, yy, zz, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(param)) and (len(xx) == len(zz))
    assert np.all(zz == 1)

    s = InteractiveSeries([S.One, sin(x), x], [(x, 0, 2 * pi)])
    s.update_data(dict())
    xx, yy, zz, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(param)) and (len(xx) == len(zz))
    assert np.all(xx == 1)

    s = InteractiveSeries([cos(x), S.One, x], [(x, 0, 2 * pi)])
    s.update_data(dict())
    xx, yy, zz, param = s.get_data()
    assert (len(xx) == len(yy)) and (len(xx) == len(param)) and (len(xx) == len(zz))
    assert np.all(yy == 1)

    # Corresponds to SurfaceOver2DRangeSeries
    s = InteractiveSeries([S.One], [(x, -2, 2), (y, -3, 3)])
    s.update_data(dict())
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)
    assert np.all(zz == 1)

    # Corresponds to ParametricSurfaceSeries
    s = InteractiveSeries([S.One, x, y], [(x, 0, 1), (y, 0, 1)])
    s.update_data(dict())
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)
    assert np.all(xx == 1)

    s = InteractiveSeries([x, S.One, y], [(x, 0, 1), (y, 0, 1)])
    s.update_data(dict())
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)
    assert np.all(yy == 1)

    s = InteractiveSeries([x, y, S.One], [(x, 0, 1), (y, 0, 1)])
    s.update_data(dict())
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)
    assert np.all(zz == 1)

    s = ComplexInteractiveSeries(S.One, (x, -5, 5), real=True, imag=False,
            modules=None)
    s.update_data(dict())
    xx, yy, zz = s.get_data()
    assert (xx.shape == yy.shape) and (xx.shape == zz.shape)

def test_interactive():
    u, x, y, z = symbols("u, x:z")

    # verify that InteractiveSeries produces the same numerical data as their
    # corresponding non-interactive series.
    def do_test(data1, data2):
        assert len(data1) == len(data2)
        for d1, d2 in zip(data1, data2):
            assert np.allclose(d1, d2)

    s1 = InteractiveSeries([u * cos(x)], [(x, -5, 5)], "", params={u: 1}, n1=50)
    s2 = LineOver1DRangeSeries(cos(x), (x, -5, 5), "", adaptive=False, n=50)
    do_test(s1.get_data(), s2.get_data())

    s1 = InteractiveSeries(
        [u * cos(x), u * sin(x)], [(x, -5, 5)], "", params={u: 1}, n1=50
    )
    s2 = Parametric2DLineSeries(cos(x), sin(x), (x, -5, 5), "", adaptive=False, n=50)
    do_test(s1.get_data(), s2.get_data())

    s1 = InteractiveSeries(
        [u * cos(x), u * sin(x), u * x], [(x, -5, 5)], "", params={u: 1}, n1=50
    )
    s2 = Parametric3DLineSeries(cos(x), sin(x), x, (x, -5, 5), "", adaptive=False, n=50)
    do_test(s1.get_data(), s2.get_data())

    s1 = InteractiveSeries(
        [cos(x ** 2 + y ** 2)],
        [(x, -3, 3), (y, -3, 3)],
        "",
        params={u: 1},
        n1=50,
        n2=50,
    )
    s2 = SurfaceOver2DRangeSeries(
        cos(x ** 2 + y ** 2), (x, -3, 3), (y, -3, 3), "", adaptive=False, n1=50, n2=50
    )
    do_test(s1.get_data(), s2.get_data())

    s1 = InteractiveSeries(
        [cos(x + y), sin(x + y), x - y],
        [(x, -3, 3), (y, -3, 3)],
        "",
        params={u: 1},
        n1=50,
        n2=50,
    )
    s2 = ParametricSurfaceSeries(
        cos(x + y),
        sin(x + y),
        x - y,
        (x, -3, 3),
        (y, -3, 3),
        "",
        adaptive=False,
        n1=50,
        n2=50,
    )
    do_test(s1.get_data(), s2.get_data())

    s1 = InteractiveSeries(
        [-u * y, u * x], [(x, -3, 3), (y, -2, 2)], "", params={u: 1}, n1=15, n2=15
    )
    s2 = Vector2DSeries(-y, x, (x, -3, 3), (y, -2, 2), "", n1=15, n2=15)
    do_test(s1.get_data(), s2.get_data())

    s1 = InteractiveSeries(
        [u * z, -u * y, u * x],
        [(x, -3, 3), (y, -2, 2), (z, -1, 1)],
        "",
        params={u: 1},
        n1=15,
        n2=15,
        n3=15,
    )
    s2 = Vector3DSeries(
        z, -y, x, (x, -3, 3), (y, -2, 2), (z, -1, 1), "", n1=15, n2=15, n3=15
    )
    do_test(s1.get_data(), s2.get_data())

    s1 = InteractiveSeries(
        [u * z, -u * y, u * x],
        [(x, -3, 3), (y, -2, 2), (z, -1, 1)],
        "",
        params={u: 1},
        slice=Plane((-1, 0, 0), (1, 0, 0)),
        n1=15,
        n2=15,
        n3=15,
    )
    s2 = SliceVector3DSeries(
        Plane((-1, 0, 0), (1, 0, 0)),
        z,
        -y,
        x,
        (x, -3, 3),
        (y, -2, 2),
        (z, -1, 1),
        "",
        n1=15,
        n2=15,
        n3=15,
    )
    do_test(s1.get_data(), s2.get_data())

    ### Test InteractiveSeries and ComplexInteractiveSeries with complex
    ### functions

    # complex function evaluated over a real line with numpy
    s1 = InteractiveSeries(
        [(z ** 2 + 1) / (z ** 2 - 1)], [(z, -3, 3)], "", n1=50,
        is_complex=True, modules=None)
    s2 = LineOver1DRangeSeries(
        (z ** 2 + 1) / (z ** 2 - 1), (z, -3, 3), "", adaptive=False,
        n=50, is_complex=True, modules=None)
    do_test(s1.get_data(), s2.get_data())

    # complex function evaluated over a real line with mpmath
    s1 = InteractiveSeries(
        [(z ** 2 + 1) / (z ** 2 - 1)], [(z, -3, 3)], "",
        n1=11, is_complex=True, modules="mpmath")
    s2 = LineOver1DRangeSeries(
        (z ** 2 + 1) / (z ** 2 - 1), (z, -3, 3), "", adaptive=False,
        n=11, is_complex=True, modules="mpmath")
    do_test(s1.get_data(), s2.get_data())

    # abs/arg values of complex function evaluated over a real line wit numpy
    expr = (z ** 2 + 1) / (z ** 2 - 1)
    s1 = InteractiveSeries(
        [expr], [(z, -3, 3)], "",
        n1=50, is_complex=True, absarg=expr, modules=None)
    s2 = LineOver1DRangeSeries(
        expr, (z, -3, 3), "", adaptive=False,
        n=50, is_complex=True, absarg=expr, modules=None)
    do_test(s1.get_data(), s2.get_data())

    # abs/arg values of complex function evaluated over a real line wit mpmath
    expr = (z ** 2 + 1) / (z ** 2 - 1)
    s1 = InteractiveSeries(
        [expr], [(z, -3, 3)], "",
        n1=50, is_complex=True, absarg=expr, modules="mpmath")
    s2 = LineOver1DRangeSeries(
        expr, (z, -3, 3), "", adaptive=False,
        n=50, is_complex=True, absarg=expr, modules="mpmath")
    do_test(s1.get_data(), s2.get_data())

    # domain coloring or 3D
    s1 = ComplexInteractiveSeries(
        u * (z ** 2 + 1) / (z ** 2 - 1), (z, -3 - 4 * I, 3 + 4 * I), "",
        n1=20, n2=20, domain_coloring=True, params = {u: 1}, modules=None
    )
    s2 = ComplexSeries(
        (z ** 2 + 1) / (z ** 2 - 1), (z, -3 - 4 * I, 3 + 4 * I), "",
        n1=20, n2=20, domain_coloring=True, modules=None
    )
    do_test(s1.get_data(), s2.get_data())

def test_complex_discretization():
    x, y, z = symbols("x:z")

    # test complex discretization for LineOver1DRangeSeries and
    # SurfaceOver2DRangeSeries and InteractiveSeries

    # is_complex=True should produce (intermediate) complex results, which are
    # later converted to floats. is_complex=False should produce (intermediate)
    # float results.
    s1 = LineOver1DRangeSeries(sqrt(x), (x, -10, 10), "s1",
            adaptive=False, is_complex=False, modules=None, n=10)
    s2 = LineOver1DRangeSeries(sqrt(x), (x, -10, 10), "s1",
            adaptive=False, is_complex=True, modules=None, n=10)
    d1 = s1._uniform_sampling(lambdify([s1.var], s1.expr, s1.modules))
    d2 = s2._uniform_sampling(lambdify([s2.var], s2.expr, s2.modules))
    assert all(isinstance(t, float) for t in d1[0])
    assert all(isinstance(t, float) for t in d1[1])
    assert all(isinstance(t, complex) for t in d2[0])
    assert all(isinstance(t, complex) for t in d2[1])
    d3 = s1.get_data()
    with warns(np.ComplexWarning):
        d4 = s2.get_data()

    # Mpmath and Numpy produces different results
    s1 = LineOver1DRangeSeries(im(sqrt(-x)), (x, -10, 10), "s1",
            adaptive=False, is_complex=True, modules=None, n=10)
    s2 = LineOver1DRangeSeries(im(sqrt(-x)), (x, -10, 10), "s1",
            adaptive=False, is_complex=True, modules="mpmath", n=10)
    d1, d2 = s1.get_data(), s2.get_data()
    assert (d1[-1][-1] < 0) and (d2[-1][-1] > 0)
    assert np.array_equal(d1[-1], -d2[-1])

    def do_test(data1, data2, compare=True):
        assert len(data1) == len(data2)
        for d1, d2 in zip(data1, data2):
            assert (d1.dtype == np.float64) and (d2.dtype == np.float64)
            if compare:
                assert np.array_equal(d1, d2)

    # using Numpy and a real discretization will produce NaN value when x<0.
    with warns(RuntimeWarning, match="invalid value encountered in sqrt"):
        s1 = LineOver1DRangeSeries(sqrt(x), (x, -10, 10), "s1",
                adaptive=False, is_complex=False, modules=None, n=20)
        s1.get_data()

    # using Numpy or Mpmath with complex discretization won't raise warnings.
    # Results between Numpy as Mpmath shoudl be really close
    s2 = LineOver1DRangeSeries(sqrt(x), (x, -10, 10), "s2",
            adaptive=False, is_complex=True, modules=None, n=20)
    s3 = LineOver1DRangeSeries(sqrt(x), (x, -10, 10), "s3",
            adaptive=False, is_complex=True, modules="mpmath", n=20)
    do_test(s2.get_data(), s3.get_data())

    # using Numpy and a real discretization will produce NaN value when x<0.
    with warns(RuntimeWarning, match="invalid value encountered in sqrt"):
        s4 = LineOver1DRangeSeries(sqrt(x), (x, -10, 10), "s4",
                adaptive=True, is_complex=False, modules=None)
        s4.get_data()

    # using Numpy or Mpmath with complex discretization won't raise warnings.
    # Results between Numpy as Mpmath shoudl be really close.
    # NOTE: changed the function because the adaptive algorithm works by
    # checking the collinearity between three points (the x, y coordinates must
    # be real). Instead, with "mpmath" the y coordinate is a complex number.
    s5 = LineOver1DRangeSeries(im(sqrt(x)), (x, -10, 10), "s5",
            adaptive=True, is_complex=True, modules=None)
    s6 = LineOver1DRangeSeries(im(sqrt(x)), (x, -10, 10), "s6",
            adaptive=True, is_complex=True, modules="mpmath")
    # can't directly compare the results because of the adaptive sampling
    do_test(s5.get_data(), s6.get_data(), False)

    # Mpmath and Numpy produces different results
    s1 = SurfaceOver2DRangeSeries(im(sqrt(-x)), (x, -5, 5), (y, -5, 5),
            is_complex=False, modules=None)
    s2 = SurfaceOver2DRangeSeries(im(sqrt(-x)), (x, -5, 5), (y, -5, 5),
            is_complex=True, modules="mpmath")
    d1, d2 = s1.get_data(), s2.get_data()
    assert (d1[-1][-1, -1] < 0) and (d2[-1][-1, -1] > 0)
    assert np.all(np.abs(d1[-1]) - np.abs(d2[-1])) < 1e-08

    # Interactive series produces the same numerical data as LineOver1DRangeSeries.
    # NOTE: InteractiveSeries doesn't support adaptive algorithm!
    s1 = LineOver1DRangeSeries(im(sqrt(-x)), (x, -10, 10), "s1",
            adaptive=False, is_complex=True, modules=None, n=10)
    s2 = InteractiveSeries([im(sqrt(-x))], [(x, -10, 10)], "s2",
            is_complex=True, modules=None, n1=10)
    s3 = InteractiveSeries([im(sqrt(-x))], [(x, -10, 10)], "s3",
            is_complex=True, modules="mpmath", n1=10)
    d1, d2, d3 = s1.get_data(), s2.get_data(), s3.get_data()
    do_test(d1, d2)
    assert np.all(np.abs(d1[-1]) - np.abs(d3[-1])) < 1e-08

    expr = sqrt(-x)
    s1 = LineOver1DRangeSeries(expr, (x, -10, 10), "s1",
            adaptive=False, is_complex=True, modules=None, n=10, absarg=expr)
    s2 = InteractiveSeries([expr], [(x, -10, 10)], "s2",
            is_complex=True, modules=None, n1=10, absarg=expr)
    s3 = InteractiveSeries([expr], [(x, -10, 10)], "s3",
            is_complex=True, modules="mpmath", n1=10, absarg=expr)
    d1, d2, d3 = s1.get_data(), s2.get_data(), s3.get_data()
    do_test(d1, d2)
    assert np.all(np.abs(d1[-1]) - np.abs(d3[-1])) < 1e-08

    # Interactive series produces the same numerical data as SurfaceOver2DRangeSeries
    s1 = SurfaceOver2DRangeSeries(im(sqrt(-x)), (x, -3, 3), (y, -3, 3),
            is_complex=True, modules="mpmath", n1=20, n2=20)
    s2 = InteractiveSeries([im(sqrt(-x))], [(x, -3, 3), (y, -3, 3)], "s2",
            is_complex=True, modules=None, n1=20, n2=20)
    s3 = InteractiveSeries([im(sqrt(-x))], [(x, -3, 3), (y, -3, 3)], "s3",
            is_complex=True, modules="mpmath", n1=20, n2=20)
    do_test(d1, d2)
    assert np.all(np.abs(d1[-1]) - np.abs(d3[-1])) < 1e-08

def test_str():
    x, y, z = symbols("x:z")

    s = LineOver1DRangeSeries(cos(x), (x, -4, 3), "test")
    assert str(s) == "cartesian line: cos(x) for x over (-4.0, 3.0)"
    s = Parametric2DLineSeries(cos(x), sin(x), (x, -4, 3), "test")
    assert str(s) == "parametric cartesian line: (cos(x), sin(x)) for x over (-4.0, 3.0)"
    s = Parametric3DLineSeries(cos(x), sin(x), x, (x, -4, 3), "test")
    assert str(s) == "3D parametric cartesian line: (cos(x), sin(x), x) for x over (-4.0, 3.0)"
    s = SurfaceOver2DRangeSeries(cos(x * y), (x, -4, 3), (y, -2, 5), "test")
    assert str(s) == "cartesian surface: cos(x*y) for x over (-4.0, 3.0) and y over (-2.0, 5.0)"
    s = ContourSeries(cos(x * y), (x, -4, 3), (y, -2, 5), "test")
    assert str(s) == "contour: cos(x*y) for x over (-4.0, 3.0) and y over (-2.0, 5.0)"
    s = ParametricSurfaceSeries(cos(x * y), sin(x * y), x * y,
        (x, -4, 3), (y, -2, 5), "test")
    assert str(s) == "parametric cartesian surface: (cos(x*y), sin(x*y), x*y) for x over (-4.0, 3.0) and y over (-2.0, 5.0)"
    s = ImplicitSeries(x < y, (x, -5, 4), (y, -3, 2), "test")
    assert str(s) == "Implicit expression: x < y for x over (-5.0, 4.0) and y over (-3.0, 2.0)"
    s = ComplexPointSeries(2 + 3 * I, "test")
    assert str(s) == "complex point 2 + 3*I"
    s = ComplexPointSeries([2 + 3 * I, 4 * I], "test")
    assert str(s) == "complex points (2 + 3*I, 4*I)"
    s = ComplexPointInteractiveSeries([2 + 3 * I], "test")
    assert str(s) == "complex interactive points: (2 + 3*I,)"
    s = ComplexPointInteractiveSeries([2 + 3 * I, 4 * I], "test")
    assert str(s) == "complex interactive points: (2 + 3*I, 4*I)"
    s = ComplexSeries(sqrt(z), (z, -2-3j, 4+5j), "test", threed=True)
    assert str(s) == "cartesian surface: sqrt(z) for re(z) over (-2.0, 4.0) and im(z) over (-3.0, 5.0)"
    s = ComplexSeries(sqrt(z), (z, -2-3j, 4+5j), "test", domain_coloring=True)
    assert str(s) == "domain coloring: sqrt(z) for re(z) over (-2.0, 4.0) and im(z) over (-3.0, 5.0)"
    s = ComplexInteractiveSeries(x * sqrt(z), (z, -2-3j, 4+5j), "test",
        threed=True, params={x: 1})
    assert str(s) == "interactive cartesian surface for expression: x*sqrt(z) over (z, (-2-3j), (4+5j)) and parameters [x, z]"
    s = ComplexInteractiveSeries(x * sqrt(z), (z, -2-3j, 4+5j), "test",
        domain_coloring=True, params={x: 1})
    assert str(s) == "interactive domain coloring for expression: x*sqrt(z) over (z, (-2-3j), (4+5j)) and parameters [x, z]"
    s = Vector2DSeries(-y, x, (x, -5, 4), (y, -3, 2), "test")
    assert str(s) == "2D vector series: [-y, x] over (x, -5.0, 4.0), (y, -3.0, 2.0)"
    s = Vector3DSeries(z, y, x, (x, -5, 4), (y, -3, 2), (z, -6, 7), "test")
    assert str(s) == "3D vector series: [z, y, x] over (x, -5.0, 4.0), (y, -3.0, 2.0), (z, -6.0, 7.0)"
    s = SliceVector3DSeries(Plane((0, 0, 0), (1, 0, 0)), z, y, x,
        (x, -5, 4), (y, -3, 2), (z, -6, 7), "test")
    assert str(s) == "sliced 3D vector series: [z, y, x] over (x, -5.0, 4.0), (y, -3.0, 2.0), (z, -6.0, 7.0) at Plane(Point3D(0, 0, 0), (1, 0, 0))"
    s = PlaneSeries(Plane((0, 0, 0), (1, 1, 1)),
        (x, -5, 4), (y, -3, 2), (z, -6, 7), "test")
    assert str(s) == "plane series of Plane(Point3D(0, 0, 0), (1, 1, 1)) over (x, -5, 4), (y, -3, 2), (z, -6, 7)"
    s = PlaneInteractiveSeries([Plane((z, 0, 0), (1, 1, 1))],
        [(x, -5, 4), (y, -3, 2), (z, -6, 7)], "test", params={z: 1})
    assert str(s) == "interactive plane series of Plane(Point3D(z, 0, 0), (1, 1, 1)) over (x, -5, 4), (y, -3, 2), (z, -6, 7) with parameters [z]"
    s = GeometrySeries(Circle(Point(0, 0), 5))
    assert str(s) == "geometry entity: Circle(Point2D(0, 0), 5)"
    s = GeometryInteractiveSeries([Circle(Point(x, 0), 5)], [], params={x: 1})
    assert str(s) == "interactive geometry entity: Circle(Point2D(x, 0), 5) with parameters [x]"

    # interactive series
    s = InteractiveSeries([z * cos(x)], [(x, -4, 3)], "test", params={z: 1})
    assert str(s) == "interactive expression: z*cos(x) with ranges (x, -4.0, 3.0) and parameters [x, z]"
    s = InteractiveSeries([z * cos(x * y)], [(x, -4, 3), (y, -2, 1)], "test",
        params={z: 1})
    assert str(s) == "interactive expression: z*cos(x*y) with ranges (x, -4.0, 3.0), (y, -2.0, 1.0) and parameters [x, y, z]"
    s = InteractiveSeries([z * cos(x * y), sin(x * y), x*y],
        [(x, -4, 3), (y, -2, 1)], "test", params={z: 1})
    assert str(s) == "interactive expression: (z*cos(x*y), sin(x*y), x*y) with ranges (x, -4.0, 3.0), (y, -2.0, 1.0) and parameters [x, y, z]"

def test_piecewise():
    x = symbols("x")

    # Test that univariate Piecewise objects are processed in such a way to
    # create multiple series, each one with the correct range

    f = Piecewise(
        (-1, x < -1),
        (x, And(-1 <= x, x < 0)),
        (x**2, And(0 <= x, x < 1)),
        (x**3, x >= 1)
    )
    s = _process_piecewise(f, (x, -5, 5), "A")
    assert len(s) == 4
    assert all(isinstance(t, LineOver1DRangeSeries) for t in s)
    assert (s[0].expr == -1) and (s[0].start == -5) and (s[0].end == -1)
    assert (s[1].expr == x) and (s[1].start == -1) and (s[1].end == 0)
    assert (s[2].expr == x**2) and (s[2].start == 0) and (s[2].end == 1)
    assert (s[3].expr == x**3) and (s[3].start == 1) and (s[3].end == 5)
    labels = ["A" + str(i + 1) for i in range(5)]
    assert all(t.label == l for t, l in zip(s, labels))

    f = Piecewise(
        (1, x < -5),
        (x, Eq(x, 0)),
        (x**2, Eq(x, 2)),
        (x**3, (x > 0) & (x < 2)),
        (x**4, True)
    )
    s = _process_piecewise(f, (x, -10, 10), "B")
    assert len(s) == 6
    assert all(isinstance(t, LineOver1DRangeSeries) for t in [s[0], s[3], s[4], s[5]])
    assert all(isinstance(t, List2DSeries) for t in [s[1], s[2]])
    assert (s[0].expr == 1) and (s[0].start == -10) and (s[0].end == -5)
    assert (np.allclose(s[1].list_x, np.array([0.])) and
        np.allclose(s[1].list_y, np.array([0.])))
    assert (np.allclose(s[2].list_x, np.array([2.])) and
        np.allclose(s[2].list_y, np.array([4.])))
    assert (s[3].expr == x**3) and (s[3].start == 0) and (s[3].end == 2)
    assert (s[4].expr == x**4) and (s[4].start == -5) and (s[4].end == 0)
    assert (s[5].expr == x**4) and (s[5].start == 2) and (s[5].end == 10)
    labels = ["B" + str(i + 1) for i in range(5)] + ["B5"]
    assert all(t.label == l for t, l in zip(s, labels))

    f = Piecewise((x, Interval(0, 1).contains(x)), (0, True))
    s = _process_piecewise(f, (x, -10, 10), "C")
    assert len(s) == 3
    assert all(isinstance(t, LineOver1DRangeSeries) for t in s)
    assert (s[0].expr == x) and (s[0].start == 0) and (s[0].end == 1)
    assert (s[1].expr == 0) and (s[1].start == -10) and (s[1].end == 0)
    assert (s[2].expr == 0) and (s[2].start == 1) and (s[2].end == 10)
    labels = ["C1", "C2", "C2"]
    assert all(t.label == l for t, l in zip(s, labels))

    f = Piecewise((x, Interval(0, 1, False, True).contains(x)), (0, True))
    s = _process_piecewise(f, (x, -10, 10), "D")
    assert len(s) == 3
    assert all(isinstance(t, LineOver1DRangeSeries) for t in s)
    assert (s[0].expr == x) and (s[0].start == 0) and (s[0].end == 1)
    assert (s[1].expr == 0) and (s[1].start == -10) and (s[1].end == 0)
    assert (s[2].expr == 0) and (s[2].start == 1) and (s[2].end == 10)
    labels = ["D1", "D2", "D2"]
    assert all(t.label == l for t, l in zip(s, labels))

    f = Piecewise((x, x < 1), (x**2, -1 <= x), (x, 3 < x))
    s = _process_piecewise(f, (x, -10, 10), "E")
    assert len(s) == 2
    assert all(isinstance(t, LineOver1DRangeSeries) for t in s)
    assert (s[0].expr == x) and (s[0].start == -10) and (s[0].end == 1)
    assert (s[1].expr == x**2) and (s[1].start == 1) and (s[1].end == 10)

    # NotImplementedError: as_set is not implemented for relationals with
    # periodic solutions
    p1 = Piecewise((cos(x), x < 0), (0, True))
    f = Piecewise((0, Eq(p1, 0)), (p1 / Abs(p1), True))
    raises(NotImplementedError, lambda: _process_piecewise(f, (x, -10, 10), "F"))

    f = Piecewise((1 - x, (x >= 0) & (x < 1)), (0, True))
    s = _process_piecewise(f, (x, -10, 10), "test")
    assert len(s) == 3
    assert all(isinstance(t, LineOver1DRangeSeries) for t in s)
    assert (s[0].expr == 1 - x) and (s[0].start == 0) and (s[0].end == 1)
    assert (s[1].expr == 0) and (s[1].start == -10) and (s[1].end == 0)
    assert (s[2].expr == 0) and (s[2].start == 1) and (s[2].end == 10)

    # The range is smaller than the function "domain"
    f = Piecewise(
        (1, x < -5),
        (x, Eq(x, 0)),
        (x**2, Eq(x, 2)),
        (x**3, (x > 0) & (x < 2)),
        (x**4, True)
    )
    s = _process_piecewise(f, (x, -3, 3), "A")
    labels = ["A2", "A3", "A4", "A5", "A5"]
    assert all(t.label == l for t, l in zip(s, labels))


import cPickle
import numpy as np
from Bio.SVDSuperimposer import SVDSuperimposer
from scipy.linalg import sqrtm, inv

def sym(w):
    return w.dot(inv(sqrtm(w.T.dot(w))))

def params2cc(parameters, return_steps=False):
    steps = []

    h2_tran = parameters[-3:]

    h2_rot = np.reshape(parameters[-12:-3], (3,3))
    helical_params = np.reshape(parameters[:-12], (2,3))

    h2_rot = sym(h2_rot)

    h1_ref = du_mean_helix
    if return_steps:
        step_coords = np.append(h1_ref,h1_ref,axis=0)
        steps.append(step_coords)

    h1_dev, h2_dev = du_pca_helix.inverse_transform(helical_params)
    h1_dev = np.reshape(h1_dev, (h1_ref.shape[0],h1_ref.shape[1]))
    h2_dev = np.reshape(h2_dev, (h1_ref.shape[0],h1_ref.shape[1]))

    h1 = h1_ref+h1_dev
    h2 = h1_ref+h2_dev

    if return_steps:
        step_coords = np.append(h1,h2,axis=0)
        steps.append(step_coords)

    h2_new = np.dot(h2, h2_rot)
    if return_steps:
        step_coords = np.append(h1,h2_new,axis=0)
        steps.append(step_coords)

    h2_new = h2_new + h2_tran
    if return_steps:
        step_coords = np.append(h1,h2_new,axis=0)
        steps.append(step_coords)

    orig_coords = np.append(h1,h2_new,axis=0)

    if return_steps:
        return steps

    return orig_coords

def cc2params(coords):

    sup=SVDSuperimposer()

    n_atoms_mono = int(coords.shape[0]/2)

    h1 = coords[:n_atoms_mono]
    h2 = coords[n_atoms_mono:]
    h1_ref = du_mean_helix

    # align h1 and h2 with mean angles to the ref helix

    sup.set(h1_ref, h1)
    sup.run()
    h1_aligned_ref = sup.get_transformed()

    sup.set(h1_ref, h2)
    sup.run()
    h2_aligned_ref = sup.get_transformed()

    # estimate parameters from pca
    # center to h1_ref
    h1_aligned_ref = h1_aligned_ref - h1_ref
    h2_aligned_ref = h2_aligned_ref - h1_ref
    # unwrap
    h1_aligned_ref = np.reshape(h1_aligned_ref, (h1_aligned_ref.shape[0]*h1_aligned_ref.shape[1]))
    h2_aligned_ref = np.reshape(h2_aligned_ref, (h2_aligned_ref.shape[0]*h2_aligned_ref.shape[1]))
    # get params
    helical_params =du_pca_helix.transform([h1_aligned_ref, h2_aligned_ref])
    h1_aligned_ref, h2_aligned_ref = du_pca_helix.inverse_transform(helical_params)
    h1_aligned_ref = np.reshape(h1_aligned_ref, (h1_ref.shape[0],h1_ref.shape[1]))
    h2_aligned_ref = np.reshape(h2_aligned_ref, (h1_ref.shape[0],h1_ref.shape[1]))

    # construct ideal helices

    h1_transformed = h1_ref+h1_aligned_ref
    h2_transformed = h1_ref+h2_aligned_ref

    # adjust hi_helix1 and hi_helix2 by the parameters

    # align h1 to h1 ideal and transform all coords

    sup.set(h1, h1_transformed)
    sup.run()
    h1_ideal = sup.get_transformed()

    # align h2_ideal to h2

    sup.set(h2, h2_transformed)
    sup.run()
    h2_ideal = sup.get_transformed()

    coords_ideal = np.append(h1_ideal,h2_ideal,axis=0)

    # align ideal coords to the ref helix
    sup.set(h1_transformed,h1_ideal)
    sup.run()
    (rot_ref, tran_ref) = sup.get_rotran()

    coords_ideal = np.dot(coords_ideal,rot_ref) + tran_ref

    h1_new = coords_ideal[:n_atoms_mono]
    h2_new = coords_ideal[n_atoms_mono:]

    sup.set(h2_new, h2_transformed)
    sup.run()
    (rot2, tran2) = sup.get_rotran()
    h2_rot = rot2.flatten()
    helical_params = helical_params.flatten()
    transform_params = np.append(helical_params,np.append(h2_rot, tran2))

    return transform_params

du_pca_helix, du_mean_helix = cPickle.load(open('helix_template.pkl', "rb"))