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

Mixed Element Assembler¶

Multi-field (block) bilinear and linear forms on Lagrange spaces — e.g. Taylor-Hood Stokes — declared field-by-field and assembled into one block sparse matrix. See the user guide chapter for a tutorial.

class MixedElementAssembler(topology: dict, *args, **kwargs)[源代码]¶

基类：Module

Assemble a multi-field bilinear form into one block sparse matrix.

Declare the fields as the class attribute fields (a list of Field), override forward() with the scalar integrand of the bilinear form, and build with from_mesh(). See the module docstring for the conventions (trial → columns, test → rows; vector fields pass tensor-valued arguments; the integrand must be bilinear).

Besides the field arguments, forward may take the same data arguments as ElementAssembler: x (coordinates), any point_data key and its grad{key}, element_data keys and scalar_data keys — plus field_data keys (and their grad{key}), which live on a field's DOFs and are interpolated with that field's own basis. point_data is interpolated with the mesh-order basis.

The assembled matrix is square of size \(N = \sum_f n_f c_f\) with the field blocks laid out in declaration order (see BlockLayout); use assembler.layout to build boundary-condition masks and to split/concatenate DOF vectors. The unchanged Condenser applies on top.

备注

Only the ReduceProjector scatter backend is supported (the SparseProjector is float32-only, see the complex-FEM ROADMAP item).
energy / from_assembler are not provided for mixed forms.
Load vectors are assembled space-aware by assemble_vector() (override forward_vector() or pass func=): same argument dispatch, linear in the test functions, same block DOF layout.

__init__(topology: dict, *args, **kwargs)[源代码]¶: Initialize internal Module state, shared by both nn.Module and ScriptModule.

property layout: BlockLayout¶: Block-DOF layout helpers (offsets, masks, split/cat, ...).

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

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

forward(*args)[源代码]¶

Scalar integrand of the bilinear form at one quadrature point.

Override in subclasses. Arguments are requested by name: the trial/test names declared in fields (values), their grad-prefixed gradients, x, and any point_data / element_data / scalar_data key. Must return a 0-d tensor.

forward_vector(*args)[源代码]¶

Scalar integrand of the linear form assembled by assemble_vector().

Override in subclasses (or pass func= to assemble_vector()). Same argument dispatch as forward(), but only test arguments (and data) may appear — e.g. a Stokes body force:

def forward_vector(self, v, x):
    return x[1] * v[0]      # f = (y, 0)

Must return a 0-d tensor, linear in the test functions.

classmethod from_mesh(mesh: Mesh, quadrature_order: int | None = None, *args, **kwargs)[源代码]¶

Build a MixedElementAssembler from a Mesh.

参数:

mesh (Mesh) -- Source mesh; its element order is the geometry (and maximum field) order.
quadrature_order (int, optional) -- Degree of exactness of the quadrature rule. Defaults to 2 * max(field.order), which integrates every product of two field values/gradients exactly on affine elements.
*args -- Additional arguments forwarded to __post_init__.
**kwargs -- Additional arguments forwarded to __post_init__.

assemble_vector(points: Tensor | None = None, func: Callable | None = None, point_data: Mapping[str, Tensor] | None = None, element_data: Mapping[str, Mapping[str, Tensor]] | Mapping[str, Tensor] | None = None, scalar_data: Mapping[str, Tensor] | None = None, batch_size: int = -1, field_data: Mapping[str, Tuple[str, Tensor]] | None = None) → Tensor[源代码]¶

Assemble a linear form into the global block load vector.

The space-aware counterpart of a load-vector assembler: the integrand (forward_vector(), or func=) is written like forward() but may reference only test arguments plus data — e.g. \(\int f \cdot v\) for the Stokes momentum equation. It is evaluated with one-hot test basis functions per field and scattered into the same block DOF layout as __call__, so the result pairs directly with the assembled matrix and the Condenser. Fields whose test arguments do not appear contribute a zero segment.

Works for every field kind, including generalized-order fields (this supersedes the interim b = M @ f(points) interpolation recipe). field_data is supported — e.g. the previous time step of a velocity living on its own field DOFs.

返回:: Dense load vector of shape \([N]\) with \(N = \sum_f n_f c_f\) (see BlockLayout).
返回类型:: Tensor

class Field(trial: str, test: str, order: int = 1, components: int = 1)[源代码]¶

Declaration of one Lagrange field of a mixed bilinear form.

参数:

trial (str) -- Name of the trial-function argument in forward; its gradient is available as "grad" + trial. Trial arguments index matrix columns.
test (str) -- Name of the test-function argument in forward; its gradient is available as "grad" + test. Test arguments index matrix rows.
order (int, optional) -- Polynomial order of the field's Lagrange space — independent of the mesh order. Orders 1 and the mesh order carry their DOFs on mesh nodes; any other order uses a topological DOF map (see the module docstring for the 3D face-DOF limitation). Default 1.
components (int, optional) -- Number of vector components \(c\). With c == 1 the field is scalar (value [], gradient [D]); with c > 1 the value is [c] and the gradient [c, D]. Default 1.

示例

Field(trial="u", test="v", order=2, components=2)   # P2 velocity
Field(trial="p", test="q", order=1)                 # P1 pressure

__init__(trial: str, test: str, order: int = 1, components: int = 1)[源代码]¶

class BlockLayout(assembler: MixedElementAssembler)[源代码]¶

Block-DOF layout helpers of a MixedElementAssembler.

The global DOF vector stacks the fields in declaration order; within a field the layout is node-major:

\[\mathrm{dof} = \mathrm{offset}_f + n_{\mathrm{local}} \cdot c_f + \mathrm{comp}.\]

Access it as assembler.layout. All tensors are read through the assembler's buffers, so the layout follows .to() / .double().

__init__(assembler: MixedElementAssembler)[源代码]¶

property names: List[str]¶: Trial names of the fields, in declaration (block) order.

property n_dofs: int¶: Total number of DOFs \(N = \sum_f n_f\, c_f\).

property offsets: Dict[str, int]¶: First global DOF of each field, keyed by trial name.

field(name: str) → Field[源代码]¶: The Field declared with trial name name.

n_nodes(name: str) → int[源代码]¶: Number of nodes (scalar DOF carriers) of field name.

node_ids(name: str) → Tensor[源代码]¶

Mesh point ids of the field's nodes, sorted — [n_f] long.

Only fields whose DOFs live on mesh nodes (order 1 or the mesh order) have this; generalized-order fields carry DOFs on vertices/edges/interiors and raise here — use points(), boundary_mask() or dof_mask(where=...) instead.

points(name: str) → Tensor[源代码]¶

Coordinates of the field's nodes — [n_f, D].

For generalized-order fields the coordinates are the images of the field's reference nodes under the isoparametric (mesh-order) geometry map, so they are exact on curved elements too.

restrict(name: str, point_data: Tensor) → Tensor[源代码]¶

Interpolate mesh-point data [n_points, ...] onto the field's nodes [n_f, ...].

A plain gather for node-carried fields; FE interpolation with the mesh-order basis for generalized-order fields.

prolong(name: str, values: Tensor) → Tensor[源代码]¶

Interpolate field-node values [n_f, ...] to all mesh points [n_points, ...].

For a field of mesh order this is the identity; for an order-1 field on a higher-order mesh the values are FE-interpolated to the midside/interior nodes (useful e.g. to plot a P1 pressure on a P2 mesh).

split(x: Tensor) → Dict[str, Tensor][源代码]¶

Split a DOF vector [N, ...] into per-field tensors.

Returns {name: [n_f, c_f, ...]}; the component axis is squeezed away for scalar fields (c_f == 1).

cat(fields: Mapping[str, Tensor | float] | None = None, **kwargs: Tensor | float) → Tensor[源代码]¶

Concatenate per-field values into one flat DOF vector [N] (inverse of split()).

Every field must be given, either as a tensor of shape [n_f, c_f] / [n_f * c_f] (or [n_f] for scalar fields) or as a python scalar that is broadcast (e.g. p=0.0).

dof_mask(name: str, node_mask: Tensor | None = None, component: int | None = None, where: Callable[[Tensor], Tensor] | None = None) → Tensor[源代码]¶

Boolean mask over all N DOFs selecting (part of) one field.

参数:

name (str) -- Trial name of the field.
node_mask (Tensor, optional) -- Boolean mask over the field's nodes [n_f], or over the mesh points [n_points] (restricted to the field's nodes automatically — node-carried fields only). None selects every node of the field.
component (int, optional) -- Restrict to a single component; None selects all.
where (Callable, optional) -- Coordinate predicate: called with the field's node coordinates [n_f, D] (see points()), must return a bool [n_f] mask. Works for every field kind; mutually exclusive with node_mask. Example: where=lambda x: x[:, 1] > 1 - 1e-6.

boundary_mask(name: str, component: int | None = None) → Tensor[源代码]¶

dof_mask of the field's boundary DOFs, detected topologically.

A facet is on the boundary iff exactly one cell references it, so no is_boundary point data or geometric tolerance is needed and edge/interior DOFs of higher-order fields are handled correctly. The per-field result is cached after the first call.

dof_index(name: str, node: int | Tensor, component: int = 0) → Tensor[源代码]¶

Global DOF index(es) of field name at mesh point id(s) node.

Node-carried fields only (order 1 or the mesh order); raises if a requested mesh point carries no DOF of this field (e.g. a midside node for an order-1 field). For generalized-order fields select DOFs via dof_mask() / boundary_mask() instead.

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

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

基类：FacetAssembler

Assemble a bilinear form over boundary facets into a sparse matrix.

Override forward with the integrand of a boundary bilinear form at a single facet-quadrature point and basis pair — e.g. u * v for a boundary mass matrix, or c * u * v with a nodal coefficient c (impedance, radiation). gradu / gradv and any point_data key are available just as in ElementAssembler.

Calling the assembler returns a SparseMatrix of shape [n_points, n_points].

示例

Boundary mass matrix on the selected boundary:

class BoundaryMass(FacetBilinearAssembler):
    def forward(self, u, v):
        return u * v

B = BoundaryMass.from_mesh(mesh, boundary_mask=inlet)(mesh.points)

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 AnisotropicLaplaceElementAssembler(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, edges: Tensor, *args, **kwargs)[源代码]¶

基类：ElementAssembler

Tensor-coefficient (anisotropic) diffusion / stiffness assembler.

Generalizes LaplaceElementAssembler to a full material tensor \(\mathbf{A}\):

\[K_{ij}^e = \int_{\Omega^e} (\mathbf{A}\, \nabla N_i) \cdot \nabla N_j \, \mathrm{d}\Omega = \int_{\Omega^e} \nabla N_i^\top \mathbf{A}\, \nabla N_j \, \mathrm{d}\Omega .\]

The coefficient A is a nodal \([|\mathcal V|, D, D]\) field passed through point_data (interpolated to each quadrature point), so it may be spatially varying, anisotropic, and complex — e.g. the diagonal coordinate-stretch tensor \(\mathrm{diag}(s_y/s_x,\, s_x/s_y)\) of a stretched-coordinate PML, or an anisotropic conductivity/permittivity. With A the identity this reduces exactly to LaplaceElementAssembler.

示例

mesh = Mesh.gen_rectangle(chara_length=0.1)
A = torch.eye(2).expand(mesh.n_points, 2, 2).clone()   # [N, D, D]
K = AnisotropicLaplaceElementAssembler.from_mesh(mesh)(
        mesh.points, point_data={"A": A})

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

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 ScaledMassElementAssembler(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, edges: Tensor, *args, **kwargs)[源代码]¶

基类：ElementAssembler

Coefficient-weighted mass assembler \(\int_\Omega c\, N_i N_j\, d\Omega\).

Like MassElementAssembler but with a scalar nodal coefficient c (from point_data, possibly spatially varying and complex). Handy for the reaction / potential term of a Helmholtz problem — e.g. the PML mass scaling \(k_0^2\,\varepsilon\, s_{\mathrm{prod}}\) (see cartesian_pml()), or a variable wave speed.

示例

c = k0**2 * eps_r * s_prod                    # complex nodal field [N]
M = ScaledMassElementAssembler.from_mesh(mesh)(mesh.points, point_data={"c": c})

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

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 DruckerPragerPlasticity(projector: ModuleDict, transformation: ModuleDict, elements: BufferDict, edges: Tensor, *args, **kwargs)[源代码]¶

基类：ElementAssembler

Drucker-Prager Plasticity Model with Linear Isotropic Hardening.

Pressure-dependent, small-strain Drucker-Prager plasticity for soils and weak rock, built on the pure constitutive primitive drucker_prager_return_mapping(). It mirrors the J2Plasticity lifecycle — per-quadrature history, previous-step state passed through element_data, and update_state() after a converged load step — so it slots into the same energy-minimisation solvers.

Yield Function (tension-positive convention):

\[f(\boldsymbol{\sigma}, \alpha) = q + \eta\, I_1 - (k + H\alpha) \leq 0\]

where \(I_1 = \mathrm{tr}(\boldsymbol{\sigma})\), \(q = \sqrt{\tfrac{3}{2}\,\mathbf{s}:\mathbf{s}}\), and the cone coefficients \(\eta, k\) are fitted to the Mohr-Coulomb triaxial-compression meridian. Because compression gives a negative \(I_1\), confinement raises the yield capacity.

参数:

material (optional) -- Object with attributes E, nu, cohesion, friction_angle, dilatancy_angle and H (e.g. tensormesh.material.FrictionalMaterial); when supplied it overrides the scalar arguments.
E (float, optional) -- Young's modulus. Default 50e6.
nu (float, optional) -- Poisson's ratio. Default 0.30.
cohesion (float, optional) -- Cohesion in Pa. Default 20e3.
friction_angle (float, optional) -- Friction angle in degrees. Default 30.0.
dilatancy_angle (float or None, optional) -- Dilatancy angle in degrees; None (default) means associated flow.
H (float, optional) -- Linear isotropic hardening modulus in Pa. Default 1e6.

history¶

Internal state variables — plastic strain (eps_p) and equivalent plastic strain (alpha) — keyed by element type.

Type:: dict

示例

from tensormesh.assemble import DruckerPragerPlasticity
from tensormesh.material import FrictionalMaterial

soil = FrictionalMaterial.from_preset("DenseSand")
model = DruckerPragerPlasticity.from_mesh(mesh, material=soil)
# In a load-stepping loop:
energy = model.energy(point_data={"displacement": u},
                      element_data=model.element_data_from_history())
energy.backward()
# After convergence:
model.update_state(u)

element_energy(graddisplacement, eps_p_n, alpha_n)[源代码]¶

Algorithmic incremental potential density at one quadrature point.

Delegates to drucker_prager_return_mapping() and returns its energy field.

update_state(u_vec)[源代码]¶: Commit per-quadrature eps_p and alpha after convergence.

element_data_from_history()[源代码]¶: Return committed history as the element_data mapping for energy.

max_alpha()[源代码]¶: Maximum committed equivalent plastic strain over all quadrature points.

mean_alpha()[源代码]¶: Mean committed equivalent plastic strain over all quadrature points.

mean_stress(u_vec)[源代码]¶: Average committed Cauchy stress over elements and quadrature points.

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