Raster Band Math and Index Calculation

# requirements.txt
rasterio==1.3.10
numpy==1.26.4
xarray==2024.2.0
rioxarray==0.15.3
dask[array]==2024.4.1
scikit-learn==1.4.2
pyproj==3.6.1

pip install rasterio==1.3.10 numpy==1.26.4 xarray==2024.2.0 \
    rioxarray==0.15.3 "dask[array]==2024.4.1" \
    scikit-learn==1.4.2 pyproj==3.6.1

sudo apt-get install gdal-bin libgdal-dev proj-bin libproj-dev

import rasterio, numpy, pyproj
print(rasterio.__version__)   # expect 1.3.x
print(numpy.__version__)      # expect 1.26.x
print(pyproj.__version__)     # expect 3.6.x

import rasterio
from rasterio.enums import Resampling

def validate_and_align(band_paths: list[str], reference_idx: int = 0) -> list[str]:
    """
    Confirm all bands share CRS, resolution, and transform.
    Returns paths unchanged if consistent; raises on mismatch.
    """
    with rasterio.open(band_paths[reference_idx]) as ref:
        ref_crs = ref.crs
        ref_transform = ref.transform
        ref_shape = (ref.height, ref.width)

    for path in band_paths[1:]:
        with rasterio.open(path) as src:
            if src.crs != ref_crs:
                raise ValueError(
                    f"CRS mismatch: {path} has {src.crs}, expected {ref_crs}. "
                    "Reproject with rasterio.warp.reproject before band math."
                )
            if (src.height, src.width) != ref_shape:
                raise ValueError(
                    f"Grid mismatch: {path} is {src.height}x{src.width}, "
                    f"expected {ref_shape[0]}x{ref_shape[1]}."
                )
            if not all(
                abs(src.transform[i] - ref_transform[i]) < 1e-8 for i in range(6)
            ):
                raise ValueError(f"Affine transform mismatch in {path}.")
    return band_paths

# Usage
paths = ["red_band.tif", "nir_band.tif"]
validate_and_align(paths)
assert len(paths) == 2

import numpy as np
import rasterio

SENTINEL2_SCALE = 10_000.0
SENTINEL2_NO_DATA = 0

def load_reflectance_band(
    path: str,
    scale: float = SENTINEL2_SCALE,
    no_data_val: int = SENTINEL2_NO_DATA,
) -> np.ndarray:
    """
    Load a single-band GeoTIFF and convert DN to [0, 1] reflectance.
    No-data pixels are returned as NaN.
    """
    with rasterio.open(path) as src:
        data = src.read(1).astype(np.float32)
    # Mask no-data before scaling
    data = np.where(data == no_data_val, np.nan, data)
    reflectance = data / scale
    # Clip to physically valid range — values outside [0, 1] are sensor artefacts
    reflectance = np.clip(reflectance, 0.0, 1.0)
    return reflectance

red = load_reflectance_band("B04.tif")
nir = load_reflectance_band("B08.tif")

assert red.dtype == np.float32
assert np.nanmax(red) <= 1.0, "Red band contains values > 1.0 after scaling"

import numpy as np
import rasterio

# Sentinel-2 Scene Classification Layer (SCL) values to exclude:
# 0=no data, 1=saturated/defective, 3=cloud shadows,
# 8=cloud medium prob, 9=cloud high prob, 10=thin cirrus, 11=snow
INVALID_SCL_CLASSES = {0, 1, 3, 8, 9, 10, 11}

def build_valid_mask(scl_path: str) -> np.ndarray:
    """
    Return a boolean array: True where pixels are valid for computation.
    """
    with rasterio.open(scl_path) as src:
        scl = src.read(1)
    invalid = np.zeros_like(scl, dtype=bool)
    for cls in INVALID_SCL_CLASSES:
        invalid |= (scl == cls)
    return ~invalid   # True = valid

def apply_mask(array: np.ndarray, valid_mask: np.ndarray) -> np.ndarray:
    """Set invalid pixels to NaN in-place."""
    out = array.copy()
    out[~valid_mask] = np.nan
    return out

valid = build_valid_mask("SCL.tif")
red_masked = apply_mask(red, valid)
nir_masked = apply_mask(nir, valid)

valid_fraction = valid.mean()
assert valid_fraction > 0.05, f"Too few valid pixels ({valid_fraction:.1%}) — check scene cloud cover."
print(f"Valid pixel fraction: {valid_fraction:.1%}")

import numpy as np

EPSILON = 1e-10  # prevents division by zero on uniform surfaces

def ndvi(red: np.ndarray, nir: np.ndarray) -> np.ndarray:
    """
    Normalized Difference Vegetation Index.
    Output range: [-1, 1]. High values indicate dense green vegetation.
    """
    return (nir - red) / (nir + red + EPSILON)

def evi(blue: np.ndarray, red: np.ndarray, nir: np.ndarray) -> np.ndarray:
    """
    Enhanced Vegetation Index — less prone to atmospheric and soil background noise than NDVI.
    Standard coefficients: G=2.5, C1=6, C2=7.5, L=1.
    """
    return 2.5 * (nir - red) / (nir + 6 * red - 7.5 * blue + 1 + EPSILON)

def ndwi(green: np.ndarray, nir: np.ndarray) -> np.ndarray:
    """
    Normalized Difference Water Index (McFeeters 1996).
    Positive values indicate open water surfaces.
    """
    return (green - nir) / (green + nir + EPSILON)

# Compute all indices in a single pass
ndvi_arr  = ndvi(red_masked, nir_masked)
evi_arr   = evi(blue_masked, red_masked, nir_masked)
ndwi_arr  = ndwi(green_masked, nir_masked)

# Sanity checks on output ranges
assert np.nanmin(ndvi_arr) >= -1.0 and np.nanmax(ndvi_arr) <= 1.0, "NDVI out of range"
assert np.nanmin(evi_arr)  >= -1.0 and np.nanmax(evi_arr)  <= 1.0, "EVI out of range"

import rasterio
from rasterio.windows import Window
import numpy as np

CHUNK = 1024  # pixels per side; adjust to fit available RAM

def compute_ndvi_chunked(
    red_path: str,
    nir_path: str,
    output_path: str,
    chunk_size: int = CHUNK,
) -> None:
    """
    Compute NDVI tile-by-tile and write to a compressed GeoTIFF.
    No full-scene array is held in memory.
    """
    with rasterio.open(red_path) as red_src, \
         rasterio.open(nir_path) as nir_src:

        profile = red_src.profile.copy()
        profile.update(dtype="float32", count=1, compress="zstd", nodata=np.nan)

        with rasterio.open(output_path, "w", **profile) as dst:
            for row_off in range(0, red_src.height, chunk_size):
                for col_off in range(0, red_src.width, chunk_size):
                    win = Window(
                        col_off, row_off,
                        min(chunk_size, red_src.width  - col_off),
                        min(chunk_size, red_src.height - row_off),
                    )
                    red_chunk = red_src.read(1, window=win).astype(np.float32) / 10_000.0
                    nir_chunk = nir_src.read(1, window=win).astype(np.float32) / 10_000.0
                    ndvi_chunk = (nir_chunk - red_chunk) / (nir_chunk + red_chunk + EPSILON)
                    dst.write(ndvi_chunk, 1, window=win)

compute_ndvi_chunked("B04.tif", "B08.tif", "ndvi_output.tif")
print("Chunked NDVI complete.")

import numpy as np
import pickle
from sklearn.preprocessing import RobustScaler

def prepare_index_for_training(
    index_array: np.ndarray,
    clip_min: float = -1.0,
    clip_max: float = 1.0,
    scaler_path: str | None = None,
) -> tuple[np.ndarray, RobustScaler]:
    """
    Clip outliers, flatten valid pixels, and fit a RobustScaler.
    Returns the scaled flat array and the fitted scaler.
    """
    clipped = np.clip(index_array, clip_min, clip_max)
    flat = clipped[~np.isnan(clipped)].reshape(-1, 1)
    scaler = RobustScaler()
    scaled = scaler.fit_transform(flat)
    if scaler_path:
        with open(scaler_path, "wb") as f:
            pickle.dump(scaler, f)
    return scaled.ravel(), scaler

ndvi_scaled, scaler = prepare_index_for_training(
    ndvi_arr, scaler_path="ndvi_scaler.pkl"
)
assert abs(np.median(ndvi_scaled)) < 0.1, "Median should be near zero after RobustScaler"

import numpy as np
import rasterio

def verify_index_output(index_path: str, expected_min: float, expected_max: float) -> None:
    """Confirm output file has valid dtype, range, and non-trivial valid pixel count."""
    with rasterio.open(index_path) as src:
        data = src.read(1).astype(np.float32)
        profile = src.profile

    assert profile["dtype"] == "float32", "Output must be float32"
    valid = data[~np.isnan(data)]
    assert len(valid) > 0, "No valid (non-NaN) pixels in output"
    assert valid.min() >= expected_min, f"Values below {expected_min}"
    assert valid.max() <= expected_max, f"Values above {expected_max}"
    print(
        f"OK — {len(valid):,} valid pixels | "
        f"mean={valid.mean():.4f} | std={valid.std():.4f}"
    )

verify_index_output("ndvi_output.tif", expected_min=-1.0, expected_max=1.0)

from concurrent.futures import ProcessPoolExecutor

tile_paths = [("B04_tile1.tif", "B08_tile1.tif", "ndvi_tile1.tif"),
              ("B04_tile2.tif", "B08_tile2.tif", "ndvi_tile2.tif")]

with ProcessPoolExecutor(max_workers=4) as pool:
    pool.map(lambda args: compute_ndvi_chunked(*args), tile_paths)