warp-lang 1.4.2__py3-none-manylinux2014_x86_64.whl → 1.5.1__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.

Potentially problematic release.

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

Files changed (166) hide show
  1. warp/__init__.py +4 -0
  2. warp/autograd.py +43 -8
  3. warp/bin/warp-clang.so +0 -0
  4. warp/bin/warp.so +0 -0
  5. warp/build.py +21 -2
  6. warp/build_dll.py +23 -6
  7. warp/builtins.py +1819 -7
  8. warp/codegen.py +197 -61
  9. warp/config.py +2 -2
  10. warp/context.py +379 -107
  11. warp/examples/assets/pixel.jpg +0 -0
  12. warp/examples/benchmarks/benchmark_cloth_paddle.py +86 -0
  13. warp/examples/benchmarks/benchmark_gemm.py +121 -0
  14. warp/examples/benchmarks/benchmark_interop_paddle.py +158 -0
  15. warp/examples/benchmarks/benchmark_tile.py +179 -0
  16. warp/examples/fem/example_adaptive_grid.py +37 -10
  17. warp/examples/fem/example_apic_fluid.py +3 -2
  18. warp/examples/fem/example_convection_diffusion_dg.py +4 -5
  19. warp/examples/fem/example_deformed_geometry.py +1 -1
  20. warp/examples/fem/example_diffusion_3d.py +47 -4
  21. warp/examples/fem/example_distortion_energy.py +220 -0
  22. warp/examples/fem/example_magnetostatics.py +127 -85
  23. warp/examples/fem/example_nonconforming_contact.py +5 -5
  24. warp/examples/fem/example_stokes.py +3 -1
  25. warp/examples/fem/example_streamlines.py +12 -19
  26. warp/examples/fem/utils.py +38 -15
  27. warp/examples/sim/example_cloth.py +4 -25
  28. warp/examples/sim/example_quadruped.py +2 -1
  29. warp/examples/tile/example_tile_convolution.py +58 -0
  30. warp/examples/tile/example_tile_fft.py +47 -0
  31. warp/examples/tile/example_tile_filtering.py +105 -0
  32. warp/examples/tile/example_tile_matmul.py +79 -0
  33. warp/examples/tile/example_tile_mlp.py +375 -0
  34. warp/fem/__init__.py +8 -0
  35. warp/fem/cache.py +16 -12
  36. warp/fem/dirichlet.py +1 -1
  37. warp/fem/domain.py +44 -1
  38. warp/fem/field/__init__.py +1 -2
  39. warp/fem/field/field.py +31 -19
  40. warp/fem/field/nodal_field.py +101 -49
  41. warp/fem/field/virtual.py +794 -0
  42. warp/fem/geometry/__init__.py +2 -2
  43. warp/fem/geometry/deformed_geometry.py +3 -105
  44. warp/fem/geometry/element.py +13 -0
  45. warp/fem/geometry/geometry.py +165 -7
  46. warp/fem/geometry/grid_2d.py +3 -6
  47. warp/fem/geometry/grid_3d.py +31 -28
  48. warp/fem/geometry/hexmesh.py +3 -46
  49. warp/fem/geometry/nanogrid.py +3 -2
  50. warp/fem/geometry/{quadmesh_2d.py → quadmesh.py} +280 -159
  51. warp/fem/geometry/tetmesh.py +2 -43
  52. warp/fem/geometry/{trimesh_2d.py → trimesh.py} +354 -186
  53. warp/fem/integrate.py +683 -261
  54. warp/fem/linalg.py +404 -0
  55. warp/fem/operator.py +101 -18
  56. warp/fem/polynomial.py +5 -5
  57. warp/fem/quadrature/quadrature.py +45 -21
  58. warp/fem/space/__init__.py +45 -11
  59. warp/fem/space/basis_function_space.py +451 -0
  60. warp/fem/space/basis_space.py +58 -11
  61. warp/fem/space/function_space.py +146 -5
  62. warp/fem/space/grid_2d_function_space.py +80 -66
  63. warp/fem/space/grid_3d_function_space.py +113 -68
  64. warp/fem/space/hexmesh_function_space.py +96 -108
  65. warp/fem/space/nanogrid_function_space.py +62 -110
  66. warp/fem/space/quadmesh_function_space.py +208 -0
  67. warp/fem/space/shape/__init__.py +45 -7
  68. warp/fem/space/shape/cube_shape_function.py +328 -54
  69. warp/fem/space/shape/shape_function.py +10 -1
  70. warp/fem/space/shape/square_shape_function.py +328 -60
  71. warp/fem/space/shape/tet_shape_function.py +269 -19
  72. warp/fem/space/shape/triangle_shape_function.py +238 -19
  73. warp/fem/space/tetmesh_function_space.py +69 -37
  74. warp/fem/space/topology.py +38 -0
  75. warp/fem/space/trimesh_function_space.py +179 -0
  76. warp/fem/utils.py +6 -331
  77. warp/jax_experimental.py +3 -1
  78. warp/native/array.h +15 -0
  79. warp/native/builtin.h +66 -26
  80. warp/native/bvh.h +4 -0
  81. warp/native/coloring.cpp +604 -0
  82. warp/native/cuda_util.cpp +68 -51
  83. warp/native/cuda_util.h +2 -1
  84. warp/native/fabric.h +8 -0
  85. warp/native/hashgrid.h +4 -0
  86. warp/native/marching.cu +8 -0
  87. warp/native/mat.h +14 -3
  88. warp/native/mathdx.cpp +59 -0
  89. warp/native/mesh.h +4 -0
  90. warp/native/range.h +13 -1
  91. warp/native/reduce.cpp +9 -1
  92. warp/native/reduce.cu +7 -0
  93. warp/native/runlength_encode.cpp +9 -1
  94. warp/native/runlength_encode.cu +7 -1
  95. warp/native/scan.cpp +8 -0
  96. warp/native/scan.cu +8 -0
  97. warp/native/scan.h +8 -1
  98. warp/native/sparse.cpp +8 -0
  99. warp/native/sparse.cu +8 -0
  100. warp/native/temp_buffer.h +7 -0
  101. warp/native/tile.h +1854 -0
  102. warp/native/tile_gemm.h +341 -0
  103. warp/native/tile_reduce.h +210 -0
  104. warp/native/volume_builder.cu +8 -0
  105. warp/native/volume_builder.h +8 -0
  106. warp/native/warp.cpp +10 -2
  107. warp/native/warp.cu +369 -15
  108. warp/native/warp.h +12 -2
  109. warp/optim/adam.py +39 -4
  110. warp/paddle.py +29 -12
  111. warp/render/render_opengl.py +140 -67
  112. warp/sim/graph_coloring.py +292 -0
  113. warp/sim/import_urdf.py +8 -8
  114. warp/sim/integrator_euler.py +4 -2
  115. warp/sim/integrator_featherstone.py +115 -44
  116. warp/sim/integrator_vbd.py +6 -0
  117. warp/sim/model.py +109 -32
  118. warp/sparse.py +1 -1
  119. warp/stubs.py +569 -4
  120. warp/tape.py +12 -7
  121. warp/tests/assets/pixel.npy +0 -0
  122. warp/tests/aux_test_instancing_gc.py +18 -0
  123. warp/tests/test_array.py +39 -0
  124. warp/tests/test_codegen.py +81 -1
  125. warp/tests/test_codegen_instancing.py +30 -0
  126. warp/tests/test_collision.py +110 -0
  127. warp/tests/test_coloring.py +251 -0
  128. warp/tests/test_context.py +34 -0
  129. warp/tests/test_examples.py +21 -5
  130. warp/tests/test_fem.py +453 -113
  131. warp/tests/test_func.py +34 -4
  132. warp/tests/test_generics.py +52 -0
  133. warp/tests/test_iter.py +68 -0
  134. warp/tests/test_lerp.py +13 -87
  135. warp/tests/test_mat_scalar_ops.py +1 -1
  136. warp/tests/test_matmul.py +6 -9
  137. warp/tests/test_matmul_lite.py +6 -11
  138. warp/tests/test_mesh_query_point.py +1 -1
  139. warp/tests/test_module_hashing.py +23 -0
  140. warp/tests/test_overwrite.py +45 -0
  141. warp/tests/test_paddle.py +27 -87
  142. warp/tests/test_print.py +56 -1
  143. warp/tests/test_smoothstep.py +17 -83
  144. warp/tests/test_spatial.py +1 -1
  145. warp/tests/test_static.py +3 -3
  146. warp/tests/test_tile.py +744 -0
  147. warp/tests/test_tile_mathdx.py +144 -0
  148. warp/tests/test_tile_mlp.py +383 -0
  149. warp/tests/test_tile_reduce.py +374 -0
  150. warp/tests/test_tile_shared_memory.py +190 -0
  151. warp/tests/test_vbd.py +12 -20
  152. warp/tests/test_volume.py +43 -0
  153. warp/tests/unittest_suites.py +19 -2
  154. warp/tests/unittest_utils.py +4 -2
  155. warp/types.py +340 -74
  156. warp/utils.py +23 -3
  157. {warp_lang-1.4.2.dist-info → warp_lang-1.5.1.dist-info}/METADATA +32 -7
  158. {warp_lang-1.4.2.dist-info → warp_lang-1.5.1.dist-info}/RECORD +161 -134
  159. {warp_lang-1.4.2.dist-info → warp_lang-1.5.1.dist-info}/WHEEL +1 -1
  160. warp/fem/field/test.py +0 -180
  161. warp/fem/field/trial.py +0 -183
  162. warp/fem/space/collocated_function_space.py +0 -102
  163. warp/fem/space/quadmesh_2d_function_space.py +0 -261
  164. warp/fem/space/trimesh_2d_function_space.py +0 -153
  165. {warp_lang-1.4.2.dist-info → warp_lang-1.5.1.dist-info}/LICENSE.md +0 -0
  166. {warp_lang-1.4.2.dist-info → warp_lang-1.5.1.dist-info}/top_level.txt +0 -0
warp/fem/field/nodal_field.py CHANGED
@@ -1,5 +1,7 @@
1
+ from typing import Any
2
+
1
3
  import warp as wp
2
- from warp.fem import cache, utils
4
+ from warp.fem import cache
3
5
  from warp.fem.space import CollocatedFunctionSpace, SpacePartition
4
6
  from warp.fem.types import NULL_NODE_INDEX, ElementIndex, Sample
5
7
 
@@ -56,58 +58,79 @@ class NodalFieldBase(DiscreteField):
56
58
  )
57
59
  pidx = self.space_partition.partition_node_index(args.eval_arg.partition_arg, nidx)
58
60
  if pidx == NULL_NODE_INDEX:
59
- return self.space.dtype(0.0)
61
+ return self.space.dof_dtype(0.0)
60
62
 
61
- return self.space.dof_mapper.dof_to_value(args.eval_arg.dof_values[pidx])
63
+ return args.eval_arg.dof_values[pidx]
62
64
 
63
65
  return read_node_value
64
66
 
65
67
  def _make_eval_inner(self):
66
68
  @cache.dynamic_func(suffix=self.name)
67
69
  def eval_inner(args: self.ElementEvalArg, s: Sample):
68
- res = self.space.element_inner_weight(
69
- args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
70
- ) * self._read_node_value(args, s.element_index, 0)
70
+ local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
71
+ res = self.space.space_value(
72
+ self._read_node_value(args, s.element_index, 0),
73
+ self.space.element_inner_weight(
74
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
75
+ ),
76
+ local_value_map,
77
+ )
78
+
71
79
  node_count = self.space.topology.element_node_count(
72
80
  args.elt_arg, args.eval_arg.topology_arg, s.element_index
73
81
  )
74
82
  for k in range(1, node_count):
75
- res += self.space.element_inner_weight(
76
- args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
77
- ) * self._read_node_value(args, s.element_index, k)
83
+ res += self.space.space_value(
84
+ self._read_node_value(args, s.element_index, k),
85
+ self.space.element_inner_weight(
86
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
87
+ ),
88
+ local_value_map,
89
+ )
78
90
  return res
79
91
 
80
92
  return eval_inner
81
93
 
82
94
  def _make_eval_grad_inner(self, world_space: bool):
83
- if not self.gradient_valid():
95
+ if not self.space.gradient_valid():
84
96
  return None
85
97
 
86
98
  @cache.dynamic_func(suffix=self.name)
87
- def eval_grad_inner_ref_space(args: self.ElementEvalArg, s: Sample):
88
- res = utils.generalized_outer(
99
+ def eval_grad_inner(args: self.ElementEvalArg, s: Sample, grad_transform: Any):
100
+ local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
101
+
102
+ res = self.space.space_gradient(
89
103
  self._read_node_value(args, s.element_index, 0),
90
104
  self.space.element_inner_weight_gradient(
91
105
  args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
92
106
  ),
107
+ local_value_map,
108
+ grad_transform,
93
109
  )
110
+
94
111
  node_count = self.space.topology.element_node_count(
95
112
  args.elt_arg, args.eval_arg.topology_arg, s.element_index
96
113
  )
97
114
  for k in range(1, node_count):
98
- res += utils.generalized_outer(
115
+ res += self.space.space_gradient(
99
116
  self._read_node_value(args, s.element_index, k),
100
117
  self.space.element_inner_weight_gradient(
101
118
  args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
102
119
  ),
120
+ local_value_map,
121
+ grad_transform,
103
122
  )
104
123
  return res
105
124
 
125
+ @cache.dynamic_func(suffix=self.name)
126
+ def eval_grad_inner_ref_space(args: self.ElementEvalArg, s: Sample):
127
+ grad_transform = 1.0
128
+ return eval_grad_inner(args, s, grad_transform)
129
+
106
130
  @cache.dynamic_func(suffix=self.name)
107
131
  def eval_grad_inner_world_space(args: self.ElementEvalArg, s: Sample):
108
132
  grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
109
- res = eval_grad_inner_ref_space(args, s)
110
- return res * grad_transform
133
+ return eval_grad_inner(args, s, grad_transform)
111
134
 
112
135
  return eval_grad_inner_world_space if world_space else eval_grad_inner_ref_space
113
136
 
@@ -118,25 +141,28 @@ class NodalFieldBase(DiscreteField):
118
141
  @cache.dynamic_func(suffix=self.name)
119
142
  def eval_div_inner(args: self.ElementEvalArg, s: Sample):
120
143
  grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
144
+ local_value_map = self.space.local_value_map_inner(args.elt_arg, s.element_index, s.element_coords)
121
145
 
122
- res = utils.generalized_inner(
146
+ res = self.space.space_divergence(
123
147
  self._read_node_value(args, s.element_index, 0),
124
148
  self.space.element_inner_weight_gradient(
125
149
  args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
126
- )
127
- * grad_transform,
150
+ ),
151
+ local_value_map,
152
+ grad_transform,
128
153
  )
129
154
 
130
155
  node_count = self.space.topology.element_node_count(
131
156
  args.elt_arg, args.eval_arg.topology_arg, s.element_index
132
157
  )
133
158
  for k in range(1, node_count):
134
- res += utils.generalized_inner(
159
+ res += self.space.space_divergence(
135
160
  self._read_node_value(args, s.element_index, k),
136
161
  self.space.element_inner_weight_gradient(
137
162
  args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
138
- )
139
- * grad_transform,
163
+ ),
164
+ local_value_map,
165
+ grad_transform,
140
166
  )
141
167
  return res
142
168
 
@@ -145,57 +171,71 @@ class NodalFieldBase(DiscreteField):
145
171
  def _make_eval_outer(self):
146
172
  @cache.dynamic_func(suffix=self.name)
147
173
  def eval_outer(args: self.ElementEvalArg, s: Sample):
148
- res = self.space.element_outer_weight(
149
- args.elt_arg,
150
- args.eval_arg.space_arg,
151
- s.element_index,
152
- s.element_coords,
153
- 0,
154
- ) * self._read_node_value(args, s.element_index, 0)
174
+ local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
175
+ res = self.space.space_value(
176
+ self._read_node_value(args, s.element_index, 0),
177
+ self.space.element_outer_weight(
178
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
179
+ ),
180
+ local_value_map,
181
+ )
182
+
155
183
  node_count = self.space.topology.element_node_count(
156
184
  args.elt_arg, args.eval_arg.topology_arg, s.element_index
157
185
  )
186
+
158
187
  for k in range(1, node_count):
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)
188
+ res += self.space.space_value(
189
+ self._read_node_value(args, s.element_index, k),
190
+ self.space.element_outer_weight(
191
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
192
+ ),
193
+ local_value_map,
194
+ )
166
195
  return res
167
196
 
168
197
  return eval_outer
169
198
 
170
199
  def _make_eval_grad_outer(self, world_space: bool):
171
- if not self.gradient_valid():
200
+ if not self.space.gradient_valid():
172
201
  return None
173
202
 
174
203
  @cache.dynamic_func(suffix=self.name)
175
- def eval_grad_outer_ref_space(args: self.ElementEvalArg, s: Sample):
176
- res = utils.generalized_outer(
204
+ def eval_grad_outer(args: self.ElementEvalArg, s: Sample, grad_transform: Any):
205
+ local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
206
+
207
+ res = self.space.space_gradient(
177
208
  self._read_node_value(args, s.element_index, 0),
178
209
  self.space.element_outer_weight_gradient(
179
210
  args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
180
211
  ),
212
+ local_value_map,
213
+ grad_transform,
181
214
  )
215
+
182
216
  node_count = self.space.topology.element_node_count(
183
217
  args.elt_arg, args.eval_arg.topology_arg, s.element_index
184
218
  )
185
219
  for k in range(1, node_count):
186
- res += utils.generalized_outer(
220
+ res += self.space.space_gradient(
187
221
  self._read_node_value(args, s.element_index, k),
188
222
  self.space.element_outer_weight_gradient(
189
223
  args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
190
224
  ),
225
+ local_value_map,
226
+ grad_transform,
191
227
  )
192
228
  return res
193
229
 
230
+ @cache.dynamic_func(suffix=self.name)
231
+ def eval_grad_outer_ref_space(args: self.ElementEvalArg, s: Sample):
232
+ grad_transform = 1.0
233
+ return eval_grad_outer_ref_space(args, s, grad_transform)
234
+
194
235
  @cache.dynamic_func(suffix=self.name)
195
236
  def eval_grad_outer_world_space(args: self.ElementEvalArg, s: Sample):
196
237
  grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
197
- res = eval_grad_outer_ref_space(args, s)
198
- return res * grad_transform
238
+ return eval_grad_outer_ref_space(args, s, grad_transform)
199
239
 
200
240
  return eval_grad_outer_world_space if world_space else eval_grad_outer_ref_space
201
241
 
@@ -206,25 +246,28 @@ class NodalFieldBase(DiscreteField):
206
246
  @cache.dynamic_func(suffix=self.name)
207
247
  def eval_div_outer(args: self.ElementEvalArg, s: Sample):
208
248
  grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
249
+ local_value_map = self.space.local_value_map_outer(args.elt_arg, s.element_index, s.element_coords)
209
250
 
210
- res = utils.generalized_inner(
251
+ res = self.space.space_divergence(
211
252
  self._read_node_value(args, s.element_index, 0),
212
253
  self.space.element_outer_weight_gradient(
213
254
  args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
214
- )
215
- * grad_transform,
255
+ ),
256
+ local_value_map,
257
+ grad_transform,
216
258
  )
217
259
 
218
260
  node_count = self.space.topology.element_node_count(
219
261
  args.elt_arg, args.eval_arg.topology_arg, s.element_index
220
262
  )
221
263
  for k in range(1, node_count):
222
- res += utils.generalized_inner(
264
+ res += self.space.space_divergence(
223
265
  self._read_node_value(args, s.element_index, k),
224
266
  self.space.element_outer_weight_gradient(
225
267
  args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
226
- )
227
- * grad_transform,
268
+ ),
269
+ local_value_map,
270
+ grad_transform,
228
271
  )
229
272
  return res
230
273
 
@@ -232,8 +275,17 @@ class NodalFieldBase(DiscreteField):
232
275
 
233
276
  def _make_set_node_value(self):
234
277
  @cache.dynamic_func(suffix=self.name)
235
- def set_node_value(args: self.EvalArg, partition_node_index: int, value: self.space.dtype):
236
- args.dof_values[partition_node_index] = self.space.dof_mapper.value_to_dof(value)
278
+ def set_node_value(
279
+ elt_arg: self.space.ElementArg,
280
+ eval_arg: self.EvalArg,
281
+ element_index: ElementIndex,
282
+ node_index_in_element: int,
283
+ partition_node_index: int,
284
+ value: self.space.dtype,
285
+ ):
286
+ eval_arg.dof_values[partition_node_index] = self.space.node_dof_value(
287
+ elt_arg, eval_arg.space_arg, element_index, node_index_in_element, value
288
+ )
237
289
 
238
290
  return set_node_value
239
291