Demand Forecasting Accelerator

src/model/train_model.py# Demand Forecasting Accelerator: Model Training # ============================================== # Trains a Gradient Boosting demand model on historical consumption + weather # and registers it in the MLflow Model Registry. Fully generic — all catalog/ # schema/table names come from bundle variables. The feature engineering mirrors # lib/demand_features.py (calendar + cyclical + degree-days + lags). # Databricks notebook source import mlflow import mlflow.sklearn import numpy as np import pandas as pd from datetime import datetime from sklearn.ensemble import GradientBoostingRegressor from sklearn.metrics import mean_absolute_error, mean_squared_error # COMMAND ---------- catalog = spark.conf.get("bundle.var.catalog") schema = spark.conf.get("bundle.var.schema") model_name = spark.conf.get("bundle.var.model_name", "demand-forecast") # COMMAND ---------- # Load 2 years of hourly consumption + weather. Adapt column names to your source. consumption = spark.sql(f""" SELECT start_time AS timestamp, SUM(energy_mwh) AS actual_mwh FROM IDENTIFIER('{catalog}.{schema}.hourly_consumption') WHERE start_time >= DATEADD(YEAR, -2, CURRENT_TIMESTAMP()) GROUP BY start_time """).toPandas() weather = spark.sql(f""" SELECT timestamp, temperature_c, wind_speed_ms, cloud_cover_pct FROM IDENTIFIER('{catalog}.{schema}.weather_observations') WHERE timestamp >= DATEADD(YEAR, -2, CURRENT_TIMESTAMP()) """).toPandas() # COMMAND ---------- # Feature engineering (see lib/demand_features.py for the reference logic + tests) df = consumption.merge(weather, on="timestamp", how="left").sort_values("timestamp") df["hour"] = df["timestamp"].dt.hour