def fit()

in fclib/fclib/models/multiple_linear_regression.py [0:0]

• 10 lines of code
• 3 McCabe index (conditional complexity)

def fit(train_df, grain_col_names, fea_col_names=[], target_col_name="target"):
    """Train multiple linear regression models with each being trained on an individual time 
    series specified by columns in grain_col_names.
    
    Args: 
        train_df (pandas.DataFrame): Training data frame including all the features
        grain_col_names (list[str]): List of the column names that specify each time series
        fea_col_names (list[str]): List of the names of columns that we want to use as input features
        target_col_name (str): Name of the target column
        
    Returns:
        dict: Dictionary including all the trained linear regression models
    """
    lr_models = {}
    if not fea_col_names:
        fea_col_names = list(train_df.columns)
        fea_col_names.remove(target_col_name)
    for name, group in train_df.groupby(grain_col_names):
        lr = LinearRegression()
        lr.fit(group[fea_col_names], group[target_col_name])
        lr_models[name] = lr
    return lr_models