make sure to access dataframe differently

src/justcause/data/utils.py

@@ -161,7 +161,10 @@ def generate_data(
  ), "Covariate function should return a dataframe with `n_samples` rows"
  else:
  indices = random_state.choice(covariates.shape[0], n_samples, replace=False)
- covariates = covariates[indices, :]
+ if isinstance(covariates, pd.DataFrame):
+ covariates = covariates.iloc[indices, :] # ensure proper access
+ else:
+ covariates = covariates[indices, :]
 
  if covariate_names is None:
  if isinstance(covariates, pd.DataFrame):

tests/test_data_generators.py

@@ -62,6 +62,14 @@ def covariates(_, *, random_state, **kwargs):
  assert len(df) == len(ihdp_cov)
  assert len(set(df.columns).intersection(set(cov_names))) == 25
 
+ # Works with DataFrame and np.array when n_samples is given
+ gen = generate_data(
+ covariates=ihdp_cov, treatment=treatment, outcomes=outcomes, n_samples=500,
+ )
+ df = list(gen)[0]
+ assert len(df) == 500
+ assert len(set(df.columns).intersection({f"x_{i}" for i in range(25)})) == 25
+
 
 def test_ihdp_generator():
  gen = multi_expo_on_ihdp(setting="multi-modal", n_replications=10)