Awesome
Optimization-Based Algebraic Multigrid Coarsening Using Reinforcement Learning (NeurIPS 2021)
Code for reproducing the experimental results of our paper: https://arxiv.org/pdf/2106.01854.pdf
Requirements
-
Python = 3.8.10
-
PyTorch = 1.9.0
-
Pytorch_Geometric = 1.7.2
-
pyamg = 4.1.0
-
networkx = 2.6.2
-
pygmsh = 7.1.6
-
gmsh = 4.8.0
Training
python main.py
Files
- Batch_Graph.py: The class for making a batch of grids to be the input to the agent's network.
- DuelingNet.py: The netwrok architecture that the agent utilizes (TAGCN network).
- Lloyd_Unstructured.py: The code for the generating Lloyd agregations on grids.
- MG_Agent.py: The class for Dueling Double DQN agent.
- Optim.py: The code for running the grid coarsinging using Lloyed aggregation.
- Solve_MG.py: The code for the two grid cycle AMG algorithm.
- Unstructured.py: The class for defining grids and making unstructured meshes.
- fem.py: Constructs a finite element discretization of a 2D Poisson problem.
- main.py: The main driver for training the RL agent and test functions.