Update AGOX examples

Author: nedtaylor

example/python_pkg/C_learn/DRSS/pca.ipynb

@@ -276,7 +276,7 @@
    "outputs": [],
    "source": [
     "# get the DRAFFLE relaxed structures and energies\n",
-    "raffle_rlxd_structures = read(\"../DRAFFLE/DResponse/rlxd_structures_seed0.traj\", index=\":\")\n",
+    "raffle_rlxd_structures = read(\"../DRAFFLE/DOutput/rlxd_structures_seed0.traj\", index=\":\")\n",
     "raffle_energies_per_atom = [structure.get_potential_energy() / len(structure) for structure in raffle_rlxd_structures]\n",
     "raffle_delta_en_per_atom = np.array(raffle_energies_per_atom) - min_energy"
    ]

example/python_pkg/agox_runs/Si-Ge_run/agox_run.py

@@ -0,0 +1,191 @@
+import matplotlib
+
+matplotlib.use("Agg")
+
+from agox import AGOX
+from agox.databases import Database
+from agox.environments import Environment
+from agox.evaluators import LocalOptimizationEvaluator
+import shephex
+from ase import build
+
+@shephex.chain()
+def raffle_agox(index=42, void=0.1, rand=0.01, walk=0.25, grow=0.25, min=1.0):
+
+
+    # Using argparse if e.g. using array-jobs on Slurm to do several independent searches.
+    # from argparse import ArgumentParser
+    # parser = ArgumentParser()
+    # parser.add_argument('-i', '--run_idx', type=int, default=0)
+    # args = parser.parse_args()
+
+    ##############################################################################
+    # System & generator settings:
+    ##############################################################################
+
+    import numpy as np
+    from agox.generators.raffle import RaffleGenerator, bounds_to_confinement
+    # from mace.calculators import mace_mp
+    from agox.utils.replica_exchange.priors import get_prior
+    from ase.io import read
+
+    ## Set seed and database-index
+    seed = index
+    database_index = index
+    np.random.seed(seed)
+
+    ## Set up the MACE calculator
+    mace = get_prior("mace-mpa")
+    # calc_params = { 'model':  '/FULL/PATH/TO/MODEL/FILE' }
+    # mace = mace_mp(**calc_params)
+
+    ## Set up the method ratio dictionary and generate a unique ID
+    method_ratio = {
+        "void": void,
+        "rand": rand,
+        "walk": walk,
+        "grow": grow,
+        "min": min
+    }
+    method_id = "_".join(f"{k}{v:g}" for k, v in sorted(method_ratio.items()))
+    print("method ratio:", method_ratio)
+    print("method_id:", method_id)
+
+    ## Set up the host structure (CHANGE THIS TO YOUR HOST STRUCTURE LOCATION)
+    template = read('${HOME}/DAGOX/DSiGe/host.traj')
+    template.calc = mace
+    template.set_pbc(True)
+
+    ## Set up the reference energies for Si and Ge
+    Si_bulk = build.bulk("Si", crystalstructure="diamond", a=5.43)
+    Si_bulk.calc = mace
+    Si_reference_energy = Si_bulk.get_potential_energy() / len(Si_bulk)
+    Ge_bulk = build.bulk("Ge", crystalstructure="diamond", a=5.65)
+    Ge_bulk.calc = mace
+    Ge_reference_energy = Ge_bulk.get_potential_energy() / len(Ge_bulk)
+
+    element_energies = {'Si': Si_reference_energy, 'Ge': Ge_reference_energy}
+    print("element reference energies:", element_energies)
+
+    ## Set up the species to place in the host structure
+    species_to_place={'Si' : 16, 'Ge': 16}
+    symbols = ""
+    for key in species_to_place.keys():
+        symbols += key
+        symbols += str(species_to_place[key])
+
+    ## Set up the bounding box for the confinement
+    confinement_corner, confinement_cell = bounds_to_confinement([[0, 0, 0.34], [1, 1, 0.52]],template)
+
+    ## Set up the environment
+    environment = Environment(
+        template=template,
+        symbols=symbols,
+        confinement_cell=confinement_cell,
+        confinement_corner=confinement_corner,
+        fix_template = False,
+    )
+
+    ## Database
+    db_path = f"../db{database_index}_{method_id}.db"  # From input argument!
+    database = Database(filename=db_path, order=5)
+    database.restore_to_memory()
+
+    ## Set up the RaffleGenerator
+    generator = RaffleGenerator(
+        **environment.get_confinement(),
+        host=template,
+        database=database,
+        environment=environment,
+        species=species_to_place,
+        element_energies=element_energies,
+        transfer_data=[Si_bulk,Ge_bulk],
+        order=1,
+        width=[0.04, np.pi/160.0, np.pi/160.0],
+        kBT=0.2,
+        n_structures=5,
+        from_host=False,
+        sampler=None,
+        train=True,
+        seed=seed,
+        method_ratio = method_ratio
+    )
+
+    ##############################################################################
+    # Structural optimisation settings:
+    ##############################################################################
+
+    ## Set up the structural optimisation evaluator (defaults to ASE's BFGS optimizer)
+    # from ase.optimize import FIRE # Uncomment to use FIRE instead of BFGS
+    evaluator = LocalOptimizationEvaluator(
+        mace,
+        # optimizer = FIRE, # Uncomment to use FIRE instead of BFGS
+        gets={"get_key": "candidates"},
+        store_trajectory=False,
+        optimizer_run_kwargs={"fmax": 0.05, "steps": 200},
+        order=3,
+        number_to_evaluate=5,
+        constraints=environment.get_constraints(),
+        fix_template = False,
+    )
+
+    ## Set up post-processors to filter the generated structures.
+    # This checks the minimum distance between atoms in the generated structures.
+    from agox.postprocessors.minimum_dist import MinimumDistPostProcess
+    c1 = 0.6
+    c2 = 5
+    minimum_dist = MinimumDistPostProcess(
+        c1=c1,
+        c2=c2,
+        order=1.5,
+    )
+    minimum_dist2 = MinimumDistPostProcess(
+        c1=c1,
+        gets = {"get_key" : "evaluated_candidates"},
+        sets = {"set_key" : "evaluated_candidates"},
+        c2=c2,
+        order=3.5,
+    )
+
+    ##############################################################################
+    # Set up and run the AGOX search:
+    ##############################################################################
+
+    ## Set up the AGOX search workflow
+    agox = AGOX(generator, minimum_dist, minimum_dist2, database, evaluator, seed=seed)
+
+    ## Run the AGOX search for N_iterations
+    agox.run(N_iterations=40)
+
+
+
+if __name__ == "__main__":
+
+    walkers = 16
+
+    ## Set up the experiments to run
+    # This is an example of how to set up a series of experiments with different
+    # method ratios and different seeds.
+    experiments = (raffle_agox(directory='method_ratio/').
+                   zip(void=[1.0, 0.1,  0.1]).
+                   zip(rand=[0.1, 0.01, 0.1]).
+                   zip(walk=[0.5, 0.25, 0.5]).
+                   zip(grow=[0.5, 0.25, 0.5]).
+                   zip(min =[1.0, 1.0,  1.0]).
+                   permute(index=range(10))
+    )
+
+    ## Uncomment for Slurm execution
+    executor = shephex.executor.SlurmExecutor.from_profile(
+        'basic',
+        safety_check=False,
+        time="09:00:00",
+        ntasks_per_node=walkers,
+        cpus_per_task=8,
+    )
+
+    # ## Uncomment for local execution
+    # executor = shephex.executor.LocalExecutor()
+
+    # Submit the experiments to the executor
+    executor.execute(experiments)

example/python_pkg/agox_runs/Si-Ge_run/tidy.py

@@ -0,0 +1,47 @@
+# This script tidies up the n-body database files by organizing them into subdirectories based on the number of bodies.
+# By running this script, then the following agox line in the command line:
+#    agox analysis method_ratios/*
+# you will get analysis of the success split up by the method ratios used in the generation of the structures.
+
+import os
+import re
+import shutil
+from math import gcd
+from functools import reduce
+
+src_dir = "method_ratio"
+dst_base = "method_ratios"
+
+# Regex to match the filename pattern
+pattern = re.compile(
+    r"^db(?P<seed>\d+)_dist_grow(?P<grow>[\d\.]+)_min(?P<min>[\d\.]+)_rand(?P<rand>[\d\.]+)_void(?P<void>[\d\.]+)_walk(?P<walk>[\d\.]+)\.db$"
+)
+
+def float_str_to_int_ratios(values):
+    # Convert all to floats
+    floats = [float(v) for v in values]
+    scale = 1000  # To avoid float precision issues
+    ints = [round(f * scale) for f in floats]
+    divisor = reduce(gcd, ints)
+    return [i // divisor for i in ints]
+
+os.makedirs(dst_base, exist_ok=True)
+
+for fname in os.listdir(src_dir):
+    match = pattern.match(fname)
+    if match:
+        parts = match.groupdict()
+        ratios = float_str_to_int_ratios([
+            parts["void"], parts["rand"], parts["walk"], parts["grow"], parts["min"]
+        ])
+        # Assign to the appropriate method
+        v, r, w, g, m = ratios
+        # exit()
+        subdir_name = f"v{v}_r{r}_w{w}_g{g}_m{m}"
+        fname_new = f"db{parts['seed']}.db" #_{subdir_name}.db"
+        subdir_path = os.path.join(dst_base, subdir_name)
+        os.makedirs(subdir_path, exist_ok=True)
+
+        src_path = os.path.join(src_dir, fname)
+        dst_path = os.path.join(subdir_path, fname_new)
+        shutil.copy2(src_path, dst_path)