Merge branch 'main' into kink

.pre-commit-config.yaml

@@ -1,43 +1,29 @@
 # See https://pre-commit.com for more information
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
-    - repo: https://github.com/pre-commit/pre-commit-hooks
-      rev: v4.4.0
-      hooks:
-          - id: trailing-whitespace
-            exclude_types: [svg]
-          - id: end-of-file-fixer
-            exclude_types: [svg]
-          - id: check-yaml
-          - id: check-added-large-files
-            args: ['--maxkb=1500']
-    - repo: https://github.com/pycqa/isort
-      rev: 5.12.0
-      hooks:
-          - id: isort
-            args: [--profile, black]
-            types: [python]
-    - repo: https://github.com/ambv/black
-      rev: 23.3.0
-      hooks:
-          - id: black
-    - repo: https://github.com/pycqa/flake8
-      rev: 6.0.0
-      hooks:
-          - id: flake8
-    - repo: https://github.com/nbQA-dev/nbQA
-      rev: 1.7.0
-      hooks:
-        - id: nbqa-black
-          # additional_dependencies: [jupytext]  # optional, only if you're using Jupytext
-        - id: nbqa-pyupgrade
-          args: ["--py37-plus"]
-        - id: nbqa-isort
-          args: ["--float-to-top"]
-    - repo: https://github.com/econchick/interrogate
-      rev: 1.5.0
-      hooks:
-        - id: interrogate
-          # needed to make excludes in pyproject.toml work
-          # see here https://github.com/econchick/interrogate/issues/60#issuecomment-735436566
-          pass_filenames: false
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.4.0
+    hooks:
+      - id: trailing-whitespace
+        exclude_types: [svg]
+      - id: end-of-file-fixer
+        exclude_types: [svg]
+      - id: check-yaml
+      - id: check-added-large-files
+        args: ['--maxkb=1500']
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: v0.1.4
+    hooks:
+      - id: ruff
+      - id: ruff-format
+  - repo: https://github.com/nbQA-dev/nbQA
+    rev: 1.7.0
+    hooks:
+      - id: nbqa-ruff
+  - repo: https://github.com/econchick/interrogate
+    rev: 1.5.0
+    hooks:
+      - id: interrogate
+        # needed to make excludes in pyproject.toml work
+        # see here https://github.com/econchick/interrogate/issues/60#issuecomment-735436566
+        pass_filenames: false

README.md

@@ -5,7 +5,7 @@
 ----
 
 ![Build](https://github.com/pymc-labs/CausalPy/workflows/ci/badge.svg)
-[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/charliermarsh/ruff/main/assets/badge/v1.json)]([https://codecov.io/gh/pymc-labs/CausalPy](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/charliermarsh/ruff/main/assets/badge/v1.json))
 [![PyPI version](https://badge.fury.io/py/CausalPy.svg)](https://badge.fury.io/py/CausalPy)
 ![GitHub Repo stars](https://img.shields.io/github/stars/pymc-labs/causalpy?style=social)
 ![Read the Docs](https://img.shields.io/readthedocs/causalpy)

causalpy/__init__.py

@@ -1,11 +1,20 @@
 import arviz as az
 
-import causalpy.pymc_experiments
-import causalpy.pymc_models
-import causalpy.skl_experiments
-import causalpy.skl_models
+from causalpy import pymc_experiments
+from causalpy import pymc_models
+from causalpy import skl_experiments
+from causalpy import skl_models
 from causalpy.version import __version__
 
 from .data import load_data
 
 az.style.use("arviz-darkgrid")
+
+__all__ = [
+    "pymc_experiments",
+    "pymc_models",
+    "skl_experiments",
+    "skl_models",
+    "load_data",
+    "__version__",
+]

docs/source/notebooks/did_skl.ipynb

@@ -13,7 +13,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import xarray as xr\n",
     "from sklearn.linear_model import LinearRegression\n",
     "\n",
     "import causalpy as cp"

docs/source/notebooks/geolift1.ipynb

@@ -346,10 +346,14 @@
     }
    ],
    "source": [
+    "formula = \"\"\"\n",
+    "    \"Denmark ~ 0 + Austria + Belgium + Bulgaria + Croatia + Cyprus + Czech_Republic\"\n",
+    "\"\"\"\n",
+    "\n",
     "result = cp.pymc_experiments.SyntheticControl(\n",
     "    df,\n",
     "    treatment_time,\n",
-    "    formula=\"Denmark ~ 0 + Austria + Belgium + Bulgaria + Croatia + Cyprus + Czech_Republic\",\n",
+    "    formula=formula,\n",
     "    model=cp.pymc_models.WeightedSumFitter(\n",
     "        sample_kwargs={\"target_accept\": 0.95, \"random_seed\": seed}\n",
     "    ),\n",

pyproject.toml

@@ -63,15 +63,7 @@ docs = [
     "statsmodels",
     "sphinxcontrib-bibtex",
 ]
-lint = [
-    "black",
-    "flake8",
-    "interrogate",
-    "isort",
-    "nbqa",
-    "pre-commit",
-    "pylint",
-]
+lint = ["interrogate", "nbqa", "pre-commit", "ruff"]
 test = ["pytest", "pytest-cov"]
 
 [metadata]
@@ -83,13 +75,6 @@ Homepage = "https://github.com/pymc-labs/CausalPy"
 "Bug Reports" = "https://github.com/pymc-labs/CausalPy/issues"
 "Source" = "https://github.com/pymc-labs/CausalPy"
 
-[tool.black]
-line_length = 88
-
-[tool.isort]
-profile = "black"
-skip_gitignore = true
-
 [tool.pytest.ini_options]
 addopts = [
     "-vv",