Skip to content

Commit

Permalink
Use common bayesian struct
Browse files Browse the repository at this point in the history
  • Loading branch information
@ShouvikGhosh2048
ShouvikGhosh2048 committed Dec 9, 2024
1 parent 792f779 commit 84631db
Show file tree
Hide file tree
Showing 4 changed files with 86 additions and 86 deletions.
40 changes: 39 additions & 1 deletion src/CRRao.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -392,9 +392,47 @@ end

Cauchit() = Cauchit(Cauchit_Link)

"""
```julia
BayesianAlgorithm
```
Abstract type representing bayesian algorithms which are used to dispatch to appropriate calls.
"""
abstract type BayesianAlgorithm end

"""
```julia
MCMC <: BayesianAlgorithm
```
A type representing MCMC algorithms.
"""
struct MCMC <: BayesianAlgorithm
sim_size::Int64
prediction_chain_start::Int64
end

MCMC() = MCMC(1000, 200)

"""
```julia
VI <: BayesianAlgorithm
```
A type representing variational inference algorithms.
"""
struct VI <: BayesianAlgorithm
distribution_sample_count::Int64
vi_max_iters::Int64
vi_samples_per_step::Int64
end

VI() = VI(1000, 10000, 100)

export LinearRegression, LogisticRegression, PoissonRegression, NegBinomRegression, Boot_Residual
export Prior_Ridge, Prior_Laplace, Prior_Cauchy, Prior_TDist, Prior_HorseShoe, Prior_Gauss
export CRRaoLink, Logit, Probit, Cloglog, Cauchit, fit
export CRRaoLink, Logit, Probit, Cloglog, Cauchit, fit, MCMC, VI
export coef, coeftable, r2, adjr2, loglikelihood, aic, bic, sigma, predict, residuals, cooksdistance, BPTest, pvalue
export FrequentistRegression, BayesianRegression

Expand Down
4 changes: 2 additions & 2 deletions src/bayesian/getter.jl
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
function predict(container::BayesianRegression{:LinearRegression}, newdata::DataFrame, prediction_chain_start::Int64 = 200)
function predict(container::BayesianRegression{:LinearRegression}, newdata::DataFrame)
X = modelmatrix(container.formula, newdata)
W = container.samples[:, prediction_chain_start:end]
W = container.samples
predictions = X * W
return vec(mean(predictions, dims=2))
end
Expand Down
104 changes: 31 additions & 73 deletions src/bayesian/linear_regression.jl
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
function linear_reg_mcmc(formula::FormulaTerm, data::DataFrame, turingModel::Function, sim_size::Int64)
formula = apply_schema(formula, schema(formula, data),RegressionModel)
function linear_reg(formula::FormulaTerm, data::DataFrame, turingModel::Function, algorithm::MCMC)
formula = apply_schema(formula, schema(formula, data), RegressionModel)
y, X = modelcols(formula, data)

if sim_size < 500
if algorithm.sim_size < 500
@warn "Simulation size should generally be atleast 500."
end
chain = sample(CRRao_rng, turingModel(X, y), NUTS(), sim_size)
params = get_params(chain[:,:,:])
chain = sample(CRRao_rng, turingModel(X, y), NUTS(), algorithm.sim_size)
params = get_params(chain[algorithm.prediction_chain_start:end,:,:])
samples = params.β
if isa(samples, Tuple)
samples = reduce(hcat, samples)
Expand All @@ -16,17 +16,17 @@ function linear_reg_mcmc(formula::FormulaTerm, data::DataFrame, turingModel::Fun
return BayesianRegression(:LinearRegression, samples, formula)
end

function linear_reg_vi(formula::FormulaTerm, data::DataFrame, turingModel::Function, max_iter::Int64)
formula = apply_schema(formula, schema(formula, data),RegressionModel)
function linear_reg(formula::FormulaTerm, data::DataFrame, turingModel::Function, algorithm::VI)
formula = apply_schema(formula, schema(formula, data), RegressionModel)
y, X = modelcols(formula, data)

if max_iter < 500
if algorithm.vi_max_iters < 500
@warn "Max iterations should generally be atleast 500."
end

model = turingModel(X, y)
dist = vi(model, ADVI(100, max_iter))
samples = rand(CRRao_rng, dist, max_iter)
dist = vi(model, ADVI(algorithm.vi_samples_per_step, algorithm.vi_max_iters))
samples = rand(CRRao_rng, dist, algorithm.distribution_sample_count)
_, symbol_to_range = bijector(model, Val(true))
samples = samples[union(symbol_to_range[:β]...), :]
return BayesianRegression(:LinearRegression, samples, formula)
Expand All @@ -39,9 +39,8 @@ fit(
data::DataFrame,
modelClass::LinearRegression,
prior::Prior_Ridge,
use_vi::Bool = false,
algorithm::BayesianAlgorithm = MCMC(),
h::Float64 = 0.01,
sim_size::Int64 = use_vi ? 20000 : 1000,
)
```
Expand Down Expand Up @@ -108,9 +107,8 @@ function fit(
data::DataFrame,
modelClass::LinearRegression,
prior::Prior_Ridge,
use_vi::Bool = false,
algorithm::BayesianAlgorithm = MCMC(),
h::Float64 = 0.01,
sim_size::Int64 = use_vi ? 20000 : 1000
)
@model LinearRegression(X, y) = begin
p = size(X, 2)
Expand All @@ -129,11 +127,7 @@ function fit(
y ~ MvNormal(X * β, σ)
end

if use_vi
return linear_reg_vi(formula, data, LinearRegression, sim_size)
else
return linear_reg_mcmc(formula, data, LinearRegression, sim_size)
end
return linear_reg(formula, data, LinearRegression, algorithm)
end

"""
Expand All @@ -143,9 +137,8 @@ fit(
data::DataFrame,
modelClass::LinearRegression,
prior::Prior_Laplace,
use_vi::Bool = false,
algorithm::BayesianAlgorithm = MCMC(),
h::Float64 = 0.01,
sim_size::Int64 = use_vi ? 20000 : 1000,
)
```
Expand Down Expand Up @@ -210,9 +203,8 @@ function fit(
data::DataFrame,
modelClass::LinearRegression,
prior::Prior_Laplace,
use_vi::Bool = false,
algorithm::BayesianAlgorithm = MCMC(),
h::Float64 = 0.01,
sim_size::Int64 = use_vi ? 20000 : 1000
)
@model LinearRegression(X, y) = begin
p = size(X, 2)
Expand All @@ -230,11 +222,7 @@ function fit(
y ~ MvNormal(X * β, σ)
end

if use_vi
return linear_reg_vi(formula, data, LinearRegression, sim_size)
else
return linear_reg_mcmc(formula, data, LinearRegression, sim_size)
end
return linear_reg(formula, data, LinearRegression, algorithm)
end

"""
Expand All @@ -244,8 +232,7 @@ fit(
data::DataFrame,
modelClass::LinearRegression,
prior::Prior_Cauchy,
use_vi::Bool = false,
sim_size::Int64 = use_vi ? 20000 : 1000,
algorithm::BayesianAlgorithm = MCMC(),
)
```
Expand Down Expand Up @@ -309,8 +296,7 @@ function fit(
data::DataFrame,
modelClass::LinearRegression,
prior::Prior_Cauchy,
use_vi::Bool = false,
sim_size::Int64 = use_vi ? 20000 : 1000
algorithm::BayesianAlgorithm = MCMC(),
)
@model LinearRegression(X, y) = begin
p = size(X, 2)
Expand All @@ -325,11 +311,7 @@ function fit(
y ~ MvNormal(X * β, σ)
end

if use_vi
return linear_reg_vi(formula, data, LinearRegression, sim_size)
else
return linear_reg_mcmc(formula, data, LinearRegression, sim_size)
end
return linear_reg(formula, data, LinearRegression, algorithm)
end

"""
Expand All @@ -339,9 +321,8 @@ fit(
data::DataFrame,
modelClass::LinearRegression,
prior::Prior_TDist,
algorithm::BayesianAlgorithm = MCMC(),
h::Float64 = 2.0,
use_vi::Bool = false,
sim_size::Int64 = use_vi ? 20000 : 1000,
)
```
Expand Down Expand Up @@ -409,9 +390,8 @@ function fit(
data::DataFrame,
modelClass::LinearRegression,
prior::Prior_TDist,
use_vi::Bool = false,
algorithm::BayesianAlgorithm = MCMC(),
h::Float64 = 2.0,
sim_size::Int64 = use_vi ? 20000 : 1000
)
@model LinearRegression(X, y) = begin
p = size(X, 2)
Expand All @@ -429,11 +409,7 @@ function fit(
y ~ MvNormal(X * β, σ)
end

if use_vi
return linear_reg_vi(formula, data, LinearRegression, sim_size)
else
return linear_reg_mcmc(formula, data, LinearRegression, sim_size)
end
return linear_reg(formula, data, LinearRegression, algorithm)
end


Expand All @@ -444,8 +420,7 @@ fit(
data::DataFrame,
modelClass::LinearRegression,
prior::Prior_HorseShoe,
use_vi::Bool = false,
sim_size::Int64 = use_vi ? 20000 : 1000,
algorithm::BayesianAlgorithm = MCMC(),
)
```
Expand Down Expand Up @@ -506,8 +481,7 @@ function fit(
data::DataFrame,
modelClass::LinearRegression,
prior::Prior_HorseShoe,
use_vi::Bool = false,
sim_size::Int64 = use_vi ? 20000 : 1000
algorithm::BayesianAlgorithm = MCMC(),
)
@model LinearRegression(X, y) = begin
p = size(X, 2)
Expand All @@ -528,11 +502,7 @@ function fit(
y ~ MvNormal( X * β, σ)
end

if use_vi
return linear_reg_vi(formula, data, LinearRegression, sim_size)
else
return linear_reg_mcmc(formula, data, LinearRegression, sim_size)
end
return linear_reg(formula, data, LinearRegression, algorithm)
end

"""
Expand All @@ -544,8 +514,7 @@ fit(
prior::Prior_Gauss,
alpha_prior_mean::Float64,
beta_prior_mean::Vector{Float64},
use_vi::Bool = false,
sim_size::Int64 = use_vi ? 20000 : 1000,
algorithm::BayesianAlgorithm = MCMC(),
)
```
Expand Down Expand Up @@ -598,8 +567,7 @@ function fit(
prior::Prior_Gauss,
alpha_prior_mean::Float64,
beta_prior_mean::Vector{Float64},
use_vi::Bool = false,
sim_size::Int64 = use_vi ? 20000 : 1000
algorithm::BayesianAlgorithm = MCMC(),
)
@model LinearRegression(X, y) = begin
p = size(X, 2)
Expand All @@ -625,11 +593,7 @@ function fit(
y ~ MvNormal(X * β, σ)
end

if use_vi
return linear_reg_vi(formula, data, LinearRegression, sim_size)
else
return linear_reg_mcmc(formula, data, LinearRegression, sim_size)
end
return linear_reg(formula, data, LinearRegression, algorithm)
end


Expand All @@ -644,8 +608,7 @@ fit(
alpha_prior_sd::Float64,
beta_prior_mean::Vector{Float64},
beta_prior_sd::Vector{Float64},
use_vi::Bool = false,
sim_size::Int64 = use_vi ? 20000 : 1000,
algorithm::BayesianAlgorithm = MCMC(),
)
```
Expand Down Expand Up @@ -699,8 +662,7 @@ function fit(
alpha_prior_sd::Float64,
beta_prior_mean::Vector{Float64},
beta_prior_sd::Vector{Float64},
use_vi::Bool = false,
sim_size::Int64 = use_vi ? 20000 : 1000
algorithm::BayesianAlgorithm = MCMC(),
)
@model LinearRegression(X, y) = begin
p = size(X, 2)
Expand Down Expand Up @@ -728,9 +690,5 @@ function fit(
y ~ MvNormal(X * β, σ)
end

if use_vi
return linear_reg_vi(formula, data, LinearRegression, sim_size)
else
return linear_reg_mcmc(formula, data, LinearRegression, sim_size)
end
return linear_reg(formula, data, LinearRegression, algorithm)
end
24 changes: 14 additions & 10 deletions test/numerical/bayesian/LinearRegression.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,24 +12,28 @@ for (prior, test_mean) in tests
# MCMC
CRRao.set_rng(StableRNG(123))
model = fit(@formula(MPG ~ HP + WT + Gear), mtcars, LinearRegression(), prior)
prediction = predict(model, mtcars)
@test mean(prediction) - 2 * std(prediction) <= test_mean && test_mean <= mean(prediction) + 2 * std(prediction)
mcmc_prediction = predict(model, mtcars)
@test mean(mcmc_prediction) - 2 * std(mcmc_prediction) <= test_mean && test_mean <= mean(mcmc_prediction) + 2 * std(mcmc_prediction)

# VI
CRRao.set_rng(StableRNG(123))
model = fit(@formula(MPG ~ HP + WT + Gear), mtcars, LinearRegression(), prior, true)
prediction = predict(model, mtcars)
@test mean(prediction) - 2 * std(prediction) <= test_mean && test_mean <= mean(prediction) + 2 * std(prediction)
model = fit(@formula(MPG ~ HP + WT + Gear), mtcars, LinearRegression(), prior, VI())
vi_prediction = predict(model, mtcars)
@test mean(vi_prediction) - 2 * std(vi_prediction) <= test_mean && test_mean <= mean(vi_prediction) + 2 * std(vi_prediction)

@test maximum(abs.(mcmc_prediction - vi_prediction)) <= 5.0
end

gauss_test = 20.0796026428345

CRRao.set_rng(StableRNG(123))
model = fit(@formula(MPG ~ HP + WT + Gear), mtcars, LinearRegression(), Prior_Gauss(), 30.0, [0.0,-3.0,1.0])
prediction = predict(model, mtcars)
@test mean(prediction) - 2 * std(prediction) <= gauss_test && gauss_test <= mean(prediction) + 2 * std(prediction)
mcmc_prediction = predict(model, mtcars)
@test mean(mcmc_prediction) - 2 * std(mcmc_prediction) <= gauss_test && gauss_test <= mean(mcmc_prediction) + 2 * std(mcmc_prediction)

CRRao.set_rng(StableRNG(123))
model = fit(@formula(MPG ~ HP + WT + Gear), mtcars, LinearRegression(), Prior_Gauss(), 30.0, [0.0,-3.0,1.0], true)
prediction = predict(model, mtcars)
@test mean(prediction) - 2 * std(prediction) <= gauss_test && gauss_test <= mean(prediction) + 2 * std(prediction)
model = fit(@formula(MPG ~ HP + WT + Gear), mtcars, LinearRegression(), Prior_Gauss(), 30.0, [0.0,-3.0,1.0], VI())
vi_prediction = predict(model, mtcars)
@test mean(vi_prediction) - 2 * std(vi_prediction) <= gauss_test && gauss_test <= mean(vi_prediction) + 2 * std(vi_prediction)

@test maximum(abs.(mcmc_prediction - vi_prediction)) <= 5.0

0 comments on commit 84631db

Please sign in to comment.