File size: 5,408 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 | """Copyright (c) Meta Platforms, Inc. and affiliates."""
"""
There are several relevant settings for DFT.
1. The MP20 settings. This is what we use to evaluate CDVAE and DiffCSP which were trained on that data.
2. The OCP settings. These will be used to evaluate models trained on larger datasets than MP20.
"""
from __future__ import annotations
from argparse import ArgumentParser, Namespace
from copy import deepcopy
from pathlib import Path
import pandas as pd
import yaml
from joblib import Parallel, delayed
from pymatgen.core import Structure
from pymatgen.io.vasp.outputs import Vasprun
from pymatgen.io.vasp.sets import MPRelaxSet
from tqdm import tqdm
from flowmm.pandas_ import filter_prerelaxed, maybe_get_missing_columns
from flowmm.pymatgen_ import COLUMNS_COMPUTATIONS, get_nary, to_structure
CHGNET_STABILITY_THRESHOLD = 0.1
def parse_vasp(path: Path | str):
path = Path(path)
vrun = Vasprun(filename=path)
vrun.final_structure
vrun.final_energy
def write_vasp_directory(
structure: Structure | dict,
path: Path | str,
potcar_spec: bool = False,
) -> None:
path = Path(path)
structure = to_structure(structure)
relax_set = MPRelaxSet(structure=structure)
try:
relax_set.write_input(
output_dir=str(path.resolve()),
make_dir_if_not_present=True,
potcar_spec=potcar_spec,
)
except (TypeError, KeyError) as exp:
print(f"{exp}")
return None
def subsample_by_nary(
df: pd.DataFrame, num_structures: int
) -> tuple[pd.DataFrame, dict[int, float]]:
assert num_structures <= len(df)
weight_map = (df.groupby("nary").count()["e_relax"] / len(df)).to_dict()
df["weights"] = df["nary"].map(weight_map)
return df.sample(num_structures, weights=df["weights"]), weight_map
def subsample(df: pd.DataFrame, num_structures: int) -> pd.DataFrame:
return df.sample(num_structures)
def get_equal_weights(df: pd.DataFrame) -> dict[int, float]:
return {i: 1 / df["nary"].max() for i in range(1, df["nary"].max() + 1)}
def main(args: Namespace) -> None:
print(f"{args.root=}")
# easier to reason with positives...
filter_converged = not args.do_not_filter_converged
# load data
print("reading json")
df = pd.read_json(args.json_in)
print(f"{len(df)=} loaded")
df = maybe_get_missing_columns(df, COLUMNS_COMPUTATIONS)
print(f"{len(df)=} get missing cols")
df = filter_prerelaxed(
df,
filter_converged=filter_converged,
maximum_nary=args.maximum_nary,
)
print(f"{len(df)=} filtered prerelaxed")
# no potential for these
df = df[~df["chemsys"].map(lambda x: "Ra" in x)]
df = df[~df["chemsys"].map(lambda x: "At" in x)]
df = df[~df["chemsys"].map(lambda x: "Po" in x)]
print(f"{len(df)=} removed structures without labels")
df["nary"] = df["chemsys"].map(get_nary)
if args.ehulls is not None:
df_hull = pd.read_json(args.ehulls)
df = df.join(df_hull, how="inner")
# filter out high energy structures
df = df[df["e_above_hull_per_atom_chgnet"] < CHGNET_STABILITY_THRESHOLD]
print("finished")
num_structures = min(args.max_num_structures, len(df))
print(f"In the end, producing {num_structures=} dft input files")
# choose subsampling stragegy
if args.subsample is None:
pass
elif args.subsample == "first":
df = df[:num_structures]
elif args.subsample == "random":
print("subsample randomly with equal weight")
df = subsample(df, num_structures)
weights = get_equal_weights(df)
elif args.subsample == "nary":
print("subsample weighted by nary")
df, weights = subsample_by_nary(df, num_structures)
else:
raise ValueError()
print(f"{len(df)=}")
# save basics
args.root.mkdir(parents=True, exist_ok=True)
with open(args.root / "args.yaml", "w") as f:
to_save = deepcopy(vars(args))
to_save["json_in"] = str(to_save["json_in"])
to_save["root"] = str(to_save["root"])
yaml.dump(to_save, f)
with open(args.root / "weights.yaml", "w") as f:
yaml.dump(weights, f)
df.drop("structure", axis=1).to_csv(args.root / "dataframe.csv")
Parallel(n_jobs=args.n_jobs)(
delayed(write_vasp_directory)(structure, args.root / f"{i:06d}")
for i, structure in tqdm(df["structure"].items())
)
if __name__ == "__main__":
parser = ArgumentParser()
parser.add_argument("json_in", type=Path, help="prerelaxed dataframe")
parser.add_argument("root", type=Path)
parser.add_argument(
"--subsample",
type=str,
choices=["first", "random", "nary"],
default="random",
)
parser.add_argument("--max_num_structures", type=int, default=10_000)
parser.add_argument(
"--maximum_nary",
type=int,
default=None,
help="Any queries to structures with higher nary are avoided.",
)
parser.add_argument(
"--ehulls",
type=Path,
default=None,
help="if given, filter to chgnet e_above_hull_per_atom < 0.01",
) # an alternative is to triage with chgnet instead of cutoff
parser.add_argument("--do_not_filter_converged", action="store_true")
parser.add_argument("--n_jobs", type=int, default=1)
args = parser.parse_args()
main(args)
|