<commit_before>import _diffcp
import cvxpy as cp
import numpy as np
import pytest
from _diffcp import Cone, ConeType
from scipy import sparse

import diffcp.cone_program as cone_prog
import diffcp.cones as cone_lib
import diffcp.utils as utils

CPP_CONES_TO_SCS = {
    ConeType.ZERO: cone_lib.EQ_DIM,
    ConeType.POS: "l",
    ConeType.SOC: "q",
    ConeType.PSD: "s",
    ConeType.EXP: "ep",
    ConeType.EXP_DUAL: "ed"
}


def test_vec_psd_dim():
    assert cone_lib.vec_psd_dim(10) == (10) * (10 + 1) / 2


def test_psd_dim():
    n = 4096
    assert cone_lib.psd_dim(cone_lib.vec_psd_dim(n)) == n


def test_in_exp():
    assert _diffcp.in_exp(np.array([0, 0, 1]))
    assert _diffcp.in_exp(np.array([-1, 0, 0]))
    assert _diffcp.in_exp(np.array([1, 1, 5]))
    assert not _diffcp.in_exp(np.array([1, 0, 0]))
    assert not _diffcp.in_exp(np.array([-1, -1, 1]))
    assert not _diffcp.in_exp(np.array([-1, 0, -1]))


def test_in_exp_dual():
    assert _diffcp.in_exp_dual(np.array([0, 1, 1]))
    assert _diffcp.in_exp_dual(np.array([-1, 1, 5]))
    assert not _diffcp.in_exp_dual(np.array([0, -1, 1]))
    assert not _diffcp.in_exp_dual(np.array([0, 1, -1]))


def test_unvec_symm():
    np.random.seed(0)
    n = 5
    x = np.random.randn(n, n)
    x = x + x.T
    np.testing.assert_allclose(
        cone_lib.unvec_symm(cone_lib.vec_symm(x), n), x)


def test_vec_symm():
    np.random.seed(0)
    n = 5
    x = np.random.randn(cone_lib.vec_psd_dim(n))
    np.testing.assert_allclose(
        cone_lib.vec_symm(cone_lib.unvec_symm(x, n)), x)


def test_proj_zero():
    np.random.seed(0)
    n = 100
    for _ in range(10):
        x = np.random.randn(n)
        np.testing.assert_allclose(
            cone_lib._proj(x, cone_lib.ZERO, dual=True), x)
        np.testing.assert_allclose(
            cone_lib._proj(x, cone_lib.ZERO, dual=False),
            np.zeros(n))


def test_proj_pos():
    np.random.seed(0)
    n = 100
    for _ in range(15):
        x = np.random.randn(n)
        p = cone_lib._proj(x, cone_lib.POS, dual=False)
        np.testing.assert_allclose(p, np.maximum(x, 0))
        np.testing.assert_allclose(
            p, cone_lib._proj(x, cone_lib.POS, dual=True))


def test_proj_soc():
    np.random.seed(0)
    n = 100
    for _ in range(15):
        x = np.random.randn(n)
        z = cp.Variable(n)
        objective = cp.Minimize(cp.sum_squares(z - x))
        constraints = [cp.norm(z[1:], 2) <= z[0]]
        prob = cp.Problem(objective, constraints)
        prob.solve(solver="SCS", eps=1e-10)
        p = cone_lib._proj(x, cone_lib.SOC, dual=False)
        np.testing.assert_allclose(
            p, np.array(z.value))
        np.testing.assert_allclose(
            p, cone_lib._proj(x, cone_lib.SOC, dual=True))


def test_proj_psd():
    np.random.seed(0)
    n = 10
    for _ in range(15):
        x = np.random.randn(n, n)
        x = x + x.T
        x_vec = cone_lib.vec_symm(x)
        z = cp.Variable((n, n), PSD=True)
        objective = cp.Minimize(cp.sum_squares(z - x))
        prob = cp.Problem(objective)
        prob.solve(solver="SCS", eps=1e-10)
        p = cone_lib.unvec_symm(
            cone_lib._proj(x_vec, cone_lib.PSD, dual=False), n)
        np.testing.assert_allclose(p, z.value, atol=1e-5, rtol=1e-5)
        np.testing.assert_allclose(p, cone_lib.unvec_symm(
            cone_lib._proj(x_vec, cone_lib.PSD, dual=True), n))


def test_proj_exp():
    np.random.seed(0)
    for _ in range(15):
        x = np.random.randn(9)
        var = cp.Variable(9)
        constr = [cp.constraints.ExpCone(var[0], var[1], var[2])]
        constr.append(cp.constraints.ExpCone(var[3], var[4], var[5]))
        constr.append(cp.constraints.ExpCone(var[6], var[7], var[8]))
        obj = cp.Minimize(cp.norm(var[0:3] - x[0:3]) +
                          cp.norm(var[3:6] - x[3:6]) +
                          cp.norm(var[6:9] - x[6:9]))
        prob = cp.Problem(obj, constr)
        prob.solve(solver="SCS", eps=1e-12, max_iters=1000)<commit_msg>20k iters<commit_after>import _diffcp
import cvxpy as cp
import numpy as np
import pytest
from _diffcp import Cone, ConeType
from scipy import sparse

import diffcp.cone_program as cone_prog
import diffcp.cones as cone_lib
import diffcp.utils as utils

CPP_CONES_TO_SCS = {
    ConeType.ZERO: cone_lib.EQ_DIM,
    ConeType.POS: "l",
    ConeType.SOC: "q",
    ConeType.PSD: "s",
    ConeType.EXP: "ep",
    ConeType.EXP_DUAL: "ed"
}


def test_vec_psd_dim():
    assert cone_lib.vec_psd_dim(10) == (10) * (10 + 1) / 2


def test_psd_dim():
    n = 4096
    assert cone_lib.psd_dim(cone_lib.vec_psd_dim(n)) == n


def test_in_exp():
    assert _diffcp.in_exp(np.array([0, 0, 1]))
    assert _diffcp.in_exp(np.array([-1, 0, 0]))
    assert _diffcp.in_exp(np.array([1, 1, 5]))
    assert not _diffcp.in_exp(np.array([1, 0, 0]))
    assert not _diffcp.in_exp(np.array([-1, -1, 1]))
    assert not _diffcp.in_exp(np.array([-1, 0, -1]))


def test_in_exp_dual():
    assert _diffcp.in_exp_dual(np.array([0, 1, 1]))
    assert _diffcp.in_exp_dual(np.array([-1, 1, 5]))
    assert not _diffcp.in_exp_dual(np.array([0, -1, 1]))
    assert not _diffcp.in_exp_dual(np.array([0, 1, -1]))


def test_unvec_symm():
    np.random.seed(0)
    n = 5
    x = np.random.randn(n, n)
    x = x + x.T
    np.testing.assert_allclose(
        cone_lib.unvec_symm(cone_lib.vec_symm(x), n), x)


def test_vec_symm():
    np.random.seed(0)
    n = 5
    x = np.random.randn(cone_lib.vec_psd_dim(n))
    np.testing.assert_allclose(
        cone_lib.vec_symm(cone_lib.unvec_symm(x, n)), x)


def test_proj_zero():
    np.random.seed(0)
    n = 100
    for _ in range(10):
        x = np.random.randn(n)
        np.testing.assert_allclose(
            cone_lib._proj(x, cone_lib.ZERO, dual=True), x)
        np.testing.assert_allclose(
            cone_lib._proj(x, cone_lib.ZERO, dual=False),
            np.zeros(n))


def test_proj_pos():
    np.random.seed(0)
    n = 100
    for _ in range(15):
        x = np.random.randn(n)
        p = cone_lib._proj(x, cone_lib.POS, dual=False)
        np.testing.assert_allclose(p, np.maximum(x, 0))
        np.testing.assert_allclose(
            p, cone_lib._proj(x, cone_lib.POS, dual=True))


def test_proj_soc():
    np.random.seed(0)
    n = 100
    for _ in range(15):
        x = np.random.randn(n)
        z = cp.Variable(n)
        objective = cp.Minimize(cp.sum_squares(z - x))
        constraints = [cp.norm(z[1:], 2) <= z[0]]
        prob = cp.Problem(objective, constraints)
        prob.solve(solver="SCS", eps=1e-10)
        p = cone_lib._proj(x, cone_lib.SOC, dual=False)
        np.testing.assert_allclose(
            p, np.array(z.value))
        np.testing.assert_allclose(
            p, cone_lib._proj(x, cone_lib.SOC, dual=True))


def test_proj_psd():
    np.random.seed(0)
    n = 10
    for _ in range(15):
        x = np.random.randn(n, n)
        x = x + x.T
        x_vec = cone_lib.vec_symm(x)
        z = cp.Variable((n, n), PSD=True)
        objective = cp.Minimize(cp.sum_squares(z - x))
        prob = cp.Problem(objective)
        prob.solve(solver="SCS", eps=1e-10)
        p = cone_lib.unvec_symm(
            cone_lib._proj(x_vec, cone_lib.PSD, dual=False), n)
        np.testing.assert_allclose(p, z.value, atol=1e-5, rtol=1e-5)
        np.testing.assert_allclose(p, cone_lib.unvec_symm(
            cone_lib._proj(x_vec, cone_lib.PSD, dual=True), n))


def test_proj_exp():
    np.random.seed(0)
    for _ in range(15):
        x = np.random.randn(9)
        var = cp.Variable(9)
        constr = [cp.constraints.ExpCone(var[0], var[1], var[2])]
        constr.append(cp.constraints.ExpCone(var[3], var[4], var[5]))
        constr.append(cp.constraints.ExpCone(var[6], var[7], var[8]))
        obj = cp.Minimize(cp.norm(var[0:3] - x[0:3]) +
                          cp.norm(var[3:6] - x[3:6]) +
                          cp.norm(var[6:9] - x[6:9]))
        prob = cp.Problem(obj, constr)
        prob.solve(solver="SCS", eps=1e-12, max_iters=20_000)
{'fraction_non_alphanumeric': 0.15455680399500624, 'fraction_numerical': 0.02933832709113608, 'mean_word_length': 2.5478299379982285, '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': '6860620', 'n_tokens_mistral': 3463, 'n_tokens_neox': 3292, 'n_words': 580}