MGPoissonSolver.cpp 20.6 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
    } else
115 116
	bp = new ArbitraryDomain(currentGeometry, orig_nr_m, hr_m, interpl);

gsell's avatar
gsell committed
117 118 119 120
    Map = 0;
    A = Teuchos::null;
    LHS = Teuchos::null;
    RHS = Teuchos::null;
121
    MLPrec = 0;
122
    prec_m = Teuchos::null;
123

gsell's avatar
gsell committed
124 125
    numBlocks_m = Options::numBlocks;
    recycleBlocks_m = Options::recycleBlocks;
126
    nLHS_m = Options::nLHS;
gsell's avatar
gsell committed
127 128
    SetupMLList();
    SetupBelosList();
129
    problem_ptr = rcp(new Belos::LinearProblem<ST, MV, OP>);
130
    // setup Belos solver
gsell's avatar
gsell committed
131
    if(numBlocks_m == 0 || recycleBlocks_m == 0)
132
        solver_ptr = rcp(new Belos::BlockCGSolMgr<double, MV, OP>());
gsell's avatar
gsell committed
133
    else
134 135
        solver_ptr = rcp(new Belos::RCGSolMgr<double, MV, OP>());
    solver_ptr->setParameters(rcp(&belosList, false));
gsell's avatar
gsell committed
136 137 138 139
    convStatusTest = rcp(new Belos::StatusTestGenResNorm<ST, MV, OP> (tol));
    convStatusTest->defineScaleForm(Belos::NormOfRHS, Belos::TwoNorm);

    //all timers used here
140 141 142 143 144 145 146 147
    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
148 149 150 151 152 153 154 155 156 157
}


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() {
158 159 160 161 162 163 164 165 166 167 168 169 170 171
    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;
172 173
    if(Map) delete Map;
    if(MLPrec) delete MLPrec; MLPrec=0;
174
    prec_m = Teuchos::null;
gsell's avatar
gsell committed
175 176 177 178 179 180
}

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

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

        extrapolateLHS();
190 191
    }
}
gsell's avatar
gsell committed
192

193 194 195 196 197
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.
198
    //we also have to redistribute LHS
199

200 201 202 203 204 205 206 207 208
    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
209

210
    //...and all previously saved LHS
211
    std::deque< Epetra_Vector >::iterator it = OldLHS.begin();
212 213 214 215 216 217 218 219 220 221
    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
222

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

247

gsell's avatar
gsell committed
248 249 250 251 252
// 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) {
253

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

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

262 263
    NDIndex<3> localId = layout_m->getLocalNDIndex();

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

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

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

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

290

291 292 293 294 295 296 297
    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();
298
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
299
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
300
    	    for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
301 302
		    if (bp->isInside(idx, idy, idz))
                        RHS->Values()[id++] = 4*M_PI*rho[idx][idy][idz].get()/scaleFactor;
gsell's avatar
gsell committed
303 304 305
            }
        }
    }
306 307 308

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

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

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

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

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

gsell's avatar
gsell committed
356 357
    std::ofstream timings;
    char filename[50];
358
    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
359 360 361
    if(Comm.MyPID() == 0) timings.open(filename, std::ios::app);
    double time = MPI_Wtime();

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

    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 <<
375
                                      solver_ptr->getNumIters() << "\t" <<
gsell's avatar
gsell committed
376 377 378 379 380 381
                                      //time <<" "<<
                                      minTime << "\t" <<
                                      maxTime << "\t" <<
                                      avgTime << "\t" <<
                                      numBlocks_m << "\t" <<
                                      recycleBlocks_m << "\t" <<
382
                                      nLHS_m << "\t" <<
gsell's avatar
gsell committed
383 384 385 386 387
                                      //OldLHS.size() <<"\t"<<
                                      std::endl;
    if(Comm.MyPID() == 0) timings.close();

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

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

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


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

    int numMyGridPoints = 0;

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

    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;

437
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
438
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
439
    	    for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
440
		    if (bp->isInside(idx, idy, idz)) {
441 442 443
                    v[0] = (double)idx;
                    v[stride] = (double)idy;
                    v[stride2] = (double)idz;
gsell's avatar
gsell committed
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 471 472
                    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);
}

473
void MGPoissonSolver::IPPLToMap3D(NDIndex<3> localId) {
474

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

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

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

493 494
    A->PutScalar(0.0);

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

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

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

        int NumEntries = 0;

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

508
        bp->getBoundaryStencil(MyGlobalElements[i], WV, EV, SV, NV, FV, BV, CV, scaleFactor);
509
        RHS->Values()[i] *= scaleFactor;
gsell's avatar
gsell committed
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 535 536

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

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

540 541 542 543 544
        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);
545
        } else {
gsell's avatar
gsell committed
546 547 548
            // off-diagonal entries
            A->InsertGlobalValues(MyGlobalElements[i], NumEntries, &Values[0], &Indices[0]);
            // diagonal entry
549
            A->InsertGlobalValues(MyGlobalElements[i], 1, &CV, MyGlobalElements + i);
550
        }
gsell's avatar
gsell committed
551 552 553
    }

    A->FillComplete();
554

gsell's avatar
gsell committed
555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577
    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();
578 579
    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
580 581 582 583 584 585
}

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;
586
    return os;
gsell's avatar
gsell committed
587 588
}

589
#endif /* HAVE_SAAMG_SOLVER */