LoadBalancer.hpp

Go to the documentation of this file.

#ifndef OPAL_LOAD_BALANCER_H
#define OPAL_LOAD_BALANCER_H
 
#include <memory>
 
#include "PartBunch/ParticleContainer.hpp"
 
template <typename T, unsigned Dim>
using ORB = ippl::OrthogonalRecursiveBisection<Field<double, Dim>, T>;
 
template <typename T, unsigned Dim>
class LoadBalancer {
    using Base          = ippl::ParticleBase<ippl::ParticleSpatialLayout<T, Dim>>;
    using FieldSolver_t = ippl::FieldSolverBase<T, Dim>;
 
private:
    double loadbalancethreshold_m;
    Field_t<Dim>* rho_m;
    VField_t<T, Dim>* E_m;
    Field<T, Dim>* phi_m;
    std::shared_ptr<ParticleContainer<T, Dim>> pc_m;
    std::shared_ptr<FieldSolver_t> fs_m;
    unsigned int loadbalancefreq_m;
    ORB<T, Dim> orb;
 
public:
    LoadBalancer(
        double lbs, std::shared_ptr<FieldContainer<T, Dim>>& fc,
        std::shared_ptr<ParticleContainer<T, Dim>>& pc, std::shared_ptr<FieldSolver_t>& fs)
        : loadbalancethreshold_m(lbs),
          rho_m(&fc->getRho()),
          E_m(&fc->getE()),
          phi_m(&fc->getPhi()),
          pc_m(pc),
          fs_m(fs) {
    }
 
    ~LoadBalancer() {
    }
 
    double getLoadBalanceThreshold() const {
        return loadbalancethreshold_m;
    }
    void setLoadBalanceThreshold(double threshold) {
        loadbalancethreshold_m = threshold;
    }
 
    Field_t<Dim>* getRho() const {
        return rho_m;
    }
    void setRho(Field_t<Dim>* rho) {
        rho_m = rho;
    }
 
    VField_t<T, Dim>* getE() const {
        return E_m;
    }
    void setE(VField_t<T, Dim>* E) {
        E_m = E;
    }
 
    Field<T, Dim>* getPhi() {
        return phi_m;
    }
    void setPhi(Field<T, Dim>* phi) {
        phi_m = phi;
    }
 
    std::shared_ptr<ParticleContainer<T, Dim>> getParticleContainer() const {
        return pc_m;
    }
    void setParticleContainer(std::shared_ptr<ParticleContainer<T, Dim>> pc) {
        pc_m = pc;
    }
 
    std::shared_ptr<FieldSolver_t> getFieldSolver() const {
        return fs_m;
    }
    void setFieldSolver(std::shared_ptr<FieldSolver_t> fs) {
        fs_m = fs;
    }
 
    void updateLayout(
        ippl::FieldLayout<Dim>* fl, ippl::UniformCartesian<T, Dim>* mesh,
        bool& isFirstRepartition) {
        // Update local fields
 
        static IpplTimings::TimerRef tupdateLayout = IpplTimings::getTimer("updateLayout");
        IpplTimings::startTimer(tupdateLayout);
        (*E_m).updateLayout(*fl);
        (*rho_m).updateLayout(*fl);
 
        if (fs_m->getStype() == "CG") {
            phi_m->updateLayout(*fl);
            phi_m->setFieldBC(phi_m->getFieldBC());
        }
 
        // Update layout with new FieldLayout
        PLayout_t<T, Dim>* layout = &pc_m->getLayout();
        (*layout).updateLayout(*fl, *mesh);
        IpplTimings::stopTimer(tupdateLayout);
        static IpplTimings::TimerRef tupdatePLayout = IpplTimings::getTimer("updatePB");
        IpplTimings::startTimer(tupdatePLayout);
        if (!isFirstRepartition) {
            pc_m->update();
        }
        IpplTimings::stopTimer(tupdatePLayout);
    }
 
    void initializeORB(ippl::FieldLayout<Dim>* fl, ippl::UniformCartesian<T, Dim>* mesh) {
        orb.initialize(*fl, *mesh, *rho_m);
    }
 
    void repartition(
        ippl::FieldLayout<Dim>* fl, ippl::UniformCartesian<T, Dim>* mesh,
        bool& isFirstRepartition) {
        // Repartition the domains
 
 
        using Base = ippl::ParticleBase<ippl::ParticleSpatialLayout<T, Dim>>;
        typename Base::particle_position_type *R;
        R = &pc_m->R;
        bool res = orb.binaryRepartition(*R, *fl, isFirstRepartition);
        if (res != true) {
            std::cout << "Could not repartition!" << std::endl;
            return;
        }
        // Update                                                                                                                                                                                                                                                                                   
        this->updateLayout(fl, mesh, isFirstRepartition);
        if constexpr (Dim == 2 || Dim == 3) {
                if (fs_m->getStype() == "FFT") {
                    std::get<FFTSolver_t<T, Dim>>(fs_m->getSolver()).setRhs(*rho_m);
                }
                if constexpr (Dim == 3) {
                        if (fs_m->getStype() == "TG") {
                            std::get<FFTTruncatedGreenSolver_t<T, Dim>>(fs_m->getSolver()).setRhs(*rho_m);
                        } else if (fs_m->getStype() == "OPEN") {
                            std::get<OpenSolver_t<T, Dim>>(fs_m->getSolver()).setRhs(*rho_m);
                        }
                    }
            }
        
    }
 
    bool balance(size_type totalP, const unsigned int nstep) {
        if (ippl::Comm->size() < 2) {
            return false;
        } else {
            int local = 0;
            std::vector<int> res(ippl::Comm->size());
            double equalPart = (double)totalP / ippl::Comm->size();
            double dev       = std::abs((double)pc_m->getLocalNum() - equalPart) / totalP;
            if (dev > loadbalancethreshold_m) {
                local = 1;
            }
            MPI_Allgather(
                &local, 1, MPI_INT, res.data(), 1, MPI_INT, ippl::Comm->getCommunicator());
 
            for (unsigned int i = 0; i < res.size(); i++) {
                if (res[i] == 1) {
                    return true;
                }
            }
            return false;
        }
    }
};
 
#endif