From de25c4282f4425873d6099bb8c1f9a36e2d86f6a Mon Sep 17 00:00:00 2001 From: ramank1137 Date: Thu, 13 Aug 2026 18:08:46 +0530 Subject: [PATCH 1/6] added changes for all the hydrology to merge it with new local-compute branch --- computing/api.py | 315 ++- computing/config.yaml | 21 +- computing/config_loader.py | 7 +- computing/hydrology_gpu/__init__.py | 1 + .../hydrology_gpu/algorithms/__init__.py | 49 + computing/hydrology_gpu/algorithms/runoff.py | 553 ++++ .../algorithms/tiled_timeseries.py | 258 ++ .../hydrology_gpu/algorithms/timeseries.py | 119 + computing/hydrology_gpu/config/__init__.py | 49 + computing/hydrology_gpu/downloads/__init__.py | 166 ++ computing/hydrology_gpu/downloads/dem.py | 95 + computing/hydrology_gpu/downloads/lulc.py | 31 + computing/hydrology_gpu/downloads/rainfall.py | 415 +++ computing/hydrology_gpu/downloads/soil.py | 17 + computing/hydrology_gpu/et_download.py | 700 +++++ computing/hydrology_gpu/lulc_mapping.py | 128 + computing/hydrology_gpu/runoff.py | 202 ++ computing/hydrology_gpu/utils.py | 399 +++ computing/hydrology_gpu/watershed_boundary.py | 448 +++ computing/mws/et_download.py | 110 + computing/mws/generate_hydrology_local.py | 2457 +++++++++++++++++ computing/mws/runoff_gpu.py | 302 ++ computing/tasks.py | 13 +- computing/urls.py | 12 + installation/environment.yml | 17 +- 25 files changed, 6841 insertions(+), 43 deletions(-) create mode 100644 computing/hydrology_gpu/__init__.py create mode 100644 computing/hydrology_gpu/algorithms/__init__.py create mode 100644 computing/hydrology_gpu/algorithms/runoff.py create mode 100644 computing/hydrology_gpu/algorithms/tiled_timeseries.py create mode 100644 computing/hydrology_gpu/algorithms/timeseries.py create mode 100644 computing/hydrology_gpu/config/__init__.py create mode 100644 computing/hydrology_gpu/downloads/__init__.py create mode 100644 computing/hydrology_gpu/downloads/dem.py create mode 100644 computing/hydrology_gpu/downloads/lulc.py create mode 100644 computing/hydrology_gpu/downloads/rainfall.py create mode 100644 computing/hydrology_gpu/downloads/soil.py create mode 100644 computing/hydrology_gpu/et_download.py create mode 100644 computing/hydrology_gpu/lulc_mapping.py create mode 100644 computing/hydrology_gpu/runoff.py create mode 100644 computing/hydrology_gpu/utils.py create mode 100644 computing/hydrology_gpu/watershed_boundary.py create mode 100644 computing/mws/et_download.py create mode 100644 computing/mws/generate_hydrology_local.py create mode 100644 computing/mws/runoff_gpu.py diff --git a/computing/api.py b/computing/api.py index 72e2bb75..18bca62f 100644 --- a/computing/api.py +++ b/computing/api.py @@ -160,7 +160,13 @@ generate_soge_vector_local as generate_soge_vector_local_task, ) from .misc.stream_order import generate_stream_order -from .mws.generate_hydrology import generate_hydrology +from .mws.generate_hydrology import generate_hydrology as generate_hydrology_gee_task +from .mws.generate_hydrology_local import ( + generate_hydrology_base_layer as generate_hydrology_base_layer_task, + generate_hydrology as generate_hydrology_local_task, +) +from .mws.et_download import et_download as et_download_task +from .mws.runoff_gpu import generate_runoff_gpu as generate_runoff_gpu_task from .mws.mws import mws_layer from .mws.mws_centroid import generate_mws_centroid_data from .mws.mws_centroid_local_compute import ( @@ -241,6 +247,27 @@ logger = logging.getLogger(__name__) +def _get_pan_india_flag(request): + value = request.data.get( + "pan_india", + request.data.get( + "pan-india", + request.data.get("panIndia", False), + ), + ) + if isinstance(value, bool): + return value + return str(value).strip().lower() in {"1", "true", "yes", "y"} + + +def _has_any_payload_field(request, fields): + return any(field in request.data for field in fields) + + +PAN_INDIA_PAYLOAD_FIELDS = ("pan_india", "pan-india", "panIndia") +PAN_INDIA_LOCATION_FIELDS = ("state", "district", "block", "tehsil", "year") + + @api_security_check(allowed_methods="POST") @schema(None) def generate_admin_boundary(request): @@ -404,66 +431,290 @@ def generate_mws_layer(request): return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) + @api_security_check(allowed_methods="POST") @schema(None) def generate_fortnightly_hydrology(request): print("Inside generate_fortnightly_hydrology") try: - state = request.data.get("state") - district = request.data.get("district") - block = request.data.get("block") - start_year = int(request.data.get("start_year")) - end_year = int(request.data.get("end_year")) - gee_account_id = request.data.get("gee_account_id") - generate_hydrology.apply_async( + return _generate_tehsil_hydrology(request, is_annual=False) + except ValueError as e: + print("Invalid request in generate_fortnightly_hydrology api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) + except Exception as e: + print("Exception in generate_fortnightly_hydrology api :: ", e) + return Response( + {"Exception": str(e)}, + status=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + +@api_view(["POST"]) +@schema(None) +def generate_annual_hydrology(request): + print("Inside generate_annual_hydrology") + try: + return _generate_tehsil_hydrology(request, is_annual=True) + except ValueError as e: + print("Invalid request in generate_annual_hydrology api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) + except Exception as e: + print("Exception in generate_annual_hydrology api :: ", e) + return Response( + {"Exception": str(e)}, + status=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + +def _generate_tehsil_hydrology(request, is_annual): + compute = _get_compute_mode(request) + if _has_any_payload_field(request, PAN_INDIA_PAYLOAD_FIELDS): + raise ValueError( + "Do not pass pan_india to the tehsil hydrology API. " + "Use /api/v1/pan-india/hydrology_fortnightly/ or " + "/api/v1/pan-india/hydrology_annual/ for Pan-India generation." + ) + + state = request.data.get("state") + district = request.data.get("district") + block = request.data.get("block") + if not all([state, district, block]): + raise ValueError("state, district, and block are required") + + if request.data.get("start_year") is None or request.data.get("end_year") is None: + raise ValueError("start_year and end_year are required") + start_year = int(request.data.get("start_year")) + end_year = int(request.data.get("end_year")) + if end_year < start_year: + raise ValueError("end_year must be greater than or equal to start_year") + + if compute == "gee": + task = generate_hydrology_gee_task.apply_async( kwargs={ "state": state, "district": district, "block": block, "start_year": start_year, "end_year": end_year, - "gee_account_id": gee_account_id, - "is_annual": False, + "gee_account_id": request.data.get("gee_account_id"), + "is_annual": is_annual, + }, + queue="nrm", + ) + source = "gee" + success = "hydrology GEE task initiated" + else: + if start_year != 2017: + raise ValueError( + "Local hydrology clipping must start from start_year=2017 because " + "the fortnightly cadence and cumulative G are anchored at 2017-07-01" + ) + task = generate_hydrology_local_task.apply_async( + kwargs={ + "state": state, + "district": district, + "block": block, + "start_year": start_year, + "end_year": end_year, + "is_annual": is_annual, + "pan_india": False, + "overwrite": request.data.get("overwrite", False), + }, + queue="nrm", + ) + source = "base_layer_clip" + success = "hydrology clipping task initiated" + + return Response( + { + "Success": success, + "task_id": task.id, + "compute": compute, + "scope": "tehsil", + "source": source, + "start_year": start_year, + "end_year": end_year, + "is_annual": bool(is_annual), + }, + status=status.HTTP_200_OK, + ) + + +def _generate_pan_india_hydrology_base_layer(request, is_annual): + compute = _get_compute_mode(request) + if compute != "local": + raise ValueError( + "Pan-India hydrology generation supports compute='local' only" + ) + if _has_any_payload_field(request, PAN_INDIA_PAYLOAD_FIELDS): + raise ValueError( + "Do not pass pan_india to the Pan-India hydrology API; " + "the /api/v1/pan-india/ route already defines the scope." + ) + forbidden_fields = [ + field for field in PAN_INDIA_LOCATION_FIELDS if field in request.data + ] + if forbidden_fields: + raise ValueError( + "Pan-India hydrology API accepts only compute, start_year, " + "end_year, and overwrite; do not pass " + f"{', '.join(forbidden_fields)}" + ) + + if request.data.get("start_year") is None or request.data.get("end_year") is None: + raise ValueError("start_year and end_year are required") + start_year = int(request.data.get("start_year")) + end_year = int(request.data.get("end_year")) + if end_year <= start_year: + raise ValueError("end_year must be greater than start_year") + if end_year != start_year + 1: + raise ValueError( + "Hydrology base-layer generation supports one hydrological year " + "at a time; use end_year=start_year+1" + ) + task = generate_hydrology_base_layer_task.apply_async( + kwargs={ + "start_year": start_year, + "end_year": end_year, + "is_annual": is_annual, + "overwrite": request.data.get("overwrite", False), + }, + queue="nrm", + ) + return Response( + { + "Success": "Pan-India hydrology task initiated", + "task_id": task.id, + "compute": compute, + "scope": "pan_india", + "start_year": start_year, + "end_year": end_year, + "year_key": f"{start_year}_{end_year}", + "is_annual": bool(is_annual), + }, + status=status.HTTP_200_OK, + ) + + +@api_view(["POST"]) +@schema(None) +def generate_pan_india_fortnightly_hydrology(request): + print("Inside generate_pan_india_fortnightly_hydrology") + try: + return _generate_pan_india_hydrology_base_layer(request, is_annual=False) + except ValueError as e: + print( + "Invalid request in generate_pan_india_fortnightly_hydrology api :: ", + e, + ) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) + except Exception as e: + print( + "Exception in generate_pan_india_fortnightly_hydrology api :: ", + e, + ) + return Response( + {"Exception": str(e)}, + status=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + +@api_view(["POST"]) +@schema(None) +def generate_pan_india_annual_hydrology(request): + print("Inside generate_pan_india_annual_hydrology") + try: + return _generate_pan_india_hydrology_base_layer(request, is_annual=True) + except ValueError as e: + print("Invalid request in generate_pan_india_annual_hydrology api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) + except Exception as e: + print("Exception in generate_pan_india_annual_hydrology api :: ", e) + return Response( + {"Exception": str(e)}, + status=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + +@api_view(["POST"]) +@schema(None) +def generate_runoff_gpu(request): + print("Inside generate_runoff_gpu") + try: + compute = _get_compute_mode(request) + if compute != "local": + raise ValueError("runoff_gpu currently supports compute='local' only") + + tehsil = request.data.get("tehsil") or request.data.get("block") + pan_india = _get_pan_india_flag(request) + task = generate_runoff_gpu_task.apply_async( + kwargs={ + "state": request.data.get("state"), + "district": request.data.get("district"), + "tehsil": tehsil, + "pan_india": pan_india, + "start_date": request.data.get("start_date"), + "end_date": request.data.get("end_date"), + "start_year": request.data.get("start_year"), + "end_year": request.data.get("end_year"), }, queue="nrm", ) return Response( - {"Success": "Successfully initiated"}, status=status.HTTP_200_OK + { + "Success": "runoff_gpu task initiated", + "task_id": task.id, + }, + status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in generate_runoff_gpu api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: - print("Exception in generate_fortnightly_hydrology api :: ", e) - return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) + print("Exception in generate_runoff_gpu api :: ", e) + return Response( + {"Exception": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR + ) @api_view(["POST"]) @schema(None) -def generate_annual_hydrology(request): - print("Inside generate_annual_hydrology") +def et_download(request): + print("Inside et_download") try: - state = request.data.get("state") - district = request.data.get("district") - block = request.data.get("block") - start_year = int(request.data.get("start_year")) - end_year = int(request.data.get("end_year")) - gee_account_id = request.data.get("gee_account_id") - generate_hydrology.apply_async( + compute = _get_compute_mode(request) + if compute != "local": + raise ValueError("et_download currently supports compute='local' only") + + pan_india = _get_pan_india_flag(request) + task = et_download_task.apply_async( kwargs={ - "state": state, - "district": district, - "block": block, - "start_year": start_year, - "end_year": end_year, - "is_annual": True, - "gee_account_id": gee_account_id, + "pan_india": pan_india, + "start_date": request.data.get("start_date"), + "end_date": request.data.get("end_date"), + "start_year": request.data.get("start_year"), + "end_year": request.data.get("end_year"), + "overwrite": request.data.get("overwrite", False), + "patch_fill": request.data.get("patch_fill", True), }, queue="nrm", ) return Response( - {"Success": "Successfully initiated"}, status=status.HTTP_200_OK + { + "Success": "et_download task initiated", + "task_id": task.id, + }, + status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in et_download api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: - print("Exception in generate_annual_hydrology api :: ", e) - return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) + print("Exception in et_download api :: ", e) + return Response( + {"Exception": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR + ) + @api_view(["POST"]) diff --git a/computing/config.yaml b/computing/config.yaml index d1d1b46d..4e648f5e 100644 --- a/computing/config.yaml +++ b/computing/config.yaml @@ -72,16 +72,19 @@ base_layers: source: s3://corestack-datasets/base_layers/static_layers/canal/canal.geojson type: file - - name: hydrological soil group - local_path: "{DATA_DIR}/base_layers/hydrological_soil_group/hydrological_soil_group.geojson" - source: s3://corestack-datasets/base_layers/static_layers/hydrological_soil_group/hydrological_soil_group.geojson - type: file - - name: mission antyodaya local_path: "{DATA_DIR}/base_layers/mission_antyodaya/mission_antyodaya.gpkg" source: "" type: file + - name: hydrological soil group + aliases: + - soil raster + - hysogs + local_path: "{DATA_DIR}/base_layers/soil/hysogs_india_250m_4326.tif" + source: "" + type: file + - name: ceew climate data local_path: "{DATA_DIR}/base_layers/ceew_climate_data/ceew_climate_data.tif" source: "" @@ -509,8 +512,16 @@ derived_layers: - name: stream_order - name: hydrology_fortnightly + filename: "deltaG_fortnight_{district}_{block}.gpkg" + local_path: "{DATA_DIR}/hydrology/hydrology_local/{state}/{district}/{block}/{filename}" + geoserver_workspace: mws_layers + layer_type: vector - name: hydrology_annual + filename: "deltaG_annual_{district}_{block}.gpkg" + local_path: "{DATA_DIR}/hydrology/hydrology_local/{state}/{district}/{block}/{filename}" + geoserver_workspace: mws_layers + layer_type: vector - name: generate_block_layer diff --git a/computing/config_loader.py b/computing/config_loader.py index 0c35461b..f1d14234 100644 --- a/computing/config_loader.py +++ b/computing/config_loader.py @@ -148,6 +148,11 @@ def _derived_output_dir(name: str) -> Path: TERRAIN_RASTER_PATH: Path = _base_layer_path("terrain") +SOIL_RASTER_PATH: Path = _base_layer_path( + "hydrological soil group", + allowed_suffixes=(".tif", ".tiff"), +) + AEZ_VECTOR_PATH: Path = _base_layer_path("aez") PRECOMPUTED_TEHSIL_WATERSHED_DIR: Path = DATA_DIR / "base_layers/tehsil_watersheds" @@ -219,7 +224,7 @@ def _derived_output_dir(name: str) -> Path: AQUIFER_VECTOR_OUTPUT_DIR: Path = _derived_output_dir("aquifer vector") SWB_VECTOR_OUTPUT_DIR: Path = _derived_output_dir("generate_swb") SOIL_TYPE_OUTPUT_DIR: Path = _derived_output_dir("soil type") - +HYDROLOGY_LOCAL_OUTPUT_DIR: Path = _derived_output_dir("hydrology_fortnightly") PAN_INDIA_DRAINAGE_LINES_GPKG_PATH = ( DATA_DIR / "base_layers/drainage_lines_pan_india.gpkg" diff --git a/computing/hydrology_gpu/__init__.py b/computing/hydrology_gpu/__init__.py new file mode 100644 index 00000000..2f6da047 --- /dev/null +++ b/computing/hydrology_gpu/__init__.py @@ -0,0 +1 @@ +"""Local GPU runoff pipeline package.""" diff --git a/computing/hydrology_gpu/algorithms/__init__.py b/computing/hydrology_gpu/algorithms/__init__.py new file mode 100644 index 00000000..38ae2d5b --- /dev/null +++ b/computing/hydrology_gpu/algorithms/__init__.py @@ -0,0 +1,49 @@ +from typing import Dict +from time import perf_counter +from .. import config as cfg +from .. import utils + +from ..utils import GeoTIFFHandler + + +def format_elapsed(seconds): + if seconds < 60: + return f"{seconds:.2f}s" + minutes, seconds = divmod(seconds, 60) + if minutes < 60: + return f"{int(minutes)}m {seconds:.2f}s" + hours, minutes = divmod(minutes, 60) + return f"{int(hours)}h {int(minutes)}m {seconds:.2f}s" + + +class GenericAlgorithm: + def __init__(self) -> None: + tif_handler = utils.tif_handler + self.tif_handler = tif_handler + self.logger = tif_handler.logger + + def load_inputs(self): + pass + + def main(self): + pass + + def save_outputs(self): + pass + + def run_timed(self, name, fn): + start_time = perf_counter() + self.logger.info("Starting %s", name) + try: + result = fn() + except Exception: + self.logger.exception("Failed %s after %s", name, format_elapsed(perf_counter() - start_time)) + raise + else: + self.logger.info("Finished %s in %s", name, format_elapsed(perf_counter() - start_time)) + return result + + def run(self): + self.run_timed("loading inputs", self.load_inputs) + self.run_timed("main algorithm", self.main) + self.run_timed("saving outputs", self.save_outputs) diff --git a/computing/hydrology_gpu/algorithms/runoff.py b/computing/hydrology_gpu/algorithms/runoff.py new file mode 100644 index 00000000..d304b207 --- /dev/null +++ b/computing/hydrology_gpu/algorithms/runoff.py @@ -0,0 +1,553 @@ +import pathlib +import warnings +import os +from natsort import natsorted +from tqdm import tqdm +from . import GenericAlgorithm, GeoTIFFHandler +import cupy as cp +import numpy as np +from ..downloads import rainfall +from .. import config as cfg +from ..lulc_mapping import ( + lulc_cache_key_for_timestamp, + map_lulc_to_dynamic_world, + nodata_lulc_value_for_source, +) + +ANTECEDENT_DAYS = 5 + +class Runoff(GenericAlgorithm): + """ + This one saves outputs in main function. It is part of a generator. It is consumed by timeseries.py + """ + def load_inputs(self): + self.rainfall_iter = rainfall.Load_from_database() + + def load_sr_inputs(self): + soil = cp.asarray(self.tif_handler.load_with_padding(cfg.SOIL_PATH)) + source = getattr(cfg, "LULC_SOURCE", "dynamicworld") + raw_lulc = cp.asarray( + self.tif_handler.load_with_padding( + cfg.LULC_PATH, + fill_value=nodata_lulc_value_for_source(source), + ) + ) + slope = cp.asarray(self.tif_handler.load_with_padding(cfg.DEMFILE_PATH)) + return soil, raw_lulc, slope + + def sr_cache_key(self, timestamp): + return lulc_cache_key_for_timestamp( + getattr(cfg, "LULC_SOURCE", "dynamicworld"), + timestamp, + ) + + def compute_sr_for_timestamp(self, soil, raw_lulc, slope, timestamp): + source = getattr(cfg, "LULC_SOURCE", "dynamicworld") + sr_key = self.sr_cache_key(timestamp) + self.logger.info("Preparing SR for LULC source=%s key=%s", source, sr_key) + mapped_lulc = map_lulc_to_dynamic_world(raw_lulc, source, timestamp) + try: + return self.compute_sr_and_CNs(soil, mapped_lulc, slope) + finally: + if mapped_lulc is not raw_lulc: + del mapped_lulc + + def main(self): + # pathlib.Path(cfg.RUNOFFS_FOLDER).mkdir(parents=True, exist_ok=True) + + antecedent_images = [] + P5_sum = None + previous_Runoff = None + soil, raw_lulc, slope = self.load_sr_inputs() + current_sr_key = None + sr1 = sr2 = sr3 = None + + for index, file in enumerate(self.rainfall_iter.main()): + # self.logger.info(f"Processing file {index}") + img = self.tif_handler.load_with_padding_inner(file['crs'], file['data'], file['bounds']) + + np.nan_to_num(img, copy=False) + + if len(antecedent_images) == ANTECEDENT_DAYS: + if P5_sum is None: + P5_sum = cp.sum(cp.stack([cp.asarray(i) for i in antecedent_images]), axis=0) + # self.logger.info("4. Initial antecedent sum creation") + + # if previous_Runoff is not None: + # P_sum += previous_Runoff + # P5_sum += previous_Runoff + # self.logger.info("6. Add previous runoff") + + P_sum = cp.asarray(img) + sr_key = self.sr_cache_key(file['timestamp']) + if sr_key != current_sr_key: + if sr1 is not None: + del sr1, sr2, sr3 + sr1, sr2, sr3 = self.compute_sr_for_timestamp( + soil, + raw_lulc, + slope, + file['timestamp'], + ) + current_sr_key = sr_key + self.logger.info("Loaded SR for LULC key %s", sr_key) + + # check_numerical_stability(P_sum, "P_sum") + # check_numerical_stability(P5_sum, "P5_sum") + # assert check_physical_range(P_sum, "P_sum", min_val=0.0) + # assert check_physical_range(P5_sum, "P5_sum", min_val=0.0) + + m1_cp = LegacyCodes.compute_M(sr1, P5_sum) + m2_cp = LegacyCodes.compute_M(sr2, P5_sum) + m3_cp = LegacyCodes.compute_M(sr3, P5_sum) + # self.logger.info("7. Compute M1,M2,M3") + + # cp.cuda.set_allocator(cp.cuda.MemoryPool().malloc) + # with cp.cuda.memory_hooks.DebugPrintHook(): + R = LegacyCodes.calculate_runoff_cupy(P_sum, P5_sum, m1_cp, m2_cp, m3_cp, sr1, sr2, sr3) + del m1_cp, m2_cp, m3_cp + # self.logger.info("8. Calculate runoff") + + # previous_Runoff = LegacyCodes.transfer_flow(destination_index, R) + + # self.test_raster(previous_Runoff, index, file) + + # positive_inf_check = cp.isposinf(cp.asarray(previous_Runoff)) + # negative_inf_check = cp.isneginf(cp.asarray(previous_Runoff)) + # nan_check= cp.isnan(cp.asarray(previous_Runoff)) + # if cp.sum(positive_inf_check) > 0 or cp.sum(negative_inf_check) > 0: + # logger.warning(f"Infinity values found in the transferred runoff data at index {index} and rainfall file {file}. They will be replaced with NaN.") + # elif cp.sum(nan_check) > 0: + # logger.warning(f"NaN values found in the transferred runoff data at index {index} and rainfall file {file}. They will be preserved as NaN.") + + # print("test") + + # self.logger.info("9. Transfer flow") + + # self.tif_handler.save_tiff(cp.asnumpy(R), os.path.join(cfg.RUNOFFS_FOLDER, f'runoff_simulation_{index}.tif')) + # self.logger.info("10. write runoff simulation result") + # runoff_rasters.append(R.get()) + self.tif_handler.save_geozarr_time(R.get(), file['timestamp'], cfg.RUNOFFS_FOLDER, "runoff") + old_img = cp.asarray(antecedent_images.pop(0)) + P5_sum = P5_sum - old_img + P_sum + antecedent_images.append(img) + del old_img, P_sum + yield (img, R, file['timestamp']) + + else: + antecedent_images.append(img) + yield (img, None, file['timestamp']) + + if sr1 is not None: + del sr1, sr2, sr3 + del soil, raw_lulc, slope, antecedent_images, P5_sum + self.logger.info("Done runoff sim") + # return runoff_rasters + + def compute_sr_and_CNs(self, soil, lulc, slope): + """ + Returns sr1, sr2 and sr3 cupy arrays. + """ + + # check_numerical_stability(soil, "soil") + # check_numerical_stability(lulc, "lulc") + # check_numerical_stability(slope, "slope") + # check_physical_range(soil, "soil", min_val=0, max_val=4) + # check_physical_range(lulc, "lulc", min_val=0, max_val=7) + # check_physical_range(slope, "slope", min_val=0.0) + + self.logger.info("Calculating SR ...") + + CN2 = LegacyCodes.compute_cn2(soil, lulc) + CN1 = LegacyCodes.compute_cn1(CN2) + CN3 = LegacyCodes.compute_cn3(CN2) + + # check_numerical_stability(CN2, "CN2") + # check_numerical_stability(CN1, "CN1") + # check_numerical_stability(CN3, "CN3") + + p1 = LegacyCodes.compute_part1(CN3, CN2) + p2 = LegacyCodes.compute_part2(slope) + + CN2a = LegacyCodes.compute_CN2a(p1, p2, CN2) + CN1a = LegacyCodes.compute_CN1a(CN2a) + CN3a = LegacyCodes.compute_CN3a(CN2a) + + # check_numerical_stability(CN2a, "CN2a") + # check_numerical_stability(CN1a, "CN1a") + # check_numerical_stability(CN3a, "CN3a") + # assert check_physical_range(CN2a, "CN2a", min_val=0.0, max_val=100.0) + # assert check_physical_range(CN1a, "CN1a", min_val=0.0, max_val=100.0) + # assert check_physical_range(CN3a, "CN3a", min_val=0.0, max_val=100.0) + + sr1 = LegacyCodes.compute_sr(CN1a) + sr2 = LegacyCodes.compute_sr(CN2a) + sr3 = LegacyCodes.compute_sr(CN3a) + + # check_numerical_stability(sr1, "sr1") + # check_numerical_stability(sr2, "sr2") + # check_numerical_stability(sr3, "sr3") + # assert check_physical_range(sr1, "sr1", min_val=0.0) + # assert check_physical_range(sr2, "sr2", min_val=0.0) + # assert check_physical_range(sr3, "sr3", min_val=0.0) + + self.logger.info("Just completed calculating SR") + + return sr1, sr2, sr3 + +def static_all_methods(cls): + for name, attr in cls.__dict__.items(): + if callable(attr): + setattr(cls, name, staticmethod(attr)) + return cls + +@static_all_methods +class LegacyCodes: + + def compute_cn2(soil: cp.ndarray, lulc: cp.ndarray) -> cp.ndarray: + """ + Compute CN2 values from soil and lulc matrices using CuPy. + + Parameters: + soil (cp.ndarray): Soil type matrix (values 0–4). + lulc (cp.ndarray): LULC class matrix (values 0–7). + + Returns: + cp.ndarray: CN2 values. + """ + # Define lookup table [soil_type][lulc_class] + LUT = cp.array([ + [ 0, 0, 0, 0, 0, 0, 0, 0], # soil 0 → CN2 = 0 (as fallback) + [ 0, 30, 39, 0, 64, 39, 82, 49], # soil 1 + [ 0, 55, 61, 0, 75, 61, 88, 69], # soil 2 + [ 0, 70, 74, 0, 82, 74, 91, 79], # soil 3 + [ 0, 77, 80, 0, 85, 80, 93, 84], # soil 4 + ], dtype=cp.int32) + + # Ensure valid bounds before indexing + soil = cp.clip(soil.astype(cp.int32), 0, 4) + lulc = cp.clip(lulc.astype(cp.int32), 0, 7) + + # Apply lookup + CN2 = LUT[soil, lulc] + + del soil, lulc, LUT + + return CN2 + + @staticmethod # this should happen for other funcs also, idk why isn't happening + def compute_cn1(CN2: cp.ndarray) -> cp.ndarray: + """ + Compute CN1 from CN2 using the formula: CN1 = -75 * CN2 / (CN2 - 175) + + Parameters: + CN2 (cp.ndarray): Curve Number 2 matrix (usually int or float) + + Returns: + cp.ndarray: CN1 values as float32 + """ + CN2 = CN2.astype(cp.float32) + denom = CN2 - 175 + # Prevent division by zero + denom = cp.where(denom == 0, cp.finfo(cp.float32).eps, denom) + CN1 = (-75 * CN2) / denom + + del denom + + return CN1 + + def compute_cn3(CN2: cp.ndarray) -> cp.ndarray: + """ + Compute CN3 from CN2 using the formula: + CN3 = CN2 * (e ** (0.00673 * (100 - CN2))) + + Parameters: + CN2 (cp.ndarray): CuPy array of Curve Number 2 values. + + Returns: + cp.ndarray: CuPy array of CN3 values. + """ + CN2 = CN2.astype(cp.float32) + exponent = 0.00673 * (100.0 - CN2) + CN3 = CN2 * cp.exp(exponent) + + del exponent, CN2 + + return CN3 + + def compute_part1(CN3: cp.ndarray, CN2: cp.ndarray) -> cp.ndarray: + CN3 = CN3.astype(cp.float32) + CN2 = CN2.astype(cp.float32) + p1 = (CN3 - CN2) / 3.0 + + del CN3, CN2 + + return p1 + + def compute_part2(slope: cp.ndarray) -> cp.ndarray: + slope = slope.astype(cp.float32) + p2 = 1.0 - 2.0 * cp.exp(-13.86 * slope) + + del slope + + return p2 + + def compute_CN2a(part1: cp.ndarray, part2: cp.ndarray, CN2: cp.ndarray) -> cp.ndarray: + # Ensure type consistency + part1 = part1.astype(cp.float32) + part2 = part2.astype(cp.float32) + CN2 = CN2.astype(cp.float32) + + CN2a = part1 * part2 + CN2 + + del part1, part2, CN2 + + return CN2a + + def compute_CN1a(CN2a: cp.ndarray) -> cp.ndarray: + CN2a = CN2a.astype(cp.float32) + CN1a = 4.2 * CN2a / (10 - 0.058 * CN2a) + + del CN2a + + return CN1a + + def compute_CN3a(CN2a: cp.ndarray) -> cp.ndarray: + CN2a = CN2a.astype(cp.float32) + CN3a = 23 * CN2a / (10 + 0.13 * CN2a) + + del CN2a + + return CN3a + + def compute_sr(CN: cp.ndarray) -> cp.ndarray: + CN = CN.astype(cp.float32) + + # Mask where CN is invalid (e.g., 0 or very small) + mask = CN <= 10 + CN = cp.where(mask, 100.0, CN) # default value or np.nan + + # Clip CN to range 30–100 + CN = cp.clip(CN, 30.0, 100.0) + + sr = (25400.0 / CN) - 254.0 + + # # Optional: mask output back where CN was invalid + # sr = cp.where(mask, -254.0, sr) + + del CN, mask + + return sr + + + + + def compute_M(sr, p): + """ + Compute M2 using CuPy, ensuring that if either sr or p is NaN, the result is NaN. + """ + nan_mask = cp.isnan(sr) | cp.isnan(p) + sqrt_term = cp.sqrt(cp.maximum(sr**2 + 4 * p * sr, 0.0)) + M2 = 0.5 * (-sr + sqrt_term) + M2[nan_mask] = cp.nan # Preserve NaN values + return M2 + + def compute_M_alt(sr, p): + # Ensure float type, potentially float64 for precision + sr = sr.astype(cp.float64) + p = p.astype(cp.float64) + + # Preserve input NaNs + nan_mask_input = cp.isnan(sr) | cp.isnan(p) + + # Calculate term inside sqrt + term = sr**2 + 4 * p * sr + + # Allow sqrt to produce NaN for negative inputs (and suppress warning) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=RuntimeWarning) # Ignore sqrt domain error warning + sqrt_term = cp.sqrt(term) # This will be NaN where term < 0 + + M_result = 0.5 * (-sr + sqrt_term) + + # Ensure input NaNs propagate, and sqrt NaNs are kept + M_result[nan_mask_input | cp.isnan(sqrt_term)] = cp.nan + + del sr, p, nan_mask_input, term, sqrt_term + + return M_result + + def sum_tif_images(images, start, end): + """ + Sum multiple TIFF images using CuPy on the GPU. + """ + return cp.sum(images[start:end], axis=0) + + def calculate_p_and_p5(file_paths): + """ + Load images, convert to CuPy, and compute total precipitation sums. + """ + images = [cp.asarray(load_tif_image(fp), dtype=cp.float32) for fp in file_paths] + total_sum = cp.sum(images, axis=0) + mid_sum = cp.sum(images[-2:], axis=0) if len(images) >= 2 else total_sum + return total_sum, mid_sum + + def calculate_runoff(P, P5, M1, M2, M3, sr1, sr2, sr3): + """ + Compute runoff using CuPy. + """ + nan_mask = cp.isnan(P) | cp.isnan(P5) | cp.isnan(M1) | cp.isnan(M2) | cp.isnan(M3) | cp.isnan(sr1) | cp.isnan(sr2) | cp.isnan(sr3) + runoff = cp.zeros_like(sr1) + mask1 = (~nan_mask) & (P >= 0.2 * sr1) & (P5 >= 0) & (P5 <= 35) + mask2 = (~nan_mask) & (P >= 0.2 * sr2) & (P5 > 35) & (P5 <= 52.5) + mask3 = (~nan_mask) & (P >= 0.2 * sr3) & (P5 > 52.5) + + runoff[mask1] = ((P[mask1] - 0.2 * sr1[mask1]) * (P[mask1] - 0.2 * sr1[mask1] + M1[mask1])) / (P[mask1] + 0.2 * sr1[mask1] + sr1[mask1] + M1[mask1]) + runoff[mask2] = ((P[mask2] - 0.2 * sr2[mask2]) * (P[mask2] - 0.2 * sr2[mask2] + M2[mask2])) / (P[mask2] + 0.2 * sr2[mask2] + sr2[mask2] + M2[mask2]) + runoff[mask3] = ((P[mask3] - 0.2 * sr3[mask3]) * (P[mask3] - 0.2 * sr3[mask3] + M3[mask3])) / (P[mask3] + 0.2 * sr3[mask3] + sr3[mask3] + M3[mask3]) + + runoff[nan_mask] = cp.nan # Restore NaNs + return runoff + + @staticmethod + def calculate_runoff_cupy(P, P5, m1, m2, m3, sr1, sr2, sr3): + """ + Calculates runoff using a CuPy implementation mirroring a GEE expression. + + Follows the logic: + Q = f(P, P5, sr, m) based on AMC I, II, III where P5 determines AMC. + Uses derived m1, m2, m3 and potential max retention sr1, sr2, sr3. + Ensures runoff >= 0 and handles NaN inputs (NaN in any input -> NaN output). + + Args: + P (cp.ndarray): Precipitation matrix. + P5 (cp.ndarray): 5-day antecedent precipitation matrix. + m1 (cp.ndarray): Derived moisture parameter for AMC I. + m2 (cp.ndarray): Derived moisture parameter for AMC II. + m3 (cp.ndarray): Derived moisture parameter for AMC III. + sr1 (cp.ndarray): Potential maximum retention for AMC I (S derived from CN1). + sr2 (cp.ndarray): Potential maximum retention for AMC II (S derived from CN2). + sr3 (cp.ndarray): Potential maximum retention for AMC III (S derived from CN3). + + Returns: + cp.ndarray: Calculated runoff matrix, with NaN where any input was NaN. + """ + # Optional: Check if inputs are indeed CuPy arrays (if function might receive others) + # P = cp.asarray(P) # etc. for all inputs + + # --- 0. Input Validation (Optional but Recommended) --- + if not P.shape == P5.shape == m1.shape == m2.shape == m3.shape == \ + sr1.shape == sr2.shape == sr3.shape: + raise ValueError("All input CuPy arrays must have the same shape.") + + # --- 1. Handle NaN Inputs: Create combined mask --- + # If any input pixel is NaN, the output for that pixel will be NaN. + + # --- 2. Calculate Intermediate Terms (Initial Abstraction) --- + Ia1 = 0.2 * sr1 + Ia2 = 0.2 * sr2 + Ia3 = 0.2 * sr3 + + # --- 3. Calculate Potential Runoff Values (Q1, Q2, Q3) --- + # Suppress potential division-by-zero or invalid value warnings as we handle them + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=RuntimeWarning) + + # # Calculate Denominators for the runoff formula + den = P + Ia1 + sr1 + m1 + num = (P - Ia1) * (P - Ia1 + m1) + Q1 = cp.where(den != 0, num / den, 0.0).astype(cp.float64) + + den = P + Ia2 + sr2 + m2 + num = (P - Ia2) * (P - Ia2 + m2) + Q2 = cp.where(den != 0, num / den, 0.0).astype(cp.float64) + + den = P + Ia3 + sr3 + m3 + num = (P - Ia3) * (P - Ia3 + m3) + Q3 = cp.where(den != 0, num / den, 0.0).astype(cp.float64) + + del den, num + + # --- 4. Define Conditions using Boolean Masks --- + # Basic precipitation conditions: P must be >= Initial Abstraction (Ia) + condP = (P >= Ia1) + condAMC = (P5 >= 0) & (P5 <= 35) + condQ_pos = (Q1 >= 0) + cond1_full = condP & condAMC & condQ_pos + + condP = (P >= Ia2) + condAMC = (P5 > 35) & (P5 <= 52.5) + condQ_pos = (Q2 >= 0) + cond2_full = condP & condAMC & condQ_pos + + condP = (P >= Ia3) + condAMC = (P5 >= 0) & (P5 > 52.5) # Corresponds to the third check in GEE ternary + condQ_pos = (Q3 >= 0) + cond3_full = condP & condAMC & condQ_pos + + del condP, condAMC, condQ_pos + del Ia1, Ia2, Ia3 + + # Antecedent Moisture Conditions based on documented P5 thresholds. + + # Runoff non-negativity conditions (Q must be >= 0) from GEE expression + + # Combine all conditions for each case + # These directly represent the full condition before the '?' in the GEE expression + + # --- 5. Apply Conditions using Nested cp.where (Mirrors GEE Ternary Logic) --- + # This structure directly implements: cond1 ? Q1 : (cond2 ? Q2 : (cond3 ? Q3 : 0)) + final_runoff = cp.where(cond1_full, Q1, # If Cond1 is true, use Q1 + cp.where(cond2_full, Q2, # Else, if Cond2 is true, use Q2 + cp.where(cond3_full, Q3, # Else, if Cond3 is true, use Q3 + 0.0))) # Else (all conditions false), use 0.0 + + del cond1_full, cond2_full, cond3_full + + # --- 6. Apply NaN Mask --- + # Ensure any pixel that had NaN in any input results in NaN output + # Note: cp.where might already propagate NaNs correctly in many cases, + # but applying the mask explicitly guarantees it. + + nan_mask = cp.isnan(P) | cp.isnan(P5) | cp.isnan(m1) | cp.isnan(m2) | cp.isnan(m3) | \ + cp.isnan(sr1) | cp.isnan(sr2) | cp.isnan(sr3) + + final_runoff[nan_mask] = cp.nan + + return final_runoff + + + def runoff_total_volume(runoff): + nan_mask = cp.isnan(runoff) + runoff[~nan_mask] = runoff[~nan_mask] * 900 + return runoff + + # Define the kernel + transfer_kernel = cp.RawKernel(r''' + extern "C" __global__ + void transfer_flow(const int* F, const float* V, float* V_out, int size) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= size) return; + + int dest = F[i]; // Get destination index + if (dest >= 0 && dest < size) { + atomicAdd(&V_out[dest], V[i]); // Transfer value to destination + } + } + ''', 'transfer_flow') + + # Function to execute the kernel + def transfer_flow(F, V): + F_cp = cp.asarray(F, dtype=cp.int32) + V_cp = cp.asarray(V, dtype=cp.float32) + V_out = cp.zeros_like(V_cp) # Initialize output matrix + + size = F_cp.size + threads_per_block = 256 + blocks_per_grid = (size + threads_per_block - 1) // threads_per_block + + + # Launch the kernel + LegacyCodes.transfer_kernel((blocks_per_grid,), (threads_per_block,), (F_cp, V_cp, V_out, size)) + + return V_out diff --git a/computing/hydrology_gpu/algorithms/tiled_timeseries.py b/computing/hydrology_gpu/algorithms/tiled_timeseries.py new file mode 100644 index 00000000..55c6959f --- /dev/null +++ b/computing/hydrology_gpu/algorithms/tiled_timeseries.py @@ -0,0 +1,258 @@ +import csv +import gc +import json +import shutil +from collections import defaultdict +from datetime import datetime +from pathlib import Path + +import cupy as cp +import numpy as np +from shapely.geometry import shape +from tqdm import tqdm + +from .. import config as cfg +from ..downloads import rainfall +from . import GenericAlgorithm +from .runoff import ANTECEDENT_DAYS, Runoff, LegacyCodes + + +def per_watershed_sum_count(mws, raster): + watershed = mws.ravel().astype(cp.int32) + values = cp.asarray(raster, dtype=cp.float32).ravel() + + mask = watershed > 0 + if int(cp.count_nonzero(mask).get()) == 0: + return [], [], [] + + watershed = watershed[mask] + values = values[mask] + finite = cp.isfinite(values) + if int(cp.count_nonzero(finite).get()) == 0: + return [], [], [] + + watershed = watershed[finite] + values = values[finite] + + sums = cp.bincount(watershed, weights=values) + counts = cp.bincount(watershed) + ids = cp.nonzero(counts)[0] + + return ids.get(), sums[ids].get(), counts[ids].get() + + +class TiledTimeSeries(GenericAlgorithm): + def __init__( + self, + tile_size: int, + series_1="Rainfall", + series_2="Runoff", + *oth_args, + **kwargs, + ) -> None: + super().__init__(*oth_args, **kwargs) + self.tile_size = tile_size + self.series_1 = series_1 + self.series_2 = series_2 + + def load_inputs(self): + with open(cfg.MICROWATERSHEDS_PATH, "r") as f: + self.mws_geojson = json.load(f) + + self.shape_records = [] + for fallback_id, feature in enumerate(self.mws_geojson["features"], start=1): + properties = feature.setdefault("properties", {}) + try: + feature_id = int(properties["id"]) + except (KeyError, TypeError, ValueError): + feature_id = fallback_id + properties["id"] = feature_id + properties.pop("timeseries", None) + geometry = shape(feature["geometry"]) + self.shape_records.append((geometry, feature_id, geometry.bounds)) + + def tile_shapes(self, tile_bounds): + left, bottom, right, top = tile_bounds + return [ + (geometry, feature_id) + for geometry, feature_id, bounds in self.shape_records + if bounds[0] < right and bounds[2] > left and bounds[1] < top and bounds[3] > bottom + ] + + @staticmethod + def _empty_series_entry(): + return [0.0, 0] + + def _series_path(self, watershed_id): + return self.series_dir / f"{watershed_id // 1000:04d}" / f"{watershed_id}.csv" + + def write_tile_series(self, tile_index, data1, data2): + watershed_ids = sorted(set(data1) | set(data2)) + if not watershed_ids: + return + for watershed_id in watershed_ids: + path = self._series_path(watershed_id) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", newline="") as f: + writer = csv.writer(f) + timestamps = set(data1[watershed_id]) | set(data2[watershed_id]) + for timestamp in sorted(timestamps): + rainfall_sum, rainfall_count = data1[watershed_id].get(timestamp, (0.0, 0)) + runoff_sum, runoff_count = data2[watershed_id].get(timestamp, (0.0, 0)) + writer.writerow((timestamp, rainfall_sum, rainfall_count, runoff_sum, runoff_count)) + self.logger.info("Stored tile %s series for %s watershed(s) in %s", tile_index, len(watershed_ids), self.series_dir) + + def add_to_series(self, running_data, watershed_raster, raster, name): + raw_date = name + dt = datetime.strptime(raw_date, "%Y%m%d_%H").isoformat() + ids, sums, counts = per_watershed_sum_count(watershed_raster, raster) + for watershed_id, value_sum, value_count in zip(ids, sums, counts): + entry = running_data[int(watershed_id)][dt] + entry[0] += float(value_sum) + entry[1] += int(value_count) + + def process_tile(self, tile_handler, tile_index, tile_count): + shapes = self.tile_shapes(tile_handler.bounds) + if not shapes: + self.logger.info("Skipping tile %s/%s with no intersecting watershed geometries", tile_index, tile_count) + return + + watershed_raster = tile_handler.rasterize_by_id(shapes) + if not np.any(watershed_raster > 0): + self.logger.info("Skipping tile %s/%s with no watershed pixels", tile_index, tile_count) + return + + self.logger.info( + "Processing tile %s/%s: size=%sx%s bounds=%s", + tile_index, + tile_count, + tile_handler.width, + tile_handler.height, + tile_handler.bounds, + ) + + watershed_cp = cp.asarray(watershed_raster) + tile_series_data1 = defaultdict(lambda: defaultdict(self._empty_series_entry)) + tile_series_data2 = defaultdict(lambda: defaultdict(self._empty_series_entry)) + runoff_algo = Runoff() + runoff_algo.tif_handler = tile_handler + + soil, raw_lulc, slope = runoff_algo.load_sr_inputs() + current_sr_key = None + sr1 = sr2 = sr3 = None + + antecedent_images = [] + p5_sum = None + rainfall_iter = rainfall.LoadTile_from_database(tile_handler) + + for index, file in enumerate(rainfall_iter.main()): + img = file["data"] + np.nan_to_num(img, copy=False) + self.add_to_series(tile_series_data1, watershed_cp, img, file["timestamp"]) + + if len(antecedent_images) == ANTECEDENT_DAYS: + if p5_sum is None: + p5_sum = cp.zeros_like(cp.asarray(antecedent_images[0], dtype=cp.float32)) + for previous_img in antecedent_images: + p5_sum = p5_sum + cp.asarray(previous_img, dtype=cp.float32) + + p_sum = cp.asarray(img, dtype=cp.float32) + sr_key = runoff_algo.sr_cache_key(file["timestamp"]) + if sr_key != current_sr_key: + if sr1 is not None: + del sr1, sr2, sr3 + sr1, sr2, sr3 = runoff_algo.compute_sr_for_timestamp( + soil, + raw_lulc, + slope, + file["timestamp"], + ) + current_sr_key = sr_key + self.logger.info("Loaded SR for LULC key %s", sr_key) + m1_cp = LegacyCodes.compute_M(sr1, p5_sum) + m2_cp = LegacyCodes.compute_M(sr2, p5_sum) + m3_cp = LegacyCodes.compute_M(sr3, p5_sum) + runoff = LegacyCodes.calculate_runoff_cupy(p_sum, p5_sum, m1_cp, m2_cp, m3_cp, sr1, sr2, sr3) + self.add_to_series(tile_series_data2, watershed_cp, runoff, file["timestamp"]) + old_img = cp.asarray(antecedent_images.pop(0), dtype=cp.float32) + p5_sum = p5_sum - old_img + p_sum + antecedent_images.append(img) + del p_sum, m1_cp, m2_cp, m3_cp, runoff, old_img + else: + antecedent_images.append(img) + + self.write_tile_series(tile_index, tile_series_data1, tile_series_data2) + if sr1 is not None: + del sr1, sr2, sr3 + del watershed_cp, soil, raw_lulc, slope, rainfall_iter, antecedent_images, p5_sum, tile_series_data1, tile_series_data2 + cp.get_default_memory_pool().free_all_blocks() + gc.collect() + + def main(self): + output_path = Path(cfg.TIMESERIES_VECTOR) + self.series_dir = output_path.parent / f"{output_path.stem}_tile_series" + shutil.rmtree(self.series_dir, ignore_errors=True) + self.series_dir.mkdir(parents=True, exist_ok=True) + + windows = list(self.tif_handler.iter_windows(self.tile_size)) + self.logger.info("Processing %s tile(s) with tile_size=%s", len(windows), self.tile_size) + for tile_index, window in enumerate(tqdm(windows, desc="Processing spatial tiles"), start=1): + tile_handler = self.tif_handler.for_window(window) + self.process_tile(tile_handler, tile_index, len(windows)) + + def save_outputs(self): + def make_output(watershed_id): + path = self._series_path(watershed_id) + if not path.exists(): + return {} + data = defaultdict(lambda: [0.0, 0, 0.0, 0]) + with path.open(newline="") as f: + for timestamp, rainfall_sum, rainfall_count, runoff_sum, runoff_count in csv.reader(f): + entry = data[timestamp] + entry[0] += float(rainfall_sum) + entry[1] += int(rainfall_count) + entry[2] += float(runoff_sum) + entry[3] += int(runoff_count) + + ret = {} + for timestamp in sorted(data): + values = {} + rainfall_sum, rainfall_count, runoff_sum, runoff_count = data[timestamp] + if rainfall_count: + values[self.series_1] = rainfall_sum / rainfall_count + if runoff_count: + values[self.series_2] = runoff_sum / runoff_count + if values: + ret[timestamp] = values + return ret + + output_path = Path(cfg.TIMESERIES_VECTOR) + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("w") as f: + f.write("{") + first_key = True + for key, value in self.mws_geojson.items(): + if key == "features": + continue + if not first_key: + f.write(",") + json.dump(key, f) + f.write(":") + json.dump(value, f, separators=(",", ":")) + first_key = False + + if not first_key: + f.write(",") + f.write('"features":[') + first_feature = True + for watershed in tqdm(self.mws_geojson["features"], desc="Writing tiled timeseries"): + watershed_id = int(watershed["properties"]["id"]) + watershed["properties"]["timeseries"] = make_output(watershed_id) + if not first_feature: + f.write(",") + json.dump(watershed, f, separators=(",", ":")) + watershed["properties"].pop("timeseries", None) + first_feature = False + f.write("]}") + + self.logger.info("Saved file to %s", cfg.TIMESERIES_VECTOR) diff --git a/computing/hydrology_gpu/algorithms/timeseries.py b/computing/hydrology_gpu/algorithms/timeseries.py new file mode 100644 index 00000000..a9a56ef2 --- /dev/null +++ b/computing/hydrology_gpu/algorithms/timeseries.py @@ -0,0 +1,119 @@ +import json +import os +import pathlib +from datetime import datetime +from pprint import pprint +from shapely.geometry import shape +from tqdm import tqdm +import cProfile +import pstats +from .runoff import Runoff +from . import GenericAlgorithm, GeoTIFFHandler +from typing import Type, Dict, Tuple, DefaultDict +from collections import defaultdict, OrderedDict +import cupy as cp +from .. import config as cfg + +# Not generic enough yet... +class TimeSeries(GenericAlgorithm): + def __init__(self, + algo: Type[GenericAlgorithm] = Runoff, + series_1="Rainfall", + series_2="Runoff", + *oth_args, **kwargs + ) -> None: + super().__init__(*oth_args, **kwargs) + self.algo = algo() + self.series_1 = series_1 + self.series_2 = series_2 + + def load_inputs(self): + self.algo.load_inputs() + # self.mws = self.tif_handler.rasterize_by_id(cfg.MICROWATERSHEDS_PATH) + # self.mws_geojson = mws.Clip(self.args, self.logger, cfg).main() + with open(cfg.MICROWATERSHEDS_PATH, 'r') as f: + self.mws_geojson = json.load(f) + + shapes = [] + for fallback_id, feature in enumerate(self.mws_geojson['features'], start=1): + properties = feature.setdefault('properties', {}) + try: + feature_id = int(properties['id']) + except (KeyError, TypeError, ValueError): + feature_id = fallback_id + properties['id'] = feature_id + shapes.append((shape(feature['geometry']), feature_id)) + self.mws = self.tif_handler.rasterize_by_id(shapes) + + def main(self): + mws_cp = cp.asarray(self.mws) + mws_series_data1 = defaultdict(list) + mws_series_data2 = defaultdict(list) + + def add_to_series(running_data, raster, name): + avg, ids = per_watershed_avg(mws_cp, raster) + raw_date = name + dt = datetime.strptime(raw_date, "%Y%m%d_%H") + for val, mws_id in zip(avg, ids): + running_data[int(mws_id)].append((val, dt.isoformat())) + + # profiler = cProfile.Profile() + # profiler.enable() + for s1, s2, name in self.algo.main(): + add_to_series(mws_series_data1, s1, name) + if s2 is not None: + add_to_series(mws_series_data2, s2, name) + + # profiler.disable() + # + # stats = pstats.Stats(profiler) + # stats.dump_stats("loop_profile.prof") + # + # self.logger.info("Profiling data saved to loop_profile.prof") + + self.mws_series_data1 = mws_series_data1 + self.mws_series_data2 = mws_series_data2 + + def save_outputs(self): + def make_output(id): + ret = dict() # apparently, dict is an OrderdDict + for (val, name) in self.mws_series_data1[id]: + ret[name] = {self.series_1: val} + for (val, name) in self.mws_series_data2[id]: + if name not in ret: + ret[name] = {self.series_2: val} + else: + ret[name][self.series_2] = val + return ret + + mws_data = self.mws_geojson + + for mws in tqdm(mws_data['features']): + id = int(mws['properties']['id']) + mws['properties']["timeseries"] = make_output(id) + + with open(cfg.TIMESERIES_VECTOR, 'w+') as f: + json.dump(mws_data, f) + + self.logger.info(f"Saved file to {cfg.TIMESERIES_VECTOR}") + + + +def per_watershed_avg(mws, raster): + runoff_raster = cp.asarray(raster) + watershed_raster = mws + + ws_flat = watershed_raster.ravel() + rf_flat = runoff_raster.ravel() + + unique_ws, inv = cp.unique(ws_flat, return_inverse=True) + + per_ws_sum = cp.bincount(inv, weights=rf_flat, minlength=unique_ws.size) + per_ws_count = cp.bincount(inv, minlength=unique_ws.size) + + per_ws_mean = per_ws_sum / per_ws_count + + # result = per_ws_mean[inv].reshape(watershed_raster.shape) + + return (per_ws_mean.get(), unique_ws.get()) + # return result diff --git a/computing/hydrology_gpu/config/__init__.py b/computing/hydrology_gpu/config/__init__.py new file mode 100644 index 00000000..4fa6a67d --- /dev/null +++ b/computing/hydrology_gpu/config/__init__.py @@ -0,0 +1,49 @@ +import os + +try: + from django.conf import settings +except ModuleNotFoundError: + settings = None + +from computing.config_loader import ( + LULC_BASE_DIR, + PROJECT_ROOT, + SOIL_RASTER_PATH, + TERRAIN_RASTER_PATH, +) + + +def _env_or_setting(name, default=""): + value = None + if settings is not None: + try: + value = getattr(settings, name, None) + except Exception: + value = None + if value in (None, ""): + value = os.environ.get(name, default) + return str(value or "").strip() + + +DATA_ROOT = PROJECT_ROOT / "data" + +# Optional. Do not commit a real project id; use env/settings when needed. +GEE_PROJECT_NAME = _env_or_setting("GEE_PROJECT_NAME") +GEE_SCALE = 30 + +# API/Celery wrappers set these per request before running hydrology. +BOUNDARY_GEOJSON_PATH = "" +MICROWATERSHEDS_PATH = BOUNDARY_GEOJSON_PATH +DEMFILE_PATH = "" +SOIL_PATH = str(SOIL_RASTER_PATH) +LULC_PATH = str(LULC_BASE_DIR / "lulc_v3_2024_2025.tif") +LULC_SOURCE = "indiasatv3" +INDIASATV3_LULC_PATH = LULC_PATH + +RAINFALL_FOLDER = "" +RUNOFFS_FOLDER = "" +TIMESERIES_VECTOR = "" + +ARG_START_DATE = "2017-07-01" +ARG_END_DATE = "2025-06-18" +TILE_SIZE = None diff --git a/computing/hydrology_gpu/downloads/__init__.py b/computing/hydrology_gpu/downloads/__init__.py new file mode 100644 index 00000000..0336b2b7 --- /dev/null +++ b/computing/hydrology_gpu/downloads/__init__.py @@ -0,0 +1,166 @@ +import os +import shutil +from pathlib import Path +from logging import Logger +from dataclasses import dataclass +import ee +import geedim +import geopandas as gpd +import requests + +from .. import config as cfg +from pydrive2.auth import GoogleAuth +from pydrive2.drive import GoogleDrive + +from .. import utils +from ..utils import GeoTIFFHandler + + +def _initialize_earth_engine(): + project = getattr(cfg, "GEE_PROJECT_NAME", "") + try: + if project: + ee.Initialize(project=project) + else: + ee.Initialize() + return + except Exception: + pass + + try: + from utilities.gee_utils import ee_initialize_safe + + ee_initialize_safe() + except Exception as exc: + print(f"Skipping Earth Engine initialization: {exc}") + + +_initialize_earth_engine() + +class GenericDownloader: + # Singleton pattern + # _instance = None + # def __new__(cls, *args, **kwargs): + # if cls._instance is None: + # cls._instance = super().__new__(cls, *args, **kwargs) + # return cls._instance + + @dataclass + class InitializationData: + gauth: GoogleAuth = None + drive: GoogleDrive = None + _init_structs = None + + def __init__(self): + self.logger = utils.tif_handler.logger + self.tif_loader = utils.tif_handler + + if GenericDownloader._init_structs is None: + GenericDownloader._init_structs = GenericDownloader.InitializationData() + + self.gauth = GenericDownloader._init_structs.gauth + self.drive = GenericDownloader._init_structs.drive + + def download_gdrive_file(self, file_id, path): + if GenericDownloader._init_structs.gauth is None: + settings_path = Path(__file__).resolve().parents[1] / "pydrive_settings.yaml" + GenericDownloader._init_structs.gauth = GoogleAuth(settings_file=str(settings_path)) + GenericDownloader._init_structs.drive = GoogleDrive(GenericDownloader._init_structs.gauth) + self.gauth = GenericDownloader._init_structs.gauth + self.drive = GenericDownloader._init_structs.drive + + file_obj = self.drive.CreateFile({'id': file_id}) + # Fetching title first to name the local file + file_obj.FetchMetadata() + self.logger.info(f"Downloading {file_obj['title']}...") + file_obj.GetContentFile(path + file_obj['title']) + return f"Finished {file_obj['title']}" + + + @staticmethod + def empty_folder(folder): + shutil.rmtree(folder) + os.mkdir(folder) + + def load_region(self): + # Load the file - GeoPandas handles FeatureCollection vs Feature automatically + gdf = gpd.read_file(cfg.BOUNDARY_GEOJSON_PATH) + + # Generalize to a single geometry (Unions everything if there are multiple features) + # Helpful if the input itself is a vector of different mws. + combined_geom = gdf.unary_union + + # Convert to Earth Engine Geometry + # __geo_interface__ is a standard way to get GeoJSON-like dicts from objects + geojson_struct = combined_geom.__geo_interface__ + region = ee.Geometry(geojson_struct) + return region + + @staticmethod + def save_from_gee(collection, region, tif_file_path): + logger = utils.tif_handler.logger if utils.tif_handler else None + os.makedirs(os.path.dirname(tif_file_path) or ".", exist_ok=True) + + try: + # 1. Attempt the fast direct download + if logger: + logger.info(f"Requesting Earth Engine download URL for {tif_file_path}") + + url = collection.getDownloadURL({ + 'format': 'GEO_TIFF', + 'scale': cfg.GEE_SCALE, + 'region': region + }) + + if logger: + logger.info(f"Downloading {tif_file_path}") + + response = requests.get(url, timeout=(30, 900)) + + # If Earth Engine says "Too Large", the status_code will not be 200 + if response.status_code == 200: + with open(tif_file_path, 'wb') as f: + f.write(response.content) + + if logger: + logger.info(f"Downloaded {tif_file_path} with `getDownloadURL`") + else: + print(f"Downloaded {tif_file_path} with `getDownloadURL`") + else: + raise ValueError(f"Image too large for direct URL (Status {response.status_code})") + + except Exception as e: + if logger: + logger.warning( + "Direct Earth Engine download failed for %s: %s. Falling back to tiled download.", + tif_file_path, + e, + ) + else: + print(f"Direct Earth Engine download failed for {tif_file_path}: {e}") + print("Falling back to tiled download.") + + # getDownloadURL uses a single Earth Engine thumbnail request, which + # fails for district-scale rasters over the 50 MB limit. geedim + # splits the same image into smaller computePixels requests and + # stitches them into one local GeoTIFF. + prepared_image = collection.gd.prepareForExport( + scale=cfg.GEE_SCALE, + region=region, + resampling="near", + ) + prepared_image.gd.toGeoTIFF( + tif_file_path, + overwrite=True, + max_tile_size=4, + max_requests=2, + max_cpus=1, + ) + + if logger: + logger.info(f"Downloaded {tif_file_path} with tiled Earth Engine download") + else: + print(f"Downloaded {tif_file_path} with tiled Earth Engine download") + + def main(self): + pass diff --git a/computing/hydrology_gpu/downloads/dem.py b/computing/hydrology_gpu/downloads/dem.py new file mode 100644 index 00000000..69373056 --- /dev/null +++ b/computing/hydrology_gpu/downloads/dem.py @@ -0,0 +1,95 @@ +""" +Currently, this calculates Slope (gradient) on GEE Servers. +DEM is not downloaded. +""" +from pathlib import Path + +import geopandas as gpd +import rasterio +from rasterio.mask import mask +import xarray +from .. import config as cfg +from . import GenericDownloader, ee +import geemap + + +def clip_local_raster(source_path, boundary_path, output_path=None, logger=None, fill_value=0): + source_path = Path(source_path) + boundary_path = Path(boundary_path) + output_path = Path(output_path or cfg.DEMFILE_PATH) + + if logger: + logger.info("Clipping local terrain raster %s to %s", source_path, output_path) + logger.warning( + "Reading local terrain raster with GTIFF_IGNORE_READ_ERRORS=YES; unreadable source tiles may be filled by GDAL" + ) + + gdf = gpd.read_file(boundary_path) + if gdf.empty: + raise ValueError(f"No geometries found in boundary file: {boundary_path}") + + output_path.parent.mkdir(parents=True, exist_ok=True) + + with rasterio.Env(GTIFF_IGNORE_READ_ERRORS="YES"): + with rasterio.open(source_path) as src: + if gdf.crs is None: + gdf = gdf.set_crs(src.crs) + elif gdf.crs != src.crs: + gdf = gdf.to_crs(src.crs) + + data, transform = mask( + src, + gdf.geometry, + crop=True, + filled=True, + nodata=fill_value, + ) + + profile = src.profile.copy() + profile.update( + driver="GTiff", + height=data.shape[1], + width=data.shape[2], + transform=transform, + nodata=fill_value, + tiled=True, + compress="ZSTD", + ZSTD_LEVEL=1, + NUM_THREADS=10, + ) + + with rasterio.open(output_path, "w", **profile) as dst: + dst.write(data) + + if logger: + logger.info("Saved clipped local terrain raster to %s", output_path) + + return output_path + + +class Downloader(GenericDownloader): + def __init__(self): + """ + Override parent class's init. As GeoTiff handler needs one tif file for CRS reference + We download DEM as this reference + """ + pass + + def main(self): + dataset = ee.Image('USGS/SRTMGL1_003') + elevation = dataset.select('elevation') + + region = self.load_region() + + # 1. Calculate slope in degrees + slope_deg = ee.Terrain.slope(elevation) + + # 2. Convert to Gradient: tan(slope * pi / 180) + slope_gradient = slope_deg.multiply(3.141592).divide(180).tan() + + elevation_clip = slope_gradient.clipToBoundsAndScale( + geometry=region, + scale=cfg.GEE_SCALE + ) + + self.save_from_gee(elevation_clip, region, cfg.DEMFILE_PATH) diff --git a/computing/hydrology_gpu/downloads/lulc.py b/computing/hydrology_gpu/downloads/lulc.py new file mode 100644 index 00000000..54ffce88 --- /dev/null +++ b/computing/hydrology_gpu/downloads/lulc.py @@ -0,0 +1,31 @@ +import geemap +import xarray +from .. import config as cfg +from . import GenericDownloader, ee +import json +import requests +from .rainfall import DownloaderBase as RainfallDownloader + +class Downloader(GenericDownloader): + """ + Currently, the LULC is mode of lulc's from start date to end date. Static + """ + def main(self): + region = self.load_region() + + end_date = ee.Date(cfg.ARG_END_DATE) + start_date = ee.Date(cfg.ARG_START_DATE) + + dw_col = (ee.ImageCollection('GOOGLE/DYNAMICWORLD/V1') + .filterDate(start_date, end_date) + .filterBounds(region) + .select('label')) + + dw_image = dw_col.reduce(ee.Reducer.mode()).rename('lulc') + + dw_clip = dw_image.clipToBoundsAndScale( + geometry=region, + scale=cfg.GEE_SCALE + ) + + self.save_from_gee(dw_clip, region, cfg.LULC_PATH) diff --git a/computing/hydrology_gpu/downloads/rainfall.py b/computing/hydrology_gpu/downloads/rainfall.py new file mode 100644 index 00000000..622dd9f7 --- /dev/null +++ b/computing/hydrology_gpu/downloads/rainfall.py @@ -0,0 +1,415 @@ +import gc +import json +import os +import pathlib +import shutil +import threading +import time +from collections import deque +from concurrent.futures import ThreadPoolExecutor +from threading import Lock +from .. import config as cfg +import xarray as xr +import requests +import concurrent.futures + +import geopandas as gpd +import zarr +from rasterio.enums import Resampling +from shapely.geometry import box +from rasterio.transform import Affine, from_origin +from rasterio.warp import reproject +from tqdm import tqdm +# import geedim as gd +import cupy as cp +import numpy as np +import pandas as pd +import cucim.skimage.transform as cimg + +from . import GenericDownloader, ee, Logger + +class DownloaderBase(GenericDownloader): + def __init__(self): + + super().__init__() + self.zarr_path = os.path.join(cfg.RAINFALL_FOLDER, "rainfall_archive.zarr") + pathlib.Path(self.zarr_path).parent.mkdir(parents=True, exist_ok=True) + + def load_region_gdf(self): + gdf = gpd.read_file(cfg.BOUNDARY_GEOJSON_PATH) + if gdf.empty: + raise ValueError(f"No geometries found in {cfg.BOUNDARY_GEOJSON_PATH}") + + if gdf.crs is None: + gdf = gdf.set_crs("EPSG:4326") + elif gdf.crs.to_epsg() != 4326: + gdf = gdf.to_crs("EPSG:4326") + return gdf + + def load_local_region_geometry(self): + gdf = self.load_region_gdf() + if hasattr(gdf.geometry, "union_all"): + return gdf.geometry.union_all() + return gdf.unary_union + + def load_buffered_bounds_geometry(self, buffer_m=12000): + """ + Build a small EE rectangle from local bounds instead of sending the full + boundary polygon to Earth Engine. Large pan-India polygons can exceed + EE's 10 MB request payload limit. + """ + gdf = self.load_region_gdf() + metric_gdf = gdf.to_crs("EPSG:3857") + minx, miny, maxx, maxy = metric_gdf.geometry.buffer(buffer_m).total_bounds + bounds_geom = gpd.GeoSeries([box(minx, miny, maxx, maxy)], crs="EPSG:3857").to_crs("EPSG:4326") + west, south, east, north = [float(value) for value in bounds_geom.total_bounds] + self.logger.info("Using local buffered boundary bounds: [%s, %s, %s, %s]", west, south, east, north) + return ee.Geometry.Rectangle([west, south, east, north], proj="EPSG:4326", geodesic=False) + + def init_zarr(self, dates, dummy_da): + """ + Initializes a Zarr store with the full time extent to allow parallel region writes. + """ + native_y, native_x = dummy_da.y.size, dummy_da.x.size + + # Create the skeleton Dataset + # We use empty/zeros but with compute=False, so no data is actually written yet + ds_skeleton = xr.Dataset( + {"precipitation": (["time", "y", "x"], + np.zeros((len(dates), native_y, native_x), dtype='float32'))}, + coords={ + "time": dates, + "y": dummy_da.y.values, + "x": dummy_da.x.values + } + ) + + # Metadata/CRS (Important for GeoZarr) + ds_skeleton.rio.write_crs(dummy_da.rio.crs, inplace=True) + ds_skeleton.rio.write_transform(dummy_da.rio.transform(), inplace=True) + + # Encoding: chunking by 1 day is standard for daily time-series + encoding = { + "precipitation": {"chunks": (1, native_y, native_x)}, + "time": { + "units": "hours since 2017-07-01 00:00:00", + "calendar": "proleptic_gregorian", + "dtype": "int64" # Ensuring integer storage for hours + } + } + + # Write ONLY metadata + ds_skeleton.to_zarr(self.zarr_path, mode='w', encoding=encoding, compute=False, zarr_format=2) + self.logger.info(f"Initialized Zarr skeleton at {self.zarr_path} with {len(dates)} slots.") + + + def save_geozarr(self, flat_data, timestamp, dummy_da, t): + """ + Saves data as a 3D (Time, Y, X) chunked array with full spatial coordinates. + """ + + # 1. Reshape flat data to native grid dimensions + # Using the shape from self.dummy_da (e.g., [Lat, Lon]) + native_y, native_x = dummy_da.y.size, dummy_da.x.size + raster_data = flat_data.reshape(native_y, native_x) + + # 2. Create the DataArray with proper spatial coords + da = xr.DataArray( + raster_data[np.newaxis, ...], # Shape: (1, Y, X) + dims=("time", "y", "x"), + coords={ + "time": [pd.to_datetime(timestamp, format='%Y%m%d_%H')], + "y": dummy_da.y.values, + "x": dummy_da.x.values + }, + name="precipitation" + ) + + # 3. Add CRS and metadata + da.rio.write_crs(dummy_da.rio.crs, inplace=True) + da.rio.write_transform(dummy_da.rio.transform(), inplace=True) + + ds = da.to_dataset().drop_vars(["y", "x", "spatial_ref"]) + + ds.to_zarr(self.zarr_path, region={"time": slice(t, t + 1)}) + + def load_geozarr(self): + """ + Loads the Zarr archive as a standard Xarray Dataset. + """ + + if not os.path.exists(self.zarr_path): + self.logger.error(f"Zarr not found at {self.zarr_path}") + return None + + # chunks={} opens it lazily using Dask + ds = xr.open_zarr(self.zarr_path, consolidated=True, chunks={}) + return ds + +class Download_to_database(DownloaderBase): + def main(self): + self.ingest_rainfall_to_zarr() + + def ingest_rainfall_to_zarr(self): + """ + Part 1: Purely fetches data, sums it on GPU, and saves to GeoZarr. + """ + self.logger.info("Starting rainfall ingestion to GeoZarr") + self.logger.info("Loading boundary bounds") + buffered_region = self.load_buffered_bounds_geometry() + + self.logger.info( + "Preparing GSMaP rainfall collection for [%s, %s)", + cfg.ARG_START_DATE, + cfg.ARG_END_DATE, + ) + rainfall_collection = ( + ee.ImageCollection('JAXA/GPM_L3/GSMaP/v6/operational') + .filterDate(cfg.ARG_START_DATE, cfg.ARG_END_DATE) + .select('hourlyPrecipRate') + ) + + # rainfall_collection = ( + # ee.ImageCollection("NASA/GPM_L3/IMERG_DAILY_V06") + # .filterDate(cfg.ARG_START_DATE, cfg.ARG_END_DATE) + # .select('total_accum') + # ) + # + # # 1. Define the time range + # start_date = ee.Date(cfg.ARG_START_DATE) + # end_date = ee.Date(cfg.ARG_END_DATE) + # + # # 2. Calculate the number of days between start and end + # n_days = end_date.difference(start_date, 'days') + # + # def sum_daily(day_offset): + # # Calculate the start and end of each 24-hour window + # start = start_date.advance(ee.Number(day_offset), 'days') + # end = start.advance(1, 'days') + # + # # Filter the collection for this specific day and sum + # daily_sum = (rainfall_collection + # .filterDate(start, end) + # .sum()) # Sums the 'hourlyPrecipRate' + # + # # Return the image with its date metadata (important for further filtering) + # return daily_sum.set({ + # 'system:time_start': start.millis(), + # 'date_string': start.format('YYYY-MM-DD') + # }) + # + # # 3. Create a sequence of days and map the function + # daily_collection = ee.ImageCollection( + # ee.List.sequence(0, n_days.subtract(1)).map(sum_daily) + # ) + + first_img = rainfall_collection.first() + self.logger.info("Fetching GSMaP native projection") + native_proj = first_img.projection() + + self.logger.info("Opening Earth Engine dataset through xarray/xee; this can take a few minutes for pan-India yearly ranges") + ds = xr.open_dataset( + rainfall_collection, + engine='ee', + projection=native_proj, + geometry=buffered_region, + fast_time_slicing=True, + ) + self.logger.info("Earth Engine dataset opened; preparing rainfall coordinates") + + da = ds['hourlyPrecipRate'].rename({'lat': 'y', 'lon': 'x'}).transpose("time", "y", "x") + total_pixels = da.y.size * da.x.size + dummy_da = da.isel(time=0) + + N = len(da.time) + K = 24 # Hours per day + self.logger.info( + "Rainfall grid has %s hourly slices and %sx%s native pixels", + N, + da.y.size, + da.x.size, + ) + + dates = pd.date_range(start=cfg.ARG_START_DATE, end=cfg.ARG_END_DATE, freq='D', inclusive='left') + self.init_zarr(dates, dummy_da) + + ASK_BUFF = 50 + WORKERS = 20 + + def process_slice(t, ticket_num): + da_slice = da.isel(time=slice(t, min(t + K, N))) + + # Extract raw data and sum on GPU + raw_data = da_slice.values + gpu_sum = cp.zeros(total_pixels, dtype=cp.float32) + + for i in range(raw_data.shape[0]): + gpu_sum += cp.asarray(raw_data[i]).ravel() + + timestamp = da_slice.time[0].dt.strftime('%Y%m%d_%H').item() + + # Save the flat sum to the database + # Assuming self.save_geozarr handles time-indexing inside the Zarr + self.save_geozarr(gpu_sum.get(), timestamp, dummy_da, ticket_num) + + with ( + concurrent.futures.ThreadPoolExecutor(max_workers=WORKERS) as executor, + tqdm(range(N), desc="Downloading Rainfall") as pbar + ): + futures = [] + + while pbar.n < pbar.total: + if len(futures) == 0: + gc.collect() + t = pbar.n + assert t % K == 0 + for i in range(ASK_BUFF): + current_hour = t+i*K + day_index = current_hour // K + if current_hour >= N: + break + futures.append(executor.submit( + process_slice, + current_hour, + day_index + )) + futures.pop().result() + pbar.update(K) + + zarr.consolidate_metadata(self.zarr_path) + self.logger.info("Ingestion complete.") + + +def transform_from_center_coords(x_values, y_values): + if len(x_values) < 2 or len(y_values) < 2: + raise ValueError("Rainfall archive must have at least two x and y coordinates") + + dx = float(np.median(np.diff(x_values))) + dy = float(np.median(np.diff(y_values))) + return Affine.translation(float(x_values[0]) - dx / 2, float(y_values[0]) - dy / 2) * Affine.scale(dx, dy) + + +def build_reference_index_map(source_shape, original_size, src_transform, src_crs, tif_handler): + source_indices = np.arange(original_size, dtype=np.int32).reshape(source_shape) + mapped_indices = np.full( + (tif_handler.height, tif_handler.width), + original_size, + dtype=np.int32, + ) + + reproject( + source=source_indices, + destination=mapped_indices, + src_transform=src_transform, + src_crs=src_crs, + dst_transform=tif_handler.transform, + dst_crs=tif_handler.crs, + dst_nodata=original_size, + resampling=Resampling.nearest, + ) + return mapped_indices + + +class Load_from_database(DownloaderBase): + + def __init__(self): + super().__init__() + self.tif_handler = self.tif_loader + self.ds = self.load_geozarr() + if self.ds is None: + raise FileNotFoundError(f"Rainfall archive not found at {self.zarr_path}") + + self.src_transform = transform_from_center_coords( + self.ds.x.values, + self.ds.y.values, + ) + self.src_crs = self.ds.rio.crs or "EPSG:4326" + self.source_shape = (self.ds.sizes["y"], self.ds.sizes["x"]) + self.original_size = self.source_shape[0] * self.source_shape[1] + self.GPU_LUT = cp.asarray(build_reference_index_map( + self.source_shape, + self.original_size, + self.src_transform, + self.src_crs, + self.tif_handler, + )) + + self.STATIC_METADATA = { + "bounds": self.tif_handler.bounds, + "transform": self.tif_handler.transform, + "crs": self.tif_handler.crs, + "shape": (self.tif_handler.height, self.tif_handler.width), + "original_size": self.original_size, + } + + def main(self): + self.logger.info("Starting rainfall load from GeoZarr onto reference grid") + + yield from self.stream_reprojected_rainfall() + + def stream_reprojected_rainfall(self): + for i in tqdm(range(len(self.ds.time))): + # Extract 2D slice and flatten for the GPU LUT + hourly_slice = self.ds.precipitation.isel(time=i) + flat_cpu = hourly_slice.values.ravel() + + # Move to GPU + gpu_src = cp.asarray(flat_cpu) + + # Append sink pixel for nodata and apply LUT + projected_buffer = cp.concatenate([gpu_src, cp.array([0], dtype=gpu_src.dtype)]) + gpu_final = projected_buffer[self.GPU_LUT] + + yield { + "timestamp": hourly_slice.time.dt.strftime('%Y%m%d_%H').item(), + "data": gpu_final.get(), + "bounds": self.STATIC_METADATA["bounds"], + "transform": self.STATIC_METADATA["transform"], + "crs": self.STATIC_METADATA["crs"] + } + + +class LoadTile_from_database(DownloaderBase): + def __init__(self, tif_handler): + super().__init__() + self.tif_handler = tif_handler + self.ds = self.load_geozarr() + if self.ds is None: + raise FileNotFoundError(f"Rainfall archive not found at {self.zarr_path}") + + self.src_transform = transform_from_center_coords( + self.ds.x.values, + self.ds.y.values, + ) + self.src_crs = self.ds.rio.crs or "EPSG:4326" + self.source_shape = (self.ds.sizes["y"], self.ds.sizes["x"]) + self.original_size = self.source_shape[0] * self.source_shape[1] + self.GPU_LUT = cp.asarray(build_reference_index_map( + self.source_shape, + self.original_size, + self.src_transform, + self.src_crs, + self.tif_handler, + )) + + def main(self): + self.logger.info("Starting tiled rainfall load from GeoZarr") + yield from self.stream_reprojected_rainfall() + + def stream_reprojected_rainfall(self): + for i in tqdm(range(len(self.ds.time)), desc="Projecting rainfall tile"): + hourly_slice = self.ds.precipitation.isel(time=i) + flat_cpu = hourly_slice.values.ravel() + gpu_src = cp.asarray(flat_cpu) + projected_buffer = cp.concatenate([gpu_src, cp.array([0], dtype=gpu_src.dtype)]) + gpu_final = projected_buffer[self.GPU_LUT] + data = gpu_final.get() + del gpu_src, projected_buffer, gpu_final + + yield { + "timestamp": hourly_slice.time.dt.strftime('%Y%m%d_%H').item(), + "data": data, + "bounds": self.tif_handler.bounds, + "transform": self.tif_handler.transform, + "crs": self.tif_handler.crs, + } diff --git a/computing/hydrology_gpu/downloads/soil.py b/computing/hydrology_gpu/downloads/soil.py new file mode 100644 index 00000000..57b448b5 --- /dev/null +++ b/computing/hydrology_gpu/downloads/soil.py @@ -0,0 +1,17 @@ +from .. import config as cfg +import xarray +from . import GenericDownloader, ee +import geemap + +class Downloader(GenericDownloader): + def main(self): + region = self.load_region() + + hsg_image = ee.Image('projects/ee-dharmisha-siddharth/assets/HYSOGs250m') + + hsg_clip = hsg_image.clipToBoundsAndScale( + geometry=region, + scale=cfg.GEE_SCALE + ) + + self.save_from_gee(hsg_clip, region, cfg.SOIL_PATH) diff --git a/computing/hydrology_gpu/et_download.py b/computing/hydrology_gpu/et_download.py new file mode 100644 index 00000000..deb5045c --- /dev/null +++ b/computing/hydrology_gpu/et_download.py @@ -0,0 +1,700 @@ +import datetime as dt +import json +import logging +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Iterable +from urllib.parse import urlencode + +import requests +from django.conf import settings +import numpy as np +import rasterio +from rasterio.fill import fillnodata + + +GESDISC_OTF_URL = "https://hydro1.gesdisc.eosdis.nasa.gov/daac-bin/OTF/HTTP_services.cgi" +EVAP_VARIABLE = "Evap_tavg" +FLDAS_FORMAT = "Y29nLw" + +FLDAS_CA_DAILY_SOURCE = "fldas_ca_daily" +FLDAS_CA_DAILY_PATCH_FILLED_SOURCE = "fldas_ca_daily_patch_filled" +FLDAS_CA_DAILY_SHORTNAME = "FLDAS_NOAHMP001_G_CA_D" +FLDAS_CA_DAILY_PRODUCT = "FLDAS_NOAHMP001_G_CA_D.001" +FLDAS_CA_DAILY_BBOX = "21,65.566,37.932,99.844" + +FLDAS_GLOBAL_MONTHLY_SOURCE = "fldas_global_monthly" +FLDAS_GLOBAL_MONTHLY_PATCH_FILLED_SOURCE = "fldas_global_monthly_patch_filled" +FLDAS_GLOBAL_MONTHLY_SHORTNAME = "FLDAS_NOAH01_C_GL_M" +FLDAS_GLOBAL_MONTHLY_PRODUCT = "FLDAS_NOAH01_C_GL_M.001" +PAN_INDIA_BBOX = "6,68,38,98" +DEFAULT_MAX_WORKERS = 3 + + +@dataclass(frozen=True) +class SourceManifest: + source: str + shortname: str + product: str + variable: str + bbox: str + temporal_resolution: str + output_folder: str + file_count: int + downloaded_count: int + skipped_count: int + failed_count: int + patch_filled_source: str | None = None + patch_filled_folder: str | None = None + patch_filled_count: int = 0 + patch_fill_skipped_count: int = 0 + patch_fill_failed_count: int = 0 + patch_fill_invalid_pixels: int = 0 + + +def _coerce_date(value, field_name): + if isinstance(value, dt.date): + return value + if value is None or str(value).strip().lower() in {"", "none", "null"}: + raise ValueError(f"{field_name} is required") + try: + return dt.datetime.strptime(str(value), "%Y-%m-%d").date() + except ValueError as exc: + raise ValueError(f"{field_name} must be in YYYY-MM-DD format") from exc + + +def _iter_days(start_date: dt.date, end_date: dt.date) -> Iterable[dt.date]: + current = start_date + while current < end_date: + yield current + current += dt.timedelta(days=1) + + +def _iter_month_starts(start_date: dt.date, end_date: dt.date) -> Iterable[dt.date]: + current = start_date.replace(day=1) + while current < end_date: + yield current + if current.month == 12: + current = current.replace(year=current.year + 1, month=1) + else: + current = current.replace(month=current.month + 1) + + +def _download_url(params): + return f"{GESDISC_OTF_URL}?{urlencode(params)}" + + +def _daily_ca_params(day: dt.date, bbox: str): + stamp = day.strftime("%Y%m%d") + return { + "FILENAME": ( + f"/data/FLDAS/{FLDAS_CA_DAILY_PRODUCT}/" + f"{day.year}/{day.month:02d}/" + f"{FLDAS_CA_DAILY_SHORTNAME}.A{stamp}.001.nc" + ), + "SERVICE": "L34RS_LDAS", + "BBOX": bbox, + "FORMAT": FLDAS_FORMAT, + "VERSION": "1.02", + "SHORTNAME": FLDAS_CA_DAILY_SHORTNAME, + "LABEL": f"{FLDAS_CA_DAILY_SHORTNAME}.A{stamp}.001.nc.SUB.tif", + "DATASET_VERSION": "001", + "VARIABLES": EVAP_VARIABLE, + } + + +def _monthly_global_params(month_start: dt.date, bbox: str): + stamp = month_start.strftime("%Y%m") + return { + "FILENAME": ( + f"/data/FLDAS/{FLDAS_GLOBAL_MONTHLY_PRODUCT}/" + f"{month_start.year}/" + f"{FLDAS_GLOBAL_MONTHLY_SHORTNAME}.A{stamp}.001.nc" + ), + "VARIABLES": EVAP_VARIABLE, + "FORMAT": FLDAS_FORMAT, + "LABEL": f"{FLDAS_GLOBAL_MONTHLY_SHORTNAME}.A{stamp}.001.nc.SUB.tif", + "SERVICE": "L34RS_LDAS", + "DATASET_VERSION": "001", + "VERSION": "1.02", + "SHORTNAME": FLDAS_GLOBAL_MONTHLY_SHORTNAME, + "BBOX": bbox, + } + + +def _gesdisc_auth(): + username = getattr(settings, "USERNAME_GESDISC", None) + password = getattr(settings, "PASSWORD_GESDISC", None) + if not username or not password: + raise ValueError("USERNAME_GESDISC and PASSWORD_GESDISC are required") + return username, password + + +def _raise_if_html_response(response, first_chunk: bytes, url: str): + content_type = response.headers.get("Content-Type", "").lower() + stripped = first_chunk.lstrip().lower() + if "text/html" not in content_type and not stripped.startswith((b" 0 and not overwrite: + logger.info("Skipping existing ET raster: %s", output_path) + return "skipped" + + output_path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = output_path.with_suffix(f"{output_path.suffix}.tmp") + logger.info("Downloading ET raster: %s", output_path) + + try: + with session.get(url, stream=True, timeout=(30, 300)) as response: + if response.status_code != 200: + body = response.content[:500].decode("utf-8", errors="replace") + raise RuntimeError( + f"GES DISC download failed with HTTP {response.status_code}. " + f"URL={url}. Response starts with: {body}" + ) + + wrote_any = False + with tmp_path.open("wb") as handle: + for chunk in response.iter_content(chunk_size=1024 * 1024): + if not chunk: + continue + if not wrote_any: + _raise_if_html_response(response, chunk, url) + wrote_any = True + handle.write(chunk) + + if not wrote_any or tmp_path.stat().st_size == 0: + raise RuntimeError(f"GES DISC returned an empty raster for URL={url}") + + tmp_path.replace(output_path) + return "downloaded" + except Exception: + tmp_path.unlink(missing_ok=True) + raise + + +def _write_json(path: Path, payload): + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=True) + + +def _download_record(record, auth, overwrite, logger, max_attempts, retry_delay_seconds): + session = requests.Session() + session.auth = auth + try: + for attempt in range(1, max_attempts + 1): + record["attempts"] = attempt + try: + status = _download_file( + session=session, + url=record["url"], + output_path=Path(record["path"]), + overwrite=overwrite, + logger=logger, + ) + record["status"] = status + record.pop("error", None) + return status + except Exception as exc: + record["error"] = str(exc) + logger.warning( + "ET raster download failed for %s on attempt %s/%s: %s", + record["path"], + attempt, + max_attempts, + exc, + ) + if attempt < max_attempts: + time.sleep(retry_delay_seconds * attempt) + + record["status"] = "failed" + logger.error( + "Skipping ET raster after %s failed attempts: %s", + max_attempts, + record["path"], + ) + return "failed" + finally: + session.close() + + +def _download_records(auth, records, overwrite, logger, max_attempts, retry_delay_seconds, max_workers): + downloaded_count = 0 + skipped_count = 0 + failed_count = 0 + if not records: + return downloaded_count, skipped_count, failed_count + + worker_count = min(max_workers, len(records)) + logger.info("Downloading %s ET rasters with %s workers", len(records), worker_count) + with ThreadPoolExecutor(max_workers=worker_count) as executor: + futures = [ + executor.submit( + _download_record, + record, + auth, + overwrite, + logger, + max_attempts, + retry_delay_seconds, + ) + for record in records + ] + for future in as_completed(futures): + status = future.result() + if status == "downloaded": + downloaded_count += 1 + elif status == "skipped": + skipped_count += 1 + elif status == "failed": + failed_count += 1 + + return downloaded_count, skipped_count, failed_count + + +def _invalid_pixel_mask(values: np.ndarray, nodata) -> np.ndarray: + invalid = ~np.isfinite(values) + if nodata is not None: + invalid |= values == nodata + invalid |= values < 0 + return invalid + + +def _fill_raster_band( + values: np.ndarray, + nodata, + *, + max_search_distance: float, + smoothing_iterations: int, +): + invalid = _invalid_pixel_mask(values, nodata) + invalid_count = int(invalid.sum()) + if invalid_count == 0: + return values, invalid_count + + valid = ~invalid + if not valid.any(): + raise ValueError("raster band has no valid pixels to fill from") + + working = values.astype("float32", copy=True) + working[invalid] = 0.0 + mask = valid.astype("uint8") + filled = fillnodata( + working, + mask=mask, + max_search_distance=max_search_distance, + smoothing_iterations=smoothing_iterations, + ) + return filled, invalid_count + + +def _patch_fill_raster( + input_path: Path, + output_path: Path, + *, + overwrite: bool, + max_search_distance: float | None, + smoothing_iterations: int, +): + if output_path.exists() and output_path.stat().st_size > 0 and not overwrite: + return "skipped", 0, None + + output_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = output_path.with_name( + f"{output_path.stem}.{int(time.time() * 1000)}.tmp{output_path.suffix}" + ) + try: + with rasterio.open(input_path) as src: + profile = src.profile.copy() + nodata = src.nodata + data = src.read() + search_distance = ( + float(max(src.width, src.height)) + if max_search_distance is None + else max_search_distance + ) + + filled = np.empty((data.shape[0], data.shape[1], data.shape[2]), dtype="float32") + invalid_pixels = 0 + for band_index in range(data.shape[0]): + filled_band, band_invalid = _fill_raster_band( + data[band_index].astype("float32", copy=False), + nodata, + max_search_distance=search_distance, + smoothing_iterations=smoothing_iterations, + ) + filled[band_index] = filled_band + invalid_pixels += band_invalid + + profile.update( + dtype="float32", + count=data.shape[0], + compress=profile.get("compress") or "deflate", + tiled=profile.get("tiled", True), + ) + with rasterio.open(temporary_path, "w", **profile) as dst: + dst.write(filled) + temporary_path.replace(output_path) + return "patch_filled", invalid_pixels, search_distance + except Exception: + temporary_path.unlink(missing_ok=True) + raise + + +def _patch_fill_record( + record, + output_root, + overwrite, + max_search_distance, + smoothing_iterations, +): + input_path = Path(record["path"]) + output_path = Path(output_root) / input_path.name + patch_record = { + key: value for key, value in record.items() if key in {"date", "month", "url"} + } + patch_record.update( + { + "source_path": str(input_path), + "path": str(output_path), + } + ) + + if not input_path.exists(): + patch_record["status"] = "failed" + patch_record["error"] = f"source raster not found: {input_path}" + return "failed", patch_record + + try: + status, invalid_pixels, search_distance = _patch_fill_raster( + input_path, + output_path, + overwrite=overwrite, + max_search_distance=max_search_distance, + smoothing_iterations=smoothing_iterations, + ) + patch_record["status"] = status + patch_record["invalid_pixels_filled"] = invalid_pixels + patch_record["max_search_distance"] = search_distance + patch_record.pop("error", None) + return status, patch_record + except Exception as exc: + patch_record["status"] = "failed" + patch_record["error"] = str(exc) + return "failed", patch_record + + +def _patch_fill_records( + source_name, + records, + output_root, + *, + overwrite, + logger, + max_workers, + max_search_distance=None, + smoothing_iterations=0, +): + patched_count = 0 + skipped_count = 0 + failed_count = 0 + invalid_pixels = 0 + patch_records = [] + if not records: + return patch_records, patched_count, skipped_count, failed_count, invalid_pixels + + output_root = Path(output_root) + output_root.mkdir(parents=True, exist_ok=True) + worker_count = min(max_workers, len(records)) + logger.info( + "Patch-filling %s ET rasters into %s with %s workers", + source_name, + output_root, + worker_count, + ) + with ThreadPoolExecutor(max_workers=worker_count) as executor: + futures = [ + executor.submit( + _patch_fill_record, + record, + output_root, + overwrite, + max_search_distance, + smoothing_iterations, + ) + for record in records + ] + for future in as_completed(futures): + status, patch_record = future.result() + patch_records.append(patch_record) + if status == "patch_filled": + patched_count += 1 + invalid_pixels += int(patch_record.get("invalid_pixels_filled") or 0) + elif status == "skipped": + skipped_count += 1 + elif status == "failed": + failed_count += 1 + logger.error( + "Patch-fill failed for %s: %s", + patch_record.get("source_path"), + patch_record.get("error"), + ) + + patch_records.sort(key=lambda item: item.get("date") or item.get("month") or "") + return patch_records, patched_count, skipped_count, failed_count, invalid_pixels + + +def download_pan_india_et_assets( + output_root, + start_date, + end_date, + *, + et_root=None, + overwrite=False, + fldas_ca_daily_bbox=FLDAS_CA_DAILY_BBOX, + fldas_global_monthly_bbox=PAN_INDIA_BBOX, + max_attempts=5, + retry_delay_seconds=5, + max_workers=DEFAULT_MAX_WORKERS, + patch_fill=True, + patch_fill_max_search_distance=None, + patch_fill_smoothing_iterations=0, + logger=None, +): + """ + Download source ET rasters used by the GEE hydrology flow. + + The Central Asia FLDAS daily product is stored for the northern/high-resolution + branch. The global monthly FLDAS product is stored for the pan-India fallback + branch currently named GLDAS in utilities/constants.py. + """ + logger = logger or logging.getLogger(__name__) + start_date = _coerce_date(start_date, "start_date") + end_date = _coerce_date(end_date, "end_date") + if end_date <= start_date: + raise ValueError("end_date must be after start_date") + if max_attempts < 1: + raise ValueError("max_attempts must be at least 1") + if max_workers < 1: + raise ValueError("max_workers must be at least 1") + if patch_fill_smoothing_iterations < 0: + raise ValueError("patch_fill_smoothing_iterations must be >= 0") + if ( + patch_fill_max_search_distance is not None + and patch_fill_max_search_distance < 0 + ): + raise ValueError("patch_fill_max_search_distance must be >= 0") + + output_root = Path(output_root) + et_root = Path(et_root) if et_root is not None else output_root / "et" + daily_root = et_root / FLDAS_CA_DAILY_SOURCE / "daily" + monthly_root = et_root / FLDAS_GLOBAL_MONTHLY_SOURCE / "monthly" + daily_patch_filled_root = et_root / FLDAS_CA_DAILY_PATCH_FILLED_SOURCE + monthly_patch_filled_root = et_root / FLDAS_GLOBAL_MONTHLY_PATCH_FILLED_SOURCE + + daily_records = [] + for day in _iter_days(start_date, end_date): + params = _daily_ca_params(day, fldas_ca_daily_bbox) + daily_records.append( + { + "date": day.isoformat(), + "path": str(daily_root / f"{day:%Y%m%d}.tif"), + "url": _download_url(params), + } + ) + + monthly_records = [] + for month_start in _iter_month_starts(start_date, end_date): + params = _monthly_global_params(month_start, fldas_global_monthly_bbox) + monthly_records.append( + { + "month": month_start.strftime("%Y-%m"), + "path": str(monthly_root / f"{month_start:%Y%m}.tif"), + "url": _download_url(params), + } + ) + + auth = _gesdisc_auth() + + logger.info( + "Downloading local ET assets for [%s, %s): %s daily CA rasters, %s global monthly rasters", + start_date, + end_date, + len(daily_records), + len(monthly_records), + ) + daily_downloaded, daily_skipped, daily_failed = _download_records( + auth, + daily_records, + overwrite, + logger, + max_attempts, + retry_delay_seconds, + max_workers, + ) + monthly_downloaded, monthly_skipped, monthly_failed = _download_records( + auth, + monthly_records, + overwrite, + logger, + max_attempts, + retry_delay_seconds, + max_workers, + ) + + daily_patch_records = [] + monthly_patch_records = [] + daily_patch_count = daily_patch_skipped = daily_patch_failed = 0 + monthly_patch_count = monthly_patch_skipped = monthly_patch_failed = 0 + daily_patch_invalid_pixels = monthly_patch_invalid_pixels = 0 + if patch_fill: + ( + daily_patch_records, + daily_patch_count, + daily_patch_skipped, + daily_patch_failed, + daily_patch_invalid_pixels, + ) = _patch_fill_records( + FLDAS_CA_DAILY_SOURCE, + daily_records, + daily_patch_filled_root, + overwrite=overwrite, + logger=logger, + max_workers=max_workers, + max_search_distance=patch_fill_max_search_distance, + smoothing_iterations=patch_fill_smoothing_iterations, + ) + ( + monthly_patch_records, + monthly_patch_count, + monthly_patch_skipped, + monthly_patch_failed, + monthly_patch_invalid_pixels, + ) = _patch_fill_records( + FLDAS_GLOBAL_MONTHLY_SOURCE, + monthly_records, + monthly_patch_filled_root, + overwrite=overwrite, + logger=logger, + max_workers=max_workers, + max_search_distance=patch_fill_max_search_distance, + smoothing_iterations=patch_fill_smoothing_iterations, + ) + + sources = [ + SourceManifest( + source=FLDAS_CA_DAILY_SOURCE, + shortname=FLDAS_CA_DAILY_SHORTNAME, + product=FLDAS_CA_DAILY_PRODUCT, + variable=EVAP_VARIABLE, + bbox=fldas_ca_daily_bbox, + temporal_resolution="daily", + output_folder=str(daily_root), + file_count=len(daily_records), + downloaded_count=daily_downloaded, + skipped_count=daily_skipped, + failed_count=daily_failed, + patch_filled_source=( + FLDAS_CA_DAILY_PATCH_FILLED_SOURCE if patch_fill else None + ), + patch_filled_folder=( + str(daily_patch_filled_root) if patch_fill else None + ), + patch_filled_count=daily_patch_count, + patch_fill_skipped_count=daily_patch_skipped, + patch_fill_failed_count=daily_patch_failed, + patch_fill_invalid_pixels=daily_patch_invalid_pixels, + ), + SourceManifest( + source=FLDAS_GLOBAL_MONTHLY_SOURCE, + shortname=FLDAS_GLOBAL_MONTHLY_SHORTNAME, + product=FLDAS_GLOBAL_MONTHLY_PRODUCT, + variable=EVAP_VARIABLE, + bbox=fldas_global_monthly_bbox, + temporal_resolution="monthly", + output_folder=str(monthly_root), + file_count=len(monthly_records), + downloaded_count=monthly_downloaded, + skipped_count=monthly_skipped, + failed_count=monthly_failed, + patch_filled_source=( + FLDAS_GLOBAL_MONTHLY_PATCH_FILLED_SOURCE if patch_fill else None + ), + patch_filled_folder=( + str(monthly_patch_filled_root) if patch_fill else None + ), + patch_filled_count=monthly_patch_count, + patch_fill_skipped_count=monthly_patch_skipped, + patch_fill_failed_count=monthly_patch_failed, + patch_fill_invalid_pixels=monthly_patch_invalid_pixels, + ), + ] + + manifest = { + "created_at": dt.datetime.now(dt.timezone.utc).isoformat(), + "output_root": str(output_root), + "et_root": str(et_root), + "date_range": { + "start_date": start_date.isoformat(), + "end_date": end_date.isoformat(), + "end_date_is_exclusive": True, + }, + "purpose": ( + "Local source rasters for replacing GEE evapotranspiration inputs in " + "computing/mws/generate_hydrology.py." + ), + "retry_policy": { + "max_attempts": max_attempts, + "retry_delay_seconds": retry_delay_seconds, + "failed_records_are_skipped": True, + }, + "download_policy": { + "max_workers": max_workers, + }, + "patch_fill_policy": { + "enabled": bool(patch_fill), + "method": "rasterio.fill.fillnodata", + "invalid_pixels": ["nodata", "nan", "inf", "negative"], + "max_search_distance": ( + "max(width, height) per raster" + if patch_fill_max_search_distance is None + else patch_fill_max_search_distance + ), + "smoothing_iterations": patch_fill_smoothing_iterations, + }, + "sources": [asdict(source) for source in sources], + "records": { + FLDAS_CA_DAILY_SOURCE: daily_records, + FLDAS_GLOBAL_MONTHLY_SOURCE: monthly_records, + FLDAS_CA_DAILY_PATCH_FILLED_SOURCE: daily_patch_records, + FLDAS_GLOBAL_MONTHLY_PATCH_FILLED_SOURCE: monthly_patch_records, + }, + } + _write_json(et_root / "manifest.json", manifest) + _write_json( + et_root / FLDAS_CA_DAILY_SOURCE / "metadata.json", + { + "source": asdict(sources[0]), + "records": daily_records, + }, + ) + _write_json( + et_root / FLDAS_GLOBAL_MONTHLY_SOURCE / "metadata.json", + { + "source": asdict(sources[1]), + "records": monthly_records, + }, + ) + return manifest diff --git a/computing/hydrology_gpu/lulc_mapping.py b/computing/hydrology_gpu/lulc_mapping.py new file mode 100644 index 00000000..bda70140 --- /dev/null +++ b/computing/hydrology_gpu/lulc_mapping.py @@ -0,0 +1,128 @@ +from datetime import datetime + +import cupy as cp + + +LULC_SOURCE_DYNAMICWORLD = "dynamicworld" +LULC_SOURCE_INDIASATV3 = "indiasatv3" +LULC_SOURCES = (LULC_SOURCE_DYNAMICWORLD, LULC_SOURCE_INDIASATV3) + +SEASON_KHARIF = "kharif" +SEASON_RABI = "rabi" +SEASON_ZAID = "zaid" +SEASON_STATIC = "static" + +DW_WATER = 0 +DW_TREES = 1 +DW_CROPS = 4 +DW_SHRUB_AND_SCRUB = 5 +DW_BUILT = 6 +DW_BARE = 7 + +INDIASAT_BACKGROUND = 0 +INDIASAT_BUILT_UP = 1 +INDIASAT_WATER_KHARIF = 2 +INDIASAT_WATER_KHARIF_RABI = 3 +INDIASAT_WATER_ALL_SEASONS = 4 +INDIASAT_TREE_FORESTS = 6 +INDIASAT_BARRENLANDS = 7 +INDIASAT_SINGLE_CROPPING = 8 +INDIASAT_SINGLE_NON_KHARIF_CROPPING = 9 +INDIASAT_DOUBLE_CROPPING = 10 +INDIASAT_TRIPLE_CROPPING = 11 +INDIASAT_SHRUB_SCRUB = 12 + + +def normalize_lulc_source(source: str) -> str: + normalized = str(source or LULC_SOURCE_DYNAMICWORLD).strip().lower() + if normalized not in LULC_SOURCES: + raise ValueError( + f"Unsupported LULC source {source!r}; expected one of {', '.join(LULC_SOURCES)}" + ) + return normalized + + +def month_from_timestamp(timestamp) -> int: + if hasattr(timestamp, "month"): + return int(timestamp.month) + + text = str(timestamp) + for fmt in ("%Y%m%d_%H", "%Y%m%d"): + try: + return datetime.strptime(text, fmt).month + except ValueError: + pass + + try: + return datetime.fromisoformat(text.replace("Z", "+00:00")).month + except ValueError as exc: + raise ValueError(f"Cannot parse rainfall timestamp {timestamp!r}") from exc + + +def season_from_month(month: int) -> str: + if month in (7, 8, 9, 10): + return SEASON_KHARIF + if month in (11, 12, 1, 2): + return SEASON_RABI + if month in (3, 4, 5, 6): + return SEASON_ZAID + raise ValueError(f"Invalid month {month!r}") + + +def lulc_cache_key_for_timestamp(source: str, timestamp) -> str: + source = normalize_lulc_source(source) + if source == LULC_SOURCE_DYNAMICWORLD: + return SEASON_STATIC + return season_from_month(month_from_timestamp(timestamp)) + + +def nodata_lulc_value_for_source(source: str) -> int: + source = normalize_lulc_source(source) + if source == LULC_SOURCE_DYNAMICWORLD: + return DW_SHRUB_AND_SCRUB + return INDIASAT_BACKGROUND + + +def map_lulc_to_dynamic_world(raw_lulc: cp.ndarray, source: str, timestamp) -> cp.ndarray: + source = normalize_lulc_source(source) + if source == LULC_SOURCE_DYNAMICWORLD: + return raw_lulc + + season = season_from_month(month_from_timestamp(timestamp)) + return map_indiasatv3_to_dynamic_world(raw_lulc, season) + + +def map_indiasatv3_to_dynamic_world(raw_lulc: cp.ndarray, season: str) -> cp.ndarray: + # Background/unknown classes stay shrub/scrub instead of becoming water. + mapped = cp.full(raw_lulc.shape, DW_SHRUB_AND_SCRUB, dtype=cp.uint8) + + mapped = cp.where(raw_lulc == INDIASAT_BUILT_UP, DW_BUILT, mapped) + mapped = cp.where(raw_lulc == INDIASAT_TREE_FORESTS, DW_TREES, mapped) + mapped = cp.where(raw_lulc == INDIASAT_BARRENLANDS, DW_BARE, mapped) + mapped = cp.where(raw_lulc == INDIASAT_SHRUB_SCRUB, DW_SHRUB_AND_SCRUB, mapped) + + water = ( + (raw_lulc == INDIASAT_WATER_KHARIF) + | (raw_lulc == INDIASAT_WATER_KHARIF_RABI) + | (raw_lulc == INDIASAT_WATER_ALL_SEASONS) + ) + mapped = cp.where(water, DW_WATER, mapped) + + mapped = cp.where( + raw_lulc == INDIASAT_SINGLE_CROPPING, + DW_CROPS if season == SEASON_KHARIF else DW_SHRUB_AND_SCRUB, + mapped, + ) + mapped = cp.where( + raw_lulc == INDIASAT_SINGLE_NON_KHARIF_CROPPING, + DW_CROPS if season == SEASON_RABI else DW_SHRUB_AND_SCRUB, + mapped, + ) + mapped = cp.where( + raw_lulc == INDIASAT_DOUBLE_CROPPING, + DW_CROPS if season in (SEASON_KHARIF, SEASON_RABI) else DW_SHRUB_AND_SCRUB, + mapped, + ) + mapped = cp.where(raw_lulc == INDIASAT_TRIPLE_CROPPING, DW_CROPS, mapped) + + return mapped.astype(cp.uint8, copy=False) diff --git a/computing/hydrology_gpu/runoff.py b/computing/hydrology_gpu/runoff.py new file mode 100644 index 00000000..4c7eced2 --- /dev/null +++ b/computing/hydrology_gpu/runoff.py @@ -0,0 +1,202 @@ +import shutil +from contextlib import contextmanager +from pathlib import Path +from time import perf_counter +from .downloads import lulc, soil +from .algorithms import tiled_timeseries, timeseries +from .downloads import rainfall +from . import config as cfg +from .utils import GeoTIFFHandler, make_logger +from . import utils +from .watershed_boundary import ( + DEFAULT_WATERSHED_ROOT, + download_boundary_path, + materialize_district_boundary, + materialize_pan_india_boundary, + materialize_state_boundary, + materialize_tehsil_boundary, +) + +logger = make_logger("runoff_only_with_rainfall.log") +PAN_INDIA_DEFAULT_TILE_SIZE = 11264 +STATE_DEFAULT_TILE_SIZE = 4096 + + +def format_elapsed(seconds): + if seconds < 60: + return f"{seconds:.2f}s" + minutes, seconds = divmod(seconds, 60) + if minutes < 60: + return f"{int(minutes)}m {seconds:.2f}s" + hours, minutes = divmod(minutes, 60) + return f"{int(hours)}h {int(minutes)}m {seconds:.2f}s" + + +@contextmanager +def timed_stage(name): + start_time = perf_counter() + logger.info("Starting %s", name) + try: + yield + except Exception: + logger.exception("Failed %s after %s", name, format_elapsed(perf_counter() - start_time)) + raise + else: + logger.info("Finished %s in %s", name, format_elapsed(perf_counter() - start_time)) + + +def validate_local_raster(path_value, option_name): + if not path_value: + return + + path = Path(path_value) + if not path.exists(): + raise ValueError(f"{option_name} path does not exist: {path}") + + if path.is_dir(): + has_tif = any( + child.is_file() and child.suffix.lower() in {".tif", ".tiff"} + for child in path.rglob("*") + ) + if not has_tif: + raise ValueError(f"{option_name} directory has no GeoTIFF files: {path}") + + +def resolve_boundary(args): + selectors = [args.pan_india, args.state, args.district, args.tehsil] + if not any(selectors): + return + if args.pan_india and any([args.state, args.district, args.tehsil]): + raise ValueError("--pan-india cannot be combined with --state, --district, or --tehsil") + if args.pan_india: + microwatersheds_path, source_paths, feature_count = materialize_pan_india_boundary( + watershed_root=args.watershed_root, + output_path=args.watershed_boundary_output, + overwrite=not args.reuse_watershed_boundary, + ) + boundary_path = download_boundary_path(microwatersheds_path) + args.boundary = str(boundary_path) + args.microwatersheds = str(microwatersheds_path) + logger.info( + "Resolved pan-India watershed boundary: sources=%s download_boundary=%s microwatersheds=%s features=%s", + len(source_paths), + boundary_path, + microwatersheds_path, + feature_count, + ) + return + + if not args.state: + raise ValueError("--state is required for watershed boundary lookup") + if args.tehsil and not args.district: + raise ValueError("--district is required when --tehsil is provided") + + if not args.district: + microwatersheds_path, source_paths, feature_count = materialize_state_boundary( + state=args.state, + watershed_root=args.watershed_root, + output_path=args.watershed_boundary_output, + overwrite=not args.reuse_watershed_boundary, + ) + boundary_path = download_boundary_path(microwatersheds_path) + args.boundary = str(boundary_path) + args.microwatersheds = str(microwatersheds_path) + logger.info( + "Resolved state watershed boundary: state=%s sources=%s download_boundary=%s microwatersheds=%s features=%s", + args.state, + len(source_paths), + boundary_path, + microwatersheds_path, + feature_count, + ) + return + + if not args.tehsil: + microwatersheds_path, source_paths, feature_count = materialize_district_boundary( + state=args.state, + district=args.district, + watershed_root=args.watershed_root, + output_path=args.watershed_boundary_output, + overwrite=not args.reuse_watershed_boundary, + ) + boundary_path = download_boundary_path(microwatersheds_path) + args.boundary = str(boundary_path) + args.microwatersheds = str(microwatersheds_path) + logger.info( + "Resolved district watershed boundary: state=%s district=%s sources=%s download_boundary=%s microwatersheds=%s features=%s", + args.state, + args.district, + len(source_paths), + boundary_path, + microwatersheds_path, + feature_count, + ) + return + + boundary_path, source_path, feature_count = materialize_tehsil_boundary( + state=args.state, + district=args.district, + tehsil=args.tehsil, + watershed_root=args.watershed_root, + output_path=args.watershed_boundary_output, + overwrite=not args.reuse_watershed_boundary, + ) + args.boundary = str(boundary_path) + args.microwatersheds = str(boundary_path) + logger.info( + "Resolved watershed boundary: state=%s district=%s tehsil=%s source=%s output=%s features=%s", + args.state, + args.district, + args.tehsil, + source_path, + boundary_path, + feature_count, + ) + +def _required_arg(args, name): + value = getattr(args, name, None) + if value is None: + raise ValueError(f"{name} is required") + return value + + +def modify_cfg(args): + cfg.BOUNDARY_GEOJSON_PATH = _required_arg(args, "boundary") + cfg.MICROWATERSHEDS_PATH = _required_arg(args, "microwatersheds") + cfg.LULC_SOURCE = args.lulc_source + if args.local_lulc is not None: + cfg.LULC_PATH = args.local_lulc + if args.local_soil is not None: + cfg.SOIL_PATH = args.local_soil + cfg.RAINFALL_FOLDER = _required_arg(args, "rainfall_folder") + cfg.RUNOFFS_FOLDER = _required_arg(args, "runoffs_folder") + if args.t: + configured_timeseries = getattr(args, "timeseries_vector", None) + if configured_timeseries: + cfg.TIMESERIES_VECTOR = Path(configured_timeseries) + else: + path_obj = Path(cfg.MICROWATERSHEDS_PATH) + new_path = path_obj.with_name(f"{path_obj.stem}_timeseries{path_obj.suffix}") + cfg.TIMESERIES_VECTOR = new_path + cfg.ARG_START_DATE = args.start + cfg.ARG_END_DATE = args.end + cfg.TILE_SIZE = args.tile_size + +def prereq(args): + downloaders = [] + if args.local_lulc: + logger.info("Using local LULC from %s; skipping downloads.lulc.Downloader", args.local_lulc) + else: + downloaders.append(lulc.Downloader) + + if args.local_soil: + logger.info("Using local soil from %s; skipping downloads.soil.Downloader", args.local_soil) + else: + downloaders.append(soil.Downloader) + + downloaders.append(rainfall.Download_to_database) + + for downloader in downloaders: + stage_name = f"prerequisite downloader: {downloader.__module__}.{downloader.__name__}" + with timed_stage(stage_name): + downloader().main() diff --git a/computing/hydrology_gpu/utils.py b/computing/hydrology_gpu/utils.py new file mode 100644 index 00000000..eccdc59b --- /dev/null +++ b/computing/hydrology_gpu/utils.py @@ -0,0 +1,399 @@ +import json +from pathlib import Path +from shapely.geometry import shape +import rasterio +from rasterio.coords import BoundingBox +from rasterio import features +from rasterio.windows import Window, bounds as window_bounds, from_bounds, transform as window_transform +import os +import logging +import numpy.typing as nptypes +import cupy as cp +import numpy as np +from tqdm import tqdm +from typing import Any, List +import xarray as xr +import rioxarray +from rasterio.enums import Resampling +from rasterio.vrt import WarpedVRT + +def load_tif_image(file_path) -> nptypes.NDArray[Any]: + """Load a .tif image efficiently and return a NumPy array (float32).""" + with rasterio.open(file_path) as src: + image = src.read(1) # Read only the first band + print(f"Loaded Raster - Shape: {image.shape}, Dtype: {image.dtype}") + return image # Kept as NumPy array for easier slicing + +def make_logger(file_name, gpu_mem_usage=False, level=logging.INFO): + os.makedirs("logs", exist_ok=True) + + # Create logger + logger = logging.getLogger("mylog") + # without level, nothing gets print. I guess level is set to WARN etc + logger.setLevel(level) + + # File handler + fh = logging.FileHandler("logs/" + file_name) + fh.setLevel(level) + + # Console handler + ch = logging.StreamHandler() + ch.setLevel(level) + + # Formatter + formatter = logging.Formatter( + "%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S" + ) + fh.setFormatter(formatter) + ch.setFormatter(formatter) + + # Add handlers + logger.addHandler(fh) + logger.addHandler(ch) + + if gpu_mem_usage: + class GPUMemUsageAdapter(logging.LoggerAdapter): + def __init__(self, logger, extra) -> None: + super().__init__(logger, extra) + # pynvml.nvmlInit() + # Get handle for the first GPU (device 0) + # Loop pynvml.nvmlDeviceGetCount() for multi-GPU setups + # self.GPU_HANDLE = pynvml.nvmlDeviceGetHandleByIndex(0) + self.mempool = cp.get_default_memory_pool() + # logger.info(f"pynvml initialized {self.GPU_HANDLE}") + + def process(self, msg, kwargs): + # mem_info = pynvml.nvmlDeviceGetMemoryInfo(self.GPU_HANDLE) + used_bytes = self.mempool.used_bytes() + # cp.cuda.runtime.memGetInfo() returns (free, total) bytes + free_device, total_device = cp.cuda.runtime.memGetInfo() + gpu_mem_str = f" | CuPy Used: {used_bytes / (1024**2):.0f}/{total_device / (1024**2):.0f} MiB" + return msg + gpu_mem_str, kwargs + + return GPUMemUsageAdapter(logger, {}) + + return logger + + +class GeoTIFFHandler: + """ + Saves tiff file with correct crs and transforms. + """ + + def __init__(self, tiff_path: str, logger: logging.Logger): + """ + Initialize by loading an existing GeoTIFF file and storing its properties. + """ + with rasterio.open(tiff_path) as src: + self.crs = src.crs + self.transform = src.transform + self.width = src.width + self.height = src.height + self.dtype = src.dtypes[0] # Get the data type of the first band + self.count = src.count # Number of bands + self.window = self._window_from_bounds + self.bounds = src.bounds + + # Read original data (optional) + # self.original_data = src.read(1) + + self.logger = logger + logger.info(f"Loaded TIFF: {tiff_path}") + logger.info(f"CRS: {self.crs}, Transform: {self.transform}, Size: {self.width}x{self.height}, Type: {self.dtype}, Window: {self.window}") + + def _window_from_bounds(self, left, bottom, right, top): + return from_bounds(left, bottom, right, top, transform=self.transform) + + def iter_windows(self, tile_size: int): + if tile_size <= 0: + raise ValueError("tile_size must be positive") + + for row_off in range(0, self.height, tile_size): + height = min(tile_size, self.height - row_off) + for col_off in range(0, self.width, tile_size): + width = min(tile_size, self.width - col_off) + yield Window(col_off, row_off, width, height) + + def for_window(self, window): + window = self._clipped_window(window, self.width, self.height) + if window is None: + raise ValueError("Window does not overlap the parent raster") + + child = self.__class__.__new__(self.__class__) + child.crs = self.crs + child.transform = window_transform(window, self.transform) + child.width = int(window.width) + child.height = int(window.height) + child.dtype = self.dtype + child.count = self.count + child.window = child._window_from_bounds + child.bounds = BoundingBox(*window_bounds(window, self.transform)) + child.logger = self.logger + return child + + @staticmethod + def _raster_paths(src_path): + path = Path(src_path) + if path.is_dir(): + paths = sorted( + child for child in path.rglob("*") + if child.is_file() and child.suffix.lower() in {".tif", ".tiff"} + ) + if not paths: + raise FileNotFoundError(f"No GeoTIFF files found under {path}") + return paths + return [path] + + @staticmethod + def _rounded_window(window): + col_off = int(np.floor(window.col_off)) + row_off = int(np.floor(window.row_off)) + col_stop = int(np.ceil(window.col_off + window.width)) + row_stop = int(np.ceil(window.row_off + window.height)) + + if col_stop <= col_off or row_stop <= row_off: + return None + + return Window(col_off, row_off, col_stop - col_off, row_stop - row_off) + + @staticmethod + def _clipped_window(window, width, height): + window = GeoTIFFHandler._rounded_window(window) + if window is None: + return None + + col_off = max(0, int(window.col_off)) + row_off = max(0, int(window.row_off)) + col_stop = min(width, int(window.col_off + window.width)) + row_stop = min(height, int(window.row_off + window.height)) + + if col_stop <= col_off or row_stop <= row_off: + return None + + return Window(col_off, row_off, col_stop - col_off, row_stop - row_off) + + def _load_raster_path_into(self, src_path, padded, fill_value): + with rasterio.open(src_path) as src: + if src.crs != self.crs: + raise ValueError(f"CRS mismatch for {src_path}; reproject first.") + + left = max(src.bounds.left, self.bounds.left) + right = min(src.bounds.right, self.bounds.right) + bottom = max(src.bounds.bottom, self.bounds.bottom) + top = min(src.bounds.top, self.bounds.top) + if left >= right or bottom >= top: + return False + + with WarpedVRT( + src, + crs=self.crs, + transform=self.transform, + width=self.width, + height=self.height, + resampling=Resampling.nearest, + ) as vrt: + data = vrt.read(1, masked=True) + + mask = np.ma.getmaskarray(data) + valid = ~mask + if not np.any(valid): + return True + + padded[valid] = data.filled(fill_value)[valid] + return True + + def save_tiff(self, new_data, output_path: str): + """ + Save a new TIFF file using the stored properties but with new data. + + :param new_data: 2D NumPy array containing new raster data. + :param output_path: Path to save the new TIFF file. + :param compression: Compression type for the TIFF file (default: "LZW"). + """ + if new_data.shape != (self.height, self.width): + raise ValueError(f"Data shape {new_data.shape} does not match the stored shape ({self.height}, {self.width})") + + # output_dir = os.path.dirname(output_path) + # if output_dir: + # os.makedirs(output_dir, exist_ok=True) + + gdal_options = { + 'tiled': True, + 'compress': 'ZSTD', # <--- Use Zstandard + # 'ZSTD_LEVEL': 6, # <--- Set to lowest level (1=fastest, 22=best) + 'ZSTD_LEVEL': 1, # <--- Set to lowest level (1=fastest, 22=best) + 'NUM_THREADS': 10 # <--- Essential for speed. we have 12 cores. + } + + with rasterio.open( + output_path, + "w", + driver="GTiff", + height=self.height, + width=self.width, + count=self.count, + dtype=new_data.dtype, + crs=self.crs, + transform=self.transform, + # compress=compression + **gdal_options + ) as dst: + dst.write(new_data, 1) # Write new data to band 1 + + self.logger.info(f"Saved new TIFF to {output_path}") + + def save_geozarr(self, new_data: np.ndarray, output_path: str, name:str="data"): + if new_data.shape != (self.height, self.width): + raise ValueError(f"Data shape {new_data.shape} does not match shape ({self.height}, {self.width})") + + # 1. Generate coordinates from transform (Top-Left + half-pixel offset) + x_coords = self.transform.c + (np.arange(self.width) + 0.5) * self.transform.a + y_coords = self.transform.f + (np.arange(self.height) + 0.5) * self.transform.e + + # 2. Wrap in Xarray + da = xr.DataArray( + new_data, + dims=("y", "x"), + coords={ + "y": y_coords, + "x": x_coords + }, + name=name + ) + + # 3. Attach CRS (Essential for GIS) + da.rio.write_crs(self.crs, inplace=True) + + # 4. Save to Zarr (Creates a directory at output_path) + da.to_dataset().to_zarr(output_path, mode="w", zarr_format=2) + + def save_geozarr_time(self, new_data: np.ndarray, time, output_path: str, name:str="data"): + if new_data.shape != (self.height, self.width): + raise ValueError(f"Data shape {new_data.shape} does not match shape ({self.height}, {self.width})") + + # 1. Generate coordinates from transform (Top-Left + half-pixel offset) + x_coords = self.transform.c + (np.arange(self.width) + 0.5) * self.transform.a + y_coords = self.transform.f + (np.arange(self.height) + 0.5) * self.transform.e + + # 2. Wrap in Xarray + da = xr.DataArray( + [new_data], + dims=("time", "y", "x"), + coords={ + "time": [time], + "y": y_coords, + "x": x_coords + }, + name=name + ) + + # 3. Attach CRS (Essential for GIS) + da.rio.write_crs(self.crs, inplace=True) + + # 4. Save to Zarr (Creates a directory at output_path) + if not os.path.exists(output_path): + da.to_dataset().to_zarr(output_path, mode="w", zarr_format=2) + else: + da.to_dataset().to_zarr(output_path, mode="a", zarr_format=2, append_dim="time") + + def save_multiband_tiff(self, output_path:str, data_arrays: list[nptypes.NDArray[Any]], compression="LZW"): + """ + Save multiple 2D NumPy arrays as bands in a single GeoTIFF. + data_arrays: list or tuple of 2D numpy arrays with identical shape. + """ + count = len(data_arrays) + dtype = data_arrays[0].dtype + + with rasterio.open( + output_path, + "w", + driver="GTiff", + height=self.height, + width=self.width, + count=count, + dtype=dtype, + crs=self.crs, + transform=self.transform, + compress=compression + ) as dst: + for i, arr in tqdm(enumerate(data_arrays, start=1)): + dst.write(arr, i) + + self.logger.info("Done writing to " + output_path) + + def load_with_padding_inner(self, crs, data: np.ndarray, src_bounds, fill_value=0): + if crs != self.crs: + raise ValueError("CRS mismatch") + + # Use your existing window logic with the bounds from the dict + ref_window = self.window(*src_bounds) + row_off, col_off = int(ref_window.row_off), int(ref_window.col_off) + + padded = np.full((self.height, self.width), fill_value, dtype=data.dtype) + + height, width = data.shape + + dest_row0 = max(0, row_off) + dest_col0 = max(0, col_off) + dest_row1 = min(self.height, dest_row0 + height) + dest_col1 = min(self.width, dest_col0 + width) + + # Paste source data if overlapping + src_row0 = max(0, -row_off) + src_col0 = max(0, -col_off) + src_row1 = src_row0 + (dest_row1 - dest_row0) + src_col1 = src_col0 + (dest_col1 - dest_col0) + + padded[dest_row0:dest_row1, dest_col0:dest_col1] = data[src_row0:src_row1, src_col0:src_col1] + + return padded + + def load_with_padding(self, src_path, fill_value=0): + paths = self._raster_paths(src_path) + + with rasterio.open(paths[0]) as first: + padded = np.full((self.height, self.width), fill_value, dtype=first.dtypes[0]) + + loaded = 0 + for path in paths: + if self._load_raster_path_into(path, padded, fill_value): + loaded += 1 + + if loaded == 0: + raise ValueError(f"No raster data from {src_path} overlaps the reference grid") + + self.logger.info("Loaded %s overlapping raster file(s) from %s", loaded, src_path) + return padded + + + def rasterize_by_id(self, shapes, fill_value=0): + """ + Rasterizes GeoJSON features using their 'id' property as the pixel value. + Not generic enough, i think. + """ + # with open(geojson_path, 'r') as f: + # geojson_data = json.load(f) + + # 1. Extract (geometry, value) pairs + # This creates a list like: [(geom1, 817103), (geom2, 818258), ...] + + # shapes = [ + # (shape(feature['geometry']), feature['properties']['id']) + # for feature in geojson_data['features'] + # ] + + # 2. Rasterize + # Note: Use a dtype large enough for your IDs (e.g., int32 or float32) + mask = rasterio.features.rasterize( + shapes=shapes, + out_shape=(self.height, self.width), + transform=self.transform, + fill=fill_value, + dtype='int32' + ) + + return mask + + +tif_handler: GeoTIFFHandler = None diff --git a/computing/hydrology_gpu/watershed_boundary.py b/computing/hydrology_gpu/watershed_boundary.py new file mode 100644 index 00000000..4e99f817 --- /dev/null +++ b/computing/hydrology_gpu/watershed_boundary.py @@ -0,0 +1,448 @@ +import csv +import re +from pathlib import Path + +import geopandas as gpd +import pandas as pd +from shapely.geometry import GeometryCollection, MultiPolygon, Polygon + +from computing.config_loader import PRECOMPUTED_TEHSIL_WATERSHED_DIR, PROJECT_ROOT + +DEFAULT_WATERSHED_ROOT = PRECOMPUTED_TEHSIL_WATERSHED_DIR +DEFAULT_BOUNDARY_OUTPUT_ROOT = PROJECT_ROOT / "data" / "hydrology_gpu" / "boundaries" +DEFAULT_PAN_INDIA_DOWNLOAD_BOUNDARY = PROJECT_ROOT / "data" / "base_layers" / "PanIndia_Boundaries" / "india_state_outer_no_islands.geojson" +PAN_INDIA_SLUG = "pan_india" + + +def normalize_name(value: str) -> str: + text = str(value).strip().lower().replace("&", " and ") + text = re.sub(r"[^a-z0-9]+", " ", text) + return re.sub(r"\s+", " ", text).strip() + + +def slugify(value: str) -> str: + text = str(value).strip().lower() + text = re.sub(r"[^a-z0-9]+", "_", text) + return re.sub(r"_+", "_", text).strip("_") + + +def default_output_path(state: str, district: str, tehsil: str) -> Path: + return ( + DEFAULT_BOUNDARY_OUTPUT_ROOT + / slugify(state) + / slugify(district) + / f"{slugify(tehsil)}.geojson" + ) + + +def default_district_output_path(state: str, district: str) -> Path: + return ( + DEFAULT_BOUNDARY_OUTPUT_ROOT + / slugify(state) + / f"{slugify(district)}.geojson" + ) + + +def default_state_output_path(state: str) -> Path: + return ( + DEFAULT_BOUNDARY_OUTPUT_ROOT + / slugify(state) + / f"{slugify(state)}.geojson" + ) + + +def default_pan_india_output_path() -> Path: + return DEFAULT_BOUNDARY_OUTPUT_ROOT / PAN_INDIA_SLUG / f"{PAN_INDIA_SLUG}.geojson" + + +def download_boundary_path(boundary_path: str | Path) -> Path: + path = Path(boundary_path) + return path.with_name(f"{path.stem}_download_boundary{path.suffix}") + + +def manifest_path(root: Path) -> Path: + return root / "tehsil_watershed_manifest.csv" + + +def load_manifest(root: Path) -> list[dict]: + path = manifest_path(root) + if not path.exists(): + raise FileNotFoundError(f"Watershed manifest not found: {path}") + + with path.open(newline="") as f: + return list(csv.DictReader(f)) + + +def manifest_relative_output(root: Path, output_path: str) -> Path | None: + if not output_path: + return None + + raw_path = Path(output_path) + if raw_path.is_absolute() and raw_path.exists(): + return raw_path + + parts = raw_path.parts + if "tehsil_watersheds" in parts: + idx = parts.index("tehsil_watersheds") + candidate = root.joinpath(*parts[idx + 1 :]) + if candidate.exists(): + return candidate + + candidate = root / raw_path + if candidate.exists(): + return candidate + + return None + + +def fallback_gpkg_path(root: Path, state: str, district: str, tehsil: str) -> Path: + return root / slugify(state) / slugify(district) / f"{slugify(tehsil)}.gpkg" + + +def find_tehsil_watershed(root: Path, state: str, district: str, tehsil: str) -> tuple[Path, dict]: + wanted_state = normalize_name(state) + wanted_district = normalize_name(district) + wanted_tehsil = normalize_name(tehsil) + + for row in load_manifest(root): + if normalize_name(row.get("state", "")) != wanted_state: + continue + if normalize_name(row.get("district", "")) != wanted_district: + continue + if normalize_name(row.get("tehsil", "")) != wanted_tehsil: + continue + + if row.get("status") != "written": + raise ValueError(f"Watershed boundary is not available for {state}/{district}/{tehsil}: {row}") + + path = manifest_relative_output(root, row.get("output_path", "")) + if path is None: + path = fallback_gpkg_path(root, state, district, tehsil) + if not path.exists(): + raise FileNotFoundError(f"Manifest matched, but watershed file does not exist: {path}") + return path, row + + path = fallback_gpkg_path(root, state, district, tehsil) + if path.exists(): + return path, {"state": state, "district": district, "tehsil": tehsil, "status": "written"} + + raise FileNotFoundError( + "Could not find tehsil watershed for " + f"state={state!r}, district={district!r}, tehsil={tehsil!r} under {root}" + ) + + +def find_district_watersheds(root: Path, state: str, district: str) -> list[tuple[Path, dict]]: + wanted_state = normalize_name(state) + wanted_district = normalize_name(district) + matches = [] + + for row in load_manifest(root): + if normalize_name(row.get("state", "")) != wanted_state: + continue + if normalize_name(row.get("district", "")) != wanted_district: + continue + if row.get("status") != "written": + continue + + path = manifest_relative_output(root, row.get("output_path", "")) + if path is None: + path = fallback_gpkg_path(root, row.get("state", state), row.get("district", district), row.get("tehsil", "")) + if path.exists(): + matches.append((path, row)) + + if not matches: + raise FileNotFoundError( + f"Could not find written district watersheds for state={state!r}, district={district!r} under {root}" + ) + + return matches + + +def find_state_watersheds(root: Path, state: str) -> list[tuple[Path, dict]]: + wanted_state = normalize_name(state) + matches = [] + + for row in load_manifest(root): + if normalize_name(row.get("state", "")) != wanted_state: + continue + if row.get("status") != "written": + continue + + path = manifest_relative_output(root, row.get("output_path", "")) + if path is None: + path = fallback_gpkg_path( + root, + row.get("state", state), + row.get("district", ""), + row.get("tehsil", ""), + ) + if path.exists(): + matches.append((path, row)) + + if not matches: + raise FileNotFoundError(f"Could not find written state watersheds for state={state!r} under {root}") + + return matches + + +def find_pan_india_watersheds(root: Path) -> list[tuple[Path, dict]]: + matches = [] + + for row in load_manifest(root): + if row.get("status") != "written": + continue + + path = manifest_relative_output(root, row.get("output_path", "")) + if path is None: + path = fallback_gpkg_path( + root, + row.get("state", ""), + row.get("district", ""), + row.get("tehsil", ""), + ) + if path.exists(): + matches.append((path, row)) + + if not matches: + raise FileNotFoundError(f"Could not find written pan-India watersheds under {root}") + + return matches + + +def prepare_boundary_gdf(gdf, state: str, district: str, tehsil: str | None, source_path: Path): + if gdf.empty: + raise ValueError(f"Watershed file has no features: {source_path}") + + if gdf.crs is None: + gdf = gdf.set_crs("EPSG:4326") + elif gdf.crs.to_epsg() != 4326: + gdf = gdf.to_crs("EPSG:4326") + + gdf = gdf[gdf.geometry.notna() & ~gdf.geometry.is_empty].copy() + if gdf.empty: + raise ValueError(f"Watershed file has no valid geometries: {source_path}") + + if "id" in gdf.columns: + gdf = gdf.rename(columns={"id": "source_id"}) + if "uid" in gdf.columns and "watershed_uid" not in gdf.columns: + gdf["watershed_uid"] = gdf["uid"].astype(str) + + gdf["selected_state"] = state + gdf["selected_district"] = district + if tehsil is not None: + gdf["selected_tehsil"] = tehsil + elif "TEHSIL" in gdf.columns: + gdf["selected_tehsil"] = gdf["TEHSIL"].astype(str) + gdf["source_gpkg"] = str(source_path) + return gdf + + +def polygon_parts(geometry): + if isinstance(geometry, Polygon): + return [geometry] + if isinstance(geometry, MultiPolygon): + return list(geometry.geoms) + if isinstance(geometry, GeometryCollection): + parts = [] + for part in geometry.geoms: + parts.extend(polygon_parts(part)) + return parts + return [] + + +def strip_inner_rings(geometry): + polygons = [] + for polygon in polygon_parts(geometry): + if not polygon.is_empty: + polygons.append(Polygon(polygon.exterior)) + + if not polygons: + raise ValueError("Could not build an outer boundary from the watershed geometries") + if len(polygons) == 1: + return polygons[0] + return MultiPolygon(polygons) + + +def write_download_boundary(gdf, destination: str | Path, state: str, district: str | None = None) -> Path: + destination = Path(destination) + union_geometry = gdf.geometry.unary_union + outer_geometry = strip_inner_rings(union_geometry) + properties = { + "id": 1, + "selected_state": state, + "boundary_role": "download_outer_boundary", + } + if district is not None: + properties["selected_district"] = district + + outer_gdf = gpd.GeoDataFrame( + [properties], + geometry=[outer_geometry], + crs=gdf.crs, + ) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(outer_gdf.to_json()) + return destination + + +def materialize_tehsil_boundary( + state: str, + district: str, + tehsil: str, + watershed_root: str | Path = DEFAULT_WATERSHED_ROOT, + output_path: str | Path | None = None, + overwrite: bool = True, +) -> tuple[Path, Path, int]: + root = Path(watershed_root) + source_path, _ = find_tehsil_watershed(root, state, district, tehsil) + destination = Path(output_path) if output_path else default_output_path(state, district, tehsil) + + if destination.exists() and not overwrite: + gdf_existing = gpd.read_file(destination) + return destination, source_path, len(gdf_existing) + + gdf = prepare_boundary_gdf(gpd.read_file(source_path), state, district, tehsil, source_path) + gdf["id"] = range(1, len(gdf) + 1) + + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(gdf.to_json()) + + return destination, source_path, len(gdf) + + +def materialize_district_boundary( + state: str, + district: str, + watershed_root: str | Path = DEFAULT_WATERSHED_ROOT, + output_path: str | Path | None = None, + overwrite: bool = True, +) -> tuple[Path, list[Path], int]: + root = Path(watershed_root) + matches = find_district_watersheds(root, state, district) + destination = Path(output_path) if output_path else default_district_output_path(state, district) + + if destination.exists() and not overwrite: + gdf_existing = gpd.read_file(destination) + return destination, [path for path, _ in matches], len(gdf_existing) + + frames = [] + source_paths = [] + for source_path, row in matches: + gdf = prepare_boundary_gdf( + gpd.read_file(source_path), + state=state, + district=district, + tehsil=row.get("tehsil") or None, + source_path=source_path, + ) + frames.append(gdf) + source_paths.append(source_path) + + if not frames: + raise ValueError(f"No non-empty watershed files found for {state}/{district}") + + combined = gpd.GeoDataFrame(pd.concat(frames, ignore_index=True), crs="EPSG:4326") + combined["id"] = range(1, len(combined) + 1) + + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(combined.to_json()) + write_download_boundary(combined, download_boundary_path(destination), state, district) + + return destination, source_paths, len(combined) + + +def materialize_state_boundary( + state: str, + watershed_root: str | Path = DEFAULT_WATERSHED_ROOT, + output_path: str | Path | None = None, + overwrite: bool = True, +) -> tuple[Path, list[Path], int]: + root = Path(watershed_root) + matches = find_state_watersheds(root, state) + destination = Path(output_path) if output_path else default_state_output_path(state) + + if destination.exists() and not overwrite: + gdf_existing = gpd.read_file(destination) + return destination, [path for path, _ in matches], len(gdf_existing) + + frames = [] + source_paths = [] + for source_path, row in matches: + district = row.get("district") or source_path.parent.name + tehsil = row.get("tehsil") or None + gdf = prepare_boundary_gdf( + gpd.read_file(source_path), + state=state, + district=district, + tehsil=tehsil, + source_path=source_path, + ) + frames.append(gdf) + source_paths.append(source_path) + + if not frames: + raise ValueError(f"No non-empty watershed files found for {state}") + + combined = gpd.GeoDataFrame(pd.concat(frames, ignore_index=True), crs="EPSG:4326") + combined["id"] = range(1, len(combined) + 1) + + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(combined.to_json()) + write_download_boundary(combined, download_boundary_path(destination), state) + + return destination, source_paths, len(combined) + + +def materialize_pan_india_boundary( + watershed_root: str | Path = DEFAULT_WATERSHED_ROOT, + output_path: str | Path | None = None, + overwrite: bool = True, + download_boundary_source: str | Path = DEFAULT_PAN_INDIA_DOWNLOAD_BOUNDARY, +) -> tuple[Path, list[Path], int]: + root = Path(watershed_root) + matches = find_pan_india_watersheds(root) + destination = Path(output_path) if output_path else default_pan_india_output_path() + + if destination.exists() and not overwrite: + download_destination = download_boundary_path(destination) + download_boundary_source = Path(download_boundary_source) + if not download_destination.exists() and download_boundary_source.exists(): + download_destination.parent.mkdir(parents=True, exist_ok=True) + download_destination.write_text(download_boundary_source.read_text()) + gdf_existing = gpd.read_file(destination) + return destination, [path for path, _ in matches], len(gdf_existing) + + frames = [] + source_paths = [] + for source_path, row in matches: + state = row.get("state") or source_path.parent.parent.name + district = row.get("district") or source_path.parent.name + tehsil = row.get("tehsil") or None + gdf = prepare_boundary_gdf( + gpd.read_file(source_path), + state=state, + district=district, + tehsil=tehsil, + source_path=source_path, + ) + frames.append(gdf) + source_paths.append(source_path) + + if not frames: + raise ValueError("No non-empty watershed files found for pan-India") + + combined = gpd.GeoDataFrame(pd.concat(frames, ignore_index=True), crs="EPSG:4326") + combined["id"] = range(1, len(combined) + 1) + + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(combined.to_json()) + + download_destination = download_boundary_path(destination) + download_boundary_source = Path(download_boundary_source) + if download_boundary_source.exists(): + download_destination.write_text(download_boundary_source.read_text()) + else: + write_download_boundary(combined, download_destination, PAN_INDIA_SLUG) + + return destination, source_paths, len(combined) diff --git a/computing/mws/et_download.py b/computing/mws/et_download.py new file mode 100644 index 00000000..4ff99576 --- /dev/null +++ b/computing/mws/et_download.py @@ -0,0 +1,110 @@ +import logging + +from nrm_app.celery import app + +from computing.config_loader import PROJECT_ROOT +from computing.hydrology_gpu.et_download import download_pan_india_et_assets + +from .runoff_gpu import _parse_bool, _resolve_dates + + +PAN_INDIA_ET_OUTPUT_ROOT = PROJECT_ROOT / "data" / "base_layers" / "hydrology" / "et" + + +def _make_logger(): + log_dir = PROJECT_ROOT / "logs" + log_dir.mkdir(parents=True, exist_ok=True) + + logger = logging.getLogger("hydrology_gpu.et_download") + logger.setLevel(logging.INFO) + logger.propagate = False + if logger.handlers: + return logger + + formatter = logging.Formatter( + "%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + file_handler = logging.FileHandler(log_dir / "et_download.log") + file_handler.setLevel(logging.INFO) + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + + stream_handler = logging.StreamHandler() + stream_handler.setLevel(logging.INFO) + stream_handler.setFormatter(formatter) + logger.addHandler(stream_handler) + + return logger + + +def run_et_download_local( + *, + pan_india=False, + start_date=None, + end_date=None, + start_year=None, + end_year=None, + overwrite=False, + patch_fill=True, +): + pan_india = _parse_bool(pan_india) + if not pan_india: + raise ValueError("et_download currently supports pan_india=true only") + + start_date, end_date, annual_start_year, annual_end_year = _resolve_dates( + start_date=start_date, + end_date=end_date, + start_year=start_year, + end_year=end_year, + ) + annual_key = f"{annual_start_year}_{annual_end_year}" + output_root = PAN_INDIA_ET_OUTPUT_ROOT / annual_key + output_root.mkdir(parents=True, exist_ok=True) + + logger = _make_logger() + manifest = download_pan_india_et_assets( + output_root=output_root, + et_root=output_root, + start_date=start_date, + end_date=end_date, + overwrite=_parse_bool(overwrite), + patch_fill=_parse_bool(patch_fill), + logger=logger, + ) + + return { + "scope": "pan_india", + "start_date": start_date, + "end_date": end_date, + "annual_key": annual_key, + "output_root": str(output_root), + "et_root": manifest["et_root"], + "sources": manifest["sources"], + "patch_fill_policy": manifest.get("patch_fill_policy"), + "manifest": str(output_root / "manifest.json"), + } + + +@app.task(bind=True) +def et_download( + self, + pan_india=False, + start_date=None, + end_date=None, + start_year=None, + end_year=None, + overwrite=False, + patch_fill=True, +): + _ = self + return run_et_download_local( + pan_india=pan_india, + start_date=start_date, + end_date=end_date, + start_year=start_year, + end_year=end_year, + overwrite=overwrite, + patch_fill=patch_fill, + ) diff --git a/computing/mws/generate_hydrology_local.py b/computing/mws/generate_hydrology_local.py new file mode 100644 index 00000000..038e562d --- /dev/null +++ b/computing/mws/generate_hydrology_local.py @@ -0,0 +1,2457 @@ +import csv +import datetime as dt +import json +import math +import os +from collections import Counter, defaultdict +from pathlib import Path + +import geopandas as gpd +import numpy as np +import pandas as pd +import rasterio +from nrm_app.celery import app +from rasterio.features import geometry_mask +from rasterio.windows import Window, from_bounds +from shapely.geometry import box, mapping + +from computing.config_loader import ( + AQUIFER_VECTOR_PATH, + HYDROLOGY_LOCAL_OUTPUT_DIR, + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + PROJECT_ROOT, +) +from computing.hydrology_gpu.watershed_boundary import ( + find_pan_india_watersheds, + find_tehsil_watershed, +) +from computing.hydrology_gpu.et_download import ( + FLDAS_CA_DAILY_PATCH_FILLED_SOURCE, + FLDAS_CA_DAILY_SOURCE, + FLDAS_GLOBAL_MONTHLY_PATCH_FILLED_SOURCE, + FLDAS_GLOBAL_MONTHLY_SOURCE, +) +from computing.local_compute_helper import ( + build_output_vector_path, + load_precomputed_watersheds, + push_local_vector_to_geoserver, + read_validated_vector_file, + write_vector_output, +) +from computing.misc.aquifer_vector_local import ( + _compute_aquifer_properties_for_watersheds, + _prepare_aquifers_for_intersection, +) +from computing.mws.runoff_gpu import ( + HYDROLOGY_OUTPUT_ROOT, + PAN_INDIA_RUNOFF_OUTPUT_ROOT, + PAN_INDIA_RUNOFF_TIMESERIES_DIR_NAME, +) +from computing.mws.et_download import PAN_INDIA_ET_OUTPUT_ROOT +from computing.utils import ( + save_layer_info_to_db, + update_layer_sync_status, +) +from utilities.gee_utils import valid_gee_text + + +GEOSERVER_WORKSPACE = "mws_layers" +LOCAL_ALGORITHM = "local_hydrology" +LOCAL_ALGORITHM_VERSION = "local-1.0" +SECONDS_PER_DAY = 86400.0 +CACHE_UID_COLUMN = "uid" +CACHE_ET_SOURCE_COLUMN = "et_source" +CACHE_ET_SOURCE_SIGNATURE_COLUMN = "et_source_signature" +CACHE_ET_ERROR_COLUMN = "et_error" +SOURCE_YEARS_COLUMN = "source_years" +FORTNIGHT_ANCHOR_DATE = dt.date(2017, 7, 1) +HYDROLOGY_BASE_LAYER_ROOT = PROJECT_ROOT / "data" / "base_layers" / "hydrology" + + +def _parse_bool(value): + if isinstance(value, bool): + return value + if value is None: + return False + return str(value).strip().lower() in {"1", "true", "yes", "y"} + + +def _normalize_location(value, field_name): + if value is None or not str(value).strip(): + raise ValueError(f"{field_name} is required") + return str(value).strip().lower() + + +def _year_key(year): + return f"{year}_{year + 1}" + + +def _source_year_for_date(day): + boundary = dt.date(day.year, 7, 1) + return day.year if day >= boundary else day.year - 1 + + +def _source_years_for_period(start_date, end_date): + years = set() + current = start_date + while current < end_date: + years.add(_source_year_for_date(current)) + current += dt.timedelta(days=1) + return sorted(years) + + +def _source_years_for_periods(periods): + years = set() + for start_date, end_date, _ in periods: + years.update(_source_years_for_period(start_date, end_date)) + return sorted(years) + + +def _cache_root(output_base_dir): + return Path(output_base_dir) / "cache" + + +def _et_cache_root(year): + return Path(PAN_INDIA_ET_OUTPUT_ROOT) / _year_key(year) / "cache" + + +def _aquifer_cache_path(output_base_dir): + return ( + Path(HYDROLOGY_BASE_LAYER_ROOT) + / "aquifer" + / "cache" + / "aquifer_by_uid.parquet" + ) + + +def _et_cache_path(output_base_dir, year, is_annual): + period = "annual" if is_annual else "fortnight" + return _et_cache_root(year) / f"{period}.parquet" + + +def _et_aggregate_root(output_base_dir, year, is_annual): + period = "annual" if is_annual else "fortnight" + return _et_cache_root(year) / "rasters" / period + + +def _base_layer_period_name(is_annual): + return "annual" if is_annual else "fortnightly" + + +def _base_layer_name(year, is_annual): + return f"hydrology_{_base_layer_period_name(is_annual)}_{_year_key(year)}" + + +def _base_layer_path(base_layer_root, year, is_annual): + return ( + Path(base_layer_root) + / _base_layer_period_name(is_annual) + / f"{_base_layer_name(year, is_annual)}.gpkg" + ) + + +def _write_parquet_atomic(frame, path): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = path.with_name(f"{path.name}.{os.getpid()}.tmp") + frame.to_parquet(temporary_path, index=False) + temporary_path.replace(path) + + +def _layer_name(district, block, is_annual): + suffix = "_".join( + [ + valid_gee_text(district.lower()), + valid_gee_text(block.lower()), + ] + ) + prefix = "deltaG_well_depth_" if is_annual else "deltaG_fortnight_" + return prefix + suffix + + +def _build_periods(year, is_annual): + start = dt.date(year, 7, 1) + end = dt.date(year + 1, 7, 1) + if is_annual: + return [(start, end, _year_key(year))] + if year < FORTNIGHT_ANCHOR_DATE.year: + raise ValueError( + "Fortnightly hydrology starts from the 2017 agricultural year" + ) + + periods = [] + current = FORTNIGHT_ANCHOR_DATE + for previous_year in range(FORTNIGHT_ANCHOR_DATE.year, year): + previous_end = dt.date(previous_year + 1, 7, 1) + while current + dt.timedelta(days=14) <= previous_end: + current += dt.timedelta(days=14) + + while current + dt.timedelta(days=14) <= end: + period_end = current + dt.timedelta(days=14) + periods.append((current, period_end, current.isoformat())) + current = period_end + return periods + + +def _resolve_base_layer_year_bounds(year=None, start_year=None, end_year=None): + if year is not None and (start_year is not None or end_year is not None): + raise ValueError("Provide start_year/end_year or year, not both") + + if year is not None: + start_year = int(year) + end_year = start_year + 1 + else: + if start_year is None or end_year is None: + raise ValueError("start_year and end_year are required") + start_year = int(start_year) + end_year = int(end_year) + + if start_year < FORTNIGHT_ANCHOR_DATE.year: + raise ValueError("Hydrology base layers can only be generated from 2017") + if end_year <= start_year: + raise ValueError("end_year must be greater than start_year") + if end_year != start_year + 1: + raise ValueError( + "Hydrology base-layer generation supports one hydrological year " + "at a time; use end_year=start_year+1" + ) + return start_year, end_year + + +def _resolve_year_inputs( + year, + hydrology_output_root=HYDROLOGY_OUTPUT_ROOT, +): + year_key = _year_key(year) + runoff_roots = [ + Path(PAN_INDIA_RUNOFF_OUTPUT_ROOT), + Path(hydrology_output_root) / "pan_india", + ] + timeseries_dir = None + for runoff_root in runoff_roots: + candidate = ( + runoff_root + / year_key + / PAN_INDIA_RUNOFF_TIMESERIES_DIR_NAME + / "pan_india_timeseries_tile_series" + ) + if candidate.exists(): + timeseries_dir = candidate + break + + et_roots = [ + Path(PAN_INDIA_ET_OUTPUT_ROOT) / year_key, + Path(hydrology_output_root) / "pan_india" / year_key / "et", + ] + et_root = next((path for path in et_roots if path.exists()), None) + + if timeseries_dir is None: + available_years = sorted( + { + path.name + for runoff_root in runoff_roots + for path in runoff_root.glob("*_*") + if ( + path + / PAN_INDIA_RUNOFF_TIMESERIES_DIR_NAME + / "pan_india_timeseries_tile_series" + ).is_dir() + } + ) + raise FileNotFoundError( + "Pan-India rainfall/runoff timeseries not found for " + f"{year_key}: " + f"{PAN_INDIA_RUNOFF_OUTPUT_ROOT / year_key / PAN_INDIA_RUNOFF_TIMESERIES_DIR_NAME / 'pan_india_timeseries_tile_series'}. " + f"Available Pan-India runoff years: {available_years or 'none'}." + ) + if et_root is None: + available_et_years = sorted( + path.name + for path in Path(PAN_INDIA_ET_OUTPUT_ROOT).glob("*_*") + if path.is_dir() + ) + raise FileNotFoundError( + "Pan-India ET folder not found for " + f"{year_key}: {PAN_INDIA_ET_OUTPUT_ROOT / year_key}. " + f"Available Pan-India ET years: {available_et_years or 'none'}. " + "Run et_download for this hydrological year first." + ) + return timeseries_dir, et_root + + +def _resolve_source_year_inputs( + *, + output_year, + periods, + hydrology_output_root=HYDROLOGY_OUTPUT_ROOT, +): + source_inputs = {} + for source_year in _source_years_for_periods(periods): + try: + source_inputs[source_year] = _resolve_year_inputs( + year=source_year, + hydrology_output_root=hydrology_output_root, + ) + except FileNotFoundError as error: + raise FileNotFoundError( + f"Hydrology output year {_year_key(output_year)} requires " + f"source year {_year_key(source_year)} because fortnight " + f"windows are anchored at {FORTNIGHT_ANCHOR_DATE.isoformat()}." + ) from error + return source_inputs + + +def _pan_india_watershed_offset( + state, + district, + block, + watershed_root=PRECOMPUTED_TEHSIL_WATERSHED_DIR, +): + root = Path(watershed_root) + target_path, target_row = find_tehsil_watershed( + root, + state, + district, + block, + ) + target_path = target_path.resolve() + offset = 0 + + matches = find_pan_india_watersheds(root) + for match_index, (source_path, row) in enumerate(matches): + feature_count = int(row.get("feature_count") or 0) + if source_path.resolve() == target_path: + return ( + matches, + match_index, + offset, + feature_count, + str(target_path), + ) + offset += feature_count + + raise ValueError( + "The requested watershed file is not part of the Pan-India runoff " + f"boundary manifest: {target_path}. Manifest row: {target_row}" + ) + + +def _pan_india_series_path(series_dir, watershed_id): + return Path(series_dir) / f"{watershed_id // 1000:04d}" / f"{watershed_id}.csv" + + +def _pan_india_series_paths_exist(series_dirs, watershed_id): + return all( + _pan_india_series_path(series_dir, watershed_id).exists() + for series_dir in series_dirs + ) + + +def _read_pan_india_series(series_dir, watershed_id): + path = _pan_india_series_path(series_dir, watershed_id) + if not path.exists(): + return {}, path + + data = defaultdict(lambda: [0.0, 0, 0.0, 0]) + with path.open(newline="") as source: + for row in csv.reader(source): + if len(row) != 5: + raise ValueError( + f"Invalid Pan-India rainfall/runoff row in {path}: {row}" + ) + timestamp, rainfall_sum, rainfall_count, runoff_sum, runoff_count = row + values = data[timestamp] + values[0] += float(rainfall_sum) + values[1] += int(rainfall_count) + values[2] += float(runoff_sum) + values[3] += int(runoff_count) + + timeseries = {} + for timestamp in sorted(data): + rainfall_sum, rainfall_count, runoff_sum, runoff_count = data[timestamp] + values = {} + if rainfall_count: + values["Rainfall"] = rainfall_sum / rainfall_count + if runoff_count: + values["Runoff"] = runoff_sum / runoff_count + if values: + timeseries[timestamp] = values + return timeseries, path + + +def _read_pan_india_series_from_dirs(series_dirs, watershed_id): + combined = {} + missing_paths = [] + for series_dir in series_dirs: + timeseries, path = _read_pan_india_series(series_dir, watershed_id) + if path.exists(): + combined.update(timeseries) + else: + missing_paths.append(path) + return combined, missing_paths + + +def _resolve_duplicate_pan_india_ids( + *, + matches, + target_match_index, + target_offset, + target_count, + unresolved_uids, + series_dirs, +): + resolved = {} + remaining = set(unresolved_uids) + source_offset = target_offset + target_count + + for source_path, row in matches[target_match_index + 1 :]: + feature_count = int(row.get("feature_count") or 0) + uid_frame = gpd.read_file( + source_path, + columns=["uid"], + ignore_geometry=True, + ) + for position, uid in enumerate(uid_frame["uid"].astype(str)): + if uid not in remaining: + continue + watershed_id = source_offset + position + 1 + if _pan_india_series_paths_exist(series_dirs, watershed_id): + resolved[uid] = watershed_id + remaining.remove(uid) + if not remaining: + break + source_offset += feature_count + + return resolved + + +def _attach_pan_india_timeseries( + watersheds_gdf, + *, + state, + district, + block, + series_dirs, + watershed_root=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + watershed_ids=None, + uid_to_watershed_id=None, +): + if "uid" not in watersheds_gdf.columns: + raise ValueError("Precomputed watershed vector must contain uid") + + series_dirs = [Path(series_dir) for series_dir in series_dirs] + if not series_dirs: + raise ValueError( + "At least one Pan-India rainfall/runoff series folder is required" + ) + + year_gdf = watersheds_gdf.copy() + year_gdf["uid"] = year_gdf["uid"].astype(str) + if year_gdf["uid"].duplicated().any(): + raise ValueError("Duplicate uid values found in precomputed watersheds") + + if uid_to_watershed_id is not None: + watershed_source = "Pan-India watershed manifest" + watershed_ids = [uid_to_watershed_id.get(uid) for uid in year_gdf["uid"]] + else: + ( + matches, + target_match_index, + offset, + expected_count, + watershed_source, + ) = _pan_india_watershed_offset( + state, + district, + block, + watershed_root=watershed_root, + ) + if len(year_gdf) != expected_count: + raise ValueError( + "Precomputed watershed count differs from the Pan-India runoff " + f"manifest for {watershed_source}: vector={len(year_gdf)}, " + f"manifest={expected_count}" + ) + + if watershed_ids is None: + watershed_ids = [offset + position + 1 for position in range(len(year_gdf))] + unresolved_positions = [ + position + for position, watershed_id in enumerate(watershed_ids) + if not _pan_india_series_paths_exist(series_dirs, watershed_id) + ] + if unresolved_positions: + unresolved_uids = { + year_gdf.iloc[position]["uid"] for position in unresolved_positions + } + duplicate_ids = _resolve_duplicate_pan_india_ids( + matches=matches, + target_match_index=target_match_index, + target_offset=offset, + target_count=expected_count, + unresolved_uids=unresolved_uids, + series_dirs=series_dirs, + ) + for position in unresolved_positions: + uid = year_gdf.iloc[position]["uid"] + if uid in duplicate_ids: + watershed_ids[position] = duplicate_ids[uid] + elif len(watershed_ids) != len(year_gdf): + raise ValueError( + "Pan-India watershed ID count differs from precomputed " + f"watersheds: ids={len(watershed_ids)}, " + f"watersheds={len(year_gdf)}" + ) + + timeseries_values = [] + missing_paths = [] + for uid, watershed_id in zip(year_gdf["uid"], watershed_ids): + if watershed_id is None: + timeseries_values.append({}) + missing_paths.append(f"uid={uid} has no Pan-India runoff series") + continue + timeseries, missing_for_uid = _read_pan_india_series_from_dirs( + series_dirs, + watershed_id, + ) + timeseries_values.append(timeseries) + missing_paths.extend(str(path) for path in missing_for_uid) + + year_gdf["timeseries"] = timeseries_values + return ( + year_gdf, + watershed_source, + missing_paths, + watershed_ids, + ) + + +def _available_pan_india_series_ids(series_dirs): + available_ids = None + for series_dir in series_dirs: + current_ids = { + int(path.stem) + for path in Path(series_dir).glob("*/*.csv") + if path.stem.isdigit() + } + available_ids = ( + current_ids if available_ids is None else available_ids & current_ids + ) + return available_ids or set() + + +def _build_pan_india_uid_index( + series_dirs, + watershed_root=PRECOMPUTED_TEHSIL_WATERSHED_DIR, +): + available_ids = _available_pan_india_series_ids(series_dirs) + if not available_ids: + raise FileNotFoundError( + "No common Pan-India rainfall/runoff CSV files were found for " + f"the requested years: {[str(path) for path in series_dirs]}" + ) + + uid_to_watershed_id = {} + offset = 0 + matches = find_pan_india_watersheds(Path(watershed_root)) + for source_path, row in matches: + feature_count = int(row.get("feature_count") or 0) + uid_frame = gpd.read_file( + source_path, + columns=["uid"], + ignore_geometry=True, + ) + if len(uid_frame) != feature_count: + raise ValueError( + "Watershed count differs from the Pan-India manifest for " + f"{source_path}: vector={len(uid_frame)}, " + f"manifest={feature_count}" + ) + for position, uid in enumerate(uid_frame["uid"].astype(str)): + watershed_id = offset + position + 1 + if watershed_id in available_ids: + uid_to_watershed_id[uid] = watershed_id + offset += feature_count + + print( + "Built Pan-India rainfall/runoff index for " + f"{len(uid_to_watershed_id)} unique watersheds" + ) + return matches, uid_to_watershed_id + + +def _read_complete_uid_cache(path, required_uids, value_columns): + path = Path(path) + if not path.exists(): + return None + + frame = pd.read_parquet(path) + required_columns = {CACHE_UID_COLUMN, *value_columns} + missing_columns = required_columns - set(frame.columns) + if missing_columns: + print( + f"Ignoring incomplete cache {path}; missing columns: " + f"{sorted(missing_columns)}" + ) + return None + + frame[CACHE_UID_COLUMN] = frame[CACHE_UID_COLUMN].astype(str) + if frame[CACHE_UID_COLUMN].duplicated().any(): + print(f"Ignoring cache with duplicate UIDs: {path}") + return None + + missing_uids = set(required_uids) - set(frame[CACHE_UID_COLUMN]) + if missing_uids: + print( + f"Ignoring incomplete cache {path}; " + f"missing {len(missing_uids)} UIDs" + ) + return None + return frame + + +def _iter_unique_watershed_partitions( + matches, + *, + allowed_uids=None, + area_limit=None, +): + seen_uids = set() + selected_matches = matches if area_limit is None else matches[: int(area_limit)] + + for source_path, row in selected_matches: + area_gdf = read_validated_vector_file( + source_path, + f"Watershed partition has no valid geometries: {source_path}", + ) + if CACHE_UID_COLUMN not in area_gdf.columns: + raise ValueError(f"Watershed partition must contain uid: {source_path}") + + area_gdf[CACHE_UID_COLUMN] = area_gdf[CACHE_UID_COLUMN].astype(str) + keep = ~area_gdf[CACHE_UID_COLUMN].isin(seen_uids) + if allowed_uids is not None: + keep &= area_gdf[CACHE_UID_COLUMN].isin(allowed_uids) + area_gdf = area_gdf.loc[keep].copy() + seen_uids.update(area_gdf[CACHE_UID_COLUMN]) + if not area_gdf.empty: + yield source_path, row, area_gdf + + +def _ensure_pan_india_aquifer_cache( + *, + matches, + required_uids, + output_base_dir, + aquifer_vector_path, + area_limit=None, +): + cache_path = _aquifer_cache_path(output_base_dir) + value_column = "weighted_avg_yeild" + cached = _read_complete_uid_cache( + cache_path, + required_uids, + [value_column], + ) + if cached is not None: + print(f"Using Pan-India aquifer cache: {cache_path}") + return cached.set_index(CACHE_UID_COLUMN), cache_path + + aquifers_gdf = read_validated_vector_file( + aquifer_vector_path, + f"Aquifer source file has no valid geometries: " f"{aquifer_vector_path}", + ) + aquifers_projected = _prepare_aquifers_for_intersection(aquifers_gdf) + records = [] + processed = 0 + for source_path, _, watersheds_gdf in _iter_unique_watershed_partitions( + matches, + allowed_uids=set(required_uids), + area_limit=area_limit, + ): + result = _compute_aquifer_properties_for_watersheds( + watersheds_gdf=watersheds_gdf[[CACHE_UID_COLUMN, "geometry"]].copy(), + aquifers_projected=aquifers_projected, + ) + records.append( + pd.DataFrame( + { + CACHE_UID_COLUMN: result[CACHE_UID_COLUMN].astype(str), + value_column: result["total_weighted_yield"].astype(float), + } + ) + ) + processed += len(result) + print( + f"Cached aquifer yield for {processed} watersheds " + f"(latest partition: {source_path})" + ) + + if not records: + raise ValueError("No watershed records were available for aquifer caching") + + cache_frame = pd.concat(records, ignore_index=True) + cache_frame = cache_frame.drop_duplicates( + CACHE_UID_COLUMN, + keep="last", + ).sort_values(CACHE_UID_COLUMN) + if area_limit is None: + missing_uids = set(required_uids) - set(cache_frame[CACHE_UID_COLUMN]) + if missing_uids: + raise ValueError( + "Aquifer cache generation missed " + f"{len(missing_uids)} required watershed UIDs" + ) + _write_parquet_atomic(cache_frame, cache_path) + print( + f"Saved Pan-India aquifer cache with {len(cache_frame)} records: " + f"{cache_path}" + ) + return cache_frame.set_index(CACHE_UID_COLUMN), cache_path + + +def _decode_timeseries(value, uid): + if isinstance(value, dict): + return value + if value is None or (isinstance(value, float) and np.isnan(value)): + raise ValueError(f"Missing rainfall/runoff timeseries for uid={uid}") + try: + return json.loads(value) + except (TypeError, json.JSONDecodeError) as error: + raise ValueError( + f"Invalid rainfall/runoff timeseries JSON for uid={uid}" + ) from error + + +def _parse_timestamp(value): + return dt.datetime.fromisoformat(str(value).replace("Z", "+00:00")).date() + + +def _aggregate_rainfall_runoff(timeseries, periods): + totals = {key: {"Precipitation": 0.0, "RunOff": 0.0} for _, _, key in periods} + for timestamp, values in timeseries.items(): + if not isinstance(values, dict): + continue + day = _parse_timestamp(timestamp) + for period_start, period_end, key in periods: + if period_start <= day < period_end: + totals[key]["Precipitation"] += float( + values.get("Rainfall", values.get("rainfall", 0.0)) or 0.0 + ) + totals[key]["RunOff"] += float( + values.get("Runoff", values.get("runoff", 0.0)) or 0.0 + ) + break + return totals + + +class _RasterZonalGrid: + def __init__(self, watersheds_gdf, reference_path): + with rasterio.open(reference_path) as src: + if src.crs is None: + raise ValueError(f"Raster CRS is missing: {reference_path}") + working_gdf = watersheds_gdf.to_crs(src.crs) + minx, miny, maxx, maxy = working_gdf.total_bounds + requested = from_bounds(minx, miny, maxx, maxy, src.transform) + col_start = max(0, math.floor(requested.col_off)) + row_start = max(0, math.floor(requested.row_off)) + col_stop = min( + src.width, + math.ceil(requested.col_off + requested.width), + ) + row_stop = min( + src.height, + math.ceil(requested.row_off + requested.height), + ) + self.window = Window( + col_start, + row_start, + col_stop - col_start, + row_stop - row_start, + ) + if self.window.width <= 0 or self.window.height <= 0: + raise ValueError( + f"Watersheds do not overlap ET raster: {reference_path}" + ) + self.transform = src.window_transform(self.window) + self.shape = (int(self.window.height), int(self.window.width)) + self.crs = src.crs + self.reference_transform = src.transform + self.reference_width = src.width + self.reference_height = src.height + self.geometries = list(working_gdf.geometry) + self._geometry_masks = None + + def read_flux(self, raster_path, negative_as_nodata): + with rasterio.open(raster_path) as src: + if ( + src.crs != self.crs + or src.width != self.reference_width + or src.height != self.reference_height + or not src.transform.almost_equals(self.reference_transform) + ): + raise ValueError( + f"ET raster grid does not match reference raster: {raster_path}" + ) + + data = src.read(1, window=self.window, masked=True) + values = np.asarray(data.filled(np.nan), dtype=np.float64) + valid = ~np.ma.getmaskarray(data) & np.isfinite(values) + if src.nodata is not None: + valid &= values != src.nodata + if negative_as_nodata: + valid &= values >= 0 + else: + values = np.where(values > 0, values, 0.0) + return values, valid + + def means(self, values, valid): + if self._geometry_masks is None: + self._geometry_masks = [] + for geometry in self.geometries: + center_mask = geometry_mask( + [mapping(geometry)], + out_shape=self.shape, + transform=self.transform, + invert=True, + all_touched=False, + ) + touched_mask = geometry_mask( + [mapping(geometry)], + out_shape=self.shape, + transform=self.transform, + invert=True, + all_touched=True, + ) + self._geometry_masks.append((center_mask, touched_mask)) + + means = [] + for center_mask, touched_mask in self._geometry_masks: + selected = center_mask & valid + if not selected.any(): + selected = touched_mask & valid + means.append(float(values[selected].mean()) if selected.any() else np.nan) + return np.asarray(means, dtype=np.float64) + + +def _aggregate_flux_rasters( + watersheds_gdf, + weighted_rasters, + *, + negative_as_nodata, +): + if not weighted_rasters: + return np.zeros(len(watersheds_gdf), dtype=np.float64) + + grid = _RasterZonalGrid(watersheds_gdf, weighted_rasters[0][0]) + total = np.zeros(grid.shape, dtype=np.float64) + has_value = np.zeros(grid.shape, dtype=bool) + + for raster_path, day_count in weighted_rasters: + values, valid = grid.read_flux( + raster_path, + negative_as_nodata=negative_as_nodata, + ) + total[valid] += values[valid] * SECONDS_PER_DAY * float(day_count) + has_value |= valid + + return grid.means(total, has_value) + + +def _folder_has_tifs(path): + path = Path(path) + return path.is_dir() and next(path.glob("*.tif"), None) is not None + + +def _preferred_et_source_root(et_root, patch_filled_source, raw_source, raw_subfolder): + et_root = Path(et_root) + patch_filled_root = et_root / patch_filled_source + if _folder_has_tifs(patch_filled_root): + return patch_filled_root + return et_root / raw_source / raw_subfolder + + +def _et_source_name_from_root(root): + root = Path(root) + if root.name in { + FLDAS_CA_DAILY_PATCH_FILLED_SOURCE, + FLDAS_GLOBAL_MONTHLY_PATCH_FILLED_SOURCE, + }: + return root.name + if root.name in {"daily", "monthly"}: + return root.parent.name + return root.name + + +def _et_source_names_from_roots(roots_by_year): + return sorted( + { + _et_source_name_from_root(root) + for root in roots_by_year.values() + } + ) + + +def _aggregate_source_folder(source_names): + source_names = sorted(set(source_names)) + if len(source_names) == 1: + return source_names[0] + return "__".join(source_names) + + +def _cache_et_sources(frame): + sources = set() + if CACHE_ET_SOURCE_COLUMN not in frame.columns: + return sources + for value in frame[CACHE_ET_SOURCE_COLUMN].dropna().astype(str): + sources.update( + source.strip() + for source in value.split(",") + if source.strip() + ) + return sources + + +def _et_source_signature(et_roots_by_year): + daily_roots = _daily_roots_by_year(et_roots_by_year) + monthly_roots = _monthly_roots_by_year(et_roots_by_year) + return "|".join( + ( + f"{_year_key(source_year)}:" + f"daily={_et_source_name_from_root(daily_roots[source_year])};" + f"monthly={_et_source_name_from_root(monthly_roots[source_year])}" + ) + for source_year in sorted(et_roots_by_year) + ) + + +def _daily_roots_by_year(et_roots_by_year): + return { + source_year: _preferred_et_source_root( + et_root, + FLDAS_CA_DAILY_PATCH_FILLED_SOURCE, + FLDAS_CA_DAILY_SOURCE, + "daily", + ) + for source_year, et_root in et_roots_by_year.items() + } + + +def _monthly_roots_by_year(et_roots_by_year): + return { + source_year: _preferred_et_source_root( + et_root, + FLDAS_GLOBAL_MONTHLY_PATCH_FILLED_SOURCE, + FLDAS_GLOBAL_MONTHLY_SOURCE, + "monthly", + ) + for source_year, et_root in et_roots_by_year.items() + } + + +def _source_root_for_date(roots_by_year, day, source_name): + source_year = _source_year_for_date(day) + root = roots_by_year.get(source_year) + if root is None: + raise FileNotFoundError( + f"{source_name} folder for source year {_year_key(source_year)} " + f"is required for {day.isoformat()}" + ) + return root + + +def _month_path(monthly_roots_by_year, day): + monthly_root = _source_root_for_date( + monthly_roots_by_year, + day, + "Monthly global ET", + ) + path = monthly_root / f"{day:%Y%m}.tif" + if not path.exists(): + raise FileNotFoundError(f"Monthly global ET raster not found: {path}") + return path + + +def _monthly_weights( + monthly_roots_by_year, + start_date, + end_date, + included_dates=None, +): + if included_dates is None: + dates = [] + current = start_date + while current < end_date: + dates.append(current) + current += dt.timedelta(days=1) + else: + dates = list(included_dates) + + month_counts = Counter(day.replace(day=1) for day in dates) + return [ + (_month_path(monthly_roots_by_year, month), day_count) + for month, day_count in sorted(month_counts.items()) + ] + + +def _daily_rasters(daily_roots_by_year, start_date, end_date): + available = [] + missing = [] + current = start_date + while current < end_date: + daily_root = _source_root_for_date( + daily_roots_by_year, + current, + "Daily CA ET", + ) + path = daily_root / f"{current:%Y%m%d}.tif" + if path.exists(): + available.append((path, 1)) + else: + missing.append(current) + current += dt.timedelta(days=1) + return available, missing + + +def _uses_daily_et(watersheds_gdf, daily_roots_by_year): + reference_path = None + for daily_root in daily_roots_by_year.values(): + reference_path = next(iter(sorted(daily_root.glob("*.tif"))), None) + if reference_path is not None: + break + if reference_path is None: + return False + + with rasterio.open(reference_path) as src: + raster_bounds = box(*src.bounds) + watersheds = watersheds_gdf.to_crs(src.crs) + return raster_bounds.covers(watersheds.geometry.union_all()) + + +def _calculate_period_et(watersheds_gdf, periods, et_roots_by_year): + daily_roots = _daily_roots_by_year(et_roots_by_year) + monthly_roots = _monthly_roots_by_year(et_roots_by_year) + daily_source = ",".join(_et_source_names_from_roots(daily_roots)) + monthly_source = ",".join(_et_source_names_from_roots(monthly_roots)) + use_daily = _uses_daily_et(watersheds_gdf, daily_roots) + result = {} + + for period_start, period_end, key in periods: + if use_daily: + daily_rasters, missing_dates = _daily_rasters( + daily_roots, + period_start, + period_end, + ) + values = _aggregate_flux_rasters( + watersheds_gdf, + daily_rasters, + negative_as_nodata=True, + ) + if missing_dates: + fallback = _aggregate_flux_rasters( + watersheds_gdf, + _monthly_weights( + monthly_roots, + period_start, + period_end, + included_dates=missing_dates, + ), + negative_as_nodata=False, + ) + values = values + fallback + else: + values = _aggregate_flux_rasters( + watersheds_gdf, + _monthly_weights(monthly_roots, period_start, period_end), + negative_as_nodata=False, + ) + + if np.isnan(values).any(): + missing_count = int(np.isnan(values).sum()) + raise ValueError( + f"ET could not be calculated for {missing_count} watersheds " + f"during period {key}" + ) + result[key] = values + + return result, daily_source if use_daily else monthly_source + + +def _write_integrated_flux_raster( + weighted_rasters, + *, + negative_as_nodata, + output_path, +): + output_path = Path(output_path) + if output_path.exists(): + return output_path + if not weighted_rasters: + return None + + output_path.parent.mkdir(parents=True, exist_ok=True) + reference_path = weighted_rasters[0][0] + with rasterio.open(reference_path) as reference: + profile = reference.profile.copy() + shape = (reference.height, reference.width) + reference_crs = reference.crs + reference_transform = reference.transform + + total = np.zeros(shape, dtype=np.float64) + has_value = np.zeros(shape, dtype=bool) + for raster_path, day_count in weighted_rasters: + with rasterio.open(raster_path) as src: + if ( + src.crs != reference_crs + or src.width != shape[1] + or src.height != shape[0] + or not src.transform.almost_equals(reference_transform) + ): + raise ValueError( + "ET raster grid does not match aggregate reference: " + f"{raster_path}" + ) + data = src.read(1, masked=True) + values = np.asarray(data.filled(np.nan), dtype=np.float64) + valid = ~np.ma.getmaskarray(data) & np.isfinite(values) + if src.nodata is not None: + valid &= values != src.nodata + if negative_as_nodata: + valid &= values >= 0 + else: + values = np.where(values > 0, values, 0.0) + total[valid] += values[valid] * SECONDS_PER_DAY * float(day_count) + has_value |= valid + + nodata = -9999.0 + output = np.where(has_value, total, nodata).astype(np.float32) + profile.update( + count=1, + dtype="float32", + nodata=nodata, + compress="deflate", + predictor=3, + ) + temporary_path = output_path.with_name( + f"{output_path.stem}.{os.getpid()}.tmp{output_path.suffix}" + ) + with rasterio.open(temporary_path, "w", **profile) as dst: + dst.write(output, 1) + temporary_path.replace(output_path) + return output_path + + +def _ensure_period_et_aggregate_rasters( + *, + periods, + et_roots_by_year, + aggregate_root, +): + daily_roots = _daily_roots_by_year(et_roots_by_year) + monthly_roots = _monthly_roots_by_year(et_roots_by_year) + daily_source_names = _et_source_names_from_roots(daily_roots) + monthly_source_names = _et_source_names_from_roots(monthly_roots) + daily_source = ",".join(daily_source_names) + monthly_source = ",".join(monthly_source_names) + daily_aggregate_folder = _aggregate_source_folder(daily_source_names) + monthly_aggregate_folder = _aggregate_source_folder(monthly_source_names) + aggregate_root = Path(aggregate_root) + result = {} + + for period_start, period_end, key in periods: + file_key = key.replace("-", "") + daily_rasters, missing_dates = _daily_rasters( + daily_roots, + period_start, + period_end, + ) + daily_path = _write_integrated_flux_raster( + daily_rasters, + negative_as_nodata=True, + output_path=aggregate_root / daily_aggregate_folder / f"{file_key}.tif", + ) + global_path = _write_integrated_flux_raster( + _monthly_weights( + monthly_roots, + period_start, + period_end, + ), + negative_as_nodata=False, + output_path=aggregate_root / monthly_aggregate_folder / f"{file_key}.tif", + ) + missing_path = None + if missing_dates: + missing_path = _write_integrated_flux_raster( + _monthly_weights( + monthly_roots, + period_start, + period_end, + included_dates=missing_dates, + ), + negative_as_nodata=False, + output_path=aggregate_root + / f"{monthly_aggregate_folder}_missing_daily" + / f"{file_key}.tif", + ) + + result[key] = { + "daily": daily_path, + "global": global_path, + "missing_daily": missing_path, + "daily_source": daily_source, + "global_source": monthly_source, + } + return result + + +def _integrated_raster_means(watersheds_gdf, raster_path, grid=None): + grid = grid or _RasterZonalGrid(watersheds_gdf, raster_path) + values, valid = grid.read_flux( + raster_path, + negative_as_nodata=True, + ) + return grid.means(values, valid) + + +def _raster_covers_watersheds(watersheds_gdf, raster_path): + if raster_path is None: + return False + with rasterio.open(raster_path) as src: + raster_bounds = box(*src.bounds) + watersheds = watersheds_gdf.to_crs(src.crs) + return raster_bounds.covers(watersheds.geometry.union_all()) + + +def _calculate_period_et_from_aggregates( + watersheds_gdf, + periods, + aggregate_rasters, +): + daily_reference = next( + ( + aggregate_rasters[key]["daily"] + for _, _, key in periods + if aggregate_rasters[key]["daily"] is not None + ), + None, + ) + use_daily = _raster_covers_watersheds( + watersheds_gdf, + daily_reference, + ) + first_key = periods[0][2] + primary_reference = ( + daily_reference if use_daily else aggregate_rasters[first_key]["global"] + ) + primary_grid = _RasterZonalGrid(watersheds_gdf, primary_reference) + missing_daily_grid = None + result = {} + + for _, _, key in periods: + paths = aggregate_rasters[key] + if use_daily: + if paths["daily"] is None: + values = np.zeros(len(watersheds_gdf), dtype=np.float64) + else: + values = _integrated_raster_means( + watersheds_gdf, + paths["daily"], + grid=primary_grid, + ) + if paths["missing_daily"] is not None: + if missing_daily_grid is None: + missing_daily_grid = _RasterZonalGrid( + watersheds_gdf, + paths["missing_daily"], + ) + values += _integrated_raster_means( + watersheds_gdf, + paths["missing_daily"], + grid=missing_daily_grid, + ) + else: + values = _integrated_raster_means( + watersheds_gdf, + paths["global"], + grid=primary_grid, + ) + + if np.isnan(values).any(): + missing_count = int(np.isnan(values).sum()) + raise ValueError( + f"Cached ET could not be calculated for {missing_count} " + f"watersheds during period {key}" + ) + result[key] = values + + source = ( + aggregate_rasters[first_key]["daily_source"] + if use_daily + else aggregate_rasters[first_key]["global_source"] + ) + return result, source + + +def _ensure_pan_india_et_cache( + *, + matches, + required_uids, + year, + is_annual, + et_roots_by_year, + output_base_dir, + area_limit=None, +): + periods = _build_periods(year, is_annual) + period_columns = [key for _, _, key in periods] + source_year_keys = ",".join( + _year_key(source_year) for source_year in sorted(et_roots_by_year) + ) + cache_path = _et_cache_path(output_base_dir, year, is_annual) + source_signature = _et_source_signature(et_roots_by_year) + expected_cache_sources = set( + _et_source_names_from_roots(_daily_roots_by_year(et_roots_by_year)) + ) | set(_et_source_names_from_roots(_monthly_roots_by_year(et_roots_by_year))) + cached = _read_complete_uid_cache( + cache_path, + required_uids, + [ + *period_columns, + CACHE_ET_SOURCE_COLUMN, + CACHE_ET_SOURCE_SIGNATURE_COLUMN, + SOURCE_YEARS_COLUMN, + ], + ) + if cached is not None: + cached_sources = _cache_et_sources(cached) + cached_signatures = set( + cached[CACHE_ET_SOURCE_SIGNATURE_COLUMN].dropna().astype(str) + ) + if cached_signatures != {source_signature}: + print( + f"Ignoring ET cache with stale source signature: {cache_path}; " + f"cache={sorted(cached_signatures)}, " + f"current={source_signature}" + ) + elif cached_sources and not cached_sources.issubset(expected_cache_sources): + print( + f"Ignoring ET cache with stale source rasters: {cache_path}; " + f"cache={sorted(cached_sources)}, " + f"current={sorted(expected_cache_sources)}" + ) + else: + print(f"Using Pan-India ET cache: {cache_path}") + return cached.set_index(CACHE_UID_COLUMN), cache_path + + aggregate_rasters = _ensure_period_et_aggregate_rasters( + periods=periods, + et_roots_by_year=et_roots_by_year, + aggregate_root=_et_aggregate_root( + output_base_dir, + year, + is_annual, + ), + ) + records = [] + processed = 0 + for source_path, _, watersheds_gdf in _iter_unique_watershed_partitions( + matches, + allowed_uids=set(required_uids), + area_limit=area_limit, + ): + record = pd.DataFrame( + { + CACHE_UID_COLUMN: watersheds_gdf[CACHE_UID_COLUMN].astype(str), + SOURCE_YEARS_COLUMN: source_year_keys, + CACHE_ET_SOURCE_SIGNATURE_COLUMN: source_signature, + } + ) + try: + et_by_period, et_source = _calculate_period_et_from_aggregates( + watersheds_gdf, + periods, + aggregate_rasters, + ) + record[CACHE_ET_SOURCE_COLUMN] = et_source + record[CACHE_ET_ERROR_COLUMN] = None + for key in period_columns: + record[key] = et_by_period[key] + except Exception as error: + record[CACHE_ET_SOURCE_COLUMN] = None + record[CACHE_ET_ERROR_COLUMN] = str(error) + for key in period_columns: + record[key] = np.nan + print( + f"ET cache unavailable for {len(record)} watersheds in " + f"{source_path}: {error}" + ) + records.append(record) + processed += len(record) + print( + f"Cached {len(period_columns)} ET period(s) for " + f"{processed} watersheds (latest partition: {source_path})" + ) + + if not records: + raise ValueError("No watershed records were available for ET caching") + + cache_frame = pd.concat(records, ignore_index=True) + cache_frame = cache_frame.drop_duplicates( + CACHE_UID_COLUMN, + keep="last", + ).sort_values(CACHE_UID_COLUMN) + if area_limit is None: + missing_uids = set(required_uids) - set(cache_frame[CACHE_UID_COLUMN]) + if missing_uids: + raise ValueError( + "ET cache generation missed " + f"{len(missing_uids)} required watershed UIDs" + ) + _write_parquet_atomic(cache_frame, cache_path) + print(f"Saved Pan-India ET cache with {len(cache_frame)} records: " f"{cache_path}") + return cache_frame.set_index(CACHE_UID_COLUMN), cache_path + + +def _add_annual_well_depth( + result_gdf, + annual_columns, + aquifer_vector_path, + aquifers_gdf=None, + aquifer_cache=None, +): + if aquifer_cache is not None: + uid_values = result_gdf[CACHE_UID_COLUMN].astype(str) + missing_uids = set(uid_values) - set(aquifer_cache.index) + if missing_uids: + raise ValueError( + "Aquifer cache is missing " f"{len(missing_uids)} watershed UIDs" + ) + result_gdf["weighted_avg_yeild"] = aquifer_cache.reindex(uid_values)[ + "weighted_avg_yeild" + ].to_numpy(dtype=float) + else: + if aquifers_gdf is None: + aquifers_gdf = read_validated_vector_file( + aquifer_vector_path, + f"Aquifer source file has no valid geometries: " + f"{aquifer_vector_path}", + ) + aquifer_result = _compute_aquifer_properties_for_watersheds( + watersheds_gdf=result_gdf[["uid", "geometry"]].copy(), + aquifers_gdf=aquifers_gdf, + ) + result_gdf["weighted_avg_yeild"] = aquifer_result[ + "total_weighted_yield" + ].to_numpy() + + for index, row in result_gdf.iterrows(): + weighted_yield = row["weighted_avg_yeild"] + for column in annual_columns: + values = json.loads(row[column]) + values["WellDepth"] = ( + values["DeltaG"] / (float(weighted_yield) * 1000.0) + if pd.notna(weighted_yield) and float(weighted_yield) > 0 + else None + ) + result_gdf.at[index, column] = json.dumps( + values, + separators=(",", ":"), + ) + + return _add_annual_net_columns(result_gdf, annual_columns) + + +def _add_annual_net_columns(result_gdf, annual_columns): + for start_index in range(len(annual_columns) - 4): + window = annual_columns[start_index : start_index + 5] + start_year = window[0].split("_")[0] + end_year = window[-1].split("_")[1][-2:] + net_column = f"Net{start_year}_{end_year}" + + def net_value(row): + well_depths = [ + json.loads(row[column]).get("WellDepth") for column in window + ] + if any(value is None for value in well_depths): + return None + return sum(float(value) for value in well_depths) + + result_gdf[net_column] = result_gdf.apply(net_value, axis=1) + + return result_gdf + + +def _run_generate_hydrology_area_local( + *, + state, + district, + block, + start_year, + end_year, + is_annual=False, + gee_account_id=None, + hydrology_output_root=HYDROLOGY_OUTPUT_ROOT, + output_base_dir=HYDROLOGY_LOCAL_OUTPUT_DIR, + aquifer_vector_path=AQUIFER_VECTOR_PATH, + push_to_geoserver=True, + sync_layer_metadata=True, + watersheds_gdf=None, + watershed_source=None, + uid_to_watershed_id=None, + aquifers_gdf=None, + aquifer_cache=None, + et_cache_by_year=None, + et_cache_paths_by_year=None, + layer_name_override=None, + write_output=True, +): + _ = gee_account_id + state = _normalize_location(state, "state") + district = _normalize_location(district, "district") + block = _normalize_location(block, "block") + start_year = int(start_year) + end_year = int(end_year) + if end_year < start_year: + raise ValueError("end_year must be greater than or equal to start_year") + + result_gdf = None + uid_to_index = None + cumulative_g = {} + period_columns = [] + et_sources = set() + input_paths = [] + pan_india_watershed_ids = None + if watersheds_gdf is None: + watersheds_gdf, watershed_source = load_precomputed_watersheds( + state=state, + district=district, + block=block, + ) + else: + watersheds_gdf = watersheds_gdf.copy() + watershed_source = watershed_source or "provided watershed partition" + + for year in range(start_year, end_year + 1): + periods = _build_periods(year, is_annual) + source_inputs = _resolve_source_year_inputs( + output_year=year, + periods=periods, + hydrology_output_root=hydrology_output_root, + ) + source_years = sorted(source_inputs) + series_dirs = [ + source_inputs[source_year][0] for source_year in source_years + ] + et_roots_by_year = { + source_year: source_inputs[source_year][1] + for source_year in source_years + } + ( + year_gdf, + runoff_watershed_source, + missing_series_paths, + pan_india_watershed_ids, + ) = _attach_pan_india_timeseries( + watersheds_gdf, + state=state, + district=district, + block=block, + series_dirs=series_dirs, + watershed_ids=pan_india_watershed_ids, + uid_to_watershed_id=uid_to_watershed_id, + ) + if missing_series_paths: + raise FileNotFoundError( + "Pan-India rainfall/runoff remains unavailable for " + f"{len(missing_series_paths)} of {len(year_gdf)} watersheds " + f"in {_year_key(year)} after checking duplicate watershed " + f"IDs. First missing files: {missing_series_paths[:3]}" + ) + + if result_gdf is None: + result_gdf = year_gdf.drop(columns=["timeseries"]).copy() + uid_to_index = {uid: index for index, uid in result_gdf["uid"].items()} + cumulative_g = {uid: 0.0 for uid in uid_to_index} + elif set(year_gdf["uid"]) != set(uid_to_index): + raise ValueError( + "Watershed uid set differs between yearly inputs: " + f"{[str(path) for path in series_dirs]}" + ) + + if et_cache_by_year is not None and year in et_cache_by_year: + et_cache = et_cache_by_year[year] + uid_values = year_gdf[CACHE_UID_COLUMN].astype(str) + missing_uids = set(uid_values) - set(et_cache.index) + if missing_uids: + raise ValueError( + f"ET cache for {_year_key(year)} is missing " + f"{len(missing_uids)} watershed UIDs" + ) + selected_et = et_cache.reindex(uid_values) + invalid_et = selected_et[ + [key for _, _, key in periods] + ].isna().any(axis=1) + if invalid_et.any(): + errors = [] + if CACHE_ET_ERROR_COLUMN in selected_et.columns: + errors = ( + selected_et.loc[invalid_et, CACHE_ET_ERROR_COLUMN] + .dropna() + .astype(str) + .unique() + .tolist() + ) + error_suffix = f" First cache errors: {errors[:3]}" if errors else "" + raise ValueError( + f"ET cache for {_year_key(year)} has no usable value for " + f"{int(invalid_et.sum())} watershed UIDs.{error_suffix}" + ) + et_by_period = { + key: selected_et[key].to_numpy(dtype=float) for _, _, key in periods + } + selected_sources = set( + selected_et[CACHE_ET_SOURCE_COLUMN].dropna().astype(str) + ) + et_sources.update(selected_sources) + et_source = ",".join(sorted(selected_sources)) + else: + et_by_period, et_source = _calculate_period_et( + year_gdf, + periods, + et_roots_by_year, + ) + et_sources.add(et_source) + input_paths.append( + { + "year": _year_key(year), + "source_years": [ + _year_key(source_year) for source_year in source_years + ], + "rainfall_runoff": [str(path) for path in series_dirs], + "rainfall_runoff_watersheds": runoff_watershed_source, + "missing_rainfall_runoff_series": len(missing_series_paths), + "et": [ + str(et_roots_by_year[source_year]) + for source_year in source_years + ], + "et_source": et_source, + "et_cache": ( + str(et_cache_paths_by_year[year]) + if et_cache_paths_by_year and year in et_cache_paths_by_year + else None + ), + } + ) + + for source_index, row in enumerate(year_gdf.itertuples(index=False)): + uid = str(row.uid) + target_index = uid_to_index[uid] + timeseries = _decode_timeseries(row.timeseries, uid) + water_balance = _aggregate_rainfall_runoff(timeseries, periods) + + for _, _, key in periods: + values = water_balance[key] + values["ET"] = float(et_by_period[key][source_index]) + values["DeltaG"] = ( + values["Precipitation"] - values["RunOff"] - values["ET"] + ) + cumulative_g[uid] += values["DeltaG"] + values["G"] = cumulative_g[uid] + result_gdf.at[target_index, key] = json.dumps( + values, + separators=(",", ":"), + ) + + period_columns.extend(key for _, _, key in periods) + + if is_annual: + result_gdf = _add_annual_well_depth( + result_gdf, + annual_columns=period_columns, + aquifer_vector_path=aquifer_vector_path, + aquifers_gdf=aquifers_gdf, + aquifer_cache=aquifer_cache, + ) + + if not write_output: + return { + "gdf": result_gdf, + "is_annual": bool(is_annual), + "start_year": start_year, + "end_year": end_year, + "period_columns": period_columns, + "period_count": len(period_columns), + "watershed_count": len(result_gdf), + "et_sources": sorted(et_sources), + "inputs": input_paths, + } + + layer_name = layer_name_override or _layer_name( + district, + block, + is_annual, + ) + output_path = build_output_vector_path( + layer_name=layer_name, + state=state, + district=district, + block=block, + output_base_dir=output_base_dir, + block_fallback="unknown_block", + ) + if output_path.exists(): + output_path.unlink() + asset_id = write_vector_output( + gdf=result_gdf, + output_path=output_path, + layer_name=layer_name, + ) + print(f"Saved local hydrology vector: {asset_id}") + + geoserver_synced = False + if push_to_geoserver: + response = push_local_vector_to_geoserver( + path=os.path.splitext(asset_id)[0], + layer_name=layer_name, + workspace=GEOSERVER_WORKSPACE, + file_type="gpkg", + ) + print(f"GeoServer response: {response}") + geoserver_synced = isinstance(response, dict) and response.get( + "status_code" + ) in (200, 201, 202) + + layer_id = None + if sync_layer_metadata: + layer_id = save_layer_info_to_db( + state=state, + district=district, + block=block, + layer_name=layer_name, + asset_id=asset_id, + dataset_name="Hydrology", + algorithm=LOCAL_ALGORITHM, + algorithm_version=LOCAL_ALGORITHM_VERSION, + misc={ + "start_date": f"{start_year}-07-01", + "end_date": f"{end_year + 1}-06-30", + "is_annual": bool(is_annual), + "is_generated_locally": True, + "et_sources": sorted(et_sources), + "inputs": input_paths, + }, + ) + if layer_id and geoserver_synced: + update_layer_sync_status( + layer_id=layer_id, + sync_to_geoserver=True, + ) + + return { + "output": asset_id, + "layer_name": layer_name, + "is_annual": bool(is_annual), + "start_year": start_year, + "end_year": end_year, + "period_count": len(period_columns), + "watershed_count": len(result_gdf), + "et_sources": sorted(et_sources), + "geoserver_synced": geoserver_synced, + "layer_id": layer_id, + } + + +def _pan_india_run_root(output_base_dir, start_year, end_year, is_annual): + period = "annual" if is_annual else "fortnight" + return Path(output_base_dir) / "pan_india" / f"{start_year}_{end_year + 1}" / period + + +def _pan_india_area_output_path( + *, + output_base_dir, + state, + district, + block, + is_annual, +): + layer_name = _layer_name( + district, + block, + is_annual, + ) + return build_output_vector_path( + layer_name=layer_name, + state=state, + district=district, + block=block, + output_base_dir=output_base_dir, + block_fallback="unknown_block", + ) + + +def _run_generate_hydrology_pan_india_local( + *, + start_year, + end_year, + is_annual, + hydrology_output_root, + output_base_dir, + aquifer_vector_path, + push_to_geoserver, + sync_layer_metadata, + overwrite=False, + area_limit=None, +): + series_dirs = [] + et_roots_by_output_year = {} + for year in range(start_year, end_year + 1): + periods = _build_periods(year, is_annual) + source_inputs = _resolve_source_year_inputs( + output_year=year, + periods=periods, + hydrology_output_root=hydrology_output_root, + ) + source_years = sorted(source_inputs) + series_dirs.extend( + source_inputs[source_year][0] for source_year in source_years + ) + et_roots_by_output_year[year] = { + source_year: source_inputs[source_year][1] + for source_year in source_years + } + series_dirs = list(dict.fromkeys(series_dirs)) + + matches, uid_to_watershed_id = _build_pan_india_uid_index( + series_dirs=series_dirs, + ) + required_uids = set(uid_to_watershed_id) + aquifer_cache = None + aquifer_cache_path = None + if is_annual: + aquifer_cache, aquifer_cache_path = _ensure_pan_india_aquifer_cache( + matches=matches, + required_uids=required_uids, + output_base_dir=output_base_dir, + aquifer_vector_path=aquifer_vector_path, + area_limit=area_limit, + ) + + et_cache_by_year = {} + et_cache_paths_by_year = {} + for year in range(start_year, end_year + 1): + et_cache, et_cache_path = _ensure_pan_india_et_cache( + matches=matches, + required_uids=required_uids, + year=year, + is_annual=is_annual, + et_roots_by_year=et_roots_by_output_year[year], + output_base_dir=output_base_dir, + area_limit=area_limit, + ) + et_cache_by_year[year] = et_cache + et_cache_paths_by_year[year] = et_cache_path + + run_root = _pan_india_run_root( + output_base_dir, + start_year, + end_year, + is_annual, + ) + layers_root = run_root / "layers" + manifest_path = run_root / "manifest.csv" + manifest_path.parent.mkdir(parents=True, exist_ok=True) + + fieldnames = [ + "state", + "district", + "block", + "source", + "watershed_count", + "duplicate_count", + "status", + "output", + "layer_name", + "geoserver_synced", + "layer_id", + "error", + ] + seen_uids = set() + written_count = 0 + skipped_count = 0 + failed_count = 0 + watershed_count = 0 + + selected_matches = matches if area_limit is None else matches[: int(area_limit)] + with manifest_path.open("w", newline="") as manifest_file: + writer = csv.DictWriter(manifest_file, fieldnames=fieldnames) + writer.writeheader() + + for source_path, row in selected_matches: + state = str(row.get("state") or "").strip().lower() + district = str(row.get("district") or "").strip().lower() + block = str(row.get("tehsil") or "").strip().lower() + manifest_row = { + "state": state, + "district": district, + "block": block, + "source": str(source_path), + "watershed_count": 0, + "duplicate_count": 0, + "status": "failed", + "output": "", + "layer_name": "", + "geoserver_synced": False, + "layer_id": "", + "error": "", + } + + try: + area_gdf = read_validated_vector_file( + source_path, + f"Watershed partition has no valid geometries: {source_path}", + ) + if "uid" not in area_gdf.columns: + raise ValueError( + f"Watershed partition must contain uid: {source_path}" + ) + area_gdf["uid"] = area_gdf["uid"].astype(str) + duplicate_mask = area_gdf["uid"].isin(seen_uids) + duplicate_count = int(duplicate_mask.sum()) + manifest_row["duplicate_count"] = duplicate_count + manifest_row["watershed_count"] = len(area_gdf) + seen_uids.update(area_gdf["uid"]) + + output_path = _pan_india_area_output_path( + output_base_dir=layers_root, + state=state, + district=district, + block=block, + is_annual=is_annual, + ) + if ( + output_path.exists() + and not overwrite + and not push_to_geoserver + and not sync_layer_metadata + ): + manifest_row["status"] = "skipped_existing" + manifest_row["output"] = str(output_path) + watershed_count += len(area_gdf) + skipped_count += 1 + writer.writerow(manifest_row) + manifest_file.flush() + continue + + layer_name = _layer_name( + district, + block, + is_annual, + ) + result = _run_generate_hydrology_area_local( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + is_annual=is_annual, + hydrology_output_root=hydrology_output_root, + output_base_dir=layers_root, + aquifer_vector_path=aquifer_vector_path, + push_to_geoserver=push_to_geoserver, + sync_layer_metadata=sync_layer_metadata, + watersheds_gdf=area_gdf, + watershed_source=str(source_path), + uid_to_watershed_id=uid_to_watershed_id, + aquifer_cache=aquifer_cache, + et_cache_by_year=et_cache_by_year, + et_cache_paths_by_year=et_cache_paths_by_year, + layer_name_override=layer_name, + ) + if push_to_geoserver and not result["geoserver_synced"]: + raise RuntimeError( + f"GeoServer upload did not succeed for {layer_name}" + ) + if sync_layer_metadata and not result["layer_id"]: + raise RuntimeError( + f"Layer metadata sync did not succeed for {layer_name}" + ) + manifest_row["status"] = "written" + manifest_row["output"] = result["output"] + manifest_row["layer_name"] = result["layer_name"] + manifest_row["geoserver_synced"] = result["geoserver_synced"] + manifest_row["layer_id"] = result["layer_id"] or "" + watershed_count += len(area_gdf) + written_count += 1 + except Exception as error: + manifest_row["error"] = str(error) + failed_count += 1 + + writer.writerow(manifest_row) + manifest_file.flush() + + return { + "scope": "pan_india", + "manifest": str(manifest_path), + "output_root": str(layers_root), + "is_annual": bool(is_annual), + "start_year": start_year, + "end_year": end_year, + "partition_count": len(selected_matches), + "written_count": written_count, + "skipped_count": skipped_count, + "failed_count": failed_count, + "watershed_count": watershed_count, + "runoff_index_size": len(uid_to_watershed_id), + "aquifer_cache": (str(aquifer_cache_path) if aquifer_cache_path else None), + "et_caches": { + _year_key(year): str(path) for year, path in et_cache_paths_by_year.items() + }, + } + + +def _run_generate_hydrology_base_layer_local( + *, + year=None, + start_year=None, + end_year=None, + is_annual, + hydrology_output_root=HYDROLOGY_OUTPUT_ROOT, + output_base_dir=HYDROLOGY_BASE_LAYER_ROOT, + cache_base_dir=HYDROLOGY_LOCAL_OUTPUT_DIR, + aquifer_vector_path=AQUIFER_VECTOR_PATH, + overwrite=False, + area_limit=None, +): + year, hydrology_end_year = _resolve_base_layer_year_bounds( + year=year, + start_year=start_year, + end_year=end_year, + ) + + overwrite = _parse_bool(overwrite) + layer_name = _base_layer_name(year, is_annual) + output_path = _base_layer_path(output_base_dir, year, is_annual) + manifest_path = output_path.with_name(f"{output_path.stem}_manifest.csv") + if output_path.exists() and not overwrite: + return { + "scope": "pan_india_base_layer", + "status": "skipped_existing", + "output": str(output_path), + "manifest": str(manifest_path) if manifest_path.exists() else None, + "layer_name": layer_name, + "start_year": year, + "end_year": hydrology_end_year, + "year_key": _year_key(year), + "is_annual": bool(is_annual), + } + + periods = _build_periods(year, is_annual) + source_inputs = _resolve_source_year_inputs( + output_year=year, + periods=periods, + hydrology_output_root=hydrology_output_root, + ) + source_years = sorted(source_inputs) + series_dirs = [source_inputs[source_year][0] for source_year in source_years] + et_roots_by_year = { + source_year: source_inputs[source_year][1] + for source_year in source_years + } + matches, uid_to_watershed_id = _build_pan_india_uid_index( + series_dirs=series_dirs, + ) + required_uids = set(uid_to_watershed_id) + aquifer_cache = None + aquifer_cache_path = None + if is_annual: + aquifer_cache, aquifer_cache_path = _ensure_pan_india_aquifer_cache( + matches=matches, + required_uids=required_uids, + output_base_dir=cache_base_dir, + aquifer_vector_path=aquifer_vector_path, + area_limit=area_limit, + ) + + et_cache, et_cache_path = _ensure_pan_india_et_cache( + matches=matches, + required_uids=required_uids, + year=year, + is_annual=is_annual, + et_roots_by_year=et_roots_by_year, + output_base_dir=cache_base_dir, + area_limit=area_limit, + ) + + output_path.parent.mkdir(parents=True, exist_ok=True) + fieldnames = [ + "state", + "district", + "block", + "source", + "watershed_count", + "status", + "error", + ] + records = [] + written_count = 0 + failed_count = 0 + watershed_count = 0 + selected_matches = matches if area_limit is None else matches[: int(area_limit)] + + with manifest_path.open("w", newline="") as manifest_file: + writer = csv.DictWriter(manifest_file, fieldnames=fieldnames) + writer.writeheader() + + for source_path, row, area_gdf in _iter_unique_watershed_partitions( + selected_matches, + allowed_uids=required_uids, + ): + state = str(row.get("state") or "unknown_state").strip().lower() + district = str(row.get("district") or "unknown_district").strip().lower() + block = str(row.get("tehsil") or Path(source_path).stem).strip().lower() + manifest_row = { + "state": state, + "district": district, + "block": block, + "source": str(source_path), + "watershed_count": len(area_gdf), + "status": "failed", + "error": "", + } + try: + result = _run_generate_hydrology_area_local( + state=state, + district=district, + block=block, + start_year=year, + end_year=year, + is_annual=is_annual, + hydrology_output_root=hydrology_output_root, + output_base_dir=output_base_dir, + aquifer_vector_path=aquifer_vector_path, + push_to_geoserver=False, + sync_layer_metadata=False, + watersheds_gdf=area_gdf, + watershed_source=str(source_path), + uid_to_watershed_id=uid_to_watershed_id, + aquifer_cache=aquifer_cache, + et_cache_by_year={year: et_cache}, + et_cache_paths_by_year={year: et_cache_path}, + write_output=False, + ) + records.append(result["gdf"]) + watershed_count += len(result["gdf"]) + written_count += 1 + manifest_row["status"] = "written" + except Exception as error: + failed_count += 1 + manifest_row["error"] = str(error) + + writer.writerow(manifest_row) + manifest_file.flush() + + if not records: + raise RuntimeError( + "Hydrology base layer generation produced no watershed records. " + f"See manifest: {manifest_path}" + ) + + combined = pd.concat(records, ignore_index=True) + result_gdf = gpd.GeoDataFrame( + combined, + geometry="geometry", + crs=records[0].crs, + ) + if output_path.exists(): + output_path.unlink() + asset_id = write_vector_output( + gdf=result_gdf, + output_path=output_path, + layer_name=layer_name, + ) + print(f"Saved local hydrology base layer: {asset_id}") + + return { + "scope": "pan_india_base_layer", + "status": "written_with_failures" if failed_count else "written", + "output": asset_id, + "manifest": str(manifest_path), + "layer_name": layer_name, + "start_year": year, + "end_year": hydrology_end_year, + "year_key": _year_key(year), + "is_annual": bool(is_annual), + "period_count": len(_build_periods(year, is_annual)), + "watershed_count": watershed_count, + "partition_count": len(selected_matches), + "written_count": written_count, + "failed_count": failed_count, + "runoff_index_size": len(uid_to_watershed_id), + "aquifer_cache": str(aquifer_cache_path) if aquifer_cache_path else None, + "et_cache": str(et_cache_path), + } + + +def _run_clip_hydrology_area_from_base_layers( + *, + state, + district, + block, + start_year, + end_year, + is_annual, + output_base_dir=HYDROLOGY_LOCAL_OUTPUT_DIR, + base_layer_root=HYDROLOGY_BASE_LAYER_ROOT, + push_to_geoserver=True, + sync_layer_metadata=True, +): + state = _normalize_location(state, "state") + district = _normalize_location(district, "district") + block = _normalize_location(block, "block") + start_year = int(start_year) + end_year = int(end_year) + if start_year != FORTNIGHT_ANCHOR_DATE.year: + raise ValueError( + "Local hydrology clipping requires start_year=2017 because the " + "fortnightly cadence is anchored at 2017-07-01" + ) + if end_year < start_year: + raise ValueError("end_year must be greater than or equal to start_year") + + watersheds_gdf, watershed_source = load_precomputed_watersheds( + state=state, + district=district, + block=block, + ) + if CACHE_UID_COLUMN not in watersheds_gdf.columns: + raise ValueError("Precomputed watershed vector must contain uid") + + result_gdf = watersheds_gdf.copy() + result_gdf[CACHE_UID_COLUMN] = result_gdf[CACHE_UID_COLUMN].astype(str) + if result_gdf[CACHE_UID_COLUMN].duplicated().any(): + raise ValueError("Duplicate uid values found in precomputed watersheds") + + uid_values = result_gdf[CACHE_UID_COLUMN].astype(str) + period_columns = [] + input_paths = [] + + for year in range(start_year, end_year + 1): + base_path = _base_layer_path(base_layer_root, year, is_annual) + if not base_path.exists(): + raise FileNotFoundError( + f"Hydrology base layer not found: {base_path}. " + "Generate it first using the /api/v1/pan-india/ hydrology API." + ) + + periods = _build_periods(year, is_annual) + year_columns = [key for _, _, key in periods] + base_frame = gpd.read_file(base_path, ignore_geometry=True) + if CACHE_UID_COLUMN not in base_frame.columns: + raise ValueError(f"Hydrology base layer must contain uid: {base_path}") + base_frame[CACHE_UID_COLUMN] = base_frame[CACHE_UID_COLUMN].astype(str) + base_frame = base_frame.drop_duplicates(CACHE_UID_COLUMN, keep="last") + base_frame = base_frame.set_index(CACHE_UID_COLUMN) + + copy_columns = list(year_columns) + if ( + is_annual + and "weighted_avg_yeild" in base_frame.columns + and "weighted_avg_yeild" not in result_gdf.columns + ): + copy_columns.append("weighted_avg_yeild") + missing_columns = sorted(set(copy_columns) - set(base_frame.columns)) + if missing_columns: + raise ValueError( + f"Hydrology base layer {base_path} is missing columns: " + f"{missing_columns}" + ) + + missing_uids = sorted(set(uid_values) - set(base_frame.index)) + if missing_uids: + raise FileNotFoundError( + f"Hydrology base layer {base_path} is missing " + f"{len(missing_uids)} watershed UIDs for this tehsil. " + f"First missing UIDs: {missing_uids[:5]}" + ) + + matched = base_frame.reindex(uid_values) + for column in copy_columns: + result_gdf[column] = matched[column].to_numpy() + + period_columns.extend(year_columns) + input_paths.append( + { + "year": _year_key(year), + "path": str(base_path), + "period_columns": year_columns, + } + ) + + if is_annual: + result_gdf = _add_annual_net_columns(result_gdf, period_columns) + + layer_name = _layer_name( + district, + block, + is_annual, + ) + output_path = build_output_vector_path( + layer_name=layer_name, + state=state, + district=district, + block=block, + output_base_dir=output_base_dir, + block_fallback="unknown_block", + ) + if output_path.exists(): + output_path.unlink() + asset_id = write_vector_output( + gdf=result_gdf, + output_path=output_path, + layer_name=layer_name, + ) + print(f"Saved clipped local hydrology vector: {asset_id}") + + geoserver_synced = False + if push_to_geoserver: + response = push_local_vector_to_geoserver( + path=os.path.splitext(asset_id)[0], + layer_name=layer_name, + workspace=GEOSERVER_WORKSPACE, + file_type="gpkg", + ) + print(f"GeoServer response: {response}") + geoserver_synced = isinstance(response, dict) and response.get( + "status_code" + ) in (200, 201, 202) + + layer_id = None + if sync_layer_metadata: + layer_id = save_layer_info_to_db( + state=state, + district=district, + block=block, + layer_name=layer_name, + asset_id=asset_id, + dataset_name="Hydrology", + algorithm=LOCAL_ALGORITHM, + algorithm_version=LOCAL_ALGORITHM_VERSION, + misc={ + "start_date": f"{start_year}-07-01", + "end_date": f"{end_year + 1}-06-30", + "is_annual": bool(is_annual), + "is_generated_locally": True, + "source": "base_layer_clip", + "watershed_source": watershed_source, + "inputs": input_paths, + }, + ) + if layer_id and geoserver_synced: + update_layer_sync_status( + layer_id=layer_id, + sync_to_geoserver=True, + ) + + return { + "output": asset_id, + "layer_name": layer_name, + "is_annual": bool(is_annual), + "start_year": start_year, + "end_year": end_year, + "period_count": len(period_columns), + "watershed_count": len(result_gdf), + "source": "base_layer_clip", + "geoserver_synced": geoserver_synced, + "layer_id": layer_id, + } + + +def run_generate_hydrology_local( + *, + state=None, + district=None, + block=None, + pan_india=False, + start_year, + end_year, + is_annual=False, + gee_account_id=None, + hydrology_output_root=HYDROLOGY_OUTPUT_ROOT, + output_base_dir=HYDROLOGY_LOCAL_OUTPUT_DIR, + aquifer_vector_path=AQUIFER_VECTOR_PATH, + push_to_geoserver=True, + sync_layer_metadata=True, + overwrite=False, +): + pan_india = _parse_bool(pan_india) + start_year = int(start_year) + end_year = int(end_year) + if end_year < start_year: + raise ValueError("end_year must be greater than or equal to start_year") + + if pan_india: + raise ValueError( + "pan_india=true is not supported on the tehsil hydrology API. " + "Use the /api/v1/pan-india/ hydrology API to generate " + "Pan-India outputs." + ) + + return _run_clip_hydrology_area_from_base_layers( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + is_annual=is_annual, + output_base_dir=output_base_dir, + push_to_geoserver=push_to_geoserver, + sync_layer_metadata=sync_layer_metadata, + ) + + +@app.task(bind=True) +def generate_hydrology( + self, + state=None, + district=None, + block=None, + pan_india=False, + start_year=None, + end_year=None, + is_annual=False, + gee_account_id=None, + overwrite=False, +): + _ = self + return run_generate_hydrology_local( + state=state, + district=district, + block=block, + pan_india=pan_india, + start_year=start_year, + end_year=end_year, + is_annual=is_annual, + gee_account_id=gee_account_id, + push_to_geoserver=True, + sync_layer_metadata=True, + overwrite=overwrite, + ) + + +@app.task(bind=True) +def generate_hydrology_base_layer( + self, + year=None, + start_year=None, + end_year=None, + is_annual=False, + overwrite=False, +): + _ = self + return _run_generate_hydrology_base_layer_local( + year=year, + start_year=start_year, + end_year=end_year, + is_annual=is_annual, + overwrite=overwrite, + ) diff --git a/computing/mws/runoff_gpu.py b/computing/mws/runoff_gpu.py new file mode 100644 index 00000000..ec2e58a2 --- /dev/null +++ b/computing/mws/runoff_gpu.py @@ -0,0 +1,302 @@ +import os +from contextlib import contextmanager +from datetime import datetime +from pathlib import Path +from types import SimpleNamespace + +from nrm_app.celery import app + +from computing.config_loader import ( + LULC_BASE_DIR, + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + PROJECT_ROOT, + SOIL_RASTER_PATH, + TERRAIN_RASTER_PATH, +) +from utilities.gee_utils import valid_gee_text + + +DATA_ROOT = PROJECT_ROOT / "data" +HYDROLOGY_OUTPUT_ROOT = DATA_ROOT / "hydrology_gpu" +PAN_INDIA_RUNOFF_OUTPUT_ROOT = DATA_ROOT / "base_layers" / "hydrology" / "runoff" +PAN_INDIA_RUNOFF_COMMONS_ROOT = PAN_INDIA_RUNOFF_OUTPUT_ROOT / "commons" +PAN_INDIA_RUNOFF_TIMESERIES_DIR_NAME = "runoff_timeseries" +DEFAULT_LOCAL_DEM_PATH = TERRAIN_RASTER_PATH +DEFAULT_LOCAL_SOIL_PATH = SOIL_RASTER_PATH +PAN_INDIA_DEFAULT_TILE_SIZE = 11264 +STATE_DEFAULT_TILE_SIZE = 4096 + + +def _is_blank(value): + return value is None or str(value).strip().lower() in {"", "none", "null"} + + +def _parse_bool(value): + if isinstance(value, bool): + return value + if value is None: + return False + return str(value).strip().lower() in {"1", "true", "yes", "y"} + + +def _parse_date(value, field_name): + if _is_blank(value): + raise ValueError(f"{field_name} is required") + try: + return datetime.strptime(str(value), "%Y-%m-%d").date() + except ValueError as exc: + raise ValueError(f"{field_name} must be in YYYY-MM-DD format") from exc + + +def _slug(value, fallback): + return valid_gee_text(str(value).strip().lower()) or fallback + + +def _resolve_dates(start_date, end_date, start_year=None, end_year=None): + if not _is_blank(start_date) and not _is_blank(end_date): + start = _parse_date(start_date, "start_date") + end = _parse_date(end_date, "end_date") + if end <= start: + raise ValueError("end_date must be after start_date") + return start.isoformat(), end.isoformat(), start.year, end.year if end.year > start.year else start.year + 1 + + if _is_blank(start_year) or _is_blank(end_year): + raise ValueError("Provide start_date and end_date, or start_year and end_year") + + start_year = int(start_year) + end_year = int(end_year) + if end_year <= start_year: + raise ValueError("end_year must be greater than start_year") + + return f"{start_year}-07-01", f"{end_year}-07-01", start_year, end_year + + +def _resolve_lulc_path(lulc_start_year, lulc_end_year): + expected_name = f"lulc_v3_{lulc_start_year}_{lulc_end_year}.tif" + lulc_path = LULC_BASE_DIR / expected_name + if not lulc_path.exists(): + raise FileNotFoundError(f"LULC raster not found for requested annual period: {lulc_path}") + return lulc_path + + +def _validate_scope(pan_india, state, district, tehsil): + if pan_india: + if any(not _is_blank(value) for value in (state, district, tehsil)): + raise ValueError("pan_india=true cannot be combined with state, district, or tehsil") + return "pan_india", None, None, None + + if _is_blank(state): + raise ValueError("state is required unless pan_india=true") + if not _is_blank(tehsil) and _is_blank(district): + raise ValueError("district is required when tehsil is provided") + + state = str(state).strip() + district = None if _is_blank(district) else str(district).strip() + tehsil = None if _is_blank(tehsil) else str(tehsil).strip() + + if tehsil: + return "tehsil", state, district, tehsil + if district: + return "district", state, district, None + return "state", state, None, None + + +def _scope_slug(scope, state, district, tehsil): + parts = [scope] + for value, fallback in ((state, "state"), (district, "district"), (tehsil, "tehsil")): + if not _is_blank(value): + parts.append(_slug(value, fallback)) + return "/".join(parts) + + +def _ensure_default_inputs_exist(): + for label, path in ( + ("Default local DEM/terrain raster", DEFAULT_LOCAL_DEM_PATH), + ("Default local soil raster", DEFAULT_LOCAL_SOIL_PATH), + ): + if not path.exists(): + raise FileNotFoundError(f"{label} not found: {path}") + + +def _build_runner_args( + *, + state, + district, + tehsil, + pan_india, + start_date, + end_date, + local_lulc_path, + annual_key, +): + scope, state, district, tehsil = _validate_scope(pan_india, state, district, tehsil) + slug_path = _scope_slug(scope, state, district, tehsil) + if scope == "pan_india": + output_root = PAN_INDIA_RUNOFF_OUTPUT_ROOT / annual_key + boundary_output = PAN_INDIA_RUNOFF_COMMONS_ROOT / "pan_india.geojson" + timeseries_output = ( + output_root + / PAN_INDIA_RUNOFF_TIMESERIES_DIR_NAME + / "pan_india_timeseries.geojson" + ) + else: + output_root = HYDROLOGY_OUTPUT_ROOT / slug_path / annual_key + boundary_output = output_root / "boundaries" / f"{slug_path.replace('/', '_')}.geojson" + timeseries_output = None + + return SimpleNamespace( + pre_req=True, + boundary=None, + t=True, + start=start_date, + end=end_date, + rainfall_folder=str(output_root / "rainfall"), + runoffs_folder=str(output_root / "runoffs"), + demfile_path=str(output_root / "dem.tif"), + pan_india=pan_india, + state=state, + district=district, + tehsil=tehsil, + watershed_root=str(PRECOMPUTED_TEHSIL_WATERSHED_DIR), + watershed_boundary_output=str(boundary_output), + timeseries_vector=str(timeseries_output) if timeseries_output else None, + reuse_watershed_boundary=True, + local_dem=str(DEFAULT_LOCAL_DEM_PATH), + local_lulc=str(local_lulc_path), + lulc_source="indiasatv3", + local_soil=str(DEFAULT_LOCAL_SOIL_PATH), + tile_size=None, + ) + + +@contextmanager +def _working_directory(path): + previous = os.getcwd() + os.chdir(path) + try: + yield + finally: + os.chdir(previous) + + +def run_runoff_gpu_local( + *, + state=None, + district=None, + tehsil=None, + pan_india=False, + start_date=None, + end_date=None, + start_year=None, + end_year=None, +): + _ensure_default_inputs_exist() + pan_india = _parse_bool(pan_india) + start_date, end_date, lulc_start_year, lulc_end_year = _resolve_dates( + start_date=start_date, + end_date=end_date, + start_year=start_year, + end_year=end_year, + ) + local_lulc_path = _resolve_lulc_path(lulc_start_year, lulc_end_year) + annual_key = f"{lulc_start_year}_{lulc_end_year}" + args = _build_runner_args( + state=state, + district=district, + tehsil=tehsil, + pan_india=pan_india, + start_date=start_date, + end_date=end_date, + local_lulc_path=local_lulc_path, + annual_key=annual_key, + ) + + with _working_directory(PROJECT_ROOT): + from computing.hydrology_gpu import runoff as hydro_runoff + from computing.hydrology_gpu.downloads import dem + + hydro_runoff.validate_local_raster(args.local_lulc, "--local-lulc") + hydro_runoff.validate_local_raster(args.local_soil, "--local-soil") + hydro_runoff.resolve_boundary(args) + if args.tile_size is None: + if args.pan_india: + args.tile_size = PAN_INDIA_DEFAULT_TILE_SIZE + elif args.state and not args.district: + args.tile_size = STATE_DEFAULT_TILE_SIZE + else: + args.tile_size = 0 + hydro_runoff.modify_cfg(args) + hydro_runoff.cfg.DEMFILE_PATH = args.local_dem if args.pan_india else args.demfile_path + + hydro_runoff.shutil.rmtree(hydro_runoff.cfg.RAINFALL_FOLDER, ignore_errors=True) + if args.pan_india: + hydro_runoff.logger.info( + "Using existing pan-India DEM/slope raster directly: %s", + hydro_runoff.cfg.DEMFILE_PATH, + ) + else: + with hydro_runoff.timed_stage(f"local DEM/slope clip from {args.local_dem}"): + dem.clip_local_raster( + args.local_dem, + hydro_runoff.cfg.BOUNDARY_GEOJSON_PATH, + hydro_runoff.cfg.DEMFILE_PATH, + hydro_runoff.logger, + ) + + with hydro_runoff.timed_stage(f"loading DEM reference grid from {hydro_runoff.cfg.DEMFILE_PATH}"): + hydro_runoff.utils.tif_handler = hydro_runoff.GeoTIFFHandler( + hydro_runoff.cfg.DEMFILE_PATH, + hydro_runoff.logger, + ) + + with hydro_runoff.timed_stage("remaining prerequisite downloads"): + hydro_runoff.prereq(args) + + hydro_runoff.shutil.rmtree(hydro_runoff.cfg.RUNOFFS_FOLDER, ignore_errors=True) + with hydro_runoff.timed_stage("runoff/timeseries processing"): + if args.tile_size: + hydro_runoff.tiled_timeseries.TiledTimeSeries(args.tile_size).run() + else: + hydro_runoff.timeseries.TimeSeries().run() + + return { + "scope": "pan_india" if args.pan_india else ("tehsil" if args.tehsil else ("district" if args.district else "state")), + "state": args.state, + "district": args.district, + "tehsil": args.tehsil, + "start_date": start_date, + "end_date": end_date, + "annual_key": annual_key, + "lulc_path": str(local_lulc_path), + "rainfall_folder": args.rainfall_folder, + "runoffs_folder": args.runoffs_folder, + "boundary": args.boundary, + "microwatersheds": args.microwatersheds, + "timeseries_vector": str(hydro_runoff.cfg.TIMESERIES_VECTOR), + "tile_size": args.tile_size, + } + + +@app.task(bind=True) +def generate_runoff_gpu( + self, + state=None, + district=None, + tehsil=None, + pan_india=False, + start_date=None, + end_date=None, + start_year=None, + end_year=None, +): + _ = self + return run_runoff_gpu_local( + state=state, + district=district, + tehsil=tehsil, + pan_india=pan_india, + start_date=start_date, + end_date=end_date, + start_year=start_year, + end_year=end_year, + ) diff --git a/computing/tasks.py b/computing/tasks.py index 71f0c21a..f0899b10 100644 --- a/computing/tasks.py +++ b/computing/tasks.py @@ -4,6 +4,12 @@ from computing.STAC_specs.stac_collection import generate_stac_collection_task from computing.bulk_layer_generation import run_pipeline +from computing.mws.et_download import et_download +from computing.mws.generate_hydrology_local import ( + generate_hydrology, + generate_hydrology_base_layer, +) +from computing.mws.runoff_gpu import generate_runoff_gpu logger = logging.getLogger(__name__) @@ -37,4 +43,9 @@ def bulk_generate_layer( ) -__all__ = ["bulk_generate_layer", "generate_stac_collection_task"] +__all__ = ["bulk_generate_layer", + "generate_stac_collection_task", + "et_download", + "generate_hydrology", + "generate_hydrology_base_layer", + "generate_runoff_gpu",] diff --git a/computing/urls.py b/computing/urls.py index b5bdfe23..29d6f862 100644 --- a/computing/urls.py +++ b/computing/urls.py @@ -20,6 +20,18 @@ name="hydrology_fortnightly", ), path("hydrology_annual/", api.generate_annual_hydrology, name="hydrology_annual"), + path( + "pan-india/hydrology_fortnightly/", + api.generate_pan_india_fortnightly_hydrology, + name="pan_india_hydrology_fortnightly", + ), + path( + "pan-india/hydrology_annual/", + api.generate_pan_india_annual_hydrology, + name="pan_india_hydrology_annual", + ), + path("runoff_gpu/", api.generate_runoff_gpu, name="runoff_gpu"), + path("et_download/", api.et_download, name="et_download"), path("lulc_for_tehsil/", api.lulc_for_tehsil, name="lulc_for_tehsil"), path("lulc_v2_river_basin/", api.lulc_v2_river_basin, name="lulc_v2_river_basin"), path("lulc_v3_river_basin/", api.lulc_v3_river_basin, name="lulc_v3_river_basin"), diff --git a/installation/environment.yml b/installation/environment.yml index ffae3d36..03a2c54b 100644 --- a/installation/environment.yml +++ b/installation/environment.yml @@ -9,6 +9,7 @@ dependencies: - pip - setuptools=80 - wheel + - docker-compose=1.29.2 - djangorestframework=3.15.2 @@ -24,6 +25,7 @@ dependencies: - seaborn=0.13.0 - contourpy=1.2.0 + - rasterio - gdal=3.6.4 - geopandas=0.14.1 - fiona=1.9.5 @@ -52,7 +54,7 @@ dependencies: - django-timezone-field==7.2.1 - djangorestframework-simplejwt==5.5.0 - djangorestframework-api-key==3.1.0 - - earthengine-api==1.5.9 + - earthengine-api==1.6.12 - geoserver-rest==2.5.3 - geojson==3.1.0 - google-api-core==2.17.0 @@ -65,7 +67,6 @@ dependencies: - google-resumable-media==2.7.0 - googleapis-common-protos==1.62.0 - docker==5.0.3 - - docker-compose==1.29.2 - python-docx==1.1.0 - pymongo==3.11.0 - xmltodict==0.13.0 @@ -73,7 +74,6 @@ dependencies: - emoji - boto3 - speechrecognition==3.14.3 - - rasterio - pystac - tqdm - geemap @@ -87,4 +87,13 @@ dependencies: - orjson - ijson - polars - - weasyprint \ No newline at end of file + - weasyprint + - cupy-cuda12x[ctk]>=13.6.0 + - cucim-cu12>=26.2.0 + - geedim==1.9.1 + - natsort>=8.4.0 + - pydrive2>=1.21.3 + - rioxarray==0.19.0 + - xee==0.0.24 + - zarr==2.18.3 + - xarray \ No newline at end of file From 56fd3457b9d5ee0e8abaa7699566f887bf2b60bc Mon Sep 17 00:00:00 2001 From: ramank1137 Date: Fri, 14 Aug 2026 11:43:42 +0530 Subject: [PATCH 2/6] adding environment changes in installation --- installation/environment.yml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/installation/environment.yml b/installation/environment.yml index 03a2c54b..44737c59 100644 --- a/installation/environment.yml +++ b/installation/environment.yml @@ -1,4 +1,4 @@ -name: corestack-backend +name: corestack-backend-test channels: - conda-forge @@ -87,6 +87,7 @@ dependencies: - orjson - ijson - polars + - numpy==1.26.4 - weasyprint - cupy-cuda12x[ctk]>=13.6.0 - cucim-cu12>=26.2.0 @@ -96,4 +97,4 @@ dependencies: - rioxarray==0.19.0 - xee==0.0.24 - zarr==2.18.3 - - xarray \ No newline at end of file + - xarray From 1df5bf9bfd4f4081300cd0fc34cf78a884d7d2b6 Mon Sep 17 00:00:00 2001 From: ramank1137 Date: Fri, 14 Aug 2026 11:56:48 +0530 Subject: [PATCH 3/6] adding changes of aquifer_vector_local that are used by hydrology --- computing/misc/aquifer_vector_local.py | 49 ++++++++++++++++---------- 1 file changed, 31 insertions(+), 18 deletions(-) diff --git a/computing/misc/aquifer_vector_local.py b/computing/misc/aquifer_vector_local.py index c67a4281..76d452b1 100644 --- a/computing/misc/aquifer_vector_local.py +++ b/computing/misc/aquifer_vector_local.py @@ -183,25 +183,13 @@ def _build_aquifer_properties(watershed_row, area_in_ha, intersections_df): ) return properties - -def _compute_aquifer_properties_for_watersheds(watersheds_gdf, aquifers_gdf): - watersheds_gdf = validate_geometry(watersheds_gdf) - if watersheds_gdf.empty: - raise ValueError("No valid watershed geometries found for local processing.") - if watersheds_gdf.crs is None: - raise ValueError("Watershed CRS is missing; cannot compute aquifer overlaps.") - +def _prepare_aquifers_for_intersection(aquifers_gdf): aquifers_gdf = validate_geometry(aquifers_gdf) if aquifers_gdf.empty: raise ValueError("Aquifer source file has no valid geometries.") if aquifers_gdf.crs is None: raise ValueError("Aquifer source CRS is missing; cannot compute overlaps.") - watersheds_result = watersheds_gdf.copy() - watersheds_result["area_in_ha"] = get_watershed_areas_in_hectares( - watersheds_result - ).astype(float) - aquifers_with_yield = aquifers_gdf.copy() aquifers_with_yield["y_value"] = aquifers_with_yield["yeild__"].apply( _map_yield_value @@ -212,8 +200,30 @@ def _compute_aquifer_properties_for_watersheds(watersheds_gdf, aquifers_gdf): if aquifers_with_yield.empty: raise ValueError("Aquifer source has no records with valid yield values.") + return aquifers_with_yield.to_crs("EPSG:6933") + +def _compute_aquifer_properties_for_watersheds( + watersheds_gdf, + aquifers_gdf=None, + aquifers_projected=None, +): + watersheds_gdf = validate_geometry(watersheds_gdf) + if watersheds_gdf.empty: + raise ValueError("No valid watershed geometries found for local processing.") + if watersheds_gdf.crs is None: + raise ValueError("Watershed CRS is missing; cannot compute aquifer overlaps.") + + if aquifers_projected is None: + if aquifers_gdf is None: + raise ValueError("Aquifer source is required for local processing.") + aquifers_projected = _prepare_aquifers_for_intersection(aquifers_gdf) + + watersheds_result = watersheds_gdf.copy() + watersheds_result["area_in_ha"] = get_watershed_areas_in_hectares( + watersheds_result + ).astype(float) watersheds_projected = watersheds_result.to_crs("EPSG:6933") - aquifers_projected = aquifers_with_yield.to_crs("EPSG:6933") + aquifer_spatial_index = aquifers_projected.sindex computed_rows = [] total = len(watersheds_projected) @@ -236,9 +246,13 @@ def _compute_aquifer_properties_for_watersheds(watersheds_gdf, aquifers_gdf): ) continue - intersecting_aquifers = aquifers_projected.loc[ - aquifers_projected.intersects(watershed_geometry) - ] + candidate_positions = sorted( + aquifer_spatial_index.query( + watershed_geometry, + predicate="intersects", + ) + ) + intersecting_aquifers = aquifers_projected.iloc[candidate_positions] intersections = [] for _, aquifer_row in intersecting_aquifers.iterrows(): @@ -295,7 +309,6 @@ def _compute_aquifer_properties_for_watersheds(watersheds_gdf, aquifers_gdf): watersheds_result[column] = computed_df[column].values return watersheds_result - def run_aquifer_vector_local( state, district, From d97dbdecf6da4b81559c2379f14cb0772b30f587 Mon Sep 17 00:00:00 2001 From: ramank1137 Date: Fri, 14 Aug 2026 15:54:45 +0530 Subject: [PATCH 4/6] adding fallback logic for if outerboundary file not present --- computing/hydrology_gpu/watershed_boundary.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/computing/hydrology_gpu/watershed_boundary.py b/computing/hydrology_gpu/watershed_boundary.py index 4e99f817..b3ff4a98 100644 --- a/computing/hydrology_gpu/watershed_boundary.py +++ b/computing/hydrology_gpu/watershed_boundary.py @@ -407,10 +407,16 @@ def materialize_pan_india_boundary( if destination.exists() and not overwrite: download_destination = download_boundary_path(destination) download_boundary_source = Path(download_boundary_source) - if not download_destination.exists() and download_boundary_source.exists(): - download_destination.parent.mkdir(parents=True, exist_ok=True) - download_destination.write_text(download_boundary_source.read_text()) - gdf_existing = gpd.read_file(destination) + gdf_existing = None + if not download_destination.exists(): + if download_boundary_source.exists(): + download_destination.parent.mkdir(parents=True, exist_ok=True) + download_destination.write_text(download_boundary_source.read_text()) + else: + gdf_existing = gpd.read_file(destination) + write_download_boundary(gdf_existing, download_destination, PAN_INDIA_SLUG) + if gdf_existing is None: + gdf_existing = gpd.read_file(destination) return destination, [path for path, _ in matches], len(gdf_existing) frames = [] From 53e9af5f9b88f9a94e4beb4548fa12e72ef5754f Mon Sep 17 00:00:00 2001 From: ramank1137 Date: Wed, 19 Aug 2026 14:57:02 +0530 Subject: [PATCH 5/6] adding changes to make the code agnostic to machine --- computing/hydrology_gpu/watershed_boundary.py | 59 +++++++++++-------- computing/mws/generate_hydrology_local.py | 11 +++- computing/mws/runoff_gpu.py | 17 +++++- installation/environment.yml | 6 +- 4 files changed, 64 insertions(+), 29 deletions(-) diff --git a/computing/hydrology_gpu/watershed_boundary.py b/computing/hydrology_gpu/watershed_boundary.py index b3ff4a98..9fc44921 100644 --- a/computing/hydrology_gpu/watershed_boundary.py +++ b/computing/hydrology_gpu/watershed_boundary.py @@ -73,6 +73,17 @@ def load_manifest(root: Path) -> list[dict]: return list(csv.DictReader(f)) +def row_feature_count(row: dict) -> int: + try: + return int(row.get("feature_count") or 0) + except (TypeError, ValueError): + return 0 + + +def rows_feature_count(rows: list[dict]) -> int: + return sum(row_feature_count(row) for row in rows) + + def manifest_relative_output(root: Path, output_path: str) -> Path | None: if not output_path: return None @@ -286,6 +297,20 @@ def write_download_boundary(gdf, destination: str | Path, state: str, district: return destination +def copy_pan_india_download_boundary(source: str | Path, destination: str | Path) -> Path: + source = Path(source) + if not source.exists(): + raise FileNotFoundError( + "Pan-India download boundary source not found: " + f"{source}. Add the canonical boundary file instead of using a fallback." + ) + + destination = Path(destination) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(source.read_text()) + return destination + + def materialize_tehsil_boundary( state: str, district: str, @@ -295,12 +320,11 @@ def materialize_tehsil_boundary( overwrite: bool = True, ) -> tuple[Path, Path, int]: root = Path(watershed_root) - source_path, _ = find_tehsil_watershed(root, state, district, tehsil) + source_path, row = find_tehsil_watershed(root, state, district, tehsil) destination = Path(output_path) if output_path else default_output_path(state, district, tehsil) if destination.exists() and not overwrite: - gdf_existing = gpd.read_file(destination) - return destination, source_path, len(gdf_existing) + return destination, source_path, row_feature_count(row) gdf = prepare_boundary_gdf(gpd.read_file(source_path), state, district, tehsil, source_path) gdf["id"] = range(1, len(gdf) + 1) @@ -323,8 +347,8 @@ def materialize_district_boundary( destination = Path(output_path) if output_path else default_district_output_path(state, district) if destination.exists() and not overwrite: - gdf_existing = gpd.read_file(destination) - return destination, [path for path, _ in matches], len(gdf_existing) + feature_count = rows_feature_count([row for _, row in matches]) + return destination, [path for path, _ in matches], feature_count frames = [] source_paths = [] @@ -363,8 +387,8 @@ def materialize_state_boundary( destination = Path(output_path) if output_path else default_state_output_path(state) if destination.exists() and not overwrite: - gdf_existing = gpd.read_file(destination) - return destination, [path for path, _ in matches], len(gdf_existing) + feature_count = rows_feature_count([row for _, row in matches]) + return destination, [path for path, _ in matches], feature_count frames = [] source_paths = [] @@ -406,18 +430,9 @@ def materialize_pan_india_boundary( if destination.exists() and not overwrite: download_destination = download_boundary_path(destination) - download_boundary_source = Path(download_boundary_source) - gdf_existing = None - if not download_destination.exists(): - if download_boundary_source.exists(): - download_destination.parent.mkdir(parents=True, exist_ok=True) - download_destination.write_text(download_boundary_source.read_text()) - else: - gdf_existing = gpd.read_file(destination) - write_download_boundary(gdf_existing, download_destination, PAN_INDIA_SLUG) - if gdf_existing is None: - gdf_existing = gpd.read_file(destination) - return destination, [path for path, _ in matches], len(gdf_existing) + copy_pan_india_download_boundary(download_boundary_source, download_destination) + feature_count = rows_feature_count([row for _, row in matches]) + return destination, [path for path, _ in matches], feature_count frames = [] source_paths = [] @@ -445,10 +460,6 @@ def materialize_pan_india_boundary( destination.write_text(combined.to_json()) download_destination = download_boundary_path(destination) - download_boundary_source = Path(download_boundary_source) - if download_boundary_source.exists(): - download_destination.write_text(download_boundary_source.read_text()) - else: - write_download_boundary(combined, download_destination, PAN_INDIA_SLUG) + copy_pan_india_download_boundary(download_boundary_source, download_destination) return destination, source_paths, len(combined) diff --git a/computing/mws/generate_hydrology_local.py b/computing/mws/generate_hydrology_local.py index 038e562d..7fcd3856 100644 --- a/computing/mws/generate_hydrology_local.py +++ b/computing/mws/generate_hydrology_local.py @@ -610,6 +610,13 @@ def _read_complete_uid_cache(path, required_uids, value_columns): return frame +def _geometry_union(geometries): + union_all = getattr(geometries, "union_all", None) + if callable(union_all): + return union_all() + return geometries.unary_union + + def _iter_unique_watershed_partitions( matches, *, @@ -1024,7 +1031,7 @@ def _uses_daily_et(watersheds_gdf, daily_roots_by_year): with rasterio.open(reference_path) as src: raster_bounds = box(*src.bounds) watersheds = watersheds_gdf.to_crs(src.crs) - return raster_bounds.covers(watersheds.geometry.union_all()) + return raster_bounds.covers(_geometry_union(watersheds.geometry)) def _calculate_period_et(watersheds_gdf, periods, et_roots_by_year): @@ -1219,7 +1226,7 @@ def _raster_covers_watersheds(watersheds_gdf, raster_path): with rasterio.open(raster_path) as src: raster_bounds = box(*src.bounds) watersheds = watersheds_gdf.to_crs(src.crs) - return raster_bounds.covers(watersheds.geometry.union_all()) + return raster_bounds.covers(_geometry_union(watersheds.geometry)) def _calculate_period_et_from_aggregates( diff --git a/computing/mws/runoff_gpu.py b/computing/mws/runoff_gpu.py index ec2e58a2..6d554a1a 100644 --- a/computing/mws/runoff_gpu.py +++ b/computing/mws/runoff_gpu.py @@ -4,6 +4,8 @@ from pathlib import Path from types import SimpleNamespace +from celery.utils.log import get_task_logger + from nrm_app.celery import app from computing.config_loader import ( @@ -25,6 +27,7 @@ DEFAULT_LOCAL_SOIL_PATH = SOIL_RASTER_PATH PAN_INDIA_DEFAULT_TILE_SIZE = 11264 STATE_DEFAULT_TILE_SIZE = 4096 +logger = get_task_logger(__name__) def _is_blank(value): @@ -289,7 +292,19 @@ def generate_runoff_gpu( start_year=None, end_year=None, ): - _ = self + logger.info( + "Starting runoff_gpu task_id=%s state=%s district=%s tehsil=%s pan_india=%s " + "start_date=%s end_date=%s start_year=%s end_year=%s", + self.request.id, + state, + district, + tehsil, + pan_india, + start_date, + end_date, + start_year, + end_year, + ) return run_runoff_gpu_local( state=state, district=district, diff --git a/installation/environment.yml b/installation/environment.yml index 44737c59..f2ff3c9c 100644 --- a/installation/environment.yml +++ b/installation/environment.yml @@ -1,6 +1,7 @@ name: corestack-backend-test channels: + - rapidsai - conda-forge - defaults @@ -47,6 +48,9 @@ dependencies: - cairo - gdk-pixbuf - libffi + - cuda-version=12.9 + - cupy=13.6.* + - cucim>=26.2.0 - pip: - django==5.2.9 @@ -89,8 +93,6 @@ dependencies: - polars - numpy==1.26.4 - weasyprint - - cupy-cuda12x[ctk]>=13.6.0 - - cucim-cu12>=26.2.0 - geedim==1.9.1 - natsort>=8.4.0 - pydrive2>=1.21.3 From 1fc47cb86d6e3a0f1f2eddb0d8f24abc15f4f73e Mon Sep 17 00:00:00 2001 From: ramank1137 Date: Wed, 19 Aug 2026 17:22:51 +0530 Subject: [PATCH 6/6] adding code to generate microwatersheds include ids, adding code to raise exception in hydrology clipping if the pan-india layers are not present. --- computing/api.py | 25 ++++++++++++++++++ computing/config.yaml | 2 +- computing/mws/generate_hydrology_local.py | 26 +++++++++++++++++++ .../store_watersheds_for_tehsils.py | 4 +-- 4 files changed, 54 insertions(+), 3 deletions(-) diff --git a/computing/api.py b/computing/api.py index 18bca62f..dfc454af 100644 --- a/computing/api.py +++ b/computing/api.py @@ -164,6 +164,7 @@ from .mws.generate_hydrology_local import ( generate_hydrology_base_layer as generate_hydrology_base_layer_task, generate_hydrology as generate_hydrology_local_task, + missing_hydrology_base_layers, ) from .mws.et_download import et_download as et_download_task from .mws.runoff_gpu import generate_runoff_gpu as generate_runoff_gpu_task @@ -509,6 +510,7 @@ def _generate_tehsil_hydrology(request, is_annual): "Local hydrology clipping must start from start_year=2017 because " "the fortnightly cadence and cumulative G are anchored at 2017-07-01" ) + _ensure_local_hydrology_base_layers(start_year, end_year, is_annual) task = generate_hydrology_local_task.apply_async( kwargs={ "state": state, @@ -540,6 +542,29 @@ def _generate_tehsil_hydrology(request, is_annual): ) +def _ensure_local_hydrology_base_layers(start_year, end_year, is_annual): + missing_layers = missing_hydrology_base_layers( + start_year=start_year, + end_year=end_year, + is_annual=is_annual, + ) + if not missing_layers: + return + + period = "annual" if is_annual else "fortnightly" + endpoint = ( + "/api/v1/pan-india/hydrology_annual/" + if is_annual + else "/api/v1/pan-india/hydrology_fortnightly/" + ) + missing_summary = ", ".join(layer["year_key"] for layer in missing_layers) + raise ValueError( + f"Missing Pan-India hydrology {period} base layer(s) for requested " + f"year(s): {missing_summary}. Generate the missing base layer(s) first " + f"using {endpoint}; Celery task was not queued." + ) + + def _generate_pan_india_hydrology_base_layer(request, is_annual): compute = _get_compute_mode(request) if compute != "local": diff --git a/computing/config.yaml b/computing/config.yaml index 4e648f5e..de847e6e 100644 --- a/computing/config.yaml +++ b/computing/config.yaml @@ -101,7 +101,7 @@ base_layers: type: file - name: admin boundaries - local_path: "{DATA_DIR}/admin-boundary/input/soi_tehsil.geojson" + local_path: "{DATA_DIR}/base_layers/admin_boundary/soi_tehsil.geojson" source: "" # Full admin-boundary archive (~8 GB, 7z) containing this file; extracted in place. gdrive_id: 1VqIhB6HrKFDkDnlk1vedcEHhh5fk4f1d diff --git a/computing/mws/generate_hydrology_local.py b/computing/mws/generate_hydrology_local.py index 7fcd3856..8f518c63 100644 --- a/computing/mws/generate_hydrology_local.py +++ b/computing/mws/generate_hydrology_local.py @@ -150,6 +150,32 @@ def _base_layer_path(base_layer_root, year, is_annual): ) +def missing_hydrology_base_layers( + *, + start_year, + end_year, + is_annual, + base_layer_root=HYDROLOGY_BASE_LAYER_ROOT, +): + start_year = int(start_year) + end_year = int(end_year) + if end_year < start_year: + raise ValueError("end_year must be greater than or equal to start_year") + + missing_layers = [] + for year in range(start_year, end_year + 1): + path = _base_layer_path(base_layer_root, year, is_annual) + if not path.is_file(): + missing_layers.append( + { + "year": year, + "year_key": _year_key(year), + "path": str(path), + } + ) + return missing_layers + + def _write_parquet_atomic(frame, path): path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) diff --git a/computing/terrain_descriptor/store_watersheds_for_tehsils.py b/computing/terrain_descriptor/store_watersheds_for_tehsils.py index 6986b4f0..dc58d21e 100644 --- a/computing/terrain_descriptor/store_watersheds_for_tehsils.py +++ b/computing/terrain_descriptor/store_watersheds_for_tehsils.py @@ -10,14 +10,14 @@ DEFAULT_MICROWATERSHED_PATH = ( "data/base_layers/Microwatershed_v2_with_details.geojson" ) -DEFAULT_TEHSIL_PATH = "data/admin-boundary/input/soi_tehsil.geojson" +DEFAULT_TEHSIL_PATH = "data/base_layers/admin_boundary/soi_tehsil.geojson" DEFAULT_OUTPUT_DIR = "data/base_layers/tehsil_watersheds" STATE_COLUMN_CANDIDATES = ["STATE", "state", "state_name", "State"] DISTRICT_COLUMN_CANDIDATES = ["District", "district", "district_name", "DISTRICT"] TEHSIL_COLUMN_CANDIDATES = ["TEHSIL", "tehsil", "tehsil_name", "block", "block_name"] MWS_UID_COLUMN_CANDIDATES = ["uid", "UID", "Uid"] -MWS_OPTIONAL_COLUMNS = ["area_in_ha", "bacode", "sbcode", "wsconc"] +MWS_OPTIONAL_COLUMNS = ["id", "area_in_ha", "bacode", "sbcode", "wsconc"] OUTPUT_FORMATS = { "geojson": ("GeoJSON", ".geojson"),