稀疏求解器¶

完成装配和边界条件凝聚之后，你便得到一个待求解的稀疏系统——线性或非线性。TensorMesh 将这项工作委托给 torch-sla，一个独立的可微稀疏线性代数库。同一份有限元代码只需改动一个关键字参数，即可在 CPU 与 GPU、迭代法与直接法求解器之间切换；非线性系统则由同一后端通过 Newton / Picard / Anderson 接口驱动。

安装 `torch-sla`¶

重要

torch-sla 是 tensormesh.sparse 的硬性、导入期依赖。缺少它该模块将无法导入。它同时也是当前及未来所有求解器工作落地的库——TensorMesh 将 torch-sla 的版本作为权威的稀疏线性代数层来跟随，我们建议保持其更新到最新。

基础 wheel 附带 CPU 技术栈（SciPy / 原生 PyTorch Krylov）。额外的后端作为可选 extras 提供：

pip install torch-sla              # CPU stack (scipy + pytorch backends)
pip install "torch-sla[cudss]"     # + NVIDIA cuDSS (fastest GPU direct)
pip install "torch-sla[pyamg]"     # + PyAMG algebraic multigrid (CPU setup)
pip install "torch-sla[all]"       # every PyPI-installable runtime backend

两个原生编译后端——strumpack（可移植多波前直接法，CPU/CUDA/ROCm）和 amgx（NVIDIA AmgX AMG/Krylov）——以预编译 wheel 的形式发布在 torch-sla 的 GitHub Releases 而非 PyPI 上；wheel 选择规则参见 torch-sla README。TensorMesh 也在自身的安装层镜像了 PyPI extras（pip install "tensormesh-fem[cudss]" 等）——详见安装。

查看可用后端¶

安装完成后，torch_sla.show_backends() 会打印一张状态表，显示当前机器上哪些后端可用，以及如何安装尚缺失的后端：

>>> import torch_sla
>>> torch_sla.show_backends()
torch-sla backend status (CUDA: available)
  scipy      [CPU]           available
  pytorch    [CPU/CUDA]      available
  cudss      [CUDA]          available
  pyamg      [CPU]           available
  amgx       [CUDA]          not available — pip install torch-sla[amgx]
  strumpack  [CPU/CUDA/ROCm] not available — pip install torch-strumpack

torch-sla 为 TensorMesh 提供了什么：

SparseTensor —— SparseMatrix 所继承的 COO 数据类型。单元装配会直接交给你一个 SparseMatrix，因此父类上的所有算子方法（@、.solve、.nonlinear_solve、.is_symmetric 等）都无需任何额外包装即可使用。
一个带有自定义反向传播（伴随稀疏求解）的统一 solve 算子，使梯度能够穿过每一个线性系统。
一个可插拔的后端层，在单一分派入口点之后涵盖 CPU、GPU、迭代法与直接法求解器。

求解 `A x = b`：`.solve()` 方法¶

权威入口点是 SparseMatrix.solve —— 原封不动地继承自 torch_sla.SparseTensor。装配本身就已返回一个 SparseMatrix，所以通常你只需要这一个调用：

K = LaplaceElementAssembler.from_mesh(mesh)()
x = K.solve(b)                              # auto backend + method
x = K.solve(b, backend="cudss")             # force NVIDIA GPU direct
x = K.solve(b, method="bicgstab")           # force a specific iterative method

对称性 / 正定性会被自动检测。 每次调用时，torch-sla 都会对矩阵计算 is_symmetric() 和 is_positive_definite()，然后为对称正定（SPD）系统选用 CG / Cholesky，否则选用 BiCGStab / LU。你不需要传入 is_spd 提示；即便传入也会被忽略。

关键关键字参数（完整列表见 torch-sla 参考文档）：

backend："auto"（根据设备、数据类型和问题规模自动选择——CPU → SciPy，CUDA → cuDSS（若可用），超大规模系统 → 原生 PyTorch Krylov），或取 "scipy"、"pytorch"、"cudss"、"strumpack"、"pyamg"、"amgx" 之一。
method："auto"（根据矩阵性质和后端自动选择），或迭代法（"cg"、"bicgstab"、"gmres"、"minres"、"lsqr"、"lsmr"）、直接分解（"lu"、"umfpack"、"cholesky"、"ldlt"）或 AMG 循环（"amg"）。并非每个后端都支持所有方法——见下表。
atol（默认 1e-10）、tol（默认 1e-12）、maxiter（默认 10000），用于迭代收敛。
verbose=True 会打印一行摘要，给出自动选择的后端 / 方法以及检测到的矩阵性质——当你想确认 "auto" 是否选中了你预期的路径时非常方便：
```
>>> x = K.solve(b, verbose=True)
[torch-sla] solve: n=1024, nnz=7168, dtype=float64, device=cpu, symmetric=True, spd=True, backend=scipy, method=lu
```

支持的后端¶

六个已验证后端（均通过与参考解的对比检查，达到/接近机器精度的相对残差）：

后端	字符串	设备	直接 / 迭代	方法
SciPy	`"scipy"`	CPU	两者	`lu`、`umfpack`、`cg`、`bicgstab`、`gmres`、`lgmres`、`minres`、`qmr`。CPU 默认。
原生 PyTorch	`"pytorch"`	CPU / CUDA / ROCm	迭代	`cg`、`bicgstab`、`gmres`、`minres`、`lsqr`、`lsmr`——带 Jacobi 预处理的 Krylov 方法，纯 torch 实现，设备无关（包括 AMD ROCm）。超大规模问题路径（> 200 万自由度）。
cuDSS	`"cudss"`	CUDA	直接	`lu`、`cholesky`、`ldlt`。NVIDIA cuDSS——最快的 GPU 直接法路径；CUDA 默认（内存允许时）。
STRUMPACK	`"strumpack"`	CPU / CUDA / ROCm	直接	`lu`——可移植多波前直接求解器（AMD 上的 GPU 直接法选项）。wheel 来自 GitHub Releases（`torch-strumpack`）。
PyAMG	`"pyamg"`	CPU 环境	迭代（AMG）	`amg`、`ruge_stuben`、`smoothed_aggregation`——面向偏微分方程系统的代数多重网格 V 循环。
AmgX	`"amgx"`	CUDA	迭代（AMG / Krylov）	`amg`、`cg`/`pcg`、`bicgstab`/`pbicgstab`、`gmres`/`fgmres`——NVIDIA AmgX。wheel 来自 GitHub Releases（`torch-amgx`）。

按问题规模选择（基于 torch-sla 的泊松基准测试，已验证至 4 亿自由度多 GPU）：

问题规模	CPU	CUDA
< 200 万自由度	`scipy` + `lu`	`cudss` + `cholesky`
200 万 – 1.69 亿自由度	`pytorch` + `cg`	`pytorch` + `cg`
> 1.69 亿自由度	`DSparseMatrix` 多进程	`DSparseMatrix` 多 GPU

经验法则：直接求解器可达到机器精度，但其内存填充量将其上限限制在约 200 万自由度；迭代路径收敛至约 1e-6，但具有近线性扩展性。迭代收敛请使用 float64。分布式行参见分布式有限元方法。

批量右端项¶

当 b 的形状为 [n_dof, n_batch] 时，torch-sla 会自动将分解的开销摊销到整个 batch 上——一次分解、n_batch 次回代——而不是对每一列独立运行一遍迭代求解器。这正是 tensormesh.dataset 机器学习工作流的主力模式。

B = torch.randn(K.shape[0], 64)             # 64 right-hand sides
X = K.solve(B)                              # one factorization

关于批量化的更全面介绍，参见批量化工作流。

非线性系统：`.nonlinear_solve()`¶

对于残差 F(u, params) = 0 关于 u 非线性的问题——超弹性、塑性、相场——请使用 SparseMatrix.nonlinear_solve（继承自 torch_sla.SparseTensor）。矩阵会被自动传入残差闭包，雅可比矩阵通过 autograd 获得，反向传播则使用伴随方法。

from torch_sla import SparseTensor  # SparseMatrix also works

# Nonlinear PDE: A @ u + u**2 = f
def residual(u, A, f):
    return A @ u + u**2 - f

f  = torch.randn(n, requires_grad=True)
u0 = torch.zeros(n)

u = A.nonlinear_solve(residual, u0, f, method='newton')

# Gradients flow via the implicit-function theorem
loss = u.sum()
loss.backward()
print(f.grad)        # dL/df

关键关键字参数：

method："newton"（默认，带 Armijo 线搜索的牛顿-拉夫森法）、"picard"（不动点迭代）或 "anderson"（Anderson 加速）。
line_search（默认 True）——开关 Armijo 回溯。
tol / atol / max_iter 用于外层非线性循环。
linear_solver / linear_method 用于选择每次内层线性求解所用的后端。

前向传播运行所选的非线性迭代；反向传播只求解单个伴随系统，因此无论前向走了多少次非线性迭代，u 关于任何 requires_grad 参数（包括 A.values）的梯度大约只额外花费一次线性求解。

更底层的入口点（遗留）¶

自 0.x 版本弃用: tensormesh.sparse.spsolve() 和 tensormesh.sparse.nonlinear_solve() 是早于 torch-sla 集成的 TensorMesh 内部回退实现。它们已计划移除；权威的求解器路径是上文介绍的 SparseMatrix（即 torch_sla.SparseTensor）方法。

tensormesh.sparse.spsolve() 是一个自由函数，它直接操作原始 COO 数组 (edata, row, col, shape, b) 而非 SparseMatrix 对象。实际上这很少用得上——装配会产生一个 SparseMatrix，凝聚也会产生一个 SparseMatrix，而 K.solve(b) 已涵盖本指南中的每一种有限元工作流。一旦剩余的调用点都迁移到 SparseMatrix.solve / torch_sla.spsolve，tensormesh.sparse.spsolve 即会退役。

tensormesh.sparse.nonlinear_solve() 是库内自带的牛顿-拉夫森辅助函数。它要求用户提供一个显式的雅可比闭包，且缺少 torch_sla.SparseTensor.nonlinear_solve 所提供的线搜索、Picard / Anderson 模式以及基于 autograd 的雅可比装配。请迁移到 A.nonlinear_solve(residual, u0, *params)。

下一步¶

时间积分 —— 将线性求解封装进时间步进循环，用于瞬态问题。
批量化工作流 —— 批量化的三个不同维度，以及各自的适用场景。
可微性 —— 伴随反向传播如何工作，以及如何把一个可学习参数接入稀疏求解。