simplification to PyMC.print_coefficients

Author: drbenvincent

causalpy/pymc_models.py

@@ -237,6 +237,23 @@ def print_row(
             formatted_val = f"{round_num(coeff_samples.mean().data, round_to)}, 94% HDI [{round_num(coeff_samples.quantile(0.03).data, round_to)}, {round_num(coeff_samples.quantile(1 - 0.03).data, round_to)}]"  # noqa: E501
             print(f"  {formatted_name}  {formatted_val}")
 
+        def print_coefficients_for_unit(
+            unit_coeffs: xr.DataArray,
+            unit_sigma: xr.DataArray,
+            labels: list,
+            round_to: int,
+        ) -> None:
+            """Print coefficients for a single unit"""
+            # Determine the width of the longest label
+            max_label_length = max(len(name) for name in labels + ["sigma"])
+
+            for name in labels:
+                coeff_samples = unit_coeffs.sel(coeffs=name)
+                print_row(max_label_length, name, coeff_samples, round_to)
+
+            # Add coefficient for measurement std
+            print_row(max_label_length, "sigma", unit_sigma, round_to)
+
         print("Model coefficients:")
         coeffs = az.extract(self.idata.posterior, var_names="beta")
 
@@ -247,32 +264,16 @@ def print_row(
             for unit in treated_units:
                 print(f"\nTreated unit: {unit}")
                 unit_coeffs = coeffs.sel(treated_units=unit)
-
-                # Determine the width of the longest label
-                max_label_length = max(len(name) for name in labels + ["sigma"])
-
-                for name in labels:
-                    coeff_samples = unit_coeffs.sel(coeffs=name)
-                    print_row(max_label_length, name, coeff_samples, round_to or 2)
-
-                # Add coefficient for measurement std for this unit
                 unit_sigma = az.extract(self.idata.posterior, var_names="sigma").sel(
                     treated_units=unit
                 )
-                print_row(max_label_length, "sigma", unit_sigma, round_to or 2)
+                print_coefficients_for_unit(
+                    unit_coeffs, unit_sigma, labels, round_to or 2
+                )
         else:
             # Single treated unit case (backward compatibility)
-            # Determine the width of the longest label
-            max_label_length = max(len(name) for name in labels + ["sigma"])
-
-            for name in labels:
-                coeff_samples = coeffs.sel(coeffs=name)
-                print_row(max_label_length, name, coeff_samples, round_to or 2)
-
-            # Add coefficient for measurement std
-            coeff_samples = az.extract(self.idata.posterior, var_names="sigma")
-            name = "sigma"
-            print_row(max_label_length, name, coeff_samples, round_to or 2)
+            unit_sigma = az.extract(self.idata.posterior, var_names="sigma")
+            print_coefficients_for_unit(coeffs, unit_sigma, labels, round_to or 2)
 
 
 class LinearRegression(PyMCModel):

docs/source/_static/interrogate_badge.svg

@@ -626,7 +626,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "23d7ea6ad64243c5ab00ac8b1ad0d19b",
+       "model_id": "d7dc0c0854004ee8b87dbe88f4484e83",
        "version_major": 2,
        "version_minor": 0
       },
@@ -848,7 +848,7 @@
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "31dce179e1a047ddbfa198b8f82d97d4",
+       "model_id": "8d6a96eecd9d4e4aba5894120d647af4",
        "version_major": 2,
        "version_minor": 0
       },
@@ -873,7 +873,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 36 seconds.\n",
+      "Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 35 seconds.\n",
       "Sampling: [beta, sigma, y_hat]\n",
       "Sampling: [y_hat]\n",
       "Sampling: [y_hat]\n",
@@ -941,7 +941,7 @@
        "<text text-anchor=\"middle\" x=\"131\" y=\"-404.74\" font-family=\"Times,serif\" font-size=\"14.00\">Data</text>\n",
        "</g>\n",
        "<!-- mu -->\n",
-       "<g id=\"node4\" class=\"node\">\n",
+       "<g id=\"node2\" class=\"node\">\n",
        "<title>mu</title>\n",
        "<polygon fill=\"none\" stroke=\"black\" points=\"193.12,-325.23 102.88,-325.23 102.88,-267.73 193.12,-267.73 193.12,-325.23\"/>\n",
        "<text text-anchor=\"middle\" x=\"148\" y=\"-307.93\" font-family=\"Times,serif\" font-size=\"14.00\">mu</text>\n",
@@ -955,15 +955,21 @@
        "<polygon fill=\"black\" stroke=\"black\" points=\"146.23,-337.34 144.08,-326.96 139.29,-336.41 146.23,-337.34\"/>\n",
        "</g>\n",
        "<!-- y_hat -->\n",
-       "<g id=\"node2\" class=\"node\">\n",
+       "<g id=\"node3\" class=\"node\">\n",
        "<title>y_hat</title>\n",
        "<ellipse fill=\"lightgrey\" stroke=\"black\" cx=\"150\" cy=\"-174.66\" rx=\"41.01\" ry=\"40.66\"/>\n",
        "<text text-anchor=\"middle\" x=\"150\" y=\"-186.11\" font-family=\"Times,serif\" font-size=\"14.00\">y_hat</text>\n",
        "<text text-anchor=\"middle\" x=\"150\" y=\"-169.61\" font-family=\"Times,serif\" font-size=\"14.00\">~</text>\n",
        "<text text-anchor=\"middle\" x=\"150\" y=\"-153.11\" font-family=\"Times,serif\" font-size=\"14.00\">Normal</text>\n",
        "</g>\n",
+       "<!-- mu&#45;&gt;y_hat -->\n",
+       "<g id=\"edge5\" class=\"edge\">\n",
+       "<title>mu&#45;&gt;y_hat</title>\n",
+       "<path fill=\"none\" stroke=\"black\" d=\"M148.47,-267.38C148.67,-255.28 148.91,-240.73 149.15,-226.83\"/>\n",
+       "<polygon fill=\"black\" stroke=\"black\" points=\"152.64,-226.94 149.31,-216.88 145.65,-226.82 152.64,-226.94\"/>\n",
+       "</g>\n",
        "<!-- y -->\n",
-       "<g id=\"node3\" class=\"node\">\n",
+       "<g id=\"node4\" class=\"node\">\n",
        "<title>y</title>\n",
        "<path fill=\"lightgrey\" stroke=\"black\" d=\"M165,-98C165,-98 135,-98 135,-98 129,-98 123,-92 123,-86 123,-86 123,-52.5 123,-52.5 123,-46.5 129,-40.5 135,-40.5 135,-40.5 165,-40.5 165,-40.5 171,-40.5 177,-46.5 177,-52.5 177,-52.5 177,-86 177,-86 177,-92 171,-98 165,-98\"/>\n",
        "<text text-anchor=\"middle\" x=\"150\" y=\"-80.7\" font-family=\"Times,serif\" font-size=\"14.00\">y</text>\n",
@@ -976,12 +982,6 @@
        "<path fill=\"none\" stroke=\"black\" d=\"M150,-133.68C150,-125.83 150,-117.6 150,-109.76\"/>\n",
        "<polygon fill=\"black\" stroke=\"black\" points=\"153.5,-109.79 150,-99.79 146.5,-109.79 153.5,-109.79\"/>\n",
        "</g>\n",
-       "<!-- mu&#45;&gt;y_hat -->\n",
-       "<g id=\"edge5\" class=\"edge\">\n",
-       "<title>mu&#45;&gt;y_hat</title>\n",
-       "<path fill=\"none\" stroke=\"black\" d=\"M148.47,-267.38C148.67,-255.28 148.91,-240.73 149.15,-226.83\"/>\n",
-       "<polygon fill=\"black\" stroke=\"black\" points=\"152.64,-226.94 149.31,-216.88 145.65,-226.82 152.64,-226.94\"/>\n",
-       "</g>\n",
        "<!-- beta -->\n",
        "<g id=\"node5\" class=\"node\">\n",
        "<title>beta</title>\n",
@@ -1014,7 +1014,7 @@
        "</svg>\n"
       ],
       "text/plain": [
-       "<graphviz.graphs.Digraph at 0x1273e0ec0>"
+       "<graphviz.graphs.Digraph at 0x106560ec0>"
       ]
      },
      "execution_count": 15,