From 7fe2608f9cb5f4561428430a726020ae568d3f13 Mon Sep 17 00:00:00 2001 From: Ahsan Ali Date: Wed, 13 Aug 2025 17:28:55 -0500 Subject: [PATCH 01/35] added the test file for KronPC --- irksome/KronPC_test.py | 313 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 313 insertions(+) create mode 100644 irksome/KronPC_test.py diff --git a/irksome/KronPC_test.py b/irksome/KronPC_test.py new file mode 100644 index 00000000..a9120254 --- /dev/null +++ b/irksome/KronPC_test.py @@ -0,0 +1,313 @@ +# suggested setting OMP_NUM_THREADS=1 to improve performance +import os +os.environ["OMP_NUM_THREADS"] = "1" + +# ----------------------------- +# Imports +# ----------------------------- +from firedrake.preconditioners.base import PCBase +from firedrake.petsc import PETSc +from firedrake import * +from firedrake.__future__ import interpolate +from irksome import MeshConstant, Dt, RadauIIA, TimeStepper +from ufl import sin, cos, pi, as_vector, exp, div, dot +import numpy as np + + +__all__ = ("KronPC",) + +# ----------------------------- +# KronPC definition +# ----------------------------- +class KronPC(PCBase): + r""" + Preconditioner applying: y = coef * (A^{-p} \otimes K^{-1}) x + - A is the Butcher matrix (RadauIIA(s)), p >= 0 (p=0 => I) + - K is assembled on the single-stage space via kron_operator \in {mass, stiffness} + - K^{-1} is approximated by a PETSc PC with prefix 'kron_sub_*' + + Options (with this PC's prefix, e.g. -fieldsplit_1_kron_*): + -kron_coef : scalar coefficient (default 1.0) + -kron_pow : exponent p for A^{-p} (default 1; p=0 gives identity) + -kron_operator : which K to assemble (default mass) + -kron_mat_type : aux matrix type for K (default aij) + # sub-PC for K^{-1}: + -kron_sub_pc_type <...> + ... (any other -kron_sub_* options forwarded to the sub PC) + """ + + # We rely on the Jacobian P being a python Mat (to get the Firedrake context) + needs_python_pmat = True + + # --- single-stage form K --- + def form(self, trial, test, operator_kind): + if operator_kind == "mass": + a = inner(trial, test) * dx + elif operator_kind == "stiffness": + # tiny mass term for regularization / maybe helpful for AMG solvers + a = inner(grad(trial), grad(test)) * dx + 1e-12 * inner(trial, test) * dx + else: + raise ValueError(f"Unknown kron_operator '{operator_kind}' (use 'mass' or 'stiffness').") + bcs = None + return a, bcs + # --- ButcherTableau implementation (not generalized yet, to be fixed) --- + def stage_mat(self, s, pow_): + if pow_ == 0: + return np.eye(s) # A^0 = I + A = RadauIIA(s).A + Ainv = np.linalg.inv(A) + if pow_ == 1: + return Ainv + return np.linalg.matrix_power(Ainv, pow_) + + @property + def num_stages(self): + return self._s + + # --- PETSc PC interface --- + def initialize(self, pc): + if pc.getType() != "python": + raise ValueError("Expecting PC type 'python' for KronPC.") + + # Prefix for options + self._prefix = (pc.getOptionsPrefix() or "") + "kron_" + opts = PETSc.Options() + + self.coef = opts.getReal(self._prefix + "coef", 1.0) + pow_ = opts.getInt(self._prefix + "pow", 1) + if pow_ < 0: + raise ValueError("kron_pow must be a nonnegative integer.") + self._pow = pow_ + + operator_kind = opts.getString(self._prefix + "operator", "mass") + mat_type = opts.getString(self._prefix + "mat_type", "aij") + + # Firedrake context & staged space + _, P = pc.getOperators() + context = P.getPythonContext() + Vbig = context.a.arguments()[0].function_space() + + # Work containers + self.work_in = Cofunction(Vbig.dual(), name="kron_work_in") # incoming algebraic vec + self.work_mid = Function(Vbig, name="kron_work_mid") # per-stage K^{-1} x + self.work_out = Function(Vbig, name="kron_work_out") # after stage mixing + + # Single-stage space (assume uniform across stages ? Need fix !!) + Vstage = Vbig.sub(0) + trial = TrialFunction(Vstage) + test = TestFunction(Vstage) + a, bcs = self.form(trial, test, operator_kind) + + # Assemble K + fc_params = getattr(context, "fc_params", None) + K = assemble(a, bcs=bcs, mat_type=mat_type, form_compiler_parameters=fc_params) + self.K = K # keep reference + + # --- Sub-PC for K^{-1} (by default using LU for now) --- + sub_pc = PETSc.PC().create(comm=pc.comm) + sub_pc.incrementTabLevel(1, parent=pc) + sub_pc.setOptionsPrefix(self._prefix + "sub_") + sub_pc.setOperators(K.M.handle) # use K for both A and P + sub_pc.setFromOptions() # allow overrides + # Force LU unless the user overrides via options: + if sub_pc.getType() is None: + sub_pc.setType("lu") + sub_pc.setUp() + self.sub_pc = sub_pc + + # Stages and dense coupling + self._s = Vbig.num_sub_spaces() + self.L = self.stage_mat(self._s, self._pow) + + def update(self, pc): + pass + + def apply(self, pc, x, y): + r""" + y = coef * (A^{-p} \otimes K^{-1}) x + """ + s = self.num_stages + + # 1) Copy PETSc Vec x -> Cofunction(Vbig.dual()) + with self.work_in.dat.vec_wo as vin: + x.copy(vin) + + # 2) Per-stage action: mid[i] = (sub_pc) in[i] + for i in range(s): + with self.work_in.subfunctions[i].dat.vec_ro as xin_i, \ + self.work_mid.subfunctions[i].dat.vec_wo as mid_i: + mid_i.set(0.0) + self.sub_pc.apply(xin_i, mid_i) + + # 3) Dense stage coupling: out[j] = coef * sum_i L[j,i] * mid[i] + for j in range(s): + self.work_out.subfunctions[j].assign(0.0) + + for j in range(s): + row = self.L[j, :] + with self.work_out.subfunctions[j].dat.vec_wo as yj: + for i in range(s): + lij = row[i] + if lij == 0.0: + continue + with self.work_mid.subfunctions[i].dat.vec_ro as ui: + yj.axpy(self.coef * lij, ui) + + # 4) Copy back to PETSc Vec y + with self.work_out.dat.vec_ro as vout: + vout.copy(y) + + def applyTranspose(self, pc, x, y): + # Leave me for now!! + self.apply(pc, x, y) + + def view(self, pc, viewer=None): + if viewer is None: + viewer = PETSc.Viewer.STDOUT(pc.comm) + viewer.printfASCII("KronPC:\n") + viewer.printfASCII(f" stages : {self.num_stages}\n") + viewer.printfASCII(f" coef : {self.coef}\n") + viewer.printfASCII(f" power (p) : {self._pow} [A^{-p}]\n") + viewer.printfASCII(" sub-PC for K^{-1} (kron_sub_*) options:\n") + if hasattr(self, "sub_pc") and self.sub_pc is not None: + self.sub_pc.view(viewer) + + def destroy(self, pc): + if hasattr(self, "sub_pc") and self.sub_pc is not None: + self.sub_pc.destroy() + self.sub_pc = None + + +# ----------------------------- +# Problem setup and solve +# ----------------------------- +if __name__ == "__main__": + # Problem parameters + mesh_size = 4 # size of mesh + nu = 1.0 / 50.0 # Viscosity + T = 2.0 # Final time + s = 2 # Radau IIA stages + dt_value = T / 10.0 # Time step + gamma = Constant(50.0) # Augmentation parameter for testing + + # Mesh and spaces + mesh = UnitSquareMesh(2 ** mesh_size, 2 ** mesh_size, quadrilateral=True) + V = VectorFunctionSpace(mesh, "CG", 2) + Q = FunctionSpace(mesh, "CG", 1) + W = V * Q + + # Time variables + t = Constant(0.0) + MC = MeshConstant(mesh) + dt = MC.Constant(dt_value) + + # Manufactured exact solution + x = SpatialCoordinate(mesh) + u_exact = 0.5 * exp(T - t) * as_vector([ + sin(pi * x[0]) * cos(pi * x[1]), + -cos(pi * x[0]) * sin(pi * x[1]) + ]) + p_exact = Constant(0.0) + + # Forcing term: Forcing term f = u_t - \nu \Delta u - u\cdot \nabla u (grad(p) = 0 because p is constant) + u_t = Dt(u_exact) + Delta_u = div(grad(u_exact)) + conv_u = dot(grad(u_exact), u_exact) + f_expr = u_t - nu * Delta_u - conv_u + + # Initial condition + w = Function(W) + (u, p) = split(w) + (u0, p0) = w.subfunctions + u0.interpolate(u_exact) + p0.assign(0.0) + + # Boundary conditions + bc = DirichletBC(W.sub(0), u_exact, "on_boundary") + + # Variational form (with augmented Lagrangian div-div term) + v, q = TestFunctions(W) + F = (inner(Dt(u), v) * dx + + nu * inner(grad(u), grad(v)) * dx + + inner(dot(grad(u), u), v) * dx + - inner(p, div(v)) * dx + + inner(div(u), q) * dx + - inner(f_expr, v) * dx + + gamma * inner(div(u), div(v)) * dx) # Augmented Lagrangian Approach + + # Pressure nullspace (constant pressure) + nsp = [(1, VectorSpaceBasis(constant=True))] + + # Group velocity/pressure stage indices + velocity_fields = ",".join(str(2*i) for i in range(s)) + pressure_fields = ",".join(str(2*i+1) for i in range(s)) + + # Solver parameters + parameters = { + # make the global operator mat-free so P is a python Mat + "mat_type": "matfree", + "pmat_type": "matfree", + + # outer solve + "ksp_type": "fgmres", + "ksp_rtol": 1e-8, + "ksp_converged_reason": None, + + # ---- Top-level: Schur upper ---- + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "schur", + "pc_fieldsplit_schur_factorization_type": "upper", + "pc_fieldsplit_schur_precondition": "selfp", + + # Map fields + "pc_fieldsplit_0_fields": velocity_fields, # (1,1) velocity stages + "pc_fieldsplit_1_fields": pressure_fields, # (2,2) pressure Schur + + # ---- (1,1) velocity block ---- + "fieldsplit_0": { + "ksp_type": "preonly", + "ksp_rtol": 1e-10, + "pc_type": "python", + "pc_python_type": "firedrake.AssembledPC", + "assembled_ksp_type":"preonly", + "assembled_pc_type": "lu", + }, + + # ---- (2,2) Schur block ---- + "fieldsplit_1": { + "ksp_type": "preonly", + "ksp_rtol": 1e-10, + "ksp_converged_reason": None, + "pc_type": "python", + "pc_python_type": "__main__.KronPC", # adjust if KronPC is imported + # KronPC options: + "kron_operator": "mass", # K = mass matrix on single-stage space + "kron_pow": 0, # A^0 = I => I \otimes (sub_pc) (MassInv-like) + # nested sub-PC on K: + "kron_sub_pc_type": "lu", + }, + } + + # Time stepper + stepper = TimeStepper( + F, + RadauIIA(s), + t, + dt, + w, + bcs=bc, + stage_type="deriv", + solver_parameters=parameters, + nullspace=nsp + ) + + # Advance in time + while float(t) < T - 1e-10: + stepper.advance() + t.assign(float(t) + float(dt)) + + # Report stats + steps, nonlinear_its, linear_its = stepper.solver_stats() + print("Total number of timesteps was %d" % (steps)) + print("Average number of nonlinear iterations per timestep was %.2f" % (nonlinear_its/steps)) + print("Average number of linear iterations per timestep was %.2f" % (linear_its/steps)) From a2c80eb73626af8c32eec13c486c1345b07a8201 Mon Sep 17 00:00:00 2001 From: "Robert C. Kirby" Date: Wed, 13 Aug 2025 22:58:05 -0500 Subject: [PATCH 02/35] working now --- irksome/KronPC_test.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/irksome/KronPC_test.py b/irksome/KronPC_test.py index a9120254..1a32e998 100644 --- a/irksome/KronPC_test.py +++ b/irksome/KronPC_test.py @@ -106,7 +106,7 @@ def initialize(self, pc): # --- Sub-PC for K^{-1} (by default using LU for now) --- sub_pc = PETSc.PC().create(comm=pc.comm) sub_pc.incrementTabLevel(1, parent=pc) - sub_pc.setOptionsPrefix(self._prefix + "sub_") + # sub_pc.setOptionsPrefix(self._prefix + "sub_") sub_pc.setOperators(K.M.handle) # use K for both A and P sub_pc.setFromOptions() # allow overrides # Force LU unless the user overrides via options: @@ -246,7 +246,6 @@ def destroy(self, pc): parameters = { # make the global operator mat-free so P is a python Mat "mat_type": "matfree", - "pmat_type": "matfree", # outer solve "ksp_type": "fgmres", @@ -257,7 +256,6 @@ def destroy(self, pc): "pc_type": "fieldsplit", "pc_fieldsplit_type": "schur", "pc_fieldsplit_schur_factorization_type": "upper", - "pc_fieldsplit_schur_precondition": "selfp", # Map fields "pc_fieldsplit_0_fields": velocity_fields, # (1,1) velocity stages @@ -279,12 +277,13 @@ def destroy(self, pc): "ksp_rtol": 1e-10, "ksp_converged_reason": None, "pc_type": "python", - "pc_python_type": "__main__.KronPC", # adjust if KronPC is imported - # KronPC options: - "kron_operator": "mass", # K = mass matrix on single-stage space - "kron_pow": 0, # A^0 = I => I \otimes (sub_pc) (MassInv-like) - # nested sub-PC on K: - "kron_sub_pc_type": "lu", + # "pc_python_type": "firedrake.MassInvPC", + "pc_python_type": "__main__.KronPC", + "kron": { + "operator": "mass", + "pow": 0, + "pc_type": "lu" + } }, } @@ -303,6 +302,7 @@ def destroy(self, pc): # Advance in time while float(t) < T - 1e-10: + # for _ in range(1): stepper.advance() t.assign(float(t) + float(dt)) From c4a9fa840b506fb8125e2064538ce3047f4982ba Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Fri, 15 Aug 2025 10:55:16 -0500 Subject: [PATCH 03/35] Test got seperated from KronPC --- irksome/KronPC_test.py | 407 +++++++++++------------------------------ 1 file changed, 104 insertions(+), 303 deletions(-) diff --git a/irksome/KronPC_test.py b/irksome/KronPC_test.py index 1a32e998..30573b67 100644 --- a/irksome/KronPC_test.py +++ b/irksome/KronPC_test.py @@ -2,312 +2,113 @@ import os os.environ["OMP_NUM_THREADS"] = "1" -# ----------------------------- -# Imports -# ----------------------------- -from firedrake.preconditioners.base import PCBase -from firedrake.petsc import PETSc from firedrake import * from firedrake.__future__ import interpolate from irksome import MeshConstant, Dt, RadauIIA, TimeStepper from ufl import sin, cos, pi, as_vector, exp, div, dot import numpy as np - -__all__ = ("KronPC",) - -# ----------------------------- -# KronPC definition -# ----------------------------- -class KronPC(PCBase): - r""" - Preconditioner applying: y = coef * (A^{-p} \otimes K^{-1}) x - - A is the Butcher matrix (RadauIIA(s)), p >= 0 (p=0 => I) - - K is assembled on the single-stage space via kron_operator \in {mass, stiffness} - - K^{-1} is approximated by a PETSc PC with prefix 'kron_sub_*' - - Options (with this PC's prefix, e.g. -fieldsplit_1_kron_*): - -kron_coef : scalar coefficient (default 1.0) - -kron_pow : exponent p for A^{-p} (default 1; p=0 gives identity) - -kron_operator : which K to assemble (default mass) - -kron_mat_type : aux matrix type for K (default aij) - # sub-PC for K^{-1}: - -kron_sub_pc_type <...> - ... (any other -kron_sub_* options forwarded to the sub PC) - """ - - # We rely on the Jacobian P being a python Mat (to get the Firedrake context) - needs_python_pmat = True - - # --- single-stage form K --- - def form(self, trial, test, operator_kind): - if operator_kind == "mass": - a = inner(trial, test) * dx - elif operator_kind == "stiffness": - # tiny mass term for regularization / maybe helpful for AMG solvers - a = inner(grad(trial), grad(test)) * dx + 1e-12 * inner(trial, test) * dx - else: - raise ValueError(f"Unknown kron_operator '{operator_kind}' (use 'mass' or 'stiffness').") - bcs = None - return a, bcs - # --- ButcherTableau implementation (not generalized yet, to be fixed) --- - def stage_mat(self, s, pow_): - if pow_ == 0: - return np.eye(s) # A^0 = I - A = RadauIIA(s).A - Ainv = np.linalg.inv(A) - if pow_ == 1: - return Ainv - return np.linalg.matrix_power(Ainv, pow_) - - @property - def num_stages(self): - return self._s - - # --- PETSc PC interface --- - def initialize(self, pc): - if pc.getType() != "python": - raise ValueError("Expecting PC type 'python' for KronPC.") - - # Prefix for options - self._prefix = (pc.getOptionsPrefix() or "") + "kron_" - opts = PETSc.Options() - - self.coef = opts.getReal(self._prefix + "coef", 1.0) - pow_ = opts.getInt(self._prefix + "pow", 1) - if pow_ < 0: - raise ValueError("kron_pow must be a nonnegative integer.") - self._pow = pow_ - - operator_kind = opts.getString(self._prefix + "operator", "mass") - mat_type = opts.getString(self._prefix + "mat_type", "aij") - - # Firedrake context & staged space - _, P = pc.getOperators() - context = P.getPythonContext() - Vbig = context.a.arguments()[0].function_space() - - # Work containers - self.work_in = Cofunction(Vbig.dual(), name="kron_work_in") # incoming algebraic vec - self.work_mid = Function(Vbig, name="kron_work_mid") # per-stage K^{-1} x - self.work_out = Function(Vbig, name="kron_work_out") # after stage mixing - - # Single-stage space (assume uniform across stages ? Need fix !!) - Vstage = Vbig.sub(0) - trial = TrialFunction(Vstage) - test = TestFunction(Vstage) - a, bcs = self.form(trial, test, operator_kind) - - # Assemble K - fc_params = getattr(context, "fc_params", None) - K = assemble(a, bcs=bcs, mat_type=mat_type, form_compiler_parameters=fc_params) - self.K = K # keep reference - - # --- Sub-PC for K^{-1} (by default using LU for now) --- - sub_pc = PETSc.PC().create(comm=pc.comm) - sub_pc.incrementTabLevel(1, parent=pc) - # sub_pc.setOptionsPrefix(self._prefix + "sub_") - sub_pc.setOperators(K.M.handle) # use K for both A and P - sub_pc.setFromOptions() # allow overrides - # Force LU unless the user overrides via options: - if sub_pc.getType() is None: - sub_pc.setType("lu") - sub_pc.setUp() - self.sub_pc = sub_pc - - # Stages and dense coupling - self._s = Vbig.num_sub_spaces() - self.L = self.stage_mat(self._s, self._pow) - - def update(self, pc): - pass - - def apply(self, pc, x, y): - r""" - y = coef * (A^{-p} \otimes K^{-1}) x - """ - s = self.num_stages - - # 1) Copy PETSc Vec x -> Cofunction(Vbig.dual()) - with self.work_in.dat.vec_wo as vin: - x.copy(vin) - - # 2) Per-stage action: mid[i] = (sub_pc) in[i] - for i in range(s): - with self.work_in.subfunctions[i].dat.vec_ro as xin_i, \ - self.work_mid.subfunctions[i].dat.vec_wo as mid_i: - mid_i.set(0.0) - self.sub_pc.apply(xin_i, mid_i) - - # 3) Dense stage coupling: out[j] = coef * sum_i L[j,i] * mid[i] - for j in range(s): - self.work_out.subfunctions[j].assign(0.0) - - for j in range(s): - row = self.L[j, :] - with self.work_out.subfunctions[j].dat.vec_wo as yj: - for i in range(s): - lij = row[i] - if lij == 0.0: - continue - with self.work_mid.subfunctions[i].dat.vec_ro as ui: - yj.axpy(self.coef * lij, ui) - - # 4) Copy back to PETSc Vec y - with self.work_out.dat.vec_ro as vout: - vout.copy(y) - - def applyTranspose(self, pc, x, y): - # Leave me for now!! - self.apply(pc, x, y) - - def view(self, pc, viewer=None): - if viewer is None: - viewer = PETSc.Viewer.STDOUT(pc.comm) - viewer.printfASCII("KronPC:\n") - viewer.printfASCII(f" stages : {self.num_stages}\n") - viewer.printfASCII(f" coef : {self.coef}\n") - viewer.printfASCII(f" power (p) : {self._pow} [A^{-p}]\n") - viewer.printfASCII(" sub-PC for K^{-1} (kron_sub_*) options:\n") - if hasattr(self, "sub_pc") and self.sub_pc is not None: - self.sub_pc.view(viewer) - - def destroy(self, pc): - if hasattr(self, "sub_pc") and self.sub_pc is not None: - self.sub_pc.destroy() - self.sub_pc = None - - -# ----------------------------- -# Problem setup and solve -# ----------------------------- -if __name__ == "__main__": - # Problem parameters - mesh_size = 4 # size of mesh - nu = 1.0 / 50.0 # Viscosity - T = 2.0 # Final time - s = 2 # Radau IIA stages - dt_value = T / 10.0 # Time step - gamma = Constant(50.0) # Augmentation parameter for testing - - # Mesh and spaces - mesh = UnitSquareMesh(2 ** mesh_size, 2 ** mesh_size, quadrilateral=True) - V = VectorFunctionSpace(mesh, "CG", 2) - Q = FunctionSpace(mesh, "CG", 1) - W = V * Q - - # Time variables - t = Constant(0.0) - MC = MeshConstant(mesh) - dt = MC.Constant(dt_value) - - # Manufactured exact solution - x = SpatialCoordinate(mesh) - u_exact = 0.5 * exp(T - t) * as_vector([ - sin(pi * x[0]) * cos(pi * x[1]), - -cos(pi * x[0]) * sin(pi * x[1]) - ]) - p_exact = Constant(0.0) - - # Forcing term: Forcing term f = u_t - \nu \Delta u - u\cdot \nabla u (grad(p) = 0 because p is constant) - u_t = Dt(u_exact) - Delta_u = div(grad(u_exact)) - conv_u = dot(grad(u_exact), u_exact) - f_expr = u_t - nu * Delta_u - conv_u - - # Initial condition - w = Function(W) - (u, p) = split(w) - (u0, p0) = w.subfunctions - u0.interpolate(u_exact) - p0.assign(0.0) - - # Boundary conditions - bc = DirichletBC(W.sub(0), u_exact, "on_boundary") - - # Variational form (with augmented Lagrangian div-div term) - v, q = TestFunctions(W) - F = (inner(Dt(u), v) * dx - + nu * inner(grad(u), grad(v)) * dx - + inner(dot(grad(u), u), v) * dx - - inner(p, div(v)) * dx - + inner(div(u), q) * dx - - inner(f_expr, v) * dx - + gamma * inner(div(u), div(v)) * dx) # Augmented Lagrangian Approach - - # Pressure nullspace (constant pressure) - nsp = [(1, VectorSpaceBasis(constant=True))] - - # Group velocity/pressure stage indices - velocity_fields = ",".join(str(2*i) for i in range(s)) - pressure_fields = ",".join(str(2*i+1) for i in range(s)) - - # Solver parameters - parameters = { - # make the global operator mat-free so P is a python Mat - "mat_type": "matfree", - - # outer solve - "ksp_type": "fgmres", - "ksp_rtol": 1e-8, - "ksp_converged_reason": None, - - # ---- Top-level: Schur upper ---- - "pc_type": "fieldsplit", - "pc_fieldsplit_type": "schur", - "pc_fieldsplit_schur_factorization_type": "upper", - - # Map fields - "pc_fieldsplit_0_fields": velocity_fields, # (1,1) velocity stages - "pc_fieldsplit_1_fields": pressure_fields, # (2,2) pressure Schur - - # ---- (1,1) velocity block ---- - "fieldsplit_0": { - "ksp_type": "preonly", - "ksp_rtol": 1e-10, - "pc_type": "python", - "pc_python_type": "firedrake.AssembledPC", - "assembled_ksp_type":"preonly", - "assembled_pc_type": "lu", - }, - - # ---- (2,2) Schur block ---- - "fieldsplit_1": { - "ksp_type": "preonly", - "ksp_rtol": 1e-10, - "ksp_converged_reason": None, - "pc_type": "python", - # "pc_python_type": "firedrake.MassInvPC", - "pc_python_type": "__main__.KronPC", - "kron": { - "operator": "mass", - "pow": 0, - "pc_type": "lu" - } - }, - } - - # Time stepper - stepper = TimeStepper( - F, - RadauIIA(s), - t, - dt, - w, - bcs=bc, - stage_type="deriv", - solver_parameters=parameters, - nullspace=nsp - ) - - # Advance in time - while float(t) < T - 1e-10: - # for _ in range(1): - stepper.advance() - t.assign(float(t) + float(dt)) - - # Report stats - steps, nonlinear_its, linear_its = stepper.solver_stats() - print("Total number of timesteps was %d" % (steps)) - print("Average number of nonlinear iterations per timestep was %.2f" % (nonlinear_its/steps)) - print("Average number of linear iterations per timestep was %.2f" % (linear_its/steps)) +from KronPC import KronPC # import the preconditioner + + +mesh_size = 4 +nu = 1.0 / 50.0 +T = 2.0 +s = 2 +dt_value = T / 10.0 +gamma = Constant(50.0) + +mesh = UnitSquareMesh(2 ** mesh_size, 2 ** mesh_size, quadrilateral=True) +V = VectorFunctionSpace(mesh, "CG", 2) +Q = FunctionSpace(mesh, "CG", 1) +W = V * Q + +t = Constant(0.0) +MC = MeshConstant(mesh) +dt = MC.Constant(dt_value) + +x = SpatialCoordinate(mesh) +u_exact = 0.5 * exp(T - t) * as_vector([ + sin(pi * x[0]) * cos(pi * x[1]), + -cos(pi * x[0]) * sin(pi * x[1]) +]) +p_exact = Constant(0.0) + +u_t = Dt(u_exact) +Delta_u = div(grad(u_exact)) +conv_u = dot(grad(u_exact), u_exact) +f_expr = u_t - nu * Delta_u - conv_u + +w = Function(W) +(u, p) = split(w) +(u0, p0) = w.subfunctions +u0.interpolate(u_exact) +p0.assign(0.0) + +bc = DirichletBC(W.sub(0), u_exact, "on_boundary") + +v, q = TestFunctions(W) +F = (inner(Dt(u), v) * dx + + nu * inner(grad(u), grad(v)) * dx + + inner(dot(grad(u), u), v) * dx + - inner(p, div(v)) * dx + + inner(div(u), q) * dx + - inner(f_expr, v) * dx + + gamma * inner(div(u), div(v)) * dx) + +nsp = [(1, VectorSpaceBasis(constant=True))] + +velocity_fields = ",".join(str(2*i) for i in range(s)) +pressure_fields = ",".join(str(2*i+1) for i in range(s)) + +parameters = { + "mat_type": "matfree", + "ksp_type": "fgmres", + "ksp_rtol": 1e-8, + "ksp_converged_reason": None, + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "schur", + "pc_fieldsplit_schur_factorization_type": "upper", + "pc_fieldsplit_0_fields": velocity_fields, + "pc_fieldsplit_1_fields": pressure_fields, + "fieldsplit_0": { + "ksp_type": "preonly", + "ksp_rtol": 1e-10, + "pc_type": "python", + "pc_python_type": "firedrake.AssembledPC", + "assembled_ksp_type": "preonly", + "assembled_pc_type": "lu", + }, + "fieldsplit_1": { + "ksp_type": "preonly", + "ksp_rtol": 1e-10, + "pc_type": "python", + "pc_python_type": "KronPC.KronPC", + "kron": { + "operator": "mass", + "pow": 1, + "pc_type": "lu" + } + }, +} + +stepper = TimeStepper( + F, + RadauIIA(s), + t, + dt, + w, + bcs=bc, + stage_type="deriv", + solver_parameters=parameters, + nullspace=nsp +) + +while float(t) < T - 1e-10: + stepper.advance() + t.assign(float(t) + float(dt)) + +steps, nonlinear_its, linear_its = stepper.solver_stats() +print("Total number of timesteps was %d" % (steps)) +print("Average number of nonlinear iterations per timestep was %.2f" % (nonlinear_its/steps)) +print("Average number of linear iterations per timestep was %.2f" % (linear_its/steps)) From 63800c7d30bd744641706aba215f2c1ec299a90a Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Fri, 15 Aug 2025 10:56:42 -0500 Subject: [PATCH 04/35] KronPC is now a seperate PC, it will go in pc.py shortly --- irksome/KronPC.py | 141 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 141 insertions(+) create mode 100644 irksome/KronPC.py diff --git a/irksome/KronPC.py b/irksome/KronPC.py new file mode 100644 index 00000000..66f5d3dc --- /dev/null +++ b/irksome/KronPC.py @@ -0,0 +1,141 @@ +# ----------------------------- +# Imports +# ----------------------------- +from firedrake.preconditioners.base import PCBase +from firedrake.petsc import PETSc +from firedrake import * +from irksome import RadauIIA +import numpy as np + +__all__ = ("KronPC",) + +# ----------------------------- +# KronPC definition +# ----------------------------- +class KronPC(PCBase): + r""" + Preconditioner applying: y = coef * (A^{-p} \otimes K^{-1}) x + - A is the Butcher matrix (RadauIIA(s)), p >= 0 (p=0 => I) + - K is assembled on the single-stage space via kron_operator \in {mass, stiffness} + - K^{-1} is approximated by a PETSc PC with prefix 'kron_sub_*' + """ + + needs_python_pmat = True + + def form(self, trial, test, operator_kind): + if operator_kind == "mass": + a = inner(trial, test) * dx + elif operator_kind == "stiffness": + a = inner(grad(trial), grad(test)) * dx + 1e-12 * inner(trial, test) * dx + else: + raise ValueError(f"Unknown kron_operator '{operator_kind}' (use 'mass' or 'stiffness').") + bcs = None + return a, bcs + + def stage_mat(self, s, pow_): + if pow_ == 0: + return np.eye(s) + A = RadauIIA(s).A + Ainv = np.linalg.inv(A) + if pow_ == 1: + return Ainv + return np.linalg.matrix_power(Ainv, pow_) + + @property + def num_stages(self): + return self._s + + def initialize(self, pc): + if pc.getType() != "python": + raise ValueError("Expecting PC type 'python' for KronPC.") + + self._prefix = (pc.getOptionsPrefix() or "") + "kron_" + opts = PETSc.Options() + + self.coef = opts.getReal(self._prefix + "coef", 1.0) + pow_ = opts.getInt(self._prefix + "pow", 1) + if pow_ < 0: + raise ValueError("kron_pow must be a nonnegative integer.") + self._pow = pow_ + + operator_kind = opts.getString(self._prefix + "operator", "mass") + mat_type = opts.getString(self._prefix + "mat_type", "aij") + + _, P = pc.getOperators() + context = P.getPythonContext() + Vbig = context.a.arguments()[0].function_space() + + self.work_in = Cofunction(Vbig.dual(), name="kron_work_in") + self.work_mid = Function(Vbig, name="kron_work_mid") + self.work_out = Function(Vbig, name="kron_work_out") + + Vstage = Vbig.sub(0) + trial = TrialFunction(Vstage) + test = TestFunction(Vstage) + a, bcs = self.form(trial, test, operator_kind) + + fc_params = getattr(context, "fc_params", None) + K = assemble(a, bcs=bcs, mat_type=mat_type, form_compiler_parameters=fc_params) + self.K = K + + sub_pc = PETSc.PC().create(comm=pc.comm) + sub_pc.incrementTabLevel(1, parent=pc) + sub_pc.setOperators(K.M.handle) + sub_pc.setFromOptions() + if sub_pc.getType() is None: + sub_pc.setType("lu") + sub_pc.setUp() + self.sub_pc = sub_pc + + self._s = Vbig.num_sub_spaces() + self.L = self.stage_mat(self._s, self._pow) + + def update(self, pc): + pass + + def apply(self, pc, x, y): + s = self.num_stages + + with self.work_in.dat.vec_wo as vin: + x.copy(vin) + + for i in range(s): + with self.work_in.subfunctions[i].dat.vec_ro as xin_i, \ + self.work_mid.subfunctions[i].dat.vec_wo as mid_i: + mid_i.set(0.0) + self.sub_pc.apply(xin_i, mid_i) + + for j in range(s): + self.work_out.subfunctions[j].assign(0.0) + + for j in range(s): + row = self.L[j, :] + with self.work_out.subfunctions[j].dat.vec_wo as yj: + for i in range(s): + lij = row[i] + if lij == 0.0: + continue + with self.work_mid.subfunctions[i].dat.vec_ro as ui: + yj.axpy(self.coef * lij, ui) + + with self.work_out.dat.vec_ro as vout: + vout.copy(y) + + def applyTranspose(self, pc, x, y): + self.apply(pc, x, y) + + def view(self, pc, viewer=None): + if viewer is None: + viewer = PETSc.Viewer.STDOUT(pc.comm) + viewer.printfASCII("KronPC:\n") + viewer.printfASCII(f" stages : {self.num_stages}\n") + viewer.printfASCII(f" coef : {self.coef}\n") + viewer.printfASCII(f" power (p) : {self._pow} [A^{-p}]\n") + viewer.printfASCII(" sub-PC for K^{-1} (kron_sub_*) options:\n") + if hasattr(self, "sub_pc") and self.sub_pc is not None: + self.sub_pc.view(viewer) + + def destroy(self, pc): + if hasattr(self, "sub_pc") and self.sub_pc is not None: + self.sub_pc.destroy() + self.sub_pc = None From 23e403c5bc7a33fb38867d5628a0fb734d7a4211 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Mon, 18 Aug 2025 09:19:15 -0500 Subject: [PATCH 05/35] KronPC_test.py script is now moved into github.com/rckirby/NSEExamples --- irksome/KronPC_test.py | 114 ----------------------------------------- 1 file changed, 114 deletions(-) delete mode 100644 irksome/KronPC_test.py diff --git a/irksome/KronPC_test.py b/irksome/KronPC_test.py deleted file mode 100644 index 30573b67..00000000 --- a/irksome/KronPC_test.py +++ /dev/null @@ -1,114 +0,0 @@ -# suggested setting OMP_NUM_THREADS=1 to improve performance -import os -os.environ["OMP_NUM_THREADS"] = "1" - -from firedrake import * -from firedrake.__future__ import interpolate -from irksome import MeshConstant, Dt, RadauIIA, TimeStepper -from ufl import sin, cos, pi, as_vector, exp, div, dot -import numpy as np - -from KronPC import KronPC # import the preconditioner - - -mesh_size = 4 -nu = 1.0 / 50.0 -T = 2.0 -s = 2 -dt_value = T / 10.0 -gamma = Constant(50.0) - -mesh = UnitSquareMesh(2 ** mesh_size, 2 ** mesh_size, quadrilateral=True) -V = VectorFunctionSpace(mesh, "CG", 2) -Q = FunctionSpace(mesh, "CG", 1) -W = V * Q - -t = Constant(0.0) -MC = MeshConstant(mesh) -dt = MC.Constant(dt_value) - -x = SpatialCoordinate(mesh) -u_exact = 0.5 * exp(T - t) * as_vector([ - sin(pi * x[0]) * cos(pi * x[1]), - -cos(pi * x[0]) * sin(pi * x[1]) -]) -p_exact = Constant(0.0) - -u_t = Dt(u_exact) -Delta_u = div(grad(u_exact)) -conv_u = dot(grad(u_exact), u_exact) -f_expr = u_t - nu * Delta_u - conv_u - -w = Function(W) -(u, p) = split(w) -(u0, p0) = w.subfunctions -u0.interpolate(u_exact) -p0.assign(0.0) - -bc = DirichletBC(W.sub(0), u_exact, "on_boundary") - -v, q = TestFunctions(W) -F = (inner(Dt(u), v) * dx - + nu * inner(grad(u), grad(v)) * dx - + inner(dot(grad(u), u), v) * dx - - inner(p, div(v)) * dx - + inner(div(u), q) * dx - - inner(f_expr, v) * dx - + gamma * inner(div(u), div(v)) * dx) - -nsp = [(1, VectorSpaceBasis(constant=True))] - -velocity_fields = ",".join(str(2*i) for i in range(s)) -pressure_fields = ",".join(str(2*i+1) for i in range(s)) - -parameters = { - "mat_type": "matfree", - "ksp_type": "fgmres", - "ksp_rtol": 1e-8, - "ksp_converged_reason": None, - "pc_type": "fieldsplit", - "pc_fieldsplit_type": "schur", - "pc_fieldsplit_schur_factorization_type": "upper", - "pc_fieldsplit_0_fields": velocity_fields, - "pc_fieldsplit_1_fields": pressure_fields, - "fieldsplit_0": { - "ksp_type": "preonly", - "ksp_rtol": 1e-10, - "pc_type": "python", - "pc_python_type": "firedrake.AssembledPC", - "assembled_ksp_type": "preonly", - "assembled_pc_type": "lu", - }, - "fieldsplit_1": { - "ksp_type": "preonly", - "ksp_rtol": 1e-10, - "pc_type": "python", - "pc_python_type": "KronPC.KronPC", - "kron": { - "operator": "mass", - "pow": 1, - "pc_type": "lu" - } - }, -} - -stepper = TimeStepper( - F, - RadauIIA(s), - t, - dt, - w, - bcs=bc, - stage_type="deriv", - solver_parameters=parameters, - nullspace=nsp -) - -while float(t) < T - 1e-10: - stepper.advance() - t.assign(float(t) + float(dt)) - -steps, nonlinear_its, linear_its = stepper.solver_stats() -print("Total number of timesteps was %d" % (steps)) -print("Average number of nonlinear iterations per timestep was %.2f" % (nonlinear_its/steps)) -print("Average number of linear iterations per timestep was %.2f" % (linear_its/steps)) From 9275f471dca5808f67d14821e228ccdf710d5699 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Mon, 18 Aug 2025 09:23:54 -0500 Subject: [PATCH 06/35] .DS_Store added in the git ignore list --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index d29ef43b..2bd21bc1 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ __pycache__/ IRKsome.egg-info/ /demos/*/*.py +.DS_Store From f5d93755a9f941e986c658a154491b4643bfc8dc Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Mon, 18 Aug 2025 09:42:35 -0500 Subject: [PATCH 07/35] update init to import KronPC --- irksome/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/irksome/__init__.py b/irksome/__init__.py index a32b97e1..69cab89d 100644 --- a/irksome/__init__.py +++ b/irksome/__init__.py @@ -31,3 +31,5 @@ from .wso_dirk_tableaux import WSODIRK # noqa: F401 from .galerkin_stepper import GalerkinTimeStepper # noqa: F401 from .discontinuous_galerkin_stepper import DiscontinuousGalerkinTimeStepper # noqa: F401 +from .KronPC import KronPC + From 4fb87ea7e5e7e307b712e7b0e50e8c6e9c8e895e Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Tue, 19 Aug 2025 12:54:55 -0500 Subject: [PATCH 08/35] KronPC is now have subclasses MassKronPC(KronPC) and StiffnessKronPC(KronPC) as needed. --- irksome/KronPC.py | 44 ++++++++++++++++++++++++++++---------------- 1 file changed, 28 insertions(+), 16 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index 66f5d3dc..37eb62ff 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -7,7 +7,7 @@ from irksome import RadauIIA import numpy as np -__all__ = ("KronPC",) +__all__ = ("KronPC","MassKronPC", "StiffnessKronPC") # ----------------------------- # KronPC definition @@ -16,21 +16,17 @@ class KronPC(PCBase): r""" Preconditioner applying: y = coef * (A^{-p} \otimes K^{-1}) x - A is the Butcher matrix (RadauIIA(s)), p >= 0 (p=0 => I) - - K is assembled on the single-stage space via kron_operator \in {mass, stiffness} - - K^{-1} is approximated by a PETSc PC with prefix 'kron_sub_*' + - K is assembled on the single-stage space by subclasses via "form(trial, test)" + - K^{-1} is approximated by a PETSc PC with prefix """ needs_python_pmat = True - def form(self, trial, test, operator_kind): - if operator_kind == "mass": - a = inner(trial, test) * dx - elif operator_kind == "stiffness": - a = inner(grad(trial), grad(test)) * dx + 1e-12 * inner(trial, test) * dx - else: - raise ValueError(f"Unknown kron_operator '{operator_kind}' (use 'mass' or 'stiffness').") - bcs = None - return a, bcs + + def form(self, trial, test): + """Return (a, bcs) for the single-stage operator K.""" + raise NotImplementedError("KronPC.form() is abstract. Use MassKronPC or StiffnessKronPC (or a custom subclass).\ + Subclass must implement 'form(trial, test)'.") def stage_mat(self, s, pow_): if pow_ == 0: @@ -58,7 +54,6 @@ def initialize(self, pc): raise ValueError("kron_pow must be a nonnegative integer.") self._pow = pow_ - operator_kind = opts.getString(self._prefix + "operator", "mass") mat_type = opts.getString(self._prefix + "mat_type", "aij") _, P = pc.getOperators() @@ -72,13 +67,15 @@ def initialize(self, pc): Vstage = Vbig.sub(0) trial = TrialFunction(Vstage) test = TestFunction(Vstage) - a, bcs = self.form(trial, test, operator_kind) + a, bcs = self.form(trial, test) fc_params = getattr(context, "fc_params", None) K = assemble(a, bcs=bcs, mat_type=mat_type, form_compiler_parameters=fc_params) self.K = K sub_pc = PETSc.PC().create(comm=pc.comm) + # Allow users to set options like kron_sub_pc_type hypre + sub_pc.setOptionsPrefix(self._prefix + "sub_") sub_pc.incrementTabLevel(1, parent=pc) sub_pc.setOperators(K.M.handle) sub_pc.setFromOptions() @@ -130,8 +127,8 @@ def view(self, pc, viewer=None): viewer.printfASCII("KronPC:\n") viewer.printfASCII(f" stages : {self.num_stages}\n") viewer.printfASCII(f" coef : {self.coef}\n") - viewer.printfASCII(f" power (p) : {self._pow} [A^{-p}]\n") - viewer.printfASCII(" sub-PC for K^{-1} (kron_sub_*) options:\n") + viewer.printfASCII(f" power (p) : {self._pow} [A^{{-p}}]\n") + viewer.printfASCII(" sub-PC for K^{-1} options:\n") if hasattr(self, "sub_pc") and self.sub_pc is not None: self.sub_pc.view(viewer) @@ -139,3 +136,18 @@ def destroy(self, pc): if hasattr(self, "sub_pc") and self.sub_pc is not None: self.sub_pc.destroy() self.sub_pc = None + +class MassKronPC(KronPC): + """K built from the mass form.""" + def form(self, trial, test): + a = inner(trial, test) * dx + bcs = None + return a, bcs + + +class StiffnessKronPC(KronPC): + """K built from the (regularized) stiffness form.""" + def form(self, trial, test): + a = inner(grad(trial), grad(test)) * dx + 1e-12 * inner(trial, test) * dx + bcs = None + return a, bcs \ No newline at end of file From 885d10f671d2b8797c31abaf128d8863507e5059 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Tue, 19 Aug 2025 12:56:54 -0500 Subject: [PATCH 09/35] update init to import KronPC subclasses --- irksome/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/irksome/__init__.py b/irksome/__init__.py index 69cab89d..de051ed4 100644 --- a/irksome/__init__.py +++ b/irksome/__init__.py @@ -31,5 +31,5 @@ from .wso_dirk_tableaux import WSODIRK # noqa: F401 from .galerkin_stepper import GalerkinTimeStepper # noqa: F401 from .discontinuous_galerkin_stepper import DiscontinuousGalerkinTimeStepper # noqa: F401 -from .KronPC import KronPC +from .KronPC import KronPC, MassKronPC, StiffnessKronPC # noqa: F401 From b38e70a38fbdbd795e7196631b1fdd78163691cf Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 20 Aug 2025 16:04:47 -0500 Subject: [PATCH 10/35] pow in the stage_mat removed. stage_mat can now whatever it wants internally. --- irksome/KronPC.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index 37eb62ff..6010af49 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -28,14 +28,13 @@ def form(self, trial, test): raise NotImplementedError("KronPC.form() is abstract. Use MassKronPC or StiffnessKronPC (or a custom subclass).\ Subclass must implement 'form(trial, test)'.") - def stage_mat(self, s, pow_): - if pow_ == 0: + def stage_mat(self, s): + p = getattr(self, "_pow", 0) + if p == 0: return np.eye(s) A = RadauIIA(s).A Ainv = np.linalg.inv(A) - if pow_ == 1: - return Ainv - return np.linalg.matrix_power(Ainv, pow_) + return Ainv if p == 1 else np.linalg.matrix_power(Ainv, p) @property def num_stages(self): @@ -85,7 +84,7 @@ def initialize(self, pc): self.sub_pc = sub_pc self._s = Vbig.num_sub_spaces() - self.L = self.stage_mat(self._s, self._pow) + self.L = self.stage_mat(self._s) def update(self, pc): pass From 2a3f0ebcfab7a206111c7a21d555dc63f2540398 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 20 Aug 2025 17:07:41 -0500 Subject: [PATCH 11/35] test added for verifying KronPC --- tests/test_kronpc.py | 89 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 tests/test_kronpc.py diff --git a/tests/test_kronpc.py b/tests/test_kronpc.py new file mode 100644 index 00000000..7f44e5fd --- /dev/null +++ b/tests/test_kronpc.py @@ -0,0 +1,89 @@ +import numpy as np +import pytest +from functools import reduce +from operator import mul + + +from firedrake import ( + UnitSquareMesh, FunctionSpace, TrialFunction, TestFunction, Cofunction, + Function, as_tensor, inner, dx, solve, errornorm, ufl +) + +def getA(ns: int): + """nonsymmetric, invertible matrix.""" + A = np.zeros((ns, ns)) + for i in range(ns): + for j in range(ns): + A[i, j] = 1.0 + 0.1*(i+1) - 0.2*(j+1) # breaks symmetry + A[i, i] += ns + return A + + +def random_rhs(Vbig, seed: int = 1234): + """random RHS in the dual of Vbig. Yes, we need the dual space here!!""" + rng = np.random.default_rng(seed) + L = Cofunction(Vbig.dual(), name="rhs") + with L.dat.vec_wo as v: + arr = rng.standard_normal(v.getSize()) + v.setValues(range(v.getSize()), arr) + v.assemble() + return L + +@pytest.mark.parametrize("ns", [2,3,4,5]) # meaningful cases are ns >= 2 +def test_mass_kron_pc(ns): + """ + Solve two ways: + (1) direct LU on the full mixed operator + (2) preonly with Python PC = MassKronPC and inner stage LU + """ + # Create mesh and function-space + msh = UnitSquareMesh(2, 2) + V = FunctionSpace(msh, "CG", 1) + + # Stage-stacked space V + Vbig = reduce(mul, [V] * ns) + + # Build global operator + A_np = getA(ns) + A = as_tensor(A_np) + uu = TrialFunction(Vbig) + vv = TestFunction(Vbig) + a = inner(ufl.dot(A, uu), vv) * dx + + # The RHS + L = random_rhs(Vbig, seed=2025) + + # Reference solution via full LU + u_ref = Function(Vbig, name="u_ref") + solve( + a == L, + u_ref, + solver_parameters={ + "ksp_type": "preonly", + "pc_type": "lu" + }, + ) + + # Exact inverse via our KronPC: + u_pc = Function(Vbig, name="u_pc") + solve( + a == L, + u_pc, + solver_parameters={ + "ksp_type": "preonly", + "pc_type": "python", + "pc_python_type": "irksome.MassKronPC", # per professor + "mat_type": "matfree", + "kron": { + "pow": 1, + "coef": 1.0, + "sub":{ + "pc_type": "lu" + }, + } + }, + appctx={"A": A_np}, + ) + + err = errornorm(u_ref, u_pc, norm_type="l2") + assert err < 1.0e-10, f"MassKronPC mismatch: ||u_ref - u_pc|| = {err:.3e} (ns={ns})" From b3ebc087c719d8b660ad0075792d033eca7694e3 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 20 Aug 2025 17:22:57 -0500 Subject: [PATCH 12/35] test added for verifying KronPC, unncessary comments removed. --- tests/test_kronpc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_kronpc.py b/tests/test_kronpc.py index 7f44e5fd..f5f1acb6 100644 --- a/tests/test_kronpc.py +++ b/tests/test_kronpc.py @@ -72,7 +72,7 @@ def test_mass_kron_pc(ns): solver_parameters={ "ksp_type": "preonly", "pc_type": "python", - "pc_python_type": "irksome.MassKronPC", # per professor + "pc_python_type": "irksome.MassKronPC", "mat_type": "matfree", "kron": { "pow": 1, From 81edd3f7c61f5f7fb99fc06cda5bb5d94fbdabe3 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Thu, 21 Aug 2025 09:34:27 -0500 Subject: [PATCH 13/35] update KronPC.py --- irksome/KronPC.py | 33 ++++++++++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index 6010af49..79db5ddb 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -4,7 +4,6 @@ from firedrake.preconditioners.base import PCBase from firedrake.petsc import PETSc from firedrake import * -from irksome import RadauIIA import numpy as np __all__ = ("KronPC","MassKronPC", "StiffnessKronPC") @@ -32,7 +31,9 @@ def stage_mat(self, s): p = getattr(self, "_pow", 0) if p == 0: return np.eye(s) - A = RadauIIA(s).A + if not hasattr(self, "_A") or self._A is None: + raise ValueError("KronPC: stage matrix A not set.") + A = self._A Ainv = np.linalg.inv(A) return Ainv if p == 1 else np.linalg.matrix_power(Ainv, p) @@ -82,8 +83,34 @@ def initialize(self, pc): sub_pc.setType("lu") sub_pc.setUp() self.sub_pc = sub_pc - self._s = Vbig.num_sub_spaces() + + # --- discover A from user-provided context (prefer appctx) --- + self._A = None + appctx = getattr(context, "appctx", None) or {} + + A_src = None + bt = appctx.get("butcher_tableau", None) + if bt is not None and hasattr(bt, "A"): + A_src = bt.A + elif "A" in appctx: + A_src = appctx["A"] + elif hasattr(context, "A"): + A_src = context.A + else: + bt = getattr(context, "butcher_tableau", None) + if bt is not None and hasattr(bt, "A"): + A_src = bt.A + + self._A = None if A_src is None else np.asarray(A_src, dtype=float) + # validate only if we actually need A + if self._pow > 0: + if self._A is None: + raise ValueError("KronPC: no stage matrix A found. " + "Set context.butcher_tableau.A or context.A before solve.") + if self._A.shape != (self._s, self._s): + raise ValueError(f"KronPC: A has shape {self._A.shape}, but s={self._s}.") + self.L = self.stage_mat(self._s) def update(self, pc): From 97b8291c20145774f3df11199cf6e968c94d1d19 Mon Sep 17 00:00:00 2001 From: "Robert C. Kirby" Date: Thu, 21 Aug 2025 09:51:14 -0500 Subject: [PATCH 14/35] comments? --- tests/test_kronpc.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_kronpc.py b/tests/test_kronpc.py index f5f1acb6..fa0464f2 100644 --- a/tests/test_kronpc.py +++ b/tests/test_kronpc.py @@ -29,6 +29,7 @@ def random_rhs(Vbig, seed: int = 1234): v.assemble() return L + @pytest.mark.parametrize("ns", [2,3,4,5]) # meaningful cases are ns >= 2 def test_mass_kron_pc(ns): """ @@ -80,7 +81,7 @@ def test_mass_kron_pc(ns): "sub":{ "pc_type": "lu" }, - } + } }, appctx={"A": A_np}, ) From de8a4a991835ed73f4cd24b3ab6a747531b619f0 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Fri, 22 Aug 2025 11:48:57 -0500 Subject: [PATCH 15/35] KronPC now does not involve any matrix power or coeff multiply --- irksome/KronPC.py | 96 ++++++++++++++++++++--------------------------- 1 file changed, 41 insertions(+), 55 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index 79db5ddb..68d2efca 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -13,34 +13,18 @@ # ----------------------------- class KronPC(PCBase): r""" - Preconditioner applying: y = coef * (A^{-p} \otimes K^{-1}) x - - A is the Butcher matrix (RadauIIA(s)), p >= 0 (p=0 => I) + Preconditioner applying: y = (L \otimes K^{-1}) x + - L is the stage matrix, provided by the user. If A is provided, we set L = A^{-1}. If neither provided, L = I_s. - K is assembled on the single-stage space by subclasses via "form(trial, test)" - K^{-1} is approximated by a PETSc PC with prefix """ - needs_python_pmat = True - def form(self, trial, test): """Return (a, bcs) for the single-stage operator K.""" raise NotImplementedError("KronPC.form() is abstract. Use MassKronPC or StiffnessKronPC (or a custom subclass).\ Subclass must implement 'form(trial, test)'.") - def stage_mat(self, s): - p = getattr(self, "_pow", 0) - if p == 0: - return np.eye(s) - if not hasattr(self, "_A") or self._A is None: - raise ValueError("KronPC: stage matrix A not set.") - A = self._A - Ainv = np.linalg.inv(A) - return Ainv if p == 1 else np.linalg.matrix_power(Ainv, p) - - @property - def num_stages(self): - return self._s - def initialize(self, pc): if pc.getType() != "python": raise ValueError("Expecting PC type 'python' for KronPC.") @@ -48,12 +32,6 @@ def initialize(self, pc): self._prefix = (pc.getOptionsPrefix() or "") + "kron_" opts = PETSc.Options() - self.coef = opts.getReal(self._prefix + "coef", 1.0) - pow_ = opts.getInt(self._prefix + "pow", 1) - if pow_ < 0: - raise ValueError("kron_pow must be a nonnegative integer.") - self._pow = pow_ - mat_type = opts.getString(self._prefix + "mat_type", "aij") _, P = pc.getOperators() @@ -74,7 +52,6 @@ def initialize(self, pc): self.K = K sub_pc = PETSc.PC().create(comm=pc.comm) - # Allow users to set options like kron_sub_pc_type hypre sub_pc.setOptionsPrefix(self._prefix + "sub_") sub_pc.incrementTabLevel(1, parent=pc) sub_pc.setOperators(K.M.handle) @@ -84,53 +61,64 @@ def initialize(self, pc): sub_pc.setUp() self.sub_pc = sub_pc self._s = Vbig.num_sub_spaces() + self.L = self._build_stage_L(self._s, context, pc) + - # --- discover A from user-provided context (prefer appctx) --- - self._A = None +# --- discover A from user-provided context (prefer appctx) --- + def _build_stage_L(self, s, context, pc): + """ + Build the stage coupling matrix L. + - If appctx/context supplies A, set L = A^{-1}. + - If appctx/context supplies butcher_tableau with .A, use that. + - Else default to I_s. + """ + # Look in appctx first appctx = getattr(context, "appctx", None) or {} - - A_src = None - bt = appctx.get("butcher_tableau", None) - if bt is not None and hasattr(bt, "A"): - A_src = bt.A - elif "A" in appctx: - A_src = appctx["A"] - elif hasattr(context, "A"): - A_src = context.A - else: - bt = getattr(context, "butcher_tableau", None) + A = appctx.get("A", None) + if A is None: + bt = appctx.get("butcher_tableau", None) if bt is not None and hasattr(bt, "A"): - A_src = bt.A - - self._A = None if A_src is None else np.asarray(A_src, dtype=float) - # validate only if we actually need A - if self._pow > 0: - if self._A is None: - raise ValueError("KronPC: no stage matrix A found. " - "Set context.butcher_tableau.A or context.A before solve.") - if self._A.shape != (self._s, self._s): - raise ValueError(f"KronPC: A has shape {self._A.shape}, but s={self._s}.") - - self.L = self.stage_mat(self._s) + A = bt.A + + # Fall back to direct attributes + if A is None: + if hasattr(context, "A"): + A = context.A + else: + bt = getattr(context, "butcher_tableau", None) + if bt is not None and hasattr(bt, "A"): + A = bt.A + + if A is not None: + A = np.asarray(A, dtype=float) + if A.shape != (s, s): + raise ValueError(f"KronPC: A has shape {A.shape}, expected {(s, s)}.") + return np.linalg.inv(A) + + # Default: identity + return np.eye(s) def update(self, pc): pass def apply(self, pc, x, y): - s = self.num_stages + s = self._s with self.work_in.dat.vec_wo as vin: x.copy(vin) + # Stagewise K^{-1} for i in range(s): with self.work_in.subfunctions[i].dat.vec_ro as xin_i, \ self.work_mid.subfunctions[i].dat.vec_wo as mid_i: mid_i.set(0.0) self.sub_pc.apply(xin_i, mid_i) + # Zero outputs for j in range(s): self.work_out.subfunctions[j].assign(0.0) + # y_stage[j] += sum_i L[j,i] * mid[i] for j in range(s): row = self.L[j, :] with self.work_out.subfunctions[j].dat.vec_wo as yj: @@ -139,7 +127,7 @@ def apply(self, pc, x, y): if lij == 0.0: continue with self.work_mid.subfunctions[i].dat.vec_ro as ui: - yj.axpy(self.coef * lij, ui) + yj.axpy(lij, ui) with self.work_out.dat.vec_ro as vout: vout.copy(y) @@ -151,9 +139,7 @@ def view(self, pc, viewer=None): if viewer is None: viewer = PETSc.Viewer.STDOUT(pc.comm) viewer.printfASCII("KronPC:\n") - viewer.printfASCII(f" stages : {self.num_stages}\n") - viewer.printfASCII(f" coef : {self.coef}\n") - viewer.printfASCII(f" power (p) : {self._pow} [A^{{-p}}]\n") + viewer.printfASCII(f" stages : {self._s}\n") viewer.printfASCII(" sub-PC for K^{-1} options:\n") if hasattr(self, "sub_pc") and self.sub_pc is not None: self.sub_pc.view(viewer) From d26ba968e22dca466bd79356df6beea228037922 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Fri, 22 Aug 2025 11:49:50 -0500 Subject: [PATCH 16/35] updated test which doesn't require any pow or coeff --- tests/test_kronpc.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tests/test_kronpc.py b/tests/test_kronpc.py index fa0464f2..77ab4332 100644 --- a/tests/test_kronpc.py +++ b/tests/test_kronpc.py @@ -75,13 +75,8 @@ def test_mass_kron_pc(ns): "pc_type": "python", "pc_python_type": "irksome.MassKronPC", "mat_type": "matfree", - "kron": { - "pow": 1, - "coef": 1.0, - "sub":{ - "pc_type": "lu" - }, - } + "kron_sub_pc_type": "lu" + }, appctx={"A": A_np}, ) From 6a2ed562b8466aaf172c12b6835d2b84070c0553 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 10 Sep 2025 15:37:15 -0500 Subject: [PATCH 17/35] added Discontinuous Galerkin (SIPG) pressure stiffness for DG spaces --- irksome/KronPC.py | 37 +++++++++++++++++++++++++++++++++++-- 1 file changed, 35 insertions(+), 2 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index 68d2efca..5be8304f 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -6,7 +6,8 @@ from firedrake import * import numpy as np -__all__ = ("KronPC","MassKronPC", "StiffnessKronPC") + +__all__ = ("KronPC","MassKronPC", "StiffnessKronPC", "SIPGStiffnessKronPC") # ----------------------------- # KronPC definition @@ -162,4 +163,36 @@ class StiffnessKronPC(KronPC): def form(self, trial, test): a = inner(grad(trial), grad(test)) * dx + 1e-12 * inner(trial, test) * dx bcs = None - return a, bcs \ No newline at end of file + return a, bcs + +class SIPGStiffnessKronPC(KronPC): + r""" + Discontinuous Galerkin (SIPG) pressure "stiffness" for DG spaces (e.g., P_{k}^{\mathrm{disc}}). + Implementation details: + - The trial/test must belong to a scalar DG space (e.g., DG(k)). + - The sub-solve for \(K^{-1}\) is handled by a PETSc PC with options prefix ``kron_sub_*``. + """ + + def form(self, trial, test): + mesh = trial.ufl_domain() + n = FacetNormal(mesh) + h = CellDiameter(mesh) + + # Polynomial degree k of the DG space + k = trial.ufl_function_space().ufl_element().degree() + + # Read base penalty C from PETSc options; default C=10 + C = PETSc.Options().getReal(self._prefix + "sipg_C", 10.0) + + # \alpha = C * (k+1)(k+2) + alpha = Constant(C * (k + 1) * (k + 2)) + a = (inner(grad(test), grad(trial)) * dx + - inner(avg(grad(trial)), jump(test, n)) * dS + - inner(avg(grad(test)), jump(trial, n)) * dS + + (alpha / avg(h)) * inner(jump(trial), jump(test)) * dS + - inner(grad(test), trial * n) * ds + - inner(grad(trial), test * n) * ds + + (alpha / h) * inner(trial, test) * ds) + + bcs = None + return a, bcs From 262e6d0a187dfed67f3a6fb2af01f562f69b181a Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 10 Sep 2025 15:37:27 -0500 Subject: [PATCH 18/35] update init to import KronPC subclasses --- irksome/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/irksome/__init__.py b/irksome/__init__.py index de051ed4..ff97c2c9 100644 --- a/irksome/__init__.py +++ b/irksome/__init__.py @@ -31,5 +31,5 @@ from .wso_dirk_tableaux import WSODIRK # noqa: F401 from .galerkin_stepper import GalerkinTimeStepper # noqa: F401 from .discontinuous_galerkin_stepper import DiscontinuousGalerkinTimeStepper # noqa: F401 -from .KronPC import KronPC, MassKronPC, StiffnessKronPC # noqa: F401 +from .KronPC import KronPC, MassKronPC, StiffnessKronPC, SIPGStiffnessKronPC # noqa: F401 From e2dd558f966445091db447c243ee2358071912ae Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Thu, 11 Sep 2025 13:41:38 -0500 Subject: [PATCH 19/35] KronPC with updated SIPG --- irksome/KronPC.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index 5be8304f..eb940928 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -161,6 +161,7 @@ def form(self, trial, test): class StiffnessKronPC(KronPC): """K built from the (regularized) stiffness form.""" def form(self, trial, test): + print("Yes I am using non SIPG") a = inner(grad(trial), grad(test)) * dx + 1e-12 * inner(trial, test) * dx bcs = None return a, bcs @@ -184,15 +185,17 @@ def form(self, trial, test): # Read base penalty C from PETSc options; default C=10 C = PETSc.Options().getReal(self._prefix + "sipg_C", 10.0) - # \alpha = C * (k+1)(k+2) - alpha = Constant(C * (k + 1) * (k + 2)) + # penal = C * (k+1)(k+2) + penal = Constant(C * (k + 1) * (k + 2)) a = (inner(grad(test), grad(trial)) * dx + + 1e-12 * inner(trial, test) * dx - inner(avg(grad(trial)), jump(test, n)) * dS - inner(avg(grad(test)), jump(trial, n)) * dS - + (alpha / avg(h)) * inner(jump(trial), jump(test)) * dS + + (penal / avg(h)) * inner(jump(trial), jump(test)) * dS - inner(grad(test), trial * n) * ds - inner(grad(trial), test * n) * ds - + (alpha / h) * inner(trial, test) * ds) + + (penal / h) * inner(trial, test) * ds) + bcs = None return a, bcs From cb8603339b6c65c023a550a32f7fb9183992abd6 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Mon, 22 Sep 2025 11:49:03 -0500 Subject: [PATCH 20/35] updated KronPC, removing previous python context import and directly use pc.getDM() --- irksome/KronPC.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index eb940928..b6ad6ce7 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -5,6 +5,7 @@ from firedrake.petsc import PETSc from firedrake import * import numpy as np +from firedrake.dmhooks import get_appctx, get_function_space __all__ = ("KronPC","MassKronPC", "StiffnessKronPC", "SIPGStiffnessKronPC") @@ -19,7 +20,7 @@ class KronPC(PCBase): - K is assembled on the single-stage space by subclasses via "form(trial, test)" - K^{-1} is approximated by a PETSc PC with prefix """ - needs_python_pmat = True + needs_python_pmat = False def form(self, trial, test): """Return (a, bcs) for the single-stage operator K.""" @@ -35,9 +36,14 @@ def initialize(self, pc): mat_type = opts.getString(self._prefix + "mat_type", "aij") - _, P = pc.getOperators() - context = P.getPythonContext() - Vbig = context.a.arguments()[0].function_space() + # _, P = pc.getOperators() + # context = P.getPythonContext() + dm = pc.getDM() + context = get_appctx(dm) + # print(type(get_appctx(pc.getDM())), getattr(get_appctx(pc.getDM()), "appctx", None)) + + # Vbig = context.a.arguments()[0].function_space() + Vbig = get_function_space(dm) self.work_in = Cofunction(Vbig.dual(), name="kron_work_in") self.work_mid = Function(Vbig, name="kron_work_mid") From 0d916344c117306facf97cbb9ee8135dd7960490 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Tue, 30 Sep 2025 13:36:35 -0500 Subject: [PATCH 21/35] update KronPC and remove unecessary print --- irksome/KronPC.py | 1 - 1 file changed, 1 deletion(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index b6ad6ce7..affe9db9 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -167,7 +167,6 @@ def form(self, trial, test): class StiffnessKronPC(KronPC): """K built from the (regularized) stiffness form.""" def form(self, trial, test): - print("Yes I am using non SIPG") a = inner(grad(trial), grad(test)) * dx + 1e-12 * inner(trial, test) * dx bcs = None return a, bcs From 097d79abab2440682f0b535330cc8c6e852d931e Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Mon, 6 Oct 2025 10:32:28 -0500 Subject: [PATCH 22/35] kron pc changes from sub-pc to sub-ksp --- irksome/KronPC.py | 66 ++++++++++++++++++++++++++++++++++------------- 1 file changed, 48 insertions(+), 18 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index affe9db9..09dbc335 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -58,15 +58,44 @@ def initialize(self, pc): K = assemble(a, bcs=bcs, mat_type=mat_type, form_compiler_parameters=fc_params) self.K = K - sub_pc = PETSc.PC().create(comm=pc.comm) - sub_pc.setOptionsPrefix(self._prefix + "sub_") - sub_pc.incrementTabLevel(1, parent=pc) - sub_pc.setOperators(K.M.handle) - sub_pc.setFromOptions() - if sub_pc.getType() is None: - sub_pc.setType("lu") - sub_pc.setUp() - self.sub_pc = sub_pc + # sub_pc = PETSc.PC().create(comm=pc.comm) + # sub_pc.setOptionsPrefix(self._prefix + "sub_") + # sub_pc.incrementTabLevel(1, parent=pc) + # sub_pc.setOperators(K.M.handle) + # sub_pc.setFromOptions() + # if sub_pc.getType() is None: + # sub_pc.setType("lu") + # sub_pc.setUp() + # self.sub_pc = sub_pc + # self._s = Vbig.num_sub_spaces() + # self.L = self._build_stage_L(self._s, context, pc) + + # NEW: create a sub-KSP (instead of a raw PC) so we can use python PCs with DM/appctx + sub_ksp = PETSc.KSP().create(comm=pc.comm) + sub_ksp.setOptionsPrefix(self._prefix + "sub_") + sub_ksp.incrementTabLevel(1, parent=pc) + + # Propagate DM (and its python context) to the sub-KSP + dmK = K.M.handle.getDM() or pc.getDM() + if dmK is not None: + sub_ksp.setDM(dmK) + # Let the sub-KSP use the DM for matrix-free PCs if needed + sub_ksp.setDMActive(False) + + # Operators and options + sub_ksp.setOperators(K.M.handle) + sub_ksp.setFromOptions() + + # Sensible defaults if the user didn't set anything: + if sub_ksp.getType() is None: + sub_ksp.setType("preonly") + if sub_ksp.getPC().getType() is None: + sub_ksp.getPC().setType("lu") + + sub_ksp.setUp() + + self.sub_ksp = sub_ksp + self.sub_pc = sub_ksp.getPC() # keep for .view() convenience self._s = Vbig.num_sub_spaces() self.L = self._build_stage_L(self._s, context, pc) @@ -116,10 +145,10 @@ def apply(self, pc, x, y): # Stagewise K^{-1} for i in range(s): - with self.work_in.subfunctions[i].dat.vec_ro as xin_i, \ + with self.work_in.subfunctions[i].dat.vec_ro as rhs_i, \ self.work_mid.subfunctions[i].dat.vec_wo as mid_i: mid_i.set(0.0) - self.sub_pc.apply(xin_i, mid_i) + self.sub_ksp.solve(rhs_i, mid_i) # Zero outputs for j in range(s): @@ -147,14 +176,15 @@ def view(self, pc, viewer=None): viewer = PETSc.Viewer.STDOUT(pc.comm) viewer.printfASCII("KronPC:\n") viewer.printfASCII(f" stages : {self._s}\n") - viewer.printfASCII(" sub-PC for K^{-1} options:\n") - if hasattr(self, "sub_pc") and self.sub_pc is not None: - self.sub_pc.view(viewer) + viewer.printfASCII(" sub-KSP for K^{-1} options:\n") + if hasattr(self, "sub_ksp") and self.sub_ksp is not None: + self.sub_ksp.view(viewer) def destroy(self, pc): - if hasattr(self, "sub_pc") and self.sub_pc is not None: - self.sub_pc.destroy() - self.sub_pc = None + if hasattr(self, "sub_ksp") and self.sub_ksp is not None: + self.sub_ksp.destroy() + self.sub_ksp = None + self.sub_pc = None class MassKronPC(KronPC): """K built from the mass form.""" @@ -193,7 +223,7 @@ def form(self, trial, test): # penal = C * (k+1)(k+2) penal = Constant(C * (k + 1) * (k + 2)) a = (inner(grad(test), grad(trial)) * dx - + 1e-12 * inner(trial, test) * dx + # + 1e-12 * inner(trial, test) * dx - inner(avg(grad(trial)), jump(test, n)) * dS - inner(avg(grad(test)), jump(trial, n)) * dS + (penal / avg(h)) * inner(jump(trial), jump(test)) * dS From 7400f5c3182a0216ce36ee8eec8a5f1a89004475 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Mon, 6 Oct 2025 11:55:09 -0500 Subject: [PATCH 23/35] update KronPC --- irksome/KronPC.py | 54 ++++++++++++++++++++++++++++++----------------- 1 file changed, 35 insertions(+), 19 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index 09dbc335..d00c0951 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -36,13 +36,9 @@ def initialize(self, pc): mat_type = opts.getString(self._prefix + "mat_type", "aij") - # _, P = pc.getOperators() - # context = P.getPythonContext() dm = pc.getDM() context = get_appctx(dm) - # print(type(get_appctx(pc.getDM())), getattr(get_appctx(pc.getDM()), "appctx", None)) - # Vbig = context.a.arguments()[0].function_space() Vbig = get_function_space(dm) self.work_in = Cofunction(Vbig.dual(), name="kron_work_in") @@ -58,17 +54,6 @@ def initialize(self, pc): K = assemble(a, bcs=bcs, mat_type=mat_type, form_compiler_parameters=fc_params) self.K = K - # sub_pc = PETSc.PC().create(comm=pc.comm) - # sub_pc.setOptionsPrefix(self._prefix + "sub_") - # sub_pc.incrementTabLevel(1, parent=pc) - # sub_pc.setOperators(K.M.handle) - # sub_pc.setFromOptions() - # if sub_pc.getType() is None: - # sub_pc.setType("lu") - # sub_pc.setUp() - # self.sub_pc = sub_pc - # self._s = Vbig.num_sub_spaces() - # self.L = self._build_stage_L(self._s, context, pc) # NEW: create a sub-KSP (instead of a raw PC) so we can use python PCs with DM/appctx sub_ksp = PETSc.KSP().create(comm=pc.comm) @@ -86,7 +71,7 @@ def initialize(self, pc): sub_ksp.setOperators(K.M.handle) sub_ksp.setFromOptions() - # Sensible defaults if the user didn't set anything: + # defaults if the user didn't set anything: if sub_ksp.getType() is None: sub_ksp.setType("preonly") if sub_ksp.getPC().getType() is None: @@ -95,7 +80,7 @@ def initialize(self, pc): sub_ksp.setUp() self.sub_ksp = sub_ksp - self.sub_pc = sub_ksp.getPC() # keep for .view() convenience + self.sub_pc = sub_ksp.getPC() self._s = Vbig.num_sub_spaces() self.L = self._build_stage_L(self._s, context, pc) @@ -139,6 +124,7 @@ def update(self, pc): def apply(self, pc, x, y): s = self._s + y.set(0.0) with self.work_in.dat.vec_wo as vin: x.copy(vin) @@ -169,7 +155,38 @@ def apply(self, pc, x, y): vout.copy(y) def applyTranspose(self, pc, x, y): - self.apply(pc, x, y) + s = self._s + y.set(0.0) + + with self.work_in.dat.vec_wo as vin: + x.copy(vin) + + # Stagewise K^{-T} + for i in range(s): + with self.work_in.subfunctions[i].dat.vec_ro as rhs_i, \ + self.work_mid.subfunctions[i].dat.vec_wo as mid_i: + mid_i.set(0.0) + # requires the sub-KSP (PC) to support transpose solves + self.sub_ksp.solveTranspose(rhs_i, mid_i) + + # Zero outputs + for j in range(s): + self.work_out.subfunctions[j].assign(0.0) + + # y_stage[j] = sum_i L^T[j,i] * mid[i] + LT = self.L.T + for j in range(s): + row = LT[j, :] + with self.work_out.subfunctions[j].dat.vec_wo as yj: + for i in range(s): + lij = float(row[i]) + if lij != 0.0: + with self.work_mid.subfunctions[i].dat.vec_ro as ui: + yj.axpy(lij, ui) + + with self.work_out.dat.vec_ro as vout: + vout.copy(y) + def view(self, pc, viewer=None): if viewer is None: @@ -223,7 +240,6 @@ def form(self, trial, test): # penal = C * (k+1)(k+2) penal = Constant(C * (k + 1) * (k + 2)) a = (inner(grad(test), grad(trial)) * dx - # + 1e-12 * inner(trial, test) * dx - inner(avg(grad(trial)), jump(test, n)) * dS - inner(avg(grad(test)), jump(trial, n)) * dS + (penal / avg(h)) * inner(jump(trial), jump(test)) * dS From b30382cdd75cf4f23b431da1d430bda8746d17b5 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 8 Oct 2025 14:34:02 -0500 Subject: [PATCH 24/35] DG stiffness matrix do have a nullspace attached to it --- irksome/KronPC.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index d00c0951..07e831db 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -250,3 +250,25 @@ def form(self, trial, test): bcs = None return a, bcs + + def initialize(self, pc): + # Assemble K and build the sub-KSP using the base class logic + super().initialize(pc) + + # If your form is *truly* pure Neumann (constant kernel), attach a constant nullspace: + # Find the stage scalar space + Vbig = get_function_space(pc.getDM()) + Vstage = Vbig.sub(0) + + # Build a constant Function on DG space + self._ns_vector = Function(Vstage, name="kron_sipg_const") # keep a ref! + self._ns_vector.assign(1.0) + + with self._ns_vector.dat.vec_ro as v: + # constant=True lets PETSc treat it as a constant field; we also provide the vector + ns = PETSc.NullSpace().create(constant=True, vectors=[v]) + + # Attach to the assembled operator(s) used by the sub-KSP (A and P) + A, P = self.sub_ksp.getOperators() + A.setNullSpace(ns) + P.setNullSpace(ns) From d6fe945694c342f931c5a2b4af81bebfc2593aed Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Mon, 20 Oct 2025 15:19:12 -0500 Subject: [PATCH 25/35] DG stiffness matrix do have a nullspace attached to it with project RHS off constant mode --- irksome/KronPC.py | 94 +++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 79 insertions(+), 15 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index 07e831db..d21bb348 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -80,7 +80,9 @@ def initialize(self, pc): sub_ksp.setUp() self.sub_ksp = sub_ksp - self.sub_pc = sub_ksp.getPC() + self.sub_pc = sub_ksp.getPC() + A, _ = self.sub_ksp.getOperators() + self._tmp_stage = A.createVecRight() self._s = Vbig.num_sub_spaces() self.L = self._build_stage_L(self._s, context, pc) @@ -229,7 +231,7 @@ class SIPGStiffnessKronPC(KronPC): def form(self, trial, test): mesh = trial.ufl_domain() n = FacetNormal(mesh) - h = CellDiameter(mesh) + h = CellVolume(mesh) # Polynomial degree k of the DG space k = trial.ufl_function_space().ufl_element().degree() @@ -251,24 +253,86 @@ def form(self, trial, test): bcs = None return a, bcs + # def initialize(self, pc): + # # Assemble K and build the sub-KSP using the base class logic + # super().initialize(pc) + + # # If your form is *truly* pure Neumann (constant kernel), attach a constant nullspace: + # # Find the stage scalar space + # Vbig = get_function_space(pc.getDM()) + # Vstage = Vbig.sub(0) + + # # Build a constant Function on DG space + # self._ns_vector = Function(Vstage, name="kron_sipg_const") # keep a ref! + # self._ns_vector.assign(1.0) + + # with self._ns_vector.dat.vec_ro as v: + # # constant=True lets PETSc treat it as a constant field; we also provide the vector + # self._ns = PETSc.NullSpace().create(constant=True, vectors=[v]) + + # # Attach to the assembled operator(s) used by the sub-KSP (A and P) + # A, P = self.sub_ksp.getOperators() + # A.setNullSpace(self._ns) + # P.setNullSpace(self._ns) + def initialize(self, pc): # Assemble K and build the sub-KSP using the base class logic super().initialize(pc) - # If your form is *truly* pure Neumann (constant kernel), attach a constant nullspace: - # Find the stage scalar space - Vbig = get_function_space(pc.getDM()) - Vstage = Vbig.sub(0) - # Build a constant Function on DG space - self._ns_vector = Function(Vstage, name="kron_sipg_const") # keep a ref! - self._ns_vector.assign(1.0) + S = get_function_space(pc.getDM()) # what the sub-KSP actually solves on + Q = S.sub(1) + + self._ns_vec = Function(Q, name="sipg_const") + self._ns_vec.assign(1.0) - with self._ns_vector.dat.vec_ro as v: - # constant=True lets PETSc treat it as a constant field; we also provide the vector - ns = PETSc.NullSpace().create(constant=True, vectors=[v]) + with self._ns_vec.dat.vec_ro as v: + self._ns = PETSc.NullSpace().create(constant=True, vectors=[v]) - # Attach to the assembled operator(s) used by the sub-KSP (A and P) A, P = self.sub_ksp.getOperators() - A.setNullSpace(ns) - P.setNullSpace(ns) + A.setNullSpace(self._ns) + P.setNullSpace(self._ns) + + + + def apply(self, pc, x, y): + s = self._s + y.set(0.0) + + # copy input into work_in + with self.work_in.dat.vec_wo as vin: + x.copy(vin) + + # stagewise K^{-1} with per-stage projections + for i in range(s): + with self.work_in.subfunctions[i].dat.vec_ro as rhs_i_ro, \ + self.work_mid.subfunctions[i].dat.vec_wo as mid_i: + + # tmp = self._tmp_stage + # tmp.copy(rhs_i_ro) + self._ns.remove(rhs_i_ro) # project RHS off constant mode + + mid_i.set(0.0) + self.sub_ksp.solve(rhs_i_ro, mid_i) + + # self._ns.remove(mid_i) # keep solution in mean-free subspace + + # zero outputs + for j in range(s): + self.work_out.subfunctions[j].assign(0.0) + + # y_j = sum_i L[j,i] * mid_i (with projection) + for j in range(s): + row = self.L[j, :] + with self.work_out.subfunctions[j].dat.vec_wo as yj: + for i in range(s): + lij = row[i] + if lij != 0.0: + with self.work_mid.subfunctions[i].dat.vec_ro as ui: + yj.axpy(lij, ui) + + + with self.work_out.dat.vec_ro as vout: + vout.copy(y) + + From f439493d7b6de42d07e3a5fd22b290a0afc52bbb Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Tue, 4 Nov 2025 11:31:31 -0600 Subject: [PATCH 26/35] fix KronPC to resolve single stage issue when using Rana preconditioner --- irksome/KronPC.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index d21bb348..bc22a11b 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -83,7 +83,11 @@ def initialize(self, pc): self.sub_pc = sub_ksp.getPC() A, _ = self.sub_ksp.getOperators() self._tmp_stage = A.createVecRight() - self._s = Vbig.num_sub_spaces() + try: + s = len(self.work_in.subfunctions) + except Exception: + s = 1 + self._s = s self.L = self._build_stage_L(self._s, context, pc) From d6665428829b210a4db4d433579ab299b37bff78 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 29 Apr 2026 14:53:23 -0500 Subject: [PATCH 27/35] update __init__.py --- irksome/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/irksome/__init__.py b/irksome/__init__.py index ddee498c..0ebb996b 100644 --- a/irksome/__init__.py +++ b/irksome/__init__.py @@ -76,6 +76,7 @@ from .discontinuous_galerkin_stepper import DiscontinuousGalerkinTimeStepper from .labeling import TimeQuadratureLabel from .stepper import TimeStepper + from .KronPC import KronPC, MassKronPC, StiffnessKronPC, SIPGStiffnessKronPC __all__ += [ "DIRKTimeStepper", From 695f0b3357aa440ac3781a810905bb5d26586786 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Tue, 5 May 2026 15:44:36 -0500 Subject: [PATCH 28/35] Update KronPC for new petsc4py setDMActive API --- irksome/KronPC.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/irksome/KronPC.py b/irksome/KronPC.py index bc22a11b..385f4221 100644 --- a/irksome/KronPC.py +++ b/irksome/KronPC.py @@ -60,12 +60,17 @@ def initialize(self, pc): sub_ksp.setOptionsPrefix(self._prefix + "sub_") sub_ksp.incrementTabLevel(1, parent=pc) - # Propagate DM (and its python context) to the sub-KSP + # Propagate DM and its python context to the sub-KSP dmK = K.M.handle.getDM() or pc.getDM() if dmK is not None: sub_ksp.setDM(dmK) - # Let the sub-KSP use the DM for matrix-free PCs if needed - sub_ksp.setDMActive(False) + + # Do not let the DM generate operators/RHS/initial guess, + # because we already set the assembled operator below. + try: + sub_ksp.setDMActive(PETSc.KSP.DMActive.ALL, False) + except TypeError: + sub_ksp.setDMActive(False) # Operators and options sub_ksp.setOperators(K.M.handle) From 16ef196ec1230a24b854dea45300968f66435f64 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Tue, 5 May 2026 15:57:40 -0500 Subject: [PATCH 29/35] Import Irksome early in KronPC test --- tests/test_kronpc.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_kronpc.py b/tests/test_kronpc.py index 77ab4332..d3a53e51 100644 --- a/tests/test_kronpc.py +++ b/tests/test_kronpc.py @@ -1,12 +1,17 @@ import numpy as np import pytest +import ufl + from functools import reduce from operator import mul +# Important: import irksome before any Firedrake/UFL forms are built/processed. +# This lets Irksome register its UFL extensions before UFL MultiFunctions run. +import irksome from firedrake import ( UnitSquareMesh, FunctionSpace, TrialFunction, TestFunction, Cofunction, - Function, as_tensor, inner, dx, solve, errornorm, ufl + Function, as_tensor, inner, dx, solve, errornorm ) def getA(ns: int): From f2eddd8c768620e260b8bb0ef8b99da43741d52f Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 6 May 2026 11:55:05 -0500 Subject: [PATCH 30/35] update kronpc and init files --- irksome/__init__.py | 2 +- irksome/{KronPC.py => kronpc.py} | 23 +---------------------- 2 files changed, 2 insertions(+), 23 deletions(-) rename irksome/{KronPC.py => kronpc.py} (92%) diff --git a/irksome/__init__.py b/irksome/__init__.py index e9aef3f8..3f1dfb1d 100644 --- a/irksome/__init__.py +++ b/irksome/__init__.py @@ -87,7 +87,7 @@ from .multistep import MultistepTimeStepper from .labeling import TimeQuadratureLabel from .stepper import TimeStepper - from .KronPC import KronPC, MassKronPC, StiffnessKronPC, SIPGStiffnessKronPC + from .kronpc import KronPC, MassKronPC, StiffnessKronPC, SIPGStiffnessKronPC __all__ += [ "DIRKTimeStepper", diff --git a/irksome/KronPC.py b/irksome/kronpc.py similarity index 92% rename from irksome/KronPC.py rename to irksome/kronpc.py index 385f4221..8daf4267 100644 --- a/irksome/KronPC.py +++ b/irksome/kronpc.py @@ -262,33 +262,12 @@ def form(self, trial, test): bcs = None return a, bcs - # def initialize(self, pc): - # # Assemble K and build the sub-KSP using the base class logic - # super().initialize(pc) - - # # If your form is *truly* pure Neumann (constant kernel), attach a constant nullspace: - # # Find the stage scalar space - # Vbig = get_function_space(pc.getDM()) - # Vstage = Vbig.sub(0) - - # # Build a constant Function on DG space - # self._ns_vector = Function(Vstage, name="kron_sipg_const") # keep a ref! - # self._ns_vector.assign(1.0) - - # with self._ns_vector.dat.vec_ro as v: - # # constant=True lets PETSc treat it as a constant field; we also provide the vector - # self._ns = PETSc.NullSpace().create(constant=True, vectors=[v]) - - # # Attach to the assembled operator(s) used by the sub-KSP (A and P) - # A, P = self.sub_ksp.getOperators() - # A.setNullSpace(self._ns) - # P.setNullSpace(self._ns) + def initialize(self, pc): # Assemble K and build the sub-KSP using the base class logic super().initialize(pc) - S = get_function_space(pc.getDM()) # what the sub-KSP actually solves on Q = S.sub(1) From c86c95985646527514cee12c319ece3d7739979e Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 6 May 2026 12:15:37 -0500 Subject: [PATCH 31/35] Fix kronpc lint --- irksome/kronpc.py | 51 +++++++++++++++++++++-------------------------- 1 file changed, 23 insertions(+), 28 deletions(-) diff --git a/irksome/kronpc.py b/irksome/kronpc.py index 8daf4267..709eed1a 100644 --- a/irksome/kronpc.py +++ b/irksome/kronpc.py @@ -8,17 +8,19 @@ from firedrake.dmhooks import get_appctx, get_function_space -__all__ = ("KronPC","MassKronPC", "StiffnessKronPC", "SIPGStiffnessKronPC") +__all__ = ("KronPC", "MassKronPC", "StiffnessKronPC", "SIPGStiffnessKronPC") # ----------------------------- # KronPC definition # ----------------------------- + + class KronPC(PCBase): r""" Preconditioner applying: y = (L \otimes K^{-1}) x - L is the stage matrix, provided by the user. If A is provided, we set L = A^{-1}. If neither provided, L = I_s. - K is assembled on the single-stage space by subclasses via "form(trial, test)" - - K^{-1} is approximated by a PETSc PC with prefix + - K^{-1} is approximated by a PETSc PC with prefix """ needs_python_pmat = False @@ -41,20 +43,19 @@ def initialize(self, pc): Vbig = get_function_space(dm) - self.work_in = Cofunction(Vbig.dual(), name="kron_work_in") - self.work_mid = Function(Vbig, name="kron_work_mid") - self.work_out = Function(Vbig, name="kron_work_out") + self.work_in = Cofunction(Vbig.dual(), name="kron_work_in") + self.work_mid = Function(Vbig, name="kron_work_mid") + self.work_out = Function(Vbig, name="kron_work_out") Vstage = Vbig.sub(0) trial = TrialFunction(Vstage) - test = TestFunction(Vstage) + test = TestFunction(Vstage) a, bcs = self.form(trial, test) fc_params = getattr(context, "fc_params", None) K = assemble(a, bcs=bcs, mat_type=mat_type, form_compiler_parameters=fc_params) self.K = K - # NEW: create a sub-KSP (instead of a raw PC) so we can use python PCs with DM/appctx sub_ksp = PETSc.KSP().create(comm=pc.comm) sub_ksp.setOptionsPrefix(self._prefix + "sub_") @@ -85,9 +86,9 @@ def initialize(self, pc): sub_ksp.setUp() self.sub_ksp = sub_ksp - self.sub_pc = sub_ksp.getPC() + self.sub_pc = sub_ksp.getPC() A, _ = self.sub_ksp.getOperators() - self._tmp_stage = A.createVecRight() + self._tmp_stage = A.createVecRight() try: s = len(self.work_in.subfunctions) except Exception: @@ -95,7 +96,6 @@ def initialize(self, pc): self._s = s self.L = self._build_stage_L(self._s, context, pc) - # --- discover A from user-provided context (prefer appctx) --- def _build_stage_L(self, s, context, pc): """ @@ -143,7 +143,7 @@ def apply(self, pc, x, y): # Stagewise K^{-1} for i in range(s): with self.work_in.subfunctions[i].dat.vec_ro as rhs_i, \ - self.work_mid.subfunctions[i].dat.vec_wo as mid_i: + self.work_mid.subfunctions[i].dat.vec_wo as mid_i: mid_i.set(0.0) self.sub_ksp.solve(rhs_i, mid_i) @@ -151,7 +151,7 @@ def apply(self, pc, x, y): for j in range(s): self.work_out.subfunctions[j].assign(0.0) - # y_stage[j] += sum_i L[j,i] * mid[i] + # y_stage[j] += sum_i L[j,i] * mid[i] for j in range(s): row = self.L[j, :] with self.work_out.subfunctions[j].dat.vec_wo as yj: @@ -175,7 +175,7 @@ def applyTranspose(self, pc, x, y): # Stagewise K^{-T} for i in range(s): with self.work_in.subfunctions[i].dat.vec_ro as rhs_i, \ - self.work_mid.subfunctions[i].dat.vec_wo as mid_i: + self.work_mid.subfunctions[i].dat.vec_wo as mid_i: mid_i.set(0.0) # requires the sub-KSP (PC) to support transpose solves self.sub_ksp.solveTranspose(rhs_i, mid_i) @@ -198,7 +198,6 @@ def applyTranspose(self, pc, x, y): with self.work_out.dat.vec_ro as vout: vout.copy(y) - def view(self, pc, viewer=None): if viewer is None: viewer = PETSc.Viewer.STDOUT(pc.comm) @@ -212,10 +211,12 @@ def destroy(self, pc): if hasattr(self, "sub_ksp") and self.sub_ksp is not None: self.sub_ksp.destroy() self.sub_ksp = None - self.sub_pc = None + self.sub_pc = None + class MassKronPC(KronPC): """K built from the mass form.""" + def form(self, trial, test): a = inner(trial, test) * dx bcs = None @@ -224,11 +225,13 @@ def form(self, trial, test): class StiffnessKronPC(KronPC): """K built from the (regularized) stiffness form.""" + def form(self, trial, test): a = inner(grad(trial), grad(test)) * dx + 1e-12 * inner(trial, test) * dx bcs = None return a, bcs - + + class SIPGStiffnessKronPC(KronPC): r""" Discontinuous Galerkin (SIPG) pressure "stiffness" for DG spaces (e.g., P_{k}^{\mathrm{disc}}). @@ -239,8 +242,8 @@ class SIPGStiffnessKronPC(KronPC): def form(self, trial, test): mesh = trial.ufl_domain() - n = FacetNormal(mesh) - h = CellVolume(mesh) + n = FacetNormal(mesh) + h = CellVolume(mesh) # Polynomial degree k of the DG space k = trial.ufl_function_space().ufl_element().degree() @@ -257,13 +260,10 @@ def form(self, trial, test): - inner(grad(test), trial * n) * ds - inner(grad(trial), test * n) * ds + (penal / h) * inner(trial, test) * ds) - bcs = None return a, bcs - - - + def initialize(self, pc): # Assemble K and build the sub-KSP using the base class logic super().initialize(pc) @@ -281,8 +281,6 @@ def initialize(self, pc): A.setNullSpace(self._ns) P.setNullSpace(self._ns) - - def apply(self, pc, x, y): s = self._s y.set(0.0) @@ -294,7 +292,7 @@ def apply(self, pc, x, y): # stagewise K^{-1} with per-stage projections for i in range(s): with self.work_in.subfunctions[i].dat.vec_ro as rhs_i_ro, \ - self.work_mid.subfunctions[i].dat.vec_wo as mid_i: + self.work_mid.subfunctions[i].dat.vec_wo as mid_i: # tmp = self._tmp_stage # tmp.copy(rhs_i_ro) @@ -319,8 +317,5 @@ def apply(self, pc, x, y): with self.work_mid.subfunctions[i].dat.vec_ro as ui: yj.axpy(lij, ui) - with self.work_out.dat.vec_ro as vout: vout.copy(y) - - From 35fdc9fae344c0b08ef7151241e566d530eb5c9f Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 6 May 2026 12:22:35 -0500 Subject: [PATCH 32/35] Fix KronPC test lint --- tests/test_kronpc.py | 37 ++++++++++++++++--------------------- 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/tests/test_kronpc.py b/tests/test_kronpc.py index d3a53e51..5544d24c 100644 --- a/tests/test_kronpc.py +++ b/tests/test_kronpc.py @@ -5,27 +5,26 @@ from functools import reduce from operator import mul -# Important: import irksome before any Firedrake/UFL forms are built/processed. -# This lets Irksome register its UFL extensions before UFL MultiFunctions run. -import irksome +import irksome # noqa: F401 from firedrake import ( UnitSquareMesh, FunctionSpace, TrialFunction, TestFunction, Cofunction, Function, as_tensor, inner, dx, solve, errornorm ) + def getA(ns: int): - """nonsymmetric, invertible matrix.""" + """Nonsymmetric, invertible matrix.""" A = np.zeros((ns, ns)) for i in range(ns): for j in range(ns): - A[i, j] = 1.0 + 0.1*(i+1) - 0.2*(j+1) # breaks symmetry - A[i, i] += ns + A[i, j] = 1.0 + 0.1 * (i + 1) - 0.2 * (j + 1) + A[i, i] += ns return A def random_rhs(Vbig, seed: int = 1234): - """random RHS in the dual of Vbig. Yes, we need the dual space here!!""" + """Random RHS in the dual of Vbig.""" rng = np.random.default_rng(seed) L = Cofunction(Vbig.dual(), name="rhs") with L.dat.vec_wo as v: @@ -35,42 +34,36 @@ def random_rhs(Vbig, seed: int = 1234): return L -@pytest.mark.parametrize("ns", [2,3,4,5]) # meaningful cases are ns >= 2 +@pytest.mark.parametrize("ns", [2, 3, 4, 5]) # meaningful cases are ns >= 2 def test_mass_kron_pc(ns): """ Solve two ways: (1) direct LU on the full mixed operator (2) preonly with Python PC = MassKronPC and inner stage LU """ - # Create mesh and function-space msh = UnitSquareMesh(2, 2) V = FunctionSpace(msh, "CG", 1) - # Stage-stacked space V Vbig = reduce(mul, [V] * ns) - # Build global operator - A_np = getA(ns) - A = as_tensor(A_np) + A_np = getA(ns) + A = as_tensor(A_np) uu = TrialFunction(Vbig) vv = TestFunction(Vbig) a = inner(ufl.dot(A, uu), vv) * dx - # The RHS L = random_rhs(Vbig, seed=2025) - # Reference solution via full LU u_ref = Function(Vbig, name="u_ref") solve( a == L, u_ref, solver_parameters={ "ksp_type": "preonly", - "pc_type": "lu" + "pc_type": "lu", }, ) - # Exact inverse via our KronPC: u_pc = Function(Vbig, name="u_pc") solve( a == L, @@ -78,13 +71,15 @@ def test_mass_kron_pc(ns): solver_parameters={ "ksp_type": "preonly", "pc_type": "python", - "pc_python_type": "irksome.MassKronPC", + "pc_python_type": "irksome.MassKronPC", "mat_type": "matfree", - "kron_sub_pc_type": "lu" - + "kron_sub_pc_type": "lu", }, appctx={"A": A_np}, ) err = errornorm(u_ref, u_pc, norm_type="l2") - assert err < 1.0e-10, f"MassKronPC mismatch: ||u_ref - u_pc|| = {err:.3e} (ns={ns})" + assert err < 1.0e-10, ( + f"MassKronPC mismatch: ||u_ref - u_pc|| = {err:.3e} " + f"(ns={ns})" + ) From 212b0042b385a4aa5d207bc2ac0c8a5d895b3579 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Wed, 6 May 2026 14:42:09 -0500 Subject: [PATCH 33/35] Use Function.zero in kronpc --- irksome/kronpc.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/irksome/kronpc.py b/irksome/kronpc.py index 709eed1a..3ee7069d 100644 --- a/irksome/kronpc.py +++ b/irksome/kronpc.py @@ -148,8 +148,7 @@ def apply(self, pc, x, y): self.sub_ksp.solve(rhs_i, mid_i) # Zero outputs - for j in range(s): - self.work_out.subfunctions[j].assign(0.0) + self.work_out.zero() # y_stage[j] += sum_i L[j,i] * mid[i] for j in range(s): From 63592d884b5491448b5d972d9a6567ac3a0bd1f5 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Thu, 7 May 2026 10:02:44 -0500 Subject: [PATCH 34/35] Reduce KronPC test stage counts --- tests/test_kronpc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_kronpc.py b/tests/test_kronpc.py index 5544d24c..d4a44620 100644 --- a/tests/test_kronpc.py +++ b/tests/test_kronpc.py @@ -34,7 +34,7 @@ def random_rhs(Vbig, seed: int = 1234): return L -@pytest.mark.parametrize("ns", [2, 3, 4, 5]) # meaningful cases are ns >= 2 +@pytest.mark.parametrize("ns", [2, 3]) # meaningful cases are ns >= 2 def test_mass_kron_pc(ns): """ Solve two ways: From a2ccd860e1d65caa8d6f603152a8ac1847482bd0 Mon Sep 17 00:00:00 2001 From: 436ahsan <436ahsan@gmail.com> Date: Thu, 7 May 2026 12:04:19 -0500 Subject: [PATCH 35/35] Parametrize kronpc tests over mass and stiffness forms --- irksome/kronpc.py | 8 ++++-- tests/test_kronpc.py | 58 ++++++++++++++++++++++++++++++++------------ 2 files changed, 49 insertions(+), 17 deletions(-) diff --git a/irksome/kronpc.py b/irksome/kronpc.py index 3ee7069d..f10bb43a 100644 --- a/irksome/kronpc.py +++ b/irksome/kronpc.py @@ -223,10 +223,14 @@ def form(self, trial, test): class StiffnessKronPC(KronPC): - """K built from the (regularized) stiffness form.""" + """K built from the regularized stiffness form.""" def form(self, trial, test): - a = inner(grad(trial), grad(test)) * dx + 1e-12 * inner(trial, test) * dx + shift = PETSc.Options().getReal(self._prefix + "stiffness_shift", 1.0e-12) + a = ( + inner(grad(trial), grad(test)) * dx + + shift * inner(trial, test) * dx + ) bcs = None return a, bcs diff --git a/tests/test_kronpc.py b/tests/test_kronpc.py index d4a44620..c0d459ab 100644 --- a/tests/test_kronpc.py +++ b/tests/test_kronpc.py @@ -9,12 +9,12 @@ from firedrake import ( UnitSquareMesh, FunctionSpace, TrialFunction, TestFunction, Cofunction, - Function, as_tensor, inner, dx, solve, errornorm + Function, as_tensor, inner, grad, dx, solve, errornorm ) def getA(ns: int): - """Nonsymmetric, invertible matrix.""" + """Nonsymmetric, invertible stage matrix.""" A = np.zeros((ns, ns)) for i in range(ns): for j in range(ns): @@ -27,19 +27,42 @@ def random_rhs(Vbig, seed: int = 1234): """Random RHS in the dual of Vbig.""" rng = np.random.default_rng(seed) L = Cofunction(Vbig.dual(), name="rhs") + with L.dat.vec_wo as v: arr = rng.standard_normal(v.getSize()) v.setValues(range(v.getSize()), arr) v.assemble() + return L -@pytest.mark.parametrize("ns", [2, 3]) # meaningful cases are ns >= 2 -def test_mass_kron_pc(ns): +def mass_form(A, uu, vv, shift): + """Stage-coupled mass form corresponding to A kron M.""" + return inner(ufl.dot(A, uu), vv) * dx + + +def stiffness_form(A, uu, vv, shift): + """Stage-coupled shifted stiffness form corresponding to A kron K.""" + Au = ufl.dot(A, uu) + return ( + inner(grad(Au), grad(vv)) * dx + + shift * inner(Au, vv) * dx + ) + + +@pytest.mark.parametrize("ns", [2, 3]) +@pytest.mark.parametrize( + ("pc_python_type", "form_builder", "shift", "tol"), + [ + ("irksome.MassKronPC", mass_form, 0.0, 1.0e-10), + ("irksome.StiffnessKronPC", stiffness_form, 1.0, 1.0e-10), + ], +) +def test_kron_pc(ns, pc_python_type, form_builder, shift, tol): """ Solve two ways: (1) direct LU on the full mixed operator - (2) preonly with Python PC = MassKronPC and inner stage LU + (2) preonly with Python KronPC and inner stage LU """ msh = UnitSquareMesh(2, 2) V = FunctionSpace(msh, "CG", 1) @@ -50,7 +73,7 @@ def test_mass_kron_pc(ns): A = as_tensor(A_np) uu = TrialFunction(Vbig) vv = TestFunction(Vbig) - a = inner(ufl.dot(A, uu), vv) * dx + a = form_builder(A, uu, vv, shift) L = random_rhs(Vbig, seed=2025) @@ -64,22 +87,27 @@ def test_mass_kron_pc(ns): }, ) + solver_parameters = { + "ksp_type": "preonly", + "pc_type": "python", + "pc_python_type": pc_python_type, + "mat_type": "matfree", + "kron_sub_pc_type": "lu", + } + + if pc_python_type == "irksome.StiffnessKronPC": + solver_parameters["kron_stiffness_shift"] = shift + u_pc = Function(Vbig, name="u_pc") solve( a == L, u_pc, - solver_parameters={ - "ksp_type": "preonly", - "pc_type": "python", - "pc_python_type": "irksome.MassKronPC", - "mat_type": "matfree", - "kron_sub_pc_type": "lu", - }, + solver_parameters=solver_parameters, appctx={"A": A_np}, ) err = errornorm(u_ref, u_pc, norm_type="l2") - assert err < 1.0e-10, ( - f"MassKronPC mismatch: ||u_ref - u_pc|| = {err:.3e} " + assert err < tol, ( + f"{pc_python_type} mismatch: ||u_ref - u_pc|| = {err:.3e} " f"(ns={ns})" )