warp-lang 1.2.2__py3-none-manylinux2014_aarch64.whl → 1.3.1__py3-none-manylinux2014_aarch64.whl

warp/fem/operator.py

@@ -1,9 +1,8 @@
-import inspect
 from typing import Any, Callable
 
 import warp as wp
 from warp.fem import utils
-from warp.fem.types import Domain, Field, Sample
+from warp.fem.types import Domain, Field, NodeIndex, Sample
 
 
 class Integrand:
@@ -15,7 +14,7 @@ class Integrand:
         self.func = func
         self.name = wp.codegen.make_full_qualified_name(self.func)
         self.module = wp.get_module(self.func.__module__)
-        self.argspec = inspect.getfullargspec(self.func)
+        self.argspec = wp.codegen.get_full_arg_spec(self.func)
 
 
 class Operator:
@@ -55,7 +54,7 @@ def position(domain: Domain, s: Sample):
     pass
 
 
-@operator(resolver=lambda dmn: dmn.eval_normal)
+@operator(resolver=lambda dmn: dmn.element_normal)
 def normal(domain: Domain, s: Sample):
     """Evaluates the element normal at the sample point `s`. Null for interior points."""
     pass
@@ -71,13 +70,12 @@ def deformation_gradient(domain: Domain, s: Sample):
 def lookup(domain: Domain, x: Any) -> Sample:
     """Looks-up the sample point corresponding to a world position `x`, projecting to the closest point on the domain.
 
-    Arg:
+    Args:
         x: world position of the point to look-up in the geometry
         guess: (optional) :class:`Sample` initial guess, may help perform the query
 
-    Notes:
-        Currently this operator is only fully supported for :class:`Grid2D` and :class:`Grid3D` geometries.
-        For :class:`TriangleMesh2D` and :class:`Tetmesh` geometries, the operator requires providing `guess`.
+    Note:
+        Currently this operator is unsupported for :class:`Hexmesh`, :class:`Quadmesh2D` and deformed geometries.
     """
     pass
 
@@ -142,7 +140,14 @@ def degree(f: Field):
 
 @operator(resolver=lambda f: f.at_node)
 def at_node(f: Field, s: Sample):
-    """For a Test or Trial field, returns a copy of the Sample `s` moved to the coordinates of the node being evaluated"""
+    """For a Test or Trial field `f`, returns a copy of the Sample `s` moved to the coordinates of the node being evaluated"""
+    pass
+
+
+@operator(resolver=lambda f: f.node_partition_index)
+def node_partition_index(f: Field, node_index: NodeIndex):
+    """For a NodalField `f`, returns the index of a given node in the fields's space partition,
+    or ``NULL_NODE_INDEX`` if it does not exists"""
     pass
 
 

warp/fem/quadrature/pic_quadrature.py

@@ -8,8 +8,6 @@ from warp.fem.utils import compress_node_indices
 
 from .quadrature import Quadrature
 
-wp.set_module_options({"enable_backward": False})
-
 
 class PicQuadrature(Quadrature):
     """Particle-based quadrature formula, using a global set of points unevenly spread out over geometry elements.
@@ -23,6 +21,7 @@ class PicQuadrature(Quadrature):
          define a global :meth:`Geometry.cell_lookup` method; currently this is only available for :class:`Grid2D` and :class:`Grid3D`.
         measures: Array containing the measure (area/volume) of each particle, used to defined the integration weights.
          If ``None``, defaults to the cell measure divided by the number of particles in the cell.
+        requires_grad: Whether gradients should be allocated for the computed quantities
         temporary_store: shared pool from which to allocate temporary arrays
     """
 
@@ -37,11 +36,14 @@ class PicQuadrature(Quadrature):
             ],
         ],
         measures: Optional["wp.array(dtype=float)"] = None,
+        requires_grad: bool = False,
         temporary_store: TemporaryStore = None,
     ):
         super().__init__(domain)
 
+        self._requires_grad = requires_grad
         self._bin_particles(positions, measures, temporary_store)
+        self._max_particles_per_cell: int = None
 
     @property
     def name(self):
@@ -82,22 +84,40 @@ class PicQuadrature(Quadrature):
         """Number of cells containing at least one particle"""
         return self._cell_count
 
+    def max_points_per_element(self):
+        if self._max_particles_per_cell is None:
+            max_ppc = wp.zeros(shape=(1,), dtype=int, device=self._cell_particle_offsets.array.device)
+            wp.launch(
+                PicQuadrature._max_particles_per_cell_kernel,
+                self._cell_particle_offsets.array.shape[0] - 1,
+                device=max_ppc.device,
+                inputs=[self._cell_particle_offsets.array, max_ppc],
+            )
+            self._max_particles_per_cell = int(max_ppc.numpy()[0])
+        return self._max_particles_per_cell
+
     @wp.func
-    def point_count(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex):
+    def point_count(elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex):
         return qp_arg.cell_particle_offsets[element_index + 1] - qp_arg.cell_particle_offsets[element_index]
 
     @wp.func
-    def point_coords(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, index: int):
+    def point_coords(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, index: int
+    ):
         particle_index = qp_arg.cell_particle_indices[qp_arg.cell_particle_offsets[element_index] + index]
         return qp_arg.particle_coords[particle_index]
 
     @wp.func
-    def point_weight(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, index: int):
+    def point_weight(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, index: int
+    ):
         particle_index = qp_arg.cell_particle_indices[qp_arg.cell_particle_offsets[element_index] + index]
         return qp_arg.particle_fraction[particle_index]
 
     @wp.func
-    def point_index(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, index: int):
+    def point_index(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, index: int
+    ):
         particle_index = qp_arg.cell_particle_indices[qp_arg.cell_particle_offsets[element_index] + index]
         return particle_index
 
@@ -158,7 +178,7 @@ class PicQuadrature(Quadrature):
             cell_index = cell_index_temp.array
 
             self._particle_coords_temp = borrow_temporary(
-                temporary_store, shape=positions.shape, dtype=Coords, device=device
+                temporary_store, shape=positions.shape, dtype=Coords, device=device, requires_grad=self._requires_grad
             )
             self._particle_coords = self._particle_coords_temp.array
 
@@ -183,7 +203,7 @@ class PicQuadrature(Quadrature):
             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
+            self.domain.geometry_element_count(), cell_index, return_unique_nodes=True, temporary_store=temporary_store
         )
 
         self._compute_fraction(cell_index, measures, temporary_store)
@@ -192,7 +212,7 @@ class PicQuadrature(Quadrature):
         device = cell_index.device
 
         self._particle_fraction_temp = borrow_temporary(
-            temporary_store, shape=cell_index.shape, dtype=float, device=device
+            temporary_store, shape=cell_index.shape, dtype=float, device=device, requires_grad=self._requires_grad
         )
         self._particle_fraction = self._particle_fraction_temp.array
 
@@ -241,3 +261,9 @@ class PicQuadrature(Quadrature):
                 ],
                 device=device,
             )
+
+    @wp.kernel
+    def _max_particles_per_cell_kernel(offsets: wp.array(dtype=int), max_count: wp.array(dtype=int)):
+        cell = wp.tid()
+        particle_count = offsets[cell + 1] - offsets[cell]
+        wp.atomic_max(max_count, 0, particle_count)

warp/fem/quadrature/quadrature.py

@@ -36,8 +36,8 @@ class Quadrature:
         """Total number of quadrature points over the domain"""
         raise NotImplementedError()
 
-    def points_per_element(self):
-        """Number of points per element if constant, or ``None`` if varying"""
+    def max_points_per_element(self):
+        """Maximum number of points per element if known, or ``None`` otherwise"""
         return None
 
     @staticmethod
@@ -61,7 +61,11 @@ class Quadrature:
 
     @staticmethod
     def point_index(
-        elt_arg: "domain.GeometryDomain.ElementArg", qp_arg: Arg, element_index: ElementIndex, qp_index: int
+        elt_arg: "domain.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"""
         raise NotImplementedError()
@@ -106,7 +110,7 @@ class RegularQuadrature(Quadrature):
     def total_point_count(self):
         return len(self.points) * self.domain.geometry_element_count()
 
-    def points_per_element(self):
+    def max_points_per_element(self):
         return self._N
 
     @property
@@ -121,7 +125,12 @@ class RegularQuadrature(Quadrature):
         N = self._N
 
         @cache.dynamic_func(suffix=self.name)
-        def point_count(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex):
+        def point_count(
+            elt_arg: self.domain.ElementArg,
+            qp_arg: self.Arg,
+            domain_element_index: ElementIndex,
+            element_index: ElementIndex,
+        ):
             return N
 
         return point_count
@@ -130,7 +139,13 @@ class RegularQuadrature(Quadrature):
         POINTS = self._POINTS
 
         @cache.dynamic_func(suffix=self.name)
-        def point_coords(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+        def point_coords(
+            elt_arg: self.domain.ElementArg,
+            qp_arg: self.Arg,
+            domain_element_index: ElementIndex,
+            element_index: ElementIndex,
+            qp_index: int,
+        ):
             return Coords(POINTS[qp_index, 0], POINTS[qp_index, 1], POINTS[qp_index, 2])
 
         return point_coords
@@ -139,7 +154,13 @@ class RegularQuadrature(Quadrature):
         WEIGHTS = self._WEIGHTS
 
         @cache.dynamic_func(suffix=self.name)
-        def point_weight(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+        def point_weight(
+            elt_arg: self.domain.ElementArg,
+            qp_arg: self.Arg,
+            domain_element_index: ElementIndex,
+            element_index: ElementIndex,
+            qp_index: int,
+        ):
             return WEIGHTS[qp_index]
 
         return point_weight
@@ -148,8 +169,14 @@ class RegularQuadrature(Quadrature):
         N = self._N
 
         @cache.dynamic_func(suffix=self.name)
-        def point_index(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
-            return N * element_index + qp_index
+        def 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 point_index
 
@@ -157,8 +184,8 @@ class RegularQuadrature(Quadrature):
 class NodalQuadrature(Quadrature):
     """Quadrature using space node points as quadrature points
 
-    Note that in contrast to the `nodal=True` flag for :func:`integrate`, this quadrature odes not make any assumption
-    about orthogonality of shape functions, and is thus safe to use for arbitrary integrands.
+    Note that in contrast to the `nodal=True` flag for :func:`integrate`, using this quadrature does not imply
+    any assumption about orthogonality of shape functions, and is thus safe to use for arbitrary integrands.
     """
 
     def __init__(self, domain: domain.GeometryDomain, space: FunctionSpace):
@@ -180,8 +207,8 @@ class NodalQuadrature(Quadrature):
     def total_point_count(self):
         return self._space.node_count()
 
-    def points_per_element(self):
-        return self._space.topology.NODES_PER_ELEMENT
+    def max_points_per_element(self):
+        return self._space.topology.MAX_NODES_PER_ELEMENT
 
     def _make_arg(self):
         @cache.dynamic_struct(suffix=self.name)
@@ -199,44 +226,67 @@ class NodalQuadrature(Quadrature):
         return arg
 
     def _make_point_count(self):
-        N = self._space.topology.NODES_PER_ELEMENT
-
         @cache.dynamic_func(suffix=self.name)
-        def point_count(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex):
-            return N
+        def point_count(
+            elt_arg: self.domain.ElementArg,
+            qp_arg: self.Arg,
+            domain_element_index: ElementIndex,
+            element_index: ElementIndex,
+        ):
+            return self._space.topology.element_node_count(elt_arg, qp_arg.topo_arg, element_index)
 
         return point_count
 
     def _make_point_coords(self):
         @cache.dynamic_func(suffix=self.name)
-        def point_coords(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+        def point_coords(
+            elt_arg: self.domain.ElementArg,
+            qp_arg: self.Arg,
+            domain_element_index: ElementIndex,
+            element_index: ElementIndex,
+            qp_index: int,
+        ):
             return self._space.node_coords_in_element(elt_arg, qp_arg.space_arg, element_index, qp_index)
 
         return point_coords
 
     def _make_point_weight(self):
         @cache.dynamic_func(suffix=self.name)
-        def point_weight(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+        def point_weight(
+            elt_arg: self.domain.ElementArg,
+            qp_arg: self.Arg,
+            domain_element_index: ElementIndex,
+            element_index: ElementIndex,
+            qp_index: int,
+        ):
             return self._space.node_quadrature_weight(elt_arg, qp_arg.space_arg, element_index, qp_index)
 
         return point_weight
 
     def _make_point_index(self):
         @cache.dynamic_func(suffix=self.name)
-        def point_index(elt_arg: self.domain.ElementArg, qp_arg: self.Arg, element_index: ElementIndex, qp_index: int):
+        def point_index(
+            elt_arg: self.domain.ElementArg,
+            qp_arg: self.Arg,
+            domain_element_index: ElementIndex,
+            element_index: ElementIndex,
+            qp_index: int,
+        ):
             return self._space.topology.element_node_index(elt_arg, qp_arg.topo_arg, element_index, qp_index)
 
         return point_index
 
 
 class ExplicitQuadrature(Quadrature):
-    """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.
+    """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.
+    Quadrature points may be provided for either the whole geometry or just the domain's elements.
 
     Args:
         domain: Domain of definition of the quadrature formula
-        points: 2d array of shape ``(domain.geometry_element-count(), points_per_cell)`` containing the coordinates of each quadrature point.
-        weights: 2d array of shape ``(domain.geometry_element-count(), points_per_cell)`` containing the weight for each quadrature point.
+        points: 2d array of shape  ``(domain.element_count(), points_per_cell)`` or ``(domain.geometry_element_count(), points_per_cell)`` containing the coordinates of each quadrature point.
+        weights: 2d array of shape ``(domain.element_count(), points_per_cell)`` or ``(domain.geometry_element_count(), points_per_cell)`` containing the weight for each quadrature point.
 
     See also: :class:`PicQuadrature`
     """
@@ -255,41 +305,78 @@ class ExplicitQuadrature(Quadrature):
         if points.shape != weights.shape:
             raise ValueError("Points and weights arrays must have the same shape")
 
+        if points.shape[0] == domain.geometry_element_count():
+            self.point_index = ExplicitQuadrature._point_index_geo
+            self.point_coords = ExplicitQuadrature._point_coords_geo
+            self.point_weight = ExplicitQuadrature._point_weight_geo
+        elif points.shape[0] == domain.element_count():
+            self.point_index = ExplicitQuadrature._point_index_domain
+            self.point_coords = ExplicitQuadrature._point_coords_domain
+            self.point_weight = ExplicitQuadrature._point_weight_domain
+        else:
+            raise NotImplementedError(
+                "The number of rows of points and weights must match the element count of either the domain or the geometry"
+            )
+
         self._points_per_cell = points.shape[1]
+        self._whole_geo = points.shape[0] == domain.geometry_element_count()
         self._points = points
         self._weights = weights
 
     @property
     def name(self):
-        return f"{self.__class__.__name__}"
+        return f"{self.__class__.__name__}_{self._whole_geo}"
 
     def total_point_count(self):
         return self._weights.size
 
-    def points_per_element(self):
+    def max_points_per_element(self):
         return self._points_per_cell
 
     @cache.cached_arg_value
     def arg_value(self, device):
         arg = self.Arg()
-        arg.points_per_cell = self._points_per_cell
         arg.points = self._points.to(device)
         arg.weights = self._weights.to(device)
 
         return arg
 
     @wp.func
-    def point_count(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex):
-        return qp_arg.points_per_cell
+    def point_count(elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex):
+        return qp_arg.points.shape[1]
 
     @wp.func
-    def point_coords(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, qp_index: int):
+    def _point_coords_domain(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+    ):
+        return qp_arg.points[domain_element_index, qp_index]
+
+    @wp.func
+    def _point_weight_domain(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+    ):
+        return qp_arg.weights[domain_element_index, qp_index]
+
+    @wp.func
+    def _point_index_domain(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+    ):
+        return qp_arg.points_per_cell * domain_element_index + qp_index
+
+    @wp.func
+    def _point_coords_geo(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+    ):
         return qp_arg.points[element_index, qp_index]
 
     @wp.func
-    def point_weight(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, qp_index: int):
+    def _point_weight_geo(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+    ):
         return qp_arg.weights[element_index, qp_index]
 
     @wp.func
-    def point_index(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, qp_index: int):
+    def _point_index_geo(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, qp_index: int
+    ):
         return qp_arg.points_per_cell * element_index + qp_index