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()