diff --git a/examples/Simulate_all_policies.ipynb b/examples/Simulate_all_policies.ipynb
index ec8628d..4de5255 100644
--- a/examples/Simulate_all_policies.ipynb
+++ b/examples/Simulate_all_policies.ipynb
@@ -60,13 +60,16 @@
     "\n",
     "candidate_dict = {\n",
     "    \"Obama 2015\": {\"policy_path\": obama2015_path, \"baseline_path\": [pre_2020_baseline], \"start_year\": 2016},\n",
-    "    \"Romney 2012\": {\"policy_path\": romney2012_path, \"baseline_path\": [pre_2020_baseline], \"start_year\": 2014}, #wanted to do 13,  but taxcalc with CPS only goes to 13\n",
+    "    \"Romney 2012\": {\"policy_path\": romney2012_path, \"baseline_path\": [pre_2020_baseline], \"start_year\": 2014}, #wanted to do 13,  but taxcalc with CPS only goes to 14\n",
     "    \"Clinton 2016\": {\"policy_path\": clinton2016_path, \"baseline_path\": [pre_2020_baseline], \"start_year\": 2017},\n",
     "    \"Trump 2016\": {\"policy_path\": trump2016_path, \"baseline_path\": [pre_2020_baseline], \"start_year\": 2017},\n",
     "    \"Biden 2020\": {\"policy_path\": biden2020_path, \"baseline_path\": [pre_2020_baseline, baseline_2020], \"start_year\": 2021},\n",
     "    \"Trump 2020\": {\"policy_path\": trump2020_path, \"baseline_path\": [pre_2020_baseline], \"start_year\": 2021},\n",
-    "    \"Harris 2024\": {\"policy_path\": trump2020_path, \"baseline_path\": [pre_2020_baseline], \"start_year\": 2025},\n",
-    "    \"Trump 2024\": {\"policy_path\": trump2020_path, \"baseline_path\": [pre_2020_baseline], \"start_year\": 2025}}"
+    "    # \"Harris 2024\": {\"policy_path\": harris2024_path, \"baseline_path\": [pre_2020_baseline, baseline_2020], \"start_year\": 2025},\n",
+    "    # \"Trump 2024\": {\"policy_path\": trump2024_path, \"baseline_path\": [pre_2020_baseline, baseline_2020], \"start_year\": 2025}\n",
+    "    \"Harris 2024\": {\"policy_path\": harris2024_path, \"baseline_path\": None, \"start_year\": 2026},  # Use 2026 for these candidates since largest diff is how treat TCJA expirations\n",
+    "    \"Trump 2024\": {\"policy_path\": trump2024_path, \"baseline_path\": None, \"start_year\": 2026}\n",
+    "    }"
    ]
   },
   {
@@ -111,7 +114,34 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "iot_all.iot[-1].mtr"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "iot_all.iot[-2].mtr"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "len(iot_all.iot)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -127,7 +157,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -144,7 +174,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -193,7 +223,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -244,7 +274,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -290,12 +320,24 @@
     "gz_plot.write_image(\n",
     "            os.path.join(path, \"gz_numerical_all.png\"),\n",
     "            scale=4\n",
+    "        )\n",
+    "\n",
+    "# Make the same plot but stop at 200_000\n",
+    "gz_plot.update_xaxes(range=[0, 200_000])\n",
+    "gz_plot.write_image(\n",
+    "            os.path.join(path, \"gz_numerical_all_200k.png\"),\n",
+    "            scale=4\n",
+    "        )\n",
+    "gz_plot.update_xaxes(range=[0, 100_000])\n",
+    "gz_plot.write_image(\n",
+    "            os.path.join(path, \"gz_numerical_all_100k.png\"),\n",
+    "            scale=4\n",
     "        )"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -342,12 +384,22 @@
     "gz_plot.write_image(\n",
     "            os.path.join(path, \"gz_numerical_democrats.png\"),\n",
     "            scale=4\n",
+    "        )\n",
+    "gz_plot.update_xaxes(range=[0, 200_000])\n",
+    "gz_plot.write_image(\n",
+    "            os.path.join(path, \"gz_numerical_democrats_200k.png\"),\n",
+    "            scale=4\n",
+    "        )\n",
+    "gz_plot.update_xaxes(range=[0, 100_000])\n",
+    "gz_plot.write_image(\n",
+    "            os.path.join(path, \"gz_numerical_democrats_100k.png\"),\n",
+    "            scale=4\n",
     "        )"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -394,12 +446,22 @@
     "gz_plot.write_image(\n",
     "            os.path.join(path, \"gz_numerical_republicans.png\"),\n",
     "            scale=4\n",
+    "        )\n",
+    "gz_plot.update_xaxes(range=[0, 200_000])\n",
+    "gz_plot.write_image(\n",
+    "            os.path.join(path, \"gz_numerical_republicans_200k.png\"),\n",
+    "            scale=4\n",
+    "        )\n",
+    "gz_plot.update_xaxes(range=[0, 100_000])\n",
+    "gz_plot.write_image(\n",
+    "            os.path.join(path, \"gz_numerical_republicans_100k.png\"),\n",
+    "            scale=4\n",
     "        )"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -419,7 +481,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -445,7 +507,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -518,7 +580,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -583,7 +645,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -617,7 +679,101 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Do experiment where hold constant g(z) - pick a candidate as baseline - then plot epsilon(z)\n",
+    "# that would recover those g(z) given the tax rates of each candidate\n",
+    "# one plot with epsilon(z) for each candidate\n",
+    "# Will pick Trump and Clinton (2016) for example\n",
+    "\n",
+    "# First, plot just their g(z)\n",
+    "fig = px.line(\n",
+    "    x=iot_all.iot[2].df().z[10:],\n",
+    "    y=[iot_all.iot[2].df().g_z_numerical[10:], iot_all.iot[3].df().g_z_numerical[10:]],\n",
+    "    labels={\"x\": \"Wages and Salaries\", \"y\": r\"$g_z$\"},\n",
+    "    )\n",
+    "fig.update_layout(\n",
+    "    template=template,\n",
+    "    legend=dict(\n",
+    "        title=\"Candidate:\",\n",
+    "        ),\n",
+    ")\n",
+    "candidate_name = [\"Obama 2012\", \"Romney 2012\"]\n",
+    "label_dict = {}\n",
+    "for i, v in enumerate(candidate_name):\n",
+    "    label_dict[\"wide_variable_\" + str(i)] = str(candidate_name[i])\n",
+    "fig.for_each_trace(lambda t: t.update(name = label_dict[t.name], legendgroup = label_dict[t.name],\n",
+    "                                      hovertemplate = t.hovertemplate.replace(t.name, label_dict[t.name])))\n",
+    "\n",
+    "fig.write_image(\n",
+    "            os.path.join(path, \"romney_obama_g_z_numerical.png\"),\n",
+    "            scale=4\n",
+    "        )\n",
+    "candidate_name = [\"Clinton 2016\", \"Trump 2016\"]\n",
+    "label_dict = {}\n",
+    "for i, v in enumerate(candidate_name):\n",
+    "    label_dict[\"wide_variable_\" + str(i)] = str(candidate_name[i])\n",
+    "fig.for_each_trace(lambda t: t.update(name = label_dict[t.name], legendgroup = label_dict[t.name],\n",
+    "                                      hovertemplate = t.hovertemplate.replace(t.name, label_dict[t.name])))\n",
+    "\n",
+    "fig.write_image(\n",
+    "            os.path.join(path, \"trump_clinton_g_z_numerical.png\"),\n",
+    "            scale=4\n",
+    "        )\n",
+    "candidate_name = [\"Biden 2020\", \"Trump 2020\"]\n",
+    "label_dict = {}\n",
+    "for i, v in enumerate(candidate_name):\n",
+    "    label_dict[\"wide_variable_\" + str(i)] = str(candidate_name[i])\n",
+    "fig.for_each_trace(lambda t: t.update(name = label_dict[t.name], legendgroup = label_dict[t.name],\n",
+    "                                      hovertemplate = t.hovertemplate.replace(t.name, label_dict[t.name])))\n",
+    "\n",
+    "fig.write_image(\n",
+    "            os.path.join(path, \"trump_biden_g_z_numerical.png\"),\n",
+    "            scale=4\n",
+    "        )\n",
+    "candidate_name = [\"Harris 2024\", \"Trump 2024\"]\n",
+    "label_dict = {}\n",
+    "for i, v in enumerate(candidate_name):\n",
+    "    label_dict[\"wide_variable_\" + str(i)] = str(candidate_name[i])\n",
+    "fig.for_each_trace(lambda t: t.update(name = label_dict[t.name], legendgroup = label_dict[t.name],\n",
+    "                                      hovertemplate = t.hovertemplate.replace(t.name, label_dict[t.name])))\n",
+    "\n",
+    "fig.write_image(\n",
+    "            os.path.join(path, \"trump_harris_g_z_numerical.png\"),\n",
+    "            scale=4\n",
+    "        )\n",
+    "# Now find the epsilon(z) that would give Trump's policies the same g(z) as Clinton\n",
+    "eti_beliefs_lw, eti_beliefs_jjz = iot.inverse_optimal_tax.find_eti(iot_all.iot[2], iot_all.iot[3], g_z_type=\"g_z\")\n",
+    "# idx = (np.where(np.absolute(eti_beliefs_jjz[1:]) < 1))[0]\n",
+    "idx = (np.where(((iot_all.iot[3].df().g_z.values) >  1.05) | ((iot_all.iot[3].df().g_z.values) < 0.95)))[0]\n",
+    "fig2 = px.line(\n",
+    "    x=iot_all.iot[2].df().z[idx],\n",
+    "    y=eti_beliefs_jjz[idx],\n",
+    "    labels={\"x\": \"Wages and Salaries\", \"y\": r\"$\\text{Implied } \\varepsilon$\"},\n",
+    "    )\n",
+    "fig2.update_layout(\n",
+    "    template=template,\n",
+    ")\n",
+    "fig2.write_image(\n",
+    "            os.path.join(path, \"trump_eti.png\"),\n",
+    "            scale=4\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "(np.where(((iot_all.iot[2].df().z.values) >  1.01) | ((iot_all.iot[2].df().z.values) < 0.99)))[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -656,125 +812,89 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Thought experiment: What beliefs about ETI do the candidates need to \n",
-    "# justify their policies if we assume they are utilitarian? \n",
-    "# If the elasticities are wildly counterfactual, we can reject this hypothesis\n",
+    "import plotly.express as px\n",
+    "import pandas as pd\n",
     "\n",
-    "eti_utilitarian = np.zeros((len(iot_all.iot), len(iot_all.iot[0].df().z)))\n",
-    "for i in range(len(iot_all.iot)):\n",
-    "    eti_utilitarian[i, :] = iot.inverse_optimal_tax.find_eti(\n",
-    "        iot_all.iot[i], g_z = np.ones(len(iot_all.iot[i].df().z))\n",
-    "        )\n",
+    "# Sample data (replace with actual data)\n",
+    "data = {\n",
+    "    \"gross_income\": [0, 25000, 50000, 75000, 100000, 120000],\n",
+    "    \"non_constant_marginal_tax_rates\": [0, 0.5, -0.2, 0.3, 0.1, 0],\n",
+    "    \"tax_base_elasticity\": [0, 1, 0.8, -0.5, -0.7, -0.6],\n",
+    "    \"remainder_terms\": [0, 0.1, -0.1, 0.05, -0.05, 0]\n",
+    "}\n",
     "\n",
-    "# Convert the 2D array to a DataFrame for plotting\n",
-    "eti_df = pd.DataFrame(eti_utilitarian.T, columns=labels)\n",
+    "df = pd.DataFrame(data)\n",
     "\n",
-    "fig_eti = px.line(\n",
-    "    eti_df,\n",
-    "    x=iot_all.iot[0].df().z,\n",
-    "    y=eti_df.columns,\n",
-    "    labels={\"x\": \"Wages and Salaries\", \"value\": r\"$\\varepsilon$\", \"variable\": \"Candidate\"},\n",
-    ")\n",
-    "fig_eti.update_layout(\n",
-    "    template=template,\n",
-    ")\n",
+    "# Create area plot\n",
+    "fig = px.area(df, x=\"gross_income\", y=[\"non_constant_marginal_tax_rates\", \"tax_base_elasticity\", \"remainder_terms\"],\n",
+    "              labels={\"value\": \"Effect\", \"variable\": \"Category\"},\n",
+    "              title=\"Area Plot of Tax and Welfare Components\",\n",
+    "              color_discrete_map={\n",
+    "                  \"non_constant_marginal_tax_rates\": \"blue\",\n",
+    "                  \"tax_base_elasticity\": \"red\",\n",
+    "                  \"remainder_terms\": \"green\"\n",
+    "              })\n",
     "\n",
-    "fig_eti.update_traces(\n",
-    "    line=dict(dash=\"dot\", color=\"blue\"),\n",
-    "    selector=dict(name=\"Obama 2015\")\n",
-    ")\n",
-    "fig_eti.update_traces(\n",
-    "    line=dict(dash=\"dot\", color=\"red\"),\n",
-    "    selector=dict(name=\"Romney 2012\")\n",
-    ")\n",
-    "fig_eti.update_traces(\n",
-    "    line=dict(dash=\"dash\", color=\"blue\"),\n",
-    "    selector=dict(name=\"Clinton 2016\")\n",
-    ")\n",
-    "fig_eti.update_traces(\n",
-    "    line=dict(dash=\"dash\", color=\"red\"),\n",
-    "    selector=dict(name=\"Trump 2016\")\n",
-    ")\n",
-    "fig_eti.update_traces(\n",
-    "    line=dict(dash=\"dashdot\", color=\"blue\"),\n",
-    "    selector=dict(name=\"Biden 2020\")\n",
-    ")\n",
-    "fig_eti.update_traces(\n",
-    "    line=dict(dash=\"dashdot\", color=\"red\"),\n",
-    "    selector=dict(name=\"Trump 2020\")\n",
-    ")\n",
-    "fig_eti.update_traces(\n",
-    "    line=dict(dash=\"solid\", color=\"blue\"),\n",
-    "    selector=dict(name=\"Harris 2024\")\n",
-    ")\n",
-    "fig_eti.update_traces(\n",
-    "    line=dict(dash=\"solid\", color=\"red\"),\n",
-    "    selector=dict(name=\"Trump 2024\")\n",
-    ")\n",
-    "fig_eti.write_image(\n",
-    "            os.path.join(path, \"eti_utilitarian.png\"),\n",
-    "            scale=4\n",
-    "        )"
+    "fig.show()\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "https://raw.githubusercontent.com/jdebacker/examples/pres_proposals/psl_examples/taxcalc/Obama2015.json\n",
+      "https://raw.githubusercontent.com/jdebacker/examples/pres_proposals/psl_examples/taxcalc/Romney2012.json\n",
+      "https://raw.githubusercontent.com/jdebacker/examples/pres_proposals/psl_examples/taxcalc/Clinton2016.json\n",
+      "https://raw.githubusercontent.com/PSLmodels/examples/main/psl_examples/taxcalc/Trump2016.json\n",
+      "https://raw.githubusercontent.com/PSLmodels/examples/main/psl_examples/taxcalc/Biden2020.json\n",
+      "https://raw.githubusercontent.com/PSLmodels/examples/main/psl_examples/taxcalc/TCJA.json\n",
+      "https://raw.githubusercontent.com/PSLmodels/examples/main/psl_examples/taxcalc/Harris2024.json\n",
+      "https://raw.githubusercontent.com/PSLmodels/examples/main/psl_examples/taxcalc/Trump2024.json\n"
+     ]
+    }
+   ],
    "source": [
-    "# Do experiment where hold constant g(z) - pick a candidate as baseline - then plot epsilon(z)\n",
-    "# that would recover those g(z) given the tax rates of each candidate\n",
-    "# one plot with epsilon(z) for each candidate\n",
-    "# Will pick Trump and Clinton (2016) for example\n",
-    "\n",
+    "policies = []\n",
+    "baseline_policies = []\n",
+    "labels = list(candidate_dict.keys())\n",
+    "# get years from start_year in candidate_dict\n",
+    "years = [v[\"start_year\"] for v in candidate_dict.values()]\n",
+    "for k, v in candidate_dict.items():\n",
+    "    # with open(v[\"policy_path\"], \"r\") as file:\n",
+    "        # json1 = file.read()\n",
+    "    json1 = v[\"policy_path\"]#json.load(open(v[\"policy_path\"]))\n",
+    "    policies.append(json1)\n",
+    "    if v[\"baseline_path\"] is None:\n",
+    "        json2 = {}\n",
+    "    else:\n",
+    "        for ii, vv in enumerate(v[\"baseline_path\"]):\n",
+    "            list_json = []\n",
+    "            # with open(vv, \"r\") as file:\n",
+    "            #     json2 = file.read()\n",
+    "            #     list_json.append(json2)\n",
+    "            json2 = v[\"baseline_path\"]#open(v[\"baseline_path\"])\n",
+    "    baseline_policies.append(list_json)\n",
     "\n",
-    "# First, plot just their g(z)\n",
-    "fig = px.line(\n",
-    "    x=iot_all.iot[2].df().z[10:],\n",
-    "    y=[iot_all.iot[2].df().g_z_numerical[10:], iot_all.iot[3].df().g_z_numerical[10:]],\n",
-    "    labels={\"x\": \"Wages and Salaries\", \"y\": r\"$g_z$\"},\n",
-    "    )\n",
+    "iot_all = iot_user.iot_comparison(\n",
+    "    policies=[policies[0]],\n",
+    "    baseline_policies=[baseline_policies[0]],\n",
+    "    mtr_smoother=\"kreg\",  # this is kreg or HSV\n",
+    "    labels=[labels[0]],\n",
+    "    years=[years[0]],\n",
+    "    data=\"CPS\"\n",
+    ")\n",
+    "fig = iot_all.JJZFig4(policy='Obama 2015')\n",
     "fig.update_layout(\n",
     "    template=template,\n",
-    "    legend=dict(\n",
-    "        title=\"Candidate:\",\n",
-    "        ),\n",
     ")\n",
-    "candidate_name = [\"Clinton 2016\", \"Trump 2016\"]\n",
-    "label_dict = {}\n",
-    "for i, v in enumerate(candidate_name):\n",
-    "    label_dict[\"wide_variable_\" + str(i)] = str(candidate_name[i])\n",
-    "fig.for_each_trace(lambda t: t.update(name = label_dict[t.name], legendgroup = label_dict[t.name],\n",
-    "                                      hovertemplate = t.hovertemplate.replace(t.name, label_dict[t.name])))\n",
-    "\n",
-    "fig.write_image(\n",
-    "            os.path.join(path, \"trump_clinton_g_z_numerical.png\"),\n",
-    "            scale=4\n",
-    "        )\n",
-    "\n",
-    "\n",
-    "# What ETI does Clinton need to justify Trump's SWW given her policy?\n",
-    "eti_clinton = iot.inverse_optimal_tax.find_eti(\n",
-    "    iot_all.iot[2], \n",
-    "    g_z=iot_all.iot[3].df().g_z)\n",
+    "# fig.update_xaxes(range=[0, 850000])\n",
     "\n",
-    "fig_eti_clinton = px.line(\n",
-    "    x=iot_all.iot[2].df().z,\n",
-    "    y=eti_clinton,\n",
-    "    labels={\"x\": \"Wages and Salaries\", \"y\": r\"$\\varepsilon$\"},\n",
-    "    )\n",
-    "fig_eti_clinton.update_layout(\n",
-    "    template=template,\n",
-    ")\n",
-    "fig_eti_clinton.update_traces(\n",
-    "    line=dict(dash=\"dash\", color=\"blue\"),\n",
-    "    selector=dict(name=\"Clinton 2016\")\n",
-    ")\n",
-    "fig_eti_clinton.write_image(\n",
-    "            os.path.join(path, \"eti_clinton.png\"),\n",
-    "            scale=4\n",
-    "        )\n"
+    "fig.show()"
    ]
   },
   {
@@ -782,30 +902,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "\n",
-    "# Converse: What elasticty does Trump need to justify Clinton's weights?\n",
-    "eti_trump = iot.inverse_optimal_tax.find_eti(\n",
-    "    iot_all.iot[3], \n",
-    "    g_z=iot_all.iot[2].df().g_z)\n",
-    "\n",
-    "fig_eti_trump = px.line(\n",
-    "    x=iot_all.iot[2].df().z,\n",
-    "    y=eti_trump,\n",
-    "    labels={\"x\": \"Wages and Salaries\", \"y\": r\"$\\varepsilon$\"},\n",
-    "    )\n",
-    "fig_eti_trump.update_layout(\n",
-    "    template=template,\n",
-    ")\n",
-    "fig_eti_trump.update_traces(\n",
-    "    line=dict(dash=\"dash\", color=\"red\"),\n",
-    "    selector=dict(name=\"Trump 2016\")\n",
-    ")\n",
-    "fig_eti_trump.write_image(\n",
-    "            os.path.join(path, \"eti_trump.png\"),\n",
-    "            scale=4\n",
-    "        )\n"
-   ]
+   "source": []
   }
  ],
  "metadata": {
@@ -824,7 +921,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.7"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,
diff --git a/iot/inverse_optimal_tax.py b/iot/inverse_optimal_tax.py
index ff5508b..c262d05 100644
--- a/iot/inverse_optimal_tax.py
+++ b/iot/inverse_optimal_tax.py
@@ -78,7 +78,12 @@ def __init__(
             self.eti = eti_spl(self.z)
         # compute marginal tax rate schedule
         self.mtr, self.mtr_prime = self.compute_mtr_dist(
-            data, weight_var, income_measure, mtr_smoother, mtr_smooth_param, kreg_bw
+            data,
+            weight_var,
+            income_measure,
+            mtr_smoother,
+            mtr_smooth_param,
+            kreg_bw,
         )
         # compute theta_z, the elasticity of the tax base
         self.theta_z = 1 + ((self.z * self.f_prime) / self.f)
@@ -111,7 +116,13 @@ def df(self):
         return df
 
     def compute_mtr_dist(
-        self, data, weight_var, income_measure, mtr_smoother, mtr_smooth_param, kreg_bw
+        self,
+        data,
+        weight_var,
+        income_measure,
+        mtr_smoother,
+        mtr_smooth_param,
+        kreg_bw,
     ):
         """
         Compute marginal tax rates over the income distribution and
@@ -372,7 +383,9 @@ def sw_weights(self):
             + ((self.theta_z * self.eti * self.mtr) / (1 - self.mtr))
             + ((self.eti * self.z * self.mtr_prime) / (1 - self.mtr) ** 2)
         )
-        integral = np.trapz(g_z * self.f, self.z) # renormalize to integrate to 1
+        integral = np.trapz(
+            g_z * self.f, self.z
+        )  # renormalize to integrate to 1
         g_z = g_z / integral
 
         # use Lockwood and Weinzierl formula, which should be equivalent but using numerical differentiation
@@ -391,7 +404,7 @@ def sw_weights(self):
         return g_z, g_z_numerical
 
 
-def find_eti(iot, g_z = None, eti_0 = 0.25):
+def find_eti(iot, g_z=None, eti_0=0.25):
     """
     This function solves for the ETI that would result in the
     policy represented via MTRs in IOT being consistent with the
@@ -409,12 +422,16 @@ def find_eti(iot, g_z = None, eti_0 = 0.25):
     Returns:
         eti_beliefs (array-like): vector of ETI beliefs over z
     """
-    
+
     if g_z is None:
         g_z = iot.g_z
-    
+
     # we solve an ODE of the form f'(z) + P(z)f(z) = Q(z)
-    P_z = 1/iot.z + iot.f_prime/iot.f + iot.mtr_prime/(iot.mtr * (1-iot.mtr))
+    P_z = (
+        1 / iot.z
+        + iot.f_prime / iot.f
+        + iot.mtr_prime / (iot.mtr * (1 - iot.mtr))
+    )
     # integrating factor for ODE: mu(z) * f'(z) + mu(z) * P(z) * f(z) = mu(z) * Q(z)
     mu_z = np.exp(np.cumsum(P_z))
     Q_z = (g_z - 1) * (1 - iot.mtr) / (iot.mtr * iot.z)
diff --git a/iot/iot_user.py b/iot/iot_user.py
index 1ffa3dc..e392d58 100644
--- a/iot/iot_user.py
+++ b/iot/iot_user.py
@@ -63,7 +63,7 @@ def __init__(
         kde_bw=None,
         mtr_smoother="kreg",
         mtr_smooth_param=1000,
-        kreg_bw=[120_000]
+        kreg_bw=[120_000],
     ):
         self.income_measure = income_measure
         self.weight_var = weight_var
@@ -204,7 +204,7 @@ def SaezFig2(self, DS2011=False, upper_bound=None):
         )
         return fig
 
-    def JJZFig4(self, policy="Current Law", var="g_z"):
+    def JJZFig4(self, policy="Current Law", var="g_z", upper_bound=500_000):
         """
         Function to plot a decomposition of the political weights, `g_z`
 
@@ -226,28 +226,29 @@ def JJZFig4(self, policy="Current Law", var="g_z"):
             g_weights = df.g_z_numerical
 
         # g1 with mtr_prime = 0
-        g1 = (
-            1
-            + +((df.theta_z * self.iot[k].eti * df.mtr) / (1 - df.mtr))
-            + ((self.iot[k].eti * df.z * 0) / (1 - df.mtr) ** 2)
+        g1 = ((df.theta_z * self.iot[k].eti * df.mtr) / (1 - df.mtr)) + (
+            (self.iot[k].eti * df.z * 0) / (1 - df.mtr) ** 2
         )
         # g2 with theta_z = 0
-        g2 = (
-            1
-            + +((0 * self.iot[k].eti * df.mtr) / (1 - df.mtr))
-            + ((self.iot[k].eti * df.z * df.mtr_prime) / (1 - df.mtr) ** 2)
+        g2 = ((0 * self.iot[k].eti * df.mtr) / (1 - df.mtr)) + (
+            (self.iot[k].eti * df.z * df.mtr_prime) / (1 - df.mtr) ** 2
         )
         integral = np.trapz(g1, df.z)
-        g1 = g1 / integral
+        # g1 = g1 / integral
         integral = np.trapz(g2, df.z)
+        # g2 = g2 / integral
         plot_df = pd.DataFrame(
             {
                 self.income_measure: df.z,
-                "Overall weight": g_weights,
-                "Tax Base Elasticity": g1,
-                "Nonconstant MTRs": g2,
+                "Overall Weight": g_weights,
+                "Tax Base Elasticity": 1 + g1,
+                "Nonconstant MTRs": 1
+                + g1
+                + g2
+                + np.abs(g1) * (np.sign(g1) != np.sign(g2)),
             }
         )
+
         fig = go.Figure()
         # add a line at y = 1
         fig.add_trace(
@@ -262,59 +263,65 @@ def JJZFig4(self, policy="Current Law", var="g_z"):
                 showlegend=False,
             )
         )
-        # fig.add_trace(
-        #     go.Scatter(
-        #         x=plot_df[self.income_measure],
-        #         y=plot_df["Nonconstant MTRs"],
-        #         fill="tonexty",  # fill area between trace1 and trace2
-        #         # fill="tozeroy",
-        #         mode="lines",
-        #         name="Nonconstant MTRs",
-        #     )
-        # )
         fig.add_trace(
             go.Scatter(
                 x=plot_df[self.income_measure],
                 y=plot_df["Tax Base Elasticity"],
-                fill="tonexty",  # fill area between trace0 and trace1
+                fill="tonexty",  # fill area from prior trace to this one
+                # fill="tozeroy",
                 mode="lines",
+                fillcolor="rgba(4,40,145,0.5)",
                 name="Tax Base Elasticity",
             )
         )
         fig.add_trace(
             go.Scatter(
-                x=plot_df[self.income_measure],
-                y=plot_df["Overall weight"],
-                fill="tonexty",
+                x=[
+                    plot_df[self.income_measure].min(),
+                    plot_df[self.income_measure].max(),
+                ],
+                y=[1, 1],
                 mode="lines",
-                name="Nonconstant MTRs",
+                line=dict(color="black", width=1, dash="dash"),
+                showlegend=False,
             )
         )
-        # add a line at y=0
         fig.add_trace(
             go.Scatter(
                 x=plot_df[self.income_measure],
-                y=plot_df["Overall weight"],
+                y=plot_df["Nonconstant MTRs"],
+                fill="tonexty",  # fill area from prior trace to this one
                 mode="lines",
-                line=dict(color="black", width=1, dash="solid"),
-                name="Overall weight",
-                showlegend=False,
+                fillcolor="rgba(229,0,0,0.5)",
+                name="Nonconstant MTRs",
             )
         )
-        # add a line at y=0
+        # Add black line for overall weight
         fig.add_trace(
             go.Scatter(
-                x=[
-                    plot_df[self.income_measure].min(),
-                    plot_df[self.income_measure].max(),
-                ],
-                y=[0, 0],
+                x=plot_df[self.income_measure],
+                y=plot_df["Overall Weight"],
                 mode="lines",
-                line=dict(color="black", width=1, dash="dash"),
+                line=dict(color="black", width=2, dash="solid"),
+                name="Overall Weight",
+                showlegend=True,
             )
         )
+        # # add a line at y=0
+        # fig.add_trace(
+        #     go.Scatter(
+        #         x=[
+        #             plot_df[self.income_measure].min(),
+        #             plot_df[self.income_measure].max(),
+        #         ],
+        #         y=[0, 0],
+        #         mode="lines",
+        #         line=dict(color="black", width=1, dash="dash"),
+        #     )
+        # )
         fig.update_layout(
             xaxis_title=OUTPUT_LABELS[self.income_measure],
             yaxis_title=r"$g_z$",
         )
+        fig.update_xaxes(range=[0, upper_bound])
         return fig