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

warp/fem/quadrature/quadrature.py

@@ -0,0 +1,294 @@
+from typing import Any
+
+import warp as wp
+
+from warp.fem import domain, cache
+from warp.fem.types import ElementIndex, Coords
+from warp.fem.space import FunctionSpace
+
+from ..polynomial import Polynomial
+
+
+class Quadrature:
+    """Interface class for quadrature rules"""
+
+    @wp.struct
+    class Arg:
+        """Structure containing arguments to be passed to device functions"""
+
+        pass
+
+    def __init__(self, domain: domain.GeometryDomain):
+        self._domain = domain
+
+    @property
+    def domain(self):
+        """Domain over which this quadrature is defined"""
+        return self._domain
+
+    def arg_value(self, device) -> "Arg":
+        """
+        Value of the argument to be passed to device
+        """
+        arg = RegularQuadrature.Arg()
+        return arg
+
+    def total_point_count(self):
+        """Total number of quadrature points over the domain"""
+        raise NotImplementedError()
+
+    def points_per_element(self):
+        """Number of points per element if constant, or ``None`` if varying"""
+        return None
+
+    @staticmethod
+    def point_count(elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex):
+        """Number of quadrature points for a given element"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def point_coords(
+        elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex, qp_index: int
+    ):
+        """Coordinates in element of the element's qp_index'th quadrature point"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def point_weight(
+        elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex, qp_index: int
+    ):
+        """Weight of the element's qp_index'th quadrature point"""
+        raise NotImplementedError()
+
+    @staticmethod
+    def point_index(
+        elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex, qp_index: int
+    ):
+        """Global index of the element's qp_index'th  quadrature point"""
+        raise NotImplementedError()
+
+    def __str__(self) -> str:
+        return self.name
+
+
+class RegularQuadrature(Quadrature):
+    """Regular quadrature formula, using a constant set of quadrature points per element"""
+
+    def __init__(
+        self,
+        domain: domain.GeometryDomain,
+        order: int,
+        family: Polynomial = None,
+    ):
+        super().__init__(domain)
+
+        self.family = family
+        self.order = order
+
+        self._element_quadrature = domain.reference_element().instantiate_quadrature(order, family)
+
+        self._N = wp.constant(len(self.points))
+
+        WeightVec = wp.vec(length=self._N, dtype=wp.float32)
+        CoordMat = wp.mat(shape=(self._N, 3), dtype=wp.float32)
+
+        self._POINTS = wp.constant(CoordMat(self.points))
+        self._WEIGHTS = wp.constant(WeightVec(self.weights))
+
+        self.point_count = self._make_point_count()
+        self.point_index = self._make_point_index()
+        self.point_coords = self._make_point_coords()
+        self.point_weight = self._make_point_weight()
+
+    @property
+    def name(self):
+        return f"{self.__class__.__name__}_{self.domain.name}_{self.family}_{self.order}"
+
+    def total_point_count(self):
+        return len(self.points) * self.domain.geometry_element_count()
+
+    def points_per_element(self):
+        return self._N
+
+    @property
+    def points(self):
+        return self._element_quadrature[0]
+
+    @property
+    def weights(self):
+        return self._element_quadrature[1]
+
+    def _make_point_count(self):
+        N = self._N
+
+        @cache.dynamic_func(suffix=self.name)
+        def point_count(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex):
+            return N
+
+        return point_count
+
+    def _make_point_coords(self):
+        POINTS = self._POINTS
+
+        @cache.dynamic_func(suffix=self.name)
+        def point_coords(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return Coords(POINTS[qp_index, 0], POINTS[qp_index, 1], POINTS[qp_index, 2])
+
+        return point_coords
+
+    def _make_point_weight(self):
+        WEIGHTS = self._WEIGHTS
+
+        @cache.dynamic_func(suffix=self.name)
+        def point_weight(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return WEIGHTS[qp_index]
+
+        return point_weight
+
+    def _make_point_index(self):
+        N = self._N
+
+        @cache.dynamic_func(suffix=self.name)
+        def point_index(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return N * element_index + qp_index
+
+        return point_index
+
+
+class NodalQuadrature(Quadrature):
+    """Quadrature using space node points as quadrature points
+
+    Note that in contrast to the `nodal=True` flag for :func:`integrate`, this quadrature odes not make any assumption
+    about orthogonality of shape functions, and is thus safe to use for arbitrary integrands.
+    """
+
+    def __init__(self, domain: domain.GeometryDomain, space: FunctionSpace):
+        super().__init__(domain)
+
+        self._space = space
+
+        self.Arg = self._make_arg()
+
+        self.point_count = self._make_point_count()
+        self.point_index = self._make_point_index()
+        self.point_coords = self._make_point_coords()
+        self.point_weight = self._make_point_weight()
+
+    @property
+    def name(self):
+        return f"{self.__class__.__name__}_{self._space.name}"
+
+    def total_point_count(self):
+        return self._space.node_count()
+
+    def points_per_element(self):
+        return self._space.topology.NODES_PER_ELEMENT
+
+    def _make_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class Arg:
+            space_arg: self._space.SpaceArg
+            topo_arg: self._space.topology.TopologyArg
+
+        return Arg
+
+    @cache.cached_arg_value
+    def arg_value(self, device):
+        arg = self.Arg()
+        arg.space_arg = self._space.space_arg_value(device)
+        arg.topo_arg = self._space.topology.topo_arg_value(device)
+        return arg
+
+    def _make_point_count(self):
+        N = self._space.topology.NODES_PER_ELEMENT
+
+        @cache.dynamic_func(suffix=self.name)
+        def point_count(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex):
+            return N
+
+        return point_count
+
+    def _make_point_coords(self):
+        @cache.dynamic_func(suffix=self.name)
+        def point_coords(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return self._space.node_coords_in_element(elt_arg, qp_arg.space_arg, element_index, qp_index)
+
+        return point_coords
+
+    def _make_point_weight(self):
+        @cache.dynamic_func(suffix=self.name)
+        def point_weight(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return self._space.node_quadrature_weight(elt_arg, qp_arg.space_arg, element_index, qp_index)
+
+        return point_weight
+
+    def _make_point_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def point_index(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+            return self._space.topology.element_node_index(elt_arg, qp_arg.topo_arg, element_index, qp_index)
+
+        return point_index
+
+
+class ExplicitQuadrature(Quadrature):
+    """Quadrature using explicit per-cell points and weights. The number of quadrature points per cell is assumed
+    to be constant and deduced from the shape of the points and weights arrays.
+
+    Args:
+        domain: Domain of definition of the quadrature formula
+        points: 2d array of shape ``(domain.geometry_element-count(), points_per_cell)`` containing the coordinates of each quadrature point.
+        weights: 2d array of shape ``(domain.geometry_element-count(), points_per_cell)`` containing the weight for each quadrature point.
+
+    See also: :class:`PicQuadrature`
+    """
+
+    @wp.struct
+    class Arg:
+        points_per_cell: int
+        points: wp.array2d(dtype=Coords)
+        weights: wp.array2d(dtype=float)
+
+    def __init__(self, domain: domain.GeometryDomain, points: "wp.array2d(dtype=Coords)", weights: "wp.array2d(dtype=float)"):
+        super().__init__(domain)
+
+        if points.shape != weights.shape:
+            raise ValueError("Points and weights arrays must have the same shape")
+
+        self._points_per_cell = points.shape[1]
+        self._points = points
+        self._weights = weights
+
+    @property
+    def name(self):
+        return f"{self.__class__.__name__}"
+
+    def total_point_count(self):
+        return self._weights.size
+
+    def points_per_element(self):
+        return self._points_per_cell
+
+    @cache.cached_arg_value
+    def arg_value(self, device):
+        arg = self.Arg()
+        arg.points_per_cell = self._points_per_cell
+        arg.points = self._points.to(device)
+        arg.weights = self._weights.to(device)
+
+        return arg
+
+    @wp.func
+    def point_count(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex):
+        return qp_arg.points_per_cell
+
+    @wp.func
+    def point_coords(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, qp_index: int):
+        return qp_arg.points[element_index, qp_index]
+
+    @wp.func
+    def point_weight(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, qp_index: int):
+        return qp_arg.weights[element_index, qp_index]
+
+    @wp.func
+    def point_index(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, qp_index: int):
+        return qp_arg.points_per_cell * element_index + qp_index

warp/fem/space/__init__.py

@@ -0,0 +1,292 @@
+from typing import Optional
+from enum import Enum
+
+import warp.fem.domain as _domain
+import warp.fem.geometry as _geometry
+import warp.fem.polynomial as _polynomial
+
+from .function_space import FunctionSpace
+from .topology import SpaceTopology
+from .basis_space import BasisSpace, PointBasisSpace
+from .collocated_function_space import CollocatedFunctionSpace
+
+from .grid_2d_function_space import (
+    GridPiecewiseConstantBasis,
+    GridBipolynomialBasisSpace,
+    GridDGBipolynomialBasisSpace,
+    GridSerendipityBasisSpace,
+    GridDGSerendipityBasisSpace,
+    GridDGPolynomialBasisSpace,
+)
+from .grid_3d_function_space import (
+    GridTripolynomialBasisSpace,
+    GridDGTripolynomialBasisSpace,
+    Grid3DPiecewiseConstantBasis,
+    Grid3DSerendipityBasisSpace,
+    Grid3DDGSerendipityBasisSpace,
+    Grid3DDGPolynomialBasisSpace,
+)
+from .trimesh_2d_function_space import (
+    Trimesh2DPiecewiseConstantBasis,
+    Trimesh2DPolynomialBasisSpace,
+    Trimesh2DDGPolynomialBasisSpace,
+    Trimesh2DNonConformingPolynomialBasisSpace,
+)
+from .tetmesh_function_space import (
+    TetmeshPiecewiseConstantBasis,
+    TetmeshPolynomialBasisSpace,
+    TetmeshDGPolynomialBasisSpace,
+    TetmeshNonConformingPolynomialBasisSpace,
+)
+from .quadmesh_2d_function_space import (
+    Quadmesh2DPiecewiseConstantBasis,
+    Quadmesh2DBipolynomialBasisSpace,
+    Quadmesh2DDGBipolynomialBasisSpace,
+    Quadmesh2DSerendipityBasisSpace,
+    Quadmesh2DDGSerendipityBasisSpace,
+    Quadmesh2DPolynomialBasisSpace,
+)
+from .hexmesh_function_space import (
+    HexmeshPiecewiseConstantBasis,
+    HexmeshTripolynomialBasisSpace,
+    HexmeshDGTripolynomialBasisSpace,
+    HexmeshSerendipityBasisSpace,
+    HexmeshDGSerendipityBasisSpace,
+    HexmeshPolynomialBasisSpace,
+)
+
+from .partition import SpacePartition, make_space_partition
+from .restriction import SpaceRestriction
+
+
+from .dof_mapper import DofMapper, IdentityMapper, SymmetricTensorMapper, SkewSymmetricTensorMapper
+
+
+def make_space_restriction(
+    space: Optional[FunctionSpace] = None,
+    space_partition: Optional[SpacePartition] = None,
+    domain: Optional[_domain.GeometryDomain] = None,
+    space_topology: Optional[SpaceTopology] = None,
+    device=None,
+    temporary_store: "Optional[warp.fem.cache.TemporaryStore]" = None,
+) -> SpaceRestriction:
+    """
+    Restricts a function space partition to a Domain, i.e. a subset of its elements.
+
+    One of `space_partition`, `space_topology`, or `space` must be provided (and will be considered in that order).
+
+    Args:
+        space: (deprecated) if neither `space_partition` nor `space_topology` are provided, the space defining the topology to restrict
+        space_partition: the subset of nodes from the space topology to consider
+        domain: the domain to restrict the space to, defaults to all cells of the space geometry or partition.
+        space_topology: the space topology to be restricted, if `space_partition` is ``None``.
+        device: device on which to perform and store computations
+        temporary_store: shared pool from which to allocate temporary arrays
+    """
+
+    if space_partition is None:
+        if space_topology is None:
+            assert space is not None
+            space_topology = space.topology
+
+        if domain is None:
+            domain = _domain.Cells(geometry=space_topology.geometry)
+
+        space_partition = make_space_partition(
+            space_topology=space_topology, geometry_partition=domain.geometry_partition
+        )
+    elif domain is None:
+        domain = _domain.Cells(geometry=space_partition.geo_partition)
+
+    return SpaceRestriction(
+        space_partition=space_partition, domain=domain, device=device, temporary_store=temporary_store
+    )
+
+
+class ElementBasis(Enum):
+    """Choice of basis function to equip individual elements"""
+
+    LAGRANGE = 0
+    """Lagrange basis functions :math:`P_k` for simplices, tensor products :math:`Q_k` for squares and cubes"""
+    SERENDIPITY = 1
+    """Serendipity elements :math:`S_k`, corresponding to Lagrange nodes with interior points removed (for degree <= 3)"""
+    NONCONFORMING_POLYNOMIAL = 2
+    """Simplex Lagrange basis functions :math:`P_{kd}` embedded into non conforming reference elements (e.g. squares or cubes). Discontinuous only."""
+
+
+def make_polynomial_basis_space(
+    geo: _geometry.Geometry,
+    degree: int = 1,
+    element_basis: Optional[ElementBasis] = None,
+    discontinuous: bool = False,
+    family: Optional[_polynomial.Polynomial] = None,
+) -> BasisSpace:
+    """
+    Equips a geometry with a polynomial basis.
+
+    Args:
+        geo: the Geometry on which to build the space
+        degree: polynomial degree of the per-element shape functions
+        discontinuous: if True, use Discontinuous Galerkin shape functions. Discontinuous is implied if degree is 0, i.e, piecewise-constant shape functions.
+        element_basis: type of basis function for the individual elements
+        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.
+
+    Returns:
+        the constructed basis space
+    """
+
+    base_geo = geo.base if isinstance(geo, _geometry.DeformedGeometry) else geo
+
+    if element_basis is None:
+        element_basis = ElementBasis.LAGRANGE
+
+    if isinstance(base_geo, _geometry.Grid2D):
+        if degree == 0:
+            return GridPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            if discontinuous:
+                return GridDGSerendipityBasisSpace(geo, degree=degree, family=family)
+            else:
+                return GridSerendipityBasisSpace(geo, degree=degree, family=family)
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return GridDGPolynomialBasisSpace(geo, degree=degree)
+
+        if discontinuous:
+            return GridDGBipolynomialBasisSpace(geo, degree=degree, family=family)
+        else:
+            return GridBipolynomialBasisSpace(geo, degree=degree, family=family)
+
+    if isinstance(base_geo, _geometry.Grid3D):
+        if degree == 0:
+            return Grid3DPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            if discontinuous:
+                return Grid3DDGSerendipityBasisSpace(geo, degree=degree, family=family)
+            else:
+                return Grid3DSerendipityBasisSpace(geo, degree=degree, family=family)
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return Grid3DDGPolynomialBasisSpace(geo, degree=degree)
+
+        if discontinuous:
+            return GridDGTripolynomialBasisSpace(geo, degree=degree, family=family)
+        else:
+            return GridTripolynomialBasisSpace(geo, degree=degree, family=family)
+
+    if isinstance(base_geo, _geometry.Trimesh2D):
+        if degree == 0:
+            return Trimesh2DPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet")
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return Trimesh2DNonConformingPolynomialBasisSpace(geo, degree=degree)
+
+        if discontinuous:
+            return Trimesh2DDGPolynomialBasisSpace(geo, degree=degree)
+        else:
+            return Trimesh2DPolynomialBasisSpace(geo, degree=degree)
+
+    if isinstance(base_geo, _geometry.Tetmesh):
+        if degree == 0:
+            return TetmeshPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 2:
+            raise NotImplementedError("Serendipity variant not implemented yet")
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return TetmeshNonConformingPolynomialBasisSpace(geo, degree=degree)
+
+        if discontinuous:
+            return TetmeshDGPolynomialBasisSpace(geo, degree=degree)
+        else:
+            return TetmeshPolynomialBasisSpace(geo, degree=degree)
+
+    if isinstance(base_geo, _geometry.Quadmesh2D):
+        if degree == 0:
+            return Quadmesh2DPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            if discontinuous:
+                return Quadmesh2DDGSerendipityBasisSpace(geo, degree=degree, family=family)
+            else:
+                return Quadmesh2DSerendipityBasisSpace(geo, degree=degree, family=family)
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return Quadmesh2DPolynomialBasisSpace(geo, degree=degree)
+
+        if discontinuous:
+            return Quadmesh2DDGBipolynomialBasisSpace(geo, degree=degree, family=family)
+        else:
+            return Quadmesh2DBipolynomialBasisSpace(geo, degree=degree, family=family)
+
+    if isinstance(base_geo, _geometry.Hexmesh):
+        if degree == 0:
+            return HexmeshPiecewiseConstantBasis(geo)
+
+        if element_basis == ElementBasis.SERENDIPITY and degree > 1:
+            if discontinuous:
+                return HexmeshDGSerendipityBasisSpace(geo, degree=degree, family=family)
+            else:
+                return HexmeshSerendipityBasisSpace(geo, degree=degree, family=family)
+
+        if element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+            return HexmeshPolynomialBasisSpace(geo, degree=degree)
+
+        if discontinuous:
+            return HexmeshDGTripolynomialBasisSpace(geo, degree=degree, family=family)
+        else:
+            return HexmeshTripolynomialBasisSpace(geo, degree=degree, family=family)
+
+    raise NotImplementedError()
+
+
+def make_collocated_function_space(
+    basis_space: BasisSpace, dtype: type = float, dof_mapper: Optional[DofMapper] = None
+) -> CollocatedFunctionSpace:
+    """
+    Constructs a function space from a basis space and a value type, such that all degrees of freedom of the value type are stored at each of the basis nodes.
+
+    Args:
+        geo: the Geometry on which to build the space
+        dtype: value type the function space. If ``dof_mapper`` is provided, the value type from the DofMapper will be used instead.
+        dof_mapper: mapping from node degrees of freedom to function values, defaults to Identity. Useful for reduced coordinates, e.g. :py:class:`SymmetricTensorMapper` maps 2x2 (resp 3x3) symmetric tensors to 3 (resp 6) degrees of freedom.
+
+    Returns:
+        the constructed function space
+    """
+    return CollocatedFunctionSpace(basis_space, dtype=dtype, dof_mapper=dof_mapper)
+
+
+def make_polynomial_space(
+    geo: _geometry.Geometry,
+    dtype: type = float,
+    dof_mapper: Optional[DofMapper] = None,
+    degree: int = 1,
+    element_basis: Optional[ElementBasis] = None,
+    discontinuous: bool = False,
+    family: Optional[_polynomial.Polynomial] = None,
+) -> CollocatedFunctionSpace:
+    """
+    Equips a geometry with a collocated, polynomial function space.
+    Equivalent to successive calls to :func:`make_polynomial_basis_space` and `make_collocated_function_space`.
+
+    Args:
+        geo: the Geometry on which to build the space
+        dtype: value type the function space. If ``dof_mapper`` is provided, the value type from the DofMapper will be used instead.
+        dof_mapper: mapping from node degrees of freedom to function values, defaults to Identity. Useful for reduced coordinates, e.g. :py:class:`SymmetricTensorMapper` maps 2x2 (resp 3x3) symmetric tensors to 3 (resp 6) degrees of freedom.
+        degree: polynomial degree of the per-element shape functions
+        discontinuous: if True, use Discontinuous Galerkin shape functions. Discontinuous is implied if degree is 0, i.e, piecewise-constant shape functions.
+        element_basis: type of basis function for the individual elements
+        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.
+
+    Returns:
+        the constructed function space
+    """
+
+    basis_space = make_polynomial_basis_space(geo, degree, element_basis, discontinuous, family)
+    return CollocatedFunctionSpace(basis_space, dtype=dtype, dof_mapper=dof_mapper)