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,672 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2023 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 numpy as np
17
+
18
+ import warp as wp
19
+ from warp.fem import cache
20
+ from warp.fem.types import Coords
21
+
22
+ from .shape_function import ShapeFunction
23
+
24
+
25
+ def _triangle_node_index(tx: int, ty: int, degree: int):
26
+ VERTEX_NODE_COUNT = 3
27
+ SIDE_INTERIOR_NODE_COUNT = degree - 1
28
+
29
+ # Index in similar order to e.g. VTK
30
+ # First vertices, then edge (counterclockwise) then interior points (recursively)
31
+
32
+ if tx == 0:
33
+ if ty == 0:
34
+ return 0
35
+ elif ty == degree:
36
+ return 2
37
+ else:
38
+ edge_index = 2
39
+ return VERTEX_NODE_COUNT + SIDE_INTERIOR_NODE_COUNT * edge_index + (SIDE_INTERIOR_NODE_COUNT - ty)
40
+ elif ty == 0:
41
+ if tx == degree:
42
+ return 1
43
+ else:
44
+ edge_index = 0
45
+ return VERTEX_NODE_COUNT + SIDE_INTERIOR_NODE_COUNT * edge_index + tx - 1
46
+ elif tx + ty == degree:
47
+ edge_index = 1
48
+ return VERTEX_NODE_COUNT + SIDE_INTERIOR_NODE_COUNT * edge_index + ty - 1
49
+
50
+ vertex_edge_node_count = 3 * degree
51
+ return vertex_edge_node_count + _triangle_node_index(tx - 1, ty - 1, degree - 3)
52
+
53
+
54
+ class TriangleShapeFunction(ShapeFunction):
55
+ VERTEX = wp.constant(0)
56
+ EDGE = wp.constant(1)
57
+ INTERIOR = wp.constant(2)
58
+
59
+ VERTEX_NODE_COUNT: int
60
+ """Number of shape function nodes per vertex"""
61
+
62
+ EDGE_NODE_COUNT: int
63
+ """Number of shape function nodes per triangle edge (excluding vertex nodes)"""
64
+
65
+ INTERIOR_NODE_COUNT: int
66
+ """Number of shape function nodes per triangle (excluding edge and vertex nodes)"""
67
+
68
+ @staticmethod
69
+ def node_type_and_index(node_index_in_elt: int):
70
+ pass
71
+
72
+ @wp.func
73
+ def _vertex_coords(vidx: int):
74
+ return wp.vec2(
75
+ float(vidx == 1),
76
+ float(vidx == 2),
77
+ )
78
+
79
+
80
+ class TrianglePolynomialShapeFunctions(TriangleShapeFunction):
81
+ def __init__(self, degree: int):
82
+ self.ORDER = wp.constant(degree)
83
+
84
+ self.NODES_PER_ELEMENT = wp.constant((degree + 1) * (degree + 2) // 2)
85
+ self.NODES_PER_SIDE = wp.constant(degree + 1)
86
+
87
+ self.VERTEX_NODE_COUNT = wp.constant(1)
88
+ self.EDGE_NODE_COUNT = wp.constant(degree - 1)
89
+ self.INTERIOR_NODE_COUNT = wp.constant(max(0, degree - 2) * max(0, degree - 1) // 2)
90
+
91
+ triangle_coords = np.empty((self.NODES_PER_ELEMENT, 2), dtype=int)
92
+
93
+ for tx in range(degree + 1):
94
+ for ty in range(degree + 1 - tx):
95
+ index = _triangle_node_index(tx, ty, degree)
96
+ triangle_coords[index] = [tx, ty]
97
+
98
+ CoordTypeVec = wp.mat(dtype=int, shape=(self.NODES_PER_ELEMENT, 2))
99
+ self.NODE_TRIANGLE_COORDS = wp.constant(CoordTypeVec(triangle_coords))
100
+
101
+ self.node_type_and_type_index = self._get_node_type_and_type_index()
102
+ self._node_triangle_coordinates = self._get_node_triangle_coordinates()
103
+
104
+ @property
105
+ def name(self) -> str:
106
+ return f"Tri_P{self.ORDER}"
107
+
108
+ def _get_node_triangle_coordinates(self):
109
+ NODE_TRIANGLE_COORDS = self.NODE_TRIANGLE_COORDS
110
+
111
+ def node_triangle_coordinates(
112
+ node_index_in_elt: int,
113
+ ):
114
+ return NODE_TRIANGLE_COORDS[node_index_in_elt]
115
+
116
+ return cache.get_func(node_triangle_coordinates, self.name)
117
+
118
+ def _get_node_type_and_type_index(self):
119
+ ORDER = self.ORDER
120
+
121
+ def node_type_and_index(
122
+ node_index_in_elt: int,
123
+ ):
124
+ if node_index_in_elt < 3:
125
+ return TrianglePolynomialShapeFunctions.VERTEX, node_index_in_elt
126
+
127
+ if node_index_in_elt < 3 * ORDER:
128
+ return TrianglePolynomialShapeFunctions.EDGE, (node_index_in_elt - 3)
129
+
130
+ return TrianglePolynomialShapeFunctions.INTERIOR, (node_index_in_elt - 3 * ORDER)
131
+
132
+ return cache.get_func(node_type_and_index, self.name)
133
+
134
+ def make_node_coords_in_element(self):
135
+ ORDER = self.ORDER
136
+
137
+ def node_coords_in_element(
138
+ node_index_in_elt: int,
139
+ ):
140
+ tri_coords = self._node_triangle_coordinates(node_index_in_elt)
141
+ cx = float(tri_coords[0]) / float(ORDER)
142
+ cy = float(tri_coords[1]) / float(ORDER)
143
+ return Coords(1.0 - cx - cy, cx, cy)
144
+
145
+ return cache.get_func(node_coords_in_element, self.name)
146
+
147
+ def make_node_quadrature_weight(self):
148
+ if self.ORDER == 3:
149
+ # P3 intrinsic quadrature
150
+ vertex_weight = 1.0 / 30
151
+ edge_weight = 0.075
152
+ interior_weight = 0.45
153
+ elif self.ORDER == 2:
154
+ # Order 1, but optimized quadrature weights for monomials of order <= 4
155
+ vertex_weight = 0.022335964126
156
+ edge_weight = 0.310997369207
157
+ interior_weight = 0.0
158
+ else:
159
+ vertex_weight = 1.0 / self.NODES_PER_ELEMENT
160
+ edge_weight = 1.0 / self.NODES_PER_ELEMENT
161
+ interior_weight = 1.0 / self.NODES_PER_ELEMENT
162
+
163
+ VERTEX_WEIGHT = wp.constant(vertex_weight)
164
+ EDGE_WEIGHT = wp.constant(edge_weight)
165
+ INTERIOR_WEIGHT = wp.constant(interior_weight)
166
+
167
+ @cache.dynamic_func(suffix=self.name)
168
+ def node_quadrature_weight(node_index_in_element: int):
169
+ node_type, type_index = self.node_type_and_type_index(node_index_in_element)
170
+
171
+ if node_type == TrianglePolynomialShapeFunctions.VERTEX:
172
+ return VERTEX_WEIGHT
173
+ elif node_type == TrianglePolynomialShapeFunctions.EDGE:
174
+ return EDGE_WEIGHT
175
+
176
+ return INTERIOR_WEIGHT
177
+
178
+ return node_quadrature_weight
179
+
180
+ def make_trace_node_quadrature_weight(self):
181
+ # Closed Newton-Cotes
182
+ if self.ORDER == 3:
183
+ vertex_weight = 1.0 / 8.0
184
+ edge_weight = 3.0 / 8.0
185
+ elif self.ORDER == 2:
186
+ vertex_weight = 1.0 / 6.0
187
+ edge_weight = 2.0 / 3.0
188
+ else:
189
+ vertex_weight = 1.0 / self.NODES_PER_SIDE
190
+ edge_weight = 1.0 / self.NODES_PER_SIDE
191
+
192
+ VERTEX_WEIGHT = wp.constant(vertex_weight)
193
+ EDGE_WEIGHT = wp.constant(edge_weight)
194
+
195
+ @cache.dynamic_func(suffix=self.name)
196
+ def trace_node_quadrature_weight(node_index_in_element: int):
197
+ node_type, type_index = self.node_type_and_type_index(node_index_in_element)
198
+
199
+ return wp.where(node_type == TrianglePolynomialShapeFunctions.VERTEX, VERTEX_WEIGHT, EDGE_WEIGHT)
200
+
201
+ return trace_node_quadrature_weight
202
+
203
+ def make_element_inner_weight(self):
204
+ ORDER = self.ORDER
205
+
206
+ def element_inner_weight_linear(
207
+ coords: Coords,
208
+ node_index_in_elt: int,
209
+ ):
210
+ return coords[node_index_in_elt]
211
+
212
+ def element_inner_weight_quadratic(
213
+ coords: Coords,
214
+ node_index_in_elt: int,
215
+ ):
216
+ node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
217
+
218
+ if node_type == TrianglePolynomialShapeFunctions.VERTEX:
219
+ # Vertex
220
+ return coords[type_index] * (2.0 * coords[type_index] - 1.0)
221
+
222
+ # Edge
223
+ c1 = type_index
224
+ c2 = (type_index + 1) % 3
225
+ return 4.0 * coords[c1] * coords[c2]
226
+
227
+ def element_inner_weight_cubic(
228
+ coords: Coords,
229
+ node_index_in_elt: int,
230
+ ):
231
+ node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
232
+
233
+ if node_type == TrianglePolynomialShapeFunctions.VERTEX:
234
+ # Vertex
235
+ return 0.5 * coords[type_index] * (3.0 * coords[type_index] - 1.0) * (3.0 * coords[type_index] - 2.0)
236
+
237
+ elif node_type == TrianglePolynomialShapeFunctions.EDGE:
238
+ # Edge
239
+ edge = type_index // 2
240
+ k = type_index - 2 * edge
241
+ c1 = (edge + k) % 3
242
+ c2 = (edge + 1 - k) % 3
243
+
244
+ return 4.5 * coords[c1] * coords[c2] * (3.0 * coords[c1] - 1.0)
245
+
246
+ # Interior
247
+ return 27.0 * coords[0] * coords[1] * coords[2]
248
+
249
+ if ORDER == 1:
250
+ return cache.get_func(element_inner_weight_linear, self.name)
251
+ elif ORDER == 2:
252
+ return cache.get_func(element_inner_weight_quadratic, self.name)
253
+ elif ORDER == 3:
254
+ return cache.get_func(element_inner_weight_cubic, self.name)
255
+
256
+ return None
257
+
258
+ def make_element_inner_weight_gradient(self):
259
+ ORDER = self.ORDER
260
+
261
+ def element_inner_weight_gradient_linear(
262
+ coords: Coords,
263
+ node_index_in_elt: int,
264
+ ):
265
+ dw_dc = wp.vec3(0.0)
266
+ dw_dc[node_index_in_elt] = 1.0
267
+
268
+ dw_du = wp.vec2(dw_dc[1] - dw_dc[0], dw_dc[2] - dw_dc[0])
269
+ return dw_du
270
+
271
+ def element_inner_weight_gradient_quadratic(
272
+ coords: Coords,
273
+ node_index_in_elt: int,
274
+ ):
275
+ node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
276
+
277
+ dw_dc = wp.vec3(0.0)
278
+
279
+ if node_type == TrianglePolynomialShapeFunctions.VERTEX:
280
+ # Vertex
281
+ dw_dc[type_index] = 4.0 * coords[type_index] - 1.0
282
+
283
+ else:
284
+ # Edge
285
+ c1 = type_index
286
+ c2 = (type_index + 1) % 3
287
+ dw_dc[c1] = 4.0 * coords[c2]
288
+ dw_dc[c2] = 4.0 * coords[c1]
289
+
290
+ dw_du = wp.vec2(dw_dc[1] - dw_dc[0], dw_dc[2] - dw_dc[0])
291
+ return dw_du
292
+
293
+ def element_inner_weight_gradient_cubic(
294
+ coords: Coords,
295
+ node_index_in_elt: int,
296
+ ):
297
+ node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
298
+
299
+ dw_dc = wp.vec3(0.0)
300
+
301
+ if node_type == TrianglePolynomialShapeFunctions.VERTEX:
302
+ # Vertex
303
+ dw_dc[type_index] = (
304
+ 0.5 * 27.0 * coords[type_index] * coords[type_index] - 9.0 * coords[type_index] + 1.0
305
+ )
306
+
307
+ elif node_type == TrianglePolynomialShapeFunctions.EDGE:
308
+ # Edge
309
+ edge = type_index // 2
310
+ k = type_index - 2 * edge
311
+ c1 = (edge + k) % 3
312
+ c2 = (edge + 1 - k) % 3
313
+
314
+ dw_dc[c1] = 4.5 * coords[c2] * (6.0 * coords[c1] - 1.0)
315
+ dw_dc[c2] = 4.5 * coords[c1] * (3.0 * coords[c1] - 1.0)
316
+
317
+ else:
318
+ # Interior
319
+ dw_dc = wp.vec3(
320
+ 27.0 * coords[1] * coords[2], 27.0 * coords[2] * coords[0], 27.0 * coords[0] * coords[1]
321
+ )
322
+
323
+ dw_du = wp.vec2(dw_dc[1] - dw_dc[0], dw_dc[2] - dw_dc[0])
324
+ return dw_du
325
+
326
+ if ORDER == 1:
327
+ return cache.get_func(element_inner_weight_gradient_linear, self.name)
328
+ elif ORDER == 2:
329
+ return cache.get_func(element_inner_weight_gradient_quadratic, self.name)
330
+ elif ORDER == 3:
331
+ return cache.get_func(element_inner_weight_gradient_cubic, self.name)
332
+
333
+ return None
334
+
335
+ def element_node_triangulation(self):
336
+ if self.ORDER == 1:
337
+ element_triangles = [[0, 1, 2]]
338
+ if self.ORDER == 2:
339
+ element_triangles = [[0, 3, 5], [3, 1, 4], [2, 5, 4], [3, 4, 5]]
340
+ elif self.ORDER == 3:
341
+ element_triangles = [
342
+ [0, 3, 8],
343
+ [3, 4, 9],
344
+ [4, 1, 5],
345
+ [8, 3, 9],
346
+ [4, 5, 9],
347
+ [8, 9, 7],
348
+ [9, 5, 6],
349
+ [6, 7, 9],
350
+ [7, 6, 2],
351
+ ]
352
+
353
+ return np.array(element_triangles)
354
+
355
+ def element_vtk_cells(self):
356
+ cells = np.arange(self.NODES_PER_ELEMENT)
357
+ if self.ORDER == 1:
358
+ cell_type = 5 # VTK_TRIANGLE
359
+ else:
360
+ cell_type = 69 # VTK_LAGRANGE_TRIANGLE
361
+ return cells[np.newaxis, :], np.array([cell_type], dtype=np.int8)
362
+
363
+
364
+ class TriangleNonConformingPolynomialShapeFunctions(ShapeFunction):
365
+ def __init__(self, degree: int):
366
+ self._tri_shape = TrianglePolynomialShapeFunctions(degree=degree)
367
+ self.ORDER = self._tri_shape.ORDER
368
+ self.NODES_PER_ELEMENT = self._tri_shape.NODES_PER_ELEMENT
369
+
370
+ self.element_node_triangulation = self._tri_shape.element_node_triangulation
371
+ self.element_vtk_cells = self._tri_shape.element_vtk_cells
372
+
373
+ # Coordinates (a, b, b) of embedded triangle
374
+ if self.ORDER == 1:
375
+ # Order 2
376
+ a = 2.0 / 3.0
377
+ elif self.ORDER == 2:
378
+ # Order 2, optimized for small intrinsic quadrature error up to degree 4
379
+ a = 0.7790771484375001
380
+ elif self.ORDER == 3:
381
+ # Order 3, optimized for small intrinsic quadrature error up to degree 6
382
+ a = 0.8429443359375002
383
+ else:
384
+ a = 1.0
385
+
386
+ b = 0.5 * (1.0 - a)
387
+ self._small_to_big = np.full((3, 3), b) + (a - b) * np.eye(3)
388
+ self._tri_scale = a - b
389
+
390
+ @property
391
+ def name(self) -> str:
392
+ return f"Tri_dP{self.ORDER}"
393
+
394
+ def make_node_quadrature_weight(self):
395
+ # Intrinsic quadrature -- precomputed integral of node shape functions
396
+ # over element. Order equal to self.ORDER
397
+
398
+ if self.ORDER == 2:
399
+ vertex_weight = 0.13743348
400
+ edge_weight = 0.19589985
401
+ interior_weight = 0.0
402
+ elif self.ORDER == 3:
403
+ vertex_weight = 0.07462578
404
+ edge_weight = 0.1019807
405
+ interior_weight = 0.16423881
406
+ else:
407
+ vertex_weight = 1.0 / self.NODES_PER_ELEMENT
408
+ edge_weight = 1.0 / self.NODES_PER_ELEMENT
409
+ interior_weight = 1.0 / self.NODES_PER_ELEMENT
410
+
411
+ VERTEX_WEIGHT = wp.constant(vertex_weight)
412
+ EDGE_WEIGHT = wp.constant(edge_weight)
413
+ INTERIOR_WEIGHT = wp.constant(interior_weight)
414
+
415
+ @cache.dynamic_func(suffix=self.name)
416
+ def node_quadrature_weight(node_index_in_element: int):
417
+ node_type, type_index = self._tri_shape.node_type_and_type_index(node_index_in_element)
418
+
419
+ if node_type == TrianglePolynomialShapeFunctions.VERTEX:
420
+ return VERTEX_WEIGHT
421
+ elif node_type == TrianglePolynomialShapeFunctions.EDGE:
422
+ return EDGE_WEIGHT
423
+
424
+ return INTERIOR_WEIGHT
425
+
426
+ return node_quadrature_weight
427
+
428
+ def make_trace_node_quadrature_weight(self):
429
+ # Non-conforming, zero measure on sides
430
+
431
+ @wp.func
432
+ def zero(node_index_in_elt: int):
433
+ return 0.0
434
+
435
+ return zero
436
+
437
+ def make_node_coords_in_element(self):
438
+ node_coords_in_tet = self._tri_shape.make_node_coords_in_element()
439
+
440
+ SMALL_TO_BIG = wp.constant(wp.mat33(self._small_to_big))
441
+
442
+ @cache.dynamic_func(suffix=self.name)
443
+ def node_coords_in_element(
444
+ node_index_in_elt: int,
445
+ ):
446
+ tri_coords = node_coords_in_tet(node_index_in_elt)
447
+ return SMALL_TO_BIG * tri_coords
448
+
449
+ return node_coords_in_element
450
+
451
+ def make_element_inner_weight(self):
452
+ tri_inner_weight = self._tri_shape.make_element_inner_weight()
453
+
454
+ BIG_TO_SMALL = wp.constant(wp.mat33(np.linalg.inv(self._small_to_big)))
455
+
456
+ @cache.dynamic_func(suffix=self.name)
457
+ def element_inner_weight(
458
+ coords: Coords,
459
+ node_index_in_elt: int,
460
+ ):
461
+ tri_coords = BIG_TO_SMALL * coords
462
+ return tri_inner_weight(tri_coords, node_index_in_elt)
463
+
464
+ return element_inner_weight
465
+
466
+ def make_element_inner_weight_gradient(self):
467
+ tri_inner_weight_gradient = self._tri_shape.make_element_inner_weight_gradient()
468
+
469
+ BIG_TO_SMALL = wp.constant(wp.mat33(np.linalg.inv(self._small_to_big)))
470
+ INV_TRI_SCALE = wp.constant(1.0 / self._tri_scale)
471
+
472
+ @cache.dynamic_func(suffix=self.name)
473
+ def element_inner_weight_gradient(
474
+ coords: Coords,
475
+ node_index_in_elt: int,
476
+ ):
477
+ tri_coords = BIG_TO_SMALL * coords
478
+ grad = tri_inner_weight_gradient(tri_coords, node_index_in_elt)
479
+ return INV_TRI_SCALE * grad
480
+
481
+ return element_inner_weight_gradient
482
+
483
+
484
+ class TriangleNedelecFirstKindShapeFunctions(TriangleShapeFunction):
485
+ value = ShapeFunction.Value.CovariantVector
486
+
487
+ def __init__(self, degree: int):
488
+ if degree != 1:
489
+ raise NotImplementedError("Only linear Nédélec implemented right now")
490
+
491
+ self.ORDER = wp.constant(degree)
492
+
493
+ self.NODES_PER_ELEMENT = wp.constant(3)
494
+ self.NODES_PER_SIDE = wp.constant(1)
495
+
496
+ self.VERTEX_NODE_COUNT = wp.constant(0)
497
+ self.EDGE_NODE_COUNT = wp.constant(1)
498
+ self.INTERIOR_NODE_COUNT = wp.constant(0)
499
+
500
+ self.node_type_and_type_index = self._get_node_type_and_type_index()
501
+
502
+ @property
503
+ def name(self) -> str:
504
+ return f"TriN1_{self.ORDER}"
505
+
506
+ def _get_node_type_and_type_index(self):
507
+ @cache.dynamic_func(suffix=self.name)
508
+ def node_type_and_index(
509
+ node_index_in_elt: int,
510
+ ):
511
+ return TriangleShapeFunction.EDGE, node_index_in_elt
512
+
513
+ return node_type_and_index
514
+
515
+ def make_node_coords_in_element(self):
516
+ @cache.dynamic_func(suffix=self.name)
517
+ def node_coords_in_element(
518
+ node_index_in_elt: int,
519
+ ):
520
+ coords = Coords(0.5)
521
+ coords[(node_index_in_elt + 2) % 3] = 0.0
522
+ return coords
523
+
524
+ return node_coords_in_element
525
+
526
+ def make_node_quadrature_weight(self):
527
+ NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
528
+
529
+ @cache.dynamic_func(suffix=self.name)
530
+ def node_quadrature_weight(node_index_in_element: int):
531
+ return 1.0 / float(NODES_PER_ELEMENT)
532
+
533
+ return node_quadrature_weight
534
+
535
+ def make_trace_node_quadrature_weight(self):
536
+ NODES_PER_SIDE = self.NODES_PER_SIDE
537
+
538
+ @cache.dynamic_func(suffix=self.name)
539
+ def trace_node_quadrature_weight(node_index_in_element: int):
540
+ return 1.0 / float(NODES_PER_SIDE)
541
+
542
+ return trace_node_quadrature_weight
543
+
544
+ @wp.func
545
+ def _vertex_coords(vidx: int):
546
+ return wp.vec2(
547
+ float(vidx == 1),
548
+ float(vidx == 2),
549
+ )
550
+
551
+ def make_element_inner_weight(self):
552
+ ORDER = self.ORDER
553
+
554
+ def element_inner_weight_linear(
555
+ coords: Coords,
556
+ node_index_in_elt: int,
557
+ ):
558
+ x = wp.vec2(coords[1], coords[2])
559
+ p = self._vertex_coords((node_index_in_elt + 2) % 3)
560
+
561
+ d = x - p
562
+ return wp.vec2(-d[1], d[0])
563
+
564
+ if ORDER == 1:
565
+ return cache.get_func(element_inner_weight_linear, self.name)
566
+
567
+ return None
568
+
569
+ def make_element_inner_weight_gradient(self):
570
+ ROT = wp.constant(wp.mat22f(0.0, -1.0, 1.0, 0.0))
571
+
572
+ def element_inner_weight_gradient_linear(
573
+ coords: Coords,
574
+ node_index_in_elt: int,
575
+ ):
576
+ return ROT
577
+
578
+ if self.ORDER == 1:
579
+ return cache.get_func(element_inner_weight_gradient_linear, self.name)
580
+
581
+ return None
582
+
583
+
584
+ class TriangleRaviartThomasShapeFunctions(TriangleShapeFunction):
585
+ value = ShapeFunction.Value.ContravariantVector
586
+
587
+ def __init__(self, degree: int):
588
+ if degree != 1:
589
+ raise NotImplementedError("Only linear Raviart-Thomas implemented right now")
590
+
591
+ self.ORDER = wp.constant(degree)
592
+
593
+ self.NODES_PER_ELEMENT = wp.constant(3)
594
+ self.NODES_PER_SIDE = wp.constant(1)
595
+
596
+ self.VERTEX_NODE_COUNT = wp.constant(0)
597
+ self.EDGE_NODE_COUNT = wp.constant(1)
598
+ self.INTERIOR_NODE_COUNT = wp.constant(0)
599
+
600
+ self.node_type_and_type_index = self._get_node_type_and_type_index()
601
+
602
+ @property
603
+ def name(self) -> str:
604
+ return f"TriRT_{self.ORDER}"
605
+
606
+ def _get_node_type_and_type_index(self):
607
+ @cache.dynamic_func(suffix=self.name)
608
+ def node_type_and_index(
609
+ node_index_in_elt: int,
610
+ ):
611
+ return TriangleShapeFunction.EDGE, node_index_in_elt
612
+
613
+ return node_type_and_index
614
+
615
+ def make_node_coords_in_element(self):
616
+ @cache.dynamic_func(suffix=self.name)
617
+ def node_coords_in_element(
618
+ node_index_in_elt: int,
619
+ ):
620
+ coords = Coords(0.5)
621
+ coords[(node_index_in_elt + 2) % 3] = 0.0
622
+ return coords
623
+
624
+ return node_coords_in_element
625
+
626
+ def make_node_quadrature_weight(self):
627
+ NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
628
+
629
+ @cache.dynamic_func(suffix=self.name)
630
+ def node_quadrature_weight(node_index_in_element: int):
631
+ return 1.0 / float(NODES_PER_ELEMENT)
632
+
633
+ return node_quadrature_weight
634
+
635
+ def make_trace_node_quadrature_weight(self):
636
+ NODES_PER_SIDE = self.NODES_PER_SIDE
637
+
638
+ @cache.dynamic_func(suffix=self.name)
639
+ def trace_node_quadrature_weight(node_index_in_element: int):
640
+ return 1.0 / float(NODES_PER_SIDE)
641
+
642
+ return trace_node_quadrature_weight
643
+
644
+ def make_element_inner_weight(self):
645
+ ORDER = self.ORDER
646
+
647
+ def element_inner_weight_linear(
648
+ coords: Coords,
649
+ node_index_in_elt: int,
650
+ ):
651
+ x = wp.vec2(coords[1], coords[2])
652
+ p = self._vertex_coords((node_index_in_elt + 2) % 3)
653
+
654
+ d = x - p
655
+ return d
656
+
657
+ if ORDER == 1:
658
+ return cache.get_func(element_inner_weight_linear, self.name)
659
+
660
+ return None
661
+
662
+ def make_element_inner_weight_gradient(self):
663
+ def element_inner_weight_gradient_linear(
664
+ coords: Coords,
665
+ node_index_in_elt: int,
666
+ ):
667
+ return wp.identity(n=2, dtype=float)
668
+
669
+ if self.ORDER == 1:
670
+ return cache.get_func(element_inner_weight_gradient_linear, self.name)
671
+
672
+ return None