added conditions to filter universe

alphamind/data/engines/sqlengine.py

@@ -135,6 +135,16 @@ class SqlEngine(object):
         if univ.exclude_codes:
             codes_set -= set(univ.exclude_codes)
 
+        if univ.filter_cond is not None:
+            query = select([UniverseTable.Code]).distinct().where(
+                and_(UniverseTable.Date == ref_date,
+                     univ.filter_cond)
+            )
+
+            cursor = self.engine.execute(query)
+            filtered_codes = {c[0] for c in cursor.fetchall()}
+            codes_set = codes_set.intersection(filtered_codes)
+
         return sorted(codes_set)
 
     def fetch_dx_return(self, ref_date, codes, expiry_date=None, horizon=1):
@@ -225,14 +235,34 @@ class SqlEngine(object):
 
 if __name__ == '__main__':
     db_url = 'mssql+pymssql://licheng:A12345678!@10.63.6.220/alpha?charset=utf8'
-    universe = Universe('zz500', ['zz500'])
 
+    from alphamind.data.dbmodel.models import Uqer
+
+    import datetime as dt
+
+    universe = Universe('zz500', ['zz500'], filter_cond=Uqer.BLEV >= 0.1)
     engine = SqlEngine(db_url)
     ref_date = '2017-07-04'
 
     codes = engine.fetch_codes(ref_date, universe)
-    df = engine.fetch_data(ref_date, factors=['EPS', 'CFinc1'], codes=codes)
 
-    print(df)
+    start = dt.datetime.now()
+    for i in range(100):
+        codes = engine.fetch_codes(ref_date, universe)
+    print(dt.datetime.now() - start)
+
+    print(codes)
+    print(len(codes))
+
+    universe = Universe('zz500', ['zz500'])
+    engine = SqlEngine(db_url)
+    ref_date = '2017-07-04'
+
+    start = dt.datetime.now()
+    for i in range(100):
+        codes = engine.fetch_codes(ref_date, universe)
+    print(dt.datetime.now() - start)
 
+    print(codes)
+    print(len(codes))
 

alphamind/data/engines/universe.py

@@ -6,6 +6,9 @@ Created on 2017-7-7
 """
 
 from typing import Iterable
+from sqlalchemy import and_
+from alphamind.data.dbmodel.models import Uqer
+from alphamind.data.dbmodel.models import Universe as UniverseTable
 
 
 class Universe(object):
@@ -15,7 +18,8 @@ class Universe(object):
                  include_universe: Iterable[str]=None,
                  exclude_universe: Iterable[str]=None,
                  include_codes: Iterable[str]=None,
-                 exclude_codes: Iterable[str]=None):
+                 exclude_codes: Iterable[str]=None,
+                 filter_cond=None):
 
         self.name = name
         self.include_universe = include_universe
@@ -23,3 +27,15 @@ class Universe(object):
         self.include_codes = include_codes
         self.exclude_codes = exclude_codes
 
+        if filter_cond is not None:
+            self.filter_cond = self.format_cond(filter_cond)
+        else:
+            self.filter_cond = None
+
+    def format_cond(self, filter_cond):
+        filter_cond = and_(
+            filter_cond,
+            Uqer.Code == UniverseTable.Code,
+            Uqer.Date == UniverseTable.Date
+        )
+        return filter_cond

alphamind/data/neutralize.py

@@ -105,3 +105,15 @@ def ls_explain(x: np.ndarray, b: np.ndarray) -> np.ndarray:
         explained[:, :, i] = b[:, i] * x
     return explained
 
+
+if __name__ == '__main__':
+    x = np.random.randn(3000)
+    y = np.random.randn(3000)
+    w = np.ones(3000)
+
+    import datetime as dt
+
+    start = dt.datetime.now()
+    for i in range(1000):
+        ls_fit(x, y, w)
+    print(dt.datetime.now() - start)