tensormesh.assemble¶

Element Assembler¶

class ElementAssembler(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, edges: Tensor, *args, **kwargs)[source]¶

Bases: Module

Assemble an element-wise bilinear form into a global sparse matrix.

ElementAssembler is an torch.nn.Module whose forward() defines the integrand of a bilinear form \(a(u, v) = \int_\Omega f(u, v)\, \mathrm{d}\Omega\). Calling the assembler integrates \(f\) over every quadrature point of every element and scatters the result into a SparseMatrix of shape \([|\mathcal V|, |\mathcal V|]\) (scalar problems) or \([|\mathcal V| \times H, |\mathcal V| \times H]\) (vector problems with \(H\) degrees of freedom per node).

Subclasses are usually built from a mesh:

from_mesh() — build from a Mesh (slower; precomputes the projection tensor \(\mathcal P_{\mathcal E}\)).
from_assembler() — share the topology of another assembler (fast).

The schematic of one __call__ is

\[K \overset{\text{bsr matrix}}{\leftarrow}\hat K_\text{global}, \qquad \hat K_{\text{global}}^{nkl} = \mathcal P_{\mathcal E}^{nhij}\, \hat K_{\text{local}}^{hklij}, \qquad \hat K_{ij} = \int_\Omega f(u, v)\, \mathrm{d}\Omega,\]

where

\(\hat K_{\text{global}}\) are the non-zero entries of the global Galerkin matrix, \(K_{\text{global}} \in \mathbb R^{|\mathcal E|\times d \times d}\);
\(\hat K_{\text{local}}\) is the per-element Galerkin block, \(K_{\text{local}} \in \mathbb R^{|\mathcal C|\times h \times h \times d \times d}\);
\(\mathcal P_{\mathcal E} \in \mathbb R_{\text{sparse}}^{|\mathcal E|\times|\mathcal C|\times h \times h}\) is the projection (assemble) tensor from local to global;
\(\mathcal C\) indexes elements/cells, \(\mathcal E\) indexes edges (the unique (i, j) pairs of basis indices), \(\mathcal V\) indexes nodes/vertices, and \(h\) is the number of basis functions per element.

Examples

Mass matrix \(M_{ij} = \int_\Omega u_i v_j\, \mathrm{d}\Omega\):

import tensormesh
class MassAssembler(tensormesh.ElementAssembler):
    def forward(self, u, v):
        return u * v
mesh = tensormesh.Mesh.gen_rectangle()
assembler = MassAssembler.from_mesh(mesh)
M = assembler(mesh.points)

Laplace stiffness \(K_{ij} = \int_\Omega \nabla u_i \cdot \nabla v_j\, \mathrm{d}\Omega\):

import tensormesh
import tensormesh.functional as F
class LaplaceAssembler(tensormesh.ElementAssembler):
    def forward(self, gradu, gradv):
        return F.dot(gradu, gradv)
mesh = tensormesh.Mesh.gen_circle()
assembler = LaplaceAssembler.from_mesh(mesh)
K = assembler(mesh.points)

projector¶

Maps each element_type to a Projector that scatters per-element contributions of shape \([|\mathcal C_e|, B_e, B_e]\) into the global edge vector of shape \([|\mathcal E|]\).

Type:: ModuleDict

transformation¶

Maps each element_type to a Transformation caching per-element Jacobians, shape values, shape gradients, and JxW at quadrature points.

Type:: ModuleDict

elements¶

Maps each element_type to its connectivity tensor of shape \([|\mathcal C|, B]\), e.g. {"triangle6": tensor([[0, 1, 2], ...])}.

Type:: BufferDict

edges¶

Long tensor of shape \([2, |\mathcal E|]\) listing the unique (row, col) index pairs of the assembled sparse matrix; produced by deduplicating all per-element basis-pair indices.

Type:: Tensor

n_points¶

Number of mesh points (size of one matrix dimension).

Type:: int

dimension¶

Spatial dimension of the mesh, one of 1, 2, 3.

Type:: int

element_types¶

Element-type strings present in the mesh, e.g. ["triangle6", "quad9"].

Type:: list[str]

__init__(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, edges: Tensor, *args, **kwargs)[source]¶: Initialize internal Module state, shared by both nn.Module and ScriptModule.

property device: device¶: Device on which the assembler’s buffers live.

property dtype: dtype¶: Floating dtype of the assembler’s buffers (float32 or float64).

type(dtype: dtype)[source]¶

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

Calculates the total potential energy of the system by integrating the element energy density over the domain.

This method is designed for energy-based variational problems (e.g., hyperelasticity, phase field, optimization). It computes the global integral:

\[E = \int_\Omega \psi(\mathbf{u}, \nabla \mathbf{u}, \dots) d\Omega\]

where \(\psi\) is the energy density function (defined in element_energy() or passed as func). The method handles efficient parallel integration using quadrature rules and PyTorch’s vmap.

How it works (Simplified):

You provide nodal data (like displacement u) in point_data.
You define an element_energy function that takes arguments like u (value) or graddisplacement (gradient) and returns the scalar energy density for a single quadrature point.
This method automatically matches your arguments, broadcasts them over all elements and quadrature points, computes the gradients if requested, and performs the numerical integration.

Parameters:

points (Tensor, optional) – The current positions of the nodes (shape: [N_nodes, Dim]). If not provided, uses the mesh’s initial points.
func (Callable, optional) –
A function that computes energy density at a single quadrature point. Signature: func(arg1, arg2, …) -> torch.Tensor (scalar). If None, uses self.element_energy.

Supported Arguments for `func`:
- `u`, `v`, … : Values interpolated from point_data (e.g., passing point_data={‘u’: …} allows requesting u).
- `gradu`, `gradv`, …: Gradients interpolated from point_data (e.g., passing point_data={‘u’: …} allows requesting gradu).
- `key` in element_data: Element-wise constant or quadrature-varying data.
- `key` in scalar_data: Global scalar constants.
point_data (Dict[str, Tensor], optional) – Nodal fields (e.g., displacement, phase field). Shape: [N_nodes, …] or [N_nodes, Dim].
element_data (Dict[str, Tensor] or Dict[str, Dict[str, Tensor]], optional) –
Data defined on elements. Can be:
- Constant per element: Shape [N_elements, …].
- Varying per quadrature point (e.g., history variables): Shape [N_elements, N_quad, …].
scalar_data (Dict[str, Tensor], optional) – Global constants (e.g., time step dt).
batch_size (int, optional) – Batch size for processing quadrature points to save memory. Default is -1 (process all at once).

Returns:

The total scalar energy (shape: []). Since this operation is differentiable, you can call .backward() on it to compute forces (gradients w.r.t points or point_data).

Return type:

Tensor

Examples

1. Hyperelasticity (Neo-Hookean Energy)

class NeoHookean(ElementAssembler):
    def element_energy(self, graddisplacement):
        # graddisplacement is [Dim, Dim] tensor at one quad point
        F = torch.eye(3) + graddisplacement
        J = torch.det(F)
        # ... compute psi ...
        return psi

model = NeoHookean.from_mesh(mesh)

# Compute energy
# Auto-differentiable w.r.t u
u = torch.zeros_like(mesh.points, requires_grad=True)
E = model.energy(point_data={"displacement": u})

# Compute Forces (Internal Force Vector)
E.backward()
F_int = u.grad

element_energy(**kwargs)[source]¶

Override this method to define the energy density at a single quadrature point.

This method is called automatically by energy() using vmap to parallelize over all quadrature points.

Parameters:

**kwargs (Tensor) –

Arguments matching the variable names requested. Common arguments:

u, v: Value of field at quadrature point.
gradu, gradv: Gradient of field at quadrature point.
element_data_key: Value of element data (constant or interpolated).

Returns:

Scalar energy density (\(\psi\)) for this quadrature point.

Return type:

Tensor

abstract forward(**kwargs)[source]¶

Define the integrand of the bilinear form at a single quadrature point.

Subclasses must override this method. The library uses torch.vmap() to lift the per-quadrature-point function over all quadrature points and all elements, so write it as if you were evaluating at one point. Parameters are dispatched by name; return values are integrated against JxW and scattered by the caller (see ElementAssembler.__call__).

Parameters:

u (Tensor, optional) – Shape value at the quadrature point — 0D tensor of shape [].
v (Tensor, optional) – Shape value at the quadrature point — 0D tensor of shape [].
gradu (Tensor, optional) – Shape gradient in physical coordinates — 1D tensor of shape [D].
gradv (Tensor, optional) – Shape gradient in physical coordinates — 1D tensor of shape [D].
x (Tensor, optional) – Physical coordinate at the quadrature point — 1D tensor of shape [D].
gradx (Tensor, optional) – Gradient of x w.r.t. reference coordinates — 2D tensor of shape [D, D].
**point_data (Tensor) – Any key passed to __call__ via point_data: if the nodal tensor has shape \([|\mathcal V|, ...]\), the value handed to forward has the trailing [...] shape, and its counterpart "grad"+key has shape [..., D].

Returns:

Either a 2D tensor of shape [B, B] (scalar problems) or a 4D tensor of shape [B, B, H, H] (vector problems with H degrees of freedom per node).

Return type:

Tensor

classmethod from_assembler(obj, *args, **kwargs)[source]¶

Build an ElementAssembler that shares topology with obj.

Much faster than from_mesh() since the projector and the cached Transformation are reused as-is.

Parameters:

obj (ElementAssembler) – An existing assembler whose mesh topology should be reused.
*args – Additional arguments forwarded to __post_init__.
**kwargs – Additional arguments forwarded to __post_init__.

Returns:

A new assembler sharing the same mesh.

Return type:

ElementAssembler

classmethod from_mesh(mesh: Mesh, quadrature_order: int = 2, project: str = 'reduce', *args, **kwargs)[source]¶

Build an ElementAssembler from a Mesh.

Slower than from_assembler() because the projection tensor \(\mathcal P_{\mathcal E}\) is precomputed from the connectivity.

Parameters:

mesh (Mesh) – Source mesh; both connectivity and points are taken from it.
quadrature_order (int, optional) – Positive integer; defaults to 2.
project ({'reduce', 'sparse'}, optional) – Backend used for the element-to-edge scatter; defaults to 'reduce'.
*args – Additional arguments forwarded to __post_init__.
**kwargs – Additional arguments forwarded to __post_init__.

Returns:

A new assembler that owns the mesh topology.

Return type:

ElementAssembler

Facet Assembler¶

class FacetAssembler(facet_mask: ModuleDict, projector: ModuleDict, transformation: ModuleDict, *args, **kwargs)[source]¶

Bases: Module

Assemble an integrand over boundary facets of a mesh.

FacetAssembler mirrors NodeAssembler but integrates over \(\partial \Omega\) instead of \(\Omega\). Override forward() to define a per-quadrature-point integrand; calling the assembler returns a flattened tensor of shape \([|\mathcal V|]\) or \([|\mathcal V| \times H]\) (vector-valued problems with \(H\) DOFs per node).

Typical uses include Neumann tractions, penalty contact, surface tension, and Robin boundary conditions.

Examples

Constant downward traction on the boundary:

import torch
from tensormesh import Mesh, FacetAssembler

class TractionAssembler(FacetAssembler):
    def forward(self, v):
        t = torch.tensor([0.0, -1.0], dtype=v.dtype, device=v.device)
        return t * v        # contribution at one quadrature point

mesh = Mesh.gen_rectangle()
f = TractionAssembler.from_mesh(mesh)(mesh.points)

projector¶

Maps each element_type to a Projector that scatters per-facet basis contributions onto the node vector.

Type:: ModuleDict

facet_mask¶

Maps each element_type to a BufferList of boolean masks marking which facets of which elements lie on the selected boundary (one mask per facet type for mixed-facet shapes, otherwise a list of one).

Type:: ModuleDict

transformation¶

Maps each element_type to its cached Transformation, providing facet_shape_val, facet_shape_grad, and FxW.

Type:: ModuleDict

n_points¶

Number of mesh points (length of the output vector for scalar problems).

Type:: int

dimension¶

Spatial dimension of the mesh, one of 1, 2, 3.

Type:: int

element_types¶

Element-type strings present in the mesh.

Type:: list[str]

__init__(facet_mask: ModuleDict, projector: ModuleDict, transformation: ModuleDict, *args, **kwargs)[source]¶: Initialize internal Module state, shared by both nn.Module and ScriptModule.

property device: device¶: Device on which the assembler’s buffers live.

property dtype: dtype¶: Floating dtype of the assembler’s buffers (float32 or float64).

type(dtype: dtype)[source]¶

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

abstract forward(*args)[source]¶

Define the facet integrand at a single quadrature point.

Subclasses must override this method. Vmap dispatches the per-quadrature-point function over all selected facets, so write it as if evaluating at one facet quadrature point. Unlike ElementAssembler.forward, the basis arguments here keep the basis dimension (the inner vmap(vmap(...)) covers facet + quadrature only, not basis).

Parameters:

u (Tensor, optional) – Shape value on the facet — 1D tensor of shape [B].
v (Tensor, optional) – Shape value on the facet — 1D tensor of shape [B].
gradu (Tensor, optional) – Shape gradient in physical coordinates — 2D tensor of shape [B, D].
gradv (Tensor, optional) – Shape gradient in physical coordinates — 2D tensor of shape [B, D].
x (Tensor, optional) – Physical coordinate at the quadrature point — 1D tensor of shape [D].
gradx (Tensor, optional) – Gradient of x w.r.t. reference coordinates — 2D tensor of shape [D, D].
**point_data (Tensor) – Any key passed to __call__ via point_data: if the nodal tensor has shape \([|\mathcal V|, ...]\), the value handed to forward has the trailing [...] shape, and its counterpart "grad"+key has shape [..., D].

Returns:

1D tensor of shape [B] (scalar problems) or 2D tensor of shape [B, H] (vector problems with H degrees of freedom per node).

Return type:

Tensor

classmethod from_assembler(obj, *args, **kwargs)[source]¶

Build a FacetAssembler sharing topology with obj.

Much faster than from_mesh() since the facet mask, projector, and cached Transformation are reused as-is.

Parameters:

obj (FacetAssembler) – An existing facet assembler whose boundary topology should be reused.
*args – Additional arguments forwarded to __post_init__.
**kwargs – Additional arguments forwarded to __post_init__.

Returns:

A new assembler sharing the same boundary topology.

Return type:

FacetAssembler

classmethod from_elements(points: Tensor, elements: Dict[str, Tensor], boundary_mask: Tensor, quadrature_order: int = 2, device: str | device = 'cpu', dtype: dtype = torch.float32, project: str = 'reduce', *args, **kwargs)[source]¶

Build a FacetAssembler from raw connectivity tensors.

Slower than from_assembler() because the boundary topology is rebuilt from scratch.

Parameters:

points (Tensor) – 2D tensor of shape \([|\mathcal V|, D]\) listing node coordinates.
elements (dict[str, Tensor]) – Connectivity keyed by element-type string, e.g. {"triangle": tensor([[0, 1, 2], [1, 2, 3]])}.
boundary_mask (Tensor) – 1D boolean tensor of shape \([|\mathcal V|]\) marking which nodes lie on the boundary; a facet is selected iff all of its corner nodes are flagged.
quadrature_order (int, optional) – Positive integer; defaults to 2.
device (device or str, optional) – Device of the assembler; defaults to "cpu".
dtype (dtype, optional) – Floating dtype; defaults to torch.float32.
project ({'reduce', 'sparse'}, optional) – Projection backend; defaults to "reduce".

Returns:

A new assembler that owns the given boundary topology.

Return type:

FacetAssembler

classmethod from_mesh(mesh: Mesh, boundary_mask: str | Tensor | None = None, quadrature_order: int = 2, project: str = 'reduce', *args, **kwargs)[source]¶

Build a FacetAssembler from a Mesh.

Slower than from_assembler() because the boundary topology is rebuilt from connectivity.

Parameters:

mesh (Mesh) – Source mesh; connectivity, points, and (default) boundary mask are read from it.
boundary_mask (str, Tensor, or None, optional) – Boundary selector. None (default) uses mesh.boundary_mask; str keys into mesh.point_data; a tensor is used verbatim and must be 1D boolean of length n_points.
quadrature_order (int, optional) – Positive integer; defaults to 2.
project ({'reduce', 'sparse'}, optional) – Projection backend; defaults to "reduce".

Returns:

A new assembler that owns the boundary topology of the mesh.

Return type:

FacetAssembler

Node Assembler¶

class NodeAssembler(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, *args, **kwargs)[source]¶

Bases: Module

Assemble an element-wise linear form into a global node vector.

NodeAssembler is the linear-form counterpart of ElementAssembler. Override forward() to define the integrand \(l(v) = \int_\Omega f(v)\, \mathrm{d}\Omega\); calling the assembler returns a 1-D torch.Tensor of shape \([|\mathcal V|]\), or \([|\mathcal V| \times H]\) for vector-valued problems with \(H\) degrees of freedom per node.

Subclasses are usually built from a mesh:

from_mesh() — build from a Mesh.
from_elements() — build from raw connectivity tensors.
from_assembler() — share topology with another assembler.

Examples

Load vector \(f_i = \int_\Omega v_i\, \mathrm{d}\Omega\):

import tensormesh
class OneAssembler(tensormesh.NodeAssembler):
    def forward(self, v):
        return v
mesh = tensormesh.Mesh.gen_rectangle()
f = OneAssembler.from_mesh(mesh)(mesh.points)

Traction-style load \(f_i = \int_\Omega \mathbf t \cdot v_i\, \mathrm{d}\Omega\):

import tensormesh
import tensormesh.functional as F
class TractionAssembler(tensormesh.NodeAssembler):
    def forward(self, v, t):
        return F.dot(t, v)
mesh = tensormesh.Mesh.gen_circle()
t = torch.ones(mesh.n_points, 2)  # unit traction in x, y
assembler = TractionAssembler.from_mesh(mesh)
f = assembler(mesh.points, point_data={"t": t})

projector¶

Maps each element_type to a Projector that scatters per-element basis contributions of shape \([|\mathcal C_e|, B_e]\) onto the node vector of shape \([|\mathcal V|]\).

Type:: ModuleDict

transformation¶

Maps each element_type to a Transformation caching shape values, shape gradients, and JxW at quadrature points.

Type:: ModuleDict

elements¶

Maps each element_type to its connectivity tensor of shape \([|\mathcal C|, B]\).

Type:: BufferDict

n_points¶

Number of mesh points (length of the output vector for scalar problems).

Type:: int

dimension¶

Spatial dimension of the mesh, one of 1, 2, 3.

Type:: int

element_types¶

Element-type strings present in the mesh.

Type:: list[str]

__init__(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, *args, **kwargs)[source]¶: Initialize internal Module state, shared by both nn.Module and ScriptModule.

property device: device¶: Device on which the assembler’s buffers live.

property dtype: dtype¶: Floating dtype of the assembler’s buffers (float32 or float64).

type(dtype: dtype)[source]¶

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:: dst_type (type or string): the desired type
Returns:: Module: self

compile(mode: Literal['default', 'reduce-overhead', 'max-autotune', 'max-autotune-no-cudagraphs', 'disable'] = 'disable', fullgraph: bool = False, dynamic: bool | None = None, backend: str = 'inductor', **kwargs) → NodeAssembler[source]¶

Enable fast mode for the assembler to speed up computation.

When compile mode is enabled, the __call__ method bypasses vmap and uses direct broadcast operations, achieving up to 5-30x speedup.

By default (mode="disable"), only the vmap bypass is enabled without torch.compile. This provides the best performance for most cases. Set mode="default" or other modes to additionally enable torch.compile.

Examples

# Enable fast mode (recommended, no torch.compile overhead)
assembler = MassAssembler.from_mesh(mesh).compile()

# Enable with torch.compile for potential additional optimization
assembler = MassAssembler.from_mesh(mesh).compile(mode="default")

# Use normally - automatically uses fast path
result = assembler(point_data={'phi': phi, 'f': f})

# Disable for debugging (can set breakpoints in forward)
assembler.reset_compile()

Parameters:

mode (str, optional) –
Compilation mode, one of:
- "disable": Only bypass vmap, no torch.compile (fastest startup, recommended)
- "default": Also use torch.compile with default settings
- "reduce-overhead": torch.compile with reduced Python overhead
- "max-autotune": torch.compile with maximum optimization
- "max-autotune-no-cudagraphs": Like max-autotune but without CUDA graphs
Default is "disable"
fullgraph (bool, optional) – Whether to compile the entire graph. Default is False
dynamic (bool or None, optional) – Whether to use dynamic shapes. Default is None (auto-detect)
backend (str, optional) – Compilation backend. Default is "inductor"
**kwargs (dict) – Additional keyword arguments passed to torch.compile

Returns:

Returns self for method chaining

Return type:

NodeAssembler

Built-in Assemblers¶

class LaplaceElementAssembler(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, edges: Tensor, *args, **kwargs)[source]¶

Bases: ElementAssembler

Laplace/Diffusion Element Assembler.

Assembles the stiffness matrix for the Laplace operator (diffusion term). This is the fundamental building block for solving elliptic PDEs such as the Poisson equation, heat equation (steady-state), and other diffusion problems.

Weak Form:

Given the Laplace equation \(-\nabla \cdot (\kappa \nabla u) = f\), the weak form is:

\[\int_{\Omega} \kappa \nabla u \cdot \nabla v \, \mathrm{d}\Omega = \int_{\Omega} f v \, \mathrm{d}\Omega\]

Element Stiffness Matrix:

For each element \(K\), the local stiffness matrix entry is:

\[K_{ij}^e = \int_{\Omega^e} \nabla N_i \cdot \nabla N_j \, \mathrm{d}\Omega\]

where \(N_i, N_j\) are the shape functions.

Implementation:

The forward method computes the integrand \(\nabla N_i \cdot \nabla N_j\) at each quadrature point, which is then integrated by the base class.

Examples

mesh = Mesh.gen_rectangle(chara_length=0.1)
K = LaplaceElementAssembler.from_mesh(mesh)(mesh.points)

forward(gradu, gradv)[source]¶

Define the integrand of the bilinear form at a single quadrature point.

Parameters:

u (Tensor, optional) – Shape value at the quadrature point — 0D tensor of shape [].
v (Tensor, optional) – Shape value at the quadrature point — 0D tensor of shape [].
gradu (Tensor, optional) – Shape gradient in physical coordinates — 1D tensor of shape [D].
gradv (Tensor, optional) – Shape gradient in physical coordinates — 1D tensor of shape [D].
x (Tensor, optional) – Physical coordinate at the quadrature point — 1D tensor of shape [D].
gradx (Tensor, optional) – Gradient of x w.r.t. reference coordinates — 2D tensor of shape [D, D].
**point_data (Tensor) – Any key passed to __call__ via point_data: if the nodal tensor has shape \([|\mathcal V|, ...]\), the value handed to forward has the trailing [...] shape, and its counterpart "grad"+key has shape [..., D].

Returns:

Either a 2D tensor of shape [B, B] (scalar problems) or a 4D tensor of shape [B, B, H, H] (vector problems with H degrees of freedom per node).

Return type:

Tensor

class MassElementAssembler(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, edges: Tensor, *args, **kwargs)[source]¶

Bases: ElementAssembler

Mass Element Assembler.

Assembles the mass matrix for finite element discretization. The mass matrix represents the \(L^2\) inner product and is essential for time-dependent problems, eigenvalue problems, and \(L^2\) projection.

Weak Form:

The mass matrix arises from terms like \(\int_\Omega u \, v \, \mathrm{d}\Omega\) in the weak formulation.

Element Mass Matrix:

For each element \(K\), the local mass matrix entry is:

\[M_{ij}^e = \int_{\Omega^e} N_i \, N_j \, \mathrm{d}\Omega\]

where \(N_i, N_j\) are the shape functions.

Applications:

Time-dependent PDEs (heat equation, wave equation)
\(L^2\) error computation: \(\|u - u_h\|_{L^2}^2 = (u-u_h)^T M (u-u_h)\)
Eigenvalue problems: \(K u = \lambda M u\)

Examples

mesh = Mesh.gen_rectangle(chara_length=0.1)
M = MassElementAssembler.from_mesh(mesh)(mesh.points)

forward(u, v)[source]¶

Define the integrand of the bilinear form at a single quadrature point.

Parameters:

u (Tensor, optional) – Shape value at the quadrature point — 0D tensor of shape [].
v (Tensor, optional) – Shape value at the quadrature point — 0D tensor of shape [].
gradu (Tensor, optional) – Shape gradient in physical coordinates — 1D tensor of shape [D].
gradv (Tensor, optional) – Shape gradient in physical coordinates — 1D tensor of shape [D].
x (Tensor, optional) – Physical coordinate at the quadrature point — 1D tensor of shape [D].
gradx (Tensor, optional) – Gradient of x w.r.t. reference coordinates — 2D tensor of shape [D, D].
**point_data (Tensor) – Any key passed to __call__ via point_data: if the nodal tensor has shape \([|\mathcal V|, ...]\), the value handed to forward has the trailing [...] shape, and its counterpart "grad"+key has shape [..., D].

Returns:

Either a 2D tensor of shape [B, B] (scalar problems) or a 4D tensor of shape [B, B, H, H] (vector problems with H degrees of freedom per node).

Return type:

Tensor

class LinearElasticityElementAssembler(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, edges: Tensor, *args, **kwargs)[source]¶

Bases: ElementAssembler

Linear Elasticity Element Assembler.

Assembles the stiffness matrix for linear elastic materials based on Hooke’s law. Suitable for small deformation analysis of isotropic materials.

Constitutive Law (Hooke’s Law):

The stress-strain relationship for isotropic linear elasticity:

\[\boldsymbol{\sigma} = \mathbf{C} : \boldsymbol{\varepsilon}\]

where the elasticity tensor \(\mathbf{C}\) is defined by:

\[C_{ijkl} = \lambda \delta_{ij} \delta_{kl} + \mu (\delta_{ik}\delta_{jl} + \delta_{il}\delta_{jk})\]

Lamé Parameters:

\[\lambda = \frac{E \nu}{(1+\nu)(1-2\nu)}, \quad \mu = \frac{E}{2(1+\nu)}\]

where \(E\) is Young’s modulus and \(\nu\) is Poisson’s ratio.

Strain Tensor:

The infinitesimal strain tensor:

\[\boldsymbol{\varepsilon} = \frac{1}{2}(\nabla \mathbf{u} + \nabla \mathbf{u}^T)\]

Weak Form:

\[\int_{\Omega} \boldsymbol{\sigma}(\mathbf{u}) : \boldsymbol{\varepsilon}(\mathbf{v}) \, \mathrm{d}\Omega = \int_{\Omega} \mathbf{f} \cdot \mathbf{v} \, \mathrm{d}\Omega + \int_{\Gamma_N} \mathbf{t} \cdot \mathbf{v} \, \mathrm{d}S\]

Strain Energy Density:

\[\Psi = \frac{1}{2} \boldsymbol{\varepsilon} : \mathbf{C} : \boldsymbol{\varepsilon} = \frac{\lambda}{2} (\mathrm{tr}\, \boldsymbol{\varepsilon})^2 + \mu \, \boldsymbol{\varepsilon} : \boldsymbol{\varepsilon}\]

Parameters:

E (float, optional) – Young’s modulus. Default: 1.0.
nu (float, optional) – Poisson’s ratio (must satisfy \(-1 < \nu < 0.5\)). Default: 0.3.

Examples

mesh = Mesh.gen_cube(chara_length=0.1)
assembler = LinearElasticityElementAssembler.from_mesh(mesh, E=210e9, nu=0.3)
K = assembler(mesh.points)

forward(gradu, gradv)[source]¶

Define the integrand of the bilinear form at a single quadrature point.

Parameters:

u (Tensor, optional) – Shape value at the quadrature point — 0D tensor of shape [].
v (Tensor, optional) – Shape value at the quadrature point — 0D tensor of shape [].
gradu (Tensor, optional) – Shape gradient in physical coordinates — 1D tensor of shape [D].
gradv (Tensor, optional) – Shape gradient in physical coordinates — 1D tensor of shape [D].
x (Tensor, optional) – Physical coordinate at the quadrature point — 1D tensor of shape [D].
gradx (Tensor, optional) – Gradient of x w.r.t. reference coordinates — 2D tensor of shape [D, D].
**point_data (Tensor) – Any key passed to __call__ via point_data: if the nodal tensor has shape \([|\mathcal V|, ...]\), the value handed to forward has the trailing [...] shape, and its counterpart "grad"+key has shape [..., D].

Returns:

Either a 2D tensor of shape [B, B] (scalar problems) or a 4D tensor of shape [B, B, H, H] (vector problems with H degrees of freedom per node).

Return type:

Tensor

element_energy(graddisplacement)[source]¶

Compute strain energy density at a quadrature point.

\[\Psi = \frac{\lambda}{2} (\mathrm{tr}\, \boldsymbol{\varepsilon})^2 + \mu \, \boldsymbol{\varepsilon} : \boldsymbol{\varepsilon}\]

Parameters:: graddisplacement (Tensor) – Displacement gradient \(\nabla \mathbf{u}\) of shape [dim, dim].
Returns:: Scalar strain energy density.
Return type:: Tensor

class NeoHookeanModel(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, edges: Tensor, *args, **kwargs)[source]¶

Bases: ElementAssembler

Neo-Hookean Hyperelastic Material Model.

A nonlinear hyperelastic constitutive model for large deformation analysis. The Neo-Hookean model is the simplest hyperelastic model and extends linear elasticity to the finite strain regime.

Deformation Gradient:

\[\mathbf{F} = \mathbf{I} + \nabla \mathbf{u}\]

Kinematic Quantities:

Right Cauchy-Green tensor: \(\mathbf{C} = \mathbf{F}^T \mathbf{F}\)
First invariant: \(I_1 = \mathrm{tr}(\mathbf{C}) = \|\mathbf{F}\|_F^2\)
Jacobian (volume ratio): \(J = \det(\mathbf{F})\)

Strain Energy Density:

\[\Psi = \frac{\mu}{2}(I_1 - d) - \mu \ln J + \frac{\lambda}{2}(\ln J)^2\]

where \(d\) is the spatial dimension (2 or 3).

First Piola-Kirchhoff Stress:

\[\mathbf{P} = \frac{\partial \Psi}{\partial \mathbf{F}} = \mu \mathbf{F} + (\lambda \ln J - \mu) \mathbf{F}^{-T}\]

Material Parameters:

\[\mu = \frac{E}{2(1+\nu)}, \quad \lambda = \frac{E\nu}{(1+\nu)(1-2\nu)}\]

Parameters:

E (float, optional) – Young’s modulus. Default: 1.0.
nu (float, optional) – Poisson’s ratio. Default: 0.3.

Notes

Requires \(J > 0\) (no element inversion). For nearly incompressible materials (\(\nu \to 0.5\)), consider a mixed formulation to avoid volumetric locking.

Examples

mesh = Mesh.gen_cube(chara_length=0.1)
model = NeoHookeanModel.from_mesh(mesh, E=1e6, nu=0.45)
E_tot = model.energy(displacement)

element_energy(graddisplacement)[source]¶

Compute Neo-Hookean strain energy density at a quadrature point.

\[\Psi = \frac{\mu}{2}(I_1 - d) - \mu \ln J + \frac{\lambda}{2}(\ln J)^2\]

Parameters:: graddisplacement (Tensor) – Displacement gradient \(\nabla \mathbf{u}\) of shape [dim, dim].
Returns:: Scalar strain energy density.
Return type:: Tensor

energy(u)[source]¶

Compute total strain energy.

\[\Pi = \int_{\Omega} \Psi(\mathbf{F}) \, \mathrm{d}\Omega\]

Parameters:: u (Tensor) – Displacement field of shape [n_nodes, dim].
Returns:: Scalar total strain energy.
Return type:: Tensor

class J2Plasticity(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, edges: Tensor, *args, **kwargs)[source]¶

Bases: ElementAssembler

J2 (von Mises) Plasticity Model with Isotropic Hardening.

Implements rate-independent J2 plasticity with linear isotropic hardening using a return-mapping algorithm. This model is suitable for metals and other ductile materials under monotonic or cyclic loading.

Yield Function (von Mises):

\[f(\boldsymbol{\sigma}, \alpha) = \|\mathbf{s}\| - \sqrt{\frac{2}{3}}(\sigma_0 + H\alpha) \leq 0\]

where:

\(\mathbf{s} = \boldsymbol{\sigma} - \frac{1}{3}\mathrm{tr}(\boldsymbol{\sigma})\mathbf{I}\) is the deviatoric stress
\(\sigma_0\) is the initial yield stress
\(H\) is the hardening modulus
\(\alpha\) is the equivalent plastic strain (internal variable)

Additive Strain Decomposition:

\[\boldsymbol{\varepsilon} = \boldsymbol{\varepsilon}^e + \boldsymbol{\varepsilon}^p\]

Flow Rule (Associated):

\[\dot{\boldsymbol{\varepsilon}}^p = \dot{\gamma} \frac{\partial f}{\partial \boldsymbol{\sigma}} = \dot{\gamma} \frac{\mathbf{s}}{\|\mathbf{s}\|}\]

Hardening Law:

\[\dot{\alpha} = \sqrt{\frac{2}{3}} \dot{\gamma}\]

Return Mapping Algorithm:

Compute trial elastic stress: \(\boldsymbol{\sigma}^{tr} = \mathbf{C}:(\boldsymbol{\varepsilon} - \boldsymbol{\varepsilon}^p_n)\)
Check yield: \(f^{tr} = \|\mathbf{s}^{tr}\| - \sqrt{2/3}(\sigma_0 + H\alpha_n)\)
If \(f^{tr} \leq 0\): elastic step, accept trial state
If \(f^{tr} > 0\): plastic correction

\[\Delta\gamma = \frac{f^{tr}}{2\mu + \frac{2}{3}H}\]

Algorithmic Incremental Potential:

\[\Psi^{alg} = \frac{K}{2}(\mathrm{tr}\,\boldsymbol{\varepsilon}^e)^2 + \mu \|\mathbf{e}^{tr}\|^2 - \frac{1}{2}(2\mu + \frac{2}{3}H)(\Delta\gamma)^2\]

where \(K = \lambda + \frac{2}{3}\mu\) is the bulk modulus and \(\mathbf{e}^{tr}\) is the deviatoric trial strain.

Parameters:

material (optional) – Material object with properties E, nu, sigma_y, H; when supplied, overrides the individual scalar arguments.
E (float, optional) – Young’s modulus. Default: 200e9 (steel).
nu (float, optional) – Poisson’s ratio. Default: 0.3.
sig0 (float, optional) – Initial yield stress. Default: 250e6.
H (float, optional) – Hardening modulus. Default: 1e9.

history¶

Internal state variables — plastic strain (\(\boldsymbol{\varepsilon}^p\)) and equivalent plastic strain (\(\alpha\)) — keyed by element type.

Type:: dict

Examples

mesh = Mesh.gen_cube(chara_length=0.05)
plasticity = J2Plasticity.from_mesh(mesh, E=200e9, nu=0.3, sig0=250e6, H=1e9)
# In time-stepping loop:
energy = plasticity.energy(point_data={"displacement": u})
energy.backward()
# After convergence:
plasticity.update_state(u)

element_energy(graddisplacement, eps_p_n, alpha_n)[source]¶

Compute algorithmic incremental potential energy density.

This implements the return-mapping algorithm at the quadrature point level.

Parameters:

graddisplacement (Tensor) – Displacement gradient \(\nabla \mathbf{u}\).
eps_p_n (Tensor) – Plastic strain from the previous step \(\boldsymbol{\varepsilon}^p_n\).
alpha_n (Tensor) – Equivalent plastic strain from the previous step \(\alpha_n\).

Returns:

Scalar incremental potential energy density.

Return type:

Tensor

update_state(u_vec)[source]¶

Update internal state variables after load-step convergence.

Call after the Newton-Raphson iteration converges to update the plastic strain and the equivalent plastic strain.

Parameters:: u_vec (Tensor) – Converged displacement field.

class ContactAssembler(facet_mask: ModuleDict, projector: ModuleDict, transformation: ModuleDict, *args, **kwargs)[source]¶

Bases: FacetAssembler

Contact/Boundary Facet Assembler.

Assembler for integrating energy contributions over boundary facets (surfaces in 3D, edges in 2D). This is useful for implementing:

Contact mechanics: Penalty or barrier methods for non-penetration constraints
Surface tension: Capillary effects in fluid-structure interaction
Pressure loads: Follower forces that remain normal to deformed surface
Robin boundary conditions: Mixed Dirichlet-Neumann conditions

Penalty Contact Formulation:

For a penalty-based contact between surface \(\Gamma\) and an obstacle:

\[\Pi_{contact} = \int_{\Gamma} \frac{\kappa}{2} \langle g_n \rangle_-^2 \, \mathrm{d}S\]

where:

\(g_n\) is the normal gap function (negative when penetrating)
\(\langle \cdot \rangle_- = \min(\cdot, 0)\) is the negative part (Macaulay bracket)
\(\kappa\) is the penalty stiffness

Usage Pattern:

Subclass ContactAssembler and implement element_energy to define the specific contact/boundary energy density.

Examples

class PenaltyContact(ContactAssembler):
    def __post_init__(self, kappa=1e6, obstacle_y=0.0):
        self.kappa = kappa
        self.obstacle_y = obstacle_y

    def element_energy(self, x):
        gap = x[..., 1] - self.obstacle_y         # y-coordinate gap
        penetration = torch.clamp(-gap, min=0.0)
        return 0.5 * self.kappa * penetration ** 2

Compute total boundary / contact energy.

Integrates element_energy over all selected boundary facets:

\[\Pi = \int_{\Gamma} \psi(\mathbf{x}, \mathbf{u}, \ldots) \, \mathrm{d}S\]

Parameters:

points (Tensor, optional) – Updated nodal coordinates; if None, the cached points are used.
func (Callable, optional) – Custom energy density to use in place of element_energy().
point_data (dict[str, Tensor], optional) – Nodal fields to interpolate at quadrature points.
element_data (dict, optional) – Element-wise data (constant or per-quadrature).
scalar_data (dict, optional) – Global scalar parameters.
batch_size (int, optional) – Quadrature-point batch size; -1 (default) means no batching.

Returns:

Scalar total boundary energy.

Return type:

Tensor

element_energy(**kwargs)[source]¶: Override this method to define the boundary energy density.

const_node_assembler(c=1)[source]¶

Factory: build a NodeAssembler for a constant body load.

Weak form:

\[f_i = \int_{\Omega} c \, N_i \, \mathrm{d}\Omega\]

This represents a uniform body force or source term.

Parameters:: c (float, optional) – Constant load value. Default: 1.
Returns:: A new NodeAssembler subclass with c baked in.
Return type:: type[NodeAssembler]

Examples

ConstLoad = const_node_assembler(c=9.81)        # gravity
f = ConstLoad.from_mesh(mesh)(mesh.points)

func_node_assembler(f=lambda x: ...)[source]¶

Factory: build a NodeAssembler for a spatially-varying load.

Weak form:

\[f_i = \int_{\Omega} f(\mathbf{x}) \, N_i \, \mathrm{d}\Omega\]

Parameters:: f (Callable) – Function returning the load value at a coordinate. Signature f(x) -> load, where x has shape [..., dim].
Returns:: A new NodeAssembler subclass with f baked in.
Return type:: type[NodeAssembler]

Examples

source = func_node_assembler(lambda x: torch.sin(np.pi * x[..., 0]))
rhs = source.from_mesh(mesh)(mesh.points)

Projector¶

Internal helper used by the assemblers to scatter element-local contributions onto global degrees of freedom. Two concrete subclasses (ReduceProjector, SparseProjector) exist; both share the same interface documented here.

class Projector(*args: Any, **kwargs: Any)[source]¶

Bases: Module

Abstract base for the element-to-global scatter operators.

A Projector consumes a tensor with leading shape from_shape (per-element / per-facet quantities) and returns a tensor with leading shape to_shape (global edge / node indexing), summing duplicates. The two concrete implementations are ReduceProjector (uses torch.Tensor.index_add_()) and SparseProjector (uses a sparse mat-vec product).