NeighborSearch/src/nearinterac.cpp

#include <iostream>
#include <NearInteraction.hpp>

template <class Real, int DIM> class MyObject {
  public:
    MyObject() {
      for (int i = 0; i < DIM; i++) coord[i] = drand48() * 3;
      rad = drand48() * 0.01;
    }

    const Real* Coord() const { return coord; }

    Real Rad() const { return rad; }

    void Pack(std::vector<char>& buff) const {
      size_t count = sizeof(MyObject<Real, DIM>);
      buff.resize(count);
      memcpy(buff.data(), this, count);
    }

    void Unpack(const std::vector<char>& buff) {
      size_t count = sizeof(MyObject<Real, DIM>);
      memcpy(this, buff.data(), count);
    }

  private:
    Real coord[DIM];
    Real rad;
};

typedef double Real;
constexpr int DIM = 3;
typedef MyObject<Real,DIM> SrcObj;
typedef MyObject<Real,DIM> TrgObj;

int main(int argc, char **argv) {
  int comm_rank = 0;
  int comm_size = 1;
#ifdef SCTL_HAVE_MPI
  MPI_Init(&argc, &argv);
  MPI_Comm mpi_comm = MPI_COMM_WORLD;
  MPI_Comm_rank(mpi_comm, &comm_rank);
  MPI_Comm_size(mpi_comm, &comm_size);
#endif

  // Generate source and target points
  srand48(comm_rank);
  std::vector<SrcObj> src_vec(10000/comm_size);
  std::vector<TrgObj> trg_vec(20000/comm_size);

  { // Compute near interactions
    sctl::Comm comm;
    #ifdef SCTL_HAVE_MPI
    comm = mpi_comm;
    #endif

    NearInteraction<Real,DIM> near_interac(comm);
    sctl::Profile::Enable(true);
    { // Repartition data
      // Setup for repartition
      sctl::Profile::Tic("RepartSetup", &comm, true);
      near_interac.SetupRepartition(src_vec, trg_vec);
      sctl::Profile::Toc();

      // Distribute source and target vectors
      sctl::Profile::Tic("Repart", &comm, true);
      std::vector<SrcObj> src_new;
      std::vector<TrgObj> trg_new;
      near_interac.ForwardScatterSrc<SrcObj>(src_vec, src_new);
      near_interac.ForwardScatterTrg<TrgObj>(trg_vec, trg_new);
      src_vec.swap(src_new);
      trg_vec.swap(trg_new);
      sctl::Profile::Toc();
    }
    { // Compute near interactions
      // Setup for near interaction
      sctl::Profile::Tic("NearSetup", &comm, true);
      near_interac.SetupNearInterac(src_vec, trg_vec);
      sctl::Profile::Toc();

      // Following code can repeat multiple times without calling Setup again as long as
      // the position and shape of the sources and the targets do not change.

      // Forward scatter
      sctl::Profile::Tic("Near", &comm, true);
      std::vector<SrcObj> src_near;
      std::vector<TrgObj> trg_near;
      near_interac.ForwardScatterSrc<SrcObj>(src_vec, src_near);
      near_interac.ForwardScatterTrg<TrgObj>(trg_vec, trg_near);

      const auto& trg_src_interac = near_interac.GetInteractionList(); // Get interaction list
      //for (auto interac : trg_src_interac) { // compute near interactions (single thread)
      //  SrcObj& t = trg_near[interac.first];
      //  TrgObj& s = src_near[interac.second];
      //  // ( compute interaction between t and s )
      //}
      #pragma omp parallel
      { // compute near interactions (multiple threads)
        int omp_p = omp_get_num_threads();
        int tid = omp_get_thread_num();
        long N = trg_src_interac.size();
        long a = (tid + 0) * N / omp_p;
        long b = (tid + 1) * N / omp_p;

        // Ensure each thread works on a different target
        if (tid > 0) while (a + 1 < N && trg_src_interac[a].first == trg_src_interac[a + 1].first) a++;
        while (b + 1 < N && trg_src_interac[b].first == trg_src_interac[b + 1].first) b++;

        for (long i = a; i < b; i++) {
          SrcObj& t = trg_near[trg_src_interac[i].first];
          TrgObj& s = src_near[trg_src_interac[i].second];
          // ( compute interaction between t and s )
        }
      }

      // Reverse scatter
      near_interac.ReverseScatterTrg<TrgObj>(trg_near, trg_vec);
      sctl::Profile::Toc();
    }
    sctl::Profile::print(&comm);
  }

#ifdef SCTL_HAVE_MPI
  MPI_Finalize();
#endif
  return 0;
}