sqs2vasp.py

#!/usr/bin/env python3
# Convert bestsqs.out generated by mcsqs (ATAT) to POSCAR file (POSCAR.sqs) for VASP
#
#     python3 sqs2vasp.py [ sqsfile (default: bestsqs.out) ]
# 
# (POSCAR.sqs will be generated)


import sys


file_str = sys.argv[1] if len(sys.argv) > 1 else 'bestsqs.out'
file_pos = 'POSCAR.sqs'

def dot(a, b):
    b_trans = list(zip(*b))    # transpose matrix
    return [[sum(ai*bi for ai, bi in zip(row, col)) for col in b_trans] for row in a]

def cvect(a, b, c, alpha, beta, gamma):
    from math import radians, sin, cos, sqrt

    af, bt, gm = map(radians, (alpha, beta, gamma))
    cy = (cos(af) - cos(bt)*cos(gm)) / sin(gm)
    vect = [[a, 0, 0],
            [b*cos(gm), b*sin(gm), 0],
            [c*cos(bt), c*cy, c*sqrt(sin(bt)**2 - cy**2)]]
    return vect

def det3(matrix):
    return ( matrix[0][0] * matrix[1][1] * matrix[2][2]
           + matrix[0][1] * matrix[1][2] * matrix[2][0]
           + matrix[0][2] * matrix[1][0] * matrix[2][1]
           - matrix[0][0] * matrix[1][2] * matrix[2][1]
           - matrix[0][1] * matrix[1][0] * matrix[2][2] 
           - matrix[0][2] * matrix[1][1] * matrix[2][0] )


with open(file_str, 'r') as f:
    line = f.readline().strip().split()
    coor = [float(x) for x in line]
    if len(coor) == 6:
        coor = cvect(*coor)
    elif len(coor) == 3:
        coor = [coor,]
        for _ in range(2):
            line = f.readline().strip().split()
            coor.append([float(x) for x in line])
    else:
        raise IOError('Unable to open lattice file defined by ATAT')

    basic = []
    for _ in range(3):
        line = f.readline().strip().split()
        basic.append([float(x) for x in line])
    
    xyzs = []
    elmts = []
    labels = []
    for line in f:
        if len(line.strip()) == 0:
            continue
        x, y, z, label = line.strip().split(maxsplit=3)
        xyz = list(map(float, [x, y, z]))
        elmt = label.split(sep=',')[0].split(sep='=')[0].strip()
        if elmt in elmts:
            index = elmts.index(elmt)
            xyzs[index].append(xyz)
            labels[index].append(label)
        else:
            elmts.append(elmt)
            xyzs.append([xyz,])
            labels.append([label,])

with open(file_pos, 'w') as f:
    header = 'Generated by sqs2vasp.py'
    scell = ['{:d}'.format(round(p)) for row in basic for p in row]
    header += '    {}   {}'.format(round(det3(basic)), '  '.join(scell))
    f.write(header + '\n')
    f.write('  1.000000\n')
    basic_coor = dot(basic, coor)
    for ib in basic_coor:
        line = '  {:>22.16f}{:>22.16f}{:>22.16f}\n'.format(*ib)
        f.write(line)
    dsp = '{:>5s}'*len(elmts) + '\n'
    f.write(dsp.format(*elmts))
    dsp = '  ' + '{:>5d}'*len(elmts) + '\n'
    f.write(dsp.format(*[len(x) for x in xyzs]))
    f.write('Cartesian\n')
    dsp = '  {:>22.16f}{:>22.16f}{:>22.16f}    {:s}\n'
    for ixyzs, ilabels in zip(xyzs, labels):
        ixyzs_coor = dot(ixyzs, coor)
        for ixyz, ilabel in zip(ixyzs_coor, ilabels):
            f.write(dsp.format(*ixyz, ilabel))