# Source: content/notes/ml/linear-models.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.linear_model import PoissonRegressor
from sklearn.metrics import mean_poisson_deviance
from sklearn.model_selection import train_test_split

rng = np.random.default_rng(4)
X = rng.normal(size=(500, 2))
exposure = rng.uniform(0.5, 3.0, size=500)
rate = np.exp(0.3 + X @ np.array([0.5, -0.25]))
counts = rng.poisson(exposure * rate)
train, test = train_test_split(np.arange(500), random_state=4)
model = PoissonRegressor(alpha=0.01, max_iter=500)
model.fit(X[train], counts[train] / exposure[train],
          sample_weight=exposure[train])
pred = exposure[test] * model.predict(X[test])
baseline = exposure[test] * counts[train].sum() / exposure[train].sum()
assert (pred > 0).all()
assert mean_poisson_deviance(counts[test], pred) < mean_poisson_deviance(counts[test], baseline)
print("Poisson deviance:", mean_poisson_deviance(counts[test], pred))