# SPDX-License-Identifier: MPL-2.0
#
# ODAT-SE -- an open framework for data analysis
# Copyright (C) 2020- The University of Tokyo
#
# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
# If a copy of the MPL was not distributed with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
from typing import Optional, TYPE_CHECKING
from io import open
import copy
import time
import itertools
import sys
import numpy as np
import odatse
import odatse.algorithm.montecarlo
from odatse.util.read_ts import read_Ts
from odatse.util.separateT import separateT, calculate_statistics_from_separated_files
from odatse.util.data_writer import DataWriter
if TYPE_CHECKING:
from mpi4py import MPI
[docs]
class Algorithm(odatse.algorithm.montecarlo.AlgorithmBase):
"""
Replica Exchange Monte Carlo (REMC) Algorithm Implementation.
This class implements the Replica Exchange Monte Carlo algorithm, also known as
Parallel Tempering. The algorithm runs multiple replicas of the system at different
temperatures and periodically attempts to swap configurations between adjacent
temperature levels.
Attributes
----------
numsteps : int
Total number of Monte Carlo steps to perform.
numsteps_exchange : int
Number of steps between exchange attempts.
numsteps_thermalization : int
Number of steps to discard for thermalization.
fx : np.ndarray
Current energy/objective function values.
istep : int
Current step number.
nreplica : int
Total number of replicas across all processes.
Tindex : np.ndarray
Temperature indices for current replicas.
rep2T : np.ndarray
Mapping from replica index to temperature index.
T2rep : np.ndarray
Mapping from temperature index to replica index.
exchange_direction : bool
Direction for attempting exchanges (alternates between True/False).
"""
# Coordinate bounds/steps live on self.statespace; the walker state is in
# self.state (see montecarlo.AlgorithmBase / state.py).
numsteps: int
numsteps_exchange: int
fx: np.ndarray
istep: int
nreplica: int
Tindex: np.ndarray
rep2T: np.ndarray
T2rep: np.ndarray
exchange_direction: bool
# Exchange-specific fields appended to the MC base checkpoint.
_checkpoint_attrs: list[str] = ["nreplica", "Tindex", "rep2T", "T2rep", "exchange_direction"]
[docs]
def __init__(
self,
info: odatse.Info,
runner: odatse.Runner = None,
run_mode: str = "initial",
) -> None:
"""
Initialize the Algorithm class.
Parameters
----------
info : odatse.Info
Information object containing algorithm parameters.
runner : odatse.Runner, optional
Runner object for executing the algorithm.
run_mode : str, optional
Mode to run the algorithm in, by default "initial".
"""
info_exchange = info.algorithm["exchange"]
nwalkers = info_exchange.get("nreplica_per_proc", 1)
super().__init__(
info=info,
runner=runner,
nwalkers=nwalkers,
run_mode=run_mode,
)
self.nreplica = odatse.mpi.algsize() * self.nwalkers
self.input_as_beta, self.betas = read_Ts(info_exchange, numT=self.nreplica)
self.numsteps = info_exchange["numsteps"]
self.numsteps_exchange = info_exchange["numsteps_exchange"]
self.numsteps_thermalization = info_exchange.get("numsteps_thermalization", int(0.1*self.numsteps))
self.export_combined_files = info_exchange.get("export_combined_files", False)
self.separate_T = info_exchange.get("separate_T", True)
[docs]
def _initialize(self) -> None:
"""
Initialize the algorithm parameters and state.
"""
# Initialize base class first
super()._initialize()
# Set up temperature indices for each walker
# Each process handles a contiguous block of temperature indices
# based on its rank and number of walkers
self.Tindex = np.arange(
(odatse.mpi.algrank() or 0) * self.nwalkers, ((odatse.mpi.algrank() or 0) + 1) * self.nwalkers
)
# Initialize mappings between replica and temperature indices
# Initially, replica i has temperature i
self.rep2T = np.arange(self.nreplica) # Maps replica index -> temperature index
self.T2rep = np.arange(self.nreplica) # Maps temperature index -> replica index
# Initialize exchange direction - alternates between True/False
# to ensure all adjacent pairs get chance to exchange
self.exchange_direction = True
self.istep = 0
self._show_parameters()
[docs]
def _run(self) -> None:
"""
Run the algorithm.
"""
# dispatch は prepare() が処理済み
# Get current beta (inverse temperature) values for each replica
beta = self.betas[self.Tindex]
# Set up output file writers
write_mode = "w" if self.mode.startswith("init") else "a"
item_list = [
"step",
"walker",
("beta" if self.input_as_beta else "T"),
"fx",
*self.label_list,
]
# Create writers for both trial moves and accepted results
fp_trial = DataWriter("trial.txt", mode=write_mode, item_list=item_list, combined=self.export_combined_files)
fp_result = DataWriter("result.txt", mode=write_mode, item_list=item_list, combined=self.export_combined_files)
self._set_writer(fp_trial, fp_result)
# For new runs, evaluate initial configuration
if self.mode.startswith("init"):
self.fx = self._evaluate(self.state)
self._write_result(fp_trial)
self._write_result(fp_result)
self.istep += 1
# Track best solution found
minidx = np.argmin(self.fx)
self.best_x = copy.copy(self.state.x[minidx, :])
self.best_fx = np.min(self.fx[minidx])
self.best_istep = 0
self.best_iwalker = 0
# Set up checkpointing intervals
next_checkpoint_step = self.istep + self.checkpoint_steps
next_checkpoint_time = time.time() + self.checkpoint_interval
# Main simulation loop
while self.istep < self.numsteps:
# Attempt replica exchange periodically
if self.istep % self.numsteps_exchange == 0:
time_sta = time.perf_counter()
if self.nreplica > 1:
self._exchange(self.exchange_direction)
# Alternate exchange direction for next attempt
if self.nreplica > 2:
self.exchange_direction = not self.exchange_direction
time_end = time.perf_counter()
self.timer["run"]["exchange"] += time_end - time_sta
# Update beta values after exchange
beta = self.betas[self.Tindex]
# Perform local Monte Carlo updates
self.local_update(beta)
self.istep += 1
# Handle checkpointing if enabled
if self.checkpoint:
time_now = time.time()
if self.istep >= next_checkpoint_step or time_now >= next_checkpoint_time:
self._save_state(self.checkpoint_file)
next_checkpoint_step = self.istep + self.checkpoint_steps
next_checkpoint_time = time_now + self.checkpoint_interval
# Clean up file handles
fp_trial.close()
fp_result.close()
print("complete main process : rank {:08d}/{:08d}".format(odatse.mpi.algrank(), odatse.mpi.algsize()))
# Save final state for possible continuation
if self.checkpoint:
self._save_state(self.checkpoint_file)
[docs]
def _exchange(self, direction: bool) -> None:
"""
Attempt temperature exchanges between replicas.
This method implements the core replica exchange logic, attempting to swap
temperatures between adjacent replicas based on the Metropolis criterion:
P(accept) = min(1, exp((β_j - β_i)(E_i - E_j)))
Parameters
----------
direction : bool
If True, attempt exchanges between even-odd pairs.
If False, attempt exchanges between odd-even pairs.
"""
if self.nwalkers == 1:
self.__exchange_single_walker(direction)
else:
self.__exchange_multi_walker(direction)
def __exchange_single_walker(self, direction: bool) -> None:
"""
Handle temperature exchanges for single walker per process case.
This method implements the exchange logic when each process has only one walker,
requiring MPI communication to coordinate exchanges between processes.
Parameters
----------
direction : bool
If True, attempt exchanges between even-odd pairs.
If False, attempt exchanges between odd-even pairs.
"""
# This path identifies each replica with exactly one MPI rank: T2rep
# entries are used directly as Send/Recv ranks, and MPI ranks are
# iterated as replica indices below (range(1, nreplica)). That identity
# holds only when there is one replica per process, i.e.
# nreplica == algsize (guaranteed here because nwalkers == 1). Guard it
# so a future change cannot silently corrupt T2rep or deadlock on a bad
# rank.
assert self.nreplica == odatse.mpi.algsize(), (
"single-walker exchange requires one replica per process "
f"(nreplica == algsize), got nreplica={self.nreplica}, "
f"algsize={odatse.mpi.algsize()}"
)
comm = odatse.mpi.algcomm()
comm.Barrier()
if direction:
if self.Tindex[0] % 2 == 0:
other_index = self.Tindex[0] + 1
is_main = True
else:
other_index = self.Tindex[0] - 1
is_main = False
else:
if self.Tindex[0] % 2 == 0:
other_index = self.Tindex[0] - 1
is_main = False
else:
other_index = self.Tindex[0] + 1
is_main = True
ibuf = np.zeros(1, dtype=np.int64)
fbuf = np.zeros(1, dtype=np.float64)
if 0 <= other_index < self.nreplica:
other_rank = self.T2rep[other_index]
if is_main:
comm.Recv(fbuf, source=other_rank, tag=1)
other_fx = fbuf[0]
beta = self.betas[self.Tindex[0]]
other_beta = self.betas[self.Tindex[0] + 1]
logp = (other_beta - beta) * (other_fx - self.fx[0])
if logp >= 0.0 or self.rng.rand() < np.exp(logp):
ibuf[0] = self.Tindex[0]
comm.Send(ibuf, dest=other_rank, tag=2)
self.Tindex[0] += 1
else:
ibuf[0] = self.Tindex[0] + 1
comm.Send(ibuf, dest=other_rank, tag=2)
else:
fbuf[0] = self.fx[0]
comm.Send(fbuf, dest=other_rank, tag=1)
comm.Recv(ibuf, source=other_rank, tag=2)
self.Tindex[0] = ibuf[0]
comm.Barrier()
if odatse.mpi.algrank() == 0:
self.T2rep[self.Tindex[0]] = odatse.mpi.algrank()
for other_rank in range(1, self.nreplica):
comm.Recv(ibuf, source=other_rank, tag=0)
self.T2rep[ibuf[0]] = other_rank
else:
ibuf[0] = self.Tindex[0]
comm.Send(ibuf, dest=0, tag=0)
comm.Bcast(self.T2rep, root=0)
def __exchange_multi_walker(self, direction: bool) -> None:
"""
Handle temperature exchanges for multiple walkers per process case.
This method implements the exchange logic when each process has multiple walkers,
requiring collective MPI operations to coordinate exchanges across all processes.
Parameters
----------
direction : bool
If True, attempt exchanges between even-odd pairs.
If False, attempt exchanges between odd-even pairs.
"""
comm = odatse.mpi.algcomm()
if odatse.mpi.algsize() > 1:
fx_all = comm.allgather(self.fx)
fx_all = np.array(fx_all).flatten()
else:
fx_all = self.fx
rep2T_diff = []
T2rep_diff = []
rank = odatse.mpi.algrank()
for irep in range(rank * self.nwalkers, (rank+1) * self.nwalkers):
iT = self.rep2T[irep]
if iT % 2 != 0:
continue
jT = iT + 1 if direction else iT - 1
if jT < 0 or jT == self.nreplica:
continue
jrep = self.T2rep[jT]
fdiff = fx_all[jrep] - fx_all[irep]
bdiff = self.betas[jT] - self.betas[iT]
logp = fdiff * bdiff
if logp >= 0.0 or self.rng.rand() < np.exp(logp):
rep2T_diff.append((irep, jT)) # this means self.rep2T[irep] = jT
rep2T_diff.append((jrep, iT))
T2rep_diff.append((iT, jrep))
T2rep_diff.append((jT, irep))
if odatse.mpi.algsize() > 1:
rep2T_diff = comm.allgather(rep2T_diff)
rep2T_diff = list(itertools.chain.from_iterable(rep2T_diff)) # flatten
T2rep_diff = comm.allgather(T2rep_diff)
T2rep_diff = list(itertools.chain.from_iterable(T2rep_diff)) # flatten
for diff in rep2T_diff:
self.rep2T[diff[0]] = diff[1]
for diff in T2rep_diff:
self.T2rep[diff[0]] = diff[1]
self.Tindex = self.rep2T[rank * self.nwalkers : (rank + 1) * self.nwalkers]
[docs]
def _prepare(self) -> None:
"""
Prepare the algorithm for execution.
"""
self.timer["run"]["submit"] = 0.0
self.timer["run"]["exchange"] = 0.0
[docs]
def _post(self) -> dict:
"""
Post-process the results of the algorithm.
"""
# Separate results by temperature if requested
if self.separate_T and not self.export_combined_files:
if odatse.mpi.algrank() == 0:
print(f"start separateT {odatse.mpi.algrank()}")
sys.stdout.flush()
# Convert beta to temperature if needed
Ts = self.betas if self.input_as_beta else 1.0 / self.betas
# Organize results by temperature
separateT(
Ts=Ts,
nwalkers=self.nwalkers,
output_dir=self.output_dir,
comm=odatse.mpi.algcomm(),
use_beta=self.input_as_beta,
buffer_size=10000,
)
calculate_statistics_from_separated_files(
Ts=Ts,
output_dir=self.output_dir,
thermalization_steps=self.numsteps_thermalization,
comm=odatse.mpi.algcomm(),
)
# Gather best results from all processes
if odatse.mpi.algsize() > 1:
# NOTE:
# ``gather`` seems not to work with many processes (say, 32) in some MPI implementation.
# ``Gather`` and ``allgather`` seem to work fine.
# Since the performance is not so important here, we use ``allgather`` for simplicity.
comm = odatse.mpi.algcomm()
best_fx = comm.allgather(self.best_fx)
best_x = comm.allgather(self.best_x)
best_istep = comm.allgather(self.best_istep)
best_iwalker = comm.allgather(self.best_iwalker)
else:
best_fx = [self.best_fx]
best_x = [self.best_x]
best_istep = [self.best_istep]
best_iwalker = [self.best_iwalker]
# Find process with best overall solution
best_rank = np.argmin(best_fx)
# Write best result to file (rank 0 only)
if odatse.mpi.algrank() == 0:
with open("best_result.txt", "w") as f:
f.write(f"nprocs = {self.nreplica}\n")
f.write(f"rank = {best_rank}\n")
f.write(f"step = {best_istep[best_rank]}\n")
f.write(f"walker = {best_iwalker[best_rank]}\n")
f.write(f"fx = {best_fx[best_rank]}\n")
for label, x in zip(self.label_list, best_x[best_rank]):
f.write(f"{label} = {x}\n")
# Print summary to stdout
print("Best Result:")
print(f" rank = {best_rank}")
print(f" step = {best_istep[best_rank]}")
print(f" walker = {best_iwalker[best_rank]}")
print(f" fx = {best_fx[best_rank]}")
for label, x in zip(self.label_list, best_x[best_rank]):
print(f" {label} = {x}")
# Return best solution information
return {
"x": best_x[best_rank],
"fx": best_fx[best_rank],
"nprocs": self.nreplica,
"rank": best_rank,
"step": best_istep[best_rank],
"walker": best_iwalker[best_rank],
}
[docs]
def _apply_state(self, data: dict, mode: str = "resume", restore_rng: bool = True) -> None:
"""Restore algorithm state from a checkpoint snapshot.
Delegates MPI validation, RNG restore, and MC-layer fields to the
base class, validates the replica count, then applies exchange-specific
fields and propagates the restored RNG to the state space.
REMC does not distinguish between resume and continue modes; ``mode``
is accepted for API consistency with PAMC and forwarded to super().
Parameters
----------
data : dict
Snapshot previously produced by ``__getstate__``.
mode : str
``"resume"`` or ``"continue"``; forwarded to the base class.
restore_rng : bool
When *True* (default) the RNG state is restored from *data*;
when *False* a fresh RNG state is kept (``--reset_rand`` mode).
"""
super()._apply_state(data, mode=mode, restore_rng=restore_rng)
assert self.nreplica == data["nreplica"]
for attr in Algorithm._checkpoint_attrs:
setattr(self, attr, data[attr])
self.statespace.rng = self.rng