One-from-group or N-from-group about abess HOT 7 CLOSED

from abess.linear import LinearRegression
from abess.datasets import make_glm_data
from sklearn.feature_selection import r_regression
from sklearn.feature_selection import SelectFromModel
import numpy as np

def abess_group(x, y, group, s):
    n, p = x.shape
    group_label = np.unique(group)
    group_num = len(group_label)

    # Marginal screening for each group to find the top s variables in each group:
    x_cor = np.abs(r_regression(x, y))
    A = []
    group_start_index = []
    for i in group_label:
        group_i_index = group == i
        group_start_index.append(np.where(group_i_index)[0].min())
        each_group_num = np.sum(group_i_index)
        A.extend(np.where(
            np.logical_and(group_i_index, x_cor >= np.sort(x_cor[group_i_index])[each_group_num - s])
        )[0].tolist())
    A = np.array(A)
    group_start_index = np.array(group_start_index)
    
    # The iteration for block-wise best-subset selection:
    model = LinearRegression()
    A1 = np.array([0] * s * group_num)
    select_var_group = group_label.repeat(s)
    while np.any(A != A1):
        A1 = A.copy()
        for i in group_label:
            x_toselect = x[:, group == i]               # the data of the i-th group
            x_fixed = x[:, A[select_var_group != i]]    # the data of the selected variables in remaining groups
            candidate_num = x_toselect.shape[1]
            x_fit = np.hstack([x_toselect, x_fixed])
            model.set_params(
                # to include the selected variables in remaining groups:
                always_select=np.arange(candidate_num, x_fit.shape[1]), 
                # to select s variables in the i-th group (leveraging warm-start):
                support_size=range(candidate_num, s+candidate_num),    
            )
            model.fit(x_fit, y)
            # update the selected s variables:
            select_result_tmp = SelectFromModel(estimator=model, prefit=True).get_support()[range(candidate_num)]
            A[select_var_group == i] = np.array(np.where(select_result_tmp)[0] + group_start_index[i])
        pass
    return A

np.random.seed(1)
p = 50
J = 10
group = np.linspace(0, 4, 5).repeat(10).astype(np.int32)
coef_ = np.array([1.0 if i % 10 == 0 or i %10 == 1 else 0.0 for i in range(50)])
dataset = make_glm_data(n=100, p=50, coef_=coef_, family='gaussian', k = 10)
print('True:', np.where(coef_ != 0.0)[0])
est = abess_group(dataset.x, dataset.y, group, s=2)
print("Estimate:", est)

True: [ 0  1 10 11 20 21 30 31 40 41]
Estimate: [ 0  1 10 11 20 21 30 31 40 41]