import torch


def generate_ellipse_segmentation(
        shape: tuple,
        center_of_mass: tuple,
        ellipse_axis: tuple
        ):
    """Generate an ellipse shaped segmentation

    :param shape: Shape the 2D image
    :param center_of_mass: Coordinates of the center of mass
    :param ellipse_axis: ellipse parameters (a, b) in equation
        `\\left(x/a\\right)^2 + \\left(y/b\\right)^2 = 1`
    """
    # Generate array of indices
    arr = torch.arange(
            shape[0] * shape[1]
            ).reshape(
                    shape[0], shape[1]
                    )

    # Transform array of indices into a segmentation
    segmentation = (
            (
                (arr // shape[1] - center_of_mass[0]) / ellipse_axis[0]

                ) ** 2 +
            (
                (arr % shape[1] - center_of_mass[1]) / ellipse_axis[1]

                ) ** 2
    ) < 1
    return segmentation