MGPoissonSolver.cpp 21.2 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 12

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

#include "Algorithms/PartBunch.h"
13
#include "Track/Track.h"
gsell's avatar
gsell committed
14 15 16 17 18
#include "Physics/Physics.h"
#include "Utilities/OpalException.h"
#include "Attributes/Attributes.h"
#include "ValueDefinitions/RealVariable.h"
#include "AbstractObjects/OpalData.h"
19
#include "Utilities/Options.h"
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 47

#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
48
#include <algorithm>
49
#include <omp.h>
gsell's avatar
gsell committed
50

51 52
using Teuchos::RCP;
using Teuchos::rcp;
gsell's avatar
gsell committed
53 54
using Physics::c;

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

gsell's avatar
gsell committed
72 73
    domain_m = layout_m->getDomain();

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

    if(itsolver == "CG") itsolver_m = AZ_cg;
80
    else if(itsolver == "BICGSTAB") itsolver_m = AZ_bicgstab;
gsell's avatar
gsell committed
81 82 83 84 85 86 87
    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;
88
    else if(precmode == "NO") precmode_m = NO;
gsell's avatar
gsell committed
89 90 91 92 93 94

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

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

102
    // Find CURRENT geometry
gsell's avatar
gsell committed
103 104
    currentGeometry = geometries_m[0];
    if(currentGeometry->getFilename() == "") {
105
        if(currentGeometry->getTopology() == "ELLIPTIC"){
106
            bp = new EllipticDomain(currentGeometry, orig_nr_m, hr_m, interpl);
107
	} else if (currentGeometry->getTopology() == "BOXCORNER") {
108
            bp = new BoxCornerDomain(currentGeometry->getA(), currentGeometry->getB(), currentGeometry->getC(), currentGeometry->getLength(),currentGeometry->getL1(), currentGeometry->getL2(), orig_nr_m, hr_m, interpl);
109
            bp->compute(itsBunch_m->get_hr());
gsell's avatar
gsell committed
110
        } else {
111 112
            throw OpalException("MGPoissonSolver::MGPoissonSolver",
                                "Geometry not known");
gsell's avatar
gsell committed
113
        }
114 115 116 117 118 119 120 121 122 123
    } 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");
        }
124
	bp = new ArbitraryDomain(currentGeometry, orig_nr_m, hr_m, interpl);
125
    }
126

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

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

    //all timers used here
150 151 152 153 154 155 156 157
    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
158 159 160 161 162 163 164 165 166 167
}


MGPoissonSolver::MGPoissonSolver(PartBunch &beam):
    layout_m(&beam.getFieldLayout()), mesh_m(&beam.getMesh()), itsBunch_m(&beam), LHS(0), Comm(Ippl::getComm()) {

    throw OpalException("MGPoissonSolver", "Copy Constructor not implemented yet");
}

MGPoissonSolver::~MGPoissonSolver() {
168 169 170 171 172 173 174 175 176 177 178 179 180 181
    if (Map) delete Map; Map = 0;
    if (MLPrec) delete MLPrec; MLPrec = 0;
    A = Teuchos::null;
    LHS = Teuchos::null;
    RHS = Teuchos::null;
    prec_m = Teuchos::null;
    solver_ptr = Teuchos::null;
    problem_ptr = Teuchos::null;
}

void MGPoissonSolver::deletePtr() {
    A = Teuchos::null;
    LHS = Teuchos::null;
    RHS = Teuchos::null;
182 183
    if(Map) delete Map;
    if(MLPrec) delete MLPrec; MLPrec=0;
184
    prec_m = Teuchos::null;
gsell's avatar
gsell committed
185 186 187 188 189 190
}

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

191
void MGPoissonSolver::computeMap(NDIndex<3> localId) {
192
    if (itsBunch_m->getLocalTrackStep()%repartFreq_m == 0){
193
        deletePtr();
gsell's avatar
gsell committed
194
        if(useRCB_m)
195
            redistributeWithRCB(localId);
196
        else
197
            IPPLToMap3D(localId);
198 199

        extrapolateLHS();
200 201
    }
}
gsell's avatar
gsell committed
202

203 204 205 206 207
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.
208
    //we also have to redistribute LHS
209

210 211 212 213 214 215 216 217 218
    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
219

220
    //...and all previously saved LHS
221
    std::deque< Epetra_Vector >::iterator it = OldLHS.begin();
222 223 224 225 226 227 228 229 230 231
    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
232

233 234 235
    // extrapolate last OldLHS.size LHS to get a new start vector
    it = OldLHS.begin();
    if(nLHS_m == 0 || OldLHS.size()==0)
236 237 238
        LHS->PutScalar(1.0);
    else if(OldLHS.size() == 1)
        *LHS = *it;
239 240 241
    else if(OldLHS.size() == 2)
        LHS->Update(2.0, *it++, -1.0, *it, 0.0);
    else if(OldLHS.size() > 2) {
242
        int n = OldLHS.size();
243 244 245
        P = rcp(new Epetra_MultiVector(*Map, nLHS_m, false));
        for(int i = 0; i < n; ++i) {
           *(*P)(i) = *it++;
246
        }
247 248 249 250
        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);
           }
251 252
        }
        *LHS = *(*P)(0);
253 254
     } else
        throw OpalException("MGPoissonSolver", "Invalid number of old LHS: " + OldLHS.size());
gsell's avatar
gsell committed
255 256
}

257

gsell's avatar
gsell committed
258 259 260 261 262
// 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) {
263

gsell's avatar
gsell committed
264 265 266
    nr_m[0] = orig_nr_m[0];
    nr_m[1] = orig_nr_m[1];
    nr_m[2] = orig_nr_m[2];
267

268 269
    bp->setGlobalMeanR(itsBunch_m->getGlobalMeanR());
    bp->setGlobalToLocalQuaternion(itsBunch_m->getGlobalToLocalQuaternion());
gsell's avatar
gsell committed
270
    bp->setNr(nr_m);
271

272 273
    NDIndex<3> localId = layout_m->getLocalNDIndex();

274
    IpplTimings::startTimer(FunctionTimer1_m);
275
    // Compute boundary intersections (only do on the first step)
276
    if(!itsBunch_m->getLocalTrackStep())
277
        bp->compute(hr, localId);
278
    IpplTimings::stopTimer(FunctionTimer1_m);
279 280

    // Define the Map
281
    INFOMSG(level2 << "* Computing Map..." << endl);
282
    IpplTimings::startTimer(FunctionTimer2_m);
283
    computeMap(localId);
284
    IpplTimings::stopTimer(FunctionTimer2_m);
285
    INFOMSG(level2 << "* Done." << endl);
gsell's avatar
gsell committed
286

287
    // Allocate the RHS with the new Epetra Map
288
    if (Teuchos::is_null(RHS))
289
        RHS = rcp(new Epetra_Vector(*Map));
290
    RHS->PutScalar(0.0);
291

292
    // get charge densities from IPPL field and store in Epetra vector (RHS)
293 294 295
    Ippl::Comm->barrier();
    std::cout << "* Node:" << Ippl::myNode() << ", Filling RHS..." << std::endl;
    Ippl::Comm->barrier();
296 297
    IpplTimings::startTimer(FunctionTimer3_m);
    int id = 0;
298
    float scaleFactor = itsBunch_m->getdT();
299

300

301 302 303 304 305 306 307
    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();
308
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
309
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
310
    	    for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
311 312
		    if (bp->isInside(idx, idy, idz))
                        RHS->Values()[id++] = 4*M_PI*rho[idx][idy][idz].get()/scaleFactor;
gsell's avatar
gsell committed
313 314 315
            }
        }
    }
316 317 318

    std::cout << "* Node:" << Ippl::myNode() << ", Number of Local Inside Points " << id << std::endl;
    Ippl::Comm->barrier();
319
    IpplTimings::stopTimer(FunctionTimer3_m);
320 321
    std::cout << "* Node:" << Ippl::myNode() << ", Done." << std::endl;
    Ippl::Comm->barrier();
322

gsell's avatar
gsell committed
323
    // build discretization matrix
324
    INFOMSG(level2 << "* Building Discretization Matrix..." << endl);
325
    IpplTimings::startTimer(FunctionTimer4_m);
326
    if(Teuchos::is_null(A))
327
        A = rcp(new Epetra_CrsMatrix(Copy, *Map,  7, true));
328
    ComputeStencil(hr, RHS);
329
    IpplTimings::stopTimer(FunctionTimer4_m);
330
    INFOMSG(level2 << "* Done." << endl);
331

gsell's avatar
gsell committed
332 333 334
#ifdef DBG_STENCIL
    EpetraExt::RowMatrixToMatlabFile("DiscrStencil.dat", *A);
#endif
335

336
    INFOMSG(level2 << "* Computing Preconditioner..." << endl);
337
    IpplTimings::startTimer(FunctionTimer5_m);
338
    if(!MLPrec) {
339 340 341 342
        MLPrec = new ML_Epetra::MultiLevelPreconditioner(*A, MLList_m);
    } else if (precmode_m == REUSE_HIERARCHY) {
        MLPrec->ReComputePreconditioner();
    } else if (precmode_m == REUSE_PREC){
343
    }
344
    IpplTimings::stopTimer(FunctionTimer5_m);
345
    INFOMSG(level2 << "* Done." << endl);
346

gsell's avatar
gsell committed
347 348
    // setup preconditioned iterative solver
    // use old LHS solution as initial guess
349
    INFOMSG(level2 << "* Final Setup of Solver..." << endl);
350
    IpplTimings::startTimer(FunctionTimer6_m);
351 352 353
    problem_ptr->setOperator(A);
    problem_ptr->setLHS(LHS);
    problem_ptr->setRHS(RHS);
354 355
    if(Teuchos::is_null(prec_m))
        prec_m = Teuchos::rcp ( new Belos::EpetraPrecOp ( rcp(MLPrec,false)));
356 357 358 359
    problem_ptr->setLeftPrec(prec_m);
    solver_ptr->setProblem( problem_ptr);
    if (!problem_ptr->isProblemSet()) {
        if (problem_ptr->setProblem() == false) {
360
            ERRORMSG("Belos::LinearProblem failed to set up correctly!" << endl);
361
        }
gsell's avatar
gsell committed
362
    }
363
    IpplTimings::stopTimer(FunctionTimer6_m);
364 365
    INFOMSG(level2 << "* Done." << endl);

gsell's avatar
gsell committed
366 367
    std::ofstream timings;
    char filename[50];
368
    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
369 370 371
    if(Comm.MyPID() == 0) timings.open(filename, std::ios::app);
    double time = MPI_Wtime();

372
    INFOMSG(level2 << "* Solving for Space Charge..." << endl);
373
    IpplTimings::startTimer(FunctionTimer7_m);
374
    	solver_ptr->solve();
375
    IpplTimings::stopTimer(FunctionTimer7_m);
376
    INFOMSG(level2 << "* Done." << endl);
gsell's avatar
gsell committed
377 378 379 380 381 382 383 384

    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 <<
385
                                      solver_ptr->getNumIters() << "\t" <<
gsell's avatar
gsell committed
386 387 388 389 390 391
                                      //time <<" "<<
                                      minTime << "\t" <<
                                      maxTime << "\t" <<
                                      avgTime << "\t" <<
                                      numBlocks_m << "\t" <<
                                      recycleBlocks_m << "\t" <<
392
                                      nLHS_m << "\t" <<
gsell's avatar
gsell committed
393 394 395 396 397
                                      //OldLHS.size() <<"\t"<<
                                      std::endl;
    if(Comm.MyPID() == 0) timings.close();

    // Store new LHS in OldLHS
398 399
    if(nLHS_m > 1) OldLHS.push_front(*(LHS.get()));
    if(OldLHS.size() > nLHS_m) OldLHS.pop_back();
gsell's avatar
gsell committed
400 401

    //now transfer solution back to IPPL grid
402 403
    IpplTimings::startTimer(FunctionTimer8_m);
    id = 0;
gsell's avatar
gsell committed
404
    rho = 0.0;
405
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
406
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
407
    	    for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
408 409
                  NDIndex<3> l(Index(idx, idx), Index(idy, idy), Index(idz, idz));
                  if (bp->isInside(idx, idy, idz))
410
                     rho.localElement(l) = LHS->Values()[id++]*scaleFactor;
gsell's avatar
gsell committed
411 412 413
            }
        }
    }
414
    IpplTimings::stopTimer(FunctionTimer8_m);
415 416 417 418 419 420

    if(itsBunch_m->getLocalTrackStep()+1 == (long long)Track::block->localTimeSteps.front()) {
        A = Teuchos::null;
        LHS = Teuchos::null;
        RHS = Teuchos::null;
        prec_m = Teuchos::null;
421
    }
gsell's avatar
gsell committed
422 423 424
}


425
void MGPoissonSolver::redistributeWithRCB(NDIndex<3> localId) {
gsell's avatar
gsell committed
426 427 428

    int numMyGridPoints = 0;

429
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
430
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
431
    	    for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
432
		    if (bp->isInside(idx, idy, idz))
gsell's avatar
gsell committed
433
                    numMyGridPoints++;
434
             }
gsell's avatar
gsell committed
435
        }
436
     }
gsell's avatar
gsell committed
437 438 439 440 441 442 443 444 445 446

    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;

447
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
448
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
449
    	    for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
450
		    if (bp->isInside(idx, idy, idz)) {
451 452 453
                    v[0] = (double)idx;
                    v[stride] = (double)idy;
                    v[stride2] = (double)idz;
gsell's avatar
gsell committed
454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482
                    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);
}

483
void MGPoissonSolver::IPPLToMap3D(NDIndex<3> localId) {
484

gsell's avatar
gsell committed
485
    int NumMyElements = 0;
486
    std::vector<int> MyGlobalElements;
gsell's avatar
gsell committed
487

488
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
489
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
490
            for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
491
		    if (bp->isInside(idx, idy, idz)) {
492
                    MyGlobalElements.push_back(bp->getIdx(idx, idy, idz));
gsell's avatar
gsell committed
493 494 495
                    NumMyElements++;
                }
            }
496 497
	}
    }
gsell's avatar
gsell committed
498 499 500 501 502
    Map = new Epetra_Map(-1, NumMyElements, &MyGlobalElements[0], 0, Comm);
}

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

503 504
    A->PutScalar(0.0);

gsell's avatar
gsell committed
505 506 507
    int NumMyElements = Map->NumMyElements();
    int *MyGlobalElements = Map->MyGlobalElements();

508 509
    std::vector<double> Values(6);
    std::vector<int> Indices(6);
gsell's avatar
gsell committed
510

511
    for(int i = 0 ; i < NumMyElements ; i++) {
gsell's avatar
gsell committed
512 513 514

        int NumEntries = 0;

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

518
        bp->getBoundaryStencil(MyGlobalElements[i], WV, EV, SV, NV, FV, BV, CV, scaleFactor);
519
        RHS->Values()[i] *= scaleFactor;
gsell's avatar
gsell committed
520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546

        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;
        }

547
    	// if matrix has already been filled (FillComplete()) we can only
gsell's avatar
gsell committed
548
        // replace entries
549

550 551 552 553 554
        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);
555
        } else {
gsell's avatar
gsell committed
556 557 558
            // off-diagonal entries
            A->InsertGlobalValues(MyGlobalElements[i], NumEntries, &Values[0], &Indices[0]);
            // diagonal entry
559
            A->InsertGlobalValues(MyGlobalElements[i], 1, &CV, MyGlobalElements + i);
560
        }
gsell's avatar
gsell committed
561 562 563
    }

    A->FillComplete();
564

gsell's avatar
gsell committed
565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587
    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();
588 589
    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
590 591 592 593 594 595
}

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;
596
    return os;
gsell's avatar
gsell committed
597 598
}

599
#endif /* HAVE_SAAMG_SOLVER */