warp-lang 1.5.1__py3-none-macosx_10_13_universal2.whl → 1.6.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 (123) hide show
  1. warp/__init__.py +5 -0
  2. warp/autograd.py +414 -191
  3. warp/bin/libwarp-clang.dylib +0 -0
  4. warp/bin/libwarp.dylib +0 -0
  5. warp/build.py +40 -12
  6. warp/build_dll.py +13 -6
  7. warp/builtins.py +1076 -480
  8. warp/codegen.py +240 -119
  9. warp/config.py +1 -1
  10. warp/context.py +298 -84
  11. warp/examples/assets/square_cloth.usd +0 -0
  12. warp/examples/benchmarks/benchmark_gemm.py +27 -18
  13. warp/examples/benchmarks/benchmark_interop_paddle.py +3 -3
  14. warp/examples/benchmarks/benchmark_interop_torch.py +3 -3
  15. warp/examples/core/example_torch.py +18 -34
  16. warp/examples/fem/example_apic_fluid.py +1 -0
  17. warp/examples/fem/example_mixed_elasticity.py +1 -1
  18. warp/examples/optim/example_bounce.py +1 -1
  19. warp/examples/optim/example_cloth_throw.py +1 -1
  20. warp/examples/optim/example_diffray.py +4 -15
  21. warp/examples/optim/example_drone.py +1 -1
  22. warp/examples/optim/example_softbody_properties.py +392 -0
  23. warp/examples/optim/example_trajectory.py +1 -3
  24. warp/examples/optim/example_walker.py +5 -0
  25. warp/examples/sim/example_cartpole.py +0 -2
  26. warp/examples/sim/example_cloth_self_contact.py +260 -0
  27. warp/examples/sim/example_granular_collision_sdf.py +4 -5
  28. warp/examples/sim/example_jacobian_ik.py +0 -2
  29. warp/examples/sim/example_quadruped.py +5 -2
  30. warp/examples/tile/example_tile_cholesky.py +79 -0
  31. warp/examples/tile/example_tile_convolution.py +2 -2
  32. warp/examples/tile/example_tile_fft.py +2 -2
  33. warp/examples/tile/example_tile_filtering.py +3 -3
  34. warp/examples/tile/example_tile_matmul.py +4 -4
  35. warp/examples/tile/example_tile_mlp.py +12 -12
  36. warp/examples/tile/example_tile_nbody.py +180 -0
  37. warp/examples/tile/example_tile_walker.py +319 -0
  38. warp/math.py +147 -0
  39. warp/native/array.h +12 -0
  40. warp/native/builtin.h +0 -1
  41. warp/native/bvh.cpp +149 -70
  42. warp/native/bvh.cu +287 -68
  43. warp/native/bvh.h +195 -85
  44. warp/native/clang/clang.cpp +5 -1
  45. warp/native/cuda_util.cpp +35 -0
  46. warp/native/cuda_util.h +5 -0
  47. warp/native/exports.h +40 -40
  48. warp/native/intersect.h +17 -0
  49. warp/native/mat.h +41 -0
  50. warp/native/mathdx.cpp +19 -0
  51. warp/native/mesh.cpp +25 -8
  52. warp/native/mesh.cu +153 -101
  53. warp/native/mesh.h +482 -403
  54. warp/native/quat.h +40 -0
  55. warp/native/solid_angle.h +7 -0
  56. warp/native/sort.cpp +85 -0
  57. warp/native/sort.cu +34 -0
  58. warp/native/sort.h +3 -1
  59. warp/native/spatial.h +11 -0
  60. warp/native/tile.h +1185 -664
  61. warp/native/tile_reduce.h +8 -6
  62. warp/native/vec.h +41 -0
  63. warp/native/warp.cpp +8 -1
  64. warp/native/warp.cu +263 -40
  65. warp/native/warp.h +19 -5
  66. warp/optim/linear.py +22 -4
  67. warp/render/render_opengl.py +124 -59
  68. warp/sim/__init__.py +6 -1
  69. warp/sim/collide.py +270 -26
  70. warp/sim/integrator_euler.py +25 -7
  71. warp/sim/integrator_featherstone.py +154 -35
  72. warp/sim/integrator_vbd.py +842 -40
  73. warp/sim/model.py +111 -53
  74. warp/stubs.py +248 -115
  75. warp/tape.py +28 -30
  76. warp/tests/aux_test_module_unload.py +15 -0
  77. warp/tests/{test_sim_grad.py → flaky_test_sim_grad.py} +104 -63
  78. warp/tests/test_array.py +74 -0
  79. warp/tests/test_assert.py +242 -0
  80. warp/tests/test_codegen.py +14 -61
  81. warp/tests/test_collision.py +2 -2
  82. warp/tests/test_examples.py +9 -0
  83. warp/tests/test_grad_debug.py +87 -2
  84. warp/tests/test_hash_grid.py +1 -1
  85. warp/tests/test_ipc.py +116 -0
  86. warp/tests/test_mat.py +138 -167
  87. warp/tests/test_math.py +47 -1
  88. warp/tests/test_matmul.py +11 -7
  89. warp/tests/test_matmul_lite.py +4 -4
  90. warp/tests/test_mesh.py +84 -60
  91. warp/tests/test_mesh_query_aabb.py +165 -0
  92. warp/tests/test_mesh_query_point.py +328 -286
  93. warp/tests/test_mesh_query_ray.py +134 -121
  94. warp/tests/test_mlp.py +2 -2
  95. warp/tests/test_operators.py +43 -0
  96. warp/tests/test_overwrite.py +2 -2
  97. warp/tests/test_quat.py +77 -0
  98. warp/tests/test_reload.py +29 -0
  99. warp/tests/test_sim_grad_bounce_linear.py +204 -0
  100. warp/tests/test_static.py +16 -0
  101. warp/tests/test_tape.py +25 -0
  102. warp/tests/test_tile.py +134 -191
  103. warp/tests/test_tile_load.py +356 -0
  104. warp/tests/test_tile_mathdx.py +61 -8
  105. warp/tests/test_tile_mlp.py +17 -17
  106. warp/tests/test_tile_reduce.py +24 -18
  107. warp/tests/test_tile_shared_memory.py +66 -17
  108. warp/tests/test_tile_view.py +165 -0
  109. warp/tests/test_torch.py +35 -0
  110. warp/tests/test_utils.py +36 -24
  111. warp/tests/test_vec.py +110 -0
  112. warp/tests/unittest_suites.py +29 -4
  113. warp/tests/unittest_utils.py +30 -11
  114. warp/thirdparty/unittest_parallel.py +2 -2
  115. warp/types.py +409 -99
  116. warp/utils.py +9 -5
  117. {warp_lang-1.5.1.dist-info → warp_lang-1.6.0.dist-info}/METADATA +68 -44
  118. {warp_lang-1.5.1.dist-info → warp_lang-1.6.0.dist-info}/RECORD +121 -110
  119. {warp_lang-1.5.1.dist-info → warp_lang-1.6.0.dist-info}/WHEEL +1 -1
  120. warp/examples/benchmarks/benchmark_tile.py +0 -179
  121. warp/native/tile_gemm.h +0 -341
  122. {warp_lang-1.5.1.dist-info → warp_lang-1.6.0.dist-info}/LICENSE.md +0 -0
  123. {warp_lang-1.5.1.dist-info → warp_lang-1.6.0.dist-info}/top_level.txt +0 -0
warp/sim/collide.py CHANGED
@@ -16,6 +16,7 @@ from warp.sim.model import Model
16
16
 
17
17
  from .model import PARTICLE_FLAG_ACTIVE, ModelShapeGeometry
18
18
 
19
+ # types of triangle's closest point to a point
19
20
  TRI_CONTACT_FEATURE_VERTEX_A = wp.constant(0)
20
21
  TRI_CONTACT_FEATURE_VERTEX_B = wp.constant(1)
21
22
  TRI_CONTACT_FEATURE_VERTEX_C = wp.constant(2)
@@ -24,6 +25,45 @@ TRI_CONTACT_FEATURE_EDGE_AC = wp.constant(4)
24
25
  TRI_CONTACT_FEATURE_EDGE_BC = wp.constant(5)
25
26
  TRI_CONTACT_FEATURE_FACE_INTERIOR = wp.constant(6)
26
27
 
28
+ # constants used to access TriMeshCollisionDetector.resize_flags
29
+ VERTEX_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(0)
30
+ TRI_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(1)
31
+ EDGE_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(2)
32
+ TRI_TRI_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(3)
33
+
34
+
35
+ @wp.func
36
+ def build_orthonormal_basis(n: wp.vec3):
37
+ """
38
+ Builds an orthonormal basis given a normal vector `n`. Return the two axes that are perpendicular to `n`.
39
+
40
+ :param n: A 3D vector (list or array-like) representing the normal vector
41
+ """
42
+ b1 = wp.vec3()
43
+ b2 = wp.vec3()
44
+ if n[2] < 0.0:
45
+ a = 1.0 / (1.0 - n[2])
46
+ b = n[0] * n[1] * a
47
+ b1[0] = 1.0 - n[0] * n[0] * a
48
+ b1[1] = -b
49
+ b1[2] = n[0]
50
+
51
+ b2[0] = b
52
+ b2[1] = n[1] * n[1] * a - 1.0
53
+ b2[2] = -n[1]
54
+ else:
55
+ a = 1.0 / (1.0 + n[2])
56
+ b = -n[0] * n[1] * a
57
+ b1[0] = 1.0 - n[0] * n[0] * a
58
+ b1[1] = b
59
+ b1[2] = -n[0]
60
+
61
+ b2[0] = b
62
+ b2[1] = 1.0 - n[1] * n[1] * a
63
+ b2[2] = -n[1]
64
+
65
+ return b1, b2
66
+
27
67
 
28
68
  @wp.func
29
69
  def triangle_closest_point_barycentric(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
@@ -1527,9 +1567,9 @@ def collide(model, state, edge_sdf_iter: int = 10, iterate_mesh_vertices: bool =
1527
1567
  # generate soft contacts for particles and shapes except ground plane (last shape)
1528
1568
  if model.particle_count and model.shape_count > 1:
1529
1569
  if requires_grad:
1530
- model.soft_contact_body_pos = wp.clone(model.soft_contact_body_pos)
1531
- model.soft_contact_body_vel = wp.clone(model.soft_contact_body_vel)
1532
- model.soft_contact_normal = wp.clone(model.soft_contact_normal)
1570
+ model.soft_contact_body_pos = wp.empty_like(model.soft_contact_body_pos)
1571
+ model.soft_contact_body_vel = wp.empty_like(model.soft_contact_body_vel)
1572
+ model.soft_contact_normal = wp.empty_like(model.soft_contact_normal)
1533
1573
  # clear old count
1534
1574
  model.soft_contact_count.zero_()
1535
1575
  wp.launch(
@@ -1626,12 +1666,12 @@ def collide(model, state, edge_sdf_iter: int = 10, iterate_mesh_vertices: bool =
1626
1666
 
1627
1667
  if model.shape_contact_pair_count or model.ground and model.shape_ground_contact_pair_count:
1628
1668
  if requires_grad:
1629
- model.rigid_contact_point0 = wp.clone(model.rigid_contact_point0)
1630
- model.rigid_contact_point1 = wp.clone(model.rigid_contact_point1)
1631
- model.rigid_contact_offset0 = wp.clone(model.rigid_contact_offset0)
1632
- model.rigid_contact_offset1 = wp.clone(model.rigid_contact_offset1)
1633
- model.rigid_contact_normal = wp.clone(model.rigid_contact_normal)
1634
- model.rigid_contact_thickness = wp.clone(model.rigid_contact_thickness)
1669
+ model.rigid_contact_point0 = wp.empty_like(model.rigid_contact_point0)
1670
+ model.rigid_contact_point1 = wp.empty_like(model.rigid_contact_point1)
1671
+ model.rigid_contact_offset0 = wp.empty_like(model.rigid_contact_offset0)
1672
+ model.rigid_contact_offset1 = wp.empty_like(model.rigid_contact_offset1)
1673
+ model.rigid_contact_normal = wp.empty_like(model.rigid_contact_normal)
1674
+ model.rigid_contact_thickness = wp.empty_like(model.rigid_contact_thickness)
1635
1675
  model.rigid_contact_count = wp.zeros_like(model.rigid_contact_count)
1636
1676
  model.rigid_contact_tids = wp.zeros_like(model.rigid_contact_tids)
1637
1677
  model.rigid_contact_shape0 = wp.empty_like(model.rigid_contact_shape0)
@@ -1679,6 +1719,18 @@ def collide(model, state, edge_sdf_iter: int = 10, iterate_mesh_vertices: bool =
1679
1719
  )
1680
1720
 
1681
1721
 
1722
+ @wp.func
1723
+ def compute_tri_aabb(
1724
+ v1: wp.vec3,
1725
+ v2: wp.vec3,
1726
+ v3: wp.vec3,
1727
+ ):
1728
+ lower = wp.min(wp.min(v1, v2), v3)
1729
+ upper = wp.max(wp.max(v1, v2), v3)
1730
+
1731
+ return lower, upper
1732
+
1733
+
1682
1734
  @wp.kernel
1683
1735
  def compute_tri_aabbs(
1684
1736
  pos: wp.array(dtype=wp.vec3),
@@ -1692,8 +1744,10 @@ def compute_tri_aabbs(
1692
1744
  v2 = pos[tri_indices[t_id, 1]]
1693
1745
  v3 = pos[tri_indices[t_id, 2]]
1694
1746
 
1695
- lower_bounds[t_id] = wp.min(wp.min(v1, v2), v3)
1696
- upper_bounds[t_id] = wp.max(wp.max(v1, v2), v3)
1747
+ upper, lower = compute_tri_aabb(v1, v2, v3)
1748
+
1749
+ lower_bounds[t_id] = lower
1750
+ upper_bounds[t_id] = upper
1697
1751
 
1698
1752
 
1699
1753
  @wp.kernel
@@ -1759,7 +1813,7 @@ def vertex_triangle_collision_detection_kernel(
1759
1813
  pos: wp.array(dtype=wp.vec3),
1760
1814
  tri_indices: wp.array(dtype=wp.int32, ndim=2),
1761
1815
  vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32),
1762
- vertex_colliding_triangles_buffer_size: wp.array(dtype=wp.int32),
1816
+ vertex_colliding_triangles_buffer_sizes: wp.array(dtype=wp.int32),
1763
1817
  triangle_colliding_vertices_offsets: wp.array(dtype=wp.int32),
1764
1818
  triangle_colliding_vertices_buffer_sizes: wp.array(dtype=wp.int32),
1765
1819
  # outputs
@@ -1792,7 +1846,7 @@ def vertex_triangle_collision_detection_kernel(
1792
1846
  vertex_colliding_triangles (array): flattened buffer of vertices' collision triangles
1793
1847
  vertex_colliding_triangles_count (array): number of triangles each vertex collides
1794
1848
  vertex_colliding_triangles_offsets (array): where each vertex' collision buffer starts
1795
- vertex_colliding_triangles_buffer_size (array): size of each vertex' collision buffer, will be modified if resizing is needed
1849
+ vertex_colliding_triangles_buffer_sizes (array): size of each vertex' collision buffer, will be modified if resizing is needed
1796
1850
  vertex_colliding_triangles_min_dist (array): each vertex' min distance to all (non-neighbor) triangles
1797
1851
  triangle_colliding_vertices (array): positions of all the triangles' collision vertices
1798
1852
  triangle_colliding_vertices_count (array): number of triangles each vertex collides
@@ -1837,7 +1891,7 @@ def vertex_triangle_collision_detection_kernel(
1837
1891
  if vertex_num_collisions < vertex_buffer_size:
1838
1892
  vertex_colliding_triangles[vertex_buffer_offset + vertex_num_collisions] = tri_index
1839
1893
  else:
1840
- resize_flags[0] = 1
1894
+ resize_flags[VERTEX_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
1841
1895
 
1842
1896
  vertex_num_collisions = vertex_num_collisions + 1
1843
1897
 
@@ -1850,14 +1904,14 @@ def vertex_triangle_collision_detection_kernel(
1850
1904
  # record v-f collision to triangle
1851
1905
  triangle_colliding_vertices[tri_buffer_offset + tri_num_collisions] = v_index
1852
1906
  else:
1853
- resize_flags[1] = 1
1907
+ resize_flags[TRI_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
1854
1908
 
1855
1909
  vertex_colliding_triangles_count[v_index] = vertex_num_collisions
1856
1910
  vertex_colliding_triangles_min_dist[v_index] = min_dis_to_tris
1857
1911
 
1858
1912
 
1859
1913
  @wp.kernel
1860
- def edge_edge_collision_detection_kernel(
1914
+ def edge_colliding_edges_detection_kernel(
1861
1915
  query_radius: float,
1862
1916
  bvh_id: wp.uint64,
1863
1917
  pos: wp.array(dtype=wp.vec3),
@@ -1876,7 +1930,7 @@ def edge_edge_collision_detection_kernel(
1876
1930
  query_radius (float):
1877
1931
  pos (array): positions of all the vertices that make up edges
1878
1932
  edge_colliding_triangles (array): flattened buffer of edges' collision edges
1879
- edge_colliding_edges_counts (array): number of edges each edge collides
1933
+ edge_colliding_edges_count (array): number of edges each edge collides
1880
1934
  edge_colliding_triangles_offsets (array): where each edge's collision buffer starts
1881
1935
  edge_colliding_triangles_buffer_size (array): size of each edge's collision buffer, will be modified if resizing is needed
1882
1936
  edge_min_dis_to_triangles (array): each vertex' min distance to all (non-neighbor) triangles
@@ -1927,7 +1981,7 @@ def edge_edge_collision_detection_kernel(
1927
1981
  if edge_num_collisions < edge_buffer_size:
1928
1982
  edge_colliding_edges[edge_buffer_offset + edge_num_collisions] = colliding_edge_index
1929
1983
  else:
1930
- resize_flags[1] = 1
1984
+ resize_flags[EDGE_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
1931
1985
 
1932
1986
  edge_num_collisions = edge_num_collisions + 1
1933
1987
 
@@ -1935,19 +1989,134 @@ def edge_edge_collision_detection_kernel(
1935
1989
  edge_colliding_edges_min_dist[e_index] = min_dis_to_edges
1936
1990
 
1937
1991
 
1992
+ @wp.kernel
1993
+ def triangle_triangle_collision_detection_kernel(
1994
+ mesh_id: wp.uint64,
1995
+ pos: wp.array(dtype=wp.vec3),
1996
+ tri_indices: wp.array(dtype=wp.int32, ndim=2),
1997
+ triangle_intersecting_triangles_offsets: wp.array(dtype=wp.int32),
1998
+ # outputs
1999
+ triangle_intersecting_triangles: wp.array(dtype=wp.int32),
2000
+ triangle_intersecting_triangles_count: wp.array(dtype=wp.int32),
2001
+ resize_flags: wp.array(dtype=wp.int32),
2002
+ ):
2003
+ tri_index = wp.tid()
2004
+ t1_v1 = tri_indices[tri_index, 0]
2005
+ t1_v2 = tri_indices[tri_index, 1]
2006
+ t1_v3 = tri_indices[tri_index, 2]
2007
+
2008
+ v1 = pos[t1_v1]
2009
+ v2 = pos[t1_v2]
2010
+ v3 = pos[t1_v3]
2011
+
2012
+ lower, upper = compute_tri_aabb(v1, v2, v3)
2013
+
2014
+ buffer_offset = triangle_intersecting_triangles_offsets[tri_index]
2015
+ buffer_size = triangle_intersecting_triangles_offsets[tri_index + 1] - buffer_offset
2016
+
2017
+ query = wp.mesh_query_aabb(mesh_id, lower, upper)
2018
+ tri_index_2 = wp.int32(0)
2019
+ intersection_count = wp.int32(0)
2020
+ while wp.mesh_query_aabb_next(query, tri_index_2):
2021
+ t2_v1 = tri_indices[tri_index_2, 0]
2022
+ t2_v2 = tri_indices[tri_index_2, 1]
2023
+ t2_v3 = tri_indices[tri_index_2, 2]
2024
+
2025
+ # filter out intersection test with neighbor triangles
2026
+ if (
2027
+ vertex_adjacent_to_triangle(t1_v1, t2_v1, t2_v2, t2_v3)
2028
+ or vertex_adjacent_to_triangle(t1_v2, t2_v1, t2_v2, t2_v3)
2029
+ or vertex_adjacent_to_triangle(t1_v3, t2_v1, t2_v2, t2_v3)
2030
+ ):
2031
+ continue
2032
+
2033
+ u1 = pos[t2_v1]
2034
+ u2 = pos[t2_v2]
2035
+ u3 = pos[t2_v3]
2036
+
2037
+ if wp.intersect_tri_tri(v1, v2, v3, u1, u2, u3):
2038
+ if intersection_count < buffer_size:
2039
+ triangle_intersecting_triangles[buffer_offset + intersection_count] = tri_index_2
2040
+ else:
2041
+ resize_flags[TRI_TRI_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
2042
+ intersection_count = intersection_count + 1
2043
+
2044
+ triangle_intersecting_triangles_count[tri_index] = intersection_count
2045
+
2046
+
2047
+ @wp.struct
2048
+ class TriMeshCollisionInfo:
2049
+ vertex_colliding_triangles: wp.array(dtype=wp.int32)
2050
+ vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32)
2051
+ vertex_colliding_triangles_buffer_sizes: wp.array(dtype=wp.int32)
2052
+ vertex_colliding_triangles_count: wp.array(dtype=wp.int32)
2053
+ vertex_colliding_triangles_min_dist: wp.array(dtype=float)
2054
+
2055
+ triangle_colliding_vertices: wp.array(dtype=wp.int32)
2056
+ triangle_colliding_vertices_offsets: wp.array(dtype=wp.int32)
2057
+ triangle_colliding_vertices_buffer_sizes: wp.array(dtype=wp.int32)
2058
+ triangle_colliding_vertices_count: wp.array(dtype=wp.int32)
2059
+ triangle_colliding_vertices_min_dist: wp.array(dtype=float)
2060
+
2061
+ edge_colliding_edges: wp.array(dtype=wp.int32)
2062
+ edge_colliding_edges_offsets: wp.array(dtype=wp.int32)
2063
+ edge_colliding_edges_buffer_sizes: wp.array(dtype=wp.int32)
2064
+ edge_colliding_edges_count: wp.array(dtype=wp.int32)
2065
+ edge_colliding_edges_min_dist: wp.array(dtype=float)
2066
+
2067
+
2068
+ @wp.func
2069
+ def get_vertex_colliding_triangles_count(col_info: TriMeshCollisionInfo, v: int):
2070
+ return wp.min(col_info.vertex_colliding_triangles_count[v], col_info.vertex_colliding_triangles_buffer_sizes[v])
2071
+
2072
+
2073
+ @wp.func
2074
+ def get_vertex_colliding_triangles(col_info: TriMeshCollisionInfo, v: int, i_collision: int):
2075
+ offset = col_info.vertex_colliding_triangles_offsets[v]
2076
+ return col_info.vertex_colliding_triangles[offset + i_collision]
2077
+
2078
+
2079
+ @wp.func
2080
+ def get_triangle_colliding_vertices_count(col_info: TriMeshCollisionInfo, tri: int):
2081
+ return wp.min(
2082
+ col_info.triangle_colliding_vertices_count[tri], col_info.triangle_colliding_vertices_buffer_sizes[tri]
2083
+ )
2084
+
2085
+
2086
+ @wp.func
2087
+ def get_triangle_colliding_vertices(col_info: TriMeshCollisionInfo, tri: int, i_collision: int):
2088
+ offset = col_info.triangle_colliding_vertices_offsets[tri]
2089
+ return col_info.triangle_colliding_vertices[offset + i_collision]
2090
+
2091
+
2092
+ @wp.func
2093
+ def get_edge_colliding_edges_count(col_info: TriMeshCollisionInfo, e: int):
2094
+ return wp.min(col_info.edge_colliding_edges_count[e], col_info.edge_colliding_edges_buffer_sizes[e])
2095
+
2096
+
2097
+ @wp.func
2098
+ def get_edge_colliding_edges(col_info: TriMeshCollisionInfo, e: int, i_collision: int):
2099
+ offset = col_info.edge_colliding_edges_offsets[e]
2100
+ return col_info.edge_colliding_edges[offset + i_collision]
2101
+
2102
+
1938
2103
  class TriMeshCollisionDetector:
1939
2104
  def __init__(
1940
2105
  self,
1941
2106
  model: Model,
2107
+ vertex_positions=None,
1942
2108
  vertex_collision_buffer_pre_alloc=8,
1943
2109
  vertex_collision_buffer_max_alloc=256,
1944
2110
  triangle_collision_buffer_pre_alloc=16,
1945
2111
  triangle_collision_buffer_max_alloc=256,
1946
2112
  edge_collision_buffer_pre_alloc=8,
1947
2113
  edge_collision_buffer_max_alloc=256,
2114
+ triangle_triangle_collision_buffer_pre_alloc=8,
2115
+ triangle_triangle_collision_buffer_max_alloc=256,
1948
2116
  edge_edge_parallel_epsilon=1e-5,
1949
2117
  ):
1950
2118
  self.model = model
2119
+ self.vertex_positions = model.particle_q if vertex_positions is None else vertex_positions
1951
2120
  self.device = model.device
1952
2121
  self.vertex_collision_buffer_pre_alloc = vertex_collision_buffer_pre_alloc
1953
2122
  self.vertex_collision_buffer_max_alloc = vertex_collision_buffer_max_alloc
@@ -1955,6 +2124,8 @@ class TriMeshCollisionDetector:
1955
2124
  self.triangle_collision_buffer_max_alloc = triangle_collision_buffer_max_alloc
1956
2125
  self.edge_collision_buffer_pre_alloc = edge_collision_buffer_pre_alloc
1957
2126
  self.edge_collision_buffer_max_alloc = edge_collision_buffer_max_alloc
2127
+ self.triangle_triangle_collision_buffer_pre_alloc = triangle_triangle_collision_buffer_pre_alloc
2128
+ self.triangle_triangle_collision_buffer_max_alloc = triangle_triangle_collision_buffer_max_alloc
1958
2129
 
1959
2130
  self.edge_edge_parallel_epsilon = edge_edge_parallel_epsilon
1960
2131
 
@@ -2021,14 +2192,44 @@ class TriMeshCollisionDetector:
2021
2192
  self.upper_bounds_edges = wp.array(shape=(model.edge_count,), dtype=wp.vec3, device=model.device)
2022
2193
  wp.launch(
2023
2194
  kernel=compute_edge_aabbs,
2024
- inputs=[model.particle_q, model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
2195
+ inputs=[self.vertex_positions, model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
2025
2196
  dim=model.edge_count,
2026
2197
  device=model.device,
2027
2198
  )
2028
2199
 
2029
2200
  self.bvh_edges = wp.Bvh(self.lower_bounds_edges, self.upper_bounds_edges)
2030
2201
 
2031
- self.resize_flags = wp.zeros(shape=(3,), dtype=wp.int32, device=self.device)
2202
+ self.resize_flags = wp.zeros(shape=(4,), dtype=wp.int32, device=self.device)
2203
+
2204
+ self.collision_info = self.get_collision_data()
2205
+
2206
+ # data for triangle-triangle intersection; they will only be initialized on demand, as triangle-triangle intersection is not needed for simulation
2207
+ self.triangle_intersecting_triangles = None
2208
+ self.triangle_intersecting_triangles_count = None
2209
+ self.triangle_intersecting_triangles_offsets = None
2210
+
2211
+ def get_collision_data(self):
2212
+ collision_info = TriMeshCollisionInfo()
2213
+
2214
+ collision_info.vertex_colliding_triangles = self.vertex_colliding_triangles
2215
+ collision_info.vertex_colliding_triangles_offsets = self.vertex_colliding_triangles_offsets
2216
+ collision_info.vertex_colliding_triangles_buffer_sizes = self.vertex_colliding_triangles_buffer_sizes
2217
+ collision_info.vertex_colliding_triangles_count = self.vertex_colliding_triangles_count
2218
+ collision_info.vertex_colliding_triangles_min_dist = self.vertex_colliding_triangles_min_dist
2219
+
2220
+ collision_info.triangle_colliding_vertices = self.triangle_colliding_vertices
2221
+ collision_info.triangle_colliding_vertices_offsets = self.triangle_colliding_vertices_offsets
2222
+ collision_info.triangle_colliding_vertices_buffer_sizes = self.triangle_colliding_vertices_buffer_sizes
2223
+ collision_info.triangle_colliding_vertices_count = self.triangle_colliding_vertices_count
2224
+ collision_info.triangle_colliding_vertices_min_dist = self.triangle_colliding_vertices_min_dist
2225
+
2226
+ collision_info.edge_colliding_edges = self.edge_colliding_edges
2227
+ collision_info.edge_colliding_edges_offsets = self.edge_colliding_edges_offsets
2228
+ collision_info.edge_colliding_edges_buffer_sizes = self.edge_colliding_edges_buffer_sizes
2229
+ collision_info.edge_colliding_edges_count = self.edge_colliding_edges_count
2230
+ collision_info.edge_colliding_edges_min_dist = self.edge_colliding_edges_min_dist
2231
+
2232
+ return collision_info
2032
2233
 
2033
2234
  def compute_collision_buffer_offsets(
2034
2235
  self, buffer_sizes: wp.array(dtype=wp.int32), offsets: wp.array(dtype=wp.int32)
@@ -2040,8 +2241,14 @@ class TriMeshCollisionDetector:
2040
2241
 
2041
2242
  offsets.assign(offsets_np)
2042
2243
 
2043
- def refit(self):
2044
- self.refit_triangles()
2244
+ def refit(self, new_pos=None):
2245
+ if new_pos is not None:
2246
+ self.vertex_positions = new_pos
2247
+ # this will automatically apply refit
2248
+ self.mesh_tris.points = new_pos
2249
+ else:
2250
+ self.refit_triangles()
2251
+
2045
2252
  self.refit_edges()
2046
2253
 
2047
2254
  def refit_triangles(self):
@@ -2050,7 +2257,7 @@ class TriMeshCollisionDetector:
2050
2257
  def refit_edges(self):
2051
2258
  wp.launch(
2052
2259
  kernel=compute_edge_aabbs,
2053
- inputs=[self.model.particle_q, self.model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
2260
+ inputs=[self.vertex_positions, self.model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
2054
2261
  dim=self.model.edge_count,
2055
2262
  device=self.model.device,
2056
2263
  )
@@ -2076,7 +2283,7 @@ class TriMeshCollisionDetector:
2076
2283
  inputs=[
2077
2284
  query_radius,
2078
2285
  self.mesh_tris.id,
2079
- self.model.particle_q,
2286
+ self.vertex_positions,
2080
2287
  self.model.tri_indices,
2081
2288
  self.vertex_colliding_triangles_offsets,
2082
2289
  self.vertex_colliding_triangles_buffer_sizes,
@@ -2098,11 +2305,11 @@ class TriMeshCollisionDetector:
2098
2305
 
2099
2306
  def edge_edge_collision_detection(self, query_radius):
2100
2307
  wp.launch(
2101
- kernel=edge_edge_collision_detection_kernel,
2308
+ kernel=edge_colliding_edges_detection_kernel,
2102
2309
  inputs=[
2103
2310
  query_radius,
2104
2311
  self.bvh_edges.id,
2105
- self.model.particle_q,
2312
+ self.vertex_positions,
2106
2313
  self.model.edge_indices,
2107
2314
  self.edge_colliding_edges_offsets,
2108
2315
  self.edge_colliding_edges_buffer_sizes,
@@ -2117,3 +2324,40 @@ class TriMeshCollisionDetector:
2117
2324
  dim=self.model.edge_count,
2118
2325
  device=self.model.device,
2119
2326
  )
2327
+
2328
+ def triangle_triangle_intersection_detection(self):
2329
+ if self.triangle_intersecting_triangles is None:
2330
+ self.triangle_intersecting_triangles = wp.zeros(
2331
+ shape=(self.model.tri_count * self.triangle_triangle_collision_buffer_pre_alloc,),
2332
+ dtype=wp.int32,
2333
+ device=self.device,
2334
+ )
2335
+
2336
+ if self.triangle_intersecting_triangles_count is None:
2337
+ self.triangle_intersecting_triangles_count = wp.array(
2338
+ shape=(self.model.tri_count,), dtype=wp.int32, device=self.device
2339
+ )
2340
+
2341
+ if self.triangle_intersecting_triangles_offsets is None:
2342
+ buffer_sizes = np.full((self.model.tri_count,), self.triangle_triangle_collision_buffer_pre_alloc)
2343
+ offsets = np.zeros((self.model.tri_count + 1,), dtype=np.int32)
2344
+ offsets[1:] = np.cumsum(buffer_sizes)
2345
+
2346
+ self.triangle_intersecting_triangles_offsets = wp.array(offsets, dtype=wp.int32, device=self.device)
2347
+
2348
+ wp.launch(
2349
+ kernel=triangle_triangle_collision_detection_kernel,
2350
+ inputs=[
2351
+ self.mesh_tris.id,
2352
+ self.vertex_positions,
2353
+ self.model.tri_indices,
2354
+ self.triangle_intersecting_triangles_offsets,
2355
+ ],
2356
+ outputs=[
2357
+ self.triangle_intersecting_triangles,
2358
+ self.triangle_intersecting_triangles_count,
2359
+ self.resize_flags,
2360
+ ],
2361
+ dim=self.model.tri_count,
2362
+ device=self.model.device,
2363
+ )
@@ -871,6 +871,7 @@ def eval_rigid_contacts(
871
871
  contact_shape0: wp.array(dtype=int),
872
872
  contact_shape1: wp.array(dtype=int),
873
873
  force_in_world_frame: bool,
874
+ friction_smoothing: float,
874
875
  # outputs
875
876
  body_f: wp.array(dtype=wp.spatial_vector),
876
877
  ):
@@ -924,6 +925,8 @@ def eval_rigid_contacts(
924
925
  n = contact_normal[tid]
925
926
  bx_a = contact_point0[tid]
926
927
  bx_b = contact_point1[tid]
928
+ r_a = wp.vec3(0.0)
929
+ r_b = wp.vec3(0.0)
927
930
  if body_a >= 0:
928
931
  X_wb_a = body_q[body_a]
929
932
  X_com_a = body_com[body_a]
@@ -990,12 +993,16 @@ def eval_rigid_contacts(
990
993
  # ft = wp.vec3(vx, 0.0, vz)
991
994
 
992
995
  # Coulomb friction (smooth, but gradients are numerically unstable around |vt| = 0)
993
- # ft = wp.normalize(vt)*wp.min(kf*wp.length(vt), abs(mu*d*ke))
994
996
  ft = wp.vec3(0.0)
995
997
  if d < 0.0:
996
- ft = wp.normalize(vt) * wp.min(kf * wp.length(vt), -mu * (fn + fd))
998
+ # use a smooth vector norm to avoid gradient instability at/around zero velocity
999
+ vs = wp.norm_huber(vt, delta=friction_smoothing)
1000
+ if vs > 0.0:
1001
+ fr = vt / vs
1002
+ ft = fr * wp.min(kf * vs, -mu * (fn + fd))
997
1003
 
998
1004
  f_total = n * (fn + fd) + ft
1005
+ # f_total = n * (fn + fd)
999
1006
  # f_total = n * fn
1000
1007
 
1001
1008
  if body_a >= 0:
@@ -1761,7 +1768,7 @@ def eval_tetrahedral_forces(model: Model, state: State, control: Control, partic
1761
1768
  )
1762
1769
 
1763
1770
 
1764
- def eval_body_contact_forces(model: Model, state: State, particle_f: wp.array):
1771
+ def eval_body_contact_forces(model: Model, state: State, particle_f: wp.array, friction_smoothing: float = 1.0):
1765
1772
  if model.rigid_contact_max and (
1766
1773
  model.ground and model.shape_ground_contact_pair_count or model.shape_contact_pair_count
1767
1774
  ):
@@ -1782,6 +1789,7 @@ def eval_body_contact_forces(model: Model, state: State, particle_f: wp.array):
1782
1789
  model.rigid_contact_shape0,
1783
1790
  model.rigid_contact_shape1,
1784
1791
  False,
1792
+ friction_smoothing,
1785
1793
  ],
1786
1794
  outputs=[state.body_f],
1787
1795
  device=model.device,
@@ -1880,7 +1888,15 @@ def eval_muscle_forces(model: Model, state: State, control: Control, body_f: wp.
1880
1888
  )
1881
1889
 
1882
1890
 
1883
- def compute_forces(model: Model, state: State, control: Control, particle_f: wp.array, body_f: wp.array, dt: float):
1891
+ def compute_forces(
1892
+ model: Model,
1893
+ state: State,
1894
+ control: Control,
1895
+ particle_f: wp.array,
1896
+ body_f: wp.array,
1897
+ dt: float,
1898
+ friction_smoothing: float = 1.0,
1899
+ ):
1884
1900
  # damped springs
1885
1901
  eval_spring_forces(model, state, particle_f)
1886
1902
 
@@ -1906,7 +1922,7 @@ def compute_forces(model: Model, state: State, control: Control, particle_f: wp.
1906
1922
  eval_particle_ground_contact_forces(model, state, particle_f)
1907
1923
 
1908
1924
  # body contacts
1909
- eval_body_contact_forces(model, state, particle_f)
1925
+ eval_body_contact_forces(model, state, particle_f, friction_smoothing=friction_smoothing)
1910
1926
 
1911
1927
  # particle shape contact
1912
1928
  eval_particle_body_contact_forces(model, state, particle_f, body_f, body_f_in_world_frame=False)
@@ -1941,12 +1957,14 @@ class SemiImplicitIntegrator(Integrator):
1941
1957
 
1942
1958
  """
1943
1959
 
1944
- def __init__(self, angular_damping: float = 0.05):
1960
+ def __init__(self, angular_damping: float = 0.05, friction_smoothing: float = 1.0):
1945
1961
  """
1946
1962
  Args:
1947
1963
  angular_damping (float, optional): Angular damping factor. Defaults to 0.05.
1964
+ friction_smoothing (float, optional): The delta value for the Huber norm (see :func:`warp.math.norm_huber`) used for the friction velocity normalization. Defaults to 1.0.
1948
1965
  """
1949
1966
  self.angular_damping = angular_damping
1967
+ self.friction_smoothing = friction_smoothing
1950
1968
 
1951
1969
  def simulate(self, model: Model, state_in: State, state_out: State, dt: float, control: Control = None):
1952
1970
  with wp.ScopedTimer("simulate", False):
@@ -1962,7 +1980,7 @@ class SemiImplicitIntegrator(Integrator):
1962
1980
  if control is None:
1963
1981
  control = model.control(clone_variables=False)
1964
1982
 
1965
- compute_forces(model, state_in, control, particle_f, body_f, dt)
1983
+ compute_forces(model, state_in, control, particle_f, body_f, dt, friction_smoothing=self.friction_smoothing)
1966
1984
 
1967
1985
  self.integrate_bodies(model, state_in, state_out, dt, self.angular_damping)
1968
1986