from stats.glm import glm_irls, FAM_BINOMIAL, LINK_PROBIT
from stats.special import ndtr
#: mtcars (R's datasets package): miles per gallon and engine shape, 32 cars
MPG = [21.0, 21.0, 22.8, 21.4, 18.7, 18.1, 14.3, 24.4, 22.8, 19.2, 17.8, 16.4, 17.3, 15.2, 10.4, 10.4,
14.7, 32.4, 30.4, 33.9, 21.5, 15.5, 15.2, 13.3, 19.2, 27.3, 26.0, 30.4, 15.8, 19.7, 15.0, 21.4]
VS = [0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0,
0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]
#: R 4.3.3, glm(vs ~ mpg, binomial(link = "probit"), epsilon = 1e-12)
R_INTERCEPT: Num = -5.09362552650073
R_SLOPE: Num = 0.246136214545281
R_DEVIANCE: Num = 25.4344760052724
#: R's convergence test is on the deviance, so the coefficients agree to about its square root
COEF_TOL: Num = 1e-6
DEVIANCE_TOL: Num = 1e-10
EPSILON: Num = 1e-12
MAX_IT: Int = 50
#: the curve's grid
MPG_LO: Num = 10.0
MPG_HI: Num = 34.0
N_GRID: Int = 97
n = len(MPG)
X = ones(n, 2)
X[:, 1] = MPG
coef, cov, dev, nulldev, iters, conv, mu, eta, w, rank = glm_irls(VS, X, ones(n), zeros(n), FAM_BINOMIAL,
LINK_PROBIT, 1.0, EPSILON, MAX_IT)
show coef
show dev
grid = linspace(MPG_LO, MPG_HI, N_GRID)
p = [ndtr(coef[0] + coef[1] * m) for m in grid]
figure("P(vs = 1) against mpg, probit fit", "mpg", "probability")
plot(grid, p, "fitted")
scatter(MPG, VS, "mtcars")
test "the coefficients and deviance are R's" {
assert_near(coef[0], R_INTERCEPT, COEF_TOL)
assert_near(coef[1], R_SLOPE, COEF_TOL)
assert_near(dev, R_DEVIANCE, DEVIANCE_TOL)
}
# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors