improving docstring

Author: JeanVanDyk

.github/workflows/ci.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.11", "3.12", "3.13"]
 
     steps:
       - uses: actions/checkout@v4

.pre-commit-config.yaml

@@ -25,7 +25,7 @@ repos:
         exclude: &exclude_pattern 'iv_weak_instruments.ipynb'
         args: ["--maxkb=1500"]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.12.10
+    rev: v0.13.1
     hooks:
       # Run the linter
       - id: ruff

causalpy/pymc_models.py

@@ -755,19 +755,29 @@ class InterventionTimeEstimator(PyMCModel):
     r"""
     Custom PyMC model to estimate the time an intervention took place.
 
-    defines the PyMC model :
+    This model implements three types of changepoints: level shift, trend change, and impulse response.
+
+    In words:
+    - `beta` represents the regression coefficients for the baseline signal `μ`.
+    - `tau` is the changepoint time, and `w` is a sigmoid function to smooth the transition.
+    - `level` and `trend` model a linear shift after the changepoint.
+    - `A` and `lambda` define an impulse response at the changepoint.
+    - `mu_in` combines the level, trend, and impulse contributions.
+    - `mu_ts` is the total mean, including baseline and intervention.
+    - `sigma` is the observation noise.
+    - Finally, `y` is drawn from a Normal distribution with mean `mu_ts` and standard deviation `sigma`.
 
     .. math::
-        \beta &\sim \mathrm{Normal}(0, 1) \\
+        \beta &\sim \mathrm{Normal}(0, 5) \\
         \mu &= \beta \cdot X\\
         \\
-        \tau &\sim \mathrm{Uniform}(0, 1) \\
+        \tau &\sim \mathrm{Uniform}(\text{lower_bound}, \text{upper_bound}) \\
         w &= sigmoid(t-\tau) \\
         \\
-        \text{level} &\sim \mathrm{Normal}(0, 1) \\
-        \text{trend} &\sim \mathrm{Normal}(0, 1) \\
-        A &\sim \mathrm{Normal}(0, 1) \\
-        \lambda &\sim \mathrm{HalfNormal}(0, 1) \\
+        \text{level} &\sim \mathrm{Normal}(0, 5) \\
+        \text{trend} &\sim \mathrm{Normal}(0, 0.5) \\
+        A &\sim \mathrm{Normal}(0, 5) \\
+        \lambda &\sim \mathrm{HalfNormal}(0, 5) \\
         \text{impulse} &= A \cdot exp(-\lambda \cdot |t-\tau|) \\
         \mu_{in} &= \text{level} + \text{trend} \cdot (t-\tau) + \text{impulse}\\
         \\

docs/source/index.md

@@ -13,52 +13,50 @@ A Python package focussing on causal inference for quasi-experiments. The packag
 To get the latest release you can use pip:
 
 ```bash
-   pip install CausalPy
+pip install CausalPy
 ```
 
 or conda:
 
 ```bash
-   conda install causalpy -c conda-forge
+conda install causalpy -c conda-forge
 ```
 
 Alternatively, if you want the very latest version of the package you can install from GitHub:
 
 ```bash
-   pip install git+https://github.com/pymc-labs/CausalPy.git
+pip install git+https://github.com/pymc-labs/CausalPy.git
 ```
 
 ## Quickstart
 
 ```python
-
-   import causalpy as cp
-   import matplotlib.pyplot as plt
-
-
-   # Import and process data
-   df = (
-      cp.load_data("drinking")
-      .rename(columns={"agecell": "age"})
-      .assign(treated=lambda df_: df_.age > 21)
-      )
-
-   # Run the analysis
-   result = cp.RegressionDiscontinuity(
-      df,
-      formula="all ~ 1 + age + treated",
-      running_variable_name="age",
-      model=cp.pymc_models.LinearRegression(),
-      treatment_threshold=21,
-      )
-
-   # Visualize outputs
-   fig, ax = result.plot();
-
-   # Get a results summary
-   result.summary()
-
-   plt.show()
+import causalpy as cp
+import matplotlib.pyplot as plt
+
+
+# Import and process data
+df = (
+    cp.load_data("drinking")
+    .rename(columns={"agecell": "age"})
+    .assign(treated=lambda df_: df_.age > 21)
+)
+
+# Run the analysis
+result = cp.RegressionDiscontinuity(
+    df,
+    formula="all ~ 1 + age + treated",
+    running_variable_name="age",
+    model=cp.pymc_models.LinearRegression(),
+    treatment_threshold=21,
+)
+
+# Visualize outputs
+fig, ax = result.plot()
+# Get a results summary
+result.summary()
+
+plt.show()
 ```
 
 ## Videos

docs/source/notebooks/its_no_treatment_time.ipynb

@@ -18,7 +18,7 @@ description = "Causal inference for quasi-experiments in Python"
 readme = "README.md"
 license = { file = "LICENSE" }
 authors = [{ name = "Ben Vincent", email = "[email protected]" }]
-requires-python = ">=3.10"
+requires-python = ">=3.11"
 
 # This field lists other packages that your project depends on to run.
 # Any package you put here will be installed by pip when your project is