fix bug in composed estimand generation

src/random_variants_test.jl

@@ -8,8 +8,6 @@ NotEnoughMatchingVariantsError(rsid, p, reltol) =
  ))
 
 
-read_snps_from_csv(path::String) = unique(CSV.read(path, DataFrame; select=[:ID, :CHR]), :ID)
-
 get_variants_to_randomize(filepath::AbstractString) = Set(open(readlines, filepath))
 
 
@@ -169,17 +167,13 @@ function find_maf_matching_random_variants(
  return variant_map
 end
 
-"""
-Each sub estimand can potentially lead to p new sub estimands. A D-dimensional composite estimand
-would lead to p^D new composed estimands. Typically p = 10, D = 9 for interactions leading to too many 
-new estimands. Instead we randomly select p for each subestimands, leading to p new composed estimands.
-"""
-function make_random_variants_estimands(Ψ::ComposedEstimand, variant_map; p=10, rng=MersenneTwister())
- newargs = Tuple(rand(rng, make_random_variants_estimands(arg, variant_map), p) for arg ∈ Ψ.args)
+
+function make_random_variants_estimands(Ψ::ComposedEstimand, variant_map)
+ newargs = Tuple(make_random_variants_estimands(arg, variant_map) for arg ∈ Ψ.args)
  return [ComposedEstimand(Ψ.f, Tuple(args)) for args ∈ zip(newargs...)]
 end
 
-function make_random_variants_estimands(Ψ::T, variant_map; kwargs...) where T <: TMLE.Estimand
+function make_random_variants_estimands(Ψ::T, variant_map) where T <: TMLE.Estimand
  transactors = keys(variant_map)
  origin_treatment_variables = keys(Ψ.treatment_values)
  # At least one trans-actor in the parameter treatments to be processed
@@ -217,8 +211,8 @@ function make_random_variants_estimands(Ψ::T, variant_map; kwargs...) where T <
  return new_estimands
 end
 
-make_random_variants_estimands(estimands, variant_map; p=10, rng=MersenneTwister(123)) = 
- vcat((make_random_variants_estimands(Ψ, variant_map; p=p, rng=rng) for Ψ in estimands)...)
+make_random_variants_estimands(estimands, variant_map) = 
+ vcat((make_random_variants_estimands(Ψ, variant_map) for Ψ in estimands)...)
 
 function generate_random_variants_parameters_and_dataset(parsed_args)
  resultsfile = parsed_args["results"]
@@ -243,7 +237,7 @@ function generate_random_variants_parameters_and_dataset(parsed_args)
  p=p, rng=rng, reltol=reltol, verbosity=verbosity
  )
  verbosity > 0 && @info string("Building new estimands from matched random variants.")
- new_estimands = make_random_variants_estimands(significant_estimands, variant_map; rng=rng, p=p)
+ new_estimands = make_random_variants_estimands(significant_estimands, variant_map)
  new_estimands = groups_ordering(new_estimands, brute_force=false, do_shuffle=false)
  serialize(out, Configuration(estimands = new_estimands))
  verbosity > 0 && @info string("Done.")

test/random_variants_test.jl

@@ -140,7 +140,7 @@ end
  trans_actors; 
  p=p, rng=rng, reltol=reltol, verbosity=0
  )
- new_estimands = NegativeControl.make_random_variants_estimands(estimands, variant_map, p=p)
+ new_estimands = NegativeControl.make_random_variants_estimands(estimands, variant_map)
  # 5 new estimands for each of the 4 input estimands
  @test length(new_estimands) == 20
  # Check 5 new estimands generated from estimand 1
@@ -156,6 +156,12 @@ end
  @test expected_new_rsids == new_rsids
  # Check 5 new estimands generated from estimand 3 (composed estimand)
  for Ψ ∈ new_estimands[11:15]
+ @test all(arg.outcome == first(Ψ.args).outcome for arg in Ψ.args)
+ treatment_settings = [arg.treatment_values for arg in Ψ.args]
+ treatment_variables = [keys(ts) for ts in treatment_settings]
+ @test length(unique(treatment_variables)) == 1
+ treatment_values = [values(ts) for ts in treatment_settings]
+ @test length(unique(treatment_values)) == length(treatment_variables)
  @test Ψ isa ComposedEstimand
  @test Ψ != estimands[3]
  end