def arm_weights()

in ax/core/batch_trial.py [0:0]

• 15 lines of code
• 8 McCabe index (conditional complexity)

    def arm_weights(self) -> MutableMapping[Arm, float]:
        """The set of arms and associated weights for the trial.

        These are constructed by merging the arms and weights from
        each generator run that is attached to the trial.
        """
        arm_weights = OrderedDict()
        if len(self._generator_run_structs) == 0 and self.status_quo is None:
            return arm_weights
        for struct in self._generator_run_structs:
            multiplier = struct.weight
            for arm, weight in struct.generator_run.arm_weights.items():
                scaled_weight = weight * multiplier
                if arm in arm_weights:
                    arm_weights[arm] += scaled_weight
                else:
                    arm_weights[arm] = scaled_weight
        if self.status_quo is not None and self._status_quo_weight_override is not None:
            # If override is specified, this is the weight the status quo gets,
            # regardless of whether it appeared in any generator runs.
            # If no override is specified, status quo does not appear in arm_weights.
            arm_weights[self.status_quo] = self._status_quo_weight_override
        return arm_weights