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

warp/fem/operator.py

@@ -17,7 +17,15 @@ from typing import Any, Callable, Dict, Optional, Set
 
 import warp as wp
 from warp.fem.linalg import skew_part, symmetric_part
-from warp.fem.types import Coords, Domain, ElementIndex, Field, NodeIndex, Sample, make_free_sample
+from warp.fem.types import (
+    Coords,
+    Domain,
+    ElementIndex,
+    Field,
+    NodeIndex,
+    Sample,
+    make_free_sample,
+)
 
 
 class Integrand:
@@ -33,7 +41,7 @@ class Integrand:
         self.kernel_options = {} if kernel_options is None else kernel_options
 
         # Operators for each field argument. This will be populated at first integrate call
-        self.operators: Dict[str, Set[Operator]] = None
+        self.operators: Optional[Dict[str, Set[Operator]]] = None
 
 
 class Operator:
@@ -41,14 +49,21 @@ class Operator:
     Operators provide syntactic sugar over Field and Domain evaluation functions and arguments
     """
 
-    def __init__(self, func: Callable, resolver: Callable, field_result: Callable = None):
+    def __init__(
+        self,
+        func: Callable,
+        resolver: Callable,
+        field_result: Optional[Callable] = None,
+        attr: Optional[str] = None,
+    ):
         self.func = func
         self.name = func.__name__
         self.resolver = resolver
+        self.attr = attr
         self.field_result = field_result
 
 
-def integrand(func: Callable = None, kernel_options: Optional[Dict[str, Any]] = None):
+def integrand(func: Optional[Callable] = None, kernel_options: Optional[Dict[str, Any]] = None):
     """Decorator for functions to be integrated (or interpolated) using warp.fem
 
     Args:
@@ -83,84 +98,124 @@ def operator(**kwargs):
 # Domain operators
 
 
-@operator(resolver=lambda dmn: dmn.element_position)
+@operator(resolver=lambda dmn: dmn.element_position, attr="geo")
 def position(domain: Domain, s: Sample):
     """Evaluates the world position of the sample point `s`"""
     pass
 
 
-@operator(resolver=lambda dmn: dmn.element_normal)
+@operator(resolver=lambda dmn: dmn.element_normal, attr="geo")
 def normal(domain: Domain, s: Sample):
     """Evaluates the element normal at the sample point `s`. Non zero if the element is a side or the geometry is embedded in a higher-dimensional space (e.g. :class:`Trimesh3D`)"""
     pass
 
 
-@operator(resolver=lambda dmn: dmn.element_deformation_gradient)
+@operator(resolver=lambda dmn: dmn.element_deformation_gradient, attr="geo")
 def deformation_gradient(domain: Domain, s: Sample):
     """Evaluates the gradient of the domain position with respect to the element reference space at the sample point `s`"""
     pass
 
 
-@operator(resolver=lambda dmn: dmn.element_lookup)
+@operator(resolver=lambda dmn: dmn.element_lookup, attr="geo")
 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.
+    """Looks-up the sample point corresponding to a world position `x`, projecting to the closest point on the geometry.
 
     Args:
-        x: world position of the point to look-up in the geometry
-        guess: (optional) :class:`Sample` initial guess, may help perform the query
+        x (vec3): world position of the point to look-up in the geometry
+        max_dist (float): maximum distance to look for a closest point
+        guess (:class:`Sample`):  initial guess, may help perform the query
+        filter_array (wp.array): Used in conjunction with `filter_target`. Only cells such that ``filter_array[element_index]==filter_target`` will be considered.
+        filter_target (Any): See `filter_array`
+    """
+    pass
+
+
+@operator(resolver=lambda dmn: dmn.element_partition_lookup)
+def partition_lookup(domain: Domain, x: Any) -> Sample:
+    """Looks-up the sample point corresponding to a world position `x`, projecting to the closest point on the geometry partition.
 
-    Note:
-        Currently this operator is unsupported for :class:`Hexmesh`, :class:`Quadmesh2D`, :class:`Quadmesh3D` and deformed geometries.
+    Args:
+        x (vec3): world position of the point to look-up in the geometry
+        max_dist (float): maximum distance to look for a closest point
     """
     pass
 
 
-@operator(resolver=lambda dmn: dmn.element_measure)
+@operator(resolver=lambda dmn: dmn.element_measure, attr="geo")
 def measure(domain: Domain, s: Sample) -> float:
     """Returns the measure (volume, area, or length) determinant of an element at a sample point `s`"""
     pass
 
 
-@operator(resolver=lambda dmn: dmn.element_measure_ratio)
+@operator(resolver=lambda dmn: dmn.element_measure_ratio, attr="geo")
 def measure_ratio(domain: Domain, s: Sample) -> float:
     """Returns the maximum ratio between the measure of this element and that of higher-dimensional neighbors."""
     pass
 
 
+@operator(resolver=lambda dmn: dmn.element_closest_point, attr="geo")
+def element_closest_point(domain: Domain, element_index: ElementIndex, x: Any) -> Sample:
+    """
+    Computes the coordinates of the closest point to a world position within a given element.
+    Returns a tuple (closest point coordinates; squared distance to the closest point)
+
+    Args:
+        element_index: Index of the element to consider
+        x: world position of the point to compute the closest point to
+    """
+    pass
+
+
+@operator(resolver=lambda dmn: dmn.element_coordinates, attr="geo")
+def element_coordinates(domain: Domain, element_index: ElementIndex, x: Any) -> Sample:
+    """Returns the coordinates in an element reference system corresponding to a work position.
+    The returned coordinates may be in the element's exterior.
+
+    Args:
+        element_index: Index of the element to consider
+        x: world position of the point to find coordinates for
+    """
+    pass
+
+
 # Operators for evaluating cell-level quantities on domains defined on sides
 
 
 @operator(
-    resolver=lambda dmn: dmn.domain_cell_arg, field_result=lambda dmn: (dmn.cell_domain(), Domain, dmn.geometry.CellArg)
+    resolver=lambda dmn: dmn.domain_cell_arg,
+    field_result=lambda dmn: (dmn.cell_domain(), Domain, dmn.cell_domain().DomainArg),
 )
 def cells(domain: Domain) -> Domain:
     """Converts a domain defined on geometry sides to a domain defined of cells."""
     pass
 
 
-@operator(resolver=lambda dmn: dmn.element_inner_cell_index)
+@operator(resolver=lambda dmn: dmn.element_inner_cell_index, attr="geo")
 def _inner_cell_index(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
     pass
 
 
-@operator(resolver=lambda dmn: dmn.element_outer_cell_index)
+@operator(resolver=lambda dmn: dmn.element_outer_cell_index, attr="geo")
 def _outer_cell_index(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
     pass
 
 
-@operator(resolver=lambda dmn: dmn.element_inner_cell_coords)
+@operator(resolver=lambda dmn: dmn.element_inner_cell_coords, attr="geo")
 def _inner_cell_coords(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
     pass
 
 
-@operator(resolver=lambda dmn: dmn.element_outer_cell_coords)
+@operator(resolver=lambda dmn: dmn.element_outer_cell_coords, attr="geo")
 def _outer_cell_coords(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
     pass
 
 
-@operator(resolver=lambda dmn: dmn.cell_to_element_coords)
+@operator(resolver=lambda dmn: dmn.cell_to_element_coords, attr="geo")
 def _cell_to_element_coords(
-    domain: Domain, side_index: ElementIndex, cell_index: ElementIndex, cell_coords: Coords
+    domain: Domain,
+    side_index: ElementIndex,
+    cell_index: ElementIndex,
+    cell_coords: Coords,
 ) -> Sample:
     pass
 
@@ -169,7 +224,8 @@ def _cell_to_element_coords(
 def to_inner_cell(domain: Domain, s: Sample):
     """Converts a :class:`Sample` defined on a side to a sample defined on the side's inner cell"""
     return make_free_sample(
-        _inner_cell_index(domain, s.element_index), _inner_cell_coords(domain, s.element_index, s.element_coords)
+        _inner_cell_index(domain, s.element_index),
+        _inner_cell_coords(domain, s.element_index, s.element_coords),
     )
 
 
@@ -177,7 +233,8 @@ def to_inner_cell(domain: Domain, s: Sample):
 def to_outer_cell(domain: Domain, s: Sample):
     """Converts a :class:`Sample` defined on a side to a sample defined on the side's outer cell"""
     return make_free_sample(
-        _outer_cell_index(domain, s.element_index), _outer_cell_coords(domain, s.element_index, s.element_coords)
+        _outer_cell_index(domain, s.element_index),
+        _outer_cell_coords(domain, s.element_index, s.element_coords),
     )
 
 
@@ -186,10 +243,26 @@ def to_cell_side(domain: Domain, cell_s: Sample, side_index: ElementIndex):
     """Converts a :class:`Sample` defined on a cell to a sample defined on one of its side.
     If the result does not lie on the side `side_index`, the resulting coordinates will be set to ``OUTSIDE``."""
     return make_free_sample(
-        side_index, _cell_to_element_coords(domain, side_index, cell_s.element_index, cell_s.element_coords)
+        side_index,
+        _cell_to_element_coords(domain, side_index, cell_s.element_index, cell_s.element_coords),
     )
 
 
+@operator(resolver=lambda dmn: dmn.element_index, attr="index")
+def element_index(domain: Domain, domain_element_index: ElementIndex):
+    """Returns the index in the geometry of the `domain_element_index`'th domain element."""
+    pass
+
+
+@operator(resolver=lambda dmn: dmn.element_partition_index, attr="index")
+def element_partition_index(domain: Domain, cell_index: ElementIndex):
+    """Returns the index of the passed cell in the domain's geometry partition, or `NULL_ELEMENT_INDEX` if not part of the partition.
+
+    :note: Currently only available for `fem.ElementKind.CELL` elements
+    """
+    pass
+
+
 # Field operators
 # On a side, inner and outer are such that normal goes from inner to outer
 
@@ -236,9 +309,33 @@ def degree(f: Field):
     pass
 
 
+@operator(resolver=lambda f: f.node_count)
+def node_count(f: Field, s: Sample):
+    """
+    Returns the number of nodes associated to the field `f` in the element containing the sample `s`
+    """
+
+
 @operator(resolver=lambda f: f.at_node)
-def at_node(f: Field, s: Sample):
-    """For a Test or Trial field `f`, returns a copy of the Sample `s` moved to the coordinates of the node being evaluated"""
+def at_node(f: Field, s: Sample, node_index_in_elt: Optional[int] = None):
+    """
+    Returns a copy of the Sample `s` moved to the coordinates of a local node of the field `f`.
+
+    If `f` is a discrete field, `node_index_in_elt` is required and indicates the element-local index of the node to consider.
+    If `f` is a Test or Trial field, `node_index_in_elt` **must not** be provided, and will be automatically set
+    to the test (resp. trial) node currently being evaluated.
+    """
+    pass
+
+
+@operator(resolver=lambda f: f.node_index)
+def node_index(f: Field, s: Sample, node_index_in_elt: Optional[int] = None):
+    """Returns the index in the function space of a local node the field `f`.
+
+    If `f` is a discrete field, `node_index_in_elt` is required and indicates the element-local index of the node to consider.
+    If `f` is a Test or Trial field, `node_index_in_elt` **must not** be provided, and will be automatically set
+    to the test (resp. trial) node currently being evaluated.
+    """
     pass
 
 

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 NULL_ELEMENT_INDEX, Coords, ElementIndex, make_free_sample
+from warp.fem.types import NULL_ELEMENT_INDEX, OUTSIDE, Coords, ElementIndex, make_free_sample
 from warp.fem.utils import compress_node_indices
 
 from .quadrature import Quadrature
@@ -32,10 +32,10 @@ class PicQuadrature(Quadrature):
     Args:
         domain: Underlying domain for the quadrature
         positions: Either an array containing the world positions of all particles, or a tuple of arrays containing
-         the cell indices and coordinates for each particle. Note that the former requires the underlying geometry to
-         define a global :meth:`Geometry.cell_lookup` method; currently this is only available for :class:`Grid2D` and :class:`Grid3D`.
+         the cell indices and coordinates for each particle.
         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.
+        max_dist: When providing world positions that fall outside of the domain's geometry partition, maximum distance to look up for embedding cells
         requires_grad: Whether gradients should be allocated for the computed quantities
         temporary_store: shared pool from which to allocate temporary arrays
     """
@@ -52,12 +52,13 @@ class PicQuadrature(Quadrature):
         ],
         measures: Optional["wp.array(dtype=float)"] = None,
         requires_grad: bool = False,
+        max_dist: float = 0.0,
         temporary_store: TemporaryStore = None,
     ):
         super().__init__(domain)
 
         self._requires_grad = requires_grad
-        self._bin_particles(positions, measures, temporary_store)
+        self._bin_particles(positions, measures, max_dist=max_dist, temporary_store=temporary_store)
         self._max_particles_per_cell: int = None
 
     @property
@@ -86,12 +87,15 @@ class PicQuadrature(Quadrature):
     @cached_arg_value
     def arg_value(self, device) -> Arg:
         arg = PicQuadrature.Arg()
-        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)
+        self.fill_arg(arg, device)
         return arg
 
+    def fill_arg(self, args: Arg, device):
+        args.cell_particle_offsets = self._cell_particle_offsets.array.to(device)
+        args.cell_particle_indices = self._cell_particle_indices.array.to(device)
+        args.particle_fraction = self._particle_fraction.to(device)
+        args.particle_coords = self.particle_coords.to(device)
+
     def total_point_count(self):
         return self.particle_coords.shape[0]
 
@@ -179,23 +183,37 @@ class PicQuadrature(Quadrature):
             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):
+    def _bin_particles(self, positions, measures, max_dist: float, temporary_store: TemporaryStore):
         if wp.types.is_array(positions):
+            device = positions.device
+            if not self.domain.supports_lookup(device):
+                raise RuntimeError(
+                    "Attempting to build a PicQuadrature from positions on a domain that does not support global lookups"
+                )
+
+            cell_lookup = self.domain.element_partition_lookup
+            cell_coordinates = self.domain.element_coordinates
+
             # Initialize from positions
-            @dynamic_kernel(suffix=f"{self.domain.name}")
+            @dynamic_kernel(suffix=self.domain.name)
             def bin_particles(
                 cell_arg_value: self.domain.ElementArg,
+                domain_index_arg_value: self.domain.ElementIndexArg,
                 positions: wp.array(dtype=positions.dtype),
+                max_dist: float,
                 cell_index: wp.array(dtype=ElementIndex),
                 cell_coords: wp.array(dtype=Coords),
             ):
                 p = wp.tid()
-                sample = self.domain.element_lookup(cell_arg_value, positions[p])
+                sample = cell_lookup(
+                    self.domain.DomainArg(cell_arg_value, domain_index_arg_value), positions[p], max_dist
+                )
 
                 cell_index[p] = sample.element_index
-                cell_coords[p] = sample.element_coords
-
-            device = positions.device
+                if sample.element_index == NULL_ELEMENT_INDEX:
+                    cell_coords[p] = Coords(OUTSIDE)
+                else:
+                    cell_coords[p] = cell_coordinates(cell_arg_value, sample.element_index, positions[p])
 
             self._cell_index_temp = borrow_temporary(temporary_store, shape=positions.shape, dtype=int, device=device)
             self.cell_indices = self._cell_index_temp.array
@@ -210,7 +228,11 @@ class PicQuadrature(Quadrature):
                 kernel=bin_particles,
                 inputs=[
                     self.domain.element_arg_value(device),
+                    self.domain.element_index_arg_value(device),
                     positions,
+                    max_dist,
+                ],
+                outputs=[
                     self.cell_indices,
                     self.particle_coords,
                 ],

warp/fem/quadrature/quadrature.py

@@ -13,7 +13,7 @@
 # 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
@@ -55,6 +55,12 @@ class Quadrature:
         arg = Quadrature.Arg()
         return arg
 
+    def fill_arg(self, arg: Arg, device):
+        """
+        Fill the argument with the value of the argument to be passed to device
+        """
+        pass
+
     def total_point_count(self):
         """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.
@@ -199,12 +205,16 @@ class _QuadratureWithRegularEvaluationPoints(Quadrature):
     """Helper subclass for quadrature formulas which use a uniform number of
     evaluations points per element. Avoids building explicit mapping"""
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "point_evaluation_index": lambda obj: obj._make_regular_point_evaluation_index(),
+        "evaluation_point_element_index": lambda obj: obj._make_regular_evaluation_point_element_index(),
+    }
+
     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()
+        cache.setup_dynamic_attributes(self, cls=__class__)
 
     ElementIndexArg = Quadrature.Arg
     element_index_arg_value = Quadrature.arg_value
@@ -254,7 +264,7 @@ class RegularQuadrature(_QuadratureWithRegularEvaluationPoints):
 
     # Cache common formulas so we do dot have to do h2d transfer for each call
     class CachedFormula:
-        _cache = {}
+        _cache: ClassVar = {}
 
         def __init__(self, element: Element, order: int, family: Polynomial):
             self.points, self.weights = element.instantiate_quadrature(order, family)
@@ -263,9 +273,12 @@ class RegularQuadrature(_QuadratureWithRegularEvaluationPoints):
         @cache.cached_arg_value
         def arg_value(self, device):
             arg = RegularQuadrature.Arg()
+            self.fill_arg(arg, device)
+            return arg
+
+        def fill_arg(self, arg: "RegularQuadrature.Arg", device):
             arg.points = wp.array(self.points, device=device, dtype=Coords)
             arg.weights = wp.array(self.weights, device=device, dtype=float)
-            return arg
 
         @staticmethod
         def get(element: Element, order: int, family: Polynomial):
@@ -277,6 +290,13 @@ class RegularQuadrature(_QuadratureWithRegularEvaluationPoints):
                 RegularQuadrature.CachedFormula._cache[key] = quadrature
                 return quadrature
 
+    _dynamic_attribute_constructors: ClassVar = {
+        "point_count": lambda obj: obj._make_point_count(),
+        "point_index": lambda obj: obj._make_point_index(),
+        "point_coords": lambda obj: obj._make_point_coords(),
+        "point_weight": lambda obj: obj._make_point_weight(),
+    }
+
     def __init__(
         self,
         domain: GeometryDomain,
@@ -289,10 +309,7 @@ class RegularQuadrature(_QuadratureWithRegularEvaluationPoints):
 
         super().__init__(domain, self._formula.count)
 
-        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()
+        cache.setup_dynamic_attributes(self)
 
     @property
     def name(self):
@@ -315,6 +332,9 @@ class RegularQuadrature(_QuadratureWithRegularEvaluationPoints):
     def arg_value(self, device):
         return self._formula.arg_value(device)
 
+    def fill_arg(self, arg: "RegularQuadrature.Arg", device):
+        self._formula.fill_arg(arg, device)
+
     def _make_point_count(self):
         N = self._formula.count
 
@@ -374,7 +394,7 @@ class RegularQuadrature(_QuadratureWithRegularEvaluationPoints):
 class NodalQuadrature(Quadrature):
     """Quadrature using space node points as quadrature points
 
-    Note that in contrast to the `nodal=True` flag for :func:`integrate`, using this quadrature does not imply
+    Note that in contrast to the `assembly="nodal"` 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.
     """
 
@@ -412,10 +432,13 @@ class NodalQuadrature(Quadrature):
     @cache.cached_arg_value
     def arg_value(self, device):
         arg = self.Arg()
-        arg.space_arg = self._space.space_arg_value(device)
-        arg.topo_arg = self._space.topology.topo_arg_value(device)
+        self.fill_arg(arg, device)
         return arg
 
+    def fill_arg(self, arg: "NodalQuadrature.Arg", device):
+        self._space.fill_space_arg(arg.space_arg, device)
+        self._space.topology.fill_topo_arg(arg.topo_arg, device)
+
     def _make_point_count(self):
         @cache.dynamic_func(suffix=self.name)
         def point_count(
@@ -533,7 +556,7 @@ class ExplicitQuadrature(_QuadratureWithRegularEvaluationPoints):
 
     @property
     def name(self):
-        return f"{self.__class__.__name__}_{self._whole_geo}"
+        return f"{self.__class__.__name__}_{self._whole_geo}_{self._points_per_cell}"
 
     def total_point_count(self):
         return self._weights.size
@@ -541,15 +564,16 @@ class ExplicitQuadrature(_QuadratureWithRegularEvaluationPoints):
     def max_points_per_element(self):
         return self._points_per_cell
 
-    @cache.cached_arg_value
     def arg_value(self, device):
         arg = self.Arg()
+        self.fill_arg(arg, device)
+        return arg
+
+    def fill_arg(self, arg: "ExplicitQuadrature.Arg", device):
         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, domain_element_index: ElementIndex, element_index: ElementIndex):
         return qp_arg.points.shape[1]

warp/fem/space/__init__.py

@@ -120,7 +120,7 @@ def make_polynomial_basis_space(
         # Degree-1 serendipity is always equivalent to Lagrange
         element_basis = ElementBasis.LAGRANGE
 
-    shape = get_shape_function(geo.reference_cell(), geo.dimension, degree, element_basis, family)
+    shape = get_shape_function(geo.reference_cell().__class__, geo.dimension, degree, element_basis, family)
 
     if discontinuous or degree == 0 or element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
         return make_discontinuous_basis_space(geo, shape)