Add POSCAR writer and str generator

This commit is contained in:
rasmusvt 2022-09-27 10:19:51 +02:00
parent 8aa8164a80
commit 61cc0343e0

View file

@ -2,6 +2,10 @@ import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import warnings
import os
import linecache
import nafuma.auxillary as aux
@ -543,3 +547,212 @@ def get_doscar_information(path):
#def get_bader_charges(poscar='POSCAR', acf='ACF.dat'):
def write_poscar(data: dict, options={}):
required_options = ['path', 'overwrite', 'scale']
default_options = {
'path': './POSCAR.vasp',
'overwrite': False,
'scale': 1.0
}
options = aux.update_options(options=options, required_options=required_options, default_options=default_options)
if os.path.isfile(options['path']) and not options['overwrite']:
warnings.warn(f"File {options['path']} already exists, and overwrite disabled. Quitting.")
return None
atom_nums = count_atoms(data)
with open(options['path'], 'w') as fout:
# Write first line
for atom in data['atoms']:
if atom_nums[atom] > 1:
fout.write(f'{atom}{atom_nums[atom]}')
else:
fout.write(f'{atom}')
fout.write('\t - Automatically generated by the NAFUMA Python package.\n')
# Write second line
fout.write(str(options['scale'])+'\n')
# Write lattice parameters
# FIXME Now only writes cells without any angles
fout.write('\t{:<09} \t {:<09} \t {:<09}\n'.format(
str(data['lattice_parameters'][0]),
str(0.0),
str(0.0),
)
)
fout.write('\t{:<09} \t {:<09} \t {:<09}\n'.format(
str(0.0),
str(data['lattice_parameters'][1]),
str(0.0),
)
)
fout.write('\t{:<09} \t {:<09} \t {:<09}\n'.format(
str(0.0),
str(0.0),
str(data['lattice_parameters'][2]),
)
)
# Write atoms
for atom in data['atoms']:
fout.write(f'\t{atom}')
fout.write('\n')
for atom in data['atoms']:
fout.write(f'\t{atom_nums[atom]}')
fout.write('\n')
fout.write('Direct\n')
for atom in data['atoms']:
for label, coords in data['coordinates'].items():
if atom in label:
fout.write('\t{:<09} \t {:<09} \t {:<09}\n'.format(
coords[0],
coords[1],
coords[2]
)
)
def apply_transformation(data, rotation=[[1, 0, 0], [0, 1, 0], [0, 0, 1]], translation=[0,0,0]):
if not isinstance(rotation, np.ndarray):
rotation = np.array(rotation)
if not rotation.shape == (3,3):
print("NOOOO!!!!")
for label in data['coordinates'].keys():
data['coordinates'][label] = rotation.dot(data['coordinates'][label])
data['coordinates'][label] = translate_back(data['coordinates'][label])
data['coordinates'][label] = data['coordinates'][label] + translation
return data
def translate_back(coords):
for i, coord in enumerate(coords):
if coord < 0:
while coords[i] < 0:
coords[i] = coords[i]+1
elif coord >= 1:
while coords[i] >= 1:
coords[i] = coords[i]-1
return coords
def count_atoms(data):
atom_nums = {}
for atom in data['atoms']:
atom_nums[atom] = 0
for label in data['coordinates'].keys():
for atom in data['atoms']:
if atom in label:
atom_nums[atom] += 1
return atom_nums
def append_data(data, new_data):
if not new_data:
return data
if not data:
data = {
'atoms': new_data['atoms'],
'coordinates': {},
'lattice_parameters': new_data['lattice_parameters']
}
atom_num = count_atoms(data)
new_coords = {}
for label, coords in data['coordinates'].items():
new_coords[label] = coords
extend_unit_cell = [0,0,0]
for label, coords in new_data['coordinates'].items():
atom = ''.join(filter(str.isalpha, label))
number = int(''.join(filter(str.isnumeric, label)))
new_number = number + atom_num[atom]
new_label = atom+str(new_number)
new_coords[new_label] = coords
for i, coord in enumerate(coords):
if coord > 1 and np.floor(coord) >= extend_unit_cell[i]:
extend_unit_cell[i] = np.floor(coord)
data['coordinates'] = new_coords
return data
def make_supercell(data, supercell):
for i, param in enumerate(data['lattice_parameters']):
data['lattice_parameters'][i] = supercell[i] * param
if supercell[i] > 0:
for label, coords in data['coordinates'].items():
data['coordinates'][label][i] = (1/supercell[i])*data['coordinates'][label][i]
def copy_data(data):
new_data = {}
new_data['atoms'] = list(data['atoms'])
new_data['coordinates'] = data['coordinates'].copy()
new_data['lattice_parameters'] = list(data['lattice_parameters'])
return new_data
def unit_vector(vector):
""" Returns the unit vector of the vector. """
return vector / np.linalg.norm(vector)
def angle_between(v1, v2):
""" Returns the angle in radians between vectors 'v1' and 'v2'::
>>> angle_between((1, 0, 0), (0, 1, 0))
1.5707963267948966
>>> angle_between((1, 0, 0), (1, 0, 0))
0.0
>>> angle_between((1, 0, 0), (-1, 0, 0))
3.141592653589793
"""
v1_u = unit_vector(v1)
v2_u = unit_vector(v2)
return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))