def test_threshold_create(self, sm_network):
expected = pd.DataFrame([[1, 1, 2, 1, 3],
[0, 0, 2, 2, 3],
[0, 1, 1, 1, 3],
[0, 0, 2, 0, 2],
[1, 0, 2, 1, 2],
[1, 0, 0, 0, 1],
[0, 1, 0, 1, 2],
[0, 0, 0, 0, 0],
[1, 1, 1, 2, 2]],
columns=["W", "A", "A_sum", "W_sum", "degree"],
index=[0, 1, 2, 3, 4, 5, 6, 7, 8])
expected["A_t2"] = np.where(expected['A_sum'] > 2, 1, 0)
expected["W_tp50"] = np.where(expected['W_sum']/expected['degree'] > 0.5, 1, 0)
tmle = NetworkTMLE(network=sm_network, exposure='A', outcome='Y')
tmle.define_threshold(variable='A', threshold=2, definition='sum')
tmle.define_threshold(variable='W', threshold=0.5, definition='mean')
created = tmle.df_restricted
# Checking that expected is the same as the created
columns = ["A", "A_t2", "W_tp50"]
pdt.assert_frame_equal(expected[columns], created[columns], check_dtype=False)
results.loc[i, 'var_' + str(p)] = stmle.conditional_se**2
results.loc[i, 'lcl_' + str(p)] = stmle.conditional_ci[0]
results.loc[i, 'ucl_' + str(p)] = stmle.conditional_ci[1]
except:
results.loc[i, 'bias_' + str(p)] = np.nan
results.loc[i, 'var_' + str(p)] = np.nan
results.loc[i, 'lcl_' + str(p)] = np.nan
results.loc[i, 'ucl_' + str(p)] = np.nan
else:
# Network TMLE
ntmle = NetworkTMLE(H,
exposure=exposure,
outcome=outcome,
degree_restrict=degree_restrict)
ntmle.define_threshold(variable='diet', threshold=3, definition='sum')
ntmle.exposure_model(gin_model)
ntmle.exposure_map_model(gsn_model,
measure=measure_gs,
distribution=distribution_gs)
ntmle.outcome_model(qn_model, distribution='normal')
for p in prop_treated: # loops through all treatment plans
try:
if shift:
z = odds_to_probability(np.exp(log_odds + p))
ntmle.fit(p=z, bound=0.01)
else:
ntmle.fit(p=p, bound=0.01)
results.loc[i, 'bias_' +
str(p)] = ntmle.marginal_outcome - truth[p]
results.loc[i, 'var_' + str(p)] = ntmle.conditional_variance