SortedSoAView.h

Go to the documentation of this file.

 
#pragma once
 
#include <algorithm>
#include <array>
#include <utility>
#include <vector>
 
#include "autopas/utils/SoAView.h"
 
namespace autopas {
 
template <class Particle_T, class Functor_T>
class SortedSoAView {
  using SoAArraysType = Particle_T::SoAArraysType;
 
 public:
  SortedSoAView(SoAView<SoAArraysType> source, const std::array<double, 3> &sortingDirection,
                SoA<SoAArraysType> &cachedSoa, std::vector<std::pair<double, size_t>> &cachedProjIdx)
      : _source(source), _sortedSoa(cachedSoa), _projIdx(cachedProjIdx) {
    const size_t n = source.size();
    if (n == 0) {
      _sortedSoa.resizeArrays(0);
      return;
    }
 
    // Step 1: project each particle onto sortingDirection and sort indices ascending.
    _projIdx.resize(n);
    const auto *xPtr = source.template begin<Particle_T::AttributeNames::posX>();
    const auto *yPtr = source.template begin<Particle_T::AttributeNames::posY>();
    const auto *zPtr = source.template begin<Particle_T::AttributeNames::posZ>();
    for (size_t i = 0; i < n; ++i) {
      _projIdx[i] = {xPtr[i] * sortingDirection[0] + yPtr[i] * sortingDirection[1] + zPtr[i] * sortingDirection[2], i};
    }
    std::sort(_projIdx.begin(), _projIdx.end(), [](const auto &a, const auto &b) { return a.first < b.first; });
 
    // Step 2: resize all arrays (value-initializes to 0, covering the computed attributes).
    _sortedSoa.resizeArrays(n);
 
    // Step 3: pack getNeededAttr() attributes from source in sorted order.
    packNeededImpl(std::make_index_sequence<Functor_T::getNeededAttr().size()>{}, n);
 
    // Step 4: zero getComputedAttr() attributes so the functor accumulates from 0.
    zeroComputedImpl(std::make_index_sequence<Functor_T::getComputedAttr().size()>{}, n);
  }
 
  SoAView<SoAArraysType> getView() { return _sortedSoa.constructView(); }
 
  void scatterBack() {
    scatterBackImpl(std::make_index_sequence<Functor_T::getComputedAttr().size()>{}, _projIdx.size());
  }
 
 private:
  template <size_t AttrIdx>
  void packAttr(size_t n) {
    constexpr size_t attr = static_cast<size_t>(Functor_T::getNeededAttr()[AttrIdx]);
    auto *dst = _sortedSoa.template begin<attr>();
    const auto *src = _source.template begin<attr>();
    for (size_t i = 0; i < n; ++i) {
      dst[i] = src[_projIdx[i].second];
    }
  }
 
  template <size_t... AttrIdxs>
  void packNeededImpl(std::index_sequence<AttrIdxs...>, size_t n) {
    (packAttr<AttrIdxs>(n), ...);
  }
 
  template <size_t AttrIdx>
  void zeroAttr(size_t n) {
    constexpr size_t attr = static_cast<size_t>(Functor_T::getComputedAttr()[AttrIdx]);
    auto *ptr = _sortedSoa.template begin<attr>();
    std::fill(ptr, ptr + n, 0);
  }
 
  template <size_t... AttrIdxs>
  void zeroComputedImpl(std::index_sequence<AttrIdxs...>, size_t n) {
    (zeroAttr<AttrIdxs>(n), ...);
  }
 
  template <size_t AttrIdx>
  void scatterAttr(size_t n) {
    constexpr size_t attr = static_cast<size_t>(Functor_T::getComputedAttr()[AttrIdx]);
    auto *orig = _source.template begin<attr>();
    const auto *sorted = _sortedSoa.template begin<attr>();
    for (size_t i = 0; i < n; ++i) {
      orig[_projIdx[i].second] += sorted[i];
    }
  }
 
  template <size_t... AttrIdxs>
  void scatterBackImpl(std::index_sequence<AttrIdxs...>, size_t n) {
    (scatterAttr<AttrIdxs>(n), ...);
  }
 
  SoAView<SoAArraysType> _source;
  SoA<SoAArraysType> &_sortedSoa;
  std::vector<std::pair<double, size_t>> &_projIdx;
};
 
}  // namespace autopas