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

warp/fem/geometry/deformed_geometry.py

@@ -0,0 +1,271 @@
+from typing import Any
+
+import warp as wp
+
+from .geometry import Geometry
+
+from warp.fem.types import Sample, ElementIndex, Coords, make_free_sample
+from warp.fem import cache
+
+_mat32 = wp.mat(shape=(3, 2), dtype=float)
+
+
+class DeformedGeometry(Geometry):
+    def __init__(self, field):
+        """Constructs a Deformed Geometry from a displacement field defined over a base geometry"""
+
+        from warp.fem.field import DiscreteField
+
+        self.field: DiscreteField = field
+        self.base = self.field.space.geometry
+        self.dimension = self.base.dimension
+
+        if not wp.types.type_is_vector(field.dtype) or wp.types.type_length(field.dtype) != self.dimension:
+            raise ValueError("Invalid value type for position field")
+
+        self.CellArg = self.field.ElementEvalArg
+
+        self.field_trace = field.trace()
+        self.SideArg = self._make_side_arg()
+        self.SideIndexArg = self.base.SideIndexArg
+
+        self.cell_count = self.base.cell_count
+        self.vertex_count = self.base.vertex_count
+        self.side_count = self.base.side_count
+        self.boundary_side_count = self.base.boundary_side_count
+        self.reference_cell = self.base.reference_cell
+        self.reference_side = self.base.reference_side
+
+        self.side_index_arg_value = self.base.side_index_arg_value
+
+        self.cell_position = self._make_cell_position()
+        self.cell_deformation_gradient = self._make_cell_deformation_gradient()
+        self.cell_inverse_deformation_gradient = self._make_cell_inverse_deformation_gradient()
+        self.cell_measure = self._make_cell_measure()
+
+        self.boundary_side_index = self.base.boundary_side_index
+
+        self.side_to_cell_arg = self._make_side_to_cell_arg()
+        self.side_position = self._make_side_position()
+        self.side_deformation_gradient = self._make_side_deformation_gradient()
+        self.side_inner_cell_index = self._make_side_inner_cell_index()
+        self.side_outer_cell_index = self._make_side_outer_cell_index()
+        self.side_inner_cell_coords = self._make_side_inner_cell_coords()
+        self.side_outer_cell_coords = self._make_side_outer_cell_coords()
+        self.side_from_cell_coords = self._make_side_from_cell_coords()
+        self.side_inner_inverse_deformation_gradient = self._make_side_inner_inverse_deformation_gradient()
+        self.side_outer_inverse_deformation_gradient = self._make_side_outer_inverse_deformation_gradient()
+        self.side_measure = self._make_side_measure()
+        self.side_measure_ratio = self._make_side_measure_ratio()
+        self.side_normal = self._make_side_normal()
+
+    @property
+    def name(self):
+        return f"DefGeo_{self.field.name}"
+
+    # Geometry device interface
+
+    @cache.cached_arg_value
+    def cell_arg_value(self, device) -> "DeformedGeometry.CellArg":
+        args = self.CellArg()
+
+        args.elt_arg = self.base.cell_arg_value(device)
+        args.eval_arg = self.field.eval_arg_value(device)
+
+        return args
+
+    def _make_cell_position(self):
+        @cache.dynamic_func(suffix=self.name)
+        def cell_position(cell_arg: self.CellArg, s: Sample):
+            return self.field.eval_inner(cell_arg, s) + self.base.cell_position(cell_arg.elt_arg, s)
+
+        return cell_position
+
+    def _make_cell_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def cell_deformation_gradient(cell_arg: self.CellArg, s: Sample):
+            return self.field.eval_reference_grad_inner(cell_arg, s) + self.base.cell_deformation_gradient(
+                cell_arg.elt_arg, s
+            )
+
+        return cell_deformation_gradient
+
+    def _make_cell_inverse_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def cell_inverse_deformation_gradient(cell_arg: self.CellArg, s: Sample):
+            return wp.inverse(self.cell_deformation_gradient(cell_arg, s))
+
+        return cell_inverse_deformation_gradient
+
+    def _make_cell_measure(self):
+        REF_MEASURE = wp.constant(self.reference_cell().measure())
+
+        @cache.dynamic_func(suffix=self.name)
+        def cell_measure(args: self.CellArg, s: Sample):
+            return wp.abs(wp.determinant(self.cell_deformation_gradient(args, s))) * REF_MEASURE
+
+        return cell_measure
+
+    @wp.func
+    def cell_normal(args: Any, s: Sample):
+        return wp.vec2(0.0)
+
+    def _make_side_arg(self):
+        @cache.dynamic_struct(suffix=self.name)
+        class SideArg:
+            base_arg: self.base.SideArg
+            trace_arg: self.field_trace.EvalArg
+            field_arg: self.field.EvalArg
+
+        return SideArg
+
+    @cache.cached_arg_value
+    def side_arg_value(self, device) -> "DeformedGeometry.SideArg":
+        args = self.SideArg()
+
+        args.base_arg = self.base.side_arg_value(device)
+        args.field_arg = self.field.eval_arg_value(device)
+        args.trace_arg = self.field_trace.eval_arg_value(device)
+
+        return args
+
+    def _make_side_position(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_position(args: self.SideArg, s: Sample):
+            trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg)
+            return self.field_trace.eval_inner(trace_arg, s) + self.base.side_position(args.base_arg, s)
+
+        return side_position
+
+    def _make_side_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_deformation_gradient(args: self.SideArg, s: Sample):
+            base_def_grad = self.base.side_deformation_gradient(args.base_arg, s)
+            trace_arg = self.field_trace.ElementEvalArg(args.base_arg, args.trace_arg)
+
+            Du = self.field_trace.eval_grad_inner(trace_arg, s)
+            return base_def_grad + Du * base_def_grad
+
+        return side_deformation_gradient
+
+    def _make_side_inner_inverse_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_inner_inverse_deformation_gradient(args: self.SideArg, s: Sample):
+            cell_index = self.side_inner_cell_index(args, s.element_index)
+            cell_coords = self.side_inner_cell_coords(args, s.element_index, s.element_coords)
+            cell_arg = self.side_to_cell_arg(args)
+            return self.cell_inverse_deformation_gradient(cell_arg, make_free_sample(cell_index, cell_coords))
+
+    def _make_side_outer_inverse_deformation_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_outer_inverse_deformation_gradient(args: self.SideArg, s: Sample):
+            cell_index = self.side_outer_cell_index(args, s.element_index)
+            cell_coords = self.side_outer_cell_coords(args, s.element_index, s.element_coords)
+            cell_arg = self.side_to_cell_arg(args)
+            return self.cell_inverse_deformation_gradient(cell_arg, make_free_sample(cell_index, cell_coords))
+
+    @wp.func
+    def _side_measure(F: wp.vec2):
+        return wp.length(F)
+
+    @wp.func
+    def _side_measure(F: _mat32):
+        Fcross = wp.vec3(
+            F[1, 0] * F[2, 1] - F[2, 0] * F[1, 1],
+            F[2, 0] * F[0, 1] - F[0, 0] * F[2, 1],
+            F[0, 0] * F[1, 1] - F[1, 0] * F[0, 1],
+        )
+        return wp.length(Fcross)
+
+    @wp.func
+    def _side_normal(F: wp.vec2):
+        return wp.normalize(wp.vec2(-F[1], F[0]))
+
+    @wp.func
+    def _side_normal(F: _mat32):
+        Fcross = wp.vec3(
+            F[1, 0] * F[2, 1] - F[2, 0] * F[1, 1],
+            F[2, 0] * F[0, 1] - F[0, 0] * F[2, 1],
+            F[0, 0] * F[1, 1] - F[1, 0] * F[0, 1],
+        )
+        return wp.normalize(Fcross)
+
+    def _make_side_measure(self):
+        REF_MEASURE = wp.constant(self.reference_side().measure())
+
+        @cache.dynamic_func(suffix=self.name)
+        def side_measure(args: self.SideArg, s: Sample):
+            F = self.side_deformation_gradient(args, s)
+            return DeformedGeometry._side_measure(F) * REF_MEASURE
+
+        return side_measure
+
+    def _make_side_measure_ratio(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_measure_ratio(args: self.SideArg, s: Sample):
+            inner = self.side_inner_cell_index(args, s.element_index)
+            outer = self.side_outer_cell_index(args, s.element_index)
+            inner_coords = self.side_inner_cell_coords(args, s.element_index, s.element_coords)
+            outer_coords = self.side_outer_cell_coords(args, s.element_index, s.element_coords)
+            cell_arg = self.side_to_cell_arg(args)
+            return self.side_measure(args, s) / wp.min(
+                self.cell_measure(cell_arg, make_free_sample(inner, inner_coords)),
+                self.cell_measure(cell_arg, make_free_sample(outer, outer_coords)),
+            )
+
+        return side_measure_ratio
+
+    def _make_side_normal(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_normal(args: self.SideArg, s: Sample):
+            F = self.side_deformation_gradient(args, s)
+            return DeformedGeometry._side_normal(F)
+
+        return side_normal
+
+    def _make_side_inner_cell_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_inner_cell_index(args: self.SideArg, side_index: ElementIndex):
+            return self.base.side_inner_cell_index(args.base_arg, side_index)
+
+        return side_inner_cell_index
+
+    def _make_side_outer_cell_index(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_outer_cell_index(args: self.SideArg, side_index: ElementIndex):
+            return self.base.side_outer_cell_index(args.base_arg, side_index)
+
+        return side_outer_cell_index
+
+    def _make_side_inner_cell_coords(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_inner_cell_coords(args: self.SideArg, side_index: ElementIndex, side_coords: Coords):
+            return self.base.side_inner_cell_coords(args.base_arg, side_index, side_coords)
+
+        return side_inner_cell_coords
+
+    def _make_side_outer_cell_coords(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_outer_cell_coords(args: self.SideArg, side_index: ElementIndex, side_coords: Coords):
+            return self.base.side_outer_cell_coords(args.base_arg, side_index, side_coords)
+
+        return side_outer_cell_coords
+
+    def _make_side_from_cell_coords(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_from_cell_coords(
+            args: self.SideArg,
+            side_index: ElementIndex,
+            cell_index: ElementIndex,
+            cell_coords: Coords,
+        ):
+            return self.base.side_from_cell_coords(args.base_arg, side_index, cell_index, cell_coords)
+
+        return side_from_cell_coords
+
+    def _make_side_to_cell_arg(self):
+        @cache.dynamic_func(suffix=self.name)
+        def side_to_cell_arg(side_arg: self.SideArg):
+            return self.CellArg(self.base.side_to_cell_arg(side_arg.base_arg), side_arg.field_arg)
+
+        return side_to_cell_arg

warp/fem/geometry/element.py

@@ -5,6 +5,10 @@ from warp.fem.polynomial import Polynomial, quadrature_1d
 
 
 class Element:
+    def measure() -> float:
+        """Measure (area, volume, ...) of the reference element"""
+        raise NotImplementedError
+
     @staticmethod
     def instantiate_quadrature(order: int, family: Polynomial) -> Tuple[List[Coords], List[float]]:
         """Returns a quadrature of a given order for a prototypical element"""
@@ -25,6 +29,10 @@ def _point_count_from_order(order: int, family: Polynomial):
 
 
 class Cube(Element):
+    @staticmethod
+    def measure() -> float:
+        return 1.0
+
     @staticmethod
     def instantiate_quadrature(order: int, family: Polynomial):
         if family is None:
@@ -40,6 +48,10 @@ class Cube(Element):
 
 
 class Square(Element):
+    @staticmethod
+    def measure() -> float:
+        return 1.0
+
     @staticmethod
     def instantiate_quadrature(order: int, family: Polynomial):
         if family is None:
@@ -55,6 +67,10 @@ class Square(Element):
 
 
 class LinearEdge(Element):
+    @staticmethod
+    def measure() -> float:
+        return 1.0
+
     @staticmethod
     def instantiate_quadrature(order: int, family: Polynomial):
         if family is None:
@@ -68,6 +84,10 @@ class LinearEdge(Element):
 
 
 class Triangle(Element):
+    @staticmethod
+    def measure() -> float:
+        return 0.5
+
     @staticmethod
     def instantiate_quadrature(order: int, family: Polynomial):
         if family is not None:
@@ -406,6 +426,10 @@ class Triangle(Element):
 
 
 class Tetrahedron(Element):
+    @staticmethod
+    def measure() -> float:
+        return 1.0 / 6.0
+
     @staticmethod
     def instantiate_quadrature(order: int, family: Polynomial):
         if family is not None:
@@ -436,7 +460,6 @@ class Tetrahedron(Element):
 
         # TODO: add Witherden and Vincent 2015,
 
-
         if order <= 1:
             weights = [1.0]
             coords = [Coords(1.0 / 4.0, 1.0 / 4.0, 1.0 / 4.0)]

warp/fem/geometry/geometry.py

@@ -3,6 +3,8 @@ from typing import Any
 import warp as wp
 
 from warp.fem.types import Sample, ElementIndex, Coords
+from warp.fem import cache
+
 from .element import Element
 
 
@@ -16,12 +18,15 @@ class Geometry:
     dimension: int = 0
 
     def cell_count(self):
+        """Number of cells in the geometry"""
         raise NotImplementedError
 
     def side_count(self):
+        """Number of sides in the geometry"""
         raise NotImplementedError
 
     def boundary_side_count(self):
+        """Number of boundary sides (sides with a single neighbour cell) in the geometry"""
         raise NotImplementedError
 
     def reference_cell(self) -> Element:
@@ -40,75 +45,142 @@ class Geometry:
         return self.name
 
     CellArg: wp.codegen.Struct
-    """Structure containing arguments to be passed to device function evaluating cell-related quantities"""
+    """Structure containing arguments to be passed to device functions evaluating cell-related quantities"""
 
     SideArg: wp.codegen.Struct
-    """Structure containing arguments to be passed to device function evaluating side-related quantities"""
+    """Structure containing arguments to be passed to device functions evaluating side-related quantities"""
+
+    SideIndexArg: wp.codegen.Struct
+    """Structure containing arguments to be passed to device functions for indexing sides"""
 
+    @staticmethod
     def cell_arg_value(self, device) -> "Geometry.CellArg":
         """Value of the arguments to be passed to cell-related device functions"""
         raise NotImplementedError
 
+    @staticmethod
     def cell_position(args: "Geometry.CellArg", s: "Sample"):
-        """Device function returning the cell position at a sample point"""
+        """Device function returning the world position of a cell sample point"""
+        raise NotImplementedError
+
+    @staticmethod
+    def cell_deformation_gradient(args: "Geometry.CellArg", s: "Sample"):
+        """Device function returning the transpose of the gradient of world position with respect to reference cell"""
         raise NotImplementedError
 
+    @staticmethod
+    def cell_inverse_deformation_gradient(args: "Geometry.CellArg", cell_index: ElementIndex, coords: Coords):
+        """Device function returning the matrix right-transforming a gradient w.r.t. cell space to a gradient w.r.t. world space
+        (i.e. the inverse deformation gradient)
+        """
+        raise NotImplementedError
+
+    @staticmethod
     def cell_lookup(args: "Geometry.CellArg", pos: Any):
         """Device function returning the cell sample point corresponding to a world position"""
         raise NotImplementedError
 
+    @staticmethod
     def cell_lookup(args: "Geometry.CellArg", pos: Any, guess: "Sample"):
         """Device function returning the cell sample point corresponding to a world position. Can use guess for faster lookup"""
         raise NotImplementedError
 
-    def cell_measure(args: "Geometry.CellArg", cell_index: ElementIndex, coords: Coords):
-        """Device function returning the measure determinant (e.g. volume, area) at a given point"""
-        raise NotImplementedError
-
+    @staticmethod
     def cell_measure(args: "Geometry.CellArg", s: "Sample"):
         """Device function returning the measure determinant (e.g. volume, area) at a given point"""
         raise NotImplementedError
 
-    def cell_measure_ratio(args: "Geometry.CellArg", s: "Sample"):
-        """Device function returning the ratio of the measure of a side to that of its neighbour cells"""
-        raise NotImplementedError
+    @wp.func
+    def cell_measure_ratio(args: Any, s: Sample):
+        return 1.0
 
+    @staticmethod
     def cell_normal(args: "Geometry.CellArg", s: "Sample"):
-        """Device function returning the element normal at a sample point"""
+        """Device function returning the element normal at a sample point.
+
+        For elements with the same dimension as the embedding space, this will be zero."""
         raise NotImplementedError
 
+    @staticmethod
     def side_arg_value(self, device) -> "Geometry.SideArg":
         """Value of the arguments to be passed to side-related device functions"""
         raise NotImplementedError
 
-    def boundary_side_index(args: "Geometry.SideArg", boundary_side_index: int):
+    @staticmethod
+    def boundary_side_index(args: "Geometry.SideIndexArg", boundary_side_index: int):
         """Device function returning the side index corresponding to a boundary side"""
         raise NotImplementedError
 
+    @staticmethod
     def side_position(args: "Geometry.SideArg", s: "Sample"):
         """Device function returning the side position at a sample point"""
         raise NotImplementedError
 
-    def side_measure(args: "Geometry.SideArg", cell_index: ElementIndex, coords: Coords):
-        """Device function returning the measure determinant (e.g. volume, area) at a given point"""
+    @staticmethod
+    def side_deformation_gradient(args: "Geometry.CellArg", s: "Sample"):
+        """Device function returning the gradient of world position with respect to reference cell"""
+        raise NotImplementedError
+
+    @staticmethod
+    def side_inner_inverse_deformation_gradient(args: "Geometry.CellArg", side_index: ElementIndex, coords: Coords):
+        """Device function returning the matrix right-transforming a gradient w.r.t. inner cell space to a gradient w.r.t. world space
+        (i.e. the inverse deformation gradient)
+        """
         raise NotImplementedError
 
+    @staticmethod
+    def side_outer_inverse_deformation_gradient(args: "Geometry.CellArg", side_index: ElementIndex, coords: Coords):
+        """Device function returning the matrix right-transforming a gradient w.r.t. outer cell space to a gradient w.r.t. world space
+        (i.e. the inverse deformation gradient)
+        """
+        raise NotImplementedError
+
+    @staticmethod
     def side_measure(args: "Geometry.SideArg", s: "Sample"):
         """Device function returning the measure determinant (e.g. volume, area) at a given point"""
         raise NotImplementedError
 
+    @staticmethod
     def side_measure_ratio(args: "Geometry.SideArg", s: "Sample"):
         """Device function returning the ratio of the measure of a side to that of its neighbour cells"""
         raise NotImplementedError
 
+    @staticmethod
     def side_normal(args: "Geometry.SideArg", s: "Sample"):
         """Device function returning the element normal at a sample point"""
         raise NotImplementedError
 
+    @staticmethod
     def side_inner_cell_index(args: "Geometry.SideArg", side_index: ElementIndex):
         """Device function returning the inner cell index for a given side"""
         raise NotImplementedError
 
+    @staticmethod
     def side_outer_cell_index(args: "Geometry.SideArg", side_index: ElementIndex):
         """Device function returning the outer cell index for a given side"""
         raise NotImplementedError
+
+    @staticmethod
+    def side_inner_cell_coords(args: "Geometry.SideArg", side_index: ElementIndex, side_coords: Coords):
+        """Device function returning the coordinates of a point on a side in the inner cell"""
+        raise NotImplementedError
+
+    @staticmethod
+    def side_outer_cell_coords(args: "Geometry.SideArg", side_index: ElementIndex, side_coords: Coords):
+        """Device function returning the coordinates of a point on a side in the outer cell"""
+        raise NotImplementedError
+
+    @staticmethod
+    def side_from_cell_coords(
+        args: "Geometry.SideArg",
+        side_index: ElementIndex,
+        element_index: ElementIndex,
+        element_coords: Coords,
+    ):
+        """Device function converting coordinates on a cell to coordinates on a side, or ``OUTSIDE``"""
+        raise NotImplementedError
+
+    @staticmethod
+    def side_to_cell_arg(side_arg: "Geometry.SideArg"):
+        """Device function converting a side-related argument value to a cell-related argument value, for promoting trace samples to the full space"""
+        raise NotImplementedError