how to use the train_x and train_y from sklearn k-fold split generator
I am using the sklearn k-fold generator to split some data 10 times. When I run the code below I expect train_x,train_y,test_x,test_y to contain all 10 splits however only the last split seems to be saved in the variables. I want to either save each split as its own data or have all 10 splits in the variables.
the code is:
def Branin(X,Y): PI = 3.14159265359 a = 1 b = 5.1 / (4 * pow(PI, 2)) c = 5 / PI r = 6 s = 10 t = 1 / (8 * PI) train_y = a * (X - b * Y ** 2 + c * Y - r) ** 2 + s * (1 - t) * torch.cos(Y) + s return train_y x,y=torch.linspace(-5, 10, 40), torch.linspace(0, 15, 40) X1 = np.column_stack((x, y)) X = torch.from_numpy(X1) Y = Branin(X[:, 0],X[:, 1]) print(X.shape) print(Y.shape) kf = KFold(n_splits=10) kf.get_n_splits(X) KFold(n_splits=10, random_state=None, shuffle=False) for train_index, test_index in kf.split(X,Y): train_x, test_x = X[train_index], X[test_index] train_y, test_y = Y[train_index], Y[test_index] print(test_x) # print("TRAIN:", train_index, "TEST:", test_index) #example of output for test_x here is the last split in the above test_x# print("test_x output",test_x)
How can I get test_x for example to be all the splits or say split five once calling it outside the generator?