def test_gradient(self):
        mesh = get_unit_2p1D_mesh_for_testing(ypts=32, zpts=32, include_particles=False)
        spatial_methods = {
            "macroscale": pybamm.FiniteVolume(),
            "current collector": pybamm.ScikitFiniteElement(),
        }
        disc = pybamm.Discretisation(mesh, spatial_methods)

        # test gradient of 5*y + 6*z
        var = pybamm.Variable("var", domain="current collector")
        disc.set_variable_slices([var])

        y = mesh["current collector"].coordinates[0, :]
        z = mesh["current collector"].coordinates[1, :]

        gradient = pybamm.grad(var)
        grad_disc = disc.process_symbol(gradient)
        grad_disc_y, grad_disc_z = grad_disc.children

        np.testing.assert_array_almost_equal(
            grad_disc_y.evaluate(None, 5 * y + 6 * z),
            5 * np.ones_like(y)[:, np.newaxis],
        )
        np.testing.assert_array_almost_equal(
            grad_disc_z.evaluate(None, 5 * y + 6 * z),
            6 * np.ones_like(z)[:, np.newaxis],
        )

        # check grad_squared positive
        eqn = pybamm.grad_squared(var)
        eqn_disc = disc.process_symbol(eqn)
        ans = eqn_disc.evaluate(None, 3 * y ** 2)
        np.testing.assert_array_less(0, ans)
    def test_disc_spatial_var(self):
        mesh = get_unit_2p1D_mesh_for_testing(ypts=4,
                                              zpts=5,
                                              include_particles=False)
        spatial_methods = {
            "macroscale": pybamm.FiniteVolume(),
            "current collector": pybamm.ScikitFiniteElement(),
        }
        disc = pybamm.Discretisation(mesh, spatial_methods)

        # discretise y and z
        y = pybamm.SpatialVariable("y", ["current collector"])
        z = pybamm.SpatialVariable("z", ["current collector"])
        y_disc = disc.process_symbol(y)
        z_disc = disc.process_symbol(z)

        # create expected meshgrid
        y_vec = np.linspace(0, 1, 4)
        z_vec = np.linspace(0, 1, 5)
        Y, Z = np.meshgrid(y_vec, z_vec)
        y_actual = np.transpose(Y).flatten()[:, np.newaxis]
        z_actual = np.transpose(Z).flatten()[:, np.newaxis]

        # spatial vars should discretise to the flattend meshgrid
        np.testing.assert_array_equal(y_disc.evaluate(), y_actual)
        np.testing.assert_array_equal(z_disc.evaluate(), z_actual)
    def test_pure_neumann_poisson(self):
        # grad^2 u = 1, du/dz = 1 at z = 1, du/dn = 0 elsewhere, u has zero average
        u = pybamm.Variable("u", domain="current collector")
        c = pybamm.Variable("c")  # lagrange multiplier
        y = pybamm.SpatialVariable("y", ["current collector"])
        z = pybamm.SpatialVariable("z", ["current collector"])

        model = pybamm.BaseModel()
        # 0*c hack otherwise gives KeyError
        model.algebraic = {
            u:
            pybamm.laplacian(u) - pybamm.source(1, u) +
            c * pybamm.DefiniteIntegralVector(u, vector_type="column"),
            c:
            pybamm.Integral(u, [y, z]) + 0 * c,
        }
        model.initial_conditions = {u: pybamm.Scalar(0), c: pybamm.Scalar(0)}
        # set boundary conditions ("negative tab" = bottom of unit square,
        # "positive tab" = top of unit square, elsewhere normal derivative is zero)
        model.boundary_conditions = {
            u: {
                "negative tab": (0, "Neumann"),
                "positive tab": (1, "Neumann")
            }
        }
        model.variables = {"c": c, "u": u}
        # create discretisation
        mesh = get_unit_2p1D_mesh_for_testing(ypts=32,
                                              zpts=32,
                                              include_particles=False)
        spatial_methods = {
            "macroscale": pybamm.FiniteVolume(),
            "current collector": pybamm.ScikitFiniteElement(),
        }
        disc = pybamm.Discretisation(mesh, spatial_methods)
        disc.process_model(model)

        # solve model
        solver = pybamm.AlgebraicSolver()
        solution = solver.solve(model)

        z = mesh["current collector"].coordinates[1, :][:, np.newaxis]
        u_exact = z**2 / 2 - 1 / 6
        np.testing.assert_array_almost_equal(solution.y[:-1],
                                             u_exact,
                                             decimal=1)
    def test_dirichlet_bcs(self):
        # manufactured solution u = a*z^2 + b*z + c
        model = pybamm.BaseModel()
        a = 3
        b = 4
        c = 5
        u = pybamm.Variable("variable", domain="current collector")
        model.algebraic = {u: -pybamm.laplacian(u) + pybamm.source(2 * a, u)}
        # set boundary conditions ("negative tab" = bottom of unit square,
        # "positive tab" = top of unit square, elsewhere normal derivative is zero)
        model.boundary_conditions = {
            u: {
                "negative tab": (pybamm.Scalar(c), "Dirichlet"),
                "positive tab": (pybamm.Scalar(a + b + c), "Dirichlet"),
            }
        }
        # bad initial guess (on purpose)
        model.initial_conditions = {u: pybamm.Scalar(1)}
        model.variables = {"u": u}
        # create discretisation
        mesh = get_unit_2p1D_mesh_for_testing(ypts=8,
                                              zpts=32,
                                              include_particles=False)
        spatial_methods = {
            "macroscale": pybamm.FiniteVolume(),
            "current collector": pybamm.ScikitFiniteElement(),
        }
        disc = pybamm.Discretisation(mesh, spatial_methods)
        disc.process_model(model)

        # solve model
        solver = pybamm.AlgebraicSolver()
        solution = solver.solve(model)

        # indepedent of y, so just check values for one y
        z = mesh["current collector"].edges["z"][:, np.newaxis]
        u_exact = a * z**2 + b * z + c
        np.testing.assert_array_almost_equal(solution.y[0:len(z)], u_exact)
    def test_manufactured_solution(self):
        mesh = get_unit_2p1D_mesh_for_testing(ypts=32, zpts=32, include_particles=False)
        spatial_methods = {
            "macroscale": pybamm.FiniteVolume(),
            "current collector": pybamm.ScikitFiniteElement(),
        }
        disc = pybamm.Discretisation(mesh, spatial_methods)

        # linear u = z (to test coordinates to degree of freedom mapping)
        var = pybamm.Variable("var", domain="current collector")
        disc.set_variable_slices([var])
        var_disc = disc.process_symbol(var)
        z_vertices = mesh["current collector"].coordinates[1, :]
        np.testing.assert_array_almost_equal(
            var_disc.evaluate(None, z_vertices), z_vertices[:, np.newaxis]
        )

        # linear u = 6*y (to test coordinates to degree of freedom mapping)
        y_vertices = mesh["current collector"].coordinates[0, :]
        np.testing.assert_array_almost_equal(
            var_disc.evaluate(None, 6 * y_vertices), 6 * y_vertices[:, np.newaxis]
        )

        # mixed u = y*z (to test coordinates to degree of freedom mapping)
        np.testing.assert_array_almost_equal(
            var_disc.evaluate(None, y_vertices * z_vertices),
            y_vertices[:, np.newaxis] * z_vertices[:, np.newaxis],
        )

        # laplace of u = sin(pi*z)
        var = pybamm.Variable("var", domain="current collector")
        eqn_zz = pybamm.laplacian(var)
        # set boundary conditions ("negative tab" = bottom of unit square,
        # "positive tab" = top of unit square, elsewhere normal derivative is zero)
        disc.bcs = {
            var.id: {
                "negative tab": (pybamm.Scalar(0), "Dirichlet"),
                "positive tab": (pybamm.Scalar(0), "Dirichlet"),
            }
        }
        disc.set_variable_slices([var])
        eqn_zz_disc = disc.process_symbol(eqn_zz)
        z_vertices = mesh["current collector"].coordinates[1, :][:, np.newaxis]
        u = np.sin(np.pi * z_vertices)
        mass = pybamm.Mass(var)
        mass_disc = disc.process_symbol(mass)
        soln = -np.pi ** 2 * u
        np.testing.assert_array_almost_equal(
            eqn_zz_disc.evaluate(None, u), mass_disc.entries @ soln, decimal=3
        )

        # laplace of u = cos(pi*y)*sin(pi*z)
        var = pybamm.Variable("var", domain="current collector")
        laplace_eqn = pybamm.laplacian(var)
        # set boundary conditions ("negative tab" = bottom of unit square,
        # "positive tab" = top of unit square, elsewhere normal derivative is zero)
        disc.bcs = {
            var.id: {
                "negative tab": (pybamm.Scalar(0), "Dirichlet"),
                "positive tab": (pybamm.Scalar(0), "Dirichlet"),
            }
        }
        disc.set_variable_slices([var])
        laplace_eqn_disc = disc.process_symbol(laplace_eqn)
        y_vertices = mesh["current collector"].coordinates[0, :][:, np.newaxis]
        z_vertices = mesh["current collector"].coordinates[1, :][:, np.newaxis]
        u = np.cos(np.pi * y_vertices) * np.sin(np.pi * z_vertices)
        mass = pybamm.Mass(var)
        mass_disc = disc.process_symbol(mass)
        soln = -np.pi ** 2 * u
        np.testing.assert_array_almost_equal(
            laplace_eqn_disc.evaluate(None, u), mass_disc.entries @ soln, decimal=2
        )