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()