diff --git a/backend/src/services/smda_access/mocked_drogon_smda_access/_mocked_stratigraphy_access.py b/backend/src/services/smda_access/mocked_drogon_smda_access/_mocked_stratigraphy_access.py index 5868a44fd..f6d1c9e30 100644 --- a/backend/src/services/smda_access/mocked_drogon_smda_access/_mocked_stratigraphy_access.py +++ b/backend/src/services/smda_access/mocked_drogon_smda_access/_mocked_stratigraphy_access.py @@ -2,7 +2,7 @@ from ..types import StratigraphicUnit, StratigraphicSurface -from ..stratigraphy_access import sort_stratigraphic_names_by_hierarchy, sort_stratigraphic_units_by_hierarchy +from ..stratigraphy_utils import sort_stratigraphic_names_by_hierarchy, sort_stratigraphic_units_by_hierarchy STRAT_UNITS = [ StratigraphicUnit( diff --git a/backend/src/services/smda_access/stratigraphy_access.py b/backend/src/services/smda_access/stratigraphy_access.py index 68277ca4b..686d90b53 100644 --- a/backend/src/services/smda_access/stratigraphy_access.py +++ b/backend/src/services/smda_access/stratigraphy_access.py @@ -1,9 +1,8 @@ -# type: ignore - -from typing import List, Dict +from typing import List from .queries.get_stratigraphic_units import get_stratigraphic_units -from .types import StratigraphicUnit, StratigraphicSurface, StratigraphicFeature +from .types import StratigraphicUnit, StratigraphicSurface +from .stratigraphy_utils import sort_stratigraphic_names_by_hierarchy, sort_stratigraphic_units_by_hierarchy class StratigraphyAccess: @@ -23,71 +22,3 @@ def get_stratigraphic_surfaces(self, stratigraphic_column_identifier: str) -> Li sorted_units = sort_stratigraphic_names_by_hierarchy(stratigraphic_units) return sorted_units - - -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} - 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) - else: - roots.append(unit) - - return roots - - -def flatten_hierarchical_structure(units: List[Dict], unit_by_id: Dict) -> List[StratigraphicUnit]: - """Flatten the hierarchical structure into a single 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)) - - return flattened_list - - -def flatten_hierarchical_structure_to_surface_name(units: List[Dict], unit_by_id: Dict) -> List[StratigraphicSurface]: - """Flatten the hierarchical structure into a single list of stratigraphical top/unit/base names, preserving the order.""" - flattened_list = [] - - for unit in units: - 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.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) - - roots = [data["unit"] for data in unit_by_id.values() if not data["unit"].strat_unit_parent] - sorted_units = flatten_hierarchical_structure_to_surface_name(roots, unit_by_id) - return sorted_units - - -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) - - roots = [data["unit"] for data in unit_by_id.values() if not data["unit"].strat_unit_parent] - sorted_units = flatten_hierarchical_structure(roots, unit_by_id) - return sorted_units diff --git a/backend/src/services/smda_access/stratigraphy_utils.py b/backend/src/services/smda_access/stratigraphy_utils.py new file mode 100644 index 000000000..2b98a97e2 --- /dev/null +++ b/backend/src/services/smda_access/stratigraphy_utils.py @@ -0,0 +1,71 @@ +from typing import List, Dict + +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} + 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) + else: + roots.append(unit) + + return roots + + +def flatten_hierarchical_structure(units: List[Dict], unit_by_id: Dict) -> List[StratigraphicUnit]: + """Flatten the hierarchical structure into a single 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)) + + return flattened_list + + +def flatten_hierarchical_structure_to_surface_name(units: List[Dict], unit_by_id: Dict) -> List[StratigraphicSurface]: + """Flatten the hierarchical structure into a single list of stratigraphical top/unit/base names, preserving the order.""" + flattened_list = [] + + for unit in units: + 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.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) + + roots = [data["unit"] for data in unit_by_id.values() if not data["unit"].strat_unit_parent] + sorted_units = flatten_hierarchical_structure_to_surface_name(roots, unit_by_id) + return sorted_units + + +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) + + roots = [data["unit"] for data in unit_by_id.values() if not data["unit"].strat_unit_parent] + sorted_units = flatten_hierarchical_structure(roots, unit_by_id) + return sorted_units