Add POSCAR writer and str generator

nafuma/dft/io.py

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