fix: update dem layer from project save

Author: lachlangrose

loopstructural/gui/modelling/model_definition/dem.py

@@ -1,22 +1,29 @@
 import os
 
 from PyQt5.QtWidgets import QWidget
+from qgis.core import QgsMapLayerProxyModel
 from qgis.PyQt import uic
 from qgis.PyQt.QtWidgets import QSizePolicy
-from qgis.core import QgsMapLayerProxyModel
+
 
 class DEMWidget(QWidget):
     def __init__(self, parent=None, data_manager=None):
         self.data_manager = data_manager
         super().__init__(parent)
         ui_path = os.path.join(os.path.dirname(__file__), "dem.ui")
         uic.loadUi(ui_path, self)
-        self.demLayerQgsMapLayerComboBox.setFilters(
-            QgsMapLayerProxyModel.RasterLayer)
+        self.demLayerQgsMapLayerComboBox.setFilters(QgsMapLayerProxyModel.RasterLayer)
         self.useDEMCheckBox.stateChanged.connect(self.onUseDEMClicked)
         self.elevationQgsDoubleSpinBox.valueChanged.connect(self.onElevationChanged)
         self.onElevationChanged()
-        
+        self.data_manager.set_dem_callback(self.set_dem_layer)
+
+    def set_dem_layer(self, layer):
+        """Set the DEM layer in the combo box."""
+        # if layer:
+        #     self.demLayerQgsMapLayerComboBox.setLayer(layer)
+        pass
+
     def onUseDEMClicked(self):
         if self.useDEMCheckBox.isChecked():
             self.demLayerQgsMapLayerComboBox.setEnabled(True)
@@ -38,11 +45,9 @@ def onDEMLayerChanged(self):
         else:
             self.data_manager.set_dem_layer(None)
         self.data_manager.set_use_dem(True)
-    
+
     def onElevationChanged(self):
         """Handle changes to the elevation value."""
         elevation = self.elevationQgsDoubleSpinBox.value()
         self.data_manager.set_elevation(elevation)
         self.data_manager.set_use_dem(False)
-
-        

loopstructural/main/data_manager.py

@@ -136,6 +136,12 @@ def set_stratigraphic_column_callback(self, callback):
         """Set the callback for when the stratigraphic column is updated."""
         self.stratigraphic_column_callback = callback
 
+    def set_dem_callback(self, callback):
+        """Set the callback for when the DEM layer is updated."""
+        self.dem_callback = callback
+        if self.dem_layer:
+            self.dem_callback(self.dem_layer)
+
     def get_bounding_box(self):
         """Get the current bounding box."""
         return self._bounding_box
@@ -155,12 +161,24 @@ def set_dem_layer(self, dem_layer):
                 log_level=2,
             )
         else:
-            self.dem_function = lambda x, y: (
-                self.dem_layer.dataProvider().sample(QgsPointXY(x, y), 1)[0]
-                if self.dem_layer
-                else np.nan
-            )
+
+            def dem_function(x, y):
+                if not self.dem_layer.isValid():
+                    self.logger(
+                        message="DEM layer is not valid, using 0.0 for elevation.",
+                        log_level=2,
+                    )
+                    return 0.0
+                return (
+                    self.dem_layer.dataProvider().sample(QgsPointXY(x, y), 1)[0]
+                    if self.dem_layer
+                    else np.nan
+                )
+
+            self.dem_function = dem_function
             self._model_manager.set_dem_function(self.dem_function)
+        if self.dem_callback:
+            self.dem_callback(self.dem_layer)
 
     def set_use_dem(self, use_dem):
         self.use_dem = use_dem
@@ -391,6 +409,8 @@ def to_dict(self):
             and structural_orientations['layer'] is not None
         ):
             structural_orientations['layer'] = structural_orientations['layer'].name()
+        if self.dem_layer is not None:
+            dem_layer_name = self.dem_layer.name()
 
         return {
             'bounding_box': self._bounding_box.to_dict(),
@@ -400,6 +420,9 @@ def to_dict(self):
             'stratigraphic_column': (
                 self._stratigraphic_column.to_dict() if self._stratigraphic_column else None
             ),
+            'dem_layer': dem_layer_name if self.dem_layer else None,
+            'use_dem': self.use_dem,
+            'elevation': self.elevation,
         }
 
     def from_dict(self, data):
@@ -413,7 +436,19 @@ def from_dict(self, data):
                 zmin=data['bounding_box']['origin'][2],
                 zmax=data['bounding_box']['maximum'][2],
             )
-
+        if 'dem_layer' in data and data['dem_layer'] is not None:
+            dem_layer = QgsProject.instance().mapLayersByName(data['dem_layer'])
+            if dem_layer:
+                self.set_dem_layer(dem_layer[0])
+            else:
+                self.logger(
+                    message=f"DEM layer '{data['dem_layer']}' not found in project.",
+                    log_level=2,
+                )
+        if 'use_dem' in data:
+            self.set_use_dem(data['use_dem'])
+        if 'elevation' in data:
+            self.set_elevation(data['elevation'])
         if 'basal_contacts' in data:
             self._basal_contacts = data['basal_contacts']
         if 'fault_traces' in data:
@@ -437,6 +472,19 @@ def update_from_dict(self, data):
             )
         else:
             self.set_bounding_box(**default_bounding_box)
+        if 'dem_layer' in data and data['dem_layer'] is not None:
+            dem_layer = QgsProject.instance().mapLayersByName(data['dem_layer'])
+            if dem_layer:
+                self.set_dem_layer(dem_layer[0])
+            else:
+                self.logger(
+                    message=f"DEM layer '{data['dem_layer']}' not found in project.",
+                    log_level=2,
+                )
+        if 'use_dem' in data:
+            self.set_use_dem(data['use_dem'])
+        if 'elevation' in data:
+            self.set_elevation(data['elevation'])
         if (
             'basal_contacts' in data
             and data['basal_contacts'] is not None