Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

expose n_jobs for rlearner #714

Merged
merged 1 commit into from
Dec 1, 2023
Merged

expose n_jobs for rlearner #714

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion causalml/inference/meta/rlearner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ def __init__(
control_name=0,
n_fold=5,
random_state=None,
cv_n_jobs=-1,
):
"""Initialize an R-learner.

Expand All @@ -52,6 +53,7 @@ def __init__(
control_name (str or int, optional): name of control group
n_fold (int, optional): the number of cross validation folds for outcome_learner
random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
cv_n_jobs (int, optional): number of parallel jobs to run for cross_val_predict. -1 means using all processors
"""
assert (learner is not None) or (
(outcome_learner is not None) and (effect_learner is not None)
Expand All @@ -71,6 +73,7 @@ def __init__(

self.random_state = random_state
self.cv = KFold(n_splits=n_fold, shuffle=True, random_state=random_state)
self.cv_n_jobs = cv_n_jobs

self.propensity = None
self.propensity_model = None
Expand Down Expand Up @@ -120,7 +123,7 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):

if verbose:
logger.info("generating out-of-fold CV outcome estimates")
yhat = cross_val_predict(self.model_mu, X, y, cv=self.cv, n_jobs=-1)
yhat = cross_val_predict(self.model_mu, X, y, cv=self.cv, n_jobs=self.cv_n_jobs)

for group in self.t_groups:
mask = (treatment == group) | (treatment == self.control_name)
Expand Down