def _clean_treatment_costs()

in spotify_confidence/samplesize/sample_size_calculator.py [0:0]

• 18 lines of code
• 7 McCabe index (conditional complexity)

    def _clean_treatment_costs(treatments, treatment_costs):
        """Validate or generate treatment cost array.

        Args:
            treatment_costs (numpy.ndarray, None): Array with same length as
                the number of treatments containing positive floats specifying
                the treatments' relative costs. None also accepted in which
                case equal relative costs are returned.
            treatments (int): Number of treatment variants in the a/b test,
                including control.

        Returns:
            numpy.ndarray: Array with each treatment's cost.

        Raises:
            TypeError: If `treatment_costs` is not None or a numpy.ndarray.
            TypeError: If the length of customs `treatment_costs` is not the
                same as the number of treatments.
            ValueError: If the values of custom `treatment_costs` are not all
                positive and sum to one.

        """
        treatments = SampleSize._clean_treatments(treatments)

        if treatment_costs is None:
            # Default equal cost of all cells
            return np.ones(treatments)

        elif (
            not (isinstance(treatment_costs, np.ndarray) or isinstance(treatment_costs, list))
            or len(treatment_costs) != treatments
        ):
            raise TypeError(
                "treatment_costs must be a list or numpy array of" "the same length as the number of treatments"
            )

        try:
            treatment_costs = np.array(treatment_costs)
            if not (treatment_costs > 0).all():
                raise ValueError("treatment_costs values must all be positive")

        except TypeError:
            raise TypeError("treatment_costs array must only contain numbers")

        return treatment_costs