INNER CODE UNIT · Python

select

dmarienko/QuantTradingResearch · models/helpers.py:65

    def select(self, selection):
        n = copy.copy(self)
        n.selection = set(n.data.keys()) & set(selection)
        return n
        
    def get_data(self, mode='prices', where='all data'):
        x, y = (self.X_price, self.Y_price) if mode.startswith('price') else (self.X_ret, self.Y_ret)
        if where.startswith('train'):
            x, y = x[:self.train_date], y[:self.train_date]
        elif where.startswith('test'):
            x, y = x[self.train_date:], y[self.train_date:]
        return (x[self.selection], y) if self.selection is not None else (x, y)
        
    def train_set(self, mode='prices'):
        return self.get_data(mode, 'train')
    
    def test_set(self, mode='prices'):
        return self.get_data(mode, 'test')

View source record →