def __init__(self, regression_type, formula, df=None):
super(RegressionModel, self).__init__(model_type='%s Regression' % regression_type)
self.formula = formula
self.df = df
# Parse formula to get dependent variable and independent variable(s).
self.y_var, self.x_vars = parse_lin_reg_formula(self.formula)
if self.df is not None:
self.response_vector = self._response_vector()
self.num_obs = self._num_obs()
self.num_xvars = self._num_xvars()
self.design_matrix = self._design_matrix
def __init__(self, formula, df=None, max_depth=50):
"""
Parameters
----------
:param formula:
:param df:
:param max_depth:
"""
super(TreeRegression, self).__init__(model_type='Tree Regression')
self.formula = formula
self.df = df
self.max_depth = max_depth
self.y_var, self.x_vars = parse_lin_reg_formula(self.formula)
self.graph = nx.DiGraph()
self.node_queue = Queue.Queue()
if self.df is not None:
self.num_obs = len(self.df.index)
self.df[self.y_var] = df[self.y_var].astype('category')
self.discrete_xvars = list(self.df.select_dtypes(include=['bool', 'object', 'category']).columns)
