# 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)