# 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](https://arxiv.org/abs/2406.00809v3) (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

```bash
uv sync
```

## Training

```bash
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:

```bash
uv run python ../data/download_suitesparse.py
```

### Offline Dataset Generation

Generate synthetic training matrices and save to disk:

```python
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](https://sparse.tamu.edu/) |

## Reference

```bibtex
@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.
