Skip to content

Commit 5d16fae

Browse files
committed
#55 add a quickstart
1 parent 8f39903 commit 5d16fae

File tree

2 files changed

+69
-0
lines changed

2 files changed

+69
-0
lines changed

README.md

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,40 @@ Alternatively, if you want the very latest version of the package you can instal
3535
pip install git+https://github.com/pymc-labs/CausalPy.git
3636
```
3737

38+
## Quickstart
39+
40+
```python
41+
from causalpy.pymc_experiments import RegressionDiscontinuity
42+
from causalpy.pymc_models import LinearRegression
43+
import pandas as pd
44+
import pathlib
45+
46+
47+
# Import and process data
48+
rd_data_path = pathlib.Path.cwd().parents[1] / "causalpy" / "data" / "drinking.csv"
49+
df = (
50+
pd.read_csv(rd_data_path)[["agecell", "all", "mva", "suicide"]]
51+
.rename(columns={"agecell": "age"})
52+
.assign(treated=lambda df_: df_.age > 21)
53+
.dropna(axis=0)
54+
)
55+
56+
# Run the analysis
57+
result = RegressionDiscontinuity(
58+
df,
59+
formula="all ~ 1 + age + treated",
60+
running_variable_name="age",
61+
prediction_model=LinearRegression(),
62+
treatment_threshold=21,
63+
)
64+
65+
# Visualize outputs
66+
fig, ax = result.plot();
67+
68+
# Get a results summary
69+
result.summary()
70+
```
71+
3872
## Roadmap
3973

4074
Plans for the repository can be seen in the [Issues](https://github.com/pymc-labs/CausalPy/issues).

docs/index.rst

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,41 @@ Alternatively, if you want the very latest version of the package you can instal
2424
pip install git+https://github.com/pymc-labs/CausalPy.git
2525
2626
27+
Quickstart
28+
----------
29+
30+
.. code-block:: python
31+
32+
from causalpy.pymc_experiments import RegressionDiscontinuity
33+
from causalpy.pymc_models import LinearRegression
34+
import pandas as pd
35+
import pathlib
36+
37+
38+
# Import and process data
39+
rd_data_path = pathlib.Path.cwd().parents[1] / "causalpy" / "data" / "drinking.csv"
40+
df = (
41+
pd.read_csv(rd_data_path)[["agecell", "all", "mva", "suicide"]]
42+
.rename(columns={"agecell": "age"})
43+
.assign(treated=lambda df_: df_.age > 21)
44+
.dropna(axis=0)
45+
)
46+
47+
# Run the analysis
48+
result = RegressionDiscontinuity(
49+
df,
50+
formula="all ~ 1 + age + treated",
51+
running_variable_name="age",
52+
prediction_model=LinearRegression(),
53+
treatment_threshold=21,
54+
)
55+
56+
# Visualize outputs
57+
fig, ax = result.plot();
58+
59+
# Get a results summary
60+
result.summary()
61+
2762
Features
2863
--------
2964

0 commit comments

Comments
 (0)