DAMASK_EICMD/python/damask/_colormap.py

import json
import functools

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import cm

from damask import Table

_eps   = 216./24389.
_kappa = 24389./27.
_ref_white = np.array([.95047, 1.00000, 1.08883])                                                   # Observer = 2, Illuminant = D65


class Colormap(mpl.colors.ListedColormap):


    @staticmethod
    def from_bounds(low,high,name='DAMASK colormap',N=256,model='rgb'):
        """
        Create a perceptually uniform colormap.

        Parameters
        ----------
        low : numpy.ndarray of shape (3)
        high : numpy.ndarray of shape (3)
        N : integer, optional
            The number of color quantization levels.
        name : str, optional
            The name of the colormap. Defaults to `DAMASK colormap`.
        model : str
            Colormodel used for low and high.

        """
        low_,high_ = map(np.array,[low,high])
        low_high = np.vstack((low_,high))
        if   model.lower() == 'rgb':
            if np.any(low_high<0) or np.any(low_high>1):
                raise ValueError(f'RGB color out of range {low} {high}.')

            low_,high_ = map(Colormap._rgb2msh,[low_,high_])

        elif model.lower() == 'hsv':
            if np.any(low_high<0) or np.any(low_high[:,1:3]>1) or np.any(low_high[:,0]>360):
                raise ValueError(f'HSV color out of range {low} {high}.')

            low_,high_ = map(Colormap._hsv2msh,[low_,high_])

        elif model.lower() == 'hsl':
            if np.any(low_high<0) or np.any(low_high[:,1:3]>1) or np.any(low_high[:,0]>360):
                raise ValueError(f'HSL color out of range {low} {high}.')

            low_,high_ = map(Colormap._hsl2msh,[low_,high_])

        elif model.lower() == 'xyz':

            low_,high_ = map(Colormap._xyz2msh,[low_,high_])

        elif model.lower() == 'lab':
            if np.any(low_high[:,0]<0):
                raise ValueError(f'CIE Lab color out of range {low} {high}.')

            low_,high_ = map(Colormap._lab2msh,[low_,high_])

        elif model.lower() == 'msh':
            pass

        else:
            raise ValueError(f'Invalid color model: {model}.')

        msh = map(functools.partial(Colormap._interpolate_msh,low=low_,high=high_),np.linspace(0,1,N))
        rgb = np.array(list(map(Colormap._msh2rgb,msh)))

        return Colormap(rgb,name=name)


    @staticmethod
    def from_predefined(name,N=256):
        """
        Select from set of predefined colormaps.

        Predefined colormaps include native matplotlib colormaps
        and common DAMASK colormaps.

        Parameters
        ----------
        name : str
            The name of the colormap.
        N : int, optional
           The number of color quantization levels. Defaults to 256.
           This parameter is not used for matplotlib colormaps
           that are of type `ListedColormap`.

        """
        # matplotlib presets
        for cat in Colormap._predefined_mpl:
            for n in cat[1]:
                if n == name:
                    colormap = cm.__dict__[name]
                    if isinstance(colormap,mpl.colors.LinearSegmentedColormap):
                        return Colormap(np.array(list(map(colormap,np.linspace(0,1,N)))),name=name)
                    else:
                        return Colormap(np.array(colormap.colors),name=name)

        # DAMASK presets
        definition = Colormap._predefined_DAMASK[name]
        return Colormap.from_bounds(definition['left'],definition['right'],name,N)


    @staticmethod
    def list_predefined():
        """List predefined colormaps by category."""
        print('DAMASK colormaps')
        print('  '+', '.join(Colormap._predefined_DAMASK.keys()))
        for cat in Colormap._predefined_mpl:
            print(f'{cat[0]}')
            print('  '+', '.join(cat[1]))


    def show(self):
        """Show colormap in window."""
        fig, ax = plt.subplots(figsize=(10,1))
        ax.set_axis_off()
        im = ax.imshow(np.broadcast_to(np.linspace(0,1,640).reshape(1,-1),(64,640)),cmap=self)      # noqa
        fig.canvas.set_window_title(self.name)
        plt.show()


    def reversed(self,name=None):
        """
        Make a reversed instance of the Colormap.

        Parameters
        ----------
        name : str, optional
            The name for the reversed colormap. If it's None
            the name will be the name of the parent colormap + "_r".

        Returns
        -------
        damask.Colormap
            The reversed colormap.

        """
        rev = super(Colormap,self).reversed(name)
        return Colormap(rev.colors,rev.name)


    def to_file(self,fname=None,format='paraview'):
        if fname is not None:
            try:
                f = open(fname,'w')
            except TypeError:
                f = fname
        else:
            f = None

        if format.lower() == 'paraview':
            Colormap._export_paraview(self,f)
        elif format.lower() == 'ascii':
            Colormap._export_ASCII(self,f)
        elif format.lower() == 'gom':
            Colormap._export_GOM(self,f)
        elif format.lower() == 'gmsh':
            Colormap._export_gmsh(self,f)
        else:
            raise ValueError('Unknown output format: {format}.')

    @staticmethod
    def _export_paraview(colormap,fhandle=None):
        """Write colormap to JSON file for Paraview."""
        colors = []
        for i,c in enumerate(np.round(colormap.colors,6).tolist()):
            colors+=[i]+c

        out = [{
                'ColorSpace':'RGB',
                'Name':colormap.name,
                'DefaultMap':True,
                'RGBPoints':colors
               }]
        if fhandle is None:
            with open(colormap.name.replace(' ','_')+'.json', 'w') as f:
                json.dump(out, f,indent=4)
        else:
            json.dump(out,fhandle,indent=4)

    @staticmethod
    def _export_ASCII(colormap,fhandle=None):
        """Write colormap to ASCII table."""
        labels = {'R':(1,),'G':(1,),'B':(1,)}
        if colormap.colors.shape[1] == 4: labels['alpha']=(1,)
        t = Table(colormap.colors,labels)

        if fhandle is None:
            with open(colormap.name.replace(' ','_')+'.txt', 'w') as f:
                t.to_ASCII(f)
        else:
            t.to_ASCII(fhandle)

    @staticmethod
    def _export_GOM(colormap,fhandle=None):
        pass
        s =(f'1 1 {colormap.name.replace(" ","_")} 9 {colormap.name.replace(" ","_")} '
            f' 0 1 0 3 0 0 -1 9 \\ 0 0 0 255 255 255 0 0 255 '
            f'30 NO_UNIT 1 1 64 64 64 255 1 0 0 0 0 0 0 3 0 {str(len(colormap.colors))}'
            ' '.join([' 0 %s 255 1'%(' '.join([str(int(x*255.0)) for x in color])) for color in reversed(colormap.colors)]))
        print(s)

    @staticmethod
    def _export_gmsh(colormap,fname=None):
        colors = colormap.colors
        colormap =('View.ColorTable = {'
                   ',\n'.join(['{%s}'%(','.join([str(x*255.0) for x in color])) for color in colors])+\
                   '}')

    @staticmethod
    def _interpolate_msh(frac,low, high):

        def rad_diff(a,b):
            return abs(a[2]-b[2])

        def adjust_hue(msh_sat, msh_unsat):
            """If saturation of one of the two colors is much less than the other, hue of the less."""
            if msh_sat[0] >= msh_unsat[0]:
               return msh_sat[2]
            else:
                hSpin = msh_sat[1]/np.sin(msh_sat[1])*np.sqrt(msh_unsat[0]**2.0-msh_sat[0]**2)/msh_sat[0]
                if msh_sat[2] < - np.pi/3.0: hSpin *= -1.0
                return msh_sat[2] + hSpin


        lo = np.array(low)
        hi = np.array(high)

        if (lo[1] > 0.05 and hi[1] > 0.05 and rad_diff(lo,hi) > np.pi/3.0):
            M_mid = max(lo[0],hi[0],88.0)
            if frac < 0.5:
                hi = np.array([M_mid,0.0,0.0])
                frac *= 2.0
            else:
                lo = np.array([M_mid,0.0,0.0])
                frac = 2.0*frac - 1.0
        if   lo[1] < 0.05 and hi[1] > 0.05:
            lo[2] = adjust_hue(hi,lo)
        elif lo[1] > 0.05 and hi[1] < 0.05:
            hi[2] = adjust_hue(lo,hi)

        return (1.0 - frac) * lo + frac * hi


    _predefined_mpl= [('Perceptually Uniform Sequential', [
                         'viridis', 'plasma', 'inferno', 'magma', 'cividis']),
                      ('Sequential', [
                         'Greys', 'Purples', 'Blues', 'Greens', 'Oranges', 'Reds',
                         'YlOrBr', 'YlOrRd', 'OrRd', 'PuRd', 'RdPu', 'BuPu',
                         'GnBu', 'PuBu', 'YlGnBu', 'PuBuGn', 'BuGn', 'YlGn']),
                      ('Sequential (2)', [
                         'binary', 'gist_yarg', 'gist_gray', 'gray', 'bone', 'pink',
                         'spring', 'summer', 'autumn', 'winter', 'cool', 'Wistia',
                         'hot', 'afmhot', 'gist_heat', 'copper']),
                      ('Diverging', [
                         'PiYG', 'PRGn', 'BrBG', 'PuOr', 'RdGy', 'RdBu',
                         'RdYlBu', 'RdYlGn', 'Spectral', 'coolwarm', 'bwr', 'seismic']),
                      ('Cyclic', ['twilight', 'twilight_shifted', 'hsv']),
                      ('Qualitative', [
                         'Pastel1', 'Pastel2', 'Paired', 'Accent',
                         'Dark2', 'Set1', 'Set2', 'Set3',
                         'tab10', 'tab20', 'tab20b', 'tab20c']),
                      ('Miscellaneous', [
                         'flag', 'prism', 'ocean', 'gist_earth', 'terrain', 'gist_stern',
                         'gnuplot', 'gnuplot2', 'CMRmap', 'cubehelix', 'brg',
                         'gist_rainbow', 'rainbow', 'jet', 'nipy_spectral', 'gist_ncar'])]

    _predefined_DAMASK = {'orientation':   {'low':   [0.933334,0.878432,0.878431],
                                            'right': [0.250980,0.007843,0.000000]},
                          'strain':        {'left':  [0.941177,0.941177,0.870588],
                                            'right': [0.266667,0.266667,0.000000]},
                          'stress':        {'left':  [0.878432,0.874511,0.949019],
                                            'right': [0.000002,0.000000,0.286275]}}

    @staticmethod
    def _hsv2rgb(hsv):
        """
        H(ue) S(aturation) V(alue) to R(red) G(reen) B(lue).

        References
        ----------
        https://www.rapidtables.com/convert/color/hsv-to-rgb.html

        """
        sextant = np.clip(int(hsv[0]/60.),0,5)
        c = hsv[1]*hsv[2]
        x = c*(1.0 - abs((hsv[0]/60.)%2 - 1.))

        return np.array([
                         [c, x, 0],
                         [x, c, 0],
                         [0, c, x],
                         [0, x, c],
                         [x, 0, c],
                         [c, 0, x],
                        ])[sextant] + hsv[2] - c

    @staticmethod
    def _rgb2hsv(rgb):
        """
        R(ed) G(reen) B(lue) to H(ue) S(aturation) V(alue).

        References
        ----------
        https://www.rapidtables.com/convert/color/rgb-to-hsv.html

        """
        C_max = rgb.max()
        C_min = rgb.min()
        Delta = C_max - C_min

        v = C_max
        s = 0. if np.isclose(C_max,0.) else Delta/C_max
        if np.isclose(Delta,0.):
            h = 0.
        elif rgb.argmax() == 0:
            h = (rgb[1]-rgb[2])/Delta%6
        elif rgb.argmax() == 1:
            h = (rgb[2]-rgb[0])/Delta + 2.
        elif rgb.argmax() == 2:
            h = (rgb[0]-rgb[1])/Delta + 4.

        h = np.clip(h,0.,6.) * 60.

        return np.array([h,s,v])


    @staticmethod
    def _hsl2rgb(hsl):
        """
        H(ue) S(aturation) L(uminance) to R(red) G(reen) B(lue).

        References
        ----------
        https://www.rapidtables.com/convert/color/hsl-to-rgb.html

        """
        sextant = np.clip(int(hsl[0]/60.),0,5)
        c = (1.0 - abs(2.0 * hsl[2] - 1.))*hsl[1]
        x = c*(1.0 - abs((hsl[0]/60.)%2 - 1.))
        m = hsl[2] - 0.5*c

        return np.array([
                         [c+m, x+m, m],
                         [x+m, c+m, m],
                         [m, c+m, x+m],
                         [m, x+m, c+m],
                         [x+m, m, c+m],
                         [c+m, m, x+m],
                        ])[sextant]

    @staticmethod
    def _rgb2hsl(rgb):
        """
        R(ed) G(reen) B(lue) to H(ue) S(aturation) L(uminance).

        References
        ----------
        https://www.rapidtables.com/convert/color/rgb-to-hsl.html

        """
        C_max = rgb.max()
        C_min = rgb.min()
        Delta = C_max - C_min

        l = np.clip((C_max + C_min)*.5,0.,1.)                                                       # noqa
        s = 0. if np.isclose(C_max,C_min) else Delta/(1.-np.abs(2*l-1.))
        if np.isclose(Delta,0.):
            h = 0.
        elif rgb.argmax() == 0:
            h = (rgb[1]-rgb[2])/Delta%6
        elif rgb.argmax() == 1:
            h = (rgb[2]-rgb[0])/Delta + 2.
        elif rgb.argmax() == 2:
            h = (rgb[0]-rgb[1])/Delta + 4.

        h = np.clip(h,0.,6.) * 60.

        return np.array([h,s,l])


    @staticmethod
    def _xyz2rgb(xyz):
        """
        CIE Xyz to R(ed) G(reen) B(lue).

        References
        ----------
        http://www.ryanjuckett.com/programming/rgb-color-space-conversion

        """
        rgb_lin = np.dot(np.array([
                                   [ 3.240969942,-1.537383178,-0.498610760],
                                   [-0.969243636, 1.875967502, 0.041555057],
                                   [ 0.055630080,-0.203976959, 1.056971514]
                                  ]),xyz)
        with np.errstate(invalid='ignore'):
            rgb = np.where(rgb_lin>0.0031308,rgb_lin**(1.0/2.4)*1.0555-0.0555,rgb_lin*12.92)

        return np.clip(rgb,0.,1.)

    @staticmethod
    def _rgb2xyz(rgb):
        """
        R(ed) G(reen) B(lue) to CIE Xyz.

        References
        ----------
        http://www.ryanjuckett.com/programming/rgb-color-space-conversion

        """
        rgb_lin = np.where(rgb>0.04045,((rgb+0.0555)/1.0555)**2.4,rgb/12.92)
        return np.dot(np.array([
                                [0.412390799,0.357584339,0.180480788],
                                [0.212639006,0.715168679,0.072192315],
                                [0.019330819,0.119194780,0.950532152]
                               ]),rgb_lin)


    @staticmethod
    def _lab2xyz(lab,ref_white=None):
        """
        CIE Lab to CIE Xyz.

        References
        ----------
        http://www.brucelindbloom.com/index.html?Eqn_Lab_to_XYZ.html

        """
        f_x = (lab[0]+16.)/116. + lab[1]/500.
        f_z = (lab[0]+16.)/116. - lab[2]/200.

        return np.array([
                         f_x**3.                if f_x**3. > _eps     else (116.*f_x-16.)/_kappa,
                         ((lab[0]+16.)/116.)**3 if lab[0]>_kappa*_eps else lab[0]/_kappa,
                         f_z**3.                if f_z**3. > _eps     else (116.*f_z-16.)/_kappa
                        ])*(ref_white if ref_white is not None else _ref_white)

    @staticmethod
    def _xyz2lab(xyz,ref_white=None):
        """
        CIE Xyz to CIE Lab.

        References
        ----------
        http://www.brucelindbloom.com/index.html?Eqn_Lab_to_XYZ.html

        """
        ref_white = ref_white if ref_white is not None else _ref_white
        f = np.where(xyz/ref_white > _eps,(xyz/ref_white)**(1./3.),(_kappa*xyz/ref_white+16.)/116.)

        return np.array([
                         116.0 *  f[1] - 16.0,
                         500.0 * (f[0] - f[1]),
                         200.0 * (f[1] - f[2])
                        ])


    @staticmethod
    def _lab2msh(lab):
        """
        CIE Lab to Msh.

        References
        ----------
        https://www.kennethmoreland.com/color-maps/ColorMapsExpanded.pdf
        https://www.kennethmoreland.com/color-maps/diverging_map.py

        """
        M = np.linalg.norm(lab)
        return np.array([
                         M,
                         np.arccos(lab[0]/M)       if M>1e-8 else 0.,
                         np.arctan2(lab[2],lab[1]) if M>1e-8 else 0.,
                        ])

    @staticmethod
    def _msh2lab(msh):
        """
        Msh to CIE Lab.

        References
        ----------
        https://www.kennethmoreland.com/color-maps/ColorMapsExpanded.pdf
        https://www.kennethmoreland.com/color-maps/diverging_map.py

        """
        return np.array([
                         msh[0] * np.cos(msh[1]),
                         msh[0] * np.sin(msh[1]) * np.cos(msh[2]),
                         msh[0] * np.sin(msh[1]) * np.sin(msh[2])
                        ])

    @staticmethod
    def _lab2rgb(lab):
        return Colormap._xyz2rgb(Colormap._lab2xyz(lab))

    @staticmethod
    def _rgb2lab(rgb):
        return Colormap._xyz2lab(Colormap._rgb2xyz(rgb))

    @staticmethod
    def _msh2rgb(msh):
        return Colormap._lab2rgb(Colormap._msh2lab(msh))

    @staticmethod
    def _rgb2msh(rgb):
        return Colormap._lab2msh(Colormap._rgb2lab(rgb))

    @staticmethod
    def _hsv2msh(hsv):
        return Colormap._rgb2msh(Colormap._hsv2rgb(hsv))

    @staticmethod
    def _hsl2msh(hsl):
        return Colormap._rgb2msh(Colormap._hsl2rgb(hsl))

    @staticmethod
    def _xyz2msh(xyz):
        return Colormap._lab2msh(Colormap._xyz2lab(xyz))