def test_to_list(self):
x = 1
y = [1]
z = tuple([1, 2, 3])
v = np.array([1, 2, 3])
for k in [x, y, z, v]:
self.assertTrue(type(to_list(k)), list)
def get_flat_data(self, graphs, targets):
"""
Expand the graph dictionary to form a list of features and targets tensors
This is useful when the model is trained on assembled graphs on the fly
Args:
graphs: (list of dictionary) list of graph dictionary for each structure
targets: (list of float or list) correpsonding target values for each structure
Returns:
tuple(node_features, edges_features, global_values, index1, index2, targets)
"""
atoms = []
bonds = []
states = []
index1 = []
index2 = []
final_targets = []
for g, t in zip(graphs, targets):
if isinstance(g, dict):
atoms.append(np.array(g['atom']))
bonds.append(np.array(g['bond']))
states.append(g['state'])
index1.append(g['index1'])
index2.append(g['index2'])
final_targets.append(to_list(t))
return atoms, bonds, states, index1, index2, final_targets
def index_rep_from_structure(structure, r=4):
"""
Take a pymatgen structure and convert it to a index-type graph representation
The graph will have node, distance, index1, index2, where node is a vector
of Z number of atoms in the structure, index1 and index2 mark the atom
indices forming the bond and separated by distance
Args:
structure: (pymatgen structure)
r: (float) distance cutoff
Returns:
(dictionary)
"""
atom_i_segment_id = [
] # index list for the center atom i for all bonds (row index)
atom_i_j_id = [] # index list for atom j
atom_number = []
all_neighbors = structure.get_all_neighbors(r, include_index=True)
distances = []
for k, n in enumerate(all_neighbors):
atom_number.append(structure.sites[k].specie.Z)
if len(n) < 1:
index = None
else:
_, distance, index = list(zip(*n))
index = np.array(index)
distance = np.array(distance)
if index is not None:
ind = np.argsort(index)
it = itemgetter(*ind)
index = it(index)
index = to_list(index)
index = [int(i) for i in index]
distance = distance[ind]
distances.append(distance)
atom_i_segment_id.extend([k] * len(index))
atom_i_j_id.extend(index)
else:
pass
if len(distances) < 1:
return None
else:
return {
'distance': np.concatenate(distances),
'index1': atom_i_segment_id,
'index2': atom_i_j_id,
'node': atom_number
}
def test_to_list(self):
x = 1
y = [1]
self.assertListEqual(to_list(x), [1])
self.assertListEqual(to_list(y), y)
