Color functions

color_function.py

from colorsys import hls_to_rgb, rgb_to_hls
import math
import cmath

CLAMP = lambda y: max(0.0, min(y, 1.0))
BLEND = lambda x, y: 0.5*x + 0.5*y
DODGE = lambda a, b: a / (1.0 - b + 1/256.0)

# gimp color balance algorithm
def balance_channel(value, l, shadows, midtones, highlights):
    a = 0.25
    b = 0.333
    scale = 0.7
    shadows    *= CLAMP((l - b) / (-a) + 0.5) * scale
    midtones   *= CLAMP((l - b) / ( a) + 0.5) * CLAMP((l + b - 1.0) / (-a) + 0.5) * scale
    highlights *= CLAMP((l + b - 1.0) / ( a) + 0.5) * scale
    value += shadows
    value += midtones
    value += highlights
    return CLAMP(value)

def balance(R, G, B, ra, rb, rc, ga, gb, gc, ba, bb, bc):
    h, l, s = rgb_to_hls(R, G, B)
    R = balance_channel(R, R, ra, rb, rc)
    G = balance_channel(G, G, ga, gb, gc)
    B = balance_channel(B, B, ba, bb, bc)
    # preserve lightness
    h2, l2, s2 = rgb_to_hls(R, G, B)
    return hls_to_rgb(h2, l, s2)

blue_orange_colors = [
  (-1.0,  0.0, 0.0, 0.0 ),
  (-0.95, 0.1, 0.2, 0.5 ),
  (-0.5,  0.0, 0.5, 1.0 ),
  (-0.05, 0.4, 0.8, 0.8 ),
  ( 0.0,  1.0, 1.0, 1.0 ),
  ( 0.05, 1.0, 0.9, 0.3 ),
  ( 0.5,  0.9, 0.5, 0.0 ),
  ( 0.95, 0.7, 0.1, 0.0 ),
  ( 1.0,  0.0, 0.0, 0.0 ),
  ( 2.0,  0.0, 0.0, 0.0 ),
]

def color_function(z, mode=0):
    z = complex(z)
    if cmath.isinf(z):
        return (1.0, 1.0, 1.0)
    if cmath.isnan(z):
        return (0.5, 0.5, 0.5)
    if mode == 0:
        H = cmath.phase(z)
        PI = math.pi
        H = (H + PI) / (2 * PI) + 0.5
        H = H - math.floor(H)
        t = abs(z)
        if t > 1e60:
            L = 1.0
        elif t < 1e-60:
            L = 0.0
        else:
            L = 1.0 - 1.0/(1.0 + t ** 0.2);
        S = 0.8
        return hls_to_rgb(H, L, S)
    if mode == 1:
        H = cmath.phase(z)
        PI = math.pi
        H = H / PI
        H = max(min(H, 1.0), -1.0)
        i = 1
        while 1:
            if blue_orange_colors[i][0] > H:
                a  = blue_orange_colors[i-1][0]
                ra = blue_orange_colors[i-1][1]
                ga = blue_orange_colors[i-1][2]
                ba = blue_orange_colors[i-1][3]
                b  = blue_orange_colors[i][0]
                rb = blue_orange_colors[i][1]
                gb = blue_orange_colors[i][2]
                bb = blue_orange_colors[i][3]
                s = (H - a) / (b - a)
                R = ra + (rb - ra) * s
                G = ga + (gb - ga) * s
                B = ba + (bb - ba) * s
                return R, G, B
            i += 1
    assert 2 <= mode <= 6
    R1, G1, B1 = color_function(z, 0)
    R2, G2, B2 = color_function(z, 1)
    R = BLEND(R1, CLAMP(DODGE(R1, R2)))
    G = BLEND(G1, CLAMP(DODGE(G1, G2)))
    B = BLEND(B1, CLAMP(DODGE(B1, B2)))
    if mode == 3:
        return balance(R, G, B, 0.0, -0.5, 0.2, 0.0, 0.0, -0.1, 0.0, -1.0, -0.2)
    elif mode == 4:
        return balance(R, G, B, 0.0, -0.5, 0.2, 0.0, 0.5, -0.1, 0.0, -0.3, -1.0)
    elif mode == 5:
        return balance(R, G, B, 0.0, -0.5, -1.0, 0.0, -0.1, -0.67, 0.0, -0.55, -0.12)
    elif mode == 6:
        return balance(R, G, B, 0.86, 0.0, 0.13, 0.57, 0.19, -0.52, 0.31, -0.30, -0.94)
    else:
        return R, G, B

if __name__ == "__main__":
    import numpy
    import mpmath
    import matplotlib.pyplot as plt

    N = 512

    def cplot(f, xa, xb, ya, yb, mode):
        Z = numpy.zeros((N, N, 3))
        for i, x in enumerate(numpy.linspace(xa, xb, N)):
            for j, y in enumerate(numpy.linspace(ya, yb, N)):
                Z[j,i] = color_function(f(x+y*1j), mode=mode)
        imgplot = plt.imshow(Z, origin='lower', interpolation='nearest')
        plt.axis('off')

    plt.subplot(221)
    f = lambda z: (mpmath.fp.qp(complex(z)**4) if abs(z) < 0.99 else mpmath.inf)
    cplot(f, -1, 1, -1, 1, mode=3)

    plt.subplot(222)
    f = lambda z: mpmath.fp.zeta(complex(z))
    cplot(f, -15, 15, -15, 15, mode=4)

    plt.subplot(223)
    f = lambda z: mpmath.fp.gamma(complex(z))
    cplot(f, -5, 5, -5, 5, mode=4)

    plt.subplot(224)
    f = lambda z: mpmath.fp.erf(complex(z))
    cplot(f, -3, 3, -3, 3, mode=6)

    plt.tight_layout()
    #plt.show()
    plt.savefig("demo.png", dpi=500)

I almost vectorized it. 😆