{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f7cef4c0",
   "metadata": {},
   "source": [
    "# A short demo of Tumult Analytics\n",
    "\n",
    "This notebook shows what the best-in-class differential privacy framework looks like :-)\n",
    "\n",
    "First, we import what we'll need for this demo, and load some data in Spark."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df4d8c47",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from math import floor\n",
    "from pyspark import SparkFiles\n",
    "from pyspark.sql import SparkSession\n",
    "from tmlt.analytics.binning_spec import BinningSpec\n",
    "from tmlt.analytics.constraints import MaxGroupsPerID, MaxRowsPerGroupPerID, MaxRowsPerID\n",
    "from tmlt.analytics.keyset import KeySet\n",
    "from tmlt.analytics.privacy_budget import PureDPBudget, RhoZCDPBudget\n",
    "from tmlt.analytics.protected_change import AddOneRow, AddRowsWithID\n",
    "from tmlt.analytics.query_builder import QueryBuilder\n",
    "from tmlt.analytics.session import Session, ColumnType\n",
    "\n",
    "# silence some Spark warnings\n",
    "import warnings\n",
    "warnings.simplefilter('ignore', UserWarning)\n",
    "warnings.simplefilter('ignore', ResourceWarning)\n",
    "warnings.simplefilter('ignore', FutureWarning)\n",
    "spark = SparkSession.builder.config('spark.driver.memory', '4g').getOrCreate()\n",
    "spark.sparkContext.setLogLevel(\"ERROR\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1925900",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download the data\n",
    "spark.sparkContext.addFile(\n",
    "    \"https://tumult-public.s3.amazonaws.com/library-members.csv\"\n",
    ")\n",
    "members_df = spark.read.csv(\n",
    "    SparkFiles.get(\"library-members.csv\"), header=True, inferSchema=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b581213a",
   "metadata": {},
   "source": [
    "This dataset lists the members of a fictional public library. Let's take a look."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12781426",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(members_df.limit(10).toPandas())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30d1a5b9",
   "metadata": {},
   "source": [
    "## How many members are in our public library?\n",
    "\n",
    "First, we initialize a Tumult Analytics **Session**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdf6adf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "session = Session.from_dataframe(\n",
    "    source_id=\"members\",\n",
    "    dataframe=members_df,                   # our data\n",
    "    protected_change=AddOneRow(),           # what we protect in our data\n",
    "    privacy_budget=PureDPBudget(epsilon=1), # what privacy budget we protect it with\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e56c1d9c",
   "metadata": {},
   "source": [
    "Second, we compute the simplest statistic there is: how many people are in our data?\n",
    "\n",
    "We do this in two steps:\n",
    "- **building** the query;\n",
    "- and **evaluating** our query on our Session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c95e34dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_count_query = QueryBuilder(\"members\").count()\n",
    "\n",
    "total_count = session.evaluate(total_count_query, PureDPBudget(0.4))\n",
    "total_count.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "543dccf7-bc39-43f8-b3a8-5d56226a1251",
   "metadata": {},
   "source": [
    "That was easy! How much budget do we have left?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40d86ff4",
   "metadata": {},
   "outputs": [],
   "source": [
    "session.remaining_privacy_budget"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7762ced",
   "metadata": {},
   "source": [
    "## Which authors are most popular?\n",
    "\n",
    "Now for something a little more complicated: finding which authors are most popular.\n",
    "\n",
    "Let's download two more datasets: \n",
    "- `checkouts_df` lists the books borrowed by library members.\n",
    "- `books_df` lists all the books that the library holds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27d0813d",
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.sparkContext.addFile(\n",
    "    \"https://tumult-public.s3.amazonaws.com/checkout-logs.csv\"\n",
    ")\n",
    "spark.sparkContext.addFile(\n",
    "    \"https://tumult-public.s3.amazonaws.com/library_books.csv\"\n",
    ")\n",
    "checkouts_df = spark.read.csv(\n",
    "    SparkFiles.get(\"checkout-logs.csv\"), header=True, inferSchema=True\n",
    ")\n",
    "books_df = spark.read.csv(\n",
    "    SparkFiles.get(\"library_books.csv\"), header=True, inferSchema=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67d78675",
   "metadata": {},
   "source": [
    "Let's take a look at the checkouts dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "027cf80b",
   "metadata": {},
   "outputs": [],
   "source": [
    "display(checkouts_df.limit(5).toPandas())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e389bf50-5c58-49be-8484-4d468cf64d65",
   "metadata": {},
   "source": [
    "Let's load it in a Session, using **privacy IDs** to protect all the contributions of each library member."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e2cbbcd-fda8-4f2e-ab23-ca775829c49e",
   "metadata": {},
   "outputs": [],
   "source": [
    "session = Session.from_dataframe(\n",
    "    source_id=\"checkouts\",\n",
    "    dataframe=checkouts_df,                   \n",
    "    protected_change=AddRowsWithID(\"member_id\"), # protecting all the contributions of any single member\n",
    "    privacy_budget=PureDPBudget(epsilon=1),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "621f1638-8b1f-4be4-aade-4c86c8e6a793",
   "metadata": {},
   "source": [
    "And we can now see who the most popular authors are, using a group-by query with **public group-by keys**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36436a69-71d3-4fd9-8b54-fd5e93fd172a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# public group-by keys\n",
    "keyset = KeySet.from_dataframe(\n",
    "    books_df.select(\"author\"),\n",
    ")\n",
    "count_query = (\n",
    "    QueryBuilder(\"checkouts\")\n",
    "    .enforce(MaxRowsPerID(20))\n",
    "    .groupby(keyset)\n",
    "    .count()\n",
    ")\n",
    "author_counts = session.evaluate(count_query, PureDPBudget(1))\n",
    "top_five_authors = author_counts.sort(\"count\", ascending=False).limit(5)\n",
    "display(top_five_authors.toPandas())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0b6b1c0",
   "metadata": {},
   "source": [
    "## What is the median age of readers of each literary genre?\n",
    "\n",
    "The `checkouts_df` table has data about book checkouts, while `members_df` has demographic information about members. \n",
    "\n",
    "To find correlations between literary genres and the age of people who read books in this genre, we need to load **multiple tables** in our Session. We will also use **zero-concentrated differential privacy** as a privacy notion, to get a better privacy/accuracy trade-off."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43bf4508",
   "metadata": {},
   "outputs": [],
   "source": [
    "session_budget = RhoZCDPBudget(rho=0.5) # zero-concentrated differential privacy budget\n",
    "id_space = \"members_id_space\"           # indicating that both tables share an identifier\n",
    "session = (\n",
    "    Session.Builder()\n",
    "    .with_privacy_budget(session_budget)\n",
    "    .with_id_space(id_space)\n",
    "    .with_private_dataframe(\n",
    "        \"checkouts\",\n",
    "        checkouts_df,\n",
    "        protected_change=AddRowsWithID( # protecting arbitrarily many rows with the same identifier…\n",
    "            id_column=\"member_id\",\n",
    "            id_space=id_space,\n",
    "        ),\n",
    "    )\n",
    "    .with_private_dataframe(\n",
    "        \"members\",\n",
    "        members_df,\n",
    "        protected_change=AddRowsWithID( # … in both tables.\n",
    "            id_column=\"id\",\n",
    "            id_space=id_space,\n",
    "        ),\n",
    "    )\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "721c9dfb",
   "metadata": {},
   "source": [
    "Each book can have *multiple* genres: we will expand the `\"checkout\"` table to put each genre on its own row, using a **flat map**, and creating a **view** with the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "139a9e44",
   "metadata": {},
   "outputs": [],
   "source": [
    "expand_genres = (\n",
    "    QueryBuilder(\"checkouts\").flat_map(\n",
    "        lambda row: [{\"genre\": genre} for genre in row[\"genres\"].split(\",\")],\n",
    "        {\"genre\": ColumnType.VARCHAR},\n",
    "        augment=True,\n",
    "    )\n",
    ")\n",
    "\n",
    "session.create_view(expand_genres, \"checkouts_single_genre\", cache=True)\n",
    "print(f\"Private dataframes: {session.private_sources}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca38bfcc",
   "metadata": {},
   "source": [
    "Then, we **join** the table we just created with the `\"members\"` table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df23ceaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "private_join = (\n",
    "    QueryBuilder(\"members\")\n",
    "    .rename({\"id\": \"member_id\"})\n",
    "    .join_private(\"checkouts_single_genre\")\n",
    ")\n",
    "session.create_view(private_join, \"checkouts_joined\", cache=True)\n",
    "print(f\"Private dataframes: {session.private_sources}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1a7a958-a845-4647-9c9f-8ee4cf45f54d",
   "metadata": {},
   "source": [
    "Finally, we can select a few genres, and compute the **median** age of readers of each genre."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9adacf69-265d-4a2c-8ead-45a02ec9d304",
   "metadata": {},
   "outputs": [],
   "source": [
    "genres = [\n",
    "    \"Mystery/thriller/crime\",\n",
    "    \"History\",\n",
    "    \"Romance\",\n",
    "    \"Fantasy\",\n",
    "    \"Classics/Literature\",\n",
    "    \"Children\",\n",
    "]\n",
    "median_age_by_genre_query = (\n",
    "    QueryBuilder(\"checkouts_joined\")\n",
    "    .enforce(MaxRowsPerID(10))\n",
    "    .groupby(KeySet.from_dict({\"genre\": genres}))\n",
    "    .median(\"age\", low=0, high=100)\n",
    ")\n",
    "median_age_by_genre = session.evaluate(median_age_by_genre_query, RhoZCDPBudget(rho=0.1)).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e212cb80",
   "metadata": {},
   "outputs": [],
   "source": [
    "ax = sns.barplot(\n",
    "    x=\"genre\",\n",
    "    y=\"age_median\",\n",
    "    data=median_age_by_genre.sort_values([\"age_median\"]),\n",
    "    color=\"#1f77b4\"\n",
    ")\n",
    "ax.set(xlabel=\"Genre\", ylabel=\"Median age\", title=\"Median age by genre\")\n",
    "ax.set_xticklabels(\n",
    "    median_age_by_genre[\"genre\"], rotation=45, horizontalalignment=\"right\"\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d21eade7-68b9-4547-a24b-a147f72763be",
   "metadata": {},
   "source": [
    "## How does reading behavior vary with age?\n",
    "\n",
    "The information about median age is somewhat coarse. Let's get a clearer picture, by looking at what genres people read most depending on their age.\n",
    "\n",
    "To do so, we will first **bin** our data into age ranges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff175986",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define age groups, binning edges at [0, 20, 40, ...,100]\n",
    "age_binspec = BinningSpec(bin_edges = [20*i for i in range(0, 6)])\n",
    "binned_age_genre_keys = KeySet.from_dict({\n",
    "    \"binned_age\": age_binspec.bins(),\n",
    "    \"genre\": genres,\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27430881-c2a3-4024-a444-551d237d07ff",
   "metadata": {},
   "source": [
    "And we can now run our query, using **advanced truncation** to optimize the differential privacy mechanism."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a460ca14",
   "metadata": {},
   "outputs": [],
   "source": [
    "genre_by_age_query = (\n",
    "    QueryBuilder(\"checkouts_joined\")\n",
    "    .bin_column(\"age\", age_binspec, name=\"binned_age\") # bin the data\n",
    "    .enforce(MaxGroupsPerID(\"binned_age\", 1))          # each person has a single age\n",
    "    .enforce(MaxGroupsPerID(\"genre\", 6))               # we are only considering 6 genres in our analysis\n",
    "    .enforce(MaxRowsPerGroupPerID(\"genre\", 1))         # we only want to count each person once per genre…\n",
    "    .enforce(MaxRowsPerGroupPerID(\"binned_age\", 6))    # … so each person will appear in at most 6 rows \n",
    "    .groupby(binned_age_genre_keys)\n",
    "    .count()\n",
    ")\n",
    "genre_by_age = session.evaluate(genre_by_age_query, RhoZCDPBudget(rho=0.4)).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea8ba9a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert binned_age to categorical\n",
    "genre_by_age[\"binned_age\"] = pd.Categorical(genre_by_age[\"binned_age\"], age_binspec.bins())\n",
    "age_counts = (\n",
    "    genre_by_age.groupby(\"binned_age\").sum().rename(columns={\"count\": \"age_count\"})\n",
    ")\n",
    "# compute percentage of each genre in each age group, replace negative values with 0\n",
    "genre_by_age_pct = genre_by_age.join(age_counts, on=\"binned_age\")\n",
    "genre_by_age_pct[\"pct\"] = genre_by_age_pct[\"count\"] / genre_by_age_pct[\"age_count\"] * 100\n",
    "genre_by_age_pct[\"pct\"] = genre_by_age_pct[\"pct\"].clip(lower=0)\n",
    "\n",
    "ax = sns.barplot(\n",
    "    x=\"binned_age\",\n",
    "    y=\"pct\",\n",
    "    order=age_binspec.bins(),\n",
    "    hue=\"genre\",\n",
    "    data=genre_by_age_pct,\n",
    ")\n",
    "ax.set(xlabel=\"Age Group\", ylabel=\"Genre (%)\", title=\"Reading Preferences by Age\")\n",
    "sns.move_legend(ax, \"upper left\", bbox_to_anchor=(1, 1), ncol=1, title=\"Genre\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
