Misc¤

def soft_binarise(array: Array, oversample: int = 3) -> Array:
    """
    Applies the CLIMB algorithm originally introduced in
    [Wong et al., 2021](https://ui.adsabs.harvard.edu/abs/2021JOSAB..38.2465W/abstract).
    Assumes a square input array.

    Parameters
    ----------
    array: Array
        The continuous array to binarise
    oversample: int = 3
        The oversample of the input array

    Returns
    -------
    soft_binary: Array
        The binarised array with soft edges
    """
    # Make sure the array can be downsampled
    if array.shape[-1] % oversample != 0:
        raise ValueError("Array shape must be a multiple of oversample")
    n_tile = array.shape[-1] // oversample

    # Reshape the array into a vector of unit cells
    tiles = array.reshape(n_tile, oversample, n_tile, oversample).transpose(0, 2, 1, 3)
    flat_tiles = tiles.reshape(-1, oversample, oversample)  # (N,b,b)
    flat_pixels = flat_tiles.reshape(flat_tiles.shape[0], -1)  # (N,P)

    # Calculate the fractional area
    M = _lsq_matrix(oversample)  # (3,P)
    coeffs = flat_pixels @ M.T  # (N,3)
    frac = vmap(_calc_area_fraction)(coeffs)  # (N,)

    # Assign the sign of the area using the mean value of the patch
    means = np.mean(flat_tiles, axis=(1, 2))
    frac = np.maximum(frac, 1 - frac)
    frac = np.where(means >= 0, frac, 1 - frac)

    # Get a mask of the positive, negative and boundary regions
    all_pos = np.all(flat_tiles > 0, axis=(1, 2))
    all_nonpos = np.all(flat_tiles <= 0, axis=(1, 2))
    any_zero = np.any(flat_tiles == 0, axis=(1, 2))

    # Map the values to the different regions
    frac = np.where(all_pos, 1.0, frac)
    frac = np.where(all_nonpos, 0.0, frac)
    frac = np.where(any_zero, (frac > 0).astype(frac.dtype), frac)

    # Clip values between 0 and 1 and reshape
    return np.clip(frac, 0, 1).reshape(n_tile, n_tile)