warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/fem/operator.py

@@ -0,0 +1,191 @@
+import inspect
+from typing import Callable, Any
+
+import warp as wp
+
+from warp.fem.types import Domain, Field, Sample
+from warp.fem import utils
+
+
+class Integrand:
+    """An integrand is a device function containing arbitrary expressions over Field and Domain variables.
+    It will get transformed to a proper warp.Function by resolving concrete Field types at call time.
+    """
+
+    def __init__(self, func: Callable):
+        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)
+
+
+class Operator:
+    """
+    Operators provide syntaxic sugar over Field and Domain evaluation functions and arguments
+    """
+
+    def __init__(self, func: Callable, resolver: Callable):
+        self.func = func
+        self.resolver = resolver
+
+
+def integrand(func: Callable):
+    """Decorator for functions to be integrated (or interpolated) using warp.fem"""
+    itg = Integrand(func)
+    itg.__doc__ = func.__doc__
+    return itg
+
+
+def operator(resolver: Callable):
+    """Decorator for functions operating on Field-like or Domain-like data inside warp.fem integrands"""
+
+    def wrap_operator(func: Callable):
+        op = Operator(func, resolver)
+        op.__doc__ = func.__doc__
+        return op
+
+    return wrap_operator
+
+
+# Domain operators
+
+
+@operator(resolver=lambda dmn: dmn.element_position)
+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, 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) -> 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, 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, s: Sample) -> float:
+    """Returns the maximum ratio between the measure of this element and that of higher-dimensional neighbours."""
+    pass
+
+
+# Field operators
+# On a side, inner and outer are such that normal goes from inner to outer
+
+
+@operator(resolver=lambda f: f.eval_inner)
+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, 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, 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 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
+
+
+@operator(resolver=lambda f: f.eval_degree)
+def degree(f: Field):
+    """Polynomial degree of a field"""
+    pass
+
+
+@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"""
+    pass
+
+
+# Common derived operators, for convenience
+
+
+@integrand
+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 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, s: Sample):
+    """Jump between inner and outer element values on an interior side. Zero for interior points or domain boundaries"""
+    return inner(f, s) - outer(f, s)
+
+
+@integrand
+def average(f: Field, s: Sample):
+    """Average between inner and outer element values"""
+    return 0.5 * (inner(f, s) + outer(f, s))
+
+
+@integrand
+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, s) - grad_outer(f, s)
+
+
+@integrand
+def grad_average(f: Field, s: Sample):
+    """Average between inner and outer element gradients"""
+    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)
+Field.call_operator = inner
+Domain.call_operator = position

warp/fem/polynomial.py

@@ -0,0 +1,213 @@
+import math
+from enum import Enum
+
+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"""
+    return family == Polynomial.LOBATTO_GAUSS_LEGENDRE or family == Polynomial.EQUISPACED_CLOSED
+
+
+def _gauss_legendre_quadrature_1d(n: int):
+    if n == 1:
+        coords = [0.0]
+        weights = [2.0]
+    elif n == 2:
+        coords = [-math.sqrt(1.0 / 3), math.sqrt(1.0 / 3)]
+        weights = [1.0, 1.0]
+    elif n == 3:
+        coords = [0.0, -math.sqrt(3.0 / 5.0), math.sqrt(3.0 / 5.0)]
+        weights = [8.0 / 9.0, 5.0 / 9.0, 5.0 / 9.0]
+    elif n == 4:
+        c_a = math.sqrt(3.0 / 7.0 - 2.0 / 7.0 * math.sqrt(6.0 / 5.0))
+        c_b = math.sqrt(3.0 / 7.0 + 2.0 / 7.0 * math.sqrt(6.0 / 5.0))
+        w_a = (18.0 + math.sqrt(30.0)) / 36.0
+        w_b = (18.0 - math.sqrt(30.0)) / 36.0
+        coords = [c_a, -c_a, c_b, -c_b]
+        weights = [w_a, w_a, w_b, w_b]
+    elif n == 5:
+        c_a = 1.0 / 3.0 * math.sqrt(5.0 - 2.0 * math.sqrt(10.0 / 7.0))
+        c_b = 1.0 / 3.0 * math.sqrt(5.0 + 2.0 * math.sqrt(10.0 / 7.0))
+        w_a = (322.0 + 13.0 * math.sqrt(70.0)) / 900.0
+        w_b = (322.0 - 13.0 * math.sqrt(70.0)) / 900.0
+        coords = [0.0, c_a, -c_a, c_b, -c_b]
+        weights = [128.0 / 225.0, w_a, w_a, w_b, w_b]
+    else:
+        raise NotImplementedError
+
+    # Shift from [-1, 1] to [0, 1]
+    weights = 0.5 * np.array(weights)
+    coords = 0.5 * np.array(coords) + 0.5
+
+    return coords, weights
+
+
+def _lobatto_gauss_legendre_quadrature_1d(n: int):
+    if n == 2:
+        coords = [-1.0, 1.0]
+        weights = [1.0, 1.0]
+    elif n == 3:
+        coords = [-1.0, 0.0, 1.0]
+        weights = [1.0 / 3.0, 4.0 / 3.0, 1.0 / 3.0]
+    elif n == 4:
+        coords = [-1.0, -1.0 / math.sqrt(5.0), 1.0 / math.sqrt(5.0), 1.0]
+        weights = [1.0 / 6.0, 5.0 / 6.0, 5.0 / 6.0, 1.0 / 6.0]
+    elif n == 5:
+        coords = [-1.0, -math.sqrt(3.0 / 7.0), 0.0, math.sqrt(3.0 / 7.0), 1.0]
+        weights = [1.0 / 10.0, 49.0 / 90.0, 32.0 / 45.0, 49.0 / 90.0, 1.0 / 10.0]
+    else:
+        raise NotImplementedError
+
+    # Shift from [-1, 1] to [0, 1]
+    weights = 0.5 * np.array(weights)
+    coords = 0.5 * np.array(coords) + 0.5
+
+    return coords, weights
+
+
+def _uniform_open_quadrature_1d(n: int):
+    step = 1.0 / (n + 1)
+    coords = np.linspace(step, 1.0 - step, n)
+    weights = np.full(n, 1.0 / (n + 1))
+
+    # Boundaries have 3/2 the weight
+    weights[0] = 1.5 / (n + 1)
+    weights[-1] = 1.5 / (n + 1)
+
+    return coords, weights
+
+
+def _uniform_closed_quadrature_1d(n: int):
+    coords = np.linspace(0.0, 1.0, n)
+    weights = np.full(n, 1.0 / (n - 1))
+
+    # Boundaries have half the weight
+    weights[0] = 0.5 / (n - 1)
+    weights[-1] = 0.5 / (n - 1)
+
+    return coords, weights
+
+
+def _open_newton_cotes_quadrature_1d(n: int):
+    step = 1.0 / (n + 1)
+    coords = np.linspace(step, 1.0 - step, n)
+
+    # Weisstein, Eric W. "Newton-Cotes Formulas." From MathWorld--A Wolfram Web Resource.
+    # https://mathworld.wolfram.com/Newton-CotesFormulas.html
+
+    if n == 1:
+        weights = np.array([1.0])
+    elif n == 2:
+        weights = np.array([0.5, 0.5])
+    elif n == 3:
+        weights = np.array([2.0, -1.0, 2.0]) / 3.0
+    elif n == 4:
+        weights = np.array([11.0, 1.0, 1.0, 11.0]) / 24.0
+    elif n == 5:
+        weights = np.array([11.0, -14.0, 26.0, -14.0, 11.0]) / 20.0
+    elif n == 6:
+        weights = np.array([611.0, -453.0, 562.0, 562.0, -453.0, 611.0]) / 1440.0
+    elif n == 7:
+        weights = np.array([460.0, -954.0, 2196.0, -2459.0, 2196.0, -954.0, 460.0]) / 945.0
+    else:
+        raise NotImplementedError
+
+    return coords, weights
+
+
+def _closed_newton_cotes_quadrature_1d(n: int):
+    coords = np.linspace(0.0, 1.0, n)
+
+    # OEIS: A093735, A093736
+
+    if n == 2:
+        weights = np.array([1.0, 1.0]) / 2.0
+    elif n == 3:
+        weights = np.array([1.0, 4.0, 1.0]) / 3.0
+    elif n == 4:
+        weights = np.array([3.0, 9.0, 9.0, 3.0]) / 8.0
+    elif n == 5:
+        weights = np.array([14.0, 64.0, 24.0, 64.0, 14.0]) / 45.0
+    elif n == 6:
+        weights = np.array([95.0 / 288.0, 125.0 / 96.0, 125.0 / 144.0, 125.0 / 144.0, 125.0 / 96.0, 95.0 / 288.0])
+    elif n == 7:
+        weights = np.array([41, 54, 27, 68, 27, 54, 41], dtype=float) / np.array(
+            [140, 35, 140, 35, 140, 35, 140], dtype=float
+        )
+    elif n == 8:
+        weights = np.array(
+            [
+                5257,
+                25039,
+                343,
+                20923,
+                20923,
+                343,
+                25039,
+                5257,
+            ]
+        ) / np.array(
+            [
+                17280,
+                17280,
+                640,
+                17280,
+                17280,
+                640,
+                17280,
+                17280,
+            ],
+            dtype=float,
+        )
+    else:
+        raise NotImplementedError
+
+    # Normalize with interval length
+    weights = weights / (n - 1)
+
+    return coords, weights
+
+
+def quadrature_1d(point_count: int, family: Polynomial):
+    """Return quadrature points and weights for the given family and point count"""
+
+    if family == Polynomial.GAUSS_LEGENDRE:
+        return _gauss_legendre_quadrature_1d(point_count)
+    if family == Polynomial.LOBATTO_GAUSS_LEGENDRE:
+        return _lobatto_gauss_legendre_quadrature_1d(point_count)
+    if family == Polynomial.EQUISPACED_CLOSED:
+        return _closed_newton_cotes_quadrature_1d(point_count)
+    if family == Polynomial.EQUISPACED_OPEN:
+        return _open_newton_cotes_quadrature_1d(point_count)
+
+    raise NotImplementedError
+
+
+def lagrange_scales(coords: np.array):
+    """Return the scaling factors for Lagrange polynomials with roots at coords"""
+    lagrange_scale = np.empty_like(coords)
+    for i in range(len(coords)):
+        deltas = coords[i] - coords
+        deltas[i] = 1.0
+        lagrange_scale[i] = 1.0 / np.prod(deltas)
+
+    return lagrange_scale

warp/fem/quadrature/__init__.py

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

warp/fem/quadrature/pic_quadrature.py

@@ -0,0 +1,245 @@
+from typing import Union, Tuple, Any, Optional
+
+import warp as wp
+
+from warp.fem.domain import GeometryDomain
+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
+
+
+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.
+
+    Useful for Particle-In-Cell and derived methods.
+
+    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`.
+        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: 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._bin_particles(positions, measures, temporary_store)
+
+    @property
+    def name(self):
+        return f"{self.__class__.__name__}"
+
+    @Quadrature.domain.setter
+    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
+        ):
+            raise RuntimeError(
+                "Cannot change the domain to that of a different Geometry and/or using different element kinds."
+            )
+
+        self._domain = domain
+
+    @wp.struct
+    class Arg:
+        cell_particle_offsets: wp.array(dtype=int)
+        cell_particle_indices: wp.array(dtype=int)
+        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.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._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(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(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(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(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)"):
+        """Fills a mask array such that all non-empty elements are set to 1, all empty elements to zero.
+
+        Args:
+            mask: Int warp array with size at least equal to `self.domain.geometry_element_count()`
+        """
+
+        wp.launch(
+            kernel=PicQuadrature._fill_mask_kernel,
+            dim=self.domain.geometry_element_count(),
+            device=mask.device,
+            inputs=[self._cell_particle_offsets.array, mask],
+        )
+
+    @wp.kernel
+    def _fill_mask_kernel(
+        element_particle_offsets: wp.array(dtype=int),
+        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)
+
+    @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()
+
+        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)
+
+    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])
+
+                cell_index[p] = sample.element_index
+                cell_coords[p] = sample.element_coords
+
+            device = positions.device
+
+            cell_index_temp = borrow_temporary(temporary_store, shape=positions.shape, dtype=int, device=device)
+            cell_index = cell_index_temp.array
+
+            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,
+            )