indices.py

""" Add Sentinel-1 and Sentinel-2 indices to a DataFrame. """
import numpy as np
from variables import s2_indices

def add_s2_indices(df, L=0.5, y=0.126, a=0.3, b=0.5):                        # pylint: disable=invalid-name
    """ Add Sentinel-2 indices to a DataFrame.

    Args:
        df (pandas.DataFrame): DataFrame to add indices to.
        L (float): Soil adjustment factor for SAVI. Default is 0.5.
        y (float): Atmosphere resistance factor for ARVI. Default is 0.126.
        a (float): Slope parameter for PVI. Default is 0.3.
        b (float): Intercept parameter for PVI. Default is 0.5.

    Returns:
        pandas.DataFrame: DataFrame with added indices.
    """

    X = df.copy()

    if X.B4.values[0] is not None:

        # Vegetation Indices
        X["ARI"] = (1 / X.B3) - (1 / X.B5)
        X["ARVI"] = (X.B8 - X.B4 - y * (X.B4 - X.B2)) / (X.B8 + X.B4 - y * (X.B4 - X.B2))
        X["AVI"] = np.sign((X.B8 * (1 - X.B4) * (X.B8 - X.B4))) * np.abs((X.B8 * (1 - X.B4) * (X.B8 - X.B4))) ** (1/3)
        X["BCC"] = X.B2 / (X.B4 + X.B3 + X.B2)
        X["CCCI"] = ((X.B8 - X.B5) / (X.B8 + X.B5)) / ((X.B8 - X.B4) / (X.B8 + X.B4))
        X["CVI"] = (X.B8 * X.B4) * (X.B3 ** 2)
        X["EVI"] = 2.5 * (X.B8 - X.B4) / (X.B8 + 6 * X.B4 - 7.5 * X.B2 + 1)
        X["ExG"] = 2 * X.B3 - X.B4 - X.B2
        X["GNDVI"] = (X.B8 - X.B3) / (X.B8 + X.B3)
        X["IKAW"] = (X.B8 - X.B2) / (X.B8 + X.B2)
        X["MCARI"] = ((X.B5 - X.B4) - 0.2 * (X.B5 - X.B3)) * (X.B5 / X.B4)
        X["MSAVI"] = 0.5 * (2 * X.B8 + 1 - (((2 * X.B8 + 1) ** 2) - 8 * (X.B8 - X.B4)) ** 0.5)
        X["MTVI2"] = (1.5 * (1.2 * (X.B8 - X.B3) - 2.5 * (X.B4 - X.B3))) / ((((2 * X.B8 + 1)**2) - (6 * X.B8 - 5 * (X.B4**0.5))-0.5) ** 0.5)
        X["NDDI"] = (((X.B8 - X.B4) / (X.B8+X.B4)) - ((X.B3 - X.B8) / (X.B3 + X.B8))) / (((X.B8 - X.B4) / (X.B8 + X.B4)) + ((X.B3 - X.B8) / (X.B3 + X.B8)))
        X["NDMI"] = (X.B8 - X.B11) / (X.B8 + X.B11)
        X["NDVI"] = (X.B8 - X.B4) / (X.B8 + X.B4) if X.B8 is not None else 0
        X["NDYI"] = (X.B3 - X.B2) / (X.B3 + X.B2)
        X["PSRI"] = (X.B4 - X.B2) / X.B6
        X["PVI"] = (X.B8 - a * X.B4 - b) / ((1 + a * 2) ** (1/2))
        X["RCC"] = X.B4 / (X.B4 + X.B3 + X.B2)
        X["RENDVI"] = (X.B6 - X.B5) / (X.B6 + X.B5)
        X["S2REP"] = 705 + 35 * ((((X.B7 + X.B4) / 2) - X.B5) / (X.B6 - X.B5))
        X["SAVI"] = (1 + L) * (X.B8 - X.B4) / (X.B8 + X.B4 + L)
        X["SI"] = ((1 - X.B2) * (1 - X.B3) * (1 - X.B4)) ** (1/3)
        X["SIPI"] = (X.B8 - X.B1) / (X.B8 - X.B4)
        X["TCARI"] = 3 * ((X.B5 - X.B4) - 0.2 * (X.B5 - X.B3) * (X.B5 / X.B4))
        X["TNDVI"] = ((X.B8 - X.B4) / (X.B8 + X.B4) + 0.5) ** (1/2)
        X["TTVI"] = 0.5 * ((865 - 740) * (X.B7 - X.B6) - (X.B8A - X.B6) * (783-740))
        X["TVI"] = (((X.B8 - X.B4) / (X.B8 + X.B4)) + 0.5) ** (1/2)
        X["VARI"] = (X.B3 - X.B4) / (X.B3 + X.B4 - X.B2)
        X["VARI700"] = (X.B5 - 1.7 * X.B4 + 0.7 * X.B2) / (X.B5 + 1.3 * X.B4 - 1.3*X.B2)

        # Water Indices
        X["AWEInsh"] = 4 * (X.B3 - X.B11) - 0.25 * X.B8 + 2.75 * X.B12
        X["AWEIsh"] = X.B2 + 2.5 * X.B3 - 1.5 * (X.B8 + X.B11) - 0.25 * X.B12
        X["MNDWI"] = (X.B3 - X.B11) / (X.B3 + X.B11)
        X["MuWIR"] = -4 * ((X.B2 - X.B3) / (X.B2 + X.B3)) + 2 * ((X.B3 - X.B8) / (X.B3 + X.B8)) + 2 * ((X.B3 - X.B12) / (X.B3 + X.B12)) - ((X.B3 - X.B11) / (X.B3 + X.B11))
        X["NDWI"] = (X.B3 - X.B8) / (X.B3 + X.B8)
        X["S2WI"] = (X.B5 - X.B12) / (X.B5 + X.B12)
        X["WI2015"] = 1.7204 + 171 * X.B3 + 3 * X.B4 - 70*X.B8 - 45*X.B11 - 71*X.B12

        # Burn / Fire Indices
        X["BAI"] = 1 / ((0.1 - X.B4) ** 2 + (0.06 - X.B8) ** 2)
        X["BAIM"] = 1 / ((0.05 - X.B8) ** 2) + ((0.2 - X.B12) ** 2) # adapted to MODIS
        X["MIRBI"] = 10 * X.B12 -9.8 * X.B11 + 2
        X["NBR"] = (X.B8 - X.B12) / (X.B8 + X.B12)
        X["NBRI"] = (X.B8 - X.B12) / (X.B8 + X.B12) # Not found
        X["NBRplus"] = (X.B12 - X.B8A - X.B3 - X.B2) / (X.B12 + X.B8A + X.B3 + X.B2)
        X["NBRSWIR"] = (X.B12 - X.B11 - 0.02) / (X.B12 + X.B11 + 0.1)
        X["NDSWIR"] = (X.B8 - X.B11)/(X.B8 + X.B11)

        # Snow Indices
        X["NBSIMS"] = 0.36 * (X.B3 + X.B4 + X.B8) - (((X.B2 + X.B12) / X.B3) + X.B11)
        X["NDSI"] = (X.B3 - X.B11) / (X.B3 + X.B11)

        # Urban Indices
        X["NBAI"] = (X.B12 - X.B11 / X.B3) / (X.B12 + X.B11 / X.B3)
        X["NDBI"] = (X.B11 - X.B8) / (X.B11 + X.B8)
        X["UI"] = (X.B12 - X.B8) / (X.B12 + X.B8)

        # Soil Indices
        X["BI"] = ((X.B11 + X.B4) - (X.B8 + X.B2)) / ((X.B11 + X.B4) + (X.B8 + X.B2))
        X["BI2"] = np.sqrt((X.B4 ** 2 + X.B3 ** 2 + X.B8 ** 2) / 3)
        X["BSI"] = ((X.B12 + X.B4) - (X.B8 + X.B2)) / ((X.B12 + X.B4) + (X.B8 + X.B2))
        X["SATVI"] = ((X.B11 - X.B4) / (X.B11 + X.B4 + L)) * (1 + L) - (X.B12 / 2)
    else:
        for index in s2_indices:
            X[index] = 0.0

    return X

def add_s1_indices(df):
    """ Add Sentinel-1 indices to a DataFrame.

    Args:
        df (pandas.DataFrame): DataFrame to add indices to.

    Returns:
        pandas.DataFrame: DataFrame with added indices.
    """

    X = df.copy()

    # Radar Indices
    X["DPDD"] = (X.VV + X.VH) / 2 ** (1/2)
    X["DpRVIVV"] = (4 * X.VH) / (X.VV + X.VH)
    X["NDPolI"] = (X.VV - X.VH) / (X.VV + X.VH)
    X["VDDPI"] = (X.VV + X.VH) / X.VV
    X["VHVVD"] = X.VH - X.VV
    X["VHVVP"] = X.VH * X.VV
    X["VHVVR"] = X.VH / X.VV
    X["VVVHD"] = X.VV - X.VH
    X["VVVHR"] = X.VV / X.VH
    X["VVVHS"] = X.VV + X.VH

    return X