File size: 4,248 Bytes
c2a61b6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """
Land-Sea Mask for ERA5 India domain.
IMPROVEMENT 4: Land-sea mask
The AI model (Swin2SR) has no knowledge of land vs ocean.
Applying super-resolution to ocean pixels creates hallucinated
fine-scale SST structure that doesn't exist physically.
This module builds a land mask at both ERA5 resolution (129x121)
and output resolution (516x484) using two methods:
Method 1 - Temperature variance proxy:
Ocean pixels have very low temporal variance (SST is smooth and slow).
Land pixels have high variance (diurnal cycle, seasons, weather).
Threshold the annual std dev to separate land from sea.
Method 2 - Known ocean bounds (backup):
Hard-code the approximate lat/lon rectangles of:
- Arabian Sea (west coast)
- Bay of Bengal (east coast)
- Indian Ocean (south)
Use as fallback or cross-validation.
The mask is applied after AI inference:
final_output[ocean] = ERA5_upsampled[ocean] (original SST preserved)
final_output[land] = AI_prediction[land] (AI enhancement applied)
This is scientifically correct: ERA5 SST is already high quality
at 0.25deg. Only land surface temperatures benefit from SR enhancement.
"""
import numpy as np
from scipy.ndimage import zoom as spz, binary_dilation
def build_land_mask(data: np.ndarray,
lat_min: float, lat_max: float,
lon_min: float, lon_max: float,
output_scale: int = 4) -> tuple:
"""
Build land mask from temporal variance of ERA5 data.
Returns:
mask_lr : (H, W) bool True=land at ERA5 resolution
mask_hr : (4H, 4W) bool True=land at output resolution
"""
T, H, W = data.shape
# Temporal std dev at each grid point across all 8784 timesteps
std_map = data.std(axis=0) # (H, W)
# Ocean = low variance (SST changes slowly)
# Land = high variance (weather, diurnal, seasonal)
# Use Otsu-like threshold: midpoint between min and median
thresh = (std_map.min() + np.percentile(std_map, 40)) / 2
land_mask_lr = std_map > thresh # True = land
# Morphological cleanup — remove isolated ocean pixels on coast
land_mask_lr = binary_dilation(land_mask_lr, iterations=1)
# Hard-mask: known open ocean bounding boxes
lats = np.linspace(lat_max, lat_min, H) # row 0 = north
lons = np.linspace(lon_min, lon_max, W)
ocean_boxes = [
# (lat_north, lat_south, lon_west, lon_east)
# Conservative boxes — only clear open ocean, not coastal zones
# Coastal land (Gujarat, Kerala, AP coast) handled by variance method
(38.0, 6.0, 68.0, 71.0), # Arabian Sea open water (far west)
(12.0, 6.0, 71.0, 79.0), # Indian Ocean (south tip, below tip of India)
(22.0, 6.0, 88.0, 98.0), # Bay of Bengal (far east open water)
]
for ln, ls, lw, le in ocean_boxes:
for r in range(H):
for c in range(W):
if ls <= lats[r] <= ln and lw <= lons[c] <= le:
land_mask_lr[r, c] = False # force ocean
# Upsample to output resolution
land_mask_hr_f = spz(land_mask_lr.astype(float), output_scale, order=0)
land_mask_hr = land_mask_hr_f > 0.5
n_land = land_mask_lr.sum()
n_total = H * W
print(f" [Mask] Land pixels: {n_land}/{n_total} "
f"({n_land/n_total*100:.1f}%) at ERA5 resolution")
print(f" [Mask] Output mask: {land_mask_hr.shape}")
return land_mask_lr, land_mask_hr
def apply_land_mask(pred_K: np.ndarray,
era5_K: np.ndarray,
mask_hr: np.ndarray) -> np.ndarray:
"""
Blend AI prediction with ERA5 using land-sea mask.
pred_K : (4H, 4W) AI prediction in Kelvin
era5_K : (H, W) ERA5 input in Kelvin
mask_hr : (4H, 4W) True=land, False=ocean
Returns: (4H, 4W) masked output
"""
# Upsample ERA5 to output resolution for ocean pixels
H, W = era5_K.shape
era5_up = spz(era5_K, pred_K.shape[0] / H, order=3)
era5_up = era5_up[:pred_K.shape[0], :pred_K.shape[1]]
# Blend: land=AI, ocean=ERA5
out = era5_up.copy()
out[mask_hr] = pred_K[mask_hr]
return out
|