def predict(self, X: Union[torch.Tensor, SliceDict], type: str = 'mean', *args, **kwargs) -> np.ndarray: """ Return an attribute of the distribution (by default the mean) as a numpy array. """ X = to_tensor(X, device=self.device, dtype=self.module_dtype_) y_out = [] for params in self.forward_iter(X, training=False): batch_size = len(params[0]) distribution_kwargs = dict( zip(self.distribution_param_names_, params)) dist = self.distribution(**distribution_kwargs) yp = getattr(dist, type) if callable(yp): yp = yp(*args, **kwargs) yp = to_numpy(yp) if yp.shape[0] != batch_size: raise RuntimeError( f"`{self.distribution.__name__}.{type}` produced a tensor whose leading dim is {yp.shape[0]}, " f"expected {batch_size}.") y_out.append(yp) y_out = np.concatenate(y_out, 0) return y_out
def estimate_laplace_params(self, X, y, **fit_params): means = torch.cat( [param.data.view(-1) for param in self.module_.parameters()]) # get loss, hessian: y_pred = self.infer(X, **fit_params) y_true = to_tensor(y, device=self.device, dtype=self.module_dtype_) loss = self.get_loss(y_pred, y_true, reduction='sum') hess = hessian(output=loss, inputs=list(self.module_.parameters()), allow_unused=True, progress=False) # create mvnorm for laplace approx: try: self.laplace_params_ = torch.distributions.MultivariateNormal( means, covariance_matrix=torch.inverse(hess)) self.converged_ = True except RuntimeError as e: if 'lapack' in str(e) or 'cholesky' in str(e): warn( "Model failed to converge; `laplace_params` cannot be estimated" ) fake_cov = (2 * means.abs().max() * torch.eye(len(hess)))**2 self.laplace_params_ = torch.distributions.MultivariateNormal( means, covariance_matrix=fake_cov) self.converged_ = False else: raise e
def predict_dataframe(self, dataframe: 'DataFrame', preprocessor: Union[ColumnTransformer, Sequence], type: str, x: torch.Tensor): """ Experimental. """ from pandas import DataFrame # check x, broadcast: x = to_tensor(x, device=self.device, dtype=self.module_dtype_) if len(x.shape) == 2: assert x.shape[1] == 1 elif len(x.shape) == 1: x = x[:, None] else: raise RuntimeError("Expected `x` to be 1D") # generate dist-param predictions: assert len(dataframe.index) == len(set(dataframe.index)) with torch.no_grad(): X = to_tensor(preprocessor.transform(dataframe), device=self.device, dtype=self.module_dtype_) params = self.infer(X) distribution_kwargs = dict( zip(self.distribution_param_names_, params)) dist = self.distribution(**distribution_kwargs) # plug x into method: pred_broadcasted = getattr(dist, type)(x) # flatten, joinable along original index: index_broadcasted = to_tensor(dataframe.index.values, device=self.device)[None, :].repeat( len(x), 1) x_broadcasted = x.repeat(1, len(dataframe.index)) return DataFrame({ 'index': index_broadcasted.view(-1).numpy(), 'x': x_broadcasted.view(-1).numpy(), type: pred_broadcasted.view(-1).numpy() })
def get_loss(self, y_pred: torch.Tensor, y_true: torch.Tensor, X: Optional[torch.Tensor] = None, training: bool = False, **kwargs): y_true = to_tensor(y_true, device=self.device, dtype=self.module_dtype_) neg_log_lik = self.criterion_(y_pred, y_true, **kwargs) penalty = self.criterion_.get_penalty(y_true=y_true, module=self.module_) return neg_log_lik + penalty
def partial_fit(self, X, y=None, classes=None, input_feature_names: Optional[Sequence[str]] = None, **fit_params): # infer number of input features if appropriate: if self.module_input_feature_names_ is None: self.module_input_feature_names_ = self._infer_input_feature_names( X, input_feature_names) # convert y to the right dtype: y = to_tensor(y, device=self.device, dtype=self.module_dtype_) return super().partial_fit(X=X, y=y, classes=classes, **fit_params)
def infer(self, x: Union[torch.Tensor, SliceDict], **fit_params): x = to_tensor(x, device=self.device, dtype=self.module_dtype_) return super().infer(x=x, **fit_params)
def km_summary(self, dataframe: 'DataFrame', preprocessor: Union[ColumnTransformer, Sequence], time_colname: str, censor_colname: str, start_time_colname: Optional[str] = None) -> 'DataFrame': """ :param dataframe: A pandas DataFrame, or a DataFrameGroupBy (i.e., the result of calling `df.groupby([...])`). :param preprocessor: Either a sklearn ColumnTransformer that takes the dataframe and returns X, or a list of column-names (such that `X = dataframe.loc[:,preprocessor].values`) :param time_colname: The column-name in the dataframe for time-to-event. :param censor_colname: The column-name in the dataframe for the censoring indicator. :param start_time_colname: Optional, the column-name in the dataframe for start-times (for left-truncation). :return: A DataFrame with kaplan-meier estimates. """ try: from pandas.core.groupby.generic import DataFrameGroupBy except ImportError as e: raise ImportError("Must install pandas for `km_summary`") from e if isinstance(dataframe, DataFrameGroupBy): df_applied = dataframe.apply(self.km_summary, preprocessor=preprocessor, time_colname=time_colname, censor_colname=censor_colname, start_time_colname=start_time_colname) index_idx = [i for i, _ in enumerate(df_applied.index.names)] return df_applied.reset_index(level=index_idx[:-1], drop=False) else: # preprocess X: if hasattr(preprocessor, 'transform'): X = preprocessor.transform(dataframe) else: X = dataframe.loc[:, preprocessor].values X = to_tensor(X, device=self.device, dtype=self.module_dtype_) # km estimate: df_km = km_summary(time=dataframe[time_colname].values, is_upper_cens=dataframe[censor_colname], lower_trunc=dataframe[start_time_colname] if start_time_colname else None) # generate predicted params, transpose as inputs to distribution: with torch.no_grad(): y_preds = self.infer(X) kwargs = { k: y_true[None, :] for k, y_true in zip(self.distribution_param_names_, y_preds) } distribution = self.distribution(**kwargs) # get unique times in distribution-friendly format: y = df_km.loc[:, ['time']].values if self.scale_y: y = self.y_scaler_.transform(y) y = to_tensor(y, device=self.device, dtype=self.module_dtype_) # b/c dist-kwargs transposed, broadcasting logic means we get array with dims: (times, dataframe_rows) observed = y[:, [0]] surv = 1. - distribution.cdf(observed) if start_time_colname: # TODO: either figure out if taking the average is valid, or emit a warning min_ltrunc = np.full_like( observed, fill_value=dataframe[start_time_colname].min()) if self.scale_y: min_ltrunc = self.y_scaler_.transform(min_ltrunc) min_ltrunc = to_tensor(min_ltrunc, device=self.device, dtype=self.module_dtype_) surv /= (1. - distribution.cdf(min_ltrunc)) # this is then reduced, collapsing across dataframe rows, so that we get a mean estimate for this dataset df_km['model_estimate'] = torch.mean(surv, dim=1) return df_km