fix wrong covariate usage in IHDP generators

inovex · Feb 18, 2020 · 51abe9d · 51abe9d
1 parent cb31410
commit 51abe9d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/justcause/data/generators/ihdp.py b/src/justcause/data/generators/ihdp.py
@@ -22,12 +22,12 @@
 
 
 def _multi_modal_effect(covariates, random_state):
- prob = expit(covariates[:, 7]) > 0.5
+ prob = expit(covariates[:, 0]) > 0.5
  return random_state.normal((3 * prob) + 1 * (1 - prob), 0.1)
 
 
 def _exponential_effect(covariates):
- return np.exp(1 + expit(covariates[:, 7]))
+ return np.exp(1 + expit(covariates[:, 0]))
 
 
 def _multi_outcome(covariates, *, random_state: RandomState, **kwargs):
@@ -48,7 +48,7 @@ def _expo_outcome(covariates, *, random_state: RandomState, **kwargs):
 
 def _treatment_assignment(covariates, *, random_state: RandomState, **kwargs):
  random_state = check_random_state(random_state)
- return random_state.binomial(1, p=expit(covariates[:, 7]))
+ return random_state.binomial(1, p=expit(covariates[:, 0]))
 
 
 def multi_expo_on_ihdp(