check_same_semterm_type(): refactor check_single_lossfun()

Author: Alexey Stukalov

src/additional_functions/helper.jl

@@ -116,35 +116,6 @@ function nonunique(values::AbstractVector)
     return res
 end
 
-# check that a model only has a single lossfun
-function check_single_lossfun(model::AbstractSemSingle; throw_error)
-    if (length(model.loss.functions) > 1) & throw_error
-        @error "The model has $(length(sem.loss.functions)) loss functions.
-            Only a single loss function is supported."
-    end
-    return isone(length(model.loss.functions))
-end
-
-# check that all models use the same single loss function
-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."
-        end
-        uniform &= isa(cur_lossfun, L)
-    end
-    return uniform
-end
-
-check_single_lossfun(model::SemEnsemble; throw_error) =
-    check_single_lossfun(model.sems...; throw_error)
-
 # scaling corrections for multigroup models
 mg_correction(::SemFIML) = 0
 mg_correction(::SemML) = 0

src/frontend/fit/fitmeasures/RMSEA.jl

@@ -27,7 +27,7 @@ RMSEA_corr_scale(::Type{<:SemML}) = -1
 RMSEA_corr_scale(::Type{<:SemWLS}) = -1
 
 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) * RMSEA_corr_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

@@ -62,6 +62,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

@@ -51,8 +51,8 @@ end
 
 function multigroup_weights(models, n)
     nsamples_total = sum(nsamples, models)
-    uniform_lossfun = check_single_lossfun(models...; throw_error = false)
-    if !uniform_lossfun
+    semloss_type = check_same_semterm_type(semterms; throw_error = false)
+    if isnothing(semloss_type)
         @info """
         Your ensemble model contains heterogeneous loss functions.
         Default weights of (#samples per group/#total samples) will be used
@@ -260,6 +260,47 @@ function sem_term(model::AbstractSem, _::Nothing = nothing)
     error("Unreachable reached")
 end
 
+# 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
+
+    _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
+    end
+
+    # 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))
 implied(model::AbstractSem, id::Nothing = nothing) = implied(sem_term(model, id))