adding placeholder to associate data with a stratigraphic unit

add ability to map group name to features e.g. Group 1: feature_name
Make arguments named only
Update stratigraphic surfaces to use new stratigraphic column

Author: lachlangrose

LoopStructural/modelling/core/geological_model.py

@@ -526,6 +526,9 @@ def set_stratigraphic_column(self, stratigraphic_column, cmap="tab20"):
                     colour=stratigraphic_column[g][u].get("colour", None),
                     thickness=thickness,
                 )
+            self.stratigraphic_column.add_unconformity(
+                name=''.join([g, 'unconformity']),
+            )
 
 
     def create_and_add_foliation(
@@ -1704,31 +1707,20 @@ def get_stratigraphic_surfaces(self, units: List[str] = [], bottoms: bool = True
         units = []
         if self.stratigraphic_column is None:
             return []
-        for group in self.stratigraphic_column.keys():
-            if group == "faults":
+        units = self.stratigraphic_column.get_isovalues()
+        for name, u in units.items():
+            if u['group'] not in self:
+                logger.warning(f"Group {u['group']} not found in model")
                 continue
-            for series in self.stratigraphic_column[group].values():
-                series['feature_name'] = group
-                units.append(series)
-        unit_table = pd.DataFrame(units)
-        for u in unit_table['feature_name'].unique():
-
-            values = unit_table.loc[unit_table['feature_name'] == u, 'min' if bottoms else 'max']
-            if 'name' not in unit_table.columns:
-                unit_table['name'] = unit_table['feature_name']
-
-            names = unit_table[unit_table['feature_name'] == u]['name']
-            values = values.loc[~np.logical_or(values == np.inf, values == -np.inf)]
-            surfaces.extend(
-                self.get_feature_by_name(u).surfaces(
-                    values.to_list(),
-                    self.bounding_box,
-                    name=names.loc[values.index].to_list(),
-                    colours=unit_table.loc[unit_table['feature_name'] == u, 'colour'].tolist()[
-                        1:
-                    ],  # we don't isosurface basement, no value
-                )
-            )
+            feature = self.get_feature_by_name(u['group'])
+            
+            surfaces.extend(feature.surfaces([u['value']],
+                                                                            self.bounding_box,
+                                                                            name=name,
+                                                                            colours=[u['colour']]
+                                                                            ))
+
+      
 
         return surfaces
 

LoopStructural/modelling/core/stratigraphic_column.py

@@ -47,7 +47,7 @@ class StratigraphicUnit(StratigraphicColumnElement):
     A class to represent a stratigraphic unit.
     """
 
-    def __init__(self, *, uuid=None, name=None, colour=None, thickness=None):
+    def __init__(self, *, uuid=None, name=None, colour=None, thickness=None, data=None):
         """
         Initializes the StratigraphicUnit with a name and an optional description.
         """
@@ -57,6 +57,7 @@ def __init__(self, *, uuid=None, name=None, colour=None, thickness=None):
             colour = rng.random(3)
         self.colour = colour
         self.thickness = thickness
+        self.data = data
         self.element_type = StratigraphicColumnElementType.UNIT
 
     def to_dict(self):
@@ -137,6 +138,18 @@ def from_dict(cls, data):
         )
         uuid = data.get("uuid", None)
         return cls(uuid=uuid, name=name, unconformity_type=unconformity_type)
+class StratigraphicGroup:
+    """
+    A class to represent a group of stratigraphic units.
+    This class is not fully implemented and serves as a placeholder for future development.
+    """
+
+    def __init__(self, name=None, units=None):
+        """
+        Initializes the StratigraphicGroup with a name and an optional list of units.
+        """
+        self.name = name
+        self.units = units if units is not None else []
 
 
 class StratigraphicColumn:
@@ -150,13 +163,15 @@ def __init__(self):
         Initializes the StratigraphicColumn with a name and a list of layers.
         """
         self.order = [StratigraphicUnit(name='Basement', colour='grey', thickness=np.inf),StratigraphicUnconformity(name='Base Unconformity', unconformity_type=UnconformityType.ERODE)]
+        self.group_mapping = {}
     def clear(self):
         """
         Clears the stratigraphic column, removing all elements.
         """
         self.order = [StratigraphicUnit(name='Basement', colour='grey', thickness=np.inf),StratigraphicUnconformity(name='Base Unconformity', unconformity_type=UnconformityType.ERODE)]
+        self.group_mapping = {}
 
-    def add_unit(self, name, colour, thickness=None, where='top'):
+    def add_unit(self, name,*, colour=None, thickness=None, where='top'):
         unit = StratigraphicUnit(name=name, colour=colour, thickness=thickness)
 
         if where == 'top':
@@ -178,7 +193,7 @@ def remove_unit(self, uuid):
                 return True
         return False
 
-    def add_unconformity(self, name, unconformity_type=UnconformityType.ERODE, where='top' ):
+    def add_unconformity(self, name, *, unconformity_type=UnconformityType.ERODE, where='top' ):
         unconformity = StratigraphicUnconformity(
             uuid=None, name=name, unconformity_type=unconformity_type
         )
@@ -226,31 +241,41 @@ def get_elements(self):
 
     def get_groups(self):
         groups = []
-        group = []
+        i=0
+        group = StratigraphicGroup(
+            name=(
+                f'Group_{i}'
+                if f'Group_{i}' not in self.group_mapping
+                else self.group_mapping[f'Group_{i}']
+            )
+        )
         for e in self.order:
             if isinstance(e, StratigraphicUnit):
-                group.append(e)
+                group.units.append(e)
             else:
-                if group:
+                if group.units:
                     groups.append(group)
-                    group = []
+                    i+=1
+                    group = StratigraphicGroup(
+                        name=(
+                            f'Group_{i}'
+                            if f'Group_{i}' not in self.group_mapping
+                            else self.group_mapping[f'Group_{i}']
+                        )
+                    )
         if group:
             groups.append(group)
         return groups
 
     def get_unitname_groups(self):
-        groups = []
+        groups = self.get_groups()
+        groups_list = []
         group = []
-        for e in self.order:
-            if isinstance(e, StratigraphicUnit):
-                group.append(e.name)
-            else:
-                if group:
-                    groups.append(group)
-                    group = []
-        if group:
-            groups.append(group)
-        return groups
+        for g in groups:
+            group = [u.name for u in g.units if isinstance(u, StratigraphicUnit)]
+            groups_list.append(group)
+        return groups_list
+        
 
     def __getitem__(self, uuid):
         """
@@ -289,7 +314,6 @@ def update_element(self, unit_data: Dict):
                 unit_data.get('unconformity_type', element.unconformity_type.value)
             )
 
-
     def __str__(self):
         """
         Returns a string representation of the stratigraphic column, listing all elements.
@@ -320,20 +344,20 @@ def from_dict(cls, data):
                 element = StratigraphicUnit.from_dict(element_data)
             column.add_element(element)
         return column
-    
+
     def get_isovalues(self) -> Dict[str, float]:
         """
         Returns a dictionary of isovalues for the stratigraphic units in the column.
         """
         surface_values = {}
         for g in reversed(self.get_groups()):
             v = 0
-            for u in g:
-                surface_values[u.name] = v
+            for u in g.units:
+                surface_values[u.name] = {'value':v,'group':g.name,'colour':u.colour}
                 v += u.thickness
         return surface_values
-    
-    def plot(self, ax=None, **kwargs):
+
+    def plot(self,*, ax=None, **kwargs):
         import matplotlib.pyplot as plt
         from matplotlib import cm
         from matplotlib.patches import Polygon