import torch


def betagamma2xyz(beta, gamma, axis=None):
    """
    Returns new tensors corresponding to angle representation from the cartesian representation.

    Args:
        beta (`torch.FloatTensor`): beta rotations about y axis.
        gamma (`torch.FloatTensor`): gamma rotations about z axis.
        axis (str, optional): cartesian axis. If None, return all axis. Defaults to None.

    Returns:
        (`torch.FloatTensor`, optional): x positions.
        (`torch.FloatTensor`, optional): y positions.
        (`torch.FloatTensor`, optional): z positions.
    """

    if axis == "x":
        return torch.sin(beta) * torch.cos(gamma)
    if axis == "y":
        return torch.sin(beta) * torch.sin(gamma)
    if axis == "z":
        return torch.cos(beta)

    x = torch.sin(beta) * torch.cos(gamma)
    y = torch.sin(beta) * torch.sin(gamma)
    z = torch.cos(beta)

    return x, y, z