tensormesh.assemble¶

Element Assembler¶

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

基类：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.

示例

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)[源代码]¶: 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)[源代码]¶

Casts all parameters and buffers to dst_type.

备注

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.

参数:

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).

返回:

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).

返回类型:

Tensor

示例

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)[源代码]¶

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.

参数:

**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).

返回:

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

返回类型:

Tensor

abstract forward(**kwargs)[源代码]¶

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__).

参数:

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].

返回:

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).

返回类型:

Tensor

classmethod from_assembler(obj, *args, **kwargs)[源代码]¶

Build an ElementAssembler that shares topology with obj.

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

参数:

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__.

返回:

A new assembler sharing the same mesh.

返回类型:

ElementAssembler

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

Build an ElementAssembler from a Mesh.

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

参数:

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__.

返回:

A new assembler that owns the mesh topology.

返回类型:

ElementAssembler

Facet Assembler¶

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

基类：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.

示例

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)[源代码]¶: 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)[源代码]¶

Casts all parameters and buffers to dst_type.

备注

This method modifies the module in-place.

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

abstract forward(*args)[源代码]¶

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).

参数:

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].

返回:

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

返回类型:

Tensor

classmethod from_assembler(obj, *args, **kwargs)[源代码]¶

Build a FacetAssembler sharing topology with obj.

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

参数:

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__.

返回:

A new assembler sharing the same boundary topology.

返回类型:

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)[源代码]¶

Build a FacetAssembler from raw connectivity tensors.

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

参数:

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".

返回:

A new assembler that owns the given boundary topology.

返回类型:

FacetAssembler

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

Build a FacetAssembler from a Mesh.

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

参数:

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".

返回:

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

返回类型:

FacetAssembler

Node Assembler¶

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

基类：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.

示例

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)[源代码]¶: 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)[源代码]¶

Casts all parameters and buffers to dst_type.

备注

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[源代码]¶

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.

示例

# 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()

参数:

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 self for method chaining

返回类型:

NodeAssembler

参见

reset_compile: Disable fast mode and use vmap path
is_compiled: Check if fast mode is enabled

flat_mode() → NodeAssembler[源代码]¶

Enable the fast broadcast-based implementation without torch.compile.

This allows to bypass vmap and use optimized broadcast operations. It is equivalent to calling compile(mode="disable").

返回:: Returns self for method chaining
返回类型:: NodeAssembler

reset_compile() → NodeAssembler[源代码]¶

Disable torch.compile and clear the compiled function cache.

This is useful for debugging or when you want to switch back to the non-compiled version.

示例

# Disable compile for debugging
assembler.reset_compile()

# Now you can set breakpoints in forward()

返回:: Returns self for method chaining
返回类型:: NodeAssembler

property is_compiled: bool¶

Check if the assembler is in compile mode.

返回:: True if compile mode is enabled, False otherwise
返回类型:: bool

abstract forward(*args)[源代码]¶

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

参数:

v (Tensor, optional) -- Shape value at the quadrature point — 0D tensor of shape [].
gradv (Tensor, optional) -- Shape gradient in physical coordinates — 1D tensor of shape [D].
x (Tensor, optional) -- Physical coordinate — 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].

返回:

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

返回类型:

Tensor

classmethod from_assembler(obj, *args, **kwargs)[源代码]¶

Build a NodeAssembler that shares topology with obj.

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

参数:

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

返回:

A new assembler sharing the same mesh.

返回类型:

NodeAssembler

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

Build a NodeAssembler from raw connectivity tensors.

Slower than from_assembler() because the projection backend is built from scratch.

参数:

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]])}.
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; "reduce" (default) uses torch.Tensor.index_add_() and is memory-efficient, "sparse" uses a sparse mat-vec product and is faster but uses more memory.

返回:

A new assembler that owns the given topology.

返回类型:

NodeAssembler

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

Build a NodeAssembler from a Mesh.

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

参数:

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) -- Projection backend; defaults to "reduce".
*args -- Additional arguments forwarded to __post_init__.
**kwargs -- Additional arguments forwarded to __post_init__.

返回:

A new assembler that owns the mesh topology.

返回类型:

NodeAssembler

Built-in Assemblers¶

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

基类：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.

示例

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

forward(gradu, gradv)[源代码]¶

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

参数:

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].

返回:

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).

返回类型:

Tensor

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

基类：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\)

示例

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

forward(u, v)[源代码]¶

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

参数:

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].

返回:

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).

返回类型:

Tensor

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

基类：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}\]

参数:

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

示例

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)[源代码]¶

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

参数:

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].

返回:

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).

返回类型:

Tensor

element_energy(graddisplacement)[源代码]¶

Compute strain energy density at a quadrature point.

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

参数:: graddisplacement (Tensor) -- Displacement gradient \(\nabla \mathbf{u}\) of shape [dim, dim].
返回:: Scalar strain energy density.
返回类型:: Tensor

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

基类：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)}\]

参数:

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

备注

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

示例

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)[源代码]¶

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\]

参数:: graddisplacement (Tensor) -- Displacement gradient \(\nabla \mathbf{u}\) of shape [dim, dim].
返回:: Scalar strain energy density.
返回类型:: Tensor

energy(u)[源代码]¶

Compute total strain energy.

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

参数:: u (Tensor) -- Displacement field of shape [n_nodes, dim].
返回:: Scalar total strain energy.
返回类型:: Tensor

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

基类：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.

参数:

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

示例

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)[源代码]¶

Compute algorithmic incremental potential energy density.

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

参数:

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\).

返回:

Scalar incremental potential energy density.

返回类型:

Tensor

update_state(u_vec)[源代码]¶

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.

参数:: u_vec (Tensor) -- Converged displacement field.

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

基类：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.

示例

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\]

参数:

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.

返回:

Scalar total boundary energy.

返回类型:

Tensor

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

const_node_assembler(c=1)[源代码]¶

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.

参数:: c (float, optional) -- Constant load value. Default: 1.
返回:: A new NodeAssembler subclass with c baked in.
返回类型:: type[NodeAssembler]

示例

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

func_node_assembler(f=lambda x: ...)[源代码]¶

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

Weak form:

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

参数:: f (Callable) -- Function returning the load value at a coordinate. Signature f(x) -> load, where x has shape [..., dim].
返回:: A new NodeAssembler subclass with f baked in.
返回类型:: type[NodeAssembler]

示例

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)[源代码]¶

基类：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).