diff --git a/docs/src/performance/mixed_differentiation.md b/docs/src/performance/mixed_differentiation.md index f33fa6ab2..a9f9abcfa 100644 --- a/docs/src/performance/mixed_differentiation.md +++ b/docs/src/performance/mixed_differentiation.md @@ -10,12 +10,9 @@ ml_term = SemML( RAMSymbolic(partable) ) -ridge_term = SemRidge( - α_ridge = 0.01, - which_ridge = params(ml_term) -) +ridge_penalty = SemRidge(partable) -model_ml_ridge = Sem(ml_term, ridge_term) +model_ml_ridge = Sem(ml_term, ridge_penalty => 0.01) model_ml_ridge_fit = fit(model_ml_ridge) ``` diff --git a/docs/src/tutorials/regularization/regularization.md b/docs/src/tutorials/regularization/regularization.md index 79e301c21..9e3f2f8b0 100644 --- a/docs/src/tutorials/regularization/regularization.md +++ b/docs/src/tutorials/regularization/regularization.md @@ -1,44 +1,41 @@ # Regularization -## Setup +At the moment, *StructuralEquationModels.jl* offers two ways to regularize SEMs: -For ridge regularization, you can simply use `SemRidge` as an additional loss function -(for example, a model with the loss functions `SemML` and `SemRidge` corresponds to ridge-regularized maximum likelihood estimation). +1. Built-in smooth penalties via [`SemParamsPenalty`](@ref) and its convenience wrappers such as + [`SemRidge`](@ref), [`SemLasso`](@ref), and [`SemElasticNet`](@ref). +2. More general proximal regularization via + [`ProximalOperators.jl`](https://github.com/JuliaFirstOrder/ProximalOperators.jl) + combined with *proximal optimization* from + [`ProximalAlgorithms.jl`](https://github.com/JuliaFirstOrder/ProximalAlgorithms.jl). + +## 1. Regularization with SemParamsPenalty terms -You can define lasso, elastic net and other forms of regularization using [`ProximalOperators.jl`](https://github.com/JuliaFirstOrder/ProximalOperators.jl) -and optimize the SEM with so-called *proximal optimization* algorithms from [`ProximalAlgorithms.jl`](https://github.com/JuliaFirstOrder/ProximalAlgorithms.jl). ```@setup reg -using StructuralEquationModels, ProximalAlgorithms, ProximalOperators +using StructuralEquationModels, ProximalAlgorithms, ProximalOperators, StenoGraphs ``` -```julia -using Pkg -Pkg.add("ProximalAlgorithms") -Pkg.add("ProximalOperators") +[`SemParamsPenalty`](@ref) is a generic regularization loss term acting on selected SEM +parameters. Its wrappers cover common penalties: -using StructuralEquationModels, ProximalAlgorithms, ProximalOperators -``` - -## Proximal optimization +- [`SemRidge`](@ref): l2 regularization +- [`SemLasso`](@ref): l1 regularization +- [`SemElasticNet`](@ref): elastic net regularization -With the *ProximalAlgorithms* package loaded, it is now possible to use the `:Proximal` -optimization engine in `SemOptimizer` for estimating regularized models. +These regularization terms are added to a model just like any other loss term. The +`=>` syntax sets the overall weight of the penalty in the objective. For +[`SemElasticNet`](@ref), its `weight2` keyword argument controls the relative strength of the +l2 part inside the elastic-net penalty, while the outer loss-term weight controls the overall +regularization strength. -```julia -SemOptimizer(; - engine = :Proximal, - algorithm = ProximalAlgorithms.PANOC(), - operator_g, - operator_h = nothing -) +```@docs +SemParamsPenalty +SemRidge +SemLasso +SemElasticNet ``` -The *proximal operator* (aka the *regularization function*) is passed as `operator_g`, see [available operators](https://juliafirstorder.github.io/ProximalOperators.jl/stable/functions/). -The `algorithm` is chosen from one of the [available algorithms](https://juliafirstorder.github.io/ProximalAlgorithms.jl/stable/guide/implemented_algorithms/). - -## First example - lasso - To show how it works, let's revisit [A first model](@ref): ```@example reg @@ -80,29 +77,78 @@ model = Sem( ) ``` -We labeled the covariances between the items because we want to regularize those: +Ridge regularization is added as an additional loss term listing the parameters we want to +regularize. In this example, we regularize the covariances between the observed variables +using the labels of the corresponding model parameters. + +```@example reg +cov_params = [:cov_15, :cov_24, :cov_26, :cov_37, :cov_48, :cov_68] + +model_ridge = + Sem( + SemML(SemObservedData(data = data), RAM(partable)), + SemRidge(partable, cov_params) => 0.02, + ) + +fit_ridge = fit(model_ridge) +``` + +The same pattern also works for lasso and elastic net: ```@example reg -cov_inds = getindex.( - Ref(param_indices(model)), - [:cov_15, :cov_24, :cov_26, :cov_37, :cov_48, :cov_68]) +lasso_term = SemLasso(partable, cov_params) +elasticnet_term = SemElasticNet(partable, cov_params; weight2 = 0.5) ``` -In the following, we fit the same model with lasso regularization of those covariances. +## 2. Proximal optimization + +*SEM.jl* uses Julia extensions mechanism to automatically enable proximal optimization when +the *ProximalAlgorithms* package is loaded. + +```julia +using Pkg +Pkg.add("ProximalAlgorithms") +Pkg.add("ProximalOperators") + +using StructuralEquationModels, ProximalAlgorithms, ProximalOperators +``` + +To use proximal optimization for SEM models fitting, specify `:Proximal` as the +optimization engine in `SemOptimizer`. + +```julia +SemOptimizer(; + engine = :Proximal, + algorithm = ProximalAlgorithms.PANOC(), + operator_g, + operator_h = nothing +) +``` + +The *proximal operator* (aka the *regularization function*) is passed as `operator_g`, +see [available operators](https://juliafirstorder.github.io/ProximalOperators.jl/stable/functions/). +The `algorithm` is chosen from one of the [available algorithms](https://juliafirstorder.github.io/ProximalAlgorithms.jl/stable/guide/implemented_algorithms/). + +### LASSO via proximal operators + +In the following example, we fit the same SEM model with lasso regularization for selected covariances. The lasso penalty is defined as ```math \sum \lambda_i \lvert \theta_i \rvert ``` -In `ProximalOperators.jl`, lasso regularization is represented by the [`NormL1`](https://juliafirstorder.github.io/ProximalOperators.jl/stable/functions/#ProximalOperators.NormL1) operator. It allows controlling the amount of -regularization individually for each SEM model parameter via the vector of hyperparameters (`λ`). -To regularize only the observed item covariances, we define `λ` as +In `ProximalOperators.jl`, lasso regularization is represented by the [`NormL1`](https://juliafirstorder.github.io/ProximalOperators.jl/stable/functions/#ProximalOperators.NormL1) +operator. It allows controlling the amount of regularization individually for each SEM model parameter +via the vector of hyperparameters (`λ`). To regularize only the observed item covariances, +we define `λ` as ```@example reg -λ = zeros(31); λ[cov_inds] .= 0.02 +cov_params = [:cov_15, :cov_24, :cov_26, :cov_37, :cov_48, :cov_68] +cov_inds = param_indices(partable, cov_params) +λ = zeros(nparams(partable)); λ[cov_inds] .= 0.02 -optimizer_lasso = SemOptimizer( +optimizer_lasso_proximal = SemOptimizer( engine = :Proximal, operator_g = NormL1(λ) ) @@ -112,7 +158,7 @@ Let's fit the regularized model ```@example reg -fit_lasso = fit(optimizer_lasso, model) +fit_lasso = fit(optimizer_lasso_proximal, model) ``` and compare the solution to unregularizted estimates: @@ -134,7 +180,7 @@ and additional keyword arguments directly to the `fit` function: fit_lasso2 = fit(model; engine = :Proximal, operator_g = NormL1(λ)) ``` -## Second example - mixed l1 and l0 regularization +### Mixed l1 and l0 regularization You can choose to penalize different parameters with different types of regularization functions. Let's use the *lasso* (*l1*) again on the covariances, but additionally penalize the error variances of @@ -161,4 +207,4 @@ Let's again compare the different results: update_partable!(partable, :estimate_mixed, fit_mixed, solution(fit_mixed)) details(partable) -``` \ No newline at end of file +``` diff --git a/src/StructuralEquationModels.jl b/src/StructuralEquationModels.jl index 0dbcd16aa..2085f1341 100644 --- a/src/StructuralEquationModels.jl +++ b/src/StructuralEquationModels.jl @@ -31,6 +31,7 @@ include("objective_gradient_hessian.jl") # helper objects and functions include("additional_functions/commutation_matrix.jl") +include("additional_functions/sparse_utils.jl") include("additional_functions/params_array.jl") # fitted objects @@ -63,9 +64,10 @@ include("implied/empty.jl") include("loss/abstract.jl") include("loss/ML/ML.jl") include("loss/ML/FIML.jl") -include("loss/regularization/ridge.jl") include("loss/WLS/WLS.jl") include("loss/constant/constant.jl") +# regularization +include("loss/regularization/params_penalty.jl") # constructor include("frontend/specification/Sem.jl") include("frontend/specification/documentation.jl") @@ -121,9 +123,14 @@ export AbstractSem, SemML, SemFIML, em_mvn, - SemRidge, - SemConstant, SemWLS, + SemConstant, + SemParamsPenalty, + SemHinge, + SemLasso, + SemNorm, + SemRidge, + SemElasticNet, loss, nsem_terms, sem_terms, diff --git a/src/additional_functions/sparse_utils.jl b/src/additional_functions/sparse_utils.jl new file mode 100644 index 000000000..76381c473 --- /dev/null +++ b/src/additional_functions/sparse_utils.jl @@ -0,0 +1,83 @@ +# generate sparse matrix with 1 in each row +function eachrow_to_col( + ::Type{T}, + column_indices::AbstractVector{Int}, + ncolumns::Integer, +) where {T} + nrows = length(column_indices) + (nrows > ncolumns) && throw( + DimensionMismatch( + "The number of rows ($nrows) cannot exceed the number of columns ($ncolumns)", + ), + ) + all(i -> 1 <= i <= ncolumns, column_indices) || + throw(ArgumentError("All column indices must be between 1 and $ncolumns")) + + sparse(eachindex(column_indices), column_indices, ones(T, nrows), nrows, ncolumns) +end + +eachrow_to_col(column_indices::AbstractVector{Int}, ncolumns::Integer) = + eachrow_to_col(Float64, column_indices, ncolumns) + +# generate sparse matrix with 1 in each column +function eachcol_to_row( + ::Type{T}, + row_indices::AbstractVector{Int}, + nrows::Integer, +) where {T} + ncols = length(row_indices) + (ncols > nrows) && throw( + DimensionMismatch( + "The number of columns ($ncols) cannot exceed the number of rows ($nrows)", + ), + ) + all(i -> 1 <= i <= nrows, row_indices) || + throw(ArgumentError("All row indices must be between 1 and $nrows")) + + sparse(row_indices, eachindex(row_indices), ones(T, ncols), nrows, ncols) +end + +eachcol_to_row(row_indices::AbstractVector{Int}, nrows::Integer) = + eachcol_to_row(Float64, row_indices, nrows) + +# non-zero indices of columns in matrix A generated by eachrow_to_col() +# if order == :rows, then the indices are in the order of the rows, +# if order == :columns, the indices are in the order of the columns +function nzcols_eachrow_to_col(F, A::SparseMatrixCSC; order::Symbol = :rows) + order ∈ [:rows, :columns] || throw(ArgumentError("order must be :rows or :columns")) + T = typeof(F(1)) + res = Vector{T}(undef, size(A, 1)) + n = 0 + for i in 1:size(A, 2) + colptr = A.colptr[i] + next_colptr = A.colptr[i+1] + if next_colptr > colptr # non-zero + @assert next_colptr - colptr == 1 + n += 1 + res[order == :rows ? A.rowval[colptr] : n] = F(i) + end + end + @assert n == size(A, 1) + return res +end + +nzcols_eachrow_to_col(A::SparseMatrixCSC; order::Symbol = :rows) = + nzcols_eachrow_to_col(identity, A, order = order) + +# same as nzcols_eachrow_to_col() +# but without assumption that each row cooresponds to exactly one column +# the order is always columns order +function nzcols(F, A::SparseMatrixCSC) + T = typeof(F(1)) + res = Vector{T}() + for i in 1:size(A, 2) + colptr = A.colptr[i] + next_colptr = A.colptr[i+1] + if next_colptr > colptr # non-zero + push!(res, F(i)) + end + end + return res +end + +nzcols(A::SparseMatrixCSC) = nzcols(identity, A) diff --git a/src/frontend/specification/RAMMatrices.jl b/src/frontend/specification/RAMMatrices.jl index d430e9c01..99733ebd7 100644 --- a/src/frontend/specification/RAMMatrices.jl +++ b/src/frontend/specification/RAMMatrices.jl @@ -22,34 +22,23 @@ vars(ram::RAMMatrices) = ram.vars isobserved_var(ram::RAMMatrices, i::Integer) = ram.F.colptr[i+1] > ram.F.colptr[i] islatent_var(ram::RAMMatrices, i::Integer) = ram.F.colptr[i+1] == ram.F.colptr[i] -# indices of observed variables in the order as they appear in ram.F rows -function observed_var_indices(ram::RAMMatrices) - obs_inds = Vector{Int}(undef, nobserved_vars(ram)) - @inbounds for i in 1:nvars(ram) - colptr = ram.F.colptr[i] - if ram.F.colptr[i+1] > colptr # is observed - obs_inds[ram.F.rowval[colptr]] = i - end - end - return obs_inds -end +# indices of observed variables, for order=:rows (default), the order is as they appear in ram.F rows +# if order=:columns, the order is as they appear in the comined variables list (ram.F columns) +observed_var_indices(ram::RAMMatrices; order::Symbol = :rows) = + nzcols_eachrow_to_col(ram.F; order) latent_var_indices(ram::RAMMatrices) = [i for i in axes(ram.F, 2) if islatent_var(ram, i)] -# observed variables in the order as they appear in ram.F rows -function observed_vars(ram::RAMMatrices) +# observed variables, if order=:rows, the order is as they appear in ram.F rows +# if order=:columns, the order is as they appear in the comined variables list (ram.F columns) +function observed_vars(ram::RAMMatrices; order::Symbol = :rows) + order ∈ [:rows, :columns] || + throw(ArgumentError("order kwarg should be :rows or :columns")) if isnothing(ram.vars) @warn "Your RAMMatrices do not contain variable names. Please make sure the order of variables in your data is correct!" return nothing else - obs_vars = Vector{Symbol}(undef, nobserved_vars(ram)) - @inbounds for (i, v) in enumerate(vars(ram)) - colptr = ram.F.colptr[i] - if ram.F.colptr[i+1] > colptr # is observed - obs_vars[ram.F.rowval[colptr]] = v - end - end - return obs_vars + return nzcols_eachrow_to_col(Base.Fix1(getindex, vars(ram)), ram.F; order = order) end end @@ -221,13 +210,7 @@ function RAMMatrices( return RAMMatrices( ParamsMatrix{T}(A_inds, A_consts, (n_vars, n_vars)), ParamsMatrix{T}(S_inds, S_consts, (n_vars, n_vars)), - sparse( - 1:n_observed, - [vars_index[var] for var in partable.observed_vars], - ones(T, n_observed), - n_observed, - n_vars, - ), + eachrow_to_col(T, [vars_index[var] for var in partable.observed_vars], n_vars), !isnothing(M_inds) ? ParamsVector{T}(M_inds, M_consts, (n_vars,)) : nothing, param_labels, vars_sorted, diff --git a/src/frontend/specification/documentation.jl b/src/frontend/specification/documentation.jl index a3a8d2659..38a31f6ed 100644 --- a/src/frontend/specification/documentation.jl +++ b/src/frontend/specification/documentation.jl @@ -37,6 +37,16 @@ Return the vector of parameter labels (in the same order as [`params`](@ref)). """ param_labels(spec::SemSpecification) = spec.param_labels +""" + param_indices(spec::SemSpecification, params::AbstractVector{Symbol}) -> Vector{Int} + +Convert parameter labels to their indices in the SEM specification. +""" +function param_indices(spec::SemSpecification, params::AbstractVector{Symbol}) + param_map = Dict(param => i for (i, param) in enumerate(SEM.param_labels(spec))) + return Int[param_map[param] for param in params] +end + """ `ParameterTable`s contain the specification of a structural equation model. diff --git a/src/loss/regularization/params_penalty.jl b/src/loss/regularization/params_penalty.jl new file mode 100644 index 000000000..e080c2154 --- /dev/null +++ b/src/loss/regularization/params_penalty.jl @@ -0,0 +1,379 @@ +# generic penalties on transformed parameter vectors + +############################################################################################ +### Types +############################################################################################ + +""" + struct SemParamsPenalty{H, TP, TA, TB, TS1, TS2, TW2, TT} <: AbstractLoss + +Generic regularization term acting on SEM parameters, optionally after an affine +transform `A * p + b`, with fixed `A` and `b`. + +Without a hinge bound (`bound = :none`), the elementwise penalty is +``|z_i|^{s_1} + w_2 * |z_i|^{s_2}``. The second term is optional. + +With `bound = :l` or `bound = :u`, the above penalty is applied to the +one-sided hinge-transformed *z* instead (``z = \\max(x - t, 0)``, or +``z = \\max(t - x, 0)``, respectively, where *t* is the `threshold` parameter). + +In addition to the generic `SemParamsPenalty` constructor, `SemNorm`, +`SemLasso`, `SemRidge`, `SemElasticNet` and `SemHinge` wrappers provide +convenient construction of common regularization penalties. + +# Constructors + SemParamsPenalty([params], [A], [b]; + shape1 = 1, shape2 = nothing, weight2 = nothing, + bound = :none, threshold = 0.0) + + SemParamsPenalty(spec::SemSpecification, params::AbstractVector, + [A::AbstractMatrix = nothing], + [b::AbstractVector = nothing]; + shape1 = 1, shape2 = nothing, + weight2 = nothing, + bound = :none, threshold = 0.0) + +# See also + +- Wrappers for `SemParamsPenalty` providing default regularization terms: + [`SemNorm`](@ref), [`SemLasso`](@ref), [`SemRidge`](@ref), [`SemElasticNet`](@ref), and + [`SemHinge`](@ref). +""" +struct SemParamsPenalty{H, TP, TA, TB, TS1, TS2, TW2, TT} <: AbstractLoss + hessianeval::ExactHessian + param_inds::TP + A::TA + At::TA + b::TB + shape1::TS1 + shape2::TS2 + weight2::TW2 + threshold::TT +end + +############################################################################ +### Constructors +############################################################################ + +function SemParamsPenalty( + param_inds::Union{AbstractVector, Nothing} = nothing, + A::Union{AbstractMatrix, Nothing} = nothing, + b::Union{AbstractVector, Nothing} = nothing; + shape1::Real = 1, + shape2::Union{Real, Nothing} = nothing, + weight2::Union{Real, Nothing} = nothing, + bound::Symbol = :none, + threshold::Real = 0.0, +) + bound ∈ (:none, :l, :u) || + throw(ArgumentError("bound must be :none, :l, or :u, $bound given")) + + shape1 > 0 || throw(ArgumentError("shape1 must be positive, got $shape1")) + isnothing(shape2) || + shape2 > 0 || + throw(ArgumentError("shape2 must be positive when provided, got $shape2")) + + isnothing(param_inds) || + !(eltype(param_inds) <: Symbol) || + throw(ArgumentError("Symbol parameter ids require a SemSpecification.")) + + isnothing(A) || + isnothing(param_inds) || + size(A, 2) == length(param_inds) || + throw( + DimensionMismatch( + "The transformation matrix columns ($(size(A, 2))) should match " * + "the number of parameters to regularize ($(length(param_inds)))", + ), + ) + + if !isnothing(b) + (!isnothing(A) || !isnothing(param_inds)) || throw( + ArgumentError( + "An intercept vector requires either a parameter subset or a transformation matrix.", + ), + ) + + nused_params = isnothing(param_inds) ? size(A, 2) : length(param_inds) + npenalty_els = isnothing(A) ? length(param_inds) : size(A, 1) + length(b) == npenalty_els || throw( + DimensionMismatch( + "The intercept length ($(length(b))) should match " * + ( + !isnothing(A) ? "the rows of the transformation matrix" : + "the number of parameters to regularize" + ) * + " ($nused_params)", + ), + ) + end + + if isnothing(shape2) + isnothing(weight2) || + throw(ArgumentError("weight2 can only be specified when shape2 is specified.")) + weight2 = nothing + else + isnothing(weight2) && + throw(ArgumentError("weight2 must be specified when shape2 is specified.")) + end + + shape1 = shape1 isa AbstractFloat && isinteger(shape1) ? Int(shape1) : shape1 + shape2 = + isnothing(shape2) ? nothing : + (shape2 isa AbstractFloat && isinteger(shape2) ? Int(shape2) : shape2) + At = !isnothing(A) ? permutedims(A) : nothing + threshold_value = float(threshold) + + return SemParamsPenalty{ + bound, + typeof(param_inds), + typeof(A), + typeof(b), + typeof(shape1), + typeof(shape2), + typeof(weight2), + typeof(threshold_value), + }( + ExactHessian(), + param_inds, + A, + At, + b, + shape1, + shape2, + weight2, + threshold_value, + ) +end + +function SemParamsPenalty( + spec::SemSpecification, + params::AbstractVector, + A::Union{AbstractMatrix, Nothing} = nothing, + b::Union{AbstractVector, Nothing} = nothing; + shape1::Real = 1, + shape2::Union{Real, Nothing} = nothing, + weight2::Union{Real, Nothing} = nothing, + bound::Symbol = :none, + threshold::Real = 0.0, +) + param_inds = eltype(params) <: Symbol ? param_indices(spec, params) : params + + isnothing(A) || + size(A, 2) == length(param_inds) || + throw( + DimensionMismatch( + "The transformation matrix columns ($(size(A, 2))) should match " * + "the parameters to regularize ($(length(param_inds)))", + ), + ) + + sel_params_mtx = eachrow_to_col(Float64, param_inds, nparams(spec)) + if !isnothing(A) + if A isa SparseMatrixCSC + # for sparse case, combine the matrix transform and + # parameter selection into a single matrix + A = convert(typeof(A), A * sel_params_mtx) + param_inds = nothing + end + else + A = sel_params_mtx + param_inds = nothing + end + + return SemParamsPenalty(param_inds, A, b; shape1, shape2, weight2, threshold, bound) +end + +# wrappers for common regularization schemes + +SemNorm(args...; shape::Real = 1) = SemParamsPenalty(args...; shape1 = shape) +SemLasso(args...; kwargs...) = SemNorm(args...; shape = 1, kwargs...) +SemRidge(args...; kwargs...) = SemNorm(args...; shape = 2, kwargs...) +SemElasticNet(args...; weight2::Real = 0.5) = + SemParamsPenalty(args...; shape1 = 1, shape2 = 2, weight2) +SemHinge(args...; bound::Symbol = :l, kwargs...) = + SemParamsPenalty(args...; bound, kwargs...) + +nparams(f::SemParamsPenalty) = + !isnothing(f.A) ? size(f.A, 2) : (isnothing(f.param_inds) ? 0 : length(f.param_inds)) + +# helpers using constant propagation for efficient computation +# for specific shapes and bounds + +_elhinge(x, ::Val{:none}, threshold) = abs(x) +_elhinge(x, ::Val{:l}, threshold) = max(x - threshold, zero(x)) +_elhinge(x, ::Val{:u}, threshold) = max(threshold - x, zero(x)) + +_elhinge_direction(x, ::Val{:none}, threshold) = sign(x) +_elhinge_direction(x, ::Val{:l}, threshold) = x > threshold ? one(x) : zero(x) +_elhinge_direction(x, ::Val{:u}, threshold) = x < threshold ? -one(x) : zero(x) + +_elhinge_isactive(x, ::Val{:none}, threshold) = true +_elhinge_isactive(x, ::Val{:l}, threshold) = x > threshold +_elhinge_isactive(x, ::Val{:u}, threshold) = x < threshold + +elnorm(x, ::Val{1}, mode, threshold) = _elhinge(x, mode, threshold) +elnorm(x, ::Val{2}, mode, threshold) = abs2(_elhinge(x, mode, threshold)) +elnorm(x, ::Val{S}, mode, threshold) where {S} = _elhinge(x, mode, threshold)^S + +elnormgrad(x, ::Val{1}, mode, threshold) = _elhinge_direction(x, mode, threshold) +elnormgrad(x, ::Val{2}, mode, threshold) = + (one(x) + one(x)) * + _elhinge(x, mode, threshold) * + _elhinge_direction(x, mode, threshold) + +elnormgrad(x, ::Val{S}, mode, threshold) where {S} = + _elhinge_isactive(x, mode, threshold) ? + S * _elhinge(x, mode, threshold)^(S - 1) * _elhinge_direction(x, mode, threshold) : + zero(x) + +elnormhdiag(x, ::Val{1}, mode, threshold) = zero(x) +elnormhdiag(x, ::Val{2}, mode, threshold) = + _elhinge_isactive(x, mode, threshold) ? (one(x) + one(x)) : zero(x) + +elnormhdiag(x, ::Val{S}, mode, threshold) where {S} = + _elhinge_isactive(x, mode, threshold) ? + S * (S - 1) * _elhinge(x, mode, threshold)^(S - 2) : zero(x) + +Base.@propagate_inbounds @inline function _elevaluate!( + objective, + elm_grad, + elm_hdiag, + x, + i, + weight, + shape, + mode, + threshold, +) + isnothing(objective) || (objective += weight * elnorm(x, shape, mode, threshold)) + isnothing(elm_grad) || (elm_grad[i] += weight * elnormgrad(x, shape, mode, threshold)) + isnothing(elm_hdiag) || + (elm_hdiag[i] += weight * elnormhdiag(x, shape, mode, threshold)) + + return objective +end + +Base.@propagate_inbounds @inline function _elevaluate!( + objective, + elm_grad, + elm_hdiag, + x, + i, + weight::Nothing, + shape, + mode, + threshold, +) + isnothing(objective) || (objective += elnorm(x, shape, mode, threshold)) + isnothing(elm_grad) || (elm_grad[i] += elnormgrad(x, shape, mode, threshold)) + isnothing(elm_hdiag) || (elm_hdiag[i] += elnormhdiag(x, shape, mode, threshold)) + + return objective +end + +function _evaluate!( + objective, + gradient, + hessian, + loss::SemParamsPenalty, + params, + shape1::Val{S1}, + shape2::Val{S2}, + mode::Val{H}, +) where {S1, S2, H} + sel_params = isnothing(loss.param_inds) ? params : params[loss.param_inds] + trf_params = isnothing(loss.A) ? sel_params : loss.A * sel_params + if !isnothing(loss.b) + trf_params === params && (trf_params = copy(trf_params)) + trf_params .+= loss.b + end + + weight2 = isnothing(loss.weight2) ? zero(eltype(trf_params)) : loss.weight2 + T = promote_type(eltype(trf_params), typeof(weight2)) + + obj = isnothing(objective) ? nothing : zero(T) + elm_grad = isnothing(gradient) ? nothing : fill!(similar(trf_params, T), zero(T)) + elm_hdiag = isnothing(hessian) ? nothing : fill!(similar(trf_params, T), zero(T)) + + @inbounds for i in eachindex(trf_params) + obj = _elevaluate!( + obj, + elm_grad, + elm_hdiag, + trf_params[i], + i, + nothing, + shape1, + mode, + loss.threshold, + ) + end + + if !isnothing(S2) + @inbounds for i in eachindex(trf_params) + obj = _elevaluate!( + obj, + elm_grad, + elm_hdiag, + trf_params[i], + i, + weight2, + shape2, + mode, + loss.threshold, + ) + end + end + + if !isnothing(gradient) + if isnothing(loss.A) + if isnothing(loss.param_inds) + copyto!(gradient, elm_grad) + else + fill!(gradient, zero(eltype(gradient))) + @inbounds gradient[loss.param_inds] .= elm_grad + end + elseif isnothing(loss.param_inds) + mul!(gradient, loss.At, elm_grad) + else + local_grad = loss.A' * elm_grad + fill!(gradient, 0) + @inbounds gradient[loss.param_inds] .= local_grad + end + end + + if !isnothing(hessian) + fill!(hessian, 0) + if !all(iszero, elm_hdiag) + @inbounds if isnothing(loss.A) + if isnothing(loss.param_inds) + hessian[diagind(hessian)] .= elm_hdiag + else + hessian[diagind(hessian)[loss.param_inds]] .= elm_hdiag + end + else + local_hessian = loss.A' * Diagonal(elm_hdiag) * loss.A + if isnothing(loss.param_inds) + copyto!(hessian, local_hessian) + else + hessian[loss.param_inds, loss.param_inds] .= local_hessian + end + end + end + end + + return isnothing(objective) ? NaN : obj +end + +evaluate!(objective, gradient, hessian, loss::SemParamsPenalty{H}, params) where {H} = + _evaluate!( + objective, + gradient, + hessian, + loss, + params, + Val(loss.shape1), + Val(loss.shape2), + Val(H), + ) diff --git a/src/loss/regularization/ridge.jl b/src/loss/regularization/ridge.jl deleted file mode 100644 index 813aff11c..000000000 --- a/src/loss/regularization/ridge.jl +++ /dev/null @@ -1,87 +0,0 @@ -# (Ridge) regularization - -############################################################################################ -### Types -############################################################################################ -""" -Ridge regularization. - -# Constructor - - SemRidge(;α_ridge, which_ridge, nparams, parameter_type = Float64, implied = nothing, kwargs...) - -# Arguments -- `α_ridge`: hyperparameter for penalty term -- `which_ridge::Vector`: Vector of parameter labels (Symbols) or indices that indicate which parameters should be regularized. -- `nparams::Int`: number of parameters of the model -- `implied::SemImplied`: implied part of the model -- `parameter_type`: type of the parameters - -# Examples -```julia -my_ridge = SemRidge(;α_ridge = 0.02, which_ridge = [:λ₁, :λ₂, :ω₂₃], nparams = 30, implied = my_implied) -``` - -# Interfaces -Analytic gradients and hessians are available. -""" -struct SemRidge{P, W1, W2, GT, HT} <: AbstractLoss - hessianeval::ExactHessian - α::P - which::W1 - which_H::W2 - - gradient::GT - hessian::HT -end - -############################################################################ -### Constructors -############################################################################ - -function SemRidge(; - α_ridge, - which_ridge, - nparams, - parameter_type = Float64, - implied = nothing, - kwargs..., -) - if eltype(which_ridge) <: Symbol - if isnothing(implied) - throw( - ArgumentError( - "When referring to parameters by label, `implied = ...` has to be specified", - ), - ) - else - par2ind = param_indices(implied) - which_ridge = getindex.(Ref(par2ind), which_ridge) - end - end - which_H = [CartesianIndex(x, x) for x in which_ridge] - return SemRidge( - ExactHessian(), - α_ridge, - which_ridge, - which_H, - zeros(parameter_type, nparams), - zeros(parameter_type, nparams, nparams), - ) -end - -############################################################################################ -### methods -############################################################################################ - -objective(ridge::SemRidge, par) = @views ridge.α * sum(abs2, par[ridge.which]) - -function gradient(ridge::SemRidge, par) - @views ridge.gradient[ridge.which] .= (2 * ridge.α) * par[ridge.which] - return ridge.gradient -end - -function hessian(ridge::SemRidge, par) - @views @. ridge.hessian[ridge.which_H] .= 2 * ridge.α - return ridge.hessian -end diff --git a/test/examples/political_democracy/by_parts.jl b/test/examples/political_democracy/by_parts.jl index ef634a59f..5e6aadf5e 100644 --- a/test/examples/political_democracy/by_parts.jl +++ b/test/examples/political_democracy/by_parts.jl @@ -12,13 +12,8 @@ model_ls_sym = model_ml_sym = Sem(specification = spec, data = dat, implied = RAMSymbolic) -model_ml_ridge = Sem( - specification = spec, - data = dat, - loss = (SemML, SemRidge), - α_ridge = 0.001, - which_ridge = 16:20, -) +model_ml_ridge = + Sem(SemML(SemObservedData(data = dat), RAM(spec)), SemRidge(16:20) => 0.001) model_ml_const = Sem( specification = spec, diff --git a/test/examples/political_democracy/constructor.jl b/test/examples/political_democracy/constructor.jl index 2efa5abeb..631691773 100644 --- a/test/examples/political_democracy/constructor.jl +++ b/test/examples/political_democracy/constructor.jl @@ -20,13 +20,8 @@ model_ls_sym = model_ml_sym = Sem(specification = spec, data = dat, implied = RAMSymbolic) -model_ml_ridge = Sem( - specification = spec, - data = dat, - loss = (SemML, SemRidge), - α_ridge = 0.001, - which_ridge = 16:20, -) +model_ml_ridge = + Sem(SemML(SemObservedData(data = dat), RAM(spec)), SemRidge(16:20) => 0.001) model_ml_const = Sem( specification = spec, diff --git a/test/examples/proximal/ridge.jl b/test/examples/proximal/ridge.jl index 61b7fa12a..4bf1ff828 100644 --- a/test/examples/proximal/ridge.jl +++ b/test/examples/proximal/ridge.jl @@ -36,13 +36,8 @@ model = Sem(specification = partable, data = dat, loss = SemML) sem_fit = fit(model) # use ridge from StructuralEquationModels -model_ridge = Sem( - specification = partable, - data = dat, - loss = (SemML, SemRidge), - α_ridge = 0.02, - which_ridge = 16:20, -) +model_ridge = + Sem(SemML(SemObservedData(data = dat), RAM(partable)), SemRidge(16:20) => 0.02) solution_ridge = fit(model_ridge) diff --git a/test/unit_tests/regularization.jl b/test/unit_tests/regularization.jl new file mode 100644 index 000000000..8847a28be --- /dev/null +++ b/test/unit_tests/regularization.jl @@ -0,0 +1,195 @@ +using SparseArrays, StructuralEquationModels, Test + +const SEM = StructuralEquationModels + +function toy_regularization_spec() + A = [0 0 :lambda1; 0 0 :lambda2; 0 0 0] + S = [:theta1 0 0; 0 :theta2 0; 0 0 :phi] + F = [1.0 0.0 0.0; 0.0 1.0 0.0] + + return SEM.RAMMatrices(; + A, + S, + F, + param_labels = [:lambda1, :lambda2, :theta1, :theta2, :phi], + vars = [:y1, :y2, :eta], + ) +end + +@testset "regularization constructors share SemParamsPenalty" begin + default_shape = SEM.SemParamsPenalty(1:2; shape2 = 2, weight2 = 0.75) + default_elastic = SEM.SemElasticNet(1:2) + spec = toy_regularization_spec() + + @test SEM.SemNorm(1:2) isa SEM.SemParamsPenalty + @test SEM.SemNorm(1:2; shape = 3) isa SEM.SemParamsPenalty + @test SEM.SemParamsPenalty(1:2; shape2 = 2, weight2 = 0.75) isa SEM.SemParamsPenalty + @test SEM.SemParamsPenalty(1:2) isa SEM.SemParamsPenalty + @test default_shape.shape1 == 1 + @test !(:weight1 in fieldnames(typeof(default_shape))) + @test SEM.SemLasso(1:2) isa SEM.SemParamsPenalty + @test SEM.SemRidge(1:2) isa SEM.SemParamsPenalty + @test SEM.SemElasticNet(1:2) isa SEM.SemParamsPenalty + @test default_elastic.shape1 == 1 + @test default_elastic.shape2 == 2 + @test default_elastic.weight2 == 0.5 + @test SEM.SemElasticNet(1:2; weight2 = 0.25).weight2 == 0.25 + @test SEM.SemHinge(1:2; bound = :u) isa SEM.SemParamsPenalty + @test SEM.SemHinge(1:2; bound = :u, shape1 = 2) isa SEM.SemParamsPenalty + @test SEM.SemHinge(spec, [:lambda1, :phi]; bound = :l) isa SEM.SemParamsPenalty + @test_throws MethodError SEM.SemElasticNet(1:2; shape1 = 3) + @test_throws MethodError SEM.SemElasticNet(1:2; shape2 = 4) + @test_throws MethodError SEM.SemNorm(1:2; shape2 = 2) + @test_throws MethodError SEM.SemNorm(1:2; weight = 0.25) + @test_throws MethodError SEM.SemParamsPenalty(1:2; weight1 = 0.25) + @test_throws TypeError SEM.SemParamsPenalty(1:2; shape1 = nothing) +end + +@testset "SemParamsPenalty | generic fractional and lower hinge branches" begin + fractional = SEM.SemParamsPenalty(1:2; shape1 = 1.5, shape2 = 3.0, weight2 = 0.5) + params = [4.0, 1.0] + grad = zeros(2) + hess = zeros(2, 2) + + obj = SEM.evaluate!(0.0, grad, hess, fractional, params) + + @test fractional.shape1 == 1.5 + @test fractional.shape2 == 3 + @test fractional.weight2 == 0.5 + @test obj ≈ 41.5 + @test grad ≈ [27.0, 3.0] + @test hess ≈ [12.375 0.0; 0.0 3.75] + + lower_hinge = SEM.SemHinge(1:3; bound = :l, threshold = 0.5, shape1 = 3) + params = [0.0, 1.0, 2.0] + grad = zeros(3) + hess = zeros(3, 3) + + obj = SEM.evaluate!(0.0, grad, hess, lower_hinge, params) + + @test obj ≈ 3.5 + @test grad ≈ [0.0, 0.75, 6.75] + @test hess ≈ [0.0 0.0 0.0; 0.0 3.0 0.0; 0.0 0.0 9.0] +end + +@testset "SemParamsPenalty | transformed elastic net" begin + penalty = SEM.SemParamsPenalty( + nothing, + [1.0 0.0 1.0; 0.0 1.0 0.0], + [0.0, 1.0]; + shape2 = 2, + weight2 = 0.5, + ) + params = [1.0, -2.0, 3.0] + grad = zeros(3) + hess = zeros(3, 3) + + obj = SEM.evaluate!(0.0, grad, hess, penalty, params) + + @test obj ≈ 13.5 + @test grad ≈ [5.0, -2.0, 5.0] + @test hess ≈ [1.0 0.0 1.0; 0.0 1.0 0.0; 1.0 0.0 1.0] +end + +@testset "SemParamsPenalty | SemSpecification constructors" begin + spec = toy_regularization_spec() + + selected = SEM.SemParamsPenalty(spec, [:lambda1, :phi]; shape1 = 2) + params = [2.0, 7.0, 11.0, 13.0, -3.0] + grad = zeros(5) + hess = zeros(5, 5) + + obj = SEM.evaluate!(0.0, grad, hess, selected, params) + + expected_hess = zeros(5, 5) + expected_hess[1, 1] = 2.0 + expected_hess[5, 5] = 2.0 + + @test selected.param_inds === nothing + @test size(selected.A) == (2, SEM.nparams(spec)) + @test SEM.nparams(selected) == SEM.nparams(spec) + @test obj ≈ 13.0 + @test grad ≈ [4.0, 0.0, 0.0, 0.0, -6.0] + @test hess ≈ expected_hess + + dense_transformed = SEM.SemParamsPenalty(spec, [:lambda1, :phi], [1.0 -1.0]; shape1 = 2) + params = [1.0, 0.0, 0.0, 0.0, 3.0] + grad = zeros(5) + hess = zeros(5, 5) + + obj = SEM.evaluate!(0.0, grad, hess, dense_transformed, params) + + expected_hess = zeros(5, 5) + expected_hess[1, 1] = 2.0 + expected_hess[1, 5] = -2.0 + expected_hess[5, 1] = -2.0 + expected_hess[5, 5] = 2.0 + + @test dense_transformed.param_inds == [1, 5] + @test size(dense_transformed.A) == (1, 2) + @test SEM.nparams(dense_transformed) == 2 + @test obj ≈ 4.0 + @test grad ≈ [-4.0, 0.0, 0.0, 0.0, 4.0] + @test hess ≈ expected_hess + + sparse_transformed = + SEM.SemParamsPenalty(spec, [:lambda1, :phi], sparse([1.0 1.0]); shape1 = 1) + params = [2.0, 0.0, 0.0, 0.0, -1.0] + grad = zeros(5) + hess = zeros(5, 5) + + obj = SEM.evaluate!(0.0, grad, hess, sparse_transformed, params) + + @test sparse_transformed.param_inds === nothing + @test sparse_transformed.A isa SparseMatrixCSC + @test size(sparse_transformed.A) == (1, SEM.nparams(spec)) + @test SEM.nparams(sparse_transformed) == SEM.nparams(spec) + @test obj ≈ 1.0 + @test grad ≈ [1.0, 0.0, 0.0, 0.0, 1.0] + @test hess == zeros(5, 5) + + @test_throws DimensionMismatch SEM.SemParamsPenalty(spec, [:lambda1, :phi], ones(1, 3)) +end + +@testset "SemParamsPenalty | lasso and hinge families" begin + params = [-1.0, 0.5, 2.0] + + lasso = SEM.SemLasso(1:3) + grad = zeros(3) + hess = zeros(3, 3) + obj = SEM.evaluate!(0.0, grad, hess, lasso, params) + + @test obj ≈ 3.5 + @test grad ≈ [-1.0, 1.0, 1.0] + @test hess == zeros(3, 3) + + hinge = SEM.SemHinge(1:3; bound = :u, threshold = 0.0) + fill!(grad, 0.0) + fill!(hess, 0.0) + obj = SEM.evaluate!(0.0, grad, hess, hinge, params) + + @test obj ≈ 1.0 + @test grad ≈ [-1.0, 0.0, 0.0] + @test hess == zeros(3, 3) + + squared_hinge = SEM.SemHinge(1:3; bound = :u, threshold = 0.0, shape1 = 2) + fill!(grad, 0.0) + fill!(hess, 0.0) + obj = SEM.evaluate!(0.0, grad, hess, squared_hinge, params) + + @test obj ≈ 1.0 + @test grad ≈ [-2.0, 0.0, 0.0] + @test hess ≈ [2.0 0.0 0.0; 0.0 0.0 0.0; 0.0 0.0 0.0] +end + +@testset "weighted loss terms wrap positional penalties" begin + ridge_term = SEM.LossTerm(SEM.SemRidge(1:2), nothing, 0.25) + hinge_term = SEM.LossTerm(SEM.SemHinge(1:2; bound = :u, shape1 = 2), nothing, 0.5) + + @test ridge_term isa SEM.LossTerm + @test hinge_term isa SEM.LossTerm + @test SEM.weight(ridge_term) == 0.25 + @test SEM.weight(hinge_term) == 0.5 + @test SEM.loss(ridge_term) isa SEM.SemParamsPenalty + @test SEM.loss(hinge_term) isa SEM.SemParamsPenalty +end diff --git a/test/unit_tests/unit_tests.jl b/test/unit_tests/unit_tests.jl index 4d7dad7cf..82782e9af 100644 --- a/test/unit_tests/unit_tests.jl +++ b/test/unit_tests/unit_tests.jl @@ -7,6 +7,7 @@ available_tests = Dict( "data_input_formats" => "SemObserved", "specification" => "SemSpecification", "model" => "Sem model", + "regularization" => "Regularization penalties", "StatsAPI" => "StatsAPI", )