Merge pull request PSLmodels#101 from jdebacker/default_plots

Merging
jdebacker · Apr 17, 2024 · 13a3da8 · 13a3da8
2 parents 4c0e3d8 + f01b59b
commit 13a3da8
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 27 deletions.
diff --git a/ogusa/bequest_transmission.py b/ogusa/bequest_transmission.py
@@ -8,7 +8,9 @@
 
 
 def get_bequest_matrix(
-    J=7, lambdas=np.array([0.25, 0.25, 0.2, 0.1, 0.1, 0.09, 0.01]), graphs=True
+    J=7,
+    lambdas=np.array([0.25, 0.25, 0.2, 0.1, 0.1, 0.09, 0.01]),
+    graphs=False,
 ):
     """
     Returns S x J matrix representing the fraction of aggregate

diff --git a/ogusa/calibrate.py b/ogusa/calibrate.py
@@ -68,6 +68,7 @@ def __init__(
                 99,
                 initial_data_year=p.start_year - 1,
                 final_data_year=p.start_year,
+                GraphDiag=False,
             )
 
             # demographics for 80 period lives (needed for getting e below)
@@ -79,6 +80,7 @@ def __init__(
                 99,
                 initial_data_year=p.start_year - 1,
                 final_data_year=p.start_year,
+                GraphDiag=False,
             )
 
             # earnings profiles

diff --git a/ogusa/labor.py b/ogusa/labor.py
@@ -147,7 +147,7 @@ def VCV_moments(cps, n, bin_weights, S):
     return VCV
 
 
-def labor_data_graphs(weighted, S, J, output_dir):
+def labor_data_graphs(weighted, S, J, output_dir=None):
     """
     Plot labor supply data.
 
@@ -166,27 +166,34 @@ def labor_data_graphs(weighted, S, J, output_dir):
     X, Y = np.meshgrid(domain, Jgrid)
     cmap2 = matplotlib.cm.get_cmap("summer")
 
-    plt.plot(domain, weighted, color="black", label="Data")
-    plt.plot(
-        np.linspace(76, 100, 23),
-        extension,
-        color="black",
-        linestyle="-.",
-        label="Extrapolation",
-    )
-    plt.plot(np.linspace(65, 76, 11), to_dot, linestyle="--", color="black")
-    plt.axvline(x=76, color="black", linestyle="--")
-    plt.xlabel(r"age-$s$")
-    plt.ylabel(r"individual labor supply $/bar{l}_s$")
-    plt.legend()
-    plt.savefig(
-        os.path.join(baseline_dir, "Demographics/labor_dist_data_withfit.png")
-    )
-
-    fig10 = plt.figure()
-    ax10 = fig10.add_subplot(projection="3d")
-    ax10.plot_surface(X, Y, lab_mat_basic.T, rstride=1, cstride=2, cmap=cmap2)
-    ax10.set_xlabel(r"age-$s$")
-    ax10.set_ylabel(r"ability type -$j$")
-    ax10.set_zlabel(r"labor $e_j(s)$")
-    plt.savefig(os.path.join(baseline_dir, "Demographics/data_labor_dist"))
+    if output_dir:
+        plt.plot(domain, weighted, color="black", label="Data")
+        plt.plot(
+            np.linspace(76, 100, 23),
+            extension,
+            color="black",
+            linestyle="-.",
+            label="Extrapolation",
+        )
+        plt.plot(
+            np.linspace(65, 76, 11), to_dot, linestyle="--", color="black"
+        )
+        plt.axvline(x=76, color="black", linestyle="--")
+        plt.xlabel(r"age-$s$")
+        plt.ylabel(r"individual labor supply $/bar{l}_s$")
+        plt.legend()
+        plt.savefig(
+            os.path.join(
+                output_dir, "Demographics/labor_dist_data_withfit.png"
+            )
+        )
+
+        fig10 = plt.figure()
+        ax10 = fig10.add_subplot(projection="3d")
+        ax10.plot_surface(
+            X, Y, lab_mat_basic.T, rstride=1, cstride=2, cmap=cmap2
+        )
+        ax10.set_xlabel(r"age-$s$")
+        ax10.set_ylabel(r"ability type -$j$")
+        ax10.set_zlabel(r"labor $e_j(s)$")
+        plt.savefig(os.path.join(output_dir, "Demographics/data_labor_dist"))
diff --git a/ogusa/transfer_distribution.py b/ogusa/transfer_distribution.py
@@ -8,7 +8,9 @@
 
 
 def get_transfer_matrix(
-    J=7, lambdas=np.array([0.25, 0.25, 0.2, 0.1, 0.1, 0.09, 0.01]), graphs=True
+    J=7,
+    lambdas=np.array([0.25, 0.25, 0.2, 0.1, 0.1, 0.09, 0.01]),
+    graphs=False,
 ):
     """
     Compute SxJ matrix representing the distribution of aggregate