Hybridisation using firedrake.SCPC

case_studies/shallow_water/blockpc/lswe_cpx.py

@@ -0,0 +1,168 @@
+import firedrake as fd
+from firedrake.petsc import PETSc
+from asQ.complex_proxy import vector as cpx
+
+import utils.shallow_water.gravity_bumps as lcase
+from utils import shallow_water as swe
+from utils.planets import earth
+from utils import units
+from utils.hybridisation import HybridisedSCPC  # noqa: F401
+
+import numpy as np
+from scipy.fft import fft, fftfreq
+
+PETSc.Sys.popErrorHandler()
+Print = PETSc.Sys.Print
+
+# get command arguments
+import argparse
+parser = argparse.ArgumentParser(
+    description='Test preconditioners for the complex linear shallow water equations.',
+    formatter_class=argparse.ArgumentDefaultsHelpFormatter
+)
+
+parser.add_argument('--dt', type=float, default=1.0, help='Timestep in hours (used to calculate the circulant eigenvalues).')
+parser.add_argument('--nt', type=int, default=16, help='Number of timesteps (used to calculate the circulant eigenvalues).')
+parser.add_argument('--theta', type=float, default=0.5, help='Parameter for implicit theta method. 0.5 for trapezium rule, 1 for backwards Euler (used to calculate the circulant eigenvalues).')
+parser.add_argument('--alpha', type=float, default=1e-3, help='Circulant parameter (used to calculate the circulant eigenvalues).')
+parser.add_argument('--eigenvalue', type=int, nargs='+', default=-1, help='Index of the circulant eigenvalues to use for the complex coefficients. -1 for all.')
+parser.add_argument('--seed', type=int, default=12345, help='Seed for the random right hand side.')
+parser.add_argument('--ref_level', type=int, default=3, help='Icosahedral sphere mesh refinement level with mesh hierarchy. 0 for no mesh hierarchy.')
+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:
+    Print(args)
+
+# eigenvalues
+nt, theta, alpha = args.nt, args.theta, args.alpha
+dt = args.dt*units.hour
+
+eigenvalues = args.eigenvalue if type(args.eigenvalue) is list else [args.eigenvalue]
+
+if len(args.eigenvalue) > 1:
+    eigenvalues = tuple(eigenvalues)
+elif args.eigenvalue[0] < 0:
+    eigenvalues = tuple(i for i in range(nt//2+1))
+
+gamma = args.alpha**(np.arange(nt)/nt)
+C1 = np.zeros(nt)
+C2 = np.zeros(nt)
+
+C1[:2] = [1/dt, -1/dt]
+C2[:2] = [theta, 1-theta]
+
+D1 = np.sqrt(nt)*fft(gamma*C1)
+D2 = np.sqrt(nt)*fft(gamma*C2)
+freqs = fftfreq(nt, dt)
+
+d1 = D1[eigenvalues[0]]
+d2 = D2[eigenvalues[0]]
+
+d1c = cpx.ComplexConstant(d1)
+d2c = cpx.ComplexConstant(d2)
+
+dhat = (d1/d2) / (1/(theta*dt))
+
+mesh = swe.create_mg_globe_mesh(ref_level=args.ref_level,
+                                coords_degree=1)
+x = fd.SpatialCoordinate(mesh)
+
+# case parameters
+g = earth.Gravity
+f = lcase.coriolis_expression(*x)
+H = lcase.H
+
+# function spaces
+V = swe.default_function_space(mesh)
+W = cpx.FunctionSpace(V)
+
+
+# shallow water equation forms
+def form_mass(u, h, v, q):
+    return swe.linear.form_mass(mesh, u, h, v, q)
+
+
+def form_function(u, h, v, q, t=None):
+    return swe.linear.form_function(mesh, g, H, f,
+                                    u, h, v, q, t)
+
+
+# random rhs
+L = fd.Cofunction(W.dual())
+
+# PETSc solver parameters
+lu_params = {
+    'ksp_type': 'preonly',
+    'pc_type': 'lu',
+    'pc_factor_mat_solver_type': 'mumps'
+}
+
+condensed_params = lu_params
+
+scpc_params = {
+    "ksp_type": 'preonly',
+    "mat_type": "matfree",
+    "pc_type": "python",
+    "pc_python_type": f"{__name__}.HybridisedSCPC",
+    "hybridscpc_condensed_field": lu_params
+}
+
+rtol = 1e-3
+params = {
+    'ksp': {
+        'monitor': None,
+        'converged_rate': None,
+        'rtol': rtol,
+    },
+}
+params.update(scpc_params)
+
+# trace component should have zero rhs
+np.random.seed(args.seed)
+L.assign(0)
+for dat in L.dat:
+    dat.data[:] = np.random.rand(*(dat.data.shape))
+
+# block forms
+M = cpx.BilinearForm(W, d1c, form_mass)
+K = cpx.BilinearForm(W, d2c, form_function)
+
+A = M + K
+
+appctx = {
+    'cpx': cpx,
+    'uref': fd.Function(W),
+    'bcs': None,
+    'tref': None,
+    'form_mass': form_mass,
+    'form_function': form_function,
+    'd1': d1c,
+    'd2': d2c,
+}
+
+wout = fd.Function(W).assign(0)
+problem = fd.LinearVariationalProblem(A, L, wout)
+solver = fd.LinearVariationalSolver(problem, appctx=appctx,
+                                    solver_parameters=params)
+
+for j in eigenvalues:
+    d1, d2 = D1[j], D2[j]
+
+    dhat = (d1/d2) / (1/(theta*dt))
+    Print(f"dhat = {np.round(dhat, 4)}")
+
+    d1c.real.assign(d1.real)
+    d1c.imag.assign(d1.imag)
+
+    d2c.real.assign(d2.real)
+    d2c.imag.assign(d2.imag)
+
+    L.assign(0)
+    for dat in L.dat:
+        dat.data[:] = np.random.rand(*(dat.data.shape))
+    wout.assign(0)
+
+    solver.solve()

case_studies/shallow_water/blockpc/lswe_rblock.py

@@ -0,0 +1,119 @@
+import firedrake as fd
+from firedrake.petsc import PETSc
+
+import utils.shallow_water.gravity_bumps as lcase
+from utils import shallow_water as swe
+from utils.planets import earth
+from utils import units
+from utils.hybridisation import HybridisedSCPC  # noqa: F401
+
+import numpy as np
+
+PETSc.Sys.popErrorHandler()
+Print = PETSc.Sys.Print
+
+
+# get command arguments
+import argparse
+parser = argparse.ArgumentParser(
+    description='Test preconditioners for the real-valued linear shallow water equations.',
+    formatter_class=argparse.ArgumentDefaultsHelpFormatter
+)
+
+parser.add_argument('--dt', type=float, default=1.0, help='Timestep in hours (used to calculate the circulant eigenvalues).')
+parser.add_argument('--theta', type=float, default=0.5, help='Parameter for implicit theta method. 0.5 for trapezium rule, 1 for backwards Euler (used to calculate the circulant eigenvalues).')
+parser.add_argument('--seed', type=int, default=12345, help='Seed for the random right hand side.')
+parser.add_argument('--ref_level', type=int, default=3, help='Icosahedral sphere mesh refinement level with mesh hierarchy. 0 for no mesh hierarchy.')
+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:
+    Print(args)
+
+# eigenvalues
+theta = args.theta
+dt = args.dt*units.hour
+
+d1 = fd.Constant(1/dt)
+d2 = fd.Constant(theta)
+
+mesh = swe.create_mg_globe_mesh(ref_level=args.ref_level,
+                                coords_degree=1)
+x = fd.SpatialCoordinate(mesh)
+
+# case parameters
+g = earth.Gravity
+f = lcase.coriolis_expression(*x)
+H = lcase.H
+
+
+# shallow water equation forms
+def form_mass(u, h, v, q):
+    return swe.linear.form_mass(mesh, u, h, v, q)
+
+
+def form_function(u, h, v, q, t=None):
+    return swe.linear.form_function(mesh, g, H, f,
+                                    u, h, v, q, t)
+
+
+V = swe.default_function_space(mesh, degree=0)
+
+# random rhs
+L = fd.Cofunction(V.dual())
+
+# PETSc solver parameters
+lu_params = {
+    'ksp_type': 'preonly',
+    'pc_type': 'lu',
+    'pc_factor_mat_solver_type': 'petsc'
+}
+
+hybrid_params = {
+    'mat_type': 'matfree',
+    'ksp_type': 'preonly',
+    'pc_type': 'python',
+    'pc_python_type': f'{__name__}.HybridisedSCPC',
+    'hybridscpc_condensed_field': lu_params
+}
+
+sparams = {
+    'ksp': {
+        'monitor': None,
+        'converged_rate': None,
+    },
+}
+sparams.update(hybrid_params)
+
+# trace component should have zero rhs
+np.random.seed(args.seed)
+L.assign(0)
+for ldat in L.dat:
+    ldat.data[:] = np.random.rand(*(ldat.data.shape))
+
+# block forms
+us = fd.TrialFunctions(V)
+vs = fd.TestFunctions(V)
+
+M = form_mass(*us, *vs)
+K = form_function(*us, *vs)
+
+A = d1*M + d2*K
+
+appctx = {
+    'uref': fd.Function(V),
+    'bcs': None,
+    'tref': None,
+    'form_mass': form_mass,
+    'form_function': form_function,
+    'dt': dt,
+    'theta': theta,
+}
+
+wout = fd.Function(V)
+problem = fd.LinearVariationalProblem(A, L, wout)
+solver = fd.LinearVariationalSolver(problem, appctx=appctx,
+                                    solver_parameters=sparams)
+solver.solve()

case_studies/shallow_water/linear_gravity_bumps.py

@@ -47,84 +47,46 @@
 
 # parameters for the implicit diagonal solve in step-(b)
 
-patch_parameters = {
-    'pc_patch': {
-        'save_operators': True,
-        'partition_of_unity': True,
-        'sub_mat_type': 'seqdense',
-        'construct_dim': 0,
-        'construct_type': 'vanka',
-        'local_type': 'additive',
-        'precompute_element_tensors': True,
-        'symmetrise_sweep': False
-    },
-    'sub': {
+from utils.hybridisation import HybridisedSCPC  # noqa: F401
+hybridscpc_parameters = {
+    "ksp_type": 'preonly',
+    "mat_type": "matfree",
+    "pc_type": "python",
+    "pc_python_type": f"{__name__}.HybridisedSCPC",
+    "hybridscpc_condensed_field": {
         'ksp_type': 'preonly',
         'pc_type': 'lu',
-        'pc_factor_shift_type': 'nonzero',
-    }
-}
-
-from utils.mg import ManifoldTransferManager  # noqa: F401
-mg_parameters = {
-    'transfer_manager': f'{__name__}.ManifoldTransferManager',
-    'levels': {
-        'ksp_type': 'gmres',
-        'ksp_max_it': 5,
-        'pc_type': 'python',
-        'pc_python_type': 'firedrake.PatchPC',
-        'patch': patch_parameters
-    },
-    'coarse': {
-        'pc_type': 'python',
-        'pc_python_type': 'firedrake.AssembledPC',
-        'assembled_pc_type': 'lu',
-        'assembled_pc_factor_mat_solver_type': 'mumps'
+        'pc_factor_mat_solver_type': 'mumps',
+        'snes': {  # reuse factorisation
+            'lag_jacobian': -2,
+            'lag_jacobian_persists': None,
+            'lag_preconditioner': -2,
+            'lag_preconditioner_persists': None,
+        }
     }
 }
 
-sparameters = {
-    'mat_type': 'matfree',
-    'ksp_type': 'fgmres',
-    'ksp': {
-        'atol': 1e-5,
-        'rtol': 1e-5,
-        'max_it': 60
-    },
-    'pc_type': 'mg',
-    'pc_mg_cycle_type': 'w',
-    'pc_mg_type': 'multiplicative',
-    'mg': mg_parameters
-}
-
 rtol = 1e-10
 atol = 1e0
 sparameters_diag = {
     'snes_type': 'ksponly',
-    'snes': {
-        'linesearch_type': 'basic',
-        'monitor': None,
-        'converged_reason': None,
-        'rtol': rtol,
-        'atol': atol,
-    },
     'mat_type': 'matfree',
     'ksp_type': 'fgmres',
     'ksp': {
         'monitor': None,
-        'converged_reason': None,
+        'converged_rate': None,
         'rtol': rtol,
         'atol': atol,
     },
     'pc_type': 'python',
     'pc_python_type': 'asQ.CirculantPC',
-    'diagfft_alpha': args.alpha,
-    'diagfft_state': 'linear',
+    'circulant_alpha': args.alpha,
+    'circulant_state': 'linear',
     'aaos_jacobian_state': 'linear',
 }
 
 for i in range(window_length):
-    sparameters_diag['diagfft_block_'+str(i)+'_'] = sparameters
+    sparameters_diag['circulant_block_'+str(i)+'_'] = hybridscpc_parameters
 
 create_mesh = partial(
     swe.create_mg_globe_mesh,