warp-lang 1.4.1__py3-none-macosx_10_13_universal2.whl → 1.5.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 (164) hide show
  1. warp/__init__.py +4 -0
  2. warp/autograd.py +43 -8
  3. warp/bin/libwarp-clang.dylib +0 -0
  4. warp/bin/libwarp.dylib +0 -0
  5. warp/build.py +21 -2
  6. warp/build_dll.py +23 -6
  7. warp/builtins.py +1920 -111
  8. warp/codegen.py +186 -62
  9. warp/config.py +2 -2
  10. warp/context.py +322 -73
  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/core/example_dem.py +2 -1
  17. warp/examples/core/example_mesh_intersect.py +3 -3
  18. warp/examples/fem/example_adaptive_grid.py +37 -10
  19. warp/examples/fem/example_apic_fluid.py +3 -2
  20. warp/examples/fem/example_convection_diffusion_dg.py +4 -5
  21. warp/examples/fem/example_deformed_geometry.py +1 -1
  22. warp/examples/fem/example_diffusion_3d.py +47 -4
  23. warp/examples/fem/example_distortion_energy.py +220 -0
  24. warp/examples/fem/example_magnetostatics.py +127 -85
  25. warp/examples/fem/example_nonconforming_contact.py +5 -5
  26. warp/examples/fem/example_stokes.py +3 -1
  27. warp/examples/fem/example_streamlines.py +12 -19
  28. warp/examples/fem/utils.py +38 -15
  29. warp/examples/optim/example_walker.py +2 -2
  30. warp/examples/sim/example_cloth.py +2 -25
  31. warp/examples/sim/example_jacobian_ik.py +6 -2
  32. warp/examples/sim/example_quadruped.py +2 -1
  33. warp/examples/tile/example_tile_convolution.py +58 -0
  34. warp/examples/tile/example_tile_fft.py +47 -0
  35. warp/examples/tile/example_tile_filtering.py +105 -0
  36. warp/examples/tile/example_tile_matmul.py +79 -0
  37. warp/examples/tile/example_tile_mlp.py +375 -0
  38. warp/fem/__init__.py +8 -0
  39. warp/fem/cache.py +16 -12
  40. warp/fem/dirichlet.py +1 -1
  41. warp/fem/domain.py +44 -1
  42. warp/fem/field/__init__.py +1 -2
  43. warp/fem/field/field.py +31 -19
  44. warp/fem/field/nodal_field.py +101 -49
  45. warp/fem/field/virtual.py +794 -0
  46. warp/fem/geometry/__init__.py +2 -2
  47. warp/fem/geometry/deformed_geometry.py +3 -105
  48. warp/fem/geometry/element.py +13 -0
  49. warp/fem/geometry/geometry.py +165 -5
  50. warp/fem/geometry/grid_2d.py +3 -6
  51. warp/fem/geometry/grid_3d.py +31 -28
  52. warp/fem/geometry/hexmesh.py +3 -46
  53. warp/fem/geometry/nanogrid.py +3 -2
  54. warp/fem/geometry/{quadmesh_2d.py → quadmesh.py} +280 -159
  55. warp/fem/geometry/tetmesh.py +2 -43
  56. warp/fem/geometry/{trimesh_2d.py → trimesh.py} +354 -186
  57. warp/fem/integrate.py +683 -261
  58. warp/fem/linalg.py +404 -0
  59. warp/fem/operator.py +101 -18
  60. warp/fem/polynomial.py +5 -5
  61. warp/fem/quadrature/quadrature.py +45 -21
  62. warp/fem/space/__init__.py +45 -11
  63. warp/fem/space/basis_function_space.py +451 -0
  64. warp/fem/space/basis_space.py +58 -11
  65. warp/fem/space/function_space.py +146 -5
  66. warp/fem/space/grid_2d_function_space.py +80 -66
  67. warp/fem/space/grid_3d_function_space.py +113 -68
  68. warp/fem/space/hexmesh_function_space.py +96 -108
  69. warp/fem/space/nanogrid_function_space.py +62 -110
  70. warp/fem/space/quadmesh_function_space.py +208 -0
  71. warp/fem/space/shape/__init__.py +45 -7
  72. warp/fem/space/shape/cube_shape_function.py +328 -54
  73. warp/fem/space/shape/shape_function.py +10 -1
  74. warp/fem/space/shape/square_shape_function.py +328 -60
  75. warp/fem/space/shape/tet_shape_function.py +269 -19
  76. warp/fem/space/shape/triangle_shape_function.py +238 -19
  77. warp/fem/space/tetmesh_function_space.py +69 -37
  78. warp/fem/space/topology.py +38 -0
  79. warp/fem/space/trimesh_function_space.py +179 -0
  80. warp/fem/utils.py +6 -331
  81. warp/jax_experimental.py +3 -1
  82. warp/native/array.h +55 -40
  83. warp/native/builtin.h +124 -43
  84. warp/native/bvh.h +4 -0
  85. warp/native/coloring.cpp +600 -0
  86. warp/native/cuda_util.cpp +14 -0
  87. warp/native/cuda_util.h +2 -1
  88. warp/native/fabric.h +8 -0
  89. warp/native/hashgrid.h +4 -0
  90. warp/native/marching.cu +8 -0
  91. warp/native/mat.h +14 -3
  92. warp/native/mathdx.cpp +59 -0
  93. warp/native/mesh.h +4 -0
  94. warp/native/range.h +13 -1
  95. warp/native/reduce.cpp +9 -1
  96. warp/native/reduce.cu +7 -0
  97. warp/native/runlength_encode.cpp +9 -1
  98. warp/native/runlength_encode.cu +7 -1
  99. warp/native/scan.cpp +8 -0
  100. warp/native/scan.cu +8 -0
  101. warp/native/scan.h +8 -1
  102. warp/native/sparse.cpp +8 -0
  103. warp/native/sparse.cu +8 -0
  104. warp/native/temp_buffer.h +7 -0
  105. warp/native/tile.h +1857 -0
  106. warp/native/tile_gemm.h +341 -0
  107. warp/native/tile_reduce.h +210 -0
  108. warp/native/volume_builder.cu +8 -0
  109. warp/native/volume_builder.h +8 -0
  110. warp/native/warp.cpp +10 -2
  111. warp/native/warp.cu +369 -15
  112. warp/native/warp.h +12 -2
  113. warp/optim/adam.py +39 -4
  114. warp/paddle.py +29 -12
  115. warp/render/render_opengl.py +137 -65
  116. warp/sim/graph_coloring.py +292 -0
  117. warp/sim/integrator_euler.py +4 -2
  118. warp/sim/integrator_featherstone.py +115 -44
  119. warp/sim/integrator_vbd.py +6 -0
  120. warp/sim/model.py +90 -17
  121. warp/stubs.py +651 -85
  122. warp/tape.py +12 -7
  123. warp/tests/assets/pixel.npy +0 -0
  124. warp/tests/aux_test_instancing_gc.py +18 -0
  125. warp/tests/test_array.py +207 -48
  126. warp/tests/test_closest_point_edge_edge.py +8 -8
  127. warp/tests/test_codegen.py +120 -1
  128. warp/tests/test_codegen_instancing.py +30 -0
  129. warp/tests/test_collision.py +110 -0
  130. warp/tests/test_coloring.py +241 -0
  131. warp/tests/test_context.py +34 -0
  132. warp/tests/test_examples.py +18 -4
  133. warp/tests/test_fabricarray.py +33 -0
  134. warp/tests/test_fem.py +453 -113
  135. warp/tests/test_func.py +48 -1
  136. warp/tests/test_generics.py +52 -0
  137. warp/tests/test_iter.py +68 -0
  138. warp/tests/test_mat_scalar_ops.py +1 -1
  139. warp/tests/test_mesh_query_point.py +5 -4
  140. warp/tests/test_module_hashing.py +23 -0
  141. warp/tests/test_paddle.py +27 -87
  142. warp/tests/test_print.py +191 -1
  143. warp/tests/test_spatial.py +1 -1
  144. warp/tests/test_tile.py +700 -0
  145. warp/tests/test_tile_mathdx.py +144 -0
  146. warp/tests/test_tile_mlp.py +383 -0
  147. warp/tests/test_tile_reduce.py +374 -0
  148. warp/tests/test_tile_shared_memory.py +190 -0
  149. warp/tests/test_vbd.py +12 -20
  150. warp/tests/test_volume.py +43 -0
  151. warp/tests/unittest_suites.py +23 -2
  152. warp/tests/unittest_utils.py +4 -0
  153. warp/types.py +339 -73
  154. warp/utils.py +22 -1
  155. {warp_lang-1.4.1.dist-info → warp_lang-1.5.0.dist-info}/METADATA +33 -7
  156. {warp_lang-1.4.1.dist-info → warp_lang-1.5.0.dist-info}/RECORD +159 -132
  157. {warp_lang-1.4.1.dist-info → warp_lang-1.5.0.dist-info}/WHEEL +1 -1
  158. warp/fem/field/test.py +0 -180
  159. warp/fem/field/trial.py +0 -183
  160. warp/fem/space/collocated_function_space.py +0 -102
  161. warp/fem/space/quadmesh_2d_function_space.py +0 -261
  162. warp/fem/space/trimesh_2d_function_space.py +0 -153
  163. {warp_lang-1.4.1.dist-info → warp_lang-1.5.0.dist-info}/LICENSE.md +0 -0
  164. {warp_lang-1.4.1.dist-info → warp_lang-1.5.0.dist-info}/top_level.txt +0 -0
@@ -8,15 +8,41 @@ from warp.fem.geometry import Grid3D
8
8
  from warp.fem.polynomial import Polynomial, is_closed, lagrange_scales, quadrature_1d
9
9
  from warp.fem.types import Coords
10
10
 
11
+ from .shape_function import ShapeFunction
11
12
  from .tet_shape_function import TetrahedronPolynomialShapeFunctions
12
13
 
13
14
 
14
- class CubeTripolynomialShapeFunctions:
15
+ class CubeShapeFunction(ShapeFunction):
15
16
  VERTEX = 0
16
17
  EDGE = 1
17
18
  FACE = 2
18
19
  INTERIOR = 3
19
20
 
21
+ @wp.func
22
+ def _vertex_coords(vidx_in_cell: int):
23
+ x = vidx_in_cell // 4
24
+ y = (vidx_in_cell - 4 * x) // 2
25
+ z = vidx_in_cell - 4 * x - 2 * y
26
+ return wp.vec3i(x, y, z)
27
+
28
+ @wp.func
29
+ def _edge_coords(type_instance: int, index_in_side: int):
30
+ return wp.vec3i(index_in_side + 1, (type_instance & 2) >> 1, type_instance & 1)
31
+
32
+ @wp.func
33
+ def _edge_axis(type_instance: int):
34
+ return type_instance >> 2
35
+
36
+ @wp.func
37
+ def _face_offset(type_instance: int):
38
+ return type_instance & 1
39
+
40
+ @wp.func
41
+ def _face_axis(type_instance: int):
42
+ return type_instance >> 1
43
+
44
+
45
+ class CubeTripolynomialShapeFunctions(CubeShapeFunction):
20
46
  def __init__(self, degree: int, family: Polynomial):
21
47
  self.family = family
22
48
 
@@ -24,6 +50,17 @@ class CubeTripolynomialShapeFunctions:
24
50
  self.NODES_PER_ELEMENT = wp.constant((degree + 1) ** 3)
25
51
  self.NODES_PER_EDGE = wp.constant(degree + 1)
26
52
 
53
+ if is_closed(self.family):
54
+ self.VERTEX_NODE_COUNT = wp.constant(1)
55
+ self.EDGE_NODE_COUNT = wp.constant(max(0, degree - 1))
56
+ self.FACE_NODE_COUNT = wp.constant(max(0, degree - 1) ** 2)
57
+ self.INTERIOR_NODE_COUNT = wp.constant(max(0, degree - 1) ** 3)
58
+ else:
59
+ self.VERTEX_NODE_COUNT = wp.constant(0)
60
+ self.EDGE_NODE_COUNT = wp.constant(0)
61
+ self.FACE_NODE_COUNT = wp.constant(0)
62
+ self.INTERIOR_NODE_COUNT = self.NODES_PER_ELEMENT
63
+
27
64
  lobatto_coords, lobatto_weight = quadrature_1d(point_count=degree + 1, family=family)
28
65
  lagrange_scale = lagrange_scales(lobatto_coords)
29
66
 
@@ -36,6 +73,10 @@ class CubeTripolynomialShapeFunctions:
36
73
  self._node_ijk = self._make_node_ijk()
37
74
  self.node_type_and_type_index = self._make_node_type_and_type_index()
38
75
 
76
+ if degree > 2:
77
+ self._face_node_ij = self._make_face_node_ij()
78
+ self._linear_face_node_index = self._make_linear_face_node_index()
79
+
39
80
  @property
40
81
  def name(self) -> str:
41
82
  return f"Cube_Q{self.ORDER}_{self.family}"
@@ -61,9 +102,35 @@ class CubeTripolynomialShapeFunctions:
61
102
 
62
103
  return cache.get_func(node_ijk, self.name)
63
104
 
105
+ def _make_face_node_ij(self):
106
+ ORDER_MINUS_ONE = wp.constant(self.ORDER - 1)
107
+
108
+ def face_node_ij(
109
+ face_node_index: int,
110
+ ):
111
+ node_i = face_node_index // ORDER_MINUS_ONE
112
+ node_j = face_node_index - ORDER_MINUS_ONE * node_i
113
+ return wp.vec2i(node_i, node_j)
114
+
115
+ return cache.get_func(face_node_ij, self.name)
116
+
117
+ def _make_linear_face_node_index(self):
118
+ ORDER_MINUS_ONE = wp.constant(self.ORDER - 1)
119
+
120
+ def linear_face_node_index(face_node_index: int, face_node_ij: wp.vec2i):
121
+ return face_node_ij[0] * ORDER_MINUS_ONE + face_node_ij[1]
122
+
123
+ return cache.get_func(linear_face_node_index, self.name)
124
+
64
125
  def _make_node_type_and_type_index(self):
65
126
  ORDER = self.ORDER
66
127
 
128
+ @cache.dynamic_func(suffix=self.name)
129
+ def node_type_and_type_index_open(
130
+ node_index_in_elt: int,
131
+ ):
132
+ return CubeShapeFunction.INTERIOR, 0, node_index_in_elt
133
+
67
134
  @cache.dynamic_func(suffix=self.name)
68
135
  def node_type_and_type_index(
69
136
  node_index_in_elt: int,
@@ -122,7 +189,7 @@ class CubeTripolynomialShapeFunctions:
122
189
  type_index = ((i - 1) * (ORDER - 1) + (j - 1)) * (ORDER - 1) + k - 1
123
190
  return CubeTripolynomialShapeFunctions.INTERIOR, 0, type_index
124
191
 
125
- return node_type_and_type_index
192
+ return node_type_and_type_index if is_closed(self.family) else node_type_and_type_index_open
126
193
 
127
194
  def make_node_coords_in_element(self):
128
195
  LOBATTO_COORDS = self.LOBATTO_COORDS
@@ -404,18 +471,13 @@ class CubeTripolynomialShapeFunctions:
404
471
  return cell_indices[np.newaxis, :], np.array([cell_type], dtype=np.int8)
405
472
 
406
473
 
407
- class CubeSerendipityShapeFunctions:
474
+ class CubeSerendipityShapeFunctions(CubeShapeFunction):
408
475
  """
409
476
  Serendipity element ~ tensor product space without interior nodes
410
477
  Edge shape functions are usual Lagrange shape functions times a bilinear function in the normal directions
411
478
  Corner shape functions are trilinear shape functions times a function of (x^{d-1} + y^{d-1})
412
479
  """
413
480
 
414
- # Node categories
415
- VERTEX = wp.constant(0)
416
- EDGE_X = wp.constant(1)
417
- EDGE_Y = wp.constant(2)
418
-
419
481
  def __init__(self, degree: int, family: Polynomial):
420
482
  if not is_closed(family):
421
483
  raise ValueError("A closed polynomial family is required to define serendipity elements")
@@ -429,6 +491,11 @@ class CubeSerendipityShapeFunctions:
429
491
  self.NODES_PER_ELEMENT = wp.constant(8 + 12 * (degree - 1))
430
492
  self.NODES_PER_EDGE = wp.constant(degree + 1)
431
493
 
494
+ self.VERTEX_NODE_COUNT = wp.constant(1)
495
+ self.EDGE_NODE_COUNT = wp.constant(degree - 1)
496
+ self.FACE_NODE_COUNT = wp.constant(0)
497
+ self.INTERIOR_NODE_COUNT = wp.constant(0)
498
+
432
499
  lobatto_coords, lobatto_weight = quadrature_1d(point_count=degree + 1, family=family)
433
500
  lagrange_scale = lagrange_scales(lobatto_coords)
434
501
 
@@ -451,48 +518,35 @@ class CubeSerendipityShapeFunctions:
451
518
  node_index_in_elt: int,
452
519
  ):
453
520
  if node_index_in_elt < 8:
454
- return CubeSerendipityShapeFunctions.VERTEX, node_index_in_elt
521
+ return CubeShapeFunction.VERTEX, node_index_in_elt, 0
455
522
 
456
- type_index = (node_index_in_elt - 8) // 3
457
- side = node_index_in_elt - 8 - 3 * type_index
458
- return CubeSerendipityShapeFunctions.EDGE_X + side, type_index
523
+ edge_index = (node_index_in_elt - 8) // 3
524
+ edge_axis = node_index_in_elt - 8 - 3 * edge_index
459
525
 
460
- return node_type_and_index
526
+ index_in_edge = edge_index // 4
527
+ edge_offset = edge_index - 4 * index_in_edge
461
528
 
462
- @wp.func
463
- def _vertex_coords(vidx_in_cell: int):
464
- x = vidx_in_cell // 4
465
- y = (vidx_in_cell - 4 * x) // 2
466
- z = vidx_in_cell - 4 * x - 2 * y
467
- return wp.vec3i(x, y, z)
529
+ return CubeShapeFunction.EDGE, 4 * edge_axis + edge_offset, index_in_edge
468
530
 
469
- @wp.func
470
- def _edge_coords(type_index: int):
471
- index_in_side = type_index // 4
472
- side_offset = type_index - 4 * index_in_side
473
- return wp.vec3i(index_in_side + 1, side_offset // 2, side_offset & 1)
474
-
475
- @wp.func
476
- def _edge_axis(node_type: int):
477
- return node_type - CubeSerendipityShapeFunctions.EDGE_X
531
+ return node_type_and_index
478
532
 
479
- @wp.func
480
- def _cube_edge_index(node_type: int, type_index: int):
481
- index_in_side = type_index // 4
482
- side_offset = type_index - 4 * index_in_side
533
+ # @wp.func
534
+ # def _cube_edge_index(node_type: int, type_index: int):
535
+ # index_in_side = type_index // 4
536
+ # side_offset = type_index - 4 * index_in_side
483
537
 
484
- return 4 * (node_type - CubeSerendipityShapeFunctions.EDGE_X) + side_offset, index_in_side
538
+ # return 4 * (node_type - CubeSerendipityShapeFunctions.EDGE_X) + side_offset, index_in_side
485
539
 
486
540
  def _get_node_lobatto_indices(self):
487
541
  ORDER = self.ORDER
488
542
 
489
543
  @cache.dynamic_func(suffix=self.name)
490
- def node_lobatto_indices(node_type: int, type_index: int):
544
+ def node_lobatto_indices(node_type: int, type_instance: int, type_index: int):
491
545
  if node_type == CubeSerendipityShapeFunctions.VERTEX:
492
- return CubeSerendipityShapeFunctions._vertex_coords(type_index) * ORDER
546
+ return CubeSerendipityShapeFunctions._vertex_coords(type_instance) * ORDER
493
547
 
494
- axis = CubeSerendipityShapeFunctions._edge_axis(node_type)
495
- local_coords = CubeSerendipityShapeFunctions._edge_coords(type_index)
548
+ axis = CubeSerendipityShapeFunctions._edge_axis(type_instance)
549
+ local_coords = CubeSerendipityShapeFunctions._edge_coords(type_instance, type_index)
496
550
 
497
551
  local_indices = wp.vec3i(local_coords[0], local_coords[1] * ORDER, local_coords[2] * ORDER)
498
552
 
@@ -507,8 +561,8 @@ class CubeSerendipityShapeFunctions:
507
561
  def node_coords_in_element(
508
562
  node_index_in_elt: int,
509
563
  ):
510
- node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
511
- node_coords = self._node_lobatto_indices(node_type, type_index)
564
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
565
+ node_coords = self._node_lobatto_indices(node_type, type_instance, type_index)
512
566
  return Coords(
513
567
  LOBATTO_COORDS[node_coords[0]], LOBATTO_COORDS[node_coords[1]], LOBATTO_COORDS[node_coords[2]]
514
568
  )
@@ -522,7 +576,7 @@ class CubeSerendipityShapeFunctions:
522
576
  def node_quadrature_weight(
523
577
  node_index_in_elt: int,
524
578
  ):
525
- node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
579
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
526
580
  if node_type == CubeSerendipityShapeFunctions.VERTEX:
527
581
  return 1.0 / float(8 * ORDER * ORDER * ORDER)
528
582
 
@@ -537,7 +591,7 @@ class CubeSerendipityShapeFunctions:
537
591
  def trace_node_quadrature_weight(
538
592
  node_index_in_elt: int,
539
593
  ):
540
- node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
594
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
541
595
  if node_type == CubeSerendipityShapeFunctions.VERTEX:
542
596
  return 0.25 / float(ORDER * ORDER)
543
597
 
@@ -560,10 +614,10 @@ class CubeSerendipityShapeFunctions:
560
614
  coords: Coords,
561
615
  node_index_in_elt: int,
562
616
  ):
563
- node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
617
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
564
618
 
565
619
  if node_type == CubeSerendipityShapeFunctions.VERTEX:
566
- node_ijk = CubeSerendipityShapeFunctions._vertex_coords(type_index)
620
+ node_ijk = CubeSerendipityShapeFunctions._vertex_coords(type_instance)
567
621
 
568
622
  cx = wp.select(node_ijk[0] == 0, coords[0], 1.0 - coords[0])
569
623
  cy = wp.select(node_ijk[1] == 0, coords[1], 1.0 - coords[1])
@@ -571,10 +625,10 @@ class CubeSerendipityShapeFunctions:
571
625
 
572
626
  w = cx * cy * cz
573
627
 
574
- if ORDER == 2:
628
+ if wp.static(ORDER == 2):
575
629
  w *= cx + cy + cz - 3.0 + LOBATTO_COORDS[1]
576
630
  return w * LAGRANGE_SCALE[0]
577
- if ORDER == 3:
631
+ if wp.static(ORDER == 3):
578
632
  w *= (
579
633
  (cx - 0.5) * (cx - 0.5)
580
634
  + (cy - 0.5) * (cy - 0.5)
@@ -583,9 +637,9 @@ class CubeSerendipityShapeFunctions:
583
637
  )
584
638
  return w * DEGREE_3_SPHERE_SCALE
585
639
 
586
- axis = CubeSerendipityShapeFunctions._edge_axis(node_type)
640
+ axis = CubeSerendipityShapeFunctions._edge_axis(type_instance)
587
641
 
588
- node_all = CubeSerendipityShapeFunctions._edge_coords(type_index)
642
+ node_all = CubeSerendipityShapeFunctions._edge_coords(type_instance, type_index)
589
643
 
590
644
  local_coords = Grid3D._world_to_local(axis, coords)
591
645
 
@@ -616,10 +670,10 @@ class CubeSerendipityShapeFunctions:
616
670
  coords: Coords,
617
671
  node_index_in_elt: int,
618
672
  ):
619
- node_type, type_index = self.node_type_and_type_index(node_index_in_elt)
673
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
620
674
 
621
675
  if node_type == CubeSerendipityShapeFunctions.VERTEX:
622
- node_ijk = CubeSerendipityShapeFunctions._vertex_coords(type_index)
676
+ node_ijk = CubeSerendipityShapeFunctions._vertex_coords(type_instance)
623
677
 
624
678
  cx = wp.select(node_ijk[0] == 0, coords[0], 1.0 - coords[0])
625
679
  cy = wp.select(node_ijk[1] == 0, coords[1], 1.0 - coords[1])
@@ -629,7 +683,7 @@ class CubeSerendipityShapeFunctions:
629
683
  gy = wp.select(node_ijk[1] == 0, 1.0, -1.0)
630
684
  gz = wp.select(node_ijk[2] == 0, 1.0, -1.0)
631
685
 
632
- if ORDER == 2:
686
+ if wp.static(ORDER == 2):
633
687
  w = cx + cy + cz - 3.0 + LOBATTO_COORDS[1]
634
688
  grad_x = cy * cz * gx * (w + cx)
635
689
  grad_y = cz * cx * gy * (w + cy)
@@ -637,7 +691,7 @@ class CubeSerendipityShapeFunctions:
637
691
 
638
692
  return wp.vec3(grad_x, grad_y, grad_z) * LAGRANGE_SCALE[0]
639
693
 
640
- if ORDER == 3:
694
+ if wp.static(ORDER == 3):
641
695
  w = (
642
696
  (cx - 0.5) * (cx - 0.5)
643
697
  + (cy - 0.5) * (cy - 0.5)
@@ -654,8 +708,8 @@ class CubeSerendipityShapeFunctions:
654
708
 
655
709
  return wp.vec3(grad_x, grad_y, grad_z) * DEGREE_3_SPHERE_SCALE
656
710
 
657
- axis = CubeSerendipityShapeFunctions._edge_axis(node_type)
658
- node_all = CubeSerendipityShapeFunctions._edge_coords(type_index)
711
+ axis = CubeSerendipityShapeFunctions._edge_axis(type_instance)
712
+ node_all = CubeSerendipityShapeFunctions._edge_coords(type_instance, type_index)
659
713
 
660
714
  local_coords = Grid3D._world_to_local(axis, coords)
661
715
 
@@ -720,7 +774,7 @@ class CubeSerendipityShapeFunctions:
720
774
  return tets, np.full(tets.shape[0], cell_type, dtype=np.int8)
721
775
 
722
776
 
723
- class CubeNonConformingPolynomialShapeFunctions:
777
+ class CubeNonConformingPolynomialShapeFunctions(ShapeFunction):
724
778
  # embeds the largest regular tet centered at (0.5, 0.5, 0.5) into the reference cube
725
779
 
726
780
  _tet_height = 2.0 / 3.0
@@ -814,3 +868,223 @@ class CubeNonConformingPolynomialShapeFunctions:
814
868
  return wp.transpose(CUBE_TO_TET) * grad
815
869
 
816
870
  return element_inner_weight_gradient
871
+
872
+
873
+ class CubeNedelecFirstKindShapeFunctions(CubeShapeFunction):
874
+ value = ShapeFunction.Value.CovariantVector
875
+
876
+ def __init__(self, degree: int):
877
+ if degree != 1:
878
+ raise NotImplementedError("Only linear Nédélec implemented right now")
879
+
880
+ self.ORDER = wp.constant(degree)
881
+ self.NODES_PER_ELEMENT = wp.constant(12)
882
+ self.NODES_PER_SIDE = wp.constant(4)
883
+
884
+ self.VERTEX_NODE_COUNT = wp.constant(0)
885
+ self.EDGE_NODE_COUNT = wp.constant(1)
886
+ self.FACE_NODE_COUNT = wp.constant(0)
887
+ self.INTERIOR_NODE_COUNT = wp.constant(0)
888
+
889
+ self.node_type_and_type_index = self._get_node_type_and_type_index()
890
+
891
+ @property
892
+ def name(self) -> str:
893
+ return f"CubeN1_{self.ORDER}"
894
+
895
+ def _get_node_type_and_type_index(self):
896
+ @cache.dynamic_func(suffix=self.name)
897
+ def node_type_and_index(
898
+ node_index_in_elt: int,
899
+ ):
900
+ return CubeShapeFunction.EDGE, node_index_in_elt, 0
901
+
902
+ return node_type_and_index
903
+
904
+ def make_node_coords_in_element(self):
905
+ @cache.dynamic_func(suffix=self.name)
906
+ def node_coords_in_element(
907
+ node_index_in_elt: int,
908
+ ):
909
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
910
+ axis = CubeShapeFunction._edge_axis(type_instance)
911
+ local_indices = CubeShapeFunction._edge_coords(type_instance, type_index)
912
+
913
+ local_coords = wp.vec3f(0.5, float(local_indices[1]), float(local_indices[2]))
914
+ return Grid3D._local_to_world(axis, local_coords)
915
+
916
+ return node_coords_in_element
917
+
918
+ def make_node_quadrature_weight(self):
919
+ NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
920
+
921
+ @cache.dynamic_func(suffix=self.name)
922
+ def node_quadrature_weight(node_index_in_element: int):
923
+ return 1.0 / float(NODES_PER_ELEMENT)
924
+
925
+ return node_quadrature_weight
926
+
927
+ def make_trace_node_quadrature_weight(self):
928
+ NODES_PER_SIDE = self.NODES_PER_SIDE
929
+
930
+ @cache.dynamic_func(suffix=self.name)
931
+ def trace_node_quadrature_weight(node_index_in_element: int):
932
+ return 1.0 / float(NODES_PER_SIDE)
933
+
934
+ return trace_node_quadrature_weight
935
+
936
+ def make_element_inner_weight(self):
937
+ @cache.dynamic_func(suffix=self.name)
938
+ def element_inner_weight(
939
+ coords: Coords,
940
+ node_index_in_elt: int,
941
+ ):
942
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
943
+
944
+ axis = CubeShapeFunction._edge_axis(type_instance)
945
+
946
+ local_coords = Grid3D._world_to_local(axis, coords)
947
+ edge_coords = CubeShapeFunction._edge_coords(type_instance, type_index)
948
+
949
+ a1 = float(2 * edge_coords[1] - 1)
950
+ a2 = float(2 * edge_coords[2] - 1)
951
+ b1 = float(1 - edge_coords[1])
952
+ b2 = float(1 - edge_coords[2])
953
+
954
+ local_w = wp.vec3((b1 + a1 * local_coords[1]) * (b2 + a2 * local_coords[2]), 0.0, 0.0)
955
+ return Grid3D._local_to_world(axis, local_w)
956
+
957
+ return element_inner_weight
958
+
959
+ def make_element_inner_weight_gradient(self):
960
+ @cache.dynamic_func(suffix=self.name)
961
+ def element_inner_weight_gradient(
962
+ coords: Coords,
963
+ node_index_in_elt: int,
964
+ ):
965
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
966
+
967
+ axis = CubeShapeFunction._edge_axis(type_instance)
968
+
969
+ local_coords = Grid3D._world_to_local(axis, coords)
970
+ edge_coords = CubeShapeFunction._edge_coords(type_instance, type_index)
971
+
972
+ a1 = float(2 * edge_coords[1] - 1)
973
+ a2 = float(2 * edge_coords[2] - 1)
974
+ b1 = float(1 - edge_coords[1])
975
+ b2 = float(1 - edge_coords[2])
976
+
977
+ local_gw = Grid3D._local_to_world(
978
+ axis, wp.vec3(0.0, a1 * (b2 + a2 * local_coords[2]), (b1 + a1 * local_coords[1]) * a2)
979
+ )
980
+
981
+ grad = wp.mat33(0.0)
982
+ grad[axis] = local_gw
983
+ return grad
984
+
985
+ return element_inner_weight_gradient
986
+
987
+
988
+ class CubeRaviartThomasShapeFunctions(CubeShapeFunction):
989
+ value = ShapeFunction.Value.ContravariantVector
990
+
991
+ def __init__(self, degree: int):
992
+ if degree != 1:
993
+ raise NotImplementedError("Only linear Raviart Thomas implemented right now")
994
+
995
+ self.ORDER = wp.constant(degree)
996
+ self.NODES_PER_ELEMENT = wp.constant(6)
997
+ self.NODES_PER_SIDE = wp.constant(1)
998
+
999
+ self.VERTEX_NODE_COUNT = wp.constant(0)
1000
+ self.EDGE_NODE_COUNT = wp.constant(0)
1001
+ self.FACE_NODE_COUNT = wp.constant(1)
1002
+ self.INTERIOR_NODE_COUNT = wp.constant(0)
1003
+
1004
+ self.node_type_and_type_index = self._get_node_type_and_type_index()
1005
+
1006
+ @property
1007
+ def name(self) -> str:
1008
+ return f"CubeRT_{self.ORDER}"
1009
+
1010
+ def _get_node_type_and_type_index(self):
1011
+ @cache.dynamic_func(suffix=self.name)
1012
+ def node_type_and_index(
1013
+ node_index_in_elt: int,
1014
+ ):
1015
+ return CubeShapeFunction.FACE, node_index_in_elt, 0
1016
+
1017
+ return node_type_and_index
1018
+
1019
+ def make_node_coords_in_element(self):
1020
+ @cache.dynamic_func(suffix=self.name)
1021
+ def node_coords_in_element(
1022
+ node_index_in_elt: int,
1023
+ ):
1024
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
1025
+ axis = CubeShapeFunction._face_axis(type_instance)
1026
+ offset = CubeShapeFunction._face_offset(type_instance)
1027
+
1028
+ coords = Coords(0.5)
1029
+ coords[axis] = float(offset)
1030
+ return coords
1031
+
1032
+ return node_coords_in_element
1033
+
1034
+ def make_node_quadrature_weight(self):
1035
+ NODES_PER_ELEMENT = self.NODES_PER_ELEMENT
1036
+
1037
+ @cache.dynamic_func(suffix=self.name)
1038
+ def node_quadrature_weight(node_index_in_element: int):
1039
+ return 1.0 / float(NODES_PER_ELEMENT)
1040
+
1041
+ return node_quadrature_weight
1042
+
1043
+ def make_trace_node_quadrature_weight(self):
1044
+ NODES_PER_SIDE = self.NODES_PER_SIDE
1045
+
1046
+ @cache.dynamic_func(suffix=self.name)
1047
+ def trace_node_quadrature_weight(node_index_in_element: int):
1048
+ return 1.0 / float(NODES_PER_SIDE)
1049
+
1050
+ return trace_node_quadrature_weight
1051
+
1052
+ def make_element_inner_weight(self):
1053
+ @cache.dynamic_func(suffix=self.name)
1054
+ def element_inner_weight(
1055
+ coords: Coords,
1056
+ node_index_in_elt: int,
1057
+ ):
1058
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
1059
+
1060
+ axis = CubeShapeFunction._face_axis(type_instance)
1061
+ offset = CubeShapeFunction._face_offset(type_instance)
1062
+
1063
+ a = float(2 * offset - 1)
1064
+ b = float(1 - offset)
1065
+
1066
+ w = wp.vec3(0.0)
1067
+ w[axis] = b + a * coords[axis]
1068
+
1069
+ return w
1070
+
1071
+ return element_inner_weight
1072
+
1073
+ def make_element_inner_weight_gradient(self):
1074
+ @cache.dynamic_func(suffix=self.name)
1075
+ def element_inner_weight_gradient(
1076
+ coords: Coords,
1077
+ node_index_in_elt: int,
1078
+ ):
1079
+ node_type, type_instance, type_index = self.node_type_and_type_index(node_index_in_elt)
1080
+
1081
+ axis = CubeShapeFunction._face_axis(type_instance)
1082
+ offset = CubeShapeFunction._face_offset(type_instance)
1083
+
1084
+ a = float(2 * offset - 1)
1085
+ grad = wp.mat33(0.0)
1086
+ grad[axis, axis] = a
1087
+
1088
+ return grad
1089
+
1090
+ return element_inner_weight_gradient
@@ -1,3 +1,5 @@
1
+ from enum import Enum
2
+
1
3
  import numpy as np
2
4
 
3
5
  import warp as wp
@@ -15,6 +17,13 @@ class ShapeFunction:
15
17
  NODES_PER_ELEMENT: int
16
18
  """Number of shape function nodes"""
17
19
 
20
+ class Value(Enum):
21
+ Scalar = 0
22
+ CovariantVector = 1
23
+ ContravariantVector = 2
24
+
25
+ value: Value = Value.Scalar
26
+
18
27
  @property
19
28
  def name(self) -> str:
20
29
  """Unique name encoding all parameters defining the shape function"""
@@ -41,7 +50,7 @@ class ShapeFunction:
41
50
  raise NotImplementedError()
42
51
 
43
52
 
44
- class ConstantShapeFunction:
53
+ class ConstantShapeFunction(ShapeFunction):
45
54
  """Shape function that is constant over the element"""
46
55
 
47
56
  def __init__(self, element: Element, space_dimension: int):