def test_hypergraph_mask_gen(self): molecules = True grammar_cache = 'tmp.pickle' grammar = 'hypergraph:' + grammar_cache # create a grammar cache inferred from our sample molecules g = HypergraphGrammar(cache_file=grammar_cache) if os.path.isfile(g.cache_file): os.remove(g.cache_file) g.strings_to_actions(get_zinc_smiles(5)) mask_gen1 = get_codec(molecules, grammar, 30).mask_gen mask_gen2 = get_codec(molecules, grammar, 30).mask_gen mask_gen1.priors = False mask_gen2.priors = True policy1 = SoftmaxRandomSamplePolicy( bias=mask_gen1.grammar.get_log_frequencies()) policy2 = SoftmaxRandomSamplePolicy() lp = [] for mg in [mask_gen1, mask_gen2]: mg.reset() mg.apply_action([None]) logit_priors = mg.action_prior_logits() # that includes any priors lp.append( torch.from_numpy(logit_priors).to(device=device, dtype=torch.float32)) dummy_model_output = torch.ones_like(lp[0]) eff_logits = [] for this_lp, policy in zip(lp, [policy1, policy2]): eff_logits.append(policy.effective_logits(dummy_model_output)) assert torch.max((eff_logits[0] - eff_logits[1]).abs()) < 1e-6
def test_hypergraph_grammar_codec(self): molecules = True input = smiles1 grammar_cache = 'tmp.pickle' grammar = 'hypergraph:' + grammar_cache # create a grammar cache inferred from our sample molecules g = HypergraphGrammar(cache_file=grammar_cache) g.strings_to_actions(smiles) self.check_codec(input, molecules, grammar)
def test_hypergraph_rpe_parser_bad_smiles(self): g = HypergraphGrammar() trees = [] for smile in bad_smiles: try: trees.append( g.normalize_tree(hypergraph_parser(MolFromSmiles(smile)))) except (AssertionError, IndexError): print('Failed for {}'.format(smile)) raise
def test_hypergraph_mask_gen(self): molecules = True grammar_cache = 'tmp.pickle' grammar = 'hypergraph:' + grammar_cache # create a grammar cache inferred from our sample molecules g = HypergraphGrammar(cache_file=grammar_cache) if os.path.isfile(g.cache_file): os.remove(g.cache_file) g.strings_to_actions(smiles) codec = get_codec(molecules, grammar, max_seq_length) self.generate_and_validate(codec)
def test_mask_gen(self): g = HypergraphGrammar() g.strings_to_actions( smiles) # that initializes g with the rules from these molecules g.calc_terminal_distance() batch_size = 10 max_rules = 50 all_actions = [] next_action = [None for _ in range(batch_size)] mask_gen = HypergraphMaskGenerator(max_rules, g) while True: try: next_masks = mask_gen(next_action) next_action = [] for mask in next_masks: inds = np.nonzero(mask)[0] next_act = random.choice(inds) next_action.append(next_act) all_actions.append(next_action) except StopIteration: break all_actions = np.array(all_actions).T # the test is that we get that far, producing valid molecules all_smiles = g.decode_from_actions(all_actions) for smile in all_smiles: print(smile)
def test_get_set_params_as_vector(self): grammar_cache = 'hyper_grammar_guac_10k_with_clique_collapse.pickle' # 'hyper_grammar.pickle' g = HypergraphGrammar.load(grammar_cache) m = CondtionalProbabilityModel(g) out = m.get_params_as_vector() out[0] = 1 m.set_params_from_vector(out) out2 = m.get_params_as_vector() assert out2[0] == out[0]
def test_hypergraph_rpe_parser(self): g = HypergraphGrammar() g.strings_to_actions(smiles) trees = [ g.normalize_tree(hypergraph_parser(MolFromSmiles(smile))) for smile in smiles ] rule_pairs = extract_popular_hypergraph_pairs(g, trees, 10) parser = HypergraphRPEParser(g, rule_pairs) collapsed_trees = [parser.parse(smile) for smile in smiles] recovered_smiles = [] for tree in collapsed_trees: graph = graph_from_graph_tree(tree) mol = to_mol(graph) recovered_smiles.append(MolToSmiles(mol)) self.assertEqual(smiles, recovered_smiles)
def test_hypergraph_rpe(self): g = HypergraphGrammar() g.strings_to_actions(smiles) tree = g.normalize_tree(hypergraph_parser(MolFromSmiles(smiles1))) num_rules_before = len(g.rules) rule_pairs = extract_popular_hypergraph_pairs(g, [tree], 10) num_rules_after = len(g.rules) tree_rules_before = len(tree.rules()) collapsed_tree = apply_hypergraph_substitution(g, tree, rule_pairs[0]) tree_rules_after = len(collapsed_tree.rules()) graph = graph_from_graph_tree(collapsed_tree) mol = to_mol(graph) recovered_smiles = MolToSmiles(mol) self.assertEqual(smiles1, recovered_smiles) self.assertGreater(num_rules_after, num_rules_before) self.assertLess(tree_rules_after, tree_rules_before)
def test_parser_roundtrip_via_indices(self): # TODO: cheating a bit here, reconstruction does fail for some smiles # chirality gets inverted sometimes, so need to run the loop twice to reconstruct the original g = HypergraphGrammar() g.delete_cache() actions = g.strings_to_actions(smiles) re_smiles = g.decode_from_actions(actions) re_actions = g.strings_to_actions(re_smiles) rere_smiles = g.decode_from_actions(re_actions) for old, new in zip(smiles, rere_smiles): old_fix = old #.replace('@@', '@')#.replace('/','\\') new_fix = new #.replace('@@', '@')#.replace('/','\\').replace('\\','') assert old_fix == new_fix
def test_cycle_length(self): g = HypergraphGrammar.load('hyper_grammar.pickle') for r in g.rules: if r is not None: g = r.to_nx() cycles = nx.minimum_cycle_basis(g) if len(cycles) > 0: maxlen = max([len(c) for c in cycles]) if maxlen > 7: print(maxlen) cc = nx.number_connected_components(g) if cc > 1: print(cc)
def get_codec(molecules, grammar, max_seq_length): if grammar is True: grammar = 'classic' # character-based models if grammar is False: if molecules: charlist = [ 'C', '(', ')', 'c', '1', '2', 'o', '=', 'O', 'N', '3', 'F', '[', '@', 'H', ']', 'n', '-', '#', 'S', 'l', '+', 's', 'B', 'r', '/', '4', '\\', '5', '6', '7', 'I', 'P', '8', ' ' ] else: charlist = [ 'x', '+', '(', ')', '1', '2', '3', '*', '/', 's', 'i', 'n', 'e', 'p', ' ' ] codec = CharacterCodec(max_len=max_seq_length, charlist=charlist) elif grammar == 'classic': if molecules: codec = CFGrammarCodec(max_len=max_seq_length, grammar=grammar_zinc, tokenizer=zinc_tokenizer) else: codec = CFGrammarCodec(max_len=max_seq_length, grammar=grammar_eq, tokenizer=eq_tokenizer) codec.mask_gen = GrammarMaskGenerator(max_seq_length, codec.grammar) elif grammar == 'new': codec = CFGrammarCodec(max_len=max_seq_length, grammar=grammar_zinc_new, tokenizer=zinc_tokenizer_new) codec.mask_gen = GrammarMaskGeneratorNew(max_seq_length, codec.grammar) elif 'hypergraph' in grammar: grammar_cache = grammar.split(':')[1] assert grammar_cache is not None, "Invalid cached hypergraph grammar file:" + str( grammar_cache) codec = HypergraphGrammar.load(grammar_cache) codec.MAX_LEN = max_seq_length codec.normalize_conditional_frequencies() codec.calc_terminal_distance( ) # just in case it wasn't initialized yet codec.mask_gen = HypergraphMaskGenerator(max_seq_length, codec) assert hasattr(codec, 'PAD_INDEX') return codec
def test_graph_from_graph_tree_idempotent(self): g = HypergraphGrammar() g.strings_to_actions(smiles) g.calc_terminal_distance() tree = g.normalize_tree(hypergraph_parser(MolFromSmiles(smiles1))) # The second call here would fail before # This was solved by copying in remove_nonterminals where the issue # was with mutating the parent tree.node state graph1 = graph_from_graph_tree(tree) graph2 = graph_from_graph_tree(tree) mol1 = to_mol(graph1) mol2 = to_mol(graph2) recovered_smiles1 = MolToSmiles(mol1) recovered_smiles2 = MolToSmiles(mol2) self.assertEqual(smiles1, recovered_smiles1) self.assertEqual(recovered_smiles1, recovered_smiles2)
settings = get_settings(molecules=True, grammar='new') thresh = 100000 # Read in the strings f = open(settings['source_data'], 'r') L = [] for line in f: line = line.strip() L.append(line) if len(L) > thresh: break f.close() fn = "rule_hypergraphs.pickle" max_rules = 50 if os.path.isfile(fn): rm = HypergraphGrammar.load(fn) else: rm = HypergraphGrammar(cache_file=fn) bad_smiles = [] for num, smile in enumerate(L[:100]): try: smile = MolToSmiles(MolFromSmiles(smile)) print(smile) mol = MolFromSmiles(smile) actions = rm.strings_to_actions([smile]) re_smile = rm.decode_from_actions(actions)[0] mol = MolFromSmiles(smile) if re_smile != smile: print("SMILES reconstruction wasn't perfect for " + smile) print(re_smile)