NvidiaWarp-GarmentCode/warp/tests/test_fem.py

# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

import math
import unittest

import numpy as np
import warp as wp
from warp.tests.unittest_utils import *


from warp.fem import Field, Sample, Domain, Coords
from warp.fem import integrand, div, grad, curl, D, normal
import warp.fem as fem

from warp.fem.types import make_free_sample
from warp.fem.geometry.closest_point import project_on_tri_at_origin, project_on_tet_at_origin
from warp.fem.geometry import DeformedGeometry
from warp.fem.space import shape
from warp.fem.cache import dynamic_kernel
from warp.fem.utils import grid_to_tets, grid_to_tris, grid_to_quads, grid_to_hexes

wp.init()


@integrand
def linear_form(s: Sample, u: Field):
    return u(s)


def test_integrate_gradient(test_case, device):
    with wp.ScopedDevice(device):
        # Grid geometry
        geo = fem.Grid2D(res=wp.vec2i(5))

        # Domain and function spaces
        domain = fem.Cells(geometry=geo)
        quadrature = fem.RegularQuadrature(domain=domain, order=3)

        scalar_space = fem.make_polynomial_space(geo, degree=3)

        u = scalar_space.make_field()
        u.dof_values = wp.zeros_like(u.dof_values, requires_grad=True)

        result = wp.empty(dtype=wp.float64, shape=(1), requires_grad=True)

        tape = wp.Tape()

        # forward pass
        with tape:
            fem.integrate(linear_form, quadrature=quadrature, fields={"u": u}, output=result)

        tape.backward(result)

        test = fem.make_test(space=scalar_space, domain=domain)
        rhs = fem.integrate(linear_form, quadrature=quadrature, fields={"u": test})

        err = np.linalg.norm(rhs.numpy() - u.dof_values.grad.numpy())
        test_case.assertLess(err, 1.0e-8)


@fem.integrand
def bilinear_field(s: fem.Sample, domain: fem.Domain):
    x = domain(s)
    return x[0] * x[1]


@fem.integrand
def grad_field(s: fem.Sample, p: fem.Field):
    return fem.grad(p, s)


def test_interpolate_gradient(test_case, device):
    with wp.ScopedDevice(device):
        # Quad mesh with single element
        # so we can test gradient with respect to vertex positions
        positions = wp.array([[0.0, 0.0], [0.0, 2.0], [2.0, 0.0], [2.0, 2.0]], dtype=wp.vec2, requires_grad=True)
        quads = wp.array([[0, 2, 3, 1]], dtype=int)
        geo = fem.Quadmesh2D(quads, positions)

        # Quadratic scalar space
        scalar_space = fem.make_polynomial_space(geo, degree=2)

        # Point-based vector space
        # So we can test gradient with respect to inteprolation point position
        point_coords = wp.array([[[0.5, 0.5, 0.0]]], dtype=fem.Coords, requires_grad=True)
        interpolation_nodes = fem.PointBasisSpace(
            fem.ExplicitQuadrature(domain=fem.Cells(geo), points=point_coords, weights=wp.array([[1.0]], dtype=float))
        )
        vector_space = fem.make_collocated_function_space(interpolation_nodes, dtype=wp.vec2)

        # Initialize scalar field with known function
        scalar_field = scalar_space.make_field()
        scalar_field.dof_values.requires_grad = True
        fem.interpolate(bilinear_field, dest=scalar_field)

        # Interpolate gradient at center point
        vector_field = vector_space.make_field()
        vector_field.dof_values.requires_grad = True
        tape = wp.Tape()
        with tape:
            fem.interpolate(grad_field, dest=vector_field, fields={"p": scalar_field})

        assert_np_equal(vector_field.dof_values.numpy(), np.array([[1.0, 1.0]]))

        vector_field.dof_values.grad.assign([1.0, 0.0])
        tape.backward()

        assert_np_equal(scalar_field.dof_values.grad.numpy(), np.array([0.0, 0.0, 0.0, 0.0, 0.0, -0.5, 0.0, 0.5, 0.0]))
        assert_np_equal(
            geo.positions.grad.numpy(),
            np.array(
                [
                    [0.25, 0.25],
                    [0.25, 0.25],
                    [-0.25, -0.25],
                    [-0.25, -0.25],
                ]
            ),
        )
        assert_np_equal(point_coords.grad.numpy(), np.array([[[0.0, 2.0, 0.0]]]))

        tape.zero()
        scalar_field.dof_values.grad.zero_()
        geo.positions.grad.zero_()
        point_coords.grad.zero_()

        vector_field.dof_values.grad.assign([0.0, 1.0])
        tape.backward()

        assert_np_equal(scalar_field.dof_values.grad.numpy(), np.array([0.0, 0.0, 0.0, 0.0, -0.5, 0.0, 0.5, 0.0, 0.0]))
        assert_np_equal(
            geo.positions.grad.numpy(),
            np.array(
                [
                    [0.25, 0.25],
                    [-0.25, -0.25],
                    [0.25, 0.25],
                    [-0.25, -0.25],
                ]
            ),
        )
        assert_np_equal(point_coords.grad.numpy(), np.array([[[2.0, 0.0, 0.0]]]))

@integrand
def vector_divergence_form(s: Sample, u: Field, q: Field):
    return div(u, s) * q(s)


@integrand
def vector_grad_form(s: Sample, u: Field, q: Field):
    return wp.dot(u(s), grad(q, s))


@integrand
def vector_boundary_form(domain: Domain, s: Sample, u: Field, q: Field):
    return wp.dot(u(s) * q(s), normal(domain, s))


def test_vector_divergence_theorem(test_case, device):
    rng = np.random.default_rng(123)

    with wp.ScopedDevice(device):
        # Grid geometry
        geo = fem.Grid2D(res=wp.vec2i(5))

        # Domain and function spaces
        interior = fem.Cells(geometry=geo)
        boundary = fem.BoundarySides(geometry=geo)

        vector_space = fem.make_polynomial_space(geo, degree=2, dtype=wp.vec2)
        scalar_space = fem.make_polynomial_space(geo, degree=1, dtype=float)

        u = vector_space.make_field()
        u.dof_values = rng.random(size=(u.dof_values.shape[0], 2))

        # Divergence theorem
        constant_one = scalar_space.make_field()
        constant_one.dof_values.fill_(1.0)

        interior_quadrature = fem.RegularQuadrature(domain=interior, order=vector_space.degree)
        boundary_quadrature = fem.RegularQuadrature(domain=boundary, order=vector_space.degree)
        div_int = fem.integrate(
            vector_divergence_form,
            quadrature=interior_quadrature,
            fields={"u": u, "q": constant_one},
            kernel_options={"enable_backward": False},
        )
        boundary_int = fem.integrate(
            vector_boundary_form,
            quadrature=boundary_quadrature,
            fields={"u": u.trace(), "q": constant_one.trace()},
            kernel_options={"enable_backward": False},
        )

        test_case.assertAlmostEqual(div_int, boundary_int, places=5)

        # Integration by parts
        q = scalar_space.make_field()
        q.dof_values = rng.random(size=q.dof_values.shape[0])

        interior_quadrature = fem.RegularQuadrature(domain=interior, order=vector_space.degree + scalar_space.degree)
        boundary_quadrature = fem.RegularQuadrature(domain=boundary, order=vector_space.degree + scalar_space.degree)
        div_int = fem.integrate(
            vector_divergence_form,
            quadrature=interior_quadrature,
            fields={"u": u, "q": q},
            kernel_options={"enable_backward": False},
        )
        grad_int = fem.integrate(
            vector_grad_form,
            quadrature=interior_quadrature,
            fields={"u": u, "q": q},
            kernel_options={"enable_backward": False},
        )
        boundary_int = fem.integrate(
            vector_boundary_form,
            quadrature=boundary_quadrature,
            fields={"u": u.trace(), "q": q.trace()},
            kernel_options={"enable_backward": False},
        )

        test_case.assertAlmostEqual(div_int + grad_int, boundary_int, places=5)


@integrand
def tensor_divergence_form(s: Sample, tau: Field, v: Field):
    return wp.dot(div(tau, s), v(s))


@integrand
def tensor_grad_form(s: Sample, tau: Field, v: Field):
    return wp.ddot(wp.transpose(tau(s)), grad(v, s))


@integrand
def tensor_boundary_form(domain: Domain, s: Sample, tau: Field, v: Field):
    return wp.dot(tau(s) * v(s), normal(domain, s))


def test_tensor_divergence_theorem(test_case, device):
    rng = np.random.default_rng(123)

    with wp.ScopedDevice(device):
        # Grid geometry
        geo = fem.Grid2D(res=wp.vec2i(5))

        # Domain and function spaces
        interior = fem.Cells(geometry=geo)
        boundary = fem.BoundarySides(geometry=geo)

        tensor_space = fem.make_polynomial_space(geo, degree=2, dtype=wp.mat22)
        vector_space = fem.make_polynomial_space(geo, degree=1, dtype=wp.vec2)

        tau = tensor_space.make_field()
        tau.dof_values = rng.random(size=(tau.dof_values.shape[0], 2, 2))

        # Divergence theorem
        constant_vec = vector_space.make_field()
        constant_vec.dof_values.fill_(wp.vec2(0.5, 2.0))

        interior_quadrature = fem.RegularQuadrature(domain=interior, order=tensor_space.degree)
        boundary_quadrature = fem.RegularQuadrature(domain=boundary, order=tensor_space.degree)
        div_int = fem.integrate(
            tensor_divergence_form,
            quadrature=interior_quadrature,
            fields={"tau": tau, "v": constant_vec},
            kernel_options={"enable_backward": False},
        )
        boundary_int = fem.integrate(
            tensor_boundary_form,
            quadrature=boundary_quadrature,
            fields={"tau": tau.trace(), "v": constant_vec.trace()},
            kernel_options={"enable_backward": False},
        )

        test_case.assertAlmostEqual(div_int, boundary_int, places=5)

        # Integration by parts
        v = vector_space.make_field()
        v.dof_values = rng.random(size=(v.dof_values.shape[0], 2))

        interior_quadrature = fem.RegularQuadrature(domain=interior, order=tensor_space.degree + vector_space.degree)
        boundary_quadrature = fem.RegularQuadrature(domain=boundary, order=tensor_space.degree + vector_space.degree)
        div_int = fem.integrate(
            tensor_divergence_form,
            quadrature=interior_quadrature,
            fields={"tau": tau, "v": v},
            kernel_options={"enable_backward": False},
        )
        grad_int = fem.integrate(
            tensor_grad_form,
            quadrature=interior_quadrature,
            fields={"tau": tau, "v": v},
            kernel_options={"enable_backward": False},
        )
        boundary_int = fem.integrate(
            tensor_boundary_form,
            quadrature=boundary_quadrature,
            fields={"tau": tau.trace(), "v": v.trace()},
            kernel_options={"enable_backward": False},
        )

        test_case.assertAlmostEqual(div_int + grad_int, boundary_int, places=5)


@integrand
def grad_decomposition(s: Sample, u: Field, v: Field):
    return wp.length_sq(grad(u, s) * v(s) - D(u, s) * v(s) - wp.cross(curl(u, s), v(s)))


def test_grad_decomposition(test_case, device):
    rng = np.random.default_rng(123)

    with wp.ScopedDevice(device):
        # Grid geometry
        geo = fem.Grid3D(res=wp.vec3i(5))

        # Domain and function spaces
        domain = fem.Cells(geometry=geo)
        quadrature = fem.RegularQuadrature(domain=domain, order=4)

        vector_space = fem.make_polynomial_space(geo, degree=2, dtype=wp.vec3)
        u = vector_space.make_field()

        u.dof_values = rng.random(size=(u.dof_values.shape[0], 3))

        err = fem.integrate(grad_decomposition, quadrature=quadrature, fields={"u": u, "v": u})
        test_case.assertLess(err, 1.0e-8)


def _gen_trimesh(N):
    x = np.linspace(0.0, 1.0, N + 1)
    y = np.linspace(0.0, 1.0, N + 1)

    positions = np.transpose(np.meshgrid(x, y, indexing="ij")).reshape(-1, 2)

    vidx = grid_to_tris(N, N)

    return wp.array(positions, dtype=wp.vec2), wp.array(vidx, dtype=int)


def _gen_quadmesh(N):
    x = np.linspace(0.0, 1.0, N + 1)
    y = np.linspace(0.0, 1.0, N + 1)

    positions = np.transpose(np.meshgrid(x, y, indexing="ij")).reshape(-1, 2)

    vidx = grid_to_quads(N, N)

    return wp.array(positions, dtype=wp.vec2), wp.array(vidx, dtype=int)


def _gen_tetmesh(N):
    x = np.linspace(0.0, 1.0, N + 1)
    y = np.linspace(0.0, 1.0, N + 1)
    z = np.linspace(0.0, 1.0, N + 1)

    positions = np.transpose(np.meshgrid(x, y, z, indexing="ij")).reshape(-1, 3)

    vidx = grid_to_tets(N, N, N)

    return wp.array(positions, dtype=wp.vec3), wp.array(vidx, dtype=int)


def _gen_hexmesh(N):
    x = np.linspace(0.0, 1.0, N + 1)
    y = np.linspace(0.0, 1.0, N + 1)
    z = np.linspace(0.0, 1.0, N + 1)

    positions = np.transpose(np.meshgrid(x, y, z, indexing="ij")).reshape(-1, 3)

    vidx = grid_to_hexes(N, N, N)

    return wp.array(positions, dtype=wp.vec3), wp.array(vidx, dtype=int)


def _launch_test_geometry_kernel(geo: fem.Geometry, device):
    @dynamic_kernel(suffix=geo.name, kernel_options={"enable_backward": False})
    def test_geo_cells_kernel(
        cell_arg: geo.CellArg,
        qps: wp.array(dtype=Coords),
        qp_weights: wp.array(dtype=float),
        cell_measures: wp.array(dtype=float),
    ):
        cell_index, q = wp.tid()

        coords = qps[q]
        s = make_free_sample(cell_index, coords)

        wp.atomic_add(cell_measures, cell_index, geo.cell_measure(cell_arg, s) * qp_weights[q])

    REF_MEASURE = geo.reference_side().measure()

    @dynamic_kernel(suffix=geo.name, kernel_options={"enable_backward": False, "max_unroll": 1})
    def test_geo_sides_kernel(
        side_arg: geo.SideArg,
        qps: wp.array(dtype=Coords),
        qp_weights: wp.array(dtype=float),
        side_measures: wp.array(dtype=float),
    ):
        side_index, q = wp.tid()

        coords = qps[q]
        s = make_free_sample(side_index, coords)

        cell_arg = geo.side_to_cell_arg(side_arg)
        inner_cell_index = geo.side_inner_cell_index(side_arg, side_index)
        outer_cell_index = geo.side_outer_cell_index(side_arg, side_index)
        inner_cell_coords = geo.side_inner_cell_coords(side_arg, side_index, coords)
        outer_cell_coords = geo.side_outer_cell_coords(side_arg, side_index, coords)

        inner_s = make_free_sample(inner_cell_index, inner_cell_coords)
        outer_s = make_free_sample(outer_cell_index, outer_cell_coords)

        pos_side = geo.side_position(side_arg, s)
        pos_inner = geo.cell_position(cell_arg, inner_s)
        pos_outer = geo.cell_position(cell_arg, outer_s)

        for k in range(type(pos_side).length):
            wp.expect_near(pos_side[k], pos_inner[k], 0.0001)
            wp.expect_near(pos_side[k], pos_outer[k], 0.0001)

        inner_side_coords = geo.side_from_cell_coords(side_arg, side_index, inner_cell_index, inner_cell_coords)
        outer_side_coords = geo.side_from_cell_coords(side_arg, side_index, outer_cell_index, outer_cell_coords)

        wp.expect_near(coords, inner_side_coords, 0.0001)
        wp.expect_near(coords, outer_side_coords, 0.0001)

        vol = geo.side_measure(side_arg, s)
        wp.atomic_add(side_measures, side_index, vol * qp_weights[q])

        # test consistency of side normal, measure, and deformation gradient
        F = geo.side_deformation_gradient(side_arg, s)
        F_det = DeformedGeometry._side_measure(F)
        wp.expect_near(F_det * REF_MEASURE, vol)

        nor = geo.side_normal(side_arg, s)
        F_cross = DeformedGeometry._side_normal(F)

        for k in range(type(pos_side).length):
            wp.expect_near(F_cross[k], nor[k], 0.0001)

    cell_measures = wp.zeros(dtype=float, device=device, shape=geo.cell_count())

    cell_quadrature = fem.RegularQuadrature(fem.Cells(geo), order=2)
    cell_qps = wp.array(cell_quadrature.points, dtype=Coords, device=device)
    cell_qp_weights = wp.array(cell_quadrature.weights, dtype=float, device=device)

    wp.launch(
        kernel=test_geo_cells_kernel,
        dim=(geo.cell_count(), cell_qps.shape[0]),
        inputs=[geo.cell_arg_value(device), cell_qps, cell_qp_weights, cell_measures],
        device=device,
    )

    side_measures = wp.zeros(dtype=float, device=device, shape=geo.side_count())

    side_quadrature = fem.RegularQuadrature(fem.Sides(geo), order=2)
    side_qps = wp.array(side_quadrature.points, dtype=Coords, device=device)
    side_qp_weights = wp.array(side_quadrature.weights, dtype=float, device=device)

    wp.launch(
        kernel=test_geo_sides_kernel,
        dim=(geo.side_count(), side_qps.shape[0]),
        inputs=[geo.side_arg_value(device), side_qps, side_qp_weights, side_measures],
        device=device,
    )

    return side_measures, cell_measures


def test_grid_2d(test_case, device):
    N = 3

    geo = fem.Grid2D(res=wp.vec2i(N))

    test_case.assertEqual(geo.cell_count(), N**2)
    test_case.assertEqual(geo.vertex_count(), (N + 1) ** 2)
    test_case.assertEqual(geo.side_count(), 2 * (N + 1) * N)
    test_case.assertEqual(geo.boundary_side_count(), 4 * N)

    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device)

    assert_np_equal(side_measures.numpy(), np.full(side_measures.shape, 1.0 / (N)), tol=1.0e-4)
    assert_np_equal(cell_measures.numpy(), np.full(cell_measures.shape, 1.0 / (N**2)), tol=1.0e-4)


def test_triangle_mesh(test_case, device):
    N = 3

    with wp.ScopedDevice(device):
        positions, tri_vidx = _gen_trimesh(N)

    geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)

    test_case.assertEqual(geo.cell_count(), 2 * (N) ** 2)
    test_case.assertEqual(geo.vertex_count(), (N + 1) ** 2)
    test_case.assertEqual(geo.side_count(), 2 * (N + 1) * N + (N**2))
    test_case.assertEqual(geo.boundary_side_count(), 4 * N)

    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device)

    assert_np_equal(cell_measures.numpy(), np.full(cell_measures.shape, 0.5 / (N**2)), tol=1.0e-4)
    test_case.assertAlmostEqual(np.sum(side_measures.numpy()), 2 * (N + 1) + N * math.sqrt(2.0), places=4)


def test_quad_mesh(test_case, device):
    N = 3

    with wp.ScopedDevice(device):
        positions, quad_vidx = _gen_quadmesh(N)

    geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)

    test_case.assertEqual(geo.cell_count(), N**2)
    test_case.assertEqual(geo.vertex_count(), (N + 1) ** 2)
    test_case.assertEqual(geo.side_count(), 2 * (N + 1) * N)
    test_case.assertEqual(geo.boundary_side_count(), 4 * N)

    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device)

    assert_np_equal(side_measures.numpy(), np.full(side_measures.shape, 1.0 / (N)), tol=1.0e-4)
    assert_np_equal(cell_measures.numpy(), np.full(cell_measures.shape, 1.0 / (N**2)), tol=1.0e-4)


def test_grid_3d(test_case, device):
    N = 3

    geo = fem.Grid3D(res=wp.vec3i(N))

    test_case.assertEqual(geo.cell_count(), (N) ** 3)
    test_case.assertEqual(geo.vertex_count(), (N + 1) ** 3)
    test_case.assertEqual(geo.side_count(), 3 * (N + 1) * N**2)
    test_case.assertEqual(geo.boundary_side_count(), 6 * N * N)
    test_case.assertEqual(geo.edge_count(), 3 * N * (N + 1) ** 2)

    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device)

    assert_np_equal(side_measures.numpy(), np.full(side_measures.shape, 1.0 / (N**2)), tol=1.0e-4)
    assert_np_equal(cell_measures.numpy(), np.full(cell_measures.shape, 1.0 / (N**3)), tol=1.0e-4)


def test_tet_mesh(test_case, device):
    N = 3

    with wp.ScopedDevice(device):
        positions, tet_vidx = _gen_tetmesh(N)

    geo = fem.Tetmesh(tet_vertex_indices=tet_vidx, positions=positions)

    test_case.assertEqual(geo.cell_count(), 5 * (N) ** 3)
    test_case.assertEqual(geo.vertex_count(), (N + 1) ** 3)
    test_case.assertEqual(geo.side_count(), 6 * (N + 1) * N**2 + (N**3) * 4)
    test_case.assertEqual(geo.boundary_side_count(), 12 * N * N)
    test_case.assertEqual(geo.edge_count(), 3 * N * (N + 1) * (2 * N + 1))

    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device)

    test_case.assertAlmostEqual(np.sum(cell_measures.numpy()), 1.0, places=4)
    test_case.assertAlmostEqual(np.sum(side_measures.numpy()), 0.5 * 6 * (N + 1) + N * 2 * math.sqrt(3.0), places=4)


def test_hex_mesh(test_case, device):
    N = 3

    with wp.ScopedDevice(device):
        positions, tet_vidx = _gen_hexmesh(N)

    geo = fem.Hexmesh(hex_vertex_indices=tet_vidx, positions=positions)

    test_case.assertEqual(geo.cell_count(), (N) ** 3)
    test_case.assertEqual(geo.vertex_count(), (N + 1) ** 3)
    test_case.assertEqual(geo.side_count(), 3 * (N + 1) * N**2)
    test_case.assertEqual(geo.boundary_side_count(), 6 * N * N)
    test_case.assertEqual(geo.edge_count(), 3 * N * (N + 1) ** 2)

    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device)

    assert_np_equal(side_measures.numpy(), np.full(side_measures.shape, 1.0 / (N**2)), tol=1.0e-4)
    assert_np_equal(cell_measures.numpy(), np.full(cell_measures.shape, 1.0 / (N**3)), tol=1.0e-4)


@integrand
def _rigid_deformation_field(s: Sample, domain: Domain, translation: wp.vec3, rotation: wp.vec3, scale: float):
    q = wp.quat_from_axis_angle(wp.normalize(rotation), wp.length(rotation))
    return translation + scale * wp.quat_rotate(q, domain(s)) - domain(s)


def test_deformed_geometry(test_case, device):
    N = 3

    with wp.ScopedDevice(device):
        positions, tet_vidx = _gen_tetmesh(N)

    geo = fem.Tetmesh(tet_vertex_indices=tet_vidx, positions=positions)

    translation = [1.0, 2.0, 3.0]
    rotation = [0.0, math.pi / 4.0, 0.0]
    scale = 2.0

    vector_space = fem.make_polynomial_space(geo, dtype=wp.vec3, degree=2)
    pos_field = vector_space.make_field()
    fem.interpolate(
        _rigid_deformation_field,
        dest=pos_field,
        values={"translation": translation, "rotation": rotation, "scale": scale},
    )

    deformed_geo = pos_field.make_deformed_geometry()

    # rigidly-deformed geometry

    test_case.assertEqual(geo.cell_count(), 5 * (N) ** 3)
    test_case.assertEqual(geo.vertex_count(), (N + 1) ** 3)
    test_case.assertEqual(geo.side_count(), 6 * (N + 1) * N**2 + (N**3) * 4)
    test_case.assertEqual(geo.boundary_side_count(), 12 * N * N)

    side_measures, cell_measures = _launch_test_geometry_kernel(deformed_geo, device)

    test_case.assertAlmostEqual(np.sum(cell_measures.numpy()), scale**3, places=4)
    test_case.assertAlmostEqual(
        np.sum(side_measures.numpy()), scale**2 * (0.5 * 6 * (N + 1) + N * 2 * math.sqrt(3.0)), places=4
    )

    @wp.kernel
    def _test_deformed_geometry_normal(
        geo_index_arg: geo.SideIndexArg, geo_arg: geo.SideArg, def_arg: deformed_geo.SideArg, rotation: wp.vec3
    ):
        i = wp.tid()
        side_index = deformed_geo.boundary_side_index(geo_index_arg, i)

        s = make_free_sample(side_index, Coords(0.5, 0.5, 0.0))
        geo_n = geo.side_normal(geo_arg, s)
        def_n = deformed_geo.side_normal(def_arg, s)

        q = wp.quat_from_axis_angle(wp.normalize(rotation), wp.length(rotation))
        wp.expect_near(wp.quat_rotate(q, geo_n), def_n, 0.001)

    wp.launch(
        _test_deformed_geometry_normal,
        dim=geo.boundary_side_count(),
        device=device,
        inputs=[
            geo.side_index_arg_value(device),
            geo.side_arg_value(device),
            deformed_geo.side_arg_value(device),
            rotation,
        ],
    )

    wp.synchronize()


@wp.kernel
def _test_closest_point_on_tri_kernel(
    e0: wp.vec2,
    e1: wp.vec2,
    points: wp.array(dtype=wp.vec2),
    sq_dist: wp.array(dtype=float),
    coords: wp.array(dtype=Coords),
):
    i = wp.tid()
    d2, c = project_on_tri_at_origin(points[i], e0, e1)
    sq_dist[i] = d2
    coords[i] = c


@wp.kernel
def _test_closest_point_on_tet_kernel(
    e0: wp.vec3,
    e1: wp.vec3,
    e2: wp.vec3,
    points: wp.array(dtype=wp.vec3),
    sq_dist: wp.array(dtype=float),
    coords: wp.array(dtype=Coords),
):
    i = wp.tid()
    d2, c = project_on_tet_at_origin(points[i], e0, e1, e2)
    sq_dist[i] = d2
    coords[i] = c


def test_closest_point_queries(test_case, device):
    # Test some simple lookup queries
    e0 = wp.vec2(2.0, 0.0)
    e1 = wp.vec2(0.0, 2.0)

    points = wp.array(
        (
            [-1.0, -1.0],
            [0.5, 0.5],
            [1.0, 1.0],
            [2.0, 2.0],
        ),
        dtype=wp.vec2,
        device=device,
    )
    expected_sq_dist = np.array([2.0, 0.0, 0.0, 2.0])
    expected_coords = np.array([[1.0, 0.0, 0.0], [0.5, 0.25, 0.25], [0.0, 0.5, 0.5], [0.0, 0.5, 0.5]])

    sq_dist = wp.empty(shape=points.shape, dtype=float, device=device)
    coords = wp.empty(shape=points.shape, dtype=Coords, device=device)
    wp.launch(
        _test_closest_point_on_tri_kernel, dim=points.shape, device=device, inputs=[e0, e1, points, sq_dist, coords]
    )

    assert_np_equal(coords.numpy(), expected_coords)
    assert_np_equal(sq_dist.numpy(), expected_sq_dist)

    # Tet

    e0 = wp.vec3(3.0, 0.0, 0.0)
    e1 = wp.vec3(0.0, 3.0, 0.0)
    e2 = wp.vec3(0.0, 0.0, 3.0)

    points = wp.array(
        (
            [-1.0, -1.0, -1.0],
            [0.5, 0.5, 0.5],
            [1.0, 1.0, 1.0],
            [2.0, 2.0, 2.0],
        ),
        dtype=wp.vec3,
        device=device,
    )
    expected_sq_dist = np.array([3.0, 0.0, 0.0, 3.0])
    expected_coords = np.array(
        [
            [0.0, 0.0, 0.0],
            [1.0 / 6.0, 1.0 / 6.0, 1.0 / 6.0],
            [1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0],
            [1.0 / 3.0, 1.0 / 3.0, 1.0 / 3.0],
        ]
    )

    sq_dist = wp.empty(shape=points.shape, dtype=float, device=device)
    coords = wp.empty(shape=points.shape, dtype=Coords, device=device)
    wp.launch(
        _test_closest_point_on_tet_kernel, dim=points.shape, device=device, inputs=[e0, e1, e2, points, sq_dist, coords]
    )

    assert_np_equal(coords.numpy(), expected_coords, tol=1.0e-4)
    assert_np_equal(sq_dist.numpy(), expected_sq_dist, tol=1.0e-4)


def test_regular_quadrature(test_case, device):
    from warp.fem.geometry.element import LinearEdge, Triangle, Polynomial

    for family in Polynomial:
        # test integrating monomials
        for degree in range(8):
            coords, weights = LinearEdge().instantiate_quadrature(degree, family=family)
            res = sum(w * pow(c[0], degree) for w, c in zip(weights, coords))
            ref = 1.0 / (degree + 1)

            test_case.assertAlmostEqual(ref, res, places=4)

        # test integrating y^k1 (1 - x)^k2 on triangle using transformation to square
        for x_degree in range(4):
            for y_degree in range(4):
                coords, weights = Triangle().instantiate_quadrature(x_degree + y_degree, family=family)
                res = 0.5 * sum(w * pow(1.0 - c[1], x_degree) * pow(c[2], y_degree) for w, c in zip(weights, coords))

                ref = 1.0 / ((x_degree + y_degree + 2) * (y_degree + 1))
                # print(x_degree, y_degree, family, len(coords), res, ref)
                test_case.assertAlmostEqual(ref, res, places=4)

    # test integrating y^k1 (1 - x)^k2 on triangle using direct formulas
    for x_degree in range(5):
        for y_degree in range(5):
            coords, weights = Triangle().instantiate_quadrature(x_degree + y_degree, family=None)
            res = 0.5 * sum(w * pow(1.0 - c[1], x_degree) * pow(c[2], y_degree) for w, c in zip(weights, coords))

            ref = 1.0 / ((x_degree + y_degree + 2) * (y_degree + 1))
            test_case.assertAlmostEqual(ref, res, places=4)


def test_dof_mapper(test_case, device):
    matrix_types = [wp.mat22, wp.mat33]

    # Symmetric mapper
    for mapping in fem.SymmetricTensorMapper.Mapping:
        for dtype in matrix_types:
            mapper = fem.SymmetricTensorMapper(dtype, mapping=mapping)
            dof_dtype = mapper.dof_dtype

            for k in range(dof_dtype._length_):
                elem = np.array(dof_dtype(0.0))
                elem[k] = 1.0
                dof_vec = dof_dtype(elem)

                mat = mapper.dof_to_value(dof_vec)
                dof_round_trip = mapper.value_to_dof(mat)

                # Check that value_to_dof(dof_to_value) is idempotent
                assert_np_equal(np.array(dof_round_trip), np.array(dof_vec))

                # Check that value is unitary for Frobenius norm 0.5 * |tau:tau|
                frob_norm2 = 0.5 * wp.ddot(mat, mat)
                test_case.assertAlmostEqual(frob_norm2, 1.0, places=6)

    # Skew-symmetric mapper
    for dtype in matrix_types:
        mapper = fem.SkewSymmetricTensorMapper(dtype)
        dof_dtype = mapper.dof_dtype

        if hasattr(dof_dtype, "_length_"):
            for k in range(dof_dtype._length_):
                elem = np.array(dof_dtype(0.0))
                elem[k] = 1.0
                dof_vec = dof_dtype(elem)

                mat = mapper.dof_to_value(dof_vec)
                dof_round_trip = mapper.value_to_dof(mat)

                # Check that value_to_dof(dof_to_value) is idempotent
                assert_np_equal(np.array(dof_round_trip), np.array(dof_vec))

                # Check that value is unitary for Frobenius norm 0.5 * |tau:tau|
                frob_norm2 = 0.5 * wp.ddot(mat, mat)
                test_case.assertAlmostEqual(frob_norm2, 1.0, places=6)
        else:
            dof_val = 1.0

            mat = mapper.dof_to_value(dof_val)
            dof_round_trip = mapper.value_to_dof(mat)

            test_case.assertAlmostEqual(dof_round_trip, dof_val)

            # Check that value is unitary for Frobenius norm 0.5 * |tau:tau|
            frob_norm2 = 0.5 * wp.ddot(mat, mat)
            test_case.assertAlmostEqual(frob_norm2, 1.0, places=6)


def test_shape_function_weight(test_case, shape: shape.ShapeFunction, coord_sampler, CENTER_COORDS):
    NODE_COUNT = shape.NODES_PER_ELEMENT
    weight_fn = shape.make_element_inner_weight()
    node_coords_fn = shape.make_node_coords_in_element()

    # Weight at node should be 1
    @dynamic_kernel(suffix=shape.name, kernel_options={"enable_backward": False})
    def node_unity_test():
        n = wp.tid()
        node_w = weight_fn(node_coords_fn(n), n)
        wp.expect_near(node_w, 1.0, places=5)

    wp.launch(node_unity_test, dim=NODE_COUNT, inputs=[])

    # Sum of node quadrature weights should be one (order 0)
    # Sum of weighted quadrature coords should be element center (order 1)
    node_quadrature_weight_fn = shape.make_node_quadrature_weight()

    @dynamic_kernel(suffix=shape.name, kernel_options={"enable_backward": False})
    def node_quadrature_unity_test():
        sum_node_qp = float(0.0)
        sum_node_qp_coords = Coords(0.0)

        for n in range(NODE_COUNT):
            w = node_quadrature_weight_fn(n)
            sum_node_qp += w
            sum_node_qp_coords += w * node_coords_fn(n)

        wp.expect_near(sum_node_qp, 1.0, 0.0001)
        wp.expect_near(sum_node_qp_coords, CENTER_COORDS, 0.0001)

    wp.launch(node_quadrature_unity_test, dim=1, inputs=[])

    @dynamic_kernel(suffix=shape.name, kernel_options={"enable_backward": False})
    def partition_of_unity_test():
        rng_state = wp.rand_init(4321, wp.tid())
        coords = coord_sampler(rng_state)

        # sum of node weights anywhere should be 1.0
        w_sum = float(0.0)
        for n in range(NODE_COUNT):
            w_sum += weight_fn(coords, n)

        wp.expect_near(w_sum, 1.0, 0.0001)

    n_samples = 100
    wp.launch(partition_of_unity_test, dim=n_samples, inputs=[])


def test_shape_function_trace(test_case, shape: shape.ShapeFunction, CENTER_COORDS):
    NODE_COUNT = shape.NODES_PER_ELEMENT
    node_coords_fn = shape.make_node_coords_in_element()

    # Sum of node quadrature weights should be one (order 0)
    # Sum of weighted quadrature coords should be element center (order 1)
    trace_node_quadrature_weight_fn = shape.make_trace_node_quadrature_weight()

    @dynamic_kernel(suffix=shape.name, kernel_options={"enable_backward": False})
    def trace_node_quadrature_unity_test():
        sum_node_qp = float(0.0)
        sum_node_qp_coords = Coords(0.0)

        for n in range(NODE_COUNT):
            coords = node_coords_fn(n)

            if wp.abs(coords[0]) < 1.0e-6:
                w = trace_node_quadrature_weight_fn(n)
                sum_node_qp += w
                sum_node_qp_coords += w * node_coords_fn(n)

        wp.expect_near(sum_node_qp, 1.0, 0.0001)
        wp.expect_near(sum_node_qp_coords, CENTER_COORDS, 0.0001)

    wp.launch(trace_node_quadrature_unity_test, dim=1, inputs=[])


def test_shape_function_gradient(test_case, shape: shape.ShapeFunction, coord_sampler, coord_delta_sampler):
    weight_fn = shape.make_element_inner_weight()
    weight_gradient_fn = shape.make_element_inner_weight_gradient()

    @dynamic_kernel(suffix=shape.name, kernel_options={"enable_backward": False})
    def finite_difference_test():
        i, n = wp.tid()
        rng_state = wp.rand_init(1234, i)

        coords = coord_sampler(rng_state)

        epsilon = 0.003
        param_delta, coords_delta = coord_delta_sampler(epsilon, rng_state)

        w_p = weight_fn(coords + coords_delta, n)
        w_m = weight_fn(coords - coords_delta, n)

        gp = weight_gradient_fn(coords + coords_delta, n)
        gm = weight_gradient_fn(coords - coords_delta, n)

        # 2nd-order finite-difference test
        # See Schroeder 2019, Practical course on computing derivatives in code
        delta_ref = w_p - w_m
        delta_est = wp.dot(gp + gm, param_delta)

        # wp.printf("%d %f %f \n", n, delta_ref, delta_est)
        wp.expect_near(delta_ref, delta_est, 0.0001)

    n_samples = 100
    wp.launch(finite_difference_test, dim=(n_samples, shape.NODES_PER_ELEMENT), inputs=[])


def test_square_shape_functions(test_case, device):
    SQUARE_CENTER_COORDS = wp.constant(Coords(0.5, 0.5, 0.0))
    SQUARE_SIDE_CENTER_COORDS = wp.constant(Coords(0.0, 0.5, 0.0))

    @wp.func
    def square_coord_sampler(state: wp.uint32):
        return Coords(wp.randf(state), wp.randf(state), 0.0)

    @wp.func
    def square_coord_delta_sampler(epsilon: float, state: wp.uint32):
        param_delta = wp.normalize(wp.vec2(wp.randf(state), wp.randf(state))) * epsilon
        return param_delta, Coords(param_delta[0], param_delta[1], 0.0)

    Q_1 = shape.SquareBipolynomialShapeFunctions(degree=1, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)
    Q_2 = shape.SquareBipolynomialShapeFunctions(degree=2, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)
    Q_3 = shape.SquareBipolynomialShapeFunctions(degree=3, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)

    test_shape_function_weight(test_case, Q_1, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_weight(test_case, Q_2, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_weight(test_case, Q_3, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_trace(test_case, Q_1, SQUARE_SIDE_CENTER_COORDS)
    test_shape_function_trace(test_case, Q_2, SQUARE_SIDE_CENTER_COORDS)
    test_shape_function_trace(test_case, Q_3, SQUARE_SIDE_CENTER_COORDS)
    test_shape_function_gradient(test_case, Q_1, square_coord_sampler, square_coord_delta_sampler)
    test_shape_function_gradient(test_case, Q_2, square_coord_sampler, square_coord_delta_sampler)
    test_shape_function_gradient(test_case, Q_3, square_coord_sampler, square_coord_delta_sampler)

    Q_1 = shape.SquareBipolynomialShapeFunctions(degree=1, family=fem.Polynomial.GAUSS_LEGENDRE)
    Q_2 = shape.SquareBipolynomialShapeFunctions(degree=2, family=fem.Polynomial.GAUSS_LEGENDRE)
    Q_3 = shape.SquareBipolynomialShapeFunctions(degree=3, family=fem.Polynomial.GAUSS_LEGENDRE)

    test_shape_function_weight(test_case, Q_1, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_weight(test_case, Q_2, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_weight(test_case, Q_3, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_gradient(test_case, Q_1, square_coord_sampler, square_coord_delta_sampler)
    test_shape_function_gradient(test_case, Q_2, square_coord_sampler, square_coord_delta_sampler)
    test_shape_function_gradient(test_case, Q_3, square_coord_sampler, square_coord_delta_sampler)

    S_2 = shape.SquareSerendipityShapeFunctions(degree=2, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)
    S_3 = shape.SquareSerendipityShapeFunctions(degree=3, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)

    test_shape_function_weight(test_case, S_2, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_weight(test_case, S_3, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_trace(test_case, S_2, SQUARE_SIDE_CENTER_COORDS)
    test_shape_function_trace(test_case, S_3, SQUARE_SIDE_CENTER_COORDS)
    test_shape_function_gradient(test_case, S_2, square_coord_sampler, square_coord_delta_sampler)
    test_shape_function_gradient(test_case, S_3, square_coord_sampler, square_coord_delta_sampler)

    P_c1 = shape.SquareNonConformingPolynomialShapeFunctions(degree=1)
    P_c2 = shape.SquareNonConformingPolynomialShapeFunctions(degree=2)
    P_c3 = shape.SquareNonConformingPolynomialShapeFunctions(degree=3)

    test_shape_function_weight(test_case, P_c1, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_weight(test_case, P_c2, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_weight(test_case, P_c3, square_coord_sampler, SQUARE_CENTER_COORDS)
    test_shape_function_gradient(test_case, P_c1, square_coord_sampler, square_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_c2, square_coord_sampler, square_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_c3, square_coord_sampler, square_coord_delta_sampler)

    wp.synchronize()


def test_cube_shape_functions(test_case, device):
    CUBE_CENTER_COORDS = wp.constant(Coords(0.5, 0.5, 0.5))
    CUBE_SIDE_CENTER_COORDS = wp.constant(Coords(0.0, 0.5, 0.5))

    @wp.func
    def cube_coord_sampler(state: wp.uint32):
        return Coords(wp.randf(state), wp.randf(state), wp.randf(state))

    @wp.func
    def cube_coord_delta_sampler(epsilon: float, state: wp.uint32):
        param_delta = wp.normalize(wp.vec3(wp.randf(state), wp.randf(state), wp.randf(state))) * epsilon
        return param_delta, param_delta

    Q_1 = shape.CubeTripolynomialShapeFunctions(degree=1, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)
    Q_2 = shape.CubeTripolynomialShapeFunctions(degree=2, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)
    Q_3 = shape.CubeTripolynomialShapeFunctions(degree=3, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)

    test_shape_function_weight(test_case, Q_1, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_weight(test_case, Q_2, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_weight(test_case, Q_3, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_trace(test_case, Q_1, CUBE_SIDE_CENTER_COORDS)
    test_shape_function_trace(test_case, Q_2, CUBE_SIDE_CENTER_COORDS)
    test_shape_function_trace(test_case, Q_3, CUBE_SIDE_CENTER_COORDS)
    test_shape_function_gradient(test_case, Q_1, cube_coord_sampler, cube_coord_delta_sampler)
    test_shape_function_gradient(test_case, Q_2, cube_coord_sampler, cube_coord_delta_sampler)
    test_shape_function_gradient(test_case, Q_3, cube_coord_sampler, cube_coord_delta_sampler)

    Q_1 = shape.CubeTripolynomialShapeFunctions(degree=1, family=fem.Polynomial.GAUSS_LEGENDRE)
    Q_2 = shape.CubeTripolynomialShapeFunctions(degree=2, family=fem.Polynomial.GAUSS_LEGENDRE)
    Q_3 = shape.CubeTripolynomialShapeFunctions(degree=3, family=fem.Polynomial.GAUSS_LEGENDRE)

    test_shape_function_weight(test_case, Q_1, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_weight(test_case, Q_2, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_weight(test_case, Q_3, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_gradient(test_case, Q_1, cube_coord_sampler, cube_coord_delta_sampler)
    test_shape_function_gradient(test_case, Q_2, cube_coord_sampler, cube_coord_delta_sampler)
    test_shape_function_gradient(test_case, Q_3, cube_coord_sampler, cube_coord_delta_sampler)

    S_2 = shape.CubeSerendipityShapeFunctions(degree=2, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)
    S_3 = shape.CubeSerendipityShapeFunctions(degree=3, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)

    test_shape_function_weight(test_case, S_2, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_weight(test_case, S_3, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_trace(test_case, S_2, CUBE_SIDE_CENTER_COORDS)
    test_shape_function_trace(test_case, S_3, CUBE_SIDE_CENTER_COORDS)
    test_shape_function_gradient(test_case, S_2, cube_coord_sampler, cube_coord_delta_sampler)
    test_shape_function_gradient(test_case, S_3, cube_coord_sampler, cube_coord_delta_sampler)

    P_c1 = shape.CubeNonConformingPolynomialShapeFunctions(degree=1)
    P_c2 = shape.CubeNonConformingPolynomialShapeFunctions(degree=2)
    P_c3 = shape.CubeNonConformingPolynomialShapeFunctions(degree=3)

    test_shape_function_weight(test_case, P_c1, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_weight(test_case, P_c2, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_weight(test_case, P_c3, cube_coord_sampler, CUBE_CENTER_COORDS)
    test_shape_function_gradient(test_case, P_c1, cube_coord_sampler, cube_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_c2, cube_coord_sampler, cube_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_c3, cube_coord_sampler, cube_coord_delta_sampler)

    wp.synchronize()


def test_tri_shape_functions(test_case, device):
    TRI_CENTER_COORDS = wp.constant(Coords(1 / 3.0, 1 / 3.0, 1 / 3.0))
    TRI_SIDE_CENTER_COORDS = wp.constant(Coords(0.0, 0.5, 0.5))

    @wp.func
    def tri_coord_sampler(state: wp.uint32):
        a = wp.randf(state)
        b = wp.randf(state)
        return Coords(1.0 - a - b, a, b)

    @wp.func
    def tri_coord_delta_sampler(epsilon: float, state: wp.uint32):
        param_delta = wp.normalize(wp.vec2(wp.randf(state), wp.randf(state))) * epsilon
        a = param_delta[0]
        b = param_delta[1]
        return param_delta, Coords(-a - b, a, b)

    P_1 = shape.Triangle2DPolynomialShapeFunctions(degree=1)
    P_2 = shape.Triangle2DPolynomialShapeFunctions(degree=2)
    P_3 = shape.Triangle2DPolynomialShapeFunctions(degree=3)

    test_shape_function_weight(test_case, P_1, tri_coord_sampler, TRI_CENTER_COORDS)
    test_shape_function_weight(test_case, P_2, tri_coord_sampler, TRI_CENTER_COORDS)
    test_shape_function_weight(test_case, P_3, tri_coord_sampler, TRI_CENTER_COORDS)
    test_shape_function_trace(test_case, P_1, TRI_SIDE_CENTER_COORDS)
    test_shape_function_trace(test_case, P_2, TRI_SIDE_CENTER_COORDS)
    test_shape_function_trace(test_case, P_3, TRI_SIDE_CENTER_COORDS)
    test_shape_function_gradient(test_case, P_1, tri_coord_sampler, tri_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_2, tri_coord_sampler, tri_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_3, tri_coord_sampler, tri_coord_delta_sampler)

    P_1d = shape.Triangle2DNonConformingPolynomialShapeFunctions(degree=1)
    P_2d = shape.Triangle2DNonConformingPolynomialShapeFunctions(degree=2)
    P_3d = shape.Triangle2DNonConformingPolynomialShapeFunctions(degree=3)

    test_shape_function_weight(test_case, P_1d, tri_coord_sampler, TRI_CENTER_COORDS)
    test_shape_function_weight(test_case, P_2d, tri_coord_sampler, TRI_CENTER_COORDS)
    test_shape_function_weight(test_case, P_3d, tri_coord_sampler, TRI_CENTER_COORDS)
    test_shape_function_gradient(test_case, P_1d, tri_coord_sampler, tri_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_2d, tri_coord_sampler, tri_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_3d, tri_coord_sampler, tri_coord_delta_sampler)

    wp.synchronize()


def test_tet_shape_functions(test_case, device):
    TET_CENTER_COORDS = wp.constant(Coords(1 / 4.0, 1 / 4.0, 1 / 4.0))
    TET_SIDE_CENTER_COORDS = wp.constant(Coords(0.0, 1.0 / 3.0, 1.0 / 3.0))

    @wp.func
    def tet_coord_sampler(state: wp.uint32):
        return Coords(wp.randf(state), wp.randf(state), wp.randf(state))

    @wp.func
    def tet_coord_delta_sampler(epsilon: float, state: wp.uint32):
        param_delta = wp.normalize(wp.vec3(wp.randf(state), wp.randf(state), wp.randf(state))) * epsilon
        return param_delta, param_delta

    P_1 = shape.TetrahedronPolynomialShapeFunctions(degree=1)
    P_2 = shape.TetrahedronPolynomialShapeFunctions(degree=2)
    P_3 = shape.TetrahedronPolynomialShapeFunctions(degree=3)

    test_shape_function_weight(test_case, P_1, tet_coord_sampler, TET_CENTER_COORDS)
    test_shape_function_weight(test_case, P_2, tet_coord_sampler, TET_CENTER_COORDS)
    test_shape_function_weight(test_case, P_3, tet_coord_sampler, TET_CENTER_COORDS)
    test_shape_function_trace(test_case, P_1, TET_SIDE_CENTER_COORDS)
    test_shape_function_trace(test_case, P_2, TET_SIDE_CENTER_COORDS)
    test_shape_function_trace(test_case, P_3, TET_SIDE_CENTER_COORDS)
    test_shape_function_gradient(test_case, P_1, tet_coord_sampler, tet_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_2, tet_coord_sampler, tet_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_3, tet_coord_sampler, tet_coord_delta_sampler)

    P_1d = shape.TetrahedronNonConformingPolynomialShapeFunctions(degree=1)
    P_2d = shape.TetrahedronNonConformingPolynomialShapeFunctions(degree=2)
    P_3d = shape.TetrahedronNonConformingPolynomialShapeFunctions(degree=3)

    test_shape_function_weight(test_case, P_1d, tet_coord_sampler, TET_CENTER_COORDS)
    test_shape_function_weight(test_case, P_2d, tet_coord_sampler, TET_CENTER_COORDS)
    test_shape_function_weight(test_case, P_3d, tet_coord_sampler, TET_CENTER_COORDS)
    test_shape_function_gradient(test_case, P_1d, tet_coord_sampler, tet_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_2d, tet_coord_sampler, tet_coord_delta_sampler)
    test_shape_function_gradient(test_case, P_3d, tet_coord_sampler, tet_coord_delta_sampler)

    wp.synchronize()


def test_point_basis(test_case, device):
    geo = fem.Grid2D(res=wp.vec2i(2))

    domain = fem.Cells(geo)

    quadrature = fem.RegularQuadrature(domain, order=2, family=fem.Polynomial.GAUSS_LEGENDRE)
    point_basis = fem.PointBasisSpace(quadrature)

    point_space = fem.make_collocated_function_space(point_basis)
    point_test = fem.make_test(point_space, domain=domain)

    # Sample at particle positions
    ones = fem.integrate(linear_form, fields={"u": point_test}, nodal=True)
    test_case.assertAlmostEqual(np.sum(ones.numpy()), 1.0, places=5)

    # Sampling outside of particle positions
    other_quadrature = fem.RegularQuadrature(domain, order=2, family=fem.Polynomial.LOBATTO_GAUSS_LEGENDRE)
    zeros = fem.integrate(linear_form, quadrature=other_quadrature, fields={"u": point_test})

    test_case.assertAlmostEqual(np.sum(zeros.numpy()), 0.0, places=5)


@fem.integrand
def _bicubic(s: Sample, domain: Domain):
    x = domain(s)
    return wp.pow(x[0], 3.0) * wp.pow(x[1], 3.0)


@fem.integrand
def _piecewise_constant(s: Sample):
    return float(s.element_index)


def test_particle_quadratures(test_case, device):
    geo = fem.Grid2D(res=wp.vec2i(2))

    domain = fem.Cells(geo)
    points, weights = domain.reference_element().instantiate_quadrature(order=4, family=fem.Polynomial.GAUSS_LEGENDRE)
    points_per_cell = len(points)

    points = points * domain.element_count()
    weights = weights * domain.element_count()

    points = wp.array(points, shape=(domain.element_count(), points_per_cell), dtype=Coords, device=device)
    weights = wp.array(weights, shape=(domain.element_count(), points_per_cell), dtype=float, device=device)

    explicit_quadrature = fem.ExplicitQuadrature(domain, points, weights)

    test_case.assertEqual(explicit_quadrature.points_per_element(), points_per_cell)
    test_case.assertEqual(explicit_quadrature.total_point_count(), points_per_cell * geo.cell_count())

    val = fem.integrate(_bicubic, quadrature=explicit_quadrature)
    test_case.assertAlmostEqual(val, 1.0 / 16, places=5)

    element_indices = wp.array([3, 3, 2], dtype=int, device=device)
    element_coords = wp.array(
        [
            [0.25, 0.5, 0.0],
            [0.5, 0.25, 0.0],
            [0.5, 0.5, 0.0],
        ],
        dtype=Coords,
        device=device,
    )

    pic_quadrature = fem.PicQuadrature(domain, positions=(element_indices, element_coords))

    test_case.assertIsNone(pic_quadrature.points_per_element())
    test_case.assertEqual(pic_quadrature.total_point_count(), 3)
    test_case.assertEqual(pic_quadrature.active_cell_count(), 2)

    val = fem.integrate(_piecewise_constant, quadrature=pic_quadrature)
    test_case.assertAlmostEqual(val, 1.25, places=5)


devices = get_test_devices()


class TestFem(unittest.TestCase):
    pass


add_function_test(TestFem, "test_regular_quadrature", test_regular_quadrature)
add_function_test(TestFem, "test_closest_point_queries", test_closest_point_queries)
add_function_test(TestFem, "test_grad_decomposition", test_grad_decomposition, devices=devices)
add_function_test(TestFem, "test_integrate_gradient", test_integrate_gradient, devices=devices)
add_function_test(TestFem, "test_interpolate_gradient", test_interpolate_gradient, devices=devices)
add_function_test(TestFem, "test_vector_divergence_theorem", test_vector_divergence_theorem, devices=devices)
add_function_test(TestFem, "test_tensor_divergence_theorem", test_tensor_divergence_theorem, devices=devices)
add_function_test(TestFem, "test_grid_2d", test_grid_2d, devices=devices)
add_function_test(TestFem, "test_triangle_mesh", test_triangle_mesh, devices=devices)
add_function_test(TestFem, "test_quad_mesh", test_quad_mesh, devices=devices)
add_function_test(TestFem, "test_grid_3d", test_grid_3d, devices=devices)
add_function_test(TestFem, "test_tet_mesh", test_tet_mesh, devices=devices)
add_function_test(TestFem, "test_hex_mesh", test_hex_mesh, devices=devices)
add_function_test(TestFem, "test_deformed_geometry", test_deformed_geometry, devices=devices)
add_function_test(TestFem, "test_dof_mapper", test_dof_mapper)
add_function_test(TestFem, "test_point_basis", test_point_basis)
add_function_test(TestFem, "test_particle_quadratures", test_particle_quadratures)


class TestFemShapeFunctions(unittest.TestCase):
    pass


add_function_test(TestFemShapeFunctions, "test_square_shape_functions", test_square_shape_functions)
add_function_test(TestFemShapeFunctions, "test_cube_shape_functions", test_cube_shape_functions)
add_function_test(TestFemShapeFunctions, "test_tri_shape_functions", test_tri_shape_functions)
add_function_test(TestFemShapeFunctions, "test_tet_shape_functions", test_tet_shape_functions)


if __name__ == "__main__":
    wp.build.clear_kernel_cache()
    unittest.main(verbosity=2)