def test_graph_generator(self):
feature = [np.random.normal(size=(3, 4)), np.random.normal(size=(2, 4))]
bond = [np.random.normal(size=(2, 5)), np.random.normal(size=(1, 5))]
glob_features = [np.random.normal(size=(1, 2)), np.random.normal(size=(1, 2))]
index1 = [np.array([0, 1]), np.array([0])]
index2 = [np.array([1, 2]), np.array([1])]
targets = np.random.normal(size=(2, 1))
gen = GraphBatchGenerator(feature, bond, glob_features, index1, index2, targets,
batch_size=2)
data = gen[0]
self.assertListEqual(list(data[0][0].shape), [1, 5, 4])
self.assertListEqual(list(data[0][1].shape), [1, 3, 5])
self.assertListEqual(list(data[0][2].shape), [1, 2, 2])
self.assertListEqual(list(data[0][3].shape), [1, 3])
self.assertListEqual(list(data[0][4].shape), [1, 3])
self.assertListEqual(list(data[1].shape), [1, 2, 1])
# Make sure it still functions if a target is not provided
gen = GraphBatchGenerator(feature, bond, glob_features, index1, index2, batch_size=2)
data = gen[0]
self.assertEqual(7, len(data)) # Should only be the inputs
self.assertListEqual(list(data[0].shape), [1, 5, 4])
# when bonds are one dimension arrays
bond = [np.random.normal(size=(2,)), np.random.normal(size=(1,))]
gen = GraphBatchGenerator(feature, bond, glob_features, index1, index2, targets,
batch_size=2)
data = gen[0]
self.assertListEqual(list(data[0][1].shape), [1, 3])
def _create_generator(self, *args, **kwargs) -> \
Union[GraphBatchDistanceConvert, GraphBatchGenerator]:
if hasattr(self.graph_converter, 'bond_converter'):
kwargs.update(
{'distance_converter': self.graph_converter.bond_converter})
return GraphBatchDistanceConvert(*args, **kwargs)
return GraphBatchGenerator(*args, **kwargs)
def test_graph_generator(self):
feature = [
np.random.normal(size=(3, 4)),
np.random.normal(size=(2, 4))
]
bond = [np.random.normal(size=(2, 5)), np.random.normal(size=(1, 5))]
glob_features = [
np.random.normal(size=(1, 2)),
np.random.normal(size=(1, 2))
]
index1 = [np.array([0, 1]), np.array([0])]
index2 = [np.array([1, 2]), np.array([1])]
targets = np.random.normal(size=(2, 1))
gen = GraphBatchGenerator(feature,
bond,
glob_features,
index1,
index2,
targets,
batch_size=2)
data = gen[0]
self.assertListEqual(list(data[0][0].shape), [1, 5, 4])
self.assertListEqual(list(data[0][1].shape), [1, 3, 5])
self.assertListEqual(list(data[0][2].shape), [1, 2, 2])
self.assertListEqual(list(data[0][3].shape), [1, 3])
self.assertListEqual(list(data[0][4].shape), [1, 3])
self.assertListEqual(list(data[1].shape), [1, 2, 1])
def _create_generator(self, *args, **kwargs):
if hasattr(self.graph_convertor, 'bond_convertor'):
kwargs.update(
{'distance_convertor': self.graph_convertor.bond_convertor})
return GraphBatchDistanceConvert(*args, **kwargs)
else:
return GraphBatchGenerator(*args, **kwargs)
def create_cached_generator(self) -> GraphBatchGenerator:
"""Generates features for all of the molecules and stores them in memory
Returns:
(GraphBatchGenerator) Graph genereator that relies on having the graphs in memory
"""
# Make all the graphs
graphs = self._generate_graphs(self.mols)
# Turn them into a fat array
inputs = self.converter.get_flat_data(graphs, self.targets)
return GraphBatchGenerator(*inputs, is_shuffle=self.is_shuffle,
batch_size=self.batch_size)
def create_cached_generator(self) -> GraphBatchGenerator:
"""Generates features for all of the molecules and stores them in memory
Returns:
(GraphBatchGenerator) Graph genereator that relies on having the graphs in memory
"""
# Make all the graphs
graphs = self._generate_graphs(self.mols)
# Turn them into a fat array
atom_features, bond_features, state_features, index1_list, index2_list, targets = \
self.converter.get_flat_data(graphs, self.targets) # type: ignore
return GraphBatchGenerator(atom_features=atom_features,
bond_features=bond_features,
state_features=state_features,
index1_list=index1_list,
index2_list=index2_list,
targets=targets,
is_shuffle=self.is_shuffle,
batch_size=self.batch_size)
def _create_generator(self, *args, **kwargs):
if self.distance_convertor is not None:
kwargs.update({'distance_convertor': self.distance_convertor})
return GraphBatchDistanceConvert(*args, **kwargs)
else:
return GraphBatchGenerator(*args, **kwargs)
