warp-lang 1.2.1__py3-none-win_amd64.whl → 1.3.0__py3-none-win_amd64.whl

warp/fabric.py

@@ -105,11 +105,14 @@ class fabricarray(noncontiguous_array_base[T]):
     # (initialized when needed)
     _vars = None
 
+    def __new__(cls, *args, **kwargs):
+        instance = super(fabricarray, cls).__new__(cls)
+        instance.deleter = None
+        return instance
+
     def __init__(self, data=None, attrib=None, dtype=Any, ndim=None):
         super().__init__(ARRAY_TYPE_FABRIC)
 
-        self.deleter = None
-
         if data is not None:
             from .context import runtime
 

warp/fem/__init__.py

@@ -1,7 +1,17 @@
 from .cache import TemporaryStore, borrow_temporary, borrow_temporary_like, set_default_temporary_store
 from .dirichlet import normalize_dirichlet_projector, project_linear_system
-from .domain import BoundarySides, Cells, FrontierSides, GeometryDomain, Sides
-from .field import DiscreteField, FieldLike, make_restriction, make_test, make_trial
+from .domain import BoundarySides, Cells, FrontierSides, GeometryDomain, Sides, Subdomain
+from .field import (
+    DiscreteField,
+    FieldLike,
+    ImplicitField,
+    NonconformingField,
+    UniformField,
+    make_discrete_field,
+    make_restriction,
+    make_test,
+    make_trial,
+)
 from .geometry import (
     ExplicitGeometryPartition,
     Geometry,
@@ -58,4 +68,4 @@ from .space import (
     make_space_partition,
     make_space_restriction,
 )
-from .types import Coords, Domain, ElementIndex, Field, Sample
+from .types import NULL_ELEMENT_INDEX, Coords, Domain, ElementIndex, Field, Sample, make_free_sample

warp/fem/cache.py

@@ -1,5 +1,7 @@
+import ast
 import bisect
 import re
+import weakref
 from copy import copy
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 
@@ -17,7 +19,7 @@ def _make_key(obj, suffix: str, use_qualified_name):
     return _key_re.sub("", f"{base_name}_{suffix}")
 
 
-def get_func(func, suffix: str, use_qualified_name: bool = False):
+def get_func(func, suffix: str, use_qualified_name: bool = False, code_transformers=None):
     key = _make_key(func, suffix, use_qualified_name)
 
     if key not in _func_cache:
@@ -28,14 +30,15 @@ def get_func(func, suffix: str, use_qualified_name: bool = False):
             module=wp.get_module(
                 func.__module__,
             ),
+            code_transformers=code_transformers,
         )
 
     return _func_cache[key]
 
 
-def dynamic_func(suffix: str, use_qualified_name=False):
+def dynamic_func(suffix: str, use_qualified_name=False, code_transformers=None):
     def wrap_func(func: Callable):
-        return get_func(func, suffix=suffix, use_qualified_name=use_qualified_name)
+        return get_func(func, suffix=suffix, use_qualified_name=use_qualified_name, code_transformers=code_transformers)
 
     return wrap_func
 
@@ -96,6 +99,92 @@ def dynamic_struct(suffix: str, use_qualified_name=False):
     return wrap_struct
 
 
+def get_argument_struct(arg_types: Dict[str, type]):
+    class Args:
+        pass
+
+    annotations = wp.codegen.get_annotations(Args)
+
+    for name, arg_type in arg_types.items():
+        setattr(Args, name, None)
+        annotations[name] = arg_type
+
+    def arg_type_name(arg_type):
+        return wp.types.get_type_code(wp.types.type_to_warp(arg_type))
+
+    try:
+        Args.__annotations__ = annotations
+    except AttributeError:
+        Args.__dict__.__annotations__ = annotations
+
+    suffix = "_".join([f"{name}_{arg_type_name(arg_type)}" for name, arg_type in annotations.items()])
+
+    return get_struct(Args, suffix=suffix)
+
+
+def populate_argument_struct(Args: wp.codegen.Struct, values: Dict[str, Any], func_name: str):
+    if values is None:
+        values = {}
+
+    value_struct_values = Args()
+    for k, v in values.items():
+        try:
+            setattr(value_struct_values, k, v)
+        except Exception as err:
+            if k not in Args.vars:
+                raise ValueError(
+                    f"Passed value argument '{k}' does not match any of the function '{func_name}' parameters"
+                ) from err
+            raise ValueError(
+                f"Passed value argument '{k}' of type '{wp.types.type_repr(v)}' is incompatible with the function '{func_name}' parameter of type '{wp.types.type_repr(Args.vars[k].type)}'"
+            ) from err
+
+    missing_values = Args.vars.keys() - values.keys()
+    if missing_values:
+        wp.utils.warn(
+            f"Missing values for parameter(s) '{', '.join(missing_values)}' of the function '{func_name}', will be zero-initialized"
+        )
+
+    return value_struct_values
+
+
+class ExpandStarredArgumentStruct(ast.NodeTransformer):
+    def __init__(
+        self,
+        structs: Dict[str, wp.codegen.Struct],
+    ):
+        self._structs = structs
+
+    @staticmethod
+    def _build_path(path, node):
+        if isinstance(node, ast.Attribute):
+            ExpandStarredArgumentStruct._build_path(path, node.value)
+            path.append(node.attr)
+        if isinstance(node, ast.Name):
+            path.append(node.id)
+        return path
+
+    def _get_expanded_struct(self, arg_node):
+        if not isinstance(arg_node, ast.Starred):
+            return None
+        path = ".".join(ExpandStarredArgumentStruct._build_path([], arg_node.value))
+        return self._structs.get(path, None)
+
+    def visit_Call(self, call: ast.Call):
+        call = self.generic_visit(call)
+
+        expanded_args = []
+        for arg in call.args:
+            struct = self._get_expanded_struct(arg)
+            if struct is None:
+                expanded_args.append(arg)
+            else:
+                expanded_args += [ast.Attribute(value=arg.value, attr=field) for field in struct.vars.keys()]
+        call.args = expanded_args
+
+        return call
+
+
 def get_integrand_function(
     integrand: "warp.fem.operator.Integrand",  # noqa: F821
     suffix: str,
@@ -103,9 +192,6 @@ def get_integrand_function(
     annotations=None,
     code_transformers=None,
 ):
-    if code_transformers is None:
-        code_transformers = []
-
     key = _make_key(integrand.func, suffix, use_qualified_name=True)
 
     if key not in _func_cache:
@@ -131,9 +217,6 @@ def get_integrand_kernel(
     if kernel_options is None:
         kernel_options = {}
 
-    if code_transformers is None:
-        code_transformers = []
-
     key = _make_key(integrand.func, suffix, use_qualified_name=True)
 
     if key not in _kernel_cache:
@@ -198,11 +281,22 @@ class Temporary:
     The temporary may also be explicitly returned to the pool before destruction using :meth:`release`.
     """
 
+    def __new__(cls, *args, **kwargs):
+        instance = super(Temporary, cls).__new__(cls)
+        instance._pool = None
+        return instance
+
     def __init__(self, array: wp.array, pool: Optional["TemporaryStore.Pool"] = None, shape=None, dtype=None):
         self._raw_array = array
         self._array_view = array
         self._pool = pool
 
+        if pool is not None and wp.context.runtime.tape is not None:
+            # Extend lifetime to that of Tape (or Pool if shorter)
+            # This is to prevent temporary arrays held in tape launch parameters to be redeemed
+            pool.hold(self)
+            weakref.finalize(wp.context.runtime.tape, TemporaryStore.Pool.stop_holding, pool, self)
+
         if shape is not None or dtype is not None:
             self._view_as(shape=shape, dtype=dtype)
 
@@ -270,6 +364,8 @@ class TemporaryStore:
             self._pool_sizes = []  # Sizes of available arrays for borrowing, ascending
             self._allocs = {}  # All allocated arrays, including borrowed ones
 
+            self._held_temporaries = set()  # Temporaries that are prevented from going out of scope
+
         def borrow(self, shape, dtype, requires_grad: bool):
             size = 1
             if isinstance(shape, int):
@@ -285,8 +381,12 @@ class TemporaryStore:
                 # Big enough array found, remove from pool
                 array = self._pool.pop(index)
                 self._pool_sizes.pop(index)
-                if requires_grad and array.grad is None:
-                    array.requires_grad = True
+                if requires_grad:
+                    if array.grad is None:
+                        array.requires_grad = True
+                    else:
+                        # Zero-out existing gradient to mimic semantics of wp.empty()
+                        array.grad.zero_()
                 return Temporary(pool=self, array=array, shape=shape, dtype=dtype)
 
             # No big enough array found, allocate new one
@@ -312,6 +412,12 @@ class TemporaryStore:
         def detach(self, array):
             del self._allocs[array.ptr]
 
+        def hold(self, temp: Temporary):
+            self._held_temporaries.add(temp)
+
+        def stop_holding(self, temp: Temporary):
+            self._held_temporaries.remove(temp)
+
     def __init__(self):
         self.clear()
 
@@ -401,3 +507,47 @@ def borrow_temporary_like(
         device=array.device,
         requires_grad=array.requires_grad,
     )
+
+
+_device_events = {}
+
+
+def capture_event(device=None):
+    """
+    Records a CUDA event on the current stream and returns it,
+    reusing previously created events if possible.
+
+    If the current device is not a CUDA device, returns ``None``.
+
+    The event can be returned to the shared per-device pool for future reuse by
+    calling :func:`synchronize_event`
+    """
+
+    device = wp.get_device(device)
+    if not device.is_cuda:
+        return None
+
+    try:
+        device_events = _device_events[device.ordinal]
+    except KeyError:
+        device_events = []
+        _device_events[device.ordinal] = device_events
+
+    with wp.ScopedDevice(device):
+        if not device_events:
+            return wp.record_event()
+
+        return wp.record_event(device_events.pop())
+
+
+def synchronize_event(event: Union[wp.Event, None]):
+    """
+    Synchronize an event created with :func:`capture_event` and returns it to the
+    per-device event pool.
+
+    If `event` is ``None``, do nothing.
+    """
+
+    if event is not None:
+        wp.synchronize_event(event)
+        _device_events[event.device.ordinal].append(event)

warp/fem/dirichlet.py

@@ -1,11 +1,10 @@
 from typing import Any, Optional
 
 import warp as wp
+from warp.fem.utils import array_axpy, symmetric_eigenvalues_qr
 from warp.sparse import BsrMatrix, bsr_assign, bsr_axpy, bsr_copy, bsr_mm, bsr_mv
 from warp.types import type_is_matrix, type_length
 
-from .utils import array_axpy
-
 
 def normalize_dirichlet_projector(projector_matrix: BsrMatrix, fixed_value: Optional[wp.array] = None):
     """
@@ -115,11 +114,34 @@ def project_linear_system(
     project_system_matrix(system_matrix, projector_matrix)
 
 
+@wp.func
+def _normalize_projector_and_value(A: Any, b: Any):
+    # Do a modal decomposition of the left and right hand side,
+    # Make lhs an orthogonal projection and apply corresponding scaling to righ-hand-side
+
+    eps = wp.trace(A) * A.dtype(1.0e-6)
+
+    diag, ev = symmetric_eigenvalues_qr(A, eps * eps)
+
+    zero = A.dtype(0)
+    A_unitary = type(A)(zero)
+    b_normalized = type(b)(zero)
+
+    for k in range(b.length):
+        if diag[k] > eps:  # ignore small eigenmodes
+            v = ev[k]
+            A_unitary += wp.outer(v, v)
+            b_normalized += wp.dot(v, b) / diag[k] * v
+
+    return A_unitary, b_normalized
+
+
 @wp.kernel
-def _normalize_dirichlet_projector_kernel(
+def _normalize_dirichlet_projector_and_values_kernel(
     offsets: wp.array(dtype=int),
     columns: wp.array(dtype=int),
     block_values: wp.array(dtype=Any),
+    fixed_values: wp.array(dtype=Any),
 ):
     row = wp.tid()
 
@@ -131,26 +153,21 @@ def _normalize_dirichlet_projector_kernel(
 
     diag = wp.lower_bound(columns, beg, end, row)
 
-    if diag < end and columns[diag] == row:
-        P = block_values[diag]
+    if diag < end:
+        if columns[diag] == row:
+            P = block_values[diag]
 
-        P_sq = P * P
-        trace_P = wp.trace(P)
-        trace_P_sq = wp.trace(P_sq)
+            P_norm, v_norm = _normalize_projector_and_value(P, fixed_values[row])
 
-        if wp.nonzero(trace_P_sq):
-            scale = trace_P / trace_P_sq
-            block_values[diag] = scale * P
-        else:
-            block_values[diag] = P - P
+            block_values[diag] = P_norm
+            fixed_values[row] = v_norm
 
 
 @wp.kernel
-def _normalize_dirichlet_projector_and_values_kernel(
+def _normalize_dirichlet_projector_kernel(
     offsets: wp.array(dtype=int),
     columns: wp.array(dtype=int),
     block_values: wp.array(dtype=Any),
-    fixed_values: wp.array(dtype=Any),
 ):
     row = wp.tid()
 
@@ -162,17 +179,9 @@ def _normalize_dirichlet_projector_and_values_kernel(
 
     diag = wp.lower_bound(columns, beg, end, row)
 
-    if diag < end and columns[diag] == row:
-        P = block_values[diag]
-
-        P_sq = P * P
-        trace_P = wp.trace(P)
-        trace_P_sq = wp.trace(P_sq)
+    if diag < end:
+        if columns[diag] == row:
+            P = block_values[diag]
 
-        if wp.nonzero(trace_P_sq):
-            scale = trace_P / trace_P_sq
-            block_values[diag] = scale * P
-            fixed_values[row] = scale * fixed_values[row]
-        else:
-            block_values[diag] = P - P
-            fixed_values[row] = fixed_values[row] - fixed_values[row]
+            P_norm, v_norm = _normalize_projector_and_value(P, type(P[0])())
+            block_values[diag] = P_norm

warp/fem/domain.py

@@ -1,15 +1,17 @@
-from enum import Enum
-from typing import Union
+from typing import Optional, Union
 
 import warp as wp
 import warp.codegen
 import warp.context
+import warp.fem.cache as cache
+import warp.fem.utils as utils
 from warp.fem.geometry import (
     Element,
     Geometry,
     GeometryPartition,
     WholeGeometryPartition,
 )
+from warp.fem.types import ElementKind
 
 GeometryOrPartition = Union[Geometry, GeometryPartition]
 
@@ -17,12 +19,6 @@ GeometryOrPartition = Union[Geometry, GeometryPartition]
 class GeometryDomain:
     """Interface class for domains, i.e. (partial) views of elements in a Geometry"""
 
-    class ElementKind(Enum):
-        """Possible kinds of elements contained in a domain"""
-
-        CELL = 0
-        SIDE = 1
-
     def __init__(self, geometry: GeometryOrPartition):
         if isinstance(geometry, GeometryPartition):
             self.geometry_partition = geometry
@@ -106,8 +102,8 @@ class Cells(GeometryDomain):
         super().__init__(geometry)
 
     @property
-    def element_kind(self) -> GeometryDomain.ElementKind:
-        return GeometryDomain.ElementKind.CELL
+    def element_kind(self) -> ElementKind:
+        return ElementKind.CELL
 
     @property
     def dimension(self) -> int:
@@ -157,7 +153,7 @@ class Cells(GeometryDomain):
         return self.geometry.cell_measure_ratio
 
     @property
-    def eval_normal(self) -> wp.Function:
+    def element_normal(self) -> wp.Function:
         return self.geometry.cell_normal
 
     @property
@@ -172,9 +168,11 @@ class Sides(GeometryDomain):
         self.geometry = geometry
         super().__init__(geometry)
 
+        self.element_lookup = None
+
     @property
-    def element_kind(self) -> GeometryDomain.ElementKind:
-        return GeometryDomain.ElementKind.SIDE
+    def element_kind(self) -> ElementKind:
+        return ElementKind.SIDE
 
     @property
     def dimension(self) -> int:
@@ -224,7 +222,7 @@ class Sides(GeometryDomain):
         return self.geometry.side_measure_ratio
 
     @property
-    def eval_normal(self) -> wp.Function:
+    def element_normal(self) -> wp.Function:
         return self.geometry.side_normal
 
 
@@ -260,3 +258,96 @@ class FrontierSides(Sides):
     @property
     def element_index(self) -> wp.Function:
         return self.geometry_partition.frontier_side_index
+
+
+class Subdomain(GeometryDomain):
+    """Subdomain -- restriction of domain to a subset of its elements"""
+
+    def __init__(
+        self,
+        domain: GeometryDomain,
+        element_mask: Optional[wp.array] = None,
+        element_indices: Optional[wp.array] = None,
+        temporary_store: Optional[cache.TemporaryStore] = None,
+    ):
+        """
+        Create a subdomain from a subset of elements.
+
+        Exactly one of `element_mask` and `element_indices` should be provided.
+
+        Args:
+            domain: the containing domain
+            element_mask: Array of length ``domain.element_count()`` indicating which elements should be included. Array values must be either ``1`` (selected) or ``0`` (not selected).
+            element_indices: Explicit array of element indices to include
+        """
+
+        super().__init__(domain.geometry_partition)
+
+        if element_indices is None:
+            if element_mask is None:
+                raise ValueError("Either 'element_mask' or 'element_indices' should be provided")
+            element_indices, _ = utils.masked_indices(mask=element_mask, temporary_store=temporary_store)
+            element_indices = element_indices.detach()
+        elif element_mask is not None:
+            raise ValueError("Only one of 'element_mask' and 'element_indices' should be provided")
+
+        self._domain = domain
+        self._element_indices = element_indices
+        self.ElementIndexArg = self._make_element_index_arg()
+        self.element_index = self._make_element_index()
+
+        # forward
+        self.ElementArg = self._domain.ElementArg
+        self.geometry_element_count = self._domain.geometry_element_count
+        self.reference_element = self._domain.reference_element
+        self.element_arg_value = self._domain.element_arg_value
+        self.element_measure = self._domain.element_measure
+        self.element_measure_ratio = self._domain.element_measure_ratio
+        self.element_position = self._domain.element_position
+        self.element_deformation_gradient = self._domain.element_deformation_gradient
+        self.element_lookup = self._domain.element_lookup
+        self.element_normal = self._domain.element_normal
+
+    @property
+    def name(self) -> str:
+        return f"{self._domain.name}_Subdomain"
+
+    def __eq__(self, other) -> bool:
+        return (
+            self.__class__ == other.__class__
+            and self.geometry_partition == other.geometry_partition
+            and self._element_indices == other._element_indices
+        )
+
+    @property
+    def element_kind(self) -> ElementKind:
+        return self._domain.element_kind
+
+    @property
+    def dimension(self) -> int:
+        return self._domain.dimension
+
+    def element_count(self) -> int:
+        return self._element_indices.shape[0]
+
+    def _make_element_index_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class ElementIndexArg:
+            domain_arg: self._domain.ElementIndexArg
+            element_indices: wp.array(dtype=int)
+
+        return ElementIndexArg
+
+    @cache.cached_arg_value
+    def element_index_arg_value(self, device: warp.context.Devicelike):
+        arg = self.ElementIndexArg()
+        arg.domain_arg = self._domain.element_index_arg_value(device)
+        arg.element_indices = self._element_indices.to(device)
+        return arg
+
+    def _make_element_index(self) -> wp.Function:
+        @cache.dynamic_func(suffix=self.name)
+        def element_index(arg: self.ElementIndexArg, index: int):
+            return self._domain.element_index(arg.domain_arg, arg.element_indices[index])
+
+        return element_index

warp/fem/field/__init__.py

@@ -3,7 +3,7 @@ from typing import Optional, Union
 from warp.fem.domain import Cells, GeometryDomain
 from warp.fem.space import FunctionSpace, SpacePartition, SpaceRestriction, make_space_partition, make_space_restriction
 
-from .field import DiscreteField, FieldLike, SpaceField
+from .field import DiscreteField, FieldLike, GeometryField, ImplicitField, NonconformingField, SpaceField, UniformField
 from .nodal_field import NodalField
 from .restriction import FieldRestriction
 from .test import TestField
@@ -85,8 +85,8 @@ def make_trial(
     """
 
     if space_restriction is not None:
-        domain = space.domain
-        space_partition = space.space_partition
+        domain = space_restriction.domain
+        space_partition = space_restriction.space_partition
 
     if space_partition is None:
         if domain is None:
@@ -98,3 +98,14 @@ def make_trial(
         domain = Cells(geometry=space_partition.geo_partition)
 
     return TrialField(space, space_partition, domain)
+
+
+def make_discrete_field(
+    space: FunctionSpace,
+    space_partition: Optional[SpacePartition] = None,
+) -> DiscreteField:
+    """Constructs  a zero-initialized discrete field over a function space or partition
+
+    See also: :meth:`warp.fem.FunctionSpace.make_field`
+    """
+    return space.make_field(space_partition=space_partition)