pipelines/xgcast_pipeline.py

from prefect import flow, task
from prefect.task_runners import ThreadPoolTaskRunner
from prefect.logging import get_run_logger
from .model_preparation_pipeline import get_consortia, save_to_file
from .soilcast_pipeline import read_input_data, data_preparation, remove_feature_inf_zero_std
import numpy as np
import pandas as pd
import yaml
import random
from xgboost import XGBRegressor
from sklearn.metrics import r2_score, mean_absolute_error


with open('config/params.yml') as file:
    config = yaml.safe_load(file)
with open('config/xgcast_params.yml') as file:
    model_params = yaml.safe_load(file)

data_availability = config['data_availability']

datetime_col = config['datetime_col']
value_col = config['value_col']
datastream_id_col = config['datastream_id_col']
datastream_name_col = config['datastream_name_col']
sensor_type_col = config['sensor_type_col']
ground_offset_col = config['ground_offset_col']

resampling_window = config['resampling_window']
resampling = int(resampling_window.split('h')[0])
num_predictions = config['num_predictions']

train_split = model_params['train_data_split']
val_split = model_params['val_data_split']


### Helper functions

def look_for_value(df, sensor, date):
    tmp = df[(df[datastream_name_col] == sensor)&(df[datetime_col] == date)][value_col]
    if len(tmp) == 1:
        return tmp.values[0]
    else:
        return None


### Pipeline nodes

@task(task_run_name="feature_selection")
def feature_selection(df, features_to_be_removed, features_to_be_added = {"irrigation", "precipitation"}):
    feature_cols = list(set(df.columns) - features_to_be_removed)
    corr_df = df[feature_cols].corr()[value_col]
    feature_cols = set(corr_df[(corr_df >= 0.2)|(corr_df <= -0.2)].sort_values().index)
    feature_cols = feature_cols | features_to_be_added
    return list(feature_cols)


@task(task_run_name="split_data")
def split_data(df: pd.DataFrame, train_data_split, val_data_split, datetime_col, id_col, feature_cols, target_col, random_seed=None):
    """Prepares the data for input to an LSTM model"""
    logger = get_run_logger()

    if not random_seed is None:
        random.seed(random_seed)

    df_t = df.sort_values(datetime_col).copy()

    n = len(df_t)
    df_t['trainvaltest_choice'] = 'test'
    df_t.loc[df_t.index[0] : df_t.index[int(train_data_split * n)-1], 'trainvaltest_choice'] = 'train'
    if train_data_split < 1:
        df_t.loc[df_t.index[int(train_data_split * n)] : df_t.index[int((train_data_split + val_data_split) * n)], 'trainvaltest_choice'] = 'val'

    # windowed and split data
    ids = df_t[id_col].drop_duplicates().values

    X_train = []
    Y_train = []
    X_val = []
    Y_val = []
    X_test = []
    Y_test = []

    for i in range(len(ids)):
        id = ids[i]
        logger.info(f'Working with id {id} ({i / len(ids) * 100}%)')
        df_wk = df_t[df_t[id_col] == id].drop([id_col], axis=1).sort_values(datetime_col).reset_index(drop=True)

        current_index = df_wk.index.values[0]
        maximum_index = df_wk.index.values[-num_predictions]

        while current_index <= maximum_index - 1:

            current_date = df_wk.loc[current_index][datetime_col]
            next_date = df_wk.loc[current_index+1][datetime_col]

            if int((next_date - current_date).days * 24 + (next_date - current_date).seconds / 3600) == resampling:
                choice = df_wk.loc[current_index]['trainvaltest_choice']
                if choice == 'train':
                    X_train.append(df_wk.loc[[current_index]][feature_cols])
                    y_candidate = df_wk.loc[current_index+1: current_index+num_predictions][[target_col]].reset_index(drop=True)
                    y_candidate = y_candidate.T
                    y_candidate.index = [current_index]
                    Y_train.append(y_candidate.rename(columns={i: f'{value_col}_{i}' for i in range(num_predictions)}))
                elif choice == 'val':
                    X_val.append(df_wk.loc[[current_index]][feature_cols])
                    y_candidate = df_wk.loc[current_index+1: current_index+num_predictions][[target_col]].reset_index(drop=True)
                    y_candidate = y_candidate.T
                    y_candidate.index = [current_index]
                    Y_val.append(y_candidate.rename(columns={i: f'{value_col}_{i}' for i in range(num_predictions)}))
                elif choice == 'test':
                    X_test.append(df_wk.loc[[current_index+1]][feature_cols])
                    y_candidate = df_wk.loc[current_index+1: current_index+num_predictions][[target_col]].reset_index(drop=True)
                    y_candidate = y_candidate.T
                    y_candidate.index = [current_index]
                    Y_test.append(y_candidate.rename(columns={i: f'{value_col}_{i}' for i in range(num_predictions)}))

            current_index += 1

    X_train = pd.concat(X_train, axis=0)
    Y_train = pd.concat(Y_train, axis=0)
    X_val = pd.concat(X_val, axis=0)
    Y_val = pd.concat(Y_val, axis=0)
    X_test = pd.concat(X_test, axis=0)
    Y_test = pd.concat(Y_test, axis=0)
    return X_train, Y_train, X_val, Y_val, X_test, Y_test


@task(task_run_name="compute_mean_std")
def compute_mean_std(df, feature_cols):
    """Computes mean and std of an input DataFrame

    Args:
        df: input dataset.
    Returns:
        mean_std: mean and std DataFrame of input dataset.
    """
    mean_df = pd.DataFrame(df[feature_cols].mean(), columns=['mean'])
    std_df = pd.DataFrame(df[feature_cols].std(), columns=['std'])

    mean_std_df = pd.concat([mean_df, std_df], axis=1)
    return mean_std_df


@task(task_run_name="normalise_df")
def normalise_df(df, mean_std_train):
    m_s_df = mean_std_train.copy()
    m_s_df.loc[m_s_df['std'] == 0, 'std'] = 1 # needed to be sure that the variables with 0 std are properly normalised to 0

    df_result = df.copy()
    for col in df_result.columns:
        if col in m_s_df.index:
            df_result[col] = (df_result[col] - m_s_df['mean'].loc[col]) / m_s_df['std'].loc[col]
        elif f'{value_col}_' in col:
            df_result[col] = (df_result[col] - m_s_df['mean'].loc[value_col]) / m_s_df['std'].loc[value_col]

    return df_result


### Flows

@flow(name='xgcast_preparation_pipeline', retries=1, task_runner=ThreadPoolTaskRunner())
def xgcast_preparation_pipeline(override_consortia=None, override_data_availability=None) -> list[str]:
    logger = get_run_logger()
    logger.info(f'Starting XGCast pipeline!')

    if not override_consortia:
        consortia = get_consortia()
    else:
        consortia = override_consortia

    for consortium_name in consortia:
        # Read data
        if override_data_availability:
            data_availability_consortium = override_data_availability
        else:
            data_availability_consortium = data_availability[consortium_name]

        full_table, crop_type_data, soil_type_data = read_input_data.submit(consortium_name, **data_availability_consortium).result()
        tensiometers, elmed, feature_cols = data_preparation.submit(full_table, crop_type_data, soil_type_data).result()

        # feature_selection(tensiometers, features_to_be_removed)
        X_train, Y_train, X_val, Y_val, X_test, Y_test = split_data.submit(tensiometers, train_split, val_split, datetime_col, datastream_name_col, feature_cols, value_col, random_seed=None).result()

        mean_std_df = compute_mean_std.submit(X_train, feature_cols).result()
        mean_std_df, feature_cols = remove_feature_inf_zero_std.submit(mean_std_df, feature_cols).result()

        X_train = normalise_df.submit(X_train[feature_cols], mean_std_df).result()
        Y_train = normalise_df.submit(Y_train, mean_std_df).result()
        X_val = normalise_df.submit(X_val[feature_cols], mean_std_df).result()
        Y_val = normalise_df.submit(Y_val, mean_std_df).result()
        X_test = normalise_df.submit(X_test[feature_cols], mean_std_df).result()
        Y_test = normalise_df.submit(Y_test, mean_std_df).result()

        # Saving

        save_to_file.submit(tensiometers, f'data//05_xgcast_input//tensiometers_{consortium_name}.parquet').result()
        save_to_file.submit(mean_std_df, f'data//05_xgcast_input//mean_std_{consortium_name}.parquet').result()
        save_to_file.submit(feature_cols, f'data//05_xgcast_input//feature_cols_{consortium_name}.pickle').result()

        save_to_file.submit(X_train, f'data//05_xgcast_input//X_train_{consortium_name}.parquet').result()
        save_to_file.submit(Y_train, f'data//05_xgcast_input//Y_train_{consortium_name}.parquet').result()
        save_to_file.submit(X_val, f'data//05_xgcast_input//X_val_{consortium_name}.parquet').result()
        save_to_file.submit(Y_val, f'data//05_xgcast_input//Y_val_{consortium_name}.parquet').result()
        save_to_file.submit(X_test, f'data//05_xgcast_input//X_test_{consortium_name}.parquet').result()
        save_to_file.submit(Y_test, f'data//05_xgcast_input//Y_test_{consortium_name}.parquet').result()


@flow(name='xgcast_model_pipeline', retries=1, task_runner=ThreadPoolTaskRunner())
def xgcast_model_pipeline(override_consortia=None) -> list[str]:
    logger = get_run_logger()
    logger.info(f'Starting XGCast model pipeline!')

    if not override_consortia:
        consortia = get_consortia()
    else:
        consortia = override_consortia

    for consortium_name in consortia:

        # # Read data
        X_train = pd.read_parquet(f'data//05_xgcast_input//X_train_{consortium_name}.parquet')
        Y_train = pd.read_parquet(f'data//05_xgcast_input//Y_train_{consortium_name}.parquet')
        X_val = pd.read_parquet(f'data//05_xgcast_input//X_val_{consortium_name}.parquet')
        Y_val = pd.read_parquet(f'data//05_xgcast_input//Y_val_{consortium_name}.parquet')
        X_test = pd.read_parquet(f'data//05_xgcast_input//X_test_{consortium_name}.parquet')
        Y_test = pd.read_parquet(f'data//05_xgcast_input//Y_test_{consortium_name}.parquet')

        xgb_params = model_params['xgb_params']
        bst = XGBRegressor(objective='reg:absoluteerror', **xgb_params)

        # fit model
        bst.fit(pd.concat([X_train, X_val], axis=0), pd.concat([Y_train, Y_val], axis=0))

        Y_pred = bst.predict(X_test)
        Y_pred = pd.DataFrame(Y_pred).rename(columns={col: f'{value_col}_{col}' for col in range(num_predictions)})

        mae = np.sqrt(mean_absolute_error(Y_test, Y_pred))
        mse = np.mean((Y_test - Y_pred) ** 2)
        r2 = r2_score(Y_test, Y_pred)

        logger.info(f'For consortium {consortium_name} -> MAE: {mae:.3f}')
        logger.info(f'For consortium {consortium_name} -> MSE: {mse:.3f}')
        logger.info(f'For consortium {consortium_name} -> R²: {r2:.3f}')

        bst.save_model(f'data//06_xgcast_output//xgcast_{consortium_name}.json')


if __name__ == "__main__":
    xgcast_preparation_pipeline()
    xgcast_model_pipeline()