blacklist_ids = list(chain(*hetatm_residues.values()))
protein_amap = np.array([i for i in range(mol.GetNumAtoms())
if i not in blacklist_ids])
protein_coords = np.array(mol.GetConformer().GetPositions())[protein_amap]
// Pocket selection based on cutoff
mask = (cdist(protein_coords, ligand_coords) <= cutoff).any(axis=1)
After Change
blacklist_ids = list(chain(*hetatm_residues.values()))
protein_amap = np.array([i for i in range(mol.GetNumAtoms())
if i not in blacklist_ids])
conf = mol.GetConformer()
protein_coords = np.array([conf.GetAtomPosition(i)
for i in range(mol.GetNumAtoms())])[protein_amap]
// Pocket selection based on cutoff
