- Adds benchmarking script in comments and shares latest numbers and plots

- Adds test for meta-learner consistency and key attributes

Author: Ic3fr0g

causalml/inference/meta/drlearner.py

@@ -239,6 +239,8 @@ def predict(
 
         # models_mu_c is fold-specific but not group-specific; predict once and reuse.
         yhat_c = np.r_[[model.predict(X) for model in self.models_mu_c]].mean(axis=0)
+        # Shared-reference dict preserves the public yhat_cs[group] API cheaply.
+        yhat_cs = {group: yhat_c for group in self.t_groups}
 
         for i, group in enumerate(self.t_groups):
             models_tau = self.models_tau[group]
@@ -264,7 +266,7 @@ def predict(
         if not return_components:
             return te
         else:
-            return te, yhat_c, yhat_ts
+            return te, yhat_cs, yhat_ts
 
     def fit_predict(
         self,
@@ -394,11 +396,11 @@ def estimate_ate(
             The mean and confidence interval (LB, UB) of the ATE estimate.
         """
         if pretrain:
-            te, yhat_c, yhat_ts = self.predict(
+            te, yhat_cs, yhat_ts = self.predict(
                 X, treatment, y, p, return_components=True
             )
         else:
-            te, yhat_c, yhat_ts = self.fit_predict(
+            te, yhat_cs, yhat_ts = self.fit_predict(
                 X, treatment, y, p, return_components=True, seed=seed
             )
         X, treatment, y = convert_pd_to_np(X, treatment, y)
@@ -427,17 +429,17 @@ def estimate_ate(
             w = (treatment_filt == group).astype(int)
             prob_treatment = float(sum(w)) / w.shape[0]
 
-            yhat_c_g = yhat_c[mask]
+            yhat_c = yhat_cs[group][mask]
             yhat_t = yhat_ts[group][mask]
             y_filt = y[mask]
 
             # SE formula is based on the lower bound formula (7) from Imbens, Guido W., and Jeffrey M. Wooldridge. 2009.
             # "Recent Developments in the Econometrics of Program Evaluation." Journal of Economic Literature
             se = np.sqrt(
                 (
-                    (y_filt[w == 0] - yhat_c_g[w == 0]).var() / (1 - prob_treatment)
+                    (y_filt[w == 0] - yhat_c[w == 0]).var() / (1 - prob_treatment)
                     + (y_filt[w == 1] - yhat_t[w == 1]).var() / prob_treatment
-                    + (yhat_t - yhat_c_g).var()
+                    + (yhat_t - yhat_c).var()
                 )
                 / y_filt.shape[0]
             )
@@ -600,6 +602,7 @@ def predict(
         yhat_c = np.r_[
             [model.predict_proba(X)[:, 1] for model in self.models_mu_c]
         ].mean(axis=0)
+        yhat_cs = {group: yhat_c for group in self.t_groups}
 
         for i, group in enumerate(self.t_groups):
             models_tau = self.models_tau[group]
@@ -625,7 +628,7 @@ def predict(
         if not return_components:
             return te
         else:
-            return te, yhat_c, yhat_ts
+            return te, yhat_cs, yhat_ts
 
 
 class XGBDRRegressor(BaseDRRegressor):

causalml/inference/meta/tlearner.py

@@ -55,6 +55,9 @@ def __init__(
         else:
             self.model_c = control_learner
 
+        # Preserve the unfitted template so repeated fit() calls always start fresh.
+        self._model_c_template = self.model_c
+
         if treatment_learner is None:
             self.model_t = deepcopy(learner)
         else:
@@ -85,10 +88,13 @@ def fit(self, X, treatment, y, p=None):
         self.models_t = {group: deepcopy(self.model_t) for group in self.t_groups}
 
         # model_c is trained on the control group, which is identical for every
-        # treatment group, so fit it once and store as a single model (not a dict).
+        # treatment group, so fit it once. Deepcopy from the unfitted template so
+        # re-calling fit() always starts from a clean state (safe with warm_start).
         control_mask = treatment == self.control_name
-        self.model_c = deepcopy(self.model_c)
+        self.model_c = deepcopy(self._model_c_template)
         self.model_c.fit(X[control_mask], y[control_mask])
+        # Expose as a shared-reference dict to preserve the public models_c API.
+        self.models_c = {group: self.model_c for group in self.t_groups}
 
         for group in self.t_groups:
             treatment_mask = treatment == group
@@ -112,6 +118,9 @@ def predict(
         yhat_ts = {}
 
         yhat_c = self.model_c.predict(X)
+        # Build a shared-reference dict so return_components callers keep the
+        # yhat_cs[group] indexing API without duplicating the underlying array.
+        yhat_cs = {group: yhat_c for group in self.t_groups}
 
         for group in self.t_groups:
             yhat_ts[group] = self.models_t[group].predict(X)
@@ -136,7 +145,7 @@ def predict(
         if not return_components:
             return te
         else:
-            return te, yhat_c, yhat_ts
+            return te, yhat_cs, yhat_ts
 
     def fit_predict(
         self,
@@ -195,6 +204,7 @@ def fit_predict(
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.model_c = deepcopy(model_c_global)
+            self.models_c = {group: self.model_c for group in self.t_groups}
             self.models_t = deepcopy(models_t_global)
 
             return (te, te_lower, te_upper)
@@ -225,9 +235,9 @@ def estimate_ate(
         """
         X, treatment, y = convert_pd_to_np(X, treatment, y)
         if pretrain:
-            te, yhat_c, yhat_ts = self.predict(X, treatment, y, return_components=True)
+            te, yhat_cs, yhat_ts = self.predict(X, treatment, y, return_components=True)
         else:
-            te, yhat_c, yhat_ts = self.fit_predict(
+            te, yhat_cs, yhat_ts = self.fit_predict(
                 X, treatment, y, return_components=True
             )
 
@@ -244,14 +254,14 @@ def estimate_ate(
             w = (treatment_filt == group).astype(int)
             prob_treatment = float(sum(w)) / w.shape[0]
 
-            yhat_c_g = yhat_c[mask]
+            yhat_c = yhat_cs[group][mask]
             yhat_t = yhat_ts[group][mask]
 
             se = np.sqrt(
                 (
-                    (y_filt[w == 0] - yhat_c_g[w == 0]).var() / (1 - prob_treatment)
+                    (y_filt[w == 0] - yhat_c[w == 0]).var() / (1 - prob_treatment)
                     + (y_filt[w == 1] - yhat_t[w == 1]).var() / prob_treatment
-                    + (yhat_t - yhat_c_g).var()
+                    + (yhat_t - yhat_c).var()
                 )
                 / y_filt.shape[0]
             )
@@ -289,6 +299,7 @@ def estimate_ate(
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.model_c = deepcopy(model_c_global)
+            self.models_c = {group: self.model_c for group in self.t_groups}
             self.models_t = deepcopy(models_t_global)
 
             return ate, ate_lower, ate_upper
@@ -371,6 +382,7 @@ def predict(
         yhat_ts = {}
 
         yhat_c = self.model_c.predict_proba(X)[:, 1]
+        yhat_cs = {group: yhat_c for group in self.t_groups}
 
         for group in self.t_groups:
             yhat_ts[group] = self.models_t[group].predict_proba(X)[:, 1]
@@ -395,7 +407,7 @@ def predict(
         if not return_components:
             return te
         else:
-            return te, yhat_c, yhat_ts
+            return te, yhat_cs, yhat_ts
 
 
 class XGBTRegressor(BaseTRegressor):

causalml/inference/meta/xlearner.py

@@ -56,6 +56,9 @@ def __init__(
         else:
             self.model_mu_c = control_outcome_learner
 
+        # Preserve the unfitted template so repeated fit() calls always start fresh.
+        self._model_mu_c_template = self.model_mu_c
+
         if treatment_outcome_learner is None:
             self.model_mu_t = deepcopy(learner)
         else:
@@ -125,14 +128,18 @@ def fit(self, X, treatment, y, p=None):
         self.vars_t = {}
 
         # model_mu_c is trained on control data, which is the same for every treatment
-        # group, so fit it once and store as a single model (not a per-group dict).
+        # group. Deepcopy from the unfitted template so re-calling fit() starts fresh.
         control_mask = treatment == self.control_name
-        self.model_mu_c = deepcopy(self.model_mu_c)
+        self.model_mu_c = deepcopy(self._model_mu_c_template)
         self.model_mu_c.fit(X[control_mask], y[control_mask])
+        # Expose as a shared-reference dict to preserve the public models_mu_c API.
+        self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
 
-        # var_c depends only on model_mu_c and control data, both constant across groups.
+        # var_c depends only on model_mu_c and control data — constant across groups.
         y_control_pred = self.model_mu_c.predict(X[control_mask])
-        var_c = (y[control_mask] - y_control_pred).var()
+        self.var_c = (y[control_mask] - y_control_pred).var()
+        # Keep vars_c dict for backward compatibility with existing callers.
+        self.vars_c = {group: self.var_c for group in self.t_groups}
 
         for group in self.t_groups:
             treatment_mask = treatment == group
@@ -141,9 +148,9 @@ def fit(self, X, treatment, y, p=None):
 
             self.models_mu_t[group].fit(X_treat, y_treat)
 
-            self.vars_c[group] = var_c
-            var_t = (y_treat - self.models_mu_t[group].predict(X_treat)).var()
-            self.vars_t[group] = var_t
+            self.vars_t[group] = (
+                y_treat - self.models_mu_t[group].predict(X_treat)
+            ).var()
 
             # Train treatment effect models using cross-group imputation
             d_c = self.models_mu_t[group].predict(X[control_mask]) - y[control_mask]
@@ -289,6 +296,7 @@ def fit_predict(
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.model_mu_c = deepcopy(model_mu_c_global)
+            self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
             self.models_mu_t = deepcopy(models_mu_t_global)
             self.models_tau_c = deepcopy(models_tau_c_global)
             self.models_tau_t = deepcopy(models_tau_t_global)
@@ -367,7 +375,7 @@ def estimate_ate(
             se = np.sqrt(
                 (
                     self.vars_t[group] / prob_treatment
-                    + self.vars_c[group] / (1 - prob_treatment)
+                    + self.var_c / (1 - prob_treatment)
                     + (p_filt * dhat_c + (1 - p_filt) * dhat_t).var()
                 )
                 / w.shape[0]
@@ -408,6 +416,7 @@ def estimate_ate(
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.model_mu_c = deepcopy(model_mu_c_global)
+            self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
             self.models_mu_t = deepcopy(models_mu_t_global)
             self.models_tau_c = deepcopy(models_tau_c_global)
             self.models_tau_t = deepcopy(models_tau_t_global)
@@ -546,12 +555,14 @@ def fit(self, X, treatment, y, p=None):
         # model_mu_c is trained on control data, which is the same for every treatment
         # group, so fit it once and store as a single model (not a per-group dict).
         control_mask = treatment == self.control_name
-        self.model_mu_c = deepcopy(self.model_mu_c)
+        self.model_mu_c = deepcopy(self._model_mu_c_template)
         self.model_mu_c.fit(X[control_mask], y[control_mask])
+        self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
 
-        # var_c depends only on model_mu_c and control data, both constant across groups.
+        # var_c depends only on model_mu_c and control data — constant across groups.
         y_control_pred = self.model_mu_c.predict_proba(X[control_mask])[:, 1]
-        var_c = (y[control_mask] - y_control_pred).var()
+        self.var_c = (y[control_mask] - y_control_pred).var()
+        self.vars_c = {group: self.var_c for group in self.t_groups}
 
         for group in self.t_groups:
             treatment_mask = treatment == group
@@ -560,11 +571,9 @@ def fit(self, X, treatment, y, p=None):
 
             self.models_mu_t[group].fit(X_treat, y_treat)
 
-            self.vars_c[group] = var_c
-            var_t = (
+            self.vars_t[group] = (
                 y_treat - self.models_mu_t[group].predict_proba(X_treat)[:, 1]
             ).var()
-            self.vars_t[group] = var_t
 
             # Train treatment effect models using cross-group imputation
             d_c = (