-
Notifications
You must be signed in to change notification settings - Fork 16
/
binsreg_illustration_plot.py
51 lines (36 loc) · 1.62 KB
/
binsreg_illustration_plot.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
################################################################################
# Binsreg: illustration file for plot
# Authors: M. D. Cattaneo, R. Crump, M. Farrell, Y. Feng and Ricardo Masini
# Last update: March 17, 2023
################################################################################
import pandas as pd
import statsmodels.formula.api as smf
from plotnine import *
from binsreg import *
######################################
###### Read the same data #############
####### used for STATA ################
#######################################
data = pd.read_csv("binsreg_sim.csv")
data.describe().T
########################################################
###### EXTERNAL PLOT USING BINSREG OUTPUT ##############
########################################################
# Run binsreg
est = binsreg('y', 'x', 'w', data=data,line = (3,3), ci=(3,3), cb=(3,3), polyreg=4)
# Extract the plotting information
result = est.data_plot[0]
# Create the figure to plot
fig = ggplot() + labs(x='X',y ='Y')
# Add the dots
fig += geom_point(data=result.dots, mapping=aes(x='x', y='fit'), color="blue", size=2, shape='o')
# Add the line
fig += geom_line(data=result.line, mapping=aes(x='x', y='fit'), color="blue", size=0.5)
# Add the CI
fig += geom_errorbar(data=result.ci, mapping=aes(x='x', ymin='ci_l', ymax='ci_r'), color="blue", size=0.5, width = 0.02, linetype='solid')
# Add the CB
fig += geom_ribbon(data=result.cb, mapping=aes(x='x', ymin='cb_l', ymax='cb_r'), fill="blue", alpha=0.2)
# Add the polyreg
fig += geom_line(data=result.poly, mapping=aes(x='x', y='fit'), color="red", size=0.5)
# Display the plot
print(fig)