warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/fem/dirichlet.py

@@ -0,0 +1,202 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Optional
+
+import warp as wp
+from warp.fem.linalg 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
+
+
+def normalize_dirichlet_projector(projector_matrix: BsrMatrix, fixed_value: Optional[wp.array] = None):
+    """
+    Scale projector so that it becomes idempotent, and apply the same scaling to fixed_value if provided
+    """
+
+    if projector_matrix.nrow < projector_matrix.nnz or projector_matrix.ncol != projector_matrix.nrow:
+        raise ValueError("Projector must be a square diagonal matrix, with at most one non-zero block per row")
+
+    # Cast blocks to matrix type if necessary
+    projector_values = projector_matrix.values
+    if not type_is_matrix(projector_values.dtype):
+        projector_values = wp.array(
+            data=None,
+            ptr=projector_values.ptr,
+            capacity=projector_values.capacity,
+            device=projector_values.device,
+            dtype=wp.mat(shape=projector_matrix.block_shape, dtype=projector_matrix.scalar_type),
+            shape=projector_values.shape[0],
+        )
+
+    if fixed_value is None:
+        wp.launch(
+            kernel=_normalize_dirichlet_projector_kernel,
+            dim=projector_matrix.nrow,
+            device=projector_values.device,
+            inputs=[projector_matrix.offsets, projector_matrix.columns, projector_values],
+        )
+
+    else:
+        if fixed_value.shape[0] != projector_matrix.nrow:
+            raise ValueError("Fixed value array must be of length equal to the number of rows of blocks")
+
+        if type_length(fixed_value.dtype) == 1:
+            # array of scalars, convert to 1d array of vectors
+            fixed_value = wp.array(
+                data=None,
+                ptr=fixed_value.ptr,
+                capacity=fixed_value.capacity,
+                device=fixed_value.device,
+                dtype=wp.vec(length=projector_matrix.block_shape[0], dtype=projector_matrix.scalar_type),
+                shape=fixed_value.shape[0],
+            )
+
+        wp.launch(
+            kernel=_normalize_dirichlet_projector_and_values_kernel,
+            dim=projector_matrix.nrow,
+            device=projector_values.device,
+            inputs=[projector_matrix.offsets, projector_matrix.columns, projector_values, fixed_value],
+        )
+
+
+def project_system_rhs(
+    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
+
+    ``rhs = (I - projector) * ( rhs - system * projector * fixed_value) + projector * fixed_value``
+    """
+
+    rhs_tmp = wp.empty_like(system_rhs)
+    rhs_tmp.assign(system_rhs)
+
+    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
+    # system_rhs = projector * fixed_value
+    array_axpy(x=rhs_tmp, y=system_rhs, alpha=1.0, beta=1.0)
+    bsr_mv(A=projector_matrix, x=rhs_tmp, y=system_rhs, alpha=-1.0, beta=1.0)
+
+
+def project_system_matrix(system_matrix: BsrMatrix, projector_matrix: BsrMatrix):
+    """Projects the right-hand-side of a linear system to enforce Dirichlet boundary conditions
+
+    ``system = (I - projector) * system * (I - projector) + projector``
+    """
+
+    complement_system = bsr_copy(system_matrix)
+    bsr_mm(x=projector_matrix, y=system_matrix, z=complement_system, alpha=-1.0, beta=1.0)
+
+    bsr_assign(dest=system_matrix, src=complement_system)
+    bsr_axpy(x=projector_matrix, y=system_matrix)
+    bsr_mm(x=complement_system, y=projector_matrix, z=system_matrix, alpha=-1.0, beta=1.0)
+
+
+def project_linear_system(
+    system_matrix: BsrMatrix,
+    system_rhs: wp.array,
+    projector_matrix: BsrMatrix,
+    fixed_value: Optional[wp.array] = None,
+    normalize_projector=True,
+):
+    """
+    Projects both the left-hand-side and right-hand-side of a linear system to enforce Dirichlet boundary conditions
+
+    If normalize_projector is True, first apply scaling so that the projector_matrix is idempotent
+    """
+    if normalize_projector:
+        normalize_dirichlet_projector(projector_matrix, fixed_value)
+
+    project_system_rhs(system_matrix, system_rhs, projector_matrix, fixed_value)
+    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_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()
+
+    beg = offsets[row]
+    end = offsets[row + 1]
+
+    if beg == end:
+        return
+
+    diag = wp.lower_bound(columns, beg, end, row)
+
+    if diag < end:
+        if columns[diag] == row:
+            P = block_values[diag]
+
+            P_norm, v_norm = _normalize_projector_and_value(P, fixed_values[row])
+
+            block_values[diag] = P_norm
+            fixed_values[row] = v_norm
+
+
+@wp.kernel
+def _normalize_dirichlet_projector_kernel(
+    offsets: wp.array(dtype=int),
+    columns: wp.array(dtype=int),
+    block_values: wp.array(dtype=Any),
+):
+    row = wp.tid()
+
+    beg = offsets[row]
+    end = offsets[row + 1]
+
+    if beg == end:
+        return
+
+    diag = wp.lower_bound(columns, beg, end, row)
+
+    if diag < end:
+        if columns[diag] == row:
+            P = block_values[diag]
+
+            P_norm, v_norm = _normalize_projector_and_value(P, type(P[0])())
+            block_values[diag] = P_norm

warp/fem/domain.py

@@ -0,0 +1,411 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Optional, Set, 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.operator import Operator
+from warp.fem.types import ElementKind
+
+GeometryOrPartition = Union[Geometry, GeometryPartition]
+
+
+class GeometryDomain:
+    """Interface class for domains, i.e. (partial) views of elements in a Geometry"""
+
+    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"""
+
+    def notify_operator_usage(self, ops: Set[Operator]):
+        """Makes the Domain aware that the operators `ops` will be applied"""
+        pass
+
+
+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) -> ElementKind:
+        return 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 element_normal(self) -> wp.Function:
+        return self.geometry.cell_normal
+
+    @property
+    def element_lookup(self) -> wp.Function:
+        return self.geometry.cell_lookup
+
+    @property
+    def domain_cell_arg(self) -> wp.Function:
+        return Cells._identity_fn
+
+    def cell_domain(self):
+        return self
+
+    @wp.func
+    def _identity_fn(x: Any):
+        return x
+
+
+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)
+
+        self.element_lookup = None
+
+    @property
+    def element_kind(self) -> ElementKind:
+        return 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 element_normal(self) -> wp.Function:
+        return self.geometry.side_normal
+
+    @property
+    def element_inner_cell_index(self) -> wp.Function:
+        return self.geometry.side_inner_cell_index
+
+    @property
+    def element_outer_cell_index(self) -> wp.Function:
+        return self.geometry.side_outer_cell_index
+
+    @property
+    def element_inner_cell_coords(self) -> wp.Function:
+        return self.geometry.side_inner_cell_coords
+
+    @property
+    def element_outer_cell_coords(self) -> wp.Function:
+        return self.geometry.side_outer_cell_coords
+
+    @property
+    def cell_to_element_coords(self) -> wp.Function:
+        return self.geometry.side_from_cell_coords
+
+    @property
+    def domain_cell_arg(self) -> wp.Function:
+        return self.geometry.side_to_cell_arg
+
+    def cell_domain(self):
+        return Cells(self.geometry_partition)
+
+
+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
+
+
+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