fix: Adding colours to surfaces

Author: lachlangrose

LoopStructural/datatypes/_surface.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Optional,Union
 import numpy as np
 import io
 from LoopStructural.utils import getLogger
@@ -11,12 +11,13 @@
 class Surface:
     vertices: np.ndarray = field(default_factory=lambda: np.array([[0, 0, 0]]))
     triangles: np.ndarray = field(default_factory=lambda: np.array([[0, 0, 0]]))
+    colour: Optional[Union[str,np.ndarray]] = field(default_factory=lambda: None)
+
     normals: Optional[np.ndarray] = None
     name: str = 'surface'
     values: Optional[np.ndarray] = None
     properties: Optional[dict] = None
     cell_properties: Optional[dict] = None
-
     @property
     def triangle_area(self):
         """_summary_

LoopStructural/modelling/core/geological_model.py

@@ -1798,7 +1798,7 @@ def get_stratigraphic_surfaces(self, units: List[str] = [], bottoms: bool = True
             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()
+                    values.to_list(), self.bounding_box, name=names.loc[values.index].to_list(), colours=unit_table.loc[unit_table['feature_name'] == u, 'colour'].tolist()
                 )
             )
 

LoopStructural/modelling/features/_base_geological_feature.py

@@ -274,6 +274,7 @@ def surfaces(
         value: Union[float, int, List[Union[float, int]]],
         bounding_box=None,
         name: Optional[Union[List[str], str]] = None,
+        colours:Optional[Union[str,np.ndarray]]=None,
     ) -> surface_list:
         """Find the surfaces of the geological feature at a given value
 
@@ -305,14 +306,14 @@ def surfaces(
             isosurfacer = LoopIsosurfacer(bounding_box, callable=callable)
             if name is None and self.name is not None:
                 name = self.name
-            surfaces = isosurfacer.fit(value, name)
+            surfaces = isosurfacer.fit(value, name,colours=colours)
         except Exception as e:
             logger.error(f"Failed to create surface for {self.name} at value {value}")
             logger.error(e)
             surfaces = []
         finally:
             self.regions = regions
-
+        
         return surfaces
 
     def scalar_field(self, bounding_box=None):

LoopStructural/utils/_surface.py

@@ -64,6 +64,7 @@ def fit(
         values: Optional[Union[list, int, float]],
         name: Optional[Union[List[str], str]] = None,
         local=False,
+        colours: Optional[List]=None,
     ) -> surface_list:
         """Extract isosurfaces from the interpolator
 
@@ -120,7 +121,9 @@ def fit(
             names = [name] * len(isovalues)
         if isinstance(name, list):
             names = name
-        for name, isovalue in zip(names, isovalues):
+        if colours is None:
+            colours = [None]*len(isovalues)
+        for name, isovalue,colour in zip(names, isovalues,colours):
             try:
                 step_vector = (self.bounding_box.maximum - self.bounding_box.origin) / (
                     np.array(self.bounding_box.nsteps) - 1
@@ -150,6 +153,7 @@ def fit(
                     normals=normals,
                     name=f"{name}_{isovalue}",
                     values=values,
+                    colour=colour,
                 )
             )
         return surfaces