pipelines/aquacrop_preparation_pipeline.py

from prefect import flow, task
from prefect.artifacts import create_table_artifact
from prefect.task_runners import ThreadPoolTaskRunner
import datetime
import json
import os
import pandas as pd
import pickle
from pipelines.preprocessing_pipeline import save_to_file
from pipelines.model_preparation_pipeline import get_consortia, get_associated_sensors, get_sensors_meteo, add_irrigation_forecast
from prefect.logging import get_run_logger
import yaml


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

with open('config/fieldsensor_irrigator_mapping_anonym.yaml') as file:
    fieldsensor_irrigator_mapping = yaml.safe_load(file)


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']

spatial_agg_method = config['spatial_agg_method']
field_agg = config['field_agg']

resampling_window = config['resampling_window']
days_weather_forecast = config['days_weather_forecast']

meteo_method = config['meteo_method']


with open('config/aquacrop_params.yml') as file:
    config_aquacrop = yaml.safe_load(file)

sensors_forecasted = config_aquacrop['sensor_forecasted']
reference_date = config_aquacrop['reference_date']

crop_per_consortium = config_aquacrop['crop_per_consortium']
planting_dates_per_consortium = config_aquacrop['planting_dates_per_consortium']
harvest_dates_per_consortium = config_aquacrop['harvest_dates_per_consortium']

irrigation_smt_per_consortium = config_aquacrop['irrigation_smt_per_consortium']
irrigation_maxirr_per_consortium = config_aquacrop['irrigation_maxirr_per_consortium']
strategy = config_aquacrop['strategy']


### Pipeline nodes

@task(task_run_name="read_input_data_{consortium_name}")
def read_input_data(consortium_name):
    field_sensor_data = pd.read_parquet(f'data//03_primary//field_sensor_data_{consortium_name}.parquet')
    irrigation_data = pd.read_parquet(f'data//03_primary//irrigation_data_{consortium_name}.parquet')
    if consortium_name != 'consortium2':
        weather_data = pd.read_parquet(f'data//03_primary//weather_data_{consortium_name}.parquet')
    else:
        weather_data = None
    historical_weather_data = pd.read_parquet(f'data//03_primary//historical_weather_data_{consortium_name}.parquet')
    forecasted_weather_data = pd.read_parquet(f'data//03_primary//forecasted_weather_data_{consortium_name}.parquet')
    with open(f'data//03_primary//crop_type_data_{consortium_name}.pickle', 'rb') as handle:
        crop_type_data = pickle.load(handle)
    if consortium_name != 'consortium2':
        soil_type_data = pd.read_parquet(f'data//03_primary//soil_type_data_{consortium_name}.parquet')
    else:
        soil_type_data = None
    locations_ids = pd.read_parquet(f'data//03_primary//locations_ids_{consortium_name}.parquet')

    info_artifact = [
        {'variable': 'Length of field sensors DataFrame', 'amount': len(field_sensor_data)},
        {'variable': 'Length of irrigation data DataFrame', 'amount': len(irrigation_data)},
        {'variable': 'Length of weather data DataFrame', 'amount': len(weather_data) if weather_data is not None else None},
        {'variable': 'Length of historical weather data DataFrame', 'amount': len(historical_weather_data)},
        {'variable': 'Length of non-NaN field sensors DataFrame', 'amount': len(field_sensor_data.dropna())},
        {'variable': 'Length of non-NaN irrigation data DataFrame', 'amount': len(irrigation_data.dropna())},
        {'variable': 'Length of non-NaN weather data DataFrame', 'amount': len(weather_data.dropna()) if weather_data is not None else None},
        {'variable': 'Length of non-NaN historical weather data DataFrame',
         'amount': len(historical_weather_data.dropna())},
    ]
    create_table_artifact(
        key=f"data-{consortium_name}-info",
        table=info_artifact,
        description="# Data info!"
    )

    return field_sensor_data, irrigation_data, weather_data, historical_weather_data, forecasted_weather_data, crop_type_data, soil_type_data, locations_ids


@task(task_run_name='merge_all_data_{consortium_name}')
def merge_all_data(consortium_name, field_sensor_data, irrigation_data, weather_data, historical_weather_data,
                   forecasted_weather_data):
    associated_irrigators = []
    for year in fieldsensor_irrigator_mapping[consortium_name]:
        tmp = pd.DataFrame(fieldsensor_irrigator_mapping[consortium_name][year].items(),
                           columns=[datastream_name_col, 'irrigator'])
        tmp['year'] = year
        associated_irrigators.append(tmp)
    associated_irrigators = pd.concat(associated_irrigators)

    field_sensor_data['year'] = field_sensor_data.index.year

    valid_names = field_sensor_data["datastream_name"].unique()
    historical_weather_data = historical_weather_data[
        historical_weather_data["datastream_name"].isin(valid_names)
    ]

    full_table = field_sensor_data.drop(columns=[datastream_id_col]).reset_index().merge(
        associated_irrigators,
        on=[datastream_name_col, 'year'],
        how='left'
    ).drop(columns=['year']).merge(
        irrigation_data.drop(columns=[datastream_id_col]).reset_index().rename(
            columns={datastream_name_col: 'irrigator', value_col: 'irrigation'}),
        on=['irrigator', datetime_col],
        how='left'
    ).merge(
        forecasted_weather_data,
        on=[datetime_col, datastream_name_col],
        how='left'
    ).merge(
        historical_weather_data.reset_index(),
        on=[datetime_col, datastream_name_col],
        how='right'
    )

    if weather_data is not None:
        full_table = full_table.merge(
            weather_data,
            on=[datetime_col, datastream_name_col],
            how='left'
        )

    full_table['irrigation'] = full_table['irrigation'].fillna(0)

    sensors = pd.DataFrame(full_table[datastream_name_col].drop_duplicates().reset_index(drop=True))
    sensors = get_associated_sensors(sensors)

    full_table = full_table.merge(
        sensors[[datastream_name_col, 'associated_sensors']],
        on=[datastream_name_col],
        how='left'
    )

    #if consortium_name == 'consortium1':
    #    full_table = full_table[full_table[datetime_col].dt.year >= 2024]

    return full_table


@task(task_run_name='add_irrigation_forecast')
def add_irrigation_forecast(full_table):
    for n in range(days_weather_forecast * 24 // int(resampling_window.split('h')[0])):
        full_table[f'forecasted_irrigation_next_{n+1}period'] = full_table['irrigation'].shift(-n)
        full_table[f'forecasted_irrigation_next_{n+1}period'] = full_table[f'forecasted_irrigation_next_{n+1}period'].fillna(0)

    return full_table


@task(task_run_name='aquacrop_preparation_{consortium_name}_{sensor_forecasted}')
def aquacrop_preparation(consortium_name, sensor_forecasted, full_table, strategy):
    full_table = full_table.sort_values(
        by=["result_time", "datastream_name"],
        ignore_index=True
    )

    # Soil Type
    if consortium_name != 'consortium2':
        soil_type_data = pd.read_parquet(f'data//03_primary//soil_type_data_{consortium_name}.parquet')
        soil_info = (
            full_table.merge(
                soil_type_data,
                on=datastream_name_col,
                how='left'
            )[[datastream_name_col, 'soil_type']]
            .drop_duplicates()
            .reset_index(drop=True)
        )
        try:
            soil_type = soil_info.loc[soil_info['datastream_name'] == sensor_forecasted, 'soil_type'].iloc[0]
            soil_type = soil_type.title().replace(" ", "")
        except:
            soil_type = 'Loam'
    else:
        soil_type = 'Loam'

    # Crop Type & Planting Date
    crop_type = crop_per_consortium[consortium_name]
    planting_date = planting_dates_per_consortium[consortium_name]
    harvest_date = harvest_dates_per_consortium[consortium_name]

    # Weather Dataset
    full_table_filtered = full_table[(full_table['datastream_name'] == sensor_forecasted)].reset_index(drop=True)
    simulation_forecast = pd.to_datetime(reference_date)
    start_of_year = pd.Timestamp(reference_date.year, 1, 1)

    full_table_filtered = full_table_filtered[
        (full_table_filtered["result_time"] >= start_of_year) &
        (full_table_filtered["result_time"] <= simulation_forecast)
    ].reset_index(drop=True)

    weather_aquacrop_observed = pd.DataFrame({
        "MinTemp": full_table_filtered["air_temperature_min_local_meteo"].fillna(full_table_filtered["temperature_2m_min"]),
        "MaxTemp": full_table_filtered["air_temperature_max_local_meteo"].fillna(full_table_filtered["temperature_2m_max"]),
        "Precipitation": full_table_filtered["precipitation_local_meteo"].fillna(full_table_filtered["precipitation"]),
        "ReferenceET": full_table_filtered["et0_fao_evapotranspiration"],
        "Date": full_table_filtered["result_time"]
    })
    weather_aquacrop_observed = weather_aquacrop_observed.drop_duplicates()

    forecasted_columns = [col for col in full_table_filtered.columns if col.startswith("forecasted_")]
    periods = sorted(list(set([col.split("_next_")[-1] for col in forecasted_columns])))
    forecasted_data = []

    for period in periods:
        row = {}
        row["Date"] = pd.to_datetime(simulation_forecast) + pd.Timedelta(days=int(period.split("period")[0]))

        # MinTemp
        col_min = f"forecasted_temperature_2m_min_next_{period}"
        if col_min in full_table_filtered.columns:
            row["MinTemp"] = full_table_filtered.loc[full_table_filtered["result_time"] == simulation_forecast, col_min].values[0]

        # MaxTemp
        col_max = f"forecasted_temperature_2m_max_next_{period}"
        if col_max in full_table_filtered.columns:
            row["MaxTemp"] = full_table_filtered.loc[full_table_filtered["result_time"] == simulation_forecast, col_max].values[0]

        # Precipitation
        col_prec = f"forecasted_precipitation_next_{period}"
        if col_prec in full_table_filtered.columns:
            row["Precipitation"] = full_table_filtered.loc[full_table_filtered["result_time"] == simulation_forecast, col_prec].values[0]

        # ReferenceET
        col_et = f"forecasted_et0_fao_evapotranspiration_next_{period}"
        if col_et in full_table_filtered.columns:
            row["ReferenceET"] = full_table_filtered.loc[full_table_filtered["result_time"] == simulation_forecast, col_et].values[0]

        forecasted_data.append(row)

    # Placeholder for the last day of the simulation (it will not be used by AquaCrop)
    last_row = forecasted_data[-1].copy()
    last_row["Date"] = last_row["Date"] + pd.Timedelta(days=1)
    forecasted_data.append(last_row)

    weather_aquacrop_forecasted = pd.DataFrame(forecasted_data)
    weather_aquacrop = pd.concat([weather_aquacrop_observed, weather_aquacrop_forecasted], ignore_index=True)
    weather_aquacrop = weather_aquacrop.sort_values("Date").reset_index(drop=True)

    # Irrigation Dataset
    if strategy == 'real' or strategy == 'hybrid':
        irrigation_aquacrop_observed = pd.DataFrame({
            "Date": full_table_filtered["result_time"],
            "Depth": full_table_filtered["irrigation"].fillna(0.0)
        })
        irrigation_aquacrop_observed = irrigation_aquacrop_observed.drop_duplicates()

        if strategy == 'real':
            forecasted_irrigation_columns = [
                col for col in full_table_filtered.columns
                if col.startswith("forecasted_irrigation_next_")
            ]

            irrigation_periods = sorted(list(set([
                col.split("forecasted_irrigation_next_")[-1]
                for col in forecasted_irrigation_columns
            ])))

            forecasted_irrigation_data = []

            for period in irrigation_periods:
                row = {}
                row["Date"] = pd.to_datetime(simulation_forecast) + pd.Timedelta(days=int(period.split("period")[0]))

                col_irrig = f"forecasted_irrigation_next_{period}"

                if col_irrig in full_table_filtered.columns:
                    row["Depth"] = full_table_filtered.loc[
                        full_table_filtered["result_time"] == simulation_forecast,
                        col_irrig
                    ].fillna(0.0).values[0]
                else:
                    row["Depth"] = 0.0

                forecasted_irrigation_data.append(row)

            irrigation_aquacrop_forecasted = pd.DataFrame(forecasted_irrigation_data)

            irrigation_aquacrop = pd.concat(
                [irrigation_aquacrop_observed, irrigation_aquacrop_forecasted],
                ignore_index=True
            )
        else:  # hybrid
            irrigation_aquacrop = irrigation_aquacrop_observed.copy()

        irrigation_aquacrop = irrigation_aquacrop.sort_values("Date").reset_index(drop=True)
        irrigation_aquacrop = irrigation_aquacrop[irrigation_aquacrop["Depth"] != 0.0].reset_index(drop=True)
    else:  # aquacrop
        irrigation_aquacrop = None

    # Simulation Dates
    month, day = map(int, planting_date.split('/'))
    planting_date_reference = datetime.date(reference_date.year, month, day)

    one_month_before = planting_date_reference - pd.DateOffset(months=1)

    filtered = weather_aquacrop[
        (weather_aquacrop['Precipitation'] > 0) &
        (weather_aquacrop['Date'] >= pd.Timestamp(one_month_before)) &
        (weather_aquacrop['Date'] < pd.Timestamp(planting_date_reference))
        ]

    if not filtered.empty:
        filtered_over_10 = filtered[filtered['Precipitation'] > 10]
        if not filtered_over_10.empty:
            last_rainy_date_raw = filtered_over_10['Date'].iloc[-1]
        else:
            max_precip = filtered['Precipitation'].max()
            last_rainy_date_raw = filtered[filtered['Precipitation'] == max_precip]['Date'].iloc[-1]
        sim_start_date = (last_rainy_date_raw + datetime.timedelta(days=1)).strftime('%Y/%m/%d')
    else:
        sim_start_date = planting_date_reference.strftime('%Y/%m/%d')

    sim_end_date_raw = weather_aquacrop['Date'].iloc[-1]
    sim_end_date = pd.to_datetime(sim_end_date_raw).strftime('%Y/%m/%d')

    # Irrigation Thresholds & Maximum Irrigation
    SMT = irrigation_smt_per_consortium[consortium_name]
    MaxIrr = irrigation_maxirr_per_consortium[consortium_name]

    return weather_aquacrop, irrigation_aquacrop, soil_type, crop_type, planting_date, harvest_date, reference_date, sim_start_date, sim_end_date, SMT, MaxIrr


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

    consortia = get_consortia()

    for sensor_forecasted in sensors_forecasted:

        for consortium_name in consortia:
            field_sensor_data, irrigation_data, weather_data, historical_weather_data, forecasted_weather_data, crop_type_data, soil_type_data, location_ids = read_input_data(consortium_name)

            sensors = field_sensor_data["datastream_name"].unique()

            if sensor_forecasted in sensors:

                if consortium_name != 'consortium2':
                    weather_data = get_sensors_meteo(
                        consortium_name=consortium_name,
                        meteo_df=weather_data,
                        location_ids=location_ids,
                        sensor_df=field_sensor_data,
                        method=meteo_method)

                full_table = merge_all_data(consortium_name, field_sensor_data, irrigation_data, weather_data, historical_weather_data, forecasted_weather_data)
                full_table = add_irrigation_forecast(full_table)

                if strategy == 'real' or strategy == 'hybrid':
                    weather_aquacrop, irrigation_aquacrop, soil_type, crop_type, planting_date, harvest_date, reference_date, sim_start_date, sim_end_date, SMT, MaxIrr = aquacrop_preparation(consortium_name, sensor_forecasted, full_table, strategy)

                    output_file = f'data//05_aquacrop_input//{consortium_name}//{sensor_forecasted}//irrigation_aquacrop.parquet'
                    os.makedirs(os.path.dirname(output_file), exist_ok=True)

                    save_to_file(
                        df=irrigation_aquacrop,
                        output_file=output_file,
                    )

                output_file = f'data//05_aquacrop_input//{consortium_name}//{sensor_forecasted}//weather_aquacrop.parquet'
                os.makedirs(os.path.dirname(output_file), exist_ok=True)

                save_to_file(
                    df=weather_aquacrop,
                    output_file=output_file,
                )

                output_file = f'data//05_aquacrop_input//{consortium_name}//{sensor_forecasted}//aquacrop_settings.json'
                os.makedirs(os.path.dirname(output_file), exist_ok=True)

                with open(output_file, 'w') as f:
                    json.dump({
                        'soil_type': soil_type,
                        'crop_type': crop_type,
                        'planting_date': planting_date,
                        'harvest_date': harvest_date,
                        'reference_date': reference_date.strftime('%Y-%m-%d'),
                        'sim_start_date': sim_start_date,
                        'sim_end_date': sim_end_date,
                        'SMT': SMT,
                        'MaxIrr': MaxIrr
                    }, f, indent=4)
            else:
                continue


if __name__ == "__main__":
    aquacrop_preparation_pipeline()