Geometry¤

def combine(
    arrays: Array,
    oversample: int = 1,
    use_sum: bool = False,
) -> Array:
    """
    Combines multiple arrays by multiplying them together, and downsampling the output.

    Parameters
    ----------
    arrays : Array
        The arrays to be combined. Should have shape (n_arrays, npix, npix).
    oversample : int = 1
        The amount to downsample the output by.
    use_sum : bool = False
        Whether to sum the arrays instead of multiplying them.

    Returns
    -------
    array : Array
        The combined array, optionally downsampled by `oversample`.
    """
    method = np.sum if use_sum else np.prod
    array = np.array(arrays)
    if oversample == 1:
        return method(array, 0)
    return dlu.downsample(method(array, 0), oversample)

def circle(coords: Array, radius: float, invert: bool = False) -> Array:
    """
    Calculates a hard-edged circle. This function is not differentiable.

    Parameters
    ----------
    coords : Array
        The coordinates to calculate the circle on.
    radius : float
        The radius of the circle.
    invert : bool = False
        Whether to invert the circle.

    Returns
    -------
    circle : Array
        The circle.
    """

    if invert:
        return (circ_distance(coords, radius) > 0).astype(float)
    return (circ_distance(coords, radius) < 0).astype(float)

def square(coords: Array, width: float, invert: bool = False) -> Array:
    """
    Calculates a hard-edged square. This function is not differentiable.

    Parameters
    ----------
    coords : Array
        The coordinates to calculate the square on.
    width : float
        The width of the square.
    invert : bool = False
        Whether to invert the square.

    Returns
    -------
    square : Array
        The square.
    """
    if invert:
        return (square_distance(coords, width) > 0).astype(float)
    return (square_distance(coords, width) < 0).astype(float)

def rectangle(
    coords: Array, width: float, height: float, invert: bool = False
) -> Array:
    """
    Calculates a hard-edged rectangle. This function is not differentiable.

    Parameters
    ----------
    coords : Array
        The coordinates to calculate the rectangle on.
    width : float
        The width of the rectangle.
    height : float
        The height of the rectangle.
    invert : bool = False
        Whether to invert the rectangle.

    Returns
    -------
    rectangle : Array
        The rectangle.
    """
    if invert:
        return (rectangle_distance(coords, width, height) > 0).astype(float)
    return (rectangle_distance(coords, width, height) < 0).astype(float)

def reg_polygon(coords: Array, rmax: float, nsides: int, invert: bool = False) -> Array:
    """
    Calculates a hard-edged regular polygon. This function is not differentiable.

    Parameters
    ----------
    coords : Array
        The coordinates to calculate the polygon on.
    rmax : float
        The radius of the polygon.
    nsides : int
        The number of sides of the polygon.
    invert : bool = False
        Whether to invert the polygon.

    Returns
    -------
    polygon : Array
        The polygon.
    """

    if invert:
        return (reg_polygon_distance(coords, nsides, rmax) > 0).astype(float)
    return (reg_polygon_distance(coords, nsides, rmax) < 0).astype(float)

def spider(coords: Array, width: float, angles: Array) -> Array:
    """
    Calculates a hard-edged spider. This function is not differentiable.

    Parameters
    ----------
    coords : Array
        The coordinates to calculate the spider on.
    width : float
        The width of the spider.
    angles : Array
        The angles of the spider in degrees.

    Returns
    -------
    spider : Array
        The spider.
    """
    angles = np.array(angles) if not isinstance(angles, np.ndarray) else angles
    calc_fn = vmap(lambda angle: spider_distance(coords, width, angle) < 0)
    spiders = ~lax.reduce(calc_fn(angles), np.array(False), lax.bitwise_or, (0,))
    return spiders.astype(float)

def soft_circle(
    coords: Array, radius: float, clip_dist: float = 0.1, invert: bool = False
) -> Array:
    """
    Calculates a soft-edged circle differentiably. The 'clip_dist' parameter defines
    the distance from the edge to 'soften' up to. A large clip_dist will result in a
    circle with a very soft edge, while a small clip_dist will result in a circle with
    a very hard edge.

    Parameters
    ----------
    coords : Array
        The coordinates to calculate the circle on.
    radius : float
        The radius of the circle.
    clip_dist : float = 0.1
        The distance from the edge to 'soften' up to.
    invert : bool = False
        Whether to invert the circle.

    Returns
    -------
    circle : Array
        The softened circle.
    """
    distances = -circ_distance(coords, radius)
    return soften(distances, clip_dist, invert)

def soft_square(
    coords: Array, width: float, clip_dist: float = 0.1, invert: bool = False
) -> Array:
    """
    Calculates a soft-edged square differentiably. The 'clip_dist' parameter defines
    the distance from the edge to 'soften' up to. A large clip_dist will result in a
    square with a very soft edge, while a small clip_dist will result in a square with
    a very hard edge.

    Parameters
    ----------
    coords : Array
        The coordinates to calculate the square on.
    width : float
        The width of the square.
    clip_dist : float = 0.1
        The distance from the edge to 'soften' up to.
    invert : bool = False
        Whether to invert the square.

    Returns
    -------
    square : Array
        The softened square.
    """
    distances = -square_distance(coords, width)
    return soften(distances, clip_dist, invert)

def soft_rectangle(
    coords: Array,
    width: float,
    height: float,
    clip_dist: float = 0.1,
    invert: bool = False,
) -> Array:
    """
    Calculates a soft-edged rectangle differentiably. The 'clip_dist' parameter defines
    the distance from the edge to 'soften' up to. A large clip_dist will result in a
    rectangle with a very soft edge, while a small clip_dist will result in a rectangle
    with a very hard edge.

    Parameters
    ----------
    coords : Array
        The coordinates to calculate the rectangle on.
    width : float
        The width of the rectangle.
    height : float
        The height of the rectangle.
    clip_dist : float = 0.1
        The distance from the edge to 'soften' up to.
    invert : bool = False
        Whether to invert the rectangle.

    Returns
    -------
    rectangle : Array
        The softened rectangle.
    """
    distances = -rectangle_distance(coords, width, height)
    return soften(distances, clip_dist, invert)

def soft_reg_polygon(
    coords: Array,
    radius: float,
    nsides: int,
    clip_dist: float = 0.1,
    invert: bool = False,
) -> Array:
    """
    Calculates a soft-edged regular polygon differentiably. The 'clip_dist' parameter
    defines the distance from the edge to 'soften' up to. A large clip_dist will result
    in a polygon with a very soft edge, while a small clip_dist will result in a polygon
    with a very hard edge.

    Parameters
    ----------
    coords : Array
        The coordinates to calculate the polygon on.
    radius : float
        The radius of the polygon.
    nsides : int
        The number of sides of the polygon.
    clip_dist : float = 0.1
        The distance from the edge to 'soften' up to.
    invert : bool = False
        Whether to invert the polygon.

    Returns
    -------
    polygon : Array
        The softened polygon.
    """
    distances = -reg_polygon_distance(coords, nsides, radius)
    return soften(distances, clip_dist, invert)

def soft_spider(
    coords: Array,
    width: float,
    angles: Array,
    clip_dist: float = 0.1,
    invert: bool = False,
) -> Array:
    """
    Calculates a soft-edged spider differentiably. The 'clip_dist' parameter defines
    the distance from the edge to 'soften' up to. A large clip_dist will result in a
    spider with a very soft edge, while a small clip_dist will result in a spider with
    a very hard edge.

    Parameters
    ----------
    coords : Array
        The coordinates to calculate the spider on.
    width : float
        The width of the spider.
    angles : Array
        The angles of the spider.
    clip_dist : float = 0.1
        The distance from the edge to 'soften' up to.
    invert : bool = False
        Whether to invert the spider.

    Returns
    -------
    spider : Array
        The softened spider.
    """
    angles = np.array(angles) if not isinstance(angles, np.ndarray) else angles
    spider_fn = vmap(lambda angle: spider_distance(coords, width, angle))
    spiders = -spider_fn(angles).min(axis=0)
    return soften(spiders, clip_dist, invert)