experimentally added new universe implementation

alphamind/data/engines/sqlengine.py

@@ -25,11 +25,7 @@ from alphamind.data.dbmodel.models import IndexComponent
 from alphamind.data.dbmodel.models import Industry
 from alphamind.data.dbmodel.models import Experimental
 from alphamind.data.dbmodel.models import RiskMaster
-from alphamind.data.dbmodel.models import RiskCovDay
-from alphamind.data.dbmodel.models import RiskCovShort
-from alphamind.data.dbmodel.models import RiskCovLong
 from alphamind.data.dbmodel.models import FullFactor
-from alphamind.data.dbmodel.models import Gogoal
 from alphamind.data.dbmodel.models import Models
 from alphamind.data.dbmodel.models import Market
 from alphamind.data.dbmodel.models import IndexMarket
@@ -42,8 +38,13 @@ from alphamind.data.transformer import Transformer
 from alphamind.model.loader import load_model
 from alphamind.formula.utilities import encode_formula
 from alphamind.formula.utilities import decode_formula
+from alphamind.data.engines.utilities import _map_factors
+from alphamind.data.engines.utilities import _map_industry_category
+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',
@@ -90,42 +91,11 @@ macro_styles = ['COUNTRY']
 
 total_risk_factors = risk_styles + industry_styles + macro_styles
 
-factor_tables = [FullFactor, Gogoal, Experimental]
-
 DEFAULT_URL = 'postgresql+psycopg2://postgres:A12345678!@10.63.6.220/alpha'
 
 DAILY_RETURN_OFFSET = 0
 
 
-def _map_risk_model_table(risk_model: str) -> tuple:
-    if risk_model == 'day':
-        return RiskCovDay, FullFactor.d_srisk
-    elif risk_model == 'short':
-        return RiskCovShort, FullFactor.s_srisk
-    elif risk_model == 'long':
-        return RiskCovLong, FullFactor.l_srisk
-    else:
-        raise ValueError("risk model name {0} is not recognized".format(risk_model))
-
-
-def _map_factors(factors: Iterable[str], used_factor_tables) -> Dict:
-    factor_cols = {}
-    excluded = {'trade_date', 'code', 'isOpen'}
-    for f in factors:
-        for t in used_factor_tables:
-            if f not in excluded and f in t.__table__.columns:
-                factor_cols[t.__table__.columns[f]] = t
-                break
-    return factor_cols
-
-
-def _map_industry_category(category: str) -> str:
-    if category == 'sw':
-        return '申万行业分类'
-    else:
-        raise ValueError("No other industry is supported at the current time")
-
-
 class SqlEngine(object):
     def __init__(self,
                  db_url: str = None):

alphamind/data/engines/universe.py

@@ -6,9 +6,18 @@ Created on 2017-7-7
 """
 
 from typing import Iterable
+import pandas as pd
 from sqlalchemy import and_
 from sqlalchemy import or_
+from sqlalchemy import select
+from sqlalchemy import join
+from sqlalchemy import outerjoin
+from PyFin.api import pyFinAssert
 from alphamind.data.dbmodel.models import Universe as UniverseTable
+from alphamind.data.dbmodel.models import FullFactor
+from alphamind.data.engines.utilities import _map_factors
+from alphamind.data.engines.utilities import factor_tables
+from alphamind.data.transformer import Transformer
 
 
 class Universe(object):
@@ -89,3 +98,68 @@ class Universe(object):
             query = and_(dates_cond, *all_and_conditions)
 
         return query
+
+
+class UniverseNew(object):
+
+    def __init__(self, name, base_universe):
+        self.name = name
+        self.base_universe = base_universe
+
+    def query(self, engine, ref_date: str, filter_cond=None) -> pd.DataFrame:
+
+        if filter_cond is None:
+            # simple case
+            query = select([UniverseTable.trade_date, UniverseTable.code]).where(
+                and_(
+                    UniverseTable.trade_date == ref_date,
+                    UniverseTable.universe.in_(self.base_universe)
+                )
+            )
+            return pd.read_sql(query, engine.engine)
+        else:
+            if isinstance(filter_cond, Transformer):
+                transformer = filter_cond
+            else:
+                transformer = Transformer(filter_cond)
+
+            dependency = transformer.dependency
+            factor_cols = _map_factors(dependency, factor_tables)
+            big_table = FullFactor
+
+            for t in set(factor_cols.values()):
+                if t.__table__.name != FullFactor.__table__.name:
+                    big_table = outerjoin(big_table, t, and_(FullFactor.trade_date == t.trade_date,
+                                                             FullFactor.code == t.code,
+                                                             FullFactor.trade_date == ref_date))
+
+            universe_cond = and_(
+                UniverseTable.trade_date == ref_date,
+                UniverseTable.universe.in_(self.base_universe)
+            )
+
+            big_table = join(big_table, UniverseTable,
+                             and_(FullFactor.trade_date == UniverseTable.trade_date,
+                                  FullFactor.code == UniverseTable.code,
+                                  universe_cond))
+
+            query = select(
+                [FullFactor.trade_date, FullFactor.code] + list(factor_cols.keys())) \
+                .select_from(big_table).distinct()
+
+            df = pd.read_sql(query, engine.engine).sort_values(['trade_date', 'code']).dropna()
+            df.set_index('trade_date', inplace=True)
+            filter_fields = transformer.names
+            pyFinAssert(len(filter_fields) == 1, ValueError, "filter fields can only be 1")
+            df = transformer.transform('code', df)
+            return df[df[filter_fields[0]] == 1].reset_index()[['trade_date', 'code']]
+
+
+if __name__ == '__main__':
+    from PyFin.api import *
+    from alphamind.data.engines.sqlengine import SqlEngine
+    engine = SqlEngine()
+    universe = UniverseNew('ss', ['hs300'])
+    print(universe.query(engine, '2017-12-21', LAST('closePrice') < 5))
+
+

alphamind/data/engines/utilities.py

+# -*- coding: utf-8 -*-
+"""
+Created on 2017-12-25
+
+@author: cheng.li
+"""
+
+from typing import Iterable
+from typing import Dict
+from alphamind.data.dbmodel.models import RiskCovDay
+from alphamind.data.dbmodel.models import RiskCovShort
+from alphamind.data.dbmodel.models import RiskCovLong
+from alphamind.data.dbmodel.models import FullFactor
+from alphamind.data.dbmodel.models import Gogoal
+from alphamind.data.dbmodel.models import Experimental
+
+
+factor_tables = [FullFactor, Gogoal, Experimental]
+
+
+def _map_risk_model_table(risk_model: str) -> tuple:
+    if risk_model == 'day':
+        return RiskCovDay, FullFactor.d_srisk
+    elif risk_model == 'short':
+        return RiskCovShort, FullFactor.s_srisk
+    elif risk_model == 'long':
+        return RiskCovLong, FullFactor.l_srisk
+    else:
+        raise ValueError("risk model name {0} is not recognized".format(risk_model))
+
+
+def _map_factors(factors: Iterable[str], used_factor_tables) -> Dict:
+    factor_cols = {}
+    excluded = {'trade_date', 'code', 'isOpen'}
+    for f in factors:
+        for t in used_factor_tables:
+            if f not in excluded and f in t.__table__.columns:
+                factor_cols[t.__table__.columns[f]] = t
+                break
+    return factor_cols
+
+
+def _map_industry_category(category: str) -> str:
+    if category == 'sw':
+        return '申万行业分类'
+    else:
+        raise ValueError("No other industry is supported at the current time")
\ No newline at end of file