check_same_semterm_type(): refactor check_single_lossfun()

Author: Alexey Stukalov

src/frontend/fit/fitmeasures/RMSEA.jl

@@ -22,7 +22,7 @@ For multigroup models, the correction proposed by J.H. Steiger is applied
 RMSEA(fit::SemFit) = RMSEA(fit, fit.model)
 
 function RMSEA(fit::SemFit, model::AbstractSem)
-    term_type = check_single_lossfun(model; throw_error = true)
+    term_type = check_same_semterm_type(model; throw_error = true)
     n = nsamples(fit) + nsem_terms(model) * multigroup_correction_scale(term_type)
     sqrt(nsem_terms(model)) * RMSEA(dof(fit), χ²(fit), n)
 end

src/frontend/fit/fitmeasures/chi2.jl

@@ -16,7 +16,7 @@ function χ²(fit::SemFit, model::AbstractSem)
     terms = sem_terms(model)
     isempty(terms) && return 0.0
 
-    L = check_single_lossfun(model; throw_error = true)
+    L = check_same_semterm_type(model; throw_error = true)
     term1 = _unwrap(loss(terms[1]))
     L = typeof(term1).name
 

src/frontend/fit/fitmeasures/minus2ll.jl

@@ -64,6 +64,6 @@ end
 ############################################################################################
 
 function minus2ll(model::AbstractSem, fit::SemFit)
-    check_single_lossfun(model; throw_error = true)
+    check_same_semterm_type(model; throw_error = true)
     sum(Base.Fix2(minus2ll, fit) ∘ _unwrap ∘ loss, sem_terms(model))
 end

src/frontend/specification/Sem.jl

@@ -62,7 +62,7 @@ multigroup_correction_scale(loss::SemLoss) =
 function multigroup_weights(semterms...)
     n = length(semterms)
     nsamples_total = sum(nsamples, semterms)
-    semloss_type = check_single_lossfun(semterms; throw_error = false)
+    semloss_type = check_same_semterm_type(semterms; throw_error = false)
     if isnothing(semloss_type)
         @info """
         Your ensemble model contains heterogeneous loss functions.
@@ -274,32 +274,46 @@ function sem_term(model::AbstractSem, id::Nothing = nothing)
     error("Unreachable reached")
 end
 
-# check that all SEM terms in the model use the same loss type
-function check_same_semloss_type(model::AbstractSem; throw_error::Bool = true)
-    terms = sem_terms(model)
+# check that all models use the same single loss function
+# returns the type of the single SEM loss function, SemLoss if there are multiple different SEM losses,
+# nothing if there are no SEM terms.
+# If throw_error=true, throws an error if there are multiple different SEM loss functions
+check_same_semterm_type(model::AbstractSem; throw_error::Bool = true) =
+    check_same_semterm_type(sem_terms(model); throw_error = throw_error)
+
+# check that all models use the same single loss function
+# returns the type of the single SEM loss function,
+# nothing if there are multiple different SEM losses or no SEM terms.
+# If throw_error=true, throws an error if there are multiple different SEM loss functions
+function check_same_semterm_type(terms::Tuple; throw_error::Bool = true)
     isempty(terms) && return nothing
 
-    L = typeof(lossfun)
-end
-
-function check_single_lossfun(models::AbstractSemSingle...; throw_error)
-    uniform = true
-    lossfun = models[1].loss.functions[1]
-    L = typeof(lossfun)
-    for (i, model) in enumerate(models)
-        uniform &= check_single_lossfun(model; throw_error = throw_error)
-        cur_lossfun = model.loss.functions[1]
-        if !isa(cur_lossfun, L) & throw_error
-            @error "Loss function for group #$i model is $(typeof(cur_lossfun)), expected $L.
-            Heterogeneous loss functions are not supported."
+    _semloss(term::SemLoss) = _unwrap(term)
+    _semloss(term::LossTerm) = _semloss(loss(term))
+    _semloss(term) = throw(ArgumentError("SemLoss term expected, $(typeof(term)) found"))
+    _semloss_label(i::Integer, _::Union{SemLoss, LossTerm{<:SemLoss, Nothing}}) = "#$i"
+    _semloss_label(i::Integer, term::LossTerm{<:SemLoss, Symbol}) = "#$i ($(SEM.id(term)))"
+
+    term1 = _semloss(terms[1])
+    L = typeof(term1).name
+
+    # check that all SemLoss terms are of the same class (ML, FIML, WLS etc), ignore typeparams
+    for (i, term) in enumerate(terms)
+        lossterm = _semloss(term)
+        @assert lossterm isa SemLoss
+        if typeof(lossterm).name != L
+            if throw_error
+                error("SemLoss term $(_semloss_label(i, term)) is $(typeof(lossterm).name), expected $L. Heterogeneous loss functions are not supported")
+            else
+                return nothing
+            end
         end
-        uniform &= isa(cur_lossfun, L)
     end
-    return uniform
-end
 
-check_single_lossfun(model::SemEnsemble; throw_error) =
-    check_single_lossfun(model.sems...; throw_error)
+    # return the type of the first SEM term
+    # note that type params of the SEM terms might be different
+    return typeof(term1)
+end
 
 # wrappers arounds a single SemLoss term
 observed(model::AbstractSem, id::Nothing = nothing) = observed(sem_term(model, id))