ENH: Add an adaptive displacement list for atoms

eon/_utils.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+from enum import Enum
+import sys
+import subprocess
+import tempfile
+import logging
+import typing as typ
+from dataclasses import dataclass
+
+from eon import fileio as eio
+from eon import atoms as eatm
+from eon import config as econf
+
+
+class ScriptType(Enum):
+    STATE = 0
+    DISP = 1
+
+@dataclass
+class ScriptConfig:
+    """Configuration for running an external atom list script."""
+
+    script_path: str
+    scratch_path: str
+    root_path: str
+
+    def __post_init__(self):
+        """
+        Post-initialization processing to validate and resolve paths.
+        This method is automatically called by the dataclass constructor.
+        """
+        for pth in (self.script_path, self.scratch_path):
+            if not pth.is_absolute():
+                pth = self.root_path / pth
+
+    @classmethod
+    def from_eon_config(cls, config: econf.ConfigClass, stype: ScriptType) -> typ.Self:
+        """
+        Factory method to create a ScriptConfig instance from the main EON config.
+        """
+        if stype == ScriptType.STATE:
+            script_path = Path(config.displace_atom_kmc_state_script)
+        elif stype == ScriptType.DISP:
+            script_path = Path(config.displace_atom_kmc_step_script)
+
+        return cls(
+            script_path=script_path,
+            scratch_path=Path(config.path_scratch),
+            root_path=Path(config.path_root),
+        )
+
+
+def gen_ids_from_con(sconf: ScriptConfig, reactant: eatm.Atoms, logger: logging.Logger):
+    script_path = Path(sconf.script_path)
+    if not script_path.is_absolute():
+        script_path = sconf.root_path / sconf.script_path
+
+    if not script_path.is_file():
+        logger.error(f"displace_atom_list_script not found: {script_path}")
+        return []
+
+    # Use a secure temporary file to pass the structure to the script
+    # The file is automatically deleted when the 'with' block is exited.
+    sconf.scratch_path.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile(
+        mode="w+", delete=True, suffix=".con", dir=sconf.scratch_path
+    ) as tmpf:
+        # Save the displaced structure to the temporary file
+        eio.savecon(tmpf, reactant)
+        # Ensure all data is written to disk before the script reads it
+        tmpf.flush()
+
+        try:
+            # Execute the script, passing the temporary filename as an argument
+            proc = subprocess.run(
+                [sys.executable, str(script_path), tmpf.name],  # Pass filename
+                capture_output=True,
+                text=True,
+                check=True,
+                cwd=sconf.root_path,
+            )
+            atom_list_str = proc.stdout.strip()
+            return atom_list_str
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Error running displace_atom_list_script '{script_path}':")
+            logger.error(f"Stderr: {e.stderr.strip()}")
+            sys.exit(1)
+        except Exception as e:
+            logger.error(
+                f"An unexpected error occurred while running the displacement script: {e}"
+            )
+            sys.exit(1)

eon/akmc.py

@@ -30,6 +30,7 @@
 from eon import superbasinscheme
 from eon import askmc
 from eon import movie
+from eon import _utils as utl
 
 def akmc(config: ConfigClass = EON_CONFIG, steps=0):
     """Poll for status of AKMC clients and possibly make KMC steps.
@@ -67,6 +68,17 @@ def akmc(config: ConfigClass = EON_CONFIG, steps=0):
     # Load metadata, the state list, and the current state.
     start_state_num, time, previous_state_num, first_run, previous_temperature = get_akmc_metadata()
 
+
+    states = get_statelist(kT)
+    current_state = states.get_state(start_state_num)
+
+    # --- START: DYNAMIC ATOM LIST SCRIPT EXECUTION ---
+    if config.displace_atom_kmc_state_script:
+        atom_list = current_state.get_displacement_atom_list(config)
+        # Overwrite the global config value with the list specific to this state.
+        config.displace_atom_list = atom_list
+    # --- END: DYNAMIC ATOM LIST SCRIPT EXECUTION ---
+
     if first_run:
         previous_temperature = config.main_temperature
 
@@ -85,8 +97,6 @@ def akmc(config: ConfigClass = EON_CONFIG, steps=0):
         # Keep the new temperature.
         previous_temperature = config.main_temperature
 
-    states = get_statelist(kT)
-    current_state = states.get_state(start_state_num)
     if previous_state_num == -1:
         previous_state = current_state
     else:
@@ -176,6 +186,8 @@ def get_superbasin_scheme(states, config):
         superbasining = superbasinscheme.TransitionCounting(config.sb_path, states, config.main_temperature / 11604.5, config.sb_tc_ntrans)
     elif config.sb_scheme == 'energy_level':
         superbasining = superbasinscheme.EnergyLevel(config.sb_path, states, config.main_temperature / 11604.5, config.sb_el_energy_increment)
+    elif config.sb_scheme == 'rate':
+        superbasining = superbasinscheme.RateThreshold(config.sb_path, states, config.main_temperature / 11604.5, config.sb_rt_rate_threshold)
     return superbasining
 
 

eon/config.py

@@ -243,6 +243,7 @@ def mpiexcepthook(type, value, traceback):
         self.saddle_dynamics_temperature = parser.getfloat('Saddle Search', 'dynamics_temperature')
         self.displace_random_weight = parser.getfloat('Saddle Search', 'displace_random_weight')
         self.displace_listed_atom_weight = parser.getfloat('Saddle Search', 'displace_listed_atom_weight')
+        self.displace_atom_kmc_state_script = parser.get('Saddle Search', 'displace_atom_kmc_state_script')
         self.displace_listed_type_weight = parser.getfloat('Saddle Search', 'displace_listed_type_weight')
         self.displace_all_listed = parser.getboolean('Saddle Search', 'displace_all_listed')
         self.displace_under_coordinated_weight = parser.getfloat('Saddle Search', 'displace_under_coordinated_weight')
@@ -310,6 +311,8 @@ def mpiexcepthook(type, value, traceback):
             self.sb_tc_ntrans = parser.getint('Coarse Graining', 'number_of_transitions')
         elif self.sb_scheme == 'energy_level':
             self.sb_el_energy_increment = parser.getfloat('Coarse Graining', 'energy_increment')
+        elif self.sb_scheme == 'rate':
+            self.sb_rt_rate_threshold = parser.getfloat('Coarse Graining', 'rate_threshold')
         self.sb_superbasin_confidence = parser.getboolean('Coarse Graining', 'superbasin_confidence')
 
         self.askmc_on = parser.getboolean('Coarse Graining','use_askmc')

eon/config.yaml

@@ -1095,6 +1095,11 @@ Saddle Search:
             kind: float
             default: 0.0
 
+        # This is run once for each new state
+        displace_atom_kmc_state_script:
+            kind: string
+            default: ''
+
         displace_under_coordinated_weight:
             kind: float
             default: 0.0
@@ -1275,6 +1280,10 @@ Coarse Graining:
             kind: float
             default: 1.5
 
+        rate_threshold:
+            kind: float
+            default: 1.0e9
+
         energy_increment:
             kind: float
             default: 0.01

eon/displace.py

@@ -3,6 +3,10 @@
 
 import os, re
 from math import cos, sin
+import sys
+import subprocess
+import tempfile
+from pathlib import Path
 import numpy
 
 from eon import atoms
@@ -362,7 +366,6 @@ def make_displacement(self):
             epicenter = self.listed_atoms
         else:
             epicenter = self.listed_atoms[numpy.random.randint(len(self.listed_atoms))]
-        logger.debug("Listed atom displacement epicenters: %s", epicenter)
         return self.get_displacement(epicenter)
 
 class ListedTypes(Displace):

eon/explorer.py

@@ -17,6 +17,7 @@
 #import kdb
 from eon import recycling
 from eon import eon_kdb as kdb
+from eon import _utils as utl
 
 from eon.config import config as EON_CONFIG
 from eon.config import ConfigClass # Typing
@@ -194,7 +195,7 @@ def make_jobs(self):
 
         # Merge potential files into invariants
         invariants = dict(invariants, **io.load_potfiles(self.config.path_pot))
-
+        atom_list_str = str(self.state.info.get("Saddle Search", "displace_atom_list", ""))
         for i in range(num_to_make):
             search = {}
             # The search dictionary contains the following key-value pairs:
@@ -216,6 +217,9 @@ def make_jobs(self):
             if self.config.saddle_method == 'dynamics' and disp_type != 'dynamics':
                 ini_changes.append( ('Saddle Search', 'method', 'min_mode') )
 
+            if atom_list_str:
+                ini_changes.append(("Saddle Search", "displace_atom_list", atom_list_str))
+
             search['config.ini'] = io.modify_config(self.config.config_path, ini_changes)
 
             if displacement:

eon/schema.py

@@ -652,6 +652,10 @@ class SaddleSearchConfig(BaseModel):
         default=[-1],
         description="The individual index should be separated by a comma. Example: 10, 20, -1 would be the 10, 20, and the last atom.",
     )
+    displace_atom_kmc_state_script: str = Field(
+        default="",
+        description="A Python script which returns a string of atoms to be displaced. This is run once for each new AKMC state.",
+    )
     displace_listed_type_weight: float = Field(
         default=0.0,
         description="Relative probability to displace with an epicenter listed in displace_type_list.",
@@ -864,23 +868,28 @@ class CoarseGrainingConfig(BaseModel):
         default="superbasin",
         description="File name for the state-specific data stored within each of the state directories.",
     )
-    superbasin_scheme: Literal["energy_level", "transition_counting"] = Field(
+    superbasin_scheme: Literal["energy_level", "transition_counting", "rate"] = Field(
         default="transition_counting",
         description="MCAMC provides a method for calculating transition rates across superbasins. An additional method is needed in order to decide when to combine states into a superbasin.",
     )
     """
     Options:
      - ``transition_counting``: Counts the number of times that the simulation has transitioned between a given pair of states. After a critical number of transitions have occurred, the pair of states are merged to form a superbasin.
-     - ``energy_level``: States are merged based on energy levels.
+     - ``energy_level``: States are merged based on energy levels filling up existing basins.
+     - ``rate``: States are merged based only on rate criteria.
     """
     max_size: int = Field(
         default=0,
         description="The maximal number of states that will be merged together. If 0, there is no limit.",
     )
     number_of_transitions: int = Field(
         default=5,
-        description="If the transition counting scheme is being used, this is the number of transitions that must occur between two states before they are merged into a superbasin.",
+        description="This is the number of transitions that must occur between two states before they are merged into a superbasin.",
     )
+    """
+    Only used if :any:`eon.schema.CoarseGrainingConfig.superbasin_scheme` is
+    ``transition_counting``.
+    """
     energy_increment: float = Field(
         default=0.01,
         description="The first time each state is visited, it is assigned an energy level first equal to the energy of the minimum. Every time the state is visited again by the Monte Carlo simulation, the energy level is increased by this amount.",
@@ -889,6 +898,14 @@ class CoarseGrainingConfig(BaseModel):
     Only used if :any:`eon.schema.CoarseGrainingConfig.superbasin_scheme` is
     ``energy_level``.
     """
+    rate_threshold: float = Field(
+        default=1e9,
+        description="Any state with a rate (in 1/time) greater than this is merged into a single basin",
+    )
+    """
+    Only used if :any:`eon.schema.CoarseGrainingConfig.superbasin_scheme` is
+    ``rate``.
+    """
     superbasin_confidence: bool = Field(
         default=True,
         description="Superbasin KMC confidence.",

eon/state.py

@@ -123,3 +123,36 @@ def get_num_procs(self):
     def get_process_table(self):
         self.load_process_table()
         return self.procs
+
+    def get_displacement_atom_list(self, config: ConfigClass) -> list[int]:
+         """
+         Gets the list of atoms for displacement.
+
+         If the list has already been generated and stored in this state's
+         info file, it is read from there. Otherwise, the displacement
+         script is run, and the result is saved to the info file for future use.
+
+         Returns:
+             A list of integer atom indices for displacement.
+         """
+         # Try to get the list from the state's own info file first.
+         # We use a try/except block in case the section or option doesn't exist yet.
+         try:
+             # The list is stored as a string, so we need to parse it.
+             atom_list_str = self.info.get("Saddle Search", "displace_atom_list")
+             return atom_list_str
+         except Exception:
+             pass
+
+         # If we are here, the list was not in the info file. Run the script.
+         from eon import _utils as utl
+         logger.info(f"State {self.number}: Generating displacement atom list for the first time.")
+         script_config = utl.ScriptConfig.from_eon_config(config, utl.ScriptType.STATE)
+         atom_list_str = utl.gen_ids_from_con(script_config, self.get_reactant(), logger)
+         if atom_list_str:
+             self.info.set("Saddle Search", "displace_atom_list", atom_list_str)
+             return atom_list_str
+         else:
+             logger.warning(f"Script for state {self.number} produced no output.")
+             self.info.set("Saddle Search", "displace_atom_list", "")
+             return []