warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/_src/fem/operator.py

@@ -0,0 +1,396 @@
+# 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, Callable, Dict, Optional, Set
+
+import warp as wp
+from warp._src.codegen import get_full_arg_spec, make_full_qualified_name
+from warp._src.fem.linalg import skew_part, symmetric_part
+from warp._src.fem.types import (
+    Coords,
+    Domain,
+    ElementIndex,
+    Field,
+    NodeIndex,
+    Sample,
+    make_free_sample,
+)
+
+
+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, kernel_options: Optional[Dict[str, Any]] = None):
+        self.func = func
+        self.name = make_full_qualified_name(self.func)
+        self.module = wp.get_module(self.func.__module__)
+        self.argspec = get_full_arg_spec(self.func)
+        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: Optional[Dict[str, Set[Operator]]] = None
+
+        # Cached kernels for each integrand call
+        self.cached_kernels = {}
+        self.cached_funcs = {}
+
+
+class Operator:
+    """
+    Operators provide syntactic sugar over Field and Domain evaluation functions and arguments
+    """
+
+    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: Optional[Callable] = None, kernel_options: Optional[Dict[str, Any]] = None):
+    """Decorator for functions to be integrated (or interpolated) using warp.fem
+
+    Args:
+        func: Decorated function
+        kernel_options: Supplemental code-generation options to be passed to the generated kernel.
+    """
+
+    if func is not None:
+        itg = Integrand(func)
+        itg.__doc__ = func.__doc__
+        return itg
+
+    def wrap_integrand(func: Callable):
+        itg = Integrand(func, kernel_options)
+        itg.__doc__ = func.__doc__
+        return itg
+
+    return wrap_integrand
+
+
+def operator(**kwargs):
+    """Decorator for functions operating on Field-like or Domain-like data inside warp.fem integrands"""
+
+    def wrap_operator(func: Callable):
+        op = Operator(func, **kwargs)
+        op.__doc__ = func.__doc__
+        return op
+
+    return wrap_operator
+
+
+# Domain operators
+
+
+@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, 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, 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, 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 geometry.
+
+    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
+        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.
+
+    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, 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, 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.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, 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, 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, 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, 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, attr="geo")
+def _cell_to_element_coords(
+    domain: Domain,
+    side_index: ElementIndex,
+    cell_index: ElementIndex,
+    cell_coords: Coords,
+) -> Sample:
+    pass
+
+
+@integrand
+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),
+    )
+
+
+@integrand
+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),
+    )
+
+
+@integrand
+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),
+    )
+
+
+@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
+
+
+@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_div_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.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`
+    """
+    pass
+
+
+@operator(resolver=lambda f: f.at_node)
+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
+
+
+@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
+
+
+# Common derived operators, for convenience
+
+
+@integrand
+def D(f: Field, s: Sample):
+    """Symmetric part of the (inner) gradient of the field at `s`"""
+    return 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 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/_src/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