Create forecast.py

app/forecast.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+import pandas as pd
+import numpy as np
+from datetime import datetime
+
+try:
+    from prophet import Prophet
+except Exception:
+    Prophet = None
+
+try:
+    from neuralprophet import NeuralProphet
+except Exception:
+    NeuralProphet = None
+
+from . import storage
+
+class SeasonalityModel:
+    def __init__(self, campaign_id: str):
+        self.campaign_id = campaign_id
+        self.model_type = None
+        self.model = None
+        self.global_mean = 0.05
+
+    def fit(self):
+        with storage.get_conn() as con:
+            df = pd.read_sql_query(
+                """
+                SELECT ts, event_type, variant_id FROM events WHERE campaign_id=?
+                """,
+                con, params=(self.campaign_id,)
+            )
+        if df.empty or Prophet is None:
+            self.model_type = "none"
+            return 
+        df["ts"] = pd.to_datetime(df["ts"], errors="coerce")
+        df = df.dropna(subset=["ts"])
+
+        df["hour"] = df["ts"].dt.floor("h")
+        agg = df.pivot_table(index="hour", columns="event_type", values="variant_id", aggfunc="count").fillna(0)
+        if "impression" not in agg:
+            agg["impression"] = 0
+        if "click" not in agg:
+            agg["click"] = 0
+        agg["ctr"] = np.where(agg["impression"] > 0, agg["click"]/agg["impression"],np.nan)
+        agg["ctr"].fillna(agg["ctr"].mean()) if not np.isna(agg["ctr"].mean())else self.global_mean
+
+        ds = agg.index.to_series().reset_index(drop=True)
+        train = pd.DataFrame({"ds":ds, "y":agg["ctr"].values})
+
+        if Prophet is not None:
+            m = Prophet(weekly_seasonality=True, daily_seasonality=True)
+            m.fit(train)
+            self.model = m
+            self.model_type = "prophet"
+        elif NeuralProphet is not None:
+            m = NeuralProphet(weekly_seasonality=True, daily_seasonality=True)
+            m.fit(train, freq='H')
+            self.model = m
+            self.model_type = "neuralprophet"
+        else:
+            self.model_type = "none"
+
+    def expected_ctr(self, context: dict) -> float:
+        if self.model_type in {None, "none"}:
+            hour = int(context.get("hour", 12))
+            base = self.global_mean
+
+            if 11 <= hour <= 13:
+                return min(0.99, base * 1.1)
+            if 20 <= hour <= 23:
+                return min(0.99, base*1.15)
+            return max(0.01, base)
+        else:
+            now = pd.Timestamp.utcnow().floor('h')
+            ds = now
+
+            if "hour" in context:
+                ds = pd.Timestamp.utcnow().floor('D') + pd.Timedelta(hours=int(context["hour"]))
+            df = pd.DataFrame({"ds":[ds]})
+            if self.model_type == "prophet"
+                yhat = self.model.predict(df)["yhat"].iloc[0]
+            else:
+                yhat = self.model.predict(df)["yhat1"].iloc[0]
+            return float(max(0.01, min(0.99,yhat)))