bootstrap: sync with Sem updates

Author: Alexey Stukalov

src/frontend/fit/standard_errors/bootstrap.jl

@@ -1,127 +1,83 @@
 """
-    se_bootstrap(sem_fit::SemFit; n_boot = 3000, data = nothing, kwargs...)
+    se_bootstrap(semfit::SemFit, n_boot::Integer = 3000, data = nothing)
 
-Return boorstrap standard errors.
+Return bootstrap standard errors.
+
+Supports both single-group and multi-group models.
+For multi-group models, each group is resampled independently.
 
 # Arguments
-- `n_boot`: number of boostrap samples
-- `data`: data to sample from. Only needed if different than the data from `sem_fit`
-- `kwargs...`: passed down to `replace_observed`
+- `n_boot`: number of bootstrap samples
+- `data`: optional new data to sample from; for multi-group models,
+    pass a `Dict{Symbol}` mapping term ids to data matrices
 """
 function se_bootstrap(
-    semfit::SemFit{Mi, So, St, Mo, O};
-    n_boot = 3000,
+    semfit::SemFit,
+    n_boot::Integer = 3000,
     data = nothing,
-    specification = nothing,
-    kwargs...,
-) where {Mi, So, St, Mo <: AbstractSemSingle, O}
+)
+    sem = model(semfit)
+    semterms = SEM.sem_terms(sem)
+
     if isnothing(data)
-        data = samples(observed(model(semfit)))
+        if length(semterms) == 1
+            # single-group: extract data matrix
+            data = samples(observed(loss(semterms[1])))
+        else
+            # multi-group: extract per-term data dict
+            data = Dict{Symbol, AbstractMatrix}(
+                id(term) => samples(observed(loss(term))) for term in semterms
+            )
+        end
     end
 
-    data = prepare_data_bootstrap(data)
-
-    start = solution(semfit)
+    fit_params = solution(semfit)
 
-    new_solution = zero(start)
-    sum = zero(start)
-    squared_sum = zero(start)
+    # accumulator of bootstrapped params and their squares
+    params_sum = zero(fit_params)
+    params_squared = zero(fit_params)
 
-    n_failed = 0.0
-
-    converged = true
+    n_conv = 0
+    n_failed = 0
 
     for _ in 1:n_boot
-        sample_data = bootstrap_sample(data)
-        new_model = replace_observed(
-            model(semfit);
-            data = sample_data,
-            specification = specification,
-            kwargs...,
-        )
-
-        new_solution .= 0.0
+        boot_data = resample_with_replacement(data)
+        boot_model = replace_observed(sem, boot_data)
 
         try
-            new_solution = solution(fit(new_model; start_val = start))
+            params_boot = solution(fit(boot_model; start_val = fit_params))
+            params_sum .+= params_boot
+            params_squared .+= params_boot .^ 2
+            n_conv += 1
         catch
             n_failed += 1
         end
-
-        @. sum += new_solution
-        @. squared_sum += new_solution^2
-
-        converged = true
     end
 
-    n_conv = n_boot - n_failed
-    sd = sqrt.(squared_sum / n_conv - (sum / n_conv) .^ 2)
-    print("Number of nonconverged models: ", n_failed, "\n")
-    return sd
-end
-
-function se_bootstrap(
-    semfit::SemFit{Mi, So, St, Mo, O};
-    n_boot = 3000,
-    data = nothing,
-    specification = nothing,
-    kwargs...,
-) where {Mi, So, St, Mo <: SemEnsemble, O}
-    models = semfit.model.sems
-    groups = semfit.model.groups
-
-    if isnothing(data)
-        data = Dict(g => samples(observed(m)) for (g, m) in zip(groups, models))
+    if n_failed > 0
+        @warn "$n_failed bootstrap attempts did not converge"
     end
 
-    data = Dict(k => prepare_data_bootstrap(data[k]) for k in keys(data))
-
-    start = solution(semfit)
-
-    new_solution = zero(start)
-    sum = zero(start)
-    squared_sum = zero(start)
-
-    n_failed = 0.0
-
-    converged = true
-
-    for _ in 1:n_boot
-        sample_data = Dict(k => bootstrap_sample(data[k]) for k in keys(data))
-        new_model = replace_observed(
-            semfit.model;
-            data = sample_data,
-            specification = specification,
-            kwargs...,
-        )
-
-        new_solution .= 0.0
-
-        try
-            new_solution = solution(fit(new_model; start_val = start))
-        catch
-            n_failed += 1
-        end
-
-        @. sum += new_solution
-        @. squared_sum += new_solution^2
+    # calculate element-wise standard errors
+    return sqrt.(params_squared ./ n_conv .- (params_sum ./ n_conv) .^ 2)
+end
 
-        converged = true
-    end
+"""
+    resample_with_replacement(data::AbstractMatrix)
+    resample_with_replacement(data::AbstractDict)
 
-    n_conv = n_boot - n_failed
-    sd = sqrt.(squared_sum / n_conv - (sum / n_conv) .^ 2)
-    print("Number of nonconverged models: ", n_failed, "\n")
-    return sd
-end
+Resample rows of a data matrix with replacement (bootstrap sample).
 
-function prepare_data_bootstrap(data)
-    return Matrix(data)
+If the dictionary of matrices is passed (for multi-group models),
+independently resamples each matrix.
+"""
+function resample_with_replacement(data::AbstractMatrix)
+    n = size(data, 1)
+    return data[rand(1:n, n), :]
 end
 
-function bootstrap_sample(data)
-    nobs = size(data, 1)
-    index_new = rand(1:nobs, nobs)
-    data_new = data[index_new, :]
-    return data_new
+function resample_with_replacement(data::AbstractDict)
+    return typeof(data)(
+        k => resample_with_replacement(v) for (k, v) in data
+    )
 end