MGPoissonSolver.cpp 22 KB
Newer Older
frey_m's avatar
frey_m committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14
//
// Class MGPoissonSolver
//   This class contains methods for solving Poisson's equation for the
//   space charge portion of the calculation.
//
//   A smoothed aggregation based AMG preconditioned iterative solver for space charge
//   \see FFTPoissonSolver
//   \warning This solver is in an EXPERIMENTAL STAGE. For reliable simulations use the FFTPoissonSolver
//
// Copyright (c) 2010 - 2013, Yves Ineichen, ETH Zürich,
//               2013 - 2015, Tülin Kaman, Paul Scherrer Institut, Villigen PSI, Switzerland
// All rights reserved
//
// Implemented as part of the PhD thesis
frey_m's avatar
frey_m committed
15
// "Toward massively parallel multi-objective optimization with application to
frey_m's avatar
frey_m committed
16 17 18 19 20 21 22 23 24 25 26 27
// particle accelerators" (https://doi.org/10.3929/ethz-a-009792359)
//
// This file is part of OPAL.
//
// OPAL is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// You should have received a copy of the GNU General Public License
// along with OPAL. If not, see <https://www.gnu.org/licenses/>.
//
28

29
//#define DBG_STENCIL
30

kraus's avatar
kraus committed
31
#include "Solvers/MGPoissonSolver.h"
32 33 34 35 36 37

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

38
#include "Track/Track.h"
gsell's avatar
gsell committed
39 40 41 42 43
#include "Physics/Physics.h"
#include "Utilities/OpalException.h"
#include "Attributes/Attributes.h"
#include "ValueDefinitions/RealVariable.h"
#include "AbstractObjects/OpalData.h"
44
#include "Utilities/Options.h"
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68

#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"

gsell's avatar
gsell committed
69 70
#include <algorithm>

71 72
using Teuchos::RCP;
using Teuchos::rcp;
gsell's avatar
gsell committed
73 74
using Physics::c;

75
MGPoissonSolver::MGPoissonSolver ( PartBunch *beam,
76 77 78 79 80 81 82
                                   Mesh_t *mesh,
                                   FieldLayout_t *fl,
                                   std::vector<BoundaryGeometry *> geometries,
                                   std::string itsolver,
                                   std::string interpl,
                                   double tol,
                                   int maxiters,
83
                                   std::string precmode):
84 85 86
    geometries_m(geometries),
    tol_m(tol),
    maxiters_m(maxiters),
87 88 89 90
    Comm(Ippl::getComm()),
    itsBunch_m(beam),
    mesh_m(mesh),
    layout_m(fl) {
91

gsell's avatar
gsell committed
92 93
    domain_m = layout_m->getDomain();

94
    for (int i = 0; i < 3; i++) {
gsell's avatar
gsell committed
95 96 97 98 99
        hr_m[i] = mesh_m->get_meshSpacing(i);
        orig_nr_m[i] = domain_m[i].length();
    }

    if(itsolver == "CG") itsolver_m = AZ_cg;
100
    else if(itsolver == "BICGSTAB") itsolver_m = AZ_bicgstab;
gsell's avatar
gsell committed
101 102 103 104 105 106 107
    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;
108
    else if(precmode == "NO") precmode_m = NO;
gsell's avatar
gsell committed
109 110 111 112 113 114

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

    hasParallelDecompositionChanged_m = true;
115
    repartFreq_m = 1000;
gsell's avatar
gsell committed
116
    useRCB_m = false;
kraus's avatar
kraus committed
117
    if(Ippl::Info->getOutputLevel() > 3)
gsell's avatar
gsell committed
118 119 120 121
        verbose_m = true;
    else
        verbose_m = false;

122
    // Find CURRENT geometry
gsell's avatar
gsell committed
123 124
    currentGeometry = geometries_m[0];
    if(currentGeometry->getFilename() == "") {
125
        if(currentGeometry->getTopology() == "ELLIPTIC"){
frey_m's avatar
frey_m committed
126
            bp_m = std::unique_ptr<IrregularDomain>(new EllipticDomain(currentGeometry, orig_nr_m, hr_m, interpl));
kraus's avatar
kraus committed
127
        } else if (currentGeometry->getTopology() == "BOXCORNER") {
frey_m's avatar
frey_m committed
128 129
            bp_m = std::unique_ptr<IrregularDomain>(new BoxCornerDomain(currentGeometry->getA(), currentGeometry->getB(), currentGeometry->getC(), currentGeometry->getLength(),currentGeometry->getL1(), currentGeometry->getL2(), orig_nr_m, hr_m, interpl));
            bp_m->compute(itsBunch_m->get_hr());
gsell's avatar
gsell committed
130
        } else {
131 132
            throw OpalException("MGPoissonSolver::MGPoissonSolver",
                                "Geometry not known");
gsell's avatar
gsell committed
133
        }
134 135 136 137 138 139 140 141 142 143
    } 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");
        }
frey_m's avatar
frey_m committed
144
        bp_m = std::unique_ptr<IrregularDomain>(new ArbitraryDomain(currentGeometry, orig_nr_m, hr_m, interpl));
145
    }
146

gsell's avatar
gsell committed
147 148 149 150
    Map = 0;
    A = Teuchos::null;
    LHS = Teuchos::null;
    RHS = Teuchos::null;
151
    MLPrec = 0;
152
    prec_m = Teuchos::null;
153

gsell's avatar
gsell committed
154 155
    numBlocks_m = Options::numBlocks;
    recycleBlocks_m = Options::recycleBlocks;
156
    nLHS_m = Options::nLHS;
gsell's avatar
gsell committed
157 158
    SetupMLList();
    SetupBelosList();
159
    problem_ptr = rcp(new Belos::LinearProblem<ST, MV, OP>);
160
    // setup Belos solver
gsell's avatar
gsell committed
161
    if(numBlocks_m == 0 || recycleBlocks_m == 0)
162
        solver_ptr = rcp(new Belos::BlockCGSolMgr<double, MV, OP>());
gsell's avatar
gsell committed
163
    else
164 165
        solver_ptr = rcp(new Belos::RCGSolMgr<double, MV, OP>());
    solver_ptr->setParameters(rcp(&belosList, false));
gsell's avatar
gsell committed
166 167 168 169
    convStatusTest = rcp(new Belos::StatusTestGenResNorm<ST, MV, OP> (tol));
    convStatusTest->defineScaleForm(Belos::NormOfRHS, Belos::TwoNorm);

    //all timers used here
170 171 172 173 174 175 176 177
    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
178 179
}

180
void MGPoissonSolver::deletePtr() {
181 182 183 184 185 186 187
    delete Map;
    Map = nullptr;
    delete MLPrec;
    MLPrec = nullptr;
    A      = Teuchos::null;
    LHS    = Teuchos::null;
    RHS    = Teuchos::null;
188 189 190
    prec_m = Teuchos::null;
}

191 192 193 194
MGPoissonSolver::~MGPoissonSolver() {
    deletePtr ();
    solver_ptr = Teuchos::null;
    problem_ptr = Teuchos::null;
gsell's avatar
gsell committed
195 196
}

197
void MGPoissonSolver::computePotential(Field_t& /*rho*/, Vector_t /*hr*/, double /*zshift*/) {
gsell's avatar
gsell committed
198 199 200
    throw OpalException("MGPoissonSolver", "not implemented yet");
}

201
void MGPoissonSolver::computeMap(NDIndex<3> localId) {
202
    if (itsBunch_m->getLocalTrackStep()%repartFreq_m == 0){
203
        deletePtr();
gsell's avatar
gsell committed
204
        if(useRCB_m)
205
            redistributeWithRCB(localId);
206
        else
207
            IPPLToMap3D(localId);
208 209

        extrapolateLHS();
210 211
    }
}
gsell's avatar
gsell committed
212

213 214 215 216 217
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.
218
    //we also have to redistribute LHS
219

220 221 222 223 224 225 226 227 228
    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
229

230
    //...and all previously saved LHS
231
    std::deque< Epetra_Vector >::iterator it = OldLHS.begin();
232 233 234 235 236 237 238 239 240 241
    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
242

243 244 245
    // extrapolate last OldLHS.size LHS to get a new start vector
    it = OldLHS.begin();
    if(nLHS_m == 0 || OldLHS.size()==0)
246 247 248
        LHS->PutScalar(1.0);
    else if(OldLHS.size() == 1)
        *LHS = *it;
249 250 251
    else if(OldLHS.size() == 2)
        LHS->Update(2.0, *it++, -1.0, *it, 0.0);
    else if(OldLHS.size() > 2) {
252
        int n = OldLHS.size();
253 254 255
        P = rcp(new Epetra_MultiVector(*Map, nLHS_m, false));
        for(int i = 0; i < n; ++i) {
           *(*P)(i) = *it++;
256
        }
257 258 259 260
        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);
           }
261 262
        }
        *LHS = *(*P)(0);
263
     } else
kraus's avatar
kraus committed
264
        throw OpalException("MGPoissonSolver", "Invalid number of old LHS: " + std::to_string(OldLHS.size()));
gsell's avatar
gsell committed
265 266
}

267

gsell's avatar
gsell committed
268 269 270 271 272
// 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) {
273

kraus's avatar
kraus committed
274
    Inform msg("OPAL ", INFORM_ALL_NODES);
gsell's avatar
gsell committed
275 276 277
    nr_m[0] = orig_nr_m[0];
    nr_m[1] = orig_nr_m[1];
    nr_m[2] = orig_nr_m[2];
278

frey_m's avatar
frey_m committed
279 280 281
    bp_m->setGlobalMeanR(itsBunch_m->getGlobalMeanR());
    bp_m->setGlobalToLocalQuaternion(itsBunch_m->getGlobalToLocalQuaternion());
    bp_m->setNr(nr_m);
282

283 284
    NDIndex<3> localId = layout_m->getLocalNDIndex();

285
    IpplTimings::startTimer(FunctionTimer1_m);
286
    // Compute boundary intersections (only do on the first step)
287
    if(!itsBunch_m->getLocalTrackStep())
frey_m's avatar
frey_m committed
288
        bp_m->compute(hr, localId);
289
    IpplTimings::stopTimer(FunctionTimer1_m);
290 291

    // Define the Map
kraus's avatar
kraus committed
292
    INFOMSG(level3 << "* Computing Map..." << endl);
293
    IpplTimings::startTimer(FunctionTimer2_m);
294
    computeMap(localId);
295
    IpplTimings::stopTimer(FunctionTimer2_m);
kraus's avatar
kraus committed
296
    INFOMSG(level3 << "* Done." << endl);
gsell's avatar
gsell committed
297

298
    // Allocate the RHS with the new Epetra Map
299
    if (Teuchos::is_null(RHS))
300
        RHS = rcp(new Epetra_Vector(*Map));
301
    RHS->PutScalar(0.0);
302

kraus's avatar
kraus committed
303 304 305 306 307 308
    // // get charge densities from IPPL field and store in Epetra vector (RHS)
    if (verbose_m) {
        Ippl::Comm->barrier();
        msg << "* Node:" << Ippl::myNode() << ", Filling RHS..." << endl;
        Ippl::Comm->barrier();
    }
309 310
    IpplTimings::startTimer(FunctionTimer3_m);
    int id = 0;
311
    float scaleFactor = itsBunch_m->getdT();
312

313

kraus's avatar
kraus committed
314 315
    if (verbose_m) {
        msg << "* Node:" << Ippl::myNode() << ", Rho for final element: " << rho[localId[0].last()][localId[1].last()][localId[2].last()].get() << endl;
316

kraus's avatar
kraus committed
317 318 319 320 321 322
        Ippl::Comm->barrier();
        msg << "* Node:" << Ippl::myNode() << ", Local nx*ny*nz = " <<  localId[2].last() *  localId[0].last() *  localId[1].last() << endl;
        msg << "* Node:" << Ippl::myNode() << ", Number of reserved local elements in RHS: " << RHS->MyLength() << endl;
        msg << "* Node:" << Ippl::myNode() << ", Number of reserved global elements in RHS: " << RHS->GlobalLength() << endl;
        Ippl::Comm->barrier();
    }
323
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
324
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
kraus's avatar
kraus committed
325
            for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
frey_m's avatar
frey_m committed
326
                if (bp_m->isInside(idx, idy, idz))
327
                        RHS->Values()[id++] = 4*M_PI*rho[idx][idy][idz].get()/scaleFactor;
gsell's avatar
gsell committed
328 329 330
            }
        }
    }
331

332
    IpplTimings::stopTimer(FunctionTimer3_m);
kraus's avatar
kraus committed
333 334 335 336 337 338
    if (verbose_m) {
        Ippl::Comm->barrier();
        msg << "* Node:" << Ippl::myNode() << ", Number of Local Inside Points " << id << endl;
        msg << "* Node:" << Ippl::myNode() << ", Done." << endl;
        Ippl::Comm->barrier();
    }
gsell's avatar
gsell committed
339
    // build discretization matrix
kraus's avatar
kraus committed
340
    INFOMSG(level3 << "* Building Discretization Matrix..." << endl);
341
    IpplTimings::startTimer(FunctionTimer4_m);
342
    if(Teuchos::is_null(A))
343
        A = rcp(new Epetra_CrsMatrix(Copy, *Map,  7, true));
344
    ComputeStencil(hr, RHS);
345
    IpplTimings::stopTimer(FunctionTimer4_m);
kraus's avatar
kraus committed
346
    INFOMSG(level3 << "* Done." << endl);
347

gsell's avatar
gsell committed
348 349 350
#ifdef DBG_STENCIL
    EpetraExt::RowMatrixToMatlabFile("DiscrStencil.dat", *A);
#endif
351

kraus's avatar
kraus committed
352
    INFOMSG(level3 << "* Computing Preconditioner..." << endl);
353
    IpplTimings::startTimer(FunctionTimer5_m);
354
    if(!MLPrec) {
355 356 357 358
        MLPrec = new ML_Epetra::MultiLevelPreconditioner(*A, MLList_m);
    } else if (precmode_m == REUSE_HIERARCHY) {
        MLPrec->ReComputePreconditioner();
    } else if (precmode_m == REUSE_PREC){
359
    }
360
    IpplTimings::stopTimer(FunctionTimer5_m);
kraus's avatar
kraus committed
361
    INFOMSG(level3 << "* Done." << endl);
362

gsell's avatar
gsell committed
363 364
    // setup preconditioned iterative solver
    // use old LHS solution as initial guess
kraus's avatar
kraus committed
365
    INFOMSG(level3 << "* Final Setup of Solver..." << endl);
366
    IpplTimings::startTimer(FunctionTimer6_m);
367 368 369
    problem_ptr->setOperator(A);
    problem_ptr->setLHS(LHS);
    problem_ptr->setRHS(RHS);
370 371
    if(Teuchos::is_null(prec_m))
        prec_m = Teuchos::rcp ( new Belos::EpetraPrecOp ( rcp(MLPrec,false)));
372 373 374 375
    problem_ptr->setLeftPrec(prec_m);
    solver_ptr->setProblem( problem_ptr);
    if (!problem_ptr->isProblemSet()) {
        if (problem_ptr->setProblem() == false) {
376
            ERRORMSG("Belos::LinearProblem failed to set up correctly!" << endl);
377
        }
gsell's avatar
gsell committed
378
    }
379
    IpplTimings::stopTimer(FunctionTimer6_m);
kraus's avatar
kraus committed
380
    INFOMSG(level3 << "* Done." << endl);
381

gsell's avatar
gsell committed
382 383
    double time = MPI_Wtime();

kraus's avatar
kraus committed
384
    INFOMSG(level3 << "* Solving for Space Charge..." << endl);
385
    IpplTimings::startTimer(FunctionTimer7_m);
kraus's avatar
kraus committed
386
    solver_ptr->solve();
387
    IpplTimings::stopTimer(FunctionTimer7_m);
kraus's avatar
kraus committed
388
    INFOMSG(level3 << "* Done." << endl);
gsell's avatar
gsell committed
389

kraus's avatar
kraus committed
390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416
    std::ofstream timings;
    if (true || verbose_m) {
        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) {
            char filename[50];
            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);
            timings.open(filename, std::ios::app);
            timings << solver_ptr->getNumIters() << "\t"
                    //<< time <<" "<<
                    << minTime << "\t"
                    << maxTime << "\t"
                    << avgTime << "\t"
                    << numBlocks_m << "\t"
                    << recycleBlocks_m << "\t"
                    << nLHS_m << "\t"
                    //<< OldLHS.size() <<"\t"<<
                    << std::endl;

            timings.close();
        }

    }
gsell's avatar
gsell committed
417
    // Store new LHS in OldLHS
418 419
    if(nLHS_m > 1) OldLHS.push_front(*(LHS.get()));
    if(OldLHS.size() > nLHS_m) OldLHS.pop_back();
gsell's avatar
gsell committed
420 421

    //now transfer solution back to IPPL grid
422 423
    IpplTimings::startTimer(FunctionTimer8_m);
    id = 0;
gsell's avatar
gsell committed
424
    rho = 0.0;
425
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
426
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
kraus's avatar
kraus committed
427
            for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
428
                  NDIndex<3> l(Index(idx, idx), Index(idy, idy), Index(idz, idz));
frey_m's avatar
frey_m committed
429
                  if (bp_m->isInside(idx, idy, idz))
430
                     rho.localElement(l) = LHS->Values()[id++]*scaleFactor;
gsell's avatar
gsell committed
431 432 433
            }
        }
    }
434
    IpplTimings::stopTimer(FunctionTimer8_m);
435 436 437 438 439 440

    if(itsBunch_m->getLocalTrackStep()+1 == (long long)Track::block->localTimeSteps.front()) {
        A = Teuchos::null;
        LHS = Teuchos::null;
        RHS = Teuchos::null;
        prec_m = Teuchos::null;
441
    }
gsell's avatar
gsell committed
442 443 444
}


445
void MGPoissonSolver::redistributeWithRCB(NDIndex<3> localId) {
gsell's avatar
gsell committed
446 447 448

    int numMyGridPoints = 0;

449
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
450
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
kraus's avatar
kraus committed
451
            for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
frey_m's avatar
frey_m committed
452
                if (bp_m->isInside(idx, idy, idz))
gsell's avatar
gsell committed
453
                    numMyGridPoints++;
454
             }
gsell's avatar
gsell committed
455
        }
456
     }
gsell's avatar
gsell committed
457 458 459 460 461 462 463 464 465 466

    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;

467
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
468
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
kraus's avatar
kraus committed
469
            for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
frey_m's avatar
frey_m committed
470
                if (bp_m->isInside(idx, idy, idz)) {
471 472 473
                    v[0] = (double)idx;
                    v[stride] = (double)idy;
                    v[stride2] = (double)idz;
gsell's avatar
gsell committed
474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494
                    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++) {
frey_m's avatar
frey_m committed
495
        MyGlobalElements.push_back(bp_m->getIdx(v[0], v[stride], v[stride2]));
gsell's avatar
gsell committed
496 497 498 499 500 501 502
        v++;
        numMyGridPoints++;
    }

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

503
void MGPoissonSolver::IPPLToMap3D(NDIndex<3> localId) {
504

gsell's avatar
gsell committed
505
    int NumMyElements = 0;
506
    std::vector<int> MyGlobalElements;
gsell's avatar
gsell committed
507

508
    for (int idz = localId[2].first(); idz <= localId[2].last(); idz++) {
509
        for (int idy = localId[1].first(); idy <= localId[1].last(); idy++) {
510
            for (int idx = localId[0].first(); idx <= localId[0].last(); idx++) {
frey_m's avatar
frey_m committed
511 512
                if (bp_m->isInside(idx, idy, idz)) {
                    MyGlobalElements.push_back(bp_m->getIdx(idx, idy, idz));
gsell's avatar
gsell committed
513 514 515
                    NumMyElements++;
                }
            }
kraus's avatar
kraus committed
516
        }
517
    }
gsell's avatar
gsell committed
518 519 520
    Map = new Epetra_Map(-1, NumMyElements, &MyGlobalElements[0], 0, Comm);
}

521
void MGPoissonSolver::ComputeStencil(Vector_t /*hr*/, Teuchos::RCP<Epetra_Vector> RHS) {
gsell's avatar
gsell committed
522

523 524
    A->PutScalar(0.0);

gsell's avatar
gsell committed
525 526 527
    int NumMyElements = Map->NumMyElements();
    int *MyGlobalElements = Map->MyGlobalElements();

528 529
    std::vector<double> Values(6);
    std::vector<int> Indices(6);
gsell's avatar
gsell committed
530

531
    for(int i = 0 ; i < NumMyElements ; i++) {
gsell's avatar
gsell committed
532 533 534

        int NumEntries = 0;

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

frey_m's avatar
frey_m committed
538
        bp_m->getBoundaryStencil(MyGlobalElements[i], WV, EV, SV, NV, FV, BV, CV, scaleFactor);
539
        RHS->Values()[i] *= scaleFactor;
gsell's avatar
gsell committed
540

frey_m's avatar
frey_m committed
541
        bp_m->getNeighbours(MyGlobalElements[i], W, E, S, N, F, B);
gsell's avatar
gsell committed
542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566
        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;
        }

kraus's avatar
kraus committed
567
        // if matrix has already been filled (FillComplete()) we can only
gsell's avatar
gsell committed
568
        // replace entries
569

570 571 572 573 574
        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);
575
        } else {
gsell's avatar
gsell committed
576 577 578
            // off-diagonal entries
            A->InsertGlobalValues(MyGlobalElements[i], NumEntries, &Values[0], &Indices[0]);
            // diagonal entry
579
            A->InsertGlobalValues(MyGlobalElements[i], 1, &CV, MyGlobalElements + i);
580
        }
gsell's avatar
gsell committed
581 582 583
    }

    A->FillComplete();
584

gsell's avatar
gsell committed
585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607
    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();
kraus's avatar
kraus committed
608 609
    *gmsg << "LBAL min = " << min << ", max = " << max << ", avg = " << avg << endl;
    *gmsg << "min nr gridpoints = " << minn << ", max nr gridpoints = " << maxn << endl;
gsell's avatar
gsell committed
610 611 612 613 614 615
}

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;
616
    return os;
gsell's avatar
gsell committed
617 618
}

619 620 621 622 623 624 625
// vi: set et ts=4 sw=4 sts=4:
// Local Variables:
// mode:c
// c-basic-offset: 4
// indent-tabs-mode: nil
// require-final-newline: nil
// End: