nics.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "671ecba2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final shapes - X: (2939, 22), y: (2939, 3)\n",
      "Positive class rate: 0.084\n"
     ]
    },
    {
     "data": {
      "application/vnd.plotly.v1+json": {
       "config": {
        "plotlyServerURL": "https://plot.ly"
       },
       "data": [
        {
         "alignmentgroup": "True",
         "hovertemplate": "variable=0<br>age=%{x}<br>value=%{y}<extra></extra>",
         "legendgroup": "0",
         "marker": {
          "color": "#2C6496",
          "pattern": {
           "shape": ""
          }
         },
         "name": "0",
         "offsetgroup": "0",
         "orientation": "v",
         "showlegend": true,
         "textposition": "auto",
         "type": "bar",
         "x": [
          16,
          17,
          18,
          19,
          20,
          21,
          22,
          23,
          24,
          25,
          26,
          27,
          28,
          29,
          30,
          31,
          32,
          33,
          34,
          35,
          36,
          37,
          38,
          39,
          40,
          41,
          42,
          43,
          44,
          45,
          46,
          47,
          48,
          49,
          50,
          51,
          52,
          53,
          54,
          55,
          56,
          57,
          58,
          59,
          60,
          61,
          62,
          63,
          64,
          65,
          66,
          67,
          68,
          69,
          70
         ],
         "xaxis": "x",
         "y": [
          0.12272046230208689,
          0.33081985439156475,
          0.1516752367078785,
          0.3772873272728147,
          0.16468009103932427,
          0.08591656144753525,
          0.07466549719671173,
          0.15777115746405443,
          0.055721381709675555,
          0.10131817932298616,
          0,
          0,
          0.04002040916023905,
          0.016137188104571,
          0,
          0.1402600892175181,
          0.04525813220638878,
          0.1608129708654496,
          0.04511183714501131,
          0.05939275941906967,
          0.03490604810205659,
          0.1154995717236862,
          0.13998911133370662,
          0,
          0,
          0.06665570148965451,
          0,
          0.060169527622779745,
          0.003022383156419279,
          0,
          0,
          0,
          0.013285449449810686,
          0.046059949355663106,
          0.02664482966877916,
          0.06620395237455799,
          0.04795735554207538,
          0.034506853345332454,
          0.04882587428663965,
          0.05067503043217135,
          0.025179332381747226,
          0.05983957519211712,
          0.06623991331869636,
          0.0540075081766612,
          0.17547062106758676,
          0.034005776823203,
          0.05179676172627123,
          0.07260665569881053,
          0.11661015350268067,
          0.09148520299870647,
          0.20846794479364292,
          0.021167258221747823,
          0.05639080905955598,
          0.030691090201591165,
          0.04208955308177593
         ],
         "yaxis": "y"
        }
       ],
       "layout": {
        "annotations": [
         {
          "showarrow": false,
          "text": "Source: PolicyEngine UK tax-benefit microsimulation model (version 2.40.2)",
          "x": 0,
          "xanchor": "left",
          "xref": "paper",
          "y": -0.2,
          "yanchor": "bottom",
          "yref": "paper"
         }
        ],
        "barmode": "relative",
        "font": {
         "color": "black",
         "family": "Roboto Serif"
        },
        "height": 600,
        "images": [
         {
          "sizex": 0.15,
          "sizey": 0.15,
          "source": "https://raw.githubusercontent.com/PolicyEngine/policyengine-app/master/src/images/logos/policyengine/blue.png",
          "x": 1.1,
          "xanchor": "right",
          "xref": "paper",
          "y": -0.2,
          "yanchor": "bottom",
          "yref": "paper"
         }
        ],
        "legend": {
         "title": {
          "text": "variable"
         },
         "tracegroupgap": 0
        },
        "margin": {
         "b": 120,
         "l": 120,
         "r": 120,
         "t": 120
        },
        "modebar": {
         "activecolor": "#F4F4F4",
         "bgcolor": "#F4F4F4",
         "color": "#F4F4F4"
        },
        "paper_bgcolor": "#F4F4F4",
        "plot_bgcolor": "#F4F4F4",
        "showlegend": false,
        "template": {
         "data": {
          "bar": [
           {
            "error_x": {
             "color": "#2a3f5f"
            },
            "error_y": {
             "color": "#2a3f5f"
            },
            "marker": {
             "line": {
              "color": "white",
              "width": 0.5
             },
             "pattern": {
              "fillmode": "overlay",
              "size": 10,
              "solidity": 0.2
             }
            },
            "type": "bar"
           }
          ],
          "barpolar": [
           {
            "marker": {
             "line": {
              "color": "white",
              "width": 0.5
             },
             "pattern": {
              "fillmode": "overlay",
              "size": 10,
              "solidity": 0.2
             }
            },
            "type": "barpolar"
           }
          ],
          "carpet": [
           {
            "aaxis": {
             "endlinecolor": "#2a3f5f",
             "gridcolor": "#C8D4E3",
             "linecolor": "#C8D4E3",
             "minorgridcolor": "#C8D4E3",
             "startlinecolor": "#2a3f5f"
            },
            "baxis": {
             "endlinecolor": "#2a3f5f",
             "gridcolor": "#C8D4E3",
             "linecolor": "#C8D4E3",
             "minorgridcolor": "#C8D4E3",
             "startlinecolor": "#2a3f5f"
            },
            "type": "carpet"
           }
          ],
          "choropleth": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "type": "choropleth"
           }
          ],
          "contour": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "contour"
           }
          ],
          "contourcarpet": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "type": "contourcarpet"
           }
          ],
          "heatmap": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "heatmap"
           }
          ],
          "heatmapgl": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "heatmapgl"
           }
          ],
          "histogram": [
           {
            "marker": {
             "pattern": {
              "fillmode": "overlay",
              "size": 10,
              "solidity": 0.2
             }
            },
            "type": "histogram"
           }
          ],
          "histogram2d": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "histogram2d"
           }
          ],
          "histogram2dcontour": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "histogram2dcontour"
           }
          ],
          "mesh3d": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "type": "mesh3d"
           }
          ],
          "parcoords": [
           {
            "line": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "parcoords"
           }
          ],
          "pie": [
           {
            "automargin": true,
            "type": "pie"
           }
          ],
          "scatter": [
           {
            "fillpattern": {
             "fillmode": "overlay",
             "size": 10,
             "solidity": 0.2
            },
            "type": "scatter"
           }
          ],
          "scatter3d": [
           {
            "line": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scatter3d"
           }
          ],
          "scattercarpet": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scattercarpet"
           }
          ],
          "scattergeo": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scattergeo"
           }
          ],
          "scattergl": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scattergl"
           }
          ],
          "scattermapbox": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scattermapbox"
           }
          ],
          "scatterpolar": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scatterpolar"
           }
          ],
          "scatterpolargl": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scatterpolargl"
           }
          ],
          "scatterternary": [
           {
            "marker": {
             "colorbar": {
              "outlinewidth": 0,
              "ticks": ""
             }
            },
            "type": "scatterternary"
           }
          ],
          "surface": [
           {
            "colorbar": {
             "outlinewidth": 0,
             "ticks": ""
            },
            "colorscale": [
             [
              0,
              "#0d0887"
             ],
             [
              0.1111111111111111,
              "#46039f"
             ],
             [
              0.2222222222222222,
              "#7201a8"
             ],
             [
              0.3333333333333333,
              "#9c179e"
             ],
             [
              0.4444444444444444,
              "#bd3786"
             ],
             [
              0.5555555555555556,
              "#d8576b"
             ],
             [
              0.6666666666666666,
              "#ed7953"
             ],
             [
              0.7777777777777778,
              "#fb9f3a"
             ],
             [
              0.8888888888888888,
              "#fdca26"
             ],
             [
              1,
              "#f0f921"
             ]
            ],
            "type": "surface"
           }
          ],
          "table": [
           {
            "cells": {
             "fill": {
              "color": "#EBF0F8"
             },
             "line": {
              "color": "white"
             }
            },
            "header": {
             "fill": {
              "color": "#C8D4E3"
             },
             "line": {
              "color": "white"
             }
            },
            "type": "table"
           }
          ]
         },
         "layout": {
          "annotationdefaults": {
           "arrowcolor": "#2a3f5f",
           "arrowhead": 0,
           "arrowwidth": 1
          },
          "autotypenumbers": "strict",
          "coloraxis": {
           "colorbar": {
            "outlinewidth": 0,
            "ticks": ""
           }
          },
          "colorscale": {
           "diverging": [
            [
             0,
             "#8e0152"
            ],
            [
             0.1,
             "#c51b7d"
            ],
            [
             0.2,
             "#de77ae"
            ],
            [
             0.3,
             "#f1b6da"
            ],
            [
             0.4,
             "#fde0ef"
            ],
            [
             0.5,
             "#f7f7f7"
            ],
            [
             0.6,
             "#e6f5d0"
            ],
            [
             0.7,
             "#b8e186"
            ],
            [
             0.8,
             "#7fbc41"
            ],
            [
             0.9,
             "#4d9221"
            ],
            [
             1,
             "#276419"
            ]
           ],
           "sequential": [
            [
             0,
             "#0d0887"
            ],
            [
             0.1111111111111111,
             "#46039f"
            ],
            [
             0.2222222222222222,
             "#7201a8"
            ],
            [
             0.3333333333333333,
             "#9c179e"
            ],
            [
             0.4444444444444444,
             "#bd3786"
            ],
            [
             0.5555555555555556,
             "#d8576b"
            ],
            [
             0.6666666666666666,
             "#ed7953"
            ],
            [
             0.7777777777777778,
             "#fb9f3a"
            ],
            [
             0.8888888888888888,
             "#fdca26"
            ],
            [
             1,
             "#f0f921"
            ]
           ],
           "sequentialminus": [
            [
             0,
             "#0d0887"
            ],
            [
             0.1111111111111111,
             "#46039f"
            ],
            [
             0.2222222222222222,
             "#7201a8"
            ],
            [
             0.3333333333333333,
             "#9c179e"
            ],
            [
             0.4444444444444444,
             "#bd3786"
            ],
            [
             0.5555555555555556,
             "#d8576b"
            ],
            [
             0.6666666666666666,
             "#ed7953"
            ],
            [
             0.7777777777777778,
             "#fb9f3a"
            ],
            [
             0.8888888888888888,
             "#fdca26"
            ],
            [
             1,
             "#f0f921"
            ]
           ]
          },
          "colorway": [
           "#636efa",
           "#EF553B",
           "#00cc96",
           "#ab63fa",
           "#FFA15A",
           "#19d3f3",
           "#FF6692",
           "#B6E880",
           "#FF97FF",
           "#FECB52"
          ],
          "font": {
           "color": "#2a3f5f"
          },
          "geo": {
           "bgcolor": "white",
           "lakecolor": "white",
           "landcolor": "white",
           "showlakes": true,
           "showland": true,
           "subunitcolor": "#C8D4E3"
          },
          "hoverlabel": {
           "align": "left"
          },
          "hovermode": "closest",
          "mapbox": {
           "style": "light"
          },
          "paper_bgcolor": "white",
          "plot_bgcolor": "white",
          "polar": {
           "angularaxis": {
            "gridcolor": "#EBF0F8",
            "linecolor": "#EBF0F8",
            "ticks": ""
           },
           "bgcolor": "white",
           "radialaxis": {
            "gridcolor": "#EBF0F8",
            "linecolor": "#EBF0F8",
            "ticks": ""
           }
          },
          "scene": {
           "xaxis": {
            "backgroundcolor": "white",
            "gridcolor": "#DFE8F3",
            "gridwidth": 2,
            "linecolor": "#EBF0F8",
            "showbackground": true,
            "ticks": "",
            "zerolinecolor": "#EBF0F8"
           },
           "yaxis": {
            "backgroundcolor": "white",
            "gridcolor": "#DFE8F3",
            "gridwidth": 2,
            "linecolor": "#EBF0F8",
            "showbackground": true,
            "ticks": "",
            "zerolinecolor": "#EBF0F8"
           },
           "zaxis": {
            "backgroundcolor": "white",
            "gridcolor": "#DFE8F3",
            "gridwidth": 2,
            "linecolor": "#EBF0F8",
            "showbackground": true,
            "ticks": "",
            "zerolinecolor": "#EBF0F8"
           }
          },
          "shapedefaults": {
           "line": {
            "color": "#2a3f5f"
           }
          },
          "ternary": {
           "aaxis": {
            "gridcolor": "#DFE8F3",
            "linecolor": "#A2B1C6",
            "ticks": ""
           },
           "baxis": {
            "gridcolor": "#DFE8F3",
            "linecolor": "#A2B1C6",
            "ticks": ""
           },
           "bgcolor": "white",
           "caxis": {
            "gridcolor": "#DFE8F3",
            "linecolor": "#A2B1C6",
            "ticks": ""
           }
          },
          "title": {
           "x": 0.05
          },
          "xaxis": {
           "automargin": true,
           "gridcolor": "#EBF0F8",
           "linecolor": "#EBF0F8",
           "ticks": "",
           "title": {
            "standoff": 15
           },
           "zerolinecolor": "#EBF0F8",
           "zerolinewidth": 2
          },
          "yaxis": {
           "automargin": true,
           "gridcolor": "#EBF0F8",
           "linecolor": "#EBF0F8",
           "ticks": "",
           "title": {
            "standoff": 15
           },
           "zerolinecolor": "#EBF0F8",
           "zerolinewidth": 2
          }
         }
        },
        "title": {
         "text": "Percentage of people who became economically active in the last 5 quarters, by <br>age (LFS) "
        },
        "uniformtext": {
         "minsize": 12,
         "mode": "hide"
        },
        "width": 800,
        "xaxis": {
         "anchor": "y",
         "domain": [
          0,
          1
         ],
         "gridcolor": "#F4F4F4",
         "ticksuffix": "",
         "title": {
          "text": "age"
         },
         "zerolinecolor": "#F4F4F4"
        },
        "yaxis": {
         "anchor": "x",
         "domain": [
          0,
          1
         ],
         "gridcolor": "#F4F4F4",
         "tickformat": ".0%",
         "ticksuffix": "",
         "title": {
          "text": "value"
         },
         "zerolinecolor": "#F4F4F4"
        }
       }
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from microimpute.comparisons import autoimpute\n",
    "from policyengine_uk import Microsimulation\n",
    "\n",
    "baseline = Microsimulation()\n",
    "baseline.subsample(20_000)\n",
    "\n",
    "df = baseline.calculate_dataframe([\n",
    "    \"age\",\n",
    "    \"gender\",\n",
    "    \"employment_status\",\n",
    "    \"pip_dl\",\n",
    "    \"pip_m\",\n",
    "], 2025)\n",
    "\n",
    "# Rough naive approach: impute 'became newly economically active' from the LFS?\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "# Load the data\n",
    "lfs = pd.read_csv(\"~/Downloads/UKDA-9133-tab/tab/lgwt22_5q_aj22_aj23_eul.tab\", sep=\"\\t\")\n",
    "\n",
    "# Create target variable as DataFrame\n",
    "inactivity_variables = [\n",
    "    \"INCAC051\",\n",
    "    \"INCAC052\",\n",
    "    \"INCAC053\",\n",
    "    \"INCAC054\",\n",
    "    \"INCAC055\",\n",
    "]\n",
    "was_inactive_at_some_point = np.any(\n",
    "    [lfs[col] >= 6 for col in inactivity_variables],\n",
    "    axis=0\n",
    ")\n",
    "became_active = np.any(\n",
    "    [lfs[col] == 1 for col in inactivity_variables],\n",
    "    axis=0\n",
    ")\n",
    "quarter_of_last_inactivity = np.argmax(\n",
    "    [lfs[col] >= 6 for col in inactivity_variables],\n",
    ")\n",
    "quarter_of_first_activity = quarter_of_last_inactivity\n",
    "length_activity_after_inactivity = (4 - quarter_of_first_activity + 1) / 4\n",
    "y_train = pd.DataFrame(dict(\n",
    "    was_inactive_at_some_point=was_inactive_at_some_point,\n",
    "    became_active_afterwards=became_active,\n",
    "    activity_length_after_inactivity=length_activity_after_inactivity * was_inactive_at_some_point * became_active\n",
    "))\n",
    "\n",
    "# Select predictor variables with clear names\n",
    "predictor_mapping = {\n",
    "    # Demographics\n",
    "    'AGE5': 'age',\n",
    "    'SEX': 'sex', \n",
    "    'MARSTA5': 'marital_status',\n",
    "    'ETUKEUL5': 'ethnicity',\n",
    "    'HIQUAL155': 'highest_qualification',\n",
    "    'GOVTOR5': 'region',\n",
    "    \n",
    "    # Employment characteristics\n",
    "    'FTPTWK5': 'full_or_part_time',\n",
    "    'SOC20M5': 'occupation_code',\n",
    "    'Inds07m5': 'industry_code',\n",
    "    'PUBLICR5': 'public_or_private_sector',\n",
    "    \n",
    "    # Income\n",
    "    'GRSSWK5': 'gross_weekly_pay',\n",
    "    'HRRATE5': 'hourly_pay_rate',\n",
    "    \n",
    "    # Household\n",
    "    'TEN15': 'housing_tenure',\n",
    "    'HDPCH195': 'num_dependent_children',\n",
    "    \n",
    "    # Education\n",
    "    'QULNOW5': 'current_qualification_studying',\n",
    "    'ENROLL5': 'enrolled_in_education',\n",
    "    \n",
    "    # Current health/disability status\n",
    "    'LNGLST5': 'has_longstanding_illness',\n",
    "    'LIMACT5': 'illness_limits_activities', \n",
    "    'DISEA5': 'disability_equality_act',\n",
    "    \n",
    "    # Previous health/disability status\n",
    "    'LNGLST1': 'had_longstanding_illness_q1',\n",
    "    'LIMACT1': 'illness_limited_activities_q1',\n",
    "    'DISEA1': 'disability_equality_act_q1'\n",
    "}\n",
    "\n",
    "# Create X_train with renamed columns\n",
    "X_train = lfs[list(predictor_mapping.keys())].rename(columns=predictor_mapping)\n",
    "\n",
    "# Get weights\n",
    "weights = lfs['LGWT22'].copy()\n",
    "\n",
    "# Convert categorical variables to strings\n",
    "categorical_vars = [\n",
    "    'sex', 'marital_status', 'ethnicity', 'highest_qualification', 'region',\n",
    "    'full_or_part_time', 'public_or_private_sector', 'housing_tenure', \n",
    "    'enrolled_in_education', 'current_qualification_studying',\n",
    "    'has_longstanding_illness', 'illness_limits_activities', 'disability_equality_act',\n",
    "    'had_longstanding_illness_q1', 'illness_limited_activities_q1', 'disability_equality_act_q1',\n",
    "    'occupation_code', 'industry_code'\n",
    "]\n",
    "\n",
    "for col in categorical_vars:\n",
    "    if col in X_train.columns:\n",
    "        X_train[col] = X_train[col].astype('Int64').astype(str).replace('<NA>', 'missing')\n",
    "\n",
    "# Remove rows with missing weights or target\n",
    "mask = weights.notna() & y_train['activity_length_after_inactivity'].notna()\n",
    "X_train = X_train[mask]\n",
    "y_train = y_train[mask]\n",
    "weights = weights[mask]\n",
    "\n",
    "\n",
    "print(f\"Final shapes - X: {X_train.shape}, y: {y_train.shape}\")\n",
    "print(f\"Positive class rate: {y_train['activity_length_after_inactivity'].mean():.3f}\")\n",
    "\n",
    "from microimpute.comparisons import autoimpute\n",
    "import plotly.express as px\n",
    "from microdf import MicroDataFrame\n",
    "\n",
    "df = pd.concat([X_train, y_train], axis=1)\n",
    "df[\"employment_income\"] = df.gross_weekly_pay.clip(lower=0) * 52\n",
    "df[\"gender\"] = df.sex.astype(int).map({1: \"MALE\", 2: \"FEMALE\"})\n",
    "df[\"weight\"] = weights\n",
    "\n",
    "df_w = MicroDataFrame(df, weights=weights)\n",
    "\n",
    "from policyengine.utils.charts import *\n",
    "\n",
    "fig = px.bar(\n",
    "    (df_w.was_inactive_at_some_point * df_w.became_active_afterwards).groupby(df_w.age).mean(),\n",
    "    color_discrete_sequence=[BLUE]\n",
    ").update_layout(\n",
    "    showlegend=False,\n",
    "    title=\"Percentage of people who became economically active in the last 5 quarters, by age (LFS)\",\n",
    "    yaxis_tickformat=\".0%\",\n",
    ")\n",
    "\n",
    "format_fig(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "c1402bfb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <link rel=\"preconnect\" href=\"https://fonts.googleapis.com\">\n",
       "    <link rel=\"preconnect\" href=\"https://fonts.gstatic.com\" crossorigin>\n",
       "    <link href=\"https://fonts.googleapis.com/css2?family=Roboto+Serif:ital,opsz,wght@0,8..144,100..900;1,8..144,100..900&display=swap\" rel=\"stylesheet\">\n",
       "    "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from policyengine.utils.charts import add_fonts\n",
    "\n",
    "add_fonts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "a75c732e",
   "metadata": {},
   "outputs": [],
   "source": [
    "imputed_vars = [\n",
    "    \"was_inactive_at_some_point\",\n",
    "    \"became_active_afterwards\",\n",
    "]\n",
    "\n",
    "predictor_vars = [\n",
    "    \"age\",\n",
    "    \"gender\",\n",
    "    \"employment_income\",\n",
    "]\n",
    "\n",
    "efrs = baseline.calculate_dataframe([\n",
    "    \"age\",\n",
    "    \"gender\",\n",
    "    \"employment_income\",\n",
    "], 2025)\n",
    "\n",
    "from quantile_forest import RandomForestQuantileRegressor as QRF\n",
    "\n",
    "df.gender = df.gender.map({\n",
    "    \"MALE\": 0,\n",
    "    \"FEMALE\": 1,\n",
    "})\n",
    "\n",
    "efrs.gender = efrs.gender.map({\n",
    "    \"MALE\": 0,\n",
    "    \"FEMALE\": 1,\n",
    "})\n",
    "\n",
    "df = df.astype(float)\n",
    "efrs = efrs.astype(float)\n",
    "\n",
    "model = QRF()\n",
    "model = model.fit(df[predictor_vars], df[imputed_vars])\n",
    "\n",
    "efrs[imputed_vars] = model.predict(efrs[predictor_vars])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "bf204593",
   "metadata": {},
   "outputs": [],
   "source": [
    "efrs[\"joined_labour_force_recently\"] = efrs.was_inactive_at_some_point * efrs.became_active_afterwards\n",
    "\n",
    "# Children not in the LFS, don't impute for them\n",
    "efrs[\"joined_labour_force_recently\"] = np.where(\n",
    "    ~efrs.age.between(16, 65),\n",
    "    0,\n",
    "    efrs[\"joined_labour_force_recently\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "cfbd833a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Impute target wage for those who joined the labour force recently\n",
    "\n",
    "target_wage_model = QRF()\n",
    "\n",
    "target_wage_model = target_wage_model.fit(\n",
    "    efrs[efrs.joined_labour_force_recently == 1][predictor_vars],\n",
    "    efrs[efrs.joined_labour_force_recently == 1].employment_income\n",
    ")\n",
    "\n",
    "efrs[\"target_wage\"] = target_wage_model.predict(efrs[predictor_vars])\n",
    "efrs[\"target_wage\"][efrs.joined_labour_force_recently == 0] = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "250dbc74",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_employer_ni(wage, exemption = False):\n",
    "    if exemption:\n",
    "        return 0\n",
    "    wages_over_threshold = np.maximum(wage - 5_000, 0)\n",
    "    return wages_over_threshold * 0.15\n",
    "\n",
    "def calculate_post_tax_salary(wage):\n",
    "    # very rough estimate of post-tax salary after it and employee ni\n",
    "    if wage <= 12_570:\n",
    "        return wage\n",
    "    elif wage <= 50_270:\n",
    "        income_tax = (wage - 12_570) * 0.2\n",
    "        employee_ni = (wage - 12_570) * 0.12\n",
    "    else:\n",
    "        income_tax = (50_270 - 12_570) * 0.2 + (wage - 50_270) * 0.4\n",
    "        employee_ni = (50_270 - 12_570) * 0.12 + (wage - 50_270) * 0.02\n",
    "\n",
    "    return wage - income_tax - employee_ni\n",
    "\n",
    "efrs[\"baseline_target_employer_ni\"] = calculate_employer_ni(efrs.target_wage)\n",
    "efrs[\"reform_target_employer_ni\"] = calculate_employer_ni(efrs.target_wage, exemption=True)\n",
    "efrs[\"change_target_employer_ni\"] = efrs.reform_target_employer_ni - efrs.baseline_target_employer_ni\n",
    "\n",
    "# 40% incidence\n",
    "\n",
    "efrs[\"change_in_target_salary\"] = -efrs.change_target_employer_ni * 0.4\n",
    "efrs[\"original_target_salary_post_tax\"] = (efrs.target_wage).apply(calculate_post_tax_salary)\n",
    "efrs[\"target_salary_post_tax\"] = (efrs.target_wage + efrs.change_in_target_salary).apply(calculate_post_tax_salary)\n",
    "efrs[\"change_in_target_salary_post_tax\"] = efrs.target_salary_post_tax - efrs.original_target_salary_post_tax\n",
    "efrs[\"rel_change_in_target_salary\"] = efrs.change_in_target_salary_post_tax / (efrs.original_target_salary_post_tax+1)\n",
    "\n",
    "rel_change_in_participation = efrs.rel_change_in_target_salary / 0.01 * 0.36\n",
    "\n",
    "efrs[\"joins_labour_force\"] = np.random.random() < rel_change_in_participation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "150e8ef0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "212.0630625"
      ]
     },
     "execution_count": 103,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(efrs.joins_labour_force * baseline.calculate(\"household_weight\", 2025, map_to=\"person\").values).sum()/1e3"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}