fix: store location tracking and lines for thickness calculators

lachlangrose · lachlangrose · commit 06e586fb70c6 · 2025-12-16T13:54:53.000+11:00
adding basal contacts unit name
diff --git a/loopstructural/gui/map2loop_tools/thickness_calculator_widget.py b/loopstructural/gui/map2loop_tools/thickness_calculator_widget.py
@@ -8,6 +8,7 @@
 from loopstructural.toolbelt.preferences import PlgOptionsManager
 
 from ...main.helpers import ColumnMatcher, get_layer_names
+from ...main.vectorLayerWrapper import addGeoDataFrameToproject
 
 
 class ThicknessCalculatorWidget(QWidget):
@@ -58,6 +59,7 @@ def __init__(self, parent=None, data_manager=None):
         self.calculatorTypeComboBox.currentIndexChanged.connect(self._on_calculator_type_changed)
         self.geologyLayerComboBox.layerChanged.connect(self._on_geology_layer_changed)
         self.structureLayerComboBox.layerChanged.connect(self._on_structure_layer_changed)
+        self.basalContactsComboBox.layerChanged.connect(self._on_basal_contacts_layer_changed)
         self.runButton.clicked.connect(self._run_calculator)
         self._guess_layers()
         # Set up field combo boxes
@@ -108,6 +110,18 @@ def _setup_field_combo_boxes(self):
         self.unitNameFieldComboBox.setLayer(self.geologyLayerComboBox.currentLayer())
         self.dipFieldComboBox.setLayer(self.structureLayerComboBox.currentLayer())
         self.dipDirFieldComboBox.setLayer(self.structureLayerComboBox.currentLayer())
+        self.basalUnitNameFieldComboBox.setLayer(self.basalContactsComboBox.currentLayer())
+
+    def _on_basal_contacts_layer_changed(self):
+        """Update field combo box when basal contacts layer changes."""
+        layer = self.basalContactsComboBox.currentLayer()
+        self.basalUnitNameFieldComboBox.setLayer(layer)
+        # Optionally auto-select a likely unit name field
+        if layer:
+            fields = [field.name() for field in layer.fields()]
+            matcher = ColumnMatcher(fields)
+            if unit_match := matcher.find_match('UNITNAME'):
+                self.basalUnitNameFieldComboBox.setField(unit_match)
 
     def _on_geology_layer_changed(self):
         """Update field combo boxes when geology layer changes."""
@@ -188,6 +202,7 @@ def _run_calculator(self):
                 'structure': self.structureLayerComboBox.currentLayer(),
                 'calculator_type': calculator_type,
                 'unit_name_field': self.unitNameFieldComboBox.currentField(),
+                'basal_unit_name_field': self.basalUnitNameFieldComboBox.currentField(),
                 'dip_field': self.dipFieldComboBox.currentField(),
                 'dipdir_field': self.dipDirFieldComboBox.currentField(),
                 'orientation_type': self.orientationTypeComboBox.currentText(),
@@ -206,19 +221,32 @@ def _run_calculator(self):
 
             # Get stratigraphic order from data_manager
             if self.data_manager and hasattr(self.data_manager, 'stratigraphic_column'):
-                strati_order = [unit['name'] for unit in self.data_manager.stratigraphic_column]
+                strati_order = [unit['name'] for unit in self.data_manager._stratigraphic_column]
                 if strati_order:
                     kwargs['stratigraphic_order'] = strati_order
 
             result = calculate_thickness(**kwargs)
 
-            for idx in result.index:
-                u = result.loc[idx, 'name']
-                thick = result.loc[idx, 'ThicknessStdDev']
+            for idx in result['thicknesses'].index:
+                u = result['thicknesses'].loc[idx, 'name']
+                thick = result['thicknesses'].loc[idx, 'ThicknessStdDev']
                 if thick > 0:
 
-                    self.data_manager.stratigraphic_column.get_unit_by_name(u).thickness = thick
-            if result is not None and not result.empty:
+                    self.data_manager._stratigraphic_column.get_unit_by_name(u).thickness = thick
+            # Save debugging files if checkbox is checked
+            if self.saveDebugCheckBox.isChecked():
+                if 'lines' in result:
+                    if result['lines'] is not None and not result['lines'].empty:
+                        addGeoDataFrameToproject(result['lines'], "Lines")
+                if 'location_tracking' in result:
+                    if (
+                        result['location_tracking'] is not None
+                        and not result['location_tracking'].empty
+                    ):
+                        addGeoDataFrameToproject(
+                            result['location_tracking'], "Thickness Location Tracking"
+                        )
+            if result is not None and not result['thicknesses'].empty:
                 QMessageBox.information(
                     self,
                     "Success",
diff --git a/loopstructural/gui/map2loop_tools/thickness_calculator_widget.ui b/loopstructural/gui/map2loop_tools/thickness_calculator_widget.ui
@@ -48,8 +48,7 @@
       </widget>
      </item>
      <item row="2" column="1">
-      <widget class="QgsMapLayerComboBox" name="geologyLayerComboBox">
-      </widget>
+      <widget class="QgsMapLayerComboBox" name="geologyLayerComboBox"/>
      </item>
      <item row="3" column="0">
       <widget class="QLabel" name="unitNameFieldLabel">
@@ -69,69 +68,66 @@
       </widget>
      </item>
      <item row="4" column="1">
-      <widget class="QgsMapLayerComboBox" name="basalContactsComboBox">
-      </widget>
+      <widget class="QgsMapLayerComboBox" name="basalContactsComboBox"/>
      </item>
-     <item row="5" column="0">
+     <item row="6" column="0">
       <widget class="QLabel" name="sampledContactsLabel">
        <property name="text">
         <string>Sampled Contacts Layer:</string>
        </property>
       </widget>
      </item>
-     <item row="5" column="1">
-      <widget class="QgsMapLayerComboBox" name="sampledContactsComboBox">
-      </widget>
+     <item row="6" column="1">
+      <widget class="QgsMapLayerComboBox" name="sampledContactsComboBox"/>
      </item>
-     <item row="6" column="0">
+     <item row="7" column="0">
       <widget class="QLabel" name="structureLayerLabel">
        <property name="text">
         <string>Structure/Orientation Layer:</string>
        </property>
       </widget>
      </item>
-     <item row="6" column="1">
-      <widget class="QgsMapLayerComboBox" name="structureLayerComboBox">
-      </widget>
+     <item row="7" column="1">
+      <widget class="QgsMapLayerComboBox" name="structureLayerComboBox"/>
      </item>
-     <item row="7" column="0">
+     <item row="8" column="0">
       <widget class="QLabel" name="dipFieldLabel">
        <property name="text">
         <string>Dip Field:</string>
        </property>
       </widget>
      </item>
-     <item row="7" column="1">
+     <item row="8" column="1">
       <widget class="QgsFieldComboBox" name="dipFieldComboBox"/>
      </item>
-     <item row="8" column="0">
+     <item row="9" column="0">
       <widget class="QLabel" name="dipDirFieldLabel">
        <property name="text">
         <string>Dip Direction Field:</string>
        </property>
       </widget>
      </item>
-     <item row="8" column="1">
+     <item row="9" column="1">
       <widget class="QgsFieldComboBox" name="dipDirFieldComboBox"/>
      </item>
-     <item row="9" column="0">
+     <item row="10" column="0">
       <widget class="QLabel" name="orientationTypeLabel">
        <property name="text">
         <string>Orientation Type:</string>
        </property>
       </widget>
      </item>
-     <item row="9" column="1">
+     <item row="10" column="1">
       <widget class="QComboBox" name="orientationTypeComboBox"/>
      </item>
-     <item row="10" column="0">
+     <item row="11" column="0">
       <widget class="QLabel" name="maxLineLengthLabel">
        <property name="text">
         <string>Max Line Length:</string>
        </property>
       </widget>
      </item>
-     <item row="10" column="1">
+     <item row="11" column="1">
       <widget class="QgsDoubleSpinBox" name="maxLineLengthSpinBox">
        <property name="maximum">
         <double>1000000.000000000000000</double>
@@ -141,6 +137,30 @@
        </property>
       </widget>
      </item>
+     <item row="12" column="0">
+      <widget class="QLabel" name="saveDebugLabel">
+       <property name="text">
+        <string>Save Debug Files:</string>
+       </property>
+      </widget>
+     </item>
+     <item row="12" column="1">
+      <widget class="QCheckBox" name="saveDebugCheckBox">
+       <property name="text">
+        <string>Save debugging files (location tracking, lines)</string>
+       </property>
+      </widget>
+     </item>
+     <item row="5" column="0">
+      <widget class="QLabel" name="basalContactUnitnameLabel">
+       <property name="text">
+        <string>Basal Contact Unitname:</string>
+       </property>
+      </widget>
+     </item>
+     <item row="5" column="1">
+      <widget class="QgsFieldComboBox" name="basalUnitNameFieldComboBox"/>
+     </item>
     </layout>
    </item>
    <item>
@@ -167,19 +187,19 @@
  </widget>
  <customwidgets>
   <customwidget>
-   <class>QgsMapLayerComboBox</class>
-   <extends>QComboBox</extends>
-   <header>qgsmaplayercombobox.h</header>
+   <class>QgsDoubleSpinBox</class>
+   <extends>QDoubleSpinBox</extends>
+   <header>qgsdoublespinbox.h</header>
   </customwidget>
   <customwidget>
    <class>QgsFieldComboBox</class>
    <extends>QComboBox</extends>
    <header>qgsfieldcombobox.h</header>
   </customwidget>
   <customwidget>
-   <class>QgsDoubleSpinBox</class>
-   <extends>QDoubleSpinBox</extends>
-   <header>qgsdoublespinbox.h</header>
+   <class>QgsMapLayerComboBox</class>
+   <extends>QComboBox</extends>
+   <header>qgsmaplayercombobox.h</header>
   </customwidget>
  </customwidgets>
  <resources/>
diff --git a/loopstructural/main/m2l_api.py b/loopstructural/main/m2l_api.py
@@ -296,6 +296,7 @@ def calculate_thickness(
     max_line_length=None,
     stratigraphic_order=None,
     updater=None,
+    basal_contacts_unit_name=None,
 ):
     """Calculate thickness using map2loop thickness calculators.
 
@@ -339,6 +340,11 @@ def calculate_thickness(
     # Convert layers to GeoDataFrames
     geology_gdf = qgsLayerToGeoDataFrame(geology)
     basal_contacts_gdf = qgsLayerToGeoDataFrame(basal_contacts)
+    basal_contacts_gdf = (
+        basal_contacts_gdf.rename(columns={basal_contacts_unit_name: 'basal_unit'})
+        if basal_contacts_unit_name
+        else basal_contacts_gdf
+    )
     sampled_contacts_gdf = qgsLayerToGeoDataFrame(sampled_contacts)
     structure_gdf = qgsLayerToDataFrame(structure)
     bounding_box = {
@@ -397,6 +403,7 @@ def calculate_thickness(
     units_unique = units.drop_duplicates(subset=[unit_name_field]).reset_index(drop=True)
     units = pd.DataFrame({'name': units_unique[unit_name_field]})
     basal_contacts_gdf['type'] = 'BASAL'  # required by calculator
+
     thickness = calculator.compute(
         units,
         stratigraphic_order,
@@ -405,7 +412,12 @@ def calculate_thickness(
         geology_gdf,
         sampled_contacts_gdf,
     )
+    res = {'thicknesses': thickness}
     if updater:
         updater(f"Thickness calculation complete: {len(thickness)} records")
-    
-    return thickness
+    if hasattr(calculator, 'lines'):
+        res['lines'] = calculator.lines
+    if hasattr(calculator, 'location_tracking'):
+        res['location_tracking'] = calculator.location_tracking
+
+    return res
