geometry_utils.compute_bond_angles_batch

geometry_utils.compute_bond_angles_batch(atom_labels, atom_coords, atom_mask, bond_atom1_idx, bond_atom2_idx, bond_mask, device)[source]

Compute all bond angles (i–j–k) in a batch.

Parameters:

  • atom_labels (List[List[str]]) – Atom labels per structure (B × N).

  • atom_coords (torch.Tensor, shape (B, N, 3)) – Cartesian coordinates of atoms.

  • atom_mask (torch.BoolTensor, shape (B, N)) – True for real atoms.

  • bond_atom1_idx (torch.LongTensor, shape (B, M)) – Index of first atom in each bond slot.

  • bond_atom2_idx (torch.LongTensor, shape (B, M)) – Index of second atom in each bond slot.

  • bond_mask (torch.BoolTensor, shape (B, M)) – True for real bonds.

  • device (torch.device) – Device for computation.

Returns:

  • angle_ids (List[List[str]]) – Per-structure list of “i–j–k” strings.

  • angles (torch.Tensor, shape (B, P_max)) – Angle values in degrees (0 where padding).

  • mask_ang (torch.BoolTensor, shape (B, P_max)) – True for real angles.

  • idx_tensor (torch.LongTensor, shape (B, P_max, 3)) – Atom index triples for each angle.

Return type:

Tuple[List[List[str]], torch.Tensor, torch.BoolTensor, torch.LongTensor]