Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add logistic regression <-> Variational Inference #156

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Add logistic regression <-> Variational Inference #156

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3ac63ff
Add `logistic_reg_vi()` function
Nandini-Jaiswal Dec 8, 2024
4106faa
Merge branch 'main' into add-logistic-regression-for-vi
Nandini-Jaiswal Dec 19, 2024
0fca4c8
Change `getters`
Nandini-Jaiswal Dec 25, 2024
312a3a1
Add docs
Nandini-Jaiswal Dec 28, 2024
80f0622
Fix `argument` position
Nandini-Jaiswal Dec 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions src/bayesian/getter.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,15 @@ function predict(container::BayesianRegressionMCMC{:LogisticRegression}, newdata
return vec(mean(container.link.link_function.(z), dims=2))
end

function predict(container::BayesianRegressionVI{:LogisticRegression}, newdata::DataFrame, number_of_samples::Int64 = 1000)
X = modelmatrix(container.formula, newdata)

W = rand(CRRao_rng, container.dist, number_of_samples)
W = W[union(container.symbol_to_range[:β]...), :]
z = X * W
return vec(mean(container.link.link_function.(z), dims=2))
end

function predict(container::BayesianRegressionMCMC{:NegativeBinomialRegression}, newdata::DataFrame, prediction_chain_start::Int64 = 200)
X = modelmatrix(container.formula, newdata)

Expand Down
118 changes: 102 additions & 16 deletions src/bayesian/logistic_regression.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::CRRaoLink, turingModel::Function, sim_size::Int64)
function logistic_reg_mcmc(formula::FormulaTerm, data::DataFrame, Link::CRRaoLink, turingModel::Function, sim_size::Int64)
formula = apply_schema(formula, schema(formula, data),RegressionModel)
y, X = modelcols(formula, data)

Expand All @@ -9,9 +9,31 @@ function logistic_reg(formula::FormulaTerm, data::DataFrame, Link::CRRaoLink, tu
return BayesianRegressionMCMC(:LogisticRegression, chain, formula, Link)
end

function logistic_reg_vi(formula::FormulaTerm, data::DataFrame, Link::CRRaoLink, turingModel::Function, max_iter::Int64)
formula = apply_schema(formula, schema(formula, data),RegressionModel)
y, X = modelcols(formula, data)

if max_iter < 500
@warn "Max iterations should generally be atleast 500."
end
model = turingModel(X, y)
dist = vi(model, ADVI(100, max_iter))
_, symbol_to_range = bijector(model, Val(true))
return BayesianRegressionVI(:LinearRegression, dist, formula, symbol_to_range)
end

"""
```julia
fit(formula::FormulaTerm, data::DataFrame, modelClass::LogisticRegression, Link::CRRaoLink, prior::Prior_Ridge, h::Float64 = 0.1, level::Float64 = 0.95, sim_size::Int64 = 1000)
fit(
formula::FormulaTerm,
data::DataFrame,
modelClass::LogisticRegression,
Link::CRRaoLink, prior::Prior_Ridge,
use_vi::Bool = false,
h::Float64 = 0.1,
level::Float64 = 0.95,
sim_size::Int64 = use_vi ? 20000 : 1000
)
```

Fit a Bayesian Logistic Regression model on the input data with a Ridge prior with the provided `Link` function.
Expand Down Expand Up @@ -185,9 +207,10 @@ function fit(
modelClass::LogisticRegression,
Link::CRRaoLink,
prior::Prior_Ridge,
use_vi::Bool = false,
h::Float64 = 0.1,
level::Float64 = 0.95,
sim_size::Int64 = 1000
sim_size::Int64 = use_vi ? 20000 : 1000
)
@model LogisticRegression(X, y) = begin
p = size(X, 2)
Expand All @@ -210,12 +233,26 @@ function fit(
end
end

return logistic_reg(formula, data, Link, LogisticRegression, sim_size)
if use_vi
return logistic_reg_vi(formula, data, Link, LogisticRegression, sim_size)
else
return logistic_reg_mcmc(formula, data, Link, LogisticRegression, sim_size)
end
end

"""
```julia
fit(formula::FormulaTerm, data::DataFrame, modelClass::LogisticRegression, Link::CRRaoLink, prior::Prior_Laplace, h::Float64 = 0.1, level::Float64 = 0.95, sim_size::Int64 = 1000)
fit(
formula::FormulaTerm,
data::DataFrame,
modelClass::LogisticRegression,
Link::CRRaoLink,
prior::Prior_Laplace,
use_vi::Bool = false,
h::Float64 = 0.1,
level::Float64 = 0.95,
sim_size::Int64 = use_vi ? 20000 : 1000
)
```

Fit a Bayesian Logistic Regression model on the input data with a Laplace prior with the provided `Link` function.
Expand Down Expand Up @@ -384,9 +421,10 @@ function fit(
modelClass::LogisticRegression,
Link::CRRaoLink,
prior::Prior_Laplace,
use_vi::Bool = false,
h::Float64 = 0.1,
level::Float64 = 0.95,
sim_size::Int64 = 1000
sim_size::Int64 = use_vi ? 20000 : 1000
)
@model LogisticRegression(X, y) = begin
p = size(X, 2)
Expand All @@ -409,12 +447,26 @@ function fit(
end
end

return logistic_reg(formula, data, Link, LogisticRegression, sim_size)
if use_vi
return logistic_reg_vi(formula, data, Link, LogisticRegression, sim_size)
else
return logistic_reg_mcmc(formula, data, Link, LogisticRegression, sim_size)
end
end

"""
```julia
fit(formula::FormulaTerm, data::DataFrame, modelClass::LogisticRegression, Link::CRRaoLink, prior::Prior_Cauchy, h::Float64 = 0.1, level::Float64 = 0.95, sim_size::Int64 = 1000)
fit(
formula::FormulaTerm,
data::DataFrame,
modelClass::LogisticRegression,
Link::CRRaoLink,
prior::Prior_Cauchy,
use_vi::Bool = false,
h::Float64 = 0.1,
level::Float64 = 0.95,
sim_size::Int64 = use_vi ? 20000 : 1000
)
```

Fit a Bayesian Logistic Regression model on the input data with a Cauchy prior with the provided `Link` function.
Expand Down Expand Up @@ -578,9 +630,10 @@ function fit(
modelClass::LogisticRegression,
Link::CRRaoLink,
prior::Prior_Cauchy,
use_vi::Bool = false,
h::Float64 = 0.1,
level::Float64 = 0.95,
sim_size::Int64 = 1000
sim_size::Int64 = use_vi ? 20000 : 1000
)
@model LogisticRegression(X, y) = begin
p = size(X, 2)
Expand All @@ -603,12 +656,26 @@ function fit(
end
end

return logistic_reg(formula, data, Link, LogisticRegression, sim_size)
if use_vi
return logistic_reg_vi(formula, data, Link, LogisticRegression, sim_size)
else
return logistic_reg_mcmc(formula, data, Link, LogisticRegression, sim_size)
end
end

"""
```julia
fit(formula::FormulaTerm, data::DataFrame, modelClass::LogisticRegression, Link::CRRaoLink, prior::Prior_TDist, h::Float64 = 1.0, level::Float64 = 0.95, sim_size::Int64 = 1000)
fit(
formula::FormulaTerm,
data::DataFrame,
modelClass::LogisticRegression,
Link::CRRaoLink,
prior::Prior_TDist,
use_vi::Bool = false,
h::Float64 = 3.0,
level::Float64 = 0.95,
sim_size::Int64 = use_vi ? 20000 : 1000
)
```

Fit a Bayesian Logistic Regression model on the input data with a T-Dist prior with the provided `Link` function.
Expand Down Expand Up @@ -797,9 +864,10 @@ function fit(
modelClass::LogisticRegression,
Link::CRRaoLink,
prior::Prior_TDist,
use_vi::Bool = false,
h::Float64 = 3.0,
level::Float64 = 0.95,
sim_size::Int64 = 1000
sim_size::Int64 = use_vi ? 20000 : 1000
)
@model LogisticRegression(X, y) = begin
p = size(X, 2)
Expand All @@ -823,14 +891,27 @@ function fit(
end
end

return logistic_reg(formula, data, Link, LogisticRegression, sim_size)
if use_vi
return logistic_reg_vi(formula, data, Link, LogisticRegression, sim_size)
else
return logistic_reg_mcmc(formula, data, Link, LogisticRegression, sim_size)
end
end



"""
```julia
fit(formula::FormulaTerm,data::DataFrame,modelClass::LogisticRegression,Link::CRRaoLink,prior::Prior_HorseShoe,level::Float64 = 0.95,sim_size::Int64 = 1000)
fit(
formula::FormulaTerm,
data::DataFrame,
modelClass::LogisticRegression,
Link::CRRaoLink,
use_vi::Bool = false,
prior::Prior_HorseShoe,
level::Float64 = 0.95,
sim_size::Int64 = use_vi ? 20000 : 1000
)
```

Fit a Bayesian Logistic Regression model on the input data with a HorseShoe prior with the provided `Link` function.
Expand Down Expand Up @@ -1033,9 +1114,10 @@ function fit(
data::DataFrame,
modelClass::LogisticRegression,
Link::CRRaoLink,
use_vi::Bool = false,
prior::Prior_HorseShoe,
level::Float64 = 0.95,
sim_size::Int64 = 1000
sim_size::Int64 = use_vi ? 20000 : 1000
)
@model LogisticRegression(X, y) = begin
p = size(X, 2)
Expand Down Expand Up @@ -1069,5 +1151,9 @@ function fit(
end
end

return logistic_reg(formula, data, Link, LogisticRegression, sim_size)
if use_vi
return logistic_reg_vi(formula, data, Link, LogisticRegression, sim_size)
else
return logistic_reg_mcmc(formula, data, Link, LogisticRegression, sim_size)
end
end
9 changes: 9 additions & 0 deletions test/numerical/bayesian/LogisticRegression.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,10 +49,19 @@ tests = [
]

for (prior, prior_testcases) in tests
# MCMC
for (link, test_mean) in prior_testcases
CRRao.set_rng(StableRNG(123))
model = fit(@formula(Vote ~ Age + Race + Income + Educate), turnout, LogisticRegression(), link, prior)
prediction = predict(model, turnout)
@test mean(prediction) - 2 * std(prediction) <= test_mean && test_mean <= mean(prediction) + 2 * std(prediction)
end

# VI
for (link, test_mean) in prior_testcases
CRRao.set_rng(StableRNG(123))
model = fit(@formula(Vote ~ Age + Race + Income + Educate), turnout, LogisticRegression(), link, prior, true)
prediction = predict(model, turnout)
@test mean(prediction) - 2 * std(prediction) <= test_mean && test_mean <= mean(prediction) + 2 * std(prediction)
end
end
Loading