diff --git a/.gitignore b/.gitignore index 65411952..c28ebc76 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,10 @@ venv/ .DS_Store .cursor/ +#IDE +.codex +.vscode/* + # Django *.log local_settings.py diff --git a/.installation_state/collectstatic.done b/.installation_state/collectstatic.done new file mode 100644 index 00000000..2812016d --- /dev/null +++ b/.installation_state/collectstatic.done @@ -0,0 +1 @@ +2026-04-13T19:43:46Z diff --git a/.installation_state/conda_env.done b/.installation_state/conda_env.done new file mode 100644 index 00000000..60785b45 --- /dev/null +++ b/.installation_state/conda_env.done @@ -0,0 +1 @@ +2026-04-13T19:43:39Z diff --git a/.installation_state/env_file.done b/.installation_state/env_file.done new file mode 100644 index 00000000..278c3835 --- /dev/null +++ b/.installation_state/env_file.done @@ -0,0 +1 @@ +2026-04-13T19:43:40Z diff --git a/.installation_state/miniconda.done b/.installation_state/miniconda.done new file mode 100644 index 00000000..e0463ce4 --- /dev/null +++ b/.installation_state/miniconda.done @@ -0,0 +1 @@ +2026-04-13T19:40:19Z diff --git a/.installation_state/postgres.done b/.installation_state/postgres.done new file mode 100644 index 00000000..c4573478 --- /dev/null +++ b/.installation_state/postgres.done @@ -0,0 +1 @@ +2026-04-13T19:40:33Z diff --git a/.installation_state/rabbitmq.done b/.installation_state/rabbitmq.done new file mode 100644 index 00000000..d868a03c --- /dev/null +++ b/.installation_state/rabbitmq.done @@ -0,0 +1 @@ +2026-04-13T19:40:34Z diff --git a/.installation_state/unzip_install.done b/.installation_state/unzip_install.done new file mode 100644 index 00000000..8fe01fb6 --- /dev/null +++ b/.installation_state/unzip_install.done @@ -0,0 +1 @@ +2026-04-13T19:40:06Z diff --git a/computing/api.py b/computing/api.py index 365a0f7b..1a9d58e2 100644 --- a/computing/api.py +++ b/computing/api.py @@ -9,7 +9,10 @@ from rest_framework.parsers import MultiPartParser, FormParser from computing.change_detection.change_detection_vector import ( - vectorise_change_detection, + vectorise_change_detection as vectorise_change_detection_gee_task, +) +from computing.change_detection.change_detection_vector_local import ( + vectorise_change_detection as vectorise_change_detection_local_task, ) from .utils import ( save_layer_info_to_db, @@ -17,7 +20,8 @@ ) from django.conf import settings from computing.STAC_specs.stac_collection import sanitize_text, STACConfig -from .lulc.lulc_vector import vectorise_lulc +from .lulc.lulc_vector import vectorise_lulc as vectorise_lulc_gee_task +from .lulc.lulc_vector_local import vectorise_lulc as vectorise_lulc_local_task from .lulc.river_basin_lulc.lulc_v2_river_basin import lulc_river_basin_v2 from .lulc.river_basin_lulc.lulc_v3_river_basin_using_v2 import lulc_river_basin_v3 from .lulc.tehsil_level.lulc_v2 import generate_lulc_v2_tehsil @@ -26,28 +30,74 @@ from .misc.ndvi_time_series import ndvi_timeseries from .misc.restoration_opportunity import generate_restoration_opportunity 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 .utils import ( Geoserver, kml_to_shp, + save_layer_info_to_db, + update_layer_sync_status, +) +from .local_compute_helper import ( + get_compute_mode as _get_compute_mode, + select_compute_task as _select_compute_task, ) from utilities.gee_utils import download_gee_layer, check_gee_task_status from django.core.files.storage import FileSystemStorage from utilities.constants import KML_PATH from .mws.mws import mws_layer from .cropping_intensity.cropping_intensity import generate_cropping_intensity +from .cropping_intensity.cropping_intesity_local import ( + generate_cropping_intensity as generate_cropping_intensity_local_task, +) from .surface_water_bodies.swb import generate_swb_layer from .drought.drought import calculate_drought -from .terrain_descriptor.terrain_clusters import generate_terrain_clusters -from .terrain_descriptor.terrain_raster_fabdem import generate_terrain_raster_clip +from .terrain_descriptor.terrain_clusters import ( + generate_terrain_clusters as generate_terrain_clusters_gee_task, +) +from .terrain_descriptor.terrain_clusters_local import ( + generate_terrain_clusters as generate_terrain_clusters_local_task, +) +from .terrain_descriptor.terrain_compute_all_local import ( + generate_terrain_compute_all as generate_terrain_compute_all_task, +) +from .terrain_descriptor.terrain_raster_fabdem import ( + generate_terrain_raster_clip as generate_terrain_raster_clip_gee_task, +) +from .terrain_descriptor.terrain_raster_fabdem_local import ( + generate_terrain_raster_clip as generate_terrain_raster_clip_local_task, +) from computing.misc.drainage_lines import clip_drainage_lines -from .lulc_X_terrain.lulc_on_slope_cluster import lulc_on_slope_cluster -from .lulc_X_terrain.lulc_on_plain_cluster import lulc_on_plain_cluster +from .lulc_X_terrain.lulc_on_slope_cluster import ( + lulc_on_slope_cluster as lulc_on_slope_cluster_gee_task, +) +from .lulc_X_terrain.lulc_on_slope_cluster_local import ( + lulc_on_slope_cluster_local as lulc_on_slope_cluster_local_task, +) +from .lulc_X_terrain.lulc_on_plain_cluster import ( + lulc_on_plain_cluster as lulc_on_plain_cluster_gee_task, +) +from .lulc_X_terrain.lulc_on_plain_cluster_local import ( + lulc_on_plain_cluster_local as lulc_on_plain_cluster_local_task, +) from .clart.clart import generate_clart_layer from .misc.admin_boundary import generate_tehsil_shape_file_data from .misc.nrega import clip_nrega_district_block -from computing.change_detection.change_detection import get_change_detection -from .lulc.lulc_v3 import clip_lulc_v3 +from computing.change_detection.change_detection import ( + get_change_detection as get_change_detection_gee_task, +) +from computing.change_detection.change_detection_local import ( + get_change_detection as get_change_detection_local_task, +) +from .lulc.lulc_v3 import clip_lulc_v3 as clip_lulc_v3_gee_task +from .lulc.lulc_v3_local import clip_lulc_v3 as clip_lulc_v3_local_task from .crop_grid.crop_grid import create_crop_grids from .tree_health.ccd import tree_health_ccd_raster from .tree_health.canopy_height import tree_health_ch_raster @@ -57,7 +107,12 @@ from .tree_health.canopy_height_vector import tree_health_ch_vector from .tree_health.ccd_vector import tree_health_ccd_vector from .plantation.site_suitability import site_suitability -from .misc.aquifer_vector import generate_aquifer_vector +from .misc.aquifer_vector import ( + generate_aquifer_vector as generate_aquifer_vector_gee_task, +) +from .misc.aquifer_vector_local import ( + generate_aquifer_vector as generate_aquifer_vector_local_task, +) from .misc.soge_vector import generate_soge_vector from .clart.fes_clart_to_geoserver import generate_fes_clart_layer from .surface_water_bodies.merge_swb_ponds import merge_swb_ponds @@ -78,6 +133,28 @@ from .mws.mws_centroid import generate_mws_centroid_data from .misc.facilities_proximity import generate_facilities_proximity_task from .STAC_specs.stac_collection import _make_celery_task as _make_stac_task +from django.conf import settings + + +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") @@ -227,61 +304,283 @@ def generate_mws_layer(request): 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( - {"Success": "Successfully initiated"}, status=status.HTTP_200_OK + {"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_fortnightly_hydrology api :: ", e) - return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) + 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_annual_hydrology(request): - print("Inside generate_annual_hydrology") +def generate_runoff_gpu(request): + print("Inside generate_runoff_gpu") 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("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": state, - "district": district, - "block": block, - "start_year": start_year, - "end_year": end_year, - "is_annual": True, - "gee_account_id": gee_account_id, + "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_annual_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 et_download(request): + print("Inside et_download") + try: + 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={ + "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": "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 et_download api :: ", e) + return Response( + {"Exception": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR + ) @api_view(["POST"]) @@ -385,13 +684,22 @@ def lulc_v3(request): start_year = request.data.get("start_year") end_year = request.data.get("end_year") gee_account_id = request.data.get("gee_account_id") - clip_lulc_v3.apply_async( + compute = _get_compute_mode(request) + task = _select_compute_task( + compute, + clip_lulc_v3_gee_task, + clip_lulc_v3_local_task, + ) + task.apply_async( args=[state, district, block, start_year, end_year, gee_account_id], queue="nrm", ) return Response( {"Success": "LULC v3 task initiated"}, status=status.HTTP_200_OK ) + except ValueError as e: + print("Invalid request in lulc_v3 api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: print("Exception in lulc_v3 api :: ", e) return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @@ -408,7 +716,13 @@ def lulc_vector(request): start_year = request.data.get("start_year") end_year = request.data.get("end_year") gee_account_id = request.data.get("gee_account_id") - vectorise_lulc.apply_async( + compute = _get_compute_mode(request) + task = _select_compute_task( + compute, + vectorise_lulc_gee_task, + vectorise_lulc_local_task, + ) + task.apply_async( args=[state, district, block, start_year, end_year, gee_account_id], queue="nrm", ) @@ -416,6 +730,9 @@ def lulc_vector(request): {"Success": "lulc_vector task initiated"}, status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in lulc_vector api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: print("Exception in lulc_vector api :: ", e) return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @@ -472,7 +789,13 @@ def generate_ci_layer(request): start_year = request.data.get("start_year") end_year = request.data.get("end_year") gee_account_id = request.data.get("gee_account_id") - generate_cropping_intensity.apply_async( + compute = _get_compute_mode(request) + task = _select_compute_task( + compute, + generate_cropping_intensity, + generate_cropping_intensity_local_task, + ) + task.apply_async( kwargs={ "state": state, "district": district, @@ -487,6 +810,9 @@ def generate_ci_layer(request): {"Success": "Cropping Intensity task initiated"}, status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in generate_cropping_intensity_layer api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: print("Exception in generate_cropping_intensity_layer api :: ", e) return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @@ -562,18 +888,49 @@ def generate_terrain_descriptor(request): district = request.data.get("district") block = request.data.get("block") gee_account_id = request.data.get("gee_account_id") - generate_terrain_clusters.apply_async( - args=[state, district, block, gee_account_id], queue="nrm" + compute = _get_compute_mode(request) + task = _select_compute_task( + compute, + generate_terrain_clusters_gee_task, + generate_terrain_clusters_local_task, ) + task.apply_async(args=[state, district, block, gee_account_id], queue="nrm") return Response( {"Success": "generate_terrain_descriptor task initiated"}, status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in generate_terrain_descriptor api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: print("Exception in generate_terrain_descriptor api :: ", e) return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) +@api_view(["POST"]) +@schema(None) +def generate_terrain_compute_all(request): + print("Inside generate_terrain_compute_all") + try: + state = request.data.get("state") + district = request.data.get("district") + block = request.data.get("block") + start_year = request.data.get("start_year") + end_year = request.data.get("end_year") + gee_account_id = request.data.get("gee_account_id") + generate_terrain_compute_all_task.apply_async( + args=[state, district, block, start_year, end_year, gee_account_id], + queue="nrm", + ) + return Response( + {"Success": "generate_terrain_compute_all task initiated"}, + status=status.HTTP_200_OK, + ) + except Exception as e: + print("Exception in generate_terrain_compute_all api :: ", e) + return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) + + @api_view(["POST"]) @schema(None) def generate_terrain_raster(request): @@ -583,7 +940,13 @@ def generate_terrain_raster(request): district = request.data.get("district") block = request.data.get("block") gee_account_id = request.data.get("gee_account_id") - generate_terrain_raster_clip.apply_async( + compute = _get_compute_mode(request) + task = _select_compute_task( + compute, + generate_terrain_raster_clip_gee_task, + generate_terrain_raster_clip_local_task, + ) + task.apply_async( kwargs={ "state": state, "district": district, @@ -597,6 +960,9 @@ def generate_terrain_raster(request): {"Success": "generate_terrain_raster task initiated"}, status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in generate_terrain_raster api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: print("Exception in generate_terrain_raster api :: ", e) return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @@ -610,10 +976,16 @@ def terrain_lulc_slope_cluster(request): state = request.data.get("state") district = request.data.get("district") block = request.data.get("block") - start_year = request.data.get("start_year") - end_year = request.data.get("end_year") + 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") - lulc_on_slope_cluster.apply_async( + compute = _get_compute_mode(request) + task = _select_compute_task( + compute, + lulc_on_slope_cluster_gee_task, + lulc_on_slope_cluster_local_task, + ) + task.apply_async( args=[state, district, block, start_year, end_year, gee_account_id], queue="nrm", ) @@ -621,6 +993,9 @@ def terrain_lulc_slope_cluster(request): {"Success": "terrain_lulc_slope_cluster task initiated"}, status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in terrain_lulc_slope_cluster api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: print("Exception in terrain_lulc_slope_cluster api :: ", e) return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @@ -634,10 +1009,16 @@ def terrain_lulc_plain_cluster(request): state = request.data.get("state") district = request.data.get("district") block = request.data.get("block") - start_year = request.data.get("start_year") - end_year = request.data.get("end_year") + 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") - lulc_on_plain_cluster.apply_async( + compute = _get_compute_mode(request) + task = _select_compute_task( + compute, + lulc_on_plain_cluster_gee_task, + lulc_on_plain_cluster_local_task, + ) + task.apply_async( args=[state, district, block, start_year, end_year, gee_account_id], queue="nrm", ) @@ -645,6 +1026,9 @@ def terrain_lulc_plain_cluster(request): {"Success": "terrain_lulc_plain_cluster task initiated"}, status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in terrain_lulc_plain_cluster api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: print("Exception in terrain_lulc_plain_cluster api :: ", e) return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @@ -679,10 +1063,16 @@ def change_detection(request): state = request.data.get("state").lower() district = request.data.get("district").lower() block = request.data.get("block").lower() - start_year = request.data.get("start_year") - end_year = request.data.get("end_year") + 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") - get_change_detection.apply_async( + compute = _get_compute_mode(request) + task = _select_compute_task( + compute, + get_change_detection_gee_task, + get_change_detection_local_task, + ) + task.apply_async( args=[state, district, block, start_year, end_year, gee_account_id], queue="nrm", ) @@ -690,6 +1080,9 @@ def change_detection(request): {"Success": "change_detection task initiated"}, status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in change_detection api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: print("Exception in change_detection api :: ", e) return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @@ -703,10 +1096,16 @@ def change_detection_vector(request): state = request.data.get("state").lower() district = request.data.get("district").lower() block = request.data.get("block").lower() - start_year = request.data.get("start_year") - end_year = request.data.get("end_year") + 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") - vectorise_change_detection.apply_async( + compute = _get_compute_mode(request) + task = _select_compute_task( + compute, + vectorise_change_detection_gee_task, + vectorise_change_detection_local_task, + ) + task.apply_async( args=[state, district, block, start_year, end_year, gee_account_id], queue="nrm", ) @@ -714,6 +1113,9 @@ def change_detection_vector(request): {"Success": "change_detection_vector task initiated"}, status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in change_detection_vector api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: print("Exception in change_detection_vector api :: ", e) return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @@ -977,13 +1379,20 @@ def aquifer_vector(request): district = request.data.get("district").lower() block = request.data.get("block").lower() gee_account_id = request.data.get("gee_account_id") - generate_aquifer_vector.apply_async( - args=[state, district, block, gee_account_id], queue="nrm" + compute = _get_compute_mode(request) + task = _select_compute_task( + compute, + generate_aquifer_vector_gee_task, + generate_aquifer_vector_local_task, ) + task.apply_async(args=[state, district, block, gee_account_id], queue="nrm") return Response( {"Success": "aquifer vector task initiated"}, status=status.HTTP_200_OK, ) + except ValueError as e: + print("Invalid request in aquifer vector api :: ", e) + return Response({"Exception": str(e)}, status=status.HTTP_400_BAD_REQUEST) except Exception as e: print("Exception in aquifer vector api :: ", e) return Response({"Exception": e}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) @@ -1795,6 +2204,34 @@ def update_layer_sync_remote(request): @authentication_classes([]) @permission_classes([AllowAny]) @schema(None) +def update_layer_sync_remote(request): + """ + Called by a local compute instance to update sync/STAC flags on a layer + record in this (prod) backend. + """ + + api_key = getattr(settings, "PROD_BACKEND_API_KEY", "") + if api_key and request.headers.get("X-Api-Key") != api_key: + return Response({"error": "Unauthorized"}, status=status.HTTP_401_UNAUTHORIZED) + + try: + d = request.data + layer_id = d.get("layer_id") + if layer_id is None: + return Response({"error": "layer_id is required"}, status=status.HTTP_400_BAD_REQUEST) + + result = update_layer_sync_status( + layer_id=layer_id, + sync_to_geoserver=d.get("sync_to_geoserver"), + is_stac_specs_generated=d.get("is_stac_specs_generated"), + ) + return Response({"layer_id": result}, status=status.HTTP_200_OK) + except Exception as e: + return Response({"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR) + + +@api_view(["POST"]) +@schema(None) def sync_layer_remote(request): """ Called by a local compute instance to persist a layer record on this (prod) backend. diff --git a/computing/apps.py b/computing/apps.py index d1f2ec46..a73777ee 100644 --- a/computing/apps.py +++ b/computing/apps.py @@ -1,6 +1,10 @@ from django.apps import AppConfig +from computing.base_layer_setup import setup_base_layers class ComputingConfig(AppConfig): default_auto_field = "django.db.models.BigAutoField" name = "computing" + + def ready(self): + setup_base_layers() diff --git a/computing/base_layer_setup.py b/computing/base_layer_setup.py new file mode 100644 index 00000000..fffbedf0 --- /dev/null +++ b/computing/base_layer_setup.py @@ -0,0 +1,243 @@ +import logging +import subprocess + +import requests +from pathlib import Path +from utilities.constants import GEOSERVER_BASE + +from computing.config_loader import ( + ADMIN_BOUNDARY_INPUT_DIR, + ADMIN_BOUNDARY_OUTPUT_DIR, + GDRIVE_ADMIN_BOUNDARY_FILE_ID as _GDRIVE_ADMIN_BOUNDARY_FILE_ID, + GDRIVE_MICROWATERSHED_FILE_ID as _GDRIVE_MICROWATERSHED_FILE_ID, + LULC_BASE_DIR as LULC_DIR, + LULC_GDRIVE_FILES as _LULC_GDRIVE_FILES, + MICROWATERSHED_PATH, + PRECOMPUTED_TEHSIL_WATERSHED_DIR as TEHSIL_WATERSHEDS_DIR, + PROJECT_ROOT, + SOI_TEHSIL_PATH, + VILLAGE_BOUNDARIES_DIR, +) + +logger = logging.getLogger(__name__) + +_SOI_WFS_PARAMS = { + "service": "WFS", + "version": "1.0.0", + "request": "GetFeature", + "typeName": "pan_india_asset:SOI_tehsil_pan_india_dataset", + "outputFormat": "application/json", +} + + +def _is_dir_populated(path: Path) -> bool: + return path.is_dir() and any(path.iterdir()) + + +def ensure_soi_tehsil(): + """ + Downloads the SOI tehsil GeoJSON from GeoServer if not already present. + This is a lightweight bootstrap; the full admin-boundary archive includes + more data but takes much longer to acquire. + """ + if SOI_TEHSIL_PATH.exists(): + logger.info("SOI tehsil layer already exists at %s, skipping.", SOI_TEHSIL_PATH) + return + + SOI_TEHSIL_PATH.parent.mkdir(parents=True, exist_ok=True) + + wfs_url = f"{GEOSERVER_BASE}pan_india_asset/ows" + logger.info("Downloading SOI tehsil layer from GeoServer...") + try: + response = requests.get( + wfs_url, params=_SOI_WFS_PARAMS, timeout=600, stream=True + ) + response.raise_for_status() + except requests.RequestException as e: + logger.error("Failed to download SOI tehsil layer: %s", e) + return + + with open(SOI_TEHSIL_PATH, "wb") as f: + for chunk in response.iter_content(chunk_size=8192): + f.write(chunk) + + logger.info("SOI tehsil layer saved to %s", SOI_TEHSIL_PATH) + + +def ensure_admin_boundary_data(): + """ + Downloads and extracts the full admin-boundary archive (~8 GB) from Google Drive. + Skipped if the input directory is already populated. + Requires `gdown` and `7z` to be available on PATH. + """ + if _is_dir_populated(ADMIN_BOUNDARY_INPUT_DIR): + logger.info("Admin boundary data already exists, skipping.") + return + + ADMIN_BOUNDARY_INPUT_DIR.mkdir(parents=True, exist_ok=True) + ADMIN_BOUNDARY_OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + + archive_path = PROJECT_ROOT / "dataset.7z" + logger.info("Downloading admin boundary data (~8 GB) from Google Drive...") + try: + subprocess.run( + ["gdown", _GDRIVE_ADMIN_BOUNDARY_FILE_ID, "-O", str(archive_path)], + check=True, + ) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + logger.error("Failed to download admin boundary archive: %s", e) + return + + logger.info("Extracting admin boundary data...") + try: + subprocess.run( + ["7z", "x", str(archive_path), f"-o{PROJECT_ROOT / 'data/admin-boundary'}"], + check=True, + ) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + logger.error("Failed to extract admin boundary archive: %s", e) + return + finally: + if archive_path.exists(): + archive_path.unlink() + + logger.info("Admin boundary data ready at %s", ADMIN_BOUNDARY_INPUT_DIR) + + +def ensure_lulc_rasters(): + """ + Downloads any missing LULC v3 yearly rasters from Google Drive. + Files already present on disk are skipped — no re-download. + Requires `gdown` on PATH. Each file is ~7-8 GB. + """ + LULC_DIR.mkdir(parents=True, exist_ok=True) + + missing = [ + (filename, file_id) + for filename, file_id in _LULC_GDRIVE_FILES + if not (LULC_DIR / filename).exists() + ] + + if not missing: + logger.info("All LULC rasters already present at %s, skipping.", LULC_DIR) + return + + logger.info( + "%d LULC raster(s) missing, downloading: %s", + len(missing), + [f for f, _ in missing], + ) + + for filename, file_id in missing: + dest = LULC_DIR / filename + logger.info("Downloading %s (~7-8 GB)...", filename) + try: + subprocess.run( + ["gdown", file_id, "-O", str(dest)], + check=True, + ) + logger.info("Saved %s", dest) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + logger.error("Failed to download %s: %s", filename, e) + if dest.exists(): + dest.unlink() + raise + + +def ensure_microwatershed(): + """ + Downloads the pan-India microwatershed GeoJSON from Google Drive if not already present. + Requires `gdown` on PATH. + Fill in _GDRIVE_MICROWATERSHED_FILE_ID above once the Drive link is available. + """ + if MICROWATERSHED_PATH.exists(): + logger.info( + "Microwatershed file already exists at %s, skipping.", MICROWATERSHED_PATH + ) + return + + if not _GDRIVE_MICROWATERSHED_FILE_ID: + logger.warning( + "Microwatershed file not found at %s and no Google Drive file ID is configured. " + "Set _GDRIVE_MICROWATERSHED_FILE_ID in base_layer_setup.py or place the file manually.", + MICROWATERSHED_PATH, + ) + return + + MICROWATERSHED_PATH.parent.mkdir(parents=True, exist_ok=True) + logger.info( + "Downloading Microwatershed_v2_with_details.geojson from Google Drive..." + ) + try: + subprocess.run( + ["gdown", _GDRIVE_MICROWATERSHED_FILE_ID, "-O", str(MICROWATERSHED_PATH)], + check=True, + ) + logger.info("Saved microwatershed file to %s", MICROWATERSHED_PATH) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + logger.error("Failed to download microwatershed file: %s", e) + if MICROWATERSHED_PATH.exists(): + MICROWATERSHED_PATH.unlink() + raise + + +def ensure_tehsil_watersheds(): + """ + Generates per-tehsil watershed .gpkg files by spatially intersecting the + microwatershed dataset against SOI tehsil boundaries. + Skipped entirely if the output directory is already populated. + Both source files (SOI tehsil + microwatershed) must exist first. + """ + if _is_dir_populated(TEHSIL_WATERSHEDS_DIR): + logger.info( + "Tehsil watershed files already present at %s, skipping.", + TEHSIL_WATERSHEDS_DIR, + ) + return + + if not SOI_TEHSIL_PATH.exists(): + logger.warning( + "Cannot generate tehsil watersheds: SOI tehsil file missing at %s.", + SOI_TEHSIL_PATH, + ) + return + + if not MICROWATERSHED_PATH.exists(): + logger.warning( + "Cannot generate tehsil watersheds: microwatershed file missing at %s.", + MICROWATERSHED_PATH, + ) + return + + TEHSIL_WATERSHEDS_DIR.mkdir(parents=True, exist_ok=True) + logger.info("Generating tehsil watershed files (this may take a while)...") + + from computing.terrain_descriptor.store_watersheds_for_tehsils import ( + generate_tehsil_watershed_copies, + ) + + generate_tehsil_watershed_copies( + microwatershed_path=str(MICROWATERSHED_PATH), + tehsil_path=str(SOI_TEHSIL_PATH), + output_dir=str(TEHSIL_WATERSHEDS_DIR), + output_format="gpkg", + overwrite=False, + ) + logger.info("Tehsil watershed files ready at %s", TEHSIL_WATERSHEDS_DIR) + + +def ensure_village_boundaries_dir(): + """ + Ensures the village boundaries directory exists. + TODO: add download logic once the source is determined. + """ + VILLAGE_BOUNDARIES_DIR.mkdir(parents=True, exist_ok=True) + + +def setup_base_layers(): + ensure_soi_tehsil() + ensure_admin_boundary_data() + # ensure_lulc_rasters() + # ensure_microwatershed() + # ensure_tehsil_watersheds() + ensure_village_boundaries_dir() diff --git a/computing/change_detection/change_detection.py b/computing/change_detection/change_detection.py index 46bf27c2..6fd98148 100644 --- a/computing/change_detection/change_detection.py +++ b/computing/change_detection/change_detection.py @@ -1,5 +1,4 @@ import ee -import copy from utilities.gee_utils import ( ee_initialize, check_task_status, @@ -117,6 +116,20 @@ def get_change_detection( return layer_at_geoserver +def _compute_then_now_modes(l1_asset_remapped, lulc_projection, roi_boundary): + if len(l1_asset_remapped) < 6: + raise ValueError( + "Change detection requires at least six yearly LULC rasters to compare the first three years against the last three years." + ) + + then = ee.ImageCollection(l1_asset_remapped[:3]).mode().reproject(lulc_projection) + now = ee.ImageCollection(l1_asset_remapped[-3:]).mode().reproject(lulc_projection) + + then = then.clip(roi_boundary.geometry()) + now = now.clip(roi_boundary.geometry()) + return now, then + + def built_up(roi_boundary, l1_asset): print("built_up function is runing") @@ -133,13 +146,9 @@ def remap_values(image): l1_asset_remapped = [remap_values(asset) for asset in l1_asset] - # Create image collections - then = ee.ImageCollection(l1_asset_remapped[:3]).mode().reproject(lulc_projection) - now = ee.ImageCollection(l1_asset_remapped[3:]).mode().reproject(lulc_projection) - - # Compute mode and clip - then = then.clip(roi_boundary.geometry()) - now = now.clip(roi_boundary.geometry()) + now, then = _compute_then_now_modes( + l1_asset_remapped, lulc_projection, roi_boundary + ) # Compute transitions trans_bu_bu = then.eq(1).And(now.eq(1)) @@ -173,13 +182,9 @@ def remap_values(image): l1_asset_remapped = [remap_values(asset) for asset in l1_asset] - # Create image collections - then = ee.ImageCollection(l1_asset_remapped[:3]).mode().reproject(lulc_projection) - now = ee.ImageCollection(l1_asset_remapped[3:]).mode().reproject(lulc_projection) - - # Compute mode and clip - then = then.clip(roi_boundary.geometry()) - now = now.clip(roi_boundary.geometry()) + now, then = _compute_then_now_modes( + l1_asset_remapped, lulc_projection, roi_boundary + ) trans_f_f = then.eq(3).And(now.eq(3)) trans_f_bu = then.eq(3).And(now.eq(1)).multiply(2) @@ -198,160 +203,23 @@ def remap_values(image): return change_deg -def change_deforestation_afforestation(roi_boundary, l1_asset, lulc_projection): - print("change_deforestation is running") - # Create an initial zero image - zero_image2 = ( - ee.Image.constant(0) - .setDefaultProjection(lulc_projection) - .clip(l1_asset[0].geometry()) - ) - - # for i in range(1, 5): - for i in range(1, len(l1_asset) - 1): - before = l1_asset[i - 1] - middle = l1_asset[i] - after = l1_asset[i + 1] - - cond1 = ( - before.eq(12) - .And(after.eq(12)) - .And( - middle.eq(6) - .Or(middle.eq(8)) - .Or(middle.eq(9)) - .Or(middle.eq(10)) - .Or(middle.eq(11)) - ) - ) - cond2 = ( - before.eq(2) - .Or(before.eq(3)) - .Or(before.eq(4)) - .And(after.eq(2).Or(after.eq(3)).Or(after.eq(4))) - .And( - middle.eq(6) - .Or(middle.eq(8)) - .Or(middle.eq(9)) - .Or(middle.eq(10)) - .Or(middle.eq(11)) - ) - ) - cond3 = before.eq(6).And(after.eq(6)).And(middle.eq(12)) - cond4 = ( - before.eq(8) - .Or(before.eq(9)) - .Or(before.eq(10)) - .Or(before.eq(11)) - .And(after.eq(8).Or(after.eq(9)).Or(after.eq(10)).Or(after.eq(11))) - .And(middle.eq(12)) - ) - cond5 = ( - before.eq(8) - .Or(before.eq(9)) - .Or(before.eq(10)) - .Or(before.eq(11)) - .And(after.eq(8).Or(after.eq(9)).Or(after.eq(10)).Or(after.eq(11))) - .And(middle.eq(7)) - ) - cond6 = ( - before.eq(6) - .And(after.eq(6)) - .And(middle.eq(8).Or(middle.eq(9)).Or(middle.eq(10)).Or(middle.eq(11))) - ) - cond7 = ( - before.eq(8) - .Or(before.eq(9)) - .Or(before.eq(10)) - .Or(before.eq(11)) - .And(after.eq(8).Or(after.eq(9)).Or(after.eq(10)).Or(after.eq(11))) - .And(middle.eq(6)) - ) - cond8 = before.eq(1).And(after.eq(1)).And(middle.eq(6)) - cond9 = before.eq(6).And(after.eq(6)).And(middle.eq(1)) - cond10 = ( - before.eq(1) - .And(after.eq(1)) - .And(middle.eq(8).Or(middle.eq(9)).Or(middle.eq(10)).Or(middle.eq(11))) - ) - cond11 = ( - before.eq(7) - .And(after.eq(7)) - .And( - middle.eq(6) - .Or(middle.eq(8)) - .Or(middle.eq(9)) - .Or(middle.eq(10)) - .Or(middle.eq(11)) - ) - ) - - zero_image2 = ( - zero_image2.add(cond1) - .add(cond2) - .add(cond3) - .add(cond4) - .add(cond5) - .add(cond6) - .add(cond7) - .add(cond8) - .add(cond9) - .add(cond10) - .add(cond11) - ) - - l1_asset_copy = copy.deepcopy(l1_asset) - for i in range(1, len(l1_asset) - 1): - # for i in range(1, 5): - before = l1_asset[i - 1] - middle = l1_asset[i] - after = l1_asset[i + 1] - - cond1 = ( - before.eq(3) - .And(middle.neq(3)) - .And(after.eq(3)) - .And((zero_image2.eq(3).Or(zero_image2.eq(4)))) - ) - cond2 = ( - before.neq(3) - .And(middle.eq(3)) - .And(after.neq(3)) - .And((zero_image2.eq(3).Or(zero_image2.eq(4)))) - ) - - middle = middle.where(cond1, 3) - middle = middle.where(cond2, before) - - l1_asset_copy[i] = middle +def change_deforestation(roi_boundary, l1_asset): + lulc_projection = l1_asset[0].projection() - # Remap values function def remap_values(image): - remapped = image.remap( + return image.remap( [1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12], [1, 2, 2, 2, 3, 5, 4, 4, 4, 4, 6], 0, "predicted_label", ).setDefaultProjection(lulc_projection) - return remapped - - l1_asset_remapped = [remap_values(asset) for asset in l1_asset_copy] - - # Create image collections - then = ee.ImageCollection(l1_asset_remapped[:3]).mode().reproject(lulc_projection) - now = ee.ImageCollection(l1_asset_remapped[3:]).mode().reproject(lulc_projection) - - # Compute mode and clip - then = then.clip(roi_boundary.geometry()) - now = now.clip(roi_boundary.geometry()) - return now, then + l1_asset_remapped = [remap_values(asset) for asset in l1_asset] -def change_deforestation(roi_boundary, l1_asset): - lulc_projection = l1_asset[0].projection() - now, then = change_deforestation_afforestation( - roi_boundary, l1_asset, lulc_projection + now, then = _compute_then_now_modes( + l1_asset_remapped, lulc_projection, roi_boundary ) + trans_fo_fo = then.eq(3).And(now.eq(3)) trans_fo_bu = then.eq(3).And(now.eq(1)).multiply(2) trans_fo_fa = then.eq(3).And(now.eq(4)).multiply(3) @@ -375,9 +243,21 @@ def change_deforestation(roi_boundary, l1_asset): def change_afforestation(roi_boundary, l1_asset): lulc_projection = l1_asset[0].projection() - now, then = change_deforestation_afforestation( - roi_boundary, l1_asset, lulc_projection + + def remap_values(image): + return image.remap( + [1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12], + [1, 2, 2, 2, 3, 5, 4, 4, 4, 4, 6], + 0, + "predicted_label", + ).setDefaultProjection(lulc_projection) + + l1_asset_remapped = [remap_values(asset) for asset in l1_asset] + + now, then = _compute_then_now_modes( + l1_asset_remapped, lulc_projection, roi_boundary ) + trans_fo_fo = then.eq(3).And(now.eq(3)) trans_bu_fo = then.eq(1).And(now.eq(3)).multiply(2) trans_fa_fo = then.eq(4).And(now.eq(3)).multiply(3) @@ -414,13 +294,9 @@ def remap_values(image): l1_asset_remapped = [remap_values(asset) for asset in l1_asset] - # Create image collections - then = ee.ImageCollection(l1_asset_remapped[:3]).mode().reproject(lulc_projection) - now = ee.ImageCollection(l1_asset_remapped[3:]).mode().reproject(lulc_projection) - - # Compute mode and clip - then = then.clip(roi_boundary.geometry()) - now = now.clip(roi_boundary.geometry()) + now, then = _compute_then_now_modes( + l1_asset_remapped, lulc_projection, roi_boundary + ) trans_do_si = then.eq(6).And(now.eq(5)) trans_tr_si = then.eq(7).And(now.eq(5)).multiply(2) diff --git a/computing/change_detection/change_detection_local.py b/computing/change_detection/change_detection_local.py new file mode 100644 index 00000000..7ad4dbba --- /dev/null +++ b/computing/change_detection/change_detection_local.py @@ -0,0 +1,599 @@ +import os +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import rasterio +from rasterio.mask import mask +from rasterio.warp import Resampling, reproject +from shapely.geometry import mapping +from utilities.gee_utils import valid_gee_text + +from nrm_app.celery import app + +from computing.config_loader import CHANGE_DETECTION_RASTER_OUTPUT_DIR as LOCAL_OUTPUT_BASE_DIR +from computing.local_compute_helper import ( + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + build_output_raster_path, + get_union_geometry, + load_precomputed_roi, + push_local_raster_to_geoserver, + resolve_lulc_raster_paths, + validate_geometry, +) +from computing.utils import save_layer_info_to_db, update_layer_sync_status +GEOSERVER_WORKSPACE = "change_detection" + +CHANGE_STAC_LAYER_NAMES = { + "Urbanization": "change_urbanization_raster", + "Degradation": "change_cropping_reduction_raster", + "Deforestation": "change_tree_cover_loss_raster", + "Afforestation": "change_tree_cover_gain_raster", + "CropIntensity": "change_cropping_intensity_raster", +} + +BUILT_UP_REMAP = { + 1: 1, + 2: 2, + 3: 2, + 4: 2, + 6: 3, + 7: 4, + 8: 3, + 9: 3, + 10: 3, + 11: 3, + 12: 4, +} + +DEGRADATION_REMAP = { + 1: 1, + 2: 2, + 3: 2, + 4: 2, + 6: 4, + 7: 5, + 8: 3, + 9: 3, + 10: 3, + 11: 3, + 12: 6, +} + +DEFORESTATION_AFFORESTATION_REMAP = { + 1: 1, + 2: 2, + 3: 2, + 4: 2, + 6: 3, + 7: 5, + 8: 4, + 9: 4, + 10: 4, + 11: 4, + 12: 6, +} + +CROP_INTENSITY_REMAP = { + 1: 1, + 2: 2, + 3: 2, + 4: 2, + 6: 3, + 7: 4, + 8: 5, + 9: 5, + 10: 6, + 11: 7, + 12: 8, +} + + +def _build_lookup_table(mapping, size=13): + lookup = np.zeros(size, dtype=np.int16) + for source_value, mapped_value in mapping.items(): + lookup[source_value] = mapped_value + return lookup + + +BUILT_UP_LOOKUP = _build_lookup_table(BUILT_UP_REMAP) +DEGRADATION_LOOKUP = _build_lookup_table(DEGRADATION_REMAP) +DEFORESTATION_AFFORESTATION_LOOKUP = _build_lookup_table( + DEFORESTATION_AFFORESTATION_REMAP +) +CROP_INTENSITY_LOOKUP = _build_lookup_table(CROP_INTENSITY_REMAP) + +CHANGE_PARAM_FUNCTIONS = { + "Urbanization": "_compute_built_up_change", + "Degradation": "_compute_degradation_change", + "Deforestation": "_compute_deforestation_change", + "Afforestation": "_compute_afforestation_change", + "CropIntensity": "_compute_crop_intensity_change", +} + +ZERO_NODATA = 0 + + +def _slug(value, fallback): + return valid_gee_text(str(value).strip().lower()) or fallback + + +def _as_lulc_int_array(array): + array = np.asarray(array) + if np.issubdtype(array.dtype, np.integer): + return array.astype(np.int16, copy=False) + + array = np.asarray(array, dtype=np.float64) + array = np.where(np.isfinite(array), array, 0) + return np.rint(array).astype(np.int16, copy=False) + + +def _remap_array(array, mapping): + source = _as_lulc_int_array(array) + if isinstance(mapping, np.ndarray): + return mapping[source] + + remapped = np.zeros(source.shape, dtype=np.int16) + for source_value, mapped_value in mapping.items(): + remapped[source == source_value] = mapped_value + return remapped + + +def _combine_transitions(shape, transitions): + result = np.zeros(shape, dtype=np.uint8) + for value, condition in transitions: + result += condition.astype(np.uint8) * np.uint8(value) + return result + + +def _base_description(district, block): + return f"change_{_slug(district, 'unknown_district')}_{_slug(block, 'unknown_block')}" + + +def _published_layer_name(district, block, param_name): + return f"{_base_description(district, block)}_{param_name}" + + +def _output_stub(district, block, param_name, start_year, end_year): + return f"{_published_layer_name(district, block, param_name)}_{start_year}_{end_year}" + + +def _select_change_detection_raster_paths(start_year, end_year): + raster_paths = resolve_lulc_raster_paths(start_year=start_year, end_year=end_year) + if len(raster_paths) < 6: + raise ValueError( + "Local change detection requires at least six yearly LULC rasters to compare the first three years against the last three years." + ) + return raster_paths[:3] + raster_paths[-3:] + + +def _build_roi_shapes_by_crs(roi_gdf, raster_paths): + roi_shapes_by_crs = {} + roi_union = get_union_geometry(roi_gdf) + if roi_union is None or roi_union.is_empty: + raise ValueError("ROI union geometry is empty for local change detection.") + + roi_shapes_by_crs[str(roi_gdf.crs)] = mapping(roi_union) + for raster_path in raster_paths: + with rasterio.open(raster_path) as src: + crs_key = str(src.crs) + if crs_key in roi_shapes_by_crs: + continue + clip_gdf = roi_gdf if src.crs is None else roi_gdf.to_crs(src.crs) + clip_union = get_union_geometry(clip_gdf) + if clip_union is None or clip_union.is_empty: + raise ValueError("ROI union geometry is empty for local change detection.") + roi_shapes_by_crs[crs_key] = mapping(clip_union) + + return roi_shapes_by_crs + + +def _load_masked_lulc_array(index, raster_path, roi_shape_by_crs): + with rasterio.open(raster_path) as src: + clipped_data, clipped_transform = mask( + src, + shapes=[roi_shape_by_crs[str(src.crs)]], + crop=True, + filled=True, + nodata=ZERO_NODATA, + ) + clipped_array = _as_lulc_int_array(clipped_data[0]) + meta = src.meta.copy() + return { + "index": index, + "raster_path": raster_path, + "array": clipped_array, + "transform": clipped_transform, + "crs": src.crs, + "meta": meta, + } + + +def _load_masked_lulc_arrays(roi_gdf, raster_paths): + roi_gdf = validate_geometry(roi_gdf) + if roi_gdf.empty: + raise ValueError("No valid ROI geometry available for local change detection.") + if roi_gdf.crs is None: + raise ValueError("ROI CRS is missing; cannot align LULC rasters.") + + roi_shapes_by_crs = _build_roi_shapes_by_crs(roi_gdf, raster_paths) + max_workers = min(len(raster_paths), max(os.cpu_count() or 1, 1)) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit( + _load_masked_lulc_array, index, raster_path, roi_shapes_by_crs + ) + for index, raster_path in enumerate(raster_paths, start=1) + ] + results = [future.result() for future in futures] + + results.sort(key=lambda item: item["index"]) + + reference = results[0] + output_meta = reference["meta"] + output_meta.update( + { + "driver": "GTiff", + "height": reference["array"].shape[0], + "width": reference["array"].shape[1], + "transform": reference["transform"], + "crs": reference["crs"], + "count": 1, + "dtype": "uint8", + "nodata": ZERO_NODATA, + "compress": "lzw", + } + ) + + arrays = [] + for result in results: + clipped_array = result["array"] + needs_alignment = ( + clipped_array.shape != (output_meta["height"], output_meta["width"]) + or result["transform"] != output_meta["transform"] + or result["crs"] != output_meta["crs"] + ) + if needs_alignment: + aligned_array = np.zeros( + (output_meta["height"], output_meta["width"]), + dtype=np.int16, + ) + reproject( + source=clipped_array, + destination=aligned_array, + source_transform=result["transform"], + source_crs=result["crs"], + destination_transform=output_meta["transform"], + destination_crs=output_meta["crs"], + src_nodata=ZERO_NODATA, + dst_nodata=ZERO_NODATA, + resampling=Resampling.nearest, + ) + clipped_array = aligned_array + + arrays.append(clipped_array.astype(np.int16, copy=False)) + print( + f"Loaded local LULC raster {result['index']}/{len(raster_paths)}: {result['raster_path']}" + ) + + return arrays, output_meta + + +def _write_change_raster(array, output_path, output_meta): + raster = np.asarray(array, dtype=np.uint8) + meta = output_meta.copy() + meta.update( + { + "driver": "GTiff", + "count": 1, + "dtype": "uint8", + "nodata": ZERO_NODATA, + "compress": "lzw", + } + ) + with rasterio.open(output_path, "w", **meta) as dst: + dst.write(raster, 1) + return str(output_path) + + +def _compute_mode_lulc_three(arrays): + first, second, third = (_as_lulc_int_array(array) for array in arrays) + mode = np.zeros(first.shape, dtype=np.int16) + + same_first = ((first == second) | (first == third)) & (first > 0) + same_second = (second == third) & (second > 0) + mode = np.where(same_first, first, mode) + mode = np.where((mode == 0) & same_second, second, mode) + + unresolved = mode == 0 + if np.any(unresolved): + sentinel = np.int16(np.iinfo(np.int16).max) + positive_stack = np.stack( + [ + np.where(first > 0, first, sentinel), + np.where(second > 0, second, sentinel), + np.where(third > 0, third, sentinel), + ], + axis=0, + ) + fallback = np.min(positive_stack, axis=0) + fallback = np.where(fallback == sentinel, 0, fallback) + mode = np.where(unresolved, fallback, mode) + + return mode + + +def _compute_then_now_modes(remapped_arrays): + if len(remapped_arrays) < 6: + raise ValueError( + "Local change detection requires at least six selected yearly LULC rasters to compare the first three years against the last three years." + ) + then = _compute_mode_lulc_three(remapped_arrays[:3]) + now = _compute_mode_lulc_three(remapped_arrays[-3:]) + return now, then + + +def _compute_built_up_change(lulc_arrays): + remapped_arrays = [_remap_array(array, BUILT_UP_LOOKUP) for array in lulc_arrays] + now, then = _compute_then_now_modes(remapped_arrays) + return _combine_transitions( + then.shape, + [ + (1, (then == 1) & (now == 1)), + (2, (then == 2) & (now == 1)), + (3, (then == 3) & (now == 1)), + (4, (then == 4) & (now == 1)), + ], + ) + + +def _compute_degradation_change(lulc_arrays): + remapped_arrays = [ + _remap_array(array, DEGRADATION_LOOKUP) for array in lulc_arrays + ] + now, then = _compute_then_now_modes(remapped_arrays) + return _combine_transitions( + then.shape, + [ + (1, (then == 3) & (now == 3)), + (2, (then == 3) & (now == 1)), + (3, (then == 3) & (now == 5)), + (4, (then == 3) & (now == 6)), + ], + ) + + +def _compute_deforestation_afforestation_modes(lulc_arrays): + remapped_arrays = [ + _remap_array(array, DEFORESTATION_AFFORESTATION_LOOKUP) + for array in lulc_arrays + ] + return _compute_then_now_modes(remapped_arrays) + + +def _build_deforestation_change(now, then): + return _combine_transitions( + then.shape, + [ + (1, (then == 3) & (now == 3)), + (2, (then == 3) & (now == 1)), + (3, (then == 3) & (now == 4)), + (4, (then == 3) & (now == 5)), + (5, (then == 3) & (now == 6)), + ], + ) + + +def _compute_deforestation_change(lulc_arrays): + now, then = _compute_deforestation_afforestation_modes(lulc_arrays) + return _build_deforestation_change(now, then) + + +def _build_afforestation_change(now, then): + return _combine_transitions( + then.shape, + [ + (1, (then == 3) & (now == 3)), + (2, (then == 1) & (now == 3)), + (3, (then == 4) & (now == 3)), + (4, (then == 5) & (now == 3)), + (5, (then == 6) & (now == 3)), + ], + ) + + +def _compute_afforestation_change(lulc_arrays): + now, then = _compute_deforestation_afforestation_modes(lulc_arrays) + return _build_afforestation_change(now, then) + + +def _compute_crop_intensity_change(lulc_arrays): + remapped_arrays = [ + _remap_array(array, CROP_INTENSITY_LOOKUP) for array in lulc_arrays + ] + now, then = _compute_then_now_modes(remapped_arrays) + return _combine_transitions( + then.shape, + [ + (1, (then == 6) & (now == 5)), + (2, (then == 7) & (now == 5)), + (3, (then == 7) & (now == 6)), + (4, (then == 5) & (now == 6)), + (5, (then == 5) & (now == 7)), + (6, (then == 6) & (now == 7)), + (7, (then == 5) & (now == 5)), + (8, (then == 6) & (now == 6)), + (9, (then == 7) & (now == 7)), + ], + ) + + +def _compute_single_change_output(param_name, lulc_arrays): + print(f"Computing local change detection raster: {param_name}") + return param_name, globals()[CHANGE_PARAM_FUNCTIONS[param_name]](lulc_arrays) + + +def _compute_forest_change_outputs(lulc_arrays): + print("Computing local change detection raster: Deforestation") + print("Computing local change detection raster: Afforestation") + now, then = _compute_deforestation_afforestation_modes(lulc_arrays) + return { + "Deforestation": _build_deforestation_change(now, then), + "Afforestation": _build_afforestation_change(now, then), + } + + +def _compute_change_outputs(lulc_arrays): + outputs = {} + cpu_workers = min(max(os.cpu_count() or 1, 1), 4) + with ThreadPoolExecutor(max_workers=cpu_workers) as executor: + futures = [ + executor.submit(_compute_single_change_output, param_name, lulc_arrays) + for param_name in ("Urbanization", "Degradation", "CropIntensity") + ] + futures.append(executor.submit(_compute_forest_change_outputs, lulc_arrays)) + + for future in futures: + result = future.result() + if isinstance(result, dict): + outputs.update(result) + else: + param_name, change_array = result + outputs[param_name] = change_array + + return outputs + + +def run_change_detection_local( + state, + district, + block, + start_year, + end_year, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + push_to_geoserver=True, + sync_layer_metadata=True, +): + state = str(state).strip().lower() + district = str(district).strip().lower() + block = str(block).strip().lower() + 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") + + roi_gdf = load_precomputed_roi( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + ) + lulc_raster_paths = _select_change_detection_raster_paths( + start_year=start_year, + end_year=end_year, + ) + lulc_arrays, output_meta = _load_masked_lulc_arrays(roi_gdf, lulc_raster_paths) + change_outputs = _compute_change_outputs(lulc_arrays) + + geoserver_statuses = [] + + for param_name, change_array in change_outputs.items(): + output_stub = _output_stub(district, block, param_name, start_year, end_year) + output_path = build_output_raster_path( + layer_name=output_stub, + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + state=state, + district=district, + block=block, + block_fallback="unknown_block", + ) + raster_path = _write_change_raster( + array=change_array, + output_path=output_path, + output_meta=output_meta, + ) + print(f"Saved local change detection raster: {raster_path}") + + published_layer_name = _published_layer_name(district, block, param_name) + if push_to_geoserver: + upload_res, style_res = push_local_raster_to_geoserver( + file_path=raster_path, + layer_name=published_layer_name, + workspace=GEOSERVER_WORKSPACE, + style_name=param_name.lower(), + ) + print(f"GeoServer upload response for {param_name}: {upload_res}") + print(f"GeoServer style response for {param_name}: {style_res}") + geoserver_statuses.append(True) + + if sync_layer_metadata: + layer_id = save_layer_info_to_db( + state=state, + district=district, + block=block, + layer_name=published_layer_name, + asset_id=raster_path, + dataset_name="Change Detection Raster", + misc={ + "start_year": start_year, + "end_year": end_year, + "is_generated_locally": True, + }, + ) + if layer_id and push_to_geoserver: + update_layer_sync_status(layer_id=layer_id, sync_to_geoserver=True) + from computing.STAC_specs import generate_STAC_layerwise + + layer_stac_generated = generate_STAC_layerwise.generate_raster_stac( + state=state, + district=district, + block=block, + layer_name=CHANGE_STAC_LAYER_NAMES[param_name], + ) + update_layer_sync_status( + layer_id=layer_id, + is_stac_specs_generated=layer_stac_generated, + ) + + return all(geoserver_statuses) if push_to_geoserver else True + + +def _get_change_detection_local_task( + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = gee_account_id + return run_change_detection_local( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + +@app.task(bind=True) +def get_change_detection( + self, + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = self + return _get_change_detection_local_task( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + gee_account_id=gee_account_id, + ) diff --git a/computing/change_detection/change_detection_vector_local.py b/computing/change_detection/change_detection_vector_local.py new file mode 100644 index 00000000..669b8c11 --- /dev/null +++ b/computing/change_detection/change_detection_vector_local.py @@ -0,0 +1,230 @@ +import os + +from nrm_app.celery import app +from utilities.gee_utils import valid_gee_text + +from computing.config_loader import ( + CHANGE_DETECTION_RASTER_OUTPUT_DIR as CHANGE_RASTER_OUTPUT_BASE_DIR, + CHANGE_DETECTION_VECTOR_OUTPUT_DIR as LOCAL_OUTPUT_BASE_DIR, +) +from computing.local_compute_helper import ( + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + build_output_raster_path, + build_output_vector_path, + compute_categorical_raster_areas_for_watersheds, + ensure_file_exists, + load_precomputed_watersheds, + push_local_vector_to_geoserver, + write_vector_output, +) +from computing.utils import ( + save_layer_info_to_db, + update_layer_sync_status, +) +GEOSERVER_WORKSPACE = "change_detection" + +CHANGE_VECTOR_CLASS_DEFINITIONS = { + "Afforestation": [ + {"value": 1, "label": "fo_fo"}, + {"value": 2, "label": "bu_fo"}, + {"value": 3, "label": "fa_fo"}, + {"value": 4, "label": "ba_fo"}, + {"value": 5, "label": "sc_fo"}, + {"value": [2, 3, 4, 5], "label": "total_aff"}, + ], + "Deforestation": [ + {"value": 1, "label": "fo_fo"}, + {"value": 2, "label": "fo_bu"}, + {"value": 3, "label": "fo_fa"}, + {"value": 4, "label": "fo_ba"}, + {"value": 5, "label": "fo_sc"}, + {"value": [2, 3, 4, 5], "label": "total_def"}, + ], + "Degradation": [ + {"value": 1, "label": "f_f"}, + {"value": 2, "label": "f_bu"}, + {"value": 3, "label": "f_ba"}, + {"value": 4, "label": "f_sc"}, + {"value": [2, 3, 4], "label": "total_deg"}, + ], + "Urbanization": [ + {"value": 1, "label": "bu_bu"}, + {"value": 2, "label": "w_bu"}, + {"value": 3, "label": "tr_bu"}, + {"value": 4, "label": "b_bu"}, + {"value": [2, 3, 4], "label": "total_urb"}, + ], + "CropIntensity": [ + {"value": 1, "label": "do_si"}, + {"value": 2, "label": "tr_si"}, + {"value": 3, "label": "tr_do"}, + {"value": 4, "label": "si_do"}, + {"value": 5, "label": "si_tr"}, + {"value": 6, "label": "do_tr"}, + {"value": 7, "label": "si_si"}, + {"value": 8, "label": "do_do"}, + {"value": 9, "label": "tr_tr"}, + {"value": [1, 2, 3, 4, 5, 6], "label": "total_change"}, + ], +} + + +def _slug(value, fallback): + return valid_gee_text(str(value).strip().lower()) or fallback + + +def _published_layer_name(district, block, param_name): + return ( + f"change_vector_{_slug(district, 'unknown_district')}_" + f"{_slug(block, 'unknown_block')}_{param_name}" + ) + + +def _output_stub(district, block, param_name, start_year, end_year): + return f"{_published_layer_name(district, block, param_name)}_{start_year}_{end_year}" + + +def _resolve_local_change_raster_path(state, district, block, param_name, start_year, end_year): + raster_path = build_output_raster_path( + layer_name=( + f"change_{_slug(district, 'unknown_district')}_" + f"{_slug(block, 'unknown_block')}_{param_name}_{int(start_year)}_{int(end_year)}" + ), + output_base_dir=CHANGE_RASTER_OUTPUT_BASE_DIR, + state=state, + district=district, + block=block, + block_fallback="unknown_block", + ) + ensure_file_exists(raster_path, f"Local change detection raster for {param_name}") + return str(raster_path) + + +def run_change_detection_vector_local( + state, + district, + block, + start_year, + end_year, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + push_to_geoserver=True, + sync_layer_metadata=True, +): + state = str(state).strip().lower() + district = str(district).strip().lower() + block = str(block).strip().lower() + start_year = int(start_year) + end_year = int(end_year) + + watersheds_gdf, watershed_source = load_precomputed_watersheds( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + ) + print(f"Watershed boundary source: {watershed_source}") + + geoserver_statuses = [] + + for param_name, class_definitions in CHANGE_VECTOR_CLASS_DEFINITIONS.items(): + raster_path = _resolve_local_change_raster_path( + state=state, + district=district, + block=block, + param_name=param_name, + start_year=start_year, + end_year=end_year, + ) + print(f"Using local change raster for {param_name}: {raster_path}") + + result_gdf = compute_categorical_raster_areas_for_watersheds( + watersheds_gdf=watersheds_gdf, + raster_path=raster_path, + class_definitions=class_definitions, + ) + + published_layer_name = _published_layer_name(district, block, param_name) + output_stub = _output_stub(district, block, param_name, start_year, end_year) + output_path = build_output_vector_path( + layer_name=output_stub, + state=state, + district=district, + block=block, + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + block_fallback="unknown_block", + ) + asset_id = write_vector_output( + gdf=result_gdf, + output_path=output_path, + layer_name=published_layer_name, + ) + print(f"Saved local change detection vector: {asset_id}") + if push_to_geoserver: + geoserver_response = push_local_vector_to_geoserver( + path=os.path.splitext(asset_id)[0], + layer_name=published_layer_name, + workspace=GEOSERVER_WORKSPACE, + file_type="gpkg", + ) + print(f"GeoServer response for {param_name}: {geoserver_response}") + if not isinstance(geoserver_response, dict) or geoserver_response.get( + "status_code" + ) not in (200, 201): + return False + geoserver_statuses.append(True) + + if sync_layer_metadata: + layer_id = save_layer_info_to_db( + state=state, + district=district, + block=block, + layer_name=published_layer_name, + asset_id=asset_id, + dataset_name="Change Detection Vector", + misc={"is_generated_locally": True}, + ) + if layer_id and push_to_geoserver: + update_layer_sync_status(layer_id=layer_id, sync_to_geoserver=True) + + return all(geoserver_statuses) if push_to_geoserver else True + + +def _vectorise_change_detection_local_task( + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = gee_account_id + return run_change_detection_vector_local( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + +@app.task(bind=True) +def vectorise_change_detection( + self, + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = self + return _vectorise_change_detection_local_task( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + gee_account_id=gee_account_id, + ) diff --git a/computing/config.yaml b/computing/config.yaml new file mode 100644 index 00000000..0b5ee911 --- /dev/null +++ b/computing/config.yaml @@ -0,0 +1,126 @@ +# Local compute dependency manifest. + +base_layers: + inputs: + - path: data/base_layers/lulc/lulc_v3_2017_2018.tif + source: google_drive + gdrive_id: 1VidwEQqkwtoHqqdUqdwURWyiGd-OteaJ + - path: data/base_layers/lulc/lulc_v3_2018_2019.tif + source: google_drive + gdrive_id: 1ZeLMAiBfolMrfEJkFnOlvb8OjSqC9vHP + - path: data/base_layers/lulc/lulc_v3_2019_2020.tif + source: google_drive + gdrive_id: 1gx5VwJCHI-WUDJIwWv48OvbybBe9y0PR + - path: data/base_layers/lulc/lulc_v3_2020_2021.tif + source: google_drive + gdrive_id: 1xbOt3-t1Ws5olq2Q88Tk32KnUVNKUXqe + - path: data/base_layers/lulc/lulc_v3_2021_2022.tif + source: google_drive + gdrive_id: 1m8ZnUBbTp-fcH_JcRTUEceRaa8WewQmz + - path: data/base_layers/lulc/lulc_v3_2022_2023.tif + source: google_drive + gdrive_id: 1_S0VESClg7s-DloAqxrfU8mLHhNSBfp7 + - path: data/base_layers/lulc/lulc_v3_2023_2024.tif + source: google_drive + gdrive_id: 1JVfl67ARRv7TPV5lyLnjoSfWDiXtvXjY + - path: data/base_layers/lulc/lulc_v3_2024_2025.tif + source: google_drive + gdrive_id: 1CPV03S47s0asEJqdAozbNOT1lkgr0YkG + + - path: data/base_layers/slope/slope_india_30m_merged.tif + source: manual + note: Pan-India slope raster used for slope/terrain classification + + - path: data/base_layers/soil/hysogs_india_250m_4326.tif + source: manual + note: Pan-India hydrologic soil group raster for runoff calculation + + - path: data/base_layers/AEZs/Agro_Ecological_Regions.shp + source: manual + note: Agro-Ecological Zone boundaries for AEZ-based cluster lookup + + - path: data/base_layers/Microwatershed_v2_with_details.geojson + source: google_drive + gdrive_id: 1gwaCrCXq1t3fWyfqcQMw1e7DTlzxoCb3 + + - path: data/base_layers/Aquifer_vector.geojson + source: manual + note: Pan-India aquifer vector layer + + - path: data/admin-boundary/input/ + source: google_drive + gdrive_id: 1VqIhB6HrKFDkDnlk1vedcEHhh5fk4f1d + note: Full admin-boundary archive (~8 GB, 7z); extracted in place + + - path: data/admin-boundary/input/soi_tehsil.geojson + source: geoserver_wfs + typename: pan_india_asset:SOI_tehsil_pan_india_dataset + note: Lightweight bootstrap; also included in the full admin-boundary archive + + - path: data/base_layers/tehsil_watersheds/ + source: derived + derived_from: + - data/base_layers/Microwatershed_v2_with_details.geojson + - data/admin-boundary/input/soi_tehsil.geojson + note: Per-tehsil watershed .gpkg files; generated by store_watersheds_for_tehsils + + - path: data/base_layers/village_boundaries/ + source: manual + note: Village boundary files; download source TBD + +local_compute_outputs: + change_detection: + # change_detection_local.py + - path: data/change_detection/change_detection_local/{state}/{district}/{block}/ + pattern: "change_{district}_{block}_{param_name}_{start_year}_{end_year}.tif" + params: + [ + Afforestation, + Deforestation, + Degradation, + Urbanization, + CropIntensity, + ] + geoserver_workspace: change_detection + + # change_detection_vector_local.py + - path: data/change_detection/change_detection_vector_local/{state}/{district}/{block}/ + pattern: "change_vector_{district}_{block}_{param_name}_{start_year}_{end_year}.gpkg" + params: + [ + Afforestation, + Deforestation, + Degradation, + Urbanization, + CropIntensity, + ] + geoserver_workspace: change_detection + + lulc: + # lulc_vector_local.py + - path: data/lulc/lulc_vector_local/{state}/{district}/{block}/ + pattern: "lulc_vector_{district}_{block}.gpkg" + geoserver_workspace: lulc_vector + + # lulc_v3_local.py + - path: data/lulc/lulc_v3_local/{state}/{district}/{block}/ + pattern: "{district}_{block}_{start_year}-07-01_{end_year}-06-30_LULCmap_10m.tif" + geoserver_workspace: LULC_v3 + + lulc_x_terrain: + # lulc_on_slope_cluster_local.py + - path: data/lulc_X_terrain/lulc_slope_clusters_local/{state}/{district}/{block}/ + pattern: "{district}_{block}_lulc_slope.gpkg" + geoserver_workspace: terrain_lulc + + misc: + # aquifer_vector_local.py + - path: data/misc/aquifer_vector_local/{state}/{district}/{block}/ + pattern: "aquifer_vector_{district}_{block}.gpkg" + geoserver_workspace: aquifer + + hydrology: + # generate_hydrology_local.py + - path: data/hydrology/hydrology_local/{state}/{district}/{block}/ + pattern: "deltaG_{period}_{district}_{block}.gpkg" + geoserver_workspace: mws_layers diff --git a/computing/config_loader.py b/computing/config_loader.py new file mode 100644 index 00000000..3462ed05 --- /dev/null +++ b/computing/config_loader.py @@ -0,0 +1,102 @@ +from pathlib import Path + +import yaml + +_CONFIG_PATH = Path(__file__).resolve().parent / "config.yaml" +PROJECT_ROOT = Path(__file__).resolve().parents[1] + + +def _load(): + with open(_CONFIG_PATH) as f: + return yaml.safe_load(f) + + +_cfg = _load() + + +def _abs(rel_path: str) -> Path: + base = rel_path.split("{")[0].rstrip("/") + return PROJECT_ROOT / base + + +def _find_input(path_suffix: str) -> dict: + for item in _cfg["base_layers"]["inputs"]: + if item["path"] == path_suffix: + return item + raise KeyError(f"No base layer input found in config.yaml for path: {path_suffix}") + + +def _output_entry(module: str, index: int = 0) -> dict: + return _cfg["local_compute_outputs"][module][index] + + +# --------------------------------------------------------------------------- +# Input paths +# --------------------------------------------------------------------------- + +LULC_BASE_DIR: Path = _abs( + next( + item["path"] + for item in _cfg["base_layers"]["inputs"] + if item["path"].startswith("data/base_layers/lulc/") + ) +).parent + +TERRAIN_RASTER_PATH: Path = PROJECT_ROOT / _find_input( + "data/base_layers/slope/slope_india_30m_merged.tif" +)["path"] + +SOIL_RASTER_PATH: Path = PROJECT_ROOT / _find_input( + "data/base_layers/soil/hysogs_india_250m_4326.tif" +)["path"] + +AEZ_VECTOR_PATH: Path = PROJECT_ROOT / _find_input( + "data/base_layers/AEZs/Agro_Ecological_Regions.shp" +)["path"] + +PRECOMPUTED_TEHSIL_WATERSHED_DIR: Path = _abs( + _find_input("data/base_layers/tehsil_watersheds/")["path"] +) + +MICROWATERSHED_PATH: Path = PROJECT_ROOT / _find_input( + "data/base_layers/Microwatershed_v2_with_details.geojson" +)["path"] + +AQUIFER_VECTOR_PATH: Path = PROJECT_ROOT / _find_input( + "data/base_layers/Aquifer_vector.geojson" +)["path"] + +SOI_TEHSIL_PATH: Path = PROJECT_ROOT / _find_input( + "data/admin-boundary/input/soi_tehsil.geojson" +)["path"] + +ADMIN_BOUNDARY_INPUT_DIR: Path = PROJECT_ROOT / "data/admin-boundary/input" +ADMIN_BOUNDARY_OUTPUT_DIR: Path = PROJECT_ROOT / "data/admin-boundary/output" +VILLAGE_BOUNDARIES_DIR: Path = PROJECT_ROOT / "data/base_layers/village_boundaries" + +# --------------------------------------------------------------------------- +# Google Drive IDs +# --------------------------------------------------------------------------- + +GDRIVE_ADMIN_BOUNDARY_FILE_ID: str = _find_input("data/admin-boundary/input/")["gdrive_id"] +GDRIVE_MICROWATERSHED_FILE_ID: str = _find_input( + "data/base_layers/Microwatershed_v2_with_details.geojson" +)["gdrive_id"] + +LULC_GDRIVE_FILES: list[tuple[str, str]] = [ + (Path(item["path"]).name, item["gdrive_id"]) + for item in _cfg["base_layers"]["inputs"] + if item["path"].startswith("data/base_layers/lulc/") and item.get("source") == "google_drive" +] + +# --------------------------------------------------------------------------- +# Output base directories +# --------------------------------------------------------------------------- + +CHANGE_DETECTION_RASTER_OUTPUT_DIR: Path = _abs(_output_entry("change_detection", 0)["path"]) +CHANGE_DETECTION_VECTOR_OUTPUT_DIR: Path = _abs(_output_entry("change_detection", 1)["path"]) +LULC_VECTOR_OUTPUT_DIR: Path = _abs(_output_entry("lulc", 0)["path"]) +LULC_V3_OUTPUT_DIR: Path = _abs(_output_entry("lulc", 1)["path"]) +LULC_SLOPE_CLUSTER_OUTPUT_DIR: Path = _abs(_output_entry("lulc_x_terrain", 0)["path"]) +AQUIFER_VECTOR_OUTPUT_DIR: Path = _abs(_output_entry("misc", 0)["path"]) +HYDROLOGY_LOCAL_OUTPUT_DIR: Path = _abs(_output_entry("hydrology", 0)["path"]) diff --git a/computing/cropping_intensity/cropping_intesity_local.py b/computing/cropping_intensity/cropping_intesity_local.py new file mode 100644 index 00000000..fddd5f62 --- /dev/null +++ b/computing/cropping_intensity/cropping_intesity_local.py @@ -0,0 +1,310 @@ +import os + +import pandas as pd +from nrm_app.celery import app +from utilities.gee_utils import valid_gee_text + +from computing.local_compute_helper import ( + LULC_BASE_DIR, + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + PROJECT_ROOT, + build_output_vector_path, + compute_categorical_raster_areas_for_watersheds, + compute_union_categorical_area_across_rasters_for_watersheds, + load_precomputed_watersheds, + resolve_lulc_raster_paths, + write_vector_output, +) +from computing.utils import ( + push_shape_to_geoserver, + save_layer_info_to_db, + update_layer_sync_status, +) + + +LOCAL_OUTPUT_BASE_DIR = ( + PROJECT_ROOT / "data/cropping_intensity/cropping_intensity_local" +) +GEOSERVER_WORKSPACE = "crop_intensity" +LOCAL_ALGORITHM = "local_cropping_intensity" +LOCAL_ALGORITHM_VERSION = "local-1.0" +INITIAL_YEAR = 2017 + +SINGLE_KHARIF = 8 +SINGLE_NON_KHARIF = 9 +DOUBLE = 10 +TRIPLE = 11 + +YEARLY_CLASS_DEFINITIONS = ( + {"value": SINGLE_KHARIF, "label_prefix": "single_kharif_cropped_area_"}, + {"value": SINGLE_NON_KHARIF, "label_prefix": "single_non_kharif_cropped_area_"}, + {"value": DOUBLE, "label_prefix": "doubly_cropped_area_"}, + {"value": TRIPLE, "label_prefix": "triply_cropped_area_"}, +) + + +def _slug(value, fallback): + return valid_gee_text(str(value).strip().lower()) or fallback + + +def _layer_name(asset_suffix, zoi_ci_asset=False): + if zoi_ci_asset: + return f"{asset_suffix}_intensity_ZOI" + return f"{asset_suffix}_intensity" + + +def _resolve_asset_suffix(state, district, block, asset_suffix): + if asset_suffix: + return asset_suffix + if state and district and block: + return ( + f"{_slug(district, 'unknown_district')}_" + f"{_slug(block, 'unknown_block')}" + ) + raise ValueError("state, district, and block are required for local cropping intensity.") + + +def _coerce_year(year_value, label): + try: + return int(year_value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{label} must be a valid integer year.") from exc + + +def _compute_yearly_area_columns(result_gdf, raster_paths, start_year): + for year, raster_path in zip(range(start_year, start_year + len(raster_paths)), raster_paths): + print(f"Computing local cropping intensity inputs for {year}-{year + 1}: {raster_path}") + year_definitions = [ + { + "value": class_definition["value"], + "label": f"{class_definition['label_prefix']}{year}", + } + for class_definition in YEARLY_CLASS_DEFINITIONS + ] + year_result = compute_categorical_raster_areas_for_watersheds( + watersheds_gdf=result_gdf, + raster_path=raster_path, + class_definitions=year_definitions, + ) + for definition in year_definitions: + result_gdf[definition["label"]] = year_result[definition["label"]].astype( + float + ) + result_gdf[f"single_cropped_area_{year}"] = ( + result_gdf[f"single_kharif_cropped_area_{year}"] + + result_gdf[f"single_non_kharif_cropped_area_{year}"] + ) + return result_gdf + + +def _compute_total_croppable_area(result_gdf, denominator_raster_paths, end_year): + output_column = f"total_cropable_area_ever_hydroyear_{INITIAL_YEAR}_{end_year}" + print( + "Computing local ever-croppable denominator " + f"for {INITIAL_YEAR}-{end_year} using {len(denominator_raster_paths)} rasters" + ) + return compute_union_categorical_area_across_rasters_for_watersheds( + watersheds_gdf=result_gdf, + raster_paths=denominator_raster_paths, + class_values=[SINGLE_KHARIF, SINGLE_NON_KHARIF, DOUBLE, TRIPLE], + output_column=output_column, + ) + + +def _compute_cropping_intensity_columns(result_gdf, start_year, end_year): + denominator_column = f"total_cropable_area_ever_hydroyear_{INITIAL_YEAR}_{end_year}" + denominator = pd.to_numeric(result_gdf[denominator_column], errors="coerce").fillna(0.0) + + for year in range(start_year, end_year + 1): + single_area = pd.to_numeric( + result_gdf[f"single_cropped_area_{year}"], errors="coerce" + ).fillna(0.0) + double_area = pd.to_numeric( + result_gdf[f"doubly_cropped_area_{year}"], errors="coerce" + ).fillna(0.0) + triple_area = pd.to_numeric( + result_gdf[f"triply_cropped_area_{year}"], errors="coerce" + ).fillna(0.0) + + nonzero_denominator = denominator > 0 + intensity = pd.Series(0.0, index=result_gdf.index, dtype=float) + intensity.loc[nonzero_denominator] = ( + single_area.loc[nonzero_denominator] / denominator.loc[nonzero_denominator] + + 2.0 * double_area.loc[nonzero_denominator] / denominator.loc[nonzero_denominator] + + 3.0 * triple_area.loc[nonzero_denominator] / denominator.loc[nonzero_denominator] + ) + result_gdf[f"cropping_intensity_{year}"] = intensity.astype(float) + + return result_gdf + + +def run_cropping_intensity_local( + state, + district, + block, + start_year, + end_year, + asset_suffix=None, + zoi_ci_asset=False, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + lulc_dir=LULC_BASE_DIR, + push_to_geoserver=True, + sync_layer_metadata=True, +): + state = str(state).strip().lower() + district = str(district).strip().lower() + block = str(block).strip().lower() + start_year = _coerce_year(start_year, "start_year") + end_year = _coerce_year(end_year, "end_year") + + if start_year > end_year: + raise ValueError("start_year cannot be greater than end_year") + if start_year < INITIAL_YEAR: + raise ValueError( + f"start_year must be greater than or equal to {INITIAL_YEAR} for local cropping intensity." + ) + + asset_suffix = _resolve_asset_suffix(state, district, block, asset_suffix) + layer_name = _layer_name(asset_suffix, zoi_ci_asset=zoi_ci_asset) + + watersheds_gdf, watershed_source = load_precomputed_watersheds( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + ) + print(f"Watershed boundary source: {watershed_source}") + + yearly_raster_paths = resolve_lulc_raster_paths( + start_year=start_year, + end_year=end_year, + lulc_dir=lulc_dir, + ) + denominator_raster_paths = resolve_lulc_raster_paths( + start_year=INITIAL_YEAR, + end_year=end_year, + lulc_dir=lulc_dir, + ) + + result_gdf = watersheds_gdf.copy() + result_gdf = _compute_yearly_area_columns( + result_gdf=result_gdf, + raster_paths=yearly_raster_paths, + start_year=start_year, + ) + result_gdf = _compute_total_croppable_area( + result_gdf=result_gdf, + denominator_raster_paths=denominator_raster_paths, + end_year=end_year, + ) + result_gdf = _compute_cropping_intensity_columns( + result_gdf=result_gdf, + start_year=start_year, + end_year=end_year, + ) + + output_path = build_output_vector_path( + layer_name=layer_name, + state=state, + district=district, + block=block, + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + block_fallback="unknown_block", + ) + asset_id = write_vector_output( + gdf=result_gdf, + output_path=output_path, + layer_name=layer_name, + ) + print(f"Saved local cropping intensity vector: {asset_id}") + + if push_to_geoserver: + geoserver_response = push_shape_to_geoserver( + os.path.splitext(asset_id)[0], + workspace=GEOSERVER_WORKSPACE, + layer_name=layer_name, + file_type="gpkg", + ) + print(f"GeoServer response: {geoserver_response}") + if not isinstance(geoserver_response, dict) or geoserver_response.get( + "status_code" + ) not in (200, 201): + return False + + 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="Cropping Intensity", + misc={ + "start_year": start_year, + "end_year": end_year, + }, + algorithm=LOCAL_ALGORITHM, + algorithm_version=LOCAL_ALGORITHM_VERSION, + ) + if layer_id and push_to_geoserver: + update_layer_sync_status(layer_id=layer_id, sync_to_geoserver=True) + + return True + + +def _generate_cropping_intensity_local_task( + state=None, + district=None, + block=None, + roi_path=None, + asset_suffix=None, + asset_folder_list=None, + app_type="MWS", + start_year=None, + end_year=None, + gee_account_id=None, + zoi_ci_asset=None, +): + _ = roi_path, asset_folder_list, app_type, gee_account_id + return run_cropping_intensity_local( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + asset_suffix=asset_suffix, + zoi_ci_asset=bool(zoi_ci_asset), + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + +@app.task(bind=True) +def generate_cropping_intensity( + self, + state=None, + district=None, + block=None, + roi_path=None, + asset_suffix=None, + asset_folder_list=None, + app_type="MWS", + start_year=None, + end_year=None, + gee_account_id=None, + zoi_ci_asset=None, +): + _ = self + return _generate_cropping_intensity_local_task( + state=state, + district=district, + block=block, + roi_path=roi_path, + asset_suffix=asset_suffix, + asset_folder_list=asset_folder_list, + app_type=app_type, + start_year=start_year, + end_year=end_year, + gee_account_id=gee_account_id, + zoi_ci_asset=zoi_ci_asset, + ) 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/local_compute_helper.py b/computing/local_compute_helper.py new file mode 100644 index 00000000..254a32ea --- /dev/null +++ b/computing/local_compute_helper.py @@ -0,0 +1,870 @@ +import os +from contextlib import ExitStack +from pathlib import Path + +import geopandas as gpd +import numpy as np +import pandas as pd +import rasterio +from rasterio.mask import mask +from shapely.geometry import mapping +from utilities.gee_utils import valid_gee_text + +from computing.utils import convert_to_zip +import logging + + +from computing.config_loader import ( + AEZ_VECTOR_PATH, + LULC_BASE_DIR, + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + PROJECT_ROOT, + TERRAIN_RASTER_PATH, +) + +PRECOMPUTED_ROI_EXTENSIONS = (".gpkg", ".geojson") +VALID_COMPUTE_TYPES = {"gee", "local"} +MIN_WATERSHED_AREA_HA = 400.0 +LULC_CLASSES = np.arange(1, 13, dtype=np.int16) + +TERRAIN_CLUSTER_CENTROIDS = np.array( + [ + [0.36255426, 0.21039965, 0.12161905, 0.17393119, 0.13149585], + [0.09171062, 0.84299211, 0.035222, 0.02172654, 0.00834873], + [0.08497599, 0.01051893, 0.23763531, 0.37992855, 0.28694122], + [0.22301813, 0.5611825, 0.08511123, 0.07314189, 0.05754624], + ], + dtype=np.float64, +) + +PLAIN_CLASSES = {5} +VALLEY_CLASSES = {1, 2, 4, 9} +HILL_SLOPES_CLASSES = {8} +RIDGE_CLASSES = {3, 7, 10, 11} +SLOPY_CLASSES = {6} + + +def _slug(value, fallback): + return valid_gee_text(str(value).strip().lower()) or fallback + + +def ensure_file_exists(path, label): + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"{label} not found: {path}") + + +def validate_geometry(gdf): + gdf = gdf[gdf.geometry.notna()].copy() + if gdf.empty: + return gdf + invalid = ~gdf.is_valid + if invalid.any(): + gdf.loc[invalid, "geometry"] = gdf.loc[invalid, "geometry"].buffer(0) + return gdf[~gdf.geometry.is_empty].copy() + + +def read_validated_vector_file(path, empty_message): + gdf = validate_geometry(gpd.read_file(path)) + if gdf.empty: + raise ValueError(empty_message) + return gdf + + +def resolve_precomputed_vector_file( + state, + district, + block, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + extensions=PRECOMPUTED_ROI_EXTENSIONS, + missing_file_label="Precomputed vector file", +): + roi_dir = Path(precomputed_roi_dir or PRECOMPUTED_TEHSIL_WATERSHED_DIR) + state_slug = _slug(state, "unknown_state") + district_slug = _slug(district, "unknown_district") + block_slug = _slug(block, "unknown_tehsil") + + expected_paths = [ + roi_dir / state_slug / district_slug / f"{block_slug}{ext}" + for ext in extensions + ] + for path in expected_paths: + if path.exists(): + return path + + raise FileNotFoundError( + f"{missing_file_label} not found. " + f"state={state}, district={district}, block={block}. " + f"Expected one of: {[str(path) for path in expected_paths]}" + ) + + +def load_precomputed_watersheds( + state, + district, + block, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, +): + watershed_path = resolve_precomputed_vector_file( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + missing_file_label="Precomputed watershed boundary file", + ) + watersheds_gdf = read_validated_vector_file( + watershed_path, + f"Precomputed watershed file has no valid geometries: {watershed_path}", + ) + print(f"Loaded watershed boundaries: {watershed_path}") + return watersheds_gdf, str(watershed_path) + + +def load_precomputed_roi( + state, + district, + block, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, +): + roi_path = resolve_precomputed_vector_file( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + missing_file_label="Precomputed tehsil watershed file", + ) + roi_gdf = read_validated_vector_file( + roi_path, + f"Precomputed ROI file has no valid geometries: {roi_path}", + ) + print(f"Loaded precomputed ROI file: {roi_path}") + return roi_gdf + + +def _build_output_dir( + output_base_dir, + state=None, + district=None, + block=None, + custom_subdir="custom", + block_fallback="unknown_tehsil", +): + output_base_dir = Path(output_base_dir) + if state and district and block: + output_dir = ( + output_base_dir + / _slug(state, "unknown_state") + / _slug(district, "unknown_district") + / _slug(block, block_fallback) + ) + else: + output_dir = output_base_dir / custom_subdir + output_dir.mkdir(parents=True, exist_ok=True) + return output_dir + + +def build_output_vector_path( + layer_name, + state, + district, + block, + output_base_dir, + block_fallback="unknown_block", +): + output_dir = _build_output_dir( + output_base_dir=output_base_dir, + state=state, + district=district, + block=block, + block_fallback=block_fallback, + ) + return output_dir / f"{layer_name}.gpkg" + + +def build_output_raster_path( + layer_name, + output_base_dir, + state=None, + district=None, + block=None, + custom_subdir="custom", + block_fallback="unknown_tehsil", +): + output_dir = _build_output_dir( + output_base_dir=output_base_dir, + state=state, + district=district, + block=block, + custom_subdir=custom_subdir, + block_fallback=block_fallback, + ) + return output_dir / f"{layer_name}.tif" + + +def write_vector_output(gdf, output_path, layer_name): + gdf.to_file(output_path, driver="GPKG", layer=layer_name) + return str(output_path) + + +def clip_raster_with_roi(roi_gdf, raster_path, output_path, raster_label="Raster"): + ensure_file_exists(raster_path, raster_label) + + with rasterio.open(raster_path) as src: + clip_gdf = roi_gdf + if src.crs and clip_gdf.crs and clip_gdf.crs != src.crs: + clip_gdf = clip_gdf.to_crs(src.crs) + + clip_gdf = validate_geometry(clip_gdf) + shapes = [ + mapping(geom) + for geom in clip_gdf.geometry + if geom is not None and not geom.is_empty + ] + if not shapes: + raise ValueError("No valid ROI geometry available for raster clipping.") + + clipped_data, clipped_transform = mask(src, shapes=shapes, crop=True) + clipped_meta = src.meta.copy() + clipped_meta.update( + { + "driver": "GTiff", + "height": clipped_data.shape[1], + "width": clipped_data.shape[2], + "transform": clipped_transform, + "compress": "lzw", + } + ) + + with rasterio.open(output_path, "w", **clipped_meta) as dst: + dst.write(clipped_data) + + return str(output_path) + + +RASTER_DECLARED_SRS = "EPSG:4326" +VECTOR_DECLARED_SRS = "EPSG:4326" + + +def _push_raster_to_geoserver_instance(geo, file_path, layer_name, workspace, style_name): + + _log = logging.getLogger(__name__) + + geo.ensure_workspace(workspace) + geo.delete_raster_store(workspace=workspace, store=layer_name) + upload_response = geo.create_coveragestore( + path=file_path, + workspace=workspace, + layer_name=layer_name, + ) + try: + geo.configure_coverage( + workspace=workspace, + store=layer_name, + coverage=layer_name, + srs=RASTER_DECLARED_SRS, + projection_policy="REPROJECT_TO_DECLARED", + enabled=True, + ) + except Exception as e: + _log.warning( + "Failed to force SRS on coverage %s:%s: %s", + workspace, + layer_name, + e, + ) + style_response = None + if style_name: + style_response = geo.publish_style( + layer_name=layer_name, + style_name=style_name, + workspace=workspace, + ) + return upload_response, style_response + + +def _verify_raster_layer(geo, workspace, layer_name): + """Return True iff a layer (not just the store) is registered on the server.""" + from utilities.geoserver_utils import GeoserverException + + try: + geo.get_layer(layer_name=layer_name, workspace=workspace) + return True + except GeoserverException: + return False + except Exception: + return False + + +def _push_vector_to_geoserver_instance(geo, path, layer_name, workspace, file_type="gpkg"): + _log = logging.getLogger(__name__) + + geo.ensure_workspace(workspace) + try: + geo.delete_vector_store(workspace=workspace, store=layer_name) + except Exception: + pass + + zip_path = convert_to_zip(path, file_type) + response = geo.create_shp_datastore( + path=zip_path, + store_name=layer_name, + workspace=workspace, + file_extension=file_type, + ) + try: + geo.configure_featuretype( + workspace=workspace, + store=layer_name, + featuretype=layer_name, + srs=VECTOR_DECLARED_SRS, + projection_policy="FORCE_DECLARED", + ) + except Exception as e: + _log.warning( + "Failed to force SRS on featuretype %s:%s: %s", + workspace, + layer_name, + e, + ) + return response + + +def push_local_vector_to_geoserver(path, layer_name, workspace, file_type="gpkg"): + import logging + from django.conf import settings + from utilities.geoserver_utils import Geoserver + + _log = logging.getLogger(__name__) + + local_geo = Geoserver() + response = _push_vector_to_geoserver_instance(local_geo, path, layer_name, workspace, file_type) + _log.info("Pushed vector %s to local GeoServer.", layer_name) + + prod_url = getattr(settings, "PROD_GEOSERVER_URL", "") + if prod_url: + try: + prod_geo = Geoserver( + service_url=prod_url, + username=settings.PROD_GEOSERVER_USERNAME, + password=settings.PROD_GEOSERVER_PASSWORD, + ) + _push_vector_to_geoserver_instance(prod_geo, path, layer_name, workspace, file_type) + _log.info("Pushed vector %s to prod GeoServer (%s).", layer_name, prod_url) + except Exception as e: + _log.error("Failed to push vector %s to prod GeoServer: %s", layer_name, e) + + return response + + +def push_local_raster_to_geoserver(file_path, layer_name, workspace, style_name=None): + import logging + from django.conf import settings + from utilities.geoserver_utils import Geoserver + + _log = logging.getLogger(__name__) + + local_geo = Geoserver() + upload_response, style_response = _push_raster_to_geoserver_instance( + local_geo, file_path, layer_name, workspace, style_name + ) + local_layer_ok = _verify_raster_layer(local_geo, workspace, layer_name) + _log.info( + "Pushed raster %s to local GeoServer (layer_exists=%s).", + layer_name, + local_layer_ok, + ) + + prod_url = getattr(settings, "PROD_GEOSERVER_URL", "") + if prod_url: + try: + prod_geo = Geoserver( + service_url=prod_url, + username=settings.PROD_GEOSERVER_USERNAME, + password=settings.PROD_GEOSERVER_PASSWORD, + ) + prod_upload, prod_style = _push_raster_to_geoserver_instance( + prod_geo, file_path, layer_name, workspace, style_name + ) + prod_layer_ok = _verify_raster_layer(prod_geo, workspace, layer_name) + _log.info( + "Pushed raster %s to prod GeoServer (%s) " + "(layer_exists=%s, upload=%s, style=%s).", + layer_name, + prod_url, + prod_layer_ok, + prod_upload, + prod_style, + ) + if not prod_layer_ok: + _log.error( + "Prod raster push for %s created store but NO layer. " + "Coverage auto-config likely failed on prod GeoServer.", + layer_name, + ) + except Exception as e: + _log.error("Failed to push raster %s to prod GeoServer: %s", layer_name, e) + + return upload_response, style_response + + + +def compute_pixel_area_grid(transform, height, width, crs): + if crs is None: + raise ValueError("Raster CRS is missing; cannot compute pixel areas.") + + if getattr(crs, "is_geographic", False): + lon_width_radians = np.deg2rad(abs(transform.a)) + row_indices = np.arange(height, dtype=np.float64) + lat_top = transform.f + (row_indices * transform.e) + lat_bottom = lat_top + transform.e + earth_radius_m = 6378137.0 + row_areas = (earth_radius_m**2) * lon_width_radians * np.abs( + np.sin(np.deg2rad(lat_top)) - np.sin(np.deg2rad(lat_bottom)) + ) + return np.broadcast_to(row_areas[:, None], (height, width)) + + pixel_area = abs(transform.a * transform.e) + return np.full((height, width), pixel_area, dtype=np.float64) + + + +def compute_categorical_raster_areas_for_watersheds( + watersheds_gdf, + raster_path, + class_definitions, +): + ensure_file_exists(raster_path, "Categorical raster") + + with rasterio.open(raster_path) as src: + working_gdf = watersheds_gdf.copy() + if working_gdf.crs is None: + raise ValueError( + "Watershed CRS is missing; cannot align with categorical raster." + ) + if src.crs and working_gdf.crs != src.crs: + working_gdf = working_gdf.to_crs(src.crs) + + nodata = src.nodata + computed_rows = [] + empty_result = { + class_definition["label"]: 0.0 + for class_definition in class_definitions + } + + total = len(working_gdf) + for index, row in enumerate(working_gdf.itertuples(index=False), start=1): + geom = row.geometry + if geom is None or geom.is_empty: + computed_rows.append(empty_result.copy()) + continue + + try: + clipped, clipped_transform = mask( + src, + [mapping(geom)], + crop=True, + filled=False, + ) + except ValueError: + computed_rows.append(empty_result.copy()) + continue + + data = clipped[0] + if data.size == 0: + computed_rows.append(empty_result.copy()) + continue + + values = np.asarray(data, dtype=np.float64) + values = np.where(np.isfinite(values), values, 0) + values = np.rint(values).astype(np.int16, copy=False) + valid_mask = ~np.ma.getmaskarray(data) + if nodata is not None and not np.isnan(nodata): + valid_mask &= values != int(round(nodata)) + + pixel_area_ha = ( + compute_pixel_area_grid( + transform=clipped_transform, + height=values.shape[0], + width=values.shape[1], + crs=src.crs, + ) + * 0.0001 + ) + + row_result = {} + for class_definition in class_definitions: + raw_values = class_definition.get("values", class_definition.get("value")) + if isinstance(raw_values, (list, tuple, set, np.ndarray)): + class_values = list(raw_values) + else: + class_values = [raw_values] + class_mask = valid_mask & np.isin(values, class_values) + row_result[class_definition["label"]] = float( + pixel_area_ha[class_mask].sum() + ) + + computed_rows.append(row_result) + + if index % 200 == 0 or index == total: + print( + f"Computed categorical raster areas for {index}/{total} watersheds" + ) + + result = watersheds_gdf.copy() + computed_df = pd.DataFrame(computed_rows) + for column in computed_df.columns: + result[column] = computed_df[column].values + return result + + +def compute_union_categorical_area_across_rasters_for_watersheds( + watersheds_gdf, + raster_paths, + class_values, + output_column, +): + if not raster_paths: + raise ValueError("At least one raster path is required for union area computation.") + + for raster_path in raster_paths: + ensure_file_exists(raster_path, "Categorical raster") + + if isinstance(class_values, (list, tuple, set, np.ndarray)): + class_values = list(class_values) + else: + class_values = [class_values] + + with ExitStack() as stack: + sources = [stack.enter_context(rasterio.open(path)) for path in raster_paths] + aligned_gdfs = [] + for src in sources: + aligned = watersheds_gdf.copy() + if aligned.crs is None: + raise ValueError( + "Watershed CRS is missing; cannot align with categorical raster." + ) + if src.crs and aligned.crs != src.crs: + aligned = aligned.to_crs(src.crs) + aligned_gdfs.append(aligned) + + computed_rows = [] + total = len(watersheds_gdf) + + for index in range(total): + union_mask = None + union_pixel_area_ha = None + + for src, aligned_gdf in zip(sources, aligned_gdfs): + geom = aligned_gdf.iloc[index].geometry + if geom is None or geom.is_empty: + continue + + try: + clipped, clipped_transform = mask( + src, + [mapping(geom)], + crop=True, + filled=False, + ) + except ValueError: + continue + + data = clipped[0] + if data.size == 0: + continue + + values = np.asarray(data, dtype=np.float64) + values = np.where(np.isfinite(values), values, 0) + values = np.rint(values).astype(np.int16, copy=False) + + valid_mask = ~np.ma.getmaskarray(data) + nodata = src.nodata + if nodata is not None and not np.isnan(nodata): + valid_mask &= values != int(round(nodata)) + + class_mask = valid_mask & np.isin(values, class_values) + if union_mask is None: + union_mask = class_mask + union_pixel_area_ha = ( + compute_pixel_area_grid( + transform=clipped_transform, + height=values.shape[0], + width=values.shape[1], + crs=src.crs, + ) + * 0.0001 + ) + else: + if union_mask.shape != class_mask.shape: + raise ValueError( + "Local categorical rasters do not align for union area computation." + ) + union_mask |= class_mask + + if union_mask is None or union_pixel_area_ha is None: + computed_rows.append({output_column: 0.0}) + else: + computed_rows.append( + {output_column: float(union_pixel_area_ha[union_mask].sum())} + ) + + if (index + 1) % 200 == 0 or (index + 1) == total: + print( + f"Computed union categorical raster areas for {index + 1}/{total} watersheds" + ) + + result = watersheds_gdf.copy() + computed_df = pd.DataFrame(computed_rows) + result[output_column] = computed_df[output_column].values + return result + + +def get_union_geometry(gdf): + if hasattr(gdf.geometry, "union_all"): + return gdf.geometry.union_all() + return gdf.geometry.unary_union + + +def resolve_clipped_terrain_raster_path( + state, + district, + block, + clipped_raster_dir, +): + layer_stub = f"{_slug(district, 'unknown_district')}_{_slug(block, 'unknown_tehsil')}_terrain_raster" + raster_path = ( + Path(clipped_raster_dir) + / _slug(state, "unknown_state") + / _slug(district, "unknown_district") + / _slug(block, "unknown_tehsil") + / f"{layer_stub}.tif" + ) + ensure_file_exists(raster_path, "Clipped terrain raster") + return str(raster_path) + + +def _compute_cluster_id( + slopy_prop, + plain_prop, + ridge_prop, + valley_prop, + hill_slopes_prop, +): + feature_vector = np.array( + [slopy_prop, plain_prop, ridge_prop, valley_prop, hill_slopes_prop], + dtype=np.float64, + ) + distances = np.sum((TERRAIN_CLUSTER_CENTROIDS - feature_vector) ** 2, axis=1) + return int(np.argmin(distances)) + + +def _fraction(values, valid_mask, class_values): + total = int(valid_mask.sum()) + if total == 0: + return 0.0 + class_mask = np.isin(values, list(class_values)) + return float((class_mask & valid_mask).sum()) / float(total) + + +def compute_terrain_properties_for_watersheds(watersheds_gdf, raster_path): + with rasterio.open(raster_path) as src: + working_gdf = watersheds_gdf.copy() + if working_gdf.crs is None: + raise ValueError( + "Watershed CRS is missing; cannot align with raster CRS." + ) + if src.crs and working_gdf.crs != src.crs: + working_gdf = working_gdf.to_crs(src.crs) + + nodata = src.nodata + computed_rows = [] + + total = len(working_gdf) + for index, row in enumerate(working_gdf.itertuples(index=False), start=1): + geom = row.geometry + if geom is None or geom.is_empty: + computed_rows.append( + { + "plain_area": 0.0, + "valley_area": 0.0, + "hill_slopes_area": 0.0, + "ridge_area": 0.0, + "slopy_area": 0.0, + "terrainClusters": -1, + } + ) + continue + + try: + clipped, _ = mask(src, [mapping(geom)], crop=True, filled=True) + except ValueError: + computed_rows.append( + { + "plain_area": 0.0, + "valley_area": 0.0, + "hill_slopes_area": 0.0, + "ridge_area": 0.0, + "slopy_area": 0.0, + "terrainClusters": -1, + } + ) + continue + + values = clipped[0] + if values.size == 0: + computed_rows.append( + { + "plain_area": 0.0, + "valley_area": 0.0, + "hill_slopes_area": 0.0, + "ridge_area": 0.0, + "slopy_area": 0.0, + "terrainClusters": -1, + } + ) + continue + + values = np.rint(values).astype(np.int16, copy=False) + valid_mask = np.ones(values.shape, dtype=bool) + if nodata is not None and not np.isnan(nodata): + valid_mask &= values != int(round(nodata)) + valid_mask &= values != 0 + + plain_prop = _fraction(values, valid_mask, PLAIN_CLASSES) + valley_prop = _fraction(values, valid_mask, VALLEY_CLASSES) + hill_slopes_prop = _fraction( + values, valid_mask, HILL_SLOPES_CLASSES + ) + ridge_prop = _fraction(values, valid_mask, RIDGE_CLASSES) + slopy_prop = _fraction(values, valid_mask, SLOPY_CLASSES) + + if valid_mask.sum() == 0: + cluster_id = -1 + else: + cluster_id = _compute_cluster_id( + slopy_prop=slopy_prop, + plain_prop=plain_prop, + ridge_prop=ridge_prop, + valley_prop=valley_prop, + hill_slopes_prop=hill_slopes_prop, + ) + + computed_rows.append( + { + "plain_area": plain_prop * 100.0, + "valley_area": valley_prop * 100.0, + "hill_slopes_area": hill_slopes_prop * 100.0, + "ridge_area": ridge_prop * 100.0, + "slopy_area": slopy_prop * 100.0, + "terrainClusters": cluster_id, + } + ) + + if index % 200 == 0 or index == total: + print( + f"Computed terrain properties for {index}/{total} watersheds" + ) + + result = watersheds_gdf.copy() + computed_df = pd.DataFrame(computed_rows) + for column in computed_df.columns: + result[column] = computed_df[column].values + return result + + +def resolve_lulc_raster_paths( + start_year, + end_year, + lulc_dir=LULC_BASE_DIR, +): + raster_paths = [] + for year in range(int(start_year), int(end_year) + 1): + raster_path = Path(lulc_dir) / f"lulc_v3_{year}_{year + 1}.tif" + ensure_file_exists(raster_path, f"LULC raster for {year}-{year + 1}") + raster_paths.append(str(raster_path)) + return raster_paths + + +def get_watershed_areas_in_hectares(watersheds_gdf): + if "area_in_ha" in watersheds_gdf.columns: + area_in_ha = pd.to_numeric( + watersheds_gdf["area_in_ha"], errors="coerce" + ) + if area_in_ha.notna().any(): + return area_in_ha + projected = watersheds_gdf.to_crs("EPSG:6933") + return projected.geometry.area / 10000.0 + + +def filter_large_watersheds( + watersheds_gdf, + min_watershed_area_ha=MIN_WATERSHED_AREA_HA, +): + area_in_ha = get_watershed_areas_in_hectares(watersheds_gdf) + filtered = watersheds_gdf.loc[area_in_ha > min_watershed_area_ha].copy() + filtered["area"] = area_in_ha.loc[filtered.index].astype(float) * 10000.0 + return filtered.reset_index(drop=True) + + +def resolve_aez_code( + watersheds_gdf, + aez_vector_path=AEZ_VECTOR_PATH, +): + ensure_file_exists(aez_vector_path, "AEZ vector") + + aez_gdf = read_validated_vector_file( + aez_vector_path, + f"AEZ vector has no valid geometries: {aez_vector_path}", + ) + + study_area = watersheds_gdf[["geometry"]].copy() + if study_area.crs is None: + raise ValueError("Watershed CRS is missing; cannot resolve AEZ.") + if aez_gdf.crs and study_area.crs != aez_gdf.crs: + study_area = study_area.to_crs(aez_gdf.crs) + + study_union = get_union_geometry(study_area) + intersecting_aez = aez_gdf.loc[aez_gdf.intersects(study_union)].copy() + if intersecting_aez.empty: + raise ValueError("No AEZ polygon intersects the watershed study area.") + + intersecting_aez["overlap_area"] = intersecting_aez.geometry.intersection( + study_union + ).area + return int( + intersecting_aez.sort_values("overlap_area", ascending=False).iloc[0][ + "ae_regcode" + ] + ) + + +def compute_mode_lulc_array(reprojected_arrays, lulc_classes=LULC_CLASSES): + stack = np.rint(np.stack(reprojected_arrays, axis=0)).astype( + np.int16, + copy=False, + ) + valid_mask = np.isfinite(stack) & (stack > 0) + counts = np.zeros( + (len(lulc_classes), stack.shape[1], stack.shape[2]), + dtype=np.uint16, + ) + + for index, lulc_class in enumerate(lulc_classes): + counts[index] = np.sum((stack == lulc_class) & valid_mask, axis=0) + + mode_index = np.argmax(counts, axis=0) + mode_values = lulc_classes[mode_index] + max_counts = np.max(counts, axis=0) + mode_values[max_counts == 0] = 0 + return mode_values + + +def get_compute_mode(request, default="local"): + compute = str(request.data.get("compute") or default).strip().lower() + if compute not in VALID_COMPUTE_TYPES: + raise ValueError("compute must be either 'gee' or 'local'") + return compute + + +def select_compute_task(compute, gee_task, local_task): + return gee_task if compute == "gee" else local_task diff --git a/computing/lulc/lulc_v3_local.py b/computing/lulc/lulc_v3_local.py new file mode 100644 index 00000000..13e273d4 --- /dev/null +++ b/computing/lulc/lulc_v3_local.py @@ -0,0 +1,267 @@ +from nrm_app.celery import app +from utilities.gee_utils import valid_gee_text + +from computing.config_loader import LULC_V3_OUTPUT_DIR as LOCAL_OUTPUT_BASE_DIR +from computing.local_compute_helper import ( + LULC_BASE_DIR, + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + build_output_raster_path, + clip_raster_with_roi, + load_precomputed_roi, + push_local_raster_to_geoserver, + read_validated_vector_file, + resolve_lulc_raster_paths, +) +from computing.models import Dataset +from computing.utils import save_layer_info_to_db, update_layer_sync_status +GEOSERVER_WORKSPACE = "LULC_v3" +GEOSERVER_STYLE = "lulc_level_3_style" +LOCAL_ALGORITHM = "local_lulc_v3_clip" +LOCAL_ALGORITHM_VERSION = "local-1.0" + + +def _slug(value, fallback): + if value is None: + return fallback + text = str(value).strip().lower() + if not text: + return fallback + return valid_gee_text(text) or fallback + + +def _resolve_roi(state, district, block, roi_path, precomputed_roi_dir): + if state and district and block: + return load_precomputed_roi( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + ) + + if not roi_path: + raise ValueError( + "For non state/district/block runs, `roi_path` must be provided." + ) + + return read_validated_vector_file( + roi_path, + f"ROI file has no valid geometries: {roi_path}", + ) + + +def _resolve_filename_prefix(district, block, asset_suffix): + if district and block: + return ( + f"{_slug(district, 'unknown_district')}_" + f"{_slug(block, 'unknown_block')}" + ) + if asset_suffix is None or not str(asset_suffix).strip(): + raise ValueError( + "For non state/district/block runs, `asset_suffix` must be provided." + ) + return _slug(asset_suffix, "custom") + + +def _build_output_stub(filename_prefix, start_year): + return ( + f"{filename_prefix}_{start_year}-07-01_" + f"{start_year + 1}-06-30_LULCmap_10m" + ) + + +def _build_layer_name(start_year, end_year, filename_prefix): + return f"LULC_{start_year}_{end_year}_{filename_prefix}_level_3" + + +def _resolve_dataset_name(): + return "LULC_v3" if Dataset.objects.filter(name="LULC_v3").exists() else "LULC_level_3" + + +def _sync_lulc_stac(layer_id, state, district, block, start_year): + from computing.STAC_specs import generate_STAC_layerwise + + layer_stac_generated = generate_STAC_layerwise.generate_raster_stac( + state=state, + district=district, + block=block, + layer_name="land_use_land_cover_raster", + start_year=str(start_year), + ) + update_layer_sync_status( + layer_id=layer_id, + is_stac_specs_generated=layer_stac_generated, + ) + + +def run_lulc_v3_local( + state=None, + district=None, + block=None, + start_year=None, + end_year=None, + roi_path=None, + asset_suffix=None, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + lulc_dir=LULC_BASE_DIR, + push_to_geoserver=True, + sync_layer_metadata=True, +): + state = str(state).strip().lower() if state else None + district = str(district).strip().lower() if district else None + block = str(block).strip().lower() if block else None + 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") + + roi_gdf = _resolve_roi( + state=state, + district=district, + block=block, + roi_path=roi_path, + precomputed_roi_dir=precomputed_roi_dir, + ) + filename_prefix = _resolve_filename_prefix(district, block, asset_suffix) + raster_paths = resolve_lulc_raster_paths( + start_year=start_year, + end_year=end_year, + lulc_dir=lulc_dir, + ) + + layer_at_geoserver = True + is_admin_run = bool(state and district and block) + + for current_year, raster_path in zip(range(start_year, end_year + 1), raster_paths): + year_start = str(current_year)[2:] + year_end = str(current_year + 1)[2:] + output_stub = _build_output_stub(filename_prefix, current_year) + output_path = build_output_raster_path( + layer_name=output_stub, + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + state=state, + district=district, + block=block, + custom_subdir=_slug(asset_suffix, "custom"), + block_fallback="unknown_block", + ) + clipped_raster_path = clip_raster_with_roi( + roi_gdf=roi_gdf, + raster_path=raster_path, + output_path=output_path, + raster_label=f"LULC raster for {current_year}-{current_year + 1}", + ) + print(f"Saved local LULC raster: {clipped_raster_path}") + + layer_name = _build_layer_name( + start_year=year_start, + end_year=year_end, + filename_prefix=filename_prefix, + ) + print(f"Prepared local LULC layer {GEOSERVER_WORKSPACE}: {layer_name}") + + layer_id = None + if sync_layer_metadata and is_admin_run: + layer_id = save_layer_info_to_db( + state=state, + district=district, + block=block, + layer_name=layer_name, + asset_id=clipped_raster_path, + dataset_name=_resolve_dataset_name(), + misc={ + "start_year": start_year, + "end_year": end_year, + "is_generated_locally": True, + }, + algorithm=LOCAL_ALGORITHM, + algorithm_version=LOCAL_ALGORITHM_VERSION, + ) + + if not push_to_geoserver: + continue + + try: + upload_res, style_res = push_local_raster_to_geoserver( + file_path=clipped_raster_path, + layer_name=layer_name, + workspace=GEOSERVER_WORKSPACE, + style_name=GEOSERVER_STYLE, + ) + print(f"GeoServer upload response for {layer_name}: {upload_res}") + print(f"GeoServer style response for {layer_name}: {style_res}") + except Exception as error: + print(f"Failed to sync local LULC raster {layer_name}: {error}") + layer_at_geoserver = False + continue + + if layer_id: + update_layer_sync_status(layer_id=layer_id, sync_to_geoserver=True) + _sync_lulc_stac( + layer_id=layer_id, + state=state, + district=district, + block=block, + start_year=current_year, + ) + + return layer_at_geoserver if push_to_geoserver else True + + +def _clip_lulc_v3_local_task( + state=None, + district=None, + block=None, + start_year=None, + end_year=None, + gee_account_id=None, + roi_path=None, + asset_folder=None, + asset_suffix=None, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + app_type="MWS", +): + _ = gee_account_id, asset_folder, app_type + return run_lulc_v3_local( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + roi_path=roi_path, + asset_suffix=asset_suffix, + precomputed_roi_dir=precomputed_roi_dir, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + +@app.task(bind=True) +def clip_lulc_v3( + self, + state=None, + district=None, + block=None, + start_year=None, + end_year=None, + gee_account_id=None, + roi_path=None, + asset_folder=None, + asset_suffix=None, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + app_type="MWS", +): + _ = self + return _clip_lulc_v3_local_task( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + gee_account_id=gee_account_id, + roi_path=roi_path, + asset_folder=asset_folder, + asset_suffix=asset_suffix, + precomputed_roi_dir=precomputed_roi_dir, + app_type=app_type, + ) diff --git a/computing/lulc/lulc_vector_local.py b/computing/lulc/lulc_vector_local.py new file mode 100644 index 00000000..e9c7628c --- /dev/null +++ b/computing/lulc/lulc_vector_local.py @@ -0,0 +1,195 @@ +import os + +from nrm_app.celery import app +from utilities.gee_utils import valid_gee_text + +from computing.config_loader import LULC_VECTOR_OUTPUT_DIR as LOCAL_OUTPUT_BASE_DIR +from computing.local_compute_helper import ( + LULC_BASE_DIR, + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + build_output_vector_path, + compute_categorical_raster_areas_for_watersheds, + load_precomputed_watersheds, + resolve_lulc_raster_paths, + write_vector_output, +) +from computing.utils import ( + push_shape_to_geoserver, + save_layer_info_to_db, + update_layer_sync_status, +) +GEOSERVER_WORKSPACE = "lulc_vector" +LOCAL_ALGORITHM = "local_lulc_vector" +LOCAL_ALGORITHM_VERSION = "local-1.0" + +LULC_VECTOR_CLASS_DEFINITIONS = ( + {"value": 1, "label": "built-up_area_"}, + {"value": 2, "label": "k_water_area_"}, + {"value": 3, "label": "kr_water_area_"}, + {"value": 4, "label": "krz_water_area_"}, + {"value": 5, "label": "cropland_area_"}, + {"value": 6, "label": "tree_forest_area_"}, + {"value": 7, "label": "barrenlands_area_"}, + {"value": 8, "label": "single_kharif_cropped_area_"}, + {"value": 9, "label": "single_non_kharif_cropped_area_"}, + {"value": 10, "label": "doubly_cropped_area_"}, + {"value": 11, "label": "triply_cropped_area_"}, + {"value": 12, "label": "shrub_scrub_area_"}, +) + + +def _slug(value, fallback): + text = str(value).strip().lower() + return valid_gee_text(text) or fallback + + +def _layer_name(district, block): + return f"lulc_vector_{_slug(district, 'unknown_district')}_{_slug(block, 'unknown_block')}" + + +def _year_class_definitions(year): + return [ + { + "value": class_definition["value"], + "label": f"{class_definition['label']}{year}", + } + for class_definition in LULC_VECTOR_CLASS_DEFINITIONS + ] + + +def run_lulc_vector_local( + state, + district, + block, + start_year, + end_year, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + lulc_dir=LULC_BASE_DIR, + push_to_geoserver=True, + sync_layer_metadata=True, +): + state = str(state).strip().lower() + district = str(district).strip().lower() + block = str(block).strip().lower() + 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") + + watersheds_gdf, watershed_source = load_precomputed_watersheds( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + ) + print(f"Watershed boundary source: {watershed_source}") + + layer_name = _layer_name(district, block) + result_gdf = watersheds_gdf.copy() + raster_paths = resolve_lulc_raster_paths( + start_year=start_year, + end_year=end_year, + lulc_dir=lulc_dir, + ) + + for year, raster_path in zip(range(start_year, end_year + 1), raster_paths): + print(f"Computing local LULC vector properties for {year}-{year + 1}: {raster_path}") + year_result = compute_categorical_raster_areas_for_watersheds( + watersheds_gdf=result_gdf, + raster_path=raster_path, + class_definitions=_year_class_definitions(year), + ) + for column in year_result.columns: + if column not in result_gdf.columns: + result_gdf[column] = year_result[column] + + output_path = build_output_vector_path( + layer_name=layer_name, + state=state, + district=district, + block=block, + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + block_fallback="unknown_block", + ) + asset_id = write_vector_output( + gdf=result_gdf, + output_path=output_path, + layer_name=layer_name, + ) + print(f"Saved local LULC vector: {asset_id}") + + if push_to_geoserver: + geoserver_response = push_shape_to_geoserver( + os.path.splitext(asset_id)[0], + workspace=GEOSERVER_WORKSPACE, + layer_name=layer_name, + file_type="gpkg", + ) + print(f"GeoServer response: {geoserver_response}") + if not isinstance(geoserver_response, dict) or geoserver_response.get( + "status_code" + ) not in (200, 201): + return False + + 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="LULC", + misc={ + "start_year": start_year, + "end_year": end_year, + "is_generated_locally": True, + }, + algorithm=LOCAL_ALGORITHM, + algorithm_version=LOCAL_ALGORITHM_VERSION, + ) + if layer_id and push_to_geoserver: + update_layer_sync_status(layer_id=layer_id, sync_to_geoserver=True) + + return True + + +def _vectorise_lulc_local_task( + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = gee_account_id + return run_lulc_vector_local( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + +@app.task(bind=True) +def vectorise_lulc( + self, + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = self + return _vectorise_lulc_local_task( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + gee_account_id=gee_account_id, + ) diff --git a/computing/lulc_X_terrain/lulc_on_plain_cluster_local.py b/computing/lulc_X_terrain/lulc_on_plain_cluster_local.py new file mode 100644 index 00000000..5427d096 --- /dev/null +++ b/computing/lulc_X_terrain/lulc_on_plain_cluster_local.py @@ -0,0 +1,399 @@ +import os +from contextlib import ExitStack + +import numpy as np +import pandas as pd +import rasterio +from rasterio.mask import mask +from rasterio.warp import Resampling, reproject +from utilities.gee_utils import valid_gee_text + +from nrm_app.celery import app + +from computing.local_compute_helper import ( + AEZ_VECTOR_PATH, + LULC_BASE_DIR, + MIN_WATERSHED_AREA_HA, + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + PROJECT_ROOT, + TERRAIN_RASTER_PATH, + build_output_vector_path as _build_output_vector_path, + compute_mode_lulc_array as _compute_mode_lulc_array, + compute_terrain_properties_for_watersheds, + ensure_file_exists as _ensure_file_exists, + filter_large_watersheds as _filter_large_watersheds, + load_precomputed_watersheds as _load_precomputed_watersheds, + resolve_aez_code as _resolve_aez_code, + resolve_lulc_raster_paths as _resolve_lulc_raster_paths, + validate_geometry as _validate_geometry, + write_vector_output as _write_output_vector, +) +from computing.utils import ( + push_shape_to_geoserver, + save_layer_info_to_db, + update_layer_sync_status, +) + +from .utils import aez_lulcXterrain_cluster_centroids + + +LOCAL_OUTPUT_BASE_DIR = PROJECT_ROOT / "data/lulc_X_terrain/lulc_plain_clusters_local" +GEOSERVER_WORKSPACE = "terrain_lulc" + +PLAIN_LULC_FIELD_MAPPING = { + "barren": 7, + "double_crop": 10, + "shrubs_scrubs": 12, + "sing_crop": 8, + "sing_non_kharif_crop": 9, + "forest": 6, + "triple_crop": 11, +} + + +def _compute_plain_lulc_properties_for_feature(geometry, terrain_src, lulc_sources): + try: + terrain_data, terrain_transform = mask( + terrain_src, + [geometry.__geo_interface__], + crop=True, + filled=True, + ) + except ValueError: + return None + + terrain_values = np.rint(terrain_data[0]).astype(np.int16, copy=False) + if terrain_values.size == 0: + return None + + terrain_valid_mask = np.ones(terrain_values.shape, dtype=bool) + if terrain_src.nodata is not None: + terrain_valid_mask &= terrain_values != int(round(terrain_src.nodata)) + terrain_valid_mask &= terrain_values > 0 + + if not terrain_valid_mask.any(): + return None + + reprojected_lulc = [] + for lulc_src in lulc_sources: + destination = np.zeros(terrain_values.shape, dtype=np.float32) + reproject( + source=rasterio.band(lulc_src, 1), + destination=destination, + src_transform=lulc_src.transform, + src_crs=lulc_src.crs, + src_nodata=lulc_src.nodata, + dst_transform=terrain_transform, + dst_crs=terrain_src.crs, + dst_nodata=0, + resampling=Resampling.nearest, + ) + reprojected_lulc.append(destination) + + lulc_mode = _compute_mode_lulc_array(reprojected_lulc) + plains_mask = terrain_valid_mask & (terrain_values == 5) + slopy_mask = terrain_valid_mask & (terrain_values == 6) + plain_plus_slope_pixels = int((plains_mask | slopy_mask).sum()) + + if plain_plus_slope_pixels == 0: + return { + "LxP_cluster": -1, + "clust_name": "No plain or slope pixels", + "barren": 0.0, + "double_crop": 0.0, + "shrubs_scrubs": 0.0, + "sing_crop": 0.0, + "sing_non_kharif_crop": 0.0, + "forest": 0.0, + "triple_crop": 0.0, + } + + proportions = {} + for field_name, lulc_class in PLAIN_LULC_FIELD_MAPPING.items(): + class_pixels = int((plains_mask & (lulc_mode == lulc_class)).sum()) + proportions[field_name] = class_pixels / float(plain_plus_slope_pixels) + + return proportions + + +def _assign_plain_clusters( + plain_watersheds_gdf, + terrain_raster_path, + plain_centroids, + lulc_raster_paths, +): + centroid_vectors = np.array( + [ + plain_centroids[str(index)]["cluster_vector"] + for index in range(len(plain_centroids)) + ], + dtype=np.float64, + ) + + computed_rows = [] + with rasterio.open(terrain_raster_path) as terrain_src, ExitStack() as stack: + lulc_sources = [ + stack.enter_context(rasterio.open(path)) + for path in lulc_raster_paths + ] + total = len(plain_watersheds_gdf) + + for index, row in enumerate( + plain_watersheds_gdf.itertuples(index=False), + start=1, + ): + properties = _compute_plain_lulc_properties_for_feature( + geometry=row.geometry, + terrain_src=terrain_src, + lulc_sources=lulc_sources, + ) + + if properties is None: + computed_rows.append( + { + "LxP_cluster": -1, + "clust_name": "No valid terrain pixels", + "barren": 0.0, + "double_crop": 0.0, + "shrubs_scrubs": 0.0, + "sing_crop": 0.0, + "sing_non_kharif_crop": 0.0, + "forest": 0.0, + "triple_crop": 0.0, + } + ) + elif properties.get("LxP_cluster") == -1: + computed_rows.append(properties) + else: + feature_vector = np.array( + [ + properties["barren"], + properties["double_crop"], + properties["shrubs_scrubs"], + properties["sing_crop"], + properties["sing_non_kharif_crop"], + properties["forest"], + properties["triple_crop"], + ], + dtype=np.float64, + ) + distances = np.sum( + (centroid_vectors - feature_vector) ** 2, + axis=1, + ) + cluster_index = int(np.argmin(distances)) + + computed_rows.append( + { + "LxP_cluster": cluster_index, + "clust_name": plain_centroids[str(cluster_index)][ + "cluster_name" + ], + "barren": properties["barren"] * 100.0, + "double_crop": properties["double_crop"] * 100.0, + "shrubs_scrubs": properties["shrubs_scrubs"] * 100.0, + "sing_crop": properties["sing_crop"] * 100.0, + "sing_non_kharif_crop": ( + properties["sing_non_kharif_crop"] * 100.0 + ), + "forest": properties["forest"] * 100.0, + "triple_crop": properties["triple_crop"] * 100.0, + } + ) + + if index % 200 == 0 or index == total: + print( + f"Computed plain LULC clusters for {index}/{total} watersheds" + ) + + result = plain_watersheds_gdf.copy() + computed_df = pd.DataFrame(computed_rows) + for column in computed_df.columns: + result[column] = computed_df[column].values + return result + + +def run_lulc_on_plain_cluster_local( + state, + district, + block, + start_year, + end_year, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + aez_vector_path=AEZ_VECTOR_PATH, + lulc_dir=LULC_BASE_DIR, + terrain_raster_path=TERRAIN_RASTER_PATH, + push_to_geoserver=True, + sync_layer_metadata=True, +): + state = str(state).strip() + district = str(district).strip() + block = str(block).strip() + start_year = int(start_year) + end_year = int(end_year) + + if start_year > end_year: + raise ValueError("start_year cannot be greater than end_year") + + _ensure_file_exists(terrain_raster_path, "Terrain raster") + lulc_raster_paths = _resolve_lulc_raster_paths( + start_year=start_year, + end_year=end_year, + lulc_dir=lulc_dir, + ) + + watersheds_gdf, watershed_source = _load_precomputed_watersheds( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + ) + watersheds_gdf = _validate_geometry(watersheds_gdf) + if watersheds_gdf.empty: + raise ValueError("No valid watershed geometries found for local processing.") + + aez_code = _resolve_aez_code( + watersheds_gdf, + aez_vector_path=aez_vector_path, + ) + plain_centroids = aez_lulcXterrain_cluster_centroids[f"aez{aez_code}"]["plains"] + + filtered_watersheds = _filter_large_watersheds(watersheds_gdf) + if filtered_watersheds.empty: + raise ValueError( + f"No watersheds larger than {MIN_WATERSHED_AREA_HA} ha found for {state}/{district}/{block}." + ) + + terrain_classified = compute_terrain_properties_for_watersheds( + watersheds_gdf=filtered_watersheds, + raster_path=str(terrain_raster_path), + ) + terrain_classified["terrain_cluster"] = terrain_classified[ + "terrainClusters" + ].astype(int) + plain_watersheds = terrain_classified.loc[ + terrain_classified["terrain_cluster"] != 2 + ].copy() + if plain_watersheds.empty: + raise ValueError( + f"No plain-cluster watersheds found for {state}/{district}/{block}." + ) + + temp_columns = [ + "terrainClusters", + "plain_area", + "valley_area", + "hill_slopes_area", + "ridge_area", + "slopy_area", + ] + plain_watersheds.drop( + columns=[ + column + for column in temp_columns + if column in plain_watersheds.columns + ], + inplace=True, + ) + + result_gdf = _assign_plain_clusters( + plain_watersheds_gdf=plain_watersheds, + terrain_raster_path=str(terrain_raster_path), + plain_centroids=plain_centroids, + lulc_raster_paths=lulc_raster_paths, + ) + + layer_name = ( + f"{valid_gee_text(str(district).strip().lower())}_" + f"{valid_gee_text(str(block).strip().lower())}_lulc_plain" + ) + output_path = _build_output_vector_path( + layer_name=layer_name, + state=state, + district=district, + block=block, + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + block_fallback="unknown_block", + ) + asset_id = _write_output_vector( + gdf=result_gdf, + output_path=output_path, + layer_name=layer_name, + ) + print(f"Saved local plain-cluster vector: {asset_id}") + print(f"Watershed boundary source: {watershed_source}") + print(f"Resolved AEZ code: {aez_code}") + + geoserver_response = None + if push_to_geoserver: + geoserver_response = push_shape_to_geoserver( + os.path.splitext(asset_id)[0], + workspace=GEOSERVER_WORKSPACE, + layer_name=layer_name, + file_type="gpkg", + ) + print(f"GeoServer response: {geoserver_response}") + + if not isinstance(geoserver_response, dict) or geoserver_response.get( + "status_code" + ) not in (200, 201): + return False + + 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="Terrain LULC", + misc={"start_year": start_year, "end_year": end_year}, + ) + if layer_id and push_to_geoserver: + update_layer_sync_status( + layer_id=layer_id, + sync_to_geoserver=True, + ) + + return True + + +def _generate_lulc_on_plain_cluster_local_task( + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = gee_account_id + return run_lulc_on_plain_cluster_local( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + +@app.task(bind=True) +def lulc_on_plain_cluster_local( + self, + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = self + return _generate_lulc_on_plain_cluster_local_task( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + gee_account_id=gee_account_id, + ) diff --git a/computing/lulc_X_terrain/lulc_on_slope_cluster_local.py b/computing/lulc_X_terrain/lulc_on_slope_cluster_local.py new file mode 100644 index 00000000..96d9ecc1 --- /dev/null +++ b/computing/lulc_X_terrain/lulc_on_slope_cluster_local.py @@ -0,0 +1,391 @@ +import os + +import numpy as np +import pandas as pd +import rasterio +from rasterio.mask import mask +from rasterio.warp import Resampling, reproject +from utilities.gee_utils import valid_gee_text + +from nrm_app.celery import app + +from computing.config_loader import LULC_SLOPE_CLUSTER_OUTPUT_DIR as LOCAL_OUTPUT_BASE_DIR +from computing.local_compute_helper import ( + AEZ_VECTOR_PATH, + LULC_BASE_DIR, + MIN_WATERSHED_AREA_HA, + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + TERRAIN_RASTER_PATH, + build_output_vector_path as _build_output_vector_path, + compute_mode_lulc_array as _compute_mode_lulc_array, + compute_terrain_properties_for_watersheds, + ensure_file_exists as _ensure_file_exists, + filter_large_watersheds as _filter_large_watersheds, + load_precomputed_watersheds as _load_precomputed_watersheds, + resolve_aez_code as _resolve_aez_code, + resolve_lulc_raster_paths as _resolve_lulc_raster_paths, + validate_geometry as _validate_geometry, + write_vector_output as _write_output_vector, +) +from computing.utils import ( + push_shape_to_geoserver, + save_layer_info_to_db, + update_layer_sync_status, +) + +from .utils import aez_lulcXterrain_cluster_centroids + + +GEOSERVER_WORKSPACE = "terrain_lulc" +SLOPE_LULC_FIELD_MAPPING = { + "barren": 7, + "double": 10, + "shrub_scrub": 12, + "sing_kharif": 8, + "sing_non_kharif": 9, + "forests": 6, + "triple": 11, +} + + +def _compute_slope_lulc_properties_for_feature(geometry, terrain_src, lulc_sources): + try: + terrain_data, terrain_transform = mask( + terrain_src, + [geometry.__geo_interface__], + crop=True, + filled=True, + ) + except ValueError: + return None + + terrain_values = np.rint(terrain_data[0]).astype(np.int16, copy=False) + if terrain_values.size == 0: + return None + + terrain_valid_mask = np.ones(terrain_values.shape, dtype=bool) + if terrain_src.nodata is not None: + terrain_valid_mask &= terrain_values != int(round(terrain_src.nodata)) + terrain_valid_mask &= terrain_values > 0 + + if not terrain_valid_mask.any(): + return None + + reprojected_lulc = [] + for lulc_src in lulc_sources: + destination = np.zeros(terrain_values.shape, dtype=np.float32) + reproject( + source=rasterio.band(lulc_src, 1), + destination=destination, + src_transform=lulc_src.transform, + src_crs=lulc_src.crs, + src_nodata=lulc_src.nodata, + dst_transform=terrain_transform, + dst_crs=terrain_src.crs, + dst_nodata=0, + resampling=Resampling.nearest, + ) + reprojected_lulc.append(destination) + + lulc_mode = _compute_mode_lulc_array(reprojected_lulc) + plains_mask = terrain_valid_mask & (terrain_values == 5) + slopy_mask = terrain_valid_mask & (terrain_values == 6) + plain_plus_slope_pixels = int((plains_mask | slopy_mask).sum()) + + if plain_plus_slope_pixels == 0: + return { + "LxS_cluster": -1, + "clust_name": "No plain or slope pixels", + "barren": 0.0, + "double": 0.0, + "shrub_scrub": 0.0, + "sing_kharif": 0.0, + "sing_non_kharif": 0.0, + "forests": 0.0, + "triple": 0.0, + } + + proportions = {} + for field_name, lulc_class in SLOPE_LULC_FIELD_MAPPING.items(): + class_pixels = int((slopy_mask & (lulc_mode == lulc_class)).sum()) + proportions[field_name] = class_pixels / float(plain_plus_slope_pixels) + + return proportions + + +def _assign_slope_clusters( + slope_watersheds_gdf, + terrain_raster_path, + slope_centroids, + lulc_raster_paths, +): + centroid_vectors = np.array( + [ + slope_centroids[str(index)]["cluster_vector"] + for index in range(len(slope_centroids)) + ], + dtype=np.float64, + ) + + computed_rows = [] + with rasterio.open(terrain_raster_path) as terrain_src: + lulc_sources = [rasterio.open(path) for path in lulc_raster_paths] + try: + total = len(slope_watersheds_gdf) + for index, row in enumerate( + slope_watersheds_gdf.itertuples(index=False), + start=1, + ): + properties = _compute_slope_lulc_properties_for_feature( + geometry=row.geometry, + terrain_src=terrain_src, + lulc_sources=lulc_sources, + ) + + if properties is None: + computed_rows.append( + { + "LxS_cluster": -1, + "clust_name": "No valid terrain pixels", + "barren": 0.0, + "double": 0.0, + "shrub_scrub": 0.0, + "sing_kharif": 0.0, + "sing_non_kharif": 0.0, + "forests": 0.0, + "triple": 0.0, + } + ) + elif properties.get("LxS_cluster") == -1: + computed_rows.append(properties) + else: + feature_vector = np.array( + [ + properties["barren"], + properties["shrub_scrub"], + properties["forests"], + ], + dtype=np.float64, + ) + distances = np.sum( + (centroid_vectors - feature_vector) ** 2, + axis=1, + ) + cluster_index = int(np.argmin(distances)) + + computed_rows.append( + { + "LxS_cluster": cluster_index, + "clust_name": slope_centroids[str(cluster_index)][ + "cluster_name" + ], + "barren": properties["barren"] * 100.0, + "double": properties["double"] * 100.0, + "shrub_scrub": properties["shrub_scrub"] * 100.0, + "sing_kharif": properties["sing_kharif"] * 100.0, + "sing_non_kharif": ( + properties["sing_non_kharif"] * 100.0 + ), + "forests": properties["forests"] * 100.0, + "triple": properties["triple"] * 100.0, + } + ) + + if index % 200 == 0 or index == total: + print( + f"Computed slope LULC clusters for {index}/{total} watersheds" + ) + finally: + for lulc_src in lulc_sources: + lulc_src.close() + + result = slope_watersheds_gdf.copy() + computed_df = pd.DataFrame(computed_rows) + for column in computed_df.columns: + result[column] = computed_df[column].values + return result + + +def run_lulc_on_slope_cluster_local( + state, + district, + block, + start_year, + end_year, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + aez_vector_path=AEZ_VECTOR_PATH, + lulc_dir=LULC_BASE_DIR, + terrain_raster_path=TERRAIN_RASTER_PATH, + push_to_geoserver=True, + sync_layer_metadata=True, +): + state = str(state).strip() + district = str(district).strip() + block = str(block).strip() + start_year = int(start_year) + end_year = int(end_year) + + if start_year > end_year: + raise ValueError("start_year cannot be greater than end_year") + + _ensure_file_exists(terrain_raster_path, "Terrain raster") + lulc_raster_paths = _resolve_lulc_raster_paths( + start_year=start_year, + end_year=end_year, + lulc_dir=lulc_dir, + ) + + watersheds_gdf, watershed_source = _load_precomputed_watersheds( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + ) + watersheds_gdf = _validate_geometry(watersheds_gdf) + if watersheds_gdf.empty: + raise ValueError("No valid watershed geometries found for local processing.") + + aez_code = _resolve_aez_code( + watersheds_gdf, + aez_vector_path=aez_vector_path, + ) + slope_centroids = aez_lulcXterrain_cluster_centroids[f"aez{aez_code}"]["slopes"] + + filtered_watersheds = _filter_large_watersheds(watersheds_gdf) + if filtered_watersheds.empty: + raise ValueError( + f"No watersheds larger than {MIN_WATERSHED_AREA_HA} ha found for {state}/{district}/{block}." + ) + + terrain_classified = compute_terrain_properties_for_watersheds( + watersheds_gdf=filtered_watersheds, + raster_path=str(terrain_raster_path), + ) + terrain_classified["terrain_cluster"] = terrain_classified[ + "terrainClusters" + ].astype(int) + slope_watersheds = terrain_classified.loc[ + terrain_classified["terrain_cluster"].isin([0, 3]) + ].copy() + if slope_watersheds.empty: + raise ValueError( + f"No slope-cluster watersheds found for {state}/{district}/{block}." + ) + + temp_columns = [ + "terrainClusters", + "plain_area", + "valley_area", + "hill_slopes_area", + "ridge_area", + "slopy_area", + ] + slope_watersheds.drop( + columns=[ + column + for column in temp_columns + if column in slope_watersheds.columns + ], + inplace=True, + ) + + result_gdf = _assign_slope_clusters( + slope_watersheds_gdf=slope_watersheds, + terrain_raster_path=str(terrain_raster_path), + slope_centroids=slope_centroids, + lulc_raster_paths=lulc_raster_paths, + ) + + layer_name = ( + f"{valid_gee_text(str(district).strip().lower())}_" + f"{valid_gee_text(str(block).strip().lower())}_lulc_slope" + ) + output_path = _build_output_vector_path( + layer_name=layer_name, + state=state, + district=district, + block=block, + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + block_fallback="unknown_block", + ) + asset_id = _write_output_vector( + gdf=result_gdf, + output_path=output_path, + layer_name=layer_name, + ) + print(f"Saved local slope-cluster vector: {asset_id}") + print(f"Watershed boundary source: {watershed_source}") + print(f"Resolved AEZ code: {aez_code}") + + if push_to_geoserver: + geoserver_response = push_shape_to_geoserver( + str(output_path.with_suffix("")), + workspace=GEOSERVER_WORKSPACE, + layer_name=layer_name, + file_type="gpkg", + ) + print(f"GeoServer response: {geoserver_response}") + + if not isinstance(geoserver_response, dict) or geoserver_response.get( + "status_code" + ) not in (200, 201): + return False + + 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="Terrain LULC", + misc={"start_year": start_year, "end_year": end_year, "is_generated_locally": True}, + ) + if layer_id and push_to_geoserver: + update_layer_sync_status( + layer_id=layer_id, + sync_to_geoserver=True, + ) + + return True + + +def _generate_lulc_on_slope_cluster_local_task( + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = gee_account_id + return run_lulc_on_slope_cluster_local( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + +@app.task(bind=True) +def lulc_on_slope_cluster_local( + self, + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = self + return _generate_lulc_on_slope_cluster_local_task( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + gee_account_id=gee_account_id, + ) diff --git a/computing/misc/aquifer_vector_local.py b/computing/misc/aquifer_vector_local.py new file mode 100644 index 00000000..47f49492 --- /dev/null +++ b/computing/misc/aquifer_vector_local.py @@ -0,0 +1,435 @@ +import os + +import pandas as pd +from utilities.gee_utils import valid_gee_text + +from nrm_app.celery import app + +from computing.config_loader import ( + AQUIFER_VECTOR_PATH, + AQUIFER_VECTOR_OUTPUT_DIR as LOCAL_OUTPUT_BASE_DIR, +) +from computing.local_compute_helper import ( + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + build_output_vector_path, + get_watershed_areas_in_hectares, + load_precomputed_watersheds, + read_validated_vector_file, + validate_geometry, + write_vector_output, +) +from computing.utils import ( + push_shape_to_geoserver, + save_layer_info_to_db, + update_layer_sync_status, +) +GEOSERVER_WORKSPACE = "aquifer" + +YIELD_VALUE_MAP = { + "": None, + "-": None, + "Upto 2%": 0.02, + "1-2%": 0.02, + "Upto 1.5%": 0.015, + "Upto 3%": 0.03, + "Upto 2.5%": 0.025, + "6 - 8%": 0.08, + "1-1.5%": 0.015, + "2-3%": 0.03, + "Upto 4%": 0.04, + "Upto 5%": 0.05, + "Upto -3.5%": 0.035, + "Upto 3 %": 0.03, + "Upto 9%": 0.09, + "1-2.5": 0.025, + "Upto 1.2%": 0.012, + "Upto 5-2%": 0.05, + "Upto 1%": 0.01, + "Up to 1.5%": 0.015, + "Upto 8%": 0.08, + "Upto 6%": 0.06, + "0.08": 0.08, + "8 - 16%": 0.16, + "Not Explored": None, + "8 - 15%": 0.15, + "6 - 10%": 0.10, + "6 - 15%": 0.15, + "8 - 20%": 0.20, + "8 - 10%": 0.10, + "6 - 12%": 0.12, + "6 - 16%": 0.16, + "8 - 12%": 0.12, + "8 - 18%": 0.18, + "Upto 3.5%": 0.035, + "Upto 15%": 0.15, + "1.5-2%": 0.02, +} + +PRINCIPAL_AQUIFERS = [ + "Laterite", + "Basalt", + "Sandstone", + "Shale", + "Limestone", + "Granite", + "Schist", + "Quartzite", + "Charnockite", + "Khondalite", + "Banded Gneissic Complex", + "Gneiss", + "Intrusive", + "Alluvium", + "None", +] + +PRINCIPAL_AQUIFER_PERCENT_COLUMNS = [ + f"principle_aq_{aquifer_name}_percent" + for aquifer_name in PRINCIPAL_AQUIFERS +] + + +def _safe_string(value): + if pd.isna(value): + return "" + return str(value).strip() + + +def _safe_int(value): + if pd.isna(value): + return None + return int(float(value)) + + +def _normalize_principal_name(value): + principal_name = _safe_string(value) + return principal_name or "None" + + +def _map_yield_value(value): + return YIELD_VALUE_MAP.get(_safe_string(value)) + + +def _build_empty_aquifer_properties(watershed_row, area_in_ha): + properties = { + "uid": watershed_row.get("uid"), + "id": watershed_row.get("id", watershed_row.get("uid")), + "area_in_ha": float(area_in_ha) if pd.notna(area_in_ha) else 0.0, + "total_weighted_yield": 0.0, + "%_area_aquifer": 0.0, + "aquifer_count": 0, + "aquifer_class": "No Data", + "Age": "", + "Lithology_": None, + "Major_Aq_1": "", + "Major_Aqui": "", + "Principal_": "", + "Recommende": None, + "yeild__": "", + "zone_m": "", + "y_value": None, + } + properties.update( + { + column_name: 0.0 + for column_name in PRINCIPAL_AQUIFER_PERCENT_COLUMNS + } + ) + return properties + + +def _build_aquifer_properties(watershed_row, area_in_ha, intersections_df): + properties = _build_empty_aquifer_properties(watershed_row, area_in_ha) + dominant_aquifer = intersections_df.sort_values( + "intersection_area_m2", + ascending=False, + ).iloc[0] + pct_by_aquifer = intersections_df.groupby("principal_name")[ + "percent_area_aquifer" + ].sum() + + for aquifer_name in PRINCIPAL_AQUIFERS: + properties[f"principle_aq_{aquifer_name}_percent"] = float( + pct_by_aquifer.get(aquifer_name, 0.0) + ) + + principal_value = _safe_string(dominant_aquifer["Principal_"]) + properties.update( + { + "total_weighted_yield": float( + intersections_df["weighted_contribution"].sum() + ), + "%_area_aquifer": float(dominant_aquifer["percent_area_aquifer"]), + "aquifer_count": int(len(intersections_df)), + "aquifer_class": ( + "Alluvium" + if principal_value == "Alluvium" + else "Hard-Rock" + ), + "Age": _safe_string(dominant_aquifer["Age"]), + "Lithology_": _safe_int(dominant_aquifer["Lithology_"]), + "Major_Aq_1": _safe_string(dominant_aquifer["Major_Aq_1"]), + "Major_Aqui": _safe_string(dominant_aquifer["Major_Aqui"]), + "Principal_": principal_value, + "Recommende": _safe_int(dominant_aquifer["Recommende"]), + "yeild__": _safe_string(dominant_aquifer["yeild__"]), + "zone_m": _safe_string(dominant_aquifer["zone_m"]), + "y_value": ( + float(dominant_aquifer["y_value"]) + if pd.notna(dominant_aquifer["y_value"]) + else None + ), + } + ) + return properties + + +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.") + + aquifers_with_yield = aquifers_gdf.copy() + aquifers_with_yield["y_value"] = aquifers_with_yield["yeild__"].apply( + _map_yield_value + ) + aquifers_with_yield = aquifers_with_yield.loc[ + aquifers_with_yield["y_value"].notna() + ].copy() + 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") + aquifer_spatial_index = aquifers_projected.sindex + + computed_rows = [] + total = len(watersheds_projected) + + for index, watershed_idx in enumerate(watersheds_projected.index, start=1): + watershed_geometry = watersheds_projected.at[watershed_idx, "geometry"] + watershed_row = watersheds_result.loc[watershed_idx] + area_in_ha = watersheds_result.at[watershed_idx, "area_in_ha"] + + if watershed_geometry is None or watershed_geometry.is_empty: + computed_rows.append( + _build_empty_aquifer_properties(watershed_row, area_in_ha) + ) + continue + + watershed_area_m2 = float(watershed_geometry.area) + if watershed_area_m2 <= 0: + computed_rows.append( + _build_empty_aquifer_properties(watershed_row, area_in_ha) + ) + continue + + 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(): + intersection_geometry = watershed_geometry.intersection( + aquifer_row.geometry + ) + if intersection_geometry.is_empty: + continue + + intersection_area_m2 = float(intersection_geometry.area) + if intersection_area_m2 <= 0: + continue + + fraction = intersection_area_m2 / watershed_area_m2 + intersections.append( + { + "intersection_area_m2": intersection_area_m2, + "percent_area_aquifer": fraction * 100.0, + "weighted_contribution": fraction * float(aquifer_row["y_value"]), + "principal_name": _normalize_principal_name( + aquifer_row["Principal_"] + ), + "Age": aquifer_row["Age"], + "Lithology_": aquifer_row["Lithology_"], + "Major_Aq_1": aquifer_row["Major_Aq_1"], + "Major_Aqui": aquifer_row["Major_Aqui"], + "Principal_": aquifer_row["Principal_"], + "Recommende": aquifer_row["Recommende"], + "yeild__": aquifer_row["yeild__"], + "zone_m": aquifer_row["zone_m"], + "y_value": aquifer_row["y_value"], + } + ) + + if intersections: + intersections_df = pd.DataFrame(intersections) + computed_rows.append( + _build_aquifer_properties( + watershed_row, + area_in_ha, + intersections_df, + ) + ) + else: + computed_rows.append( + _build_empty_aquifer_properties(watershed_row, area_in_ha) + ) + + if index % 200 == 0 or index == total: + print(f"Computed aquifer properties for {index}/{total} watersheds") + + computed_df = pd.DataFrame(computed_rows) + for column in computed_df.columns: + watersheds_result[column] = computed_df[column].values + return watersheds_result + + +def run_aquifer_vector_local( + state, + district, + block, + aquifer_vector_path=AQUIFER_VECTOR_PATH, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + push_to_geoserver=True, + sync_layer_metadata=True, +): + state = str(state).strip().lower() + district = str(district).strip().lower() + block = str(block).strip().lower() + + watersheds_gdf, watershed_source = load_precomputed_watersheds( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + ) + aquifers_gdf = read_validated_vector_file( + aquifer_vector_path, + f"Aquifer source file has no valid geometries: {aquifer_vector_path}", + ) + + result_gdf = _compute_aquifer_properties_for_watersheds( + watersheds_gdf=watersheds_gdf, + aquifers_gdf=aquifers_gdf, + ) + + layer_name = ( + f"aquifer_vector_{valid_gee_text(district.lower())}_" + f"{valid_gee_text(block.lower())}" + ) + output_path = build_output_vector_path( + layer_name=layer_name, + state=state, + district=district, + block=block, + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + block_fallback="unknown_block", + ) + asset_id = write_vector_output( + gdf=result_gdf, + output_path=output_path, + layer_name=layer_name, + ) + print(f"Saved local aquifer vector: {asset_id}") + print(f"Watershed boundary source: {watershed_source}") + + if push_to_geoserver: + geoserver_response = push_shape_to_geoserver( + os.path.splitext(asset_id)[0], + workspace=GEOSERVER_WORKSPACE, + layer_name=layer_name, + file_type="gpkg", + ) + print(f"GeoServer response: {geoserver_response}") + if not isinstance(geoserver_response, dict) or geoserver_response.get( + "status_code" + ) not in (200, 201): + return False + + if sync_layer_metadata: + from computing.STAC_specs import generate_STAC_layerwise + + layer_id = save_layer_info_to_db( + state=state, + district=district, + block=block, + layer_name=layer_name, + asset_id=asset_id, + dataset_name="Aquifer", + misc={"is_generated_locally": True}, + ) + if layer_id and push_to_geoserver: + update_layer_sync_status(layer_id=layer_id, sync_to_geoserver=True) + layer_stac_generated = generate_STAC_layerwise.generate_vector_stac( + state=state, + district=district, + block=block, + layer_name="aquifer_vector", + ) + update_layer_sync_status( + layer_id=layer_id, + is_stac_specs_generated=layer_stac_generated, + ) + + return True + + +def _generate_aquifer_vector_local_task( + state, + district, + block, + gee_account_id=None, +): + _ = gee_account_id + return run_aquifer_vector_local( + state=state, + district=district, + block=block, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + +@app.task(bind=True) +def generate_aquifer_vector( + self, + state, + district, + block, + gee_account_id=None, +): + _ = self + return _generate_aquifer_vector_local_task( + state=state, + district=district, + block=block, + gee_account_id=gee_account_id, + ) 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/soil_health/soil_health.py b/computing/soil_health/soil_health.py new file mode 100644 index 00000000..c11cb120 --- /dev/null +++ b/computing/soil_health/soil_health.py @@ -0,0 +1,146 @@ +import os +import geopandas as gpd + +from computing.soil_health.soil_health_helper import nutrient_stats_for_geometries +from utilities.gee_utils import valid_gee_text +from computing.local_compute_helper import ( + PROJECT_ROOT, + build_output_vector_path, + clip_raster_with_roi, + build_output_raster_path, + read_validated_vector_file, + validate_geometry, + write_vector_output, + push_local_vector_to_geoserver, + push_local_raster_to_geoserver, +) + +ROI_PATH = str(PROJECT_ROOT / "data/yalburga_mws.json") +LOCAL_OUTPUT_BASE_DIR = "data/soil_health" +GEOSERVER_STYLE = "" +GEOSERVER_WORKSPACE = "soil_health" +NUTRIENTS = ["N", "K", "P", "OC"] +NUTRIENT_PERCENTILES = (5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 95) + + +def clip_soil_health_raster( + state=None, + district=None, + block=None, + asset_suffix=None, + roi=None, + precomputed_roi_dir=None, + push_to_geoserver=True, + sync_layer_metadata=False, +): + + asset_suffix, roi_gdf = get_roi(asset_suffix, block, district, roi, state) + layer_name = f"{asset_suffix}_soil_health_raster" + geoserver_statuses = [] + for nutrient in NUTRIENTS: + SOIL_MAP_PATH = str( + PROJECT_ROOT + / f"data/soil_health/AEZ_3_{nutrient}_20260311.tif" # TODO Replace with pan India local path + ) + output_raster_path = build_output_raster_path( + layer_name=f"{layer_name}_{nutrient}", + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + state=state, + district=district, + block=block, + ) + + clip_raster_with_roi( + roi_gdf, SOIL_MAP_PATH, output_raster_path, raster_label="Raster" + ) + + if push_to_geoserver: + upload_res, style_res = push_local_raster_to_geoserver( + file_path=output_raster_path, + layer_name=f"{layer_name}_{nutrient}", + workspace=GEOSERVER_WORKSPACE, + ) + print(f"GeoServer upload response for {nutrient}: {upload_res}") + geoserver_statuses.append(True) + + return all(geoserver_statuses) if push_to_geoserver else True # TODO Add Stac specs + + +def get_roi(asset_suffix, block, district, roi, state): + if state and district and block: + asset_suffix = f"{valid_gee_text(str(district).lower())}_{valid_gee_text(str(block).lower())}" + # roi_gdf = load_precomputed_roi( + # state=state, + # district=district, + # block=block, + # precomputed_roi_dir=precomputed_roi_dir, + # ) + roi_gdf = gpd.read_file(ROI_PATH) + else: + if not roi or not asset_suffix: + raise ValueError( + "For non state/district/block runs, both `roi` and `asset_suffix` are required." + ) + + roi_gdf = read_validated_vector_file( + roi, + f"ROI file has no valid geometries: {roi}", + ) + roi_gdf = validate_geometry(roi_gdf) + return asset_suffix, roi_gdf + + +def vectorize_soil_health( + state=None, + district=None, + block=None, + asset_suffix=None, + roi=None, + percentiles=NUTRIENT_PERCENTILES, + push_to_geoserver=True, +): + + asset_suffix, roi_gdf = get_roi(asset_suffix, block, district, roi, state) + layer_name = f"{asset_suffix}_soil_health" + + for nutrient in NUTRIENTS: + # This produces one output feature per ROI geometry with Nitrogen summary columns. + raster_path = build_output_raster_path( + layer_name=f"{layer_name}_raster_{nutrient}", + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + state=state, + district=district, + block=block, + ) + + result_gdf = nutrient_stats_for_geometries( + roi_gdf=roi_gdf, + raster_path=raster_path, + percentiles=tuple(percentiles), + nutrient=nutrient, + ) + + output_path = build_output_vector_path( + layer_name=f"{layer_name}_vector_{nutrient}", + state=state, + district=district, + block=block, + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + ) + write_vector_output( + gdf=result_gdf, + output_path=output_path, + layer_name=layer_name, + ) + print(f"Saved soil health vector: {output_path}") + + if push_to_geoserver: + geoserver_response = push_local_vector_to_geoserver( + path=os.path.splitext(output_path)[0], + layer_name=layer_name, + workspace=GEOSERVER_WORKSPACE, + file_type="gpkg", + ) + print(f"GeoServer response: {geoserver_response}") + + # TODO Add Stac specs diff --git a/computing/soil_health/soil_health_helper.py b/computing/soil_health/soil_health_helper.py new file mode 100644 index 00000000..8d96acb5 --- /dev/null +++ b/computing/soil_health/soil_health_helper.py @@ -0,0 +1,79 @@ +import numpy as np +import rasterio +from rasterio.mask import mask +from shapely.geometry import mapping + +from computing.local_compute_helper import ensure_file_exists + + +def nutrient_stats_for_geometries(roi_gdf, raster_path, percentiles, nutrient): + ensure_file_exists(raster_path, "Clipped soil health raster") + + with rasterio.open(raster_path) as src: + # Raster masking expects geometries in the raster CRS, so reproject a + # working copy and preserve the original ROI CRS for the output vector. + working_gdf = roi_gdf.copy() + if working_gdf.crs is None: + raise ValueError( + "ROI CRS is missing; cannot align with soil health raster." + ) + if src.crs and working_gdf.crs != src.crs: + working_gdf = working_gdf.to_crs(src.crs) + + nodata = src.nodata + rows = [] + total = len(working_gdf) + for index, row in enumerate(working_gdf.itertuples(index=False), start=1): + geom = row.geometry + if geom is None or geom.is_empty: + values = np.array([], dtype=np.float64) + else: + # Clip the raster to one ROI/watershed polygon. `filled=False` + # keeps rasterio's mask, which lets us drop pixels outside the geometry. + try: + clipped, _ = mask( + src, + [mapping(geom)], + crop=True, + filled=False, + ) + except ValueError: + clipped = None + + if clipped is None or clipped.size == 0: + values = np.array([], dtype=np.float64) + else: + data = clipped[0] + valid_mask = ~np.ma.getmaskarray(data) + values = np.asarray(data, dtype=np.float64) + + # Ignore pixels outside the polygon, nodata pixels, and any + # non-finite values before calculating nutrient statistics. + valid_mask &= np.isfinite(values) + if nodata is not None and not np.isnan(nodata): + valid_mask &= values != float(nodata) + values = values[valid_mask] + + stats = {f"{nutrient}_count": int(values.size)} + if values.size == 0: + # Use NULLs in the vector table when a polygon has no valid raster pixels. + stats[f"{nutrient}_mean"] = None + for percentile in percentiles: + stats[f"{nutrient}_p{percentile:02d}"] = None + else: + stats[f"{nutrient}_mean"] = float(np.mean(values)) + percentile_values = np.percentile(values, percentiles) + print(f"Percentile values: {percentile_values}") + for percentile, value in zip(percentiles, percentile_values): + print(f"Percentile: {percentile}, Value: {value}") + stats[f"{nutrient}_p{percentile:02d}"] = float(value) + rows.append(stats) + + if index % 200 == 0 or index == total: + print(f"Computed soil health stats for {index}/{total} geometries") + + result = roi_gdf.copy() + # Attach the computed columns back to the original ROI geometries for output. + for column in rows[0] if rows else []: + result[column] = [row[column] for row in rows] + return result diff --git a/computing/spei/__init__.py b/computing/spei/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/computing/spei/drought/__init__.py b/computing/spei/drought/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/computing/spei/drought/download_chirps_local.py b/computing/spei/drought/download_chirps_local.py new file mode 100644 index 00000000..e2bbbb3e --- /dev/null +++ b/computing/spei/drought/download_chirps_local.py @@ -0,0 +1,395 @@ +""" +Download drought inputs locally from Google Earth Engine. + +This script downloads pan-India or single-state GeoTIFFs locally for: + - UCSB-CHG/CHIRPS/DAILY precipitation + - MODIS/061/MOD16A2GF PET + +It downloads one small GeoTIFF per time step for each dataset. + +Examples: + main(aoi="india", datasets=["both"]) + main(aoi="india", datasets=["chirps", "modis_pet"], start_date="2004-01-01", end_date="2023-12-31") + main(aoi="state", state="Madhya Pradesh", datasets=["chirps"]) +""" + +from __future__ import annotations + +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +import ee +import requests + +from utilities.gee_utils import ee_initialize + +CHIRPS_COLLECTION = "UCSB-CHG/CHIRPS/DAILY" +MODIS_PET_COLLECTION = "MODIS/061/MOD16A2GF" +DEFAULT_PROJECT = "ee-corestackdev" +DATASET_CHOICES = ("chirps", "modis_pet", "both") + + +def initialize_earth_engine(project: str) -> None: + # Uses the repo's shared Earth Engine initialization helper. + # If auth/project errors happen, debug utilities.gee_utils.ee_initialize first. + ee_initialize() + print("Earth Engine initialized.") + + +def get_aoi(aoi_type: str, state_name: str) -> ee.FeatureCollection: + # GAUL level1 contains Indian state boundaries. For pan-India, keep all + # Indian state features and use their combined geometry downstream. + # admin = ee.FeatureCollection("FAO/GAUL/2015/level1").filter( + # ee.Filter.eq("ADM0_NAME", "India") + # ) + admin = ee.FeatureCollection( + "projects/ext-datasets/assets/datasets/State_pan_india" + ) + + if aoi_type == "india": + admin = admin.filter(ee.Filter.neq("Name", "Andaman & Nicobar")).filter( + ee.Filter.neq("Name", "Lakshadweep") + ) + return admin.union() + return admin.filter(ee.Filter.eq("Name", state_name)) + + +def date_range( + start_date: ee.Date, end_date_exclusive: ee.Date, unit: str +) -> list[ee.Date]: + # Builds monthly/daily date anchors. end_date_exclusive is intentionally + # advanced by one day in build_dataset_image so user input is inclusive. + count = end_date_exclusive.difference(start_date, unit).round().getInfo() + return [start_date.advance(offset, unit) for offset in range(count)] + + +def collection_dates( + collection: ee.ImageCollection, start_date: ee.Date, end_date_exclusive: ee.Date +) -> list[ee.Date]: + # Native cadence should use the source image timestamps. This is important + # for MODIS PET, which is 8-day rather than daily. + millis = ( + collection.filterDate(start_date, end_date_exclusive) + .aggregate_array("system:time_start") + .getInfo() + ) + return [ee.Date(value) for value in millis] + + +def expand_datasets(selected: list[str]) -> list[str]: + print("Inside expand_datasets") + if isinstance(selected, str): + selected = [selected] + if "both" in selected: + return ["chirps", "modis_pet"] + return selected + + +def make_chirps_image( + chirps: ee.ImageCollection, + aoi: ee.FeatureCollection, + start_date: ee.Date, + frequency: str, +) -> tuple[ee.Image, str]: + # CHIRPS is daily precipitation. Monthly mode sums daily precipitation + # into one monthly image before local download. + if frequency in ("daily", "native"): + end_date = start_date.advance(1, "day") + label = start_date.format("YYYYMMdd").getInfo() + image = chirps.filterDate(start_date, end_date).first() + else: + end_date = start_date.advance(1, "month") + label = start_date.format("YYYYMM").getInfo() + image = chirps.filterDate(start_date, end_date).sum() + + image = ( + ee.Image(image) + .select("precipitation") + .rename("precipitation") + .clip(aoi) + .toFloat() + ) + return image, label + + +def make_modis_pet_image( + modis_pet: ee.ImageCollection, + aoi: ee.FeatureCollection, + start_date: ee.Date, + frequency: str, + target_projection: ee.Projection | None = None, +) -> tuple[ee.Image, str]: + # MOD16A2GF PET is not daily. The PET band has a 0.1 scale factor, applied + # here so downloaded rasters contain real PET values. + if frequency == "daily": + raise ValueError("MODIS PET is not daily. Use --frequency monthly or native.") + + if frequency == "native": + end_date = start_date.advance(8, "day") + label = start_date.format("YYYYMMdd").getInfo() + image = modis_pet.filterDate(start_date, end_date).first() + else: + end_date = start_date.advance(1, "month") + label = start_date.format("YYYYMM").getInfo() + image = modis_pet.filterDate(start_date, end_date).sum() + image = image.setDefaultProjection(modis_pet.first().projection()) + + image = ee.Image(image).multiply(0.1).rename("PET") + + if target_projection is not None: + # Match the original GEE P-PET script: + # PET.reduceResolution(mean).reproject(crs=P.projection()) + image = image.reduceResolution( + reducer=ee.Reducer.mean(), maxPixels=65536 + ).reproject(crs=target_projection) + + image = image.clip(aoi).toFloat() + return image, label + + +def download_image( + image: ee.Image, + aoi: ee.FeatureCollection, + output_path: Path, + scale: int, + crs: str | None = "EPSG:4326", + retries: int = 3, +) -> None: + # Earth Engine's direct download endpoint has a request-size limit, so this + # function downloads only one time step at a time. + params = { + "scale": scale, + "region": aoi.geometry(), + "format": "GEO_TIFF", + } + if crs: + params["crs"] = crs + + for attempt in range(1, retries + 1): + try: + # Generate a signed EE URL, then stream the GeoTIFF bytes to disk. + url = image.getDownloadURL(params) + with requests.get(url, stream=True, timeout=300) as response: + response.raise_for_status() + output_path.parent.mkdir(parents=True, exist_ok=True) + with output_path.open("wb") as handle: + for chunk in response.iter_content(chunk_size=1024 * 1024): + if chunk: + handle.write(chunk) + return + except Exception: + # Remove partial files so retries/reruns do not treat corrupt files + # as valid completed downloads. + if output_path.exists(): + output_path.unlink() + if attempt == retries: + raise + time.sleep(2 * attempt) + + +def validate_inputs(aoi: str, datasets: list[str] | str, frequency: str) -> None: + # Keep validation near main() so notebook calls fail early and clearly. + print("Inside validate_inputs") + if aoi not in ("india", "state"): + raise ValueError("aoi must be 'india' or 'state'.") + if frequency not in ("daily", "monthly", "native"): + raise ValueError("frequency must be 'daily', 'monthly', or 'native'.") + + selected = [datasets] if isinstance(datasets, str) else datasets + invalid_datasets = set(selected) - set(DATASET_CHOICES) + if invalid_datasets: + raise ValueError(f"Invalid dataset(s): {sorted(invalid_datasets)}") + + +def dataset_label(dataset: str) -> str: + # Output labels are used in file names and final raster band descriptions. + if dataset == "chirps": + return "CHIRPS" + if dataset == "modis_pet": + return "MODIS_PET" + raise ValueError(f"Unknown dataset: {dataset}") + + +def build_dataset_image( + aoi: str, + state: str, + dataset: str, + start_date: str, + end_date: str, + frequency: str, +) -> tuple[ee.ImageCollection, ee.FeatureCollection, str, list[tuple[ee.Date, str]]]: + # Prepares the collection, AOI, and date labels for exactly one dataset. + # CHIRPS and MODIS PET stay separate from this point through final output. + start_ee_date = ee.Date(start_date) + end_date_exclusive = ee.Date(end_date).advance(1, "day") + + region = get_aoi(aoi, state) + chirps = ee.ImageCollection(CHIRPS_COLLECTION).select("precipitation") + modis_pet = ee.ImageCollection(MODIS_PET_COLLECTION).select("PET") + + aoi_label = "India" if aoi == "india" else state + + if dataset == "chirps": + unit = "day" if frequency in ("daily", "native") else "month" + dates = date_range(start_ee_date, end_date_exclusive, unit) + collection = chirps + name = dataset_label(dataset) + else: + if frequency == "daily": + raise ValueError( + "MODIS PET is not daily. Use --frequency monthly or native." + ) + if frequency == "native": + dates = collection_dates(modis_pet, start_ee_date, end_date_exclusive) + else: + dates = date_range(start_ee_date, end_date_exclusive, "month") + collection = modis_pet + name = dataset_label(dataset) + + labeled_dates = [] + for date in dates: + # These labels become the downloaded GeoTIFF file names. + date_format = "YYYYMMdd" if frequency in ("daily", "native") else "YYYYMM" + label = date.format(date_format).getInfo() + labeled_dates.append((date, label)) + + print( + f"Preparing {len(labeled_dates)} {frequency} {name} image(s) " + f"for {aoi_label}" + ) + return collection, region, name, labeled_dates + + +def download_dataset_images( + aoi: str, + state: str, + dataset: str, + start_date: str, + end_date: str, + frequency: str, + output_dir: str, + scale: int, + crs: str | None, + sleep: float, + overwrite: bool, + max_workers: int, +) -> None: + # End-to-end local workflow for one dataset: + # download one small GeoTIFF per time step and keep those files on disk. + aoi_label = "India" if aoi == "india" else state + safe_aoi = aoi_label.replace(" ", "_") + name = dataset_label(dataset) + dataset_output_dir = Path(output_dir) / safe_aoi / frequency / dataset + + collection, region, _, labeled_dates = build_dataset_image( + aoi=aoi, + state=state, + dataset=dataset, + start_date=start_date, + end_date=end_date, + frequency=frequency, + ) + + download_jobs = [] + target_projection = None + if dataset == "modis_pet" and frequency == "monthly": + # Use CHIRPS monthly precipitation projection as the output grid for PET. + # This keeps local PET files aligned with CHIRPS and reproduces the + # original GEE reduceResolution(mean)->reproject(P.projection()) step. + first_date = labeled_dates[0][0] + first_chirps, _ = make_chirps_image( + ee.ImageCollection(CHIRPS_COLLECTION).select("precipitation"), + region, + first_date, + frequency, + ) + target_projection = first_chirps.projection() + + for index, (current_date, label) in enumerate(labeled_dates, start=1): + # Keeping one time step per request avoids the EE direct-download size + # error seen when trying to download a large multiband image directly. + band_name = f"{name}_{label}" + single_path = dataset_output_dir / f"{band_name}.tif" + + if single_path.exists() and not overwrite: + print(f"[{index}/{len(labeled_dates)}] exists: {single_path.name}") + else: + download_jobs.append((index, current_date, band_name, single_path)) + + if max_workers <= 1: + for index, current_date, band_name, single_path in download_jobs: + if dataset == "chirps": + image, _ = make_chirps_image(collection, region, current_date, frequency) + else: + image, _ = make_modis_pet_image( + collection, region, current_date, frequency, target_projection + ) + print(f"[{index}/{len(labeled_dates)}] downloading: {single_path.name}") + download_image(image.rename(band_name), region, single_path, scale, crs) + time.sleep(sleep) + else: + # Parallel downloads are the main speed-up. Keep max_workers modest + # because Earth Engine may throttle too many simultaneous URL requests. + def download_one(job: tuple[int, ee.Date, str, Path]) -> tuple[int, str]: + index, current_date, band_name, single_path = job + if dataset == "chirps": + image, _ = make_chirps_image(collection, region, current_date, frequency) + else: + image, _ = make_modis_pet_image( + collection, region, current_date, frequency, target_projection + ) + download_image(image.rename(band_name), region, single_path, scale, crs) + if sleep: + time.sleep(sleep) + return index, single_path.name + + print(f"Downloading {len(download_jobs)} file(s) with {max_workers} workers") + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [executor.submit(download_one, job) for job in download_jobs] + for future in as_completed(futures): + index, file_name = future.result() + print(f"[{index}/{len(labeled_dates)}] downloaded: {file_name}") + + print(f"Downloaded {len(labeled_dates)} {name} image(s) to {dataset_output_dir}") + + +def main( + aoi: str = "india", + datasets: list[str] | str | None = None, + start_date: str = "2004-01-01", + end_date: str = "2023-12-31", + frequency: str = "monthly", + state: str = "Madhya Pradesh", + project: str = DEFAULT_PROJECT, + output_dir: str = "data/drought_inputs", + sleep: float = 0.2, + max_workers: int = 4, + overwrite: bool = False, +) -> None: + # Public entry point for notebooks/scripts. datasets=["both"] downloads + # separate CHIRPS and MODIS PET time-step GeoTIFFs. + selected_datasets = datasets or ["both"] + validate_inputs(aoi, selected_datasets, frequency) + initialize_earth_engine(project) + + expanded_datasets = expand_datasets(selected_datasets) + for dataset in expanded_datasets: + scale = 5500 + crs = "EPSG:4326" + download_dataset_images( + aoi=aoi, + state=state, + dataset=dataset, + start_date=start_date, + end_date=end_date, + frequency=frequency, + output_dir=output_dir, + scale=scale, + crs=crs, + sleep=sleep, + overwrite=overwrite, + max_workers=max_workers, + ) + + print("Done.") diff --git a/computing/spei/drought/drought_spei.R b/computing/spei/drought/drought_spei.R new file mode 100644 index 00000000..61d0c15c --- /dev/null +++ b/computing/spei/drought/drought_spei.R @@ -0,0 +1,341 @@ +# ============================================================================= +# SPEI Pipeline — Step 2 (Single State) +# Read multiband P-PET GeoTIFF, compute SPEI-1/3/12 pixel-wise, +# write 3 multiband output GeoTIFFs with named bands. +# ============================================================================= + +# ============================================================================= +# INSTALL (Run once if needed) +# ============================================================================= +# install.packages(c("SPEI", "raster"), repos="https://cloud.r-project.org") + +# ============================================================================= +# LIBRARIES +# ============================================================================= +library(SPEI) +library(raster) + +# ============================================================================= +# MAIN FUNCTION +# ============================================================================= +run_spei_pipeline <- function( + state_safe, + input_file, + output_dir +) { + + paste("Starting R script") + + # ------------------------------------------------------------------------- + # Create output directory + # ------------------------------------------------------------------------- + if (!dir.exists(output_dir)) { + dir.create(output_dir, recursive = TRUE) + } + + # ------------------------------------------------------------------------- + # Resume check + # ------------------------------------------------------------------------- + out_check <- file.path( + output_dir, + paste0("SPEI12_", state_safe, ".tif") + ) + + if (file.exists(out_check)) { + stop( + paste( + "Already processed:", + state_safe, + "— delete output files to rerun." + ) + ) + } + + # ========================================================================= + # SPEI FUNCTION + # + # Input: + # x = 240-length vector of monthly P-PET values + # + # Output: + # 340-length vector: + # [1:240] SPEI-1 + # [241:320] SPEI-3 seasonal months only + # [321:340] SPEI-12 annual only + # ========================================================================= + spei_function <- function(x, ...) { + + tryCatch({ + + if (all(is.na(x))) { + return(rep(NA, 340)) + } + + pixel_ts <- ts( + x, + start = c(2004, 1), + frequency = 12 + ) + + spei1_all <- as.vector( + spei( + pixel_ts, + 1, + distribution = "log-Logistic", + na.rm = TRUE + )$fitted + ) + + spei3_all <- as.vector( + spei( + pixel_ts, + 3, + distribution = "log-Logistic", + na.rm = TRUE + )$fitted + ) + + spei12_all <- as.vector( + spei( + pixel_ts, + 12, + distribution = "log-Logistic", + na.rm = TRUE + )$fitted + ) + + if (length(spei1_all) != 240) { + stop("Incorrect output length.") + } + + # ----------------------------------------------------------------- + # SPEI-3: keep only Mar, Jun, Sep, Dec + # ----------------------------------------------------------------- + seasonal_idx <- which( + ((seq_along(spei3_all) - 1) %% 12 + 1) %in% + c(3, 6, 9, 12) + ) + + spei3_sel <- spei3_all[seasonal_idx] + + # ----------------------------------------------------------------- + # SPEI-12: keep only December + # ----------------------------------------------------------------- + annual_idx <- seq(12, 240, by = 12) + + spei12_sel <- spei12_all[annual_idx] + + return( + c( + spei1_all, + spei3_sel, + spei12_sel + ) + ) + + }, error = function(e) { + + return(rep(NA, 340)) + + }) + } + + # ========================================================================= + # LOAD INPUT + # ========================================================================= + cat(paste("Loading:", input_file, "\n")) + + p_pet_brick <- brick(input_file) + + cat( + paste( + "Loaded", + nlayers(p_pet_brick), + "bands (expected 240)\n" + ) + ) + + # ========================================================================= + # COMPUTE BLOCK BY BLOCK + # ========================================================================= + cat("Running SPEI computation...\n") + + temp_file <- file.path( + output_dir, + paste0(state_safe, "_temp.tif") + ) + + result_brick <- brick( + p_pet_brick, + nl = 340 + ) + + result_brick <- writeStart( + result_brick, + filename = temp_file, + overwrite = TRUE + ) + + bs <- blockSize(p_pet_brick) + + for (i in 1:bs$n) { + + v <- getValues( + p_pet_brick, + row = bs$row[i], + nrows = bs$nrows[i] + ) + + res <- t( + apply(v, 1, spei_function) + ) + + writeValues( + result_brick, + res, + bs$row[i] + ) + + cat( + paste( + "Chunk", + i, + "/", + bs$n, + "\n" + ) + ) + } + + result_brick <- writeStop(result_brick) + + cat("Computation complete.\n") + + # ========================================================================= + # GENERATE BAND NAMES + # ========================================================================= + spei1_names <- paste0( + "y", + rep(2004:2023, each = 12), + "_m", + sprintf("%02d", rep(1:12, 20)) + ) + + spei3_names <- paste0( + "y", + rep(2004:2023, each = 4), + "_m", + sprintf("%02d", rep(c(3, 6, 9, 12), 20)) + ) + + spei12_names <- paste0( + "y", + 2004:2023 + ) + + # ========================================================================= + # SPLIT OUTPUTS + # ========================================================================= + cat("Saving output files...\n") + + all_b <- brick(temp_file) + + spei1_brick <- all_b[[1:240]] + spei3_brick <- all_b[[241:320]] + spei12_brick <- all_b[[321:340]] + + names(spei1_brick) <- spei1_names + names(spei3_brick) <- spei3_names + names(spei12_brick) <- spei12_names + + # ========================================================================= + # WRITE OUTPUTS + # ========================================================================= + writeRaster( + spei1_brick, + file.path( + output_dir, + paste0("SPEI1_", state_safe, ".tif") + ), + format = "GTiff", + overwrite = TRUE, + NAflag = -9999 + ) + + writeRaster( + spei3_brick, + file.path( + output_dir, + paste0("SPEI3_", state_safe, ".tif") + ), + format = "GTiff", + overwrite = TRUE, + NAflag = -9999 + ) + + writeRaster( + spei12_brick, + file.path( + output_dir, + paste0("SPEI12_", state_safe, ".tif") + ), + format = "GTiff", + overwrite = TRUE, + NAflag = -9999 + ) + + # ========================================================================= + # CLEANUP + # ========================================================================= + file.remove(temp_file) + + # ========================================================================= + # DONE + # ========================================================================= + cat( + paste0( + "\n✅ Done. Output files saved to: ", + output_dir, + "\n" + ) + ) + + cat( + paste0( + " SPEI1_", + state_safe, + ".tif — ", + nlayers(spei1_brick), + " bands\n" + ) + ) + + cat( + paste0( + " SPEI3_", + state_safe, + ".tif — ", + nlayers(spei3_brick), + " bands\n" + ) + ) + + cat( + paste0( + " SPEI12_", + state_safe, + ".tif — ", + nlayers(spei12_brick), + " bands\n" + ) + ) +} + +# ============================================================================= +# FUNCTION CALL +# ============================================================================= +run_spei_pipeline( + state_safe = "Madhya_Pradesh", + input_file = "data/drought_inputs/Madhya_Pradesh/monthly/ppet/P_PET_Madhya_Pradesh_monthly_multiband.tif", + output_dir = "data/drought_inputs/Madhya_Pradesh/monthly/ppet" +) \ No newline at end of file diff --git a/computing/spei/drought/export_ppet_single_state.py b/computing/spei/drought/export_ppet_single_state.py new file mode 100644 index 00000000..c3c7f52a --- /dev/null +++ b/computing/spei/drought/export_ppet_single_state.py @@ -0,0 +1,170 @@ +# ============================================================================= +# SPEI Pipeline - Step 1 (Local P-PET) +# +# Uses GeoTIFFs downloaded by download_chirps_local.py instead of reading: +# ee.ImageCollection("UCSB-CHG/CHIRPS/DAILY") +# ee.ImageCollection("MODIS/061/MOD16A2GF").select("PET") +# +# Output: one local multiband GeoTIFF with one P-PET band per month. +# Band names follow the original script: y{year}_m{month}, e.g. y2015_m06. +# ============================================================================= + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import rasterio +from rasterio.enums import Resampling +from rasterio.warp import reproject + + +# --- CONFIG --- +state_name = "Madhya_Pradesh" +start_year = 2004 +end_year = 2023 + +input_root = Path("data/drought_inputs") +output_dir = Path("data/drought_inputs") / state_name / "monthly" / "ppet" +output_path = output_dir / f"P_PET_{state_name}_monthly_multiband.tif" +OUTPUT_NODATA = -9999.0 + + +def find_monthly_file(dataset_dir: Path, prefix: str, year: int, month: int) -> Path: + """Find a downloaded monthly GeoTIFF, allowing nested folders from old runs.""" + label = f"{year}{month:02d}" + matches = sorted(dataset_dir.rglob(f"{prefix}_{label}.tif")) + if not matches: + raise FileNotFoundError( + f"Missing {prefix}_{label}.tif under {dataset_dir}. " + "Run download_chirps_local.py first." + ) + return matches[0] + + +def read_masked_band(dataset: rasterio.DatasetReader) -> np.ma.MaskedArray: + """Read band 1 and mask nodata plus any nan/inf values.""" + data = dataset.read(1, masked=True).astype("float32") + data = np.ma.masked_invalid(data) + return np.ma.masked_where(np.abs(data) > 1.0e20, data) + + +def reproject_modis_to_chirps_grid( + modis_path: Path, + chirps_dataset: rasterio.DatasetReader, + log_metadata: bool = False, +) -> np.ma.MaskedArray: + """Reproject/resample MODIS PET onto the CHIRPS projection and pixel grid.""" + with rasterio.open(modis_path) as modis_src: + if log_metadata: + print( + " MODIS -> CHIRPS grid:", + f"modis_crs={modis_src.crs}", + f"chirps_crs={chirps_dataset.crs}", + f"modis_shape=({modis_src.height}, {modis_src.width})", + f"chirps_shape=({chirps_dataset.height}, {chirps_dataset.width})", + ) + + if ( + modis_src.crs == chirps_dataset.crs + and modis_src.transform == chirps_dataset.transform + and modis_src.width == chirps_dataset.width + and modis_src.height == chirps_dataset.height + ): + return read_masked_band(modis_src) + + pet_source = read_masked_band(modis_src) + pet_on_chirps_grid = np.full( + (chirps_dataset.height, chirps_dataset.width), + OUTPUT_NODATA, + dtype="float32", + ) + + # MODIS PET is finer (~500m) than CHIRPS (~5500m), so average is the + # right downsampling behavior for a monthly total/continuous variable. + reproject( + source=pet_source.filled(OUTPUT_NODATA), + destination=pet_on_chirps_grid, + src_transform=modis_src.transform, + src_crs=modis_src.crs, + dst_transform=chirps_dataset.transform, + dst_crs=chirps_dataset.crs, + dst_width=chirps_dataset.width, + dst_height=chirps_dataset.height, + src_nodata=OUTPUT_NODATA, + dst_nodata=OUTPUT_NODATA, + init_dest_nodata=True, + resampling=Resampling.average, + ) + return np.ma.masked_where( + (pet_on_chirps_grid == OUTPUT_NODATA) | ~np.isfinite(pet_on_chirps_grid), + pet_on_chirps_grid, + ) + + +def main( + state: str = state_name, + start: int = start_year, + end: int = end_year, + data_root: Path | str = input_root, + output: Path | str | None = output_path, +) -> Path: + data_root = Path(data_root) + chirps_dir = data_root / state / "monthly" / "chirps" + modis_dir = data_root / state / "monthly" / "modis_pet" + + output_file = Path(output) if output else data_root / state / "ppet" / ( + f"P_PET_{state}_monthly_multiband.tif" + ) + output_file.parent.mkdir(parents=True, exist_ok=True) + + months = [(year, month) for year in range(start, end + 1) for month in range(1, 13)] + first_chirps = find_monthly_file(chirps_dir, "CHIRPS", *months[0]) + + with rasterio.open(first_chirps) as template: + profile = template.profile.copy() + profile.update( + count=len(months), + dtype="float32", + nodata=OUTPUT_NODATA, + compress="lzw", + BIGTIFF="IF_SAFER", + ) + + with rasterio.open(output_file, "w", **profile) as dst: + for band_index, (year, month) in enumerate(months, start=1): + chirps_path = find_monthly_file(chirps_dir, "CHIRPS", year, month) + modis_path = find_monthly_file(modis_dir, "MODIS_PET", year, month) + + with rasterio.open(chirps_path) as chirps_src: + precipitation = read_masked_band(chirps_src) + + # MODIS PET downloaded by download_chirps_local.py already has + # the 0.1 scale factor applied, so do not multiply again here. + pet = reproject_modis_to_chirps_grid( + modis_path, chirps_src, log_metadata=(band_index == 1) + ) + + precipitation_data = precipitation.filled(np.nan).astype("float32") + pet_data = pet.filled(np.nan).astype("float32") + + valid_mask = np.isfinite(precipitation_data) & np.isfinite(pet_data) + ppet_data = np.full( + precipitation_data.shape, OUTPUT_NODATA, dtype="float32" + ) + with np.errstate(invalid="ignore", over="ignore"): + ppet_data[valid_mask] = ( + precipitation_data[valid_mask] - pet_data[valid_mask] + ) + + ppet = np.ma.masked_where( + (ppet_data == OUTPUT_NODATA) | ~np.isfinite(ppet_data), ppet_data + ) + band_name = f"y{year}_m{month:02d}" + dst.write(ppet.filled(OUTPUT_NODATA).astype("float32"), band_index) + dst.set_band_description(band_index, band_name) + invalid_count = int(np.ma.count_masked(ppet)) + print(f"Prepared: {band_name} nodata_pixels={invalid_count}") + + print(f"\nWrote {len(months)} P-PET band(s): {output_file}") + return output_file diff --git a/computing/spei/drought/spei_runner.py b/computing/spei/drought/spei_runner.py new file mode 100644 index 00000000..0f528c4c --- /dev/null +++ b/computing/spei/drought/spei_runner.py @@ -0,0 +1,37 @@ +from pathlib import Path +import subprocess + + +BASE_DIR = Path(__file__).resolve().parent.parent + +R_SCRIPT = BASE_DIR / "drought" / "drought_spei.R" + + +def run_spei_pipeline(state_safe=None, input_file=None, output_dir=None): + + command = [ + "Rscript", + str(R_SCRIPT), + state_safe, + input_file, + output_dir, + ] + + print("COMMAND:", command) + + result = subprocess.run( + command, + capture_output=True, + text=True, + ) + + print("RETURN CODE:", result.returncode) + print("STDOUT:\n", result.stdout) + print("STDERR:\n", result.stderr) + + if result.returncode != 0: + raise Exception( + f" R Script Failed COMMAND:{command} RETURN CODE:{result.returncode} STDOUT:{result.stdout} STDERR:{result.stderr}" + ) + + return result.stdout diff --git a/computing/tasks.py b/computing/tasks.py new file mode 100644 index 00000000..eb0b1000 --- /dev/null +++ b/computing/tasks.py @@ -0,0 +1,13 @@ +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 + +__all__ = [ + "et_download", + "generate_hydrology", + "generate_hydrology_base_layer", + "generate_runoff_gpu", +] diff --git a/computing/terrain_descriptor/store_watersheds_for_tehsils.py b/computing/terrain_descriptor/store_watersheds_for_tehsils.py new file mode 100644 index 00000000..530a8ef8 --- /dev/null +++ b/computing/terrain_descriptor/store_watersheds_for_tehsils.py @@ -0,0 +1,394 @@ +import argparse +import csv +import os +import re +import time + +import geopandas as gpd + + +DEFAULT_MICROWATERSHED_PATH = ( + "data/base_layers/Microwatershed_v2_with_details.geojson" +) +DEFAULT_TEHSIL_PATH = "data/admin-boundary/input/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"] + +OUTPUT_FORMATS = { + "geojson": ("GeoJSON", ".geojson"), + "gpkg": ("GPKG", ".gpkg"), +} + + +def valid_gee_text(description): + description = re.sub(r"[^a-zA-Z0-9 ,:;_-]", "", description) + return description.replace(" ", "_") + + +def _normalize_text(value): + return valid_gee_text(str(value).strip().lower()) + +def _find_matching_column(columns, candidates): + columns_list = list(columns) + lower_lookup = {col.lower(): col for col in columns_list} + for candidate in candidates: + if candidate.lower() in lower_lookup: + return lower_lookup[candidate.lower()] + + compact_lookup = { + re.sub(r"[^a-z0-9]", "", col.lower()): col for col in columns_list + } + for candidate in candidates: + compact = re.sub(r"[^a-z0-9]", "", candidate.lower()) + if compact in compact_lookup: + return compact_lookup[compact] + return None + + +def _validate_geometry(gdf, fix_invalid=False): + gdf = gdf[gdf.geometry.notna()].copy() + if gdf.empty: + return gdf + + gdf = gdf[~gdf.geometry.is_empty].copy() + if gdf.empty: + return gdf + + if fix_invalid: + invalid = ~gdf.is_valid + if invalid.any(): + gdf.loc[invalid, "geometry"] = gdf.loc[invalid, "geometry"].buffer(0) + gdf = gdf[gdf.geometry.notna() & ~gdf.geometry.is_empty].copy() + return gdf + + +def _prepare_tehsil_boundaries(tehsil_path, state=None, district=None, tehsil=None): + tehsil_gdf = gpd.read_file(tehsil_path) + + state_col = "STATE" if "STATE" in tehsil_gdf.columns else None + district_col = "District" if "District" in tehsil_gdf.columns else None + tehsil_col = "TEHSIL" if "TEHSIL" in tehsil_gdf.columns else None + + if state_col is None: + state_col = _find_matching_column(tehsil_gdf.columns, STATE_COLUMN_CANDIDATES) + if district_col is None: + district_col = _find_matching_column(tehsil_gdf.columns, DISTRICT_COLUMN_CANDIDATES) + if tehsil_col is None: + tehsil_col = _find_matching_column(tehsil_gdf.columns, TEHSIL_COLUMN_CANDIDATES) + + if not all([state_col, district_col, tehsil_col]): + raise ValueError( + "Could not identify STATE/District/TEHSIL columns in tehsil file. " + f"Available columns: {list(tehsil_gdf.columns)}" + ) + + if state: + tehsil_gdf = tehsil_gdf[ + tehsil_gdf[state_col].fillna("").map(_normalize_text) + == _normalize_text(state) + ] + if district: + tehsil_gdf = tehsil_gdf[ + tehsil_gdf[district_col].fillna("").map(_normalize_text) + == _normalize_text(district) + ] + if tehsil: + tehsil_gdf = tehsil_gdf[ + tehsil_gdf[tehsil_col].fillna("").map(_normalize_text) + == _normalize_text(tehsil) + ] + + tehsil_gdf = _validate_geometry(tehsil_gdf, fix_invalid=True) + if tehsil_gdf.empty: + raise ValueError("No matching tehsil records found after applying filters.") + + if tehsil_gdf.crs is None: + tehsil_gdf = tehsil_gdf.set_crs("EPSG:4326") + + tehsil_gdf["__state_norm"] = tehsil_gdf[state_col].map(_normalize_text) + tehsil_gdf["__district_norm"] = tehsil_gdf[district_col].map(_normalize_text) + tehsil_gdf["__tehsil_norm"] = tehsil_gdf[tehsil_col].map(_normalize_text) + tehsil_gdf["__key"] = ( + tehsil_gdf["__state_norm"] + + "||" + + tehsil_gdf["__district_norm"] + + "||" + + tehsil_gdf["__tehsil_norm"] + ) + + meta = ( + tehsil_gdf.groupby("__key", as_index=False) + .agg( + { + state_col: "first", + district_col: "first", + tehsil_col: "first", + } + ) + .rename( + columns={ + state_col: "STATE", + district_col: "District", + tehsil_col: "TEHSIL", + } + ) + ) + + dissolved = tehsil_gdf[["__key", "geometry"]].dissolve(by="__key").reset_index() + dissolved = dissolved.merge(meta, on="__key", how="left") + dissolved = _validate_geometry(dissolved, fix_invalid=True) + + return dissolved + + +def _load_microwatersheds(microwatershed_path, fix_invalid=False): + microwatershed_gdf = gpd.read_file(microwatershed_path) + if microwatershed_gdf.empty: + raise ValueError("Microwatershed file is empty.") + + uid_col = _find_matching_column( + microwatershed_gdf.columns, MWS_UID_COLUMN_CANDIDATES + ) + if not uid_col: + raise ValueError( + "Could not identify microwatershed UID column. " + f"Available columns: {list(microwatershed_gdf.columns)}" + ) + + keep_columns = ["geometry", uid_col] + for col in MWS_OPTIONAL_COLUMNS: + if col in microwatershed_gdf.columns: + keep_columns.append(col) + + microwatershed_gdf = microwatershed_gdf[keep_columns].copy() + microwatershed_gdf = microwatershed_gdf.rename(columns={uid_col: "uid"}) + microwatershed_gdf = _validate_geometry(microwatershed_gdf, fix_invalid=fix_invalid) + + if microwatershed_gdf.crs is None: + microwatershed_gdf = microwatershed_gdf.set_crs("EPSG:4326") + + if microwatershed_gdf.empty: + raise ValueError("No valid microwatershed geometry found.") + + return microwatershed_gdf + + +def _save_subset(subset_gdf, output_path, driver): + if driver == "GPKG": + subset_gdf.to_file(output_path, driver=driver, layer="watersheds") + else: + subset_gdf.to_file(output_path, driver=driver) + + +def generate_tehsil_watershed_copies( + microwatershed_path=DEFAULT_MICROWATERSHED_PATH, + tehsil_path=DEFAULT_TEHSIL_PATH, + output_dir=DEFAULT_OUTPUT_DIR, + output_format="gpkg", + overwrite=False, + include_empty=False, + fix_invalid_mws=False, + state=None, + district=None, + tehsil=None, +): + if output_format not in OUTPUT_FORMATS: + raise ValueError( + f"Unsupported format: {output_format}. Supported: {sorted(OUTPUT_FORMATS)}" + ) + + driver, extension = OUTPUT_FORMATS[output_format] + + if not os.path.exists(microwatershed_path): + raise FileNotFoundError(f"Microwatershed file not found: {microwatershed_path}") + if not os.path.exists(tehsil_path): + raise FileNotFoundError(f"Tehsil file not found: {tehsil_path}") + + os.makedirs(output_dir, exist_ok=True) + start_time = time.time() + + print("Loading tehsil boundaries...") + tehsil_boundaries = _prepare_tehsil_boundaries( + tehsil_path=tehsil_path, + state=state, + district=district, + tehsil=tehsil, + ) + print(f"Prepared {len(tehsil_boundaries)} unique tehsil boundaries.") + + print("Loading microwatersheds (this can take time once)...") + microwatersheds = _load_microwatersheds( + microwatershed_path=microwatershed_path, + fix_invalid=fix_invalid_mws, + ) + print(f"Loaded {len(microwatersheds)} microwatershed features.") + + if tehsil_boundaries.crs != microwatersheds.crs: + tehsil_boundaries = tehsil_boundaries.to_crs(microwatersheds.crs) + + print("Building microwatershed spatial index...") + mws_sindex = microwatersheds.sindex + + manifest_rows = [] + total = len(tehsil_boundaries) + written_count = 0 + skipped_count = 0 + + for idx, row in enumerate(tehsil_boundaries.itertuples(index=False), start=1): + state_name = row.STATE + district_name = row.District + tehsil_name = row.TEHSIL + tehsil_geom = row.geometry + + try: + candidate_ids = mws_sindex.query(tehsil_geom, predicate="intersects") + except TypeError: + candidate_ids = list(mws_sindex.intersection(tehsil_geom.bounds)) + + if len(candidate_ids) == 0: + intersection_gdf = microwatersheds.iloc[0:0].copy() + else: + candidates = microwatersheds.iloc[list(candidate_ids)] + intersection_gdf = candidates[candidates.intersects(tehsil_geom)].copy() + + feature_count = len(intersection_gdf) + state_dir = valid_gee_text(str(state_name).strip().lower()) or "unknown_state" + district_dir = ( + valid_gee_text(str(district_name).strip().lower()) or "unknown_district" + ) + tehsil_file = ( + valid_gee_text(str(tehsil_name).strip().lower()) + or f"unknown_tehsil_{idx:05d}" + ) + + output_subdir = os.path.join(output_dir, state_dir, district_dir) + os.makedirs(output_subdir, exist_ok=True) + out_path = os.path.join(output_subdir, f"{tehsil_file}{extension}") + + row_status = "empty" + if feature_count > 0 or include_empty: + if feature_count > 0: + intersection_gdf["STATE"] = state_name + intersection_gdf["District"] = district_name + intersection_gdf["TEHSIL"] = tehsil_name + + if overwrite or not os.path.exists(out_path): + _save_subset(intersection_gdf, out_path, driver) + row_status = "written" + written_count += 1 + else: + row_status = "exists" + skipped_count += 1 + else: + skipped_count += 1 + + manifest_rows.append( + { + "state": state_name, + "district": district_name, + "tehsil": tehsil_name, + "feature_count": feature_count, + "status": row_status, + "output_path": out_path if (feature_count > 0 or include_empty) else "", + } + ) + + if idx % 100 == 0 or idx == total: + print(f"Processed {idx}/{total} tehsils...") + + manifest_path = os.path.join(output_dir, "tehsil_watershed_manifest.csv") + with open(manifest_path, "w", newline="", encoding="utf-8") as csvfile: + writer = csv.DictWriter( + csvfile, + fieldnames=[ + "state", + "district", + "tehsil", + "feature_count", + "status", + "output_path", + ], + ) + writer.writeheader() + writer.writerows(manifest_rows) + + elapsed = time.time() - start_time + print("Done.") + print(f"Output directory: {output_dir}") + print(f"Manifest: {manifest_path}") + print(f"Tehsils processed: {total}") + print(f"Files written: {written_count}") + print(f"Skipped: {skipped_count}") + print(f"Elapsed time: {elapsed:.2f} seconds") + + return manifest_path + + +def _build_parser(): + parser = argparse.ArgumentParser( + description=( + "Precompute and store microwatersheds for each tehsil based on spatial " + "intersection with tehsil boundaries." + ) + ) + parser.add_argument( + "--microwatershed-path", + default=DEFAULT_MICROWATERSHED_PATH, + help="Path to Microwatershed GeoJSON", + ) + parser.add_argument( + "--tehsil-path", + default=DEFAULT_TEHSIL_PATH, + help="Path to tehsil boundary GeoJSON", + ) + parser.add_argument( + "--output-dir", + default=DEFAULT_OUTPUT_DIR, + help="Directory to save per-tehsil watershed files", + ) + parser.add_argument( + "--format", + choices=sorted(OUTPUT_FORMATS.keys()), + default="gpkg", + help="Output file format", + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Overwrite existing output files", + ) + parser.add_argument( + "--include-empty", + action="store_true", + help="Also write empty files for tehsils with no intersecting microwatersheds", + ) + parser.add_argument( + "--fix-invalid-mws", + action="store_true", + help="Fix invalid microwatershed geometries (slower, use only if needed)", + ) + parser.add_argument("--state", help="Optional state filter") + parser.add_argument("--district", help="Optional district filter") + parser.add_argument("--tehsil", help="Optional tehsil filter") + return parser + + +if __name__ == "__main__": + args = _build_parser().parse_args() + generate_tehsil_watershed_copies( + microwatershed_path=args.microwatershed_path, + tehsil_path=args.tehsil_path, + output_dir=args.output_dir, + output_format=args.format, + overwrite=args.overwrite, + include_empty=args.include_empty, + fix_invalid_mws=args.fix_invalid_mws, + state=args.state, + district=args.district, + tehsil=args.tehsil, + ) diff --git a/computing/terrain_descriptor/terrain_clusters_local.py b/computing/terrain_descriptor/terrain_clusters_local.py new file mode 100644 index 00000000..bc788982 --- /dev/null +++ b/computing/terrain_descriptor/terrain_clusters_local.py @@ -0,0 +1,181 @@ +import os +import shutil +import tempfile +import zipfile + +from nrm_app.celery import app +from utilities.gee_utils import valid_gee_text + +from computing.local_compute_helper import ( + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + build_output_vector_path, + compute_terrain_properties_for_watersheds, + load_precomputed_watersheds, + resolve_clipped_terrain_raster_path, + write_vector_output, +) + + +CLIPPED_TERRAIN_RASTER_DIR = "data/terrain/fabdem_local" +LOCAL_CLUSTER_OUTPUT_DIR = "data/terrain/terrain_clusters_local" +GEOSERVER_WORKSPACE = "terrain" + + +def _push_vector_to_geoserver(gdf, layer_name): + from utilities.geoserver_utils import Geoserver + + temp_dir = tempfile.mkdtemp(prefix=f"{layer_name}_") + gpkg_path = os.path.join(temp_dir, f"{layer_name}.gpkg") + zip_path = os.path.join(temp_dir, f"{layer_name}.zip") + + try: + gdf.to_file(gpkg_path, driver="GPKG", layer=layer_name) + with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf: + zipf.write(gpkg_path, arcname=os.path.basename(gpkg_path)) + + geo = Geoserver() + geo.delete_vector_store(workspace=GEOSERVER_WORKSPACE, store=layer_name) + response = geo.create_shp_datastore( + path=zip_path, + store_name=layer_name, + workspace=GEOSERVER_WORKSPACE, + file_extension="gpkg", + ) + return response + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + +def run_terrain_clusters_local( + state, + district, + block, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + clipped_raster_dir=CLIPPED_TERRAIN_RASTER_DIR, + push_to_geoserver=True, + sync_layer_metadata=False, +): + state = str(state).strip() + district = str(district).strip() + block = str(block).strip() + + layer_name = ( + f"{valid_gee_text(str(district).strip().lower())}_" + f"{valid_gee_text(str(block).strip().lower())}_cluster" + ) + + watersheds_gdf, watershed_source = load_precomputed_watersheds( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + ) + clipped_raster_path = resolve_clipped_terrain_raster_path( + state=state, + district=district, + block=block, + clipped_raster_dir=clipped_raster_dir, + ) + print(f"Using clipped terrain raster: {clipped_raster_path}") + + terrain_clusters_gdf = compute_terrain_properties_for_watersheds( + watersheds_gdf=watersheds_gdf, + raster_path=clipped_raster_path, + ) + output_path = build_output_vector_path( + layer_name=layer_name, + state=state, + district=district, + block=block, + output_base_dir=LOCAL_CLUSTER_OUTPUT_DIR, + block_fallback="unknown_tehsil", + ) + local_vector_path = write_vector_output( + gdf=terrain_clusters_gdf, + output_path=output_path, + layer_name=layer_name, + ) + print(f"Saved local terrain cluster vector: {local_vector_path}") + + if push_to_geoserver: + try: + geoserver_response = _push_vector_to_geoserver( + gdf=terrain_clusters_gdf, + layer_name=layer_name, + ) + print(f"GeoServer response: {geoserver_response}") + except Exception as error: + print(f"Failed to sync terrain clusters vector to GeoServer: {error}") + return False + + if sync_layer_metadata: + from computing.STAC_specs import generate_STAC_layerwise + from computing.utils import save_layer_info_to_db, update_layer_sync_status + + layer_id = save_layer_info_to_db( + state=state, + district=district, + block=block, + layer_name=layer_name, + asset_id=local_vector_path, + dataset_name="Terrain Vector", + algorithm="FABDEM", + algorithm_version="2.0", + ) + if layer_id: + update_layer_sync_status(layer_id=layer_id, sync_to_geoserver=True) + layer_stac_generated = generate_STAC_layerwise.generate_vector_stac( + state=state, + district=district, + block=block, + layer_name="terrain_vector", + ) + update_layer_sync_status( + layer_id=layer_id, + is_stac_specs_generated=layer_stac_generated, + ) + + print(f"Completed terrain cluster computation for {state}/{district}/{block}") + print(f"Watershed boundary source: {watershed_source}") + return True + + +def _generate_terrain_clusters_task( + state, + district, + block, + gee_account_id=None, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + clipped_raster_dir=CLIPPED_TERRAIN_RASTER_DIR, +): + _ = gee_account_id + return run_terrain_clusters_local( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + clipped_raster_dir=clipped_raster_dir, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + +@app.task(bind=True) +def generate_terrain_clusters( + self, + state, + district, + block, + gee_account_id=None, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, + clipped_raster_dir=CLIPPED_TERRAIN_RASTER_DIR, +): + _ = self + return _generate_terrain_clusters_task( + state=state, + district=district, + block=block, + gee_account_id=gee_account_id, + precomputed_roi_dir=precomputed_roi_dir, + clipped_raster_dir=clipped_raster_dir, + ) diff --git a/computing/terrain_descriptor/terrain_compute_all_local.py b/computing/terrain_descriptor/terrain_compute_all_local.py new file mode 100644 index 00000000..27e98847 --- /dev/null +++ b/computing/terrain_descriptor/terrain_compute_all_local.py @@ -0,0 +1,255 @@ +from pathlib import Path + +from nrm_app.celery import app + +from computing.local_compute_helper import PRECOMPUTED_TEHSIL_WATERSHED_DIR +from computing.lulc_X_terrain.lulc_on_plain_cluster_local import ( + run_lulc_on_plain_cluster_local, +) +from computing.lulc_X_terrain.lulc_on_slope_cluster_local import ( + run_lulc_on_slope_cluster_local, +) +from computing.terrain_descriptor.terrain_clusters_local import ( + run_terrain_clusters_local, +) +from computing.terrain_descriptor.terrain_raster_fabdem_local import ( + run_terrain_raster_fabdem_local, +) +from utilities.gee_utils import valid_gee_text + + +def _run_step(step_name, step_func, **kwargs): + print(f"Starting step: {step_name}") + step_result = step_func(**kwargs) + print(f"Completed step: {step_name} -> {step_result}") + if step_result is False: + raise RuntimeError(f"{step_name} failed") + return step_result + + +def _is_missing_block(block): + if block is None: + return True + return str(block).strip().lower() in {"", "null", "none"} + + +def _resolve_blocks_for_district( + state, + district, + precomputed_roi_dir=PRECOMPUTED_TEHSIL_WATERSHED_DIR, +): + district_dir = ( + Path(precomputed_roi_dir) + / (valid_gee_text(str(state).strip().lower()) or "unknown_state") + / (valid_gee_text(str(district).strip().lower()) or "unknown_district") + ) + if not district_dir.exists(): + raise FileNotFoundError( + f"District watershed directory not found: {district_dir}" + ) + + asset_block_slugs = sorted( + { + path.stem + for path in district_dir.iterdir() + if path.is_file() and path.suffix.lower() in {".gpkg", ".geojson"} + } + ) + if not asset_block_slugs: + raise FileNotFoundError( + f"No watershed files found for district: {state}/{district}" + ) + + canonical_block_names = {} + try: + from geoadmin.models import DistrictSOI, StateSOI, TehsilSOI + + state_obj = StateSOI.objects.get(state_name__iexact=state) + district_obj = DistrictSOI.objects.get( + district_name__iexact=district, + state=state_obj, + ) + tehsil_names = TehsilSOI.objects.filter(district=district_obj).values_list( + "tehsil_name", + flat=True, + ) + canonical_block_names = { + ( + valid_gee_text(str(tehsil_name).strip().lower()) + or "unknown_block" + ): tehsil_name + for tehsil_name in tehsil_names + } + except Exception as error: + print( + "Unable to resolve canonical block names from DB. " + f"Falling back to asset slugs. Error: {error}" + ) + + return [ + canonical_block_names.get(block_slug, block_slug) + for block_slug in asset_block_slugs + ] + + +def _run_terrain_compute_for_block( + state, + district, + block, + start_year, + end_year, +): + state = str(state).strip() + district = str(district).strip() + block = str(block).strip() + start_year = int(start_year) + end_year = int(end_year) + + results = {} + + results["terrain_raster"] = _run_step( + "terrain_raster", + run_terrain_raster_fabdem_local, + state=state, + district=district, + block=block, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + results["terrain_vector"] = _run_step( + "terrain_vector", + run_terrain_clusters_local, + state=state, + district=district, + block=block, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + results["terrain_lulc_slope"] = _run_step( + "terrain_lulc_slope", + run_lulc_on_slope_cluster_local, + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + results["terrain_lulc_plain"] = _run_step( + "terrain_lulc_plain", + run_lulc_on_plain_cluster_local, + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + print( + f"Completed local terrain compute-all flow for {state}/{district}/{block}: {results}" + ) + return results + + +def run_terrain_compute_all_local( + state, + district, + block, + start_year, + end_year, +): + state = str(state).strip() + district = str(district).strip() + start_year = int(start_year) + end_year = int(end_year) + + if _is_missing_block(block): + block_names = _resolve_blocks_for_district( + state=state, + district=district, + ) + district_results = {} + success_count = 0 + + for block_name in block_names: + try: + district_results[block_name] = _run_terrain_compute_for_block( + state=state, + district=district, + block=block_name, + start_year=start_year, + end_year=end_year, + ) + success_count += 1 + except Exception as error: + print( + f"Failed district-wide terrain compute-all for block {block_name}: {error}" + ) + district_results[block_name] = {"error": str(error)} + + summary = { + "scope": "district", + "state": state, + "district": district, + "total_blocks": len(block_names), + "successful_blocks": success_count, + "failed_blocks": len(block_names) - success_count, + "blocks": district_results, + } + print(f"Completed district-wide terrain compute-all flow: {summary}") + + if success_count == 0: + raise RuntimeError( + f"Terrain compute-all failed for every block in {state}/{district}" + ) + return summary + + return _run_terrain_compute_for_block( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + ) + + +def _generate_terrain_compute_all_local_task( + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = gee_account_id + return run_terrain_compute_all_local( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + ) + + +@app.task(bind=True) +def generate_terrain_compute_all( + self, + state, + district, + block, + start_year, + end_year, + gee_account_id=None, +): + _ = self + return _generate_terrain_compute_all_local_task( + state=state, + district=district, + block=block, + start_year=start_year, + end_year=end_year, + gee_account_id=gee_account_id, + ) diff --git a/computing/terrain_descriptor/terrain_raster_fabdem_local.py b/computing/terrain_descriptor/terrain_raster_fabdem_local.py new file mode 100644 index 00000000..2d4d253e --- /dev/null +++ b/computing/terrain_descriptor/terrain_raster_fabdem_local.py @@ -0,0 +1,166 @@ +from utilities.gee_utils import valid_gee_text + +from nrm_app.celery import app + +from computing.local_compute_helper import ( + PRECOMPUTED_TEHSIL_WATERSHED_DIR, + TERRAIN_RASTER_PATH, + build_output_raster_path, + clip_raster_with_roi, + load_precomputed_roi, + push_local_raster_to_geoserver, + read_validated_vector_file, +) + + +LOCAL_OUTPUT_BASE_DIR = "data/terrain/fabdem_local" +GEOSERVER_STYLE = "terrain_raster" +GEOSERVER_WORKSPACE = "terrain" + + +def run_terrain_raster_fabdem_local( + state=None, + district=None, + block=None, + asset_suffix=None, + roi=None, + precomputed_roi_dir=None, + push_to_geoserver=True, + sync_layer_metadata=False, +): + if state and district and block: + layer_name = ( + f"{valid_gee_text(str(district).strip().lower())}_" + f"{valid_gee_text(str(block).strip().lower())}_terrain_raster" + ) + roi_gdf = load_precomputed_roi( + state=state, + district=district, + block=block, + precomputed_roi_dir=precomputed_roi_dir, + ) + else: + if not roi or not asset_suffix: + raise ValueError( + "For non state/district/block runs, both `roi` and `asset_suffix` are required." + ) + layer_name = f"{asset_suffix}_terrain_raster".lower() + roi_gdf = read_validated_vector_file( + roi, + f"ROI file has no valid geometries: {roi}", + ) + + output_raster_path = build_output_raster_path( + layer_name=layer_name, + output_base_dir=LOCAL_OUTPUT_BASE_DIR, + state=state, + district=district, + block=block, + ) + clipped_raster_path = clip_raster_with_roi( + roi_gdf=roi_gdf, + raster_path=TERRAIN_RASTER_PATH, + output_path=output_raster_path, + ) + + print(f"Local clipped FABDEM raster written to: {clipped_raster_path}") + + if push_to_geoserver: + try: + upload_res, style_res = push_local_raster_to_geoserver( + file_path=clipped_raster_path, + layer_name=layer_name, + workspace=GEOSERVER_WORKSPACE, + style_name=GEOSERVER_STYLE, + ) + print(f"GeoServer upload response: {upload_res}") + print(f"GeoServer style response: {style_res}") + except Exception as error: + print(f"Failed to sync local FABDEM raster to GeoServer: {error}") + return False + + if sync_layer_metadata and state and district and block: + from computing.STAC_specs import generate_STAC_layerwise + from computing.utils import save_layer_info_to_db, update_layer_sync_status + + layer_id = save_layer_info_to_db( + state=state, + district=district, + block=block, + layer_name=layer_name, + asset_id=clipped_raster_path, + dataset_name="Terrain Raster", + algorithm="FABDEM", + algorithm_version="2.0", + ) + + if layer_id: + update_layer_sync_status(layer_id=layer_id, sync_to_geoserver=True) + print("sync to geoserver flag is updated") + + layer_stac_generated = generate_STAC_layerwise.generate_raster_stac( + state=state, + district=district, + block=block, + layer_name="terrain_raster", + ) + update_layer_sync_status( + layer_id=layer_id, + is_stac_specs_generated=layer_stac_generated, + ) + + return True + + +def _generate_terrain_raster_clip_task( + state=None, + district=None, + block=None, + gee_account_id=None, + asset_suffix=None, + asset_folder=None, + proj_id=None, + roi=None, + precomputed_roi_dir=None, + app_type="MWS", +): + _ = gee_account_id, asset_folder, proj_id, app_type + return run_terrain_raster_fabdem_local( + state=state, + district=district, + block=block, + asset_suffix=asset_suffix, + roi=roi, + precomputed_roi_dir=precomputed_roi_dir, + push_to_geoserver=True, + sync_layer_metadata=True, + ) + + +@app.task(bind=True) +def generate_terrain_raster_clip( + self, + state=None, + district=None, + block=None, + gee_account_id=None, + asset_suffix=None, + asset_folder=None, + proj_id=None, + roi=None, + precomputed_roi_dir=None, + app_type="MWS", +): + _ = self + return _generate_terrain_raster_clip_task( + state=state, + district=district, + block=block, + gee_account_id=gee_account_id, + asset_suffix=asset_suffix, + asset_folder=asset_folder, + proj_id=proj_id, + roi=roi, + precomputed_roi_dir=precomputed_roi_dir, + app_type=app_type, + ) diff --git a/computing/urls.py b/computing/urls.py index 19e8bd3f..706f554e 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"), @@ -40,6 +52,11 @@ api.generate_terrain_descriptor, name="generate_terrain_descriptor", ), + path( + "generate_terrain_compute_all/", + api.generate_terrain_compute_all, + name="generate_terrain_compute_all", + ), path( "generate_terrain_raster/", api.generate_terrain_raster, diff --git a/computing/utils.py b/computing/utils.py index 2590f3b9..c4038942 100644 --- a/computing/utils.py +++ b/computing/utils.py @@ -1,8 +1,12 @@ +import logging import os import geopandas as gpd import fiona import copy +import requests + +logger = logging.getLogger(__name__) from computing.models import Layer, Dataset from geoadmin.models import ( @@ -37,6 +41,8 @@ from shapely.geometry import shape from shapely.validation import explain_validity import zipfile +from django.conf import settings + from datetime import datetime, timedelta from django.conf import settings import logging @@ -607,138 +613,72 @@ def merge_with_ndvi(pair): sync_project_fc_to_geoserver(merged_final, proj_obj.name, layer_name, "waterrej") -def save_layer_info_to_db( - state, - district, - block, - layer_name, - asset_id, - dataset_name, - sync_to_geoserver=False, - layer_version="1.0", - algorithm=None, - algorithm_version="1.0", - misc=None, - is_override=False, -): - print("inside the save_layer_info_to_db function") +def _get_prod_backend_url(): + return getattr(settings, "PROD_BACKEND_URL", "").rstrip("/") - dataset = Dataset.objects.get(name=dataset_name) - try: - state_obj = StateSOI.objects.get(state_name__iexact=state) - district_obj = DistrictSOI.objects.get( - district_name__iexact=district, state=state_obj - ) - block_obj = TehsilSOI.objects.get( - tehsil_name__iexact=block, district=district_obj - ) - except Exception as e: - print("Error fetching in state district block:", e) - return +def _get_prod_api_key(): + return getattr(settings, "PROD_BACKEND_API_KEY", "") - is_public = is_asset_public(asset_id) - # Check if there’s an existing layer - existing_layer = ( - Layer.objects.filter( - dataset=dataset, - layer_name=layer_name, - state=state_obj, - district=district_obj, - block=block_obj, - ) - .order_by("-layer_version") - .first() - ) +def _sync_layer_to_prod_db(payload: dict): + prod_url = _get_prod_backend_url() + if not prod_url: + return None - if existing_layer: - if existing_layer.algorithm_version != algorithm_version: - # Algorithm version changed --> create new record with incremented layer_version - new_layer_version = str(float(existing_layer.layer_version) + 1) - print( - f"Algorithm version changed. Creating new layer version: {new_layer_version}" - ) - layer_obj = Layer.objects.create( - dataset=dataset, - layer_name=layer_name, - state=state_obj, - district=district_obj, - block=block_obj, - layer_version=new_layer_version, - algorithm=algorithm, - algorithm_version=algorithm_version, - is_sync_to_geoserver=sync_to_geoserver, - is_public_gee_asset=is_public, - is_override=is_override, - misc=misc, - gee_asset_path=asset_id, - ) - else: - # Algorithm version is same --> update existing layer - print("Algorithm version same. Updating existing layer.") - for field, value in { - "algorithm": algorithm, - "algorithm_version": algorithm_version, - "is_sync_to_geoserver": sync_to_geoserver, - "is_public_gee_asset": is_public, - "is_override": is_override, - "misc": misc, - "gee_asset_path": asset_id, - }.items(): - setattr(existing_layer, field, value) - existing_layer.save() - layer_obj = existing_layer - else: - # No existing record --> create a new one - print("No existing layer found. Creating new one.") - layer_obj = Layer.objects.create( - dataset=dataset, - layer_name=layer_name, - state=state_obj, - district=district_obj, - block=block_obj, - layer_version=layer_version, - algorithm=algorithm, - algorithm_version=algorithm_version, - is_sync_to_geoserver=sync_to_geoserver, - is_public_gee_asset=is_public, - is_override=is_override, - misc=misc, - gee_asset_path=asset_id, + endpoint = prod_url + "/api/v1/sync_layer_remote/" + try: + response = requests.post( + endpoint, + json=payload, + headers={"X-Api-Key": _get_prod_api_key()}, + timeout=30, ) + if response.status_code not in (200, 201): + logger.warning( + "Prod DB sync returned %s for layer %s: %s", + response.status_code, + payload.get("layer_name"), + response.text, + ) + return None + layer_id = response.json().get("layer_id") + logger.info("Layer %s synced to prod DB (id=%s).", payload.get("layer_name"), layer_id) + return layer_id + except requests.RequestException as e: + logger.error("Failed to sync layer %s to prod DB: %s", payload.get("layer_name"), e) + return None - print(f"Saved layer info (id={layer_obj.id}, version={layer_obj.layer_version})") - return layer_obj.id +def _update_layer_sync_remote(layer_id, sync_to_geoserver=None, is_stac_specs_generated=None): + prod_url = _get_prod_backend_url() + if not prod_url or layer_id is None: + return -def update_layer_sync_status( - layer_id, sync_to_geoserver=None, is_stac_specs_generated=None -): + endpoint = prod_url + "/api/v1/update_layer_sync_remote/" + payload = { + "layer_id": layer_id, + "sync_to_geoserver": sync_to_geoserver, + "is_stac_specs_generated": is_stac_specs_generated, + } try: - layer_obj = Layer.objects.filter(id=layer_id) - if sync_to_geoserver is not None: - updated_count = layer_obj.update(is_sync_to_geoserver=sync_to_geoserver) - - if updated_count > 0: - print( - f"Updated sync status to {sync_to_geoserver} for layer ID: {layer_id}" - ) - return layer_id - - if is_stac_specs_generated is not None: - updated_count = layer_obj.update( - is_stac_specs_generated=is_stac_specs_generated + response = requests.post( + endpoint, + json=payload, + headers={"X-Api-Key": _get_prod_api_key()}, + timeout=30, + ) + if response.status_code not in (200, 201): + logger.warning( + "Prod layer sync status update returned %s for layer %s: %s", + response.status_code, + layer_id, + response.text, ) - - if updated_count > 0: - print( - f"Updated sync status to {is_stac_specs_generated} for layer ID: {layer_id}" - ) - return layer_id - - except Exception as e: - print(f"Error updating layer sync status: {e}") + else: + logger.info("Layer sync status updated on prod DB for id=%s.", layer_id) + except requests.RequestException as e: + logger.error("Failed to update layer sync status on prod DB for id=%s: %s", layer_id, e) def get_existing_end_year(dataset_name, layer_name): @@ -1022,73 +962,6 @@ def compute_for_feature(feature): -def _get_prod_backend_url(): - return getattr(settings, "PROD_BACKEND_URL", "").rstrip("/") - - -def _get_prod_api_key(): - return getattr(settings, "PROD_BACKEND_API_KEY", "") - - -def _sync_layer_to_prod_db(payload: dict): - prod_url = _get_prod_backend_url() - if not prod_url: - return None - - endpoint = prod_url + "/api/v1/computing/sync_layer_remote/" - try: - response = requests.post( - endpoint, - json=payload, - headers={"X-Api-Key": _get_prod_api_key()}, - timeout=30, - ) - if response.status_code not in (200, 201): - logger.warning( - "Prod DB sync returned %s for layer %s: %s", - response.status_code, - payload.get("layer_name"), - response.text, - ) - return None - layer_id = response.json().get("layer_id") - logger.info("Layer %s synced to prod DB (id=%s).", payload.get("layer_name"), layer_id) - return layer_id - except requests.RequestException as e: - logger.error("Failed to sync layer %s to prod DB: %s", payload.get("layer_name"), e) - return None - - -def _update_layer_sync_remote(layer_id, sync_to_geoserver=None, is_stac_specs_generated=None): - prod_url = _get_prod_backend_url() - if not prod_url or layer_id is None: - return - - endpoint = prod_url + "/api/v1/computing/update_layer_sync_remote/" - payload = { - "layer_id": layer_id, - "sync_to_geoserver": sync_to_geoserver, - "is_stac_specs_generated": is_stac_specs_generated, - } - try: - response = requests.post( - endpoint, - json=payload, - headers={"X-Api-Key": _get_prod_api_key()}, - timeout=30, - ) - if response.status_code not in (200, 201): - logger.warning( - "Prod layer sync status update returned %s for layer %s: %s", - response.status_code, - layer_id, - response.text, - ) - else: - logger.info("Layer sync status updated on prod DB for id=%s.", layer_id) - except requests.RequestException as e: - logger.error("Failed to update layer sync status on prod DB for id=%s: %s", layer_id, e) - def save_layer_info_to_db( state, @@ -1105,7 +978,7 @@ def save_layer_info_to_db( is_override=False, ): if _get_prod_backend_url(): - return _sync_layer_to_prod_db({ + layer_id = _sync_layer_to_prod_db({ "state": state, "district": district, "block": block, @@ -1119,6 +992,11 @@ def save_layer_info_to_db( "misc": misc, "is_override": is_override, }) + if layer_id is not None: + return layer_id + logger.warning( + "Prod DB sync failed for layer %s — falling back to local DB write.", layer_name + ) print("inside the save_layer_info_to_db function") @@ -1176,13 +1054,14 @@ def save_layer_info_to_db( else: # Algorithm version is same --> update existing layer print("Algorithm version same. Updating existing layer.") + merged_misc = {**(existing_layer.misc or {}), **(misc or {})} for field, value in { "algorithm": algorithm, "algorithm_version": algorithm_version, "is_sync_to_geoserver": sync_to_geoserver, "is_public_gee_asset": is_public, "is_override": is_override, - "misc": misc, + "misc": merged_misc, "gee_asset_path": asset_id, }.items(): setattr(existing_layer, field, value) diff --git a/dpr/api.py b/dpr/api.py index 1c116ea2..307b9b5d 100644 --- a/dpr/api.py +++ b/dpr/api.py @@ -449,9 +449,8 @@ def download_mws_report(request): @api_view(["GET"]) -@auth_free @schema(None) -@api_security_check(auth_type="Auth_free") +@auth_free def generate_tehsil_report(request): try: # ? district, block, mwsId diff --git a/installation/environment.yml b/installation/environment.yml index a566ba42..cbc60142 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 @@ -48,7 +50,7 @@ dependencies: - django-timezone-field==7.2.1 - djangorestframework-simplejwt==5.5.0 - djangorestframework-api-key==3.1.0 - - earthengine-api==0.1.389 + - earthengine-api==1.6.12 - geoserver-rest==2.5.3 - geojson==3.1.0 - google-api-core==2.17.0 @@ -61,7 +63,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 @@ -69,7 +70,6 @@ dependencies: - emoji - boto3 - speechrecognition==3.14.3 - - rasterio - pystac - tqdm - geemap @@ -79,4 +79,13 @@ dependencies: - pymannkendall - selenium - webdriver_manager - - celery \ No newline at end of file + - celery + - 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 diff --git a/nrm_app/settings.py b/nrm_app/settings.py index d06a3d21..abf03bc4 100755 --- a/nrm_app/settings.py +++ b/nrm_app/settings.py @@ -378,6 +378,13 @@ def resolve_env_path(name, default="", *, trailing_sep=False): PROD_BACKEND_URL = env("PROD_BACKEND_URL", default="") PROD_BACKEND_API_KEY = env("PROD_BACKEND_API_KEY", default="") +PROD_GEOSERVER_URL = env("PROD_GEOSERVER_URL", default="") +PROD_GEOSERVER_USERNAME = env("PROD_GEOSERVER_USERNAME", default="") +PROD_GEOSERVER_PASSWORD = env("PROD_GEOSERVER_PASSWORD", default="") + +PROD_BACKEND_URL = env("PROD_BACKEND_URL", default="") +PROD_BACKEND_API_KEY = env("PROD_BACKEND_API_KEY", default="") + CE_BUCKET_URL = env("CE_BUCKET_URL") EARTH_DATA_USER = env("EARTH_DATA_USER") diff --git a/utilities/active_loc_layer_generation.py b/utilities/active_loc_layer_generation.py new file mode 100644 index 00000000..7eaf1144 --- /dev/null +++ b/utilities/active_loc_layer_generation.py @@ -0,0 +1,110 @@ +import requests +import os +import json + +# Load the locations data from the JSON file +with open("/home/cfpt-jedi/developer/repos/core-stack-backend/utilities/active_location_for_layer_gen.json", "r") as file: + locations = json.load(file) + +# Base URL for API request + + +# Function to make the GET request and save the file +def download_excel(state, district, block): + # Constructing the URL with query parameters + base_url = "http://127.0.0.1:8000/api/v1/add_new_layer_data_to_excel/" + params = { + "state": state, + "district": district, + "block": block, + "workspace": "nrega_assets", + } + + print(f"Generation Excel for {state} {district} {block}") + + try: + # Make GET request to the API + response = requests.get(base_url, params=params) + + # Check if the request was successful + if response.status_code == 200: + # Ensure directory exists + # os.makedirs(os.path.dirname(save_path), exist_ok=True) + + # Saving the file (Excel file) + # with open(save_path, 'wb') as file: + # file.write(response.content) + print("File saved to:") + else: + print( + f"Failed to fetch data for {state} - {district} - {block}. Status code: {response.status_code}" + ) + except Exception as e: + print(f"An error occurred: {e}") + + +def restoration_vector(state, district, block): + # Updated base URL + base_url = "http://localhost:8000/api/v1/generate_fabdem_raster/" + + headers = { + "Authorization": "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ0b2tlbl90eXBlIjoiYWNjZXNzIiwiZXhwIjoxNzg0NTQwNTQ2LCJpYXQiOjE3NzY3NjQ1NDYsImp0aSI6IjIxYmVmZTMwYWE2MzQ4ZDM4YmE5M2EyNGQ3MjdhYzgyIiwidXNlcl9pZCI6NH0.FHqq8dRaI39L5zIDKotjYIlwnnc8JWyi9B1vDLWBhZQ", + "Content-Type": "application/json", + } + + # base_url = "https://geoserver.core-stack.org/api/v1/change_detection/" + # base_url = "https://geoserver.core-stack.org/api/v1/generate_distance_nearest_DL/" + # base_url = "https://geoserver.core-stack.org/api/v1/generate_ci_layer/" + # base_url = "https://geoserver.core-stack.org/api/v1/generate_facilities_proximity/" + + # headers = { + # "Authorization": "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ0b2tlbl90eXBlIjoiYWNjZXNzIiwiZXhwIjoxNzc5NzE0OTU5LCJpYXQiOjE3NzE5Mzg5NTksImp0aSI6IjA5OGU0MzRjYWE0ZTQ2ZDVhYmE3MTc1YjljMDAwODQ3IiwidXNlcl9pZCI6NH0.BjE41D07gARbRXCdKOcVxJjpsVMZDxVYg7oaN4ABxCg", + # "Content-Type": "application/json", + # } + + # Request bodyS + body = { + "state": state, + "district": district, + "block": block, + "start_year": 2017, + "end_year": 2024, + "compute": "local" + } + + print(f"Generating Layer for {state}, {district}, {block}") + + try: + # Make POST request to the API with headers and JSON body + response = requests.post(base_url, headers=headers, json=body) + + # Check if the request was successful + if response.status_code == 200: + print(f"Response: {response.text}") + return True + else: + print( + f"Failed to fetch data for {state} - {district} - {block}. Status code: {response.status_code}" + ) + print(f"Response: {response.text}") + return False + except Exception as e: + print(f"An error occurred: {e}") + return False + + +# Loop through the JSON data to fetch state, district, and block and download files +for entry in locations: + # state = entry["state"] + # district = entry["district"] + # block = entry["block"] + state = entry["state"] + district = entry["district"] + block = entry["block"] + + # Define the save path (e.g., saving as a .xlsx file in a folder based on district) + # save_path = os.path.join("downloaded_files", state, district, f"{block}.xlsx") + + # Download and save the Excel file + #download_excel(state, district, block) + restoration_vector(state, district, block) diff --git a/utilities/active_location_for_layer_gen.json b/utilities/active_location_for_layer_gen.json new file mode 100644 index 00000000..54dcc84d --- /dev/null +++ b/utilities/active_location_for_layer_gen.json @@ -0,0 +1,232 @@ +[ +{ + "state": "Telangana", + "district": "Vikarabad", + "block": "Bomraspet" + }, + { + "state": "Telangana", + "district": "Vikarabad", + "block": "Doma" + }, + { + "state": "Telangana", + "district": "Vikarabad", + "block": "Doulatabad" + }, + { + "state": "Uttar Pradesh", + "district": "Bahraich", + "block": "Bahraich" + }, + { + "state": "Uttar Pradesh", + "district": "Bahraich", + "block": "Nanpara" + }, + { + "state": "Uttar Pradesh", + "district": "Ballia", + "block": "Bairia" + }, + { + "state": "Uttar Pradesh", + "district": "Ballia", + "block": "Ballia" + }, + { + "state": "Uttar Pradesh", + "district": "Ballia", + "block": "Bansdih" + }, + { + "state": "Uttar Pradesh", + "district": "Ballia", + "block": "Belthara Road" + }, + { + "state": "Uttar Pradesh", + "district": "Ballia", + "block": "Rasra" + }, + { + "state": "Uttar Pradesh", + "district": "Ballia", + "block": "Sikandarpur" + }, + { + "state": "Uttar Pradesh", + "district": "Balrampur", + "block": "Balrampur" + }, + { + "state": "Uttar Pradesh", + "district": "Balrampur", + "block": "Tulsipur" + }, + { + "state": "Uttar Pradesh", + "district": "Bara banki", + "block": "Fatehpur" + }, + { + "state": "Uttar Pradesh", + "district": "Basti", + "block": "Basti" + }, + { + "state": "Uttar Pradesh", + "district": "Basti", + "block": "Bhanpur" + }, + { + "state": "Uttar Pradesh", + "district": "Basti", + "block": "Harraiya" + }, + { + "state": "Uttar Pradesh", + "district": "Basti", + "block": "Rudhauli" + }, + { + "state": "Uttar Pradesh", + "district": "Bhadohi", + "block": "Aurai" + }, + { + "state": "Uttar Pradesh", + "district": "Bhadohi", + "block": "Gyanpur" + }, + { + "state": "Uttar Pradesh", + "district": "Chitrakoot", + "block": "Karwi" + }, + { + "state": "Uttar Pradesh", + "district": "Chitrakoot", + "block": "Mau" + }, + { + "state": "Uttar Pradesh", + "district": "Deoria", + "block": "Deoria" + }, + { + "state": "Uttar Pradesh", + "district": "Fatehpur", + "block": "Binki" + }, + { + "state": "Uttar Pradesh", + "district": "Fatehpur", + "block": "Fatehpur" + }, + { + "state": "Uttar Pradesh", + "district": "Fatehpur", + "block": "Khaga" + }, + { + "state": "Uttar Pradesh", + "district": "Hamirpur", + "block": "Maudaha" + }, + { + "state": "Uttar Pradesh", + "district": "Jaunpur", + "block": "Badlapur" + }, + { + "state": "Uttar Pradesh", + "district": "Jhansi", + "block": "Jhansi" + }, + { + "state": "Uttar Pradesh", + "district": "Kaushambi", + "block": "Chail" + }, + { + "state": "Uttar Pradesh", + "district": "Lucknow", + "block": "Lucknow" + }, + { + "state": "Uttar Pradesh", + "district": "Mau", + "block": "Ghosi" + }, + { + "state": "Uttar Pradesh", + "district": "Mau", + "block": "Madhuban" + }, + { + "state": "Uttar Pradesh", + "district": "Mau", + "block": "Mau" + }, + { + "state": "Uttar Pradesh", + "district": "Mau", + "block": "Muhammadabad" + }, + { + "state": "Uttar Pradesh", + "district": "Mirzapur", + "block": "Marihan" + }, + { + "state": "Uttar Pradesh", + "district": "Mirzapur", + "block": "Mirzapur" + }, + { + "state": "Uttar Pradesh", + "district": "Shrawasti", + "block": "Bhinga" + }, + { + "state": "West Bengal", + "district": "Bankura", + "block": "Bankura" + }, + { + "state": "West Bengal", + "district": "Darjiling", + "block": "Karsiyang" + }, + { + "state": "West Bengal", + "district": "Darjiling", + "block": "Siliguri" + }, + { + "state": "West Bengal", + "district": "North Twenty-Four Parganas", + "block": "Basirhat" + }, + { + "state": "West Bengal", + "district": "Purba Medinipur", + "block": "Tamluk" + }, + { + "state": "West Bengal", + "district": "Puruliya", + "block": "Raghunathpur" + }, + { + "state": "West Bengal", + "district": "South Twenty-Four Parganas", + "block": "Canning" + }, + { + "state": "West Bengal", + "district": "South Twenty-Four Parganas", + "block": "Kakdwip" + } +] diff --git a/utilities/geoserver_utils.py b/utilities/geoserver_utils.py index c55289a1..78b4599a 100644 --- a/utilities/geoserver_utils.py +++ b/utilities/geoserver_utils.py @@ -1,8 +1,11 @@ # inbuilt libraries import io +import logging import os from typing import List, Optional, Set +logger = logging.getLogger(__name__) + # third-party libraries import requests from xmltodict import parse, unparse @@ -246,6 +249,22 @@ def create_workspace(self, workspace: str): else: raise GeoserverException(r.status_code, r.content) + def ensure_workspace(self, workspace: str): + """ + Creates the workspace if it does not already exist. + Never raises — a failed workspace creation is non-fatal; the upload + attempt that follows will surface the real error if the workspace is + truly missing. + """ + try: + url = "{}/rest/workspaces/{}.json".format(self.service_url, workspace) + r = self._requests("get", url) + if r.status_code == 200: + return + self.create_workspace(workspace) + except Exception as e: + logger.warning("ensure_workspace(%s) failed (proceeding anyway): %s", workspace, e) + def delete_workspace(self, workspace: str): """ @@ -320,7 +339,7 @@ def delete_vector_store(self, workspace, store): - store: Name of the vector datastore to delete """ print("inside delete_vector_store") - url = f"{GEOSERVER_URL}/rest/workspaces/{workspace}/datastores/{store}?recurse=true" + url = f"{self.service_url}/rest/workspaces/{workspace}/datastores/{store}?recurse=true" resp = self._requests(method="delete", url=url) if resp.status_code in [200, 202]: print(f"Vector store '{store}' deleted successfully.") @@ -416,14 +435,13 @@ def create_coveragestore( headers = {"content-type": content_type, "Accept": "application/json"} - r = None with open(path, "rb") as f: r = self._requests(method="put", url=url, data=f, headers=headers) - if r.status_code == 201: - return r.json() - else: - raise GeoserverException(r.status_code, r.content) + if r.status_code in (200, 201): + return r.json() + else: + raise GeoserverException(r.status_code, r.content) def upload_raster( self, @@ -513,6 +531,81 @@ def publish_time_dimension_to_coveragestore( else: raise GeoserverException(r.status_code, r.content) + def configure_coverage( + self, + workspace: str, + store: str, + coverage: Optional[str] = None, + srs: str = "EPSG:4326", + projection_policy: str = "FORCE_DECLARED", + enabled: bool = True, + ): + """ + Force a coverage's declared SRS, projection policy, and enabled state. + + Use this right after ``create_coveragestore`` when the source raster + may have missing/unresolvable CRS metadata (e.g. broken PROJ on the + writer side). ``FORCE_DECLARED`` tells GeoServer to treat the declared + SRS as authoritative instead of relying on the file's native CRS. + """ + coverage_name = coverage or store + url = "{0}/rest/workspaces/{1}/coveragestores/{2}/coverages/{3}".format( + self.service_url, workspace, store, coverage_name + ) + body = ( + "" + f"{coverage_name}" + f"{coverage_name}" + f"{str(enabled).lower()}" + f"{srs}" + f"{projection_policy}" + "" + ) + r = self._requests( + method="put", + url=url, + data=body, + headers={"content-type": "application/xml"}, + ) + if r.status_code in (200, 201): + return r.status_code + raise GeoserverException(r.status_code, r.content) + + def configure_featuretype( + self, + workspace: str, + store: str, + featuretype: Optional[str] = None, + srs: str = "EPSG:4326", + projection_policy: str = "FORCE_DECLARED", + ): + """ + Force a featuretype's declared SRS and projection policy. + + Use after publishing a vector store when the source file may have + missing or mismatched CRS metadata. ``FORCE_DECLARED`` treats the + declared SRS as authoritative without reprojecting coordinates. + """ + ft_name = featuretype or store + url = "{0}/rest/workspaces/{1}/datastores/{2}/featuretypes/{3}".format( + self.service_url, workspace, store, ft_name + ) + body = ( + "" + f"{srs}" + f"{projection_policy}" + "" + ) + r = self._requests( + method="put", + url=url, + data=body, + headers={"content-type": "application/xml"}, + ) + if r.status_code in (200, 201): + return r.status_code + raise GeoserverException(r.status_code, r.content) + # delete coveragestore(raster) def delete_raster_store(self, workspace, store): """ @@ -523,7 +616,7 @@ def delete_raster_store(self, workspace, store): - store: Name of the vector datastore to delete """ print("inside delete_raster_store") - url = f"{GEOSERVER_URL}/rest/workspaces/{workspace}/coveragestores/{store}?recurse=true" + url = f"{self.service_url}/rest/workspaces/{workspace}/coveragestores/{store}?recurse=true" resp = self._requests(method="delete", url=url) if resp.status_code in [200, 202]: print(f"Raster store '{store}' deleted successfully.")