"""Ready-made assemblers for the most common FEM forms.
This module is the "battery" tier of :mod:`tensormesh.assemble`: each class
below is a worked subclass of one of the three base assemblers
(:class:`ElementAssembler`, :class:`NodeAssembler`, :class:`FacetAssembler`)
covering a textbook form — Laplace, mass, linear elasticity, Neo-Hookean
hyperelasticity, J2 plasticity, contact, and constant/function loads.
Subclass these (or use the two factory functions at the bottom) to skip
re-deriving the weak form for standard physics.
"""
import inspect
from typing import Optional, Dict, Union, Callable
import numpy as np
import torch
from .element_assembler import ElementAssembler, InputBroadcast
from .facet_assembler import FacetAssembler
from .node_assembler import NodeAssembler
from ..functional.elasticity import voigt_shape_grad, voigt_stiffness
from ..vmap import vmap
[文档]
class LaplaceElementAssembler(ElementAssembler):
r"""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 :math:`-\nabla \cdot (\kappa \nabla u) = f`, the weak form is:
.. math::
\int_{\Omega} \kappa \nabla u \cdot \nabla v \, \mathrm{d}\Omega = \int_{\Omega} f v \, \mathrm{d}\Omega
**Element Stiffness Matrix:**
For each element :math:`K`, the local stiffness matrix entry is:
.. math::
K_{ij}^e = \int_{\Omega^e} \nabla N_i \cdot \nabla N_j \, \mathrm{d}\Omega
where :math:`N_i, N_j` are the shape functions.
**Implementation:**
The ``forward`` method computes the integrand :math:`\nabla N_i \cdot \nabla N_j`
at each quadrature point, which is then integrated by the base class.
Examples
--------
.. code-block:: python
mesh = Mesh.gen_rectangle(chara_length=0.1)
K = LaplaceElementAssembler.from_mesh(mesh)(mesh.points)
"""
[文档]
def forward(self, gradu, gradv):
return gradu @ gradv
[文档]
class MassElementAssembler(ElementAssembler):
r"""Mass Element Assembler.
Assembles the mass matrix for finite element discretization. The mass matrix
represents the :math:`L^2` inner product and is essential for time-dependent
problems, eigenvalue problems, and :math:`L^2` projection.
**Weak Form:**
The mass matrix arises from terms like :math:`\int_\Omega u \, v \, \mathrm{d}\Omega`
in the weak formulation.
**Element Mass Matrix:**
For each element :math:`K`, the local mass matrix entry is:
.. math::
M_{ij}^e = \int_{\Omega^e} N_i \, N_j \, \mathrm{d}\Omega
where :math:`N_i, N_j` are the shape functions.
**Applications:**
- Time-dependent PDEs (heat equation, wave equation)
- :math:`L^2` error computation: :math:`\|u - u_h\|_{L^2}^2 = (u-u_h)^T M (u-u_h)`
- Eigenvalue problems: :math:`K u = \lambda M u`
Examples
--------
.. code-block:: python
mesh = Mesh.gen_rectangle(chara_length=0.1)
M = MassElementAssembler.from_mesh(mesh)(mesh.points)
"""
[文档]
def forward(self, u, v):
return u * v
[文档]
class LinearElasticityElementAssembler(ElementAssembler):
r"""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:
.. math::
\boldsymbol{\sigma} = \mathbf{C} : \boldsymbol{\varepsilon}
where the elasticity tensor :math:`\mathbf{C}` is defined by:
.. math::
C_{ijkl} = \lambda \delta_{ij} \delta_{kl} + \mu (\delta_{ik}\delta_{jl} + \delta_{il}\delta_{jk})
**Lamé Parameters:**
.. math::
\lambda = \frac{E \nu}{(1+\nu)(1-2\nu)}, \quad \mu = \frac{E}{2(1+\nu)}
where :math:`E` is Young's modulus and :math:`\nu` is Poisson's ratio.
**Strain Tensor:**
The infinitesimal strain tensor:
.. math::
\boldsymbol{\varepsilon} = \frac{1}{2}(\nabla \mathbf{u} + \nabla \mathbf{u}^T)
**Weak Form:**
.. math::
\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:**
.. math::
\Psi = \frac{1}{2} \boldsymbol{\varepsilon} : \mathbf{C} : \boldsymbol{\varepsilon}
= \frac{\lambda}{2} (\mathrm{tr}\, \boldsymbol{\varepsilon})^2 + \mu \, \boldsymbol{\varepsilon} : \boldsymbol{\varepsilon}
Parameters
----------
E : float, optional
Young's modulus. Default: ``1.0``.
nu : float, optional
Poisson's ratio (must satisfy :math:`-1 < \nu < 0.5`). Default: ``0.3``.
Examples
--------
.. code-block:: python
mesh = Mesh.gen_cube(chara_length=0.1)
assembler = LinearElasticityElementAssembler.from_mesh(mesh, E=210e9, nu=0.3)
K = assembler(mesh.points)
"""
def __post_init__(self, E=1.0, nu=0.3):
self.E = E
self.nu = nu
[文档]
def forward(self, gradu, gradv):
dim = gradu.shape[0]
Ba = voigt_shape_grad(gradu)
Bb = voigt_shape_grad(gradv)
C = voigt_stiffness(self.E, self.nu, dim)
C = C.to(dtype=gradu.dtype, device=gradu.device)
return Ba.T @ C @ Bb
[文档]
def element_energy(self, graddisplacement):
r"""Compute strain energy density at a quadrature point.
.. math::
\Psi = \frac{\lambda}{2} (\mathrm{tr}\, \boldsymbol{\varepsilon})^2
+ \mu \, \boldsymbol{\varepsilon} : \boldsymbol{\varepsilon}
Parameters
----------
graddisplacement : torch.Tensor
Displacement gradient :math:`\nabla \mathbf{u}` of shape ``[dim, dim]``.
Returns
-------
torch.Tensor
Scalar strain energy density.
"""
grad_u = graddisplacement
dim = grad_u.shape[-1]
# Strain epsilon = 0.5 (grad_u + grad_u.T)
eps = 0.5 * (grad_u + grad_u.transpose(-1, -2))
# Lame parameters
mu = self.E / (2 * (1 + self.nu))
lam = (self.E * self.nu) / ((1 + self.nu) * (1 - 2 * self.nu))
# Trace squared
tr_eps = eps.diagonal(dim1=-2, dim2=-1).sum(-1)
vol_term = 0.5 * lam * (tr_eps ** 2)
# Double dot product eps : eps
eps_sq = (eps * eps).sum(dim=(-2, -1))
dev_term = mu * eps_sq
energy = vol_term + dev_term
return energy
[文档]
class NeoHookeanModel(ElementAssembler):
r"""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:**
.. math::
\mathbf{F} = \mathbf{I} + \nabla \mathbf{u}
**Kinematic Quantities:**
- Right Cauchy-Green tensor: :math:`\mathbf{C} = \mathbf{F}^T \mathbf{F}`
- First invariant: :math:`I_1 = \mathrm{tr}(\mathbf{C}) = \|\mathbf{F}\|_F^2`
- Jacobian (volume ratio): :math:`J = \det(\mathbf{F})`
**Strain Energy Density:**
.. math::
\Psi = \frac{\mu}{2}(I_1 - d) - \mu \ln J + \frac{\lambda}{2}(\ln J)^2
where :math:`d` is the spatial dimension (2 or 3).
**First Piola-Kirchhoff Stress:**
.. math::
\mathbf{P} = \frac{\partial \Psi}{\partial \mathbf{F}}
= \mu \mathbf{F} + (\lambda \ln J - \mu) \mathbf{F}^{-T}
**Material Parameters:**
.. math::
\mu = \frac{E}{2(1+\nu)}, \quad \lambda = \frac{E\nu}{(1+\nu)(1-2\nu)}
Parameters
----------
E : float, optional
Young's modulus. Default: ``1.0``.
nu : float, optional
Poisson's ratio. Default: ``0.3``.
Notes
-----
Requires :math:`J > 0` (no element inversion). For nearly incompressible
materials (:math:`\nu \to 0.5`), consider a mixed formulation to avoid
volumetric locking.
Examples
--------
.. code-block:: python
mesh = Mesh.gen_cube(chara_length=0.1)
model = NeoHookeanModel.from_mesh(mesh, E=1e6, nu=0.45)
E_tot = model.energy(displacement)
"""
def __post_init__(self, E=1.0, nu=0.3):
self.mu = E / (2 * (1 + nu))
self.lam = E * nu / ((1 + nu) * (1 - 2 * nu))
[文档]
def element_energy(self, graddisplacement):
r"""Compute Neo-Hookean strain energy density at a quadrature point.
.. math::
\Psi = \frac{\mu}{2}(I_1 - d) - \mu \ln J + \frac{\lambda}{2}(\ln J)^2
Parameters
----------
graddisplacement : torch.Tensor
Displacement gradient :math:`\nabla \mathbf{u}` of shape ``[dim, dim]``.
Returns
-------
torch.Tensor
Scalar strain energy density.
"""
grad_u = graddisplacement
dim = grad_u.shape[-1]
# F = I + grad_u
F = torch.eye(dim, device=grad_u.device, dtype=grad_u.dtype) + grad_u
# Invariants
J = torch.linalg.det(F)
I1 = (F * F).sum() # tr(F^T F)
log_J = torch.log(J)
psi = (self.mu / 2) * (I1 - dim) - self.mu * log_J + (self.lam / 2) * (log_J ** 2)
return psi
[文档]
def energy(self, u):
r"""Compute total strain energy.
.. math::
\Pi = \int_{\Omega} \Psi(\mathbf{F}) \, \mathrm{d}\Omega
Parameters
----------
u : torch.Tensor
Displacement field of shape ``[n_nodes, dim]``.
Returns
-------
torch.Tensor
Scalar total strain energy.
"""
return super().energy(point_data={"displacement": u})
[文档]
class J2Plasticity(ElementAssembler):
r"""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):**
.. math::
f(\boldsymbol{\sigma}, \alpha) = \|\mathbf{s}\| - \sqrt{\frac{2}{3}}(\sigma_0 + H\alpha) \leq 0
where:
- :math:`\mathbf{s} = \boldsymbol{\sigma} - \frac{1}{3}\mathrm{tr}(\boldsymbol{\sigma})\mathbf{I}` is the deviatoric stress
- :math:`\sigma_0` is the initial yield stress
- :math:`H` is the hardening modulus
- :math:`\alpha` is the equivalent plastic strain (internal variable)
**Additive Strain Decomposition:**
.. math::
\boldsymbol{\varepsilon} = \boldsymbol{\varepsilon}^e + \boldsymbol{\varepsilon}^p
**Flow Rule (Associated):**
.. math::
\dot{\boldsymbol{\varepsilon}}^p = \dot{\gamma} \frac{\partial f}{\partial \boldsymbol{\sigma}}
= \dot{\gamma} \frac{\mathbf{s}}{\|\mathbf{s}\|}
**Hardening Law:**
.. math::
\dot{\alpha} = \sqrt{\frac{2}{3}} \dot{\gamma}
**Return Mapping Algorithm:**
1. Compute trial elastic stress: :math:`\boldsymbol{\sigma}^{tr} = \mathbf{C}:(\boldsymbol{\varepsilon} - \boldsymbol{\varepsilon}^p_n)`
2. Check yield: :math:`f^{tr} = \|\mathbf{s}^{tr}\| - \sqrt{2/3}(\sigma_0 + H\alpha_n)`
3. If :math:`f^{tr} \leq 0`: elastic step, accept trial state
4. If :math:`f^{tr} > 0`: plastic correction
.. math::
\Delta\gamma = \frac{f^{tr}}{2\mu + \frac{2}{3}H}
**Algorithmic Incremental Potential:**
.. math::
\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 :math:`K = \lambda + \frac{2}{3}\mu` is the bulk modulus and
:math:`\mathbf{e}^{tr}` is the deviatoric trial strain.
Parameters
----------
material : optional
Material object with properties ``E``, ``nu``, ``sigma_y``, ``H``;
when supplied, overrides the individual scalar arguments.
E : float, optional
Young's modulus. Default: ``200e9`` (steel).
nu : float, optional
Poisson's ratio. Default: ``0.3``.
sig0 : float, optional
Initial yield stress. Default: ``250e6``.
H : float, optional
Hardening modulus. Default: ``1e9``.
Attributes
----------
history : dict
Internal state variables — plastic strain (:math:`\boldsymbol{\varepsilon}^p`)
and equivalent plastic strain (:math:`\alpha`) — keyed by element type.
Examples
--------
.. code-block:: python
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)
"""
def __post_init__(self, material=None, E=200e9, nu=0.3, sig0=250e6, H=1e9):
if material is not None:
self.E = material.E
self.nu = material.nu
self.sig0 = material.sigma_y
self.H = material.H
else:
self.E = E
self.nu = nu
self.sig0 = sig0
self.H = H
# Lamé parameters
self.mu = self.E / (2 * (1 + self.nu))
self.lam = (self.E * self.nu) / ((1 + self.nu) * (1 - 2 * self.nu))
self.bulk = self.lam + 2/3 * self.mu
# Initialize History Variables
self.history = {}
for etype, trans in self.transformation.items():
n_elem = trans.n_elements
n_quad = trans.n_quadrature
# Plastic strain tensor (trace is 0 for J2)
eps_p = torch.zeros((n_elem, n_quad, 3, 3), device=self.device, dtype=self.dtype)
# Equivalent plastic strain
alpha = torch.zeros((n_elem, n_quad), device=self.device, dtype=self.dtype)
self.history[etype] = {'eps_p': eps_p, 'alpha': alpha}
[文档]
def element_energy(self, graddisplacement, eps_p_n, alpha_n):
r"""Compute algorithmic incremental potential energy density.
This implements the return-mapping algorithm at the quadrature point level.
Parameters
----------
graddisplacement : torch.Tensor
Displacement gradient :math:`\nabla \mathbf{u}`.
eps_p_n : torch.Tensor
Plastic strain from the previous step :math:`\boldsymbol{\varepsilon}^p_n`.
alpha_n : torch.Tensor
Equivalent plastic strain from the previous step :math:`\alpha_n`.
Returns
-------
torch.Tensor
Scalar incremental potential energy density.
"""
grad_u = graddisplacement
dim = grad_u.shape[0]
# Construct 3D strain tensor
if dim == 2:
eps_2d = 0.5 * (grad_u + grad_u.T)
eps = torch.nn.functional.pad(eps_2d, (0, 1, 0, 1))
else:
eps = 0.5 * (grad_u + grad_u.T)
# Trial Elastic Step
eps_tr = eps - eps_p_n
tr_eps_tr = eps_tr.diagonal(dim1=-2, dim2=-1).sum(-1)
dev_eps_tr = eps_tr - (tr_eps_tr / 3.0) * torch.eye(3, device=grad_u.device, dtype=grad_u.dtype)
norm_dev_eps_tr = torch.norm(dev_eps_tr)
# Trial yield criterion
norm_s_tr = 2 * self.mu * norm_dev_eps_tr
radius = np.sqrt(2/3) * (self.sig0 + self.H * alpha_n)
f_tr = norm_s_tr - radius
# Volumetric energy
vol_energy = 0.5 * self.bulk * (tr_eps_tr**2)
# Plastic multiplier
denom = 2 * self.mu + (2/3) * self.H
d_gamma = torch.clamp(f_tr, min=0.0) / denom
# Deviatoric energy with plastic correction
dev_energy = self.mu * (norm_dev_eps_tr**2) - 0.5 * denom * (d_gamma**2)
psi = vol_energy + dev_energy
return psi
[文档]
def update_state(self, u_vec):
r"""Update internal state variables after load-step convergence.
Call after the Newton-Raphson iteration converges to update the
plastic strain and the equivalent plastic strain.
Parameters
----------
u_vec : torch.Tensor
Converged displacement field.
"""
with torch.no_grad():
for etype, trans in self.transformation.items():
cells = trans.elements
u_elem = u_vec[cells]
grad_u = torch.einsum('bqkx,bku->bqux', trans.shape_grad, u_elem)
dim = grad_u.shape[-1]
eps = torch.zeros(grad_u.shape[:2] + (3, 3), device=u_vec.device, dtype=u_vec.dtype)
if dim == 2:
grad_u_2d = grad_u
eps[..., :2, :2] = 0.5 * (grad_u_2d + grad_u_2d.transpose(-1, -2))
else:
eps = 0.5 * (grad_u + grad_u.transpose(-1, -2))
hist = self.history[etype]
eps_p_n = hist['eps_p']
alpha_n = hist['alpha']
# Trial
eps_tr = eps - eps_p_n
tr_eps_tr = eps_tr.diagonal(dim1=-2, dim2=-1).sum(-1)
dev_eps_tr = eps_tr - (tr_eps_tr.unsqueeze(-1).unsqueeze(-1) / 3.0) * torch.eye(3, device=u_vec.device, dtype=u_vec.dtype)
norm_dev_eps_tr = torch.norm(dev_eps_tr, dim=(-2, -1))
norm_s_tr = 2 * self.mu * norm_dev_eps_tr
radius = np.sqrt(2/3) * (self.sig0 + self.H * alpha_n)
f_tr = norm_s_tr - radius
d_gamma = torch.clamp(f_tr, min=0.0) / (2 * self.mu + (2/3) * self.H)
# Update direction
norm_safe = torch.where(norm_dev_eps_tr < 1e-12, torch.ones_like(norm_dev_eps_tr), norm_dev_eps_tr)
n_tensor = dev_eps_tr / norm_safe.unsqueeze(-1).unsqueeze(-1)
yield_mask = f_tr > 0
d_gamma_masked = d_gamma * yield_mask.float()
hist['eps_p'] += d_gamma_masked.unsqueeze(-1).unsqueeze(-1) * n_tensor
hist['alpha'] += np.sqrt(2/3) * d_gamma_masked
[文档]
def const_node_assembler(c = 1):
r"""Factory: build a :class:`NodeAssembler` for a constant body load.
**Weak form:**
.. math::
f_i = \int_{\Omega} c \, N_i \, \mathrm{d}\Omega
This represents a uniform body force or source term.
Parameters
----------
c : float, optional
Constant load value. Default: ``1``.
Returns
-------
type[NodeAssembler]
A new :class:`NodeAssembler` subclass with ``c`` baked in.
Examples
--------
.. code-block:: python
ConstLoad = const_node_assembler(c=9.81) # gravity
f = ConstLoad.from_mesh(mesh)(mesh.points)
"""
class ConstNodeAssembler(NodeAssembler):
r"""Constant load node assembler.
.. math::
f = \int_{\Omega} c\cdot v \mathrm{d}\Omega
"""
def __post_init__(self, c=c):
self.c = c
def forward(self, v):
f = self.c * v
return f
return ConstNodeAssembler
[文档]
def func_node_assembler(f=lambda x: x):
r"""Factory: build a :class:`NodeAssembler` for a spatially-varying load.
**Weak form:**
.. math::
f_i = \int_{\Omega} f(\mathbf{x}) \, N_i \, \mathrm{d}\Omega
Parameters
----------
f : Callable
Function returning the load value at a coordinate. Signature
``f(x) -> load``, where ``x`` has shape ``[..., dim]``.
Returns
-------
type[NodeAssembler]
A new :class:`NodeAssembler` subclass with ``f`` baked in.
Examples
--------
.. code-block:: python
source = func_node_assembler(lambda x: torch.sin(np.pi * x[..., 0]))
rhs = source.from_mesh(mesh)(mesh.points)
"""
class FuncNodeAssembler(NodeAssembler):
r"""Function-based load node assembler.
.. math::
f = \int_{\Omega} f(\mathbf{x}) \, v \, \mathrm{d}\Omega
"""
def __post_init__(self, f=f):
self.f = f
def forward(self, x, v):
f = self.f(x) * v
return f
return FuncNodeAssembler
__all__ = ["LaplaceElementAssembler", "MassElementAssembler", "LinearElasticityElementAssembler", "NeoHookeanModel", "J2Plasticity", "ContactAssembler", "const_node_assembler", "func_node_assembler"]