299 lines
9.6 KiB
Python
299 lines
9.6 KiB
Python
"""
|
|
Museum visitors vs. city population: regression model + FastAPI service.
|
|
|
|
Run locally (from the repo root):
|
|
|
|
uvicorn api.museum_api:app --reload
|
|
|
|
or run this file directly (e.g. from the IDE):
|
|
|
|
python api/museum_api.py
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
if __package__ in (None, ""):
|
|
# Started as a script (python museum_api.py): put the repo root on the
|
|
# path so the `api` package imports below resolve.
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
from dataclasses import dataclass, field
|
|
from typing import Literal
|
|
import numpy as np
|
|
import pandas as pd
|
|
from app import data_setup, constants
|
|
from fastapi import FastAPI, HTTPException, Query
|
|
from pydantic import BaseModel, Field
|
|
from sklearn.linear_model import LinearRegression
|
|
from sklearn.model_selection import KFold, cross_val_score
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
REQUIRED_COLUMNS = ["Name", "Visitors", "City", "Country", "Population"]
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Data + model
|
|
# --------------------------------------------------------------------------- #
|
|
def clean(df: pd.DataFrame) -> pd.DataFrame:
|
|
"""Validate columns, coerce numerics, drop rows that can't be modelled."""
|
|
missing = set(REQUIRED_COLUMNS) - set(df.columns)
|
|
if missing:
|
|
raise ValueError(f"Missing columns: {sorted(missing)}")
|
|
|
|
out = df[REQUIRED_COLUMNS].copy()
|
|
for col in ("Visitors", "Population"):
|
|
# Handles strings like "1,234,567"
|
|
out[col] = pd.to_numeric(
|
|
out[col].astype(str).str.replace(r"[,\s]", "", regex=True),
|
|
errors="coerce",
|
|
)
|
|
out = out.dropna(subset=["Visitors", "Population"])
|
|
out = out[(out["Visitors"] > 0) & (out["Population"] > 0)]
|
|
return out.reset_index(drop=True)
|
|
|
|
|
|
@dataclass
|
|
class FittedModel:
|
|
reg: LinearRegression
|
|
scale: Literal["log", "linear"]
|
|
data: pd.DataFrame # the rows actually used for fitting
|
|
n: int
|
|
r2: float
|
|
cv_r2: float | None
|
|
pearson_r: float
|
|
spearman_rho: float
|
|
aggregate_by_city: bool
|
|
predictions: np.ndarray = field(repr=False)
|
|
|
|
def _x(self, population: np.ndarray) -> np.ndarray:
|
|
x = np.asarray(population, dtype=float).reshape(-1, 1)
|
|
return np.log10(x) if self.scale == "log" else x
|
|
|
|
def predict(self, population) -> np.ndarray:
|
|
y = self.reg.predict(self._x(population))
|
|
return 10 ** y if self.scale == "log" else y
|
|
|
|
@property
|
|
def slope(self) -> float:
|
|
return float(self.reg.coef_[0])
|
|
|
|
@property
|
|
def intercept(self) -> float:
|
|
return float(self.reg.intercept_)
|
|
|
|
def summary(self) -> dict:
|
|
s = {
|
|
"n_samples": self.n,
|
|
"scale": self.scale,
|
|
"aggregate_by_city": self.aggregate_by_city,
|
|
"slope": self.slope,
|
|
"intercept": self.intercept,
|
|
"r2": self.r2,
|
|
"cv_r2_5fold": self.cv_r2,
|
|
"pearson_r": self.pearson_r,
|
|
"spearman_rho": self.spearman_rho,
|
|
}
|
|
if self.scale == "log":
|
|
s["equation"] = (
|
|
f"visitors = {10 ** self.intercept:,.1f} * population^{self.slope:.3f}"
|
|
)
|
|
s["interpretation"] = (
|
|
f"A 10x larger city is associated with "
|
|
f"{10 ** self.slope:.2f}x the visitors."
|
|
)
|
|
else:
|
|
s["equation"] = f"visitors = {self.slope:.4f} * population + {self.intercept:,.0f}"
|
|
s["interpretation"] = (
|
|
f"Each additional 1M residents is associated with "
|
|
f"{self.slope * 1e6:,.0f} more visitors."
|
|
)
|
|
return s
|
|
|
|
|
|
def fit_model(
|
|
df: pd.DataFrame,
|
|
scale: Literal["log", "linear"] = "log",
|
|
aggregate_by_city: bool = False,
|
|
) -> FittedModel:
|
|
"""
|
|
Fit visitors ~ population.
|
|
|
|
scale="log" (default) fits log10(visitors) ~ log10(population). Both
|
|
variables are heavily right-skewed, so a log-log fit is usually far
|
|
better behaved than a linear one and the slope reads as an elasticity.
|
|
|
|
aggregate_by_city=True sums visitors per (City, Country) first, so a
|
|
city with many museums counts once rather than once per museum.
|
|
"""
|
|
data = clean(df)
|
|
if aggregate_by_city:
|
|
data = (
|
|
data.groupby(["City", "Country"], as_index=False)
|
|
.agg(Visitors=("Visitors", "sum"),
|
|
Population=("Population", "first"),
|
|
Name=("Name", lambda s: ", ".join(s)))
|
|
)
|
|
if len(data) < 3:
|
|
raise ValueError(f"Need at least 3 usable rows, got {len(data)}")
|
|
|
|
x_raw = data["Population"].to_numpy(float)
|
|
y_raw = data["Visitors"].to_numpy(float)
|
|
X = (np.log10(x_raw) if scale == "log" else x_raw).reshape(-1, 1)
|
|
y = np.log10(y_raw) if scale == "log" else y_raw
|
|
|
|
reg = LinearRegression().fit(X, y)
|
|
|
|
cv_r2 = None
|
|
if len(data) >= 10:
|
|
cv = KFold(n_splits=5, shuffle=True, random_state=0)
|
|
cv_r2 = float(cross_val_score(LinearRegression(), X, y, cv=cv, scoring="r2").mean())
|
|
|
|
xs, ys = pd.Series(X.ravel()), pd.Series(y)
|
|
model = FittedModel(
|
|
reg=reg,
|
|
scale=scale,
|
|
data=data,
|
|
n=len(data),
|
|
r2=float(reg.score(X, y)),
|
|
cv_r2=cv_r2,
|
|
pearson_r=float(xs.corr(ys)),
|
|
# Spearman = Pearson on ranks (avoids a scipy dependency)
|
|
spearman_rho=float(pd.Series(x_raw).rank().corr(pd.Series(y_raw).rank())),
|
|
aggregate_by_city=aggregate_by_city,
|
|
predictions=np.empty(0),
|
|
)
|
|
model.predictions = model.predict(x_raw)
|
|
return model
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# API schemas
|
|
# --------------------------------------------------------------------------- #
|
|
class MuseumRecord(BaseModel):
|
|
Name: str
|
|
Visitors: float = Field(gt=0)
|
|
City: str
|
|
Country: str
|
|
Population: float = Field(gt=0)
|
|
|
|
|
|
class TrainRequest(BaseModel):
|
|
records: list[MuseumRecord] = Field(min_length=3)
|
|
scale: Literal["log", "linear"] = "log"
|
|
aggregate_by_city: bool = False
|
|
|
|
|
|
class PredictRequest(BaseModel):
|
|
populations: list[float] = Field(min_length=1)
|
|
|
|
|
|
class Prediction(BaseModel):
|
|
population: float
|
|
predicted_visitors: float
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# App factory
|
|
# --------------------------------------------------------------------------- #
|
|
def create_app(
|
|
df: pd.DataFrame | None = None,
|
|
scale: Literal["log", "linear"] = "log",
|
|
aggregate_by_city: bool = False,
|
|
) -> FastAPI:
|
|
app = FastAPI(title="Museum Visitors Regression", version="1.0")
|
|
app.state.model = fit_model(df, scale, aggregate_by_city) if df is not None else None
|
|
|
|
def get_model() -> FittedModel:
|
|
if app.state.model is None:
|
|
raise HTTPException(503, "No model trained yet. POST /train first.")
|
|
return app.state.model
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "model_loaded": app.state.model is not None}
|
|
|
|
@app.get("/model")
|
|
def model_summary():
|
|
return get_model().summary()
|
|
|
|
@app.post("/train")
|
|
def train(req: TrainRequest):
|
|
df_new = pd.DataFrame([r.model_dump() for r in req.records])
|
|
try:
|
|
app.state.model = fit_model(df_new, req.scale, req.aggregate_by_city)
|
|
except ValueError as e:
|
|
raise HTTPException(422, str(e))
|
|
return app.state.model.summary()
|
|
|
|
@app.get("/predict", response_model=Prediction)
|
|
def predict_one(population: float = Query(gt=0)):
|
|
m = get_model()
|
|
return Prediction(population=population,
|
|
predicted_visitors=float(m.predict([population])[0]))
|
|
|
|
@app.post("/predict", response_model=list[Prediction])
|
|
def predict_many(req: PredictRequest):
|
|
if any(p <= 0 for p in req.populations):
|
|
raise HTTPException(422, "Populations must be > 0")
|
|
m = get_model()
|
|
preds = m.predict(req.populations)
|
|
return [Prediction(population=p, predicted_visitors=float(v))
|
|
for p, v in zip(req.populations, preds)]
|
|
|
|
@app.get("/data")
|
|
def data():
|
|
"""Every row the model was fitted on, with its prediction."""
|
|
m = get_model()
|
|
d = m.data.copy()
|
|
d["Predicted"] = m.predictions
|
|
return d.to_dict(orient="records")
|
|
|
|
@app.get("/residuals")
|
|
def residuals(
|
|
top: int = Query(10, ge=1, le=500),
|
|
order: Literal["over", "under"] = "over",
|
|
):
|
|
"""Museums that most over/under-perform what their city size predicts."""
|
|
m = get_model()
|
|
d = m.data.copy()
|
|
d["Predicted"] = m.predictions
|
|
d["Ratio"] = d["Visitors"] / d["Predicted"] # >1 = outperforms city size
|
|
d = d.sort_values("Ratio", ascending=(order == "under")).head(top)
|
|
return d[["Name", "City", "Country", "Population",
|
|
"Visitors", "Predicted", "Ratio"]].to_dict(orient="records")
|
|
|
|
return app
|
|
|
|
|
|
def _load_data() -> pd.DataFrame | None:
|
|
path = constants.MUSEUM_DATA_FILE
|
|
|
|
if not path.exists():
|
|
print("Fetching and conforming museum data and city population...")
|
|
museum_data = data_setup.get_museum_data()
|
|
museum_data.to_csv(path)
|
|
else:
|
|
museum_data = pd.read_csv(path)
|
|
|
|
return museum_data
|
|
|
|
|
|
# Module-level app so `uvicorn api.museum_api:app` works.
|
|
app = create_app(
|
|
_load_data(),
|
|
scale=os.environ.get("MUSEUM_SCALE", "log"), # type: ignore[arg-type]
|
|
aggregate_by_city=os.environ.get("MUSEUM_AGG_CITY", "0") == "1",
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(
|
|
app,
|
|
host=os.environ.get("HOST", "127.0.0.1"),
|
|
port=int(os.environ.get("PORT", "8000")),
|
|
) |