def record(self, verbose: int):
self.trainer.treatment_net.train(False)
if self.trainer.covariate_net is not None:
self.trainer.covariate_net.train(False)
n_train_data = self.train_data_t.treatment.size()[0]
n_val_data = self.validation_data_t.treatment.size()[0]
n_test_data = self.test_data_t.treatment.size()[0]
with torch.no_grad():
treatment_train_feature = self.trainer.treatment_net(self.train_data_t.treatment)
treatment_val_feature = self.trainer.treatment_net(self.validation_data_t.treatment)
treatment_test_feature = self.trainer.treatment_net(self.test_data_t.treatment)
covariate_train_feature = None
covariate_val_feature = None
covariate_test_feature = None
if self.trainer.covariate_net is not None:
covariate_train_feature = self.trainer.covariate_net(self.train_data_t.covariate)
covariate_val_feature = self.trainer.covariate_net(self.validation_data_t.covariate)
covariate_test_feature = self.trainer.covariate_net(self.test_data_t.covariate)
# stage2
feature = DFIVModel.augment_stage2_feature(treatment_train_feature,
covariate_train_feature,
self.trainer.add_intercept)
weight = fit_linear(self.train_data_t.outcome, feature, self.trainer.lam)
insample_pred = linear_reg_pred(feature, weight)
insample_loss = torch.norm(self.train_data_t.outcome - insample_pred) ** 2 / n_train_data
val_feature = DFIVModel.augment_stage2_feature(treatment_val_feature,
covariate_val_feature,
self.trainer.add_intercept)
outsample_pred = linear_reg_pred(val_feature, weight)
outsample_loss = torch.norm(self.validation_data_t.outcome - outsample_pred) ** 2 / n_val_data
# eval for test
test_feature = DFIVModel.augment_stage2_feature(treatment_test_feature,
covariate_test_feature,
self.trainer.add_intercept)
test_pred = linear_reg_pred(test_feature, weight)
test_loss = torch.norm(self.test_data_t.structural - test_pred) ** 2 / n_test_data
if verbose >= 1:
logger.info(f"insample_loss:{insample_loss.item()}")
logger.info(f"outsample_loss:{outsample_loss.item()}")
logger.info(f"test_loss:{test_loss.item()}")
def predict_t(self, treatment: torch.Tensor,
covariate: Optional[torch.Tensor]):
treatment_feature = self.treatment_net(treatment)
covariate_feature = None
if self.covariate_net:
covariate_feature = self.covariate_net(covariate)
feature = DFLModel.augment_feature(treatment_feature,
covariate_feature,
self.add_intercept)
return linear_reg_pred(feature, self.weight_mat)
def fit_dfl(treatment_feature: torch.Tensor,
covariate_feature: Optional[torch.Tensor],
outcome_t: torch.Tensor, lam: float, add_intercept: bool):
# stage1
feature = DFLModel.augment_feature(treatment_feature,
covariate_feature, add_intercept)
weight = fit_linear(outcome_t, feature, lam)
pred = linear_reg_pred(feature, weight)
loss = torch.norm((outcome_t - pred))**2 + lam * torch.norm(weight)**2
return dict(weight=weight, loss=loss)
def fit_2sls(
treatment_1st_feature: torch.Tensor,
instrumental_1st_feature: torch.Tensor,
instrumental_2nd_feature: torch.Tensor,
covariate_2nd_feature: Optional[torch.Tensor],
outcome_2nd_t: torch.Tensor,
lam1: float,
lam2: float,
add_stage1_intercept: bool,
add_stage2_intercept: bool,
):
# stage1
feature = DFIVModel.augment_stage1_feature(instrumental_1st_feature,
add_stage1_intercept)
stage1_weight = fit_linear(treatment_1st_feature, feature, lam1)
# predicting for stage 2
feature = DFIVModel.augment_stage1_feature(instrumental_2nd_feature,
add_stage1_intercept)
predicted_treatment_feature = linear_reg_pred(feature, stage1_weight)
# stage2
feature = DFIVModel.augment_stage2_feature(predicted_treatment_feature,
covariate_2nd_feature,
add_stage2_intercept)
stage2_weight = fit_linear(outcome_2nd_t, feature, lam2)
pred = linear_reg_pred(feature, stage2_weight)
stage2_loss = torch.norm(
(outcome_2nd_t - pred))**2 + lam2 * torch.norm(stage2_weight)**2
return dict(stage1_weight=stage1_weight,
predicted_treatment_feature=predicted_treatment_feature,
stage2_weight=stage2_weight,
stage2_loss=stage2_loss)
def update_covariate_net(self, train_1st_data: TrainDataSetTorch,
train_2nd_data: TrainDataSetTorch, verbose: int):
# have instrumental features
self.instrumental_net.train(False)
instrumental_1st_feature = self.instrumental_net(
train_1st_data.instrumental).detach()
instrumental_2nd_feature = self.instrumental_net(
train_2nd_data.instrumental).detach()
self.treatment_net.train(False)
treatment_1st_feature = self.treatment_net(
train_1st_data.treatment).detach()
feature = DFIVModel.augment_stage1_feature(instrumental_1st_feature,
self.add_stage1_intercept)
stage1_weight = fit_linear(treatment_1st_feature, feature, self.lam1)
# predicting for stage 2
feature = DFIVModel.augment_stage1_feature(instrumental_2nd_feature,
self.add_stage1_intercept)
predicted_treatment_feature = linear_reg_pred(feature,
stage1_weight).detach()
self.covariate_net.train(True)
for i in range(self.covariate_iter):
self.covariate_opt.zero_grad()
covariate_feature = self.covariate_net(train_2nd_data.covariate)
feature = DFIVModel.augment_stage2_feature(
predicted_treatment_feature, covariate_feature,
self.add_stage2_intercept)
loss = linear_reg_loss(train_2nd_data.outcome, feature, self.lam2)
loss.backward()
if verbose >= 2:
logger.info(f"update covariate: {loss.item()}")
self.covariate_opt.step()
def record(self, verbose: int):
self.trainer.treatment_net.train(False)
self.trainer.instrumental_net.train(False)
if self.trainer.covariate_net is not None:
self.trainer.covariate_net.train(False)
n_1st_data = self.train_1st_data_t.treatment.size()[0]
n_2nd_data = self.train_2nd_data_t.treatment.size()[0]
n_val_data = self.validation_data_t.treatment.size()[0]
n_test_data = self.test_data_t.treatment.size()[0]
with torch.no_grad():
treatment_1st_feature = self.trainer.treatment_net(
self.train_1st_data_t.treatment)
treatment_val_feature = self.trainer.treatment_net(
self.validation_data_t.treatment)
treatment_test_feature = self.trainer.treatment_net(
self.test_data_t.treatment)
instrumental_1st_feature = self.trainer.instrumental_net(
self.train_1st_data_t.instrumental)
instrumental_2nd_feature = self.trainer.instrumental_net(
self.train_2nd_data_t.instrumental)
instrumental_val_feature = self.trainer.instrumental_net(
self.validation_data_t.instrumental)
covariate_2nd_feature = None
covariate_val_feature = None
covariate_test_feature = None
if self.trainer.covariate_net is not None:
covariate_2nd_feature = self.trainer.covariate_net(
self.train_2nd_data_t.covariate)
covariate_val_feature = self.trainer.covariate_net(
self.validation_data_t.covariate)
covariate_test_feature = self.trainer.covariate_net(
self.test_data_t.covariate)
# stage1
feature = DFIVModel.augment_stage1_feature(
instrumental_1st_feature, self.trainer.add_stage1_intercept)
stage1_weight = fit_linear(treatment_1st_feature, feature,
self.trainer.lam1)
insample_1st_pred = linear_reg_pred(feature, stage1_weight)
stage1_insample = torch.norm(treatment_1st_feature -
insample_1st_pred)**2 / n_1st_data
val_feature = DFIVModel.augment_stage1_feature(
instrumental_val_feature, self.trainer.add_stage1_intercept)
outsample_1st_pred = linear_reg_pred(val_feature, stage1_weight)
stage1_outsample = torch.norm(treatment_val_feature -
outsample_1st_pred)**2 / n_val_data
# predicting for stage 2
feature = DFIVModel.augment_stage1_feature(
instrumental_2nd_feature, self.trainer.add_stage1_intercept)
predicted_treatment_feature = linear_reg_pred(
feature, stage1_weight)
# stage2
feature = DFIVModel.augment_stage2_feature(
predicted_treatment_feature, covariate_2nd_feature,
self.trainer.add_stage2_intercept)
stage2_weight = fit_linear(self.train_2nd_data_t.outcome, feature,
self.trainer.lam2)
insample_2nd_pred = linear_reg_pred(feature, stage2_weight)
stage2_insample = torch.norm(self.train_2nd_data_t.outcome -
insample_2nd_pred)**2 / n_2nd_data
val_feature = DFIVModel.augment_stage2_feature(
outsample_1st_pred, covariate_val_feature,
self.trainer.add_stage2_intercept)
outsample_2nd_pred = linear_reg_pred(val_feature, stage2_weight)
stage2_outsample = torch.norm(self.validation_data_t.outcome -
outsample_2nd_pred)**2 / n_val_data
# eval for test
test_feature = DFIVModel.augment_stage2_feature(
treatment_test_feature, covariate_test_feature,
self.trainer.add_stage2_intercept)
test_pred = linear_reg_pred(test_feature, stage2_weight)
test_loss = torch.norm(self.test_data_t.structural -
test_pred)**2 / n_test_data
if verbose >= 1:
logger.info(f"stage1_insample:{stage1_insample.item()}")
logger.info(f"stage1_outsample:{stage1_outsample.item()}")
logger.info(f"stage2_insample:{stage2_insample.item()}")
logger.info(f"stage2_outsample:{stage2_outsample.item()}")
logger.info(f"test_loss:{test_loss.item()}")