msd-backtest 0.1.2__tar.gz → 0.1.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (35) hide show
  1. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/PKG-INFO +3 -2
  2. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/__init__.py +12 -9
  3. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/run/cli.py +25 -21
  4. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/run/msd.py +4 -1
  5. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/run/report.py +4 -4
  6. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/run/review.py +16 -9
  7. msd_backtest-0.1.4/mbt/select/__init__.py +18 -0
  8. msd_backtest-0.1.4/mbt/select/cli.py +152 -0
  9. msd_backtest-0.1.4/mbt/select/msd.py +124 -0
  10. msd_backtest-0.1.4/mbt/select/selector.py +89 -0
  11. msd_backtest-0.1.4/mbt/select/stocks.py +75 -0
  12. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/msd_backtest.egg-info/PKG-INFO +3 -2
  13. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/msd_backtest.egg-info/SOURCES.txt +5 -0
  14. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/msd_backtest.egg-info/requires.txt +2 -1
  15. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/pyproject.toml +5 -5
  16. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/tests/test_alpha.py +1 -1
  17. msd_backtest-0.1.4/tests/test_selector.py +140 -0
  18. msd_backtest-0.1.2/mbt/select/__init__.py +0 -0
  19. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/README.md +0 -0
  20. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/account.py +0 -0
  21. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/data_provider.py +0 -0
  22. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/run/__init__.py +0 -0
  23. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/run/backend.py +0 -0
  24. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/run/runner.py +0 -0
  25. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/run/server.py +0 -0
  26. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/mbt/strategy.py +0 -0
  27. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/msd_backtest.egg-info/dependency_links.txt +0 -0
  28. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/msd_backtest.egg-info/top_level.txt +0 -0
  29. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/setup.cfg +0 -0
  30. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/tests/test_account.py +0 -0
  31. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/tests/test_load_strategy.py +0 -0
  32. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/tests/test_report.py +0 -0
  33. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/tests/test_runner.py +0 -0
  34. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/tests/test_strategy.py +0 -0
  35. {msd_backtest-0.1.2 → msd_backtest-0.1.4}/tests/test_strategy_loading_complex.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: msd-backtest
3
- Version: 0.1.2
3
+ Version: 0.1.4
4
4
  Summary: MSD backtest framework
5
5
  Requires-Python: >=3.11
6
6
  Description-Content-Type: text/markdown
@@ -10,7 +10,8 @@ Requires-Dist: py-alpha-lib>=0.2; extra == "alpha"
10
10
  Provides-Extra: msd
11
11
  Requires-Dist: pymsd>=0.1; extra == "msd"
12
12
  Requires-Dist: requests>=2.32; extra == "msd"
13
- Requires-Dist: pandas>=3.0; extra == "msd"
13
+ Requires-Dist: polars>=1.0.0; extra == "msd"
14
+ Requires-Dist: xlsxwriter>=3.2.9; extra == "msd"
14
15
  Provides-Extra: server
15
16
  Requires-Dist: fastapi[standard]>=0.135.1; extra == "server"
16
17
  Requires-Dist: msd-backtest[cli]; extra == "server"
@@ -9,20 +9,23 @@ from .account import (
9
9
  ACTION_KEEP,
10
10
  )
11
11
  from .run.runner import Runner
12
+ from .select import Selector, SelectorDataProvider
12
13
 
13
14
 
14
15
  __all__ = [
15
- "Context",
16
- "Strategy",
17
- "Runner",
18
- "Account",
19
- "Fee",
20
- "DataProvider",
21
- "Operation",
22
- "Position",
23
16
  "ACTION_BUY",
24
17
  "ACTION_BUY_NO_CASH",
18
+ "ACTION_KEEP",
25
19
  "ACTION_SELL",
26
20
  "ACTION_SELL_NO_HOLD",
27
- "ACTION_KEEP",
21
+ "Account",
22
+ "Context",
23
+ "DataProvider",
24
+ "Fee",
25
+ "Operation",
26
+ "Position",
27
+ "Runner",
28
+ "Selector",
29
+ "SelectorDataProvider",
30
+ "Strategy",
28
31
  ]
@@ -9,7 +9,7 @@ from mbt.run.msd import load_data
9
9
  import logging
10
10
  from typing import Tuple
11
11
  import argparse
12
- import pandas as pd
12
+ import polars as pl
13
13
  import orjson
14
14
  import os
15
15
  import numpy as np
@@ -111,7 +111,9 @@ def parse_args() -> Tuple[RunRequest, str, str, Fee]:
111
111
  return req, args.msd, args.output, fee
112
112
 
113
113
 
114
- def float_format(f: float) -> str:
114
+ def float_format(f: float | None) -> str:
115
+ if f is None or (isinstance(f, float) and np.isnan(f)):
116
+ return ""
115
117
  s = f"{f:.6f}"
116
118
  if "." in s:
117
119
  s = s.rstrip("0").rstrip(".")
@@ -134,24 +136,26 @@ def build_output(reports: list[Report], output: str):
134
136
  logger.info(f"output kind: {ext}")
135
137
  if ext == "csv":
136
138
  for report in reports:
137
- df = pd.DataFrame(report.detailed(with_symbol=True))
138
- df.to_csv(output, index=False, float_format=float_format) # type: ignore
139
+ df = pl.DataFrame(report.detailed(with_symbol=True))
140
+ float_cols = [col for col, dtype in zip(df.columns, df.dtypes) if dtype in (pl.Float32, pl.Float64)]
141
+ df = df.with_columns([
142
+ pl.col(col).map_elements(float_format, return_dtype=pl.String)
143
+ for col in float_cols
144
+ ])
145
+ df.write_csv(output)
139
146
  elif ext == "xlsx":
140
- with pd.ExcelWriter(output) as writer:
141
- df_metrics = pd.DataFrame([r.metrics(True) for r in reports])
142
- df_metrics.to_excel(writer, sheet_name="metrics", index=False)
143
- for report in reports:
144
- data = report.detailed()
145
- data['price'] = report.price
146
- data['fw_price'] = report.fw_price
147
- data['bw_price'] = report.bw_price
148
- df = pd.DataFrame(data)
149
- df.to_excel(
150
- writer,
151
- sheet_name=report.symbol,
152
- index=False,
153
- float_format="%.6f",
154
- )
147
+ import xlsxwriter
148
+ wb = xlsxwriter.Workbook(output)
149
+ df_metrics = pl.DataFrame([r.metrics(True) for r in reports])
150
+ df_metrics.write_excel(workbook=wb, worksheet="metrics")
151
+ for report in reports:
152
+ data = report.detailed()
153
+ data['price'] = report.price
154
+ data['fw_price'] = report.fw_price
155
+ data['bw_price'] = report.bw_price
156
+ df = pl.DataFrame(data)
157
+ df.write_excel(workbook=wb, worksheet=report.symbol)
158
+ wb.close()
155
159
  elif ext == "json":
156
160
  data = build_json_report(reports)
157
161
  logger.info("build json")
@@ -163,8 +167,8 @@ def build_output(reports: list[Report], output: str):
163
167
 
164
168
 
165
169
  def echo_reports(reports: list[Report]):
166
- df_metrics = pd.DataFrame([r.metrics(True) for r in reports])
167
- print(df_metrics.to_json(orient='records', indent=2))
170
+ metrics_list = [r.metrics(True) for r in reports]
171
+ print(orjson.dumps(metrics_list, option=orjson.OPT_INDENT_2 | orjson.OPT_SERIALIZE_NUMPY).decode())
168
172
 
169
173
 
170
174
  def main():
@@ -28,7 +28,7 @@ def load_data(
28
28
  DataProvider and symbols
29
29
  """
30
30
 
31
- client = pymsd.create_msd_pandas(msd_host)
31
+ client = pymsd.create_msd_polars(msd_host)
32
32
 
33
33
  dfs = client.load(
34
34
  objs=symbols,
@@ -45,6 +45,9 @@ def load_data(
45
45
  return DataProvider({"ts": np.array([])}, symbols=[]), []
46
46
 
47
47
  data = client.adaptor.to_numpy(data)
48
+ for col in ["dividend", "transfer_shares", "right_shares", "right_price"]:
49
+ if col not in data:
50
+ data[col] = np.zeros_like(data["close"])
48
51
 
49
52
  dp = DataProvider(data, symbols=symbols, next_hook=next_hook_al)
50
53
 
@@ -1,6 +1,6 @@
1
1
  from mbt import Account
2
2
  from mbt import DataProvider
3
- import pandas as pd
3
+ import polars as pl
4
4
  import numpy as np
5
5
  import dataclasses
6
6
 
@@ -49,9 +49,9 @@ class Report:
49
49
  self.ror_original = np.round(uniform_ror(dp.slice("bw_price", s)), 4)
50
50
 
51
51
  # Clean NaNs by forward filling and filling leading NaNs with 0
52
- self.ror = pd.Series(self.ror).ffill().fillna(0.0).values
53
- self.ror_hold = pd.Series(self.ror_hold).ffill().fillna(0.0).values
54
- self.ror_original = pd.Series(self.ror_original).ffill().fillna(0.0).values
52
+ self.ror = pl.Series("", self.ror).fill_nan(None).fill_null(strategy="forward").fill_null(0.0).to_numpy()
53
+ self.ror_hold = pl.Series("", self.ror_hold).fill_nan(None).fill_null(strategy="forward").fill_null(0.0).to_numpy()
54
+ self.ror_original = pl.Series("", self.ror_original).fill_nan(None).fill_null(strategy="forward").fill_null(0.0).to_numpy()
55
55
 
56
56
  self.actions = dp.slice("actions", s)
57
57
  self.price = dp.slice("price", s)
@@ -2,7 +2,7 @@ from mbt.run.report import Report
2
2
  import orjson
3
3
  import sys
4
4
  import argparse
5
- import pandas as pd
5
+ import polars as pl
6
6
 
7
7
 
8
8
 
@@ -29,10 +29,17 @@ def parse_args():
29
29
 
30
30
  def do_stat(args):
31
31
  reports = load_reports(args.report)
32
- df = pd.DataFrame([r.metrics(True) for r in reports.values()])
33
- df.set_index("symbol", inplace=True)
32
+ df = pl.DataFrame([r.metrics(True) for r in reports.values()])
33
+ if "symbol" in df.columns:
34
+ df = df.drop("symbol")
34
35
  stat = df.describe()
35
- print(stat.to_json())
36
+ stat_col = stat.columns[0]
37
+ stats = stat[stat_col].to_list()
38
+ res = {}
39
+ for col in stat.columns[1:]:
40
+ vals = stat[col].to_list()
41
+ res[col] = {s: v for s, v in zip(stats, vals) if s != "null_count" and v is not None}
42
+ print(orjson.dumps(res).decode())
36
43
 
37
44
  def do_detail(args):
38
45
  reports = load_reports(args.report)
@@ -44,12 +51,12 @@ def do_detail(args):
44
51
 
45
52
  def do_sort(args):
46
53
  reports = load_reports(args.report)
47
- df = pd.DataFrame([{"symbol":r.symbol, args.metric: getattr(r, args.metric, 0)} for r in reports.values()])
48
- df.set_index("symbol", inplace=True)
49
- df = df.sort_values(by=args.metric, ascending=args.order == "asc")
54
+ records = [{"symbol": r.symbol, "val": getattr(r, args.metric, 0)} for r in reports.values()]
55
+ records.sort(key=lambda x: x["val"], reverse=args.order != "asc")
50
56
  if args.limit:
51
- df = df.head(args.limit)
52
- print(df.to_json(orient="columns"))
57
+ records = records[:args.limit]
58
+ res = {args.metric: {r["symbol"]: r["val"] for r in records}}
59
+ print(orjson.dumps(res).decode())
53
60
 
54
61
  def load_reports(report_path: str) -> dict[str, Report]:
55
62
  with open(report_path, "r") as f:
@@ -0,0 +1,18 @@
1
+ from .selector import Selector, SelectorDataProvider
2
+ from .msd import MsdSelectorDataProvider
3
+ from .stocks import ALL_STOCKS, A_STOCKS, A_STOCKS_EXCLUDE_ST, FOUNDS, KIND_INDEX, KIND_FUND, KIND_STOCK, STATUS_NORMAL, STATUS_ST, STATUS_STAR_ST
4
+
5
+ __all__ = [
6
+ 'ALL_STOCKS',
7
+ 'A_STOCKS',
8
+ 'A_STOCKS_EXCLUDE_ST',
9
+ 'FOUNDS',
10
+ 'KIND_INDEX',
11
+ 'KIND_FUND',
12
+ 'STATUS_NORMAL',
13
+ 'STATUS_ST',
14
+ 'STATUS_STAR_ST',
15
+ 'Selector',
16
+ 'SelectorDataProvider',
17
+ 'MsdSelectorDataProvider',
18
+ ]
@@ -0,0 +1,152 @@
1
+ import logging
2
+ from dotenv import load_dotenv
3
+
4
+ load_dotenv(override=True)
5
+
6
+ from typing import Any
7
+ from mbt.select import Selector, SelectorDataProvider, MsdSelectorDataProvider
8
+ import mbt.select.stocks
9
+ import logging
10
+ from typing import Tuple
11
+ import argparse
12
+ import os
13
+ import sys
14
+ import importlib.util
15
+
16
+ logging.basicConfig(
17
+ format="%(asctime)s %(name)s %(levelname)s %(message)s", level=logging.INFO
18
+ )
19
+
20
+ logger = logging.getLogger("selector")
21
+ logger.setLevel(logging.DEBUG)
22
+
23
+
24
+ def expand_symbols(symbols: list[str]) -> list[str]:
25
+ """read as file when input is a file path, and expand it. otherwise, return the input as a list of symbols"""
26
+ expanded = set()
27
+ for s in symbols:
28
+ if os.path.isfile(s):
29
+ with open(s, "r") as f:
30
+ lines = filter(
31
+ lambda x: len(x) > 0 and not x.startswith("#"),
32
+ map(lambda x: x.strip(), f.readlines()),
33
+ )
34
+ expanded.update(lines)
35
+ else:
36
+ match s:
37
+ case "A_STOCKS_EXCLUDE_ST":
38
+ expanded.update(mbt.select.stocks.A_STOCKS_EXCLUDE_ST)
39
+ case "A_STOCKS":
40
+ expanded.update(mbt.select.stocks.A_STOCKS)
41
+ case "FOUNDS":
42
+ expanded.update(mbt.select.stocks.FOUNDS)
43
+ case _:
44
+ expanded.add(s)
45
+
46
+ l = list(expanded)
47
+ l.sort()
48
+ return l
49
+
50
+
51
+ def load_selector(py_file: str, **kwargs) -> Selector:
52
+ path = os.path.abspath(py_file)
53
+ dir_path = os.path.dirname(path)
54
+ if dir_path not in sys.path:
55
+ sys.path.insert(0, dir_path)
56
+
57
+ module_name = os.path.splitext(os.path.basename(path))[0]
58
+ spec = importlib.util.spec_from_file_location(module_name, path)
59
+ if spec and spec.loader:
60
+ module = importlib.util.module_from_spec(spec)
61
+ spec.loader.exec_module(module)
62
+ for obj in vars(module).values():
63
+ if isinstance(obj, type) and issubclass(obj, Selector) and obj is not Selector:
64
+ return obj(**kwargs) if kwargs is not None else obj()
65
+ raise ValueError(f"Could not find a Selector class in {py_file}")
66
+
67
+
68
+
69
+ def parse_args(args: list[str]) -> dict[str, Any]:
70
+ if args is None:
71
+ return {}
72
+ kwargs = {}
73
+
74
+ def try_bool(s: str) -> Tuple[bool, bool]:
75
+ match s.lower():
76
+ case "true" | "t" | "yes" | "y" | "on":
77
+ return True, True
78
+ case "false" | "f" | "no" | "n" | "off":
79
+ return False, True
80
+ case _:
81
+ return False, False
82
+
83
+ def try_int(s: str) -> Tuple[int, bool]:
84
+ try:
85
+ return int(s), True
86
+ except ValueError:
87
+ return 0, False
88
+
89
+ def try_float(s: str) -> Tuple[float, bool]:
90
+ try:
91
+ return float(s), True
92
+ except ValueError:
93
+ return 0.0, False
94
+
95
+ def try_float(s: str) -> Tuple[float, bool]:
96
+ try:
97
+ return float(s), True
98
+ except ValueError:
99
+ return 0.0, False
100
+
101
+ def try_value(s: str) -> Any:
102
+ for try_func in [try_bool, try_int, try_float]:
103
+ v, ok = try_func(s)
104
+ if ok:
105
+ return v
106
+ return s
107
+
108
+ for arg in args:
109
+ kv = arg.split(arg, 1)
110
+ if len(kv) != 2:
111
+ logger.warning(f'invalid arg "{arg}"')
112
+ continue
113
+ key, value = kv
114
+ kwargs[key] = try_value(value)
115
+ return kwargs
116
+
117
+ def main():
118
+ parser = argparse.ArgumentParser(description="Run a selector")
119
+ parser.add_argument("selector", help="Selector Python file")
120
+ parser.add_argument("-a", "--arg", nargs='*', help="any init arguments with format key=value")
121
+ parser.add_argument(
122
+ "-m", "--msd", type=str, default=os.environ.get("MSD_HOST", "http://localhost:50510"), help="MSD host, env $MSD_HOST"
123
+ )
124
+ parser.add_argument("-s", "--symbols", nargs="+", default=[], help="Symbols to run")
125
+ args = parser.parse_args()
126
+ kwargs = parse_args(args.arg)
127
+ selector = load_selector(args.selector, **kwargs)
128
+ logger.info(f"selector loaded: {selector}")
129
+
130
+ dp = MsdSelectorDataProvider(args.msd)
131
+ logger.info("data provider created")
132
+
133
+ init_symbols = expand_symbols(args.symbols)
134
+
135
+ stocks = selector.execute(dp, init_symbols)
136
+ logger.info("selector finished")
137
+ factors = getattr(selector, 'factors', None)
138
+ if factors is not None:
139
+ print(factors)
140
+ else:
141
+ for stock in stocks:
142
+ print(stock)
143
+
144
+ if __name__ == "__main__":
145
+ main()
146
+
147
+
148
+
149
+
150
+
151
+
152
+
@@ -0,0 +1,124 @@
1
+ import pymsd
2
+ from .selector import SelectorDataProvider
3
+ import polars as pl
4
+ import alpha as al
5
+ import pymsd
6
+ import logging
7
+
8
+ logger = logging.getLogger("selector")
9
+
10
+
11
+ class MsdSelectorDataProvider(SelectorDataProvider):
12
+
13
+ def __init__(self, msd_host: str):
14
+ super().__init__()
15
+ self._msd_host = msd_host
16
+ self._client = pymsd.create_msd_polars(msd_host)
17
+
18
+
19
+ def load_kline(self, symbols: list[str], lastN: int = 100) -> dict[str, pl.DataFrame]:
20
+ if lastN == 1:
21
+ dividend_lastN = 1
22
+ shares_lastN = 1
23
+ else:
24
+ dividend_lastN = max(lastN//20, 5)
25
+ shares_lastN = max(lastN//10, 5)
26
+
27
+ logger.debug(f"start load kline for {len(symbols)} symbols, lastN={lastN}")
28
+ dfs = self._client.load(
29
+ objs=symbols,
30
+ tables=["stock_kline_1d", "stock_dividend", "stock_shares"],
31
+ join={"stock_dividend": "zero", "*": "backward"},
32
+ start=[lastN, dividend_lastN, shares_lastN],
33
+ end=self.date,
34
+ )
35
+ logger.debug(f"finish load kline for {len(symbols)} symbols")
36
+
37
+
38
+ # Apply forward adjustment factor to price related columns
39
+ al.set_ctx(groups=1, flags=al.FLAG_SKIP_NAN)
40
+ for symbol, data in dfs.items():
41
+ if len(data["ts"]) == 1:
42
+ # Skip if there is only one day of data
43
+ continue
44
+ dividends = data["dividend"].to_numpy()
45
+ transfer_shares = data["transfer_shares"].to_numpy()
46
+ right_shares = data["right_shares"].to_numpy()
47
+ right_price = data["right_price"].to_numpy()
48
+ dfs[symbol] = data.with_columns(
49
+ close=al.FW_SPLIT(data["close"].to_numpy(), dividends, transfer_shares, right_shares, right_price),
50
+ open=al.FW_SPLIT(data["open"].to_numpy(), dividends, transfer_shares, right_shares, right_price),
51
+ high=al.FW_SPLIT(data["high"].to_numpy(), dividends, transfer_shares, right_shares, right_price),
52
+ low=al.FW_SPLIT(data["low"].to_numpy(), dividends, transfer_shares, right_shares, right_price)
53
+ )
54
+
55
+ return dfs
56
+
57
+ def load_financial(self, symbols: list[str], fields: list[str], only_year: bool = True, lastN: int = 3) -> dict[str, pl.DataFrame]:
58
+ if only_year:
59
+ lastN = lastN * 4 + 1
60
+
61
+ logger.debug(f"start load financial for {len(symbols)} symbols, lastN={lastN}")
62
+ dfs = self._client.load(
63
+ objs=symbols,
64
+ tables=["stock_financial"],
65
+ start=lastN,
66
+ end=self.date,
67
+ fields=fields,
68
+ )
69
+ logger.debug(f"finish load financial for {len(symbols)} symbols")
70
+
71
+ data = {}
72
+ for symbol, v in dfs.items():
73
+ df = v['stock_financial']
74
+ # df = df.with_columns(
75
+ # pl.col("ts").map_elements(lambda d: int(d.year), return_dtype=pl.Int32).alias("f_year"),
76
+ # pl.col("ts").map_elements(lambda d: int(d.month/3), return_dtype=pl.Int32).alias("f_quarter")
77
+ # )
78
+ if only_year:
79
+ data[symbol] = df.filter(pl.col('ts').dt.month() == 12)
80
+ else:
81
+ data[symbol] = df
82
+ return data
83
+
84
+ def load_snapshot(self, symbols: list[str], fin_fields: list[str] = [], fin_only_year: bool = True) -> pl.DataFrame:
85
+
86
+ dfs = self._client.load(
87
+ objs=symbols,
88
+ tables=["stock_kline_1d", "stock_shares"],
89
+ join={"stock_shares": "backward"},
90
+ start=1,
91
+ end=self.date
92
+ )
93
+
94
+ if len(fin_fields) > 0:
95
+ for obj, fin_df in self.load_financial(symbols, fin_fields, fin_only_year, 1).items():
96
+ if obj in dfs:
97
+ if len(fin_df) > 1:
98
+ fin_df = fin_df.tail(1)
99
+ dfs[obj] = pl.concat([dfs[obj], fin_df.rename({"ts": "f_ts"})], how='horizontal_extend')
100
+
101
+ return pl.concat([
102
+ df.with_columns(
103
+ pl.lit(obj).alias('obj')
104
+ ) for obj, df in dfs.items()
105
+ ])
106
+
107
+
108
+ if __name__ == "__main__":
109
+ import os
110
+ import logging
111
+ from .stocks import A_STOCKS_EXCLUDE_ST
112
+ logging.basicConfig(level=logging.INFO)
113
+ msd_host = os.environ.get("MSD_HOST", "http://localhost:50511")
114
+ if not msd_host:
115
+ raise Exception("MSD_HOST is not set")
116
+ dp = MsdSelectorDataProvider(msd_host)
117
+ klines = dp.load_kline(["SH600000"], 200)
118
+ financial = dp.load_financial(["SH600000"], ["f001", "f007"], False, 12)
119
+ #print(klines)
120
+ print(financial)
121
+
122
+ snap = dp.load_snapshot(A_STOCKS_EXCLUDE_ST, ["f137"])
123
+ snap.write_csv("dev/snap.csv")
124
+
@@ -0,0 +1,89 @@
1
+ from uvicorn import logging
2
+ import polars as pl
3
+ import logging
4
+ import time
5
+
6
+ logger = logging.getLogger("selector")
7
+
8
+
9
+ class SelectorDataProvider:
10
+
11
+ date: str|None = None
12
+
13
+ """
14
+ Provide data to Selector
15
+ """
16
+
17
+ def load_kline(self, symbols: list[str], lastN: int = 100) -> dict[str, pl.DataFrame]:
18
+ """
19
+ load `lastN` kline bars for `symbols`
20
+
21
+ Returns:
22
+ A dict of [symbol] -> pl.DataFrame of lastN kline bars
23
+ DataFrame must have columns
24
+ ['ts', 'open', 'high', 'low', 'close', 'amount', 'volume', 'total_shares', 'tradable_shares']
25
+ All prices had been forward-adjusted (split, bonus, etc. already reflected).
26
+ """
27
+ pass
28
+
29
+ def load_financial(self, symbols: list[str], fields: list[str], only_year: bool = True, lastN: int = 12) -> dict[str, pl.DataFrame]:
30
+ """
31
+ load `lastN` financial data for `symbols`
32
+ Financial data is provided on a quarterly basis, with dates of 03-30, 06-30, 09-30, and 12-31.
33
+ When the only_year parameter is set to true, only the annual data (12-31) is selected.
34
+
35
+ Returns:
36
+ A dict of [symbol] -> pl.DataFrame of lastN financial data
37
+ DataFrame columns ['ts', *fields]
38
+ """
39
+ pass
40
+
41
+ def load_snapshot(self, symbols: list[str], fin_fields: list[str] = [], fin_only_year: bool = True) -> pl.DataFrame:
42
+ """
43
+ Take a snapshot for the given symbols at 'snapshot_date', this is more faster than
44
+ load_kline and load_financial separately.
45
+
46
+ User usually use this function to get the latest data for symbols.
47
+
48
+ Args:
49
+ symbols: symbols to snapshot
50
+ fin_fields: fields to snapshot from financial data
51
+ fin_only_year: whether to only snapshot annual financial data
52
+
53
+ Returns:
54
+ A DataFrame of snapshot for the given symbols, with columns combined from
55
+ `load_kline` and `load_financial` except the financial data's 'ts` column will be renamed to 'f_ts'
56
+ """
57
+ pass
58
+
59
+
60
+ class Selector:
61
+
62
+ dp: SelectorDataProvider
63
+
64
+ def __init__(self) :
65
+ pass
66
+
67
+ def execute(self, dp: SelectorDataProvider, init_stocks: list[str]) -> list[str]:
68
+ self.dp = dp
69
+ stocks = init_stocks
70
+ t0 = time.time()
71
+ logger.info(f"start execute selector, {self.__class__.__name__}")
72
+ for i in range(100):
73
+ fn_name = f'step{i:02d}'
74
+ fn = getattr(self, fn_name, None)
75
+ if fn is None:
76
+ break
77
+ t1 = time.time()
78
+ logger.info(f"start {fn_name} with {len(stocks)} stocks")
79
+ if len(fn.__doc__.strip()) > 0:
80
+ logger.info(fn.__doc__.strip())
81
+ stocks = fn(stocks)
82
+ logger.info(f"done {fn_name}, used {time.time() - t1:0.3f}s, remain {len(stocks)}")
83
+ logger.info(f"done execute selector {self.__class__.__name__}, used {time.time() - t0:0.3f}s, remain {len(stocks)}")
84
+ return stocks
85
+
86
+
87
+
88
+
89
+
@@ -0,0 +1,75 @@
1
+ import orjson
2
+ import polars as pl
3
+ import os
4
+
5
+
6
+
7
+ KIND_STOCK = 1
8
+ KIND_FUND = 2
9
+ KIND_INDEX = 3
10
+
11
+ STATUS_NORMAL = 0
12
+ STATUS_ST = 1
13
+ STATUS_STAR_ST = 2
14
+
15
+ def parse_obj_kind(kind: str) -> int:
16
+ match kind.lower():
17
+ case 'stock':
18
+ return KIND_STOCK
19
+ case 'fund':
20
+ return KIND_FUND
21
+ case 'index':
22
+ return KIND_INDEX
23
+ return 0
24
+
25
+
26
+ def parse_obj_status(name: str) -> int:
27
+ if 'ST' in name:
28
+ return 1
29
+ if '*ST' in name:
30
+ return 2
31
+ return 0
32
+
33
+
34
+ def load_all_stocks(path: str) -> pl.DataFrame:
35
+ with open(path) as fp:
36
+ content = orjson.loads(fp.read())
37
+ df = pl.DataFrame(list(content.values()), schema={
38
+ 'symbol': pl.String,
39
+ 'name': pl.String,
40
+ "volume_unit": pl.UInt8,
41
+ "price_decimal": pl.String,
42
+ 'kind': pl.String,
43
+ })
44
+
45
+ df = df.with_columns([
46
+ pl.col('kind').map_elements(parse_obj_kind, return_dtype=pl.Int8),
47
+ pl.col('name').map_elements(parse_obj_status, return_dtype=pl.Int8).alias("status"),
48
+ ])
49
+ return df
50
+
51
+
52
+
53
+ def where_clause(symbols: list[str], col: str = 'obj') -> str:
54
+ sql = ','.join([f'"{s}"' for s in symbols])
55
+ sql = f"{col} IN ({sql})"
56
+ return sql
57
+
58
+ ALL_STOCKS = load_all_stocks(os.environ.get("MSD_STOCKS_FILE", 'etc/stocks.json'))
59
+ A_STOCKS = ALL_STOCKS.filter(pl.col("kind") == KIND_STOCK).get_column("symbol").to_list()
60
+ A_STOCKS_EXCLUDE_ST = ALL_STOCKS.filter(
61
+ (pl.col("kind") == KIND_STOCK) &
62
+ (pl.col("status") == STATUS_NORMAL)
63
+ ).get_column("symbol").to_list()
64
+ FOUNDS = ALL_STOCKS.filter(pl.col("kind") == KIND_FUND).get_column("symbol").to_list()
65
+
66
+
67
+ if __name__ == "__main__":
68
+ print(ALL_STOCKS)
69
+ print(len(A_STOCKS))
70
+ print(len(FOUNDS))
71
+ print(len(A_STOCKS_EXCLUDE_ST))
72
+
73
+ with open("a.txt", "w") as fp:
74
+ for symbol in A_STOCKS_EXCLUDE_ST:
75
+ fp.write(f"{symbol}\n")
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: msd-backtest
3
- Version: 0.1.2
3
+ Version: 0.1.4
4
4
  Summary: MSD backtest framework
5
5
  Requires-Python: >=3.11
6
6
  Description-Content-Type: text/markdown
@@ -10,7 +10,8 @@ Requires-Dist: py-alpha-lib>=0.2; extra == "alpha"
10
10
  Provides-Extra: msd
11
11
  Requires-Dist: pymsd>=0.1; extra == "msd"
12
12
  Requires-Dist: requests>=2.32; extra == "msd"
13
- Requires-Dist: pandas>=3.0; extra == "msd"
13
+ Requires-Dist: polars>=1.0.0; extra == "msd"
14
+ Requires-Dist: xlsxwriter>=3.2.9; extra == "msd"
14
15
  Provides-Extra: server
15
16
  Requires-Dist: fastapi[standard]>=0.135.1; extra == "server"
16
17
  Requires-Dist: msd-backtest[cli]; extra == "server"
@@ -13,6 +13,10 @@ mbt/run/review.py
13
13
  mbt/run/runner.py
14
14
  mbt/run/server.py
15
15
  mbt/select/__init__.py
16
+ mbt/select/cli.py
17
+ mbt/select/msd.py
18
+ mbt/select/selector.py
19
+ mbt/select/stocks.py
16
20
  msd_backtest.egg-info/PKG-INFO
17
21
  msd_backtest.egg-info/SOURCES.txt
18
22
  msd_backtest.egg-info/dependency_links.txt
@@ -23,5 +27,6 @@ tests/test_alpha.py
23
27
  tests/test_load_strategy.py
24
28
  tests/test_report.py
25
29
  tests/test_runner.py
30
+ tests/test_selector.py
26
31
  tests/test_strategy.py
27
32
  tests/test_strategy_loading_complex.py
@@ -12,7 +12,8 @@ tabulate>=0.10.0
12
12
  [msd]
13
13
  pymsd>=0.1
14
14
  requests>=2.32
15
- pandas>=3.0
15
+ polars>=1.0.0
16
+ xlsxwriter>=3.2.9
16
17
 
17
18
  [server]
18
19
  fastapi[standard]>=0.135.1
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "msd-backtest"
3
- version = "0.1.2"
3
+ version = "0.1.4"
4
4
  description = "MSD backtest framework"
5
5
  readme = "README.md"
6
6
  requires-python = ">=3.11"
@@ -12,7 +12,7 @@ excel = ["openpyxl>=3.1.5"]
12
12
 
13
13
  [project.optional-dependencies]
14
14
  alpha = ["py-alpha-lib>=0.2"]
15
- msd = ["pymsd>=0.1", "requests>=2.32", "pandas>=3.0"]
15
+ msd = ["pymsd>=0.1", "requests>=2.32", "polars>=1.0.0", "xlsxwriter>=3.2.9"]
16
16
  server = ["fastapi[standard]>=0.135.1", "msd-backtest[cli]"]
17
17
  cli = [
18
18
  "orjson>=3.11.9",
@@ -33,6 +33,6 @@ addopts = "--benchmark-skip"
33
33
  [tool.setuptools.packages.find]
34
34
  include = ["mbt*"]
35
35
 
36
- # [tool.uv.sources]
37
- # py-alpha-lib = { path = "../alpha-lib-rs/py-binding", editable = true }
38
- # pymsd = { path = "../msd-rs2/bindings/python", editable = true }
36
+ [tool.uv.sources]
37
+ py-alpha-lib = { path = "../alpha-lib-rs/py-binding", editable = true }
38
+ pymsd = { path = "../msd-rs2/bindings/python", editable = true }
@@ -22,7 +22,7 @@ def load_data_msd(
22
22
  msd_host: str, symbols: list[str], start: str, end: str
23
23
  ) -> tuple[DataProvider, list[str]]:
24
24
 
25
- client = pymsd.create_msd_pandas(msd_host)
25
+ client = pymsd.create_msd_polars(msd_host)
26
26
 
27
27
  dfs = client.load(
28
28
  objs=symbols,
@@ -0,0 +1,140 @@
1
+ from datetime import datetime, timedelta
2
+ import polars as pl
3
+ import numpy as np
4
+ import mbt
5
+ from mbt.select import SelectorDataProvider
6
+ from strategy.growth_large_cap_ma15_selector import GrowthLargeCapMA15Selector
7
+
8
+ class MockSelectorDataProvider(SelectorDataProvider):
9
+ def __init__(self, kline_data, financial_data):
10
+ super().__init__()
11
+ self.kline_data = kline_data
12
+ self.financial_data = financial_data
13
+
14
+ def load_kline(self, symbols: list[str], lastN: int = 100) -> dict[str, pl.DataFrame]:
15
+ return {sym: self.kline_data[sym] for sym in symbols if sym in self.kline_data}
16
+
17
+ def load_financial(self, symbols: list[str], fields: list[str], only_year: bool = True, lastN: int = 12) -> dict[str, pl.DataFrame]:
18
+ return {sym: self.financial_data[sym] for sym in symbols if sym in self.financial_data}
19
+
20
+ def load_snapshot(self, symbols: list[str], fin_fields: list[str] = [], fin_only_year: bool = True, snapshot_date: str = '') -> pl.DataFrame:
21
+ rows = []
22
+ for sym in symbols:
23
+ if sym not in self.kline_data:
24
+ continue
25
+ k_df = self.kline_data[sym]
26
+ fin_df = self.financial_data.get(sym)
27
+
28
+ close = k_df["close"][-1]
29
+ total_shares = k_df["total_shares"][-1]
30
+ row_dict = {
31
+ "obj": sym,
32
+ "close": close,
33
+ "total_shares": total_shares,
34
+ }
35
+ if fin_df is not None and len(fin_df) > 0:
36
+ fin_last = fin_df.tail(1)
37
+ for f in fin_fields:
38
+ if f in fin_last.columns:
39
+ row_dict[f] = fin_last[f][0]
40
+ rows.append(row_dict)
41
+
42
+ if not rows:
43
+ schema = {
44
+ "obj": pl.String,
45
+ "close": pl.Float64,
46
+ "total_shares": pl.Float64,
47
+ }
48
+ for f in fin_fields:
49
+ schema[f] = pl.Float64
50
+ return pl.DataFrame([], schema=schema)
51
+
52
+ return pl.DataFrame(rows)
53
+
54
+ def test_growth_large_cap_ma15_selector():
55
+ # Define dates
56
+ dates_fin = [datetime(2023, 12, 31), datetime(2024, 12, 31), datetime(2025, 12, 31)]
57
+ end_date = datetime(2026, 3, 6)
58
+ dates_k = [end_date - timedelta(days=i) for i in range(19, -1, -1)]
59
+
60
+ # 1. SZ000001: All filters pass
61
+ # Net Profit: 100 -> 120 -> 150 (growing & positive)
62
+ df_fin_1 = pl.DataFrame({"ts": dates_fin, "f137": [100.0, 120.0, 150.0]})
63
+ # Close price drops below MA15 on the last day
64
+ close_1 = np.ones(20) * 12.0
65
+ close_1[-2] = 13.0
66
+ close_1[-1] = 8.0
67
+ df_k_1 = pl.DataFrame({
68
+ "ts": dates_k,
69
+ "open": close_1,
70
+ "high": close_1,
71
+ "low": close_1,
72
+ "close": close_1,
73
+ "volume": np.ones(20) * 1000,
74
+ "amount": np.ones(20) * 12000,
75
+ "total_shares": np.ones(20) * 1000,
76
+ "tradable_shares": np.ones(20) * 1000,
77
+ })
78
+
79
+ # 2. SZ000002: Fails financial growth (150 -> 120 -> 100)
80
+ df_fin_2 = pl.DataFrame({"ts": dates_fin, "f137": [150.0, 120.0, 100.0]})
81
+ df_k_2 = df_k_1.clone()
82
+
83
+ # 3. SZ000003: Fails large cap (market cap is small)
84
+ df_fin_3 = df_fin_1.clone()
85
+ df_k_3 = df_k_1.clone()
86
+ df_k_3 = df_k_3.with_columns(pl.lit(10.0).alias("total_shares"))
87
+
88
+ # 4. SZ000004: Fails MA15 breakdown (stays at 12.0, close is never below MA15)
89
+ df_fin_4 = df_fin_1.clone()
90
+ close_4 = np.ones(20) * 12.0
91
+ df_k_4 = pl.DataFrame({
92
+ "ts": dates_k,
93
+ "open": close_4,
94
+ "high": close_4,
95
+ "low": close_4,
96
+ "close": close_4,
97
+ "volume": np.ones(20) * 1000,
98
+ "amount": np.ones(20) * 12000,
99
+ "total_shares": np.ones(20) * 1000,
100
+ "tradable_shares": np.ones(20) * 1000,
101
+ })
102
+
103
+ # 5. SZ000005: Fails PE filter (PE is too high, cap=40000, net_profit=150, PE=266.67)
104
+ df_fin_5 = df_fin_1.clone()
105
+ df_k_5 = df_k_1.clone()
106
+ df_k_5 = df_k_5.with_columns(pl.lit(5000.0).alias("total_shares"))
107
+
108
+ financial_data = {
109
+ "SZ000001": df_fin_1,
110
+ "SZ000002": df_fin_2,
111
+ "SZ000003": df_fin_3,
112
+ "SZ000004": df_fin_4,
113
+ "SZ000005": df_fin_5,
114
+ }
115
+
116
+ kline_data = {
117
+ "SZ000001": df_k_1,
118
+ "SZ000002": df_k_2,
119
+ "SZ000003": df_k_3,
120
+ "SZ000004": df_k_4,
121
+ "SZ000005": df_k_5,
122
+ }
123
+
124
+ dp = MockSelectorDataProvider(kline_data, financial_data)
125
+
126
+ # Use top 50% percentile
127
+ selector = GrowthLargeCapMA15Selector(large_cap_percentile=0.5, break_days=3)
128
+
129
+ selected = selector.execute(dp, ["SZ000001", "SZ000002", "SZ000003", "SZ000004", "SZ000005"])
130
+
131
+ assert selected == ["SZ000001"]
132
+
133
+ # Check factors
134
+ factors = selector.factors
135
+ assert len(factors) == 1
136
+ assert factors[0, "symbol"] == "SZ000001"
137
+ assert factors[0, "net_profit_y3"] == 150.0
138
+ assert factors[0, "market_cap"] == 8000.0
139
+ assert factors[0, "pe"] == 8000.0 / 150.0
140
+ assert factors[0, "close"] == 8.0
File without changes
File without changes
File without changes