Update decomposition plot to take numerical or analytical weights

.github/workflows/build_and_test.yml

@@ -1,4 +1,4 @@
-name: Build and test [Python 3.9, 3.10, 3.11]
+name: Build and test [Python 3.10, 3.11]
 
 on: [push, pull_request]
 
@@ -7,7 +7,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.9, "3.10", "3.11"]
+        python-version: ["3.10", "3.11"]
 
     steps:
       - name: Checkout

iot/inverse_optimal_tax.py

@@ -95,6 +95,7 @@ def df(self):
         dict_out = {
             "z": self.z,
             "f": self.f,
+            "F": self.F,
             "f_prime": self.f_prime,
             "mtr": self.mtr,
             "mtr_prime": self.mtr_prime,
@@ -336,7 +337,7 @@ def sw_weights(self):
             + ((self.eti * self.z * self.mtr_prime) / (1 - self.mtr) ** 2)
         )
         integral = np.trapz(g_z, self.z)
-        g_z = g_z / integral
+        # g_z = g_z / integral
         # use Lockwood and Weinzierl formula, which should be equivalent but using numerical differentiation
         bracket_term = (
             1
@@ -348,7 +349,7 @@ def sw_weights(self):
         d_dz_bracket = np.append(d_dz_bracket, d_dz_bracket[-1])
         g_z_numerical = -(1 / self.f) * d_dz_bracket
         integral = np.trapz(g_z_numerical, self.z)
-        g_z_numerical = g_z_numerical / integral
+        # g_z_numerical = g_z_numerical / integral
         return g_z, g_z_numerical
 
 

iot/iot_user.py

@@ -104,9 +104,9 @@ def __init__(
                     income_measure=income_measure,
                     weight_var=weight_var,
                     eti=eti,
-                    bandwidth=bandwidth,
-                    lower_bound=lower_bound,
-                    upper_bound=upper_bound,
+                    # bandwidth=bandwidth,
+                    # lower_bound=lower_bound,
+                    # upper_bound=upper_bound,
                     dist_type=dist_type,
                     kde_bw=kde_bw,
                     mtr_smoother=mtr_smoother,
@@ -202,39 +202,74 @@ def SaezFig2(self, DS2011=False, upper_bound=None):
         )
         return fig
 
-    def JJZFig4(self, policy="Current Law"):
+    def JJZFig4(self, policy="Current Law", var="g_z"):
+        """
+        Function to plot a decomposition of the political weights, `g_z`
+
+        Args:
+            policy (str): policy to plot
+            var (str): variable to plot against income
+                Variable options are:
+                 * 'g_z' for analytically derived weights
+                 * 'g_z_numeric' for numerically derived weights
+
+        Returns:
+            fig (plotly.express figure): figure with the decomposition
+        """
         k = self.labels.index(policy)
         df = self.iot[k].df()
+        if var == "g_z":
+            g_weights = df.g_z
+        else:
+            g_weights = df.g_z_numerical
+
         # g1 with mtr_prime = 0
         g1 = (
-            0
-            + ((df.theta_z * self.iot[k].eti * df.mtr) / (1 - df.mtr))
+            1
+            + +((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 = (
-            0
-            + ((0 * self.iot[k].eti * df.mtr) / (1 - df.mtr))
+            1
+            + +((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
+        integral = np.trapz(g2, df.z)
         plot_df = pd.DataFrame(
             {
                 self.income_measure: df.z,
-                "Overall weight": df.g_z,
-                "Tax Base Elasticity": df.g_z - g1,
-                "Nonconstant MTRs": df.g_z - g1 - g2,
+                "Overall weight": g_weights,
+                "Tax Base Elasticity": g1,
+                "Nonconstant MTRs": g2,
             }
         )
         fig = go.Figure()
+        # add a line at y = 1
         fig.add_trace(
             go.Scatter(
-                x=plot_df[self.income_measure],
-                y=plot_df["Overall weight"],
-                fill=None,
+                x=[
+                    plot_df[self.income_measure].min(),
+                    plot_df[self.income_measure].max(),
+                ],
+                y=[1, 1],
                 mode="lines",
-                name="Overall weight",
+                line=dict(color="black", width=1, dash="dash"),
+                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],
@@ -247,12 +282,35 @@ def JJZFig4(self, policy="Current Law"):
         fig.add_trace(
             go.Scatter(
                 x=plot_df[self.income_measure],
-                y=plot_df["Nonconstant MTRs"],
-                fill="tonexty",  # fill area between trace1 and trace2
+                y=plot_df["Overall weight"],
+                fill="tonexty",
                 mode="lines",
                 name="Nonconstant MTRs",
             )
         )
+        # add a line at y=0
+        fig.add_trace(
+            go.Scatter(
+                x=plot_df[self.income_measure],
+                y=plot_df["Overall weight"],
+                mode="lines",
+                line=dict(color="black", width=1, dash="solid"),
+                name="Overall weight",
+                showlegend=False,
+            )
+        )
+        # 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$",