static char help[] = "Test case: 2 ksp with KSPSetComputeOperators.\n\n";

#include <petscdm.h>
#include <petscdmda.h>
#include <petscksp.h>
#include <petsctime.h>


PetscErrorCode ComputeMatrix1(KSP ksp, Mat J,Mat jac, void *ctx);
PetscErrorCode ComputeMatrix2(KSP ksp, Mat J,Mat jac, void *ctx);
PetscErrorCode setupRhs(DM da, Vec b);


int main(int argc,char **argv)
{
    KSP            ksp1, ksp2;
    PC             pc1,pc2;
    DM             da;
    PetscReal      norm;
    PetscErrorCode ierr;
    Vec         x1,b1,x2,b2;
    Mat         J;
    PetscScalar val;
    
    ierr     = PetscInitialize(&argc, &argv, NULL, help);CHKERRQ(ierr);
    
    ierr = KSPCreate(PETSC_COMM_WORLD,&ksp1);CHKERRQ(ierr);
    ierr = KSPCreate(PETSC_COMM_WORLD,&ksp2);CHKERRQ(ierr);
    ierr = DMDACreate3d(PETSC_COMM_WORLD,DM_BOUNDARY_NONE,DM_BOUNDARY_NONE,DM_BOUNDARY_NONE,DMDA_STENCIL_STAR,16,16,16,PETSC_DECIDE,PETSC_DECIDE,PETSC_DECIDE,1,1,0,0,0,&da);CHKERRQ(ierr);
    ierr = DMDASetInterpolationType(da, DMDA_Q0);CHKERRQ(ierr);
    ierr = KSPSetDM(ksp1,da);CHKERRQ(ierr);
    ierr = KSPSetDM(ksp2,da);CHKERRQ(ierr);
    
    ierr = DMCreateGlobalVector(da, &b1);CHKERRQ(ierr);
    ierr = DMCreateGlobalVector(da, &b2);CHKERRQ(ierr);
    ierr = DMCreateGlobalVector(da, &x1);CHKERRQ(ierr);
    ierr = DMCreateGlobalVector(da, &x2);CHKERRQ(ierr);
    ierr = setupRhs(da, b1);CHKERRQ(ierr);
    ierr = setupRhs(da, b2);CHKERRQ(ierr);
    
    
    ierr = KSPSetComputeOperators(ksp1,ComputeMatrix1,NULL);CHKERRQ(ierr);
    ierr = KSPSetComputeOperators(ksp2,ComputeMatrix2,NULL);CHKERRQ(ierr);
    ierr = KSPSetType(ksp1, KSPGMRES);CHKERRQ(ierr);
    ierr = KSPSetType(ksp2, KSPGMRES);CHKERRQ(ierr);
    ierr = KSPGetPC(ksp1,&pc1);CHKERRQ(ierr);
    ierr = KSPGetPC(ksp2,&pc2);CHKERRQ(ierr);
    ierr = PCSetType(pc1, PCJACOBI);CHKERRQ(ierr);
    ierr = PCSetType(pc2, PCJACOBI);CHKERRQ(ierr);
    
    ierr = KSPSolve(ksp1,b1,x1);CHKERRQ(ierr);
    ierr = KSPSolve(ksp2,b2,x2);CHKERRQ(ierr);
    
    return 0;
}

PetscErrorCode ComputeMatrix1(KSP ksp, Mat J,Mat jac, void *ctx)
{
    PetscErrorCode ierr;
    PetscInt       i,j,k,mx,my,mz,xm,ym,zm,xs,ys,zs,size;
    PetscScalar    v[7],Hx,Hy,Hz;
    MatStencil     row, col[7];
    DM             da;
    
    PetscFunctionBeginUser;
    ierr    = PetscPrintf(PETSC_COMM_WORLD,"1\n");CHKERRQ(ierr);
    ierr    = KSPGetDM(ksp,&da);CHKERRQ(ierr);
    ierr    = DMDAGetInfo(da,0,&mx,&my,&mz,0,0,0,0,0,0,0,0,0);CHKERRQ(ierr);
    Hx      = 1.0 / (PetscReal)(mx-1);
    Hy      = 1.0 / (PetscReal)(my-1);
    Hz      = 1.0 / (PetscReal)(mz-1);
    ierr    = DMDAGetCorners(da,&xs,&ys,&zs,&xm,&ym,&zm);CHKERRQ(ierr);
    for (k=zs; k<zs+zm; k++) {
        for (j=ys; j<ys+ym; j++) {
            for (i=xs; i<xs+xm; i++) {
                size = 1;
                row.i = i; row.j = j; row.k = k;
                col[0].i = i;   col[0].j = j;   col[0].k = k;   v[0] = 1;
                ierr = MatSetValuesStencil(jac,1,&row,size,col,v,INSERT_VALUES);CHKERRQ(ierr);
            }
        }
    }
    
    ierr = MatAssemblyBegin(jac,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
    ierr = MatAssemblyEnd(jac,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PetscErrorCode ComputeMatrix2(KSP ksp, Mat J,Mat jac, void *ctx)
{
    PetscErrorCode ierr;
    ierr    = PetscPrintf(PETSC_COMM_WORLD,"2\n");CHKERRQ(ierr);
    PetscInt       i,j,k,mx,my,mz,xm,ym,zm,xs,ys,zs,size;
    PetscScalar    v[7],Hx,Hy,Hz;
    MatStencil     row, col[7];
    DM             da;
    
    PetscFunctionBeginUser;
    
    ierr    = KSPGetDM(ksp,&da);CHKERRQ(ierr);
    ierr    = DMDAGetInfo(da,0,&mx,&my,&mz,0,0,0,0,0,0,0,0,0);CHKERRQ(ierr);
    Hx      = 1.0 / (PetscReal)(mx-1);
    Hy      = 1.0 / (PetscReal)(my-1);
    Hz      = 1.0 / (PetscReal)(mz-1);
    ierr    = DMDAGetCorners(da,&xs,&ys,&zs,&xm,&ym,&zm);CHKERRQ(ierr);
    for (k=zs; k<zs+zm; k++) {
        for (j=ys; j<ys+ym; j++) {
            for (i=xs; i<xs+xm; i++) {
                size = 1;
                row.i = i; row.j = j; row.k = k;
                col[0].i = i;   col[0].j = j;   col[0].k = k;   v[0] = 1;
                ierr = MatSetValuesStencil(jac,1,&row,size,col,v,INSERT_VALUES);CHKERRQ(ierr);
            }
        }
    }
    
    ierr = MatAssemblyBegin(jac,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
    ierr = MatAssemblyEnd(jac,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
    PetscFunctionReturn(0);
}

PetscErrorCode setupRhs(DM da, Vec b)
{
    PetscErrorCode ierr;
    PetscInt       i,j,k,mx,my,mz,xm,ym,zm,xs,ys,zs;
    PetscScalar    Hx,Hy,Hz;
    PetscScalar    ***array;
    
    PetscFunctionBeginUser;
    ierr    = DMDAGetInfo(da,0,&mx,&my,&mz,0,0,0,0,0,0,0,0,0);CHKERRQ(ierr);
    Hx      = 1.0 / (PetscReal)(mx-1);
    Hy      = 1.0 / (PetscReal)(my-1);
    Hz      = 1.0 / (PetscReal)(mz-1);
    ierr    = DMDAGetCorners(da,&xs,&ys,&zs,&xm,&ym,&zm);CHKERRQ(ierr);
    ierr    = DMDAVecGetArray(da, b, &array);CHKERRQ(ierr);
    for (k=zs; k<zs+zm; k++) {
        for (j=ys; j<ys+ym; j++) {
            for (i=xs; i<xs+xm; i++) {
                array[k][j][i] = 1;
            }
        }
    }
    ierr = DMDAVecRestoreArray(da, b, &array);CHKERRQ(ierr);
    ierr = VecAssemblyBegin(b);CHKERRQ(ierr);
    ierr = VecAssemblyEnd(b);CHKERRQ(ierr);
    PetscFunctionReturn(0);
}