warp-lang 1.6.2__py3-none-macosx_10_13_universal2.whl → 1.7.1__py3-none-macosx_10_13_universal2.whl

warp/fem/quadrature/pic_quadrature.py

@@ -18,7 +18,7 @@ from typing import Any, Optional, Tuple, Union
 import warp as wp
 from warp.fem.cache import TemporaryStore, borrow_temporary, cached_arg_value, dynamic_kernel
 from warp.fem.domain import GeometryDomain
-from warp.fem.types import Coords, ElementIndex, make_free_sample
+from warp.fem.types import NULL_ELEMENT_INDEX, Coords, ElementIndex, make_free_sample
 from warp.fem.utils import compress_node_indices
 
 from .quadrature import Quadrature
@@ -68,10 +68,10 @@ class PicQuadrature(Quadrature):
     def domain(self, domain: GeometryDomain):
         # Allow changing the quadrature domain as long as underlying geometry and element kind are the same
         if self.domain is not None and (
-            domain.geometry != self.domain.geometry or domain.element_kind != self.domain.element_kind
+            domain.element_kind != self.domain.element_kind or domain.geometry.base != self.domain.geometry.base
         ):
             raise RuntimeError(
-                "Cannot change the domain to that of a different Geometry and/or using different element kinds."
+                "The new domain must use the same base geometry and kind of elements as the current one."
             )
 
         self._domain = domain
@@ -89,11 +89,11 @@ class PicQuadrature(Quadrature):
         arg.cell_particle_offsets = self._cell_particle_offsets.array.to(device)
         arg.cell_particle_indices = self._cell_particle_indices.array.to(device)
         arg.particle_fraction = self._particle_fraction.to(device)
-        arg.particle_coords = self._particle_coords.to(device)
+        arg.particle_coords = self.particle_coords.to(device)
         return arg
 
     def total_point_count(self):
-        return self._particle_coords.shape[0]
+        return self.particle_coords.shape[0]
 
     def active_cell_count(self):
         """Number of cells containing at least one particle"""
@@ -136,6 +136,12 @@ class PicQuadrature(Quadrature):
         particle_index = qp_arg.cell_particle_indices[qp_arg.cell_particle_offsets[element_index] + index]
         return particle_index
 
+    @wp.func
+    def point_evaluation_index(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, index: int
+    ):
+        return qp_arg.cell_particle_offsets[element_index] + index
+
     def fill_element_mask(self, mask: "wp.array(dtype=int)"):
         """Fills a mask array such that all non-empty elements are set to 1, all empty elements to zero.
 
@@ -156,7 +162,7 @@ class PicQuadrature(Quadrature):
         element_mask: wp.array(dtype=int),
     ):
         i = wp.tid()
-        element_mask[i] = wp.select(element_particle_offsets[i] == element_particle_offsets[i + 1], 1, 0)
+        element_mask[i] = wp.where(element_particle_offsets[i] == element_particle_offsets[i + 1], 0, 1)
 
     @wp.kernel
     def _compute_uniform_fraction(
@@ -167,9 +173,11 @@ class PicQuadrature(Quadrature):
         p = wp.tid()
 
         cell = cell_index[p]
-        cell_particle_count = cell_particle_offsets[cell + 1] - cell_particle_offsets[cell]
-
-        cell_fraction[p] = 1.0 / float(cell_particle_count)
+        if cell == NULL_ELEMENT_INDEX:
+            cell_fraction[p] = 0.0
+        else:
+            cell_particle_count = cell_particle_offsets[cell + 1] - cell_particle_offsets[cell]
+            cell_fraction[p] = 1.0 / float(cell_particle_count)
 
     def _bin_particles(self, positions, measures, temporary_store: TemporaryStore):
         if wp.types.is_array(positions):
@@ -189,13 +197,13 @@ class PicQuadrature(Quadrature):
 
             device = positions.device
 
-            cell_index_temp = borrow_temporary(temporary_store, shape=positions.shape, dtype=int, device=device)
-            cell_index = cell_index_temp.array
+            self._cell_index_temp = borrow_temporary(temporary_store, shape=positions.shape, dtype=int, device=device)
+            self.cell_indices = self._cell_index_temp.array
 
             self._particle_coords_temp = borrow_temporary(
                 temporary_store, shape=positions.shape, dtype=Coords, device=device, requires_grad=self._requires_grad
             )
-            self._particle_coords = self._particle_coords_temp.array
+            self.particle_coords = self._particle_coords_temp.array
 
             wp.launch(
                 dim=positions.shape[0],
@@ -203,25 +211,28 @@ class PicQuadrature(Quadrature):
                 inputs=[
                     self.domain.element_arg_value(device),
                     positions,
-                    cell_index,
-                    self._particle_coords,
+                    self.cell_indices,
+                    self.particle_coords,
                 ],
                 device=device,
             )
 
         else:
-            cell_index, self._particle_coords = positions
-            if cell_index.shape != self._particle_coords.shape:
+            self.cell_indices, self.particle_coords = positions
+            if self.cell_indices.shape != self.particle_coords.shape:
                 raise ValueError("Cell index and coordinates arrays must have the same shape")
 
-            cell_index_temp = None
+            self._cell_index_temp = None
             self._particle_coords_temp = None
 
         self._cell_particle_offsets, self._cell_particle_indices, self._cell_count, _ = compress_node_indices(
-            self.domain.geometry_element_count(), cell_index, return_unique_nodes=True, temporary_store=temporary_store
+            self.domain.geometry_element_count(),
+            self.cell_indices,
+            return_unique_nodes=True,
+            temporary_store=temporary_store,
         )
 
-        self._compute_fraction(cell_index, measures, temporary_store)
+        self._compute_fraction(self.cell_indices, measures, temporary_store)
 
     def _compute_fraction(self, cell_index, measures, temporary_store: TemporaryStore):
         device = cell_index.device
@@ -260,9 +271,13 @@ class PicQuadrature(Quadrature):
                 cell_fraction: wp.array(dtype=float),
             ):
                 p = wp.tid()
-                sample = make_free_sample(cell_index[p], cell_coords[p])
 
-                cell_fraction[p] = measures[p] / self.domain.element_measure(cell_arg_value, sample)
+                cell = cell_index[p]
+                if cell == NULL_ELEMENT_INDEX:
+                    cell_fraction[p] = 0.0
+                else:
+                    sample = make_free_sample(cell_index[p], cell_coords[p])
+                    cell_fraction[p] = measures[p] / self.domain.element_measure(cell_arg_value, sample)
 
             wp.launch(
                 dim=measures.shape[0],
@@ -271,7 +286,7 @@ class PicQuadrature(Quadrature):
                     self.domain.element_arg_value(device),
                     measures,
                     cell_index,
-                    self._particle_coords,
+                    self.particle_coords,
                     self._particle_fraction,
                 ],
                 device=device,

warp/fem/quadrature/quadrature.py

@@ -13,17 +13,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
+from typing import Any, Optional
 
 import warp as wp
-from warp.fem import cache, domain
+from warp.fem import cache
+from warp.fem.domain import GeometryDomain
 from warp.fem.geometry import Element
 from warp.fem.space import FunctionSpace
-from warp.fem.types import Coords, ElementIndex
+from warp.fem.types import NULL_ELEMENT_INDEX, Coords, ElementIndex, QuadraturePointIndex
 
 from ..polynomial import Polynomial
 
 
+@wp.struct
+class QuadraturePointElementIndex:
+    domain_element_index: ElementIndex
+    qp_index_in_element: int
+
+
 class Quadrature:
     """Interface class for quadrature rules"""
 
@@ -33,7 +40,7 @@ class Quadrature:
 
         pass
 
-    def __init__(self, domain: domain.GeometryDomain):
+    def __init__(self, domain: GeometryDomain):
         self._domain = domain
 
     @property
@@ -45,52 +52,197 @@ class Quadrature:
         """
         Value of the argument to be passed to device
         """
-        arg = RegularQuadrature.Arg()
+        arg = Quadrature.Arg()
         return arg
 
     def total_point_count(self):
-        """Total number of quadrature points over the domain"""
+        """Number of unique quadrature points that can be indexed by this rule.
+        Returns a number such that `point_index()` is always smaller than this number.
+        """
         raise NotImplementedError()
 
+    def evaluation_point_count(self):
+        """Number of quadrature points that needs to be evaluated, mostly for internal purposes.
+        If the indexing scheme is sparse, or if a quadrature point is shared among multiple elements
+        (e.g, nodal quadrature), `evaluation_point_count` may be different than `total_point_count()`.
+        Returns a number such that `evaluation_point_index()` is always smaller than this number.
+        """
+        return self.total_point_count()
+
     def max_points_per_element(self):
         """Maximum number of points per element if known, or ``None`` otherwise"""
         return None
 
     @staticmethod
-    def point_count(elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex):
+    def point_count(
+        elt_arg: "GeometryDomain.ElementArg",
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        geo_element_index: ElementIndex,
+    ):
         """Number of quadrature points for a given element"""
         raise NotImplementedError()
 
     @staticmethod
     def point_coords(
-        elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex, qp_index: int
+        elt_arg: "GeometryDomain.ElementArg",
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        geo_element_index: ElementIndex,
+        element_qp_index: int,
     ):
         """Coordinates in element of the element's qp_index'th quadrature point"""
         raise NotImplementedError()
 
     @staticmethod
     def point_weight(
-        elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex, qp_index: int
+        elt_arg: "GeometryDomain.ElementArg",
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        geo_element_index: ElementIndex,
+        element_qp_index: int,
     ):
         """Weight of the element's qp_index'th quadrature point"""
         raise NotImplementedError()
 
     @staticmethod
     def point_index(
-        elt_arg: "domain.GeometryDomain.ElementArg",
+        elt_arg: "GeometryDomain.ElementArg",
+        qp_arg: Arg,
+        domain_element_index: ElementIndex,
+        geo_element_index: ElementIndex,
+        element_qp_index: int,
+    ):
+        """
+        Global index of the element's qp_index'th  quadrature point.
+        May be shared among elements.
+        This is what determines `qp_index` in integrands' `Sample` arguments.
+        """
+        raise NotImplementedError()
+
+    @staticmethod
+    def point_evaluation_index(
+        elt_arg: "GeometryDomain.ElementArg",
         qp_arg: Arg,
         domain_element_index: ElementIndex,
         geo_element_index: ElementIndex,
         element_qp_index: int,
     ):
-        """Global index of the element's qp_index'th  quadrature point"""
+        """Quadrature point index according to evaluation order.
+        Quadrature points for distinct elements must have different evaluation indices.
+        Mostly for internal/parallelization purposes.
+        """
         raise NotImplementedError()
 
     def __str__(self) -> str:
         return self.name
 
+    # By default cache the mapping from evaluation point indices to domain elements
+
+    ElementIndexArg = wp.array(dtype=QuadraturePointElementIndex)
+
+    @cache.cached_arg_value
+    def element_index_arg_value(self, device):
+        """Builds a map from quadrature point evaluation indices to their index in the element to which they belong"""
+
+        @cache.dynamic_kernel(f"{self.name}{self.domain.name}")
+        def quadrature_point_element_indices(
+            qp_arg: self.Arg,
+            domain_arg: self.domain.ElementArg,
+            domain_index_arg: self.domain.ElementIndexArg,
+            result: wp.array(dtype=QuadraturePointElementIndex),
+        ):
+            domain_element_index = wp.tid()
+            element_index = self.domain.element_index(domain_index_arg, domain_element_index)
+
+            qp_point_count = self.point_count(domain_arg, qp_arg, domain_element_index, element_index)
+            for k in range(qp_point_count):
+                qp_eval_index = self.point_evaluation_index(domain_arg, qp_arg, domain_element_index, element_index, k)
+                result[qp_eval_index] = QuadraturePointElementIndex(domain_element_index, k)
+
+        null_qp_index = QuadraturePointElementIndex()
+        null_qp_index.domain_element_index = NULL_ELEMENT_INDEX
+        result = wp.full(
+            value=null_qp_index,
+            shape=(self.evaluation_point_count()),
+            dtype=QuadraturePointElementIndex,
+            device=device,
+        )
+        wp.launch(
+            quadrature_point_element_indices,
+            device=result.device,
+            dim=self.domain.element_count(),
+            inputs=[
+                self.arg_value(result.device),
+                self.domain.element_arg_value(result.device),
+                self.domain.element_index_arg_value(result.device),
+                result,
+            ],
+        )
+
+        return result
+
+    @wp.func
+    def evaluation_point_element_index(
+        element_index_arg: wp.array(dtype=QuadraturePointElementIndex),
+        qp_eval_index: QuadraturePointIndex,
+    ):
+        """Maps from quadrature point evaluation indices to their index in the element to which they belong
+        If the quadrature point does not exist, should return NULL_ELEMENT_INDEX as the domain element index
+        """
+
+        element_index = element_index_arg[qp_eval_index]
+        return element_index.domain_element_index, element_index.qp_index_in_element
+
+
+class _QuadratureWithRegularEvaluationPoints(Quadrature):
+    """Helper subclass for quadrature formulas which use a uniform number of
+    evaluations points per element. Avoids building explicit mapping"""
+
+    def __init__(self, domain: GeometryDomain, N: int):
+        super().__init__(domain)
+        self._EVALUATION_POINTS_PER_ELEMENT = N
+
+        self.point_evaluation_index = self._make_regular_point_evaluation_index()
+        self.evaluation_point_element_index = self._make_regular_evaluation_point_element_index()
 
-class RegularQuadrature(Quadrature):
+    ElementIndexArg = Quadrature.Arg
+    element_index_arg_value = Quadrature.arg_value
+
+    def evaluation_point_count(self):
+        return self.domain.element_count() * self._EVALUATION_POINTS_PER_ELEMENT
+
+    def _make_regular_point_evaluation_index(self):
+        N = self._EVALUATION_POINTS_PER_ELEMENT
+
+        @cache.dynamic_func(suffix=f"{self.name}")
+        def evaluation_point_index(
+            elt_arg: self.domain.ElementArg,
+            qp_arg: self.Arg,
+            domain_element_index: ElementIndex,
+            element_index: ElementIndex,
+            qp_index: int,
+        ):
+            return N * domain_element_index + qp_index
+
+        return evaluation_point_index
+
+    def _make_regular_evaluation_point_element_index(self):
+        N = self._EVALUATION_POINTS_PER_ELEMENT
+
+        @cache.dynamic_func(suffix=f"{N}")
+        def quadrature_evaluation_point_element_index(
+            qp_arg: Quadrature.Arg,
+            qp_index: QuadraturePointIndex,
+        ):
+            domain_element_index = qp_index // N
+            index_in_element = qp_index - domain_element_index * N
+            return domain_element_index, index_in_element
+
+        return quadrature_evaluation_point_element_index
+
+
+class RegularQuadrature(_QuadratureWithRegularEvaluationPoints):
     """Regular quadrature formula, using a constant set of quadrature points per element"""
 
     @wp.struct
@@ -127,16 +279,15 @@ class RegularQuadrature(Quadrature):
 
     def __init__(
         self,
-        domain: domain.GeometryDomain,
+        domain: GeometryDomain,
         order: int,
         family: Polynomial = None,
     ):
-        super().__init__(domain)
-
+        self._formula = RegularQuadrature.CachedFormula.get(domain.reference_element(), order, family)
         self.family = family
         self.order = order
 
-        self._formula = RegularQuadrature.CachedFormula.get(domain.reference_element(), order, family)
+        super().__init__(domain, self._formula.count)
 
         self.point_count = self._make_point_count()
         self.point_index = self._make_point_index()
@@ -227,17 +378,18 @@ class NodalQuadrature(Quadrature):
     any assumption about orthogonality of shape functions, and is thus safe to use for arbitrary integrands.
     """
 
-    def __init__(self, domain: domain.GeometryDomain, space: FunctionSpace):
-        super().__init__(domain)
-
+    def __init__(self, domain: Optional[GeometryDomain], space: FunctionSpace):
         self._space = space
 
+        super().__init__(domain)
+
         self.Arg = self._make_arg()
 
         self.point_count = self._make_point_count()
         self.point_index = self._make_point_index()
         self.point_coords = self._make_point_coords()
         self.point_weight = self._make_point_weight()
+        self.point_evaluation_index = self._make_point_evaluation_index()
 
     @property
     def name(self):
@@ -315,8 +467,26 @@ class NodalQuadrature(Quadrature):
 
         return point_index
 
+    def evaluation_point_count(self):
+        return self.domain.element_count() * self._space.topology.MAX_NODES_PER_ELEMENT
 
-class ExplicitQuadrature(Quadrature):
+    def _make_point_evaluation_index(self):
+        N = self._space.topology.MAX_NODES_PER_ELEMENT
+
+        @cache.dynamic_func(suffix=self.name)
+        def evaluation_point_index(
+            elt_arg: self.domain.ElementArg,
+            qp_arg: self.Arg,
+            domain_element_index: ElementIndex,
+            element_index: ElementIndex,
+            qp_index: int,
+        ):
+            return N * domain_element_index + qp_index
+
+        return evaluation_point_index
+
+
+class ExplicitQuadrature(_QuadratureWithRegularEvaluationPoints):
     """Quadrature using explicit per-cell points and weights.
 
     The number of quadrature points per cell is assumed to be constant and deduced from the shape of the points and weights arrays.
@@ -336,11 +506,7 @@ class ExplicitQuadrature(Quadrature):
         points: wp.array2d(dtype=Coords)
         weights: wp.array2d(dtype=float)
 
-    def __init__(
-        self, domain: domain.GeometryDomain, points: "wp.array2d(dtype=Coords)", weights: "wp.array2d(dtype=float)"
-    ):
-        super().__init__(domain)
-
+    def __init__(self, domain: GeometryDomain, points: "wp.array2d(dtype=Coords)", weights: "wp.array2d(dtype=float)"):
         if points.shape != weights.shape:
             raise ValueError("Points and weights arrays must have the same shape")
 
@@ -358,7 +524,10 @@ class ExplicitQuadrature(Quadrature):
             )
 
         self._points_per_cell = points.shape[1]
+
         self._whole_geo = points.shape[0] == domain.geometry_element_count()
+
+        super().__init__(domain, self._points_per_cell)
         self._points = points
         self._weights = weights
 

warp/fem/space/__init__.py

@@ -112,7 +112,7 @@ def make_polynomial_basis_space(
         the constructed basis space
     """
 
-    base_geo = geo.base if isinstance(geo, _geometry.DeformedGeometry) else geo
+    base_geo = geo.base
 
     if element_basis is None:
         element_basis = ElementBasis.LAGRANGE

warp/fem/space/basis_function_space.py

@@ -19,7 +19,7 @@ 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.types import Coords, ElementIndex, make_free_sample
+from warp.fem.types import NULL_QP_INDEX, Coords, ElementIndex, make_free_sample
 
 from .basis_space import BasisSpace
 from .dof_mapper import DofMapper, IdentityMapper
@@ -305,7 +305,9 @@ class VectorValuedFunctionSpace(FunctionSpace):
             space_value: self.dtype,
         ):
             coords = self.node_coords_in_element(elt_arg, space_arg, element_index, node_index_in_elt)
-            weight = self.element_inner_weight(elt_arg, space_arg, element_index, coords, node_index_in_elt)
+            weight = self.element_inner_weight(
+                elt_arg, space_arg, element_index, coords, node_index_in_elt, NULL_QP_INDEX
+            )
             local_value_map = self.local_value_map_inner(elt_arg, element_index, coords)
 
             unit_value = local_value_map * weight

warp/fem/space/basis_space.py

@@ -21,7 +21,14 @@ import warp as wp
 from warp.fem import cache
 from warp.fem.geometry import Geometry
 from warp.fem.quadrature import Quadrature
-from warp.fem.types import NULL_ELEMENT_INDEX, Coords, ElementIndex, make_free_sample
+from warp.fem.types import (
+    NULL_ELEMENT_INDEX,
+    NULL_QP_INDEX,
+    Coords,
+    ElementIndex,
+    QuadraturePointIndex,
+    make_free_sample,
+)
 
 from .shape import ShapeFunction
 from .topology import RegularDiscontinuousSpaceTopology, SpaceTopology
@@ -235,6 +242,7 @@ class ShapeBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             if wp.static(self.value == ShapeFunction.Value.Scalar):
                 return shape_element_inner_weight(coords, node_index_in_elt)
@@ -254,6 +262,7 @@ class ShapeBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             if wp.static(self.value == ShapeFunction.Value.Scalar):
                 return shape_element_inner_weight_gradient(coords, node_index_in_elt)
@@ -373,6 +382,7 @@ class TraceBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             cell_index, index_in_cell = self.topology.inner_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
@@ -381,13 +391,7 @@ class TraceBasisSpace(BasisSpace):
             cell_coords = self.geometry.side_inner_cell_coords(geo_side_arg, element_index, coords)
 
             geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
-            return cell_inner_weight(
-                geo_cell_arg,
-                basis_arg,
-                cell_index,
-                cell_coords,
-                index_in_cell,
-            )
+            return cell_inner_weight(geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell, NULL_QP_INDEX)
 
         return trace_element_inner_weight
 
@@ -401,6 +405,7 @@ class TraceBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             cell_index, index_in_cell = self.topology.outer_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
@@ -409,13 +414,7 @@ class TraceBasisSpace(BasisSpace):
             cell_coords = self.geometry.side_outer_cell_coords(geo_side_arg, element_index, coords)
 
             geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
-            return cell_outer_weight(
-                geo_cell_arg,
-                basis_arg,
-                cell_index,
-                cell_coords,
-                index_in_cell,
-            )
+            return cell_outer_weight(geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell, NULL_QP_INDEX)
 
         return trace_element_outer_weight
 
@@ -429,6 +428,7 @@ class TraceBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             cell_index, index_in_cell = self.topology.inner_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
@@ -436,7 +436,9 @@ class TraceBasisSpace(BasisSpace):
 
             cell_coords = self.geometry.side_inner_cell_coords(geo_side_arg, element_index, coords)
             geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
-            return cell_inner_weight_gradient(geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell)
+            return cell_inner_weight_gradient(
+                geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell, NULL_QP_INDEX
+            )
 
         return trace_element_inner_weight_gradient
 
@@ -450,6 +452,7 @@ class TraceBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             cell_index, index_in_cell = self.topology.outer_cell_index(geo_side_arg, element_index, node_index_in_elt)
             if cell_index == NULL_ELEMENT_INDEX:
@@ -457,7 +460,9 @@ class TraceBasisSpace(BasisSpace):
 
             cell_coords = self.geometry.side_outer_cell_coords(geo_side_arg, element_index, coords)
             geo_cell_arg = self.geometry.side_to_cell_arg(geo_side_arg)
-            return cell_outer_weight_gradient(geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell)
+            return cell_outer_weight_gradient(
+                geo_cell_arg, basis_arg, cell_index, cell_coords, index_in_cell, NULL_QP_INDEX
+            )
 
         return trace_element_outer_weight_gradient
 
@@ -624,11 +629,12 @@ class PointBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             qp_coord = self._quadrature.point_coords(
                 elt_arg, basis_arg, element_index, element_index, node_index_in_elt
             )
-            return wp.select(wp.length_sq(coords - qp_coord) < _DIRAC_INTEGRATION_RADIUS, 0.0, 1.0)
+            return wp.where(wp.length_sq(coords - qp_coord) < _DIRAC_INTEGRATION_RADIUS, 1.0, 0.0)
 
         return element_inner_weight
 
@@ -642,6 +648,7 @@ class PointBasisSpace(BasisSpace):
             element_index: ElementIndex,
             coords: Coords,
             node_index_in_elt: int,
+            qp_index: QuadraturePointIndex,
         ):
             return gradient_vec(0.0)
 

warp/fem/space/hexmesh_function_space.py

@@ -237,13 +237,13 @@ class HexmeshSpaceTopology(SpaceTopology):
                     hex_edge = _CUBE_TO_HEX_EDGE[type_instance]
                     v0 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 0]]
                     v1 = geo_arg.hex_vertex_indices[element_index, EDGE_VERTEX_INDICES[hex_edge, 1]]
-                    return wp.select(v0 > v1, 1.0, -1.0)
+                    return wp.where(v0 > v1, -1.0, 1.0)
 
             if wp.static(FACE_NODE_COUNT > 0):
                 if node_type == CubeShapeFunction.FACE:
                     face_index_and_ori = topo_arg.hex_face_indices[element_index, type_instance]
                     flip = face_index_and_ori[1] & 1
-                    return wp.select(flip == 0, -1.0, 1.0)
+                    return wp.where(flip == 0, 1.0, -1.0)
 
             return 1.0