def test_cp_pde_with_wrong_boundary_constraint_length():
    diff_eq = DiffusionEquation(2)
    mesh = Mesh([(0., 5.), (-5., 5.)], [.1, .2])
    static_bcs = [(DirichletBoundaryCondition(lambda x, t: np.zeros((13, 1)),
                                              is_static=True), ) * 2] * 2
    with pytest.raises(ValueError):
        ConstrainedProblem(diff_eq, mesh, static_bcs)

    dynamic_bcs = [
        (DirichletBoundaryCondition(lambda x, t: np.zeros((13, 1))), ) * 2
    ] * 2
    cp = ConstrainedProblem(diff_eq, mesh, dynamic_bcs)
    with pytest.raises(ValueError):
        cp.create_boundary_constraints(True, 0.)
def test_cp_3d_pde():
    mesh = Mesh([(2., 6.), (-3., 3.), (10., 12.)], [.1, .2, .5])

    assert mesh.shape(True) == (41, 31, 5)
    assert mesh.shape(False) == (40, 30, 4)

    diff_eq = WaveEquation(3)
    cp = ConstrainedProblem(
        diff_eq, mesh,
        ((DirichletBoundaryCondition(
            vectorize_bc_function(lambda x, t: (999., None)), is_static=True),
          NeumannBoundaryCondition(
              vectorize_bc_function(lambda x, t: (None, None)),
              is_static=True)),
         (DirichletBoundaryCondition(
             vectorize_bc_function(lambda x, t: (0., 0.)), is_static=True),
          NeumannBoundaryCondition(lambda x, t: np.full((len(x), 2), t))),
         (NeumannBoundaryCondition(lambda x, t: -x[:, :2] * x[:, 1:3],
                                   is_static=True),
          DirichletBoundaryCondition(
              vectorize_bc_function(lambda x, t: (-999., None))))))

    assert cp.y_shape(True) == (41, 31, 5, 2)
    assert cp.y_shape(False) == (40, 30, 4, 2)

    assert not cp.are_all_boundary_conditions_static
    assert cp.are_there_boundary_conditions_on_y

    assert cp.static_y_vertex_constraints.shape == (2, )

    y = np.full(cp._y_vertices_shape, -1)
    cp.static_y_vertex_constraints[0].apply(y[..., :1])
    cp.static_y_vertex_constraints[1].apply(y[..., 1:])

    assert np.all(y[0, 1:, :, 0] == 999.)
    assert np.all(y[:, 0, :, 0] == 0.)
    assert np.all(y[1:, 1:, :, 0] == -1.)
    assert np.all(y[:, 0, :, 1] == 0.)
    assert np.all(y[:, 1:, :, 1] == -1.)

    vertex_boundary_constraints = cp.static_boundary_vertex_constraints
    cell_boundary_constraints = cp.static_boundary_cell_constraints

    for y_boundary_constraints in \
            [vertex_boundary_constraints[0], cell_boundary_constraints[0]]:
        assert y_boundary_constraints.shape == (3, 2)
        assert y_boundary_constraints[0, 0][0] is not None
        assert y_boundary_constraints[0, 1][0] is not None
        assert y_boundary_constraints[0, 0][1] is None
        assert y_boundary_constraints[0, 1][1] is None
        assert y_boundary_constraints[1, 0][0] is not None
        assert y_boundary_constraints[1, 1][0] is not None
        assert y_boundary_constraints[1, 0][1] is None
        assert y_boundary_constraints[1, 1][1] is None
        assert y_boundary_constraints[2, 0][0] is None
        assert y_boundary_constraints[2, 1][0] is None
        assert y_boundary_constraints[2, 0][1] is None
        assert y_boundary_constraints[2, 1][1] is None

    for d_y_boundary_constraints in \
            [vertex_boundary_constraints[1], cell_boundary_constraints[1]]:
        assert d_y_boundary_constraints.shape == (3, 2)
        assert d_y_boundary_constraints[0, 0][0] is None
        assert d_y_boundary_constraints[0, 1][0] is None
        assert d_y_boundary_constraints[0, 0][1] is not None
        assert d_y_boundary_constraints[0, 1][1] is not None
        assert d_y_boundary_constraints[1, 0][0] is None
        assert d_y_boundary_constraints[1, 1][0] is None
        assert d_y_boundary_constraints[1, 0][1] is None
        assert d_y_boundary_constraints[1, 1][1] is None
        assert d_y_boundary_constraints[2, 0][0] is not None
        assert d_y_boundary_constraints[2, 1][0] is not None
        assert d_y_boundary_constraints[2, 0][1] is None
        assert d_y_boundary_constraints[2, 1][1] is None

    new_vertex_boundary_constraints = cp.create_boundary_constraints(True, 1.)
    new_y_boundary_constraints = new_vertex_boundary_constraints[0]
    new_d_y_boundary_constraints = new_vertex_boundary_constraints[1]
    assert new_y_boundary_constraints[2, 0][1] is not None
    assert new_y_boundary_constraints[2, 1][1] is not None
    assert new_d_y_boundary_constraints[1, 0][1] is not None
    assert new_d_y_boundary_constraints[1, 1][1] is not None

    d_y_boundary = np.full((41, 1, 5, 2), np.nan)
    new_d_y_boundary_constraints[1, 0][1].apply(d_y_boundary[..., :1])
    new_d_y_boundary_constraints[1, 1][1].apply(d_y_boundary[..., 1:])
    assert np.all(d_y_boundary == 1.)

    new_y_vertex_constraints = \
        cp.create_y_vertex_constraints(new_y_boundary_constraints)
    assert new_y_vertex_constraints.shape == (2, )

    y = np.full(cp._y_vertices_shape, -1)
    new_y_vertex_constraints[0].apply(y[..., :1])
    new_y_vertex_constraints[1].apply(y[..., 1:])

    assert np.all(y[0, 1:, :-1, 0] == 999.)
    assert np.all(y[:, 0, :-1, 0] == 0.)
    assert np.all(y[:, :, -1, 0] == -999.)
    assert np.all(y[1:, 1:, :-1, 0] == -1.)
    assert np.all(y[:, 0, :, 1] == 0.)
    assert np.all(y[:, 1:, :, 1] == -1.)