def get_optimized_structure_cell_opt(self, xyztraj, cell_opt_out,
directory):
"""
:param xyztraj: output xxx.xyz trajectory file
:param cell_opt_out: output file of cell opt
:param directory: directory to put the optimized structure
Note: deal with CELL_OPT
"""
with open(cell_opt_out, 'r') as fin:
cell_opt_out_lines = fin.readlines()
for i in range(len(cell_opt_out_lines)):
#if len(cell_opt_out_lines[i].split()) == 0:
# continue
if "GEOMETRY OPTIMIZATION COMPLETED" in cell_opt_out_lines[i]:
a = []
b = []
c = []
for j in range(3):
a.append(float(cell_opt_out_lines[i + 6].split()[4 + j]))
b.append(float(cell_opt_out_lines[i + 7].split()[4 + j]))
c.append(float(cell_opt_out_lines[i + 8].split()[4 + j]))
cell = []
cell.append(a)
cell.append(b)
cell.append(c)
traj_lines = []
with open(xyztraj, 'r') as fin:
for line in fin:
if len(line.split()) == 0:
continue
traj_lines.append(line)
natom = int(traj_lines[0].split()[0])
nimage = int(len(traj_lines) / (natom + 2))
with open(os.path.join(directory, "cell-opt-optimized.xyz"),
'w') as fout:
fout.write("%d\n" % natom)
fout.write(
"cell: %.9f %.9f %.9f | %.9f %.9f %.9f | %.9f %.9f %.9f\n" % (
cell[0][0],
cell[0][1],
cell[0][2],
cell[1][0],
cell[1][1],
cell[1][2],
cell[2][0],
cell[2][1],
cell[2][2],
))
for i in range((natom + 2) * (nimage - 1) + 2,
(natom + 2) * nimage):
fout.write(traj_lines[i])
a = read_structure(
filepath=os.path.join(directory, "cell-opt-optimized.xyz"))
write_structure(structure=a,
filepath=os.path.join(directory,
"cell-opt-optimized.cif"))
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-i",
"--input",
type=str,
required=True,
help="input PARCHG file")
parser.add_argument("--output-structure",
type=str,
default="parchg.cif",
help="output stucture contained in PARCHG")
parser.add_argument("-o",
"--output",
type=str,
default="stm",
help="prefix of the output image file name")
parser.add_argument("--levels",
type=int,
default=10,
help="levels of the color map or color bar")
parser.add_argument(
"-z",
"--z",
type=float,
default=1,
help=
"a value between 0 and 1, indicat height of in z direction to print the plot"
)
parser.add_argument("--cmap",
type=str,
default="gray",
choices=[
"gray", "hot", "afmhot", "Spectral", "plasma",
"magma", "hsv", "rainbow", "brg"
])
# ==========================================================
# transfer parameters from the arg subparser to static_run setting
# ==========================================================
args = parser.parse_args()
parchg_filepath = args.input
with open(parchg_filepath, "r") as fin:
parchg = fin.readlines()
for i in range(len(parchg)):
if len(parchg[i].split()) == 0:
first_blank_line = i
break
first_augmentation_line = None
for i in range(len(parchg)):
if "augmentation" in parchg[i]:
first_augmentation_line = i
break
os.system("mkdir -p /tmp/pymatflow/")
with open("/tmp/pymatflow/POSCAR", "w") as fout:
for i in range(first_blank_line):
fout.write(parchg[i])
structure = read_structure("/tmp/pymatflow/POSCAR")
write_structure(structure=structure, filepath=args.output_structure)
ngxf = int(parchg[first_blank_line + 1].split()[0])
ngyf = int(parchg[first_blank_line + 1].split()[1])
ngzf = int(parchg[first_blank_line + 1].split()[2])
if first_augmentation_line == None:
#data = np.loadtxt(chg[first_blank_line+2:])
tmp_str = "".join(parchg[first_blank_line + 2:])
data = np.fromstring(tmp_str, sep="\n")
else:
#data = np.loadtxt(chg[first_blank_line+2:first_augmentation_line])
tmp_str = "".join(parchg[first_blank_line + 2:first_augmentation_line])
data = np.fromstring(tmp_str, sep="\n")
data = data.reshape(ngzf, ngyf, ngxf)
# -------------------------------------------------------
# gray scale image only for z direction
# may not work for triclinic and monoclinic crystal system
# -------------------------------------------------------
zi = int((data.shape[0] - 1) * args.z)
#img = data[i, ::-1, ::]
img = data[zi, ::, ::]
img = (img - img.min()) / (img.max() - img.min()) * 255
# need to do a transform when the cell is not Orthorhombic
# skew the image
a = np.array(structure.cell[0])
b = np.array(structure.cell[1])
cosangle = a.dot(b) / (np.linalg.norm(a) * np.linalg.norm(b))
angle = np.arccos(cosangle) * 180 / np.pi
ax = plt.axes() #plt.figure()
n1 = ngxf
n2 = ngyf
n1_right = n1
n1_left = -(n2 * np.tan((angle - 90) / 180 * np.pi))
#im = ax.imshow(img, cmap="gray", extent=[n1_left, n1_right, 0, n2], interpolation="none", origin="lower", clip_on=True)
im = ax.imshow(img,
cmap="gray",
extent=[0, n1, 0, n2],
interpolation="none",
origin="lower",
clip_on=True)
#im = plt.imshow(data[i, :, :], cmap="gray")
trans_data = mtransforms.Affine2D().skew_deg(90 - angle, 0) + ax.transData
im.set_transform(trans_data)
# display intended extent of the image
x1, x2, y1, y2 = im.get_extent()
# do not view the line, but it is needed to be plot so the intended image is dispalyed completely
ax.plot([x1, x2, x2, x1, x1], [y1, y1, y2, y2, y1],
"y--",
transform=trans_data,
visible=False)
#ax.set_xlim(n1_left, n1_right)
#ax.set_ylim(0, n2)
plt.colorbar(im)
ax.autoscale()
plt.tight_layout()
plt.savefig(args.output + "-z-%f.grayscale.png" % args.z)
plt.close()
# -----------------------------------------------------------------------------
# 2D contour plot
#------------------------------------------------------------------------------
nx = np.linspace(0, 1, ngxf)
ny = np.linspace(0, 1, ngyf)
X, Y = np.meshgrid(
nx, ny
) # now this Mesh grid cannot be used directly, we have to calc the real x y for it
for xi in range(len(nx)):
for yi in range(len(ny)):
X[yi, xi] = structure.cell[0][0] * nx[xi] + structure.cell[1][
0] * ny[yi]
Y[yi, xi] = structure.cell[0][1] * nx[xi] + structure.cell[1][
1] * ny[yi]
Z = data[zi, :, :]
Z = (Z - Z.min()) / (Z.max() - Z.min()) * 255
# fill color, three color are divided into three layer(6)
# cmap = plt.cm.hot means using thermostat plot(graduated red yellow)
#cset = plt.contourf(X, Y, Z, levels=args.levels, cmap=plt.cm.hot)
cset = plt.contourf(X, Y, Z, levels=args.levels, cmap=args.cmap)
contour = plt.contour(X, Y, Z, levels=[20, 40], colors='k')
plt.colorbar(cset)
plt.autoscale()
plt.tight_layout()
#plt.show()
plt.axis("equal") # set axis equally spaced
plt.xlabel('x')
plt.ylabel('y')
plt.savefig(args.output + ".2d-contour-z-%f.png" % args.z)
plt.close()
def get_optimized_structure_geo_opt(self, xyztraj, geo_opt_inp, directory):
"""
:param xyztraj: output xxx.xyz trajectory file
:param geo_opt_inp: input file for geo opt
:param directory: directory to put the optimized structure
Note: deal with GEO_OPT
"""
with open(geo_opt_inp, 'r') as fin:
geo_opt_inp_lines = fin.readlines()
for i in range(len(geo_opt_inp_lines)):
if len(geo_opt_inp_lines[i].split(
)) > 0 and geo_opt_inp_lines[i].split()[0].upper() == "&CELL":
j = 1
while not ("&END" in geo_opt_inp_lines[i + j].upper().split()
and "CELL"
in geo_opt_inp_lines[i + j].upper().split()):
if len(geo_opt_inp_lines[i + j].split(
)) > 0 and geo_opt_inp_lines[i +
j].split()[0].upper() == "A":
a = []
for k in range(3):
a.append(
float(geo_opt_inp_lines[i + j].split()[k + 1]))
elif len(geo_opt_inp_lines[i + j].split(
)) > 0 and geo_opt_inp_lines[i +
j].split()[0].upper() == "B":
b = []
for k in range(3):
b.append(
float(geo_opt_inp_lines[i + j].split()[k + 1]))
elif len(geo_opt_inp_lines[i + j].split(
)) > 0 and geo_opt_inp_lines[i +
j].split()[0].upper() == "C":
c = []
for k in range(3):
c.append(
float(geo_opt_inp_lines[i + j].split()[k + 1]))
j += 1
break
cell = []
cell.append(a)
cell.append(b)
cell.append(c)
traj_lines = []
with open(xyztraj, 'r') as fin:
for line in fin:
if len(line.split()) == 0:
continue
traj_lines.append(line)
natom = int(traj_lines[0].split()[0])
nimage = int(len(traj_lines) / (natom + 2))
with open(os.path.join(directory, "geo-opt-optimized.xyz"),
'w') as fout:
fout.write("%d\n" % natom)
fout.write(
"cell: %.9f %.9f %.9f | %.9f %.9f %.9f | %.9f %.9f %.9f\n" % (
cell[0][0],
cell[0][1],
cell[0][2],
cell[1][0],
cell[1][1],
cell[1][2],
cell[2][0],
cell[2][1],
cell[2][2],
))
for i in range((natom + 2) * (nimage - 1) + 2,
(natom + 2) * nimage):
fout.write(traj_lines[i])
a = read_structure(
filepath=os.path.join(directory, "geo-opt-optimized.xyz"))
write_structure(structure=a,
filepath=os.path.join(directory,
"geo-opt-optimized.cif"))
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-i",
"--input",
type=str,
required=True,
help="input cube file")
parser.add_argument("--output-structure",
type=str,
default="cube.cif",
help="output stucture contained in PARCHG")
parser.add_argument("-o",
"--output",
type=str,
default="cube",
help="prefix of the output image file name")
parser.add_argument("--levels",
type=int,
default=10,
help="levels of the color map or color bar")
parser.add_argument(
"-z",
"--z",
type=float,
default=1,
help=
"a value between 0 and 1, indicat height of in z direction to print the plot"
)
parser.add_argument("--cmap",
type=str,
default="gray",
choices=[
"gray", "hot", "afmhot", "Spectral", "plasma",
"magma", "hsv", "rainbow", "brg"
])
# ==========================================================
# transfer parameters from the arg subparser to static_run setting
# ==========================================================
args = parser.parse_args()
cube_filepath = args.input
with open(cube_filepath, "r") as fin:
cube = fin.readlines()
ngridx = int(cube[3].split()[0])
ngridy = int(cube[4].split()[0])
ngridz = int(cube[5].split()[0])
# read structure info
natom = abs(int(
cube[2].split()
[0])) # it might be negative, if MO infor are included in cube file
bohr_to_angstrom = 0.529177249
structure = crystal()
structure.cell = []
for i in range(3):
tmp = []
for j in range(3):
tmp.append(
int(cube[i + 3].split()[0]) *
float(cube[i + 3].split()[j + 1]) * bohr_to_angstrom)
structure.cell.append(tmp)
atoms_list = []
for i in range(natom):
atomic_number = int(cube[i + 6].split()[0])
for e in elem:
if elem[e].number == atomic_number:
label = e
atoms_list.append([
label,
float(cube[i + 6].split()[2]) * bohr_to_angstrom,
float(cube[i + 6].split()[3]) * bohr_to_angstrom,
float(cube[i + 6].split()[4]) * bohr_to_angstrom,
])
structure.get_atoms(atoms_list)
# end read structure info
write_structure(structure=structure, filepath=args.output_structure)
# read grid value
tmp_str = "".join(cube[natom + 6:])
data = np.fromstring(tmp_str, sep="\n")
#data = data.reshape(ngridz, ngridy, ngridx)
# grid data in cube is iterated in different compared to *CHG* of vasp
data = data.reshape(ngridx, ngridy, ngridz)
# charge data in cube file is in shape (ngridx, ngridy, ngridz)
# while charge in *CHG* file is in shape (ngzf, ngyf, ngxf)
# they are different!
# -------------------------------------------------------
# matrix image only for z direction
# may not work for triclinic and monoclinic crystal system
# -------------------------------------------------------
zi = int((data.shape[2] - 1) * args.z)
img = data[::, ::,
zi].T #.reshape(ngridy, ngridx) should be transpose here but not reshape
img = (img - img.min()) / (img.max() - img.min()) * 255
# need to do a transform when the cell is not Orthorhombic
# skew the image
a = np.array(structure.cell[0])
b = np.array(structure.cell[1])
cosangle = a.dot(b) / (np.linalg.norm(a) * np.linalg.norm(b))
angle = np.arccos(cosangle) * 180 / np.pi
ax = plt.axes() #plt.figure()
n1 = ngridx
n2 = ngridy
n1_right = n1
n1_left = -(n2 * np.tan((angle - 90) / 180 * np.pi))
#im = ax.imshow(img, cmap="gray", extent=[0, n1, 0, n2], interpolation="none", origin="lower", clip_on=True)
im = ax.imshow(img,
cmap=args.cmap,
extent=[0, n1, 0, n2],
interpolation="none",
origin="lower",
clip_on=True)
#im = plt.imshow(data[i, :, :], cmap="gray")
trans_data = mtransforms.Affine2D().skew_deg(90 - angle, 0) + ax.transData
im.set_transform(trans_data)
# display intended extent of the image
x1, x2, y1, y2 = im.get_extent()
# do not view the line, but it is needed to be plot so the intended image is dispalyed completely
ax.plot([x1, x2, x2, x1, x1], [y1, y1, y2, y2, y1],
"y--",
transform=trans_data,
visible=False)
#ax.set_xlim(n1_left, n1_right)
#ax.set_ylim(0, n2)
plt.colorbar(im)
ax.autoscale()
plt.tight_layout()
plt.savefig(args.output + "-z-%f.matrix-image.png" % args.z)
plt.close()
# -----------------------------------------------------------------------------
# 2D contour plot
#------------------------------------------------------------------------------
nx = np.linspace(0, 1, ngridx)
ny = np.linspace(0, 1, ngridy)
X, Y = np.meshgrid(
nx, ny
) # now this Mesh grid cannot be used directly, we have to calc the real x y for it
for xi in range(len(nx)):
for yi in range(len(ny)):
X[yi, xi] = structure.cell[0][0] * nx[xi] + structure.cell[1][
0] * ny[yi]
Y[yi, xi] = structure.cell[0][1] * nx[xi] + structure.cell[1][
1] * ny[yi]
Z = data[::, ::,
zi].T #.reshape(ngridy, ngridx) should be transpose here but not reshape
Z = (Z - Z.min()) / (Z.max() - Z.min()) * 255
# fill color, three color are divided into three layer(6)
# cmap = plt.cm.hot means using thermostat plot(graduated red yellow)
#cset = plt.contourf(X, Y, Z, levels=args.levels, cmap=plt.cm.hot)
#cset = plt.contourf(X, Y, Z, levels=args.levels, cmap=plt.cm.gray)
cset = plt.contourf(X, Y, Z, levels=args.levels, cmap=args.cmap)
contour = plt.contour(X, Y, Z, levels=[20, 40], colors='k')
plt.colorbar(cset)
plt.autoscale()
plt.tight_layout()
#plt.show()
plt.axis("equal") # set axis equally spaced
plt.xlabel('x')
plt.ylabel('y')
plt.savefig(args.output + ".2d-contour-z-%f.png" % args.z)
plt.close()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-i",
"--input",
type=str,
required=True,
help="input cube file, -i xxx.cube")
parser.add_argument("--output-structure",
type=str,
default="chg",
help="output stucture contained in cube file")
parser.add_argument("-o",
"--output",
type=str,
default="chg",
help="prefix of the output image file name")
parser.add_argument("--levels",
type=int,
default=10,
help="levels of the color map or color bar")
parser.add_argument(
"-z",
"--z",
type=float,
default=1,
help=
"a value between 0 and 1, indicat height of in z direction to print the plot"
)
parser.add_argument("--cmap",
type=str,
default="gray",
choices=[
"gray", "hot", "afmhot", "Spectral", "plasma",
"magma", "hsv", "rainbow", "brg"
])
parser.add_argument(
"--abscissa",
type=str,
nargs="+",
default=["a", "b", "c"],
choices=["a", "b", "c"],
help="choose the direction to do the dimension reduction")
# ==========================================================
# transfer parameters from the arg subparser to static_run setting
# ==========================================================
args = parser.parse_args()
cube_filepath = args.input
with open(cube_filepath, "r") as fin:
cube = fin.readlines()
bohr_to_angstrom = 0.529177249
# read structure info
natom = abs(int(
cube[2].split()
[0])) # it might be negative, if MO infor are included in cube file
structure = crystal()
structure.cell = []
for i in range(3):
tmp = []
for j in range(3):
tmp.append(
int(cube[i + 3].split()[0]) *
float(cube[i + 3].split()[j + 1]) * bohr_to_angstrom)
structure.cell.append(tmp)
atoms_list = []
for i in range(natom):
atomic_number = int(cube[i + 6].split()[0])
for e in elem:
if elem[e].number == atomic_number:
label = e
atoms_list.append([
label,
float(cube[i + 6].split()[2]) * bohr_to_angstrom,
float(cube[i + 6].split()[3]) * bohr_to_angstrom,
float(cube[i + 6].split()[4]) * bohr_to_angstrom,
])
structure.get_atoms(atoms_list)
# end read structure info
write_structure(structure=structure,
filepath=args.output_structure + ".cif")
a = np.linalg.norm(structure.cell[0])
b = np.linalg.norm(structure.cell[1])
c = np.linalg.norm(structure.cell[2])
# assume three cube file have the same ngridx ngridy and ngridz
ngridx = int(cube[3].split()[0])
ngridy = int(cube[4].split()[0])
ngridz = int(cube[5].split()[0])
# read grid value
tmp_str = "".join(cube[natom + 6:])
data = np.fromstring(tmp_str, sep="\n")
# grid data in cube is iterated in different compared to *CHG* of vasp
data = data.reshape(ngridx, ngridy, ngridz)
# charge data in cube file is in shape (ngridx, ngridy, ngridz)
# while charge in *CHG* file is in shape (ngzf, ngyf, ngxf)
# they are different!
# data dimension reduction
# the unit of value is actually not physical now!
#
# cell_volume are in unit of Angstrom^3
cell_volume = np.dot(
np.cross(np.array(structure.cell[0]), np.array(structure.cell[1])),
np.array(structure.cell[2]))
cell_volume_per_unit = cell_volume / (ngridx * ngridy * ngridz)
# value in cube file are \rho(r)_of_electrons in unit of e/Bohr^3
# namely number of electrons each Borh^3
# so we have to convert it to e/Angstrom^23, through divide it by borh_to_angstrom**3
total_electrons = np.sum(data) * cell_volume_per_unit / bohr_to_angstrom**3
#
print("======================================================\n")
print(" Information collected\n")
print("------------------------------------------------------\n")
print("cell volume: %f (A^3)\n" % cell_volume)
print("total electrons: %f\n" % total_electrons)
# data is in unit of e/Bohr^3
# we will build data_red_? to be in unit of e/Anstrom, namely number of electrons per Angstrom
# to do this we have to time the volume density with bohr_to_angstrom^-3
data_red_a = []
data_red_b = []
data_red_c = []
if "c" in args.abscissa:
len_ci = c / ngridz
for ci in range(data.shape[2]):
tmp = 0
for bi in range(data.shape[1]):
tmp += np.sum(data[:, bi, ci])
nelect_ci = tmp * cell_volume_per_unit / bohr_to_angstrom**3
rho_line = nelect_ci / len_ci
data_red_c.append(rho_line)
if "b" in args.abscissa:
len_bi = b / ngridy
for bi in range(data.shape[1]):
tmp = 0
for ai in range(data.shape[0]):
tmp += np.sum(data[ai, bi, :])
nelect_bi = tmp * cell_volume_per_unit / bohr_to_angstrom**3
rho_line = nelect_bi / len_bi
data_red_b.append(rho_line)
if "a" in args.abscissa:
len_ai = a / ngridx
for ai in range(data.shape[0]):
tmp = 0
for ci in range(data.shape[2]):
tmp += np.sum(data[ai, :, ci])
nelect_ai = tmp * cell_volume_per_unit / bohr_to_angstrom**3
rho_line = nelect_ai / len_ai
data_red_a.append(rho_line)
# output the data and make the plot
if "c" in args.abscissa:
with open(args.output + ".1d.c.data", 'w') as fout:
fout.write(
"#c(angstrom) rho(e) (number of electron per Angstrom)\n")
c_coord = np.linspace(0, c, len(data_red_c))
for i in range(len(data_red_c)):
fout.write("%f %f\n" % (c_coord[i], data_red_c[i]))
plt.plot(np.linspace(0, c, len(data_red_c)), data_red_c)
plt.ylabel(r"$\rho (e/\AA)$")
plt.tight_layout()
plt.savefig(args.output + ".1d.c.png")
plt.close()
if "b" in args.abscissa:
with open(args.output + ".1d.b.data", 'w') as fout:
fout.write(
"#b(angstrom) rho(e) (number of electron per Angstrom)\n")
b_coord = np.linspace(0, b, len(data_red_b))
for i in range(len(data_red_b)):
fout.write("%f %f\n" % (b_coord[i], data_red_b[i]))
plt.plot(np.linspace(0, b, len(data_red_b)), data_red_b)
plt.ylabel(r"$\rho (e/\AA)$")
plt.tight_layout()
plt.savefig(args.output + ".1d.b.png")
plt.close()
if "a" in args.abscissa:
with open(args.output + ".1d.a.data", 'w') as fout:
fout.write(
"#a(angstrom) rho(e) (number of electron per Angstrom)\n")
a_coord = np.linspace(0, a, len(data_red_a))
for i in range(len(data_red_a)):
fout.write("%f %f\n" % (a_coord[i], data_red_a[i]))
plt.plot(np.linspace(0, a, len(data_red_a)), data_red_a)
plt.ylabel(r"$\rho (e/\AA)$")
plt.tight_layout()
plt.savefig(args.output + ".1d.a.png")
plt.close()
1].split()[1]),
float(poscars[direct_begin_lines[i] + j +
1].split()[2])
])))
atoms_list.append(
[element_list[j], cartesian[0], cartesian[1], cartesian[2]])
image.get_atoms(atoms_list)
images.append(image)
# write structure files
os.system("mkdir -p %s" % args.directory)
from pymatflow.cmd.structflow import write_structure
if args.format == "cif":
for i in range(len(images)):
write_structure(structure=images[i],
filepath=os.path.join(args.directory,
"%d.cif" % (i + 1)))
elif args.format == "xsd":
for i in range(len(images)):
write_structure(structure=images[i],
filepath=os.path.join(args.directory,
"%d.xsd" % (i + 1)))
elif args.format == "xsf":
for i in range(len(images)):
write_structure(structure=images[i],
filepath=os.path.join(args.directory,
"%d.xsf" % (i + 1)))
else:
pass
# end write structure files
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-i",
"--images",
type=str,
nargs=2,
required=True,
help="the initial and final structure file")
parser.add_argument("-n",
"--nimage",
type=int,
default=None,
required=True,
help="number of inter images")
parser.add_argument("-m",
"--moving-atom",
type=int,
nargs="+",
required=True,
help="specifying the moving atoms, index start from 0")
parser.add_argument("-d",
"--directory",
type=str,
default="./",
help="directory to put the generated images")
parser.add_argument("--frac",
type=int,
default=1,
choices=[0, 1],
help="1(default): use faractional, 0: use cartesian")
# ==============================================================
args = parser.parse_args()
initial = read_structure(args.images[0])
final = read_structure(args.images[1])
inter_images = interpolate(initial=initial,
final=final,
nimage=args.nimage,
moving_atom=args.moving_atom)
os.system("mkdir -p %s" % os.path.join(args.directory, "%.2d" % (0)))
write_structure(structure=initial,
filepath=os.path.join(args.directory, "%.2d/POSCAR" % (0)),
frac=args.frac)
os.system("mkdir -p %s" % os.path.join(args.directory, "%.2d" %
(args.nimage + 1)))
write_structure(structure=final,
filepath=os.path.join(args.directory,
"%.2d/POSCAR" % (args.nimage + 1)),
frac=args.frac)
for i in range(len(inter_images)):
os.system("mkdir -p %s" % os.path.join(args.directory, "%.2d" %
(i + 1)))
write_structure(structure=inter_images[i],
filepath=os.path.join(args.directory,
"%.2d/POSCAR" % (i + 1)),
frac=args.frac)
print("===========================================\n")
print("generate inter images for neb calculation\n")
print("===========================================\n")
print("-------------------------------------------\n")
def main():
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers(
dest="driver",
title="subcommands",
description="choose one and only one subcommand")
subparser = subparsers.add_parser(
"1d", help="dimension reduction of CHG* to one dimensional")
subparser.add_argument("-i",
"--input",
type=str,
required=True,
help="input vasp *CHG* file, -i *CHG* ")
subparser.add_argument("--output-structure",
type=str,
default="chg",
help="output stucture contained in *CHG*")
subparser.add_argument("-o",
"--output",
type=str,
default="chg",
help="prefix of the output image file name")
subparser.add_argument("--levels",
type=int,
default=10,
help="levels of the color map or color bar")
subparser.add_argument(
"-z",
"--z",
type=float,
default=1,
help=
"a value between 0 and 1, indicat height of in z direction to print the plot"
)
subparser.add_argument("--cmap",
type=str,
default="gray",
choices=[
"gray", "hot", "afmhot", "Spectral", "plasma",
"magma", "hsv", "rainbow", "brg"
])
subparser.add_argument(
"--abscissa",
type=str,
nargs="+",
default=["a", "b", "c"],
choices=["a", "b", "c"],
help="choose the direction to do the dimension reduction")
# add
subparser = subparsers.add_parser("add", help="add CHG*")
subparser.add_argument(
"-i",
"--input",
type=str,
nargs="+",
required=True,
help="input vasp *CHG* file, -i CHG1 CHG2 CHG3 ... ")
subparser.add_argument("--output-structure",
type=str,
default="chg",
help="output stucture contained in *CHG*")
subparser.add_argument("-o",
"--output",
type=str,
default="chg-merged",
help="prefix of the output chg file name")
# slice
subparser = subparsers.add_parser("slice",
help="slice CHG* to get 2d DATA")
subparser.add_argument("-i",
"--input",
type=str,
required=True,
help="input vasp *CHG* file, -i *CHG* ")
subparser.add_argument("--output-structure",
type=str,
default="chg-slice",
help="output stucture contained in *CHG*")
subparser.add_argument("-o",
"--output",
type=str,
default="chg",
help="prefix of the output image file name")
subparser.add_argument("--levels",
type=int,
default=10,
help="levels of the color map or color bar")
subparser.add_argument(
"-z",
"--z",
type=float,
default=1,
help=
"a value between 0 and 1, indicat height of in z direction to print the plot"
)
subparser.add_argument("--cmap",
type=str,
default="gray",
choices=[
"gray", "hot", "afmhot", "Spectral", "plasma",
"magma", "hsv", "rainbow", "brg"
])
subparser.add_argument(
"--abscissa",
type=str,
nargs="+",
default=["a", "b", "c"],
choices=["a", "b", "c"],
help="choose the direction to do the dimension reduction")
# ==========================================================
# transfer parameters from the arg subparser to static_run setting
# ==========================================================
args = parser.parse_args()
if args.driver == "1d":
chg_filepath = args.input
with open(chg_filepath, "r") as fin:
chg = fin.readlines()
for j in range(len(chg)):
if len(chg[j].split()) == 0:
first_blank_line = j
break
first_augmentation_line = None
for j in range(len(chg)):
if "augmentation" in chg[j]:
first_augmentation_line = j
break
os.system("mkdir -p /tmp/pymatflow/")
with open("/tmp/pymatflow/POSCAR", "w") as fout:
for j in range(first_blank_line):
fout.write(chg[j])
structure = read_structure("/tmp/pymatflow/POSCAR")
write_structure(structure=structure,
filepath=args.output_structure + ".cif")
a = np.linalg.norm(structure.cell[0])
b = np.linalg.norm(structure.cell[1])
c = np.linalg.norm(structure.cell[2])
# assume three *CHG* have the same ngxf and ngyf ngzf
ngxf = int(chg[first_blank_line + 1].split()[0])
ngyf = int(chg[first_blank_line + 1].split()[1])
ngzf = int(chg[first_blank_line + 1].split()[2])
if first_augmentation_line == None:
tmp_str = "".join(chg[first_blank_line + 2:])
data = np.fromstring(tmp_str, sep="\n").reshape(ngzf, ngyf, ngxf)
#data = data.reshape(ngzf, ngyf, ngxf)
# data dimension reduction
# the unit of value is actually not physical now!
cell_volume = np.dot(
np.cross(np.array(structure.cell[0]), np.array(structure.cell[1])),
np.array(structure.cell[2]))
cell_volume_per_unit = cell_volume / (ngzf * ngyf * ngxf)
# value in Vasp *CHG* are \rho(r)_of_electrons * Volume_of_cell, so we should divide it by cell_volume here and time it with cell_volume_per_unit
# to get the number of electrons per divided unit
total_electrons = np.sum(data) / cell_volume * cell_volume_per_unit
#
print("======================================================\n")
print(" Information collected\n")
print("------------------------------------------------------\n")
print("cell volume: %f (A^3)\n" % cell_volume)
print("total electrons: %f\n" % total_electrons)
# unit of data_red_? is e/Anstrom, namely number of electrons per Angstrom
data_red_a = []
data_red_b = []
data_red_c = []
if "c" in args.abscissa:
factor = cell_volume_per_unit / cell_volume
len_ci = c / ngzf
for ci in range(data.shape[0]):
tmp = 0
for bi in range(data.shape[1]):
tmp += np.sum(data[ci, bi, :])
nelect_ci = tmp * factor
rho_line = nelect_ci / len_ci
data_red_c.append(rho_line)
if "b" in args.abscissa:
factor = cell_volume_per_unit / cell_volume
len_bi = b / ngyf
for bi in range(data.shape[1]):
tmp = 0
for ai in range(data.shape[2]):
tmp += np.sum(data[:, bi, ai])
nelect_bi = tmp * factor
rho_line = nelect_bi / len_bi
data_red_b.append(rho_line)
if "a" in args.abscissa:
factor = cell_volume_per_unit / cell_volume
len_ai = a / ngxf
for ai in range(data.shape[2]):
tmp = 0
for ci in range(data.shape[0]):
tmp += np.sum(data[ci, :, ai])
nelect_ai = tmp * factor
rho_line = nelect_ai / len_ai
data_red_a.append(rho_line)
# output the data and make the plot
if "c" in args.abscissa:
with open(args.output + ".1d.c.data", 'w') as fout:
fout.write(
"#c(angstrom) rho(e) (number of electron per Angstrom)\n")
c_coord = np.linspace(0, c, len(data_red_c))
for i in range(len(data_red_c)):
fout.write("%f %f\n" % (c_coord[i], data_red_c[i]))
plt.plot(np.linspace(0, c, len(data_red_c)), data_red_c)
plt.ylabel(r"$\rho (e/\AA)$")
plt.tight_layout()
plt.savefig(args.output + ".1d.c.png")
plt.close()
if "b" in args.abscissa:
with open(args.output + ".1d.b.data", 'w') as fout:
fout.write(
"#b(angstrom) rho(e) (number of electron per Angstrom)\n")
b_coord = np.linspace(0, b, len(data_red_b))
for i in range(len(data_red_b)):
fout.write("%f %f\n" % (b_coord[i], data_red_b[i]))
plt.plot(np.linspace(0, b, len(data_red_b)), data_red_b)
plt.ylabel(r"$\rho (e/\AA)$")
plt.tight_layout()
plt.savefig(args.output + ".1d.b.png")
plt.close()
if "a" in args.abscissa:
with open(args.output + ".1d.a.data", 'w') as fout:
fout.write(
"#a(angstrom) rho(e) (number of electron per Angstrom)\n")
a_coord = np.linspace(0, a, len(data_red_a))
for i in range(len(data_red_a)):
fout.write("%f %f\n" % (a_coord[i], data_red_a[i]))
plt.plot(np.linspace(0, a, len(data_red_a)), data_red_a)
plt.ylabel(r"$\rho (e/\AA)$")
plt.tight_layout()
plt.savefig(args.output + ".1d.a.png")
plt.close()
elif args.driver == "add":
chgs = []
for item in args.input:
chgs.append(vasp_chg(item))
pass
elif args.driver == "slice":
chg = vasp_chg(args.input)
zi = int((chg.data.shape[0] - 1) * args.z)
img = chg.data[zi, ::, ::]
img = (img - img.min()) / (img.max() - img.min()) * 255
# need to do a transform when the cell is not Orthorhombic
# skew the image
a = np.array(chg.structure.cell[0])
b = np.array(chg.structure.cell[1])
cosangle = a.dot(b) / (np.linalg.norm(a) * np.linalg.norm(b))
angle = np.arccos(cosangle) * 180 / np.pi
ax = plt.axes()
n1 = chg.ngxf
n2 = chg.ngyf
n1_right = n1
n1_left = -(n2 * np.tan((angle - 90) / 180 * np.pi))
#im = ax.imshow(img, cmap="gray", extent=[n1_left, n1_right, 0, n2], interpolation="none", origin="lower", clip_on=True)
im = ax.imshow(img,
cmap="gray",
extent=[0, n1, 0, n2],
interpolation="none",
origin="lower",
clip_on=True)
#im = plt.imshow(data[i, :, :], cmap="gray")
trans_data = mtransforms.Affine2D().skew_deg(90 - angle,
0) + ax.transData
im.set_transform(trans_data)
# display intended extent of the image
x1, x2, y1, y2 = im.get_extent()
# do not view the line, but it is needed to be plot so the intended image is dispalyed completely
ax.plot([x1, x2, x2, x1, x1], [y1, y1, y2, y2, y1],
"y--",
transform=trans_data,
visible=False)
#ax.set_xlim(n1_left, n1_right)
#ax.set_ylim(0, n2)
plt.colorbar(im)
ax.autoscale()
plt.tight_layout()
plt.savefig(args.output + "-z-%f.grayscale.png" % args.z)
plt.close()
# -----------------------------------------------------------------------------
# 2D contour plot
#------------------------------------------------------------------------------
nx = np.linspace(0, 1, chg.ngxf)
ny = np.linspace(0, 1, chg.ngyf)
X, Y = np.meshgrid(
nx, ny
) # now this mesh grid cannot be used directly, we have to calc the real x y for it
for xi in range(len(nx)):
for yi in range(len(ny)):
X[yi, xi] = chg.structure.cell[0][0] * nx[
xi] + chg.structure.cell[1][0] * ny[yi]
Y[yi, xi] = chg.structure.cell[0][1] * nx[
xi] + chg.structure.cell[1][1] * ny[yi]
density_z = chg.data[zi, :, :] / chg.cell_volume
# export data
with open(args.output + '.slice-z.%f.data' % args.z, "w") as fout:
fout.write("# x y val(e/Angstrom^3)\n")
for xi in range(len(nx)):
for yi in range(len(ny)):
fout.write("%f %f %f\n" %
(X[xi, yi], Y[xi, yi], density_z[xi, yi]))
Z = chg.data[zi, :, :]
Z = (Z - Z.min()) / (Z.max() - Z.min()) * 255
# fill color, three color are divided into three layer(6)
# cmap = plt.cm.hot means using thermostat plot(graduated red yellow)
#cset = plt.contourf(x, y, z, levels=args.levels, cmap=plt.cm.hot)
#cset = plt.contourf(x, y, z, levels=args.levels, cmap=plt.cm.gray)
cset = plt.contourf(X, Y, Z, levels=args.levels, cmap=args.cmap)
contour = plt.contour(X, Y, Z, levels=[20, 40], colors='k')
plt.colorbar(cset)
plt.autoscale()
plt.tight_layout()
plt.axis("equal") # set axis equally spaced
#plt.show()
plt.xlabel('x')
plt.ylabel('y')
plt.savefig(args.output + ".2d-slice-z-%f.png" % args.z)
plt.close()