warp-lang 1.7.2__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__py3-none-manylinux_2_34_aarch64.whl

warp/tests/test_fem.py

@@ -476,11 +476,8 @@ def _test_geo_cells(
     wp.atomic_add(cell_measures, s.element_index, fem.measure(domain, s) * s.qp_weight)
 
 
-@fem.integrand(kernel_options={"enable_backward": False})
-def _test_cell_lookup(
-    s: fem.Sample,
-    domain: fem.Domain,
-):
+@fem.integrand(kernel_options={"enable_backward": False, "max_unroll": 2})
+def _test_cell_lookup(s: fem.Sample, domain: fem.Domain, cell_filter: wp.array(dtype=int)):
     pos = domain(s)
 
     s_guess = fem.lookup(domain, pos, s)
@@ -491,6 +488,23 @@ def _test_cell_lookup(
     wp.expect_eq(s_noguess.element_index, s.element_index)
     wp.expect_near(domain(s_noguess), pos, 0.001)
 
+    # Filtered lookup
+    max_dist = 10.0
+    filter_target = 1
+    s_filter = fem.lookup(domain, pos, max_dist, cell_filter, filter_target)
+    wp.expect_eq(s_filter.element_index, 0)
+
+    if s.element_index != 0:
+        # test closest point optimality
+        pos_f = domain(s_filter)
+        pos_f += 0.1 * (pos - pos_f)
+        coord_proj, _sq_dist = fem.element_closest_point(domain, s_filter.element_index, pos_f)
+        wp.expect_near(coord_proj, s_filter.element_coords, 0.001)
+
+    # test that extrapolated coordinates yield bak correct position
+    s_filter.element_coords = fem.element_coordinates(domain, s_filter.element_index, pos)
+    wp.expect_near(domain(s_filter), pos, 0.001)
+
 
 @fem.integrand(kernel_options={"enable_backward": False, "max_unroll": 1})
 def _test_geo_sides(
@@ -520,6 +534,7 @@ def _test_geo_sides(
 
     wp.expect_near(coords, inner_side_s.element_coords, 0.0001)
     wp.expect_near(coords, outer_side_s.element_coords, 0.0001)
+    wp.expect_near(coords, fem.element_coordinates(domain, side_index, pos_side), 0.001)
 
     area = fem.measure(domain, s)
     wp.atomic_add(side_measures, side_index, area * s.qp_weight)
@@ -544,7 +559,7 @@ def _test_side_normals(
         wp.expect_near(F_cross[k], nor[k], 0.0001)
 
 
-def _launch_test_geometry_kernel(geo: fem.Geometry, device, test_cell_lookup: bool = True):
+def _launch_test_geometry_kernel(geo: fem.Geometry, device):
     cell_measures = wp.zeros(dtype=float, device=device, shape=geo.cell_count())
     cell_quadrature = fem.RegularQuadrature(fem.Cells(geo), order=2)
 
@@ -557,11 +572,11 @@ def _launch_test_geometry_kernel(geo: fem.Geometry, device, test_cell_lookup: bo
             quadrature=cell_quadrature,
             values={"cell_measures": cell_measures},
         )
-        if test_cell_lookup:
-            fem.interpolate(
-                _test_cell_lookup,
-                quadrature=cell_quadrature,
-            )
+
+        cell_filter = np.zeros(geo.cell_count(), dtype=int)
+        cell_filter[0] = 1
+        cell_filter = wp.array(cell_filter, dtype=int)
+        fem.interpolate(_test_cell_lookup, quadrature=cell_quadrature, values={"cell_filter": cell_filter})
 
         fem.interpolate(
             _test_geo_sides,
@@ -637,14 +652,14 @@ def test_quad_mesh(test, device):
     with wp.ScopedDevice(device):
         positions, quad_vidx = _gen_quadmesh(N)
 
-    geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions)
+    geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions, build_bvh=True)
 
     test.assertEqual(geo.cell_count(), N**2)
     test.assertEqual(geo.vertex_count(), (N + 1) ** 2)
     test.assertEqual(geo.side_count(), 2 * (N + 1) * N)
     test.assertEqual(geo.boundary_side_count(), 4 * N)
 
-    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device, test_cell_lookup=False)
+    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)
@@ -655,14 +670,14 @@ def test_quad_mesh(test, device):
     positions = np.hstack((positions, np.ones((positions.shape[0], 1))))
     positions = wp.array(positions, device=device, dtype=wp.vec3)
 
-    geo = fem.Quadmesh3D(quad_vertex_indices=quad_vidx, positions=positions)
+    geo = fem.Quadmesh3D(quad_vertex_indices=quad_vidx, positions=positions, build_bvh=True)
 
     test.assertEqual(geo.cell_count(), N**2)
     test.assertEqual(geo.vertex_count(), (N + 1) ** 2)
     test.assertEqual(geo.side_count(), 2 * (N + 1) * N)
     test.assertEqual(geo.boundary_side_count(), 4 * N)
 
-    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device, test_cell_lookup=False)
+    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)
@@ -711,7 +726,7 @@ def test_hex_mesh(test, device):
     with wp.ScopedDevice(device):
         positions, tet_vidx = _gen_hexmesh(N)
 
-    geo = fem.Hexmesh(hex_vertex_indices=tet_vidx, positions=positions)
+    geo = fem.Hexmesh(hex_vertex_indices=tet_vidx, positions=positions, build_bvh=True)
 
     test.assertEqual(geo.cell_count(), (N) ** 3)
     test.assertEqual(geo.vertex_count(), (N + 1) ** 3)
@@ -719,7 +734,7 @@ def test_hex_mesh(test, device):
     test.assertEqual(geo.boundary_side_count(), 6 * N * N)
     test.assertEqual(geo.edge_count(), 3 * N * (N + 1) ** 2)
 
-    side_measures, cell_measures = _launch_test_geometry_kernel(geo, device, test_cell_lookup=False)
+    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)
@@ -844,7 +859,8 @@ def test_deformed_geometry(test, device):
         test.assertEqual(geo.side_count(), 6 * (N + 1) * N**2 + (N**3) * 4)
         test.assertEqual(geo.boundary_side_count(), 12 * N * N)
 
-        side_measures, cell_measures = _launch_test_geometry_kernel(deformed_geo, device, test_cell_lookup=False)
+        deformed_geo.build_bvh()
+        side_measures, cell_measures = _launch_test_geometry_kernel(deformed_geo, device)
 
         test.assertAlmostEqual(
             np.sum(cell_measures.numpy()), scale**3, places=4, msg=f"cell_measures = {cell_measures.numpy()}"
@@ -1497,7 +1513,7 @@ def test_point_basis(test, device):
     point_test = fem.make_test(point_space, domain=domain)
 
     # Sample at particle positions
-    ones = fem.integrate(linear_form, fields={"u": point_test}, nodal=True)
+    ones = fem.integrate(linear_form, fields={"u": point_test}, assembly="nodal")
     test.assertAlmostEqual(np.sum(ones.numpy()), 1.0, places=5)
 
     # Sampling outside of particle positions
@@ -1613,12 +1629,23 @@ def test_particle_quadratures(test, device):
     assert_np_equal(measures.grad.numpy(), np.full(3, 4.0))  # == 1.0 / cell_area
 
 
-@fem.integrand
-def _value_at_node(s: fem.Sample, f: fem.Field, values: wp.array(dtype=float)):
+@fem.integrand(kernel_options={"enable_backward": False})
+def _value_at_node(domain: fem.Domain, s: fem.Sample, f: fem.Field, values: wp.array(dtype=float)):
+    # lookup at node is ambiguous, check that partition_lookup retains sample on current partition
+    s_partition = fem.partition_lookup(domain, domain(s))
+    wp.expect_eq(s.element_index, s_partition.element_index)
+    wp.expect_neq(fem.operator.element_partition_index(domain, s.element_index), fem.NULL_ELEMENT_INDEX)
+
     node_index = fem.operator.node_partition_index(f, s.qp_index)
     return values[node_index]
 
 
+@fem.integrand(kernel_options={"enable_backward": False})
+def _test_node_index(s: fem.Sample, u: fem.Field):
+    wp.expect_eq(fem.node_index(u, s), s.qp_index)
+    return 0.0
+
+
 def test_nodal_quadrature(test, device):
     geo = fem.Grid2D(res=wp.vec2i(2))
 
@@ -1635,15 +1662,18 @@ def test_nodal_quadrature(test, device):
 
     # test accessing data associated to a given node
 
-    piecewise_constant_space = fem.make_polynomial_space(geo, degree=0)
-    geo_partition = fem.LinearGeometryPartition(geo, 3, 4)
-    space_partition = fem.make_space_partition(piecewise_constant_space, geo_partition)
+    piecewise_constant_space = fem.make_polynomial_space(geo, degree=1)
+    geo_partition = fem.LinearGeometryPartition(geo, 2, 4)
+    assert geo_partition.cell_count() == 1
+
+    space_partition = fem.make_space_partition(piecewise_constant_space, geo_partition, with_halo=False)
+
     field = fem.make_discrete_field(piecewise_constant_space, space_partition=space_partition)
 
     partition_domain = fem.Cells(geo_partition)
     partition_nodal_quadrature = fem.NodalQuadrature(partition_domain, piecewise_constant_space)
 
-    partition_node_values = wp.array([5.0], dtype=float)
+    partition_node_values = wp.full(value=5.0, shape=space_partition.node_count(), dtype=float)
     val = fem.integrate(
         _value_at_node,
         quadrature=partition_nodal_quadrature,
@@ -1652,6 +1682,9 @@ def test_nodal_quadrature(test, device):
     )
     test.assertAlmostEqual(val, 5.0 / geo.cell_count(), places=5)
 
+    u_test = fem.make_test(space)
+    fem.integrate(_test_node_index, assembly="nodal", fields={"u": u_test})
+
 
 @wp.func
 def aniso_bicubic_fn(x: wp.vec2, scale: wp.vec2):

warp/tests/test_func.py

@@ -259,7 +259,25 @@ def test_return_annotation_none() -> None:
     user_func_return_none()
 
 
-devices = get_test_devices()
+@wp.func
+def divide_by_zero(x: float):
+    return x / 0.0
+
+
+@wp.func
+def normalize_vector(vec_a: wp.vec3):
+    return wp.normalize(vec_a)
+
+
+# This pair is to test the situation where one overload throws an error, but a second one works.
+@wp.func
+def divide_by_zero_overload(x: wp.float32):
+    return x / 0
+
+
+@wp.func
+def divide_by_zero_overload(x: wp.float64):
+    return wp.div(x, 0.0)
 
 
 class TestFunc(unittest.TestCase):
@@ -425,6 +443,29 @@ class TestFunc(unittest.TestCase):
         ):
             a * b
 
+    def test_cpython_call_user_function_with_error(self):
+        # Actually the following also includes a ZeroDivisionError in the message due to exception chaining,
+        # but I don't know how to test for that.
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "Error calling function 'divide_by_zero'. No version succeeded. "
+            "See above for the error from the last version that was tried.",
+        ):
+            divide_by_zero(1.0)
+
+    def test_cpython_call_user_function_with_overloads(self):
+        self.assertEqual(divide_by_zero_overload(1.0), math.inf)
+
+    def test_cpython_call_user_function_with_wrong_argument_types(self):
+        with self.assertRaisesRegex(
+            RuntimeError,
+            "Error calling function 'normalize_vector'. No version succeeded. "
+            "See above for the error from the last version that was tried.",
+        ):
+            normalize_vector(1.0)
+
+
+devices = get_test_devices()
 
 add_kernel_test(TestFunc, kernel=test_overload_func, name="test_overload_func", dim=1, devices=devices)
 add_function_test(TestFunc, func=test_return_func, name="test_return_func", devices=devices)

warp/tests/test_grad.py

@@ -518,7 +518,7 @@ def test_mesh_grad(test, device):
             pos_np[i, j] += eps
             fd_grad[i, j] = (f1 - f2) / (2 * eps)
 
-    assert np.allclose(ad_grad, fd_grad, atol=1e-3)
+    np.testing.assert_allclose(ad_grad, fd_grad, atol=1e-3)
 
 
 @wp.func

warp/tests/test_lerp.py

@@ -119,9 +119,7 @@ def test_lerp(test, device):
                 a = wp.array([test_data.a], dtype=data_type, device=device, requires_grad=True)
                 b = wp.array([test_data.b], dtype=data_type, device=device, requires_grad=True)
                 t = wp.array([test_data.t], dtype=float, device=device, requires_grad=True)
-                out = wp.array(
-                    [0] * wp.types.type_length(data_type), dtype=data_type, device=device, requires_grad=True
-                )
+                out = wp.array([0] * wp.types.type_size(data_type), dtype=data_type, device=device, requires_grad=True)
 
                 with wp.Tape() as tape:
                     wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)