warp-lang 1.4.2__py3-none-manylinux2014_aarch64.whl → 1.5.0__py3-none-manylinux2014_aarch64.whl

warp/fem/geometry/__init__.py

@@ -12,6 +12,6 @@ from .partition import (
     LinearGeometryPartition,
     WholeGeometryPartition,
 )
-from .quadmesh_2d import Quadmesh2D
+from .quadmesh import Quadmesh, Quadmesh2D, Quadmesh3D
 from .tetmesh import Tetmesh
-from .trimesh_2d import Trimesh2D
+from .trimesh import Trimesh, Trimesh2D, Trimesh3D

warp/fem/geometry/deformed_geometry.py

@@ -1,8 +1,6 @@
-from typing import Any
-
 import warp as wp
 from warp.fem import cache
-from warp.fem.types import Coords, ElementIndex, Sample, make_free_sample
+from warp.fem.types import Coords, ElementIndex, Sample
 
 from .geometry import Geometry
 
@@ -53,8 +51,6 @@ class DeformedGeometry(Geometry):
 
         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
 
@@ -66,11 +62,8 @@ class DeformedGeometry(Geometry):
         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()
+
+        self._make_default_dependent_implementations()
 
     @property
     def name(self):
@@ -111,26 +104,6 @@ class DeformedGeometry(Geometry):
 
         return cell_deformation_gradient if self._relative else cell_deformation_gradient_absolute
 
-    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:
@@ -182,81 +155,6 @@ class DeformedGeometry(Geometry):
 
         return side_deformation_gradient if self._relative else side_deformation_gradient_absolute
 
-    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):

warp/fem/geometry/element.py

@@ -5,6 +5,9 @@ from warp.fem.types import Coords
 
 
 class Element:
+    dimension = 0
+    """Intrinsic dimension of the element"""
+
     def measure() -> float:
         """Measure (area, volume, ...) of the reference element"""
         raise NotImplementedError
@@ -33,6 +36,8 @@ def _point_count_from_order(order: int, family: Polynomial):
 
 
 class Cube(Element):
+    dimension = 3
+
     @staticmethod
     def measure() -> float:
         return 1.0
@@ -52,6 +57,8 @@ class Cube(Element):
 
 
 class Square(Element):
+    dimension = 2
+
     @staticmethod
     def measure() -> float:
         return 1.0
@@ -71,6 +78,8 @@ class Square(Element):
 
 
 class LinearEdge(Element):
+    dimension = 1
+
     @staticmethod
     def measure() -> float:
         return 1.0
@@ -88,6 +97,8 @@ class LinearEdge(Element):
 
 
 class Triangle(Element):
+    dimension = 2
+
     @staticmethod
     def measure() -> float:
         return 0.5
@@ -430,6 +441,8 @@ class Triangle(Element):
 
 
 class Tetrahedron(Element):
+    dimension = 3
+
     @staticmethod
     def measure() -> float:
         return 1.0 / 6.0

warp/fem/geometry/geometry.py

@@ -1,10 +1,13 @@
 from typing import Any
 
 import warp as wp
-from warp.fem.types import Coords, ElementIndex, Sample
+from warp.fem import cache
+from warp.fem.types import Coords, ElementIndex, Sample, make_free_sample
 
 from .element import Element
 
+_mat32 = wp.mat(shape=(3, 2), dtype=float)
+
 
 class Geometry:
     """
@@ -35,6 +38,11 @@ class Geometry:
         """Prototypical element for a side"""
         raise NotImplementedError
 
+    @property
+    def cell_dimension(self) -> int:
+        """Manifold dimension of the geometry cells"""
+        return self.reference_cell().dimension
+
     @property
     def name(self) -> str:
         return self.__class__.__name__
@@ -67,7 +75,7 @@ class Geometry:
         raise NotImplementedError
 
     @staticmethod
-    def cell_inverse_deformation_gradient(args: "Geometry.CellArg", cell_index: ElementIndex, coords: Coords):
+    def cell_inverse_deformation_gradient(args: "Geometry.CellArg", s: "Sample"):
         """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)
         """
@@ -115,12 +123,12 @@ class Geometry:
         raise NotImplementedError
 
     @staticmethod
-    def side_deformation_gradient(args: "Geometry.CellArg", s: "Sample"):
-        """Device function returning the gradient of world position with respect to reference cell"""
+    def side_deformation_gradient(args: "Geometry.SideArg", s: "Sample"):
+        """Device function returning the gradient of world position with respect to reference side"""
         raise NotImplementedError
 
     @staticmethod
-    def side_inner_inverse_deformation_gradient(args: "Geometry.CellArg", side_index: ElementIndex, coords: Coords):
+    def side_inner_inverse_deformation_gradient(args: "Geometry.Siderg", 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)
         """
@@ -182,3 +190,155 @@ class Geometry:
     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
+
+    # Default implementations for dependent quantities
+    # Can be overridden in derived classes if more efficient implementations exist
+
+    def _make_default_dependent_implementations(self):
+        self.cell_inverse_deformation_gradient = self._make_cell_inverse_deformation_gradient()
+        self.cell_measure = self._make_cell_measure()
+        self.cell_normal = self._make_cell_normal()
+
+        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()
+
+    @wp.func
+    def _element_measure(F: wp.vec2):
+        return wp.length(F)
+
+    @wp.func
+    def _element_measure(F: wp.vec3):
+        return wp.length(F)
+
+    @wp.func
+    def _element_measure(F: _mat32):
+        Ft = wp.transpose(F)
+        Fcross = wp.cross(Ft[0], Ft[1])
+        return wp.length(Fcross)
+
+    @wp.func
+    def _element_measure(F: wp.mat33):
+        return wp.abs(wp.determinant(F))
+
+    @wp.func
+    def _element_measure(F: wp.mat22):
+        return wp.abs(wp.determinant(F))
+
+    @wp.func
+    def _element_normal(F: wp.vec2):
+        return wp.normalize(wp.vec2(-F[1], F[0]))
+
+    @wp.func
+    def _element_normal(F: _mat32):
+        Ft = wp.transpose(F)
+        Fcross = wp.cross(Ft[0], Ft[1])
+        return wp.normalize(Fcross)
+
+    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):
+            F = self.cell_deformation_gradient(args, s)
+            return Geometry._element_measure(F) * REF_MEASURE
+
+        return cell_measure
+
+    def _make_cell_normal(self):
+        cell_dim = self.reference_cell().dimension
+        geo_dim = self.dimension
+        normal_vec = wp.vec(length=geo_dim, dtype=float)
+
+        @cache.dynamic_func(suffix=self.name)
+        def zero_normal(args: self.CellArg, s: Sample):
+            return normal_vec(0.0)
+
+        @cache.dynamic_func(suffix=self.name)
+        def cell_hyperplane_normal(args: self.CellArg, s: Sample):
+            F = self.cell_deformation_gradient(args, s)
+            return Geometry._element_normal(F)
+
+        if cell_dim == geo_dim:
+            return zero_normal
+        if cell_dim == geo_dim - 1:
+            return cell_hyperplane_normal
+
+        return None
+
+    def _make_cell_inverse_deformation_gradient(self):
+        cell_dim = self.reference_cell().dimension
+        geo_dim = self.dimension
+
+        @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))
+
+        @cache.dynamic_func(suffix=self.name)
+        def cell_pseudoinverse_deformation_gradient(cell_arg: self.CellArg, s: Sample):
+            F = self.cell_deformation_gradient(cell_arg, s)
+            Ft = wp.transpose(F)
+            return wp.inverse(Ft * F) * Ft
+
+        return cell_inverse_deformation_gradient if cell_dim == geo_dim else cell_pseudoinverse_deformation_gradient
+
+    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 Geometry._element_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):
+        side_dim = self.reference_side().dimension
+        geo_dim = self.dimension
+
+        @cache.dynamic_func(suffix=self.name)
+        def hyperplane_normal(args: self.SideArg, s: Sample):
+            F = self.side_deformation_gradient(args, s)
+            return Geometry._element_normal(F)
+
+        if side_dim == geo_dim - 1:
+            return hyperplane_normal
+
+        return None
+
+    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))
+
+        return side_inner_inverse_deformation_gradient
+
+    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))
+
+        return side_outer_inverse_deformation_gradient

warp/fem/geometry/grid_2d.py

@@ -29,7 +29,7 @@ class Grid2D(Geometry):
         Args:
             res: Resolution of the grid along each dimension
             bounds_lo: Position of the lower bound of the axis-aligned grid
-            bounds_up: Position of the upper bound of the axis-aligned grid
+            bounds_hi: Position of the upper bound of the axis-aligned grid
         """
 
         if bounds_lo is None:
@@ -158,10 +158,7 @@ class Grid2D(Geometry):
             return Grid2D.Side(axis, origin)
 
         axis_side_index = side_index - 2 * arg.cell_count
-        if axis_side_index < arg.axis_offsets[1]:
-            axis = 0
-        else:
-            axis = 1
+        axis = wp.select(axis_side_index < arg.axis_offsets[1], 1, 0)
 
         altitude = arg.cell_arg.res[Grid2D.ROTATION[axis, 0]]
         longitude = axis_side_index - arg.axis_offsets[axis]
@@ -230,7 +227,7 @@ class Grid2D(Geometry):
 
         args.axis_offsets = wp.vec2i(
             0,
-            self.res[0],
+            self.res[1],
         )
         args.cell_count = self.cell_count()
         args.cell_arg = self.cell_arg_value(device)

warp/fem/geometry/grid_3d.py

@@ -23,17 +23,13 @@ class Grid3D(Geometry):
 
     dimension = 3
 
-    Permutation = wp.types.matrix(shape=(3, 3), dtype=int)
-    LOC_TO_WORLD = wp.constant(Permutation(0, 1, 2, 1, 2, 0, 2, 0, 1))
-    WORLD_TO_LOC = wp.constant(Permutation(0, 1, 2, 2, 0, 1, 1, 2, 0))
-
     def __init__(self, res: wp.vec3i, bounds_lo: Optional[wp.vec3] = None, bounds_hi: Optional[wp.vec3] = None):
         """Constructs a dense 3D grid
 
         Args:
             res: Resolution of the grid along each dimension
             bounds_lo: Position of the lower bound of the axis-aligned grid
-            bounds_up: Position of the upper bound of the axis-aligned grid
+            bounds_hi: Position of the upper bound of the axis-aligned grid
         """
 
         if bounds_lo is None:
@@ -148,28 +144,32 @@ class Grid3D(Geometry):
     @wp.func
     def _world_to_local(axis: int, vec: Any):
         return type(vec)(
-            vec[Grid3D.LOC_TO_WORLD[axis, 0]],
-            vec[Grid3D.LOC_TO_WORLD[axis, 1]],
-            vec[Grid3D.LOC_TO_WORLD[axis, 2]],
+            vec[axis],
+            vec[(axis + 1) % 3],
+            vec[(axis + 2) % 3],
         )
 
     @wp.func
     def _local_to_world(axis: int, vec: Any):
         return type(vec)(
-            vec[Grid3D.WORLD_TO_LOC[axis, 0]],
-            vec[Grid3D.WORLD_TO_LOC[axis, 1]],
-            vec[Grid3D.WORLD_TO_LOC[axis, 2]],
+            vec[(2 * axis) % 3],
+            vec[(2 * axis + 1) % 3],
+            vec[(2 * axis + 2) % 3],
         )
 
+    @wp.func
+    def _local_to_world_axis(axis: int, loc_index: Any):
+        return (axis + loc_index) % 3
+
     @wp.func
     def side_index(arg: SideArg, side: Side):
-        alt_axis = Grid3D.LOC_TO_WORLD[side.axis, 0]
+        alt_axis = Grid3D._local_to_world_axis(side.axis, 0)
         if side.origin[0] == arg.cell_arg.res[alt_axis]:
             # Upper-boundary side
             longitude = side.origin[1]
             latitude = side.origin[2]
 
-            latitude_res = arg.cell_arg.res[Grid3D.LOC_TO_WORLD[side.axis, 2]]
+            latitude_res = arg.cell_arg.res[Grid3D._local_to_world_axis(side.axis, 2)]
             lat_long = latitude_res * longitude + latitude
 
             return 3 * arg.cell_count + arg.axis_offsets[side.axis] + lat_long
@@ -179,24 +179,27 @@ class Grid3D(Geometry):
 
     @wp.func
     def get_side(arg: SideArg, side_index: ElementIndex):
+        res = arg.cell_arg.res
+
         if side_index < 3 * arg.cell_count:
             axis = side_index // arg.cell_count
             cell_index = side_index - axis * arg.cell_count
-            origin = Grid3D._world_to_local(axis, Grid3D.get_cell(arg.cell_arg.res, cell_index))
-            return Grid3D.Side(axis, origin)
+            origin_loc = Grid3D._world_to_local(axis, Grid3D.get_cell(res, cell_index))
+            return Grid3D.Side(axis, origin_loc)
 
+        axis_offsets = arg.axis_offsets
         axis_side_index = side_index - 3 * arg.cell_count
-        if axis_side_index < arg.axis_offsets[1]:
+        if axis_side_index < axis_offsets[1]:
             axis = 0
-        elif axis_side_index < arg.axis_offsets[2]:
+        elif axis_side_index < axis_offsets[2]:
             axis = 1
         else:
             axis = 2
 
-        altitude = arg.cell_arg.res[Grid3D.LOC_TO_WORLD[axis, 0]]
+        altitude = res[Grid3D._local_to_world_axis(axis, 0)]
 
-        lat_long = axis_side_index - arg.axis_offsets[axis]
-        latitude_res = arg.cell_arg.res[Grid3D.LOC_TO_WORLD[axis, 2]]
+        lat_long = axis_side_index - axis_offsets[axis]
+        latitude_res = res[Grid3D._local_to_world_axis(axis, 2)]
 
         longitude = lat_long // latitude_res
         latitude = lat_long - longitude * latitude_res
@@ -299,7 +302,7 @@ class Grid3D(Geometry):
             axis = 2
 
         lat_long = axis_side_index - args.axis_offsets[axis]
-        latitude_res = args.cell_arg.res[Grid3D.LOC_TO_WORLD[axis, 2]]
+        latitude_res = args.cell_arg.res[Grid3D._local_to_world_axis(axis, 2)]
 
         longitude = lat_long // latitude_res
         latitude = lat_long - longitude * latitude_res
@@ -347,14 +350,14 @@ class Grid3D(Geometry):
     @wp.func
     def side_measure(args: SideArg, s: Sample):
         side = Grid3D.get_side(args, s.element_index)
-        long_axis = Grid3D.LOC_TO_WORLD[side.axis, 1]
-        lat_axis = Grid3D.LOC_TO_WORLD[side.axis, 2]
+        long_axis = Grid3D._local_to_world_axis(side.axis, 1)
+        lat_axis = Grid3D._local_to_world_axis(side.axis, 2)
         return args.cell_arg.cell_size[long_axis] * args.cell_arg.cell_size[lat_axis]
 
     @wp.func
     def side_measure_ratio(args: SideArg, s: Sample):
         side = Grid3D.get_side(args, s.element_index)
-        alt_axis = Grid3D.LOC_TO_WORLD[side.axis, 0]
+        alt_axis = Grid3D._local_to_world_axis(side.axis, 0)
         return 1.0 / args.cell_arg.cell_size[alt_axis]
 
     @wp.func
@@ -381,7 +384,7 @@ class Grid3D(Geometry):
     def side_outer_cell_index(arg: SideArg, side_index: ElementIndex):
         side = Grid3D.get_side(arg, side_index)
 
-        alt_axis = Grid3D.LOC_TO_WORLD[side.axis, 0]
+        alt_axis = Grid3D._local_to_world_axis(side.axis, 0)
 
         outer_alt = wp.select(
             side.origin[0] == arg.cell_arg.res[alt_axis], side.origin[0], arg.cell_arg.res[alt_axis] - 1
@@ -406,7 +409,7 @@ class Grid3D(Geometry):
     def side_outer_cell_coords(args: SideArg, side_index: ElementIndex, side_coords: Coords):
         side = Grid3D.get_side(args, side_index)
 
-        alt_axis = Grid3D.LOC_TO_WORLD[side.axis, 0]
+        alt_axis = Grid3D._local_to_world_axis(side.axis, 0)
         outer_alt = wp.select(side.origin[0] == args.cell_arg.res[alt_axis], 0.0, 1.0)
 
         side_coord0 = wp.select(side.origin[0] == 0, side_coords[0], 1.0 - side_coords[0])
@@ -424,8 +427,8 @@ class Grid3D(Geometry):
         cell = Grid3D.get_cell(args.cell_arg.res, element_index)
 
         if float(side.origin[0] - cell[side.axis]) == element_coords[side.axis]:
-            long_axis = Grid3D.LOC_TO_WORLD[side.axis, 1]
-            lat_axis = Grid3D.LOC_TO_WORLD[side.axis, 2]
+            long_axis = Grid3D._local_to_world_axis(side.axis, 1)
+            lat_axis = Grid3D._local_to_world_axis(side.axis, 2)
             long_coord = element_coords[long_axis]
             long_coord = wp.select(side.origin[0] == 0, long_coord, 1.0 - long_coord)
             return Coords(long_coord, element_coords[lat_axis], 0.0)

warp/fem/geometry/hexmesh.py

@@ -7,7 +7,7 @@ from warp.fem.cache import (
     borrow_temporary_like,
     cached_arg_value,
 )
-from warp.fem.types import OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
+from warp.fem.types import OUTSIDE, Coords, ElementIndex, Sample
 
 from .element import Cube, Square
 from .geometry import Geometry
@@ -142,6 +142,8 @@ class Hexmesh(Geometry):
         self._edge_count = 0
         self._build_topology(temporary_store)
 
+        self._make_default_dependent_implementations()
+
     def cell_count(self):
         return self.hex_vertex_indices.shape[0]
 
@@ -246,18 +248,6 @@ class Hexmesh(Geometry):
             + wp.outer(cell_arg.positions[hex_idx[7]], wp.vec3(-w_p[1] * w_p[2], w_m[0] * w_p[2], w_m[0] * w_p[1]))
         )
 
-    @wp.func
-    def cell_inverse_deformation_gradient(cell_arg: CellArg, s: Sample):
-        return wp.inverse(Hexmesh.cell_deformation_gradient(cell_arg, s))
-
-    @wp.func
-    def cell_measure(args: CellArg, s: Sample):
-        return wp.abs(wp.determinant(Hexmesh.cell_deformation_gradient(args, s)))
-
-    @wp.func
-    def cell_normal(args: CellArg, s: Sample):
-        return wp.vec3(0.0)
-
     @cached_arg_value
     def side_index_arg_value(self, device) -> SideIndexArg:
         args = self.SideIndexArg()
@@ -319,39 +309,6 @@ class Hexmesh(Geometry):
         v1, v2 = Hexmesh._side_deformation_vecs(args, s.element_index, s.element_coords)
         return _mat32(v1, v2)
 
-    @wp.func
-    def side_inner_inverse_deformation_gradient(args: SideArg, s: Sample):
-        cell_index = Hexmesh.side_inner_cell_index(args, s.element_index)
-        cell_coords = Hexmesh.side_inner_cell_coords(args, s.element_index, s.element_coords)
-        return Hexmesh.cell_inverse_deformation_gradient(args.cell_arg, make_free_sample(cell_index, cell_coords))
-
-    @wp.func
-    def side_outer_inverse_deformation_gradient(args: SideArg, s: Sample):
-        cell_index = Hexmesh.side_outer_cell_index(args, s.element_index)
-        cell_coords = Hexmesh.side_outer_cell_coords(args, s.element_index, s.element_coords)
-        return Hexmesh.cell_inverse_deformation_gradient(args.cell_arg, make_free_sample(cell_index, cell_coords))
-
-    @wp.func
-    def side_measure(args: SideArg, s: Sample):
-        v1, v2 = Hexmesh._side_deformation_vecs(args, s.element_index, s.element_coords)
-        return wp.length(wp.cross(v1, v2))
-
-    @wp.func
-    def side_measure_ratio(args: SideArg, s: Sample):
-        inner = Hexmesh.side_inner_cell_index(args, s.element_index)
-        outer = Hexmesh.side_outer_cell_index(args, s.element_index)
-        inner_coords = Hexmesh.side_inner_cell_coords(args, s.element_index, s.element_coords)
-        outer_coords = Hexmesh.side_outer_cell_coords(args, s.element_index, s.element_coords)
-        return Hexmesh.side_measure(args, s) / wp.min(
-            Hexmesh.cell_measure(args.cell_arg, make_free_sample(inner, inner_coords)),
-            Hexmesh.cell_measure(args.cell_arg, make_free_sample(outer, outer_coords)),
-        )
-
-    @wp.func
-    def side_normal(args: SideArg, s: Sample):
-        v1, v2 = Hexmesh._side_deformation_vecs(args, s.element_index, s.element_coords)
-        return wp.normalize(wp.cross(v1, v2))
-
     @wp.func
     def side_inner_cell_index(arg: SideArg, side_index: ElementIndex):
         return arg.face_hex_indices[side_index][0]

warp/fem/geometry/nanogrid.py

@@ -4,6 +4,7 @@ import numpy as np
 
 import warp as wp
 from warp.fem import cache, utils
+from warp.fem.linalg import basis_element
 from warp.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, Sample, make_free_sample
 
 from .element import Cube, Square
@@ -300,8 +301,8 @@ class Nanogrid(Geometry):
 
     @wp.func
     def _face_tangent_vecs(cell_grid: wp.uint64, axis: int, flip: int):
-        u_axis = utils.unit_element(wp.vec3(), (axis + 1 + flip) % 3)
-        v_axis = utils.unit_element(wp.vec3(), (axis + 2 - flip) % 3)
+        u_axis = basis_element(wp.vec3(), (axis + 1 + flip) % 3)
+        v_axis = basis_element(wp.vec3(), (axis + 2 - flip) % 3)
 
         return wp.volume_index_to_world_dir(cell_grid, u_axis), wp.volume_index_to_world_dir(cell_grid, v_axis)