text: import numpy as np
  import pandas as pd
  
  from sklearn.base import clone
  
  import warnings
  
  
  class CounterfactualUnitSelector:
      '''
      A highly experimental implementation of the counterfactual unit selection
      model proposed by Li and Pearl (2019).
  
      Parameters
      ----------
      learner : object
          The base learner used to estimate the segment probabilities.
  
      nevertaker_payoff : float
          The payoff from targeting a never-taker
  
      alwaystaker_payoff : float
          The payoff from targeting an always-taker
  
      complier_payoff : float
          The payoff from targeting a complier
  
      defier_payoff : float
          The payoff from targeting a defier
  
      organic_conversion : float, optional (default=None)
          The organic conversion rate in the population without an intervention.
          If None, the organic conversion rate is obtained from tne control group.
  
          NB: The organic conversion in the control group is not always the same
          as the organic conversion rate without treatment.
  
      data : DataFrame
          A pandas DataFrame containing the features, treatment assignment
          indicator and the outcome of interest.
  
      treatment : string
          A string corresponding to the name of the treatment column. The
          assumed coding in the column is 1 for treatment and 0 for control.
  
      outcome : string
          A string corresponding to the name of the outcome column. The assumed
          coding in the column is 1 for conversion and 0 for no conversion.
  
      References
      ----------
      Li, Ang, and Judea Pearl. 2019. “Unit Selection Based on Counterfactual
      Logic.” https://ftp.cs.ucla.edu/pub/stat_ser/r488.pdf.
      '''
  
      def __init__(self, learner, nevertaker_payoff, alwaystaker_payoff,
                   complier_payoff, defier_payoff, organic_conversion=None):
  
          self.learner = learner
          self.nevertaker_payoff = nevertaker_payoff
          self.alwaystaker_payoff = alwaystaker_payoff
          self.complier_payoff = complier_payoff
          self.defier_payoff = defier_payoff
          self.organic_conversion = organic_conversion
  
      def fit(self, data, treatment, outcome):
          '''
          Fits the class.
          '''
  
          if self._gain_equality_check():
  
              self._fit_segment_model(data, treatment, outcome)
  
          else:
  
              self._fit_segment_model(data, treatment, outcome)
              self._fit_condprob_models(data, treatment, outcome)
  
      def predict(self, data, treatment, outcome):
          '''
          Predicts an individual-level payoff. If gain equality is satisfied, uses
          the exact function; if not, uses the midpoint between bounds.
          '''
  
          if self._gain_equality_check():
  
              est_payoff = self._get_exact_benefit(data, treatment, outcome)
  
          else:
  
              est_payoff = self._obj_func_midp(data, treatment, outcome)
  
          return est_payoff
  
      def _gain_equality_check(self):
          '''
          Checks if gain equality is satisfied. If so, the optimization task can
          be simplified.
          '''
  
          return self.complier_payoff + self.defier_payoff == \
              self.alwaystaker_payoff + self.nevertaker_payoff
  
      @staticmethod
      def _make_segments(data, treatment, outcome):
          '''
          Constructs the following segments:
  
          * AC = Pr(Y = 1, W = 1 /mid X)
          * AD = Pr(Y = 1, W = 0 /mid X)
          * ND = Pr(Y = 0, W = 1 /mid X)
          * ND = Pr(Y = 0, W = 0 /mid X)
  
          where the names of the outcomes correspond the combinations of
          the relevant segments, eg AC = Always-taker or Complier.
          '''
  
          segments = np.empty(data.shape[0], dtype='object')
  
          segments[(data[treatment] == 1) & (data[outcome] == 1)] = 'AC'
          segments[(data[treatment] == 0) & (data[outcome] == 1)] = 'AD'
          segments[(data[treatment] == 1) & (data[outcome] == 0)] = 'ND'
          segments[(data[treatment] == 0) & (data[outcome] == 0)] = 'NC'
  
          return segments
  
      def _fit_segment_model(self, data, treatment, outcome):
          '''
          Fits a classifier for estimating the probabilities for the unit
          segment combinations.        
          '''
  
          model = clone(self.learner)
  
          X = data.drop([treatment, outcome], axis=1)
          y = self._make_segments(data, treatment, outcome)
  
          self.segment_model = model.fit(X, y)
  
      def _fit_condprob_models(self, data, treatment, outcome):
          '''
          Fits two classifiers to estimate conversion probabilities conditional
          on the treatment.
          '''
  
          trt_learner = clone(self.learner)
          ctr_learner = clone(self.learner)
  
          treated = data[treatment] == 1
  
          X = data.drop([treatment, outcome], axis=1)
          y = data['outcome']
  
          self.trt_model = trt_learner.fit(X[treated], y[treated])
          self.ctr_model = ctr_learner.fit(X[~treated], y[~treated])
  
      def _get_exact_benefit(self, data, treatment, outcome):
          '''
          Calculates the exact benefit function of Theorem 4 in Li and Pearl (2019).
          Returns the exact benefit.
          '''
          beta = self.complier_payoff
          gamma = self.alwaystaker_payoff
          theta = self.nevertaker_payoff
  
          X = data.drop([treatment, outcome], axis=1)
  
          segment_prob = self.segment_model.predict_proba(X)
          segment_name = self.segment_model.classes_
  
          benefit = (beta - theta) * segment_prob[:, segment_name == 'AC'] + \
              (gamma - beta) * segment_prob[:, segment_name == 'AD'] + theta
  
          return benefit
  
      def _obj_func_midp(self, data, treatment, outcome):
          '''
          Calculates bounds for the objective function. Returns the midpoint
          between bounds.
  
          Parameters
          ----------
          pr_y1_w1 : float
              The probability of conversion given treatment assignment.
  
          pr_y1_w0 : float
              The probability of conversion given control assignment.
  
          pr_y0_w1 : float
              The probability of no conversion fiven treatment assignment
              (1 - pr_y1_w1).
  
          pr_y0_w0 : float
              The probability of no conversion given control assignment
              (1 - pr_1y_w0)
  
          pr_y1w1_x : float
              Probability of complier or always-taker given X.
  
          pr_y0w0_x : float
              Probability of complier or never-taker given X.
  
          pr_y1w0_x : float
              Probability of defier or always-taker given X.
  
          pr_y0w1_x : float
              Probability of never-taker or defier given X.
  
          pr_y_x : float
              Organic probability of conversion.
          '''
  
          X = data.drop([treatment, outcome], axis=1)
  
          beta = self.complier_payoff
          gamma = self.alwaystaker_payoff
          theta = self.nevertaker_payoff
          delta = self.defier_payoff
  
          pr_y0_w1, pr_y1_w1 = np.split(self.trt_model.predict_proba(X),
                                        indices_or_sections=2, axis=1)
          pr_y0_w0, pr_y1_w0 = np.split(self.ctr_model.predict_proba(X),
                                        indices_or_sections=2, axis=1)
  
          segment_prob = self.segment_model.predict_proba(X)
          segment_name = self.segment_model.classes_
  
          pr_y1w1_x = segment_prob[:, segment_name == 'AC']
          pr_y0w0_x = segment_prob[:, segment_name == 'NC']
          pr_y1w0_x = segment_prob[:, segment_name == 'AD']
          pr_y0w1_x = segment_prob[:, segment_name == 'ND']
  
          if self.organic_conversion is not None:
  
              pr_y_x = self.organic_conversion
  
          else:
  
              pr_y_x = pr_y1_w0
              warnings.warn(
                  'Probability of organic conversion estimated from control observations.')
  
          p1 = (beta - theta) * pr_y1_w1 + delta * pr_y1_w0 + theta * pr_y0_w0
          p2 = gamma * pr_y1_w1 + delta * pr_y0_w1 + (beta - gamma) * pr_y0_w0
          p3 = (gamma - delta) * pr_y1_w1 + delta * pr_y1_w0 + theta * \
              pr_y0_w0 + (beta - gamma - theta + delta) * (pr_y1w1_x + pr_y0w0_x)
          p4 = (beta - theta) * pr_y1_w1 - (beta - gamma - theta) * pr_y1_w0 + \
              theta * pr_y0_w0 + (beta - gamma - theta + delta) * \
              (pr_y1w0_x + pr_y0w1_x)
          p5 = (gamma - delta) * pr_y1_w1 + delta * pr_y1_w0 + theta * pr_y0_w0
          p6 = (beta - theta) * pr_y1_w1 - (beta - gamma - theta) * pr_y1_w0 + \
              theta * pr_y0_w0
          p7 = (gamma - delta) * pr_y1_w1 - (beta - gamma - theta) * pr_y1_w0 + \
              theta * pr_y0_w0 + (beta - gamma - theta + delta) * pr_y_x
          p8 = (beta - theta) * pr_y1_w1 + delta * pr_y1_w0 + theta * \
              pr_y0_w0 - (beta - gamma - theta + delta) * pr_y_x
  
          params_1 = np.concatenate((p1, p2, p3, p4), axis=1)
          params_2 = np.concatenate((p5, p6, p7, p8), axis=1)
  
          sigma = beta - gamma - theta + delta
  
          if sigma < 0:
              lower_bound = np.max(params_1, axis=1)
              upper_bound = np.min(params_2, axis=1)
  
          elif sigma > 0:
              lower_bound = np.max(params_2, axis=1)
              upper_bound = np.min(params_1, axis=1)
  
          return (lower_bound + upper_bound) / 2
  