From 8c948df1d16aa0dd35a1373117266c4a4b07527b Mon Sep 17 00:00:00 2001 From: Chain-Frost Date: Fri, 7 Nov 2025 14:11:09 +0800 Subject: [PATCH] Add location filter to POMM combine workflow --- ryan-scripts/TUFLOW-python/POMM_combine.py | 7 +++++++ ryan_library/functions/tuflow/pomm_combine.py | 15 ++++++++++++++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/ryan-scripts/TUFLOW-python/POMM_combine.py b/ryan-scripts/TUFLOW-python/POMM_combine.py index d57e0dad..b7441c24 100644 --- a/ryan-scripts/TUFLOW-python/POMM_combine.py +++ b/ryan-scripts/TUFLOW-python/POMM_combine.py @@ -9,6 +9,10 @@ console_log_level = "INFO" # or "DEBUG" +# Update this tuple to restrict processing to specific PO/Location values. +# Leave empty to include every location found in the POMM files. +LOCATIONS_TO_INCLUDE: tuple[str, ...] = () + def main() -> None: """Wrapper script to merge POMM results. @@ -24,10 +28,13 @@ def main() -> None: if not change_working_directory(target_dir=script_directory): return + locations_to_include: tuple[str, ...] | None = LOCATIONS_TO_INCLUDE or None + main_processing( paths_to_process=[script_directory], include_data_types=["POMM"], console_log_level=console_log_level, + locations_to_include=locations_to_include, ) print() print_library_version() diff --git a/ryan_library/functions/tuflow/pomm_combine.py b/ryan_library/functions/tuflow/pomm_combine.py index 7027d602..6a29c826 100644 --- a/ryan_library/functions/tuflow/pomm_combine.py +++ b/ryan_library/functions/tuflow/pomm_combine.py @@ -1,6 +1,8 @@ """Modern POMM combination utilities.""" +from collections.abc import Collection from pathlib import Path + import pandas as pd from loguru import logger @@ -8,6 +10,7 @@ collect_files, process_files_in_parallel, ) +from ryan_library.processors.tuflow.base_processor import BaseProcessor from ryan_library.processors.tuflow.processor_collection import ProcessorCollection from ryan_library.functions.file_utils import ensure_output_directory from ryan_library.functions.misc_functions import ExcelExporter @@ -19,12 +22,15 @@ def main_processing( paths_to_process: list[Path], include_data_types: list[str] | None = None, console_log_level: str = "INFO", + locations_to_include: Collection[str] | None = None, ) -> None: """Generate merged culvert data and export the results.""" if include_data_types is None: include_data_types = ["POMM"] + normalized_locations: frozenset[str] = BaseProcessor.normalize_locations(locations_to_include) + with setup_logger(console_log_level=console_log_level) as log_queue: csv_file_list: list[Path] = collect_files( paths_to_process=paths_to_process, @@ -35,7 +41,14 @@ def main_processing( logger.info("No valid files found to process.") return - results_set: ProcessorCollection = process_files_in_parallel(file_list=csv_file_list, log_queue=log_queue) + results_set: ProcessorCollection = process_files_in_parallel( + file_list=csv_file_list, + log_queue=log_queue, + location_filter=normalized_locations if normalized_locations else None, + ) + + if normalized_locations: + results_set.filter_locations(normalized_locations) export_results(results=results_set) logger.info("End of POMM results combination processing")