comparisons.py 文件源码

def precalculate_factors(self, data, sort=True):
    """Initializes the factor variable.

    Args:
      data: A pandas dataframe.
      sort: Boolean indicating whether or not the conditions should be sorted.

    Raises:
      ValueError: The baseline key isn't found.
    """

    self.factors, condition_keys = pd.factorize(data[self.condition_column],
                                                sort=sort)

    self.alternate_indices = [
        ii for ii, label in enumerate(condition_keys)
        if self._include_base or label != self.baseline_key
    ]
    self.alternate_keys = condition_keys[self.alternate_indices]

    if any(condition_keys == self.baseline_key):
      self.baseline_index = np.where(condition_keys == self.baseline_key)[0][0]
    else:
      raise ValueError("Baseline value {} not present in column {}".format(
          self.baseline_key, self.condition_column))

    self._baseline_mask = (self.factors == self.baseline_index)
    self._alternate_masks = {}
    for ii in self.alternate_indices:
      self._alternate_masks[ii] = (self.factors == ii)