geoextract-gee 0.2.0__py3-none-any.whl

geoextract_gee/asset_manager.py

@@ -0,0 +1,128 @@
+import hashlib
+import time
+import uuid
+from pathlib import Path
+
+import ee
+from google.cloud import storage
+
+from .datasets.common import gee_setting
+
+def get_file_md5(file_path: Path) -> str:
+    """Calculates the MD5 hash of a local file."""
+    hash_md5 = hashlib.md5()
+    with open(file_path, "rb") as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            hash_md5.update(chunk)
+    return hash_md5.hexdigest()
+
+def check_asset_exists_and_matches(asset_id: str, local_md5: str) -> bool:
+    """Checks if an EE asset exists and has a matching local_md5 property."""
+    try:
+        asset = ee.data.getAsset(asset_id)
+        asset_md5 = asset.get('properties', {}).get('local_md5')
+        return asset_md5 == local_md5
+    except ee.ee_exception.EEException:
+        # Asset doesn't exist or we don't have access
+        return False
+
+def upload_to_gcs(local_path: Path, bucket_name: str, logger) -> str:
+    """Uploads a local file to GCS and returns the gs:// URI."""
+    client = storage.Client()
+    bucket = client.bucket(bucket_name)
+    # Put them in a staging folder so they don't clutter the export directory
+    dest_blob_name = f"gee_staging/{local_path.name}"
+    blob = bucket.blob(dest_blob_name)
+    
+    logger.info(f"Uploading {local_path.name} to gs://{bucket_name}/{dest_blob_name}...")
+    blob.upload_from_filename(str(local_path))
+    return f"gs://{bucket_name}/{dest_blob_name}"
+
+def start_ingestion(gcs_uri: str, asset_id: str, local_md5: str):
+    """Triggers an Earth Engine ingestion task from a GCS URI."""
+    manifest = {
+        "name": asset_id,
+        "tilesets": [
+            {"id": "tileset_1", "sources": [{"uris": [gcs_uri]}]}
+        ],
+        "properties": {
+            "local_md5": local_md5
+        }
+    }
+    request_id = str(uuid.uuid4())
+    ee.data.startIngestion(request_id, manifest)
+
+def sync_cropmasks(params, unique_paths: set) -> dict:
+    """
+    Synchronizes local TIFF cropmasks with Earth Engine.
+    Returns a dict mapping the local Path string to the final GEE Asset ID.
+    """
+    params.logger.info(f"Synchronizing {len(unique_paths)} unique cropmasks with Earth Engine...")
+    
+    project_name = gee_setting(params, "project")
+    bucket_name = gee_setting(params, "bucket")
+    
+    # 1. Ensure the parent folder exists in GEE
+    folder_id = f"projects/{project_name}/assets/geoextract"
+    try:
+        ee.data.createAsset({'type': 'FOLDER'}, folder_id)
+        params.logger.info(f"Created new GEE folder: {folder_id}")
+    except ee.ee_exception.EEException:
+        pass # Already exists
+        
+    cropmask_map = {}
+    pending_ingestions = []
+    
+    # 2. Check and stage files
+    for path in unique_paths:
+        path_obj = Path(path)
+        asset_name = path_obj.stem # Removes the path and .tif extension
+        asset_id = f"{folder_id}/{asset_name}"
+        cropmask_map[str(path)] = asset_id
+        
+        if not path_obj.exists():
+            params.logger.warning(f"Cropmask file not found locally: {path_obj}")
+            continue
+            
+        md5 = get_file_md5(path_obj)
+        
+        # Check if it already exists in GEE and matches our hash
+        if check_asset_exists_and_matches(asset_id, md5):
+            params.logger.debug(f"Asset {asset_name} already up to date in GEE.")
+            continue
+            
+        # It needs to be uploaded
+        gcs_uri = upload_to_gcs(path_obj, bucket_name, params.logger)
+        
+        params.logger.info(f"Triggering ingestion for {asset_name}...")
+        start_ingestion(gcs_uri, asset_id, md5)
+        pending_ingestions.append(asset_id)
+        
+    # 3. Block and wait for ingestions to finish
+    if pending_ingestions:
+        params.logger.info(f"Waiting for {len(pending_ingestions)} cropmasks to finish ingesting...")
+        
+        while True:
+            ops = ee.data.listOperations()
+            active_ingestions = 0
+            
+            for op in ops:
+                metadata = op.get('metadata', {})
+                state = metadata.get('state')
+                
+                # Check if it's an image ingestion task
+                if metadata.get('type') == 'INGEST_IMAGE':
+                    if state in ['PENDING', 'RUNNING']:
+                        active_ingestions += 1
+                    elif state == 'FAILED':
+                        # Optionally check if it's one of ours, but usually safe to log
+                        params.logger.error(f"An ingestion failed: {op.get('name')}")
+                        
+            if active_ingestions == 0:
+                params.logger.info("All cropmask ingestions have completed.")
+                break
+                
+            params.logger.info(f"{active_ingestions} assets still ingesting. Waiting 30s...")
+            time.sleep(30)
+            
+    return cropmask_map

geoextract_gee/datasets/aef.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import ee
+from .common import TaskConfig, transform_geometry
+
+# --- Dataset Constants ---
+AEF_COLLECTION_ID = "GOOGLE/SATELLITE_EMBEDDING/V1/ANNUAL"
+
+def _aef_stats_feature_batch(
+    avg_img: ee.Image,
+    config: TaskConfig,
+    mask_threshold_percent: float,
+) -> ee.FeatureCollection:
+    """
+    One row per region representing the multi-year average of 
+    the 64 AEF embedding bands over the crop-masked area.
+    """
+    avg_img = ee.Image(avg_img)
+
+    afi_thresh = ee.Number(float(mask_threshold_percent * 100))
+    fc_bounds = config.geometry_r.geometry().bounds(1)
+    w_raw = ee.Image(config.cropmask_asset).float().clip(fc_bounds)
+    w = w_raw.updateMask(w_raw.gt(afi_thresh))
+
+    # Mask the AEF image to the crop area
+    new_names = [f"aef_{i+1}" for i in range(64)]
+    
+    p = avg_img.rename(new_names).float().clip(fc_bounds)
+    pf = p.updateMask(w.mask())
+
+    reducer_common_parms = {
+        "crs": config.reduce_crs,
+        "scale": config.reduce_scale,
+        "maxPixels": 1e9,
+        "tileScale": 4,
+        "bestEffort": True,
+    }
+
+    def _process_region(feat):
+        geom = feat.geometry()
+
+        # Returns a dict like {"A00": 0.123, "A01": 0.456, ...}
+        d_mean = pf.reduceRegion(
+            reducer=ee.Reducer.mean().unweighted(),
+            geometry=geom,
+            **reducer_common_parms
+        )
+        
+        props = ee.Dictionary({
+            "region_label": feat.get("region_label"),
+            "region_id": feat.get("region_id")
+        }).combine(d_mean)
+
+        return ee.Feature(None, props)
+
+    return ee.FeatureCollection(config.geometry_r).map(_process_region)
+
+
+def create_task(config: TaskConfig, mask_threshold_percent: float) -> ee.batch.Task:
+    """Creates a GEE Export task for AEF (AlphaEarth Foundations) 64-band data."""
+    
+    # AEF has data from 2018 onwards
+    start_year = int(config.date_from.split('-')[0])
+    end_year = int(config.date_to.split('-')[0])
+    
+    start_year = max(2018, start_year)
+    end_year = min(2025, end_year - 1)
+    
+    # If the user requested a year completely outside AEF bounds (e.g., 2026-2027),
+    # we fall back to the global 2018-2025 average so geomerge still gets the columns.
+    # This mimics extract_EO using the cached aef_avg.tif for future years
+    if start_year > end_year:
+        start_year = 2018
+        end_year = 2025
+        
+    # Annual images are at 01-01T00:00:00Z, 
+    ic = ee.ImageCollection(AEF_COLLECTION_ID).filterDate(f"{start_year}-01-01", f"{end_year}-01-02")
+    
+    avg_image = ic.mean()
+
+    # ensure projection equals dataset projection to decrease reducers cost
+    aef_proj = ee.Projection(
+        ee.Algorithms.If(
+            ic.size().gt(0), 
+            ic.first().select("A00").projection(), 
+            ee.Projection('EPSG:4326')
+        )
+    )
+    
+    config.geometry_r = transform_geometry(config.geometry_ee, aef_proj)
+    config.reduce_crs = aef_proj.crs()
+    
+    # Force scale to 5600m (0.05 degrees) to match local equivalents
+    # and leverage GEE Image Pyramids for the 64-band reduction
+    config.reduce_scale = 5600
+
+    # We only want one row per region
+    # So wrap the single average image in a single-feature FeatureCollection.
+    col = _aef_stats_feature_batch(avg_image, config, mask_threshold_percent)
+
+    # just region_label and 64 AEF bands
+    selectors = ["region_label"] + [f"aef_{i+1}" for i in range(64)]
+
+    task_desc = config.task_desc
+    
+    return ee.batch.Export.table.toCloudStorage(
+        collection=col,
+        description=task_desc[:100],
+        bucket=config.export_bucket,
+        fileNamePrefix=config.export_prefix,
+        fileFormat="CSV",
+        selectors=selectors,
+    )

geoextract_gee/datasets/agera5.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+import ee
+from .common import TaskConfig, transform_geometry
+
+# --- Dataset Constants ---
+AGERA5_COLLECTION_ID = "projects/climate-engine-pro/assets/ce-ag-era5-v2/daily"
+
+# var -> band. Only variables the tiles backend also extracts, so both backends
+# give the same columns (docs/decisions/001).
+AGERA5_BANDS = {
+    "agera5_snow_thickness_lwe": "Snow_Thickness_LWE_Mean_24h",
+    "agera5_dewpoint": "Dew_Point_Temperature_2m_Mean_24h",
+}
+
+def _agera5_daily_stats_feature_batch(
+    img: ee.Image,
+    config: TaskConfig,
+    band_name: str,
+    mask_threshold_percent: float,
+) -> ee.FeatureCollection:
+    """
+    One row per day per region, aligned with geoextract geom_extract + arr_stats.
+    """
+    img = ee.Image(img)
+
+    afi_thresh = ee.Number(float(mask_threshold_percent * 100))
+    fc_bounds = config.geometry_r.geometry().bounds(1)
+    w_raw = ee.Image(config.cropmask_asset).float().clip(fc_bounds)
+    w = w_raw.updateMask(w_raw.gt(afi_thresh))
+
+
+    p = (
+        img.select(band_name)
+        .float()
+        .clip(fc_bounds)
+    )
+    
+    # Same masks and units as the tiles backend (geoprepare extract/stats.py)
+    if config.var == 'agera5_snow_thickness_lwe':
+        p = p.updateMask(p.gte(0))
+    elif config.var == 'agera5_dewpoint':
+        # Kelvin -> deg C; below 150 K is fill, not weather
+        p = p.updateMask(p.gte(150)).subtract(273.15)
+        
+    pf = p.updateMask(w.mask())
+    pf_and_w = pf.addBands(w.rename("weight"))
+
+    if config.include_audit:
+        stack_audit = ee.Image.cat([
+            ee.Image.constant(1).clip(fc_bounds).rename("total_w"),
+            ee.Image.constant(1).updateMask(p.mask()).rename("valid_p_w"),
+            w.rename("w_all"),
+            w.updateMask(pf.mask()).rename("w_used"),
+            pf.pow(2).multiply(w).rename("p2w"),
+        ]).float()
+    else:
+        stack_audit = None
+
+    reducer_common_parms = {
+        "crs": config.reduce_crs,
+        "scale": config.reduce_scale,
+        "maxPixels": 1e9,
+        "tileScale": 4,
+        "bestEffort": True,
+    }
+
+    def _process_region(feat):
+        geom = feat.geometry()
+
+        # Common stats: mean, min, max, median, count (valid after mask)
+        d_mean = pf_and_w.reduceRegion(
+            reducer=ee.Reducer.mean().splitWeights(),
+            geometry=geom,
+            **reducer_common_parms
+        )
+        
+        # Combined min, max, median, and unweighted count
+        d_stats = pf.reduceRegion(
+            reducer=ee.Reducer.minMax()
+                .combine(ee.Reducer.median().unweighted(), sharedInputs=True)
+                .combine(ee.Reducer.count().unweighted(), sharedInputs=True),
+            geometry=geom,
+            **reducer_common_parms
+        )
+
+        valid_after = d_stats.getNumber(f"{band_name}_count")
+        mean_raw = d_mean.getNumber("mean")
+        
+        props = ee.Dictionary({
+            "date": img.date().format("YYYY-MM-dd"),
+            "region_label": feat.get("region_label"),
+            "region_id": feat.get("region_id"),
+            "stats_mean": mean_raw,
+            "stats_min": d_stats.get(f"{band_name}_min"),
+            "stats_max": d_stats.get(f"{band_name}_max"),
+            "stats_median": d_stats.get(f"{band_name}_median"),
+            "stats_count": valid_after,
+        })
+
+        if config.include_audit:
+            d_audit = stack_audit.reduceRegion(
+                reducer=ee.Reducer.sum().unweighted(),
+                geometry=geom,
+                **reducer_common_parms
+            )
+
+            w_used_sum = d_audit.getNumber("w_used")
+            sum_p2w = d_audit.getNumber("p2w")
+            
+            # Weighted Std Dev logic
+            var_raw = ee.Number(
+                ee.Algorithms.If(
+                    w_used_sum.gt(0),
+                    sum_p2w.divide(w_used_sum).subtract(mean_raw.pow(2)),
+                    0,
+                )
+            )
+            std_mm = ee.Number(
+                ee.Algorithms.If(
+                    w_used_sum.lte(0),
+                    None,
+                    ee.Algorithms.If(valid_after.gt(1), var_raw.max(0).sqrt(), 0),
+                )
+            )
+
+            props = props.combine(ee.Dictionary({
+                "counts_total": d_audit.getNumber("total_w"),
+                "counts_valid_data": d_audit.getNumber("valid_p_w"),
+                "counts_valid_data_after_masking": valid_after,
+                "counts_weight_sum": d_audit.getNumber("w_all"),
+                "counts_weight_sum_used": w_used_sum,
+                "stats_std": std_mm,
+            }))
+
+        return ee.Feature(None, props)
+
+    return ee.FeatureCollection(config.geometry_r).map(_process_region)
+
+
+def create_task(config: TaskConfig, mask_threshold_percent: float) -> ee.batch.Task:
+    """Creates a GEE Export task for AgERA5 data."""
+    ic = (
+        ee.ImageCollection(AGERA5_COLLECTION_ID)
+        .filterDate(config.date_from, config.date_to)
+        .filterBounds(config.geometry_ee)
+    )
+
+    band_name = AGERA5_BANDS[config.var]
+
+    # ensure projection equals dataset projection to decrease reducers cost
+    agera5_proj = ee.Projection(
+        ee.Algorithms.If(
+            ic.size().gt(0), 
+            ic.first().select(band_name).projection(), 
+            ee.Projection('EPSG:4326')
+        )
+    )
+    
+    config.geometry_r = transform_geometry(config.geometry_ee, agera5_proj)
+    config.reduce_crs = agera5_proj.crs()
+    config.reduce_scale = agera5_proj.nominalScale()
+
+    col = ic.map(lambda im: _agera5_daily_stats_feature_batch(im, config, band_name, mask_threshold_percent)).flatten()
+
+    selectors = ["date", "region_label", "stats_mean", "stats_min", "stats_max", "stats_median", "stats_count"]
+    if config.include_audit:
+        selectors += [
+            "counts_total", "counts_valid_data", "counts_valid_data_after_masking",
+            "counts_weight_sum", "counts_weight_sum_used", "stats_std"
+        ]
+
+    task_desc = config.task_desc
+    
+    return ee.batch.Export.table.toCloudStorage(
+        collection=col,
+        description=task_desc[:100],
+        bucket=config.export_bucket,
+        fileNamePrefix=config.export_prefix,
+        fileFormat="CSV",
+        selectors=selectors,
+    )

geoextract_gee/datasets/chirps.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+import ee
+from .common import TaskConfig, transform_geometry
+
+# --- Dataset Constants ---
+CHIRPS_COLLECTION_ID = "UCSB-CHC/CHIRPS/V3/DAILY_SAT"
+
+def _chirps_daily_stats_feature_batch(
+    img: ee.Image,
+    config: TaskConfig,
+    mask_threshold_percent: float,
+) -> ee.FeatureCollection:
+    """
+    One row per day per region, aligned with geoextract geom_extract + arr_stats (chirps + AFI).
+    """
+    img = ee.Image(img)
+
+    afi_thresh = ee.Number(float(mask_threshold_percent * 100))
+    fc_bounds = config.geometry_r.geometry().bounds(1)
+    w_raw = ee.Image(config.cropmask_asset).float().clip(fc_bounds)
+    w = w_raw.updateMask(w_raw.gt(afi_thresh))
+
+    p = (
+        img.select("precipitation")
+        .float()
+        .multiply(100)
+        .int16()
+        .divide(100)
+        .clip(fc_bounds)
+    )
+    pf = p.updateMask(p.gte(0)).updateMask(w.mask())
+    pf_and_w = pf.addBands(w.rename("weight"))
+
+    if config.include_audit:
+        stack_audit = ee.Image.cat([
+            ee.Image.constant(1).clip(fc_bounds).rename("total_w"),
+            ee.Image.constant(1).updateMask(p.mask()).rename("valid_p_w"),
+            w.rename("w_all"),
+            w.updateMask(pf.mask()).rename("w_used"),
+            pf.pow(2).multiply(w).rename("p2w"),
+        ]).float()
+    else:
+        stack_audit = None
+
+    reducer_common_parms = {
+        "crs": config.reduce_crs,
+        "scale": config.reduce_scale,
+        "maxPixels": 1e9,
+        "tileScale": 4,
+        "bestEffort": True,
+    }
+
+    def _process_region(feat):
+        geom = feat.geometry()
+
+        # Common stats: mean, min, max, median, count (valid after mask)
+        d_mean = pf_and_w.reduceRegion(
+            reducer=ee.Reducer.mean().splitWeights(),
+            geometry=geom,
+            **reducer_common_parms
+        )
+        
+        # Combined min, max, median, and unweighted count
+        d_stats = pf.reduceRegion(
+            reducer=ee.Reducer.minMax()
+                .combine(ee.Reducer.median().unweighted(), sharedInputs=True)
+                .combine(ee.Reducer.count().unweighted(), sharedInputs=True),
+            geometry=geom,
+            **reducer_common_parms
+        )
+
+        valid_after = d_stats.getNumber("precipitation_count")
+        mean_raw = d_mean.getNumber("mean")
+        
+        props = ee.Dictionary({
+            "date": img.date().format("YYYY-MM-dd"),
+            "region_label": feat.get("region_label"),
+            "region_id": feat.get("region_id"),
+            "stats_mean": mean_raw,
+            "stats_min": d_stats.get("precipitation_min"),
+            "stats_max": d_stats.get("precipitation_max"),
+            "stats_median": d_stats.get("precipitation_median"),
+            "stats_count": valid_after,
+        })
+
+        if config.include_audit:
+            d_audit = stack_audit.reduceRegion(
+                reducer=ee.Reducer.sum().unweighted(),
+                geometry=geom,
+                **reducer_common_parms
+            )
+
+            w_used_sum = d_audit.getNumber("w_used")
+            sum_p2w = d_audit.getNumber("p2w")
+            
+            # Weighted Std Dev logic
+            var_raw = ee.Number(
+                ee.Algorithms.If(
+                    w_used_sum.gt(0),
+                    sum_p2w.divide(w_used_sum).subtract(mean_raw.pow(2)),
+                    0,
+                )
+            )
+            std_mm = ee.Number(
+                ee.Algorithms.If(
+                    w_used_sum.lte(0),
+                    None,
+                    ee.Algorithms.If(valid_after.gt(1), var_raw.max(0).sqrt(), 0),
+                )
+            )
+
+            props = props.combine(ee.Dictionary({
+                "counts_total": d_audit.getNumber("total_w"),
+                "counts_valid_data": d_audit.getNumber("valid_p_w"),
+                "counts_valid_data_after_masking": valid_after,
+                "counts_weight_sum": d_audit.getNumber("w_all"),
+                "counts_weight_sum_used": w_used_sum,
+                "stats_std": std_mm,
+            }))
+
+        return ee.Feature(None, props)
+
+    return ee.FeatureCollection(config.geometry_r).map(_process_region)
+
+
+def create_task(config: TaskConfig, mask_threshold_percent: float) -> ee.batch.Task:
+    """Creates a GEE Export task for CHIRPS data."""
+    ic = (
+        ee.ImageCollection(CHIRPS_COLLECTION_ID)
+        .filterDate(config.date_from, config.date_to)
+        .filterBounds(config.geometry_ee)
+    )
+
+    # ensure projection equals dataset projection to decrease reducers cost
+    chirps_proj = ee.Projection(
+        ee.Algorithms.If(
+            ic.size().gt(0), 
+            ic.first().select("precipitation").projection(), 
+            ee.Projection('EPSG:4326')
+        )
+    )
+    
+    config.geometry_r = transform_geometry(config.geometry_ee, chirps_proj)
+    config.reduce_crs = chirps_proj.crs()
+    config.reduce_scale = chirps_proj.nominalScale()
+
+    col = ic.map(lambda im: _chirps_daily_stats_feature_batch(im, config, mask_threshold_percent)).flatten()
+
+    selectors = ["date", "region_label", "stats_mean", "stats_min", "stats_max", "stats_median", "stats_count"]
+    if config.include_audit:
+        selectors += [
+            "counts_total", "counts_valid_data", "counts_valid_data_after_masking",
+            "counts_weight_sum", "counts_weight_sum_used", "stats_std"
+        ]
+
+    # GEE task description only allows letters, numbers, -, and _
+    # task_desc is now centralized in TaskConfig
+    task_desc = config.task_desc
+    
+    return ee.batch.Export.table.toCloudStorage(
+        collection=col,
+        description=task_desc[:100],
+        bucket=config.export_bucket,
+        fileNamePrefix=config.export_prefix,
+        fileFormat="CSV",
+        selectors=selectors,
+    )