warp-lang 1.1.0__py3-none-macosx_10_13_universal2.whl → 1.2.0__py3-none-macosx_10_13_universal2.whl

warp/tests/test_volume.py

@@ -6,14 +6,13 @@
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 import unittest
+from typing import Any
 
 import numpy as np
 
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-
 
 # float volume tests
 @wp.kernel
@@ -30,6 +29,7 @@ def test_volume_lookup_f(volume: wp.uint64, points: wp.array(dtype=wp.vec3)):
     k = int(p[2])
 
     expect_eq(wp.volume_lookup_f(volume, i, j, k), expected)
+    expect_eq(wp.volume_lookup(volume, i, j, k, dtype=wp.float32), expected)
 
 
 @wp.kernel
@@ -45,6 +45,7 @@ def test_volume_sample_closest_f(volume: wp.uint64, points: wp.array(dtype=wp.ve
         expected = 10.0
 
     expect_eq(wp.volume_sample_f(volume, p, wp.Volume.CLOSEST), expected)
+    expect_eq(wp.volume_sample(volume, p, wp.Volume.CLOSEST, dtype=wp.float32), expected)
 
     q = wp.volume_index_to_world(volume, p)
     q_inv = wp.volume_world_to_index(volume, q)
@@ -62,6 +63,7 @@ def test_volume_sample_linear_f(volume: wp.uint64, points: wp.array(dtype=wp.vec
         return  # not testing against background values
 
     expect_near(wp.volume_sample_f(volume, p, wp.Volume.LINEAR), expected, 2.0e-4)
+    expect_near(wp.volume_sample(volume, p, wp.Volume.LINEAR, dtype=wp.float32), expected, 2.0e-4)
 
 
 @wp.kernel
@@ -86,6 +88,13 @@ def test_volume_sample_grad_linear_f(volume: wp.uint64, points: wp.array(dtype=w
     expect_near(grad[1], expected_gy, 2.0e-4)
     expect_near(grad[2], expected_gz, 2.0e-4)
 
+    val = wp.volume_sample_grad(volume, p, wp.Volume.LINEAR, grad, dtype=wp.float32)
+
+    expect_near(val, expected_val, 2.0e-4)
+    expect_near(grad[0], expected_gx, 2.0e-4)
+    expect_near(grad[1], expected_gy, 2.0e-4)
+    expect_near(grad[2], expected_gz, 2.0e-4)
+
 
 @wp.kernel
 def test_volume_sample_local_f_linear_values(
@@ -162,6 +171,7 @@ def test_volume_lookup_v(volume: wp.uint64, points: wp.array(dtype=wp.vec3)):
     k = int(p[2])
 
     expect_eq(wp.volume_lookup_v(volume, i, j, k), expected)
+    expect_eq(wp.volume_lookup(volume, i, j, k, dtype=wp.vec3), expected)
 
 
 @wp.kernel
@@ -177,6 +187,7 @@ def test_volume_sample_closest_v(volume: wp.uint64, points: wp.array(dtype=wp.ve
         expected = wp.vec3(10.8, -4.13, 10.26)
 
     expect_eq(wp.volume_sample_v(volume, p, wp.Volume.CLOSEST), expected)
+    expect_eq(wp.volume_sample(volume, p, wp.Volume.CLOSEST, dtype=wp.vec3), expected)
 
     q = wp.volume_index_to_world(volume, p)
     q_inv = wp.volume_world_to_index(volume, q)
@@ -196,6 +207,30 @@ def test_volume_sample_linear_v(volume: wp.uint64, points: wp.array(dtype=wp.vec
         return  # not testing against background values
 
     expect_near(wp.volume_sample_v(volume, p, wp.Volume.LINEAR), expected, 2.0e-4)
+    expect_near(wp.volume_sample(volume, p, wp.Volume.LINEAR, dtype=wp.vec3), expected, 2.0e-4)
+
+
+@wp.kernel
+def test_volume_sample_grad_linear_v(volume: wp.uint64, points: wp.array(dtype=wp.vec3)):
+    tid = wp.tid()
+
+    p = points[tid]
+
+    if abs(p[0]) > 10.0 or abs(p[1]) > 10.0 or abs(p[2]) > 10.0:
+        return  # not testing against background values
+
+    expected_val = wp.vec3(
+        p[0] + 2.0 * p[1] + 3.0 * p[2], 4.0 * p[0] + 5.0 * p[1] + 6.0 * p[2], 7.0 * p[0] + 8.0 * p[1] + 9.0 * p[2]
+    )
+    expected_grad = wp.mat33(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0)
+
+    grad = wp.mat33(0.0)
+    val = wp.volume_sample_grad(volume, p, wp.Volume.LINEAR, grad, dtype=wp.vec3)
+
+    expect_near(val, expected_val, 2.0e-4)
+    expect_near(grad[0], expected_grad[0], 2.0e-4)
+    expect_near(grad[1], expected_grad[1], 2.0e-4)
+    expect_near(grad[2], expected_grad[2], 2.0e-4)
 
 
 @wp.kernel
@@ -233,6 +268,7 @@ def test_volume_lookup_i(volume: wp.uint64, points: wp.array(dtype=wp.vec3)):
         expected = 10
 
     expect_eq(wp.volume_lookup_i(volume, i, j, k), expected)
+    expect_eq(wp.volume_lookup(volume, i, j, k, dtype=wp.int32), expected)
 
 
 @wp.kernel
@@ -248,6 +284,7 @@ def test_volume_sample_i(volume: wp.uint64, points: wp.array(dtype=wp.vec3)):
         expected = 10
 
     expect_eq(wp.volume_sample_i(volume, p), expected)
+    expect_eq(wp.volume_sample(volume, p, wp.Volume.CLOSEST, dtype=wp.int32), expected)
 
     q = wp.volume_index_to_world(volume, p)
     q_inv = wp.volume_world_to_index(volume, q)
@@ -293,9 +330,7 @@ def test_volume_store_f(volume: wp.uint64, points: wp.array(dtype=wp.vec3), valu
     j = int(p[1])
     k = int(p[2])
 
-    # NB: Writing outside the allocated domain overwrites the background value of the Volume
-    if abs(i) <= 11 and abs(j) <= 11 and abs(k) <= 11:
-        wp.volume_store_f(volume, i, j, k, float(i + 100 * j + 10000 * k))
+    wp.volume_store(volume, i, j, k, float(i + 100 * j + 10000 * k))
     values[tid] = wp.volume_lookup_f(volume, i, j, k)
 
 
@@ -308,9 +343,7 @@ def test_volume_store_v(volume: wp.uint64, points: wp.array(dtype=wp.vec3), valu
     j = int(p[1])
     k = int(p[2])
 
-    # NB: Writing outside the allocated domain overwrites the background value of the Volume
-    if abs(i) <= 11 and abs(j) <= 11 and abs(k) <= 11:
-        wp.volume_store_v(volume, i, j, k, p)
+    wp.volume_store(volume, i, j, k, p)
     values[tid] = wp.volume_lookup_v(volume, i, j, k)
 
 
@@ -323,9 +356,7 @@ def test_volume_store_i(volume: wp.uint64, points: wp.array(dtype=wp.vec3), valu
     j = int(p[1])
     k = int(p[2])
 
-    # NB: Writing outside the allocated domain overwrites the background value of the Volume
-    if abs(i) <= 11 and abs(j) <= 11 and abs(k) <= 11:
-        wp.volume_store_i(volume, i, j, k, i + 100 * j + 10000 * k)
+    wp.volume_store(volume, i, j, k, i + 100 * j + 10000 * k)
     values[tid] = wp.volume_lookup_i(volume, i, j, k)
 
 
@@ -355,6 +386,7 @@ volume_paths = {
     "float": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/test_grid.nvdb")),
     "int32": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/test_int32_grid.nvdb")),
     "vec3f": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/test_vec_grid.nvdb")),
+    "index": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/test_index_grid.nvdb")),
     "torus": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/torus.nvdb")),
     "float_write": os.path.abspath(os.path.join(os.path.dirname(__file__), "assets/test_grid.nvdb")),
 }
@@ -492,7 +524,7 @@ def test_volume_sample_linear_v_gradient(test, device):
             )
         tape.backward(values)
 
-        grad_expected = np.array([6.0, 15.0, 24.0])
+        grad_expected = np.array([12.0, 15.0, 18.0])
         grad_computed = tape.gradients[uvws].numpy()[0]
         np.testing.assert_allclose(grad_computed, grad_expected, rtol=1e-4)
 
@@ -506,7 +538,7 @@ def test_volume_sample_linear_v_gradient(test, device):
             )
         tape.backward(values)
 
-        grad_expected = np.array([6.0, 15.0, 24.0]) / 0.25
+        grad_expected = np.array([12.0, 15.0, 18.0]) / 0.25
         grad_computed = tape.gradients[xyzs].numpy()[0]
         np.testing.assert_allclose(grad_computed, grad_expected, rtol=1e-4)
 
@@ -604,6 +636,188 @@ def test_volume_introspection(test, device):
             np.testing.assert_equal(test_volume_tiles, tiles_sorted)
             np.testing.assert_equal([0.25] * 3, voxel_size)
 
+            voxel_count = volume.get_voxel_count()
+            voxels_actual = volume.get_voxels().numpy()
+            assert voxel_count == voxels_actual.shape[0]
+
+            # Voxel coordinates should be unique
+            voxels_unique = np.unique(voxels_actual, axis=0)
+            assert voxel_count == voxels_unique.shape[0]
+
+            # Get back tiles from voxels, should match get_tiles()
+            voxel_tiles = 8 * (voxels_unique // 8)
+            voxel_tiles_sorted = voxel_tiles[np.lexsort(voxel_tiles.T[::-1])]
+            voxel_tiles_unique = np.unique(voxel_tiles_sorted, axis=0)
+
+            np.testing.assert_equal(voxel_tiles_unique, tiles_sorted)
+
+
+def test_volume_multiple_grids(test, device):
+    volume = volumes["index"][device.alias]
+
+    volume_2 = volume.load_next_grid()
+
+    test.assertIsNotNone(volume_2)
+
+    test.assertNotEqual(volume.id, volume_2.id)
+    test.assertNotEqual(volume.get_voxel_count(), volume_2.get_voxel_count())
+
+    test.assertEqual(volume.get_grid_info().grid_count, volume_2.get_grid_info().grid_count)
+    test.assertEqual(volume.get_grid_info().grid_index + 1, volume_2.get_grid_info().grid_index)
+
+    volume_3 = volume_2.load_next_grid()
+    test.assertIsNone(volume_3)
+
+
+def test_volume_feature_array(test, device):
+    volume = volumes["index"][device.alias]
+
+    test.assertEqual(volume.get_feature_array_count(), 1)
+
+    array = volume.feature_array(0, dtype=wp.uint64)
+    test.assertEqual(array.device, device)
+    test.assertEqual(array.dtype, wp.uint64)
+
+    # fVDB convention, data starts with array ndim + shape
+    np.testing.assert_equal(array.numpy()[0:4], [3, volume.get_voxel_count(), 2, 3])
+
+
+@wp.kernel
+def fill_leaf_values_kernel(volume: wp.uint64, ijk: wp.array2d(dtype=wp.int32), values: wp.array(dtype=Any)):
+    tid = wp.tid()
+
+    i = ijk[tid, 0]
+    j = ijk[tid, 1]
+    k = ijk[tid, 2]
+
+    expect_eq(tid, wp.volume_lookup_index(volume, i, j, k))
+
+    values[tid] = wp.volume_lookup(volume, i, j, k, dtype=values.dtype)
+
+
+@wp.kernel
+def test_volume_sample_index_kernel(
+    volume: wp.uint64,
+    points: wp.array(dtype=wp.vec3),
+    values: wp.array(dtype=Any),
+    background: wp.array(dtype=Any),
+    sampled_values: wp.array(dtype=Any),
+):
+    tid = wp.tid()
+    p = points[tid]
+
+    ref = wp.volume_sample(volume, p, wp.Volume.LINEAR, dtype=values.dtype)
+    sampled_values[tid] = wp.volume_sample_index(volume, p, wp.Volume.LINEAR, values, background[0])
+    expect_eq(sampled_values[tid], ref)
+
+
+@wp.kernel
+def test_volume_sample_grad_index_kernel(
+    volume: wp.uint64,
+    points: wp.array(dtype=wp.vec3),
+    values: wp.array(dtype=Any),
+    background: wp.array(dtype=Any),
+    sampled_values: wp.array(dtype=Any),
+    sampled_grads: wp.array(dtype=Any),
+):
+    tid = wp.tid()
+    p = points[tid]
+
+    ref_grad = sampled_grads.dtype()
+    ref = wp.volume_sample_grad(volume, p, wp.Volume.LINEAR, ref_grad, dtype=values.dtype)
+
+    grad = type(ref_grad)()
+    sampled_values[tid] = wp.volume_sample_grad_index(volume, p, wp.Volume.LINEAR, values, background[0], grad)
+    expect_eq(sampled_values[tid], ref)
+
+    expect_eq(grad[0], ref_grad[0])
+    expect_eq(grad[1], ref_grad[1])
+    expect_eq(grad[2], ref_grad[2])
+    sampled_grads[tid] = grad
+
+
+def test_volume_sample_index(test, device):
+    points = rng.uniform(-10.0, 10.0, size=(100, 3))
+    points[0:10, 0] += 100.0  # ensure some points are over unallocated voxels
+    uvws = wp.array(points, dtype=wp.vec3, device=device)
+
+    bg_values = {
+        "float": 10.0,
+        "vec3f": wp.vec3(10.8, -4.13, 10.26),
+    }
+    grad_types = {
+        "float": wp.vec3,
+        "vec3f": wp.mat33,
+    }
+
+    for volume_names in ("float", "vec3f"):
+        with test.subTest(volume_names=volume_names):
+            volume = volumes[volume_names][device.alias]
+
+            ijk = volume.get_voxels()
+
+            values = wp.empty(shape=volume.get_voxel_count(), dtype=volume.dtype, device=device, requires_grad=True)
+
+            vid = wp.uint64(volume.id)
+            wp.launch(fill_leaf_values_kernel, dim=values.shape, inputs=[vid, ijk, values], device=device)
+
+            sampled_values = wp.empty(shape=points.shape[0], dtype=volume.dtype, device=device, requires_grad=True)
+            background = wp.array([bg_values[volume_names]], dtype=volume.dtype, device=device, requires_grad=True)
+
+            tape = wp.Tape()
+            with tape:
+                wp.launch(
+                    test_volume_sample_index_kernel,
+                    dim=points.shape[0],
+                    inputs=[vid, uvws, values, background, sampled_values],
+                    device=device,
+                )
+
+            sampled_values.grad.fill_(1.0)
+            tape.backward()
+
+            # test adjoint w.r.t voxel and background value arrays
+            # we should have sum(sampled_values) = sum(adj_values * values) + (adj_background * background)
+            sum_sampled_values = np.sum(sampled_values.numpy(), axis=0)
+            sum_values_adj = np.sum(values.numpy() * values.grad.numpy(), axis=0)
+            sum_background_adj = background.numpy()[0] * background.grad.numpy()[0]
+
+            np.testing.assert_allclose(sum_sampled_values, sum_values_adj + sum_background_adj, rtol=1.0e-3)
+
+            tape.reset()
+
+            sampled_grads = wp.empty(
+                shape=points.shape[0], dtype=grad_types[volume_names], device=device, requires_grad=True
+            )
+
+            with tape:
+                wp.launch(
+                    test_volume_sample_grad_index_kernel,
+                    dim=points.shape[0],
+                    inputs=[vid, uvws, values, background, sampled_values, sampled_grads],
+                    device=device,
+                )
+
+            sampled_values.grad.fill_(1.0)
+            tape.backward()
+
+            # we should have sum(sampled_values) = sum(adj_values * values) + (adj_background * background)
+            sum_sampled_values = np.sum(sampled_values.numpy(), axis=0)
+            sum_values_adj = np.sum(values.numpy() * values.grad.numpy(), axis=0)
+            sum_background_adj = background.numpy()[0] * background.grad.numpy()[0]
+            np.testing.assert_allclose(sum_sampled_values, sum_values_adj + sum_background_adj, rtol=1.0e-3)
+
+            tape.zero()
+            sampled_values.grad.fill_(0.0)
+            sampled_grads.grad.fill_(1.0)
+            tape.backward()
+
+            # we should have sum(sampled_grad, axes=(0, -1)) = sum(adj_values * values) + (adj_background * background)
+            sum_sampled_grads = np.sum(np.sum(sampled_grads.numpy(), axis=0), axis=-1)
+            sum_values_adj = np.sum(values.numpy() * values.grad.numpy(), axis=0)
+            sum_background_adj = background.numpy()[0] * background.grad.numpy()[0]
+            np.testing.assert_allclose(sum_sampled_grads, sum_values_adj + sum_background_adj, rtol=1.0e-3)
+
 
 def test_volume_from_numpy(test, device):
     # Volume.allocate_from_tiles() is only available with CUDA
@@ -657,6 +871,9 @@ add_function_test(TestVolume, "test_volume_introspection", test_volume_introspec
 add_function_test(
     TestVolume, "test_volume_from_numpy", test_volume_from_numpy, devices=get_selected_cuda_test_devices()
 )
+add_function_test(TestVolume, "test_volume_multiple_grids", test_volume_multiple_grids, devices=devices)
+add_function_test(TestVolume, "test_volume_feature_array", test_volume_feature_array, devices=devices)
+add_function_test(TestVolume, "test_volume_sample_index", test_volume_sample_index, devices=devices)
 
 points = {}
 points_jittered = {}
@@ -715,6 +932,13 @@ for device in devices:
         inputs=[volumes["vec3f"][device.alias].id, points_jittered[device.alias]],
         devices=[device.alias],
     )
+    add_kernel_test(
+        TestVolume,
+        test_volume_sample_grad_linear_v,
+        dim=len(point_grid),
+        inputs=[volumes["vec3f"][device.alias].id, points_jittered[device.alias]],
+        devices=[device.alias],
+    )
 
     add_kernel_test(
         TestVolume,

warp/tests/test_volume_write.py

@@ -12,9 +12,6 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
-wp.init()
-# wp.config.cache_kernels = False
-
 
 # Volume write tests
 @wp.kernel
@@ -166,6 +163,10 @@ def test_volume_allocation(test, device):
         points_in_world_space=True,
         device=device,
     )
+
+    assert wp.types.types_equal(volume_a.dtype, wp.float32)
+    assert wp.types.types_equal(volume_b.dtype, wp.float32)
+
     points = wp.array(points_ref, dtype=wp.vec3, device=device)
     values_a = wp.empty(num_points, dtype=wp.float32, device=device)
     values_b = wp.empty(num_points, dtype=wp.float32, device=device)
@@ -201,6 +202,10 @@ def test_volume_allocate_by_tiles_f(test, device):
     points_ws_d = wp.array(points_ws, dtype=wp.vec3, device=device)
     volume_a = wp.Volume.allocate_by_tiles(points_is_d, voxel_size, background_value, translation, device=device)
     volume_b = wp.Volume.allocate_by_tiles(points_ws_d, voxel_size, background_value, translation, device=device)
+
+    assert wp.types.types_equal(volume_a.dtype, wp.float32)
+    assert wp.types.types_equal(volume_b.dtype, wp.float32)
+
     values_a = wp.empty(num_tiles * 512, dtype=wp.float32, device=device)
     values_b = wp.empty(num_tiles * 512, dtype=wp.float32, device=device)
 
@@ -229,6 +234,9 @@ def test_volume_allocate_by_tiles_v(test, device):
 
     points_d = wp.array(points_is, dtype=wp.int32, device=device)
     volume = wp.Volume.allocate_by_tiles(points_d, 0.1, wp.vec3(1, 2, 3), device=device)
+
+    assert wp.types.types_equal(volume.dtype, wp.vec3)
+
     values = wp.empty(len(points_d) * 512, dtype=wp.vec3, device=device)
 
     wp.launch(test_volume_tile_store_v, dim=len(points_d), inputs=[volume.id, points_d], device=device)
@@ -238,6 +246,72 @@ def test_volume_allocate_by_tiles_v(test, device):
     np.testing.assert_equal(values_res, values_ref)
 
 
+def test_volume_allocate_by_tiles_index(test, device):
+    num_tiles = 10
+    rng = np.random.default_rng(101215)
+    tiles = rng.integers(-512, 512, size=(num_tiles, 3), dtype=np.int32)
+    points_is = tiles * 8
+
+    points_d = wp.array(points_is, dtype=wp.int32, device=device)
+    volume = wp.Volume.allocate_by_tiles(points_d, 0.1, bg_value=None, device=device)
+
+    assert volume.is_index
+
+    vol_tiles = volume.get_tiles().numpy() / 8
+    vol_tile_sorted = vol_tiles[np.lexsort(vol_tiles.T[::-1])]
+    vol_tile_unique = np.unique(vol_tile_sorted, axis=0)
+
+    tile_sorted = tiles[np.lexsort(tiles.T[::-1])]
+    tile_unique = np.unique(tile_sorted, axis=0)
+
+    np.testing.assert_equal(tile_unique, vol_tile_unique)
+
+
+def test_volume_allocation_from_voxels(test, device):
+    point_count = 387
+    rng = np.random.default_rng(101215)
+
+    # Create from world-space points
+    points = wp.array(rng.uniform(5.0, 10.0, size=(point_count, 3)), dtype=float, device=device)
+
+    volume = wp.Volume.allocate_by_voxels(
+        voxel_points=points, voxel_size=0.25, translation=(0.0, 5.0, 10.0), device=device
+    )
+
+    assert volume.is_index
+
+    test.assertNotEqual(volume.id, 0)
+
+    test.assertAlmostEqual(volume.get_voxel_size(), (0.25, 0.25, 0.25))
+    voxel_count = volume.get_voxel_count()
+    test.assertGreaterEqual(point_count, voxel_count)
+    test.assertGreaterEqual(voxel_count, 1)
+
+    voxels = volume.get_voxels()
+
+    # Check that world-to-index transform has been correctly applied
+    voxel_low = np.min(voxels.numpy(), axis=0)
+    voxel_up = np.max(voxels.numpy(), axis=0)
+    np.testing.assert_array_less([19, -1, -21], voxel_low)
+    np.testing.assert_array_less(voxel_up, [41, 21, 1])
+
+    # Recreate the volume from ijk coords
+    volume_from_ijk = wp.Volume.allocate_by_voxels(
+        voxel_points=voxels, voxel_size=0.25, translation=(0.0, 5.0, 10.0), device=device
+    )
+
+    assert volume_from_ijk.is_index
+
+    assert volume_from_ijk.get_voxel_count() == voxel_count
+    ijk_voxels = volume_from_ijk.get_voxels().numpy()
+
+    voxels = voxels.numpy()
+    voxel_sorted = voxels[np.lexsort(voxels.T[::-1])]
+    ijk_voxel_sorted = ijk_voxels[np.lexsort(ijk_voxels.T[::-1])]
+
+    np.testing.assert_equal(voxel_sorted, ijk_voxel_sorted)
+
+
 devices = get_selected_cuda_test_devices()
 
 
@@ -248,6 +322,15 @@ class TestVolumeWrite(unittest.TestCase):
 add_function_test(TestVolumeWrite, "test_volume_allocation", test_volume_allocation, devices=devices)
 add_function_test(TestVolumeWrite, "test_volume_allocate_by_tiles_f", test_volume_allocate_by_tiles_f, devices=devices)
 add_function_test(TestVolumeWrite, "test_volume_allocate_by_tiles_v", test_volume_allocate_by_tiles_v, devices=devices)
+add_function_test(
+    TestVolumeWrite, "test_volume_allocate_by_tiles_index", test_volume_allocate_by_tiles_index, devices=devices
+)
+add_function_test(
+    TestVolumeWrite,
+    "test_volume_allocation_from_voxels",
+    test_volume_allocation_from_voxels,
+    devices=devices,
+)
 
 
 if __name__ == "__main__":

warp/tests/unittest_serial.py

@@ -5,28 +5,29 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-import warp as wp
-from warp.tests.unittest_utils import TeamCityTestRunner
+import unittest
 
-wp.init()
+import warp as wp
+import warp.tests.unittest_suites
 
 
 def run_suite() -> bool:
     """Run a test suite"""
 
-    import warp.tests.unittest_suites
-
     # force rebuild of all kernels
     wp.build.clear_kernel_cache()
     print("Cleared Warp kernel cache")
 
-    runner = TeamCityTestRunner(verbosity=2, failfast=False)
+    runner = unittest.TextTestRunner(verbosity=2, failfast=True)
+
+    # Can swap out different suites
+    suite = warp.tests.unittest_suites.default_suite()
+    # suite = warp.tests.unittest_suites.auto_discover_suite()
+    # suite = warp.tests.unittest_suites.kit_suite()
 
-    # Can swap out warp.tests.unittest_suites.explicit_suite()
-    suite = warp.tests.unittest_suites.auto_discover_suite()
     print(f"Test suite has {suite.countTestCases()} tests")
 
-    ret = not runner.run(suite, "WarpTests").wasSuccessful()
+    ret = not runner.run(suite).wasSuccessful()
     return ret
 
 

warp/tests/unittest_suites.py

@@ -82,7 +82,7 @@ def compare_unittest_suites(
     return test_suite
 
 
-def default_suite(test_loader: unittest.TestLoader):
+def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
     """Example of a manually constructed test suite.
 
     Intended to be modified to create additional test suites
@@ -144,6 +144,7 @@ def default_suite(test_loader: unittest.TestLoader):
     from warp.tests.test_mesh_query_ray import TestMeshQueryRay
     from warp.tests.test_mlp import TestMLP
     from warp.tests.test_model import TestModel
+    from warp.tests.test_module_hashing import TestModuleHashing
     from warp.tests.test_modules_lite import TestModuleLite
     from warp.tests.test_multigpu import TestMultiGPU
     from warp.tests.test_noise import TestNoise
@@ -234,6 +235,7 @@ def default_suite(test_loader: unittest.TestLoader):
         TestMeshQueryRay,
         TestMLP,
         TestModel,
+        TestModuleHashing,
         TestModuleLite,
         TestMultiGPU,
         TestNoise,
@@ -271,7 +273,7 @@ def default_suite(test_loader: unittest.TestLoader):
     return _create_suite_from_test_classes(test_loader, test_classes)
 
 
-def kit_suite(test_loader: unittest.TestLoader):
+def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
     """Tries to mimic the test suite used for testing omni.warp.core in Kit
 
     Requires manual updates with test_ext.py for now.
@@ -300,6 +302,7 @@ def kit_suite(test_loader: unittest.TestLoader):
     from warp.tests.test_mesh_query_aabb import TestMeshQueryAABBMethods
     from warp.tests.test_mesh_query_point import TestMeshQueryPoint
     from warp.tests.test_mesh_query_ray import TestMeshQueryRay
+    from warp.tests.test_module_hashing import TestModuleHashing
     from warp.tests.test_modules_lite import TestModuleLite
     from warp.tests.test_noise import TestNoise
     from warp.tests.test_operators import TestOperators
@@ -342,6 +345,7 @@ def kit_suite(test_loader: unittest.TestLoader):
         TestMeshQueryAABBMethods,
         TestMeshQueryPoint,
         TestMeshQueryRay,
+        TestModuleHashing,
         TestModuleLite,
         TestNoise,
         TestOperators,