def main():
parser = argparse.ArgumentParser()
parser.add_argument("-i",
"--input",
type=str,
required=True,
help="input structure file")
args = parser.parse_args()
# if no argument passed to matflow
if len(sys.argv) == 1:
# display help message when no args provided
parser.print_help()
sys.exit(1)
a = read_structure(filepath=args.input)
for i in range(1, 10):
supercell = a.build_supercell([i, i, i])
new_structure = crystal()
new_structure.get_cell_atoms(cell=supercell["cell"],
atoms=supercell["atoms"])
# calculate the madelung constant
# get center atom
all_x = [atoms.x for atom in new_structure.atoms]
all_y = [atoms.y for atom in new_structure.atoms]
all_z = [atoms.z for atom in new_structure.atoms]
def read_chg(self, chg_filepath):
with open(chg_filepath, "r") as fin:
self.lines = fin.readlines()
for j in range(len(self.lines)):
if len(self.lines[j].split()) == 0:
first_blank_line = j
break
first_augmentation_line = None
for j in range(len(self.lines)):
if "augmentation" in self.lines[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(self.lines[j])
self.structure = read_structure("/tmp/pymatflow/POSCAR")
# assume three *CHG* have the same ngxf and ngyf ngzf
self.ngxf = int(self.lines[first_blank_line + 1].split()[0])
self.ngyf = int(self.lines[first_blank_line + 1].split()[1])
self.ngzf = int(self.lines[first_blank_line + 1].split()[2])
if first_augmentation_line == None:
tmp_str = "".join(self.lines[first_blank_line + 2:])
self.data = np.fromstring(tmp_str, sep="\n").reshape(
self.ngzf, self.ngyf, self.ngxf)
# the unit of value is actually not physical now!
self.cell_volume = np.dot(
np.cross(np.array(self.structure.cell[0]),
np.array(self.structure.cell[1])),
np.array(self.structure.cell[2]))
self.cell_volume_per_unit = self.cell_volume / (self.ngzf * self.ngyf *
self.ngxf)
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 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"))
print("%f %f %f %f %f %f\n" %
(c_elastic_n_m2[4, 0], c_elastic_n_m2[4, 1], c_elastic_n_m2[4, 2],
c_elastic_n_m2[4, 3], c_elastic_n_m2[4, 4], c_elastic_n_m2[4, 5]))
print("%f %f %f %f %f %f\n" %
(c_elastic_n_m2[5, 0], c_elastic_n_m2[5, 1], c_elastic_n_m2[5, 2],
c_elastic_n_m2[5, 3], c_elastic_n_m2[5, 4], c_elastic_n_m2[5, 5]))
print(
"Piezoelectric stress tensor of 2D materials (z as surface direction)\n"
)
print("Piezoelectric stess tensor in 2D (pC/m)\n")
print("e_ij(2D)\n")
print("e11 e12 e16\n")
print("e21 e22 e26\n")
print("e31 e32 e36\n")
structure = read_structure(filepath=args.poscar)
c = np.linalg.norm(np.array(structure.cell[2]))
e_total_2d_pc_m = []
for i in [1, 2, 3]:
row = []
for j in [1, 2, 6]:
row.append(e_total[i - 1, j - 1])
e_total_2d_pc_m.append(row)
e_total_2d_pc_m = np.array(e_total_2d_pc_m) * 1.0e+12 * c * 1.0e-10 # pC/m
print(
"%f %f %f\n" %
(e_total_2d_pc_m[0, 0], e_total_2d_pc_m[0, 1], e_total_2d_pc_m[0, 2]))
print(
"%f %f %f\n" %
(e_total_2d_pc_m[1, 0], e_total_2d_pc_m[1, 1], e_total_2d_pc_m[1, 2]))
print(
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input", type=str, required=True,
help="input structure file")
parser.add_argument("-o", "--output", type=str, default="./contour",
help="prefix of the output image file name")
parser.add_argument("--atoms", nargs='+', type=int,
help="list of fixed atoms, index start from 1, default is all atoms")
parser.add_argument("--around-z", type=float, nargs=3,
help="select atoms around specified z in Angstrom with tolerance, like this --around-z 10 -0.5 0.5")
parser.add_argument("--dgrid3d", type=int, nargs=3,
default=[100, 100, 4],
help="used by gnuplot to set dgrid3d int, int, int")
parser.add_argument("--levels", type=int, default=10,
help="levels of the color map or color bar")
parser.add_argument("--ngridx", type=int, default=200,
help="ngridx to plot the contour for irregularly spaced data")
parser.add_argument("--ngridy", type=int, default=200,
help="ngridy to plot the contour for irregularly spaced data")
# ==========================================================
# transfer parameters from the arg subparser to static_run setting
# ==========================================================
args = parser.parse_args()
structure = read_structure(args.input)
if args.atoms == None and args.around_z == None:
atoms_index_from_1 = range(1, len(structure.atoms)+1)
elif args.atoms == None and args.around_z != None:
atoms_index_from_1 = []
for i in range(len(structure.atoms)):
if structure.atoms[i].z > (args.around_z[0] + args.around_z[1]) and structure.atoms[i].z < (args.around_z[0] + args.around_z[2]):
atoms_index_from_1.append(i+1)
elif args.atoms != None and args.around_z == None:
atoms_index_from_1 = args.atoms
elif args.atoms != None and args.around_z != None:
print("======================================================\n")
print(" Warning\n")
print("--atoms and --around-z can not be used at the same time.")
sys.exit(1)
data = []
for i in atoms_index_from_1:
#X.append(structure.atoms[i-1].x)
#Y.append(structure.atoms[i-1].y)
#Z.append(structure.atoms[i-1].z)
data.append([structure.atoms[i-1].x, structure.atoms[i-1].y, structure.atoms[i-1].z])
with open(args.output+".data", "w") as fout:
fout.write("# format: x y z\n")
for d in data:
fout.write("%f\t%f\t%f\n" % (d[0], d[1], d[2]))
with open("plot.gnuplot", "w") as fout:
fout.write("set term gif\n")
fout.write("set dgrid3d %d, %d, %d\n" % (args.dgrid3d[0], args.dgrid3d[1], args.dgrid3d[2]))
fout.write("set contour base\n")
fout.write("unset key\n")
fout.write("set autoscale\n")
fout.write("set xlabel \"X\"\n")
fout.write("set ylabel \"Y\"\n")
fout.write("set zlabel \"Z\"\n")
#fout.write("set output \"%s\"\n" % (args.output+".gif"))
#fout.write("set multiplot layout 1, 2\n")
#fout.write("set origin 0, 0\n")
fout.write("\n\n")
fout.write("set surface\n")
fout.write("set pm3d hidden3d\n")
fout.write("set output \"%s\"\n" % (args.output+".surf"+".gif"))
fout.write("splot '%s' u 1:2:3 w pm3d\n" % (args.output+".data"))
#fout.write("set origin 0.5, 0\n")
fout.write("\n\n")
fout.write("set pm3d map\n")
fout.write("set surface\n")
fout.write("set output \"%s\"\n" % (args.output+".map"+".gif"))
fout.write("splot '%s' u 1:2:3 w pm3d\n" % (args.output+".data"))
os.system("gnuplot plot.gnuplot")
# ----------------
# 2D contour plot
#-----------------
# Contour plot of irregularly spaced data
# data is irregularly spaced data
# we can refer to https://matplotlib.org/3.2.1/gallery/images_contours_and_fields/irregulardatagrid.html
# to see how to build contour plot of such kind of data
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.tri as tri
data_np = np.array(data)
# 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=10, cmap=plt.cm.hot)
#contour = plt.contour(X, Y, Z, colors='k')
#plt.colorbar(cset)
#plt.autoscale()
#plt.tight_layout()
#plt.show()
#plt.xlabel('x')
#plt.ylabel('y')
#plt.savefig(args.output+".2d-contour.png")
#plt.close()
ngridx = args.ngridx
ngridy = args.ngridy
x = data_np[:, 0]
y = data_np[:, 1]
z = data_np[:, 2]
#fig, (ax1, ax2) = plt.subplots(nrows=2)
fig = plt.figure()
# -----------------------
# Interpolation on a grid
# -----------------------
# A contour plot of irregularly spaced data coordinates
# via interpolation on a grid.
# Create grid values first.
xi = np.linspace(np.min(x), np.max(x), ngridx)
yi = np.linspace(np.min(y), np.max(y), ngridy)
# Linearly interpolate the data (x, y) on a grid defined by (xi, yi).
triang = tri.Triangulation(x, y)
interpolator = tri.LinearTriInterpolator(triang, z)
Xi, Yi = np.meshgrid(xi, yi)
zi = interpolator(Xi, Yi)
# Note that scipy.interpolate provides means to interpolate data on a grid
# as well. The following would be an alternative to the four lines above:
#from scipy.interpolate import griddata
#zi = griddata((x, y), z, (xi[None,:], yi[:,None]), method='linear')
#ax1.contour(xi, yi, zi, levels=args.levels, linewidths=0.5, colors='k')
#cntr1 = ax1.contourf(xi, yi, zi, levels=args.levels, cmap="RdBu_r")
#fig.colorbar(cntr1, ax=ax1)
#ax1.plot(x, y, 'ko', ms=3)
#ax1.set(xlim=(-2, 2), ylim=(-2, 2))
#ax1.set_title('grid and contour (%d points, %d grid points)' %
#(npts, ngridx * ngridy))
#ax1.axis('equal') # set x y axis equal spaced
# ----------
# Tricontour
# ----------
# Directly supply the unordered, irregularly spaced coordinates
# to tricontour.
#ax2.tricontour(x, y, z, levels=args.levels, linewidths=0.5, colors='k')
plt.tricontour(x, y, z, levels=args.levels, linewidths=0.5, colors='k')
#cntr2 = ax2.tricontourf(x, y, z, levels=args.levels, cmap="RdBu_r")
cntr2 = plt.tricontourf(x, y, z, levels=args.levels, cmap="RdBu_r")
#fig.colorbar(cntr2, ax=ax2)
fig.colorbar(cntr2)
#ax2.plot(x, y, 'ko', ms=3)
#plt.plot(x, y, "ko", ms=3)
#ax2.set(xlim=(-2, 2), ylim=(-2, 2))
#ax2.set_title('tricontour (%d points)' % npts)
#ax2.axis('equal') # set x y axis equal spaced
plt.axis("equal")
#plt.subplots_adjust(hspace=0.5)
plt.autoscale()
plt.show()
fig.savefig(args.output+".2d-contour.png")
# =======================
# =======================
plt.close()
ax = plt.axes(projection='3d')
cset = ax.plot_trisurf(x, y, z, cmap='rainbow')
plt.colorbar(cset)
plt.autoscale()
plt.tight_layout()
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
plt.savefig(args.output+".3d-trisurf.png")
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"-i",
"--input",
type=str,
nargs=2,
required=True,
help=
"input initial and final structure file, eg. -i initial.cif final.cif")
parser.add_argument("-o",
"--output",
type=str,
default="./contour-diff",
help="prefix of the output image file name")
parser.add_argument(
"--atoms",
nargs='+',
type=int,
help="list of fixed atoms, index start from 1, default is all atoms")
parser.add_argument(
"--around-z",
type=float,
nargs=3,
help=
"select atoms around specified z in Angstrom with tolerance, like this --around-z 10 -0.5 0.5"
)
parser.add_argument("--levels",
type=int,
default=10,
help="levels of the color map or color bar")
parser.add_argument(
"--ngridx",
type=int,
default=200,
help="ngridx to plot the contour for irregularly spaced data")
parser.add_argument(
"--ngridy",
type=int,
default=200,
help="ngridy to plot the contour for irregularly spaced data")
parser.add_argument(
"--diff",
type=str,
default="z",
choices=["x", "y", "z", "xyz"],
help=
"choose to plot the diff of x or y or z, or xyz which means displacement"
)
# ==========================================================
# transfer parameters from the arg subparser to static_run setting
# ==========================================================
args = parser.parse_args()
initial_structure = read_structure(args.input[0])
final_structure = read_structure(args.input[1])
if args.atoms == None and args.around_z == None:
atoms_index_from_1 = range(1, len(initial_structure.atoms) + 1)
elif args.atoms == None and args.around_z != None:
atoms_index_from_1 = []
for i in range(len(initial_structure.atoms)):
if initial_structure.atoms[i].z > (
args.around_z[0] +
args.around_z[1]) and initial_structure.atoms[i].z < (
args.around_z[0] + args.around_z[2]):
atoms_index_from_1.append(i + 1)
elif args.atoms != None and args.around_z == None:
atoms_index_from_1 = args.atoms
elif args.atoms != None and args.around_z != None:
print("======================================================\n")
print(" Warning\n")
print("--atoms and --around-z can not be used at the same time.")
sys.exit(1)
initial_data = []
for i in atoms_index_from_1:
#X.append(structure.atoms[i-1].x)
#Y.append(structure.atoms[i-1].y)
#Z.append(structure.atoms[i-1].z)
initial_data.append([
initial_structure.atoms[i - 1].x, initial_structure.atoms[i - 1].y,
initial_structure.atoms[i - 1].z
])
final_data = []
for i in atoms_index_from_1:
#X.append(structure.atoms[i-1].x)
#Y.append(structure.atoms[i-1].y)
#Z.append(structure.atoms[i-1].z)
final_data.append([
final_structure.atoms[i - 1].x, final_structure.atoms[i - 1].y,
final_structure.atoms[i - 1].z
])
# data: [x, y, z, diff(x|y|z|xyz)]
data = copy.deepcopy(final_data)
for i in range(len(final_data)):
if args.diff == "x":
diff = final_data[i][0] - initial_data[i][0]
elif args.diff == "y":
diff = final_data[i][1] - initial_data[i][1]
elif args.diff == "z":
diff = final_data[i][2] - initial_data[i][2]
elif args.diff == "xyz":
diff = np.sqrt((final_data[i][0] - initial_data[i][0])**2 +
(final_data[i][1] - initial_data[i][1])**2 +
(final_data[i][2] - initial_data[i][2])**2)
else:
pass
data[i].append(diff)
with open(args.output + ".data", "w") as fout:
fout.write("# format: x y z diff(%s)\n" % (args.diff))
for d in data:
fout.write("%f\t%f\t%f\t%f\n" % (d[0], d[1], d[2], d[3]))
# ----------------
# 2D contour plot
#-----------------
# Contour plot of irregularly spaced data
# data is irregularly spaced data
# we can refer to https://matplotlib.org/3.2.1/gallery/images_contours_and_fields/irregulardatagrid.html
# to see how to build contour plot of such kind of data
data_np = np.array(data)
# 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=10, cmap=plt.cm.hot)
#contour = plt.contour(X, Y, Z, colors='k')
#plt.colorbar(cset)
#plt.autoscale()
#plt.tight_layout()
#plt.show()
#plt.xlabel('x')
#plt.ylabel('y')
#plt.savefig(args.output+".2d-contour.png")
#plt.close()
ngridx = args.ngridx
ngridy = args.ngridy
x = data_np[:, 0] # x
y = data_np[:, 1] # y
z = data_np[:, 3] # diff(x|y|z)
#fig, (ax1, ax2) = plt.subplots(nrows=2)
fig = plt.figure()
# -----------------------
# Interpolation on a grid
# -----------------------
# A contour plot of irregularly spaced data coordinates
# via interpolation on a grid.
# Create grid values first.
xi = np.linspace(np.min(x), np.max(x), ngridx)
yi = np.linspace(np.min(y), np.max(y), ngridy)
# Linearly interpolate the data (x, y) on a grid defined by (xi, yi).
triang = tri.Triangulation(x, y)
interpolator = tri.LinearTriInterpolator(triang, z)
Xi, Yi = np.meshgrid(xi, yi)
zi = interpolator(Xi, Yi)
# Note that scipy.interpolate provides means to interpolate data on a grid
# as well. The following would be an alternative to the four lines above:
#from scipy.interpolate import griddata
#zi = griddata((x, y), z, (xi[None,:], yi[:,None]), method='linear')
#ax1.contour(xi, yi, zi, levels=args.levels, linewidths=0.5, colors='k')
#cntr1 = ax1.contourf(xi, yi, zi, levels=args.levels, cmap="RdBu_r")
#fig.colorbar(cntr1, ax=ax1)
#ax1.plot(x, y, 'ko', ms=3)
#ax1.set(xlim=(-2, 2), ylim=(-2, 2))
#ax1.set_title('grid and contour (%d points, %d grid points)' %
#(npts, ngridx * ngridy))
#ax1.axis('equal') # set x y axis equal spaced
# ----------
# Tricontour
# ----------
# Directly supply the unordered, irregularly spaced coordinates
# to tricontour.
#ax2.tricontour(x, y, z, levels=args.levels, linewidths=0.5, colors='k')
plt.tricontour(x, y, z, levels=args.levels, linewidths=0.5, colors='k')
#cntr2 = ax2.tricontourf(x, y, z, levels=args.levels, cmap="RdBu_r")
cntr2 = plt.tricontourf(x, y, z, levels=args.levels, cmap="RdBu_r")
#fig.colorbar(cntr2, ax=ax2)
fig.colorbar(cntr2)
#ax2.plot(x, y, 'ko', ms=3)
#plt.plot(x, y, "ko", ms=3)
#ax2.set(xlim=(-2, 2), ylim=(-2, 2))
#ax2.set_title('tricontour (%d points)' % npts)
#ax2.axis('equal') # set x y axis equal spaced
plt.axis("equal")
#plt.subplots_adjust(hspace=0.5)
plt.autoscale()
plt.show()
fig.savefig(args.output + ".2d-contour.png")
]
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("-i",
"--input",
type=str,
help="the input structure file")
parser.add_argument("--symprec",
type=float,
default=1.0e-7,
help="symmetry precision")
args = parser.parse_args()
a = read_structure(filepath=args.input)
print("===========================================\n")
print("calculated using spglib\n")
print("===========================================\n")
print("spacegroup is : %s\n" %
get_spacegroup_and_kpath(structure=a, symprec=args.symprec)[0])
print("Warning:\n")
print("the result is not guaranteed to be correct\n")
print("-------------------------------------------\n")
print("suggested k path calculated using seekpath:\n")
print(get_spacegroup_and_kpath(structure=a, symprec=args.symprec)[1])
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 main():
parser = argparse.ArgumentParser()
parser.add_argument("-i",
"--input",
type=str,
required=True,
help="input structure file")
parser.add_argument("-o",
"--output",
type=str,
default="kpath-from-seekpath.txt",
help="the output kpoitns file")
parser.add_argument(
"--join",
type=int,
default=15,
help=
"default number of kpoint to connect the connected high symmetry k point"
)
# ===============================================================================
args = parser.parse_args()
structure = read_structure(filepath=args.input)
lattice = structure.cell
positions = []
numbers = []
a = np.sqrt(structure.cell[0][0]**2 + structure.cell[0][1]**2 +
structure.cell[0][2]**2)
b = np.sqrt(structure.cell[1][0]**2 + structure.cell[1][1]**2 +
structure.cell[1][2]**2)
c = np.sqrt(structure.cell[2][0]**2 + structure.cell[2][1]**2 +
structure.cell[2][2]**2)
frac_coord = structure.get_fractional()
for i in range(len(frac_coord)):
positions.append(
[frac_coord[i][1], frac_coord[i][2],
frac_coord[i][3]]) # must be fractional coordinates
numbers.append(element[frac_coord[i][0]].number)
cell = (lattice, positions, numbers)
kpoints_seekpath = seekpath.get_path(cell)
kpath = []
# [[kx, ky, kz, label, connect_indicator], ...] like [[0.0, 0.0, 0.0, 'GAMMA', 15], ...]
# if connect_indicator in a kpoint is an integer, then it will connect to the following point
# through the number of kpoints defined by connect_indicator.
# if connect_indicator in a kpoint is '|', then it will not connect to the following point,
kpath.append([
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"][0][0]]
[0]),
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"][0][0]]
[1]),
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"][0][0]]
[2]),
kpoints_seekpath["path"][0][0],
args.join,
])
kpath.append([
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"][0][1]]
[0]),
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"][0][1]]
[1]),
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"][0][1]]
[2]),
kpoints_seekpath["path"][0][1],
args.
join, # first set to args.join here, and it will be reset by the following path by judge whether it is connected
])
for i in range(1, len(kpoints_seekpath["path"])):
if kpoints_seekpath["path"][i][0] == kpoints_seekpath["path"][i -
1][1]:
kpath[-1][4] = args.join
kpath.append([
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"]
[i][1]][0]),
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"]
[i][1]][1]),
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"]
[i][1]][2]),
kpoints_seekpath["path"][i][1],
args.join,
])
else:
kpath[-1][4] = "|"
kpath.append([
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"]
[i][0]][0]),
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"]
[i][0]][1]),
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"]
[i][0]][2]),
kpoints_seekpath["path"][i][0],
args.join,
])
kpath.append([
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"]
[i][1]][0]),
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"]
[i][1]][1]),
float(kpoints_seekpath["point_coords"][kpoints_seekpath["path"]
[i][1]][2]),
kpoints_seekpath["path"][i][1],
args.join,
])
#
with open(args.output, 'w') as fout:
fout.write("%d\n" % len(kpath))
for kpoint in kpath:
fout.write(
"%f %f %f #%s %s\n" %
(kpoint[0], kpoint[1], kpoint[2], kpoint[3], str(kpoint[4])))
#
print("===========================================\n")
print("calculated using seekpath\n")
print("===========================================\n")
print("Warning:\n")
print("the result is not guaranteed to be correct\n")
print("-------------------------------------------\n")
print("suggested k path calculated using seekpath:\n")
print(kpoints_seekpath)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input", type=str, nargs="+", required=True,
help="input structure file and vasprun.xml, you can specify two vasprun.xml(one for scf and one for nscf). eg. -i structure.cif vasprun.xml.scf vasprun.xml.nscf")
parser.add_argument("-o", "--output", type=str, default="./contour-ldos",
help="prefix of the output image file name")
parser.add_argument("--atoms", nargs='+', type=int,
help="list of fixed atoms, index start from 1, default is all atoms")
parser.add_argument("--around-z", type=float, nargs=3,
help="select atoms around specified z in Angstrom with tolerance, like this --around-z 10 -0.5 0.5")
parser.add_argument("--levels", type=int, default=10,
help="levels of the color map or color bar")
parser.add_argument("--ngridx", type=int, default=200,
help="ngridx to plot the contour for irregularly spaced data")
parser.add_argument("--ngridy", type=int, default=200,
help="ngridy to plot the contour for irregularly spaced data")
# ==========================================================
# transfer parameters from the arg subparser to static_run setting
# ==========================================================
args = parser.parse_args()
structure = read_structure(args.input[0])
pdos = post_pdos()
if len(args.input) == 2:
efermi_from = "nscf"
pdos.get_vasprun(args.input[1])
pdos.get_efermi(vasprun=args.input[1])
elif len(args.input) == 3:
efermi_from = "scf"
pdos.get_vasprun(args.input[2])
pdos.get_efermi(vasprun=args.input[1])
#pdos.export(directory=args.directory, engine=args.engine, plotrange=args.plotrange)
if args.atoms == None and args.around_z == None:
atoms_index_from_1 = range(1, len(structure.atoms)+1)
elif args.atoms == None and args.around_z != None:
atoms_index_from_1 = []
for i in range(len(structure.atoms)):
if structure.atoms[i].z > (args.around_z[0] + args.around_z[1]) and structure.atoms[i].z < (args.around_z[0] + args.around_z[2]):
atoms_index_from_1.append(i+1)
elif args.atoms != None and args.around_z == None:
atoms_index_from_1 = args.atoms
elif args.atoms != None and args.around_z != None:
print("======================================================\n")
print(" Warning\n")
print("--atoms and --around-z can not be used at the same time.")
sys.exit(1)
# data: [x, y, z, ldos]
data = []
for i in atoms_index_from_1:
ldos = 0.0
#if pdos.magnetic_status == "non-soc-ispin-1":
# pdos.data[i-1]
#elif pdos.magnetic_status == "non-soc-ispin-2":
# pdos.data[i-1]
#elif pdos.magnetic_status == "soc-ispin-1" or pdos.magnetic_status == "soc-ispin-2":
# pdos.data[i-1]
for item in pdos.data[i-1]:
if item == "ion":
continue
for itm in pdos.data[i-1][item]:
if itm == "energy":
continue
ldos += sum(pdos.data[i-1][item][itm])
data.append([structure.atoms[i-1].x, structure.atoms[i-1].y, structure.atoms[i-1].z, ldos])
with open(args.output+".data", "w") as fout:
fout.write("# format: x y z ldos\n")
for d in data:
fout.write("%f\t%f\t%f\t%f\n" % (d[0], d[1], d[2], d[3]))
# ----------------
# 2D contour plot
#-----------------
# Contour plot of irregularly spaced data
# data is irregularly spaced data
# we can refer to https://matplotlib.org/3.2.1/gallery/images_contours_and_fields/irregulardatagrid.html
# to see how to build contour plot of such kind of data
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.tri as tri
data_np = np.array(data)
# 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=10, cmap=plt.cm.hot)
#contour = plt.contour(X, Y, Z, colors='k')
#plt.colorbar(cset)
#plt.autoscale()
#plt.tight_layout()
#plt.show()
#plt.xlabel('x')
#plt.ylabel('y')
#plt.savefig(args.output+".2d-contour.png")
#plt.close()
ngridx = args.ngridx
ngridy = args.ngridy
x = data_np[:, 0] # x
y = data_np[:, 1] # y
z = data_np[:, 3] # ldos
#fig, (ax1, ax2) = plt.subplots(nrows=2)
fig = plt.figure()
# -----------------------
# Interpolation on a grid
# -----------------------
# A contour plot of irregularly spaced data coordinates
# via interpolation on a grid.
# Create grid values first.
xi = np.linspace(np.min(x), np.max(x), ngridx)
yi = np.linspace(np.min(y), np.max(y), ngridy)
# Linearly interpolate the data (x, y) on a grid defined by (xi, yi).
triang = tri.Triangulation(x, y)
interpolator = tri.LinearTriInterpolator(triang, z)
Xi, Yi = np.meshgrid(xi, yi)
zi = interpolator(Xi, Yi)
# Note that scipy.interpolate provides means to interpolate data on a grid
# as well. The following would be an alternative to the four lines above:
#from scipy.interpolate import griddata
#zi = griddata((x, y), z, (xi[None,:], yi[:,None]), method='linear')
#ax1.contour(xi, yi, zi, levels=args.levels, linewidths=0.5, colors='k')
#cntr1 = ax1.contourf(xi, yi, zi, levels=args.levels, cmap="RdBu_r")
#fig.colorbar(cntr1, ax=ax1)
#ax1.plot(x, y, 'ko', ms=3)
#ax1.set(xlim=(-2, 2), ylim=(-2, 2))
#ax1.set_title('grid and contour (%d points, %d grid points)' %
#(npts, ngridx * ngridy))
#ax1.axis('equal') # set x y axis equal spaced
# ----------
# Tricontour
# ----------
# Directly supply the unordered, irregularly spaced coordinates
# to tricontour.
#ax2.tricontour(x, y, z, levels=args.levels, linewidths=0.5, colors='k')
plt.tricontour(x, y, z, levels=args.levels, linewidths=0.5, colors='k')
#cntr2 = ax2.tricontourf(x, y, z, levels=args.levels, cmap="RdBu_r")
cntr2 = plt.tricontourf(x, y, z, levels=args.levels, cmap="RdBu_r")
#fig.colorbar(cntr2, ax=ax2)
fig.colorbar(cntr2)
#ax2.plot(x, y, 'ko', ms=3)
#plt.plot(x, y, "ko", ms=3)
#ax2.set(xlim=(-2, 2), ylim=(-2, 2))
#ax2.set_title('tricontour (%d points)' % npts)
#ax2.axis('equal') # set x y axis equal spaced
plt.axis("equal")
#plt.subplots_adjust(hspace=0.5)
plt.autoscale()
plt.show()
fig.savefig(args.output+".2d-contour.png")
plt.close()
ax = plt.axes(projection='3d')
cset = ax.plot_trisurf(x, y, z, cmap='rainbow')
plt.colorbar(cset)
plt.autoscale()
plt.tight_layout()
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('ldos')
plt.savefig(args.output+".3d-trisurf.png")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("-i",
"--input",
type=str,
required=True,
help="input structure file")
parser.add_argument("-o",
"--output",
type=str,
default="output.xsd",
help="prefix of the output image file name")
parser.add_argument(
"--levels",
type=int,
default=3,
help="number of levels in color map or alpha channel, etc.")
# ==========================================================
# transfer parameters from the arg subparser to static_run setting
# ==========================================================
args = parser.parse_args()
# read xsd file
xsd = parse(args.input)
structure = read_structure(args.input)
# ID of Atom3D in xsd file start from 4
imap = xsd.getroot().find("AtomisticTreeRoot").find("SymmetrySystem").find(
"MappingSet").find("MappingFamily").find("IdentityMapping")
atoms = imap.findall("Atom3d")
bonds = imap.findall("Bond")
natom = len(atoms)
nbond = len(bonds)
bonds_length = []
for bond in bonds:
connects = []
connects.append(int(bond.attrib["Connects"].split(",")[0]))
connects.append(int(bond.attrib["Connects"].split(",")[1]))
connect_0_in_atoms = False
connect_1_in_atoms = False
for i in range(len(atoms)):
if int(atoms[i].attrib["ID"]) == connects[0]:
connect_0_in_atoms = True
connect_0_atom_i = i
break
for i in range(len(atoms)):
if int(atoms[i].attrib["ID"]) == connects[1]:
connect_1_in_atoms = True
connect_1_atom_i = i
break
if connect_0_in_atoms == True and connect_1_in_atoms == True:
x0 = float(atoms[connect_0_atom_i].attrib["XYZ"].split(",")[0])
y0 = float(atoms[connect_0_atom_i].attrib["XYZ"].split(",")[1])
z0 = float(atoms[connect_0_atom_i].attrib["XYZ"].split(",")[2])
x1 = float(atoms[connect_1_atom_i].attrib["XYZ"].split(",")[0])
y1 = float(atoms[connect_1_atom_i].attrib["XYZ"].split(",")[1])
z1 = float(atoms[connect_1_atom_i].attrib["XYZ"].split(",")[2])
# transform from fractional coord to cartesian
latcell = np.array(structure.cell)
convmat = latcell.T
x0, y0, z0 = list(convmat.dot(np.array([x0, y0, z0])))
x1, y1, z1 = list(convmat.dot(np.array([x1, y1, z1])))
length = np.sqrt((x1 - x0)**2 + (y1 - y0)**2 + (z1 - z0)**2)
bonds_length.append(length)
else:
bonds_length.append(None)
#
bonds_length_norm = []
tmp = copy.deepcopy(bonds_length)
while None in tmp:
tmp.remove(None)
#print(tmp)
bond_max = max(tmp)
bond_min = min(tmp)
print(bond_max, bond_min)
for i in range(len(bonds_length)):
if bonds_length[i] == None:
bonds_length_norm.append(None)
else:
bonds_length_norm.append(
(bonds_length[i] - bond_min) / (bond_max - bond_min))
print(bonds_length_norm)
# ------------------------------------------------------------------------
# now bonds length are alreay handled
# we should use different scheme to set the color of them
# ------------------------------------------------------------------------
#Scheme I
# set vriridis cmap color
from matplotlib import cm
#cmap = cm.get_cmap("viridis", args.levels)
#cmap = cm.get_cmap("Blues", args.levels)
#cmap = cm.get_cmap("Greens", args.levels)
#cmap = cm.get_cmap("Reds", args.levels)
#cmap = cm.get_cmap("gist_heat", args.levels)
#cmap = cm.get_cmap("hsv", args.levels)
cmap = cm.get_cmap("rainbow", args.levels)
#cmap = cm.get_cmap("winter", args.levels)
for i in range(len(bonds)):
if bonds_length_norm[i] == None:
pass
else:
# there are two ways to get the color for one value, if you use n color in cmap
# if you pass a integer (0<=value<=(n-1)) to cmap(), you will get the color of it
# or you can pass a float within([0.0, 1.0]) to cmap(), you will get the color of it
# as bonds_length_norm is normalized to (0.0-1.0), we could directly pass it to cmap() to get the right color
color = cmap(bonds_length_norm[i])
# cmap return RGB in range of (0, 1), we have to multiply it with 255 to get a value between range(0, 255)
bonds[i].set(
"Color", "%f, %f, %f, %f" %
(color[0] * 255, color[1] * 255, color[2] * 255, color[3]))
# plot the cmap
# the plot is with no use, what we want is the corresponding color bar.
tmp_norm = copy.deepcopy(bonds_length_norm)
while None in tmp_norm:
tmp_norm.remove(None)
tmp_value = [
tmp_norm[i] * (bond_max - bond_min) + bond_min
for i in range(len(tmp_norm))
]
color_bar = plt.scatter(tmp_norm, tmp_norm, c=tmp_value, cmap=cmap)
plt.colorbar(color_bar)
plt.savefig(args.output + ".colorbar.png")
plt.close()
"""
#Scheme II
for i in range(len(bonds)):
if bonds_length_norm[i] == None:
pass
else:
bonds[i].set("Color", "%f, %f, %f, %f" % (35.3, 35.3, 35.3, bonds_length_norm[3]))
"""
# write xsd file
xsd.write(args.output)
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()