write_ncmaps

#!/usr/bin/env python3
#============================================================
#    write_ncmaps
#    Copyright (C) 2021  Thomas Lavergne
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License along
#    with this program; if not, write to the Free Software Foundation, Inc.,
#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#============================================================

import sys
import os
from glob import glob

# Import python modules holding colormaps
try:
    # https://matplotlib.org/stable/tutorials/colors/colormaps.html
    from matplotlib import cm
    from matplotlib import pylab
    has_mlib = True
except ImportError:
    has_mlib = False

try:
    # https://matplotlib.org/cmocean
    import cmocean
    has_cmocean = True
except ImportError:
    has_cmocean = False

try:
    # https://www.fabiocrameri.ch/colourmaps
    # https://pypi.org/project/cmcrameri
    import cmcrameri
    has_cmcrameri = True
except ImportError:
    has_cmcrameri = False

try:
    #  https://cmasher.readthedocs.io/index.html
    import cmasher as cmr
    from cmasher import cm as cmrcm
    has_cmasher = True
except ImportError:
    has_cmasher = False

# We could probaby do without numpy, but it is convenient
import numpy as np

def write_one_ncmap(cmap, name, odir):
    ''' Write a colormap into a .ncmap file that ncview will understand. '''
    x = np.arange(0,256) / 255.
    f = 255 * cmap(x)
    fname = name + '.ncmap'
    fname = os.path.join(odir,fname)
    if not os.path.exists(odir):
        os.makedirs(odir)
        print("Make directory {}".format(odir,))

    print("Write {}".format(fname,))
    np.savetxt(fname,f[:,:3],['%d','%d','%d'])

def write_mlib_cmaps(odir, ):
    ''' Write the colormaps from matplotlib. '''
    # Perceptually Uniform Sequential colormaps
    cmaps = ['viridis', 'inferno', 'magma', 'plasma', 'cividis']
    # Add some more colormaps (diverging and cyclic)
    cmaps.extend(['PiYG', 'PRGn', 'RdBu', 'coolwarm', 'bwr', 'twilight', 'twilight_shifted'])

    for cmap in cmaps:
        cmo = pylab.get_cmap(cmap)
        write_one_ncmap(cmo, cmap, odir)

    return cmaps

def write_cmocean_cmaps(odir,):
    ''' Write the colormaps from cmocean. '''
    cmaps = cmocean.cm.cmapnames
    for cmap in cmaps:
        cmo = cmocean.cm.cmap_d[cmap]
        write_one_ncmap(cmo, cmap, odir)

    return cmaps

def write_cmcrameri_cmaps(odir,):
    ''' Write the colormaps from cmcrameri. '''
    try:
        # old interface (at least before 1.4)
        cmap_dict = cmcrameri.cm.crameri_cmap
    except AttributeError:
        cmap_dict = cmcrameri.cm.cmaps

    cmaps = [cma for cma in cmap_dict.keys() if not cma.endswith('_r') and not cma.endswith('S')]
    for cmap in cmaps:
        cmo = cmap_dict[cmap]
        write_one_ncmap(cmo, cmap, odir)

    return cmaps

def write_cmasher_cmaps(odir,):
    ''' Write the colomaps from cmasher. '''
    all_cmaps = []
    for cmtype in ('sequential', 'diverging', 'cyclic'):
        cmaps = [ t for t in cmrcm.cmap_cd[cmtype].keys() if not t.endswith('_r')]
        for cmap in cmaps:
            cmo = cmrcm.cmap_cd[cmtype][cmap]
            write_one_ncmap(cmo, cmap, odir)
        all_cmaps.extend(cmaps,)

    return all_cmaps

def write_ncviewrc_header(ncviewrc,):
    ''' Create .ncviewrc file and Write a header (2 lines). '''

    with open(ncviewrc, 'w') as _:
        _.write('-1 "STRINGLIST_SAVE_FILE_VERSION" INT 1\n')
        _.write('0 "NCVIEW_STATE_FILE_VERSION" INT 1\n')

def _file_len(fname,):
    ''' Count number of lines in a (small) text file. '''
    i = 0
    with open(fname) as f:
        for i, l in enumerate(f):
            pass
    return i + 1

def write_cmaps_ncviewrc(ncviewrc, cmaps, enabled=True):
    ''' Append colormaps to the .ncviewrc file. '''

    # the first field of each line is a line counter (starting at -1)
    #   so we must first count the number of lines in the file.
    c = _file_len(ncviewrc)
    c -= 1

    with open(ncviewrc, 'a') as _:
        for ic, cmap in enumerate(cmaps,):
            _.write('{} "CMAP_{}" INT {}\n'.format(c+ic,cmap,int(enabled)))

if __name__ == '__main__':
    import argparse
    from argparse import RawDescriptionHelpFormatter

    # default directory hidden in HOME
    from os.path import expanduser
    home = expanduser("~")
    default_dir = os.path.join(home,'.ncmaps')

    # Get commandline parameters
    p = argparse.ArgumentParser("write_ncmaps", formatter_class=argparse.ArgumentDefaultsHelpFormatter,
                                description='Prepare scientific colormaps so that they can be used from ncview.')
    p.add_argument('--dir', help='Directory where to write the colormaps.', default=default_dir)
    args = p.parse_args()

    # Use the full path to the output directory
    args.dir = os.path.realpath(args.dir)

    if not has_mlib and not has_cmocean and not has_cmcrameri:
        sys.exit("Found no colormap to write.")

    # write all the .ncmaps files
    all_cmaps = []
    if has_mlib:
        all_cmaps.extend(write_mlib_cmaps(args.dir,))
    else:
        print("Skip colormaps from matplotlib (module not found).")

    if has_cmocean:
        all_cmaps.extend(write_cmocean_cmaps(args.dir,))
    else:
        print("Skip colormaps from cmocean (module not found).")

    if has_cmcrameri:
        all_cmaps.extend(write_cmcrameri_cmaps(args.dir,))
    else:
        print("Skip colormaps from cmcrameri (module not found).")

    # In ncview (Ncview 2.1.8 David W. Pierce  8 March 2017)
    #   a bug is triggered if the number of colormaps available
    #   to ncview through NCVIEWBASE is larger than 79.
    #   (see https://github.com/TomLav/ncmaps/issues/5#issuecomment-832060948)
    #
    # We can thus not support the cmasher colormaps at present. The code
    #   to handler cmasher is kept but de-activated. 
    #
    #if has_cmasher:
    #    all_cmaps.extend(write_cmasher_cmaps(args.dir,))
    #else:
    #    print("Skip colormaps from cmasher (module not found).")


    # the next step is to write a .ncviewrc file, to give the new
    #   colormaps priority and even disable the legacy ones. It is
    #   key that the .ncviewrc lists exactly the content of the
    #   directory with .ncmaps files.
    patt = os.path.join(args.dir,'*.ncmap')
    list_ncmap_files = glob(patt)
    if len(patt) == 0:
        sys.exit("Internal error: cannot find the .ncmaps files we just wrote.")
    ncmaps_in_dir = []
    for ncmap_file in list_ncmap_files:
        ncf = os.path.basename(ncmap_file)
        ncmaps_in_dir.append(ncf.replace('.ncmap',''))

    # This sort() will make the colormaps to be listed alphabetically
    #   in the .ncviewrc file (and thus ncview).
    ncmaps_in_dir.sort(key=str.lower)

    # Check that we are below the limit set by the bug in ncview (see above).
    if len(ncmaps_in_dir) >= 80:
        print("WARNING ! There are now {} .ncmap files in {}. This might trigger a bug in ncview (>= 80).".format(len(ncmaps_in_dir), args.dir))

    # prepare a .ncviewrc file
    ncviewrc = os.path.join(args.dir, '.ncviewrc')
    write_ncviewrc_header(ncviewrc)
    # write the new colormaps to .ncviewrc
    write_cmaps_ncviewrc(ncviewrc,ncmaps_in_dir)
    # disable the legacy ncview colormaps
    legacy_cmaps = ('3gauss', 'detail', 'ssec', 'bright', 'banded', 'rainbow', 'jaisnb', \
                    'jaisnc', 'jaisnd', 'blu_red', 'manga', 'jet', 'wheel', '3saw', \
                    'bw', 'default', 'extrema', 'helix', 'helix2', 'hotres')
    write_cmaps_ncviewrc(ncviewrc,legacy_cmaps,enabled=False)

    print("")
    print("Success! You can now:")
    print(" 1. Have 'export NCVIEWBASE={}' in your .profile so that ncview finds the new colormaps.".format(args.dir))
    print(' 2. Copy {} to your HOME to disable the legacy colormaps.'.format(ncviewrc,))