warp-lang 1.0.0b2__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.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.
Files changed (269) 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.so +0 -0
  57. warp/bin/warp.so +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/fem/field/discrete_field.py +0 -80
  257. warp/fem/space/nodal_function_space.py +0 -233
  258. warp/tests/test_all.py +0 -223
  259. warp/tests/test_array_scan.py +0 -60
  260. warp/tests/test_base.py +0 -208
  261. warp/tests/test_unresolved_func.py +0 -7
  262. warp/tests/test_unresolved_symbol.py +0 -7
  263. warp_lang-1.0.0b2.dist-info/METADATA +0 -26
  264. warp_lang-1.0.0b2.dist-info/RECORD +0 -378
  265. /warp/tests/{test_compile_consts_dummy.py → aux_test_compile_consts_dummy.py} +0 -0
  266. /warp/tests/{test_reference_reference.py → aux_test_reference_reference.py} +0 -0
  267. /warp/tests/{test_square.py → aux_test_square.py} +0 -0
  268. {warp_lang-1.0.0b2.dist-info → warp_lang-1.0.0b6.dist-info}/LICENSE.md +0 -0
  269. {warp_lang-1.0.0b2.dist-info → warp_lang-1.0.0b6.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,489 @@
1
+ from typing import Optional
2
+
3
+ import warp as wp
4
+
5
+ from warp.fem.types import ElementIndex, Coords, make_free_sample
6
+ from warp.fem.geometry import Geometry
7
+ from warp.fem.quadrature import Quadrature
8
+ from warp.fem import cache
9
+
10
+ from .topology import SpaceTopology, DiscontinuousSpaceTopology
11
+ from .shape import ShapeFunction
12
+
13
+
14
+ class BasisSpace:
15
+ """Interface class for defining a scalar-valued basis over a geometry.
16
+
17
+ A basis space makes it easy to define multiple function spaces sharing the same basis (and thus nodes) but with different valuation functions;
18
+ however, it is not a required ingredient of a function space.
19
+
20
+ See also: :func:`make_polynomial_basis_space`, :func:`make_collocated_function_space`
21
+ """
22
+
23
+ @wp.struct
24
+ class BasisArg:
25
+ """Argument structure to be passed to device functions"""
26
+
27
+ pass
28
+
29
+ def __init__(self, topology: SpaceTopology):
30
+ self._topology = topology
31
+
32
+ self.NODES_PER_ELEMENT = self._topology.NODES_PER_ELEMENT
33
+
34
+ @property
35
+ def topology(self) -> SpaceTopology:
36
+ """Underlying topology of the basis space"""
37
+ return self._topology
38
+
39
+ @property
40
+ def geometry(self) -> Geometry:
41
+ """Underlying geometry of the basis space"""
42
+ return self._topology.geometry
43
+
44
+ def basis_arg_value(self, device) -> "BasisArg":
45
+ """Value for the argument structure to be passed to device functions"""
46
+ return BasisSpace.BasisArg()
47
+
48
+ # Helpers for generating node positions
49
+
50
+ def node_positions(self, out: Optional[wp.array] = None) -> wp.array:
51
+ """Returns a temporary array containing the world position for each node"""
52
+
53
+ NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
54
+
55
+ pos_type = cache.cached_vec_type(length=self.geometry.dimension, dtype=float)
56
+
57
+ node_coords_in_element = self.make_node_coords_in_element()
58
+
59
+ @cache.dynamic_kernel(suffix=self.name, kernel_options={"max_unroll": 4, "enable_backward": False})
60
+ def fill_node_positions(
61
+ geo_cell_arg: self.geometry.CellArg,
62
+ basis_arg: self.BasisArg,
63
+ topo_arg: self.topology.TopologyArg,
64
+ node_positions: wp.array(dtype=pos_type),
65
+ ):
66
+ element_index = wp.tid()
67
+
68
+ for n in range(NODES_PER_ELEMENT):
69
+ node_index = self.topology.element_node_index(geo_cell_arg, topo_arg, element_index, n)
70
+ coords = node_coords_in_element(geo_cell_arg, basis_arg, element_index, n)
71
+
72
+ sample = make_free_sample(element_index, coords)
73
+ pos = self.geometry.cell_position(geo_cell_arg, sample)
74
+
75
+ node_positions[node_index] = pos
76
+
77
+ shape = (self.topology.node_count(),)
78
+ if out is None:
79
+ node_positions = wp.empty(
80
+ shape=shape,
81
+ dtype=pos_type,
82
+ )
83
+ else:
84
+ if out.shape != shape or not wp.types.types_equal(pos_type, out.dtype):
85
+ raise ValueError(
86
+ f"Out node positions array must have shape {shape} and data type {wp.types.type_repr(pos_type)}"
87
+ )
88
+ node_positions = out
89
+
90
+ wp.launch(
91
+ dim=self.geometry.cell_count(),
92
+ kernel=fill_node_positions,
93
+ inputs=[
94
+ self.geometry.cell_arg_value(device=node_positions.device),
95
+ self.basis_arg_value(device=node_positions.device),
96
+ self.topology.topo_arg_value(device=node_positions.device),
97
+ node_positions,
98
+ ],
99
+ )
100
+
101
+ return node_positions
102
+
103
+ def make_node_coords_in_element(self):
104
+ raise NotImplementedError()
105
+
106
+ def make_node_quadrature_weight(self):
107
+ raise NotImplementedError()
108
+
109
+ def make_element_inner_weight(self):
110
+ raise NotImplementedError()
111
+
112
+ def make_element_outer_weight(self):
113
+ return self.make_element_inner_weight()
114
+
115
+ def make_element_inner_weight_gradient(self):
116
+ raise NotImplementedError()
117
+
118
+ def make_element_outer_weight_gradient(self):
119
+ return self.make_element_inner_weight_gradient()
120
+
121
+ def make_trace_node_quadrature_weight(self):
122
+ raise NotImplementedError()
123
+
124
+ def trace(self) -> "TraceBasisSpace":
125
+ return TraceBasisSpace(self)
126
+
127
+
128
+ class ShapeBasisSpace(BasisSpace):
129
+ """Base class for defining shape-function-based basis spaces."""
130
+
131
+ def __init__(self, topology: SpaceTopology, shape: ShapeFunction):
132
+ super().__init__(topology)
133
+ self._shape = shape
134
+
135
+ self.ORDER = self._shape.ORDER
136
+
137
+ if hasattr(shape, "element_node_triangulation"):
138
+ self.node_triangulation = self._node_triangulation
139
+ if hasattr(shape, "element_node_tets"):
140
+ self.node_tets = self._node_tets
141
+ if hasattr(shape, "element_node_hexes"):
142
+ self.node_hexes = self._node_hexes
143
+
144
+ @property
145
+ def shape(self) -> ShapeFunction:
146
+ """Shape functions used for defining individual element basis"""
147
+ return self._shape
148
+
149
+ @property
150
+ def name(self):
151
+ return f"{self.topology.name}_{self._shape.name}"
152
+
153
+ def make_node_coords_in_element(self):
154
+ shape_node_coords_in_element = self._shape.make_node_coords_in_element()
155
+
156
+ @cache.dynamic_func(suffix=self.name)
157
+ def node_coords_in_element(
158
+ elt_arg: self.geometry.CellArg,
159
+ basis_arg: self.BasisArg,
160
+ element_index: ElementIndex,
161
+ node_index_in_elt: int,
162
+ ):
163
+ return shape_node_coords_in_element(node_index_in_elt)
164
+
165
+ return node_coords_in_element
166
+
167
+ def make_node_quadrature_weight(self):
168
+ shape_node_quadrature_weight = self._shape.make_node_quadrature_weight()
169
+
170
+ @cache.dynamic_func(suffix=self.name)
171
+ def node_quadrature_weight(
172
+ elt_arg: self.geometry.CellArg,
173
+ basis_arg: self.BasisArg,
174
+ element_index: ElementIndex,
175
+ node_index_in_elt: int,
176
+ ):
177
+ return shape_node_quadrature_weight(node_index_in_elt)
178
+
179
+ return node_quadrature_weight
180
+
181
+ def make_element_inner_weight(self):
182
+ shape_element_inner_weight = self._shape.make_element_inner_weight()
183
+
184
+ @cache.dynamic_func(suffix=self.name)
185
+ def element_inner_weight(
186
+ elt_arg: self.geometry.CellArg,
187
+ basis_arg: self.BasisArg,
188
+ element_index: ElementIndex,
189
+ coords: Coords,
190
+ node_index_in_elt: int,
191
+ ):
192
+ return shape_element_inner_weight(coords, node_index_in_elt)
193
+
194
+ return element_inner_weight
195
+
196
+ def make_element_inner_weight_gradient(self):
197
+ shape_element_inner_weight_gradient = self._shape.make_element_inner_weight_gradient()
198
+
199
+ @cache.dynamic_func(suffix=self.name)
200
+ def element_inner_weight_gradient(
201
+ elt_arg: self.geometry.CellArg,
202
+ basis_arg: self.BasisArg,
203
+ element_index: ElementIndex,
204
+ coords: Coords,
205
+ node_index_in_elt: int,
206
+ ):
207
+ return shape_element_inner_weight_gradient(coords, node_index_in_elt)
208
+
209
+ return element_inner_weight_gradient
210
+
211
+ def make_trace_node_quadrature_weight(self, trace_basis):
212
+ shape_trace_node_quadrature_weight = self._shape.make_trace_node_quadrature_weight()
213
+
214
+ @cache.dynamic_func(suffix=self.name)
215
+ def trace_node_quadrature_weight(
216
+ geo_side_arg: trace_basis.geometry.SideArg,
217
+ basis_arg: trace_basis.BasisArg,
218
+ element_index: ElementIndex,
219
+ node_index_in_elt: int,
220
+ ):
221
+ neighbour_elem, index_in_neighbour = trace_basis.topology.neighbor_cell_index(
222
+ geo_side_arg, element_index, node_index_in_elt
223
+ )
224
+ return shape_trace_node_quadrature_weight(index_in_neighbour)
225
+
226
+ return trace_node_quadrature_weight
227
+
228
+ def _node_triangulation(self):
229
+ element_node_indices = self._topology.element_node_indices().numpy()
230
+ element_triangles = self._shape.element_node_triangulation()
231
+
232
+ tri_indices = element_node_indices[:, element_triangles].reshape(-1, 3)
233
+ return tri_indices
234
+
235
+ def _node_tets(self):
236
+ element_node_indices = self._topology.element_node_indices().numpy()
237
+ element_tets = self._shape.element_node_tets()
238
+
239
+ tet_indices = element_node_indices[:, element_tets].reshape(-1, 4)
240
+ return tet_indices
241
+
242
+ def _node_hexes(self):
243
+ element_node_indices = self._topology.element_node_indices().numpy()
244
+ element_hexes = self._shape.element_node_hexes()
245
+
246
+ hex_indices = element_node_indices[:, element_hexes].reshape(-1, 8)
247
+ return hex_indices
248
+
249
+
250
+ class TraceBasisSpace(BasisSpace):
251
+ """Auto-generated trace space evaluating the cell-defined basis on the geometry sides"""
252
+
253
+ def __init__(self, basis: BasisSpace):
254
+ super().__init__(basis.topology.trace())
255
+
256
+ self.ORDER = basis.ORDER
257
+
258
+ self._basis = basis
259
+ self.BasisArg = self._basis.BasisArg
260
+ self.basis_arg_value = self._basis.basis_arg_value
261
+
262
+ @property
263
+ def name(self):
264
+ return f"{self._basis.name}_Trace"
265
+
266
+ def make_node_coords_in_element(self):
267
+ node_coords_in_cell = self._basis.make_node_coords_in_element()
268
+
269
+ @cache.dynamic_func(suffix=self._basis.name)
270
+ def trace_node_coords_in_element(
271
+ geo_side_arg: self.geometry.SideArg,
272
+ basis_arg: self.BasisArg,
273
+ element_index: ElementIndex,
274
+ node_index_in_elt: int,
275
+ ):
276
+ neighbour_elem, index_in_neighbour = self.topology.neighbor_cell_index(
277
+ geo_side_arg, element_index, node_index_in_elt
278
+ )
279
+ geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
280
+ neighbour_coords = node_coords_in_cell(
281
+ geo_cell_arg,
282
+ basis_arg,
283
+ neighbour_elem,
284
+ index_in_neighbour,
285
+ )
286
+
287
+ return self.geometry.side_from_cell_coords(geo_side_arg, element_index, neighbour_elem, neighbour_coords)
288
+
289
+ return trace_node_coords_in_element
290
+
291
+ def make_node_quadrature_weight(self):
292
+ return self._basis.make_trace_node_quadrature_weight(self)
293
+
294
+ def make_element_inner_weight(self):
295
+ cell_inner_weight = self._basis.make_element_inner_weight()
296
+
297
+ @cache.dynamic_func(suffix=self._basis.name)
298
+ def trace_element_inner_weight(
299
+ geo_side_arg: self.geometry.SideArg,
300
+ basis_arg: self.BasisArg,
301
+ element_index: ElementIndex,
302
+ coords: Coords,
303
+ node_index_in_elt: int,
304
+ ):
305
+ cell_index, index_in_cell = self.topology.inner_cell_index(geo_side_arg, element_index, node_index_in_elt)
306
+ if index_in_cell < 0:
307
+ return 0.0
308
+
309
+ cell_coords = self.geometry.side_inner_cell_coords(geo_side_arg, element_index, coords)
310
+
311
+ geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
312
+ return cell_inner_weight(
313
+ geo_cell_arg,
314
+ basis_arg,
315
+ cell_index,
316
+ cell_coords,
317
+ index_in_cell,
318
+ )
319
+
320
+ return trace_element_inner_weight
321
+
322
+ def make_element_outer_weight(self):
323
+ cell_outer_weight = self._basis.make_element_outer_weight()
324
+
325
+ @cache.dynamic_func(suffix=self._basis.name)
326
+ def trace_element_outer_weight(
327
+ geo_side_arg: self.geometry.SideArg,
328
+ basis_arg: self.BasisArg,
329
+ element_index: ElementIndex,
330
+ coords: Coords,
331
+ node_index_in_elt: int,
332
+ ):
333
+ cell_index, index_in_cell = self.topology.outer_cell_index(geo_side_arg, element_index, node_index_in_elt)
334
+ if index_in_cell < 0:
335
+ return 0.0
336
+
337
+ cell_coords = self.geometry.side_outer_cell_coords(geo_side_arg, element_index, coords)
338
+
339
+ geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
340
+ return cell_outer_weight(
341
+ geo_cell_arg,
342
+ basis_arg,
343
+ cell_index,
344
+ cell_coords,
345
+ index_in_cell,
346
+ )
347
+
348
+ return trace_element_outer_weight
349
+
350
+ def make_element_inner_weight_gradient(self):
351
+ cell_inner_weight_gradient = self._basis.make_element_inner_weight_gradient()
352
+ grad_vec_type = wp.vec(length=self.geometry.dimension, dtype=float)
353
+
354
+ @cache.dynamic_func(suffix=self._basis.name)
355
+ def trace_element_inner_weight_gradient(
356
+ geo_side_arg: self.geometry.SideArg,
357
+ basis_arg: self.BasisArg,
358
+ element_index: ElementIndex,
359
+ coords: Coords,
360
+ node_index_in_elt: int,
361
+ ):
362
+ cell_index, index_in_cell = self.topology.inner_cell_index(geo_side_arg, element_index, node_index_in_elt)
363
+ if index_in_cell < 0:
364
+ return grad_vec_type(0.0)
365
+
366
+ cell_coords = self.geometry.side_inner_cell_coords(geo_side_arg, element_index, coords)
367
+ geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
368
+ return cell_inner_weight_gradient(geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell)
369
+
370
+ return trace_element_inner_weight_gradient
371
+
372
+ def make_element_outer_weight_gradient(self):
373
+ cell_outer_weight_gradient = self._basis.make_element_outer_weight_gradient()
374
+ grad_vec_type = wp.vec(length=self.geometry.dimension, dtype=float)
375
+
376
+ @cache.dynamic_func(suffix=self._basis.name)
377
+ def trace_element_outer_weight_gradient(
378
+ geo_side_arg: self.geometry.SideArg,
379
+ basis_arg: self.BasisArg,
380
+ element_index: ElementIndex,
381
+ coords: Coords,
382
+ node_index_in_elt: int,
383
+ ):
384
+ cell_index, index_in_cell = self.topology.outer_cell_index(geo_side_arg, element_index, node_index_in_elt)
385
+ if index_in_cell < 0:
386
+ return grad_vec_type(0.0)
387
+
388
+ cell_coords = self.geometry.side_outer_cell_coords(geo_side_arg, element_index, coords)
389
+ geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
390
+ return cell_outer_weight_gradient(geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell)
391
+
392
+ return trace_element_outer_weight_gradient
393
+
394
+ def __eq__(self, other: "TraceBasisSpace") -> bool:
395
+ return self._topo == other._topo
396
+
397
+
398
+ class PointBasisSpace(BasisSpace):
399
+ """An unstructured :class:`BasisSpace` that is non-zero at a finite set of points only.
400
+
401
+ The node locations and nodal quadrature weights are defined by a :class:`Quadrature` formula.
402
+ """
403
+
404
+ def __init__(self, quadrature: Quadrature):
405
+ self._quadrature = quadrature
406
+
407
+ if quadrature.points_per_element() is None:
408
+ raise NotImplementedError("Varying number of points per element is not supported yet")
409
+
410
+ topology = DiscontinuousSpaceTopology(
411
+ geometry=quadrature.domain.geometry, nodes_per_element=quadrature.points_per_element()
412
+ )
413
+ super().__init__(topology)
414
+
415
+ self.BasisArg = quadrature.Arg
416
+ self.basis_arg_value = quadrature.arg_value
417
+ self.ORDER = 0
418
+
419
+ self.make_element_outer_weight = self.make_element_inner_weight
420
+ self.make_element_outer_weight_gradient = self.make_element_outer_weight_gradient
421
+
422
+ @property
423
+ def name(self):
424
+ return f"{self._quadrature.name}_Point"
425
+
426
+ def make_node_coords_in_element(self):
427
+ @cache.dynamic_func(suffix=self.name)
428
+ def node_coords_in_element(
429
+ elt_arg: self._quadrature.domain.ElementArg,
430
+ basis_arg: self.BasisArg,
431
+ element_index: ElementIndex,
432
+ node_index_in_elt: int,
433
+ ):
434
+ return self._quadrature.point_coords(elt_arg, basis_arg, element_index, node_index_in_elt)
435
+
436
+ return node_coords_in_element
437
+
438
+ def make_node_quadrature_weight(self):
439
+ @cache.dynamic_func(suffix=self.name)
440
+ def node_quadrature_weight(
441
+ elt_arg: self._quadrature.domain.ElementArg,
442
+ basis_arg: self.BasisArg,
443
+ element_index: ElementIndex,
444
+ node_index_in_elt: int,
445
+ ):
446
+ return self._quadrature.point_weight(elt_arg, basis_arg, element_index, node_index_in_elt)
447
+
448
+ return node_quadrature_weight
449
+
450
+ def make_element_inner_weight(self):
451
+ @cache.dynamic_func(suffix=self.name)
452
+ def element_inner_weight(
453
+ elt_arg: self._quadrature.domain.ElementArg,
454
+ basis_arg: self.BasisArg,
455
+ element_index: ElementIndex,
456
+ coords: Coords,
457
+ node_index_in_elt: int,
458
+ ):
459
+ qp_coord = self._quadrature.point_coords(elt_arg, basis_arg, element_index, node_index_in_elt)
460
+ return wp.select(wp.length_sq(coords - qp_coord) < 0.001, 0.0, 1.0)
461
+
462
+ return element_inner_weight
463
+
464
+ def make_element_inner_weight_gradient(self):
465
+ gradient_vec = cache.cached_vec_type(length=self.geometry.dimension, dtype=float)
466
+
467
+ @cache.dynamic_func(suffix=self.name)
468
+ def element_inner_weight_gradient(
469
+ elt_arg: self._quadrature.domain.ElementArg,
470
+ basis_arg: self.BasisArg,
471
+ element_index: ElementIndex,
472
+ coords: Coords,
473
+ node_index_in_elt: int,
474
+ ):
475
+ return gradient_vec(0.0)
476
+
477
+ return element_inner_weight_gradient
478
+
479
+ def make_trace_node_quadrature_weight(self, trace_basis):
480
+ @cache.dynamic_func(suffix=self.name)
481
+ def trace_node_quadrature_weight(
482
+ elt_arg: trace_basis.geometry.SideArg,
483
+ basis_arg: trace_basis.BasisArg,
484
+ element_index: ElementIndex,
485
+ node_index_in_elt: int,
486
+ ):
487
+ return 0.0
488
+
489
+ return trace_node_quadrature_weight
@@ -0,0 +1,105 @@
1
+ from typing import Any, Optional
2
+
3
+ import warp as wp
4
+
5
+ from warp.fem.types import DofIndex, get_node_coord
6
+ from warp.fem.geometry import GeometryPartition
7
+ from warp.fem import cache, utils
8
+
9
+
10
+ from .function_space import FunctionSpace
11
+ from .dof_mapper import DofMapper, IdentityMapper
12
+ from .partition import make_space_partition, SpacePartition
13
+
14
+ from .basis_space import BasisSpace
15
+
16
+
17
+ class CollocatedFunctionSpace(FunctionSpace):
18
+ """Function space where values are collocated at nodes"""
19
+
20
+ def __init__(self, basis: BasisSpace, dtype: type = float, dof_mapper: DofMapper = None):
21
+ super().__init__(topology=basis.topology)
22
+
23
+ self.dof_mapper = IdentityMapper(dtype) if dof_mapper is None else dof_mapper
24
+ self.dtype = self.dof_mapper.value_dtype
25
+ self.dof_dtype = self.dof_mapper.dof_dtype
26
+ self.VALUE_DOF_COUNT = self.dof_mapper.DOF_SIZE
27
+
28
+ self._basis = basis
29
+ self.SpaceArg = self._basis.BasisArg
30
+
31
+ self.ORDER = self._basis.ORDER
32
+
33
+ self.unit_dof_value = self._make_unit_dof_value(self.dof_mapper)
34
+
35
+ self.node_coords_in_element = self._basis.make_node_coords_in_element()
36
+ self.node_quadrature_weight = self._basis.make_node_quadrature_weight()
37
+ self.element_inner_weight = self._basis.make_element_inner_weight()
38
+ self.element_inner_weight_gradient = self._basis.make_element_inner_weight_gradient()
39
+ self.element_outer_weight = self._basis.make_element_outer_weight()
40
+ self.element_outer_weight_gradient = self._basis.make_element_outer_weight_gradient()
41
+
42
+ # For backward compatibility
43
+ if hasattr(basis, "node_grid"):
44
+ self.node_grid = basis.node_grid
45
+ if hasattr(basis, "node_triangulation"):
46
+ self.node_triangulation = basis.node_triangulation
47
+ if hasattr(basis, "node_tets"):
48
+ self.node_tets = basis.node_tets
49
+ if hasattr(basis, "node_hexes"):
50
+ self.node_hexes = basis.node_hexes
51
+
52
+ def space_arg_value(self, device):
53
+ return self._basis.basis_arg_value(device)
54
+
55
+ @property
56
+ def name(self):
57
+ return f"{self._basis.name}_{self.dof_mapper}".replace(".", "_")
58
+
59
+ @property
60
+ def degree(self):
61
+ """Maximum polynomial degree of the underlying basis"""
62
+ return self.ORDER
63
+
64
+ def make_field(
65
+ self,
66
+ space_partition: Optional[SpacePartition] = None,
67
+ ) -> "wp.fem.field.NodalField":
68
+
69
+ from warp.fem.field import NodalField
70
+
71
+ if space_partition is None:
72
+ space_partition = make_space_partition(space_topology=self.topology)
73
+
74
+ return NodalField(space=self, space_partition=space_partition)
75
+
76
+ def _make_unit_dof_value(self, dof_mapper: DofMapper):
77
+ @cache.dynamic_func(suffix=self.name)
78
+ def unit_dof_value(geo_arg: self.topology.ElementArg, space_arg: self.SpaceArg, dof: DofIndex):
79
+ return dof_mapper.dof_to_value(utils.unit_element(dof_mapper.dof_dtype(0.0), get_node_coord(dof)))
80
+
81
+ return unit_dof_value
82
+
83
+ def node_count(self):
84
+ return self.topology.node_count()
85
+
86
+ def node_positions(self, out:Optional[wp.array] = None) -> wp.array:
87
+ return self._basis.node_positions(out=out)
88
+
89
+ def trace(self) -> "CollocatedFunctionSpace":
90
+ return CollocatedFunctionSpaceTrace(self)
91
+
92
+
93
+ class CollocatedFunctionSpaceTrace(CollocatedFunctionSpace):
94
+ """Trace of a :class:`CollocatedFunctionSpace`"""
95
+
96
+ def __init__(self, space: CollocatedFunctionSpace):
97
+ self._space = space
98
+ super().__init__(space._basis.trace(), space.dtype, space.dof_mapper)
99
+
100
+ @property
101
+ def name(self):
102
+ return f"{self._space.name}_Trace"
103
+
104
+ def __eq__(self, other: "CollocatedFunctionSpaceTrace") -> bool:
105
+ return self._space == other._space
@@ -5,7 +5,7 @@ import math
5
5
  import warp as wp
6
6
  import warp.types
7
7
 
8
- from warp.fem.types import vec6
8
+ vec6 = wp.types.vector(length=6, dtype=wp.float32)
9
9
 
10
10
  _SQRT_2 = wp.constant(math.sqrt(2.0))
11
11
  _SQRT_3 = wp.constant(math.sqrt(3.0))
@@ -65,7 +65,7 @@ class SymmetricTensorMapper(DofMapper):
65
65
  first the three diagonal terms, then off-diagonal coefficients"""
66
66
  DB16 = 1
67
67
  """Ordering that also separates normal from tangential coefficients:
68
- first trace, then other diagonal terms, then off-diagonal coefficients.
68
+ first trace, then other diagonal terms, then off-diagonal coefficients.
69
69
  See [Daviet and Bertails-Descoubes 2016]"""
70
70
 
71
71
  def __init__(self, dtype: type, mapping: Mapping = Mapping.VOIGT):
@@ -187,3 +187,50 @@ class SymmetricTensorMapper(DofMapper):
187
187
  f = 0.5 * (val[1, 0] + val[0, 1])
188
188
 
189
189
  return vec6(a, b, c, d, e, f)
190
+
191
+
192
+ class SkewSymmetricTensorMapper(DofMapper):
193
+ """Orthonormal isomorphism from R^{n (n-1)} to nxn skew-symmetric tensors,
194
+ using usual L2 norm for vectors and half Frobenius norm, (tau : tau)/2 for tensors.
195
+ """
196
+
197
+ def __init__(self, dtype: type):
198
+ self.value_dtype = dtype
199
+
200
+ if dtype == wp.mat22:
201
+ self.dof_dtype = float
202
+ self.DOF_SIZE = wp.constant(1)
203
+ self.dof_to_value = SkewSymmetricTensorMapper.dof_to_value_2d
204
+ self.value_to_dof = SkewSymmetricTensorMapper.value_to_dof_2d
205
+ elif dtype == wp.mat33:
206
+ self.dof_dtype = wp.vec3
207
+ self.DOF_SIZE = wp.constant(3)
208
+ self.dof_to_value = SkewSymmetricTensorMapper.dof_to_value_3d
209
+ self.value_to_dof = SkewSymmetricTensorMapper.value_to_dof_3d
210
+ else:
211
+ raise ValueError("Unsupported value dtype: ", dtype)
212
+
213
+ def __str__(self):
214
+ return f"{self.__class__.__name__}_{self.DOF_SIZE}"
215
+
216
+ @wp.func
217
+ def dof_to_value_2d(dof: float):
218
+ return wp.mat22(0.0, -dof, dof, 0.0)
219
+
220
+ @wp.func
221
+ def value_to_dof_2d(val: wp.mat22):
222
+ return 0.5 * (val[1, 0] - val[0, 1])
223
+
224
+ @wp.func
225
+ def dof_to_value_3d(dof: wp.vec3):
226
+ a = dof[0]
227
+ b = dof[1]
228
+ c = dof[2]
229
+ return wp.mat33(0.0, -c, b, c, 0.0, -a, -b, a, 0.0)
230
+
231
+ @wp.func
232
+ def value_to_dof_3d(val: wp.mat33):
233
+ a = 0.5 * (val[2, 1] - val[1, 2])
234
+ b = 0.5 * (val[0, 2] - val[2, 0])
235
+ c = 0.5 * (val[1, 0] - val[0, 1])
236
+ return wp.vec3(a, b, c)