warp-lang 1.7.0__py3-none-manylinux_2_34_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 (429) hide show
  1. warp/__init__.py +139 -0
  2. warp/__init__.pyi +1 -0
  3. warp/autograd.py +1142 -0
  4. warp/bin/warp-clang.so +0 -0
  5. warp/bin/warp.so +0 -0
  6. warp/build.py +557 -0
  7. warp/build_dll.py +405 -0
  8. warp/builtins.py +6855 -0
  9. warp/codegen.py +3969 -0
  10. warp/config.py +158 -0
  11. warp/constants.py +57 -0
  12. warp/context.py +6812 -0
  13. warp/dlpack.py +462 -0
  14. warp/examples/__init__.py +24 -0
  15. warp/examples/assets/bear.usd +0 -0
  16. warp/examples/assets/bunny.usd +0 -0
  17. warp/examples/assets/cartpole.urdf +110 -0
  18. warp/examples/assets/crazyflie.usd +0 -0
  19. warp/examples/assets/cube.usd +0 -0
  20. warp/examples/assets/nonuniform.usd +0 -0
  21. warp/examples/assets/nv_ant.xml +92 -0
  22. warp/examples/assets/nv_humanoid.xml +183 -0
  23. warp/examples/assets/nvidia_logo.png +0 -0
  24. warp/examples/assets/pixel.jpg +0 -0
  25. warp/examples/assets/quadruped.urdf +268 -0
  26. warp/examples/assets/rocks.nvdb +0 -0
  27. warp/examples/assets/rocks.usd +0 -0
  28. warp/examples/assets/sphere.usd +0 -0
  29. warp/examples/assets/square_cloth.usd +0 -0
  30. warp/examples/benchmarks/benchmark_api.py +389 -0
  31. warp/examples/benchmarks/benchmark_cloth.py +296 -0
  32. warp/examples/benchmarks/benchmark_cloth_cupy.py +96 -0
  33. warp/examples/benchmarks/benchmark_cloth_jax.py +105 -0
  34. warp/examples/benchmarks/benchmark_cloth_numba.py +161 -0
  35. warp/examples/benchmarks/benchmark_cloth_numpy.py +85 -0
  36. warp/examples/benchmarks/benchmark_cloth_paddle.py +94 -0
  37. warp/examples/benchmarks/benchmark_cloth_pytorch.py +94 -0
  38. warp/examples/benchmarks/benchmark_cloth_taichi.py +120 -0
  39. warp/examples/benchmarks/benchmark_cloth_warp.py +153 -0
  40. warp/examples/benchmarks/benchmark_gemm.py +164 -0
  41. warp/examples/benchmarks/benchmark_interop_paddle.py +166 -0
  42. warp/examples/benchmarks/benchmark_interop_torch.py +166 -0
  43. warp/examples/benchmarks/benchmark_launches.py +301 -0
  44. warp/examples/benchmarks/benchmark_tile_load_store.py +103 -0
  45. warp/examples/browse.py +37 -0
  46. warp/examples/core/example_cupy.py +86 -0
  47. warp/examples/core/example_dem.py +241 -0
  48. warp/examples/core/example_fluid.py +299 -0
  49. warp/examples/core/example_graph_capture.py +150 -0
  50. warp/examples/core/example_marching_cubes.py +194 -0
  51. warp/examples/core/example_mesh.py +180 -0
  52. warp/examples/core/example_mesh_intersect.py +211 -0
  53. warp/examples/core/example_nvdb.py +182 -0
  54. warp/examples/core/example_raycast.py +111 -0
  55. warp/examples/core/example_raymarch.py +205 -0
  56. warp/examples/core/example_render_opengl.py +193 -0
  57. warp/examples/core/example_sample_mesh.py +300 -0
  58. warp/examples/core/example_sph.py +411 -0
  59. warp/examples/core/example_torch.py +211 -0
  60. warp/examples/core/example_wave.py +269 -0
  61. warp/examples/fem/example_adaptive_grid.py +286 -0
  62. warp/examples/fem/example_apic_fluid.py +423 -0
  63. warp/examples/fem/example_burgers.py +261 -0
  64. warp/examples/fem/example_convection_diffusion.py +178 -0
  65. warp/examples/fem/example_convection_diffusion_dg.py +204 -0
  66. warp/examples/fem/example_deformed_geometry.py +172 -0
  67. warp/examples/fem/example_diffusion.py +196 -0
  68. warp/examples/fem/example_diffusion_3d.py +225 -0
  69. warp/examples/fem/example_diffusion_mgpu.py +220 -0
  70. warp/examples/fem/example_distortion_energy.py +228 -0
  71. warp/examples/fem/example_magnetostatics.py +240 -0
  72. warp/examples/fem/example_mixed_elasticity.py +291 -0
  73. warp/examples/fem/example_navier_stokes.py +261 -0
  74. warp/examples/fem/example_nonconforming_contact.py +298 -0
  75. warp/examples/fem/example_stokes.py +213 -0
  76. warp/examples/fem/example_stokes_transfer.py +262 -0
  77. warp/examples/fem/example_streamlines.py +352 -0
  78. warp/examples/fem/utils.py +1000 -0
  79. warp/examples/interop/example_jax_callable.py +116 -0
  80. warp/examples/interop/example_jax_ffi_callback.py +132 -0
  81. warp/examples/interop/example_jax_kernel.py +205 -0
  82. warp/examples/optim/example_bounce.py +266 -0
  83. warp/examples/optim/example_cloth_throw.py +228 -0
  84. warp/examples/optim/example_diffray.py +561 -0
  85. warp/examples/optim/example_drone.py +870 -0
  86. warp/examples/optim/example_fluid_checkpoint.py +497 -0
  87. warp/examples/optim/example_inverse_kinematics.py +182 -0
  88. warp/examples/optim/example_inverse_kinematics_torch.py +191 -0
  89. warp/examples/optim/example_softbody_properties.py +400 -0
  90. warp/examples/optim/example_spring_cage.py +245 -0
  91. warp/examples/optim/example_trajectory.py +227 -0
  92. warp/examples/sim/example_cartpole.py +143 -0
  93. warp/examples/sim/example_cloth.py +225 -0
  94. warp/examples/sim/example_cloth_self_contact.py +322 -0
  95. warp/examples/sim/example_granular.py +130 -0
  96. warp/examples/sim/example_granular_collision_sdf.py +202 -0
  97. warp/examples/sim/example_jacobian_ik.py +244 -0
  98. warp/examples/sim/example_particle_chain.py +124 -0
  99. warp/examples/sim/example_quadruped.py +203 -0
  100. warp/examples/sim/example_rigid_chain.py +203 -0
  101. warp/examples/sim/example_rigid_contact.py +195 -0
  102. warp/examples/sim/example_rigid_force.py +133 -0
  103. warp/examples/sim/example_rigid_gyroscopic.py +115 -0
  104. warp/examples/sim/example_rigid_soft_contact.py +140 -0
  105. warp/examples/sim/example_soft_body.py +196 -0
  106. warp/examples/tile/example_tile_cholesky.py +87 -0
  107. warp/examples/tile/example_tile_convolution.py +66 -0
  108. warp/examples/tile/example_tile_fft.py +55 -0
  109. warp/examples/tile/example_tile_filtering.py +113 -0
  110. warp/examples/tile/example_tile_matmul.py +85 -0
  111. warp/examples/tile/example_tile_mlp.py +383 -0
  112. warp/examples/tile/example_tile_nbody.py +199 -0
  113. warp/examples/tile/example_tile_walker.py +327 -0
  114. warp/fabric.py +355 -0
  115. warp/fem/__init__.py +106 -0
  116. warp/fem/adaptivity.py +508 -0
  117. warp/fem/cache.py +572 -0
  118. warp/fem/dirichlet.py +202 -0
  119. warp/fem/domain.py +411 -0
  120. warp/fem/field/__init__.py +125 -0
  121. warp/fem/field/field.py +619 -0
  122. warp/fem/field/nodal_field.py +326 -0
  123. warp/fem/field/restriction.py +37 -0
  124. warp/fem/field/virtual.py +848 -0
  125. warp/fem/geometry/__init__.py +32 -0
  126. warp/fem/geometry/adaptive_nanogrid.py +857 -0
  127. warp/fem/geometry/closest_point.py +84 -0
  128. warp/fem/geometry/deformed_geometry.py +221 -0
  129. warp/fem/geometry/element.py +776 -0
  130. warp/fem/geometry/geometry.py +362 -0
  131. warp/fem/geometry/grid_2d.py +392 -0
  132. warp/fem/geometry/grid_3d.py +452 -0
  133. warp/fem/geometry/hexmesh.py +911 -0
  134. warp/fem/geometry/nanogrid.py +571 -0
  135. warp/fem/geometry/partition.py +389 -0
  136. warp/fem/geometry/quadmesh.py +663 -0
  137. warp/fem/geometry/tetmesh.py +855 -0
  138. warp/fem/geometry/trimesh.py +806 -0
  139. warp/fem/integrate.py +2335 -0
  140. warp/fem/linalg.py +419 -0
  141. warp/fem/operator.py +293 -0
  142. warp/fem/polynomial.py +229 -0
  143. warp/fem/quadrature/__init__.py +17 -0
  144. warp/fem/quadrature/pic_quadrature.py +299 -0
  145. warp/fem/quadrature/quadrature.py +591 -0
  146. warp/fem/space/__init__.py +228 -0
  147. warp/fem/space/basis_function_space.py +468 -0
  148. warp/fem/space/basis_space.py +667 -0
  149. warp/fem/space/dof_mapper.py +251 -0
  150. warp/fem/space/function_space.py +309 -0
  151. warp/fem/space/grid_2d_function_space.py +177 -0
  152. warp/fem/space/grid_3d_function_space.py +227 -0
  153. warp/fem/space/hexmesh_function_space.py +257 -0
  154. warp/fem/space/nanogrid_function_space.py +201 -0
  155. warp/fem/space/partition.py +367 -0
  156. warp/fem/space/quadmesh_function_space.py +223 -0
  157. warp/fem/space/restriction.py +179 -0
  158. warp/fem/space/shape/__init__.py +143 -0
  159. warp/fem/space/shape/cube_shape_function.py +1105 -0
  160. warp/fem/space/shape/shape_function.py +133 -0
  161. warp/fem/space/shape/square_shape_function.py +926 -0
  162. warp/fem/space/shape/tet_shape_function.py +834 -0
  163. warp/fem/space/shape/triangle_shape_function.py +672 -0
  164. warp/fem/space/tetmesh_function_space.py +271 -0
  165. warp/fem/space/topology.py +424 -0
  166. warp/fem/space/trimesh_function_space.py +194 -0
  167. warp/fem/types.py +99 -0
  168. warp/fem/utils.py +420 -0
  169. warp/jax.py +187 -0
  170. warp/jax_experimental/__init__.py +16 -0
  171. warp/jax_experimental/custom_call.py +351 -0
  172. warp/jax_experimental/ffi.py +698 -0
  173. warp/jax_experimental/xla_ffi.py +602 -0
  174. warp/math.py +244 -0
  175. warp/native/array.h +1145 -0
  176. warp/native/builtin.h +1800 -0
  177. warp/native/bvh.cpp +492 -0
  178. warp/native/bvh.cu +791 -0
  179. warp/native/bvh.h +554 -0
  180. warp/native/clang/clang.cpp +536 -0
  181. warp/native/coloring.cpp +613 -0
  182. warp/native/crt.cpp +51 -0
  183. warp/native/crt.h +362 -0
  184. warp/native/cuda_crt.h +1058 -0
  185. warp/native/cuda_util.cpp +646 -0
  186. warp/native/cuda_util.h +307 -0
  187. warp/native/error.cpp +77 -0
  188. warp/native/error.h +36 -0
  189. warp/native/exports.h +1878 -0
  190. warp/native/fabric.h +245 -0
  191. warp/native/hashgrid.cpp +311 -0
  192. warp/native/hashgrid.cu +87 -0
  193. warp/native/hashgrid.h +240 -0
  194. warp/native/initializer_array.h +41 -0
  195. warp/native/intersect.h +1230 -0
  196. warp/native/intersect_adj.h +375 -0
  197. warp/native/intersect_tri.h +339 -0
  198. warp/native/marching.cpp +19 -0
  199. warp/native/marching.cu +514 -0
  200. warp/native/marching.h +19 -0
  201. warp/native/mat.h +2220 -0
  202. warp/native/mathdx.cpp +87 -0
  203. warp/native/matnn.h +343 -0
  204. warp/native/mesh.cpp +266 -0
  205. warp/native/mesh.cu +404 -0
  206. warp/native/mesh.h +1980 -0
  207. warp/native/nanovdb/GridHandle.h +366 -0
  208. warp/native/nanovdb/HostBuffer.h +590 -0
  209. warp/native/nanovdb/NanoVDB.h +6624 -0
  210. warp/native/nanovdb/PNanoVDB.h +3390 -0
  211. warp/native/noise.h +859 -0
  212. warp/native/quat.h +1371 -0
  213. warp/native/rand.h +342 -0
  214. warp/native/range.h +139 -0
  215. warp/native/reduce.cpp +174 -0
  216. warp/native/reduce.cu +364 -0
  217. warp/native/runlength_encode.cpp +79 -0
  218. warp/native/runlength_encode.cu +61 -0
  219. warp/native/scan.cpp +47 -0
  220. warp/native/scan.cu +53 -0
  221. warp/native/scan.h +23 -0
  222. warp/native/solid_angle.h +466 -0
  223. warp/native/sort.cpp +251 -0
  224. warp/native/sort.cu +277 -0
  225. warp/native/sort.h +33 -0
  226. warp/native/sparse.cpp +378 -0
  227. warp/native/sparse.cu +524 -0
  228. warp/native/spatial.h +657 -0
  229. warp/native/svd.h +702 -0
  230. warp/native/temp_buffer.h +46 -0
  231. warp/native/tile.h +2584 -0
  232. warp/native/tile_reduce.h +264 -0
  233. warp/native/vec.h +1426 -0
  234. warp/native/volume.cpp +501 -0
  235. warp/native/volume.cu +67 -0
  236. warp/native/volume.h +969 -0
  237. warp/native/volume_builder.cu +477 -0
  238. warp/native/volume_builder.h +52 -0
  239. warp/native/volume_impl.h +70 -0
  240. warp/native/warp.cpp +1082 -0
  241. warp/native/warp.cu +3636 -0
  242. warp/native/warp.h +381 -0
  243. warp/optim/__init__.py +17 -0
  244. warp/optim/adam.py +163 -0
  245. warp/optim/linear.py +1137 -0
  246. warp/optim/sgd.py +112 -0
  247. warp/paddle.py +407 -0
  248. warp/render/__init__.py +18 -0
  249. warp/render/render_opengl.py +3518 -0
  250. warp/render/render_usd.py +784 -0
  251. warp/render/utils.py +160 -0
  252. warp/sim/__init__.py +65 -0
  253. warp/sim/articulation.py +793 -0
  254. warp/sim/collide.py +2395 -0
  255. warp/sim/graph_coloring.py +300 -0
  256. warp/sim/import_mjcf.py +790 -0
  257. warp/sim/import_snu.py +227 -0
  258. warp/sim/import_urdf.py +579 -0
  259. warp/sim/import_usd.py +894 -0
  260. warp/sim/inertia.py +324 -0
  261. warp/sim/integrator.py +242 -0
  262. warp/sim/integrator_euler.py +1997 -0
  263. warp/sim/integrator_featherstone.py +2101 -0
  264. warp/sim/integrator_vbd.py +2048 -0
  265. warp/sim/integrator_xpbd.py +3292 -0
  266. warp/sim/model.py +4791 -0
  267. warp/sim/particles.py +121 -0
  268. warp/sim/render.py +427 -0
  269. warp/sim/utils.py +428 -0
  270. warp/sparse.py +2057 -0
  271. warp/stubs.py +3333 -0
  272. warp/tape.py +1203 -0
  273. warp/tests/__init__.py +1 -0
  274. warp/tests/__main__.py +4 -0
  275. warp/tests/assets/curlnoise_golden.npy +0 -0
  276. warp/tests/assets/mlp_golden.npy +0 -0
  277. warp/tests/assets/pixel.npy +0 -0
  278. warp/tests/assets/pnoise_golden.npy +0 -0
  279. warp/tests/assets/spiky.usd +0 -0
  280. warp/tests/assets/test_grid.nvdb +0 -0
  281. warp/tests/assets/test_index_grid.nvdb +0 -0
  282. warp/tests/assets/test_int32_grid.nvdb +0 -0
  283. warp/tests/assets/test_vec_grid.nvdb +0 -0
  284. warp/tests/assets/torus.nvdb +0 -0
  285. warp/tests/assets/torus.usda +105 -0
  286. warp/tests/aux_test_class_kernel.py +34 -0
  287. warp/tests/aux_test_compile_consts_dummy.py +18 -0
  288. warp/tests/aux_test_conditional_unequal_types_kernels.py +29 -0
  289. warp/tests/aux_test_dependent.py +29 -0
  290. warp/tests/aux_test_grad_customs.py +29 -0
  291. warp/tests/aux_test_instancing_gc.py +26 -0
  292. warp/tests/aux_test_module_unload.py +23 -0
  293. warp/tests/aux_test_name_clash1.py +40 -0
  294. warp/tests/aux_test_name_clash2.py +40 -0
  295. warp/tests/aux_test_reference.py +9 -0
  296. warp/tests/aux_test_reference_reference.py +8 -0
  297. warp/tests/aux_test_square.py +16 -0
  298. warp/tests/aux_test_unresolved_func.py +22 -0
  299. warp/tests/aux_test_unresolved_symbol.py +22 -0
  300. warp/tests/cuda/__init__.py +0 -0
  301. warp/tests/cuda/test_async.py +676 -0
  302. warp/tests/cuda/test_ipc.py +124 -0
  303. warp/tests/cuda/test_mempool.py +233 -0
  304. warp/tests/cuda/test_multigpu.py +169 -0
  305. warp/tests/cuda/test_peer.py +139 -0
  306. warp/tests/cuda/test_pinned.py +84 -0
  307. warp/tests/cuda/test_streams.py +634 -0
  308. warp/tests/geometry/__init__.py +0 -0
  309. warp/tests/geometry/test_bvh.py +200 -0
  310. warp/tests/geometry/test_hash_grid.py +221 -0
  311. warp/tests/geometry/test_marching_cubes.py +74 -0
  312. warp/tests/geometry/test_mesh.py +316 -0
  313. warp/tests/geometry/test_mesh_query_aabb.py +399 -0
  314. warp/tests/geometry/test_mesh_query_point.py +932 -0
  315. warp/tests/geometry/test_mesh_query_ray.py +311 -0
  316. warp/tests/geometry/test_volume.py +1103 -0
  317. warp/tests/geometry/test_volume_write.py +346 -0
  318. warp/tests/interop/__init__.py +0 -0
  319. warp/tests/interop/test_dlpack.py +729 -0
  320. warp/tests/interop/test_jax.py +371 -0
  321. warp/tests/interop/test_paddle.py +800 -0
  322. warp/tests/interop/test_torch.py +1001 -0
  323. warp/tests/run_coverage_serial.py +39 -0
  324. warp/tests/sim/__init__.py +0 -0
  325. warp/tests/sim/disabled_kinematics.py +244 -0
  326. warp/tests/sim/flaky_test_sim_grad.py +290 -0
  327. warp/tests/sim/test_collision.py +604 -0
  328. warp/tests/sim/test_coloring.py +258 -0
  329. warp/tests/sim/test_model.py +224 -0
  330. warp/tests/sim/test_sim_grad_bounce_linear.py +212 -0
  331. warp/tests/sim/test_sim_kinematics.py +98 -0
  332. warp/tests/sim/test_vbd.py +597 -0
  333. warp/tests/test_adam.py +163 -0
  334. warp/tests/test_arithmetic.py +1096 -0
  335. warp/tests/test_array.py +2972 -0
  336. warp/tests/test_array_reduce.py +156 -0
  337. warp/tests/test_assert.py +250 -0
  338. warp/tests/test_atomic.py +153 -0
  339. warp/tests/test_bool.py +220 -0
  340. warp/tests/test_builtins_resolution.py +1298 -0
  341. warp/tests/test_closest_point_edge_edge.py +327 -0
  342. warp/tests/test_codegen.py +810 -0
  343. warp/tests/test_codegen_instancing.py +1495 -0
  344. warp/tests/test_compile_consts.py +215 -0
  345. warp/tests/test_conditional.py +252 -0
  346. warp/tests/test_context.py +42 -0
  347. warp/tests/test_copy.py +238 -0
  348. warp/tests/test_ctypes.py +638 -0
  349. warp/tests/test_dense.py +73 -0
  350. warp/tests/test_devices.py +97 -0
  351. warp/tests/test_examples.py +482 -0
  352. warp/tests/test_fabricarray.py +996 -0
  353. warp/tests/test_fast_math.py +74 -0
  354. warp/tests/test_fem.py +2003 -0
  355. warp/tests/test_fp16.py +136 -0
  356. warp/tests/test_func.py +454 -0
  357. warp/tests/test_future_annotations.py +98 -0
  358. warp/tests/test_generics.py +656 -0
  359. warp/tests/test_grad.py +893 -0
  360. warp/tests/test_grad_customs.py +339 -0
  361. warp/tests/test_grad_debug.py +341 -0
  362. warp/tests/test_implicit_init.py +411 -0
  363. warp/tests/test_import.py +45 -0
  364. warp/tests/test_indexedarray.py +1140 -0
  365. warp/tests/test_intersect.py +73 -0
  366. warp/tests/test_iter.py +76 -0
  367. warp/tests/test_large.py +177 -0
  368. warp/tests/test_launch.py +411 -0
  369. warp/tests/test_lerp.py +151 -0
  370. warp/tests/test_linear_solvers.py +193 -0
  371. warp/tests/test_lvalue.py +427 -0
  372. warp/tests/test_mat.py +2089 -0
  373. warp/tests/test_mat_lite.py +122 -0
  374. warp/tests/test_mat_scalar_ops.py +2913 -0
  375. warp/tests/test_math.py +178 -0
  376. warp/tests/test_mlp.py +282 -0
  377. warp/tests/test_module_hashing.py +258 -0
  378. warp/tests/test_modules_lite.py +44 -0
  379. warp/tests/test_noise.py +252 -0
  380. warp/tests/test_operators.py +299 -0
  381. warp/tests/test_options.py +129 -0
  382. warp/tests/test_overwrite.py +551 -0
  383. warp/tests/test_print.py +339 -0
  384. warp/tests/test_quat.py +2315 -0
  385. warp/tests/test_rand.py +339 -0
  386. warp/tests/test_reload.py +302 -0
  387. warp/tests/test_rounding.py +185 -0
  388. warp/tests/test_runlength_encode.py +196 -0
  389. warp/tests/test_scalar_ops.py +105 -0
  390. warp/tests/test_smoothstep.py +108 -0
  391. warp/tests/test_snippet.py +318 -0
  392. warp/tests/test_sparse.py +582 -0
  393. warp/tests/test_spatial.py +2229 -0
  394. warp/tests/test_special_values.py +361 -0
  395. warp/tests/test_static.py +592 -0
  396. warp/tests/test_struct.py +734 -0
  397. warp/tests/test_tape.py +204 -0
  398. warp/tests/test_transient_module.py +93 -0
  399. warp/tests/test_triangle_closest_point.py +145 -0
  400. warp/tests/test_types.py +562 -0
  401. warp/tests/test_utils.py +588 -0
  402. warp/tests/test_vec.py +1487 -0
  403. warp/tests/test_vec_lite.py +80 -0
  404. warp/tests/test_vec_scalar_ops.py +2327 -0
  405. warp/tests/test_verify_fp.py +100 -0
  406. warp/tests/tile/__init__.py +0 -0
  407. warp/tests/tile/test_tile.py +780 -0
  408. warp/tests/tile/test_tile_load.py +407 -0
  409. warp/tests/tile/test_tile_mathdx.py +208 -0
  410. warp/tests/tile/test_tile_mlp.py +402 -0
  411. warp/tests/tile/test_tile_reduce.py +447 -0
  412. warp/tests/tile/test_tile_shared_memory.py +247 -0
  413. warp/tests/tile/test_tile_view.py +173 -0
  414. warp/tests/unittest_serial.py +47 -0
  415. warp/tests/unittest_suites.py +427 -0
  416. warp/tests/unittest_utils.py +468 -0
  417. warp/tests/walkthrough_debug.py +93 -0
  418. warp/thirdparty/__init__.py +0 -0
  419. warp/thirdparty/appdirs.py +598 -0
  420. warp/thirdparty/dlpack.py +145 -0
  421. warp/thirdparty/unittest_parallel.py +570 -0
  422. warp/torch.py +391 -0
  423. warp/types.py +5230 -0
  424. warp/utils.py +1137 -0
  425. warp_lang-1.7.0.dist-info/METADATA +516 -0
  426. warp_lang-1.7.0.dist-info/RECORD +429 -0
  427. warp_lang-1.7.0.dist-info/WHEEL +5 -0
  428. warp_lang-1.7.0.dist-info/licenses/LICENSE.md +202 -0
  429. warp_lang-1.7.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,932 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import math
17
+ import unittest
18
+
19
+ import numpy as np
20
+
21
+ import warp as wp
22
+ import warp.examples
23
+ from warp.tests.unittest_utils import *
24
+
25
+
26
+ @wp.kernel
27
+ def sample_mesh_query(
28
+ mesh: wp.uint64,
29
+ query_points: wp.array(dtype=wp.vec3),
30
+ query_faces: wp.array(dtype=int),
31
+ query_signs: wp.array(dtype=float),
32
+ query_dist: wp.array(dtype=float),
33
+ ):
34
+ tid = wp.tid()
35
+
36
+ face_index = int(0)
37
+ face_u = float(0.0)
38
+ face_v = float(0.0)
39
+ sign = float(0.0)
40
+
41
+ max_dist = 10012.0
42
+
43
+ p = query_points[tid]
44
+
45
+ wp.mesh_query_point(mesh, p, max_dist, sign, face_index, face_u, face_v)
46
+
47
+ cp = wp.mesh_eval_position(mesh, face_index, face_u, face_v)
48
+
49
+ query_signs[tid] = sign
50
+ query_faces[tid] = face_index
51
+ query_dist[tid] = wp.length(cp - p)
52
+
53
+ query = wp.mesh_query_point(mesh, p, max_dist)
54
+ wp.expect_eq(query.sign, sign)
55
+ wp.expect_eq(query.face, face_index)
56
+ wp.expect_eq(query.u, face_u)
57
+ wp.expect_eq(query.v, face_v)
58
+
59
+
60
+ @wp.kernel
61
+ def sample_mesh_query_no_sign(
62
+ mesh: wp.uint64,
63
+ query_points: wp.array(dtype=wp.vec3),
64
+ query_faces: wp.array(dtype=int),
65
+ query_dist: wp.array(dtype=float),
66
+ ):
67
+ tid = wp.tid()
68
+
69
+ face_index = int(0)
70
+ face_u = float(0.0)
71
+ face_v = float(0.0)
72
+
73
+ max_dist = 10012.0
74
+
75
+ p = query_points[tid]
76
+
77
+ wp.mesh_query_point_no_sign(mesh, p, max_dist, face_index, face_u, face_v)
78
+
79
+ cp = wp.mesh_eval_position(mesh, face_index, face_u, face_v)
80
+
81
+ query_faces[tid] = face_index
82
+ query_dist[tid] = wp.length(cp - p)
83
+
84
+ query = wp.mesh_query_point_no_sign(mesh, p, max_dist)
85
+ wp.expect_eq(query.face, face_index)
86
+ wp.expect_eq(query.u, face_u)
87
+ wp.expect_eq(query.v, face_v)
88
+
89
+
90
+ @wp.kernel
91
+ def sample_mesh_query_sign_normal(
92
+ mesh: wp.uint64,
93
+ query_points: wp.array(dtype=wp.vec3),
94
+ query_faces: wp.array(dtype=int),
95
+ query_signs: wp.array(dtype=float),
96
+ query_dist: wp.array(dtype=float),
97
+ ):
98
+ tid = wp.tid()
99
+
100
+ face_index = int(0)
101
+ face_u = float(0.0)
102
+ face_v = float(0.0)
103
+ sign = float(0.0)
104
+
105
+ max_dist = 10012.0
106
+
107
+ p = query_points[tid]
108
+
109
+ wp.mesh_query_point_sign_normal(mesh, p, max_dist, sign, face_index, face_u, face_v)
110
+
111
+ cp = wp.mesh_eval_position(mesh, face_index, face_u, face_v)
112
+
113
+ query_signs[tid] = sign
114
+ query_faces[tid] = face_index
115
+ query_dist[tid] = wp.length(cp - p)
116
+
117
+ query = wp.mesh_query_point_sign_normal(mesh, p, max_dist)
118
+ wp.expect_eq(query.sign, sign)
119
+ wp.expect_eq(query.face, face_index)
120
+ wp.expect_eq(query.u, face_u)
121
+ wp.expect_eq(query.v, face_v)
122
+
123
+
124
+ @wp.kernel
125
+ def sample_mesh_query_sign_winding_number(
126
+ mesh: wp.uint64,
127
+ query_points: wp.array(dtype=wp.vec3),
128
+ query_faces: wp.array(dtype=int),
129
+ query_signs: wp.array(dtype=float),
130
+ query_dist: wp.array(dtype=float),
131
+ ):
132
+ tid = wp.tid()
133
+
134
+ face_index = int(0)
135
+ face_u = float(0.0)
136
+ face_v = float(0.0)
137
+ sign = float(0.0)
138
+
139
+ max_dist = 10012.0
140
+
141
+ p = query_points[tid]
142
+
143
+ wp.mesh_query_point_sign_winding_number(mesh, p, max_dist, sign, face_index, face_u, face_v)
144
+
145
+ cp = wp.mesh_eval_position(mesh, face_index, face_u, face_v)
146
+
147
+ query_signs[tid] = sign
148
+ query_faces[tid] = face_index
149
+ query_dist[tid] = wp.length(cp - p)
150
+
151
+ query = wp.mesh_query_point_sign_winding_number(mesh, p, max_dist)
152
+ wp.expect_eq(query.sign, sign)
153
+ wp.expect_eq(query.face, face_index)
154
+ wp.expect_eq(query.u, face_u)
155
+ wp.expect_eq(query.v, face_v)
156
+
157
+
158
+ @wp.func
159
+ def triangle_closest_point(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
160
+ ab = b - a
161
+ ac = c - a
162
+ ap = p - a
163
+
164
+ d1 = wp.dot(ab, ap)
165
+ d2 = wp.dot(ac, ap)
166
+
167
+ if d1 <= 0.0 and d2 <= 0.0:
168
+ return wp.vec2(1.0, 0.0)
169
+
170
+ bp = p - b
171
+ d3 = wp.dot(ab, bp)
172
+ d4 = wp.dot(ac, bp)
173
+
174
+ if d3 >= 0.0 and d4 <= d3:
175
+ return wp.vec2(0.0, 1.0)
176
+
177
+ vc = d1 * d4 - d3 * d2
178
+ v = d1 / (d1 - d3)
179
+ if vc <= 0.0 and d1 >= 0.0 and d3 <= 0.0:
180
+ return wp.vec2(1.0 - v, v)
181
+
182
+ cp = p - c
183
+ d5 = wp.dot(ab, cp)
184
+ d6 = wp.dot(ac, cp)
185
+
186
+ if d6 >= 0.0 and d5 <= d6:
187
+ return wp.vec2(0.0, 0.0)
188
+
189
+ vb = d5 * d2 - d1 * d6
190
+ w = d2 / (d2 - d6)
191
+ if vb <= 0.0 and d2 >= 0.0 and d6 <= 0.0:
192
+ return wp.vec2(1.0 - w, 0.0)
193
+
194
+ va = d3 * d6 - d5 * d4
195
+ w = (d4 - d3) / ((d4 - d3) + (d5 - d6))
196
+ if va <= 0.0 and (d4 - d3) >= 0.0 and (d5 - d6) >= 0.0:
197
+ return wp.vec2(0.0, 1.0 - w)
198
+
199
+ denom = 1.0 / (va + vb + vc)
200
+ v = vb * denom
201
+ w = vc * denom
202
+ u = 1.0 - v - w
203
+
204
+ return wp.vec2(u, v)
205
+
206
+
207
+ @wp.func
208
+ def solid_angle(v0: wp.vec3, v1: wp.vec3, v2: wp.vec3, p: wp.vec3):
209
+ a = v0 - p
210
+ b = v1 - p
211
+ c = v2 - p
212
+
213
+ a_len = wp.length(a)
214
+ b_len = wp.length(b)
215
+ c_len = wp.length(c)
216
+
217
+ det = wp.dot(a, wp.cross(b, c))
218
+ den = a_len * b_len * c_len + wp.dot(a, b) * c_len + wp.dot(b, c) * a_len + wp.dot(c, a) * b_len
219
+
220
+ return 2.0 * wp.atan2(det, den)
221
+
222
+
223
+ @wp.kernel
224
+ def sample_mesh_brute(
225
+ tri_points: wp.array(dtype=wp.vec3),
226
+ tri_indices: wp.array(dtype=int),
227
+ tri_count: int,
228
+ query_points: wp.array(dtype=wp.vec3),
229
+ query_faces: wp.array(dtype=int),
230
+ query_signs: wp.array(dtype=float),
231
+ query_dist: wp.array(dtype=float),
232
+ ):
233
+ tid = wp.tid()
234
+
235
+ min_face = int(0)
236
+ min_dist = float(1.0e6)
237
+
238
+ sum_solid_angle = float(0.0)
239
+
240
+ p = query_points[tid]
241
+
242
+ for i in range(0, tri_count):
243
+ a = tri_points[tri_indices[i * 3 + 0]]
244
+ b = tri_points[tri_indices[i * 3 + 1]]
245
+ c = tri_points[tri_indices[i * 3 + 2]]
246
+
247
+ sum_solid_angle += solid_angle(a, b, c, p)
248
+
249
+ bary = triangle_closest_point(a, b, c, p)
250
+ u = bary[0]
251
+ v = bary[1]
252
+
253
+ cp = u * a + v * b + (1.0 - u - v) * c
254
+ cp_dist = wp.length(cp - p)
255
+
256
+ if cp_dist < min_dist:
257
+ min_dist = cp_dist
258
+ min_face = i
259
+
260
+ # for an inside point, the sum of the solid angle should be 4PI
261
+ # for an outside point, the sum should be 0
262
+ query_faces[tid] = min_face
263
+ query_signs[tid] = sum_solid_angle
264
+ query_dist[tid] = min_dist
265
+
266
+
267
+ # constructs a grid of evenly spaced particles
268
+ def particle_grid(dim_x, dim_y, dim_z, lower, radius, jitter):
269
+ rng = np.random.default_rng(123)
270
+
271
+ points = np.meshgrid(np.linspace(0, dim_x, dim_x), np.linspace(0, dim_y, dim_y), np.linspace(0, dim_z, dim_z))
272
+ points_t = np.array((points[0], points[1], points[2])).T * radius * 2.0 + np.array(lower)
273
+ points_t = points_t + rng.random(points_t.shape) * radius * jitter
274
+
275
+ return points_t.reshape((-1, 3))
276
+
277
+
278
+ # triangulate a list of polygon face indices
279
+ def triangulate(face_counts, face_indices):
280
+ num_tris = np.sum(np.subtract(face_counts, 2))
281
+ num_tri_vtx = num_tris * 3
282
+ tri_indices = np.zeros(num_tri_vtx, dtype=int)
283
+ ctr = 0
284
+ wedgeIdx = 0
285
+
286
+ for nb in face_counts:
287
+ for i in range(nb - 2):
288
+ tri_indices[ctr] = face_indices[wedgeIdx]
289
+ tri_indices[ctr + 1] = face_indices[wedgeIdx + i + 1]
290
+ tri_indices[ctr + 2] = face_indices[wedgeIdx + i + 2]
291
+ ctr += 3
292
+ wedgeIdx += nb
293
+
294
+ return tri_indices
295
+
296
+
297
+ @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
298
+ def test_mesh_query_point(test, device):
299
+ from pxr import Usd, UsdGeom
300
+
301
+ mesh = Usd.Stage.Open(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "spiky.usd")))
302
+ mesh_geom = UsdGeom.Mesh(mesh.GetPrimAtPath("/Cube/Cube"))
303
+
304
+ mesh_counts = mesh_geom.GetFaceVertexCountsAttr().Get()
305
+ mesh_indices = mesh_geom.GetFaceVertexIndicesAttr().Get()
306
+
307
+ tri_indices = triangulate(mesh_counts, mesh_indices)
308
+
309
+ mesh_points = wp.array(np.array(mesh_geom.GetPointsAttr().Get()), dtype=wp.vec3, device=device)
310
+ mesh_indices = wp.array(np.array(tri_indices), dtype=int, device=device)
311
+
312
+ if device.is_cpu:
313
+ constructors = ["sah", "median"]
314
+ else:
315
+ constructors = ["sah", "median", "lbvh"]
316
+
317
+ for constructor in constructors:
318
+ # create mesh
319
+ mesh = wp.Mesh(
320
+ points=mesh_points,
321
+ velocities=None,
322
+ indices=mesh_indices,
323
+ support_winding_number=True,
324
+ bvh_constructor=constructor,
325
+ )
326
+
327
+ p = particle_grid(32, 32, 32, np.array([-1.1, -1.1, -1.1]), 0.05, 0.0)
328
+
329
+ query_count = len(p)
330
+ query_points = wp.array(p, dtype=wp.vec3, device=device)
331
+
332
+ signs_query = wp.zeros(query_count, dtype=float, device=device)
333
+ faces_query = wp.zeros(query_count, dtype=int, device=device)
334
+ dist_query = wp.zeros(query_count, dtype=float, device=device)
335
+
336
+ faces_query_no_sign = wp.zeros(query_count, dtype=int, device=device)
337
+ dist_query_no_sign = wp.zeros(query_count, dtype=float, device=device)
338
+
339
+ signs_query_normal = wp.zeros(query_count, dtype=float, device=device)
340
+ faces_query_normal = wp.zeros(query_count, dtype=int, device=device)
341
+ dist_query_normal = wp.zeros(query_count, dtype=float, device=device)
342
+
343
+ signs_query_winding_number = wp.zeros(query_count, dtype=float, device=device)
344
+ faces_query_winding_number = wp.zeros(query_count, dtype=int, device=device)
345
+ dist_query_winding_number = wp.zeros(query_count, dtype=float, device=device)
346
+
347
+ signs_brute = wp.zeros(query_count, dtype=float, device=device)
348
+ faces_brute = wp.zeros(query_count, dtype=int, device=device)
349
+ dist_brute = wp.zeros(query_count, dtype=float, device=device)
350
+
351
+ wp.launch(
352
+ kernel=sample_mesh_query,
353
+ dim=query_count,
354
+ inputs=[mesh.id, query_points, faces_query, signs_query, dist_query],
355
+ device=device,
356
+ )
357
+
358
+ wp.launch(
359
+ kernel=sample_mesh_query_no_sign,
360
+ dim=query_count,
361
+ inputs=[mesh.id, query_points, faces_query_no_sign, dist_query_no_sign],
362
+ device=device,
363
+ )
364
+
365
+ wp.launch(
366
+ kernel=sample_mesh_query_sign_normal,
367
+ dim=query_count,
368
+ inputs=[mesh.id, query_points, faces_query_normal, signs_query_normal, dist_query_normal],
369
+ device=device,
370
+ )
371
+
372
+ wp.launch(
373
+ kernel=sample_mesh_query_sign_winding_number,
374
+ dim=query_count,
375
+ inputs=[
376
+ mesh.id,
377
+ query_points,
378
+ faces_query_winding_number,
379
+ signs_query_winding_number,
380
+ dist_query_winding_number,
381
+ ],
382
+ device=device,
383
+ )
384
+
385
+ wp.launch(
386
+ kernel=sample_mesh_brute,
387
+ dim=query_count,
388
+ inputs=[
389
+ mesh_points,
390
+ mesh_indices,
391
+ int(len(mesh_indices) / 3),
392
+ query_points,
393
+ faces_brute,
394
+ signs_brute,
395
+ dist_brute,
396
+ ],
397
+ device=device,
398
+ )
399
+
400
+ signs_query = signs_query.numpy()
401
+ faces_query = faces_query.numpy()
402
+ dist_query = dist_query.numpy()
403
+
404
+ faces_query_no_sign = faces_query_no_sign.numpy()
405
+ dist_query_no_sign = dist_query_no_sign.numpy()
406
+
407
+ signs_query_normal = signs_query_normal.numpy()
408
+ faces_query_normal = faces_query_normal.numpy()
409
+ dist_query_normal = dist_query_normal.numpy()
410
+
411
+ signs_query_winding_number = signs_query_winding_number.numpy()
412
+ faces_query_winding_number = faces_query_winding_number.numpy()
413
+ dist_query_winding_number = dist_query_winding_number.numpy()
414
+
415
+ signs_brute = signs_brute.numpy()
416
+ faces_brute = faces_brute.numpy()
417
+ dist_brute = dist_brute.numpy()
418
+
419
+ query_points = query_points.numpy()
420
+
421
+ inside_query = [[0.0, 0.0, 0.0]]
422
+ inside_query_normal = [[0.0, 0.0, 0.0]]
423
+ inside_query_winding_number = [[0.0, 0.0, 0.0]]
424
+ inside_brute = [[0.0, 0.0, 0.0]]
425
+
426
+ for i in range(query_count):
427
+ if signs_query[i] < 0.0:
428
+ inside_query.append(query_points[i].tolist())
429
+
430
+ if signs_query_normal[i] < 0.0:
431
+ inside_query_normal.append(query_points[i].tolist())
432
+
433
+ if signs_query_winding_number[i] < 0.0:
434
+ inside_query_winding_number.append(query_points[i].tolist())
435
+
436
+ if signs_brute[i] > math.pi * 2.0:
437
+ inside_brute.append(query_points[i].tolist())
438
+
439
+ inside_query = np.array(inside_query)
440
+ inside_query_normal = np.array(inside_query_normal)
441
+ inside_query_winding_number = np.array(inside_query_winding_number)
442
+ inside_brute = np.array(inside_brute)
443
+
444
+ # import warp.render
445
+
446
+ # stage = warp.render.UsdRenderer("tests/outputs/test_mesh_query_point.usd")
447
+
448
+ # radius = 0.1
449
+ # stage.begin_frame(0.0)
450
+ # stage.render_mesh(points=mesh_points.numpy(), indices=mesh_indices.numpy(), name="mesh")
451
+ # stage.render_points(points=inside_query, radius=radius, name="query")
452
+ # stage.render_points(points=inside_brute, radius=radius, name="brute")
453
+ # stage.render_points(points=query_points, radius=radius, name="all")
454
+ # stage.end_frame()
455
+
456
+ # stage.save()
457
+
458
+ test.assertTrue(len(inside_query) == len(inside_brute))
459
+ test.assertTrue(len(inside_query_normal) == len(inside_brute))
460
+ test.assertTrue(len(inside_query_winding_number) == len(inside_brute))
461
+
462
+ tolerance = 1.5e-4
463
+ dist_error = np.max(np.abs(dist_query - dist_brute))
464
+ sign_error = np.max(np.abs(inside_query - inside_brute))
465
+
466
+ test.assertTrue(dist_error < tolerance, f"mesh_query_point dist_error is {dist_error} which is >= {tolerance}")
467
+ test.assertTrue(sign_error < tolerance, f"mesh_query_point sign_error is {sign_error} which is >= {tolerance}")
468
+
469
+ dist_error = np.max(np.abs(dist_query_no_sign - dist_brute))
470
+
471
+ test.assertTrue(
472
+ dist_error < tolerance, f"mesh_query_point_no_sign dist_error is {dist_error} which is >= {tolerance}"
473
+ )
474
+
475
+ dist_error = np.max(np.abs(dist_query_normal - dist_brute))
476
+ sign_error = np.max(np.abs(inside_query_normal - inside_brute))
477
+
478
+ test.assertTrue(
479
+ dist_error < tolerance, f"mesh_query_point_sign_normal dist_error is {dist_error} which is >= {tolerance}"
480
+ )
481
+ test.assertTrue(
482
+ sign_error < tolerance, f"mesh_query_point_sign_normal sign_error is {sign_error} which is >= {tolerance}"
483
+ )
484
+
485
+ dist_error = np.max(np.abs(dist_query_winding_number - dist_brute))
486
+ sign_error = np.max(np.abs(inside_query_winding_number - inside_brute))
487
+
488
+ test.assertTrue(
489
+ dist_error < tolerance,
490
+ f"mesh_query_point_sign_winding_number dist_error is {dist_error} which is >= {tolerance}",
491
+ )
492
+ test.assertTrue(
493
+ sign_error < tolerance,
494
+ f"mesh_query_point_sign_winding_number sign_error is {sign_error} which is >= {tolerance}",
495
+ )
496
+
497
+
498
+ @wp.kernel
499
+ def mesh_query_point_loss(
500
+ mesh: wp.uint64,
501
+ query_points: wp.array(dtype=wp.vec3),
502
+ projected_points: wp.array(dtype=wp.vec3),
503
+ loss: wp.array(dtype=float),
504
+ ):
505
+ tid = wp.tid()
506
+
507
+ face_index = int(0)
508
+ face_u = float(0.0)
509
+ face_v = float(0.0)
510
+ sign = float(0.0)
511
+
512
+ max_dist = 10012.0
513
+
514
+ p = query_points[tid]
515
+
516
+ wp.mesh_query_point(mesh, p, max_dist, sign, face_index, face_u, face_v)
517
+ q = wp.mesh_eval_position(mesh, face_index, face_u, face_v)
518
+
519
+ projected_points[tid] = q
520
+
521
+ dist = wp.length(wp.sub(p, q))
522
+ loss[tid] = dist
523
+
524
+
525
+ @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
526
+ def test_adj_mesh_query_point(test, device):
527
+ from pxr import Usd, UsdGeom
528
+
529
+ mesh = Usd.Stage.Open(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "torus.usda")))
530
+ mesh_geom = UsdGeom.Mesh(mesh.GetPrimAtPath("/World/Torus"))
531
+
532
+ mesh_counts = mesh_geom.GetFaceVertexCountsAttr().Get()
533
+ mesh_indices = mesh_geom.GetFaceVertexIndicesAttr().Get()
534
+
535
+ tri_indices = triangulate(mesh_counts, mesh_indices)
536
+
537
+ mesh_points = wp.array(np.array(mesh_geom.GetPointsAttr().Get()), dtype=wp.vec3, device=device)
538
+ mesh_indices = wp.array(np.array(tri_indices), dtype=int, device=device)
539
+
540
+ if device.is_cpu:
541
+ constructors = ["sah", "median"]
542
+ else:
543
+ constructors = ["sah", "median", "lbvh"]
544
+
545
+ for constructor in constructors:
546
+ # test tri
547
+ # print("Testing Single Triangle")
548
+ # mesh_points = wp.array(np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 2.0, 0.0]]), dtype=wp.vec3, device=device)
549
+ # mesh_indices = wp.array(np.array([0,1,2]), dtype=int, device=device)
550
+
551
+ # create mesh
552
+ mesh = wp.Mesh(
553
+ points=mesh_points,
554
+ velocities=None,
555
+ indices=mesh_indices,
556
+ support_winding_number=True,
557
+ bvh_constructor=constructor,
558
+ )
559
+
560
+ # p = particle_grid(32, 32, 32, np.array([-5.0, -5.0, -5.0]), 0.1, 0.1)*100.0
561
+ p = wp.vec3(50.0, 50.0, 50.0)
562
+
563
+ tape = wp.Tape()
564
+
565
+ # analytic gradients
566
+ with tape:
567
+ query_points = wp.array(p, dtype=wp.vec3, device=device, requires_grad=True)
568
+ projected_points = wp.zeros(n=1, dtype=wp.vec3, device=device)
569
+ loss = wp.zeros(n=1, dtype=float, device=device, requires_grad=True)
570
+
571
+ wp.launch(
572
+ kernel=mesh_query_point_loss,
573
+ dim=1,
574
+ inputs=[mesh.id, query_points, projected_points, loss],
575
+ device=device,
576
+ )
577
+
578
+ tape.backward(loss=loss)
579
+ analytic = tape.gradients[query_points].numpy().flatten()
580
+
581
+ # numeric gradients
582
+ eps = 1.0e-3
583
+ loss_values = []
584
+ numeric = np.zeros(3)
585
+
586
+ offset_query_points = [
587
+ wp.vec3(p[0] - eps, p[1], p[2]),
588
+ wp.vec3(p[0] + eps, p[1], p[2]),
589
+ wp.vec3(p[0], p[1] - eps, p[2]),
590
+ wp.vec3(p[0], p[1] + eps, p[2]),
591
+ wp.vec3(p[0], p[1], p[2] - eps),
592
+ wp.vec3(p[0], p[1], p[2] + eps),
593
+ ]
594
+
595
+ for i in range(6):
596
+ q = offset_query_points[i]
597
+
598
+ query_points = wp.array(q, dtype=wp.vec3, device=device)
599
+ projected_points = wp.zeros(n=1, dtype=wp.vec3, device=device)
600
+ loss = wp.zeros(n=1, dtype=float, device=device)
601
+
602
+ wp.launch(
603
+ kernel=mesh_query_point_loss,
604
+ dim=1,
605
+ inputs=[mesh.id, query_points, projected_points, loss],
606
+ device=device,
607
+ )
608
+
609
+ loss_values.append(loss.numpy()[0])
610
+
611
+ for i in range(3):
612
+ l_0 = loss_values[i * 2]
613
+ l_1 = loss_values[i * 2 + 1]
614
+ gradient = (l_1 - l_0) / (2.0 * eps)
615
+ numeric[i] = gradient
616
+
617
+ error = ((analytic - numeric) * (analytic - numeric)).sum(axis=0)
618
+
619
+ tolerance = 1.0e-3
620
+ test.assertTrue(error < tolerance, f"error is {error} which is >= {tolerance}")
621
+
622
+
623
+ @wp.kernel
624
+ def sample_furthest_points(mesh: wp.uint64, query_points: wp.array(dtype=wp.vec3), query_result: wp.array(dtype=float)):
625
+ tid = wp.tid()
626
+
627
+ p = query_points[tid]
628
+
629
+ face = int(0)
630
+ bary_u = float(0.0)
631
+ bary_v = float(0.0)
632
+
633
+ if wp.mesh_query_furthest_point_no_sign(mesh, p, 0.0, face, bary_u, bary_v):
634
+ closest = wp.mesh_eval_position(mesh, face, bary_u, bary_v)
635
+
636
+ query_result[tid] = wp.length_sq(p - closest)
637
+
638
+ query = wp.mesh_query_furthest_point_no_sign(mesh, p, 0.0)
639
+ wp.expect_eq(query.face, face)
640
+ wp.expect_eq(query.u, bary_u)
641
+ wp.expect_eq(query.v, bary_v)
642
+
643
+
644
+ @wp.kernel
645
+ def sample_furthest_points_brute(
646
+ mesh_points: wp.array(dtype=wp.vec3), query_points: wp.array(dtype=wp.vec3), query_result: wp.array(dtype=float)
647
+ ):
648
+ tid = wp.tid()
649
+
650
+ p = query_points[tid]
651
+ max_dist_sq = float(0.0)
652
+
653
+ for i in range(mesh_points.shape[0]):
654
+ dist_sq = wp.length_sq(p - mesh_points[i])
655
+
656
+ if dist_sq > max_dist_sq:
657
+ max_dist_sq = dist_sq
658
+
659
+ query_result[tid] = max_dist_sq
660
+
661
+
662
+ @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
663
+ def test_mesh_query_furthest_point(test, device):
664
+ from pxr import Usd, UsdGeom
665
+
666
+ mesh = Usd.Stage.Open(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "assets", "spiky.usd")))
667
+ mesh_geom = UsdGeom.Mesh(mesh.GetPrimAtPath("/Cube/Cube"))
668
+
669
+ mesh_counts = mesh_geom.GetFaceVertexCountsAttr().Get()
670
+ mesh_indices = mesh_geom.GetFaceVertexIndicesAttr().Get()
671
+
672
+ tri_indices = triangulate(mesh_counts, mesh_indices)
673
+
674
+ mesh_points = wp.array(np.array(mesh_geom.GetPointsAttr().Get()), dtype=wp.vec3, device=device)
675
+ mesh_indices = wp.array(np.array(tri_indices), dtype=int, device=device)
676
+
677
+ if device.is_cpu:
678
+ constructors = ["sah", "median"]
679
+ else:
680
+ constructors = ["sah", "median", "lbvh"]
681
+
682
+ for constructor in constructors:
683
+ # create mesh
684
+ mesh = wp.Mesh(points=mesh_points, indices=mesh_indices, bvh_constructor=constructor)
685
+
686
+ p = particle_grid(32, 32, 32, np.array([-1.1, -1.1, -1.1]), 0.05, 0.0)
687
+
688
+ query_count = len(p)
689
+ query_points = wp.array(p, dtype=wp.vec3, device=device)
690
+
691
+ dist_query = wp.zeros(query_count, dtype=float, device=device)
692
+ dist_brute = wp.zeros(query_count, dtype=float, device=device)
693
+
694
+ wp.launch(sample_furthest_points, dim=query_count, inputs=[mesh.id, query_points, dist_query], device=device)
695
+ wp.launch(
696
+ sample_furthest_points_brute, dim=query_count, inputs=[mesh_points, query_points, dist_brute], device=device
697
+ )
698
+
699
+ assert_np_equal(dist_query.numpy(), dist_brute.numpy(), tol=1.0e-3)
700
+
701
+
702
+ @wp.func
703
+ def triangle_closest_point_for_test(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
704
+ ab = b - a
705
+ ac = c - a
706
+ ap = p - a
707
+
708
+ d1 = wp.dot(ab, ap)
709
+ d2 = wp.dot(ac, ap)
710
+ if d1 <= 0.0 and d2 <= 0.0:
711
+ bary = wp.vec3(1.0, 0.0, 0.0)
712
+ return a, bary
713
+
714
+ bp = p - b
715
+ d3 = wp.dot(ab, bp)
716
+ d4 = wp.dot(ac, bp)
717
+ if d3 >= 0.0 and d4 <= d3:
718
+ bary = wp.vec3(0.0, 1.0, 0.0)
719
+ return b, bary
720
+
721
+ cp = p - c
722
+ d5 = wp.dot(ab, cp)
723
+ d6 = wp.dot(ac, cp)
724
+ if d6 >= 0.0 and d5 <= d6:
725
+ bary = wp.vec3(0.0, 0.0, 1.0)
726
+ return c, bary
727
+
728
+ vc = d1 * d4 - d3 * d2
729
+ if vc <= 0.0 and d1 >= 0.0 and d3 <= 0.0:
730
+ v = d1 / (d1 - d3)
731
+ bary = wp.vec3(1.0 - v, v, 0.0)
732
+ return a + v * ab, bary
733
+
734
+ vb = d5 * d2 - d1 * d6
735
+ if vb <= 0.0 and d2 >= 0.0 and d6 <= 0.0:
736
+ v = d2 / (d2 - d6)
737
+ bary = wp.vec3(1.0 - v, 0.0, v)
738
+ return a + v * ac, bary
739
+
740
+ va = d3 * d6 - d5 * d4
741
+ if va <= 0.0 and (d4 - d3) >= 0.0 and (d5 - d6) >= 0.0:
742
+ v = (d4 - d3) / ((d4 - d3) + (d5 - d6))
743
+ bary = wp.vec3(0.0, 1.0 - v, v)
744
+ return b + v * (c - b), bary
745
+
746
+ denom = 1.0 / (va + vb + vc)
747
+ v = vb * denom
748
+ w = vc * denom
749
+ bary = wp.vec3(1.0 - v - w, v, w)
750
+ return a + v * ab + w * ac, bary
751
+
752
+
753
+ def load_mesh(model_name="bunny"):
754
+ from pxr import Usd, UsdGeom
755
+
756
+ usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), model_name + ".usd"))
757
+ usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/" + model_name))
758
+
759
+ vertices = np.array(usd_geom.GetPointsAttr().Get())
760
+ faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
761
+
762
+ return vertices, faces
763
+
764
+
765
+ @wp.kernel
766
+ def point_query_aabb_and_closest(
767
+ query_radius: float,
768
+ mesh_id: wp.uint64,
769
+ pts: wp.array(dtype=wp.vec3),
770
+ pos: wp.array(dtype=wp.vec3),
771
+ tri_indices: wp.array(dtype=wp.int32, ndim=2),
772
+ query_results_num_collisions: wp.array(dtype=wp.int32),
773
+ query_results_min_dist: wp.array(dtype=float),
774
+ query_results_closest_point_velocity: wp.array(dtype=wp.vec3),
775
+ ):
776
+ p_index = wp.tid()
777
+ p = pts[p_index]
778
+
779
+ lower = wp.vec3(p[0] - query_radius, p[1] - query_radius, p[2] - query_radius)
780
+ upper = wp.vec3(p[0] + query_radius, p[1] + query_radius, p[2] + query_radius)
781
+
782
+ closest_query = wp.mesh_query_point_no_sign(mesh_id, p, query_radius)
783
+ if closest_query.result:
784
+ closest_p = wp.mesh_eval_position(mesh_id, closest_query.face, closest_query.u, closest_query.v)
785
+ closest_p_vel = wp.mesh_eval_velocity(mesh_id, closest_query.face, closest_query.u, closest_query.v)
786
+
787
+ query_results_min_dist[p_index] = wp.length(closest_p - p)
788
+ query_results_closest_point_velocity[p_index] = closest_p_vel
789
+
790
+ query = wp.mesh_query_aabb(mesh_id, lower, upper)
791
+
792
+ tri_index = wp.int32(0)
793
+ num_collisions = wp.int32(0)
794
+ min_dis_to_tris = query_radius
795
+ while wp.mesh_query_aabb_next(query, tri_index):
796
+ t1 = tri_indices[tri_index, 0]
797
+ t2 = tri_indices[tri_index, 1]
798
+ t3 = tri_indices[tri_index, 2]
799
+
800
+ u1 = pos[t1]
801
+ u2 = pos[t2]
802
+ u3 = pos[t3]
803
+
804
+ closest_p1, barycentric1 = triangle_closest_point_for_test(u1, u2, u3, p)
805
+
806
+ dis = wp.length(closest_p1 - p)
807
+
808
+ if dis < query_radius:
809
+ num_collisions = num_collisions + 1
810
+
811
+ query_results_num_collisions[p_index] = num_collisions
812
+
813
+
814
+ @unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
815
+ def test_set_mesh_points(test, device):
816
+ if device.is_cpu:
817
+ constructors = ["sah", "median"]
818
+ else:
819
+ constructors = ["sah", "median", "lbvh"]
820
+
821
+ rng = np.random.default_rng(123)
822
+
823
+ models = ["bunny", "nonuniform"]
824
+
825
+ for model in models:
826
+ vs, fs = load_mesh(model)
827
+ vertices1 = wp.array(vs, dtype=wp.vec3, device=device)
828
+ velocities1_np = rng.standard_normal(size=(vertices1.shape[0], 3))
829
+ velocities1 = wp.array(velocities1_np, dtype=wp.vec3, device=device)
830
+
831
+ faces = wp.array(fs, dtype=wp.int32, device=device)
832
+
833
+ n = 1000
834
+ query_radius = 0.2
835
+ pts1 = wp.array(rng.standard_normal(size=(n, 3)), dtype=wp.vec3, device=device)
836
+
837
+ query_results_num_cols1 = wp.zeros(n, dtype=wp.int32, device=device)
838
+ query_results_min_dist1 = wp.zeros(n, dtype=float, device=device)
839
+ query_results_closest_point_velocity1 = wp.zeros(n, dtype=wp.vec3, device=device)
840
+
841
+ for constructor in constructors:
842
+ mesh = wp.Mesh(vertices1, faces, velocities=velocities1, bvh_constructor=constructor)
843
+ fs_2D = faces.reshape((-1, 3))
844
+
845
+ wp.launch(
846
+ kernel=point_query_aabb_and_closest,
847
+ inputs=[
848
+ query_radius,
849
+ mesh.id,
850
+ pts1,
851
+ vertices1,
852
+ fs_2D,
853
+ query_results_num_cols1,
854
+ query_results_min_dist1,
855
+ query_results_closest_point_velocity1,
856
+ ],
857
+ dim=n,
858
+ device=device,
859
+ )
860
+
861
+ shift = rng.standard_normal(size=3)
862
+
863
+ vs_higher = vs + shift
864
+ vertices2 = wp.array(vs_higher, dtype=wp.vec3, device=device)
865
+
866
+ velocities2_np = velocities1_np + shift[None, ...]
867
+ velocities2 = wp.array(velocities2_np, dtype=wp.vec3, device=device)
868
+
869
+ pts2 = wp.array(pts1.numpy() + shift, dtype=wp.vec3, device=device)
870
+
871
+ mesh.points = vertices2
872
+ mesh.velocities = velocities2
873
+
874
+ query_results_num_cols2 = wp.zeros(n, dtype=wp.int32, device=device)
875
+ query_results_min_dist2 = wp.zeros(n, dtype=float, device=device)
876
+ query_results_closest_point_velocity2 = wp.array([shift for i in range(n)], dtype=wp.vec3, device=device)
877
+
878
+ wp.launch(
879
+ kernel=point_query_aabb_and_closest,
880
+ inputs=[
881
+ query_radius,
882
+ mesh.id,
883
+ pts2,
884
+ vertices2,
885
+ fs_2D,
886
+ query_results_num_cols2,
887
+ query_results_min_dist2,
888
+ query_results_closest_point_velocity2,
889
+ ],
890
+ dim=n,
891
+ device=device,
892
+ )
893
+
894
+ test.assertTrue((query_results_num_cols1.numpy() == query_results_num_cols2.numpy()).all())
895
+ test.assertTrue(((query_results_min_dist1.numpy() - query_results_min_dist2.numpy()) < 1e-5).all())
896
+
897
+
898
+ devices = get_test_devices()
899
+
900
+
901
+ class TestMeshQueryPoint(unittest.TestCase):
902
+ def test_mesh_query_point_codegen_adjoints_with_select(self):
903
+ def kernel_fn(
904
+ mesh: wp.uint64,
905
+ ):
906
+ v = wp.vec3(0.0, 0.0, 0.0)
907
+ d = 1e-6
908
+
909
+ if True:
910
+ query_1 = wp.mesh_query_point(mesh, v, d)
911
+ query_2 = wp.mesh_query_point_no_sign(mesh, v, d)
912
+ query_3 = wp.mesh_query_furthest_point_no_sign(mesh, v, d)
913
+ query_4 = wp.mesh_query_point_sign_normal(mesh, v, d)
914
+ query_5 = wp.mesh_query_point_sign_winding_number(mesh, v, d)
915
+ else:
916
+ query_1 = wp.mesh_query_point(mesh, v, d)
917
+ query_2 = wp.mesh_query_point_no_sign(mesh, v, d)
918
+ query_3 = wp.mesh_query_furthest_point_no_sign(mesh, v, d)
919
+ query_4 = wp.mesh_query_point_sign_normal(mesh, v, d)
920
+ query_5 = wp.mesh_query_point_sign_winding_number(mesh, v, d)
921
+
922
+ wp.Kernel(func=kernel_fn)
923
+
924
+
925
+ add_function_test(TestMeshQueryPoint, "test_mesh_query_point", test_mesh_query_point, devices=devices)
926
+ add_function_test(TestMeshQueryPoint, "test_mesh_query_furthest_point", test_mesh_query_furthest_point, devices=devices)
927
+ add_function_test(TestMeshQueryPoint, "test_adj_mesh_query_point", test_adj_mesh_query_point, devices=devices)
928
+ add_function_test(TestMeshQueryPoint, "test_set_mesh_points", test_set_mesh_points, devices=devices)
929
+
930
+ if __name__ == "__main__":
931
+ wp.clear_kernel_cache()
932
+ unittest.main(verbosity=2)