training_code/data/dataset.py

from __future__ import annotations
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset

from data.preprocessing import (
    load_table,
    decode_bytes_in_object_cols,
    find_daily_cols,
    daily_to_cumulative_weekly,
    DEFAULT_WEATHER_AGG_RULES,
)

class YieldDataset(Dataset):
    def __init__(
        self,
        data_file: str | Path,
        weather_vars: list[str],
        soil_vars: list[str],
        split: str = "all",
        seed: int = 1234,
        crop: str | None = None,
        years: str | None = None,
        crop_col: str = "crop",
        yield_col: str = "yield",
        field_col: str = "farm_field",
        year_col: str = "year",
        time_agg: str = "weekly_cumulative",
        split_strategy: str = "field",
        val_split: float = 0.2,
        test_split: float = 0.2,
        require_yield: bool = True,
    ):
        self.weather_vars = weather_vars
        self.soil_vars = soil_vars
        self.crop_col = crop_col
        self.yield_col = yield_col
        self.field_col = field_col
        self.year_col = year_col
        self.time_agg = time_agg
        self.crop_map = {"corn": 0, "maize": 0, "soybean": 1, "soy": 1}

        self.w_mean = None
        self.w_std = None
        self.s_mean = None
        self.s_std = None

        df = load_table(data_file)
        df = decode_bytes_in_object_cols(df)
        df.columns = [str(c).strip() for c in df.columns]

        if yield_col != "yield" and yield_col in df.columns:
            df = df.rename(columns={yield_col: "yield"})
            self.yield_col = "yield"

        required = [crop_col, field_col, year_col] + soil_vars
        if require_yield:
            required.append(self.yield_col)

        missing = [c for c in required if c not in df.columns]
        if missing:
            raise ValueError(f"Missing required columns: {missing[:20]}")

        if require_yield:
            df[self.yield_col] = pd.to_numeric(df[self.yield_col], errors="coerce")
            df = df.dropna(subset=[self.yield_col]).reset_index(drop=True)
        else:
            if self.yield_col not in df.columns:
                df[self.yield_col] = np.nan

        df[crop_col] = df[crop_col].astype(str).str.strip().str.lower()
        if crop:
            df = df[df[crop_col] == crop.strip().lower()].reset_index(drop=True)

        if years:
            year_list = [int(y.strip()) for y in years.split(",") if y.strip()]
            df = df[df[year_col].isin(year_list)].reset_index(drop=True)

        self.weather_cols_by_var = {}

        for v in weather_vars:
            cols = find_daily_cols(df, v)

            if not cols:
                examples = [c for c in df.columns if str(c).startswith(v)][:10]
                raise ValueError(
                    f"No indexed columns found for weather var '{v}'. "
                    f"Examples: {examples}"
                )

            self.weather_cols_by_var[v] = cols


        if self.time_agg == "weekly":
            # Input data is already weekly.
            self.K = len(
                self.weather_cols_by_var[weather_vars[0]]
            )

        elif self.time_agg == "weekly_cumulative":
            # Existing behavior for daily input.
            self.K = len(
                daily_to_cumulative_weekly(
                    df.loc[
                        0,
                        self.weather_cols_by_var[weather_vars[0]]
                    ].to_numpy(dtype=np.float32),
                    agg=DEFAULT_WEATHER_AGG_RULES.get(
                        weather_vars[0],
                        "mean",
                    ),
                )
            )

        else:
            raise ValueError(
                "time_agg must be 'weekly' or 'weekly_cumulative'"
            )

        df = df.reset_index(drop=True)
        self.df_full = df
        self.indices = self._make_split_indices(df, split, seed, split_strategy, val_split, test_split)

        self.df = df
        self.soil_arr = df[soil_vars].to_numpy(np.float32)
        self.crop_arr = df[crop_col].astype(str).str.strip().str.lower().to_numpy()
        self.y_arr = df[self.yield_col].to_numpy(np.float32)

    def _make_split_indices(self, df, split, seed, split_strategy, val_split, test_split):
        rng = np.random.default_rng(seed)

        if split == "all" or split_strategy == "none":
            return np.arange(len(df))

        if split_strategy == "field":
            fields = df[self.field_col].astype(str).unique()
            rng.shuffle(fields)

            n = len(fields)
            n_test = max(1, int(test_split * n)) if n >= 3 else 0
            n_val = max(1, int(val_split * n)) if n >= 3 else 0
            n_train = n - n_val - n_test

            train_fields = set(fields[:n_train])
            val_fields = set(fields[n_train:n_train + n_val])
            test_fields = set(fields[n_train + n_val:])

            if split == "train":
                keep = train_fields
            elif split == "val":
                keep = val_fields
            elif split == "test":
                keep = test_fields
            else:
                raise ValueError("split must be train, val, test, or all")

            return df.index[df[self.field_col].astype(str).isin(keep)].to_numpy()

        if split_strategy == "random":
            idx = np.arange(len(df))
            rng.shuffle(idx)
            n = len(idx)
            n_test = int(test_split * n)
            n_val = int(val_split * n)
            n_train = n - n_val - n_test

            if split == "train":
                return idx[:n_train]
            if split == "val":
                return idx[n_train:n_train + n_val]
            if split == "test":
                return idx[n_train + n_val:]

        raise ValueError(f"Unknown split_strategy={split_strategy}")

    def set_normalization(self, w_mean, w_std, s_mean, s_std):
        self.w_mean = np.asarray(w_mean, dtype=np.float32)
        self.w_std = np.asarray(w_std, dtype=np.float32)
        self.s_mean = np.asarray(s_mean, dtype=np.float32)
        self.s_std = np.asarray(s_std, dtype=np.float32)

    def __len__(self):
        return len(self.indices)

    def __getitem__(self, idx):
        ridx = int(self.indices[idx])
        row = self.df.loc[ridx]

        weather_vars = []

        for v in self.weather_vars:

            values = row[
                self.weather_cols_by_var[v]
            ].to_numpy(dtype=np.float32)

            if self.time_agg == "weekly":

                # Data is already weekly.
                # Use the 52 weekly values exactly as supplied.
                seq = values

            elif self.time_agg == "weekly_cumulative":

                # Existing behavior for datasets containing daily weather.
                agg = DEFAULT_WEATHER_AGG_RULES.get(v, "mean")

                seq = daily_to_cumulative_weekly(
                    values,
                    agg=agg,
                    week_len=7,
                )

            weather_vars.append(seq)


        if self.time_agg == "weekly":

            # [K, number_of_weather_variables]
            weather = np.stack(
                weather_vars,
                axis=1,
            ).astype(np.float32)

        else:

            # Existing daily -> weekly+cumulative representation
            weather = np.stack(
                weather_vars,
                axis=1,
            ).astype(np.float32)

            weather = weather.reshape(
                weather.shape[0],
                -1,
            ).astype(np.float32)
        
        soil = self.soil_arr[ridx].astype(np.float32)

        if self.w_mean is not None:
            weather = np.where(np.isnan(weather), self.w_mean[None, :], weather)
            weather = (weather - self.w_mean[None, :]) / self.w_std[None, :]

        if self.s_mean is not None:
            soil = np.where(np.isnan(soil), self.s_mean, soil)
            soil = (soil - self.s_mean) / self.s_std

        crop = str(self.crop_arr[ridx]).strip().lower()
        crop_id = self.crop_map.get(crop, 0)

        return {
            "weather": torch.from_numpy(weather),
            "soil": torch.from_numpy(soil),
            "crop_id": torch.tensor(crop_id, dtype=torch.long),
            "yield": torch.tensor(float(self.y_arr[ridx]), dtype=torch.float32),
            "farm_field": str(row[self.field_col]),
            "year": int(row[self.year_col]),
        }