warp-lang 1.0.0b2__py3-none-win_amd64.whl → 1.0.0b6__py3-none-win_amd64.whl

warp/fem/operator.py

@@ -51,31 +51,46 @@ def operator(resolver: Callable):
 
 
 @operator(resolver=lambda dmn: dmn.element_position)
-def position(domain: Domain, x: Sample):
-    """Evaluates the world position of the sample point x"""
+def position(domain: Domain, s: Sample):
+    """Evaluates the world position of the sample point `s`"""
     pass
 
 
 @operator(resolver=lambda dmn: dmn.eval_normal)
-def normal(domain: Domain, x: Sample):
-    """Evaluates the element normal at the sample point x. Null for interior points."""
+def normal(domain: Domain, s: Sample):
+    """Evaluates the element normal at the sample point `s`. Null for interior points."""
+    pass
+
+
+@operator(resolver=lambda dmn: dmn.element_deformation_gradient)
+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)
-def lookup(domain: Domain, x: Any):
-    """Look-ups a sample point from a world position, projecting to the closest point on the domain"""
+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:
+        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`.
+    """
     pass
 
 
 @operator(resolver=lambda dmn: dmn.element_measure)
-def measure(domain: Domain, sample: Sample):
-    """Returns the measure (volume, area, or length) of an element"""
+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)
-def measure_ratio(domain: Domain, sample: Sample):
+def measure_ratio(domain: Domain, s: Sample) -> float:
     """Returns the maximum ratio between the measure of this element and that of higher-dimensional neighbours."""
     pass
 
@@ -85,26 +100,38 @@ def measure_ratio(domain: Domain, sample: Sample):
 
 
 @operator(resolver=lambda f: f.eval_inner)
-def inner(f: Field, x: Sample):
-    """Evaluates the field at a sample point x. On oriented sides, use the inner element"""
+def inner(f: Field, s: Sample):
+    """Evaluates the field at a sample point `s`. On oriented sides, uses the inner element"""
     pass
 
 
 @operator(resolver=lambda f: f.eval_grad_inner)
-def grad(f: Field, x: Sample):
-    """Evaluates the field gradient at a sample point x. On oriented sides, use the inner element"""
+def grad(f: Field, s: Sample):
+    """Evaluates the field gradient at a sample point `s`. On oriented sides, uses the inner element"""
+    pass
+
+
+@operator(resolver=lambda f: f.eval_div_inner)
+def div(f: Field, s: Sample):
+    """Evaluates the field divergence at a sample point `s`. On oriented sides, uses the inner element"""
     pass
 
 
 @operator(resolver=lambda f: f.eval_outer)
-def outer(f: Field, x: Sample):
-    """Evaluates the field at a sample point x. On oriented sides, use the outer element. On interior points and on domain boundaries, this is equivalent to inner."""
+def outer(f: Field, s: Sample):
+    """Evaluates the field at a sample point `s`. On oriented sides, uses the outer element. On interior points and on domain boundaries, this is equivalent to :func:`inner`."""
+    pass
+
+
+@operator(resolver=lambda f: f.eval_grad_outer)
+def grad_outer(f: Field, s: Sample):
+    """Evaluates the field gradient at a sample point `s`. On oriented sides, uses the outer element. On interior points and on domain boundaries, this is equivalent to :func:`grad`."""
     pass
 
 
 @operator(resolver=lambda f: f.eval_grad_outer)
-def grad_outer(f: Field, x: Sample):
-    """Evaluates the field gradient at a sample point x. On oriented sides, use the outer element. On interior points and on domain boundaries, this is equivalent to grad."""
+def div_outer(f: Field, s: Sample):
+    """Evaluates the field divergence at a sample point `s`. On oriented sides, uses the outer element. On interior points and on domain boundaries, this is equivalent to :func:`div`."""
     pass
 
 
@@ -124,39 +151,39 @@ def at_node(f: Field, s: Sample):
 
 
 @integrand
-def D(f: Field, x: Sample):
-    """Symmetric part of the (inner) gradient of the field at x"""
-    return utils.symmetric_part(grad(f, x))
+def D(f: Field, s: Sample):
+    """Symmetric part of the (inner) gradient of the field at `s`"""
+    return utils.symmetric_part(grad(f, s))
 
 
 @integrand
-def div(f: Field, x: Sample):
-    """(Inner) divergence of the field at x"""
-    return wp.trace(grad(f, x))
+def curl(f: Field, s: Sample):
+    """Skew part of the (inner) gradient of the field at `s`, as a vector such that ``wp.cross(curl(u), v) = skew(grad(u)) v``"""
+    return utils.skew_part(grad(f, s))
 
 
 @integrand
-def jump(f: Field, x: Sample):
+def jump(f: Field, s: Sample):
     """Jump between inner and outer element values on an interior side. Zero for interior points or domain boundaries"""
-    return inner(f, x) - outer(f, x)
+    return inner(f, s) - outer(f, s)
 
 
 @integrand
-def average(f: Field, x: Sample):
+def average(f: Field, s: Sample):
     """Average between inner and outer element values"""
-    return 0.5 * (inner(f, x) + outer(f, x))
+    return 0.5 * (inner(f, s) + outer(f, s))
 
 
 @integrand
-def grad_jump(f: Field, x: Sample):
+def grad_jump(f: Field, s: Sample):
     """Jump between inner and outer element gradients on an interior side. Zero for interior points or domain boundaries"""
-    return grad(f, x) - grad_outer(f, x)
+    return grad(f, s) - grad_outer(f, s)
 
 
 @integrand
-def grad_average(f: Field, x: Sample):
+def grad_average(f: Field, s: Sample):
     """Average between inner and outer element gradients"""
-    return 0.5 * (grad(f, x) + grad_outer(f, x))
+    return 0.5 * (grad(f, s) + grad_outer(f, s))
 
 
 # Set default call operators for argument types, so that field(s) = inner(field, s) and domain(s) = position(domain, s)

warp/fem/polynomial.py

@@ -5,11 +5,22 @@ import numpy as np
 
 
 class Polynomial(Enum):
+    """Polynomial family defining interpolation nodes over an interval"""
+
     GAUSS_LEGENDRE = 0
+    """Gauss--Legendre 1D polynomial family (does not include endpoints)"""
+    
     LOBATTO_GAUSS_LEGENDRE = 1
+    """Lobatto--Gauss--Legendre 1D polynomial family (includes endpoints)"""
+
     EQUISPACED_CLOSED = 2
+    """Closed 1D polynomial family with uniformly distributed nodes (includes endpoints)"""
+
     EQUISPACED_OPEN = 3
+    """Open 1D polynomial family with uniformly distributed nodes (does not include endpoints)"""
 
+    def __str__(self):
+        return self.name
 
 def is_closed(family: Polynomial):
     """Whether the polynomial roots include interval endpoints"""

warp/fem/quadrature/__init__.py

@@ -1,2 +1,2 @@
-from .quadrature import Quadrature, RegularQuadrature
+from .quadrature import Quadrature, RegularQuadrature, NodalQuadrature, ExplicitQuadrature
 from .pic_quadrature import PicQuadrature

warp/fem/quadrature/pic_quadrature.py

@@ -1,8 +1,11 @@
+from typing import Union, Tuple, Any, Optional
+
 import warp as wp
 
 from warp.fem.domain import GeometryDomain
-from warp.fem.types import ElementIndex, Coords
+from warp.fem.types import ElementIndex, Coords, make_free_sample
 from warp.fem.utils import compress_node_indices
+from warp.fem.cache import cached_arg_value, TemporaryStore, borrow_temporary, dynamic_kernel
 
 from .quadrature import Quadrature
 
@@ -11,23 +14,36 @@ wp.set_module_options({"enable_backward": False})
 
 
 class PicQuadrature(Quadrature):
-    """Particle-based quadrature formula, using a global set of points irregularely spread out over geometry elements.
+    """Particle-based quadrature formula, using a global set of points unevenly spread out over geometry elements.
 
     Useful for Particle-In-Cell and derived methods.
+
+    Args:
+        domain: Undelying domain for the qaudrature
+        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`.
+        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.
+        temporary_store: shared pool from which to allocate temporary arrays
     """
 
     def __init__(
         self,
         domain: GeometryDomain,
-        positions: "wp.array()",
-        measures: "wp.array(dtype=float)",
+        positions: Union[
+            "wp.array(dtype=wp.vecXd)",
+            Tuple[
+                "wp.array(dtype=ElementIndex)",
+                "wp.array(dtype=Coords)",
+            ],
+        ],
+        measures: Optional["wp.array(dtype=float)"] = None,
+        temporary_store: TemporaryStore = None,
     ):
         super().__init__(domain)
 
-        self.positions = positions
-        self.measures = measures
-
-        self._bin_particles()
+        self._bin_particles(positions, measures, temporary_store)
 
     @property
     def name(self):
@@ -52,34 +68,39 @@ class PicQuadrature(Quadrature):
         particle_fraction: wp.array(dtype=float)
         particle_coords: wp.array(dtype=Coords)
 
+    @cached_arg_value
     def arg_value(self, device) -> Arg:
         arg = PicQuadrature.Arg()
-        arg.cell_particle_offsets = self._cell_particle_offsets.to(device)
-        arg.cell_particle_indices = self._cell_particle_indices.to(device)
+        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)
         return arg
 
     def total_point_count(self):
-        return self.positions.shape[0]
+        return self._particle_coords.shape[0]
+
+    def active_cell_count(self):
+        """Number of cells containing at least one particle"""
+        return self._cell_count
 
     @wp.func
-    def point_count(arg: Arg, element_index: ElementIndex):
-        return arg.cell_particle_offsets[element_index + 1] - arg.cell_particle_offsets[element_index]
+    def point_count(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex):
+        return qp_arg.cell_particle_offsets[element_index + 1] - qp_arg.cell_particle_offsets[element_index]
 
     @wp.func
-    def point_coords(arg: Arg, element_index: ElementIndex, index: int):
-        particle_index = arg.cell_particle_indices[arg.cell_particle_offsets[element_index] + index]
-        return arg.particle_coords[particle_index]
+    def point_coords(elt_arg: Any, qp_arg: Arg, 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(arg: Arg, element_index: ElementIndex, index: int):
-        particle_index = arg.cell_particle_indices[arg.cell_particle_offsets[element_index] + index]
-        return arg.particle_fraction[particle_index]
+    def point_weight(elt_arg: Any, qp_arg: Arg, 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(arg: Arg, element_index: ElementIndex, index: int):
-        particle_index = arg.cell_particle_indices[arg.cell_particle_offsets[element_index] + index]
+    def point_index(elt_arg: Any, qp_arg: Arg, element_index: ElementIndex, index: int):
+        particle_index = qp_arg.cell_particle_indices[qp_arg.cell_particle_offsets[element_index] + index]
         return particle_index
 
     def fill_element_mask(self, mask: "wp.array(dtype=int)"):
@@ -93,7 +114,7 @@ class PicQuadrature(Quadrature):
             kernel=PicQuadrature._fill_mask_kernel,
             dim=self.domain.geometry_element_count(),
             device=mask.device,
-            inputs=[self._cell_particle_offsets, mask],
+            inputs=[self._cell_particle_offsets.array, mask],
         )
 
     @wp.kernel
@@ -104,50 +125,121 @@ class PicQuadrature(Quadrature):
         i = wp.tid()
         element_mask[i] = wp.select(element_particle_offsets[i] == element_particle_offsets[i + 1], 1, 0)
 
-    def _bin_particles(self):
-        from warp.fem import cache
+    @wp.kernel
+    def _compute_uniform_fraction(
+        cell_index: wp.array(dtype=ElementIndex),
+        cell_particle_offsets: wp.array(dtype=int),
+        cell_fraction: wp.array(dtype=float),
+    ):
+        p = wp.tid()
 
-        def bin_particles_fn(
-            cell_arg_value: self.domain.ElementArg,
-            positions: wp.array(dtype=self.positions.dtype),
-            measures: wp.array(dtype=float),
-            cell_index: wp.array(dtype=ElementIndex),
-            cell_coords: wp.array(dtype=Coords),
-            cell_fraction: wp.array(dtype=float),
-        ):
-            p = wp.tid()
-            sample = self.domain.element_lookup(cell_arg_value, positions[p])
+        cell = cell_index[p]
+        cell_particle_count = cell_particle_offsets[cell + 1] - cell_particle_offsets[cell]
 
-            cell_index[p] = sample.element_index
+        cell_fraction[p] = 1.0 / float(cell_particle_count)
 
-            cell_coords[p] = sample.element_coords
-            cell_fraction[p] = measures[p] / self.domain.element_measure(cell_arg_value, sample)
+    def _bin_particles(self, positions, measures, temporary_store: TemporaryStore):
+        if wp.types.is_array(positions):
+            # Initialize from positions
+            @dynamic_kernel(suffix=f"{self.domain.name}")
+            def bin_particles(
+                cell_arg_value: self.domain.ElementArg,
+                positions: wp.array(dtype=positions.dtype),
+                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])
 
-        bin_particles = cache.get_kernel(
-            bin_particles_fn,
-            suffix=f"{self.domain.name}",
-        )
+                cell_index[p] = sample.element_index
+                cell_coords[p] = sample.element_coords
 
-        device = self.positions.device
+            device = positions.device
 
-        cell_index = wp.empty(shape=self.positions.shape, dtype=int, device=device)
-        self._particle_coords = wp.empty(shape=self.positions.shape, dtype=Coords, device=device)
-        self._particle_fraction = wp.empty(shape=self.positions.shape, dtype=float, device=device)
+            cell_index_temp = borrow_temporary(temporary_store, shape=positions.shape, dtype=int, device=device)
+            cell_index = cell_index_temp.array
 
-        wp.launch(
-            dim=self.positions.shape[0],
-            kernel=bin_particles,
-            inputs=[
-                self.domain.element_arg_value(device),
-                self.positions,
-                self.measures,
-                cell_index,
-                self._particle_coords,
-                self._particle_fraction,
-            ],
-            device=device,
-        )
+            self._particle_coords_temp = borrow_temporary(
+                temporary_store, shape=positions.shape, dtype=Coords, device=device
+            )
+            self._particle_coords = self._particle_coords_temp.array
+
+            wp.launch(
+                dim=positions.shape[0],
+                kernel=bin_particles,
+                inputs=[
+                    self.domain.element_arg_value(device),
+                    positions,
+                    cell_index,
+                    self._particle_coords,
+                ],
+                device=device,
+            )
+
+        else:
+            cell_index, self._particle_coords = positions
+            if cell_index.shape != self._particle_coords.shape:
+                raise ValueError("Cell index and coordinates arrays must have the same shape")
+
+            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
         )
+
+        self._compute_fraction(cell_index, measures, temporary_store)
+
+    def _compute_fraction(self, cell_index, measures, temporary_store: TemporaryStore):
+        device = cell_index.device
+
+        self._particle_fraction_temp = borrow_temporary(
+            temporary_store, shape=cell_index.shape, dtype=float, device=device
+        )
+        self._particle_fraction = self._particle_fraction_temp.array
+
+        if measures is None:
+            # Split fraction uniformly over all particles in cell
+
+            wp.launch(
+                dim=cell_index.shape,
+                kernel=PicQuadrature._compute_uniform_fraction,
+                inputs=[
+                    cell_index,
+                    self._cell_particle_offsets.array,
+                    self._particle_fraction,
+                ],
+                device=device,
+            )
+
+        else:
+            # Fraction from particle measure
+
+            if measures.shape != cell_index.shape:
+                raise ValueError("Measures should be an 1d array or length equal to particle count")
+
+            @dynamic_kernel(suffix=f"{self.domain.name}")
+            def compute_fraction(
+                cell_arg_value: self.domain.ElementArg,
+                measures: wp.array(dtype=float),
+                cell_index: wp.array(dtype=ElementIndex),
+                cell_coords: wp.array(dtype=Coords),
+                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)
+
+            wp.launch(
+                dim=measures.shape[0],
+                kernel=compute_fraction,
+                inputs=[
+                    self.domain.element_arg_value(device),
+                    measures,
+                    cell_index,
+                    self._particle_coords,
+                    self._particle_fraction,
+                ],
+                device=device,
+            )