Heckman 两阶段

Heckman 样本选择模型两阶段估计。

核心代码

    def _fit_heckman_two_step(self, decimals: int, title: str) -> Any:
        if not self.y_var:
            raise ValueError(_('请先选择因变量'))
        if not self.selection_var:
            raise ValueError(_('Heckman 两阶段需要设置样本是否被观察到的选择变量'))

        outcome_vars = self._unique_preserve_order(self.x_vars)
        exclusion_vars = self._unique_preserve_order(self.selection_exclusion_vars)
        cluster_var = self.se_options.get('cluster_var')

        if not outcome_vars:
            raise ValueError(_('Heckman 两阶段至少需要一个结果方程解释变量'))

        selection_vars = self._unique_preserve_order(outcome_vars + exclusion_vars)
        if not exclusion_vars:
            warnings.warn(
                _('当前 Heckman 选择方程没有额外排除限制变量，模型主要依赖 Probit 非线性识别。实证论文中通常建议加入至少一个只进入选择方程的变量。'),
                UserWarning
            )

        stage1_cols = [self.selection_var] + selection_vars
        if cluster_var:
            stage1_cols.append(cluster_var)
        stage1_cols = self._unique_preserve_order([col for col in stage1_cols if col])

        stage1_df = self.data[stage1_cols].dropna().copy()
        if stage1_df.empty:
            raise ValueError(_('选择方程可用样本为空，请检查选择变量和协变量缺失情况'))

        selection_values = set(pd.Series(stage1_df[self.selection_var]).dropna().unique().tolist())
        if not selection_values.issubset({0, 1, False, True}):
            raise ValueError(_('Heckman 的选择变量必须是 0/1 二元变量，其中 1 表示结果变量可被观察'))
        stage1_df[self.selection_var] = stage1_df[self.selection_var].astype(int)

        stage1_y = stage1_df[self.selection_var]
        stage1_x = sm.add_constant(stage1_df[selection_vars], has_constant='add')

        try:
            stage1_model = Probit(stage1_y, stage1_x)
            self.selection_result = stage1_model.fit(disp=0)
        except Exception as exc:
            raise ValueError(_('Heckman 第一阶段 Probit 估计失败: %(error)s') % {'error': str(exc)})

        stage1_index = pd.Index(stage1_df.index)
        stage1_linear = pd.Series(
            np.asarray(stage1_x @ self.selection_result.params),
            index=stage1_index,
            name='selection_index'
        )
        cdf_values = np.clip(norm.cdf(stage1_linear), 1e-8, 1 - 1e-8)
        imr_selected = pd.Series(norm.pdf(stage1_linear) / cdf_values, index=stage1_index, name='IMR')

        stage2_cols = [self.y_var, self.selection_var] + outcome_vars
        if cluster_var:
            stage2_cols.append(cluster_var)
        stage2_cols = self._unique_preserve_order([col for col in stage2_cols if col])

        stage2_df = self.data[stage2_cols].copy()
        stage2_df = stage2_df.join(imr_selected, how='inner')
        stage2_df[self.selection_var] = stage2_df[self.selection_var].astype(float)
        stage2_df = stage2_df[stage2_df[self.selection_var] == 1]
        stage2_df = stage2_df.dropna(subset=[self.y_var] + outcome_vars + ['IMR'])

        if stage2_df.empty:
            raise ValueError(_('Heckman 第二阶段没有可用样本，请确认选择变量=1 的观测中因变量和解释变量不缺失'))
        if len(stage2_df) <= len(outcome_vars) + 2:
            raise ValueError(_('Heckman 第二阶段有效样本量过少，无法稳定估计结果方程'))

        stage2_y = stage2_df[self.y_var]
        stage2_x = sm.add_constant(stage2_df[outcome_vars + ['IMR']], has_constant='add')
        stage2_model = sm.OLS(stage2_y, stage2_x)
        self.result = self._fit_statsmodels_model(stage2_model, stage2_df)

        if self.se_options.get('type', 'iid') == 'iid':
            self._apply_heckman_twostep_covariance(
                stage2_x=stage2_x,
                stage2_y=stage2_y,
                stage1_x=stage1_x,
                stage1_index=stage1_index
            )
        else:
            self._params = self.result.params.copy()
            self._std_errors = self.result.bse.copy()
            self._test_stats = self.result.tvalues.copy()
            self._pvalues = self.result.pvalues.copy()

        imr_p = self._pvalues.get('IMR', np.nan)
        f_value = self._safe_stat(getattr(self.result, 'fvalue', np.nan), max_abs=1e12)
        self.model_stats = {
            'N': int(len(stage1_df)),
            'Selected-N': int(len(stage2_df)),
            'R2': self._safe_stat(getattr(self.result, 'rsquared', np.nan), max_abs=10),
            'Adj-R2': self._safe_stat(getattr(self.result, 'rsquared_adj', np.nan), max_abs=10),
            'F': f_value,
            'IMR-p': imr_p,
            'Selection-Pseudo-R2': getattr(self.selection_result, 'prsquared', np.nan),
        }

        self.display_x_vars = outcome_vars + ['Inverse Mills Ratio']
        self.table_custom_rows = [
            {'label': 'Selection correction', 'value': 'Yes'}
        ]

        self.raw_output = (
            "Heckman Two-Step - Outcome Equation\n"
            + self.result.summary().as_text()
            + "\n\n"
            + "Heckman Two-Step - Selection Equation (Probit)\n"
            + self.selection_result.summary().as_text()
        )
        self.diagnostic_html = self._build_heckman_diagnostic_html(
            stage1_sample=len(stage1_df),
            stage2_sample=len(stage2_df),
            selection_vars=selection_vars,
            exclusion_vars=exclusion_vars,
            decimals=decimals,
            title=title
        )

        return self.result