Latest STM workflow with Bdg and STM tools

aiida_kkr/tools/tools_STM_scan.py

@@ -228,3 +228,234 @@ def create_combined_potential_node_cf(add_position, host_remote, imp_potential_n
     pot_combined_node = create_combined_potential_node(add_position, host_remote, imp_potential_node)
 
     return pot_combined_node
+
+#####################################################################
+# STM pathfinder
+
+
+def STM_pathfinder(host_structure):
+    #from aiida_kkr.tools import find_parent_structure
+    from ase.spacegroup import Spacegroup
+    """This function is used to help visualize the scanned positions
+       and the symmetries that are present in the system            
+    """
+
+    """
+    inputs::
+    host_struture : RemoteData : The Remote data contains all the information needed to create the path to scan
+
+    outputs::
+    struc_info : Dict  : Dictionary containing the structural information of the film
+    matrices   : Array : Array containing the matrices that generate the symmetries of the system
+
+    """
+
+    def info_creation(structure):
+        from ase.spacegroup import get_spacegroup
+        # List of the Bravais vectors
+        vec_list = structure.cell.tolist()
+
+        # Find the Bravais vectors that are in plane vectors
+        plane_vectors = {'plane_vectors': [], 'space_group': ''}
+        for vec in vec_list:
+            # Is this sufficient to find all the in-plane vectors?
+            if vec[2] == 0:
+                plane_vectors['plane_vectors'].append(vec[:2])
+
+        space_symmetry = get_spacegroup(ase_struc)
+        plane_vectors['space_group'] = space_symmetry.no
+
+        return plane_vectors
+
+    def symmetry_finder(struc_info):
+        # Here we get the symmetry operations that are possible
+        symmetry_matrices = Spacegroup(struc_info['space_group'])
+
+        # Reduce the dimensionality, we only want the 2D matrices
+        matrices = []
+        for element in symmetry_matrices.get_rotations():
+            matrices.append(element[:2, :2])
+
+        # Uniqueness of the elements must be preserved
+        unique_matrices = []
+        for matrix in matrices:
+            if not any(np.array_equal(matrix, m) for m in unique_matrices):
+                unique_matrices.append(matrix)
+
+        return unique_matrices
+
+    struc = find_parent_structure(host_structure)
+    # clone the structure since it has already been saved in AiiDA and cannot be modified
+    supp_struc = struc.clone()
+
+    # If the structure is not periodic in every direction we force it to be.
+    supp_struc.pbc = (True, True, True)
+
+    # ASE struc
+    ase_struc = supp_struc.get_ase()
+
+    # Structural informations are stored here
+    struc_info = info_creation(ase_struc)
+
+    # The structural informations are then used to find the symmetries of the system
+    symm_matrices = symmetry_finder(struc_info)
+
+    return struc_info, symm_matrices
+
+
+@engine.calcfunction
+def STM_pathfinder_cf(host_structure):
+    """
+    Calcfunction that gives back the structural information of the film, and the symmetries of the system
+    """
+
+    struc_info, symm_matrices = STM_pathfinder(host_structure)
+
+    return struc_info, symm_matrices
+
+
+###################################################################
+# lattice generation (function of lattice plot)
+
+
+def lattice_generation(x_len, y_len, rot, vec):
+    import math
+    """
+    x_len : int  : value to create points between - x and x.
+    y_len : int  : value to create points between - y and y.
+    rot   : list : list of the rotation matrices given by the symmetry of the system.
+    vec   : list : list containing the two Bravais vectors.
+    """
+
+    # Here we create a grid in made of points whic are the linear combination of the lattice vectors
+    lattice_points = []
+
+    for i in range(-x_len, x_len + 1):
+        lattice_points_col = []
+        for j in range(-y_len, y_len + 1):
+            p = [i * x + j * y for x, y in zip(vec[0], vec[1])]
+            lattice_points_col.append(p)
+        lattice_points.append(lattice_points_col)
+
+    # Eliminiatio of the symmetrical sites
+    points_to_eliminate = []
+
+    for i in range(-x_len, x_len + 1):
+        for j in range(-y_len, y_len + 1):
+            if ( 
+               (lattice_points[i][j][0] > 0 or math.isclose(lattice_points[i][j][0],0, abs_tol=1e-3)) and
+               (lattice_points[i][j][1] > 0 or math.isclose(lattice_points[i][j][1],0, abs_tol=1e-3))
+                ):
+                for element in rot[1:]:
+                    point = np.dot(element, lattice_points[i][j])
+                    if point[0] >= 0 and point[1] >= 0:
+                        continue
+                    else:
+                        points_to_eliminate.append(point)
+
+    points_to_scan = []
+
+    for i in range(-x_len, x_len + 1):
+        for j in range(-y_len, y_len + 1):
+            eliminate = False
+            for elem in points_to_eliminate:
+                # Since there can be some numerical error in the dot product we use the isclose function
+                if (
+                    math.isclose(elem[0], lattice_points[i][j][0], abs_tol=1e-4) and
+                    math.isclose(elem[1], lattice_points[i][j][1], abs_tol=1e-4)
+                ):
+                    eliminate = True
+            if not eliminate:
+                points_to_scan.append(lattice_points[i][j])
+
+    return points_to_eliminate, points_to_scan
+
+
+###############################################################################################
+# lattice plot
+
+
+def lattice_plot(plane_vectors, symm_vec, symm_matrices, grid_length_x, grid_length_y):
+    #from aiida_kkr.tools.tools_STM_scan import lattice_generation
+    import matplotlib.pyplot as plt
+    from matplotlib.lines import Line2D
+
+    origin = np.array([[0, 0], [0, 0]])
+    # Generation of the points to plot
+    unused, used = lattice_generation(grid_length_x, grid_length_y, symm_matrices, plane_vectors)
+
+    # Plotting of the points
+    for element in unused:
+        plt.scatter(element[0], element[1], marker='s', s=130, c='#33638DFF')
+
+    for element in used:
+        plt.scatter(element[0], element[1], marker='D', s=130, c='#FDE725FF')
+
+    # Plot of the crystal symmetry directions, tag must be activated.
+    if symm_vec:
+        import numpy.linalg as lin
+
+        for element in symm_matrices:
+            eig_val, eig_vec = lin.eig(element)
+
+        for element in eig_vec:
+            plt.quiver(
+                *origin, element[0], element[1], alpha=1, color='#B8DE29FF', angles='xy', scale_units='xy', scale=1
+            )
+
+    # Plot of the Bravais lattice
+    for element in plane_vectors:
+        plt.quiver(*origin, element[0], element[1], color='#3CBB75FF', angles='xy', scale_units='xy', scale=1)
+
+    legend_elements = [
+        Line2D([0], [0], color='#33638DFF', lw=2, label='Unscanned Sites', marker='s'),
+        Line2D([0], [0], color='#FDE725FF', lw=2, label='Scanned Sites', marker='D'),
+        Line2D([0], [0], color='#3CBB75FF', lw=2, label='Bravais lattice'),
+    ]
+    plt.legend(handles=legend_elements, bbox_to_anchor=(0.75, -0.15))
+
+    plt.title('Lattice plot and symmetry directions')
+    plt.ylabel('y direction')
+    plt.xlabel('x direction')
+    #plt.xticks(np.arange(-grid_length, grid_length, float(plane_vectors[0][0])))
+    #plt.set_cmap(cmap)
+    plt.grid(linestyle='--')
+    plt.show()
+
+
+##########################################################
+# find linear combination coefficients
+
+
+def find_linear_combination_coefficients(plane_vectors, vectors):
+    from operator import itemgetter
+    """This helper function takes the planar vectors and a list of vectors
+       and return the coefficients in the base of the planar vectors"""
+
+    # Formulate the system of equations Ax = b
+    A = np.vstack((plane_vectors[0], plane_vectors[1])).T
+
+    # Solve the system of equations using least squares method
+    data = []
+    for element in vectors:
+        b = element
+        # We use the least square mean error procedure to evaulate the units of da and db
+        # lstsq returns: coeff, residue, rank, and singular value
+        # We only need the coefficient.
+        data.append(np.linalg.lstsq(A, b, rcond=None))
+
+    indices = []
+    for element in data:
+        supp = []
+        for elem in element[0]:
+            # Here we round to an integer, this is because of the numerical error
+            # which is present inside the calculation.
+            supp.append(round(elem))
+        indices.append(supp)
+
+    # Before returning the indices, we reorder them first from the lowest to the highest valued
+    # on the x axis and then from the lowest to the highest on the y axis.
+
+    indices = sorted(indices, key=itemgetter(0, 1))
+
+    return indices