Merge pull request #51 from olivierlabayle/fix_adapt

update adapt
TARGENE · Mar 19, 2022 · 687e543 · 687e543 · olivierlabayle · Mar 19, 2022
2 parents b375237 + 7094b21
commit 687e543
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/src/utils.jl b/src/utils.jl
@@ -21,8 +21,9 @@ Base.merge(ndt₁::AbstractNode, ndt₂::AbstractNode) =
 
 Adapts the type of the treatment variable passed to the G learner
 """
-adapt(T::NamedTuple{<:Any, NTuple{1, Z}}) where Z = T[1]
-adapt(T) = T
+adapt(T) =
+    size(Tables.columnnames(T), 1) == 1 ? Tables.getcolumn(T, 1) : T
+
 adapt(T::AbstractNode) = node(adapt, T)
 
 

diff --git a/test/model.jl b/test/model.jl
@@ -96,13 +96,13 @@ end
 @testset "Test multiple targets with missing data" begin
     n = 100
     rng = StableRNG(123)
-    T = (t=categorical(rand(rng, [0, 1], n)),)
-    W = Tables.table(rand(rng, n, 2))
+    T = Tables.table(categorical(rand(rng, [0, 1], n)))
+    W = (w=rand(rng, n),)
     Y = (
         y₁ = vcat(rand(rng, n-10), repeat([missing], 10)),
         y₂ = vcat(repeat([missing], 20), rand(rng, n-20))
     )
-    query = Query((t=0,), (t=1,))
+    query = Query((Column1=0,), (Column1=1,))
     Q̅ = MLJ.DeterministicConstantRegressor()
     G = ConstantClassifier()