main.cpp

#include "box.hpp"
#include "fmm.hpp"
#include "matrix.hpp"
#include "timer.hpp"
#include <vector>
#include <iostream>
#include <cstdlib>

constexpr int dim = 3;
using Vector_t = Vector<double, dim>;


int main(int argc, char** argv)
{
    MPI_Init(&argc, &argv);

    int size, rank;
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    // The direct comparison is O(N^2), so N is worth lowering for a quick check.
    // The box orientation belongs on the command line because the field is the
    // one quantity that is not invariant under it.
    const size_t N = argc > 1 ? std::strtoul(argv[1], nullptr, 10) : 200000;
    const int vertical_axis = argc > 2 ? std::atoi(argv[2]) : dim - 1;
    const double phideg = argc > 3 ? std::atof(argv[3]) : 0;

    // Initialize random positions in a cube [0, 1]^3
    svector<Vector_t> positions(get_shm_comm());
    std::vector<Vector_t> vector_charges;
    positions.reserve(N);
    vector_charges.reserve(N);
    for (size_t i=0; i<N; ++i) {
        if (positions.root()) {
            positions.emplace_back(Vector_t{static_cast<double>(std::rand())/RAND_MAX,
                                            static_cast<double>(std::rand())/RAND_MAX,
                                            static_cast<double>(std::rand())/RAND_MAX});
        }
        vector_charges.push_back({static_cast<double>(std::rand())/RAND_MAX,
                                  static_cast<double>(std::rand())/RAND_MAX,
                                  static_cast<double>(std::rand())/RAND_MAX});
    }
    positions.sync_all();

    // Define simulation domain by a cube
    Box<double, dim, Vector> box(positions, vertical_axis, phideg);
    if (rank == 0) {
        std::cout << box << std::endl;
    }

    // Adaptively partition the simulation domain
    int max_level = 20;
    int max_particles_per_node = 128;
    OctreePartitioner<double, Vector> partitioner(positions, box);
    partitioner.refine(max_level, max_particles_per_node);
    if (rank == 0) {
        std::cout << partitioner << std::endl;
    }

    // Precompute FMM
    constexpr int expansion_order = 8;
    FMM3D<double, dim, expansion_order, Vector, true> fmm(partitioner, MPI_COMM_WORLD);

    // Calculation via FMM
    std::vector<Vector_t> vector_potentials_fmm(N);
    std::vector<Vector_t> vector_fields_fmm(N);

    tic("FMM potential calculation");
    fmm.rinv(vector_charges.data(), vector_potentials_fmm.data());
    MPI_Allreduce(MPI_IN_PLACE, vector_potentials_fmm.data(), dim*N, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
    toc("FMM potential calculation");

    tic("FMM field calculation");
    fmm.rinv<true>(vector_charges.data(), vector_fields_fmm.data());
    MPI_Allreduce(MPI_IN_PLACE, vector_fields_fmm.data(), dim*N, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
    toc("FMM field calculation");

    // The far field through shared buffers, against the plain call with the
    // same sources supplied by rank 0 alone.
    std::vector<Vector_t> Q_plain(N);
    std::vector<Vector_t> U_plain(fmm.n_result());
    svector<Vector_t> Q_shared(MPI_COMM_WORLD, N);
    svector<Vector_t> U_shared(MPI_COMM_WORLD, fmm.n_result());
    if (rank == 0) {
        Q_plain = vector_charges;
        std::copy(vector_charges.cbegin(), vector_charges.cend(), Q_shared.data());
    }
    tic("FMM far field, plain");
    fmm.rinv_far(Q_plain.data(), U_plain.data());
    toc("FMM far field, plain");
    tic("FMM far field, shared");
    fmm.rinv_far(Q_shared, U_shared);
    toc("FMM far field, shared");
    if (rank == 0) {
        double diff = 0;
        for (int i=0; i<fmm.n_result(); ++i) {
            diff = std::max(diff, (U_plain[i] - U_shared[i]).norm());
        }
        std::cout << "Max difference between plain and shared far field = " << diff << std::endl;
    }

    // Calculation via direct method
    svector<Vector_t> vector_potentials_direct(MPI_COMM_WORLD, N);
    svector<Vector_t> vector_fields_direct(MPI_COMM_WORLD, N);

    tic("Direct potential calculation");
    for (size_t i=begin(N, size, rank); i<end(N, size, rank); ++i) {
        for (size_t j=0; j<N; ++j) {
            if (i != j) {
                const auto R = positions[i] - positions[j];
                vector_potentials_direct[i] += vector_charges[j]/R.norm();
            }
        }
    }
    vector_potentials_direct.sync(); vector_potentials_direct.template allreduce<double>();
    toc("Direct potential calculation");

    tic("Direct field calculation");
    for (size_t i=begin(N, size, rank); i<end(N, size, rank); ++i) {
        for (size_t j=0; j<N; ++j) {
            if (i != j) {
                const auto R = positions[i] - positions[j];
                vector_fields_direct[i] += vector_charges[j].cross(R)/std::pow(R.norm(), 3);
            }
        }
    }
    vector_fields_direct.sync(); vector_fields_direct.template allreduce<double>();
    toc("Direct field calculation");

    // Calculate error
    if (rank == 0) {
        double numer_potential = 0;
        double denom_potential = 0;
        double numer_field = 0;
        double denom_field = 0;
        for (size_t i=0; i<N; ++i) {
            numer_potential += std::pow((vector_potentials_direct[i] - vector_potentials_fmm[i]).norm(), 2);
            denom_potential += std::pow(vector_potentials_direct[i].norm(), 2);
            numer_field += std::pow((vector_fields_direct[i] - vector_fields_fmm[i]).norm(), 2);
            denom_field += std::pow(vector_fields_direct[i].norm(), 2);
        }
        std::cout << "Relative l2 error of potential = " << std::sqrt(numer_potential/denom_potential) << std::endl;
        std::cout << "Relative l2 error of field = " << std::sqrt(numer_field/denom_field) << std::endl;
    }

    MPI_Finalize();
}