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/__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' 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 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() diff --git a/foundry/evaluation/marginal_effects.py b/foundry/evaluation/marginal_effects.py index 0a99f09..6ee565e 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) @@ -267,13 +272,14 @@ 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, 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 ( @@ -286,17 +292,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: @@ -342,6 +360,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) + @@ -349,14 +374,21 @@ def plot(self, theme(figure_size=(8, 6), subplots_adjust={'wspace': 0.10}) ) if isinstance(data[x.replace('_binned', '')].dtype, pd.CategoricalDtype): + if group_by_prediction_col: + raise RuntimeError("Please explicitly map 'prediction_col' to color or facet.") plot += geom_col() else: - plot += geom_line() + 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: @@ -381,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) @@ -389,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 66e0e5a..a6b83ad 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 @@ -117,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): """ @@ -143,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., @@ -346,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 @@ -682,8 +697,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, @@ -699,23 +714,95 @@ 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() + 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, + 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)``, 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``. + """ + 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._sample_coef_mvnorm(n_iters=n_iters) + orig = copy.deepcopy(self.module_) + try: + out = [] + 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: + 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 = lambda x: np.concatenate(x, 1) + + return collate_fn(out) + def _get_penalty(self) -> torch.Tensor: """ Get penalty on sum(log_prob) based on the module-weights and self.penalty. @@ -778,26 +865,65 @@ 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[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.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 prefixes, 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_``. + """ + orig = copy.deepcopy(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_ + if start_ != 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, + 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: + 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) 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 = [] + 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 def family_from_string(family: str, y: Optional[dict] = None) -> Family: 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/interactions.py b/foundry/preprocessing/sklearn/interactions.py index b426af4..bdc08c6 100644 --- a/foundry/preprocessing/sklearn/interactions.py +++ b/foundry/preprocessing/sklearn/interactions.py @@ -145,8 +145,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 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.") diff --git a/foundry/util.py b/foundry/util.py index 8f47718..6fd9c99 100644 --- a/foundry/util.py +++ b/foundry/util.py @@ -11,6 +11,25 @@ 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 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) + + def transpose_last_dims(x: torch.Tensor) -> torch.Tensor: args = list(range(len(x.shape))) args[-2], args[-1] = args[-1], args[-2] @@ -116,22 +135,6 @@ class FitFailedException(RuntimeError): pass -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) - - class SliceDict(dict): """ Adapted from https://github.com/skorch-dev/skorch/blob/baf0580/skorch/helper.py#L20 diff --git a/setup.py b/setup.py index b03df1f..9940700 100644 --- a/setup.py +++ b/setup.py @@ -20,12 +20,14 @@ 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', ], - 'docs' : [ + 'docs': [ 'requests' ] }