-
Notifications
You must be signed in to change notification settings - Fork 0
/
structure.py
131 lines (108 loc) · 3.94 KB
/
structure.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
from ase.io import gen
import pandas as pd
import numpy as np
from ase.geometry.analysis import Analysis
def getslab(struct):
"""
Input:
struct: structre from which we will trim unbound species (Atoms object)
Output:
baseslab: structure with unbound species trimmed
"""
adjmat = Analysis(struct, bothways = True).adjacency_matrix[0]
numnodes = adjmat.shape[0]
g = Graph(numnodes)
for i in range(numnodes):
for j in range(numnodes):
if adjmat[i,j]:
g.addEdge(i,j)
cc = g.connectedComponents()
maingraph = np.array([i for i in cc if 0 in i][0])
return struct[[atom.index for atom in struct if atom.index in maingraph]]
def getFragIndices(struct, check = False):
"""
Input:
struct: structure (Atoms object)
check: would you like to check max connectivity?
Output:
array of indices for fragments
"""
a = Analysis(struct, bothways = True)
adjmat = a.adjacency_matrix[0].toarray()
if check:
maxbonds = {'Ar': 0, 'Si': 6, 'F': 1, 'N': 4, 'H': 1, 'C': 5}
for i, adjrow in enumerate(adjmat):
elem = struct[i].symbol
while np.sum(adjrow) > maxbonds[elem] + 1: # +1 because adjmat[i,i] = 1
distances = {atom.index: struct.get_distance(i, atom.index) for atom in struct if adjrow[atom.index]}
for j, a in enumerate(adjrow):
if j in distances.keys() and distances[j] == np.max(list(distances.values())):
adjrow[j] = 0
adjmat[j, i] = 0 # Delete both directions of the edge
adjmat[i] = adjrow
numnodes = adjmat.shape[0] # Adjacency matrix is NxN, N = #atoms
g = Graph(numnodes)
for i in range(numnodes):
for j in range(numnodes):
if adjmat[i,j]:
g.addEdge(i,j)
cc = g.connectedComponents()
fragIndices = np.array([i for i in cc if 0 not in i])
return fragIndices
def getslabs(data, directory, useInputs = False):
"""
Utility for getting and writing slab files from readData (utils.py) function
data is the df from readData function or any df with (struct, in) and (struct, out) columns
"""
if useInputs:
slabSource = data['struct']['in']
else:
slabSource = data['struct']['out']
dataDir = directory
slabs = {}
# # to generate slabs
for key, value in slabSource.iteritems():
slabs[key] = getslab(value)
for key, value in slabs.items():
gen.write_gen(dataDir + "slab{}.gen".format(key), value)
# to read slabs
for key in data.index:
slabs[key] = gen.read_gen(dataDir + "slab{}.gen".format(key))
if useInputs:
data.loc[:, ('struct', 'inslab')] = pd.Series(slabs)
else:
data.loc[:, ('struct', 'outslab')] = pd.Series(slabs)
class Graph:
# init function to declare class variables
def __init__(self, V):
self.V = V
self.adj = [[] for i in range(V)]
def DFSUtil(self, temp, v, visited):
# Mark the current vertex as visited
visited[v] = True
# Store the vertex to list
temp.append(v)
# Repeat for all vertices adjacent
# to this vertex v
for i in self.adj[v]:
if visited[i] == False:
# Update the list
temp = self.DFSUtil(temp, i, visited)
return temp
# method to add an undirected edge
def addEdge(self, v, w):
self.adj[v].append(w)
self.adj[w].append(v)
# Method to retrieve connected components
# in an undirected graph
def connectedComponents(self):
visited = []
cc = []
for i in range(self.V):
visited.append(False)
for v in range(self.V):
if visited[v] == False:
temp = []
cc.append(self.DFSUtil(temp, v, visited))
return cc
# This code is contributed by Abhishek Valsan