warp-lang 1.7.2__py3-none-macosx_10_13_universal2.whl → 1.8.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 (181) hide show
  1. warp/__init__.py +3 -1
  2. warp/__init__.pyi +3489 -1
  3. warp/autograd.py +45 -122
  4. warp/bin/libwarp-clang.dylib +0 -0
  5. warp/bin/libwarp.dylib +0 -0
  6. warp/build.py +241 -252
  7. warp/build_dll.py +125 -26
  8. warp/builtins.py +1907 -384
  9. warp/codegen.py +257 -101
  10. warp/config.py +12 -1
  11. warp/constants.py +1 -1
  12. warp/context.py +657 -223
  13. warp/dlpack.py +1 -1
  14. warp/examples/benchmarks/benchmark_cloth.py +2 -2
  15. warp/examples/benchmarks/benchmark_tile_sort.py +155 -0
  16. warp/examples/core/example_sample_mesh.py +1 -1
  17. warp/examples/core/example_spin_lock.py +93 -0
  18. warp/examples/core/example_work_queue.py +118 -0
  19. warp/examples/fem/example_adaptive_grid.py +5 -5
  20. warp/examples/fem/example_apic_fluid.py +1 -1
  21. warp/examples/fem/example_burgers.py +1 -1
  22. warp/examples/fem/example_convection_diffusion.py +9 -6
  23. warp/examples/fem/example_darcy_ls_optimization.py +489 -0
  24. warp/examples/fem/example_deformed_geometry.py +1 -1
  25. warp/examples/fem/example_diffusion.py +2 -2
  26. warp/examples/fem/example_diffusion_3d.py +1 -1
  27. warp/examples/fem/example_distortion_energy.py +1 -1
  28. warp/examples/fem/example_elastic_shape_optimization.py +387 -0
  29. warp/examples/fem/example_magnetostatics.py +5 -3
  30. warp/examples/fem/example_mixed_elasticity.py +5 -3
  31. warp/examples/fem/example_navier_stokes.py +11 -9
  32. warp/examples/fem/example_nonconforming_contact.py +5 -3
  33. warp/examples/fem/example_streamlines.py +8 -3
  34. warp/examples/fem/utils.py +9 -8
  35. warp/examples/interop/example_jax_ffi_callback.py +2 -2
  36. warp/examples/optim/example_drone.py +1 -1
  37. warp/examples/sim/example_cloth.py +1 -1
  38. warp/examples/sim/example_cloth_self_contact.py +48 -54
  39. warp/examples/tile/example_tile_block_cholesky.py +502 -0
  40. warp/examples/tile/example_tile_cholesky.py +2 -1
  41. warp/examples/tile/example_tile_convolution.py +1 -1
  42. warp/examples/tile/example_tile_filtering.py +1 -1
  43. warp/examples/tile/example_tile_matmul.py +1 -1
  44. warp/examples/tile/example_tile_mlp.py +2 -0
  45. warp/fabric.py +7 -7
  46. warp/fem/__init__.py +5 -0
  47. warp/fem/adaptivity.py +1 -1
  48. warp/fem/cache.py +152 -63
  49. warp/fem/dirichlet.py +2 -2
  50. warp/fem/domain.py +136 -6
  51. warp/fem/field/field.py +141 -99
  52. warp/fem/field/nodal_field.py +85 -39
  53. warp/fem/field/virtual.py +97 -52
  54. warp/fem/geometry/adaptive_nanogrid.py +91 -86
  55. warp/fem/geometry/closest_point.py +13 -0
  56. warp/fem/geometry/deformed_geometry.py +102 -40
  57. warp/fem/geometry/element.py +56 -2
  58. warp/fem/geometry/geometry.py +323 -22
  59. warp/fem/geometry/grid_2d.py +157 -62
  60. warp/fem/geometry/grid_3d.py +116 -20
  61. warp/fem/geometry/hexmesh.py +86 -20
  62. warp/fem/geometry/nanogrid.py +166 -86
  63. warp/fem/geometry/partition.py +59 -25
  64. warp/fem/geometry/quadmesh.py +86 -135
  65. warp/fem/geometry/tetmesh.py +47 -119
  66. warp/fem/geometry/trimesh.py +77 -270
  67. warp/fem/integrate.py +107 -52
  68. warp/fem/linalg.py +25 -58
  69. warp/fem/operator.py +124 -27
  70. warp/fem/quadrature/pic_quadrature.py +36 -14
  71. warp/fem/quadrature/quadrature.py +40 -16
  72. warp/fem/space/__init__.py +1 -1
  73. warp/fem/space/basis_function_space.py +66 -46
  74. warp/fem/space/basis_space.py +17 -4
  75. warp/fem/space/dof_mapper.py +1 -1
  76. warp/fem/space/function_space.py +2 -2
  77. warp/fem/space/grid_2d_function_space.py +4 -1
  78. warp/fem/space/hexmesh_function_space.py +4 -2
  79. warp/fem/space/nanogrid_function_space.py +3 -1
  80. warp/fem/space/partition.py +11 -2
  81. warp/fem/space/quadmesh_function_space.py +4 -1
  82. warp/fem/space/restriction.py +5 -2
  83. warp/fem/space/shape/__init__.py +10 -8
  84. warp/fem/space/tetmesh_function_space.py +4 -1
  85. warp/fem/space/topology.py +52 -21
  86. warp/fem/space/trimesh_function_space.py +4 -1
  87. warp/fem/utils.py +53 -8
  88. warp/jax.py +1 -2
  89. warp/jax_experimental/ffi.py +12 -17
  90. warp/jax_experimental/xla_ffi.py +37 -24
  91. warp/math.py +171 -1
  92. warp/native/array.h +99 -0
  93. warp/native/builtin.h +174 -31
  94. warp/native/coloring.cpp +1 -1
  95. warp/native/exports.h +118 -63
  96. warp/native/intersect.h +3 -3
  97. warp/native/mat.h +5 -10
  98. warp/native/mathdx.cpp +11 -5
  99. warp/native/matnn.h +1 -123
  100. warp/native/quat.h +28 -4
  101. warp/native/sparse.cpp +121 -258
  102. warp/native/sparse.cu +181 -274
  103. warp/native/spatial.h +305 -17
  104. warp/native/tile.h +583 -72
  105. warp/native/tile_radix_sort.h +1108 -0
  106. warp/native/tile_reduce.h +237 -2
  107. warp/native/tile_scan.h +240 -0
  108. warp/native/tuple.h +189 -0
  109. warp/native/vec.h +6 -16
  110. warp/native/warp.cpp +36 -4
  111. warp/native/warp.cu +574 -51
  112. warp/native/warp.h +47 -74
  113. warp/optim/linear.py +5 -1
  114. warp/paddle.py +7 -8
  115. warp/py.typed +0 -0
  116. warp/render/render_opengl.py +58 -29
  117. warp/render/render_usd.py +124 -61
  118. warp/sim/__init__.py +9 -0
  119. warp/sim/collide.py +252 -78
  120. warp/sim/graph_coloring.py +8 -1
  121. warp/sim/import_mjcf.py +4 -3
  122. warp/sim/import_usd.py +11 -7
  123. warp/sim/integrator.py +5 -2
  124. warp/sim/integrator_euler.py +1 -1
  125. warp/sim/integrator_featherstone.py +1 -1
  126. warp/sim/integrator_vbd.py +751 -320
  127. warp/sim/integrator_xpbd.py +1 -1
  128. warp/sim/model.py +265 -260
  129. warp/sim/utils.py +10 -7
  130. warp/sparse.py +303 -166
  131. warp/tape.py +52 -51
  132. warp/tests/cuda/test_conditional_captures.py +1046 -0
  133. warp/tests/cuda/test_streams.py +1 -1
  134. warp/tests/geometry/test_volume.py +2 -2
  135. warp/tests/interop/test_dlpack.py +9 -9
  136. warp/tests/interop/test_jax.py +0 -1
  137. warp/tests/run_coverage_serial.py +1 -1
  138. warp/tests/sim/disabled_kinematics.py +2 -2
  139. warp/tests/sim/{test_vbd.py → test_cloth.py} +296 -113
  140. warp/tests/sim/test_collision.py +159 -51
  141. warp/tests/sim/test_coloring.py +15 -1
  142. warp/tests/test_array.py +254 -2
  143. warp/tests/test_array_reduce.py +2 -2
  144. warp/tests/test_atomic_cas.py +299 -0
  145. warp/tests/test_codegen.py +142 -19
  146. warp/tests/test_conditional.py +47 -1
  147. warp/tests/test_ctypes.py +0 -20
  148. warp/tests/test_devices.py +8 -0
  149. warp/tests/test_fabricarray.py +4 -2
  150. warp/tests/test_fem.py +58 -25
  151. warp/tests/test_func.py +42 -1
  152. warp/tests/test_grad.py +1 -1
  153. warp/tests/test_lerp.py +1 -3
  154. warp/tests/test_map.py +481 -0
  155. warp/tests/test_mat.py +1 -24
  156. warp/tests/test_quat.py +6 -15
  157. warp/tests/test_rounding.py +10 -38
  158. warp/tests/test_runlength_encode.py +7 -7
  159. warp/tests/test_smoothstep.py +1 -1
  160. warp/tests/test_sparse.py +51 -2
  161. warp/tests/test_spatial.py +507 -1
  162. warp/tests/test_struct.py +2 -2
  163. warp/tests/test_tuple.py +265 -0
  164. warp/tests/test_types.py +2 -2
  165. warp/tests/test_utils.py +24 -18
  166. warp/tests/tile/test_tile.py +420 -1
  167. warp/tests/tile/test_tile_mathdx.py +518 -14
  168. warp/tests/tile/test_tile_reduce.py +213 -0
  169. warp/tests/tile/test_tile_shared_memory.py +130 -1
  170. warp/tests/tile/test_tile_sort.py +117 -0
  171. warp/tests/unittest_suites.py +4 -6
  172. warp/types.py +462 -308
  173. warp/utils.py +647 -86
  174. {warp_lang-1.7.2.dist-info → warp_lang-1.8.0.dist-info}/METADATA +20 -6
  175. {warp_lang-1.7.2.dist-info → warp_lang-1.8.0.dist-info}/RECORD +178 -166
  176. warp/stubs.py +0 -3381
  177. warp/tests/sim/test_xpbd.py +0 -399
  178. warp/tests/test_mlp.py +0 -282
  179. {warp_lang-1.7.2.dist-info → warp_lang-1.8.0.dist-info}/WHEEL +0 -0
  180. {warp_lang-1.7.2.dist-info → warp_lang-1.8.0.dist-info}/licenses/LICENSE.md +0 -0
  181. {warp_lang-1.7.2.dist-info → warp_lang-1.8.0.dist-info}/top_level.txt +0 -0
@@ -78,7 +78,7 @@ def vertex_triangle_collision_detection_brute_force(
78
78
 
79
79
 
80
80
  @wp.kernel
81
- def validate_vertex_collisions(
81
+ def vertex_triangle_collision_detection_brute_force_no_triangle_buffers(
82
82
  query_radius: float,
83
83
  bvh_id: wp.uint64,
84
84
  pos: wp.array(dtype=wp.vec3),
@@ -88,33 +88,82 @@ def validate_vertex_collisions(
88
88
  vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32),
89
89
  vertex_colliding_triangles_buffer_size: wp.array(dtype=wp.int32),
90
90
  vertex_colliding_triangles_min_dist: wp.array(dtype=float),
91
+ triangle_colliding_vertices_min_dist: wp.array(dtype=float),
91
92
  resize_flags: wp.array(dtype=wp.int32),
92
93
  ):
93
94
  v_index = wp.tid()
94
95
  v = pos[v_index]
95
96
 
96
- num_cols = vertex_colliding_triangles_count[v_index]
97
- offset = vertex_colliding_triangles_offsets[v_index]
98
- min_dis = vertex_colliding_triangles_min_dist[v_index]
99
- for col in range(wp.min(num_cols, vertex_colliding_triangles_buffer_size[v_index])):
100
- tri_index = vertex_colliding_triangles[offset + col]
101
-
97
+ vertex_num_collisions = wp.int32(0)
98
+ min_dis_to_tris = query_radius
99
+ for tri_index in range(tri_indices.shape[0]):
102
100
  t1 = tri_indices[tri_index, 0]
103
101
  t2 = tri_indices[tri_index, 1]
104
102
  t3 = tri_indices[tri_index, 2]
105
- # wp.expect_eq(vertex_on_triangle(v_index, t1, t2, t3), False)
103
+ if vertex_adjacent_to_triangle(v_index, t1, t2, t3):
104
+ continue
106
105
 
107
106
  u1 = pos[t1]
108
107
  u2 = pos[t2]
109
108
  u3 = pos[t3]
110
109
 
111
110
  closest_p, bary, feature_type = triangle_closest_point(u1, u2, u3, v)
111
+
112
112
  dis = wp.length(closest_p - v)
113
- wp.expect_eq(dis < query_radius, True)
114
- wp.expect_eq(dis >= min_dis, True)
115
113
 
116
- # wp.printf("vertex %d, offset %d, num cols %d, colliding with triangle: %d, dis: %f\n",
117
- # v_index, offset, num_cols, tri_index, dis)
114
+ if dis < query_radius:
115
+ vertex_num_collisions = vertex_num_collisions + 1
116
+ min_dis_to_tris = wp.min(dis, min_dis_to_tris)
117
+
118
+ wp.atomic_min(triangle_colliding_vertices_min_dist, tri_index, dis)
119
+
120
+ vertex_colliding_triangles_count[v_index] = vertex_num_collisions
121
+ vertex_colliding_triangles_min_dist[v_index] = min_dis_to_tris
122
+
123
+
124
+ @wp.kernel
125
+ def validate_vertex_collisions(
126
+ query_radius: float,
127
+ bvh_id: wp.uint64,
128
+ pos: wp.array(dtype=wp.vec3),
129
+ tri_indices: wp.array(dtype=wp.int32, ndim=2),
130
+ vertex_colliding_triangles: wp.array(dtype=wp.int32),
131
+ vertex_colliding_triangles_count: wp.array(dtype=wp.int32),
132
+ vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32),
133
+ vertex_colliding_triangles_buffer_size: wp.array(dtype=wp.int32),
134
+ vertex_colliding_triangles_min_dist: wp.array(dtype=float),
135
+ resize_flags: wp.array(dtype=wp.int32),
136
+ ):
137
+ v_index = wp.tid()
138
+ v = pos[v_index]
139
+
140
+ num_cols = vertex_colliding_triangles_count[v_index]
141
+ offset = vertex_colliding_triangles_offsets[v_index]
142
+ min_dis = vertex_colliding_triangles_min_dist[v_index]
143
+ for col in range(vertex_colliding_triangles_buffer_size[v_index]):
144
+ vertex_index = vertex_colliding_triangles[2 * (offset + col)]
145
+ tri_index = vertex_colliding_triangles[2 * (offset + col) + 1]
146
+ if col < num_cols:
147
+ t1 = tri_indices[tri_index, 0]
148
+ t2 = tri_indices[tri_index, 1]
149
+ t3 = tri_indices[tri_index, 2]
150
+ # wp.expect_eq(vertex_on_triangle(v_index, t1, t2, t3), False)
151
+
152
+ u1 = pos[t1]
153
+ u2 = pos[t2]
154
+ u3 = pos[t3]
155
+
156
+ closest_p, bary, feature_type = triangle_closest_point(u1, u2, u3, v)
157
+ dis = wp.length(closest_p - v)
158
+ wp.expect_eq(dis < query_radius, True)
159
+ wp.expect_eq(dis >= min_dis, True)
160
+ wp.expect_eq(v_index == vertex_colliding_triangles[2 * (offset + col)], True)
161
+
162
+ # wp.printf("vertex %d, offset %d, num cols %d, colliding with triangle: %d, dis: %f\n",
163
+ # v_index, offset, num_cols, tri_index, dis)
164
+ else:
165
+ wp.expect_eq(vertex_index == -1, True)
166
+ wp.expect_eq(tri_index == -1, True)
118
167
 
119
168
 
120
169
  @wp.kernel
@@ -242,30 +291,35 @@ def validate_edge_collisions(
242
291
  num_cols = edge_colliding_edges_count[e0_index]
243
292
  offset = edge_colliding_edges_offsets[e0_index]
244
293
  min_dist = edge_colliding_edges_min_dist[e0_index]
245
- for col in range(wp.min(num_cols, edge_colliding_edges_buffer_sizes[e0_index])):
246
- e1_index = edge_colliding_edges[offset + col]
294
+ for col in range(edge_colliding_edges_buffer_sizes[e0_index]):
295
+ e1_index = edge_colliding_edges[2 * (offset + col) + 1]
247
296
 
248
- e1_v0 = edge_indices[e1_index, 2]
249
- e1_v1 = edge_indices[e1_index, 3]
297
+ if col < num_cols:
298
+ e1_v0 = edge_indices[e1_index, 2]
299
+ e1_v1 = edge_indices[e1_index, 3]
250
300
 
251
- if e0_v0 == e1_v0 or e0_v0 == e1_v1 or e0_v1 == e1_v0 or e0_v1 == e1_v1:
252
- wp.expect_eq(False, True)
301
+ if e0_v0 == e1_v0 or e0_v0 == e1_v1 or e0_v1 == e1_v0 or e0_v1 == e1_v1:
302
+ wp.expect_eq(False, True)
253
303
 
254
- e1_v0_pos = pos[e1_v0]
255
- e1_v1_pos = pos[e1_v1]
304
+ e1_v0_pos = pos[e1_v0]
305
+ e1_v1_pos = pos[e1_v1]
256
306
 
257
- st = wp.closest_point_edge_edge(e0_v0_pos, e0_v1_pos, e1_v0_pos, e1_v1_pos, edge_edge_parallel_epsilon)
258
- s = st[0]
259
- t = st[1]
260
- c1 = e0_v0_pos + (e0_v1_pos - e0_v0_pos) * s
261
- c2 = e1_v0_pos + (e1_v1_pos - e1_v0_pos) * t
307
+ st = wp.closest_point_edge_edge(e0_v0_pos, e0_v1_pos, e1_v0_pos, e1_v1_pos, edge_edge_parallel_epsilon)
308
+ s = st[0]
309
+ t = st[1]
310
+ c1 = e0_v0_pos + (e0_v1_pos - e0_v0_pos) * s
311
+ c2 = e1_v0_pos + (e1_v1_pos - e1_v0_pos) * t
262
312
 
263
- dist = wp.length(c2 - c1)
313
+ dist = wp.length(c2 - c1)
264
314
 
265
- wp.expect_eq(dist >= min_dist, True)
315
+ wp.expect_eq(dist >= min_dist, True)
316
+ wp.expect_eq(e0_index == edge_colliding_edges[2 * (offset + col)], True)
317
+ else:
318
+ wp.expect_eq(e1_index == -1, True)
319
+ wp.expect_eq(edge_colliding_edges[2 * (offset + col)] == -1, True)
266
320
 
267
321
 
268
- def init_model(vs, fs, device):
322
+ def init_model(vs, fs, device, record_triangle_contacting_vertices=True):
269
323
  vertices = [wp.vec3(v) for v in vs]
270
324
 
271
325
  builder = wp.sim.ModelBuilder()
@@ -283,7 +337,7 @@ def init_model(vs, fs, device):
283
337
  )
284
338
  model = builder.finalize(device=device)
285
339
 
286
- collision_detector = TriMeshCollisionDetector(model=model)
340
+ collision_detector = TriMeshCollisionDetector(model=model, record_triangle_contacting_vertices=True)
287
341
 
288
342
  return model, collision_detector
289
343
 
@@ -304,7 +358,8 @@ def get_data():
304
358
  def test_vertex_triangle_collision(test, device):
305
359
  vertices, faces = get_data()
306
360
 
307
- model, collision_detector = init_model(vertices, faces, device)
361
+ # record triangle contacting vertices
362
+ model, collision_detector = init_model(vertices, faces, device, True)
308
363
 
309
364
  rs = [1e-2, 2e-2, 5e-2, 1e-1]
310
365
 
@@ -316,6 +371,42 @@ def test_vertex_triangle_collision(test, device):
316
371
  triangle_colliding_vertices_count_1 = collision_detector.triangle_colliding_vertices_count.numpy()
317
372
  triangle_min_dis_1 = collision_detector.triangle_colliding_vertices_min_dist.numpy()
318
373
 
374
+ wp.launch(
375
+ kernel=validate_vertex_collisions,
376
+ inputs=[
377
+ query_radius,
378
+ collision_detector.bvh_tris.id,
379
+ collision_detector.model.particle_q,
380
+ collision_detector.model.tri_indices,
381
+ collision_detector.vertex_colliding_triangles,
382
+ collision_detector.vertex_colliding_triangles_count,
383
+ collision_detector.vertex_colliding_triangles_offsets,
384
+ collision_detector.vertex_colliding_triangles_buffer_sizes,
385
+ collision_detector.vertex_colliding_triangles_min_dist,
386
+ collision_detector.resize_flags,
387
+ ],
388
+ dim=model.particle_count,
389
+ device=device,
390
+ )
391
+
392
+ wp.launch(
393
+ kernel=validate_triangle_collisions,
394
+ inputs=[
395
+ query_radius,
396
+ collision_detector.bvh_tris.id,
397
+ collision_detector.model.particle_q,
398
+ collision_detector.model.tri_indices,
399
+ collision_detector.triangle_colliding_vertices,
400
+ collision_detector.triangle_colliding_vertices_count,
401
+ collision_detector.triangle_colliding_vertices_offsets,
402
+ collision_detector.triangle_colliding_vertices_buffer_sizes,
403
+ collision_detector.triangle_colliding_vertices_min_dist,
404
+ collision_detector.resize_flags,
405
+ ],
406
+ dim=model.tri_count,
407
+ device=model.device,
408
+ )
409
+
319
410
  wp.launch(
320
411
  kernel=init_triangle_collision_data_kernel,
321
412
  inputs=[
@@ -325,7 +416,7 @@ def test_vertex_triangle_collision(test, device):
325
416
  collision_detector.resize_flags,
326
417
  ],
327
418
  dim=model.tri_count,
328
- device=device,
419
+ device=model.device,
329
420
  )
330
421
 
331
422
  wp.launch(
@@ -348,7 +439,7 @@ def test_vertex_triangle_collision(test, device):
348
439
  collision_detector.resize_flags,
349
440
  ],
350
441
  dim=model.particle_count,
351
- device=device,
442
+ device=model.device,
352
443
  )
353
444
 
354
445
  vertex_colliding_triangles_count_2 = collision_detector.vertex_colliding_triangles_count.numpy()
@@ -362,6 +453,18 @@ def test_vertex_triangle_collision(test, device):
362
453
  assert_np_equal(triangle_colliding_vertices_count_2, triangle_colliding_vertices_count_1)
363
454
  assert_np_equal(vertex_min_dis_2, vertex_min_dis_1)
364
455
 
456
+ # do not record triangle contacting vertices
457
+ model, collision_detector = init_model(vertices, faces, device, False)
458
+
459
+ rs = [1e-2, 2e-2, 5e-2, 1e-1]
460
+
461
+ for query_radius in rs:
462
+ collision_detector.vertex_triangle_collision_detection(query_radius)
463
+ vertex_colliding_triangles_count_1 = collision_detector.vertex_colliding_triangles_count.numpy()
464
+ vertex_min_dis_1 = collision_detector.vertex_colliding_triangles_min_dist.numpy()
465
+
466
+ triangle_min_dis_1 = collision_detector.triangle_colliding_vertices_min_dist.numpy()
467
+
365
468
  wp.launch(
366
469
  kernel=validate_vertex_collisions,
367
470
  inputs=[
@@ -381,24 +484,31 @@ def test_vertex_triangle_collision(test, device):
381
484
  )
382
485
 
383
486
  wp.launch(
384
- kernel=validate_triangle_collisions,
487
+ kernel=vertex_triangle_collision_detection_brute_force_no_triangle_buffers,
385
488
  inputs=[
386
489
  query_radius,
387
490
  collision_detector.bvh_tris.id,
388
491
  collision_detector.model.particle_q,
389
492
  collision_detector.model.tri_indices,
390
- collision_detector.triangle_colliding_vertices,
391
- collision_detector.triangle_colliding_vertices_count,
392
- collision_detector.triangle_colliding_vertices_offsets,
393
- collision_detector.triangle_colliding_vertices_buffer_sizes,
493
+ collision_detector.vertex_colliding_triangles,
494
+ collision_detector.vertex_colliding_triangles_count,
495
+ collision_detector.vertex_colliding_triangles_offsets,
496
+ collision_detector.vertex_colliding_triangles_buffer_sizes,
497
+ collision_detector.vertex_colliding_triangles_min_dist,
394
498
  collision_detector.triangle_colliding_vertices_min_dist,
395
499
  collision_detector.resize_flags,
396
500
  ],
397
- dim=model.tri_count,
398
- device=device,
501
+ dim=model.particle_count,
502
+ device=model.device,
399
503
  )
400
504
 
401
- wp.synchronize_device(device)
505
+ vertex_colliding_triangles_count_2 = collision_detector.vertex_colliding_triangles_count.numpy()
506
+ vertex_min_dis_2 = collision_detector.vertex_colliding_triangles_min_dist.numpy()
507
+ triangle_min_dis_2 = collision_detector.triangle_colliding_vertices_min_dist.numpy()
508
+
509
+ assert_np_equal(vertex_colliding_triangles_count_2, vertex_colliding_triangles_count_1)
510
+ assert_np_equal(triangle_min_dis_2, triangle_min_dis_1)
511
+ assert_np_equal(vertex_min_dis_2, vertex_min_dis_1)
402
512
 
403
513
 
404
514
  @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
@@ -416,7 +526,7 @@ def test_edge_edge_collision(test, device):
416
526
  edge_min_dist_1 = collision_detector.edge_colliding_edges_min_dist.numpy()
417
527
 
418
528
  wp.launch(
419
- kernel=edge_edge_collision_detection_brute_force,
529
+ kernel=validate_edge_collisions,
420
530
  inputs=[
421
531
  query_radius,
422
532
  collision_detector.bvh_edges.id,
@@ -432,18 +542,12 @@ def test_edge_edge_collision(test, device):
432
542
  collision_detector.edge_colliding_edges_min_dist,
433
543
  collision_detector.resize_flags,
434
544
  ],
435
- dim=model.edge_count,
545
+ dim=model.particle_count,
436
546
  device=device,
437
547
  )
438
548
 
439
- edge_colliding_edges_count_2 = collision_detector.edge_colliding_edges_count.numpy()
440
- edge_min_dist_2 = collision_detector.edge_colliding_edges_min_dist.numpy()
441
-
442
- assert_np_equal(edge_colliding_edges_count_2, edge_colliding_edges_count_1)
443
- assert_np_equal(edge_min_dist_2, edge_min_dist_1)
444
-
445
549
  wp.launch(
446
- kernel=validate_edge_collisions,
550
+ kernel=edge_edge_collision_detection_brute_force,
447
551
  inputs=[
448
552
  query_radius,
449
553
  collision_detector.bvh_edges.id,
@@ -459,11 +563,15 @@ def test_edge_edge_collision(test, device):
459
563
  collision_detector.edge_colliding_edges_min_dist,
460
564
  collision_detector.resize_flags,
461
565
  ],
462
- dim=model.particle_count,
566
+ dim=model.edge_count,
463
567
  device=device,
464
568
  )
465
569
 
466
- wp.synchronize_device(device)
570
+ edge_colliding_edges_count_2 = collision_detector.edge_colliding_edges_count.numpy()
571
+ edge_min_dist_2 = collision_detector.edge_colliding_edges_min_dist.numpy()
572
+
573
+ assert_np_equal(edge_colliding_edges_count_2, edge_colliding_edges_count_1)
574
+ assert_np_equal(edge_min_dist_1, edge_min_dist_2)
467
575
 
468
576
 
469
577
  def test_particle_collision(test, device):
@@ -48,6 +48,19 @@ def color_lattice_grid(num_x, num_y):
48
48
  return color_groups
49
49
 
50
50
 
51
+ def test_coloring_corner_case(test, device):
52
+ builder_1 = wp.sim.ModelBuilder()
53
+ builder_1.color()
54
+ test.assertTrue(len(builder_1.particle_color_groups) == 0)
55
+
56
+ builder_2 = wp.sim.ModelBuilder()
57
+ builder_2.add_particle(pos=wp.vec3(0, 0, 0), vel=wp.vec3(0, 0, 0), mass=1.0)
58
+ builder_2.add_particle(pos=wp.vec3(1, 0, 0), vel=wp.vec3(0, 0, 0), mass=1.0)
59
+ builder_2.color()
60
+ test.assertTrue(len(builder_2.particle_color_groups) == 1)
61
+ test.assertTrue(builder_2.particle_color_groups[0].shape[0] == 2)
62
+
63
+
51
64
  @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
52
65
  def test_coloring_trimesh(test, device):
53
66
  from pxr import Usd, UsdGeom
@@ -227,7 +240,7 @@ def test_combine_coloring(test, device):
227
240
 
228
241
  particle_number_colored = np.full((model.particle_count), -1, dtype=int)
229
242
  particle_colors = np.full((model.particle_count), -1, dtype=int)
230
- for color, color_group in enumerate(model.particle_coloring):
243
+ for color, color_group in enumerate(model.particle_color_groups):
231
244
  particle_number_colored[color_group.numpy()] += 1
232
245
  particle_colors[color_group.numpy()] = color
233
246
 
@@ -252,6 +265,7 @@ class TestColoring(unittest.TestCase):
252
265
 
253
266
  add_function_test(TestColoring, "test_coloring_trimesh", test_coloring_trimesh, devices=devices)
254
267
  add_function_test(TestColoring, "test_combine_coloring", test_combine_coloring, devices=devices)
268
+ add_function_test(TestColoring, "test_coloring_corner_case", test_coloring_corner_case, devices=devices)
255
269
 
256
270
  if __name__ == "__main__":
257
271
  wp.clear_kernel_cache()
warp/tests/test_array.py CHANGED
@@ -412,7 +412,7 @@ def test_slicing(test, device):
412
412
  assert_array_equal(wp_arr[:5], wp.array(np_arr[:5], dtype=int, device=device))
413
413
  assert_array_equal(wp_arr[1:5], wp.array(np_arr[1:5], dtype=int, device=device))
414
414
  assert_array_equal(wp_arr[-9:-5:1], wp.array(np_arr[-9:-5:1], dtype=int, device=device))
415
- assert_array_equal(wp_arr[:5,], wp.array(np_arr[:5], dtype=int, device=device)) # noqa: E231
415
+ assert_array_equal(wp_arr[:5,], wp.array(np_arr[:5], dtype=int, device=device))
416
416
 
417
417
 
418
418
  def test_view(test, device):
@@ -2370,6 +2370,257 @@ def test_array_from_cai(test, device):
2370
2370
  assert_np_equal(arr_warp.numpy(), np.array([[2, 1, 1], [1, 0, 0], [1, 0, 0]]))
2371
2371
 
2372
2372
 
2373
+ def test_array_from_data(test, device):
2374
+ with wp.ScopedDevice(device):
2375
+ # =========================================
2376
+ # scalars, reshaping
2377
+
2378
+ data = np.arange(12, dtype=np.float32).reshape((3, 4))
2379
+ src = wp.array(data)
2380
+
2381
+ assert src.device == device
2382
+
2383
+ dtypes = [Any, wp.float32]
2384
+ shapes = [None, (3, 4), (12,), (3, 2, 2)]
2385
+
2386
+ for dtype in dtypes:
2387
+ for shape in shapes:
2388
+ with test.subTest(msg=f"scalar, dtype={dtype}, shape={shape}"):
2389
+ dst = wp.array(src, dtype=dtype, shape=shape)
2390
+ assert dst.device == src.device
2391
+ if dtype is Any:
2392
+ assert dst.dtype == src.dtype
2393
+ else:
2394
+ assert dst.dtype == dtype
2395
+ if shape is None:
2396
+ assert dst.shape == src.shape
2397
+ assert_np_equal(dst.numpy(), data)
2398
+ else:
2399
+ assert dst.shape == shape
2400
+ assert_np_equal(dst.numpy(), data.reshape(shape))
2401
+
2402
+ # =========================================
2403
+ # vectors, reshaping
2404
+
2405
+ with test.subTest(msg="vector, single"):
2406
+ data = np.arange(3, dtype=np.float32)
2407
+ src = wp.array(data)
2408
+ dst = wp.array(src, dtype=wp.vec3)
2409
+ assert dst.dtype == wp.vec3
2410
+ assert dst.shape == (1,)
2411
+ assert_np_equal(dst.numpy(), data.reshape((1, 3)))
2412
+
2413
+ with test.subTest(msg="vector, multiple in 1d"):
2414
+ data = np.arange(12, dtype=np.float32)
2415
+ src = wp.array(data)
2416
+ dst = wp.array(src, dtype=wp.vec3)
2417
+ assert dst.dtype == wp.vec3
2418
+ assert dst.shape == (4,)
2419
+ assert_np_equal(dst.numpy(), data.reshape((4, 3)))
2420
+
2421
+ with test.subTest(msg="vector, singles in 2d"):
2422
+ data = np.arange(12, dtype=np.float32).reshape((4, 3))
2423
+ src = wp.array(data)
2424
+ dst = wp.array(src, dtype=wp.vec3)
2425
+ assert dst.dtype == wp.vec3
2426
+ assert dst.shape == (4,)
2427
+ assert_np_equal(dst.numpy(), data.reshape((4, 3)))
2428
+
2429
+ with test.subTest(msg="vector, multiples in 2d"):
2430
+ data = np.arange(24, dtype=np.float32).reshape((4, 6))
2431
+ src = wp.array(data)
2432
+ dst = wp.array(src, dtype=wp.vec3)
2433
+ assert dst.dtype == wp.vec3
2434
+ assert dst.shape == (4, 2)
2435
+ assert_np_equal(dst.numpy(), data.reshape((4, 2, 3)))
2436
+
2437
+ with test.subTest(msg="vector, singles in 2d, reshape"):
2438
+ data = np.arange(12, dtype=np.float32).reshape((4, 3))
2439
+ src = wp.array(data)
2440
+ dst = wp.array(src, dtype=wp.vec3, shape=(2, 2))
2441
+ assert dst.dtype == wp.vec3
2442
+ assert dst.shape == (2, 2)
2443
+ assert_np_equal(dst.numpy(), data.reshape((2, 2, 3)))
2444
+
2445
+ with test.subTest(msg="vector, multiples in 2d, reshape"):
2446
+ data = np.arange(24, dtype=np.float32).reshape((4, 6))
2447
+ src = wp.array(data)
2448
+ dst = wp.array(src, dtype=wp.vec3, shape=(2, 2, 2))
2449
+ assert dst.dtype == wp.vec3
2450
+ assert dst.shape == (2, 2, 2)
2451
+ assert_np_equal(dst.numpy(), data.reshape((2, 2, 2, 3)))
2452
+
2453
+ # =========================================
2454
+ # matrices, reshaping
2455
+
2456
+ with test.subTest(msg="matrix, single in 2d"):
2457
+ # one 2x2 matrix in a 2d array
2458
+ data = np.arange(4, dtype=np.float32).reshape((2, 2))
2459
+ src = wp.array(data)
2460
+ dst = wp.array(src, dtype=wp.mat22)
2461
+ assert dst.dtype == wp.mat22
2462
+ assert dst.shape == (1,)
2463
+ assert_np_equal(dst.numpy(), data.reshape((1, 2, 2)))
2464
+
2465
+ with test.subTest(msg="matrix, single in 1d"):
2466
+ # 2x2 matrix in a 1d array
2467
+ data = np.arange(4, dtype=np.float32)
2468
+ src = wp.array(data)
2469
+ dst = wp.array(src, dtype=wp.mat22)
2470
+ assert dst.dtype == wp.mat22
2471
+ assert dst.shape == (1,)
2472
+ assert_np_equal(dst.numpy(), data.reshape((1, 2, 2)))
2473
+
2474
+ with test.subTest(msg="matrix, multiples in 1d"):
2475
+ # 3 2x2 matrices in a 1d array
2476
+ data = np.arange(12, dtype=np.float32)
2477
+ src = wp.array(data)
2478
+ dst = wp.array(src, dtype=wp.mat22)
2479
+ assert dst.dtype == wp.mat22
2480
+ assert dst.shape == (3,)
2481
+ assert_np_equal(dst.numpy(), data.reshape((3, 2, 2)))
2482
+
2483
+ with test.subTest(msg="matrix, multiples in 1d, reshape"):
2484
+ # 4 2x2 matrices in a 1d array
2485
+ data = np.arange(16, dtype=np.float32)
2486
+ src = wp.array(data)
2487
+ dst = wp.array(src, dtype=wp.mat22, shape=(4,))
2488
+ assert dst.dtype == wp.mat22
2489
+ assert dst.shape == (4,)
2490
+ assert_np_equal(dst.numpy(), data.reshape((4, 2, 2)))
2491
+
2492
+ with test.subTest(msg="matrix, multiples in 2d"):
2493
+ # 3 2x2 matrices in a 2d array
2494
+ data = np.arange(12, dtype=np.float32).reshape((3, 4))
2495
+ src = wp.array(data)
2496
+ dst = wp.array(src, dtype=wp.mat22)
2497
+ assert dst.dtype == wp.mat22
2498
+ assert dst.shape == (3,)
2499
+ assert_np_equal(dst.numpy(), data.reshape((3, 2, 2)))
2500
+
2501
+ with test.subTest(msg="matrix, multiples in 2d, reshape"):
2502
+ # 4 2x2 matrices in a 2d array
2503
+ data = np.arange(16, dtype=np.float32).reshape((4, 4))
2504
+ src = wp.array(data)
2505
+ dst = wp.array(src, dtype=wp.mat22, shape=(2, 2))
2506
+ assert dst.dtype == wp.mat22
2507
+ assert dst.shape == (2, 2)
2508
+ assert_np_equal(dst.numpy(), data.reshape((2, 2, 2, 2)))
2509
+
2510
+ with test.subTest(msg="matrix, multiples in 3d"):
2511
+ # 3 2x2 matrices in a 3d array
2512
+ data = np.arange(12, dtype=np.float32).reshape((3, 2, 2))
2513
+ src = wp.array(data)
2514
+ dst = wp.array(src, dtype=wp.mat22)
2515
+ assert dst.dtype == wp.mat22
2516
+ assert dst.shape == (3,)
2517
+ assert_np_equal(dst.numpy(), data.reshape((3, 2, 2)))
2518
+
2519
+ with test.subTest(msg="matrix, multiples in 3d, reshape"):
2520
+ # 4 2x2 matrices in a 3d array
2521
+ data = np.arange(16, dtype=np.float32).reshape((4, 2, 2))
2522
+ src = wp.array(data)
2523
+ dst = wp.array(src, dtype=wp.mat22, shape=(2, 2))
2524
+ assert dst.dtype == wp.mat22
2525
+ assert dst.shape == (2, 2)
2526
+ assert_np_equal(dst.numpy(), data.reshape((2, 2, 2, 2)))
2527
+
2528
+ # =========================================
2529
+ # vectors and matrices in strided arrays
2530
+
2531
+ with test.subTest(msg="vector, singles in 2d, strided"):
2532
+ # 4 vec3 in strided 2d array
2533
+ data = np.arange(20, dtype=np.float32).reshape((4, 5))
2534
+ src = wp.array(data)[:, 2:] # source with strides
2535
+ dst = wp.array(src, dtype=wp.vec3)
2536
+ assert dst.dtype == wp.vec3
2537
+ assert dst.shape == (4,)
2538
+ expected = np.array(
2539
+ [
2540
+ [2, 3, 4],
2541
+ [7, 8, 9],
2542
+ [12, 13, 14],
2543
+ [17, 18, 19],
2544
+ ],
2545
+ dtype=np.float32,
2546
+ )
2547
+ assert_np_equal(dst.numpy(), expected)
2548
+
2549
+ with test.subTest(msg="vector, multiples in 2d, strided"):
2550
+ # 4 vec3 in strided 2d array
2551
+ data = np.arange(14, dtype=np.float32).reshape((2, 7))
2552
+ src = wp.array(data)[:, 1:] # source with strides
2553
+ dst = wp.array(src, dtype=wp.vec3)
2554
+ assert dst.dtype == wp.vec3
2555
+ assert dst.shape == (2, 2)
2556
+ expected = np.array(
2557
+ [
2558
+ [
2559
+ [1, 2, 3],
2560
+ [4, 5, 6],
2561
+ ],
2562
+ [
2563
+ [8, 9, 10],
2564
+ [11, 12, 13],
2565
+ ],
2566
+ ],
2567
+ dtype=np.float32,
2568
+ )
2569
+ assert_np_equal(dst.numpy(), expected)
2570
+
2571
+ with test.subTest(msg="matrix, multiples in 2d, strided"):
2572
+ # 3 2x2 matrices in a 2d array
2573
+ data = np.arange(15, dtype=np.float32).reshape((3, 5))
2574
+ src = wp.array(data)[:, 1:] # source with strides
2575
+ dst = wp.array(src, dtype=wp.mat22)
2576
+ assert dst.dtype == wp.mat22
2577
+ assert dst.shape == (3,)
2578
+ expected = np.array(
2579
+ [
2580
+ [
2581
+ [1, 2],
2582
+ [3, 4],
2583
+ ],
2584
+ [
2585
+ [6, 7],
2586
+ [8, 9],
2587
+ ],
2588
+ [
2589
+ [11, 12],
2590
+ [13, 14],
2591
+ ],
2592
+ ],
2593
+ dtype=np.float32,
2594
+ )
2595
+ assert_np_equal(dst.numpy(), expected)
2596
+
2597
+ with test.subTest(msg="matrix, multiples in 3d, strided"):
2598
+ # 3 2x2 matrices in a 3d array
2599
+ data = np.arange(18, dtype=np.float32).reshape((3, 3, 2))
2600
+ src = wp.array(data)[:, 1:] # source with strides
2601
+ dst = wp.array(src, dtype=wp.mat22)
2602
+ assert dst.dtype == wp.mat22
2603
+ assert dst.shape == (3,)
2604
+ expected = np.array(
2605
+ [
2606
+ [
2607
+ [2, 3],
2608
+ [4, 5],
2609
+ ],
2610
+ [
2611
+ [8, 9],
2612
+ [10, 11],
2613
+ ],
2614
+ [
2615
+ [14, 15],
2616
+ [16, 17],
2617
+ ],
2618
+ ],
2619
+ dtype=np.float32,
2620
+ )
2621
+ assert_np_equal(dst.numpy(), expected)
2622
+
2623
+
2373
2624
  @wp.kernel
2374
2625
  def inplace_add_1d(x: wp.array(dtype=float), y: wp.array(dtype=float)):
2375
2626
  i = wp.tid()
@@ -2604,7 +2855,7 @@ def test_array_inplace_non_diff_ops(test, device):
2604
2855
  wp.launch(inplace_div_1d, N, inputs=[x1, y1], device=device)
2605
2856
  assert_np_equal(x1.numpy(), np.full(N, fill_value=2.0, dtype=float))
2606
2857
 
2607
- for dtype in wp.types.non_atomic_types + (wp.vec2b, wp.vec2ub, wp.vec2s, wp.vec2us, uint16vec3):
2858
+ for dtype in (*wp.types.non_atomic_types, wp.vec2b, wp.vec2ub, wp.vec2s, wp.vec2us, uint16vec3):
2608
2859
  x = wp.full(N, value=0, dtype=dtype, device=device)
2609
2860
  y = wp.full(N, value=1, dtype=dtype, device=device)
2610
2861
 
@@ -2943,6 +3194,7 @@ add_function_test(TestArray, "test_alloc_strides", test_alloc_strides, devices=d
2943
3194
  add_function_test(TestArray, "test_casting", test_casting, devices=devices)
2944
3195
  add_function_test(TestArray, "test_array_len", test_array_len, devices=devices)
2945
3196
  add_function_test(TestArray, "test_cuda_interface_conversion", test_cuda_interface_conversion, devices=devices)
3197
+ add_function_test(TestArray, "test_array_from_data", test_array_from_data, devices=devices)
2946
3198
 
2947
3199
  try:
2948
3200
  import torch
@@ -28,7 +28,7 @@ def make_test_array_sum(dtype):
28
28
  def test_array_sum(test, device):
29
29
  rng = np.random.default_rng(123)
30
30
 
31
- cols = wp.types.type_length(dtype)
31
+ cols = wp.types.type_size(dtype)
32
32
 
33
33
  values_np = rng.random(size=(N, cols))
34
34
  values = wp.array(values_np, device=device, dtype=dtype)
@@ -77,7 +77,7 @@ def make_test_array_inner(dtype):
77
77
  def test_array_inner(test, device):
78
78
  rng = np.random.default_rng(123)
79
79
 
80
- cols = wp.types.type_length(dtype)
80
+ cols = wp.types.type_size(dtype)
81
81
 
82
82
  a_np = rng.random(size=(N, cols))
83
83
  b_np = rng.random(size=(N, cols))