SemParamsPenalty loss term for flexible regularization

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)
 ```

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)
-```
+```

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,

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)

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,

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.