********* PETSc ********* PETSc_ is a suite of data structures and routines for the scalable (parallel) solution of scientific applications modeled by partial differential equations. .. _PETSc: https://www.mcs.anl.gov/petsc/ ------------ Installation ------------ Please follow the following steps:: >>> pip3 install pymgrit >>> pip3 install petsc4py ------------------------ 2D heat equation example ------------------------ `example_petsc_heat_2d.py`_ .. _example_petsc_heat_2d.py: https://github.com/pymgrit/pymgrit/blob/master/examples/petsc4py/example_heat_2d_petsc4py.py The following example demonstrates the combination of PyMGRIT and petsc4py_ for the forced 2D heat equation with .. _petsc4py: https://www.mcs.anl.gov/petsc/petsc4py-current/docs/ .. math:: u_t - \Delta u = b(x,y,t) \;\; \text{ on } \; [0,1]\times[0,1]\times(t_0,t_{end}], with :math:`u(x,y, t_0) = u_0(x,y)`, homogeneous Dirichlet boundary condition and forcing term :math:`b(x,y,t)` such that the exact solution is given by .. math:: u(x,y,t)=\sin(2\pi x)sin(2\pi y)cos(t) \;\; \text{ on } \; [0,1]\times[0,1]\times[t_0,t_{end}]. Code ^^^^ :: import pathlib import os import copy import numpy as np from mpi4py import MPI from pymgrit.core.split import split_communicator from pymgrit.core.mgrit import Mgrit from pymgrit.core.grid_transfer import GridTransfer from pymgrit.core.grid_transfer_copy import GridTransferCopy from pymgrit.core.vector import Vector as PymgritVector from pymgrit.core.application import Application try: from petsc4py import PETSc except ImportError as e: import sys sys.exit("This examples requires petsc4py.") class VectorPetsc(PymgritVector): """ Vector class for PETSc vectors """ def __init__(self, values: PETSc.Vec) -> None: """ Constructor. :param values: PETSc.Vec with approximation """ if isinstance(values, PETSc.Vec): self.values = copy.deepcopy(values) else: raise Exception('Wrong datatype') def __add__(self, other: '__class__') -> '__class__': """ Addition of two vector objects (self and other) :param other: vector object to be added to self :return: sum of vector object self and input object other """ return VectorPetsc(self.get_values() + other.get_values()) def __sub__(self, other: '__class__') -> '__class__': """ Subtraction of two vector objects (self and other) :param other: vector object to be subtracted from self :return: difference of vector object self and input object other """ return VectorPetsc(self.get_values() - other.get_values()) def __mul__(self, other) -> 'VectorMachine': """ Multiplication of a vector object and a float (self and other) :param other: object to be multiplied with self :return: difference of vector object self and input object other """ return VectorPetsc(self.get_values() * other) def norm(self) -> float: """ Norm of a vector object :return: Frobenius-norm of vector object """ return self.values.norm(PETSc.NormType.FROBENIUS) def clone(self) -> '__class__': """ Initialize vector object with copied values :rtype: vector object with zero values """ return VectorPetsc(self.get_values()) def clone_zero(self) -> '__class__': """ Initialize vector object with zeros :rtype: vector object with zero values """ return VectorPetsc(self.get_values() * 0) def clone_rand(self) -> '__class__': """ Initialize vector object with random values :rtype: vector object with random values """ tmp = VectorPetsc(self.get_values()) return tmp def set_values(self, values: PETSc.Vec) -> None: """ Set vector data :param values: values for vector object """ self.values = values def get_values(self) -> PETSc.Vec: """ Get vector data :return: values of vector object """ return self.values def pack(self) -> np.ndarray: """ Pack data :return: values of vector object """ return self.values.getArray() def unpack(self, values: np.ndarray) -> None: """ Unpack and set data :param values: values for vector object """ self.values.setArray(values) class HeatPetsc(Application): """ 2D heat equation application with Dirichlet BCs in [0,1]x[0,1] """ def __init__(self, dmda: PETSc.DMDA, comm_x: MPI.Comm, freq: int, a: float, rtol: float = 1e-10, atol: float = 1e-10, max_it: int = 100, *args, **kwargs) -> None: """ Constructor :param dmda: PETSc DMDA grid :param comm_x: space communicator :param freq: frequency :param a: diffusion coefficient :param rtol: spatial solver tolerance :param atol: spatial solver tolerance :param max_it: spatial solver max iter :param args: :param kwargs: """ super().__init__(*args, **kwargs) self.dmda = dmda self.mx, self.my = self.dmda.sizes self.dx = 1.0 / (self.mx - 1) self.dy = 1.0 / (self.my - 1) (self.xs, self.xe), (self.ys, self.ye) = self.dmda.ranges self.freq = freq self.a = a self.space_disc = self.get_a() self.id = self.get_id() self.vector_template = VectorPetsc(self.dmda.createGlobalVec()) self.vector_t_start = VectorPetsc(self.u_exact(0).get_values()) # setup solver self.ksp = PETSc.KSP() self.ksp.create(comm=comm_x) self.ksp.setType('gmres') pc = self.ksp.getPC() pc.setType('none') self.ksp.setFromOptions() self.ksp.setTolerances(rtol=rtol, atol=atol, max_it=max_it) def step(self, u_start: VectorPetsc, t_start: float, t_stop: float) -> VectorPetsc: """ Time integration routine for 2D heat equation example problem :param u_start: approximate solution for the input time t_start :param t_start: time associated with the input approximate solution u_start :param t_stop: time to evolve the input approximate solution to :return: approximate solution at input time t_stop """ result = self.dmda.createGlobalVec() self.ksp.setOperators((t_stop - t_start) * self.space_disc + self.id) self.ksp.solve(self.compute_rhs(u_start=u_start, t_start=t_start, t_stop=t_stop), result) return VectorPetsc(result) def get_a(self) -> PETSc.Mat: """ Define spatial discretization matrix for 2D heat equation Second-order central finite differences with matrix stencil [ -f_y ] [-f_x 2(f_x + f_y) -f_x] [ -f_y ] with f_x = (a / dx^2) and f_y = (a / dy^2) :return: spatial discretization """ A = self.dmda.createMatrix() A.setType('aij') A.setFromOptions() A.setPreallocationNNZ((5, 5)) A.setUp() fx = self.a / self.dx ** 2 fy = self.a / self.dy ** 2 A.zeroEntries() row = PETSc.Mat.Stencil() col = PETSc.Mat.Stencil() for j in range(self.ys, self.ye): for i in range(self.xs, self.xe): row.index = (i, j) row.field = 0 if i == 0 or j == 0 or i == self.mx - 1 or j == self.my - 1: A.setValueStencil(row, row, 0.0) else: diag = 2 * (fx + fy) for index, value in [ ((i, j - 1), -fx), ((i - 1, j), -fy), ((i, j), diag), ((i + 1, j), -fy), ((i, j + 1), -fx), ]: col.index = index col.field = 0 A.setValueStencil(row, col, value) A.assemble() return A def get_id(self) -> PETSc.Mat: """ Define identity matrix :return: identity matrix """ Id = self.dmda.createMatrix() Id.setType('aij') Id.setFromOptions() Id.setPreallocationNNZ((1, 1)) Id.setUp() Id.zeroEntries() row = PETSc.Mat.Stencil() (xs, xe), (ys, ye) = self.dmda.ranges for j in range(ys, ye): for i in range(xs, xe): row.index = (i, j) row.field = 0 Id.setValueStencil(row, row, 1.0) Id.assemble() return Id def compute_rhs(self, u_start: VectorPetsc, t_start: float, t_stop: float) -> PETSc.Vec: """ Right-hand side of spatial system :param u_start: approximate solution for the input time t_start :param t_start: time associated with the input approximate solution u_start :param t_stop: time to evolve the input approximate solution to :return: right-hand side of spatial system at each time step in case of implicit time integration """ b = self.dmda.createGlobalVec() u = u_start.get_values() ba = self.dmda.getVecArray(b) ua = self.dmda.getVecArray(u) ba[self.xs:self.xe, self.ys:self.ye] = ua[self.xs:self.xe, self.ys:self.ye] + ( t_stop - t_start) * self.rhs(t_stop) return b def rhs(self, t_stop: float) -> np.ndarray: """ Right hand side :param t_stop: time point :return: right hand side """ xv, yv = np.meshgrid(range(self.xs, self.xe), range(self.ys, self.ye), indexing='ij') res = -np.sin(np.pi * self.freq * xv * self.dx) * \ np.sin(np.pi * self.freq * yv * self.dy) * \ (np.sin(t_stop) - self.a * 2.0 * (np.pi * self.freq) ** 2 * np.cos(t_stop)) return res def u_exact(self, t: float) -> VectorPetsc: """ Exact solution :param t: time point :return: exact solution at point t """ u = self.dmda.createGlobalVec() xa = self.dmda.getVecArray(u) for i in range(self.xs, self.xe): for j in range(self.ys, self.ye): xa[i, j] = np.sin(np.pi * self.freq * i * self.dx) * \ np.sin(np.pi * self.freq * j * self.dy) * np.cos(t) return VectorPetsc(u) class GridTransferPetsc(GridTransfer): """ Grid Transfer class between PETSc DMDA grids """ def __init__(self, fine_prob: PETSc.DMDA, coarse_prob: PETSc.DMDA) -> None: """ Constructor :param fine_prob: :param coarse_prob: """ super().__init__() self.coarse_prob = coarse_prob self.fine_prob = fine_prob self.interp, _ = self.coarse_prob.createInterpolation(fine_prob) self.inject = self.coarse_prob.createInjection(fine_prob) def restriction(self, u: VectorPetsc) -> VectorPetsc: """ Restriction :param u: fine approximation :return: coarse approximation """ u_coarse = self.coarse_prob.createGlobalVec() self.inject.mult(u.get_values(), u_coarse) return VectorPetsc(u_coarse) def interpolation(self, u: VectorPetsc) -> VectorPetsc: """ Interpolation :param u: coarse approximation :return: fine approximation """ u_fine = self.fine_prob.createGlobalVec() self.interp.mult(u.get_values(), u_fine) return VectorPetsc(u_fine) def main(): def output_fcn(self): # Set path to solution path = 'results/petsc' # Create path if not existing pathlib.Path(path).mkdir(parents=True, exist_ok=True) # Save solution with corresponding time point to file np.save(path + '/petsc' + str(self.comm_time_rank) + str(self.comm_space_rank), [[[self.t[0][i], self.comm_space_rank, self.u[0][i].get_values().getArray()] for i in self.index_local[0]]]) # Split the communicator into space and time communicator comm_world = MPI.COMM_WORLD comm_x, comm_t = split_communicator(comm_world, 2) # Create PETSc DMDA grids nx = 129 ny = 129 dmda_coarse = PETSc.DMDA().create([nx, ny], stencil_width=1, comm=comm_x) dmda_fine = dmda_coarse.refine() # Set up the problem heat_petsc_0 = HeatPetsc(dmda=dmda_fine, comm_x=comm_x, freq=1, a=1.0, t_start=0, t_stop=1, nt=33) heat_petsc_1 = HeatPetsc(dmda=dmda_coarse, comm_x=comm_x, freq=1, a=1.0, t_interval=heat_petsc_0.t[::2]) heat_petsc_2 = HeatPetsc(dmda=dmda_coarse, comm_x=comm_x, freq=1, a=1.0, t_interval=heat_petsc_1.t[::2]) # Setup three-level MGRIT solver with the space and time communicators and # solve the problem mgrit = Mgrit(problem=[heat_petsc_0, heat_petsc_1, heat_petsc_2], transfer=[GridTransferPetsc(fine_prob=dmda_fine, coarse_prob=dmda_coarse), GridTransferCopy()], comm_time=comm_t, comm_space=comm_x, output_fcn=output_fcn) info = mgrit.solve() import time if comm_t.Get_rank() == 0: time.sleep(1) sol = [] path = 'results/petsc/' for filename in os.listdir(path): data = np.load(path + filename, allow_pickle=True).tolist()[0] sol += data sol = [item for item in sol if item[1] == comm_x.Get_rank()] sol.sort(key=lambda tup: tup[0]) u_e = heat_petsc_0.u_exact(t=heat_petsc_0.t[-1]).get_values().getArray() diff = sol[-1][2] - u_e print('Difference at time point', heat_petsc_0.t[-1], ':', np.linalg.norm(diff, np.inf), '(space rank',comm_x.Get_rank() , ')') if __name__ == '__main__': main()