replace_observed(): simplify & refactor

remove update_observed!()

Author: Alexey Stukalov

ext/SEMNLOptExt/NLopt.jl

@@ -107,13 +107,6 @@ function SemOptimizerNLopt(;
     )
 end
 
-############################################################################################
-### Recommended methods
-############################################################################################
-
-SEM.update_observed(optimizer::SemOptimizerNLopt, observed::SemObserved; kwargs...) =
-    optimizer
-
 ############################################################################################
 ### additional methods
 ############################################################################################

ext/SEMProximalOptExt/ProximalAlgorithms.jl

@@ -34,13 +34,6 @@ SemOptimizerProximal(;
 
 SEM.sem_optimizer_subtype(::Val{:Proximal}) = SemOptimizerProximal
 
-############################################################################################
-### Recommended methods
-############################################################################################
-
-SEM.update_observed(optimizer::SemOptimizerProximal, observed::SemObserved; kwargs...) =
-    optimizer
-
 ############################################################################
 ### Model fitting
 ############################################################################

src/additional_functions/simulation.jl

@@ -1,118 +1,3 @@
-"""
-    (1) replace_observed(model::AbstractSemSingle; kwargs...)
-
-    (2) replace_observed(model::AbstractSemSingle, observed; kwargs...)
-
-    (3) replace_observed(model::SemEnsemble; column = :group, weights = nothing, kwargs...)
-
-Return a new model with swaped observed part.
-
-# Arguments
-- `model::AbstractSemSingle`: model to swap the observed part of.
-- `kwargs`: additional keyword arguments; typically includes `data` and `specification`
-- `observed`: Either an object of subtype of `SemObserved` or a subtype of `SemObserved`
-
-# For SemEnsemble models:
-- `column`: if a DataFrame is passed as `data = ...`, which column signifies the group?
-- `weights`: how to weight the different sub-models,
-    defaults to number of samples per group in the new data
-- `kwargs`: has to be a dict with keys equal to the group names.
-    For `data` can also be a DataFrame with `column` containing the group information,
-    and for `specification` can also be an `EnsembleParameterTable`.
-
-# Examples
-See the online documentation on [Replace observed data](@ref).
-"""
-function replace_observed end
-
-"""
-    update_observed(to_update, observed::SemObserved; kwargs...)
-
-Update a `SemImplied`, `SemLossFunction` or `SemOptimizer` object to use a `SemObserved` object.
-
-# Examples
-See the online documentation on [Replace observed data](@ref).
-
-# Implementation
-You can provide a method for this function when defining a new type, for more information
-on this see the online developer documentation on [Update observed data](@ref).
-"""
-function update_observed end
-
-############################################################################################
-# change observed (data) without reconstructing the whole model
-############################################################################################
-
-# don't change non-SEM terms
-replace_observed(loss::AbstractLoss; kwargs...) = loss
-
-# use the same observed type as before
-replace_observed(loss::SemLoss; kwargs...) =
-    replace_observed(loss, typeof(SEM.observed(loss)).name.wrapper; kwargs...)
-
-# construct a new observed type
-replace_observed(loss::SemLoss, observed_type; kwargs...) =
-    replace_observed(loss, observed_type(; kwargs...); kwargs...)
-
-function replace_observed(loss::SemLoss, new_observed::SemObserved; kwargs...)
-    kwargs = Dict{Symbol, Any}(kwargs...)
-    old_observed = SEM.observed(loss)
-    implied = SEM.implied(loss)
-
-    # get field types
-    kwargs[:observed_type] = typeof(new_observed)
-    kwargs[:old_observed_type] = typeof(old_observed)
-
-    # update implied
-    new_implied = update_observed(implied, new_observed; kwargs...)
-    kwargs[:implied] = new_implied
-    kwargs[:implied_type] = typeof(new_implied)
-    kwargs[:nparams] = nparams(new_implied)
-
-    # update loss
-    return update_observed(loss, new_observed; kwargs...)
-end
-
-replace_observed(loss::LossTerm; kwargs...) =
-    LossTerm(replace_observed(loss.loss; kwargs...), loss.id, loss.weight)
-
-function replace_observed(sem::Sem; kwargs...)
-    updated_terms = Tuple(replace_observed(term; kwargs...) for term in loss_terms(sem))
-    return Sem(updated_terms...)
-end
-
-function replace_observed(
-    emodel::SemEnsemble;
-    column = :group,
-    weights = nothing,
-    kwargs...,
-)
-    kwargs = Dict{Symbol, Any}(kwargs...)
-    # allow for EnsembleParameterTable to be passed as specification
-    if haskey(kwargs, :specification) && isa(kwargs[:specification], EnsembleParameterTable)
-        kwargs[:specification] = convert(Dict{Symbol, RAMMatrices}, kwargs[:specification])
-    end
-    # allow for DataFrame with group variable "column" to be passed as new data
-    if haskey(kwargs, :data) && isa(kwargs[:data], DataFrame)
-        kwargs[:data] = Dict(
-            group =>
-                select(filter(r -> r[column] == group, kwargs[:data]), Not(column)) for
-            group in emodel.groups
-        )
-    end
-    # update each model for new data
-    models = emodel.sems
-    new_models = Tuple(
-        replace_observed(m; group_kwargs(g, kwargs)...) for
-        (m, g) in zip(models, emodel.groups)
-    )
-    return SemEnsemble(new_models...; weights = weights, groups = emodel.groups)
-end
-
-function group_kwargs(g, kwargs)
-    return Dict(k => kwargs[k][g] for k in keys(kwargs))
-end
-
 ############################################################################################
 # simulate data
 ############################################################################################

src/frontend/specification/Sem.jl

@@ -396,12 +396,108 @@ function get_SemLoss(loss, observed, implied; kwargs...)
     end
 end
 
-function update_observed(sem::Sem, new_observed; kwargs...)
-    new_terms = Tuple(
-        update_observed(lossterm.loss, new_observed; kwargs...) for
-        lossterm in loss_terms(sem)
+##############################################################
+# replace_observed: Sem level
+##############################################################
+
+"""
+    replace_observed(model::Sem, observed::SemObserved)
+    replace_observed(model::Sem, data::AbstractDict{Symbol})
+    replace_observed(model::Sem, data::AbstractDataFrame; [semterm_column])
+    replace_observed(loss::SemLoss, observed::SemObserved)
+    replace_observed(loss::SemLoss, data::Union{AbstractMatrix, DataFrame})
+
+Construct a new SEM model or SEM loss with replaced observed data.
+
+The SEM structure (implied covariance, loss type) is preserved;
+only the observed data is swapped.
+
+# Single-term models
+
+Pass a `SemObserved` object, a data matrix, or a `DataFrame`:
+```julia
+replace_observed(model, new_data_matrix)
+replace_observed(model, new_sem_observed)
+replace_observed(model, new_df)
+```
+
+# Multi-term models
+
+Pass a `Dict{Symbol}` mapping term ids to data or `SemObserved` objects:
+```julia
+replace_observed(model, Dict(:g1 => data1, :g2 => data2))
+```
+
+Or pass a `DataFrame` with a `semterm_column` identifying the group:
+```julia
+replace_observed(model, new_df; semterm_column = :group)
+```
+"""
+function replace_observed end
+
+function replace_observed(sem::Sem, data::Union{SemObserved, AbstractMatrix})
+    nsem_terms(sem) > 1 && throw(
+        ArgumentError(
+            "Model contains $(nsem_terms(sem)) SEM terms. " *
+            "Use a Dict{Symbol} or a DataFrame with `semterm_column` to provide per-term data.",
+        ),
+    )
+    updated_terms = Tuple(replace_observed(term, data) for term in loss_terms(sem))
+    return Sem(updated_terms...)
+end
+
+function replace_observed(sem::Sem, data::AbstractDict{Symbol})
+    term_ids = Set(begin
+        tid= id(term)
+        isnothing(tid) && throw(
+            ArgumentError(
+                "Multigroup replace_observed requires all SEM terms to have ids.",
+            ),
+        )
+        end for term in loss_terms(sem) if issemloss(term)
+    )
+    # check for extra ids
+    for tid in keys(data)
+        if tid ∉ term_ids
+            @warn "Data provided for term id :$tid, but no SEM term with this id exists in the model"
+        end
+    end
+
+    updated_terms = map(loss_terms(sem)) do term
+        issemloss(term) || return term
+        tid = id(term)
+        term_data = get(data, tid, nothing)
+        isnothing(term_data) && throw(
+            ArgumentError("No data provided for SEM term :$tid"),
+        )
+        return replace_observed(term, term_data)
+    end
+    return Sem(Tuple(updated_terms)...)
+end
+
+function replace_observed(
+    sem::Sem,
+    data::AbstractDataFrame;
+    semterm_column::Union{Symbol, Nothing} = nothing,
+)
+    if isnothing(semterm_column)
+        # single-term shortcut
+        nsem_terms(sem) > 1 && throw(
+            ArgumentError(
+                "Model contains $(nsem_terms(sem)) SEM terms. " *
+                "Provide `semterm_column` to specify which DataFrame column identifies the groups.",
+            ),
+        )
+        updated_terms = Tuple(replace_observed(term, data) for term in loss_terms(sem))
+        return Sem(updated_terms...)
+    end
+
+    # multi-term: split DataFrame by semterm_column
+    terms_data = Dict(
+        g[semterm_column] => group_data
+        for (g, group_data) in pairs(groupby(data, semterm_column))
     )
-    return Sem(new_terms...)
+    return replace_observed(sem, terms_data)
 end
 
 ##############################################################

src/implied/RAM/generic.jl

@@ -181,15 +181,3 @@ function update!(targets::EvaluationTargets, implied::RAM, params)
         mul!(implied.μ, implied.F⨉I_A⁻¹, implied.M)
     end
 end
-
-############################################################################################
-### Recommended methods
-############################################################################################
-
-function update_observed(implied::RAM, observed::SemObserved; kwargs...)
-    if nobserved_vars(observed) == size(implied.Σ, 1)
-        return implied
-    else
-        return RAM(; observed = observed, kwargs...)
-    end
-end

src/implied/RAM/symbolic.jl

@@ -192,18 +192,6 @@ function update!(targets::EvaluationTargets, implied::RAMSymbolic, par)
     end
 end
 
-############################################################################################
-### Recommended methods
-############################################################################################
-
-function update_observed(implied::RAMSymbolic, observed::SemObserved; kwargs...)
-    if nobserved_vars(observed) == size(implied.Σ, 1)
-        return implied
-    else
-        return RAMSymbolic(; observed = observed, kwargs...)
-    end
-end
-
 ############################################################################################
 ### additional functions
 ############################################################################################

src/implied/empty.jl

@@ -46,9 +46,3 @@ end
 ############################################################################################
 
 update!(targets::EvaluationTargets, implied::ImpliedEmpty, par) = nothing
-
-############################################################################################
-### Recommended methods
-############################################################################################
-
-update_observed(implied::ImpliedEmpty, observed::SemObserved; kwargs...) = implied

src/loss/ML/FIML.jl

@@ -159,14 +159,6 @@ function evaluate!(objective, gradient, hessian, loss::SemFIML, params)
 
     return objective
 end
-
-############################################################################################
-### Recommended methods
-############################################################################################
-
-update_observed(loss::SemFIML, observed::SemObserved; kwargs...) =
-    SemFIML(; observed = observed, kwargs...)
-
 ############################################################################################
 ### additional functions
 ############################################################################################

src/loss/ML/ML.jl

@@ -235,18 +235,3 @@ function non_posdef_return(par)
         return typemax(eltype(par))
     end
 end
-
-############################################################################################
-### recommended methods
-############################################################################################
-
-update_observed(loss::SemML, observed::SemObservedMissing; kwargs...) =
-    error("ML estimation does not work with missing data - use FIML instead")
-
-function update_observed(loss::SemML, observed::SemObserved; kwargs...)
-    if (obs_cov(loss) == obs_cov(observed)) && (obs_mean(loss) == obs_mean(observed))
-        return loss # no change
-    else
-        return SemML(observed, loss.implied; kwargs...)
-    end
-end

src/loss/ML/abstract.jl

@@ -40,3 +40,34 @@ function check_observed_vars(observed::SemObserved, implied::SemImplied)
 end
 
 check_observed_vars(sem::SemLoss) = check_observed_vars(observed(sem), implied(sem))
+
+############################################################################################
+# replace_observed: SemLoss, AbstractLoss, LossTerm
+############################################################################################
+
+function replace_observed(loss::SemLoss, new_observed::SemObserved)
+    old_obs = SEM.observed(loss)
+    observed_vars(old_obs) == observed_vars(new_observed) || throw(
+        ArgumentError(
+            "observed_vars of the new data do not match the model: " *
+            "expected $(observed_vars(old_obs)), got $(observed_vars(new_observed))",
+        ),
+    )
+    return typeof(loss).name.wrapper(new_observed, SEM.implied(loss))
+end
+
+function replace_observed(loss::SemLoss, data::Union{AbstractMatrix, DataFrame})
+    old_obs = SEM.observed(loss)
+    new_observed = typeof(old_obs).name.wrapper(
+        data = data,
+        observed_vars = observed_vars(old_obs),
+    )
+    return replace_observed(loss, new_observed)
+end
+
+# non-SEM loss terms are unchanged
+replace_observed(loss::AbstractLoss, ::Any) = loss
+
+# LossTerm: delegate to inner loss
+replace_observed(term::LossTerm, data) =
+    LossTerm(replace_observed(loss(term), data), id(term), weight(term))