warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

warp/fem/types.py

@@ -1,13 +1,9 @@
 import warp as wp
 
-
-vec1i = wp.types.vector(length=1, dtype=wp.int32)
-vec2i = wp.types.vector(length=2, dtype=wp.int32)
-vec3i = wp.types.vector(length=3, dtype=wp.int32)
-vec4i = wp.types.vector(length=4, dtype=wp.int32)
-vec8i = wp.types.vector(length=8, dtype=wp.int32)
-
-vec6 = wp.types.vector(length=6, dtype=wp.float32)
+# kept to avoid breaking existing example code, no longer used internally
+vec2i = wp.vec2i
+vec3i = wp.vec3i
+vec4i = wp.vec4i
 
 Coords = wp.vec3
 OUTSIDE = wp.constant(-1.0e8)
@@ -20,7 +16,7 @@ NULL_ELEMENT_INDEX = wp.constant(-1)
 NULL_QP_INDEX = wp.constant(-1)
 NULL_NODE_INDEX = wp.constant(-1)
 
-DofIndex = vec2i
+DofIndex = wp.vec2i
 """Opaque descriptor for indexing degrees of freedom within elements"""
 NULL_DOF_INDEX = wp.constant(DofIndex(-1, -1))
 
@@ -59,12 +55,18 @@ class Sample:
     """For bilinear form assembly, index of the trial degree-of-freedom currently being considered"""
 
 
+@wp.func
+def make_free_sample(element_index: ElementIndex, element_coords: Coords):
+    """Returns a :class:`Sample` that is not associated to any quadrature point or dof"""
+    return Sample(element_index, element_coords, NULL_QP_INDEX, 0.0, NULL_DOF_INDEX, NULL_DOF_INDEX)
+
+
 class Field:
     """
     Tag for field-like integrand arguments
     """
 
-    call_operator: "wp.fem.Operator" = None  # Set in operator.py
+    call_operator: "warp.fem.operator.Operator" = None  # Set in operator.py
 
 
 class Domain:
@@ -72,4 +74,4 @@ class Domain:
     Tag for domain-like integrand arguments
     """
 
-    call_operator: "wp.fem.Operator" = None  # Set in operator.py
+    call_operator: "warp.fem.operator.Operator" = None  # Set in operator.py

warp/fem/utils.py

@@ -1,83 +1,123 @@
 from typing import Any, Tuple
 
-import warp as wp
-from warp.utils import radix_sort_pairs, runlength_encode, array_scan
+import numpy as np
 
-from .types import vec6
+import warp as wp
+from warp.fem.cache import (
+    Temporary,
+    TemporaryStore,
+    borrow_temporary,
+    borrow_temporary_like,
+)
+from warp.utils import array_scan, radix_sort_pairs, runlength_encode
 
 
 @wp.func
-def generalized_outer(x: Any, y: wp.vec2):
+def generalized_outer(x: Any, y: Any):
+    """Generalized outer product allowing for the first argument to be a scalar"""
     return wp.outer(x, y)
 
 
 @wp.func
-def generalized_outer(x: Any, y: wp.vec3):
-    return wp.outer(x, y)
+def generalized_outer(x: wp.float32, y: wp.vec2):
+    return x * y
 
 
 @wp.func
-def generalized_outer(x: Any, y: wp.float32):
+def generalized_outer(x: wp.float32, y: wp.vec3):
     return x * y
 
 
 @wp.func
-def unit_element(template_type: wp.float32, coord: int):
-    return 1.0
+def generalized_inner(x: Any, y: Any):
+    """Generalized inner product allowing for the first argument to be a tensor"""
+    return wp.dot(x, y)
 
 
 @wp.func
-def unit_element(template_type: wp.vec2, coord: int):
-    t = wp.vec2(0.0)
-    t[coord] = 1.0
-    return t
+def generalized_inner(x: wp.mat22, y: wp.vec2):
+    return x[0] * y[0] + x[1] * y[1]
 
 
 @wp.func
-def unit_element(template_type: wp.vec3, coord: int):
-    t = wp.vec3(0.0)
-    t[coord] = 1.0
-    return t
+def generalized_inner(x: wp.mat33, y: wp.vec3):
+    return x[0] * y[0] + x[1] * y[1] + x[2] * y[2]
+
+
+@wp.func
+def apply_right(x: Any, y: Any):
+    """Performs x y multiplication with y a square matrix and x either a row-vector or a matrix.
+    Will be removed once native @ operator is implemented.
+    """
+    return x * y
 
 
 @wp.func
-def unit_element(template_type: vec6, coord: int):
-    t = vec6(0.0)
+def apply_right(x: wp.vec2, y: wp.mat22):
+    return x[0] * y[0] + x[1] * y[1]
+
+
+@wp.func
+def apply_right(x: wp.vec3, y: wp.mat33):
+    return x[0] * y[0] + x[1] * y[1] + x[2] * y[2]
+
+
+@wp.func
+def unit_element(template_type: Any, coord: int):
+    """Returns a instance of `template_type` with a single coordinate set to 1 in the canonical basis"""
+
+    t = type(template_type)(0.0)
     t[coord] = 1.0
     return t
 
 
+@wp.func
+def unit_element(template_type: wp.float32, coord: int):
+    return 1.0
+
+
 @wp.func
 def unit_element(template_type: wp.mat22, coord: int):
     t = wp.mat22(0.0)
-    t[coord // 2, coord % 2] = 1.0
+    row = coord // 2
+    col = coord - 2 * row
+    t[row, col] = 1.0
     return t
 
 
 @wp.func
 def unit_element(template_type: wp.mat33, coord: int):
-    t = wp.mat22(0.0)
-    t[coord // 3, coord % 3] = 1.0
+    t = wp.mat33(0.0)
+    row = coord // 3
+    col = coord - 3 * row
+    t[row, col] = 1.0
     return t
 
 
 @wp.func
-def symmetric_part(x: wp.mat22):
-    off_diag = 0.5 * (x[0, 1] + x[1, 0])
-    return wp.mat22(x[0, 0], off_diag, off_diag, x[1, 1])
+def symmetric_part(x: Any):
+    """Symmetric part of a square tensor"""
+    return 0.5 * (x + wp.transpose(x))
 
 
 @wp.func
-def symmetric_part(x: wp.mat33):
-    d = 0.5 * (x[1, 2] + x[2, 1])
-    e = 0.5 * (x[2, 0] + x[0, 2])
-    f = 0.5 * (x[0, 1] + x[1, 0])
-    return wp.mat33(x[0, 0], f, e, f, x[1, 1], d, e, d, x[2, 2])
+def skew_part(x: wp.mat22):
+    """Skew part of a 2x2 tensor as corresponding rotation angle"""
+    return 0.5 * (x[1, 0] - x[0, 1])
+
+
+@wp.func
+def skew_part(x: wp.mat33):
+    """Skew part of a 3x3 tensor as the corresponding rotation vector"""
+    a = 0.5 * (x[2, 1] - x[1, 2])
+    b = 0.5 * (x[0, 2] - x[2, 0])
+    c = 0.5 * (x[1, 0] - x[0, 1])
+    return wp.vec3(a, b, c)
 
 
 def compress_node_indices(
-    node_count: int, node_indices: wp.array(dtype=int)
-) -> Tuple[wp.array, wp.array, int, wp.array]:
+    node_count: int, node_indices: wp.array(dtype=int), temporary_store: TemporaryStore = None
+) -> Tuple[Temporary, Temporary, int, Temporary]:
     """
     Compress an unsorted list of node indices into:
      - a node_offsets array, giving for each node the start offset of corresponding indices in sorted_array_indices
@@ -87,8 +127,14 @@ def compress_node_indices(
     """
 
     index_count = node_indices.size
-    sorted_node_indices = wp.empty(2 * index_count, dtype=int, device=node_indices.device)
-    sorted_array_indices = wp.empty_like(sorted_node_indices)
+
+    sorted_node_indices_temp = borrow_temporary(
+        temporary_store, shape=2 * index_count, dtype=int, device=node_indices.device
+    )
+    sorted_array_indices_temp = borrow_temporary_like(sorted_node_indices_temp, temporary_store)
+
+    sorted_node_indices = sorted_node_indices_temp.array
+    sorted_array_indices = sorted_array_indices_temp.array
 
     wp.copy(dest=sorted_node_indices, src=node_indices, count=index_count)
 
@@ -104,14 +150,44 @@ def compress_node_indices(
     radix_sort_pairs(sorted_node_indices, sorted_array_indices, count=index_count)
 
     # Build prefix sum of number of elements per node
-    unique_node_indices = wp.empty(n=index_count, dtype=int, device=node_indices.device)
-    node_element_counts = wp.empty(n=index_count, dtype=int, device=node_indices.device)
-    unique_node_count = runlength_encode(
-        sorted_node_indices, unique_node_indices, node_element_counts, value_count=index_count
+    unique_node_indices_temp = borrow_temporary(
+        temporary_store, shape=index_count, dtype=int, device=node_indices.device
+    )
+    node_element_counts_temp = borrow_temporary(
+        temporary_store, shape=index_count, dtype=int, device=node_indices.device
+    )
+
+    unique_node_indices = unique_node_indices_temp.array
+    node_element_counts = node_element_counts_temp.array
+
+    unique_node_count_dev = borrow_temporary(temporary_store, shape=(1,), dtype=int, device=sorted_node_indices.device)
+    runlength_encode(
+        sorted_node_indices,
+        unique_node_indices,
+        node_element_counts,
+        value_count=index_count,
+        run_count=unique_node_count_dev.array,
     )
 
+    # Transfer unique node count to host
+    if node_indices.device.is_cuda:
+        unique_node_count_host = borrow_temporary(temporary_store, shape=(1,), dtype=int, pinned=True, device="cpu")
+        wp.copy(src=unique_node_count_dev.array, dest=unique_node_count_host.array, count=1)
+        wp.synchronize_stream(wp.get_stream(node_indices.device))
+        unique_node_count_dev.release()
+        unique_node_count = int(unique_node_count_host.array.numpy()[0])
+        unique_node_count_host.release()
+    else:
+        unique_node_count = int(unique_node_count_dev.array.numpy()[0])
+        unique_node_count_dev.release()
+
     # Scatter seen run counts to global array of element count per node
-    node_offsets = wp.zeros(shape=(node_count + 1), device=node_element_counts.device, dtype=int)
+    node_offsets_temp = borrow_temporary(
+        temporary_store, shape=(node_count + 1), device=node_element_counts.device, dtype=int
+    )
+    node_offsets = node_offsets_temp.array
+
+    node_offsets.zero_()
     wp.launch(
         kernel=_scatter_node_counts,
         dim=unique_node_count,
@@ -122,51 +198,47 @@ def compress_node_indices(
     # Prefix sum of number of elements per node
     array_scan(node_offsets, node_offsets, inclusive=True)
 
-    return node_offsets, sorted_array_indices, unique_node_count, unique_node_indices
-
-
-_pinned_temp_count_buffer = {}
-
+    sorted_node_indices_temp.release()
+    node_element_counts_temp.release()
 
-def _get_pinned_temp_count_buffer(device):
-    device = str(device)
-    if device not in _pinned_temp_count_buffer:
-        _pinned_temp_count_buffer[device] = wp.empty(shape=(1,), dtype=int, pinned=True, device="cpu")
+    return node_offsets_temp, sorted_array_indices_temp, unique_node_count, unique_node_indices_temp
 
-    return _pinned_temp_count_buffer[device]
 
-
-def masked_indices(mask: wp.array(dtype=int), missing_index=-1) -> Tuple[wp.array, wp.array]:
+def masked_indices(
+    mask: wp.array, missing_index=-1, temporary_store: TemporaryStore = None
+) -> Tuple[Temporary, Temporary]:
     """
     From an array of boolean masks (must be either 0 or 1), returns:
       - The list of indices for which the mask is 1
       - A map associating to each element of the input mask array its local index if non-zero, or missing_index if zero.
     """
 
-    offsets = wp.empty_like(mask)
+    offsets_temp = borrow_temporary_like(mask, temporary_store)
+    offsets = offsets_temp.array
 
     wp.utils.array_scan(mask, offsets, inclusive=True)
 
     # Get back total counts on host
     if offsets.device.is_cuda:
-        masked_count = _get_pinned_temp_count_buffer(offsets.device)
-        wp.copy(dest=masked_count, src=offsets, src_offset=offsets.shape[0] - 1, count=1)
-        wp.synchronize_stream(wp.get_stream())
-        masked_count = int(masked_count.numpy()[0])
+        masked_count_temp = borrow_temporary(temporary_store, shape=1, dtype=int, pinned=True, device="cpu")
+        wp.copy(dest=masked_count_temp.array, src=offsets, src_offset=offsets.shape[0] - 1, count=1)
+        wp.synchronize_stream(wp.get_stream(offsets.device))
+        masked_count = int(masked_count_temp.array.numpy()[0])
+        masked_count_temp.release()
     else:
         masked_count = int(offsets.numpy()[-1])
 
     # Convert counts to indices
-    indices = wp.empty(n=masked_count, device=mask.device, dtype=int)
+    indices_temp = borrow_temporary(temporary_store, shape=masked_count, device=mask.device, dtype=int)
 
     wp.launch(
         kernel=_masked_indices_kernel,
         dim=offsets.shape,
-        inputs=[missing_index, mask, offsets, indices, offsets],
+        inputs=[missing_index, mask, offsets, indices_temp.array, offsets],
         device=mask.device,
     )
 
-    return indices, offsets
+    return indices_temp, offsets_temp
 
 
 def array_axpy(x: wp.array, y: wp.array, alpha: float = 1.0, beta: float = 1.0):
@@ -177,7 +249,7 @@ def array_axpy(x: wp.array, y: wp.array, alpha: float = 1.0, beta: float = 1.0):
     alpha = dtype(alpha)
     beta = dtype(beta)
 
-    if x.dtype != y.dtype or x.shape != y.shape or x.device != y.device:
+    if not wp.types.types_equal(x.dtype, y.dtype) or x.shape != y.shape or x.device != y.device:
         raise ValueError("x and y arrays must have same dat atype, shape and device")
 
     wp.launch(kernel=_array_axpy_kernel, dim=x.shape, device=x.device, inputs=[x, y, alpha, beta])
@@ -218,3 +290,206 @@ def _masked_indices_kernel(
 def _array_axpy_kernel(x: wp.array(dtype=Any), y: wp.array(dtype=Any), alpha: Any, beta: Any):
     i = wp.tid()
     y[i] = beta * y[i] + alpha * x[i]
+
+
+def grid_to_tris(Nx: int, Ny: int):
+    """Constructs a triangular mesh topology by dividing each cell of a dense 2D grid into two triangles.
+
+    The resulting triangles will be oriented counter-clockwise assuming that `y` is the fastest moving index direction
+
+    Args:
+        Nx: Resolution of the grid along `x` dimension
+        Ny: Resolution of the grid along `y` dimension
+
+    Returns:
+        Array of shape (2 * Nx * Ny, 3) containing vertex indices for each triangle
+    """
+
+    cx, cy = np.meshgrid(np.arange(Nx, dtype=int), np.arange(Ny, dtype=int), indexing="ij")
+
+    vidx = np.transpose(
+        np.array(
+            [
+                (Ny + 1) * cx + cy,
+                (Ny + 1) * (cx + 1) + cy,
+                (Ny + 1) * (cx + 1) + (cy + 1),
+                (Ny + 1) * cx + cy,
+                (Ny + 1) * (cx + 1) + (cy + 1),
+                (Ny + 1) * (cx) + (cy + 1),
+            ]
+        )
+    ).reshape((-1, 3))
+
+    return vidx
+
+
+def grid_to_tets(Nx: int, Ny: int, Nz: int):
+    """Constructs a tetrahedral mesh topology by diving each cell of a dense 3D grid into five tetrahedrons
+
+    The resulting tets have positive volume assuming that `z` is the fastest moving index direction
+
+    Args:
+        Nx: Resolution of the grid along `x` dimension
+        Ny: Resolution of the grid along `y` dimension
+        Nz: Resolution of the grid along `z` dimension
+
+    Returns:
+        Array of shape (5 * Nx * Ny * Nz, 4) containing vertex indices for each tet
+    """
+
+    # Global node indices for each cell
+    cx, cy, cz = np.meshgrid(
+        np.arange(Nx, dtype=int), np.arange(Ny, dtype=int), np.arange(Nz, dtype=int), indexing="ij"
+    )
+
+    grid_vidx = np.array(
+        [
+            (Ny + 1) * (Nz + 1) * cx + (Nz + 1) * cy + cz,
+            (Ny + 1) * (Nz + 1) * cx + (Nz + 1) * cy + cz + 1,
+            (Ny + 1) * (Nz + 1) * cx + (Nz + 1) * (cy + 1) + cz,
+            (Ny + 1) * (Nz + 1) * cx + (Nz + 1) * (cy + 1) + cz + 1,
+            (Ny + 1) * (Nz + 1) * (cx + 1) + (Nz + 1) * cy + cz,
+            (Ny + 1) * (Nz + 1) * (cx + 1) + (Nz + 1) * cy + cz + 1,
+            (Ny + 1) * (Nz + 1) * (cx + 1) + (Nz + 1) * (cy + 1) + cz,
+            (Ny + 1) * (Nz + 1) * (cx + 1) + (Nz + 1) * (cy + 1) + cz + 1,
+        ]
+    )
+
+    # decompose grid cells into 5 tets
+    tet_vidx = np.array(
+        [
+            [0, 1, 2, 4],
+            [3, 2, 1, 7],
+            [5, 1, 7, 4],
+            [6, 7, 4, 2],
+            [4, 1, 2, 7],
+        ]
+    )
+
+    # Convert to 3d index coordinates
+    vidx_coords = np.array(
+        [
+            [0, 0, 0],
+            [0, 0, 1],
+            [0, 1, 0],
+            [0, 1, 1],
+            [1, 0, 0],
+            [1, 0, 1],
+            [1, 1, 0],
+            [1, 1, 1],
+        ]
+    )
+    tet_coords = vidx_coords[tet_vidx]
+
+    # Symmetry bits for each cell
+    ox, oy, oz = np.meshgrid(
+        np.arange(Nx, dtype=int) % 2, np.arange(Ny, dtype=int) % 2, np.arange(Nz, dtype=int) % 2, indexing="ij"
+    )
+    tet_coords = np.broadcast_to(tet_coords, shape=(*ox.shape, *tet_coords.shape))
+
+    # Flip coordinates according to symmetry
+    ox_bk = np.broadcast_to(ox.reshape(*ox.shape, 1, 1), tet_coords.shape[:-1])
+    oy_bk = np.broadcast_to(oy.reshape(*oy.shape, 1, 1), tet_coords.shape[:-1])
+    oz_bk = np.broadcast_to(oz.reshape(*oz.shape, 1, 1), tet_coords.shape[:-1])
+
+    tet_coords_x = tet_coords[..., 0] ^ ox_bk
+    tet_coords_y = tet_coords[..., 1] ^ oy_bk
+    tet_coords_z = tet_coords[..., 2] ^ oz_bk
+
+    # Back to local vertex indices
+    corner_indices = 4 * tet_coords_x + 2 * tet_coords_y + tet_coords_z
+
+    # Now go from cell-local to global node indices
+    # There must be a nicer way than this, but for small grids this works
+
+    corner_indices = corner_indices.reshape(-1, 4)
+
+    grid_vidx = grid_vidx.reshape((8, -1, 1))
+    grid_vidx = np.broadcast_to(grid_vidx, shape=(8, grid_vidx.shape[1], 5))
+    grid_vidx = grid_vidx.reshape((8, -1))
+
+    node_indices = np.arange(corner_indices.shape[0])
+    tet_grid_vidx = np.transpose(
+        [
+            grid_vidx[corner_indices[:, 0], node_indices],
+            grid_vidx[corner_indices[:, 1], node_indices],
+            grid_vidx[corner_indices[:, 2], node_indices],
+            grid_vidx[corner_indices[:, 3], node_indices],
+        ]
+    )
+
+    return tet_grid_vidx
+
+
+def grid_to_quads(Nx: int, Ny: int):
+    """Constructs a quadrilateral mesh topology from a dense 2D grid
+
+    The resulting quads will be indexed counter-clockwise
+
+    Args:
+        Nx: Resolution of the grid along `x` dimension
+        Ny: Resolution of the grid along `y` dimension
+
+    Returns:
+        Array of shape (Nx * Ny, 4) containing vertex indices for each quadrilateral
+    """
+
+    quad_vtx = np.array(
+        [
+            [0, 0],
+            [1, 0],
+            [1, 1],
+            [0, 1],
+        ]
+    ).T
+
+    quads = np.stack(np.meshgrid(np.arange(0, Nx), np.arange(0, Ny), indexing="ij"))
+
+    quads_vtx_shape = (*quads.shape, quad_vtx.shape[1])
+    quads_vtx = np.broadcast_to(quads.reshape(*quads.shape, 1), quads_vtx_shape) + np.broadcast_to(
+        quad_vtx.reshape(2, 1, 1, quad_vtx.shape[1]), quads_vtx_shape
+    )
+
+    quad_vtx_indices = quads_vtx[0] * (Ny + 1) + quads_vtx[1]
+
+    return quad_vtx_indices.reshape(-1, 4)
+
+
+def grid_to_hexes(Nx: int, Ny: int, Nz: int):
+    """Constructs a hexahedral mesh topology from a dense 3D grid
+
+    The resulting hexes will be indexed following usual convention assuming that `z` is the fastest moving index direction
+    (counter-clockwise bottom vertices, then counter-clockwise top vertices)
+
+    Args:
+        Nx: Resolution of the grid along `x` dimension
+        Ny: Resolution of the grid along `y` dimension
+        Nz: Resolution of the grid along `z` dimension
+
+    Returns:
+        Array of shape (Nx * Ny * Nz, 8) containing vertex indices for each hexaedron
+    """
+
+    hex_vtx = np.array(
+        [
+            [0, 0, 0],
+            [1, 0, 0],
+            [1, 1, 0],
+            [0, 1, 0],
+            [0, 0, 1],
+            [1, 0, 1],
+            [1, 1, 1],
+            [0, 1, 1],
+        ]
+    ).T
+
+    hexes = np.stack(np.meshgrid(np.arange(0, Nx), np.arange(0, Ny), np.arange(0, Nz), indexing="ij"))
+
+    hexes_vtx_shape = (*hexes.shape, hex_vtx.shape[1])
+    hexes_vtx = np.broadcast_to(hexes.reshape(*hexes.shape, 1), hexes_vtx_shape) + np.broadcast_to(
+        hex_vtx.reshape(3, 1, 1, 1, hex_vtx.shape[1]), hexes_vtx_shape
+    )
+
+    hexes_vtx_indices = hexes_vtx[0] * (Nz + 1) * (Ny + 1) + hexes_vtx[1] * (Nz + 1) + hexes_vtx[2]
+
+    return hexes_vtx_indices.reshape(-1, 8)