[Shizheng-Wen]

What

ElementAssembler.__call__, ElementAssembler.energy and NodeAssembler.__call__ now take quadrature_data. Its tensors have shape [n_element, n_quadrature, ...], or {element_type: tensor} on meshes with several element types. Their value at the current quadrature point reaches forward by name, like any other data key:

class Tangent(ElementAssembler):          # K_ab = ∫ ∇N_a · C : ∇N_b
    def forward(self, gradu, gradv, C):
        return torch.einsum("j,ijkl,l->ik", gradu, C, gradv)

class InternalForce(NodeAssembler):       # f_a = ∫ σ ∇N_a
    def forward(self, gradv, sigma):
        return sigma @ gradv

K = Tangent.from_mesh(mesh)(quadrature_data={"C": C_alg})                 # C_alg: [E, Q, 3, 3, 3, 3]
f = InternalForce.from_mesh(mesh)(quadrature_data={"sigma": s}, batch_size=-1)

History-dependent materials need this for a Newton solve. Plasticity with a consistent tangent and a stress that depends on internal variables stored at the quadrature points is the case at hand. So far only energy() could see per-point data: it guessed the layout of element_data from shape[1] == n_quadrature, which is ambiguous when a per-element dimension equals the number of points. __call__ handed the whole [n_quadrature, ...] slice to forward, and NodeAssembler took no element data at all.

Details

• Checks. Shapes are checked up front. A key that also appears in point_data, element_data or scalar_data, or that is a built-in name (u, v, x, …), is rejected.
• Batching. The data is split together with the quadrature points when batch_size > 0.
• Compiled path. A compiled NodeAssembler falls back to the vmap path for calls with quadrature data.
• Bug fixed. energy() now slices per-point element_data along quadrature batches too. Before, batch_size > 0 with per-point history failed with a shape error.
• Docs. The forms user guide (EN + zh) lists the new dict in the data-plumbing table and the feature matrix.

Tests

tests/assemble/test_quadrature_data.py (new) checks:

• the element matrix and the node vector against manual einsum assembly from Transformation.shape_grad/JxW;
• that a constant isotropic tangent reproduces LinearElasticityElementAssembler;
• that batch sizes 1, 2 and 4 agree with the unbatched result;
• energy() with quadrature_data and with per-point element_data in batches, including its gradient with respect to the data;
• error messages;
• CUDA.

Full suite on Clariden (aarch64) in a CPU conda env with gmsh: 548 passed, 18 skipped, 0 failed. Main gives 539 passed, 17 skipped; the extra skip is the CUDA test. The new tests also pass on a GH200 node, CUDA test included. The docs build adds no warnings (EN and zh: 21 → 21).

🤖 Generated with Claude Code

https://claude.ai/code/session_01Dst8zhPLnxVpZMBhR4Pvfs