def get_feature_dict(all_cols, cats=None):
    """
    This helper function makes it easy to loop though columns in a dataframe
    and group onehot encoded columns together.

    :param all_cols: all columns of a dataframe
    :type all_cols: list
    :param cats: categorical columns that have been onehotencoded, defaults to None
    :type cats: list, optional
    :return: returns a dict with as key all original (not one hot encoded) columns
        and as items a list of columns associated with that column.

        e.g. {'Age': ['Age'],
              'Gender' : ['Gender_Male', 'Gender_Female']}
    :rtype: dict
    """
    feature_dict = {}

    if cats is None:
        return {col: [col] for col in all_cols}

    for col in cats:
        cat_cols = [c for c in all_cols if c.startswith(col)]
        if len(cat_cols) > 1:
            feature_dict[col] = cat_cols

    # add all the individual features
    other_cols = list(set(all_cols) - set([item for sublist in list(feature_dict.values())
                                                for item in sublist]))

    for col in other_cols:
        feature_dict[col] = [col]
    return feature_dict