Why isn't there an additional broadcast dimension over the atom indices? Won't this code only calculate pairwise distances between the same atom types in different residues instead of all pairwise distances between two residues?
def get_interface_residues(positions, atom_mask, asym_id, interface_threshold):
coord_diff = positions[..., None, None, :, :] - positions[..., None, :, :, None, :]
pairwise_dists = torch.sqrt(torch.sum(coord_diff ** 2, dim=-1))
diff_chain_mask = (asym_id[..., None, :] != asym_id[..., :, None]).float()
pair_mask = atom_mask[..., None, None, :] * atom_mask[..., None, :, :, None]
mask = (diff_chain_mask[..., None, None] * pair_mask).bool()
pairwise_dists = torch.where(mask, pairwise_dists, torch.inf)
min_dist_per_res = pairwise_dists.amin(dim=(-1, -2))
valid_interfaces = torch.sum((min_dist_per_res < interface_threshold).float(), dim=-1)
interface_residues_idxs = torch.nonzero(valid_interfaces, as_tuple=True)[0]
return interface_residues_idxs
Why isn't there an additional broadcast dimension over the atom indices? Won't this code only calculate pairwise distances between the same atom types in different residues instead of all pairwise distances between two residues?
openfold/openfold/data/data_transforms_multimer.py
Line 320 in e938c18
The following code broadcasts over both the sequence length and atoms: