warp-lang 1.4.2__py3-none-manylinux2014_x86_64.whl → 1.5.0__py3-none-manylinux2014_x86_64.whl

warp/fem/linalg.py

@@ -0,0 +1,404 @@
+from typing import Any
+
+import warp as wp
+
+
+@wp.func
+def generalized_outer(x: Any, y: Any):
+    """Generalized outer product allowing for the first argument to be a scalar"""
+    return wp.outer(x, y)
+
+
+@wp.func
+def generalized_outer(x: wp.float32, y: wp.vec2):
+    return x * y
+
+
+@wp.func
+def generalized_outer(x: wp.float32, y: wp.vec3):
+    return x * y
+
+
+@wp.func
+def generalized_inner(x: Any, y: Any):
+    """Generalized inner product allowing for the first argument to be a tensor"""
+    return wp.dot(x, y)
+
+
+@wp.func
+def generalized_inner(x: float, y: float):
+    return x * y
+
+
+@wp.func
+def generalized_inner(x: wp.mat22, y: wp.vec2):
+    return x[0] * y[0] + x[1] * y[1]
+
+
+@wp.func
+def generalized_inner(x: wp.mat33, y: wp.vec3):
+    return x[0] * y[0] + x[1] * y[1] + x[2] * y[2]
+
+
+@wp.func
+def basis_element(template_type: Any, coord: int):
+    """Returns a instance of `template_type` with a single coordinate set to 1 in the canonical basis"""
+
+    t = type(template_type)(0.0)
+    t[coord] = 1.0
+    return t
+
+
+@wp.func
+def basis_element(template_type: wp.float32, coord: int):
+    return 1.0
+
+
+@wp.func
+def basis_element(template_type: wp.mat22, coord: int):
+    t = wp.mat22(0.0)
+    row = coord // 2
+    col = coord - 2 * row
+    t[row, col] = 1.0
+    return t
+
+
+@wp.func
+def basis_element(template_type: wp.mat33, coord: int):
+    t = wp.mat33(0.0)
+    row = coord // 3
+    col = coord - 3 * row
+    t[row, col] = 1.0
+    return t
+
+
+@wp.func
+def basis_coefficient(val: wp.float32, i: int):
+    return val
+
+
+@wp.func
+def basis_coefficient(val: Any, i: int):
+    return val[i]
+
+
+@wp.func
+def basis_coefficient(val: wp.vec2, i: int, j: int):
+    # treat as row vector
+    return val[j]
+
+
+@wp.func
+def basis_coefficient(val: wp.vec3, i: int, j: int):
+    # treat as row vector
+    return val[j]
+
+
+@wp.func
+def basis_coefficient(val: Any, i: int, j: int):
+    return val[i, j]
+
+
+@wp.func
+def basis_coefficient(template_type: wp.mat33, coord: int):
+    t = wp.mat33(0.0)
+    row = coord // 3
+    col = coord - 3 * row
+    t[row, col] = 1.0
+    return t
+
+
+@wp.func
+def symmetric_part(x: Any):
+    """Symmetric part of a square tensor"""
+    return 0.5 * (x + wp.transpose(x))
+
+
+@wp.func
+def spherical_part(x: wp.mat22):
+    """Spherical part of a square tensor"""
+    return 0.5 * wp.trace(x) * wp.identity(n=2, dtype=float)
+
+
+@wp.func
+def spherical_part(x: wp.mat33):
+    """Spherical part of a square tensor"""
+    return (wp.trace(x) / 3.0) * wp.identity(n=3, dtype=float)
+
+
+@wp.func
+def skew_part(x: wp.mat22):
+    """Skew part of a 2x2 tensor as corresponding rotation angle"""
+    return 0.5 * (x[1, 0] - x[0, 1])
+
+
+@wp.func
+def skew_part(x: wp.mat33):
+    """Skew part of a 3x3 tensor as the corresponding rotation vector"""
+    a = 0.5 * (x[2, 1] - x[1, 2])
+    b = 0.5 * (x[0, 2] - x[2, 0])
+    c = 0.5 * (x[1, 0] - x[0, 1])
+    return wp.vec3(a, b, c)
+
+
+@wp.func
+def householder_qr_decomposition(A: Any):
+    """
+    QR decomposition of a square matrix using Householder reflections
+
+    Returns Q and R such that Q R = A, Q orthonormal (such that QQ^T = Id), R upper triangular
+    """
+
+    x = type(A[0])()
+    Q = wp.identity(n=type(x).length, dtype=A.dtype)
+
+    zero = x.dtype(0.0)
+    two = x.dtype(2.0)
+
+    for i in range(type(x).length):
+        for k in range(type(x).length):
+            x[k] = wp.select(k < i, A[k, i], zero)
+
+        alpha = wp.length(x) * wp.sign(x[i])
+        x[i] += alpha
+        two_over_x_sq = wp.select(alpha == zero, two / wp.length_sq(x), zero)
+
+        A -= wp.outer(two_over_x_sq * x, x * A)
+        Q -= wp.outer(Q * x, two_over_x_sq * x)
+
+    return Q, A
+
+
+@wp.func
+def householder_make_hessenberg(A: Any):
+    """Transforms a square matrix to Hessenberg form (single lower diagonal) using Householder reflections
+
+    Returns:
+        Q and H such that Q H Q^T = A, Q orthonormal, H under Hessenberg form
+        If A is symmetric, H will be tridiagonal
+    """
+
+    x = type(A[0])()
+    Q = wp.identity(n=type(x).length, dtype=A.dtype)
+
+    zero = x.dtype(0.0)
+    two = x.dtype(2.0)
+
+    for i in range(1, type(x).length):
+        for k in range(type(x).length):
+            x[k] = wp.select(k < i, A[k, i - 1], zero)
+
+        alpha = wp.length(x) * wp.sign(x[i])
+        x[i] += alpha
+        two_over_x_sq = wp.select(alpha == zero, two / wp.length_sq(x), zero)
+
+        # apply on both sides
+        A -= wp.outer(two_over_x_sq * x, x * A)
+        A -= wp.outer(A * x, two_over_x_sq * x)
+        Q -= wp.outer(Q * x, two_over_x_sq * x)
+
+    return Q, A
+
+
+@wp.func
+def solve_triangular(R: Any, b: Any):
+    """Solves for R x = b where R is an upper triangular matrix
+
+    Returns x
+    """
+    zero = b.dtype(0)
+    x = type(b)(b.dtype(0))
+    for i in range(b.length, 0, -1):
+        j = i - 1
+        r = b[j] - wp.dot(R[j], x)
+        x[j] = wp.select(R[j, j] == zero, r / R[j, j], zero)
+
+    return x
+
+
+@wp.func
+def inverse_qr(A: Any):
+    # Computes a square matrix inverse using QR factorization
+
+    Q, R = householder_qr_decomposition(A)
+
+    A_inv = type(A)()
+    for i in range(type(A[0]).length):
+        A_inv[i] = solve_triangular(R, Q[i])  # ith column of Q^T
+
+    return wp.transpose(A_inv)
+
+
+@wp.func
+def _wilkinson_shift(a: Any, b: Any, c: Any, tol: Any):
+    # Wilkinson shift: estimate eigenvalue of 2x2 symmetric matrix [a, c, c, b]
+    d = (a - b) * type(tol)(0.5)
+    return b + d - wp.sign(d) * wp.sqrt(d * d + c * c)
+
+
+@wp.func
+def _givens_rotation(a: Any, b: Any):
+    # Givens rotation [[c -s], [s c]] such that sa+cb =0
+    zero = type(a)(0.0)
+    one = type(a)(1.0)
+
+    b2 = b * b
+    if b2 == zero:
+        # id rotation
+        return one, zero
+
+    scale = one / wp.sqrt(a * a + b2)
+    return a * scale, -b * scale
+
+
+@wp.func
+def tridiagonal_symmetric_eigenvalues_qr(D: Any, L: Any, Q: Any, tol: Any):
+    """
+    Computes the eigenvalues and eigen vectors of a symmetric tridiagonal matrix using the
+    Symmetric tridiagonal QR algorithm with implicit Wilkinson shift
+
+    Args:
+        D: Main diagonal of the matrix
+        L: Lower diagonal of the matrix, indexed such that L[i] = A[i+1, i]
+        Q: Initialization for the eigenvectors, useful if a pre-transformation has been applied, otherwise may be identity
+        tol: Tolerance for the diagonalization residual (Linf norm of off-diagonal over diagonal terms)
+
+    Returns a tuple (D: vector of eigenvalues, P: matrix with one eigenvector per row) such that A = P^T D P
+
+
+    Ref: Arbenz P, Numerical Methods for Solving Large Scale Eigenvalue Problems, Chapter 4 (QR algorithm, Mar 13, 2018)
+    """
+
+    two = D.dtype(2.0)
+
+    # so that we can use the type length in expressions
+    # this will prevent unrolling by warp, but should be ok for native code
+    m = int(0)
+    for _ in range(type(D).length):
+        m += 1
+
+    start = int(0)
+    y = D.dtype(0.0)  # moving buldge
+    x = D.dtype(0.0)  # coeff atop buldge
+
+    for _ in range(32 * m):  # failsafe, usually converges faster than that
+        # Iterate over all independent (deflated) blocks
+        end = int(-1)
+
+        for k in range(m - 1):
+            if k >= end:
+                # Check if new block is starting
+                if k == end or wp.abs(L[k]) <= tol * (wp.abs(D[k]) + wp.abs(D[k + 1])):
+                    continue
+
+                # Find end of block
+                start = k
+                end = start + 1
+                while end + 1 < m:
+                    if wp.abs(L[end]) <= tol * (wp.abs(D[end + 1]) + wp.abs(D[end])):
+                        break
+                    end += 1
+
+                # Wilkinson shift (an eigenvalue of the last 2x2 block)
+                shift = _wilkinson_shift(D[end - 1], D[end], L[end - 1], tol)
+
+                # start with eliminating lower diag of first column of shifted matrix
+                # (i.e. first step of explicit QR factorization)
+                # Then all further steps eliminate the buldge (second diag) of the non-shifted matrix
+                x = D[start] - shift
+                y = L[start]
+
+            c, s = _givens_rotation(x, y)
+
+            # Apply Givens rotation on both sides of tridiagonal matrix
+
+            # middle block
+            d = D[k] - D[k + 1]
+            z = (two * c * L[k] + d * s) * s
+            D[k] -= z
+            D[k + 1] += z
+            L[k] = d * c * s + (c * c - s * s) * L[k]
+
+            if k > start:
+                L[k - 1] = c * x - s * y
+
+            x = L[k]
+            y = -s * L[k + 1]  # new buldge
+            L[k + 1] *= c
+
+            # apply givens rotation on left of Q
+            # note: Q is transposed compared to usual impls, as Warp makes it easier to index rows
+            Qk0 = Q[k]
+            Qk1 = Q[k + 1]
+            Q[k] = c * Qk0 - s * Qk1
+            Q[k + 1] = c * Qk1 + s * Qk0
+
+        if end <= 0:
+            # We did nothing, so diagonalization must have been achieved
+            break
+
+    return D, Q
+
+
+@wp.func
+def symmetric_eigenvalues_qr(A: Any, tol: Any):
+    """
+    Computes the eigenvalues and eigen vectors of a square symmetric matrix A using the QR algorithm
+
+    Args:
+        A: square symmetric matrix
+        tol: Tolerance for the diagonalization residual (Linf norm of off-diagonal over diagonal terms)
+
+    Returns a tuple (D: vector of eigenvalues, P: matrix with one eigenvector per row) such that A = P^T D P
+    """
+
+    # Put A under Hessenberg form (tridiagonal)
+    Q, H = householder_make_hessenberg(A)
+
+    # tridiagonal storage for H
+    D = wp.get_diag(H)
+    L = type(D)(A.dtype(0.0))
+    for i in range(1, type(D).length):
+        L[i - 1] = H[i, i - 1]
+
+    Qt = wp.transpose(Q)
+    ev, P = tridiagonal_symmetric_eigenvalues_qr(D, L, Qt, tol)
+    return ev, P
+
+
+def array_axpy(x: wp.array, y: wp.array, alpha: float = 1.0, beta: float = 1.0):
+    """Performs y = alpha*x + beta*y"""
+
+    dtype = wp.types.type_scalar_type(y.dtype)
+
+    alpha = dtype(alpha)
+    beta = dtype(beta)
+
+    if x.shape != y.shape or x.device != y.device:
+        raise ValueError("x and y arrays must have the same shape and device")
+
+    # array_axpy requires a custom adjoint; unfortunately we cannot use `wp.func_grad`
+    # as generic functions are not supported yet. Instead we use a non-differentiable kernel
+    # and record a custom adjoint function on the tape.
+
+    # temporarily disable tape to avoid printing warning that kernel is not differentiable
+    (tape, wp.context.runtime.tape) = (wp.context.runtime.tape, None)
+    wp.launch(kernel=_array_axpy_kernel, dim=x.shape, device=x.device, inputs=[x, y, alpha, beta])
+    wp.context.runtime.tape = tape
+
+    if tape is not None and (x.requires_grad or y.requires_grad):
+
+        def backward_axpy():
+            # adj_x += adj_y * alpha
+            # adj_y = adj_y * beta
+            array_axpy(x=y.grad, y=x.grad, alpha=alpha, beta=1.0)
+            if beta != 1.0:
+                array_axpy(x=y.grad, y=y.grad, alpha=0.0, beta=beta)
+
+        tape.record_func(backward_axpy, arrays=[x, y])
+
+
+@wp.kernel(enable_backward=False)
+def _array_axpy_kernel(x: wp.array(dtype=Any), y: wp.array(dtype=Any), alpha: Any, beta: Any):
+    i = wp.tid()
+    y[i] = beta * y[i] + alpha * y.dtype(x[i])

warp/fem/operator.py

@@ -1,8 +1,8 @@
-from typing import Any, Callable
+from typing import Any, Callable, Dict, Optional, Set
 
 import warp as wp
-from warp.fem import utils
-from warp.fem.types import Domain, Field, NodeIndex, Sample
+from warp.fem.linalg import skew_part, symmetric_part
+from warp.fem.types import Coords, Domain, ElementIndex, Field, NodeIndex, Sample, make_free_sample
 
 
 class Integrand:
@@ -10,35 +10,55 @@ class Integrand:
     It will get transformed to a proper warp.Function by resolving concrete Field types at call time.
     """
 
-    def __init__(self, func: Callable):
+    def __init__(self, func: Callable, kernel_options: Optional[Dict[str, Any]] = None):
         self.func = func
         self.name = wp.codegen.make_full_qualified_name(self.func)
         self.module = wp.get_module(self.func.__module__)
         self.argspec = wp.codegen.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: Dict[str, Set[Operator]] = None
 
 
 class Operator:
     """
-    Operators provide syntaxic sugar over Field and Domain evaluation functions and arguments
+    Operators provide syntactic sugar over Field and Domain evaluation functions and arguments
     """
 
-    def __init__(self, func: Callable, resolver: Callable):
+    def __init__(self, func: Callable, resolver: Callable, field_result: Callable = None):
         self.func = func
+        self.name = func.__name__
         self.resolver = resolver
+        self.field_result = field_result
+
+
+def integrand(func: 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.
+    """
 
-def integrand(func: Callable):
-    """Decorator for functions to be integrated (or interpolated) using warp.fem"""
-    itg = Integrand(func)
-    itg.__doc__ = func.__doc__
-    return itg
+    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
 
-def operator(resolver: Callable):
+    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, resolver)
+        op = Operator(func, **kwargs)
         op.__doc__ = func.__doc__
         return op
 
@@ -56,7 +76,7 @@ def position(domain: Domain, s: Sample):
 
 @operator(resolver=lambda dmn: dmn.element_normal)
 def normal(domain: Domain, s: Sample):
-    """Evaluates the element normal at the sample point `s`. Null for interior points."""
+    """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
 
 
@@ -75,7 +95,7 @@ def lookup(domain: Domain, x: Any) -> Sample:
         guess: (optional) :class:`Sample` initial guess, may help perform the query
 
     Note:
-        Currently this operator is unsupported for :class:`Hexmesh`, :class:`Quadmesh2D` and deformed geometries.
+        Currently this operator is unsupported for :class:`Hexmesh`, :class:`Quadmesh2D`, :class:`Quadmesh3D` and deformed geometries.
     """
     pass
 
@@ -88,10 +108,73 @@ def measure(domain: Domain, s: Sample) -> float:
 
 @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."""
+    """Returns the maximum ratio between the measure of this element and that of higher-dimensional neighbors."""
+    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)
+)
+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)
+def _inner_cell_index(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
+    pass
+
+
+@operator(resolver=lambda dmn: dmn.element_outer_cell_index)
+def _outer_cell_index(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
     pass
 
 
+@operator(resolver=lambda dmn: dmn.element_inner_cell_coords)
+def _inner_cell_coords(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
+    pass
+
+
+@operator(resolver=lambda dmn: dmn.element_outer_cell_coords)
+def _outer_cell_coords(domain: Domain, side_index: ElementIndex, side_coords: Coords) -> Sample:
+    pass
+
+
+@operator(resolver=lambda dmn: dmn.cell_to_element_coords)
+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)
+    )
+
+
 # Field operators
 # On a side, inner and outer are such that normal goes from inner to outer
 
@@ -157,13 +240,13 @@ def node_partition_index(f: Field, node_index: NodeIndex):
 @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))
+    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 utils.skew_part(grad(f, s))
+    return skew_part(grad(f, s))
 
 
 @integrand

warp/fem/polynomial.py

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