warp-lang 1.4.2__py3-none-macosx_10_13_universal2.whl → 1.5.0__py3-none-macosx_10_13_universal2.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of warp-lang might be problematic. Click here for more details.

Files changed (158) hide show
  1. warp/__init__.py +4 -0
  2. warp/autograd.py +43 -8
  3. warp/bin/libwarp-clang.dylib +0 -0
  4. warp/bin/libwarp.dylib +0 -0
  5. warp/build.py +21 -2
  6. warp/build_dll.py +23 -6
  7. warp/builtins.py +1783 -2
  8. warp/codegen.py +177 -45
  9. warp/config.py +2 -2
  10. warp/context.py +321 -73
  11. warp/examples/assets/pixel.jpg +0 -0
  12. warp/examples/benchmarks/benchmark_cloth_paddle.py +86 -0
  13. warp/examples/benchmarks/benchmark_gemm.py +121 -0
  14. warp/examples/benchmarks/benchmark_interop_paddle.py +158 -0
  15. warp/examples/benchmarks/benchmark_tile.py +179 -0
  16. warp/examples/fem/example_adaptive_grid.py +37 -10
  17. warp/examples/fem/example_apic_fluid.py +3 -2
  18. warp/examples/fem/example_convection_diffusion_dg.py +4 -5
  19. warp/examples/fem/example_deformed_geometry.py +1 -1
  20. warp/examples/fem/example_diffusion_3d.py +47 -4
  21. warp/examples/fem/example_distortion_energy.py +220 -0
  22. warp/examples/fem/example_magnetostatics.py +127 -85
  23. warp/examples/fem/example_nonconforming_contact.py +5 -5
  24. warp/examples/fem/example_stokes.py +3 -1
  25. warp/examples/fem/example_streamlines.py +12 -19
  26. warp/examples/fem/utils.py +38 -15
  27. warp/examples/sim/example_cloth.py +2 -25
  28. warp/examples/sim/example_quadruped.py +2 -1
  29. warp/examples/tile/example_tile_convolution.py +58 -0
  30. warp/examples/tile/example_tile_fft.py +47 -0
  31. warp/examples/tile/example_tile_filtering.py +105 -0
  32. warp/examples/tile/example_tile_matmul.py +79 -0
  33. warp/examples/tile/example_tile_mlp.py +375 -0
  34. warp/fem/__init__.py +8 -0
  35. warp/fem/cache.py +16 -12
  36. warp/fem/dirichlet.py +1 -1
  37. warp/fem/domain.py +44 -1
  38. warp/fem/field/__init__.py +1 -2
  39. warp/fem/field/field.py +31 -19
  40. warp/fem/field/nodal_field.py +101 -49
  41. warp/fem/field/virtual.py +794 -0
  42. warp/fem/geometry/__init__.py +2 -2
  43. warp/fem/geometry/deformed_geometry.py +3 -105
  44. warp/fem/geometry/element.py +13 -0
  45. warp/fem/geometry/geometry.py +165 -5
  46. warp/fem/geometry/grid_2d.py +3 -6
  47. warp/fem/geometry/grid_3d.py +31 -28
  48. warp/fem/geometry/hexmesh.py +3 -46
  49. warp/fem/geometry/nanogrid.py +3 -2
  50. warp/fem/geometry/{quadmesh_2d.py → quadmesh.py} +280 -159
  51. warp/fem/geometry/tetmesh.py +2 -43
  52. warp/fem/geometry/{trimesh_2d.py → trimesh.py} +354 -186
  53. warp/fem/integrate.py +683 -261
  54. warp/fem/linalg.py +404 -0
  55. warp/fem/operator.py +101 -18
  56. warp/fem/polynomial.py +5 -5
  57. warp/fem/quadrature/quadrature.py +45 -21
  58. warp/fem/space/__init__.py +45 -11
  59. warp/fem/space/basis_function_space.py +451 -0
  60. warp/fem/space/basis_space.py +58 -11
  61. warp/fem/space/function_space.py +146 -5
  62. warp/fem/space/grid_2d_function_space.py +80 -66
  63. warp/fem/space/grid_3d_function_space.py +113 -68
  64. warp/fem/space/hexmesh_function_space.py +96 -108
  65. warp/fem/space/nanogrid_function_space.py +62 -110
  66. warp/fem/space/quadmesh_function_space.py +208 -0
  67. warp/fem/space/shape/__init__.py +45 -7
  68. warp/fem/space/shape/cube_shape_function.py +328 -54
  69. warp/fem/space/shape/shape_function.py +10 -1
  70. warp/fem/space/shape/square_shape_function.py +328 -60
  71. warp/fem/space/shape/tet_shape_function.py +269 -19
  72. warp/fem/space/shape/triangle_shape_function.py +238 -19
  73. warp/fem/space/tetmesh_function_space.py +69 -37
  74. warp/fem/space/topology.py +38 -0
  75. warp/fem/space/trimesh_function_space.py +179 -0
  76. warp/fem/utils.py +6 -331
  77. warp/jax_experimental.py +3 -1
  78. warp/native/array.h +15 -0
  79. warp/native/builtin.h +66 -26
  80. warp/native/bvh.h +4 -0
  81. warp/native/coloring.cpp +600 -0
  82. warp/native/cuda_util.cpp +14 -0
  83. warp/native/cuda_util.h +2 -1
  84. warp/native/fabric.h +8 -0
  85. warp/native/hashgrid.h +4 -0
  86. warp/native/marching.cu +8 -0
  87. warp/native/mat.h +14 -3
  88. warp/native/mathdx.cpp +59 -0
  89. warp/native/mesh.h +4 -0
  90. warp/native/range.h +13 -1
  91. warp/native/reduce.cpp +9 -1
  92. warp/native/reduce.cu +7 -0
  93. warp/native/runlength_encode.cpp +9 -1
  94. warp/native/runlength_encode.cu +7 -1
  95. warp/native/scan.cpp +8 -0
  96. warp/native/scan.cu +8 -0
  97. warp/native/scan.h +8 -1
  98. warp/native/sparse.cpp +8 -0
  99. warp/native/sparse.cu +8 -0
  100. warp/native/temp_buffer.h +7 -0
  101. warp/native/tile.h +1857 -0
  102. warp/native/tile_gemm.h +341 -0
  103. warp/native/tile_reduce.h +210 -0
  104. warp/native/volume_builder.cu +8 -0
  105. warp/native/volume_builder.h +8 -0
  106. warp/native/warp.cpp +10 -2
  107. warp/native/warp.cu +369 -15
  108. warp/native/warp.h +12 -2
  109. warp/optim/adam.py +39 -4
  110. warp/paddle.py +29 -12
  111. warp/render/render_opengl.py +137 -65
  112. warp/sim/graph_coloring.py +292 -0
  113. warp/sim/integrator_euler.py +4 -2
  114. warp/sim/integrator_featherstone.py +115 -44
  115. warp/sim/integrator_vbd.py +6 -0
  116. warp/sim/model.py +88 -15
  117. warp/stubs.py +569 -4
  118. warp/tape.py +12 -7
  119. warp/tests/assets/pixel.npy +0 -0
  120. warp/tests/aux_test_instancing_gc.py +18 -0
  121. warp/tests/test_array.py +39 -0
  122. warp/tests/test_codegen.py +81 -1
  123. warp/tests/test_codegen_instancing.py +30 -0
  124. warp/tests/test_collision.py +110 -0
  125. warp/tests/test_coloring.py +241 -0
  126. warp/tests/test_context.py +34 -0
  127. warp/tests/test_examples.py +18 -4
  128. warp/tests/test_fem.py +453 -113
  129. warp/tests/test_func.py +13 -0
  130. warp/tests/test_generics.py +52 -0
  131. warp/tests/test_iter.py +68 -0
  132. warp/tests/test_mat_scalar_ops.py +1 -1
  133. warp/tests/test_mesh_query_point.py +1 -1
  134. warp/tests/test_module_hashing.py +23 -0
  135. warp/tests/test_paddle.py +27 -87
  136. warp/tests/test_print.py +56 -1
  137. warp/tests/test_spatial.py +1 -1
  138. warp/tests/test_tile.py +700 -0
  139. warp/tests/test_tile_mathdx.py +144 -0
  140. warp/tests/test_tile_mlp.py +383 -0
  141. warp/tests/test_tile_reduce.py +374 -0
  142. warp/tests/test_tile_shared_memory.py +190 -0
  143. warp/tests/test_vbd.py +12 -20
  144. warp/tests/test_volume.py +43 -0
  145. warp/tests/unittest_suites.py +19 -2
  146. warp/tests/unittest_utils.py +4 -0
  147. warp/types.py +338 -72
  148. warp/utils.py +22 -1
  149. {warp_lang-1.4.2.dist-info → warp_lang-1.5.0.dist-info}/METADATA +33 -7
  150. {warp_lang-1.4.2.dist-info → warp_lang-1.5.0.dist-info}/RECORD +153 -126
  151. {warp_lang-1.4.2.dist-info → warp_lang-1.5.0.dist-info}/WHEEL +1 -1
  152. warp/fem/field/test.py +0 -180
  153. warp/fem/field/trial.py +0 -183
  154. warp/fem/space/collocated_function_space.py +0 -102
  155. warp/fem/space/quadmesh_2d_function_space.py +0 -261
  156. warp/fem/space/trimesh_2d_function_space.py +0 -153
  157. {warp_lang-1.4.2.dist-info → warp_lang-1.5.0.dist-info}/LICENSE.md +0 -0
  158. {warp_lang-1.4.2.dist-info → warp_lang-1.5.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,18 @@
1
+ # Copyright (c) 2024 NVIDIA CORPORATION. All rights reserved.
2
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
3
+ # and proprietary rights in and to this software, related documentation
4
+ # and any modifications thereto. Any use, reproduction, disclosure or
5
+ # distribution of this software and related documentation without an express
6
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
+
8
+ """Helper module used in test_codegen_instancing.py"""
9
+
10
+ import warp as wp
11
+
12
+
13
+ def create_kernel_closure(value: int):
14
+ @wp.kernel
15
+ def k(a: wp.array(dtype=int)):
16
+ a[0] = value
17
+
18
+ return k
warp/tests/test_array.py CHANGED
@@ -6,6 +6,7 @@
6
6
  # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
7
 
8
8
  import unittest
9
+ from typing import Any
9
10
 
10
11
  import numpy as np
11
12
 
@@ -2429,6 +2430,16 @@ def inplace_add_rhs(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.ar
2429
2430
  wp.atomic_add(z, 0, a)
2430
2431
 
2431
2432
 
2433
+ vec9 = wp.vec(length=9, dtype=float)
2434
+
2435
+
2436
+ @wp.kernel
2437
+ def inplace_add_custom_vec(x: wp.array(dtype=vec9), y: wp.array(dtype=vec9)):
2438
+ i = wp.tid()
2439
+ x[i] += y[i]
2440
+ x[i] += y[i]
2441
+
2442
+
2432
2443
  def test_array_inplace_diff_ops(test, device):
2433
2444
  N = 3
2434
2445
  x1 = wp.ones(N, dtype=float, requires_grad=True, device=device)
@@ -2537,6 +2548,18 @@ def test_array_inplace_diff_ops(test, device):
2537
2548
 
2538
2549
  assert_np_equal(x.grad.numpy(), np.ones(1, dtype=float))
2539
2550
  assert_np_equal(y.grad.numpy(), np.ones(1, dtype=float))
2551
+ tape.reset()
2552
+
2553
+ x = wp.zeros(1, dtype=vec9, requires_grad=True, device=device)
2554
+ y = wp.ones(1, dtype=vec9, requires_grad=True, device=device)
2555
+
2556
+ with tape:
2557
+ wp.launch(inplace_add_custom_vec, 1, inputs=[x, y], device=device)
2558
+
2559
+ tape.backward(grads={x: wp.ones_like(x)})
2560
+
2561
+ assert_np_equal(x.numpy(), np.full((1, 9), 2.0, dtype=float))
2562
+ assert_np_equal(y.grad.numpy(), np.full((1, 9), 2.0, dtype=float))
2540
2563
 
2541
2564
 
2542
2565
  @wp.kernel
@@ -2551,6 +2574,15 @@ def inplace_div_1d(x: wp.array(dtype=float), y: wp.array(dtype=float)):
2551
2574
  x[i] /= y[i]
2552
2575
 
2553
2576
 
2577
+ @wp.kernel
2578
+ def inplace_add_non_atomic_types(x: wp.array(dtype=Any), y: wp.array(dtype=Any)):
2579
+ i = wp.tid()
2580
+ x[i] += y[i]
2581
+
2582
+
2583
+ uint16vec3 = wp.vec(length=3, dtype=wp.uint16)
2584
+
2585
+
2554
2586
  def test_array_inplace_non_diff_ops(test, device):
2555
2587
  N = 3
2556
2588
  x1 = wp.full(N, value=10.0, dtype=float, device=device)
@@ -2564,6 +2596,13 @@ def test_array_inplace_non_diff_ops(test, device):
2564
2596
  wp.launch(inplace_div_1d, N, inputs=[x1, y1], device=device)
2565
2597
  assert_np_equal(x1.numpy(), np.full(N, fill_value=2.0, dtype=float))
2566
2598
 
2599
+ for dtype in wp.types.non_atomic_types + (wp.vec2b, wp.vec2ub, wp.vec2s, wp.vec2us, uint16vec3):
2600
+ x = wp.full(N, value=0, dtype=dtype, device=device)
2601
+ y = wp.full(N, value=1, dtype=dtype, device=device)
2602
+
2603
+ wp.launch(inplace_add_non_atomic_types, N, inputs=[x, y], device=device)
2604
+ assert_np_equal(x.numpy(), y.numpy())
2605
+
2567
2606
 
2568
2607
  @wp.kernel
2569
2608
  def inc_scalar(a: wp.array(dtype=float)):
@@ -7,11 +7,27 @@
7
7
 
8
8
  import sys
9
9
  import unittest
10
+ from typing import Tuple
10
11
 
11
12
  import warp as wp
12
13
  from warp.tests.unittest_utils import *
13
14
 
14
15
 
16
+ @wp.kernel
17
+ def test_expect():
18
+ a = 1.0
19
+ a += 2.0
20
+
21
+ wp.expect_eq(123, 123)
22
+ wp.expect_neq(123, 234)
23
+
24
+ wp.expect_eq(wp.vec2(1.0, 2.0), wp.vec2(1.0, 2.0))
25
+ wp.expect_neq(wp.vec2(1.0, 2.0), wp.vec2(2.0, 3.0))
26
+
27
+ wp.expect_eq(wp.mat22(1.0, 2.0, 3.0, 4.0), wp.mat22(1.0, 2.0, 3.0, 4.0))
28
+ wp.expect_neq(wp.mat22(1.0, 2.0, 3.0, 4.0), wp.mat22(2.0, 3.0, 4.0, 5.0))
29
+
30
+
15
31
  @wp.kernel
16
32
  def test_rename():
17
33
  a = 0
@@ -574,6 +590,64 @@ def test_error_mutating_constant_in_dynamic_loop(test, device):
574
590
  assert_np_equal(mats.numpy(), np.zeros((1, 3, 3)))
575
591
 
576
592
 
593
+ def test_error_return_annotation_mismatch(test, device):
594
+ @wp.func
595
+ def foo_1(x: wp.int32) -> wp.int16:
596
+ return wp.int8(x)
597
+
598
+ def kernel_1_fn():
599
+ x = foo_1(123)
600
+
601
+ @wp.func
602
+ def foo_2(x: int) -> int:
603
+ return (x + x, x * x)
604
+
605
+ def kernel_2_fn():
606
+ x = foo_2(123)
607
+
608
+ @wp.func
609
+ def foo_3(x: int) -> Tuple[int, int]:
610
+ return (x, 1.23)
611
+
612
+ def kernel_3_fn():
613
+ x, y = foo_3(123)
614
+
615
+ @wp.func
616
+ def foo_4(x: int) -> Tuple[int, int, int]:
617
+ return (x + x, x * x)
618
+
619
+ def kernel_4_fn():
620
+ x, y, z = foo_4(123)
621
+
622
+ kernel = wp.Kernel(func=kernel_1_fn)
623
+ with test.assertRaisesRegex(
624
+ wp.codegen.WarpCodegenError,
625
+ r"The function `foo_1` has its return type annotated as `int16` but the code returns a value of type `int8`.",
626
+ ):
627
+ wp.launch(kernel, dim=1, device=device)
628
+
629
+ kernel = wp.Kernel(func=kernel_2_fn)
630
+ with test.assertRaisesRegex(
631
+ wp.codegen.WarpCodegenError,
632
+ r"The function `foo_2` has its return type annotated as `int` but the code returns 2 values.",
633
+ ):
634
+ wp.launch(kernel, dim=1, device=device)
635
+
636
+ kernel = wp.Kernel(func=kernel_3_fn)
637
+ with test.assertRaisesRegex(
638
+ wp.codegen.WarpCodegenError,
639
+ r"The function `foo_3` has its return type annotated as `Tuple\[int, int\]` but the code returns a tuple with types `\(int32, float32\)`.",
640
+ ):
641
+ wp.launch(kernel, dim=1, device=device)
642
+
643
+ kernel = wp.Kernel(func=kernel_4_fn)
644
+ with test.assertRaisesRegex(
645
+ wp.codegen.WarpCodegenError,
646
+ r"The function `foo_4` has its return type annotated as a tuple of 3 elements but the code returns 2 values.",
647
+ ):
648
+ wp.launch(kernel, dim=1, device=device)
649
+
650
+
577
651
  @wp.kernel
578
652
  def test_call_syntax():
579
653
  expected_pow = 16.0
@@ -622,6 +696,7 @@ class TestCodeGen(unittest.TestCase):
622
696
 
623
697
  devices = get_test_devices()
624
698
 
699
+ add_kernel_test(TestCodeGen, name="test_expect", kernel=test_expect, dim=1, devices=devices)
625
700
  add_kernel_test(TestCodeGen, name="test_inplace", kernel=test_inplace, dim=1, devices=devices)
626
701
  add_kernel_test(TestCodeGen, name="test_rename", kernel=test_rename, dim=1, devices=devices)
627
702
  add_kernel_test(TestCodeGen, name="test_constant", kernel=test_constant, inputs=[1.0], dim=1, devices=devices)
@@ -758,7 +833,12 @@ add_function_test(
758
833
  name="test_error_mutating_constant_in_dynamic_loop",
759
834
  devices=devices,
760
835
  )
761
-
836
+ add_function_test(
837
+ TestCodeGen,
838
+ func=test_error_return_annotation_mismatch,
839
+ name="test_error_return_annotation_mismatch",
840
+ devices=devices,
841
+ )
762
842
  add_kernel_test(TestCodeGen, name="test_call_syntax", kernel=test_call_syntax, dim=1, devices=devices)
763
843
  add_kernel_test(TestCodeGen, name="test_shadow_builtin", kernel=test_shadow_builtin, dim=1, devices=devices)
764
844
  add_kernel_test(TestCodeGen, name="test_while_condition_eval", kernel=test_while_condition_eval, dim=1, devices=devices)
@@ -1287,6 +1287,35 @@ def test_module_mark_modified(test, device):
1287
1287
  # =======================================================================
1288
1288
 
1289
1289
 
1290
+ def test_garbage_collection(test, device):
1291
+ """Test that dynamically generated kernels without user references are not retained in the module."""
1292
+
1293
+ # use a helper module with a known kernel count
1294
+ import warp.tests.aux_test_instancing_gc as gc_test_module
1295
+
1296
+ with wp.ScopedDevice(device):
1297
+ a = wp.zeros(1, dtype=int)
1298
+
1299
+ for i in range(10):
1300
+ # create a unique kernel on each iteration
1301
+ k = gc_test_module.create_kernel_closure(i)
1302
+
1303
+ # import gc
1304
+ # gc.collect()
1305
+
1306
+ # since we don't keep references to the previous kernels,
1307
+ # they should be garbage-collected and not appear in the module
1308
+ k.module.load(device=device)
1309
+ test.assertEqual(len(k.module.live_kernels), 1)
1310
+
1311
+ # test the kernel
1312
+ wp.launch(k, dim=1, inputs=[a])
1313
+ test.assertEqual(a.numpy()[0], i)
1314
+
1315
+
1316
+ # =======================================================================
1317
+
1318
+
1290
1319
  class TestCodeGenInstancing(unittest.TestCase):
1291
1320
  pass
1292
1321
 
@@ -1450,6 +1479,7 @@ add_function_test(TestCodeGenInstancing, func=test_create_kernel_loop, name="tes
1450
1479
  add_function_test(
1451
1480
  TestCodeGenInstancing, func=test_module_mark_modified, name="test_module_mark_modified", devices=devices
1452
1481
  )
1482
+ add_function_test(TestCodeGenInstancing, func=test_garbage_collection, name="test_garbage_collection", devices=devices)
1453
1483
 
1454
1484
 
1455
1485
  if __name__ == "__main__":
@@ -7,8 +7,11 @@
7
7
 
8
8
  import unittest
9
9
 
10
+ import warp as wp
10
11
  import warp.examples
12
+ import warp.sim
11
13
  from warp.sim.collide import *
14
+ from warp.sim.integrator_euler import eval_triangles_contact
12
15
  from warp.tests.unittest_utils import *
13
16
 
14
17
 
@@ -474,12 +477,119 @@ def test_edge_edge_collision(test, device):
474
477
  devices = get_test_devices()
475
478
 
476
479
 
480
+ def test_particle_collision(test, device):
481
+ with wp.ScopedDevice(device):
482
+ contact_radius = 1.23
483
+ builder1 = wp.sim.ModelBuilder()
484
+ builder1.add_cloth_grid(
485
+ pos=wp.vec3(0.0, 0.0, 0.0),
486
+ rot=wp.quat_identity(),
487
+ vel=wp.vec3(0.0, 0.0, 0.0),
488
+ dim_x=100,
489
+ dim_y=100,
490
+ cell_x=0.1,
491
+ cell_y=0.1,
492
+ mass=0.1,
493
+ particle_radius=contact_radius,
494
+ )
495
+
496
+ cloth_grid = builder1.finalize()
497
+ cloth_grid_particle_radius = cloth_grid.particle_radius.numpy()
498
+ assert_np_equal(cloth_grid_particle_radius, np.full(cloth_grid_particle_radius.shape, contact_radius), tol=1e-5)
499
+
500
+ vertices = [
501
+ [2.0, 0.0, 0.0],
502
+ [2.0, 2.0, 0.0],
503
+ [0.0, 0.0, 0.0],
504
+ [1.0, 0.0, 1.0],
505
+ [1.0, 1.0, 1.0],
506
+ [0.0, 0.0, 1.0],
507
+ ]
508
+ vertices = [wp.vec3(v) for v in vertices]
509
+ faces = [0, 1, 2, 3, 4, 5]
510
+
511
+ builder2 = wp.sim.ModelBuilder()
512
+ builder2.add_cloth_mesh(
513
+ pos=wp.vec3(0.0, 0.0, 0.0),
514
+ rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
515
+ scale=1.0,
516
+ vertices=vertices,
517
+ indices=faces,
518
+ tri_ke=1e4,
519
+ tri_ka=1e4,
520
+ tri_kd=1e-5,
521
+ edge_ke=10,
522
+ edge_kd=0.0,
523
+ vel=wp.vec3(0.0, 0.0, 0.0),
524
+ density=0.1,
525
+ particle_radius=contact_radius,
526
+ )
527
+ cloth_mesh = builder2.finalize()
528
+ cloth_mesh_particle_radius = cloth_mesh.particle_radius.numpy()
529
+ assert_np_equal(cloth_mesh_particle_radius, np.full(cloth_mesh_particle_radius.shape, contact_radius), tol=1e-5)
530
+
531
+ state = cloth_mesh.state()
532
+ particle_f = wp.zeros_like(state.particle_q)
533
+ wp.launch(
534
+ kernel=eval_triangles_contact,
535
+ dim=cloth_mesh.tri_count * cloth_mesh.particle_count,
536
+ inputs=[
537
+ cloth_mesh.particle_count,
538
+ state.particle_q,
539
+ state.particle_qd,
540
+ cloth_mesh.tri_indices,
541
+ cloth_mesh.tri_materials,
542
+ cloth_mesh.particle_radius,
543
+ ],
544
+ outputs=[particle_f],
545
+ )
546
+ test.assertTrue((np.linalg.norm(particle_f.numpy(), axis=1) != 0).all())
547
+
548
+ builder3 = wp.sim.ModelBuilder()
549
+ builder3.add_cloth_mesh(
550
+ pos=wp.vec3(0.0, 0.0, 0.0),
551
+ rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
552
+ scale=1.0,
553
+ vertices=vertices,
554
+ indices=faces,
555
+ tri_ke=1e4,
556
+ tri_ka=1e4,
557
+ tri_kd=1e-5,
558
+ edge_ke=10,
559
+ edge_kd=0.0,
560
+ vel=wp.vec3(0.0, 0.0, 0.0),
561
+ density=0.1,
562
+ particle_radius=0.5,
563
+ )
564
+ cloth_mesh_2 = builder3.finalize()
565
+ cloth_mesh_2_particle_radius = cloth_mesh_2.particle_radius.numpy()
566
+ assert_np_equal(cloth_mesh_2_particle_radius, np.full(cloth_mesh_2_particle_radius.shape, 0.5), tol=1e-5)
567
+
568
+ state_2 = cloth_mesh_2.state()
569
+ particle_f_2 = wp.zeros_like(cloth_mesh_2.particle_q)
570
+ wp.launch(
571
+ kernel=eval_triangles_contact,
572
+ dim=cloth_mesh_2.tri_count * cloth_mesh_2.particle_count,
573
+ inputs=[
574
+ cloth_mesh_2.particle_count,
575
+ state_2.particle_q,
576
+ state_2.particle_qd,
577
+ cloth_mesh_2.tri_indices,
578
+ cloth_mesh_2.tri_materials,
579
+ cloth_mesh_2.particle_radius,
580
+ ],
581
+ outputs=[particle_f_2],
582
+ )
583
+ test.assertTrue((np.linalg.norm(particle_f_2.numpy(), axis=1) == 0).all())
584
+
585
+
477
586
  class TestCollision(unittest.TestCase):
478
587
  pass
479
588
 
480
589
 
481
590
  add_function_test(TestCollision, "test_vertex_triangle_collision", test_vertex_triangle_collision, devices=devices)
482
591
  add_function_test(TestCollision, "test_edge_edge_collision", test_vertex_triangle_collision, devices=devices)
592
+ add_function_test(TestCollision, "test_particle_collision", test_particle_collision, devices=devices)
483
593
 
484
594
  if __name__ == "__main__":
485
595
  wp.clear_kernel_cache()
@@ -0,0 +1,241 @@
1
+ # Copyright (c) 2024 NVIDIA CORPORATION. All rights reserved.
2
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
3
+ # and proprietary rights in and to this software, related documentation
4
+ # and any modifications thereto. Any use, reproduction, disclosure or
5
+ # distribution of this software and related documentation without an express
6
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
+
8
+
9
+ import numpy as np
10
+
11
+ import warp as wp
12
+ import warp.examples
13
+ import warp.sim
14
+ from warp.sim.graph_coloring import ColoringAlgorithm, construct_trimesh_graph_edges, validate_graph_coloring
15
+ from warp.tests.unittest_utils import *
16
+
17
+
18
+ def color_lattice_grid(num_x, num_y):
19
+ colors = []
20
+ for _ in range(4):
21
+ colors.append([])
22
+
23
+ for xi in range(num_x + 1):
24
+ for yi in range(num_y + 1):
25
+ node_dx = yi * (num_x + 1) + xi
26
+
27
+ a = 1 if xi % 2 else 0
28
+ b = 1 if yi % 2 else 0
29
+
30
+ c = b * 2 + a
31
+
32
+ colors[c].append(node_dx)
33
+
34
+ color_groups = [np.array(group) for group in colors]
35
+
36
+ return color_groups
37
+
38
+
39
+ @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
40
+ def test_coloring_trimesh(test, device):
41
+ from pxr import Usd, UsdGeom
42
+
43
+ with wp.ScopedDevice(device):
44
+ usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), "bunny.usd"))
45
+ usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
46
+
47
+ vertices = np.array(usd_geom.GetPointsAttr().Get())
48
+ faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
49
+
50
+ builder = wp.sim.ModelBuilder()
51
+
52
+ builder.add_cloth_mesh(
53
+ pos=wp.vec3(0.0, 0.0, 0.0),
54
+ rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
55
+ scale=1.0,
56
+ vertices=[wp.vec3(p) for p in vertices],
57
+ indices=faces.flatten(),
58
+ vel=wp.vec3(0.0, 0.0, 0.0),
59
+ density=0.02,
60
+ )
61
+
62
+ model = builder.finalize()
63
+
64
+ particle_colors = wp.empty(shape=(model.particle_count), dtype=int, device="cpu")
65
+
66
+ edge_indices_cpu = wp.array(model.edge_indices.numpy()[:, 2:], dtype=int, device="cpu")
67
+
68
+ # coloring without bending
69
+ num_colors_greedy = wp.context.runtime.core.graph_coloring(
70
+ model.particle_count,
71
+ edge_indices_cpu.__ctype__(),
72
+ ColoringAlgorithm.GREEDY.value,
73
+ particle_colors.__ctype__(),
74
+ )
75
+ wp.launch(
76
+ kernel=validate_graph_coloring,
77
+ inputs=[edge_indices_cpu, particle_colors],
78
+ dim=edge_indices_cpu.shape[0],
79
+ device="cpu",
80
+ )
81
+
82
+ num_colors_mcs = wp.context.runtime.core.graph_coloring(
83
+ model.particle_count,
84
+ edge_indices_cpu.__ctype__(),
85
+ ColoringAlgorithm.MCS.value,
86
+ particle_colors.__ctype__(),
87
+ )
88
+ wp.launch(
89
+ kernel=validate_graph_coloring,
90
+ inputs=[edge_indices_cpu, particle_colors],
91
+ dim=edge_indices_cpu.shape[0],
92
+ device="cpu",
93
+ )
94
+
95
+ # coloring with bending
96
+ edge_indices_cpu_with_bending = construct_trimesh_graph_edges(model.edge_indices, True)
97
+ num_colors_greedy = wp.context.runtime.core.graph_coloring(
98
+ model.particle_count,
99
+ edge_indices_cpu_with_bending.__ctype__(),
100
+ ColoringAlgorithm.GREEDY.value,
101
+ particle_colors.__ctype__(),
102
+ )
103
+ wp.context.runtime.core.balance_coloring(
104
+ model.particle_count,
105
+ edge_indices_cpu_with_bending.__ctype__(),
106
+ num_colors_greedy,
107
+ 1.1,
108
+ particle_colors.__ctype__(),
109
+ )
110
+ wp.launch(
111
+ kernel=validate_graph_coloring,
112
+ inputs=[edge_indices_cpu_with_bending, particle_colors],
113
+ dim=edge_indices_cpu_with_bending.shape[0],
114
+ device="cpu",
115
+ )
116
+
117
+ num_colors_mcs = wp.context.runtime.core.graph_coloring(
118
+ model.particle_count,
119
+ edge_indices_cpu_with_bending.__ctype__(),
120
+ ColoringAlgorithm.MCS.value,
121
+ particle_colors.__ctype__(),
122
+ )
123
+ wp.context.runtime.core.balance_coloring(
124
+ model.particle_count,
125
+ edge_indices_cpu_with_bending.__ctype__(),
126
+ num_colors_mcs,
127
+ 1.1,
128
+ particle_colors.__ctype__(),
129
+ )
130
+ wp.launch(
131
+ kernel=validate_graph_coloring,
132
+ inputs=[edge_indices_cpu_with_bending, particle_colors],
133
+ dim=edge_indices_cpu_with_bending.shape[0],
134
+ device="cpu",
135
+ )
136
+
137
+
138
+ @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
139
+ def test_combine_coloring(test, device):
140
+ from pxr import Usd, UsdGeom
141
+
142
+ with wp.ScopedDevice(device):
143
+ builder1 = wp.sim.ModelBuilder()
144
+ usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), "bunny.usd"))
145
+ usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
146
+
147
+ vertices = np.array(usd_geom.GetPointsAttr().Get())
148
+ faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
149
+
150
+ builder1.add_cloth_mesh(
151
+ pos=wp.vec3(0.0, 0.0, 0.0),
152
+ rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
153
+ scale=1.0,
154
+ vertices=[wp.vec3(p) for p in vertices],
155
+ indices=faces.flatten(),
156
+ vel=wp.vec3(0.0, 0.0, 0.0),
157
+ density=0.02,
158
+ )
159
+
160
+ builder1.add_cloth_grid(
161
+ pos=wp.vec3(0.0, 4.0, 0.0),
162
+ rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
163
+ vel=wp.vec3(0.0, 0.0, 0.0),
164
+ dim_x=50,
165
+ dim_y=100,
166
+ cell_x=0.1,
167
+ cell_y=0.1,
168
+ mass=0.1,
169
+ fix_left=True,
170
+ )
171
+ builder1.color()
172
+
173
+ builder2 = wp.sim.ModelBuilder()
174
+ builder2.add_cloth_grid(
175
+ pos=wp.vec3(0.0, 4.0, 0.0),
176
+ rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
177
+ vel=wp.vec3(0.0, 0.0, 0.0),
178
+ dim_x=50,
179
+ dim_y=100,
180
+ cell_x=0.1,
181
+ cell_y=0.1,
182
+ mass=0.1,
183
+ # to include bending in coloring
184
+ edge_ke=100000,
185
+ fix_left=True,
186
+ )
187
+ builder2.color()
188
+
189
+ builder3 = wp.sim.ModelBuilder()
190
+ builder3.add_cloth_grid(
191
+ pos=wp.vec3(0.0, 4.0, 0.0),
192
+ rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
193
+ vel=wp.vec3(0.0, 0.0, 0.0),
194
+ dim_x=50,
195
+ dim_y=100,
196
+ cell_x=0.1,
197
+ cell_y=0.1,
198
+ mass=0.1,
199
+ fix_left=True,
200
+ )
201
+
202
+ builder3.set_coloring(
203
+ color_lattice_grid(50, 100),
204
+ )
205
+
206
+ builder1.add_builder(builder2)
207
+ builder1.add_builder(builder3)
208
+
209
+ model = builder2.finalize()
210
+
211
+ particle_number_colored = np.full((model.particle_count), -1, dtype=int)
212
+ particle_colors = np.full((model.particle_count), -1, dtype=int)
213
+ for color, color_group in enumerate(model.particle_coloring):
214
+ particle_number_colored[color_group.numpy()] += 1
215
+ particle_colors[color_group.numpy()] = color
216
+
217
+ # all particles has been colored exactly once
218
+ assert_np_equal(particle_number_colored, 0)
219
+
220
+ edge_indices_cpu = wp.array(model.edge_indices.numpy()[:, 2:], dtype=int, device="cpu")
221
+ wp.launch(
222
+ kernel=validate_graph_coloring,
223
+ inputs=[edge_indices_cpu, wp.array(particle_colors, dtype=int, device="cpu")],
224
+ dim=edge_indices_cpu.shape[0],
225
+ device="cpu",
226
+ )
227
+
228
+
229
+ devices = get_test_devices()
230
+
231
+
232
+ class TestColoring(unittest.TestCase):
233
+ pass
234
+
235
+
236
+ add_function_test(TestColoring, "test_coloring_trimesh", test_coloring_trimesh, devices=devices)
237
+ add_function_test(TestColoring, "test_combine_coloring", test_combine_coloring, devices=devices)
238
+
239
+ if __name__ == "__main__":
240
+ wp.clear_kernel_cache()
241
+ unittest.main(verbosity=2)
@@ -0,0 +1,34 @@
1
+ # Copyright (c) 2024 NVIDIA CORPORATION. All rights reserved.
2
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
3
+ # and proprietary rights in and to this software, related documentation
4
+ # and any modifications thereto. Any use, reproduction, disclosure or
5
+ # distribution of this software and related documentation without an express
6
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
+
8
+ import unittest
9
+ from typing import List, Tuple
10
+
11
+ import warp as wp
12
+ from warp.tests.unittest_utils import *
13
+
14
+
15
+ def test_context_type_str(test, device):
16
+ assert wp.context.type_str(List[int]) == "List[int]"
17
+ assert wp.context.type_str(List[float]) == "List[float]"
18
+
19
+ assert wp.context.type_str(Tuple[int]) == "Tuple[int]"
20
+ assert wp.context.type_str(Tuple[float]) == "Tuple[float]"
21
+ assert wp.context.type_str(Tuple[int, float]) == "Tuple[int, float]"
22
+ assert wp.context.type_str(Tuple[int, ...]) == "Tuple[int, ...]"
23
+
24
+
25
+ class TestContext(unittest.TestCase):
26
+ pass
27
+
28
+
29
+ add_function_test(TestContext, "test_context_type_str", test_context_type_str)
30
+
31
+
32
+ if __name__ == "__main__":
33
+ wp.clear_kernel_cache()
34
+ unittest.main(verbosity=2)