refactorize input validation for fit method

Author: getzze

src/robustlinearmodel.jl

@@ -95,6 +95,8 @@ function StatsAPI.fit(
     M1<:Union{Missing,<:Real},
     M2<:Union{Missing,<:Real},
 }
+    extra_args = filter_model_extra_arguments(M, kwargs)
+
     X_ismissing = eltype(X) >: Missing
     y_ismissing = eltype(y) >: Missing
     if any([y_ismissing, X_ismissing])
@@ -105,34 +107,53 @@ function StatsAPI.fit(
             )
             throw(ArgumentError(msg))
         end
-        X, y, _ = missing_omit(X, y)
+        X, y, nonmissings = missing_omit(X, y)
+
+    else
+        nonmissings = trues(size(y))
+    end
+    # drop (X,y) missing rows in extra_args
+    if all(nonmissings)
+        extra_args = NamedTuple(var => _missing_omit(val) for (var, val) in pairs(extra_args))
+    else
+        rows = findall(nonmissings)
+        extra_args = NamedTuple(var => _missing_omit(view(val, rows)) for (var, val) in pairs(extra_args))
     end
 
     # Make sure X and y have the same float eltype
     pX, py = promote_to_same_float(X, y)
+    # Make sure extra values in keyword argument have the same float eltype
+    T = eltype(py)
+    extra_args = NamedTuple(var => convert_vec_to_float(T, val) for (var, val) in pairs(extra_args))
+
+    kwargs = (; kwargs..., extra_args...)
     return fit(M, pX, py, args...; kwargs...)
 end
 
 ## Convert from formula-data to modelmatrix-response calling form
 ## the `fit` method must allow the `wts`, `contrasts` and `__formula` keyword arguments
-## Specialize to allow other keyword arguments (offset, precision...) to be taken from
-## a column of the dataframe.
+## Specialize the `model_extra_arguments` method to allow other keyword arguments
+## (offset, precision...) to be taken from a column of the dataframe.
 function StatsAPI.fit(
     ::Type{M},
     f::FormulaTerm,
     data,
     args...;
     dropmissing::Bool=false,
-    wts::Union{Nothing,Symbol,FPVector}=nothing,
     contrasts::AbstractDict{Symbol,Any}=Dict{Symbol,Any}(),
     kwargs...,
 ) where {M<:AbstractRobustModel}
     # Extract arrays from data using formula
-    f, y, X, extra = modelframe(f, data, contrasts, dropmissing, M; wts=wts)
+    f, y, X, extra = modelframe(M, f, data, contrasts, dropmissing; kwargs...)
     # Call the `fit` method with arrays
     pX, py = promote_to_same_float(X, y)
+    # Make sure extra values in keyword argument have the same float eltype
+    T = eltype(py)
+    extra = NamedTuple(var => convert_vec_to_float(T, val) for (var, val) in pairs(extra))
+
+    kwargs = (; kwargs..., extra...)
     return fit(
-        M, pX, py, args...; wts=extra.wts, contrasts=contrasts, __formula=f, kwargs...
+        M, pX, py, args...; contrasts=contrasts, __formula=f, kwargs...
     )
 end
 
@@ -442,6 +463,19 @@ The arguments `X` and `y` can be a `Matrix` and a `Vector` or a `Formula` and a
 rlm(X, y, args...; kwargs...) = fit(RobustLinearModel, X, y, args...; kwargs...)
 
 
+"""
+    model_extra_arguments(::Type{M}) where {M<:RobustLinearModel}
+
+Get the names of extra array arguments that are used by the model.
+For RobustLinearModel, [:wts, :offset].
+
+Returns an array of extra arguments used by the model.
+"""
+function model_extra_arguments(::Type{M}; kwargs...) where {M<:RobustLinearModel}
+    return [:wts, :offset]
+end
+
+
 """
     fit(::Type{M},
         X::Union{AbstractMatrix{T},SparseMatrixCSC{T}},
@@ -617,35 +651,6 @@ function StatsAPI.fit(
     return dofit ? fit!(m; fitargs...) : m
 end
 
-## Convert from formula-data to modelmatrix-response calling form
-## the `fit` method must allow the `wts`, `offset`, `contrasts` and `__formula` keyword arguments
-function StatsAPI.fit(
-    ::Type{M},
-    f::FormulaTerm,
-    data,
-    args...;
-    dropmissing::Bool=false,
-    wts::Union{Nothing,Symbol,FPVector}=nothing,
-    offset::Union{Nothing,Symbol,FPVector}=nothing,
-    contrasts::AbstractDict{Symbol,Any}=Dict{Symbol,Any}(),
-    kwargs...,
-) where {M<:RobustLinearModel}
-    # Extract arrays from data using formula
-    f, y, X, extra = modelframe(f, data, contrasts, dropmissing, M; wts=wts, offset=offset)
-    # Call the `fit` method with arrays
-    return fit(
-        M,
-        X,
-        y,
-        args...;
-        wts=extra.wts,
-        offset=extra.offset,
-        contrasts=contrasts,
-        __formula=f,
-        kwargs...,
-    )
-end
-
 
 """
     fit(::Type{M}, X, y; kwarg...) where {M<:RobustLinearModel}

src/tools.jl

@@ -15,6 +15,13 @@ function promote_to_same_float(X::AbstractMatrix, y::AbstractVector)
     return convert.(T, X)::MT, convert.(T, y)::VT
 end
 
+
+function convert_vec_to_float(::Type{T}, v::AbstractVector{<:Real}) where {T<:AbstractFloat}
+    VT = AbstractVector{T}
+    return convert.(T, v)::VT
+end
+
+
 _missing_omit(x::AbstractArray{T}) where {T} = copyto!(similar(x, nonmissingtype(T)), x)
 
 function StatsModels.missing_omit(X::AbstractMatrix, y::AbstractVector)
@@ -73,6 +80,62 @@ end
 ################################################
 
 const ModelFrameType = Tuple{FormulaTerm,<:AbstractVector,<:AbstractMatrix,NamedTuple}
+const AllowedExtraArgType = Union{Nothing,Symbol,Union{AbstractVector{<:Real},AbstractVector{Union{Missing,<:Real}}}}
+
+
+"""
+    model_extra_arguments(::Type{M}) where {M<:AbstractRobustModel}
+
+Get the names of extra array arguments (of the same size than the response ``y``)
+that are used by the model.
+For general AbstractRobustModel, `wts` is a possible keyword argument.
+
+This method should specialize to the different models to include other arguments (offset, invvar...)
+
+Returns an array of extra arguments used by the model.
+"""
+function model_extra_arguments(::Type{M}) where {M<:AbstractRobustModel}
+    return [:wts]
+end
+
+
+"""
+    filter_model_extra_arguments(::Type{M}, kwargs) where {M<:AbstractRobustModel}
+
+Filter the kwargs to keep only the extra arguments used by the model.
+The values are checked to be of type Union{Nothing,Symbol,AbstractVector{<:Real}}.
+
+For general AbstractRobustModel, `wts` keyword argument is fetched from the data table.
+
+Returns a Dict of (key, value) with key, extra arguments used by the model, and value,
+given by `kwargs`.
+"""
+function filter_model_extra_arguments(
+    ::Type{M},
+    kwargs::Union{Dict{Symbol,Any}, Base.Pairs, NamedTuple},
+) where {M<:AbstractRobustModel}
+    allowed = model_extra_arguments(M)
+
+    extra = Dict{Symbol, AllowedExtraArgType}()
+    for (k, val) in pairs(kwargs)
+        s = Symbol(k)
+        if !(s in allowed)
+            continue
+        end
+        if !isa(val, AllowedExtraArgType)
+            msg = (
+                "extra argument does not have a compatible type, should be Nothing, " *
+                "a Symbol or a real array: $(typeof(val))"
+            )
+            @warn(msg)
+            continue
+        end
+        # Add to dict
+        extra[s] = val
+    end
+    return extra
+end
+
 
 """
     modelframe(f::FormulaTerm, data, contrasts::AbstractDict, ::Type{M}; kwargs...) where M
@@ -84,15 +147,17 @@ are extracted from the `data` Table using the formula `f`.
 Adapted from GLM.jl
 """
 function modelframe(
-    f::FormulaTerm, data, contrasts::AbstractDict, dropmissing::Bool, ::Type{M}; kwargs...
+    ::Type{M}, f::FormulaTerm, data, contrasts::AbstractDict, dropmissing::Bool; kwargs...
 )::ModelFrameType where {M<:AbstractRobustModel}
     # Check is a Table
     Tables.istable(data) ||
         throw(ArgumentError("expected data in a Table, got $(typeof(data))"))
     t = Tables.columntable(data)
 
-    # Check columns exist
+    # Get columns
     cols = collect(termvars(f))
+
+    # Check columns exist
     msg = ""
     for col in cols
         msg *= checkcol(t, col)
@@ -101,7 +166,12 @@ function modelframe(
         end
     end
     msg != "" && throw(ArgumentError("Error with formula term names.\n" * msg))
-    for val in Base.values(kwargs)
+
+    # Get extra columns
+    extra_args = filter_model_extra_arguments(M, kwargs)
+
+    # Check extra columns exist
+    for val in values(extra_args)
         if isa(val, Symbol)
             msg = checkcol(t, val)
             msg != "" && throw(ArgumentError("Error with extra column name.\n" * msg))
@@ -132,15 +202,19 @@ function modelframe(
     # response and model matrix
     ## Do not copy the arrays!
     y, X = modelcols(f, t)
-    extra_vec = NamedTuple(var => (
-        if isa(val, Symbol)
-            t[val]
-        elseif isnothing(val)
-            similar(y, 0)
-        else
-            val
+
+    extra_vec = NamedTuple(
+        var => begin
+            if isa(val, Symbol)
+                t[val]
+            elseif isnothing(val)
+                similar(y, 0)
+            else
+                val
+            end
         end
-    ) for (var, val) in pairs(kwargs))
+        for (var, val) in pairs(extra_args)
+    )
 
     return f, y, X, extra_vec
 end