replace_observed(): support kwargs

Author: Alexey Stukalov

src/frontend/finite_diff.jl

@@ -2,8 +2,8 @@ _unwrap(wrapper::SemFiniteDiff) = wrapper.model
 params(wrapper::SemFiniteDiff) = params(wrapper.model)
 loss_terms(wrapper::SemFiniteDiff) = loss_terms(wrapper.model)
 
-replace_observed(wrapper::SemFiniteDiff, data) =
-    SemFiniteDiff(replace_observed(wrapper.model, data))
+replace_observed(wrapper::SemFiniteDiff, data; kwargs...) =
+    SemFiniteDiff(replace_observed(wrapper.model, data; kwargs...))
 
 FiniteDiffLossWrappers = Union{LossFiniteDiff, SemLossFiniteDiff}
 
@@ -12,16 +12,17 @@ _unwrap(wrapper::FiniteDiffLossWrappers) = wrapper.loss
 implied(wrapper::FiniteDiffLossWrappers) = implied(_unwrap(wrapper))
 observed(wrapper::FiniteDiffLossWrappers) = observed(_unwrap(wrapper))
 
-replace_observed(wrapper::LossFiniteDiff, data) =
-    LossFiniteDiff(replace_observed(_unwrap(wrapper), data))
+replace_observed(wrapper::LossFiniteDiff, data; kwargs...) =
+    LossFiniteDiff(replace_observed(_unwrap(wrapper), data; kwargs...))
 
-replace_observed(wrapper::SemLossFiniteDiff, new_observed::SemObserved) =
-    SemLossFiniteDiff(replace_observed(_unwrap(wrapper), new_observed))
+replace_observed(wrapper::SemLossFiniteDiff, new_observed::SemObserved; kwargs...) =
+    SemLossFiniteDiff(replace_observed(_unwrap(wrapper), new_observed; kwargs...))
 
 replace_observed(
     wrapper::SemLossFiniteDiff,
-    data::Union{AbstractMatrix, DataFrame},
-) = SemLossFiniteDiff(replace_observed(_unwrap(wrapper), data))
+    data::Union{AbstractMatrix, DataFrame};
+    kwargs...,
+) = SemLossFiniteDiff(replace_observed(_unwrap(wrapper), data; kwargs...))
 
 FiniteDiffWrapper(model::AbstractSem) = SemFiniteDiff(model)
 FiniteDiffWrapper(loss::AbstractLoss) = LossFiniteDiff(loss)

src/frontend/specification/Sem.jl

@@ -475,18 +475,19 @@ replace_observed(model, new_df; semterm_column = :group)
 """
 function replace_observed end
 
-function replace_observed(sem::Sem, data::Union{SemObserved, AbstractMatrix})
+function replace_observed(sem::Sem, data::Union{SemObserved, AbstractMatrix}; kwargs...)
     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))
+    updated_terms =
+        Tuple(replace_observed(term, data; kwargs...) for term in loss_terms(sem))
     return Sem(updated_terms...)
 end
 
-function replace_observed(sem::Sem, data::AbstractDict{Symbol})
+function replace_observed(sem::Sem, data::AbstractDict{Symbol}; kwargs...)
     term_ids = Set(
         if !isnothing(id(term))
             id(term)
@@ -507,20 +508,20 @@ function replace_observed(sem::Sem, data::AbstractDict{Symbol})
         term_data = get(data, tid, nothing)
         isnothing(term_data) &&
             throw(ArgumentError("No data provided for SEM term :$tid"))
-        return replace_observed(term, term_data)
+        return replace_observed(term, term_data; kwargs...)
     end
     return Sem(Tuple(updated_terms)...)
 end
 
-function replace_observed(sem::Sem, data::AbstractVector)
+function replace_observed(sem::Sem, data::AbstractVector; kwargs...)
     nsem = nsem_terms(sem)
     nsem == length(data) || throw(
         ArgumentError(
             "Length of data ($(length(data))) does not match number of SEM terms ($nsem)",
         ),
     )
     updated_terms = map(enumerate(loss_terms(sem))) do (i, term)
-        issemloss(term) ? replace_observed(term, data[i]) : term
+        issemloss(term) ? replace_observed(term, data[i]; kwargs...) : term
     end
     return Sem(Tuple(updated_terms)...)
 end
@@ -529,6 +530,7 @@ function replace_observed(
     sem::Sem,
     data::AbstractDataFrame;
     semterm_column::Union{Symbol, Nothing} = nothing,
+    kwargs...,
 )
     if isnothing(semterm_column)
         # single-term shortcut
@@ -538,7 +540,8 @@ function replace_observed(
                 "Provide `semterm_column` to specify which DataFrame column identifies the groups.",
             ),
         )
-        updated_terms = Tuple(replace_observed(term, data) for term in loss_terms(sem))
+        updated_terms =
+            Tuple(replace_observed(term, data; kwargs...) for term in loss_terms(sem))
         return Sem(updated_terms...)
     end
 
@@ -547,7 +550,7 @@ function replace_observed(
         g[semterm_column] => group_data for
         (g, group_data) in pairs(groupby(data, semterm_column))
     )
-    return replace_observed(sem, terms_data)
+    return replace_observed(sem, terms_data; kwargs...)
 end
 
 ##############################################################

src/loss/abstract.jl

@@ -45,27 +45,28 @@ check_observed_vars(sem::SemLoss) = check_observed_vars(observed(sem), implied(s
 # replace_observed: SemLoss, AbstractLoss, LossTerm
 ############################################################################################
 
-function replace_observed(loss::SemLoss, new_observed::SemObserved)
+function replace_observed(loss::SemLoss, new_observed::SemObserved; kwargs...)
     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))",
         ),
     )
+    # the default replace_observed() does not pass through kwargs to the ctor
     return typeof(loss).name.wrapper(new_observed, SEM.implied(loss))
 end
 
-function replace_observed(loss::SemLoss, data::Union{AbstractMatrix, DataFrame})
+function replace_observed(loss::SemLoss, data::Union{AbstractMatrix, DataFrame}; kwargs...)
     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)
+    return replace_observed(loss, new_observed; kwargs...)
 end
 
 # non-SEM loss terms are unchanged
-replace_observed(loss::AbstractLoss, ::Any) = loss
+replace_observed(loss::AbstractLoss, ::Any; kwargs...) = loss
 
 # LossTerm: delegate to inner loss
-replace_observed(term::LossTerm, data) =
-    LossTerm(replace_observed(loss(term), data), id(term), weight(term))
+replace_observed(term::LossTerm, data; kwargs...) =
+    LossTerm(replace_observed(loss(term), data; kwargs...), id(term), weight(term))