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

warp/fem/space/basis_space.py

@@ -4,9 +4,10 @@ import warp as wp
 
 from warp.fem.types import ElementIndex, Coords, make_free_sample
 from warp.fem.geometry import Geometry
-from warp.fem import cache, utils
+from warp.fem.quadrature import Quadrature
+from warp.fem import cache
 
-from .topology import SpaceTopology
+from .topology import SpaceTopology, DiscontinuousSpaceTopology
 from .shape import ShapeFunction
 
 
@@ -25,19 +26,10 @@ class BasisSpace:
 
         pass
 
-    def __init__(self, topology: SpaceTopology, shape: ShapeFunction):
+    def __init__(self, topology: SpaceTopology):
         self._topology = topology
-        self._shape = shape
 
         self.NODES_PER_ELEMENT = self._topology.NODES_PER_ELEMENT
-        self.ORDER = self._shape.ORDER
-
-        if hasattr(shape, "element_node_triangulation"):
-            self.node_triangulation = self._node_triangulation
-        if hasattr(shape, "element_node_tets"):
-            self.node_tets = self._node_tets
-        if hasattr(shape, "element_node_hexes"):
-            self.node_hexes = self._node_hexes
 
     @property
     def topology(self) -> SpaceTopology:
@@ -49,19 +41,10 @@ class BasisSpace:
         """Underlying geometry of the basis space"""
         return self._topology.geometry
 
-    @property
-    def shape(self) -> ShapeFunction:
-        """Shape functions used for defining individual element basis"""
-        return self._shape
-
     def basis_arg_value(self, device) -> "BasisArg":
         """Value for the argument structure to be passed to device functions"""
         return BasisSpace.BasisArg()
 
-    @property
-    def name(self):
-        return f"{self.topology.name}_{self._shape.name}"
-
     # Helpers for generating node positions
 
     def node_positions(self, out: Optional[wp.array] = None) -> wp.array:
@@ -117,6 +100,56 @@ class BasisSpace:
 
         return node_positions
 
+    def make_node_coords_in_element(self):
+        raise NotImplementedError()
+
+    def make_node_quadrature_weight(self):
+        raise NotImplementedError()
+
+    def make_element_inner_weight(self):
+        raise NotImplementedError()
+
+    def make_element_outer_weight(self):
+        return self.make_element_inner_weight()
+
+    def make_element_inner_weight_gradient(self):
+        raise NotImplementedError()
+
+    def make_element_outer_weight_gradient(self):
+        return self.make_element_inner_weight_gradient()
+
+    def make_trace_node_quadrature_weight(self):
+        raise NotImplementedError()
+
+    def trace(self) -> "TraceBasisSpace":
+        return TraceBasisSpace(self)
+
+
+class ShapeBasisSpace(BasisSpace):
+    """Base class for defining shape-function-based basis spaces."""
+
+    def __init__(self, topology: SpaceTopology, shape: ShapeFunction):
+        super().__init__(topology)
+        self._shape = shape
+
+        self.ORDER = self._shape.ORDER
+
+        if hasattr(shape, "element_node_triangulation"):
+            self.node_triangulation = self._node_triangulation
+        if hasattr(shape, "element_node_tets"):
+            self.node_tets = self._node_tets
+        if hasattr(shape, "element_node_hexes"):
+            self.node_hexes = self._node_hexes
+
+    @property
+    def shape(self) -> ShapeFunction:
+        """Shape functions used for defining individual element basis"""
+        return self._shape
+
+    @property
+    def name(self):
+        return f"{self.topology.name}_{self._shape.name}"
+
     def make_node_coords_in_element(self):
         shape_node_coords_in_element = self._shape.make_node_coords_in_element()
 
@@ -156,14 +189,10 @@ class BasisSpace:
             coords: Coords,
             node_index_in_elt: int,
         ):
-            a = self.BasisArg()
             return shape_element_inner_weight(coords, node_index_in_elt)
 
         return element_inner_weight
 
-    def make_element_outer_weight(self):
-        return self.make_element_inner_weight()
-
     def make_element_inner_weight_gradient(self):
         shape_element_inner_weight_gradient = self._shape.make_element_inner_weight_gradient()
 
@@ -179,11 +208,22 @@ class BasisSpace:
 
         return element_inner_weight_gradient
 
-    def make_element_outer_weight_gradient(self):
-        return self.make_element_inner_weight_gradient()
+    def make_trace_node_quadrature_weight(self, trace_basis):
+        shape_trace_node_quadrature_weight = self._shape.make_trace_node_quadrature_weight()
 
-    def trace(self) -> "TraceBasisSpace":
-        return TraceBasisSpace(self)
+        @cache.dynamic_func(suffix=self.name)
+        def trace_node_quadrature_weight(
+            geo_side_arg: trace_basis.geometry.SideArg,
+            basis_arg: trace_basis.BasisArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            neighbour_elem, index_in_neighbour = trace_basis.topology.neighbor_cell_index(
+                geo_side_arg, element_index, node_index_in_elt
+            )
+            return shape_trace_node_quadrature_weight(index_in_neighbour)
+
+        return trace_node_quadrature_weight
 
     def _node_triangulation(self):
         element_node_indices = self._topology.element_node_indices().numpy()
@@ -211,7 +251,9 @@ class TraceBasisSpace(BasisSpace):
     """Auto-generated trace space evaluating the cell-defined basis on the geometry sides"""
 
     def __init__(self, basis: BasisSpace):
-        super().__init__(basis.topology.trace(), basis.shape)
+        super().__init__(basis.topology.trace())
+
+        self.ORDER = basis.ORDER
 
         self._basis = basis
         self.BasisArg = self._basis.BasisArg
@@ -247,21 +289,7 @@ class TraceBasisSpace(BasisSpace):
         return trace_node_coords_in_element
 
     def make_node_quadrature_weight(self):
-        shape_trace_node_quadrature_weight = self._shape.make_trace_node_quadrature_weight()
-
-        @cache.dynamic_func(suffix=self._basis.name)
-        def trace_node_quadrature_weight(
-            geo_side_arg: self.geometry.SideArg,
-            basis_arg: self.BasisArg,
-            element_index: ElementIndex,
-            node_index_in_elt: int,
-        ):
-            neighbour_elem, index_in_neighbour = self.topology.neighbor_cell_index(
-                geo_side_arg, element_index, node_index_in_elt
-            )
-            return shape_trace_node_quadrature_weight(index_in_neighbour)
-
-        return trace_node_quadrature_weight
+        return self._basis.make_trace_node_quadrature_weight(self)
 
     def make_element_inner_weight(self):
         cell_inner_weight = self._basis.make_element_inner_weight()
@@ -337,9 +365,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_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)
 
         return trace_element_inner_weight_gradient
 
@@ -361,11 +387,103 @@ 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)
 
         return trace_element_outer_weight_gradient
 
     def __eq__(self, other: "TraceBasisSpace") -> bool:
         return self._topo == other._topo
+
+
+class PointBasisSpace(BasisSpace):
+    """An unstructured :class:`BasisSpace` that is non-zero at a finite set of points only.
+
+    The node locations and nodal quadrature weights are defined by a :class:`Quadrature` formula.
+    """
+
+    def __init__(self, quadrature: Quadrature):
+        self._quadrature = quadrature
+
+        if quadrature.points_per_element() is None:
+            raise NotImplementedError("Varying number of points per element is not supported yet")
+
+        topology = DiscontinuousSpaceTopology(
+            geometry=quadrature.domain.geometry, nodes_per_element=quadrature.points_per_element()
+        )
+        super().__init__(topology)
+
+        self.BasisArg = quadrature.Arg
+        self.basis_arg_value = quadrature.arg_value
+        self.ORDER = 0
+
+        self.make_element_outer_weight = self.make_element_inner_weight
+        self.make_element_outer_weight_gradient = self.make_element_outer_weight_gradient
+
+    @property
+    def name(self):
+        return f"{self._quadrature.name}_Point"
+
+    def make_node_coords_in_element(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_coords_in_element(
+            elt_arg: self._quadrature.domain.ElementArg,
+            basis_arg: self.BasisArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            return self._quadrature.point_coords(elt_arg, basis_arg, element_index, node_index_in_elt)
+
+        return node_coords_in_element
+
+    def make_node_quadrature_weight(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_quadrature_weight(
+            elt_arg: self._quadrature.domain.ElementArg,
+            basis_arg: self.BasisArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            return self._quadrature.point_weight(elt_arg, basis_arg, element_index, node_index_in_elt)
+
+        return node_quadrature_weight
+
+    def make_element_inner_weight(self):
+        @cache.dynamic_func(suffix=self.name)
+        def element_inner_weight(
+            elt_arg: self._quadrature.domain.ElementArg,
+            basis_arg: self.BasisArg,
+            element_index: ElementIndex,
+            coords: Coords,
+            node_index_in_elt: int,
+        ):
+            qp_coord = self._quadrature.point_coords(elt_arg, basis_arg, element_index, node_index_in_elt)
+            return wp.select(wp.length_sq(coords - qp_coord) < 0.001, 0.0, 1.0)
+
+        return element_inner_weight
+
+    def make_element_inner_weight_gradient(self):
+        gradient_vec = cache.cached_vec_type(length=self.geometry.dimension, dtype=float)
+
+        @cache.dynamic_func(suffix=self.name)
+        def element_inner_weight_gradient(
+            elt_arg: self._quadrature.domain.ElementArg,
+            basis_arg: self.BasisArg,
+            element_index: ElementIndex,
+            coords: Coords,
+            node_index_in_elt: int,
+        ):
+            return gradient_vec(0.0)
+
+        return element_inner_weight_gradient
+
+    def make_trace_node_quadrature_weight(self, trace_basis):
+        @cache.dynamic_func(suffix=self.name)
+        def trace_node_quadrature_weight(
+            elt_arg: trace_basis.geometry.SideArg,
+            basis_arg: trace_basis.BasisArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+        ):
+            return 0.0
+
+        return trace_node_quadrature_weight

warp/fem/space/grid_2d_function_space.py

@@ -7,7 +7,7 @@ from warp.fem.geometry import Grid2D
 from warp.fem import cache
 
 from .topology import SpaceTopology, DiscontinuousSpaceTopologyMixin, forward_base_topology
-from .basis_space import BasisSpace, TraceBasisSpace
+from .basis_space import ShapeBasisSpace, TraceBasisSpace
 
 from .shape import ShapeFunction, ConstantShapeFunction
 from .shape import (
@@ -44,7 +44,7 @@ class Grid2DDiscontinuousSpaceTopology(
     pass
 
 
-class Grid2DBasisSpace(BasisSpace):
+class Grid2DBasisSpace(ShapeBasisSpace):
     def __init__(self, topology: Grid2DSpaceTopology, shape: ShapeFunction):
         super().__init__(topology, shape)
 

warp/fem/space/grid_3d_function_space.py

@@ -7,7 +7,7 @@ from warp.fem.geometry import Grid3D
 from warp.fem import cache
 
 from .topology import SpaceTopology, DiscontinuousSpaceTopologyMixin, forward_base_topology
-from .basis_space import BasisSpace, TraceBasisSpace
+from .basis_space import ShapeBasisSpace, TraceBasisSpace
 
 from .shape import ShapeFunction, ConstantShapeFunction
 from .shape.cube_shape_function import (
@@ -45,7 +45,7 @@ class Grid3DDiscontinuousSpaceTopology(
     pass
 
 
-class Grid3DBasisSpace(BasisSpace):
+class Grid3DBasisSpace(ShapeBasisSpace):
     def __init__(self, topology: Grid3DSpaceTopology, shape: ShapeFunction):
         super().__init__(topology, shape)
 

warp/fem/space/hexmesh_function_space.py

@@ -6,7 +6,7 @@ from warp.fem.geometry import Hexmesh
 from warp.fem import cache
 
 from .topology import SpaceTopology, DiscontinuousSpaceTopologyMixin, forward_base_topology
-from .basis_space import BasisSpace, TraceBasisSpace
+from .basis_space import ShapeBasisSpace, TraceBasisSpace
 
 from .shape import ShapeFunction, ConstantShapeFunction
 from .shape import (
@@ -121,7 +121,7 @@ class HexmeshDiscontinuousSpaceTopology(
         super().__init__(mesh, shape.NODES_PER_ELEMENT)
 
 
-class HexmeshBasisSpace(BasisSpace):
+class HexmeshBasisSpace(ShapeBasisSpace):
     def __init__(self, topology: HexmeshSpaceTopology, shape: ShapeFunction):
         super().__init__(topology, shape)
 

warp/fem/space/partition.py

@@ -1,15 +1,18 @@
 from typing import Any, Optional, Union
 
 import warp as wp
-
+from warp.fem.cache import (
+    TemporaryStore,
+    borrow_temporary,
+    borrow_temporary_like,
+    cached_arg_value,
+)
 from warp.fem.geometry import GeometryPartition, WholeGeometryPartition
-from warp.fem.utils import compress_node_indices, _iota_kernel
 from warp.fem.types import NULL_NODE_INDEX
-from warp.fem.cache import cached_arg_value, TemporaryStore, borrow_temporary, borrow_temporary_like
+from warp.fem.utils import _iota_kernel, compress_node_indices
 
-from .topology import SpaceTopology
 from .function_space import FunctionSpace
-
+from .topology import SpaceTopology
 
 wp.set_module_options({"enable_backward": False})
 
@@ -272,7 +275,7 @@ class NodePartition(SpacePartition):
 
         if device.is_cuda:
             # TODO switch to synchronize_event once available
-            wp.synchronize_stream(wp.get_stream())
+            wp.synchronize_stream(wp.get_stream(device))
 
         category_offsets.release()
 

warp/fem/space/quadmesh_2d_function_space.py

@@ -6,7 +6,7 @@ from warp.fem.geometry import Quadmesh2D
 from warp.fem import cache
 
 from .topology import SpaceTopology, DiscontinuousSpaceTopologyMixin, forward_base_topology
-from .basis_space import BasisSpace, TraceBasisSpace
+from .basis_space import ShapeBasisSpace, TraceBasisSpace
 
 from .shape import ShapeFunction, ConstantShapeFunction
 from .shape import (
@@ -116,7 +116,7 @@ class Quadmesh2DDiscontinuousSpaceTopology(
         super().__init__(mesh, shape.NODES_PER_ELEMENT)
 
 
-class Quadmesh2DBasisSpace(BasisSpace):
+class Quadmesh2DBasisSpace(ShapeBasisSpace):
     def __init__(self, topology: Quadmesh2DSpaceTopology, shape: ShapeFunction):
         super().__init__(topology, shape)
 

warp/fem/space/shape/cube_shape_function.py

@@ -41,6 +41,7 @@ class CubeTripolynomialShapeFunctions:
         self.LOBATTO_COORDS = wp.constant(NodeVec(lobatto_coords))
         self.LOBATTO_WEIGHT = wp.constant(NodeVec(lobatto_weight))
         self.LAGRANGE_SCALE = wp.constant(NodeVec(lagrange_scale))
+        self.ORDER_PLUS_ONE = wp.constant(self.ORDER + 1)
 
         self._node_ijk = self._make_node_ijk()
         self.node_type_and_type_index = self._make_node_type_and_type_index()
@@ -57,21 +58,21 @@ class CubeTripolynomialShapeFunctions:
         return wp.vec3(float(x), float(y), float(z))
 
     def _make_node_ijk(self):
-        ORDER = self.ORDER
+        ORDER_PLUS_ONE = self.ORDER_PLUS_ONE
 
         def node_ijk(
             node_index_in_elt: int,
         ):
-            node_i = node_index_in_elt // ((ORDER + 1) * (ORDER + 1))
-            node_jk = node_index_in_elt - (ORDER + 1) * (ORDER + 1) * node_i
-            node_j = node_jk // (ORDER + 1)
-            node_k = node_jk - (ORDER + 1) * node_j
+            node_i = node_index_in_elt // (ORDER_PLUS_ONE * ORDER_PLUS_ONE)
+            node_jk = node_index_in_elt - ORDER_PLUS_ONE * ORDER_PLUS_ONE * node_i
+            node_j = node_jk // ORDER_PLUS_ONE
+            node_k = node_jk - ORDER_PLUS_ONE * node_j
             return node_i, node_j, node_k
 
         return cache.get_func(node_ijk, self.name)
 
     def _make_node_type_and_type_index(self):
-        ORDER = wp.constant(self.ORDER)
+        ORDER = self.ORDER
 
         @cache.dynamic_func(suffix=self.name)
         def node_type_and_type_index(
@@ -190,13 +191,21 @@ class CubeTripolynomialShapeFunctions:
         ):
             return 0.25
 
+        def trace_node_quadrature_weight_open(
+            node_index_in_elt: int,
+        ):
+            return 0.0
+
+        if not is_closed(self.family):
+            return cache.get_func(trace_node_quadrature_weight_open, self.name)
+
         if ORDER == 1:
             return cache.get_func(trace_node_quadrature_weight_linear, self.name)
 
         return cache.get_func(trace_node_quadrature_weight, self.name)
 
     def make_element_inner_weight(self):
-        ORDER = self.ORDER
+        ORDER_PLUS_ONE = self.ORDER_PLUS_ONE
         LOBATTO_COORDS = self.LOBATTO_COORDS
         LAGRANGE_SCALE = self.LAGRANGE_SCALE
 
@@ -207,7 +216,7 @@ class CubeTripolynomialShapeFunctions:
             node_i, node_j, node_k = self._node_ijk(node_index_in_elt)
 
             w = float(1.0)
-            for k in range(ORDER + 1):
+            for k in range(ORDER_PLUS_ONE):
                 if k != node_i:
                     w *= coords[0] - LOBATTO_COORDS[k]
                 if k != node_j:
@@ -230,13 +239,13 @@ class CubeTripolynomialShapeFunctions:
             wz = (1.0 - coords[2]) * (1.0 - v[2]) + v[2] * coords[2]
             return wx * wy * wz
 
-        if ORDER == 1:
+        if self.ORDER == 1 and is_closed(self.family):
             return cache.get_func(element_inner_weight_linear, self.name)
 
         return cache.get_func(element_inner_weight, self.name)
 
     def make_element_inner_weight_gradient(self):
-        ORDER = self.ORDER
+        ORDER_PLUS_ONE = self.ORDER_PLUS_ONE
         LOBATTO_COORDS = self.LOBATTO_COORDS
         LAGRANGE_SCALE = self.LAGRANGE_SCALE
 
@@ -249,7 +258,7 @@ class CubeTripolynomialShapeFunctions:
             prefix_xy = float(1.0)
             prefix_yz = float(1.0)
             prefix_zx = float(1.0)
-            for k in range(ORDER + 1):
+            for k in range(ORDER_PLUS_ONE):
                 if k != node_i:
                     prefix_yz *= coords[0] - LOBATTO_COORDS[k]
                 if k != node_j:
@@ -265,7 +274,7 @@ class CubeTripolynomialShapeFunctions:
             grad_y = float(0.0)
             grad_z = float(0.0)
 
-            for k in range(ORDER + 1):
+            for k in range(ORDER_PLUS_ONE):
                 if k != node_i:
                     delta_x = coords[0] - LOBATTO_COORDS[k]
                     grad_x = grad_x * delta_x + prefix_x
@@ -308,7 +317,7 @@ class CubeTripolynomialShapeFunctions:
 
             return wp.vec3(dx * wy * wz, dy * wz * wx, dz * wx * wy)
 
-        if ORDER == 1:
+        if self.ORDER == 1 and is_closed(self.family):
             return cache.get_func(element_inner_weight_gradient_linear, self.name)
 
         return cache.get_func(element_inner_weight_gradient, self.name)
@@ -356,6 +365,7 @@ class CubeSerendipityShapeFunctions:
         self.LOBATTO_COORDS = wp.constant(NodeVec(lobatto_coords))
         self.LOBATTO_WEIGHT = wp.constant(NodeVec(lobatto_weight))
         self.LAGRANGE_SCALE = wp.constant(NodeVec(lagrange_scale))
+        self.ORDER_PLUS_ONE = wp.constant(self.ORDER + 1)
 
         self.node_type_and_type_index = self._get_node_type_and_type_index()
         self._node_lobatto_indices = self._get_node_lobatto_indices()
@@ -466,6 +476,7 @@ class CubeSerendipityShapeFunctions:
 
     def make_element_inner_weight(self):
         ORDER = self.ORDER
+        ORDER_PLUS_ONE = self.ORDER_PLUS_ONE
 
         LOBATTO_COORDS = self.LOBATTO_COORDS
         LAGRANGE_SCALE = self.LAGRANGE_SCALE
@@ -511,7 +522,7 @@ class CubeSerendipityShapeFunctions:
             w *= wp.select(node_all[1] == 0, local_coords[1], 1.0 - local_coords[1])
             w *= wp.select(node_all[2] == 0, local_coords[2], 1.0 - local_coords[2])
 
-            for k in range(ORDER + 1):
+            for k in range(ORDER_PLUS_ONE):
                 if k != node_all[0]:
                     w *= local_coords[0] - LOBATTO_COORDS[k]
             w *= LAGRANGE_SCALE[node_all[0]]
@@ -522,6 +533,7 @@ class CubeSerendipityShapeFunctions:
 
     def make_element_inner_weight_gradient(self):
         ORDER = self.ORDER
+        ORDER_PLUS_ONE = self.ORDER_PLUS_ONE
         LOBATTO_COORDS = self.LOBATTO_COORDS
         LAGRANGE_SCALE = self.LAGRANGE_SCALE
 
@@ -585,7 +597,7 @@ class CubeSerendipityShapeFunctions:
             w_alt = LAGRANGE_SCALE[node_all[0]]
             g_alt = float(0.0)
             prefix_alt = LAGRANGE_SCALE[node_all[0]]
-            for k in range(ORDER + 1):
+            for k in range(ORDER_PLUS_ONE):
                 if k != node_all[0]:
                     delta_alt = local_coords[0] - LOBATTO_COORDS[k]
                     w_alt *= delta_alt