Correlator.cpp

Go to the documentation of this file.

/*
 * Copyright (C) 2010-2026 The ESPResSo project
 *
 * This file is part of ESPResSo.
 *
 * ESPResSo is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * ESPResSo is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include "Correlator.hpp"
 
#include <utils/Vector.hpp>
#include <utils/math/sqr.hpp>
#include <utils/serialization/multi_array.hpp>
 
#include <boost/archive/binary_iarchive.hpp>
#include <boost/archive/binary_oarchive.hpp>
#include <boost/iostreams/device/array.hpp>
#include <boost/iostreams/stream.hpp>
#include <boost/serialization/string.hpp>
#include <boost/serialization/vector.hpp>
 
#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <cstddef>
#include <functional>
#include <numeric>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>
 
namespace Accumulators {
/** Compress computing arithmetic mean: A_compressed=(A1+A2)/2 */
std::vector<double> compress_linear(std::vector<double> const &A1,
                                    std::vector<double> const &A2) {
  assert(A1.size() == A2.size());
  std::vector<double> A_compressed(A1.size());
 
  std::ranges::transform(A1, A2, A_compressed.begin(),
                         [](double a, double b) { return 0.5 * (a + b); });
 
  return A_compressed;
}
 
/** Compress discarding the 1st argument and return the 2nd */
std::vector<double>
compress_discard1([[maybe_unused]] std::vector<double> const &A1,
                  [[maybe_unused]] std::vector<double> const &A2) {
  assert(A1.size() == A2.size());
  std::vector<double> A_compressed(A2);
  return A_compressed;
}
 
/** Compress discarding the 2nd argument and return the 1st */
std::vector<double>
compress_discard2([[maybe_unused]] std::vector<double> const &A1,
                  [[maybe_unused]] std::vector<double> const &A2) {
  assert(A1.size() == A2.size());
  std::vector<double> A_compressed(A1);
  return A_compressed;
}
 
std::vector<double> scalar_product(std::vector<double> const &A,
                                   std::vector<double> const &B,
                                   Utils::Vector3d const &) {
  if (A.size() != B.size()) {
    throw std::runtime_error(
        "Error in scalar product: The vector sizes do not match");
  }
 
  auto const result = std::inner_product(A.begin(), A.end(), B.begin(), 0.0);
  return {result};
}
 
std::vector<double> componentwise_product(std::vector<double> const &A,
                                          std::vector<double> const &B,
                                          Utils::Vector3d const &) {
  std::vector<double> C(A.size());
  if (A.size() != B.size()) {
    throw std::runtime_error(
        "Error in componentwise product: The vector sizes do not match");
  }
 
  std::ranges::transform(A, B, C.begin(), std::multiplies<>());
 
  return C;
}
 
std::vector<double> tensor_product(std::vector<double> const &A,
                                   std::vector<double> const &B,
                                   Utils::Vector3d const &) {
  std::vector<double> C(A.size() * B.size());
  auto C_it = C.begin();
 
  for (double a : A) {
    for (double b : B) {
      *(C_it++) = a * b;
    }
  }
 
  return C;
}
 
std::vector<double> square_distance_componentwise(std::vector<double> const &A,
                                                  std::vector<double> const &B,
                                                  Utils::Vector3d const &) {
  if (A.size() != B.size()) {
    throw std::runtime_error(
        "Error in square distance componentwise: The vector sizes do not "
        "match.");
  }
 
  std::vector<double> C(A.size());
 
  std::ranges::transform(A, B, C.begin(), [](double a, double b) -> double {
    return Utils::sqr(a - b);
  });
 
  return C;
}
 
// note: the argument name wsquare denotes that its value is w^2 while the user
// sets w
std::vector<double> fcs_acf(std::vector<double> const &A,
                            std::vector<double> const &B,
                            Utils::Vector3d const &wsquare) {
  if (A.size() != B.size()) {
    throw std::runtime_error(
        "Error in fcs_acf: The vector sizes do not match.");
  }
 
  auto const C_size = A.size() / 3u;
  assert(3u * C_size == A.size());
 
  std::vector<double> C{};
  C.reserve(C_size);
 
  for (std::size_t i = 0u; i < C_size; i++) {
    auto acc = 0.;
    for (std::size_t j = 0u; j < 3u; j++) {
      auto const a = A[3u * i + j];
      auto const b = B[3u * i + j];
      acc -= Utils::sqr(a - b) / wsquare[j];
    }
    C.emplace_back(std::exp(acc));
  }
 
  return C;
}
 
void Correlator::initialize_operations() {
  // Class members are assigned via the initializer list
 
  if (m_tau_lin == 1) { // use the default
    m_tau_lin = static_cast<int>(std::ceil(m_tau_max / m_dt));
    m_tau_lin += m_tau_lin % 2;
  }
 
  if (m_tau_lin < 2) {
    throw std::runtime_error("tau_lin must be >= 2");
  }
 
  if (m_tau_lin % 2) {
    throw std::runtime_error("tau_lin must be divisible by 2");
  }
 
  if (m_tau_max <= m_dt) {
    throw std::runtime_error("tau_max must be >= delta_t (delta_N too large)");
  }
  // set hierarchy depth which can accommodate at least m_tau_max
  if ((m_tau_max / m_dt) < m_tau_lin) {
    m_hierarchy_depth = 1;
  } else {
    auto const operand = (m_tau_max / m_dt) / double(m_tau_lin - 1);
    assert(operand > 0.);
    m_hierarchy_depth = static_cast<int>(std::ceil(1. + std::log2(operand)));
  }
 
  assert(A_obs);
  assert(B_obs);
  dim_A = A_obs->n_values();
  dim_B = B_obs->n_values();
 
  if (dim_A == 0u) {
    throw std::runtime_error("dimension of first observable has to be >= 1");
  }
  if (dim_B == 0u) {
    throw std::runtime_error("dimension of second observable has to be >= 1");
  }
 
  // choose the correlation operation
  if (corr_operation_name == "componentwise_product") {
    m_dim_corr = dim_A;
    m_shape = A_obs->shape();
    corr_operation = &componentwise_product;
    m_correlation_args = Utils::Vector3d{0, 0, 0};
  } else if (corr_operation_name == "tensor_product") {
    m_dim_corr = dim_A * dim_B;
    m_shape.clear();
    m_shape.emplace_back(dim_A);
    m_shape.emplace_back(dim_B);
    corr_operation = &tensor_product;
    m_correlation_args = Utils::Vector3d{0, 0, 0};
  } else if (corr_operation_name == "square_distance_componentwise") {
    m_dim_corr = dim_A;
    m_shape = A_obs->shape();
    corr_operation = &square_distance_componentwise;
    m_correlation_args = Utils::Vector3d{0, 0, 0};
  } else if (corr_operation_name == "fcs_acf") {
    // note: user provides w=(wx,wy,wz) but we want to use
    // wsquare=(wx^2,wy^2,wz^2)
    if (not(m_correlation_args_input > Utils::Vector3d::broadcast(0.))) {
      throw std::runtime_error("missing parameter for fcs_acf: w_x w_y w_z");
    }
    m_correlation_args = Utils::hadamard_product(m_correlation_args_input,
                                                 m_correlation_args_input);
    if (dim_A % 3u)
      throw std::runtime_error("dimA must be divisible by 3 for fcs_acf");
    m_dim_corr = dim_A / 3u;
    m_shape = A_obs->shape();
    if (m_shape.back() != 3u)
      throw std::runtime_error(
          "the last dimension of dimA must be 3 for fcs_acf");
    m_shape.pop_back();
    corr_operation = &fcs_acf;
  } else if (corr_operation_name == "scalar_product") {
    m_dim_corr = 1u;
    m_shape.clear();
    m_shape.emplace_back(1u);
    corr_operation = &scalar_product;
    m_correlation_args = Utils::Vector3d{0, 0, 0};
  } else {
    throw std::invalid_argument("correlation operation '" +
                                corr_operation_name + "' not implemented");
  }
 
  // Choose the compression function
  if (compressA_name == "discard2") {
    compressA = &compress_discard2;
  } else if (compressA_name == "discard1") {
    compressA = &compress_discard1;
  } else if (compressA_name == "linear") {
    compressA = &compress_linear;
  } else {
    throw std::invalid_argument("unknown compression method '" +
                                compressA_name + "' for first observable");
  }
 
  if (compressB_name == "discard2") {
    compressB = &compress_discard2;
  } else if (compressB_name == "discard1") {
    compressB = &compress_discard1;
  } else if (compressB_name == "linear") {
    compressB = &compress_linear;
  } else {
    throw std::invalid_argument("unknown compression method '" +
                                compressB_name + "' for second observable");
  }
}
 
void Correlator::initialize_buffers() {
  using index_type = decltype(result)::index;
 
  A.resize(std::array<int, 2>{{m_hierarchy_depth, m_tau_lin + 1}});
  std::fill_n(A.data(), A.num_elements(), std::vector<double>(dim_A, 0));
  B.resize(std::array<int, 2>{{m_hierarchy_depth, m_tau_lin + 1}});
  std::fill_n(B.data(), B.num_elements(), std::vector<double>(dim_B, 0));
 
  n_data = 0;
  A_accumulated_average = std::vector<double>(dim_A, 0);
  B_accumulated_average = std::vector<double>(dim_B, 0);
 
  auto const n_result = n_values();
  n_sweeps = std::vector<std::size_t>(n_result, 0);
  n_vals = std::vector<long>(m_hierarchy_depth, 0);
 
  result.resize(std::array<std::size_t, 2>{{n_result, m_dim_corr}});
  for (index_type i = 0; i < static_cast<index_type>(n_result); i++) {
    for (index_type j = 0; j < static_cast<index_type>(m_dim_corr); j++) {
      result[i][j] = 0.;
    }
  }
 
  newest = std::vector<long>(m_hierarchy_depth, m_tau_lin);
 
  tau.resize(n_result);
  for (int i = 0; i < m_tau_lin + 1; i++) {
    tau[i] = i;
  }
 
  for (int j = 1; j < m_hierarchy_depth; j++) {
    for (int k = 0; k < m_tau_lin / 2; k++) {
      tau[m_tau_lin + 1 + (j - 1) * m_tau_lin / 2 + k] =
          (k + (m_tau_lin / 2) + 1) * (1 << j);
    }
  }
}
 
void Correlator::compress_kernel(long lowest_level, long highest_level) {
  auto const tau = static_cast<long>(m_tau_lin);
  for (long i = highest_level; i >= lowest_level; i--) {
    // We increase the index indicating the newest on level i+1 by one (plus
    // folding)
    newest[i + 1l] = (newest[i + 1l] + 1l) % (tau + 1l);
    n_vals[i + 1l] += 1l;
 
    A[i + 1l][newest[i + 1l]] =
        (*compressA)(A[i][(newest[i] + 1l) % (tau + 1l)],
                     A[i][(newest[i] + 2l) % (tau + 1l)]);
    B[i + 1l][newest[i + 1l]] =
        (*compressB)(B[i][(newest[i] + 1l) % (tau + 1l)],
                     B[i][(newest[i] + 2l) % (tau + 1l)]);
  }
  newest[lowest_level] = (newest[lowest_level] + 1l) % (tau + 1l);
}
 
void Correlator::correlate_kernel(long lowest_level, long highest_level) {
  using index_type = decltype(result)::index;
  auto const tau = static_cast<long>(m_tau_lin);
  auto const half_tau = (tau + 1l) / 2l + 1l;
  // We only need to update correlation estimates for the higher levels
  for (long i = lowest_level + 1l; i < highest_level + 2l; i++) {
    for (long j = half_tau; j < std::min(tau + 1l, n_vals[i]); j++) {
      auto const index_new = newest[i];
      auto const index_old = (newest[i] - j + tau + 1l) % (tau + 1l);
      auto const index_res =
          tau + (i - 1l) * tau / 2l + (j - tau / 2l + 1l) - 1l;
      auto const temp = (corr_operation)(A[i][index_old], B[i][index_new],
                                         m_correlation_args);
      assert(temp.size() == m_dim_corr);
      n_sweeps[index_res]++;
      for (index_type k = 0; k < static_cast<index_type>(m_dim_corr); k++) {
        result[index_res][k] += temp[k];
      }
    }
  }
}
 
void Correlator::update(boost::mpi::communicator const &comm) {
  if (finalized) {
    throw std::runtime_error(
        "No data can be added after finalize() was called.");
  }
 
  if (comm.rank() != 0) {
    // worker nodes just need to update the observables and exit
    A_obs->operator()(comm);
    if (A_obs != B_obs) {
      B_obs->operator()(comm);
    }
 
    return;
  }
 
  // We must now go through the hierarchy and make sure there is space for the
  // new datapoint. For every hierarchy level we have to decide if it is
  // necessary to move something
  long highest_level_to_compress = -1l;
 
  t++;
 
  // Let's find out how far we have to go back in the hierarchy to make space
  // for the new value
  {
    auto const max_depth = m_hierarchy_depth - 1;
    int i = 0;
    while (true) {
      if (i >= max_depth or n_vals[i] <= m_tau_lin) {
        break;
      }
      auto const modulo = 1 << (i + 1);
      auto const remainder = (t - (m_tau_lin + 1) * (modulo - 1) - 1) % modulo;
      if (remainder != 0) {
        break;
      }
      highest_level_to_compress++;
      i++;
    }
  }
 
  // Now we know we must make space on the levels 0..highest_level_to_compress
  // Now let's compress the data level by level.
  compress_kernel(0l, highest_level_to_compress);
 
  n_vals[0]++;
 
  A[0][newest[0]] = A_obs->operator()(comm);
  if (A_obs != B_obs) {
    B[0][newest[0]] = B_obs->operator()(comm);
  } else {
    B[0][newest[0]] = A[0][newest[0]];
  }
 
  // Now we update the cumulated averages and variances of A and B
  n_data++;
  for (std::size_t k = 0; k < dim_A; k++) {
    A_accumulated_average[k] += A[0][newest[0]][k];
  }
 
  for (std::size_t k = 0; k < dim_B; k++) {
    B_accumulated_average[k] += B[0][newest[0]][k];
  }
 
  using index_type = decltype(result)::index;
  auto const tau = static_cast<long>(m_tau_lin);
  // Now update the lowest level correlation estimates
  for (long j = 0l; j < std::min(tau + 1l, n_vals[0]); j++) {
    auto const index_new = newest[0];
    auto const index_old = (newest[0] - j + tau + 1l) % (tau + 1l);
    auto const temp =
        (corr_operation)(A[0][index_old], B[0][index_new], m_correlation_args);
    assert(temp.size() == m_dim_corr);
 
    n_sweeps[j]++;
    for (index_type k = 0; k < static_cast<index_type>(m_dim_corr); k++) {
      result[j][k] += temp[k];
    }
  }
  // Now for the higher ones
  correlate_kernel(0l, highest_level_to_compress);
}
 
int Correlator::finalize(boost::mpi::communicator const &comm) {
  if (finalized) {
    throw std::runtime_error("Correlator::finalize() can only be called once.");
  }
  // We must now go through the hierarchy and make sure there is space for the
  // new datapoint. For every hierarchy level we have to decide if it is
  // necessary to move something
 
  // mark the correlation as finalized
  finalized = true;
 
  // worker nodes don't need to do anything
  if (comm.rank() != 0) {
    return 0;
  }
 
  auto const tau = static_cast<long>(m_tau_lin);
  for (long ll = 0; ll < static_cast<long>(m_hierarchy_depth - 1); ll++) {
    long vals_ll; // number of values remaining in the lowest level
    if (n_vals[ll] > tau + 1l)
      vals_ll = tau + n_vals[ll] % 2l;
    else
      vals_ll = n_vals[ll];
 
    while (vals_ll) {
      // Check, if we will want to push the value from the lowest level
      auto highest_level_to_compress = (vals_ll % 2l) ? ll : -1l;
 
      // Let's find out how far we have to go back in the hierarchy to make
      // space for the new value
      {
        auto const max_depth = static_cast<long>(m_hierarchy_depth - 1);
        long i = ll + 1l; // lowest level for which to check for compression
        while (highest_level_to_compress > -1l) {
          if (i >= max_depth or n_vals[i] % 2l == 0l or n_vals[i] <= tau) {
            break;
          }
          highest_level_to_compress++;
          i++;
        }
      }
      vals_ll--;
 
      // Now we know we must make space on the levels
      // ll..highest_level_to_compress
      // Now let's compress the data level by level.
      compress_kernel(ll, highest_level_to_compress);
      correlate_kernel(ll, highest_level_to_compress);
    }
  }
  return 0;
}
 
std::vector<double> Correlator::get_correlation() {
  using index_type = decltype(result)::index;
  auto const n_result = n_values();
  std::vector<double> res(n_result * m_dim_corr);
 
  for (std::size_t i = 0; i < n_result; i++) {
    auto const index = static_cast<index_type>(m_dim_corr * i);
    for (index_type k = 0; k < static_cast<index_type>(m_dim_corr); k++) {
      if (n_sweeps[i]) {
        res[index + k] = result[static_cast<index_type>(i)][k] /
                         static_cast<double>(n_sweeps[i]);
      }
    }
  }
  return res;
}
 
std::vector<double> Correlator::get_lag_times() const {
  std::vector<double> res(n_values());
  std::ranges::transform(tau, res.begin(),
                         [dt = m_dt](auto const &a) { return a * dt; });
  return res;
}
 
std::string Correlator::get_internal_state() const {
  std::stringstream ss;
  boost::archive::binary_oarchive oa(ss);
 
  oa << t;
  oa << m_dt;
  oa << m_shape;
  oa << m_correlation_args_input;
  oa << A;
  oa << B;
  oa << result;
  oa << n_sweeps;
  oa << n_vals;
  oa << newest;
  oa << A_accumulated_average;
  oa << B_accumulated_average;
  oa << n_data;
 
  return ss.str();
}
 
void Correlator::set_internal_state(std::string const &state) {
  namespace iostreams = boost::iostreams;
  iostreams::array_source src(state.data(), state.size());
  iostreams::stream<iostreams::array_source> ss(src);
  boost::archive::binary_iarchive ia(ss);
 
  ia >> t;
  ia >> m_dt;
  ia >> m_shape;
  ia >> m_correlation_args_input;
  ia >> A;
  ia >> B;
  ia >> result;
  ia >> n_sweeps;
  ia >> n_vals;
  ia >> newest;
  ia >> A_accumulated_average;
  ia >> B_accumulated_average;
  ia >> n_data;
  initialize_operations();
  m_system = nullptr;
}
 
} // namespace Accumulators