warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/fem/domain.py

@@ -0,0 +1,263 @@
+from typing import Union
+from enum import Enum
+
+import warp as wp
+import warp.codegen
+import warp.context
+
+from warp.fem.geometry import (
+    Element,
+    Geometry,
+    GeometryPartition,
+    WholeGeometryPartition,
+)
+
+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
+        else:
+            self.geometry_partition = WholeGeometryPartition(geometry)
+
+        self.geometry = self.geometry_partition.geometry
+
+    @property
+    def name(self) -> str:
+        return f"{self.geometry_partition.name}_{self.__class__.__name__}"
+
+    def __str__(self) -> str:
+        return self.name
+
+    def __eq__(self, other) -> bool:
+        return self.__class__ == other.__class__ and self.geometry_partition == other.geometry_partition
+
+    @property
+    def element_kind(self) -> ElementKind:
+        """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
+
+    def element_count(self) -> int:
+        """Number of elements in the domain"""
+        raise NotImplementedError
+
+    def geometry_element_count(self) -> int:
+        """Number of elements in the underlying geometry"""
+        return self.geometry.cell_count()
+
+    def reference_element(self) -> Element:
+        """Protypical element"""
+        raise NotImplementedError
+
+    def element_index_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+        """Value of the argument to be passed to device functions"""
+        raise NotImplementedError
+
+    def element_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+        """Value of the argument to be passed to device functions"""
+        raise NotImplementedError
+
+    ElementIndexArg: warp.codegen.Struct
+    """Structure containing arguments to be passed to device functions computing element indices"""
+
+    element_index: wp.Function
+    """Device function for retrieving an ElementIndex from a linearized index"""
+
+    ElementArg: warp.codegen.Struct
+    """Structure containing arguments to be passed to device functions computing element geometry"""
+
+    element_measure: wp.Function
+    """Device function returning the measure determinant (e.g. volume, area) at a given point"""
+
+    element_measure_ratio: wp.Function
+    """Device function returning the ratio of the measure of a side to that of its neighbour cells"""
+
+    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"""
+
+    element_lookup: wp.Function
+    """Device function returning the sample point corresponding to a world position"""
+
+
+class Cells(GeometryDomain):
+    """A Domain containing all cells of the geometry or geometry partition"""
+
+    def __init__(self, geometry: GeometryOrPartition):
+        super().__init__(geometry)
+
+    @property
+    def element_kind(self) -> GeometryDomain.ElementKind:
+        return GeometryDomain.ElementKind.CELL
+
+    @property
+    def dimension(self) -> int:
+        return self.geometry.dimension
+
+    def reference_element(self) -> Element:
+        return self.geometry.reference_cell()
+
+    def element_count(self) -> int:
+        return self.geometry_partition.cell_count()
+
+    def geometry_element_count(self) -> int:
+        return self.geometry.cell_count()
+
+    @property
+    def ElementIndexArg(self) -> warp.codegen.Struct:
+        return self.geometry_partition.CellArg
+
+    def element_index_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+        return self.geometry_partition.cell_arg_value(device)
+
+    @property
+    def element_index(self) -> wp.Function:
+        return self.geometry_partition.cell_index
+
+    def element_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+        return self.geometry.cell_arg_value(device)
+
+    @property
+    def ElementArg(self) -> warp.codegen.Struct:
+        return self.geometry.CellArg
+
+    @property
+    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
+
+    @property
+    def element_measure_ratio(self) -> wp.Function:
+        return self.geometry.cell_measure_ratio
+
+    @property
+    def eval_normal(self) -> wp.Function:
+        return self.geometry.cell_normal
+
+    @property
+    def element_lookup(self) -> wp.Function:
+        return self.geometry.cell_lookup
+
+
+class Sides(GeometryDomain):
+    """A Domain containing all (interior and boundary) sides of the geometry or geometry partition"""
+
+    def __init__(self, geometry: GeometryOrPartition):
+        self.geometry = geometry
+        super().__init__(geometry)
+
+    @property
+    def element_kind(self) -> GeometryDomain.ElementKind:
+        return GeometryDomain.ElementKind.SIDE
+
+    @property
+    def dimension(self) -> int:
+        return self.geometry.dimension - 1
+
+    def reference_element(self) -> Element:
+        return self.geometry.reference_side()
+
+    def element_count(self) -> int:
+        return self.geometry_partition.side_count()
+
+    def geometry_element_count(self) -> int:
+        return self.geometry.side_count()
+
+    @property
+    def ElementIndexArg(self) -> warp.codegen.Struct:
+        return self.geometry_partition.SideArg
+
+    def element_index_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+        return self.geometry_partition.side_arg_value(device)
+
+    @property
+    def element_index(self) -> wp.Function:
+        return self.geometry_partition.side_index
+
+    @property
+    def ElementArg(self) -> warp.codegen.Struct:
+        return self.geometry.SideArg
+
+    def element_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+        return self.geometry.side_arg_value(device)
+
+    @property
+    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
+
+    @property
+    def element_measure_ratio(self) -> wp.Function:
+        return self.geometry.side_measure_ratio
+
+    @property
+    def eval_normal(self) -> wp.Function:
+        return self.geometry.side_normal
+
+
+class BoundarySides(Sides):
+    """A Domain containing boundary sides of the geometry or geometry partition"""
+
+    def __init__(self, geometry: GeometryOrPartition):
+        super().__init__(geometry)
+
+    def element_count(self) -> int:
+        return self.geometry_partition.boundary_side_count()
+
+    def geometry_element_count(self) -> int:
+        return self.geometry.boundary_side_count()
+
+    @property
+    def element_index(self) -> wp.Function:
+        return self.geometry_partition.boundary_side_index
+
+
+class FrontierSides(Sides):
+    """A Domain containing frontier sides of the geometry partition (sides shared with at least another partition)"""
+
+    def __init__(self, geometry: GeometryOrPartition):
+        super().__init__(geometry)
+
+    def element_count(self) -> int:
+        return self.geometry_partition.frontier_side_count()
+
+    def geometry_element_count(self) -> int:
+        raise RuntimeError("Frontier sides not defined at the geometry level")
+
+    @property
+    def element_index(self) -> wp.Function:
+        return self.geometry_partition.frontier_side_index

warp/fem/field/__init__.py

@@ -0,0 +1,101 @@
+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 .field import DiscreteField, SpaceField, FieldLike
+from .restriction import FieldRestriction
+from .test import TestField
+from .trial import TrialField
+
+from .nodal_field import NodalField
+
+
+def make_restriction(
+    field: DiscreteField,
+    space_restriction: Optional[SpaceRestriction] = None,
+    domain: Optional[GeometryDomain] = None,
+    device=None,
+) -> FieldRestriction:
+    """
+    Restricts a discrete field to a subset of elements.
+
+    Args:
+        field: the discrete field to restrict
+        space_restriction: the function space restriction defining the subset of elements to consider
+        domain: if ``space_restriction`` is not provided, the :py:class:`Domain` defining the subset of elements to consider
+        device: Warp device on which to perform and store computations
+
+    Returns:
+        the field restriction
+    """
+
+    if space_restriction is None:
+        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: 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
+        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 space_restriction is None:
+        space_restriction = make_space_restriction(
+            space_topology=space.topology, space_partition=space_partition, domain=domain, device=device
+        )
+
+    return TestField(space_restriction=space_restriction, space=space)
+
+
+def make_trial(
+    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_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 space_restriction is not None:
+        domain = space.domain
+        space_partition = space.space_partition
+
+    if space_partition is None:
+        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)

warp/fem/field/field.py

@@ -0,0 +1,149 @@
+from typing import Any
+
+import warp as wp
+
+from warp.fem.types import Sample
+from warp.fem.space import FunctionSpace, SpacePartition
+from warp.fem.geometry import Geometry, DeformedGeometry
+
+
+class FieldLike:
+    """Base class for integrable fields"""
+
+    EvalArg: wp.codegen.Struct
+    """Structure containing field-level arguments passed to device functions for field evaluation"""
+
+    ElementEvalArg: wp.codegen.Struct
+    """Structure combining geometry-level and field-level arguments passed to device functions for field evaluation"""
+
+    def eval_arg_value(self, device) -> "EvalArg":
+        """Value of the field-level arguments to be passed to device functions"""
+        raise NotImplementedError
+
+    @property
+    def degree(self) -> int:
+        """Polynomial degree of the field, used to estimate necessary quadrature order"""
+        raise NotImplementedError
+
+    @property
+    def name(self) -> str:
+        raise NotImplementedError
+
+    @property
+    def __str__(self) -> str:
+        return self.name
+
+    def eval_arg_value(self, device):
+        """Value of arguments to be passed to device functions"""
+        raise NotImplementedError
+
+    @staticmethod
+    def eval_inner(args: "ElementEvalArg", s: "Sample"):
+        """Device function evaluating the inner field value at a sample point"""
+        raise NotImplementedError
+
+    @staticmethod
+    def eval_grad_inner(args: "ElementEvalArg", s: "Sample"):
+        """Device function evaluating the inner field gradient at a sample point"""
+        raise NotImplementedError
+
+    @staticmethod
+    def eval_div_inner(args: "ElementEvalArg", s: "Sample"):
+        """Device function evaluating the inner field divergence at a sample point"""
+        raise NotImplementedError
+
+    @staticmethod
+    def eval_outer(args: "ElementEvalArg", s: "Sample"):
+        """Device function evaluating the outer field value at a sample point"""
+        raise NotImplementedError
+
+    @staticmethod
+    def eval_grad_outer(args: "ElementEvalArg", s: "Sample"):
+        """Device function evaluating the outer field gradient at a sample point"""
+        raise NotImplementedError
+
+    @staticmethod
+    def eval_div_outer(args: "ElementEvalArg", s: "Sample"):
+        """Device function evaluating the outer field divergence at a sample point"""
+        raise NotImplementedError
+
+
+class SpaceField(FieldLike):
+    """Base class for fields defined over a function space"""
+
+    def __init__(self, space: FunctionSpace, space_partition: SpacePartition):
+        self._space = space
+        self._space_partition = space_partition
+
+    @property
+    def space(self) -> FunctionSpace:
+        return self._space
+
+    @property
+    def space_partition(self) -> SpacePartition:
+        return self._space_partition
+
+    @property
+    def degree(self) -> int:
+        return self.space.degree
+
+    @property
+    def dtype(self) -> type:
+        return self.space.dtype
+
+    @property
+    def dof_dtype(self) -> type:
+        return self.space.dof_dtype
+
+    def gradient_valid(self) -> bool:
+        """Whether gradient operator can be computed. Only for scalar and vector fields as higher-order tensors are not support yet"""
+        return not wp.types.type_is_matrix(self.dtype)
+
+    def divergence_valid(self) -> bool:
+        """Whether divergence of this field can be computed. Only for vector and tensor fields with same dimension as embedding geometry"""
+        if wp.types.type_is_vector(self.dtype):
+            return wp.types.type_length(self.dtype) == self.space.geometry.dimension
+        if wp.types.type_is_matrix(self.dtype):
+            return self.dtype._shape_[0] == self.space.geometry.dimension
+        return False
+
+    def _make_eval_degree(self):
+        ORDER = self.space.ORDER
+        from warp.fem import cache
+
+        @cache.dynamic_func(suffix=self.name)
+        def degree(args: self.ElementEvalArg):
+            return ORDER
+
+        return degree
+
+
+class DiscreteField(SpaceField):
+    """Explicitly-valued field defined over a partition of a discrete function space"""
+
+    @property
+    def dof_values(self) -> wp.array:
+        """Array of degrees of freedom values"""
+        raise NotImplementedError
+
+    @dof_values.setter
+    def dof_values(self, values: wp.array):
+        """Sets degrees of freedom values from an array"""
+        raise NotImplementedError
+
+    def trace(self) -> "DiscreteField":
+        """Trace of this field over a lower-dimensional function space"""
+        raise NotImplementedError
+
+    @staticmethod
+    def set_node_value(args: "FieldLike.EvalArg", node_index: int, value: Any):
+        """Device function setting the value at given node"""
+        raise NotImplementedError
+
+    @property
+    def name(self) -> str:
+        return f"{self.__class__.__qualname__}_{self.space.name}_{self.space_partition.name}"
+
+    def make_deformed_geometry(self) -> Geometry:
+        """Returns a deformed version of the underlying geometry using this field's values as displacement"""
+        return DeformedGeometry(self)