warp-lang 1.0.0b2__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

warp/fem/field/test.py

@@ -4,110 +4,178 @@ from warp.fem.space import SpaceRestriction, FunctionSpace
 from warp.fem.types import Sample, get_node_index_in_element
 from warp.fem import utils, cache
 
+from .field import SpaceField
 
-class TestField:
-    """Field defined over a space restriction that can be used as a test function"""
 
-    def __init__(self, space_restriction: SpaceRestriction):
+class TestField(SpaceField):
+    """Field defined over a space restriction that can be used as a test function.
+
+    In order to reuse computations, it is possible to define the test field using a SpaceRestriction
+    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):
+        if space_restriction.domain.dimension == space.dimension - 1:
+            space = space.trace()
+
+        if space_restriction.domain.dimension != space.dimension:
+            raise ValueError("Incompatible space and domain dimensions")
+
+        if space.topology != space_restriction.space_topology:
+            raise ValueError("Incompatible space and space partition topologies")
+
+        super().__init__(space, space_restriction.space_partition)
+
         self.space_restriction = space_restriction
-        self.space_partition = self.space_restriction.space_partition
-        self.space = self.space_restriction.space
         self.domain = self.space_restriction.domain
-        self.name = self.space.name + "Test"
 
-        self.eval_degree = TestField._make_eval_degree(self.space)
-        self.eval_inner = TestField._make_eval_inner(self.space)
-        self.eval_grad_inner = TestField._make_eval_grad_inner(self.space)
-        self.eval_outer = TestField._make_eval_outer(self.space)
-        self.eval_grad_outer = TestField._make_eval_grad_outer(self.space)
-        self.at_node = TestField._make_at_node(self.space)
+        self.EvalArg = self.space.SpaceArg
+        self.ElementEvalArg = self._make_element_eval_arg()
+
+        self.eval_degree = self._make_eval_degree()
+        self.eval_inner = self._make_eval_inner()
+        self.eval_grad_inner = self._make_eval_grad_inner()
+        self.eval_div_inner = self._make_eval_div_inner()
+        self.eval_outer = self._make_eval_outer()
+        self.eval_grad_outer = self._make_eval_grad_outer()
+        self.eval_div_outer = self._make_eval_div_outer()
+        self.at_node = self._make_at_node()
 
     @property
-    def EvalArg(self) -> wp.codegen.Struct:
-        return self.space.SpaceArg
+    def name(self) -> str:
+        return self.space.name + "Test"
 
     def eval_arg_value(self, device) -> wp.codegen.StructInstance:
         return self.space.space_arg_value(device)
 
-    @staticmethod
-    def _make_eval_degree(space: FunctionSpace):
-        ORDER = space.ORDER
+    def _make_element_eval_arg(self):
+        from warp.fem import cache
 
-        def degree(args: space.SpaceArg):
-            return ORDER
+        @cache.dynamic_struct(suffix=self.name)
+        class ElementEvalArg:
+            elt_arg: self.domain.ElementArg
+            eval_arg: self.EvalArg
 
-        return cache.get_func(degree, space)
+        return ElementEvalArg
 
-    @staticmethod
-    def _make_eval_inner(space: FunctionSpace):
-        def eval_test(args: space.SpaceArg, s: Sample):
-            weight = space.element_inner_weight(
-                args,
+    def _make_eval_inner(self):
+        @cache.dynamic_func(suffix=self.name)
+        def eval_test_inner(args: self.ElementEvalArg, s: Sample):
+            weight = self.space.element_inner_weight(
+                args.elt_arg,
+                args.eval_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(s.test_dof),
             )
-            return weight * space.unit_dof_value(args, s.test_dof)
+            return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof)
 
-        return cache.get_func(eval_test, space.name)
+        return eval_test_inner
 
-    @staticmethod
-    def _make_eval_grad_inner(space: FunctionSpace):
-        if wp.types.type_is_matrix(space.dtype):
-            # There is no Warp high-order tensor type to represent matrix gradients
+    def _make_eval_grad_inner(self):
+        if not self.gradient_valid():
             return None
 
-        def eval_nabla_test_inner(args: space.SpaceArg, s: Sample):
-            nabla_weight = space.element_inner_weight_gradient(
-                args,
+        @cache.dynamic_func(suffix=self.name)
+        def eval_nabla_test_inner(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_inner_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(s.test_dof),
             )
+            grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
             return utils.generalized_outer(
-                nabla_weight,
-                space.unit_dof_value(args, s.test_dof),
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
+                utils.apply_right(nabla_weight, grad_transform),
+            )
+
+        return eval_nabla_test_inner
+
+    def _make_eval_div_inner(self):
+        if not self.divergence_valid():
+            return None
+
+        @cache.dynamic_func(suffix=self.name)
+        def eval_div_test_inner(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_inner_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(s.test_dof),
+            )
+            grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
+            return utils.generalized_inner(
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
+                utils.apply_right(nabla_weight, grad_transform),
             )
 
-        return cache.get_func(eval_nabla_test_inner, space.name)
+        return eval_div_test_inner
 
-    @staticmethod
-    def _make_eval_outer(space: FunctionSpace):
-        def eval_test_outer(args: space.SpaceArg, s: Sample):
-            weight = space.element_outer_weight(
-                args,
+    def _make_eval_outer(self):
+        @cache.dynamic_func(suffix=self.name)
+        def eval_test_outer(args: self.ElementEvalArg, s: Sample):
+            weight = self.space.element_outer_weight(
+                args.elt_arg,
+                args.eval_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(s.test_dof),
             )
-            return weight * space.unit_dof_value(args, s.test_dof)
+            return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof)
 
-        return cache.get_func(eval_test_outer, space.name)
+        return eval_test_outer
 
-    @staticmethod
-    def _make_eval_grad_outer(space: FunctionSpace):
-        if wp.types.type_is_matrix(space.dtype):
-            # There is no Warp high-order tensor type to represent matrix gradients
+    def _make_eval_grad_outer(self):
+        if not self.gradient_valid():
             return None
 
-        def eval_nabla_test(args: space.SpaceArg, s: Sample):
-            nabla_weight = space.element_outer_weight_gradient(
-                args,
+        @cache.dynamic_func(suffix=self.name)
+        def eval_nabla_test_outer(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_outer_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(s.test_dof),
             )
+            grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
             return utils.generalized_outer(
-                nabla_weight,
-                space.unit_dof_value(args, s.test_dof),
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
+                utils.apply_right(nabla_weight, grad_transform),
             )
 
-        return cache.get_func(eval_nabla_test, space.name)
+        return eval_nabla_test_outer
 
-    @staticmethod
-    def _make_at_node(space: FunctionSpace):
-        def at_node(args: space.SpaceArg, s: Sample):
-            node_coords = space.node_coords_in_element(args, s.element_index, get_node_index_in_element(s.test_dof))
+    def _make_eval_div_outer(self):
+        if not self.divergence_valid():
+            return None
+
+        @cache.dynamic_func(suffix=self.name)
+        def eval_div_test_outer(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_outer_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(s.test_dof),
+            )
+            grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
+            return utils.generalized_inner(
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.test_dof),
+                utils.apply_right(nabla_weight, grad_transform),
+            )
+
+        return eval_div_test_outer
+
+    def _make_at_node(self):
+        @cache.dynamic_func(suffix=self.name)
+        def at_node(args: self.ElementEvalArg, s: Sample):
+            node_coords = self.space.node_coords_in_element(
+                args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(s.test_dof)
+            )
             return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
 
-        return cache.get_func(at_node, space.name)
+        return at_node

warp/fem/field/trial.py

@@ -4,8 +4,10 @@ from warp.fem.space import FunctionSpace, SpacePartition
 from warp.fem.types import Sample, get_node_index_in_element
 from warp.fem import utils, cache
 
+from .field import SpaceField
 
-class TrialField:
+
+class TrialField(SpaceField):
     """Field defined over a domain that can be used as a trial function"""
 
     def __init__(
@@ -14,113 +16,168 @@ class TrialField:
         space_partition: SpacePartition,
         domain: GeometryDomain,
     ):
-        if domain.dimension() == space.DIMENSION - 1:
+        if domain.dimension == space.dimension - 1:
             space = space.trace()
 
-        if domain.dimension() != space.DIMENSION:
+        if domain.dimension != space.dimension:
             raise ValueError("Incompatible space and domain dimensions")
 
-        self.space = space
+        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
-        self.space_partition = space_partition
-        self.name = self.space.name + "Trial"
 
-        self.eval_degree = TrialField._make_eval_degree(self.space)
-        self.eval_inner = TrialField._make_eval_inner(self.space)
-        self.eval_grad_inner = TrialField._make_eval_grad_inner(self.space)
-        self.eval_outer = TrialField._make_eval_outer(self.space)
-        self.eval_grad_outer = TrialField._make_eval_grad_outer(self.space)
-        self.at_node = TrialField._make_at_node(self.space)
+        self.EvalArg = self.space.SpaceArg
+        self.ElementEvalArg = self._make_element_eval_arg()
+
+        self.eval_degree = self._make_eval_degree()
+        self.eval_inner = self._make_eval_inner()
+        self.eval_grad_inner = self._make_eval_grad_inner()
+        self.eval_div_inner = self._make_eval_div_inner()
+        self.eval_outer = self._make_eval_outer()
+        self.eval_grad_outer = self._make_eval_grad_outer()
+        self.eval_div_outer = self._make_eval_div_outer()
+        self.at_node = self._make_at_node()
 
     def partition_node_count(self) -> int:
+        """Returns the number of nodes in the associated space topology partition"""
         return self.space_partition.node_count()
 
     @property
-    def EvalArg(self) -> wp.codegen.Struct:
-        return self.space.SpaceArg
+    def name(self) -> str:
+        return self.space.name + "Trial"
 
     def eval_arg_value(self, device) -> wp.codegen.StructInstance:
         return self.space.space_arg_value(device)
 
-    @staticmethod
-    def _make_eval_degree(space: FunctionSpace):
-        ORDER = space.ORDER
-
-        def degree(args: space.SpaceArg):
-            return ORDER
+    def _make_element_eval_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class ElementEvalArg:
+            elt_arg: self.domain.ElementArg
+            eval_arg: self.EvalArg
 
-        return cache.get_func(degree, space)
+        return ElementEvalArg
 
-    @staticmethod
-    def _make_eval_inner(space: FunctionSpace):
-        def eval_trial(args: space.SpaceArg, s: Sample):
-            weight = space.element_inner_weight(
-                args,
+    def _make_eval_inner(self):
+        @cache.dynamic_func(suffix=self.name)
+        def eval_trial_inner(args: self.ElementEvalArg, s: Sample):
+            weight = self.space.element_inner_weight(
+                args.elt_arg,
+                args.eval_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(s.trial_dof),
             )
-            return weight * space.unit_dof_value(args, s.trial_dof)
+            return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
 
-        return cache.get_func(eval_trial, space.name)
+        return eval_trial_inner
 
-    @staticmethod
-    def _make_eval_grad_inner(space: FunctionSpace):
-        if wp.types.type_is_matrix(space.dtype):
-            # There is no Warp high-order tensor type to represent matrix gradients
+    def _make_eval_grad_inner(self):
+        if not self.gradient_valid():
             return None
 
-        def eval_nabla_trial(args: space.SpaceArg, s: Sample):
-            nabla_weight = space.element_inner_weight_gradient(
-                args,
+        @cache.dynamic_func(suffix=self.name)
+        def eval_nabla_trial_inner(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_inner_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(s.trial_dof),
             )
+            grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
             return utils.generalized_outer(
-                nabla_weight,
-                space.unit_dof_value(args, s.trial_dof),
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
+                utils.apply_right(nabla_weight, grad_transform),
+            )
+
+        return eval_nabla_trial_inner
+
+    def _make_eval_div_inner(self):
+        if not self.divergence_valid():
+            return None
+
+        @cache.dynamic_func(suffix=self.name)
+        def eval_div_trial_inner(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_inner_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(s.trial_dof),
+            )
+            grad_transform = self.space.element_inner_reference_gradient_transform(args.elt_arg, s)
+            return utils.generalized_inner(
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
+                utils.apply_right(nabla_weight, grad_transform),
             )
 
-        return cache.get_func(eval_nabla_trial, space.name)
+        return eval_div_trial_inner
 
-    @staticmethod
-    def _make_eval_outer(space: FunctionSpace):
-        def eval_trial_outer(args: space.SpaceArg, s: Sample):
-            weight = space.element_outer_weight(
-                args,
+    def _make_eval_outer(self):
+        @cache.dynamic_func(suffix=self.name)
+        def eval_trial_outer(args: self.ElementEvalArg, s: Sample):
+            weight = self.space.element_outer_weight(
+                args.elt_arg,
+                args.eval_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(s.trial_dof),
             )
-            return weight * space.unit_dof_value(args, s.trial_dof)
+            return weight * self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof)
 
-        return cache.get_func(eval_trial_outer, space.name)
+        return eval_trial_outer
 
-    @staticmethod
-    def _make_eval_grad_outer(space: FunctionSpace):
-        if wp.types.type_is_matrix(space.dtype):
-            # There is no Warp high-order tensor type to represent matrix gradients
+    def _make_eval_grad_outer(self):
+        if not self.gradient_valid():
             return None
 
-        def eval_nabla_trial_outer(args: space.SpaceArg, s: Sample):
-            nabla_weight = space.element_outer_weight_gradient(
-                args,
+        @cache.dynamic_func(suffix=self.name)
+        def eval_nabla_trial_outer(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_outer_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
                 s.element_index,
                 s.element_coords,
                 get_node_index_in_element(s.trial_dof),
             )
+            grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
             return utils.generalized_outer(
-                nabla_weight,
-                space.unit_dof_value(args, s.trial_dof),
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
+                utils.apply_right(nabla_weight, grad_transform),
+            )
+
+        return eval_nabla_trial_outer
+
+    def _make_eval_div_outer(self):
+        if not self.divergence_valid():
+            return None
+
+        @cache.dynamic_func(suffix=self.name)
+        def eval_div_trial_outer(args: self.ElementEvalArg, s: Sample):
+            nabla_weight = self.space.element_outer_weight_gradient(
+                args.elt_arg,
+                args.eval_arg,
+                s.element_index,
+                s.element_coords,
+                get_node_index_in_element(s.trial_dof),
+            )
+            grad_transform = self.space.element_outer_reference_gradient_transform(args.elt_arg, s)
+            return utils.generalized_inner(
+                self.space.unit_dof_value(args.elt_arg, args.eval_arg, s.trial_dof),
+                utils.apply_right(nabla_weight, grad_transform),
             )
 
-        return cache.get_func(eval_nabla_trial_outer, space.name)
+        return eval_div_trial_outer
 
-    @staticmethod
-    def _make_at_node(space: FunctionSpace):
-        def at_node(args: space.SpaceArg, s: Sample):
-            node_coords = space.node_coords_in_element(args, s.element_index, get_node_index_in_element(s.trial_dof))
+    def _make_at_node(self):
+        @cache.dynamic_func(suffix=self.name)
+        def at_node(args: self.ElementEvalArg, s: Sample):
+            node_coords = self.space.node_coords_in_element(
+                args.elt_arg, args.eval_arg, s.element_index, get_node_index_in_element(s.trial_dof)
+            )
             return Sample(s.element_index, node_coords, s.qp_index, s.qp_weight, s.test_dof, s.trial_dof)
 
-        return cache.get_func(at_node, space.name)
+        return at_node

warp/fem/geometry/__init__.py

@@ -1,5 +1,7 @@
 from .element import Element
 from .geometry import Geometry
+from .deformed_geometry import DeformedGeometry
+
 from .partition import (
     GeometryPartition,
     WholeGeometryPartition,
@@ -12,4 +14,6 @@ from .grid_2d import Grid2D
 from .grid_3d import Grid3D
 
 from .trimesh_2d import Trimesh2D
-from .tetmesh import Tetmesh
+from .tetmesh import Tetmesh
+from .quadmesh_2d import Quadmesh2D
+from .hexmesh import Hexmesh