from rdkit.Chem import rdFingerprintGenerator
from rdkit.Chem import DataStructs
from rdkit.Chem import SDMolSupplier, SDWriter
def similarity_filter_tversky(ref_smiles, sdf_filename, output_filename):
# 从参考分子 SMILES 字符串中构建 RDKit 分子对象
ref = Chem.MolFromSmiles(ref_smiles)
if ref is None:
print("无法解析参考分子的 SMILES")
return
# 创建 RDKit 分子指纹生成器
fpgen = rdFingerprintGenerator.GetMorganGenerator(radius=2, fpSize=1024)
# 从 SDF 文件中读取分子
suppl = SDMolSupplier(sdf_filename)
if suppl is None:
print("无法读取 SDF 文件")
return
# 创建输出 SDF 文件的写入器
writer = SDWriter(output_filename)
# 遍历每个分子
for mol in suppl:
if mol is None:
print("无法解析 SDF 文件中的分子")
continue
# 计算参考分子和当前分子的 Tversky 相似性
fp1 = fpgen.GetCountFingerprint(ref)
fp2 = fpgen.GetCountFingerprint(mol)
simi = DataStructs.cDataStructs.TverskySimilarity(fp1, fp2, 0.05, 0.95)
# 如果相似性大于 0.4，则将该分子写入输出文件
if simi > 0.4:
writer.write(mol)
print(f"分子与参考分子相似性: {simi}, 分子已保留")
else:
print(f"分子与参考分子相似性: {simi}, 分子被过滤")
writer.close()
print("处理完成")
# 使用示例
reference_smiles = "C1(C2=CC=CC=C2)=CC=CC=C1"
sdf_file = "input.sdf"
output_sdf_file = "output_tversky.sdf"
similarity_filter_tversky(reference_smiles, sdf_file, output_sdf_file)
