从函数输出向Pandas Dataframe添加新列

时间:2018-06-01 12:02:41

标签: python pandas numpy linear-regression

我写了函数来估计简单线性回归的参数。该功能产生多个输出。功能输入是两个列表。另外,我从我导出两个列表的地方获得了初始DataFrame df。

我想将函数中的一些输出作为新列添加到初始DataFrame中,或者在函数外部添加新列表。

例如:

def predict(X,Y):
     beta1 = sum([(X[i] - mean_X)*(Y[i] - mean_Y) for i in range(len(X))]) / sum([(X[i] - mean_X)**2 for i in range(len(X))])
     beta0 = mean_Y - beta1 * mean_X

     y_hat = [beta0 + beta1*X[i] for i in range(len(X))]

     return df.assign(prediction = y_hat)

这里,mean_X和mean_Y分别是列表X和列表Y的样本平均值。

我还尝试 numpy.insert()将y_hat添加到非初始DataFrame中,但添加到X中,我将其转换为numpy数组。

我没有成功达到预期的结果,所以有人可以帮助我吗?

2 个答案:

答案 0 :(得分:2)

据我了解您的问题,您希望在现有/新专栏中使用您的功能。如果是这种情况,这是一种方法。如果没有,那么让我知道,我会删除答案。感谢

import pandas as pd

def Somefunction(x, y):
  a = 2 *x
  b = 3 * y
  return pd.Series([a, b], index= ['YourColumn1', 'YourColumn2'])





df = pd.read_csv('YourFile')

df = df.join(df.apply(lambda x: 
  Somefunction(x['ColumnYouWantToApplyFunctionReturnValue a'], 
  x['ColumnYouWantToApplyFunctionReturnValue B']), axis=1))

答案 1 :(得分:1)

您的代码似乎不太清楚。 mean_Xmean_Y变量是什么?

编辑:添加了变量声明。

无论如何,这是一个简单的建议:

import numpy as np
def predict(X, Y, df):
    mean_X = np.mean(X)
    mean_Y = np.mean(Y)
    beta1 = sum([(X[i] - mean_X)*(Y[i] - mean_Y) for i in range(len(X))]) / sum([(X[i] - mean_X)**2 for i in range(len(X))])
    beta0 = mean_Y - beta1 * mean_X
    y_hat = [beta0 + beta1*X[i] for i in range(len(X))]
    df['prediction'] = y_hat
    return df

更聪明的方法是使用在DataFrame上调用的apply()函数。