nice colorgradients using SR2Lab (or OKLab)

color_fun.py

import numpy as np
from matplotlib import colors
import matplotlib.pyplot as plt


def lch_to_lab(LCh):
    a = np.cos(LCh[2] * np.pi / 180) * LCh[1]
    b = np.sin(LCh[2] * np.pi / 180) * LCh[1]
    return np.array((LCh[0], a, b))


def lab_to_lch(Lab):
    C = np.sqrt(Lab[1] ** 2 + Lab[2] ** 2)
    h = np.arctan2(Lab[2], Lab[1]) * 180 / np.pi
    return np.array((Lab[0], C, h))


def linear_srgb_to_oklab(rgb):
    rgb2lms = np.array(
        [
            [0.4122214708, 0.5363325363, 0.0514459929],
            [0.2119034982, 0.6806995451, 0.1073969566],
            [0.0883024619, 0.2817188376, 0.6299787005],
        ]
    )
    lms = rgb2lms @ rgb
    lms2oklab = np.array(
        [
            [0.2104542553, 0.7936177850, -0.0040720468],
            [1.9779984951, -2.4285922050, 0.4505937099],
            [0.0259040371, 0.7827717662, -0.8086757660],
        ]
    )
    return lms2oklab @ np.cbrt(lms)


def oklab_to_linear_srgb(Lab):
    oklab2lms_ = np.array(
        [
            [1, 0.3963377774, 0.2158037573],
            [1, -0.1055613458, -0.0638541728],
            [1, -0.0894841775, -1.2914855480],
        ]
    )
    lms = np.power(oklab2lms_ @ Lab, 3)
    lms2srgb = np.array(
        [
            [4.0767416621, -3.3077115913, 0.2309699292],
            [-1.2684380046, 2.6097574011, -0.3413193965],
            [-0.0041960863, -0.7034186147, 1.7076147010],
        ]
    )
    return lms2srgb @ lms


# SR2LAB from https://gist.github.com/jjrv/b27d0840b4438502f9cad2a0f9edeabc
def linearize(x):
    return np.where(x <= 0.04045, x / 12.92, np.power((x + 0.055) / 1.055, 2.4))


def delinearize(x):
    return np.where(x <= 0.0031308, x * 12.92, np.power(x, 1 / 2.4) * 1.055 - 0.055)


def srcube(x):
    return np.where(x <= 8, x * 27 / 24389, np.power((x + 16) / 116, 3))


def srroot(x):
    return np.where(
        x <= 216 / 24389, x * 24389 / 2700, 1.16 * np.power(x, 1 / 3) - 0.16
    )


def srlab2rgb(Lab):
    M1 = np.array([[100, 9.04127, 4.56344], [100, -5.33159, -2.69178], [100, 0, -58]])
    xyz = srcube(M1 @ Lab)
    M2 = np.array(
        [
            [5.435679, -4.599131, 0.163593],
            [-1.16809, 2.327977, -0.159798],
            [0.03784, -0.198564, 1.160644],
        ]
    )
    return delinearize(M2 @ xyz)


def rgb2srlab(rgb):
    M1 = np.array(
        [
            [0.32053, 0.63692, 0.04256],
            [0.161987, 0.756636, 0.081376],
            [0.017228, 0.10866, 0.874112],
        ]
    )
    xyz = srroot(M1 @ linearize(rgb))
    M2 = np.array(
        [
            [0.37095, 0.629054, -0.000008],
            [6.634684, -7.505078, +0.870328],
            [0.639569, 1.084576, -1.724152],
        ]
    )
    return M2 @ xyz


# to_rgb = lambda c: oklab_to_linear_srgb(lch_to_lab(c))
# to_lch = lambda c: lab_to_lch(linear_srgb_to_oklab(c))

to_rgb = lambda c: srlab2rgb(lch_to_lab(c))
to_lch = lambda c: lab_to_lch(rgb2srlab(c))


def hue_map(h):
    if not isinstance(h, str):
        return h
    colordict = {
        " ": np.nan,
        ".": np.nan,
        "r": 0,
        "g": 120,
        "b": 240,
        "c": 180,
        "m": 300,
        "y": 60,
        "R": 360,
        "G": 480,
        "B": 240 + 360,
        "C": 180 + 360,
        "M": 300 + 360,
        "Y": 60 + 360,
    }
    return [colordict[k] for k in list(h)]


def N_expand(y, N):
    y = np.ascontiguousarray(y)
    keep = ~np.isnan(y)
    x = np.cumsum(np.where(keep, 1, -0.9999999))
    return np.interp(np.linspace(x[0], x[-1], N), x[keep], y[keep])


def colorgradient(
    L=[0.1, 1], C=[0.7, 0], h="bc", N=150, L_adjust=0.1, C_adjust=0.1, name="bc"
):
    # L: luma
    # C: chroma (saturation)
    # h: hue (0-360)
    L = N_expand(L, N)
    C = N_expand(C, N)
    h = N_expand(hue_map(h), N) % 360
    for ii in range(15):
        cmap = np.array((L, C, h))
        rgb = to_rgb(cmap)
        maxrgb = np.max(rgb)
        minrgb = np.min(rgb)
        # fix clipping somewhat!
        if minrgb < 0:
            L = (L - minrgb * L_adjust) / (1 - minrgb * L_adjust)
            C = C * (1 - C_adjust)
        if maxrgb > 1:
            L = L * (1 - L_adjust)  # /(maxrgb*L_adjust+(1-L_adjust))
            C = C * (1 - C_adjust)
        else:
            # print(ii, np.mean(L), np.mean(C))
            break
    rgb = np.clip(rgb, 0, 1)  # finally clip
    cmap = colors.LinearSegmentedColormap.from_list(name, rgb.T)
    return cmap


def muted_blue(h="b"):
    return colorgradient(
        L=[0.9, 0.1], C=[0, 0.1, 0.3, 0.1, 0], h=h, L_adjust=0.01, name="muted_blue"
    )


def rainbow(L=0.5, C=0.5, h="rR"):
    return colorgradient(L=L, C=C, h="rR", L_adjust=0, name="rainbow")


def divergent(h="bc.yr", Cmax=0.7, Lmin=0.1):
    return colorgradient(
        L=[Lmin, 1, Lmin], C=[Cmax, 0, Cmax], h=h, L_adjust=0.01, name="divergent"
    )


def mimosa_to_viola(N=4):
    mimosa = to_lch(np.array([239 / 255, 192 / 255, 80 / 255]))
    viola = to_lch(np.array([166 / 255, 146 / 255, 186 / 255]))
    return colorgradient(
        [mimosa[0], viola[0]],
        [mimosa[1], viola[1]],
        [mimosa[2], viola[2] + 360],
        N=N,
        name="mimosa_to_viola",
    )