import os
from collections import OrderedDict
from pathlib import Path
from shutil import copyfileobj
import numpy as np
from ase.db import connect
from ase.io.vasp import write_vasp
from storq.tools import uniquify
[docs]
class Writers:
[docs]
def write_persistence_db(self, atoms=None, data=None):
"""Database containing persistence information for storq calculations.
Contains atoms, results, parameters and confiugration for easy restart.
Note that the datbase should only contain a single row.
"""
# Get the atoms object from the calculator
if not data:
data = {
"path": str(self.directory),
"parameters": self.parameters,
"jobid": self.jobid,
"conf": self.conf,
}
with connect(self.db) as db:
try:
if atoms:
db.update(1, atoms, data=data)
else:
db.update(1, data=data)
except FileNotFoundError:
db.write(atoms, data=data)
[docs]
def write_potcar(self):
"""Write POTCAR file. """
potcars_small = self.get_potcars()
with open(self.potcar, "wb") as fpotcar:
for pot_small in potcars_small.values():
with open(pot_small, "rb") as fsmall:
copyfileobj(fsmall, fpotcar)
[docs]
def write_poscar(self):
"""Write POSCAR file. """
symbs_sort = np.asarray(self.atoms.get_chemical_symbols())[self.isort]
symbs_uniq = uniquify(symbs_sort)
symbol_count = [(s, (symbs_sort == s).sum()) for s in symbs_uniq]
write_vasp(
str(self.poscar), self.atoms[self.isort], symbol_count=symbol_count, vasp5=True
)
[docs]
def write_incar(self):
"""Write INCAR file.
Boolean values are written as .TRUE./.FALSE.; integers, floats and
strings are written out as is; lists and tuples are written out as space
separated values.
Parameters
----------
fname : str
output file name
"""
special_kwargs = {
"xc",
"pp",
"setups",
"kpts",
"ldau_luj",
}
incar_keys = list(set(self.parameters) - special_kwargs)
with open(self.incar, "w") as f:
f.write("INCAR created by Storq\n")
for key in sorted(incar_keys):
val = self.parameters[key]
key = " " + key.upper()
if val is None: # Skip None values, they are used to delete tags.
pass
elif isinstance(val, bool):
s = ".TRUE." if val else ".FALSE."
f.write(f"{key} = {s}\n")
elif isinstance(val, (list, tuple, np.ndarray)):
s = " ".join([str(x) for x in val])
f.write(f"{key} = {s}\n")
else:
f.write(f"{key} = {val}\n")
[docs]
def write_kpoints(self):
"""Write KPOINTS file.
The KPOINTS file format is as follows:
* line 1: a comment
* line 2: number of kpoints
* n <= 0 Automatic kpoint generation
* n > 0 explicit number of kpoints
* line 3: kpt format
* if n > 0:
C,c,K,k = cartesian coordinates
anything else = reciprocal coordinates
* if n <= 0:
M,m,G,g for Monkhorst-Pack or Gamma grid
anything else is a special case
* line 4:
* if n <= 0, the Monkhorst-Pack grid
* if n > 0, then one line per kpoint
* line 5: if n <= 0 it is the gamma shift
"""
# handle the input args
p = self.parameters
kpts = p["kpts"] # guaranteed to be set
# we are using automatic or semi-automatic mode
if "size" in kpts or "spacing" in kpts:
if kpts["gamma"]:
mode = "Gamma"
else:
mode = "Monkhorst-Pack"
num_kpts = 0
elif "auto" in kpts:
mode = "Auto"
num_kpts = 0
# we are in an explicit listing mode
elif "list" in kpts:
if kpts.get("cartesian", False):
mode = "Cartesian"
else:
mode = "Reciprocal"
num_kpts = len(kpts["list"])
# we are in line mode
elif "lines" in kpts:
mode = "Line-mode"
num_lines = len(kpts["lines"])
# write the file
with open(self.kpoints, "w") as f:
# line 1: comment
# hide some info here to restore kpts to right state on restart
comment = "KPOINTS created by Atomic Simulation Environment"
if "spacing" in kpts:
comment = "spacing {0:.4f}".format(kpts["spacing"])
if kpts.get("surface", False) is True:
comment += " surface"
f.write(comment + "\n")
# line 2: number of kpts
if mode == "Line-mode":
f.write("{}\n".format(kpts.get("intersections", 10)))
else:
f.write("{}\n".format(num_kpts))
# line 3: the mode
f.write(mode + "\n")
# line 4+: k-point specification
if mode in ["Monkhorst-Pack", "Gamma"]:
if "size" in kpts:
kpts_size = kpts["size"]
elif "spacing" in kpts:
kpts_size = [1, 1, 1]
rec_cell = self.atoms.cell.reciprocal()
kpts_size[0] = int(
max(
1,
np.ceil(
np.linalg.norm(rec_cell[0])
* 2.0
* np.pi
/ kpts["spacing"]
),
)
)
kpts_size[1] = int(
max(
1,
np.ceil(
np.linalg.norm(rec_cell[1])
* 2.0
* np.pi
/ kpts["spacing"]
),
)
)
kpts_size[2] = int(
max(
1,
np.ceil(
np.linalg.norm(rec_cell[2])
* 2.0
* np.pi
/ kpts["spacing"]
),
)
)
if kpts.get("surface", False) is False:
f.write("{0} {1} {2}\n".format(*tuple(kpts_size)))
else:
f.write("{0} {1} 1\n".format(*tuple(kpts_size)))
elif mode == "Auto":
f.write("{0}\n".format(kpts["auto"]))
elif mode == "Line-mode":
if kpts.get("caretesian", False):
f.write("Cartesian\n")
else:
f.write("Reciprocal\n")
for n in range(0, num_lines, 2):
f.write("{0} {1} {2}\n".format(*tuple(kpts["lines"][n])))
f.write("{0} {1} {2}\n".format(*tuple(kpts["lines"][n + 1])))
f.write("\n")
elif mode in ["Cartesian", "Reciprocal"]:
for n in range(0, num_kpts):
f.write("{0} {1} {2} {3}\n".format(*tuple(kpts["list"][n])))
# if we are in 'semi' automatic mode and have a shift
if mode in ["Monkhorst-Pack", "Gamma"]:
f.write(
"{0} {1} {2}\n".format(*tuple(kpts.get("shift", [0.0, 0.0, 0.0])))
)