Skip to content

A
alpha-mind
Commit 9ace5e9b authored by Dr.李's avatar Dr.李
Browse files
Options

update universe

parent 2c46b551
Hide whitespace changes
Inline Side-by-side
Showing with 153 additions and 120 deletions
alphamind/data/engines/universe.py
View file @ 9ace5e9b
...@@ -5,142 +5,175 @@ Created on 2017-7-7 ...@@ -5,142 +5,175 @@ Created on 2017-7-7
@author: cheng.li @author: cheng.li
""" """
from typing import Iterable import sys
import abc
import pandas as pd import pandas as pd
from simpleutils.miscellaneous import list_eq
from sqlalchemy import and_ from sqlalchemy import and_
from sqlalchemy import or_ from sqlalchemy import or_
from sqlalchemy import not_
from sqlalchemy import select from sqlalchemy import select
from sqlalchemy import join
from sqlalchemy import outerjoin
from alphamind.data.dbmodel.models import Universe as UniverseTable from alphamind.data.dbmodel.models import Universe as UniverseTable
from alphamind.data.engines.utilities import _map_factors
from alphamind.data.dbmodel.models import Market
from alphamind.data.engines.utilities import factor_tables
from alphamind.data.transformer import Transformer
from alphamind.utilities import encode
from alphamind.utilities import decode
class Universe(object):
def __init__(self,
name: str,
base_universe: Iterable,
exclude_universe: Iterable = None,
special_codes: Iterable = None,
filter_cond=None):
self.name = name
self.base_universe = sorted([u.lower() for u in base_universe]) if base_universe else []
self.exclude_universe = sorted([u.lower() for u in exclude_universe]) if exclude_universe else []
self.special_codes = sorted(special_codes) if special_codes else []
self.filter_cond = filter_cond
def __eq__(self, rhs):
return self.name == rhs.name \
and list_eq(self.base_universe, rhs.base_universe) \
and list_eq(self.exclude_universe, rhs.exclude_universe) \
and list_eq(self.special_codes, rhs.special_codes) \
and str(self.filter_cond) == str(rhs.filter_cond)
@property
def is_filtered(self):
return True if self.filter_cond is not None else False
def _query_statements(self, start_date, end_date, dates):
or_conditions = []
for u in self.base_universe:
or_conditions.append(
getattr(UniverseTable, u) == 1
)
if self.special_codes:
or_conditions.append(UniverseTable.code.in_(self.special_codes))
query = or_(
*or_conditions
)
and_conditions = []
if self.exclude_universe:
for u in self.exclude_universe:
and_conditions.append( getattr(UniverseTable, u) != 1)
class BaseUniverse(metaclass=abc.ABCMeta):
@abc.abstractmethod
def condition(self):
pass
def __add__(self, rhs):
return OrUniverse(self, rhs)
def __sub__(self, rhs):
return XorUniverse(self, rhs)
def __and__(self, rhs):
return AndUniverse(self, rhs)
def __or__(self, rhs):
return OrUniverse(self, rhs)
def isin(self, rhs):
return AndUniverse(self, rhs)
@abc.abstractmethod
def save(self):
pass
@classmethod
def load(cls, u_desc: dict):
pass
def query(self, engine, start_date: str = None, end_date: str = None, dates=None):
query = select([UniverseTable.trade_date, UniverseTable.code]).where(
self._query_statements(start_date, end_date, dates)
)
return pd.read_sql(query, engine.engine)
def _query_statements(self, start_date: str = None, end_date: str = None, dates=None):
return and_( return and_(
query, self.condition(),
*and_conditions, UniverseTable.trade_date.in_(dates) if dates else UniverseTable.trade_date.between(start_date, end_date)
UniverseTable.trade_date.in_(dates) if dates else UniverseTable.trade_date.between(start_date, end_date), )
class Universe(BaseUniverse):
def __init__(self, u_name: str):
self.u_name = u_name
def condition(self):
return getattr(UniverseTable, self.u_name) == 1
def save(self):
return dict(
u_type=self.__class__.__name__,
u_name=self.u_name
)
@classmethod
def load(cls, u_desc: dict):
return cls(u_name=u_desc['u_name'])
def __eq__(self, other):
return self.u_name == other.u_name
class OrUniverse(BaseUniverse):
def __init__(self, lhs: BaseUniverse, rhs: BaseUniverse):
self.lhs = lhs
self.rhs = rhs
def condition(self):
return or_(self.lhs.condition(), self.rhs.condition())
def save(self):
return dict(
u_type=self.__class__.__name__,
lhs=self.lhs.save(),
rhs=self.rhs.save()
)
@classmethod
def load(cls, u_desc: dict):
lhs = u_desc['lhs']
rhs = u_desc['rhs']
return cls(
lhs=getattr(sys.modules[__name__], lhs['u_type']).load(lhs),
rhs=getattr(sys.modules[__name__], rhs['u_type']).load(rhs),
) )
def query(self, engine, start_date: str = None, end_date: str = None, dates=None) -> pd.DataFrame: def __eq__(self, other):
return self.lhs == other.lhs and self.rhs == other.rhs and isinstance(other, OrUniverse)
universe_cond = self._query_statements(start_date, end_date, dates)
if self.filter_cond is None: class AndUniverse(BaseUniverse):
# simple case def __init__(self, lhs: BaseUniverse, rhs: BaseUniverse):
query = select([UniverseTable.trade_date, UniverseTable.code]).where( self.lhs = lhs
universe_cond self.rhs = rhs
)
return pd.read_sql(query, engine.engine) def condition(self):
else: return and_(self.lhs.condition(), self.rhs.condition())
if self.filter_cond is not None:
if isinstance(self.filter_cond, Transformer):
transformer = self.filter_cond
else:
transformer = Transformer(self.filter_cond)
dependency = transformer.dependency
factor_cols = _map_factors(dependency, factor_tables)
big_table = Market
for t in set(factor_cols.values()):
if t.__table__.name != Market.__table__.name:
big_table = outerjoin(big_table, t, and_(Market.trade_date == t.trade_date,
Market.code == t.code,
Market.trade_date.in_(
dates) if dates else Market.trade_date.between(
start_date, end_date)))
big_table = join(big_table, UniverseTable,
and_(Market.trade_date == UniverseTable.trade_date,
Market.code == UniverseTable.code,
universe_cond))
query = select(
[Market.trade_date, Market.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)
df = df[df[filter_fields[0]] == 1].reset_index()[['trade_date', 'code']]
return df
def save(self): def save(self):
return dict( return dict(
name=self.name, u_type=self.__class__.__name__,
base_universe=self.base_universe, lhs=self.lhs.save(),
exclude_universe=self.exclude_universe, rhs=self.rhs.save()
special_codes=self.special_codes,
filter_cond=encode(self.filter_cond)
) )
@classmethod @classmethod
def load(cls, universe_desc: dict): def load(cls, u_desc: dict):
name = universe_desc['name'] lhs = u_desc['lhs']
base_universe = universe_desc['base_universe'] rhs = u_desc['rhs']
exclude_universe = universe_desc['exclude_universe'] return cls(
special_codes = universe_desc['special_codes'] lhs=getattr(sys.modules[__name__], lhs['u_type']).load(lhs),
filter_cond = decode(universe_desc['filter_cond']) rhs=getattr(sys.modules[__name__], rhs['u_type']).load(rhs),
)
return cls(name=name,
base_universe=base_universe, def __eq__(self, other):
exclude_universe=exclude_universe, return self.lhs == other.lhs and self.rhs == other.rhs and isinstance(other, AndUniverse)
special_codes=special_codes,
filter_cond=filter_cond)
class XorUniverse(BaseUniverse):
def __init__(self, lhs: BaseUniverse, rhs: BaseUniverse):
self.lhs = lhs
self.rhs = rhs
def condition(self):
return and_(self.lhs.condition(), not_(self.rhs.condition()))
def save(self):
return dict(
u_type=self.__class__.__name__,
lhs=self.lhs.save(),
rhs=self.rhs.save()
)
@classmethod
def load(cls, u_desc: dict):
lhs = u_desc['lhs']
rhs = u_desc['rhs']
return cls(
lhs=getattr(sys.modules[__name__], lhs['u_type']).load(lhs),
rhs=getattr(sys.modules[__name__], rhs['u_type']).load(rhs),
)
def __eq__(self, other):
return self.lhs == other.lhs and self.rhs == other.rhs and isinstance(other, XorUniverse)
def load_universe(u_desc: dict):
u_name = u_desc['u_type']
if u_name == 'Universe':
return Universe.load(u_desc)
elif u_name == 'OrUniverse':
return OrUniverse.load(u_desc)
elif u_name == 'AndUniverse':
return AndUniverse.load(u_desc)
elif u_name == 'XorUniverse':
return XorUniverse.load(u_desc)
if __name__ == '__main__': if __name__ == '__main__':
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment