Save color ramps (#62)

* Add functionality for load and save button in color widget

imodqgis/utils/color.py

@@ -0,0 +1,36 @@
+from PyQt5.QtGui import QColor
+
+from qgis.core import (
+    QgsGradientStop,
+    QgsGradientColorRamp
+)
+import numpy as np
+
+from typing import List
+
+
+def create_colorramp(
+        boundaries: List[float],
+        colors: List[QColor],
+        discrete: bool = False,
+    ):
+    """
+    Manually construct colorramp from boundaries and colors. The stops
+    determined by the createColorRamp method appear to be broken.
+    """
+
+    # For some reason discrete colormaps require the last color also added as
+    # stop
+    if discrete:
+        indices_stops = slice(1, None)
+    else:
+        indices_stops = slice(1, -1)
+
+    bound_arr = np.array(boundaries)
+    boundaries_norm = (bound_arr-bound_arr[0])/(bound_arr[-1]-bound_arr[0])
+    stops = [
+        QgsGradientStop(stop, color) for stop, color in zip(
+            boundaries_norm[indices_stops], colors[indices_stops]
+        )
+    ]
+    return QgsGradientColorRamp(colors[0], colors[-1], discrete, stops)

imodqgis/widgets/pseudocolor_widget.py

@@ -1,7 +1,7 @@
 # Copyright © 2021 Deltares
 # SPDX-License-Identifier: GPL-2.0-or-later
 #
-from typing import Dict, Union
+from typing import Dict, Union, List
 
 import numpy as np
 from PyQt5.QtCore import Qt
@@ -11,6 +11,7 @@
     QCheckBox,
     QComboBox,
     QGridLayout,
+    QFileDialog,
     QHBoxLayout,
     QLabel,
     QLineEdit,
@@ -22,12 +23,17 @@
 )
 from qgis.core import (
     QgsColorRampShader,
+    QgsGradientColorRamp,
+    QgsGradientStop,
+    QgsRasterRendererUtils
 )
 from qgis.gui import (
     QgsColorRampButton,
     QgsColorSwatchDelegate,
     QgsTreeWidgetItemObject,
 )
+from imodqgis.utils.color import create_colorramp
+
 
 Number = Union[int, float]
 
@@ -172,8 +178,20 @@ def maximum(self):
         else:
             return np.nanmax(self.data)
 
-    def classify(self) -> None:
+    def _set_color_items_in_table(self, boundaries: List, colors: List):
         self.table.clear()
+        for boundary, color in zip(boundaries, colors):
+            new_item = QgsTreeWidgetItemObject(self.table)
+            new_item.setData(0, Qt.ItemDataRole.DisplayRole, float(boundary))
+            new_item.setData(1, Qt.ItemDataRole.EditRole, color)
+            new_item.setText(2, "")
+            new_item.setFlags(
+                Qt.ItemIsEnabled | Qt.ItemIsEditable | Qt.ItemIsSelectable
+            )
+            new_item.itemEdited.connect(self.item_edited)
+        self.format_labels()
+
+    def classify(self) -> None:
         shader_mode = CLASSIFICATION_MODE[self.classification_box.currentText()]
         shader_type = SHADER_TYPES[self.interpolation_box.currentText()]
 
@@ -196,16 +214,7 @@ def classify(self) -> None:
                 boundaries = np.linspace(self.minimum(), self.maximum(), n_class)
 
         colors = [ramp.color(f) for f in np.linspace(0.0, 1.0, n_class)]
-        for boundary, color in zip(boundaries, colors):
-            new_item = QgsTreeWidgetItemObject(self.table)
-            new_item.setData(0, Qt.ItemDataRole.DisplayRole, float(boundary))
-            new_item.setData(1, Qt.ItemDataRole.EditRole, color)
-            new_item.setText(2, "")
-            new_item.setFlags(
-                Qt.ItemIsEnabled | Qt.ItemIsEditable | Qt.ItemIsSelectable
-            )
-            new_item.itemEdited.connect(self.item_edited)
-        self.format_labels()
+        self._set_color_items_in_table(boundaries, colors)
 
     def format_labels(self):
         shader_type = SHADER_TYPES[self.interpolation_box.currentText()]
@@ -253,10 +262,33 @@ def remove_selection(self):
             self.table.takeTopLevelItem(self.table.indexOfTopLevelItem(item))
 
     def load_classes(self):
-        pass
+        path, _ = QFileDialog.getOpenFileName(self, "Load colormap", "", "*.txt")
+        # Load colorsmap
+        load_succeeded, color_ramp_items, _, load_errors = QgsRasterRendererUtils.parseColorMapFile(path)
+        if not load_succeeded:
+            raise ValueError(f"Encountered the following errors while parsing color map file: {load_errors}")
+        # Unpack values
+        boundaries, colors = zip(*[(c_ramp_item.value, c_ramp_item.color) for c_ramp_item in color_ramp_items])
+        # Set color items in table
+        self._set_color_items_in_table(boundaries, colors)
+        # Set color ramp
+        colorramp = create_colorramp(boundaries, colors, discrete=False)
+        self.color_ramp_button.setColorRamp(colorramp)
+        return
 
     def save_classes(self):
-        pass
+        """
+        Save colors to a QGIS colormap textfile. This stores colors and
+        corresponding values.
+        """
+        path, _ = QFileDialog.getSaveFileName(self, "Save colormap", "", "*.txt")
+        shader = self.shader()
+        shader_type = SHADER_TYPES[self.interpolation_box.currentText()]
+        color_ramp_items = shader.colorRampItemList()
+        save_succeeded = QgsRasterRendererUtils.saveColorMapFile(path, color_ramp_items, shader_type)
+        if not save_succeeded:
+            raise ValueError(f"Error saving color map file {path}")
+        return
 
     def labels(self) -> Dict[str, str]:
         label_dict = {}

imodqgis/widgets/unique_color_widget.py

@@ -1,15 +1,17 @@
 # Copyright © 2021 Deltares
 # SPDX-License-Identifier: GPL-2.0-or-later
 #
-from typing import Dict
+from typing import Dict, List
 
 import numpy as np
 import pandas as pd
+import json
 from PyQt5.QtCore import Qt
 from PyQt5.QtGui import QColor
 from PyQt5.QtWidgets import (
     QAbstractItemView,
     QHBoxLayout,
+    QFileDialog,
     QLabel,
     QPushButton,
     QTreeWidget,
@@ -24,6 +26,7 @@
     QgsColorSwatchDelegate,
     QgsTreeWidgetItemObject,
 )
+from imodqgis.utils.color import create_colorramp
 
 
 class ImodUniqueColorShader:
@@ -85,14 +88,14 @@ def __init__(self, parent=None):
 
     def set_data(self, data: np.ndarray):
         self.data = data
-        self.classify()
+        # Extend list of colors with colors from colorramp button if more data
+        # points available than in loaded file.
+        colors = self.get_colors_from_ramp_button()
+        self.set_legend(colors)
 
-    def classify(self) -> None:
-        self.table.clear()
+    def set_legend(self, colors) -> None:
         uniques = pd.Series(self.data).dropna().unique()
-        n_class = uniques.size
-        ramp = self.color_ramp_button.colorRamp()
-        colors = [ramp.color(f) for f in np.linspace(0.0, 1.0, n_class)]
+        self.table.clear()
         for value, color in zip(uniques, colors):
             new_item = QgsTreeWidgetItemObject(self.table)
             # Make sure to convert from numpy type to Python type with .item()
@@ -107,6 +110,54 @@ def classify(self) -> None:
                 Qt.ItemIsEnabled | Qt.ItemIsEditable | Qt.ItemIsSelectable
             )
 
+    def _get_cyclic_normalized_midpoints(self, n_elements, cycle_size):
+        """
+        Create array with length n_elements which cycles through normalized
+        midpoint values, used to fetch values on unique colors colormap.
+
+        Example
+        -------
+        >>> self._get_cyclic_normalized_midpoints(n_elements=6, cycle_size=4)
+        array([0.125, 0.375, 0.625, 0.875, 0.125, 0.375)]
+        """
+        return ((np.arange(n_elements) % cycle_size) + 0.5) / cycle_size
+
+    def _needs_cyclic_colorramp(self):
+        ramp = self.color_ramp_button.colorRamp()
+        if ramp.type() in ["colorbrewer", "random"]:
+            return True
+        elif hasattr(ramp, "isDiscrete") and ramp.isDiscrete():
+            return True
+        else:
+            return False
+
+    def _count_discrete_colors(self):
+        """
+        The discrete gradient has one stop more than colors, whereas the
+        colorbrewer colorramp has as many stops as colors
+        """
+        ramp = self.color_ramp_button.colorRamp()
+        if ramp.type() in ["gradient"]:
+            return ramp.count() - 1
+        else:
+            return ramp.count()
+
+    def get_colors_from_ramp_button(self) -> List[QColor]:
+        uniques = pd.Series(self.data).dropna().unique()
+        n_class = uniques.size
+        ramp = self.color_ramp_button.colorRamp()
+        if self._needs_cyclic_colorramp():
+            n_colors = self._count_discrete_colors()
+            values_colors = self._get_cyclic_normalized_midpoints(n_class, n_colors)
+        else:
+            values_colors = np.linspace(0.0, 1.0, n_class)
+        return [ramp.color(f) for f in values_colors]
+
+    def classify(self) -> None:
+        self.table.clear()
+        colors = self.get_colors_from_ramp_button()
+        self.set_legend(colors)
+
     def add_class(self) -> None:
         new_item = QgsTreeWidgetItemObject(self.table)
         new_item.setData(0, Qt.ItemDataRole.DisplayRole, 0)
@@ -143,10 +194,35 @@ def remove_selection(self) -> None:
             self.table.takeTopLevelItem(self.table.indexOfTopLevelItem(item))
 
     def load_classes(self) -> None:
-        pass
+        path, _ = QFileDialog.getOpenFileName(self, "Load colors", "", "*.json")
+        # Load colors
+        with open(path, "r") as file:
+            rgb_values = json.load(file)
+        colors = [QColor(*rgb) for rgb in rgb_values]
+        # Set colorramp button
+        boundaries = np.linspace(0.0, 1.0, len(colors)+1)
+        color_ramp = create_colorramp(boundaries, colors, discrete=True)
+        self.color_ramp_button.setColorRamp(color_ramp)
+        # Set colors in table
+        table_iter = range(self.table.topLevelItemCount())
+        for i, color in zip(table_iter, colors):
+            item = self.table.topLevelItem(i)
+            item.setData(1, Qt.ItemDataRole.EditRole, color)
 
     def save_classes(self) -> None:
-        pass
+        """
+        Save colors to a .json file. The unique color widget saves only colors
+        without corresponding values, as values are usually too unique (e.g.
+        name of borehole) to be used for different datasets. QGIS has no
+        standard functionality for writing purely colors to file, hence we save
+        to a json.
+        """
+        path, _ = QFileDialog.getSaveFileName(self, "Save colors", "", "*.json")
+        colors = self.colors().values()
+        rgb_values = [c.getRgb() for c in colors]
+
+        with open(path, "w") as file:
+            json.dump(rgb_values, file)
 
     def shader(self) -> ImodUniqueColorShader:
         values = []