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

warp/fem/__init__.py

@@ -1,9 +1,27 @@
+from .geometry import Geometry, Grid2D, Trimesh2D, Quadmesh2D, Grid3D, Tetmesh, Hexmesh
+from .geometry import GeometryPartition, LinearGeometryPartition, ExplicitGeometryPartition
 
+from .space import FunctionSpace, make_polynomial_space, ElementBasis
+from .space import BasisSpace, PointBasisSpace, make_polynomial_basis_space, make_collocated_function_space
+from .space import DofMapper, SkewSymmetricTensorMapper, SymmetricTensorMapper
+from .space import SpaceTopology, SpacePartition, SpaceRestriction, make_space_partition, make_space_restriction
 
-    
+from .domain import GeometryDomain, Cells, Sides, BoundarySides, FrontierSides
+from .quadrature import Quadrature, RegularQuadrature, NodalQuadrature, ExplicitQuadrature, PicQuadrature
+from .polynomial import Polynomial
 
-# Shared DOFs : vertices or edges (sides in 3D)
-# DG: Interior DOFs only 
-# Unique identifier -> vertex or edge, side
-# Element -> DOF id
-    
+from .field import FieldLike, DiscreteField, make_test, make_trial, make_restriction
+
+from .integrate import integrate, interpolate
+
+from .operator import integrand
+from .operator import position, normal, lookup, measure, measure_ratio, deformation_gradient
+from .operator import inner, grad, div, outer, grad_outer, div_outer
+from .operator import degree, at_node
+from .operator import D, curl, jump, average, grad_jump, grad_average
+
+from .types import Sample, Field, Domain, Coords, ElementIndex
+
+from .dirichlet import project_linear_system, normalize_dirichlet_projector
+
+from .cache import TemporaryStore, set_default_temporary_store, borrow_temporary, borrow_temporary_like

warp/fem/cache.py

@@ -1,22 +1,26 @@
-from typing import Callable, Optional
+from typing import Callable, Optional, Union, Tuple, Dict, Any
+from copy import copy
+import bisect
+import re
 
-import warp as wp
 
-from warp.fem.operator import Integrand
+import warp as wp
 
-import re
 
 _kernel_cache = dict()
 _struct_cache = dict()
 _func_cache = dict()
 
-
 _key_re = re.compile("[^0-9a-zA-Z_]+")
 
 
-def get_func(func, suffix=""):
-    key = f"{func.__name__}_{suffix}"
-    key = _key_re.sub("", key)
+def _make_key(obj, suffix: str, use_qualified_name):
+    base_name = f"{obj.__module__}.{obj.__qualname__}" if use_qualified_name else obj.__name__
+    return _key_re.sub("", f"{base_name}_{suffix}")
+
+
+def get_func(func, suffix: str, use_qualified_name: bool = False):
+    key = _make_key(func, suffix, use_qualified_name)
 
     if key not in _func_cache:
         _func_cache[key] = wp.Function(
@@ -31,23 +35,49 @@ def get_func(func, suffix=""):
     return _func_cache[key]
 
 
-def get_kernel(func, suffix=""):
-    module = wp.get_module(func.__module__)
-    key = func.__name__ + "_" + suffix
-    key = _key_re.sub("", key)
+def dynamic_func(suffix: str, use_qualified_name=False):
+    def wrap_func(func: Callable):
+        return get_func(func, suffix=suffix, use_qualified_name=use_qualified_name)
+
+    return wrap_func
+
+
+def get_kernel(
+    func,
+    suffix: str,
+    use_qualified_name: bool = False,
+    kernel_options: Dict[str, Any] = {},
+):
+    key = _make_key(func, suffix, use_qualified_name)
 
     if key not in _kernel_cache:
+        # Avoid creating too long file names -- can lead to issues on Windows
+        # We could hash the key, but prefer to keep it human-readable
+        module_name = f"{func.__module__}.dyn.{key}"
+        module_name = module_name[:128] if len(module_name) > 128 else module_name
+        module = wp.get_module(module_name)
+        module.options = copy(wp.get_module(func.__module__).options)
+        module.options.update(kernel_options)
         _kernel_cache[key] = wp.Kernel(func=func, key=key, module=module)
     return _kernel_cache[key]
 
 
-def get_struct(Fields):
-    module = wp.get_module(Fields.__module__)
-    key = _key_re.sub("", Fields.__qualname__)
+def dynamic_kernel(suffix: str, use_qualified_name=False, kernel_options: Dict[str, Any] = {}):
+    def wrap_kernel(func: Callable):
+        return get_kernel(func, suffix=suffix, use_qualified_name=use_qualified_name, kernel_options=kernel_options)
+
+    return wrap_kernel
+
+
+def get_struct(struct: type, suffix: str, use_qualified_name: bool = False):
+    key = _make_key(struct, suffix, use_qualified_name)
+    # used in codegen
+    struct.__qualname__ = key
 
     if key not in _struct_cache:
+        module = wp.get_module(struct.__module__)
         _struct_cache[key] = wp.codegen.Struct(
-            cls=Fields,
+            cls=struct,
             key=key,
             module=module,
         )
@@ -55,18 +85,25 @@ def get_struct(Fields):
     return _struct_cache[key]
 
 
+def dynamic_struct(suffix: str, use_qualified_name=False):
+    def wrap_struct(struct: type):
+        return get_struct(struct, suffix=suffix, use_qualified_name=use_qualified_name)
+
+    return wrap_struct
+
+
 def get_integrand_function(
-    integrand: Integrand,
+    integrand: "warp.fem.operator.Integrand",
     suffix: str,
+    func=None,
     annotations=None,
     code_transformers=[],
 ):
-    key = integrand.name + suffix
-    key = _key_re.sub("", key)
+    key = _make_key(integrand.func, suffix, use_qualified_name=True)
 
     if key not in _func_cache:
         _func_cache[key] = wp.Function(
-            func=integrand.func,
+            func=integrand.func if func is None else func,
             key=key,
             namespace="",
             module=integrand.module,
@@ -78,19 +115,275 @@ def get_integrand_function(
 
 
 def get_integrand_kernel(
-    integrand: Integrand,
+    integrand: "warp.fem.operator.Integrand",
     suffix: str,
     kernel_fn: Optional[Callable] = None,
+    kernel_options: Dict[str, Any] = {},
     code_transformers=[],
 ):
-    module = wp.get_module(f"{integrand.module.name}.{integrand.name}")
-    module.options = integrand.module.options
-    key = integrand.name + "_" + suffix
-    key = _key_re.sub("", key)
+    key = _make_key(integrand.func, suffix, use_qualified_name=True)
 
     if key not in _kernel_cache:
         if kernel_fn is None:
             return None
 
+        module = wp.get_module(f"{integrand.module.name}.{integrand.name}")
+        module.options = copy(integrand.module.options)
+        module.options.update(kernel_options)
+
         _kernel_cache[key] = wp.Kernel(func=kernel_fn, key=key, module=module, code_transformers=code_transformers)
     return _kernel_cache[key]
+
+
+def cached_arg_value(func: Callable):
+    """Decorator to be applied to member methods assembling Arg structs, so that the result gets
+    automatically cached for the lifetime of the parent object
+    """
+
+    cache_attr = f"_{func.__name__}_cache"
+
+    def get_arg(obj, device):
+        if not hasattr(obj, cache_attr):
+            setattr(obj, cache_attr, {})
+
+        cache = getattr(obj, cache_attr, {})
+
+        device = wp.get_device(device)
+        if device.ordinal not in cache:
+            cache[device.ordinal] = func(obj, device)
+
+        return cache[device.ordinal]
+
+    return get_arg
+
+
+_cached_vec_types = {}
+_cached_mat_types = {}
+
+
+def cached_vec_type(length, dtype):
+    key = (length, dtype)
+    if key not in _cached_vec_types:
+        _cached_vec_types[key] = wp.vec(length=length, dtype=dtype)
+
+    return _cached_vec_types[key]
+
+
+def cached_mat_type(shape, dtype):
+    key = (*shape, dtype)
+    if key not in _cached_mat_types:
+        _cached_mat_types[key] = wp.mat(shape=shape, dtype=dtype)
+
+    return _cached_mat_types[key]
+
+
+class Temporary:
+    """Handle over a temporary array from a :class:`TemporaryStore`.
+
+    The array will be automatically returned to the temporary pool for re-use upon destruction of this object, unless
+    the temporary is explicitly detached from the pool using :meth:`detach`.
+    The temporary may also be explicitly returned to the pool before destruction using :meth:`release`.
+    """
+
+    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 shape is not None or dtype is not None:
+            self._view_as(shape=shape, dtype=dtype)
+
+    def detach(self) -> wp.array:
+        """Detaches the temporary so it is never returned to the pool"""
+        if self._pool is not None:
+            self._pool.detach(self._raw_array)
+
+        self._pool = None
+        return self._array_view
+
+    def release(self):
+        """Returns the temporary array to the pool"""
+        if self._pool is not None:
+            self._pool.redeem(self._raw_array)
+
+        self._pool = None
+
+    @property
+    def array(self) -> wp.array:
+        """View of the array with desired shape and data type."""
+        return self._array_view
+
+    def _view_as(self, shape, dtype) -> "Temporary":
+        def _view_reshaped_truncated(array):
+            return wp.types.array(
+                ptr=array.ptr,
+                dtype=dtype,
+                shape=shape,
+                device=array.device,
+                pinned=array.pinned,
+                capacity=array.capacity,
+                copy=False,
+                owner=False,
+                grad=None if array.grad is None else _view_reshaped_truncated(array.grad),
+            )
+
+        self._array_view = _view_reshaped_truncated(self._raw_array)
+        return self
+
+    def __del__(self):
+        self.release()
+
+
+class TemporaryStore:
+    """
+    Shared pool of temporary arrays that will be persisted and reused across invocations of ``warp.fem`` functions.
+
+    A :class:`TemporaryStore` instance may either be passed explicitly to ``warp.fem`` functions that accept such an argument, for instance :func:`.integrate.integrate`,
+    or can be set globally as the default store using :func:`set_default_temporary_store`.
+
+    By default, there is no default temporary store, so that temporary allocations are not persisted.
+    """
+
+    _default_store: "TemporaryStore" = None
+
+    class Pool:
+        def __init__(self, dtype, device, pinned: bool):
+            self.dtype = dtype
+            self.device = device
+            self.pinned = pinned
+
+            self._pool = []  # Currently available arrays for borrowing, ordered by size
+            self._pool_sizes = []  # Sizes of available arrays for borrowing, ascending
+            self._allocs = {}  # All allocated arrays, including borrowed ones
+
+        def borrow(self, shape, dtype, requires_grad: bool):
+            size = 1
+            if isinstance(shape, int):
+                shape = (shape,)
+            for d in shape:
+                size *= d
+
+            index = bisect.bisect_left(
+                a=self._pool_sizes,
+                x=size,
+            )
+            if index < len(self._pool):
+                # 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
+                return Temporary(pool=self, array=array, shape=shape, dtype=dtype)
+
+            # No big enough array found, allocate new one
+            if len(self._pool) > 0:
+                grow_factor = 1.5
+                size = max(int(self._pool_sizes[-1] * grow_factor), size)
+
+            array = wp.empty(
+                shape=(size,), dtype=self.dtype, pinned=self.pinned, device=self.device, requires_grad=requires_grad
+            )
+            self._allocs[array.ptr] = array
+            return Temporary(pool=self, array=array, shape=shape, dtype=dtype)
+
+        def redeem(self, array):
+            # Insert back array into available pool
+            index = bisect.bisect_left(
+                a=self._pool_sizes,
+                x=array.size,
+            )
+            self._pool.insert(index, array)
+            self._pool_sizes.insert(index, array.size)
+
+        def detach(self, array):
+            del self._allocs[array.ptr]
+
+    def __init__(self):
+        self.clear()
+
+    def clear(self):
+        self._temporaries = {}
+
+    def borrow(self, shape, dtype, pinned: bool = False, device=None, requires_grad: bool = False) -> Temporary:
+        dtype = wp.types.type_to_warp(dtype)
+        device = wp.get_device(device)
+
+        type_length = wp.types.type_length(dtype)
+        key = (dtype._type_, type_length, pinned, device.ordinal)
+
+        pool = self._temporaries.get(key, None)
+        if pool is None:
+            value_type = (
+                cached_vec_type(length=type_length, dtype=wp.types.type_scalar_type(dtype))
+                if type_length > 1
+                else dtype
+            )
+            pool = TemporaryStore.Pool(value_type, device, pinned=pinned)
+            self._temporaries[key] = pool
+
+        return pool.borrow(dtype=dtype, shape=shape, requires_grad=requires_grad)
+
+
+def set_default_temporary_store(temporary_store: Optional[TemporaryStore]):
+    """Globally sets the default :class:`TemporaryStore` instance to use for temporary allocations in ``warp.fem`` functions.
+
+    If the default temporary store is set to ``None``, temporary allocations are not persisted unless a :class:`TemporaryStore` is provided at a per-function granularity.
+    """
+
+    TemporaryStore._default_store = temporary_store
+
+
+def borrow_temporary(
+    temporary_store: Optional[TemporaryStore],
+    shape: Union[int, Tuple[int]],
+    dtype: type,
+    pinned: bool = False,
+    requires_grad: bool = False,
+    device=None,
+) -> Temporary:
+    """
+    Borrows and returns a temporary array with specified attributes from a shared pool.
+
+    If an array with sufficient capacity and matching desired attributes is already available in the pool, it will be returned.
+    Otherwise, a new allocation wil be performed.
+
+    Args:
+        temporary_store: the shared pool to borrow the temporary from. If `temporary_store` is ``None``, the global default temporary store, if set, will be used.
+        shape: desired dimensions for the temporary array
+        dtype: desired data type for the temporary array
+        pinned: whether a pinned allocation is desired
+        device: device on which the momory should be allocated; if ``None``, the current device will be used.
+    """
+
+    if temporary_store is None:
+        temporary_store = TemporaryStore._default_store
+
+    if temporary_store is None:
+        return Temporary(
+            array=wp.empty(shape=shape, dtype=dtype, pinned=pinned, device=device, requires_grad=requires_grad)
+        )
+
+    return temporary_store.borrow(shape=shape, dtype=dtype, device=device, pinned=pinned, requires_grad=requires_grad)
+
+
+def borrow_temporary_like(
+    array: Union[wp.array, Temporary],
+    temporary_store: Optional[TemporaryStore],
+) -> Temporary:
+    """
+    Borrows and returns a temporary array with the same attributes as another array or temporary.
+
+    Args:
+        array: Warp or temporary array to read the desired attributes from
+        temporary_store: the shared pool to borrow the temporary from. If `temporary_store` is ``None``, the global default temporary store, if set, will be used.
+    """
+    if isinstance(array, Temporary):
+        array = array.array
+    return borrow_temporary(
+        temporary_store=temporary_store,
+        shape=array.shape,
+        dtype=array.dtype,
+        pinned=array.pinned,
+        device=array.device,
+        requires_grad=array.requires_grad,
+    )

warp/fem/dirichlet.py

@@ -62,7 +62,7 @@ def normalize_dirichlet_projector(projector_matrix: BsrMatrix, fixed_value: Opti
 
 
 def project_system_rhs(
-    system_matrix: BsrMatrix, system_rhs: wp.array, projector_matrix: BsrMatrix, fixed_value: wp.array
+    system_matrix: BsrMatrix, system_rhs: wp.array, projector_matrix: BsrMatrix, fixed_value: Optional[wp.array] = None
 ):
     """Projects the right-hand-side of a linear system to enforce Dirichlet boundary conditions
 
@@ -72,7 +72,11 @@ def project_system_rhs(
     rhs_tmp = wp.empty_like(system_rhs)
     rhs_tmp.assign(system_rhs)
 
-    bsr_mv(A=projector_matrix, x=fixed_value, y=system_rhs, alpha=1.0, beta=0.0)
+    if fixed_value is None:
+        system_rhs.zero_()
+    else:
+        bsr_mv(A=projector_matrix, x=fixed_value, y=system_rhs, alpha=1.0, beta=0.0)
+
     bsr_mv(A=system_matrix, x=system_rhs, y=rhs_tmp, alpha=-1.0, beta=1.0)
 
     # here rhs_tmp = system_rhs - system_matrix * projector * fixed_value
@@ -99,7 +103,7 @@ def project_linear_system(
     system_matrix: BsrMatrix,
     system_rhs: wp.array,
     projector_matrix: BsrMatrix,
-    fixed_value: wp.array,
+    fixed_value: Optional[wp.array] = None,
     normalize_projector=True,
 ):
     """

warp/fem/domain.py

@@ -47,6 +47,7 @@ class GeometryDomain:
         """Kind of elements that this domain contains (cells or sides)"""
         raise NotImplementedError
 
+    @property
     def dimension(self) -> int:
         """Dimension of the elements of the domain"""
         raise NotImplementedError
@@ -89,6 +90,9 @@ class GeometryDomain:
     element_position: wp.Function
     """Device function returning the element position at a sample point"""
 
+    element_deformation_gradient: wp.Function
+    """Device function returning the gradient of the position with respect to the element's reference space"""
+
     element_normal: wp.Function
     """Device function returning the element normal at a sample point"""
 
@@ -106,6 +110,7 @@ class Cells(GeometryDomain):
     def element_kind(self) -> GeometryDomain.ElementKind:
         return GeometryDomain.ElementKind.CELL
 
+    @property
     def dimension(self) -> int:
         return self.geometry.dimension
 
@@ -140,6 +145,10 @@ class Cells(GeometryDomain):
     def element_position(self) -> wp.Function:
         return self.geometry.cell_position
 
+    @property
+    def element_deformation_gradient(self) -> wp.Function:
+        return self.geometry.cell_deformation_gradient
+
     @property
     def element_measure(self) -> wp.Function:
         return self.geometry.cell_measure
@@ -168,6 +177,7 @@ class Sides(GeometryDomain):
     def element_kind(self) -> GeometryDomain.ElementKind:
         return GeometryDomain.ElementKind.SIDE
 
+    @property
     def dimension(self) -> int:
         return self.geometry.dimension - 1
 
@@ -202,6 +212,10 @@ class Sides(GeometryDomain):
     def element_position(self) -> wp.Function:
         return self.geometry.side_position
 
+    @property
+    def element_deformation_gradient(self) -> wp.Function:
+        return self.geometry.side_deformation_gradient
+
     @property
     def element_measure(self) -> wp.Function:
         return self.geometry.side_measure

warp/fem/field/__init__.py

@@ -3,15 +3,13 @@ from typing import Union, Optional
 from warp.fem.domain import GeometryDomain, Cells
 from warp.fem.space import FunctionSpace, SpaceRestriction, SpacePartition, make_space_partition, make_space_restriction
 
-from .discrete_field import DiscreteField
+from .field import DiscreteField, SpaceField, FieldLike
 from .restriction import FieldRestriction
 from .test import TestField
 from .trial import TrialField
 
 from .nodal_field import NodalField
 
-FieldLike = Union[DiscreteField, FieldRestriction, TestField, TrialField]
-
 
 def make_restriction(
     field: DiscreteField,
@@ -33,77 +31,71 @@ def make_restriction(
     """
 
     if space_restriction is None:
-        space_restriction = make_space_restriction(
-            space=field.space, space_partition=field.space_partition, domain=domain, device=device
-        )
+        space_restriction = make_space_restriction(space_partition=field.space_partition, domain=domain, device=device)
 
     return FieldRestriction(field=field, space_restriction=space_restriction)
 
 
 def make_test(
-    space: Union[FunctionSpace, SpaceRestriction] = None,
-    space_partition: SpacePartition = None,
-    domain: GeometryDomain = None,
+    space: FunctionSpace,
+    space_restriction: Optional[SpaceRestriction] = None,
+    space_partition: Optional[SpacePartition] = None,
+    domain: Optional[GeometryDomain] = None,
     device=None,
 ) -> TestField:
     """
     Constructs a test field over a function space or its restriction
 
     Args:
-        space: the function space or function space restriction
-        space_partition: if ``space`` is a whole function space, the optional subset of node indices to consider
-        domain: if ``space`` is a whole function space, the optional subset of elements to consider
+        space: the function space
+        space_restriction: restriction of the space topology to a domain
+        space_partition: if `space_restriction` is ``None``, the optional subset of node indices to consider
+        domain: if `space_restriction` is ``None``, optional subset of elements to consider
         device: Warp device on which to perform and store computations
 
     Returns:
         the test field
     """
 
-    if not isinstance(space, SpaceRestriction):
-        if space is None:
-            space = space_partition.space
-
-        if domain is None:
-            domain = Cells(geometry=space.geometry)
-
-        if space_partition is None:
-            space_partition = make_space_partition(space, domain.geometry_partition)
-
-        space = make_space_restriction(space=space, space_partition=space_partition, domain=domain, device=device)
+    if space_restriction is None:
+        space_restriction = make_space_restriction(
+            space_topology=space.topology, space_partition=space_partition, domain=domain, device=device
+        )
 
-    return TestField(space)
+    return TestField(space_restriction=space_restriction, space=space)
 
 
 def make_trial(
-    space: Union[FunctionSpace, SpaceRestriction] = None,
-    space_partition: SpacePartition = None,
-    domain: GeometryDomain = None,
+    space: FunctionSpace,
+    space_restriction: Optional[SpaceRestriction] = None,
+    space_partition: Optional[SpacePartition] = None,
+    domain: Optional[GeometryDomain] = None,
 ) -> TrialField:
     """
     Constructs a trial field over a function space or partition
 
     Args:
         space: the function space or function space restriction
-        space_partition: if ``space`` is a whole function space, the optional subset of node indices to consider
-        domain: if ``space`` is a whole function space, the optional subset of elements to consider
+        space_restriction: restriction of the space topology to a domain
+        space_partition: if `space_restriction` is ``None``, the optional subset of node indices to consider
+        domain: if `space_restriction` is ``None``, optional subset of elements to consider
         device: Warp device on which to perform and store computations
 
     Returns:
         the trial field
     """
 
-    if isinstance(space, SpaceRestriction):
+    if space_restriction is not None:
         domain = space.domain
-        space = space.space
         space_partition = space.space_partition
 
-    if space is None:
-        space = space_partition.space
-
-    if domain is None:
-        domain = Cells(geometry=space.geometry)
-
     if space_partition is None:
-        space_partition = make_space_partition(space, domain.geometry_partition)
+        if domain is None:
+            domain = Cells(geometry=space.geometry)
+        space_partition = make_space_partition(
+            space_topology=space.topology, geometry_partition=domain.geometry_partition
+        )
+    elif domain is None:
+        domain = Cells(geometry=space_partition.geo_partition)
 
     return TrialField(space, space_partition, domain)