From 70e7ab952d1daffc70b0a314c75cf396dcbcd584 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Wed, 10 Jan 2024 11:43:13 +0000 Subject: [PATCH 01/18] nonlinear gauss siedel for galewsky --- case_studies/gauss_siedel/galewsky.py | 230 ++++++++++++++++++++++++++ 1 file changed, 230 insertions(+) create mode 100644 case_studies/gauss_siedel/galewsky.py diff --git a/case_studies/gauss_siedel/galewsky.py b/case_studies/gauss_siedel/galewsky.py new file mode 100644 index 00000000..12e64d81 --- /dev/null +++ b/case_studies/gauss_siedel/galewsky.py @@ -0,0 +1,230 @@ +from firedrake.petsc import PETSc + +import asQ +import firedrake as fd # noqa: F401 +from utils import units +from utils.planets import earth +import utils.shallow_water as swe +from utils.shallow_water import galewsky + +from functools import partial +from math import sqrt + +PETSc.Sys.popErrorHandler() + +# get command arguments +import argparse +parser = argparse.ArgumentParser( + description='Galewsky testcase for ParaDiag solver using fully implicit SWE solver.', + formatter_class=argparse.ArgumentDefaultsHelpFormatter +) + +parser.add_argument('--ref_level', type=int, default=2, help='Refinement level of icosahedral grid.') +parser.add_argument('--nchunks', type=int, default=4, help='Number of chunks to solve simultaneously.') +parser.add_argument('--nsweeps', type=int, default=4, help='Number of nonlinear sweeps.') +parser.add_argument('--nslices', type=int, default=2, help='Number of time-slices per time-window.') +parser.add_argument('--slice_length', type=int, default=2, help='Number of timesteps per time-slice.') +parser.add_argument('--alpha', type=float, default=0.0001, help='Circulant coefficient.') +parser.add_argument('--dt', type=float, default=0.5, help='Timestep in hours.') +parser.add_argument('--filename', type=str, default='galewsky', help='Name of output vtk files') +parser.add_argument('--metrics_dir', type=str, default='metrics', help='Directory to save paradiag metrics to.') +parser.add_argument('--print_res', action='store_true', help='Print the residuals of each timestep at each iteration.') +parser.add_argument('--show_args', action='store_true', help='Output all the arguments.') + +args = parser.parse_known_args() +args = args[0] + +if args.show_args: + PETSc.Sys.Print(args) + +PETSc.Sys.Print('') +PETSc.Sys.Print('### === --- Setting up --- === ###') +PETSc.Sys.Print('') + +# time steps + +time_partition = tuple((args.slice_length for _ in range(args.nslices))) +chunk_length = sum(time_partition) + +dt = args.dt*units.hour + + +# alternative operator to precondition blocks + +def aux_form_function(u, h, v, q, t): + mesh = v.ufl_domain() + coords = fd.SpatialCoordinate(mesh) + + gravity = earth.Gravity + coriolis = swe.earth_coriolis_expression(*coords) + + H = galewsky.H0 + + return swe.linear.form_function(mesh, gravity, H, coriolis, + u, h, v, q, t) + + +block_appctx = { + 'aux_form_function': aux_form_function +} + +# parameters for the implicit diagonal solve in step-(b) +factorisation_params = { + 'ksp_type': 'preonly', + # 'pc_factor_mat_ordering_type': 'rcm', + 'pc_factor_reuse_ordering': None, + 'pc_factor_reuse_fill': None, +} + +lu_params = {'pc_type': 'lu', 'pc_factor_mat_solver_type': 'mumps'} +lu_params.update(factorisation_params) + +aux_pc = { + 'snes_lag_preconditioner': -2, + 'snes_lag_preconditioner_persists': None, + 'pc_type': 'python', + 'pc_python_type': 'asQ.AuxiliaryBlockPC', + 'aux': lu_params, +} + + +sparameters = { + 'mat_type': 'matfree', + 'ksp_type': 'fgmres', + 'ksp': { + 'atol': 1e-5, + 'rtol': 1e-5, + 'max_it': 40, + 'converged_maxits': None, + }, +} +sparameters.update(aux_pc) + +atol = 1e4 +patol = sqrt(chunk_length)*atol +sparameters_diag = { + 'snes': { + 'monitor': None, + # 'converged_reason': None, + 'atol': 1e-0, + 'rtol': 1e-10, + 'stol': 1e-12, + # 'ksp_ew': None, + # 'ksp_ew_version': 1, + # 'ksp_ew_threshold': 1e-2, + 'max_it': 1, + 'convergence_test': 'skip', + }, + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'ksp': { + # 'monitor': None, + # 'converged_reason': None, + 'rtol': 1e-5, + 'atol': 1e-0, + }, + 'pc_type': 'python', + 'pc_python_type': 'asQ.DiagFFTPC', + 'diagfft_alpha': args.alpha, + 'diagfft_state': 'window', + 'aaos_jacobian_state': 'current' +} + +for i in range(chunk_length): + sparameters_diag['diagfft_block_'+str(i)+'_'] = sparameters + +create_mesh = partial( + swe.create_mg_globe_mesh, + ref_level=args.ref_level, + coords_degree=1) # remove coords degree once UFL issue with gradient of cell normals fixed + +# check convergence of each timestep + + +def post_function_callback(aaosolver, X, F): + if args.print_res: + residuals = asQ.SharedArray(time_partition, + comm=aaosolver.ensemble.ensemble_comm) + # all-at-once residual + res = aaosolver.aaoform.F + for i in range(res.nlocal_timesteps): + with res[i].dat.vec_ro as vec: + residuals.dlocal[i] = vec.norm() + residuals.synchronise() + PETSc.Sys.Print('') + PETSc.Sys.Print([f"{r:.4e}" for r in residuals.data()]) + PETSc.Sys.Print('') + + +PETSc.Sys.Print('### === --- Calculating parallel solution --- === ###') + +appctx = {'block_appctx': block_appctx} + +miniapp = swe.ShallowWaterMiniApp(gravity=earth.Gravity, + topography_expression=galewsky.topography_expression, + velocity_expression=galewsky.velocity_expression, + depth_expression=galewsky.depth_expression, + reference_depth=galewsky.H0, + reference_state=True, + create_mesh=create_mesh, + dt=dt, theta=0.5, + time_partition=time_partition, + appctx=appctx, + paradiag_sparameters=sparameters_diag, + file_name='output/'+args.filename, + post_function_callback=post_function_callback, + record_diagnostics={'cfl': True, 'file': False}) + +paradiag = miniapp.paradiag +aaosolver = paradiag.solver +aaofunc = aaosolver.aaofunc +aaoform = aaosolver.aaoform + +chunks = tuple(aaofunc.copy() for _ in range(args.nchunks)) + +chunk_ids = [i for i in range(args.nchunks)] + +nconverged = 0 +for j in range(args.nsweeps): + PETSc.Sys.Print('') + PETSc.Sys.Print(f'### === --- Calculating nonlinear sweep {j} --- === ###') + PETSc.Sys.Print('') + + for i in range(args.nchunks): + PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') + + # 1) load chunk i solution + aaofunc.assign(chunks[i]) + + # 2) set ic from previous chunk + if i > 0: + chunks[i-1].bcast_field(-1, aaofunc.initial_condition) + + # 3) one iteration of chunk i + aaosolver.solve() + + # 4) save chunk i solution + chunks[i].assign(aaofunc) + + PETSc.Sys.Print("") + + # check if first chunk is converged + aaoform.assemble(chunks[0]) + with aaoform.F.global_vec_ro() as rvec: + res = rvec.norm() + if res < patol: + PETSc.Sys.Print(">>> First chunk has converged. Rotating chunks <<<") + nconverged += 1 + + # shuffle chunks down + for i in range(args.nchunks-1): + chunks[i].assign(chunks[i+1]) + + # reset last chunk to start from end of series + chunks[-1].bcast_field(-1, chunks[-1].initial_condition) + chunks[-1].assign(chunks[-1].initial_condition) + +PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") +PETSc.Sys.Print(f"Number of sweeps: {args.nsweeps}") +PETSc.Sys.Print(f"Number of chunks converged: {nconverged}") +PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/args.nsweeps}") From 12020473b9f41651113358b4e8a949638b896ec0 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 25 Jan 2024 15:48:49 +0000 Subject: [PATCH 02/18] correct tolerances for nonlinear gs --- case_studies/gauss_siedel/galewsky.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/case_studies/gauss_siedel/galewsky.py b/case_studies/gauss_siedel/galewsky.py index 12e64d81..5861ce00 100644 --- a/case_studies/gauss_siedel/galewsky.py +++ b/case_studies/gauss_siedel/galewsky.py @@ -90,7 +90,7 @@ def aux_form_function(u, h, v, q, t): sparameters = { 'mat_type': 'matfree', - 'ksp_type': 'fgmres', + 'ksp_type': 'gmres', 'ksp': { 'atol': 1e-5, 'rtol': 1e-5, @@ -106,13 +106,13 @@ def aux_form_function(u, h, v, q, t): 'snes': { 'monitor': None, # 'converged_reason': None, - 'atol': 1e-0, + 'atol': patol, 'rtol': 1e-10, 'stol': 1e-12, # 'ksp_ew': None, # 'ksp_ew_version': 1, # 'ksp_ew_threshold': 1e-2, - 'max_it': 1, + 'max_it': 2, 'convergence_test': 'skip', }, 'mat_type': 'matfree', @@ -121,7 +121,7 @@ def aux_form_function(u, h, v, q, t): # 'monitor': None, # 'converged_reason': None, 'rtol': 1e-5, - 'atol': 1e-0, + 'atol': patol, }, 'pc_type': 'python', 'pc_python_type': 'asQ.DiagFFTPC', From 2624c0204815b1f5bfdd1fc702658adc1143dbbb Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 26 Jan 2024 12:03:34 +0000 Subject: [PATCH 03/18] converge multiple chunks for nonlinear gauss siedel and output block iteration counts --- case_studies/gauss_siedel/galewsky.py | 50 ++++++++++++++++++--------- 1 file changed, 33 insertions(+), 17 deletions(-) diff --git a/case_studies/gauss_siedel/galewsky.py b/case_studies/gauss_siedel/galewsky.py index 5861ce00..78800d0d 100644 --- a/case_studies/gauss_siedel/galewsky.py +++ b/case_studies/gauss_siedel/galewsky.py @@ -22,9 +22,11 @@ parser.add_argument('--ref_level', type=int, default=2, help='Refinement level of icosahedral grid.') parser.add_argument('--nchunks', type=int, default=4, help='Number of chunks to solve simultaneously.') parser.add_argument('--nsweeps', type=int, default=4, help='Number of nonlinear sweeps.') +parser.add_argument('--nsmooth', type=int, default=1, help='Number of nonlinear iterations per chunk at each sweep.') +parser.add_argument('--nchecks', type=int, default=1, help='Maximum number of chunks allowed to converge after each sweep.') parser.add_argument('--nslices', type=int, default=2, help='Number of time-slices per time-window.') parser.add_argument('--slice_length', type=int, default=2, help='Number of timesteps per time-slice.') -parser.add_argument('--alpha', type=float, default=0.0001, help='Circulant coefficient.') +parser.add_argument('--alpha', type=float, default=1e-4, help='Circulant coefficient.') parser.add_argument('--dt', type=float, default=0.5, help='Timestep in hours.') parser.add_argument('--filename', type=str, default='galewsky', help='Name of output vtk files') parser.add_argument('--metrics_dir', type=str, default='metrics', help='Directory to save paradiag metrics to.') @@ -87,20 +89,19 @@ def aux_form_function(u, h, v, q, t): 'aux': lu_params, } - sparameters = { 'mat_type': 'matfree', 'ksp_type': 'gmres', 'ksp': { 'atol': 1e-5, 'rtol': 1e-5, - 'max_it': 40, + 'max_it': 25, 'converged_maxits': None, }, } sparameters.update(aux_pc) -atol = 1e4 +atol = 1e2 patol = sqrt(chunk_length)*atol sparameters_diag = { 'snes': { @@ -112,7 +113,7 @@ def aux_form_function(u, h, v, q, t): # 'ksp_ew': None, # 'ksp_ew_version': 1, # 'ksp_ew_threshold': 1e-2, - 'max_it': 2, + 'max_it': args.nsmooth, 'convergence_test': 'skip', }, 'mat_type': 'matfree', @@ -120,6 +121,8 @@ def aux_form_function(u, h, v, q, t): 'ksp': { # 'monitor': None, # 'converged_reason': None, + # 'max_it': 1, + # 'convergence_test': 'skip', 'rtol': 1e-5, 'atol': patol, }, @@ -209,22 +212,35 @@ def post_function_callback(aaosolver, X, F): PETSc.Sys.Print("") # check if first chunk is converged - aaoform.assemble(chunks[0]) - with aaoform.F.global_vec_ro() as rvec: - res = rvec.norm() - if res < patol: - PETSc.Sys.Print(">>> First chunk has converged. Rotating chunks <<<") - nconverged += 1 + for i in range(args.nchecks): + aaoform.assemble(chunks[0]) + with aaoform.F.global_vec_ro() as rvec: + res = rvec.norm() + if res < patol: + PETSc.Sys.Print(f">>> Chunk {i} has converged. Rotating chunks <<<") + nconverged += 1 + + # shuffle chunks down + for i in range(args.nchunks-1): + chunks[i].assign(chunks[i+1]) - # shuffle chunks down - for i in range(args.nchunks-1): - chunks[i].assign(chunks[i+1]) + # reset last chunk to start from end of series + chunks[-1].bcast_field(-1, chunks[-1].initial_condition) + chunks[-1].assign(chunks[-1].initial_condition) - # reset last chunk to start from end of series - chunks[-1].bcast_field(-1, chunks[-1].initial_condition) - chunks[-1].assign(chunks[-1].initial_condition) +pc_block_its = aaosolver.jacobian.pc.block_iterations +pc_block_its.synchronise() +pc_block_its = pc_block_its.data(deepcopy=True) +pc_block_its = pc_block_its/args.nchunks +niterations = args.nsweeps*args.nsmooth + +PETSc.Sys.Print('') PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") PETSc.Sys.Print(f"Number of sweeps: {args.nsweeps}") PETSc.Sys.Print(f"Number of chunks converged: {nconverged}") PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/args.nsweeps}") +PETSc.Sys.Print(f"Number of sweeps per converged chunk: {args.nsweeps/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Number of iterations per converged chunk: {niterations/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f'Block iterations: {pc_block_its}') +PETSc.Sys.Print(f'Block iterations per block solve: {pc_block_its/niterations}') From 26c43e6c67f11d15de42290886d4db1b6a6e7432 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 29 Jan 2024 13:30:26 +0000 Subject: [PATCH 04/18] switch for nonlinear jacobi or gauss-siedel switch --- case_studies/gauss_siedel/galewsky.py | 83 +++++++++++++++++++-------- 1 file changed, 58 insertions(+), 25 deletions(-) diff --git a/case_studies/gauss_siedel/galewsky.py b/case_studies/gauss_siedel/galewsky.py index 78800d0d..57c1fb76 100644 --- a/case_studies/gauss_siedel/galewsky.py +++ b/case_studies/gauss_siedel/galewsky.py @@ -20,10 +20,15 @@ ) parser.add_argument('--ref_level', type=int, default=2, help='Refinement level of icosahedral grid.') +parser.add_argument('--base_level', type=int, default=2, help='Refinement level of coarse grid.') +parser.add_argument('--nwindows', type=int, default=1, help='Total number of time-windows.') +parser.add_argument('--nlmethod', type=str, default='gs', choices=['gs', 'jac'], help='Nonlinear method. "gs" for Gauss-Siedel or "jac" for Jacobi.') parser.add_argument('--nchunks', type=int, default=4, help='Number of chunks to solve simultaneously.') parser.add_argument('--nsweeps', type=int, default=4, help='Number of nonlinear sweeps.') parser.add_argument('--nsmooth', type=int, default=1, help='Number of nonlinear iterations per chunk at each sweep.') parser.add_argument('--nchecks', type=int, default=1, help='Maximum number of chunks allowed to converge after each sweep.') +parser.add_argument('--ninitialise', type=int, default=0, help='Number of sweeps before checking convergence.') +parser.add_argument('--atol', type=float, default=1e0, help='Average atol of each timestep.') parser.add_argument('--nslices', type=int, default=2, help='Number of time-slices per time-window.') parser.add_argument('--slice_length', type=int, default=2, help='Number of timesteps per time-slice.') parser.add_argument('--alpha', type=float, default=1e-4, help='Circulant coefficient.') @@ -91,22 +96,23 @@ def aux_form_function(u, h, v, q, t): sparameters = { 'mat_type': 'matfree', - 'ksp_type': 'gmres', + 'ksp_type': 'fgmres', 'ksp': { 'atol': 1e-5, 'rtol': 1e-5, - 'max_it': 25, + 'max_it': 30, 'converged_maxits': None, }, } sparameters.update(aux_pc) -atol = 1e2 +atol = args.atol patol = sqrt(chunk_length)*atol sparameters_diag = { 'snes': { + 'linesearch_type': 'basic', 'monitor': None, - # 'converged_reason': None, + 'converged_reason': None, 'atol': patol, 'rtol': 1e-10, 'stol': 1e-12, @@ -121,9 +127,9 @@ def aux_form_function(u, h, v, q, t): 'ksp': { # 'monitor': None, # 'converged_reason': None, - # 'max_it': 1, - # 'convergence_test': 'skip', - 'rtol': 1e-5, + # 'max_it': 2, + # 'converged_maxits': None, + 'rtol': 1e-2, 'atol': patol, }, 'pc_type': 'python', @@ -139,6 +145,7 @@ def aux_form_function(u, h, v, q, t): create_mesh = partial( swe.create_mg_globe_mesh, ref_level=args.ref_level, + base_level=args.base_level, coords_degree=1) # remove coords degree once UFL issue with gradient of cell normals fixed # check convergence of each timestep @@ -185,41 +192,51 @@ def post_function_callback(aaosolver, X, F): chunks = tuple(aaofunc.copy() for _ in range(args.nchunks)) -chunk_ids = [i for i in range(args.nchunks)] - nconverged = 0 +PETSc.Sys.Print('') for j in range(args.nsweeps): - PETSc.Sys.Print('') PETSc.Sys.Print(f'### === --- Calculating nonlinear sweep {j} --- === ###') PETSc.Sys.Print('') + # 1) set ics from previous iteration of previous chunk + if (args.nlmethod == 'jac'): + for i in range(1, args.nchunks): + chunks[i-1].bcast_field(-1, chunks[i].initial_condition) + for i in range(args.nchunks): PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') - # 1) load chunk i solution - aaofunc.assign(chunks[i]) + # 1) set ic from current iteration of previous chunk + if (args.nlmethod == 'gs'): + if i > 0: + chunks[i-1].bcast_field(-1, chunks[i].initial_condition) - # 2) set ic from previous chunk - if i > 0: - chunks[i-1].bcast_field(-1, aaofunc.initial_condition) + # 2) load chunk i solution + aaofunc.assign(chunks[i]) # 3) one iteration of chunk i aaosolver.solve() + PETSc.Sys.Print("") # 4) save chunk i solution chunks[i].assign(aaofunc) - PETSc.Sys.Print("") - - # check if first chunk is converged + # check if first chunks have converged for i in range(args.nchecks): + if j < args.ninitialise: + break + aaoform.assemble(chunks[0]) with aaoform.F.global_vec_ro() as rvec: res = rvec.norm() + if res < patol: - PETSc.Sys.Print(f">>> Chunk {i} has converged. Rotating chunks <<<") + PETSc.Sys.Print(f">>> Chunk {str(i).rjust(2)} converged with function norm = {res:.6e} <<<") nconverged += 1 + # make sure we calculate the most up to date residual for the next chunk + chunks[0].bcast_field(-1, chunks[1].initial_condition) + # shuffle chunks down for i in range(args.nchunks-1): chunks[i].assign(chunks[i+1]) @@ -228,19 +245,35 @@ def post_function_callback(aaosolver, X, F): chunks[-1].bcast_field(-1, chunks[-1].initial_condition) chunks[-1].assign(chunks[-1].initial_condition) + else: # stop checking at first unconverged chunk + break + + PETSc.Sys.Print('') + PETSc.Sys.Print(f">>> Converged chunks: {int(nconverged)}.") + converged_time = nconverged*chunk_length*args.dt + PETSc.Sys.Print(f">>> Converged time: {converged_time} hours.") + PETSc.Sys.Print('') + + if nconverged >= args.nwindows: + PETSc.Sys.Print(f"Finished iterating to {args.nwindows}.") + break + +nsweeps = j + pc_block_its = aaosolver.jacobian.pc.block_iterations pc_block_its.synchronise() pc_block_its = pc_block_its.data(deepcopy=True) pc_block_its = pc_block_its/args.nchunks -niterations = args.nsweeps*args.nsmooth +niterations = nsweeps*args.nsmooth -PETSc.Sys.Print('') PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") -PETSc.Sys.Print(f"Number of sweeps: {args.nsweeps}") -PETSc.Sys.Print(f"Number of chunks converged: {nconverged}") -PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/args.nsweeps}") -PETSc.Sys.Print(f"Number of sweeps per converged chunk: {args.nsweeps/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Maximum number of sweeps: {args.nsweeps}") +PETSc.Sys.Print(f"Actual number of sweeps: {nsweeps}") +PETSc.Sys.Print(f"Number of chunks converged: {int(nconverged)}") +PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/nsweeps}") +PETSc.Sys.Print(f"Number of sweeps per converged chunk: {nsweeps/nconverged if nconverged else 'n/a'}") PETSc.Sys.Print(f"Number of iterations per converged chunk: {niterations/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Number of timesteps per iteration: {nconverged*chunk_length/niterations}") PETSc.Sys.Print(f'Block iterations: {pc_block_its}') PETSc.Sys.Print(f'Block iterations per block solve: {pc_block_its/niterations}') From d3f710c59e9af7e8d224d9740e5da1b67a09d601 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 30 Jan 2024 08:08:51 +0000 Subject: [PATCH 05/18] split serial/parallel gauss-siedel galewsky --- case_studies/gauss_siedel/galewsky.py | 54 +--- case_studies/gauss_siedel/galewsky_serial.py | 279 +++++++++++++++++++ 2 files changed, 293 insertions(+), 40 deletions(-) create mode 100644 case_studies/gauss_siedel/galewsky_serial.py diff --git a/case_studies/gauss_siedel/galewsky.py b/case_studies/gauss_siedel/galewsky.py index 57c1fb76..3284649c 100644 --- a/case_studies/gauss_siedel/galewsky.py +++ b/case_studies/gauss_siedel/galewsky.py @@ -198,55 +198,29 @@ def post_function_callback(aaosolver, X, F): PETSc.Sys.Print(f'### === --- Calculating nonlinear sweep {j} --- === ###') PETSc.Sys.Print('') - # 1) set ics from previous iteration of previous chunk - if (args.nlmethod == 'jac'): - for i in range(1, args.nchunks): - chunks[i-1].bcast_field(-1, chunks[i].initial_condition) - + # 1) one smoothing application on each chunk for i in range(args.nchunks): PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') - # 1) set ic from current iteration of previous chunk - if (args.nlmethod == 'gs'): - if i > 0: - chunks[i-1].bcast_field(-1, chunks[i].initial_condition) - - # 2) load chunk i solution - aaofunc.assign(chunks[i]) + global_comm.Barrier() + if chunk_id == i: + aaosolver.solve() + global_comm.Barrier() - # 3) one iteration of chunk i - aaosolver.solve() PETSc.Sys.Print("") - # 4) save chunk i solution - chunks[i].assign(aaofunc) - - # check if first chunks have converged - for i in range(args.nchecks): - if j < args.ninitialise: - break - - aaoform.assemble(chunks[0]) - with aaoform.F.global_vec_ro() as rvec: - res = rvec.norm() + # 2) update ics of each chunk from previous chunk + update_ics(global_aaofunc, chunk_aaofunc) - if res < patol: - PETSc.Sys.Print(f">>> Chunk {str(i).rjust(2)} converged with function norm = {res:.6e} <<<") - nconverged += 1 - - # make sure we calculate the most up to date residual for the next chunk - chunks[0].bcast_field(-1, chunks[1].initial_condition) - - # shuffle chunks down - for i in range(args.nchunks-1): - chunks[i].assign(chunks[i+1]) + # 3) shuffle earlier chunks if converged + if j < args.ninitialise: + break - # reset last chunk to start from end of series - chunks[-1].bcast_field(-1, chunks[-1].initial_condition) - chunks[-1].assign(chunks[-1].initial_condition) + update_chunk_residuals(chunk_aaofunc, chunk_residuals) + nshuffle = count_converged(chunk_residuals, args.nchecks) + shuffle_chunks(global_aaofunc, chunk_aaofunc, nshuffle) - else: # stop checking at first unconverged chunk - break + nconverged += nshuffle PETSc.Sys.Print('') PETSc.Sys.Print(f">>> Converged chunks: {int(nconverged)}.") diff --git a/case_studies/gauss_siedel/galewsky_serial.py b/case_studies/gauss_siedel/galewsky_serial.py new file mode 100644 index 00000000..57c1fb76 --- /dev/null +++ b/case_studies/gauss_siedel/galewsky_serial.py @@ -0,0 +1,279 @@ +from firedrake.petsc import PETSc + +import asQ +import firedrake as fd # noqa: F401 +from utils import units +from utils.planets import earth +import utils.shallow_water as swe +from utils.shallow_water import galewsky + +from functools import partial +from math import sqrt + +PETSc.Sys.popErrorHandler() + +# get command arguments +import argparse +parser = argparse.ArgumentParser( + description='Galewsky testcase for ParaDiag solver using fully implicit SWE solver.', + formatter_class=argparse.ArgumentDefaultsHelpFormatter +) + +parser.add_argument('--ref_level', type=int, default=2, help='Refinement level of icosahedral grid.') +parser.add_argument('--base_level', type=int, default=2, help='Refinement level of coarse grid.') +parser.add_argument('--nwindows', type=int, default=1, help='Total number of time-windows.') +parser.add_argument('--nlmethod', type=str, default='gs', choices=['gs', 'jac'], help='Nonlinear method. "gs" for Gauss-Siedel or "jac" for Jacobi.') +parser.add_argument('--nchunks', type=int, default=4, help='Number of chunks to solve simultaneously.') +parser.add_argument('--nsweeps', type=int, default=4, help='Number of nonlinear sweeps.') +parser.add_argument('--nsmooth', type=int, default=1, help='Number of nonlinear iterations per chunk at each sweep.') +parser.add_argument('--nchecks', type=int, default=1, help='Maximum number of chunks allowed to converge after each sweep.') +parser.add_argument('--ninitialise', type=int, default=0, help='Number of sweeps before checking convergence.') +parser.add_argument('--atol', type=float, default=1e0, help='Average atol of each timestep.') +parser.add_argument('--nslices', type=int, default=2, help='Number of time-slices per time-window.') +parser.add_argument('--slice_length', type=int, default=2, help='Number of timesteps per time-slice.') +parser.add_argument('--alpha', type=float, default=1e-4, help='Circulant coefficient.') +parser.add_argument('--dt', type=float, default=0.5, help='Timestep in hours.') +parser.add_argument('--filename', type=str, default='galewsky', help='Name of output vtk files') +parser.add_argument('--metrics_dir', type=str, default='metrics', help='Directory to save paradiag metrics to.') +parser.add_argument('--print_res', action='store_true', help='Print the residuals of each timestep at each iteration.') +parser.add_argument('--show_args', action='store_true', help='Output all the arguments.') + +args = parser.parse_known_args() +args = args[0] + +if args.show_args: + PETSc.Sys.Print(args) + +PETSc.Sys.Print('') +PETSc.Sys.Print('### === --- Setting up --- === ###') +PETSc.Sys.Print('') + +# time steps + +time_partition = tuple((args.slice_length for _ in range(args.nslices))) +chunk_length = sum(time_partition) + +dt = args.dt*units.hour + + +# alternative operator to precondition blocks + +def aux_form_function(u, h, v, q, t): + mesh = v.ufl_domain() + coords = fd.SpatialCoordinate(mesh) + + gravity = earth.Gravity + coriolis = swe.earth_coriolis_expression(*coords) + + H = galewsky.H0 + + return swe.linear.form_function(mesh, gravity, H, coriolis, + u, h, v, q, t) + + +block_appctx = { + 'aux_form_function': aux_form_function +} + +# parameters for the implicit diagonal solve in step-(b) +factorisation_params = { + 'ksp_type': 'preonly', + # 'pc_factor_mat_ordering_type': 'rcm', + 'pc_factor_reuse_ordering': None, + 'pc_factor_reuse_fill': None, +} + +lu_params = {'pc_type': 'lu', 'pc_factor_mat_solver_type': 'mumps'} +lu_params.update(factorisation_params) + +aux_pc = { + 'snes_lag_preconditioner': -2, + 'snes_lag_preconditioner_persists': None, + 'pc_type': 'python', + 'pc_python_type': 'asQ.AuxiliaryBlockPC', + 'aux': lu_params, +} + +sparameters = { + 'mat_type': 'matfree', + 'ksp_type': 'fgmres', + 'ksp': { + 'atol': 1e-5, + 'rtol': 1e-5, + 'max_it': 30, + 'converged_maxits': None, + }, +} +sparameters.update(aux_pc) + +atol = args.atol +patol = sqrt(chunk_length)*atol +sparameters_diag = { + 'snes': { + 'linesearch_type': 'basic', + 'monitor': None, + 'converged_reason': None, + 'atol': patol, + 'rtol': 1e-10, + 'stol': 1e-12, + # 'ksp_ew': None, + # 'ksp_ew_version': 1, + # 'ksp_ew_threshold': 1e-2, + 'max_it': args.nsmooth, + 'convergence_test': 'skip', + }, + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'ksp': { + # 'monitor': None, + # 'converged_reason': None, + # 'max_it': 2, + # 'converged_maxits': None, + 'rtol': 1e-2, + 'atol': patol, + }, + 'pc_type': 'python', + 'pc_python_type': 'asQ.DiagFFTPC', + 'diagfft_alpha': args.alpha, + 'diagfft_state': 'window', + 'aaos_jacobian_state': 'current' +} + +for i in range(chunk_length): + sparameters_diag['diagfft_block_'+str(i)+'_'] = sparameters + +create_mesh = partial( + swe.create_mg_globe_mesh, + ref_level=args.ref_level, + base_level=args.base_level, + coords_degree=1) # remove coords degree once UFL issue with gradient of cell normals fixed + +# check convergence of each timestep + + +def post_function_callback(aaosolver, X, F): + if args.print_res: + residuals = asQ.SharedArray(time_partition, + comm=aaosolver.ensemble.ensemble_comm) + # all-at-once residual + res = aaosolver.aaoform.F + for i in range(res.nlocal_timesteps): + with res[i].dat.vec_ro as vec: + residuals.dlocal[i] = vec.norm() + residuals.synchronise() + PETSc.Sys.Print('') + PETSc.Sys.Print([f"{r:.4e}" for r in residuals.data()]) + PETSc.Sys.Print('') + + +PETSc.Sys.Print('### === --- Calculating parallel solution --- === ###') + +appctx = {'block_appctx': block_appctx} + +miniapp = swe.ShallowWaterMiniApp(gravity=earth.Gravity, + topography_expression=galewsky.topography_expression, + velocity_expression=galewsky.velocity_expression, + depth_expression=galewsky.depth_expression, + reference_depth=galewsky.H0, + reference_state=True, + create_mesh=create_mesh, + dt=dt, theta=0.5, + time_partition=time_partition, + appctx=appctx, + paradiag_sparameters=sparameters_diag, + file_name='output/'+args.filename, + post_function_callback=post_function_callback, + record_diagnostics={'cfl': True, 'file': False}) + +paradiag = miniapp.paradiag +aaosolver = paradiag.solver +aaofunc = aaosolver.aaofunc +aaoform = aaosolver.aaoform + +chunks = tuple(aaofunc.copy() for _ in range(args.nchunks)) + +nconverged = 0 +PETSc.Sys.Print('') +for j in range(args.nsweeps): + PETSc.Sys.Print(f'### === --- Calculating nonlinear sweep {j} --- === ###') + PETSc.Sys.Print('') + + # 1) set ics from previous iteration of previous chunk + if (args.nlmethod == 'jac'): + for i in range(1, args.nchunks): + chunks[i-1].bcast_field(-1, chunks[i].initial_condition) + + for i in range(args.nchunks): + PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') + + # 1) set ic from current iteration of previous chunk + if (args.nlmethod == 'gs'): + if i > 0: + chunks[i-1].bcast_field(-1, chunks[i].initial_condition) + + # 2) load chunk i solution + aaofunc.assign(chunks[i]) + + # 3) one iteration of chunk i + aaosolver.solve() + PETSc.Sys.Print("") + + # 4) save chunk i solution + chunks[i].assign(aaofunc) + + # check if first chunks have converged + for i in range(args.nchecks): + if j < args.ninitialise: + break + + aaoform.assemble(chunks[0]) + with aaoform.F.global_vec_ro() as rvec: + res = rvec.norm() + + if res < patol: + PETSc.Sys.Print(f">>> Chunk {str(i).rjust(2)} converged with function norm = {res:.6e} <<<") + nconverged += 1 + + # make sure we calculate the most up to date residual for the next chunk + chunks[0].bcast_field(-1, chunks[1].initial_condition) + + # shuffle chunks down + for i in range(args.nchunks-1): + chunks[i].assign(chunks[i+1]) + + # reset last chunk to start from end of series + chunks[-1].bcast_field(-1, chunks[-1].initial_condition) + chunks[-1].assign(chunks[-1].initial_condition) + + else: # stop checking at first unconverged chunk + break + + PETSc.Sys.Print('') + PETSc.Sys.Print(f">>> Converged chunks: {int(nconverged)}.") + converged_time = nconverged*chunk_length*args.dt + PETSc.Sys.Print(f">>> Converged time: {converged_time} hours.") + PETSc.Sys.Print('') + + if nconverged >= args.nwindows: + PETSc.Sys.Print(f"Finished iterating to {args.nwindows}.") + break + +nsweeps = j + +pc_block_its = aaosolver.jacobian.pc.block_iterations +pc_block_its.synchronise() +pc_block_its = pc_block_its.data(deepcopy=True) +pc_block_its = pc_block_its/args.nchunks + +niterations = nsweeps*args.nsmooth + +PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") +PETSc.Sys.Print(f"Maximum number of sweeps: {args.nsweeps}") +PETSc.Sys.Print(f"Actual number of sweeps: {nsweeps}") +PETSc.Sys.Print(f"Number of chunks converged: {int(nconverged)}") +PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/nsweeps}") +PETSc.Sys.Print(f"Number of sweeps per converged chunk: {nsweeps/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Number of iterations per converged chunk: {niterations/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Number of timesteps per iteration: {nconverged*chunk_length/niterations}") +PETSc.Sys.Print(f'Block iterations: {pc_block_its}') +PETSc.Sys.Print(f'Block iterations per block solve: {pc_block_its/niterations}') From 0180e04348edea9dfbecb954458c8171eacd7b87 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 20 Feb 2024 20:05:41 +0000 Subject: [PATCH 06/18] dg_advection with aaos gauss-siedel --- .../advection/dg_advection_gauss_siedel.py | 257 ++++++++++++++++++ 1 file changed, 257 insertions(+) create mode 100644 examples/advection/dg_advection_gauss_siedel.py diff --git a/examples/advection/dg_advection_gauss_siedel.py b/examples/advection/dg_advection_gauss_siedel.py new file mode 100644 index 00000000..5c14ae21 --- /dev/null +++ b/examples/advection/dg_advection_gauss_siedel.py @@ -0,0 +1,257 @@ +from math import pi, cos, sin, sqrt + +import firedrake as fd +from firedrake.petsc import PETSc +import asQ + +import argparse + +parser = argparse.ArgumentParser( + description='ParaDiag timestepping for scalar advection of a Gaussian bump in a periodic square with DG in space and implicit-theta in time. Based on the Firedrake DG advection example https://www.firedrakeproject.org/demos/DG_advection.py.html', + formatter_class=argparse.ArgumentDefaultsHelpFormatter +) +parser.add_argument('--nx', type=int, default=16, help='Number of cells along each square side.') +parser.add_argument('--cfl', type=float, default=0.8, help='Convective CFL number.') +parser.add_argument('--angle', type=float, default=pi/6, help='Angle of the convective velocity.') +parser.add_argument('--degree', type=int, default=1, help='Degree of the scalar spaces.') +parser.add_argument('--theta', type=float, default=0.5, help='Parameter for the implicit theta timestepping method.') +parser.add_argument('--width', type=float, default=0.2, help='Width of the Gaussian bump.') +parser.add_argument('--nwindows', type=int, default=1, help='Total number of time-windows.') +parser.add_argument('--nchunks', type=int, default=4, help='Number of chunks to solve simultaneously.') +parser.add_argument('--nsweeps', type=int, default=4, help='Number of nonlinear sweeps.') +parser.add_argument('--nsmooth', type=int, default=1, help='Number of nonlinear iterations per chunk at each sweep.') +parser.add_argument('--nchecks', type=int, default=1, help='Maximum number of chunks allowed to converge after each sweep.') +parser.add_argument('--ninitialise', type=int, default=0, help='Number of sweeps before checking convergence.') +parser.add_argument('--insurance_freq', type=int, default=0, help='Frequency of sweeps where no convergence is allowed.') +parser.add_argument('--atol', type=float, default=1e-8, help='Average atol of each timestep.') +parser.add_argument('--nslices', type=int, default=2, help='Number of time-slices per time-window.') +parser.add_argument('--slice_length', type=int, default=2, help='Number of timesteps per time-slice.') +parser.add_argument('--alpha', type=float, default=0.0001, help='Circulant coefficient.') +parser.add_argument('--show_args', action='store_true', help='Output all the arguments.') + +args = parser.parse_known_args() +args = args[0] + +if args.show_args: + PETSc.Sys.Print(args) + +# The time partition describes how many timesteps are included on each time-slice of the ensemble +# Here we use the same number of timesteps on each slice, but they can be different + +time_partition = tuple(args.slice_length for _ in range(args.nslices)) +chunk_length = sum(time_partition) + +# Calculate the timestep from the CFL number +umax = 1. +dx = 1./args.nx +dt = args.cfl*dx/umax + +# The Ensemble with the spatial and time communicators +global_comm = fd.COMM_WORLD +ensemble = asQ.create_ensemble(time_partition, comm=global_comm) + +# # # === --- domain --- === # # # + +# The mesh needs to be created with the spatial communicator +mesh = fd.PeriodicUnitSquareMesh(args.nx, args.nx, quadrilateral=True, comm=ensemble.comm) + +# We use a discontinuous Galerkin space for the advected scalar +# and a continuous Galerkin space for the advecting velocity field +V = fd.FunctionSpace(mesh, "DQ", args.degree) +W = fd.VectorFunctionSpace(mesh, "CG", args.degree+1) + +# # # === --- initial conditions --- === # # # + +x, y = fd.SpatialCoordinate(mesh) + + +def radius(x, y): + return fd.sqrt(pow(x-0.5, 2) + pow(y-0.5, 2)) + + +def gaussian(x, y): + return fd.exp(-0.5*pow(radius(x, y)/args.width, 2)) + + +# The scalar initial conditions are a Gaussian bump centred at (0.5, 0.5) +q0 = fd.Function(V, name="scalar_initial") +q0.interpolate(1 + gaussian(x, y)) + +# The advecting velocity field is constant and directed at an angle to the x-axis +u = fd.Function(W, name='velocity') +u.interpolate(fd.as_vector((umax*cos(args.angle), umax*sin(args.angle)))) + +# We create an all-at-once function representing the timeseries solution of the scalar + +aaofunc = asQ.AllAtOnceFunction(ensemble, time_partition, V) +aaofunc.assign(q0) + + +# # # === --- finite element forms --- === # # # + + +# The time-derivative mass form for the scalar advection equation. +# asQ assumes that the mass form is linear so here +# q is a TrialFunction and phi is a TestFunction +def form_mass(q, phi): + return phi*q*fd.dx + + +# The DG advection form for the scalar advection equation. +# asQ assumes that the function form is nonlinear so here +# q is a Function and phi is a TestFunction +def form_function(q, phi, t): + # upwind switch + n = fd.FacetNormal(mesh) + un = fd.Constant(0.5)*(fd.dot(u, n) + abs(fd.dot(u, n))) + + # integration over element volume + int_cell = q*fd.div(phi*u)*fd.dx + + # integration over internal facets + int_facet = (phi('+')-phi('-'))*(un('+')*q('+')-un('-')*q('-'))*fd.dS + + return int_facet - int_cell + + +# Construct the all-at-once form representing the coupled equations for +# the implicit-theta method at every timestep of the timeseries. + +aaoform = asQ.AllAtOnceForm(aaofunc, dt, args.theta, + form_mass, form_function) + +# # # === --- PETSc solver parameters --- === # # # + + +# The PETSc solver parameters used to solve the +# blocks in step (b) of inverting the ParaDiag matrix. +block_parameters = { + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps' +} + +# The PETSc solver parameters for solving the all-at-once system. +# The python preconditioner 'asQ.DiagFFTPC' applies the ParaDiag matrix. +# +# The equation is linear so we can use 'snes_type': 'ksponly' and +# use your favourite Krylov method (if a Krylov method is used on +# the blocks then the outer Krylov method must be either flexible +# or Richardson). + +patol = sqrt(chunk_length)*args.atol +solver_parameters = { + 'snes_type': 'ksponly', + 'snes_convergence_test': 'skip', + 'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'ksp': { + 'monitor': None, + 'converged_rate': None, + 'rtol': 1e-100, + 'atol': patol, + 'stol': 1e-12, + }, + 'pc_type': 'python', + 'pc_python_type': 'asQ.DiagFFTPC', + 'diagfft_alpha': args.alpha, + 'diagfft_state': 'linear', + 'aaos_jacobian_state': 'linear', +} + +# We need to add a block solver parameters dictionary for each block. +# Here they are all the same but they could be different. +for i in range(aaofunc.ntimesteps): + solver_parameters['diagfft_block_'+str(i)+'_'] = block_parameters + +# Create a solver object to set up and solve the (possibly nonlinear) problem +# for the timeseries in the all-at-once function. +aaosolver = asQ.AllAtOnceSolver(aaoform, aaofunc, + solver_parameters=solver_parameters) + +chunks = tuple(aaofunc.copy() for _ in range(args.nchunks)) + +nconverged = 0 +PETSc.Sys.Print('') + +for j in range(args.nsweeps): + PETSc.Sys.Print(f'### === --- Calculating nonlinear sweep {j} --- === ###') + PETSc.Sys.Print('') + + # only iterate chunks that the wavefront has reached + active_chunks = min(j+1, args.nchunks) + + for i in range(active_chunks): + PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') + + # 1) load chunk i solution + aaofunc.assign(chunks[i]) + + # 2) one iteration of chunk i + aaosolver.solve() + PETSc.Sys.Print("") + + # 3) save chunk i solution + chunks[i].assign(aaofunc) + + # 4) pass forward ic to next chunk + for i in range(active_chunks): + if i < args.nchunks-1: + chunks[i].bcast_field(-1, chunks[i+1].initial_condition) + + # skip convergence test during the ramp up and periodically after + ramp_up = j < (args.nchunks + args.ninitialise) + insurance_step = (j % args.insurance_freq) == 0 if args.insurance_freq > 1 else False + + if ramp_up or insurance_step: + continue + + # check if first chunk has converged + aaoform.assemble(chunks[0]) + with aaoform.F.global_vec_ro() as rvec: + res = rvec.norm() + + if res < patol: + PETSc.Sys.Print(f">>> Chunk 0 converged with function norm = {res:.6e} <<<") + nconverged += 1 + + # pass on the most up to date initial condition for the next chunk + chunks[0].bcast_field(-1, chunks[1].initial_condition) + + # shuffle chunks down + for i in range(args.nchunks-1): + chunks[i].assign(chunks[i+1]) + + # reset last chunk to start from end of series + chunks[-1].bcast_field(-1, chunks[-1].initial_condition) + chunks[-1].assign(chunks[-1].initial_condition) + + PETSc.Sys.Print('') + PETSc.Sys.Print(f">>> Converged chunks: {int(nconverged)}.") + converged_time = nconverged*chunk_length*dt + PETSc.Sys.Print(f">>> Converged time: {converged_time} hours.") + PETSc.Sys.Print('') + + if nconverged >= args.nwindows: + PETSc.Sys.Print(f"Finished iterating to {args.nwindows}.") + break + +nsweeps = j + +pc_block_its = aaosolver.jacobian.pc.block_iterations +pc_block_its.synchronise() +pc_block_its = pc_block_its.data(deepcopy=True) +pc_block_its = pc_block_its/args.nchunks + +niterations = nsweeps*args.nsmooth + +PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") +PETSc.Sys.Print(f"Maximum number of sweeps: {args.nsweeps}") +PETSc.Sys.Print(f"Actual number of sweeps: {nsweeps}") +PETSc.Sys.Print(f"Number of chunks converged: {int(nconverged)}") +PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/nsweeps}") +PETSc.Sys.Print(f"Number of sweeps per converged chunk: {nsweeps/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Number of iterations per converged chunk: {niterations/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Number of timesteps per iteration: {nconverged*chunk_length/niterations}") +PETSc.Sys.Print(f'Block iterations: {pc_block_its}') +PETSc.Sys.Print(f'Block iterations per block solve: {pc_block_its/niterations}') From 4404d45af370de1d0dd1e150eeb0a680c1872566 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 22 Feb 2024 10:51:17 +0000 Subject: [PATCH 07/18] split_ensemble functions from JacobiPC branch --- asQ/ensemble.py | 82 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 asQ/ensemble.py diff --git a/asQ/ensemble.py b/asQ/ensemble.py new file mode 100644 index 00000000..740ddd9a --- /dev/null +++ b/asQ/ensemble.py @@ -0,0 +1,82 @@ +from firedrake import COMM_WORLD, Ensemble +from pyop2.mpi import internal_comm, decref + +__all__ = ['create_ensemble', 'split_ensemble', 'EnsembleConnector'] + + +def create_ensemble(time_partition, comm=COMM_WORLD): + ''' + Create an Ensemble for the given slice partition. + Checks that the number of slices and the size of the communicator are compatible. + + :arg time_partition: a list of integers, the number of timesteps on each time-rank + :arg comm: the global communicator for the ensemble + ''' + nslices = 1 if type(time_partition) is int else len(time_partition) + nranks = comm.size + + if nranks % nslices != 0: + raise ValueError("Number of time slices must be exact factor of number of MPI ranks") + + nspatial_domains = nranks//nslices + + return Ensemble(comm, nspatial_domains) + + +def split_ensemble(ensemble, split_size): + """ + Split an Ensemble into multiple smaller Ensembles which share the same + spatial communicators `ensemble.comm`. + + Each smaller Ensemble returned is defined over a contiguous subset of the + members of the large Ensemble. + + :arg ensemble: the large Ensemble to split. + :arg split_size: the number of members in each smaller Ensemble. + """ + if (ensemble.ensemble_comm.size % split_size) != 0: + msg = "Ensemble size must be integer multiple of split_size" + raise ValueError(msg) + + # which split are we part of? + split_rank = ensemble.ensemble_comm.rank // split_size + + # create split_ensemble.global_comm + split_comm = ensemble.global_comm.Split(color=split_rank, + key=ensemble.global_comm.rank) + + return EnsembleConnector(split_comm, ensemble.comm, split_size) + + +class EnsembleConnector(Ensemble): + def __init__(self, global_comm, local_comm, nmembers): + """ + An Ensemble created from provided spatial communicators (ensemble.comm). + + :arg global_comm: global communicator the Ensemble is defined over. + :arg local_comm: communicator to use for the Ensemble.comm member. + :arg nmembers: number of Ensemble members (ensemble.ensemble_comm.size). + """ + if nmembers*local_comm.size != global_comm.size: + msg = "The global ensemble must have the same number of ranks as the sum of the local comms" + raise ValueError(msg) + + self.global_comm = global_comm + self._global_comm = internal_comm(self.global_comm) + + self.comm = local_comm + self._comm = internal_comm(self.comm) + + self.ensemble_comm = self.global_comm.Split(color=self.comm.rank, + key=global_comm.rank) + + self._ensemble_comm = internal_comm(self.ensemble_comm) + + def __del__(self): + if hasattr(self, "ensemble_comm"): + self.ensemble_comm.Free() + del self.ensemble_comm + for comm_name in ["_global_comm", "_comm", "_ensemble_comm"]: + if hasattr(self, comm_name): + comm = getattr(self, comm_name) + decref(comm) From dd275b7d8710a06c31fbada9629d95b3ea147746 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 11 Mar 2024 08:22:07 +0000 Subject: [PATCH 08/18] continuing gauss_seidel parallelisation --- case_studies/gauss_siedel/galewsky.py | 116 +++++++++++++------------- 1 file changed, 56 insertions(+), 60 deletions(-) diff --git a/case_studies/gauss_siedel/galewsky.py b/case_studies/gauss_siedel/galewsky.py index 3284649c..af6790ef 100644 --- a/case_studies/gauss_siedel/galewsky.py +++ b/case_studies/gauss_siedel/galewsky.py @@ -7,7 +7,7 @@ import utils.shallow_water as swe from utils.shallow_water import galewsky -from functools import partial +import numpy as np from math import sqrt PETSc.Sys.popErrorHandler() @@ -50,24 +50,44 @@ # time steps -time_partition = tuple((args.slice_length for _ in range(args.nslices))) -chunk_length = sum(time_partition) +chunk_partition = tuple((args.slice_length for _ in range(args.nslices))) +chunk_length = sum(chunk_partition) +total_timesteps = chunk_length*args.nchunks +total_slices = args.nslices*args.nchunks + +global_comm = fd.COMM_WORLD +global_time_partition = tuple((args.slice_length for _ in range(total_slices))) +global_ensemble = asQ.create_ensemble(global_time_partition, global_comm) +chunk_ensemble = asQ.split_ensemble(global_ensemble, args.nslices) + +# which chunk are we? +chunk_id = global_ensemble.ensemble_comm.rank // args.nslices dt = args.dt*units.hour +mesh = swe.create_mg_globe_mesh( + ref_level=args.ref_level, base_level=args.base_level, + coords_degree=1, comm=chunk_ensemble.comm) +coords = fd.SpatialCoordinate(mesh) # alternative operator to precondition blocks +g = earth.Gravity +f = swe.earth_coriolis_expression(*coords) +H = galewsky.H0 +b = galewsky.topography_expression(*coords) + -def aux_form_function(u, h, v, q, t): - mesh = v.ufl_domain() - coords = fd.SpatialCoordinate(mesh) +def form_mass(u, h, v, q): + return swe.nonlinear.form_mass(mesh, u, h, v, q) - gravity = earth.Gravity - coriolis = swe.earth_coriolis_expression(*coords) - H = galewsky.H0 +def form_function(u, h, v, q, t=None): + return swe.nonlinear.form_function(mesh, g, b, f, + u, h, v, q, t) - return swe.linear.form_function(mesh, gravity, H, coriolis, + +def aux_form_function(u, h, v, q, t=None): + return swe.linear.form_function(mesh, g, H, f, u, h, v, q, t) @@ -142,55 +162,38 @@ def aux_form_function(u, h, v, q, t): for i in range(chunk_length): sparameters_diag['diagfft_block_'+str(i)+'_'] = sparameters -create_mesh = partial( - swe.create_mg_globe_mesh, - ref_level=args.ref_level, - base_level=args.base_level, - coords_degree=1) # remove coords degree once UFL issue with gradient of cell normals fixed +appctx = {'block_appctx': block_appctx} -# check convergence of each timestep +# function spaces and initial conditions +W = swe.default_function_space(mesh) -def post_function_callback(aaosolver, X, F): - if args.print_res: - residuals = asQ.SharedArray(time_partition, - comm=aaosolver.ensemble.ensemble_comm) - # all-at-once residual - res = aaosolver.aaoform.F - for i in range(res.nlocal_timesteps): - with res[i].dat.vec_ro as vec: - residuals.dlocal[i] = vec.norm() - residuals.synchronise() - PETSc.Sys.Print('') - PETSc.Sys.Print([f"{r:.4e}" for r in residuals.data()]) - PETSc.Sys.Print('') +winitial = fd.Function(W) +uinitial, hinitial = winitial.subfunctions +uinitial.project(galewsky.velocity_expression(*coords)) +hinitial.interpolate(galewsky.depth_expression(*coords)) +# all at once solver -PETSc.Sys.Print('### === --- Calculating parallel solution --- === ###') +chunk_aaofunc = asQ.AllAtOnceFunction(chunk_ensemble, chunk_partition, W) +chunk_aaofunc.assign(winitial) -appctx = {'block_appctx': block_appctx} +theta = 0.5 +chunk_aaoform = asQ.AllAtOnceForm(chunk_aaofunc, dt, theta, + form_mass, form_function) + +chunk_solver = asQ.AllAtOnceSolver(chunk_aaoform, chunk_aaofunc, + solver_parameters=sparameters_diag, + appctx=appctx) + +# which part of total timeseries is each chunk currently? +chunk_indexes = np.array((i for i in range(args.nchunks)), dtype=int) -miniapp = swe.ShallowWaterMiniApp(gravity=earth.Gravity, - topography_expression=galewsky.topography_expression, - velocity_expression=galewsky.velocity_expression, - depth_expression=galewsky.depth_expression, - reference_depth=galewsky.H0, - reference_state=True, - create_mesh=create_mesh, - dt=dt, theta=0.5, - time_partition=time_partition, - appctx=appctx, - paradiag_sparameters=sparameters_diag, - file_name='output/'+args.filename, - post_function_callback=post_function_callback, - record_diagnostics={'cfl': True, 'file': False}) - -paradiag = miniapp.paradiag -aaosolver = paradiag.solver -aaofunc = aaosolver.aaofunc -aaoform = aaosolver.aaoform - -chunks = tuple(aaofunc.copy() for _ in range(args.nchunks)) +# which chunks are currently at the beginning/end of the sweep? +first_chunk = 0 +last_chunk = args.nchunks - 1 + +PETSc.Sys.Print('### === --- Calculating parallel solution --- === ###') nconverged = 0 PETSc.Sys.Print('') @@ -204,7 +207,7 @@ def post_function_callback(aaosolver, X, F): global_comm.Barrier() if chunk_id == i: - aaosolver.solve() + chunk_solver.solve() global_comm.Barrier() PETSc.Sys.Print("") @@ -234,11 +237,6 @@ def post_function_callback(aaosolver, X, F): nsweeps = j -pc_block_its = aaosolver.jacobian.pc.block_iterations -pc_block_its.synchronise() -pc_block_its = pc_block_its.data(deepcopy=True) -pc_block_its = pc_block_its/args.nchunks - niterations = nsweeps*args.nsmooth PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") @@ -249,5 +247,3 @@ def post_function_callback(aaosolver, X, F): PETSc.Sys.Print(f"Number of sweeps per converged chunk: {nsweeps/nconverged if nconverged else 'n/a'}") PETSc.Sys.Print(f"Number of iterations per converged chunk: {niterations/nconverged if nconverged else 'n/a'}") PETSc.Sys.Print(f"Number of timesteps per iteration: {nconverged*chunk_length/niterations}") -PETSc.Sys.Print(f'Block iterations: {pc_block_its}') -PETSc.Sys.Print(f'Block iterations per block solve: {pc_block_its/niterations}') From 0242daac22353abb476dec80931d80281a3b8aa9 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 11 Mar 2024 08:22:52 +0000 Subject: [PATCH 09/18] I can't spell German --- case_studies/{gauss_siedel => gauss_seidel}/galewsky.py | 0 case_studies/{gauss_siedel => gauss_seidel}/galewsky_serial.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename case_studies/{gauss_siedel => gauss_seidel}/galewsky.py (100%) rename case_studies/{gauss_siedel => gauss_seidel}/galewsky_serial.py (100%) diff --git a/case_studies/gauss_siedel/galewsky.py b/case_studies/gauss_seidel/galewsky.py similarity index 100% rename from case_studies/gauss_siedel/galewsky.py rename to case_studies/gauss_seidel/galewsky.py diff --git a/case_studies/gauss_siedel/galewsky_serial.py b/case_studies/gauss_seidel/galewsky_serial.py similarity index 100% rename from case_studies/gauss_siedel/galewsky_serial.py rename to case_studies/gauss_seidel/galewsky_serial.py From 384ec50e105f686cafc3eacc5832a2c1ba003caa Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 11 Mar 2024 19:04:54 +0000 Subject: [PATCH 10/18] add ensemble file to asQ/__init__ --- asQ/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/asQ/__init__.py b/asQ/__init__.py index 4f645719..76246a3e 100644 --- a/asQ/__init__.py +++ b/asQ/__init__.py @@ -2,5 +2,6 @@ from asQ.paradiag import * # noqa: F401 from asQ.diag_preconditioner import * # noqa: F401 from asQ.allatonce import * # noqa: F401 +from asQ.ensemble import * # noqa: F401 from asQ.post import * # noqa: F401 from asQ import complex_proxy # noqa: F401 From fcacf5664a1a515ea78fc3c467d27df6eff29e9f Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 11 Mar 2024 19:05:45 +0000 Subject: [PATCH 11/18] initialisation loop for gauss-seidel --- case_studies/gauss_seidel/galewsky.py | 132 +++++++++++++++++++------- 1 file changed, 96 insertions(+), 36 deletions(-) diff --git a/case_studies/gauss_seidel/galewsky.py b/case_studies/gauss_seidel/galewsky.py index af6790ef..176d6186 100644 --- a/case_studies/gauss_seidel/galewsky.py +++ b/case_studies/gauss_seidel/galewsky.py @@ -187,22 +187,53 @@ def aux_form_function(u, h, v, q, t=None): appctx=appctx) # which part of total timeseries is each chunk currently? -chunk_indexes = np.array((i for i in range(args.nchunks)), dtype=int) +chunk_indexes = np.array([i for i in range(args.nchunks)], dtype=int) # which chunks are currently at the beginning/end of the sweep? first_chunk = 0 last_chunk = args.nchunks - 1 +# update chunk ics from previous chunk +uprev = fd.Function(W) + + +def update_chunk_halos(uhalo): + chunk_begin = chunk_aaofunc.layout.is_local(0) + chunk_end = chunk_aaofunc.layout.is_local(-1) + + global_rank = global_ensemble.ensemble_comm.rank + global_size = global_ensemble.ensemble_comm.size + + # ring communication so the first chunk can + # pick up after last chunk after convergence + + # send forward + if chunk_end: + dst = (global_rank + 1) % global_size + global_ensemble.send(chunk_aaofunc[-1], dest=dst, tag=dst) + + # recv previous + if chunk_begin: + src = (global_rank - 1) % global_size + global_ensemble.recv(uhalo, source=src, tag=global_rank) + + # broadcast new ics to all ranks + chunk_ensemble.bcast(uhalo) + + PETSc.Sys.Print('### === --- Calculating parallel solution --- === ###') +PETSc.Sys.Print('') nconverged = 0 + +PETSc.Sys.Print('### === --- Initialising all chunks --- === ###') PETSc.Sys.Print('') -for j in range(args.nsweeps): - PETSc.Sys.Print(f'### === --- Calculating nonlinear sweep {j} --- === ###') +for j in range(args.nchunks): + PETSc.Sys.Print(f' === --- Initial nonlinear sweep {j} --- === ') PETSc.Sys.Print('') - # 1) one smoothing application on each chunk - for i in range(args.nchunks): + # only smooth chunks that the first sweep has reached + for i in range(j+1): PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') global_comm.Barrier() @@ -212,38 +243,67 @@ def aux_form_function(u, h, v, q, t=None): PETSc.Sys.Print("") - # 2) update ics of each chunk from previous chunk - update_ics(global_aaofunc, chunk_aaofunc) - - # 3) shuffle earlier chunks if converged - if j < args.ninitialise: - break + # propogate solution + update_chunk_halos(uprev) - update_chunk_residuals(chunk_aaofunc, chunk_residuals) - nshuffle = count_converged(chunk_residuals, args.nchecks) - shuffle_chunks(global_aaofunc, chunk_aaofunc, nshuffle) + # update initial conditions + if chunk_id != 0: + chunk_aaofunc.initial_condition.assign(uprev) - nconverged += nshuffle + # initial guess before first sweep is persistence forecast + if chunk_id > j: + chunk_aaofunc.assign(chunk_aaofunc.initial_condition) - PETSc.Sys.Print('') - PETSc.Sys.Print(f">>> Converged chunks: {int(nconverged)}.") - converged_time = nconverged*chunk_length*args.dt - PETSc.Sys.Print(f">>> Converged time: {converged_time} hours.") - PETSc.Sys.Print('') - - if nconverged >= args.nwindows: - PETSc.Sys.Print(f"Finished iterating to {args.nwindows}.") - break - -nsweeps = j - -niterations = nsweeps*args.nsmooth +PETSc.Sys.Print('### === --- All chunks initialised --- === ###') +PETSc.Sys.Print('') -PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") -PETSc.Sys.Print(f"Maximum number of sweeps: {args.nsweeps}") -PETSc.Sys.Print(f"Actual number of sweeps: {nsweeps}") -PETSc.Sys.Print(f"Number of chunks converged: {int(nconverged)}") -PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/nsweeps}") -PETSc.Sys.Print(f"Number of sweeps per converged chunk: {nsweeps/nconverged if nconverged else 'n/a'}") -PETSc.Sys.Print(f"Number of iterations per converged chunk: {niterations/nconverged if nconverged else 'n/a'}") -PETSc.Sys.Print(f"Number of timesteps per iteration: {nconverged*chunk_length/niterations}") +# for j in range(args.nsweeps): +# PETSc.Sys.Print(f' === --- Calculating nonlinear sweep {j} --- === ') +# PETSc.Sys.Print('') +# +# # 1) one smoothing application on each chunk +# for i in range(args.nchunks): +# PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') +# +# global_comm.Barrier() +# if chunk_id == i: +# chunk_solver.solve() +# global_comm.Barrier() +# +# PETSc.Sys.Print("") +# +# # 2) update ics of each chunk from previous chunk +# update_chunk_ics() +# +# # 3) shuffle earlier chunks if converged +# if j < args.ninitialise: +# break +# +# update_chunk_residuals(chunk_aaofunc, chunk_residuals) +# nshuffle = count_converged(chunk_residuals, args.nchecks) +# shuffle_chunks(global_aaofunc, chunk_aaofunc, nshuffle) +# +# nconverged += nshuffle +# +# PETSc.Sys.Print('') +# PETSc.Sys.Print(f">>> Converged chunks: {int(nconverged)}.") +# converged_time = nconverged*chunk_length*args.dt +# PETSc.Sys.Print(f">>> Converged time: {converged_time} hours.") +# PETSc.Sys.Print('') +# +# if nconverged >= args.nwindows: +# PETSc.Sys.Print(f"Finished iterating to {args.nwindows}.") +# break +# +# nsweeps = j +# +# niterations = nsweeps*args.nsmooth +# +# PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") +# PETSc.Sys.Print(f"Maximum number of sweeps: {args.nsweeps}") +# PETSc.Sys.Print(f"Actual number of sweeps: {nsweeps}") +# PETSc.Sys.Print(f"Number of chunks converged: {int(nconverged)}") +# PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/nsweeps}") +# PETSc.Sys.Print(f"Number of sweeps per converged chunk: {nsweeps/nconverged if nconverged else 'n/a'}") +# PETSc.Sys.Print(f"Number of iterations per converged chunk: {niterations/nconverged if nconverged else 'n/a'}") +# PETSc.Sys.Print(f"Number of timesteps per iteration: {nconverged*chunk_length/niterations}") From d32e9ebdfcfe66a58051aad0f55c220e718e1e14 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 11 Mar 2024 20:02:28 +0000 Subject: [PATCH 12/18] parallel gauss-seidel? --- case_studies/gauss_seidel/galewsky.py | 153 ++++++++++++++++---------- 1 file changed, 93 insertions(+), 60 deletions(-) diff --git a/case_studies/gauss_seidel/galewsky.py b/case_studies/gauss_seidel/galewsky.py index 176d6186..4fc9e4bc 100644 --- a/case_studies/gauss_seidel/galewsky.py +++ b/case_studies/gauss_seidel/galewsky.py @@ -1,12 +1,13 @@ from firedrake.petsc import PETSc import asQ -import firedrake as fd # noqa: F401 +import firedrake as fd from utils import units from utils.planets import earth import utils.shallow_water as swe from utils.shallow_water import galewsky +from time import sleep # noqa: F401 import numpy as np from math import sqrt @@ -28,7 +29,7 @@ parser.add_argument('--nsmooth', type=int, default=1, help='Number of nonlinear iterations per chunk at each sweep.') parser.add_argument('--nchecks', type=int, default=1, help='Maximum number of chunks allowed to converge after each sweep.') parser.add_argument('--ninitialise', type=int, default=0, help='Number of sweeps before checking convergence.') -parser.add_argument('--atol', type=float, default=1e0, help='Average atol of each timestep.') +parser.add_argument('--atol', type=float, default=1e5, help='Average atol of each timestep.') parser.add_argument('--nslices', type=int, default=2, help='Number of time-slices per time-window.') parser.add_argument('--slice_length', type=int, default=2, help='Number of timesteps per time-slice.') parser.add_argument('--alpha', type=float, default=1e-4, help='Circulant coefficient.') @@ -186,9 +187,15 @@ def aux_form_function(u, h, v, q, t=None): solver_parameters=sparameters_diag, appctx=appctx) -# which part of total timeseries is each chunk currently? +# which chunk is holding which part of the total timeseries? chunk_indexes = np.array([i for i in range(args.nchunks)], dtype=int) +# we need to make this an array so we can send it via mpi +convergence_flag = np.array([False], dtype=bool) + +# first mpi rank on each chunk (assumes all chunks are equal size): +chunk_root = lambda c: c*chunk_ensemble.global_comm.size + # which chunks are currently at the beginning/end of the sweep? first_chunk = 0 last_chunk = args.nchunks - 1 @@ -228,20 +235,24 @@ def update_chunk_halos(uhalo): PETSc.Sys.Print('### === --- Initialising all chunks --- === ###') PETSc.Sys.Print('') +sleep_time = 0.1 for j in range(args.nchunks): PETSc.Sys.Print(f' === --- Initial nonlinear sweep {j} --- === ') PETSc.Sys.Print('') # only smooth chunks that the first sweep has reached - for i in range(j+1): - PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') + # for i in range(j+1): + # PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') - global_comm.Barrier() - if chunk_id == i: - chunk_solver.solve() - global_comm.Barrier() + # global_comm.Barrier(); sleep(sleep_time) + # if chunk_id == i: + # chunk_solver.solve() + # global_comm.Barrier(); sleep(sleep_time) - PETSc.Sys.Print("") + # PETSc.Sys.Print("") + + if chunk_id < j+1: + chunk_solver.solve() # propogate solution update_chunk_halos(uprev) @@ -257,53 +268,75 @@ def update_chunk_halos(uhalo): PETSc.Sys.Print('### === --- All chunks initialised --- === ###') PETSc.Sys.Print('') -# for j in range(args.nsweeps): -# PETSc.Sys.Print(f' === --- Calculating nonlinear sweep {j} --- === ') -# PETSc.Sys.Print('') -# -# # 1) one smoothing application on each chunk -# for i in range(args.nchunks): -# PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') -# -# global_comm.Barrier() -# if chunk_id == i: -# chunk_solver.solve() -# global_comm.Barrier() -# -# PETSc.Sys.Print("") -# -# # 2) update ics of each chunk from previous chunk -# update_chunk_ics() -# -# # 3) shuffle earlier chunks if converged -# if j < args.ninitialise: -# break -# -# update_chunk_residuals(chunk_aaofunc, chunk_residuals) -# nshuffle = count_converged(chunk_residuals, args.nchecks) -# shuffle_chunks(global_aaofunc, chunk_aaofunc, nshuffle) -# -# nconverged += nshuffle -# -# PETSc.Sys.Print('') -# PETSc.Sys.Print(f">>> Converged chunks: {int(nconverged)}.") -# converged_time = nconverged*chunk_length*args.dt -# PETSc.Sys.Print(f">>> Converged time: {converged_time} hours.") -# PETSc.Sys.Print('') -# -# if nconverged >= args.nwindows: -# PETSc.Sys.Print(f"Finished iterating to {args.nwindows}.") -# break -# -# nsweeps = j -# -# niterations = nsweeps*args.nsmooth -# -# PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") -# PETSc.Sys.Print(f"Maximum number of sweeps: {args.nsweeps}") -# PETSc.Sys.Print(f"Actual number of sweeps: {nsweeps}") -# PETSc.Sys.Print(f"Number of chunks converged: {int(nconverged)}") -# PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/nsweeps}") -# PETSc.Sys.Print(f"Number of sweeps per converged chunk: {nsweeps/nconverged if nconverged else 'n/a'}") -# PETSc.Sys.Print(f"Number of iterations per converged chunk: {niterations/nconverged if nconverged else 'n/a'}") -# PETSc.Sys.Print(f"Number of timesteps per iteration: {nconverged*chunk_length/niterations}") +for j in range(args.nsweeps): + PETSc.Sys.Print(f' === --- Calculating nonlinear sweep {j} --- === ') + PETSc.Sys.Print('') + + # 1) one smoothing application on each chunk + # for i in range(args.nchunks): + # PETSc.Sys.Print(f' --- Calculating chunk {i} on solver {chunk_indexes[i]} --- ') + + # global_comm.Barrier(); sleep(sleep_time) + # if chunk_id == chunk_indexes[i]: + # chunk_solver.solve() + # global_comm.Barrier(); sleep(sleep_time) + + # PETSc.Sys.Print("") + chunk_solver.solve() + + # 2) update ics of each chunk from previous chunk + update_chunk_halos(uprev) + + earliest_chunk = (chunk_id == chunk_indexes[0]) + + # everyone uses latest ic guesses except chunk 0 (uses 'exact' ic) + if not earliest_chunk: + chunk_aaofunc.initial_condition.assign(uprev) + + # 3) check convergence of earliest chunk + + if earliest_chunk: + chunk_aaoform.assemble() + with chunk_aaoform.F.global_vec_ro() as rvec: + res = rvec.norm() + convergence_flag[0] = (res < patol) + + # rank 0 on the earliest chunk tells everyone if they've converged + global_ensemble.global_comm.Bcast(convergence_flag, + root=chunk_root(chunk_indexes[0])) + + # 4) shuffle if converged + if convergence_flag[0]: + # earliest chunk becomes last chunk + if earliest_chunk: + chunk_aaofunc.assign(uprev) + + # update record of which chunk is in which position + for i in range(args.nchunks): + chunk_indexes[i] = (chunk_indexes[i] + 1) % args.nchunks + + nconverged += 1 + + PETSc.Sys.Print('') + PETSc.Sys.Print(f">>> Converged chunks: {int(nconverged)}.") + converged_time = nconverged*chunk_length*args.dt + PETSc.Sys.Print(f">>> Converged time: {converged_time} hours.") + PETSc.Sys.Print('') + + if nconverged >= args.nwindows: + PETSc.Sys.Print(f"Finished iterating to {args.nwindows} windows.") + PETSc.Sys.Print('') + break + +nsweeps = j + +niterations = nsweeps*args.nsmooth + +PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") +PETSc.Sys.Print(f"Maximum number of sweeps: {args.nsweeps}") +PETSc.Sys.Print(f"Actual number of sweeps: {nsweeps}") +PETSc.Sys.Print(f"Number of chunks converged: {int(nconverged)}") +PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/nsweeps}") +PETSc.Sys.Print(f"Number of sweeps per converged chunk: {nsweeps/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Number of iterations per converged chunk: {niterations/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Number of timesteps per iteration: {nconverged*chunk_length/niterations}") From e24161dbaa34a86edabaa0012a2c8ad3169d6ce5 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Tue, 12 Mar 2024 10:21:30 +0000 Subject: [PATCH 13/18] switch between serial and parallel gauss-seidel sweeps --- case_studies/gauss_seidel/galewsky.py | 118 +++++++++++++------------- 1 file changed, 61 insertions(+), 57 deletions(-) diff --git a/case_studies/gauss_seidel/galewsky.py b/case_studies/gauss_seidel/galewsky.py index 4fc9e4bc..d049d55e 100644 --- a/case_studies/gauss_seidel/galewsky.py +++ b/case_studies/gauss_seidel/galewsky.py @@ -16,27 +16,22 @@ # get command arguments import argparse parser = argparse.ArgumentParser( - description='Galewsky testcase for ParaDiag solver using fully implicit SWE solver.', + description='Galewsky testcase for ParaDiag solver using a pipelined nonlinear Gauss-Seidel method.', formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser.add_argument('--ref_level', type=int, default=2, help='Refinement level of icosahedral grid.') parser.add_argument('--base_level', type=int, default=2, help='Refinement level of coarse grid.') parser.add_argument('--nwindows', type=int, default=1, help='Total number of time-windows.') -parser.add_argument('--nlmethod', type=str, default='gs', choices=['gs', 'jac'], help='Nonlinear method. "gs" for Gauss-Siedel or "jac" for Jacobi.') parser.add_argument('--nchunks', type=int, default=4, help='Number of chunks to solve simultaneously.') parser.add_argument('--nsweeps', type=int, default=4, help='Number of nonlinear sweeps.') parser.add_argument('--nsmooth', type=int, default=1, help='Number of nonlinear iterations per chunk at each sweep.') -parser.add_argument('--nchecks', type=int, default=1, help='Maximum number of chunks allowed to converge after each sweep.') -parser.add_argument('--ninitialise', type=int, default=0, help='Number of sweeps before checking convergence.') parser.add_argument('--atol', type=float, default=1e5, help='Average atol of each timestep.') parser.add_argument('--nslices', type=int, default=2, help='Number of time-slices per time-window.') parser.add_argument('--slice_length', type=int, default=2, help='Number of timesteps per time-slice.') parser.add_argument('--alpha', type=float, default=1e-4, help='Circulant coefficient.') parser.add_argument('--dt', type=float, default=0.5, help='Timestep in hours.') -parser.add_argument('--filename', type=str, default='galewsky', help='Name of output vtk files') -parser.add_argument('--metrics_dir', type=str, default='metrics', help='Directory to save paradiag metrics to.') -parser.add_argument('--print_res', action='store_true', help='Print the residuals of each timestep at each iteration.') +parser.add_argument('--serial', action='store_true', help='Calculate each chunk in serial.') parser.add_argument('--show_args', action='store_true', help='Output all the arguments.') args = parser.parse_known_args() @@ -117,7 +112,7 @@ def aux_form_function(u, h, v, q, t=None): sparameters = { 'mat_type': 'matfree', - 'ksp_type': 'fgmres', + 'ksp_type': 'gmres', 'ksp': { 'atol': 1e-5, 'rtol': 1e-5, @@ -154,7 +149,7 @@ def aux_form_function(u, h, v, q, t=None): 'atol': patol, }, 'pc_type': 'python', - 'pc_python_type': 'asQ.DiagFFTPC', + 'pc_python_type': 'asQ.CirculantPC', 'diagfft_alpha': args.alpha, 'diagfft_state': 'window', 'aaos_jacobian_state': 'current' @@ -187,8 +182,8 @@ def aux_form_function(u, h, v, q, t=None): solver_parameters=sparameters_diag, appctx=appctx) -# which chunk is holding which part of the total timeseries? -chunk_indexes = np.array([i for i in range(args.nchunks)], dtype=int) +# which chunk_id is holding which part of the total timeseries? +chunk_indexes = np.array([*range(args.nchunks)], dtype=int) # we need to make this an array so we can send it via mpi convergence_flag = np.array([False], dtype=bool) @@ -196,9 +191,8 @@ def aux_form_function(u, h, v, q, t=None): # first mpi rank on each chunk (assumes all chunks are equal size): chunk_root = lambda c: c*chunk_ensemble.global_comm.size -# which chunks are currently at the beginning/end of the sweep? -first_chunk = 0 -last_chunk = args.nchunks - 1 +# am I the chunk with the earliest timesteps? +earliest = lambda: (chunk_id == chunk_indexes[0]) # update chunk ics from previous chunk uprev = fd.Function(W) @@ -214,15 +208,15 @@ def update_chunk_halos(uhalo): # ring communication so the first chunk can # pick up after last chunk after convergence - # send forward + # send forward last step of chunk if chunk_end: - dst = (global_rank + 1) % global_size - global_ensemble.send(chunk_aaofunc[-1], dest=dst, tag=dst) + dest = (global_rank + 1) % global_size + global_ensemble.send(chunk_aaofunc[-1], dest=dest, tag=dest) - # recv previous + # recv updated ics from previous chunk if chunk_begin: - src = (global_rank - 1) % global_size - global_ensemble.recv(uhalo, source=src, tag=global_rank) + source = (global_rank - 1) % global_size + global_ensemble.recv(uhalo, source=source, tag=global_rank) # broadcast new ics to all ranks chunk_ensemble.bcast(uhalo) @@ -235,33 +229,37 @@ def update_chunk_halos(uhalo): PETSc.Sys.Print('### === --- Initialising all chunks --- === ###') PETSc.Sys.Print('') -sleep_time = 0.1 +sleep_time = 0.01 for j in range(args.nchunks): PETSc.Sys.Print(f' === --- Initial nonlinear sweep {j} --- === ') PETSc.Sys.Print('') # only smooth chunks that the first sweep has reached - # for i in range(j+1): - # PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') - # global_comm.Barrier(); sleep(sleep_time) - # if chunk_id == i: - # chunk_solver.solve() - # global_comm.Barrier(); sleep(sleep_time) + if args.serial: + for i in range(j+1): + PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') - # PETSc.Sys.Print("") + global_comm.Barrier() + sleep(sleep_time) + if chunk_id == i: + chunk_solver.solve() + global_comm.Barrier() + sleep(sleep_time) - if chunk_id < j+1: - chunk_solver.solve() + PETSc.Sys.Print("") + else: + if chunk_id < j+1: + chunk_solver.solve() # propogate solution update_chunk_halos(uprev) - # update initial conditions + # update initial condition guess for later chunks if chunk_id != 0: chunk_aaofunc.initial_condition.assign(uprev) - # initial guess before first sweep is persistence forecast + # initial guess in front of first sweep is persistence forecast if chunk_id > j: chunk_aaofunc.assign(chunk_aaofunc.initial_condition) @@ -273,29 +271,32 @@ def update_chunk_halos(uhalo): PETSc.Sys.Print('') # 1) one smoothing application on each chunk - # for i in range(args.nchunks): - # PETSc.Sys.Print(f' --- Calculating chunk {i} on solver {chunk_indexes[i]} --- ') + if args.serial: + for i in range(args.nchunks): + PETSc.Sys.Print(f' --- Calculating chunk {i} on solver {chunk_indexes[i]} --- ') - # global_comm.Barrier(); sleep(sleep_time) - # if chunk_id == chunk_indexes[i]: - # chunk_solver.solve() - # global_comm.Barrier(); sleep(sleep_time) + global_comm.Barrier() + sleep(sleep_time) + if chunk_id == chunk_indexes[i]: + chunk_solver.solve() + global_comm.Barrier() + sleep(sleep_time) - # PETSc.Sys.Print("") - chunk_solver.solve() + PETSc.Sys.Print("") + + else: + chunk_solver.solve() # 2) update ics of each chunk from previous chunk update_chunk_halos(uprev) - earliest_chunk = (chunk_id == chunk_indexes[0]) - - # everyone uses latest ic guesses except chunk 0 (uses 'exact' ic) - if not earliest_chunk: + # everyone uses latest ic guesses, except chunk + # with earliest timesteps (already has 'exact' ic) + if not earliest(): chunk_aaofunc.initial_condition.assign(uprev) # 3) check convergence of earliest chunk - - if earliest_chunk: + if earliest(): chunk_aaoform.assemble() with chunk_aaoform.F.global_vec_ro() as rvec: res = rvec.norm() @@ -305,29 +306,32 @@ def update_chunk_halos(uhalo): global_ensemble.global_comm.Bcast(convergence_flag, root=chunk_root(chunk_indexes[0])) - # 4) shuffle if converged + # update and report if convergence_flag[0]: - # earliest chunk becomes last chunk - if earliest_chunk: - chunk_aaofunc.assign(uprev) - - # update record of which chunk is in which position - for i in range(args.nchunks): - chunk_indexes[i] = (chunk_indexes[i] + 1) % args.nchunks - nconverged += 1 - PETSc.Sys.Print('') - PETSc.Sys.Print(f">>> Converged chunks: {int(nconverged)}.") converged_time = nconverged*chunk_length*args.dt + PETSc.Sys.Print('') + PETSc.Sys.Print(f">>> Converged chunks: {nconverged}.") PETSc.Sys.Print(f">>> Converged time: {converged_time} hours.") PETSc.Sys.Print('') + # 4) stop iterating if we've reached the end if nconverged >= args.nwindows: PETSc.Sys.Print(f"Finished iterating to {args.nwindows} windows.") PETSc.Sys.Print('') break + # 5) shuffle and restart if we haven't reached the end + if convergence_flag[0]: + # earliest chunk_id becomes last chunk + if earliest(): + chunk_aaofunc.assign(uprev) + + # update record of which chunk_id is in which position + for i in range(args.nchunks): + chunk_indexes[i] = (chunk_indexes[i] + 1) % args.nchunks + nsweeps = j niterations = nsweeps*args.nsmooth From 5768808d82b0559d2e5df34edd4a76667245204b Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Wed, 13 Mar 2024 16:00:50 +0000 Subject: [PATCH 14/18] remove duplicate asQ.ensemble import --- asQ/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/asQ/__init__.py b/asQ/__init__.py index 173d3587..1e5fb660 100644 --- a/asQ/__init__.py +++ b/asQ/__init__.py @@ -3,6 +3,5 @@ from asQ.paradiag import * # noqa: F401 from asQ.preconditioners import * # noqa: F401 from asQ.allatonce import * # noqa: F401 -from asQ.ensemble import * # noqa: F401 from asQ.post import * # noqa: F401 from asQ import complex_proxy # noqa: F401 From 1a5db26b0ded6feb84a52666074ffc1e838cd07b Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Wed, 13 Mar 2024 16:01:41 +0000 Subject: [PATCH 15/18] advection gs example --- case_studies/gauss_seidel/advection.py | 347 +++++++++++++++++++++++++ 1 file changed, 347 insertions(+) create mode 100644 case_studies/gauss_seidel/advection.py diff --git a/case_studies/gauss_seidel/advection.py b/case_studies/gauss_seidel/advection.py new file mode 100644 index 00000000..c3f611ed --- /dev/null +++ b/case_studies/gauss_seidel/advection.py @@ -0,0 +1,347 @@ +from firedrake.petsc import PETSc + +import asQ +import firedrake as fd + +from time import sleep # noqa: F401 +import numpy as np +from math import sqrt + +PETSc.Sys.popErrorHandler() + +# get command arguments +import argparse +parser = argparse.ArgumentParser( + description='DG scalar advection testcase for ParaDiag solver using a pipelined nonlinear Gauss-Seidel method.', + formatter_class=argparse.ArgumentDefaultsHelpFormatter +) + +parser.add_argument('--nx', type=int, default=16, help='Number of cells along each square side.') +parser.add_argument('--cfl', type=float, default=0.8, help='Convective CFL number.') +parser.add_argument('--angle', type=float, default=pi/6, help='Angle of the convective velocity.') +parser.add_argument('--degree', type=int, default=1, help='Degree of the scalar spaces.') +parser.add_argument('--width', type=float, default=0.2, help='Width of the Gaussian bump.') +parser.add_argument('--nwindows', type=int, default=1, help='Total number of time-windows.') +parser.add_argument('--nchunks', type=int, default=4, help='Number of chunks to solve simultaneously.') +parser.add_argument('--nsweeps', type=int, default=4, help='Number of nonlinear sweeps.') +parser.add_argument('--nsmooth', type=int, default=1, help='Number of nonlinear iterations per chunk at each sweep.') +parser.add_argument('--atol', type=float, default=1e5, help='Average atol of each timestep.') +parser.add_argument('--nslices', type=int, default=2, help='Number of time-slices per time-window.') +parser.add_argument('--slice_length', type=int, default=2, help='Number of timesteps per time-slice.') +parser.add_argument('--alpha', type=float, default=1e-1, help='Circulant coefficient.') +parser.add_argument('--theta', type=float, default=0.5, help='Parameter for the implicit theta timestepping method.') +parser.add_argument('--dt', type=float, default=0.5, help='Timestep in hours.') +parser.add_argument('--serial', action='store_true', help='Calculate each chunk in serial.') +parser.add_argument('--show_args', action='store_true', help='Output all the arguments.') + +args = parser.parse_known_args() +args = args[0] + +if args.show_args: + PETSc.Sys.Print(args) + +PETSc.Sys.Print('') +PETSc.Sys.Print('### === --- Setting up --- === ###') +PETSc.Sys.Print('') + +# time steps + +chunk_partition = tuple((args.slice_length for _ in range(args.nslices))) +chunk_length = sum(chunk_partition) +total_timesteps = chunk_length*args.nchunks +total_slices = args.nslices*args.nchunks + +global_comm = fd.COMM_WORLD +global_time_partition = tuple((args.slice_length for _ in range(total_slices))) +global_ensemble = asQ.create_ensemble(global_time_partition, global_comm) +chunk_ensemble = asQ.split_ensemble(global_ensemble, args.nslices) + +# which chunk are we? +chunk_id = global_ensemble.ensemble_comm.rank // args.nslices + +# Calculate the timestep from the CFL number +umax = 1. +dx = 1./args.nx +dt = args.cfl*dx/umax + +# # # === --- domain and FE spaces --- === # # # + +mesh = fd.PeriodicUnitSquareMesh(args.nx, args.nx, quadrilateral=True, comm=chunk_ensemble.comm) + +V = fd.FunctionSpace(mesh, "DQ", args.degree) + +# # # === --- initial conditions --- === # # # + +x, y = fd.SpatialCoordinate(mesh) + + +def radius(x, y): + return fd.sqrt(pow(x-0.5, 2) + pow(y-0.5, 2)) + + +def gaussian(x, y): + return fd.exp(-0.5*pow(radius(x, y)/args.width, 2)) + + +# Gaussian bump centred at (0.5, 0.5) +q0 = fd.Function(V, name="scalar_initial") +q0.interpolate(1 + gaussian(x, y)) + +# The advecting velocity at angle to the x-axis +u = fd.Constant(fd.as_vector((umax*cos(args.angle), umax*sin(args.angle)))) + +# # # === --- finite element forms --- === # # # + + +# The time-derivative mass form for the scalar advection equation. +# asQ assumes that the mass form is linear so here +# q is a TrialFunction and phi is a TestFunction +def form_mass(q, phi): + return phi*q*fd.dx + + +# The DG advection form for the scalar advection equation. +# asQ assumes that the function form is nonlinear so here +# q is a Function and phi is a TestFunction +def form_function(q, phi, t): + # upwind switch + n = fd.FacetNormal(mesh) + un = fd.Constant(0.5)*(fd.dot(u, n) + abs(fd.dot(u, n))) + + # integration over element volume + int_cell = q*fd.div(phi*u)*fd.dx + + # integration over internal facets + int_facet = (phi('+')-phi('-'))*(un('+')*q('+')-un('-')*q('-'))*fd.dS + + return int_facet - int_cell + + +# # # === --- PETSc solver parameters --- === # # # + +# parameters for the implicit diagonal solve in step-(b) +block_parameters = { + 'ksp_type': 'preonly', + 'pc_type': 'lu', + 'pc_factor_mat_solver_type': 'mumps' + # 'pc_factor_mat_ordering_type': 'rcm', +} + +atol = args.atol +patol = sqrt(chunk_length)*atol +sparameters_diag = { + 'snes': { + 'linesearch_type': 'ksponly', + 'monitor': None, + 'converged_reason': None, + 'atol': patol, + 'rtol': 1e-10, + 'stol': 1e-12, + # 'ksp_ew': None, + # 'ksp_ew_version': 1, + # 'ksp_ew_threshold': 1e-2, + 'max_it': args.nsmooth, + 'convergence_test': 'skip', + }, + 'mat_type': 'matfree', + 'ksp_type': 'richardson', + 'ksp': { + 'rtol': 1e-2, + 'atol': patol, + }, + 'pc_type': 'python', + 'pc_python_type': 'asQ.CirculantPC', + 'diagfft_alpha': args.alpha, + 'diagfft_state': 'window', + 'aaos_jacobian_state': 'current' +} + +for i in range(chunk_length): + sparameters_diag['diagfft_block_'+str(i)+'_'] = block_parameters + +appctx = {'block_appctx': block_appctx} + +# function spaces and initial conditions + +W = swe.default_function_space(mesh) + +winitial = fd.Function(W) +uinitial, hinitial = winitial.subfunctions +uinitial.project(galewsky.velocity_expression(*coords)) +hinitial.interpolate(galewsky.depth_expression(*coords)) + +# all at once solver + +chunk_aaofunc = asQ.AllAtOnceFunction(chunk_ensemble, chunk_partition, W) +chunk_aaofunc.assign(winitial) + +theta = 0.5 +chunk_aaoform = asQ.AllAtOnceForm(chunk_aaofunc, dt, theta, + form_mass, form_function) + +chunk_solver = asQ.AllAtOnceSolver(chunk_aaoform, chunk_aaofunc, + solver_parameters=sparameters_diag, + appctx=appctx) + +# which chunk_id is holding which part of the total timeseries? +chunk_indexes = np.array([*range(args.nchunks)], dtype=int) + +# we need to make this an array so we can send it via mpi +convergence_flag = np.array([False], dtype=bool) + +# first mpi rank on each chunk (assumes all chunks are equal size): +chunk_root = lambda c: c*chunk_ensemble.global_comm.size + +# am I the chunk with the earliest timesteps? +earliest = lambda: (chunk_id == chunk_indexes[0]) + +# update chunk ics from previous chunk +uprev = fd.Function(W) + + +def update_chunk_halos(uhalo): + chunk_begin = chunk_aaofunc.layout.is_local(0) + chunk_end = chunk_aaofunc.layout.is_local(-1) + + global_rank = global_ensemble.ensemble_comm.rank + global_size = global_ensemble.ensemble_comm.size + + # ring communication so the first chunk can + # pick up after last chunk after convergence + + # send forward last step of chunk + if chunk_end: + dest = (global_rank + 1) % global_size + global_ensemble.send(chunk_aaofunc[-1], dest=dest, tag=dest) + + # recv updated ics from previous chunk + if chunk_begin: + source = (global_rank - 1) % global_size + global_ensemble.recv(uhalo, source=source, tag=global_rank) + + # broadcast new ics to all ranks + chunk_ensemble.bcast(uhalo) + + +PETSc.Sys.Print('### === --- Calculating parallel solution --- === ###') +PETSc.Sys.Print('') + +nconverged = 0 + +PETSc.Sys.Print('### === --- Initialising all chunks --- === ###') +PETSc.Sys.Print('') +sleep_time = 0.01 +for j in range(args.nchunks): + PETSc.Sys.Print(f' === --- Initial nonlinear sweep {j} --- === ') + PETSc.Sys.Print('') + + # only smooth chunks that the first sweep has reached + + if args.serial: + for i in range(j+1): + PETSc.Sys.Print(f' --- Calculating chunk {i} --- ') + + global_comm.Barrier() + sleep(sleep_time) + if chunk_id == i: + chunk_solver.solve() + global_comm.Barrier() + sleep(sleep_time) + + PETSc.Sys.Print("") + else: + if chunk_id < j+1: + chunk_solver.solve() + + # propogate solution + update_chunk_halos(uprev) + + # update initial condition guess for later chunks + if chunk_id != 0: + chunk_aaofunc.initial_condition.assign(uprev) + + # initial guess in front of first sweep is persistence forecast + if chunk_id > j: + chunk_aaofunc.assign(chunk_aaofunc.initial_condition) + +PETSc.Sys.Print('### === --- All chunks initialised --- === ###') +PETSc.Sys.Print('') + +for j in range(args.nsweeps): + PETSc.Sys.Print(f' === --- Calculating nonlinear sweep {j} --- === ') + PETSc.Sys.Print('') + + # 1) one smoothing application on each chunk + if args.serial: + for i in range(args.nchunks): + PETSc.Sys.Print(f' --- Calculating chunk {i} on solver {chunk_indexes[i]} --- ') + + global_comm.Barrier() + sleep(sleep_time) + if chunk_id == chunk_indexes[i]: + chunk_solver.solve() + global_comm.Barrier() + sleep(sleep_time) + + PETSc.Sys.Print("") + + else: + chunk_solver.solve() + + # 2) update ics of each chunk from previous chunk + update_chunk_halos(uprev) + + # everyone uses latest ic guesses, except chunk + # with earliest timesteps (already has 'exact' ic) + if not earliest(): + chunk_aaofunc.initial_condition.assign(uprev) + + # 3) check convergence of earliest chunk + if earliest(): + chunk_aaoform.assemble() + with chunk_aaoform.F.global_vec_ro() as rvec: + res = rvec.norm() + convergence_flag[0] = (res < patol) + + # rank 0 on the earliest chunk tells everyone if they've converged + global_ensemble.global_comm.Bcast(convergence_flag, + root=chunk_root(chunk_indexes[0])) + + # update and report + if convergence_flag[0]: + nconverged += 1 + + converged_time = nconverged*chunk_length*args.dt + PETSc.Sys.Print('') + PETSc.Sys.Print(f">>> Converged chunks: {nconverged}.") + PETSc.Sys.Print(f">>> Converged time: {converged_time} hours.") + PETSc.Sys.Print('') + + # 4) stop iterating if we've reached the end + if nconverged >= args.nwindows: + PETSc.Sys.Print(f"Finished iterating to {args.nwindows} windows.") + PETSc.Sys.Print('') + break + + # 5) shuffle and restart if we haven't reached the end + if convergence_flag[0]: + # earliest chunk_id becomes last chunk + if earliest(): + chunk_aaofunc.assign(uprev) + + # update record of which chunk_id is in which position + for i in range(args.nchunks): + chunk_indexes[i] = (chunk_indexes[i] + 1) % args.nchunks + +nsweeps = j + +niterations = nsweeps*args.nsmooth + +PETSc.Sys.Print(f"Number of chunks: {args.nchunks}") +PETSc.Sys.Print(f"Maximum number of sweeps: {args.nsweeps}") +PETSc.Sys.Print(f"Actual number of sweeps: {nsweeps}") +PETSc.Sys.Print(f"Number of chunks converged: {int(nconverged)}") +PETSc.Sys.Print(f"Number of chunks converged per sweep: {nconverged/nsweeps}") +PETSc.Sys.Print(f"Number of sweeps per converged chunk: {nsweeps/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Number of iterations per converged chunk: {niterations/nconverged if nconverged else 'n/a'}") +PETSc.Sys.Print(f"Number of timesteps per iteration: {nconverged*chunk_length/niterations}") From ddd6383b40c8305cbfd744b4251e1c0b05db7a99 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Thu, 14 Mar 2024 11:15:34 +0000 Subject: [PATCH 16/18] update EnsembleConnector with new pyop2.internal_comm interface --- asQ/ensemble.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/asQ/ensemble.py b/asQ/ensemble.py index 740ddd9a..9079e26c 100644 --- a/asQ/ensemble.py +++ b/asQ/ensemble.py @@ -62,15 +62,15 @@ def __init__(self, global_comm, local_comm, nmembers): raise ValueError(msg) self.global_comm = global_comm - self._global_comm = internal_comm(self.global_comm) + self._comm = internal_comm(self.global_comm, self) self.comm = local_comm - self._comm = internal_comm(self.comm) + self._spatial_comm = internal_comm(self.comm, self) self.ensemble_comm = self.global_comm.Split(color=self.comm.rank, key=global_comm.rank) - self._ensemble_comm = internal_comm(self.ensemble_comm) + self._ensemble_comm = internal_comm(self.ensemble_comm, self) def __del__(self): if hasattr(self, "ensemble_comm"): From 707f6d90da6312db39a13c7e2d542b7f2741d6b4 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Thu, 14 Mar 2024 11:35:59 +0000 Subject: [PATCH 17/18] advection gauss-seidel script with error check vs serial --- case_studies/gauss_seidel/advection.py | 57 ++++++++++++++------------ 1 file changed, 31 insertions(+), 26 deletions(-) diff --git a/case_studies/gauss_seidel/advection.py b/case_studies/gauss_seidel/advection.py index c3f611ed..c9cc4841 100644 --- a/case_studies/gauss_seidel/advection.py +++ b/case_studies/gauss_seidel/advection.py @@ -5,7 +5,7 @@ from time import sleep # noqa: F401 import numpy as np -from math import sqrt +from math import sqrt, pi, cos, sin PETSc.Sys.popErrorHandler() @@ -25,12 +25,11 @@ parser.add_argument('--nchunks', type=int, default=4, help='Number of chunks to solve simultaneously.') parser.add_argument('--nsweeps', type=int, default=4, help='Number of nonlinear sweeps.') parser.add_argument('--nsmooth', type=int, default=1, help='Number of nonlinear iterations per chunk at each sweep.') -parser.add_argument('--atol', type=float, default=1e5, help='Average atol of each timestep.') +parser.add_argument('--atol', type=float, default=1e-6, help='Average atol of each timestep.') parser.add_argument('--nslices', type=int, default=2, help='Number of time-slices per time-window.') parser.add_argument('--slice_length', type=int, default=2, help='Number of timesteps per time-slice.') parser.add_argument('--alpha', type=float, default=1e-1, help='Circulant coefficient.') parser.add_argument('--theta', type=float, default=0.5, help='Parameter for the implicit theta timestepping method.') -parser.add_argument('--dt', type=float, default=0.5, help='Timestep in hours.') parser.add_argument('--serial', action='store_true', help='Calculate each chunk in serial.') parser.add_argument('--show_args', action='store_true', help='Output all the arguments.') @@ -119,32 +118,36 @@ def form_function(q, phi, t): # # # === --- PETSc solver parameters --- === # # # +snes_linear_params = { + 'type': 'ksponly', + 'lag_jacobian': -2, + 'lag_jacobian_persists': None, + 'lag_preconditioner': -2, + 'lag_preconditioner_persists': None, +} + # parameters for the implicit diagonal solve in step-(b) block_parameters = { + 'snes': snes_linear_params, 'ksp_type': 'preonly', 'pc_type': 'lu', 'pc_factor_mat_solver_type': 'mumps' - # 'pc_factor_mat_ordering_type': 'rcm', } atol = args.atol patol = sqrt(chunk_length)*atol sparameters_diag = { 'snes': { - 'linesearch_type': 'ksponly', 'monitor': None, 'converged_reason': None, 'atol': patol, 'rtol': 1e-10, 'stol': 1e-12, - # 'ksp_ew': None, - # 'ksp_ew_version': 1, - # 'ksp_ew_threshold': 1e-2, 'max_it': args.nsmooth, 'convergence_test': 'skip', }, 'mat_type': 'matfree', - 'ksp_type': 'richardson', + 'ksp_type': 'preonly', 'ksp': { 'rtol': 1e-2, 'atol': patol, @@ -152,36 +155,27 @@ def form_function(q, phi, t): 'pc_type': 'python', 'pc_python_type': 'asQ.CirculantPC', 'diagfft_alpha': args.alpha, - 'diagfft_state': 'window', - 'aaos_jacobian_state': 'current' + 'diagfft_state': 'linear', + 'aaos_jacobian_state': 'linear' } +sparameters_diag['snes'].update(snes_linear_params) for i in range(chunk_length): sparameters_diag['diagfft_block_'+str(i)+'_'] = block_parameters -appctx = {'block_appctx': block_appctx} - # function spaces and initial conditions -W = swe.default_function_space(mesh) - -winitial = fd.Function(W) -uinitial, hinitial = winitial.subfunctions -uinitial.project(galewsky.velocity_expression(*coords)) -hinitial.interpolate(galewsky.depth_expression(*coords)) - # all at once solver -chunk_aaofunc = asQ.AllAtOnceFunction(chunk_ensemble, chunk_partition, W) -chunk_aaofunc.assign(winitial) +chunk_aaofunc = asQ.AllAtOnceFunction(chunk_ensemble, chunk_partition, V) +chunk_aaofunc.assign(q0) theta = 0.5 chunk_aaoform = asQ.AllAtOnceForm(chunk_aaofunc, dt, theta, form_mass, form_function) chunk_solver = asQ.AllAtOnceSolver(chunk_aaoform, chunk_aaofunc, - solver_parameters=sparameters_diag, - appctx=appctx) + solver_parameters=sparameters_diag) # which chunk_id is holding which part of the total timeseries? chunk_indexes = np.array([*range(args.nchunks)], dtype=int) @@ -196,7 +190,7 @@ def form_function(q, phi, t): earliest = lambda: (chunk_id == chunk_indexes[0]) # update chunk ics from previous chunk -uprev = fd.Function(W) +uprev = fd.Function(V) def update_chunk_halos(uhalo): @@ -311,7 +305,7 @@ def update_chunk_halos(uhalo): if convergence_flag[0]: nconverged += 1 - converged_time = nconverged*chunk_length*args.dt + converged_time = nconverged*chunk_length*dt PETSc.Sys.Print('') PETSc.Sys.Print(f">>> Converged chunks: {nconverged}.") PETSc.Sys.Print(f">>> Converged time: {converged_time} hours.") @@ -333,6 +327,17 @@ def update_chunk_halos(uhalo): for i in range(args.nchunks): chunk_indexes[i] = (chunk_indexes[i] + 1) % args.nchunks +global_comm.Barrier() +sleep(sleep_time) +if earliest() and chunk_aaofunc.layout.is_local(-1): + from utils.serial import SerialMiniApp + serialapp = SerialMiniApp(dt, theta, q0, form_mass, form_function, block_parameters) + serialapp.solve(nt=nconverged*chunk_length) + PETSc.Sys.Print(f"serial error: {fd.errornorm(serialapp.w0, chunk_aaofunc[-1])}", comm=chunk_ensemble.comm) + PETSc.Sys.Print('') +global_comm.Barrier() +sleep(sleep_time) + nsweeps = j niterations = nsweeps*args.nsmooth From 00704a4a89c231b54360af327fd6ff36ce7c36f1 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Tue, 19 Mar 2024 15:19:39 +0000 Subject: [PATCH 18/18] chunk snes outputs to seperate files --- case_studies/gauss_seidel/advection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/case_studies/gauss_seidel/advection.py b/case_studies/gauss_seidel/advection.py index c9cc4841..1b0ecf3d 100644 --- a/case_studies/gauss_seidel/advection.py +++ b/case_studies/gauss_seidel/advection.py @@ -138,8 +138,8 @@ def form_function(q, phi, t): patol = sqrt(chunk_length)*atol sparameters_diag = { 'snes': { - 'monitor': None, - 'converged_reason': None, + 'monitor': f":snes_monitor_chunk{chunk_id}.log", + 'converged_reason': f":snes_converged_reason_chunk{chunk_id}.log", 'atol': patol, 'rtol': 1e-10, 'stol': 1e-12,