def test_get_nnz(self):
import scipy.sparse
M = pyamgx.Matrix().create(self.rsrc)
M.upload_CSR(scipy.sparse.csr_matrix(
np.array([[0., 1.], [2., 3.]])))
assert(M.get_nnz() == 3)
M.destroy()
def test_upload(self):
M = pyamgx.Matrix()
M.create(self.rsrc)
M.upload(np.array([0, 1, 3], dtype=np.int32),
np.array([1, 0, 1], dtype=np.int32),
np.array([1, 2, 3], dtype=np.float64))
M.destroy()
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
import pyamgx
pyamgx.initialize()
self.cfg = pyamgx.Config()
conf_string = f"""{{
"config_version": 2,
"solver": {{
"solver": "BICGSTAB",
"max_iters": {self.max_iter},
"monitor_residual": 1,
"tolerance": {self.tolerance},
"norm": "L2",
"print_solve_stats": 0,
"obtain_timings": 0,
"print_grid_stats": 0
}}
}}"""
self.cfg.create(conf_string)
self.resources = pyamgx.Resources().create_simple(self.cfg)
self._rhs = pyamgx.Vector().create(self.resources)
self._phi_vec = pyamgx.Vector().create(self.resources).upload(
self.phi_vec)
self._matrix = pyamgx.Matrix().create(self.resources).upload_CSR(
self.A.tocsr())
self._solver = pyamgx.Solver().create(self.resources, self.cfg)
self._solver.setup(self._matrix)
def test_upload_rectangular_device(self):
M = pyamgx.Matrix()
M.create(self.rsrc)
M.upload(cp.array([0, 1, 3], dtype=np.int32),
cp.array([1, 0, 2], dtype=np.int32),
cp.array([1, 2, 3], dtype=np.float64))
M.destroy()
def __init__(self,
config_dict,
tolerance=1e-10,
iterations=2000,
precon=None,
smoother=None,
**kwargs):
"""
:Parameters:
- `config_dict`: Dictionary specifying AMGX configuration options.
- `tolerance`: The required error tolerance.
- `iterations`: The maximum number of iterative steps to perform.
- `precon`: Preconditioner to use.
- `smoother`: Smoother to use.
- `kwargs` - Keyword arguments specifying other solver options (see AMGX Reference Manual).
"""
# update solver config:
config_dict["solver"]["tolerance"] = tolerance
config_dict["solver"]["max_iters"] = iterations
if precon:
config_dict["solver"]["preconditioner"] = precon
if smoother:
config_dict["solver"]["smoother"] = smoother
config_dict["solver"].update(kwargs)
# create AMGX objects:
self.cfg = pyamgx.Config().create_from_dict(config_dict)
self.resources = pyamgx.Resources().create_simple(self.cfg)
self.x_gpu = pyamgx.Vector().create(self.resources)
self.b_gpu = pyamgx.Vector().create(self.resources)
self.A_gpu = pyamgx.Matrix().create(self.resources)
self.solver = pyamgx.Solver().create(self.resources, self.cfg)
super(PyAMGXSolver, self).__init__(tolerance=tolerance,
iterations=iterations)
def test_upload_CSR_device(self):
M = pyamgx.Matrix()
M.create(self.rsrc)
M.upload_CSR(
cp.sparse.csr_matrix(
scipy.sparse.csr_matrix(np.array([[1., 2.], [3., 4]]))))
M.destroy()
def test_upload_CSR_rectangular(self):
M = pyamgx.Matrix()
M.create(self.rsrc)
M.upload_CSR(
scipy.sparse.csr_matrix(np.array([[1., 2., 3.], [4., 5., 6.]])))
M.destroy()
def test_get_iterations_number(self):
M = pyamgx.Matrix().create(self.rsrc)
x = pyamgx.Vector().create(self.rsrc)
b = pyamgx.Vector().create(self.rsrc)
M.upload(np.array([0, 1, 2, 3], dtype=np.int32),
np.array([0, 1, 2], dtype=np.int32), np.array([1., 1., 1.]))
x.upload(np.zeros(3, dtype=np.float64))
b.upload(np.array([1., 2., 4.], dtype=np.float64))
self.cfg.create_from_dict({'monitor_residual': 1, 'max_iters': 0})
solver = pyamgx.Solver().create(self.rsrc, self.cfg)
solver.setup(M)
solver.solve(b, x, zero_initial_guess=True)
assert (solver.iterations_number == 0)
solver.destroy()
self.cfg.create_from_dict({'max_iters': 1})
solver = pyamgx.Solver().create(self.rsrc, self.cfg)
solver.setup(M)
solver.solve(b, x)
assert (solver.iterations_number == 1)
solver.destroy()
M.destroy()
x.destroy()
b.destroy()
def test_get_residual(self):
M = pyamgx.Matrix().create(self.rsrc)
x = pyamgx.Vector().create(self.rsrc)
b = pyamgx.Vector().create(self.rsrc)
M.upload(np.array([0, 1, 2, 3], dtype=np.int32),
np.array([0, 1, 2], dtype=np.int32), np.array([1., 1., 1.]))
x.upload(np.zeros(3, dtype=np.float64))
b.upload(np.array([1., 2., 4.], dtype=np.float64))
self.cfg.create_from_dict({
'monitor_residual': 1,
'store_res_history': 1,
'tolerance': 1e-14
})
solver = pyamgx.Solver().create(self.rsrc, self.cfg)
solver.setup(M)
solver.solve(b, x)
assert (solver.get_residual() <= 1e-14)
niter = solver.iterations_number
assert (solver.get_residual() == solver.get_residual(niter))
solver.destroy()
M.destroy()
x.destroy()
b.destroy()
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 test_upload_zero_rows(self):
M = pyamgx.Matrix()
M.create(self.rsrc)
matrix = cp.sparse.csr_matrix(
scipy.sparse.csr_matrix(
np.array([[1., 2., 0.], [3., 4., 0.], [0., 0., 0.]])))
M.upload(matrix.indptr, matrix.indices, matrix.data)
M.destroy()
def test_upload_CSR_zero_rows(self):
M = pyamgx.Matrix()
M.create(self.rsrc)
M.upload_CSR(
cp.sparse.csr_matrix(
scipy.sparse.csr_matrix(
np.array([[1., 2., 0.], [3., 4., 0.], [0., 0., 0.]]))))
M.destroy()
def create(self):
self.cfg = pyamgx.Config().create_from_dict(self.config_dict)
self.resources = pyamgx.Resources().create_simple(self.cfg)
self.x_gpu = pyamgx.Vector().create(self.resources)
self.b_gpu = pyamgx.Vector().create(self.resources)
self.A_gpu = pyamgx.Matrix().create(self.resources)
self.solver = pyamgx.Solver().create(self.resources, self.cfg)
return self
def test_replace_coefficients(self):
import scipy.sparse
M = pyamgx.Matrix().create(self.rsrc)
M.upload_CSR(scipy.sparse.csr_matrix(
np.array([[0., 1.], [2., 3.]])))
M.replace_coefficients(
np.array([1., 0., 3.]))
M.destroy()
def test_upload_rectangular(self):
M = pyamgx.Matrix()
M.create(self.rsrc)
with pytest.raises(ValueError):
M.upload(
np.array([0, 1, 3], dtype=np.intc),
np.array([1, 0, 2], dtype=np.intc),
np.array([1, 2, 3], dtype=np.float64))
M.destroy()
def test_get_sizes(self):
M = pyamgx.Matrix()
M.create(self.rsrc)
M.upload(np.array([0, 1, 3], dtype=np.int32),
np.array([1, 0, 1], dtype=np.int32),
np.array([1, 2, 3], dtype=np.float64))
n, block_dims = M.get_size()
assert (n == 2)
assert (block_dims == (1, 1))
M.destroy()
def test_solve_rhs_solution_dim_mismatch(self):
M = pyamgx.Matrix().create(self.rsrc)
x = pyamgx.Vector().create(self.rsrc)
b = pyamgx.Vector().create(self.rsrc)
M.upload(np.array([0, 1, 2, 3], dtype=np.int32),
np.array([0, 1, 2], dtype=np.int32), np.array([1., 1., 1.]))
# RHS - solution mismatch
x.upload(np.zeros(4, dtype=np.float64))
b.upload(np.array([1., 2., 3], dtype=np.float64))
solver = pyamgx.Solver().create(self.rsrc, self.cfg)
solver.setup(M)
with pytest.raises(ValueError):
solver.solve(b, x)
solver.destroy()
M.destroy()
x.destroy()
b.destroy()
def test_solve_defaults(self):
M = pyamgx.Matrix().create(self.rsrc)
x = pyamgx.Vector().create(self.rsrc)
b = pyamgx.Vector().create(self.rsrc)
M.upload(np.array([0, 1, 2, 3], dtype=np.int32),
np.array([0, 1, 2], dtype=np.int32), np.array([1., 1., 1.]))
x.upload(np.zeros(3, dtype=np.float64))
b.upload(np.array([1., 2., 4.], dtype=np.float64))
solver = pyamgx.Solver().create(self.rsrc, self.cfg)
solver.setup(M)
solver.solve(b, x)
sol = np.zeros(3, dtype=np.float64)
x.download(sol)
assert_allclose(sol, np.array([1., 2., 4.]))
solver.destroy()
M.destroy()
x.destroy()
b.destroy()
def __init__(self,
config_dict,
tolerance=1e-10,
iterations=2000,
precon=None,
smoother=None,
**kwargs):
"""
Parameters
----------
config_dict : dict
AMGX configuration options
tolerance : float
Required error tolerance.
iterations : int
Maximum number of iterative steps to perform.
precon : ~fipy.solvers.pyamgx.preconditioners.preconditioners.Preconditioner, optional
smoother : ~fipy.solvers.pyamgx.smoothers.smoothers.Smoother, optional
**kwargs
Other AMGX solver options
"""
# update solver config:
config_dict["solver"]["tolerance"] = tolerance
config_dict["solver"]["max_iters"] = iterations
if precon:
config_dict["solver"]["preconditioner"] = precon
if smoother:
config_dict["solver"]["smoother"] = smoother
config_dict["solver"].update(kwargs)
# create AMGX objects:
self.cfg = pyamgx.Config().create_from_dict(config_dict)
self.resources = pyamgx.Resources().create_simple(self.cfg)
self.x_gpu = pyamgx.Vector().create(self.resources)
self.b_gpu = pyamgx.Vector().create(self.resources)
self.A_gpu = pyamgx.Matrix().create(self.resources)
self.solver = pyamgx.Solver().create(self.resources, self.cfg)
super(PyAMGXSolver, self).__init__(tolerance=tolerance,
iterations=iterations)
def test_solve_matrix_rectangular(self):
M = pyamgx.Matrix().create(self.rsrc)
x = pyamgx.Vector().create(self.rsrc)
b = pyamgx.Vector().create(self.rsrc)
''' Matrix:
1, 2, 0
2, 1, 0
'''
M.upload(np.array([0, 2, 4], dtype=np.int32),
np.array([0, 1, 0, 1], dtype=np.int32),
np.array([1., 2., 2., 1.], dtype=np.float64))
x.upload(np.zeros(3, dtype=np.float64))
b.upload(np.array([1, 2], dtype=np.float64))
solver = pyamgx.Solver().create(self.rsrc, self.cfg)
solver.setup(M)
with pytest.raises(ValueError):
solver.solve(b, x)
solver.destroy()
M.destroy()
x.destroy()
b.destroy()
def test_upload_CSR_singular(self):
M = pyamgx.Matrix()
M.create(self.rsrc)
M.upload_CSR(
scipy.sparse.csr_matrix(np.zeros([3, 3], dtype=np.float64)))
M.destroy()
def test_create_and_destroy(self):
M = pyamgx.Matrix()
M.create(self.rsrc)
M.destroy()
"determinism_flag": 1,
"exception_handling": 1,
"solver": {
"monitor_residual": 1,
"solver": "BICGSTAB",
"convergence": "RELATIVE_INI_CORE",
"preconditioner": {
"solver": "NOSOLVER"
}
}
})
rsc = pyamgx.Resources().create_simple(cfg)
# Create matrices and vectors:
A = pyamgx.Matrix().create(rsc)
b = pyamgx.Vector().create(rsc)
x = pyamgx.Vector().create(rsc)
# 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)
def __init__(self, A, linear_solver, env):
if linear_solver is 'lu':
self.A = A.tocsc()
elif linear_solver is 'amgx':
import pyamgx
pyamgx.initialize()
hA = A.tocsr()
AMGX_CONFIG_FILE_NAME = 'amgx_config/PCGF_CLASSICAL_AGGRESSIVE_PMIS_JACOBI.json'
if False:
cfg = pyamgx.Config().create_from_file(AMGX_CONFIG_FILE_NAME)
else:
cfg = pyamgx.Config().create_from_dict({
"config_version": 2,
"determinism_flag": 0,
"solver": {
"preconditioner": {
"print_grid_stats": c.print_stats,
"algorithm": "AGGREGATION",
"print_vis_data": 0,
"solver": "AMG",
"smoother": {
"relaxation_factor": 0.8,
"scope": "jacobi",
"solver": "BLOCK_JACOBI",
"monitor_residual": 0,
"print_solve_stats": 0
},
"print_solve_stats": 0,
"presweeps": 2,
"selector": "SIZE_2",
"coarse_solver": "NOSOLVER",
"max_iters": 2,
"monitor_residual": 0,
"store_res_history": 0,
"scope": "amg_solver",
"max_levels": 1000,
"postsweeps": 2,
"cycle": "V"
},
"solver": "PCGF",
"print_solve_stats": c.print_stats,
"obtain_timings": c.print_stats,
"max_iters": c.max_iters,
"monitor_residual": 1,
"convergence": "RELATIVE_INI",
"scope": "main",
"tolerance": c.err_tol,
"norm": "L2"
}
})
rsc = pyamgx.Resources().create_simple(cfg)
mode = 'dDDI'
# Create solver:
self.amgx = pyamgx.Solver().create(rsc, cfg, mode)
# Create matrices and vectors:
self.d_A = pyamgx.Matrix().create(rsc, mode)
self.d_x = pyamgx.Vector().create(rsc, mode)
self.d_b = pyamgx.Vector().create(rsc, mode)
self.d_A.upload_CSR(hA)
# Setup and solve system:
# self.amgx.setup(d_A)
## Clean up:
#A.destroy()
#x.destroy()
#b.destroy()
#self.amgx.destroy()
#rsc.destroy()
#cfg.destroy()
#pyamgx.finalize()
elif linear_solver is 'cg' or 'amg':
pass
else:
raise ValueError("Invalid solver choice.")
def __init__(self, A, linear_solver, env):
if linear_solver is 'lu':
self.A = A.tocsc()
self.lu = splu(self.A)
elif linear_solver is 'cg':
self.A = A.tocsr()
elif linear_solver in ['cudaCG']:
self.A = A.tocsr()
self.dData = env.cuda.to_device(self.A.data)
self.dPtr = env.cuda.to_device(self.A.indptr)
self.dInd = env.cuda.to_device(self.A.indices)
self.cuSparseDescr = env.cuSparse.matdescr()
elif linear_solver in ['cudaPCG']:
self.A = A.tocsr()
self.Adescr = env.cuSparse.matdescr()
self.Adata = env.cuda.to_device(self.A.data)
self.Aptr = env.cuda.to_device(self.A.indptr)
self.Aind = env.cuda.to_device(self.A.indices)
A_t = self.A.copy()
A_t = -A_t # Make it positive definite
A_t.data = np.where(A_t.nonzero()[0] >= A_t.nonzero()[1], A_t.data,
0.)
A_t.eliminate_zeros()
A_t_descr = env.cuSparse.matdescr(matrixtype='S', fillmode='L')
info = env.cuSparse.csrsv_analysis(trans='N', m=A_t.shape[0], nnz=A_t.nnz, \
descr=A_t_descr, csrVal=A_t.data, \
csrRowPtr=A_t.indptr, csrColInd=A_t.indices)
env.cuSparse.csric0(trans='N', m=A_t.shape[0], \
descr=A_t_descr, csrValM=A_t.data, csrRowPtrA=A_t.indptr,\
csrColIndA=A_t.indices, info=info)
self.L = A_t
self.Lmv_descr = env.cuSparse.matdescr()
# self.Lsv_descr = cuSparse.matdescr(matrixtype='T', fillmode='L')
self.Lsv_descr = env.cuSparse.matdescr(matrixtype='T')
self.Ldata = env.cuda.to_device(self.L.data)
self.Lptr = env.cuda.to_device(self.L.indptr)
self.Lind = env.cuda.to_device(self.L.indices)
self.Lsv_info = env.cuSparse.csrsv_analysis(trans='N', m=self.L.shape[0], \
nnz=self.L.nnz, descr=self.Lsv_descr, csrVal=self.Ldata, \
csrRowPtr=self.Lptr, csrColInd=self.Lind)
self.LT = self.L.transpose()
self.LT.tocsr()
self.LTmv_descr = env.cuSparse.matdescr()
# self.LTsv_descr = env.cuSparse.matdescr(matrixtype='T', fillmode='U')
self.LTsv_descr = env.cuSparse.matdescr()
self.LTdata = env.cuda.to_device(self.LT.data)
self.LTptr = env.cuda.to_device(self.LT.indptr)
self.LTind = env.cuda.to_device(self.LT.indices)
self.LTsv_info = env.cuSparse.csrsv_analysis(trans='T', m=self.L.shape[0], \
nnz=self.L.nnz, descr=self.Lsv_descr, csrVal=self.Ldata, \
csrRowPtr=self.Lptr, csrColInd=self.Lind)
elif linear_solver is 'pcg':
self.A = A.tocsr()
A_t = -self.D2s
elif linear_solver is 'amg':
self.A = A.tocsr()
self.A_spd = self.A.copy()
self.A_spd = -self.A_spd
self.B = np.ones((self.A_spd.shape[0], 1), dtype=self.A_spd.dtype)
self.BH = self.B.copy()
if self.A_spd.shape[0] in [40962, 163842, 655362]:
self.A_amg = rootnode_solver(self.A_spd, B=self.B, BH=self.BH,
strength=('evolution', {'epsilon': 2.0, 'k': 2, 'proj_type': 'l2'}),
smooth=('energy', {'weighting': 'local', 'krylov': 'cg', 'degree': 2, 'maxiter': 3}),
improve_candidates=[('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 4}), \
None, None, None, None, None, None, None, None, None, None, \
None, None, None, None],
aggregate="standard",
presmoother=('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 1}),
postsmoother=('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 1}),
max_levels=15,
max_coarse=300,
coarse_solver="pinv")
elif self.A_spd.shape[0] in [81920, 327680, 1310720, 5242880]:
self.A_amg = rootnode_solver(self.A_spd, B=self.B, BH=self.BH,
strength=('evolution', {'epsilon': 4.0, 'k': 2, 'proj_type': 'l2'}),
smooth=('energy', {'weighting': 'local', 'krylov': 'cg', 'degree': 2, 'maxiter': 3}),
improve_candidates=[('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 4}), \
None, None, None, None, None, None, None, None, None, None, \
None, None, None, None],
aggregate="standard",
presmoother=('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 1}),
postsmoother=('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 1}),
max_levels=15,
max_coarse=300,
coarse_solver="pinv")
elif self.A_spd.shape[0] in [2621442]:
self.A_amg = rootnode_solver(self.A_spd, B=self.B, BH=self.BH,
strength=('evolution', {'epsilon': 4.0, 'k': 2, 'proj_type': 'l2'}),
smooth=('energy', {'weighting': 'local', 'krylov': 'cg', 'degree': 3, 'maxiter': 4}),
improve_candidates=[('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 4}), \
None, None, None, None, None, None, None, None, None, None, \
None, None, None, None],
aggregate="standard",
presmoother=('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 1}),
postsmoother=('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 1}),
max_levels=15,
max_coarse=300,
coarse_solver="pinv")
else:
print("Unknown matrix. Using a generic AMG solver")
self.A_amg = rootnode_solver(self.A_spd, B=self.B, BH=self.BH,
strength=('evolution', {'epsilon': 2.0, 'k': 2, 'proj_type': 'l2'}),
smooth=('energy', {'weighting': 'local', 'krylov': 'cg', 'degree': 2, 'maxiter': 3}),
improve_candidates=[('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 4}), \
None, None, None, None, None, None, None, None, None, None, \
None, None, None, None],
aggregate="standard",
presmoother=('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 1}),
postsmoother=('block_gauss_seidel', {'sweep': 'symmetric', 'iterations': 1}),
max_levels=15,
max_coarse=300,
coarse_solver="pinv")
elif linear_solver is 'amgx':
import pyamgx
pyamgx.initialize()
hA = A.tocsr()
if hA.nnz * 1. / hA.shape[0] > 5.5: # Primary mesh
AMGX_CONFIG_FILE_NAME = 'amgx_config/PCGF_CLASSICAL_AGGRESSIVE_PMIS_JACOBI.json'
if hA.nnz * 1. / hA.shape[0] < 5.5: # Dual mesh
AMGX_CONFIG_FILE_NAME = 'amgx_config/PCGF_CLASSICAL_AGGRESSIVE_PMIS.json'
else:
print(
'Error: cannot determine primary or dual mesh, not sure which config to use.'
)
cfg = pyamgx.Config().create_from_file(AMGX_CONFIG_FILE_NAME)
rsc = pyamgx.Resources().create_simple(cfg)
mode = 'dDDI'
# Create solver:
self.amgx = pyamgx.Solver().create(rsc, cfg, mode)
# Create matrices and vectors:
d_A = pyamgx.Matrix().create(rsc, mode)
self.d_x = pyamgx.Vector().create(rsc, mode)
self.d_b = pyamgx.Vector().create(rsc, mode)
d_A.upload(hA.indptr, hA.indices, hA.data)
# Setup and solve system:
self.amgx.setup(d_A)
## Clean up:
#A.destroy()
#x.destroy()
#b.destroy()
#self.amgx.destroy()
#rsc.destroy()
#cfg.destroy()
#pyamgx.finalize()
else:
raise ValueError("Invalid solver choice.")
def __init__(self, vc, g, c):
# load appropriate module for working with objects on CPU / GPU
if c.use_gpu:
if not c.linear_solver is 'amgx':
raise ValueError("Invalid solver choice.")
import cupy as xp
from cupyx.scipy.sparse import coo_matrix, csc_matrix, csr_matrix, eye, diags, bmat
areaCell_cpu = g.areaCell.get()
else:
if c.linear_solver is 'amgx':
raise ValueError("Invalid solver choice.")
import numpy as xp
from scipy.sparse import coo_matrix, csc_matrix, csr_matrix, eye, diags, bmat
areaCell_cpu = g.areaCell
# Construct matrix blocks of the coupled elliptic system
# A diagonal matrix representing scaling by cell areas
mAreaCell = diags(g.areaCell, 0, format='csr')
mAreaCell_phi = mAreaCell.copy()
mAreaCell_phi[0, 0] = 0.
#mAreaCell_phi.eliminate_zeros( )
if c.on_a_global_sphere:
mAreaCell_psi = mAreaCell_phi.copy()
else:
areaCell_psi = g.areaCell.copy()
areaCell_psi[g.cellBoundary - 1] = 0.
mAreaCell_psi = diags(areaCell_psi, 0, format='csr')
#mAreaCell_psi.eliminate_zeros( )
## Construct the coefficient matrix for the coupled elliptic
## system for psi and phi, using the normal vector
# Left, row 1
self.AMC = mAreaCell_psi * vc.mVertex2cell * vc.mCurl_t
self.AC = mAreaCell_psi * vc.mCurl_v
#self.AMC.eliminate_zeros( )
self.AMC.sort_indices()
#self.AC.eliminate_zeros( )
self.AC.sort_indices()
# Left, row 2
self.AMD = mAreaCell_phi * vc.mVertex2cell * vc.mDiv_t
self.AD = mAreaCell_phi * vc.mDiv_v
#self.AMD.eliminate_zeros( )
self.AMD.sort_indices()
#self.AD.eliminate_zeros( )
self.AD.sort_indices()
# Right, col 2
self.GN = vc.mGrad_tn * vc.mCell2vertex_n
#self.GN.eliminate_zeros( )
self.GN.sort_indices()
# Right, col 1
self.SN = vc.mSkewgrad_nd * vc.mCell2vertex_psi
#self.SN.eliminate_zeros( )
self.SN.sort_indices()
## Construct an artificial thickness vector
thickness_edge = 100 * (10. + xp.random.rand(g.nEdges))
self.thicknessInv = 1. / thickness_edge
# self.mThicknessInv = eye(g.nEdges)
# self.mThicknessInv.data[0,:] = 1./thickness_edge
# Copy certain matrices over from VectorCalculus; maybe unnecessary
# in the future.
self.mSkewgrad_td = vc.mSkewgrad_td.copy()
self.mGrad_n_n = vc.mGrad_n_n.copy()
if c.linear_solver is 'amgx':
import pyamgx
pyamgx.initialize()
err_tol = c.err_tol * 1e-5 * np.mean(areaCell_cpu) * np.sqrt(
g.nCells) # For vorticity
cfg1 = pyamgx.Config().create_from_dict({
"config_version": 2,
"determinism_flag": 0,
"solver": {
"preconditioner": {
"print_grid_stats": c.print_stats,
"algorithm": "AGGREGATION",
"print_vis_data": 0,
"solver": "AMG",
"smoother": {
"relaxation_factor": 0.8,
"scope": "jacobi",
"solver": "BLOCK_JACOBI",
"monitor_residual": 0,
"print_solve_stats": 0
},
"print_solve_stats": 0,
"presweeps": 2,
"selector": "SIZE_2",
"coarse_solver": "NOSOLVER",
"max_iters": 2,
"monitor_residual": 0,
"store_res_history": 0,
"scope": "amg_solver",
"max_levels": 100,
"postsweeps": 2,
"cycle": "V"
},
"solver": "PCGF",
"print_solve_stats": c.print_stats,
"obtain_timings": c.print_stats,
"max_iters": c.max_iters,
"monitor_residual": 1,
"convergence": "ABSOLUTE",
"scope": "main",
"tolerance": err_tol,
"norm": "L2"
}
})
# Smaller error tolerance for divergence because geophysical flows
# are largely nondivergent
err_tol = c.err_tol * 1e-6 * np.mean(areaCell_cpu) * np.sqrt(
g.nCells)
cfg2 = pyamgx.Config().create_from_dict({
"config_version": 2,
"determinism_flag": 0,
"solver": {
"preconditioner": {
"print_grid_stats": c.print_stats,
"algorithm": "AGGREGATION",
"print_vis_data": 0,
"solver": "AMG",
"smoother": {
"relaxation_factor": 0.8,
"scope": "jacobi",
"solver": "BLOCK_JACOBI",
"monitor_residual": 0,
"print_solve_stats": 0
},
"print_solve_stats": 0,
"presweeps": 2,
"selector": "SIZE_2",
"coarse_solver": "NOSOLVER",
"max_iters": 2,
"monitor_residual": 0,
"store_res_history": 0,
"scope": "amg_solver",
"max_levels": 100,
"postsweeps": 2,
"cycle": "V"
},
"solver": "PCGF",
"print_solve_stats": c.print_stats,
"obtain_timings": c.print_stats,
"max_iters": c.max_iters,
"monitor_residual": 1,
"convergence": "ABSOLUTE",
"scope": "main",
"tolerance": err_tol,
"norm": "L2"
}
})
rsc1 = pyamgx.Resources().create_simple(cfg1)
rsc2 = pyamgx.Resources().create_simple(cfg2)
mode = 'dDDI'
# Create solver:
self.slv11 = pyamgx.Solver().create(rsc1, cfg1, mode)
self.slv22 = pyamgx.Solver().create(rsc2, cfg2, mode)
# Create matrices and vectors:
self.d_A11 = pyamgx.Matrix().create(rsc1, mode)
self.d_x = pyamgx.Vector().create(rsc1, mode)
self.d_b1 = pyamgx.Vector().create(rsc1, mode)
self.d_A22 = pyamgx.Matrix().create(rsc2, mode)
self.d_y = pyamgx.Vector().create(rsc2, mode)
self.d_b2 = pyamgx.Vector().create(rsc2, mode)
elif c.linear_solver is 'amg':
from pyamg import rootnode_solver
elif c.linear_solver is 'lu':
pass
else:
raise ValueError("Invalid solver choice.")