better print output alignment when printing model coefficient labels …

…and values
pymc-labs · Jun 19, 2024 · 22f5008 · 22f5008
1 parent 6ca0842
commit 22f5008
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 37 deletions.
diff --git a/causalpy/skl_experiments.py b/causalpy/skl_experiments.py
@@ -62,8 +62,15 @@ def print_coefficients(self, round_to=None) -> None:
             Number of decimals used to round results. Defaults to 2. Use "None" to return raw numbers.
         """
         print("Model coefficients:")
+        # Determine the width of the longest label
+        max_label_length = max(len(name) for name in self.labels)
+        # Print each coefficient with formatted alignment
         for name, val in zip(self.labels, self.model.coef_[0]):
-            print(f"\t{name}\t\t{round_num(val, round_to)}")
+            # Left-align the name
+            formatted_name = f"{name:<{max_label_length}}"
+            # Right-align the value with width 10
+            formatted_val = f"{round_num(val, round_to):>10}"
+            print(f"  {formatted_name}\t{formatted_val}")
 
 
 class PrePostFit(ExperimentalDesign, PrePostFitDataValidator):
@@ -729,7 +736,7 @@ def plot(self, round_to=None):
         ax.legend(fontsize=LEGEND_FONT_SIZE)
         return (fig, ax)
 
-    def summary(self, round_to=None):
+    def summary(self, round_to=None) -> None:
         """
         Print text output summarising the results
 
@@ -742,4 +749,5 @@ def summary(self, round_to=None):
         print(f"Threshold on running variable: {self.treatment_threshold}")
         print("\nResults:")
         print(f"Discontinuity at threshold = {self.discontinuity_at_threshold:.2f}")
+        print("\n")
         self.print_coefficients(round_to)
diff --git a/docs/source/notebooks/did_skl.ipynb b/docs/source/notebooks/did_skl.ipynb
@@ -79,10 +79,10 @@
       "Results:\n",
       "Causal impact = 0.504\n",
       "Model coefficients:\n",
-      "\tIntercept\t\t0\n",
-      "\tpost_treatment[T.True]\t\t0.986\n",
-      "\tgroup\t\t0.162\n",
-      "\tgroup:post_treatment[T.True]\t\t0.504\n"
+      "  Intercept                   \t         0\n",
+      "  post_treatment[T.True]      \t     0.986\n",
+      "  group                       \t     0.162\n",
+      "  group:post_treatment[T.True]\t     0.504\n"
      ]
     }
    ],

diff --git a/docs/source/notebooks/its_skl.ipynb b/docs/source/notebooks/its_skl.ipynb
@@ -211,19 +211,19 @@
       "==================================Pre-Post Fit==================================\n",
       "Formula: y ~ 1 + t + C(month)\n",
       "Model coefficients:\n",
-      "\tIntercept\t\t0\n",
-      "\tC(month)[T.2]\t\t2.85\n",
-      "\tC(month)[T.3]\t\t1.16\n",
-      "\tC(month)[T.4]\t\t7.15\n",
-      "\tC(month)[T.5]\t\t15\n",
-      "\tC(month)[T.6]\t\t24.8\n",
-      "\tC(month)[T.7]\t\t18.2\n",
-      "\tC(month)[T.8]\t\t33.5\n",
-      "\tC(month)[T.9]\t\t16.2\n",
-      "\tC(month)[T.10]\t\t9.19\n",
-      "\tC(month)[T.11]\t\t6.28\n",
-      "\tC(month)[T.12]\t\t0.578\n",
-      "\tt\t\t0.21\n"
+      "  Intercept     \t         0\n",
+      "  C(month)[T.2] \t      2.85\n",
+      "  C(month)[T.3] \t      1.16\n",
+      "  C(month)[T.4] \t      7.15\n",
+      "  C(month)[T.5] \t        15\n",
+      "  C(month)[T.6] \t      24.8\n",
+      "  C(month)[T.7] \t      18.2\n",
+      "  C(month)[T.8] \t      33.5\n",
+      "  C(month)[T.9] \t      16.2\n",
+      "  C(month)[T.10]\t      9.19\n",
+      "  C(month)[T.11]\t      6.28\n",
+      "  C(month)[T.12]\t     0.578\n",
+      "  t             \t      0.21\n"
      ]
     }
    ],

diff --git a/docs/source/notebooks/rd_skl.ipynb b/docs/source/notebooks/rd_skl.ipynb
@@ -184,9 +184,9 @@
       "Results:\n",
       "Discontinuity at threshold = 0.19\n",
       "Model coefficients:\n",
-      "\tIntercept\t\t0\n",
-      "\ttreated[T.True]\t\t0.19\n",
-      "\tx\t\t1.23\n"
+      "  Intercept      \t         0\n",
+      "  treated[T.True]\t      0.19\n",
+      "  x              \t      1.23\n"
      ]
     }
    ],
@@ -264,10 +264,10 @@
       "Results:\n",
       "Discontinuity at threshold = 0.92\n",
       "Model coefficients:\n",
-      "\tIntercept\t\t0\n",
-      "\ttreated[T.True]\t\t2.47\n",
-      "\tx\t\t1.32\n",
-      "\tx:treated[T.True]\t\t-3.11\n"
+      "  Intercept        \t         0\n",
+      "  treated[T.True]  \t      2.47\n",
+      "  x                \t      1.32\n",
+      "  x:treated[T.True]\t     -3.11\n"
      ]
     }
    ],
@@ -467,14 +467,14 @@
       "Results:\n",
       "Discontinuity at threshold = 0.41\n",
       "Model coefficients:\n",
-      "\tIntercept\t\t0\n",
-      "\ttreated[T.True]\t\t0.407\n",
-      "\tbs(x, df=6)[0]\t\t-0.594\n",
-      "\tbs(x, df=6)[1]\t\t-1.07\n",
-      "\tbs(x, df=6)[2]\t\t0.278\n",
-      "\tbs(x, df=6)[3]\t\t1.65\n",
-      "\tbs(x, df=6)[4]\t\t1.03\n",
-      "\tbs(x, df=6)[5]\t\t0.567\n"
+      "  Intercept      \t         0\n",
+      "  treated[T.True]\t     0.407\n",
+      "  bs(x, df=6)[0] \t    -0.594\n",
+      "  bs(x, df=6)[1] \t     -1.07\n",
+      "  bs(x, df=6)[2] \t     0.278\n",
+      "  bs(x, df=6)[3] \t      1.65\n",
+      "  bs(x, df=6)[4] \t      1.03\n",
+      "  bs(x, df=6)[5] \t     0.567\n"
      ]
     }
    ],

diff --git a/docs/source/notebooks/rd_skl_drinking.ipynb b/docs/source/notebooks/rd_skl_drinking.ipynb
@@ -108,10 +108,12 @@
       "\n",
       "Results:\n",
       "Discontinuity at threshold = 7.66\n",
+      "\n",
+      "\n",
       "Model coefficients:\n",
-      "\tIntercept\t\t0\n",
-      "\ttreated[T.True]\t\t7.7\n",
-      "\tage\t\t-0.97\n"
+      "  Intercept      \t         0\n",
+      "  treated[T.True]\t       7.7\n",
+      "  age            \t     -0.97\n"
      ]
     }
    ],