using NamedTuple instead of a Dict in flat_params

JuliaAI · Aug 14, 2023 · 121e9f3 · 121e9f3
1 parent c3f0f99
commit 121e9f3
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 16 deletions.
diff --git a/src/MLJModelInterface.jl b/src/MLJModelInterface.jl
@@ -80,7 +80,7 @@ end
 export UnivariateFinite
 
 # parameter_inspection:
-export params, flat_params
+export params
 
 # model constructor + metadata
 export @mlj_model, metadata_pkg, metadata_model

diff --git a/src/parameter_inspection.jl b/src/parameter_inspection.jl
@@ -54,7 +54,7 @@ not a hard requirement.
 
 """
 flat_params(m; prefix="") = flat_params(m, Val(isamodel(m)); prefix=prefix)
-flat_params(m, ::Val{false}; prefix="") = Dict(prefix=>m)
+flat_params(m, ::Val{false}; prefix="") = NamedTuple{(Symbol(prefix),), Tuple{Any}}((m,))
 function flat_params(m, ::Val{true}; prefix="")
     fields = propertynames(m)
     prefix = prefix == "" ? "" : prefix * "__"

diff --git a/test/parameter_inspection.jl b/test/parameter_inspection.jl
@@ -1,18 +1,18 @@
 using Test
 using MLJModelInterface
 
-struct Opaque <: Model
+struct Opaque
     a::Int
 end
 
-struct Transparent <: Model
+struct Transparent
     A::Int
     B::Opaque
 end
 
 MLJModelInterface.istransparent(::Transparent) = true
 
-struct Dummy <: Model
+struct Dummy <: MLJType
     t::Transparent
     o::Opaque
     n::Integer
@@ -27,20 +27,37 @@ end
     @test params(m) == (
         t = (
             A = 6,
-            B = (
-                a = 5,
-            )
-        ),
-        o = (
-            a = 7,
+            B = Opaque(5)
         ),
+        o = Opaque(7),
         n = 42
     )
-    @test flat_params(m) == Dict(
-        "o__a" => 7,
-        "t__A" => 6,
-        "t__B__a" => 5,
-        "n" => 42
+end
+
+struct ChildModel <: Model
+    x::Int
+    y::String
+end
+
+struct ParentModel <: Model
+    x::Int
+    y::String
+    first_child::ChildModel
+    second_child::ChildModel
+end
+
+@testset "flat_params method" begin
+
+    m = ParentModel(1, "parent", ChildModel(2, "child1"),
+        ChildModel(3, "child2"))
+
+    @test MLJModelInterface.flat_params(m) == (
+        x = 1,
+        y = "parent",
+        first_child__x = 2,
+        first_child__y = "child1",
+        second_child__x = 3,
+        second_child__y = "child2"
     )
 end
 true