From 2c862c09ba5057f0b68e65fa2c6d297be7220adc Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 17 Feb 2023 10:33:06 -0600 Subject: [PATCH 01/23] Add posterior_predict and better support in marginalEffects.plot --- foundry/evaluation/marginal_effects.py | 35 +++++-- foundry/glm/glm.py | 121 +++++++++++++++++++++---- 2 files changed, 133 insertions(+), 23 deletions(-) diff --git a/foundry/evaluation/marginal_effects.py b/foundry/evaluation/marginal_effects.py index 68699ea..d4e6e8d 100644 --- a/foundry/evaluation/marginal_effects.py +++ b/foundry/evaluation/marginal_effects.py @@ -268,7 +268,7 @@ def __call__(self, def to_dataframe(self) -> pd.DataFrame: if self._dataframe is None: raise RuntimeError("This `MarginalEffects()` needs to be called on a dataset first.") - return self._dataframe + return self._dataframe.copy() def plot(self, x: Optional[str] = None, @@ -287,17 +287,29 @@ def plot(self, facets = list(facets or []) data = self.to_dataframe() + + # available default features: + available_default_features = list(self.config['vary_features']) + list(self.config['groupby_features']) + + group_by_prediction_col = False if len(self.config['pred_colnames']) > 1: + assert 'prediction_col' not in data.columns data = data.melt( id_vars=[c for c in data if c not in self.config['pred_colnames']], - var_name='prediction', + var_name='prediction_col', value_name='predicted', value_vars=self.config['pred_colnames'] ) - facets.append('prediction') - - # available default features: - available_default_features = list(self.config['vary_features']) + list(self.config['groupby_features']) + if (data['prediction_col'].nunique() > 9) and ('class' not in self.config['pred_colnames'][0]): + # a common use-case for multiple prediction columns is predict_posterior, where we have 100s of + # predictions and it doesn't make sense to map them to color/facets + group_by_prediction_col = True + warn( + "Many prediction columns, will map to group but not color/facets; you can override this by setting " + "color='prediction_col' or facets='prediction_col'" + ) + else: + available_default_features.append('prediction_col') # if x is set, that's not an available default feature: if x: @@ -343,6 +355,13 @@ def plot(self, warn(f"Adding {available_default_features} to `facets`") facets.extend(available_default_features) + if group_by_prediction_col: + if aes_kwargs['group'] == '1': + aes_kwargs['group'] = 'prediction_col' + else: + data['_group'] = data[aes_kwargs['group']].astype('str') + data['prediction_col'].astype('str') + aes_kwargs['group'] = '_group' + plot = ( ggplot(data, aes(**aes_kwargs)) + geom_hline(yintercept=0) + @@ -350,9 +369,11 @@ def plot(self, theme(figure_size=(8, 6), subplots_adjust={'wspace': 0.10}) ) if pd.api.types.is_categorical_dtype(data[x.replace('_binned', '')]): + if group_by_prediction_col: + raise RuntimeError("Please explicitly map 'prediction_col' to color or facet.") plot += geom_col() else: - plot += geom_line() + plot += geom_line(alpha=.50 if group_by_prediction_col else 1) if include_actuals is None: include_actuals = 'actual' in self._dataframe.columns diff --git a/foundry/glm/glm.py b/foundry/glm/glm.py index 06822f0..c0f8cee 100644 --- a/foundry/glm/glm.py +++ b/foundry/glm/glm.py @@ -1,3 +1,4 @@ +import copy import os from time import sleep from typing import Union, Sequence, Optional, Callable, Tuple, Dict @@ -695,23 +696,83 @@ def predict(self, else: type = 'mean' + x_dict, kwargs = self._prepare_predict(X=X, kwargs_as_is=kwargs_as_is, **kwargs) + + result = self._predict_from_x_dict(x_dict=x_dict, type_=type, **kwargs) + + return result + + @torch.inference_mode() + def _prepare_predict(self, X: ModelMatrix, kwargs_as_is: Union[bool, dict] = False, **kwargs) -> Tuple[dict, dict]: x_dict, *_ = self._build_model_mats(X=X, y=None) if 'validate_args' in kwargs: x_dict['validate_args'] = kwargs.pop('validate_args') + + if not isinstance(kwargs_as_is, dict): + kwargs_as_is = {k: kwargs_as_is for k in kwargs} + for k in list(kwargs): + if not kwargs_as_is.get(k, False) and is_array(kwargs[k]): + kwargs[k] = to_2d(to_tensor(kwargs[k], **get_to_kwargs(self.module_))) + + return x_dict, kwargs + + @torch.inference_mode() + def _predict_from_x_dict(self, x_dict: dict, type_: str, **kwargs) -> np.ndarray: dist_kwargs = self._get_family_kwargs(**x_dict) dist = self.family(**dist_kwargs) - result = getattr(dist, type) + result = getattr(dist, type_) if callable(result): - if not isinstance(kwargs_as_is, dict): - kwargs_as_is = {k: kwargs_as_is for k in kwargs} - for k in list(kwargs): - if not kwargs_as_is.get(k, False) and is_array(kwargs[k]): - kwargs[k] = to_2d(to_tensor(kwargs[k], **get_to_kwargs(self.module_))) result = result(**kwargs) elif kwargs: - warn(f"Ignoring {set(kwargs)}, `{dist.__class__.__name__}.{type}` not callable.") + warn(f"Ignoring {set(kwargs)}, `{type(dist).__name__}.{type_}` not callable.") return result.numpy() + @torch.no_grad() + def posterior_predict(self, + X: ModelMatrix, + type: str = 'mean', + kwargs_as_is: Union[bool, dict] = False, + n_iters: int = 500, + collate_fn: Optional[callable] = None, + **kwargs) -> np.ndarray: + """ + Generate predictions from the model's (MAP) posterior. + + :param X: An array or dictionary of arrays. + :param type: The type of the prediction -- i.e. the attribute to be extracted from the resulting + ``torch.Distribution``. Default is to predict the mean of the distribution (note this differs from the default + in ``predict`` since that method follows sklearn behavior for classifiers and predicts the argmax). + :param kwargs_as_is: See documentation in ``predict()`` + :param n_iters: How many samples to draw, default 500. + :param collate_fn: A function that will take the sampled output as a list of length ``n_iters`` and + stack/concatenate into an array. Default is to create an array that is ``orig_shape + (n_iters,)``, unless + the output from ``predict`` has shape ``(n_records, 1)``, i which case will create an array with + ``(n_records, n_iters)``. + :param kwargs: Keyword arguments to pass if ``type`` is callable. See also ``kwargs_as_is``. + :return: A ndarray of predictions, see ``collate_fn``. + """ + if self._coef_mvnorm_ is None: + raise RuntimeError("Must call ``estimate_laplace_coefs()`` first.") + x_dict, kwargs = self._prepare_predict(X=X, kwargs_as_is=kwargs_as_is, **kwargs) + + samples = self._coef_mvnorm_.sample((n_iters,)) + orig = copy.deepcopy(self.module_) + try: + out = [] + for sample in tqdm(samples.T, delay=10): + self._set_all_params(sample) + out.append(self._predict_from_x_dict(x_dict, type_=type, **kwargs)) + finally: + self.module_ = orig + + if collate_fn is None: + if len(out[0].shape) > 2 or out[0].shape[-1] > 1: + collate_fn = lambda x: np.stack(x, -1) + else: + collate_fn = np.concatenate + + return collate_fn(out) + def _get_penalty(self) -> torch.Tensor: """ Get penalty on sum(log_prob) based on the module-weights and self.penalty. @@ -774,26 +835,54 @@ def estimate_laplace_coefs(self, X: ModelMatrix, y: ModelMatrix, sample_weight: fake_cov = torch.diag(torch.diag(hess).pow(-1).clip(min=1E-5)) self._coef_mvnorm_ = torch.distributions.MultivariateNormal(means, covariance_matrix=fake_cov) - def _estimate_laplace_coefs(self, - X: ModelMatrix, - y: ModelMatrix, - sample_weight: Optional[np.ndarray] = None - ) -> Tuple[Sequence[str], torch.Tensor, torch.Tensor]: + def _get_all_params_and_names(self) -> Tuple[Sequence[np.ndarray], Sequence[torch.Tensor]]: + """ + Get all params as a list, and their names as a matching list of ndarrays w/same-shape + """ all_param_names = [] all_params = [] for dp in self.family.params: for nm, param_values in self.module_[dp].named_parameters(): param_names = self._module_param_names_[dp][nm] assert param_names.shape == param_values.shape, f"param_names.shape!=param_values.shape for {dp}.{nm}" - all_param_names.extend(f"{dp}__{pnm}" for pnm in param_names.reshape(-1)) - all_params.append(param_values) # TODO: any way to assert reshape(-1) matches internals of hessian? - means = torch.cat([p.reshape(-1) for p in all_params]) + all_param_names.append(param_names) + all_params.append(param_values) + + return all_param_names, all_params + + @torch.no_grad() + def _set_all_params(self, unrolled: torch.Tensor): + """ + ``_estimate_laplace_coefs()`` provides names, means, hess that are 'unrolled' from the nested structure in + ``self.module_``. This function takes a 1d tensor with that same unrolled structure, and puts it back into + ``self.module_``. + """ + start_ = 0 + for dp in self.family.params: + for nm, param_values in self.module_[dp].named_parameters(): + end_ = start_ + param_values.numel() + param_values[:] = unrolled[start_:end_] + start_ = end_ + + def _estimate_laplace_coefs(self, + X: ModelMatrix, + y: ModelMatrix, + sample_weight: Optional[np.ndarray] = None + ) -> Tuple[Sequence[str], torch.Tensor, torch.Tensor]: + # get all params as a list, and their names as a matching list: + all_param_names, all_params = self._get_all_params_and_names() + # get the log-prob -> hessian: x_dict, lp_dict = self._build_model_mats(X, y, sample_weight, include_y=True) log_prob = self.get_log_prob(x_dict=x_dict, lp_dict=lp_dict, mean=False) hess = hessian(output=-log_prob.squeeze(), inputs=all_params, allow_unused=True, progress=False) - return all_param_names, means, hess + # flatten out the params and names so that names, means, and hess all match: + # TODO: any way to assert reshape(-1) matches internals of `hessian()`? + means = torch.cat([p.reshape(-1) for p in all_params]) + all_param_names_flat = list(np.concatenate([pn.reshape(-1) for pn in all_param_names])) + + return all_param_names_flat, means, hess def family_from_string(family: str, y: Optional[dict] = None) -> Family: From b20fbcadaa9e3397a59f37bd66559ad30c230fc6 Mon Sep 17 00:00:00 2001 From: Jacob Date: Wed, 15 Nov 2023 21:55:52 -0600 Subject: [PATCH 02/23] fix issues in posterior_predict and ME integration --- docs/README.py | 12 +++-- foundry/evaluation/marginal_effects.py | 65 +++++++++++++++----------- foundry/glm/glm.py | 48 +++++++++++++++---- 3 files changed, 86 insertions(+), 39 deletions(-) diff --git a/docs/README.py b/docs/README.py index 77f8b7c..1d21875 100644 --- a/docs/README.py +++ b/docs/README.py @@ -6,7 +6,7 @@ # extension: .py # format_name: percent # format_version: '1.3' -# jupytext_version: 1.14.4 +# jupytext_version: 1.15.2 # kernelspec: # display_name: Python 3 (ipykernel) # language: python @@ -110,15 +110,17 @@ # By default, the `Glm` will estimate not just the parameters of our model, but also the uncertainty associated with them. We can access a dataframe of these with the `coef_dataframe_` attribute: # %% -df_coefs = glm[-1].coef_dataframe_ + +# %% +df_coefs = glm[-1].coef_dataframe_.copy() +df_coefs = df_coefs[~df_coefs['name'].str.startswith('probs_bias')].reset_index(drop=True) +df_coefs[['param', 'trans', 'term']] = df_coefs['name'].str.split('__', n=3, expand=True) df_coefs # %% [markdown] # Using this, it's easy to plot our model-coefficients: # %% -df_coefs[['param', 'trans', 'term']] = df_coefs['name'].str.split('__', n=3, expand=True) - df_coefs[df_coefs['name'].str.contains('page_feat')].plot('term', 'estimate', kind='bar', yerr='se') df_coefs[df_coefs['name'].str.contains('user_agent_platform')].plot('term', 'estimate', kind='bar', yerr='se') @@ -129,7 +131,7 @@ # %% glm_me = MarginalEffects(glm) -glm_me.fit( +glm_me( X=df_val_expanded, y=df_val_expanded['is_click'], vary_features=['page_feat3'] diff --git a/foundry/evaluation/marginal_effects.py b/foundry/evaluation/marginal_effects.py index 147a54f..b9f36c7 100644 --- a/foundry/evaluation/marginal_effects.py +++ b/foundry/evaluation/marginal_effects.py @@ -73,6 +73,11 @@ def raw(col: str) -> Callable: class MarginalEffects: def __init__(self, pipeline: Pipeline, predict_method: Optional[str] = None, quiet: bool = False): self.pipeline = pipeline + if predict_method is not None: + warn( + "Passing `predict_method` to `MarginalEffects` is deprecated, pass to `__call__` instead.", + DeprecationWarning + ) self.predict_method = predict_method self._dataframe = None self.config = None @@ -114,6 +119,7 @@ def __call__(self, vary_features_aggfun: Union[str, dict, Callable] = 'mean', marginalize_aggfun: Union[str, dict, Callable, None] = 'downsample100000', y_aggfun: Union[str, Callable] = 'mean', + predict_method: Optional[str] = None, **predict_kwargs) -> 'MarginalEffects': """ Prepare a dataframe/plot showing how predictions vary when one or more features are varied, holding @@ -217,12 +223,13 @@ def __call__(self, continue df_vary_grid = df_vary_grid.merge(df_mapping, on=binned_fname) + predict_method = predict_method or self.predict_method + self.config = {'pred_colnames': []} if marginalize_aggfun: - df_me = df_vary_grid.merge(df_no_vary, how='left', on=groupby_colnames) - for col, preds in self.get_predictions(X=df_me, **predict_kwargs).items(): - df_me[col] = preds - self.config['pred_colnames'].append(col) + df_me_base = df_vary_grid.merge(df_no_vary, how='left', on=groupby_colnames) + df_me = self.add_predictions(df_me_base, predict_method=predict_method, **predict_kwargs) + self.config['pred_colnames'].extend(set(df_me.columns) - set(df_me_base.columns)) else: pred_colnames = set() if df_no_vary.shape[0] > 100_000 and not self.quiet: @@ -234,11 +241,9 @@ def __call__(self, df_me = [] for _df_vary_chunk in chunks: - - _df_merged = _df_vary_chunk.merge(df_no_vary, how='left', on=groupby_colnames) - for col, preds in self.get_predictions(X=_df_merged, **predict_kwargs).items(): - _df_merged[col] = preds - pred_colnames.add(col) + _df_merged_base = _df_vary_chunk.merge(df_no_vary, how='left', on=groupby_colnames) + _df_merged = self.add_predictions(_df_merged_base, predict_method=predict_method, **predict_kwargs) + pred_colnames.update(set(_df_merged.columns) - set(_df_merged_base.columns)) _df_collapsed = (_df_merged .groupby(groupby_colnames + list(vary_features), observed=False) @@ -273,7 +278,8 @@ def plot(self, x: Optional[str] = None, color: Optional[str] = None, facets: Optional[Sequence[str]] = None, - include_actuals: Optional[bool] = None) -> 'ggplot': + include_actuals: Optional[bool] = None, + line_alpha: Optional[float] = None) -> 'ggplot': try: from plotnine import ( @@ -372,12 +378,17 @@ def plot(self, raise RuntimeError("Please explicitly map 'prediction_col' to color or facet.") plot += geom_col() else: - plot += geom_line(alpha=.50 if group_by_prediction_col else 1) + if line_alpha is None: + line_alpha = max( + 1 / (data[aes_kwargs['group']].nunique() if aes_kwargs['group'] in data.columns else 1), + .01 + ) + plot += geom_line(alpha=line_alpha) if include_actuals is None: - include_actuals = 'actual' in self._dataframe.columns + include_actuals = 'actual' in data.columns if include_actuals: - if 'actual' in self._dataframe.columns: + if 'actual' in data.columns: plot += geom_point(aes(y='actual', size='n')) plot += ylab("predicted (line) & actual (dots)") else: @@ -402,7 +413,8 @@ def _get_maybe_binned_features(X: pd.DataFrame, binned_fname = f'{fname}_binned' yield fname, binned_fname, binned_feature - def get_predictions(self, X: pd.DataFrame, **kwargs) -> Dict[str, np.ndarray]: + def add_predictions(self, X: pd.DataFrame, predict_method: Optional[str], **kwargs) -> pd.DataFrame: + Xorig = X # avoid warning from ColumnTransformer: X = X.reindex(columns=self.all_column_names_in) @@ -410,33 +422,34 @@ def get_predictions(self, X: pd.DataFrame, **kwargs) -> Dict[str, np.ndarray]: prep_steps_, estimator_ = self.pipeline[0:-1], self.pipeline[-1] # default behavior: use proba if it's available - if self.predict_method is None: + if predict_method is None: if hasattr(estimator_, 'predict_proba'): - self.predict_method = 'predict_proba' + predict_method = 'predict_proba' else: - self.predict_method = 'predict' - predictions = getattr(estimator_, self.predict_method)(prep_steps_.transform(X), **kwargs) + predict_method = 'predict' + predfun = getattr(estimator_, predict_method) + Xt = prep_steps_.transform(X) + predictions = predfun(Xt, **kwargs) # validate output: if len(predictions.shape) > 1 and predictions.shape[1] > 1: assert len(predictions.shape) == 2 - if self.predict_method == 'predict_proba' and predictions.shape[1] == 2: + if predict_method == 'predict_proba' and predictions.shape[1] == 2: # handle common case of 2-class prediction, only plot p(positive-class) predictions = predictions[:, 1] elif not isinstance(predictions, pd.DataFrame): predictions = pd.DataFrame(predictions) - if self.predict_method == 'predict_proba': + if predict_method == 'predict_proba': predictions.columns = [f'class{i}' for i in range(predictions.shape[1])] else: predictions.columns = [f'pred{i}' for i in range(predictions.shape[1])] - # handle multi-output: - if isinstance(predictions, pd.DataFrame): - predictions = {k: v.values for k, v in predictions.to_dict(orient='series').items()} - else: - predictions = {'predicted': to_1d(predictions)} + # handle single or multi-output: + if not isinstance(predictions, pd.DataFrame): + predictions = pd.DataFrame({'predicted': to_1d(predictions)}) - return predictions + predictions.index = Xorig.index + return Xorig.join(predictions) @staticmethod def _standardize_maybe_binned(data: pd.DataFrame, features: Collection[Union[str, Binned]]) -> Dict[str, Binned]: diff --git a/foundry/glm/glm.py b/foundry/glm/glm.py index d3558e4..234f239 100644 --- a/foundry/glm/glm.py +++ b/foundry/glm/glm.py @@ -118,6 +118,17 @@ family_names['gaussian'] = family_names['normal'] family_names['mvnorm'] = family_names['multivariate_normal'] +_posterior_predictive_cache = {} + + +class _glm_pp: + + def __enter__(self): + _posterior_predictive_cache['_enabled'] = True + + def __exit__(self, exc_type, exc_val, exc_tb): + _posterior_predictive_cache.clear() + class Glm(BaseEstimator): """ @@ -144,6 +155,8 @@ class Glm(BaseEstimator): """ family_names = family_names + fixed_posterior = _glm_pp() + def __init__(self, family: Union[str, Family], penalty: Union[float, Sequence[float], Dict[str, float]] = 0., @@ -730,6 +743,18 @@ def _predict_from_x_dict(self, x_dict: dict, type_: str, **kwargs) -> np.ndarray warn(f"Ignoring {set(kwargs)}, `{type(dist).__name__}.{type_}` not callable.") return result.numpy() + def _sample_coef_mvnorm(self, n_iters: int, use_cache: bool = None) -> torch.Tensor: + if use_cache is None: + use_cache = _posterior_predictive_cache.get('_enabled', False) + + if not use_cache: + return self._coef_mvnorm_.sample((n_iters,)) + + key = (id(self._coef_mvnorm_), n_iters) + if key not in _posterior_predictive_cache: + _posterior_predictive_cache[key] = self._sample_coef_mvnorm(n_iters=n_iters, use_cache=False) + return _posterior_predictive_cache[key] + @torch.no_grad() def posterior_predict(self, X: ModelMatrix, @@ -749,7 +774,7 @@ def posterior_predict(self, :param n_iters: How many samples to draw, default 500. :param collate_fn: A function that will take the sampled output as a list of length ``n_iters`` and stack/concatenate into an array. Default is to create an array that is ``orig_shape + (n_iters,)``, unless - the output from ``predict`` has shape ``(n_records, 1)``, i which case will create an array with + the output from ``predict`` has shape ``(n_records, 1)``, in which case will create an array with ``(n_records, n_iters)``. :param kwargs: Keyword arguments to pass if ``type`` is callable. See also ``kwargs_as_is``. :return: A ndarray of predictions, see ``collate_fn``. @@ -758,11 +783,11 @@ def posterior_predict(self, raise RuntimeError("Must call ``estimate_laplace_coefs()`` first.") x_dict, kwargs = self._prepare_predict(X=X, kwargs_as_is=kwargs_as_is, **kwargs) - samples = self._coef_mvnorm_.sample((n_iters,)) + samples = self._sample_coef_mvnorm(n_iters=n_iters) orig = copy.deepcopy(self.module_) try: out = [] - for sample in tqdm(samples.T, delay=10): + for sample in tqdm(samples, delay=10): self._set_all_params(sample) out.append(self._predict_from_x_dict(x_dict, type_=type, **kwargs)) finally: @@ -772,7 +797,7 @@ def posterior_predict(self, if len(out[0].shape) > 2 or out[0].shape[-1] > 1: collate_fn = lambda x: np.stack(x, -1) else: - collate_fn = np.concatenate + collate_fn = lambda x: np.concatenate(x, 1) return collate_fn(out) @@ -838,20 +863,22 @@ def estimate_laplace_coefs(self, X: ModelMatrix, y: ModelMatrix, sample_weight: fake_cov = torch.diag(torch.diag(hess).pow(-1).clip(min=1E-5)) self._coef_mvnorm_ = torch.distributions.MultivariateNormal(means, covariance_matrix=fake_cov) - def _get_all_params_and_names(self) -> Tuple[Sequence[np.ndarray], Sequence[torch.Tensor]]: + def _get_all_params_and_names(self) -> Tuple[Sequence[str], Sequence[np.ndarray], Sequence[torch.Tensor]]: """ Get all params as a list, and their names as a matching list of ndarrays w/same-shape """ + prefixes = [] all_param_names = [] all_params = [] for dp in self.family.params: for nm, param_values in self.module_[dp].named_parameters(): + prefixes.append(f"{dp}_{nm}") param_names = self._module_param_names_[dp][nm] assert param_names.shape == param_values.shape, f"param_names.shape!=param_values.shape for {dp}.{nm}" all_param_names.append(param_names) all_params.append(param_values) - return all_param_names, all_params + return prefixes, all_param_names, all_params @torch.no_grad() def _set_all_params(self, unrolled: torch.Tensor): @@ -866,6 +893,8 @@ def _set_all_params(self, unrolled: torch.Tensor): end_ = start_ + param_values.numel() param_values[:] = unrolled[start_:end_] start_ = end_ + if start_ != len(unrolled): + raise ValueError(f"Expected unrolled to have length {start_:,}, but was {len(unrolled):,}") def _estimate_laplace_coefs(self, X: ModelMatrix, @@ -873,7 +902,7 @@ def _estimate_laplace_coefs(self, sample_weight: Optional[np.ndarray] = None ) -> Tuple[Sequence[str], torch.Tensor, torch.Tensor]: # get all params as a list, and their names as a matching list: - all_param_names, all_params = self._get_all_params_and_names() + prefixes, all_param_names, all_params = self._get_all_params_and_names() # get the log-prob -> hessian: x_dict, lp_dict = self._build_model_mats(X, y, sample_weight, include_y=True) @@ -883,7 +912,10 @@ def _estimate_laplace_coefs(self, # flatten out the params and names so that names, means, and hess all match: # TODO: any way to assert reshape(-1) matches internals of `hessian()`? means = torch.cat([p.reshape(-1) for p in all_params]) - all_param_names_flat = list(np.concatenate([pn.reshape(-1) for pn in all_param_names])) + + all_param_names_flat = [] + for prefix, pns in zip(prefixes, all_param_names): + all_param_names_flat.extend(f"{prefix}__{pnm}" for pnm in pns.reshape(-1)) return all_param_names_flat, means, hess From 6473125851038e5d79186d4bfb382d8d35e1343a Mon Sep 17 00:00:00 2001 From: Jacob Date: Thu, 16 Nov 2023 08:58:29 -0600 Subject: [PATCH 03/23] restore original if issue --- foundry/glm/glm.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/foundry/glm/glm.py b/foundry/glm/glm.py index 234f239..4a28b34 100644 --- a/foundry/glm/glm.py +++ b/foundry/glm/glm.py @@ -887,6 +887,7 @@ def _set_all_params(self, unrolled: torch.Tensor): ``self.module_``. This function takes a 1d tensor with that same unrolled structure, and puts it back into ``self.module_``. """ + orig = copy.deepcopy(self.module_) start_ = 0 for dp in self.family.params: for nm, param_values in self.module_[dp].named_parameters(): @@ -894,7 +895,10 @@ def _set_all_params(self, unrolled: torch.Tensor): param_values[:] = unrolled[start_:end_] start_ = end_ if start_ != len(unrolled): - raise ValueError(f"Expected unrolled to have length {start_:,}, but was {len(unrolled):,}") + self.module_ = orig + raise ValueError( + f"Expected unrolled to have length {start_:,}, but was {len(unrolled):,}. Restored original params." + ) def _estimate_laplace_coefs(self, X: ModelMatrix, From 7fad90a706a20334c66eaea6cbba2963ae858c55 Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 1 Dec 2023 14:41:15 -0600 Subject: [PATCH 04/23] fix bug in how simulate_data's random_state is used for torch --- foundry/data.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/foundry/data.py b/foundry/data.py index f1903e9..acd1fba 100644 --- a/foundry/data.py +++ b/foundry/data.py @@ -139,6 +139,7 @@ def simulate_data(family: str, n_samples: int, n_features: int, predict_params: Optional[Sequence] = None, + random_state: Optional[np.random.RandomState] = None, **kwargs) -> pd.DataFrame: """ :param family: A family-alias, see ``Glm.family_names``. @@ -151,6 +152,10 @@ def simulate_data(family: str, :return: A dataframe with predictors in format ``x_{i}`` and target ``target``. The ground-truth predictor- coefficients are stored in the ``attrs`` attribute of the dataframe. """ + random_state = random_state or np.random.randint(1e6) + if not isinstance(random_state, np.random.RandomState): + random_state = np.random.RandomState(random_state) + if isinstance(family, str): family = family_from_string(family) if predict_params is None: @@ -174,6 +179,7 @@ def simulate_data(family: str, n_targets=n_targets, noise=0, coef=True, + random_state=random_state, **kwargs ) # undo squeezing: @@ -204,10 +210,8 @@ def simulate_data(family: str, distribution = family(**family_kwargs) # set torch random-state: - random_state = kwargs.get('random_state') with torch.random.fork_rng(): - if random_state is not None: - torch.manual_seed(random_state if isinstance(random_state, int) else hash(random_state)) + torch.manual_seed(hash(random_state.get_state()[1].tobytes())) # sample from distribution for target: out['target'] = distribution.sample().numpy() From 239d3e08a6ff7a9cc6ba5bea2b6e7d8970f061d6 Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 15 Dec 2023 15:32:32 -0600 Subject: [PATCH 05/23] update pin so notebook isnt broken --- setup.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index b03df1f..ea481d6 100644 --- a/setup.py +++ b/setup.py @@ -20,7 +20,9 @@ extras_require={ 'dev': [ 'jupytext', - 'notebook', + # pinning these is needed because notebook>6 isn't working with jupytext + 'notebook==6.5.4', + 'traitlets==5.9' 'pytest', 'requests', 'plotnine', From 8abde8a1d4e18bae93a49585de8f8b1790c8d4de Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 15 Dec 2023 15:35:56 -0600 Subject: [PATCH 06/23] missing comma --- setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index ea481d6..9940700 100644 --- a/setup.py +++ b/setup.py @@ -22,12 +22,12 @@ 'jupytext', # pinning these is needed because notebook>6 isn't working with jupytext 'notebook==6.5.4', - 'traitlets==5.9' + 'traitlets==5.9', 'pytest', 'requests', 'plotnine', ], - 'docs' : [ + 'docs': [ 'requests' ] } From 33e847ad4c01a56c8f515a5da3aa1ec12a0f88c0 Mon Sep 17 00:00:00 2001 From: Jacob Date: Tue, 25 Feb 2025 13:13:36 -0600 Subject: [PATCH 07/23] add todo --- foundry/glm/glm.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/foundry/glm/glm.py b/foundry/glm/glm.py index 4a28b34..b284c5e 100644 --- a/foundry/glm/glm.py +++ b/foundry/glm/glm.py @@ -360,6 +360,7 @@ def _fit(self, # 'closure' for torch.optimizer.Optimizer.step method: def closure(): self.optimizer_.zero_grad() + # TODO: if distribution fails validation due to nans/infs, will not currently trigger FitFailedException loss = -self.get_log_prob(x_dict, lp_dict) if is_invalid(loss): self._fit_failed += 1 @@ -695,8 +696,8 @@ def predict(self, :param X: An array or dictionary of arrays. :param type: The type of the prediction -- i.e. the attribute to be extracted from the resulting ``torch.Distribution``. The default depends on the family. If ``self.family.supports_predict_proba``, and the - distribution doesn't have a ``total_count`` parametere (e.g. multinomial), then this method will predict the - class. Otherwise, this distribution will predict the mean of the distribution. + distribution doesn't have a ``total_count`` parameter (e.g. multinomial), then this method will predict the + class. Otherwise, this method will predict the mean of the distribution. :param kwargs_as_is: If the ``type`` is callable, then kwargs that are arrays will be converted to tensors, and if 1D then will be unsqueezed to 2D (unsqueezing is to avoid accidental broadcasting, wherein (e.g.) a distribution with batch_shape of N*1 receives a ``value`` w/shape (N,1), resulting in a (N,N) tensor, @@ -794,7 +795,7 @@ def posterior_predict(self, self.module_ = orig if collate_fn is None: - if len(out[0].shape) > 2 or out[0].shape[-1] > 1: + if len(out[0].shape) != 2 or out[0].shape[-1] > 1: collate_fn = lambda x: np.stack(x, -1) else: collate_fn = lambda x: np.concatenate(x, 1) From bc54771e132061d33d66ff16a77b30842ae698b8 Mon Sep 17 00:00:00 2001 From: Jacob Date: Thu, 24 Apr 2025 15:41:15 -0500 Subject: [PATCH 08/23] update as_transformer --- foundry/preprocessing/__init__.py | 1 - foundry/preprocessing/sklearn/__init__.py | 2 +- foundry/preprocessing/sklearn/utils.py | 29 +++++++++-------------- 3 files changed, 12 insertions(+), 20 deletions(-) diff --git a/foundry/preprocessing/__init__.py b/foundry/preprocessing/__init__.py index 40ee4eb..3eda0a9 100644 --- a/foundry/preprocessing/__init__.py +++ b/foundry/preprocessing/__init__.py @@ -5,7 +5,6 @@ FunctionTransformer, SimpleImputer, as_transformer, - identity, ) from .categorical import ToCategorical from .dates import FourierFeatures diff --git a/foundry/preprocessing/sklearn/__init__.py b/foundry/preprocessing/sklearn/__init__.py index fbd318e..d2ce3c2 100644 --- a/foundry/preprocessing/sklearn/__init__.py +++ b/foundry/preprocessing/sklearn/__init__.py @@ -2,4 +2,4 @@ from .column_dropper import ColumnDropper from .dataframe_transformer import DataFrameTransformer from .interactions import InteractionFeatures -from .utils import as_transformer, identity +from .utils import as_transformer diff --git a/foundry/preprocessing/sklearn/utils.py b/foundry/preprocessing/sklearn/utils.py index 2602224..ccc1bab 100644 --- a/foundry/preprocessing/sklearn/utils.py +++ b/foundry/preprocessing/sklearn/utils.py @@ -1,36 +1,29 @@ -from warnings import warn +from typing import Union, Sequence from sklearn.base import TransformerMixin -from sklearn.pipeline import make_pipeline +from sklearn.pipeline import make_pipeline, Pipeline from .backports import FunctionTransformer -def identity(x): - return x - - -def as_transformer(x) -> TransformerMixin: +def as_transformer( + x: Union[TransformerMixin, Sequence[TransformerMixin], callable, None], + **kwargs +) -> Union[TransformerMixin, Pipeline]: """ Standardize a transformer, function, or list of these into a transformer. Lists are converted to sklearn.pipelines. - - # TODO: consider deprecating, only use-case is `as_transformer(x)` which has no real advantages - over `FunctionTransformer(x)` """ - warn( - "as_transformer may be removed in future versions of foundry. Consider using" - "foundry.preprocessing.FunctionTransformer", - category=DeprecationWarning - ) + if callable(x): + return FunctionTransformer(x, kw_args=kwargs) + if kwargs: + raise ValueError("`kwargs` can only be passed if x is a function.") if x is None: - return as_transformer(identity) + return FunctionTransformer() if hasattr(x, '__iter__') and not isinstance(x, str): return make_pipeline(*[as_transformer(xi) for xi in x]) if hasattr(x, 'transform'): return x - elif callable(x): - return FunctionTransformer(x) else: raise TypeError(f"{type(x).__name__} does not have a `transform()` method.") From 2b6db1d1110890f0ff8ed010bf0c0f439b0ed9d1 Mon Sep 17 00:00:00 2001 From: Jacob Date: Mon, 26 Jan 2026 21:29:16 -0600 Subject: [PATCH 09/23] handle new sklearn indexing behavior --- foundry/util.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/foundry/util.py b/foundry/util.py index bbd9abb..c9337c8 100644 --- a/foundry/util.py +++ b/foundry/util.py @@ -151,6 +151,10 @@ def __getitem__(self, sl: Union[int, str, slice]) -> Union['SliceDict', ArrayTyp ) if isinstance(sl, str): return super(SliceDict, self).__getitem__(sl) + if isinstance(sl, tuple) and len(sl) == 2 and sl[-1] is Ellipsis: + # array[(ind, Ellipsis)] and array[ind] should be equivalent for ndarrays, but the former will break + # pandas types. sklearn _array_indexing previously did the latter but switched to the former + sl = sl[0] return SliceDict(**{k: (v[sl] if hasattr(v, 'shape') else v) for k, v in self.items()}) def __setitem__(self, key: str, value: ArrayType): From 10fb7c410b7dda4735e0c3b919941a091eb3e159 Mon Sep 17 00:00:00 2001 From: Jacob Date: Tue, 27 Jan 2026 09:26:12 -0600 Subject: [PATCH 10/23] make sure sklearn indexes SliceDict with pandas indexing if values are pandas otherwise non-default index will cause unexpected behavior --- foundry/util.py | 34 +++++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/foundry/util.py b/foundry/util.py index c9337c8..100eeb9 100644 --- a/foundry/util.py +++ b/foundry/util.py @@ -135,8 +135,23 @@ def __init__(self, **kwargs): else: self._len = lengths[0] + # sklearn checks if it should use pandas indexing by checking if there's an iloc attribute + if self.is_pandas: + self.__dict__['iloc'] = True + else: + self.__dict__.pop('iloc', None) + super().__init__(**kwargs) + @property + def is_pandas(self) -> bool: + is_pandas = [hasattr(v, 'iloc') for v in self.values() if hasattr(v, 'shape')] + any_pandas = any(is_pandas) + all_pandas = all(is_pandas) + if any_pandas and not all_pandas: + raise ValueError("Currenlty SliceDict does not support a mix of pandas and non-pandas") + return all(is_pandas) + @staticmethod def _standardize_val(val): return np.asarray(val) if is_array(val) and not hasattr(val, 'shape') else val @@ -151,11 +166,8 @@ def __getitem__(self, sl: Union[int, str, slice]) -> Union['SliceDict', ArrayTyp ) if isinstance(sl, str): return super(SliceDict, self).__getitem__(sl) - if isinstance(sl, tuple) and len(sl) == 2 and sl[-1] is Ellipsis: - # array[(ind, Ellipsis)] and array[ind] should be equivalent for ndarrays, but the former will break - # pandas types. sklearn _array_indexing previously did the latter but switched to the former - sl = sl[0] - return SliceDict(**{k: (v[sl] if hasattr(v, 'shape') else v) for k, v in self.items()}) + cls = type(self) + return cls(**{k: (v[sl] if hasattr(v, 'shape') else v) for k, v in self.items()}) def __setitem__(self, key: str, value: ArrayType): value = self._standardize_val(value) @@ -181,6 +193,18 @@ def __setitem__(self, key: str, value: ArrayType): super().__setitem__(key, value) + # sklearn checks if it should use pandas indexing by checking if there's an iloc attribute + if self.is_pandas: + self.__dict__['iloc'] = True + else: + self.__dict__.pop('iloc', None) + + def take(self, indices, axis: int = 0, **kwargs) -> 'SliceDict': + if axis: + raise ValueError("Only axis=0 is supported") + cls = type(self) + return cls(**{k: (v.take(indices, axis, **kwargs) if hasattr(v, 'shape') else v) for k, v in self.items()}) + def update(self, kwargs: dict): for key, value in kwargs.items(): self.__setitem__(key, value) From 134c49e5c6bf71031dfabbca0a18a07e11f29618 Mon Sep 17 00:00:00 2001 From: Jacob Date: Sat, 23 May 2026 08:53:52 -0500 Subject: [PATCH 11/23] avoid pandas deprecation warning for sparse --- foundry/preprocessing/sklearn/interactions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/foundry/preprocessing/sklearn/interactions.py b/foundry/preprocessing/sklearn/interactions.py index 73c0c07..e32fc07 100644 --- a/foundry/preprocessing/sklearn/interactions.py +++ b/foundry/preprocessing/sklearn/interactions.py @@ -142,8 +142,8 @@ def _sparse_safe_multiply(old_vals: pd.Series, new_vals: pd.Series) -> Union[Spa # because sparse-arrays allow for any fill-value, they don't leverage the fact that, if fill_value=0, # sparse*sparse only needs to capture the intersection of the two; instead, they fill the union. - old_is_sparse = pd.api.types.is_sparse(old_vals) - new_is_sparse = pd.api.types.is_sparse(new_vals) + old_is_sparse = isinstance(old_vals.dtype, pd.SparseDtype) + new_is_sparse = isinstance(new_vals.dtype, pd.SparseDtype) if new_is_sparse and old_is_sparse: index_intersection = old_vals.sp_index.intersect(new_vals.sp_index) assert new_vals.fill_value == old_vals.fill_value == 0 From 8b72e6b81ce9a5f91e72c3c12c93e3320f348c11 Mon Sep 17 00:00:00 2001 From: Jacob Date: Thu, 16 Jul 2026 15:11:14 -0500 Subject: [PATCH 12/23] fix merge --- foundry/preprocessing/sklearn/interactions.py | 3 + foundry/uplift/__init__.py | 1 + foundry/uplift/meta_learners/__init__.py | 3 + foundry/uplift/meta_learners/s_learner.py | 114 ++++++++ foundry/uplift/meta_learners/t_learner.py | 80 ++++++ foundry/uplift/meta_learners/x_learner.py | 133 +++++++++ foundry/uplift/util.py | 255 ++++++++++++++++++ foundry/util.py | 16 ++ 8 files changed, 605 insertions(+) create mode 100644 foundry/uplift/__init__.py create mode 100644 foundry/uplift/meta_learners/__init__.py create mode 100644 foundry/uplift/meta_learners/s_learner.py create mode 100644 foundry/uplift/meta_learners/t_learner.py create mode 100644 foundry/uplift/meta_learners/x_learner.py create mode 100644 foundry/uplift/util.py diff --git a/foundry/preprocessing/sklearn/interactions.py b/foundry/preprocessing/sklearn/interactions.py index e32fc07..bdc08c6 100644 --- a/foundry/preprocessing/sklearn/interactions.py +++ b/foundry/preprocessing/sklearn/interactions.py @@ -95,6 +95,9 @@ def fit(self, X: pd.DataFrame, y=None) -> 'InteractionFeatures': return self def transform(self, X: pd.DataFrame, y=None) -> pd.DataFrame: + if not X.index.is_unique: + raise ValueError(f"{type(self).__name__} only works when X has a unique index; try resetting the index.") + orig_cols = set(X.columns) new_cols = {} diff --git a/foundry/uplift/__init__.py b/foundry/uplift/__init__.py new file mode 100644 index 0000000..4d03b2e --- /dev/null +++ b/foundry/uplift/__init__.py @@ -0,0 +1 @@ +from .meta_learners import SLearner, TLearner, XLearner diff --git a/foundry/uplift/meta_learners/__init__.py b/foundry/uplift/meta_learners/__init__.py new file mode 100644 index 0000000..c224da8 --- /dev/null +++ b/foundry/uplift/meta_learners/__init__.py @@ -0,0 +1,3 @@ +from .t_learner import TLearner +from .s_learner import SLearner +from .x_learner import XLearner diff --git a/foundry/uplift/meta_learners/s_learner.py b/foundry/uplift/meta_learners/s_learner.py new file mode 100644 index 0000000..121f914 --- /dev/null +++ b/foundry/uplift/meta_learners/s_learner.py @@ -0,0 +1,114 @@ +import warnings +from typing import Tuple, Union, overload, Literal, Optional + +import numpy as np +import pandas as pd +from sklearn.base import BaseEstimator, clone + +from foundry.util import SliceDict, to_1d, safe_predict +from ..util import get_qini_curve, get_cumulative_gain_score + + +class SLearner(BaseEstimator): + """ + An S-learner. Please note that current implementation assumes randomized treatment/control! + + :param estimator: Any instance that supports the sklearn API (fit/predict and can call ``clone()`` on it). + :param include_interaction: Whether to include X * treatment interaction terms. + """ + estimator_: Optional[BaseEstimator] = None + + def __init__(self, estimator: BaseEstimator, include_interaction: bool = True) -> None: + self.estimator = estimator + self.include_interaction = include_interaction + + def fit(self, X: Union[pd.DataFrame, np.ndarray], y: SliceDict, **fit_kwargs) -> "SLearner": + y_arr, treatment_ind = self._normalize_y(y) + + X_aug = self._augment_with_treatment(X, treatment_ind) + + self.estimator_ = clone(self.estimator).fit(X=X_aug, y=y_arr, **fit_kwargs) + return self + + + @overload + def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[False] = ..., **predict_kwargs) -> np.ndarray: ... + @overload + def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[True], **predict_kwargs) -> Tuple[np.ndarray, np.ndarray]: ... + + def predict(self, X, return_components=False, **predict_kwargs): + X_t = self._augment_with_treatment(X, np.ones(len(X), dtype=bool)) + X_c = self._augment_with_treatment(X, np.zeros(len(X), dtype=bool)) + + yhat_t = safe_predict(self.estimator_, X=X_t, **predict_kwargs) + yhat_c = safe_predict(self.estimator_, X=X_c, **predict_kwargs) + + if return_components: + return yhat_t, yhat_c + + return yhat_t - yhat_c + + def score( + self, + X: Union[pd.DataFrame, np.ndarray], + y: SliceDict, + sample_weight: Union[np.ndarray, None] = None, + method: str = 'qini', + normalize: bool = True, + **kwargs, + ) -> float: + y_arr, treatment_ind = self._normalize_y(y) + if sample_weight is not None: + raise NotImplementedError + pred = self.predict(X=X) + if method == 'cumulative_gain': + return get_cumulative_gain_score( + y_true=y_arr, + treatment=treatment_ind, + score=pred, + **kwargs, + ) + qini = get_qini_curve( + y_true=y_arr, + treatment=treatment_ind, + score=pred, + normalize=normalize, + **kwargs, + ) + random_area = np.linspace(0, qini[-1], qini.shape[0]).sum() + return (np.nansum(qini) - random_area) / qini.shape[0] + + def _augment_with_treatment( + self, + X: Union[pd.DataFrame, np.ndarray], + treatment_ind: np.ndarray, + ) -> Union[pd.DataFrame, np.ndarray]: + treatment_col = treatment_ind.astype(int) + + if isinstance(X, pd.DataFrame): + X_aug = X.copy() + X_aug["treatment"] = treatment_col + + if self.include_interaction: + for col in X.columns: + X_aug[f"{col}_x_treatment"] = X[col] * treatment_col + + return X_aug + + else: + treatment_col = treatment_col.reshape(-1, 1) + + if self.include_interaction: + interaction_terms = X * treatment_col + return np.hstack([X, treatment_col, interaction_terms]) + + return np.hstack([X, treatment_col]) + + @staticmethod + def _normalize_y(y: SliceDict) -> Tuple[np.ndarray, np.ndarray]: + y = y.copy() + y_arr = to_1d(np.asanyarray(y.pop("value"))) + treatment_ind = to_1d(np.asanyarray(y.pop("is_treatment")).astype(bool)) + if len(y.keys()): + warnings.warn(f"Unused keys in ``y``: {set(y)}") + return y_arr, treatment_ind diff --git a/foundry/uplift/meta_learners/t_learner.py b/foundry/uplift/meta_learners/t_learner.py new file mode 100644 index 0000000..4e50f46 --- /dev/null +++ b/foundry/uplift/meta_learners/t_learner.py @@ -0,0 +1,80 @@ +import warnings +from typing import Optional, Tuple, Union, overload, Literal + +import numpy as np +import pandas as pd +from sklearn.base import BaseEstimator, clone + +from foundry.util import SliceDict, to_1d, safe_predict +from ..util import get_qini_curve, get_cumulative_gain_score + + +class TLearner(BaseEstimator): + """ + A T-learner. Please note that current implementation assumes randomized treatment/control! + + :param estimator: Any instance that supports the sklearn API (fit/predict and can call ``clone()`` on it). + """ + treatment_est_: Optional[BaseEstimator] = None + control_est_: Optional[BaseEstimator] = None + + def __init__(self, estimator: BaseEstimator) -> None: + self.estimator = estimator + + def fit(self, X: Union[pd.DataFrame, np.ndarray], y: SliceDict, **fit_kwargs) -> "TLearner": + y_arr, treatment_ind = self._normalize_y(y) + self.treatment_est_ = clone(self.estimator).fit(X=X[treatment_ind], y=y_arr[treatment_ind], **fit_kwargs) + self.control_est_ = clone(self.estimator).fit(X[~treatment_ind], y_arr[~treatment_ind], **fit_kwargs) + + return self + + @overload + def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[False] = ..., **predict_kwargs) -> np.ndarray: ... + @overload + def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[True], **predict_kwargs) -> Tuple[np.ndarray, np.ndarray]: ... + + def predict(self, X, return_components=False, **predict_kwargs): + yhat_t = safe_predict(self.treatment_est_, X=X, **predict_kwargs) + yhat_c = safe_predict(self.control_est_, X=X, **predict_kwargs) + if return_components: + return yhat_t, yhat_c + return yhat_t - yhat_c + + def score( + self, + X: Union[pd.DataFrame, np.ndarray], + y: SliceDict, + sample_weight: Optional[np.ndarray] = None, + method: str = 'qini', + normalize: bool = True, + **kwargs, + ) -> float: + y_arr, treatment_ind = self._normalize_y(y) + if sample_weight is not None: + raise NotImplementedError + pred = self.predict(X=X) + if method == 'cumulative_gain': + return get_cumulative_gain_score( + y_true=y_arr, + treatment=treatment_ind, + score=pred, + **kwargs, + ) + qini = get_qini_curve( + y_true=y_arr, + treatment=treatment_ind, + score=pred, + normalize=normalize, + **kwargs + ) + random_area = np.linspace(0, qini[-1], qini.shape[0]).sum() + return (np.nansum(qini) - random_area) / qini.shape[0] + + @staticmethod + def _normalize_y(y: SliceDict) -> Tuple[np.ndarray, np.ndarray]: + y = y.copy() + y_arr = to_1d(np.asanyarray(y.pop('value'))) + treatment_ind = to_1d(np.asanyarray(y.pop('is_treatment')).astype(bool)) + if len(y.keys()): + warnings.warn(f"Unused keys in ``y``: {set(y)}") + return y_arr, treatment_ind diff --git a/foundry/uplift/meta_learners/x_learner.py b/foundry/uplift/meta_learners/x_learner.py new file mode 100644 index 0000000..89e19b5 --- /dev/null +++ b/foundry/uplift/meta_learners/x_learner.py @@ -0,0 +1,133 @@ +import warnings +from typing import Any, Dict, Optional, Tuple, Union, overload, Literal + +import numpy as np +import pandas as pd +from sklearn.base import BaseEstimator, clone + +from foundry.util import SliceDict, to_1d, safe_predict +from ..util import get_qini_curve, get_cumulative_gain_score + + +class XLearner(BaseEstimator): + """ + An X-learner. Please note that current implementation assumes randomized treatment/control! + Adapted from https://matheusfacure.github.io/python-causality-handbook/21-Meta-Learners.html + and the original paper by Kunzel et al. (2019): https://arxiv.org/abs/1706.03461 + + :param first_stage_estimator: Any instance that supports the sklearn API (fit/predict and can call ``clone()`` on it). + :param second_stage_estimator: Any instance that supports the sklearn API (fit/predict and can call ``clone()`` on it). + :param propensity_estimator: Any instance that supports the sklearn API (fit/predict and can call ``clone()`` on it). + :param first_stage_fit_params: Optional dict of kwargs passed to ``fit()`` for the first-stage models. + :param second_stage_fit_params: Optional dict of kwargs passed to ``fit()`` for the second-stage models. + :param propensity_fit_params: Optional dict of kwargs passed to ``fit()`` for the propensity model. + """ + first_treatment_est_: Optional[BaseEstimator] = None + first_control_est_: Optional[BaseEstimator] = None + second_treatment_est_: Optional[BaseEstimator] = None + second_control_est_: Optional[BaseEstimator] = None + propensity_est_: Optional[BaseEstimator] = None + + def __init__( + self, + first_stage_estimator: BaseEstimator, + second_stage_estimator: BaseEstimator, + propensity_estimator: BaseEstimator, + first_stage_fit_params: Optional[Dict[str, Any]] = None, + second_stage_fit_params: Optional[Dict[str, Any]] = None, + propensity_fit_params: Optional[Dict[str, Any]] = None, + ) -> None: + self.first_stage_estimator = first_stage_estimator + self.second_stage_estimator = second_stage_estimator + self.propensity_estimator = propensity_estimator + self.first_stage_fit_params = first_stage_fit_params + self.second_stage_fit_params = second_stage_fit_params + self.propensity_fit_params = propensity_fit_params + + def fit(self, X: Union[pd.DataFrame, np.ndarray], y: SliceDict) -> "XLearner": + y_arr, treatment_ind = self._normalize_y(y) + + _first_stage_fit_params = self.first_stage_fit_params or {} + _second_stage_fit_params = self.second_stage_fit_params or {} + _propensity_fit_params = self.propensity_fit_params or {} + + self.first_control_est_ = clone(self.first_stage_estimator).fit( + X[~treatment_ind], y_arr[~treatment_ind], **_first_stage_fit_params + ) + self.first_treatment_est_ = clone(self.first_stage_estimator).fit( + X[treatment_ind], y_arr[treatment_ind], **_first_stage_fit_params + ) + + self.propensity_est_ = clone(self.propensity_estimator).fit( + X, treatment_ind, **_propensity_fit_params + ) + + imputed_te = np.where( + treatment_ind, + y_arr - safe_predict(self.first_control_est_, X), + safe_predict(self.first_treatment_est_, X) - y_arr, + ) + + self.second_control_est_ = clone(self.second_stage_estimator).fit( + X[~treatment_ind], imputed_te[~treatment_ind], **_second_stage_fit_params + ) + self.second_treatment_est_ = clone(self.second_stage_estimator).fit( + X[treatment_ind], imputed_te[treatment_ind], **_second_stage_fit_params + ) + + return self + + @overload + def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[False] = ..., **predict_kwargs) -> np.ndarray: ... + @overload + def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[True], **predict_kwargs) -> Tuple[np.ndarray, np.ndarray]: ... + + def predict(self, X, return_components=False, **predict_kwargs): + p_treatment = safe_predict(self.propensity_est_, X) + p_control = 1 - p_treatment + + tau0 = safe_predict(self.second_control_est_, X, **predict_kwargs) + tau1 = safe_predict(self.second_treatment_est_, X, **predict_kwargs) + + if return_components: + return tau0, tau1 + return p_treatment * tau0 + p_control * tau1 + + def score( + self, + X: Union[pd.DataFrame, np.ndarray], + y: SliceDict, + sample_weight: Optional[np.ndarray] = None, + method: str = 'qini', + normalize: bool = True, + **kwargs, + ) -> float: + y_arr, treatment_ind = self._normalize_y(y) + if sample_weight is not None: + raise NotImplementedError + pred = self.predict(X=X) + if method == 'cumulative_gain': + return get_cumulative_gain_score( + y_true=y_arr, + treatment=treatment_ind, + score=pred, + **kwargs, + ) + qini = get_qini_curve( + y_true=y_arr, + treatment=treatment_ind, + score=pred, + normalize=normalize, + **kwargs, + ) + random_area = np.linspace(0, qini[-1], qini.shape[0]).sum() + return (np.nansum(qini) - random_area) / qini.shape[0] + + @staticmethod + def _normalize_y(y: SliceDict) -> Tuple[np.ndarray, np.ndarray]: + y = y.copy() + y_arr = to_1d(np.asanyarray(y.pop('value'))) + treatment_ind = to_1d(np.asanyarray(y.pop('is_treatment')).astype(bool)) + if len(y.keys()): + warnings.warn(f"Unused keys in ``y``: {set(y)}") + return y_arr, treatment_ind diff --git a/foundry/uplift/util.py b/foundry/uplift/util.py new file mode 100644 index 0000000..6a6b980 --- /dev/null +++ b/foundry/uplift/util.py @@ -0,0 +1,255 @@ +import numpy as np +import pandas as pd +from sklearn.pipeline import Pipeline + + +def get_qini_curve(y_true: np.ndarray, + treatment: np.ndarray, + score: np.ndarray, + min_n_per: int = 1, + normalize: bool = True) -> np.ndarray: + """ + Adapted from https://www.uplift-modeling.com/en/latest/_modules/sklift/metrics/metrics.html#qini_curve + + :param y_true: The true values + :param treatment: A treatment indicator (boolean). + :param score: The uplift score predicted for each record. + :param min_n_per: Minimum number of treatment and control records. For example, ``min_n_per=2`` means no qini + calculations until both (1) treatment has at least 2 records **and** (2) control has at least 2 records. + :param normalize: Whether to normalize to 0-1. + :return: np.ndarray with + """ + + y_true = np.asarray(y_true) + score = np.asarray(score) + treatment = np.asarray(treatment, dtype='bool') + + desc_score_indices = np.argsort(score, kind="mergesort")[::-1] + y_true = y_true[desc_score_indices] + treatment = treatment[desc_score_indices] + uplift = score[desc_score_indices] + + distinct_value_indices = np.where(np.diff(uplift))[0] + threshold_indices = np.concatenate([distinct_value_indices, [uplift.size - 1]]) + + cumu_num_trmnt = np.cumsum(treatment)[threshold_indices] + y_trmnt = np.cumsum(np.where(treatment, y_true, 0))[threshold_indices] + + cumu_num_all = threshold_indices + 1 + cumu_num_ctrl = cumu_num_all - cumu_num_trmnt + y_ctrl = np.cumsum(np.where(~treatment, y_true, 0))[threshold_indices] + + mask = (cumu_num_trmnt >= min_n_per) & (cumu_num_ctrl >= min_n_per) + ratio = cumu_num_trmnt[mask] / cumu_num_ctrl[mask] + curve_values = y_trmnt[mask] - y_ctrl[mask] * ratio + # TODO + # if num_all.size == 0 or curve_values[0] != 0 or num_all[0] != 0: + # # Add an extra threshold position if necessary + # # to make sure that the curve starts at (0, 0) + # curve_values = np.r_[0, curve_values] + + out = np.full(len(cumu_num_all), np.nan) + out[mask] = curve_values + if normalize: + out /= out[-1] + + return out + + +def qini_scorer(estimator, X, y, normalize=True) -> float: + if isinstance(estimator, Pipeline): + X_transformed = estimator[:-1].transform(X) + return estimator[-1].score(X_transformed, y, method='qini', normalize=normalize) + return estimator.score(X, y, method='qini', normalize=normalize) + + +# ── Elasticity helpers ──────────────────────────────────────────────────────── +# Adapted from https://matheusfacure.github.io/python-causality-handbook/21-Meta-Learners.html + +def elast(data, y: str, t: str) -> float: + """ + OLS slope of y on t -- equivalent to the ATE under randomization. + Matches the @curry elast from the causality handbook. + """ + t_vals = np.asarray(data[t], dtype=float) + y_vals = np.asarray(data[y], dtype=float) + num = np.sum((t_vals - t_vals.mean()) * (y_vals - y_vals.mean())) + den = np.sum((t_vals - t_vals.mean()) ** 2) + return float(num / den) if den != 0 else np.nan + + +def elast_ci(data, y: str, t: str, z: float = 1.96) -> np.ndarray: + """95% confidence interval around the elasticity estimate.""" + n = len(data) + t_bar = np.asarray(data[t], dtype=float).mean() + beta1 = elast(data, y, t) + beta0 = np.asarray(data[y], dtype=float).mean() - beta1 * t_bar + e = np.asarray(data[y], dtype=float) - (beta0 + beta1 * np.asarray(data[t], dtype=float)) + se = np.sqrt(((1 / (n - 2)) * np.sum(e ** 2)) / + np.sum((np.asarray(data[t], dtype=float) - t_bar) ** 2)) + return np.array([beta1 - z * se, beta1 + z * se]) + + +# ── Cumulative gain / elasticity curves ─────────────────────────────────────────── +# Adapted from https://matheusfacure.github.io/python-causality-handbook/21-Meta-Learners.html +# Code here: https://github.com/matheusfacure/python-causality-handbook/blob/master/causal-inference-for-the-brave-and-true/nb21.py + +def cumulative_gain(dataset, prediction: str, y: str, t: str, + min_periods: int = 30, steps: int = 100) -> np.ndarray: + """ + Cumulative gain curve -- matches the causality handbook implementation. + Returns a 1-D array of length ~steps suitable for plt.plot(). + """ + size = dataset.shape[0] + ordered_df = dataset.sort_values(prediction, ascending=False).reset_index(drop=True) + n_rows = list(range(min_periods, size, size // steps)) + [size] + return np.array([0.0] + [elast(ordered_df.head(rows), y, t) * (rows / size) + for rows in n_rows]) + + +def cumulative_gain_ci(dataset, prediction: str, y: str, t: str, + min_periods: int = 30, steps: int = 100) -> np.ndarray: + """ + Cumulative gain curve with 95% CI bands. + Returns an (N, 2) array of [lower, upper] at each step. + """ + size = dataset.shape[0] + ordered_df = dataset.sort_values(prediction, ascending=False).reset_index(drop=True) + n_rows = list(range(min_periods, size, size // steps)) + [size] + return np.array([[0.0, 0.0]] + [elast_ci(ordered_df.head(rows), y, t) * (rows / size) + for rows in n_rows]) + + +def cumulative_elast_curve_ci(dataset, prediction: str, y: str, t: str, + min_periods: int = 30, steps: int = 100) -> np.ndarray: + """ + Cumulative elasticity curve with 95% CI (not multiplied by rows/size). + Returns an (N, 2) array of [lower, upper] at each step. + """ + size = dataset.shape[0] + ordered_df = dataset.sort_values(prediction, ascending=False).reset_index(drop=True) + n_rows = list(range(min_periods, size, size // steps)) + [size] + return np.array([elast_ci(ordered_df.head(rows), y, t) for rows in n_rows]) + + +# ── Cumulative gain scalar score ────────────────────────────────────────────── + +def get_cumulative_gain_score(y_true: np.ndarray, + treatment: np.ndarray, + score: np.ndarray, + steps: int = 100, + min_periods: int = 30) -> float: + """ + Compute the area between the cumulative gain curve and the random baseline + diagonal as a scalar model quality score. + + The cumulative gain curve sorts observations by predicted CATE descending, + then at each population fraction k computes elast(top k%) * k -- the + expected gain from treating only the top k% of players. The random baseline + is a straight line from (0, 0) to (1, ATE). A model that ranks the most + treatment-responsive players first arcs above the baseline early, yielding + a positive score. A model with no ranking ability tracks the diagonal and + scores near zero. + + :param y_true: Binary outcome array (0/1). + :param treatment: Binary treatment indicator (0/1 or bool). + :param score: Predicted CATE or uplift score for each observation. Higher + values are assumed to indicate higher predicted treatment responsiveness. + :param steps: Number of evaluation points along the population fraction + axis. More steps give a smoother curve and more precise AUC estimate. + Default is 100. + :param min_periods: Minimum number of observations required before + computing elasticity at a given population fraction. Avoids unstable + estimates at very small sample sizes. Default is 30. + :return: float -- area between the cumulative gain curve and the random + baseline. Positive = model beats random targeting. Returns np.nan if + fewer than 2 finite evaluation points exist. + """ + n = len(y_true) + order = np.argsort(score)[::-1] + y_s = y_true[order] + t_s = treatment[order] + + n_rows = list(range(min_periods, n, max(1, n // steps))) + [n] + + sorted_df = pd.DataFrame({'y': y_s, 't': t_s}) + full_df = pd.DataFrame({'y': y_true, 't': treatment}) + + gains = np.array([elast(sorted_df.head(k), 'y', 't') * (k / n) for k in n_rows]) + xs = np.array([k / n for k in n_rows]) + baseline = elast(full_df, 'y', 't') + baseline_curve = xs * baseline + + finite = np.isfinite(gains) & np.isfinite(baseline_curve) + if finite.sum() < 2: + return np.nan + + return float( + np.trapezoid(gains[finite], xs[finite]) - + np.trapezoid(baseline_curve[finite], xs[finite]) + ) + + +def cumulative_gain_scorer(estimator, X, y) -> float: + """ + Scorer callable with signature ``(estimator, X, y)`` for use with + ``sklearn.model_selection.cross_validate`` or ``GridSearchCV``. + + :param estimator: A fitted estimator with a ``score(X, y, method=...)`` method. + :param X: Feature matrix passed to ``estimator.score``. + :param y: SliceDict with keys ``'value'`` (outcome) and ``'is_treatment'`` + (treatment indicator), passed to ``estimator.score``. + :return: float -- area between the cumulative gain curve and the random + baseline. Positive values indicate the model beats random targeting. + + Example (cross_validate) + ------------------------ + :: + + from foundry.uplift.util import cumulative_gain_scorer + from sklearn.model_selection import StratifiedKFold, cross_validate + + skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42) + splits = list(skf.split(df, t)) + + scores_dict = cross_validate( + model, + X=df, + y=SliceDict(value=df['outcome'], is_treatment=df['treatment']), + cv=splits, + scoring=cumulative_gain_scorer, + return_train_score=True, + ) + + print(scores_dict['test_score']) + print(scores_dict['train_score']) + + Example (GridSearchCV) + ---------------------- + :: + + from foundry.uplift.util import cumulative_gain_scorer + from sklearn.model_selection import GridSearchCV + + gs = GridSearchCV( + estimator=model, + param_grid={'tlearner__estimator__max_depth': [2, 3, 4]}, + scoring=cumulative_gain_scorer, + cv=splits, + refit=True, + ) + + gs.fit( + df, + SliceDict(value=df['outcome'], is_treatment=df['treatment']), + ) + + print(f"Best params: {gs.best_params_}") + print(f"Best score: {gs.best_score_:.4f}") + """ + # unwrap pipeline if needed to reach the learner's score method + if isinstance(estimator, Pipeline): + # transform X through all steps except the last + X_transformed = estimator[:-1].transform(X) + return estimator[-1].score(X_transformed, y, method='cumulative_gain') + return estimator.score(X, y, method='cumulative_gain') diff --git a/foundry/util.py b/foundry/util.py index 0160f13..1ec2a75 100644 --- a/foundry/util.py +++ b/foundry/util.py @@ -11,6 +11,22 @@ ModelMatrix = Union[np.ndarray, pd.DataFrame, Dict[str, Union[np.ndarray, pd.DataFrame]]] +def safe_predict(estimator, *args, **kwargs) -> np.ndarray: + if hasattr(estimator, 'predict_proba'): + try: + out = estimator.predict_proba(*args, **kwargs) + except NotImplementedError: + out = None + + if out is not None: + if out.shape[1] == 2: + out = out[:, 1] + elif out.shape[1] > 2: + raise NotImplementedError("Multi-class predict_proba not supported.") + return out + return estimator.predict(*args, **kwargs) + + def transpose_last_dims(x: torch.Tensor) -> torch.Tensor: args = list(range(len(x.shape))) args[-2], args[-1] = args[-1], args[-2] From 67f2506ac80a0be3ef92451a35634c67d549e12f Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 17 Jul 2026 10:47:19 -0500 Subject: [PATCH 13/23] fix merge --- foundry/glm/glm.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/foundry/glm/glm.py b/foundry/glm/glm.py index b284c5e..a6b83ad 100644 --- a/foundry/glm/glm.py +++ b/foundry/glm/glm.py @@ -446,13 +446,14 @@ def _get_family_kwargs(self, **kwargs) -> dict: def _get_xdict(self, X: ModelMatrix, sparse_threshold: float) -> Dict[str, torch.Tensor]: _to_kwargs = get_to_kwargs(self.module_) - Xdict = self.to_slice_dict_.transform(X) - # convert to tensors: - for nm in list(Xdict): - if is_array(Xdict[nm]): - Xdict[nm] = to_tensor(Xdict[nm], sparse_threshold=sparse_threshold, **_to_kwargs) + Xdict = SliceDict(**{ + nm: (to_tensor(v, sparse_threshold=sparse_threshold, **_to_kwargs) if is_array(v) else v) + for nm, v in self.to_slice_dict_.transform(X).items() + }) + # validate: + for nm in list(Xdict): if nm in self.family.params: # model-mat params assert len(Xdict[nm].shape) == 2, f"len(X['{nm}'].shape)!=2" From 1edaad81a3ff3bad2e7ba3022ae38340d21d1d1c Mon Sep 17 00:00:00 2001 From: Jacob Date: Mon, 20 Jul 2026 15:06:50 -0500 Subject: [PATCH 14/23] add metalearner base class --- foundry/uplift/meta_learners/base.py | 17 +++++++++++++++++ foundry/uplift/meta_learners/s_learner.py | 20 +++++--------------- foundry/uplift/meta_learners/t_learner.py | 19 +++++-------------- foundry/uplift/meta_learners/x_learner.py | 21 ++++++--------------- 4 files changed, 33 insertions(+), 44 deletions(-) create mode 100644 foundry/uplift/meta_learners/base.py diff --git a/foundry/uplift/meta_learners/base.py b/foundry/uplift/meta_learners/base.py new file mode 100644 index 0000000..ab6878c --- /dev/null +++ b/foundry/uplift/meta_learners/base.py @@ -0,0 +1,17 @@ +import warnings + +import numpy as np +from foundry.util import SliceDict, to_1d +from sklearn.base import BaseEstimator + + +class MetaLearner(BaseEstimator): + @classmethod + def normalize_y(cls, y: SliceDict) -> tuple[np.ndarray, np.ndarray]: + y = y.copy() + y_arr = to_1d(np.asanyarray(y.pop("value"))) + treatment_ind = to_1d(np.asanyarray(y.pop("is_treatment")).astype(bool)) + # TODO: validate treatment is binary? + if len(y): + warnings.warn(f"Unused keys in ``y``: {set(y)}") + return y_arr, treatment_ind diff --git a/foundry/uplift/meta_learners/s_learner.py b/foundry/uplift/meta_learners/s_learner.py index 121f914..cb5e8ea 100644 --- a/foundry/uplift/meta_learners/s_learner.py +++ b/foundry/uplift/meta_learners/s_learner.py @@ -1,15 +1,15 @@ -import warnings from typing import Tuple, Union, overload, Literal, Optional import numpy as np import pandas as pd from sklearn.base import BaseEstimator, clone -from foundry.util import SliceDict, to_1d, safe_predict +from foundry.util import SliceDict, safe_predict from ..util import get_qini_curve, get_cumulative_gain_score +from .base import MetaLearner -class SLearner(BaseEstimator): +class SLearner(MetaLearner): """ An S-learner. Please note that current implementation assumes randomized treatment/control! @@ -23,14 +23,13 @@ def __init__(self, estimator: BaseEstimator, include_interaction: bool = True) - self.include_interaction = include_interaction def fit(self, X: Union[pd.DataFrame, np.ndarray], y: SliceDict, **fit_kwargs) -> "SLearner": - y_arr, treatment_ind = self._normalize_y(y) + y_arr, treatment_ind = self.normalize_y(y) X_aug = self._augment_with_treatment(X, treatment_ind) self.estimator_ = clone(self.estimator).fit(X=X_aug, y=y_arr, **fit_kwargs) return self - @overload def predict(self, X: Union[pd.DataFrame, np.ndarray], return_components: Literal[False] = ..., **predict_kwargs) -> np.ndarray: ... @overload @@ -57,7 +56,7 @@ def score( normalize: bool = True, **kwargs, ) -> float: - y_arr, treatment_ind = self._normalize_y(y) + y_arr, treatment_ind = self.normalize_y(y) if sample_weight is not None: raise NotImplementedError pred = self.predict(X=X) @@ -103,12 +102,3 @@ def _augment_with_treatment( return np.hstack([X, treatment_col, interaction_terms]) return np.hstack([X, treatment_col]) - - @staticmethod - def _normalize_y(y: SliceDict) -> Tuple[np.ndarray, np.ndarray]: - y = y.copy() - y_arr = to_1d(np.asanyarray(y.pop("value"))) - treatment_ind = to_1d(np.asanyarray(y.pop("is_treatment")).astype(bool)) - if len(y.keys()): - warnings.warn(f"Unused keys in ``y``: {set(y)}") - return y_arr, treatment_ind diff --git a/foundry/uplift/meta_learners/t_learner.py b/foundry/uplift/meta_learners/t_learner.py index 4e50f46..162a286 100644 --- a/foundry/uplift/meta_learners/t_learner.py +++ b/foundry/uplift/meta_learners/t_learner.py @@ -1,15 +1,15 @@ -import warnings from typing import Optional, Tuple, Union, overload, Literal import numpy as np import pandas as pd from sklearn.base import BaseEstimator, clone -from foundry.util import SliceDict, to_1d, safe_predict +from foundry.util import SliceDict, safe_predict from ..util import get_qini_curve, get_cumulative_gain_score +from .base import MetaLearner -class TLearner(BaseEstimator): +class TLearner(MetaLearner): """ A T-learner. Please note that current implementation assumes randomized treatment/control! @@ -22,7 +22,7 @@ def __init__(self, estimator: BaseEstimator) -> None: self.estimator = estimator def fit(self, X: Union[pd.DataFrame, np.ndarray], y: SliceDict, **fit_kwargs) -> "TLearner": - y_arr, treatment_ind = self._normalize_y(y) + y_arr, treatment_ind = self.normalize_y(y) self.treatment_est_ = clone(self.estimator).fit(X=X[treatment_ind], y=y_arr[treatment_ind], **fit_kwargs) self.control_est_ = clone(self.estimator).fit(X[~treatment_ind], y_arr[~treatment_ind], **fit_kwargs) @@ -49,7 +49,7 @@ def score( normalize: bool = True, **kwargs, ) -> float: - y_arr, treatment_ind = self._normalize_y(y) + y_arr, treatment_ind = self.normalize_y(y) if sample_weight is not None: raise NotImplementedError pred = self.predict(X=X) @@ -69,12 +69,3 @@ def score( ) random_area = np.linspace(0, qini[-1], qini.shape[0]).sum() return (np.nansum(qini) - random_area) / qini.shape[0] - - @staticmethod - def _normalize_y(y: SliceDict) -> Tuple[np.ndarray, np.ndarray]: - y = y.copy() - y_arr = to_1d(np.asanyarray(y.pop('value'))) - treatment_ind = to_1d(np.asanyarray(y.pop('is_treatment')).astype(bool)) - if len(y.keys()): - warnings.warn(f"Unused keys in ``y``: {set(y)}") - return y_arr, treatment_ind diff --git a/foundry/uplift/meta_learners/x_learner.py b/foundry/uplift/meta_learners/x_learner.py index 89e19b5..0286957 100644 --- a/foundry/uplift/meta_learners/x_learner.py +++ b/foundry/uplift/meta_learners/x_learner.py @@ -1,15 +1,15 @@ -import warnings from typing import Any, Dict, Optional, Tuple, Union, overload, Literal import numpy as np import pandas as pd from sklearn.base import BaseEstimator, clone -from foundry.util import SliceDict, to_1d, safe_predict +from foundry.util import SliceDict, safe_predict from ..util import get_qini_curve, get_cumulative_gain_score +from .base import MetaLearner -class XLearner(BaseEstimator): +class XLearner(MetaLearner): """ An X-learner. Please note that current implementation assumes randomized treatment/control! Adapted from https://matheusfacure.github.io/python-causality-handbook/21-Meta-Learners.html @@ -45,7 +45,7 @@ def __init__( self.propensity_fit_params = propensity_fit_params def fit(self, X: Union[pd.DataFrame, np.ndarray], y: SliceDict) -> "XLearner": - y_arr, treatment_ind = self._normalize_y(y) + y_arr, treatment_ind = self.normalize_y(y) _first_stage_fit_params = self.first_stage_fit_params or {} _second_stage_fit_params = self.second_stage_fit_params or {} @@ -86,7 +86,7 @@ def predict(self, X, return_components=False, **predict_kwargs): p_treatment = safe_predict(self.propensity_est_, X) p_control = 1 - p_treatment - tau0 = safe_predict(self.second_control_est_, X, **predict_kwargs) + tau0 = safe_predict(self.second_control_est_, X, **predict_kwargs) tau1 = safe_predict(self.second_treatment_est_, X, **predict_kwargs) if return_components: @@ -102,7 +102,7 @@ def score( normalize: bool = True, **kwargs, ) -> float: - y_arr, treatment_ind = self._normalize_y(y) + y_arr, treatment_ind = self.normalize_y(y) if sample_weight is not None: raise NotImplementedError pred = self.predict(X=X) @@ -122,12 +122,3 @@ def score( ) random_area = np.linspace(0, qini[-1], qini.shape[0]).sum() return (np.nansum(qini) - random_area) / qini.shape[0] - - @staticmethod - def _normalize_y(y: SliceDict) -> Tuple[np.ndarray, np.ndarray]: - y = y.copy() - y_arr = to_1d(np.asanyarray(y.pop('value'))) - treatment_ind = to_1d(np.asanyarray(y.pop('is_treatment')).astype(bool)) - if len(y.keys()): - warnings.warn(f"Unused keys in ``y``: {set(y)}") - return y_arr, treatment_ind From 9a133f1241cfcab05a5d149e371ed2ee9b3d2fc8 Mon Sep 17 00:00:00 2001 From: Jacob Date: Wed, 22 Jul 2026 12:06:11 -0500 Subject: [PATCH 15/23] safe_predict allows len(array.shape)==1 --- foundry/util.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/foundry/util.py b/foundry/util.py index 1ec2a75..a81ed8f 100644 --- a/foundry/util.py +++ b/foundry/util.py @@ -19,10 +19,13 @@ def safe_predict(estimator, *args, **kwargs) -> np.ndarray: out = None if out is not None: - if out.shape[1] == 2: - out = out[:, 1] - elif out.shape[1] > 2: - raise NotImplementedError("Multi-class predict_proba not supported.") + if len(out.shape) == 2: + if out.shape[1] == 2: + out = out[:, 1] + elif out.shape[1] > 2: + raise NotImplementedError("Multi-class ``predict_proba`` not supported.") + elif len(out.shape) > 2: + raise RuntimeError(f"Expected 1d or 2d but {estimator:} returned shape {out.shape}.") return out return estimator.predict(*args, **kwargs) From e8a56453274f1c7d4ae46da2edc4b83805d73e06 Mon Sep 17 00:00:00 2001 From: Jacob Date: Wed, 22 Jul 2026 12:06:31 -0500 Subject: [PATCH 16/23] add iloc to SliceDict after super.init call --- foundry/util.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/foundry/util.py b/foundry/util.py index a81ed8f..25431fe 100644 --- a/foundry/util.py +++ b/foundry/util.py @@ -154,13 +154,11 @@ def __init__(self, **kwargs): else: self._len = lengths[0] + super().__init__(**kwargs) + # sklearn checks if it should use pandas indexing by checking if there's an iloc attribute if self.is_pandas: - self.__dict__['iloc'] = True - else: - self.__dict__.pop('iloc', None) - - super().__init__(**kwargs) + self.iloc = True @property def is_pandas(self) -> bool: From a2ac02360560cb61ae393f8a092f9d3aabd7dc38 Mon Sep 17 00:00:00 2001 From: Jacob Date: Wed, 22 Jul 2026 13:11:00 -0500 Subject: [PATCH 17/23] add an __array__ method for clearer error message when using default sklearn scoring --- foundry/util.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/foundry/util.py b/foundry/util.py index 25431fe..6fd9c99 100644 --- a/foundry/util.py +++ b/foundry/util.py @@ -176,6 +176,15 @@ def _standardize_val(val): def __len__(self) -> int: return self._len + def __array__(self) -> np.ndarray: + array = self.get('__array__', None) + if array is None: + raise RuntimeError( + f"Tried to convert a {type(self).__name__} to an array, which is ambiguous. If you'd like this to work," + f" you should add a key '__array__', which can have its value passed to np.asarray." + ) + return np.asarray(array) + def __getitem__(self, sl: Union[int, str, slice]) -> Union['SliceDict', ArrayType]: if isinstance(sl, int): raise ValueError( @@ -234,7 +243,7 @@ def __repr__(self): def shape(self): return (self._len,) - def copy(self): + def copy(self) -> 'SliceDict': return type(self)(**self) def fromkeys(self, *args, **kwargs): From 06dcc7c6a7ec9bb428be97f314468c51b91c4456 Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 7 Aug 2026 14:45:31 -0500 Subject: [PATCH 18/23] add more supported estimators --- foundry/evaluation/one_se_rule.py | 105 +++++++++++++++++++----------- 1 file changed, 68 insertions(+), 37 deletions(-) diff --git a/foundry/evaluation/one_se_rule.py b/foundry/evaluation/one_se_rule.py index be97269..6201f4f 100644 --- a/foundry/evaluation/one_se_rule.py +++ b/foundry/evaluation/one_se_rule.py @@ -3,6 +3,19 @@ import numpy as np from sklearn.base import BaseEstimator +from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor +from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor + +try: + import lightgbm as lgb + + _LGBM_DEFAULTS = lgb.LGBMRegressor().get_params() +except ImportError: + lgb = None + _LGBM_DEFAULTS = {"n_estimators": 100, "learning_rate": 0.1, "max_depth": -1, "num_leaves": 31} + +_DT_DEFAULTS = {**DecisionTreeClassifier().get_params(), **DecisionTreeRegressor().get_params()} +_RF_DEFAULTS = {**RandomForestClassifier().get_params(), **RandomForestRegressor().get_params()} class OneSeRule: @@ -126,51 +139,69 @@ def complexity(i): @classmethod def _estimator_to_complexity_func(cls, estimator: BaseEstimator) -> Callable: - estimator_to_complexity_func = [] + if isinstance(estimator, + (DecisionTreeClassifier, DecisionTreeRegressor)): + return cls._decision_tree_complexity + + if isinstance(estimator, + (RandomForestClassifier, RandomForestRegressor)): + return cls._random_forest_complexity - _lgb_complexity = cls._lgb_complexity() - if _lgb_complexity: - estimator_to_complexity_func.append(_lgb_complexity) + if lgb is not None and isinstance(estimator, + (lgb.LGBMRegressor, lgb.LGBMClassifier, lgb.LGBMModel, lgb.LGBMRanker)): + return cls._lgbm_complexity - params_to_complexity = next( - (func for types, func in estimator_to_complexity_func if isinstance(estimator, types)), - None + raise ValueError( + "Do not know how to define complexity for {}, please provide a callable that takes params and " + "returns a float".format(type(estimator).__name__) ) - if params_to_complexity is None: + + @staticmethod + def _decision_tree_complexity(max_depth=_DT_DEFAULTS["max_depth"], + max_leaf_nodes=_DT_DEFAULTS["max_leaf_nodes"], + **kwargs) -> float: + unexpected = set(kwargs) - set(_DT_DEFAULTS) + if unexpected: raise ValueError( - "Do not know how to define complexity for {}, please provide a callable that takes params and " - "returns a float".format(type(estimator).__name__) + f"Got params not recognized by the DecisionTree sklearn API defaults: {unexpected}. " ) + return _treelike_complexity(max_depth=max_depth, max_leaf_nodes=max_leaf_nodes, **kwargs) - return params_to_complexity + @staticmethod + def _random_forest_complexity(max_depth=_RF_DEFAULTS["max_depth"], + max_leaf_nodes=_RF_DEFAULTS["max_leaf_nodes"], + **kwargs) -> float: + unexpected = set(kwargs) - set(_RF_DEFAULTS) + if unexpected: + raise ValueError( + f"Got params not recognized by the RandomForest sklearn API defaults: {unexpected}. " + ) + return _treelike_complexity(max_depth=max_depth, max_leaf_nodes=max_leaf_nodes, **kwargs) @staticmethod - def _lgb_complexity() -> Optional[tuple[tuple, Callable]]: - try: - import lightgbm as lgb - except ImportError: - return None - - _LGBM_DEFAULTS = lgb.LGBMRegressor().get_params() - - def _lgbm_complexity(n_estimators=_LGBM_DEFAULTS["n_estimators"], - learning_rate=_LGBM_DEFAULTS["learning_rate"], - max_depth=_LGBM_DEFAULTS["max_depth"], - num_leaves=_LGBM_DEFAULTS["num_leaves"], - **kwargs) -> float: - unexpected = set(kwargs) - set(_LGBM_DEFAULTS) - if unexpected: - raise ValueError( - f"Got params not recognized by the LGBM sklearn API defaults: {unexpected}. " - "This usually means either the `prefix` is wrong, or you're using a native " - "LightGBM alias (e.g. `min_data_in_leaf`, `feature_fraction`, `lambda_l1`) " - "instead of the sklearn wrapper's canonical param name." - ) - cap = np.inf if (max_depth is None or max_depth <= 0) else 2 ** max_depth - effective_iterations = n_estimators * learning_rate - return effective_iterations * min(cap, num_leaves) - - return (lgb.LGBMRegressor, lgb.LGBMClassifier, lgb.LGBMModel, lgb.LGBMRanker), _lgbm_complexity + def _lgbm_complexity(n_estimators=_LGBM_DEFAULTS["n_estimators"], + learning_rate=_LGBM_DEFAULTS["learning_rate"], + max_depth=_LGBM_DEFAULTS["max_depth"], + num_leaves=_LGBM_DEFAULTS["num_leaves"], + **kwargs) -> float: + unexpected = set(kwargs) - set(_LGBM_DEFAULTS) + if unexpected: + raise ValueError( + f"Got params not recognized by the LGBM sklearn API defaults: {unexpected}. " + "Are you using a native LightGBM alias (e.g. `min_data_in_leaf`, `feature_fraction`, `lambda_l1`) " + "instead of the sklearn wrapper's canonical param name?" + ) + cap = np.inf if (max_depth is None or max_depth <= 0) else 2 ** max_depth + effective_iterations = n_estimators * learning_rate + return effective_iterations * min(cap, num_leaves) + + +def _treelike_complexity(**kwargs) -> float: + max_depth = kwargs.pop("max_depth", None) + max_leaf_nodes = kwargs.pop("max_leaf_nodes", None) + depth_cap = np.inf if (max_depth is None or max_depth <= 0) else 2 ** max_depth + leaf_cap = np.inf if max_leaf_nodes is None else max_leaf_nodes + return min(depth_cap, leaf_cap) def _default_complexity_reduce_fun(complexities: Sequence[float]) -> float: From 656757844aad82ff68db3400871816e058c96137 Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 7 Aug 2026 14:47:46 -0500 Subject: [PATCH 19/23] check for need for complexity_reduce_fun later --- foundry/evaluation/one_se_rule.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/foundry/evaluation/one_se_rule.py b/foundry/evaluation/one_se_rule.py index 6201f4f..238d806 100644 --- a/foundry/evaluation/one_se_rule.py +++ b/foundry/evaluation/one_se_rule.py @@ -71,11 +71,6 @@ def __init__(self, else: estimators = {prefix: estimator} - if complexity_reduce_fun is None: - if len(estimators) == 1: - complexity_reduce_fun = _default_complexity_reduce_fun - else: - raise ValueError("If multiple estimators are passed, must supply `complexity_reduce_fun`.") self.complexity_reduce_fun = complexity_reduce_fun self.params_to_complexity_funs = {} @@ -119,6 +114,11 @@ def __call__(self, cv_results: dict[str, np.ndarray]) -> int: f"{list(self.params_to_complexity_funs)} matched any params in this grid " f"(got keys {any_row_keys}). Nothing would be measured for complexity." ) + if self.complexity_reduce_fun is None: + if len(active_prefixes) == 1: + self.complexity_reduce_fun = _default_complexity_reduce_fun + else: + raise ValueError("If multiple estimators are passed, must supply `complexity_reduce_fun`.") def complexity(i): params_this_row = cv_results["params"][i] From 37b3ac3eb94f44c71cbab3142dff6de4d76fb44c Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 7 Aug 2026 14:52:03 -0500 Subject: [PATCH 20/23] defer check --- foundry/evaluation/one_se_rule.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/foundry/evaluation/one_se_rule.py b/foundry/evaluation/one_se_rule.py index 238d806..1917c0d 100644 --- a/foundry/evaluation/one_se_rule.py +++ b/foundry/evaluation/one_se_rule.py @@ -104,6 +104,11 @@ def __call__(self, cv_results: dict[str, np.ndarray]) -> int: active_prefixes = {} for prefix, func in self.params_to_complexity_funs.items(): if any(k.startswith(prefix) for k in any_row_keys): + if func is None: + raise ValueError( + "Do not know how to define complexity for '{}', please provide a callable that takes params and" + " returns a float".format(prefix) + ) active_prefixes[prefix] = func elif self.verbose: print(f"`{prefix}` does not appear to be getting tuned so will not contribute to complexity calcs") @@ -138,7 +143,7 @@ def complexity(i): return min(candidates, key=complexity) @classmethod - def _estimator_to_complexity_func(cls, estimator: BaseEstimator) -> Callable: + def _estimator_to_complexity_func(cls, estimator: BaseEstimator) -> Optional[Callable]: if isinstance(estimator, (DecisionTreeClassifier, DecisionTreeRegressor)): return cls._decision_tree_complexity @@ -151,10 +156,7 @@ def _estimator_to_complexity_func(cls, estimator: BaseEstimator) -> Callable: (lgb.LGBMRegressor, lgb.LGBMClassifier, lgb.LGBMModel, lgb.LGBMRanker)): return cls._lgbm_complexity - raise ValueError( - "Do not know how to define complexity for {}, please provide a callable that takes params and " - "returns a float".format(type(estimator).__name__) - ) + return None @staticmethod def _decision_tree_complexity(max_depth=_DT_DEFAULTS["max_depth"], From 3207bfb9da7d6a54b30e63917fc458c85fd2e3f6 Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 7 Aug 2026 14:57:08 -0500 Subject: [PATCH 21/23] allow specifying the threshold directly --- foundry/evaluation/one_se_rule.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/foundry/evaluation/one_se_rule.py b/foundry/evaluation/one_se_rule.py index 1917c0d..6eb842c 100644 --- a/foundry/evaluation/one_se_rule.py +++ b/foundry/evaluation/one_se_rule.py @@ -39,15 +39,18 @@ class OneSeRule: per-sub-estimator complexity floats and reduces them to one float (or tuple) for the combined model -- e.g. `math.prod` for a hurdle model, where the branches compose serially (partition refinement), which multiplication captures better than summing independent contributors. + :param score_threshold: Rather than use the one-SE rule, simply specify a threshold for the score directly. """ def __init__(self, estimator: Union[BaseEstimator, Callable], sub_estimators: Union[str, Sequence[str]] = (), complexity_reduce_fun: Optional[Callable] = None, + score_threshold: Optional[float] = None, verbose: bool = True): self.verbose = verbose + self.score_threshold = score_threshold prefix, estimator = self._peel_pipeline(estimator) @@ -91,13 +94,16 @@ def _peel_pipeline(estimator) -> tuple[str, BaseEstimator]: def __call__(self, cv_results: dict[str, np.ndarray]) -> int: mean_scores = np.array(cv_results["mean_test_score"]) - std_scores = np.array(cv_results["std_test_score"]) - n_splits = sum(1 for k in cv_results if k.startswith("split") and k.endswith("_test_score")) - se = std_scores / np.sqrt(n_splits) + if self.score_threshold is not None: + threshold = self.score_threshold + else: + std_scores = np.array(cv_results["std_test_score"]) + n_splits = sum(1 for k in cv_results if k.startswith("split") and k.endswith("_test_score")) + se = std_scores / np.sqrt(n_splits) + best_idx = np.argmax(mean_scores) + threshold = mean_scores[best_idx] - se[best_idx] - best_idx = np.argmax(mean_scores) - threshold = mean_scores[best_idx] - se[best_idx] candidates = [i for i, s in enumerate(mean_scores) if s >= threshold] any_row_keys = set(cv_results["params"][0]) From 53d5a824befc79c9c5623ed49227fa5b530cb234 Mon Sep 17 00:00:00 2001 From: Jacob Date: Tue, 11 Aug 2026 12:03:09 -0500 Subject: [PATCH 22/23] Create composite_model.py --- foundry/composite_model.py | 212 +++++++++++++++++++++++++++++++++++++ 1 file changed, 212 insertions(+) create mode 100644 foundry/composite_model.py diff --git a/foundry/composite_model.py b/foundry/composite_model.py new file mode 100644 index 0000000..a951c33 --- /dev/null +++ b/foundry/composite_model.py @@ -0,0 +1,212 @@ +import keyword +from typing import Sequence, Union, Collection +from warnings import warn + +import numpy as np +import pandas as pd +from foundry.util import to_1d, safe_predict +from sklearn.utils import _safe_indexing +from sklearn.utils.metaestimators import _BaseComposition +from sklearn.base import BaseEstimator, clone + + +class CompositeModel(_BaseComposition): + r""" + A model whose predictions are generated by combining the predictions of sub-models in a deterministic fashion (such + as a hurdle model). + + Subclasses should implement a ``sub_model_names`` class attribute, the ``_combine_predictions()`` method, and the + ``_score()`` method. + + Declaring a subclass needs nothing but ``sub_model_names`` and the two abstract methods -- no ``__init__``:: + + >>> class HurdleModel(CompositeModel): + ... sub_model_names = ['event_est', 'value_given_event_est'] + ... + ... def _combine_predictions(self, components): + ... return components['event_est'] * components['value_given_event_est'] + ... + ... def _score(self, X, y, sample_weight=None): + ... from sklearn.metrics import r2_score + ... return r2_score(y['value_given_event_est'], self.predict(X), sample_weight=sample_weight) + + Subclassing magic + ------------------ + Setting ``sub_model_names`` on a subclass triggers ``__init_subclass__``, which does two things: + + 1. For each name in ``sub_model_names``, adds two read-only properties: ```` (the *unfitted* template, + read from ``self.sub_models``) and ``_`` (the *fitted* sub-model, read from ``self.sub_models_`` -- + only available after calling ``fit()``, per sklearn's trailing-underscore convention). + 2. If the subclass does *not* define its own ``__init__``, one is generated automatically with a named + parameter for each entry in ``sub_model_names`` -- so subclasses get an explicit, autocomplete-friendly + constructor for free, equivalent to calling ``super().__init__(dict(...))`` by hand. Writing your own + ``__init__`` (e.g. for extra validation or extra arguments) opts out of this entirely; nothing below + applies in that case. + + Even though no ``__init__`` was written, one was generated, so construction reads like an ordinary, explicit + constructor -- not the generic ``sub_models=`` form:: + + >>> from sklearn.linear_model import LogisticRegression, LinearRegression + >>> model = HurdleModel( + ... event_est=LogisticRegression(), + ... value_given_event_est=LinearRegression(), + ... ) + >>> model.event_est # unfitted template, via the auto-added property + >>> # LogisticRegression() + + After fitting, the fitted sub-models are available, and because the generated ``__init__`` signature lines up + exactly with what ``get_params()`` reports, ``clone()``/``get_params()``/``set_params()`` all work normally -- + no extra effort required despite the friendlier constructor:: + + >>> model.fit(X_train, y={'event_est': y_event, 'value_given_event_est': y_value}) + >>> model.event_est_ # the *fitted* estimator (distinct object from model.event_est) + >>> # LogisticRegression() + >>> from sklearn.base import clone + >>> clone(model) # round-trips correctly through get_params() -> __init__ + >>> # HurdleModel(event_est=LogisticRegression(), value_given_event_est=LinearRegression()) + """ + sub_model_names: Sequence[str] + + def __init__(self, sub_models: Union[dict, Sequence[tuple]]): + if isinstance(sub_models, dict): + sub_models = list(sub_models.items()) + self.sub_models = sub_models + + def predict(self, X, **kwargs): + components = self.predict_components(X, **kwargs) + return self._combine_predictions(components) + + def _combine_predictions(self, components): + raise NotImplementedError + + def predict_components(self, X, **kwargs): + kwargs_per_model = _organize_owned_kwargs(kwargs, possible_owners=set(self.sub_model_names)) + out = {} + for model_name, model in self.sub_models_: + out[model_name] = to_1d(safe_predict(model, X, **kwargs_per_model[model_name])) + return out + + def fit(self, X, y, verbose=False, **kwargs): + if kwargs.get('sample_weight', None) is None and isinstance(y, dict) and 'sample_weight' in y: + y = y.copy() + kwargs['sample_weight'] = y.pop('sample_weight') + + sub_model_names, _ = zip(*self.sub_models) + kwargs_per_model = _organize_owned_kwargs(kwargs, possible_owners=set(sub_model_names)) + + self.sub_models_ = [] + for nm in self.sub_model_names: + if verbose: + print(f"Fitting {nm}...") + model_template = dict(self.sub_models).get(nm, None) + if model_template is None: + raise ValueError(f"A model for `{nm}` was not provided at init") + + this_y = y[nm] + _is_valid = self._get_valid_mask(this_y) + this_y = _safe_indexing(this_y, _is_valid) + this_X = _safe_indexing(X, _is_valid) + this_kwargs = kwargs_per_model[nm].copy() + for k in list(this_kwargs): + if k == 'sample_weight' or k.endswith('__sample_weight'): + this_kwargs[k] = _safe_indexing(this_kwargs[k], _is_valid) + + fitted = clone(model_template).fit(this_X, this_y, **this_kwargs) + self.sub_models_.append((nm, fitted)) + return self + + def _get_valid_mask(self, y): + return np.array(pd.notnull(y)) + + def get_params(self, deep: bool = True) -> dict: + return self._get_params('sub_models', deep=deep) + + def set_params(self, **kwargs): + self._set_params('sub_models', **kwargs) + return self + + def score(self, X, y, sample_weight=None, score_sub_models: bool = False): + score = self._score(X=X, y=y, sample_weight=sample_weight) + if not score_sub_models: + return score + scores = {'combined': score} + for nm, model in self.sub_models_: + y_true = y[nm] + mask = self._get_valid_mask(y_true) + scores[nm] = model.score( + _safe_indexing(X, mask), + _safe_indexing(y_true, mask), + sample_weight=None if sample_weight is None else _safe_indexing(sample_weight, mask) + ) + return scores + + def _score(self, X, y, sample_weight=None): + raise NotImplementedError + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + for sub_model_name in cls.sub_model_names: + setattr(cls, sub_model_name, property(_GetHelper(sub_model_name, 'sub_models'))) + setattr(cls, sub_model_name + '_', property(_GetHelper(sub_model_name, 'sub_models_'))) + + if '__init__' not in cls.__dict__: + cls.__init__ = _build_nice_init(cls) + + +class _GetHelper: + def __init__(self, sub_model_name: str, dict_attr_name: str): + self.sub_model_name = sub_model_name + self.dict_attr_name = dict_attr_name + + def __call__(self, instance) -> BaseEstimator: + return dict(getattr(instance, self.dict_attr_name))[self.sub_model_name] + + +_RESERVED_NAMES = { + 'fit', 'predict', 'predict_components', 'score', + 'get_params', 'set_params', 'sub_models', 'sub_models_', 'sub_model_names', +} + + +def _build_nice_init(cls): + names = list(cls.sub_model_names) + if len(set(names)) != len(names): + raise ValueError(f"{cls.__name__}.sub_model_names has duplicate entries: {names}") + for name in names: + if not name.isidentifier() or keyword.iskeyword(name): + raise ValueError(f"{cls.__name__}: sub-model name {name!r} is not a valid Python identifier") + if name in _RESERVED_NAMES: + raise ValueError(f"{cls.__name__}: sub-model name {name!r} collides with the CompositeModel API") + + params_sig = ", ".join(names) + dict_body = ", ".join(f"{n}={n}" for n in names) + src = ( + f"def __init__(self, {params_sig}):\n" + f" CompositeModel.__init__(self, dict({dict_body}))\n" + ) + namespace = {'CompositeModel': CompositeModel} + filename = f"" + exec(compile(src, filename, 'exec'), namespace) + fn = namespace['__init__'] + fn.__qualname__ = f"{cls.__qualname__}.__init__" + fn.__doc__ = ( + f"Auto-generated from `sub_model_names` by CompositeModel.__init_subclass__.\n" + f"Equivalent to: super().__init__(dict({dict_body}))" + ) + return fn + + +def _organize_owned_kwargs(kwargs: dict, possible_owners: Collection[str]) -> dict[str, dict]: + possible_owners = set(possible_owners) + kwargs_per_owner = {owner: {} for owner in possible_owners} + for k, v in kwargs.items(): + owner, _, subk = k.partition('__') + if subk: + if owner in possible_owners: + kwargs_per_owner[owner].update({subk: v}) + else: + warn(f"{k} startswith `{owner}__`, but that's not an option: {possible_owners}. Will be ignored.") + else: + for owner in possible_owners: + kwargs_per_owner[owner].update({k: v}) + return kwargs_per_owner From 533fa22d8a6fd90cc818561b12405bc09cffd5aa Mon Sep 17 00:00:00 2001 From: Jacob Date: Tue, 11 Aug 2026 12:03:25 -0500 Subject: [PATCH 23/23] Update __init__.py --- foundry/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/foundry/__init__.py b/foundry/__init__.py index e913eb4..c5b683c 100644 --- a/foundry/__init__.py +++ b/foundry/__init__.py @@ -1 +1 @@ -__version__ = '0.2.10' +__version__ = '0.2.11'