Skip to content

Commit

Permalink
fixed formatting issues
Browse files Browse the repository at this point in the history
Signed-off-by: Amit Sharma <amit_sharma@live.com>
  • Loading branch information
amit-sharma committed Nov 4, 2023
1 parent 34c18aa commit 9d7dba7
Showing 1 changed file with 12 additions and 19 deletions.
31 changes: 12 additions & 19 deletions tests/causal_estimators/test_two_stage_regression_estimator.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
import numpy as np
import pandas as pd
import pytest
from pytest import mark

from dowhy.causal_estimators.two_stage_regression_estimator import TwoStageRegressionEstimator
from dowhy import CausalModel
import numpy as np
import pandas as pd

from dowhy.causal_estimators.two_stage_regression_estimator import TwoStageRegressionEstimator

from .base import TestEstimator

Expand Down Expand Up @@ -124,29 +123,23 @@ def test_frontdoor_estimator(self):
target "X"
]
]
""".replace('\n', '')
""".replace(
"\n", ""
)

N_SAMPLES = 10000
# Generate the data
U = np.random.randn(N_SAMPLES)
X = np.random.randn(N_SAMPLES) + 0.3*U
Z = 0.7*X + 0.3*np.random.randn(N_SAMPLES)
Y = 0.65*Z + 0.2*U
X = np.random.randn(N_SAMPLES) + 0.3 * U
Z = 0.7 * X + 0.3 * np.random.randn(N_SAMPLES)
Y = 0.65 * Z + 0.2 * U

# Data to df
df = pd.DataFrame(np.vstack([X, Z, Y]).T, columns=['X', 'Z', 'Y'])
df = pd.DataFrame(np.vstack([X, Z, Y]).T, columns=["X", "Z", "Y"])

# Create a model
model = CausalModel(
data=df,
treatment='X',
outcome='Y',
graph=graph
)
model = CausalModel(data=df, treatment="X", outcome="Y", graph=graph)
estimand = model.identify_effect(proceed_when_unidentifiable=True)
# Estimate the effect with front-door
estimate = model.estimate_effect(
identified_estimand=estimand,
method_name='frontdoor.two_stage_regression'
)
estimate = model.estimate_effect(identified_estimand=estimand, method_name="frontdoor.two_stage_regression")
assert estimate.value == pytest.approx(0.45, 0.025)

0 comments on commit 9d7dba7

Please sign in to comment.