tutorials/balance_transformations_and_formulas.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "collapsed": true,
        "originalKey": "b09e243b-4de0-4211-aee8-80ed3b05499a",
        "showInput": false
      },
      "source": [
        "# balance: transformations and formulas\n",
        "\n",
        "This tutorial focuses on the ways transformations, formulas and penalty can be included in your pre-processing of the coveriates before adjusting for them."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "72611c31-352b-40f0-93f2-49fcbf582b50",
        "showInput": false
      },
      "source": [
        "## Example dataset - preparing the objects\n",
        "\n",
        "The following is a toy simulated dataset.\n",
        "\n",
        "For a more basic walkthrough of the elements in the next code block, please take a look at the tutorial: [balance Quickstart: Analyzing and adjusting the bias on a simulated toy dataset](https://import-balance.org/docs/tutorials/quickstart/)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609343074,
        "executionStopTime": 1669609343457,
        "originalKey": "a379ae12-4f6e-47ae-85c8-c44ffeced3d4",
        "requestMsgId": "db56900f-d26b-4d2b-b924-cfc82e6a43e1",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "from balance import load_data\n",
        "target_df, sample_df = load_data()\n",
        "from balance import Sample\n",
        "sample = Sample.from_frame(sample_df, outcome_columns=[\"happiness\"])\n",
        "target = Sample.from_frame(target_df, outcome_columns=[\"happiness\"])\n",
        "sample_with_target = sample.set_target(target)\n",
        "sample_with_target"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "8e655fe2-bd3a-45ef-831d-db1f53ae9d47",
        "showInput": false
      },
      "source": [
        "# Transformations\n",
        "\n",
        "## Basic usage: manipulating existing variables"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "583f5117-d3a0-4145-b9fa-9e372131e503",
        "showInput": false
      },
      "source": [
        "When trying to understand what an adjustment does, we can look at the model_coef items collected from the diagnostics method."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609343652,
        "executionStopTime": 1669609347257,
        "originalKey": "c10e715c-84ad-416e-a2dc-38e76af17144",
        "requestMsgId": "ec109b15-c89b-43da-b8e5-43a95716265c",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "adjusted = sample_with_target.adjust(\n",
        "    # method=\"ipw\", # default method\n",
        "    # transformations=None,\n",
        "    # formula=None,\n",
        "    # penalty_factor=None, # all 1s\n",
        "    # max_de=None,\n",
        ")\n",
        "adj_diag = adjusted.diagnostics()\n",
        "adj_diag.query(\"metric == 'model_coef'\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "05b1bd1a-b590-4634-9cbe-7cb252fd5882",
        "showInput": false
      },
      "source": [
        "As we can see from the glm coefficients, the age and gender groups got an extra NA column. And the income variable is bucketed into 10 buckets.\n",
        "\n",
        "We can change these defaults by deciding on the specific transformation we want.\n",
        "\n",
        "Let's start with NO transformations.\n",
        "\n",
        "The transformation argument accepts either a dict or None. None indicates no transformations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609347595,
        "executionStopTime": 1669609351280,
        "originalKey": "0462dc3e-b1e0-4c96-89b9-c6cb45379c25",
        "requestMsgId": "00817f39-3f2d-4914-b652-fa42e79408db",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "adjusted = sample_with_target.adjust(\n",
        "    # method=\"ipw\",\n",
        "    transformations=None,\n",
        "    # formula=formula,\n",
        "    # penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "adj_diag = adjusted.diagnostics()\n",
        "adj_diag.query(\"metric == 'model_coef'\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "d8e36900-2380-4ca0-97bf-c6a50dca13ff",
        "showInput": false
      },
      "source": [
        "In this setting, income was treated as a numeric variable, with no transformations (e.g.: bucketing) on it.\n",
        "Regardless of the transformations, the model matrix made sure to turn the gender and age_group into dummy variables (including a column for NA).\n",
        "\n",
        "\n",
        "Next we can fit a simple transformation.\n",
        "\n",
        "Let's say we wanted to bucket age_groups groups that are smaller than 25% of the data, and use different bucketing on income, here is how we'd do it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609351521,
        "executionStopTime": 1669609354592,
        "originalKey": "83eefdfe-c792-4ca6-aea9-c59c8a33b2fd",
        "requestMsgId": "1cde8a1a-50aa-4b0d-a73e-6da7c200d0bf",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "from balance.util import fct_lump, quantize\n",
        "\n",
        "transformations = {\n",
        "    \"age_group\": lambda x: fct_lump(x, 0.25),\n",
        "    \"gender\": lambda x: x,\n",
        "    \"income\": lambda x: quantize(x.fillna(x.mean()), q=3),\n",
        "}\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    # method=\"ipw\",\n",
        "    transformations=transformations,\n",
        "    # formula=formula,\n",
        "    # penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "adj_diag = adjusted.diagnostics()\n",
        "adj_diag.query(\"metric == 'model_coef'\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "41379897-7487-4ff6-821f-ce66ceb363bc",
        "showInput": false
      },
      "source": [
        "As we can see - we managed to change the bucket sizes of income to have only 3 buckets, and we lumped the age_group to two groups (and collapsed together \"small\" buckets into the _lumped_other bucket).\n",
        "\n",
        "Lastly, notice that if we omit a variable from transformations, it will not be available for the model construction (This behavior might change in the future)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609354770,
        "executionStopTime": 1669609356802,
        "originalKey": "9a0b7eab-1a74-476e-878e-f34b3bb77376",
        "requestMsgId": "ecc2096b-08ed-4f76-998e-762f6b2c4c0a",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "transformations = {\n",
        "    # \"age_group\": lambda x: fct_lump(x, 0.25),\n",
        "    \"gender\": lambda x: x,\n",
        "    # \"income\": lambda x: quantize(x.fillna(x.mean()), q=3),\n",
        "}\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    # method=\"ipw\",\n",
        "    transformations=transformations,\n",
        "    # formula=formula,\n",
        "    # penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "adj_diag = adjusted.diagnostics()\n",
        "adj_diag.query(\"metric == 'model_coef'\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "6f3b10c0-b649-45ea-b18c-50e5e7f79e23",
        "showInput": false
      },
      "source": [
        "As we can see, only gender was included in the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609356964,
        "executionStopTime": 1669609356972,
        "originalKey": "62c7954b-9f63-4244-a214-b253623da71b",
        "requestMsgId": "2298603a-2546-4595-80bc-459b8f1e73fc",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "# TODO: add more examples about how add_na works\n",
        "# TODO: add more examples about rare values in categorical variables and how they are grouped together. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "d2d22978-1f8f-460e-a80d-6f511a2144aa",
        "showInput": false
      },
      "source": [
        "## Creating new variables\n",
        "\n",
        "In the next example we will create several new transformations of income. \n",
        "\n",
        "The info gives information on which variables were added, which were transformed, and what is the final variables in the output.\n",
        "\n",
        "The x in the lambda function can have one of two meanings:\n",
        "1. When the keys in the dict match the exact names of the variables in the DataFrame (e.g.: \"income\"), then the lambda function treats x as the pandas.Series of that variable.\n",
        "2. If the name of the key does NOT exist in the DataFrame (e.g.: \"income_squared\"), then x will become the DataFrame of the data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609357125,
        "executionStopTime": 1669609362009,
        "originalKey": "d032d811-f409-4f68-ab99-fdeb8d556f5d",
        "requestMsgId": "806707e9-8ab5-43af-90df-4f264321619a",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "from balance.util import fct_lump, quantize\n",
        "\n",
        "transformations = {\n",
        "    \"age_group\": lambda x: x,\n",
        "    \"gender\": lambda x: x,\n",
        "    \"income\": lambda x: x,\n",
        "    \"income_squared\": lambda x: x.income**2,\n",
        "    \"income_buckets\": lambda x: quantize(x.income.fillna(x.income.mean()), q=3),\n",
        "}\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    # method=\"ipw\",\n",
        "    transformations=transformations,\n",
        "    # formula=formula,\n",
        "    # penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "adj_diag = adjusted.diagnostics()\n",
        "adj_diag.query(\"metric == 'model_coef'\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "2c6e9a86-19e5-4881-8516-c595426fd1bd",
        "showInput": false
      },
      "source": [
        "# Formula\n",
        "\n",
        "The formula can accept a list of strings indicating how to combine the transformed variables together. It follows [the formula notation from patsy](https://patsy.readthedocs.io/en/latest/formulas.html).\n",
        "\n",
        "For example, we can have an interaction between age_group and gender:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609362246,
        "executionStopTime": 1669609365906,
        "originalKey": "4cf5e54c-19d8-4eed-84a2-2496868c5df5",
        "requestMsgId": "46bea74d-e966-4410-bc90-1c89cd724c00",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "from balance.util import fct_lump_by, quantize\n",
        "\n",
        "transformations = {\n",
        "    \"age_group\": lambda x: x,\n",
        "    \"gender\": lambda x: x,\n",
        "    \"income\": lambda x: quantize(x.fillna(x.mean()), q=20),\n",
        "}\n",
        "formula = [\"age_group * gender\"]\n",
        "# the penalty is per elemnt in the list of formula:\n",
        "# penalty_factor = [0.1, 0.1, 0.1]\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    method=\"ipw\",\n",
        "    transformations=transformations,\n",
        "    formula=formula,\n",
        "    # penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "adj_diag = adjusted.diagnostics()\n",
        "adj_diag.query(\"metric == 'model_coef'\")\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "f15a3528-3e8e-480b-825f-86c82f2cc710",
        "showInput": false
      },
      "source": [
        "As we can see, the formula makes it so that we have combinations of age_group and gender, as well as a main effects of age_group and gender. Since income was not in the formula, it is not included in the model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "74074181-c1e2-45b3-9fcd-781ae760f190",
        "showInput": false
      },
      "source": [
        "# Formula and penalty_factor\n",
        "\n",
        "The formula can be provided as several strings, and then the penalty factor can indicate how much the model should focus to adjust to that element of the formula. Larger penalty factors means that element will be less corrected.\n",
        "\n",
        "The next two examples shows how in one case we focus on correcting for income, and in the second case we focus to correct for age and gender."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609366055,
        "executionStopTime": 1669609368854,
        "originalKey": "83034e4c-f94e-40ca-9a12-7425f4ca010b",
        "requestMsgId": "efd734cc-08e0-4eee-b754-bbe0a0a67e06",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "transformations = {\n",
        "    \"age_group\": lambda x: x,\n",
        "    \"gender\": lambda x: x,\n",
        "    \"income\": lambda x: x,\n",
        "}\n",
        "formula = [\"age_group + gender\", \"income\"]\n",
        "# the penalty is per elemnt in the list of formula:\n",
        "penalty_factor = [10, 0.1]\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    method=\"ipw\",\n",
        "    transformations=transformations,\n",
        "    formula=formula,\n",
        "    penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "adj_diag = adjusted.diagnostics()\n",
        "adj_diag.query(\"metric == 'model_coef'\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "4552f04e-c72e-460a-916b-f924bbecc043",
        "showInput": false
      },
      "source": [
        "The above example corrected more to income. As we can see, age and gender got 0 correction (since their penalty was so high). Let's now over correct for age and gender:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609369011,
        "executionStopTime": 1669609371780,
        "originalKey": "ab7d1b29-a2fc-449b-8b85-9360540f7bdf",
        "requestMsgId": "26d537bb-12b5-4a1a-b919-b632c08806cc",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "transformations = {\n",
        "    \"age_group\": lambda x: x,\n",
        "    \"gender\": lambda x: x,\n",
        "    \"income\": lambda x: x,\n",
        "}\n",
        "formula = [\"age_group + gender\", \"income\"]\n",
        "# the penalty is per elemnt in the list of formula:\n",
        "penalty_factor = [0.1, 10]  # this is flipped\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    method=\"ipw\",\n",
        "    transformations=transformations,\n",
        "    formula=formula,\n",
        "    penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "adj_diag = adjusted.diagnostics()\n",
        "adj_diag.query(\"metric == 'model_coef'\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "168e0e8c-1b51-4a0a-8309-b8ea1a137c8b",
        "showInput": false
      },
      "source": [
        "In the above case, income basically got 0 correction.\n",
        "\n",
        "We can add two versions of income, and give each of them a higher penalty than the age and gender:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609371954,
        "executionStopTime": 1669609376038,
        "originalKey": "ee24f18c-2d2a-493c-8969-d2b0bbdddb5e",
        "requestMsgId": "76e1c453-1db7-49a7-8c46-fb244b1fb1f5",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "from balance.util import fct_lump_by, quantize\n",
        "\n",
        "transformations = {\n",
        "    \"age_group\": lambda x: x,\n",
        "    \"gender\": lambda x: x,\n",
        "    \"income\": lambda x: x,\n",
        "    \"income_buckets\": lambda x: quantize(x.income.fillna(x.mean()), q=4),\n",
        "}\n",
        "formula = [\"age_group + gender\", \"income\", \"income_buckets\"]\n",
        "# the penalty is per elemnt in the list of formula:\n",
        "penalty_factor = [1, 2, 2]\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    method=\"ipw\",\n",
        "    transformations=transformations,\n",
        "    formula=formula,\n",
        "    penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "adj_diag = adjusted.diagnostics()\n",
        "adj_diag.query(\"metric == 'model_coef'\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "84df8ccc-d0f5-4f7f-be08-12205658de7e",
        "showInput": false
      },
      "source": [
        "Another way is to create a formula for several variations of each variable, and give each a penalty of 1. For example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609376232,
        "executionStopTime": 1669609380501,
        "originalKey": "8162c4bb-6ed3-4b5f-a226-a9fc305676d4",
        "requestMsgId": "15962ce8-8d7f-4511-9a12-9d882459c179",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "from balance.util import fct_lump_by, quantize\n",
        "\n",
        "transformations = {\n",
        "    \"age_group\": lambda x: x,\n",
        "    \"gender\": lambda x: x,\n",
        "    \"income\": lambda x: x,\n",
        "    \"income_buckets\": lambda x: quantize(x.income.fillna(x.mean()), q=4),\n",
        "}\n",
        "formula = [\"age_group\", \"gender\", \"income + income_buckets\"]\n",
        "# the penalty is per elemnt in the list of formula:\n",
        "penalty_factor = [1, 1, 1]\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    method=\"ipw\",\n",
        "    transformations=transformations,\n",
        "    formula=formula,\n",
        "    penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "adj_diag = adjusted.diagnostics()\n",
        "adj_diag.query(\"metric == 'model_coef'\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "15170c34-1ff5-477f-a9fd-de5fee3238b0",
        "showInput": false
      },
      "source": [
        "# The impact of trnasformations and formulas\n",
        "\n",
        "## ipw\n",
        "\n",
        "Using the above can have an impact on the final design effect, ASMD, and outcome. Here are several simple examples."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609380664,
        "executionStopTime": 1669609384993,
        "originalKey": "a4e9eb09-8b27-4c9b-9678-e16240705f38",
        "requestMsgId": "2e25549c-4b70-4d8e-822f-d99983ed3b67",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "# Defaults from the package\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "print(adjusted.summary())\n",
        "print(adjusted.outcomes().summary())\n",
        "adjusted.covars().plot(library = \"seaborn\", dist_type = \"kde\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609385192,
        "executionStopTime": 1669609390221,
        "originalKey": "b7c6e9c0-d82f-49ec-a11c-75d83630d8c7",
        "requestMsgId": "43602680-1dda-4060-87ff-aae3a8d31dd1",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "# No transformations at all\n",
        "\n",
        "# transformations = None is just like using:\n",
        "# transformations = {\n",
        "#     \"age_group\": lambda x: x,\n",
        "#     \"gender\": lambda x: x,\n",
        "#     \"income\": lambda x: x,\n",
        "# }\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    method=\"ipw\",\n",
        "    transformations=None,\n",
        "    # formula=formula,\n",
        "    # penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "print(adjusted.summary())\n",
        "print(adjusted.outcomes().summary())\n",
        "adjusted.covars().plot(library = \"seaborn\", dist_type = \"kde\")\n",
        "\n",
        "# slightly smaller design effect, slightly better ASMD reduction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609390371,
        "executionStopTime": 1669609394673,
        "originalKey": "ce89fc82-887f-43a7-8b99-1bb7b702cc20",
        "requestMsgId": "a876ee51-2015-4a70-9ec0-284ba42e1f10",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "# No transformations at all\n",
        "transformations = None\n",
        "# But passing a squared term of income to the formula:\n",
        "formula = [\"age_group + gender + income + income**2\"]\n",
        "# the penalty is per elemnt in the list of formula:\n",
        "# penalty_factor = [1]\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    method=\"ipw\",\n",
        "    transformations=transformations,\n",
        "    formula=formula,\n",
        "    # penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "print(adjusted.summary())\n",
        "print(adjusted.outcomes().summary())\n",
        "adjusted.covars().plot(library = \"seaborn\", dist_type = \"kde\")\n",
        "\n",
        "# Adding income**2 to the formula led to lower Deff but also lower ASMD reduction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609394826,
        "executionStopTime": 1669609400912,
        "originalKey": "c4d31c9b-66a6-4ded-919d-7fb5e0945cda",
        "requestMsgId": "c72f9725-a0e0-4991-9e5c-6b09b5f3b9f3",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "transformations = {\n",
        "    \"age_group\": lambda x: x,\n",
        "    \"gender\": lambda x: x,\n",
        "    \"income\": lambda x: x,\n",
        "    \"income_buckets\": lambda x: quantize(x.income.fillna(x.income.mean()), q=20),\n",
        "}\n",
        "formula = [\"age_group + gender\", \"income_buckets\"]\n",
        "# the penalty is per elemnt in the list of formula:\n",
        "penalty_factor = [1, 0.1]\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    method=\"ipw\",\n",
        "    transformations=transformations,\n",
        "    formula=formula,\n",
        "    penalty_factor=penalty_factor,\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "print(adjusted.summary())\n",
        "print(adjusted.outcomes().summary())\n",
        "adjusted.covars().plot(library = \"seaborn\", dist_type = \"kde\")\n",
        "\n",
        "# By adding income_buckets and using it instead of income, as well as putting more weight in it in terms of penalty\n",
        "# we managed to correct income quite well, but at the expense of age and gender."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609401059,
        "executionStopTime": 1669609401069,
        "originalKey": "c45c122f-e795-46c7-90ad-b9d36571d508",
        "requestMsgId": "c7359660-8eca-4f26-b4f6-da9b164ad8f4",
        "showInput": true
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "customInput": null,
        "originalKey": "b23f9d06-8425-4d7a-b32f-26fd62257b0f",
        "showInput": false
      },
      "source": [
        "## CBPS\n",
        "\n",
        "Let's see if we can improve on CBPS a bit."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609401234,
        "executionStopTime": 1669609410847,
        "originalKey": "ca83ebdb-efe9-47ba-bc97-736af4b0cbbb",
        "requestMsgId": "9be6d7e3-9020-4bad-8281-b6abef2ede7a",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "# Defaults from the package\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    method = \"cbps\",\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "print(adjusted.summary())\n",
        "print(adjusted.outcomes().summary())\n",
        "adjusted.covars().plot(library = \"seaborn\", dist_type = \"kde\")\n",
        "\n",
        "# CBPS already corrects a lot. Let's see if we can make it correct a tiny bit more."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609410997,
        "executionStopTime": 1669609420455,
        "originalKey": "98f9c652-c0ec-46e0-b831-7e98d7100e7f",
        "requestMsgId": "c3f92542-2f14-486b-b347-5c3314c2d39b",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "# No transformations at all\n",
        "transformations = {\n",
        "    \"age_group\": lambda x: x,\n",
        "    \"gender\": lambda x: x,\n",
        "    # \"income\": lambda x: x,\n",
        "    \"income_log\": lambda x: np.log(x.income.fillna(x.income.mean())),\n",
        "    \"income_buckets\": lambda x: quantize(x.income.fillna(x.income.mean()), q=5),\n",
        "}\n",
        "formula = [\"age_group + gender + income_log * income_buckets\"]\n",
        "\n",
        "adjusted = sample_with_target.adjust(\n",
        "    method=\"cbps\",\n",
        "    transformations=transformations,\n",
        "    formula=formula,\n",
        "    # penalty_factor=penalty_factor, # CBPS seems to ignore the penalty factor.\n",
        "    # max_de=None,\n",
        ")\n",
        "\n",
        "print(adjusted.summary())\n",
        "print(adjusted.outcomes().summary())\n",
        "adjusted.covars().plot(library=\"seaborn\", dist_type=\"kde\")\n",
        "\n",
        "# Trying various transformations gives slightly different results (some effect on the outcome, Deff and ASMD) - but nothing too major here."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false,
        "customInput": null,
        "customOutput": null,
        "executionStartTime": 1669609420600,
        "executionStopTime": 1669609420608,
        "originalKey": "debbdc2e-63ed-4cc6-85f4-e5a9ae522cfa",
        "requestMsgId": "723571a9-86e1-400b-895f-6414e49d4368",
        "showInput": true
      },
      "outputs": [],
      "source": [
        "# Sessions info\n",
        "import session_info\n",
        "session_info.show(html=False, dependencies=True)"
      ]
    }
  ],
  "metadata": {
    "custom": {
      "cells": [],
      "metadata": {
        "bento_stylesheets": {
          "bento/extensions/flow/main.css": true,
          "bento/extensions/kernel_selector/main.css": true,
          "bento/extensions/kernel_ui/main.css": true,
          "bento/extensions/new_kernel/main.css": true,
          "bento/extensions/system_usage/main.css": true,
          "bento/extensions/theme/main.css": true
        },
        "captumWidgetMessage": {},
        "dataExplorerConfig": {},
        "fileHeader": "",
        "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.6.3rc1+"
        },
        "outputWidgetContext": {}
      },
      "nbformat": 4,
      "nbformat_minor": 2
    },
    "indentAmount": 2
  },
  "nbformat": 4,
  "nbformat_minor": 2
}