Generalize and refactor colorblindfriendly.py (#155)

jaredsampson · web-flow · commit 6410edc6587b · 2025-02-10T13:00:15.000-05:00
diff --git a/colorblindfriendly.py b/colorblindfriendly.py
@@ -16,7 +16,7 @@
     reference to consult when making all kinds of figures, not just those made
     using PyMOL.
 
-    The colors are:
+    The "colorblind" color palette includes:
 
     * cb_black
     * cb_orange
@@ -27,20 +27,38 @@
     * cb_vermillion (also: cb_red, cb_redorange, cb_red_orange)
     * cb_reddish_purple (also: cb_rose, cb_violet, cb_magenta)
 
+    Also added are two palettes from matplotlib, "viridis" and "magma", which
+    are designed to be perceptually uniform in both color and black-and-white
+    printouts.  These are available as "viridis[1-11]", "magma[1-11]".
+
 USAGE
 
+    With the PyMOL Script Repo installed and importable, import the module and
+    set the colors:
+
+    ```
     import colorblindfriendly as cbf
 
     # Add the new colors
     cbf.set_colors()
     color myObject, cb_red
 
-    # Replace built-in colors with cbf ones
+    # Replace built-in colors of same names with cbf ones
     cbf.set_colors(replace=True)
     color myOtherObject, yellow   # actually cb_yellow
 
     # Add a `cb_colors` menu item to the OpenGL GUI ([C] menu in the right panel)
     cbf.add_menu()
+    ```
+
+    Or, to use without installing, run the script directly from GitHub.  This
+    will add the colors and install GUI palette menus for  all three default
+    color palettes:
+
+    ```
+    run https://github.com/Pymol-Scripts/Pymol-script-repo/blob/master/colorblindfriendly.py
+    color myObject, cb_red
+    ```
 
 REQUIREMENTS
 
@@ -53,7 +71,7 @@
 
 LICENSE
 
-Copyright (c) 2014-2017 Jared Sampson
+Copyright (c) 2014-2025 Jared Sampson
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
@@ -75,6 +93,13 @@
 
 CHANGELOG
 
+    0.4.0   Add Palette and PaletteColor NamedTuples for cleaner declaration of
+            color palettes.  [2025-02-10]
+
+    0.3.0   Generalize the way colors and menus are efined and added, to
+            enable the use of additional color palettes. Add Viridis and Magma
+            palettes (contributed by Yehudi Bloch).  [2021-10-27]
+
     0.2.0   Complete overhaul for PyMOL 2.0 with conversion to module format.
             Now, setting the new `cb_*` color values requires a call to the
            `set_colors()` function after import.  You can also now add a
@@ -83,98 +108,163 @@
 
 '''
 from __future__ import print_function
+import math
+from typing import NamedTuple, Optional
 
 __author__ = 'Jared Sampson'
-__version__ = '0.2.0'
+__version__ = '0.4.0'
 
 import pymol
 from pymol import cmd
 
 
-# Color blind-friendly color list based on information found at:
-# http://jfly.iam.u-tokyo.ac.jp/html/color_blind/#pallet
-CB_COLORS = {
-    'black': {
-        'rgb': [0, 0, 0],
-        'alt': None,
-    },
-    'orange': {
-        'rgb': [230, 159, 0],
-        'alt': None,
-    },
-    'sky_blue': {
-        'rgb': [86, 180, 233],
-        'alt': ['skyblue', 'light_blue', 'lightblue'],
-    },
-    'bluish_green': {
-        'rgb': [0, 158, 115],
-        'alt': ['bluishgreen', 'green'],
-    },
-    'yellow': {
-        'rgb': [240, 228, 66],
-        'alt': None,
-    },
-    'blue': {
-        'rgb': [0, 114, 178],
-        'alt': None,
-    },
-    'vermillion': {
-        'rgb': [213, 94, 0],
-        'alt': ['red', 'red_orange', 'redorange'],
-    },
-    'reddish_purple': {
-        'rgb': [204, 121, 167],
-        'alt': ['reddishpurple', 'rose', 'violet', 'magenta'],
-    },
-}
+class PaletteColor(NamedTuple):
+    '''Named tuple for storing color information.'''
+    name: str
+    rgb: tuple[int, int, int]
+    alt_names: Optional[list[str]] = None
+    # Allow code to be set explicitly in palette definition. This is helpful
+    # for very dark colors, to allow contrast against the dark menu background.
+    short_code: Optional[str] = None  # for GUI menu
 
+    def all_names(self):
+        '''Return a list of all names for this color.'''
+        names = [self.name]
+        if self.alt_names:
+            names.extend(self.alt_names)
+        return names
 
-def set_colors(replace=False):
-    '''Add the color blind-friendly colors to PyMOL.'''
-    # Track the added colors
-    added_colors = []
+    def get_short_code(self):
+        '''Return a 3-digit string approximating the RGB color.'''
+        if self.short_code:
+            return self.short_code
+        return ''.join([str(math.floor(x / 256 * 10)) for x in self.rgb])
 
-    for color, properties in CB_COLORS.items():
-        # RGB tuple shortcut
-        rgb = properties['rgb']
 
-        # Get the primary and alternate color names into a single list
-        names = [color]
-        if properties['alt']:
-            names.extend(properties['alt'])
+class Palette(NamedTuple):
+    '''Named tuple for storing palette information.'''
+    name: str
+    colors: list[PaletteColor]
+    prefix: str = ''
 
-        # Set the colors
-        for name in names:
-            # Set the cb_color
-            cb_name = 'cb_{}'.format(name)
-            cmd.set_color(cb_name, rgb)
+    def install(self):
+        '''Install the palette, adding colors and the GUI menu.'''
+        PALETTES_MAP[self.name] = self
+        add_menu(self.name)
 
-            # Optionally replace built-in colors
-            if replace:
-                cmd.set_color(name, rgb)
-                spacer = (20 - len(name)) * ' '
-                added_colors.append('    {}{}{}'.format(name, spacer, cb_name))
-            else:
-                added_colors.append('    {}'.format(cb_name))
 
-    # Notify user of newly available colors
-    print('\nColor blind-friendly colors are now available:')
-    print('\n'.join(added_colors))
-    print('')
+# Color blind-friendly color list based on information found at:
+# http://jfly.iam.u-tokyo.ac.jp/html/color_blind/#pallet
+CB_COLORS = [
+    PaletteColor('red', (213, 94, 0),
+                 ['vermillion', 'red_orange', 'redorange']),
+    PaletteColor('orange', (230, 159, 0)),
+    PaletteColor('yellow', (240, 228, 66)),
+    PaletteColor('green', (0, 158, 115),
+                 ['bluish_green', 'bluishgreen']),
+    PaletteColor('light_blue', (86, 180, 233),
+                 ['lightblue', 'sky_blue', 'skyblue']),
+    PaletteColor('blue', (0, 114, 178)),
+    PaletteColor('violet', (204, 121, 167),
+                 ['reddish_purple', 'reddishpurple', 'rose', 'magenta']),
+    PaletteColor('black', (0, 0, 0), short_code='222'),
+]
+CB_PALETTE = Palette('colorblind', CB_COLORS, prefix='cb_')
+
+# Viridis and Magma palettes contributed by Yehudi Bloch, originally
+# developed by Stéfan van der Walt and Nathaniel Smith for matplotlib.
+# https://matplotlib.org/stable/users/prev_whats_new/whats_new_1.5.html
+VIRIDIS_COLORS = [
+    PaletteColor('viridis1',  (253, 231,  36)),
+    PaletteColor('viridis2',  (186, 222,  39)),
+    PaletteColor('viridis3',  (121, 209,  81)),
+    PaletteColor('viridis4',  ( 66, 190, 113)),
+    PaletteColor('viridis5',  ( 34, 167, 132)),
+    PaletteColor('viridis6',  ( 32, 143, 140)),
+    PaletteColor('viridis7',  ( 41, 120, 142)),
+    PaletteColor('viridis8',  ( 52,  94, 141)),
+    PaletteColor('viridis9',  ( 64,  67, 135)),
+    PaletteColor('viridis10', ( 72,  35, 116)),
+    PaletteColor('viridis11', ( 68,   1,  84)),
+]
+VIRIDIS_PALETTE = Palette('viridis', VIRIDIS_COLORS)
+
+MAGMA_COLORS = [
+    PaletteColor('magma1',  (251, 252, 191)),
+    PaletteColor('magma2',  (253, 205, 114)),
+    PaletteColor('magma3',  (253, 159, 108)),
+    PaletteColor('magma4',  (246, 110,  91)),
+    PaletteColor('magma5',  (221,  73, 104)),
+    PaletteColor('magma6',  (181,  54, 121)),
+    PaletteColor('magma7',  (140,  41, 128)),
+    PaletteColor('magma8',  ( 99,  25, 127)),
+    PaletteColor('magma9',  ( 59,  15, 111)),
+    PaletteColor('magma10', ( 20,  13,  53)),
+    PaletteColor('magma11', (  0,   0,   3)),
+]
+MAGMA_PALETTE = Palette('magma', MAGMA_COLORS)
+
+PALETTES_MAP = {
+    CB_PALETTE.name: CB_PALETTE,
+    VIRIDIS_PALETTE.name: VIRIDIS_PALETTE,
+    MAGMA_PALETTE.name: MAGMA_PALETTE,
+}
+
+
+def _get_palettes(palette_name: Optional[str] = None):
+    '''Return the desired Palette(s).'''
+    if palette_name is None:
+        return PALETTES_MAP.values()
+    if palette_name not in PALETTES_MAP:
+        raise ValueError(f'Palette "{palette_name}" not found.')
+    else:
+        return [PALETTES_MAP[palette_name]]
 
 
-def add_menu():
-    '''Add a color blind-friendly list of colors to the PyMOL OpenGL menu.'''
+def set_colors(palette=None, replace=False):
+    '''Add the color blind-friendly colors to PyMOL.'''
+    palettes = _get_palettes(palette)
+    for palette in palettes:
+        added_colors = []
+        for color in palette.colors:
+            # RGB tuple shortcut
+            rgb = color.rgb
+
+            # Set the colors
+            for name in color.all_names():
+                if palette.prefix:
+                    use_name = f'{palette.prefix}{name}'
+                else:
+                    use_name = name
+                cmd.set_color(use_name, rgb)
+
+                # Optionally replace built-in colors
+                if replace:
+                    cmd.set_color(name, rgb)
+                    # FIXME hard-coded column width
+                    spacer = (20 - len(name)) * ' '
+                    added_colors.append(f'    {name}{spacer}{use_name}')
+                else:
+                    added_colors.append('    {}'.format(use_name))
+
+        # Notify user of newly available colors
+        print(f'These {palette.name} colors are now available:')
+        print('\n'.join(added_colors))
+
+
+def _add_palette_menu(palette: Palette):
+    '''Add a color palette to the PyMOL OpenGL menu.'''
 
     # Make sure cb_colors are installed.
-    print('Checking for colorblindfriendly colors...')
+    print(f'Checking for {palette.name} colors...')
     try:
-        if cmd.get_color_index('cb_red') == -1:
-            # mimic pre-1.7.4 behavior
-            raise pymol.CmdException
+        for color in palette.colors:
+            if cmd.get_color_index(color.name) == -1:
+                # mimic pre-1.7.4 behavior
+                raise pymol.CmdException
     except pymol.CmdException:
-        print('Adding colorblindfriendly colors...')
-        set_colors()
+        print(f'Adding {palette.name} palette colors...')
+        set_colors(palette=palette.name)
 
     # Abort if PyMOL is too old.
     try:
@@ -184,35 +274,50 @@ def add_menu():
         return
 
     # Add the menu
-    print('Adding cb_colors menu...')
+    print(f'Adding {palette.name} menu...')
     # mimic pymol.menu.all_colors_list format
     # first color in list is used for menu item color
-    cb_colors = ('cb_colors', [
-        ('830', 'cb_red'),
-        ('064', 'cb_green'),
-        ('046', 'cb_blue'),
-        ('882', 'cb_yellow'),
-        ('746', 'cb_magenta'),
-        ('368', 'cb_skyblue'),
-        ('860', 'cb_orange'),
-    ])
+
+    # Menu item for each color in the menu should be a tuple in the form
+    #    ('999', 'color_name')
+    # where '999' is a string representing the 0-255 RGB color converted to
+    # a 0-9 integer RGB format (i.e. 1000 colors).
+    color_tuples = [
+        (color.get_short_code(), palette.prefix + color.name)
+        for color in palette.colors
+    ]
+    menu_colors = (palette.name, color_tuples)
+
     # First `pymol` is the program instance, second is the Python module
     all_colors_list = pymol.pymol.menu.all_colors_list
-    if cb_colors in all_colors_list:
-        print('Menu was already added!')
+    if menu_colors in all_colors_list:
+        print(f'  - Menu for {palette.name} was already added!')
     else:
-        all_colors_list.append(cb_colors)
-    print('  done.')
+        all_colors_list.append(menu_colors)
+    print('    done.\n')
+
 
+def add_menu(palette_name=None):
+    '''Add the specified color palettes to the PyMOL OpenGL menu.'''
+    palettes = _get_palettes(palette_name)
+    for palette in palettes:
+        _add_palette_menu(palette)
 
-def remove_menu():
-    '''Remove the cb_colors menu.'''
+
+def remove_menu(palette_name=None):
+    '''Remove the color palette menu(s).'''
+    palettes = _get_palettes(palette_name)
     all_colors_list = pymol.pymol.menu.all_colors_list
-    if all_colors_list[-1][0] == 'cb_colors':
-        all_colors_list.pop()
-        print('The `cb_colors` menu has been removed.')
-    else:
-        print('The `cb_colors` menu was not found! Aborting.')
+    for palette in palettes:
+        initial_length = len(all_colors_list)
+        all_colors_list[:] = [
+            color_menu for color_menu in all_colors_list
+            if color_menu[0] != palette.name
+        ]
+        if len(all_colors_list) == initial_length:
+            print(f'No menu for {palette.name} palette found. Nothing deleted.')
+        else:
+            print(f'Deleted menu for {palette.name} palette.')
 
 
 if __name__ == "pymol":
