Update show_tree.py

wukan1986 · Mar 22, 2024 · bdeb1dc · bdeb1dc
1 parent 07088ab
commit bdeb1dc
Showing 1 changed file with 16 additions and 11 deletions.
diff --git a/examples/show_tree.py b/examples/show_tree.py
@@ -2,29 +2,34 @@
 from sympy import numbered_symbols
 
 from examples.sympy_define import *
+from expr_codegen.codes import sources_to_exprs
 from expr_codegen.dag import zero_outdegree
-from expr_codegen.expr import string_to_exprs
 from expr_codegen.model import create_dag_exprs, init_dag_exprs, draw_expr_tree, merge_nodes_1, merge_nodes_2
 from expr_codegen.tool import ExprTool
 
+from polars_ta.prefix.tdx import *  # noqa
+from polars_ta.prefix.ta import *  # noqa
+from polars_ta.prefix.wq import *  # noqa
+from polars_ta.prefix.cdl import *  # noqa
+
 RETURNS, VWAP, = symbols('RETURNS, VWAP, ', cls=Symbol)
 
-exprs_src = {
-    "alpha_001": (
-            cs_rank(ts_arg_max(signed_power(if_else((RETURNS < 0), ts_std_dev(RETURNS, 20), CLOSE), 2.), 5)) - 0.5),
-    "alpha_002": (-1 * ts_corr(cs_rank(ts_delta(log(VOLUME), 2)), cs_rank(((CLOSE - OPEN) / OPEN)), 6)),
-    "alpha_003": (-1 * ts_corr(cs_rank(OPEN), cs_rank(VOLUME), 10)),
-    "alpha_004": (-1 * ts_rank(cs_rank(LOW), 9)),
-    "alpha_005": (cs_rank((OPEN - (ts_sum(VWAP, 10) / 10))) * (-1 * abs_(cs_rank((CLOSE - VWAP))))),
-    "alpha_006": -1 * ts_corr(OPEN, VOLUME, 10),
-}
+# exprs_src = {
+#     "alpha_001": (
+#             cs_rank(ts_arg_max(signed_power(if_else((RETURNS < 0), ts_std_dev(RETURNS, 20), CLOSE), 2.), 5)) - 0.5),
+#     "alpha_002": (-1 * ts_corr(cs_rank(ts_delta(log(VOLUME), 2)), cs_rank(((CLOSE - OPEN) / OPEN)), 6)),
+#     "alpha_003": (-1 * ts_corr(cs_rank(OPEN), cs_rank(VOLUME), 10)),
+#     "alpha_004": (-1 * ts_rank(cs_rank(LOW), 9)),
+#     "alpha_005": (cs_rank((OPEN - (ts_sum(VWAP, 10) / 10))) * (-1 * abs_(cs_rank((CLOSE - VWAP))))),
+#     "alpha_006": -1 * ts_corr(OPEN, VOLUME, 10),
+# }
 
 # 表达式设置
 exprs_src = """
 alpha_101=(CLOSE - OPEN) / ((HIGH - LOW) + 0.001)
 alpha_201=alpha_101+CLOSE # 中间变量示例
 """
-exprs_src = string_to_exprs(exprs_src, globals().copy())
+raw, exprs_src = sources_to_exprs(globals().copy(), exprs_src)
 
 tool = ExprTool()
 # 子表达式在前，原表式在最后