🚀flag to decide if the initial states will be pen

src/StateSpaceLearning.jl

@@ -36,14 +36,16 @@ export fit_model, forecast
     - `hyperparameter_selection::String`: Information criteria method for hyperparameter selection (default: "aic").
     - `adalasso_coef::Float64`: AdaLasso adjustment coefficient (default: 0.1).
     - `select_exogenous::Bool`: Flag to select exogenous variables. When false the penalty factor for these variables will be set to 0 (default: true).
+    - `penalize_initial_states::Bool`: Flag to penalize initial states. When false the penalty factor for these variables will be set to 0 (default: true).
 
     # Returns
     - `Output`: Output object containing model information, coefficients, residuals, etc.
 
 """
 function fit_model(y::Vector{Fl}; model_type::String="Basic Structural", Exogenous_X::Union{Matrix{Fl}, Missing}=missing,
                     estimation_procedure::String="AdaLasso", s::Int64=12, outlier::Bool=false, stabilize_ζ::Int64=0,
-                    α::Float64=0.1, hyperparameter_selection::String="aic", adalasso_coef::Float64=0.1, select_exogenous::Bool=true)::Output where Fl
+                    α::Float64=0.1, hyperparameter_selection::String="aic", adalasso_coef::Float64=0.1, 
+                    select_exogenous::Bool=true, penalize_initial_states::Bool=true)::Output where Fl
 
     T = length(y)
     @assert T > s "Time series must be longer than the seasonal period"
@@ -62,7 +64,7 @@ function fit_model(y::Vector{Fl}; model_type::String="Basic Structural", Exogeno
 
     components_indexes  = get_components_indexes(T, s, Exogenous_X, outlier, model_type, stabilize_ζ)
 
-    coefs, estimation_ϵ = fit_estimation_procedure(estimation_procedure, Estimation_X, estimation_y, α, hyperparameter_selection, components_indexes, adalasso_coef, select_exogenous)
+    coefs, estimation_ϵ = fit_estimation_procedure(estimation_procedure, Estimation_X, estimation_y, α, hyperparameter_selection, components_indexes, adalasso_coef, select_exogenous, penalize_initial_states)
 
     components          = build_components(Estimation_X, coefs, components_indexes)
 

src/estimation_procedure/adalasso.jl

@@ -14,6 +14,7 @@
     - `components_indexes::Dict{String, Vector{Int64}}`: Dictionary containing indexes for different components.
     - `adalasso_coef::Float64`: AdaLasso adjustment coefficient (default: 0.1).
     - `select_exogenous::Bool`: Flag for selecting exogenous variables. When false the penalty factor for these variables will be set to 0.
+    - `penalize_initial_states::Bool`: Flag to penalize initial states. When false the penalty factor for these variables will be set to 0.
 
     # Returns
     - `Tuple{Vector{Float64}, Vector{Float64}}`: Tuple containing coefficients and residuals of the fitted AdaLasso model.
@@ -22,7 +23,7 @@
 function fit_adalasso(Estimation_X::Matrix{Tl}, estimation_y::Vector{Fl}, α::Float64, 
                         hyperparameter_selection::String, 
                         components_indexes::Dict{String, Vector{Int64}},
-                        adalasso_coef::Float64, select_exogenous::Bool)::Tuple{Vector{Float64}, Vector{Float64}} where {Tl, Fl}
+                        adalasso_coef::Float64, select_exogenous::Bool, penalize_initial_states::Bool)::Tuple{Vector{Float64}, Vector{Float64}} where {Tl, Fl}
 
     penalty_factor = ones(size(Estimation_X, 2) - 1); penalty_factor[components_indexes["initial_states"][2:end] .- 1] .= 0
     coefs, _  = fit_lasso(Estimation_X, estimation_y, α, hyperparameter_selection, select_exogenous, components_indexes; penalty_factor = penalty_factor, intercept = false)
@@ -40,6 +41,7 @@ function fit_adalasso(Estimation_X::Matrix{Tl}, estimation_y::Vector{Fl}, α::Fl
             end
         end
     end 
+    !penalize_initial_states ? penalty_factor[components_indexes["initial_states"][2:end] .- 1] .= 0 : nothing
     return fit_lasso(Estimation_X, estimation_y, α, hyperparameter_selection, select_exogenous, components_indexes; penalty_factor=penalty_factor)
 end
 

src/estimation_procedure/estimation_utils.jl

@@ -62,18 +62,19 @@ end
     - `components_indexes::Dict{String, Vector{Int64}}`: Dictionary containing indexes for different components.
     - `adalasso_coef::Float64`: AdaLasso adjustment coefficient (default: 0.1).
     - `select_exogenous::Bool`: Flag for selecting exogenous variables. When false the penalty factor for these variables will be set to 0.
+    - `penalize_initial_states::Bool`: Flag to penalize initial states. When false the penalty factor for these variables will be set to 0.
 
     # Returns
     - `Tuple{Vector{Float64}, Vector{Float64}}`: Tuple containing coefficients and residuals of the fitted estimation procedure.
 
 """
 function fit_estimation_procedure(estimation_procedure::String, Estimation_X::Matrix{Tl}, estimation_y::Vector{Fl}, α::Float64, 
                                     hyperparameter_selection::String, components_indexes::Dict{String, Vector{Int64}}, 
-                                    adalasso_coef::Float64, select_exogenous::Bool)::Tuple{Vector{Float64}, Vector{Float64}} where {Tl, Fl}
+                                    adalasso_coef::Float64, select_exogenous::Bool, penalize_initial_states::Bool)::Tuple{Vector{Float64}, Vector{Float64}} where {Tl, Fl}
 
     estimation_arguments     = Dict("Lasso" => (Estimation_X, estimation_y, α, hyperparameter_selection, select_exogenous, components_indexes),
                                 "AdaLasso" => (Estimation_X, estimation_y, α, hyperparameter_selection, 
-                                               components_indexes, adalasso_coef, select_exogenous))
+                                               components_indexes, adalasso_coef, select_exogenous, penalize_initial_states))
 
     available_estimation = Dict("Lasso" => fit_lasso,
                                 "AdaLasso" => fit_adalasso)

test/estimation_procedure/adalasso.jl

@@ -6,11 +6,16 @@
 
     estimation_y = Estimation_X*rand(size(Estimation_X, 2)) + rand(10).*5
 
-    coefs1, ϵ1 = StateSpaceLearning.fit_adalasso(Estimation_X, estimation_y, 0.1, "aic", components_indexes, 0.1, true)
+    coefs1, ϵ1 = StateSpaceLearning.fit_adalasso(Estimation_X, estimation_y, 0.1, "aic", components_indexes, 0.1, true, true)
     @test length(coefs1) == 43
     @test length(ϵ1) == 10
 
-    coefs2, ϵ2 = StateSpaceLearning.fit_adalasso(Estimation_X, estimation_y, 0.1, "aic", components_indexes, 10000.0, true)
+    coefs1, ϵ1 = StateSpaceLearning.fit_adalasso(Estimation_X, estimation_y, 0.1, "aic", components_indexes, 0.1, true, false)
+    @test length(coefs1) == 43
+    @test length(ϵ1) == 10
+    @test all(coefs1[components_indexes["initial_states"][2:end] .- 1] .!= 0)
+
+    coefs2, ϵ2 = StateSpaceLearning.fit_adalasso(Estimation_X, estimation_y, 0.1, "aic", components_indexes, 10000.0, true, true)
     coefs_lasso, ϵ_lasso = StateSpaceLearning.fit_lasso(Estimation_X, estimation_y, 0.1, "aic", true, components_indexes; intercept = true)
     @test all(isapprox.(coefs2, coefs_lasso; atol = 1e-3))
     @test all(isapprox.(ϵ2, ϵ_lasso; atol = 1e-3))

test/estimation_procedure/estimation_utils.jl

@@ -32,11 +32,11 @@ end
     Estimation_X = StateSpaceLearning.create_X("Basic Structural", 10, 3, Exogenous_X, true, 0)
     estimation_y = Estimation_X*rand(size(Estimation_X, 2)) + rand(10).*5
 
-    coefs1, ϵ1 = StateSpaceLearning.fit_estimation_procedure("Lasso", Estimation_X, estimation_y, 0.1, "aic", components_indexes, 0.1, true)
+    coefs1, ϵ1 = StateSpaceLearning.fit_estimation_procedure("Lasso", Estimation_X, estimation_y, 0.1, "aic", components_indexes, 0.1, true, true)
     @test length(coefs1) == 43
     @test length(ϵ1) == 10
 
-    coefs2, ϵ2 = StateSpaceLearning.fit_estimation_procedure("AdaLasso", Estimation_X, estimation_y, 0.1, "aic", components_indexes, 0.1, true)
+    coefs2, ϵ2 = StateSpaceLearning.fit_estimation_procedure("AdaLasso", Estimation_X, estimation_y, 0.1, "aic", components_indexes, 0.1, true, true)
     @test length(coefs2) == 43
     @test length(ϵ2) == 10
 end