def graph_reward_fn(smiles, predictor, tokens, device, fn):
node_attributes = {}
node_attributes['valence'] = Attribute('node',
'valence',
one_hot=True,
values=[1, 2, 3, 4, 5, 6, 7])
node_attributes['charge'] = Attribute('node',
'charge',
one_hot=True,
values=[-1, 0, 1, 2, 3, 4])
node_attributes['hybridization'] = Attribute(
'node', 'hybridization', one_hot=True, values=[0, 1, 2, 3, 4, 5, 6, 7])
node_attributes['aromatic'] = Attribute('node',
'aromatic',
one_hot=True,
values=[0, 1])
node_attributes['atom_element'] = Attribute('node',
'atom_element',
one_hot=True,
values=list(range(11)))
adj_matrix, node_feature_matrix = process_graphs(smiles,
node_attributes,
get_atomic_attributes,
edge_attributes=None,
get_bond_attributes=None)
adj_matrix = np.array(adj_matrix)
node_feature_matrix = np.array(node_feature_matrix)
adj_tensor = torch.from_numpy(adj_matrix).to(dtype=torch.float32,
device=device)
node_feature_tensor = torch.from_numpy(node_feature_matrix).to(
dtype=torch.float32, device=device)
inp = (node_feature_tensor, adj_tensor)
prediction = predictor(inp, eval=True)
if predictor.task == "classification":
prediction = torch.argmax(prediction, dim=1)
rewards = fn(prediction.to(torch.float))
return rewards
53: 9
}
if atomic_num in atomic_mapping.keys():
attr_dict['atom_element'] = atomic_mapping[atomic_num]
else:
attr_dict['atom_element'] = 10
attr_dict['valence'] = atom.GetTotalValence()
attr_dict['charge'] = atom.GetFormalCharge()
attr_dict['hybridization'] = atom.GetHybridization().real
attr_dict['aromatic'] = int(atom.GetIsAromatic())
return attr_dict
node_attributes = {}
node_attributes['valence'] = Attribute('node',
'valence',
one_hot=True,
values=[1, 2, 3, 4, 5, 6])
node_attributes['charge'] = Attribute('node',
'charge',
one_hot=True,
values=[-1, 0, 1, 2, 3, 4])
node_attributes['hybridization'] = Attribute('node',
'hybridization',
one_hot=True,
values=[0, 1, 2, 3, 4, 5, 6, 7])
node_attributes['aromatic'] = Attribute('node',
'aromatic',
one_hot=True,
values=[0, 1])
node_attributes['atom_element'] = Attribute('node',
'atom_element',
def get_atomic_attributes(atom):
atomic_num = atom.GetAtomicNum()
attr_dict = dict(atom_element=atomic_num)
return attr_dict
def get_edge_attributes(bond):
attr_dict = dict()
attr_dict['bond_type'] = bond.GetBondTypeAsDouble()
return attr_dict
node_attributes = dict(atom_element=Attribute('node',
'atom_element',
one_hot=False), )
edge_attributes = dict(bond_type=Attribute('edge', 'bond_type', one_hot=False))
restrict_min_atoms = 10
restrict_max_atoms = 50
train_dataset = BFSGraphDataset(
get_atomic_attributes,
node_attributes,
'./benchmark_datasets/chembl_full/small_chembl.smi',
cols_to_read=[0, 1],
get_bond_attributes=get_edge_attributes,
edge_attributes=edge_attributes,
delimiter=',',
random_order=True,