diff --git a/backend/src/services/smda_access/stratigraphy_utils.py b/backend/src/services/smda_access/stratigraphy_utils.py
index f4eb665cd..83b44ef99 100644
--- a/backend/src/services/smda_access/stratigraphy_utils.py
+++ b/backend/src/services/smda_access/stratigraphy_utils.py
@@ -1,72 +1,84 @@
-from typing import List, Dict
+from typing import List, Dict, Optional
+from pydantic import BaseModel
 
 from .types import StratigraphicUnit, StratigraphicSurface, StratigraphicFeature
 
 
-def create_hierarchical_structure(strat_units: List[StratigraphicUnit]) -> List[StratigraphicUnit]:
-    """Organizes the stratigraphic units into a hierarchical nested list based on parent relationships."""
-    unit_by_id = {unit.identifier: {"unit": unit, "children": []} for unit in strat_units}
+class HierarchicalStratigraphicUnit(BaseModel):
+    """A stratigraphic unit within a hierarchical structure, i.e a multi-level stratigraphical column"""
+
+    unit: StratigraphicUnit
+    children: Optional[List["HierarchicalStratigraphicUnit"]]
+
+
+def create_hierarchical_structure(strat_units: List[StratigraphicUnit]) -> List[HierarchicalStratigraphicUnit]:
+    """Create a hierarchical structure of stratigraphic units using strat_unit_parent.
+    On Drogon this will result in the following structure:
+    - VOLANTIS GP.
+        - Valysar Fm.
+        - Therys Fm.
+        - Volon Fm
+    """
+    unit_by_id = {unit.identifier: HierarchicalStratigraphicUnit(unit=unit, children=[]) for unit in strat_units}
     roots = []
 
     for unit in strat_units:
         if unit.strat_unit_parent and unit.strat_unit_parent in unit_by_id:
             parent = unit_by_id[unit.strat_unit_parent]
-            parent["children"].append(unit)  # type: ignore
+            parent.children.append(unit_by_id[unit.identifier])
         else:
-            roots.append(unit)
+            roots.append(unit_by_id[unit.identifier])
 
     return roots
 
 
-def flatten_hierarchical_structure(units: List[StratigraphicUnit], unit_by_id: Dict) -> List[StratigraphicUnit]:
-    """Flatten the hierarchical structure into a single list of stratigraphic units, preserving the order."""
+def flatten_hierarchical_structure(units: List[HierarchicalStratigraphicUnit]) -> List[StratigraphicUnit]:
+    """Flattens the hierarchy of stratigraphic units to a list of stratigraphic units preserving the order."""
     flattened_list = []
 
-    for unit in units:
-        flattened_list.append(unit)
-        if unit.identifier in unit_by_id:
-            flattened_list.extend(flatten_hierarchical_structure(unit_by_id[unit.identifier]["children"], unit_by_id))
+    for hierarchical_unit in units:
+        flattened_list.append(hierarchical_unit.unit)
+        flattened_list.extend(flatten_hierarchical_structure(hierarchical_unit.children))
+
     return flattened_list
 
 
 def flatten_hierarchical_structure_to_surface_name(
-    units: List[StratigraphicUnit], unit_by_id: Dict
+    units: List[HierarchicalStratigraphicUnit],
 ) -> List[StratigraphicSurface]:
-    """Flatten the hierarchical structure into a single list of stratigraphical top/unit/base names, preserving the order."""
+    """Flattens the hierarchy of stratigraphic units to a list of stratigraphic surfaces preserving the order.
+    On Drogon this will result in the following list:
+    - Volantis GP. Top
+    - Volantis GP.
+    - Valysar Fm. Top
+    - Valysar Fm.
+    - Therys Fm. Top
+    - Therys Fm.
+    - Volon Fm. Top
+    - Volon Fm.
+    - Volantis GP. Base
+    """
     flattened_list = []
 
-    for unit in units:
+    for hierarchical_unit in units:
+        unit = hierarchical_unit.unit
         flattened_list.append(StratigraphicSurface(name=unit.top, feature=StratigraphicFeature.HORIZON))
         flattened_list.append(StratigraphicSurface(name=unit.identifier, feature=StratigraphicFeature.ZONE))
-        if unit.identifier in unit_by_id:
-            flattened_list.extend(
-                flatten_hierarchical_structure_to_surface_name(unit_by_id[unit.identifier]["children"], unit_by_id)
-            )
+        flattened_list.extend(flatten_hierarchical_structure_to_surface_name(hierarchical_unit.children))
         flattened_list.append(StratigraphicSurface(name=unit.base, feature=StratigraphicFeature.HORIZON))
+
     return flattened_list
 
 
 def sort_stratigraphic_names_by_hierarchy(strat_units: List[StratigraphicUnit]) -> List[StratigraphicSurface]:
-    """Sort stratigraphic top/unit/base by hierarchy."""
-    unit_by_id = {unit.identifier: {"unit": unit, "children": []} for unit in strat_units}
-
-    for unit in strat_units:
-        if unit.strat_unit_parent and unit.strat_unit_parent in unit_by_id:
-            unit_by_id[unit.strat_unit_parent]["children"].append(unit)  # type: ignore
-
-    roots = [unit for unit in strat_units if not unit.strat_unit_parent]
-    sorted_units = flatten_hierarchical_structure_to_surface_name(roots, unit_by_id)
-    return sorted_units
+    """Creates hierarchy then flattens to a list of surfaces."""
+    hierarchical_units = create_hierarchical_structure(strat_units)
+    sorted_surfaces = flatten_hierarchical_structure_to_surface_name(hierarchical_units)
+    return sorted_surfaces
 
 
 def sort_stratigraphic_units_by_hierarchy(strat_units: List[StratigraphicUnit]) -> List[StratigraphicUnit]:
-    """Sort stratigraphic units by hierarchy."""
-    unit_by_id = {unit.identifier: {"unit": unit, "children": []} for unit in strat_units}
-
-    for unit in strat_units:
-        if unit.strat_unit_parent and unit.strat_unit_parent in unit_by_id:
-            unit_by_id[unit.strat_unit_parent]["children"].append(unit)  # type: ignore
-
-    roots = [unit for unit in strat_units if not unit.strat_unit_parent]
-    sorted_units = flatten_hierarchical_structure(roots, unit_by_id)
+    """Creates hierarchy then flattens to a list of units."""
+    hierarchical_units = create_hierarchical_structure(strat_units)
+    sorted_units = flatten_hierarchical_structure(hierarchical_units)
     return sorted_units