warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/_src/fem/dirichlet.py

@@ -0,0 +1,188 @@
+# 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._src.fem.linalg import array_axpy, symmetric_eigenvalues_qr
+from warp._src.types import type_is_matrix, type_size
+from warp.sparse import BsrMatrix, bsr_assign, bsr_axpy, bsr_block_index, bsr_copy, bsr_mm, bsr_mv
+
+
+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.ncol != projector_matrix.nrow:
+        raise ValueError("Projector must be a square block-diagonal matrix")
+
+    # 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_size(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 right-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()
+
+    diag = bsr_block_index(row, row, offsets, columns)
+
+    if diag != -1:
+        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()
+
+    diag = bsr_block_index(row, row, offsets, columns)
+
+    if diag != -1:
+        P = block_values[diag]
+
+        P_norm, v_norm = _normalize_projector_and_value(P, type(P[0])())
+        block_values[diag] = P_norm

warp/{fem → _src/fem}/domain.py

@@ -17,18 +17,18 @@ from functools import cached_property
 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 (
+import warp._src.fem.cache as cache
+import warp._src.fem.utils as utils
+from warp._src.codegen import Struct, StructInstance
+from warp._src.context import Devicelike
+from warp._src.fem.geometry import (
     Element,
     Geometry,
     GeometryPartition,
     WholeGeometryPartition,
 )
-from warp.fem.operator import Operator
-from warp.fem.types import NULL_ELEMENT_INDEX, ElementKind
+from warp._src.fem.operator import Operator
+from warp._src.fem.types import NULL_ELEMENT_INDEX, ElementKind
 
 GeometryOrPartition = Union[Geometry, GeometryPartition]
 
@@ -44,7 +44,7 @@ class GeometryDomain:
 
         self.geometry = self.geometry_partition.geometry
 
-    @property
+    @cached_property
     def name(self) -> str:
         return f"{self.geometry_partition.name}_{self.__class__.__name__}"
 
@@ -73,18 +73,28 @@ class GeometryDomain:
         return self.geometry.cell_count()
 
     def reference_element(self) -> Element:
-        """Protypical element"""
+        """Type of reference element"""
         raise NotImplementedError
 
-    def element_index_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+    def element_index_arg_value(self, device: Devicelike) -> StructInstance:
         """Value of the argument to be passed to device functions"""
-        raise NotImplementedError
+        args = self.ElementIndexArg()
+        self.fill_element_index_arg(args, device)
+        return args
+
+    def fill_element_index_arg(self, arg: "GeometryDomain.ElementIndexArg", device: Devicelike):
+        arg.assign(self.element_index_arg_value(device))
 
-    def element_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+    def element_arg_value(self, device: Devicelike) -> StructInstance:
         """Value of the argument to be passed to device functions"""
-        raise NotImplementedError
+        args = self.ElementArg()
+        self.fill_element_arg(args, device)
+        return args
+
+    def fill_element_arg(self, arg: "GeometryDomain.ElementArg", device: Devicelike):
+        arg.assign(self.element_arg_value(device))
 
-    ElementIndexArg: warp.codegen.Struct
+    ElementIndexArg: Struct
     """Structure containing arguments to be passed to device functions computing element indices"""
 
     element_index: wp.Function
@@ -93,7 +103,7 @@ class GeometryDomain:
     element_partition_index: wp.Function
     """Device function for retrieving linearized index in the domain's partition from an ElementIndex"""
 
-    ElementArg: warp.codegen.Struct
+    ElementArg: Struct
     """Structure containing arguments to be passed to device functions computing element geometry"""
 
     element_measure: wp.Function
@@ -164,13 +174,13 @@ class Cells(GeometryDomain):
         return self.geometry.cell_count()
 
     @property
-    def ElementIndexArg(self) -> warp.codegen.Struct:
+    def ElementIndexArg(self) -> Struct:
         return self.geometry_partition.CellArg
 
-    def element_index_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+    def element_index_arg_value(self, device: Devicelike) -> StructInstance:
         return self.geometry_partition.cell_arg_value(device)
 
-    def fill_element_index_arg(self, arg: ElementIndexArg, device: warp.context.Devicelike):
+    def fill_element_index_arg(self, arg: ElementIndexArg, device: Devicelike):
         self.geometry_partition.fill_cell_arg(arg, device)
 
     @property
@@ -181,14 +191,14 @@ class Cells(GeometryDomain):
     def element_partition_index(self) -> wp.Function:
         return self.geometry_partition.partition_cell_index
 
-    def element_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+    def element_arg_value(self, device: Devicelike) -> StructInstance:
         return self.geometry.cell_arg_value(device)
 
-    def fill_element_arg(self, arg: "ElementArg", device: warp.context.Devicelike):
+    def fill_element_arg(self, arg: "ElementArg", device: Devicelike):
         self.geometry.fill_cell_arg(arg, device)
 
     @property
-    def ElementArg(self) -> warp.codegen.Struct:
+    def ElementArg(self) -> Struct:
         return self.geometry.CellArg
 
     @property
@@ -223,7 +233,7 @@ class Cells(GeometryDomain):
     def element_lookup(self) -> wp.Function:
         return self.geometry.cell_lookup
 
-    @property
+    @cached_property
     def element_partition_lookup(self) -> wp.Function:
         pos_type = cache.cached_vec_type(self.geometry.dimension, dtype=float)
 
@@ -295,13 +305,13 @@ class Sides(GeometryDomain):
         return self.geometry.side_count()
 
     @property
-    def ElementIndexArg(self) -> warp.codegen.Struct:
+    def ElementIndexArg(self) -> Struct:
         return self.geometry_partition.SideArg
 
-    def element_index_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+    def element_index_arg_value(self, device: Devicelike) -> StructInstance:
         return self.geometry_partition.side_arg_value(device)
 
-    def fill_element_index_arg(self, arg: "ElementIndexArg", device: warp.context.Devicelike):
+    def fill_element_index_arg(self, arg: "ElementIndexArg", device: Devicelike):
         self.geometry_partition.fill_side_arg(arg, device)
 
     @property
@@ -309,13 +319,13 @@ class Sides(GeometryDomain):
         return self.geometry_partition.side_index
 
     @property
-    def ElementArg(self) -> warp.codegen.Struct:
+    def ElementArg(self) -> Struct:
         return self.geometry.SideArg
 
-    def element_arg_value(self, device: warp.context.Devicelike) -> warp.codegen.StructInstance:
+    def element_arg_value(self, device: Devicelike) -> StructInstance:
         return self.geometry.side_arg_value(device)
 
-    def fill_element_arg(self, arg: "ElementArg", device: warp.context.Devicelike):
+    def fill_element_arg(self, arg: "ElementArg", device: Devicelike):
         self.geometry.fill_side_arg(arg, device)
 
     @property
@@ -498,12 +508,12 @@ class Subdomain(GeometryDomain):
         return ElementIndexArg
 
     @cache.cached_arg_value
-    def element_index_arg_value(self, device: warp.context.Devicelike):
+    def element_index_arg_value(self, device: Devicelike):
         arg = self.ElementIndexArg()
         self.fill_element_index_arg(arg, device)
         return arg
 
-    def fill_element_index_arg(self, arg: "GeometryDomain.ElementIndexArg", device: warp.context.Devicelike):
+    def fill_element_index_arg(self, arg: "GeometryDomain.ElementIndexArg", device: Devicelike):
         self._domain.fill_element_index_arg(arg.domain_arg, device)
         arg.element_indices = self._element_indices.to(device)
 

warp/_src/fem/field/__init__.py

@@ -0,0 +1,131 @@
+# 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 Optional, Union
+
+from warp._src.fem.domain import Cells, GeometryDomain
+from warp._src.fem.space import (
+    FunctionSpace,
+    SpacePartition,
+    SpaceRestriction,
+    make_space_partition,
+    make_space_restriction,
+)
+
+from .field import DiscreteField, FieldLike, GeometryField, ImplicitField, NonconformingField, SpaceField, UniformField
+from .nodal_field import NodalField
+from .restriction import FieldRestriction
+from .virtual import LocalTestField, LocalTrialField, TestField, TrialField
+
+
+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_restriction.domain
+        space_partition = space_restriction.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)
+
+
+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)