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,571 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024 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
+ from typing import Optional
17
+
18
+ import numpy as np
19
+
20
+ import warp as wp
21
+ from warp.fem import cache, utils
22
+ from warp.fem.linalg import basis_element
23
+ from warp.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
24
+
25
+ from .element import Cube, Square
26
+ from .geometry import Geometry
27
+
28
+ # Flag used for building edge/face grids to disambiguiate axis within the grid
29
+ # Morton indexing allows for
30
+ GRID_AXIS_FLAG = wp.constant(wp.int32(1 << 20))
31
+
32
+ FACE_AXIS_MASK = wp.constant(wp.uint8((1 << 2) - 1))
33
+ FACE_INNER_OFFSET_BIT = wp.constant(wp.uint8(2))
34
+ FACE_OUTER_OFFSET_BIT = wp.constant(wp.uint8(3))
35
+
36
+
37
+ @wp.func
38
+ def _add_axis_flag(ijk: wp.vec3i, axis: int):
39
+ coord = ijk[axis]
40
+ ijk[axis] = wp.where(coord < 0, coord & (~GRID_AXIS_FLAG), coord | GRID_AXIS_FLAG)
41
+ return ijk
42
+
43
+
44
+ @wp.func
45
+ def _extract_axis_flag(ijk: wp.vec3i):
46
+ for ax in range(3):
47
+ coord = ijk[ax]
48
+ if coord < 0:
49
+ if (ijk[ax] & GRID_AXIS_FLAG) == 0:
50
+ ijk[ax] = ijk[ax] | GRID_AXIS_FLAG
51
+ return ax, ijk
52
+ else:
53
+ if (ijk[ax] & GRID_AXIS_FLAG) != 0:
54
+ ijk[ax] = ijk[ax] & (~GRID_AXIS_FLAG)
55
+ return ax, ijk
56
+
57
+ return -1, ijk
58
+
59
+
60
+ @wp.struct
61
+ class NanogridCellArg:
62
+ # Utility device functions
63
+ cell_grid: wp.uint64
64
+ cell_ijk: wp.array(dtype=wp.vec3i)
65
+ inverse_transform: wp.mat33
66
+ cell_volume: float
67
+
68
+
69
+ @wp.struct
70
+ class NanogridSideArg:
71
+ # Utility device functions
72
+ cell_arg: NanogridCellArg
73
+ face_ijk: wp.array(dtype=wp.vec3i)
74
+ face_flags: wp.array(dtype=wp.uint8)
75
+ face_areas: wp.vec3
76
+
77
+
78
+ class Nanogrid(Geometry):
79
+ """Sparse grid geometry"""
80
+
81
+ dimension = 3
82
+
83
+ def __init__(self, grid: wp.Volume, temporary_store: Optional[cache.TemporaryStore] = None):
84
+ """
85
+ Constructs a sparse grid geometry from an in-memory NanoVDB volume.
86
+
87
+ Args:
88
+ grid: The NanoVDB volume. Any type is accepted, but for indexing efficiency an index grid is recommended.
89
+ If `grid` is an 'on' index grid, cells will be created for active voxels only, otherwise cells will
90
+ be created for all leaf voxels.
91
+ temporary_store: shared pool from which to allocate temporary arrays
92
+ """
93
+
94
+ self._cell_grid = grid
95
+ self._cell_grid_info = grid.get_grid_info()
96
+
97
+ device = grid.device
98
+
99
+ cell_count = grid.get_voxel_count()
100
+ self._cell_ijk = wp.array(shape=(cell_count,), dtype=wp.vec3i, device=device)
101
+ grid.get_voxels(out=self._cell_ijk)
102
+
103
+ self._node_grid = _build_node_grid(self._cell_ijk, grid, temporary_store)
104
+ node_count = self._node_grid.get_voxel_count()
105
+ self._node_ijk = wp.array(shape=(node_count,), dtype=wp.vec3i, device=device)
106
+ self._node_grid.get_voxels(out=self._node_ijk)
107
+
108
+ self._face_grid = None
109
+ self._face_ijk = None
110
+
111
+ self._edge_grid = None
112
+ self._edge_count = 0
113
+
114
+ @property
115
+ def cell_grid(self) -> wp.Volume:
116
+ return self._cell_grid
117
+
118
+ @property
119
+ def vertex_grid(self) -> wp.Volume:
120
+ return self._node_grid
121
+
122
+ @property
123
+ def face_grid(self) -> wp.Volume:
124
+ self._ensure_face_grid()
125
+ return self._face_grid
126
+
127
+ @property
128
+ def edge_grid(self) -> wp.Volume:
129
+ self._ensure_edge_grid()
130
+ return self._edge_grid
131
+
132
+ def cell_count(self):
133
+ return self._cell_ijk.shape[0]
134
+
135
+ def vertex_count(self):
136
+ return self._node_ijk.shape[0]
137
+
138
+ def side_count(self):
139
+ self._ensure_face_grid()
140
+ return self._face_ijk.shape[0]
141
+
142
+ def boundary_side_count(self):
143
+ self._ensure_face_grid()
144
+ return self._boundary_face_indices.shape[0]
145
+
146
+ def edge_count(self):
147
+ self._ensure_edge_grid()
148
+ return self._edge_count
149
+
150
+ def reference_cell(self) -> Cube:
151
+ return Cube()
152
+
153
+ def reference_side(self) -> Square:
154
+ return Square()
155
+
156
+ CellArg = NanogridCellArg
157
+
158
+ @cache.cached_arg_value
159
+ def cell_arg_value(self, device) -> CellArg:
160
+ args = self.CellArg()
161
+ args.cell_grid = self._cell_grid.id
162
+ args.cell_ijk = self._cell_ijk
163
+
164
+ transform = np.array(self._cell_grid_info.transform_matrix).reshape(3, 3)
165
+ args.inverse_transform = wp.mat33f(np.linalg.inv(transform))
166
+ args.cell_volume = abs(np.linalg.det(transform))
167
+
168
+ return args
169
+
170
+ @wp.func
171
+ def cell_position(args: CellArg, s: Sample):
172
+ uvw = wp.vec3(args.cell_ijk[s.element_index]) + s.element_coords
173
+ return wp.volume_index_to_world(args.cell_grid, uvw - wp.vec3(0.5))
174
+
175
+ @wp.func
176
+ def cell_deformation_gradient(args: CellArg, s: Sample):
177
+ return wp.inverse(args.inverse_transform)
178
+
179
+ @wp.func
180
+ def cell_inverse_deformation_gradient(args: CellArg, s: Sample):
181
+ return args.inverse_transform
182
+
183
+ @wp.func
184
+ def cell_lookup(args: CellArg, pos: wp.vec3):
185
+ uvw = wp.volume_world_to_index(args.cell_grid, pos) + wp.vec3(0.5)
186
+ ijk = wp.vec3i(int(wp.floor(uvw[0])), int(wp.floor(uvw[1])), int(wp.floor(uvw[2])))
187
+ cell_index = wp.volume_lookup_index(args.cell_grid, ijk[0], ijk[1], ijk[2])
188
+
189
+ coords = uvw - wp.vec3(ijk)
190
+ if cell_index == -1:
191
+ if wp.min(coords) == 0.0 or wp.max(coords) == 1.0:
192
+ il = wp.where(coords[0] > 0.5, 0, -1)
193
+ jl = wp.where(coords[1] > 0.5, 0, -1)
194
+ kl = wp.where(coords[2] > 0.5, 0, -1)
195
+
196
+ for n in range(8):
197
+ ni = n >> 2
198
+ nj = (n & 2) >> 1
199
+ nk = n & 1
200
+ nijk = ijk + wp.vec3i(ni + il, nj + jl, nk + kl)
201
+
202
+ coords = uvw - wp.vec3(nijk)
203
+ if wp.min(coords) >= 0.0 and wp.max(coords) <= 1.0:
204
+ cell_index = wp.volume_lookup_index(args.cell_grid, nijk[0], nijk[1], nijk[2])
205
+ if cell_index != -1:
206
+ return make_free_sample(cell_index, coords)
207
+
208
+ return make_free_sample(NULL_ELEMENT_INDEX, Coords(OUTSIDE))
209
+
210
+ return make_free_sample(cell_index, coords)
211
+
212
+ @wp.func
213
+ def _project_on_voxel_at_origin(coords: wp.vec3):
214
+ proj_coords = wp.min(wp.max(coords, wp.vec3(0.0)), wp.vec3(1.0))
215
+ return wp.length_sq(coords - proj_coords), proj_coords
216
+
217
+ @wp.func
218
+ def cell_lookup(args: CellArg, pos: wp.vec3, guess: Sample):
219
+ s_global = Nanogrid.cell_lookup(args, pos)
220
+
221
+ if s_global.element_index != NULL_ELEMENT_INDEX:
222
+ return s_global
223
+
224
+ closest_voxel = int(NULL_ELEMENT_INDEX)
225
+ closest_coords = Coords(OUTSIDE)
226
+ closest_dist = float(1.0e8)
227
+
228
+ # project to closest in stencil
229
+ uvw = wp.volume_world_to_index(args.cell_grid, pos) + wp.vec3(0.5)
230
+ cell_ijk = args.cell_ijk[guess.element_index]
231
+ for ni in range(-1, 2):
232
+ for nj in range(-1, 2):
233
+ for nk in range(-1, 2):
234
+ nijk = cell_ijk + wp.vec3i(ni, nj, nk)
235
+ cell_idx = wp.volume_lookup_index(args.cell_grid, nijk[0], nijk[1], nijk[2])
236
+ dist, coords = Nanogrid._project_on_voxel_at_origin(uvw - wp.vec3(nijk))
237
+ if cell_idx != -1 and dist <= closest_dist:
238
+ closest_dist = dist
239
+ closest_voxel = cell_idx
240
+ closest_coords = coords
241
+
242
+ return make_free_sample(closest_voxel, closest_coords)
243
+
244
+ @wp.func
245
+ def cell_measure(args: CellArg, s: Sample):
246
+ return args.cell_volume
247
+
248
+ @wp.func
249
+ def cell_normal(args: CellArg, s: Sample):
250
+ return wp.vec3(0.0)
251
+
252
+ SideArg = NanogridSideArg
253
+
254
+ @cache.cached_arg_value
255
+ def side_arg_value(self, device) -> SideArg:
256
+ self._ensure_face_grid()
257
+
258
+ args = self.SideArg()
259
+ args.cell_arg = self.cell_arg_value(device)
260
+ args.face_ijk = self._face_ijk.to(device)
261
+ args.face_flags = self._face_flags.to(device)
262
+ transform = np.array(self._cell_grid_info.transform_matrix).reshape(3, 3)
263
+ args.face_areas = wp.vec3(
264
+ tuple(np.linalg.norm(np.cross(transform[:, k - 2], transform[:, k - 1])) for k in range(3))
265
+ )
266
+
267
+ return args
268
+
269
+ @wp.struct
270
+ class SideIndexArg:
271
+ boundary_face_indices: wp.array(dtype=int)
272
+
273
+ @cache.cached_arg_value
274
+ def side_index_arg_value(self, device) -> SideIndexArg:
275
+ self._ensure_face_grid()
276
+
277
+ args = self.SideIndexArg()
278
+ args.boundary_face_indices = self._boundary_face_indices.to(device)
279
+ return args
280
+
281
+ @wp.func
282
+ def boundary_side_index(args: SideIndexArg, boundary_side_index: int):
283
+ return args.boundary_face_indices[boundary_side_index]
284
+
285
+ @wp.func
286
+ def _side_to_cell_coords(axis: int, inner: float, side_coords: Coords):
287
+ uvw = wp.vec3()
288
+ uvw[axis] = inner
289
+ uvw[(axis + 1) % 3] = side_coords[0]
290
+ uvw[(axis + 2) % 3] = side_coords[1]
291
+ return uvw
292
+
293
+ @wp.func
294
+ def _get_face_axis(flags: wp.uint8):
295
+ return wp.int32(flags & FACE_AXIS_MASK)
296
+
297
+ @wp.func
298
+ def _get_face_inner_offset(flags: wp.uint8):
299
+ return wp.int32(flags >> FACE_INNER_OFFSET_BIT) & 1
300
+
301
+ @wp.func
302
+ def _get_face_outer_offset(flags: wp.uint8):
303
+ return wp.int32(flags >> FACE_OUTER_OFFSET_BIT) & 1
304
+
305
+ @wp.func
306
+ def side_position(args: SideArg, s: Sample):
307
+ ijk = args.face_ijk[s.element_index]
308
+ axis = Nanogrid._get_face_axis(args.face_flags[s.element_index])
309
+
310
+ uvw = wp.vec3(ijk) + Nanogrid._side_to_cell_coords(axis, 0.0, s.element_coords)
311
+
312
+ cell_grid = args.cell_arg.cell_grid
313
+ return wp.volume_index_to_world(cell_grid, uvw - wp.vec3(0.5))
314
+
315
+ @wp.func
316
+ def _face_tangent_vecs(cell_grid: wp.uint64, axis: int, flip: int):
317
+ u_axis = basis_element(wp.vec3(), (axis + 1 + flip) % 3)
318
+ v_axis = basis_element(wp.vec3(), (axis + 2 - flip) % 3)
319
+
320
+ return wp.volume_index_to_world_dir(cell_grid, u_axis), wp.volume_index_to_world_dir(cell_grid, v_axis)
321
+
322
+ @wp.func
323
+ def side_deformation_gradient(args: SideArg, s: Sample):
324
+ flags = args.face_flags[s.element_index]
325
+ axis = Nanogrid._get_face_axis(flags)
326
+ flip = Nanogrid._get_face_inner_offset(flags)
327
+ v1, v2 = Nanogrid._face_tangent_vecs(args.cell_arg.cell_grid, axis, flip)
328
+ return wp.matrix_from_cols(v1, v2)
329
+
330
+ @wp.func
331
+ def side_inner_inverse_deformation_gradient(args: SideArg, s: Sample):
332
+ return Nanogrid.cell_inverse_deformation_gradient(args.cell_arg, s)
333
+
334
+ @wp.func
335
+ def side_outer_inverse_deformation_gradient(args: SideArg, s: Sample):
336
+ return Nanogrid.cell_inverse_deformation_gradient(args.cell_arg, s)
337
+
338
+ @wp.func
339
+ def side_measure(args: SideArg, s: Sample):
340
+ axis = Nanogrid._get_face_axis(args.face_flags[s.element_index])
341
+ return args.face_areas[axis]
342
+
343
+ @wp.func
344
+ def side_measure_ratio(args: SideArg, s: Sample):
345
+ axis = Nanogrid._get_face_axis(args.face_flags[s.element_index])
346
+ return args.face_areas[axis] / args.cell_arg.cell_volume
347
+
348
+ @wp.func
349
+ def side_normal(args: SideArg, s: Sample):
350
+ flags = args.face_flags[s.element_index]
351
+ axis = Nanogrid._get_face_axis(flags)
352
+ flip = Nanogrid._get_face_inner_offset(flags)
353
+
354
+ v1, v2 = Nanogrid._face_tangent_vecs(args.cell_arg.cell_grid, axis, flip)
355
+ return wp.cross(v1, v2) / args.face_areas[axis]
356
+
357
+ @wp.func
358
+ def side_inner_cell_index(args: SideArg, side_index: ElementIndex):
359
+ ijk = args.face_ijk[side_index]
360
+ flags = args.face_flags[side_index]
361
+ axis = Nanogrid._get_face_axis(flags)
362
+ offset = Nanogrid._get_face_inner_offset(flags)
363
+
364
+ ijk[axis] += offset - 1
365
+ cell_grid = args.cell_arg.cell_grid
366
+
367
+ return wp.volume_lookup_index(cell_grid, ijk[0], ijk[1], ijk[2])
368
+
369
+ @wp.func
370
+ def side_outer_cell_index(args: SideArg, side_index: ElementIndex):
371
+ ijk = args.face_ijk[side_index]
372
+ flags = args.face_flags[side_index]
373
+ axis = Nanogrid._get_face_axis(flags)
374
+ offset = Nanogrid._get_face_outer_offset(flags)
375
+
376
+ ijk[axis] -= offset
377
+ cell_grid = args.cell_arg.cell_grid
378
+
379
+ return wp.volume_lookup_index(cell_grid, ijk[0], ijk[1], ijk[2])
380
+
381
+ @wp.func
382
+ def side_inner_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
383
+ flags = args.face_flags[side_index]
384
+ axis = Nanogrid._get_face_axis(flags)
385
+ offset = float(Nanogrid._get_face_inner_offset(flags))
386
+ return Nanogrid._side_to_cell_coords(axis, 1.0 - offset, side_coords)
387
+
388
+ @wp.func
389
+ def side_outer_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
390
+ flags = args.face_flags[side_index]
391
+ axis = Nanogrid._get_face_axis(flags)
392
+ offset = float(Nanogrid._get_face_outer_offset(flags))
393
+ return Nanogrid._side_to_cell_coords(axis, offset, side_coords)
394
+
395
+ @wp.func
396
+ def side_from_cell_coords(
397
+ args: SideArg,
398
+ side_index: ElementIndex,
399
+ element_index: ElementIndex,
400
+ element_coords: Coords,
401
+ ):
402
+ flags = args.face_flags[side_index]
403
+ axis = Nanogrid._get_face_axis(flags)
404
+
405
+ cell_ijk = args.cell_arg.cell_ijk[element_index]
406
+ side_ijk = args.face_ijk[side_index]
407
+
408
+ on_side = float(side_ijk[axis] - cell_ijk[axis]) == element_coords[axis]
409
+
410
+ return wp.where(
411
+ on_side, Coords(element_coords[(axis + 1) % 3], element_coords[(axis + 2) % 3], 0.0), Coords(OUTSIDE)
412
+ )
413
+
414
+ @wp.func
415
+ def side_to_cell_arg(side_arg: SideArg):
416
+ return side_arg.cell_arg
417
+
418
+ def _build_face_grid(self, temporary_store: Optional[cache.TemporaryStore] = None):
419
+ device = self._cell_grid.device
420
+ self._face_grid = _build_face_grid(self._cell_ijk, self._cell_grid, temporary_store)
421
+ face_count = self._face_grid.get_voxel_count()
422
+ self._face_ijk = wp.array(shape=(face_count,), dtype=wp.vec3i, device=device)
423
+ self._face_grid.get_voxels(out=self._face_ijk)
424
+
425
+ self._face_flags = wp.array(shape=(face_count,), dtype=wp.uint8, device=device)
426
+ boundary_face_mask = cache.borrow_temporary(temporary_store, shape=(face_count,), dtype=wp.int32, device=device)
427
+
428
+ wp.launch(
429
+ _build_face_flags,
430
+ dim=face_count,
431
+ device=device,
432
+ inputs=[self._cell_grid.id, self._face_ijk, self._face_flags, boundary_face_mask.array],
433
+ )
434
+ boundary_face_indices, _ = utils.masked_indices(boundary_face_mask.array)
435
+ self._boundary_face_indices = boundary_face_indices.detach()
436
+
437
+ def _build_edge_grid(self, temporary_store: Optional[cache.TemporaryStore] = None):
438
+ self._edge_grid = _build_edge_grid(self._cell_ijk, self._cell_grid, temporary_store)
439
+ self._edge_count = self._edge_grid.get_voxel_count()
440
+
441
+ def _ensure_face_grid(self):
442
+ if self._face_ijk is None:
443
+ self._build_face_grid()
444
+
445
+ def _ensure_edge_grid(self):
446
+ if self._edge_grid is None:
447
+ self._build_edge_grid()
448
+
449
+
450
+ @wp.kernel
451
+ def _cell_node_indices(
452
+ cell_ijk: wp.array(dtype=wp.vec3i),
453
+ node_ijk: wp.array2d(dtype=wp.vec3i),
454
+ ):
455
+ cell, n = wp.tid()
456
+ node_ijk[cell, n] = cell_ijk[cell] + wp.vec3i((n & 4) >> 2, (n & 2) >> 1, n & 1)
457
+
458
+
459
+ @wp.kernel
460
+ def _cell_face_indices(
461
+ cell_ijk: wp.array(dtype=wp.vec3i),
462
+ node_ijk: wp.array2d(dtype=wp.vec3i),
463
+ ):
464
+ cell = wp.tid()
465
+ ijk = cell_ijk[cell]
466
+ node_ijk[cell, 0] = _add_axis_flag(ijk, 0)
467
+ node_ijk[cell, 1] = _add_axis_flag(ijk, 1)
468
+ node_ijk[cell, 2] = _add_axis_flag(ijk, 2)
469
+
470
+ node_ijk[cell, 3] = _add_axis_flag(ijk + wp.vec3i(1, 0, 0), 0)
471
+ node_ijk[cell, 4] = _add_axis_flag(ijk + wp.vec3i(0, 1, 0), 1)
472
+ node_ijk[cell, 5] = _add_axis_flag(ijk + wp.vec3i(0, 0, 1), 2)
473
+
474
+
475
+ @wp.kernel
476
+ def _cell_edge_indices(
477
+ cell_ijk: wp.array(dtype=wp.vec3i),
478
+ edge_ijk: wp.array2d(dtype=wp.vec3i),
479
+ ):
480
+ cell = wp.tid()
481
+ ijk = cell_ijk[cell]
482
+ edge_ijk[cell, 0] = _add_axis_flag(ijk, 0)
483
+ edge_ijk[cell, 1] = _add_axis_flag(ijk, 1)
484
+ edge_ijk[cell, 2] = _add_axis_flag(ijk, 2)
485
+
486
+ edge_ijk[cell, 3] = _add_axis_flag(ijk + wp.vec3i(0, 1, 0), 0)
487
+ edge_ijk[cell, 4] = _add_axis_flag(ijk + wp.vec3i(0, 0, 1), 1)
488
+ edge_ijk[cell, 5] = _add_axis_flag(ijk + wp.vec3i(1, 0, 0), 2)
489
+
490
+ edge_ijk[cell, 6] = _add_axis_flag(ijk + wp.vec3i(0, 1, 1), 0)
491
+ edge_ijk[cell, 7] = _add_axis_flag(ijk + wp.vec3i(1, 0, 1), 1)
492
+ edge_ijk[cell, 8] = _add_axis_flag(ijk + wp.vec3i(1, 1, 0), 2)
493
+
494
+ edge_ijk[cell, 9] = _add_axis_flag(ijk + wp.vec3i(0, 0, 1), 0)
495
+ edge_ijk[cell, 10] = _add_axis_flag(ijk + wp.vec3i(1, 0, 0), 1)
496
+ edge_ijk[cell, 11] = _add_axis_flag(ijk + wp.vec3i(0, 1, 0), 2)
497
+
498
+
499
+ def _build_node_grid(cell_ijk, grid: wp.Volume, temporary_store: cache.TemporaryStore):
500
+ cell_count = cell_ijk.shape[0]
501
+
502
+ cell_nodes = cache.borrow_temporary(temporary_store, shape=(cell_count, 8), dtype=wp.vec3i, device=cell_ijk.device)
503
+ wp.launch(
504
+ _cell_node_indices, dim=cell_nodes.array.shape, inputs=[cell_ijk, cell_nodes.array], device=cell_ijk.device
505
+ )
506
+ node_grid = wp.Volume.allocate_by_voxels(
507
+ cell_nodes.array.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
508
+ )
509
+
510
+ return node_grid
511
+
512
+
513
+ def _build_face_grid(cell_ijk, grid: wp.Volume, temporary_store: cache.TemporaryStore):
514
+ cell_count = cell_ijk.shape[0]
515
+
516
+ cell_faces = cache.borrow_temporary(temporary_store, shape=(cell_count, 6), dtype=wp.vec3i, device=cell_ijk.device)
517
+ wp.launch(_cell_face_indices, dim=cell_count, inputs=[cell_ijk, cell_faces.array], device=cell_ijk.device)
518
+ face_grid = wp.Volume.allocate_by_voxels(
519
+ cell_faces.array.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
520
+ )
521
+
522
+ return face_grid
523
+
524
+
525
+ def _build_edge_grid(cell_ijk, grid: wp.Volume, temporary_store: cache.TemporaryStore):
526
+ cell_count = cell_ijk.shape[0]
527
+
528
+ cell_edges = cache.borrow_temporary(temporary_store, shape=(cell_count, 12), dtype=wp.vec3i, device=cell_ijk.device)
529
+ wp.launch(_cell_edge_indices, dim=cell_count, inputs=[cell_ijk, cell_edges.array], device=cell_ijk.device)
530
+ edge_grid = wp.Volume.allocate_by_voxels(
531
+ cell_edges.array.flatten(), voxel_size=grid.get_voxel_size(), device=cell_ijk.device
532
+ )
533
+
534
+ return edge_grid
535
+
536
+
537
+ @wp.func
538
+ def _make_face_flags(axis: int, plus_cell_index: int, minus_cell_index: int):
539
+ plus_boundary = wp.uint8(wp.where(plus_cell_index == -1, 1, 0)) << FACE_OUTER_OFFSET_BIT
540
+ minus_boundary = wp.uint8(wp.where(minus_cell_index == -1, 1, 0)) << FACE_INNER_OFFSET_BIT
541
+
542
+ return wp.uint8(axis) | plus_boundary | minus_boundary
543
+
544
+
545
+ @wp.func
546
+ def _get_boundary_mask(flags: wp.uint8):
547
+ return int((flags >> FACE_OUTER_OFFSET_BIT) | (flags >> FACE_INNER_OFFSET_BIT)) & 1
548
+
549
+
550
+ @wp.kernel
551
+ def _build_face_flags(
552
+ cell_grid: wp.uint64,
553
+ face_ijk: wp.array(dtype=wp.vec3i),
554
+ face_flags: wp.array(dtype=wp.uint8),
555
+ boundary_face_mask: wp.array(dtype=int),
556
+ ):
557
+ face = wp.tid()
558
+
559
+ axis, ijk = _extract_axis_flag(face_ijk[face])
560
+
561
+ ijk_minus = ijk
562
+ ijk_minus[axis] -= 1
563
+
564
+ plus_cell_index = wp.volume_lookup_index(cell_grid, ijk[0], ijk[1], ijk[2])
565
+ minus_cell_index = wp.volume_lookup_index(cell_grid, ijk_minus[0], ijk_minus[1], ijk_minus[2])
566
+
567
+ face_ijk[face] = ijk
568
+
569
+ flags = _make_face_flags(axis, plus_cell_index, minus_cell_index)
570
+ face_flags[face] = flags
571
+ boundary_face_mask[face] = _get_boundary_mask(flags)