<commit_before>from hazan import *
from hazan_penalties import *
import scipy
import numpy as np

# Do a simple test of General SDP Solver with binary search

def test1():
    """
    A simple semidefinite program

    max x_11 + x_22
    subject to
        x_11 + 2 x_22 == 1
        x_12 == x_21 == 0
        Tr(X) = x_11 + x_22 <= 1
        X semidefinite

    The solution to this problem is

        X = [[1, 0],
             [0, 0]]

    from Lagrange multiplier.
    """
    dim = 2
    N_iter = 400
    eps = 1e-2
    g = GeneralSDPHazanSolver()
    def h(X):
        return X[0,0] + X[1,1]
    def gradh(X):
        return np.eye(2)
    As = []
    bs = []
    Cs = [np.array([[1., 0.],
                    [0., 2.]]),
          np.array([[0., 1.],
                    [0., 0.]]),
          np.array([[0., 0.],
                    [1., 0.]])]
    ds = [1., 0., 0.]
    Fs = []
    gradFs = []
    Gs = []
    gradGs = []
    U = 2
    L = 0
    R = 1.
    upper, lower, X_upper, X_lower, SUCCEED = g.solve(h, gradh, As, bs,
                Cs, ds, Fs, gradFs, Gs, gradGs, eps, dim, R, U, L, N_iter)
    print
    print "General SDP Solver Finished"
    print "SUCCEED: ", SUCCEED
    print "upper: ", upper
    print "X_upper:\n", X_upper
    print "X_lower:\n", X_lower
    print "lower: ", lower

def testSchurComplement():
    """
    Specifies a Schur complement SDP of form

    minimize -log det Q^{-1}
          ------------
         |D-Q   A     |
    X =  |A.T D^{-1}  |
         |           Q|
          ------------
    X is PSD
    """
    dim = 1
    cdim = 3 * dim
    g = GeneralSDPHazanSolver()
    As, bs, Cs, ds, = [], [], [], []
    Fs, gradFs, Gs, gradGs = [], [], [], []

    # Generate initial data
    D = np.eye(dim)
    Dinv = np.linalg.inv(D)
    A = 0.5 * np.eye(dim)
    """ We neeed to enforce constant equalities in X.
      -------------
     |D-_   A    0 |
C =  |A.T D^{-1} 0 |
     | 0    0    _ |
      -------------
    """
    C = np.zeros((cdim, cdim))
    A_1_cds = (0, dim, dim, 2*dim)
    A_T_1_cds = (dim, 2*dim, 0, dim)
    Dinv_1_cds  = (dim, 2*dim, dim, 2*dim)
    constraints = [(A_1_cds, A), (A_T_1_cds, A.T),
            (Dinv_1_cds, Dinv)]
    def const_regions(X):
        return many_batch_equals(X, constraints)
    def grad_const_regions(X):
        return grad_many_batch_equals(X, constraints)
    Gs.append(const_regions)
    gradGs.append(grad_const_regions)

    """ We need to constraint linear inequalities
          -----------
         |D-Q        |
    C =  |           |
         |         Q |
          -----------
    """
    Q_cds = (2*dim, 3*dim, 2*dim, 3*dim)
    block_1_Q_cds = (0, dim, 0, dim)
    linear_constraints = [(-1., Q_cds, D, block_1_Q_cds)]

    def linear_regions(X):
        return many_batch_linear_equals(X, linear_constraints)
    def grad_linear_regions(X):
        return grad_many_batch_linear_equals(X, linear_constraints)
    Gs.append(linear_regions)
    gradGs.append(grad_linear_regions)

    # log det Q^{-1} = - log det Q
    def h(X):
        Q = get_entries(X, Q_cds)
        #D_Q = get_entries(X, block_1_Q_cds)
        #block_1_Q = -D_Q + D
        #return -np.log(np.linalg.det(Q)) - np.log(np.linalg.det(block_1_Q))
        return np.log(np.linalg.det(Q))
    # grad log det Q^{-1} = -Q^{-1} (see Boyd and Vandenberge, A4.1)
    def gradh(X):
        grad = np.zeros(np.shape(X))
        Q = get_entries(X, Q_cds)
        # Look into avoiding this computation if possible
        eigs_Q = np.linalg.eigh(Q)[0]
        max_eig_Q = np.amax(eigs_Q)
        max_eig_Qinv = 1./max_eig_Q
        Qinv = np.linalg.inv(Q)
        # Scale Qinv down to unit spectral norm
        Qinv = (1/max_eig_Qinv) * Qinv
        # Scale Qinv to have spectral norm the same as Q
        Qinv *= max_eig_Q
        gradQ = Qinv
        set_entries(grad, Q_cds, gradQ)
        set_entries(grad, block_1_Q_cds, gradQ)
        return grad

    D_upper = np.trace(D)
    D_inv_upper = np.trace(Dinv)
    R = (D_upper + D_inv_upper + D_upper)
    L = 0
    U = 25
    eps = 3e-2
    N_iter = 150
    X_init = np.zeros((cdim, cdim))
    Q_init = 0.2 * np.eye(dim)
    set_entries(X_init, Q_cds, Q_init)
    set_entries(X_init, block_1_Q_cds, D - Q_init)
    import pdb
    pdb.set_trace()
    upper, lower, X_upper, X_lower, SUCCEED = g.solve(h, gradh, As, bs,
                Cs, ds, Fs, gradFs, Gs, gradGs, eps, cdim, R, U, L,
                N_iter, X_init=X_init)
    print "X_lower\n", X_lower
    if X_lower != None:
        print "h(X_lower)\n", h(X_lower)
    print "X_upper\n", X_upper
    if X_upper != None:<commit_msg>Changed testSchurComplement to match corrected Q-problem<commit_after>from hazan import *
from hazan_penalties import *
import scipy
import numpy as np

# Do a simple test of General SDP Solver with binary search

def test1():
    """
    A simple semidefinite program

    max x_11 + x_22
    subject to
        x_11 + 2 x_22 == 1
        x_12 == x_21 == 0
        Tr(X) = x_11 + x_22 <= 1
        X semidefinite

    The solution to this problem is

        X = [[1, 0],
             [0, 0]]

    from Lagrange multiplier.
    """
    dim = 2
    N_iter = 400
    eps = 1e-2
    g = GeneralSDPHazanSolver()
    def h(X):
        return X[0,0] + X[1,1]
    def gradh(X):
        return np.eye(2)
    As = []
    bs = []
    Cs = [np.array([[1., 0.],
                    [0., 2.]]),
          np.array([[0., 1.],
                    [0., 0.]]),
          np.array([[0., 0.],
                    [1., 0.]])]
    ds = [1., 0., 0.]
    Fs = []
    gradFs = []
    Gs = []
    gradGs = []
    U = 2
    L = 0
    R = 1.
    upper, lower, X_upper, X_lower, SUCCEED = g.solve(h, gradh, As, bs,
                Cs, ds, Fs, gradFs, Gs, gradGs, eps, dim, R, U, L, N_iter)
    print
    print "General SDP Solver Finished"
    print "SUCCEED: ", SUCCEED
    print "upper: ", upper
    print "X_upper:\n", X_upper
    print "X_lower:\n", X_lower
    print "lower: ", lower

def testSchurComplement():
    """
    Specifies a Schur complement convex program of form

    minimize -log det R
          --------------
         |D-ADA.T  I    |
    X =  |   I     R    |
         |            R |
          --------------
    X is PSD
    """
    dim = 1
    cdim = 3 * dim
    g = GeneralSDPHazanSolver()
    As, bs, Cs, ds, = [], [], [], []
    Fs, gradFs, Gs, gradGs = [], [], [], []

    # Generate initial data
    D = np.eye(dim)
    Dinv = np.linalg.inv(D)
    A = 0.5 * np.eye(dim)
    """ We neeed to enforce constant equalities in X.
      -------------
     |D-ADA.T   I    0 |
C =  | I        _    0 |
     | 0        0    _ |
      -------------
    """
    C = np.zeros((cdim, cdim))
    D_ADA_T_cds = (0, dim, 0, dim)
    D_ADA_T = D - np.dot(A, np.dot(D, A.T))
    I_1_cds = (0, dim, dim, 2*dim)
    I_2_cds = (dim, 2*dim, 0, dim)
    constraints = [(D_ADA_T_cds, D_ADA_T), (I_1_cds, np.eye(dim)),
            (I_2_cds, np.eye(dim))]
    # Add zero constraints
    constraints += [((2*dim, 3*dim, 0, 2*dim), np.zeros((dim, 2*dim))),
            ((0, 2*dim, 2*dim, 3*dim), np.zeros((2*dim, dim)))]
    def const_regions(X):
        return many_batch_equals(X, constraints)
    def grad_const_regions(X):
        return grad_many_batch_equals(X, constraints)
    Gs.append(const_regions)
    gradGs.append(grad_const_regions)

    """ We need to constraint linear inequalities
          ----------
         |          |
    C =  |     R    |
         |        R |
          ----------
    """
    R_cds = (2*dim, 3*dim, 2*dim, 3*dim)
    block_1_R_cds = (dim, 2*dim, dim, 2*dim)
    linear_constraints = [(1., R_cds, np.zeros((dim,dim)), block_1_R_cds)]

    def linear_regions(X):
        return many_batch_linear_equals(X, linear_constraints)
    def grad_linear_regions(X):
        return grad_many_batch_linear_equals(X, linear_constraints)
    Gs.append(linear_regions)
    gradGs.append(grad_linear_regions)

    # - log det R
    def h(X):
        R = get_entries(X, R_cds)
        return -np.log(np.linalg.det(R))
    # grad - log det R = -R^{-1} (see Boyd and Vandenberge, A4.1)
    # Note -R^{-1} = -Q
    def gradh(X):
        grad = np.zeros(np.shape(X))
        R = get_entries(X, R_cds)
        # Look into avoiding this computation if possible
        eigs_R = np.linalg.eigh(R)[0]
        max_eig_R = np.amax(eigs_R)
        max_eig_Q = 1./max_eig_R
        Q = np.linalg.inv(R)
        # Scale Q down to unit spectral norm
        Q = (1/max_eig_Q) * Q
        # Scale Q to have spectral norm the same as Q
        #Qinv *= max_eig_Q
        gradR = Q
        set_entries(grad, R_cds, gradR)
        set_entries(grad, block_1_R_cds, gradR)
        return grad

    #D_upper = np.trace(D)
    #D_inv_upper = np.trace(Dinv)
    #R = (D_upper + D_inv_upper + D_inv_upper)
    R = 100 # Just going to set this high...
    L = 0
    U = 25
    eps = 3e-2
    N_iter = 150
    X_init = np.zeros((cdim, cdim))
    Q_init = 0.2 * np.eye(dim)
    R_init = np.linalg.inv(Q_init)
    set_entries(X_init, R_cds, R_init)
    set_entries(X_init, block_1_R_cds, R_init)
    import pdb
    pdb.set_trace()
    upper, lower, X_upper, X_lower, SUCCEED = g.solve(h, gradh, As, bs,
                Cs, ds, Fs, gradFs, Gs, gradGs, eps, cdim, R, U, L,
                N_iter, X_init=X_init)
    print "X_lower\n", X_lower
    if X_lower != None:
        print "h(X_lower)\n", h(X_lower)
    print "X_upper\n", X_upper
    if X_upper != None:
{'fraction_non_alphanumeric': 0.15017462165308498, 'fraction_numerical': 0.023494549687797652, 'mean_word_length': 1.7696365767878077, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6264011', 'n_tokens_mistral': 3909, 'n_tokens_neox': 3662, 'n_words': 907}