Upload 21 files

training_code/config/config.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+from dataclasses import dataclass, asdict
+from pathlib import Path
+import json
+
+@dataclass
+class TrainConfig:
+    mode: str = "train_eval"  # train_eval | train | eval | predict
+
+    train_file: Path | None = None
+    val_file: Path | None = None
+    test_file: Path | None = None
+    predict_file: Path | None = None
+    checkpoint_path: Path | None = None
+    output_csv: Path | None = None
+    single_sample_json: Path | None = None
+
+    weather_vars: list[str] | None = None
+    soil_vars: list[str] | None = None
+
+    crop_col: str = "crop"
+    yield_col: str = "yield"
+    field_col: str = "farm_field"
+    year_col: str = "year"
+
+    years: str | None = None
+    crop: str | None = None
+    time_agg: str = "weekly_cumulative"
+
+    train_cutoffs: list[int] | None = None
+    eval_cutoffs: list[int] | None = None
+    predict_cutoff: int | None = None
+
+    d_model: int = 64
+    nhead: int = 4
+    num_layers: int = 4
+    dim_ff: int = 128
+    dropout: float = 0.4
+    pool: str = "last"
+    use_crop: bool = True
+    crop_emb_dim: int = 8
+
+    epochs: int = 10
+    lr: float = 1e-4
+    batch_size: int = 64
+    weight_decay: float = 1e-3
+    seed: int = 1234
+    early_stop_patience: int = 3
+
+    split_strategy: str = "field"  # field | random | none
+    val_split: float = 0.2
+    test_split: float = 0.2
+
+    expt_name: str = "test"
+    out_dir: Path = Path("outputs")
+    log_dir: Path = Path("runs")
+
+    def validate(self) -> None:
+        if self.mode in {"train", "train_eval"}:
+            if self.train_file is None:
+                raise ValueError("train_file is required for training.")
+            if not self.weather_vars:
+                raise ValueError("weather_vars is required for training.")
+            if not self.soil_vars:
+                raise ValueError("soil_vars is required for training.")
+
+        if self.mode in {"eval", "predict"} and self.checkpoint_path is None:
+            raise ValueError("checkpoint_path is required for eval/predict.")
+
+        if self.mode == "eval" and self.test_file is None:
+            raise ValueError("test_file is required for eval.")
+
+        if self.mode == "predict":
+            if self.predict_file is None and self.single_sample_json is None:
+                raise ValueError("predict_file or single_sample_json is required.")
+            if self.predict_cutoff is None:
+                raise ValueError("predict_cutoff is required for predict.")
+
+    def save(self) -> None:
+        self.out_dir.mkdir(parents=True, exist_ok=True)
+        with open(self.out_dir / "config.json", "w") as f:
+            json.dump(asdict(self), f, indent=2, default=str)

training_code/data/dataset.py

@@ -0,0 +1,260 @@
+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]),
+        }

training_code/data/preprocessing.py

@@ -0,0 +1,119 @@
+import re
+import math
+import numpy as np
+import pandas as pd
+
+DEFAULT_WEATHER_AGG_RULES = {
+    "dayl": "mean",
+    "prcp": "sum",
+    "srad": "mean",
+    "tmax": "mean",
+    "tmin": "mean",
+    "vp": "mean",
+    "tmean": "mean",
+    "gdd": "sum",
+    "precip_3day_avg_perday": "mean",
+    "precip_7day_avg_perday": "mean",
+    "precip_14day_avg_perday": "mean",
+}
+
+def load_table(path):
+    path = str(path)
+    if path.endswith(".h5") or path.endswith(".hdf5"):
+        return pd.read_hdf(path)
+    if path.endswith(".csv"):
+        return pd.read_csv(path)
+    if path.endswith(".parquet"):
+        return pd.read_parquet(path)
+    raise ValueError(f"Unsupported file type: {path}")
+
+def decode_bytes_in_object_cols(df: pd.DataFrame) -> pd.DataFrame:
+    out = df.copy()
+    for c in out.select_dtypes(include=["object"]).columns:
+        if len(out) > 0 and isinstance(out[c].iloc[0], (bytes, bytearray)):
+            out[c] = out[c].str.decode("utf-8")
+    return out
+
+def find_daily_cols(df: pd.DataFrame, var: str) -> list[str]:
+    pat_bracket = re.compile(rf"^{re.escape(var)}.*?_\[(\d+)\]$")
+    pat_plain = re.compile(rf"^{re.escape(var)}.*?_(\d+)$")
+    hits = []
+    for c in df.columns:
+        s = str(c)
+        m = pat_bracket.match(s) or pat_plain.match(s)
+        if m:
+            hits.append((int(m.group(1)), c))
+    hits.sort(key=lambda x: x[0])
+    return [c for _, c in hits]
+
+# def daily_to_cumulative_weekly(daily: np.ndarray, agg: str, week_len: int = 7) -> np.ndarray:
+#     daily = daily.astype(np.float32)
+#     T = daily.shape[0]
+#     K = int(math.ceil(T / week_len))
+#     out = np.zeros(K, dtype=np.float32)
+
+#     for w in range(K):
+#         e = min((w + 1) * week_len, T)
+#         chunk = daily[:e]
+#         out[w] = np.nansum(chunk) if agg == "sum" else np.nanmean(chunk)
+
+#     return out
+
+def daily_to_cumulative_weekly(
+    daily: np.ndarray,
+    agg: str,
+    week_len: int = 7,
+) -> np.ndarray:
+    daily = daily.astype(np.float32)
+    T = daily.shape[0]
+    K = int(math.ceil(T / week_len))
+
+    weekly = np.zeros(K, dtype=np.float32)
+    cumulative = np.zeros(K, dtype=np.float32)
+
+    for w in range(K):
+        s = w * week_len
+        e = min((w + 1) * week_len, T)
+
+        week_chunk = daily[s:e]
+        cumulative_chunk = daily[:e]
+
+        if agg == "sum":
+            weekly[w] = np.nansum(week_chunk)
+            cumulative[w] = np.nansum(cumulative_chunk)
+        else:
+            weekly[w] = np.nanmean(week_chunk)
+            cumulative[w] = np.nanmean(cumulative_chunk)
+
+    return np.stack([weekly, cumulative], axis=1).astype(np.float32)
+
+def compute_x_stats(dataset, max_samples=20000, seed=0):
+    rng = np.random.default_rng(seed)
+    n = min(len(dataset), max_samples)
+    idxs = rng.choice(len(dataset), size=n, replace=False) if n < len(dataset) else np.arange(len(dataset))
+
+    weather, soil = [], []
+    for i in idxs:
+        item = dataset[int(i)]
+        weather.append(item["weather"].numpy())
+        soil.append(item["soil"].numpy())
+
+    W_all = np.stack(weather).astype(np.float32)
+    S_all = np.stack(soil).astype(np.float32)
+
+    w_mean = np.nanmean(W_all, axis=(0, 1))
+    w_std = np.nanstd(W_all, axis=(0, 1)) + 1e-6
+    s_mean = np.nanmean(S_all, axis=0)
+    s_std = np.nanstd(S_all, axis=0) + 1e-6
+
+    w_mean = np.where(np.isfinite(w_mean), w_mean, 0.0)
+    w_std = np.where((np.isfinite(w_std)) & (w_std > 1e-6), w_std, 1.0)
+    s_mean = np.where(np.isfinite(s_mean), s_mean, 0.0)
+    s_std = np.where((np.isfinite(s_std)) & (s_std > 1e-6), s_std, 1.0)
+
+    return w_mean, w_std, s_mean, s_std
+
+def compute_y_stats(dataset):
+    ys = np.array([float(dataset[i]["yield"]) for i in range(len(dataset))], dtype=np.float32)
+    #ys_log = np.log1p(np.clip(ys, 0.0, None))
+    return float(ys.mean()), float(ys.std() + 1e-6)

training_code/examples/sample_input_weekly.json

@@ -0,0 +1,399 @@
+{
+  "crop": "corn",
+  "weather_format": "weekly",
+  "cutoff": 52,
+  "weather": {
+    "prcp": [
+      0.0,
+      0.701298713684082,
+      0.4545454680919647,
+      2.155844211578369,
+      0.9350649118423462,
+      0.11688311398029327,
+      0.0,
+      3.753246784210205,
+      0.15584415197372437,
+      2.2337663173675537,
+      0.0,
+      0.4285714328289032,
+      2.402597427368164,
+      2.3896102905273438,
+      5.753246784210205,
+      0.9220778942108154,
+      0.0,
+      1.1298701763153076,
+      1.2077921628952026,
+      1.2207791805267334,
+      2.441558361053467,
+      5.402597427368164,
+      2.909090995788574,
+      3.6883115768432617,
+      5.1688313484191895,
+      5.233766078948975,
+      7.389610290527344,
+      4.883116722106934,
+      1.5194804668426514,
+      0.8311688303947449,
+      4.051948070526123,
+      0.03896103799343109,
+      0.0,
+      5.350649356842041,
+      2.0389609336853027,
+      4.103896141052246,
+      0.012987012974917889,
+      15.49350643157959,
+      1.1168831586837769,
+      3.142857074737549,
+      6.155844211578369,
+      0.3246753215789795,
+      1.6753246784210205,
+      0.5714285969734192,
+      2.8831169605255127,
+      0.10389610379934311,
+      0.9610389471054077,
+      1.1168831586837769,
+      0.0,
+      0.0,
+      0.051948051899671555,
+      4.181818008422852
+    ],
+    "srad": [
+      187.0129852294922,
+      181.23635864257812,
+      184.8935089111328,
+      179.7818145751953,
+      255.91688537597656,
+      267.4701232910156,
+      323.158447265625,
+      282.09869384765625,
+      347.76104736328125,
+      274.28570556640625,
+      384.70648193359375,
+      332.7168884277344,
+      384.0831298828125,
+      455.7298583984375,
+      359.5636291503906,
+      525.6727294921875,
+      584.2701416015625,
+      522.7221069335938,
+      385.9532470703125,
+      420.5298767089844,
+      386.36883544921875,
+      335.12725830078125,
+      327.6467590332031,
+      329.35064697265625,
+      319.7090759277344,
+      341.6103820800781,
+      374.150634765625,
+      387.20001220703125,
+      379.0545349121094,
+      412.1766357421875,
+      353.08050537109375,
+      397.0493469238281,
+      376.4363708496094,
+      312.2701416015625,
+      312.4363708496094,
+      313.7662353515625,
+      341.07012939453125,
+      249.10130310058594,
+      292.32208251953125,
+      190.7532501220703,
+      189.13246154785156,
+      301.589599609375,
+      214.69090270996094,
+      198.7324676513672,
+      148.0311737060547,
+      213.77662658691406,
+      168.64414978027344,
+      116.23896026611328,
+      148.4051971435547,
+      169.22596740722656,
+      130.86753845214844,
+      157.7818145751953
+    ],
+    "swe": [
+      18.909090042114258,
+      18.753246307373047,
+      23.220779418945312,
+      32.93506622314453,
+      37.24675369262695,
+      41.45454406738281,
+      41.45454406738281,
+      51.22077941894531,
+      67.84415435791016,
+      62.181819915771484,
+      61.610389709472656,
+      59.844154357910156,
+      54.80519485473633,
+      71.94805145263672,
+      78.9610366821289,
+      100.98701477050781,
+      84.467529296875,
+      50.54545593261719,
+      7.636363506317139,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      0.0,
+      1.402597427368164,
+      0.0,
+      0.0,
+      0.0,
+      9.454545021057129
+    ],
+    "tmax": [
+      -13.350648880004883,
+      -3.694805145263672,
+      -2.4740259647369385,
+      -0.1428571492433548,
+      -7.298701286315918,
+      -10.746753692626953,
+      -1.3701298236846924,
+      -2.844155788421631,
+      5.285714149475098,
+      0.7272727489471436,
+      4.675324440002441,
+      3.948051929473877,
+      4.324675559997559,
+      -0.4285714328289032,
+      3.9285714626312256,
+      8.292207717895508,
+      16.785715103149414,
+      23.915584564208984,
+      19.662338256835938,
+      25.34415626525879,
+      28.39610481262207,
+      28.39610481262207,
+      24.863636016845703,
+      27.415584564208984,
+      25.766233444213867,
+      28.10389518737793,
+      28.441558837890625,
+      30.61688232421875,
+      26.87013053894043,
+      26.402597427368164,
+      26.571428298950195,
+      29.324674606323242,
+      29.714284896850586,
+      25.02597427368164,
+      25.441558837890625,
+      23.467533111572266,
+      28.746753692626953,
+      20.162338256835938,
+      15.733766555786133,
+      12.071428298950195,
+      8.357142448425293,
+      12.409090995788574,
+      11.629870414733887,
+      9.253246307373047,
+      -0.012987012974917889,
+      -0.6623376607894897,
+      0.9155844449996948,
+      -2.0,
+      -3.344155788421631,
+      0.5064935088157654,
+      1.7402597665786743,
+      -2.0568182468414307
+    ],
+    "tmin": [
+      -23.694805145263672,
+      -14.48051929473877,
+      -12.051947593688965,
+      -7.941558361053467,
+      -19.376623153686523,
+      -21.123376846313477,
+      -14.344156265258789,
+      -14.02597427368164,
+      -7.7207794189453125,
+      -7.7272725105285645,
+      -6.344155788421631,
+      -4.116883277893066,
+      -5.246753215789795,
+      -11.681818008422852,
+      -3.909090995788574,
+      -3.649350643157959,
+      1.4220778942108154,
+      8.98051929473877,
+      8.870129585266113,
+      10.454545021057129,
+      15.129870414733887,
+      16.363636016845703,
+      14.461038589477539,
+      17.175325393676758,
+      16.701297760009766,
+      18.720779418945312,
+      17.922077178955078,
+      19.519479751586914,
+      16.305194854736328,
+      14.642857551574707,
+      15.694805145263672,
+      17.422077178955078,
+      18.220779418945312,
+      15.175324440002441,
+      15.558441162109375,
+      13.49350643157959,
+      17.201297760009766,
+      11.655843734741211,
+      5.7727274894714355,
+      4.551948070526123,
+      2.2272727489471436,
+      -0.16233766078948975,
+      1.7727272510528564,
+      1.0324674844741821,
+      -5.954545497894287,
+      -10.149351119995117,
+      -6.623376846313477,
+      -7.201298713684082,
+      -10.422078132629395,
+      -8.344156265258789,
+      -4.707792282104492,
+      -9.943181991577148
+    ],
+    "vp": [
+      100.2597427368164,
+      259.22076416015625,
+      287.2727355957031,
+      344.41558837890625,
+      137.14285278320312,
+      115.84415435791016,
+      211.42857360839844,
+      215.06494140625,
+      359.48052978515625,
+      358.4415588378906,
+      377.6623229980469,
+      454.5454406738281,
+      409.8701171875,
+      249.87013244628906,
+      451.4285583496094,
+      438.4415588378906,
+      540.7792358398438,
+      887.792236328125,
+      1041.0389404296875,
+      1108.052001953125,
+      1534.54541015625,
+      1750.6492919921875,
+      1589.6103515625,
+      1931.4285888671875,
+      1902.337646484375,
+      2178.181884765625,
+      2062.857177734375,
+      2275.32470703125,
+      1851.947998046875,
+      1670.1298828125,
+      1802.077880859375,
+      1990.1298828125,
+      2097.662353515625,
+      1738.1817626953125,
+      1810.9090576171875,
+      1587.012939453125,
+      2006.7532958984375,
+      1414.0260009765625,
+      942.337646484375,
+      852.467529296875,
+      731.9480590820312,
+      614.5454711914062,
+      715.8441772460938,
+      658.7012939453125,
+      416.1038818359375,
+      297.1428527832031,
+      392.7272644042969,
+      367.7922058105469,
+      289.35064697265625,
+      336.6233825683594,
+      435.3246765136719,
+      295.4545593261719
+    ]
+  },
+  "soil": {
+    "bdod_mean_0-5cm": 127.71450805664062,
+    "bdod_mean_5-15cm": 148.86965942382812,
+    "bdod_mean_15-30cm": 150.70364379882812,
+    "bdod_mean_30-60cm": 158.2148895263672,
+    "bdod_mean_60-100cm": 162.8789825439453,
+    "bdod_mean_100-200cm": 165.12025451660156,
+    "cec_mean_0-5cm": 252.34054565429688,
+    "cec_mean_5-15cm": 174.4131317138672,
+    "cec_mean_15-30cm": 142.52432250976562,
+    "cec_mean_30-60cm": 129.1745147705078,
+    "cec_mean_60-100cm": 119.97528839111328,
+    "cec_mean_100-200cm": 103.29729461669922,
+    "cfvo_mean_0-5cm": 21.431659698486328,
+    "cfvo_mean_5-15cm": 22.905019760131836,
+    "cfvo_mean_15-30cm": 21.820077896118164,
+    "cfvo_mean_30-60cm": 24.725868225097656,
+    "cfvo_mean_60-100cm": 34.54826354980469,
+    "cfvo_mean_100-200cm": 45.90888214111328,
+    "clay_mean_0-5cm": 150.02239990234375,
+    "clay_mean_5-15cm": 149.1436309814453,
+    "clay_mean_15-30cm": 150.81776428222656,
+    "clay_mean_30-60cm": 150.43707275390625,
+    "clay_mean_60-100cm": 134.29266357421875,
+    "clay_mean_100-200cm": 125.47721862792969,
+    "nitrogen_mean_0-5cm": 4723.78857421875,
+    "nitrogen_mean_5-15cm": 2148.763671875,
+    "nitrogen_mean_15-30cm": 1729.8802490234375,
+    "nitrogen_mean_30-60cm": 776.644775390625,
+    "nitrogen_mean_60-100cm": 547.14208984375,
+    "nitrogen_mean_100-200cm": 534.0015258789062,
+    "ocd_mean_0-5cm": 435.4494323730469,
+    "ocd_mean_5-15cm": 255.3752899169922,
+    "ocd_mean_15-30cm": 181.84786987304688,
+    "ocd_mean_30-60cm": 93.57837677001953,
+    "ocd_mean_60-100cm": 60.80463409423828,
+    "ocd_mean_100-200cm": 44.90888214111328,
+    "ocs_mean_0-5cm": 57.79845428466797,
+    "ocs_mean_5-15cm": 57.79845428466797,
+    "ocs_mean_15-30cm": 57.79845428466797,
+    "ocs_mean_30-60cm": 57.79845428466797,
+    "ocs_mean_60-100cm": 57.79845428466797,
+    "ocs_mean_100-200cm": 57.79845428466797,
+    "phh2o_mean_0-5cm": 61.812355041503906,
+    "phh2o_mean_5-15cm": 61.73050308227539,
+    "phh2o_mean_15-30cm": 61.63166046142578,
+    "phh2o_mean_30-60cm": 62.09111785888672,
+    "phh2o_mean_60-100cm": 64.0540542602539,
+    "phh2o_mean_100-200cm": 69.0687255859375,
+    "sand_mean_0-5cm": 625.2216186523438,
+    "sand_mean_5-15cm": 626.4702758789062,
+    "sand_mean_15-30cm": 631.22314453125,
+    "sand_mean_30-60cm": 649.113525390625,
+    "sand_mean_60-100cm": 697.4996337890625,
+    "sand_mean_100-200cm": 712.781494140625,
+    "silt_mean_0-5cm": 224.77297973632812,
+    "silt_mean_5-15cm": 224.3876495361328,
+    "silt_mean_15-30cm": 217.95135498046875,
+    "silt_mean_30-60cm": 200.43319702148438,
+    "silt_mean_60-100cm": 168.21775817871094,
+    "silt_mean_100-200cm": 161.7498016357422,
+    "soc_mean_0-5cm": 446.14508056640625,
+    "soc_mean_5-15cm": 223.34988403320312,
+    "soc_mean_15-30cm": 144.84561157226562,
+    "soc_mean_30-60cm": 107.50271606445312,
+    "soc_mean_60-100cm": 53.54926300048828,
+    "soc_mean_100-200cm": 47.3972053527832
+  }
+}

training_code/hf/auto.py

@@ -0,0 +1,15 @@
+from transformers import AutoConfig, AutoModel
+from hf.configuration_yield import YieldConfig
+from hf.modeling_yield import YieldForRegression
+
+
+def register_yield_autoclass():
+    try:
+        AutoConfig.register("yield-weather-soil", YieldConfig)
+    except ValueError:
+        pass
+
+    try:
+        AutoModel.register(YieldConfig, YieldForRegression)
+    except ValueError:
+        pass

training_code/hf/configuration_yield.py

@@ -0,0 +1,58 @@
+from transformers import PretrainedConfig
+
+
+class YieldConfig(PretrainedConfig):
+    model_type = "yield-weather-soil"
+
+    def __init__(
+        self,
+        weather_vars=None,
+        soil_vars=None,
+        w_mean=None,
+        w_std=None,
+        s_mean=None,
+        s_std=None,
+        y_mean=None,
+        y_std=None,
+        K=None,
+        W=None,
+        S=None,
+        train_cutoffs=None,
+        eval_cutoffs=None,
+        d_model=128,
+        nhead=4,
+        num_layers=4,
+        dim_ff=256,
+        dropout=0.3,
+        pool="mean",
+        use_crop=True,
+        crop_emb_dim=8,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.weather_vars = weather_vars
+        self.soil_vars = soil_vars
+
+        self.w_mean = w_mean
+        self.w_std = w_std
+        self.s_mean = s_mean
+        self.s_std = s_std
+        self.y_mean = y_mean
+        self.y_std = y_std
+
+        self.K = K
+        self.W = W
+        self.S = S
+
+        self.train_cutoffs = train_cutoffs
+        self.eval_cutoffs = eval_cutoffs
+
+        self.d_model = d_model
+        self.nhead = nhead
+        self.num_layers = num_layers
+        self.dim_ff = dim_ff
+        self.dropout = dropout
+        self.pool = pool
+        self.use_crop = use_crop
+        self.crop_emb_dim = crop_emb_dim

training_code/hf/modeling_yield.py

@@ -0,0 +1,91 @@
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+import torch
+from torch import nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from dataclasses import dataclass
+
+from hf.configuration_yield import YieldConfig
+from models.unimodal_ws_crossattn import UnimodalWS_CrossAttn_TemporalTF
+
+
+@dataclass
+class YieldModelOutput(ModelOutput):
+    loss: torch.Tensor | None = None
+    logits: torch.Tensor | None = None
+    predictions: torch.Tensor | None = None
+
+
+class YieldForRegression(PreTrainedModel):
+    config_class = YieldConfig
+    base_model_prefix = "yield_model"
+
+    def __init__(self, config: YieldConfig):
+        super().__init__(config)
+
+        self.yield_model = UnimodalWS_CrossAttn_TemporalTF(
+            w_dim=config.W,
+            soil_dim=config.S,
+            d_model=config.d_model,
+            nhead=config.nhead,
+            num_layers=config.num_layers,
+            dim_ff=config.dim_ff,
+            dropout=config.dropout,
+            use_crop=config.use_crop,
+            crop_emb_dim=config.crop_emb_dim,
+            max_weeks=max(32, config.K),
+            pool=config.pool,
+        )
+
+        self.post_init()
+
+    def forward(
+        self,
+        weather,
+        soil,
+        crop_id,
+        labels=None,
+        horizon_idx=None,
+        causal=True,
+        return_sequence=False,
+        return_dict=True,
+    ):
+        if horizon_idx is None:
+            horizon_idx = weather.shape[1]
+
+        logits = self.yield_model(
+            weather,
+            soil,
+            crop_id,
+            horizon_idx=horizon_idx,
+            causal=causal,
+            return_sequence=return_sequence,
+        )
+
+        y_mean = torch.tensor(self.config.y_mean, device=logits.device, dtype=logits.dtype)
+        y_std = torch.tensor(self.config.y_std, device=logits.device, dtype=logits.dtype)
+
+        #predictions = torch.expm1(logits * y_std + y_mean)
+        predictions = logits * y_std + y_mean
+
+        loss = None
+        if labels is not None:
+            # labels_log = torch.log1p(torch.clamp(labels, min=0.0))
+            # labels_norm = (labels_log - y_mean) / y_std
+            # loss = nn.functional.mse_loss(logits, labels_norm)
+            labels_norm = (labels - y_mean) / y_std
+            loss = nn.functional.mse_loss(logits, labels_norm)
+
+        if not return_dict:
+            return (loss, logits, predictions)
+
+        return YieldModelOutput(
+            loss=loss,
+            logits=logits,
+            predictions=predictions,
+        )

training_code/models/unimodal_ws_crossattn.py

@@ -0,0 +1,171 @@
+import torch
+import torch.nn as nn
+
+
+class MLP(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int, hidden: int = 128, dropout: float = 0.05):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_dim, hidden//2),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden//2, hidden),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden, out_dim)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+class UnimodalWS_CrossAttn_TemporalTF(nn.Module):
+    def __init__(
+        self,
+        w_dim: int,
+        soil_dim: int,
+        d_model: int = 128,
+        nhead: int = 4,
+        num_layers: int = 2,
+        dim_ff: int = 256,
+        dropout: float = 0.05,
+        use_crop: bool = True,
+        crop_emb_dim: int = 8,
+        max_weeks: int = 32,
+        pool: str = "mean", # Changed default to mean
+        farm_emb_dim: int = 8,
+        horizon_emb_dim: int = 16,
+    ):
+        super().__init__()
+        assert pool in ["last", "cls", "mean"]
+        self.pool = pool
+        self.use_crop = use_crop
+        self.max_weeks = max_weeks
+
+        crop_in = crop_emb_dim if use_crop else 0
+        if use_crop:
+            self.crop_emb = nn.Embedding(2, crop_emb_dim)
+
+        self.horizon_emb = nn.Embedding(max_weeks + 1, horizon_emb_dim)
+        self.horizon_proj = nn.Linear(horizon_emb_dim, d_model)
+        
+        # Hard Positional Signal: Week ID projection
+        self.week_proj = nn.Linear(1, d_model) 
+
+        self.weather_enc = MLP(w_dim + crop_in, d_model, hidden=d_model, dropout=dropout)
+        self.soil_enc = MLP(soil_dim + crop_in, d_model, hidden=d_model, dropout=dropout)
+
+        self.ws_cross_attn = nn.MultiheadAttention(
+            embed_dim=d_model,
+            num_heads=nhead,
+            dropout=dropout,
+            batch_first=True,
+        )
+        self.ws_ln = nn.LayerNorm(d_model)
+
+        self.pos_emb = nn.Parameter(torch.zeros(1, max_weeks + 1, d_model))
+        nn.init.trunc_normal_(self.pos_emb, std=0.02)
+
+        if pool == "cls":
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
+            nn.init.trunc_normal_(self.cls_token, std=0.02)
+
+        enc_layer = nn.TransformerEncoderLayer(
+            d_model=d_model, nhead=nhead, dim_feedforward=dim_ff,
+            dropout=dropout, activation="gelu", batch_first=True, norm_first=True,
+        )
+        self.temporal_tf = nn.TransformerEncoder(enc_layer, num_layers=num_layers)
+
+        # Head now takes [Pooled Sequence + Horizon Embedding + Soil Skip Connection]
+        self.head = nn.Sequential(
+            nn.LayerNorm(d_model + horizon_emb_dim + d_model), 
+            nn.Linear(d_model + horizon_emb_dim + d_model, d_model),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_model, 1),
+        )
+
+    def forward(
+        self,
+        weather,
+        soil,
+        crop_id,
+        farm_field_encoded=None, 
+        week_mask=None,
+        causal=False,
+        return_sequence=False,
+        horizon_idx=None,
+    ):
+        B, t, W = weather.shape
+        device = weather.device
+
+        # 1. Horizon Processing
+        if horizon_idx is None: horizon_idx = torch.tensor(t, device=device).expand(B)
+        if not torch.is_tensor(horizon_idx): horizon_idx = torch.tensor(horizon_idx, device=device).expand(B)
+        
+        h_idx = horizon_idx.long().clamp(min=1, max=self.max_weeks)
+        h_emb = self.horizon_emb(h_idx)
+        h_tok = self.horizon_proj(h_emb) # [B, D]
+
+        # 2. Input Encoding
+        if self.use_crop:
+            c = self.crop_emb(crop_id)
+            weather_in = torch.cat([weather, c[:, None, :].expand(-1, t, -1)], dim=-1)
+            soil_in = torch.cat([soil, c], dim=-1)
+        else:
+            weather_in, soil_in = weather, soil
+
+        w_tok = self.weather_enc(weather_in)
+        s_encoded = self.soil_enc(soil_in) # [B, D] - Saved for skip connection later
+        
+        # 3. Cross Attention & Horizon 
+        # attn_out, _ = self.ws_cross_attn(query=w_tok, key=s_encoded[:, None, :], value=s_encoded[:, None, :])
+        # fused = self.ws_ln(w_tok + attn_out)
+
+        fused = self.ws_ln(w_tok + s_encoded[:, None, :])
+        fused = fused + h_tok[:, None, :] # Horizon aware tokens
+
+        # 4. Temporal Transformer
+        # Add Hard Week Signal + Learned Pos Emb
+        weeks = torch.arange(1, t + 1, device=device).float() / self.max_weeks
+        week_signal = self.week_proj(weeks[None, :, None].expand(B, -1, -1))
+        
+        x = fused + week_signal + self.pos_emb[:, :t, :]
+        
+        if self.pool == "cls":
+            x = torch.cat([self.cls_token.expand(B, -1, -1), x], dim=1)
+            t_idx = t + 1
+        else: t_idx = t
+
+        src_key_padding_mask = ~week_mask.bool() if week_mask is not None else None
+        #h = self.temporal_tf(x, src_key_padding_mask=src_key_padding_mask)
+
+        attn_mask = None
+        if causal:
+            L = x.size(1)
+            attn_mask = torch.triu(
+                torch.ones(L, L, device=device, dtype=torch.bool),
+                diagonal=1,
+            )
+
+        h = self.temporal_tf(
+            x,
+            mask=attn_mask,
+            src_key_padding_mask=src_key_padding_mask,
+        )
+
+        # 5. Global Mean Pooling
+        if self.pool == "mean":
+            if week_mask is not None:
+                mask = week_mask.unsqueeze(-1).float()
+                pooled = (h * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
+            else:
+                pooled = h.mean(dim=1)
+        elif self.pool == "cls":
+            pooled = h[:, 0, :]
+        else: # "last"
+            pooled = h[:, -1, :]
+
+        # 6. Final Fusion (Pooled Sequence + Horizon + Soil Skip)
+        out = torch.cat([pooled, h_emb, s_encoded], dim=-1)
+        return self.head(out).squeeze(-1)
+

training_code/requirements.txt

@@ -0,0 +1,153 @@
+annotated-types==0.7.0
+anyio==4.11.0
+argon2-cffi==25.1.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+ase==3.22.1
+asttokens==2.2.1
+async-lru==2.0.5
+attrs==25.4.0
+babel==2.17.0
+backcall==0.2.0
+beautifulsoup4==4.14.2
+bleach==6.3.0
+blosc2==4.3.3
+certifi==2026.6.17
+charset-normalizer==3.4.7
+comm==0.2.3
+contourpy==1.0.7
+cycler==0.11.0
+debugpy==1.8.17
+decorator==5.1.1
+defusedxml==0.7.1
+exceptiongroup==1.3.1
+executing==1.2.0
+fastjsonschema==2.21.2
+filelock==3.29.0
+fonttools==4.39.3
+fqdn==1.5.1
+fsspec==2026.4.0
+googledrivedownloader==0.4
+h11==0.16.0
+h5py==3.8.0
+hf-xet==1.5.1
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==0.32.4
+idna==3.18
+ipdb==0.13.7
+ipykernel==7.1.0
+ipython==8.12.0
+ipywidgets==8.1.8
+isodate==0.6.1
+isoduration==20.11.0
+jedi==0.18.2
+Jinja2==3.1.6
+joblib==1.2.0
+json5==0.12.1
+jsonschema==4.25.1
+jsonschema-specifications==2025.9.1
+jupyter==1.1.1
+jupyter-console==6.6.3
+jupyter-events==0.12.0
+jupyter-lsp==2.3.0
+jupyter_client==8.6.3
+jupyter_core==5.9.1
+jupyter_server==2.17.0
+jupyter_server_terminals==0.5.3
+jupyterlab==4.5.0
+jupyterlab_pygments==0.3.0
+jupyterlab_server==2.28.0
+jupyterlab_widgets==3.0.16
+kiwisolver==1.4.4
+lark==1.3.1
+llvmlite==0.39.1
+MarkupSafe==3.0.3
+matplotlib==3.10.7
+matplotlib-inline==0.1.6
+mistune==3.1.4
+mpmath==1.3.0
+msgpack==1.2.1
+nbclient==0.10.2
+nbconvert==7.16.6
+nbformat==5.10.4
+ndindex==1.10.1
+nest-asyncio==1.6.0
+networkx==3.4.2
+notebook==7.5.0
+notebook_shim==0.2.4
+numba==0.56.4
+numexpr==2.14.1
+numpy==2.2.6
+nvidia-cufile-cu12==1.13.1.3
+nvidia-cusparselt-cu12==0.7.1
+nvidia-nvjitlink-cu12==12.8.93
+nvidia-nvshmem-cu12==3.3.20
+overrides==7.7.0
+packaging==26.0
+pandas==2.3.3
+pandocfilters==1.5.1
+parso==0.8.3
+pexpect==4.8.0
+pickleshare==0.7.5
+pillow==12.0.0
+prometheus_client==0.23.1
+prompt-toolkit==3.0.38
+psutil==7.1.3
+ptyprocess==0.7.0
+pure-eval==0.2.2
+py-cpuinfo==9.0.0
+pydantic==2.13.4
+pydantic_core==2.46.4
+Pygments==2.14.0
+pyparsing==3.0.9
+python-dateutil==2.8.2
+python-json-logger==4.0.0
+python-louvain==0.16
+pytz==2023.3
+PyYAML==6.0.3
+pyzmq==27.1.0
+rdflib==6.3.2
+referencing==0.37.0
+regex==2026.6.28
+requests==2.34.2
+rfc3339-validator==0.1.4
+rfc3986-validator==0.1.1
+rfc3987-syntax==1.1.0
+rpds-py==0.29.0
+safetensors==0.8.0
+scikit-learn==1.7.2
+scipy==1.15.3
+seaborn==0.12.2
+Send2Trash==1.8.3
+six==1.16.0
+sklearn==0.0.post1
+sniffio==1.3.1
+soupsieve==2.8
+stack-data==0.6.2
+sympy==1.14.0
+tables==3.10.1
+terminado==0.18.1
+threadpoolctl==3.1.0
+tinycss2==1.4.0
+tokenizers==0.21.1
+toml==0.10.2
+tomli==2.3.0
+torch==2.12.1+cpu
+torch-geometric==1.7.0
+torchcodec==0.8.1
+torchvision==0.27.1+cpu
+tornado==6.5.2
+tqdm==4.68.3
+traitlets==5.9.0
+transformers==4.48.3
+typing-inspection==0.4.2
+typing_extensions==4.15.0
+tzdata==2025.2
+uri-template==1.3.0
+urllib3==2.7.0
+wcwidth==0.2.6
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+widgetsnbextension==4.0.15

training_code/scripts/evaluate_hf.py

@@ -0,0 +1,92 @@
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+import argparse
+import json
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoModel
+
+from data.dataset import YieldDataset
+from training.engine import evaluate
+from hf.auto import register_yield_autoclass
+
+
+def parse_int_list(x):
+    return [int(v.strip()) for v in x.split(",") if v.strip()]
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--hf_model_dir", required=True)
+    p.add_argument("--test_file", required=True)
+    p.add_argument("--cutoffs", default=None)
+    p.add_argument("--batch_size", type=int, default=64)
+    p.add_argument("--output_csv", default="eval_predictions.csv")
+    p.add_argument("--metrics_json", default="eval_metrics.json")
+    p.add_argument(
+        "--time_agg",
+        default="weekly",
+        choices=["weekly", "weekly_cumulative"],
+    )
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    register_yield_autoclass()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = AutoModel.from_pretrained(args.hf_model_dir).to(device)
+    model.eval()
+    cfg = model.config
+
+    cutoffs = parse_int_list(args.cutoffs) if args.cutoffs else cfg.eval_cutoffs
+
+    test_ds = YieldDataset(
+        data_file=args.test_file,
+        weather_vars=cfg.weather_vars,
+        soil_vars=cfg.soil_vars,
+        split="all",
+        seed=1234,
+        crop=None,
+        years=None,
+        time_agg=args.time_agg,
+    )
+    test_ds.set_normalization(
+        cfg.w_mean,
+        cfg.w_std,
+        cfg.s_mean,
+        cfg.s_std,
+    )
+
+    loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False)
+
+    metrics, rows = evaluate(
+        model=model,
+        loader=loader,
+        device=device,
+        y_mean=cfg.y_mean,
+        y_std=cfg.y_std,
+        cutoffs=cutoffs,
+    )
+
+    print(json.dumps(metrics, indent=2))
+
+    Path(args.output_csv).parent.mkdir(parents=True, exist_ok=True)
+    pd.DataFrame(rows).to_csv(args.output_csv, index=False)
+
+    with open(args.metrics_json, "w") as f:
+        json.dump(metrics, f, indent=2)
+
+    print(f"Saved predictions to {args.output_csv}")
+    print(f"Saved metrics to {args.metrics_json}")
+
+
+if __name__ == "__main__":
+    main()

training_code/scripts/inference_hf.py

@@ -0,0 +1,197 @@
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+import argparse
+import json
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoModel
+
+from data.dataset import YieldDataset
+from data.preprocessing import daily_to_cumulative_weekly, DEFAULT_WEATHER_AGG_RULES
+from hf.auto import register_yield_autoclass
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--hf_model_dir", required=True)
+    p.add_argument("--input_file", default=None)
+    p.add_argument("--single_sample_json", default=None)
+    p.add_argument("--cutoff", type=int, required=True)
+    p.add_argument("--batch_size", type=int, default=64)
+    p.add_argument("--output_csv", default="inference_predictions.csv")
+    return p.parse_args()
+
+
+def load_single_sample_json(path, cfg, cutoff):
+    with open(path, "r") as f:
+        sample = json.load(f)
+
+    weather_format = sample.get("weather_format", "weekly_cumulative")
+
+    weather_cols = []
+    for v in cfg.weather_vars:
+        if v not in sample["weather"]:
+            raise ValueError(f"Missing weather variable in JSON: {v}")
+
+        arr = np.asarray(sample["weather"][v], dtype=np.float32)
+
+        if weather_format == "daily":
+            agg = DEFAULT_WEATHER_AGG_RULES.get(v, "mean")
+            arr = daily_to_cumulative_weekly(arr, agg=agg, week_len=7)
+        elif weather_format in ("weekly", "weekly_cumulative"):
+            pass
+        else:
+            raise ValueError(
+                "weather_format must be 'daily', 'weekly', "
+                "or 'weekly_cumulative'."
+            )
+
+        weather_cols.append(arr)
+
+    lengths = [len(x) for x in weather_cols]
+    if len(set(lengths)) != 1:
+        raise ValueError(f"Weather variable lengths do not match after aggregation: {lengths}")
+
+    weather = np.stack(weather_cols, axis=1).astype(np.float32)
+
+    soil = []
+    for v in cfg.soil_vars:
+        if v not in sample["soil"]:
+            raise ValueError(f"Missing soil variable in JSON: {v}")
+        soil.append(float(sample["soil"][v]))
+
+    soil = np.asarray(soil, dtype=np.float32)
+
+    w_mean = np.asarray(cfg.w_mean, dtype=np.float32)
+    w_std = np.asarray(cfg.w_std, dtype=np.float32)
+    s_mean = np.asarray(cfg.s_mean, dtype=np.float32)
+    s_std = np.asarray(cfg.s_std, dtype=np.float32)
+
+    weather = np.where(np.isnan(weather), w_mean[None, :], weather)
+    weather = (weather - w_mean[None, :]) / w_std[None, :]
+
+    soil = np.where(np.isnan(soil), s_mean, soil)
+    soil = (soil - s_mean) / s_std
+
+    crop_map = {"corn": 0, "maize": 0, "soybean": 1, "soy": 1}
+    crop = str(sample.get("crop", "corn")).strip().lower()
+    crop_id = crop_map.get(crop, 0)
+
+    t_eff = min(cutoff, weather.shape[0])
+
+    return {
+        "weather": torch.from_numpy(weather[:t_eff]).unsqueeze(0),
+        "soil": torch.from_numpy(soil).unsqueeze(0),
+        "crop_id": torch.tensor([crop_id], dtype=torch.long),
+        "t_eff": t_eff,
+    }
+
+
+@torch.no_grad()
+def main():
+    args = parse_args()
+    register_yield_autoclass()
+
+    if args.input_file is None and args.single_sample_json is None:
+        raise ValueError("Provide either --input_file or --single_sample_json.")
+
+    if args.input_file is not None and args.single_sample_json is not None:
+        raise ValueError("Use only one: --input_file OR --single_sample_json.")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = AutoModel.from_pretrained(args.hf_model_dir).to(device)
+    model.eval()
+    cfg = model.config
+
+    if args.single_sample_json is not None:
+        sample = load_single_sample_json(args.single_sample_json, cfg, args.cutoff)
+
+        weather = sample["weather"].to(device)
+        soil = sample["soil"].to(device)
+        crop_id = sample["crop_id"].to(device)
+        t_eff = sample["t_eff"]
+
+        out = model(
+            weather=weather,
+            soil=soil,
+            crop_id=crop_id,
+            horizon_idx=t_eff,
+            causal=True,
+            return_sequence=False,
+        )
+
+        pred = float(out.predictions.item())
+
+        Path(args.output_csv).parent.mkdir(parents=True, exist_ok=True)
+        pd.DataFrame([{
+            "sample_idx": 0,
+            "cutoff": int(args.cutoff),
+            "y_pred": pred,
+        }]).to_csv(args.output_csv, index=False)
+
+        print(f"Predicted yield: {pred:.4f}")
+        print(f"Saved inference prediction to {args.output_csv}")
+        return
+
+    ds = YieldDataset(
+        data_file=args.input_file,
+        weather_vars=cfg.weather_vars,
+        soil_vars=cfg.soil_vars,
+        split="all",
+        seed=1234,
+        crop=None,
+        years=None,
+        require_yield=False,
+    )
+
+    ds.set_normalization(
+        cfg.w_mean,
+        cfg.w_std,
+        cfg.s_mean,
+        cfg.s_std,
+    )
+
+    loader = DataLoader(ds, batch_size=args.batch_size, shuffle=False)
+
+    rows = []
+    sample_idx = 0
+
+    for batch in loader:
+        weather = batch["weather"].to(device)
+        soil = batch["soil"].to(device)
+        crop_id = batch["crop_id"].to(device)
+
+        t_eff = min(args.cutoff, weather.size(1))
+
+        out = model(
+            weather=weather[:, :t_eff, :],
+            soil=soil,
+            crop_id=crop_id,
+            horizon_idx=t_eff,
+            causal=True,
+            return_sequence=False,
+        )
+
+        for pred in out.predictions.detach().cpu().numpy().tolist():
+            rows.append({
+                "sample_idx": sample_idx,
+                "cutoff": int(args.cutoff),
+                "y_pred": float(pred),
+            })
+            sample_idx += 1
+
+    Path(args.output_csv).parent.mkdir(parents=True, exist_ok=True)
+    pd.DataFrame(rows).to_csv(args.output_csv, index=False)
+
+    print(f"Saved inference predictions to {args.output_csv}")
+
+
+if __name__ == "__main__":
+    main()

training_code/scripts/prepare_cornbelt.py

@@ -0,0 +1,710 @@
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+
+# ============================================================
+# Configuration
+# ============================================================
+
+INPUT_CSV = Path("khaki_multi_crop_yield.csv")
+OUTPUT_DIR = Path("data/cornbelt")
+
+TRAIN_YEARS = [2013, 2014, 2015, 2016]
+VAL_YEAR = 2017
+TEST_YEAR = 2018
+
+
+# Public dataset weather mapping
+WEATHER_MAP = {
+    "prcp": 1,
+    "srad": 2,
+    "swe": 3,
+    "tmax": 4,
+    "tmin": 5,
+    "vp": 6,
+}
+
+
+SOIL_MEASUREMENTS = [
+    "bdod",
+    "cec",
+    "cfvo",
+    "clay",
+    "nitrogen",
+    "ocd",
+    "ocs",
+    "phh2o",
+    "sand",
+    "silt",
+    "soc",
+]
+
+
+SOIL_DEPTHS = [
+    "0-5cm",
+    "5-15cm",
+    "15-30cm",
+    "30-60cm",
+    "60-100cm",
+    "100-200cm",
+]
+
+
+SOIL_VARS = [
+    f"{measurement}_mean_{depth}"
+    for measurement in SOIL_MEASUREMENTS
+    for depth in SOIL_DEPTHS
+]
+
+
+# ============================================================
+# Load
+# ============================================================
+
+print(f"Reading {INPUT_CSV}")
+
+df = pd.read_csv(INPUT_CSV)
+
+print("Raw shape:", df.shape)
+
+
+# ============================================================
+# Verify required columns
+# ============================================================
+
+if "corn_yield" not in df.columns:
+    raise ValueError("corn_yield column not found")
+
+
+required_meta = [
+    "loc_ID",
+    "year",
+    "State",
+    "County",
+    "lat",
+    "lng",
+]
+
+missing_meta = [
+    c for c in required_meta
+    if c not in df.columns
+]
+
+if missing_meta:
+    raise ValueError(
+        f"Missing metadata columns: {missing_meta}"
+    )
+
+
+missing_soil = [
+    c for c in SOIL_VARS
+    if c not in df.columns
+]
+
+if missing_soil:
+    raise ValueError(
+        f"Missing soil columns: {missing_soil}"
+    )
+
+
+# Verify all 52 weekly values exist for each weather variable
+missing_weather = []
+
+for weather_name, source_idx in WEATHER_MAP.items():
+    for week in range(1, 53):
+        src = f"W_{source_idx}_{week}"
+
+        if src not in df.columns:
+            missing_weather.append(src)
+
+if missing_weather:
+    raise ValueError(
+        f"Missing weather columns: {missing_weather[:20]}"
+    )
+
+
+# ============================================================
+# Convert year and corn yield to numeric
+# ============================================================
+
+df["year"] = pd.to_numeric(
+    df["year"],
+    errors="coerce",
+)
+
+df["corn_yield"] = pd.to_numeric(
+    df["corn_yield"],
+    errors="coerce",
+)
+
+
+# ============================================================
+# IMPORTANT:
+# Filter to corn samples from 2013-2018 BEFORE constructing out
+# ============================================================
+
+df = df[
+    (df["year"] >= 2013) &
+    (df["year"] <= 2018)
+].copy()
+
+
+# Remove samples without valid corn yield or year
+df = df.dropna(
+    subset=[
+        "corn_yield",
+        "year",
+    ]
+).reset_index(drop=True)
+
+
+df["year"] = df["year"].astype(int)
+
+
+print()
+print("===================================")
+print("FILTERED SOURCE DATA")
+print("===================================")
+
+print(
+    "Years retained:",
+    sorted(df["year"].unique())
+)
+
+print(
+    "Rows after 2013-2018 + corn yield filter:",
+    len(df)
+)
+
+print("\nSamples per year:")
+print(
+    df["year"]
+    .value_counts()
+    .sort_index()
+)
+
+
+expected_years = {
+    2013,
+    2014,
+    2015,
+    2016,
+    2017,
+    2018,
+}
+
+assert set(df["year"].unique()) == expected_years
+
+
+# ============================================================
+# Construct metadata block
+# ============================================================
+
+metadata_df = pd.DataFrame(
+    {
+        "crop": "corn",
+
+        # County/location acts as the sample spatial identifier.
+        # It is metadata required by the current dataset interface.
+        "farm_field": (
+            "county_" +
+            df["loc_ID"].astype(str)
+        ),
+
+        "year": df["year"].astype(int),
+
+        "yield": df[
+            "corn_yield"
+        ].astype(np.float32),
+
+        "loc_ID": df["loc_ID"],
+
+        "state": df[
+            "State"
+        ].astype(str),
+
+        "county": df[
+            "County"
+        ].astype(str),
+
+        "lat": pd.to_numeric(
+            df["lat"],
+            errors="coerce",
+        ).astype(np.float32),
+
+        "lng": pd.to_numeric(
+            df["lng"],
+            errors="coerce",
+        ).astype(np.float32),
+    }
+)
+
+
+# ============================================================
+# Weather block
+#
+# IMPORTANT:
+# - No interpolation
+# - No daily conversion
+# - No aggregation
+#
+# Each source W_x_1 ... W_x_52 is copied directly.
+#
+# Example:
+# W_1_1 -> prcp_0
+# W_1_2 -> prcp_1
+# ...
+# W_1_52 -> prcp_51
+# ============================================================
+
+weather_data = {}
+
+
+for weather_name, source_idx in WEATHER_MAP.items():
+
+    for week in range(1, 53):
+
+        src = f"W_{source_idx}_{week}"
+
+        # Zero-based temporal indexing used by current loader
+        dst = f"{weather_name}_{week - 1}"
+
+        weather_data[dst] = pd.to_numeric(
+            df[src],
+            errors="coerce",
+        ).astype(np.float32)
+
+
+weather_df = pd.DataFrame(
+    weather_data,
+    index=df.index,
+)
+
+
+# ============================================================
+# Soil block
+# ============================================================
+
+soil_data = {}
+
+
+for col in SOIL_VARS:
+
+    soil_data[col] = pd.to_numeric(
+        df[col],
+        errors="coerce",
+    ).astype(np.float32)
+
+
+soil_df = pd.DataFrame(
+    soil_data,
+    index=df.index,
+)
+
+
+# ============================================================
+# Combine all blocks
+# ============================================================
+
+out = pd.concat(
+    [
+        metadata_df.reset_index(drop=True),
+        weather_df.reset_index(drop=True),
+        soil_df.reset_index(drop=True),
+    ],
+    axis=1,
+)
+
+
+# Replace inf values with NaN.
+# Normalization / missing-value handling remains in train_hf.py
+# and YieldDataset.
+out = out.replace(
+    [np.inf, -np.inf],
+    np.nan,
+)
+
+
+# ============================================================
+# Critical safety check
+# ============================================================
+
+print()
+print("===================================")
+print("FINAL MODEL DATASET")
+print("===================================")
+
+print(
+    "Shape:",
+    out.shape,
+)
+
+print(
+    "Years:",
+    sorted(out["year"].unique())
+)
+
+print(
+    "Rows:",
+    len(out)
+)
+
+
+assert set(out["year"].unique()) == {
+    2013,
+    2014,
+    2015,
+    2016,
+    2017,
+    2018,
+}
+
+
+assert len(out) == len(df)
+
+
+# ============================================================
+# Chronological split
+#
+# Train: 2013-2016
+# Val:   2017
+# Test:  2018
+# ============================================================
+
+train_df = out[
+    out["year"].isin(TRAIN_YEARS)
+].reset_index(drop=True)
+
+
+val_df = out[
+    out["year"] == VAL_YEAR
+].reset_index(drop=True)
+
+
+test_df = out[
+    out["year"] == TEST_YEAR
+].reset_index(drop=True)
+
+
+# ============================================================
+# Split reporting
+# ============================================================
+
+print()
+print("===================================")
+print("SPLITS")
+print("===================================")
+
+
+print("\nTRAIN")
+
+print(
+    train_df[
+        "year"
+    ].value_counts().sort_index()
+)
+
+print(
+    "Samples:",
+    len(train_df),
+)
+
+print(
+    "Counties:",
+    train_df[
+        "farm_field"
+    ].nunique(),
+)
+
+
+print("\nVALIDATION")
+
+print(
+    val_df[
+        "year"
+    ].value_counts().sort_index()
+)
+
+print(
+    "Samples:",
+    len(val_df),
+)
+
+print(
+    "Counties:",
+    val_df[
+        "farm_field"
+    ].nunique(),
+)
+
+
+print("\nTEST")
+
+print(
+    test_df[
+        "year"
+    ].value_counts().sort_index()
+)
+
+print(
+    "Samples:",
+    len(test_df),
+)
+
+print(
+    "Counties:",
+    test_df[
+        "farm_field"
+    ].nunique(),
+)
+
+
+# ============================================================
+# Split sanity checks
+# ============================================================
+
+assert set(
+    train_df["year"].unique()
+) == {
+    2013,
+    2014,
+    2015,
+    2016,
+}
+
+
+assert set(
+    val_df["year"].unique()
+) == {
+    2017,
+}
+
+
+assert set(
+    test_df["year"].unique()
+) == {
+    2018,
+}
+
+
+assert (
+    len(train_df)
+    + len(val_df)
+    + len(test_df)
+    == len(out)
+)
+
+
+# No row should appear in more than one split
+assert set(train_df.index).isdisjoint(
+    set(range(
+        len(train_df),
+        len(train_df) + len(val_df)
+    ))
+)
+
+
+print()
+print("All year/split checks passed.")
+
+
+# ============================================================
+# Weather sanity checks
+# ============================================================
+
+weather_vars = list(
+    WEATHER_MAP.keys()
+)
+
+
+for weather_var in weather_vars:
+
+    cols = [
+        c
+        for c in out.columns
+        if c.startswith(
+            f"{weather_var}_"
+        )
+    ]
+
+    assert len(cols) == 52, (
+        f"{weather_var}: expected 52 weekly "
+        f"columns, found {len(cols)}"
+    )
+
+
+print()
+print("Weather variables:")
+print(weather_vars)
+
+print(
+    "Weeks per weather variable:",
+    52,
+)
+
+print(
+    "Total weather columns:",
+    52 * len(weather_vars),
+)
+
+
+# ============================================================
+# Soil sanity checks
+# ============================================================
+
+assert len(SOIL_VARS) == 66
+
+assert all(
+    col in out.columns
+    for col in SOIL_VARS
+)
+
+
+print(
+    "Number of soil variables:",
+    len(SOIL_VARS),
+)
+
+
+# ============================================================
+# Check expected tensor dimensions
+# ============================================================
+
+print()
+print("Expected model input dimensions:")
+
+print(
+    "weather = [52, 6]"
+)
+
+print(
+    "soil    = [66]"
+)
+
+
+# ============================================================
+# Yield summaries
+# ============================================================
+
+print()
+print("===================================")
+print("YIELD SUMMARY")
+print("===================================")
+
+
+print("\nTrain:")
+print(
+    train_df[
+        "yield"
+    ].describe()
+)
+
+
+print("\nValidation:")
+print(
+    val_df[
+        "yield"
+    ].describe()
+)
+
+
+print("\nTest:")
+print(
+    test_df[
+        "yield"
+    ].describe()
+)
+
+
+# ============================================================
+# Save HDF5 files
+# ============================================================
+
+OUTPUT_DIR.mkdir(
+    parents=True,
+    exist_ok=True,
+)
+
+
+train_path = OUTPUT_DIR / "train.h5"
+val_path = OUTPUT_DIR / "val.h5"
+test_path = OUTPUT_DIR / "test.h5"
+
+
+train_df.to_hdf(
+    train_path,
+    key="data",
+    mode="w",
+)
+
+
+val_df.to_hdf(
+    val_path,
+    key="data",
+    mode="w",
+)
+
+
+test_df.to_hdf(
+    test_path,
+    key="data",
+    mode="w",
+)
+
+
+print()
+print("===================================")
+print("SAVED")
+print("===================================")
+
+print(train_path)
+print(val_path)
+print(test_path)
+
+
+# ============================================================
+# Reload files to verify they were written correctly
+# ============================================================
+
+train_check = pd.read_hdf(train_path)
+val_check = pd.read_hdf(val_path)
+test_check = pd.read_hdf(test_path)
+
+
+assert len(train_check) == len(train_df)
+assert len(val_check) == len(val_df)
+assert len(test_check) == len(test_df)
+
+
+assert set(
+    train_check["year"].unique()
+) == {
+    2013,
+    2014,
+    2015,
+    2016,
+}
+
+
+assert set(
+    val_check["year"].unique()
+) == {
+    2017,
+}
+
+
+assert set(
+    test_check["year"].unique()
+) == {
+    2018,
+}
+
+
+print()
+print("HDF5 reload verification passed.")
+
+print()
+print("Done.")
+

training_code/scripts/train_hf.py

@@ -0,0 +1,381 @@
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from tqdm.auto import tqdm
+import time
+import argparse
+import json
+import random
+import pandas as pd
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoConfig, AutoModel
+
+from config.config import TrainConfig
+from data.dataset import YieldDataset
+from data.preprocessing import compute_x_stats, compute_y_stats
+from training.engine import evaluate
+from hf.configuration_yield import YieldConfig
+from hf.auto import register_yield_autoclass
+
+
+def parse_list(x):
+    return [v.strip() for v in x.split(",") if v.strip()]
+
+
+def parse_int_list(x):
+    return [int(v.strip()) for v in x.split(",") if v.strip()]
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--train_file", required=True)
+    p.add_argument("--val_file", default=None)
+    p.add_argument("--test_file", default=None)
+    p.add_argument("--weather_vars", required=True)
+    p.add_argument("--soil_vars", required=True)
+    p.add_argument("--crop", default=None)
+    p.add_argument("--years", default=None)
+    p.add_argument("--train_cutoffs", default="4,8,12,16,22")
+    p.add_argument("--eval_cutoffs", default=None)
+    p.add_argument("--out_dir", default="outputs/yield_hf_model")
+    p.add_argument("--epochs", type=int, default=10)
+    p.add_argument("--batch_size", type=int, default=64)
+    p.add_argument("--lr", type=float, default=3e-4)
+    p.add_argument("--weight_decay", type=float, default=1e-4)
+    p.add_argument("--seed", type=int, default=1234)
+    p.add_argument(
+        "--time_agg",
+        default="weekly_cumulative",
+        choices=["weekly", "weekly_cumulative"],
+    )
+    return p.parse_args()
+
+
+def train_one_epoch_hf(model, loader, optimizer, device, cutoffs):
+    model.train()
+
+    total_loss = 0.0
+    n_total = 0
+    se_raw = 0.0
+    n = 0
+
+    pbar = tqdm(
+        loader,
+        desc="Training",
+        leave=False,
+    )
+
+    for batch in pbar:
+        weather = batch["weather"].to(device)
+        soil = batch["soil"].to(device)
+        crop_id = batch["crop_id"].to(device)
+        labels = batch["yield"].to(device)
+
+        # t = random.choice(cutoffs)
+        # t_eff = min(t, weather.size(1))
+
+        # out = model(
+        #     weather=weather[:, :t_eff, :],
+        #     soil=soil,
+        #     crop_id=crop_id,
+        #     labels=labels,
+        #     horizon_idx=t_eff,
+        #     causal=True,
+        #     return_sequence=False,
+        # )
+
+        # loss = out.loss
+
+        losses = []
+        preds_for_rmse = []
+
+        for t in cutoffs:
+            t_eff = min(t, weather.size(1))
+
+            out = model(
+                weather=weather[:, :t_eff, :],
+                soil=soil,
+                crop_id=crop_id,
+                labels=labels,
+                horizon_idx=t_eff,
+                causal=True,
+                return_sequence=False,
+            )
+
+            losses.append(out.loss)
+            preds_for_rmse.append(out.predictions)
+
+        loss = torch.stack(losses).mean()
+
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
+        optimizer.step()
+
+        total_loss += loss.item() * labels.size(0)
+        n_total += labels.size(0)
+
+        with torch.no_grad():
+            # pred = out.predictions
+            # se_raw += ((pred - labels) ** 2).sum().item()
+            # n += labels.numel()
+            pred = torch.stack(preds_for_rmse, dim=0).mean(dim=0)
+            se_raw += ((pred - labels) ** 2).sum().item()
+            n += labels.numel()
+
+    rmse = (se_raw / max(n, 1)) ** 0.5
+
+    return {
+        "loss": total_loss / max(n_total, 1),
+        "rmse": rmse,
+    }
+
+
+def save_json(path, obj):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w") as f:
+        json.dump(obj, f, indent=2, default=str)
+
+
+def main():
+    print("Starting training script...")
+    args = parse_args()
+    register_yield_autoclass()
+
+    cfg = TrainConfig(
+        mode="train_eval",
+        train_file=Path(args.train_file),
+        val_file=Path(args.val_file) if args.val_file else None,
+        test_file=Path(args.test_file) if args.test_file else None,
+        weather_vars=parse_list(args.weather_vars),
+        soil_vars=parse_list(args.soil_vars),
+        crop=args.crop,
+        years=args.years,
+        train_cutoffs=parse_int_list(args.train_cutoffs),
+        eval_cutoffs=parse_int_list(args.eval_cutoffs) if args.eval_cutoffs else parse_int_list(args.train_cutoffs),
+        out_dir=Path(args.out_dir),
+        epochs=args.epochs,
+        batch_size=args.batch_size,
+        lr=args.lr,
+        weight_decay=args.weight_decay,
+        seed=args.seed,
+    )
+
+    cfg.validate()
+    cfg.save()
+
+    torch.manual_seed(cfg.seed)
+    random.seed(cfg.seed)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    train_ds = YieldDataset(
+        data_file=cfg.train_file,
+        weather_vars=cfg.weather_vars,
+        soil_vars=cfg.soil_vars,
+        split="train" if cfg.val_file is None else "all",
+        seed=cfg.seed,
+        crop=cfg.crop,
+        years=cfg.years,
+        split_strategy=cfg.split_strategy,
+        val_split=cfg.val_split,
+        test_split=cfg.test_split,
+        time_agg=args.time_agg,
+    )
+
+    if cfg.val_file:
+        val_ds = YieldDataset(
+            data_file=cfg.val_file,
+            weather_vars=cfg.weather_vars,
+            soil_vars=cfg.soil_vars,
+            split="all",
+            seed=cfg.seed,
+            crop=cfg.crop,
+            years=cfg.years,
+            time_agg=args.time_agg,
+        )
+    else:
+        val_ds = YieldDataset(
+            data_file=cfg.train_file,
+            weather_vars=cfg.weather_vars,
+            soil_vars=cfg.soil_vars,
+            split="val",
+            seed=cfg.seed,
+            crop=cfg.crop,
+            years=cfg.years,
+            split_strategy=cfg.split_strategy,
+            val_split=cfg.val_split,
+            test_split=cfg.test_split,
+            time_agg=args.time_agg,
+        )
+
+    w_mean, w_std, s_mean, s_std = compute_x_stats(train_ds, seed=cfg.seed)
+    y_mean, y_std = compute_y_stats(train_ds)
+
+    train_ds.set_normalization(w_mean, w_std, s_mean, s_std)
+    val_ds.set_normalization(w_mean, w_std, s_mean, s_std)
+
+    b0 = train_ds[0]
+    K, W = b0["weather"].shape
+    S = b0["soil"].shape[0]
+
+    hf_config = YieldConfig(
+        weather_vars=cfg.weather_vars,
+        soil_vars=cfg.soil_vars,
+        w_mean=w_mean.tolist(),
+        w_std=w_std.tolist(),
+        s_mean=s_mean.tolist(),
+        s_std=s_std.tolist(),
+        y_mean=float(y_mean),
+        y_std=float(y_std),
+        K=int(K),
+        W=int(W),
+        S=int(S),
+        train_cutoffs=cfg.train_cutoffs,
+        eval_cutoffs=cfg.eval_cutoffs,
+        d_model=cfg.d_model,
+        nhead=cfg.nhead,
+        num_layers=cfg.num_layers,
+        dim_ff=cfg.dim_ff,
+        dropout=cfg.dropout,
+        pool=cfg.pool,
+        use_crop=cfg.use_crop,
+        crop_emb_dim=cfg.crop_emb_dim,
+    )
+
+    model = AutoModel.from_config(hf_config).to(device)
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=cfg.lr,
+        weight_decay=cfg.weight_decay,
+    )
+
+    train_loader = DataLoader(train_ds, batch_size=cfg.batch_size, shuffle=True, num_workers=0)
+    val_loader = DataLoader(val_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=0)
+
+    best_val = float("inf")
+    best_metrics = None
+    patience = 5
+    bad_epochs = 0
+
+    for epoch in range(1, cfg.epochs + 1):
+        epoch_start = time.time()
+
+        train_metrics = train_one_epoch_hf(
+            model=model,
+            loader=train_loader,
+            optimizer=optimizer,
+            device=device,
+            cutoffs=cfg.train_cutoffs,
+        )
+
+        val_metrics_by_t, _ = evaluate(
+            model=model,
+            loader=val_loader,
+            device=device,
+            y_mean=y_mean,
+            y_std=y_std,
+            cutoffs=cfg.eval_cutoffs,
+        )
+
+        epoch_time = time.time() - epoch_start
+        remaining = (cfg.epochs - epoch) * epoch_time
+
+        print(
+            f"\nEpoch {epoch}/{cfg.epochs}"
+            f" | time={epoch_time/60:.1f} min"
+            f" | ETA={remaining/60:.1f} min"
+        )
+        val_rmse = sum(m["rmse"] for m in val_metrics_by_t.values()) / len(val_metrics_by_t)
+
+        print(
+            f"Epoch {epoch:03d} | "
+            f"train_rmse={train_metrics['rmse']:.4f} | "
+            f"val_rmse={val_rmse:.4f}"
+        )
+
+        print("\nValidation Metrics")
+        print("-" * 42)
+        print(f"{'Cutoff':>8} {'RMSE':>10} {'Bias':>10} {'R²':>10}")
+
+        for cutoff in cfg.eval_cutoffs:
+            m = val_metrics_by_t[cutoff]
+            print(
+                f"{cutoff:>8}"
+                f"{m['rmse']:>10.3f}"
+                f"{m['bias']:>10.3f}"
+                f"{m['r2']:>10.3f}"
+            )
+
+        if val_rmse < best_val:
+            bad_epochs = 0
+            best_val = val_rmse
+            best_metrics = {
+                "epoch": epoch,
+                "val_rmse": val_rmse,
+                "val_by_cutoff": val_metrics_by_t,
+            }
+
+            model.save_pretrained(cfg.out_dir)
+            hf_config.save_pretrained(cfg.out_dir)
+            save_json(cfg.out_dir / "metrics.json", best_metrics)
+
+            print(f"Saved best HF model to {cfg.out_dir}")
+        else:
+            bad_epochs += 1
+            print(f"No improvement for {bad_epochs}/{patience} epochs.")
+
+            if bad_epochs >= patience:
+                print("Early stopping.")
+                break
+
+    print("Best metrics:")
+    print(json.dumps(best_metrics, indent=2))
+
+    if cfg.test_file is not None:
+        print("\nRunning final test using best HF model...")
+
+        best_model = AutoModel.from_pretrained(cfg.out_dir).to(device)
+        best_model.eval()
+
+        test_ds = YieldDataset(
+            data_file=cfg.test_file,
+            weather_vars=cfg.weather_vars,
+            soil_vars=cfg.soil_vars,
+            split="all",
+            seed=cfg.seed,
+            crop=cfg.crop,
+            years=cfg.years,
+            time_agg=args.time_agg,
+        )
+
+        test_ds.set_normalization(w_mean, w_std, s_mean, s_std)
+
+        test_loader = DataLoader(test_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=0)
+
+        test_metrics, test_rows = evaluate(
+            model=best_model,
+            loader=test_loader,
+            device=device,
+            y_mean=y_mean,
+            y_std=y_std,
+            cutoffs=cfg.eval_cutoffs,
+        )
+
+        print("Final test metrics:")
+        print(json.dumps(test_metrics, indent=2))
+
+        save_json(cfg.out_dir / "test_metrics.json", test_metrics)
+        pd.DataFrame(test_rows).to_csv(cfg.out_dir / "test_predictions.csv", index=False)
+
+        print(f"Saved final test outputs to {cfg.out_dir}")
+
+
+if __name__ == "__main__":
+    main()

training_code/training.slurm

@@ -0,0 +1,28 @@
+#!/bin/bash
+#SBATCH --job-name=yield-cornbelt
+#SBATCH --mail-type=ALL
+#SBATCH --mail-user=sridhar.86@buckeyemail.osu.edu
+#SBATCH --time=08:00:00
+#SBATCH --account=PAS2699
+#SBATCH --output=yield-training.%j.out
+#SBATCH --mem=100G
+#SBATCH --gpus-per-node=1
+
+module load python/3.12
+
+source activate yield_hf
+
+python scripts/train_hf.py \
+  --train_file data/cornbelt/train.h5 \
+  --val_file data/cornbelt/val.h5 \
+  --test_file data/cornbelt/test.h5 \
+  --weather_vars prcp,srad,swe,tmax,tmin,vp \
+  --soil_vars bdod_mean_0-5cm,bdod_mean_5-15cm,bdod_mean_15-30cm,bdod_mean_30-60cm,bdod_mean_60-100cm,bdod_mean_100-200cm,cec_mean_0-5cm,cec_mean_5-15cm,cec_mean_15-30cm,cec_mean_30-60cm,cec_mean_60-100cm,cec_mean_100-200cm,cfvo_mean_0-5cm,cfvo_mean_5-15cm,cfvo_mean_15-30cm,cfvo_mean_30-60cm,cfvo_mean_60-100cm,cfvo_mean_100-200cm,clay_mean_0-5cm,clay_mean_5-15cm,clay_mean_15-30cm,clay_mean_30-60cm,clay_mean_60-100cm,clay_mean_100-200cm,nitrogen_mean_0-5cm,nitrogen_mean_5-15cm,nitrogen_mean_15-30cm,nitrogen_mean_30-60cm,nitrogen_mean_60-100cm,nitrogen_mean_100-200cm,ocd_mean_0-5cm,ocd_mean_5-15cm,ocd_mean_15-30cm,ocd_mean_30-60cm,ocd_mean_60-100cm,ocd_mean_100-200cm,ocs_mean_0-5cm,ocs_mean_5-15cm,ocs_mean_15-30cm,ocs_mean_30-60cm,ocs_mean_60-100cm,ocs_mean_100-200cm,phh2o_mean_0-5cm,phh2o_mean_5-15cm,phh2o_mean_15-30cm,phh2o_mean_30-60cm,phh2o_mean_60-100cm,phh2o_mean_100-200cm,sand_mean_0-5cm,sand_mean_5-15cm,sand_mean_15-30cm,sand_mean_30-60cm,sand_mean_60-100cm,sand_mean_100-200cm,silt_mean_0-5cm,silt_mean_5-15cm,silt_mean_15-30cm,silt_mean_30-60cm,silt_mean_60-100cm,silt_mean_100-200cm,soc_mean_0-5cm,soc_mean_5-15cm,soc_mean_15-30cm,soc_mean_30-60cm,soc_mean_60-100cm,soc_mean_100-200cm \
+  --crop corn \
+  --time_agg weekly \
+  --train_cutoffs 20,24,28,32,36,40,44,48,52 \
+  --eval_cutoffs 20,24,28,32,36,40,44,48,52 \
+  --epochs 30 \
+  --lr 3e-5 \
+  --batch_size 32 \
+  --out_dir checkpoints

training_code/training/engine.py

@@ -0,0 +1,122 @@
+import math
+import random
+import torch
+import torch.nn.functional as F
+import numpy as np
+from sklearn.metrics import r2_score
+from tqdm.auto import tqdm
+
+def train_one_epoch(model, loader, optimizer, device, y_mean, y_std, cutoffs):
+    model.train()
+    total_loss = 0.0
+    n_total = 0
+
+    se_norm = 0.0
+    se_raw = 0.0
+    n = 0
+
+    for batch in loader:
+        weather = batch["weather"].to(device)
+        soil = batch["soil"].to(device)
+        crop_id = batch["crop_id"].to(device)
+        y = batch["yield"].to(device)
+
+        t = random.choice(cutoffs)
+        t_eff = min(t, weather.size(1))
+
+        y_log = torch.log1p(torch.clamp(y, min=0.0))
+        y_norm = (y_log - y_mean) / y_std
+
+        yhat_norm = model(
+            weather[:, :t_eff, :],
+            soil,
+            crop_id,
+            horizon_idx=t_eff,
+            causal=True,
+            return_sequence=False,
+        )
+
+        loss = F.mse_loss(yhat_norm, y_norm)
+
+        optimizer.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
+        optimizer.step()
+
+        total_loss += loss.item() * y.size(0)
+        n_total += y.size(0)
+
+        with torch.no_grad():
+            yhat = torch.expm1(yhat_norm * y_std + y_mean)
+            se_norm += ((yhat_norm - y_norm) ** 2).sum().item()
+            se_raw += ((yhat - y) ** 2).sum().item()
+            n += y.numel()
+
+    return {
+        "loss_norm_mse": total_loss / max(n_total, 1),
+        "norm_rmse": math.sqrt(se_norm / max(n, 1)),
+        "rmse": math.sqrt(se_raw / max(n, 1)),
+    }
+
+@torch.no_grad()
+def evaluate(model, loader, device, y_mean, y_std, cutoffs):
+    model.eval()
+    results = {}
+    all_rows = []
+
+    for t in sorted(cutoffs):
+        y_true_all = []
+        y_pred_all = []
+
+        pbar = tqdm(
+            loader,
+            desc="Validation",
+            leave=False,
+        )
+
+        for batch in pbar:
+            weather = batch["weather"].to(device)
+            soil = batch["soil"].to(device)
+            crop_id = batch["crop_id"].to(device)
+            y = batch["yield"].to(device)
+
+            t_eff = min(t, weather.size(1))
+
+            out = model(
+                weather=weather[:, :t_eff, :],
+                soil=soil,
+                crop_id=crop_id,
+                horizon_idx=t_eff,
+                causal=True,
+                return_sequence=False,
+            )
+
+            if hasattr(out, "predictions"):
+                yhat = out.predictions
+            else:
+                #yhat = torch.expm1(out * y_std + y_mean)
+                yhat = out * y_std + y_mean
+
+            y_true_all.extend(y.detach().cpu().numpy().tolist())
+            y_pred_all.extend(yhat.detach().cpu().numpy().tolist())
+
+        y_true = np.asarray(y_true_all, dtype=np.float32)
+        y_pred = np.asarray(y_pred_all, dtype=np.float32)
+        err = y_pred - y_true
+
+        results[int(t)] = {
+            "rmse": float(np.sqrt(np.mean(err ** 2))),
+            "bias": float(np.mean(err)),
+            "r2": float(r2_score(y_true, y_pred)) if len(y_true) > 1 else float("nan"),
+        }
+
+        for i, (yt, yp) in enumerate(zip(y_true, y_pred)):
+            all_rows.append({
+                "cutoff": int(t),
+                "sample_idx": i,
+                "y_true": float(yt),
+                "y_pred": float(yp),
+                "error": float(yp - yt),
+            })
+
+    return results, all_rows