Coordinates¤

def cart2polar(coordinates: Array) -> Array:
    """
    Converts the input (x, y) Cartesian coordinates into (r, phi) polar
    coordinates.

    Parameters
    ----------
    coordinates : Array
        The (x, y) Cartesian coordinates to be converted into polar
        coordinates.

    Returns
    -------
    coordinates : Array
        The input Cartesian coordinates converted into (r, phi) polar
        coordinates.
    """
    x, y = coordinates
    return np.array([np.hypot(x, y), np.arctan2(y, x)])

def polar2cart(coordinates: Array) -> Array:
    """
    Converts the input (r, phi) polar coordinates into (x, y) Cartesian
    coordinates.

    Parameters
    ----------
    coordinates : Array
        The (r, phi) polar coordinates to be converted into Cartesian
        coordinates.

    Returns
    -------
    coordinates : Array
        The input polar coordinates converted into (x, y) Cartesian
        coordinates.
    """
    r, phi = coordinates
    return np.array([r * np.cos(phi), r * np.sin(phi)])

def pixel_coords(
    npixels: int,
    diameter: float = None,
    radius: float = None,
    pixel_scale: float = None,
    polar: bool = False,
    fft_style: bool = False,
) -> Array:
    """
    Returns a paraxial set of 2d coordinates for each pixel center.

    Parameters
    ----------
    npixels : int
        The output size of the coordinates array to generate.
    diameter : float = None
        The diameter of the coordinates array to generate.
    radius : float = None
        The radius of the coordinates array to generate.
    pixel_scale : float = None
        The pixel scale of the coordinates array to generate.
    polar : bool = False
        Output the coordinates in polar (r, phi) coordinates.
    fft_style : bool = False
        If True, use FFT-style centering. For even npixels this produces integer
        centered coordinates. For odd npixels this is identical to the default.

    Returns
    -------
    coordinates : Array
        The array of pixel center coordinates.
    """
    supplied = sum(value is not None for value in (diameter, radius, pixel_scale))
    if supplied != 1:
        raise ValueError(
            "Exactly one of diameter, radius, or pixel_scale must be provided."
        )

    if diameter is not None:
        pixscale = diameter / npixels
    elif radius is not None:
        pixscale = 2 * radius / npixels
    else:
        pixscale = pixel_scale

    # Default: symmetric pixel-center coordinates (half-integer for even N)
    offsets = (0.0, 0.0)

    # FFT-style: shift by +0.5 pixel for even N so coordinates become integer-centered
    if fft_style and (npixels % 2 == 0):
        offsets = (pixscale / 2.0, pixscale / 2.0)

    # Generate the coordinates
    coords = nd_coords((npixels,) * 2, (pixscale,) * 2, offsets=offsets, indexing="xy")

    # Convert to polar if requested
    if polar:
        return cart2polar(coords)
    return coords

def nd_coords(
    npixels: int | tuple[int, ...],
    pixel_scales: float | tuple[float, ...] = 1.0,
    offsets: float | tuple[float, ...] = 0.0,
    indexing: str = "xy",
) -> Array:
    """
    Returns a set of nd pixel center coordinates, with an optional offset. Each
    dimension can have a different number of pixels, pixel scale and offset by passing
    in tuples of values: `nd_coords((10, 10), (1, 2), (0, 1))`. pixel scale and offset
    can also be passed in as floats to apply those values to all dimensions, i.e.:
    `nd_coords((10, 10), 1, 0)`.

    The indexing argument is the same as in numpy.meshgrid, i.e.: giving the
    string ‘ij’ returns a meshgrid with matrix indexing, while ‘xy’ returns a
    meshgrid with Cartesian indexing. In the 2-D case with inputs of length M
    and N, the outputs are of shape (N, M) for ‘xy’ indexing and (M, N) for
    ‘ij’ indexing. In the 3-D case with inputs of length M, N and P, outputs
    are of shape (N, M, P) for ‘xy’ indexing and (M, N, P) for ‘ij’ indexing.

    Parameters
    ----------
    npixels : int | tuple[int, ...]
        The number of pixels in each dimension.
    pixel_scales : float | tuple[float, ...] = 1.0
        The pixel_scales of each dimension. If a tuple, the length
        of the tuple must match the number of dimensions. If a float, the same
        scale is applied to all dimensions.
    offsets : float | tuple[float, ...] = 0.0
        The offset of the pixel centers in each dimension. If a tuple, the
        length of the tuple must match the number of dimensions. If a float,
        the same offset is applied to all dimensions.
        set to 0.
    indexing : str = 'xy'
        The indexing of the output. Default is 'xy'. See numpy.meshgrid for
        more details.

    Returns
    -------
    coordinates : Array
        The positions of the pixel centers in the given dimensions.
    """
    if indexing not in ["xy", "ij"]:
        raise ValueError("indexing must be either 'xy' or 'ij'.")

    # Validate npixels and ensure tuple
    npixels = _cast_tuple(npixels, "npixels")

    # Get the number of dimensions
    ndim = max(
        len(npixels), _input_len(pixel_scales, "mean"), _input_len(offsets, "std")
    )

    # Validate and ensure tuples
    pixel_scales = _cast_scalar(pixel_scales, ndim, "pixel_scales")
    offsets = _cast_scalar(offsets, ndim, "offsets")

    if len(npixels) != ndim:
        npixels *= ndim

    def pixel_fn(n, offset, scale):
        start = -(n - 1) / 2 * scale - offset
        end = (n - 1) / 2 * scale - offset
        return np.linspace(start, end, n)

    # Generate the linear edges of each axes
    lin_pixels = jtu.map(pixel_fn, npixels, offsets, pixel_scales)

    # Output (x, y) for 2d, else in order.
    positions = np.array(np.meshgrid(*lin_pixels, indexing=indexing))

    # Squeeze the output in case of 1d input
    return np.squeeze(positions)

def translate_coords(coords: Array, translation: Array) -> Array:
    """
    Translates the coordinates by to a new centre. Translation must have shape (2,).

    Parameters
    ----------
    coords : Array
        The input coordinates to translate.
    translation : Array
        The translation to apply to the coordinates.

    Returns
    -------
    coords : Array
        The translated coordinates.
    """
    return coords - translation[:, None, None]

def compress_coords(coords: Array, compress: Array) -> Array:
    """
    Compresses the coordinates by a given factor. Compress must have shape (2,).

    Parameters
    ----------
    coords : Array
        The input coordinates to compress.
    compress : Array
        The compression to apply to the coordinates.

    Returns
    -------
    coords : Array
        The compressed coordinates.
    """
    return coords * compress[:, None, None]

def shear_coords(coords: Array, shear: Array) -> Array:
    """
    Shears the coordinates by a given factor. Shear must have shape (2,).

    Parameters
    ----------
    coords : Array
        The input coordinates to shear.
    shear : Array
        The shear to apply to the coordinates.

    Returns
    -------
    coords : Array
        The sheared coordinates.
    """
    trans_coords = np.transpose(coords, (0, 2, 1))
    return coords + trans_coords * shear[:, None, None]

def rotate_coords(coords: Array, rotation: float) -> Array:
    """
    Rotates the coordinates by a given angle.

    Parameters
    ----------
    coords : Array
        The input coordinates to rotate.
    rotation : float, radians
        The rotation to apply to the coordinates.

    Returns
    -------
    coords : Array
        The rotated coordinates.
    """

    x, y = coords
    new_x = np.cos(-rotation) * x + np.sin(-rotation) * y
    new_y = -np.sin(-rotation) * x + np.cos(-rotation) * y
    return np.array([new_x, new_y])

def gen_powers(degree: int):
    """
    Generates the powers required for a 2d polynomial

    Parameters
    ----------
    degree : int
        Maximum power to generate

    Returns
    -------
    xpows : Array
        x axis powers
    ypows : Array
        y axis powers
    """
    n = triangular_number(degree)
    vals = np.arange(n)

    # Ypows
    tris = triangular_number(np.arange(degree))
    ydiffs = np.repeat(tris, np.arange(1, degree + 1))
    ypows = vals - ydiffs

    # Xpows
    tris = triangular_number(np.arange(1, degree + 1))
    xdiffs = np.repeat(n - np.flip(tris), np.arange(degree, 0, -1))
    xpows = np.flip(vals - xdiffs)

    return np.array([xpows, ypows])

def distort_coords(coords: Array, coeffs: Array, pows: Array):
    """
    Apply a 2D polynomial distortion to some coordinates

    Parameters
    ----------
    coords : Array
        Input coordinates to distort
    coeffs : Array
        Distortion polynomial coefficients
    pows : Array
        Distortion polynomial powers

    Returns
    -------
    distorted_coords : Array
        Coords with the distortion applied
    """
    pow_base = np.multiply(*(coords[:, None, ...] ** pows[..., None, None]))
    distortion = np.sum(coeffs[..., None, None] * pow_base[None, ...], axis=1)
    return coords + distortion