From a423c40049a42e948ba48bbcc31a5eed4957b60a Mon Sep 17 00:00:00 2001 From: Jared Callaham Date: Wed, 20 Mar 2024 04:54:33 -0400 Subject: [PATCH] Linear stability analysis functionality (#176) * Helmholtz eigenvalue problem example * [WIP] testing direct stability on cylinder * Link to old code with working SLEPc call * Derive linear control term * Simple test problems for KS algorithm * Arnoldi example in Firedrake with heat equation * Debugging KS in Firedrake * Working Krylov-Schur on heat equation * Add linearized NS solver * Add base flow forcing terms * Fixed Arnoldi iteration with Navier-Stokes mass matrix * Tweaking Arnoldi config * Testing shift-invert on heat eqution * Working cylinder stability with spectral transformation * Add spectral transformation script * Sorting ritz vals * Partial cleanup of cyl_st script * Cleaned up cyl_st. Save before making the analysis generic * Output formatting * Non-crashing stability analysis for cavity * Working shift-invert * Add general-purpose 'residual' method * Filtering duplicate eigenvalues * Working general functions with filtered eigenvalues * Shift-invert Arnoldi test on cavity * [WIP] adding CLI args to cyl script * Clean up examples/cylinder/stability/ * Combined Arnoldi and K-S into one function * Move operatorto __matmul__ * Pre-assemble linear operator * Separate code for linear stability and linear operators * Cavity stability analysis * Fix render method for cavity and step * Move common functionality to utils.stability * Black * yapf * Fix codespell * Undo re-running gmsh: * Remove examples/cavity/stability/ * Final cleanup of examples/ * codespell complaint --- examples/cavity/stability.py | 179 +++++++++++ examples/cylinder/lti_system.py | 65 ++++ examples/cylinder/pressure-probes.py | 1 + examples/cylinder/stability.py | 166 +++++++++++ hydrogym/firedrake/__init__.py | 2 +- hydrogym/firedrake/envs/cavity/flow.py | 30 +- hydrogym/firedrake/envs/cylinder/flow.py | 11 +- hydrogym/firedrake/envs/pinball/flow.py | 11 +- hydrogym/firedrake/envs/step/flow.py | 58 +++- hydrogym/firedrake/flow.py | 251 +++++++++++++++- hydrogym/firedrake/solvers/__init__.py | 3 +- hydrogym/firedrake/solvers/base.py | 16 +- hydrogym/firedrake/solvers/bdf_ext.py | 143 ++++++--- hydrogym/firedrake/solvers/integrate.py | 3 +- hydrogym/firedrake/solvers/stabilization.py | 42 +++ hydrogym/firedrake/utils/__init__.py | 2 + hydrogym/firedrake/utils/eig.py | 310 ++++++++++++++++++++ hydrogym/firedrake/utils/linalg.py | 150 ++++++---- hydrogym/firedrake/utils/stability.py | 125 ++++++++ 19 files changed, 1417 insertions(+), 151 deletions(-) create mode 100644 examples/cavity/stability.py create mode 100644 examples/cylinder/lti_system.py create mode 100644 examples/cylinder/stability.py create mode 100644 hydrogym/firedrake/utils/eig.py create mode 100644 hydrogym/firedrake/utils/stability.py diff --git a/examples/cavity/stability.py b/examples/cavity/stability.py new file mode 100644 index 0000000..b378940 --- /dev/null +++ b/examples/cavity/stability.py @@ -0,0 +1,179 @@ +""" +To get the leading mode, run with: + +``` +mpiexec -np 20 python stability.py --output-dir eig_output --sigma 1.0+11j +``` +""" + +import argparse +import os + +# from cyl_common import base_checkpoint, evec_checkpoint, flow +from typing import NamedTuple + +import firedrake as fd +import numpy as np + +import hydrogym.firedrake as hgym + +parser = argparse.ArgumentParser( + description="Stability analysis of the open cavity flow.") +parser.add_argument( + "--mesh", + default="fine", + type=str, + help="Identifier for the mesh resolution", +) +parser.add_argument( + "--reynolds", + default=7500.0, + type=float, + help="Reynolds number of the flow", +) +parser.add_argument( + "--krylov-dim", + default=100, + type=int, + help="Dimension of the Krylov subspace (number of Arnoldi vectors)", +) +parser.add_argument( + "--tol", + default=1e-10, + type=float, + help="Tolerance to use for determining converged eigenvalues.", +) +parser.add_argument( + "--schur", + action="store_true", + dest="schur", + default=False, + help="Use Krylov-Schur iteration to restart the Arnoldi process.", +) +parser.add_argument( + "--no-adjoint", + action="store_true", + default=False, + help="Skip computing the adjoint modes.", +) +parser.add_argument( + "--sigma", + default=0.0, + type=complex, + help="Shift for the shift-invert Arnoldi method.", +) +parser.add_argument( + "--base-flow", + type=str, + help="Path to the HDF5 checkpoint containing the base flow.", +) +parser.add_argument( + "--output-dir", + type=str, + default="eig_output", + help="Directory in which output files will be stored.", +) +if __name__ == "__main__": + args = parser.parse_args() + + mesh = args.mesh + m = args.krylov_dim + sigma = args.sigma + tol = args.tol + Re = args.reynolds + + output_dir = args.output_dir + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + velocity_order = 2 + stabilization = "none" + + flow = hgym.Cavity( + Re=Re, + mesh=mesh, + velocity_order=velocity_order, + ) + + dof = flow.mixed_space.dim() + + hgym.print("|--------------------------------------------------|") + hgym.print("| Linear stability analysis of the open cavity flow |") + hgym.print("|--------------------------------------------------|") + hgym.print(f"Reynolds number: {Re}") + hgym.print(f"Krylov dimension: {m}") + hgym.print(f"Spectral shift: {sigma}") + hgym.print(f"Include adjoint modes: {not args.no_adjoint}") + hgym.print(f"Krylov-Schur restart: {args.schur}") + hgym.print(f"Total dof: {dof} --- dof/rank: {int(dof/fd.COMM_WORLD.size)}") + hgym.print("") + + if args.base_flow: + hgym.print(f"Loading base flow from checkpoint {args.base_flow}...") + flow.load_checkpoint(args.base_flow) + + else: + + # Since this flow is at high Reynolds number we have to + # ramp to get the steady state + hgym.print("Solving the steady-state problem...") + + steady_solver = hgym.NewtonSolver( + flow, + stabilization=stabilization, + solver_parameters={"snes_monitor": None}) + Re_init = [500, 1000, 2000, 4000, Re] + + for i, Re in enumerate(Re_init): + flow.Re.assign(Re) + hgym.print(f"Steady solve at Re={Re_init[i]}") + qB = steady_solver.solve() + + flow.save_checkpoint(f"{args.output_dir}/base.h5") + + hgym.print("Computing direct modes...") + dir_results = hgym.utils.stability_analysis( + flow, sigma, m, tol, schur_restart=args.schur, adjoint=False) + np.save(f"{output_dir}/raw_evals", dir_results.raw_evals) + + # Save checkpoints + evals = dir_results.evals + eval_data = np.zeros((len(evals), 3), dtype=np.float64) + eval_data[:, 0] = evals.real + eval_data[:, 1] = evals.imag + eval_data[:, 2] = dir_results.residuals + np.save(f"{output_dir}/evals", eval_data) + + with fd.CheckpointFile(f"{output_dir}/evecs.h5", "w") as chk: + for i in range(len(evals)): + chk.save_mesh(flow.mesh) + dir_results.evecs_real[i].rename(f"evec_{i}_re") + chk.save_function(dir_results.evecs_real[i]) + dir_results.evecs_imag[i].rename(f"evec_{i}_im") + chk.save_function(dir_results.evecs_imag[i]) + + if not args.no_adjoint: + hgym.print("Computing adjoint modes...") + adj_results = hgym.utils.stability_analysis( + flow, sigma, m, tol, adjoint=True) + + # Save checkpoints + evals = adj_results.evals + eval_data = np.zeros((len(evals), 3), dtype=np.float64) + eval_data[:, 0] = evals.real + eval_data[:, 1] = evals.imag + eval_data[:, 2] = adj_results.residuals + np.save(f"{output_dir}/adj_evals", eval_data) + + with fd.CheckpointFile(f"{output_dir}/adj_evecs.h5", "w") as chk: + for i in range(len(evals)): + chk.save_mesh(flow.mesh) + adj_results.evecs_real[i].rename(f"evec_{i}_re") + chk.save_function(adj_results.evecs_real[i]) + adj_results.evecs_imag[i].rename(f"evec_{i}_im") + chk.save_function(adj_results.evecs_imag[i]) + + hgym.print( + "NOTE: If there is a warning following this, ignore it. It is raised by PETSc " + "CLI argument handling and not argparse. It does not indicate that any CLI " + "arguments are ignored by this script.") diff --git a/examples/cylinder/lti_system.py b/examples/cylinder/lti_system.py new file mode 100644 index 0000000..2c1e56c --- /dev/null +++ b/examples/cylinder/lti_system.py @@ -0,0 +1,65 @@ +"""Illustrate constructing an LTI system from a flow configuration. + +This is a work in progress - so far it's just the base flow and the control +vector. Ultimately it will need LinearOperator functionality to have the form + +```` +x' = A*x + B*u +y = C*x + D*u +``` + +where `x` is a firedrake.Function and `u` and `y` are numpy arrays. +""" +import os + +import firedrake as fd +import matplotlib.pyplot as plt +from firedrake.pyplot import tripcolor +from ufl import dx, inner + +import hydrogym.firedrake as hgym + +show_plots = False +output_dir = "output" +if not os.path.exists(output_dir): + os.makedirs(output_dir) + +flow = hgym.RotaryCylinder(Re=100, mesh="medium") + +# 1. Compute base flow +steady_solver = hgym.NewtonSolver(flow) +qB = steady_solver.solve() +qB.rename("qB") + +# Check lift/drag +print(flow.compute_forces(qB)) + +if show_plots: + vort = flow.vorticity(qB.subfunctions[0]) + fig, ax = plt.subplots(1, 1, figsize=(6, 3)) + tripcolor(vort, axes=ax, cmap="RdBu", vmin=-2, vmax=2) + +# 2. Derive flow field associated with actuation BC +# See Barbagallo et al. (2009) for details on the "lifting" procedure +F = steady_solver.steady_form(qB) # Nonlinear variational form +J = fd.derivative(F, qB) # Jacobian with automatic differentiation + +flow.linearize_bcs(mixed=True) +flow.set_control([1.0]) +bcs = flow.collect_bcs() + +# Solve steady, inhomogeneous problem +qC = fd.Function(flow.mixed_space, name="qC") +v, s = fd.TestFunctions(flow.mixed_space) +zero = inner(fd.Constant((0.0, 0.0)), v) * dx +fd.solve(J == zero, qC, bcs=bcs) + +if show_plots: + vort = flow.vorticity(qC.subfunctions[0]) + fig, ax = plt.subplots(1, 1, figsize=(6, 3)) + tripcolor(vort, axes=ax, cmap="RdBu", vmin=-2, vmax=2) + +with fd.CheckpointFile(f"{output_dir}/lin_fields.h5", "w") as chk: + chk.save_mesh(flow.mesh) + chk.save_function(qB) + chk.save_function(qC) diff --git a/examples/cylinder/pressure-probes.py b/examples/cylinder/pressure-probes.py index 70abf8a..1c98b54 100644 --- a/examples/cylinder/pressure-probes.py +++ b/examples/cylinder/pressure-probes.py @@ -1,6 +1,7 @@ """ Simulate the flow around the cylinder with evenly spaced surface pressure probes. """ + import os import matplotlib.pyplot as plt import numpy as np diff --git a/examples/cylinder/stability.py b/examples/cylinder/stability.py new file mode 100644 index 0000000..eedf933 --- /dev/null +++ b/examples/cylinder/stability.py @@ -0,0 +1,166 @@ +import argparse +import os + +import firedrake as fd +import numpy as np + +import hydrogym.firedrake as hgym + +parser = argparse.ArgumentParser( + description="Stability analysis of the Re=100 cylinder wake.") +parser.add_argument( + "--mesh", + default="medium", + type=str, + help='Identifier for the mesh resolution. Options: ["medium", "fine"]', +) +parser.add_argument( + "--reynolds", + default=100.0, + type=float, + help="Reynolds number of the flow", +) +parser.add_argument( + "--krylov-dim", + default=100, + type=int, + help="Dimension of the Krylov subspace (number of Arnoldi vectors)", +) +parser.add_argument( + "--tol", + default=1e-10, + type=float, + help="Tolerance to use for determining converged eigenvalues.", +) +parser.add_argument( + "--schur", + action="store_true", + dest="schur", + default=False, + help="Use Krylov-Schur iteration to restart the Arnoldi process.", +) +parser.add_argument( + "--no-adjoint", + action="store_true", + default=False, + help="Skip computing the adjoint modes.", +) +parser.add_argument( + "--sigma", + default=0.0, + type=complex, + help="Shift for the shift-invert Arnoldi method.", +) +parser.add_argument( + "--base-flow", + type=str, + help="Path to the HDF5 checkpoint containing the base flow.", +) +parser.add_argument( + "--output-dir", + type=str, + default="eig_output", + help="Directory in which output files will be stored.", +) +if __name__ == "__main__": + args = parser.parse_args() + + mesh = args.mesh + m = args.krylov_dim + sigma = args.sigma + tol = args.tol + Re = args.reynolds + + output_dir = args.output_dir + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + velocity_order = 2 + stabilization = "none" + + flow = hgym.Cylinder( + Re=Re, + mesh=mesh, + velocity_order=velocity_order, + ) + + hgym.print("|------------------------------------------------|") + hgym.print("| Linear stability analysis of the cylinder wake |") + hgym.print("|------------------------------------------------|") + hgym.print(f"Reynolds number: {Re}") + hgym.print(f"Krylov dimension: {m}") + hgym.print(f"Spectral shift: {sigma}") + hgym.print(f"Include adjoint modes: {not args.no_adjoint}") + hgym.print(f"Krylov-Schur restart: {args.schur}") + hgym.print("") + + if args.base_flow: + hgym.print(f"Loading base flow from checkpoint {args.base_flow}...") + flow.load_checkpoint(args.base_flow) + + else: + hgym.print("Solving the steady-state problem for the cylinder base flow...") + + steady_solver = hgym.NewtonSolver( + flow, + stabilization=stabilization, + solver_parameters={"snes_monitor": None}) + if Re > 50: + hgym.print("Solving steady-state problem at Re=50...") + flow.Re.assign(50) + steady_solver.solve() + + hgym.print( + f"Solving steady-state problem at target Reynolds number Re={Re}...") + flow.Re.assign(Re) + steady_solver.solve() + CL, CD = flow.compute_forces() + hgym.print(f"Lift: {CL:0.3e}, Drag: {CD:0.3e}") + flow.save_checkpoint(f"{args.output_dir}/base.h5") + + hgym.print("Computing direct modes...") + dir_results = hgym.utils.stability_analysis( + flow, sigma, m, tol, schur_restart=args.schur, adjoint=False) + np.save(f"{output_dir}/raw_evals", dir_results.raw_evals) + + # Save checkpoints + evals = dir_results.evals + eval_data = np.zeros((len(evals), 3), dtype=np.float64) + eval_data[:, 0] = evals.real + eval_data[:, 1] = evals.imag + eval_data[:, 2] = dir_results.residuals + np.save(f"{output_dir}/evals", eval_data) + + with fd.CheckpointFile(f"{output_dir}/evecs.h5", "w") as chk: + for i in range(len(evals)): + chk.save_mesh(flow.mesh) + dir_results.evecs_real[i].rename(f"evec_{i}_re") + chk.save_function(dir_results.evecs_real[i]) + dir_results.evecs_imag[i].rename(f"evec_{i}_im") + chk.save_function(dir_results.evecs_imag[i]) + + if not args.no_adjoint: + hgym.print("Computing adjoint modes...") + adj_results = hgym.utils.stability_analysis( + flow, sigma, m, tol, adjoint=True) + + # Save checkpoints + evals = adj_results.evals + eval_data = np.zeros((len(evals), 3), dtype=np.float64) + eval_data[:, 0] = evals.real + eval_data[:, 1] = evals.imag + eval_data[:, 2] = adj_results.residuals + np.save(f"{output_dir}/adj_evals", eval_data) + + with fd.CheckpointFile(f"{output_dir}/adj_evecs.h5", "w") as chk: + for i in range(len(evals)): + chk.save_mesh(flow.mesh) + adj_results.evecs_real[i].rename(f"evec_{i}_re") + chk.save_function(adj_results.evecs_real[i]) + adj_results.evecs_imag[i].rename(f"evec_{i}_im") + chk.save_function(adj_results.evecs_imag[i]) + + hgym.print( + "NOTE: If there is a warning following this, ignore it. It is raised by PETSc " + "CLI argument handling and not argparse. It does not indicate that any CLI " + "arguments are ignored by this script.") diff --git a/hydrogym/firedrake/__init__.py b/hydrogym/firedrake/__init__.py index 347a1aa..14d1260 100644 --- a/hydrogym/firedrake/__init__.py +++ b/hydrogym/firedrake/__init__.py @@ -2,7 +2,7 @@ from .actuator import DampedActuator from .flow import FlowConfig, ObservationFunction, ScaledDirichletBC -from .solvers import NewtonSolver, SemiImplicitBDF, integrate +from .solvers import LinearizedBDF, NewtonSolver, SemiImplicitBDF, integrate from .utils import io, is_rank_zero, linalg, modeling, print from .envs import Cylinder, RotaryCylinder, Pinball, Cavity, Step # isort:skip diff --git a/hydrogym/firedrake/envs/cavity/flow.py b/hydrogym/firedrake/envs/cavity/flow.py index b92d482..04c1b65 100644 --- a/hydrogym/firedrake/envs/cavity/flow.py +++ b/hydrogym/firedrake/envs/cavity/flow.py @@ -55,8 +55,11 @@ def configure_observations(self, return supported_obs_types[obs_type] - def init_bcs(self): - V, Q = self.function_spaces(mixed=True) + def init_bcs(self, function_spaces=None): + if function_spaces is None: + V, Q = self.function_spaces(mixed=True) + else: + V, Q = function_spaces # Define static boundary conditions self.U_inf = fd.Constant((1.0, 0.0)) @@ -88,9 +91,9 @@ def collect_bcu(self): def collect_bcp(self): return [self.bcp_outflow] - def linearize_bcs(self): + def linearize_bcs(self, function_spaces=None): self.reset_controls() - self.init_bcs() + self.init_bcs(function_spaces=function_spaces) self.bcu_inflow.set_value(fd.Constant((0, 0))) def wall_stress_sensor(self, q=None): @@ -111,22 +114,21 @@ def evaluate_objective(self, q=None, qB=None): KE = 0.5 * fd.assemble(fd.inner(u - uB, u - uB) * fd.dx) return KE - # TODO: Rendering function needs to be revisited as this is only a hot-fix def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): + _fig, ax = plt.subplots(1, 1, figsize=(6, 3)) if clim is None: - clim = (-2, 2) + clim = (-10, 10) if levels is None: - levels = np.linspace(*clim, 10) - vort = fd.project(fd.curl(self.u), self.pressure_space) - im = tricontourf( - vort, - cmap=cmap, + levels = np.linspace(*clim, 20) + tricontourf( + self.vorticity(), levels=levels, vmin=clim[0], vmax=clim[1], extend="both", + cmap=cmap, **kwargs, ) - - cyl = plt.Circle((0, 0), 0.5, edgecolor="k", facecolor="gray") - im.axes.add_artist(cyl) + ax.set_xlim([-0.5, 2.5]) + ax.set_ylim([-1, 0.5]) + ax.set_facecolor("grey") diff --git a/hydrogym/firedrake/envs/cylinder/flow.py b/hydrogym/firedrake/envs/cylinder/flow.py index c7d91bb..fd02b10 100644 --- a/hydrogym/firedrake/envs/cylinder/flow.py +++ b/hydrogym/firedrake/envs/cylinder/flow.py @@ -63,8 +63,11 @@ def configure_observations(self, return supported_obs_types[obs_type] - def init_bcs(self): - V, Q = self.function_spaces(mixed=True) + def init_bcs(self, function_spaces=None): + if function_spaces is None: + V, Q = self.function_spaces(mixed=True) + else: + V, Q = function_spaces # Define the static boundary conditions self.U_inf = fd.Constant((1.0, 0.0)) @@ -166,8 +169,8 @@ def shear_force(self, q: fd.Function = None) -> float: # fd.solve(A == zero, f, bcs=self.collect_bcs()) # return f - def linearize_bcs(self): - self.reset_controls() + def linearize_bcs(self, function_spaces=None): + self.reset_controls(function_spaces=function_spaces) self.bcu_inflow.set_value(fd.Constant((0, 0))) self.bcu_freestream.set_value(fd.Constant(0.0)) diff --git a/hydrogym/firedrake/envs/pinball/flow.py b/hydrogym/firedrake/envs/pinball/flow.py index b96d06b..888d837 100644 --- a/hydrogym/firedrake/envs/pinball/flow.py +++ b/hydrogym/firedrake/envs/pinball/flow.py @@ -32,8 +32,11 @@ class Pinball(FlowConfig): MESH_DIR = os.path.abspath(f"{__file__}/..") - def init_bcs(self): - V, Q = self.function_spaces(mixed=True) + def init_bcs(self, function_spaces=None): + if function_spaces is None: + V, Q = self.function_spaces(mixed=True) + else: + V, Q = function_spaces # Define the static boundary conditions self.U_inf = fd.Constant((1.0, 0.0)) @@ -100,8 +103,9 @@ def compute_forces(self, q: fd.Function = None) -> Iterable[float]: CD = [fd.assemble(2 * force[0] * ds(cyl)) for cyl in self.CYLINDER] return CL, CD - def linearize_bcs(self): + def linearize_bcs(self, function_spaces=None): self.reset_controls() + self.init_bcs(function_spaces=function_spaces) self.bcu_inflow.set_value(fd.Constant((0, 0))) self.bcu_freestream.set_value(fd.Constant(0.0)) @@ -113,7 +117,6 @@ def evaluate_objective(self, q=None): CL, CD = self.compute_forces(q=q) return sum(CD) - # TODO: Needs to be revisited as the self calls here look hella suss def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): if clim is None: clim = (-2, 2) diff --git a/hydrogym/firedrake/envs/step/flow.py b/hydrogym/firedrake/envs/step/flow.py index 13a268a..ef47f87 100644 --- a/hydrogym/firedrake/envs/step/flow.py +++ b/hydrogym/firedrake/envs/step/flow.py @@ -12,6 +12,19 @@ class Step(FlowConfig): + """Backwards-facing step + + Notes on meshes: + - "medium" - outlet at L=25 (110k elements) + - "fine" - outlet at L=25 (223k elements) + - "m1" - "medium" resolution with outlet at L=10 (66k elements) + - "m2" - "medium" resolution with outlet at L=15 (81k elements) + - "m3" - "medium" resolution with outlet at L=20 (95k elements) + - "m4" - "fine" resolution with outlet at L=10 (130k elements) + - "m5" - "fine" resolution with outlet at L=15 (162k elements) + - "m6" - "fine" resolution with outlet at L=20 (193k elements) + """ + DEFAULT_REYNOLDS = 600 DEFAULT_MESH = "fine" DEFAULT_DT = 1e-2 @@ -78,8 +91,11 @@ def configure_observations(self, return supported_obs_types[obs_type] - def init_bcs(self): - V, Q = self.function_spaces(mixed=True) + def init_bcs(self, function_spaces=None): + if function_spaces is None: + V, Q = self.function_spaces(mixed=True) + else: + V, Q = function_spaces # Define static boundary conditions self.U_inf = ufl.as_tensor( @@ -111,9 +127,9 @@ def advance_time(self, dt, control=None): return super().advance_time(dt, control) - def linearize_bcs(self): + def linearize_bcs(self, function_spaces=None): self.reset_controls() - self.init_bcs() + self.init_bcs(function_spaces=function_spaces) self.bcu_inflow.set_value(fd.Constant((0, 0))) def collect_bcu(self): @@ -144,22 +160,34 @@ def evaluate_objective(self, q=None, qB=None): KE = 0.5 * fd.assemble(fd.inner(u - uB, u - uB) * fd.dx) return KE - # TODO: Rendering function needs to be revisited as this is only a hot-fix - def render(self, mode="human", clim=None, levels=None, cmap="RdBu", **kwargs): + def render( + self, + mode="human", + axes=None, + clim=None, + levels=None, + cmap="RdBu", + xlim=None, + **kwargs, + ): if clim is None: - clim = (-2, 2) + clim = (-5, 5) + if xlim is None: + xlim = [-2, 10] if levels is None: - levels = np.linspace(*clim, 10) - vort = fd.project(fd.curl(self.u), self.pressure_space) - im = tricontourf( - vort, - cmap=cmap, + levels = np.linspace(*clim, 20) + if axes is None: + _fix, axes = plt.subplots(1, 1, figsize=(12, 2)) + tricontourf( + self.vorticity(), levels=levels, vmin=clim[0], vmax=clim[1], extend="both", + cmap=cmap, + axes=axes, **kwargs, ) - - cyl = plt.Circle((0, 0), 0.5, edgecolor="k", facecolor="gray") - im.axes.add_artist(cyl) + axes.set_xlim(*xlim) + axes.set_ylim([-0.5, 0.5]) + axes.set_facecolor("grey") diff --git a/hydrogym/firedrake/flow.py b/hydrogym/firedrake/flow.py index fc6cbe5..5595362 100644 --- a/hydrogym/firedrake/flow.py +++ b/hydrogym/firedrake/flow.py @@ -9,11 +9,11 @@ from firedrake.__future__ import interpolate from mpi4py import MPI from numpy.typing import ArrayLike -from ufl import curl, dot, inner, nabla_grad, sqrt, sym +from ufl import curl, div, dot, inner, nabla_grad, sqrt, sym from ..core import ActuatorBase, PDEBase from .actuator import DampedActuator -from .utils import print +from .utils.linalg import DirectOperator, InverseOperator class ScaledDirichletBC(fd.DirichletBC): @@ -102,8 +102,8 @@ def load_checkpoint(self, filename: str, idx=None, read_mesh=True): # If the checkpoint saved on a different function space, # approximate the same field on the current function space # by projecting the checkpoint field onto the current space - V_chk = f_load.function_space().ufl_element() - V_self = f_self.function_space().ufl_element() + V_chk = f_load.function_space() + V_self = f_self.function_space() if V_chk.ufl_element() != V_self.ufl_element(): f_load = fd.project(f_load, V_self) @@ -173,7 +173,7 @@ def create_actuator(self, tau=None) -> ActuatorBase: tau = self.TAU return DampedActuator(1 / tau) - def reset_controls(self): + def reset_controls(self, function_spaces=None): """Reset the controls to a zero state Note that this is broken out from `reset` because @@ -183,7 +183,7 @@ def reset_controls(self): TODO: Allow for different kinds of actuators """ self.actuators = [self.create_actuator() for _ in range(self.num_inputs)] - self.init_bcs() + self.init_bcs(function_spaces=function_spaces) @property def nu(self): @@ -250,6 +250,30 @@ def sigma(self, u, p) -> ufl.Form: """Newtonian stress tensor""" return 2 * self.nu * self.epsilon(u) - p * fd.Identity(len(u)) + def residual(self, q, q_test=None): + """Nonlinear residual for the incompressible Navier-Stokes equations. + + Returns a UFL form F(u, p, v, s) = 0, where (u, p) is the trial function + and (v, s) is the test function. This residual is also the right-hand side + of the unsteady equations. + + A linearized form can be constructed by calling: + ``` + F = flow.residual((uB, pB), (v, s)) + J = fd.derivative(F, qB, q_trial) + ``` + """ + (u, p) = q + if q_test is None: + (v, s) = fd.TestFunctions(self.mixed_space) + else: + (v, s) = q_test + + sigma, epsilon = self.sigma, self.epsilon + F = (-inner(dot(u, nabla_grad(u)), v) * dx - + inner(sigma(u, p), epsilon(v)) * dx + inner(div(u), s) * dx) + return F + @pyadjoint.no_annotations def max_cfl(self, dt) -> float: """Estimate of maximum CFL number""" @@ -284,11 +308,32 @@ def set_control(self, act: ArrayLike = None): self.MAX_CONTROL) self.bcu_actuation[i].set_scale(u) - def dot(self, q1: fd.Function, q2: fd.Function) -> float: - """Energy inner product between two fields""" - u1 = q1.subfunctions[0] - u2 = q2.subfunctions[0] - return fd.assemble(inner(u1, u2) * dx) + def inner_product( + self, + q1: fd.Function, + q2: fd.Function, + assemble=True, + augmented=False, + ): + """Energy inner product for the Navier-Stokes equations. + + `augmented` is used to specify whether the function space is + extended to represent complex numbers. In this case the inner + product is the L2 norm of the real and imaginary parts. + """ + if not augmented: + (u, _) = fd.split(q1) + (v, _) = fd.split(q2) + M = inner(u, v) * dx + + else: + (u_re, _, u_im, _) = fd.split(q1) + (v_re, _, v_im, _) = fd.split(q2) + M = (inner(u_re, v_re) + inner(u_im, v_im)) * dx + + if not assemble: + return M + return fd.assemble(M) def velocity_probe(self, probes, q: fd.Function = None) -> list[float]: """Probe velocity in the wake. @@ -316,3 +361,187 @@ def vorticity_probe(self, probes, q: fd.Function = None) -> list[float]: u = q.subfunctions[0] vort = self.vorticity(u=u) return np.stack(vort.at(probes)) + + def linearize(self, + qB=None, + adjoint=False, + sigma=0.0, + inverse=False, + solver_parameters=None): + if sigma != 0.0 and not inverse: + raise ValueError("Must use `inverse=True` with spectral shift.") + + if qB is None: + qB = self.q + + if not inverse: + return self._jacobian_operator(qB, adjoint=adjoint) + + if sigma.imag == 0.0: + return self._real_shift_inv_operator( + qB, sigma, adjoint=adjoint, solver_parameters=solver_parameters) + + return self._complex_shift_inv_operator( + qB, sigma, adjoint=adjoint, solver_parameters=solver_parameters) + + # TODO: Test this. This is just here as an extension of the work done with + # shift-inverse-type operators, but hasn't been directly tested yet. Really + # it should be used for linearized time-stepping. + def _jacobian_operator(self, qB, adjoint=False): + """Construct the Jacobian operator for the Navier-Stokes equations. + + A matrix-vector product for this operator is equivalent to evaluating + the Jacobian of the Navier-Stokes equations at a given state. + """ + # Set up function spaces + fn_space = self.mixed_space + (uB, pB) = fd.split(qB) + q_trial = fd.TrialFunction(fn_space) + q_test = fd.TestFunction(fn_space) + (v, s) = fd.split(q_test) + + # Collect boundary conditions + self.linearize_bcs() + bcs = self.collect_bcs() + + # Linear form expressing the RHS of the Navier-Stokes without time derivative + # For a steady solution this is F(qB) = 0. + F = self.residual((uB, pB), q_test=(v, s)) + + # The Jacobian of F is the bilinear form J(qB, q_test) = dF/dq(qB) @ q_test + J = fd.derivative(F, qB, q_trial) + + A = DirectOperator(J, bcs, fn_space) + + if adjoint: + A = A.T + + return A + + def _real_shift_inv_operator(self, + qB, + sigma, + adjoint=False, + solver_parameters=None): + """Construct a shift-inverse Arnoldi iterator with real (or zero) shift. + + The shift-inverse iteration solves the matrix pencil + (J - sigma * M) @ v1 = M @ v0 for v1, where J is the Jacobian of the + Navier-Stokes equations, and M is the mass matrix. + """ + fn_space = self.mixed_space + (uB, pB) = fd.split(qB) + q_trial = fd.TrialFunction(fn_space) + q_test = fd.TestFunction(fn_space) + (v, s) = fd.split(q_test) + + # Collect boundary conditions + self.linearize_bcs() + bcs = self.collect_bcs() + + def M(q): + """Mass matrix for the Navier-Stokes equations""" + return self.inner_product(q, q_test, assemble=False) + + # Linear form expressing the RHS of the Navier-Stokes without time derivative + # For a steady solution this is F(qB) = 0. + F = self.residual((uB, pB), q_test=(v, s)) + + if sigma != 0.0: + F -= sigma * M(qB) + + # The Jacobian of F is the bilinear form J(qB, q_test) = dF/dq(qB) @ q_test + J = fd.derivative(F, qB, q_trial) + + A = InverseOperator(J, M, bcs, fn_space, solver_parameters) + + if adjoint: + A = A.T + + return A + + def _complex_shift_inv_operator(self, + qB, + sigma, + adjoint=False, + solver_parameters=None): + """Construct a shift-inverse Arnoldi iterator with complex-valued shift. + + The shifted operator is `A = (J - sigma * M)` + + For sigma = (sr, si), the real and imaginary parts of A are + A = (J - sr * M, -si * M) + The system solve is A @ v1 = M @ v0, so for complex vectors v = (vr, vi): + (Ar + 1j * Ai) @ (v1r + 1j * v1i) = M @ (v0r + 1j * v0i) + Since we can't do complex analysis without re-building PETSc, instead we treat this + as a block system: + ``` + [Ar, Ai] [v1r] = [M 0] [v0r] + [-Ai, Ar] [v1i] [0 M] [v0i] + ``` + + Note that this will be more expensive than real- or zero-shifted iteration, + since there are twice as many degrees of freedom. However, it will tend to + converge faster for the eigenvalues of interest. + + The shift-inverse iteration solves the matrix pencil + (J - sigma * M) @ v1 = M @ v0 for v1, where J is the Jacobian of the + Navier-Stokes equations, and M is the mass matrix. + """ + fn_space = self.mixed_space + W = fn_space * fn_space + V1, Q1, V2, Q2 = W + + # Set the boundary conditions for each function space + # These will be identical + self.linearize_bcs(function_spaces=(V1, Q1)) + bcs1 = self.collect_bcs() + + self.linearize_bcs(function_spaces=(V2, Q2)) + bcs2 = self.collect_bcs() + + bcs = [*bcs1, *bcs2] + + # Since the base flow is used to construct the Navier-Stokes Jacobian + # which is used on the diagonal block for both real and imaginary components, + # we have to duplicate the base flow for both components. This does NOT + # mean that the base flow literally has an imaginary component + qB_aug = fd.Function(W) + uB_re, pB_re, uB_im, pB_im = fd.split(qB_aug) + qB_aug.subfunctions[0].interpolate(qB.subfunctions[0]) + qB_aug.subfunctions[1].interpolate(qB.subfunctions[1]) + qB_aug.subfunctions[2].interpolate(qB.subfunctions[0]) + qB_aug.subfunctions[3].interpolate(qB.subfunctions[1]) + + # Create trial and test functions + q_trial = fd.TrialFunction(W) + q_test = fd.TestFunction(W) + (v_re, s_re, v_im, s_im) = fd.split(q_test) + + # Construct the nonlinear residual for the Navier-Stokes equations + F_re = self.residual((uB_re, pB_re), q_test=(v_re, s_re)) + F_im = self.residual((uB_im, pB_im), q_test=(v_im, s_im)) + + def _inner(u, v): + return ufl.inner(u, v) * ufl.dx + + # Shift each block of the linear form appropriately + F11 = F_re - sigma.real * _inner(uB_re, v_re) + F22 = F_im - sigma.real * _inner(uB_im, v_im) + F12 = sigma.imag * _inner(uB_im, v_re) + F21 = -sigma.imag * _inner(uB_re, v_im) + F = F11 + F22 + F12 + F21 + + # Differentiate to get the bilinear form for the Jacobian + J = fd.derivative(F, qB_aug, q_trial) + + def M(q): + """Mass matrix for the Navier-Stokes equations""" + return self.inner_product(q, q_test, assemble=False, augmented=True) + + A = InverseOperator(J, M, bcs, W, solver_parameters) + + if adjoint: + A = A.T + + return A diff --git a/hydrogym/firedrake/solvers/__init__.py b/hydrogym/firedrake/solvers/__init__.py index df67d8a..3e1f98b 100644 --- a/hydrogym/firedrake/solvers/__init__.py +++ b/hydrogym/firedrake/solvers/__init__.py @@ -1,9 +1,10 @@ from .base import NewtonSolver -from .bdf_ext import SemiImplicitBDF +from .bdf_ext import LinearizedBDF, SemiImplicitBDF from .integrate import integrate __all__ = [ "NewtonSolver", "SemiImplicitBDF", + "LinearizedBDF", "integrate", ] diff --git a/hydrogym/firedrake/solvers/base.py b/hydrogym/firedrake/solvers/base.py index d5f1f8a..67e7334 100644 --- a/hydrogym/firedrake/solvers/base.py +++ b/hydrogym/firedrake/solvers/base.py @@ -34,7 +34,7 @@ def solve(self, q: fd.Function = None): self.flow.init_bcs() - F = self.steady_form(q) # Nonlinear variational form + F = self.steady_form(fd.split(q)) # Nonlinear variational form J = fd.derivative(F, q) # Jacobian with automatic differentiation bcs = self.flow.collect_bcs() @@ -45,14 +45,14 @@ def solve(self, q: fd.Function = None): return q.copy(deepcopy=True) - def steady_form(self, q: fd.Function): - (u, p) = fd.split(q) - (v, s) = fd.TestFunctions(self.flow.mixed_space) + def steady_form(self, q: fd.Function, q_test=None): + (u, p) = q + if q_test is None: + (v, s) = fd.TestFunctions(self.flow.mixed_space) + else: + (v, s) = q_test - F = ( - inner(dot(u, nabla_grad(u)), v) * dx + - inner(self.flow.sigma(u, p), self.flow.epsilon(v)) * dx + - inner(div(u), s) * dx) + F = self.flow.residual((u, p), q_test=(v, s)) stab = self.stabilization_type( self.flow, q_trial=(u, p), diff --git a/hydrogym/firedrake/solvers/bdf_ext.py b/hydrogym/firedrake/solvers/bdf_ext.py index 58b1c5e..756243d 100644 --- a/hydrogym/firedrake/solvers/bdf_ext.py +++ b/hydrogym/firedrake/solvers/bdf_ext.py @@ -5,6 +5,11 @@ from .base import NavierStokesTransientSolver from .stabilization import ns_stabilization +__all__ = [ + "SemiImplicitBDF", + "LinearizedBDF", +] + _alpha_BDF = [1.0, 3.0 / 2.0, 11.0 / 6.0] _beta_BDF = [ [1.0], @@ -23,7 +28,7 @@ class SemiImplicitBDF(NavierStokesTransientSolver): def __init__( self, flow: FlowConfig, - dt: float = None, + dt: float, order: int = 3, stabilization: str = "default", rtol=1e-6, @@ -65,40 +70,7 @@ def initialize_functions(self): for u in self.u_prev: u.assign(flow.q.subfunctions[0]) - def _make_order_k_solver(self, k): - # Setup functions and spaces - flow = self.flow - h = fd.Constant(self.dt) - - (u, p) = self.q_trial - (v, s) = self.q_test - - # Combinations of functions for form construction - k_idx = k - 1 - # The "wind" w is the extrapolation estimate of u[n+1] - w = sum(beta * u_n for beta, u_n in zip(_beta_EXT[k_idx], self.u_prev)) - u_BDF = sum(beta * u_n for beta, u_n in zip(_beta_BDF[k_idx], self.u_prev)) - alpha_k = _alpha_BDF[k_idx] - u_t = (alpha_k * u - u_BDF) / h # BDF estimate of time derivative - - # Semi-implicit weak form - weak_form = ( - dot(u_t, v) * dx + dot(dot(w, nabla_grad(u)), v) * dx + - inner(flow.sigma(u, p), flow.epsilon(v)) * dx + dot(div(u), s) * dx - - dot(self.f, v) * dx) - - # Stabilization (SUPG, GLS, etc.) - stab = self.StabilizationType( - flow, - self.q_trial, - self.q_test, - wind=w, - dt=self.dt, - u_t=u_t, - f=self.f, - ) - weak_form = stab.stabilize(weak_form) - + def _make_petsc_solver(self, weak_form): # Construct variational problem and PETSc solver q = self.flow.q a = lhs(weak_form) @@ -115,9 +87,11 @@ def _make_order_k_solver(self, k): "pc_fieldsplit_type": "schur", "pc_fieldsplit_schur_fact_type": "full", "pc_fieldsplit_schur_precondition": "selfp", + # # Default preconditioner for inv(A) # (ilu in serial, bjacobi in parallel) "fieldsplit_0_ksp_type": "preonly", + # # Single multigrid cycle preconditioner for inv(S) "fieldsplit_1_ksp_type": "preonly", "fieldsplit_1_pc_type": "hypre", @@ -127,6 +101,44 @@ def _make_order_k_solver(self, k): bdf_prob, solver_parameters=solver_parameters) return petsc_solver + def _stabilize_weak_form(self, weak_form, u_t, wind, f=None): + # Stabilization (SUPG, GLS, etc.) + stab = self.StabilizationType( + self.flow, + self.q_trial, + self.q_test, + wind=wind, + dt=self.dt, + u_t=u_t, + f=f, + ) + return stab.stabilize(weak_form) + + def _make_order_k_solver(self, k): + # Setup functions and spaces + flow = self.flow + h = fd.Constant(self.dt) + + (u, p) = self.q_trial + (v, s) = self.q_test + + # Combinations of functions for form construction + k_idx = k - 1 + # The "wind" w is the extrapolation estimate of u[n+1] + w = sum(beta * u_n for beta, u_n in zip(_beta_EXT[k_idx], self.u_prev)) + u_BDF = sum(beta * u_n for beta, u_n in zip(_beta_BDF[k_idx], self.u_prev)) + alpha_k = _alpha_BDF[k_idx] + u_t = (alpha_k * u - u_BDF) / h # BDF estimate of time derivative + + # Semi-implicit weak form + weak_form = ( + dot(u_t, v) * dx + dot(dot(w, nabla_grad(u)), v) * dx + + inner(flow.sigma(u, p), flow.epsilon(v)) * dx + dot(div(u), s) * dx - + dot(self.f, v) * dx) + + weak_form = self._stabilize_weak_form(weak_form, u_t, wind=w, f=self.f) + return self._make_petsc_solver(weak_form) + def initialize_operators(self): self.flow.init_bcs() self.petsc_solver = self._make_order_k_solver(self.k) @@ -156,3 +168,64 @@ def step(self, iter, control=None): self.u_prev[0].assign(self.flow.q.subfunctions[0]) return self.flow + + +class LinearizedBDF(SemiImplicitBDF): + + def __init__(self, *args, qB: fd.Function, **kwargs): + self.qB = qB + stabilization = kwargs.pop("stabilization", "none").split("_") + if stabilization[0] != "linearized": + stabilization = ["linearized", stabilization[0]] + stabilization = "_".join(stabilization) + super().__init__(*args, stabilization=stabilization, **kwargs) + + def _make_order_k_solver(self, k): + # Setup functions and spaces + flow = self.flow + sigma, epsilon = flow.sigma, flow.epsilon + + h = fd.Constant(self.dt) + + flow.linearize_bcs() + + (uB, pB) = self.qB.subfunctions + (u, p) = self.q_trial + (v, s) = self.q_test + + # Combinations of functions for form construction + k_idx = k - 1 + u_BDF = sum(beta * u_n for beta, u_n in zip(_beta_BDF[k_idx], self.u_prev)) + alpha_k = _alpha_BDF[k_idx] + u_t = (alpha_k * u - u_BDF) / h # BDF estimate of time derivative + + # Semi-implicit weak form + # Note that the base flow terms are added to the RHS in `self.f` + weak_form = ( + dot(u_t, v) * dx + dot(dot(uB, nabla_grad(u)), v) * dx + + dot(dot(u, nabla_grad(uB)), v) * dx + + inner(sigma(u, p), epsilon(v)) * dx + dot(div(u), s) * dx - + dot(self.f, v) * dx + # # Base flow forcing (will be zero if base flow is steady solution) + # + inner(dot(uB, nabla_grad(uB)), v) * dx + # + inner(sigma(uB, pB), epsilon(v)) * dx + ) + + # RHS forcing term for stabilization (also zero if d(qB)/dt = 0) + f_stab = self.f # - dot(uB, nabla_grad(uB)) + div(sigma(uB, pB)) + + weak_form = self._stabilize_weak_form(weak_form, u_t, wind=uB, f=f_stab) + return self._make_petsc_solver(weak_form) + + # def initialize_functions(self): + # # Add to the RHS forcing with base flow terms + # super().initialize_functions() + + # sigma, epsilon = self.flow.sigma, self.flow.epsilon + # (uB, pB) = self.qB.subfunctions + # (v, _) = self.q_test + + # self.f += ( + # inner(sigma(uB, pB), epsilon(v)) * dx + # - inner(dot(uB, nabla_grad(uB)), v) * dx + # ) diff --git a/hydrogym/firedrake/solvers/integrate.py b/hydrogym/firedrake/solvers/integrate.py index 7068c4e..20c6188 100644 --- a/hydrogym/firedrake/solvers/integrate.py +++ b/hydrogym/firedrake/solvers/integrate.py @@ -1,9 +1,10 @@ -from .bdf_ext import SemiImplicitBDF +from .bdf_ext import LinearizedBDF, SemiImplicitBDF __all__ = ["integrate"] METHODS = { "BDF": SemiImplicitBDF, + "LinearizedBDF": LinearizedBDF, } diff --git a/hydrogym/firedrake/solvers/stabilization.py b/hydrogym/firedrake/solvers/stabilization.py index dae0c53..7dd8d85 100644 --- a/hydrogym/firedrake/solvers/stabilization.py +++ b/hydrogym/firedrake/solvers/stabilization.py @@ -115,8 +115,50 @@ def Lv(self): return dot(w, nabla_grad(v)) - div(sigma(v, s)) +class LinearizedNSStabilization(UpwindNSStabilization): + + @property + def Lu(self): + (u, p) = self.q_trial + + uB = self.wind + sigma = self.flow.sigma + + Lu = dot(uB, nabla_grad(u)) + dot(u, nabla_grad(uB)) - div(sigma(u, p)) + + if self.u_t is not None: + Lu += self.u_t + + if self.f is not None: + Lu -= self.f + + return Lu + + +class LinearizedSUPG(LinearizedNSStabilization): + + @property + def Lv(self): + (v, _) = self.q_test + uB = self.wind + return dot(uB, nabla_grad(v)) + dot(v, nabla_grad(uB)) + + +class LinearizedGLS(LinearizedNSStabilization): + + @property + def Lv(self): + (v, s) = self.q_test + uB = self.wind + sigma = self.flow.sigma + return dot(uB, nabla_grad(v)) + dot(v, nabla_grad(uB)) - div(sigma(v, s)) + + ns_stabilization = { "none": NavierStokesStabilization, "supg": SUPG, "gls": GLS, + "linearized_none": NavierStokesStabilization, + "linearized_supg": LinearizedSUPG, + "linearized_gls": LinearizedGLS, } diff --git a/hydrogym/firedrake/utils/__init__.py b/hydrogym/firedrake/utils/__init__.py index b39b4cb..ff87239 100644 --- a/hydrogym/firedrake/utils/__init__.py +++ b/hydrogym/firedrake/utils/__init__.py @@ -1,3 +1,5 @@ +from .eig import eig +from .stability import stability_analysis from .utils import get_array, is_rank_zero, print, set_from_array, white_noise # from . import io, linalg, modeling diff --git a/hydrogym/firedrake/utils/eig.py b/hydrogym/firedrake/utils/eig.py new file mode 100644 index 0000000..ee8dec0 --- /dev/null +++ b/hydrogym/firedrake/utils/eig.py @@ -0,0 +1,310 @@ +import dataclasses +from functools import partial +from typing import Callable + +import firedrake as fd +import numpy as np +from scipy import linalg + +from .linalg import LinearOperator +from .utils import print as parprint + +__all__ = ["eig"] + + +def sorted_eig(H, which, inverse=True): + evals, evecs = linalg.eig(H) + if inverse: + evals = 1 / evals # Undo spectral shift + + x = { + "lm": -abs(evals), + "sm": abs(evals), + "lr": -evals.real, + "sr": evals.real, + "li": -evals.imag, + "si": evals.imag, + }[which] + sort_idx = np.argsort(x) + + evals = evals[sort_idx] + evecs = evecs[:, sort_idx] + return evals, evecs + + +@dataclasses.dataclass +class ArnoldiIterator: + A: LinearOperator + inner_product: Callable + random_vec: Callable = None + print_evals: bool = True + which: str = "lr" + inverse: bool = True + + def __call__(self, v0, m, restart=None): + return self._arnoldi_factorization( + v0, + m=m, + restart=restart, + ) + + def _arnoldi_factorization( + self, + v0, + m=100, + restart=None, + ): + """Run Arnoldi iteration. + + Note that `which`, and `inverse` here are only used for printing the eigenvalues. + Otherwise these are not used as part of the Arnoldi factorization. If + `print_evals=False` these will be ignored. Also, the printed eigenvalues do not + include any spectral transformation. + + Given a linear operator `A` to calculate a matrix-vector product `f = A @ v` + (or inverse iteration `f = A^{-1} @ v`, shifted matrix pencil iteration, etc.), + an inner product `inner_product(u, v)`, and a starting vector `v0`, this + function computes an m-stage Arnoldi decomposition of the matrix `A` (or matrix + pencil (A, M) for mass matrix M) and returns the resulting Krylov vectors and + upper Hessenberg matrix. + """ + if restart is None: + V = [v0.copy(deepcopy=True).assign(0.0) for _ in range(m + 1)] + H = np.zeros((m, m)) + start_idx = 1 + V[0].assign(v0) + else: + V, H, start_idx = restart + + f = v0.copy(deepcopy=True) + for j in range(start_idx, m + 1): + # Apply iteration A @ f = M @ v + # This is the iteration f = (A^{-1} @ M) @ v + # which is inverse to M @ f = A @ v + # Stores the result in `f` + f = self.A @ V[j - 1] + + # Gram-Schmidt + h = np.zeros(j) + for i in range(j): + h[i] = self.inner_product(V[i], f) + f.assign(f - V[i] * h[i]) + + # Fill in the upper Hessenberg matrix + H[:j, j - 1] = h + + # Norm of the orthogonal residual + beta = np.sqrt(self.inner_product(f, f)) + + if j < m: + H[j, j - 1] = beta + + # Normalize the orthogonal residual for use as a new Krylov vector + V[j].assign(f / beta) + + # DEBUG: Print the eigenvalues + if self.print_evals: + ritz_vals, ritz_vecs = sorted_eig( + H[:j, :j], self.which, inverse=self.inverse) + res = abs(beta * ritz_vecs[-1]) + parprint("\n*************************************") + parprint(f"****** Arnoldi iteration {j} ********") + parprint("*************************************") + parprint("| evals | residuals |") + parprint("|---------------------|-------------|") + for i in range(min(10, j)): + parprint(f"| {ritz_vals[i]:.2e} | {res[i]:.2e} |") + + return V, H, beta + + +def _eig( + arnoldi: ArnoldiIterator, + v0=None, + schur_restart=False, + n_evals=10, + krylov_dim=100, + sort=None, + tol=1e-10, + delta=0.1, + maxiter=None, + rng_seed=None, +): + """Compute eigenvalues and eigenvectors using Arnoldi iteration. + + This function is written to be generic, i.e. not specific to `FlowConfig` + or Navier-Stokes. + """ + m = krylov_dim + if v0 is None: + v0 = arnoldi.random_vec(rng_seed) + + # TODO: Handle sort functions better - use a lookup of pre-defined + # functions for "lr", "sm", etc. + if sort is None: + # Decreasing magnitude + def sort(x): + return abs(x) > 1.0 - delta + + fn_space = v0.function_space() + + restart = None + is_converged = False + k = 0 # Iteration count (for logging) + while not is_converged: + V, H, beta = arnoldi(v0, m, restart=restart) + + # Check for convergence based on Arnoldi residuals + evals, ritz_vecs = sorted_eig(H, which="lr", inverse=True) + residuals = abs(beta * ritz_vecs[-1]) + + converged = [] + for i in range(m): + if residuals[i] < tol: + converged.append((evals[i], ritz_vecs[:, i], residuals[i])) + + parprint(f"Iteration {k}: {len(converged)} converged") + + # If enough eigenvalues have converged, return them and exit + if ((not schur_restart) or (len(converged) >= n_evals) or + (maxiter is not None and k >= maxiter)): + is_converged = True + parprint(f"Converged: {len(converged)} eigenvalues") + + p = len(converged) + evals = np.zeros(p, dtype=np.complex128) + residuals = np.zeros(p) + evecs_real = [fd.Function(fn_space) for _ in range(p)] + evecs_imag = [fd.Function(fn_space) for _ in range(p)] + for i, (mu, y, res) in enumerate(converged): + evals[i] = mu + residuals[i] = res + evecs_real[i].assign(sum(V[j] * y[j].real for j in range(m))) + evecs_imag[i].assign(sum(V[j] * y[j].imag for j in range(m))) + return evals, evecs_real, evecs_imag, residuals + + # Schur decomposition + S, Q, p = linalg.schur(H, sort=sort) + parprint(f" Schur form: {p} retained eigenvalues") + + # Keep the "wanted" part of the Schur form. These will be the eigenvalues + # that are closest to the unit circle (least stable). + S[p:, :p] = 0 + S[:p, p:] = 0 + S[p:, p:] = 0 + + # Re-order the Krylov basis. Since Q is real, these fields are also real. + Vp = [fd.Function(fn_space) for _ in range(m + 1)] + v0.assign(sum(V[j] * Q[j, p] for j in range(m))) + for i in range(p): + Vp[i].assign(sum(V[j] * Q[j, i] for j in range(m))) + + # Update the matrix with the "b" vector for residuals + b = np.zeros(m) + b[-1] = beta + b = Q.T @ b + S[p, :p] = b[:p] + Vp[p].assign(V[-1]) # Restart from the last Krylov basis vector + + # p += 1 #?? + + # Restart with the wanted part of the Krylov basis + restart = (Vp, S, p) + k += 1 + + +def _make_real_inv_iterator(flow, sigma, adjoint=False, solver_parameters=None): + + # Set up function spaces + fn_space = flow.mixed_space + A = flow.linearize( + sigma=sigma, + inverse=True, + adjoint=adjoint, + solver_parameters=solver_parameters) + + _inner_product = partial(flow.inner_product, augmented=False) + + def _random_vec(rng_seed=None): + rng = fd.RandomGenerator(fd.PCG64(seed=rng_seed)) + v0 = rng.standard_normal(fn_space) + alpha = fd.sqrt(_inner_product(v0, v0)) + v0.assign(v0 / alpha) + return v0 + + return ArnoldiIterator( + A, + _inner_product, + _random_vec, + which="lr", + inverse=True, + ) + + +def _make_complex_inv_iterator(flow, + sigma, + adjoint=False, + solver_parameters=None): + fn_space = flow.mixed_space + W = fn_space * fn_space + A = flow.linearize( + sigma=sigma, + inverse=True, + adjoint=adjoint, + solver_parameters=solver_parameters) + + _inner_product = partial(flow.inner_product, augmented=True) + + def _random_vec(rng_seed=None): + rng = fd.RandomGenerator(fd.PCG64(seed=rng_seed)) + v0 = rng.standard_normal(W) + alpha = fd.sqrt(_inner_product(v0, v0)) + v0.assign(v0 / alpha) + return v0 + + return ArnoldiIterator( + A, + _inner_product, + _random_vec, + which="lr", + inverse=True, + ) + + +def _make_st_iterator(flow, sigma=0.0, adjoint=False, solver_parameters=None): + """Construct a spectrally-transformed (shift-invert) Arnoldi iterator""" + if sigma.imag == 0.0: + return _make_real_inv_iterator( + flow, sigma.real, adjoint=adjoint, solver_parameters=solver_parameters) + return _make_complex_inv_iterator( + flow, sigma, adjoint=adjoint, solver_parameters=solver_parameters) + + +def eig( + flow, + v0=None, + sigma=0.0, + adjoint=False, + schur_restart=False, + n_evals=10, + krylov_dim=100, + sort=None, + tol=1e-10, + delta=0.1, + maxiter=None, + rng_seed=None, +): + """Eigenvalue decomposition using a shift-invert Arnoldi method.""" + arnoldi = _make_st_iterator(flow, sigma=sigma, adjoint=adjoint) + return _eig( + arnoldi, + v0=v0, + schur_restart=schur_restart, + n_evals=n_evals, + krylov_dim=krylov_dim, + sort=sort, + tol=tol, + delta=delta, + maxiter=maxiter, + rng_seed=rng_seed, + ) diff --git a/hydrogym/firedrake/utils/linalg.py b/hydrogym/firedrake/utils/linalg.py index abd6c1c..0ce8f28 100644 --- a/hydrogym/firedrake/utils/linalg.py +++ b/hydrogym/firedrake/utils/linalg.py @@ -1,4 +1,8 @@ -from typing import Iterable +from __future__ import annotations + +import abc +import dataclasses +from typing import TYPE_CHECKING, Callable, Iterable import firedrake as fd import numpy as np @@ -6,70 +10,102 @@ from firedrake import logging from scipy import sparse -# Type suggestions -from hydrogym.firedrake import FlowConfig +if TYPE_CHECKING: + from hydrogym.firedrake import FlowConfig +MUMPS_SOLVER_PARAMETERS = { + "mat_type": "aij", + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", +} -def adjoint(L): - args = L.arguments() - L_adj = ufl.adjoint(L, reordered_arguments=(args[0], args[1])) - return L_adj +class LinearOperator(metaclass=abc.ABCMeta): -EPS_PARAMETERS = { - "eps_gen_non_hermitian": None, - "eps_target": "0", - "eps_type": "krylovschur", - "eps_largest_real": True, - "st_type": "sinvert", - "st_pc_factor_shift_type": "NONZERO", - "eps_tol": 1e-10, -} + @abc.abstractmethod + def __matmul__(self, v: fd.Function): + """Return the matrix-vector product A @ v.""" + pass + + +class DirectOperator(LinearOperator): + J: ufl.Form + bcs: list[fd.DirichletBC] + function_space: fd.FunctionSpace + + def __matmul__(self, v: fd.Function): + f_bar = fd.assemble(self.J @ v, bcs=self.bcs) # Cofunction + return fd.Function(self.function_space, val=f_bar.dat) + @property + def T(self) -> InverseOperator: + """Return the adjoint operator. -def eig(A, M, num_eigenvalues=1, sigma=None, options={}): - """ - Compute eigendecomposition of the matrix pencil `(A, M)` using SLEPc, - where `A` is the dynamics matrix and `M` is the mass matrix. + This will solve the matrix pencil A^T @ f = M^T @ v0 for f. + """ + cls = self.__class__ + args = self.J.arguments() + JT = ufl.adjoint(self.J, reordered_arguments=(args[0], args[1])) + return cls(JT, self.bcs, self.function_space) - The default behavior is to use a shift-invert transformation to avoid inverting `M`, - which is singular in the case of the incompressible Navier-Stokes equations. - Ultimately this computes `[A - sigma * M]^-1` instead of `M^-1*A`, making it somewhat sensitive - to the choice of the shift `sigma`. This should be chosen near eigenvalues of interest + +@dataclasses.dataclass +class InverseOperator(LinearOperator): + """A simple wrapper for the inverse of a matrix pencil. + + Note that this object will own the output Function unless + `copy_output=True` is set. This is memory-efficient, but could + lead to confusion if the output reference is modified. The Arnoldi + iteration algorithm is written so this isn't a problem. """ - import numpy as np - from firedrake import COMM_WORLD - from firedrake.petsc import PETSc - from slepc4py import SLEPc - - assert (PETSc.ScalarType == np.complex128 - ), "Complex PETSc configuration required for stability analysis" - - assert not ( - (sigma is not None) and ("eps_target" in options) - ), "Shift value specified twice: use either `sigma` or `options['eps_target']` (behavior is the same)" - - slepc_options = EPS_PARAMETERS.copy() - slepc_options.update(options) - if sigma is not None: - slepc_options.update({"eps_target": f"{sigma.real}+{sigma.imag}i"}) - - # SLEPc Setup - opts = PETSc.Options() - for key, val in slepc_options.items(): - opts.setValue(key, val) - - es = SLEPc.EPS().create(comm=COMM_WORLD) - es.setDimensions(num_eigenvalues) - es.setOperators(A, M) - es.setFromOptions() - es.solve() - - nconv = es.getConverged() - vr, vi = A.getVecs() - - evals = np.array([es.getEigenpair(i, vr, vi) for i in range(nconv)]) - return evals, es + + J: ufl.Form + M: Callable[[fd.Function], ufl.Form] + bcs: list[fd.DirichletBC] + function_space: fd.FunctionSpace + solver_parameters: dict = None + copy_output: bool = False + + def __post_init__(self): + self._f = fd.Function(self.function_space) + self._v = fd.Function(self.function_space) + if self.solver_parameters is None: + self.solver_parameters = MUMPS_SOLVER_PARAMETERS + + self._problem = fd.LinearVariationalProblem( + self.J, self.M(self._v), self._f, bcs=self.bcs, constant_jacobian=True) + self._solver = fd.LinearVariationalSolver( + self._problem, solver_parameters=self.solver_parameters) + + @property + def T(self) -> InverseOperator: + """Return the adjoint operator. + + This will solve the matrix pencil A^T @ f = M^T @ v0 for f. + """ + cls = self.__class__ + args = self.J.arguments() + JT = ufl.adjoint(self.J, reordered_arguments=(args[0], args[1])) + return cls(JT, self.M, self.bcs, self.function_space, + self.solver_parameters) + + def __matmul__(self, v0): + """Solve the matrix pencil A @ f = M @ v0 for f. + + This is equivalent to the "inverse iteration" f = (A^{-1} @ M) @ v0 + """ + self._v.assign(v0) + self._solver.solve() + if self.copy_output: + return self._f.copy(deepcopy=True) + return self._f + + +def adjoint(L): + args = L.arguments() + L_adj = ufl.adjoint(L, reordered_arguments=(args[0], args[1])) + return L_adj def define_inner_product(mass_matrix): diff --git a/hydrogym/firedrake/utils/stability.py b/hydrogym/firedrake/utils/stability.py new file mode 100644 index 0000000..450a22f --- /dev/null +++ b/hydrogym/firedrake/utils/stability.py @@ -0,0 +1,125 @@ +from typing import NamedTuple + +import firedrake as fd +import numpy as np + +from .eig import eig +from .utils import print as parprint + +__all__ = ["stability_analysis"] + + +class StabilityResults(NamedTuple): + evals: np.ndarray + evecs_real: list[fd.Function] + evecs_imag: list[fd.Function] + residuals: np.ndarray + raw_evals: np.ndarray + + +def _filter_evals(evals, sigma, residuals, tol): + # The shifted eigenvalues have imaginary parts +/- (wi - si), but there + # are duplicates at +/- (wi + si). We want to discard these duplicates. + # We can check by doing a positive and negative shift by si and checking + # for duplicates. + keep = np.ones_like(evals, dtype=bool) + for i in range(len(evals)): + if not keep[i]: + continue + if residuals[i] > tol: + parprint(f"Not converged: {evals[i]}: {residuals[i]}") + keep[i] = False + continue + near_zero = ( + abs(evals[i].imag - sigma.imag) < tol or + abs(evals[i].imag + sigma.imag) < tol) + if near_zero: + parprint(f"Found near zero: {evals[i]}") + keep[i] = False + continue + for j in range(i + 1, len(evals)): + real_close = abs(evals[i].real - evals[j].real) < tol + shift_imag_close = ( + abs(evals[i].imag - evals[j].imag - 2 * sigma.imag) < tol or + abs(evals[i].imag - evals[j].imag + 2 * sigma.imag) < tol) + if real_close and shift_imag_close: + parprint(f"Found duplicate: {evals[i]} and {evals[j]}") + # Keep whichever has the smaller residual + if residuals[i] < residuals[j]: + keep[j] = False + else: + keep[i] = False + return keep + + +def stability_analysis( + flow, + sigma=0.0, + krylov_dim=100, + tol=1e-6, + adjoint=False, + schur_restart=False, + schur_delta=0.1, + n_evals=12, +): + """Linear stability analysis of the flow. + + Args: + flow: The FlowConfiguration to analyze. + sigma: + The shift for the shift-invert Arnoldi method. The algorithm will converge + most quickly if the shift is close to the eigenvalues of interest. + m: The dimension of the Krylov subspace (number of Arnoldi vectors). + tol: Tolerance to use for determining converged eigenvalues. + adjoint: If True, compute the adjoint modes along with the direct modes. + schur_restart: + If True, use Krylov-Schur iteration to restart the Arnoldi process. + schur_delta: + The stability margin to use when determining which Schur eigenvalues + to keep in Krylov-Schur iteration. Ignored if `schur_restart` is False. + n_evals: + The number of eigenvalues to converge in order to terminate Krylov-Schur + iteration. Ignored if `schur_restart` is False. + """ + + evals, evecs_real, evecs_imag, residuals = eig( + flow, + sigma=sigma, + adjoint=adjoint, + schur_restart=schur_restart, + krylov_dim=krylov_dim, + tol=tol, + n_evals=n_evals, + sort=lambda x: x.real > -schur_delta, + ) + + raw_evals = evals.copy() + + # Remove duplicate eigenvalues (only needed when sigma is complex) + if sigma.imag != 0: + keep = _filter_evals(evals, sigma, residuals, tol) + + else: + # Just filter by non-converged + keep = residuals < tol + + parprint(f"{np.sum(keep)} / {len(keep)} converged:") + + evals = evals[keep] + evals = np.where(evals.imag > 0, evals + sigma, evals + sigma.conjugate()) + + residuals = residuals[keep] + evecs_real = [evecs_real[i] for i in range(len(keep)) if keep[i]] + evecs_imag = [evecs_imag[i] for i in range(len(keep)) if keep[i]] + + n_save = min(len(evals), 32) + for _eval in evals[:n_save]: + parprint(_eval) + + return StabilityResults( + evals=evals, + evecs_real=evecs_real, + evecs_imag=evecs_imag, + residuals=residuals, + raw_evals=raw_evals, + )