matrixpfn 0.2.0

GNN-based learned preconditioner for sparse linear systems

MatrixPFN

GNN-based learned preconditioner for sparse linear systems. A Graph Neural Network learns to approximate the inverse matrix application $A^{-1}r$, then serves as a preconditioner inside an FGMRES solver.

Based on the GNP paper (Graph Neural Preconditioner), extended with context-based learning to generalize across unseen matrices at runtime.

Method

Sparse Matrix A
      │
      ▼
ContextResGCN ◄── context pairs (x, Ax)
      │
      ▼
MatrixPFN.apply(r) ≈ A⁻¹r
      │
      ▼
FGMRES(A, b, preconditioner=MatrixPFN)
      │
      ▼
Solution x

The key contribution over GNP is ContextResGCN: the GCN receives context pairs $(x, Ax)$ that encode information about the current matrix, enabling a single trained model to precondition unseen matrices via set_matrix() at runtime. FGMRES (Flexible GMRES) is required because the neural preconditioner is nonlinear.

Requirements

• Python 3.13
• CUDA-capable GPU (recommended)

Setup

uv sync

Training

uv run python experiments/poc_experiment.py

Evaluation

SuiteSparse Matrices

Download the 867 evaluation matrices (square, real, non-SPD, 1K–100K rows, <2M nnz) matching the GNP paper:

uv run python ../data/download_suitesparse.py

Offline Dataset Generation

Generate synthetic training matrices and save to disk:

from matrixpfn.generator import GeneratorConfig, build_training_registry
from matrixpfn.generator.offline import OfflineGenerationRunner

config = GeneratorConfig(grid_size=32)
registry = build_training_registry(config, device)
runner = OfflineGenerationRunner(registry, output_dir, num_matrices_per_domain=1000, num_context_pairs=10)
runner.run()

Architecture

src/matrixpfn/
├── nn/                  # ContextResGCN — the core GNN
├── solver/              # FGMRES + Arnoldi decomposition
├── precond/             # Preconditioner implementations
│   ├── matrix_pfn.py    #   Neural preconditioner (wraps ContextResGCN)
│   ├── jacobi.py        #   Diagonal scaling
│   ├── ilu.py           #   Incomplete LU factorization
│   ├── amg.py           #   Algebraic Multigrid (SA + AIR)
│   ├── block_jacobi.py  #   Block-diagonal LU
│   └── gmres_preconditioner.py  # Inner-outer GMRES
└── generator/           # Synthetic matrix generators (13 domains)
    ├── base.py          #   MatrixDomain enum, BatchMatrixData, MatrixGenerator protocol
    ├── registry.py      #   MatrixGeneratorRegistry + factory functions
    └── offline/         #   Save-to-disk generation pipeline

Matrix Domains

Domain	Type	Source
Diffusion	5-point stencil PDE	Synthetic
Diffusion-Advection	PDE with convection	Synthetic
Graph Laplacian	Barabási-Albert model	Synthetic
Elasticity	2D plane stress	Synthetic
Stokes	Saddle-point block structure	Synthetic
SBM	Stochastic Block Model	Synthetic
Spectral Stress	Condition number scaling	Synthetic
Variable Diffusion	Multi-grid-size, material jumps	Synthetic
Variable Advection	Variable grid + advection	Synthetic
Enhanced Diffusion	Anisotropic, holes, permutation	Synthetic
Enhanced Advection	Enhanced diffusion + advection	Synthetic
Fast Graph Laplacian	igraph-based Barabási-Albert	Synthetic
SuiteSparse	Real-world sparse matrices	SuiteSparse Collection

Reference

@article{chen2024gnp,
  title={Graph Neural Preconditioners for Iterative Solutions of Sparse Linear Systems},
  author={Chen, Jie and Hua, Yousef and Mukherjee, Subhadeep and Bai, Yu},
  journal={arXiv preprint arXiv:2406.00809},
  year={2024}
}

License

This project is part of a bachelor's thesis and is not licensed for redistribution.