GroupSchedule.h

Go to the documentation of this file.

/*
** This file is part of eOn.
**
** SPDX-License-Identifier: BSD-3-Clause
**
** Copyright (c) 2010--present, eOn Development Team
** All rights reserved.
**
** Repo:
** https://github.com/TheochemUI/eOn
*/
#pragma once
 
#include "eon/PbcSimd.h"
 
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <map>
#include <vector>
 
namespace eonc {

class GroupSchedule {
public:
  explicit GroupSchedule(int groups) : groups_{std::max(groups, 1)} {}
 
  [[nodiscard]] int groups() const noexcept { return groups_; }

  [[nodiscard]] int defaultHome(std::int64_t key, long position) const {
    const std::int64_t base = key >= 0 ? key : position;
    return static_cast<int>(((base % groups_) + groups_) % groups_);
  }

  [[nodiscard]] std::vector<int> assign(const std::vector<std::int64_t> &keys) {
    const long m = static_cast<long>(keys.size());
    const long cap = (m + groups_ - 1) / groups_;
    std::vector<long> load(static_cast<size_t>(groups_), 0);
    std::vector<int> route(static_cast<size_t>(m), -1);
    for (long j = 0; j < m; ++j) {
      const int home = homeOf(keys[static_cast<size_t>(j)], j);
      if (load[static_cast<size_t>(home)] < cap) {
        route[static_cast<size_t>(j)] = home;
        ++load[static_cast<size_t>(home)];
      }
    }
    for (long j = 0; j < m; ++j) {
      if (route[static_cast<size_t>(j)] >= 0)
        continue;
      const auto least = std::min_element(load.begin(), load.end());
      const int g = static_cast<int>(least - load.begin());
      route[static_cast<size_t>(j)] = g;
      ++*least;
    }
    for (long j = 0; j < m; ++j)
      home_[keys[static_cast<size_t>(j)]] = route[static_cast<size_t>(j)];
    return route;
  }

  void record(std::int64_t key, const double *positions, long n3,
              const double *cell = nullptr) {
    auto &stored = last_[key];
    stored.assign(positions, positions + n3);
    if (cell)
      cells_[key].assign(cell, cell + 9);
    else
      cells_.erase(key);
  }

  [[nodiscard]] int nearest(const double *positions, long n3,
                            const double *cell = nullptr) const {
    double best = std::numeric_limits<double>::infinity();
    int group = 0;
    for (const auto &[key, stored] : last_) {
      if (static_cast<long>(stored.size()) != n3)
        continue;
      const auto recordedCell = cells_.find(key);
      if (cell && recordedCell != cells_.end() &&
          !std::equal(recordedCell->second.begin(), recordedCell->second.end(),
                      cell))
        continue;
      std::vector<double> differences(static_cast<size_t>(n3));
      for (long i = 0; i < n3; ++i) {
        differences[static_cast<size_t>(i)] =
            stored[static_cast<size_t>(i)] - positions[i];
      }
      if (cell &&
          std::any_of(cell, cell + 9, [](double x) { return x != 0.0; }))
        pbc::closestDisplacements(differences.data(), differences.size(), cell);
      double d2 = 0.0;
      for (double d : differences)
        d2 += d * d;
      const auto home = home_.find(key);
      if (d2 < best && home != home_.end()) {
        best = d2;
        group = home->second;
      }
    }
    return group;
  }

  void settle(std::int64_t key, int group) { home_[key] = group; }
 
private:
  [[nodiscard]] int homeOf(std::int64_t key, long position) const {
    const auto it = home_.find(key);
    return it != home_.end() ? it->second : defaultHome(key, position);
  }
 
  int groups_;
  std::map<std::int64_t, int> home_;
  std::map<std::int64_t, std::vector<double>> last_;
  std::map<std::int64_t, std::vector<double>> cells_;
};
 
} // namespace eonc