make submodules visible in __init__.py, changed API, reran notebooks

Author: drbenvincent

causalpy/__init__.py

@@ -0,0 +1,6 @@
+import causalpy.pymc_experiments
+import causalpy.pymc_models
+import causalpy.skl_experiments
+import causalpy.skl_models
+
+from .data import load_data

causalpy/data/__init__.py

@@ -0,0 +1,4 @@
+"""Code for loading datasets."""
+from .datasets import load_data
+
+__all__ = ["load_data"]

causalpy/data/datasets.py

@@ -0,0 +1,33 @@
+import os
+import pathlib
+
+import pandas as pd
+
+import causalpy as cp
+
+DATASETS = {
+    "banks": {"filename": "banks.csv"},
+    "did": {"filename": "did.csv"},
+    "drinking": {"filename": "drinking.csv"},
+    "its": {"filename": "its.csv"},
+    "its simple": {"filename": "its_simple.csv"},
+    "rd": {"filename": "regression_discontinuity.csv"},
+    "sc": {"filename": "synthetic_control.csv"},
+}
+
+
+def get_data_home():
+    """Return the path of the data directory"""
+    return pathlib.Path(cp.__file__).parents[1] / "causalpy" / "data"
+
+
+def load_data(dataset: str = None):
+
+    if dataset in DATASETS:
+
+        data_dir = get_data_home()
+        datafile = DATASETS[dataset]
+        file_path = data_dir / datafile["filename"]
+        return pd.read_csv(file_path)
+    else:
+        raise ValueError(f"Dataset {dataset} not found!")

docs/notebooks/did_pymc.ipynb

@@ -13,8 +13,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pandas as pd\n",
-    "import pathlib\n",
+    "import causalpy as cp\n",
     "import arviz as az"
    ]
   },
@@ -27,28 +26,13 @@
     "az.style.use(\"arviz-darkgrid\")"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Load data"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
-    "did_data_path = pathlib.Path.cwd().parents[1] / \"causalpy\" / \"data\" / \"did.csv\"\n",
-    "data = pd.read_csv(did_data_path)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Run the analysis"
+    "data = cp.load_data(\"did\")"
    ]
   },
   {
@@ -57,26 +41,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from causalpy.skl_experiments import DifferenceInDifferences\n",
     "from sklearn.linear_model import LinearRegression\n",
     "\n",
-    "# NOTE: `treated` is a deterministic function of `t` and `group`. So add this function into the formula.\n",
     "\n",
-    "result = DifferenceInDifferences(\n",
+    "result = cp.skl_experiments.DifferenceInDifferences(\n",
     "    data,\n",
     "    formula=\"y ~ 1 + group + t + treated:group\",\n",
     "    time_variable_name=\"t\",\n",
     "    prediction_model=LinearRegression(),\n",
     ")"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Examine the results"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 5,

docs/notebooks/did_pymc_banks.ipynb

@@ -17,9 +17,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pandas as pd\n",
-    "import pathlib\n",
-    "import arviz as az"
+    "import arviz as az\n",
+    "import causalpy as cp"
    ]
   },
   {
@@ -31,45 +30,29 @@
     "az.style.use(\"arviz-darkgrid\")"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Load data"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
-    "rd_data_path = pathlib.Path.cwd().parents[1] / \"causalpy\" / \"data\" / \"drinking.csv\"\n",
     "df = (\n",
-    "    pd.read_csv(rd_data_path)[[\"agecell\", \"all\", \"mva\", \"suicide\"]]\n",
+    "    cp.load_data(\"drinking\")\n",
     "    .rename(columns={\"agecell\": \"age\"})\n",
     "    .assign(treated=lambda df_: df_.age > 21)\n",
     "    .dropna(axis=0)\n",
     ")"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Linear model"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from causalpy.skl_experiments import RegressionDiscontinuity\n",
     "from sklearn.linear_model import LinearRegression\n",
     "\n",
-    "result = RegressionDiscontinuity(\n",
+    "result = cp.skl_experiments.RegressionDiscontinuity(\n",
     "    df,\n",
     "    formula=\"all ~ 1 + age + treated\",\n",
     "    running_variable_name=\"age\",\n",