xref: /petsc/src/binding/petsc4py/demo/hpddm/hpddm.py (revision dc97da86830806a16cfaacfab740baa2e54c8135)
1*dc97da86SPierre Jolivet# This script demonstrates solving a symmetric positive definite linear system using PETSc and HPDDM preconditioner
2*dc97da86SPierre Jolivet# It must be run with exactly 4 MPI processes
3*dc97da86SPierre Jolivet# Example run:
4*dc97da86SPierre Jolivet#   mpirun -n 4 python3 hpddm.py -pc_hpddm_levels_1_sub_pc_type lu -pc_hpddm_levels_1_eps_nev 20 -pc_hpddm_levels_1_eps_threshold 0.1 -ksp_monitor
5*dc97da86SPierre Jolivet# For more options, see ${PETSC_DIR}/src/ksp/ksp/tutorials/ex76.c
6*dc97da86SPierre Jolivet
7*dc97da86SPierre Jolivetimport sys
8*dc97da86SPierre Jolivetimport petsc4py
9*dc97da86SPierre Jolivet# Initialize PETSc with command-line arguments
10*dc97da86SPierre Jolivetpetsc4py.init(sys.argv)
11*dc97da86SPierre Jolivetfrom petsc4py import PETSc
12*dc97da86SPierre Jolivet
13*dc97da86SPierre Jolivet# Get the rank of the current process
14*dc97da86SPierre Jolivetrank = PETSc.COMM_WORLD.getRank()
15*dc97da86SPierre Jolivet# Ensure that the script is run with exactly 4 processes
16*dc97da86SPierre Jolivetif PETSc.COMM_WORLD.getSize() != 4:
17*dc97da86SPierre Jolivet  if rank == 0:
18*dc97da86SPierre Jolivet    print("This example requires 4 processes")
19*dc97da86SPierre Jolivet  quit()
20*dc97da86SPierre Jolivet
21*dc97da86SPierre Jolivet# Load directory for input data
22*dc97da86SPierre Jolivetload_dir = PETSc.Options().getString("load_dir", "${DATAFILESPATH}/matrices/hpddm/GENEO")
23*dc97da86SPierre Jolivet
24*dc97da86SPierre Jolivet# Load an index set (IS) from binary file
25*dc97da86SPierre Jolivetsizes = PETSc.IS().load(PETSc.Viewer().createBinary(f"{load_dir}/sizes_{rank}_4.dat", "r", comm = PETSc.COMM_SELF))
26*dc97da86SPierre Jolivet# Get indices from the loaded IS
27*dc97da86SPierre Jolivetidx   = sizes.getIndices()
28*dc97da86SPierre Jolivet
29*dc97da86SPierre Jolivet# Create a PETSc matrix object
30*dc97da86SPierre JolivetA = PETSc.Mat().create()
31*dc97da86SPierre Jolivet# Set the local and global sizes of the matrix
32*dc97da86SPierre JolivetA.setSizes([[idx[0], idx[2]], [idx[1], idx[3]]])
33*dc97da86SPierre Jolivet# Configure matrix using runtime options
34*dc97da86SPierre JolivetA.setFromOptions()
35*dc97da86SPierre Jolivet# Load matrix A from binary file
36*dc97da86SPierre JolivetA = A.load(PETSc.Viewer().createBinary(f"{load_dir}/A.dat", "r", comm = PETSc.COMM_WORLD))
37*dc97da86SPierre Jolivet
38*dc97da86SPierre Jolivet# Load an index set (IS) from binary file
39*dc97da86SPierre Jolivetaux_IS  = PETSc.IS().load(PETSc.Viewer().createBinary(f"{load_dir}/is_{rank}_4.dat", "r", comm = PETSc.COMM_SELF))
40*dc97da86SPierre Jolivet# Set the block size of the index set
41*dc97da86SPierre Jolivetaux_IS.setBlockSize(2)
42*dc97da86SPierre Jolivet# Load the Neumann matrix of the current process
43*dc97da86SPierre Jolivetaux_Mat = PETSc.Mat().load(PETSc.Viewer().createBinary(f"{load_dir}/Neumann_{rank}_4.dat", "r", comm = PETSc.COMM_SELF))
44*dc97da86SPierre Jolivet
45*dc97da86SPierre Jolivet# Create and configure the linear solver (KSP) and preconditioner (PC)
46*dc97da86SPierre Jolivetksp = PETSc.KSP(PETSc.COMM_WORLD).create()
47*dc97da86SPierre Jolivetpc = ksp.getPC()
48*dc97da86SPierre Jolivet# Use HPDDM as the preconditioner
49*dc97da86SPierre Jolivetpc.setType(PETSc.PC.Type.HPDDM)
50*dc97da86SPierre Jolivet# Set the auxiliary matrix and index set for the local-to-global numbering
51*dc97da86SPierre Jolivetpc.setHPDDMAuxiliaryMat(aux_IS, aux_Mat)
52*dc97da86SPierre Jolivet# Inform HPDDM that the auxiliary matrix is the local Neumann matrix
53*dc97da86SPierre Jolivetpc.setHPDDMHasNeumannMat(True)
54*dc97da86SPierre Jolivet# Apply any command-line options (which may override options from the source file)
55*dc97da86SPierre Jolivetksp.setFromOptions()
56*dc97da86SPierre Jolivet# Set the system matrix (Amat = Pmat)
57*dc97da86SPierre Jolivetksp.setOperators(A)
58*dc97da86SPierre Jolivet
59*dc97da86SPierre Jolivet# Create RHS (b) and solution (x) vectors, set random values to b, and solve the system
60*dc97da86SPierre Jolivetb, x = A.createVecs()
61*dc97da86SPierre Jolivetb.setRandom()
62*dc97da86SPierre Jolivetksp.solve(b, x)
63*dc97da86SPierre Jolivet
64*dc97da86SPierre Jolivet# Output grid and operator complexities on rank 0
65*dc97da86SPierre Jolivetgc, oc = pc.getHPDDMComplexities()
66*dc97da86SPierre Jolivetif rank == 0:
67*dc97da86SPierre Jolivet  print("grid complexity = ", gc, ", operator complexity = ", oc, sep = "")
68