Merge branch 'main' into vs_module

Author: NathanielF

.github/workflows/ci.yml

@@ -24,9 +24,9 @@ jobs:
       - name: Setup environment
         run: pip install -e .[test]
       - name: Run doctests
-        run: pytest --doctest-modules --ignore=causalpy/tests/ causalpy/ --config-file=causalpy/tests/conftest.py
+        run: pytest --doctest-modules --ignore=causalpy/tests/ causalpy/ --config-file=causalpy/tests/conftest.py --no-cov
       - name: Run extra tests
-        run: pytest docs/source/.codespell/test_notebook_to_markdown.py
+        run: pytest docs/source/.codespell/test_notebook_to_markdown.py --no-cov
       - name: Run tests
         run: pytest --cov-report=xml --no-cov-on-fail
       - name: Check codespell for notebooks

.github/workflows/uml.yml

@@ -58,16 +58,25 @@ jobs:
       - name: Create PR for changes
         if: steps.changes.outputs.changes_exist == 'true'
         run: |
-          git checkout -b update-uml-diagrams
+          BRANCH_NAME="update-uml-diagrams"
+
+          # Delete the branch locally and remotely if it exists
+          git branch -D $BRANCH_NAME 2>/dev/null || true
+          git push origin --delete $BRANCH_NAME 2>/dev/null || true
+
+          # Create and push the new branch
+          git checkout -b $BRANCH_NAME
           git commit -m "Update UML Diagrams"
-          git push -u origin update-uml-diagrams
+          git push -u origin $BRANCH_NAME
+
+          # Create PR (will fail gracefully if PR already exists)
           gh pr create \
             --base main \
             --title "Update UML Diagrams" \
             --body "This PR updates the UML diagrams
             This PR was created automatically by the [UML workflow](https://github.com/pymc-labs/CausalPy/blob/main/.github/workflows/uml.yml).
             See the logs [here](https://github.com/pymc-labs/CausalPy/actions/workflows/uml.yml) for more details." \
             --label "no releasenotes" \
-            --reviewer drbenvincent
+            --reviewer drbenvincent || echo "PR may already exist"
         env:
           GH_TOKEN: ${{ github.token }}

.gitignore

@@ -14,3 +14,5 @@ dist/
 docs/build/
 docs/jupyter_execute/
 docs/source/api/generated/
+
+.cursor/

.pre-commit-config.yaml

@@ -21,16 +21,24 @@ repos:
       - id: end-of-file-fixer
         exclude_types: [svg]
       - id: check-yaml
+      - id: check-toml
+      - id: check-json
       - id: check-added-large-files
         exclude: &exclude_pattern '(iv_weak_instruments|its_lift_test)\.ipynb'
         args: ["--maxkb=1500"]
+      - id: check-merge-conflict
+      - id: check-case-conflict
+      - id: mixed-line-ending
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.5
+    rev: v0.14.8
     hooks:
       # Run the linter
       - id: ruff
         types_or: [ python, pyi, jupyter ]
         args: [ --fix ]
+        # Exclude docs/ to avoid applying strict linting rules to example notebooks
+        # Remove this exclusion if you want to enforce strict rules on documentation
+        exclude: ^docs/
       # Run the formatter
       - id: ruff-format
         types_or: [ python, pyi, jupyter ]
@@ -49,9 +57,13 @@ repos:
           # Support pyproject.toml configuration
           - tomli
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.18.2
+    rev: v1.19.0
     hooks:
       - id: mypy
         args: [--ignore-missing-imports]
         files: ^causalpy/
         additional_dependencies: [numpy>=1.20, pandas-stubs]
+  - repo: https://github.com/abravalheri/validate-pyproject
+    rev: v0.24.1
+    hooks:
+      - id: validate-pyproject

AGENTS.md

@@ -41,6 +41,15 @@
 - **Custom exceptions**: Use project-specific exceptions from `causalpy.custom_exceptions`: `FormulaException`, `DataException`, `BadIndexException`
 - **File organization**: Experiments in `causalpy/experiments/`, PyMC models in `causalpy/pymc_models.py`, scikit-learn models in `causalpy/skl_models.py`
 
+## Code quality checks
+
+- **Before committing**: Always run `pre-commit run --all-files` to ensure all checks pass (linting, formatting, type checking)
+- **Quick check**: Run `ruff check causalpy/` for fast linting feedback during development
+- **Auto-fix**: Run `ruff check --fix causalpy/` to automatically fix many linting issues
+- **Format**: Run `ruff format causalpy/` to format code according to project standards
+- **Linting rules**: Project uses strict linting (F, B, UP, C4, SIM, I) to catch bugs and enforce modern Python patterns
+- **Note**: Documentation notebooks in `docs/` are excluded from strict linting rules
+
 ## Type Checking
 
 - **Tool**: MyPy

causalpy/data/PISA18sampleScale.csv

@@ -98,4 +98,4 @@ PV1READ,Female,ESCS,METASUM,PERFEED,JOYREAD,MASTGOAL,ADAPTIVITY,TEACHINT,SCREADD
 1.382758437,0,1.112755087,1.024329216,-1.363017023,0.546695917,-0.300835494,0.488206717,0.990856172,0.19103636,-0.063690149
 -0.180166117,0,-0.903784153,1.416120964,-0.789077206,0.3589234,-2.209705037,-1.353064842,-0.668887592,0.689393741,-0.929939088
 -0.138452609,0,-1.523831485,-0.171666648,0.021680557,-0.050770988,1.521829065,-0.8111262,-0.039406676,0.66554889,-0.851380406
-0.907727459,1,0.115773982,1.024329216,1.478217432,0.461175761,-0.873789642,0.080450276,-0.668887592,-0.265544842,-0.063690149
+0.907727459,1,0.115773982,1.024329216,1.478217432,0.461175761,-0.873789642,0.080450276,-0.668887592,-0.265544842,-0.063690149

causalpy/data/simulate_data.py

@@ -308,7 +308,7 @@ def impact(x: np.ndarray) -> np.ndarray:
 
 def generate_ancova_data(
     N: int = 200,
-    pre_treatment_means: np.ndarray = np.array([10, 12]),
+    pre_treatment_means: np.ndarray | None = None,
     treatment_effect: int = 2,
     sigma: int = 1,
 ) -> pd.DataFrame:
@@ -324,6 +324,8 @@ def generate_ancova_data(
     ... )
     >>> df.to_csv(pathlib.Path.cwd() / "ancova_data.csv", index=False)  # doctest: +SKIP
     """
+    if pre_treatment_means is None:
+        pre_treatment_means = np.array([10, 12])
     group = np.random.choice(2, size=N)
     pre = np.random.normal(loc=pre_treatment_means[group])
     post = pre + treatment_effect * group + np.random.normal(size=N) * sigma

causalpy/experiments/base.py

@@ -16,7 +16,7 @@
 """
 
 from abc import abstractmethod
-from typing import Any, Literal, Union
+from typing import Any, Literal
 
 import arviz as az
 import matplotlib.pyplot as plt
@@ -54,7 +54,7 @@ class BaseExperiment:
     supports_bayes: bool
     supports_ols: bool
 
-    def __init__(self, model: Union[PyMCModel, RegressorMixin] | None = None) -> None:
+    def __init__(self, model: PyMCModel | RegressorMixin | None = None) -> None:
         # Ensure we've made any provided Scikit Learn model (as identified as being type
         # RegressorMixin) compatible with CausalPy by appending our custom methods.
         if isinstance(model, RegressorMixin):
@@ -141,7 +141,7 @@ def get_plot_data_ols(self, *args: Any, **kwargs: Any) -> pd.DataFrame:
 
     def effect_summary(
         self,
-        window: Union[Literal["post"], tuple, slice] = "post",
+        window: Literal["post"] | tuple | slice = "post",
         direction: Literal["increase", "decrease", "two-sided"] = "increase",
         alpha: float = 0.05,
         cumulative: bool = True,

causalpy/experiments/diff_in_diff.py

@@ -15,8 +15,6 @@
 Difference in differences
 """
 
-from typing import Union
-
 import arviz as az
 import numpy as np
 import pandas as pd
@@ -98,7 +96,7 @@ def __init__(
         time_variable_name: str,
         group_variable_name: str,
         post_treatment_variable_name: str = "post_treatment",
-        model: Union[PyMCModel, RegressorMixin] | None = None,
+        model: PyMCModel | RegressorMixin | None = None,
         **kwargs: dict,
     ) -> None:
         super().__init__(model=model)
@@ -234,7 +232,7 @@ def __init__(
         elif isinstance(self.model, RegressorMixin):
             # This is the coefficient on the interaction term
             # Store the coefficient into dictionary {intercept:value}
-            coef_map = dict(zip(self.labels, self.model.get_coeffs()))
+            coef_map = dict(zip(self.labels, self.model.get_coeffs(), strict=False))
             # Create and find the interaction term based on the values user provided
             interaction_term = (
                 f"{self.group_variable_name}:{self.post_treatment_variable_name}"

causalpy/experiments/instrumental_variable.py

@@ -238,7 +238,9 @@ def get_naive_OLS_fit(self) -> None:
         ols_reg = sk_lin_reg().fit(self.X, self.y)
         beta_params = list(ols_reg.coef_[0][1:])
         beta_params.insert(0, ols_reg.intercept_[0])
-        self.ols_beta_params = dict(zip(self._x_design_info.column_names, beta_params))
+        self.ols_beta_params = dict(
+            zip(self._x_design_info.column_names, beta_params, strict=False)
+        )
         self.ols_reg = ols_reg
 
     def plot(self, *args, **kwargs) -> None:  # type: ignore[override]