tensormesh.operator¶

Condenser¶

class Condenser(dirichlet_mask: Tensor, dirichlet_value: Tensor | None = None)[source]¶

Bases: Module

Static-condensation operator for Dirichlet boundary conditions.

Partitions a global system \(K u = f\) into inner (free) DOFs and outer (Dirichlet) DOFs and condenses the prescribed values into the right-hand side:

\[K_{ii}\, u_i = f_i - K_{io}\, u_o.\]

Parameters:

dirichlet_mask (Tensor) – 1D boolean tensor of shape \([n_{\text{dof}}]\). True marks DOFs whose value is prescribed.
dirichlet_value (Tensor, optional) – 1D tensor of shape \([n_{\text{dof}}]\) (a full vector — only the entries where dirichlet_mask is True are read) or \([n_{\text{outer\_dof}}]\) (already restricted to the boundary). Defaults to all zeros.

dirichlet_mask¶

Boolean mask of shape \([n_{\text{dof}}]\).

Type:: Tensor

dirichlet_value¶

Prescribed values restricted to the boundary, shape \([n_{\text{outer\_dof}}]\).

Type:: Tensor

inner_row, inner_col

Row/column indices of the inner block \(K_{ii}\) in local inner-DOF numbering. Populated lazily on the first call.

Type:: Tensor or None

ou2in_row, ou2in_col

Row/column indices of the coupling block \(K_{io}\) in local numbering. Populated lazily.

Type:: Tensor or None

is_inner_edge, is_ou2in_edge

Boolean masks over the matrix’s COO edge list selecting the \(K_{ii}\) / \(K_{io}\) entries.

Type:: Tensor or None

is_inner_dof, is_outer_dof

Boolean masks over the global DOFs.

Type:: Tensor or None

inner_shape, ou2in_shape

Shapes of \(K_{ii}\) and \(K_{io}\).

Type:: tuple of int or None

n_inner_dof, n_outer_dof, n_dof

DOF counts.

Type:: int or None

layout_hash¶

Sparsity-pattern hash cached from the first input matrix; used to detect a pattern change on subsequent calls.

Type:: int or None

K_ou2in¶

Cached \(K_{io}\) block; reused by condense_rhs().

Type:: SparseMatrix or None

Notes

Condenser is a torch.nn.Module. All tensor-valued attributes (dirichlet_mask, dirichlet_value, and the lazily computed index buffers) are registered as PyTorch buffers, so condenser.to(device) / condenser.cuda() / condenser.cpu() move them together with the input system.

The first call to __call__ lazily computes the inner / outer edge masks and caches them on the instance. Subsequent calls reuse the cached layout as long as the input SparseMatrix has the same sparsity pattern (checked via matrix.layout_hash). The lazy buffers are registered with persistent=False so they are not saved into state_dict.

Examples

import torch
from tensormesh import Mesh, Condenser
from tensormesh.assemble import LaplaceElementAssembler

mesh = Mesh.gen_rectangle(chara_length=0.2)
K    = LaplaceElementAssembler.from_mesh(mesh)()
f    = torch.ones(mesh.n_points, dtype=mesh.dtype)

# Homogeneous Dirichlet on the whole boundary
condenser = Condenser(mesh.boundary_mask)

# Condense: returns (K_inner, f_inner) — note this is __call__,
# NOT a separate "condense_matrix" method.
K_inner, f_inner = condenser(K, f)

# Solve the inner system and recover the full solution
u_inner = K_inner.solve(f_inner)
u       = condenser.recover(u_inner)

For time-dependent boundary data, update the prescribed values between solves via update_dirichlet(), then re-condense only the right-hand side with condense_rhs() (cheaper than rebuilding K_inner).

__init__(dirichlet_mask: Tensor, dirichlet_value: Tensor | None = None)[source]¶: Initialize internal Module state, shared by both nn.Module and ScriptModule.

update_dirichlet(dirichlet_value: Tensor)[source]¶

Replace the cached prescribed boundary values.

Useful for time-dependent or parameter-swept problems where only the right-hand side changes between solves; the cached \(K_{io}\) block (populated by __call__) is preserved.

Parameters:: dirichlet_value (Tensor) – 1D tensor of shape \([n_{\text{dof}}]\) or \([n_{\text{outer\_dof}}]\), with the same conventions as the dirichlet_value argument to __init__().

condense_rhs(rhs: Tensor) → Tensor[source]¶

Condense the right-hand side only, reusing the cached matrix layout.

\[f_i \leftarrow f_i - K_{io}\, u_o.\]

Use this after a first __call__ to re-condense f when the matrix is unchanged but the load vector changes (e.g. between time steps).

Parameters:: rhs (Tensor) – Global right-hand side of shape \([n_{\text{dof}}, \ldots]\).
Returns:: Condensed right-hand side of shape \([n_{\text{inner\_dof}}, \ldots]\).
Return type:: Tensor
Raises:: AssertionError – If __call__ has not been invoked yet: the operator has no cached \(K_{io}\) block to apply.

recover(u: Tensor) → Tensor[source]¶

Recover the full-DOF solution from an inner-DOF solution.

Scatters the condensed solution u back into the free-DOF slots and writes the prescribed boundary values into the constrained slots.

Parameters:: u (Tensor) – Inner-system solution of shape \([n_{\text{inner\_dof}}, \ldots]\).
Returns:: Full-system solution of shape \([n_{\text{dof}}, \ldots]\).
Return type:: Tensor

restrict(f: Tensor) → Tensor[source]¶

Project a full-DOF vector down to inner DOFs.

Pure linear restriction \(f_i \leftarrow f|_{\text{inner}}\), with no Dirichlet-value correction. Use this when the right-hand side has no implicit Dirichlet contribution to subtract — for example, the per-stage right-hand side of a time-integration scheme such as tensormesh.ode.ImplicitLinearRungeKutta, where the time-derivative at a Dirichlet DOF is zero by construction and so the \(-K_{io}\,u_o\) term in condense_rhs() would over-apply the boundary correction.

Unlike Condenser.__call__ / condense_rhs(), restrict does not require the matrix layout to be cached first: it only needs dirichlet_mask.

Parameters:: f (Tensor) – Full-DOF vector of shape \([n_{\text{dof}}, \ldots]\).
Returns:: Inner-DOF vector of shape \([n_{\text{inner\_dof}}, \ldots]\).
Return type:: Tensor

prolong(f_inner: Tensor) → Tensor[source]¶

Lift an inner-DOF vector up to full DOF with zeros on the boundary.

Pure linear prolongation: inner entries are scattered into the free-DOF slots, constrained slots are filled with zero — not with dirichlet_value. Use this when the quantity being lifted should vanish on the boundary regardless of the prescribed Dirichlet value, e.g. the per-stage slope of a time integrator (since a fixed-value DOF has zero time-derivative).

Like restrict(), prolong only needs dirichlet_mask and does not require the matrix layout to be cached first.

Parameters:: f_inner (Tensor) – Inner-DOF vector of shape \([n_{\text{inner\_dof}}, \ldots]\).
Returns:: Full-DOF vector of shape \([n_{\text{dof}}, \ldots]\) with zeros in the constrained slots.
Return type:: Tensor