warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__py3-none-manylinux_2_34_aarch64.whl

warp/fem/space/basis_function_space.py

@@ -13,13 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Optional
+from typing import Any, ClassVar, Optional
 
 import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Geometry
-from warp.fem.linalg import basis_element, generalized_inner, generalized_outer
+from warp.fem.linalg import generalized_inner, generalized_outer
 from warp.fem.types import NULL_QP_INDEX, Coords, ElementIndex, make_free_sample
+from warp.fem.utils import type_basis_element
 
 from .basis_space import BasisSpace
 from .dof_mapper import DofMapper, IdentityMapper
@@ -30,6 +31,21 @@ from .partition import SpacePartition, make_space_partition
 class CollocatedFunctionSpace(FunctionSpace):
     """Function space where values are collocated at nodes"""
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "node_basis_element": lambda obj: obj._make_node_basis_element(),
+        "value_basis_element": lambda obj: obj._make_value_basis_element(),
+        "node_coords_in_element": lambda obj: obj._basis.make_node_coords_in_element(),
+        "node_quadrature_weight": lambda obj: obj._basis.make_node_quadrature_weight(),
+        "element_inner_weight": lambda obj: obj._basis.make_element_inner_weight(),
+        "element_inner_weight_gradient": lambda obj: obj._basis.make_element_inner_weight_gradient(),
+        "element_outer_weight": lambda obj: obj._basis.make_element_outer_weight(),
+        "element_outer_weight_gradient": lambda obj: obj._basis.make_element_outer_weight_gradient(),
+        "space_value": lambda obj: obj._make_space_value(),
+        "space_gradient": lambda obj: obj._make_space_gradient(),
+        "space_divergence": lambda obj: obj._make_space_divergence(),
+        "node_dof_value": lambda obj: obj._make_node_dof_value(),
+    }
+
     @wp.struct
     class LocalValueMap:
         pass
@@ -47,24 +63,11 @@ class CollocatedFunctionSpace(FunctionSpace):
 
         self.SpaceArg = self._basis.BasisArg
         self.space_arg_value = self._basis.basis_arg_value
+        self.fill_space_arg = self._basis.fill_basis_arg
 
         self.ORDER = self._basis.ORDER
 
-        self.node_basis_element = self._make_node_basis_element()
-        self.value_basis_element = self._make_value_basis_element()
-
-        self.node_coords_in_element = self._basis.make_node_coords_in_element()
-        self.node_quadrature_weight = self._basis.make_node_quadrature_weight()
-        self.element_inner_weight = self._basis.make_element_inner_weight()
-        self.element_inner_weight_gradient = self._basis.make_element_inner_weight_gradient()
-        self.element_outer_weight = self._basis.make_element_outer_weight()
-        self.element_outer_weight_gradient = self._basis.make_element_outer_weight_gradient()
-
-        self.space_value = self._make_space_value()
-        self.space_gradient = self._make_space_gradient()
-        self.space_divergence = self._make_space_divergence()
-
-        self.node_dof_value = self._make_node_dof_value()
+        cache.setup_dynamic_attributes(self)
 
         # For backward compatibility
         if hasattr(basis, "node_grid"):
@@ -100,11 +103,8 @@ class CollocatedFunctionSpace(FunctionSpace):
         return CollocatedFunctionSpaceTrace(self)
 
     def _make_node_basis_element(self):
-        @cache.dynamic_func(suffix=self.name)
-        def node_basis_element(dof_coord: int):
-            return basis_element(self.dof_dtype(0.0), dof_coord)
-
-        return node_basis_element
+        basis_element = type_basis_element(self.dof_dtype)
+        return basis_element
 
     def _make_value_basis_element(self):
         @cache.dynamic_func(suffix=self.name)
@@ -196,19 +196,34 @@ class CollocatedFunctionSpaceTrace(CollocatedFunctionSpace):
 class VectorValuedFunctionSpace(FunctionSpace):
     """Function space whose values are vectors"""
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "value_basis_element": lambda obj: obj._make_value_basis_element(),
+        "node_coords_in_element": lambda obj: obj._basis.make_node_coords_in_element(),
+        "node_quadrature_weight": lambda obj: obj._basis.make_node_quadrature_weight(),
+        "element_inner_weight": lambda obj: obj._basis.make_element_inner_weight(),
+        "element_inner_weight_gradient": lambda obj: obj._basis.make_element_inner_weight_gradient(),
+        "element_outer_weight": lambda obj: obj._basis.make_element_outer_weight(),
+        "element_outer_weight_gradient": lambda obj: obj._basis.make_element_outer_weight_gradient(),
+        "space_value": lambda obj: obj._make_space_value(),
+        "space_gradient": lambda obj: obj._make_space_gradient(),
+        "space_divergence": lambda obj: obj._make_space_divergence(),
+        "node_dof_value": lambda obj: obj._make_node_dof_value(),
+    }
+
     def __init__(self, basis: BasisSpace):
         self._basis = basis
 
         super().__init__(topology=basis.topology)
 
         self.dtype = cache.cached_vec_type(self.geometry.dimension, dtype=float)
-        self.dof_dtype = float
+        self.dof_dtype = wp.float32
 
         self.VALUE_DOF_COUNT = self.geometry.dimension
         self.NODE_DOF_COUNT = 1
 
         self.SpaceArg = self._basis.BasisArg
         self.space_arg_value = self._basis.basis_arg_value
+        self.fill_space_arg = self._basis.fill_basis_arg
 
         self.ORDER = self._basis.ORDER
 
@@ -216,20 +231,7 @@ class VectorValuedFunctionSpace(FunctionSpace):
             shape=(self.geometry.dimension, self.geometry.cell_dimension), dtype=float
         )
 
-        self.value_basis_element = self._make_value_basis_element()
-
-        self.node_coords_in_element = self._basis.make_node_coords_in_element()
-        self.node_quadrature_weight = self._basis.make_node_quadrature_weight()
-        self.element_inner_weight = self._basis.make_element_inner_weight()
-        self.element_inner_weight_gradient = self._basis.make_element_inner_weight_gradient()
-        self.element_outer_weight = self._basis.make_element_outer_weight()
-        self.element_outer_weight_gradient = self._basis.make_element_outer_weight_gradient()
-
-        self.space_value = self._make_space_value()
-        self.space_gradient = self._make_space_gradient()
-        self.space_divergence = self._make_space_divergence()
-
-        self.node_dof_value = self._make_node_dof_value()
+        cache.setup_dynamic_attributes(self, cls=__class__)
 
     @property
     def name(self):
@@ -254,9 +256,11 @@ class VectorValuedFunctionSpace(FunctionSpace):
         return 1.0
 
     def _make_value_basis_element(self):
+        basis_element = type_basis_element(self.dtype)
+
         @cache.dynamic_func(suffix=self.name)
         def value_basis_element(dof_coord: int, value_map: Any):
-            return value_map * basis_element(self.dtype(0.0), dof_coord)
+            return value_map * basis_element(dof_coord)
 
         return value_basis_element
 
@@ -319,11 +323,15 @@ class VectorValuedFunctionSpace(FunctionSpace):
 class CovariantFunctionSpace(VectorValuedFunctionSpace):
     """Function space whose values are covariant vectors"""
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "local_value_map_inner": lambda obj: obj._make_local_value_map(),
+        "local_value_map_outer": lambda obj: obj.local_value_map_inner,
+    }
+
     def __init__(self, basis: BasisSpace):
         super().__init__(basis)
 
-        self.local_value_map_inner = self._make_local_value_map()
-        self.local_value_map_outer = self.local_value_map_inner
+        cache.setup_dynamic_attributes(self, cls=__class__)
 
     def trace(self) -> "CovariantFunctionSpaceTrace":
         return CovariantFunctionSpaceTrace(self)
@@ -348,12 +356,16 @@ class CovariantFunctionSpace(VectorValuedFunctionSpace):
 class CovariantFunctionSpaceTrace(VectorValuedFunctionSpace):
     """Trace of a :class:`CovariantFunctionSpace`"""
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "local_value_map_inner": lambda obj: obj._make_local_value_map_inner(),
+        "local_value_map_outer": lambda obj: obj._make_local_value_map_outer(),
+    }
+
     def __init__(self, space: VectorValuedFunctionSpace):
         self._space = space
         super().__init__(space._basis.trace())
 
-        self.local_value_map_inner = self._make_local_value_map_inner()
-        self.local_value_map_outer = self._make_local_value_map_outer()
+        cache.setup_dynamic_attributes(self, cls=__class__)
 
     @property
     def name(self):
@@ -396,11 +408,15 @@ class CovariantFunctionSpaceTrace(VectorValuedFunctionSpace):
 class ContravariantFunctionSpace(VectorValuedFunctionSpace):
     """Function space whose values are contravariant vectors"""
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "local_value_map_inner": lambda obj: obj._make_local_value_map(),
+        "local_value_map_outer": lambda obj: obj.local_value_map_inner,
+    }
+
     def __init__(self, basis: BasisSpace):
         super().__init__(basis)
 
-        self.local_value_map_inner = self._make_local_value_map()
-        self.local_value_map_outer = self.local_value_map_inner
+        cache.setup_dynamic_attributes(self, cls=__class__)
 
     def trace(self) -> "ContravariantFunctionSpaceTrace":
         return ContravariantFunctionSpaceTrace(self)
@@ -421,12 +437,16 @@ class ContravariantFunctionSpace(VectorValuedFunctionSpace):
 class ContravariantFunctionSpaceTrace(VectorValuedFunctionSpace):
     """Trace of a :class:`ContravariantFunctionSpace`"""
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "local_value_map_inner": lambda obj: obj._make_local_value_map_inner(),
+        "local_value_map_outer": lambda obj: obj._make_local_value_map_outer(),
+    }
+
     def __init__(self, space: ContravariantFunctionSpace):
         self._space = space
         super().__init__(space._basis.trace())
 
-        self.local_value_map_inner = self._make_local_value_map_inner()
-        self.local_value_map_outer = self._make_local_value_map_outer()
+        cache.setup_dynamic_attributes(self, cls=__class__)
 
     @property
     def name(self):

warp/fem/space/basis_space.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import ClassVar, Optional
 
 import numpy as np
 
@@ -71,6 +71,9 @@ class BasisSpace:
         """Value for the argument structure to be passed to device functions"""
         return BasisSpace.BasisArg()
 
+    def fill_basis_arg(self, arg, device):
+        pass
+
     # Helpers for generating node positions
 
     def node_positions(self, out: Optional[wp.array] = None) -> wp.array:
@@ -174,6 +177,7 @@ class ShapeBasisSpace(BasisSpace):
         if self.value is not ShapeFunction.Value.Scalar:
             self.BasisArg = self.topology.TopologyArg
             self.basis_arg_value = self.topology.topo_arg_value
+            self.fill_basis_arg = self.topology.fill_topo_arg
 
         self.ORDER = self._shape.ORDER
 
@@ -335,6 +339,7 @@ class TraceBasisSpace(BasisSpace):
         self._basis = basis
         self.BasisArg = self._basis.BasisArg
         self.basis_arg_value = self._basis.basis_arg_value
+        self.fill_basis_arg = self._basis.fill_basis_arg
 
     @property
     def name(self):
@@ -507,6 +512,12 @@ def make_discontinuous_basis_space(geometry: Geometry, shape: ShapeFunction):
 class UnstructuredPointTopology(SpaceTopology):
     """Topology for unstructured points defined from quadrature formula. See :class:`PointBasisSpace`"""
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "element_node_index": lambda obj: obj._make_element_node_index(),
+        "element_node_count": lambda obj: obj._make_element_node_count(),
+        "side_neighbor_node_counts": lambda obj: obj._make_side_neighbor_node_counts(),
+    }
+
     def __init__(self, quadrature: Quadrature):
         if quadrature.max_points_per_element() is None:
             raise ValueError("Quadrature must define a maximum number of points per element")
@@ -516,12 +527,12 @@ class UnstructuredPointTopology(SpaceTopology):
 
         self._quadrature = quadrature
         self.TopologyArg = quadrature.Arg
+        self.topo_arg_value = quadrature.arg_value
+        self.fill_topo_arg = quadrature.fill_arg
 
         super().__init__(quadrature.domain.geometry, max_nodes_per_element=quadrature.max_points_per_element())
 
-        self.element_node_index = self._make_element_node_index()
-        self.element_node_count = self._make_element_node_count()
-        self.side_neighbor_node_counts = self._make_side_neighbor_node_counts()
+        cache.setup_dynamic_attributes(self, cls=__class__)
 
     def node_count(self):
         return self._quadrature.total_point_count()
@@ -582,6 +593,8 @@ class PointBasisSpace(BasisSpace):
 
         self.BasisArg = quadrature.Arg
         self.basis_arg_value = quadrature.arg_value
+        self.fill_basis_arg = quadrature.fill_arg
+
         self.ORDER = 0
 
         self.make_element_outer_weight = self.make_element_inner_weight

warp/fem/space/dof_mapper.py

@@ -57,7 +57,7 @@ class IdentityMapper(DofMapper):
         self.value_dtype = dtype
         self.dof_dtype = dtype
 
-        size = warp.types.type_length(dtype)
+        size = warp.types.type_size(dtype)
         self.DOF_SIZE = wp.constant(size)
 
     @wp.func

warp/fem/space/function_space.py

@@ -122,13 +122,13 @@ class FunctionSpace:
         raise NotImplementedError
 
     def gradient_valid(self) -> bool:
-        """Whether gradient operator can be computed. Only for scalar and vector fields as higher-order tensors are not support yet"""
+        """Whether gradient operator can be computed. Only for scalar and vector fields as higher-order tensors are not supported yet"""
         return not wp.types.type_is_matrix(self.dtype)
 
     def divergence_valid(self) -> bool:
         """Whether divergence of this field can be computed. Only for vector and tensor fields with same dimension as embedding geometry"""
         if wp.types.type_is_vector(self.dtype):
-            return wp.types.type_length(self.dtype) == self.geometry.dimension
+            return wp.types.type_size(self.dtype) == self.geometry.dimension
         if wp.types.type_is_matrix(self.dtype):
             return self.dtype._shape_[0] == self.geometry.dimension
         return False

warp/fem/space/grid_2d_function_space.py

@@ -41,6 +41,9 @@ class Grid2DSpaceTopology(SpaceTopology):
     def topo_arg_value(self, device):
         return self.geometry.side_arg_value(device)
 
+    def fill_topo_arg(self, arg: Grid2D.SideArg, device):
+        self.geometry.fill_side_arg(arg, device)
+
     def node_count(self) -> int:
         return (
             self.geometry.vertex_count() * self._shape.VERTEX_NODE_COUNT
@@ -84,7 +87,7 @@ class Grid2DSpaceTopology(SpaceTopology):
                 axis = 1 - (node_type - SquareShapeFunction.EDGE_X)
 
                 cell = Grid2D.get_cell(cell_arg.res, element_index)
-                origin = wp.vec2i(cell[Grid2D.ROTATION[axis, 0]] + type_instance, cell[Grid2D.ROTATION[axis, 1]])
+                origin = Grid2D.orient(axis, cell) + wp.vec2i(type_instance, 0)
 
                 side = Grid2D.Side(axis, origin)
                 side_index = Grid2D.side_index(topo_arg, side)

warp/fem/space/hexmesh_function_space.py

@@ -85,13 +85,15 @@ class HexmeshSpaceTopology(SpaceTopology):
     @cache.cached_arg_value
     def topo_arg_value(self, device):
         arg = HexmeshTopologyArg()
+        self.fill_topo_arg(arg, device)
+        return arg
+
+    def fill_topo_arg(self, arg: HexmeshTopologyArg, device):
         arg.hex_edge_indices = self._hex_edge_indices.to(device)
         arg.hex_face_indices = self._hex_face_indices.to(device)
-
         arg.vertex_count = self._mesh.vertex_count()
         arg.face_count = self._mesh.side_count()
         arg.edge_count = self._edge_count
-        return arg
 
     def _compute_hex_face_indices(self):
         self._hex_face_indices = wp.empty(

warp/fem/space/nanogrid_function_space.py

@@ -72,7 +72,10 @@ class NanogridSpaceTopology(SpaceTopology):
     @cache.cached_arg_value
     def topo_arg_value(self, device):
         arg = NanogridTopologyArg()
+        self.fill_topo_arg(arg, device)
+        return arg
 
+    def fill_topo_arg(self, arg, device):
         arg.vertex_grid = self._vertex_grid
         arg.face_grid = self._face_grid
         arg.edge_grid = self._edge_grid
@@ -80,7 +83,6 @@ class NanogridSpaceTopology(SpaceTopology):
         arg.vertex_count = self._grid.vertex_count()
         arg.face_count = self._face_count
         arg.edge_count = self._edge_count
-        return arg
 
     def _make_element_node_index(self):
         element_node_index_generic = self._make_element_node_index_generic()

warp/fem/space/partition.py

@@ -50,6 +50,9 @@ class SpacePartition:
     def partition_arg_value(self, device):
         pass
 
+    def fill_partition_arg(self, arg, device):
+        pass
+
     @staticmethod
     def partition_node_index(args: "PartitionArg", space_node_index: int):
         """Returns the index in the partition of a function space node, or ``NULL_NODE_INDEX`` if it does not exist"""
@@ -93,12 +96,15 @@ class WholeSpacePartition(SpacePartition):
     def partition_arg_value(self, device):
         return WholeSpacePartition.PartitionArg()
 
+    def fill_partition_arg(self, arg, device):
+        pass
+
     @wp.func
     def partition_node_index(args: Any, space_node_index: int):
         return space_node_index
 
     def __eq__(self, other: SpacePartition) -> bool:
-        return isinstance(other, SpacePartition) and self.space_topology == other.space_topology
+        return isinstance(other, WholeSpacePartition) and self.space_topology == other.space_topology
 
     @property
     def name(self) -> str:
@@ -160,9 +166,12 @@ class NodePartition(SpacePartition):
     @cache.cached_arg_value
     def partition_arg_value(self, device):
         arg = NodePartition.PartitionArg()
-        arg.space_to_partition = self._space_to_partition.array.to(device)
+        self.fill_partition_arg(arg, device)
         return arg
 
+    def fill_partition_arg(self, arg, device):
+        arg.space_to_partition = self._space_to_partition.array.to(device)
+
     @wp.func
     def partition_node_index(args: PartitionArg, space_node_index: int):
         return args.space_to_partition[space_node_index]

warp/fem/space/quadmesh_function_space.py

@@ -55,13 +55,16 @@ class QuadmeshSpaceTopology(SpaceTopology):
     @cache.cached_arg_value
     def topo_arg_value(self, device):
         arg = Quadmesh2DTopologyArg()
+        self.fill_topo_arg(arg, device)
+        return arg
+
+    def fill_topo_arg(self, arg: Quadmesh2DTopologyArg, device):
         arg.quad_edge_indices = self._quad_edge_indices.to(device)
         arg.edge_vertex_indices = self._mesh.edge_vertex_indices.to(device)
 
         arg.vertex_count = self._mesh.vertex_count()
         arg.edge_count = self._mesh.side_count()
         arg.cell_count = self._mesh.cell_count()
-        return arg
 
     def _compute_quad_edge_indices(self):
         self._quad_edge_indices = wp.empty(

warp/fem/space/restriction.py

@@ -144,13 +144,16 @@ class SpaceRestriction:
         dof_indices_in_element: wp.array(dtype=int)
 
     @cached_arg_value
-    def node_arg(self, device):
+    def node_arg_value(self, device):
         arg = SpaceRestriction.NodeArg()
+        self.fill_node_arg(arg, device)
+        return arg
+
+    def fill_node_arg(self, arg: NodeArg, device):
         arg.dof_element_offsets = self._dof_partition_element_offsets.array.to(device)
         arg.dof_element_indices = self._dof_element_indices.array.to(device)
         arg.dof_partition_indices = self._dof_partition_indices.array.to(device)
         arg.dof_indices_in_element = self._dof_indices_in_element.array.to(device)
-        return arg
 
     @wp.func
     def node_partition_index(args: NodeArg, restriction_node_index: int):

warp/fem/space/shape/__init__.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 from enum import Enum
 from typing import Optional
 
@@ -67,8 +68,9 @@ class ElementBasis(Enum):
     """Raviart-Thomas H(div) shape functions. Should be used with contravariant function space."""
 
 
+@functools.lru_cache(maxsize=None)
 def get_shape_function(
-    element: _element.Element,
+    element_class: type,
     space_dimension: int,
     degree: int,
     element_basis: ElementBasis,
@@ -78,7 +80,7 @@ def get_shape_function(
     Equips a reference element with a shape function basis.
 
     Args:
-        element: the reference element on which to build the shape function
+        element_class: the type of reference element on which to build the shape function
         space_dimension: the dimension of the embedding space
         degree: polynomial degree of the per-element shape functions
         element_basis: type of basis function for the individual elements
@@ -89,12 +91,12 @@ def get_shape_function(
     """
 
     if degree == 0:
-        return ConstantShapeFunction(element, space_dimension)
+        return ConstantShapeFunction(element_class(), space_dimension)
 
     if family is None:
         family = Polynomial.LOBATTO_GAUSS_LEGENDRE
 
-    if isinstance(element, _element.Square):
+    if issubclass(element_class, _element.Square):
         if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
             return SquareNedelecFirstKindShapeFunctions(degree=degree)
         if element_basis == ElementBasis.RAVIART_THOMAS:
@@ -105,7 +107,7 @@ def get_shape_function(
             return SquareSerendipityShapeFunctions(degree=degree, family=family)
 
         return SquareBipolynomialShapeFunctions(degree=degree, family=family)
-    if isinstance(element, _element.Triangle):
+    if issubclass(element_class, _element.Triangle):
         if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
             return TriangleNedelecFirstKindShapeFunctions(degree=degree)
         if element_basis == ElementBasis.RAVIART_THOMAS:
@@ -117,7 +119,7 @@ def get_shape_function(
 
         return TrianglePolynomialShapeFunctions(degree=degree)
 
-    if isinstance(element, _element.Cube):
+    if issubclass(element_class, _element.Cube):
         if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
             return CubeNedelecFirstKindShapeFunctions(degree=degree)
         if element_basis == ElementBasis.RAVIART_THOMAS:
@@ -128,7 +130,7 @@ def get_shape_function(
             return CubeSerendipityShapeFunctions(degree=degree, family=family)
 
         return CubeTripolynomialShapeFunctions(degree=degree, family=family)
-    if isinstance(element, _element.Tetrahedron):
+    if issubclass(element_class, _element.Tetrahedron):
         if element_basis == ElementBasis.NEDELEC_FIRST_KIND:
             return TetrahedronNedelecFirstKindShapeFunctions(degree=degree)
         if element_basis == ElementBasis.RAVIART_THOMAS:
@@ -140,4 +142,4 @@ def get_shape_function(
 
         return TetrahedronPolynomialShapeFunctions(degree=degree)
 
-    return NotImplementedError("Unrecognized element type")
+    raise NotImplementedError(f"Unrecognized element type {element_class.__name__}")

warp/fem/space/tetmesh_function_space.py

@@ -75,6 +75,10 @@ class TetmeshSpaceTopology(SpaceTopology):
     @cache.cached_arg_value
     def topo_arg_value(self, device):
         arg = TetmeshTopologyArg()
+        self.fill_topo_arg(arg, device)
+        return arg
+
+    def fill_topo_arg(self, arg: TetmeshTopologyArg, device):
         arg.tet_face_indices = self._tet_face_indices.to(device)
         arg.tet_edge_indices = self._tet_edge_indices.to(device)
         arg.face_vertex_indices = self._mesh.face_vertex_indices.to(device)
@@ -83,7 +87,6 @@ class TetmeshSpaceTopology(SpaceTopology):
         arg.vertex_count = self._mesh.vertex_count()
         arg.face_count = self._mesh.side_count()
         arg.edge_count = self._edge_count
-        return arg
 
     def _compute_tet_face_indices(self):
         self._tet_face_indices = wp.empty(