02_generate_case_study_files.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6f802687-416a-4b38-82d3-ab387dc16c38",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Generate simulated data to link\n",
    "\n",
    "In this case study, we imagine running a person linkage on a simulated 2030 Census Unedited File (CUF) --\n",
    "see the main notebook for more details, including references used throughout this notebook.\n",
    "This notebook creates simulated input (CUF) and reference files approximating what would be used in such a process,\n",
    "emulating the methods found in publicly available descriptions of the Census Bureau's primary person linkage system, PVS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c4407b8-a1d5-4d38-a2a1-820fd078487f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Query planning is now on by default, but it has some rough edges.\n",
    "# See https://github.com/dask/dask/issues/10995 for general discussion\n",
    "# and https://github.com/dask/dask-expr/issues/1060 for the particular\n",
    "# issue I ran into.\n",
    "try:\n",
    "    import dask\n",
    "\n",
    "    dask.config.set({\"dataframe.query-planning\": False})\n",
    "except ImportError:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "737efe65-8d84-4276-b39f-83b2c2cd4cff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, time\n",
    "import warnings\n",
    "import datetime\n",
    "\n",
    "# Importing pandas for access, regardless of whether we are using it as the compute engine\n",
    "import pandas\n",
    "import numpy\n",
    "from IPython.display import display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4b26e57-7bf4-458a-bb14-2da532d4361b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "114fdbee-8c21-40f1-a5b2-4a1d2c2743b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from person_linkage_case_study_utils import distributed_compute, utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "028e3f70-dce9-46ef-9610-722c895e4a78",
   "metadata": {},
   "outputs": [],
   "source": [
    "warnings.simplefilter(action=\"ignore\", category=FutureWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ac996f3-1270-4211-a829-728e3b2ba6f7",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "# DO NOT EDIT if this notebook is not called 02_generate_case_study_files.ipynb!\n",
    "# This notebook is designed to be run with papermill; this cell is tagged 'parameters'\n",
    "data_to_use = \"small_sample\"\n",
    "input_dir = \"output/01_generate_pseudopeople_simulated_datasets\"\n",
    "output_dir = \"output/02_generate_case_study_files\"\n",
    "compute_engine = \"pandas\"\n",
    "compute_engine_num_workers = 10\n",
    "compute_engine_cpus_per_worker = 2\n",
    "compute_engine_threads_per_worker = 1\n",
    "compute_engine_memory_per_worker = \"10GB\"\n",
    "queue = None\n",
    "account = None\n",
    "# NOTE: This is, as Dask requests, a directory local to the compute node.\n",
    "# But IHME's cluster doesn't support this very well -- it can be small-ish,\n",
    "# full of stuff from other users, etc.\n",
    "compute_engine_local_directory = (\n",
    "    f\"/tmp/{os.environ['USER']}_{int(time.time())}_person_linkage_case_study\"\n",
    ")\n",
    "compute_engine_log_directory = f\"{output_dir}/{data_to_use}/logs\"\n",
    "walltime = None\n",
    "compute_engine_memory_constrained = True\n",
    "scheduler = \"slurm\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65653e90-5589-4ae9-917e-b1082cf475b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "if compute_engine != \"pandas\":\n",
    "    utils.ensure_empty(compute_engine_local_directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79ff29f5-473d-4dab-9e7f-453f284ad8d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(datetime.datetime.now())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3bb7031-36f0-4567-842b-3ad15c8dd274",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_ops, pd = distributed_compute.start_compute_engine(\n",
    "    compute_engine,\n",
    "    num_workers=compute_engine_num_workers,\n",
    "    cpus_per_worker=compute_engine_cpus_per_worker,\n",
    "    threads_per_worker=compute_engine_threads_per_worker,\n",
    "    memory_per_worker=compute_engine_memory_per_worker,\n",
    "    worker_walltime=walltime,\n",
    "    local_directory=compute_engine_local_directory,\n",
    "    log_directory=compute_engine_log_directory,\n",
    "    memory_constrained=compute_engine_memory_constrained,\n",
    "    scheduler=scheduler,\n",
    "    queue=queue,\n",
    "    account=account,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d617fe6-f34e-4d10-a3c3-96c7e37cdb8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_file_with_ground_truth(file_name, file, ground_truth):\n",
    "    # Check that file and ground truth have the same records\n",
    "    assert (\n",
    "        # record_id is unique\n",
    "        len(file)\n",
    "        == len(df_ops.drop_duplicates(file[[\"record_id\"]]))\n",
    "    )\n",
    "    assert len(ground_truth) == len(\n",
    "        file[[\"record_id\"]].merge(\n",
    "            ground_truth[[\"record_id\"]], on=\"record_id\", how=\"inner\"\n",
    "        )\n",
    "    )\n",
    "\n",
    "    file_path = f\"{output_dir}/{data_to_use}/{file_name}.parquet\"\n",
    "    utils.remove_path(file_path)\n",
    "    file.to_parquet(file_path)\n",
    "\n",
    "    ground_truth_path = f\"{output_dir}/{data_to_use}/{file_name}_ground_truth.parquet\"\n",
    "    utils.remove_path(ground_truth_path)\n",
    "    ground_truth.to_parquet(ground_truth_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a08bfa0-f644-4faa-b303-ef24d9a32b86",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Load pseudopeople simulated datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e91d26b9-b129-453b-84d0-51d2e366003f",
   "metadata": {},
   "source": [
    "### Record ID tracking (data lineage)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87aab756-5f18-4c62-a014-a4ff88e4f342",
   "metadata": {},
   "source": [
    "We do a little bit of work here to enable tracking the \"ground truth\" (the simulant IDs from\n",
    "pseudopeople).\n",
    "We give each pseudopeople record/row a unique identifier for tracking, and then we immediately\n",
    "separate the ground truth information (which we would not have if we were using real data)\n",
    "from the rest of the columns (which we would have).\n",
    "The ground truth is only used in the specific \"ground truth\" section of this notebook,\n",
    "to help avoid accidentally leaking information into the case study.\n",
    "\n",
    "Since we also combine/aggregate pseudopeople records as part of the process of generating the\n",
    "simulated reference files, ground truth is a bit more complicated than you might imagine.\n",
    "For example, the ground truth may tell us that a single row in a reference file is actually\n",
    "a composite of several individuals, because even the deterministic linkage (by SSN) we use\n",
    "here is not without error.\n",
    "\n",
    "We handle this by tracking *all* source records used in the construction of each record in\n",
    "our reference files.\n",
    "This is achieved by having a table mapping composite record IDs to the \"source record IDs\"\n",
    "(IDs of records that were directly generated by pseudopeople).\n",
    "When we aggregate records, this is combined accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "090283e4-4e2f-453c-8896-154689ed5d8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_unique_id_col(df, col_name=\"unique_id\", value_prefix=\"\"):\n",
    "    if compute_engine == \"pandas\" or compute_engine.startswith(\"modin\"):\n",
    "        return (\n",
    "            df.reset_index()\n",
    "            .rename(columns={\"index\": col_name})\n",
    "            .assign(**{col_name: lambda df: value_prefix + df[col_name].astype(str)})\n",
    "        )\n",
    "    elif compute_engine.startswith(\"dask\"):\n",
    "        # Can use cumsum as in https://stackoverflow.com/a/60852409/ if it needs\n",
    "        # to be incrementing, but we just need uniqueness\n",
    "\n",
    "        def add_id_to_partition(df_part, partition_info=None):\n",
    "            return df_part.assign(**{col_name: range(len(df_part))}).assign(\n",
    "                **{\n",
    "                    col_name: lambda x: (\n",
    "                        value_prefix\n",
    "                        + str(\n",
    "                            partition_info[\"number\"]\n",
    "                            if partition_info is not None\n",
    "                            else 0\n",
    "                        )\n",
    "                        + \"_\"\n",
    "                        + x[col_name].astype(str)\n",
    "                    ).astype(\"large_string[pyarrow]\")\n",
    "                }\n",
    "            )\n",
    "\n",
    "        df = df.map_partitions(add_id_to_partition)\n",
    "\n",
    "        return df_ops.ensure_large_string_capacity(df)\n",
    "    else:\n",
    "        raise ValueError()\n",
    "\n",
    "\n",
    "def add_unique_record_id(df, dataset_name):\n",
    "    return add_unique_id_col(df, col_name=\"record_id\", value_prefix=f\"{dataset_name}_\")\n",
    "\n",
    "\n",
    "# Initializes a table listing the pairs between record_ids and source record_ids.\n",
    "# Should only be called on \"source records\"; that is, records that\n",
    "# come directly out of pseudopeople.\n",
    "def record_id_to_single_source_record_pairs(df, source_col=\"record_id\"):\n",
    "    if source_col == \"record_id\":\n",
    "        # We can't have duplicate column names, so we make a new column\n",
    "        # literally called 'source_col'\n",
    "        df = df.assign(source_col=lambda df: df[source_col])\n",
    "        source_col = \"source_col\"\n",
    "\n",
    "    return df[[\"record_id\", source_col]].rename(\n",
    "        columns={source_col: \"source_record_id\"}\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cb6694d-1ee4-4826-9064-e8354e00bded",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Operations that aggregate records, combining the source record pairs\n",
    "# between all records that are aggregated\n",
    "\n",
    "\n",
    "def merge_preserving_source_records(\n",
    "    dfs, source_record_pairings, new_record_id_prefix, *args, **kwargs\n",
    "):\n",
    "    assert len(dfs) == len(source_record_pairings)\n",
    "    for df in dfs:\n",
    "        assert \"record_id\" in df.columns\n",
    "\n",
    "    on = kwargs.get(\"on\", None)\n",
    "    if on is not None:\n",
    "        # If there are nulls in any of the merge columns, they can't match to anything\n",
    "        dfs = [df.dropna(subset=on, how=\"any\") for df in dfs]\n",
    "\n",
    "    result = dfs[0]\n",
    "    source_record_pairs = source_record_pairings[0]\n",
    "    dfs_and_source_record_pairs_to_combine = list(\n",
    "        zip(dfs[1:], source_record_pairings[1:])\n",
    "    )\n",
    "    for index, (df_to_merge, source_record_pairs_to_merge) in enumerate(\n",
    "        dfs_and_source_record_pairs_to_combine\n",
    "    ):\n",
    "        result = result.merge(df_to_merge, *args, **kwargs)\n",
    "        if index == len(dfs_and_source_record_pairs_to_combine) - 1:\n",
    "            # Since this is the last step, these are the record_ids that will actually be returned\n",
    "            accumulate_step_record_id_prefix = new_record_id_prefix\n",
    "        else:\n",
    "            # A dummy intermediate -- this shouldn't be exposed to the user\n",
    "            accumulate_step_record_id_prefix = f\"merge_iter_{index}\"\n",
    "\n",
    "        result = df_ops.rebalance(\n",
    "            add_unique_record_id(result, accumulate_step_record_id_prefix)\n",
    "        )\n",
    "        source_record_pairs = df_ops.concat(\n",
    "            [\n",
    "                # The pairs that were already in result\n",
    "                source_record_pairs.rename(columns={\"record_id\": \"record_id_x\"})\n",
    "                .merge(\n",
    "                    result[[\"record_id\", \"record_id_x\"]].dropna(how=\"any\"),\n",
    "                    on=\"record_id_x\",\n",
    "                )\n",
    "                .drop(columns=[\"record_id_x\"]),\n",
    "                # The new ones\n",
    "                source_record_pairs_to_merge.rename(\n",
    "                    columns={\"record_id\": \"record_id_y\"}\n",
    "                )\n",
    "                .merge(\n",
    "                    result[[\"record_id\", \"record_id_y\"]].dropna(how=\"any\"),\n",
    "                    on=\"record_id_y\",\n",
    "                )\n",
    "                .drop(columns=[\"record_id_y\"]),\n",
    "            ]\n",
    "        )\n",
    "        result = result.drop(columns=[\"record_id_x\", \"record_id_y\"])\n",
    "\n",
    "    return result, source_record_pairs\n",
    "\n",
    "\n",
    "def dedupe_preserving_source_records(\n",
    "    df, source_record_pairs, columns_to_dedupe, new_record_id_prefix\n",
    "):\n",
    "    result = df_ops.drop_duplicates(df[columns_to_dedupe])\n",
    "    result = add_unique_record_id(result, new_record_id_prefix)\n",
    "    df_to_result_mapping = (\n",
    "        df[[\"record_id\"] + columns_to_dedupe]\n",
    "        .rename(columns={\"record_id\": \"record_id_pre_dedupe\"})\n",
    "        .merge(result, on=columns_to_dedupe)[[\"record_id\", \"record_id_pre_dedupe\"]]\n",
    "    )\n",
    "    result_source_record_pairs = (\n",
    "        source_record_pairs.rename(columns={\"record_id\": \"record_id_pre_dedupe\"})\n",
    "        .merge(df_to_result_mapping, on=\"record_id_pre_dedupe\")\n",
    "        .drop(columns=[\"record_id_pre_dedupe\"])\n",
    "    )\n",
    "    return result, result_source_record_pairs\n",
    "\n",
    "\n",
    "def concat_preserving_source_records(dfs, source_record_pairings, new_record_id_prefix):\n",
    "    dfs = [df.rename(columns={\"record_id\": \"record_id_pre_concat\"}) for df in dfs]\n",
    "    result = df_ops.concat(dfs, ignore_index=True)\n",
    "    result = add_unique_record_id(result, new_record_id_prefix)\n",
    "\n",
    "    record_id_mapping = result[[\"record_id\", \"record_id_pre_concat\"]]\n",
    "    validate_kwarg = (\n",
    "        {\n",
    "            \"validate\": \"m:1\",\n",
    "        }\n",
    "        if not compute_engine.startswith(\"dask\")\n",
    "        else {}\n",
    "    )  # Dask doesn't support validate\n",
    "\n",
    "    all_source_record_pairings = df_ops.persist(\n",
    "        df_ops.concat(source_record_pairings, ignore_index=False)\n",
    "    )\n",
    "    result_source_record_pairs = df_ops.persist(\n",
    "        all_source_record_pairings.rename(columns={\"record_id\": \"record_id_pre_concat\"})\n",
    "        .merge(record_id_mapping, on=\"record_id_pre_concat\", **validate_kwarg)\n",
    "        .drop(columns=[\"record_id_pre_concat\"])\n",
    "    )\n",
    "\n",
    "    if compute_engine.startswith(\"dask\"):\n",
    "        # Manual alternative to the 'validate' kwarg\n",
    "        assert len(result_source_record_pairs) == len(all_source_record_pairings)\n",
    "\n",
    "    return result.drop(columns=[\"record_id_pre_concat\"]), result_source_record_pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebe6d46d-4188-492a-a586-dda76a5e546b",
   "metadata": {},
   "source": [
    "### Custom noise types\n",
    "\n",
    "A few noise types are really important here, but not implemented in pseudopeople, so we implement them as \"add-ons.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96eacdf9-c72b-4b38-8882-76cd238555c2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def split_columns_incorrectly(\n",
    "    df: pandas.DataFrame, columns: list[str], row_probability: float, seed: int = 1234\n",
    "):\n",
    "    rng = numpy.random.default_rng(seed)\n",
    "\n",
    "    to_split_incorrectly = pandas.Series(\n",
    "        rng.random(len(df)) < row_probability, index=df.index\n",
    "    )\n",
    "    full_name = (\n",
    "        df.loc[to_split_incorrectly, columns]\n",
    "        .fillna(\"\")\n",
    "        .agg(\" \".join, axis=1)\n",
    "        .str.strip()\n",
    "    )\n",
    "\n",
    "    num_spaces_or_hyphens = full_name.str.count(\"[ -]\")\n",
    "    split_name = full_name.str.split(\"[ -]\", regex=True, expand=True).fillna(\"\").values\n",
    "\n",
    "    split_points = numpy.ceil(\n",
    "        numpy.sort(rng.random((len(full_name), len(columns) - 1)), axis=1)\n",
    "        * num_spaces_or_hyphens.values.reshape((len(full_name), 1))\n",
    "    ).astype(int)\n",
    "    string_start_points_1d = (\n",
    "        numpy.arange(split_name.shape[0], dtype=int) * split_name.shape[1]\n",
    "    )\n",
    "    string_end_points_1d = numpy.append(\n",
    "        string_start_points_1d[1:], [split_name.shape[0] * split_name.shape[1]]\n",
    "    )\n",
    "    last_split_points_1d = None\n",
    "\n",
    "    for split_index, column_name in enumerate(columns):\n",
    "        # Map into 1D space\n",
    "        new_column = split_name.copy().reshape(-1)\n",
    "\n",
    "        if last_split_points_1d is not None:\n",
    "            erase_before_slices = tuple(\n",
    "                slice(string_start_point, last_split_point)\n",
    "                for string_start_point, last_split_point in zip(\n",
    "                    string_start_points_1d, last_split_points_1d\n",
    "                )\n",
    "            )\n",
    "        else:\n",
    "            erase_before_slices = tuple()\n",
    "\n",
    "        if split_index < split_points.shape[1]:\n",
    "            split_points_1d = (\n",
    "                numpy.arange(split_name.shape[0], dtype=int) * split_name.shape[1]\n",
    "                + split_points[:, split_index]\n",
    "            )\n",
    "            erase_after_slices = tuple(\n",
    "                slice(split_point, string_end_point)\n",
    "                for split_point, string_end_point in zip(\n",
    "                    split_points_1d, string_end_points_1d\n",
    "                )\n",
    "            )\n",
    "        else:\n",
    "            erase_after_slices = tuple()\n",
    "\n",
    "        new_column[numpy.r_[erase_before_slices + erase_after_slices]] = \"\"\n",
    "        new_column[new_column != \"\"] += \" \"\n",
    "        df.loc[to_split_incorrectly, column_name] = (\n",
    "            new_column.reshape(split_name.shape)\n",
    "        ).sum(axis=1)\n",
    "        df.loc[to_split_incorrectly, column_name] = df.loc[\n",
    "            to_split_incorrectly, column_name\n",
    "        ].str.strip()\n",
    "\n",
    "        last_split_points_1d = split_points_1d\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dbb8945-e5c7-4dc3-8387-023dbaef3fa0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def shuffle_columns(\n",
    "    df: pandas.DataFrame, columns: list[str], row_probability: float, seed: int = 1234\n",
    "):\n",
    "    rng = numpy.random.default_rng(seed)\n",
    "\n",
    "    to_shuffle = pandas.Series(rng.random(len(df)) < row_probability, index=df.index)\n",
    "    original_values = df.loc[to_shuffle, columns].values\n",
    "    random_ordering = numpy.argsort(rng.random(original_values.shape), axis=1)\n",
    "    df.loc[to_shuffle, columns] = numpy.take_along_axis(\n",
    "        original_values, random_ordering, axis=1\n",
    "    )\n",
    "\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c87d4bd4-95fa-4931-a8a9-4887dcf03d55",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def truncate_column(\n",
    "    df: pandas.DataFrame,\n",
    "    column: str,\n",
    "    cell_probability: float,\n",
    "    truncation_length: int,\n",
    "    seed: int = 1234,\n",
    "):\n",
    "    rng = numpy.random.default_rng(seed)\n",
    "    to_truncate = pandas.Series(rng.random(len(df)) < cell_probability, index=df.index)\n",
    "    df.loc[to_truncate, \"last_name\"] = df.loc[to_truncate, \"last_name\"].str[\n",
    "        :truncation_length\n",
    "    ]\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14c682ac-0533-4418-a9b8-2753aa42727b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def apply_custom_noise_type(df, noise_type, *args, seed: int = 1234, **kwargs):\n",
    "    if compute_engine.startswith(\"dask\"):\n",
    "        return df.map_partitions(\n",
    "            lambda partition, partition_info=None: noise_type(\n",
    "                partition,\n",
    "                *args,\n",
    "                seed=(\n",
    "                    seed\n",
    "                    + (partition_info[\"number\"] if partition_info is not None else 1)\n",
    "                    * 10_000\n",
    "                ),\n",
    "                **kwargs,\n",
    "            ),\n",
    "            meta=df,\n",
    "        )\n",
    "    else:\n",
    "        return noise_type(\n",
    "            df,\n",
    "            *args,\n",
    "            seed=seed,\n",
    "            **kwargs,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fe8e46b-f936-4704-b120-9bbcbce6c3e9",
   "metadata": {},
   "source": [
    "### 1040 tax filings\n",
    "\n",
    "We assume that the last 5 years of taxes would be available and used in the construction of the reference files -- see section about reference files below.\n",
    "\n",
    "Note that these are retrieved by *tax* year, so the 2029 taxes would be available in early 2030\n",
    "(around when our hypothetical case study is taking place)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8157a22b-45f3-4fc1-803e-4670a554b23f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tax_years = list(range(2025, 2030))\n",
    "tax_years"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8df3663-c547-48c1-b8e1-669ebfbdaeec",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Combine 1040 for all years.\n",
    "simulated_taxes_1040 = df_ops.concat(\n",
    "    [\n",
    "        df_ops.read_parquet(\n",
    "            f\"{input_dir}/{data_to_use}/simulated_taxes_1040_{year}.parquet\",\n",
    "            columns=[\n",
    "                \"simulant_id\",\n",
    "                \"ssn\",\n",
    "                \"first_name\",\n",
    "                \"middle_initial\",\n",
    "                \"last_name\",\n",
    "                \"mailing_address_street_number\",\n",
    "                \"mailing_address_street_name\",\n",
    "                \"mailing_address_unit_number\",\n",
    "                \"mailing_address_po_box\",\n",
    "                \"mailing_address_city\",\n",
    "                \"mailing_address_state\",\n",
    "                \"mailing_address_zipcode\",\n",
    "            ],\n",
    "        )\n",
    "        for year in tax_years\n",
    "    ],\n",
    "    ignore_index=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f72c1ac-dbf4-40dc-acdc-10d406938d19",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_taxes_1040 = apply_custom_noise_type(\n",
    "    simulated_taxes_1040,\n",
    "    split_columns_incorrectly,\n",
    "    [\"first_name\", \"middle_initial\", \"last_name\"],\n",
    "    row_probability=0.01,\n",
    "    seed=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25a0dfbe-f3be-4847-9f26-c3f9f68e508b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# \"... many of the [IRS] records contain only the first four letters of the last name.\"\n",
    "# (Brown et al. 2023, p.30, footnote 19)\n",
    "# Note that this truncation only matters for ITIN PIKing since for SSNs that are present in SSA we use name from SSA.\n",
    "PROPORTION_OF_IRS_RECORDS_WITH_TRUNCATION = (\n",
    "    0.4  # is this a good guess at \"many\" in the quote above?\n",
    ")\n",
    "simulated_taxes_1040 = apply_custom_noise_type(\n",
    "    simulated_taxes_1040,\n",
    "    truncate_column,\n",
    "    \"last_name\",\n",
    "    cell_probability=PROPORTION_OF_IRS_RECORDS_WITH_TRUNCATION,\n",
    "    truncation_length=4,\n",
    "    seed=2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de591252-0bce-4e65-9269-708c02028a52",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_taxes_1040 = apply_custom_noise_type(\n",
    "    simulated_taxes_1040,\n",
    "    shuffle_columns,\n",
    "    [\"first_name\", \"middle_initial\", \"last_name\"],\n",
    "    row_probability=0.01,\n",
    "    seed=3,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8dd6b79-ee7a-4942-b6fd-bfa0804c1880",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_taxes_1040 = df_ops.persist(\n",
    "    add_unique_record_id(simulated_taxes_1040, \"simulated_1040\")\n",
    ")\n",
    "simulated_taxes_1040_source_record_pairs = df_ops.persist(\n",
    "    record_id_to_single_source_record_pairs(simulated_taxes_1040)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76c53fed-1cf8-4ec1-ae97-10c39d12d92e",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_taxes_1040_ground_truth = simulated_taxes_1040[[\"record_id\", \"simulant_id\"]]\n",
    "simulated_taxes_1040 = simulated_taxes_1040.drop(columns=[\"simulant_id\"])\n",
    "simulated_taxes_1040, simulated_taxes_1040_ground_truth = df_ops.persist(\n",
    "    simulated_taxes_1040, simulated_taxes_1040_ground_truth\n",
    ")\n",
    "simulated_taxes_1040"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8bad94f-c803-41e0-bf17-9ecfdb56cdbd",
   "metadata": {},
   "source": [
    "### W2/1099 tax filings\n",
    "\n",
    "We assume that the last 5 years of taxes would be available and used in the construction of the reference files --\n",
    "see section about reference files below.\n",
    "\n",
    "Note that these are retrieved by *tax* year, so the 2029 taxes would be available in early 2030\n",
    "(around when our hypothetical case study is taking place)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c072c69-f7c7-4bf1-b99e-48c3257063ba",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Combine W2/1099 for all years.\n",
    "simulated_w2_1099 = df_ops.concat(\n",
    "    [\n",
    "        df_ops.read_parquet(\n",
    "            f\"{input_dir}/{data_to_use}/simulated_taxes_w2_and_1099_{year}.parquet\",\n",
    "            columns=[\n",
    "                \"simulant_id\",\n",
    "                \"ssn\",\n",
    "                \"first_name\",\n",
    "                \"middle_initial\",\n",
    "                \"last_name\",\n",
    "                \"mailing_address_street_number\",\n",
    "                \"mailing_address_street_name\",\n",
    "                \"mailing_address_unit_number\",\n",
    "                \"mailing_address_po_box\",\n",
    "                \"mailing_address_city\",\n",
    "                \"mailing_address_state\",\n",
    "                \"mailing_address_zipcode\",\n",
    "            ],\n",
    "        )\n",
    "        for year in tax_years\n",
    "    ],\n",
    "    ignore_index=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b1e35ee-d2dd-4fa0-ac4c-abc3cac65cfc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_w2_1099 = apply_custom_noise_type(\n",
    "    simulated_w2_1099,\n",
    "    split_columns_incorrectly,\n",
    "    [\"first_name\", \"middle_initial\", \"last_name\"],\n",
    "    row_probability=0.01,\n",
    "    seed=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e66f562-b58f-4f6f-821a-6260e00737db",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_w2_1099 = apply_custom_noise_type(\n",
    "    simulated_w2_1099,\n",
    "    truncate_column,\n",
    "    \"last_name\",\n",
    "    cell_probability=PROPORTION_OF_IRS_RECORDS_WITH_TRUNCATION,\n",
    "    truncation_length=4,\n",
    "    seed=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb6310e4-65ac-40f5-b060-791b58d42676",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_w2_1099 = apply_custom_noise_type(\n",
    "    simulated_w2_1099,\n",
    "    shuffle_columns,\n",
    "    [\"first_name\", \"middle_initial\", \"last_name\"],\n",
    "    row_probability=0.01,\n",
    "    seed=6,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d63dfd5-3cca-4bf0-8683-6b9382acfe56",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_w2_1099 = df_ops.persist(\n",
    "    add_unique_record_id(simulated_w2_1099, \"simulated_w2_1099\")\n",
    ")\n",
    "simulated_w2_1099_source_record_pairs = df_ops.persist(\n",
    "    record_id_to_single_source_record_pairs(simulated_w2_1099)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f54aee-ad8f-45db-abfa-25e4d027f236",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_w2_1099_ground_truth = simulated_w2_1099[[\"record_id\", \"simulant_id\"]]\n",
    "simulated_w2_1099 = simulated_w2_1099.drop(columns=[\"simulant_id\"])\n",
    "simulated_w2_1099, simulated_w2_1099_ground_truth = df_ops.persist(\n",
    "    simulated_w2_1099, simulated_w2_1099_ground_truth\n",
    ")\n",
    "simulated_w2_1099"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32e84952-9adb-4fcd-8a43-0ec0b76bbd5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_taxes, simulated_taxes_source_record_pairs = concat_preserving_source_records(\n",
    "    [simulated_taxes_1040, simulated_w2_1099],\n",
    "    [simulated_taxes_1040_source_record_pairs, simulated_w2_1099_source_record_pairs],\n",
    "    new_record_id_prefix=\"simulated_taxes\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7977c68-f5ce-4aa6-b641-fb4e724a57df",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_taxes_1040_with_itins = simulated_taxes_1040[\n",
    "    simulated_taxes_1040.ssn.notnull() & simulated_taxes_1040.ssn.str.startswith(\"9\")\n",
    "]\n",
    "simulated_taxes_1040_with_itins_source_record_pairs = (\n",
    "    simulated_taxes_1040_source_record_pairs.merge(\n",
    "        simulated_taxes_1040_with_itins[[\"record_id\"]], on=\"record_id\", how=\"inner\"\n",
    "    )\n",
    ")\n",
    "simulated_taxes_1040_with_itins, simulated_taxes_1040_with_itins_source_record_pairs = (\n",
    "    df_ops.persist(\n",
    "        simulated_taxes_1040_with_itins,\n",
    "        simulated_taxes_1040_with_itins_source_record_pairs,\n",
    "    )\n",
    ")\n",
    "\n",
    "%xdel simulated_taxes_1040\n",
    "%xdel simulated_taxes_1040_source_record_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38e94242-ff62-4d3c-ba5b-a94e3ba0a6e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%xdel simulated_w2_1099\n",
    "%xdel simulated_w2_1099_source_record_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "201d14e0-53e3-4e1c-99f2-9cfaec0d852e",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_taxes, simulated_taxes_source_record_pairs = df_ops.persist(\n",
    "    simulated_taxes, simulated_taxes_source_record_pairs\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed6967cb-dfa6-4b43-98b9-e590d266f55f",
   "metadata": {},
   "source": [
    "### 2030 Census Unedited File (CUF)\n",
    "\n",
    "For now, we gloss over the data schema for addresses.\n",
    "We don't know how addresses would be formatted in the CUF (and it's hard to guess, because\n",
    "address is not part of the Census form), but it likely would have some of these fields\n",
    "(street number, street name, etc) combined.\n",
    "\n",
    "While PVS input files do not in general have names split into first, middle, and last,\n",
    "I am guessing the CUF **would** have first name, middle initial, last name (which is how pseudopeople\n",
    "generates it), because that [matches the Census questionnaire](https://www2.census.gov/programs-surveys/decennial/2020/technical-documentation/questionnaires-and-instructions/questionnaires/2020-informational-questionnaire-english_DI-Q1.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e98619b-2ae6-4c89-a82e-9f975cdc9caf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "simulated_census_2030 = df_ops.read_parquet(\n",
    "    f\"{input_dir}/{data_to_use}/simulated_census_2030.parquet\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "606c2fc7-553e-445f-b295-3f33dc584f8e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_census_2030 = apply_custom_noise_type(\n",
    "    simulated_census_2030,\n",
    "    split_columns_incorrectly,\n",
    "    [\"first_name\", \"middle_initial\", \"last_name\"],\n",
    "    row_probability=0.01,\n",
    "    seed=7,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1907e12c-f19b-4805-b571-f2855d971d3b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_census_2030 = apply_custom_noise_type(\n",
    "    simulated_census_2030,\n",
    "    shuffle_columns,\n",
    "    [\"first_name\", \"middle_initial\", \"last_name\"],\n",
    "    row_probability=0.01,\n",
    "    seed=8,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "164bbda6-29a1-4add-b5dd-5c666d552ed8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_census_2030 = df_ops.persist(\n",
    "    add_unique_record_id(simulated_census_2030, \"simulated_census_2030\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89e992c8-d7e1-4b63-8f44-be6dd89808ed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_census_2030_ground_truth = simulated_census_2030[[\"record_id\", \"simulant_id\"]]\n",
    "simulated_census_2030 = simulated_census_2030.drop(columns=[\"simulant_id\"])\n",
    "simulated_census_2030, simulated_census_2030_ground_truth = df_ops.persist(\n",
    "    simulated_census_2030, simulated_census_2030_ground_truth\n",
    ")\n",
    "simulated_census_2030"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e361daa-1ff6-4bd1-a5a8-21d0cc191e36",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "save_file_with_ground_truth(\n",
    "    \"simulated_census_2030\", simulated_census_2030, simulated_census_2030_ground_truth\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "580c2493-25a6-4a47-a82b-cc78c5d3b5d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%xdel simulated_census_2030"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62aee3f5-01ab-4e55-af67-f0cf656d9af0",
   "metadata": {},
   "source": [
    "### Simulated SSA Numident"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d6496b3-0c2c-4739-acad-2c3c81056001",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "simulated_ssa_numident = df_ops.read_parquet(\n",
    "    f\"{input_dir}/{data_to_use}/simulated_ssa_numident.parquet\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "224d38c5-89f0-44cb-96b5-452799f363de",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_ssa_numident = apply_custom_noise_type(\n",
    "    simulated_ssa_numident,\n",
    "    split_columns_incorrectly,\n",
    "    [\"first_name\", \"middle_name\", \"last_name\"],\n",
    "    row_probability=0.01,\n",
    "    seed=9,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "842af49f-cd65-4963-ba27-e7de57638a81",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_ssa_numident = apply_custom_noise_type(\n",
    "    simulated_ssa_numident,\n",
    "    shuffle_columns,\n",
    "    [\"first_name\", \"middle_name\", \"last_name\"],\n",
    "    row_probability=0.01,\n",
    "    seed=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "828c67a3-23f9-4fb7-8737-6ce29f9f80bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "if data_to_use == \"ri\":\n",
    "    # Remove simulants who were never observed in Rhode Island -- this is unrealistic to real life!\n",
    "    # We only do this to create a medium-sized set of reference files for computational testing.\n",
    "    simulated_ssa_numident = (\n",
    "        simulated_ssa_numident.merge(\n",
    "            df_ops.drop_duplicates(\n",
    "                simulated_taxes_1040_ground_truth[[\"simulant_id\"]]\n",
    "            ).assign(observed_in_taxes_1040=1),\n",
    "            on=\"simulant_id\",\n",
    "            how=\"left\",\n",
    "        )\n",
    "        .merge(\n",
    "            df_ops.drop_duplicates(\n",
    "                simulated_w2_1099_ground_truth[[\"simulant_id\"]]\n",
    "            ).assign(observed_in_taxes_w2_1099=1),\n",
    "            on=\"simulant_id\",\n",
    "            how=\"left\",\n",
    "        )\n",
    "        .merge(\n",
    "            df_ops.drop_duplicates(\n",
    "                simulated_census_2030_ground_truth[[\"simulant_id\"]]\n",
    "            ).assign(observed_in_census_2030=1),\n",
    "            on=\"simulant_id\",\n",
    "            how=\"left\",\n",
    "        )\n",
    "        .assign(\n",
    "            observed_ever=lambda df: (df.observed_in_taxes_1040 == 1)\n",
    "            | (df.observed_in_taxes_w2_1099 == 1)\n",
    "            | (df.observed_in_census_2030 == 1)\n",
    "        )\n",
    "        .pipe(lambda df: df[df.observed_ever])\n",
    "        .drop(\n",
    "            columns=[\n",
    "                \"observed_in_taxes_1040\",\n",
    "                \"observed_in_taxes_w2_1099\",\n",
    "                \"observed_in_census_2030\",\n",
    "                \"observed_ever\",\n",
    "            ]\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37933d8a-b708-4061-a246-866c2a1ef55e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%xdel simulated_census_2030_ground_truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1516962f-1527-4f51-b5e6-16eefdcbaf26",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_ssa_numident = df_ops.persist(\n",
    "    add_unique_record_id(simulated_ssa_numident, \"simulated_ssa_numident\")\n",
    ")\n",
    "simulated_ssa_numident_source_record_pairs = df_ops.persist(\n",
    "    record_id_to_single_source_record_pairs(simulated_ssa_numident)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07b4e12e-21c7-4f2c-bfa8-cb78a0e595d0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_ssa_numident_ground_truth = simulated_ssa_numident[\n",
    "    [\"record_id\", \"simulant_id\"]\n",
    "]\n",
    "simulated_ssa_numident = simulated_ssa_numident.drop(columns=[\"simulant_id\"])\n",
    "simulated_ssa_numident, simulated_ssa_numident_ground_truth = df_ops.persist(\n",
    "    simulated_ssa_numident, simulated_ssa_numident_ground_truth\n",
    ")\n",
    "simulated_ssa_numident"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "383bcc71-4e12-4626-9e4f-b232d6571c06",
   "metadata": {},
   "outputs": [],
   "source": [
    "source_record_ground_truth = df_ops.persist(\n",
    "    df_ops.concat(\n",
    "        [\n",
    "            simulated_ssa_numident_ground_truth,\n",
    "            simulated_taxes_1040_ground_truth,\n",
    "            simulated_w2_1099_ground_truth,\n",
    "        ]\n",
    "    ).rename(columns={\"record_id\": \"source_record_id\"})\n",
    ")\n",
    "source_record_ground_truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31af5d87-1cd6-41dd-b46d-79cea5a3e99e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%xdel simulated_taxes_1040_ground_truth\n",
    "%xdel simulated_w2_1099_ground_truth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d891f010-d434-46f9-bfd4-b17efc4eeac7",
   "metadata": {},
   "source": [
    "## Create reference files\n",
    "\n",
    "> The Census Numident – all Social Security Administration (SSA) Numident SSN records are\n",
    "  edited (collapsed) to produce a Census Numident file that contains “one best-data record” for\n",
    "  each SSN. All variants of name information for each SSN are retained in the Alternate Name\n",
    "  Numident file, while all variants of date of birth data are retained in the Alternate DOB\n",
    "  Numident. The SSN-PIK crosswalk file is used to attach a corresponding unique PIK value for\n",
    "  each SSN value in the Census Numident file."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa4fa52d-580f-4557-92a9-6f237d8b45f5",
   "metadata": {},
   "source": [
    "### Census Numident\n",
    "\n",
    "Luque and Wagner, p. 4:\n",
    "  \n",
    "> The SSA Numident file contains all transactions ever recorded against any single SSN - with each entry\n",
    "representing an addition or change (such as name changes) to the SSN record. This file is edited to\n",
    "create the **Census Numident**, which contains one record for each SSN. Each SSN record in the Census\n",
    "Numident contains name, DOB, sex, race, place of birth, parents’ name, citizenship status and date of death information.\n",
    "\n",
    "and in footnote 5:\n",
    "\n",
    "> Name edits, DOB reconciliation, and race identifiers are some of the edits conducted to produce this Numident\n",
    "file. **The resulting Numident file contains the most recent name and DOB data.**\n",
    "\n",
    "We are missing quite a few columns, since they are missing in pseudopeople's SSA Numident: race, place of birth, parents' name,\n",
    "citizenship status.\n",
    "However, I'm pretty sure none of these are used in linking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bde9378-288c-47ae-a739-6af8e81a5eef",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fill_dates(df, fill_with):\n",
    "    return (\n",
    "        # Replace invalid dates with nans\n",
    "        pd.to_datetime(df.event_date, format=\"%Y%m%d\", errors=\"coerce\").fillna(\n",
    "            pandas.to_datetime(\"2100-01-01\" if fill_with == \"latest\" else \"1900-01-01\")\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "def best_data_from_columns(df, columns, best_is_latest=True):\n",
    "    # We don't want to throw out events with a missing/invalid date, so we'll fill them with the value *least* likely to be chosen\n",
    "    # (earlier than all values if taking the latest, later than all values if taking the earliest).\n",
    "    fill_with = \"earliest\" if best_is_latest else \"latest\"\n",
    "\n",
    "    result = (\n",
    "        df\n",
    "        # Without mutating the existing date column, get one that is actually\n",
    "        # a date type and can be used for sorting.\n",
    "        # Note: we actually convert this to an integer for sorting purposes, because Modin was having trouble\n",
    "        # sorting by it as an actual datetime\n",
    "        .assign(\n",
    "            event_date_for_sort=lambda df: fill_dates(df, fill_with=fill_with).astype(\n",
    "                numpy.int64\n",
    "            )\n",
    "            // 10**9\n",
    "        )\n",
    "        .dropna(subset=columns, how=\"all\")\n",
    "        .pipe(\n",
    "            lambda df: df_ops.drop_duplicates(\n",
    "                df,\n",
    "                subset=\"ssn\",\n",
    "                sort_col=\"event_date_for_sort\",\n",
    "                keep=(\"last\" if best_is_latest else \"first\"),\n",
    "            )\n",
    "        )[[\"record_id\", \"ssn\"] + columns]\n",
    "    )\n",
    "\n",
    "    return result, record_id_to_single_source_record_pairs(result)\n",
    "\n",
    "\n",
    "best_name, best_name_source_record_pairs = best_data_from_columns(\n",
    "    simulated_ssa_numident,\n",
    "    columns=[\"first_name\", \"middle_name\", \"last_name\"],\n",
    ")\n",
    "\n",
    "best_date_of_birth, best_date_of_birth_source_record_pairs = best_data_from_columns(\n",
    "    simulated_ssa_numident,\n",
    "    columns=[\"date_of_birth\"],\n",
    ")\n",
    "\n",
    "best_date_of_death, best_date_of_death_source_record_pairs = best_data_from_columns(\n",
    "    simulated_ssa_numident[simulated_ssa_numident.event_type == \"death\"],\n",
    "    columns=[\"event_date\"],\n",
    ")\n",
    "best_date_of_death = best_date_of_death.rename(columns={\"event_date\": \"date_of_death\"})\n",
    "\n",
    "simulated_census_numident, simulated_census_numident_source_record_pairs = (\n",
    "    df_ops.persist(\n",
    "        merge_preserving_source_records(\n",
    "            [best_name, best_date_of_birth, best_date_of_death],\n",
    "            [\n",
    "                best_name_source_record_pairs,\n",
    "                best_date_of_birth_source_record_pairs,\n",
    "                best_date_of_death_source_record_pairs,\n",
    "            ],\n",
    "            new_record_id_prefix=\"simulated_census_numident\",\n",
    "            on=[\"ssn\"],\n",
    "            how=\"left\",\n",
    "        )\n",
    "    )\n",
    ")\n",
    "simulated_census_numident"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd465ebe-e5e2-47ab-a9e0-ae9d21d476da",
   "metadata": {},
   "source": [
    "### Alternate Name Numident\n",
    "\n",
    "Wagner and Layne, p. 9:\n",
    "\n",
    ">  All variants of name information for each SSN are retained in the Alternate Name\n",
    "Numident file..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "573250ce-0525-45de-9df9-33d8ea757161",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "(\n",
    "    simulated_alternate_name_numident,\n",
    "    simulated_alternate_name_numident_source_record_pairs,\n",
    ") = df_ops.persist(\n",
    "    dedupe_preserving_source_records(\n",
    "        simulated_ssa_numident,\n",
    "        simulated_ssa_numident_source_record_pairs,\n",
    "        columns_to_dedupe=[\"ssn\", \"first_name\", \"middle_name\", \"last_name\"],\n",
    "        new_record_id_prefix=\"simulated_alternate_name_numident\",\n",
    "    )\n",
    ")\n",
    "simulated_alternate_name_numident"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "739b9b20-dde1-4079-ae90-9c5608d7a683",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_ops.compute(\n",
    "    df_ops.groupby_agg_small_groups(\n",
    "        simulated_alternate_name_numident, by=\"ssn\", agg_func=lambda x: x.size()\n",
    "    ).describe()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d67dd34f-2edf-4dcd-a2d2-80f421c24884",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Show alternate names for the same simulated SSN\n",
    "simulated_alternate_name_numident.merge(\n",
    "    df_ops.groupby_agg_small_groups(\n",
    "        simulated_alternate_name_numident[[\"ssn\"]],\n",
    "        by=\"ssn\",\n",
    "        agg_func=lambda x: x.size(),\n",
    "    )\n",
    "    .pipe(lambda x: x[x > 1])\n",
    "    .reset_index()[[\"ssn\"]],\n",
    "    on=\"ssn\",\n",
    "    how=\"inner\",\n",
    ").sort_values(\"ssn\").tail(100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b44f0a25-79c3-47c2-8474-f6afd205f7e2",
   "metadata": {},
   "source": [
    "### Alternate DOB Numident\n",
    "\n",
    "Wagner and Layne, p. 9:\n",
    "\n",
    "> ... while all variants of date of birth data are retained in the Alternate DOB\n",
    "Numident."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "128fe5fb-57b6-4ccb-9096-e6aafd5452ee",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "(\n",
    "    simulated_alternate_dob_numident,\n",
    "    simulated_alternate_dob_numident_source_record_pairs,\n",
    ") = df_ops.persist(\n",
    "    dedupe_preserving_source_records(\n",
    "        simulated_ssa_numident,\n",
    "        simulated_ssa_numident_source_record_pairs,\n",
    "        columns_to_dedupe=[\"ssn\", \"date_of_birth\"],\n",
    "        new_record_id_prefix=\"simulated_alternate_dob_numident\",\n",
    "    )\n",
    ")\n",
    "simulated_alternate_dob_numident"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85796613-6206-4587-b13f-f13cfab9fb89",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_ops.compute(\n",
    "    df_ops.groupby_agg_small_groups(\n",
    "        simulated_alternate_dob_numident, by=\"ssn\", agg_func=lambda x: x.size()\n",
    "    ).describe()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb4537c0-d04c-42a3-8a66-8d12e56af33e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_alternate_dob_numident.merge(\n",
    "    df_ops.groupby_agg_small_groups(\n",
    "        simulated_alternate_dob_numident[[\"ssn\"]], by=\"ssn\", agg_func=lambda x: x.size()\n",
    "    )\n",
    "    .pipe(lambda x: x[x > 1])\n",
    "    .reset_index()[[\"ssn\"]],\n",
    "    on=\"ssn\",\n",
    "    how=\"inner\",\n",
    ").sort_values(\"ssn\").tail(100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9040aaad-1228-4921-bc95-ad3011f1bdf3",
   "metadata": {},
   "source": [
    "### Name/DOB Reference File\n",
    "\n",
    "Wagner and Layne, p. 9:\n",
    "\n",
    "> The Name and DOB Reference files are reformatted versions of the Census Numident\n",
    "and includes **all possible combinations of alternate names and dates of birth, as well as\n",
    "ITIN data**. All of the reference files contain SSN/ITIN and the corresponding PIK. When\n",
    "an input record is linked to a reference file, the corresponding PIK is assigned. Table 1\n",
    "presents the number of observations in each of the reference files.\n",
    "\n",
    "A slightly confusing point: sometimes the Name and DOB reference files are described\n",
    "as one and the same thing, and sometimes as separate.\n",
    "I believe this is because **they differ only in how they are \"cut\" for the PVS process:**\n",
    "the name reference file is cut by first and last initial,\n",
    "while the DOB reference file is cut by month and day of birth.\n",
    "\n",
    "This is described in Wagner and Layne, p.15:\n",
    "\n",
    "> The [DOBSearch] module matches against a re-split\n",
    "version of the Numident Name Reference file, splitting the data based on month and day\n",
    "of birth.\n",
    "\n",
    "Since we handle the logic of \"cutting\" in the linkage process itself, we generate\n",
    "a single reference file here.\n",
    "\n",
    "Note that unlike for addresses, and unlike for the pre-processing of *input* files\n",
    "(as opposed to reference files), there is no explicit nickname processing/correction here.\n",
    "I am fairly sure that is accurate to the PVS, which I believe assumes that nicknames\n",
    "would not be present in SSA/tax records (or at least, that the real name would appear\n",
    "at least once in these records)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29cc9e45-d862-484a-b5ff-4c12ff6b00c9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "(\n",
    "    simulated_name_dob_numident_records,\n",
    "    simulated_name_dob_numident_records_source_record_pairs,\n",
    ") = df_ops.persist(\n",
    "    merge_preserving_source_records(\n",
    "        [simulated_alternate_name_numident, simulated_alternate_dob_numident],\n",
    "        [\n",
    "            simulated_alternate_name_numident_source_record_pairs,\n",
    "            simulated_alternate_dob_numident_source_record_pairs,\n",
    "        ],\n",
    "        on=[\"ssn\"],\n",
    "        how=\"left\",\n",
    "        new_record_id_prefix=\"name_dob_numident_records\",\n",
    "    )\n",
    ")\n",
    "simulated_name_dob_numident_records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eb4add5-6f0b-4d43-ac31-6e085a678e06",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Show alternate name/DOB for the same simulated SSN\n",
    "simulated_name_dob_numident_records.merge(\n",
    "    df_ops.groupby_agg_small_groups(\n",
    "        simulated_name_dob_numident_records[[\"ssn\"]],\n",
    "        by=\"ssn\",\n",
    "        agg_func=lambda x: x.size(),\n",
    "    )\n",
    "    .pipe(lambda x: x[x > 1])\n",
    "    .reset_index()[[\"ssn\"]],\n",
    "    on=\"ssn\",\n",
    "    how=\"inner\",\n",
    ").sort_values(\"ssn\").tail(100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06083d51-f83c-4cd0-92cf-e1364220456b",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Incorporating people with ITINs\n",
    "\n",
    "Individual Taxpayer Identification Numbers (ITINs) can be issued to people who are required to file\n",
    "federal taxes but are not eligible for a Social Security Number.\n",
    "The most common reason for this is being an undocumented immigrant and therefore not being authorized\n",
    "to work in the United States.\n",
    "\n",
    "People without SSNs used to be impossible to assign PIKs to.\n",
    "In 2011 the NORC report stated (p. 38, footnote 19):\n",
    "\n",
    "> NORC understands that the Census Bureau has undertaken an effort to enhance the PVS reference files with IRS\n",
    "files that include Individual Taxpayer Identification Numbers (ITIN). For those people who are required to file a tax\n",
    "return but do not have, and may not want an SSN—such as a non-U.S. citizen—the IRS issues the taxpayer an ITIN.\n",
    "This enhancement to the PVS reference file may help to match more non-U.S citizens.\n",
    "\n",
    "By 2014 (Wagner and Layne, p. 5):\n",
    "\n",
    "> One of the key enhancements [made in recent years] increased the coverage of the reference files by\n",
    "including records for persons with Individual Taxpayer Identification Numbers assigned\n",
    "by the Internal Revenue Service (ITINs) to [along with?] the SSN-based Numident data. \n",
    "\n",
    "I have not found a specific description of how ITIN records are constructed in any of the\n",
    "publicly-available sources.\n",
    "This may be because it is straightforward, or because the tax data schema is confidential.\n",
    "I assume that only IRS data is used, since no other data source that I am aware of would\n",
    "report ITIN.\n",
    "\n",
    "It is stated that the ITIN records are created directly from tax filings and not\n",
    "from ITIN applications (Brown et al. p. 29, footnote 16), which is convenient\n",
    "because the tax filing data is what we can simulate with\n",
    "pseudopeople:\n",
    "\n",
    "> The NUMIDENT provides the PII on the SSN-holder from the issuing agency (SSA), and that PII is used in SSN\n",
    "verification. **For ITINs, the Census Bureau does not have access to the ITIN applications** to the issuing agency (IRS),\n",
    "so name and DOB verification of ITINs is less reliable.\n",
    "\n",
    "\"Less reliable\" is a bit confusing here, because as stated above when generating\n",
    "the simulated tax data, IRS data should not contain date of birth at all.\n",
    "Here, we have stayed true to this by omitting it entirely.\n",
    "\n",
    "We assume that only 1040 filings would be used for this purpose; we wouldn't expect ITINs to\n",
    "show up on employer-filed W-2/1099 forms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a580a0a7-3a43-4cd9-b427-1f72920fbcd6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Analogous to the process of getting alternate names and dates of birth\n",
    "# from SSA, we retain all versions of the name from taxes.\n",
    "simulated_name_for_itins, simulated_name_for_itins_source_record_pairs = df_ops.persist(\n",
    "    dedupe_preserving_source_records(\n",
    "        simulated_taxes_1040_with_itins.rename(\n",
    "            columns={\"middle_initial\": \"middle_name\"}\n",
    "        ),\n",
    "        simulated_taxes_1040_with_itins_source_record_pairs,\n",
    "        columns_to_dedupe=[\"ssn\", \"first_name\", \"middle_name\", \"last_name\"],\n",
    "        new_record_id_prefix=\"simulated_name_for_itins\",\n",
    "    )\n",
    ")\n",
    "simulated_name_for_itins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77482b40-4980-4e30-8ebf-bdcaea5023ad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df_ops.compute(\n",
    "    df_ops.groupby_agg_small_groups(\n",
    "        simulated_name_for_itins, by=\"ssn\", agg_func=lambda x: x.size()\n",
    "    ).describe()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ff43555-25ea-4f23-aa31-46f82f3405e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# How many records in the simulated name/DOB numident are in the ITIN range?\n",
    "# With normal pseudopeople settings, this shouldn't happen, because SSN in SSA is noiseless.\n",
    "# If this number is >0, we may be adding ITIN records as more alternates to existing records.\n",
    "df_ops.compute(simulated_name_dob_numident_records.ssn.str.startswith(\"9\").sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed5e69b5-1d43-4ea3-8fc2-b88eae1a2d92",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "(\n",
    "    simulated_name_dob_reference_file,\n",
    "    simulated_name_dob_reference_file_source_record_pairs,\n",
    ") = df_ops.persist(\n",
    "    concat_preserving_source_records(\n",
    "        [simulated_name_dob_numident_records, simulated_name_for_itins],\n",
    "        [\n",
    "            simulated_name_dob_numident_records_source_record_pairs,\n",
    "            simulated_name_for_itins_source_record_pairs,\n",
    "        ],\n",
    "        new_record_id_prefix=\"simulated_name_dob_reference_file\",\n",
    "    )\n",
    ")\n",
    "simulated_name_dob_reference_file"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9816d72-0a20-45d2-be2a-eb47872008da",
   "metadata": {},
   "source": [
    "### GeoBase Reference File\n",
    "\n",
    "Wagner and Layne, p. 9:\n",
    "\n",
    "> PVS creates three other sets of reference\n",
    "files containing Numident data: the **GeoBase Reference File**, the Name Reference File,\n",
    "and the DOB Reference file.\n",
    "The GeoBase Reference File appends addresses from administrative records attached\n",
    "to Numident data, including all possible combinations of alternate names and dates of\n",
    "birth for SSN. Addresses from administrative records are edited and processed through\n",
    "commercial software product to clean and standardize address data. ITIN data is also\n",
    "incorporated into the Geobase.\n",
    "\n",
    "Luque and Wagner, p. 5:\n",
    "\n",
    "> Reference files contain data from the Numident file enhanced with address\n",
    "data obtained from federal AR [administrative records] files.<sup>8</sup>\n",
    "The reference files, thus, contain all variants of a person’s name, DOB,\n",
    "and sex, as well as current and recent addresses. These reference files are\n",
    "referred to as the (PVS) Geobase reference file since addresses (a geographic component)\n",
    "are appended to each person record.<sup>9</sup> It is important to note that there are\n",
    "multiple Geobase reference files that are created depending on the vintage of the\n",
    "incoming file to be processed through PVS.\n",
    "\n",
    "> <sup>8</sup> Namely, data from the IRS, Department of Housing and Urban Development,\n",
    "several files from the Department of Health and Human Services, and Selective Service.\n",
    "\n",
    "> <sup>9</sup> In particular, the address data is cleaned and standardized and used\n",
    "to construct a variable called GEOKEY. The GEOKEY variable is constructed as a subset\n",
    "of the full address, and then is appended to the Numident data to create the\n",
    "PVS Geobase Reference file.\n",
    "\n",
    "We only have IRS data to use for addresses, and specifically only W-2/1099 data,\n",
    "which is a limitation of this case study.\n",
    "I can't find a concrete definition of \"recent\" -- as noted above, we use 5 years\n",
    "of IRS data.\n",
    "This may be longer than the true window, but this may end up making up for\n",
    "the lack of non-IRS data sources, and get us closer to a realistic number of\n",
    "alternate addresses.\n",
    "\n",
    "Also, our address data comes out of pseudopeople already parsed into address parts\n",
    "like street name, etc.\n",
    "For more realism, pseudopeople should output a single string that we have to (imperfectly) parse apart.\n",
    "\n",
    "I haven't been able to find out more about what kind of \"subset\" the geokey is.\n",
    "It is unclear to me why geokey is \"interesting\" since it is just derived from the\n",
    "address parts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59bc80de-7ef0-4826-aed0-6289c12dfcb4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "address_cols = [c for c in simulated_taxes.columns if \"mailing_address\" in c]\n",
    "\n",
    "\n",
    "def standardize_address_part(column):\n",
    "    return (\n",
    "        column\n",
    "        # Remove leading or trailing whitespace\n",
    "        .str.strip()\n",
    "        # Turn any strings of consecutive whitespace into a single space\n",
    "        .str.replace(\"\\s+\", \" \", regex=True)\n",
    "        # Normalize case\n",
    "        .str.upper()\n",
    "        # Normalize the word street as described in the example quoted above\n",
    "        # In reality, there would be many rules like this\n",
    "        .str.replace(\"\\b(STREET|STR)\\b\", \"ST\", regex=True)\n",
    "        # Make sure missingness is represented consistently\n",
    "        .replace(\"\", numpy.nan)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d17a30e9-307b-42bb-95be-60d54f96b092",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_tax_addresses = (\n",
    "    simulated_taxes\n",
    "    # Can only link these to the other files if they have an SSN\n",
    "    .dropna(subset=[\"ssn\"])[[\"record_id\", \"ssn\"] + address_cols]\n",
    "    .assign(\n",
    "        **{c: lambda df, c=c: standardize_address_part(df[c]) for c in address_cols}\n",
    "    )\n",
    "    .pipe(df_ops.ensure_large_string_capacity)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edb08490-9c5d-4952-b12c-fb63b8fd3856",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_addresses_by_ssn, simulated_addresses_by_ssn_source_record_pairs = (\n",
    "    df_ops.persist(\n",
    "        dedupe_preserving_source_records(\n",
    "            simulated_tax_addresses,\n",
    "            simulated_taxes_source_record_pairs,\n",
    "            columns_to_dedupe=[\"ssn\"] + address_cols,\n",
    "            new_record_id_prefix=\"addresses_by_ssn\",\n",
    "        )\n",
    "    )\n",
    ")\n",
    "simulated_addresses_by_ssn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77f18f0f-35a7-46d3-98e6-3f029d758b0d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "num_addresses = (\n",
    "    df_ops.groupby_agg_small_groups(\n",
    "        simulated_addresses_by_ssn, by=\"ssn\", agg_func=lambda x: x.size()\n",
    "    )\n",
    "    .rename(\"size\")\n",
    "    .reset_index()\n",
    "    .sort_values(\"size\")\n",
    ")\n",
    "num_addresses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0289de32-9d4c-4c4f-a7c0-b85ed743c7ce",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Show some SSNs with a lot of address variation\n",
    "df_ops.compute(\n",
    "    pd.merge(\n",
    "        num_addresses.tail(10), simulated_addresses_by_ssn, on=\"ssn\", how=\"inner\"\n",
    "    ).sort_values(\"ssn\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d859fee-278c-4f8e-b959-527de4f13c48",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Rough estimate of how many rows we should have in our reference file, once we do this Cartesian product\n",
    "df_ops.compute(\n",
    "    len(simulated_name_dob_reference_file)\n",
    "    * df_ops.groupby_agg_small_groups(\n",
    "        simulated_addresses_by_ssn, by=\"ssn\", agg_func=lambda x: x.size()\n",
    "    ).mean()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7f8d3f4-eeda-468d-adf1-e6c6cf9aad8a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "(\n",
    "    simulated_geobase_reference_file,\n",
    "    simulated_geobase_reference_file_source_record_pairs,\n",
    ") = df_ops.persist(\n",
    "    merge_preserving_source_records(\n",
    "        [simulated_name_dob_reference_file, simulated_addresses_by_ssn],\n",
    "        [\n",
    "            simulated_name_dob_reference_file_source_record_pairs,\n",
    "            simulated_addresses_by_ssn_source_record_pairs,\n",
    "        ],\n",
    "        on=[\"ssn\"],\n",
    "        how=\"left\",\n",
    "        new_record_id_prefix=\"simulated_geobase_reference_file\",\n",
    "    )\n",
    ")\n",
    "simulated_geobase_reference_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c15c50c-7301-43cf-940d-6aff2814c0c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Actual number of rows\n",
    "len(simulated_geobase_reference_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9115e73-135d-46fe-8e09-a7526d688d05",
   "metadata": {},
   "source": [
    "## Track ground truth for reference files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b98c0689-1241-41c7-a429-1f9f9ab03ebd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def get_simulants_of_source_records(source_record_pairs, filter_record_ids=None):\n",
    "    if filter_record_ids is not None:\n",
    "        source_record_pairs = source_record_pairs.pipe(filter_record_ids)\n",
    "\n",
    "    result = (\n",
    "        df_ops.drop_duplicates(source_record_pairs)\n",
    "        .merge(source_record_ground_truth, on=\"source_record_id\")[\n",
    "            [\"record_id\", \"simulant_id\"]\n",
    "        ]\n",
    "        .pipe(df_ops.drop_duplicates)\n",
    "    )\n",
    "    result = result.merge(\n",
    "        df_ops.groupby_agg_small_groups(\n",
    "            result,\n",
    "            by=\"record_id\",\n",
    "            agg_func=lambda x: x.simulant_id.nunique().rename(\"n_unique_simulants\"),\n",
    "        ).reset_index(),\n",
    "        on=\"record_id\",\n",
    "        how=\"left\",\n",
    "    )\n",
    "    return result\n",
    "\n",
    "\n",
    "def get_ground_truth_for_records(source_record_pairs):\n",
    "    result = df_ops.persist(get_simulants_of_source_records(source_record_pairs))\n",
    "\n",
    "    result_nunique_describe = df_ops.compute(result[\"n_unique_simulants\"].describe())\n",
    "\n",
    "    if numpy.isclose(result_nunique_describe.loc[\"max\"], 1):\n",
    "        print(\"No collisions\")\n",
    "        return result\n",
    "\n",
    "    print(\"Collisions:\")\n",
    "    display(result_nunique_describe)\n",
    "\n",
    "    print(\"Simulated tax records with the most collisions:\")\n",
    "    most_collisions_record_id = (\n",
    "        result.sort_values(\"n_unique_simulants\", ascending=False)\n",
    "        .head(1)\n",
    "        .record_id.iloc[0]\n",
    "    )\n",
    "    most_collisions_source_record_ids = df_ops.compute(\n",
    "        df_ops.drop_duplicates(\n",
    "            source_record_pairs[\n",
    "                source_record_pairs.record_id == most_collisions_record_id\n",
    "            ][[\"source_record_id\"]]\n",
    "        )\n",
    "    )\n",
    "    tax_record_ids = most_collisions_source_record_ids.merge(\n",
    "        df_ops.compute(\n",
    "            df_ops.drop_duplicates(\n",
    "                simulated_taxes_source_record_pairs[\n",
    "                    simulated_taxes_source_record_pairs.source_record_id.isin(\n",
    "                        most_collisions_source_record_ids.source_record_id\n",
    "                    )\n",
    "                ][[\"record_id\", \"source_record_id\"]]\n",
    "            )\n",
    "        ),\n",
    "        on=\"source_record_id\",\n",
    "        how=\"inner\",\n",
    "    )\n",
    "    most_collisions_ground_truth = df_ops.compute(\n",
    "        source_record_ground_truth[\n",
    "            source_record_ground_truth.source_record_id.isin(\n",
    "                most_collisions_source_record_ids.source_record_id\n",
    "            )\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    most_collisions_tax_filings = tax_record_ids.merge(\n",
    "        df_ops.compute(\n",
    "            simulated_taxes[simulated_taxes.record_id.isin(tax_record_ids.record_id)],\n",
    "        ),\n",
    "        on=\"record_id\",\n",
    "        how=\"left\",\n",
    "    ).merge(\n",
    "        most_collisions_ground_truth,\n",
    "        on=\"source_record_id\",\n",
    "        how=\"left\",\n",
    "    )\n",
    "    display(most_collisions_tax_filings)\n",
    "\n",
    "    print(\"Simulated SSA records with the most collisions:\")\n",
    "    display(\n",
    "        df_ops.compute(\n",
    "            most_collisions_source_record_ids.merge(\n",
    "                df_ops.compute(\n",
    "                    simulated_ssa_numident[\n",
    "                        simulated_ssa_numident.record_id.isin(\n",
    "                            most_collisions_source_record_ids.source_record_id\n",
    "                        )\n",
    "                    ].rename(columns={\"record_id\": \"source_record_id\"})\n",
    "                ),\n",
    "                on=\"source_record_id\",\n",
    "                how=\"inner\",\n",
    "            ).merge(most_collisions_ground_truth, on=\"source_record_id\", how=\"left\")\n",
    "        )\n",
    "    )\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cbe7a99-116e-4626-b423-fd0d0f87cdda",
   "metadata": {},
   "source": [
    "### Simulated Census Numident"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86e6be4b-238d-4aa5-876c-0a72b38ff2c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(simulated_census_numident_source_record_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43eb7b26-6eda-4b87-b63e-3550c84fb473",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_census_numident_ground_truth = get_ground_truth_for_records(\n",
    "    simulated_census_numident_source_record_pairs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a7888b4-6e0b-4f76-9160-bdbf4334576f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_census_numident_ground_truth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14e1fa1c-6d37-43cf-b612-e869925e391c",
   "metadata": {},
   "source": [
    "### Simulated Alternate Name Numident"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cd90ac4-dc11-4b89-9772-16816d21dee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(simulated_alternate_name_numident_source_record_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da850e46-ba56-4eba-bbd1-de58b4bf307f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_alternate_name_numident_ground_truth = get_ground_truth_for_records(\n",
    "    simulated_alternate_name_numident_source_record_pairs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82a182bc-af0e-4ea4-a5b6-8eeb4c907f6d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_alternate_name_numident_ground_truth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed608b32-c590-40a8-b12a-93bc22b74ae2",
   "metadata": {},
   "source": [
    "### Alternate DOB Numident"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49f2c07a-29d3-42c2-a5cb-27d087a62cff",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(simulated_alternate_dob_numident_source_record_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11bea745-7445-4725-80e5-a7452ea924a1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_alternate_dob_numident_ground_truth = get_ground_truth_for_records(\n",
    "    simulated_alternate_dob_numident_source_record_pairs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1f40e9e-5d42-4b73-b731-bafcb826d471",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_alternate_dob_numident_ground_truth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a72beb7-0bf0-4531-bccc-9a89219322a7",
   "metadata": {},
   "source": [
    "### Simulated Name/DOB Reference File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64ea878b-1b7e-4714-b9ad-f0192d3c1e1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(simulated_name_dob_reference_file_source_record_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "538109c5-a088-4f33-a907-cfd6d1bd6368",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_name_dob_reference_file_ground_truth = get_ground_truth_for_records(\n",
    "    simulated_name_dob_reference_file_source_record_pairs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "997af054-23dc-4ea4-90fa-1a8777712f2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_name_dob_reference_file_ground_truth"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec1a0a88-b046-48a1-a230-d12a576c9be3",
   "metadata": {},
   "source": [
    "### Simulated GeoBase Reference File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c016c8c-3130-4b7a-a024-df4c003ebf44",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(simulated_geobase_reference_file_source_record_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bffe19a-e2a8-4624-9d16-79296a7d0481",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_geobase_reference_file_ground_truth = get_ground_truth_for_records(\n",
    "    simulated_geobase_reference_file_source_record_pairs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95cb8c5d-c919-4a93-9f6f-60c8df385e91",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "simulated_geobase_reference_file_ground_truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "794b921e-85b2-4107-b528-c00b0cbd703c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%xdel simulated_taxes\n",
    "%xdel simulated_taxes_source_record_pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26dbbbf8-ddc6-4ecb-9743-9036bd03c79e",
   "metadata": {},
   "source": [
    "### Get ground truth by SSN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19ad9a6e-df82-4482-8d49-2128a07a2a41",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_ssn_simulant_pairs = df_ops.persist(\n",
    "    df_ops.concat(\n",
    "        [\n",
    "            simulated_census_numident[[\"record_id\", \"ssn\"]].merge(\n",
    "                simulated_census_numident_ground_truth, on=\"record_id\"\n",
    "            ),\n",
    "            simulated_alternate_name_numident[[\"record_id\", \"ssn\"]].merge(\n",
    "                simulated_alternate_name_numident_ground_truth, on=\"record_id\"\n",
    "            ),\n",
    "            simulated_alternate_dob_numident[[\"record_id\", \"ssn\"]].merge(\n",
    "                simulated_alternate_dob_numident_ground_truth, on=\"record_id\"\n",
    "            ),\n",
    "            simulated_name_dob_reference_file[[\"record_id\", \"ssn\"]].merge(\n",
    "                simulated_name_dob_reference_file_ground_truth, on=\"record_id\"\n",
    "            ),\n",
    "            simulated_geobase_reference_file[[\"record_id\", \"ssn\"]].merge(\n",
    "                simulated_geobase_reference_file_ground_truth, on=\"record_id\"\n",
    "            ),\n",
    "        ]\n",
    "    )[[\"ssn\", \"simulant_id\"]].pipe(df_ops.drop_duplicates)\n",
    ")\n",
    "simulated_ssn_simulant_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a75ebbe1-f0a6-402f-8833-b7bdbeff591a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# How many simulated SSNs have different simulant IDs\n",
    "# contributing to them?\n",
    "df_ops.compute(\n",
    "    (\n",
    "        df_ops.groupby_agg_small_groups(\n",
    "            simulated_ssn_simulant_pairs,\n",
    "            by=\"ssn\",\n",
    "            agg_func=lambda x: x.simulant_id.nunique(),\n",
    "        )\n",
    "        > 1\n",
    "    ).sum()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbaef894-d1e4-4425-8fac-58c96ef7708d",
   "metadata": {},
   "source": [
    "## Save results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e55ff1f7-9332-419b-8a5b-306572d4cad5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "files = {\n",
    "    \"simulated_census_numident\": (\n",
    "        simulated_census_numident,\n",
    "        simulated_census_numident_ground_truth,\n",
    "    ),\n",
    "    \"simulated_alternate_name_numident\": (\n",
    "        simulated_alternate_name_numident,\n",
    "        simulated_alternate_name_numident_ground_truth,\n",
    "    ),\n",
    "    \"simulated_alternate_dob_numident\": (\n",
    "        simulated_alternate_dob_numident,\n",
    "        simulated_alternate_dob_numident_ground_truth,\n",
    "    ),\n",
    "    \"simulated_geobase_reference_file\": (\n",
    "        simulated_geobase_reference_file,\n",
    "        simulated_geobase_reference_file_ground_truth,\n",
    "    ),\n",
    "    \"simulated_name_dob_reference_file\": (\n",
    "        simulated_name_dob_reference_file,\n",
    "        simulated_name_dob_reference_file_ground_truth,\n",
    "    ),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ca11cf1-963c-448e-8961-261f9ddc2143",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_reference_files = [\n",
    "    simulated_census_numident,\n",
    "    simulated_alternate_name_numident,\n",
    "    simulated_alternate_dob_numident,\n",
    "    simulated_geobase_reference_file,\n",
    "    simulated_name_dob_reference_file,\n",
    "]\n",
    "# TODO: Rename the ssn column to explicitly include itins, since this is confusing\n",
    "simulated_all_ssns_itins_in_reference_files = df_ops.concat(\n",
    "    [df[[\"ssn\"]] for df in simulated_reference_files], ignore_index=True\n",
    ")\n",
    "simulated_ssn_to_pik = simulated_all_ssns_itins_in_reference_files.pipe(\n",
    "    df_ops.drop_duplicates\n",
    ").pipe(add_unique_id_col, col_name=\"pik\")[[\"ssn\", \"pik\"]]\n",
    "simulated_ssn_to_pik"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be53e5b9-2820-4d74-b6c0-0c6cd696a8c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_pik_simulant_pairs = simulated_ssn_simulant_pairs.merge(\n",
    "    simulated_ssn_to_pik, on=\"ssn\", how=\"inner\"\n",
    ")[[\"pik\", \"simulant_id\"]]\n",
    "simulated_pik_simulant_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "977d392f-b0af-4e8e-b644-e58949768e35",
   "metadata": {},
   "outputs": [],
   "source": [
    "%xdel simulated_ssn_simulant_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64cbf611-0ac1-48fa-bd91-efe3cd8a23a1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for file_name, (file, ground_truth) in files.items():\n",
    "    file = df_ops.persist(file.merge(simulated_ssn_to_pik, on=\"ssn\", how=\"left\"))\n",
    "    assert df_ops.compute(file.pik.notnull().all())\n",
    "\n",
    "    save_file_with_ground_truth(file_name, file, ground_truth)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "834b4cdb-3b83-4f83-b53a-3a67b23e9ff6",
   "metadata": {},
   "outputs": [],
   "source": [
    "simulated_pik_simulant_pairs_path = (\n",
    "    f\"{output_dir}/{data_to_use}/simulated_pik_simulant_pairs.parquet\"\n",
    ")\n",
    "utils.remove_path(simulated_pik_simulant_pairs_path)\n",
    "simulated_pik_simulant_pairs.to_parquet(simulated_pik_simulant_pairs_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aa94833-9286-427e-9d9a-257ba17bf678",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(datetime.datetime.now())"
   ]
  }
 ],
 "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",
   "version": "3.10.14"
  },
  "nbdime-conflicts": {
   "local_diff": [
    {
     "diff": [
      {
       "diff": [
        {
         "diff": [
          {
           "key": 24,
           "op": "addrange",
           "valuelist": "3"
          },
          {
           "key": 24,
           "length": 7,
           "op": "removerange"
          }
         ],
         "key": 0,
         "op": "patch"
        }
       ],
       "key": "display_name",
       "op": "patch"
      },
      {
       "diff": [
        {
         "diff": [
          {
           "key": 24,
           "op": "addrange",
           "valuelist": "3"
          },
          {
           "key": 24,
           "length": 7,
           "op": "removerange"
          }
         ],
         "key": 0,
         "op": "patch"
        }
       ],
       "key": "name",
       "op": "patch"
      }
     ],
     "key": "kernelspec",
     "op": "patch"
    }
   ],
   "remote_diff": [
    {
     "diff": [
      {
       "diff": [
        {
         "diff": [
          {
           "key": 24,
           "length": 6,
           "op": "removerange"
          }
         ],
         "key": 0,
         "op": "patch"
        }
       ],
       "key": "display_name",
       "op": "patch"
      },
      {
       "diff": [
        {
         "diff": [
          {
           "key": 24,
           "length": 6,
           "op": "removerange"
          }
         ],
         "key": 0,
         "op": "patch"
        }
       ],
       "key": "name",
       "op": "patch"
      }
     ],
     "key": "kernelspec",
     "op": "patch"
    }
   ]
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}