RL_H_diffusion/gen_data_graph.py at master · htang113/RL_H_diffusion · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
from rlmd.configuration import configuration
from rlmd.trajectory import trajectory
from rlmd.step import environment
from rlmd.train_graph import ContextBandit
from rlmd.action_space import actions
from rlmd.action_space_v3 import actions as actions_v3
from rlmd.logger import setup_logger
from rgnn.models.reaction_models import PaiNN
from rgnn.models.dqn import ReactionDQN
import numpy as np
from ase import io

import os
import warnings


warnings.filterwarnings("ignore", category=UserWarning)

task = "data/Vrandom_mace_256/traj0"
horizon = 20
n_traj = 100

species = ["Cr", "Co", "Ni"]

# gcnn = PaiNN(species=species)
# reaction_model = ReactionDQN(gcnn)
# trainer = ContextBandit(reaction_model, temperature=1000, lr=5e-5)

pool = ["data/POSCARs_256/POSCAR_" + str(i) for i in range(0, 100)]

traj_list = []
if task not in os.listdir():
    os.makedirs(task, exist_ok=True)
# Configure logging
log_filename = f"{task}/logger.log"  # Define your log filename
opt_log_filenmae = f"{task}" + "/log"
if os.path.exists(opt_log_filenmae):
    os.remove(opt_log_filenmae)

logger = setup_logger("RL", log_filename)
# new_pool = []
# for filename in pool:
#     atoms = io.read(filename)
#     if len(atoms) < 500:
#         new_pool.append(filename)
# logger.info(f"Original pool num: {len(pool)}, Filtered pool num: {len(new_pool)}")


if "traj" not in os.listdir(task):
    os.mkdir(task + "/traj")
if "model" not in os.listdir(task):
    os.mkdir(task + "/model")

for epoch in range(n_traj):
    conf = configuration()
    file = pool[np.random.randint(len(pool))]
    conf.load(file)
    logger.info("epoch = " + str(epoch) + ":  " + file)
    conf.set_potential(platform="mace")
    env = environment(conf, logfile=task + "/log", max_iter=100)
    env.relax(accuracy=0.1)
    traj_list.append(trajectory(1, 0))
    for tstep in range(horizon):
        action_space = actions_v3(conf)
        act_id = np.random.choice(len(action_space))
        action = action_space[act_id]
        info = {
            "act": act_id,
            "act_probs": [],
            "act_space": action_space,
            "state": conf.atoms.copy(),
            "E_min": conf.potential(),
        }

        E_next, fail = env.step(action, accuracy=0.05)
        if not fail:
            E_s, freq, fail = env.saddle(
                action[0], n_points=8, accuracy=0.07, platform="mace"
            )
            info["E_s"], info["log_freq"] = E_s, freq
        else:
            info["E_s"] = 0
            info["log_freq"] = 0
            logger.info("fail step 1")

        info["next"], info["fail"], info["E_next"] = conf.atoms.copy(), fail, E_next
        traj_list[-1].add(info)
        if fail:
            logger.info("fail")
        if tstep % 10 == 0 and tstep > 0:
            logger.info("    t = " + str(tstep))

    try:
        traj_list[epoch].save(task + "/traj/traj" + str(epoch))
    except:
        logger.info("saving failure")