anypinn.lib.diff

Differential operators for Physics-Informed Neural Networks.

Composable utilities built on torch.autograd.grad for computing first-order, higher-order, and mixed partial derivatives without re-implementing autograd boilerplate in every constraint.

All operators default to create_graph=True so their outputs are differentiable — required when used inside loss functions that must back-propagate. Pass create_graph=False for detached results (e.g. visualisation or adaptive sampling).

Note

For even higher performance on large-batch Hessians or Jacobians, torch.func.jacrev / torch.func.hessian can be composed with torch.vmap. The operators here intentionally use the simpler autograd.grad path for broad compatibility and torch.compile friendliness.

divergence(v: Tensor, x: Tensor, *, create_graph: bool = True) -> Tensor

Compute the divergence \(\nabla \cdot v = \sum_i \partial v_i / \partial x_i\).

Parameters:

Name	Type	Description	Default
v	Tensor	Vector field values, shape (n, d) matching x.shape[1].	required
x	Tensor	Input coordinates, shape (n, d) with requires_grad=True.	required
create_graph	bool	Keep the result in the computation graph (default True).	True

Returns:

Type	Description
Tensor	Tensor of shape (n, 1).

Source code in src/anypinn/lib/diff.py

def divergence(
    v: Tensor,
    x: Tensor,
    *,
    create_graph: bool = True,
) -> Tensor:
    """Compute the divergence $\\nabla \\cdot v = \\sum_i \\partial v_i / \\partial x_i$.

    Args:
        v: Vector field values, shape ``(n, d)`` matching ``x.shape[1]``.
        x: Input coordinates, shape ``(n, d)`` with ``requires_grad=True``.
        create_graph: Keep the result in the computation graph (default ``True``).

    Returns:
        Tensor of shape ``(n, 1)``.
    """
    ndim = x.shape[1]
    if v.shape[1] != ndim:
        raise ValueError(
            f"v has {v.shape[1]} components but x has {ndim} dimensions; they must match"
        )

    div = torch.zeros(x.shape[0], 1, device=x.device, dtype=x.dtype)
    for i in range(ndim):
        dvi_dxi = grad(v[:, i : i + 1], x, create_graph=True)[:, i : i + 1]
        div = div + dvi_dxi

    if not create_graph:
        div = div.detach()
    return div

grad(u: Tensor, x: Tensor, *, create_graph: bool = True) -> Tensor

Compute the full gradient \(\nabla u\) with respect to coordinates \(x\).

Parameters:

Name	Type	Description	Default
u	Tensor	Scalar field values, shape (n,) or (n, 1).	required
x	Tensor	Input coordinates, shape (n, d) with requires_grad=True.	required
create_graph	bool	Keep the result in the computation graph (default True).	True

Returns:

Type	Description
Tensor	Tensor of shape (n, d) with rows
Tensor	\([\partial u/\partial x_0, \ldots, \partial u/\partial x_{d-1}]\).

Source code in src/anypinn/lib/diff.py

def grad(
    u: Tensor,
    x: Tensor,
    *,
    create_graph: bool = True,
) -> Tensor:
    """Compute the full gradient $\\nabla u$ with respect to coordinates $x$.

    Args:
        u: Scalar field values, shape ``(n,)`` or ``(n, 1)``.
        x: Input coordinates, shape ``(n, d)`` with ``requires_grad=True``.
        create_graph: Keep the result in the computation graph (default ``True``).

    Returns:
        Tensor of shape ``(n, d)`` with rows
        $[\\partial u/\\partial x_0, \\ldots, \\partial u/\\partial x_{d-1}]$.
    """
    (grad_u,) = torch.autograd.grad(
        u.reshape(-1).sum(),
        x,
        create_graph=create_graph,
    )
    return grad_u

hessian(u: Tensor, x: Tensor, *, create_graph: bool = True) -> Tensor

Compute the Hessian matrix \(H[u]\) where \(H_{ij} = \partial^2 u / (\partial x_i\partial x_j)\).

Computes the first-order gradient once, then differentiates each component to build each row — d + 1 autograd calls total.

Parameters:

Name	Type	Description	Default
u	Tensor	Scalar field values, shape (n,) or (n, 1).	required
x	Tensor	Input coordinates, shape (n, d) with requires_grad=True.	required
create_graph	bool	Keep the result in the computation graph (default True).	True

Returns:

Type	Description
Tensor	Tensor of shape (n, d, d).

Source code in src/anypinn/lib/diff.py

def hessian(
    u: Tensor,
    x: Tensor,
    *,
    create_graph: bool = True,
) -> Tensor:
    """Compute the Hessian matrix $H[u]$ where
    $H_{ij} = \\partial^2 u / (\\partial x_i\\partial x_j)$.

    Computes the first-order gradient once, then differentiates each
    component to build each row — ``d + 1`` autograd calls total.

    Args:
        u: Scalar field values, shape ``(n,)`` or ``(n, 1)``.
        x: Input coordinates, shape ``(n, d)`` with ``requires_grad=True``.
        create_graph: Keep the result in the computation graph (default ``True``).

    Returns:
        Tensor of shape ``(n, d, d)``.
    """
    grad_u = grad(u, x, create_graph=True)
    ndim = x.shape[1]
    rows = []
    for i in range(ndim):
        row = grad(grad_u[:, i : i + 1], x, create_graph=True)
        rows.append(row)

    H = torch.stack(rows, dim=1)

    if not create_graph:
        H = H.detach()
    return H

laplacian(u: Tensor, x: Tensor, *, create_graph: bool = True) -> Tensor

Compute the Laplacian \(\nabla^2 u = \sum_i \partial^2 u / \partial x_i^2\).

Computes the full first-order gradient once and then differentiates each component — d + 1 autograd calls total for d dimensions.

Parameters:

Name	Type	Description	Default
u	Tensor	Scalar field values, shape (n,) or (n, 1).	required
x	Tensor	Input coordinates, shape (n, d) with requires_grad=True.	required
create_graph	bool	Keep the result in the computation graph (default True).	True

Returns:

Type	Description
Tensor	Tensor of shape (n, 1).

Source code in src/anypinn/lib/diff.py

def laplacian(
    u: Tensor,
    x: Tensor,
    *,
    create_graph: bool = True,
) -> Tensor:
    """Compute the Laplacian $\\nabla^2 u = \\sum_i \\partial^2 u / \\partial x_i^2$.

    Computes the full first-order gradient once and then differentiates
    each component — ``d + 1`` autograd calls total for ``d`` dimensions.

    Args:
        u: Scalar field values, shape ``(n,)`` or ``(n, 1)``.
        x: Input coordinates, shape ``(n, d)`` with ``requires_grad=True``.
        create_graph: Keep the result in the computation graph (default ``True``).

    Returns:
        Tensor of shape ``(n, 1)``.
    """
    grad_u = grad(u, x, create_graph=True)
    ndim = x.shape[1]
    lap = torch.zeros(x.shape[0], 1, device=x.device, dtype=x.dtype)
    for i in range(ndim):
        d2u_dxi2 = grad(grad_u[:, i : i + 1], x, create_graph=True)[:, i : i + 1]
        lap = lap + d2u_dxi2

    if not create_graph:
        lap = lap.detach()
    return lap

mixed_partial(u: Tensor, x: Tensor, dims: tuple[int, ...], *, create_graph: bool = True) -> Tensor

Compute a mixed derivative \(\partial^k u / (\partial x_{d_0} \partial x_{d_1} \cdots)\).

Derivatives are applied left-to-right: first differentiate w.r.t. dims[0], then the result w.r.t. dims[1], and so on.

Parameters:

Name	Type	Description	Default
u	Tensor	Scalar field values, shape (n,) or (n, 1).	required
x	Tensor	Input coordinates, shape (n, d) with requires_grad=True.	required
dims	tuple[int, ...]	Dimension indices to differentiate along, in order.	required
create_graph	bool	Keep the result in the computation graph (default True).	True

Returns:

Type	Description
Tensor	Tensor of shape (n, 1).

Source code in src/anypinn/lib/diff.py

def mixed_partial(
    u: Tensor,
    x: Tensor,
    dims: tuple[int, ...],
    *,
    create_graph: bool = True,
) -> Tensor:
    """Compute a mixed derivative $\\partial^k u / (\\partial x_{d_0} \\partial x_{d_1} \\cdots)$.

    Derivatives are applied left-to-right: first differentiate w.r.t.
    ``dims[0]``, then the result w.r.t. ``dims[1]``, and so on.

    Args:
        u: Scalar field values, shape ``(n,)`` or ``(n, 1)``.
        x: Input coordinates, shape ``(n, d)`` with ``requires_grad=True``.
        dims: Dimension indices to differentiate along, in order.
        create_graph: Keep the result in the computation graph (default ``True``).

    Returns:
        Tensor of shape ``(n, 1)``.
    """
    if len(dims) == 0:
        raise ValueError("dims must contain at least one dimension index")

    result = grad(u, x, create_graph=True)[:, dims[0] : dims[0] + 1]
    for d in dims[1:]:
        result = grad(result, x, create_graph=True)[:, d : d + 1]

    if not create_graph:
        result = result.detach()
    return result

partial(u: Tensor, x: Tensor, dim: int, *, order: int = 1, create_graph: bool = True) -> Tensor

Compute the order-k derivative \(\partial^k u / \partial x_d^k\) along one dimension.

Parameters:

Name	Type	Description	Default
u	Tensor	Scalar field values, shape (n,) or (n, 1).	required
x	Tensor	Input coordinates, shape (n, d) with requires_grad=True.	required
dim	int	Spatial dimension index to differentiate along.	required
order	int	Derivative order (≥ 1, default 1).	1
create_graph	bool	Keep the result in the computation graph (default True).	True

Returns:

Type	Description
Tensor	Tensor of shape (n, 1).

Source code in src/anypinn/lib/diff.py

def partial(
    u: Tensor,
    x: Tensor,
    dim: int,
    *,
    order: int = 1,
    create_graph: bool = True,
) -> Tensor:
    """Compute the order-k derivative $\\partial^k u / \\partial x_d^k$ along one dimension.

    Args:
        u: Scalar field values, shape ``(n,)`` or ``(n, 1)``.
        x: Input coordinates, shape ``(n, d)`` with ``requires_grad=True``.
        dim: Spatial dimension index to differentiate along.
        order: Derivative order (≥ 1, default 1).
        create_graph: Keep the result in the computation graph (default ``True``).

    Returns:
        Tensor of shape ``(n, 1)``.
    """
    if order < 1:
        raise ValueError(f"order must be >= 1, got {order}")

    result = grad(u, x, create_graph=True)[:, dim : dim + 1]
    for _ in range(order - 1):
        result = grad(result, x, create_graph=True)[:, dim : dim + 1]

    if not create_graph:
        result = result.detach()
    return result