def predict( self, data: MRData = None, slope_quantile: Dict[str, float] = None, ref_cov: Tuple[str, Any] = None, ): if data is None: data = self.data1 data._sort_by_data_id() pred1 = self.model1.predict(data, slope_quantile=slope_quantile, ref_cov=ref_cov) return self.model2.predict(data) + pred1
def test_data_id(df, study_id): if study_id is not None: df['study_id'] = study_id col_study_id = 'study_id' else: col_study_id = None d = MRData() d.load_df(df, col_obs='obs', col_obs_se='obs_se', col_covs=['cov0', 'cov1', 'cov2'], col_study_id=col_study_id) d._sort_by_data_id() assert np.allclose(d.obs, df['obs']) assert np.allclose(d.obs_se, df['obs_se']) for i in range(3): assert np.allclose(d.covs[f'cov{i}'], df[f'cov{i}'])