colormaps.py

'''
Author: Yotam Gingold <yotam (strudel) yotamgingold.com>
License: Public Domain [CC0](http://creativecommons.org/publicdomain/zero/1.0/)
'''

from __future__ import print_function, division

from numpy import *
from skimage import color

def univariate_with_white( ts, rgb_color ):
    '''
    Given a numpy.array of values `ts` between 0 and 1,
    and an RGB color `rgb_color` as a triplet of three numbers between 0 and 1,
    returns the RGB color triplets that linearly blend between `rgb_color` and a
    grey color according to `ts`, where `rgb_color` corresponds to t=1.
    The grey color is chosen to be as bright as possible without producing colors outside
    the space of valid colors.
    The shape of the result is `t.shape + (3,)`.
    
    NOTE: If `rgb_color` is white, the result will be constant for all `ts`.
    
    NOTE: RGB values are linear sRGB.
    
    If you want to choose the second color manually, you can try
    `univariate_between_two_colors_lab()` or `univariate_between_two_colors_lch()` below.
    Use `univariate_between_two_colors_lab()` if you want the second color to be a grey
    or don't care about hue shifts (and don't pass `equal_luminance = True`).
    You may have to experiment to find a grey color that doesn't result
    in out-of-range colors. (The function will print a warning if you do.)
    You can look for legal endpoints on this page. Scroll down to
    "Now, let's grab some colors", switch it to C/L, adjust the hue on the slider to the
    right, and then drag the endpoints around.
        http://vis4.net/blog/posts/avoid-equidistant-hsv-colors/
    NOTE: If you want the second color to be something other than grey, you should use
          `univariate_between_two_colors_lch()`.
    '''
    
    rgb_color0 = brightest_white_with_a_valid_line_to_rgb_color( rgb_color )
    return univariate_between_two_colors_lab( ts, rgb_color0, rgb_color )

def univariate_with_black( ts, rgb_color ):
    '''
    Given a numpy.array of values `ts` between 0 and 1,
    and an RGB color `rgb_color` as a triplet of three numbers between 0 and 1,
    returns the RGB color triplets that linearly blend between `rgb_color` and
    black according to `ts`, where `rgb_color` corresponds to t=1.
    The shape of the result is `t.shape + (3,)`.
    
    NOTE: If `rgb_color` is black, the result will be constant for all `ts`.
    
    NOTE: RGB values are linear sRGB.
    
    NOTE: This will never produce invalid colors that have to be clipped.
    
    If you want to choose the second color manually, you can try
    `univariate_between_two_colors_lab()` or `univariate_between_two_colors_lch()` below.
    Use `univariate_between_two_colors_lab()` if you want the second color to be a grey
    or don't care about hue shifts (and don't pass `equal_luminance = True`).
    You may have to experiment to find a grey color that doesn't result
    in out-of-range colors. (The function will print a warning if you do.)
    You can look for legal endpoints on this page. Scroll down to
    "Now, let's grab some colors", switch it to C/L, adjust the hue on the slider to the
    right, and then drag the endpoints around.
        http://vis4.net/blog/posts/avoid-equidistant-hsv-colors/
    NOTE: If you want the second color to be something other than grey, you should use
          `univariate_between_two_colors_lch()`.
    '''
    
    return univariate_between_two_colors_lab( ts, (0,0,0), rgb_color )

def univariate_with_grey( ts, rgb_color ):
    '''
    Given a numpy.array of values `ts` between 0 and 1,
    and an RGB color `rgb_color` as a triplet of three numbers between 0 and 1,
    returns the RGB color triplets that linearly blend between `rgb_color` and a
    grey color with the same luminance as `rgb_color` according to `ts`,
    where `rgb_color` corresponds to t=1.
    The shape of the result is `t.shape + (3,)`.
    
    Effectively, this colormap has all its variation in the hue and saturation.
    This is not necessarily a good thing, since our brightness perception is
    much stronger.
    Try `univariate_with_white()` above.
    
    NOTE: If `rgb_color` is grey, the result will be constant for all `ts`.
    
    NOTE: RGB values are linear sRGB.
    
    If you want to choose the second color manually, you can try
    `univariate_between_two_colors_lab()` or `univariate_between_two_colors_lch()` below.
    Use `univariate_between_two_colors_lab()` if you want the second color to be a grey
    or don't care about hue shifts (and don't pass `equal_luminance = True`).
    You may have to experiment to find a grey color that doesn't result
    in out-of-range colors. (The function will print a warning if you do.)
    You can look for legal endpoints on this page. Scroll down to
    "Now, let's grab some colors", switch it to C/L, adjust the hue on the slider to the
    right, and then drag the endpoints around.
        http://vis4.net/blog/posts/avoid-equidistant-hsv-colors/
    NOTE: If you want the second color to be something other than grey, you should use
          `univariate_between_two_colors_lch()`.
    '''
    
    ### 1 Convert `rgb_color` into a Lab-space color.
    ### 2 Scale the ab channel by `ts`.
    ### 3 Convert the result to RGB-space.
    
    ts = asfarray(ts)
    rgb_color = asfarray( rgb_color )
    
    assert asarray( rgb_color ).shape == (3,)
    
    ### 1
    lab_color = color.rgb2lab( rgb_color.reshape(1,1,3) )[0,0]
    
    ### 2
    ## Flatten t. Store the original shape first.
    ts_original_shape = ts.shape
    ts = ts.ravel()
    N = len(ts)
    
    ## Repeat lab_color so that it has shape Nx3.
    result = tile( lab_color[newaxis,:], (N,1) )
    
    ## Scale the ab channel by t.
    result[:,1:] *= ts[:,newaxis]
    
    ### 3
    ## Clip to 0,1
    result_rgb = color.lab2rgb( result[newaxis,...] ).clip(0,1)
    compare( color.rgb2lab( result_rgb )[0], result )
    
    ## Reshape the result
    result_rgb = result_rgb.reshape( ts_original_shape + (3,) )
    
    return result_rgb

def univariate_between_two_colors_lab( ts, rgb_color0, rgb_color1, equal_luminance = None ):
    '''
    Given a numpy.array of values `ts` between 0 and 1,
    and an RGB color for the 0 value `rgb_color0` as a triplet of three numbers between 0 and 1,
    and an RGB color for the 1 value `rgb_color1` as a triplet of three numbers between 0 and 1,
    returns the RGB color triplets that linearly blend between `rgb_color0` and `rgb_color1`
    according to `ts`.
    The shape of the result is `t.shape + (3,)`.
    
    If the optional parameter `equal_luminance` is True, the resulting colors will all have
    the luminance of the darker of the two colors.
    (The darker luminance is chosen so that the interpolation won't result in RGB colors
    outside [0,1].)
    This is not necessarily a good thing, since humans perceive luminance the best.
    But it will almost surely avoid the interpolation passing outside of the space
    of valid colors.
    
    You can look for legal endpoints on this page. Scroll down to
    "Now, let's grab some colors", switch it to C/L, adjust the hue on the slider to the
    right, and then drag the endpoints around.
        http://vis4.net/blog/posts/avoid-equidistant-hsv-colors/
    NOTE: If you want the second color to be something other than grey, you should use
          `univariate_between_two_colors_lch()`.
    
    NOTE: When neither of `rgb_color0` and `rgb_color1` are grey
          (including white and black), the interpolated values will
          be desaturated (greyer). To avoid this problem,
          use univariate_between_two_colors_lch() instead.
    
    NOTE: RGB values are linear sRGB.
    '''
    
    if equal_luminance is None: equal_luminance = False
    
    ### 1 Convert `rgb_color0` and `rgb_color1` into Lab-space colors.
    ### 2 Linearly interpolate them according to `ts`.
    ### 3 Convert the result to RGB-space.
    
    ts = asfarray(ts)
    rgb_color0 = asfarray( rgb_color0 )
    rgb_color1 = asfarray( rgb_color1 )
    
    assert asarray( rgb_color0 ).shape == (3,)
    assert asarray( rgb_color1 ).shape == (3,)
    
    ### 1
    lab_color0 = color.rgb2lab( rgb_color0.reshape(1,1,3) )[0,0]
    lab_color1 = color.rgb2lab( rgb_color1.reshape(1,1,3) )[0,0]
    if equal_luminance:
        Lmin = min( lab_color0[0], lab_color1[0] )
        lab_color0[0] = Lmin
        lab_color1[0] = Lmin
    
    ### 2
    ## Flatten t. Store the original shape first.
    ts_original_shape = ts.shape
    ts = ts.ravel()
    N = len(ts)
    
    ## Linearly interpolate
    result = lab_color0[newaxis,:] + ts[:,newaxis]*(lab_color1[newaxis,:] - lab_color0[newaxis,:])
    
    ### 3
    ## Clip to 0,1
    result_rgb = color.lab2rgb( result[newaxis,...] ).clip(0,1)
    compare( color.rgb2lab( result_rgb )[0], result )
    
    ## Reshape the result
    result_rgb = result_rgb.reshape( ts_original_shape + (3,) )
    
    return result_rgb

def univariate_between_two_colors_lch( ts, rgb_color0, rgb_color1, equal_luminance = None ):
    '''
    Given a numpy.array of values `ts` between 0 and 1,
    and an RGB color for the 0 value `rgb_color0` as a triplet of three numbers between 0 and 1,
    and an RGB color for the 1 value `rgb_color1` as a triplet of three numbers between 0 and 1,
    returns the RGB color triplets that linearly blend between `rgb_color0` and `rgb_color1`
    according to `ts`.
    The shape of the result is `t.shape + (3,)`.
    
    If the optional parameter `equal_luminance` is True, the resulting colors will all have
    the luminance of the darker of the two colors.
    (The darker luminance is chosen so that the interpolation won't result in RGB colors
    outside [0,1].)
    This is not necessarily a good thing, since humans perceive luminance the best.
    But it will almost surely avoid the interpolation passing outside of the space
    of valid colors.
    
    You can look for legal endpoints on this page. Scroll down to
    "Now, let's grab some colors", switch it to C/L, adjust the hue on the slider to the
    right, and then drag the endpoints around.
        http://vis4.net/blog/posts/avoid-equidistant-hsv-colors/
    NOTE: If you want the second color to be grey, you should use
          `univariate_between_two_colors_lab()`.
    
    NOTE: This function should not be used with a grey (including white and black),
          since the hue of grey is undefined.
    
    NOTE: RGB values are linear sRGB.
    '''
    
    if equal_luminance is None: equal_luminance = False
    
    ### 1 Convert `rgb_color0` and `rgb_color1` into Lab-space colors.
    ### 2 Linearly interpolate them according to `ts`.
    ### 3 Convert the result to RGB-space.
    
    ts = asfarray(ts)
    rgb_color0 = asfarray( rgb_color0 )
    rgb_color1 = asfarray( rgb_color1 )
    
    assert asarray( rgb_color0 ).shape == (3,)
    assert asarray( rgb_color1 ).shape == (3,)
    
    ### 1
    lch_color0 = color.lab2lch( color.rgb2lab( rgb_color0.reshape(1,1,3) ) )[0,0]
    lch_color1 = color.lab2lch( color.rgb2lab( rgb_color1.reshape(1,1,3) ) )[0,0]
    if equal_luminance:
        Lmin = min( lch_color0[0], lch_color1[0] )
        lch_color0[0] = Lmin
        lch_color1[0] = Lmin
    
    ### 2
    ## Flatten t. Store the original shape first.
    ts_original_shape = ts.shape
    ts = ts.ravel()
    N = len(ts)
    
    ## Linearly interpolate
    def lerp( b, a, t ): return b + t*(a - b)
    
    result = lerp( lch_color0[newaxis,:], lch_color1[newaxis,:], ts[:,newaxis] )
    ## Hue is in radians [0,2*pi]. We have to consider the before and after colors being
    ## more than pi apart.
    if lch_color1[2] - lch_color0[2] > pi:
        result[:,2] = fmod( lerp( lch_color0[2] + 2*pi, lch_color1[2], ts ), 2*pi )
    elif lch_color0[2] - lch_color1[2] > pi:
        result[:,2] = fmod( lerp( lch_color0[2], lch_color1[2] + 2*pi, ts ), 2*pi )
    
    ## There are only two colors. This general version is unnecessary.
    #mark_result0_plus_2pi = ( result1[:,2] - result0[:,2] > pi )
    #mark_result1_plus_2pi = ( result0[:,2] - result1[:,2] > pi )
    #result[:,2][ mark_result0_plus_2pi ] = fmod( lerp( result0[:,2] + 2*pi, result1[:,2], ts[:,newaxis] ), 2*pi )[ mark_result0_plus_2pi ]
    #result[:,2][ mark_result1_plus_2pi ] = fmod( lerp( result0[:,2], result1[:,2] + 2*pi, ts[:,newaxis] ), 2*pi )[ mark_result1_plus_2pi ]
    
    if equal_luminance:
        result[:,0] = min( lch_color0[0], lch_color1[0] )
    
    ### 3
    result_rgb = color.lab2rgb( color.lch2lab( result[newaxis,...] ) ).clip(0,1)
    compare( color.rgb2lab( result_rgb )[0], color.lch2lab( result[newaxis,...] )[0] )
    
    ## Reshape the result
    result_rgb = result_rgb.reshape( ts_original_shape + (3,) )
    
    return result_rgb

def compare( lab0, lab1 ):
    assert len( lab0.shape ) == 2 and lab0.shape[1] == 3
    assert len( lab1.shape ) == 2 and lab1.shape[1] == 3
    
    diff = lab0 - lab1
    diff = sqrt( ( diff**2 ).sum( axis = 1 ) )
    
    ## The just-noticeable difference in Lab-space is around 2.3
    if diff.max() < 2: return
    
    print( "Clipping: max Lab-space difference:", diff.max(), "and average Lab-space difference:", average( diff ), "(just-noticeable difference is ~2.3)" )

def brightest_white_with_a_valid_line_to_rgb_color( rgb_color ):
    '''
    Given an RGB color `rgb_color` as a triplet of three numbers between 0 and 1,
    returns the brightest white RGB color such that a line to `rgb_color` lies
    entirely within the space of valid colors.
    '''
    
    ### 1 Convert `rgb_color` to Lab space.
    ### 2 Create all greys whose luminance goes from `rgb_color`'s to 1.
    ### 3 Create an image whose first column are the greys, last column is `rgb_color`,
    ###   and whose intermediate columns blend.
    ### 4 Convert the image to rgb space.
    ### 5 Find the bottommost row whose values are all in [0,1].
    
    rgb_color = asfarray( rgb_color )
    
    assert asarray( rgb_color ).shape == (3,)
    
    ### 1
    lab_color = color.rgb2lab( rgb_color.reshape(1,1,3) )[0,0]
    
    
    ### 2
    num_greys = 101
    greys = zeros( (num_greys,3), dtype = lab_color.dtype )
    greys[:,0] = linspace( lab_color[0], 100., num_greys )
    
    
    ### 3
    num_blends = 100
    ts = linspace( 0, 1, num_blends )
    
    ## Linearly interpolate between `lab_color` and `greys`
    blends = greys[:,newaxis,:] + ts[newaxis,:,newaxis]*(lab_color[newaxis,newaxis,:] - greys[:,newaxis,:])
    
    
    ### 4
    blends_rgb = color.lab2rgb( blends )
    
    
    ### 5
    eps = 1e-8
    rows_that_didnt_clip = logical_and( blends_rgb > -eps, blends_rgb < 1 + eps ).all( axis = 2 ).all( axis = 1 )
    brightest_index = max( where( rows_that_didnt_clip )[0] )
    
    lab_result = greys[ brightest_index ]
    return color.lab2rgb( lab_result[newaxis,newaxis,:] )[0,0]

def test_with_white( color = None, name = None ):
    ## A 1D linear ramp 1000-pixels long.
    ts = linspace( 0, 1, 1000 )
    result = univariate_with_white( ts, color )
    import skimage.io
    ## Save it as an image, repeating each value 100 times in each column of the image.
    outname = "colormaps_test_%s_to_white.png" % name
    skimage.io.imsave( outname, tile( result[newaxis,:,:], ( 100, 1, 1 ) ) )
    print( "Saved:", outname )

def test_with_black( color = None, name = None ):
    ## A 1D linear ramp 1000-pixels long.
    ts = linspace( 0, 1, 1000 )
    result = univariate_with_black( ts, color )
    import skimage.io
    ## Save it as an image, repeating each value 100 times in each column of the image.
    outname = "colormaps_test_%s_to_black.png" % name
    skimage.io.imsave( outname, tile( result[newaxis,:,:], ( 100, 1, 1 ) ) )
    print( "Saved:", outname )

def test_with_grey( color = None, name = None ):
    ## A 1D linear ramp 1000-pixels long.
    ts = linspace( 0, 1, 1000 )
    result = univariate_with_grey( ts, color )
    import skimage.io
    ## Save it as an image, repeating each value 100 times in each column of the image.
    outname = "colormaps_test_%s_to_grey.png" % name
    skimage.io.imsave( outname, tile( result[newaxis,:,:], ( 100, 1, 1 ) ) )
    print( "Saved:", outname )

def test_two_colors_lab( color0 = None, name0 = None, color1 = None, name1 = None, equal_luminance = None ):
    ## A 1D linear ramp 1000-pixels long.
    ts = linspace( 0, 1, 1000 )
    result = univariate_between_two_colors_lab( ts, color0, color1, equal_luminance )
    import skimage.io
    ## Save it as an image, repeating each value 100 times in each column of the image.
    outname = "colormaps_test_%s_to_%s_lab%s.png" % ( name0, name1, '' if not equal_luminance else '_equal_luminance' )
    skimage.io.imsave( outname, tile( result[newaxis,:,:], ( 100, 1, 1 ) ) )
    print( "Saved:", outname )

def test_two_colors_lch( color0 = None, name0 = None, color1 = None, name1 = None, equal_luminance = None ):
    ## A 1D linear ramp 1000-pixels long.
    ts = linspace( 0, 1, 1000 )
    result = univariate_between_two_colors_lch( ts, color0, color1, equal_luminance )
    import skimage.io
    ## Save it as an image, repeating each value 100 times in each column of the image.
    outname = "colormaps_test_%s_to_%s_lch%s.png" % ( name0, name1, '' if not equal_luminance else '_equal_luminance' )
    skimage.io.imsave( outname, tile( result[newaxis,:,:], ( 100, 1, 1 ) ) )
    print( "Saved:", outname )

if __name__ == '__main__':
     test_with_white( color = (1,0,0), name = 'red' )
     test_with_black( color = (1,0,0), name = 'red' )
     test_with_grey( color = (1,0,0), name = 'red' )
     test_with_white( color = (0,1,0), name = 'green' )
     test_with_black( color = (0,1,0), name = 'green' )
     test_with_grey( color = (0,1,0), name = 'green' )
     test_with_white( color = (0,0,1), name = 'blue' )
     test_with_black( color = (0,0,1), name = 'blue' )
     test_with_grey( color = (0,0,1), name = 'blue' )
     
     test_two_colors_lab( color0 = (1,1,1), name0 = 'white', color1 = (1,0,0), name1 = 'red' )
     test_two_colors_lab( color0 = (1,1,1), name0 = 'white', color1 = (1,0,0), name1 = 'red', equal_luminance = True )
     test_two_colors_lab( color0 = (1,1,1), name0 = 'white', color1 = (0,1,0), name1 = 'green' )
     test_two_colors_lab( color0 = (1,1,1), name0 = 'white', color1 = (0,1,0), name1 = 'green', equal_luminance = True )
     test_two_colors_lab( color0 = (1,1,1), name0 = 'white', color1 = (0,0,1), name1 = 'blue' )
     test_two_colors_lab( color0 = (1,1,1), name0 = 'white', color1 = (0,0,1), name1 = 'blue', equal_luminance = True )
     
     test_two_colors_lab( color0 = (1,0,0), name0 = 'red', color1 = (0,0,1), name1 = 'blue' )
     test_two_colors_lch( color0 = (1,0,0), name0 = 'red', color1 = (0,0,1), name1 = 'blue' )
     test_two_colors_lab( color0 = (1,0,0), name0 = 'red', color1 = (0,0,1), name1 = 'blue', equal_luminance = True )
     test_two_colors_lch( color0 = (1,0,0), name0 = 'red', color1 = (0,0,1), name1 = 'blue', equal_luminance = True )
     
     test_two_colors_lab( color0 = (.6,.2,.2), name0 = 'reddish', color1 = (.2,.2,.6), name1 = 'blueish' )
     test_two_colors_lch( color0 = (.6,.2,.2), name0 = 'reddish', color1 = (.2,.2,.6), name1 = 'blueish' )
     test_two_colors_lab( color0 = (.6,.2,.2), name0 = 'reddish', color1 = (.2,.2,.6), name1 = 'blueish', equal_luminance = True )
     test_two_colors_lch( color0 = (.6,.2,.2), name0 = 'reddish', color1 = (.2,.2,.6), name1 = 'blueish', equal_luminance = True )
     
     test_two_colors_lab( color0 = (.9,.9,.1), name0 = 'yellow', color1 = (.9,.1,.9), name1 = 'magenta' )
     test_two_colors_lch( color0 = (.9,.9,.1), name0 = 'yellow', color1 = (.9,.1,.9), name1 = 'magenta' )
     test_two_colors_lab( color0 = (.9,.9,.1), name0 = 'yellow', color1 = (.9,.1,.9), name1 = 'magenta', equal_luminance = True )
     test_two_colors_lch( color0 = (.9,.9,.1), name0 = 'yellow', color1 = (.9,.1,.9), name1 = 'magenta', equal_luminance = True )