Skip to content

Commit

Permalink
redo swe complex-proxy script and add composite pc
Browse files Browse the repository at this point in the history
  • Loading branch information
@JHopeCollins
JHopeCollins committed May 3, 2024
1 parent c91ec03 commit c4a00e4
Showing 1 changed file with 85 additions and 27 deletions.
112 changes: 85 additions & 27 deletions case_studies/shallow_water/blockpc/swe_cpx.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,9 @@ def initialize(self, pc):
v = fd.TestFunction(V)
_, q = fd.split(v)

# doesn't actually matter what the form is because
# the HybridisedSCPC pulls form_{mass,function} off
# the appctx and uses those.
Af = fd.inner(w, v)*fd.dx
Lf = fd.inner(q, self.xf)*fd.dx

Expand Down Expand Up @@ -129,10 +132,10 @@ def read_checkpoint(checkpoint_name, funcname, index, ref_level=0):
parser.add_argument('--nt', type=int, default=16, help='Number of timesteps (used to calculate the circulant eigenvalues).')
parser.add_argument('--theta', type=float, default=0.5, help='Parameter for implicit theta method. 0.5 for trapezium rule, 1 for backwards Euler (used to calculate the circulant eigenvalues).')
parser.add_argument('--alpha', type=float, default=1e-3, help='Circulant parameter (used to calculate the circulant eigenvalues).')
parser.add_argument('--eigenvalue', type=int, default=0, help='Index of the circulant eigenvalues to use for the complex coefficients.')
parser.add_argument('--eigenvalue', type=int, default=-1, help='Index of the circulant eigenvalues to use for the complex coefficients. -1 for all')
parser.add_argument('--seed', type=int, default=12345, help='Seed for the random right hand side.')
parser.add_argument('--nrhs', type=int, default=1, help='Number of random right hand sides to solve for.')
parser.add_argument('--method', type=str, default='mg', choices=['lu', 'mg', 'lswe', 'hybr', 'schur'], help='Preconditioning method to use.')
parser.add_argument('--method', type=str, default='mg', choices=['lu', 'mg', 'lswe', 'hybr', 'schur', 'composite'], help='Preconditioning method to use.')
parser.add_argument('--rtol', type=float, default=1e-5, help='Relative tolerance for solution of each block.')
parser.add_argument('--foutname', type=str, default='iterations', help='Name of output file to write iteration counts.')
parser.add_argument('--checkpoint', type=str, default='swe_series', help='Name of checkpoint file.')
Expand All @@ -149,10 +152,18 @@ def read_checkpoint(checkpoint_name, funcname, index, ref_level=0):
if args.show_args:
Print(args)

verbosity = 0
if args.v:
verbosity = 1
if args.vv:
verbosity = 2

mesh, u, f, b = read_checkpoint(args.checkpoint, args.funcname,
args.index, args.ref_level)

x = fd.SpatialCoordinate(mesh)
outward_normals = fd.CellNormal(mesh)
facet_normals = fd.FacetNormal(mesh)

# case parameters
g = earth.Gravity
Expand All @@ -171,7 +182,7 @@ def form_function(u, h, v, q, t=None):
u, h, v, q, t)


def aux_form_function(u, h, v, q, t):
def aux_form_function(u, h, v, q, t=None):
return swe.linear.form_function(mesh, g, H, f,
u, h, v, q, t)

Expand Down Expand Up @@ -214,9 +225,9 @@ def aux_form_function(u, h, v, q, t):
A = M + K

# random rhs
# use one-form because Cofunction doesn't work with fieldsplit
l = fd.Function(W)
L = fd.inner(fd.TestFunction(W), l)*fd.dx
# L = fd.Cofunction(W.dual())

# PETSc solver parameters

Expand All @@ -237,7 +248,7 @@ def aux_form_function(u, h, v, q, t):
'mg': {
'levels': {
'ksp_type': 'gmres',
'ksp_max_it': 3,
'ksp_max_it': 4,
'pc_type': 'python',
'pc_python_type': 'firedrake.PatchPC',
'patch': {
Expand Down Expand Up @@ -281,24 +292,64 @@ def aux_form_function(u, h, v, q, t):
"aux": lu_params,
}

# using fgmres on the schur complement results in an
# outer residual drop of approximate the same factor
# as the schur complement residual.
schurhybr_sparams = {
'pc_type': 'fieldsplit',
'pc_fieldsplit_type': 'schur',
'pc_fieldsplit_schur_fact_type': 'full',
'fieldsplit_0': {
'ksp_type': 'preonly',
'pc_type': 'lu',
'pc_factor_mat_solver_type': 'mumps',
},
'fieldsplit_0': lu_params,
'fieldsplit_1': {
'ksp_type': 'gmres',
'ksp_type': 'fgmres',
'ksp_converged_rate': None,
'gmres_restart': 60,
'ksp_rtol': 1e-2,
'ksp_rtol': 0.9*args.rtol,
'pc_type': 'python',
'pc_python_type': f'{__name__}.ApproxHybridPC',
},
}

# linear operator is much more effective with gmres
# but this means we need fgmres on the patch smoother
composite_sparams = {
"pc_type": "composite",
"pc_composite_type": "multiplicative",
"pc_composite_pcs": "ksp,ksp",
"sub_0": {
"ksp_ksp_type": "gmres",
"ksp_ksp_max_it": 2,
"ksp_ksp_convergence_test": 'skip',
"ksp_ksp_converged_maxits": None,
"ksp": aux_sparams,
},
"sub_1_ksp": {
"ksp_type": "fgmres",
"ksp_max_it": 1,
"ksp_convergence_test": 'skip',
"ksp_converged_maxits": None,
"pc_type": "python",
"pc_python_type": "firedrake.PatchPC",
'patch': {
'pc_patch': {
'save_operators': True,
'partition_of_unity': True,
'sub_mat_type': 'seqdense',
'construct_dim': 0,
'construct_type': 'star',
'local_type': 'additive',
'precompute_element_tensors': True,
'symmetrise_sweep': False
},
'sub': {
'ksp_type': 'preonly',
'pc_type': 'lu',
'pc_factor_shift_type': 'nonzero',
}
}
}
}

sparams = {
'ksp_rtol': args.rtol,
'ksp_type': 'fgmres',
Expand All @@ -314,20 +365,22 @@ def aux_form_function(u, h, v, q, t):
sparams.update(hybridization_sparams)
elif args.method == 'schur':
sparams.update(schurhybr_sparams)
elif args.method == 'composite':
sparams.update(composite_sparams)
else:
raise ValueError(f"Unknown method {args.method}")

if args.v or args.vv:
if verbosity > 0:
sparams["ksp_converged_rate"] = None
if args.vv:
if verbosity > 1:
sparams["ksp_monitor"] = None

appctx = {
'cpx': cpx,
'uref': wref,
'tref': None,
'd1': d1c,
'd2': d2c,
'd1': d1,
'd2': d2,
'bcs': [],
'form_mass': form_mass,
'form_function': form_function,
Expand All @@ -339,31 +392,36 @@ def aux_form_function(u, h, v, q, t):

problem = fd.LinearVariationalProblem(A, L, wout,
constant_jacobian=True)
solver = fd.LinearVariationalSolver(problem, appctx=appctx,
options_prefix="block",
solver_parameters=sparams)

fstr = lambda x, n: str(round(x, n)).rjust(3+n)

neigs = args.nt//2+1
neigs = args.nt//2+1 if args.eigenvalue < 0 else 1
eigs = [*range(neigs)] if args.eigenvalue < 0 else [args.eigenvalue,]
nits = np.zeros((neigs, args.nrhs), dtype=int)
for i in range(neigs):

freq = freqs[i]
d1 = D1[i]
d2 = D2[i]
k = eigs[i]
freq = freqs[k]
d1 = D1[k]
d2 = D2[k]
dhat = (d1/d2) / (1/(theta*dt))

appctx['d1'] = d1
appctx['d2'] = d2

d1r.assign(d1.real)
d1i.assign(d1.imag)

d2r.assign(d2.real)
d2i.assign(d2.imag)

solver = fd.LinearVariationalSolver(
problem, appctx=appctx,
options_prefix="block",
solver_parameters=sparams)
solver.invalidate_jacobian()

if args.v or args.vv:
PETSc.Sys.Print(f"Eigenvalues {str(i).rjust(3)}:\n d1 = {np.round(d1,4)}, d2 = {np.round(d2,4)}, dhat = {np.round(dhat,4)}")
if verbosity > 0:
PETSc.Sys.Print(f"Eigenvalues {str(k).rjust(3)}:\n d1 = {np.round(d1,4)}, d2 = {np.round(d2,4)}, dhat = {np.round(dhat,4)}")

np.random.seed(args.seed)
for j in range(args.nrhs):
Expand All @@ -382,7 +440,7 @@ def aux_form_function(u, h, v, q, t):
for i in range(neigs):
PETSc.Sys.Print(f"Eigenvalue {str(i).rjust(2)} iterations: mean = {fstr(meanit[i], 1)} | max = {str(maxit[i]).rjust(2)} | min = {str(minit[i]).rjust(2)} | std = {fstr(stdit[i], 2)}")

PETSc.Sys.Print(f"max iterations: {round(max(meanit), 3)}, min iterations: {round(min(meanit), 3)}")
PETSc.Sys.Print(f"min iterations: {round(min(meanit), 3)}, max iterations: {round(max(meanit), 3)}")
PETSc.Sys.Print([round(it, 3) for it in np.mean(nits, axis=1)])

if fd.COMM_WORLD.rank == 0:
Expand Down

0 comments on commit c4a00e4

Please sign in to comment.