def fetch_trial_data()

in ax/metrics/factorial.py [0:0]

• 29 lines of code
• 4 McCabe index (conditional complexity)

    def fetch_trial_data(self, trial: BaseTrial, **kwargs: Any) -> Data:
        if not isinstance(trial, BatchTrial):
            raise ValueError("Factorial metric can only fetch data for batch trials.")
        if not trial.status.expecting_data:
            raise ValueError("Can only fetch data if trial is expecting data.")

        data = []
        normalized_arm_weights = trial.normalized_arm_weights()
        for name, arm in trial.arms_by_name.items():
            weight = normalized_arm_weights[arm]
            mean, sem = evaluation_function(
                parameterization=arm.parameters,
                weight=weight,
                coefficients=self.coefficients,
                batch_size=self.batch_size,
                noise_var=self.noise_var,
            )
            n = np.random.binomial(self.batch_size, weight)
            data.append(
                {
                    "arm_name": name,
                    "metric_name": self.name,
                    "mean": mean,
                    "sem": sem,
                    "trial_index": trial.index,
                    "n": n,
                    "frac_nonnull": mean,
                }
            )
        return Data(df=pd.DataFrame(data))