warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

warp/tests/test_fem.py

@@ -1,31 +1,31 @@
-# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
+# 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.test_base import *
+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 *
-from warp.fem.geometry import Grid2D, Trimesh2D, Tetmesh
+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.space import make_polynomial_space, SymmetricTensorMapper
-from warp.fem.field import make_test
-from warp.fem.domain import Cells
-from warp.fem.integrate import integrate
-from warp.fem.operator import integrand
-from warp.fem.quadrature import RegularQuadrature
-from warp.fem.utils import unit_element
+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()
 
-wp.config.mode = "debug"
-wp.config.verify_cuda = True
-
 
 @integrand
 def linear_form(s: Sample, u: Field):
@@ -35,13 +35,13 @@ def linear_form(s: Sample, u: Field):
 def test_integrate_gradient(test_case, device):
     with wp.ScopedDevice(device):
         # Grid geometry
-        geo = Grid2D(res=vec2i(5))
+        geo = fem.Grid2D(res=wp.vec2i(5))
 
         # Domain and function spaces
-        domain = Cells(geometry=geo)
-        quadrature = RegularQuadrature(domain=domain, order=3)
+        domain = fem.Cells(geometry=geo)
+        quadrature = fem.RegularQuadrature(domain=domain, order=3)
 
-        scalar_space = make_polynomial_space(geo, degree=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)
@@ -52,37 +52,305 @@ def test_integrate_gradient(test_case, device):
 
         # forward pass
         with tape:
-            integrate(linear_form, quadrature=quadrature, fields={"u": u}, output=result)
+            fem.integrate(linear_form, quadrature=quadrature, fields={"u": u}, output=result)
 
         tape.backward(result)
 
-        test = make_test(space=scalar_space, domain=domain)
-        rhs = integrate(linear_form, quadrature=quadrature, fields={"u": test})
+        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)
 
-    cx, cy = np.meshgrid(np.arange(N, dtype=int), np.arange(N, dtype=int), indexing="ij")
-
-    vidx = np.transpose(
-        np.array(
-            [
-                (N + 1) * cx + cy,
-                (N + 1) * (cx + 1) + cy,
-                (N + 1) * (cx + 1) + (cy + 1),
-                (N + 1) * cx + cy,
-                (N + 1) * (cx + 1) + (cy + 1),
-                (N + 1) * (cx) + (cy + 1),
-            ]
-        )
-    ).reshape((-1, 3))
+    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)
 
@@ -94,82 +362,132 @@ def _gen_tetmesh(N):
 
     positions = np.transpose(np.meshgrid(x, y, z, indexing="ij")).reshape(-1, 3)
 
-    # Global node indices for each cell
-    cx, cy, cz = np.meshgrid(np.arange(N, dtype=int), np.arange(N, dtype=int), np.arange(N, dtype=int), indexing="ij")
-    grid_vidx = np.array(
-        [
-            (N + 1) ** 2 * cx + (N + 1) * cy + cz,
-            (N + 1) ** 2 * cx + (N + 1) * cy + cz + 1,
-            (N + 1) ** 2 * cx + (N + 1) * (cy + 1) + cz,
-            (N + 1) ** 2 * cx + (N + 1) * (cy + 1) + cz + 1,
-            (N + 1) ** 2 * (cx + 1) + (N + 1) * cy + cz,
-            (N + 1) ** 2 * (cx + 1) + (N + 1) * cy + cz + 1,
-            (N + 1) ** 2 * (cx + 1) + (N + 1) * (cy + 1) + cz,
-            (N + 1) ** 2 * (cx + 1) + (N + 1) * (cy + 1) + cz + 1,
-        ]
-    )
+    vidx = grid_to_tets(N, N, N)
 
-    # decompose grid cells into 5 tets
-    tet_vidx = np.array(
-        [
-            [0, 1, 2, 4],
-            [3, 2, 1, 7],
-            [5, 1, 7, 4],
-            [6, 7, 4, 2],
-            [4, 1, 2, 7],
-        ]
-    )
+    return wp.array(positions, dtype=wp.vec3), wp.array(vidx, dtype=int)
 
-    # Convert to 3d index coordinates
-    vidx_coords = np.array(
-        [
-            [0, 0, 0],
-            [0, 0, 1],
-            [0, 1, 0],
-            [0, 1, 1],
-            [1, 0, 0],
-            [1, 0, 1],
-            [1, 1, 0],
-            [1, 1, 1],
-        ]
-    )
-    tet_coords = vidx_coords[tet_vidx]
 
-    # Symmetry bits for each cell
-    ox, oy, oz = np.meshgrid(
-        np.arange(N, dtype=int) % 2, np.arange(N, dtype=int) % 2, np.arange(N, dtype=int) % 2, indexing="ij"
+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,
     )
-    tet_coords = np.broadcast_to(tet_coords, shape=(*ox.shape, *tet_coords.shape))
-
-    # Flip coordinates according to symmetry
-    ox_bk = np.broadcast_to(ox.reshape(*ox.shape, 1, 1), tet_coords.shape[:-1])
-    oy_bk = np.broadcast_to(oy.reshape(*ox.shape, 1, 1), tet_coords.shape[:-1])
-    oz_bk = np.broadcast_to(oz.reshape(*ox.shape, 1, 1), tet_coords.shape[:-1])
-    tet_coords_x = tet_coords[..., 0] ^ ox_bk
-    tet_coords_y = tet_coords[..., 1] ^ oy_bk
-    tet_coords_z = tet_coords[..., 2] ^ oz_bk
-
-    # Back to local vertex indices
-    corner_indices = 4 * tet_coords_x + 2 * tet_coords_y + tet_coords_z
-
-    # Now go from cell-local to global node indices
-    # There must be a nicer way than this, but for example purposes this works
-    corner_indices = corner_indices.reshape(-1, 4)
-    grid_vidx = grid_vidx.reshape((8, -1, 1))
-    grid_vidx = np.broadcast_to(grid_vidx, shape=(8, grid_vidx.shape[1], 5))
-    grid_vidx = grid_vidx.reshape((8, -1))
-
-    node_indices = np.arange(corner_indices.shape[0])
-    tet_grid_vidx = np.transpose(
-        [
-            grid_vidx[corner_indices[:, 0], node_indices],
-            grid_vidx[corner_indices[:, 1], node_indices],
-            grid_vidx[corner_indices[:, 2], node_indices],
-            grid_vidx[corner_indices[:, 3], node_indices],
-        ]
+
+    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 wp.array(positions, dtype=wp.vec3), wp.array(tet_grid_vidx, dtype=int)
+    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):
@@ -178,13 +496,54 @@ def test_triangle_mesh(test_case, device):
     with wp.ScopedDevice(device):
         positions, tri_vidx = _gen_trimesh(N)
 
-    geo = Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions)
+    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
@@ -192,12 +551,109 @@ def test_tet_mesh(test_case, device):
     with wp.ScopedDevice(device):
         positions, tet_vidx = _gen_tetmesh(N)
 
-    geo = Tetmesh(tet_vertex_indices=tet_vidx, positions=positions)
+    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
@@ -273,7 +729,14 @@ def test_closest_point_queries(test_case, device):
         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]])
+    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)
@@ -281,8 +744,8 @@ def test_closest_point_queries(test_case, device):
         _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.e-4)
-    assert_np_equal(sq_dist.numpy(), expected_sq_dist, tol = 1.e-4)
+    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):
@@ -319,9 +782,11 @@ def test_regular_quadrature(test_case, device):
 
 def test_dof_mapper(test_case, device):
     matrix_types = [wp.mat22, wp.mat33]
-    for mapping in SymmetricTensorMapper.Mapping:
+
+    # Symmetric mapper
+    for mapping in fem.SymmetricTensorMapper.Mapping:
         for dtype in matrix_types:
-            mapper = SymmetricTensorMapper(dtype, mapping=mapping)
+            mapper = fem.SymmetricTensorMapper(dtype, mapping=mapping)
             dof_dtype = mapper.dof_dtype
 
             for k in range(dof_dtype._length_):
@@ -339,23 +804,468 @@ def test_dof_mapper(test_case, device):
                 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)
 
-def register(parent):
-    devices = get_test_devices()
 
-    class TestFem(parent):
-        pass
+class TestFemShapeFunctions(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_integrate_gradient", test_integrate_gradient, devices=devices)
-    add_function_test(TestFem, "test_triangle_mesh", test_triangle_mesh, devices=devices)
-    add_function_test(TestFem, "test_tet_mesh", test_tet_mesh, devices=devices)
-    add_function_test(TestFem, "test_dof_mapper", test_dof_mapper)
 
-    return TestFem
+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__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)