def pyamgx_solve(A, b, config = None, x0 = None):
'''
Uses the (experimental) pyamgx Python bindings to the Nvidia AMGX library to solve the system Ax=b on the GPU using multigrid.
A: CSR format sparse matrix
b: numpy array
config: AMGX config. See AMGX github for details.
x0: numpy array. Initial guess for the mgrid algorithm.
Outputs a numpy array containing the solution to the equation system.
'''
pyamgx.initialize()
#pyamgx.register_print_callback(lambda msg: print(''))
try:#try-except block to call pyamgx.finalize() in the case an error occurs - subsequent calls with good inputs will fail otherwise
if config is None:
#default config copied directly from https://github.com/NVIDIA/AMGX/blob/master/core/configs/AMG_CLASSICAL_CG.json
config = get_default_pyamgx_config()
config = pyamgx.Config().create_from_dict(config)
elif isinstance(config, dict):
config = pyamgx.Config().create_from_dict(config)
resources = pyamgx.Resources()
resources.create_simple(config)
#Allocate memory for variables on GPU
A_pyamgx = pyamgx.Matrix()
A_pyamgx.create(resources, mode='dDDI')
A_pyamgx.upload_CSR(A)
if not isinstance(b,np.ndarray):
b = np.array(b)
b = b.astype(np.float64)
b_pyamgx = pyamgx.Vector()
b_pyamgx.create(resources, mode='dDDI')
b_pyamgx.upload(b)
x = pyamgx.Vector().create(resources)
x0 = x0 if x0 is not None else np.zeros(b.shape,dtype=b.dtype)
x.upload(x0)
#Solve system
solver = pyamgx.Solver()
solver.create(resources, config)
solver.setup(A_pyamgx)
solver.solve(b_pyamgx, x)
rval = x.download()
print(solver.get_residual())
#Cleanup to prevent GPU memory leak
solver.destroy()
A_pyamgx.destroy()
b_pyamgx.destroy()
x.destroy()
resources.destroy()
config.destroy()
pyamgx.finalize()
return rval
except:
pyamgx.finalize()
raise(RuntimeError('pyamgx variable creation or solver error. See stack trace.'))
def __del__(self):
import pyamgx
self._solver.destroy()
self._rhs.destroy()
self._phi_vec.destroy()
self._matrix.destroy()
self.resources.destroy()
self.cfg.destroy()
pyamgx.finalize()
def teardown_class(self):
self.rsrc.destroy()
self.cfg.destroy()
pyamgx.finalize()
def teardown(self):
pyamgx.finalize()
def teardown(self):
self.cfg.destroy()
pyamgx.finalize()
# Create solver:
solver = pyamgx.Solver().create(rsc, cfg)
# Upload system:
M = sparse.csr_matrix(np.random.rand(5, 5))
rhs = np.random.rand(5)
sol = np.zeros(5, dtype=np.float64)
A.upload_CSR(M)
b.upload(rhs)
x.upload(sol)
# Setup and solve system:
solver.setup(A)
solver.solve(b, x)
# Download solution
x.download(sol)
print("pyamgx solution: ", sol)
print("scipy solution: ", splinalg.spsolve(M, rhs))
# Clean up:
A.destroy()
x.destroy()
b.destroy()
solver.destroy()
rsc.destroy()
cfg.destroy()
pyamgx.finalize()
def teardown_class(cls):
pyamgx.finalize()