warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__py3-none-manylinux_2_34_aarch64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of warp-lang might be problematic. Click here for more details.

Files changed (181) hide show
  1. warp/__init__.py +3 -1
  2. warp/__init__.pyi +3489 -1
  3. warp/autograd.py +45 -122
  4. warp/bin/warp-clang.so +0 -0
  5. warp/bin/warp.so +0 -0
  6. warp/build.py +241 -252
  7. warp/build_dll.py +125 -26
  8. warp/builtins.py +1907 -384
  9. warp/codegen.py +257 -101
  10. warp/config.py +12 -1
  11. warp/constants.py +1 -1
  12. warp/context.py +657 -223
  13. warp/dlpack.py +1 -1
  14. warp/examples/benchmarks/benchmark_cloth.py +2 -2
  15. warp/examples/benchmarks/benchmark_tile_sort.py +155 -0
  16. warp/examples/core/example_sample_mesh.py +1 -1
  17. warp/examples/core/example_spin_lock.py +93 -0
  18. warp/examples/core/example_work_queue.py +118 -0
  19. warp/examples/fem/example_adaptive_grid.py +5 -5
  20. warp/examples/fem/example_apic_fluid.py +1 -1
  21. warp/examples/fem/example_burgers.py +1 -1
  22. warp/examples/fem/example_convection_diffusion.py +9 -6
  23. warp/examples/fem/example_darcy_ls_optimization.py +489 -0
  24. warp/examples/fem/example_deformed_geometry.py +1 -1
  25. warp/examples/fem/example_diffusion.py +2 -2
  26. warp/examples/fem/example_diffusion_3d.py +1 -1
  27. warp/examples/fem/example_distortion_energy.py +1 -1
  28. warp/examples/fem/example_elastic_shape_optimization.py +387 -0
  29. warp/examples/fem/example_magnetostatics.py +5 -3
  30. warp/examples/fem/example_mixed_elasticity.py +5 -3
  31. warp/examples/fem/example_navier_stokes.py +11 -9
  32. warp/examples/fem/example_nonconforming_contact.py +5 -3
  33. warp/examples/fem/example_streamlines.py +8 -3
  34. warp/examples/fem/utils.py +9 -8
  35. warp/examples/interop/example_jax_ffi_callback.py +2 -2
  36. warp/examples/optim/example_drone.py +1 -1
  37. warp/examples/sim/example_cloth.py +1 -1
  38. warp/examples/sim/example_cloth_self_contact.py +48 -54
  39. warp/examples/tile/example_tile_block_cholesky.py +502 -0
  40. warp/examples/tile/example_tile_cholesky.py +2 -1
  41. warp/examples/tile/example_tile_convolution.py +1 -1
  42. warp/examples/tile/example_tile_filtering.py +1 -1
  43. warp/examples/tile/example_tile_matmul.py +1 -1
  44. warp/examples/tile/example_tile_mlp.py +2 -0
  45. warp/fabric.py +7 -7
  46. warp/fem/__init__.py +5 -0
  47. warp/fem/adaptivity.py +1 -1
  48. warp/fem/cache.py +152 -63
  49. warp/fem/dirichlet.py +2 -2
  50. warp/fem/domain.py +136 -6
  51. warp/fem/field/field.py +141 -99
  52. warp/fem/field/nodal_field.py +85 -39
  53. warp/fem/field/virtual.py +97 -52
  54. warp/fem/geometry/adaptive_nanogrid.py +91 -86
  55. warp/fem/geometry/closest_point.py +13 -0
  56. warp/fem/geometry/deformed_geometry.py +102 -40
  57. warp/fem/geometry/element.py +56 -2
  58. warp/fem/geometry/geometry.py +323 -22
  59. warp/fem/geometry/grid_2d.py +157 -62
  60. warp/fem/geometry/grid_3d.py +116 -20
  61. warp/fem/geometry/hexmesh.py +86 -20
  62. warp/fem/geometry/nanogrid.py +166 -86
  63. warp/fem/geometry/partition.py +59 -25
  64. warp/fem/geometry/quadmesh.py +86 -135
  65. warp/fem/geometry/tetmesh.py +47 -119
  66. warp/fem/geometry/trimesh.py +77 -270
  67. warp/fem/integrate.py +107 -52
  68. warp/fem/linalg.py +25 -58
  69. warp/fem/operator.py +124 -27
  70. warp/fem/quadrature/pic_quadrature.py +36 -14
  71. warp/fem/quadrature/quadrature.py +40 -16
  72. warp/fem/space/__init__.py +1 -1
  73. warp/fem/space/basis_function_space.py +66 -46
  74. warp/fem/space/basis_space.py +17 -4
  75. warp/fem/space/dof_mapper.py +1 -1
  76. warp/fem/space/function_space.py +2 -2
  77. warp/fem/space/grid_2d_function_space.py +4 -1
  78. warp/fem/space/hexmesh_function_space.py +4 -2
  79. warp/fem/space/nanogrid_function_space.py +3 -1
  80. warp/fem/space/partition.py +11 -2
  81. warp/fem/space/quadmesh_function_space.py +4 -1
  82. warp/fem/space/restriction.py +5 -2
  83. warp/fem/space/shape/__init__.py +10 -8
  84. warp/fem/space/tetmesh_function_space.py +4 -1
  85. warp/fem/space/topology.py +52 -21
  86. warp/fem/space/trimesh_function_space.py +4 -1
  87. warp/fem/utils.py +53 -8
  88. warp/jax.py +1 -2
  89. warp/jax_experimental/ffi.py +12 -17
  90. warp/jax_experimental/xla_ffi.py +37 -24
  91. warp/math.py +171 -1
  92. warp/native/array.h +99 -0
  93. warp/native/builtin.h +174 -31
  94. warp/native/coloring.cpp +1 -1
  95. warp/native/exports.h +118 -63
  96. warp/native/intersect.h +3 -3
  97. warp/native/mat.h +5 -10
  98. warp/native/mathdx.cpp +11 -5
  99. warp/native/matnn.h +1 -123
  100. warp/native/quat.h +28 -4
  101. warp/native/sparse.cpp +121 -258
  102. warp/native/sparse.cu +181 -274
  103. warp/native/spatial.h +305 -17
  104. warp/native/tile.h +583 -72
  105. warp/native/tile_radix_sort.h +1108 -0
  106. warp/native/tile_reduce.h +237 -2
  107. warp/native/tile_scan.h +240 -0
  108. warp/native/tuple.h +189 -0
  109. warp/native/vec.h +6 -16
  110. warp/native/warp.cpp +36 -4
  111. warp/native/warp.cu +574 -51
  112. warp/native/warp.h +47 -74
  113. warp/optim/linear.py +5 -1
  114. warp/paddle.py +7 -8
  115. warp/py.typed +0 -0
  116. warp/render/render_opengl.py +58 -29
  117. warp/render/render_usd.py +124 -61
  118. warp/sim/__init__.py +9 -0
  119. warp/sim/collide.py +252 -78
  120. warp/sim/graph_coloring.py +8 -1
  121. warp/sim/import_mjcf.py +4 -3
  122. warp/sim/import_usd.py +11 -7
  123. warp/sim/integrator.py +5 -2
  124. warp/sim/integrator_euler.py +1 -1
  125. warp/sim/integrator_featherstone.py +1 -1
  126. warp/sim/integrator_vbd.py +751 -320
  127. warp/sim/integrator_xpbd.py +1 -1
  128. warp/sim/model.py +265 -260
  129. warp/sim/utils.py +10 -7
  130. warp/sparse.py +303 -166
  131. warp/tape.py +52 -51
  132. warp/tests/cuda/test_conditional_captures.py +1046 -0
  133. warp/tests/cuda/test_streams.py +1 -1
  134. warp/tests/geometry/test_volume.py +2 -2
  135. warp/tests/interop/test_dlpack.py +9 -9
  136. warp/tests/interop/test_jax.py +0 -1
  137. warp/tests/run_coverage_serial.py +1 -1
  138. warp/tests/sim/disabled_kinematics.py +2 -2
  139. warp/tests/sim/{test_vbd.py → test_cloth.py} +296 -113
  140. warp/tests/sim/test_collision.py +159 -51
  141. warp/tests/sim/test_coloring.py +15 -1
  142. warp/tests/test_array.py +254 -2
  143. warp/tests/test_array_reduce.py +2 -2
  144. warp/tests/test_atomic_cas.py +299 -0
  145. warp/tests/test_codegen.py +142 -19
  146. warp/tests/test_conditional.py +47 -1
  147. warp/tests/test_ctypes.py +0 -20
  148. warp/tests/test_devices.py +8 -0
  149. warp/tests/test_fabricarray.py +4 -2
  150. warp/tests/test_fem.py +58 -25
  151. warp/tests/test_func.py +42 -1
  152. warp/tests/test_grad.py +1 -1
  153. warp/tests/test_lerp.py +1 -3
  154. warp/tests/test_map.py +481 -0
  155. warp/tests/test_mat.py +1 -24
  156. warp/tests/test_quat.py +6 -15
  157. warp/tests/test_rounding.py +10 -38
  158. warp/tests/test_runlength_encode.py +7 -7
  159. warp/tests/test_smoothstep.py +1 -1
  160. warp/tests/test_sparse.py +51 -2
  161. warp/tests/test_spatial.py +507 -1
  162. warp/tests/test_struct.py +2 -2
  163. warp/tests/test_tuple.py +265 -0
  164. warp/tests/test_types.py +2 -2
  165. warp/tests/test_utils.py +24 -18
  166. warp/tests/tile/test_tile.py +420 -1
  167. warp/tests/tile/test_tile_mathdx.py +518 -14
  168. warp/tests/tile/test_tile_reduce.py +213 -0
  169. warp/tests/tile/test_tile_shared_memory.py +130 -1
  170. warp/tests/tile/test_tile_sort.py +117 -0
  171. warp/tests/unittest_suites.py +4 -6
  172. warp/types.py +462 -308
  173. warp/utils.py +647 -86
  174. {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/METADATA +20 -6
  175. {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/RECORD +178 -166
  176. warp/stubs.py +0 -3381
  177. warp/tests/sim/test_xpbd.py +0 -399
  178. warp/tests/test_mlp.py +0 -282
  179. {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/WHEEL +0 -0
  180. {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/licenses/LICENSE.md +0 -0
  181. {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/top_level.txt +0 -0
warp/fem/domain.py CHANGED
@@ -13,6 +13,7 @@
13
13
  # See the License for the specific language governing permissions and
14
14
  # limitations under the License.
15
15
 
16
+ from functools import cached_property
16
17
  from typing import Any, Optional, Set, Union
17
18
 
18
19
  import warp as wp
@@ -27,7 +28,7 @@ from warp.fem.geometry import (
27
28
  WholeGeometryPartition,
28
29
  )
29
30
  from warp.fem.operator import Operator
30
- from warp.fem.types import ElementKind
31
+ from warp.fem.types import NULL_ELEMENT_INDEX, ElementKind
31
32
 
32
33
  GeometryOrPartition = Union[Geometry, GeometryPartition]
33
34
 
@@ -89,6 +90,9 @@ class GeometryDomain:
89
90
  element_index: wp.Function
90
91
  """Device function for retrieving an ElementIndex from a linearized index"""
91
92
 
93
+ element_partition_index: wp.Function
94
+ """Device function for retrieving linearized index in the domain's partition from an ElementIndex"""
95
+
92
96
  ElementArg: warp.codegen.Struct
93
97
  """Structure containing arguments to be passed to device functions computing element geometry"""
94
98
 
@@ -107,13 +111,34 @@ class GeometryDomain:
107
111
  element_normal: wp.Function
108
112
  """Device function returning the element normal at a sample point"""
109
113
 
114
+ element_closest_point: wp.Function
115
+ """Device function returning the coordinates of the closest point in a given element to a world position"""
116
+
117
+ element_coordinates: wp.Function
118
+ """Device function returning the coordinates corresponding to a world position in a given element reference system"""
119
+
110
120
  element_lookup: wp.Function
111
- """Device function returning the sample point corresponding to a world position"""
121
+ """Device function returning the sample point in the domain's geometry corresponding to a world position"""
122
+
123
+ element_partition_lookup: wp.Function
124
+ """Device function returning the sample point in the domain's geometry partition corresponding to a world position"""
112
125
 
113
126
  def notify_operator_usage(self, ops: Set[Operator]):
114
127
  """Makes the Domain aware that the operators `ops` will be applied"""
115
128
  pass
116
129
 
130
+ @cached_property
131
+ def DomainArg(self):
132
+ return self._make_domain_arg()
133
+
134
+ def _make_domain_arg(self):
135
+ @cache.dynamic_struct(suffix=self.name)
136
+ class DomainArg:
137
+ geo: self.ElementArg
138
+ index: self.ElementIndexArg
139
+
140
+ return DomainArg
141
+
117
142
 
118
143
  class Cells(GeometryDomain):
119
144
  """A Domain containing all cells of the geometry or geometry partition"""
@@ -145,13 +170,23 @@ class Cells(GeometryDomain):
145
170
  def element_index_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
146
171
  return self.geometry_partition.cell_arg_value(device)
147
172
 
173
+ def fill_element_index_arg(self, arg: ElementIndexArg, device: warp.context.Devicelike):
174
+ self.geometry_partition.fill_cell_arg(arg, device)
175
+
148
176
  @property
149
177
  def element_index(self) -> wp.Function:
150
178
  return self.geometry_partition.cell_index
151
179
 
180
+ @property
181
+ def element_partition_index(self) -> wp.Function:
182
+ return self.geometry_partition.partition_cell_index
183
+
152
184
  def element_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
153
185
  return self.geometry.cell_arg_value(device)
154
186
 
187
+ def fill_element_arg(self, arg: "ElementArg", device: warp.context.Devicelike):
188
+ self.geometry.fill_cell_arg(arg, device)
189
+
155
190
  @property
156
191
  def ElementArg(self) -> warp.codegen.Struct:
157
192
  return self.geometry.CellArg
@@ -176,10 +211,46 @@ class Cells(GeometryDomain):
176
211
  def element_normal(self) -> wp.Function:
177
212
  return self.geometry.cell_normal
178
213
 
214
+ @property
215
+ def element_closest_point(self) -> wp.Function:
216
+ return self.geometry.cell_closest_point
217
+
218
+ @property
219
+ def element_coordinates(self) -> wp.Function:
220
+ return self.geometry.cell_coordinates
221
+
179
222
  @property
180
223
  def element_lookup(self) -> wp.Function:
181
224
  return self.geometry.cell_lookup
182
225
 
226
+ @property
227
+ def element_partition_lookup(self) -> wp.Function:
228
+ pos_type = cache.cached_vec_type(self.geometry.dimension, dtype=float)
229
+
230
+ @cache.dynamic_func(suffix=self.geometry_partition.name)
231
+ def is_in_partition(args: self.ElementIndexArg, cell_index: int):
232
+ return self.geometry_partition.partition_cell_index(args, cell_index) != NULL_ELEMENT_INDEX
233
+
234
+ filtered_cell_lookup = self.geometry.make_filtered_cell_lookup(filter_func=is_in_partition)
235
+
236
+ # overloads
237
+ filter_target = True
238
+ pos_type = cache.cached_vec_type(self.geometry.dimension, dtype=float)
239
+
240
+ @cache.dynamic_func(suffix=self.name)
241
+ def cell_partition_lookup(args: self.DomainArg, pos: pos_type, max_dist: float):
242
+ return filtered_cell_lookup(args.geo, pos, max_dist, args.index, filter_target)
243
+
244
+ @cache.dynamic_func(suffix=self.name)
245
+ def cell_partition_lookup(args: self.DomainArg, pos: pos_type):
246
+ max_dist = 0.0
247
+ return filtered_cell_lookup(args.geo, pos, max_dist, args.index, filter_target)
248
+
249
+ return cell_partition_lookup
250
+
251
+ def supports_lookup(self, device):
252
+ return self.geometry.supports_cell_lookup(device)
253
+
183
254
  @property
184
255
  def domain_cell_arg(self) -> wp.Function:
185
256
  return Cells._identity_fn
@@ -200,6 +271,11 @@ class Sides(GeometryDomain):
200
271
  super().__init__(geometry)
201
272
 
202
273
  self.element_lookup = None
274
+ self.element_partition_lookup = None
275
+ self.element_filtered_lookup = None
276
+
277
+ def supports_lookup(self, device):
278
+ return False
203
279
 
204
280
  @property
205
281
  def element_kind(self) -> ElementKind:
@@ -225,6 +301,9 @@ class Sides(GeometryDomain):
225
301
  def element_index_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
226
302
  return self.geometry_partition.side_arg_value(device)
227
303
 
304
+ def fill_element_index_arg(self, arg: "ElementIndexArg", device: warp.context.Devicelike):
305
+ self.geometry_partition.fill_side_arg(arg, device)
306
+
228
307
  @property
229
308
  def element_index(self) -> wp.Function:
230
309
  return self.geometry_partition.side_index
@@ -236,6 +315,9 @@ class Sides(GeometryDomain):
236
315
  def element_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
237
316
  return self.geometry.side_arg_value(device)
238
317
 
318
+ def fill_element_arg(self, arg: "ElementArg", device: warp.context.Devicelike):
319
+ self.geometry.fill_side_arg(arg, device)
320
+
239
321
  @property
240
322
  def element_position(self) -> wp.Function:
241
323
  return self.geometry.side_position
@@ -256,6 +338,14 @@ class Sides(GeometryDomain):
256
338
  def element_normal(self) -> wp.Function:
257
339
  return self.geometry.side_normal
258
340
 
341
+ @property
342
+ def element_closest_point(self) -> wp.Function:
343
+ return self.geometry.side_closest_point
344
+
345
+ @property
346
+ def element_coordinates(self) -> wp.Function:
347
+ return self.geometry.side_coordinates
348
+
259
349
  @property
260
350
  def element_inner_cell_index(self) -> wp.Function:
261
351
  return self.geometry.side_inner_cell_index
@@ -276,9 +366,18 @@ class Sides(GeometryDomain):
276
366
  def cell_to_element_coords(self) -> wp.Function:
277
367
  return self.geometry.side_from_cell_coords
278
368
 
279
- @property
369
+ @cached_property
280
370
  def domain_cell_arg(self) -> wp.Function:
281
- return self.geometry.side_to_cell_arg
371
+ CellDomainArg = self.cell_domain().DomainArg
372
+
373
+ @cache.dynamic_func(suffix=self.name)
374
+ def domain_cell_arg(x: self.DomainArg):
375
+ return CellDomainArg(
376
+ self.geometry.side_to_cell_arg(x.geo),
377
+ self.geometry_partition.side_to_cell_arg(x.index),
378
+ )
379
+
380
+ return domain_cell_arg
282
381
 
283
382
  def cell_domain(self):
284
383
  return Cells(self.geometry_partition)
@@ -359,11 +458,13 @@ class Subdomain(GeometryDomain):
359
458
  self.geometry_element_count = self._domain.geometry_element_count
360
459
  self.reference_element = self._domain.reference_element
361
460
  self.element_arg_value = self._domain.element_arg_value
461
+ self.fill_element_arg = self._domain.fill_element_arg
362
462
  self.element_measure = self._domain.element_measure
363
463
  self.element_measure_ratio = self._domain.element_measure_ratio
364
464
  self.element_position = self._domain.element_position
365
465
  self.element_deformation_gradient = self._domain.element_deformation_gradient
366
466
  self.element_lookup = self._domain.element_lookup
467
+ self.element_partition_lookup = self._domain.element_partition_lookup
367
468
  self.element_normal = self._domain.element_normal
368
469
 
369
470
  @property
@@ -399,13 +500,42 @@ class Subdomain(GeometryDomain):
399
500
  @cache.cached_arg_value
400
501
  def element_index_arg_value(self, device: warp.context.Devicelike):
401
502
  arg = self.ElementIndexArg()
402
- arg.domain_arg = self._domain.element_index_arg_value(device)
403
- arg.element_indices = self._element_indices.to(device)
503
+ self.fill_element_index_arg(arg, device)
404
504
  return arg
405
505
 
506
+ def fill_element_index_arg(self, arg: "GeometryDomain.ElementIndexArg", device: warp.context.Devicelike):
507
+ self._domain.fill_element_index_arg(arg.domain_arg, device)
508
+ arg.element_indices = self._element_indices.to(device)
509
+
406
510
  def _make_element_index(self) -> wp.Function:
407
511
  @cache.dynamic_func(suffix=self.name)
408
512
  def element_index(arg: self.ElementIndexArg, index: int):
409
513
  return self._domain.element_index(arg.domain_arg, arg.element_indices[index])
410
514
 
411
515
  return element_index
516
+
517
+ def _make_element_partition_index(self) -> wp.Function:
518
+ @cache.dynamic_func(suffix=self.name)
519
+ def element_partition_index(arg: self.ElementIndexArg, element_index: int):
520
+ return self._domain.element_partition_index(arg.domain_arg, element_index)
521
+
522
+ return element_partition_index
523
+
524
+ def supports_lookup(self, device):
525
+ return self._domain.supports_lokup(device)
526
+
527
+ def cell_domain(self):
528
+ return self._domain.cell_domain()
529
+
530
+ @cached_property
531
+ def domain_cell_arg(self) -> wp.Function:
532
+ CellDomainArg = self.cell_domain().DomainArg
533
+
534
+ @cache.dynamic_func(suffix=self.name)
535
+ def domain_cell_arg(x: self.DomainArg):
536
+ return CellDomainArg(
537
+ self.geometry.side_to_cell_arg(x.geo),
538
+ self.geometry_partition.side_to_cell_arg(x.index.domain_arg),
539
+ )
540
+
541
+ return domain_cell_arg
warp/fem/field/field.py CHANGED
@@ -13,7 +13,7 @@
13
13
  # See the License for the specific language governing permissions and
14
14
  # limitations under the License.
15
15
 
16
- from typing import Any, Dict, Optional, Set
16
+ from typing import Any, ClassVar, Dict, Optional, Set
17
17
 
18
18
  import warp as wp
19
19
  from warp.fem import cache
@@ -22,6 +22,7 @@ from warp.fem.geometry import DeformedGeometry, Geometry
22
22
  from warp.fem.operator import Operator, integrand
23
23
  from warp.fem.space import FunctionSpace, SpacePartition
24
24
  from warp.fem.types import NULL_ELEMENT_INDEX, ElementKind, Sample
25
+ from warp.fem.utils import type_zero_element
25
26
 
26
27
 
27
28
  class FieldLike:
@@ -37,6 +38,10 @@ class FieldLike:
37
38
  """Value of the field-level arguments to be passed to device functions"""
38
39
  raise NotImplementedError
39
40
 
41
+ def fill_eval_arg(self, arg: "FieldLike.EvalArg", device):
42
+ """Fill the field-level arguments to be passed to device functions"""
43
+ raise NotImplementedError
44
+
40
45
  @property
41
46
  def degree(self) -> int:
42
47
  """Polynomial degree of the field, used to estimate necessary quadrature order"""
@@ -50,10 +55,6 @@ class FieldLike:
50
55
  def __str__(self) -> str:
51
56
  return self.name
52
57
 
53
- def eval_arg_value(self, device):
54
- """Value of arguments to be passed to device functions"""
55
- raise NotImplementedError
56
-
57
58
  def gradient_valid(self) -> bool:
58
59
  """Whether the gradient operator is implemented for this field."""
59
60
  return False
@@ -139,6 +140,51 @@ class GeometryField(FieldLike):
139
140
  """
140
141
  return DeformedGeometry(self, relative=relative)
141
142
 
143
+ @property
144
+ def gradient_dtype(self):
145
+ """Return type of the (world space) gradient operator. Assumes self.gradient_valid()"""
146
+ if wp.types.type_is_matrix(self.dtype):
147
+ return None
148
+
149
+ if wp.types.type_is_vector(self.dtype):
150
+ return cache.cached_mat_type(
151
+ shape=(wp.types.type_size(self.dtype), self.geometry.dimension),
152
+ dtype=wp.types.type_scalar_type(self.dtype),
153
+ )
154
+ if wp.types.type_is_quaternion(self.dtype):
155
+ return cache.cached_mat_type(
156
+ shape=(4, self.geometry.dimension),
157
+ dtype=wp.types.type_scalar_type(self.dtype),
158
+ )
159
+ return cache.cached_vec_type(length=self.geometry.dimension, dtype=wp.types.type_scalar_type(self.dtype))
160
+
161
+ @property
162
+ def reference_gradient_dtype(self):
163
+ """Return type of the reference space gradient operator. Assumes self.gradient_valid()"""
164
+ if wp.types.type_is_matrix(self.dtype):
165
+ return None
166
+
167
+ if wp.types.type_is_vector(self.dtype):
168
+ return cache.cached_mat_type(
169
+ shape=(wp.types.type_size(self.dtype), self.geometry.cell_dimension),
170
+ dtype=wp.types.type_scalar_type(self.dtype),
171
+ )
172
+ if wp.types.type_is_quaternion(self.dtype):
173
+ return cache.cached_mat_type(
174
+ shape=(4, self.geometry.cell_dimension),
175
+ dtype=wp.types.type_scalar_type(self.dtype),
176
+ )
177
+ return cache.cached_vec_type(length=self.geometry.cell_dimension, dtype=wp.types.type_scalar_type(self.dtype))
178
+
179
+ @property
180
+ def divergence_dtype(self):
181
+ """Return type of the divergence operator. Assumes self.divergence_valid()"""
182
+ if wp.types.type_is_vector(self.dtype):
183
+ return wp.types.type_scalar_type(self.dtype)
184
+ if wp.types.type_is_matrix(self.dtype):
185
+ return cache.cached_vec_type(length=self.dtype._shape_[1], dtype=wp.types.type_scalar_type(self.dtype))
186
+ return None
187
+
142
188
 
143
189
  class SpaceField(GeometryField):
144
190
  """Base class for fields defined over a function space"""
@@ -178,33 +224,6 @@ class SpaceField(GeometryField):
178
224
  def dof_dtype(self) -> type:
179
225
  return self.space.dof_dtype
180
226
 
181
- @property
182
- def gradient_dtype(self):
183
- """Return type of the (world space) gradient operator. Assumes self.gradient_valid()"""
184
- if wp.types.type_is_vector(self.dtype):
185
- return cache.cached_mat_type(
186
- shape=(wp.types.type_length(self.dtype), self.geometry.dimension),
187
- dtype=wp.types.type_scalar_type(self.dtype),
188
- )
189
- return cache.cached_vec_type(length=self.geometry.dimension, dtype=wp.types.type_scalar_type(self.dtype))
190
-
191
- @property
192
- def reference_gradient_dtype(self):
193
- """Return type of the reference space gradient operator. Assumes self.gradient_valid()"""
194
- if wp.types.type_is_vector(self.dtype):
195
- return cache.cached_mat_type(
196
- shape=(wp.types.type_length(self.dtype), self.geometry.cell_dimension),
197
- dtype=wp.types.type_scalar_type(self.dtype),
198
- )
199
- return cache.cached_vec_type(length=self.geometry.cell_dimension, dtype=wp.types.type_scalar_type(self.dtype))
200
-
201
- @property
202
- def divergence_dtype(self):
203
- """Return type of the divergence operator. Assumes self.gradient_valid()"""
204
- if wp.types.type_is_vector(self.dtype):
205
- return wp.types.type_scalar_type(self.dtype)
206
- return cache.cached_vec_type(length=self.dtype._shape_[1], dtype=wp.types.type_scalar_type(self.dtype))
207
-
208
227
  def _make_eval_degree(self):
209
228
  ORDER = self.space.ORDER
210
229
 
@@ -251,6 +270,19 @@ class ImplicitField(GeometryField):
251
270
  degree: Optional hint for automatic determination of quadrature orders when integrating this field
252
271
  """
253
272
 
273
+ _dynamic_attribute_constructors: ClassVar = {
274
+ "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
275
+ "eval_degree": lambda obj: obj._make_eval_degree(),
276
+ "eval_inner": lambda obj: obj._make_eval_func(obj._func),
277
+ "eval_grad_inner": lambda obj: obj._make_eval_func(obj._grad_func),
278
+ "eval_div_inner": lambda obj: obj._make_eval_func(obj._div_func),
279
+ "eval_reference_grad_inner": lambda obj: obj._make_eval_reference_grad(),
280
+ "eval_outer": lambda obj: obj.eval_inner,
281
+ "eval_grad_outer": lambda obj: obj.eval_grad_inner,
282
+ "eval_div_outer": lambda obj: obj.eval_div_inner,
283
+ "eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
284
+ }
285
+
254
286
  def __init__(
255
287
  self,
256
288
  domain: GeometryDomain,
@@ -284,18 +316,7 @@ class ImplicitField(GeometryField):
284
316
  self.EvalArg = cache.get_argument_struct(arg_types)
285
317
  self.values = values
286
318
 
287
- self.ElementEvalArg = self._make_element_eval_arg()
288
- self.eval_degree = self._make_eval_degree()
289
-
290
- self.eval_inner = self._make_eval_func(func)
291
- self.eval_grad_inner = self._make_eval_func(grad_func)
292
- self.eval_div_inner = self._make_eval_func(div_func)
293
- self.eval_reference_grad_inner = self._make_eval_reference_grad()
294
-
295
- self.eval_outer = self.eval_inner
296
- self.eval_grad_outer = self.eval_grad_inner
297
- self.eval_div_outer = self.eval_div_inner
298
- self.eval_reference_grad_outer = self.eval_reference_grad_inner
319
+ cache.setup_dynamic_attributes(self)
299
320
 
300
321
  @property
301
322
  def values(self):
@@ -303,6 +324,7 @@ class ImplicitField(GeometryField):
303
324
 
304
325
  @values.setter
305
326
  def values(self, v):
327
+ self._values = v
306
328
  self._func_arg = cache.populate_argument_struct(self.EvalArg, v, self._func.func.__name__)
307
329
 
308
330
  @property
@@ -316,6 +338,9 @@ class ImplicitField(GeometryField):
316
338
  def eval_arg_value(self, device):
317
339
  return self._func_arg
318
340
 
341
+ def fill_eval_arg(self, arg, device):
342
+ cache.populate_argument_struct(self.EvalArg, self._values, self._func.func.__name__, arg)
343
+
319
344
  @property
320
345
  def degree(self) -> int:
321
346
  return self._degree
@@ -324,6 +349,12 @@ class ImplicitField(GeometryField):
324
349
  def name(self) -> str:
325
350
  return f"Implicit_{self.domain.name}_{self.degree}_{self.EvalArg.key}"
326
351
 
352
+ def gradient_valid(self) -> bool:
353
+ return self._grad_func is not None
354
+
355
+ def divergence_valid(self) -> bool:
356
+ return self._div_func is not None
357
+
327
358
  def _make_eval_func(self, func):
328
359
  if func is None:
329
360
  return None
@@ -387,6 +418,20 @@ class UniformField(GeometryField):
387
418
  value: Uniform value over the domain
388
419
  """
389
420
 
421
+ _dynamic_attribute_constructors: ClassVar = {
422
+ "EvalArg": lambda obj: obj._make_eval_arg(),
423
+ "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
424
+ "eval_degree": lambda obj: obj._make_eval_degree(),
425
+ "eval_inner": lambda obj: obj._make_eval_inner(),
426
+ "eval_grad_inner": lambda obj: obj._make_eval_zero(obj.gradient_dtype),
427
+ "eval_div_inner": lambda obj: obj._make_eval_zero(obj.divergence_dtype),
428
+ "eval_reference_grad_inner": lambda obj: obj._make_eval_zero(obj.reference_gradient_dtype),
429
+ "eval_outer": lambda obj: obj.eval_inner,
430
+ "eval_grad_outer": lambda obj: obj.eval_grad_inner,
431
+ "eval_div_outer": lambda obj: obj.eval_div_inner,
432
+ "eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
433
+ }
434
+
390
435
  def __init__(self, domain: GeometryDomain, value: Any):
391
436
  self.domain = domain
392
437
 
@@ -396,30 +441,7 @@ class UniformField(GeometryField):
396
441
  self.dtype = wp.types.type_to_warp(type(value))
397
442
  self._value = self.dtype(value)
398
443
 
399
- scalar_type = wp.types.type_scalar_type(self.dtype)
400
- if wp.types.type_is_vector(self.dtype):
401
- grad_type = wp.mat(shape=(wp.types.type_length(self.dtype), self.geometry.dimension), dtype=scalar_type)
402
- div_type = scalar_type
403
- elif wp.types.type_is_matrix(self.dtype):
404
- grad_type = None
405
- div_type = wp.vec(length=(wp.types.type_length(self.dtype) // self.geometry.dimension), dtype=scalar_type)
406
- else:
407
- div_type = None
408
- grad_type = wp.vec(length=self.geometry.dimension, dtype=scalar_type)
409
-
410
- self.EvalArg = self._make_eval_arg()
411
- self.ElementEvalArg = self._make_element_eval_arg()
412
- self.eval_degree = self._make_eval_degree()
413
-
414
- self.eval_inner = self._make_eval_inner()
415
- self.eval_grad_inner = self._make_eval_zero(grad_type)
416
- self.eval_div_inner = self._make_eval_zero(div_type)
417
- self.eval_reference_grad_inner = self.eval_grad_inner
418
-
419
- self.eval_outer = self.eval_inner
420
- self.eval_grad_outer = self.eval_grad_inner
421
- self.eval_div_outer = self.eval_div_inner
422
- self.eval_reference_grad_outer = self.eval_reference_grad_inner
444
+ cache.setup_dynamic_attributes(self)
423
445
 
424
446
  @property
425
447
  def value(self):
@@ -444,10 +466,19 @@ class UniformField(GeometryField):
444
466
  arg.value = self.value
445
467
  return arg
446
468
 
469
+ def fill_eval_arg(self, arg, device):
470
+ arg.value = self.value
471
+
447
472
  @property
448
473
  def degree(self) -> int:
449
474
  return 0
450
475
 
476
+ def gradient_valid(self) -> bool:
477
+ return self.gradient_dtype is not None
478
+
479
+ def divergence_valid(self) -> bool:
480
+ return self.divergence_dtype is not None
481
+
451
482
  @property
452
483
  def name(self) -> str:
453
484
  return f"Uniform{self.domain.name}_{wp.types.get_type_code(self.dtype)}"
@@ -463,11 +494,11 @@ class UniformField(GeometryField):
463
494
  if dtype is None:
464
495
  return None
465
496
 
466
- scalar_type = wp.types.type_scalar_type(dtype)
497
+ zero_element = type_zero_element(dtype)
467
498
 
468
499
  @cache.dynamic_func(suffix=f"{self.name}_{wp.types.get_type_code(dtype)}")
469
500
  def eval_zero(args: self.ElementEvalArg, s: Sample):
470
- return dtype(scalar_type(0.0))
501
+ return zero_element()
471
502
 
472
503
  return eval_zero
473
504
 
@@ -511,6 +542,20 @@ class NonconformingField(GeometryField):
511
542
 
512
543
  _LOOKUP_EPS = wp.constant(1.0e-6)
513
544
 
545
+ _dynamic_attribute_constructors: ClassVar = {
546
+ "EvalArg": lambda obj: obj._make_eval_arg(),
547
+ "ElementEvalArg": lambda obj: obj._make_element_eval_arg(),
548
+ "eval_degree": lambda obj: obj._make_eval_degree(),
549
+ "eval_inner": lambda obj: obj._make_nonconforming_eval("eval_inner"),
550
+ "eval_grad_inner": lambda obj: obj._make_nonconforming_eval("eval_grad_inner"),
551
+ "eval_div_inner": lambda obj: obj._make_nonconforming_eval("eval_div_inner"),
552
+ "eval_reference_grad_inner": lambda obj: obj._make_eval_reference_grad(),
553
+ "eval_outer": lambda obj: obj.eval_inner,
554
+ "eval_grad_outer": lambda obj: obj.eval_grad_inner,
555
+ "eval_div_outer": lambda obj: obj.eval_div_inner,
556
+ "eval_reference_grad_outer": lambda obj: obj.eval_reference_grad_inner,
557
+ }
558
+
514
559
  def __init__(self, domain: GeometryDomain, field: DiscreteField, background: Any = 0.0):
515
560
  self.domain = domain
516
561
 
@@ -523,20 +568,7 @@ class NonconformingField(GeometryField):
523
568
  raise ValueError("Background field must be conforming to the domain")
524
569
  self.background = background
525
570
 
526
- self.EvalArg = self._make_eval_arg()
527
- self.ElementEvalArg = self._make_element_eval_arg()
528
- self.eval_degree = self._make_eval_degree()
529
-
530
- self.eval_inner = self._make_nonconforming_eval("eval_inner")
531
- self.eval_grad_inner = self._make_nonconforming_eval("eval_grad_inner")
532
- self.eval_div_inner = self._make_nonconforming_eval("eval_div_inner")
533
- self.eval_reference_grad_inner = self._make_eval_reference_grad()
534
-
535
- # Nonconforming evaluation is position based, does not handle discontinuous fields
536
- self.eval_outer = self.eval_inner
537
- self.eval_grad_outer = self.eval_grad_inner
538
- self.eval_div_outer = self.eval_div_inner
539
- self.eval_reference_grad_outer = self.eval_reference_grad_inner
571
+ cache.setup_dynamic_attributes(self)
540
572
 
541
573
  @property
542
574
  def geometry(self) -> Geometry:
@@ -546,19 +578,26 @@ class NonconformingField(GeometryField):
546
578
  def element_kind(self) -> ElementKind:
547
579
  return self.domain.element_kind
548
580
 
549
- @cache.cached_arg_value
550
581
  def eval_arg_value(self, device):
551
582
  arg = self.EvalArg()
552
- arg.field_cell_eval_arg = self.field.ElementEvalArg()
553
- arg.field_cell_eval_arg.elt_arg = self.field.geometry.cell_arg_value(device)
554
- arg.field_cell_eval_arg.eval_arg = self.field.eval_arg_value(device)
555
- arg.background_arg = self.background.eval_arg_value(device)
583
+ self.fill_eval_arg(arg, device)
556
584
  return arg
557
585
 
586
+ def fill_eval_arg(self, arg, device):
587
+ self.field.fill_eval_arg(arg.field_cell_eval_arg.eval_arg, device)
588
+ self.field.geometry.fill_cell_arg(arg.field_cell_eval_arg.elt_arg, device)
589
+ self.background.fill_eval_arg(arg.background_arg, device)
590
+
558
591
  @property
559
592
  def degree(self) -> int:
560
593
  return self.field.degree
561
594
 
595
+ def gradient_valid(self) -> bool:
596
+ return self.field.gradient_valid() and self.background.gradient_valid()
597
+
598
+ def divergence_valid(self) -> bool:
599
+ return self.field.divergence_valid() and self.background.divergence_valid()
600
+
562
601
  @property
563
602
  def name(self) -> str:
564
603
  return f"{self.domain.name}_{self.field.name}_{self.background.name}"
@@ -570,20 +609,23 @@ class NonconformingField(GeometryField):
570
609
  if field_eval is None or bg_eval is None:
571
610
  return None
572
611
 
612
+ cell_lookup = self.field.geometry.cell_lookup
613
+
573
614
  @cache.dynamic_func(suffix=f"{eval_func_name}_{self.name}")
574
615
  def eval_nc(args: self.ElementEvalArg, s: Sample):
575
616
  pos = self.domain.element_position(args.elt_arg, s)
576
617
  cell_arg = args.eval_arg.field_cell_eval_arg.elt_arg
577
- nonconforming_s = self.field.geometry.cell_lookup(cell_arg, pos)
578
- if (
579
- nonconforming_s.element_index == NULL_ELEMENT_INDEX
580
- or wp.length_sq(pos - self.field.geometry.cell_position(cell_arg, nonconforming_s))
581
- > NonconformingField._LOOKUP_EPS
582
- ):
583
- return bg_eval(self.background.ElementEvalArg(args.elt_arg, args.eval_arg.background_arg), s)
584
- return field_eval(
585
- self.field.ElementEvalArg(cell_arg, args.eval_arg.field_cell_eval_arg.eval_arg), nonconforming_s
586
- )
618
+ nonconforming_s = cell_lookup(cell_arg, pos, NonconformingField._LOOKUP_EPS)
619
+ if nonconforming_s.element_index != NULL_ELEMENT_INDEX:
620
+ if (
621
+ wp.length_sq(pos - self.field.geometry.cell_position(cell_arg, nonconforming_s))
622
+ <= NonconformingField._LOOKUP_EPS
623
+ ):
624
+ return field_eval(
625
+ self.field.ElementEvalArg(cell_arg, args.eval_arg.field_cell_eval_arg.eval_arg), nonconforming_s
626
+ )
627
+
628
+ return bg_eval(self.background.ElementEvalArg(args.elt_arg, args.eval_arg.background_arg), s)
587
629
 
588
630
  return eval_nc
589
631