feat: add stratigraphic age application functionality with color ramp selection

Author: lachlangrose

loopstructural/gui/modelling/stratigraphic_column/stratigraphic_column.py

@@ -1,9 +1,11 @@
 from LoopStructural.modelling.core.stratigraphic_column import StratigraphicColumnElementType
-from qgis.core import QgsMapLayerProxyModel
+from qgis.core import QgsMapLayerProxyModel, QgsStyle
 from qgis.gui import QgsFieldComboBox, QgsMapLayerComboBox
 from qgis.PyQt.QtWidgets import (
     QAbstractItemView,
+    QComboBox,
     QHBoxLayout,
+    QLabel,
     QListWidget,
     QListWidgetItem,
     QMessageBox,
@@ -94,6 +96,26 @@ def __init__(self, parent=None, data_manager=None):
         applyColoursButton.clicked.connect(self.apply_colours_to_layer)
         layout.addWidget(applyColoursButton)
 
+        # Colour ramp picker + apply stratigraphic age button
+        ageRow = QHBoxLayout()
+        ageRow.addWidget(QLabel("Colour ramp:"))
+        self.strat_ageColorRampComboBox = QComboBox()
+        ramp_names = sorted(QgsStyle().defaultStyle().colorRampNames())
+        self.strat_ageColorRampComboBox.addItems(ramp_names)
+        default_ramp_index = self.strat_ageColorRampComboBox.findText('Viridis')
+        if default_ramp_index >= 0:
+            self.strat_ageColorRampComboBox.setCurrentIndex(default_ramp_index)
+        ageRow.addWidget(self.strat_ageColorRampComboBox)
+        layout.addLayout(ageRow)
+
+        applyAgeButton = QPushButton("Apply Stratigraphic Age to Map Layer")
+        applyAgeButton.setToolTip(
+            "Write a 'strat_order' field (0 = first unit in the column) onto "
+            "the selected layer above and style it with a graduated colour ramp."
+        )
+        applyAgeButton.clicked.connect(self.apply_age_to_layer)
+        layout.addWidget(applyAgeButton)
+
         self._guess_units_layer()
         self._restore_units_layer_selection()
 
@@ -272,6 +294,38 @@ def apply_colours_to_layer(self):
                 "Could not apply colours. The stratigraphic column may have no units.",
             )
 
+    def apply_age_to_layer(self):
+        """Write the stratigraphic order onto the selected units layer and style it by a graduated ramp."""
+        if not self.data_manager:
+            print("Error: Data manager is not initialized.")
+            return
+        layer = self.unitsLayerComboBox.currentLayer()
+        field_name = self.unitsLayerFieldComboBox.currentField()
+        if layer is None or not field_name:
+            QMessageBox.warning(
+                self,
+                "Apply Stratigraphic Age to Map Layer",
+                "Please select a units layer and unit name field above.",
+            )
+            return
+        ramp_name = self.strat_ageColorRampComboBox.currentText()
+        applied = self.data_manager.apply_stratigraphic_age_to_layer(
+            layer, field_name, ramp_name=ramp_name
+        )
+        if applied:
+            QMessageBox.information(
+                self,
+                "Apply Stratigraphic Age to Map Layer",
+                f"Applied stratigraphic age and graduated styling to layer '{layer.name()}'.",
+            )
+        else:
+            QMessageBox.warning(
+                self,
+                "Apply Stratigraphic Age to Map Layer",
+                "Could not apply stratigraphic age. The stratigraphic column may have no "
+                "units, or no features matched a stratigraphic unit.",
+            )
+
     def add_unit(self, *, unit_data=None, create_new=True):
         if unit_data is None:
             unit_data = {'type': 'unit', 'name': ''}

loopstructural/main/data_manager.py

@@ -5,9 +5,12 @@
 from qgis.core import (
     QgsCategorizedSymbolRenderer,
     QgsCoordinateReferenceSystem,
+    QgsGraduatedSymbolRenderer,
     QgsPointXY,
     QgsProject,
     QgsRendererCategory,
+    QgsRendererRange,
+    QgsStyle,
     QgsSymbol,
     QgsVectorLayer,
 )
@@ -17,6 +20,7 @@
 from LoopStructural.datatypes import BoundingBox
 from LoopStructural.modelling.core.stratigraphic_column import StratigraphicColumnElementType
 
+from .m2l_api import paint_stratigraphic_order
 from .vectorLayerWrapper import qgsLayerToGeoDataFrame
 
 
@@ -417,6 +421,81 @@ def apply_stratigraphic_colours_to_layer(self, layer, field_name):
         self.logger(message=f"Applied stratigraphic column colours to layer '{layer.name()}'.")
         return True
 
+    def apply_stratigraphic_age_to_layer(self, layer, field_name, ramp_name=None):
+        """Write the stratigraphic order onto a layer and style it with a graduated colour ramp.
+
+        Writes an integer 'strat_order' field to ``layer`` (0 = first unit in
+        the stratigraphic column) matched via ``field_name``, then applies a
+        graduated renderer over that field.
+
+        Parameters
+        ----------
+        layer : QgsVectorLayer
+            The layer to update (e.g. the geological units/geology layer).
+        field_name : str
+            Name of the field on ``layer`` holding the stratigraphic unit name.
+        ramp_name : str, optional
+            Name of a QGIS colour ramp (from QgsStyle) to use for the
+            graduated renderer. Falls back to any available ramp if not found.
+
+        Returns
+        -------
+        bool
+            True if the field was written and the renderer applied, False otherwise.
+        """
+        if layer is None or not field_name:
+            self.logger(message="No layer/unit name field set, cannot apply stratigraphic age.")
+            return False
+
+        unit_names = self.get_stratigraphic_unit_names()
+        if not unit_names:
+            self.logger(message="Stratigraphic column has no units, cannot apply stratigraphic age.")
+            return False
+
+        age_field_name = "strat_order"
+        try:
+            paint_stratigraphic_order(layer, unit_names, field_name)
+        except Exception as err:
+            self.logger(message=f"Failed to write stratigraphic order onto layer: {err}")
+            return False
+
+        unique_values = set()
+        for feature in layer.getFeatures():
+            value = feature[age_field_name]
+            if value is None or (hasattr(value, 'isNull') and value.isNull()):
+                continue
+            unique_values.add(int(value))
+        unique_values = sorted(unique_values)
+
+        if not unique_values:
+            self.logger(
+                message="No features matched a stratigraphic unit, cannot style layer by age."
+            )
+            return False
+
+        style = QgsStyle().defaultStyle()
+        ramp = style.colorRamp(ramp_name) if ramp_name else None
+        if ramp is None:
+            ramp_names = style.colorRampNames()
+            if ramp_names:
+                ramp = style.colorRamp(ramp_names[0])
+
+        n = len(unique_values)
+        ranges = []
+        for i, value in enumerate(unique_values):
+            symbol = QgsSymbol.defaultSymbol(layer.geometryType())
+            if ramp is not None:
+                symbol.setColor(ramp.color(i / (n - 1) if n > 1 else 0))
+            ranges.append(QgsRendererRange(value - 0.5, value + 0.5, symbol, str(value)))
+
+        layer.setRenderer(QgsGraduatedSymbolRenderer(age_field_name, ranges))
+        layer.triggerRepaint()
+        self.logger(
+            message=f"Applied stratigraphic age field '{age_field_name}' and graduated "
+            f"styling to layer '{layer.name()}'."
+        )
+        return True
+
     def get_stratigraphic_unit_names(self):
         """Get the names of the stratigraphic units in the column."""
         units = []