warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__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 (271) hide show
  1. docs/conf.py +17 -5
  2. examples/env/env_ant.py +1 -1
  3. examples/env/env_cartpole.py +1 -1
  4. examples/env/env_humanoid.py +1 -1
  5. examples/env/env_usd.py +4 -1
  6. examples/env/environment.py +8 -9
  7. examples/example_dem.py +34 -33
  8. examples/example_diffray.py +364 -337
  9. examples/example_fluid.py +32 -23
  10. examples/example_jacobian_ik.py +97 -93
  11. examples/example_marching_cubes.py +6 -16
  12. examples/example_mesh.py +6 -16
  13. examples/example_mesh_intersect.py +16 -14
  14. examples/example_nvdb.py +14 -16
  15. examples/example_raycast.py +14 -13
  16. examples/example_raymarch.py +16 -23
  17. examples/example_render_opengl.py +19 -10
  18. examples/example_sim_cartpole.py +82 -78
  19. examples/example_sim_cloth.py +45 -48
  20. examples/example_sim_fk_grad.py +51 -44
  21. examples/example_sim_fk_grad_torch.py +47 -40
  22. examples/example_sim_grad_bounce.py +108 -133
  23. examples/example_sim_grad_cloth.py +99 -113
  24. examples/example_sim_granular.py +5 -6
  25. examples/{example_sim_sdf_shape.py → example_sim_granular_collision_sdf.py} +37 -26
  26. examples/example_sim_neo_hookean.py +51 -55
  27. examples/example_sim_particle_chain.py +4 -4
  28. examples/example_sim_quadruped.py +126 -81
  29. examples/example_sim_rigid_chain.py +54 -61
  30. examples/example_sim_rigid_contact.py +66 -70
  31. examples/example_sim_rigid_fem.py +3 -3
  32. examples/example_sim_rigid_force.py +1 -1
  33. examples/example_sim_rigid_gyroscopic.py +3 -4
  34. examples/example_sim_rigid_kinematics.py +28 -39
  35. examples/example_sim_trajopt.py +112 -110
  36. examples/example_sph.py +9 -8
  37. examples/example_wave.py +7 -7
  38. examples/fem/bsr_utils.py +30 -17
  39. examples/fem/example_apic_fluid.py +85 -69
  40. examples/fem/example_convection_diffusion.py +97 -93
  41. examples/fem/example_convection_diffusion_dg.py +142 -149
  42. examples/fem/example_convection_diffusion_dg0.py +141 -136
  43. examples/fem/example_deformed_geometry.py +146 -0
  44. examples/fem/example_diffusion.py +115 -84
  45. examples/fem/example_diffusion_3d.py +116 -86
  46. examples/fem/example_diffusion_mgpu.py +102 -79
  47. examples/fem/example_mixed_elasticity.py +139 -100
  48. examples/fem/example_navier_stokes.py +175 -162
  49. examples/fem/example_stokes.py +143 -111
  50. examples/fem/example_stokes_transfer.py +186 -157
  51. examples/fem/mesh_utils.py +59 -97
  52. examples/fem/plot_utils.py +138 -17
  53. tools/ci/publishing/build_nodes_info.py +54 -0
  54. warp/__init__.py +4 -3
  55. warp/__init__.pyi +1 -0
  56. warp/bin/warp-clang.dll +0 -0
  57. warp/bin/warp.dll +0 -0
  58. warp/build.py +5 -3
  59. warp/build_dll.py +29 -9
  60. warp/builtins.py +836 -492
  61. warp/codegen.py +864 -553
  62. warp/config.py +3 -1
  63. warp/context.py +389 -172
  64. warp/fem/__init__.py +24 -6
  65. warp/fem/cache.py +318 -25
  66. warp/fem/dirichlet.py +7 -3
  67. warp/fem/domain.py +14 -0
  68. warp/fem/field/__init__.py +30 -38
  69. warp/fem/field/field.py +149 -0
  70. warp/fem/field/nodal_field.py +244 -138
  71. warp/fem/field/restriction.py +8 -6
  72. warp/fem/field/test.py +127 -59
  73. warp/fem/field/trial.py +117 -60
  74. warp/fem/geometry/__init__.py +5 -1
  75. warp/fem/geometry/deformed_geometry.py +271 -0
  76. warp/fem/geometry/element.py +24 -1
  77. warp/fem/geometry/geometry.py +86 -14
  78. warp/fem/geometry/grid_2d.py +112 -54
  79. warp/fem/geometry/grid_3d.py +134 -65
  80. warp/fem/geometry/hexmesh.py +953 -0
  81. warp/fem/geometry/partition.py +85 -33
  82. warp/fem/geometry/quadmesh_2d.py +532 -0
  83. warp/fem/geometry/tetmesh.py +451 -115
  84. warp/fem/geometry/trimesh_2d.py +197 -92
  85. warp/fem/integrate.py +534 -268
  86. warp/fem/operator.py +58 -31
  87. warp/fem/polynomial.py +11 -0
  88. warp/fem/quadrature/__init__.py +1 -1
  89. warp/fem/quadrature/pic_quadrature.py +150 -58
  90. warp/fem/quadrature/quadrature.py +209 -57
  91. warp/fem/space/__init__.py +230 -53
  92. warp/fem/space/basis_space.py +489 -0
  93. warp/fem/space/collocated_function_space.py +105 -0
  94. warp/fem/space/dof_mapper.py +49 -2
  95. warp/fem/space/function_space.py +90 -39
  96. warp/fem/space/grid_2d_function_space.py +149 -496
  97. warp/fem/space/grid_3d_function_space.py +173 -538
  98. warp/fem/space/hexmesh_function_space.py +352 -0
  99. warp/fem/space/partition.py +129 -76
  100. warp/fem/space/quadmesh_2d_function_space.py +369 -0
  101. warp/fem/space/restriction.py +46 -34
  102. warp/fem/space/shape/__init__.py +15 -0
  103. warp/fem/space/shape/cube_shape_function.py +738 -0
  104. warp/fem/space/shape/shape_function.py +103 -0
  105. warp/fem/space/shape/square_shape_function.py +611 -0
  106. warp/fem/space/shape/tet_shape_function.py +567 -0
  107. warp/fem/space/shape/triangle_shape_function.py +429 -0
  108. warp/fem/space/tetmesh_function_space.py +132 -1039
  109. warp/fem/space/topology.py +295 -0
  110. warp/fem/space/trimesh_2d_function_space.py +104 -742
  111. warp/fem/types.py +13 -11
  112. warp/fem/utils.py +335 -60
  113. warp/native/array.h +120 -34
  114. warp/native/builtin.h +101 -72
  115. warp/native/bvh.cpp +73 -325
  116. warp/native/bvh.cu +406 -23
  117. warp/native/bvh.h +22 -40
  118. warp/native/clang/clang.cpp +1 -0
  119. warp/native/crt.h +2 -0
  120. warp/native/cuda_util.cpp +8 -3
  121. warp/native/cuda_util.h +1 -0
  122. warp/native/exports.h +1522 -1243
  123. warp/native/intersect.h +19 -4
  124. warp/native/intersect_adj.h +8 -8
  125. warp/native/mat.h +76 -17
  126. warp/native/mesh.cpp +33 -108
  127. warp/native/mesh.cu +114 -18
  128. warp/native/mesh.h +395 -40
  129. warp/native/noise.h +272 -329
  130. warp/native/quat.h +51 -8
  131. warp/native/rand.h +44 -34
  132. warp/native/reduce.cpp +1 -1
  133. warp/native/sparse.cpp +4 -4
  134. warp/native/sparse.cu +163 -155
  135. warp/native/spatial.h +2 -2
  136. warp/native/temp_buffer.h +18 -14
  137. warp/native/vec.h +103 -21
  138. warp/native/warp.cpp +2 -1
  139. warp/native/warp.cu +28 -3
  140. warp/native/warp.h +4 -3
  141. warp/render/render_opengl.py +261 -109
  142. warp/sim/__init__.py +1 -2
  143. warp/sim/articulation.py +385 -185
  144. warp/sim/import_mjcf.py +59 -48
  145. warp/sim/import_urdf.py +15 -15
  146. warp/sim/import_usd.py +174 -102
  147. warp/sim/inertia.py +17 -18
  148. warp/sim/integrator_xpbd.py +4 -3
  149. warp/sim/model.py +330 -250
  150. warp/sim/render.py +1 -1
  151. warp/sparse.py +625 -152
  152. warp/stubs.py +341 -309
  153. warp/tape.py +9 -6
  154. warp/tests/__main__.py +3 -6
  155. warp/tests/assets/curlnoise_golden.npy +0 -0
  156. warp/tests/assets/pnoise_golden.npy +0 -0
  157. warp/tests/{test_class_kernel.py → aux_test_class_kernel.py} +9 -1
  158. warp/tests/aux_test_conditional_unequal_types_kernels.py +21 -0
  159. warp/tests/{test_dependent.py → aux_test_dependent.py} +2 -2
  160. warp/tests/{test_reference.py → aux_test_reference.py} +1 -1
  161. warp/tests/aux_test_unresolved_func.py +14 -0
  162. warp/tests/aux_test_unresolved_symbol.py +14 -0
  163. warp/tests/disabled_kinematics.py +239 -0
  164. warp/tests/run_coverage_serial.py +31 -0
  165. warp/tests/test_adam.py +103 -106
  166. warp/tests/test_arithmetic.py +94 -74
  167. warp/tests/test_array.py +82 -101
  168. warp/tests/test_array_reduce.py +57 -23
  169. warp/tests/test_atomic.py +64 -28
  170. warp/tests/test_bool.py +22 -12
  171. warp/tests/test_builtins_resolution.py +1292 -0
  172. warp/tests/test_bvh.py +18 -18
  173. warp/tests/test_closest_point_edge_edge.py +54 -57
  174. warp/tests/test_codegen.py +165 -134
  175. warp/tests/test_compile_consts.py +28 -20
  176. warp/tests/test_conditional.py +108 -24
  177. warp/tests/test_copy.py +10 -12
  178. warp/tests/test_ctypes.py +112 -88
  179. warp/tests/test_dense.py +21 -14
  180. warp/tests/test_devices.py +98 -0
  181. warp/tests/test_dlpack.py +75 -75
  182. warp/tests/test_examples.py +237 -0
  183. warp/tests/test_fabricarray.py +22 -24
  184. warp/tests/test_fast_math.py +15 -11
  185. warp/tests/test_fem.py +1034 -124
  186. warp/tests/test_fp16.py +23 -16
  187. warp/tests/test_func.py +187 -86
  188. warp/tests/test_generics.py +194 -49
  189. warp/tests/test_grad.py +123 -181
  190. warp/tests/test_grad_customs.py +176 -0
  191. warp/tests/test_hash_grid.py +35 -34
  192. warp/tests/test_import.py +10 -23
  193. warp/tests/test_indexedarray.py +24 -25
  194. warp/tests/test_intersect.py +18 -9
  195. warp/tests/test_large.py +141 -0
  196. warp/tests/test_launch.py +14 -41
  197. warp/tests/test_lerp.py +64 -65
  198. warp/tests/test_lvalue.py +493 -0
  199. warp/tests/test_marching_cubes.py +12 -13
  200. warp/tests/test_mat.py +517 -2898
  201. warp/tests/test_mat_lite.py +115 -0
  202. warp/tests/test_mat_scalar_ops.py +2889 -0
  203. warp/tests/test_math.py +103 -9
  204. warp/tests/test_matmul.py +304 -69
  205. warp/tests/test_matmul_lite.py +410 -0
  206. warp/tests/test_mesh.py +60 -22
  207. warp/tests/test_mesh_query_aabb.py +21 -25
  208. warp/tests/test_mesh_query_point.py +111 -22
  209. warp/tests/test_mesh_query_ray.py +12 -24
  210. warp/tests/test_mlp.py +30 -22
  211. warp/tests/test_model.py +92 -89
  212. warp/tests/test_modules_lite.py +39 -0
  213. warp/tests/test_multigpu.py +88 -114
  214. warp/tests/test_noise.py +12 -11
  215. warp/tests/test_operators.py +16 -20
  216. warp/tests/test_options.py +11 -11
  217. warp/tests/test_pinned.py +17 -18
  218. warp/tests/test_print.py +32 -11
  219. warp/tests/test_quat.py +275 -129
  220. warp/tests/test_rand.py +18 -16
  221. warp/tests/test_reload.py +38 -34
  222. warp/tests/test_rounding.py +50 -43
  223. warp/tests/test_runlength_encode.py +168 -20
  224. warp/tests/test_smoothstep.py +9 -11
  225. warp/tests/test_snippet.py +143 -0
  226. warp/tests/test_sparse.py +261 -63
  227. warp/tests/test_spatial.py +276 -243
  228. warp/tests/test_streams.py +110 -85
  229. warp/tests/test_struct.py +268 -63
  230. warp/tests/test_tape.py +39 -21
  231. warp/tests/test_torch.py +90 -86
  232. warp/tests/test_transient_module.py +10 -12
  233. warp/tests/test_types.py +363 -0
  234. warp/tests/test_utils.py +451 -0
  235. warp/tests/test_vec.py +354 -2050
  236. warp/tests/test_vec_lite.py +73 -0
  237. warp/tests/test_vec_scalar_ops.py +2099 -0
  238. warp/tests/test_volume.py +418 -376
  239. warp/tests/test_volume_write.py +124 -134
  240. warp/tests/unittest_serial.py +35 -0
  241. warp/tests/unittest_suites.py +291 -0
  242. warp/tests/unittest_utils.py +342 -0
  243. warp/tests/{test_misc.py → unused_test_misc.py} +13 -5
  244. warp/tests/{test_debug.py → walkthough_debug.py} +3 -17
  245. warp/thirdparty/appdirs.py +36 -45
  246. warp/thirdparty/unittest_parallel.py +589 -0
  247. warp/types.py +622 -211
  248. warp/utils.py +54 -393
  249. warp_lang-1.0.0b6.dist-info/METADATA +238 -0
  250. warp_lang-1.0.0b6.dist-info/RECORD +409 -0
  251. {warp_lang-1.0.0b2.dist-info → warp_lang-1.0.0b6.dist-info}/WHEEL +1 -1
  252. examples/example_cache_management.py +0 -40
  253. examples/example_multigpu.py +0 -54
  254. examples/example_struct.py +0 -65
  255. examples/fem/example_stokes_transfer_3d.py +0 -210
  256. warp/bin/warp-clang.so +0 -0
  257. warp/bin/warp.so +0 -0
  258. warp/fem/field/discrete_field.py +0 -80
  259. warp/fem/space/nodal_function_space.py +0 -233
  260. warp/tests/test_all.py +0 -223
  261. warp/tests/test_array_scan.py +0 -60
  262. warp/tests/test_base.py +0 -208
  263. warp/tests/test_unresolved_func.py +0 -7
  264. warp/tests/test_unresolved_symbol.py +0 -7
  265. warp_lang-1.0.0b2.dist-info/METADATA +0 -26
  266. warp_lang-1.0.0b2.dist-info/RECORD +0 -380
  267. /warp/tests/{test_compile_consts_dummy.py → aux_test_compile_consts_dummy.py} +0 -0
  268. /warp/tests/{test_reference_reference.py → aux_test_reference_reference.py} +0 -0
  269. /warp/tests/{test_square.py → aux_test_square.py} +0 -0
  270. {warp_lang-1.0.0b2.dist-info → warp_lang-1.0.0b6.dist-info}/LICENSE.md +0 -0
  271. {warp_lang-1.0.0b2.dist-info → warp_lang-1.0.0b6.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,149 @@
1
+ from typing import Any
2
+
3
+ import warp as wp
4
+
5
+ from warp.fem.types import Sample
6
+ from warp.fem.space import FunctionSpace, SpacePartition
7
+ from warp.fem.geometry import Geometry, DeformedGeometry
8
+
9
+
10
+ class FieldLike:
11
+ """Base class for integrable fields"""
12
+
13
+ EvalArg: wp.codegen.Struct
14
+ """Structure containing field-level arguments passed to device functions for field evaluation"""
15
+
16
+ ElementEvalArg: wp.codegen.Struct
17
+ """Structure combining geometry-level and field-level arguments passed to device functions for field evaluation"""
18
+
19
+ def eval_arg_value(self, device) -> "EvalArg":
20
+ """Value of the field-level arguments to be passed to device functions"""
21
+ raise NotImplementedError
22
+
23
+ @property
24
+ def degree(self) -> int:
25
+ """Polynomial degree of the field, used to estimate necessary quadrature order"""
26
+ raise NotImplementedError
27
+
28
+ @property
29
+ def name(self) -> str:
30
+ raise NotImplementedError
31
+
32
+ @property
33
+ def __str__(self) -> str:
34
+ return self.name
35
+
36
+ def eval_arg_value(self, device):
37
+ """Value of arguments to be passed to device functions"""
38
+ raise NotImplementedError
39
+
40
+ @staticmethod
41
+ def eval_inner(args: "ElementEvalArg", s: "Sample"):
42
+ """Device function evaluating the inner field value at a sample point"""
43
+ raise NotImplementedError
44
+
45
+ @staticmethod
46
+ def eval_grad_inner(args: "ElementEvalArg", s: "Sample"):
47
+ """Device function evaluating the inner field gradient at a sample point"""
48
+ raise NotImplementedError
49
+
50
+ @staticmethod
51
+ def eval_div_inner(args: "ElementEvalArg", s: "Sample"):
52
+ """Device function evaluating the inner field divergence at a sample point"""
53
+ raise NotImplementedError
54
+
55
+ @staticmethod
56
+ def eval_outer(args: "ElementEvalArg", s: "Sample"):
57
+ """Device function evaluating the outer field value at a sample point"""
58
+ raise NotImplementedError
59
+
60
+ @staticmethod
61
+ def eval_grad_outer(args: "ElementEvalArg", s: "Sample"):
62
+ """Device function evaluating the outer field gradient at a sample point"""
63
+ raise NotImplementedError
64
+
65
+ @staticmethod
66
+ def eval_div_outer(args: "ElementEvalArg", s: "Sample"):
67
+ """Device function evaluating the outer field divergence at a sample point"""
68
+ raise NotImplementedError
69
+
70
+
71
+ class SpaceField(FieldLike):
72
+ """Base class for fields defined over a function space"""
73
+
74
+ def __init__(self, space: FunctionSpace, space_partition: SpacePartition):
75
+ self._space = space
76
+ self._space_partition = space_partition
77
+
78
+ @property
79
+ def space(self) -> FunctionSpace:
80
+ return self._space
81
+
82
+ @property
83
+ def space_partition(self) -> SpacePartition:
84
+ return self._space_partition
85
+
86
+ @property
87
+ def degree(self) -> int:
88
+ return self.space.degree
89
+
90
+ @property
91
+ def dtype(self) -> type:
92
+ return self.space.dtype
93
+
94
+ @property
95
+ def dof_dtype(self) -> type:
96
+ return self.space.dof_dtype
97
+
98
+ def gradient_valid(self) -> bool:
99
+ """Whether gradient operator can be computed. Only for scalar and vector fields as higher-order tensors are not support yet"""
100
+ return not wp.types.type_is_matrix(self.dtype)
101
+
102
+ def divergence_valid(self) -> bool:
103
+ """Whether divergence of this field can be computed. Only for vector and tensor fields with same dimension as embedding geometry"""
104
+ if wp.types.type_is_vector(self.dtype):
105
+ return wp.types.type_length(self.dtype) == self.space.geometry.dimension
106
+ if wp.types.type_is_matrix(self.dtype):
107
+ return self.dtype._shape_[0] == self.space.geometry.dimension
108
+ return False
109
+
110
+ def _make_eval_degree(self):
111
+ ORDER = self.space.ORDER
112
+ from warp.fem import cache
113
+
114
+ @cache.dynamic_func(suffix=self.name)
115
+ def degree(args: self.ElementEvalArg):
116
+ return ORDER
117
+
118
+ return degree
119
+
120
+
121
+ class DiscreteField(SpaceField):
122
+ """Explicitly-valued field defined over a partition of a discrete function space"""
123
+
124
+ @property
125
+ def dof_values(self) -> wp.array:
126
+ """Array of degrees of freedom values"""
127
+ raise NotImplementedError
128
+
129
+ @dof_values.setter
130
+ def dof_values(self, values: wp.array):
131
+ """Sets degrees of freedom values from an array"""
132
+ raise NotImplementedError
133
+
134
+ def trace(self) -> "DiscreteField":
135
+ """Trace of this field over a lower-dimensional function space"""
136
+ raise NotImplementedError
137
+
138
+ @staticmethod
139
+ def set_node_value(args: "FieldLike.EvalArg", node_index: int, value: Any):
140
+ """Device function setting the value at given node"""
141
+ raise NotImplementedError
142
+
143
+ @property
144
+ def name(self) -> str:
145
+ return f"{self.__class__.__qualname__}_{self.space.name}_{self.space_partition.name}"
146
+
147
+ def make_deformed_geometry(self) -> Geometry:
148
+ """Returns a deformed version of the underlying geometry using this field's values as displacement"""
149
+ return DeformedGeometry(self)
@@ -1,193 +1,299 @@
1
1
  import warp as wp
2
2
 
3
- from warp.fem.space import NodalFunctionSpace, SpacePartition
3
+ from warp.fem.space import CollocatedFunctionSpace, SpacePartition
4
4
  from warp.fem import cache, utils
5
5
  from warp.fem.types import Sample, ElementIndex, NULL_NODE_INDEX
6
6
 
7
- from .discrete_field import DiscreteField
7
+ from .field import DiscreteField
8
8
 
9
9
 
10
- class NodalField(DiscreteField):
11
- def __init__(self, space_partition: SpacePartition, space: NodalFunctionSpace = None):
12
- if space is None:
13
- space = space_partition.space
10
+ class NodalFieldBase(DiscreteField):
11
+ """Base class for nodal field and nodal field traces. Does not hold values"""
14
12
 
13
+ def __init__(self, space: CollocatedFunctionSpace, space_partition: SpacePartition):
15
14
  super().__init__(space, space_partition)
16
15
 
17
- self._dof_values = wp.zeros(n=self.space_partition.node_count(), dtype=self.dof_dtype)
18
-
19
- self.EvalArg = NodalField._make_eval_arg(self.space, self.space_partition)
20
- self.eval_degree = DiscreteField._make_eval_degree(self.EvalArg, self.space)
21
- self.set_node_value = NodalField._make_set_node_value(self.EvalArg, self.space)
22
-
23
- read_node_value = NodalField._make_read_node_value(self.EvalArg, self.space, self.space_partition)
16
+ self.EvalArg = self._make_eval_arg()
17
+ self.ElementEvalArg = self._make_element_eval_arg()
18
+ self.eval_degree = DiscreteField._make_eval_degree(self)
24
19
 
25
- self.eval_inner = NodalField._make_eval_inner(self.EvalArg, self.space, read_node_value)
26
- self.eval_outer = NodalField._make_eval_outer(self.EvalArg, self.space, read_node_value)
27
- self.eval_grad_inner = NodalField._make_eval_grad_inner(self.EvalArg, self.space, read_node_value)
28
- self.eval_grad_outer = NodalField._make_eval_grad_outer(self.EvalArg, self.space, read_node_value)
20
+ self._read_node_value = self._make_read_node_value()
29
21
 
30
- def eval_arg_value(self, device):
31
- arg = self.EvalArg()
32
- arg.space_arg = self.space.space_arg_value(device)
33
- arg.partition_arg = self.space_partition.partition_arg_value(device)
34
- arg.dof_values = self._dof_values.to(device)
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()
35
30
 
36
- return arg
31
+ self.set_node_value = self._make_set_node_value()
37
32
 
38
- @property
39
- def dof_values(self):
40
- return self._dof_values
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)
41
60
 
42
- @dof_values.setter
43
- def dof_values(self, values):
44
- if isinstance(values, wp.array):
45
- self._dof_values = values
46
- else:
47
- self._dof_values = wp.array(values, dtype=self.dof_dtype)
61
+ return self.space.dof_mapper.dof_to_value(args.eval_arg.dof_values[pidx])
48
62
 
49
- class Trace(DiscreteField):
50
- def __init__(self, field):
51
- self._field = field
52
- super().__init__(field.space.trace(), field.space_partition)
63
+ return read_node_value
53
64
 
54
- self.EvalArg = field.EvalArg
55
- self.eval_degree = DiscreteField._make_eval_degree(self.EvalArg, self.space)
56
- self.eval_arg_value = field.eval_arg_value
65
+ def _make_eval_inner(self):
66
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
57
67
 
58
- self.set_node_value = field.set_node_value
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
59
78
 
60
- read_node_value = NodalField._make_read_node_value(self.EvalArg, self.space, self.space_partition)
79
+ return eval_inner
61
80
 
62
- self.eval_inner = NodalField._make_eval_inner(self.EvalArg, self.space, read_node_value)
63
- self.eval_outer = NodalField._make_eval_outer(self.EvalArg, self.space, read_node_value)
64
- self.eval_grad_inner = NodalField._make_eval_grad_inner(self.EvalArg, self.space, read_node_value)
65
- self.eval_grad_outer = NodalField._make_eval_grad_outer(self.EvalArg, self.space, read_node_value)
81
+ def _make_eval_grad_inner(self, world_space: bool):
82
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
66
83
 
67
- def trace(self) -> Trace:
68
- trace_field = NodalField.Trace(self)
69
- return trace_field
84
+ if not self.gradient_valid():
85
+ return None
70
86
 
71
- @staticmethod
72
- def _make_eval_arg(space: NodalFunctionSpace, space_partition: SpacePartition):
73
- from warp.fem import cache
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
74
103
 
75
- class EvalArg:
76
- space_arg: space.SpaceArg
77
- partition_arg: space_partition.PartitionArg
78
- dof_values: wp.array(dtype=space.dof_dtype)
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)
79
109
 
80
- EvalArg.__qualname__ = f"{space.name}_{space_partition.name}_FieldEvalArg"
81
- return cache.get_struct(EvalArg)
110
+ return eval_grad_inner_world_space if world_space else eval_grad_inner_ref_space
82
111
 
83
- @staticmethod
84
- def _make_set_node_value(EvalArg, space: NodalFunctionSpace):
85
- def set_node_value(args: EvalArg, partition_node_index: int, value: space.dtype):
86
- args.dof_values[partition_node_index] = space.dof_mapper.value_to_dof(value)
112
+ def _make_eval_div_inner(self):
113
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
87
114
 
88
- return cache.get_func(set_node_value, space)
115
+ if not self.divergence_valid():
116
+ return None
89
117
 
90
- @staticmethod
91
- def _make_read_node_value(EvalArg, space: NodalFunctionSpace, space_partition: SpacePartition):
92
- def read_node_value(args: EvalArg, geo_element_index: ElementIndex, node_index_in_elt: int):
93
- nidx = space.element_node_index(args.space_arg, geo_element_index, node_index_in_elt)
94
- pidx = space_partition.partition_node_index(args.partition_arg, nidx)
95
- if pidx == NULL_NODE_INDEX:
96
- return space.dtype(0.0)
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
+ )
97
131
 
98
- return space.dof_mapper.dof_to_value(args.dof_values[pidx])
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
99
143
 
100
- return cache.get_func(read_node_value, f"{space}_{space_partition}")
144
+ return eval_div_inner
101
145
 
102
- @staticmethod
103
- def _make_eval_inner(
104
- EvalArg,
105
- space: NodalFunctionSpace,
106
- read_node_value: wp.Function,
107
- ):
108
- NODES_PER_ELEMENT = space.NODES_PER_ELEMENT
146
+ def _make_eval_outer(self):
147
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
109
148
 
110
- def eval_inner(args: EvalArg, s: Sample):
111
- res = space.element_inner_weight(args.space_arg, s.element_index, s.element_coords, 0) * read_node_value(
112
- args, s.element_index, 0
113
- )
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)
114
158
  for k in range(1, NODES_PER_ELEMENT):
115
- res += space.element_inner_weight(
116
- args.space_arg, s.element_index, s.element_coords, k
117
- ) * read_node_value(args, s.element_index, k)
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)
118
166
  return res
119
167
 
120
- return cache.get_func(eval_inner, read_node_value.key)
168
+ return eval_outer
121
169
 
122
- @staticmethod
123
- def _make_eval_grad_inner(
124
- EvalArg,
125
- space: NodalFunctionSpace,
126
- read_node_value: wp.Function,
127
- ):
128
- NODES_PER_ELEMENT = space.NODES_PER_ELEMENT
170
+ def _make_eval_grad_outer(self, world_space: bool):
171
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
129
172
 
130
- if wp.types.type_is_matrix(space.dtype):
131
- # There is no Warp high-order tensor type to represent matrix gradients
173
+ if not self.gradient_valid():
132
174
  return None
133
175
 
134
- def eval_grad_inner(args: EvalArg, s: Sample):
176
+ @cache.dynamic_func(suffix=self.name)
177
+ def eval_grad_outer_ref_space(args: self.ElementEvalArg, s: Sample):
135
178
  res = utils.generalized_outer(
136
- space.element_inner_weight_gradient(args.space_arg, s.element_index, s.element_coords, 0),
137
- read_node_value(args, s.element_index, 0),
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
+ ),
138
183
  )
139
-
140
184
  for k in range(1, NODES_PER_ELEMENT):
141
185
  res += utils.generalized_outer(
142
- space.element_inner_weight_gradient(args.space_arg, s.element_index, s.element_coords, k),
143
- read_node_value(args, s.element_index, k),
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
+ ),
144
190
  )
145
191
  return res
146
192
 
147
- return cache.get_func(eval_grad_inner, read_node_value.key)
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)
148
198
 
149
- @staticmethod
150
- def _make_eval_outer(
151
- EvalArg,
152
- space: NodalFunctionSpace,
153
- read_node_value: wp.Function,
154
- ):
155
- NODES_PER_ELEMENT = space.NODES_PER_ELEMENT
199
+ return eval_grad_outer_world_space if world_space else eval_grad_outer_ref_space
156
200
 
157
- def eval_outer(args: EvalArg, s: Sample):
158
- res = space.element_outer_weight(args.space_arg, s.element_index, s.element_coords, 0) * read_node_value(
159
- args, s.element_index, 0
160
- )
161
- for k in range(1, NODES_PER_ELEMENT):
162
- res += space.element_outer_weight(
163
- args.space_arg, s.element_index, s.element_coords, k
164
- ) * read_node_value(args, s.element_index, k)
165
- return res
166
-
167
- return cache.get_func(eval_outer, read_node_value.key)
168
-
169
- @staticmethod
170
- def _make_eval_grad_outer(
171
- EvalArg,
172
- space: NodalFunctionSpace,
173
- read_node_value: wp.Function,
174
- ):
175
- NODES_PER_ELEMENT = space.NODES_PER_ELEMENT
201
+ def _make_eval_div_outer(self):
202
+ NODES_PER_ELEMENT = self.space.topology.NODES_PER_ELEMENT
176
203
 
177
- if wp.types.type_is_matrix(space.dtype):
178
- # There is no Warp high-order tensor type to represent matrix gradients
204
+ if not self.divergence_valid():
179
205
  return None
180
206
 
181
- def eval_grad_outer(args: EvalArg, s: Sample):
182
- res = utils.generalized_outer(
183
- space.element_outer_weight_gradient(args.space_arg, s.element_index, s.element_coords, 0),
184
- read_node_value(args, s.element_index, 0),
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
+ ),
185
219
  )
186
220
  for k in range(1, NODES_PER_ELEMENT):
187
- res += utils.generalized_outer(
188
- space.element_outer_weight_gradient(args.space_arg, s.element_index, s.element_coords, k),
189
- read_node_value(args, s.element_index, k),
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
+ ),
190
229
  )
191
230
  return res
192
231
 
193
- return cache.get_func(eval_grad_outer, read_node_value.key)
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
@@ -1,19 +1,21 @@
1
1
  from warp.fem.space import SpaceRestriction
2
- from .discrete_field import DiscreteField
2
+ from .field import DiscreteField
3
3
 
4
4
 
5
5
  class FieldRestriction:
6
- """Restriction of a space to a given GeometryDomain"""
6
+ """Restriction of a discrete field to a given GeometryDomain"""
7
7
 
8
8
  def __init__(self, space_restriction: SpaceRestriction, field: DiscreteField):
9
-
10
- if field.space.DIMENSION - 1 == space_restriction.space.DIMENSION:
9
+ if field.space.dimension - 1 == space_restriction.space_topology.dimension:
11
10
  field = field.trace()
12
11
 
13
- if field.space.DIMENSION != space_restriction.space.DIMENSION:
12
+ if field.space.dimension != space_restriction.space_topology.dimension:
14
13
  raise ValueError("Incompatible space and field dimensions")
15
14
 
15
+ if field.space.topology != space_restriction.space_topology:
16
+ raise ValueError("Incompatible field and space restriction topologies")
17
+
16
18
  self.space_restriction = space_restriction
17
- self.space = self.space_restriction.space
18
19
  self.domain = self.space_restriction.domain
19
20
  self.field = field
21
+ self.space = self.field.space