axis_rotation.py

import numpy as np

def axis_rotation(str_file, angle, axis):
    with open(str_file, 'r') as data:
        fl = data.readlines()
        
        
        scale_factor = float(fl[1])
        
        lattice_vectors = scale_factor*np.array([list(map(float, i.split())) for i in fl[2:5]])
        
        Coord_mat = np.array([list(map(float, fl[8:][i].split())) for i in range(len(fl[8:]))]).T
        

        theta = np.pi*angle/180
        
        Rx = np.array([[1, 0, 0], [0, np.cos(theta), -np.sin(theta)], [0, np.sin(theta), np.cos(theta)]])
        Ry = np.array([[np.cos(theta), 0, -np.sin(theta)], [0, 1, 0], [np.sin(theta), 0, np.cos(theta)]])
        Rz = np.array([[np.cos(theta), -np.sin(theta), 0], [np.sin(theta), np.cos(theta), 0], [0, 0, 1]])
        
        if axis=="X":
            rot_mat = Rx
        elif axis=="Y":
            rot_mat = Ry
        elif axis=="Z":
            rot_mat = Rz
        else:
            print("Type error")
        
        new_Coord = np.dot(rot_mat, Coord_mat).T
        

        with open('POSCAR_rot'+str(axis)+str(angle), 'w') as output:
            output.write("generated by sampad's python code" +'\n')
            output.write("1.0" + '\n')
            
            for lv in lattice_vectors:
                output.write('\t'+"{:.16f}".format(lv[0])+'\t'+"{:.16f}".format(lv[1])+'\t'+ "{:.16f}".format(lv[2])+'\n')

            output.writelines(fl[5:8])
            for nc in new_Coord:
                output.write('\t'+"{:.16f}".format(nc[0])+'\t'+"{:.16f}".format(nc[1])+'\t'+"{:.16f}".format(nc[2])+'\n')
    print("\nPOSCAR_rot"+str(axis)+str(angle)+ " is written")            
    return rot_mat, new_Coord

structure_file = input("Enter the structure file in POSCAR format\n")
angle = eval(input("Enter the angle of rotation\n"))
axis = input("Enter the axis of rotation in capital: X or Y or Z\n")

axis_rotation(structure_file, angle, axis)
#print(axis_rotation(structure_file, angle, axis))