integrate new traits: predict_scitype, etc

Author: ablaom

src/MLJModelInterface.jl

@@ -4,6 +4,10 @@ const MODEL_TRAITS = [
     :input_scitype,
     :output_scitype,
     :target_scitype,
+    :training_scitype,
+    :predict_scitype,
+    :transform_scitype,
+    :inverse_transform_scitype,
     :is_pure_julia,
     :package_name,
     :package_license,
@@ -18,6 +22,7 @@ const MODEL_TRAITS = [
     :name,
     :is_supervised,
     :prediction_type,
+    :abstract_type,
     :implemented_methods,
     :hyperparameters,
     :hyperparameter_types,

src/model_traits.jl

@@ -16,3 +16,73 @@ StatisticalTraits.prediction_type(::Type{<:Interval})      = :interval
 implemented_methods(M::Type) = implemented_methods(get_interface_mode(), M)
 implemented_methods(model) = implemented_methods(typeof(model))
 implemented_methods(::LightInterface, M) = errlight("implemented_methods")
+
+for M in ABSTRACT_MODEL_SUBTYPES
+    @eval(StatisticalTraits.abstract_type(::Type{<:$M}) = $M)
+end
+
+StatisticalTraits.training_scitype(M::Type{<:Model}) = input_scitype(M)
+StatisticalTraits.training_scitype(::Type{<:Static}) = Tuple{}
+function StatisticalTraits.training_scitype(M::Type{<:Supervised})
+    I = input_scitype(M)
+    T = target_scitype(M)
+    ret = Tuple{I,T}
+    if supports_weights(M)
+        W = AbstractVector{Union{Continuous,Count}} # weight scitype
+        return Union{ret,Tuple{I,T,W}}
+    elseif supports_class_weights(M)
+        W = AbstractDict{Finite,Union{Continuous,Count}}
+        return Union{ret,Tuple{I,T,W}}
+    end
+    return ret
+end
+
+StatisticalTraits.transform_scitype(M::Type{<:Unsupervised}) =
+    output_scitype(M)
+
+StatisticalTraits.inverse_transform_scitype(M::Type{<:Unsupervised}) =
+    input_scitype(M)
+
+StatisticalTraits.predict_scitype(M::Type{<:Deterministic}) = target_scitype(M)
+
+
+## FALLBACKS FOR `predict_scitype` FOR `Probabilistic` MODELS
+
+# This seems less than ideal but should reduce the number of `Unknown`
+# in `prediction_type` for models which, historically, have not
+# implemented the trait.
+
+StatisticalTraits.predict_scitype(M::Type{<:Probabilistic}) =
+    _density(target_scitype(M))
+
+_density(::Any) = Unknown
+for T in [:Continuous, :Count, :Textual]
+    eval(quote
+         _density(::Type{AbstractArray{$T,D}}) where D =
+         AbstractArray{Density{$T},D}
+         end)
+end
+
+for T in [:Finite,
+          :Multiclass,
+          :OrderedFactor,
+          :Infinite,
+          :Continuous,
+          :Count,
+          :Textual]
+    eval(quote
+         _density(::Type{AbstractArray{<:$T,D}}) where D =
+         AbstractArray{Density{<:$T},D}
+         _density(::Type{Table($T)}) = Table(Density{$T})
+         end)
+end
+
+for T in [:Finite, :Multiclass, :OrderedFactor]
+    eval(quote
+         _density(::Type{AbstractArray{<:$T{N},D}}) where {N,D} =
+         AbstractArray{Density{<:$T{N}},D}
+         _density(::Type{AbstractArray{$T{N},D}}) where {N,D} =
+         AbstractArray{Density{$T{N}},D}
+         _density(::Type{Table($T{N})}) where N = Table(Density{$T{N}})
+         end)
+end

test/model_traits.jl

@@ -79,3 +79,67 @@ import .Fruit
     @test docstring(Float64) == "Float64"
     @test docstring(Fruit.Banana) == "Banana"
 end
+
+@testset "`_density` - helper for predict_scitype fallback" begin
+    for T in [Continuous, Count, Textual]
+        @test M._density(AbstractArray{T,3}) ==
+            AbstractArray{Density{T},3}
+    end
+
+    for T in [Finite,
+              Multiclass,
+              OrderedFactor,
+              Infinite,
+              Continuous,
+              Count,
+              Textual]
+        @test M._density(AbstractVector{<:T}) ==
+            AbstractVector{Density{<:T}}
+        @test M._density(Table(T)) == Table(Density{T})
+    end
+
+    for T in [Finite, Multiclass, OrderedFactor]
+        @test M._density(AbstractArray{<:T{2},3}) ==
+            AbstractArray{Density{<:T{2}},3}
+        @test M._density(AbstractArray{T{2},3}) ==
+            AbstractArray{Density{T{2}},3}
+        @test M._density(Table(T{2})) == Table(Density{T{2}})
+    end
+end
+
+@mlj_model mutable struct P2 <: Probabilistic end
+M.target_scitype(::Type{<:P2}) = AbstractVector{<:Multiclass}
+M.input_scitype(::Type{<:P2}) = Table(Continuous)
+
+@mlj_model mutable struct U2 <: Unsupervised end
+M.output_scitype(::Type{<:U2}) = AbstractVector{<:Multiclass}
+M.input_scitype(::Type{<:U2}) = Table(Continuous)
+
+@mlj_model mutable struct S2 <: Static end
+M.output_scitype(::Type{<:S2}) = AbstractVector{<:Multiclass}
+M.input_scitype(::Type{<:S2}) = Table(Continuous)
+
+@testset "operation scitypes" begin
+    @test predict_scitype(P2()) == AbstractVector{Density{<:Multiclass}}
+    @test transform_scitype(P2()) == Unknown
+    @test transform_scitype(U2()) == AbstractVector{<:Multiclass}
+    @test inverse_transform_scitype(U2()) == Table(Continuous)
+    @test predict_scitype(U2()) == Unknown
+    @test transform_scitype(S2()) == AbstractVector{<:Multiclass}
+    @test inverse_transform_scitype(S2()) == Table(Continuous)
+end
+
+@testset "abstract_type, training_scitype" begin
+    @test abstract_type(P2()) == Probabilistic
+    @test abstract_type(S1()) == Supervised
+    @test abstract_type(U1()) == Unsupervised
+    @test abstract_type(D1()) == Deterministic
+    @test abstract_type(P1()) == Probabilistic
+
+    @test training_scitype(P2()) ==
+        Tuple{Table(Continuous),AbstractVector{<:Multiclass}}
+    @test training_scitype(U2()) == Table(Continuous)
+    @test training_scitype(S2()) == Tuple{}
+end
+
+true