xref: /petsc/src/dm/tests/ex30.c (revision ebead697dbf761eb322f829370bbe90b3bd93fa3)
1 static char help[] = "Tests DMSLICED operations\n\n";
2 
3 #include <petscdmsliced.h>
4 
5 int main(int argc,char *argv[])
6 {
7   char           mat_type[256] = MATAIJ; /* default matrix type */
8   MPI_Comm       comm;
9   PetscMPIInt    rank,size;
10   DM             slice;
11   PetscInt       i,bs=1,N=5,n,m,rstart,ghosts[2],*d_nnz,*o_nnz,dfill[4]={1,0,0,1},ofill[4]={1,1,1,1};
12   PetscReal      alpha   =1,K=1,rho0=1,u0=0,sigma=0.2;
13   PetscBool      useblock=PETSC_TRUE;
14   PetscScalar    *xx;
15   Mat            A;
16   Vec            x,b,lf;
17 
18   PetscFunctionBeginUser;
19   PetscCall(PetscInitialize(&argc,&argv,0,help));
20   comm = PETSC_COMM_WORLD;
21   PetscCallMPI(MPI_Comm_size(comm,&size));
22   PetscCallMPI(MPI_Comm_rank(comm,&rank));
23 
24   PetscOptionsBegin(comm,0,"Options for DMSliced test",0);
25   {
26     PetscCall(PetscOptionsInt("-n","Global number of nodes","",N,&N,NULL));
27     PetscCall(PetscOptionsInt("-bs","Block size (1 or 2)","",bs,&bs,NULL));
28     if (bs != 1) {
29       PetscCheck(bs == 2,PETSC_COMM_WORLD,PETSC_ERR_SUP,"Block size must be 1 or 2");
30       PetscCall(PetscOptionsReal("-alpha","Inverse time step for wave operator","",alpha,&alpha,NULL));
31       PetscCall(PetscOptionsReal("-K","Bulk modulus of compressibility","",K,&K,NULL));
32       PetscCall(PetscOptionsReal("-rho0","Reference density","",rho0,&rho0,NULL));
33       PetscCall(PetscOptionsReal("-u0","Reference velocity","",u0,&u0,NULL));
34       PetscCall(PetscOptionsReal("-sigma","Width of Gaussian density perturbation","",sigma,&sigma,NULL));
35       PetscCall(PetscOptionsBool("-block","Use block matrix assembly","",useblock,&useblock,NULL));
36     }
37     PetscCall(PetscOptionsString("-sliced_mat_type","Matrix type to use (aij or baij)","",mat_type,mat_type,sizeof(mat_type),NULL));
38   }
39   PetscOptionsEnd();
40 
41   /* Split ownership, set up periodic grid in 1D */
42   n         = PETSC_DECIDE;
43   PetscCall(PetscSplitOwnership(comm,&n,&N));
44   rstart    = 0;
45   PetscCallMPI(MPI_Scan(&n,&rstart,1,MPIU_INT,MPI_SUM,comm));
46   rstart   -= n;
47   ghosts[0] = (N+rstart-1)%N;
48   ghosts[1] = (rstart+n)%N;
49 
50   PetscCall(PetscMalloc2(n,&d_nnz,n,&o_nnz));
51   for (i=0; i<n; i++) {
52     if (size > 1 && (i==0 || i==n-1)) {
53       d_nnz[i] = 2;
54       o_nnz[i] = 1;
55     } else {
56       d_nnz[i] = 3;
57       o_nnz[i] = 0;
58     }
59   }
60   PetscCall(DMSlicedCreate(comm,bs,n,2,ghosts,d_nnz,o_nnz,&slice)); /* Currently does not copy X_nnz so we can't free them until after DMSlicedGetMatrix */
61 
62   if (!useblock) PetscCall(DMSlicedSetBlockFills(slice,dfill,ofill)); /* Irrelevant for baij formats */
63   PetscCall(DMSetMatType(slice,mat_type));
64   PetscCall(DMCreateMatrix(slice,&A));
65   PetscCall(PetscFree2(d_nnz,o_nnz));
66   PetscCall(MatSetOption(A,MAT_NEW_NONZERO_ALLOCATION_ERR,PETSC_TRUE));
67 
68   PetscCall(DMCreateGlobalVector(slice,&x));
69   PetscCall(VecDuplicate(x,&b));
70 
71   PetscCall(VecGhostGetLocalForm(x,&lf));
72   PetscCall(VecGetSize(lf,&m));
73   PetscCheck(m == (n+2)*bs,PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"size of local form %" PetscInt_FMT ", expected %" PetscInt_FMT,m,(n+2)*bs);
74   PetscCall(VecGetArray(lf,&xx));
75   for (i=0; i<n; i++) {
76     PetscInt        row[2],col[9],im,ip;
77     PetscScalar     v[12];
78     const PetscReal xref = 2.0*(rstart+i)/N - 1; /* [-1,1] */
79     const PetscReal h    = 1.0/N;                /* grid spacing */
80     im = (i==0) ? n : i-1;
81     ip = (i==n-1) ? n+1 : i+1;
82     switch (bs) {
83     case 1:                     /* Laplacian with periodic boundaries */
84       col[0] = im;         col[1] = i;        col[2] = ip;
85       v[0]   = -h;           v[1] = 2*h;        v[2] = -h;
86       PetscCall(MatSetValuesLocal(A,1,&i,3,col,v,INSERT_VALUES));
87       xx[i]  = PetscSinReal(xref*PETSC_PI);
88       break;
89     case 2:                     /* Linear acoustic wave operator in variables [rho, u], central differences, periodic, timestep 1/alpha */
90       v[0] = -0.5*u0;   v[1] = -0.5*K;      v[2] = alpha; v[3] = 0;       v[4] = 0.5*u0;    v[5] = 0.5*K;
91       v[6] = -0.5/rho0; v[7] = -0.5*u0;     v[8] = 0;     v[9] = alpha;   v[10] = 0.5/rho0; v[11] = 0.5*u0;
92       if (useblock) {
93         row[0] = i; col[0] = im; col[1] = i; col[2] = ip;
94         PetscCall(MatSetValuesBlockedLocal(A,1,row,3,col,v,INSERT_VALUES));
95       } else {
96         row[0] = 2*i; row[1] = 2*i+1;
97         col[0] = 2*im; col[1] = 2*im+1; col[2] = 2*i; col[3] = 2*ip; col[4] = 2*ip+1;
98         v[3]   = v[4]; v[4] = v[5];                                                     /* pack values in first row */
99         PetscCall(MatSetValuesLocal(A,1,row,5,col,v,INSERT_VALUES));
100         col[2] = 2*i+1;
101         v[8]   = v[9]; v[9] = v[10]; v[10] = v[11];                                     /* pack values in second row */
102         PetscCall(MatSetValuesLocal(A,1,row+1,5,col,v+6,INSERT_VALUES));
103       }
104       /* Set current state (gaussian density perturbation) */
105       xx[2*i]   = 0.2*PetscExpReal(-PetscSqr(xref)/(2*PetscSqr(sigma)));
106       xx[2*i+1] = 0;
107       break;
108     default: SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"not implemented for block size %" PetscInt_FMT,bs);
109     }
110   }
111   PetscCall(VecRestoreArray(lf,&xx));
112   PetscCall(VecGhostRestoreLocalForm(x,&lf));
113   PetscCall(MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY));
114   PetscCall(MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY));
115 
116   PetscCall(MatMult(A,x,b));
117   PetscCall(MatView(A,PETSC_VIEWER_STDOUT_WORLD));
118   PetscCall(VecView(x,PETSC_VIEWER_STDOUT_WORLD));
119   PetscCall(VecView(b,PETSC_VIEWER_STDOUT_WORLD));
120 
121   /* Update the ghosted values, view the result on rank 0. */
122   PetscCall(VecGhostUpdateBegin(b,INSERT_VALUES,SCATTER_FORWARD));
123   PetscCall(VecGhostUpdateEnd(b,INSERT_VALUES,SCATTER_FORWARD));
124   if (rank == 0) {
125     PetscCall(VecGhostGetLocalForm(b,&lf));
126     PetscCall(PetscViewerASCIIPrintf(PETSC_VIEWER_STDOUT_SELF,"Local form of b on rank 0, last two nodes are ghost nodes\n"));
127     PetscCall(VecView(lf,PETSC_VIEWER_STDOUT_SELF));
128     PetscCall(VecGhostRestoreLocalForm(b,&lf));
129   }
130 
131   PetscCall(DMDestroy(&slice));
132   PetscCall(VecDestroy(&x));
133   PetscCall(VecDestroy(&b));
134   PetscCall(MatDestroy(&A));
135   PetscCall(PetscFinalize());
136   return 0;
137 }
138 
139 /*TEST
140 
141    test:
142       nsize: 2
143       args: -bs 2 -block 0 -sliced_mat_type baij -alpha 10 -u0 0.1
144 
145    test:
146       suffix: 2
147       nsize: 2
148       args: -bs 2 -block 1 -sliced_mat_type aij -alpha 10 -u0 0.1
149 
150    test:
151       suffix: 3
152       nsize: 2
153       args: -bs 2 -block 0 -sliced_mat_type aij -alpha 10 -u0 0.1
154 
155 TEST*/
156