Contour a 2D array based on area

area_contours.py

import numpy as np
import matplotlib.pyplot as plt

def area_contour_image(img: np.ndarray, quantiles: list[float], ax=None, x=None, y=None, **contour_kwds):
    normalized_img = (img - np.min(img)) / (np.max(img) - np.min(img))
    flattened_img = normalized_img.flatten()

    # sort pixels: small areas first, then large areas
    sorted_values = np.sort(flattened_img)[::-1]

    cumulative_sum = np.cumsum(sorted_values)
    cumulative_sum /= cumulative_sum.max()

    # find indices in the cumulative sum that correpond to the levels specified
    indices = np.searchsorted(cumulative_sum, quantiles)

    # Threshold values for the desired percentages and draw contours
    thresholds = np.sort(sorted_values[indices])
    if x is None or y is None:
        return ax.contour(normalized_img, levels=thresholds, **contour_kwds)
    else:
        return ax.contour(x, y, normalized_img, levels=thresholds, **contour_kwds)

def test():
    x = np.linspace(-10, 10, num=1000)
    y = x.copy()
    xx, yy = np.meshgrid(x, y)
    
    # sort of gaussian
    sigx, sigy = 2, 5
    z = 5 * np.exp(-(xx**2 / sigx**2 + yy**2 / sigy**2))
    fig, ax = plt.subplots(layout='constrained', figsize=(6, 6))

    ax.pcolormesh(xx, yy, z, cmap='viridis')
    # optionally use the return value
    contours = area_contour_image(z, [0.05, 0.3, 0.5, 0.8], ax=ax, x=x, y=y, cmap='Reds')

    ax.set(xlabel='x', ylabel='y', title='Area contours')
    plt.show()


if __name__ == '__main__':
    test()