Skip to content

Commit de9dd6e

Browse files
anevolbapanevolbap
committed
chore: enforce stricter ruff on docs notebooks
Removes the docs/ exclusion from .pre-commit-config.yaml ruff hook and adds [tool.ruff.lint.per-file-ignores] for B018 on *.ipynb (notebooks display dataframes via bare expressions on the last cell line). Auto-fixes 27 violations across 14 notebooks (B006, B905, F541, C408, F401, I001, F811) and applies two manual fixes: - iv_vs_priors.ipynb: rename unused loop var label -> _label (B007) - staggered_did_pymc.ipynb: move matplotlib.ticker import to cell top (E402) Closes #574
1 parent deb8774 commit de9dd6e

16 files changed

Lines changed: 65 additions & 52 deletions

‎.pre-commit-config.yaml‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,6 @@ repos:
3636
- id: ruff
3737
types_or: [ python, pyi, jupyter ]
3838
args: [ --fix ]
39-
# Exclude docs/ to avoid applying strict linting rules to example notebooks
40-
# Remove this exclusion if you want to enforce strict rules on documentation
41-
exclude: ^docs/
4239
# Run the formatter
4340
- id: ruff-format
4441
types_or: [ python, pyi, jupyter ]

‎docs/source/knowledgebase/structural_causal_models.ipynb‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -929,7 +929,9 @@
929929
" return beta_outcome, beta_treatment\n",
930930
"\n",
931931
"\n",
932-
"def make_joint_model(X, Y, T, coords, priors_type=\"normal\", priors={}):\n",
932+
"def make_joint_model(X, Y, T, coords, priors_type=\"normal\", priors=None):\n",
933+
" if priors is None:\n",
934+
" priors = {}\n",
933935
" p = X.shape[1]\n",
934936
" p0 = 5.0 # pick an expected number of nonzero coeffs\n",
935937
" sigma_est = 1.0\n",
@@ -1663,6 +1665,7 @@
16631665
" \"tight_rho\",\n",
16641666
" \"tight_rho_spike_slab\",\n",
16651667
" ],\n",
1668+
" strict=False,\n",
16661669
"):\n",
16671670
" ax.axvline(3, linestyle=\"--\", color=\"k\", label=\"True Treatment Effect\")\n",
16681671
" ax.axhline(0, linestyle=\"--\", color=\"red\", label=\"True rho\")\n",
@@ -2572,6 +2575,7 @@
25722575
" \"tight_rho\",\n",
25732576
" \"tight_rho_spike_slab\",\n",
25742577
" ],\n",
2578+
" strict=False,\n",
25752579
"):\n",
25762580
" ax.axvline(3, linestyle=\"--\", color=\"k\", label=\"True Treatment Effect\")\n",
25772581
" ax.axhline(0.6, linestyle=\"--\", color=\"red\", label=\"True rho\")\n",
@@ -3497,7 +3501,9 @@
34973501
"az.plot_posterior(idata_binary_model, var_names=\"alpha\", ax=axs[2])\n",
34983502
"az.plot_posterior(idata_binary_bart_treatment_cate, var_names=\"alpha\", ax=axs[3])\n",
34993503
"for ax, title in zip(\n",
3500-
" axs, [\"bart_treatment\", \"bart_outcome\", \"no_bart\", \"bart_treatment_cate\"]\n",
3504+
" axs,\n",
3505+
" [\"bart_treatment\", \"bart_outcome\", \"no_bart\", \"bart_treatment_cate\"],\n",
3506+
" strict=False,\n",
35013507
"):\n",
35023508
" ax.axvline(3, linestyle=\"--\", color=\"k\")\n",
35033509
" ax.set_title(f\"Model: {title}\")\n",

‎docs/source/notebooks/geolift1.ipynb‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,6 @@
3535
"outputs": [],
3636
"source": [
3737
"import arviz as az\n",
38-
"import matplotlib.dates as mdates\n",
3938
"import pandas as pd\n",
4039
"\n",
4140
"import causalpy as cp"

‎docs/source/notebooks/inv_prop_latent.ipynb‎

Lines changed: 31 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -245,17 +245,19 @@
245245
" T_data,\n",
246246
" Y_data,\n",
247247
" coords,\n",
248-
" priors={\n",
249-
" \"beta_\": [0, 1],\n",
250-
" \"beta_trt\": [0, 1],\n",
251-
" \"alpha_trt\": [0, 1],\n",
252-
" \"alpha_outcome\": [0, 1],\n",
253-
" \"sigma\": 1,\n",
254-
" \"beta_ps\": [0, 1],\n",
255-
" },\n",
248+
" priors=None,\n",
256249
" noncentred=True,\n",
257250
" normal_outcome=True,\n",
258251
"):\n",
252+
" if priors is None:\n",
253+
" priors = {\n",
254+
" \"beta_\": [0, 1],\n",
255+
" \"beta_trt\": [0, 1],\n",
256+
" \"alpha_trt\": [0, 1],\n",
257+
" \"alpha_outcome\": [0, 1],\n",
258+
" \"sigma\": 1,\n",
259+
" \"beta_ps\": [0, 1],\n",
260+
" }\n",
259261
" with pm.Model(coords=coords) as model:\n",
260262
" X_data_trt = pm.Data(\"X\", X_trt, dims=(\"obs\", \"beta_trt\"))\n",
261263
" X_data_outcome = pm.Data(\"X_outcome\", X_outcome, dims=(\"obs\", \"betas\"))\n",
@@ -606,14 +608,16 @@
606608
" X_trt,\n",
607609
" T_data,\n",
608610
" coords,\n",
609-
" priors={\n",
610-
" \"beta_\": [0, 1],\n",
611-
" \"beta_trt\": [0, 1],\n",
612-
" \"alpha_trt\": [0, 1],\n",
613-
" \"alpha_outcome\": [0, 1],\n",
614-
" },\n",
611+
" priors=None,\n",
615612
" noncentred=True,\n",
616613
"):\n",
614+
" if priors is None:\n",
615+
" priors = {\n",
616+
" \"beta_\": [0, 1],\n",
617+
" \"beta_trt\": [0, 1],\n",
618+
" \"alpha_trt\": [0, 1],\n",
619+
" \"alpha_outcome\": [0, 1],\n",
620+
" }\n",
617621
" with pm.Model(coords=coords) as model_trt:\n",
618622
" X_data_trt = pm.Data(\"X\", X_trt, dims=(\"obs\", \"betas_trt\"))\n",
619623
" T_data_ = pm.Data(\"T\", T_data, dims=\"obs\")\n",
@@ -649,20 +653,22 @@
649653
" X_outcome,\n",
650654
" Y_data,\n",
651655
" coords,\n",
652-
" priors={\n",
653-
" \"beta_\": [0, 1],\n",
654-
" \"beta_trt\": [0, 1],\n",
655-
" \"alpha_trt\": [0, 1],\n",
656-
" \"alpha_outcome\": [0, 1],\n",
657-
" \"sigma\": 1,\n",
658-
" \"beta_ps\": [0, 1],\n",
659-
" },\n",
656+
" priors=None,\n",
660657
" noncentred=True,\n",
661658
" spline_component=False,\n",
662659
" propensity_score_idata=None,\n",
663660
" normal_outcome=True,\n",
664661
" winsorize_boundary=0.0,\n",
665662
"):\n",
663+
" if priors is None:\n",
664+
" priors = {\n",
665+
" \"beta_\": [0, 1],\n",
666+
" \"beta_trt\": [0, 1],\n",
667+
" \"alpha_trt\": [0, 1],\n",
668+
" \"alpha_outcome\": [0, 1],\n",
669+
" \"sigma\": 1,\n",
670+
" \"beta_ps\": [0, 1],\n",
671+
" }\n",
666672
" propensity_scores = az.extract(propensity_score_idata)[\"p\"]\n",
667673
" with pm.Model(coords=coords) as model_outcome:\n",
668674
" X_data_outcome = pm.Data(\"X_outcome\", X_outcome, dims=(\"obs\", \"betas\"))\n",
@@ -908,13 +914,11 @@
908914
" X_outcome,\n",
909915
" Y_data,\n",
910916
" coords,\n",
911-
" priors={\n",
912-
" \"beta_\": [0, 1],\n",
913-
" \"alpha_outcome\": [0, 1],\n",
914-
" \"sigma\": 1,\n",
915-
" },\n",
917+
" priors=None,\n",
916918
" noncentred=True,\n",
917919
"):\n",
920+
" if priors is None:\n",
921+
" priors = {\"beta_\": [0, 1], \"alpha_outcome\": [0, 1], \"sigma\": 1}\n",
918922
" with pm.Model(coords=coords) as reg_model:\n",
919923
" X_data_outcome = pm.Data(\"X_outcome\", X_outcome, dims=(\"obs\", \"betas\"))\n",
920924
" Y_data_ = pm.Data(\"Y\", Y_data, dims=\"obs\")\n",

‎docs/source/notebooks/its_lift_test.ipynb‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -416,7 +416,7 @@
416416
"channels = [\"tv_brand_awareness\", \"tv_promo\", \"digital_paid\", \"radio\"]\n",
417417
"colors = [\"#1f77b4\", \"#d62728\", \"#2ca02c\", \"#ff7f0e\"]\n",
418418
"\n",
419-
"for ax, channel, color in zip(axes, channels, colors):\n",
419+
"for ax, channel, color in zip(axes, channels, colors, strict=False):\n",
420420
" ax.plot(df_spend.index, df_spend[channel], color=color, linewidth=1.5)\n",
421421
" ax.axvline(\n",
422422
" intervention_start,\n",

‎docs/source/notebooks/its_pymc_comparative.ipynb‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -491,7 +491,7 @@
491491
" f\"Correlation: {corr:.3f}\",\n",
492492
" transform=ax.transAxes,\n",
493493
" verticalalignment=\"top\",\n",
494-
" bbox=dict(boxstyle=\"round\", facecolor=\"wheat\", alpha=0.5),\n",
494+
" bbox={\"boxstyle\": \"round\", \"facecolor\": \"wheat\", \"alpha\": 0.5},\n",
495495
")\n",
496496
"plt.tight_layout()\n",
497497
"plt.show()"

‎docs/source/notebooks/iv_vs_priors.ipynb‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3067,7 +3067,7 @@
30673067
"ax1 = axs[1]\n",
30683068
"y_positions = np.arange(len(elas))\n",
30693069
"\n",
3070-
"for i, (label, values) in enumerate(elas.items()):\n",
3070+
"for i, (_label, values) in enumerate(elas.items()):\n",
30713071
" y = np.full_like(values, i, dtype=float)\n",
30723072
" ax.scatter(values, y, alpha=0.4)\n",
30733073
"\n",

‎docs/source/notebooks/iv_weak_instruments.ipynb‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -832,7 +832,9 @@
832832
"]\n",
833833
"\n",
834834
"\n",
835-
"def make_reg_model(covariate_df, target=\"log_wage\", prior_beta=[0, 1]):\n",
835+
"def make_reg_model(covariate_df, target=\"log_wage\", prior_beta=None):\n",
836+
" if prior_beta is None:\n",
837+
" prior_beta = [0, 1]\n",
836838
" coords = {\"covariates\": list(covariate_df.columns)}\n",
837839
" with pm.Model(coords=coords) as reg_model:\n",
838840
" X = pm.Data(\"X\", covariate_df)\n",
@@ -3285,7 +3287,7 @@
32853287
" axs[0].hist(np.exp(y), bins=30, alpha=0.2, color=\"blue\", label=\"Observed Histogram\")\n",
32863288
" axs[0].set_xlabel(\"Predicted Wage\", fontsize=9)\n",
32873289
" axs[0].legend()\n",
3288-
" for c, ed in zip([\"red\", \"purple\", \"blue\"], [2, 10, 18]):\n",
3290+
" for c, ed in zip([\"red\", \"purple\", \"blue\"], [2, 10, 18], strict=False):\n",
32893291
" temp = covariate_df.copy()\n",
32903292
" temp[\"education\"] = ed\n",
32913293
" means = []\n",

‎docs/source/notebooks/panel_fixed_effects.ipynb‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -896,12 +896,12 @@
896896
"print(\"=\" * 60)\n",
897897
"print(\"COMPARISON: One-Way vs Two-Way Fixed Effects\")\n",
898898
"print(\"=\" * 60)\n",
899-
"print(f\"\\nTrue treatment effect: 3.0\")\n",
900-
"print(f\"\\nOne-way (unit FE only):\")\n",
899+
"print(\"\\nTrue treatment effect: 3.0\")\n",
900+
"print(\"\\nOne-way (unit FE only):\")\n",
901901
"print(\n",
902902
" f\" Treatment coefficient: {result_demeaned.model.idata.posterior['beta'].sel(coeffs='treatment').mean().values:.3f}\"\n",
903903
")\n",
904-
"print(f\"\\nTwo-way (unit + time FE):\")\n",
904+
"print(\"\\nTwo-way (unit + time FE):\")\n",
905905
"print(\n",
906906
" f\" Treatment coefficient: {result_twoway.model.idata.posterior['beta'].sel(coeffs='treatment').mean().values:.3f}\"\n",
907907
")\n",

‎docs/source/notebooks/piecewise_its_pymc.ipynb‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
"metadata": {},
1414
"outputs": [],
1515
"source": [
16-
"import arviz as az\n",
1716
"import matplotlib.pyplot as plt\n",
1817
"import numpy as np\n",
1918
"\n",

0 commit comments

Comments
 (0)