Spaces:
Sleeping
Sleeping
| """Predictor sets and feature engineering for the yield-prediction models. | |
| M1, M2, M3 are linear regressions on increasingly rich plot-level | |
| features. RF is a Random Forest baseline that uses a seven-feature | |
| predictor set including squared and total-area aggregate features. | |
| """ | |
| from __future__ import annotations | |
| import pandas as pd | |
| PREDICTORS_M1: tuple[str, ...] = ("total_area_in2",) | |
| PREDICTORS_M2: tuple[str, ...] = ( | |
| "area_B_in2", | |
| "area_A1_in2", | |
| "area_A2_in2", | |
| "area_OV_in2", | |
| ) | |
| PREDICTORS_M3: tuple[str, ...] = ( | |
| "area_B_in2", | |
| "area_A1_in2", | |
| "area_A2_in2", | |
| "area_OV_in2", | |
| "mean_area", | |
| ) | |
| # Wider feature set for the Random Forest baseline: four class areas plus | |
| # mean_area, total_area_sq, and total_area_in2. | |
| PREDICTORS_RF: tuple[str, ...] = ( | |
| "area_B_in2", | |
| "area_A1_in2", | |
| "area_A2_in2", | |
| "area_OV_in2", | |
| "mean_area", | |
| "total_area_sq", | |
| "total_area_in2", | |
| ) | |
| RESPONSE: str = "obs_sum_B_A1_A2_OV" | |
| RF_KWARGS: dict = { | |
| "n_estimators": 500, | |
| "random_state": 42, | |
| "max_depth": None, | |
| "min_samples_split": 2, | |
| "min_samples_leaf": 1, | |
| "n_jobs": -1, | |
| } | |
| def add_model_features(df: pd.DataFrame) -> pd.DataFrame: | |
| """Recompute total_area_in2 from class areas, add mean_area and total_area_sq. | |
| total_area_in2 is overwritten with the sum of the four class-area | |
| columns (this can differ from the upstream aggregate if any tubers | |
| have a NaN class). mean_area = total_area_in2 / n_tubers, and | |
| total_area_sq = total_area_in2 ** 2. | |
| """ | |
| df = df.copy() | |
| df["total_area_in2"] = df[list(PREDICTORS_M2)].sum(axis=1) | |
| df["mean_area"] = df["total_area_in2"] / df["n_tubers"] | |
| df["total_area_sq"] = df["total_area_in2"] ** 2 | |
| return df | |