Integrated Brier Score for sklearn's GridSearchCV

Question

I have a custom scorer function whose inputs depend on the specific train and validation fold, additionally, the estimator's .predict_survival_function output is also needed. To give a more concrete example:

I am trying to run a GridSearch for a Random Survival Forest (scikit-survival package) with the Integrated Brier Score (IBS) as the scoring method. The challenge is in the fact that the domain of the IBS is data- (and therefore, fold-) specific as it relies on the Kaplan-Meyer estimate at some point. Moreover, the .predict_survival_function method needs to be called every time during the scoring evaluation step and not only at the end of it.

It seems that I managed to to deal with the first issue by creating the following function:

def IB_time_interval(y_train, y_test):
    y_times_tr = [i[2] for i in y_train]
    y_times_te = [i[2] for i in y_test]

    T1 = np.percentile(y_times_tr, 5, interpolation='higher')
    T2 = np.percentile(y_times_tr, 95, interpolation='lower')
    T3 = np.percentile(y_times_te, 5, interpolation='higher')
    T4 = np.percentile(y_times_te, 95, interpolation='lower')
    
    return np.linspace(np.maximum(T1,T3), np.minimum(T2, T4))

That is robust enough to work throughout all the folds. However, I am unable to retrieve the estimator's predictions during the grid search phase, as a non-fitted copy of it seems to passed instead every time the custom scorer function is called. The workaround that I I tried is to re-fit the estimator inside the scoring function, but not only this is conceptually wrong, it also raises errors.

The custom scorer function looks like the following:

def IB_scorer(y_true, y_pred, times=times_linspace, y=y, clf=rsf):
    
    rsf.fit(X_train,y_train) #<--- = conceptually wrong
    survs_test = rsf.predict_survival_function(X_test, return_array=False) #<---
    T1, T2 = survs_test[0].x.min(), survs_test[0].x.max()
    mask2 = np.logical_or(times >= T2, times < T1) # mask outer interval            
    times = times[~mask2]
    #preds has shape (n_y-s, n_times)
    preds_test = np.asarray([[fn(t) for t in times] for fn in survs_test])
    return integrated_brier_score(y, y_true, preds_test, times)

and I create the scoring object immediately afterwards:

trial_IB_scorer = make_scorer(IB_scorer, greater_is_better=False)

Any suggestions? It would be great to be able to use GridSearch with more complex scoring functions, especially for the survival analysis case!

PS. I will paste the rest of the minimal working code here:

import numpy as np
from sksurv.ensemble import RandomSurvivalForest
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV
from sksurv.metrics import integrated_brier_score


from sksurv.datasets import load_breast_cancer
X, y = load_breast_cancer()
X = X.drop(["er", "grade"], axis=1)
y_cens = np.array([i[0] for i in y]) #censoring status 1 or 0

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size= 0.2,
                                                    shuffle=True,
                                                    random_state=0,
                                                    stratify = y_cens)

param_grid = {
            'max_depth': [4, 20, None],
            'max_features': ["sqrt", None]
            }

rsf = RandomSurvivalForest(n_jobs=1, random_state=0)

times_linspace = IB_time_interval(y_train, y_test)


clf = GridSearchCV(rsf, param_grid, refit=True, n_jobs=1, 
                    scoring=trial_IB_scorer)

clf.fit(X_train, y_train)
print("final score clf", clf.score(X_train, y_train))
print(clf.best_params_)