fix counterfactual treatment with order added

TARGENE · Apr 8, 2023 · e7ee755 · e7ee755
1 parent ca20ac8
commit e7ee755
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/src/utils.jl b/src/utils.jl
@@ -121,7 +121,7 @@ function counterfactualTreatment(vals, T)
  Tnames = Tables.columnnames(T)
  n = nrows(T)
  NamedTuple{Tnames}(
- [categorical(repeat([vals[i]], n), levels=levels(Tables.getcolumn(T, name)))
+ [categorical(repeat([vals[i]], n), levels=levels(Tables.getcolumn(T, name)), ordered=isordered(Tables.getcolumn(T, name)))
  for (i, name) in enumerate(Tnames)])
 end
 

diff --git a/test/utils.jl b/test/utils.jl
@@ -141,14 +141,17 @@ end
 @testset "Test counterfactualTreatment" begin
  vals = (true, "a")
  T = (
- t₁ = categorical([true, false, false]),
+ t₁ = categorical([true, false, false], ordered=true),
  t₂ = categorical(["a", "a", "c"])
  )
  cfT = TMLE.counterfactualTreatment(vals, T)
  @test cfT == (
  t₁ = categorical([true, true, true]),
  t₂ = categorical(["a", "a", "a"])
  )
+ @test isordered(cfT.t₁)
+ @test !isordered(cfT.t₂)
+
 end
 
 @testset "Test compute_covariate" begin