warp-lang 1.3.2__py3-none-manylinux2014_aarch64.whl → 1.4.0__py3-none-manylinux2014_aarch64.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 (107) hide show
  1. warp/__init__.py +6 -0
  2. warp/autograd.py +59 -6
  3. warp/bin/warp.so +0 -0
  4. warp/build_dll.py +8 -10
  5. warp/builtins.py +126 -4
  6. warp/codegen.py +435 -53
  7. warp/config.py +1 -1
  8. warp/context.py +678 -403
  9. warp/dlpack.py +2 -0
  10. warp/examples/benchmarks/benchmark_cloth.py +10 -0
  11. warp/examples/core/example_render_opengl.py +12 -10
  12. warp/examples/fem/example_adaptive_grid.py +251 -0
  13. warp/examples/fem/example_apic_fluid.py +1 -1
  14. warp/examples/fem/example_diffusion_3d.py +2 -2
  15. warp/examples/fem/example_magnetostatics.py +1 -1
  16. warp/examples/fem/example_streamlines.py +1 -0
  17. warp/examples/fem/utils.py +23 -4
  18. warp/examples/sim/example_cloth.py +50 -6
  19. warp/fem/__init__.py +2 -0
  20. warp/fem/adaptivity.py +493 -0
  21. warp/fem/field/field.py +2 -1
  22. warp/fem/field/nodal_field.py +18 -26
  23. warp/fem/field/test.py +4 -4
  24. warp/fem/field/trial.py +4 -4
  25. warp/fem/geometry/__init__.py +1 -0
  26. warp/fem/geometry/adaptive_nanogrid.py +843 -0
  27. warp/fem/geometry/nanogrid.py +55 -28
  28. warp/fem/space/__init__.py +1 -1
  29. warp/fem/space/nanogrid_function_space.py +69 -35
  30. warp/fem/utils.py +113 -107
  31. warp/jax_experimental.py +28 -15
  32. warp/native/array.h +0 -1
  33. warp/native/builtin.h +103 -6
  34. warp/native/bvh.cu +2 -0
  35. warp/native/cuda_util.cpp +14 -0
  36. warp/native/cuda_util.h +2 -0
  37. warp/native/error.cpp +4 -2
  38. warp/native/exports.h +99 -17
  39. warp/native/mat.h +97 -0
  40. warp/native/mesh.cpp +36 -0
  41. warp/native/mesh.cu +51 -0
  42. warp/native/mesh.h +1 -0
  43. warp/native/quat.h +43 -0
  44. warp/native/spatial.h +6 -0
  45. warp/native/vec.h +74 -0
  46. warp/native/warp.cpp +2 -1
  47. warp/native/warp.cu +10 -3
  48. warp/native/warp.h +8 -1
  49. warp/paddle.py +382 -0
  50. warp/sim/__init__.py +1 -0
  51. warp/sim/collide.py +519 -0
  52. warp/sim/integrator_euler.py +18 -5
  53. warp/sim/integrator_featherstone.py +5 -5
  54. warp/sim/integrator_vbd.py +1026 -0
  55. warp/sim/model.py +49 -23
  56. warp/stubs.py +459 -0
  57. warp/tape.py +2 -0
  58. warp/tests/aux_test_dependent.py +1 -0
  59. warp/tests/aux_test_name_clash1.py +32 -0
  60. warp/tests/aux_test_name_clash2.py +32 -0
  61. warp/tests/aux_test_square.py +1 -0
  62. warp/tests/test_array.py +222 -0
  63. warp/tests/test_async.py +3 -3
  64. warp/tests/test_atomic.py +6 -0
  65. warp/tests/test_closest_point_edge_edge.py +93 -1
  66. warp/tests/test_codegen.py +62 -15
  67. warp/tests/test_codegen_instancing.py +1457 -0
  68. warp/tests/test_collision.py +486 -0
  69. warp/tests/test_compile_consts.py +3 -28
  70. warp/tests/test_dlpack.py +170 -0
  71. warp/tests/test_examples.py +22 -8
  72. warp/tests/test_fast_math.py +10 -4
  73. warp/tests/test_fem.py +64 -0
  74. warp/tests/test_func.py +46 -0
  75. warp/tests/test_implicit_init.py +49 -0
  76. warp/tests/test_jax.py +58 -0
  77. warp/tests/test_mat.py +84 -0
  78. warp/tests/test_mesh_query_point.py +188 -0
  79. warp/tests/test_module_hashing.py +40 -0
  80. warp/tests/test_multigpu.py +3 -3
  81. warp/tests/test_overwrite.py +8 -0
  82. warp/tests/test_paddle.py +852 -0
  83. warp/tests/test_print.py +89 -0
  84. warp/tests/test_quat.py +111 -0
  85. warp/tests/test_reload.py +31 -1
  86. warp/tests/test_scalar_ops.py +2 -0
  87. warp/tests/test_static.py +412 -0
  88. warp/tests/test_streams.py +64 -3
  89. warp/tests/test_struct.py +4 -4
  90. warp/tests/test_torch.py +24 -0
  91. warp/tests/test_triangle_closest_point.py +137 -0
  92. warp/tests/test_types.py +1 -1
  93. warp/tests/test_vbd.py +386 -0
  94. warp/tests/test_vec.py +143 -0
  95. warp/tests/test_vec_scalar_ops.py +139 -0
  96. warp/tests/test_volume.py +30 -0
  97. warp/tests/unittest_suites.py +12 -0
  98. warp/tests/unittest_utils.py +9 -5
  99. warp/thirdparty/dlpack.py +3 -1
  100. warp/types.py +157 -34
  101. warp/utils.py +37 -14
  102. {warp_lang-1.3.2.dist-info → warp_lang-1.4.0.dist-info}/METADATA +10 -8
  103. {warp_lang-1.3.2.dist-info → warp_lang-1.4.0.dist-info}/RECORD +106 -94
  104. warp/tests/test_point_triangle_closest_point.py +0 -143
  105. {warp_lang-1.3.2.dist-info → warp_lang-1.4.0.dist-info}/LICENSE.md +0 -0
  106. {warp_lang-1.3.2.dist-info → warp_lang-1.4.0.dist-info}/WHEEL +0 -0
  107. {warp_lang-1.3.2.dist-info → warp_lang-1.4.0.dist-info}/top_level.txt +0 -0
warp/sim/__init__.py CHANGED
@@ -16,6 +16,7 @@ from .inertia import transform_inertia
16
16
  from .integrator import Integrator, integrate_bodies, integrate_particles
17
17
  from .integrator_euler import SemiImplicitIntegrator
18
18
  from .integrator_featherstone import FeatherstoneIntegrator
19
+ from .integrator_vbd import VBDIntegrator
19
20
  from .integrator_xpbd import XPBDIntegrator
20
21
  from .model import (
21
22
  GEO_BOX,
warp/sim/collide.py CHANGED
@@ -9,10 +9,21 @@
9
9
  Collision handling functions and kernels.
10
10
  """
11
11
 
12
+ import numpy as np
13
+
12
14
  import warp as wp
15
+ from warp.sim.model import Model
13
16
 
14
17
  from .model import PARTICLE_FLAG_ACTIVE, ModelShapeGeometry
15
18
 
19
+ TRI_CONTACT_FEATURE_VERTEX_A = wp.constant(0)
20
+ TRI_CONTACT_FEATURE_VERTEX_B = wp.constant(1)
21
+ TRI_CONTACT_FEATURE_VERTEX_C = wp.constant(2)
22
+ TRI_CONTACT_FEATURE_EDGE_AB = wp.constant(3)
23
+ TRI_CONTACT_FEATURE_EDGE_AC = wp.constant(4)
24
+ TRI_CONTACT_FEATURE_EDGE_BC = wp.constant(5)
25
+ TRI_CONTACT_FEATURE_FACE_INTERIOR = wp.constant(6)
26
+
16
27
 
17
28
  @wp.func
18
29
  def triangle_closest_point_barycentric(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
@@ -62,6 +73,74 @@ def triangle_closest_point_barycentric(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp
62
73
  return wp.vec3(1.0 - v - w, v, w)
63
74
 
64
75
 
76
+ @wp.func
77
+ def triangle_closest_point(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
78
+ """
79
+ feature_type type:
80
+ TRI_CONTACT_FEATURE_VERTEX_A
81
+ TRI_CONTACT_FEATURE_VERTEX_B
82
+ TRI_CONTACT_FEATURE_VERTEX_C
83
+ TRI_CONTACT_FEATURE_EDGE_AB : at edge A-B
84
+ TRI_CONTACT_FEATURE_EDGE_AC : at edge A-C
85
+ TRI_CONTACT_FEATURE_EDGE_BC : at edge B-C
86
+ TRI_CONTACT_FEATURE_FACE_INTERIOR
87
+ """
88
+ ab = b - a
89
+ ac = c - a
90
+ ap = p - a
91
+
92
+ d1 = wp.dot(ab, ap)
93
+ d2 = wp.dot(ac, ap)
94
+ if d1 <= 0.0 and d2 <= 0.0:
95
+ feature_type = TRI_CONTACT_FEATURE_VERTEX_A
96
+ bary = wp.vec3(1.0, 0.0, 0.0)
97
+ return a, bary, feature_type
98
+
99
+ bp = p - b
100
+ d3 = wp.dot(ab, bp)
101
+ d4 = wp.dot(ac, bp)
102
+ if d3 >= 0.0 and d4 <= d3:
103
+ feature_type = TRI_CONTACT_FEATURE_VERTEX_B
104
+ bary = wp.vec3(0.0, 1.0, 0.0)
105
+ return b, bary, feature_type
106
+
107
+ cp = p - c
108
+ d5 = wp.dot(ab, cp)
109
+ d6 = wp.dot(ac, cp)
110
+ if d6 >= 0.0 and d5 <= d6:
111
+ feature_type = TRI_CONTACT_FEATURE_VERTEX_C
112
+ bary = wp.vec3(0.0, 0.0, 1.0)
113
+ return c, bary, feature_type
114
+
115
+ vc = d1 * d4 - d3 * d2
116
+ if vc <= 0.0 and d1 >= 0.0 and d3 <= 0.0:
117
+ v = d1 / (d1 - d3)
118
+ feature_type = TRI_CONTACT_FEATURE_EDGE_AB
119
+ bary = wp.vec3(1.0 - v, v, 0.0)
120
+ return a + v * ab, bary, feature_type
121
+
122
+ vb = d5 * d2 - d1 * d6
123
+ if vb <= 0.0 and d2 >= 0.0 and d6 <= 0.0:
124
+ v = d2 / (d2 - d6)
125
+ feature_type = TRI_CONTACT_FEATURE_EDGE_AC
126
+ bary = wp.vec3(1.0 - v, 0.0, v)
127
+ return a + v * ac, bary, feature_type
128
+
129
+ va = d3 * d6 - d5 * d4
130
+ if va <= 0.0 and (d4 - d3) >= 0.0 and (d5 - d6) >= 0.0:
131
+ v = (d4 - d3) / ((d4 - d3) + (d5 - d6))
132
+ feature_type = TRI_CONTACT_FEATURE_EDGE_BC
133
+ bary = wp.vec3(0.0, 1.0 - v, v)
134
+ return b + v * (c - b), bary, feature_type
135
+
136
+ denom = 1.0 / (va + vb + vc)
137
+ v = vb * denom
138
+ w = vc * denom
139
+ feature_type = TRI_CONTACT_FEATURE_FACE_INTERIOR
140
+ bary = wp.vec3(1.0 - v - w, v, w)
141
+ return a + v * ab + w * ac, bary, feature_type
142
+
143
+
65
144
  @wp.func
66
145
  def sphere_sdf(center: wp.vec3, radius: float, p: wp.vec3):
67
146
  return wp.length(p - center) - radius
@@ -1598,3 +1677,443 @@ def collide(model, state, edge_sdf_iter: int = 10, iterate_mesh_vertices: bool =
1598
1677
  ],
1599
1678
  device=model.device,
1600
1679
  )
1680
+
1681
+
1682
+ @wp.kernel
1683
+ def compute_tri_aabbs(
1684
+ pos: wp.array(dtype=wp.vec3),
1685
+ tri_indices: wp.array(dtype=wp.int32, ndim=2),
1686
+ lower_bounds: wp.array(dtype=wp.vec3),
1687
+ upper_bounds: wp.array(dtype=wp.vec3),
1688
+ ):
1689
+ t_id = wp.tid()
1690
+
1691
+ v1 = pos[tri_indices[t_id, 0]]
1692
+ v2 = pos[tri_indices[t_id, 1]]
1693
+ v3 = pos[tri_indices[t_id, 2]]
1694
+
1695
+ lower_bounds[t_id] = wp.min(wp.min(v1, v2), v3)
1696
+ upper_bounds[t_id] = wp.max(wp.max(v1, v2), v3)
1697
+
1698
+
1699
+ @wp.kernel
1700
+ def compute_edge_aabbs(
1701
+ pos: wp.array(dtype=wp.vec3),
1702
+ edge_indices: wp.array(dtype=wp.int32, ndim=2),
1703
+ lower_bounds: wp.array(dtype=wp.vec3),
1704
+ upper_bounds: wp.array(dtype=wp.vec3),
1705
+ ):
1706
+ e_id = wp.tid()
1707
+
1708
+ v1 = pos[edge_indices[e_id, 2]]
1709
+ v2 = pos[edge_indices[e_id, 3]]
1710
+
1711
+ lower_bounds[e_id] = wp.min(v1, v2)
1712
+ upper_bounds[e_id] = wp.max(v1, v2)
1713
+
1714
+
1715
+ @wp.func
1716
+ def tri_is_neighbor(a_1: wp.int32, a_2: wp.int32, a_3: wp.int32, b_1: wp.int32, b_2: wp.int32, b_3: wp.int32):
1717
+ tri_is_neighbor = (
1718
+ a_1 == b_1
1719
+ or a_1 == b_2
1720
+ or a_1 == b_3
1721
+ or a_2 == b_1
1722
+ or a_2 == b_2
1723
+ or a_2 == b_3
1724
+ or a_3 == b_1
1725
+ or a_3 == b_2
1726
+ or a_3 == b_3
1727
+ )
1728
+
1729
+ return tri_is_neighbor
1730
+
1731
+
1732
+ @wp.func
1733
+ def vertex_adjacent_to_triangle(v: wp.int32, a: wp.int32, b: wp.int32, c: wp.int32):
1734
+ return v == a or v == b or v == c
1735
+
1736
+
1737
+ @wp.kernel
1738
+ def init_triangle_collision_data_kernel(
1739
+ query_radius: float,
1740
+ # outputs
1741
+ triangle_colliding_vertices_count: wp.array(dtype=wp.int32),
1742
+ triangle_colliding_vertices_min_dist: wp.array(dtype=float),
1743
+ resize_flags: wp.array(dtype=wp.int32),
1744
+ ):
1745
+ tri_index = wp.tid()
1746
+
1747
+ triangle_colliding_vertices_count[tri_index] = 0
1748
+ triangle_colliding_vertices_min_dist[tri_index] = query_radius
1749
+
1750
+ if tri_index == 0:
1751
+ for i in range(3):
1752
+ resize_flags[i] = 0
1753
+
1754
+
1755
+ @wp.kernel
1756
+ def vertex_triangle_collision_detection_kernel(
1757
+ query_radius: float,
1758
+ mesh_id: wp.uint64,
1759
+ pos: wp.array(dtype=wp.vec3),
1760
+ tri_indices: wp.array(dtype=wp.int32, ndim=2),
1761
+ vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32),
1762
+ vertex_colliding_triangles_buffer_size: wp.array(dtype=wp.int32),
1763
+ triangle_colliding_vertices_offsets: wp.array(dtype=wp.int32),
1764
+ triangle_colliding_vertices_buffer_sizes: wp.array(dtype=wp.int32),
1765
+ # outputs
1766
+ vertex_colliding_triangles: wp.array(dtype=wp.int32),
1767
+ vertex_colliding_triangles_count: wp.array(dtype=wp.int32),
1768
+ vertex_colliding_triangles_min_dist: wp.array(dtype=float),
1769
+ triangle_colliding_vertices: wp.array(dtype=wp.int32),
1770
+ triangle_colliding_vertices_count: wp.array(dtype=wp.int32),
1771
+ triangle_colliding_vertices_min_dist: wp.array(dtype=float),
1772
+ resize_flags: wp.array(dtype=wp.int32),
1773
+ ):
1774
+ """
1775
+ This function applies discrete collision detection between vertices and triangles. It uses pre-allocated spaces to
1776
+ record the collision data. This collision detector works both ways, i.e., it records vertices' colliding triangles to
1777
+ vertex_colliding_triangles, and records each triangles colliding vertices to triangle_colliding_vertices.
1778
+
1779
+ This function assumes that all the vertices are on triangles, and can be indexed from the pos argument.
1780
+
1781
+ Note:
1782
+
1783
+ The collision date buffer is pre-allocated and cannot be changed during collision detection, therefore, the space
1784
+ may not be enough. If the space is not enough to record all the collision information, the function will set a
1785
+ certain element in resized_flag to be true. The user can reallocate the buffer based on vertex_colliding_triangles_count
1786
+ and vertex_colliding_triangles_count.
1787
+
1788
+ Attributes:
1789
+ mesh_id (int): the mesh id you want to collide with
1790
+ query_radius (float): the contact radius. vertex-triangle pairs whose distance are less than this will get detected
1791
+ pos (array): positions of all the vertices that make up triangles
1792
+ vertex_colliding_triangles (array): flattened buffer of vertices' collision triangles
1793
+ vertex_colliding_triangles_count (array): number of triangles each vertex collides
1794
+ 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
1796
+ vertex_colliding_triangles_min_dist (array): each vertex' min distance to all (non-neighbor) triangles
1797
+ triangle_colliding_vertices (array): positions of all the triangles' collision vertices
1798
+ triangle_colliding_vertices_count (array): number of triangles each vertex collides
1799
+ triangle_colliding_vertices_offsets (array): where each triangle's collision buffer starts
1800
+ triangle_colliding_vertices_buffer_sizes (array): size of each triangle's collision buffer, will be modified if resizing is needed
1801
+ triangle_colliding_vertices_min_dist (array): each triangle's min distance to all (non-self) vertices
1802
+ resized_flag (array): size == 3, (vertex_buffer_resize_required, triangle_buffer_resize_required, edge_buffer_resize_required)
1803
+ """
1804
+
1805
+ v_index = wp.tid()
1806
+ v = pos[v_index]
1807
+
1808
+ lower = wp.vec3(v[0] - query_radius, v[1] - query_radius, v[2] - query_radius)
1809
+ upper = wp.vec3(v[0] + query_radius, v[1] + query_radius, v[2] + query_radius)
1810
+
1811
+ query = wp.mesh_query_aabb(mesh_id, lower, upper)
1812
+
1813
+ tri_index = wp.int32(0)
1814
+ vertex_num_collisions = wp.int32(0)
1815
+ min_dis_to_tris = query_radius
1816
+ while wp.mesh_query_aabb_next(query, tri_index):
1817
+ t1 = tri_indices[tri_index, 0]
1818
+ t2 = tri_indices[tri_index, 1]
1819
+ t3 = tri_indices[tri_index, 2]
1820
+ if vertex_adjacent_to_triangle(v_index, t1, t2, t3):
1821
+ continue
1822
+
1823
+ u1 = pos[t1]
1824
+ u2 = pos[t2]
1825
+ u3 = pos[t3]
1826
+
1827
+ closest_p, bary, feature_type = triangle_closest_point(u1, u2, u3, v)
1828
+
1829
+ dist = wp.length(closest_p - v)
1830
+
1831
+ if dist < query_radius:
1832
+ vertex_buffer_offset = vertex_colliding_triangles_offsets[v_index]
1833
+ vertex_buffer_size = vertex_colliding_triangles_offsets[v_index + 1] - vertex_buffer_offset
1834
+
1835
+ # record v-f collision to vertex
1836
+ min_dis_to_tris = wp.min(min_dis_to_tris, dist)
1837
+ if vertex_num_collisions < vertex_buffer_size:
1838
+ vertex_colliding_triangles[vertex_buffer_offset + vertex_num_collisions] = tri_index
1839
+ else:
1840
+ resize_flags[0] = 1
1841
+
1842
+ vertex_num_collisions = vertex_num_collisions + 1
1843
+
1844
+ wp.atomic_min(triangle_colliding_vertices_min_dist, tri_index, dist)
1845
+ tri_buffer_size = triangle_colliding_vertices_buffer_sizes[tri_index]
1846
+ tri_num_collisions = wp.atomic_add(triangle_colliding_vertices_count, tri_index, 1)
1847
+
1848
+ if tri_num_collisions < tri_buffer_size:
1849
+ tri_buffer_offset = triangle_colliding_vertices_offsets[tri_index]
1850
+ # record v-f collision to triangle
1851
+ triangle_colliding_vertices[tri_buffer_offset + tri_num_collisions] = v_index
1852
+ else:
1853
+ resize_flags[1] = 1
1854
+
1855
+ vertex_colliding_triangles_count[v_index] = vertex_num_collisions
1856
+ vertex_colliding_triangles_min_dist[v_index] = min_dis_to_tris
1857
+
1858
+
1859
+ @wp.kernel
1860
+ def edge_edge_collision_detection_kernel(
1861
+ query_radius: float,
1862
+ bvh_id: wp.uint64,
1863
+ pos: wp.array(dtype=wp.vec3),
1864
+ edge_indices: wp.array(dtype=wp.int32, ndim=2),
1865
+ edge_colliding_edges_offsets: wp.array(dtype=wp.int32),
1866
+ edge_colliding_edges_buffer_sizes: wp.array(dtype=wp.int32),
1867
+ edge_edge_parallel_epsilon: float,
1868
+ # outputs
1869
+ edge_colliding_edges: wp.array(dtype=wp.int32),
1870
+ edge_colliding_edges_count: wp.array(dtype=wp.int32),
1871
+ edge_colliding_edges_min_dist: wp.array(dtype=float),
1872
+ resize_flags: wp.array(dtype=wp.int32),
1873
+ ):
1874
+ """
1875
+ bvh_id (int): the bvh id you want to do collision detection on
1876
+ query_radius (float):
1877
+ pos (array): positions of all the vertices that make up edges
1878
+ edge_colliding_triangles (array): flattened buffer of edges' collision edges
1879
+ edge_colliding_edges_counts (array): number of edges each edge collides
1880
+ edge_colliding_triangles_offsets (array): where each edge's collision buffer starts
1881
+ edge_colliding_triangles_buffer_size (array): size of each edge's collision buffer, will be modified if resizing is needed
1882
+ edge_min_dis_to_triangles (array): each vertex' min distance to all (non-neighbor) triangles
1883
+ resized_flag (array): size == 3, (vertex_buffer_resize_required, triangle_buffer_resize_required, edge_buffer_resize_required)
1884
+ """
1885
+ e_index = wp.tid()
1886
+
1887
+ e0_v0 = edge_indices[e_index, 2]
1888
+ e0_v1 = edge_indices[e_index, 3]
1889
+
1890
+ e0_v0_pos = pos[e0_v0]
1891
+ e0_v1_pos = pos[e0_v1]
1892
+
1893
+ lower = wp.min(e0_v0_pos, e0_v1_pos)
1894
+ upper = wp.max(e0_v0_pos, e0_v1_pos)
1895
+
1896
+ lower = wp.vec3(lower[0] - query_radius, lower[1] - query_radius, lower[2] - query_radius)
1897
+ upper = wp.vec3(upper[0] + query_radius, upper[1] + query_radius, upper[2] + query_radius)
1898
+
1899
+ query = wp.bvh_query_aabb(bvh_id, lower, upper)
1900
+
1901
+ colliding_edge_index = wp.int32(0)
1902
+ edge_num_collisions = wp.int32(0)
1903
+ min_dis_to_edges = query_radius
1904
+ while wp.bvh_query_next(query, colliding_edge_index):
1905
+ e1_v0 = edge_indices[colliding_edge_index, 2]
1906
+ e1_v1 = edge_indices[colliding_edge_index, 3]
1907
+
1908
+ if e0_v0 == e1_v0 or e0_v0 == e1_v1 or e0_v1 == e1_v0 or e0_v1 == e1_v1:
1909
+ continue
1910
+
1911
+ e1_v0_pos = pos[e1_v0]
1912
+ e1_v1_pos = pos[e1_v1]
1913
+
1914
+ st = wp.closest_point_edge_edge(e0_v0_pos, e0_v1_pos, e1_v0_pos, e1_v1_pos, edge_edge_parallel_epsilon)
1915
+ s = st[0]
1916
+ t = st[1]
1917
+ c1 = e0_v0_pos + (e0_v1_pos - e0_v0_pos) * s
1918
+ c2 = e1_v0_pos + (e1_v1_pos - e1_v0_pos) * t
1919
+
1920
+ dist = wp.length(c1 - c2)
1921
+ if dist < query_radius:
1922
+ edge_buffer_offset = edge_colliding_edges_offsets[e_index]
1923
+ edge_buffer_size = edge_colliding_edges_offsets[e_index + 1] - edge_buffer_offset
1924
+
1925
+ # record e-e collision to e0, and leave e1; e1 will detect this collision from its own thread
1926
+ min_dis_to_edges = wp.min(min_dis_to_edges, dist)
1927
+ if edge_num_collisions < edge_buffer_size:
1928
+ edge_colliding_edges[edge_buffer_offset + edge_num_collisions] = colliding_edge_index
1929
+ else:
1930
+ resize_flags[1] = 1
1931
+
1932
+ edge_num_collisions = edge_num_collisions + 1
1933
+
1934
+ edge_colliding_edges_count[e_index] = edge_num_collisions
1935
+ edge_colliding_edges_min_dist[e_index] = min_dis_to_edges
1936
+
1937
+
1938
+ class TriMeshCollisionDetector:
1939
+ def __init__(
1940
+ self,
1941
+ model: Model,
1942
+ vertex_collision_buffer_pre_alloc=8,
1943
+ vertex_collision_buffer_max_alloc=256,
1944
+ triangle_collision_buffer_pre_alloc=16,
1945
+ triangle_collision_buffer_max_alloc=256,
1946
+ edge_collision_buffer_pre_alloc=8,
1947
+ edge_collision_buffer_max_alloc=256,
1948
+ edge_edge_parallel_epsilon=1e-5,
1949
+ ):
1950
+ self.model = model
1951
+ self.device = model.device
1952
+ self.vertex_collision_buffer_pre_alloc = vertex_collision_buffer_pre_alloc
1953
+ self.vertex_collision_buffer_max_alloc = vertex_collision_buffer_max_alloc
1954
+ self.triangle_collision_buffer_pre_alloc = triangle_collision_buffer_pre_alloc
1955
+ self.triangle_collision_buffer_max_alloc = triangle_collision_buffer_max_alloc
1956
+ self.edge_collision_buffer_pre_alloc = edge_collision_buffer_pre_alloc
1957
+ self.edge_collision_buffer_max_alloc = edge_collision_buffer_max_alloc
1958
+
1959
+ self.edge_edge_parallel_epsilon = edge_edge_parallel_epsilon
1960
+
1961
+ self.mesh_tris = wp.Mesh(model.particle_q, model.tri_indices.flatten())
1962
+
1963
+ # collision detections results
1964
+
1965
+ # vertex collision buffers
1966
+ self.vertex_colliding_triangles = wp.zeros(
1967
+ shape=(model.particle_count * self.vertex_collision_buffer_pre_alloc,), dtype=wp.int32, device=self.device
1968
+ )
1969
+ self.vertex_colliding_triangles_count = wp.array(
1970
+ shape=(model.particle_count,), dtype=wp.int32, device=self.device
1971
+ )
1972
+ self.vertex_colliding_triangles_min_dist = wp.array(
1973
+ shape=(model.particle_count,), dtype=float, device=self.device
1974
+ )
1975
+ self.vertex_colliding_triangles_buffer_sizes = wp.full(
1976
+ shape=(model.particle_count,),
1977
+ value=self.vertex_collision_buffer_pre_alloc,
1978
+ dtype=wp.int32,
1979
+ device=self.device,
1980
+ )
1981
+ self.vertex_colliding_triangles_offsets = wp.array(
1982
+ shape=(model.particle_count + 1,), dtype=wp.int32, device=self.device
1983
+ )
1984
+ self.compute_collision_buffer_offsets(
1985
+ self.vertex_colliding_triangles_buffer_sizes, self.vertex_colliding_triangles_offsets
1986
+ )
1987
+
1988
+ # triangle collision buffers
1989
+ self.triangle_colliding_vertices = wp.zeros(
1990
+ shape=(model.tri_count * self.triangle_collision_buffer_pre_alloc,), dtype=wp.int32, device=self.device
1991
+ )
1992
+ self.triangle_colliding_vertices_count = wp.zeros(shape=(model.tri_count,), dtype=wp.int32, device=self.device)
1993
+ self.triangle_colliding_vertices_buffer_sizes = wp.full(
1994
+ shape=(model.tri_count,), value=self.triangle_collision_buffer_pre_alloc, dtype=wp.int32, device=self.device
1995
+ )
1996
+ self.triangle_colliding_vertices_min_dist = wp.array(shape=(model.tri_count,), dtype=float, device=self.device)
1997
+
1998
+ self.triangle_colliding_vertices_offsets = wp.array(
1999
+ shape=(model.tri_count + 1,), dtype=wp.int32, device=self.device
2000
+ )
2001
+ self.compute_collision_buffer_offsets(
2002
+ self.triangle_colliding_vertices_buffer_sizes, self.triangle_colliding_vertices_offsets
2003
+ )
2004
+
2005
+ # edge collision buffers
2006
+ self.edge_colliding_edges = wp.zeros(
2007
+ shape=(model.edge_count * self.edge_collision_buffer_pre_alloc,), dtype=wp.int32, device=self.device
2008
+ )
2009
+ self.edge_colliding_edges_count = wp.zeros(shape=(model.edge_count,), dtype=wp.int32, device=self.device)
2010
+ self.edge_colliding_edges_buffer_sizes = wp.full(
2011
+ shape=(model.edge_count,),
2012
+ value=self.edge_collision_buffer_pre_alloc,
2013
+ dtype=wp.int32,
2014
+ device=self.device,
2015
+ )
2016
+ self.edge_colliding_edges_offsets = wp.array(shape=(model.edge_count + 1,), dtype=wp.int32, device=self.device)
2017
+ self.compute_collision_buffer_offsets(self.edge_colliding_edges_buffer_sizes, self.edge_colliding_edges_offsets)
2018
+ self.edge_colliding_edges_min_dist = wp.array(shape=(model.edge_count,), dtype=float, device=self.device)
2019
+
2020
+ self.lower_bounds_edges = wp.array(shape=(model.edge_count,), dtype=wp.vec3, device=model.device)
2021
+ self.upper_bounds_edges = wp.array(shape=(model.edge_count,), dtype=wp.vec3, device=model.device)
2022
+ wp.launch(
2023
+ kernel=compute_edge_aabbs,
2024
+ inputs=[model.particle_q, model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
2025
+ dim=model.edge_count,
2026
+ device=model.device,
2027
+ )
2028
+
2029
+ self.bvh_edges = wp.Bvh(self.lower_bounds_edges, self.upper_bounds_edges)
2030
+
2031
+ self.resize_flags = wp.zeros(shape=(3,), dtype=wp.int32, device=self.device)
2032
+
2033
+ def compute_collision_buffer_offsets(
2034
+ self, buffer_sizes: wp.array(dtype=wp.int32), offsets: wp.array(dtype=wp.int32)
2035
+ ):
2036
+ assert offsets.size == buffer_sizes.size + 1
2037
+ offsets_np = np.empty(shape=(offsets.size,), dtype=np.int32)
2038
+ offsets_np[1:] = np.cumsum(buffer_sizes.numpy())[:]
2039
+ offsets_np[0] = 0
2040
+
2041
+ offsets.assign(offsets_np)
2042
+
2043
+ def refit(self):
2044
+ self.refit_triangles()
2045
+ self.refit_edges()
2046
+
2047
+ def refit_triangles(self):
2048
+ self.mesh_tris.refit()
2049
+
2050
+ def refit_edges(self):
2051
+ wp.launch(
2052
+ kernel=compute_edge_aabbs,
2053
+ inputs=[self.model.particle_q, self.model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
2054
+ dim=self.model.edge_count,
2055
+ device=self.model.device,
2056
+ )
2057
+ self.bvh_edges.refit()
2058
+
2059
+ def vertex_triangle_collision_detection(self, query_radius):
2060
+ wp.launch(
2061
+ kernel=init_triangle_collision_data_kernel,
2062
+ inputs=[
2063
+ query_radius,
2064
+ ],
2065
+ outputs=[
2066
+ self.triangle_colliding_vertices_count,
2067
+ self.triangle_colliding_vertices_min_dist,
2068
+ self.resize_flags,
2069
+ ],
2070
+ dim=self.model.tri_count,
2071
+ device=self.model.device,
2072
+ )
2073
+
2074
+ wp.launch(
2075
+ kernel=vertex_triangle_collision_detection_kernel,
2076
+ inputs=[
2077
+ query_radius,
2078
+ self.mesh_tris.id,
2079
+ self.model.particle_q,
2080
+ self.model.tri_indices,
2081
+ self.vertex_colliding_triangles_offsets,
2082
+ self.vertex_colliding_triangles_buffer_sizes,
2083
+ self.triangle_colliding_vertices_offsets,
2084
+ self.triangle_colliding_vertices_buffer_sizes,
2085
+ ],
2086
+ outputs=[
2087
+ self.vertex_colliding_triangles,
2088
+ self.vertex_colliding_triangles_count,
2089
+ self.vertex_colliding_triangles_min_dist,
2090
+ self.triangle_colliding_vertices,
2091
+ self.triangle_colliding_vertices_count,
2092
+ self.triangle_colliding_vertices_min_dist,
2093
+ self.resize_flags,
2094
+ ],
2095
+ dim=self.model.particle_count,
2096
+ device=self.model.device,
2097
+ )
2098
+
2099
+ def edge_edge_collision_detection(self, query_radius):
2100
+ wp.launch(
2101
+ kernel=edge_edge_collision_detection_kernel,
2102
+ inputs=[
2103
+ query_radius,
2104
+ self.bvh_edges.id,
2105
+ self.model.particle_q,
2106
+ self.model.edge_indices,
2107
+ self.edge_colliding_edges_offsets,
2108
+ self.edge_colliding_edges_buffer_sizes,
2109
+ self.edge_edge_parallel_epsilon,
2110
+ ],
2111
+ outputs=[
2112
+ self.edge_colliding_edges,
2113
+ self.edge_colliding_edges_count,
2114
+ self.edge_colliding_edges_min_dist,
2115
+ self.resize_flags,
2116
+ ],
2117
+ dim=self.model.edge_count,
2118
+ device=self.model.device,
2119
+ )
@@ -450,6 +450,9 @@ def eval_bending(
450
450
  k = indices[tid, 2]
451
451
  l = indices[tid, 3]
452
452
 
453
+ if i == -1 or j == -1 or k == -1 or l == -1:
454
+ return
455
+
453
456
  rest_angle = rest[tid]
454
457
 
455
458
  x1 = x[i]
@@ -758,6 +761,7 @@ def eval_particle_contacts(
758
761
  contact_body_vel: wp.array(dtype=wp.vec3),
759
762
  contact_normal: wp.array(dtype=wp.vec3),
760
763
  contact_max: int,
764
+ body_f_in_world_frame: bool,
761
765
  # outputs
762
766
  particle_f: wp.array(dtype=wp.vec3),
763
767
  body_f: wp.array(dtype=wp.spatial_vector),
@@ -806,7 +810,11 @@ def eval_particle_contacts(
806
810
  body_v = wp.spatial_bottom(body_v_s)
807
811
 
808
812
  # compute the body velocity at the particle position
809
- bv = body_v + wp.cross(body_w, r) + wp.transform_vector(X_wb, contact_body_vel[tid])
813
+ bv = body_v + wp.transform_vector(X_wb, contact_body_vel[tid])
814
+ if body_f_in_world_frame:
815
+ bv += wp.cross(body_w, bx)
816
+ else:
817
+ bv += wp.cross(body_w, r)
810
818
 
811
819
  # relative velocity
812
820
  v = pv - bv
@@ -837,12 +845,14 @@ def eval_particle_contacts(
837
845
  ft = wp.normalize(vt) * wp.min(kf * wp.length(vt), abs(mu * c * ke))
838
846
 
839
847
  f_total = fn + (fd + ft)
840
- t_total = wp.cross(r, f_total)
841
848
 
842
849
  wp.atomic_sub(particle_f, particle_index, f_total)
843
850
 
844
851
  if body_index >= 0:
845
- wp.atomic_add(body_f, body_index, wp.spatial_vector(t_total, f_total))
852
+ if body_f_in_world_frame:
853
+ wp.atomic_sub(body_f, body_index, wp.spatial_vector(wp.cross(bx, f_total), f_total))
854
+ else:
855
+ wp.atomic_add(body_f, body_index, wp.spatial_vector(wp.cross(r, f_total), f_total))
846
856
 
847
857
 
848
858
  @wp.kernel
@@ -1811,7 +1821,9 @@ def eval_body_joint_forces(model: Model, state: State, control: Control, body_f:
1811
1821
  )
1812
1822
 
1813
1823
 
1814
- def eval_particle_body_contact_forces(model: Model, state: State, particle_f: wp.array, body_f: wp.array):
1824
+ def eval_particle_body_contact_forces(
1825
+ model: Model, state: State, particle_f: wp.array, body_f: wp.array, body_f_in_world_frame: bool = False
1826
+ ):
1815
1827
  if model.particle_count and model.shape_count > 1:
1816
1828
  wp.launch(
1817
1829
  kernel=eval_particle_contacts,
@@ -1838,6 +1850,7 @@ def eval_particle_body_contact_forces(model: Model, state: State, particle_f: wp
1838
1850
  model.soft_contact_body_vel,
1839
1851
  model.soft_contact_normal,
1840
1852
  model.soft_contact_max,
1853
+ body_f_in_world_frame,
1841
1854
  ],
1842
1855
  # outputs
1843
1856
  outputs=[particle_f, body_f],
@@ -1894,7 +1907,7 @@ def compute_forces(model: Model, state: State, control: Control, particle_f: wp.
1894
1907
  eval_body_contact_forces(model, state, particle_f)
1895
1908
 
1896
1909
  # particle shape contact
1897
- eval_particle_body_contact_forces(model, state, particle_f, body_f)
1910
+ eval_particle_body_contact_forces(model, state, particle_f, body_f, body_f_in_world_frame=False)
1898
1911
 
1899
1912
  # muscles
1900
1913
  if False:
@@ -592,7 +592,7 @@ def jcalc_integrate(
592
592
  p_s = wp.vec3(joint_q[coord_start + 0], joint_q[coord_start + 1], joint_q[coord_start + 2])
593
593
 
594
594
  # linear vel of origin (note q/qd switch order of linear angular elements)
595
- # note we are converting the body twist in the space frame (w_s, v_s) to compute center of mass velcity
595
+ # note we are converting the body twist in the space frame (w_s, v_s) to compute center of mass velocity
596
596
  dpdt_s = v_s + wp.cross(w_s, p_s)
597
597
 
598
598
  # quat and quat derivative
@@ -1015,7 +1015,7 @@ def eval_rigid_tau(
1015
1015
  def eval_rigid_jacobian(
1016
1016
  articulation_start: wp.array(dtype=int),
1017
1017
  articulation_J_start: wp.array(dtype=int),
1018
- joint_parent: wp.array(dtype=int),
1018
+ joint_ancestor: wp.array(dtype=int),
1019
1019
  joint_qd_start: wp.array(dtype=int),
1020
1020
  joint_S_s: wp.array(dtype=wp.spatial_vector),
1021
1021
  # outputs
@@ -1051,7 +1051,7 @@ def eval_rigid_jacobian(
1051
1051
  for k in range(6):
1052
1052
  J[J_offset + dense_index(articulation_dof_count, row_start + k, col)] = S[k]
1053
1053
 
1054
- j = joint_parent[j]
1054
+ j = joint_ancestor[j]
1055
1055
 
1056
1056
 
1057
1057
  @wp.func
@@ -1621,7 +1621,7 @@ class FeatherstoneIntegrator(Integrator):
1621
1621
  eval_particle_ground_contact_forces(model, state_in, particle_f)
1622
1622
 
1623
1623
  # particle shape contact
1624
- eval_particle_body_contact_forces(model, state_in, particle_f, body_f)
1624
+ eval_particle_body_contact_forces(model, state_in, particle_f, body_f, body_f_in_world_frame=True)
1625
1625
 
1626
1626
  # muscles
1627
1627
  if False:
@@ -1769,7 +1769,7 @@ class FeatherstoneIntegrator(Integrator):
1769
1769
  inputs=[
1770
1770
  model.articulation_start,
1771
1771
  self.articulation_J_start,
1772
- model.joint_parent,
1772
+ model.joint_ancestor,
1773
1773
  model.joint_qd_start,
1774
1774
  state_aug.joint_S_s,
1775
1775
  ],