Add button to automatically apply RGB presets (#2558)

* Add button to apply preset RGB options

Add missing import

* Add changelog

* Add API method to apply presets

* Fix codestyle

* Add test

* Change button color

* Don't roll our own

* Don't show in mixed case

* Don't hard code numbers in the conditional statements

* Improve logic for loop

Co-authored-by: Kyle Conroy <kyleconroy@gmail.com>

* Put presets in chromatic order

---------

Co-authored-by: Kyle Conroy <kyleconroy@gmail.com>

CHANGES.rst

@@ -17,6 +17,8 @@ New Features
 
 - Reorder viewer and layer settings in Plot Options. [#2543, #2557]
 
+- Add button in Plot Options to apply preset RBG options to visible layers when in Monochromatic mode. [#2558]
+
 Cubeviz
 ^^^^^^^
 

jdaviz/configs/default/plugins/plot_options/plot_options.py

@@ -1,3 +1,4 @@
+import math
 import os
 import matplotlib
 import numpy as np
@@ -567,7 +568,7 @@ def user_api(self):
                        'image_contrast', 'image_bias',
                        'contour_visible', 'contour_mode',
                        'contour_min', 'contour_max', 'contour_nlevels', 'contour_custom_levels',
-                       'stretch_curve_visible']
+                       'stretch_curve_visible', 'apply_RGB_presets']
 
         return PluginUserApi(self, expose)
 
@@ -605,6 +606,61 @@ def vue_set_value(self, data):
         value = data.get('value')
         setattr(self, attr_name, value)
 
+    def apply_RGB_presets(self):
+        """
+        Applies preset colors, opacities, and stretch settings to all visible layers
+        (in all viewers) when in Monochromatic mode.
+        """
+
+        if (self.image_color_mode_value != "One color per layer" or
+                self.image_color_mode_sync['mixed']):
+            raise ValueError("RGB presets can only be applied if color mode is Monochromatic.")
+        # Preselected colors we want to use for 5 or less layers
+        preset_colors = [self.swatches_palette[0][0],
+                         self.swatches_palette[1][0],
+                         "#00FF00",
+                         "#0000FF",
+                         self.swatches_palette[4][1]
+                         ]
+
+        # Switch back to this at the end
+        initial_layer = self.layer_selected
+
+        # Filter out subset layers
+        image_layers = [layer for layer in self.layer.choices if "Subset" not in layer]
+        visible_layers = []
+        for layer in image_layers:
+            self.layer_selected = layer
+            if self.image_visible.value:
+                visible_layers.append(layer)
+
+        # Set opacity to something that seems sensible
+        n_visible = len(visible_layers)
+        default_opacity = 1
+        if n_visible > 2:
+            default_opacity = 1 / math.log2(n_visible)
+        # Sample along a colormap if we have too many layers
+        if n_visible == 2:
+            preset_colors = [preset_colors[0], preset_colors[3]]
+        elif n_visible == 3:
+            preset_colors = [preset_colors[0], preset_colors[2], preset_colors[3]]
+        elif n_visible > len(preset_colors):
+            cmap = matplotlib.colormaps['gist_rainbow'].resampled(n_visible)
+            preset_colors = [matplotlib.colors.to_hex(cmap(i), keep_alpha=True) for
+                             i in range(n_visible)]
+
+        for i in range(n_visible):
+            self.layer_selected = visible_layers[i]
+            self.image_opacity_value = default_opacity
+            self.image_color_value = preset_colors[i]
+            self.stretch_function_value = "arcsinh"
+            self.stretch_preset_value = 99
+
+        self.layer_selected = initial_layer
+
+    def vue_apply_RGB_presets(self, data):
+        self.apply_RGB_presets()
+
     @observe('is_active', 'layer_selected', 'viewer_selected',
              'stretch_hist_zoom_limits')
     @skip_if_no_updates_since_last_active()

jdaviz/configs/default/plugins/plot_options/plot_options.vue

@@ -68,6 +68,14 @@
       </glue-state-sync-wrapper>
     </div>
 
+    <div v-if="image_color_mode_value === 'One color per layer' && !image_color_mode_sync['mixed']">
+      <v-row justify="end">
+        <j-tooltip tooltipcontent="Apply preset RGB colors, scaling, and opacities to visible layers">
+          <v-btn color="accent" text @click="apply_RGB_presets">Assign Presets</v-btn>
+        </j-tooltip>
+      </v-row>
+    </div>
+
     <!-- GENERAL:AXES -->
     <glue-state-sync-wrapper v-if="axes_visible_sync.in_subscribed_states && config !== 'imviz'" :sync="axes_visible_sync" :multiselect="multiselect" @unmix-state="unmix_state('axes_visible')">
       <v-switch
@@ -181,7 +189,7 @@
         </v-menu>
       </div>
     </glue-state-sync-wrapper>
-    
+
     <glue-state-sync-wrapper v-if="line_visible_value  && (!marker_visible_sync.in_subscribed_states || marker_visible_value)" :sync="line_width_sync" :multiselect="multiselect" @unmix-state="unmix_state('line_width')">
       <glue-float-field label="Line Width" :value.sync="line_width_value" />
     </glue-state-sync-wrapper>
@@ -520,7 +528,7 @@
     <!-- IMAGE:CONTOUR -->
     <j-plugin-section-header v-if="contour_visible_sync.in_subscribed_states">Contours</j-plugin-section-header>
     <div style="display: grid"> <!-- overlay container -->
-      <div style="grid-area: 1/1">        
+      <div style="grid-area: 1/1">
         <glue-state-sync-wrapper :sync="contour_visible_sync" :multiselect="multiselect" @unmix-state="unmix_state('contour_visible')">
           <span>
             <v-btn icon @click.stop="contour_visible_value = !contour_visible_value">
@@ -568,7 +576,7 @@
           </div>
           <div v-else>
             <glue-state-sync-wrapper :sync="contour_custom_levels_sync" :multiselect="multiselect" @unmix-state="unmix_state('contour_levels')">
-              <v-text-field 
+              <v-text-field
                 label="Contour Levels"
                 :value="contour_custom_levels_txt"
                 @focus="contour_custom_levels_focus"
@@ -580,7 +588,7 @@
       </div>
       <div v-if="contour_spinner"
            class="text-center"
-           style="grid-area: 1/1; 
+           style="grid-area: 1/1;
                   z-index:2;
                   margin-left: -24px;
                   margin-right: -24px;

jdaviz/configs/default/plugins/plot_options/tests/test_plot_options.py

@@ -2,6 +2,7 @@
 import pytest
 from astropy import units as u
 from astropy.nddata import NDData
+import matplotlib
 from numpy.testing import assert_allclose
 from photutils.datasets import make_4gaussians_image
 
@@ -250,3 +251,57 @@ def test_update_knots_mismatched_length():
     stretch = SplineStretch()
     with pytest.raises(ValueError):
         stretch.update_knots(x=[0, 0.1, 0.2], y=[0, 0.05])
+
+
+def test_apply_presets(imviz_helper):
+    arr = np.arange(36).reshape(6, 6)
+    po = imviz_helper.plugins['Plot Options']
+
+    preset_colors = [po._obj.swatches_palette[0][0], po._obj.swatches_palette[1][0],
+                     "#00FF00", "#0000FF", po._obj.swatches_palette[4][1]]
+
+    # Test applying presets with < 6 layers
+
+    for i in range(4):
+        imviz_helper.load_data(arr, data_label=f"array_{i}")
+
+    po.image_color_mode = "Monochromatic"
+    po.apply_RGB_presets()
+
+    for i in range(4):
+        po.layer = f"array_{i}"
+        assert po.stretch_function.value == "arcsinh"
+        assert po.stretch_preset.value == 99
+        assert po.image_color.value == preset_colors[i]
+
+    # Test applying presets with > 5 layers
+
+    for i in range(4):
+        imviz_helper.load_data(arr, data_label=f"array_{i+4}")
+
+    po.layer = "array_5"
+    po.image_visible = False
+    po.layer = "array_3"
+
+    po.apply_RGB_presets()
+
+    assert po.layer.selected == "array_3"
+
+    colorbar_colors = matplotlib.colormaps['gist_rainbow'].resampled(7)
+    color_ind = 0
+
+    def _rgb_to_hex(rgb):
+        rgb = [int(x * 255) for x in rgb]
+        return f"#{rgb[0]:02x}{rgb[1]:02x}{rgb[2]:02x}{rgb[3]:02x}"
+
+    for i in range(8):
+        po.layer = f"array_{i}"
+        if i == 5:
+            # Make sure this one didn't change
+            assert po.stretch_function.value == "linear"
+        else:
+            assert po.stretch_function.value == "arcsinh"
+            assert po.stretch_preset.value == 99
+            assert po.image_color.value == matplotlib.colors.to_hex(colorbar_colors(color_ind),
+                                                                    keep_alpha=True)
+            color_ind += 1