lasso_regression_large.py

#!/usr/bin/env python
# coding: utf-8

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rc
from scipy.optimize import minimize_scalar, minimize
import datetime
import os
import networkx as nx
import pandas as pd 
from gurobipy import*
from collections import Counter, defaultdict
from scipy.io import loadmat
from time import process_time
from matplotlib import ticker
import warnings
warnings.filterwarnings("ignore")
from autograd import grad as grad_a
from sklearn.preprocessing import StandardScaler


# # Oracles

# In[2]:



def line_search(x, d, gamma_max,func):
    
    #line-search using Brent's rule in scipy
    
    '''
    Minimizes f over [x, y], i.e., f(x+gamma*d) as a function of scalar gamma in [0,gamma_max]
    '''

    def fun(gamma):
        ls = x + gamma*d
        return func(ls)

    res = minimize_scalar(fun, bounds=(0, gamma_max), method='bounded')

    gamma = res.x
    ls = x + gamma*d        
    return ls, gamma


def segment_search(f, grad_f, x, y, tol=1e-6, stepsize=True):
    
    #line-search using golden-section rule coded from scratch
    
    '''
    Minimizes f over [x, y], i.e., f(x+gamma*(y-x)) as a function of scalar gamma in [0,1]
    '''
    
    # restrict segment of search to [x, y]
    d = (y-x).copy()
    left, right = x.copy(), y.copy()
    
    # if the minimum is at an endpoint
    if np.dot(d, grad_f(x))*np.dot(d, grad_f(y)) >= 0:
        if f(y) <= f(x):
            return y, 1
        else:
            return x, 0
    
    # apply golden-section method to segment
    gold = (1+np.sqrt(5))/2
    improv = np.inf
    while improv > tol:
        old_left, old_right = left, right
        new = left+(right-left)/(1+gold)
        probe = new+(right-new)/2
        if f(probe) <= f(new):
            left, right = new, right
        else:
            left, right = left, probe
        improv = np.linalg.norm(f(right)-f(old_right))+np.linalg.norm(f(left)-f(old_left))
    x_min = (left+right)/2
    
    # compute step size gamma
    gamma = 0
    if stepsize == True:
        for i in range(len(d)):
            if d[i] != 0:
                gamma = (x_min[i]-x[i])/d[i]
                break
                
    return x_min, gamma


def linear_oracle(A,b, c):
    
    '''
    General form LP solver that solves min c^Tx subject to Ax <= b
    '''

    m = Model("opt")
    n = len(A.T)

    #define variables
    x = []
    for i in range(n):
        x.append(m.addVar(lb=-GRB.INFINITY, name='x_{}'.format(i)))

    m.update()              

    objExp = np.dot(np.array(x),c)
    m.setObjective(objExp, GRB.MINIMIZE)
    m.update()

    #feasibility constraints
    for i in range(len(A)):
        m.addConstr(np.dot(np.array(x),A[i]),'<=', b[i])

    m.update()

    #convex hull constraint

    m.update()

    #optimize
    m.setParam( 'OutputFlag', False )
    m.optimize()

    return np.array([i.x for i in x])



def active_constraints_oracle(A,b,x):
    
    '''
    finding active constraints for polytope of the form Ax <= b
    '''
    b_prime = []
    A_prime = []
    A_not = []
    b_not = []
    for i in range(len(A)):
        if np.abs(np.round(np.dot(A[i],x),4) - b[i]) < 0.01:
            A_prime.append(list(A[i]))
            b_prime.append(b[i])
        else:
            A_not.append(list(A[i]))
            b_not.append(b[i])
            
    return np.array(A_prime),np.array(b_prime),np.array(A_not),np.array(b_not)


def max_step_size(A,b,x, d):
    
    
    '''
    finding maximum step-size: argmax{\delta | x + \delta d \in P} where P is given by Ax <= b and d is a 
    feasible direction at x. Here we use Gurobi as a black box solver.
    '''

    m = Model("opt")
    n = len(A.T)

    lam = m.addVar(lb=-GRB.INFINITY, name='lam')

    m.update()              

    objExp = lam
    m.setObjective(objExp, GRB.MAXIMIZE)
    
    m.update()

    #feasibility constraints
    
    for i in range(len(A)):
        m.addConstr(np.dot(np.array(x),A[i]) + lam* np.dot(np.array(d),A[i]),'<=', b[i])
        
    m.update()

    #optimize
    m.setParam( 'OutputFlag', False )
    m.optimize()
    
    return lam.x
    

def max_step_size_search(A,b,x, d):
    
    '''
    finding maximum step-size: argmax{\delta | x + \delta d \in P} where P is given by Ax <= b and d is a 
    feasible direction at x. Here we search over set of inactive constraints that might get violated.
    '''

    A_J,b_J = active_constraints_oracle(A,b,x)[2:]
    slack = b_J - np.dot(A_J,x)
    
    try:
        if len(d) > 1:
            denom = np.dot(A_J,d)
            excess = [slack[i]/denom[i] for i in range(len(slack)) if denom[i] > 0]
            gamma_max = min(excess)
            return gamma_max
    except:
        return min(slack)


def shadow_oracle(A,b,x, d):
     
    '''
    Projecting d on set of feasible directions for polytope given by Ax <= b
    '''
    
    #find matrix of active constraints
    A_I = active_constraints_oracle(A,b,x)[0]

    m = Model("opt")
    n = len(A_I.T)

    z = []
    for i in range(n):
        z.append(m.addVar(lb=-GRB.INFINITY, name='z_{}'.format(i)))

    m.update()              

    objExp = quicksum(np.square(np.array([d[k] -z[k] for k in range(n)])))
   
    m.setObjective(objExp, GRB.MINIMIZE)
    m.update()

    #feasibility constraints
    for i in range(len(A_I)):
        m.addConstr(np.dot(np.array(z),A_I[i]),'<=', 0)
        
    m.update()

    #optimize
    m.setParam( 'OutputFlag', False )
    m.optimize()

    der = [i.x for i in z]

    return np.array(der)


def shadow_oracle_other(A,b,x, grad,M,vert_rep,vertices):
     
    '''
    Computing projection if we know an upper bound M on the value at which normal cone is not changing
    '''
    
    epsilon = M/2
    
    if vert_rep == True:
        g = projection_oracle_vertices(vertices,x - epsilon*grad)
        der = (g -x)/epsilon
        
    else:
        projection_oracle(x - epsilon*grad,A,b)
        der = (g -x)/epsilon

    return der

    return np.array(der)


def project_Null(A,d):
    
    '''
       Function to project the direction d onto the nullspace of matrix A
    '''
    
    if len(A) > 1:
        return np.dot(np.eye(len(A.T)) - np.matmul(A.T,np.matmul(np.linalg.pinv(np.matmul(A,A.T)),A)),d)
        
    else:
        A = A[0]
        return d - np.dot(A,d)*A/np.dot(A,A)
    
    

def in_face_shadow_oracle (A,b,x,d):
    
    '''
       Function to compute in-face shadow, i.e. project the direction d onto minimal face of x for polytope given by Ax <= b
    '''
    
    A_I = active_constraints_oracle(A,b,x)[0]
    d_hat = project_Null(A_I,d)
    
    return d_hat

def in_face_shadow_oracle_other(A,b,x, d):
     
    '''
        Function to compute in-face shadow, i.e. project the direction d onto minimal face of x for polytope given by Ax <= b
    '''
    
    #find matrix of active constraints
    A_I = active_constraints_oracle(A,b,x)[0]

    m = Model("opt")
    n = len(A_I.T)

    z = []
    for i in range(n):
        z.append(m.addVar(lb=-GRB.INFINITY, name='z_{}'.format(i)))

    m.update()              

    objExp = quicksum(np.square(np.array([d[k] -z[k] for k in range(n)])))
   
    m.setObjective(objExp, GRB.MINIMIZE)
    m.update()

    #feasibility constraints
    for i in range(len(A_I)):
        m.addConstr(np.dot(np.array(z),A_I[i]),'==', 0)
        
    m.update()

    #optimize
    m.setParam( 'OutputFlag', False )
    m.optimize()

    der = [i.x for i in z]

    return np.array(der)



def find_lam_hat(A_I, x_0,x,d,d_hat,gamma_total):
    
    
    '''
       Function to compute the maximum movement along the in-face shadow using a linear program
    '''

    m = Model("opt")
    n = len(A_I.T)

    #define variables
    lam = m.addVar(name='lam')

    #define variables
    mu = []
    for i in range(len(A_I)):
        mu.append(m.addVar(lb=0, name='mu_{}'.format(i)))

    m.update()

    m.setObjective(lam, GRB.MAXIMIZE)

    m.update()

    #feasibility constraints
    for i in range(n):
        m.addConstr(np.dot(np.array(mu),A_I.T[i]),'==', x_0[i] + lam*d[i] - x[i] -(lam - gamma_total)*d_hat[i])

    m.update()
    
    m.setParam( 'OutputFlag', False )
    m.setParam( 'DualReductions', 0)

    #optimize
    m.optimize()
    
    if m.status == GRB.INFEASIBLE:
        return 0
    else:
        return lam.x - gamma_total

def trace_in_face(A,b,x_0,x,d,d_hat, gamma_total, lmo):
    
    '''
       Function to run the trace-in-face procedure
    '''
    
    A_I = active_constraints_oracle(A,b,x)[0]

    #line-search along d_hat
    if all(np.round(d_hat,6) != 0):
        gamma_max = max_step_size_search(A,b,x, d_hat)
    
        #check weather this optimal by checking first order optimality
        normal = x_0 + (gamma_max +gamma_total)*d - x - gamma_max*d_hat
        if np.dot(normal, LO(-normal) - x - gamma_max*d_hat) <= 0:
            lam_hat =  gamma_max
        else:                     #if not solve Lp to find max step size where normal cone does not change
            lam_hat = find_lam_hat(A_I, x_0,x,d,d_hat,gamma_total) 
    else:
        lam_hat = find_lam_hat(A_I, x_0,x,d,d_hat,gamma_total) 
    
    return lam_hat


# # Algorithms - FW Variants

# # FW

# In[3]:


def FW(x, lmo, epsilon,func,grad_f, f_tol, time_tol):
    
    #record primal gap, function value, and time every iteration
    now=datetime.datetime.now()
    primal_gap = []
    function_value=[func(x)]
    time = [0]
    f_improv = np.inf

    #initialize iteration count
    t = 0    

    while f_improv > f_tol and time[-1] < time_tol:
        
        start = process_time()
        
        #compute gradient
        grad = grad_f(x)

        #solve linear subproblem and compute FW direction
        v = lmo(grad)
        d_FW = v-x

        #If primal gap is small enough - terminate
        if np.dot(-grad,d_FW) <= epsilon:
            break
        else:
            #update convergence data
            primal_gap.append(np.dot(-grad,d_FW))

        #Update next iterate by doing a feasible line-search
        x, gamma = segment_search(func, grad_f, x, v)
        
        end = process_time()
        
        time.append(time[t] + end - start)
        f_improv = function_value[-1] - func(x)
        function_value.append(func(x))
        
        t+=1
        
    return x, function_value, time,t,primal_gap


# # AFW 

# In[4]:


#Function to compute away vertex
def away_step(grad, S):
    
    '''
    Compute away vertex by searching over current active set
    '''
    
    costs = {}
    
    for k,v in S.items():
        cost = np.dot(k,grad)
        costs[cost] = [k,v]
    vertex, alpha = costs[max(costs.keys())]  
    return vertex,alpha

#Function to update active set
def update_S(S,gamma, Away, vertex):
    
    '''
    Update convex decompistion of active step after every iteration 
    '''
    
    S = S.copy()
    vertex = tuple(vertex)
    
    if not Away:
        if vertex not in S.keys():
            S[vertex] = gamma
        else:
            S[vertex] *= (1-gamma)
            S[vertex] += gamma
            
        for k in S.keys():
            if k != vertex:
                S[k] *= (1-gamma)
    else:
        for k in S.keys():
            if k != vertex:
                S[k] *= (1+gamma)
            else:
                S[k] *= (1+gamma)
                S[k] -= gamma
    return {k:v for k,v in S.items() if np.round(v,3) > 0}


#AFW Algorithm
def AFW(x, lmo, epsilon,func,grad_f, f_tol, time_tol):
    
    #record primal gap, function value, and time every iteration
    now=datetime.datetime.now()
    primal_gap = []
    function_value=[func(x)]
    time = [0]
    f_improv = np.inf

    #initialize starting point and active set
    t = 0    
    S={tuple(x): 1}

    while f_improv > f_tol and time[-1] < time_tol:
        
        start = process_time()
        
        #compute gradient
        grad = grad_f(x)

        #solve linear subproblem and compute FW direction
        v = lmo(grad)
        d_FW = v-x

        #If primal gap is small enough - terminate
        if np.dot(-grad,d_FW) <= epsilon:
            break
        else:
            #update convergence data
            primal_gap.append(np.dot(-grad,d_FW))

        #Compute away vertex and direction
        a,alpha_a = away_step(grad, S)
        d_A = x - a

        #Check if FW gap is greater than away gap
        if np.dot(-grad,d_FW) >= np.dot(-grad,d_A):
            #choose FW direction
            d = d_FW
            vertex = v
            gamma_max = 1
            Away = False
        else:
            #choose Away direction
            d = d_A
            vertex = a
            gamma_max = alpha_a/(1-alpha_a)
            Away = True

        #Update next iterate by doing a feasible line-search
        x, gamma = segment_search(func, grad_f, x, x + gamma_max *d)

        #update active set based on direction chosen
        S = update_S(S,gamma, Away, vertex)
        
        end = process_time()
        time.append(time[t] + end - start)
        
        f_improv = function_value[-1] - func(x)
        function_value.append(func(x))
        
        t+=1
        
    return x, function_value, time,t,primal_gap


# # Pairwise FW

# In[5]:


def update_S_PW(S,gamma, FW_vertex, Away_vertex):
    
    '''
    Update convex decompistion of active step after pairwise direction is chosen 
    '''
    
    vertex1 = tuple(FW_vertex)
    vertex2 = tuple(Away_vertex)
    S = S.copy()
    
    if vertex1 not in S.keys():
        S[vertex1] = gamma
    else:
        S[vertex1] += gamma
    
    S[vertex2] -= gamma
    return {k:v for k,v in S.items() if np.round(v,4) > 0}


#PFW Algorithm
def PFW(x, lmo, epsilon,func,grad_f, f_tol, time_tol):
    
    #record primal gap, function value, and time every iteration

    primal_gap = []
    function_value=[func(x)]
    time = [0]
    f_improv = np.inf

    #initialize starting point and active set
    t = 0    
    S={tuple(x): 1}

    while f_improv > f_tol and time[-1] < time_tol:
        
        start = process_time()
        
        #compute gradient
        grad = grad_f(x)

        #solve linear subproblem and compute FW direction
        v = lmo(grad)
        d_FW = v-x

        #If primal gap is small enough - terminate
        if np.dot(-grad,d_FW) <= epsilon:
            break
        else:
            #update convergence data
            primal_gap.append(np.dot(-grad,d_FW))

        #Compute away vertex and direction
        a,alpha_a = away_step(grad, S)
        d_A = x - a

        #Pairwise step
        d =  d_FW + d_A
        gamma_max = alpha_a

        #Update next iterate by doing a feasible line-search
        x, gamma = line_search(x, d, gamma_max,func)

        #update active set
        S = update_S_PW(S,gamma, v, a)
        
        end = process_time()
        time.append(time[t] + end - start)
        f_improv = function_value[-1] - func(x)
        function_value.append(func(x))
        
        t+=1
        
    return x, function_value, time,t,primal_gap


# # Decoposition Invaraint FW (DICG)

# In[6]:


#DICG algorithm
def DICG(x, lmo, feasibility_oracle, epsilon,func,grad_f, f_tol, time_tol):
    
    #record primal gap, function value, and time every iteration

    primal_gap = []
    function_value=[func(x)]
    time = [0]
    f_improv = np.inf

    #initialize starting point and active set
    t = 0    

    while f_improv > f_tol and time[-1] < time_tol:
        
        start = process_time()
        
        #compute gradient
        grad = grad_f(x)

        #solve linear subproblem and compute FW direction
        v = lmo(grad)
        d_FW = v-x

        #If primal gap is small enough - terminate
        if np.dot(-grad,d_FW) <= epsilon:
            break
        else:
            #update convergence data
            primal_gap.append(np.dot(-grad,d_FW))

        #create new gradient to find best away vertex
        g = np.array([grad[i] if x[i] > 0 else -9e9 for i in range(len(x))])
        a = lmo(-g)
        d_A = x - a

        #Pairwise step
        d =  d_FW + d_A
        gamma_max = feasibility_oracle(x,d)

        #Update next iterate by doing a feasible line-search
        x, gamma = segment_search(func, grad_f, x, x+gamma_max*d)
        
        end = process_time()
        time.append(time[t] + end - start)
        f_improv = function_value[-1] - func(x)
        function_value.append(func(x))
        
        t+=1
        
    return x, function_value, time,t,primal_gap


# # Projected Gradient variants

# # Shadow Walk

# In[85]:


def trace_PW_curve_app(x,grad,shadow,feasibility_oracle,func):

    '''
    trace the piecewise linear projection curve approximately by only taking shadow steps
    '''
    
    count_in_face = 0
    count_shadow = 0
    t_all = 0
    t_all_but_in_face = 0
    t_shadow = 0
    t_lam = 0
    t_none = 0
    
    gamma_total = 0
    d_pi = shadow(x,-grad)
    
    
    while np.dot(d_pi,d_pi)**0.5 > 0.00001:
    
        t1 = process_time()
        
        t2,t3,t4,t5 = 0,0,0,0
        
        # only take shadow steps
        count_shadow +=1
        gamma_max = feasibility_oracle(x, d_pi)
        x, gamma = segment_search(func, grad_f, x, x + gamma_max *d_pi)

        #update total step-size accrued and shadow
        gamma_total += gamma*gamma_max
        
        #break if sufficient progress
        if abs(gamma - 1) > 0.001:
            t6,t7 = 0,0
            break
        else:
            t6 = process_time()
            d_pi = shadow(x,-grad)
            t7 = process_time()
            
        t8 = process_time()
        t_all += t8 - t1 - (t3 - t2) - (t5 - t4) - (t7 - t6)
        t_all_but_in_face += t8 - t1 - (t5 - t4) - (t7 - t6)
        t_none += t8 - t1
        t_shadow += t8 - t1 - (t5 - t4)
        t_lam += t8 - t1 - (t7 - t6)
        print(func(x))
    return x,count_in_face,count_shadow,(t_none, t_lam, t_shadow, t_all_but_in_face, t_all),d_pi



def trace_PW_curve(x,grad,shadow,in_face_shadow,in_face_trace,feasibility_oracle,func,grad_f):
    
    '''
    trace the piecewise linear projection curve
    '''

    count_in_face = 0
    count_shadow = 0
    t_all = 0
    t_all_but_in_face = 0
    t_shadow = 0
    t_lam = 0
    t_none = 0
    
    gamma_total = 0
    d_pi = shadow(x,-grad)
    x_0 = x.copy()
    
    while np.dot(d_pi,d_pi)**0.5 > 1e-6:
    
        t1 = process_time()
        #print(np.dot(x_0 - gamma_total*grad - x,d_pi),func(x))
        
        if abs(np.dot(x_0 - gamma_total*grad - x,d_pi)) > 1e-6:
            
            #in-face step
            t2 = process_time()
            d_hat = in_face_shadow(x,-grad)
            t3 = process_time()
            
            count_in_face +=1
            
            t4 = process_time()
            gamma_max = in_face_trace(x_0,x,-grad, gamma_total,d_hat)
            
            if np.round(gamma_max,4) > 0:
                t5 = process_time()
                if all(np.round(d_hat,6) == 0):
                    gamma = 1
                else:
                    x, gamma = segment_search(func, grad_f, x, x + gamma_max *d_hat)
                
            else:
                t2,t3,t4,t5 = 0,0,0,0
                #shadow step
                count_shadow +=1
                gamma_max = feasibility_oracle(x, d_pi)
                x, gamma = segment_search(func, grad_f, x, x + gamma_max*d_pi)

        else:
            t2,t3,t4,t5 = 0,0,0,0
            #shadow step
            count_shadow +=1
            gamma_max = feasibility_oracle(x, d_pi)
            x, gamma = segment_search(func, grad_f, x, x + gamma_max*d_pi)
        
        #update total step-size accrued and shadow
        gamma_total += gamma*gamma_max
        
        #break if sufficient progress
        if abs(gamma - 1) > 0.001:
            t6,t7 = 0,0
            break
        else:
            t6 = process_time()
            d_pi = shadow(x,-grad)
            t7 = process_time()
        
        t8 = process_time()
        t_all += t8 - t1 - (t3 - t2) - (t5 - t4) - (t7 - t6)
        t_all_but_in_face += t8 - t1 - (t5 - t4) - (t7 - t6)
        t_none += t8 - t1
        t_shadow += t8 - t1 - (t5 - t4)
        t_lam += t8 - t1 - (t7 - t6)
        
    return x,count_in_face,count_shadow,(t_none, t_lam, t_shadow, t_all_but_in_face, t_all),d_pi


#shadow descent algorithm:
def SD(x, lmo, shadow,in_face_shadow,in_face_trace, feasibility_oracle, epsilon,func,grad_f, f_tol, time_tol,tol,shadow_tol=1e-6):
    
    #record primal gap, function value, and time every iteration

    primal_gap = []
    function_value=[func(x)]
    time1 = [0]
    time2 = [0]
    f_improv = np.inf
    counts_shadow = []
    counts_in_face = []

    #initialize starting point and active set
    t = 0    

    while f_improv > f_tol and time1[-1] < time_tol:
        
        start = process_time()
        
        #compute gradient
        grad = grad_f(x)

        #solve linear subproblem and compute FW direction
        s = lmo(grad)
        d_FW = s-x

        #If primal gap is small enough - terminate
        if np.dot(-grad,d_FW) <= epsilon:
            break
        else:
            #update convergence data
            primal_gap.append(np.dot(-grad,d_FW))
            
        end = process_time()
        
        x, c1,c2, t_pw,d_pi = trace_PW_curve(x,grad,shadow,in_face_shadow,in_face_trace,feasibility_oracle,func,grad_f)
        #x, c1,c2, t_pw,d_pi = trace_PW_curve_app(x,grad,shadow,feasibility_oracle,func)
        
        if np.dot(d_pi,d_pi)**0.5 < shadow_tol:
            break
        else:
            counts_in_face.append(c1)
            counts_shadow.append(c2)

            time1.append(time1[t] + end - start + t_pw[0])
            time2.append( time2[t] + end - start + t_pw[-2])

            f_improv = function_value[-1] - func(x)
            function_value.append(func(x))
            t+=1
        
        #print(func(x))
        
    return x, function_value, time1,time2,t,primal_gap,counts_shadow,counts_in_face


# # Shadow CG

# In[38]:


#shadow conditional gradients
def SCG(x, lmo, shadow,in_face_shadow,in_face_trace, feasibility_oracle, epsilon,func,grad_f, f_tol, time_tol,tol,shadow_tol=1e-6):
    
    #record primal gap, function value, and time every iteration

    primal_gap = []
    function_value=[func(x)]
    time1 = [0]
    time2 = [0]
    f_improv = np.inf
    counts_shadow = []
    counts_in_face = []
    FW = []

    #initialize starting point and active set
    t = 0    

    while f_improv > f_tol and time1[-1] < time_tol:
        
        t1 = process_time()
        #compute gradient
        grad = grad_f(x)

        #solve linear subproblem and compute FW direction
        v = lmo(grad)
        d_FW = v-x

        #If primal gap is small enough - terminate
        if np.dot(-grad,d_FW) <= epsilon:
            break
        else:
            #update convergence data
            primal_gap.append(np.dot(-grad,d_FW))
            
        #Compute directional derivative
        t2 = process_time()
        d_pi = shadow(x,-grad)
        t3 = process_time()
        
        #Check if FW gap direction is better than normalized shadow
        if np.dot(-grad,d_FW) >= np.dot(-grad,d_pi/(np.dot(d_pi,d_pi)**0.5)):
            
            t4 = process_time()
            
            #choose FW direction
            gamma_max = 1
            x, gamma = segment_search(func, grad_f, x, x + gamma_max*d_FW)
            FW.append(1)
            
            counts_in_face.append(0)
            counts_shadow.append(1)
            t5,t_pw = process_time(),[0]*5
            
        else: 
            #trace the projections curve
            x, c1,c2, t_pw,d_pi = trace_PW_curve(x,grad,shadow,in_face_shadow,in_face_trace,feasibility_oracle,func,grad_f)
            #x, c1,c2, t_pw,d_pi = trace_PW_curve_app(x,grad,shadow,feasibility_oracle,func)
            FW.append(0)
            
            counts_in_face.append(c1)
            counts_shadow.append(c2)
            t4,t5 = 0,0
            
        if np.dot(d_pi,d_pi)**0.5 < shadow_tol:
            break
        else:
            time1.append(time1[t] + (t3 - t1) + (t5-t4) + t_pw[0])
            time2.append(time2[t] + (t2 - t1) + (t5-t4) + t_pw[-2])
            f_improv = function_value[-1] - func(x)
            function_value.append(func(x))
            t+=1
            
        #print(func(x))
        
    return x, function_value, time1,time2,t,primal_gap,counts_shadow,counts_in_face,FW


# # PGD

# In[26]:


# PDG  algorithm:
def PGD(x, lmo,proj,L, epsilon,func,grad_f, f_tol, time_tol):
    
    #record primal gap, function value, and time every iteration
    now=datetime.datetime.now()
    primal_gap = []
    function_value=[func(x)]
    time = [0]
    f_improv = np.inf
    counts = []

    #initialize starting point and active set
    t = 0    

    while f_improv > f_tol and time[-1] < time_tol:
        #compute gradient
        grad = grad_f(x)

        #solve linear subproblem and compute FW direction
        s = lmo(grad)
        d_FW = s-x

        #If primal gap is small enough - terminate
        if np.dot(-grad,d_FW) <= epsilon:
            break
        else:
            #update convergence data
            primal_gap.append(np.dot(-grad,d_FW))

        # take a 1/L step-size in negative gradient direction and then project back
        x = proj(x,grad,L)
        
        later=datetime.datetime.now()
        time.append((later - now).total_seconds())
        f_improv = function_value[-1] - func(x)
        function_value.append(func(x))
        t+=1
        
    return x, function_value, time,t,primal_gap


# # Plotting

# In[27]:


styles = {'AFW':'-', 'PFW':'--', 'DICG':'-', 'Shadow':'-',
          'PGD':'--', 'Shadow CG':'-'}
colors = {'AFW':'tab:blue', 'PFW':'tab:blue', 'DICG':'tab:red', 'Shadow':'tab:green',
          'PGD':'tab:orange', 'Shadow CG':'tab:orange'}
labels = {'AFW':'AFW', 'PFW':'PFW', 'DICG':'DICG', 'PGD':'PGD',
          'Shadow':'Shadow Walk', 'Shadow CG':'Shadow CG'}


# In[98]:


def plotter(res, styles, colors, labels,oracle = False,fstar = None, log = True, sci = False, outfilename=None):
    
    '''
    plot primal gap specifying whether we want to assume oracle access for shadow or not
    '''
    
    fig = plt.figure()
    fig.set_figheight(3)
    fig.set_figwidth(10)

    plt.subplot(1, 2, 1)
    
    if fstar:
        for alg in res.keys():
            plt.plot(np.arange(len(res[alg][1])), res[alg][1] - fstar, linestyle=styles[alg], color=colors[alg],
                     label=labels[alg])
    else:
        for alg in res.keys():
            plt.plot(np.arange(len(res[alg][1])), res[alg][1], linestyle=styles[alg], color=colors[alg],
                     label=labels[alg])
    plt.legend()
    plt.ylabel(r'$\log (f(\mathbf{x}_t) - f(\mathbf{x}^*))$')
    plt.xlabel('Iteration')
    plt.tick_params(which='minor', left=False)
    if log:
        plt.yscale('log')
    if sci:
        plt.ticklabel_format(axis = 'y',style='sci',scilimits = (0,0))
    plt.grid(linestyle=':')
    #plt.xlim(0,30)
    #plt.ylim(10e-7,10e-1)

    plt.subplot(1, 2, 2)
    if fstar:
        for alg in res.keys():
            if oracle and alg in ['Shadow', 'Shadow CG']:
                plt.plot(res[alg][3], res[alg][1] - fstar, linestyle=styles[alg], color=colors[alg])
                plt.xlabel('Wall-clock time assuming shadow and in-face oracle (s)')
            else:
                plt.xlabel('Wall-clock time (s)')
                plt.plot(res[alg][2], res[alg][1] - fstar, linestyle=styles[alg], color=colors[alg])
    else:
        for alg in res.keys():
            
            if oracle and alg in ['Shadow', 'Shadow CG']:
                plt.plot(res[alg][3], res[alg][1], linestyle=styles[alg], color=colors[alg])
                plt.xlabel('Wall-clock time assuming shadow and in-face oracle (s)')
            else:
                plt.plot(res[alg][2], res[alg][1], linestyle=styles[alg], color=colors[alg])
                plt.xlabel('Wall-clock time (s)')
                
    plt.ylabel(r'$\log (f(\mathbf{x}_t) - f(\mathbf{x}^*))$')
    plt.tick_params(which='both', left=False, labelleft=False)
    plt.grid(linestyle=':')
    if log:
        plt.yscale('log')
    if sci:
        plt.ticklabel_format(axis = 'y',style='sci',scilimits = (0,0))

    if outfilename is not None:
        plt.savefig(outfilename, dpi=200, bbox_inches='tight')
        
    plt.show()


def plotter_gaps(res, styles, colors, labels,oracle = False, log = True, sci = False, outfilename=None):
    
    
    '''
    plot duality gap specifying whether we want to assume oracle access for shadow or not
    '''
    
    fig = plt.figure()
    fig.set_figheight(3)
    fig.set_figwidth(10)

    plt.subplot(1, 2, 1)
    for alg in res.keys():
        if alg in ['Shadow', 'Shadow CG']:
            plt.plot(np.arange(len(res[alg][5])), res[alg][5], linestyle=styles[alg], color=colors[alg],
                     label=labels[alg])
        else:
            plt.plot(np.arange(len(res[alg][4])), res[alg][4], linestyle=styles[alg], color=colors[alg],
                     label=labels[alg])
    plt.legend()
    plt.ylabel(r'$\log(\mathrm{Duality}$' + ' gap)')
    plt.xlabel('Iteration')
    plt.tick_params(which='minor', left=False)
    if log:
        plt.yscale('log')
    if sci:
        plt.ticklabel_format(axis = 'y',style='sci',scilimits = (0,0))
    plt.grid(linestyle=':')

    plt.subplot(1, 2, 2)
    for alg in res.keys():
        if oracle and alg in ['Shadow', 'Shadow CG']:
            plt.plot(res[alg][3][1:], res[alg][5], linestyle=styles[alg], color=colors[alg])
            plt.xlabel('Wall-clock time assuming shadow and in-face oracle (s)')
        elif alg in ['Shadow', 'Shadow CG']:
            plt.plot(res[alg][2][1:], res[alg][5], linestyle=styles[alg], color=colors[alg])
            plt.xlabel('Wall-clock time (s)')
        else:
            plt.plot(res[alg][2][1:], res[alg][4], linestyle=styles[alg], color=colors[alg])
            plt.xlabel('Wall-clock time (s)')
    if log:
        plt.yscale('log')
    if sci:
        plt.ticklabel_format(axis = 'y',style='sci',scilimits = (0,0))
    plt.ylabel(r'$\log(\mathrm{Duality}$' + ' gap)')
    plt.tick_params(which='both', left=False, labelleft=False)
    plt.grid(linestyle=':')

    if outfilename is not None:
        plt.savefig(outfilename, dpi=200, bbox_inches='tight')
        
    plt.show()

def plotter_break_point(res, styles, colors, labels, m = False,outfilename=None):
    
    '''
    plot number of oracle calls made by shadow
    '''
    
    fig = plt.figure()
    fig.set_figheight(3)
    fig.set_figwidth(4.5)
    
    li = ['Shadow', 'Shadow CG']
    
    for alg in li:
        plt.plot(np.arange(len(res[alg][6])), np.array(res[alg][6]), linestyle=styles[alg], color=colors[alg],
                 label=labels[alg])
    if m:
        plt.axhline(m,  color='black', linestyle='--', label = r'$m$')
        

    plt.legend()
    #plt.ylabel('Number of Iterations spent in '+r'$\mathrm{Trace}(\mathbf{x}_t, \nabla f(\mathbf{x}_t))$')
    plt.ylabel('Number of calls to shadow oracle')
    plt.xlabel('Iteration')
    plt.tick_params(which='minor', left=False)
    plt.grid(linestyle=':')

    if outfilename is not None:
        plt.savefig(outfilename, dpi=200, bbox_inches='tight')
        
    plt.show()

    
def plotter_cum(res, styles, colors, labels, m = False,outfilename=None):
    
    '''
    plot cumulative number of shadow steps taken
    '''
    
    fig = plt.figure()
    fig.set_figheight(3)
    fig.set_figwidth(4.5)
    
    li = ['Shadow', 'Shadow CG']
    
    for alg in li:
        if alg == 'Shadow':   
            data = res[alg][6].copy()
            plt.plot(np.arange(len(data)), np.cumsum(np.array(data)+1), linestyle=styles[alg], color=colors[alg],
                 label=labels[alg])
        else:
            data = []
            for i,j in enumerate(res[alg][6]):
                if res[alg][-1][i] == 0:
                    data.append(j)
                else:
                    data.append(0)
            plt.plot(np.arange(len(data)), np.cumsum(np.array(data)), linestyle=styles[alg], color=colors[alg],
                 label=labels[alg])

    plt.legend()
    plt.ylabel('Cumulative number of shadow steps')
    plt.xlabel('Iteration')
    plt.tick_params(which='minor', left=False)
    plt.grid(linestyle=':')

    if outfilename is not None:
        plt.savefig(outfilename, dpi=200, bbox_inches='tight')
        
    plt.show()
    
    
def plotter_break_point_in_face(res, styles, colors, labels, m = False,outfilename=None):
    
    '''
    plot number of oracle calls made by in-face shadow
    '''
    
    fig = plt.figure()
    fig.set_figheight(3)
    fig.set_figwidth(4.5)
    
    li = ['Shadow', 'Shadow CG']
    
    for alg in li:
        plt.plot(np.arange(len(res[alg][7])), np.array(res[alg][7]), linestyle=styles[alg], color=colors[alg],
                 label=labels[alg])
    if m:
        plt.axhline(m,  color='black', linestyle='--', label = r'$m$')
        

    plt.legend()
    #plt.ylabel('Number of Iterations spent in '+r'$\mathrm{Trace}(\mathbf{x}_t, \nabla f(\mathbf{x}_t))$')
    plt.ylabel('Number of calls to in-face shadow oracle')
    plt.xlabel('Iteration')
    plt.tick_params(which='minor', left=False)
    plt.grid(linestyle=':')

    if outfilename is not None:
        plt.savefig(outfilename, dpi=200, bbox_inches='tight')
        
    plt.show()

    
def plotter_cum_in_face(res, styles, colors, labels, m = False,outfilename=None):
    
    '''
    plot cumulative number of in-face shadow steps taken
    '''
    
    fig = plt.figure()
    fig.set_figheight(3)
    fig.set_figwidth(4.5)
    
    li = ['Shadow', 'Shadow CG']
    
    for alg in li:
        data = res[alg][7].copy()
        plt.plot(np.arange(len(data)), np.cumsum(np.array(data)), linestyle=styles[alg], color=colors[alg],
             label=labels[alg])

    plt.legend()
    plt.ylabel('Cumulative number of in-face shadow steps')
    plt.xlabel('Iteration')
    plt.tick_params(which='minor', left=False)
    plt.grid(linestyle=':')

    if outfilename is not None:
        plt.savefig(outfilename, dpi=200, bbox_inches='tight')
        
    plt.show()

    
    
def plotter_wout_PGD(res, styles, colors, labels,oracle = False,fstar = None, log = True, sci = False, outfilename=None):
    
    
    '''
    plot primal gap without PGD -  specifying whether we want to assume oracle access for shadow or not
    '''
    
    
    fig = plt.figure()
    fig.set_figheight(3)
    fig.set_figwidth(10)

    plt.subplot(1, 2, 1)
    
    if fstar:
        for alg in res.keys():
            if alg != 'PGD':
                plt.plot(np.arange(len(res[alg][1])), res[alg][1] - fstar, linestyle=styles[alg], color=colors[alg],
                         label=labels[alg])
    else:
        for alg in res.keys():
            if alg != 'PGD':
                plt.plot(np.arange(len(res[alg][1])), res[alg][1], linestyle=styles[alg], color=colors[alg],
                     label=labels[alg])
    plt.legend()
    plt.ylabel(r'$\log (f(\mathbf{x}_t) - f(\mathbf{x}^*))$')
    plt.xlabel('Iteration')
    plt.tick_params(which='minor', left=False)
    if log:
        plt.yscale('log')
    if sci:
        plt.ticklabel_format(axis = 'y',style='sci',scilimits = (0,0))
    plt.grid(linestyle=':')
    #plt.xlim(0,30)
    #plt.ylim(10e-7,10e-1)


    plt.subplot(1, 2, 2)
    if fstar:
        for alg in res.keys():
            if alg != 'PGD':
                if oracle and alg in ['Shadow', 'Shadow CG']:
                    plt.plot(res[alg][3], res[alg][1] - fstar, linestyle=styles[alg], color=colors[alg])
                    plt.xlabel('Wall-clock time assuming shadow and in-face oracle (s)')
                else:
                    plt.xlabel('Wall-clock time (s)')
                    plt.plot(res[alg][2], res[alg][1] - fstar, linestyle=styles[alg], color=colors[alg])
    else:
        for alg in res.keys():
            if alg != 'PGD':
                if oracle and alg in ['Shadow', 'Shadow CG']:
                    plt.plot(res[alg][3], res[alg][1], linestyle=styles[alg], color=colors[alg])
                    plt.xlabel('Wall-clock time assuming shadow and in-face shadow oracle (s)')
                else:
                    plt.xlabel('Wall-clock time (s)')
                    plt.plot(res[alg][2], res[alg][1], linestyle=styles[alg], color=colors[alg])

    plt.ylabel(r'$\log (f(\mathbf{x}_t) - f(\mathbf{x}^*))$')
    plt.tick_params(which='both', left=False, labelleft=False)
    plt.grid(linestyle=':')
    if log:
        plt.yscale('log')
    if sci:
        plt.ticklabel_format(axis = 'y',style='sci',scilimits = (0,0))

    if outfilename is not None:
        plt.savefig(outfilename, dpi=200, bbox_inches='tight')
        
    plt.show()



def plotter_gaps_wout_PGD(res, styles, colors, labels,oracle = False, log = True, sci = False, outfilename=None):
    
        
    '''
    plot duality gap without PGD -  specifying whether we want to assume oracle access for shadow or not
    '''
    
    
    fig = plt.figure()
    fig.set_figheight(3)
    fig.set_figwidth(10)

    plt.subplot(1, 2, 1)
    for alg in res.keys():
        if alg != 'PGD':
            if alg in ['Shadow', 'Shadow CG']:
                plt.plot(np.arange(len(res[alg][5])), res[alg][5], linestyle=styles[alg], color=colors[alg],
                         label=labels[alg])
            else:
                plt.plot(np.arange(len(res[alg][4])), res[alg][4], linestyle=styles[alg], color=colors[alg],
                         label=labels[alg])
    plt.legend()
    plt.ylabel(r'$\log(\mathrm{Duality}$' + ' gap)')
    plt.xlabel('Iteration')
    plt.tick_params(which='minor', left=False)
    if log:
        plt.yscale('log')
    if sci:
        plt.ticklabel_format(axis = 'y',style='sci',scilimits = (0,0))
    plt.grid(linestyle=':')

    plt.subplot(1, 2, 2)
    for alg in res.keys():
        if alg != 'PGD':
            if oracle and alg in ['Shadow', 'Shadow CG']:
                plt.plot(res[alg][3][1:], res[alg][5], linestyle=styles[alg], color=colors[alg])
                plt.xlabel('Wall-clock time assuming shadow and in-face oracle (s)')
            elif alg in ['Shadow', 'Shadow CG']:
                plt.xlabel('Wall-clock time (s)')
                plt.plot(res[alg][2][1:], res[alg][5], linestyle=styles[alg], color=colors[alg])
            else:
                plt.xlabel('Wall-clock time (s)')
                plt.plot(res[alg][2][1:], res[alg][4], linestyle=styles[alg], color=colors[alg])
    if log:
        plt.yscale('log')
    if sci:
        plt.ticklabel_format(axis = 'y',style='sci',scilimits = (0,0))
    plt.ylabel(r'$\log(\mathrm{Duality}$' + ' gap)')
    plt.tick_params(which='both', left=False, labelleft=False)
    plt.grid(linestyle=':')

    if outfilename is not None:
        plt.savefig(outfilename, dpi=200, bbox_inches='tight')
        
    plt.show()


# # Shadow and Projection oracles using FW

# In[12]:


#projection oracle using FW
def proj_FW(x_0,y,lmo,f_tol,time_tol):
    
    '''
    use the FW algorithm to project y onto polytope P using x_0 as a starting point
    '''

    def f_proj(z):
        return np.dot(z - y,z - y)

    def grad_f_proj(z):
        return 2*(z - y)


    return FW(x_0, lmo, 0,f_proj,grad_f_proj, f_tol, time_tol)[0]


def shadow_oracle_FW_binary(x_0,grad_0,lmo,start,tol,f_tol,time_tol):
    
    
    '''
    use the FW algorithm to compute shadow of gradinet by doing binary search to find the breakpoint
    '''

    def f_proj(z):
        return np.dot(z - x_0 + epsilon*grad_0,z - x_0 + epsilon*grad_0)

    def grad_f_proj(z):
        return 2*(z - x_0 + epsilon*grad_0)

    epsilon = start

    while True:

        z1 = FW(x_0, lmo, 0,f_proj,grad_f_proj, f_tol, time_tol)[0]
        d_pi = (z1 - x_0)/epsilon

        epsilon = epsilon / 2
        z2 = FW(x_0, lmo, 0,f_proj,grad_f_proj, f_tol, time_tol)[0]

        if np.dot( np.round(x_0 + epsilon*d_pi -z2,5), np.round(x_0 + epsilon*d_pi - z2,5))**0.5  <= tol:
            break
        else:
            epsilon = epsilon/2

    return d_pi


# # Computational Experiments

# # Figures 8, 9 and 10: Lasso Regression (large instance)

# In[108]:


#specify problem dimension and noise parameter
m, n, sigma = 50, 100, 0.1
np.random.seed(11)

#generate problem data
x_star = np.zeros(n)
rand = np.random.randint(1,n,int(0.25*n))
for i in rand[:int(len(rand)*0.5)]:
    x_star[i] = 1
    
for i in rand[:int(len(rand)*0.5)]:
    x_star[i] = -1

A = np.random.normal(0.5,0.25,(m, n))
y = np.dot(A, x_star)+sigma*np.random.randn(m)


#define objective function and gradient
f = lambda z: np.linalg.norm(y-np.dot(A, z[:n]-z[n:]), 2)**2

def grad_f(x):
    return np.concatenate([2*np.matmul(A.T,np.dot(A, x[:n]-x[n:]) - y),-2*np.matmul(A.T,np.dot(A, x[:n]-x[n:]) - y)])


#define linear optimization oracle
tau = np.linalg.norm(x_star, 1)
V = tau*np.identity(2*n)
x = V[np.random.randint(len(V))].copy()
lmo = lambda g: V[np.argmin(g)]
    
#generate constraint matrix:
B = np.vstack((np.ones(2*n), -np.ones(2*n), -np.eye(2*n)))
e = [tau] + [-tau] + [0]*(2*n)
#polytope is now given by Bx <= e

#choose desired feasibility oracle 
#since our polytope is a scaled simplex, the feasibility oracle could simply be the following
def feasibility_oracle (x,d):
    eta = np.min([-x[i]/d[i] if d[i] < 0 else np.inf for i in range(len(x))])
    eta = min(eta, 1)
    return eta

#choose desired shadow oracle - here for example we can gurobi to compute shadow
def shadow(x,d):
    return shadow_oracle(B,e,x, d)

#choose in-face shadow oracle
def in_face_shadow(x,d):
    return in_face_shadow_oracle(B,e,x,d)
    
#choose in-face trace oracle to compute maximum in-face movement
def in_face_trace(x_0,x,d, gamma_total,d_hat):
    return trace_in_face(B,e,x_0,x,d,d_hat, gamma_total, lmo)

f_tol,time_tol,epsilon,tol = 0.00001, np.inf, 0.1, 0.001


# In[103]:


res_lasso_sparse_large = {}
print('AFW')
res_lasso_sparse_large['AFW'] = AFW(x, lmo, epsilon,f,grad_f, f_tol, time_tol)
print('PFW')
res_lasso_sparse_large['PFW'] = PFW(x, lmo, epsilon,f,grad_f, f_tol, time_tol=np.inf)
print('DICG')
res_lasso_sparse_large['DICG'] = DICG(x, lmo, feasibility_oracle, epsilon,f,grad_f, f_tol, time_tol)
print('Shadow')
res_lasso_sparse_large['Shadow'] = SD(x, lmo, shadow,in_face_shadow,in_face_trace,feasibility_oracle, epsilon,f,grad_f, f_tol, time_tol, tol)
print('Shadow CG')
res_lasso_sparse_large['Shadow CG'] = SCG(x, lmo, shadow,in_face_shadow,in_face_trace, feasibility_oracle, epsilon,f,grad_f, f_tol, time_tol, tol)


# In[109]:


f_star = f(res_lasso_sparse_large['Shadow'][0])

#plot primal gaps and duality gaps without assuming oracle access to shadow
plotter_wout_PGD(res_lasso_sparse_large, styles, colors, labels, False, f_star, log = True, outfilename=None)
plotter_gaps_wout_PGD(res_lasso_sparse_large, styles, colors, labels, False,log = True, outfilename=None)

#plot primal gaps and duality gaps assuming oracle access to shadow and leaving PGD out for a better comparison
plotter_wout_PGD(res_lasso_sparse_large, styles, colors, labels, True, f_star, log = True,outfilename=None)
plotter_gaps_wout_PGD(res_lasso_sparse_large, styles, colors, labels, True,log = True, outfilename=None)

#plot number of shadow oracle calls made and access cumulative number of shadow steps taken by Shadow-CG and Shadow-Walk
plotter_break_point(res_lasso_sparse_large, styles, colors, labels, outfilename=None)
plotter_cum(res_lasso_sparse_large, styles, colors, labels,  outfilename=None)

#plot number of in-face shadow oracle calls made and access cumulative number of in-face shadow steps 
#taken by Shadow-CG and Shadow-Walk _in_face
plotter_break_point_in_face(res_lasso_sparse_large, styles, colors, labels, outfilename=None)
plotter_cum_in_face(res_lasso_sparse_large, styles, colors, labels,  outfilename=None)