import torch


def compute_angle(xyz, angle_list):
    """
    Compute angles between atoms.
    Args:
        xyz (torch.Tensor): coordinates of the atoms.
        angle_list (torch.LongTensor): directed indices
            of sets of three atoms that are all in each
            other's neighborhood.
    Returns:
        angle (torch.Tensor): tensor of angles for each
            element in the angle list.

    """

    # points from j -> i
    r_ji = xyz[angle_list[:, 0]] - xyz[angle_list[:, 1]]
    # points from j -> k
    r_jk = xyz[angle_list[:, 2]] - xyz[angle_list[:, 1]]

    x = torch.sum(r_ji * r_jk, dim=-1)
    y = torch.cross(r_ji, r_jk)
    y = torch.norm(y, dim=-1)
    angle = torch.atan2(y, x)

    return angle