Могу ли я использовать API формулы Statsmodels в Spark Pandas UDF - PullRequest
Новая
       68

Могу ли я использовать API формулы Statsmodels в Spark Pandas UDF

0 голосов
/ 04 августа 2020

В этом блоге Databricks есть пример Pandas UDF со Statsmodels ссылка 

import statsmodels.api as sm
# df has four columns: id, y, x1, x2

group_column = 'id'
y_column = 'y'
x_columns = ['x1', 'x2']
schema = df.select(group_column, *x_columns).schema

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
# Input/output are both a pandas.DataFrame
def ols(pdf):
    group_key = pdf[group_column].iloc[0]
    y = pdf[y_column]
    X = pdf[x_columns]
      X = sm.add_constant(X)
    model = sm.OLS(y, X).fit()

    return pd.DataFrame([[group_key] + [model.params[i] for i in   x_columns]], columns=[group_column] + x_columns)

beta = df.groupby(group_column).apply(ols)

Как я могу использовать API формулы Statsmodels для воссоздания того же кода? Более конкретно, я хотел бы определить Pandas UDF, чтобы его входы были 

formula: a string that specifies the R-style regression formula
df:      a Pandas DataFrame

Например, 


formula='y ~ x1 + x2 + x1:x2 - 1 '

df=pdf

def ols_formula(formula,df):
    import statsmodels.formula.api as smf

    model = smf.ols(formula, df).fit()

    return model
...