diff --git a/ippisite/ippidb/utils.py b/ippisite/ippidb/utils.py index 4dca5e17da3b7d2414dcc3337786d6d2fcca89db..fda288b6535ca7c90f0160c537431b0f3a65b35c 100644 --- a/ippisite/ippidb/utils.py +++ b/ippisite/ippidb/utils.py @@ -9,8 +9,8 @@ def mol2smi(mol_string): class FingerPrinter(object): - def __init__(self, name): - self.fingerprinter = ob.OBFingerprint.FindFingerprint("FP4") + def __init__(self, name="FP4"): + self.fingerprinter = ob.OBFingerprint.FindFingerprint(name) self._smiles_parser = ob.OBConversion() self._smiles_parser.SetInFormat("smi") @@ -26,19 +26,16 @@ class FingerPrinter(object): self.fingerprinter.GetFingerprint(self.parse_smiles(smiles), fp) return fp - def fp_list(self, smiles_list): - "generate a list of fingerprints from a list of SMILES strings" - return [self.fp(smiles) for smiles in smiles_list] + def fp_dict(self, smiles_dict): + "generate a dict of {compound id: fingerprint} from a dict of {compound id: fingerprint}" + return {compound_id: self.fp(smiles_entry) for compound_id, smiles_entry in smiles_dict.items()} - def tanimoto_fps(self, smiles_query, fp_list): - "perform a tanimoto similarity search using a smiles query string on a fingerprints list" + def tanimoto_fps(self, smiles_query, fp_dict): + "perform a tanimoto similarity search using a smiles query string on a dict of {compound id: fingerprint}" fp_query = self.fp(smiles_query) - db_tanimotos = [] - for fp_entry in fp_list: - db_tanimotos.append(self.fingerprinter.Tanimoto(fp_query, fp_entry)) - return db_tanimotos - - def tanimoto_smiles(self, query_smiles, smiles_list): - "perform a tanimoto similarity search using a smiles query on a smiles list" - fps = self.fp_list(smiles_list) - return self.tanimoto_fps(query_smiles, fps) \ No newline at end of file + return {compound_id: self.fingerprinter.Tanimoto(fp_query, fp_entry) for compound_id, fp_entry in fp_dict.items()} + + def tanimoto_smiles(self, query_smiles, smiles_dict): + "perform a tanimoto similarity search using a smiles query on a dict of {compound id: SMILES}" + fp_dict = self.fp_dict(smiles_dict) + return self.tanimoto_fps(query_smiles, fp_dict) \ No newline at end of file