def get_dataframe_from_replay()

in aepsych/server.py [0:0]

• 42 lines of code
• 12 McCabe index (conditional complexity)

    def get_dataframe_from_replay(self, uuid_of_replay=None):
        if uuid_of_replay is None:
            records = self.db.get_master_records()
            if len(records) > 0:
                uuid_of_replay = records[-1].experiment_id
            else:
                raise RuntimeError("Server has no experiment records!")

        recs = self.db.get_replay_for(uuid_of_replay)

        strats = self.get_strats_from_replay(uuid_of_replay)

        out = pd.DataFrame(
            [
                self._flatten_tell_record(rec)
                for rec in recs
                if rec.message_type == "tell"
            ]
        )

        # flatten any final nested lists
        def _flatten(x):
            return x[0] if len(x) == 1 else x

        for col in out.columns:
            if out[col].dtype == object:
                out.loc[:, col] = out[col].apply(_flatten)

        n_tell_records = len(out)
        n_strat_datapoints = 0
        post_means = []
        post_vars = []

        # collect posterior means and vars
        for strat in strats:
            if strat.has_model:
                post_mean, post_var = strat.predict(strat.x)
                n_tell_records = len(out)
                n_strat_datapoints += len(post_mean)
                post_means.extend(post_mean.detach().numpy())
                post_vars.extend(post_var.detach().numpy())

        if n_tell_records == n_strat_datapoints:
            out["post_mean"] = post_means
            out["post_var"] = post_vars
        else:
            logger.warn(
                f"Number of tell records ({n_tell_records}) does not match "
                + f"number of datapoints in strat ({n_strat_datapoints}) "
                + "cowardly refusing to populate GP mean and var to dataframe!"
            )
        return out