using new universe

alphamind/data/engines/sqlengine.py

@@ -43,7 +43,6 @@ from alphamind.data.engines.utilities import _map_risk_model_table
 from alphamind.data.engines.utilities import factor_tables
 from PyFin.api import advanceDateByCalendar
 
-
 risk_styles = ['BETA',
                'MOMENTUM',
                'SIZE',
@@ -204,11 +203,7 @@ class SqlEngine(object):
         end_date = advanceDateByCalendar('china.sse', end_date,
                                          str(1 + horizon + offset + DAILY_RETURN_OFFSET) + 'b').strftime('%Y-%m-%d')
 
-        cond = universe.query_range(start_date, end_date)
-        big_table = join(Market, UniverseTable,
-                         and_(Market.trade_date == UniverseTable.trade_date,
-                              Market.code == UniverseTable.code,
-                              cond))
+        codes = universe.query(self, start_date, end_date, dates)
 
         stats = func.sum(self.ln_func(1. + Market.chgPct)).over(
             partition_by=Market.code,
@@ -216,9 +211,13 @@ class SqlEngine(object):
             rows=(1 + offset + DAILY_RETURN_OFFSET, 1 + horizon + offset + DAILY_RETURN_OFFSET)).label('dx')
 
         query = select([Market.trade_date, Market.code, stats]) \
-            .select_from(big_table)
+            .where(
+            and_(Market.trade_date.in_(codes.trade_date.dt.strftime('%Y-%m-%d').unique()),
+                 Market.code.in_(codes.code.tolist()))
+        )
 
         df = pd.read_sql(query, self.session.bind).dropna()
+        df = pd.merge(df, codes, how='inner', on=['trade_date', 'code'])
 
         if dates:
             df = df[df.trade_date.isin(dates)]
@@ -355,7 +354,7 @@ class SqlEngine(object):
         else:
             factor_cols = _map_factors(dependency, factor_tables)
 
-        cond = universe.query_range(start_date, end_date, dates)
+        codes = universe.query(self, start_date, end_date, dates)
 
         big_table = FullFactor
 
@@ -370,16 +369,15 @@ class SqlEngine(object):
                                                              FullFactor.code == t.code,
                                                              FullFactor.trade_date.between(start_date, end_date)))
 
-        big_table = join(big_table, UniverseTable,
-                         and_(FullFactor.trade_date == UniverseTable.trade_date,
-                              FullFactor.code == UniverseTable.code,
-                              cond))
-
         query = select(
             [FullFactor.trade_date, FullFactor.code, FullFactor.isOpen] + list(factor_cols.keys())) \
-            .select_from(big_table).distinct()
+            .select_from(big_table).where(
+            and_(FullFactor.trade_date.in_(codes.trade_date.dt.strftime('%Y-%m-%d').unique()),
+                 FullFactor.code.in_(codes.code.tolist()))
+        ).distinct()
 
-        df = pd.read_sql(query, self.engine).sort_values(['trade_date', 'code'])
+        df = pd.read_sql(query, self.engine)
+        df = pd.merge(df, codes, how='inner', on=['trade_date', 'code']).sort_values(['trade_date', 'code'])
 
         if external_data is not None:
             df = pd.merge(df, external_data, on=['trade_date', 'code']).dropna()
@@ -479,17 +477,16 @@ class SqlEngine(object):
 
         risk_exposure_cols = [FullFactor.__table__.columns[f] for f in total_risk_factors if f not in set(excluded)]
 
-        cond = universe.query_range(start_date, end_date, dates)
-        big_table = join(FullFactor, UniverseTable,
-                         and_(FullFactor.trade_date == UniverseTable.trade_date,
-                              FullFactor.code == UniverseTable.code,
-                              cond))
+        codes = universe.query(self, start_date, end_date, dates)
 
         query = select(
-            [FullFactor.trade_date, FullFactor.code, special_risk_col] + risk_exposure_cols) \
-            .select_from(big_table).distinct()
+            [FullFactor.trade_date, FullFactor.code, special_risk_col] + risk_exposure_cols).where(
+            and_(FullFactor.trade_date.in_(codes.trade_date.dt.strftime('%Y-%m-%d').unique()),
+                 FullFactor.code.in_(codes.code.tolist()))
+        ).distinct()
 
         risk_exp = pd.read_sql(query, self.engine)
+        risk_exp = pd.merge(risk_exp, codes, how='inner', on=['trade_date', 'code']).sort_values(['trade_date', 'code'])
 
         return risk_cov, risk_exp
 
@@ -519,28 +516,20 @@ class SqlEngine(object):
                              dates: Iterable[str] = None,
                              category: str = 'sw'):
         industry_category_name = _map_industry_category(category)
-        cond = universe.query_range(start_date, end_date, dates)
-
-        if dates:
-            big_table = join(Industry, UniverseTable,
-                             and_(Industry.trade_date == UniverseTable.trade_date,
-                                  Industry.code == UniverseTable.code,
-                                  Industry.industry == industry_category_name,
-                                  Industry.trade_date.in_(dates),
-                                  cond))
-        else:
-            big_table = join(Industry, UniverseTable,
-                             and_(Industry.trade_date == UniverseTable.trade_date,
-                                  Industry.code == UniverseTable.code,
-                                  Industry.industry == industry_category_name,
-                                  Industry.trade_date.between(start_date, end_date),
-                                  cond))
+        codes = universe.query(self, start_date, end_date, dates)
 
         query = select([Industry.trade_date,
                         Industry.code,
                         Industry.industryID1.label('industry_code'),
-                        Industry.industryName1.label('industry')]).select_from(big_table).distinct()
-        return pd.read_sql(query, self.engine)
+                        Industry.industryName1.label('industry')]).where(
+            and_(Industry.trade_date.in_(codes.trade_date.dt.strftime('%Y-%m-%d').unique()),
+                 Industry.code.in_(codes.code.tolist()),
+                 Industry.industry == industry_category_name)
+        ).distinct()
+
+        df = pd.read_sql(query, self.engine)
+        df = pd.merge(df, codes, how='inner', on=['trade_date', 'code']).sort_values(['trade_date', 'code'])
+        return df
 
     def fetch_data(self, ref_date: str,
                    factors: Iterable[str],
@@ -778,9 +767,14 @@ class SqlEngine(object):
 
 
 if __name__ == '__main__':
+
+    universe = Universe('ss', ['hs300'])
+
     engine = SqlEngine()
 
-    df = pd.DataFrame(dict(trade_date=['2017-11-24'],
-                           portfolio_name=['test']))
+    df = engine.fetch_factor_range(universe,
+                                   ['closePrice'],
+                                   start_date='2012-12-21',
+                                   end_date='2017-12-25')
 
-    engine.insert_portfolio_schedule(df)
+    print(df)
\ No newline at end of file

alphamind/data/engines/universe.py

@@ -22,87 +22,6 @@ from alphamind.data.transformer import Transformer
 
 class Universe(object):
 
-    def __init__(self,
-                 name,
-                 base_universe: Iterable[str] = None,
-                 include_universe: Iterable[str] = None,
-                 exclude_universe: Iterable[str] = None,
-                 include_codes: Iterable[str] = None,
-                 exclude_codes: Iterable[str] = None):
-
-        self.name = name
-        self.base_universe = base_universe
-        self.include_universe = include_universe
-        self.exclude_universe = exclude_universe
-        self.include_codes = include_codes
-        self.exclude_codes = exclude_codes
-
-    def _create_condition(self):
-
-        all_and_conditions = []
-
-        univ_in = UniverseTable.universe.in_(self.base_universe)
-        all_and_conditions.append(univ_in)
-
-        if self.exclude_universe:
-            univ_out = UniverseTable.universe.notin_(self.exclude_universe)
-            all_and_conditions.append(univ_out)
-
-        if self.exclude_codes:
-            codes_out = UniverseTable.code.notin_(self.exclude_codes)
-            all_and_conditions.append(codes_out)
-
-        all_or_conditions = []
-        if self.include_universe:
-            univ_in = UniverseTable.universe.in_(self.include_universe)
-            all_or_conditions.append(univ_in)
-
-        if self.include_codes:
-            codes_in = UniverseTable.code.in_(self.include_codes)
-            all_or_conditions.append(codes_in)
-
-        return all_and_conditions, all_or_conditions
-
-    def query(self, ref_date):
-        all_and_conditions, all_or_conditions = self._create_condition()
-
-        if all_or_conditions:
-            query = and_(
-                UniverseTable.trade_date == ref_date,
-                or_(
-                    and_(*all_and_conditions),
-                    *all_or_conditions
-                )
-            )
-        else:
-            query = and_(
-                UniverseTable.trade_date == ref_date,
-                *all_and_conditions
-            )
-
-        return query
-
-    def query_range(self, start_date=None, end_date=None, dates=None):
-        all_and_conditions, all_or_conditions = self._create_condition()
-        dates_cond = UniverseTable.trade_date.in_(dates) if dates else UniverseTable.trade_date.between(start_date,
-                                                                                                        end_date)
-
-        if all_or_conditions:
-            query = and_(
-                dates_cond,
-                or_(
-                    and_(*all_and_conditions),
-                    *all_or_conditions
-                )
-            )
-        else:
-            query = and_(dates_cond, *all_and_conditions)
-
-        return query
-
-
-class UniverseNew(object):
-
     def __init__(self, name, base_universe, filter_cond=None):
         self.name = name
         self.base_universe = base_universe