Class for solving an eigenvalue problem with fenics and slepc, accounting for bcs

fenics-eigensolver.py

from dolfin import (dx, Constant, assemble_system, TestFunction, 
                    as_backend_type, PETScOptions, Function, plot, File, MPI)
from matplotlib.pyplot import subplots
from slepc4py import SLEPc
import numpy as np

class EigenSolver(object):
    
    def __init__(self, 
                 a_k, 
                 a_m,
                 u,
                 bcs=None, 
                 slepc_options={'eps_max_it':100},
                 option_prefix=None
                ):
        self.slepc_options = slepc_options
        if option_prefix:
            self.E.setOptionsPrefix(option_prefix)
        self.V = u.function_space()
        l = Constant(0)*TestFunction(self.V)*dx
        if bcs:
            self.K_bc, _ = assemble_system(a_k,l,bcs)
            self.M_bc, _ = assemble_system(a_m,l,bcs)
            try:
                [bc.zero(self.M_bc) for bc in bcs]
            except:
                bcs.zero(self.M_bc)
        else: 
            self.K_bc, _ = assemble_system(a_k,l)
            self.M_bc, _ = assemble_system(a_m,l)
        self.E = self.eigensolver_setup()
        self.set_operators(self.K_bc,self.M_bc)
    
    def eigensolver_setup(self):
        E = SLEPc.EPS()
        E.create()        
        E.setType(SLEPc.EPS.Type.KRYLOVSCHUR)
        E.setProblemType(SLEPc.EPS.ProblemType.GHEP)
        #E.setWhichEigenpairs(E.Which.SMALLEST_MAGNITUDE)
        E.setWhichEigenpairs(E.Which.TARGET_MAGNITUDE)
        E.setTarget(0) 
        st = E.getST()
        st.setType('sinvert')
        return E
    
    def solve(self, n_eig):
        E = self.E
        E.setDimensions(n_eig)
        self.set_options(self.slepc_options)
        E.setFromOptions()
        E.solve()
        # print info
        its = E.getIterationNumber()
        eps_type = E.getType()
        self.nev, ncv, mpd = E.getDimensions()
        tol, maxit = E.getTolerances()
        self.nconv = E.getConverged()
        print("Solution method: {:s}, stopping condition: tol={:.4g}, maxit={:d}".format(eps_type,tol, maxit))
        print("Number of converged/requested eigenvalues with {:d} iterations  : {:d}/{:d}".format(its,self.nconv,self.nev))
        return self.nconv, its
    
    def set_operators(self,K,M):
        try:
            self.E.setOperators(K,M)
        except:   
            K_petsc = as_backend_type(K).mat()
            M_petsc = as_backend_type(M).mat()
            self.E.setOperators(K_petsc,M_petsc)
        
    def set_options(self,slepc_options):
        print("---- setting additional slepc options -----")
        for (opt, value) in slepc_options.items():
            print("    ",opt,":",value)
            PETScOptions.set(opt,value) 
        print("-------------------------------------------")
        self.E.setFromOptions()
        
    def plot_operators(self):
        fig, axs = subplots(nrows=1, ncols=2,)
        ax = axs[0]
        ax.spy(self.K_bc.array())
        ax.set_title('stiffness matrix', y=1.2)
        ax = axs[1]
        ax.spy(self.M_bc.array())
        ax.set_title('mass matrix',y=1.2);
        return fig, axs
        
    def get_eigenpair(self,i):
        u_r = Function(self.V)
        u_im = Function(self.V)
        eig = self.E.getEigenpair(i,  u_r.vector().vec(), u_im.vector().vec())
        return eig, u_r
        
    def get_eigenvalues(self,n):
        eigenvalues = [] 
        for i in range(n):
            eig, u_r = self.get_eigenpair(i)
            eigenvalues.append(eig.real)
        return np.array(eigenvalues)
    
    def get_eigenpairs(self,n):
        eigenvalues = [] 
        eigenvectors = [] 
        for i in range(n):
            eig, u_r = self.get_eigenpair(i)
            eigenvalues.append(eig.real)
            eigenvectors.append(u_r)
        return np.array(eigenvalues), eigenvectors
    
    def save_eigenvectors(self,n,file_name="output/modes.pvd"):
        eigenvalues = [] 
        eigenvectors = [] 
        eig, u_r = self.get_eigenpair(0)
        file = File(u_r.function_space().mesh().mpi_comm(),file_name)
        for i in range(n):
            eig, u_r = self.get_eigenpair(i)
            u_r.rename("mode","mode")
            file.write(u_r,eig.real)
        return file_name
    
    def plot_eigenpair(self,i):
        eig, u_r = self.get_eigenpair(i)
        p = plot(u_r,title="Eigenvector {:d} with eigenvalue {:6g}".format(i,eig))
        return p 
    
if __name__ == "__main__":

    import dolfin
    import ufl
    import numpy as np
    import matplotlib.pyplot as plt
    dolfin.parameters["use_petsc_signal_handler"] = True
    dolfin.parameters["form_compiler"]["cpp_optimize"] = True
    dolfin.parameters["form_compiler"]["representation"] = "uflacs"
    plt.style.use('seaborn')

    n = 100
    mesh = dolfin.UnitSquareMesh(n,n)
    V = dolfin.FunctionSpace(mesh,'CG',1)
    u = dolfin.Function(V)
    ut = dolfin.TestFunction(V)
    v = dolfin.TrialFunction(V)
    dx = dolfin.Measure("dx",domain=mesh)
    bc = dolfin.DirichletBC(V,0,"on_boundary") 
    a_k = ufl.dot(ufl.grad(ut),ufl.grad(v))*dx
    a_m = ut*v*dx
    eig_solver = EigenSolver(a_k, a_m, u, bc)
    eig_solver.plot_operators()
    plt.savefig("operators.png")
    ncv, it = eig_solver.solve(10)
    eigs = eig_solver.get_eigenvalues(ncv)
    plt.figure()
    plt.plot(eigs,'o')
    plt.title("Eigenvalues")
    plt.savefig("eigenvalues.png")

    eig_solver.save_eigenvectors(ncv)
    for i in range(ncv): 
        plt.figure()
        eig_solver.plot_eigenpair(i)
        plt.savefig("mode-{:d}.png".format(i))