Upload submission.py

submission.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+import numpy as np
+import pandas as pd
+from sklearn.ensemble import GradientBoostingRegressor
+import emflow as ef
+
+
+class QuantileRegressionPredictor(ef.Predictor):
+    def __init__(self, name="quantile-regression-ar"):
+        self.name = name
+        self.lags = [1, 24]
+        self.quantile = 0.5
+        self.models = {}
+
+    def _prepare_features(self, df: pd.DataFrame, col: str):
+        series = df[col]
+        # Use timezone-normalized timestamps to derive time features
+        # Note: input_df index is expected to be a DatetimeIndex
+        hour = df.index.hour
+        day_of_week = df.index.dayofweek
+        month = df.index.month
+
+        X = pd.DataFrame(
+            {
+                "hour": hour,
+                "day_of_week": day_of_week,
+                "month": month,
+                "lag_1h": series.shift(1),
+                "lag_24h": series.shift(24),
+            },
+            index=df.index,
+        )
+        return X
+
+    def train(self, train_df: pd.DataFrame):
+        if isinstance(train_df, pd.Series):
+            train_df = train_df.to_frame()
+
+        # Filter for Stockholm-Observatoriekullen A as per user's specific interest
+        # though the contract implies handling all columns passed.
+        # Given the "swedish-temperatures:ar" context, we handle all.
+        for col in train_df.columns:
+            X = self._prepare_features(train_df, col)
+            y = train_df[col]
+
+            # Combine and drop NaNs
+            data = pd.concat([X, y], axis=1).dropna()
+            if len(data) < 100:  # Heuristic for sufficient data
+                self.models[col] = None
+                continue
+
+            X_train = data.drop(columns=[col])
+            y_train = data[col]
+
+            model = GradientBoostingRegressor(
+                loss="quantile",
+                alpha=self.quantile,
+                n_estimators=100,
+                max_depth=5,
+                random_state=42,
+            )
+            model.fit(X_train, y_train)
+            self.models[col] = model
+        return self
+
+    def predict(self, input_df: pd.DataFrame):
+        if isinstance(input_df, pd.Series):
+            input_df = input_df.to_frame()
+
+        preds = {}
+        for col in input_df.columns:
+            model = self.models.get(col)
+            if model is None:
+                preds[col] = np.full(len(input_df), np.nan)
+                continue
+
+            X = self._prepare_features(input_df, col)
+            # We can't dropna here because we need a value for the last row
+            # even if it has NaNs in current values (but lags should be there)
+            # The contract says last timestamp is the target to forecast.
+            # Only the current value at last timestamp is NaN. Lags should be fine.
+
+            # Fill NaNs in features with a neutral value or previous if necessary
+            # but usually for the last row, shift(1) of NaN is the value at T-1.
+            out = np.full(len(input_df), np.nan)
+
+            # Identify rows where we have all features
+            valid_mask = X.notna().all(axis=1)
+            if valid_mask.any():
+                out[valid_mask] = model.predict(X[valid_mask])
+
+            preds[col] = out
+
+        return pd.DataFrame(preds, index=input_df.index, columns=input_df.columns)
+
+
+model = QuantileRegressionPredictor()