<filename>generalized_rates/datasets/load_compas.py<gh_stars>0
# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Loads COMPAS dataset, processes it and saves it to output folder.

Adapted from code provided by <NAME> (<EMAIL>). The COMPAS
data file can be downloaded from:
https://github.com/propublica/compas-analysis/blob/master/compas-scores-two-years.csv
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl import app
from absl import flags

import numpy as np
import pandas as pd
from sklearn import model_selection


flags.DEFINE_float("test_fraction", 1.0 / 3.0,
                   "Fraction of overall dataset that constitutes the test set.")
flags.DEFINE_float("vali_fraction", 1.0 / 3.0,
                   "Fraction of train set that constitutes the validation set.")
flags.DEFINE_string("data_file", "compas-scores-two-years.csv",
                    "Path to COMPAS dataset csv file.")
flags.DEFINE_string("output_directory", "datasets/",
                    "Path to store processed dataset.")


FLAGS = flags.FLAGS


def load_data():
  """Load and process dataset from provided data path."""
  df = pd.read_csv(FLAGS.data_file)

  # Filter relevant features.
  features = [
      "age", "c_charge_degree", "race", "score_text", "sex", "priors_count",
      "days_b_screening_arrest", "decile_score", "is_recid", "two_year_recid"]
  df = df[features]
  df = df[df.days_b_screening_arrest <= 30]
  df = df[df.days_b_screening_arrest >= -30]
  df = df[df.is_recid != -1]
  df = df[df.c_charge_degree != "O"]
  df = df[df.score_text != "N/A"]

  # Divide features into: continuous, categorical and those that are continuous,
  # but need to be converted to categorical.
  categorical_features = ["c_charge_degree", "race", "score_text", "sex"]
  continuous_to_categorical_features = ["age", "decile_score", "priors_count"]

  # Bucketize features in continuous_to_categorical_features.
  for feature in continuous_to_categorical_features:
    if feature == "priors_count":
      bins = list(np.percentile(df[feature], [0, 50, 70, 80, 90, 100]))
    else:
      bins = [0] + list(np.percentile(df[feature], [20, 40, 60, 80, 90, 100]))
    df[feature] = pd.cut(df[feature], bins, labels=False)

  # Binarize all categorical features (including the ones bucketized above).
  df = pd.get_dummies(df, columns=categorical_features +
                      continuous_to_categorical_features)

  # Fill values for decile scores and prior counts feature buckets.
  to_fill = [u"decile_score_0", u"decile_score_1", u"decile_score_2",
             u"decile_score_3", u"decile_score_4", u"decile_score_5"]
  for i in range(len(to_fill) - 1):
    df[to_fill[i]] = df[to_fill[i:]].max(axis=1)
  to_fill = [u"priors_count_0.0", u"priors_count_1.0", u"priors_count_2.0",
             u"priors_count_3.0", u"priors_count_4.0"]
  for i in range(len(to_fill) - 1):
    df[to_fill[i]] = df[to_fill[i:]].max(axis=1)

  # Get the labels (two year recidivism) and groups (female defendants).
  labels = df["two_year_recid"]
  groups = df["sex_Female"]

  # Retain all features other than "two_year_recid" and "is_recid".
  df.drop(columns=["two_year_recid", "is_recid"], inplace=True)

  return df.to_numpy(), labels.to_numpy(), groups.to_numpy()


def main(argv):
  if len(argv) > 1:
    raise app.UsageError("Too many command-line arguments.")

  # Load and pre-process COMPAS dataset.
  features, labels, groups = load_data()

  # Split dataset indices into train and test.
  train_indices, test_indices = model_selection.train_test_split(
      np.arange(features.shape[0]), test_size=FLAGS.test_fraction)
  # Split train indices further into train and validation,
  train_indices, vali_indices = model_selection.train_test_split(
      train_indices, test_size=FLAGS.vali_fraction)

  # Split features, labels and groups for train, test and validation sets,
  x_train = features[train_indices, :]
  y_train = labels[train_indices]
  z_train = groups[train_indices]

  x_test = features[test_indices, :]
  y_test = labels[test_indices]
  z_test = groups[test_indices]

  x_vali = features[vali_indices, :]
  y_vali = labels[vali_indices]
  z_vali = groups[vali_indices]

  train_set = x_train, y_train, z_train
  vali_set = x_vali, y_vali, z_vali
  test_set = x_test, y_test, z_test

  # Save processed dataset.
  with open(FLAGS.output_directory + "COMPAS.npy", "wb") as f:
    np.save(f, (train_set, vali_set, test_set), allow_pickle=True)

if __name__ == "__main__":
  app.run(main)
