Minor quality-of-life and performance enhancements

README.md

@@ -11,6 +11,10 @@ cuda version: 12.1
 
 Note that in most cases, different version of packages should also work. We list exactly the versions in our calculations in case version inconsistency issue occurs. If users intend to run the program on a cpu device, the cuda package is not needed.
 
+Optional cuEquivariance acceleration requires PyTorch 2.5.1 (CUDA 12.4),
+`mace-torch==0.3.16`, `cuequivariance==0.6.1`,
+`cuequivariance-torch==0.6.1`, and `e3nn==0.4.4`.
+
 ## Installation guide
 ### Create conda envrionment
 

configs/deploy_drl_lss.toml

@@ -14,6 +14,7 @@ dqn = true
 [deploy]
 horizon = 10 
 mode = "lss" 
+output_name = "XDATCAR{episode}"
 n_episodes = 2 
 n_poscars = 100 
 poscar_dir = "/path/to/poscars" # Change this to the directory where your POSCAR files are stored.

configs/deploy_drl_tks.toml

@@ -14,6 +14,7 @@ dqn = false
 [deploy]
 horizon = 10 
 mode = "tks" 
+output_name = "XDATCAR{episode}"
 n_episodes = 2 
 n_poscars = 100 
 poscar_dir = "/path/to/poscars" # Change this to the directory where your POSCAR files are stored.

configs/deploy_mcmc_example.toml

@@ -0,0 +1,37 @@
+# Example MCMC deployment. Copy this file and edit the paths for your system.
+# Settings not shown here retain their historical defaults.
+task = "test_run_mcmc"
+
+[logger]
+filename = "deploy"
+name = "Deploy"
+
+[deploy]
+horizon = 1000
+mode = "mcmc"
+# Omit output_name to retain the legacy XDATCAR{episode} filenames.
+output_name = "XDATCAR-mcmc-{episode}"
+n_episodes = 1
+n_poscars = 1
+poscar_dir = "/path/to/poscars"
+
+[deploy.simulation_params]
+temperature = 300
+# global_swap is the vacancy-free mode. Legacy values such as "all" and
+# "vacancy_only" remain supported.
+action_mode = "global_swap"
+n_sweeps = 1
+# Larger values reduce output I/O by buffering more frames before flushing.
+trajectory_buffer_size = 100
+
+[deploy.calc_info]
+cutoff = 4.0
+device = "cuda"
+platform = "mace"
+# Zero disables post-swap relaxation; set a positive value for the legacy
+# relaxation workflow.
+max_iter = 10
+relax_accuracy = 0.01
+relax_log = "relax.log"
+# Optional and backwards-compatible; requires the cuEquivariance environment.
+enable_cueq = true

configs/gen_time_dataset.toml

@@ -14,6 +14,7 @@ dqn = false
 [deploy]
 horizon = 5
 mode = "tks"
+output_name = "XDATCAR{episode}"
 n_episodes = 2
 [deploy.simulation_params]
 temperature = 900

demo/config.toml

@@ -14,6 +14,7 @@ dqn = true
 [deploy]
 horizon = 20 
 mode = "lss" 
+output_name = "XDATCAR0"
 n_episodes = 1
 n_poscars = 1 
 poscar_dir = "poscars" # Change this to the directory where your POSCAR files are stored.

rlsim/actions/action.py

@@ -73,6 +73,18 @@ def get_action_space_mcmc(config, lattice_parameter=3.615, action_mode="vacancy_
     Returns:
         action_space: List of [site, vector] for vacancy jumps, or [site, site] for swaps.
     """
+    if action_mode == "global_swap":
+        numbers = config.atoms.get_atomic_numbers()
+        species = np.unique(numbers)
+        if len(species) < 2:
+            raise ValueError(
+                "action_mode='global_swap' requires at least two chemical species"
+            )
+        species_pair = np.random.choice(species, size=2, replace=False)
+        first = np.random.choice(np.flatnonzero(numbers == species_pair[0]))
+        second = np.random.choice(np.flatnonzero(numbers == species_pair[1]))
+        return [[int(first), int(second)]]
+
     cell = np.array(config.atoms.get_cell())
     pbc = config.atoms.get_pbc()
 

rlsim/cli/main.py

@@ -8,6 +8,7 @@
 from rlsim.drl.train import train_DQN
 from rlsim.time.train import TimeTrainer
 from rlsim.utils.logger import setup_logger
+from rlsim.utils.output import artifact_filename
 
 
 @click.command()
@@ -45,9 +46,11 @@ def main(simulation, config_name):
         config = toml.load(f)
         task = config.pop("task")
         logger_config = config.pop("logger")
+    output_name = config.get("deploy", {}).get("output_name")
     if task not in os.listdir():
         os.makedirs(task, exist_ok=True)
-    log_filename = f"{task}/{logger_config['filename']}.log"
+    log_name = artifact_filename(output_name, f"{logger_config['filename']}.log")
+    log_filename = os.path.join(task, log_name)
     logger = setup_logger(logger_config["name"], log_filename)
     if simulation == "rl-train":
         train_DQN(task, logger, config)
@@ -67,4 +70,3 @@ def main(simulation, config_name):
     else:
         raise click.UsageError(f"Unsupported simulation type: {simulation}. Please use 'rl-train', 'rl-deploy' or 'time-train'.")
 
-

rlsim/drl/deploy.py

@@ -7,12 +7,16 @@
 
 from rlsim.drl.simulator import RLSimulator
 from rlsim.environment import Environment
+from rlsim.utils.output import artifact_filename, trajectory_filename
 
 
 def deploy_RL(task, logger, config, atoms_traj=None):
     logger.info(f"Deploy DRL in: {os.path.realpath(task)}")
-    toml.dump(config, open(f"{task}/config_copied.toml", "w"))
     deploy_config = config["deploy"]
+    output_name = deploy_config.get("output_name")
+    config_copy = os.path.join(task, artifact_filename(output_name, "config_copied.toml"))
+    with open(config_copy, "w") as file:
+        toml.dump(config, file)
     model_config = config.get("model", None)
     if model_config is not None:
         model = registry.get_model_class(model_config["@name"]).load(f"{model_config['model_path']}")
@@ -21,8 +25,10 @@ def deploy_RL(task, logger, config, atoms_traj=None):
         model = None
         model_params = None
 
+    output_name = deploy_config.pop("output_name", None)
     calc_params = deploy_config.pop("calc_info")
-    calc_params.update({"relax_log": f"{task}/{calc_params['relax_log']}"})
+    relax_log = artifact_filename(output_name, calc_params["relax_log"])
+    calc_params.update({"relax_log": os.path.join(task, relax_log)})
     horizon = deploy_config.pop("horizon")
     simulation_mode = deploy_config.pop("mode")
     simulation_params = deploy_config.pop("simulation_params")
@@ -56,22 +62,22 @@ def deploy_RL(task, logger, config, atoms_traj=None):
 
     # if simulation_mode == "lss" or simulation_mode == "mcmc":
     El = []
-    output_file = str(task) + "/converge.json"
+    output_file = os.path.join(task, artifact_filename(output_name, "converge.json"))
     if simulation_mode != "mcmc" and simulation_mode != "mmc":
         Ql = []
-        output_file_q = str(task) + "/q_values.json"
+        output_file_q = os.path.join(task, artifact_filename(output_name, "q_values.json"))
+        soutput_file_sro_chosen = os.path.join(task, artifact_filename(output_name, "SRO.json"))
         if sro_pixel is not None:
             SRO_values_list = []
             SROlist = []
-            output_file_sro = str(task) + "/sro_values.json"
-            soutput_file_sro_chosen = str(task) + "/SRO.json"
+            output_file_sro = os.path.join(task, artifact_filename(output_name, "sro_values.json"))
     if simulation_mode == "mcmc" or simulation_mode == "mmc":
         SROlist = []
-        output_file_sro_accepted = str(task) + "/SRO.json"
+        output_file_sro_accepted = os.path.join(task, artifact_filename(output_name, "SRO.json"))
     if simulation_mode == "tks":
         Tl = []
         Cl = []
-        output_file_tks = str(task) + "/diffuse.json"
+        output_file_tks = os.path.join(task, artifact_filename(output_name, "diffuse.json"))
     for u in range(n_episodes):
         if deploy_config.get("all_episodes", False):
             logger.info(f"Episode: {u} (Serial)")
@@ -86,10 +92,13 @@ def deploy_RL(task, logger, config, atoms_traj=None):
                                 q_params=model_params,
                                 sro_pixel=sro_pixel,
                                 **simulation_params)
-        atoms_traj = str(task) + "/XDATCAR" + str(u)
+        trajectory_name = trajectory_filename(output_name, u, n_episodes)
+        atoms_traj = os.path.join(task, trajectory_name)
+        final_atoms = os.path.join(task, artifact_filename(trajectory_name, "last_atoms"))
         outputs = simulator.run(horizon=horizon,
                                 logger=logger,
                                 atoms_traj=atoms_traj,
+                                final_atoms=final_atoms,
                                 mode=simulation_mode,
                                 **simulation_params)
         El.append(outputs[0])

rlsim/drl/simulator.py

@@ -8,6 +8,7 @@
 import torch
 import torch.nn as nn
 from ase import Atoms, io
+from ase.io.vasp import _write_xdatcar_config
 from ase.neb import NEB
 from ase.optimize import BFGS, FIRE, MDMin
 from numpy.linalg import norm
@@ -23,6 +24,15 @@
 ENERGY_DIFF_LIMIT = 1.5  # in eV
 
 
+def _flush_xdatcar(atoms_traj, frames):
+    if not frames:
+        return
+    with open(atoms_traj, "a") as trajectory:
+        for index, atoms in frames:
+            _write_xdatcar_config(trajectory, atoms, index)
+    frames.clear()
+
+
 class RLSimulator:
     def __init__(self,
                  environment: Environment,
@@ -33,7 +43,7 @@ def __init__(self,
                  **kwargs
                  ):
         self.env = environment
-        self.calculator = self.env.get_calculator(**self.env.calc_params)
+        self.calculator = self.env.calculator
         self.q_params = q_params
         self.model = model
         self.device = self.env.calc_params["device"]
@@ -171,20 +181,21 @@ def run(self,
             logger,
             atoms_traj: str,
             mode: str = 'lss',
+            final_atoms: str | None = None,
             **simulation_params):
         io.write(atoms_traj, self.env.atoms, format="vasp-xdatcar")
 
         if mode == "lss":
-            outputs = self.run_LSS(horizon, atoms_traj, logger, **simulation_params)
+            outputs = self.run_LSS(horizon, atoms_traj, logger, final_atoms=final_atoms, **simulation_params)
         elif mode == "mcmc" or mode == "mmc":
-            outputs = self.run_MCMC(horizon, atoms_traj, logger, **simulation_params)
+            outputs = self.run_MCMC(horizon, atoms_traj, logger, final_atoms=final_atoms, **simulation_params)
         elif mode == "tks":
             assert not self.q_params["dqn"], "TKS is only available for dqn==False."
-            outputs = self.run_TKS(horizon, atoms_traj, logger, **simulation_params)
+            outputs = self.run_TKS(horizon, atoms_traj, logger, final_atoms=final_atoms, **simulation_params)
         logger.info("Simulation finished.")
         return outputs
 
-    def run_LSS(self, horizon, atoms_traj, logger, **simulation_params):
+    def run_LSS(self, horizon, atoms_traj, logger, final_atoms=None, **simulation_params):
         if simulation_params.get("annealing_time", None) is not None:
             T_scheduler = ThermalAnnealing(total_horizon=horizon,
                                            annealing_time=simulation_params["annealing_time"],
@@ -231,11 +242,11 @@ def run_LSS(self, horizon, atoms_traj, logger, **simulation_params):
                 logger.info(
                     f"Step: {tstep}, T: {new_T:.0f}, E: {Elist[-1]:.3f}"
                 )
-        last_atoms_filename = atoms_traj.replace("XDATCAR", "last_atoms")
+        last_atoms_filename = final_atoms or atoms_traj.replace("XDATCAR", "last_atoms")
         io.write(last_atoms_filename, self.env.atoms, format="vasp")
         return (Elist, Qlist, SROlist, action_idx_list)
 
-    def run_TKS(self, horizon, atoms_traj, logger, **simulation_params):
+    def run_TKS(self, horizon, atoms_traj, logger, final_atoms=None, **simulation_params):
         tlist = [0]
         clist = [self.env.atoms.get_positions()[-1].tolist()]
         temperature = simulation_params["temperature"]
@@ -282,21 +293,26 @@ def run_TKS(self, horizon, atoms_traj, logger, **simulation_params):
                 logger.info(
                     f"Step: {tstep}, T: {temperature:.0f}, E: {Elist[-1]:.3f}"
                 )
-        last_atoms_filename = atoms_traj.replace("XDATCAR", "last_atoms")
+        last_atoms_filename = final_atoms or atoms_traj.replace("XDATCAR", "last_atoms")
         io.write(last_atoms_filename, self.env.atoms, format="vasp")
         return (Elist, Qlist, SROlist, action_idx_list, tlist, clist)
 
-    def run_MCMC(self, horizon, atoms_traj, logger, **simulation_params):
+    def run_MCMC(self, horizon, atoms_traj, logger, final_atoms=None, **simulation_params):
         logger.info(f"Action mode: {simulation_params.get('action_mode', 'vacancy_only')}")
+        trajectory_buffer_size = simulation_params.get("trajectory_buffer_size", 100)
+        if not isinstance(trajectory_buffer_size, int) or trajectory_buffer_size < 1:
+            raise ValueError("trajectory_buffer_size must be a positive integer")
+        trajectory_buffer = []
+        next_configuration = 2
         if simulation_params.get("annealing_time", None) is not None:
             T_scheduler = ThermalAnnealing(total_horizon=horizon,
                                            annealing_time=simulation_params["annealing_time"],
                                            T_start=simulation_params["T_start"],
                                            T_end=simulation_params["T_end"])
         atoms = self.env.atoms.copy()
         atoms, _ = self.env.relax(atoms)
-        E0 = atoms.get_potential_energy()
-        Elist = [E0] # Record for every ten steps
+        energy = atoms.get_potential_energy()
+        Elist = [energy] # Record for every ten steps
         if self.save_sro or self.sro_pixel is not None:
             atoms = self.env.atoms.copy()
             sro = get_sro_from_atoms(atoms)
@@ -310,37 +326,49 @@ def run_MCMC(self, horizon, atoms_traj, logger, **simulation_params):
             else:
                 new_T = simulation_params["temperature"]
             n_sweeps = simulation_params.get("n_sweeps", 1)
-            energy, _, count = self.mcmc_sweep(n_sweeps=n_sweeps, temperature=new_T, action_mode=simulation_params.get("action_mode", "vacancy_only"))
+            energy, _, count = self.mcmc_sweep(n_sweeps=n_sweeps, temperature=new_T, energy=energy, action_mode=simulation_params.get("action_mode", "vacancy_only"))
             if self.save_sro or self.sro_pixel is not None:
                 atoms = self.env.atoms.copy()
                 sro = get_sro_from_atoms(atoms)
                 SROlist.append(sro.tolist())
-            io.write(atoms_traj, self.env.atoms, format="vasp-xdatcar", append=True)
+            trajectory_buffer.append((next_configuration, self.env.atoms.copy()))
+            next_configuration += 1
+            if len(trajectory_buffer) >= trajectory_buffer_size:
+                _flush_xdatcar(atoms_traj, trajectory_buffer)
 
             if tstep % 10 == 0 or tstep == horizon - 1:
-                atoms = self.env.atoms.copy()
-                atoms, _ = self.env.relax(atoms)
-                Elist.append(atoms.get_potential_energy())
+                Elist.append(energy)
                 logger.info(
                     f"Step: {tstep} | Sweep: {count}/{n_sweeps}| T: {new_T:.0f} K | E: {Elist[-1]:.3f} eV"
                 )
-        last_atoms_filename = atoms_traj.replace("XDATCAR", "last_atoms")
+        _flush_xdatcar(atoms_traj, trajectory_buffer)
+        last_atoms_filename = final_atoms or atoms_traj.replace("XDATCAR", "last_atoms")
         io.write(last_atoms_filename, self.env.atoms, format="vasp")
         return (Elist, SROlist)
 
-    def mcmc_sweep(self, n_sweeps, temperature, action_mode="vacancy_only"):
+    def mcmc_sweep(self, n_sweeps, temperature, action_mode="vacancy_only", energy=None):
+        if energy is None:
+            atoms = self.env.atoms.copy()
+            atoms, _ = self.env.relax(atoms)
+            energy = atoms.get_potential_energy()
         accept = False
         count = 0
         while not accept and count < n_sweeps:
-            energy, accept = self.mcmc_step(temperature, action_mode=action_mode)
+            energy, accept = self.mcmc_step(
+                temperature, action_mode=action_mode, energy=energy
+            )
             count += 1
             self.total_mcmc_step += 1
         return energy, accept, count
 
-    def mcmc_step(self, temperature, action_mode):
+    def mcmc_step(self, temperature, action_mode, energy=None):
         accept = False
         base_prob = 0.0
         kT = temperature * 8.617 * 10**-5
+        if energy is None:
+            atoms = self.env.atoms.copy()
+            atoms, _ = self.env.relax(atoms)
+            energy = atoms.get_potential_energy()
         action_space_length = 0
         while action_space_length == 0:
             if self.kwargs.get("lattice_parameter", None) is not None:
@@ -388,16 +416,13 @@ def mcmc_step(self, temperature, action_mode):
             action_space_length = len(action_space)
 
         action = random.choice(action_space)
-        initial_atoms = self.env.atoms.copy()
-        initial_atoms, fail_min = self.env.relax(initial_atoms)
         self.env.step(action)
         next_atoms = self.env.atoms.copy()
         next_atoms, fail_next = self.env.relax(next_atoms)
-        E_prev = initial_atoms.get_potential_energy()
+        E_prev = energy
         E_next = next_atoms.get_potential_energy()
-        fail = fail_min + fail_next
 
-        if not fail:
+        if not fail_next:
             energy_diff = (E_next - E_prev) # / self.env.n_atom
             if energy_diff < 0.0:
                 base_prob = 1.0  # Automatically accept for sufficiently negative energy diff
@@ -458,4 +483,3 @@ def convert_to_graph_list(atoms: Atoms, actions: List[List[float]]) -> List[Reac
     # dataset = ReactionDataset(dataset_list)
 
     return dataset_list
-