warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/fem/polynomial.py

@@ -0,0 +1,229 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from enum import Enum
+
+import numpy as np
+
+
+class Polynomial(Enum):
+    """Polynomial family defining interpolation nodes over an interval"""
+
+    GAUSS_LEGENDRE = "GL"
+    """Gauss--Legendre 1D polynomial family (does not include endpoints)"""
+
+    LOBATTO_GAUSS_LEGENDRE = "LGL"
+    """Lobatto--Gauss--Legendre 1D polynomial family (includes endpoints)"""
+
+    EQUISPACED_CLOSED = "closed"
+    """Closed 1D polynomial family with uniformly distributed nodes (includes endpoints)"""
+
+    EQUISPACED_OPEN = "open"
+    """Open 1D polynomial family with uniformly distributed nodes (does not include endpoints)"""
+
+    def __str__(self):
+        return self.value
+
+
+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,17 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .pic_quadrature import PicQuadrature
+from .quadrature import ExplicitQuadrature, NodalQuadrature, Quadrature, RegularQuadrature

warp/fem/quadrature/pic_quadrature.py

@@ -0,0 +1,299 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+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.utils import compress_node_indices
+
+from .quadrature import Quadrature
+
+
+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.
+        requires_grad: Whether gradients should be allocated for the computed quantities
+        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,
+        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):
+        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.element_kind != self.domain.element_kind or domain.geometry.base != self.domain.geometry.base
+        ):
+            raise RuntimeError(
+                "The new domain must use the same base geometry and kind of elements as the current one."
+            )
+
+        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
+
+    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, 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, 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, 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, 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
+
+    @wp.func
+    def point_evaluation_index(
+        elt_arg: Any, qp_arg: Arg, domain_element_index: ElementIndex, element_index: ElementIndex, index: int
+    ):
+        return qp_arg.cell_particle_offsets[element_index] + 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.where(element_particle_offsets[i] == element_particle_offsets[i + 1], 0, 1)
+
+    @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]
+        if cell == NULL_ELEMENT_INDEX:
+            cell_fraction[p] = 0.0
+        else:
+            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
+
+            self._cell_index_temp = borrow_temporary(temporary_store, shape=positions.shape, dtype=int, device=device)
+            self.cell_indices = self._cell_index_temp.array
+
+            self._particle_coords_temp = borrow_temporary(
+                temporary_store, shape=positions.shape, dtype=Coords, device=device, requires_grad=self._requires_grad
+            )
+            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,
+                    self.cell_indices,
+                    self.particle_coords,
+                ],
+                device=device,
+            )
+
+        else:
+            self.cell_indices, self.particle_coords = positions
+            if self.cell_indices.shape != self.particle_coords.shape:
+                raise ValueError("Cell index and coordinates arrays must have the same shape")
+
+            self._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(),
+            self.cell_indices,
+            return_unique_nodes=True,
+            temporary_store=temporary_store,
+        )
+
+        self._compute_fraction(self.cell_indices, 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, requires_grad=self._requires_grad
+        )
+        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()
+
+                cell = cell_index[p]
+                if cell == NULL_ELEMENT_INDEX:
+                    cell_fraction[p] = 0.0
+                else:
+                    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,
+            )
+
+    @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)