def get_treatment_costs()

in causalml/optimize/utils.py [0:0]

• 35 lines of code
• 3 McCabe index (conditional complexity)

def get_treatment_costs(treatment, control_name, cc_dict, ic_dict):
    '''
    Set the conversion and impression costs based on a dict of parameters.

    Calculate the actual cost of targeting a user with the actual treatment
    group using the above parameters.

    Params
    ------
    treatment : array, shape = (num_samples, )
        Treatment array.

    control_name, str
        Control group name as string.

    cc_dict : dict
        Dict containing the conversion cost for each treatment.

    ic_dict
        Dict containing the impression cost for each treatment.

    Returns
    -------
    conversion_cost : ndarray, shape = (num_samples, num_treatments)
        An array of conversion costs for each treatment.

    impression_cost : ndarray, shape = (num_samples, num_treatments)
        An array of impression costs for each treatment.

    conditions : list, len = len(set(treatment))
        A list of experimental conditions.
    '''

    # Set the conversion costs of the treatments
    conversion_cost = np.zeros((len(treatment), len(cc_dict.keys())))
    for idx, dict_key in enumerate(cc_dict.keys()):
        conversion_cost[:, idx] = cc_dict.get(dict_key)

    # Set the impression costs of the treatments
    impression_cost = np.zeros((len(treatment), len(ic_dict.keys())))
    for idx, dict_key in enumerate(ic_dict.keys()):
        impression_cost[:, idx] = ic_dict.get(dict_key)

    # Get a sorted list of conditions
    conditions = list(set(treatment))
    conditions.remove(control_name)
    conditions_sorted = sorted(conditions)
    conditions_sorted.insert(0, control_name)

    return conversion_cost, impression_cost, conditions_sorted