def structure2input(structure, r=4, distance_convertor=None, **kwargs):
"""
Take a pymatgen structure and convert it to a index-type graph representation as model input
:param structure: (pymatgen structure)
:param r: (float) cutoff radius
:param state_attributes: (list) a list of state attributes
:param distance_convertor: (object) convert numeric distance values into a vector as bond features
:return: (dictionary) inputs for model.predict
"""
graph = structure2graph(structure, r=r)
if distance_convertor is None:
centers = kwargs.get('centers', np.linspace(0, 6, 100))
width = kwargs.get('width', 0.5)
distance_convertor = GaussianDistance(centers, width)
gnode = [0] * len(structure)
gbond = [0] * len(graph['index1'])
return [
expand_1st(graph['node']),
expand_1st(distance_convertor.convert(graph['distance'])),
expand_1st(np.array(graph['state'])),
expand_1st(np.array(graph['index1'])),
expand_1st(np.array(graph['index2'])),
expand_1st(np.array(gnode)),
expand_1st(np.array(gbond)),
]
def __getitem__(self, index):
batch_index = self.mol_index[index * self.batch_size:(index + 1) *
self.batch_size]
it = itemgetter(*batch_index)
# get atom features from a batch of structures
feature_list_temp = itemgetter_list(self.atom_features, batch_index)
# get atom's structure id
gnode = []
for i, j in enumerate(feature_list_temp):
gnode += [i] * len(j)
# get bond features from a batch of structures
connection_list_temp = itemgetter_list(self.bond_features, batch_index)
# get bond's structure id
gbond = []
for i, j in enumerate(connection_list_temp):
gbond += [i] * len(j)
global_list_temp = itemgetter_list(self.state_features, batch_index)
# assemble atom features together
feature_list_temp = np.concatenate(feature_list_temp, axis=0)
feature_list_temp = self.process_atom_feature(feature_list_temp)
# assemble bond feature together
connection_list_temp = np.concatenate(connection_list_temp, axis=0)
connection_list_temp = self.process_bond_feature(connection_list_temp)
# assemble state feature together
global_list_temp = np.concatenate(global_list_temp, axis=0)
global_list_temp = self.process_state_feature(global_list_temp)
# assemble bond indices
index1_temp = it(self.index1_list)
index2_temp = it(self.index2_list)
index1 = []
index2 = []
offset_ind = 0
for ind1, ind2 in zip(index1_temp, index2_temp):
index1 += [i + offset_ind for i in ind1]
index2 += [i + offset_ind for i in ind2]
offset_ind += (max(ind1) + 1)
# get targets
target_temp = it(self.targets)
target_temp = np.atleast_2d(target_temp)
return [
expand_1st(feature_list_temp),
expand_1st(connection_list_temp),
expand_1st(global_list_temp),
expand_1st(index1),
expand_1st(index2),
expand_1st(gnode),
expand_1st(gbond)
], expand_1st(target_temp)
def graph_to_input(self, graph):
"""
Turns a graph into model input
"""
gnode = [0] * len(graph['atom'])
gbond = [0] * len(graph['index1'])
return [
expand_1st(self.atom_convertor.convert(graph['atom'])),
expand_1st(self.bond_convertor.convert(graph['bond'])),
expand_1st(np.array(graph['state'])),
expand_1st(np.array(graph['index1'])),
expand_1st(np.array(graph['index2'])),
expand_1st(np.array(gnode)),
expand_1st(np.array(gbond)),
]
def predict_structure(self, structure):
"""
Predict property from structure
:param structure:
:return:
"""
self.assert_graph_convertor()
graph = self.graph_convertor(structure)
gnode = [0] * len(graph['node'])
gbond = [0] * len(graph['index1'])
inp = [
expand_1st(graph['node']),
expand_1st(self.distance_convertor.convert(graph['distance'])),
expand_1st(np.array(graph['state'])),
expand_1st(np.array(graph['index1'])),
expand_1st(np.array(graph['index2'])),
expand_1st(np.array(gnode)),
expand_1st(np.array(gbond)),
]
return self.yscaler.inverse_transform(
self.predict(inp).reshape((-1, 1))).ravel()
def test_expand_dim(self):
x = np.array([1, 2, 3])
self.assertListEqual(list(expand_1st(x).shape), [1, 3])