fga.py

import os
import traceback
import itertools
import numpy as np
from math import ceil
from ase.build.tools import sort
from ase.io import read
from fga.util import *

norm = np.linalg.norm

def cut_cell(unit_atom,xy_par):
    """
    The function fits the atoms object in the given cell/parallelogram. It will only work if the cell is generated by multiplying
    the unit cell by a valid projection matrix
    
    :unit_atom: (ase atoms object) to be cut
    :xy_cell: four corners of the parallelogram along which the atoms object is cut
    :return: (cut ase atoms object)
    
    """   
    atoms = unit_atom.copy()

    original_cell = atoms.cell
    new_cell = [[xy_par[1][0],xy_par[1][1],0],[xy_par[2][0],xy_par[2][1],0],[0,0,original_cell[2][2]]]
    
    #Repeat and orient atom before cutting
    mag_new_cell = norm((new_cell[0][0] + new_cell[1][0], new_cell[0][1] + new_cell[1][1]))
    min_len      = min(norm((atoms.cell[0][0], atoms.cell[0][1])),norm((atoms.cell[1][0], atoms.cell[1][1])))
    repetition   = int(ceil(1.5*mag_new_cell/min_len))
    atoms *= (repetition,repetition,1)
    
    atoms.set_cell(new_cell)
    atoms = square_up(atoms)

    #Setup rough box
    cell = atoms.cell
    max_x = max(0, cell[0][0], cell[1][0], cell[0][0] + cell[1][0]) + 0.5
    min_x = min(0, cell[0][0], cell[1][0], cell[0][0] + cell[1][0]) - 0.5
    max_y = max(0, cell[0][1], cell[1][1], cell[0][1] + cell[1][1]) + 0.5
    min_y = min(0, cell[0][1], cell[1][1], cell[0][1] + cell[1][1]) - 0.5
    
    #Delete atoms not in the rough box
    valid = [] 
    for x in range(len(atoms)):
        x_value = atoms.get_positions()[x][0]
        y_value = atoms.get_positions()[x][1]
        if x_value  < min_x or x_value > max_x or y_value > max_y or y_value < min_y:
            continue
        else:
            valid.append(x)
    del atoms[[i for i in range(len(atoms)) if i not in valid]]
    
    #Rotate axis to make parallel
    atoms.rotate('x', [1,0,0], rotate_cell=True)
    atoms = remove_doubles(atoms)
    atoms.translate((0.1, 0.1, 0.1)) 
    atoms = remove_doubles(atoms)
    atoms = sort(atoms,tags=None)
    return atoms

def convert_matrix_parallelogram(unit_cell,matrix,max_size=40,max_length_ratio=1.5,unrotate=False):
    """
    Multiply unit cell with the projection matrix.
    Accept projected cell if within certain parameters and then rotate along x axis
    :unit_cell: (2x2 matrix) [x_vector,y_vector]
    :matrix: [[i,j],[k,l]] Projection matrix 
    :max_size: (float) maximum size of cell diagonal
    :max_length_ratio: (float) Maximum ratio of x_vector to y_vector
    :unrotate: (Boolean) to return rotated or unrotated films
    
    :return: [corner1,corner2,corner3.corner4]
             Four corners of the parallelogram (corner = (x,y))
    """
    xy_cell = unit_cell[0:2,0:2]
    mat = np.matmul(matrix,xy_cell)
    new_cell = [(0,0),(mat[0][0],mat[0][1]),(mat[1][0],mat[1][1]), (mat[0][0]+mat[1][0],mat[0][1]+mat[1][1])]

    
    if norm(new_cell[2]) == 0:
        #print("divide by zero")
        return False
    len_ratio = norm(new_cell[1])/norm(new_cell[2])
    
    if norm(new_cell[3]) > max_size or norm((new_cell[1][0]-new_cell[2][0], new_cell[1][1], new_cell[2][1])) > max_size:
        #print("doesnt satisfy max_size")
        return False

    if len_ratio > max_length_ratio or len_ratio < 1/max_length_ratio:
        #print("doesnt satisfy max_ratio")
        return False
      
    if abs(clock_angle(new_cell[1], new_cell[2])) >= np.pi/2: #Rotate cells which are obtuse
        new_cell = [(0, 0), (-1*new_cell[1][0] , -1*new_cell[1][1]), (new_cell[3][0] - 1*new_cell[1][0], new_cell[3][1] -1*new_cell[1][1]), (new_cell[2][0] - 1*new_cell[1][0], new_cell[2][1] -1*new_cell[1][1])]

    #Check if the parallelofram isnt too small
    a = np.array([new_cell[1][0], new_cell[1][1]])
    b = np.array([new_cell[2][0], new_cell[2][1]])
    if abs(np.cross(a, b)) < 0.0001:
        #print("its a line")
        return False
    elif np.cross(a, b) < 0: #To find left handed cells 
        new_cell = [(0, 0), (new_cell[2][0] , new_cell[2][1]), (new_cell[1][0], new_cell[1][1]), (new_cell[3][0], new_cell[3][1])]
    
    #Save unrotated cell for atom orientations
    unrotate_cell = new_cell
    
    #Rotate along x-axis
    theta = -1*clock_angle((1,0), new_cell[1])
    rotatedPolygon = []
    for corner in new_cell:
        rotatedPolygon.append((corner[0]*np.cos(theta)-corner[1]*np.sin(theta),corner[0]*np.sin(theta)+corner[1]*np.cos(theta)))
    new_cell = rotatedPolygon

    #Final unnecessary check
    if abs(new_cell[2][1]) < 0.0001 and abs(new_cell[3][1]) < 0.0001:
        #print("values too small")
        return False
    
    if unrotate:
        return unrotate_cell
    else:
        return new_cell

def make_unique_cells(search_size,unit_atom_cell,max_size=40,max_length_ratio=1.3,tol=1e-4):
    """
    Make unique cells for the given range of projection matrices
    :search_size: (int) search size decides the maxumum repetition of cell in xy direction
    :unit_atom_cell: (2x2 matrix) unit cell to be projected
    :max_size: (float) maximum size of cell diagonal
    :max_length_ratio: (float) Maximum ratio of x_vector to y_vector
    :unrotate: (Boolean) to return rotated or unrotated films
    
    :return: Dictionary of unique cells
             cells[i][0] --> [corner1,corner2,corner3.corner4] Four corners of the parallelogram (corner = (x,y))
             cells[i][1] --> [[i,j],[k,l]] Projection matrix
    """
    
    #Make all projection matrix and multiply them by the unit cell
    cells = []
    for i,j in itertools.product(range(-search_size,search_size+1,1),range(-search_size,search_size+1,1)):
        for k,l in itertools.product(range(search_size+1),range(search_size+1)):
            if (i == 0 and j == 0) or (l == 0 and k == 0):
                continue
            else:
                matrix = [[i,j],[k,l]]
                cell = convert_matrix_parallelogram(unit_atom_cell,matrix,max_size=max_size,max_length_ratio=max_length_ratio,unrotate=False)
                if cell:
                    cells.append([cell,matrix])            
    if not cells:
        print("No cells found")
        return False
    #Check for repetitions
    val = 0 
    clean_cells = [cells[0]]
    for og_cell in cells:
        for new_cell in clean_cells:
            i=og_cell[0]
            j=new_cell[0]
            if check(i[1][0], j[1][0]) and check(i[1][1], j[1][1]) and check(i[2][0], j[2][0]) and check(i[2][1], j[2][1]):
                val = 1
                break
        if val != 1:
            clean_cells.append(og_cell)
        val = 0
        
    return clean_cells

def make_valid_pairs(substrate_cells,film_cells,unit_film,unit_substrate,max_ratio_change=0.5,max_num_of_atoms=200,max_theta_dialation=0.5,remove_same_magnitude=False):
    """
    Try to fit the film paralleograms on substrate parallelograms
    :substrate_cells: Output from make_unique_cells() for substrate
    :film_cells: Output from make_unique_cells() for film
    :unit_film: ase atoms object of film
    :unit_substrate: ase atoms object of substrate
    :max_ratio_change: (float) Maximum ratio to be looked on the either side of perfect fit
    :max_num_of_atoms: (int) Maximum number of atoms allowed in cell
    :max_theta_dialation: (float) Maximum difference in between the inner angle of two parallelograms
    :remove_same_magnitude: (Boolean) whether cells with same diagonal magnitude be removed
    
    :return: dictionary of valid pairs
            pair[0] --> Four corners of substrate parallelogram 'sub_cell'
            pair[1] --> Projection matrix to get the sub parallelogram 'sub_matrix'
            pair[2] --> Four corners of film parallelogram 'film_cell'
            pair[3] --> Projection matrix to get the film parallelogram 'film_matrix'
            
    """
    sub_area = get_area(unit_substrate.cell)
    film_area = get_area(unit_film.cell)
    ratio = sub_area/film_area
    max_error = max_error = calc_error(unit_film.cell,(1 + max_ratio_change) * unit_film.cell)
    pairs = []
    #loop over all pairs
    for sub,film in itertools.product(substrate_cells,film_cells):
        sx, sy = sub[0][1], sub[0][2] #Substrate - X and Y vector
        fx, fy = film[0][1], film[0][2] #Film - X and Y vector
        sub_scale = round(get_area([sx, sy])/sub_area, 2) #Scale of substrate modification
        film_scale = round(get_area([fx, fy])/film_area, 2) #Scale of film modification
        if ratio - max_ratio_change <= film_scale/sub_scale <= ratio + max_ratio_change: #Check Ratio
            error = calc_error([sx, sy],[fx, fy])
            if  error <= max_error: #Check Error
                total_atoms = film_scale*len(unit_film)+sub_scale*len(unit_substrate)
                theta_dilat  = (clock_angle((1,0), sub[0][2]) - clock_angle((1,0), film[0][2]))/clock_angle((1,0), sub[0][2])*100
                if total_atoms < max_num_of_atoms and abs(theta_dilat) < max_theta_dialation: #Check dilations and number of atoms
                    x_dist = distortion(sx, fx)
                    y_dist = distortion(sy, fy)
                    mag = norm(sub[0][3])
                    is_looping=True
                    if pairs and remove_same_magnitude: #Check for cells with same magnitude
                        for pair in pairs:
                            if check(pair['magnitude'],mag):
                                is_looping=False
                                break
                    if is_looping:
                        pairs.append({'sub_cell': sub[0],'sub_matrix': sub[1],
                                      'film_cell': film[0],'film_matrix': film[1],
                                      'ratio': round(film_scale/sub_scale, 4),
                                      'film_scale': film_scale, 'sub_scale': sub_scale,
                                      'x_distortion': x_dist,'y_distortion': y_dist, 
                                      'theta_dilation':theta_dilat,
                                      'error':error, 'magnitude':mag})
    return pairs

def put_film_substrate(film,sub,buffer=2,vacuum=10):
    """
    Fit the cut film on cut substrate
    :film: ase atoms object
    :sub: ase atoms object
    :buffer: distance between film and substrate
    :vacuum: the vacuum between z images
    """
    
    #Substract the vacuum from substrate
    min_z_sub = np.min(sub.get_positions()[:, 2])
    for atom in sub:
        atom.position[2] = atom.position[2] - min_z_sub
    
    #Set the substrate unit cell ad film unit cell
    film.set_cell([sub.cell[0],sub.cell[1],film.cell[2]], scale_atoms=True)
    max_z_sub = np.max(sub.get_positions()[:, 2])
    min_z_film = np.min(film.get_positions()[:, 2])
    for atom in film:
        atom.position[2] = atom.position[2] - min_z_film + max_z_sub + buffer
    
    #Put everything together
    sub += film
    sub.set_scaled_positions(sub.get_scaled_positions())
    sub.center(axis=2, vacuum=vacuum)
    sub = sort(sub,tags=None)

    return sub

def one_conbination_run(unit_substrate_path,unit_film_path,sub_search_size=4,film_search_size=4
             ,max_size=40,max_length_ratio=1.5,max_ratio_change=0.5,max_num_of_atoms=200
             ,max_theta_dialation=0.5,remove_same_magnitude=True,buffer=2,vacuum=10):
    
    """
    For the given pair of substrate and film, make all the pairs
    :unit_substrate_path: (str) Path to substrate POSCAR
    :unit_film_path: (str) Path to film POSCAR
    :sub_search_size: (int) maxumum repetition of cell in xy direction for the substrate
    :film_search_size: (int) maxumum repetition of cell in xy direction for the film
    """
    
    sub_name =  str(unit_substrate_path.split("/")[-1][6:])
    film_name = str(unit_film_path.split("/")[-1][6:])
    print("Making cells for ",film_name," on ",sub_name)
    
    base_dir = "./Generated_films/Substrate" + sub_name + "/Film" + film_name
    try:
        os.makedirs(base_dir)
    except Exception as e: 
        print("Warning! The library exists")
        pass
    
    print("Substrate Search space :",sub_search_size)
    print("Film Search space :",film_search_size)
    unit_substrate = read(unit_substrate_path)
    unit_film = read(unit_film_path)
    substrate_cells = make_unique_cells(sub_search_size,unit_substrate.cell,max_size=max_size,max_length_ratio=max_length_ratio)
    print("Total substrate cells found :",len(substrate_cells))
    if not substrate_cells:
        print("None of the substrate cells satisfy parameters")
        print("------------------------------------------- \n")
        return
    film_cells = make_unique_cells(film_search_size,unit_film.cell,max_size=max_size,max_length_ratio=max_length_ratio)
    print("Total film cells found :",len(film_cells))
    if not film_cells:
        print("None of the film cells satisfy parameters")
        print("------------------------------------------- \n")
        return
    pairs = make_valid_pairs(substrate_cells,film_cells,unit_film,unit_substrate,
                             max_ratio_change=max_ratio_change,max_num_of_atoms=max_num_of_atoms,
                             max_theta_dialation=max_theta_dialation,remove_same_magnitude=remove_same_magnitude)
    if not pairs:
        print("No valid pairs found. Exiting loop")
        print("------------------------------------------- \n")
        return
    pairs = sorted(pairs, key=lambda d: d['error'])
    print("Total pairs found :",len(pairs))
    print("Minimum Ratio :",pairs[0]['film_scale'],"film on",pairs[0]['sub_scale'],"substrate")
    for pair in pairs:
        film_cell = convert_matrix_parallelogram(unit_film.cell,pair['film_matrix'],max_size=max_size,max_length_ratio=max_length_ratio,unrotate=True)
        sub_cell = convert_matrix_parallelogram(unit_substrate.cell,pair['sub_matrix'],max_size=max_size,max_length_ratio=max_length_ratio,unrotate=True)
        film = cut_cell(unit_film,film_cell)
        sub = cut_cell(unit_substrate,sub_cell)
        atoms = put_film_substrate(film,sub,buffer=buffer,vacuum=vacuum)
        fs1 = int(pair['film_scale'])
        ss1 = int(pair['sub_scale'])
        path = base_dir+'/'+str(round(fs1/ss1,2))+'_'+str(fs1)+'_'+str(ss1)+'_'+str(int(pair['magnitude']))
        poscar_path = path + '/POSCAR'
        txt = cell_summary(pair,film,sub,film_name,sub_name)
        try:
            os.makedirs(path)
            atoms.write(poscar_path,format='vasp')
            with open(path + '/cell_summary.txt', 'w') as f:
                f.write(txt)
        except Exception as e:
            #print(e)
            continue
    print("------------------------------------------- \n")  
    return

def get_best_cell(unit_substrate_path,unit_film_path,sub_search_size=4,film_search_size=4
             ,max_size=40,max_length_ratio=1.5,max_ratio_change=0.5,max_num_of_atoms=200
             ,max_theta_dialation=0.5,remove_same_magnitude=True,buffer=2,vacuum=10):
    
    """
    For the given pair of substrate and film, make all the pairs
    :unit_substrate_path: (str) Path to substrate POSCAR
    :unit_film_path: (str) Path to film POSCAR
    :sub_search_size: (int) maxumum repetition of cell in xy direction for the substrate
    :film_search_size: (int) maxumum repetition of cell in xy direction for the film
    """
    sub_name =  str(unit_substrate_path.split("/")[-1][6:])
    film_name = str(unit_film_path.split("/")[-1][6:])
    
    unit_substrate = read(unit_substrate_path)
    unit_film = read(unit_film_path)
    substrate_cells = make_unique_cells(sub_search_size,unit_substrate.cell,max_size=max_size,max_length_ratio=max_length_ratio)
    if not substrate_cells:
        print("None of the substrate cells satisfy parameters")
        print("------------------------------------------- \n")
        return
    film_cells = make_unique_cells(film_search_size,unit_film.cell,max_size=max_size,max_length_ratio=max_length_ratio)
    if not film_cells:
        print("None of the film cells satisfy parameters")
        print("------------------------------------------- \n")
        return
    pairs = make_valid_pairs(substrate_cells,film_cells,unit_film,unit_substrate,
                             max_ratio_change=max_ratio_change,max_num_of_atoms=max_num_of_atoms,
                             max_theta_dialation=max_theta_dialation,remove_same_magnitude=remove_same_magnitude)
    pairs = sorted(pairs, key=lambda d: d['error'])
    if not pairs:
        print("No valid pairs found. Exiting loop")
        print("------------------------------------------- \n")
        return
    pair = pairs[0]
    film_cell = convert_matrix_parallelogram(unit_film.cell,pair['film_matrix'],max_size=max_size,max_length_ratio=max_length_ratio,unrotate=True)
    sub_cell = convert_matrix_parallelogram(unit_substrate.cell,pair['sub_matrix'],max_size=max_size,max_length_ratio=max_length_ratio,unrotate=True)
    film = cut_cell(unit_film,film_cell)
    sub = cut_cell(unit_substrate,sub_cell)
    atoms = put_film_substrate(film,sub,buffer=buffer,vacuum=vacuum)
    print("------------------------------------------- \n")
    print(cell_summary(pair,film,sub,film_name,sub_name))
    print("------------------------------------------- \n")
    return atoms

def cell_summary(pair,film,sub,film_name,sub_name):
    """
    The function writes a text file with all the necessary information
    :pair: Output from make_valid_pairs()
    :film: ase atoms object for film
    :sub: ase atoms object for substrate
    :film_name: (str) name of film
    :sub_name: (str) name of substrate
    
    :return: (str)
    """
    
    txt  = film_name + " on " + sub_name + '\n'
    txt += "Film Scale        : " + str(pair['film_scale']) + '\n'
    txt += "Sub Scale         : " + str(pair['sub_scale']) + '\n'
    txt += "Number of Atoms   : " + str(len(sub)) + '\n'
    txt += "Chemical Formula  : " + str(sub.get_chemical_formula()) +'\n'
    txt += "Area of Cell      : " + str(round(get_area(sub.cell),3)) + " Angstroms^2" + '\n'
    txt += "Magnitude of Cell : " + str(round(pair['magnitude'], 3)) + " Angstroms" + '\n'
    txt += "Ratio             : " + str(round(pair['film_scale']/pair['sub_scale'], 3)) + '\n'
    txt += "Substrate:        : " + str(pair['sub_matrix']) +'\n'
    txt += "Film              : " + str(pair['film_matrix']) +'\n'
    txt += "X-Distortion      : " + str(pair['x_distortion']) + " % " + '\n' 
    txt += "Y-Distortion      : " + str(pair['y_distortion']) + " % " + '\n'
    txt += "Theta-Dilation    : " + str(round(pair['theta_dilation'],3)) + " % " + '\n'

    return txt