warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.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 (300) hide show
  1. warp/__init__.py +10 -4
  2. warp/__init__.pyi +1 -0
  3. warp/bin/warp-clang.dll +0 -0
  4. warp/bin/warp.dll +0 -0
  5. warp/build.py +5 -3
  6. warp/build_dll.py +29 -9
  7. warp/builtins.py +868 -507
  8. warp/codegen.py +1074 -638
  9. warp/config.py +3 -3
  10. warp/constants.py +6 -0
  11. warp/context.py +715 -222
  12. warp/fabric.py +326 -0
  13. warp/fem/__init__.py +27 -0
  14. warp/fem/cache.py +389 -0
  15. warp/fem/dirichlet.py +181 -0
  16. warp/fem/domain.py +263 -0
  17. warp/fem/field/__init__.py +101 -0
  18. warp/fem/field/field.py +149 -0
  19. warp/fem/field/nodal_field.py +299 -0
  20. warp/fem/field/restriction.py +21 -0
  21. warp/fem/field/test.py +181 -0
  22. warp/fem/field/trial.py +183 -0
  23. warp/fem/geometry/__init__.py +19 -0
  24. warp/fem/geometry/closest_point.py +70 -0
  25. warp/fem/geometry/deformed_geometry.py +271 -0
  26. warp/fem/geometry/element.py +744 -0
  27. warp/fem/geometry/geometry.py +186 -0
  28. warp/fem/geometry/grid_2d.py +373 -0
  29. warp/fem/geometry/grid_3d.py +435 -0
  30. warp/fem/geometry/hexmesh.py +953 -0
  31. warp/fem/geometry/partition.py +376 -0
  32. warp/fem/geometry/quadmesh_2d.py +532 -0
  33. warp/fem/geometry/tetmesh.py +840 -0
  34. warp/fem/geometry/trimesh_2d.py +577 -0
  35. warp/fem/integrate.py +1616 -0
  36. warp/fem/operator.py +191 -0
  37. warp/fem/polynomial.py +213 -0
  38. warp/fem/quadrature/__init__.py +2 -0
  39. warp/fem/quadrature/pic_quadrature.py +245 -0
  40. warp/fem/quadrature/quadrature.py +294 -0
  41. warp/fem/space/__init__.py +292 -0
  42. warp/fem/space/basis_space.py +489 -0
  43. warp/fem/space/collocated_function_space.py +105 -0
  44. warp/fem/space/dof_mapper.py +236 -0
  45. warp/fem/space/function_space.py +145 -0
  46. warp/fem/space/grid_2d_function_space.py +267 -0
  47. warp/fem/space/grid_3d_function_space.py +306 -0
  48. warp/fem/space/hexmesh_function_space.py +352 -0
  49. warp/fem/space/partition.py +350 -0
  50. warp/fem/space/quadmesh_2d_function_space.py +369 -0
  51. warp/fem/space/restriction.py +160 -0
  52. warp/fem/space/shape/__init__.py +15 -0
  53. warp/fem/space/shape/cube_shape_function.py +738 -0
  54. warp/fem/space/shape/shape_function.py +103 -0
  55. warp/fem/space/shape/square_shape_function.py +611 -0
  56. warp/fem/space/shape/tet_shape_function.py +567 -0
  57. warp/fem/space/shape/triangle_shape_function.py +429 -0
  58. warp/fem/space/tetmesh_function_space.py +292 -0
  59. warp/fem/space/topology.py +295 -0
  60. warp/fem/space/trimesh_2d_function_space.py +221 -0
  61. warp/fem/types.py +77 -0
  62. warp/fem/utils.py +495 -0
  63. warp/native/array.h +147 -44
  64. warp/native/builtin.h +122 -149
  65. warp/native/bvh.cpp +73 -325
  66. warp/native/bvh.cu +406 -23
  67. warp/native/bvh.h +34 -43
  68. warp/native/clang/clang.cpp +13 -8
  69. warp/native/crt.h +2 -0
  70. warp/native/cuda_crt.h +5 -0
  71. warp/native/cuda_util.cpp +15 -3
  72. warp/native/cuda_util.h +3 -1
  73. warp/native/cutlass/tools/library/scripts/conv2d_operation.py +463 -0
  74. warp/native/cutlass/tools/library/scripts/conv3d_operation.py +321 -0
  75. warp/native/cutlass/tools/library/scripts/gemm_operation.py +988 -0
  76. warp/native/cutlass/tools/library/scripts/generator.py +4625 -0
  77. warp/native/cutlass/tools/library/scripts/library.py +799 -0
  78. warp/native/cutlass/tools/library/scripts/manifest.py +402 -0
  79. warp/native/cutlass/tools/library/scripts/pycutlass/docs/source/conf.py +96 -0
  80. warp/native/cutlass/tools/library/scripts/pycutlass/profile/conv/conv2d_f16_sm80.py +106 -0
  81. warp/native/cutlass/tools/library/scripts/pycutlass/profile/gemm/gemm_f32_sm80.py +91 -0
  82. warp/native/cutlass/tools/library/scripts/pycutlass/setup.py +80 -0
  83. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/__init__.py +48 -0
  84. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/arguments.py +118 -0
  85. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/c_types.py +241 -0
  86. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/compiler.py +432 -0
  87. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/conv2d_operation.py +631 -0
  88. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/epilogue.py +1026 -0
  89. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/frontend.py +104 -0
  90. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/gemm_operation.py +1276 -0
  91. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/library.py +744 -0
  92. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/memory_manager.py +74 -0
  93. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/operation.py +110 -0
  94. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/parser.py +619 -0
  95. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/reduction_operation.py +398 -0
  96. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/tensor_ref.py +70 -0
  97. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/__init__.py +4 -0
  98. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/conv2d_testbed.py +646 -0
  99. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/gemm_grouped_testbed.py +235 -0
  100. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/gemm_testbed.py +557 -0
  101. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/profiler.py +70 -0
  102. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/type_hint.py +39 -0
  103. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/utils/__init__.py +1 -0
  104. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/utils/device.py +76 -0
  105. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/utils/reference_model.py +255 -0
  106. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/__init__.py +0 -0
  107. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py +201 -0
  108. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +177 -0
  109. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py +98 -0
  110. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py +95 -0
  111. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_few_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.py +163 -0
  112. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_fixed_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.py +187 -0
  113. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py +309 -0
  114. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +54 -0
  115. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py +96 -0
  116. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py +107 -0
  117. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_strided_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +253 -0
  118. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py +97 -0
  119. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +242 -0
  120. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py +96 -0
  121. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py +107 -0
  122. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/run_all_tests.py +10 -0
  123. warp/native/cutlass/tools/library/scripts/pycutlass/test/frontend/test_frontend.py +146 -0
  124. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/__init__.py +0 -0
  125. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_bf16_sm80.py +96 -0
  126. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_f16_sm80.py +447 -0
  127. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_f32_sm80.py +146 -0
  128. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_f64_sm80.py +102 -0
  129. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_grouped_sm80.py +203 -0
  130. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_s8_sm80.py +229 -0
  131. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/run_all_tests.py +9 -0
  132. warp/native/cutlass/tools/library/scripts/pycutlass/test/unit/test_sm80.py +453 -0
  133. warp/native/cutlass/tools/library/scripts/rank_2k_operation.py +398 -0
  134. warp/native/cutlass/tools/library/scripts/rank_k_operation.py +387 -0
  135. warp/native/cutlass/tools/library/scripts/rt.py +796 -0
  136. warp/native/cutlass/tools/library/scripts/symm_operation.py +400 -0
  137. warp/native/cutlass/tools/library/scripts/trmm_operation.py +407 -0
  138. warp/native/cutlass_gemm.cu +5 -3
  139. warp/native/exports.h +1240 -952
  140. warp/native/fabric.h +228 -0
  141. warp/native/hashgrid.cpp +4 -4
  142. warp/native/hashgrid.h +22 -2
  143. warp/native/intersect.h +22 -7
  144. warp/native/intersect_adj.h +8 -8
  145. warp/native/intersect_tri.h +1 -1
  146. warp/native/marching.cu +157 -161
  147. warp/native/mat.h +80 -19
  148. warp/native/matnn.h +2 -2
  149. warp/native/mesh.cpp +33 -108
  150. warp/native/mesh.cu +114 -23
  151. warp/native/mesh.h +446 -46
  152. warp/native/noise.h +272 -329
  153. warp/native/quat.h +51 -8
  154. warp/native/rand.h +45 -35
  155. warp/native/range.h +6 -2
  156. warp/native/reduce.cpp +1 -1
  157. warp/native/reduce.cu +10 -12
  158. warp/native/runlength_encode.cu +6 -10
  159. warp/native/scan.cu +8 -11
  160. warp/native/sparse.cpp +4 -4
  161. warp/native/sparse.cu +164 -154
  162. warp/native/spatial.h +2 -2
  163. warp/native/temp_buffer.h +14 -30
  164. warp/native/vec.h +107 -23
  165. warp/native/volume.h +120 -0
  166. warp/native/warp.cpp +560 -30
  167. warp/native/warp.cu +431 -44
  168. warp/native/warp.h +13 -4
  169. warp/optim/__init__.py +1 -0
  170. warp/optim/linear.py +922 -0
  171. warp/optim/sgd.py +92 -0
  172. warp/render/render_opengl.py +335 -119
  173. warp/render/render_usd.py +11 -11
  174. warp/sim/__init__.py +2 -2
  175. warp/sim/articulation.py +385 -185
  176. warp/sim/collide.py +8 -0
  177. warp/sim/import_mjcf.py +297 -106
  178. warp/sim/import_urdf.py +389 -210
  179. warp/sim/import_usd.py +198 -97
  180. warp/sim/inertia.py +17 -18
  181. warp/sim/integrator_euler.py +14 -8
  182. warp/sim/integrator_xpbd.py +158 -16
  183. warp/sim/model.py +795 -291
  184. warp/sim/render.py +3 -3
  185. warp/sim/utils.py +3 -0
  186. warp/sparse.py +640 -150
  187. warp/stubs.py +606 -267
  188. warp/tape.py +61 -10
  189. warp/tests/__main__.py +3 -6
  190. warp/tests/assets/curlnoise_golden.npy +0 -0
  191. warp/tests/assets/pnoise_golden.npy +0 -0
  192. warp/tests/{test_class_kernel.py → aux_test_class_kernel.py} +9 -1
  193. warp/tests/aux_test_conditional_unequal_types_kernels.py +21 -0
  194. warp/tests/{test_dependent.py → aux_test_dependent.py} +2 -2
  195. warp/tests/{test_reference.py → aux_test_reference.py} +1 -1
  196. warp/tests/aux_test_unresolved_func.py +14 -0
  197. warp/tests/aux_test_unresolved_symbol.py +14 -0
  198. warp/tests/disabled_kinematics.py +239 -0
  199. warp/tests/run_coverage_serial.py +31 -0
  200. warp/tests/test_adam.py +103 -106
  201. warp/tests/test_arithmetic.py +128 -74
  202. warp/tests/test_array.py +212 -97
  203. warp/tests/test_array_reduce.py +57 -23
  204. warp/tests/test_atomic.py +64 -28
  205. warp/tests/test_bool.py +99 -0
  206. warp/tests/test_builtins_resolution.py +1292 -0
  207. warp/tests/test_bvh.py +42 -18
  208. warp/tests/test_closest_point_edge_edge.py +54 -57
  209. warp/tests/test_codegen.py +208 -130
  210. warp/tests/test_compile_consts.py +28 -20
  211. warp/tests/test_conditional.py +108 -24
  212. warp/tests/test_copy.py +10 -12
  213. warp/tests/test_ctypes.py +112 -88
  214. warp/tests/test_dense.py +21 -14
  215. warp/tests/test_devices.py +98 -0
  216. warp/tests/test_dlpack.py +75 -75
  217. warp/tests/test_examples.py +277 -0
  218. warp/tests/test_fabricarray.py +955 -0
  219. warp/tests/test_fast_math.py +15 -11
  220. warp/tests/test_fem.py +1271 -0
  221. warp/tests/test_fp16.py +53 -19
  222. warp/tests/test_func.py +187 -86
  223. warp/tests/test_generics.py +194 -49
  224. warp/tests/test_grad.py +178 -109
  225. warp/tests/test_grad_customs.py +176 -0
  226. warp/tests/test_hash_grid.py +52 -37
  227. warp/tests/test_import.py +10 -23
  228. warp/tests/test_indexedarray.py +32 -31
  229. warp/tests/test_intersect.py +18 -9
  230. warp/tests/test_large.py +141 -0
  231. warp/tests/test_launch.py +14 -41
  232. warp/tests/test_lerp.py +64 -65
  233. warp/tests/test_linear_solvers.py +154 -0
  234. warp/tests/test_lvalue.py +493 -0
  235. warp/tests/test_marching_cubes.py +12 -13
  236. warp/tests/test_mat.py +517 -2898
  237. warp/tests/test_mat_lite.py +115 -0
  238. warp/tests/test_mat_scalar_ops.py +2889 -0
  239. warp/tests/test_math.py +103 -9
  240. warp/tests/test_matmul.py +305 -69
  241. warp/tests/test_matmul_lite.py +410 -0
  242. warp/tests/test_mesh.py +71 -14
  243. warp/tests/test_mesh_query_aabb.py +41 -25
  244. warp/tests/test_mesh_query_point.py +140 -22
  245. warp/tests/test_mesh_query_ray.py +39 -22
  246. warp/tests/test_mlp.py +30 -22
  247. warp/tests/test_model.py +92 -89
  248. warp/tests/test_modules_lite.py +39 -0
  249. warp/tests/test_multigpu.py +88 -114
  250. warp/tests/test_noise.py +12 -11
  251. warp/tests/test_operators.py +16 -20
  252. warp/tests/test_options.py +11 -11
  253. warp/tests/test_pinned.py +17 -18
  254. warp/tests/test_print.py +32 -11
  255. warp/tests/test_quat.py +275 -129
  256. warp/tests/test_rand.py +18 -16
  257. warp/tests/test_reload.py +38 -34
  258. warp/tests/test_rounding.py +50 -43
  259. warp/tests/test_runlength_encode.py +168 -20
  260. warp/tests/test_smoothstep.py +9 -11
  261. warp/tests/test_snippet.py +143 -0
  262. warp/tests/test_sparse.py +261 -63
  263. warp/tests/test_spatial.py +276 -243
  264. warp/tests/test_streams.py +110 -85
  265. warp/tests/test_struct.py +268 -63
  266. warp/tests/test_tape.py +39 -21
  267. warp/tests/test_torch.py +118 -89
  268. warp/tests/test_transient_module.py +12 -13
  269. warp/tests/test_types.py +614 -0
  270. warp/tests/test_utils.py +494 -0
  271. warp/tests/test_vec.py +354 -2050
  272. warp/tests/test_vec_lite.py +73 -0
  273. warp/tests/test_vec_scalar_ops.py +2099 -0
  274. warp/tests/test_volume.py +457 -293
  275. warp/tests/test_volume_write.py +124 -134
  276. warp/tests/unittest_serial.py +35 -0
  277. warp/tests/unittest_suites.py +341 -0
  278. warp/tests/unittest_utils.py +568 -0
  279. warp/tests/unused_test_misc.py +71 -0
  280. warp/tests/{test_debug.py → walkthough_debug.py} +3 -17
  281. warp/thirdparty/appdirs.py +36 -45
  282. warp/thirdparty/unittest_parallel.py +549 -0
  283. warp/torch.py +9 -6
  284. warp/types.py +1089 -366
  285. warp/utils.py +93 -387
  286. warp_lang-0.11.0.dist-info/METADATA +238 -0
  287. warp_lang-0.11.0.dist-info/RECORD +332 -0
  288. {warp_lang-0.10.1.dist-info → warp_lang-0.11.0.dist-info}/WHEEL +1 -1
  289. warp/tests/test_all.py +0 -219
  290. warp/tests/test_array_scan.py +0 -60
  291. warp/tests/test_base.py +0 -208
  292. warp/tests/test_unresolved_func.py +0 -7
  293. warp/tests/test_unresolved_symbol.py +0 -7
  294. warp_lang-0.10.1.dist-info/METADATA +0 -21
  295. warp_lang-0.10.1.dist-info/RECORD +0 -188
  296. /warp/tests/{test_compile_consts_dummy.py → aux_test_compile_consts_dummy.py} +0 -0
  297. /warp/tests/{test_reference_reference.py → aux_test_reference_reference.py} +0 -0
  298. /warp/tests/{test_square.py → aux_test_square.py} +0 -0
  299. {warp_lang-0.10.1.dist-info → warp_lang-0.11.0.dist-info}/LICENSE.md +0 -0
  300. {warp_lang-0.10.1.dist-info → warp_lang-0.11.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,299 @@
1
+ import warp as wp
2
+
3
+ from warp.fem.space import CollocatedFunctionSpace, SpacePartition
4
+ from warp.fem import cache, utils
5
+ from warp.fem.types import Sample, ElementIndex, NULL_NODE_INDEX
6
+
7
+ from .field import DiscreteField
8
+
9
+
10
+ class NodalFieldBase(DiscreteField):
11
+ """Base class for nodal field and nodal field traces. Does not hold values"""
12
+
13
+ def __init__(self, space: CollocatedFunctionSpace, space_partition: SpacePartition):
14
+ super().__init__(space, space_partition)
15
+
16
+ self.EvalArg = self._make_eval_arg()
17
+ self.ElementEvalArg = self._make_element_eval_arg()
18
+ self.eval_degree = DiscreteField._make_eval_degree(self)
19
+
20
+ self._read_node_value = self._make_read_node_value()
21
+
22
+ self.eval_inner = self._make_eval_inner()
23
+ self.eval_outer = self._make_eval_outer()
24
+ self.eval_grad_inner = self._make_eval_grad_inner(world_space=True)
25
+ self.eval_grad_outer = self._make_eval_grad_outer(world_space=True)
26
+ self.eval_reference_grad_inner = self._make_eval_grad_inner(world_space=False)
27
+ self.eval_reference_grad_outer = self._make_eval_grad_outer(world_space=False)
28
+ self.eval_div_inner = self._make_eval_div_inner()
29
+ self.eval_div_outer = self._make_eval_div_outer()
30
+
31
+ self.set_node_value = self._make_set_node_value()
32
+
33
+ def _make_eval_arg(self):
34
+ @cache.dynamic_struct(suffix=self.name)
35
+ class EvalArg:
36
+ dof_values: wp.array(dtype=self.space.dof_dtype)
37
+ space_arg: self.space.SpaceArg
38
+ topology_arg: self.space.topology.TopologyArg
39
+ partition_arg: self.space_partition.PartitionArg
40
+
41
+ return EvalArg
42
+
43
+ def _make_element_eval_arg(self):
44
+ @cache.dynamic_struct(suffix=self.name)
45
+ class ElementEvalArg:
46
+ elt_arg: self.space.topology.ElementArg
47
+ eval_arg: self.EvalArg
48
+
49
+ return ElementEvalArg
50
+
51
+ def _make_read_node_value(self):
52
+ @cache.dynamic_func(suffix=self.name)
53
+ def read_node_value(args: self.ElementEvalArg, geo_element_index: ElementIndex, node_index_in_elt: int):
54
+ nidx = self.space.topology.element_node_index(
55
+ args.elt_arg, args.eval_arg.topology_arg, geo_element_index, node_index_in_elt
56
+ )
57
+ pidx = self.space_partition.partition_node_index(args.eval_arg.partition_arg, nidx)
58
+ if pidx == NULL_NODE_INDEX:
59
+ return self.space.dtype(0.0)
60
+
61
+ return self.space.dof_mapper.dof_to_value(args.eval_arg.dof_values[pidx])
62
+
63
+ return read_node_value
64
+
65
+ def _make_eval_inner(self):
66
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
67
+
68
+ @cache.dynamic_func(suffix=self.name)
69
+ def eval_inner(args: self.ElementEvalArg, s: Sample):
70
+ res = self.space.element_inner_weight(
71
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
72
+ ) * self._read_node_value(args, s.element_index, 0)
73
+ for k in range(1, NODES_PER_ELEMENT):
74
+ res += self.space.element_inner_weight(
75
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
76
+ ) * self._read_node_value(args, s.element_index, k)
77
+ return res
78
+
79
+ return eval_inner
80
+
81
+ def _make_eval_grad_inner(self, world_space: bool):
82
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
83
+
84
+ if not self.gradient_valid():
85
+ return None
86
+
87
+ @cache.dynamic_func(suffix=self.name)
88
+ def eval_grad_inner_ref_space(args: self.ElementEvalArg, s: Sample):
89
+ res = utils.generalized_outer(
90
+ self._read_node_value(args, s.element_index, 0),
91
+ self.space.element_inner_weight_gradient(
92
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
93
+ ),
94
+ )
95
+ for k in range(1, NODES_PER_ELEMENT):
96
+ res += utils.generalized_outer(
97
+ self._read_node_value(args, s.element_index, k),
98
+ self.space.element_inner_weight_gradient(
99
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
100
+ ),
101
+ )
102
+ return res
103
+
104
+ @cache.dynamic_func(suffix=self.name)
105
+ def eval_grad_inner_world_space(args: self.ElementEvalArg, s: Sample):
106
+ grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
107
+ res = eval_grad_inner_ref_space(args, s)
108
+ return utils.apply_right(res, grad_transform)
109
+
110
+ return eval_grad_inner_world_space if world_space else eval_grad_inner_ref_space
111
+
112
+ def _make_eval_div_inner(self):
113
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
114
+
115
+ if not self.divergence_valid():
116
+ return None
117
+
118
+ @cache.dynamic_func(suffix=self.name)
119
+ def eval_div_inner(args: self.ElementEvalArg, s: Sample):
120
+ grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
121
+
122
+ res = utils.generalized_inner(
123
+ self._read_node_value(args, s.element_index, 0),
124
+ utils.apply_right(
125
+ self.space.element_inner_weight_gradient(
126
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
127
+ ),
128
+ grad_transform,
129
+ ),
130
+ )
131
+
132
+ for k in range(1, NODES_PER_ELEMENT):
133
+ res += utils.generalized_inner(
134
+ self._read_node_value(args, s.element_index, k),
135
+ utils.apply_right(
136
+ self.space.element_inner_weight_gradient(
137
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
138
+ ),
139
+ grad_transform,
140
+ ),
141
+ )
142
+ return res
143
+
144
+ return eval_div_inner
145
+
146
+ def _make_eval_outer(self):
147
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
148
+
149
+ @cache.dynamic_func(suffix=self.name)
150
+ def eval_outer(args: self.ElementEvalArg, s: Sample):
151
+ res = self.space.element_outer_weight(
152
+ args.elt_arg,
153
+ args.eval_arg.space_arg,
154
+ s.element_index,
155
+ s.element_coords,
156
+ 0,
157
+ ) * self._read_node_value(args, s.element_index, 0)
158
+ for k in range(1, NODES_PER_ELEMENT):
159
+ res += self.space.element_outer_weight(
160
+ args.elt_arg,
161
+ args.eval_arg.space_arg,
162
+ s.element_index,
163
+ s.element_coords,
164
+ k,
165
+ ) * self._read_node_value(args, s.element_index, k)
166
+ return res
167
+
168
+ return eval_outer
169
+
170
+ def _make_eval_grad_outer(self, world_space: bool):
171
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
172
+
173
+ if not self.gradient_valid():
174
+ return None
175
+
176
+ @cache.dynamic_func(suffix=self.name)
177
+ def eval_grad_outer_ref_space(args: self.ElementEvalArg, s: Sample):
178
+ res = utils.generalized_outer(
179
+ self._read_node_value(args, s.element_index, 0),
180
+ self.space.element_outer_weight_gradient(
181
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
182
+ ),
183
+ )
184
+ for k in range(1, NODES_PER_ELEMENT):
185
+ res += utils.generalized_outer(
186
+ self._read_node_value(args, s.element_index, k),
187
+ self.space.element_outer_weight_gradient(
188
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
189
+ ),
190
+ )
191
+ return res
192
+
193
+ @cache.dynamic_func(suffix=self.name)
194
+ def eval_grad_outer_world_space(args: self.ElementEvalArg, s: Sample):
195
+ grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
196
+ res = eval_grad_outer_ref_space(args, s)
197
+ return utils.apply_right(res, grad_transform)
198
+
199
+ return eval_grad_outer_world_space if world_space else eval_grad_outer_ref_space
200
+
201
+ def _make_eval_div_outer(self):
202
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
203
+
204
+ if not self.divergence_valid():
205
+ return None
206
+
207
+ @cache.dynamic_func(suffix=self.name)
208
+ def eval_div_outer(args: self.ElementEvalArg, s: Sample):
209
+ grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
210
+
211
+ res = utils.generalized_inner(
212
+ self._read_node_value(args, s.element_index, 0),
213
+ utils.apply_right(
214
+ self.space.element_outer_weight_gradient(
215
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
216
+ ),
217
+ grad_transform,
218
+ ),
219
+ )
220
+ for k in range(1, NODES_PER_ELEMENT):
221
+ res += utils.generalized_inner(
222
+ self._read_node_value(args, s.element_index, k),
223
+ utils.apply_right(
224
+ self.space.element_outer_weight_gradient(
225
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
226
+ ),
227
+ grad_transform,
228
+ ),
229
+ )
230
+ return res
231
+
232
+ return eval_div_outer
233
+
234
+ def _make_set_node_value(self):
235
+ @cache.dynamic_func(suffix=self.name)
236
+ def set_node_value(args: self.EvalArg, partition_node_index: int, value: self.space.dtype):
237
+ args.dof_values[partition_node_index] = self.space.dof_mapper.value_to_dof(value)
238
+
239
+ return set_node_value
240
+
241
+
242
+ class NodalField(NodalFieldBase):
243
+ """A field holding values for all degrees of freedom at each node of the underlying function space partition
244
+
245
+ See also: warp.fem.space.CollocatedFunctionSpace.make_field
246
+ """
247
+
248
+ def __init__(self, space: CollocatedFunctionSpace, space_partition: SpacePartition):
249
+ if space.topology != space_partition.space_topology:
250
+ raise ValueError("Incompatible space and space partition topologies")
251
+
252
+ super().__init__(space, space_partition)
253
+
254
+ self._dof_values = wp.zeros(n=self.space_partition.node_count(), dtype=self.dof_dtype)
255
+
256
+ def eval_arg_value(self, device):
257
+ arg = self.EvalArg()
258
+ arg.dof_values = self._dof_values.to(device)
259
+ arg.space_arg = self.space.space_arg_value(device)
260
+ arg.partition_arg = self.space_partition.partition_arg_value(device)
261
+ arg.topology_arg = self.space.topology.topo_arg_value(device)
262
+
263
+ return arg
264
+
265
+ @property
266
+ def dof_values(self) -> wp.array:
267
+ """Returns a warp array containing the values at all degrees of freedom of the underlying space partition"""
268
+ return self._dof_values
269
+
270
+ @dof_values.setter
271
+ def dof_values(self, values):
272
+ """Sets the degrees-of-freedom values
273
+
274
+ Args:
275
+ values: Array that is convertible to a warp array of length ``self.space_partition.node_count()`` and data type ``self.space.dof_dtype``
276
+ """
277
+
278
+ if isinstance(values, wp.array):
279
+ self._dof_values = values
280
+ else:
281
+ self._dof_values = wp.array(values, dtype=self.dof_dtype)
282
+
283
+ class Trace(NodalFieldBase):
284
+ def __init__(self, field):
285
+ self._field = field
286
+ super().__init__(field.space.trace(), field.space_partition)
287
+
288
+ def eval_arg_value(self, device):
289
+ arg = self.EvalArg()
290
+ arg.dof_values = self._field.dof_values.to(device)
291
+ arg.space_arg = self.space.space_arg_value(device)
292
+ arg.partition_arg = self.space_partition.partition_arg_value(device)
293
+ arg.topology_arg = self.space.topology.topo_arg_value(device)
294
+
295
+ return arg
296
+
297
+ def trace(self) -> Trace:
298
+ trace_field = NodalField.Trace(self)
299
+ return trace_field
@@ -0,0 +1,21 @@
1
+ from warp.fem.space import SpaceRestriction
2
+ from .field import DiscreteField
3
+
4
+
5
+ class FieldRestriction:
6
+ """Restriction of a discrete field to a given GeometryDomain"""
7
+
8
+ def __init__(self, space_restriction: SpaceRestriction, field: DiscreteField):
9
+ if field.space.dimension - 1 == space_restriction.space_topology.dimension:
10
+ field = field.trace()
11
+
12
+ if field.space.dimension != space_restriction.space_topology.dimension:
13
+ raise ValueError("Incompatible space and field dimensions")
14
+
15
+ if field.space.topology != space_restriction.space_topology:
16
+ raise ValueError("Incompatible field and space restriction topologies")
17
+
18
+ self.space_restriction = space_restriction
19
+ self.domain = self.space_restriction.domain
20
+ self.field = field
21
+ self.space = self.field.space
warp/fem/field/test.py ADDED
@@ -0,0 +1,181 @@
1
+ import warp as wp
2
+
3
+ from warp.fem.space import SpaceRestriction, FunctionSpace
4
+ from warp.fem.types import Sample, get_node_index_in_element
5
+ from warp.fem import utils, cache
6
+
7
+ from .field import SpaceField
8
+
9
+
10
+ class TestField(SpaceField):
11
+ """Field defined over a space restriction that can be used as a test function.
12
+
13
+ In order to reuse computations, it is possible to define the test field using a SpaceRestriction
14
+ defined for a different value type than the test function value type, as long as the node topology is similar.
15
+ """
16
+
17
+ def __init__(self, space_restriction: SpaceRestriction, space: FunctionSpace):
18
+ if space_restriction.domain.dimension == space.dimension - 1:
19
+ space = space.trace()
20
+
21
+ if space_restriction.domain.dimension != space.dimension:
22
+ raise ValueError("Incompatible space and domain dimensions")
23
+
24
+ if space.topology != space_restriction.space_topology:
25
+ raise ValueError("Incompatible space and space partition topologies")
26
+
27
+ super().__init__(space, space_restriction.space_partition)
28
+
29
+ self.space_restriction = space_restriction
30
+ self.domain = self.space_restriction.domain
31
+
32
+ self.EvalArg = self.space.SpaceArg
33
+ self.ElementEvalArg = self._make_element_eval_arg()
34
+
35
+ self.eval_degree = self._make_eval_degree()
36
+ self.eval_inner = self._make_eval_inner()
37
+ self.eval_grad_inner = self._make_eval_grad_inner()
38
+ self.eval_div_inner = self._make_eval_div_inner()
39
+ self.eval_outer = self._make_eval_outer()
40
+ self.eval_grad_outer = self._make_eval_grad_outer()
41
+ self.eval_div_outer = self._make_eval_div_outer()
42
+ self.at_node = self._make_at_node()
43
+
44
+ @property
45
+ def name(self) -> str:
46
+ return self.space.name + "Test"
47
+
48
+ def eval_arg_value(self, device) -> wp.codegen.StructInstance:
49
+ return self.space.space_arg_value(device)
50
+
51
+ def _make_element_eval_arg(self):
52
+ from warp.fem import cache
53
+
54
+ @cache.dynamic_struct(suffix=self.name)
55
+ class ElementEvalArg:
56
+ elt_arg: self.domain.ElementArg
57
+ eval_arg: self.EvalArg
58
+
59
+ return ElementEvalArg
60
+
61
+ def _make_eval_inner(self):
62
+ @cache.dynamic_func(suffix=self.name)
63
+ def eval_test_inner(args: self.ElementEvalArg, s: Sample):
64
+ weight = self.space.element_inner_weight(
65
+ args.elt_arg,
66
+ args.eval_arg,
67
+ s.element_index,
68
+ s.element_coords,
69
+ get_node_index_in_element(s.test_dof),
70
+ )
71
+ return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof)
72
+
73
+ return eval_test_inner
74
+
75
+ def _make_eval_grad_inner(self):
76
+ if not self.gradient_valid():
77
+ return None
78
+
79
+ @cache.dynamic_func(suffix=self.name)
80
+ def eval_nabla_test_inner(args: self.ElementEvalArg, s: Sample):
81
+ nabla_weight = self.space.element_inner_weight_gradient(
82
+ args.elt_arg,
83
+ args.eval_arg,
84
+ s.element_index,
85
+ s.element_coords,
86
+ get_node_index_in_element(s.test_dof),
87
+ )
88
+ grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
89
+ return utils.generalized_outer(
90
+ self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
91
+ utils.apply_right(nabla_weight, grad_transform),
92
+ )
93
+
94
+ return eval_nabla_test_inner
95
+
96
+ def _make_eval_div_inner(self):
97
+ if not self.divergence_valid():
98
+ return None
99
+
100
+ @cache.dynamic_func(suffix=self.name)
101
+ def eval_div_test_inner(args: self.ElementEvalArg, s: Sample):
102
+ nabla_weight = self.space.element_inner_weight_gradient(
103
+ args.elt_arg,
104
+ args.eval_arg,
105
+ s.element_index,
106
+ s.element_coords,
107
+ get_node_index_in_element(s.test_dof),
108
+ )
109
+ grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
110
+ return utils.generalized_inner(
111
+ self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
112
+ utils.apply_right(nabla_weight, grad_transform),
113
+ )
114
+
115
+ return eval_div_test_inner
116
+
117
+ def _make_eval_outer(self):
118
+ @cache.dynamic_func(suffix=self.name)
119
+ def eval_test_outer(args: self.ElementEvalArg, s: Sample):
120
+ weight = self.space.element_outer_weight(
121
+ args.elt_arg,
122
+ args.eval_arg,
123
+ s.element_index,
124
+ s.element_coords,
125
+ get_node_index_in_element(s.test_dof),
126
+ )
127
+ return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof)
128
+
129
+ return eval_test_outer
130
+
131
+ def _make_eval_grad_outer(self):
132
+ if not self.gradient_valid():
133
+ return None
134
+
135
+ @cache.dynamic_func(suffix=self.name)
136
+ def eval_nabla_test_outer(args: self.ElementEvalArg, s: Sample):
137
+ nabla_weight = self.space.element_outer_weight_gradient(
138
+ args.elt_arg,
139
+ args.eval_arg,
140
+ s.element_index,
141
+ s.element_coords,
142
+ get_node_index_in_element(s.test_dof),
143
+ )
144
+ grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
145
+ return utils.generalized_outer(
146
+ self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
147
+ utils.apply_right(nabla_weight, grad_transform),
148
+ )
149
+
150
+ return eval_nabla_test_outer
151
+
152
+ def _make_eval_div_outer(self):
153
+ if not self.divergence_valid():
154
+ return None
155
+
156
+ @cache.dynamic_func(suffix=self.name)
157
+ def eval_div_test_outer(args: self.ElementEvalArg, s: Sample):
158
+ nabla_weight = self.space.element_outer_weight_gradient(
159
+ args.elt_arg,
160
+ args.eval_arg,
161
+ s.element_index,
162
+ s.element_coords,
163
+ get_node_index_in_element(s.test_dof),
164
+ )
165
+ grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
166
+ return utils.generalized_inner(
167
+ self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
168
+ utils.apply_right(nabla_weight, grad_transform),
169
+ )
170
+
171
+ return eval_div_test_outer
172
+
173
+ def _make_at_node(self):
174
+ @cache.dynamic_func(suffix=self.name)
175
+ def at_node(args: self.ElementEvalArg, s: Sample):
176
+ node_coords = self.space.node_coords_in_element(
177
+ args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(s.test_dof)
178
+ )
179
+ return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
180
+
181
+ return at_node
@@ -0,0 +1,183 @@
1
+ import warp as wp
2
+ from warp.fem.domain import GeometryDomain
3
+ from warp.fem.space import FunctionSpace, SpacePartition
4
+ from warp.fem.types import Sample, get_node_index_in_element
5
+ from warp.fem import utils, cache
6
+
7
+ from .field import SpaceField
8
+
9
+
10
+ class TrialField(SpaceField):
11
+ """Field defined over a domain that can be used as a trial function"""
12
+
13
+ def __init__(
14
+ self,
15
+ space: FunctionSpace,
16
+ space_partition: SpacePartition,
17
+ domain: GeometryDomain,
18
+ ):
19
+ if domain.dimension == space.dimension - 1:
20
+ space = space.trace()
21
+
22
+ if domain.dimension != space.dimension:
23
+ raise ValueError("Incompatible space and domain dimensions")
24
+
25
+ if not space.topology.is_derived_from(space_partition.space_topology):
26
+ raise ValueError("Incompatible space and space partition topologies")
27
+
28
+ super().__init__(space, space_partition)
29
+
30
+ self.domain = domain
31
+
32
+ self.EvalArg = self.space.SpaceArg
33
+ self.ElementEvalArg = self._make_element_eval_arg()
34
+
35
+ self.eval_degree = self._make_eval_degree()
36
+ self.eval_inner = self._make_eval_inner()
37
+ self.eval_grad_inner = self._make_eval_grad_inner()
38
+ self.eval_div_inner = self._make_eval_div_inner()
39
+ self.eval_outer = self._make_eval_outer()
40
+ self.eval_grad_outer = self._make_eval_grad_outer()
41
+ self.eval_div_outer = self._make_eval_div_outer()
42
+ self.at_node = self._make_at_node()
43
+
44
+ def partition_node_count(self) -> int:
45
+ """Returns the number of nodes in the associated space topology partition"""
46
+ return self.space_partition.node_count()
47
+
48
+ @property
49
+ def name(self) -> str:
50
+ return self.space.name + "Trial"
51
+
52
+ def eval_arg_value(self, device) -> wp.codegen.StructInstance:
53
+ return self.space.space_arg_value(device)
54
+
55
+ def _make_element_eval_arg(self):
56
+ @cache.dynamic_struct(suffix=self.name)
57
+ class ElementEvalArg:
58
+ elt_arg: self.domain.ElementArg
59
+ eval_arg: self.EvalArg
60
+
61
+ return ElementEvalArg
62
+
63
+ def _make_eval_inner(self):
64
+ @cache.dynamic_func(suffix=self.name)
65
+ def eval_trial_inner(args: self.ElementEvalArg, s: Sample):
66
+ weight = self.space.element_inner_weight(
67
+ args.elt_arg,
68
+ args.eval_arg,
69
+ s.element_index,
70
+ s.element_coords,
71
+ get_node_index_in_element(s.trial_dof),
72
+ )
73
+ return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
74
+
75
+ return eval_trial_inner
76
+
77
+ def _make_eval_grad_inner(self):
78
+ if not self.gradient_valid():
79
+ return None
80
+
81
+ @cache.dynamic_func(suffix=self.name)
82
+ def eval_nabla_trial_inner(args: self.ElementEvalArg, s: Sample):
83
+ nabla_weight = self.space.element_inner_weight_gradient(
84
+ args.elt_arg,
85
+ args.eval_arg,
86
+ s.element_index,
87
+ s.element_coords,
88
+ get_node_index_in_element(s.trial_dof),
89
+ )
90
+ grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
91
+ return utils.generalized_outer(
92
+ self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
93
+ utils.apply_right(nabla_weight, grad_transform),
94
+ )
95
+
96
+ return eval_nabla_trial_inner
97
+
98
+ def _make_eval_div_inner(self):
99
+ if not self.divergence_valid():
100
+ return None
101
+
102
+ @cache.dynamic_func(suffix=self.name)
103
+ def eval_div_trial_inner(args: self.ElementEvalArg, s: Sample):
104
+ nabla_weight = self.space.element_inner_weight_gradient(
105
+ args.elt_arg,
106
+ args.eval_arg,
107
+ s.element_index,
108
+ s.element_coords,
109
+ get_node_index_in_element(s.trial_dof),
110
+ )
111
+ grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
112
+ return utils.generalized_inner(
113
+ self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
114
+ utils.apply_right(nabla_weight, grad_transform),
115
+ )
116
+
117
+ return eval_div_trial_inner
118
+
119
+ def _make_eval_outer(self):
120
+ @cache.dynamic_func(suffix=self.name)
121
+ def eval_trial_outer(args: self.ElementEvalArg, s: Sample):
122
+ weight = self.space.element_outer_weight(
123
+ args.elt_arg,
124
+ args.eval_arg,
125
+ s.element_index,
126
+ s.element_coords,
127
+ get_node_index_in_element(s.trial_dof),
128
+ )
129
+ return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
130
+
131
+ return eval_trial_outer
132
+
133
+ def _make_eval_grad_outer(self):
134
+ if not self.gradient_valid():
135
+ return None
136
+
137
+ @cache.dynamic_func(suffix=self.name)
138
+ def eval_nabla_trial_outer(args: self.ElementEvalArg, s: Sample):
139
+ nabla_weight = self.space.element_outer_weight_gradient(
140
+ args.elt_arg,
141
+ args.eval_arg,
142
+ s.element_index,
143
+ s.element_coords,
144
+ get_node_index_in_element(s.trial_dof),
145
+ )
146
+ grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
147
+ return utils.generalized_outer(
148
+ self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
149
+ utils.apply_right(nabla_weight, grad_transform),
150
+ )
151
+
152
+ return eval_nabla_trial_outer
153
+
154
+ def _make_eval_div_outer(self):
155
+ if not self.divergence_valid():
156
+ return None
157
+
158
+ @cache.dynamic_func(suffix=self.name)
159
+ def eval_div_trial_outer(args: self.ElementEvalArg, s: Sample):
160
+ nabla_weight = self.space.element_outer_weight_gradient(
161
+ args.elt_arg,
162
+ args.eval_arg,
163
+ s.element_index,
164
+ s.element_coords,
165
+ get_node_index_in_element(s.trial_dof),
166
+ )
167
+ grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
168
+ return utils.generalized_inner(
169
+ self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
170
+ utils.apply_right(nabla_weight, grad_transform),
171
+ )
172
+
173
+ return eval_div_trial_outer
174
+
175
+ def _make_at_node(self):
176
+ @cache.dynamic_func(suffix=self.name)
177
+ def at_node(args: self.ElementEvalArg, s: Sample):
178
+ node_coords = self.space.node_coords_in_element(
179
+ args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(s.trial_dof)
180
+ )
181
+ return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
182
+
183
+ return at_node