File size: 5,951 Bytes
4947683 | 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | """Copyright (c) Meta Platforms, Inc. and affiliates."""
from __future__ import annotations
import time
from argparse import ArgumentParser, Namespace
from functools import partial
from pathlib import Path
import numpy as np
import pandas as pd
import submitit
import torch
from chgnet.model import CHGNet
from flowmm.chgnet_ import RelaxationData, prerelax_with_chgnet
from flowmm.pymatgen_ import cdvae_to_structure, diffcsp_to_structure, get_get_structure
#
# Here's a really nice script for CHGNet
# https://github.com/janosh/matbench-discovery/blob/main/models/chgnet/test_chgnet.py
#
def wait_for_jobs_to_finish(jobs: list, sleep_time_s: int = 5) -> None:
# wait for the job to be finished
num_finished = 0
print(f"number of jobs: {len(jobs):02d}")
while num_finished < len(jobs):
time.sleep(sleep_time_s)
num_finished = sum(job.done() for job in jobs)
print(f"number finished: {num_finished:02d}", flush=True, end="\r")
print("")
print("jobs done!")
return None
def prerelax_multiple(
gen_file: Path,
steps: int,
index: np.ndarray,
json_file: Path,
) -> pd.DataFrame:
"""cdvae structures, max relaxation steps, list of indexes -> json file"""
chgnet = CHGNet.load()
device = next(chgnet.parameters()).device
gen = torch.load(gen_file, map_location=device)
for k, v in gen.items():
if isinstance(v, torch.Tensor):
gen[k] = v.squeeze(0)
rd = RelaxationData()
for i in index.tolist():
get_structure = get_get_structure(gen, clip_atom_types=False)
try:
structure = get_structure(gen, i)
pair = prerelax_with_chgnet(structure=structure, chgnet=chgnet, steps=steps)
rd.index.append(i)
rd.e_gen.append(pair.energies[0])
rd.e_relax.append(pair.energies[1])
rd.n_to_relax.append(pair.n_steps_to_relax)
rd.rms_dist.append(pair.rms_dist)
rd.matched.append(pair.match)
rd.converged.append(True if pair.n_steps_to_relax < steps else False)
rd.exception.append(False)
rd.num_sites.append(structure.num_sites)
rd.structure.append(pair.structure_dicts[-1])
except Exception as exp:
print(exp)
rd.index.append(i)
rd.e_gen.append(None)
rd.e_relax.append(None)
rd.n_to_relax.append(pd.NA)
rd.rms_dist.append(None)
rd.matched.append(False)
rd.converged.append(False)
rd.exception.append(True)
rd.num_sites.append(pd.NA)
rd.structure.append({})
# dataclasses.fields(rd) produces blank :(
# dataclasses.asdict(rd) also produces blank
pre_pandas = dict(
(field.name, getattr(rd, field.name))
for field in getattr(rd, "__dataclass_fields__").values()
)
if "batch_indices" in gen.keys():
pre_pandas["index"] = gen["batch_indices"][index].tolist()
else:
print("no batch_indices given so assuming order is correct")
df = pd.DataFrame.from_dict(pre_pandas).set_index("index")
df.loc[:, "e_delta"] = np.abs(df.loc[:, "e_relax"] - df.loc[:, "e_gen"])
df.to_json(json_file)
return df
def main(args: Namespace) -> None:
# get file
gen_file = Path(args.path_to_structures)
print(f"reading file: {str(gen_file.resolve())}")
# determine how many structures there are
num_structures = torch.load(gen_file)["num_atoms"].squeeze().numel()
if args.num_structures is not None:
assert args.num_structures <= num_structures
num_structures = args.num_structures
print(f"{num_structures=}")
# split structures into parts
index = np.arange(num_structures)
indexes = np.array_split(index, args.num_jobs)
indexes = [i for i in indexes if i.size > 0] # filter out blanks
# compute the energies
files = [
Path(args.log_dir) / f"{index.min():07d}-{index.max():07d}.json"
for index in indexes
]
# set cluster
cluster = "local" if args.slurm_partition is None else "slurm"
if args.debug:
cluster = "debug"
executor = submitit.AutoExecutor(
folder=args.log_dir,
cluster=cluster,
)
executor.update_parameters(
slurm_array_parallelism=args.num_jobs,
nodes=1,
slurm_ntasks_per_node=1,
cpus_per_task=4,
gpus_per_node=0,
timeout_min=args.timeout_min,
slurm_partition=args.slurm_partition,
)
doit = partial(prerelax_multiple, gen_file, args.steps)
jobs = executor.map_array(doit, indexes, files)
wait_for_jobs_to_finish(jobs)
# write to file
df = pd.concat([job.result() for job in jobs])
df.to_json(Path(args.path_to_json))
print(f"wrote file to: ")
print(f"{args.path_to_json}")
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument(
"path_to_structures",
type=Path,
help="path to the eval_pt, typically `eval_for_dft.pt`, file from cdvae / diffcsp / rfm",
)
parser.add_argument("path_to_json", type=Path, help="output")
parser.add_argument("log_dir", type=Path)
parser.add_argument(
"--num_jobs",
type=int,
default=1,
help="number of jobs to divide structures between, only works when --slurm_partition is set",
)
parser.add_argument("-n", "--num_structures", type=int, default=None)
parser.add_argument(
"--steps", type=int, default=1_500, help="maximum number of steps in relaxation"
)
parser.add_argument("--timeout_min", type=int, default=300)
parser.add_argument("--slurm_partition", type=str, default=None)
parser.add_argument(
"--debug",
action="store_true",
help="run the prerelaxations sequentially in the same process as this script.",
)
args = parser.parse_args()
main(args)
|