vivien/causal-simulator

1import numpy as np2import pandas as pd3from sklearn import linear_model4import streamlit as st5from st_cytoscape import cytoscape6 7num_observations = 10_0008generating_model = """9Z <-- N10C1 <-- N11C21 <-- N12C31 <-- N13C33 <-- N14C22 <-- C21 + N15C32 <-- C31 + C33 + N16X <-- Z + C1 + C21 + C31 + N17M <-- X + N18Y <-- M + C1 + C22 + C33 + N19C4 <-- X + Y + N20"""21 22 23def rewrite(x):24    if x == "N":25        return "np.random.randn(num_observations)"26    else:27        return f'd["{x}"]'28 29 30@st.cache_data31def generate_data():32    np.random.seed(seed=0)33    d = {}34    nodes = set()35    edges = set()36    for line in generating_model.split("\n"):37        if " <-- " in line:38            left, right = line.split(" <-- ")39            right_terms = right.split(" + ")40            nodes.add(left)41            for node in right_terms:42                if node != "N":43                    nodes.add(node)44                    edges.add((node, left))45            formula = f"{rewrite(left)} = {' + '.join(list(map(rewrite, right_terms)))}"46            exec(formula)47    return pd.DataFrame.from_dict(d), nodes, edges48 49 50df, nodes, edges = generate_data()51elements = []52for node in nodes:53    elements.append(54        {55            "data": {"id": node},56            "selected": node == "X",57            "selectable": node not in ["X", "Y"],58        }59    )60for edge in edges:61    elements.append(62        {63            "data": {64                "source": edge[0],65                "target": edge[1],66                "id": f"{edge[0]}-{edge[1]}",67            },68            "selectable": False,69        }70    )71stylesheet = [72    {"selector": "node", "style": {"label": "data(id)", "width": 20, "height": 20}},73    {74        "selector": "edge",75        "style": {76            "width": 2,77            "curve-style": "bezier",78            "target-arrow-shape": "triangle",79        },80    },81]82 83layout = {"name": "fcose", "animationDuration": 0}84layout["alignmentConstraint"] = {"horizontal": [["Z", "X", "M", "Y"]]}85layout["relativePlacementConstraint"] = [{"left": "X", "right": "Y"}]86layout["relativePlacementConstraint"].append({"top": "C1", "bottom": "X"})87layout["relativePlacementConstraint"].append({"top": "C21", "bottom": "X"})88layout["relativePlacementConstraint"].append({"top": "X", "bottom": "C4"})89layout["relativePlacementConstraint"].append({"top": "X", "bottom": "C31"})90layout["nodeRepulsion"] = 5000091 92st.sidebar.title("Causal simulator")93 94st.sidebar.markdown(95    """96**Estimating the effect of a variable X** (e.g. vaccination status) **on another variable Y** (e.g. symptoms) **may require controlling for other variables** (e.g. age if age increases both access to the vaccine and the risks of getting sick).97 98This demo illustrates that **the choice of the variables to control for critically depends on the causal relationships between the variables**.99 100*Inspired by [The Book of Why](http://bayes.cs.ucla.edu/WHY/) by Judea Pearl and Dana Mackenzie and built by [Vivien](https://twitter.com/vivien000000) with [Streamlit](https://streamlit.io/), [Cytoscape.js](https://js.cytoscape.org/) and [scikit-learn](https://scikit-learn.org/stable/)*101"""102)103 104st.subheader("Data generating process")105 106st.markdown(107    """10810,000 observations have been generated for the variables mentioned in the causal graph below. The values for each variable were derived from the values of its parents in the causal graph as follows:109"""110)111st.latex(112    "U = \sum_{V \in \mathrm{\ Parents}(U)} V + \epsilon_U  \quad \mathrm{where} \quad \epsilon_U \overset{\mathrm{i.i.d.}}{\sim} \mathcal{N}(0, 1)"113)114st.subheader("Results of controlling for certain variables")115 116st.markdown(117    "We are using a linear regression to estimate the effect on Y of increasing X by one unit. In the causal graph below, **select the variables to include in the regression and see whether the regression coefficient for X matches the expected value (1, given the data generating process)**."118)119 120 121def add_smiley(x):122    return x + (" 😀" if np.abs(float(x) - 1) < 0.1 else " 😨")123 124 125col1, col2 = st.columns(2)126order = ["X"] + sorted([n for n in nodes if n not in ["X", "Y"]])127results = {v: "" for v in order}128with col1:129    selected = cytoscape(130        elements,131        stylesheet,132        height="450px",133        layout=layout,134        selection_type="additive",135        user_panning_enabled=False,136        user_zooming_enabled=False,137        key="graph",138    )139try:140    selected_nodes = [n for n in order if n in selected["nodes"]]141    regr = linear_model.LinearRegression()142    regr.fit(df[selected_nodes], df[["Y"]])143    for i in range(len(regr.feature_names_in_)):144        results[regr.feature_names_in_[i]] = "%.3f" % regr.coef_[0, i]145    results["X"] = add_smiley(results["X"])146    with col2:147        table = "<table style='margin: auto;'><tbody>"148        table += "<tr><td colspan=2><b>Regression coefficients</b></td></tr>"149        for k in order:150            table += f"<tr><td>{k}</td><td>{results[k]}</td></tr>"151        table += "</table></tbody>"152        st.markdown(table, unsafe_allow_html=True)153except TypeError:154    pass155 156 157@st.cache_data158def compute_instrumental_variable():159    regr = linear_model.LinearRegression()160    regr2 = linear_model.LinearRegression()161    regr.fit(df[["Z"]], df[["Y"]])162    regr2.fit(df[["Z"]], df[["X"]])163    return "%.3f" % (regr.coef_[0, 0] / regr2.coef_[0, 0])164 165 166@st.cache_data167def compute_front_door():168    regr = linear_model.LinearRegression()169    regr2 = linear_model.LinearRegression()170    regr.fit(df[["M", "X"]], df[["Y"]])171    regr2.fit(df[["X"]], df[["M"]])172    index_m = [i for i in range(2) if regr.feature_names_in_[i] == "M"][0]173    return "%.3f" % (regr.coef_[0, index_m] * regr2.coef_[0, 0])174 175 176result_instrumental_variable = add_smiley(compute_instrumental_variable())177result_front_door = add_smiley(compute_front_door())178 179with st.expander("What variables should be controlled for?"):180    st.markdown(181        """The ***back-door criterion*** (*Book of Why*, chapter 4) provides sufficient conditions for variables to be adequate controls:182- The non-causal paths between X and Y going through C1, C21 and C22 need to be blocked by controlling for **C1**, as well as for **C21 or C22** or both;183- C4 and C32 are *colliders* (common consequences of their two respective neighbors). The 2 non-causal paths going through them are then blocked as long as C4 and C32 are not controlled for. Therefore, C4, C31, C32 and C33 should not be controlled for. However, if C32 is controlled for, C31 or C33 or both should be controlled for to block the non-causal path now open;184- M is on a causal path from X to Y and should not be controlled for (so that the causal effect of X on Y is not masked);185- Z is only a cause of X. Controlling for it is useless.186"""187    )188with st.expander("Bonus: front-door criterion and instrumental variable"):189    st.markdown(190        f"""191If using the back-door criterion as described above is impossible, we can deduce from the causal diagram that two alternative methods are available:192- the use of the ***instrumental variable*** Z (result: {result_instrumental_variable})193- the ***front-door criterion*** (*Book of Why*, chapter 5) with M as a mediator (result: {result_front_door})194"""195    )196