ensure random_state is used

inovex · Jan 16, 2020 · 91ed69c · 91ed69c
1 parent 9e672c8
commit 91ed69c
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
@@ -21,9 +21,9 @@
 OptRandState = Optional[Union[int, RandomState]]
 
 
-def _multi_modal_effect(covariates):
+def _multi_modal_effect(covariates, random_state):
  prob = expit(covariates[:, 7]) > 0.5
- return np.random.normal((3 * prob) + 1 * (1 - prob), 0.1)
+ return random_state.normal((3 * prob) + 1 * (1 - prob), 0.1)
 
 
 def _exponential_effect(covariates):
@@ -33,7 +33,7 @@ def _exponential_effect(covariates):
 def _multi_outcome(covariates, *, random_state: RandomState, **kwargs):
  random_state = check_random_state(random_state)
  y_0 = random_state.normal(0, 0.2, size=len(covariates))
- y_1 = y_0 + _multi_modal_effect(covariates)
+ y_1 = y_0 + _multi_modal_effect(covariates, random_state)
  mu_0, mu_1 = y_0, y_1
  return mu_0, mu_1, y_0, y_1