warp-lang 1.3.2__py3-none-macosx_10_13_universal2.whl → 1.4.0__py3-none-macosx_10_13_universal2.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 (107) hide show
  1. warp/__init__.py +6 -0
  2. warp/autograd.py +59 -6
  3. warp/bin/libwarp.dylib +0 -0
  4. warp/build_dll.py +8 -10
  5. warp/builtins.py +126 -4
  6. warp/codegen.py +435 -53
  7. warp/config.py +1 -1
  8. warp/context.py +678 -403
  9. warp/dlpack.py +2 -0
  10. warp/examples/benchmarks/benchmark_cloth.py +10 -0
  11. warp/examples/core/example_render_opengl.py +12 -10
  12. warp/examples/fem/example_adaptive_grid.py +251 -0
  13. warp/examples/fem/example_apic_fluid.py +1 -1
  14. warp/examples/fem/example_diffusion_3d.py +2 -2
  15. warp/examples/fem/example_magnetostatics.py +1 -1
  16. warp/examples/fem/example_streamlines.py +1 -0
  17. warp/examples/fem/utils.py +23 -4
  18. warp/examples/sim/example_cloth.py +50 -6
  19. warp/fem/__init__.py +2 -0
  20. warp/fem/adaptivity.py +493 -0
  21. warp/fem/field/field.py +2 -1
  22. warp/fem/field/nodal_field.py +18 -26
  23. warp/fem/field/test.py +4 -4
  24. warp/fem/field/trial.py +4 -4
  25. warp/fem/geometry/__init__.py +1 -0
  26. warp/fem/geometry/adaptive_nanogrid.py +843 -0
  27. warp/fem/geometry/nanogrid.py +55 -28
  28. warp/fem/space/__init__.py +1 -1
  29. warp/fem/space/nanogrid_function_space.py +69 -35
  30. warp/fem/utils.py +113 -107
  31. warp/jax_experimental.py +28 -15
  32. warp/native/array.h +0 -1
  33. warp/native/builtin.h +103 -6
  34. warp/native/bvh.cu +2 -0
  35. warp/native/cuda_util.cpp +14 -0
  36. warp/native/cuda_util.h +2 -0
  37. warp/native/error.cpp +4 -2
  38. warp/native/exports.h +99 -17
  39. warp/native/mat.h +97 -0
  40. warp/native/mesh.cpp +36 -0
  41. warp/native/mesh.cu +51 -0
  42. warp/native/mesh.h +1 -0
  43. warp/native/quat.h +43 -0
  44. warp/native/spatial.h +6 -0
  45. warp/native/vec.h +74 -0
  46. warp/native/warp.cpp +2 -1
  47. warp/native/warp.cu +10 -3
  48. warp/native/warp.h +8 -1
  49. warp/paddle.py +382 -0
  50. warp/sim/__init__.py +1 -0
  51. warp/sim/collide.py +519 -0
  52. warp/sim/integrator_euler.py +18 -5
  53. warp/sim/integrator_featherstone.py +5 -5
  54. warp/sim/integrator_vbd.py +1026 -0
  55. warp/sim/model.py +49 -23
  56. warp/stubs.py +459 -0
  57. warp/tape.py +2 -0
  58. warp/tests/aux_test_dependent.py +1 -0
  59. warp/tests/aux_test_name_clash1.py +32 -0
  60. warp/tests/aux_test_name_clash2.py +32 -0
  61. warp/tests/aux_test_square.py +1 -0
  62. warp/tests/test_array.py +222 -0
  63. warp/tests/test_async.py +3 -3
  64. warp/tests/test_atomic.py +6 -0
  65. warp/tests/test_closest_point_edge_edge.py +93 -1
  66. warp/tests/test_codegen.py +62 -15
  67. warp/tests/test_codegen_instancing.py +1457 -0
  68. warp/tests/test_collision.py +486 -0
  69. warp/tests/test_compile_consts.py +3 -28
  70. warp/tests/test_dlpack.py +170 -0
  71. warp/tests/test_examples.py +22 -8
  72. warp/tests/test_fast_math.py +10 -4
  73. warp/tests/test_fem.py +64 -0
  74. warp/tests/test_func.py +46 -0
  75. warp/tests/test_implicit_init.py +49 -0
  76. warp/tests/test_jax.py +58 -0
  77. warp/tests/test_mat.py +84 -0
  78. warp/tests/test_mesh_query_point.py +188 -0
  79. warp/tests/test_module_hashing.py +40 -0
  80. warp/tests/test_multigpu.py +3 -3
  81. warp/tests/test_overwrite.py +8 -0
  82. warp/tests/test_paddle.py +852 -0
  83. warp/tests/test_print.py +89 -0
  84. warp/tests/test_quat.py +111 -0
  85. warp/tests/test_reload.py +31 -1
  86. warp/tests/test_scalar_ops.py +2 -0
  87. warp/tests/test_static.py +412 -0
  88. warp/tests/test_streams.py +64 -3
  89. warp/tests/test_struct.py +4 -4
  90. warp/tests/test_torch.py +24 -0
  91. warp/tests/test_triangle_closest_point.py +137 -0
  92. warp/tests/test_types.py +1 -1
  93. warp/tests/test_vbd.py +386 -0
  94. warp/tests/test_vec.py +143 -0
  95. warp/tests/test_vec_scalar_ops.py +139 -0
  96. warp/tests/test_volume.py +30 -0
  97. warp/tests/unittest_suites.py +12 -0
  98. warp/tests/unittest_utils.py +9 -5
  99. warp/thirdparty/dlpack.py +3 -1
  100. warp/types.py +157 -34
  101. warp/utils.py +37 -14
  102. {warp_lang-1.3.2.dist-info → warp_lang-1.4.0.dist-info}/METADATA +10 -8
  103. {warp_lang-1.3.2.dist-info → warp_lang-1.4.0.dist-info}/RECORD +106 -94
  104. warp/tests/test_point_triangle_closest_point.py +0 -143
  105. {warp_lang-1.3.2.dist-info → warp_lang-1.4.0.dist-info}/LICENSE.md +0 -0
  106. {warp_lang-1.3.2.dist-info → warp_lang-1.4.0.dist-info}/WHEEL +0 -0
  107. {warp_lang-1.3.2.dist-info → warp_lang-1.4.0.dist-info}/top_level.txt +0 -0
warp/fem/adaptivity.py ADDED
@@ -0,0 +1,493 @@
1
+ from typing import List, Optional, Tuple
2
+
3
+ import numpy as np
4
+
5
+ import warp as wp
6
+ from warp.fem import cache
7
+ from warp.fem.domain import Cells
8
+ from warp.fem.field import GeometryField
9
+ from warp.fem.geometry import AdaptiveNanogrid
10
+ from warp.fem.integrate import interpolate
11
+ from warp.fem.operator import integrand, lookup
12
+ from warp.fem.types import NULL_ELEMENT_INDEX, Domain, Field, Sample
13
+
14
+
15
+ def adaptive_nanogrid_from_hierarchy(
16
+ grids: List[wp.Volume], grading: Optional[str] = None, temporary_store: Optional[cache.TemporaryStore] = None
17
+ ) -> AdaptiveNanogrid:
18
+ """
19
+ Constructs a :class:`warp.fem.AdaptiveNanogrid` from a non-overlapping grid hierarchy.
20
+
21
+ Warning: The result is undefined if there are partial overlaps between levels, that is if a cell at level `l` is only partially covered by cells at levels `l-1` or lower.
22
+
23
+ Args:
24
+ grids: List of sparse Volumes, from finest to coarsest
25
+ grading: Supplementary grading condition, may be ``None``, "face" or "vertex"; see :func:`enforce_nanogrid_grading`
26
+ temporary_store: Storage for temporary allocations
27
+ """
28
+ if not grids:
29
+ raise ValueError("No grids to build from!")
30
+
31
+ level_count = len(grids)
32
+
33
+ device = grids[0].device
34
+
35
+ # Concatenate voxels for each grid
36
+ voxel_counts = [grid.get_voxel_count() for grid in grids]
37
+
38
+ voxel_offsets = np.cumsum(np.array([0] + voxel_counts))
39
+ merged_ijks = cache.borrow_temporary(temporary_store, dtype=wp.vec3i, shape=int(voxel_offsets[-1]), device=device)
40
+ for l in range(level_count):
41
+ voxel_count = voxel_counts[l]
42
+ grid_voxels = cache.borrow_temporary(temporary_store, shape=(voxel_count,), dtype=wp.vec3i, device=device)
43
+ grids[l].get_voxels(out=grid_voxels.array)
44
+
45
+ wp.launch(
46
+ _fill_hierarchy_merged_ijk,
47
+ dim=voxel_count,
48
+ device=device,
49
+ inputs=[l, voxel_offsets[l], grid_voxels.array, merged_ijks.array],
50
+ )
51
+
52
+ # Allocate merged grid
53
+ grid_info = grids[0].get_grid_info()
54
+ cell_grid = wp.Volume.allocate_by_voxels(
55
+ merged_ijks.array,
56
+ transform=grid_info.transform_matrix,
57
+ translation=grid_info.translation,
58
+ device=device,
59
+ )
60
+
61
+ # Get unique voxel and corresponding level
62
+ cell_count = cell_grid.get_voxel_count()
63
+ cell_ijk = cache.borrow_temporary(temporary_store, shape=(cell_count,), dtype=wp.vec3i, device=device)
64
+ cell_level = wp.array(shape=(cell_count,), dtype=wp.uint8, device=device)
65
+
66
+ cell_grid.get_voxels(out=cell_ijk.array)
67
+
68
+ cell_grid_ids = wp.array([grid.id for grid in grids], dtype=wp.uint64, device=device)
69
+ wp.launch(
70
+ _fill_hierarchy_cell_level,
71
+ device=device,
72
+ dim=cell_count,
73
+ inputs=[level_count, cell_grid_ids, cell_ijk.array, cell_level],
74
+ )
75
+
76
+ cell_grid, cell_level = enforce_nanogrid_grading(
77
+ cell_grid, cell_level, level_count=level_count, grading=grading, temporary_store=temporary_store
78
+ )
79
+
80
+ return AdaptiveNanogrid(cell_grid, cell_level=cell_level, level_count=level_count, temporary_store=temporary_store)
81
+
82
+
83
+ def adaptive_nanogrid_from_field(
84
+ coarse_grid: wp.Volume,
85
+ level_count: int,
86
+ refinement_field: GeometryField,
87
+ samples_per_voxel: int = 64,
88
+ grading: Optional[str] = None,
89
+ temporary_store: Optional[cache.TemporaryStore] = None,
90
+ ) -> AdaptiveNanogrid:
91
+ """
92
+ Constructs a :class:`warp.fem.AdaptiveNanogrid` from a coarse grid and a refinement field.
93
+
94
+ Args:
95
+ coarse_grid: Base grid from which to start refining. No voxels will be added outside of the base grid.
96
+ level_count: Maximum number of refinement levels
97
+ refinement_field: Scalar field used as a refinement oracle. If the returned value is negative, the corresponding voxel will be carved out.
98
+ Positive values indicate the desired refinement with 0.0 corresponding to the finest level and 1.0 to the coarsest level.
99
+ samples_per_voxel: How many samples to use for evaluating the refinement field within each voxel
100
+ grading: Supplementary grading condition, may be ``None``, "face" or "vertex"; see :func:`enforce_nanogrid_grading`
101
+ temporary_store: Storage for temporary allocations
102
+ """
103
+
104
+ device = coarse_grid.device
105
+ cell_count = coarse_grid.get_voxel_count()
106
+
107
+ cell_ijk = cache.borrow_temporary(temporary_store, shape=(cell_count,), dtype=wp.vec3i, device=device)
108
+ cell_level = cache.borrow_temporary(temporary_store, shape=(cell_count,), dtype=wp.uint8, device=device)
109
+
110
+ cell_level.array.fill_(level_count - 1)
111
+ coarse_grid.get_voxels(out=cell_ijk.array)
112
+
113
+ domain = Cells(refinement_field.geometry)
114
+
115
+ fine_count = cache.borrow_temporary(temporary_store, dtype=int, shape=1, device=device)
116
+ fine_count.array.zero_()
117
+
118
+ for _ in range(level_count):
119
+ cell_count = cell_ijk.array.shape[0]
120
+ cell_refinement = cache.borrow_temporary(temporary_store, shape=(cell_count,), dtype=wp.int8, device=device)
121
+
122
+ with wp.ScopedDevice(device):
123
+ interpolate(
124
+ _count_refined_voxels,
125
+ domain=domain,
126
+ dim=cell_count,
127
+ fields={"field": refinement_field},
128
+ values={
129
+ "sample_count": samples_per_voxel,
130
+ "level_count": level_count,
131
+ "coarse_grid": coarse_grid.id,
132
+ "coarse_ijk": cell_ijk.array,
133
+ "coarse_level": cell_level.array,
134
+ "coarse_refinement": cell_refinement.array,
135
+ "fine_count": fine_count.array,
136
+ },
137
+ )
138
+
139
+ fine_shape = int(fine_count.array.numpy()[0])
140
+ fine_ijk = cache.borrow_temporary(temporary_store, shape=fine_shape, dtype=wp.vec3i, device=device)
141
+ fine_level = cache.borrow_temporary(temporary_store, shape=fine_shape, dtype=wp.uint8, device=device)
142
+
143
+ wp.launch(
144
+ _fill_refined_voxels,
145
+ dim=cell_count,
146
+ device=device,
147
+ inputs=[
148
+ cell_ijk.array,
149
+ cell_level.array,
150
+ cell_refinement.array,
151
+ fine_count.array,
152
+ fine_ijk.array,
153
+ fine_level.array,
154
+ ],
155
+ )
156
+
157
+ # Fine is the new coarse
158
+ cell_ijk = fine_ijk
159
+ cell_level = fine_level
160
+
161
+ wp.launch(_adjust_refined_ijk, dim=fine_shape, device=device, inputs=[cell_ijk.array, cell_level.array])
162
+
163
+ # We now have our refined voxels, allocate the grid
164
+ coarse_info = coarse_grid.get_grid_info()
165
+ fine_scale = 1.0 / (1 << (level_count - 1))
166
+ fine_transform = coarse_info.transform_matrix * fine_scale
167
+ fine_translation = coarse_info.translation + (fine_scale - 1.0) * 0.5 * wp.vec3(coarse_grid.get_voxel_size())
168
+ fine_grid = wp.Volume.allocate_by_voxels(
169
+ cell_ijk.array, translation=fine_translation, transform=fine_transform, device=device
170
+ )
171
+
172
+ # Reorder cell_levels (voxels will have moved)
173
+ fine_count = fine_grid.get_voxel_count()
174
+ fine_level = wp.array(dtype=wp.uint8, shape=fine_count, device=device)
175
+ wp.launch(
176
+ _fill_refined_level,
177
+ dim=fine_count,
178
+ device=device,
179
+ inputs=[fine_grid.id, cell_ijk.array, cell_level.array, fine_level],
180
+ )
181
+
182
+ fine_grid, fine_level = enforce_nanogrid_grading(
183
+ fine_grid, fine_level, level_count=level_count, grading=grading, temporary_store=temporary_store
184
+ )
185
+
186
+ return AdaptiveNanogrid(fine_grid, cell_level=fine_level, level_count=level_count, temporary_store=temporary_store)
187
+
188
+
189
+ def enforce_nanogrid_grading(
190
+ cell_grid: wp.Volume,
191
+ cell_level: wp.array,
192
+ level_count: int,
193
+ grading: Optional[str] = None,
194
+ temporary_store: Optional[cache.TemporaryStore] = None,
195
+ ) -> Tuple[wp.Volume, wp.array]:
196
+ """
197
+ Refines an adaptive grid such that if satisfies a grading condition.
198
+
199
+ Arguments are similar to the :class:`warp.fem.AdaptiveNanogrid` constructor, with the
200
+ addition of the `grading` condition which can be:
201
+ - "face": two cells sharing a common face must have a level difference of at most 1
202
+ - "vertex": two cells sharing a common vertex must have a level difference of at most 1
203
+ - "none" or ``None``: no grading condition
204
+
205
+ Returns the refined grid and levels
206
+ """
207
+
208
+ if not grading or grading == "none" or level_count <= 2:
209
+ # skip
210
+ return cell_grid, cell_level
211
+
212
+ device = cell_grid.device
213
+
214
+ grid_info = cell_grid.get_grid_info()
215
+
216
+ fine_count = cache.borrow_temporary(temporary_store, shape=(1,), dtype=int, device=device)
217
+ grading_kernel = _count_ungraded_faces if grading == "face" else _count_ungraded_vertices
218
+
219
+ for _ in range(level_count - 2):
220
+ cell_count = cell_grid.get_voxel_count()
221
+ cell_ijk = cache.borrow_temporary(temporary_store, shape=(cell_count,), dtype=wp.vec3i, device=device)
222
+ cell_grid.get_voxels(out=cell_ijk.array)
223
+
224
+ refinement = cache.borrow_temporary(temporary_store, shape=(cell_count,), dtype=int, device=device)
225
+ refinement.array.zero_()
226
+
227
+ wp.launch(
228
+ grading_kernel,
229
+ dim=cell_count,
230
+ device=device,
231
+ inputs=[cell_grid.id, cell_ijk.array, cell_level, level_count, refinement.array],
232
+ )
233
+
234
+ fine_count.array.fill_(cell_count)
235
+ wp.launch(
236
+ _count_graded_cells,
237
+ dim=cell_count,
238
+ device=device,
239
+ inputs=[
240
+ refinement.array,
241
+ fine_count.array,
242
+ ],
243
+ )
244
+
245
+ # Add new coordinates
246
+ fine_shape = int(fine_count.array.numpy()[0])
247
+ if fine_shape == cell_count:
248
+ break
249
+
250
+ fine_ijk = cache.borrow_temporary(temporary_store, shape=fine_shape, dtype=wp.vec3i, device=device)
251
+ fine_level = cache.borrow_temporary(temporary_store, shape=fine_shape, dtype=wp.uint8, device=device)
252
+
253
+ wp.launch(
254
+ _fill_graded_cells,
255
+ dim=cell_count,
256
+ device=device,
257
+ inputs=[
258
+ cell_ijk.array,
259
+ cell_level,
260
+ refinement.array,
261
+ fine_count.array,
262
+ fine_ijk.array,
263
+ fine_level.array,
264
+ ],
265
+ )
266
+
267
+ # Rebuild grid and levels
268
+ cell_grid = wp.Volume.allocate_by_voxels(
269
+ fine_ijk.array, translation=grid_info.translation, transform=grid_info.transform_matrix, device=device
270
+ )
271
+ cell_level = wp.empty(fine_shape, dtype=wp.uint8, device=device)
272
+ wp.launch(
273
+ _fill_refined_level,
274
+ dim=fine_shape,
275
+ device=device,
276
+ inputs=[cell_grid.id, fine_ijk.array, fine_level.array, cell_level],
277
+ )
278
+
279
+ return cell_grid, cell_level
280
+
281
+
282
+ @wp.kernel
283
+ def _count_ungraded_faces(
284
+ cell_grid: wp.uint64,
285
+ cell_ijk: wp.array(dtype=wp.vec3i),
286
+ cell_level: wp.array(dtype=wp.uint8),
287
+ level_count: int,
288
+ refinement: wp.array(dtype=wp.int32),
289
+ ):
290
+ cell = wp.tid()
291
+
292
+ ijk = cell_ijk[cell]
293
+ level = int(cell_level[cell])
294
+
295
+ for axis in range(3):
296
+ for j in range(2):
297
+ nijk = ijk
298
+ nijk[axis] += (2 * j - 1) << level
299
+
300
+ nidx = AdaptiveNanogrid.find_cell(cell_grid, nijk, level_count, cell_level)
301
+ if nidx != -1:
302
+ n_level = cell_level[nidx]
303
+ if n_level > (level + 1):
304
+ wp.atomic_add(refinement, nidx, 1)
305
+
306
+
307
+ @wp.kernel
308
+ def _count_ungraded_vertices(
309
+ cell_grid: wp.uint64,
310
+ cell_ijk: wp.array(dtype=wp.vec3i),
311
+ cell_level: wp.array(dtype=wp.uint8),
312
+ level_count: int,
313
+ refinement: wp.array(dtype=wp.int32),
314
+ ):
315
+ cell = wp.tid()
316
+
317
+ ijk = cell_ijk[cell]
318
+ level = int(cell_level[cell])
319
+
320
+ for i in range(-1, 2):
321
+ for j in range(-1, 2):
322
+ for k in range(-1, 2):
323
+ nijk = ijk + AdaptiveNanogrid.fine_ijk(wp.vec3i(i, j, k), level)
324
+
325
+ nidx = AdaptiveNanogrid.find_cell(cell_grid, nijk, level_count, cell_level)
326
+ if nidx != -1:
327
+ n_level = cell_level[nidx]
328
+ if n_level > (level + 1):
329
+ wp.atomic_add(refinement, nidx, 1)
330
+
331
+
332
+ @wp.kernel
333
+ def _count_graded_cells(
334
+ refinement: wp.array(dtype=wp.int32),
335
+ fine_count: wp.array(dtype=int),
336
+ ):
337
+ cell = wp.tid()
338
+ if refinement[cell] > 0:
339
+ wp.atomic_add(fine_count, 0, 7)
340
+
341
+
342
+ @wp.kernel
343
+ def _fill_graded_cells(
344
+ coarse_ijk: wp.array(dtype=wp.vec3i),
345
+ coarse_level: wp.array(dtype=wp.uint8),
346
+ coarse_refinement: wp.array(dtype=wp.int32),
347
+ fine_count: wp.array(dtype=int),
348
+ fine_ijk: wp.array(dtype=wp.vec3i),
349
+ fine_level: wp.array(dtype=wp.uint8),
350
+ ):
351
+ cell = wp.tid()
352
+ ijk = coarse_ijk[cell]
353
+ level = int(coarse_level[cell])
354
+ refinement = wp.min(1, coarse_refinement[cell])
355
+
356
+ count = wp.select(refinement > 0, 1, 8)
357
+ offset = wp.atomic_sub(fine_count, 0, count) - count
358
+
359
+ f_level = level - refinement
360
+ for k in range(count):
361
+ f_ijk = ijk + AdaptiveNanogrid.fine_ijk(wp.vec3i(k >> 2, (k & 2) >> 1, k & 1), f_level)
362
+ fine_ijk[offset + k] = f_ijk
363
+ fine_level[offset + k] = wp.uint8(f_level)
364
+
365
+
366
+ @integrand
367
+ def _sample_refinement(
368
+ rng: wp.uint32,
369
+ sample_count: int,
370
+ ijk: wp.vec3i,
371
+ domain: Domain,
372
+ field: Field,
373
+ coarse_grid: wp.uint64,
374
+ cur_level: int,
375
+ level_count: int,
376
+ ):
377
+ min_level = level_count
378
+
379
+ scale = 1.0 / float(1 << (level_count - 1 - cur_level))
380
+ uvw = wp.vec3(ijk) * scale + wp.vec3(0.5 * (scale - 1.0))
381
+
382
+ for _ in range(sample_count):
383
+ trial_uvw = uvw + wp.sample_unit_cube(rng) * scale
384
+ pos = wp.volume_index_to_world(coarse_grid, trial_uvw)
385
+ field_s = lookup(domain, pos)
386
+ if field_s.element_index != NULL_ELEMENT_INDEX:
387
+ sampled_level = wp.min(level_count - 1, int(wp.floor(field(field_s) * float(level_count))))
388
+ if sampled_level >= 0:
389
+ min_level = wp.min(sampled_level, min_level)
390
+
391
+ return wp.select(min_level < level_count, -1, cur_level - wp.clamp(min_level, 0, cur_level))
392
+
393
+
394
+ @integrand
395
+ def _count_refined_voxels(
396
+ s: Sample,
397
+ domain: Domain,
398
+ field: Field,
399
+ sample_count: int,
400
+ level_count: int,
401
+ coarse_grid: wp.uint64,
402
+ coarse_ijk: wp.array(dtype=wp.vec3i),
403
+ coarse_level: wp.array(dtype=wp.uint8),
404
+ coarse_refinement: wp.array(dtype=wp.int8),
405
+ fine_count: wp.array(dtype=int),
406
+ ):
407
+ cell = s.qp_index
408
+
409
+ ijk = coarse_ijk[cell]
410
+ cur_level = int(coarse_level[cell])
411
+
412
+ seed = (cur_level << 24) ^ (ijk[0] << 12) ^ (ijk[1] << 6) ^ ijk[2]
413
+ rng = wp.rand_init(seed)
414
+
415
+ refinement = _sample_refinement(rng, sample_count, ijk, domain, field, coarse_grid, cur_level, level_count)
416
+
417
+ coarse_refinement[cell] = wp.int8(refinement)
418
+ if refinement >= 0:
419
+ wp.atomic_add(fine_count, 0, wp.select(refinement > 0, 1, 8))
420
+
421
+
422
+ @wp.kernel
423
+ def _fill_refined_voxels(
424
+ coarse_ijk: wp.array(dtype=wp.vec3i),
425
+ coarse_level: wp.array(dtype=wp.uint8),
426
+ coarse_refinement: wp.array(dtype=wp.int8),
427
+ fine_count: wp.array(dtype=int),
428
+ fine_ijk: wp.array(dtype=wp.vec3i),
429
+ fine_level: wp.array(dtype=wp.uint8),
430
+ ):
431
+ cell = wp.tid()
432
+ ijk = coarse_ijk[cell]
433
+ level = int(coarse_level[cell])
434
+ refinement = wp.min(1, int(coarse_refinement[cell]))
435
+
436
+ if refinement >= 0:
437
+ count = wp.select(refinement > 0, 1, 8)
438
+ offset = wp.atomic_sub(fine_count, 0, count) - count
439
+
440
+ f_level = level - refinement
441
+ for k in range(count):
442
+ f_ijk = AdaptiveNanogrid.fine_ijk(ijk, refinement) + wp.vec3i(k >> 2, (k & 2) >> 1, k & 1)
443
+ fine_ijk[offset + k] = f_ijk
444
+ fine_level[offset + k] = wp.uint8(f_level)
445
+
446
+
447
+ @wp.kernel
448
+ def _adjust_refined_ijk(
449
+ cell_ijk: wp.array(dtype=wp.vec3i),
450
+ cell_level: wp.array(dtype=wp.uint8),
451
+ ):
452
+ cell = wp.tid()
453
+ cell_ijk[cell] = AdaptiveNanogrid.fine_ijk(cell_ijk[cell], int(cell_level[cell]))
454
+
455
+
456
+ @wp.kernel
457
+ def _fill_refined_level(
458
+ fine_grid: wp.uint64,
459
+ cell_ijk: wp.array(dtype=wp.vec3i),
460
+ cell_level: wp.array(dtype=wp.uint8),
461
+ fine_level: wp.array(dtype=wp.uint8),
462
+ ):
463
+ cell = wp.tid()
464
+ ijk = cell_ijk[cell]
465
+ level = cell_level[cell]
466
+
467
+ fine_level[wp.volume_lookup_index(fine_grid, ijk[0], ijk[1], ijk[2])] = level
468
+
469
+
470
+ @wp.kernel
471
+ def _fill_hierarchy_merged_ijk(
472
+ level: int, level_offset: int, level_ijk: wp.array(dtype=wp.vec3i), merged_ijk: wp.array(dtype=wp.vec3i)
473
+ ):
474
+ cell = wp.tid()
475
+ merged_ijk[cell + level_offset] = AdaptiveNanogrid.fine_ijk(level_ijk[cell], level)
476
+
477
+
478
+ @wp.kernel
479
+ def _fill_hierarchy_cell_level(
480
+ level_count: int,
481
+ grids: wp.array(dtype=wp.uint64),
482
+ cells_ijk: wp.array(dtype=wp.vec3i),
483
+ levels: wp.array(dtype=wp.uint8),
484
+ ):
485
+ cell = wp.tid()
486
+ ijk = cells_ijk[cell]
487
+
488
+ for k in range(level_count):
489
+ idx = wp.volume_lookup_index(grids[k], ijk[0], ijk[1], ijk[2])
490
+ if idx != -1:
491
+ levels[cell] = wp.uint8(k)
492
+ break
493
+ ijk = AdaptiveNanogrid.coarse_ijk(ijk, 1)
warp/fem/field/field.py CHANGED
@@ -114,7 +114,8 @@ class GeometryField(FieldLike):
114
114
  """Returns a deformed version of the underlying geometry, with positions displaced according to this field's values.
115
115
 
116
116
  Args:
117
- relative: If ``True``, the field is intepreted as a relative displacement over the original geometry. If ``False``, the field values are intepreted as absolute positions.
117
+ relative: If ``True``, the field is interpreted as a relative displacement over the original geometry.
118
+ If ``False``, the field values are interpreted as absolute positions.
118
119
 
119
120
  """
120
121
  return DeformedGeometry(self, relative=relative)
@@ -107,7 +107,7 @@ class NodalFieldBase(DiscreteField):
107
107
  def eval_grad_inner_world_space(args: self.ElementEvalArg, s: Sample):
108
108
  grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
109
109
  res = eval_grad_inner_ref_space(args, s)
110
- return utils.apply_right(res, grad_transform)
110
+ return res * grad_transform
111
111
 
112
112
  return eval_grad_inner_world_space if world_space else eval_grad_inner_ref_space
113
113
 
@@ -121,12 +121,10 @@ class NodalFieldBase(DiscreteField):
121
121
 
122
122
  res = utils.generalized_inner(
123
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
- ),
124
+ self.space.element_inner_weight_gradient(
125
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
126
+ )
127
+ * grad_transform,
130
128
  )
131
129
 
132
130
  node_count = self.space.topology.element_node_count(
@@ -135,12 +133,10 @@ class NodalFieldBase(DiscreteField):
135
133
  for k in range(1, node_count):
136
134
  res += utils.generalized_inner(
137
135
  self._read_node_value(args, s.element_index, k),
138
- utils.apply_right(
139
- self.space.element_inner_weight_gradient(
140
- args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
141
- ),
142
- grad_transform,
143
- ),
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,
144
140
  )
145
141
  return res
146
142
 
@@ -199,7 +195,7 @@ class NodalFieldBase(DiscreteField):
199
195
  def eval_grad_outer_world_space(args: self.ElementEvalArg, s: Sample):
200
196
  grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
201
197
  res = eval_grad_outer_ref_space(args, s)
202
- return utils.apply_right(res, grad_transform)
198
+ return res * grad_transform
203
199
 
204
200
  return eval_grad_outer_world_space if world_space else eval_grad_outer_ref_space
205
201
 
@@ -213,12 +209,10 @@ class NodalFieldBase(DiscreteField):
213
209
 
214
210
  res = utils.generalized_inner(
215
211
  self._read_node_value(args, s.element_index, 0),
216
- utils.apply_right(
217
- self.space.element_outer_weight_gradient(
218
- args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
219
- ),
220
- grad_transform,
221
- ),
212
+ self.space.element_outer_weight_gradient(
213
+ args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, 0
214
+ )
215
+ * grad_transform,
222
216
  )
223
217
 
224
218
  node_count = self.space.topology.element_node_count(
@@ -227,12 +221,10 @@ class NodalFieldBase(DiscreteField):
227
221
  for k in range(1, node_count):
228
222
  res += utils.generalized_inner(
229
223
  self._read_node_value(args, s.element_index, k),
230
- utils.apply_right(
231
- self.space.element_outer_weight_gradient(
232
- args.elt_arg, args.eval_arg.space_arg, s.element_index, s.element_coords, k
233
- ),
234
- grad_transform,
235
- ),
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,
236
228
  )
237
229
  return res
238
230
 
warp/fem/field/test.py CHANGED
@@ -87,7 +87,7 @@ class TestField(SpaceField):
87
87
  grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
88
88
  return utils.generalized_outer(
89
89
  self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
90
- utils.apply_right(nabla_weight, grad_transform),
90
+ nabla_weight * grad_transform,
91
91
  )
92
92
 
93
93
  return eval_nabla_test_inner
@@ -108,7 +108,7 @@ class TestField(SpaceField):
108
108
  grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
109
109
  return utils.generalized_inner(
110
110
  self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
111
- utils.apply_right(nabla_weight, grad_transform),
111
+ nabla_weight * grad_transform,
112
112
  )
113
113
 
114
114
  return eval_div_test_inner
@@ -143,7 +143,7 @@ class TestField(SpaceField):
143
143
  grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
144
144
  return utils.generalized_outer(
145
145
  self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
146
- utils.apply_right(nabla_weight, grad_transform),
146
+ nabla_weight * grad_transform,
147
147
  )
148
148
 
149
149
  return eval_nabla_test_outer
@@ -164,7 +164,7 @@ class TestField(SpaceField):
164
164
  grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
165
165
  return utils.generalized_inner(
166
166
  self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
167
- utils.apply_right(nabla_weight, grad_transform),
167
+ nabla_weight * grad_transform,
168
168
  )
169
169
 
170
170
  return eval_div_test_outer
warp/fem/field/trial.py CHANGED
@@ -90,7 +90,7 @@ class TrialField(SpaceField):
90
90
  grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
91
91
  return utils.generalized_outer(
92
92
  self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
93
- utils.apply_right(nabla_weight, grad_transform),
93
+ nabla_weight * grad_transform,
94
94
  )
95
95
 
96
96
  return eval_nabla_trial_inner
@@ -111,7 +111,7 @@ class TrialField(SpaceField):
111
111
  grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
112
112
  return utils.generalized_inner(
113
113
  self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
114
- utils.apply_right(nabla_weight, grad_transform),
114
+ nabla_weight * grad_transform,
115
115
  )
116
116
 
117
117
  return eval_div_trial_inner
@@ -146,7 +146,7 @@ class TrialField(SpaceField):
146
146
  grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
147
147
  return utils.generalized_outer(
148
148
  self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
149
- utils.apply_right(nabla_weight, grad_transform),
149
+ nabla_weight * grad_transform,
150
150
  )
151
151
 
152
152
  return eval_nabla_trial_outer
@@ -167,7 +167,7 @@ class TrialField(SpaceField):
167
167
  grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
168
168
  return utils.generalized_inner(
169
169
  self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
170
- utils.apply_right(nabla_weight, grad_transform),
170
+ nabla_weight * grad_transform,
171
171
  )
172
172
 
173
173
  return eval_div_trial_outer
@@ -1,3 +1,4 @@
1
+ from .adaptive_nanogrid import AdaptiveNanogrid
1
2
  from .deformed_geometry import DeformedGeometry
2
3
  from .element import Element
3
4
  from .geometry import Geometry