MGPoissonSolver.cpp 20.7 KB
Newer Older
1
#ifdef HAVE_SAAMG_SOLVER
2
//#define DBG_STENCIL
3

kraus's avatar
kraus committed
4
#include "Solvers/MGPoissonSolver.h"
5 6 7 8 9 10 11

#include "Structure/BoundaryGeometry.h"
#include "ArbitraryDomain.h"
#include "EllipticDomain.h"
#include "BoxCornerDomain.h"
//#include "RectangularDomain.h"

12
#include "Track/Track.h"
gsell's avatar
gsell committed
13 14 15 16 17
#include "Physics/Physics.h"
#include "Utilities/OpalException.h"
#include "Attributes/Attributes.h"
#include "ValueDefinitions/RealVariable.h"
#include "AbstractObjects/OpalData.h"
18
#include "Utilities/Options.h"
19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-local-typedefs"
#include "Epetra_Map.h"
#include "Epetra_Vector.h"
#include "Epetra_CrsMatrix.h"
#include "Epetra_Operator.h"
#include "EpetraExt_RowMatrixOut.h"
#include "Epetra_Import.h"

#include "Teuchos_CommandLineProcessor.hpp"

#include "BelosLinearProblem.hpp"
#include "BelosRCGSolMgr.hpp"
#include "BelosEpetraAdapter.hpp"
#include "BelosBlockCGSolMgr.hpp"

#include "ml_MultiLevelPreconditioner.h"
#include "ml_MultiLevelOperator.h"
#include "ml_epetra_utils.h"

#include "Isorropia_Exception.hpp"
#include "Isorropia_Epetra.hpp"
#include "Isorropia_EpetraRedistributor.hpp"
#include "Isorropia_EpetraPartitioner.hpp"

#pragma GCC diagnostic pop

gsell's avatar
gsell committed
47 48
#include <algorithm>

49 50
using Teuchos::RCP;
using Teuchos::rcp;
gsell's avatar
gsell committed
51 52
using Physics::c;

53
MGPoissonSolver::MGPoissonSolver ( PartBunch *beam,
54 55 56 57 58 59 60
                                   Mesh_t *mesh,
                                   FieldLayout_t *fl,
                                   std::vector<BoundaryGeometry *> geometries,
                                   std::string itsolver,
                                   std::string interpl,
                                   double tol,
                                   int maxiters,
61
                                   std::string precmode):
62 63 64
    geometries_m(geometries),
    tol_m(tol),
    maxiters_m(maxiters),
65 66 67 68
    Comm(Ippl::getComm()),
    itsBunch_m(beam),
    mesh_m(mesh),
    layout_m(fl) {
69

gsell's avatar
gsell committed
70 71
    domain_m = layout_m->getDomain();

72
    for (int i = 0; i < 3; i++) {
gsell's avatar
gsell committed
73 74 75 76 77
        hr_m[i] = mesh_m->get_meshSpacing(i);
        orig_nr_m[i] = domain_m[i].length();
    }

    if(itsolver == "CG") itsolver_m = AZ_cg;
78
    else if(itsolver == "BICGSTAB") itsolver_m = AZ_bicgstab;
gsell's avatar
gsell committed
79 80 81 82 83 84 85
    else if(itsolver == "GMRES") itsolver_m = AZ_gmres;
    else throw OpalException("MGPoissonSolver", "No valid iterative solver selected!");

    precmode_m = STD_PREC;
    if(precmode == "STD") precmode_m = STD_PREC;
    else if(precmode == "HIERARCHY") precmode_m = REUSE_HIERARCHY;
    else if(precmode == "REUSE") precmode_m = REUSE_PREC;
86
    else if(precmode == "NO") precmode_m = NO;
gsell's avatar
gsell committed
87 88 89 90 91 92

    tolerableIterationsCount_m = 2;
    numIter_m = -1;
    forcePreconditionerRecomputation_m = false;

    hasParallelDecompositionChanged_m = true;
93
    repartFreq_m = 1000;
gsell's avatar
gsell committed
94
    useRCB_m = false;
95
    if(Ippl::Info->getOutputLevel() > 1)
gsell's avatar
gsell committed
96 97 98 99
        verbose_m = true;
    else
        verbose_m = false;

100
    // Find CURRENT geometry
gsell's avatar
gsell committed
101 102
    currentGeometry = geometries_m[0];
    if(currentGeometry->getFilename() == "") {
103
        if(currentGeometry->getTopology() == "ELLIPTIC"){
104
            bp = new EllipticDomain(currentGeometry, orig_nr_m, hr_m, interpl);
105
	} else if (currentGeometry->getTopology() == "BOXCORNER") {
106
            bp = new BoxCornerDomain(currentGeometry->getA(), currentGeometry->getB(), currentGeometry->getC(), currentGeometry->getLength(),currentGeometry->getL1(), currentGeometry->getL2(), orig_nr_m, hr_m, interpl);
107
            bp->compute(itsBunch_m->get_hr());
gsell's avatar
gsell committed
108
        } else {
109 110
            throw OpalException("MGPoissonSolver::MGPoissonSolver",
                                "Geometry not known");
gsell's avatar
gsell committed
111
        }
112 113 114 115 116 117 118 119 120 121
    } else {
        NDIndex<3> localId = layout_m->getLocalNDIndex();
        if (localId[0].length() != domain_m[0].length() ||
            localId[1].length() != domain_m[1].length()) {
            throw OpalException("ArbitraryDomain::compute",
                                "The class ArbitraryDomain only works with parallelization\n"
                                "in z-direction.\n"
                                "Please set PARFFTX=FALSE, PARFFTY=FALSE, PARFFTT=TRUE in \n"
                                "the definition of the field solver in the input file.\n");
        }
122
	bp = new ArbitraryDomain(currentGeometry, orig_nr_m, hr_m, interpl);
123
    }
124

gsell's avatar
gsell committed
125 126 127 128
    Map = 0;
    A = Teuchos::null;
    LHS = Teuchos::null;
    RHS = Teuchos::null;
129
    MLPrec = 0;
130
    prec_m = Teuchos::null;
131

gsell's avatar
gsell committed
132 133
    numBlocks_m = Options::numBlocks;
    recycleBlocks_m = Options::recycleBlocks;
134
    nLHS_m = Options::nLHS;
gsell's avatar
gsell committed
135 136
    SetupMLList();
    SetupBelosList();
137
    problem_ptr = rcp(new Belos::LinearProblem<ST, MV, OP>);
138
    // setup Belos solver
gsell's avatar
gsell committed
139
    if(numBlocks_m == 0 || recycleBlocks_m == 0)
140
        solver_ptr = rcp(new Belos::BlockCGSolMgr<double, MV, OP>());
gsell's avatar
gsell committed
141
    else
142 143
        solver_ptr = rcp(new Belos::RCGSolMgr<double, MV, OP>());
    solver_ptr->setParameters(rcp(&belosList, false));
gsell's avatar
gsell committed
144 145 146 147
    convStatusTest = rcp(new Belos::StatusTestGenResNorm<ST, MV, OP> (tol));
    convStatusTest->defineScaleForm(Belos::NormOfRHS, Belos::TwoNorm);

    //all timers used here
148 149 150 151 152 153 154 155
    FunctionTimer1_m = IpplTimings::getTimer("BGF-IndexCoordMap");
    FunctionTimer2_m = IpplTimings::getTimer("computeMap");
    FunctionTimer3_m = IpplTimings::getTimer("IPPL to RHS");
    FunctionTimer4_m = IpplTimings::getTimer("ComputeStencil");
    FunctionTimer5_m = IpplTimings::getTimer("ML");
    FunctionTimer6_m = IpplTimings::getTimer("Setup");
    FunctionTimer7_m = IpplTimings::getTimer("CG");
    FunctionTimer8_m = IpplTimings::getTimer("LHS to IPPL");
gsell's avatar
gsell committed
156 157
}

158 159 160 161 162
void MGPoissonSolver::deletePtr() {
    if(Map) delete Map;
    Map = 0;
    if(MLPrec) delete MLPrec;
    MLPrec=0;
163 164 165 166 167 168
    A = Teuchos::null;
    LHS = Teuchos::null;
    RHS = Teuchos::null;
    prec_m = Teuchos::null;
}

169 170 171 172
MGPoissonSolver::~MGPoissonSolver() {
    deletePtr ();
    solver_ptr = Teuchos::null;
    problem_ptr = Teuchos::null;
gsell's avatar
gsell committed
173 174 175 176 177 178
}

void MGPoissonSolver::computePotential(Field_t &rho, Vector_t hr, double zshift) {
    throw OpalException("MGPoissonSolver", "not implemented yet");
}

179
void MGPoissonSolver::computeMap(NDIndex<3> localId) {
180
    if (itsBunch_m->getLocalTrackStep()%repartFreq_m == 0){
181
        deletePtr();
gsell's avatar
gsell committed
182
        if(useRCB_m)
183
            redistributeWithRCB(localId);
184
        else
185
            IPPLToMap3D(localId);
186 187

        extrapolateLHS();
188 189
    }
}
gsell's avatar
gsell committed
190

191 192 193 194 195
void MGPoissonSolver::extrapolateLHS() {
// Aitken-Neville
// Pi0 (x) := yi , i = 0 : n
// Pik (x) := (x − xi ) Pi+1,k−1(x) − (x − xi+k ) Pi,k−1(x) /(xi+k − xi )
// k = 1, . . . , n, i = 0, . . . , n − k.
196
    //we also have to redistribute LHS
197

198 199 200 201 202 203 204 205 206
    if (Teuchos::is_null(LHS)){
        LHS = rcp(new Epetra_Vector(*Map));
        LHS->PutScalar(1.0);
    } else {
        RCP<Epetra_Vector> tmplhs = rcp(new Epetra_Vector(*Map));
        Epetra_Import importer(*Map, LHS->Map());
        tmplhs->Import(*LHS, importer, Add);
        LHS = tmplhs;
    }
gsell's avatar
gsell committed
207

208
    //...and all previously saved LHS
209
    std::deque< Epetra_Vector >::iterator it = OldLHS.begin();
210 211 212 213 214 215 216 217 218 219
    if(OldLHS.size() > 0) {
        int n = OldLHS.size();
        for(int i = 0; i < n; ++i) {
            Epetra_Vector tmplhs = Epetra_Vector(*Map);
            Epetra_Import importer(*Map, it->Map());
            tmplhs.Import(*it, importer, Add);
            *it = tmplhs;
            ++it;
        }
    }
gsell's avatar
gsell committed
220

221 222 223
    // extrapolate last OldLHS.size LHS to get a new start vector
    it = OldLHS.begin();
    if(nLHS_m == 0 || OldLHS.size()==0)
224 225 226
        LHS->PutScalar(1.0);
    else if(OldLHS.size() == 1)
        *LHS = *it;
227 228 229
    else if(OldLHS.size() == 2)
        LHS->Update(2.0, *it++, -1.0, *it, 0.0);
    else if(OldLHS.size() > 2) {
230
        int n = OldLHS.size();
231 232 233
        P = rcp(new Epetra_MultiVector(*Map, nLHS_m, false));
        for(int i = 0; i < n; ++i) {
           *(*P)(i) = *it++;
234
        }
235 236 237 238
        for(int k = 1; k < n; ++k) {
           for(int i = 0; i < n - k; ++i) {
              (*P)(i)->Update(-(i + 1) / (float)k, *(*P)(i + 1), (i + k + 1) / (float)k);
           }
239 240
        }
        *LHS = *(*P)(0);
241
     } else
242
	    throw OpalException("MGPoissonSolver", "Invalid number of old LHS: " + std::to_string(OldLHS.size()));
gsell's avatar
gsell committed
243 244
}

245

gsell's avatar
gsell committed
246 247 248 249 250
// given a charge-density field rho and a set of mesh spacings hr,
// compute the electric field and put in eg by solving the Poisson's equation
// XXX: use matrix stencil in computation directly (no Epetra, define operators
// on IPPL GRID)
void MGPoissonSolver::computePotential(Field_t &rho, Vector_t hr) {
251

gsell's avatar
gsell committed
252 253 254
    nr_m[0] = orig_nr_m[0];
    nr_m[1] = orig_nr_m[1];
    nr_m[2] = orig_nr_m[2];
255

256 257
    bp->setGlobalMeanR(itsBunch_m->getGlobalMeanR());
    bp->setGlobalToLocalQuaternion(itsBunch_m->getGlobalToLocalQuaternion());
gsell's avatar
gsell committed
258
    bp->setNr(nr_m);
259

260 261
    NDIndex<3> localId = layout_m->getLocalNDIndex();

262
    IpplTimings::startTimer(FunctionTimer1_m);
263
    // Compute boundary intersections (only do on the first step)
264
    if(!itsBunch_m->getLocalTrackStep())
265
        bp->compute(hr, localId);
266
    IpplTimings::stopTimer(FunctionTimer1_m);
267 268

    // Define the Map
269
    INFOMSG(level2 << "* Computing Map..." << endl);
270
    IpplTimings::startTimer(FunctionTimer2_m);
271
    computeMap(localId);
272
    IpplTimings::stopTimer(FunctionTimer2_m);
273
    INFOMSG(level2 << "* Done." << endl);
gsell's avatar
gsell committed
274

275
    // Allocate the RHS with the new Epetra Map
276
    if (Teuchos::is_null(RHS))
277
        RHS = rcp(new Epetra_Vector(*Map));
278
    RHS->PutScalar(0.0);
279

280
    // get charge densities from IPPL field and store in Epetra vector (RHS)
281 282 283
    Ippl::Comm->barrier();
    std::cout << "* Node:" << Ippl::myNode() << ", Filling RHS..." << std::endl;
    Ippl::Comm->barrier();
284 285
    IpplTimings::startTimer(FunctionTimer3_m);
    int id = 0;
286
    float scaleFactor = itsBunch_m->getdT();
287

288

289 290 291 292 293 294 295
    std::cout << "* Node:" << Ippl::myNode() << ", Rho for final element: " << rho[localId[0].last()][localId[1].last()][localId[2].last()].get() << std::endl;

    Ippl::Comm->barrier();
    std::cout << "* Node:" << Ippl::myNode() << ", Local nx*ny*nz = " <<  localId[2].last() *  localId[0].last() *  localId[1].last() << std::endl;
    std::cout << "* Node:" << Ippl::myNode() << ", Number of reserved local elements in RHS: " << RHS->MyLength() << std::endl;
    std::cout << "* Node:" << Ippl::myNode() << ", Number of reserved global elements in RHS: " << RHS->GlobalLength() << std::endl;
    Ippl::Comm->barrier();
296
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
297
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
298
    	    for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
299 300
		    if (bp->isInside(idx, idy, idz))
                        RHS->Values()[id++] = 4*M_PI*rho[idx][idy][idz].get()/scaleFactor;
gsell's avatar
gsell committed
301 302 303
            }
        }
    }
304 305 306

    std::cout << "* Node:" << Ippl::myNode() << ", Number of Local Inside Points " << id << std::endl;
    Ippl::Comm->barrier();
307
    IpplTimings::stopTimer(FunctionTimer3_m);
308 309
    std::cout << "* Node:" << Ippl::myNode() << ", Done." << std::endl;
    Ippl::Comm->barrier();
310

gsell's avatar
gsell committed
311
    // build discretization matrix
312
    INFOMSG(level2 << "* Building Discretization Matrix..." << endl);
313
    IpplTimings::startTimer(FunctionTimer4_m);
314
    if(Teuchos::is_null(A))
315
        A = rcp(new Epetra_CrsMatrix(Copy, *Map,  7, true));
316
    ComputeStencil(hr, RHS);
317
    IpplTimings::stopTimer(FunctionTimer4_m);
318
    INFOMSG(level2 << "* Done." << endl);
319

gsell's avatar
gsell committed
320 321 322
#ifdef DBG_STENCIL
    EpetraExt::RowMatrixToMatlabFile("DiscrStencil.dat", *A);
#endif
323

324
    INFOMSG(level2 << "* Computing Preconditioner..." << endl);
325
    IpplTimings::startTimer(FunctionTimer5_m);
326
    if(!MLPrec) {
327 328 329 330
        MLPrec = new ML_Epetra::MultiLevelPreconditioner(*A, MLList_m);
    } else if (precmode_m == REUSE_HIERARCHY) {
        MLPrec->ReComputePreconditioner();
    } else if (precmode_m == REUSE_PREC){
331
    }
332
    IpplTimings::stopTimer(FunctionTimer5_m);
333
    INFOMSG(level2 << "* Done." << endl);
334

gsell's avatar
gsell committed
335 336
    // setup preconditioned iterative solver
    // use old LHS solution as initial guess
337
    INFOMSG(level2 << "* Final Setup of Solver..." << endl);
338
    IpplTimings::startTimer(FunctionTimer6_m);
339 340 341
    problem_ptr->setOperator(A);
    problem_ptr->setLHS(LHS);
    problem_ptr->setRHS(RHS);
342 343
    if(Teuchos::is_null(prec_m))
        prec_m = Teuchos::rcp ( new Belos::EpetraPrecOp ( rcp(MLPrec,false)));
344 345 346 347
    problem_ptr->setLeftPrec(prec_m);
    solver_ptr->setProblem( problem_ptr);
    if (!problem_ptr->isProblemSet()) {
        if (problem_ptr->setProblem() == false) {
348
            ERRORMSG("Belos::LinearProblem failed to set up correctly!" << endl);
349
        }
gsell's avatar
gsell committed
350
    }
351
    IpplTimings::stopTimer(FunctionTimer6_m);
352 353
    INFOMSG(level2 << "* Done." << endl);

gsell's avatar
gsell committed
354 355
    std::ofstream timings;
    char filename[50];
356
    sprintf(filename, "timing_MX%d_MY%d_MZ%d_nProc%d_recB%d_numB%d_nLHS%d", orig_nr_m[0], orig_nr_m[1], orig_nr_m[2], Comm.NumProc(), recycleBlocks_m, numBlocks_m, nLHS_m);
gsell's avatar
gsell committed
357 358 359
    if(Comm.MyPID() == 0) timings.open(filename, std::ios::app);
    double time = MPI_Wtime();

360
    INFOMSG(level2 << "* Solving for Space Charge..." << endl);
361
    IpplTimings::startTimer(FunctionTimer7_m);
362
    	solver_ptr->solve();
363
    IpplTimings::stopTimer(FunctionTimer7_m);
364
    INFOMSG(level2 << "* Done." << endl);
gsell's avatar
gsell committed
365 366 367 368 369 370 371 372

    time = MPI_Wtime() - time;
    double minTime, maxTime, avgTime;
    Comm.MinAll(&time, &minTime, 1);
    Comm.MaxAll(&time, &maxTime, 1);
    Comm.SumAll(&time, &avgTime, 1);
    avgTime /= Comm.NumProc();
    if(Comm.MyPID() == 0) timings <<
373
                                      solver_ptr->getNumIters() << "\t" <<
gsell's avatar
gsell committed
374 375 376 377 378 379
                                      //time <<" "<<
                                      minTime << "\t" <<
                                      maxTime << "\t" <<
                                      avgTime << "\t" <<
                                      numBlocks_m << "\t" <<
                                      recycleBlocks_m << "\t" <<
380
                                      nLHS_m << "\t" <<
gsell's avatar
gsell committed
381 382 383 384 385
                                      //OldLHS.size() <<"\t"<<
                                      std::endl;
    if(Comm.MyPID() == 0) timings.close();

    // Store new LHS in OldLHS
386 387
    if(nLHS_m > 1) OldLHS.push_front(*(LHS.get()));
    if(OldLHS.size() > nLHS_m) OldLHS.pop_back();
gsell's avatar
gsell committed
388 389

    //now transfer solution back to IPPL grid
390 391
    IpplTimings::startTimer(FunctionTimer8_m);
    id = 0;
gsell's avatar
gsell committed
392
    rho = 0.0;
393
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
394
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
395
    	    for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
396 397
                  NDIndex<3> l(Index(idx, idx), Index(idy, idy), Index(idz, idz));
                  if (bp->isInside(idx, idy, idz))
398
                     rho.localElement(l) = LHS->Values()[id++]*scaleFactor;
gsell's avatar
gsell committed
399 400 401
            }
        }
    }
402
    IpplTimings::stopTimer(FunctionTimer8_m);
403 404 405 406 407 408

    if(itsBunch_m->getLocalTrackStep()+1 == (long long)Track::block->localTimeSteps.front()) {
        A = Teuchos::null;
        LHS = Teuchos::null;
        RHS = Teuchos::null;
        prec_m = Teuchos::null;
409
    }
gsell's avatar
gsell committed
410 411 412
}


413
void MGPoissonSolver::redistributeWithRCB(NDIndex<3> localId) {
gsell's avatar
gsell committed
414 415 416

    int numMyGridPoints = 0;

417
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
418
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
419
    	    for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
420
		    if (bp->isInside(idx, idy, idz))
gsell's avatar
gsell committed
421
                    numMyGridPoints++;
422
             }
gsell's avatar
gsell committed
423
        }
424
     }
gsell's avatar
gsell committed
425 426 427 428 429 430 431 432 433 434

    Epetra_BlockMap bmap(-1, numMyGridPoints, 1, 0, Comm);
    Teuchos::RCP<const Epetra_MultiVector> coords = Teuchos::rcp(new Epetra_MultiVector(bmap, 3, false));

    double *v;
    int stride, stride2;

    coords->ExtractView(&v, &stride);
    stride2 = 2 * stride;

435
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
436
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
437
    	    for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
438
		    if (bp->isInside(idx, idy, idz)) {
439 440 441
                    v[0] = (double)idx;
                    v[stride] = (double)idy;
                    v[stride2] = (double)idz;
gsell's avatar
gsell committed
442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470
                    v++;
                }
            }
        }
    }

    Teuchos::ParameterList paramlist;
    paramlist.set("Partitioning Method", "RCB");
    Teuchos::ParameterList &sublist = paramlist.sublist("ZOLTAN");
    sublist.set("RCB_RECTILINEAR_BLOCKS", "1");
    sublist.set("DEBUG_LEVEL", "1");

    Teuchos::RCP<Isorropia::Epetra::Partitioner> part = Teuchos::rcp(new Isorropia::Epetra::Partitioner(coords, paramlist));
    Isorropia::Epetra::Redistributor rd(part);
    Teuchos::RCP<Epetra_MultiVector> newcoords = rd.redistribute(*coords);

    newcoords->ExtractView(&v, &stride);
    stride2 = 2 * stride;
    numMyGridPoints = 0;
    std::vector<int> MyGlobalElements;
    for(int i = 0; i < newcoords->MyLength(); i++) {
        MyGlobalElements.push_back(bp->getIdx(v[0], v[stride], v[stride2]));
        v++;
        numMyGridPoints++;
    }

    Map = new Epetra_Map(-1, numMyGridPoints, &MyGlobalElements[0], 0, Comm);
}

471
void MGPoissonSolver::IPPLToMap3D(NDIndex<3> localId) {
472

gsell's avatar
gsell committed
473
    int NumMyElements = 0;
474
    std::vector<int> MyGlobalElements;
gsell's avatar
gsell committed
475

476
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
477
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
478
            for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
479
		    if (bp->isInside(idx, idy, idz)) {
480
                    MyGlobalElements.push_back(bp->getIdx(idx, idy, idz));
gsell's avatar
gsell committed
481 482 483
                    NumMyElements++;
                }
            }
484 485
	}
    }
gsell's avatar
gsell committed
486 487 488 489 490
    Map = new Epetra_Map(-1, NumMyElements, &MyGlobalElements[0], 0, Comm);
}

void MGPoissonSolver::ComputeStencil(Vector_t hr, Teuchos::RCP<Epetra_Vector> RHS) {

491 492
    A->PutScalar(0.0);

gsell's avatar
gsell committed
493 494 495
    int NumMyElements = Map->NumMyElements();
    int *MyGlobalElements = Map->MyGlobalElements();

496 497
    std::vector<double> Values(6);
    std::vector<int> Indices(6);
gsell's avatar
gsell committed
498

499
    for(int i = 0 ; i < NumMyElements ; i++) {
gsell's avatar
gsell committed
500 501 502

        int NumEntries = 0;

503
        double WV, EV, SV, NV, FV, BV, CV, scaleFactor=1.0;
504
        int W, E, S, N, F, B;
gsell's avatar
gsell committed
505

506
        bp->getBoundaryStencil(MyGlobalElements[i], WV, EV, SV, NV, FV, BV, CV, scaleFactor);
507
        RHS->Values()[i] *= scaleFactor;
gsell's avatar
gsell committed
508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534

        bp->getNeighbours(MyGlobalElements[i], W, E, S, N, F, B);
        if(E != -1) {
            Indices[NumEntries] = E;
            Values[NumEntries++] = EV;
        }
        if(W != -1) {
            Indices[NumEntries] = W;
            Values[NumEntries++] = WV;
        }
        if(S != -1) {
            Indices[NumEntries] = S;
            Values[NumEntries++] = SV;
        }
        if(N != -1) {
            Indices[NumEntries] = N;
            Values[NumEntries++] = NV;
        }
        if(F != -1) {
            Indices[NumEntries] = F;
            Values[NumEntries++] = FV;
        }
        if(B != -1) {
            Indices[NumEntries] = B;
            Values[NumEntries++] = BV;
        }

535
    	// if matrix has already been filled (FillComplete()) we can only
gsell's avatar
gsell committed
536
        // replace entries
537

538 539 540 541 542
        if(A->Filled()) {
            // off-diagonal entries
            A->ReplaceGlobalValues(MyGlobalElements[i], NumEntries, &Values[0], &Indices[0]);
            // diagonal entry
            A->ReplaceGlobalValues(MyGlobalElements[i], 1, &CV, MyGlobalElements + i);
543
        } else {
gsell's avatar
gsell committed
544 545 546
            // off-diagonal entries
            A->InsertGlobalValues(MyGlobalElements[i], NumEntries, &Values[0], &Indices[0]);
            // diagonal entry
547
            A->InsertGlobalValues(MyGlobalElements[i], 1, &CV, MyGlobalElements + i);
548
        }
gsell's avatar
gsell committed
549 550 551
    }

    A->FillComplete();
552

gsell's avatar
gsell committed
553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575
    A->OptimizeStorage();
}

void MGPoissonSolver::printLoadBalanceStats() {

    //compute some load balance statistics
    size_t myNumPart = Map->NumMyElements();
    size_t NumPart = Map->NumGlobalElements() * 1.0 / Comm.NumProc();
    double imbalance = 1.0;
    if(myNumPart >= NumPart)
        imbalance += (myNumPart - NumPart) / NumPart;
    else
        imbalance += (NumPart - myNumPart) / NumPart;

    double max = 0.0, min = 0.0, avg = 0.0;
    int minn = 0, maxn = 0;
    MPI_Reduce(&imbalance, &min, 1, MPI_DOUBLE, MPI_MIN, 0, Ippl::getComm());
    MPI_Reduce(&imbalance, &max, 1, MPI_DOUBLE, MPI_MAX, 0, Ippl::getComm());
    MPI_Reduce(&imbalance, &avg, 1, MPI_DOUBLE, MPI_SUM, 0, Ippl::getComm());
    MPI_Reduce(&myNumPart, &minn, 1, MPI_INT, MPI_MIN, 0, Ippl::getComm());
    MPI_Reduce(&myNumPart, &maxn, 1, MPI_INT, MPI_MAX, 0, Ippl::getComm());

    avg /= Comm.NumProc();
576 577
    if(Comm.MyPID() == 0) *gmsg << "LBAL min = " << min << ", max = " << max << ", avg = " << avg << endl;
    if(Comm.MyPID() == 0) *gmsg << "min nr gridpoints = " << minn << ", max nr gridpoints = " << maxn << endl;
gsell's avatar
gsell committed
578 579 580 581 582 583
}

Inform &MGPoissonSolver::print(Inform &os) const {
    os << "* *************** M G P o i s s o n S o l v e r ************************************ " << endl;
    os << "* h " << hr_m << '\n';
    os << "* ********************************************************************************** " << endl;
584
    return os;
gsell's avatar
gsell committed
585 586
}

587
#endif /* HAVE_SAAMG_SOLVER */