Как определить субъекты в прогнозируемых результатах МЖС в Statsmodels?

Question

Я делаю линейную регрессию, используя Statsmodels в блокноте Jupyter. Данные находятся в DataFrame, называемом train_base, где столбец id идентифицирует каждый уникальный предмет моей базы данных. Train_base выглядит так:

id     y     x0     x1     x2
a123   20     8      1      3
b789   33     8      3      2
d782   77     9      6      5      

Основной кусок кода показан ниже. Обратите внимание, что я использую другую базу под названием "test_base", чтобы делать прогнозы, и эта база также имеет ту же структуру, что и "train_base", за исключением столбца "y":

results = smf.ols('y ~ x0 + x1 + x2', data=train_base).fit()
predictions = results.predict(test_base)
predictions.head()

Прогнозы такие:

0   -0.054789
1   -0.036042
2   -0.043962
3   -0.135725
4   -0.409129
dtype: float64

Мне кажется, что первый столбец, показанный в прогнозах, является индексом оригинальной train_base (я прав?). Так как мне нужно определить прогнозируемое значение для каждого индивидуума моей тестовой базы, что мне нужно сделать, чтобы в прогнозах был столбец "id"?

Answer 1

С некоторыми тонкими исправлениями в ответе @Leda Grasiele (включение внутренних квадратных скобок и использование функции транспонирования) я пришел к правильному коду, показанному ниже:

predict = pd.DataFrame([results.predict(test_base), train_base['id']]).T
predict.columns = ['predictions', 'id']
predict

Answer 2

Первое: да, столбец id является индексом. Второе: у меня нет доступа к вашим данным для проверки моего предложения, но я думаю, что может сработать следующее (или что-то подобное):

predict = pd.DataFrame(results.predict(test_base), train_base['id'])
predict.columns = ['predict']
predict

Я думаю, что это может сработать, если каждое из значений, которые появляются в ваших прогнозах, связано с каждым из индекса id, поскольку целью этого кода является создание кадра данных, состоящего из результатов прогноза и индекса id.