warp-lang 1.10.0__py3-none-macosx_11_0_arm64.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 (468) hide show
  1. warp/__init__.py +334 -0
  2. warp/__init__.pyi +5856 -0
  3. warp/_src/__init__.py +14 -0
  4. warp/_src/autograd.py +1077 -0
  5. warp/_src/build.py +620 -0
  6. warp/_src/build_dll.py +642 -0
  7. warp/_src/builtins.py +10555 -0
  8. warp/_src/codegen.py +4361 -0
  9. warp/_src/config.py +178 -0
  10. warp/_src/constants.py +59 -0
  11. warp/_src/context.py +8352 -0
  12. warp/_src/dlpack.py +464 -0
  13. warp/_src/fabric.py +362 -0
  14. warp/_src/fem/__init__.py +14 -0
  15. warp/_src/fem/adaptivity.py +510 -0
  16. warp/_src/fem/cache.py +689 -0
  17. warp/_src/fem/dirichlet.py +190 -0
  18. warp/_src/fem/domain.py +553 -0
  19. warp/_src/fem/field/__init__.py +131 -0
  20. warp/_src/fem/field/field.py +703 -0
  21. warp/_src/fem/field/nodal_field.py +403 -0
  22. warp/_src/fem/field/restriction.py +39 -0
  23. warp/_src/fem/field/virtual.py +1021 -0
  24. warp/_src/fem/geometry/__init__.py +32 -0
  25. warp/_src/fem/geometry/adaptive_nanogrid.py +782 -0
  26. warp/_src/fem/geometry/closest_point.py +99 -0
  27. warp/_src/fem/geometry/deformed_geometry.py +277 -0
  28. warp/_src/fem/geometry/element.py +854 -0
  29. warp/_src/fem/geometry/geometry.py +693 -0
  30. warp/_src/fem/geometry/grid_2d.py +478 -0
  31. warp/_src/fem/geometry/grid_3d.py +539 -0
  32. warp/_src/fem/geometry/hexmesh.py +956 -0
  33. warp/_src/fem/geometry/nanogrid.py +660 -0
  34. warp/_src/fem/geometry/partition.py +483 -0
  35. warp/_src/fem/geometry/quadmesh.py +597 -0
  36. warp/_src/fem/geometry/tetmesh.py +762 -0
  37. warp/_src/fem/geometry/trimesh.py +588 -0
  38. warp/_src/fem/integrate.py +2507 -0
  39. warp/_src/fem/linalg.py +385 -0
  40. warp/_src/fem/operator.py +398 -0
  41. warp/_src/fem/polynomial.py +231 -0
  42. warp/_src/fem/quadrature/__init__.py +17 -0
  43. warp/_src/fem/quadrature/pic_quadrature.py +318 -0
  44. warp/_src/fem/quadrature/quadrature.py +665 -0
  45. warp/_src/fem/space/__init__.py +248 -0
  46. warp/_src/fem/space/basis_function_space.py +499 -0
  47. warp/_src/fem/space/basis_space.py +681 -0
  48. warp/_src/fem/space/dof_mapper.py +253 -0
  49. warp/_src/fem/space/function_space.py +312 -0
  50. warp/_src/fem/space/grid_2d_function_space.py +179 -0
  51. warp/_src/fem/space/grid_3d_function_space.py +229 -0
  52. warp/_src/fem/space/hexmesh_function_space.py +255 -0
  53. warp/_src/fem/space/nanogrid_function_space.py +199 -0
  54. warp/_src/fem/space/partition.py +435 -0
  55. warp/_src/fem/space/quadmesh_function_space.py +222 -0
  56. warp/_src/fem/space/restriction.py +221 -0
  57. warp/_src/fem/space/shape/__init__.py +152 -0
  58. warp/_src/fem/space/shape/cube_shape_function.py +1107 -0
  59. warp/_src/fem/space/shape/shape_function.py +134 -0
  60. warp/_src/fem/space/shape/square_shape_function.py +928 -0
  61. warp/_src/fem/space/shape/tet_shape_function.py +829 -0
  62. warp/_src/fem/space/shape/triangle_shape_function.py +674 -0
  63. warp/_src/fem/space/tetmesh_function_space.py +270 -0
  64. warp/_src/fem/space/topology.py +461 -0
  65. warp/_src/fem/space/trimesh_function_space.py +193 -0
  66. warp/_src/fem/types.py +114 -0
  67. warp/_src/fem/utils.py +488 -0
  68. warp/_src/jax.py +188 -0
  69. warp/_src/jax_experimental/__init__.py +14 -0
  70. warp/_src/jax_experimental/custom_call.py +389 -0
  71. warp/_src/jax_experimental/ffi.py +1286 -0
  72. warp/_src/jax_experimental/xla_ffi.py +658 -0
  73. warp/_src/marching_cubes.py +710 -0
  74. warp/_src/math.py +416 -0
  75. warp/_src/optim/__init__.py +14 -0
  76. warp/_src/optim/adam.py +165 -0
  77. warp/_src/optim/linear.py +1608 -0
  78. warp/_src/optim/sgd.py +114 -0
  79. warp/_src/paddle.py +408 -0
  80. warp/_src/render/__init__.py +14 -0
  81. warp/_src/render/imgui_manager.py +291 -0
  82. warp/_src/render/render_opengl.py +3638 -0
  83. warp/_src/render/render_usd.py +939 -0
  84. warp/_src/render/utils.py +162 -0
  85. warp/_src/sparse.py +2718 -0
  86. warp/_src/tape.py +1208 -0
  87. warp/_src/thirdparty/__init__.py +0 -0
  88. warp/_src/thirdparty/appdirs.py +598 -0
  89. warp/_src/thirdparty/dlpack.py +145 -0
  90. warp/_src/thirdparty/unittest_parallel.py +676 -0
  91. warp/_src/torch.py +393 -0
  92. warp/_src/types.py +5888 -0
  93. warp/_src/utils.py +1695 -0
  94. warp/autograd.py +33 -0
  95. warp/bin/libwarp-clang.dylib +0 -0
  96. warp/bin/libwarp.dylib +0 -0
  97. warp/build.py +29 -0
  98. warp/build_dll.py +24 -0
  99. warp/codegen.py +24 -0
  100. warp/constants.py +24 -0
  101. warp/context.py +33 -0
  102. warp/dlpack.py +24 -0
  103. warp/examples/__init__.py +24 -0
  104. warp/examples/assets/bear.usd +0 -0
  105. warp/examples/assets/bunny.usd +0 -0
  106. warp/examples/assets/cube.usd +0 -0
  107. warp/examples/assets/nonuniform.usd +0 -0
  108. warp/examples/assets/nvidia_logo.png +0 -0
  109. warp/examples/assets/pixel.jpg +0 -0
  110. warp/examples/assets/rocks.nvdb +0 -0
  111. warp/examples/assets/rocks.usd +0 -0
  112. warp/examples/assets/sphere.usd +0 -0
  113. warp/examples/assets/square_cloth.usd +0 -0
  114. warp/examples/benchmarks/benchmark_api.py +389 -0
  115. warp/examples/benchmarks/benchmark_cloth.py +296 -0
  116. warp/examples/benchmarks/benchmark_cloth_cupy.py +96 -0
  117. warp/examples/benchmarks/benchmark_cloth_jax.py +105 -0
  118. warp/examples/benchmarks/benchmark_cloth_numba.py +161 -0
  119. warp/examples/benchmarks/benchmark_cloth_numpy.py +85 -0
  120. warp/examples/benchmarks/benchmark_cloth_paddle.py +94 -0
  121. warp/examples/benchmarks/benchmark_cloth_pytorch.py +94 -0
  122. warp/examples/benchmarks/benchmark_cloth_taichi.py +120 -0
  123. warp/examples/benchmarks/benchmark_cloth_warp.py +153 -0
  124. warp/examples/benchmarks/benchmark_gemm.py +164 -0
  125. warp/examples/benchmarks/benchmark_interop_paddle.py +166 -0
  126. warp/examples/benchmarks/benchmark_interop_torch.py +166 -0
  127. warp/examples/benchmarks/benchmark_launches.py +301 -0
  128. warp/examples/benchmarks/benchmark_tile_load_store.py +103 -0
  129. warp/examples/benchmarks/benchmark_tile_sort.py +155 -0
  130. warp/examples/browse.py +37 -0
  131. warp/examples/core/example_cupy.py +86 -0
  132. warp/examples/core/example_dem.py +241 -0
  133. warp/examples/core/example_fluid.py +299 -0
  134. warp/examples/core/example_graph_capture.py +150 -0
  135. warp/examples/core/example_marching_cubes.py +195 -0
  136. warp/examples/core/example_mesh.py +180 -0
  137. warp/examples/core/example_mesh_intersect.py +211 -0
  138. warp/examples/core/example_nvdb.py +182 -0
  139. warp/examples/core/example_raycast.py +111 -0
  140. warp/examples/core/example_raymarch.py +205 -0
  141. warp/examples/core/example_render_opengl.py +290 -0
  142. warp/examples/core/example_sample_mesh.py +300 -0
  143. warp/examples/core/example_sph.py +411 -0
  144. warp/examples/core/example_spin_lock.py +93 -0
  145. warp/examples/core/example_torch.py +211 -0
  146. warp/examples/core/example_wave.py +269 -0
  147. warp/examples/core/example_work_queue.py +118 -0
  148. warp/examples/distributed/example_jacobi_mpi.py +506 -0
  149. warp/examples/fem/example_adaptive_grid.py +286 -0
  150. warp/examples/fem/example_apic_fluid.py +469 -0
  151. warp/examples/fem/example_burgers.py +261 -0
  152. warp/examples/fem/example_convection_diffusion.py +181 -0
  153. warp/examples/fem/example_convection_diffusion_dg.py +225 -0
  154. warp/examples/fem/example_darcy_ls_optimization.py +489 -0
  155. warp/examples/fem/example_deformed_geometry.py +172 -0
  156. warp/examples/fem/example_diffusion.py +196 -0
  157. warp/examples/fem/example_diffusion_3d.py +225 -0
  158. warp/examples/fem/example_diffusion_mgpu.py +225 -0
  159. warp/examples/fem/example_distortion_energy.py +228 -0
  160. warp/examples/fem/example_elastic_shape_optimization.py +387 -0
  161. warp/examples/fem/example_magnetostatics.py +242 -0
  162. warp/examples/fem/example_mixed_elasticity.py +293 -0
  163. warp/examples/fem/example_navier_stokes.py +263 -0
  164. warp/examples/fem/example_nonconforming_contact.py +300 -0
  165. warp/examples/fem/example_stokes.py +213 -0
  166. warp/examples/fem/example_stokes_transfer.py +262 -0
  167. warp/examples/fem/example_streamlines.py +357 -0
  168. warp/examples/fem/utils.py +1047 -0
  169. warp/examples/interop/example_jax_callable.py +146 -0
  170. warp/examples/interop/example_jax_ffi_callback.py +132 -0
  171. warp/examples/interop/example_jax_kernel.py +232 -0
  172. warp/examples/optim/example_diffray.py +561 -0
  173. warp/examples/optim/example_fluid_checkpoint.py +497 -0
  174. warp/examples/tile/example_tile_block_cholesky.py +502 -0
  175. warp/examples/tile/example_tile_cholesky.py +88 -0
  176. warp/examples/tile/example_tile_convolution.py +66 -0
  177. warp/examples/tile/example_tile_fft.py +55 -0
  178. warp/examples/tile/example_tile_filtering.py +113 -0
  179. warp/examples/tile/example_tile_matmul.py +85 -0
  180. warp/examples/tile/example_tile_mcgp.py +191 -0
  181. warp/examples/tile/example_tile_mlp.py +385 -0
  182. warp/examples/tile/example_tile_nbody.py +199 -0
  183. warp/fabric.py +24 -0
  184. warp/fem/__init__.py +173 -0
  185. warp/fem/adaptivity.py +26 -0
  186. warp/fem/cache.py +30 -0
  187. warp/fem/dirichlet.py +24 -0
  188. warp/fem/field/__init__.py +24 -0
  189. warp/fem/field/field.py +26 -0
  190. warp/fem/geometry/__init__.py +21 -0
  191. warp/fem/geometry/closest_point.py +31 -0
  192. warp/fem/linalg.py +38 -0
  193. warp/fem/operator.py +32 -0
  194. warp/fem/polynomial.py +29 -0
  195. warp/fem/space/__init__.py +22 -0
  196. warp/fem/space/basis_space.py +24 -0
  197. warp/fem/space/shape/__init__.py +68 -0
  198. warp/fem/space/topology.py +24 -0
  199. warp/fem/types.py +24 -0
  200. warp/fem/utils.py +32 -0
  201. warp/jax.py +29 -0
  202. warp/jax_experimental/__init__.py +29 -0
  203. warp/jax_experimental/custom_call.py +29 -0
  204. warp/jax_experimental/ffi.py +39 -0
  205. warp/jax_experimental/xla_ffi.py +24 -0
  206. warp/marching_cubes.py +24 -0
  207. warp/math.py +37 -0
  208. warp/native/array.h +1687 -0
  209. warp/native/builtin.h +2327 -0
  210. warp/native/bvh.cpp +562 -0
  211. warp/native/bvh.cu +826 -0
  212. warp/native/bvh.h +555 -0
  213. warp/native/clang/clang.cpp +541 -0
  214. warp/native/coloring.cpp +622 -0
  215. warp/native/crt.cpp +51 -0
  216. warp/native/crt.h +568 -0
  217. warp/native/cuda_crt.h +1058 -0
  218. warp/native/cuda_util.cpp +677 -0
  219. warp/native/cuda_util.h +313 -0
  220. warp/native/error.cpp +77 -0
  221. warp/native/error.h +36 -0
  222. warp/native/exports.h +2023 -0
  223. warp/native/fabric.h +246 -0
  224. warp/native/hashgrid.cpp +311 -0
  225. warp/native/hashgrid.cu +89 -0
  226. warp/native/hashgrid.h +240 -0
  227. warp/native/initializer_array.h +41 -0
  228. warp/native/intersect.h +1253 -0
  229. warp/native/intersect_adj.h +375 -0
  230. warp/native/intersect_tri.h +348 -0
  231. warp/native/mat.h +5189 -0
  232. warp/native/mathdx.cpp +93 -0
  233. warp/native/matnn.h +221 -0
  234. warp/native/mesh.cpp +266 -0
  235. warp/native/mesh.cu +406 -0
  236. warp/native/mesh.h +2097 -0
  237. warp/native/nanovdb/GridHandle.h +533 -0
  238. warp/native/nanovdb/HostBuffer.h +591 -0
  239. warp/native/nanovdb/NanoVDB.h +6246 -0
  240. warp/native/nanovdb/NodeManager.h +323 -0
  241. warp/native/nanovdb/PNanoVDB.h +3390 -0
  242. warp/native/noise.h +859 -0
  243. warp/native/quat.h +1664 -0
  244. warp/native/rand.h +342 -0
  245. warp/native/range.h +145 -0
  246. warp/native/reduce.cpp +174 -0
  247. warp/native/reduce.cu +363 -0
  248. warp/native/runlength_encode.cpp +79 -0
  249. warp/native/runlength_encode.cu +61 -0
  250. warp/native/scan.cpp +47 -0
  251. warp/native/scan.cu +55 -0
  252. warp/native/scan.h +23 -0
  253. warp/native/solid_angle.h +466 -0
  254. warp/native/sort.cpp +251 -0
  255. warp/native/sort.cu +286 -0
  256. warp/native/sort.h +35 -0
  257. warp/native/sparse.cpp +241 -0
  258. warp/native/sparse.cu +435 -0
  259. warp/native/spatial.h +1306 -0
  260. warp/native/svd.h +727 -0
  261. warp/native/temp_buffer.h +46 -0
  262. warp/native/tile.h +4124 -0
  263. warp/native/tile_radix_sort.h +1112 -0
  264. warp/native/tile_reduce.h +838 -0
  265. warp/native/tile_scan.h +240 -0
  266. warp/native/tuple.h +189 -0
  267. warp/native/vec.h +2199 -0
  268. warp/native/version.h +23 -0
  269. warp/native/volume.cpp +501 -0
  270. warp/native/volume.cu +68 -0
  271. warp/native/volume.h +970 -0
  272. warp/native/volume_builder.cu +483 -0
  273. warp/native/volume_builder.h +52 -0
  274. warp/native/volume_impl.h +70 -0
  275. warp/native/warp.cpp +1143 -0
  276. warp/native/warp.cu +4604 -0
  277. warp/native/warp.h +358 -0
  278. warp/optim/__init__.py +20 -0
  279. warp/optim/adam.py +24 -0
  280. warp/optim/linear.py +35 -0
  281. warp/optim/sgd.py +24 -0
  282. warp/paddle.py +24 -0
  283. warp/py.typed +0 -0
  284. warp/render/__init__.py +22 -0
  285. warp/render/imgui_manager.py +29 -0
  286. warp/render/render_opengl.py +24 -0
  287. warp/render/render_usd.py +24 -0
  288. warp/render/utils.py +24 -0
  289. warp/sparse.py +51 -0
  290. warp/tape.py +24 -0
  291. warp/tests/__init__.py +1 -0
  292. warp/tests/__main__.py +4 -0
  293. warp/tests/assets/curlnoise_golden.npy +0 -0
  294. warp/tests/assets/mlp_golden.npy +0 -0
  295. warp/tests/assets/pixel.npy +0 -0
  296. warp/tests/assets/pnoise_golden.npy +0 -0
  297. warp/tests/assets/spiky.usd +0 -0
  298. warp/tests/assets/test_grid.nvdb +0 -0
  299. warp/tests/assets/test_index_grid.nvdb +0 -0
  300. warp/tests/assets/test_int32_grid.nvdb +0 -0
  301. warp/tests/assets/test_vec_grid.nvdb +0 -0
  302. warp/tests/assets/torus.nvdb +0 -0
  303. warp/tests/assets/torus.usda +105 -0
  304. warp/tests/aux_test_class_kernel.py +34 -0
  305. warp/tests/aux_test_compile_consts_dummy.py +18 -0
  306. warp/tests/aux_test_conditional_unequal_types_kernels.py +29 -0
  307. warp/tests/aux_test_dependent.py +29 -0
  308. warp/tests/aux_test_grad_customs.py +29 -0
  309. warp/tests/aux_test_instancing_gc.py +26 -0
  310. warp/tests/aux_test_module_aot.py +7 -0
  311. warp/tests/aux_test_module_unload.py +23 -0
  312. warp/tests/aux_test_name_clash1.py +40 -0
  313. warp/tests/aux_test_name_clash2.py +40 -0
  314. warp/tests/aux_test_reference.py +9 -0
  315. warp/tests/aux_test_reference_reference.py +8 -0
  316. warp/tests/aux_test_square.py +16 -0
  317. warp/tests/aux_test_unresolved_func.py +22 -0
  318. warp/tests/aux_test_unresolved_symbol.py +22 -0
  319. warp/tests/cuda/__init__.py +0 -0
  320. warp/tests/cuda/test_async.py +676 -0
  321. warp/tests/cuda/test_conditional_captures.py +1147 -0
  322. warp/tests/cuda/test_ipc.py +124 -0
  323. warp/tests/cuda/test_mempool.py +233 -0
  324. warp/tests/cuda/test_multigpu.py +169 -0
  325. warp/tests/cuda/test_peer.py +139 -0
  326. warp/tests/cuda/test_pinned.py +84 -0
  327. warp/tests/cuda/test_streams.py +691 -0
  328. warp/tests/geometry/__init__.py +0 -0
  329. warp/tests/geometry/test_bvh.py +335 -0
  330. warp/tests/geometry/test_hash_grid.py +259 -0
  331. warp/tests/geometry/test_marching_cubes.py +294 -0
  332. warp/tests/geometry/test_mesh.py +318 -0
  333. warp/tests/geometry/test_mesh_query_aabb.py +392 -0
  334. warp/tests/geometry/test_mesh_query_point.py +935 -0
  335. warp/tests/geometry/test_mesh_query_ray.py +323 -0
  336. warp/tests/geometry/test_volume.py +1103 -0
  337. warp/tests/geometry/test_volume_write.py +346 -0
  338. warp/tests/interop/__init__.py +0 -0
  339. warp/tests/interop/test_dlpack.py +730 -0
  340. warp/tests/interop/test_jax.py +1673 -0
  341. warp/tests/interop/test_paddle.py +800 -0
  342. warp/tests/interop/test_torch.py +1001 -0
  343. warp/tests/run_coverage_serial.py +39 -0
  344. warp/tests/test_adam.py +162 -0
  345. warp/tests/test_arithmetic.py +1096 -0
  346. warp/tests/test_array.py +3756 -0
  347. warp/tests/test_array_reduce.py +156 -0
  348. warp/tests/test_assert.py +303 -0
  349. warp/tests/test_atomic.py +336 -0
  350. warp/tests/test_atomic_bitwise.py +209 -0
  351. warp/tests/test_atomic_cas.py +312 -0
  352. warp/tests/test_bool.py +220 -0
  353. warp/tests/test_builtins_resolution.py +732 -0
  354. warp/tests/test_closest_point_edge_edge.py +327 -0
  355. warp/tests/test_codegen.py +974 -0
  356. warp/tests/test_codegen_instancing.py +1495 -0
  357. warp/tests/test_compile_consts.py +215 -0
  358. warp/tests/test_conditional.py +298 -0
  359. warp/tests/test_context.py +35 -0
  360. warp/tests/test_copy.py +319 -0
  361. warp/tests/test_ctypes.py +618 -0
  362. warp/tests/test_dense.py +73 -0
  363. warp/tests/test_devices.py +127 -0
  364. warp/tests/test_enum.py +136 -0
  365. warp/tests/test_examples.py +424 -0
  366. warp/tests/test_fabricarray.py +998 -0
  367. warp/tests/test_fast_math.py +72 -0
  368. warp/tests/test_fem.py +2204 -0
  369. warp/tests/test_fixedarray.py +229 -0
  370. warp/tests/test_fp16.py +136 -0
  371. warp/tests/test_func.py +501 -0
  372. warp/tests/test_future_annotations.py +100 -0
  373. warp/tests/test_generics.py +656 -0
  374. warp/tests/test_grad.py +893 -0
  375. warp/tests/test_grad_customs.py +339 -0
  376. warp/tests/test_grad_debug.py +341 -0
  377. warp/tests/test_implicit_init.py +411 -0
  378. warp/tests/test_import.py +45 -0
  379. warp/tests/test_indexedarray.py +1140 -0
  380. warp/tests/test_intersect.py +103 -0
  381. warp/tests/test_iter.py +76 -0
  382. warp/tests/test_large.py +177 -0
  383. warp/tests/test_launch.py +411 -0
  384. warp/tests/test_lerp.py +151 -0
  385. warp/tests/test_linear_solvers.py +223 -0
  386. warp/tests/test_lvalue.py +427 -0
  387. warp/tests/test_map.py +526 -0
  388. warp/tests/test_mat.py +3515 -0
  389. warp/tests/test_mat_assign_copy.py +178 -0
  390. warp/tests/test_mat_constructors.py +573 -0
  391. warp/tests/test_mat_lite.py +122 -0
  392. warp/tests/test_mat_scalar_ops.py +2913 -0
  393. warp/tests/test_math.py +212 -0
  394. warp/tests/test_module_aot.py +287 -0
  395. warp/tests/test_module_hashing.py +258 -0
  396. warp/tests/test_modules_lite.py +70 -0
  397. warp/tests/test_noise.py +252 -0
  398. warp/tests/test_operators.py +299 -0
  399. warp/tests/test_options.py +129 -0
  400. warp/tests/test_overwrite.py +551 -0
  401. warp/tests/test_print.py +408 -0
  402. warp/tests/test_quat.py +2653 -0
  403. warp/tests/test_quat_assign_copy.py +145 -0
  404. warp/tests/test_rand.py +339 -0
  405. warp/tests/test_reload.py +303 -0
  406. warp/tests/test_rounding.py +157 -0
  407. warp/tests/test_runlength_encode.py +196 -0
  408. warp/tests/test_scalar_ops.py +133 -0
  409. warp/tests/test_smoothstep.py +108 -0
  410. warp/tests/test_snippet.py +318 -0
  411. warp/tests/test_sparse.py +845 -0
  412. warp/tests/test_spatial.py +2859 -0
  413. warp/tests/test_spatial_assign_copy.py +160 -0
  414. warp/tests/test_special_values.py +361 -0
  415. warp/tests/test_static.py +640 -0
  416. warp/tests/test_struct.py +901 -0
  417. warp/tests/test_tape.py +242 -0
  418. warp/tests/test_transient_module.py +93 -0
  419. warp/tests/test_triangle_closest_point.py +192 -0
  420. warp/tests/test_tuple.py +361 -0
  421. warp/tests/test_types.py +615 -0
  422. warp/tests/test_utils.py +594 -0
  423. warp/tests/test_vec.py +1408 -0
  424. warp/tests/test_vec_assign_copy.py +143 -0
  425. warp/tests/test_vec_constructors.py +325 -0
  426. warp/tests/test_vec_lite.py +80 -0
  427. warp/tests/test_vec_scalar_ops.py +2327 -0
  428. warp/tests/test_verify_fp.py +100 -0
  429. warp/tests/test_version.py +75 -0
  430. warp/tests/tile/__init__.py +0 -0
  431. warp/tests/tile/test_tile.py +1519 -0
  432. warp/tests/tile/test_tile_atomic_bitwise.py +403 -0
  433. warp/tests/tile/test_tile_cholesky.py +608 -0
  434. warp/tests/tile/test_tile_load.py +724 -0
  435. warp/tests/tile/test_tile_mathdx.py +156 -0
  436. warp/tests/tile/test_tile_matmul.py +179 -0
  437. warp/tests/tile/test_tile_mlp.py +400 -0
  438. warp/tests/tile/test_tile_reduce.py +950 -0
  439. warp/tests/tile/test_tile_shared_memory.py +376 -0
  440. warp/tests/tile/test_tile_sort.py +121 -0
  441. warp/tests/tile/test_tile_view.py +173 -0
  442. warp/tests/unittest_serial.py +47 -0
  443. warp/tests/unittest_suites.py +430 -0
  444. warp/tests/unittest_utils.py +469 -0
  445. warp/tests/walkthrough_debug.py +95 -0
  446. warp/torch.py +24 -0
  447. warp/types.py +51 -0
  448. warp/utils.py +31 -0
  449. warp_lang-1.10.0.dist-info/METADATA +459 -0
  450. warp_lang-1.10.0.dist-info/RECORD +468 -0
  451. warp_lang-1.10.0.dist-info/WHEEL +5 -0
  452. warp_lang-1.10.0.dist-info/licenses/LICENSE.md +176 -0
  453. warp_lang-1.10.0.dist-info/licenses/licenses/Gaia-LICENSE.txt +6 -0
  454. warp_lang-1.10.0.dist-info/licenses/licenses/appdirs-LICENSE.txt +22 -0
  455. warp_lang-1.10.0.dist-info/licenses/licenses/asset_pixel_jpg-LICENSE.txt +3 -0
  456. warp_lang-1.10.0.dist-info/licenses/licenses/cuda-LICENSE.txt +1582 -0
  457. warp_lang-1.10.0.dist-info/licenses/licenses/dlpack-LICENSE.txt +201 -0
  458. warp_lang-1.10.0.dist-info/licenses/licenses/fp16-LICENSE.txt +28 -0
  459. warp_lang-1.10.0.dist-info/licenses/licenses/libmathdx-LICENSE.txt +220 -0
  460. warp_lang-1.10.0.dist-info/licenses/licenses/llvm-LICENSE.txt +279 -0
  461. warp_lang-1.10.0.dist-info/licenses/licenses/moller-LICENSE.txt +16 -0
  462. warp_lang-1.10.0.dist-info/licenses/licenses/nanovdb-LICENSE.txt +2 -0
  463. warp_lang-1.10.0.dist-info/licenses/licenses/nvrtc-LICENSE.txt +1592 -0
  464. warp_lang-1.10.0.dist-info/licenses/licenses/svd-LICENSE.txt +23 -0
  465. warp_lang-1.10.0.dist-info/licenses/licenses/unittest_parallel-LICENSE.txt +21 -0
  466. warp_lang-1.10.0.dist-info/licenses/licenses/usd-LICENSE.txt +213 -0
  467. warp_lang-1.10.0.dist-info/licenses/licenses/windingnumber-LICENSE.txt +21 -0
  468. warp_lang-1.10.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,398 @@
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
+ from typing import Any, Callable, Dict, Optional, Set
17
+
18
+ import warp as wp
19
+ from warp._src.codegen import get_full_arg_spec, make_full_qualified_name
20
+ from warp._src.fem.linalg import skew_part, symmetric_part
21
+ from warp._src.fem.types import (
22
+ Coords,
23
+ Domain,
24
+ ElementIndex,
25
+ Field,
26
+ NodeIndex,
27
+ Sample,
28
+ make_free_sample,
29
+ )
30
+
31
+ _wp_module_name_ = "warp.fem.operator"
32
+
33
+
34
+ class Integrand:
35
+ """An integrand is a device function containing arbitrary expressions over Field and Domain variables.
36
+ It will get transformed to a proper warp.Function by resolving concrete Field types at call time.
37
+ """
38
+
39
+ def __init__(self, func: Callable, kernel_options: Optional[Dict[str, Any]] = None):
40
+ self.func = func
41
+ self.name = make_full_qualified_name(self.func)
42
+ self.module = wp.get_module(self.func.__module__)
43
+ self.argspec = get_full_arg_spec(self.func)
44
+ self.kernel_options = {} if kernel_options is None else kernel_options
45
+
46
+ # Operators for each field argument. This will be populated at first integrate call
47
+ self.operators: Optional[Dict[str, Set[Operator]]] = None
48
+
49
+ # Cached kernels for each integrand call
50
+ self.cached_kernels = {}
51
+ self.cached_funcs = {}
52
+
53
+
54
+ class Operator:
55
+ """
56
+ Operators provide syntactic sugar over Field and Domain evaluation functions and arguments
57
+ """
58
+
59
+ def __init__(
60
+ self,
61
+ func: Callable,
62
+ resolver: Callable,
63
+ field_result: Optional[Callable] = None,
64
+ attr: Optional[str] = None,
65
+ ):
66
+ self.func = func
67
+ self.name = func.__name__
68
+ self.resolver = resolver
69
+ self.attr = attr
70
+ self.field_result = field_result
71
+
72
+
73
+ def integrand(func: Optional[Callable] = None, kernel_options: Optional[Dict[str, Any]] = None):
74
+ """Decorator for functions to be integrated (or interpolated) using warp.fem
75
+
76
+ Args:
77
+ func: Decorated function
78
+ kernel_options: Supplemental code-generation options to be passed to the generated kernel.
79
+ """
80
+
81
+ if func is not None:
82
+ itg = Integrand(func)
83
+ itg.__doc__ = func.__doc__
84
+ return itg
85
+
86
+ def wrap_integrand(func: Callable):
87
+ itg = Integrand(func, kernel_options)
88
+ itg.__doc__ = func.__doc__
89
+ return itg
90
+
91
+ return wrap_integrand
92
+
93
+
94
+ def operator(**kwargs):
95
+ """Decorator for functions operating on Field-like or Domain-like data inside warp.fem integrands"""
96
+
97
+ def wrap_operator(func: Callable):
98
+ op = Operator(func, **kwargs)
99
+ op.__doc__ = func.__doc__
100
+ return op
101
+
102
+ return wrap_operator
103
+
104
+
105
+ # Domain operators
106
+
107
+
108
+ @operator(resolver=lambda dmn: dmn.element_position, attr="geo")
109
+ def position(domain: Domain, s: Sample):
110
+ """Evaluates the world position of the sample point `s`"""
111
+ pass
112
+
113
+
114
+ @operator(resolver=lambda dmn: dmn.element_normal, attr="geo")
115
+ def normal(domain: Domain, s: Sample):
116
+ """Evaluates the element normal at the sample point `s`. Non zero if the element is a side or the geometry is embedded in a higher-dimensional space (e.g. :class:`Trimesh3D`)"""
117
+ pass
118
+
119
+
120
+ @operator(resolver=lambda dmn: dmn.element_deformation_gradient, attr="geo")
121
+ def deformation_gradient(domain: Domain, s: Sample):
122
+ """Evaluates the gradient of the domain position with respect to the element reference space at the sample point `s`"""
123
+ pass
124
+
125
+
126
+ @operator(resolver=lambda dmn: dmn.element_lookup, attr="geo")
127
+ def lookup(domain: Domain, x: Any) -> Sample:
128
+ """Looks-up the sample point corresponding to a world position `x`, projecting to the closest point on the geometry.
129
+
130
+ Args:
131
+ x (vec3): world position of the point to look-up in the geometry
132
+ max_dist (float): maximum distance to look for a closest point
133
+ guess (:class:`Sample`): initial guess, may help perform the query
134
+ filter_array (wp.array): Used in conjunction with `filter_target`. Only cells such that ``filter_array[element_index]==filter_target`` will be considered.
135
+ filter_target (Any): See `filter_array`
136
+ """
137
+ pass
138
+
139
+
140
+ @operator(resolver=lambda dmn: dmn.element_partition_lookup)
141
+ def partition_lookup(domain: Domain, x: Any) -> Sample:
142
+ """Looks-up the sample point corresponding to a world position `x`, projecting to the closest point on the geometry partition.
143
+
144
+ Args:
145
+ x (vec3): world position of the point to look-up in the geometry
146
+ max_dist (float): maximum distance to look for a closest point
147
+ """
148
+ pass
149
+
150
+
151
+ @operator(resolver=lambda dmn: dmn.element_measure, attr="geo")
152
+ def measure(domain: Domain, s: Sample) -> float:
153
+ """Returns the measure (volume, area, or length) determinant of an element at a sample point `s`"""
154
+ pass
155
+
156
+
157
+ @operator(resolver=lambda dmn: dmn.element_measure_ratio, attr="geo")
158
+ def measure_ratio(domain: Domain, s: Sample) -> float:
159
+ """Returns the maximum ratio between the measure of this element and that of higher-dimensional neighbors."""
160
+ pass
161
+
162
+
163
+ @operator(resolver=lambda dmn: dmn.element_closest_point, attr="geo")
164
+ def element_closest_point(domain: Domain, element_index: ElementIndex, x: Any) -> Sample:
165
+ """
166
+ Computes the coordinates of the closest point to a world position within a given element.
167
+ Returns a tuple (closest point coordinates; squared distance to the closest point)
168
+
169
+ Args:
170
+ element_index: Index of the element to consider
171
+ x: world position of the point to compute the closest point to
172
+ """
173
+ pass
174
+
175
+
176
+ @operator(resolver=lambda dmn: dmn.element_coordinates, attr="geo")
177
+ def element_coordinates(domain: Domain, element_index: ElementIndex, x: Any) -> Sample:
178
+ """Returns the coordinates in an element reference system corresponding to a work position.
179
+ The returned coordinates may be in the element's exterior.
180
+
181
+ Args:
182
+ element_index: Index of the element to consider
183
+ x: world position of the point to find coordinates for
184
+ """
185
+ pass
186
+
187
+
188
+ # Operators for evaluating cell-level quantities on domains defined on sides
189
+
190
+
191
+ @operator(
192
+ resolver=lambda dmn: dmn.domain_cell_arg,
193
+ field_result=lambda dmn: (dmn.cell_domain(), Domain, dmn.cell_domain().DomainArg),
194
+ )
195
+ def cells(domain: Domain) -> Domain:
196
+ """Converts a domain defined on geometry sides to a domain defined of cells."""
197
+ pass
198
+
199
+
200
+ @operator(resolver=lambda dmn: dmn.element_inner_cell_index, attr="geo")
201
+ def _inner_cell_index(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
202
+ pass
203
+
204
+
205
+ @operator(resolver=lambda dmn: dmn.element_outer_cell_index, attr="geo")
206
+ def _outer_cell_index(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
207
+ pass
208
+
209
+
210
+ @operator(resolver=lambda dmn: dmn.element_inner_cell_coords, attr="geo")
211
+ def _inner_cell_coords(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
212
+ pass
213
+
214
+
215
+ @operator(resolver=lambda dmn: dmn.element_outer_cell_coords, attr="geo")
216
+ def _outer_cell_coords(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
217
+ pass
218
+
219
+
220
+ @operator(resolver=lambda dmn: dmn.cell_to_element_coords, attr="geo")
221
+ def _cell_to_element_coords(
222
+ domain: Domain,
223
+ side_index: ElementIndex,
224
+ cell_index: ElementIndex,
225
+ cell_coords: Coords,
226
+ ) -> Sample:
227
+ pass
228
+
229
+
230
+ @integrand
231
+ def to_inner_cell(domain: Domain, s: Sample):
232
+ """Converts a :class:`Sample` defined on a side to a sample defined on the side's inner cell"""
233
+ return make_free_sample(
234
+ _inner_cell_index(domain, s.element_index),
235
+ _inner_cell_coords(domain, s.element_index, s.element_coords),
236
+ )
237
+
238
+
239
+ @integrand
240
+ def to_outer_cell(domain: Domain, s: Sample):
241
+ """Converts a :class:`Sample` defined on a side to a sample defined on the side's outer cell"""
242
+ return make_free_sample(
243
+ _outer_cell_index(domain, s.element_index),
244
+ _outer_cell_coords(domain, s.element_index, s.element_coords),
245
+ )
246
+
247
+
248
+ @integrand
249
+ def to_cell_side(domain: Domain, cell_s: Sample, side_index: ElementIndex):
250
+ """Converts a :class:`Sample` defined on a cell to a sample defined on one of its side.
251
+ If the result does not lie on the side `side_index`, the resulting coordinates will be set to ``OUTSIDE``."""
252
+ return make_free_sample(
253
+ side_index,
254
+ _cell_to_element_coords(domain, side_index, cell_s.element_index, cell_s.element_coords),
255
+ )
256
+
257
+
258
+ @operator(resolver=lambda dmn: dmn.element_index, attr="index")
259
+ def element_index(domain: Domain, domain_element_index: ElementIndex):
260
+ """Returns the index in the geometry of the `domain_element_index`'th domain element."""
261
+ pass
262
+
263
+
264
+ @operator(resolver=lambda dmn: dmn.element_partition_index, attr="index")
265
+ def element_partition_index(domain: Domain, cell_index: ElementIndex):
266
+ """Returns the index of the passed cell in the domain's geometry partition, or `NULL_ELEMENT_INDEX` if not part of the partition.
267
+
268
+ :note: Currently only available for `fem.ElementKind.CELL` elements
269
+ """
270
+ pass
271
+
272
+
273
+ # Field operators
274
+ # On a side, inner and outer are such that normal goes from inner to outer
275
+
276
+
277
+ @operator(resolver=lambda f: f.eval_inner)
278
+ def inner(f: Field, s: Sample):
279
+ """Evaluates the field at a sample point `s`. On oriented sides, uses the inner element"""
280
+ pass
281
+
282
+
283
+ @operator(resolver=lambda f: f.eval_grad_inner)
284
+ def grad(f: Field, s: Sample):
285
+ """Evaluates the field gradient at a sample point `s`. On oriented sides, uses the inner element"""
286
+ pass
287
+
288
+
289
+ @operator(resolver=lambda f: f.eval_div_inner)
290
+ def div(f: Field, s: Sample):
291
+ """Evaluates the field divergence at a sample point `s`. On oriented sides, uses the inner element"""
292
+ pass
293
+
294
+
295
+ @operator(resolver=lambda f: f.eval_outer)
296
+ def outer(f: Field, s: Sample):
297
+ """Evaluates the field at a sample point `s`. On oriented sides, uses the outer element. On interior points and on domain boundaries, this is equivalent to :func:`inner`."""
298
+ pass
299
+
300
+
301
+ @operator(resolver=lambda f: f.eval_grad_outer)
302
+ def grad_outer(f: Field, s: Sample):
303
+ """Evaluates the field gradient at a sample point `s`. On oriented sides, uses the outer element. On interior points and on domain boundaries, this is equivalent to :func:`grad`."""
304
+ pass
305
+
306
+
307
+ @operator(resolver=lambda f: f.eval_div_outer)
308
+ def div_outer(f: Field, s: Sample):
309
+ """Evaluates the field divergence at a sample point `s`. On oriented sides, uses the outer element. On interior points and on domain boundaries, this is equivalent to :func:`div`."""
310
+ pass
311
+
312
+
313
+ @operator(resolver=lambda f: f.eval_degree)
314
+ def degree(f: Field):
315
+ """Polynomial degree of a field"""
316
+ pass
317
+
318
+
319
+ @operator(resolver=lambda f: f.node_count)
320
+ def node_count(f: Field, s: Sample):
321
+ """
322
+ Returns the number of nodes associated to the field `f` in the element containing the sample `s`
323
+ """
324
+ pass
325
+
326
+
327
+ @operator(resolver=lambda f: f.at_node)
328
+ def at_node(f: Field, s: Sample, node_index_in_elt: Optional[int] = None):
329
+ """
330
+ Returns a copy of the Sample `s` moved to the coordinates of a local node of the field `f`.
331
+
332
+ If `f` is a discrete field, `node_index_in_elt` is required and indicates the element-local index of the node to consider.
333
+ If `f` is a Test or Trial field, `node_index_in_elt` **must not** be provided, and will be automatically set
334
+ to the test (resp. trial) node currently being evaluated.
335
+ """
336
+ pass
337
+
338
+
339
+ @operator(resolver=lambda f: f.node_index)
340
+ def node_index(f: Field, s: Sample, node_index_in_elt: Optional[int] = None):
341
+ """Returns the index in the function space of a local node the field `f`.
342
+
343
+ If `f` is a discrete field, `node_index_in_elt` is required and indicates the element-local index of the node to consider.
344
+ If `f` is a Test or Trial field, `node_index_in_elt` **must not** be provided, and will be automatically set
345
+ to the test (resp. trial) node currently being evaluated.
346
+ """
347
+ pass
348
+
349
+
350
+ @operator(resolver=lambda f: f.node_partition_index)
351
+ def node_partition_index(f: Field, node_index: NodeIndex):
352
+ """For a NodalField `f`, returns the index of a given node in the fields's space partition,
353
+ or ``NULL_NODE_INDEX`` if it does not exists"""
354
+ pass
355
+
356
+
357
+ # Common derived operators, for convenience
358
+
359
+
360
+ @integrand
361
+ def D(f: Field, s: Sample):
362
+ """Symmetric part of the (inner) gradient of the field at `s`"""
363
+ return symmetric_part(grad(f, s))
364
+
365
+
366
+ @integrand
367
+ def curl(f: Field, s: Sample):
368
+ """Skew part of the (inner) gradient of the field at `s`, as a vector such that ``wp.cross(curl(u), v) = skew(grad(u)) v``"""
369
+ return skew_part(grad(f, s))
370
+
371
+
372
+ @integrand
373
+ def jump(f: Field, s: Sample):
374
+ """Jump between inner and outer element values on an interior side. Zero for interior points or domain boundaries"""
375
+ return inner(f, s) - outer(f, s)
376
+
377
+
378
+ @integrand
379
+ def average(f: Field, s: Sample):
380
+ """Average between inner and outer element values"""
381
+ return 0.5 * (inner(f, s) + outer(f, s))
382
+
383
+
384
+ @integrand
385
+ def grad_jump(f: Field, s: Sample):
386
+ """Jump between inner and outer element gradients on an interior side. Zero for interior points or domain boundaries"""
387
+ return grad(f, s) - grad_outer(f, s)
388
+
389
+
390
+ @integrand
391
+ def grad_average(f: Field, s: Sample):
392
+ """Average between inner and outer element gradients"""
393
+ return 0.5 * (grad(f, s) + grad_outer(f, s))
394
+
395
+
396
+ # Set default call operators for argument types, so that field(s) = inner(field, s) and domain(s) = position(domain, s)
397
+ Field.call_operator = inner
398
+ Domain.call_operator = position
@@ -0,0 +1,231 @@
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 math
17
+ from enum import Enum
18
+
19
+ import numpy as np
20
+
21
+ _wp_module_name_ = "warp.fem.polynomial"
22
+
23
+
24
+ class Polynomial(Enum):
25
+ """Polynomial family defining interpolation nodes over an interval"""
26
+
27
+ GAUSS_LEGENDRE = "GL"
28
+ """Gauss--Legendre 1D polynomial family (does not include endpoints)"""
29
+
30
+ LOBATTO_GAUSS_LEGENDRE = "LGL"
31
+ """Lobatto--Gauss--Legendre 1D polynomial family (includes endpoints)"""
32
+
33
+ EQUISPACED_CLOSED = "closed"
34
+ """Closed 1D polynomial family with uniformly distributed nodes (includes endpoints)"""
35
+
36
+ EQUISPACED_OPEN = "open"
37
+ """Open 1D polynomial family with uniformly distributed nodes (does not include endpoints)"""
38
+
39
+ def __str__(self):
40
+ return self.value
41
+
42
+
43
+ def is_closed(family: Polynomial):
44
+ """Whether the polynomial roots include interval endpoints"""
45
+ return family == Polynomial.LOBATTO_GAUSS_LEGENDRE or family == Polynomial.EQUISPACED_CLOSED
46
+
47
+
48
+ def _gauss_legendre_quadrature_1d(n: int):
49
+ if n == 1:
50
+ coords = [0.0]
51
+ weights = [2.0]
52
+ elif n == 2:
53
+ coords = [-math.sqrt(1.0 / 3), math.sqrt(1.0 / 3)]
54
+ weights = [1.0, 1.0]
55
+ elif n == 3:
56
+ coords = [0.0, -math.sqrt(3.0 / 5.0), math.sqrt(3.0 / 5.0)]
57
+ weights = [8.0 / 9.0, 5.0 / 9.0, 5.0 / 9.0]
58
+ elif n == 4:
59
+ c_a = math.sqrt(3.0 / 7.0 - 2.0 / 7.0 * math.sqrt(6.0 / 5.0))
60
+ c_b = math.sqrt(3.0 / 7.0 + 2.0 / 7.0 * math.sqrt(6.0 / 5.0))
61
+ w_a = (18.0 + math.sqrt(30.0)) / 36.0
62
+ w_b = (18.0 - math.sqrt(30.0)) / 36.0
63
+ coords = [c_a, -c_a, c_b, -c_b]
64
+ weights = [w_a, w_a, w_b, w_b]
65
+ elif n == 5:
66
+ c_a = 1.0 / 3.0 * math.sqrt(5.0 - 2.0 * math.sqrt(10.0 / 7.0))
67
+ c_b = 1.0 / 3.0 * math.sqrt(5.0 + 2.0 * math.sqrt(10.0 / 7.0))
68
+ w_a = (322.0 + 13.0 * math.sqrt(70.0)) / 900.0
69
+ w_b = (322.0 - 13.0 * math.sqrt(70.0)) / 900.0
70
+ coords = [0.0, c_a, -c_a, c_b, -c_b]
71
+ weights = [128.0 / 225.0, w_a, w_a, w_b, w_b]
72
+ else:
73
+ raise NotImplementedError
74
+
75
+ # Shift from [-1, 1] to [0, 1]
76
+ weights = 0.5 * np.array(weights)
77
+ coords = 0.5 * np.array(coords) + 0.5
78
+
79
+ return coords, weights
80
+
81
+
82
+ def _lobatto_gauss_legendre_quadrature_1d(n: int):
83
+ if n == 2:
84
+ coords = [-1.0, 1.0]
85
+ weights = [1.0, 1.0]
86
+ elif n == 3:
87
+ coords = [-1.0, 0.0, 1.0]
88
+ weights = [1.0 / 3.0, 4.0 / 3.0, 1.0 / 3.0]
89
+ elif n == 4:
90
+ coords = [-1.0, -1.0 / math.sqrt(5.0), 1.0 / math.sqrt(5.0), 1.0]
91
+ weights = [1.0 / 6.0, 5.0 / 6.0, 5.0 / 6.0, 1.0 / 6.0]
92
+ elif n == 5:
93
+ coords = [-1.0, -math.sqrt(3.0 / 7.0), 0.0, math.sqrt(3.0 / 7.0), 1.0]
94
+ weights = [1.0 / 10.0, 49.0 / 90.0, 32.0 / 45.0, 49.0 / 90.0, 1.0 / 10.0]
95
+ else:
96
+ raise NotImplementedError
97
+
98
+ # Shift from [-1, 1] to [0, 1]
99
+ weights = 0.5 * np.array(weights)
100
+ coords = 0.5 * np.array(coords) + 0.5
101
+
102
+ return coords, weights
103
+
104
+
105
+ def _uniform_open_quadrature_1d(n: int):
106
+ step = 1.0 / (n + 1)
107
+ coords = np.linspace(step, 1.0 - step, n)
108
+ weights = np.full(n, 1.0 / (n + 1))
109
+
110
+ # Boundaries have 3/2 the weight
111
+ weights[0] = 1.5 / (n + 1)
112
+ weights[-1] = 1.5 / (n + 1)
113
+
114
+ return coords, weights
115
+
116
+
117
+ def _uniform_closed_quadrature_1d(n: int):
118
+ coords = np.linspace(0.0, 1.0, n)
119
+ weights = np.full(n, 1.0 / (n - 1))
120
+
121
+ # Boundaries have half the weight
122
+ weights[0] = 0.5 / (n - 1)
123
+ weights[-1] = 0.5 / (n - 1)
124
+
125
+ return coords, weights
126
+
127
+
128
+ def _open_newton_cotes_quadrature_1d(n: int):
129
+ step = 1.0 / (n + 1)
130
+ coords = np.linspace(step, 1.0 - step, n)
131
+
132
+ # Weisstein, Eric W. "Newton-Cotes Formulas." From MathWorld--A Wolfram Web Resource.
133
+ # https://mathworld.wolfram.com/Newton-CotesFormulas.html
134
+
135
+ if n == 1:
136
+ weights = np.array([1.0])
137
+ elif n == 2:
138
+ weights = np.array([0.5, 0.5])
139
+ elif n == 3:
140
+ weights = np.array([2.0, -1.0, 2.0]) / 3.0
141
+ elif n == 4:
142
+ weights = np.array([11.0, 1.0, 1.0, 11.0]) / 24.0
143
+ elif n == 5:
144
+ weights = np.array([11.0, -14.0, 26.0, -14.0, 11.0]) / 20.0
145
+ elif n == 6:
146
+ weights = np.array([611.0, -453.0, 562.0, 562.0, -453.0, 611.0]) / 1440.0
147
+ elif n == 7:
148
+ weights = np.array([460.0, -954.0, 2196.0, -2459.0, 2196.0, -954.0, 460.0]) / 945.0
149
+ else:
150
+ raise NotImplementedError
151
+
152
+ return coords, weights
153
+
154
+
155
+ def _closed_newton_cotes_quadrature_1d(n: int):
156
+ coords = np.linspace(0.0, 1.0, n)
157
+
158
+ # OEIS: A093735, A093736
159
+
160
+ if n == 2:
161
+ weights = np.array([1.0, 1.0]) / 2.0
162
+ elif n == 3:
163
+ weights = np.array([1.0, 4.0, 1.0]) / 3.0
164
+ elif n == 4:
165
+ weights = np.array([3.0, 9.0, 9.0, 3.0]) / 8.0
166
+ elif n == 5:
167
+ weights = np.array([14.0, 64.0, 24.0, 64.0, 14.0]) / 45.0
168
+ elif n == 6:
169
+ weights = np.array([95.0 / 288.0, 125.0 / 96.0, 125.0 / 144.0, 125.0 / 144.0, 125.0 / 96.0, 95.0 / 288.0])
170
+ elif n == 7:
171
+ weights = np.array([41, 54, 27, 68, 27, 54, 41], dtype=float) / np.array(
172
+ [140, 35, 140, 35, 140, 35, 140], dtype=float
173
+ )
174
+ elif n == 8:
175
+ weights = np.array(
176
+ [
177
+ 5257,
178
+ 25039,
179
+ 343,
180
+ 20923,
181
+ 20923,
182
+ 343,
183
+ 25039,
184
+ 5257,
185
+ ]
186
+ ) / np.array(
187
+ [
188
+ 17280,
189
+ 17280,
190
+ 640,
191
+ 17280,
192
+ 17280,
193
+ 640,
194
+ 17280,
195
+ 17280,
196
+ ],
197
+ dtype=float,
198
+ )
199
+ else:
200
+ raise NotImplementedError
201
+
202
+ # Normalize with interval length
203
+ weights = weights / (n - 1)
204
+
205
+ return coords, weights
206
+
207
+
208
+ def quadrature_1d(point_count: int, family: Polynomial):
209
+ """Return quadrature points and weights for the given family and point count"""
210
+
211
+ if family == Polynomial.GAUSS_LEGENDRE:
212
+ return _gauss_legendre_quadrature_1d(point_count)
213
+ if family == Polynomial.LOBATTO_GAUSS_LEGENDRE:
214
+ return _lobatto_gauss_legendre_quadrature_1d(point_count)
215
+ if family == Polynomial.EQUISPACED_CLOSED:
216
+ return _closed_newton_cotes_quadrature_1d(point_count)
217
+ if family == Polynomial.EQUISPACED_OPEN:
218
+ return _open_newton_cotes_quadrature_1d(point_count)
219
+
220
+ raise NotImplementedError
221
+
222
+
223
+ def lagrange_scales(coords: np.array):
224
+ """Return the scaling factors for Lagrange polynomials with roots at coords"""
225
+ lagrange_scale = np.empty_like(coords)
226
+ for i in range(len(coords)):
227
+ deltas = coords[i] - coords
228
+ deltas[i] = 1.0
229
+ lagrange_scale[i] = 1.0 / np.prod(deltas)
230
+
231
+ return lagrange_scale
@@ -0,0 +1,17 @@
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
+ from .pic_quadrature import PicQuadrature
17
+ from .quadrature import ExplicitQuadrature, NodalQuadrature, Quadrature, RegularQuadrature