warp-lang 1.9.0__py3-none-manylinux_2_34_aarch64.whl → 1.10.0__py3-none-manylinux_2_34_aarch64.whl

warp/{fem → _src/fem}/field/nodal_field.py

@@ -16,13 +16,15 @@
 from typing import Any, ClassVar
 
 import warp as wp
-from warp.fem import cache
-from warp.fem.space import CollocatedFunctionSpace, SpacePartition
-from warp.fem.types import NULL_NODE_INDEX, ElementIndex, Sample
-from warp.fem.utils import type_zero_element
+from warp._src.fem import cache
+from warp._src.fem.space import FunctionSpace, SpacePartition
+from warp._src.fem.types import NULL_NODE_INDEX, ElementIndex, Sample
+from warp._src.fem.utils import type_zero_element
 
 from .field import DiscreteField
 
+_wp_module_name_ = "warp.fem.field.nodal_field"
+
 
 class NodalFieldBase(DiscreteField):
     """Base class for nodal field and nodal field traces. Does not hold values"""
@@ -47,7 +49,7 @@ class NodalFieldBase(DiscreteField):
         "at_node": lambda obj: obj._make_at_node(),
     }
 
-    def __init__(self, space: CollocatedFunctionSpace, space_partition: SpacePartition):
+    def __init__(self, space: FunctionSpace, space_partition: SpacePartition):
         super().__init__(space, space_partition)
         cache.setup_dynamic_attributes(self)
 
@@ -325,10 +327,10 @@ class NodalFieldBase(DiscreteField):
 class NodalField(NodalFieldBase):
     """A field holding values for all degrees of freedom at each node of the underlying function space partition
 
-    See also: warp.fem.space.CollocatedFunctionSpace.make_field
+    See also: warp.fem.space.FunctionSpace.make_field
     """
 
-    def __init__(self, space: CollocatedFunctionSpace, space_partition: SpacePartition):
+    def __init__(self, space: FunctionSpace, space_partition: SpacePartition):
         if space.topology != space_partition.space_topology:
             raise ValueError("Incompatible space and space partition topologies")
 
@@ -336,10 +338,9 @@ class NodalField(NodalFieldBase):
 
         self._dof_values = wp.zeros(n=self.space_partition.node_count(), dtype=self.dof_dtype)
 
+    @cache.cached_arg_value
     def eval_arg_value(self, device):
-        arg = self.EvalArg()
-        self.fill_eval_arg(arg, device)
-        return arg
+        return super().eval_arg_value(device)
 
     def fill_eval_arg(self, arg, device):
         arg.dof_values = self._dof_values.to(device)
@@ -347,6 +348,25 @@ class NodalField(NodalFieldBase):
         self.space_partition.fill_partition_arg(arg.partition_arg, device)
         self.space.topology.fill_topo_arg(arg.topology_arg, device)
 
+    def rebind(self, space: FunctionSpace, space_partition: SpacePartition):
+        """Rebind the field to a new space partition and space.
+        The new space partition and space must be of similar types as the current ones
+        """
+
+        if space_partition.name != self.space_partition.name or space.name != self.space.name:
+            raise ValueError("Incompatible space and/or space partition")
+
+        self._space = space
+        self._space_partition = space_partition
+
+        node_count = space_partition.node_count()
+        if node_count < self._dof_values.shape[0]:
+            self._dof_values = self._dof_values[:node_count]
+        elif node_count > self._dof_values.shape[0]:
+            self._dof_values = wp.empty(n=node_count, dtype=self.dof_dtype)
+
+        self.eval_arg_value.invalidate(self)
+
     @property
     def dof_values(self) -> wp.array:
         """Returns a warp array containing the values at all degrees of freedom of the underlying space partition"""
@@ -365,16 +385,13 @@ class NodalField(NodalFieldBase):
         else:
             self._dof_values = wp.array(values, dtype=self.dof_dtype)
 
+        self.eval_arg_value.invalidate(self)
+
     class Trace(NodalFieldBase):
         def __init__(self, field):
             self._field = field
             super().__init__(field.space.trace(), field.space_partition)
 
-        def eval_arg_value(self, device):
-            arg = self.EvalArg()
-            self.fill_eval_arg(arg, device)
-            return arg
-
         def fill_eval_arg(self, arg, device):
             arg.dof_values = self._field.dof_values.to(device)
             self.space.fill_space_arg(arg.space_arg, device)

warp/{fem → _src/fem}/field/restriction.py

@@ -13,10 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from warp.fem.space import SpaceRestriction
+from warp._src.fem.space import SpaceRestriction
 
 from .field import DiscreteField
 
+_wp_module_name_ = "warp.fem.field.restriction"
+
 
 class FieldRestriction:
     """Restriction of a discrete field to a given GeometryDomain"""

warp/{fem → _src/fem}/field/virtual.py

@@ -13,16 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import cached_property
 from typing import Any, ClassVar, Dict, Optional, Set
 
 import warp as wp
-import warp.fem.operator as operator
-from warp.fem import cache
-from warp.fem.domain import GeometryDomain
-from warp.fem.linalg import basis_coefficient, generalized_inner, generalized_outer
-from warp.fem.quadrature import Quadrature
-from warp.fem.space import FunctionSpace, SpacePartition, SpaceRestriction
-from warp.fem.types import (
+import warp._src.fem.operator as operator
+from warp._src.fem import cache
+from warp._src.fem.domain import GeometryDomain
+from warp._src.fem.linalg import basis_coefficient, generalized_inner, generalized_outer
+from warp._src.fem.quadrature import Quadrature
+from warp._src.fem.space import FunctionSpace, SpacePartition, SpaceRestriction
+from warp._src.fem.types import (
     NULL_ELEMENT_INDEX,
     NULL_NODE_INDEX,
     DofIndex,
@@ -32,10 +33,12 @@ from warp.fem.types import (
     get_node_coord,
     get_node_index_in_element,
 )
-from warp.fem.utils import type_zero_element
+from warp._src.fem.utils import type_zero_element
 
 from .field import SpaceField
 
+_wp_module_name_ = "warp.fem.field.virtual"
+
 
 class AdjointField(SpaceField):
     """Adjoint of a discrete field with respect to its degrees of freedom"""
@@ -55,28 +58,44 @@ class AdjointField(SpaceField):
         "node_index": lambda obj: obj._make_node_index(),
     }
 
-    def __init__(self, space: FunctionSpace, space_partition: SpacePartition):
+    def __init__(self, space: FunctionSpace, space_partition: SpacePartition, domain: GeometryDomain):
         super().__init__(space, space_partition=space_partition)
 
         self.node_dof_count = self.space.NODE_DOF_COUNT
         self.value_dof_count = self.space.VALUE_DOF_COUNT
+        self.domain = domain
 
         cache.setup_dynamic_attributes(self)
 
-    @property
+    @cached_property
     def name(self) -> str:
         return f"{self.__class__.__name__}{self.space.name}{self._space_partition.name}"
 
     @cache.cached_arg_value
     def eval_arg_value(self, device):
-        arg = self.EvalArg()
-        self.fill_eval_arg(arg, device)
-        return arg
+        return super().eval_arg_value(device)
 
     def fill_eval_arg(self, arg, device):
         self.space.fill_space_arg(arg.space_arg, device)
         self.space.topology.fill_topo_arg(arg.topo_arg, device)
 
+    def rebind(self, space: FunctionSpace, space_partition: SpacePartition, domain: GeometryDomain):
+        """Rebind the field to a new space partition, space and domain.
+        The new space topology and space must be of similar types as the current ones
+        """
+
+        if (
+            space_partition.space_topology.name != self.space_partition.space_topology.name
+            or space.name != self.space.name
+        ):
+            raise ValueError("Incompatible space and/or space partition")
+
+        self._space = space
+        self._space_partition = space_partition
+        self.domain = domain
+
+        self.eval_arg_value.invalidate(self)
+
     def _make_eval_arg(self):
         @cache.dynamic_struct(suffix=self.name)
         class EvalArg:
@@ -254,7 +273,7 @@ class TestField(AdjointField):
     defined for a different value type than the test function value type, as long as the node topology is similar.
     """
 
-    def __init__(self, space_restriction: SpaceRestriction, space: FunctionSpace):
+    def __init__(self, space: FunctionSpace, space_restriction: SpaceRestriction):
         if space_restriction.domain.dimension == space.dimension - 1:
             space = space.trace()
 
@@ -264,10 +283,17 @@ class TestField(AdjointField):
         if space.topology != space_restriction.space_topology:
             raise ValueError("Incompatible space and space partition topologies")
 
-        super().__init__(space, space_restriction.space_partition)
+        super().__init__(space, space_restriction.space_partition, space_restriction.domain)
 
         self.space_restriction = space_restriction
-        self.domain = space_restriction.domain
+
+    def rebind(self, space: FunctionSpace, space_restriction: SpaceRestriction):
+        """Rebind the test field to a new space restriction and space.
+        The new space restriction and space must be of a similar type as the current ones
+        """
+
+        super().rebind(space, space_restriction.space_partition, space_restriction.domain)
+        self.space_restriction = space_restriction
 
     @wp.func
     def _get_dof(s: Sample):
@@ -292,8 +318,7 @@ class TrialField(AdjointField):
         if not space.topology.is_derived_from(space_partition.space_topology):
             raise ValueError("Incompatible space and space partition topologies")
 
-        super().__init__(space, space_partition)
-        self.domain = domain
+        super().__init__(space, space_partition, domain)
 
     def partition_node_count(self) -> int:
         """Returns the number of nodes in the associated space topology partition"""
@@ -410,14 +435,11 @@ class LocalAdjointField(SpaceField):
     def name(self) -> str:
         return f"{self.global_field.name}_Taylor{self._dof_suffix}"
 
-    def eval_arg_value(self, device):
-        return LocalAdjointField.EvalArg()
-
     def fill_eval_arg(self, arg, device):
         pass
 
     def _make_element_eval_arg(self):
-        from warp.fem import cache
+        from warp._src.fem import cache
 
         @cache.dynamic_struct(suffix=self.name)
         class ElementEvalArg:
@@ -631,7 +653,7 @@ def make_linear_dispatch_kernel(
         return qp, elem_offset, qp_point_count, element_index, test_element_index
 
     @cache.dynamic_kernel(
-        f"{test.name}_{quadrature.name}_{wp.types.get_type_code(accumulate_dtype)}_{tile_size}",
+        (test.name, quadrature.name, cache.pod_type_key(accumulate_dtype), tile_size),
         kernel_options=kernel_options,
     )
     def dispatch_linear_kernel_fn(
@@ -646,6 +668,9 @@ def make_linear_dispatch_kernel(
         local_node_index, lane = wp.tid()
 
         node_index = space_restriction.node_partition_index(test_arg, local_node_index)
+        if node_index == NULL_NODE_INDEX:
+            return
+
         element_beg, element_end = space_restriction.node_element_range(test_arg, node_index)
 
         val_sum = res_vec()
@@ -816,7 +841,7 @@ def make_bilinear_dispatch_kernel(
     val_t = cache.cached_mat_type(shape=(test.node_dof_count, trial.node_dof_count), dtype=accumulate_dtype)
 
     @cache.dynamic_kernel(
-        f"{trial.name}_{test.name}_{quadrature.name}{wp.types.get_type_code(accumulate_dtype)}_{tile_size}",
+        (trial.name, test.name, quadrature.name, cache.pod_type_key(accumulate_dtype), tile_size),
         kernel_options=kernel_options,
     )
     def dispatch_bilinear_kernel_fn(
@@ -841,9 +866,12 @@ def make_bilinear_dispatch_kernel(
         element_index = domain.element_index(domain_index_arg, test_element_index.domain_element_index)
         test_node = test_element_index.node_index_in_element
 
-        element_trial_node_count = trial.space.topology.element_node_count(
-            domain_arg, trial_topology_arg, element_index
-        )
+        if element_index == NULL_ELEMENT_INDEX:
+            element_trial_node_count = 0
+        else:
+            element_trial_node_count = trial.space.topology.element_node_count(
+                domain_arg, trial_topology_arg, element_index
+            )
 
         if trial_node >= element_trial_node_count:
             block_offset = test_node_offset * MAX_NODES_PER_ELEMENT + trial_node

warp/_src/fem/geometry/__init__.py

@@ -0,0 +1,32 @@
+# 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 .adaptive_nanogrid import AdaptiveNanogrid
+from .deformed_geometry import DeformedGeometry
+from .element import Element
+from .geometry import Geometry
+from .grid_2d import Grid2D
+from .grid_3d import Grid3D
+from .hexmesh import Hexmesh
+from .nanogrid import Nanogrid
+from .partition import (
+    ExplicitGeometryPartition,
+    GeometryPartition,
+    LinearGeometryPartition,
+    WholeGeometryPartition,
+)
+from .quadmesh import Quadmesh, Quadmesh2D, Quadmesh3D
+from .tetmesh import Tetmesh
+from .trimesh import Trimesh, Trimesh2D, Trimesh3D