def test_admm_on_dumb():
    adj1, adj2 = dumb_admm_system()
    xymap1 = []
    xymap2 = []
    uymap1 = [0, 1]
    uymap2 = [0, 1]
    yzmap1 = [0, 1]
    yzmap2 = [1, 0]
    data1 = AdjointAdmmData(adj1, xymap1, uymap1, yzmap1, numpy.zeros(2), rho=.5)
    data2 = AdjointAdmmData(adj2, xymap2, uymap2, yzmap2, numpy.zeros(2), rho=.5)

    sp1 = AdjointAdmmSubProblem(data1, solver=AcceleratedGradient())
    sp2 = AdjointAdmmSubProblem(data2, solver=AcceleratedGradient())

    def objective(problem):
        return (problem.z[0] - 1) ** 2 + (problem.z[1] - 2) ** 2 + sum(
            [sp.data.admm_cost() for sp in problem.sub_problems])

    problem = Admm([sp2, sp1], n_iters=40, objective=objective)
    problem.solve()
    assert_array_almost_equal(problem.z, [1.0, 2.0])
def test_full_admm():
    l, r = left_system(), right_system()
    lxy = [1]
    luy = []
    rxy = []
    ruy = [0]
    lyz = [0]
    ryz = [0]
    ldata = AdjointAdmmData(l, lxy, luy, lyz, zeros(1))
    rdata = AdjointAdmmData(r, rxy, ruy, ryz, zeros(1))
    lsp = AdjointAdmmSubProblem(ldata, solver=PositiveGSReal())
    rsp = AdjointAdmmSubProblem(rdata, solver=PositiveGSReal())

    def objective(problem):
        control = problem.sub_problems[0].data.control
        return full_system().j(full_system().fs(control), control) + sum(
            [sp.data.admm_cost() for sp in problem.sub_problems])

    problem = Admm([lsp, rsp], objective=objective, n_iters=25)
    problem.solve()
    assert_almost_equal(problem.sub_problems[0].data.control[0], 410. / 21, delta=.05)
    #problem.plot()
def test_split_network():
    scen, control = test_ctm.small_scenario()
    set_cells_left = [1]
    set_ramps_left = [1]
    adjoint_left = ctm_admm.AdmmCTMAdjoint(scen, set_cells_left, set_ramps_left)

    set_cells_right = [0]
    set_ramps_right = [0]
    adjoint_right = ctm_admm.AdmmCTMAdjoint(scen, set_cells_right, set_ramps_right)
    T = scen.T
    xy_map_left = range(T) + range(2 * T, 3 * T)
    xy_map_right = range(T, 2 * T) + range(3 * T, 4 * T)
    uy_map_left = range(4 * T)
    uy_map_right = range(4 * T)
    yz_map_left = range(2 * T, 4 * T) + range(2 * T) + range(4 * T, 6 * T)
    yz_map_right = range(4 * T, 6 * T) + range(2 * T) + range(2 * T, 4 * T)
    sub_problems = [
                       AdjointAdmmSubProblem(AdjointAdmmData(adjoint_left, xy_map_left, uy_map_left, yz_map_left,
                                                             ctm_admm.AdmmCTMAdjoint.no_control(scen, set_cells_left,
                                                                                                set_ramps_left),
                                                             rho=20.),
                                             solver=AcceleratedGradient(15, 2000, 0.0, 1.0))] + [
                       AdjointAdmmSubProblem(AdjointAdmmData(adjoint_right, xy_map_right, uy_map_right, yz_map_right,
                                                             ctm_admm.AdmmCTMAdjoint.no_control(scen, set_cells_right,
                                                                                                set_ramps_right),
                                                             rho=20.), solver=AcceleratedGradient(15, 2000, 0.0, 1.0))
                   ]
    sln = Admm(sub_problems, n_iters=100)
    sln.solve()
    control = sln.z[:2 * T].reshape(T, 2).transpose()
    state = CTMSimulator().simulate(scen, control)
    state.plot()
    ttt = scen.ttt(state)
    print control
    print ttt
    sln.plot()
def test_single_admm_system():
    adj1, adj2 = dumb_admm_system()
    xymap = []
    uymap = [0,1]
    yzmap = [0,1]

    data = AdjointAdmmData(adj1, xymap, uymap, yzmap, numpy.zeros(2))
    sp = AdjointAdmmSubProblem(data, solver=AcceleratedGradient())

    def objective(problem):
        return (problem.z[0] - 1) ** 2 + problem.sub_problems[0].data.admm_cost()

    problem = Admm([sp], n_iters=20, objective=objective)
    problem.solve()
    assert_almost_equal(problem.z[0], 1.0)

    data = AdjointAdmmData(adj2, xymap, uymap, yzmap, numpy.zeros(2))
    sp = AdjointAdmmSubProblem(data, solver=AcceleratedGradient())

    def objective(problem):
        return (problem.z[0] - 2) ** 2 + problem.sub_problems[0].data.admm_cost()

    problem = Admm([sp], n_iters=20, objective=objective)
    problem.solve()
    assert_almost_equal(problem.z[0], 2.0)

    uymap = [0]
    yzmap = [0]

    data = AdjointAdmmData(adj1, xymap, uymap, yzmap, numpy.zeros(2))
    sp = AdjointAdmmSubProblem(data, solver=AcceleratedGradient())

    def objective(problem):
        return (problem.z[0] - 1) ** 2 + problem.sub_problems[0].data.admm_cost()

    problem = Admm([sp], n_iters=20, objective=objective)
    problem.solve()
    assert_almost_equal(problem.z[0], 1.0)
    assert_equal(len(problem.z), 1)

    data = AdjointAdmmData(adj2, xymap, uymap, yzmap, numpy.zeros(2))
    sp = AdjointAdmmSubProblem(data, solver=AcceleratedGradient())

    def objective(problem):
        return (problem.z[0] - 2) ** 2 + problem.sub_problems[0].data.admm_cost()

    problem = Admm([sp], n_iters=20, objective=objective)
    problem.solve()
    assert_almost_equal(problem.z[0], 2.0)
    assert_equal(len(problem.z), 1)