ReplicaExchangeJob.cpp

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
*/
#include "eon/ReplicaExchangeJob.h"
#include "eon/BaseStructures.h"
#include "eon/Dynamics.h"
#include "eon/HelperFunctions.h"
#include "eon/Matter.h"
#include "eon/PotCapabilities.h"
#include "eon/RandomNumbers.h"
#include "eon/SafeMath.h"
 
#include <algorithm>
#include <cmath>
#include <format>
#include <fstream>
#include <stdexcept>
#include <string>
#include <thread>
 
namespace eonc {
 
std::vector<std::string> ReplicaExchangeJob::run() {
  std::string posFilename =
      eonc::helpers::getRelevantFile(params.main_options().conFilename);
  pos = std::make_shared<Matter>(pot, params);
  if (!eonc::io::io_ok(pos->con2matter(posFilename))) {
    QUILL_LOG_CRITICAL(log, "Failed to load {}", posFilename);
    throw std::runtime_error("failed to load " + posFilename);
  }
  (void)runFromMatter(pos);
  return returnFiles;
}
 
std::shared_ptr<Matter>
ReplicaExchangeJob::runFromMatter(std::shared_ptr<Matter> initial) {
  if (!initial) {
    throw std::runtime_error("ReplicaExchangeJob::runFromMatter: null Matter");
  }
  pos = initial;
  pos->setPotential(pot);
 
  auto &rex = ParametersLoadAccess::replica_exchange_options(params);
  if (rex.replicas < 1) {
    throw std::invalid_argument(
        "ReplicaExchangeJob: replica_exchange.replicas must be >= 1");
  }
  if (rex.temperature_low <= 0.0) {
    rex.temperature_low = params.main_options().temperature > 0.0
                              ? params.main_options().temperature
                              : 300.0;
  }
  if (rex.temperature_high <= rex.temperature_low) {
    rex.temperature_high = rex.temperature_low * 1.5;
  }
 
  long samplingSteps =
      static_cast<long>(params.replica_exchange_options().sampling_time /
                            params.dynamics_options().time_step +
                        0.5);
  long exchangePeriodSteps =
      static_cast<long>(params.replica_exchange_options().exchange_period /
                            params.dynamics_options().time_step +
                        0.5);
  const double kB = params.constants().kB;
  if (samplingSteps <= 0)
    samplingSteps = 1;
  if (exchangePeriodSteps <= 0)
    exchangePeriodSteps = 1;
 
  QUILL_LOG_DEBUG(log, "Running Replica Exchange");
 
  long refForceCalls = PotRegistry::get().total_force_calls();
 
  const long nReplicas = params.replica_exchange_options().replicas;
  std::vector<std::shared_ptr<Matter>> replica(nReplicas);
  std::vector<std::unique_ptr<Dynamics>> replicaDynamics(nReplicas);
 
  // Per-image potentials: each replica gets its own potential instance when
  // the potential requires it (e.g. ML potentials with internal state)
  const bool perImage = pot->needsPerImageInstance();
  for (long i = 0; i < nReplicas; i++) {
    auto replicaPot = perImage ? eonc::helpers::makePotential(params) : pot;
    replica[i] = std::make_shared<Matter>(replicaPot, params);
    *replica[i] = *pos;
    replica[i]->setPotential(replicaPot);
    replicaDynamics[i] = std::make_unique<Dynamics>(replica[i].get(), params);
  }
 
  std::vector<double> replicaTemperature(nReplicas);
 
  QUILL_LOG_DEBUG(log, "Temperature distribution:");
  if (nReplicas < 2) {
    replicaTemperature[0] = params.replica_exchange_options().temperature_low;
    replicaDynamics[0]->setTemperature(replicaTemperature[0]);
    replicaDynamics[0]->setThermalVelocity();
  } else if (params.replica_exchange_options().temperature_distribution ==
             "linear") {
    for (long i = 0; i < nReplicas; i++) {
      replicaTemperature[i] =
          params.replica_exchange_options().temperature_low +
          static_cast<double>(i) / static_cast<double>(nReplicas - 1) *
              (params.replica_exchange_options().temperature_high -
               params.replica_exchange_options().temperature_low);
      replicaDynamics[i]->setTemperature(replicaTemperature[i]);
      replicaDynamics[i]->setThermalVelocity();
    }
  } else if (params.replica_exchange_options().temperature_distribution ==
             "exponential") {
    double kTemp = std::log(params.replica_exchange_options().temperature_high /
                            params.replica_exchange_options().temperature_low) /
                   static_cast<double>(nReplicas - 1);
    for (long i = 0; i < nReplicas; i++) {
      replicaTemperature[i] =
          params.replica_exchange_options().temperature_low *
          std::exp(kTemp * static_cast<double>(i));
      replicaDynamics[i]->setTemperature(replicaTemperature[i]);
      replicaDynamics[i]->setThermalVelocity();
      QUILL_LOG_DEBUG(log, "replica: {} temperature {:.0f}", i + 1,
                      replicaTemperature[i]);
    }
  } else {
    throw std::invalid_argument(
        "replica_exchange.temperature_distribution must be linear or "
        "exponential");
  }
 
  QUILL_LOG_DEBUG(
      log, "Replica Exchange sampling for {:.0f} fs; {} steps; {} replicas.",
      params.replica_exchange_options().sampling_time * 10.18, samplingSteps,
      params.replica_exchange_options().replicas);
 
  // Parallel replica dynamics when enabled and potential supports it
  const bool canParallel = params.main_options().parallel &&
                           (eonc::potAllowsSharedInstance(*pot) || perImage);
 
  for (long step = 1; step <= samplingSteps; step++) {
    if (canParallel && nReplicas > 1) {
      // Parallel: each replica runs its MD step in a separate thread
      std::vector<std::thread> threads;
      threads.reserve(nReplicas);
      for (long i = 0; i < nReplicas; i++) {
        threads.emplace_back([&, i] {
          // New std::thread, new TLS ran2. Split the stream by replica and
          // step so Langevin/Andersen noise is not identical across replicas.
          const long userSeed = params.main_options().randomSeed;
          const long base = (userSeed > 0) ? userSeed : 1;
          eonc::rng::random(base + (step + 1) * 10007 + (i + 1));
          replicaDynamics[i]->oneStep();
        });
      }
      for (auto &t : threads) {
        t.join();
      }
    } else {
      for (long i = 0; i < nReplicas; i++) {
        replicaDynamics[i]->oneStep();
      }
    }
 
    // Metropolis replica exchange
    if (nReplicas >= 2 && (step % exchangePeriodSteps) == 0) {
      for (long trial = 0;
           trial < params.replica_exchange_options().exchange_trials; trial++) {
        long i = eonc::rng::randomInt(0, nReplicas - 2);
        double energyLow = replica[i]->getPotentialEnergy();
        double energyHigh = replica[i + 1]->getPotentialEnergy();
        double kbTLow = kB * replicaTemperature[i];
        double kbTHigh = kB * replicaTemperature[i + 1];
        double pAcc =
            std::min(1.0, std::exp((energyHigh - energyLow) *
                                   (eonc::safemath::safe_recip(kbTHigh, 0.0) -
                                    eonc::safemath::safe_recip(kbTLow, 0.0))));
        double rnd = eonc::rng::randomDouble();
        QUILL_LOG_INFO(log,
                       "step: {} trial swap, i {}, elow: {:.5f}, ehigh: "
                       "{:.5f}, pAcc: {:.5f}, rand: {}",
                       step, i, energyLow, energyHigh, pAcc, rnd);
        if (rnd < pAcc) {
          QUILL_LOG_INFO(log, "swap");
          // Swap configurations, not pointers. Dynamics stays on its T slot.
          auto potLow = replica[i]->getPotential();
          auto potHigh = replica[i + 1]->getPotential();
          Matter tmp = *replica[i];
          *replica[i] = *replica[i + 1];
          *replica[i + 1] = tmp;
          replica[i]->setPotential(potLow);
          replica[i + 1]->setPotential(potHigh);
          replicaDynamics[i]->setThermalVelocity();
          replicaDynamics[i + 1]->setThermalVelocity();
        } else {
          QUILL_LOG_INFO(log, "no swap");
        }
      }
    }
  }
 
  forceCalls = PotRegistry::get().total_force_calls() - refForceCalls;
  if (!replica.empty()) {
    *pos = *replica[0];
  }
  saveData();
  return pos;
}
 
void ReplicaExchangeJob::saveData() {
  std::string resultsFilename("results.dat");
  returnFiles.push_back(resultsFilename);
 
  std::ofstream out(resultsFilename, std::ios::binary);
  if (out) {
    out << std::format("{} termination_reason\n", 0);
    out << "GOOD termination_reason_text\n";
    out << "replica_exchange job_type\n";
    out << std::format("{} random_seed\n", params.main_options().randomSeed);
    out << std::format(
        "{} potential_type\n",
        magic_enum::enum_name<PotType>(params.potential_options().potential));
    out << std::format("{} force_calls_sampling\n", forceCalls);
  }
 
  std::string posFilename("pos_out.con");
  returnFiles.push_back(posFilename);
  if (pos) {
    if (!eonc::io::io_ok(pos->matter2con(posFilename))) {
      QUILL_LOG_ERROR(log, "Failed to write {}", posFilename);
    }
  }
}
 
} // namespace eonc