[ { "path": ".github/workflows/python-publish.yml", "content": "# This workflow will upload a Python Package using Twine when a release is created\n# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries\n\n# This workflow uses actions that are not certified by GitHub.\n# They are provided by a third-party and are governed by\n# separate terms of service, privacy policy, and support\n# documentation.\n\nname: Upload Python Package\n\non:\n release:\n types: [published]\n\npermissions:\n contents: read\n\njobs:\n deploy:\n\n runs-on: ubuntu-latest\n\n steps:\n - uses: actions/checkout@v3\n - name: Set up Python\n uses: actions/setup-python@v3\n with:\n python-version: '3.x'\n - name: Install dependencies\n run: |\n python -m pip install --upgrade pip\n pip install build\n - name: Build package\n run: python -m build\n - name: Publish package\n uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29\n with:\n user: __token__\n password: ${{ secrets.PYPI_API_TOKEN }}\n" }, { "path": ".gitignore", "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\ncover/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\n.pybuilder/\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n# For a library or package, you might want to ignore these files since the code is\n# intended to run in multiple environments; otherwise, check them in:\n# .python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# poetry\n# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.\n# This is especially recommended for binary packages to ensure reproducibility, and is more\n# commonly ignored for libraries.\n# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control\n#poetry.lock\n\n# pdm\n# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.\n#pdm.lock\n# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it\n# in version control.\n# https://pdm.fming.dev/#use-with-ide\n.pdm.toml\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# Cython debug symbols\ncython_debug/\n\n# PyCharm\n# JetBrains specific template is maintained in a separate JetBrains.gitignore that can\n# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore\n# and can be added to the global gitignore or merged into this file. For a more nuclear\n# option (not recommended) you can uncomment the following to ignore the entire idea folder.\n#.idea/\n" }, { "path": ".readthedocs.yaml", "content": "# Read the Docs configuration file for MkDocs projects\n# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details\n\n# Required\nversion: 2\n\n# Set the version of Python and other tools you might need\nbuild:\n os: ubuntu-22.04\n tools:\n python: \"3.11\"\n\nmkdocs:\n configuration: mkdocs.yml\n\n# Optionally declare the Python requirements required to build your docs\npython:\n install:\n - requirements: requirements-docs.txt" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2023 wukan\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "# polars_ta\n\nTechnical Indicator Operators Rewritten in `polars`.\n\nWe provide wrappers for some functions (like `TA-Lib`) that are not `pl.Expr` alike.\n\n## How to Install\n\n### Using `pip`\n\n```commandline\npip install -i https://pypi.org/simple --upgrade polars_ta\npip install -i https://pypi.tuna.tsinghua.edu.cn/simple --upgrade polars_ta # Mirror in China\n```\n\n### Build from Source\n\n```commandline\ngit clone --depth=1 https://github.com/wukan1986/polars_ta.git\ncd polars_ta\npython -m build\ncd dist\npip install polars_ta-0.1.2-py3-none-any.whl\n```\n\n### How to Install TA-Lib\n\nNon-official `TA-Lib` wheels can be downloaded from `https://github.com/cgohlke/talib-build/releases`\n\n## Usage\n\nSee `examples` folder.\n\n```python\n# We need to modify the function name by prefixing `ts_` before using them in `expr_coodegen`\nfrom polars_ta.prefix.tdx import *\n# Import functions from `wq`\nfrom polars_ta.prefix.wq import *\n\n# Example\ndf = df.with_columns([\n # Load from `wq`\n *[ts_returns(CLOSE, i).alias(f'ROCP_{i:03d}') for i in (1, 3, 5, 10, 20, 60, 120)],\n *[ts_mean(CLOSE, i).alias(f'SMA_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_std_dev(CLOSE, i).alias(f'STD_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_max(HIGH, i).alias(f'HHV_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_min(LOW, i).alias(f'LLV_{i:03d}') for i in (5, 10, 20, 60, 120)],\n\n # Load from `tdx`\n *[ts_RSI(CLOSE, i).alias(f'RSI_{i:03d}') for i in (6, 12, 24)],\n])\n```\n\nWhen both `min_samples` and `MIN_SAMPLES` are set, `min_samples` takes precedence. default value is `None`.\n\n```python\nimport polars_ta\n\n# Global settings. Priority Low\npolars_ta.MIN_SAMPLES = 1\n\n# High priority\nts_mean(CLOSE, 10, min_samples=1)\n```\n\n## How We Designed This\n\n1. We use `Expr` instead of `Series` to avoid using `Series` in the calculation. Functions are no longer methods of class.\n2. Use `wq` first. It mimics `WorldQuant Alpha` and strives to be consistent with them.\n3. Use `ta` otherwise. It is a `polars`-style version of `TA-Lib`. It tries to reuse functions from `wq`.\n4. Use `tdx` last. It also tries to import functions from `wq` and `ta`.\n5. We keep the same signature and parameters as the original `TA-Lib` in `talib`.\n6. If there is a naming conflict, we suggest calling `wq`, `ta`, `tdx`, `talib` in order. The higher the priority, the closer the implementation is to `Expr`.\n\n## Comparison of Our Indicators and Others\n\nSee [compare](compare.md)\n\n## Handling Null/NaN Values\n\nSee [nan_to_null](nan_to_null.md)\n\n## Debugging\n\n```commandline\ngit clone --depth=1 https://github.com/wukan1986/polars_ta.git\ncd polars_ta\npip install -e .\n```\n\nNotice:\nIf you have added some functions in `ta` or `tdx`, please run `prefix_ta.py` or `prefix_tdx.py` inside the `tools` folder to generate the corrected Python script (with the prefix added).\nThis is required to use in `expr_codegen`.\n\n## Reference\n\n- https://github.com/pola-rs/polars\n- https://github.com/TA-Lib/ta-lib\n- https://github.com/twopirllc/pandas-ta\n- https://github.com/bukosabino/ta\n- https://github.com/peerchemist/finta\n- https://github.com/wukan1986/ta_cn\n- https://support.worldquantbrain.com/hc/en-us/community/posts/20278408956439-从价量看技术指标总结-Technical-Indicator-\n- https://platform.worldquantbrain.com/learn/operators/operators\n\n# polars_ta\n\n基于`polars`的算子库。实现量化投研中常用的技术指标、数据处理等函数。对于不易翻译成`Expr`的库（如：`TA-Lib`）也提供了函数式调用的封装\n\n## 安装\n\n### 在线安装\n\n```commandline\npip install -i https://pypi.org/simple --upgrade polars_ta # 官方源\npip install -i https://pypi.tuna.tsinghua.edu.cn/simple --upgrade polars_ta # 国内镜像源\n```\n\n### 源码安装\n\n```commandline\ngit clone --depth=1 https://github.com/wukan1986/polars_ta.git\ncd polars_ta\npython -m build\ncd dist\npip install polars_ta-0.1.2-py3-none-any.whl\n```\n\n### TA-Lib安装\n\nWindows用户不会安装可从`https://github.com/cgohlke/talib-build/releases` 下载对应版本whl文件\n\n## 使用方法\n\n参考`examples`目录即可，例如：\n\n```python\n# 如果需要在`expr_codegen`中使用，需要有`ts_`等前权，这里导入提供了前缀\nfrom polars_ta.prefix.tdx import *\n# 导入wq公式\nfrom polars_ta.prefix.wq import *\n\n# 演示生成大量指标\ndf = df.with_columns([\n # 从wq中导入指标\n *[ts_returns(CLOSE, i).alias(f'ROCP_{i:03d}') for i in (1, 3, 5, 10, 20, 60, 120)],\n *[ts_mean(CLOSE, i).alias(f'SMA_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_std_dev(CLOSE, i).alias(f'STD_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_max(HIGH, i).alias(f'HHV_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_min(LOW, i).alias(f'LLV_{i:03d}') for i in (5, 10, 20, 60, 120)],\n\n # 从tdx中导入指标\n *[ts_RSI(CLOSE, i).alias(f'RSI_{i:03d}') for i in (6, 12, 24)],\n])\n```\n\n当`min_samples`和`MIN_SAMPLES`都设置时，以`min_samples`为准，默认值为`None`\n\n```python\nimport polars_ta\n\n# 全局设置。优先级低\npolars_ta.MIN_SAMPLES = 1\n\n# 指定函数。优先级高\nts_mean(CLOSE, 10, min_samples=1)\n\n```\n\n## 设计原则\n\n1. 调用方法由`成员函数`换成`独立函数`。输入输出使用`Expr`，避免使用`Series`\n2. 优先实现`wq`公式，它仿`WorldQuant Alpha`公式，与官网尽量保持一致。如果部分功能实现在此更合适将放在此处\n3. 其次实现`ta`公式，它相当于`TA-Lib`的`polars`风格的版本。优先从`wq`中导入更名\n4. 最后实现`tdx`公式，它也是优先从`wq`和`ta`中导入\n5. `talib`的函数名与参数与原版`TA-Lib`完全一致\n6. 如果出现了命名冲突，建议调用优先级为`wq`、`ta`、`tdx`、`talib`。因为优先级越高，实现方案越接近于`Expr`\n\n## 指标区别\n\n请参考[compare](compare.md)\n\n## 空值处理\n\n请参考[nan_to_null](nan_to_null.md)\n\n## 开发调试\n\n```commandline\ngit clone --depth=1 https://github.com/wukan1986/polars_ta.git\ncd polars_ta\npip install -e .\n```\n\n注意：如果你在`ta`或`tdx`中添加了新的函数，请再运行`tools`下的`prefix_ta.py`或`prefix_tdx.py`，用于生成对应的前缀文件。前缀文件方便在`expr_codegen`中使用\n\n## 文档生成\n\n```commandline\npip install -r requirements-docs.txt\nmkdocs build\n```\n\n文档生成在`site`目录下，其中的`llms-full.txt`可以作为大语言模型的知识库导入。\n\n也可以通过以下链接导入：\nhttps://polars-ta.readthedocs.io/en/latest/llms-full.txt\n\n## 提示词\n由于`llms-full.txt`信息不适合做提示词，所以`tools/prompt.py`提供了生成更简洁算子清单的功能。\n\n用户也可以直接使用`prompt.txt`(欢迎提示词工程专家帮忙改进，做的更准确)\n\n## 参考\n\n- https://github.com/pola-rs/polars\n- https://github.com/TA-Lib/ta-lib\n- https://github.com/twopirllc/pandas-ta\n- https://github.com/bukosabino/ta\n- https://github.com/peerchemist/finta\n- https://github.com/wukan1986/ta_cn\n- https://support.worldquantbrain.com/hc/en-us/community/posts/20278408956439-从价量看技术指标总结-Technical-Indicator-\n\n" }, { "path": "compare.md", "content": "# Differences in Indicators\n\n## 1. EMA alike indicators\n\n### 1.1 EMA\n\n1. EMA(CLOSE, 10), `talib.set_compatibility(0)` is the default, and is equivalent to`EXPMEMA`\n - The first not nan value starts as `talib.SMA(CLOSE, 10)`\n2. EMA(CLOSE, 10), `talib.set_compatibility(1)`\n - The first not nan value starts as `CLOSE`\n\nSince the logic of `EMA` calculation has changed in `TA-Lib`'s `compatibility mode 0`, it is complex and inefficient to implement it using expressions.\n\nWe only implement `compatibility mode 1`.\nIt just so happens that the Chinese stock software only implements `compatibility mode 1`.\nYou can compare the full data with `TA-Lib` for unit testing.\n\nIndicators affected by `EMA` include `MACD, DEMA, TEMA, TRIX, T3`, etc.\n\n### 1.2 Chinese Version of SMA(X, N, M)\n\nIn essence, it is `RMA compatibility mode 0`, that is,\nthe first valid value is the moving average. And then the difference is `alpha`\n\n1. `SMA(X, N, M) = X.ema_mean(alpha=M/N)`\n2. `RMA(X, N) = X.ema_mean(alpha=1/N) = SMA(X, N, 1)`\n3. `EMA(x, N) = X.ema_mean(alpha=2/(N+1)) = X.ema_mean(span=N)`\n\nRefer to: [ewm_mean](https://pola-rs.github.io/polars/py-polars/html/reference/expressions/api/polars.Expr.ewm_mean.html#polars.Expr.ewm_mean)\n\nIn this case, we use `RMA compatibility mode 1`.\nThere exists error in switching between the two modes.\nSo please provide enough length of data.\nThe data in the later part can be used for unit testing.\n\nIndicators affected by `SMA` include `ATR, RSI`, etc.\n\n\n### 1.3 Moving Sum\n\nSome indicators including `ADX` require the first value as a `SUM` rather than `SMA`. We will implement them later by `ema_mean(alpha=1/N)`\n\n\n## 2. MAX/MIN\n\n1. In package `wq`, `max_/min_` are cross-sectional operators, `ts_max/ts_min` are time series indicators\n2. In `talib`, `MAX/MIN` are time series indicators, without cross-sectional operators.\n3. In `ta`, to mimic `talib`, `MAX/MIN` are time series indicators, without cross-sectional operators.\n4. In `tdx`, `MAX/MIN` are cross-sectional operators, `HHV/LLV` are time series indicators.\n\n\n# 指标区别\n\n## 1. EMA系列\n\n### 1.1 EMA指标\n\n1. EMA(CLOSE, 10)，`talib.set_compatibility(0)`，此为默认设置，等价于`EXPMEMA`\n - 第一个有效值为`talib.SMA(CLOSE, 10)`\n2. EMA(CLOSE, 10)，`talib.set_compatibility(1)`\n - 第一个有效值为`CLOSE`\n\n由于`TA-Lib`的`兼容模式0`在`EMA`计算时逻辑发生了变动，用表达式实现起来复杂、计算效率低。\n所以本库只实现`兼容模式1`。正好国内股票软件其实也只实现了`兼容模式1`。可以全量数据与`TA-Lib`进行单元测试比较\n\n因`EMA`受影响的指标有`MACD, DEMA, TEMA, TRIX, T3`等。\n\n### 1.2 中国版SMA(X, N, M)\n\n本质上是国外的`RMA 兼容模式0`，即第一个有效值为移动平均，然后就是`alpha`区别\n\n1. `SMA(X, N, M) = X.ema_mean(alpha=M/N)`\n2. `RMA(X, N) = X.ema_mean(alpha=1/N) = SMA(X, N, 1)`\n3. `EMA(x, N) = X.ema_mean(alpha=2/(N+1)) = X.ema_mean(span=N)`\n\n换算关系可参考: [ewm_mean](https://pola-rs.github.io/polars/py-polars/html/reference/expressions/api/polars.Expr.ewm_mean.html#polars.Expr.ewm_mean)\n\n遇到这种情况，本项目还是用`RMA 兼容模式1`来代替，数据误差由大到小，所以请预先提供一定长度的数据。后面一段的数据可以单元测试\n\n受影响的的指标有`ATR, RSI`等\n\n### 1.3 移动求和\n\n`ADX`等一类的指标第一个有效值算法为`SUM`，而不是`SMA`，之后使用`ema_mean(alpha=1/N)`。此类暂不实现\n\n## 2. MAX/MIN等指标\n\n1. 在`wq`中，`max_/min_`横向算子，`ts_max/ts_min`时序指标\n2. 在`talib`中, `MAX/MIN`时序指标，没有横向算子\n3. 在`ta`中，由于要模仿`talib`，所以有`MAX/MIN`时序指标，也没有横向算子\n4. 在`tdx`中，`MAX/MIN`横向算子，`HHV/LLV`时序指标\n\n" }, { "path": "deprecated/demo_ta1.py", "content": "\"\"\"\nThisis how polars implements calling third-party packages\n以下是polars提供的实现调用第三方库的方案\nexpr.ta.func\n\n\"\"\"\nimport polars as pl\n\nfrom polars_ta.utils.helper import TaLibHelper # noqa\n\ndf = pl.DataFrame(\n {\n \"A\": [5, None, 3, 2, 1],\n \"B\": [5, 4, None, 2, 1],\n \"C\": [5, 4, 3, 2, 1],\n }\n)\n\ndf = df.with_columns([\n # single input single ouput, no need to handle null/nan values\n # 一输入一输出，不需处理空值\n pl.col('A').ta.COS().alias('COS'),\n # single input, multi output\n # 一输入多输出\n pl.col('A').ta.BBANDS(timeperiod=2, skip_nan=True, schema=['upperband', 'middleband', 'lowerband']).alias('BBANDS'),\n # multi input, single output\n # 多输入一输出\n pl.struct(['A', 'B', 'C']).ta.ATR(timeperiod=2, skip_nan=True).alias('ATR'),\n # multi input, multi output\n # 多输入多输出\n pl.struct(['A', 'B']).ta.AROON(timeperiod=2, skip_nan=True, schema=('aroondown', 'aroonup'), schema_format='XX_{}_YY').alias('AROON'),\n # multi input, single output\n # 多输入一输出\n pl.struct(['A', 'B']).ta.AROON(timeperiod=2, skip_nan=True, output_idx=1).alias('aroonup1'),\n # call third-party packages\n # 调用另一库\n pl.col('A').bn.move_rank(window=2, skip_nan=False).alias('move_rank'),\n])\n\nprint(df)\n\ndf = df.unnest('BBANDS', 'AROON')\nprint(df)\n" }, { "path": "deprecated/demo_ta2.py", "content": "\"\"\"\nThis is how we wrap upon polars's code\n以下是polars提供的方案基础上封装的方案\nfunc(expr)\n\n\"\"\"\nimport polars as pl\n\nfrom polars_ta.utils.wrapper import init\n\ndf = pl.DataFrame(\n {\n \"A\": [5, None, 3, 2, 1],\n \"B\": [5, 4, None, 2, 1],\n \"C\": [5, 4, 3, 2, 1],\n }\n)\n\n# register to global variables. You may directly use them.\n# IDE might give warnings about not defined.\n# 已经注册到全局，可直接使用。但IDE中没有智能提示\ninit(to_globals=True, name_format='{}')\n\ndf = df.with_columns([\n # single input single ouput, no need to handle null/nan values\n # 一输入一输出，不需处理空值\n COS(pl.col('A')).alias('COS'),\n\n # multi input, single output\n # 多输入一输出\n ATR(pl.struct(['A', 'B', 'C']), timeperiod=2).alias('ATR1'),\n ATR(pl.col('A'), pl.col('B'), pl.col('C'), 2, skip_nan=True).alias('ATR2'),\n\n # single input, multi output, add prefix to the column names passing to `BBANDS`\n # 一输入多输出，可通过prefix为多输出添加前缀\n BBANDS(pl.col('A'), timeperiod=2, skip_nan=True, schema_format='bbands_{}').alias('BBANDS'),\n\n # multi input multi output, set column name directly by `schema`\n # 多输入多输出。可通过schema直接添加\n AROON(pl.struct(['A', 'B']), timeperiod=2, skip_nan=True, schema=('aroondown', 'aroonup')).alias('AROON'),\n\n])\n\nprint(df)\n\ndf = df.unnest('BBANDS', 'AROON')\nprint(df)\n\n# another way of calling init\n# 另一种调用方法\nt = init(to_globals=False, name_format='ts_{}')\ndf = df.with_columns([\n # single input single ouput, no need to handle null/nan values\n # 一输入一输出，不需处理空值\n t.ts_COS(pl.col('A')).alias('ts_COS'),\n])\n\nprint(df)\n" }, { "path": "deprecated/helper.py", "content": "\"\"\"\nWe follow the spirit of\n\nhttps://github.com/pola-rs/polars/issues/9261\n\nand adjusted it to fit our needs. The usage is as follows\n\nexpr.ta.func(..., skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False)\n\nskip_nan: bool\n reduces speed\noutput_idx: int\n select a single column when outputing multiple columns\nschema: list or tuple\n assign column names for multiple output columns\nschema_format: str\n assign data types for multiple output columns\nnan_to_null: bool\n replace all nan to null when outputing\n\n\n\n以下是在github issues中 @cmdlineluser 提供的 注册命名空间 方案\n\nhttps://github.com/pola-rs/polars/issues/9261\n\n本人做了一定的调整。使用方法如下\n\nexpr.ta.func(..., skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False)\n\nskip_nan: bool\n 是否跳过空值。可以处理停牌无数据的问题，但会降低运行速度\noutput_idx: int\n 多列输出时，选择只输出其中一列\nschema: list or tuple\n 返回为多列时，会组装成struct，可以提前设置子列的名字\nschema_format: str\n 为子列名指定格式\nnan_to_null: bool\n 返回值是否改成null\n\n其它参数按**位置参数**和**命名参数**输入皆可\n\"\"\"\nimport numpy as np\nfrom polars import Expr, Float64, DataFrame, api\nfrom polars import Series, Struct\n\n\ndef func_wrap_mn(func, cols,\n *args,\n skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False,\n **kwargs):\n \"\"\"multiple input multiple output, compatible with single input single output\n 多输入多输出，兼容一输入一输出\n\n Parameters\n ----------\n func\n cols\n args\n skip_nan\n output_idx\n schema\n schema_format\n nan_to_null\n kwargs\n\n Returns\n -------\n Series\n\n \"\"\"\n if cols.dtype.base_type() == Struct:\n # struct(['A', 'B']).ta.AROON\n _cols = [cols.struct[field] for field in cols.dtype.to_schema()]\n else:\n # col('A').ta.BBANDS\n _cols = [cols]\n\n if skip_nan:\n _cols = [c.cast(Float64).to_numpy() for c in _cols]\n _cols = np.vstack(_cols)\n\n # move nan to head\n idx1 = (~np.isnan(_cols).any(axis=0)).argsort(kind='stable')\n # restore nan\n idx2 = idx1.argsort(kind='stable')\n\n _cols = [_cols[i, idx1] for i in range(len(_cols))]\n result = func(*_cols, *args, **kwargs)\n\n if isinstance(result, tuple):\n result = tuple([_[idx2] for _ in result])\n else:\n result = result[idx2]\n else:\n result = func(*_cols, *args, **kwargs)\n\n if isinstance(result, tuple):\n if output_idx is None:\n # struct(['A').ta.BBANDS\n # you need alias outside, finally unnest\n if schema is None:\n schema = [f'column_{i}' for i in range(len(result))]\n\n schema = [schema_format.format(name) for name in schema]\n # nan_to_null is not effective for struct\n # nan_to_null 对struct中的nan无效\n return DataFrame(result, schema=schema, nan_to_null=nan_to_null).to_struct('')\n # output only one column\n if 0 <= output_idx < len(result):\n return Series(result[output_idx], nan_to_null=nan_to_null)\n elif not isinstance(result, Series):\n # col('A').bn.move_rank\n return Series(result, nan_to_null=nan_to_null)\n else:\n # col('A').ta.COS\n if nan_to_null:\n return result.fill_nan(None)\n else:\n return result\n\n\ndef func_wrap_11(func, cols,\n *args,\n skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False,\n **kwargs):\n \"\"\"single input single output is faster than multiple input multiple output\n 一输入，一输出。处理速度能更快一些\"\"\"\n _cols = cols\n\n if skip_nan:\n _cols = _cols.cast(Float64).to_numpy()\n\n # move nan to head\n idx1 = (~np.isnan(_cols)).argsort(kind='stable')\n # restore nan\n idx2 = idx1.argsort(kind='stable')\n\n _cols = _cols[idx1]\n result = func(_cols, *args, **kwargs)\n\n result = result[idx2]\n else:\n result = func(_cols, *args, **kwargs)\n\n if isinstance(result, Series):\n if nan_to_null:\n return result.fill_nan(None)\n else:\n return result\n else:\n return Series(result, nan_to_null=nan_to_null)\n\n\nclass FuncHelper:\n def __init__(self, expr: Expr, lib=None, wrap=None) -> None:\n \"\"\"initialization\n\n Parameters\n ----------\n expr\n lib:\n third-party packages required\n wrap:\n wrapper for third-party packages\n\n\n 初始化\n\n Parameters\n ----------\n expr\n lib:\n 需要调用的第三方库\n wrap:\n 对第三方库的封装方法\n\n \"\"\"\n object.__setattr__(self, '_expr', expr)\n object.__setattr__(self, '_lib', lib)\n object.__setattr__(self, '_wrap', wrap)\n\n def __getattribute__(self, name: str):\n _expr: Expr = object.__getattribute__(self, '_expr')\n _lib = object.__getattribute__(self, '_lib')\n _wrap = object.__getattribute__(self, '_wrap')\n _func = getattr(_lib, name)\n return (\n lambda *args, skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False, **kwargs:\n _expr.map_batches(\n lambda cols: _wrap(_func, cols,\n *args,\n skip_nan=skip_nan, output_idx=output_idx, schema=schema, schema_format=schema_format, nan_to_null=nan_to_null,\n **kwargs)\n )\n )\n\n\n@api.register_expr_namespace('ta')\nclass TaLibHelper(FuncHelper):\n def __init__(self, expr: Expr) -> None:\n import talib as ta\n super().__init__(expr, ta, func_wrap_mn)\n\n\n@api.register_expr_namespace('bn')\nclass BottleneckHelper(FuncHelper):\n def __init__(self, expr: Expr) -> None:\n import bottleneck as bn\n super().__init__(expr, bn, func_wrap_11)\n" }, { "path": "deprecated/pandas_.py", "content": "from functools import lru_cache\nfrom typing import Tuple\n\nimport numpy as np\nfrom numba import jit\nfrom pandas._libs.window.aggregations import roll_kurt as _roll_kurt\nfrom pandas._libs.window.aggregations import roll_rank as _roll_rank\nfrom polars import Series\n\n\"\"\"\nWhen converting float32 to float64 before computing. Either use\nx.cast(Float64).to_numpy()\nor\nx.to_numpy().astype(float)\n\nThe second one is faster\n\n在计算前需要将float32转成float64，有以下两种方法\nx.cast(Float64).to_numpy()\nx.to_numpy().astype(float)\n\n第二种方法更快\n\"\"\"\n\n\n@lru_cache\n@jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef get_window_bounds(\n num_values: int = 0,\n window_size: int = 10,\n) -> Tuple[np.ndarray, np.ndarray]:\n end = np.arange(1, num_values + 1, dtype=np.int64)\n start = end - window_size\n start = np.clip(start, 0, num_values)\n return start, end\n\n\ndef roll_rank(x: Series, d: int, minp: int, pct: bool = True, method: str = 'average', ascending: bool = True):\n start, end = get_window_bounds(len(x), d)\n \"\"\"\n https://github.com/pandas-dev/pandas/blob/main/pandas/_libs/window/aggregations.pyx#L1281\n\n def roll_rank(const float64_t[:] values, ndarray[int64_t] start,\n ndarray[int64_t] end, int64_t minp, bint percentile,\n str method, bint ascending) -> np.ndarray:\n\n O(N log(window)) implementation using skip list\n \"\"\"\n ret = _roll_rank(x.to_numpy().astype(float), start, end, minp, pct, method, ascending)\n return Series(ret, nan_to_null=True)\n\n\ndef roll_kurt(x, d: int, minp: int):\n start, end = get_window_bounds(len(x), d)\n \"\"\"\n https://github.com/pandas-dev/pandas/blob/main/pandas/_libs/window/aggregations.pyx#L803\n\n def roll_kurt(ndarray[float64_t] values, ndarray[int64_t] start,\n ndarray[int64_t] end, int64_t minp) -> np.ndarray:\n \"\"\"\n ret = _roll_kurt(x.to_numpy().astype(float), start, end, minp)\n return Series(ret, nan_to_null=True)\n" }, { "path": "deprecated/wrapper.py", "content": "\"\"\"\nAnother wrapper for raised github issue\n\nbefore:\nexpr.ta.func(..., skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False)\n\nnow:\nfunc(expr, ..., skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False)\n\nThis wrapper allows the prefix expression to be easily used in genetic algorithm tools\n\n此处是对github issues上提出的 注册命名空间 方案的再封装\n\n之前的使用方法为\nexpr.ta.func(..., skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False)\n\n封装后方法为\nfunc(expr, ..., skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False)\n\n此种封装方法的优点是前缀表达式方便输入到遗传算法工具包中使用\n\"\"\"\n\nfrom functools import wraps\n\nimport talib as _talib\nfrom polars import Expr\nfrom polars import struct\nfrom talib import abstract as _abstract\n\nfrom polars_ta.utils.helper import TaLibHelper\n\n_ = TaLibHelper\n\n\ndef ta_func(func, func_name, input_names, output_names,\n *args,\n skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False,\n **kwargs):\n \"\"\"\n\n Parameters\n ----------\n func\n func_name\n input_names: list of str\n 函数输入名\n output_names: list of str\n 函数输出名\n args\n 位置参数\n skip_nan\n output_idx\n schema\n schema_format\n nan_to_null\n kwargs\n 命名参数\n\n Returns\n -------\n\n \"\"\"\n exprs = [arg for arg in args if isinstance(arg, Expr)]\n param = [arg for arg in args if not isinstance(arg, Expr)]\n\n if len(exprs) == 1:\n ef = getattr(exprs[0].ta, func_name)\n else:\n ef = getattr(struct(*exprs).ta, func_name)\n\n return ef(*param,\n skip_nan=skip_nan, output_idx=output_idx, schema=schema, schema_format=schema_format, nan_to_null=nan_to_null,\n **kwargs)\n\n\ndef ta_decorator(func, func_name, input_names, output_names):\n @wraps(func)\n def decorated(*args,\n skip_nan=False, output_idx=None, schema=None, schema_format='{}', nan_to_null=False,\n **kwargs):\n return ta_func(func, func_name, input_names, output_names,\n *args,\n skip_nan=skip_nan, output_idx=output_idx, schema=schema, schema_format=schema_format, nan_to_null=nan_to_null,\n **kwargs)\n\n return decorated\n\n\ndef init(to_globals=False, name_format='{}'):\n \"\"\"初始化环境\n\n Parameters\n ----------\n to_globals: bool\n 是否注册到全局变量中\n name_format: str\n 函数名格式\n\n Returns\n -------\n class\n 对象实例，可调用其内的方法\n\n \"\"\"\n\n class TA_LIB:\n pass\n\n lib = TA_LIB()\n for i, func_name in enumerate(_talib.get_functions()):\n \"\"\"talib遍历\"\"\"\n _ta_func = getattr(_talib, func_name)\n info = _abstract.Function(func_name).info\n output_names = info['output_names']\n input_names = info['input_names']\n\n # 创建函数\n f = ta_decorator(_ta_func, func_name, input_names, output_names)\n\n name = name_format.format(func_name)\n setattr(lib, name, f)\n\n if to_globals:\n from inspect import currentframe\n frame = currentframe().f_back\n frame.f_globals[name] = f\n\n return lib\n" }, { "path": "docs/index.md", "content": "# polars_ta\n\n尽量接近`WorldQuant Alpha101`风格函数，但部分又做了一定的调整\n\n例如：\n\n1. x.abs().log()\n - 可利用`IDE`的自动补全，输入方便\n - 默认是一个输入，多输入要通过参数列表。输入不统一\n2. log(abs_(x))\n - 都是通过参数列表，输入统一\n - 一层套一层，正好对应表达式树，直接可用于遗传规划\n - `abs`与`python`内置函数冲突，使用`abs_`代替,同样还有`and_`、`int_`、`max_`等\n\n## References\n\n- https://platform.worldquantbrain.com/learn/operators/operators\n- https://github.com/TA-Lib/ta-lib" }, { "path": "docs/ta/index.md", "content": "用`polars`函数模仿`TA-Lib`的函数" }, { "path": "docs/ta/momentum.md", "content": "::: polars_ta.ta.momentum" }, { "path": "docs/ta/operators.md", "content": "::: polars_ta.ta.operators" }, { "path": "docs/ta/overlap.md", "content": "::: polars_ta.ta.overlap" }, { "path": "docs/ta/price.md", "content": "::: polars_ta.ta.price" }, { "path": "docs/ta/statistic.md", "content": "::: polars_ta.ta.statistic" }, { "path": "docs/ta/transform.md", "content": "::: polars_ta.ta.transform" }, { "path": "docs/ta/volatility.md", "content": "::: polars_ta.ta.volatility" }, { "path": "docs/ta/volume.md", "content": "::: polars_ta.ta.volume" }, { "path": "docs/talib/index.md", "content": "直接调用`TA-Lib`, 多输出时返回`struct`, 可以使用`.struct[0]`取到对应字段的`Series`.\n\n::: polars_ta.talib" }, { "path": "docs/tdx/arithmetic.md", "content": "::: polars_ta.tdx.arithmetic" }, { "path": "docs/tdx/choice.md", "content": "::: polars_ta.tdx.choice" }, { "path": "docs/tdx/energy.md", "content": "::: polars_ta.tdx.energy" }, { "path": "docs/tdx/logical.md", "content": "::: polars_ta.tdx.logical" }, { "path": "docs/tdx/moving_average.md", "content": "::: polars_ta.tdx.moving_average" }, { "path": "docs/tdx/over_bought_over_sold.md", "content": "::: polars_ta.tdx.over_bought_over_sold" }, { "path": "docs/tdx/pattern.md", "content": "::: polars_ta.tdx.pattern" }, { "path": "docs/tdx/pattern_feature.md", "content": "::: polars_ta.tdx.pattern_feature" }, { "path": "docs/tdx/pressure_support.md", "content": "::: polars_ta.tdx.pressure_support" }, { "path": "docs/tdx/reference.md", "content": "::: polars_ta.tdx.reference" }, { "path": "docs/tdx/statistic.md", "content": "::: polars_ta.tdx.statistic" }, { "path": "docs/tdx/times.md", "content": "::: polars_ta.tdx.times" }, { "path": "docs/tdx/trend.md", "content": "::: polars_ta.tdx.trend" }, { "path": "docs/tdx/trend_feature.md", "content": "::: polars_ta.tdx.trend_feature" }, { "path": "docs/tdx/volume.md", "content": "::: polars_ta.tdx.volume" }, { "path": "docs/wq/arithmetic.md", "content": "::: polars_ta.wq.arithmetic" }, { "path": "docs/wq/cross_sectional.md", "content": "::: polars_ta.wq.cross_sectional" }, { "path": "docs/wq/half_life.md", "content": "::: polars_ta.wq.half_life" }, { "path": "docs/wq/logical.md", "content": "::: polars_ta.wq.logical" }, { "path": "docs/wq/preprocess.md", "content": "::: polars_ta.wq.preprocess" }, { "path": "docs/wq/time_series.md", "content": "::: polars_ta.wq.time_series" }, { "path": "docs/wq/transformational.md", "content": "::: polars_ta.wq.transformational" }, { "path": "docs/wq/vector.md", "content": "::: polars_ta.wq.vector" }, { "path": "examples/alpha101.py", "content": "from datetime import datetime\n\nimport pandas as pd\nimport polars as pl\n\nfrom polars_ta.prefix.tdx import *\nfrom polars_ta.prefix.wq import *\n\n# definition of your features\n# 因子定义\nOPEN, HIGH, LOW, CLOSE, VOLUME, AMOUNT, = pl.col('OPEN'), pl.col('HIGH'), pl.col('LOW'), pl.col('CLOSE'), pl.col('VOLUME'), pl.col('AMOUNT'),\nRETURNS, VWAP, CAP, = pl.col('RETURNS'), pl.col('VWAP'), pl.col('CAP'),\nADV5, ADV10, ADV15, ADV20, ADV30, ADV40, ADV50, ADV60, ADV81, ADV120, ADV150, ADV180, = (\n pl.col('ADV5'), pl.col('ADV10'), pl.col('ADV15'), pl.col('ADV20'),\n pl.col('ADV30'), pl.col('ADV40'), pl.col('ADV50'), pl.col('ADV60'),\n pl.col('ADV81'), pl.col('ADV120'), pl.col('ADV150'), pl.col('ADV180'),)\nSECTOR, INDUSTRY, SUBINDUSTRY, = pl.col('SECTOR'), pl.col('INDUSTRY'), pl.col('SUBINDUSTRY'),\n\n# you can only call column-alike factor directly. Since ts_, cs_ and gp_ etc. need to be used with group_by, please use expr_codegen project for more complex formula\n# 列因子才可以直接调用，而ts_、cs_和gp_等公式需要套用group_by，复杂公式请使用expr_codegen项目\nalpha_041 = (((HIGH * LOW) ** 0.5) - VWAP)\nalpha_054 = ((-1 * ((LOW - CLOSE) * (OPEN ** 5))) / ((LOW - HIGH) * (CLOSE ** 5)))\nalpha_101 = ((CLOSE - OPEN) / ((HIGH - LOW) + 0.001))\n\n# simulate 5000 stocks for 10 years\n# 模拟5000支股票10年\nASSET_COUNT = 5000\nDATE_COUNT = 250 * 10\nDATE = pd.date_range(datetime(2020, 1, 1), periods=DATE_COUNT).repeat(ASSET_COUNT)\nASSET = [f'A{i:04d}' for i in range(ASSET_COUNT)] * DATE_COUNT\n\n# your test data. multi-asset multi-feature\n# 测试数据。多资产多特征\ndf = pl.DataFrame(\n {\n 'date': DATE,\n 'asset': ASSET,\n \"OPEN\": np.random.rand(DATE_COUNT * ASSET_COUNT),\n \"HIGH\": np.random.rand(DATE_COUNT * ASSET_COUNT),\n \"LOW\": np.random.rand(DATE_COUNT * ASSET_COUNT),\n \"CLOSE\": np.random.rand(DATE_COUNT * ASSET_COUNT),\n \"VWAP\": np.random.rand(DATE_COUNT * ASSET_COUNT),\n \"FILTER\": np.tri(DATE_COUNT, ASSET_COUNT, k=-2).reshape(-1),\n }\n)\n\n# filter out some data. test if it will raise error when the length is not enough\n# 每行数据量不同，测试是否会因为长度不足报错\ndf = df.filter(pl.col(\"FILTER\") == 1)\n\n# some simple features defined in Alpha101. They don't contain groupby in time-series or cross-section\n# 部分Alpha101计算。不涉及时序和横截面，可直接计算\ndf = df.with_columns(\n alpha_041=alpha_041,\n alpha_054=alpha_054,\n alpha_101=alpha_101,\n)\n\n\ndef func_ts_date(df: pl.DataFrame) -> pl.DataFrame:\n # ensure sorted before computing time-series features\n # 时序指标计算前一定要保证有序\n df = df.sort(by=['date'])\n # Example of generating features\n # 演示生成大量指标\n df = df.with_columns([\n # load from sub-package `wq`\n # 从wq中导入指标\n *[ts_returns(CLOSE, i).alias(f'ROCP_{i:03d}') for i in (1, 3, 5, 10, 20, 60, 120)],\n *[ts_mean(CLOSE, i).alias(f'SMA_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_std_dev(CLOSE, i).alias(f'STD_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_max(HIGH, i).alias(f'HHV_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_min(LOW, i).alias(f'LLV_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_rank(CLOSE, i).alias(f'RANK_{i:03d}') for i in (5, 10, 20, 60, 120)],\n *[ts_arg_max(HIGH, i).alias(f'HHVBAR_{i:03d}') for i in (5, 10, 20, 60, 120)],\n\n # load from sub-package `tdx`\n # 从tdx中导入指标\n *[ts_RSI(CLOSE, i).alias(f'RSI_{i:03d}') for i in (6, 12, 24)],\n ])\n\n return df\n\n\n# ensure grouped before applying to multi-asset data\n# 多资产需要先按资产分组\ndf = df.group_by('asset').map_groups(func_ts_date)\n\nprint(df)\n" }, { "path": "examples/demo_ta3.py", "content": "\"\"\"\nThis is how we wrap upon polars's code\n以下是polars提供的方案基础上封装的方案\nfunc(expr)\n\n\"\"\"\nimport polars as pl\n\nfrom polars_ta.talib import *\n\ndf = pl.DataFrame(\n {\n \"A\": [5, None, 3, 2, 1],\n \"B\": [5, 4, None, 2, 1],\n \"C\": [5, 4, 3, 2, 1],\n }\n)\n\ndf = df.with_columns([\n # single input single ouput, no need to handle null/nan values\n # # 一输入一输出，不需处理空值\n COS(pl.col('A')).alias('COS'),\n # 多输入一输出\n ATR(pl.col('A'), pl.col('B'), pl.col('C'), 2).alias('ATR2'),\n # 原为一输入多输出，通过ret_idx指定一输出\n BBANDS(pl.col('A'), timeperiod=2).alias('BBANDS'),\n # 原为多输入多输出，通过ret_idx指定一输出\n AROON('A', 'B', timeperiod=2).alias('AROON'),\n])\n\nprint(df)\n" }, { "path": "mkdocs.yml", "content": "site_name: polars_ta API References\nsite_url: https://polars-ta.readthedocs.io/en/latest/\n\nnav:\n - Home: index.md\n - polars_ta.wq:\n - wq/arithmetic.md\n - wq/cross_sectional.md\n - wq/logical.md\n - wq/preprocess.md\n - wq/time_series.md\n - wq/transformational.md\n - wq/vector.md\n - polars_ta.ta:\n - ta/index.md\n - ta/momentum.md\n - ta/operators.md\n - ta/overlap.md\n - ta/price.md\n - ta/statistic.md\n - ta/transform.md\n - ta/volatility.md\n - ta/volume.md\n - polars_ta.tdx:\n - tdx/arithmetic.md\n - tdx/choice.md\n - tdx/energy.md\n - tdx/logical.md\n - tdx/moving_average.md\n - tdx/over_bought_over_sold.md\n - tdx/pattern.md\n - tdx/pattern_feature.md\n - tdx/pressure_support.md\n - tdx/reference.md\n - tdx/statistic.md\n - tdx/trend.md\n - tdx/trend_feature.md\n - tdx/volume.md\n - polars_ta.talib: talib/index.md\n\ntheme:\n name: material\n features:\n - navigation.tracking\n - navigation.instant\n - navigation.tabs\n - navigation.tabs.sticky\n - navigation.footer\n - navigation.indexes\n - content.tabs.link\n - content.code.copy\n\nstrict: true\n\nplugins:\n - search\n - mkdocstrings:\n default_handler: python\n handlers:\n python:\n paths: [ . ]\n options:\n summary: true\n show_root_heading: true\n show_if_no_docstring: true\n show_source: false\n show_signature_annotations: true\n docstring_style: numpy\n - llmstxt:\n full_output: llms-full.txt\n sections:\n Usage documentation:\n - index.md\n - wq/*.md\n - ta/*.md\n - tdx/*.md\n - talib/*.md\n\n" }, { "path": "nan_to_null.md", "content": "# nan_to_null\n\nIt is important to handle nan and null values correctly. In many cases, null values may occur. Such as:\n\n1. Stock suspension, data missing, etc.\n2. Calculation exceptions, such as division by 0, log non-positive, etc.\n\n\n## Null (None in Python) and NaN are different!\n\n```python\nNone == None # True # (It should be `None is None`)\nnp.nan == np.nan # False\n```\n\nBut in `pandas/numpy` many do not distinguish between `None` and `np.nan`, so you should always use `is_null/is_nan` functions to be safe.\n\n1. In `pandas` 1.x, the representation of null values is:\n - floating point uses `nan`,\n - string uses `None`,\n - time uses `pd.NaT`,\n - integer has no null value representation, it can be converted to a floating-point number.\n2. In `pandas` 2.x, the backend can use `Arrow`. It has two storage areas, one is the original data, and the other is the validity bitmap array, which marks whether it is `null`, so there is no data type restriction.\n3. In `polars`, the backend is `Arrow`.\n\nWith a bitmap array, counting `null` values is faster than traversing the original array.\n\n## `polars`\n\n1. Most functions adapt to `null` but not `nan`, the best way should be using `fill_nan(None)` first.\n2. To adapt to third-party packages, such as `TA-Lib`, when calling `to_numpy`, `null` will be automatically converted to `nan`, and the return value should be `pl.Series(, nan_to_null=True)`.\n3. For the case of `null/nan` appearing in the middle, there is currently no good way to handle it.\n\n## Reference\n\n1. https://docs.pola.rs/user-guide/expressions/missing-data/#null-and-nan-values\n2. https://pandas.pydata.org/docs/user_guide/missing_data.html\n\n\n\n\n\n# nan_to_null\n\n空值的处理是一件非常头疼的事性。多种情况下可能会出现空值。如：\n\n1. 股票停牌、数据缺失 等\n2. 计算异常，如：除0，log非正数 等\n\n## None与NaN区别大不同\n\n比较结果不同\n\n```python\nNone == None # True # 注意None is None才是规范的写法\nnp.nan == np.nan # False\n```\n\n但在`pandas/numpy`中将`None`当成`np.nan`使用，这类容易混淆的地方，一定要使用`is_null/is_nan`函数才稳妥\n\n1. 在pandas 1.x版本中，空值表示方法有：浮点用`nan`，字符串用`None`，时间用`pd.NaT`，整型没有空值表示方法，只能转成浮点\n2. 在pandas 2.x版本中，后端可用`Arrow`。它有两块存储区域，一块还是原数据，另一块则是有效性位图数组，由它标记是否为`null`, 所以没有数据类型限制\n3. 在polars中后端是`Arrow`\n\n有位图数组，所以统计`null`数量肯定比遍历原数组要快\n\n## polars\n\n1. 大部分函数只适配了`null`，所以对于`nan`，最合适的处理方法是`fill_nan(None)`后再处理\n2. 需要调用第三方函数时，以`TA-Lib`为例，`to_numpy`时会自动`null`转`nan`，返回值需`pl.Series(, nan_to_null=True)`\n3. 对于中段出现`null/nan`的情况，目前还没有很好的处理方法\n\n# 参考\n\n1. https://docs.pola.rs/user-guide/expressions/missing-data/#null-and-nan-values\n2. https://pandas.pydata.org/docs/user_guide/missing_data.html\n" }, { "path": "point_in_time.md", "content": "# Point In Time\n\nEssentially it is due to the fix of historical erroneous data. Sometimes we even need to fix the historical market data.\n\n1. Market data：date\\asset\\features\\[update_time]\n2. Financial statement data：date\\asset\\features\\update_time\n\n`date`: time of the report,\n`update_time`: time of the release\n\n## Update of Financial Statement Data\n\nThere are two ways to download the data.\n\n1. Download by `date`, just specify the 4 report periods of each year. But if the historical report period is changed in the new record, and the user cannot know which report period it is, the entire data needs to be downloaded.\n2. Download by `update_time`, which can perfectly solve the above problem. But the release date is not fixed, and holidays may also occur. Downloading by day is inefficient, so it needs to be downloaded period by period.\n\n## Storage of Financial Statement Data\n\nSince the data is downloaded by `update_time`, it should be also stored by `update_time`, which is convenient for reading and updating.\n\n## Storage of Historical Market Data\n\nIn market data, we do not need to distinguish the `date` and `update_time`.\n\nTwo ways to modify market data:\n\n1. Similar to financial statement data, only add without modifying.\n - Stored by day, is the record appended to an extra file or added to the file of that day?\n - Because the storage is ordered by `update_time`, it should be added to the extra last file.\n2. Directly modify the original data. In market data, the `update_time` is generally omitted, so the update method is more often to replace the old data.\n\n## Dealing with data from PIT\n\n1. The data update time may be on weekends. But the ideal way is to handle it the same as market data, so the `date` needs to be moved to Friday, and the `update_time` remains the same as the real time.\n2. The PIT processing is divided into three steps.\n - filter out the corresponding time of the DF,\n - calculate the various time series indicators on each DF separately (or not),\n - take the latest part of each DF for merging.\n\nAfter this, the time series calculation cannot be executed again, because it will introduce future data.\n\n\n# Point In Time\n\n本质上是因发现历史数据有问题，需要进行修改而产生的。我们使用的行情数据其实也会有修改需求。所以是否能将其一起讨论呢？\n\n1. 行情数据：date\\asset\\features\\[update_time]\n2. 财务数据：date\\asset\\features\\update_time\n\ndate表示报告期、update_time表示公布日期\n\n## 财务数据的更新\n\n1. 按报告期date下载，只要指定每年的4个报告期即可下载。但新记录中如果改了历史报告期。而用户又无法知道是哪个报告期，得全下\n2. 按公布日update_time下载，能完美的解决上面的问题。但公布日不固定、节假日也会发生，按日下载又效率低，所以得按时间段下载\n\n## 财务数据的存储\n\n由于已经选定用update_time来下载数据，所以存储时也按update_time来分文件名，这样也方便读取和更新\n\n## 行情数据的更新与存储\n\n行情数据中，date其本质就是报告期，又因更新时间等于报告期，所以被省略了，同股票date一般也不会重复，如有重复，必有update_time字段做区分，否则数据有误。\n\n行情数据有两种修改方法:\n\n1. 仿财务数据，只添加不修改。（其实财务数据发布后在第二天交易开始前发生了更新可以直接修改）\n - 按日分文件存储，记录是添加到最后一个文件中，还是添加到那天的文件呢？\n - 因为存储是按update_time有序，所以应当添加到最后\n2. 直接修改原数据。在行情数据由于一般省略了update_time，所以更新方式也更多是覆盖\n\n## PIT数据的处理\n\n1. 数据更新时间可能是周末。但理想方式是与行情一样处理，所以date需要移动到周五,update_time保持真实时间\n2. PIT处理分三步，第一步过滤出对应时间的DF、第二步在每个DF上分别计算各时序指标（也可不算），第三步取每个DF中最新部分进行合并\n\n不能在第三步后再执行时序计算了，因为这会引入未来数据" }, { "path": "polars_ta/__init__.py", "content": "from typing import Optional\n\nfrom ._version import __version__\n\n# 默认最小误差值\nTA_EPSILON: float = 1e-8\n# 默认最小样本数量\nMIN_SAMPLES: Optional[int] = None\n" }, { "path": "polars_ta/_version.py", "content": "__version__ = \"0.5.17\"\n" }, { "path": "polars_ta/candles/__init__.py", "content": "from polars_ta.candles.cdl1 import * # noqa\nfrom polars_ta.candles.cdl1_limit import * # noqa\nfrom polars_ta.candles.cdl2 import * # noqa\n" }, { "path": "polars_ta/candles/cdl1.py", "content": "\"\"\"\nIn this file we use\nopen_: Expr, high: Expr, low: Expr, close: Expr\nfor all parameters\nto ensure the similar function signatures when calling them\n\n本文件中参数全用\nopen_: Expr, high: Expr, low: Expr, close: Expr\n统一的好处是在使用时不用考虑函数调用区别\n\"\"\"\nfrom polars import Expr\nfrom polars import max_horizontal, min_horizontal\n\nfrom polars_ta import TA_EPSILON\n\n\n# https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_utility.h#L327\ndef real_body(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"实体\"\"\"\n return (close - open_).abs()\n\n\ndef upper_shadow(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"上影\"\"\"\n return high - max_horizontal(open_, close)\n\n\ndef lower_shadow(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"下影\"\"\"\n return min_horizontal(open_, close) - low\n\n\ndef high_low_range(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"总长\"\"\"\n return high - low\n\n\ndef upper_body(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"实体上沿\"\"\"\n return max_horizontal(open_, close)\n\n\ndef lower_body(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"实体下沿\"\"\"\n return min_horizontal(open_, close)\n\n\ndef shadows(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"阴影\"\"\"\n return high_low_range(open_, high, low, close) - real_body(open_, high, low, close)\n\n\ndef efficiency_ratio(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"\n abs(close-open) / (2 * (high-low) - abs(close-open) + TA_EPSILON)\n K线内的市场效率。两个总长减去一个实体长就是路程\n\n 比较粗略的计算市场效率的方法。丢失了部分路程信息，所以结果会偏大\n \"\"\"\n displacement = real_body(open_, high, low, close)\n distance = 2 * high_low_range(open_, high, low, close) - displacement\n return displacement / (distance + TA_EPSILON)\n\n\n# ====================================\n\ndef candle_color(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"K线颜色\"\"\"\n return close >= open_\n\n\ndef four_price_doji(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"一字\"\"\"\n return low >= (high - TA_EPSILON)\n\n\ndef doji(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"十字(含一字、T字)\"\"\"\n return (open_ - close).abs() <= TA_EPSILON\n\n\ndef dragonfly(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"正T字\"\"\"\n return doji(open_, high, low, close) & (low < close - TA_EPSILON)\n\n\ndef gravestone(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"倒T字\"\"\"\n return doji(open_, high, low, close) & (high > close + TA_EPSILON)\n\n# def candle_range(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n# return real_body(open_, high, low, close)\n# return high_low_range(open_, high, low, close)\n# return shadows(open_, high, low, close)\n#\n#\n# def candle_average(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n# \"\"\"平均值\"\"\"\n# return high_low_range(open_, high, low, close)\n" }, { "path": "polars_ta/candles/cdl1_limit.py", "content": "\"\"\"\nIn this file we use\nopen_: Expr, high: Expr, low: Expr, close: Expr, high_limit: Expr\nopen_: Expr, high: Expr, low: Expr, close: Expr, low_limit: Expr\nfor all parameters\nto ensure the similar function signatures when calling them\n\n本文件中参数全用\nopen_: Expr, high: Expr, low: Expr, close: Expr, high_limit: Expr\nopen_: Expr, high: Expr, low: Expr, close: Expr, low_limit: Expr\n统一的好处是在使用时不用考虑函数调用区别\n\"\"\"\nfrom polars import Expr\n\nfrom polars_ta import TA_EPSILON\nfrom polars_ta.candles.cdl1 import four_price_doji, dragonfly, gravestone\n\n\ndef limit_up(price: Expr, high_limit: Expr) -> Expr:\n \"\"\"\n unified handles for\n reached `high_limit` at `open_`\n reached `high_limit` at `close`\n reached `high_limit` at `high` but `close` < `high_limit` (at least once reached `high_limit` in this day)\n 开盘 涨停\n 收盘 涨停\n 最高 曾涨停\n \"\"\"\n return price >= high_limit - TA_EPSILON\n\n\ndef limit_up_at_open(open_: Expr, high: Expr, low: Expr, close: Expr, high_limit: Expr) -> Expr:\n \"\"\"\n reached `high_limit` at `open_`\n 开盘 涨停\n \"\"\"\n return limit_up(open_, high_limit)\n\n\ndef limit_up_at_close(open_: Expr, high: Expr, low: Expr, close: Expr, high_limit: Expr) -> Expr:\n \"\"\"\n reached `high_limit` at `close`\n 收盘 涨停\n \"\"\"\n return limit_up(close, high_limit)\n\n\ndef limit_up_at_high(open_: Expr, high: Expr, low: Expr, close: Expr, high_limit: Expr) -> Expr:\n \"\"\"\n reached `high_limit` at `high` but `close` < `high_limit` (at least once reached `high_limit` in this day)\n 曾涨停\n \"\"\"\n return limit_up(high, high_limit) & ~limit_up(close, high_limit)\n\n\ndef limit_up_four_price_doji(open_: Expr, high: Expr, low: Expr, close: Expr, high_limit: Expr) -> Expr:\n \"\"\"\n `low` = `high` = `open_` = `close` = `high_limit`\n 一字涨停\n \"\"\"\n return four_price_doji(open_, high, low, close) & limit_up(close, high_limit)\n\n\ndef limit_up_dragonfly(open_: Expr, high: Expr, low: Expr, close: Expr, high_limit: Expr) -> Expr:\n \"\"\"\n `low` < `high` = `open_` = `close` = `high_limit`\n T字涨停\"\"\"\n return limit_up(close, high_limit) & dragonfly(open_, high, low, close)\n\n\n# ======================================\ndef limit_down(price: Expr, low_limit: Expr) -> Expr:\n \"\"\"\n unified handles for\n reached `low_limit` at `open_`\n reached `low_limit` at `close`\n reached `low_limit` at `high` but `close` > `low_limit` (at least once reached `low_limit` in this day)\n 开盘 跌停\n 收盘 跌停\n 最低 曾跌停\n \"\"\"\n return price <= low_limit + TA_EPSILON\n\n\ndef limit_down_at_open(open_: Expr, high: Expr, low: Expr, close: Expr, low_limit: Expr) -> Expr:\n \"\"\"\n reached `low_limit` at `open_`\n 开盘 跌停\n \"\"\"\n return limit_down(open_, low_limit)\n\n\ndef limit_down_at_close(open_: Expr, high: Expr, low: Expr, close: Expr, low_limit: Expr) -> Expr:\n \"\"\"\n reached `low_limit` at `close`\n 收盘 跌停\n \"\"\"\n return limit_down(close, low_limit)\n\n\ndef limit_down_at_high(open_: Expr, high: Expr, low: Expr, close: Expr, low_limit: Expr) -> Expr:\n \"\"\"\n reached `low_limit` at `high` but `close` > `low_limit` (at least once reached `low_limit` in this day)\n 曾跌停\n \"\"\"\n return limit_down(low, low_limit) & ~limit_down(close, low_limit)\n\n\ndef limit_down_four_price_doji(open_: Expr, high: Expr, low: Expr, close: Expr, low_limit: Expr) -> Expr:\n \"\"\"\n `low` = `high` = `open_` = `close` = `low_limit`\n 一字跌停\n \"\"\"\n return four_price_doji(open_, high, low, close) & limit_down(close, low_limit)\n\n\ndef limit_down_gravestone(open_: Expr, high: Expr, low: Expr, close: Expr, low_limit: Expr) -> Expr:\n \"\"\"\n `high` > `low` = `open_` = `close` = `low_limit`\n T字跌停\n \"\"\"\n return limit_down(open_, low_limit) & gravestone(open_, high, low, close)\n" }, { "path": "polars_ta/candles/cdl2.py", "content": "\"\"\"\nIn this file we use\nopen_: Expr, high: Expr, low: Expr, close: Expr\nfor all parameters\nto ensure the similar function signatures when calling them\n\n本文件中参数全用\nopen_: Expr, high: Expr, low: Expr, close: Expr\n统一的好处是在使用时不用考虑函数调用区别\n\"\"\"\nfrom polars import Expr\n\nfrom polars_ta.candles.cdl1 import lower_body, upper_body\n\n\n# https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_utility.h#L360\ndef ts_gap_up(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"跳空高开\"\"\"\n return low > high.shift(1)\n\n\ndef ts_gap_down(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"跳空低开\"\"\"\n return high < low.shift(1)\n\n\ndef ts_real_body_gap_up(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"实体跳空高开\"\"\"\n return lower_body(open_, high, low, close) > upper_body(open_, high, low, close).shift(1)\n\n\ndef ts_real_body_gap_down(open_: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"实体跳空低开\"\"\"\n return upper_body(open_, high, low, close) < lower_body(open_, high, low, close).shift(1)\n" }, { "path": "polars_ta/labels/__init__.py", "content": "from polars_ta.labels.future import * # noqa\n" }, { "path": "polars_ta/labels/_nb.py", "content": "import numpy as np\nfrom numba import jit\nfrom numpy import full\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _triple_barrier(close: np.ndarray, high: np.ndarray, low: np.ndarray, window: int, take_profit: float, stop_loss: float) -> np.ndarray:\n \"\"\"三重障碍打标法\"\"\"\n out = full(close.shape[0], np.nan, dtype=np.float64)\n for i in range(close.shape[0] - window + 1):\n entry_price = close[i]\n if np.isnan(entry_price):\n # out[i] = 0\n continue\n upper_barrier = entry_price * (1 + take_profit)\n lower_barrier = entry_price * (1 - stop_loss)\n for j in range(i + 1, i + window):\n hit_upper = high[j] >= upper_barrier\n hit_lower = low[j] <= lower_barrier\n if hit_upper and hit_lower:\n # TODO 同一天无法知道是先触发止损还是先触发止盈\n # 1. 假定离收盘价远的先触发\n if high[j] - close[j] > close[j] - low[j]:\n out[i] = 1 # 最高价更远，触发止盈\n else:\n out[i] = -1 # 最低价更远，触发止损\n\n # out[i] = -1 # 2. 简化处理认为先触发止损\n break\n if hit_upper:\n out[i] = 1 # 止盈\n break\n if hit_lower:\n out[i] = -1 # 止损\n break\n else:\n # out[i] = 0 # 1. 时间到了触发平仓\n out[i] = np.sign(close[j] / entry_price - 1) # 2. 时间到了触发平仓\n\n return out\n" }, { "path": "polars_ta/labels/future.py", "content": "\"\"\"\n\n由于标签的定义比较灵活，所以以下代码主要用于参考\n\nNotes\n-----\n标签都是未来数据,在机器学习中，只能用于`y`,不能用于`X`。\n\nReferences\n----------\nhttps://mp.weixin.qq.com/s/XtgYezFsslOfW-QyIMr0VA\nhttps://github.com/Rachnog/Advanced-Deep-Trading/blob/master/bars-labels-diff/Labeling.ipynb\n\n\"\"\"\nfrom polars import Expr, struct, Float64\n\nfrom polars_ta.labels._nb import _triple_barrier\nfrom polars_ta.utils.numba_ import batches_i2_o1, struct_to_numpy\nfrom polars_ta.wq import cut, ts_delay, ts_log_diff, log\n\n\ndef ts_log_return(close: Expr, n: int = 5) -> Expr:\n \"\"\"将未来数据当成卖出价后移到买入价位置，计算对数收益率\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 10, 11, 12, 9, 12, 13],\n }).with_columns(\n out1=ts_log_return(pl.col('a'), 3),\n out2=_ts_log_return(pl.col('a'), 3),\n )\n\n shape: (7, 3)\n ┌──────┬───────────┬───────────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ f64 ┆ f64 │\n ╞══════╪═══════════╪═══════════╡\n │ null ┆ null ┆ null │\n │ 10 ┆ -0.105361 ┆ -0.105361 │\n │ 11 ┆ 0.087011 ┆ 0.087011 │\n │ 12 ┆ 0.080043 ┆ 0.080043 │\n │ 9 ┆ null ┆ null │\n │ 12 ┆ null ┆ null │\n │ 13 ┆ null ┆ null │\n └──────┴───────────┴───────────┘\n ```\n\n \"\"\"\n # return (close.shift(-n) / close).log()\n return log(ts_delay(close, -n) / close)\n\n\ndef _ts_log_return(close: Expr, n: int = 5) -> Expr:\n \"\"\"计算对数收益率，但将结果后移\n\n 如果打标签方式复杂，这种最终结果后移的方法更方便\n \"\"\"\n # return (close / close.shift(n)).log().shift(-n)\n return ts_delay(ts_log_diff(close, n), -n)\n\n\ndef ts_simple_return(close: Expr, n: int = 5, threshold: float = 0.0, *more_threshold) -> Expr:\n \"\"\"简单收益率标签。支持二分类、三分类等。对收益率使用`cut`进行分类\n \n Parameters\n ----------\n close\n n:int\n 未来n天\n threshold:float\n 收益率阈值，小于该值为0，大于等于该值为1\n more_threshold:float\n 更多的阈值，用于三分类等。小于该值为1，大于等于该值为2，以此类推\n\n Returns\n -------\n Expr\n 标签列, 类型为UInt32, 取值为0, 1, 2, ...\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 10., 9.99, 9., 10., 11., 11.],\n }).with_columns(\n out1=label_simple_return(pl.col('a'), 1, 0),\n out2=label_simple_return(pl.col('a'), 1, -0.001, 0.001),\n )\n\n shape: (7, 3)\n ┌──────┬──────┬──────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ f64 ┆ u32 ┆ u32 │\n ╞══════╪══════╪══════╡\n │ null ┆ null ┆ null │\n │ 10.0 ┆ 0 ┆ 0 │\n │ 9.99 ┆ 0 ┆ 0 │\n │ 9.0 ┆ 1 ┆ 2 │\n │ 10.0 ┆ 1 ┆ 2 │\n │ 11.0 ┆ 0 ┆ 1 │\n │ 11.0 ┆ null ┆ null │\n └──────┴──────┴──────┘\n\n \"\"\"\n return cut(close.pct_change(n).shift(-n), threshold, *more_threshold)\n\n\ndef ts_triple_barrier(close: Expr, high: Expr, low: Expr, d: int = 5, take_profit: float = 0.1, stop_loss: float = 0.05) -> Expr:\n \"\"\"三重障碍打标法\n\n Parameters\n ----------\n close:Expr\n 收盘价\n high:Expr\n 最高价\n low:Expr\n 最低价\n d:int\n 时间窗口\n take_profit:float\n 止盈比例\n stop_loss:float\n 止损比例\n\n Returns\n -------\n Expr\n 标签列。取值为-1止损, 1止盈，0时间到期\n\n Notes\n -----\n 1. `high`, `low`在粗略情况下可用`close`代替\n 2. 时间到期时，根据盈亏返回不同的标签\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n \"close\": [np.nan, 1, 1, 1.0],\n \"high\": [np.nan, 1, 1.1, 1],\n \"low\": [np.nan, 1, 1, 0.95],\n }).with_columns(\n out=ts_triple_barrier(pl.col(\"close\"), pl.col(\"high\"), pl.col(\"low\"), 2, 0.1, 0.05)\n )\n\n shape: (4, 4)\n ┌───────┬──────┬──────┬──────┐\n │ close ┆ high ┆ low ┆ out │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ f64 ┆ f64 ┆ f64 ┆ f64 │\n ╞═══════╪══════╪══════╪══════╡\n │ NaN ┆ NaN ┆ NaN ┆ null │\n │ 1.0 ┆ 1.0 ┆ 1.0 ┆ 1.0 │\n │ 1.0 ┆ 1.1 ┆ 1.0 ┆ -1.0 │\n │ 1.0 ┆ 1.0 ┆ 0.95 ┆ null │\n └───────┴──────┴──────┴──────┘\n ```\n\n \"\"\"\n return struct(f0=close, f1=high, f2=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3), _triple_barrier, d, take_profit, stop_loss), return_dtype=Float64)\n" }, { "path": "polars_ta/noise.py", "content": "\"\"\"\nThe following comes from Trading Systems and Methods, Chapter 1, Measuring Noise\n\nWe do not provide ATR and STD here.\n\n\nATR与STD也是一种度量波动的方法，这里不再提供\n\n以下方法来自于Trading Systems and Methods, Chapter 1, Measuring Noise\n\nReferences\n----------\nhttps://zhuanlan.zhihu.com/p/544744582\n\n\"\"\"\nimport numpy as np\nfrom polars import Expr\n\n\ndef ts_efficiency_ratio(close: Expr, timeperiod: int = 14) -> Expr:\n \"\"\"效率系数。值越大，噪音越小。最大值为1，最小值为0\n\n 本质上是位移除以路程\n \"\"\"\n t1 = close.diff(timeperiod).abs()\n t2 = close.diff(1).abs().rolling_sum(timeperiod)\n return t1 / t2\n\n\ndef ts_price_density(high: Expr, low: Expr, timeperiod: int = 14) -> Expr:\n \"\"\"价格密度。值越大，噪音越大\n\n 如果K线高低相连，上涨为1，下跌也为1\n 如果K线高低平行，值大于1，最大为timeperiod\n \"\"\"\n t1 = (high - low).rolling_sum(timeperiod)\n t2 = high.rolling_max(timeperiod) - low.rolling_min(timeperiod)\n return t1 / t2\n\n\ndef ts_fractal_dimension(high: Expr, low: Expr, close: Expr, timeperiod: int = 14) -> Expr:\n \"\"\"分形维度。值越大，噪音越大\"\"\"\n n1 = (1 / timeperiod) ** 2\n n2 = np.log(2)\n n3 = np.log(2 * timeperiod)\n\n t1 = high.rolling_max(timeperiod) - low.rolling_min(timeperiod)\n t2 = close.diff(1)\n\n L = (n1 + t2 / t1).sqrt().rolling_sum(timeperiod)\n return 1 + (L.log() + n2) / n3\n" }, { "path": "polars_ta/performance/__init__.py", "content": "" }, { "path": "polars_ta/performance/drawdown.py", "content": "from polars import Expr\n\n\ndef ts_max_drawdown(close: Expr) -> Expr:\n \"\"\"最大回撤\"\"\"\n return close - close.cum_max()\n\n\ndef ts_max_drawdown_rate(close: Expr) -> Expr:\n \"\"\"最大回撤率\"\"\"\n return close / close.cum_max() - 1\n" }, { "path": "polars_ta/performance/returns.py", "content": "from polars import Expr\n\nfrom polars_ta.wq.arithmetic import log1p, expm1\n\n\ndef ts_cum_return(close: Expr) -> Expr:\n \"\"\"\n cumulative return\n 累计收益\n \"\"\"\n\n return close / close.drop_nulls().first()\n\n\ndef simple_to_log_return(x: Expr) -> Expr:\n \"\"\"简单收益率 转 对数收益率\"\"\"\n return log1p(x)\n\n\ndef log_to_simple_return(x: Expr) -> Expr:\n \"\"\"对数收益率 转 简单收益率\"\"\"\n return expm1(x)\n" }, { "path": "polars_ta/prefix/__init__.py", "content": "" }, { "path": "polars_ta/prefix/cdl.py", "content": "from polars_ta.candles import * # noqa\n" }, { "path": "polars_ta/prefix/labels.py", "content": "from polars_ta.labels import * # noqa\n" }, { "path": "polars_ta/prefix/reports.py", "content": "from polars_ta.reports import * # noqa\n" }, { "path": "polars_ta/prefix/ta.py", "content": "# this code is auto generated by tools/prefix_ta.py\n\nfrom polars_ta.ta.momentum import APO as ts_APO # noqa\nfrom polars_ta.ta.momentum import AROON as ts_AROON # noqa\nfrom polars_ta.ta.momentum import MACD as ts_MACD # noqa\nfrom polars_ta.ta.momentum import MOM as ts_MOM # noqa\nfrom polars_ta.ta.momentum import PPO as ts_PPO # noqa\nfrom polars_ta.ta.momentum import ROC as ts_ROC # noqa\nfrom polars_ta.ta.momentum import ROCP as ts_ROCP # noqa\nfrom polars_ta.ta.momentum import ROCR as ts_ROCR # noqa\nfrom polars_ta.ta.momentum import ROCR100 as ts_ROCR100 # noqa\nfrom polars_ta.ta.momentum import RSI as ts_RSI # noqa\nfrom polars_ta.ta.momentum import RSV as ts_RSV # noqa\nfrom polars_ta.ta.momentum import STOCHF as ts_STOCHF # noqa\nfrom polars_ta.ta.momentum import TRIX as ts_TRIX # noqa\nfrom polars_ta.ta.momentum import WILLR as ts_WILLR # noqa\nfrom polars_ta.ta.operators import ADD # noqa\nfrom polars_ta.ta.operators import DIV # noqa\nfrom polars_ta.ta.operators import MAX as ts_MAX # noqa\nfrom polars_ta.ta.operators import MAXINDEX as ts_MAXINDEX # noqa\nfrom polars_ta.ta.operators import MIN as ts_MIN # noqa\nfrom polars_ta.ta.operators import MININDEX as ts_MININDEX # noqa\nfrom polars_ta.ta.operators import MUL # noqa\nfrom polars_ta.ta.operators import SUB # noqa\nfrom polars_ta.ta.operators import SUM as ts_SUM # noqa\nfrom polars_ta.ta.overlap import BBANDS # noqa\nfrom polars_ta.ta.overlap import DEMA as ts_DEMA # noqa\nfrom polars_ta.ta.overlap import EMA as ts_EMA # noqa\nfrom polars_ta.ta.overlap import KAMA as ts_KAMA # noqa\nfrom polars_ta.ta.overlap import MIDPOINT as ts_MIDPOINT # noqa\nfrom polars_ta.ta.overlap import MIDPRICE as ts_MIDPRICE # noqa\nfrom polars_ta.ta.overlap import RMA as ts_RMA # noqa\nfrom polars_ta.ta.overlap import SMA as ts_SMA # noqa\nfrom polars_ta.ta.overlap import TEMA as ts_TEMA # noqa\nfrom polars_ta.ta.overlap import TRIMA as ts_TRIMA # noqa\nfrom polars_ta.ta.overlap import WMA as ts_WMA # noqa\nfrom polars_ta.ta.price import AVGPRICE # noqa\nfrom polars_ta.ta.price import MEDPRICE # noqa\nfrom polars_ta.ta.price import TYPPRICE # noqa\nfrom polars_ta.ta.price import WCLPRICE # noqa\nfrom polars_ta.ta.statistic import BETA as ts_BETA # noqa\nfrom polars_ta.ta.statistic import CORREL as ts_CORREL # noqa\nfrom polars_ta.ta.statistic import LINEARREG as ts_LINEARREG # noqa\nfrom polars_ta.ta.statistic import LINEARREG_ANGLE as ts_LINEARREG_ANGLE # noqa\nfrom polars_ta.ta.statistic import LINEARREG_INTERCEPT as ts_LINEARREG_INTERCEPT # noqa\nfrom polars_ta.ta.statistic import LINEARREG_SLOPE as ts_LINEARREG_SLOPE # noqa\nfrom polars_ta.ta.statistic import STDDEV as ts_STDDEV # noqa\nfrom polars_ta.ta.statistic import TSF as ts_TSF # noqa\nfrom polars_ta.ta.statistic import VAR as ts_VAR # noqa\nfrom polars_ta.ta.transform import ACOS # noqa\nfrom polars_ta.ta.transform import ASIN # noqa\nfrom polars_ta.ta.transform import ATAN # noqa\nfrom polars_ta.ta.transform import CEIL # noqa\nfrom polars_ta.ta.transform import COS # noqa\nfrom polars_ta.ta.transform import COSH # noqa\nfrom polars_ta.ta.transform import EXP # noqa\nfrom polars_ta.ta.transform import FLOOR # noqa\nfrom polars_ta.ta.transform import LN # noqa\nfrom polars_ta.ta.transform import LOG10 # noqa\nfrom polars_ta.ta.transform import SIN # noqa\nfrom polars_ta.ta.transform import SINH # noqa\nfrom polars_ta.ta.transform import SQRT # noqa\nfrom polars_ta.ta.transform import TAN # noqa\nfrom polars_ta.ta.transform import TANH # noqa\nfrom polars_ta.ta.volatility import ATR as ts_ATR # noqa\nfrom polars_ta.ta.volatility import NATR as ts_NATR # noqa\nfrom polars_ta.ta.volatility import TRANGE as ts_TRANGE # noqa\nfrom polars_ta.ta.volume import AD as ts_AD # noqa\nfrom polars_ta.ta.volume import ADOSC as ts_ADOSC # noqa\nfrom polars_ta.ta.volume import OBV as ts_OBV # noqa" }, { "path": "polars_ta/prefix/talib.py", "content": "# this code is auto generated by tools/prefix_talib.py\n \nfrom polars_ta.talib import HT_DCPERIOD as ts_HT_DCPERIOD # noqa\nfrom polars_ta.talib import HT_DCPHASE as ts_HT_DCPHASE # noqa\nfrom polars_ta.talib import HT_PHASOR as ts_HT_PHASOR # noqa\nfrom polars_ta.talib import HT_SINE as ts_HT_SINE # noqa\nfrom polars_ta.talib import HT_TRENDMODE as ts_HT_TRENDMODE # noqa\nfrom polars_ta.talib import ADD # noqa\nfrom polars_ta.talib import DIV # noqa\nfrom polars_ta.talib import MAX # noqa\nfrom polars_ta.talib import MAXINDEX # noqa\nfrom polars_ta.talib import MIN # noqa\nfrom polars_ta.talib import MININDEX # noqa\nfrom polars_ta.talib import MINMAX # noqa\nfrom polars_ta.talib import MINMAXINDEX # noqa\nfrom polars_ta.talib import MULT # noqa\nfrom polars_ta.talib import SUB # noqa\nfrom polars_ta.talib import SUM # noqa\nfrom polars_ta.talib import ACOS # noqa\nfrom polars_ta.talib import ASIN # noqa\nfrom polars_ta.talib import ATAN # noqa\nfrom polars_ta.talib import CEIL # noqa\nfrom polars_ta.talib import COS # noqa\nfrom polars_ta.talib import COSH # noqa\nfrom polars_ta.talib import EXP # noqa\nfrom polars_ta.talib import FLOOR # noqa\nfrom polars_ta.talib import LN # noqa\nfrom polars_ta.talib import LOG10 # noqa\nfrom polars_ta.talib import SIN # noqa\nfrom polars_ta.talib import SINH # noqa\nfrom polars_ta.talib import SQRT # noqa\nfrom polars_ta.talib import TAN # noqa\nfrom polars_ta.talib import TANH # noqa\nfrom polars_ta.talib import ADX as ts_ADX # noqa\nfrom polars_ta.talib import ADXR as ts_ADXR # noqa\nfrom polars_ta.talib import APO as ts_APO # noqa\nfrom polars_ta.talib import AROON as ts_AROON # noqa\nfrom polars_ta.talib import AROONOSC as ts_AROONOSC # noqa\nfrom polars_ta.talib import BOP as ts_BOP # noqa\nfrom polars_ta.talib import CCI as ts_CCI # noqa\nfrom polars_ta.talib import CMO as ts_CMO # noqa\nfrom polars_ta.talib import DX as ts_DX # noqa\nfrom polars_ta.talib import MACD as ts_MACD # noqa\nfrom polars_ta.talib import MACDEXT as ts_MACDEXT # noqa\nfrom polars_ta.talib import MACDFIX as ts_MACDFIX # noqa\nfrom polars_ta.talib import MFI as ts_MFI # noqa\nfrom polars_ta.talib import MINUS_DI as ts_MINUS_DI # noqa\nfrom polars_ta.talib import MINUS_DM as ts_MINUS_DM # noqa\nfrom polars_ta.talib import MOM as ts_MOM # noqa\nfrom polars_ta.talib import PLUS_DI as ts_PLUS_DI # noqa\nfrom polars_ta.talib import PLUS_DM as ts_PLUS_DM # noqa\nfrom polars_ta.talib import PPO as ts_PPO # noqa\nfrom polars_ta.talib import ROC as ts_ROC # noqa\nfrom polars_ta.talib import ROCP as ts_ROCP # noqa\nfrom polars_ta.talib import ROCR as ts_ROCR # noqa\nfrom polars_ta.talib import ROCR100 as ts_ROCR100 # noqa\nfrom polars_ta.talib import RSI as ts_RSI # noqa\nfrom polars_ta.talib import STOCH as ts_STOCH # noqa\nfrom polars_ta.talib import STOCHF as ts_STOCHF # noqa\nfrom polars_ta.talib import STOCHRSI as ts_STOCHRSI # noqa\nfrom polars_ta.talib import TRIX as ts_TRIX # noqa\nfrom polars_ta.talib import ULTOSC as ts_ULTOSC # noqa\nfrom polars_ta.talib import WILLR as ts_WILLR # noqa\nfrom polars_ta.talib import BBANDS as ts_BBANDS # noqa\nfrom polars_ta.talib import DEMA as ts_DEMA # noqa\nfrom polars_ta.talib import EMA as ts_EMA # noqa\nfrom polars_ta.talib import HT_TRENDLINE as ts_HT_TRENDLINE # noqa\nfrom polars_ta.talib import KAMA as ts_KAMA # noqa\nfrom polars_ta.talib import MA as ts_MA # noqa\nfrom polars_ta.talib import MAMA as ts_MAMA # noqa\nfrom polars_ta.talib import MAVP as ts_MAVP # noqa\nfrom polars_ta.talib import MIDPOINT as ts_MIDPOINT # noqa\nfrom polars_ta.talib import MIDPRICE as ts_MIDPRICE # noqa\nfrom polars_ta.talib import SAR as ts_SAR # noqa\nfrom polars_ta.talib import SAREXT as ts_SAREXT # noqa\nfrom polars_ta.talib import SMA as ts_SMA # noqa\nfrom polars_ta.talib import T3 as ts_T3 # noqa\nfrom polars_ta.talib import TEMA as ts_TEMA # noqa\nfrom polars_ta.talib import TRIMA as ts_TRIMA # noqa\nfrom polars_ta.talib import WMA as ts_WMA # noqa\nfrom polars_ta.talib import CDL2CROWS as ts_CDL2CROWS # noqa\nfrom polars_ta.talib import CDL3BLACKCROWS as ts_CDL3BLACKCROWS # noqa\nfrom polars_ta.talib import CDL3INSIDE as ts_CDL3INSIDE # noqa\nfrom polars_ta.talib import CDL3LINESTRIKE as ts_CDL3LINESTRIKE # noqa\nfrom polars_ta.talib import CDL3OUTSIDE as ts_CDL3OUTSIDE # noqa\nfrom polars_ta.talib import CDL3STARSINSOUTH as ts_CDL3STARSINSOUTH # noqa\nfrom polars_ta.talib import CDL3WHITESOLDIERS as ts_CDL3WHITESOLDIERS # noqa\nfrom polars_ta.talib import CDLABANDONEDBABY as ts_CDLABANDONEDBABY # noqa\nfrom polars_ta.talib import CDLADVANCEBLOCK as ts_CDLADVANCEBLOCK # noqa\nfrom polars_ta.talib import CDLBELTHOLD as ts_CDLBELTHOLD # noqa\nfrom polars_ta.talib import CDLBREAKAWAY as ts_CDLBREAKAWAY # noqa\nfrom polars_ta.talib import CDLCLOSINGMARUBOZU as ts_CDLCLOSINGMARUBOZU # noqa\nfrom polars_ta.talib import CDLCONCEALBABYSWALL as ts_CDLCONCEALBABYSWALL # noqa\nfrom polars_ta.talib import CDLCOUNTERATTACK as ts_CDLCOUNTERATTACK # noqa\nfrom polars_ta.talib import CDLDARKCLOUDCOVER as ts_CDLDARKCLOUDCOVER # noqa\nfrom polars_ta.talib import CDLDOJI as ts_CDLDOJI # noqa\nfrom polars_ta.talib import CDLDOJISTAR as ts_CDLDOJISTAR # noqa\nfrom polars_ta.talib import CDLDRAGONFLYDOJI as ts_CDLDRAGONFLYDOJI # noqa\nfrom polars_ta.talib import CDLENGULFING as ts_CDLENGULFING # noqa\nfrom polars_ta.talib import CDLEVENINGDOJISTAR as ts_CDLEVENINGDOJISTAR # noqa\nfrom polars_ta.talib import CDLEVENINGSTAR as ts_CDLEVENINGSTAR # noqa\nfrom polars_ta.talib import CDLGAPSIDESIDEWHITE as ts_CDLGAPSIDESIDEWHITE # noqa\nfrom polars_ta.talib import CDLGRAVESTONEDOJI as ts_CDLGRAVESTONEDOJI # noqa\nfrom polars_ta.talib import CDLHAMMER as ts_CDLHAMMER # noqa\nfrom polars_ta.talib import CDLHANGINGMAN as ts_CDLHANGINGMAN # noqa\nfrom polars_ta.talib import CDLHARAMI as ts_CDLHARAMI # noqa\nfrom polars_ta.talib import CDLHARAMICROSS as ts_CDLHARAMICROSS # noqa\nfrom polars_ta.talib import CDLHIGHWAVE as ts_CDLHIGHWAVE # noqa\nfrom polars_ta.talib import CDLHIKKAKE as ts_CDLHIKKAKE # noqa\nfrom polars_ta.talib import CDLHIKKAKEMOD as ts_CDLHIKKAKEMOD # noqa\nfrom polars_ta.talib import CDLHOMINGPIGEON as ts_CDLHOMINGPIGEON # noqa\nfrom polars_ta.talib import CDLIDENTICAL3CROWS as ts_CDLIDENTICAL3CROWS # noqa\nfrom polars_ta.talib import CDLINNECK as ts_CDLINNECK # noqa\nfrom polars_ta.talib import CDLINVERTEDHAMMER as ts_CDLINVERTEDHAMMER # noqa\nfrom polars_ta.talib import CDLKICKING as ts_CDLKICKING # noqa\nfrom polars_ta.talib import CDLKICKINGBYLENGTH as ts_CDLKICKINGBYLENGTH # noqa\nfrom polars_ta.talib import CDLLADDERBOTTOM as ts_CDLLADDERBOTTOM # noqa\nfrom polars_ta.talib import CDLLONGLEGGEDDOJI as ts_CDLLONGLEGGEDDOJI # noqa\nfrom polars_ta.talib import CDLLONGLINE as ts_CDLLONGLINE # noqa\nfrom polars_ta.talib import CDLMARUBOZU as ts_CDLMARUBOZU # noqa\nfrom polars_ta.talib import CDLMATCHINGLOW as ts_CDLMATCHINGLOW # noqa\nfrom polars_ta.talib import CDLMATHOLD as ts_CDLMATHOLD # noqa\nfrom polars_ta.talib import CDLMORNINGDOJISTAR as ts_CDLMORNINGDOJISTAR # noqa\nfrom polars_ta.talib import CDLMORNINGSTAR as ts_CDLMORNINGSTAR # noqa\nfrom polars_ta.talib import CDLONNECK as ts_CDLONNECK # noqa\nfrom polars_ta.talib import CDLPIERCING as ts_CDLPIERCING # noqa\nfrom polars_ta.talib import CDLRICKSHAWMAN as ts_CDLRICKSHAWMAN # noqa\nfrom polars_ta.talib import CDLRISEFALL3METHODS as ts_CDLRISEFALL3METHODS # noqa\nfrom polars_ta.talib import CDLSEPARATINGLINES as ts_CDLSEPARATINGLINES # noqa\nfrom polars_ta.talib import CDLSHOOTINGSTAR as ts_CDLSHOOTINGSTAR # noqa\nfrom polars_ta.talib import CDLSHORTLINE as ts_CDLSHORTLINE # noqa\nfrom polars_ta.talib import CDLSPINNINGTOP as ts_CDLSPINNINGTOP # noqa\nfrom polars_ta.talib import CDLSTALLEDPATTERN as ts_CDLSTALLEDPATTERN # noqa\nfrom polars_ta.talib import CDLSTICKSANDWICH as ts_CDLSTICKSANDWICH # noqa\nfrom polars_ta.talib import CDLTAKURI as ts_CDLTAKURI # noqa\nfrom polars_ta.talib import CDLTASUKIGAP as ts_CDLTASUKIGAP # noqa\nfrom polars_ta.talib import CDLTHRUSTING as ts_CDLTHRUSTING # noqa\nfrom polars_ta.talib import CDLTRISTAR as ts_CDLTRISTAR # noqa\nfrom polars_ta.talib import CDLUNIQUE3RIVER as ts_CDLUNIQUE3RIVER # noqa\nfrom polars_ta.talib import CDLUPSIDEGAP2CROWS as ts_CDLUPSIDEGAP2CROWS # noqa\nfrom polars_ta.talib import CDLXSIDEGAP3METHODS as ts_CDLXSIDEGAP3METHODS # noqa\nfrom polars_ta.talib import AVGPRICE # noqa\nfrom polars_ta.talib import MEDPRICE # noqa\nfrom polars_ta.talib import TYPPRICE # noqa\nfrom polars_ta.talib import WCLPRICE # noqa\nfrom polars_ta.talib import BETA as ts_BETA # noqa\nfrom polars_ta.talib import CORREL as ts_CORREL # noqa\nfrom polars_ta.talib import LINEARREG as ts_LINEARREG # noqa\nfrom polars_ta.talib import LINEARREG_ANGLE as ts_LINEARREG_ANGLE # noqa\nfrom polars_ta.talib import LINEARREG_INTERCEPT as ts_LINEARREG_INTERCEPT # noqa\nfrom polars_ta.talib import LINEARREG_SLOPE as ts_LINEARREG_SLOPE # noqa\nfrom polars_ta.talib import STDDEV as ts_STDDEV # noqa\nfrom polars_ta.talib import TSF as ts_TSF # noqa\nfrom polars_ta.talib import VAR as ts_VAR # noqa\nfrom polars_ta.talib import ATR as ts_ATR # noqa\nfrom polars_ta.talib import NATR as ts_NATR # noqa\nfrom polars_ta.talib import TRANGE as ts_TRANGE # noqa\nfrom polars_ta.talib import AD as ts_AD # noqa\nfrom polars_ta.talib import ADOSC as ts_ADOSC # noqa\nfrom polars_ta.talib import OBV as ts_OBV # noqa" }, { "path": "polars_ta/prefix/tdx.py", "content": "# this code is auto generated by tools/prefix_tdx.py\n\nfrom polars_ta.tdx.arithmetic import ABS # noqa\nfrom polars_ta.tdx.arithmetic import ACOS # noqa\nfrom polars_ta.tdx.arithmetic import ADD # noqa\nfrom polars_ta.tdx.arithmetic import ASIN # noqa\nfrom polars_ta.tdx.arithmetic import ATAN # noqa\nfrom polars_ta.tdx.arithmetic import BETWEEN # noqa\nfrom polars_ta.tdx.arithmetic import CEILING # noqa\nfrom polars_ta.tdx.arithmetic import COS # noqa\nfrom polars_ta.tdx.arithmetic import EXP # noqa\nfrom polars_ta.tdx.arithmetic import FLOOR # noqa\nfrom polars_ta.tdx.arithmetic import FRACPART # noqa\nfrom polars_ta.tdx.arithmetic import LN # noqa\nfrom polars_ta.tdx.arithmetic import LOG # noqa\nfrom polars_ta.tdx.arithmetic import MAX # noqa\nfrom polars_ta.tdx.arithmetic import MIN # noqa\nfrom polars_ta.tdx.arithmetic import MOD # noqa\nfrom polars_ta.tdx.arithmetic import POW # noqa\nfrom polars_ta.tdx.arithmetic import REVERSE # noqa\nfrom polars_ta.tdx.arithmetic import ROUND # noqa\nfrom polars_ta.tdx.arithmetic import ROUND2 # noqa\nfrom polars_ta.tdx.arithmetic import SGN # noqa\nfrom polars_ta.tdx.arithmetic import SIGN # noqa\nfrom polars_ta.tdx.arithmetic import SIN # noqa\nfrom polars_ta.tdx.arithmetic import SQRT # noqa\nfrom polars_ta.tdx.arithmetic import SUB # noqa\nfrom polars_ta.tdx.arithmetic import TAN # noqa\nfrom polars_ta.tdx.choice import IF # noqa\nfrom polars_ta.tdx.choice import IFF # noqa\nfrom polars_ta.tdx.choice import IFN # noqa\nfrom polars_ta.tdx.choice import VALUEWHEN as ts_VALUEWHEN # noqa\nfrom polars_ta.tdx.energy import BRAR_AR as ts_BRAR_AR # noqa\nfrom polars_ta.tdx.energy import BRAR_BR as ts_BRAR_BR # noqa\nfrom polars_ta.tdx.energy import CR as ts_CR # noqa\nfrom polars_ta.tdx.energy import MASS as ts_MASS # noqa\nfrom polars_ta.tdx.energy import PSY as ts_PSY # noqa\nfrom polars_ta.tdx.logical import ALL as ts_ALL # noqa\nfrom polars_ta.tdx.logical import ANY as ts_ANY # noqa\nfrom polars_ta.tdx.logical import CROSS as ts_CROSS # noqa\nfrom polars_ta.tdx.logical import DOWNNDAY as ts_DOWNNDAY # noqa\nfrom polars_ta.tdx.logical import EVERY as ts_EVERY # noqa\nfrom polars_ta.tdx.logical import EXIST as ts_EXIST # noqa\nfrom polars_ta.tdx.logical import EXISTR as ts_EXISTR # noqa\nfrom polars_ta.tdx.logical import LAST as ts_LAST # noqa\nfrom polars_ta.tdx.logical import LONGCROSS as ts_LONGCROSS # noqa\nfrom polars_ta.tdx.logical import NDAY as ts_NDAY # noqa\nfrom polars_ta.tdx.logical import NOT # noqa\nfrom polars_ta.tdx.logical import UPNDAY as ts_UPNDAY # noqa\nfrom polars_ta.tdx.moving_average import BBI as ts_BBI # noqa\nfrom polars_ta.tdx.over_bought_over_sold import ATR as ts_ATR # noqa\nfrom polars_ta.tdx.over_bought_over_sold import BIAS as ts_BIAS # noqa\nfrom polars_ta.tdx.over_bought_over_sold import CCI as ts_CCI # noqa\nfrom polars_ta.tdx.over_bought_over_sold import KDJ as ts_KDJ # noqa\nfrom polars_ta.tdx.over_bought_over_sold import MFI as ts_MFI # noqa\nfrom polars_ta.tdx.over_bought_over_sold import MTM as ts_MTM # noqa\nfrom polars_ta.tdx.over_bought_over_sold import RSI as ts_RSI # noqa\nfrom polars_ta.tdx.over_bought_over_sold import RSV as ts_RSV # noqa\nfrom polars_ta.tdx.pattern import ts_WINNER_COST # noqa\nfrom polars_ta.tdx.pattern_feature import 仙人指路 as ts_仙人指路 # noqa\nfrom polars_ta.tdx.pattern_feature import 低开大阳线 as ts_低开大阳线 # noqa\nfrom polars_ta.tdx.pattern_feature import 低点搜寻 as ts_低点搜寻 # noqa\nfrom polars_ta.tdx.pattern_feature import 出水芙蓉 as ts_出水芙蓉 # noqa\nfrom polars_ta.tdx.pattern_feature import 出水芙蓉II as ts_出水芙蓉II # noqa\nfrom polars_ta.tdx.pattern_feature import 剑 as ts_剑 # noqa\nfrom polars_ta.tdx.pattern_feature import 单阳不破选股 as ts_单阳不破选股 # noqa\nfrom polars_ta.tdx.pattern_feature import 四串阳 as ts_四串阳 # noqa\nfrom polars_ta.tdx.pattern_feature import 四串阴 as ts_四串阴 # noqa\nfrom polars_ta.tdx.pattern_feature import 回补跳空向上缺口 as ts_回补跳空向上缺口 # noqa\nfrom polars_ta.tdx.pattern_feature import 均线多头排列 as ts_均线多头排列 # noqa\nfrom polars_ta.tdx.pattern_feature import 均线空头排列 as ts_均线空头排列 # noqa\nfrom polars_ta.tdx.pattern_feature import 天量法则 as ts_天量法则 # noqa\nfrom polars_ta.tdx.pattern_feature import 强势整理 as ts_强势整理 # noqa\nfrom polars_ta.tdx.pattern_feature import 揉搓线 as ts_揉搓线 # noqa\nfrom polars_ta.tdx.pattern_feature import 早晨之星 as ts_早晨之星 # noqa\nfrom polars_ta.tdx.pattern_feature import 旭日初升 as ts_旭日初升 # noqa\nfrom polars_ta.tdx.pattern_feature import 突破 as ts_突破 # noqa\nfrom polars_ta.tdx.pattern_feature import 老鸭头 as ts_老鸭头 # noqa\nfrom polars_ta.tdx.pattern_feature import 蜻蜓点水 as ts_蜻蜓点水 # noqa\nfrom polars_ta.tdx.pattern_feature import 跳空缺口选股 as ts_跳空缺口选股 # noqa\nfrom polars_ta.tdx.pattern_feature import 近日创历史新低 as ts_近日创历史新低 # noqa\nfrom polars_ta.tdx.pattern_feature import 近日创历史新高 as ts_近日创历史新高 # noqa\nfrom polars_ta.tdx.pattern_feature import 高开大阴线 as ts_高开大阴线 # noqa\nfrom polars_ta.tdx.pattern_feature import 鸳鸯底 as ts_鸳鸯底 # noqa\nfrom polars_ta.tdx.pressure_support import BOLL as ts_BOLL # noqa\nfrom polars_ta.tdx.pressure_support import BOLL_M as ts_BOLL_M # noqa\nfrom polars_ta.tdx.reference import BARSLAST as ts_BARSLAST # noqa\nfrom polars_ta.tdx.reference import BARSLASTCOUNT as ts_BARSLASTCOUNT # noqa\nfrom polars_ta.tdx.reference import BARSSINCE as ts_BARSSINCE # noqa\nfrom polars_ta.tdx.reference import BARSSINCEN as ts_BARSSINCEN # noqa\nfrom polars_ta.tdx.reference import COUNT as ts_COUNT # noqa\nfrom polars_ta.tdx.reference import CUMSUM as ts_CUMSUM # noqa\nfrom polars_ta.tdx.reference import DIFF as ts_DIFF # noqa\nfrom polars_ta.tdx.reference import DMA as ts_DMA # noqa\nfrom polars_ta.tdx.reference import EMA as ts_EMA # noqa\nfrom polars_ta.tdx.reference import EXPMA as ts_EXPMA # noqa\nfrom polars_ta.tdx.reference import EXPMEMA as ts_EXPMEMA # noqa\nfrom polars_ta.tdx.reference import FILTER as ts_FILTER # noqa\nfrom polars_ta.tdx.reference import HHV as ts_HHV # noqa\nfrom polars_ta.tdx.reference import HHVBARS as ts_HHVBARS # noqa\nfrom polars_ta.tdx.reference import HOD as ts_HOD # noqa\nfrom polars_ta.tdx.reference import LLV as ts_LLV # noqa\nfrom polars_ta.tdx.reference import LLVBARS as ts_LLVBARS # noqa\nfrom polars_ta.tdx.reference import LOD as ts_LOD # noqa\nfrom polars_ta.tdx.reference import LOWRANGE # noqa\nfrom polars_ta.tdx.reference import MA as ts_MA # noqa\nfrom polars_ta.tdx.reference import MAX # noqa\nfrom polars_ta.tdx.reference import MEMA as ts_MEMA # noqa\nfrom polars_ta.tdx.reference import MIN # noqa\nfrom polars_ta.tdx.reference import MULAR as ts_MULAR # noqa\nfrom polars_ta.tdx.reference import RANGE # noqa\nfrom polars_ta.tdx.reference import REF as ts_REF # noqa\nfrom polars_ta.tdx.reference import REFX as ts_REFX # noqa\nfrom polars_ta.tdx.reference import SMA_CN as ts_SMA_CN # noqa\nfrom polars_ta.tdx.reference import SUM as ts_SUM # noqa\nfrom polars_ta.tdx.reference import SUMIF as ts_SUMIF # noqa\nfrom polars_ta.tdx.reference import TMA as ts_TMA # noqa\nfrom polars_ta.tdx.reference import TR as ts_TR # noqa\nfrom polars_ta.tdx.reference import WMA as ts_WMA # noqa\nfrom polars_ta.tdx.statistic import AVEDEV as ts_AVEDEV # noqa\nfrom polars_ta.tdx.statistic import COVAR as ts_COVAR # noqa\nfrom polars_ta.tdx.statistic import DEVSQ as ts_DEVSQ # noqa\nfrom polars_ta.tdx.statistic import RELATE as ts_RELATE # noqa\nfrom polars_ta.tdx.statistic import SLOPE as ts_SLOPE # noqa\nfrom polars_ta.tdx.statistic import STD as ts_STD # noqa\nfrom polars_ta.tdx.statistic import STDDEV as ts_STDDEV # noqa\nfrom polars_ta.tdx.statistic import STDP as ts_STDP # noqa\nfrom polars_ta.tdx.statistic import VAR as ts_VAR # noqa\nfrom polars_ta.tdx.statistic import VARP as ts_VARP # noqa\nfrom polars_ta.tdx.statistic import ts_up_stat # noqa\nfrom polars_ta.tdx.times import FROMOPEN # noqa\nfrom polars_ta.tdx.times import FROMOPEN_1 # noqa\nfrom polars_ta.tdx.trend import ADX as ts_ADX # noqa\nfrom polars_ta.tdx.trend import ADXR as ts_ADXR # noqa\nfrom polars_ta.tdx.trend import DPO as ts_DPO # noqa\nfrom polars_ta.tdx.trend import EMV as ts_EMV # noqa\nfrom polars_ta.tdx.trend import MINUS_DI as ts_MINUS_DI # noqa\nfrom polars_ta.tdx.trend import MINUS_DM as ts_MINUS_DM # noqa\nfrom polars_ta.tdx.trend import PLUS_DI as ts_PLUS_DI # noqa\nfrom polars_ta.tdx.trend import PLUS_DM as ts_PLUS_DM # noqa\nfrom polars_ta.tdx.trend_feature import N天内出现以涨停收盘 as ts_N天内出现以涨停收盘 # noqa\nfrom polars_ta.tdx.trend_feature import N天内出现涨停 as ts_N天内出现涨停 # noqa\nfrom polars_ta.tdx.trend_feature import N天内有跳空向上缺口 as ts_N天内有跳空向上缺口 # noqa\nfrom polars_ta.tdx.trend_feature import N天内经常涨停 as ts_N天内经常涨停 # noqa\nfrom polars_ta.tdx.trend_feature import N日内上涨多于下跌 as ts_N日内上涨多于下跌 # noqa\nfrom polars_ta.tdx.trend_feature import N日内下跌多于上涨 as ts_N日内下跌多于上涨 # noqa\nfrom polars_ta.tdx.trend_feature import N日内创新低 as ts_N日内创新低 # noqa\nfrom polars_ta.tdx.trend_feature import N日内创新高 as ts_N日内创新高 # noqa\nfrom polars_ta.tdx.trend_feature import N日内阳线多于阴线 as ts_N日内阳线多于阴线 # noqa\nfrom polars_ta.tdx.trend_feature import N日内阴线多于阳线 as ts_N日内阴线多于阳线 # noqa\nfrom polars_ta.tdx.trend_feature import 下跌多日再放量上涨 as ts_下跌多日再放量上涨 # noqa\nfrom polars_ta.tdx.trend_feature import 价量渐低后阳包阴 as ts_价量渐低后阳包阴 # noqa\nfrom polars_ta.tdx.trend_feature import 单日放量 as ts_单日放量 # noqa\nfrom polars_ta.tdx.trend_feature import 小步碎阳 as ts_小步碎阳 # noqa\nfrom polars_ta.tdx.trend_feature import 平台整理 as ts_平台整理 # noqa\nfrom polars_ta.tdx.trend_feature import 拉升后多日调整 as ts_拉升后多日调整 # noqa\nfrom polars_ta.tdx.trend_feature import 持续放量 as ts_持续放量 # noqa\nfrom polars_ta.tdx.trend_feature import 持续缩量 as ts_持续缩量 # noqa\nfrom polars_ta.tdx.trend_feature import 放量上攻 as ts_放量上攻 # noqa\nfrom polars_ta.tdx.trend_feature import 昨日底部十字星 as ts_昨日底部十字星 # noqa\nfrom polars_ta.tdx.trend_feature import 温和放量上攻 as ts_温和放量上攻 # noqa\nfrom polars_ta.tdx.trend_feature import 突然放量 as ts_突然放量 # noqa\nfrom polars_ta.tdx.trend_feature import 突破长期盘整 as ts_突破长期盘整 # noqa\nfrom polars_ta.tdx.trend_feature import 跳空高开或低开 as ts_跳空高开或低开 # noqa\nfrom polars_ta.tdx.trend_feature import 连续N天收阳线 as ts_连续N天收阳线 # noqa\nfrom polars_ta.tdx.trend_feature import 连续N天收阴线 as ts_连续N天收阴线 # noqa\nfrom polars_ta.tdx.trend_feature import 间隔放量 as ts_间隔放量 # noqa\nfrom polars_ta.tdx.trend_feature import 阶段放量 as ts_阶段放量 # noqa\nfrom polars_ta.tdx.trend_feature import 阶段缩量 as ts_阶段缩量 # noqa\nfrom polars_ta.tdx.volume import OBV as ts_OBV # noqa\nfrom polars_ta.tdx.volume import VR as ts_VR # noqa" }, { "path": "polars_ta/prefix/vec.py", "content": "# this code is auto generated by tools/prefix_vec.py\n\nfrom polars_ta.wq.vector import vec_avg as cs_vec_avg # noqa\nfrom polars_ta.wq.vector import vec_choose as cs_vec_choose # noqa\nfrom polars_ta.wq.vector import vec_count as cs_vec_count # noqa\nfrom polars_ta.wq.vector import vec_ir as cs_vec_ir # noqa\nfrom polars_ta.wq.vector import vec_kurtosis as cs_vec_kurtosis # noqa\nfrom polars_ta.wq.vector import vec_l2_norm as cs_vec_l2_norm # noqa\nfrom polars_ta.wq.vector import vec_max as cs_vec_max # noqa\nfrom polars_ta.wq.vector import vec_median as cs_vec_median # noqa\nfrom polars_ta.wq.vector import vec_min as cs_vec_min # noqa\nfrom polars_ta.wq.vector import vec_norm as cs_vec_norm # noqa\nfrom polars_ta.wq.vector import vec_percentage as cs_vec_percentage # noqa\nfrom polars_ta.wq.vector import vec_powersum as cs_vec_powersum # noqa\nfrom polars_ta.wq.vector import vec_range as cs_vec_range # noqa\nfrom polars_ta.wq.vector import vec_skewness as cs_vec_skewness # noqa\nfrom polars_ta.wq.vector import vec_stddev as cs_vec_stddev # noqa\nfrom polars_ta.wq.vector import vec_sum as cs_vec_sum # noqa" }, { "path": "polars_ta/prefix/wq.py", "content": "from polars_ta.wq import * # noqa\n" }, { "path": "polars_ta/reports/__init__.py", "content": "from polars_ta.reports.cicc import * # noqa\n" }, { "path": "polars_ta/reports/cicc.py", "content": "from polars import Expr\n\nfrom polars_ta.wq import ts_corr, ts_zscore, ts_std_dev, ts_regression_slope\n\n\ndef ts_RSRS_R2(high: Expr, low: Expr, n: int = 18, m: int = 600) -> Expr:\n \"\"\"光大RSRS指标，中金QRS指标，R^2调整\n\n References\n ----------\n 中金：金融工程视角下的技术择时艺术\n \"\"\"\n a = ts_corr(high, low, n)\n return ts_zscore(ts_std_dev(high, n) / ts_std_dev(low, n) * a, m) * (a ** 2)\n\n\ndef ts_RSRS(high: Expr, low: Expr, n: int = 18, m: int = 600) -> Expr:\n \"\"\"光大RSRS指标，中金QRS指标\n\n References\n ----------\n 中金：金融工程视角下的技术择时艺术\n \"\"\"\n return ts_zscore(ts_regression_slope(high, low, n), m)\n" }, { "path": "polars_ta/ta/README.md", "content": "# polars_ta.ta\n\n1. Files in this folder mimic `talib`, and implement `polars` versions for the same functions\n2. Since we reduce the functino calls between `Python` and `C` code, it should be faster than `talib`.\n3. We first try to import from `ta`, then from `wq`, and only implement the function if it is not available.\n4. When there is a circular dependency, we use `polars` instead.\n\n\n1. 本文件夹中模仿`talib`，实现同名函数的`polars`版\n2. 由于减少了python与c来回调用，理论上比直接调用`talib`快\n3. 优先从`ta`中导入，然后从`wq`中导入，没有的才实现\n4. 出现循环依赖时，使用`polars`" }, { "path": "polars_ta/ta/__init__.py", "content": "from polars_ta.ta.momentum import * # noqa\nfrom polars_ta.ta.operators import * # noqa\nfrom polars_ta.ta.overlap import * # noqa\nfrom polars_ta.ta.price import * # noqa\nfrom polars_ta.ta.statistic import * # noqa\nfrom polars_ta.ta.transform import * # noqa\nfrom polars_ta.ta.volatility import * # noqa\nfrom polars_ta.ta.volume import * # noqa\n" }, { "path": "polars_ta/ta/momentum.py", "content": "from polars import Expr, when, struct\n\nfrom polars_ta import TA_EPSILON\nfrom polars_ta.ta.operators import MAX\nfrom polars_ta.ta.operators import MIN\nfrom polars_ta.ta.overlap import EMA\nfrom polars_ta.ta.overlap import RMA\nfrom polars_ta.ta.overlap import SMA\nfrom polars_ta.wq.arithmetic import max_\nfrom polars_ta.wq.time_series import ts_delta, ts_arg_max, ts_arg_min\nfrom polars_ta.wq.time_series import ts_returns\n\n\ndef APO(close: Expr, fastperiod: int = 12, slowperiod: int = 26, matype: int = 0) -> Expr:\n if matype == 0:\n return SMA(close, fastperiod) - SMA(close, slowperiod)\n else:\n return EMA(close, fastperiod) - EMA(close, slowperiod)\n\n\ndef AROON(high: Expr, low: Expr, timeperiod: int = 14) -> Expr:\n \"\"\"\n 下轨:(N-LLVBARS(L,N))/N*100,COLORGREEN;\n 上轨:(N-HHVBARS(H,N))/N*100,COLORRED;\n\n Notes\n -----\n You cannot use pd.Series.rolling().arg_max() with reverse order, which leads to a larger result when there are two or more high points\n so we don't use pandas\n pd.Series.rolling().arg_max()没有逆序，导致出现两个及以上最高点时，结果偏大\n\n \"\"\"\n aroondown = 1 - ts_arg_min(low, timeperiod, reverse=True) / timeperiod\n aroonup = 1 - ts_arg_max(high, timeperiod, reverse=True) / timeperiod\n\n return struct(aroondown=aroondown, aroonup=aroonup)\n\n\ndef MACD(close: Expr, fastperiod: int = 12, slowperiod: int = 26, signalperiod: int = 9) -> Expr:\n \"\"\"MACD\n\n Notes\n -----\n When counting how many data we have\n we refer to `slowperiod`, while `talib.MACD` refers to `fastperiod`\n\n talib.MACD有效数据按fastperiod，而本项目按slowperiod\n\n \"\"\"\n macd = EMA(close, fastperiod) - EMA(close, slowperiod)\n macdsignal = EMA(macd, signalperiod)\n macdhist = (macd - macdsignal) # * 2 # 中国版多了乘2\n return struct(macd=macd, macdsignal=macdsignal, macdhist=macdhist)\n\n\ndef MOM(close: Expr, timeperiod: int = 10) -> Expr:\n \"\"\"MOM = (price - prevPrice) [Momentum]\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_MOM.c#L200\n\n \"\"\"\n return ts_delta(close, timeperiod)\n\n\ndef PPO(close: Expr, fastperiod: int = 12, slowperiod: int = 26, matype: int = 0) -> Expr:\n if matype == 0:\n return SMA(close, fastperiod) / SMA(close, slowperiod) - 1\n else:\n return EMA(close, fastperiod) / EMA(close, slowperiod) - 1\n\n\ndef ROC(close: Expr, timeperiod: int = 10) -> Expr:\n \"\"\"ROC = ((price/prevPrice)-1)*100 [Rate of change]\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_ROC.c#L200\n\n \"\"\"\n return ROCP(close, timeperiod) * 100\n\n\ndef ROCP(close: Expr, timeperiod: int = 10) -> Expr:\n \"\"\"ROCP = (price-prevPrice)/prevPrice [Rate of change Percentage]\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_ROCP.c#L202\n\n \"\"\"\n return ts_returns(close, timeperiod)\n\n\ndef ROCR(close: Expr, timeperiod: int = 10) -> Expr:\n \"\"\"ROCR = (price/prevPrice) [Rate of change ratio]\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_ROCR.c#L203\n\n \"\"\"\n return close / close.shift(timeperiod)\n\n\ndef ROCR100(close: Expr, timeperiod: int = 10) -> Expr:\n \"\"\"ROCR100 = (price/prevPrice)*100 [Rate of change ratio 100 Scale]\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_ROCR100.c#L203\n\n \"\"\"\n return ROCR(close, timeperiod) * 100\n\n\ndef RSI(close: Expr, timeperiod: int = 14) -> Expr:\n dif = close.diff().fill_null(0)\n return RMA(max_(dif, 0), timeperiod) / (RMA(dif.abs(), timeperiod) + TA_EPSILON) # * 100\n\n\ndef STOCHF(high: Expr, low: Expr, close: Expr, fastk_period: int = 5, fastd_period: int = 3) -> Expr:\n fastk = RSV(high, low, close, fastk_period)\n fastd = SMA(fastk, fastd_period)\n return struct(fastk=fastk, fastd=fastd)\n\n\ndef TRIX(close: Expr, timeperiod: int = 30) -> Expr:\n EMA1 = EMA(close, timeperiod)\n EMA2 = EMA(EMA1, timeperiod)\n EMA3 = EMA(EMA2, timeperiod)\n return ROCP(EMA3, 1)\n\n\ndef RSV(high: Expr, low: Expr, close: Expr, timeperiod: int = 5) -> Expr:\n \"\"\"RSV=STOCHF_FASTK\n\n (Today's Close - LowestLow)\n FASTK(Kperiod) = --------------------------- * 100\n (HighestHigh - LowestLow)\n\n Notes\n -----\n RSV = STOCHF_FASTK。没乘以100\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_STOCHF.c#L279\n\n \"\"\"\n a = MAX(high, timeperiod)\n b = MIN(low, timeperiod)\n\n # return (close - b) / (a - b + TA_EPSILON)\n return when(a != b).then((close - b) / (a - b)).otherwise(0)\n\n\ndef WILLR(high: Expr, low: Expr, close: Expr, timeperiod: int = 14) -> Expr:\n \"\"\"威廉指标\n\n Notes\n -----\n WILLR=1-RSV\n\n References\n ----------\n - https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_WILLR.c#L294\n - https://www.investopedia.com/terms/w/williamsr.asp\n - https://school.stockcharts.com/doku.php?id=technical_indicators:williams_r\n\n\n \"\"\"\n a = MAX(high, timeperiod)\n b = MIN(low, timeperiod)\n\n # return (a - close) / (a - b + TA_EPSILON)\n return when(a != b).then((a - close) / (a - b)).otherwise(0)\n" }, { "path": "polars_ta/ta/operators.py", "content": "\"\"\"\n通过`import`直接导入或更名的函数\n\n```python\nfrom polars_ta.wq.arithmetic import add as ADD # noqa\nfrom polars_ta.wq.arithmetic import divide as DIV # noqa\nfrom polars_ta.wq.arithmetic import multiply as MUL # noqa\nfrom polars_ta.wq.arithmetic import subtract as SUB # noqa\nfrom polars_ta.wq.time_series import ts_max as MAX # noqa\nfrom polars_ta.wq.time_series import ts_min as MIN # noqa\nfrom polars_ta.wq.time_series import ts_sum as SUM # noqa\n```\n\n\"\"\"\nfrom polars import Expr\n\nfrom polars_ta.wq.arithmetic import add as ADD # noqa\nfrom polars_ta.wq.arithmetic import divide as DIV # noqa\nfrom polars_ta.wq.arithmetic import multiply as MUL # noqa\nfrom polars_ta.wq.arithmetic import subtract as SUB # noqa\nfrom polars_ta.wq.time_series import ts_arg_max\nfrom polars_ta.wq.time_series import ts_arg_min\nfrom polars_ta.wq.time_series import ts_max as MAX # noqa\nfrom polars_ta.wq.time_series import ts_min as MIN # noqa\nfrom polars_ta.wq.time_series import ts_sum as SUM # noqa\n\n\ndef MAXINDEX(close: Expr, timeperiod: int = 30) -> Expr:\n \"\"\"\n\n Notes\n -----\n Comparing to `ts_arg_max` this also marks the abs. position of the max value\n\n 与ts_arg_max的区别是，标记了每个区间最大值的绝对位置，可用来画图标记\n\n\n Examples\n --------\n ```python\n from polars_ta.ta import MAXINDEX as ta_MAXINDEX\n from polars_ta.talib import MAXINDEX as talib_MAXINDEX\n from polars_ta.wq import ts_arg_max\n\n df = pl.DataFrame({\n 'a': [6, 2, 8, 5, 9, 4],\n }).with_columns(\n out1=ts_arg_max(pl.col('a'), 3),\n out2=ta_MAXINDEX(pl.col('a'), 3),\n out3=talib_MAXINDEX(pl.col('a'), 3),\n )\n shape: (6, 4)\n ┌─────┬──────┬──────┬──────┐\n │ a ┆ out1 ┆ out2 ┆ out3 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ u16 ┆ i64 ┆ i32 │\n ╞═════╪══════╪══════╪══════╡\n │ 6 ┆ null ┆ null ┆ 0 │\n │ 2 ┆ null ┆ null ┆ 0 │\n │ 8 ┆ 0 ┆ 2 ┆ 2 │\n │ 5 ┆ 1 ┆ 2 ┆ 2 │\n │ 9 ┆ 0 ┆ 4 ┆ 4 │\n │ 4 ┆ 1 ┆ 4 ┆ 4 │\n └─────┴──────┴──────┴──────┘\n ```\n \"\"\"\n a = close.cum_count()\n b = ts_arg_max(close, timeperiod)\n return a - b - 1\n\n\ndef MININDEX(close: Expr, timeperiod: int = 30) -> Expr:\n \"\"\"\n\n\n\n Examples\n --------\n ```python\n from polars_ta.ta import MININDEX as ta_MININDEX\n from polars_ta.talib import MININDEX as talib_MININDEX\n from polars_ta.wq import ts_arg_min\n\n df = pl.DataFrame({\n 'a': [6, 2, 8, 5, 9, 4],\n\n }).with_columns(\n out1=ts_arg_min(pl.col('a'), 3),\n out2=ta_MININDEX(pl.col('a'), 3),\n out3=talib_MININDEX(pl.col('a'), 3),\n )\n shape: (6, 4)\n ┌─────┬──────┬──────┬──────┐\n │ a ┆ out1 ┆ out2 ┆ out3 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ u16 ┆ i64 ┆ i32 │\n ╞═════╪══════╪══════╪══════╡\n │ 6 ┆ null ┆ null ┆ 0 │\n │ 2 ┆ null ┆ null ┆ 0 │\n │ 8 ┆ 1 ┆ 1 ┆ 1 │\n │ 5 ┆ 2 ┆ 1 ┆ 1 │\n │ 9 ┆ 1 ┆ 3 ┆ 3 │\n │ 4 ┆ 0 ┆ 5 ┆ 5 │\n └─────┴──────┴──────┴──────┘\n ```\n \"\"\"\n a = close.cum_count()\n b = ts_arg_min(close, timeperiod)\n return a - b - 1\n" }, { "path": "polars_ta/ta/overlap.py", "content": "from math import ceil, floor\n\nfrom polars import Expr, struct\n\nfrom polars_ta.ta.operators import MAX\nfrom polars_ta.ta.operators import MIN\nfrom polars_ta.ta.statistic import STDDEV\nfrom polars_ta.wq.time_series import ts_decay_linear as WMA # noqa\nfrom polars_ta.wq.time_series import ts_mean as SMA # noqa\n\n\ndef BBANDS(close: Expr, timeperiod: float = 5.0, nbdevup: float = 2.0, nbdevdn: float = 2.0, matype: float = 0.0) -> Expr:\n middleband = SMA(close, timeperiod)\n stddev = STDDEV(close, timeperiod)\n return struct(upperband=middleband + stddev * nbdevup, middleband=middleband, lowerband=middleband - stddev * nbdevdn)\n\n\ndef DEMA(close: Expr, timeperiod: int = 30) -> Expr:\n EMA1 = EMA(close, timeperiod)\n EMA2 = EMA(EMA1, timeperiod)\n return EMA1 * 2 - EMA2\n\n\ndef EMA(close: Expr, timeperiod: int = 30) -> Expr:\n \"\"\"\n\n References\n ----------\n https://pola-rs.github.io/polars/py-polars/html/reference/expressions/api/polars.Expr.ewm_mean.html#polars.Expr.ewm_mean\n\n \"\"\"\n # 相当于alpha=2/(1+timeperiod)\n return close.ewm_mean(span=timeperiod, adjust=False, min_samples=timeperiod)\n\n\ndef KAMA(close: Expr, timeperiod: int = 30) -> Expr:\n raise\n\n\ndef MIDPOINT(close: Expr, timeperiod: int = 14) -> Expr:\n \"\"\"MIDPOINT = (Highest Value + Lowest Value)/2\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_MIDPOINT.c#L198\n\n \"\"\"\n return (MAX(close, timeperiod) + MIN(close, timeperiod)) / 2\n\n\ndef MIDPRICE(high: Expr, low: Expr, timeperiod: int = 14) -> Expr:\n \"\"\"MIDPRICE = (Highest High + Lowest Low)/2\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_MIDPRICE.c#L202\n\n \"\"\"\n return (MAX(high, timeperiod) + MIN(low, timeperiod)) / 2\n\n\ndef RMA(close: Expr, timeperiod: int = 30) -> Expr:\n \"\"\"TA-Lib does not provide this algorithm explicitly, it is just put here for convenience\n TA-Lib没有明确的提供此算法，这里只是为了调用方便而放在此处\n\n References\n ----------\n https://pola-rs.github.io/polars/py-polars/html/reference/expressions/api/polars.Expr.ewm_mean.html#polars.Expr.ewm_mean\n https://github.com/twopirllc/pandas-ta/blob/main/pandas_ta/overlap/rma.py\n\n \"\"\"\n return close.ewm_mean(alpha=1 / timeperiod, adjust=False, min_samples=timeperiod)\n\n\ndef TEMA(close: Expr, timeperiod: int = 30) -> Expr:\n \"\"\"\n\n Notes\n -----\n todo: if the nesting level is too deep, maybe call talib.TEMA directly\n 嵌套层数过多，也许直接调用talib.TEMA更快\n\n \"\"\"\n EMA1 = EMA(close, timeperiod)\n EMA2 = EMA(EMA1, timeperiod)\n EMA3 = EMA(EMA2, timeperiod)\n return (EMA1 - EMA2) * 3 + EMA3\n\n\ndef TRIMA(close: Expr, timeperiod: int = 30) -> Expr:\n SMA1 = SMA(close, ceil(timeperiod / 2))\n return SMA(SMA1, floor(timeperiod / 2) + 1)\n" }, { "path": "polars_ta/ta/price.py", "content": "from polars import Expr\n\n\ndef AVGPRICE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"(open + high + low + close) / 4\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_AVGPRICE.c#L187\n\n \"\"\"\n return (open + high + low + close) / 4\n\n\ndef MEDPRICE(high: Expr, low: Expr) -> Expr:\n \"\"\"(high + low) / 2\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_MEDPRICE.c#L180\n\n \"\"\"\n return (high + low) / 2\n\n\ndef TYPPRICE(high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"(high + low + close) / 3\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_TYPPRICE.c#L185\n\n \"\"\"\n return (high + low + close) / 3\n\n\ndef WCLPRICE(high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"(high + low + close * 2) / 4\n\n References\n ----------\n https://github.com/TA-Lib/ta-lib/blob/main/src/ta_func/ta_WCLPRICE.c#L184\n\n \"\"\"\n return (high + low + close * 2) / 4\n" }, { "path": "polars_ta/ta/statistic.py", "content": "\"\"\"\n通过`import`直接导入或更名的函数\n\n```python\nfrom polars_ta.wq.time_series import ts_corr as CORREL # noqa\n```\n\n\"\"\"\n\nfrom polars import Expr\n\nfrom polars_ta.wq.time_series import ts_corr as CORREL # noqa\nfrom polars_ta.wq.time_series import ts_std_dev\n\n\ndef BETA(high: Expr, low: Expr, timeperiod: int = 5) -> Expr:\n raise\n\n\ndef LINEARREG(close: Expr, timeperiod: int = 14) -> Expr:\n raise\n\n\ndef LINEARREG_ANGLE(close: Expr, timeperiod: int = 14) -> Expr:\n raise\n\n\ndef LINEARREG_INTERCEPT(close: Expr, timeperiod: int = 14) -> Expr:\n raise\n\n\ndef LINEARREG_SLOPE(close: Expr, timeperiod: int = 14) -> Expr:\n raise\n\n\ndef STDDEV(close: Expr, timeperiod: int = 5, nbdev: float = 1) -> Expr:\n return ts_std_dev(close, timeperiod, ddof=0) * nbdev\n\n\ndef TSF(close: Expr, timeperiod: int = 14) -> Expr:\n raise\n\n\ndef VAR(close: Expr, timeperiod: int = 5, nbdev: float = 1) -> Expr:\n return close.rolling_var(timeperiod, ddof=0) * nbdev\n" }, { "path": "polars_ta/ta/transform.py", "content": "\"\"\"\n通过`import`直接导入或更名的函数\n\n```python\nfrom polars_ta.wq.arithmetic import arc_cos as ACOS # noqa\nfrom polars_ta.wq.arithmetic import arc_sin as ASIN # noqa\nfrom polars_ta.wq.arithmetic import arc_tan as ATAN # noqa\nfrom polars_ta.wq.arithmetic import ceiling as CEIL # noqa\nfrom polars_ta.wq.arithmetic import cos as COS # noqa\nfrom polars_ta.wq.arithmetic import cosh as COSH # noqa\nfrom polars_ta.wq.arithmetic import exp as EXP # noqa\nfrom polars_ta.wq.arithmetic import floor as FLOOR # noqa\nfrom polars_ta.wq.arithmetic import log as LN # noqa\nfrom polars_ta.wq.arithmetic import log10 as LOG10 # noqa\nfrom polars_ta.wq.arithmetic import sin as SIN # noqa\nfrom polars_ta.wq.arithmetic import sinh as SINH # noqa\nfrom polars_ta.wq.arithmetic import sqrt as SQRT # noqa\nfrom polars_ta.wq.arithmetic import tan as TAN # noqa\nfrom polars_ta.wq.arithmetic import tanh as TANH # noqa\n```\n\n\"\"\"\nfrom polars_ta.wq.arithmetic import arc_cos as ACOS # noqa\nfrom polars_ta.wq.arithmetic import arc_sin as ASIN # noqa\nfrom polars_ta.wq.arithmetic import arc_tan as ATAN # noqa\nfrom polars_ta.wq.arithmetic import ceiling as CEIL # noqa\nfrom polars_ta.wq.arithmetic import cos as COS # noqa\nfrom polars_ta.wq.arithmetic import cosh as COSH # noqa\nfrom polars_ta.wq.arithmetic import exp as EXP # noqa\nfrom polars_ta.wq.arithmetic import floor as FLOOR # noqa\nfrom polars_ta.wq.arithmetic import log as LN # noqa\nfrom polars_ta.wq.arithmetic import log10 as LOG10 # noqa\nfrom polars_ta.wq.arithmetic import sin as SIN # noqa\nfrom polars_ta.wq.arithmetic import sinh as SINH # noqa\nfrom polars_ta.wq.arithmetic import sqrt as SQRT # noqa\nfrom polars_ta.wq.arithmetic import tan as TAN # noqa\nfrom polars_ta.wq.arithmetic import tanh as TANH # noqa\n" }, { "path": "polars_ta/ta/volatility.py", "content": "from polars import Expr, max_horizontal\n\nfrom polars_ta.ta.overlap import RMA\n\n\ndef TRANGE(high: Expr, low: Expr, close: Expr) -> Expr:\n \"\"\"\n\n Notes\n -----\n the 0-th position is `x` rather than `nan` in talib\n\n 第0位置为max(x,nan,nan)=x,比talib多一个值\n\n \"\"\"\n prev_close = close.shift(1)\n tr1 = high - low\n tr2 = (high - prev_close).abs()\n tr3 = (low - prev_close).abs()\n return max_horizontal(tr1, tr2, tr3)\n\n\ndef ATR(high: Expr, low: Expr, close: Expr, timeperiod: int = 14) -> Expr:\n \"\"\"\"\"\"\n return RMA(TRANGE(high, low, close), timeperiod)\n\n\ndef NATR(high: Expr, low: Expr, close: Expr, timeperiod: int = 14) -> Expr:\n \"\"\"\n\n Notes\n -----\n talib.ATR multiples another 100\n\n talib.ATR版相当于多乘了100\n\n \"\"\"\n return ATR(high, low, close, timeperiod) / close\n" }, { "path": "polars_ta/ta/volume.py", "content": "from polars import Expr, when\n\nfrom polars_ta.ta.overlap import EMA\n\n\ndef AD(high: Expr, low: Expr, close: Expr, volume: Expr) -> Expr:\n ad = when(high != low).then(((close - low) - (high - close)) / (high - low)).otherwise(0)\n return (ad * volume).cum_sum()\n\n\ndef ADOSC(high: Expr, low: Expr, close: Expr, volume: Expr, fastperiod: int = 3, slowperiod: int = 10) -> Expr:\n ad = AD(high, low, close, volume)\n return EMA(ad, fastperiod) - EMA(ad, slowperiod)\n\n\ndef OBV(close: Expr, volume: Expr) -> Expr:\n \"\"\"\"\"\"\n # using volume for the first value will be exactly the same as talib.OBV\n # 第一个值用volume就与talib.OBV完全一样了\n obv = close.diff().sign().fill_null(1) * volume\n return obv.cum_sum()\n" }, { "path": "polars_ta/talib/README.md", "content": "# polars_ta.talib\n\nInside this package, files are generated by `tools.codegen_talib2`.\nIt is a wrapper of `talib` functions, with the following features:\n\n1. Input and output are `Expr` instead of `Series`\n2. ~~Add skipna feature (not efficient, will update when `polars` support backward fill)~~\n\n本包由`tools.codegen_talib2`自动生成，是对`talib`代码的封装。实现了以下功能\n1. 输入输出由`Series`改`Expr`\n2. ~~添加跳过空值功能（效率不高，等`polars`支持反向填充，此部分将更新）~~\n\n" }, { "path": "polars_ta/talib/__init__.py", "content": "# generated by codegen_talib.py\nimport talib as _ta\nfrom polars import Expr, struct, Struct, Field, Float64, Int32\n\nfrom polars_ta.utils.numba_ import batches_i1_o1, batches_i1_o2, batches_i2_o1, batches_i2_o2, struct_to_numpy\n\n\ndef HT_DCPERIOD(close: Expr) -> Expr: # ['real']\n \"\"\"HT_DCPERIOD(ndarray real)\n\nHT_DCPERIOD(real)\n\nHilbert Transform - Dominant Cycle Period (Cycle Indicators)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.HT_DCPERIOD), return_dtype=Float64)\n\n\ndef HT_DCPHASE(close: Expr) -> Expr: # ['real']\n \"\"\"HT_DCPHASE(ndarray real)\n\nHT_DCPHASE(real)\n\nHilbert Transform - Dominant Cycle Phase (Cycle Indicators)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.HT_DCPHASE), return_dtype=Float64)\n\n\ndef HT_PHASOR(close: Expr) -> Expr: # ['inphase', 'quadrature']\n \"\"\"HT_PHASOR(ndarray real)\n\nHT_PHASOR(real)\n\nHilbert Transform - Phasor Components (Cycle Indicators)\n\nInputs:\n real: (any ndarray)\nOutputs:\n inphase\n quadrature\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return close.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), _ta.HT_PHASOR), return_dtype=dtype)\n\n\ndef HT_SINE(close: Expr) -> Expr: # ['sine', 'leadsine']\n \"\"\"HT_SINE(ndarray real)\n\nHT_SINE(real)\n\nHilbert Transform - SineWave (Cycle Indicators)\n\nInputs:\n real: (any ndarray)\nOutputs:\n sine\n leadsine\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return close.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), _ta.HT_SINE), return_dtype=dtype)\n\n\ndef HT_TRENDMODE(close: Expr) -> Expr: # ['integer']\n \"\"\"HT_TRENDMODE(ndarray real)\n\nHT_TRENDMODE(real)\n\nHilbert Transform - Trend vs Cycle Mode (Cycle Indicators)\n\nInputs:\n real: (any ndarray)\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.HT_TRENDMODE), return_dtype=Int32)\n\n\ndef ADD(high: Expr, low: Expr) -> Expr: # ['real']\n \"\"\"ADD(ndarray real0, ndarray real1)\n\nADD(real0, real1)\n\nVector Arithmetic Add (Math Operators)\n\nInputs:\n real0: (any ndarray)\n real1: (any ndarray)\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.ADD), return_dtype=Float64)\n\n\ndef DIV(high: Expr, low: Expr) -> Expr: # ['real']\n \"\"\"DIV(ndarray real0, ndarray real1)\n\nDIV(real0, real1)\n\nVector Arithmetic Div (Math Operators)\n\nInputs:\n real0: (any ndarray)\n real1: (any ndarray)\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.DIV), return_dtype=Float64)\n\n\ndef MAX(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"MAX(ndarray real, int timeperiod=-0x80000000)\n\nMAX(real[, timeperiod=?])\n\nHighest value over a specified period (Math Operators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.MAX, timeperiod), return_dtype=Float64)\n\n\ndef MAXINDEX(close: Expr, timeperiod: float = 30.0) -> Expr: # ['integer']\n \"\"\"MAXINDEX(ndarray real, int timeperiod=-0x80000000)\n\nMAXINDEX(real[, timeperiod=?])\n\nIndex of highest value over a specified period (Math Operators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.MAXINDEX, timeperiod), return_dtype=Int32)\n\n\ndef MIN(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"MIN(ndarray real, int timeperiod=-0x80000000)\n\nMIN(real[, timeperiod=?])\n\nLowest value over a specified period (Math Operators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.MIN, timeperiod), return_dtype=Float64)\n\n\ndef MININDEX(close: Expr, timeperiod: float = 30.0) -> Expr: # ['integer']\n \"\"\"MININDEX(ndarray real, int timeperiod=-0x80000000)\n\nMININDEX(real[, timeperiod=?])\n\nIndex of lowest value over a specified period (Math Operators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.MININDEX, timeperiod), return_dtype=Int32)\n\n\ndef MINMAX(close: Expr, timeperiod: float = 30.0) -> Expr: # ['min', 'max']\n \"\"\"MINMAX(ndarray real, int timeperiod=-0x80000000)\n\nMINMAX(real[, timeperiod=?])\n\nLowest and highest values over a specified period (Math Operators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n min\n max\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return close.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), _ta.MINMAX, timeperiod), return_dtype=dtype)\n\n\ndef MINMAXINDEX(close: Expr, timeperiod: float = 30.0) -> Expr: # ['minidx', 'maxidx']\n \"\"\"MINMAXINDEX(ndarray real, int timeperiod=-0x80000000)\n\nMINMAXINDEX(real[, timeperiod=?])\n\nIndexes of lowest and highest values over a specified period (Math Operators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n minidx\n maxidx\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return close.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), _ta.MINMAXINDEX, timeperiod), return_dtype=dtype)\n\n\ndef MULT(high: Expr, low: Expr) -> Expr: # ['real']\n \"\"\"MULT(ndarray real0, ndarray real1)\n\nMULT(real0, real1)\n\nVector Arithmetic Mult (Math Operators)\n\nInputs:\n real0: (any ndarray)\n real1: (any ndarray)\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.MULT), return_dtype=Float64)\n\n\ndef SUB(high: Expr, low: Expr) -> Expr: # ['real']\n \"\"\"SUB(ndarray real0, ndarray real1)\n\nSUB(real0, real1)\n\nVector Arithmetic Subtraction (Math Operators)\n\nInputs:\n real0: (any ndarray)\n real1: (any ndarray)\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.SUB), return_dtype=Float64)\n\n\ndef SUM(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"SUM(ndarray real, int timeperiod=-0x80000000)\n\nSUM(real[, timeperiod=?])\n\nSummation (Math Operators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.SUM, timeperiod), return_dtype=Float64)\n\n\ndef ACOS(close: Expr) -> Expr: # ['real']\n \"\"\"ACOS(ndarray real)\n\nACOS(real)\n\nVector Trigonometric ACos (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.ACOS), return_dtype=Float64)\n\n\ndef ASIN(close: Expr) -> Expr: # ['real']\n \"\"\"ASIN(ndarray real)\n\nASIN(real)\n\nVector Trigonometric ASin (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.ASIN), return_dtype=Float64)\n\n\ndef ATAN(close: Expr) -> Expr: # ['real']\n \"\"\"ATAN(ndarray real)\n\nATAN(real)\n\nVector Trigonometric ATan (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.ATAN), return_dtype=Float64)\n\n\ndef CEIL(close: Expr) -> Expr: # ['real']\n \"\"\"CEIL(ndarray real)\n\nCEIL(real)\n\nVector Ceil (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.CEIL), return_dtype=Float64)\n\n\ndef COS(close: Expr) -> Expr: # ['real']\n \"\"\"COS(ndarray real)\n\nCOS(real)\n\nVector Trigonometric Cos (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.COS), return_dtype=Float64)\n\n\ndef COSH(close: Expr) -> Expr: # ['real']\n \"\"\"COSH(ndarray real)\n\nCOSH(real)\n\nVector Trigonometric Cosh (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.COSH), return_dtype=Float64)\n\n\ndef EXP(close: Expr) -> Expr: # ['real']\n \"\"\"EXP(ndarray real)\n\nEXP(real)\n\nVector Arithmetic Exp (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.EXP), return_dtype=Float64)\n\n\ndef FLOOR(close: Expr) -> Expr: # ['real']\n \"\"\"FLOOR(ndarray real)\n\nFLOOR(real)\n\nVector Floor (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.FLOOR), return_dtype=Float64)\n\n\ndef LN(close: Expr) -> Expr: # ['real']\n \"\"\"LN(ndarray real)\n\nLN(real)\n\nVector Log Natural (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.LN), return_dtype=Float64)\n\n\ndef LOG10(close: Expr) -> Expr: # ['real']\n \"\"\"LOG10(ndarray real)\n\nLOG10(real)\n\nVector Log10 (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.LOG10), return_dtype=Float64)\n\n\ndef SIN(close: Expr) -> Expr: # ['real']\n \"\"\"SIN(ndarray real)\n\nSIN(real)\n\nVector Trigonometric Sin (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.SIN), return_dtype=Float64)\n\n\ndef SINH(close: Expr) -> Expr: # ['real']\n \"\"\"SINH(ndarray real)\n\nSINH(real)\n\nVector Trigonometric Sinh (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.SINH), return_dtype=Float64)\n\n\ndef SQRT(close: Expr) -> Expr: # ['real']\n \"\"\"SQRT(ndarray real)\n\nSQRT(real)\n\nVector Square Root (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.SQRT), return_dtype=Float64)\n\n\ndef TAN(close: Expr) -> Expr: # ['real']\n \"\"\"TAN(ndarray real)\n\nTAN(real)\n\nVector Trigonometric Tan (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.TAN), return_dtype=Float64)\n\n\ndef TANH(close: Expr) -> Expr: # ['real']\n \"\"\"TANH(ndarray real)\n\nTANH(real)\n\nVector Trigonometric Tanh (Math Transform)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.TANH), return_dtype=Float64)\n\n\ndef ADX(high: Expr, low: Expr, close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"ADX(ndarray high, ndarray low, ndarray close, int timeperiod=-0x80000000)\n\nADX(high, low, close[, timeperiod=?])\n\nAverage Directional Movement Index (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.ADX, timeperiod), return_dtype=Float64)\n\n\ndef ADXR(high: Expr, low: Expr, close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"ADXR(ndarray high, ndarray low, ndarray close, int timeperiod=-0x80000000)\n\nADXR(high, low, close[, timeperiod=?])\n\nAverage Directional Movement Index Rating (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.ADXR, timeperiod), return_dtype=Float64)\n\n\ndef APO(close: Expr, fastperiod: float = 12.0, slowperiod: float = 26.0, matype: float = 0.0) -> Expr: # ['real']\n \"\"\"APO(ndarray real, int fastperiod=-0x80000000, int slowperiod=-0x80000000, int matype=0)\n\nAPO(real[, fastperiod=?, slowperiod=?, matype=?])\n\nAbsolute Price Oscillator (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n fastperiod: 12\n slowperiod: 26\n matype: 0 (Simple Moving Average)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.APO, fastperiod, slowperiod, matype), return_dtype=Float64)\n\n\ndef AROON(high: Expr, low: Expr, timeperiod: float = 14.0) -> Expr: # ['aroondown', 'aroonup']\n \"\"\"AROON(ndarray high, ndarray low, int timeperiod=-0x80000000)\n\nAROON(high, low[, timeperiod=?])\n\nAroon (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low']\nParameters:\n timeperiod: 14\nOutputs:\n aroondown\n aroonup\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o2(struct_to_numpy(xx, 2, dtype=float), _ta.AROON, timeperiod), return_dtype=dtype)\n\n\ndef AROONOSC(high: Expr, low: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"AROONOSC(ndarray high, ndarray low, int timeperiod=-0x80000000)\n\nAROONOSC(high, low[, timeperiod=?])\n\nAroon Oscillator (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.AROONOSC, timeperiod), return_dtype=Float64)\n\n\ndef BOP(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['real']\n \"\"\"BOP(ndarray open, ndarray high, ndarray low, ndarray close)\n\nBOP(open, high, low, close)\n\nBalance Of Power (Momentum Indicators)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n real\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.BOP), return_dtype=Float64)\n\n\ndef CCI(high: Expr, low: Expr, close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"CCI(ndarray high, ndarray low, ndarray close, int timeperiod=-0x80000000)\n\nCCI(high, low, close[, timeperiod=?])\n\nCommodity Channel Index (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.CCI, timeperiod), return_dtype=Float64)\n\n\ndef CMO(close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"CMO(ndarray real, int timeperiod=-0x80000000)\n\nCMO(real[, timeperiod=?])\n\nChande Momentum Oscillator (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.CMO, timeperiod), return_dtype=Float64)\n\n\ndef DX(high: Expr, low: Expr, close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"DX(ndarray high, ndarray low, ndarray close, int timeperiod=-0x80000000)\n\nDX(high, low, close[, timeperiod=?])\n\nDirectional Movement Index (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.DX, timeperiod), return_dtype=Float64)\n\n\ndef MACD(close: Expr, fastperiod: float = 12.0, slowperiod: float = 26.0, signalperiod: float = 9.0) -> Expr: # ['macd', 'macdsignal', 'macdhist']\n \"\"\"MACD(ndarray real, int fastperiod=-0x80000000, int slowperiod=-0x80000000, int signalperiod=-0x80000000)\n\nMACD(real[, fastperiod=?, slowperiod=?, signalperiod=?])\n\nMoving Average Convergence/Divergence (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n fastperiod: 12\n slowperiod: 26\n signalperiod: 9\nOutputs:\n macd\n macdsignal\n macdhist\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(3)])\n return close.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), _ta.MACD, fastperiod, slowperiod, signalperiod), return_dtype=dtype)\n\n\ndef MACDEXT(close: Expr, fastperiod: float = 12.0, fastmatype: float = 0.0, slowperiod: float = 26.0, slowmatype: float = 0.0, signalperiod: float = 9.0, signalmatype: float = 0.0) -> Expr: # ['macd', 'macdsignal', 'macdhist']\n \"\"\"MACDEXT(ndarray real, int fastperiod=-0x80000000, int fastmatype=0, int slowperiod=-0x80000000, int slowmatype=0, int signalperiod=-0x80000000, int signalmatype=0)\n\nMACDEXT(real[, fastperiod=?, fastmatype=?, slowperiod=?, slowmatype=?, signalperiod=?, signalmatype=?])\n\nMACD with controllable MA type (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n fastperiod: 12\n fastmatype: 0\n slowperiod: 26\n slowmatype: 0\n signalperiod: 9\n signalmatype: 0\nOutputs:\n macd\n macdsignal\n macdhist\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(3)])\n return close.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), _ta.MACDEXT, fastperiod, fastmatype, slowperiod, slowmatype, signalperiod, signalmatype), return_dtype=dtype)\n\n\ndef MACDFIX(close: Expr, signalperiod: float = 9.0) -> Expr: # ['macd', 'macdsignal', 'macdhist']\n \"\"\"MACDFIX(ndarray real, int signalperiod=-0x80000000)\n\nMACDFIX(real[, signalperiod=?])\n\nMoving Average Convergence/Divergence Fix 12/26 (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n signalperiod: 9\nOutputs:\n macd\n macdsignal\n macdhist\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(3)])\n return close.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), _ta.MACDFIX, signalperiod), return_dtype=dtype)\n\n\ndef MFI(high: Expr, low: Expr, close: Expr, volume: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"MFI(ndarray high, ndarray low, ndarray close, ndarray volume, int timeperiod=-0x80000000)\n\nMFI(high, low, close, volume[, timeperiod=?])\n\nMoney Flow Index (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close', 'volume']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close, f3=volume).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.MFI, timeperiod), return_dtype=Float64)\n\n\ndef MINUS_DI(high: Expr, low: Expr, close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"MINUS_DI(ndarray high, ndarray low, ndarray close, int timeperiod=-0x80000000)\n\nMINUS_DI(high, low, close[, timeperiod=?])\n\nMinus Directional Indicator (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.MINUS_DI, timeperiod), return_dtype=Float64)\n\n\ndef MINUS_DM(high: Expr, low: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"MINUS_DM(ndarray high, ndarray low, int timeperiod=-0x80000000)\n\nMINUS_DM(high, low[, timeperiod=?])\n\nMinus Directional Movement (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.MINUS_DM, timeperiod), return_dtype=Float64)\n\n\ndef MOM(close: Expr, timeperiod: float = 10.0) -> Expr: # ['real']\n \"\"\"MOM(ndarray real, int timeperiod=-0x80000000)\n\nMOM(real[, timeperiod=?])\n\nMomentum (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 10\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.MOM, timeperiod), return_dtype=Float64)\n\n\ndef PLUS_DI(high: Expr, low: Expr, close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"PLUS_DI(ndarray high, ndarray low, ndarray close, int timeperiod=-0x80000000)\n\nPLUS_DI(high, low, close[, timeperiod=?])\n\nPlus Directional Indicator (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.PLUS_DI, timeperiod), return_dtype=Float64)\n\n\ndef PLUS_DM(high: Expr, low: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"PLUS_DM(ndarray high, ndarray low, int timeperiod=-0x80000000)\n\nPLUS_DM(high, low[, timeperiod=?])\n\nPlus Directional Movement (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.PLUS_DM, timeperiod), return_dtype=Float64)\n\n\ndef PPO(close: Expr, fastperiod: float = 12.0, slowperiod: float = 26.0, matype: float = 0.0) -> Expr: # ['real']\n \"\"\"PPO(ndarray real, int fastperiod=-0x80000000, int slowperiod=-0x80000000, int matype=0)\n\nPPO(real[, fastperiod=?, slowperiod=?, matype=?])\n\nPercentage Price Oscillator (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n fastperiod: 12\n slowperiod: 26\n matype: 0 (Simple Moving Average)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.PPO, fastperiod, slowperiod, matype), return_dtype=Float64)\n\n\ndef ROC(close: Expr, timeperiod: float = 10.0) -> Expr: # ['real']\n \"\"\"ROC(ndarray real, int timeperiod=-0x80000000)\n\nROC(real[, timeperiod=?])\n\nRate of change : ((real/prevPrice)-1)*100 (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 10\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.ROC, timeperiod), return_dtype=Float64)\n\n\ndef ROCP(close: Expr, timeperiod: float = 10.0) -> Expr: # ['real']\n \"\"\"ROCP(ndarray real, int timeperiod=-0x80000000)\n\nROCP(real[, timeperiod=?])\n\nRate of change Percentage: (real-prevPrice)/prevPrice (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 10\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.ROCP, timeperiod), return_dtype=Float64)\n\n\ndef ROCR(close: Expr, timeperiod: float = 10.0) -> Expr: # ['real']\n \"\"\"ROCR(ndarray real, int timeperiod=-0x80000000)\n\nROCR(real[, timeperiod=?])\n\nRate of change ratio: (real/prevPrice) (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 10\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.ROCR, timeperiod), return_dtype=Float64)\n\n\ndef ROCR100(close: Expr, timeperiod: float = 10.0) -> Expr: # ['real']\n \"\"\"ROCR100(ndarray real, int timeperiod=-0x80000000)\n\nROCR100(real[, timeperiod=?])\n\nRate of change ratio 100 scale: (real/prevPrice)*100 (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 10\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.ROCR100, timeperiod), return_dtype=Float64)\n\n\ndef RSI(close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"RSI(ndarray real, int timeperiod=-0x80000000)\n\nRSI(real[, timeperiod=?])\n\nRelative Strength Index (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.RSI, timeperiod), return_dtype=Float64)\n\n\ndef STOCH(high: Expr, low: Expr, close: Expr, fastk_period: float = 5.0, slowk_period: float = 3.0, slowk_matype: float = 0.0, slowd_period: float = 3.0, slowd_matype: float = 0.0) -> Expr: # ['slowk', 'slowd']\n \"\"\"STOCH(ndarray high, ndarray low, ndarray close, int fastk_period=-0x80000000, int slowk_period=-0x80000000, int slowk_matype=0, int slowd_period=-0x80000000, int slowd_matype=0)\n\nSTOCH(high, low, close[, fastk_period=?, slowk_period=?, slowk_matype=?, slowd_period=?, slowd_matype=?])\n\nStochastic (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n fastk_period: 5\n slowk_period: 3\n slowk_matype: 0\n slowd_period: 3\n slowd_matype: 0\nOutputs:\n slowk\n slowd\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o2(struct_to_numpy(xx, 3, dtype=float), _ta.STOCH, fastk_period, slowk_period, slowk_matype, slowd_period, slowd_matype), return_dtype=dtype)\n\n\ndef STOCHF(high: Expr, low: Expr, close: Expr, fastk_period: float = 5.0, fastd_period: float = 3.0, fastd_matype: float = 0.0) -> Expr: # ['fastk', 'fastd']\n \"\"\"STOCHF(ndarray high, ndarray low, ndarray close, int fastk_period=-0x80000000, int fastd_period=-0x80000000, int fastd_matype=0)\n\nSTOCHF(high, low, close[, fastk_period=?, fastd_period=?, fastd_matype=?])\n\nStochastic Fast (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n fastk_period: 5\n fastd_period: 3\n fastd_matype: 0\nOutputs:\n fastk\n fastd\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o2(struct_to_numpy(xx, 3, dtype=float), _ta.STOCHF, fastk_period, fastd_period, fastd_matype), return_dtype=dtype)\n\n\ndef STOCHRSI(close: Expr, timeperiod: float = 14.0, fastk_period: float = 5.0, fastd_period: float = 3.0, fastd_matype: float = 0.0) -> Expr: # ['fastk', 'fastd']\n \"\"\"STOCHRSI(ndarray real, int timeperiod=-0x80000000, int fastk_period=-0x80000000, int fastd_period=-0x80000000, int fastd_matype=0)\n\nSTOCHRSI(real[, timeperiod=?, fastk_period=?, fastd_period=?, fastd_matype=?])\n\nStochastic Relative Strength Index (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 14\n fastk_period: 5\n fastd_period: 3\n fastd_matype: 0\nOutputs:\n fastk\n fastd\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return close.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), _ta.STOCHRSI, timeperiod, fastk_period, fastd_period, fastd_matype), return_dtype=dtype)\n\n\ndef TRIX(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"TRIX(ndarray real, int timeperiod=-0x80000000)\n\nTRIX(real[, timeperiod=?])\n\n1-day Rate-Of-Change (ROC) of a Triple Smooth EMA (Momentum Indicators)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.TRIX, timeperiod), return_dtype=Float64)\n\n\ndef ULTOSC(high: Expr, low: Expr, close: Expr, timeperiod1: float = 7.0, timeperiod2: float = 14.0, timeperiod3: float = 28.0) -> Expr: # ['real']\n \"\"\"ULTOSC(ndarray high, ndarray low, ndarray close, int timeperiod1=-0x80000000, int timeperiod2=-0x80000000, int timeperiod3=-0x80000000)\n\nULTOSC(high, low, close[, timeperiod1=?, timeperiod2=?, timeperiod3=?])\n\nUltimate Oscillator (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n timeperiod1: 7\n timeperiod2: 14\n timeperiod3: 28\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.ULTOSC, timeperiod1, timeperiod2, timeperiod3), return_dtype=Float64)\n\n\ndef WILLR(high: Expr, low: Expr, close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"WILLR(ndarray high, ndarray low, ndarray close, int timeperiod=-0x80000000)\n\nWILLR(high, low, close[, timeperiod=?])\n\nWilliams' %R (Momentum Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.WILLR, timeperiod), return_dtype=Float64)\n\n\ndef BBANDS(close: Expr, timeperiod: float = 5.0, nbdevup: float = 2.0, nbdevdn: float = 2.0, matype: float = 0.0) -> Expr: # ['upperband', 'middleband', 'lowerband']\n \"\"\"BBANDS(ndarray real, int timeperiod=-0x80000000, double nbdevup=-4e37, double nbdevdn=-4e37, int matype=0)\n\nBBANDS(real[, timeperiod=?, nbdevup=?, nbdevdn=?, matype=?])\n\nBollinger Bands (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 5\n nbdevup: 2.0\n nbdevdn: 2.0\n matype: 0 (Simple Moving Average)\nOutputs:\n upperband\n middleband\n lowerband\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(3)])\n return close.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), _ta.BBANDS, timeperiod, nbdevup, nbdevdn, matype), return_dtype=dtype)\n\n\ndef DEMA(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"DEMA(ndarray real, int timeperiod=-0x80000000)\n\nDEMA(real[, timeperiod=?])\n\nDouble Exponential Moving Average (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.DEMA, timeperiod), return_dtype=Float64)\n\n\ndef EMA(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"EMA(ndarray real, int timeperiod=-0x80000000)\n\nEMA(real[, timeperiod=?])\n\nExponential Moving Average (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.EMA, timeperiod), return_dtype=Float64)\n\n\ndef HT_TRENDLINE(close: Expr) -> Expr: # ['real']\n \"\"\"HT_TRENDLINE(ndarray real)\n\nHT_TRENDLINE(real)\n\nHilbert Transform - Instantaneous Trendline (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.HT_TRENDLINE), return_dtype=Float64)\n\n\ndef KAMA(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"KAMA(ndarray real, int timeperiod=-0x80000000)\n\nKAMA(real[, timeperiod=?])\n\nKaufman Adaptive Moving Average (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.KAMA, timeperiod), return_dtype=Float64)\n\n\ndef MA(close: Expr, timeperiod: float = 30.0, matype: float = 0.0) -> Expr: # ['real']\n \"\"\"MA(ndarray real, int timeperiod=-0x80000000, int matype=0)\n\nMA(real[, timeperiod=?, matype=?])\n\nMoving average (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\n matype: 0 (Simple Moving Average)\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.MA, timeperiod, matype), return_dtype=Float64)\n\n\ndef MAMA(close: Expr, fastlimit: float = 0.5, slowlimit: float = 0.05) -> Expr: # ['mama', 'fama']\n \"\"\"MAMA(ndarray real, double fastlimit=-4e37, double slowlimit=-4e37)\n\nMAMA(real[, fastlimit=?, slowlimit=?])\n\nMESA Adaptive Moving Average (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n fastlimit: 0.5\n slowlimit: 0.05\nOutputs:\n mama\n fama\"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return close.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), _ta.MAMA, fastlimit, slowlimit), return_dtype=dtype)\n\n\ndef MAVP(close: Expr, periods: Expr, minperiod: float = 2.0, maxperiod: float = 30.0, matype: float = 0.0) -> Expr: # ['real']\n \"\"\"MAVP(ndarray real, ndarray periods, int minperiod=-0x80000000, int maxperiod=-0x80000000, int matype=0)\n\nMAVP(real, periods[, minperiod=?, maxperiod=?, matype=?])\n\nMoving average with variable period (Overlap Studies)\n\nInputs:\n real: (any ndarray)\n periods: (any ndarray)\nParameters:\n minperiod: 2\n maxperiod: 30\n matype: 0 (Simple Moving Average)\nOutputs:\n real\"\"\"\n return struct(f0=close, f1=periods).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.MAVP, minperiod, maxperiod, matype), return_dtype=Float64)\n\n\ndef MIDPOINT(close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"MIDPOINT(ndarray real, int timeperiod=-0x80000000)\n\nMIDPOINT(real[, timeperiod=?])\n\nMidPoint over period (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.MIDPOINT, timeperiod), return_dtype=Float64)\n\n\ndef MIDPRICE(high: Expr, low: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"MIDPRICE(ndarray high, ndarray low, int timeperiod=-0x80000000)\n\nMIDPRICE(high, low[, timeperiod=?])\n\nMidpoint Price over period (Overlap Studies)\n\nInputs:\n prices: ['high', 'low']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.MIDPRICE, timeperiod), return_dtype=Float64)\n\n\ndef SAR(high: Expr, low: Expr, acceleration: float = 0.02, maximum: float = 0.2) -> Expr: # ['real']\n \"\"\"SAR(ndarray high, ndarray low, double acceleration=0.02, double maximum=0.2)\n\nSAR(high, low[, acceleration=?, maximum=?])\n\nParabolic SAR (Overlap Studies)\n\nInputs:\n prices: ['high', 'low']\nParameters:\n acceleration: 0.02\n maximum: 0.2\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.SAR, acceleration, maximum), return_dtype=Float64)\n\n\ndef SAREXT(high: Expr, low: Expr, startvalue: float = 0.0, offsetonreverse: float = 0.0, accelerationinitlong: float = 0.02, accelerationlong: float = 0.02, accelerationmaxlong: float = 0.2, accelerationinitshort: float = 0.02, accelerationshort: float = 0.02, accelerationmaxshort: float = 0.2) -> Expr: # ['real']\n \"\"\"SAREXT(ndarray high, ndarray low, double startvalue=-4e37, double offsetonreverse=-4e37, double accelerationinitlong=-4e37, double accelerationlong=-4e37, double accelerationmaxlong=-4e37, double accelerationinitshort=-4e37, double accelerationshort=-4e37, double accelerationmaxshort=-4e37)\n\nSAREXT(high, low[, startvalue=?, offsetonreverse=?, accelerationinitlong=?, accelerationlong=?, accelerationmaxlong=?, accelerationinitshort=?, accelerationshort=?, accelerationmaxshort=?])\n\nParabolic SAR - Extended (Overlap Studies)\n\nInputs:\n prices: ['high', 'low']\nParameters:\n startvalue: 0.0\n offsetonreverse: 0.0\n accelerationinitlong: 0.02\n accelerationlong: 0.02\n accelerationmaxlong: 0.2\n accelerationinitshort: 0.02\n accelerationshort: 0.02\n accelerationmaxshort: 0.2\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.SAREXT, startvalue, offsetonreverse, accelerationinitlong, accelerationlong, accelerationmaxlong, accelerationinitshort, accelerationshort, accelerationmaxshort), return_dtype=Float64)\n\n\ndef SMA(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"SMA(ndarray real, int timeperiod=-0x80000000)\n\nSMA(real[, timeperiod=?])\n\nSimple Moving Average (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.SMA, timeperiod), return_dtype=Float64)\n\n\ndef T3(close: Expr, timeperiod: float = 5.0, vfactor: float = 0.7) -> Expr: # ['real']\n \"\"\"T3(ndarray real, int timeperiod=-0x80000000, double vfactor=-4e37)\n\nT3(real[, timeperiod=?, vfactor=?])\n\nTriple Exponential Moving Average (T3) (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 5\n vfactor: 0.7\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.T3, timeperiod, vfactor), return_dtype=Float64)\n\n\ndef TEMA(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"TEMA(ndarray real, int timeperiod=-0x80000000)\n\nTEMA(real[, timeperiod=?])\n\nTriple Exponential Moving Average (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.TEMA, timeperiod), return_dtype=Float64)\n\n\ndef TRIMA(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"TRIMA(ndarray real, int timeperiod=-0x80000000)\n\nTRIMA(real[, timeperiod=?])\n\nTriangular Moving Average (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.TRIMA, timeperiod), return_dtype=Float64)\n\n\ndef WMA(close: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"WMA(ndarray real, int timeperiod=-0x80000000)\n\nWMA(real[, timeperiod=?])\n\nWeighted Moving Average (Overlap Studies)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.WMA, timeperiod), return_dtype=Float64)\n\n\ndef CDL2CROWS(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDL2CROWS(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDL2CROWS(open, high, low, close)\n\nTwo Crows (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDL2CROWS), return_dtype=Int32)\n\n\ndef CDL3BLACKCROWS(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDL3BLACKCROWS(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDL3BLACKCROWS(open, high, low, close)\n\nThree Black Crows (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDL3BLACKCROWS), return_dtype=Int32)\n\n\ndef CDL3INSIDE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDL3INSIDE(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDL3INSIDE(open, high, low, close)\n\nThree Inside Up/Down (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDL3INSIDE), return_dtype=Int32)\n\n\ndef CDL3LINESTRIKE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDL3LINESTRIKE(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDL3LINESTRIKE(open, high, low, close)\n\nThree-Line Strike (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDL3LINESTRIKE), return_dtype=Int32)\n\n\ndef CDL3OUTSIDE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDL3OUTSIDE(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDL3OUTSIDE(open, high, low, close)\n\nThree Outside Up/Down (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDL3OUTSIDE), return_dtype=Int32)\n\n\ndef CDL3STARSINSOUTH(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDL3STARSINSOUTH(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDL3STARSINSOUTH(open, high, low, close)\n\nThree Stars In The South (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDL3STARSINSOUTH), return_dtype=Int32)\n\n\ndef CDL3WHITESOLDIERS(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDL3WHITESOLDIERS(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDL3WHITESOLDIERS(open, high, low, close)\n\nThree Advancing White Soldiers (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDL3WHITESOLDIERS), return_dtype=Int32)\n\n\ndef CDLABANDONEDBABY(open: Expr, high: Expr, low: Expr, close: Expr, penetration: float = 0.3) -> Expr: # ['integer']\n \"\"\"CDLABANDONEDBABY(ndarray open, ndarray high, ndarray low, ndarray close, double penetration=0.3)\n\nCDLABANDONEDBABY(open, high, low, close[, penetration=?])\n\nAbandoned Baby (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nParameters:\n penetration: 0.3\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLABANDONEDBABY, penetration), return_dtype=Int32)\n\n\ndef CDLADVANCEBLOCK(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLADVANCEBLOCK(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLADVANCEBLOCK(open, high, low, close)\n\nAdvance Block (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLADVANCEBLOCK), return_dtype=Int32)\n\n\ndef CDLBELTHOLD(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLBELTHOLD(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLBELTHOLD(open, high, low, close)\n\nBelt-hold (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLBELTHOLD), return_dtype=Int32)\n\n\ndef CDLBREAKAWAY(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLBREAKAWAY(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLBREAKAWAY(open, high, low, close)\n\nBreakaway (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLBREAKAWAY), return_dtype=Int32)\n\n\ndef CDLCLOSINGMARUBOZU(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLCLOSINGMARUBOZU(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLCLOSINGMARUBOZU(open, high, low, close)\n\nClosing Marubozu (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLCLOSINGMARUBOZU), return_dtype=Int32)\n\n\ndef CDLCONCEALBABYSWALL(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLCONCEALBABYSWALL(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLCONCEALBABYSWALL(open, high, low, close)\n\nConcealing Baby Swallow (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLCONCEALBABYSWALL), return_dtype=Int32)\n\n\ndef CDLCOUNTERATTACK(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLCOUNTERATTACK(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLCOUNTERATTACK(open, high, low, close)\n\nCounterattack (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLCOUNTERATTACK), return_dtype=Int32)\n\n\ndef CDLDARKCLOUDCOVER(open: Expr, high: Expr, low: Expr, close: Expr, penetration: float = 0.5) -> Expr: # ['integer']\n \"\"\"CDLDARKCLOUDCOVER(ndarray open, ndarray high, ndarray low, ndarray close, double penetration=0.5)\n\nCDLDARKCLOUDCOVER(open, high, low, close[, penetration=?])\n\nDark Cloud Cover (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nParameters:\n penetration: 0.5\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLDARKCLOUDCOVER, penetration), return_dtype=Int32)\n\n\ndef CDLDOJI(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLDOJI(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLDOJI(open, high, low, close)\n\nDoji (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLDOJI), return_dtype=Int32)\n\n\ndef CDLDOJISTAR(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLDOJISTAR(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLDOJISTAR(open, high, low, close)\n\nDoji Star (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLDOJISTAR), return_dtype=Int32)\n\n\ndef CDLDRAGONFLYDOJI(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLDRAGONFLYDOJI(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLDRAGONFLYDOJI(open, high, low, close)\n\nDragonfly Doji (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLDRAGONFLYDOJI), return_dtype=Int32)\n\n\ndef CDLENGULFING(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLENGULFING(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLENGULFING(open, high, low, close)\n\nEngulfing Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLENGULFING), return_dtype=Int32)\n\n\ndef CDLEVENINGDOJISTAR(open: Expr, high: Expr, low: Expr, close: Expr, penetration: float = 0.3) -> Expr: # ['integer']\n \"\"\"CDLEVENINGDOJISTAR(ndarray open, ndarray high, ndarray low, ndarray close, double penetration=0.3)\n\nCDLEVENINGDOJISTAR(open, high, low, close[, penetration=?])\n\nEvening Doji Star (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nParameters:\n penetration: 0.3\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLEVENINGDOJISTAR, penetration), return_dtype=Int32)\n\n\ndef CDLEVENINGSTAR(open: Expr, high: Expr, low: Expr, close: Expr, penetration: float = 0.3) -> Expr: # ['integer']\n \"\"\"CDLEVENINGSTAR(ndarray open, ndarray high, ndarray low, ndarray close, double penetration=0.3)\n\nCDLEVENINGSTAR(open, high, low, close[, penetration=?])\n\nEvening Star (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nParameters:\n penetration: 0.3\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLEVENINGSTAR, penetration), return_dtype=Int32)\n\n\ndef CDLGAPSIDESIDEWHITE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLGAPSIDESIDEWHITE(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLGAPSIDESIDEWHITE(open, high, low, close)\n\nUp/Down-gap side-by-side white lines (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLGAPSIDESIDEWHITE), return_dtype=Int32)\n\n\ndef CDLGRAVESTONEDOJI(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLGRAVESTONEDOJI(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLGRAVESTONEDOJI(open, high, low, close)\n\nGravestone Doji (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLGRAVESTONEDOJI), return_dtype=Int32)\n\n\ndef CDLHAMMER(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLHAMMER(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLHAMMER(open, high, low, close)\n\nHammer (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLHAMMER), return_dtype=Int32)\n\n\ndef CDLHANGINGMAN(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLHANGINGMAN(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLHANGINGMAN(open, high, low, close)\n\nHanging Man (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLHANGINGMAN), return_dtype=Int32)\n\n\ndef CDLHARAMI(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLHARAMI(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLHARAMI(open, high, low, close)\n\nHarami Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLHARAMI), return_dtype=Int32)\n\n\ndef CDLHARAMICROSS(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLHARAMICROSS(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLHARAMICROSS(open, high, low, close)\n\nHarami Cross Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLHARAMICROSS), return_dtype=Int32)\n\n\ndef CDLHIGHWAVE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLHIGHWAVE(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLHIGHWAVE(open, high, low, close)\n\nHigh-Wave Candle (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLHIGHWAVE), return_dtype=Int32)\n\n\ndef CDLHIKKAKE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLHIKKAKE(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLHIKKAKE(open, high, low, close)\n\nHikkake Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLHIKKAKE), return_dtype=Int32)\n\n\ndef CDLHIKKAKEMOD(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLHIKKAKEMOD(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLHIKKAKEMOD(open, high, low, close)\n\nModified Hikkake Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLHIKKAKEMOD), return_dtype=Int32)\n\n\ndef CDLHOMINGPIGEON(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLHOMINGPIGEON(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLHOMINGPIGEON(open, high, low, close)\n\nHoming Pigeon (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLHOMINGPIGEON), return_dtype=Int32)\n\n\ndef CDLIDENTICAL3CROWS(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLIDENTICAL3CROWS(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLIDENTICAL3CROWS(open, high, low, close)\n\nIdentical Three Crows (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLIDENTICAL3CROWS), return_dtype=Int32)\n\n\ndef CDLINNECK(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLINNECK(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLINNECK(open, high, low, close)\n\nIn-Neck Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLINNECK), return_dtype=Int32)\n\n\ndef CDLINVERTEDHAMMER(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLINVERTEDHAMMER(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLINVERTEDHAMMER(open, high, low, close)\n\nInverted Hammer (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLINVERTEDHAMMER), return_dtype=Int32)\n\n\ndef CDLKICKING(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLKICKING(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLKICKING(open, high, low, close)\n\nKicking (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLKICKING), return_dtype=Int32)\n\n\ndef CDLKICKINGBYLENGTH(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLKICKINGBYLENGTH(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLKICKINGBYLENGTH(open, high, low, close)\n\nKicking - bull/bear determined by the longer marubozu (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLKICKINGBYLENGTH), return_dtype=Int32)\n\n\ndef CDLLADDERBOTTOM(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLLADDERBOTTOM(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLLADDERBOTTOM(open, high, low, close)\n\nLadder Bottom (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLLADDERBOTTOM), return_dtype=Int32)\n\n\ndef CDLLONGLEGGEDDOJI(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLLONGLEGGEDDOJI(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLLONGLEGGEDDOJI(open, high, low, close)\n\nLong Legged Doji (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLLONGLEGGEDDOJI), return_dtype=Int32)\n\n\ndef CDLLONGLINE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLLONGLINE(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLLONGLINE(open, high, low, close)\n\nLong Line Candle (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLLONGLINE), return_dtype=Int32)\n\n\ndef CDLMARUBOZU(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLMARUBOZU(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLMARUBOZU(open, high, low, close)\n\nMarubozu (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLMARUBOZU), return_dtype=Int32)\n\n\ndef CDLMATCHINGLOW(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLMATCHINGLOW(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLMATCHINGLOW(open, high, low, close)\n\nMatching Low (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLMATCHINGLOW), return_dtype=Int32)\n\n\ndef CDLMATHOLD(open: Expr, high: Expr, low: Expr, close: Expr, penetration: float = 0.5) -> Expr: # ['integer']\n \"\"\"CDLMATHOLD(ndarray open, ndarray high, ndarray low, ndarray close, double penetration=0.5)\n\nCDLMATHOLD(open, high, low, close[, penetration=?])\n\nMat Hold (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nParameters:\n penetration: 0.5\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLMATHOLD, penetration), return_dtype=Int32)\n\n\ndef CDLMORNINGDOJISTAR(open: Expr, high: Expr, low: Expr, close: Expr, penetration: float = 0.3) -> Expr: # ['integer']\n \"\"\"CDLMORNINGDOJISTAR(ndarray open, ndarray high, ndarray low, ndarray close, double penetration=0.3)\n\nCDLMORNINGDOJISTAR(open, high, low, close[, penetration=?])\n\nMorning Doji Star (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nParameters:\n penetration: 0.3\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLMORNINGDOJISTAR, penetration), return_dtype=Int32)\n\n\ndef CDLMORNINGSTAR(open: Expr, high: Expr, low: Expr, close: Expr, penetration: float = 0.3) -> Expr: # ['integer']\n \"\"\"CDLMORNINGSTAR(ndarray open, ndarray high, ndarray low, ndarray close, double penetration=0.3)\n\nCDLMORNINGSTAR(open, high, low, close[, penetration=?])\n\nMorning Star (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nParameters:\n penetration: 0.3\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLMORNINGSTAR, penetration), return_dtype=Int32)\n\n\ndef CDLONNECK(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLONNECK(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLONNECK(open, high, low, close)\n\nOn-Neck Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLONNECK), return_dtype=Int32)\n\n\ndef CDLPIERCING(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLPIERCING(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLPIERCING(open, high, low, close)\n\nPiercing Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLPIERCING), return_dtype=Int32)\n\n\ndef CDLRICKSHAWMAN(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLRICKSHAWMAN(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLRICKSHAWMAN(open, high, low, close)\n\nRickshaw Man (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLRICKSHAWMAN), return_dtype=Int32)\n\n\ndef CDLRISEFALL3METHODS(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLRISEFALL3METHODS(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLRISEFALL3METHODS(open, high, low, close)\n\nRising/Falling Three Methods (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLRISEFALL3METHODS), return_dtype=Int32)\n\n\ndef CDLSEPARATINGLINES(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLSEPARATINGLINES(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLSEPARATINGLINES(open, high, low, close)\n\nSeparating Lines (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLSEPARATINGLINES), return_dtype=Int32)\n\n\ndef CDLSHOOTINGSTAR(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLSHOOTINGSTAR(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLSHOOTINGSTAR(open, high, low, close)\n\nShooting Star (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLSHOOTINGSTAR), return_dtype=Int32)\n\n\ndef CDLSHORTLINE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLSHORTLINE(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLSHORTLINE(open, high, low, close)\n\nShort Line Candle (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLSHORTLINE), return_dtype=Int32)\n\n\ndef CDLSPINNINGTOP(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLSPINNINGTOP(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLSPINNINGTOP(open, high, low, close)\n\nSpinning Top (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLSPINNINGTOP), return_dtype=Int32)\n\n\ndef CDLSTALLEDPATTERN(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLSTALLEDPATTERN(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLSTALLEDPATTERN(open, high, low, close)\n\nStalled Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLSTALLEDPATTERN), return_dtype=Int32)\n\n\ndef CDLSTICKSANDWICH(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLSTICKSANDWICH(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLSTICKSANDWICH(open, high, low, close)\n\nStick Sandwich (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLSTICKSANDWICH), return_dtype=Int32)\n\n\ndef CDLTAKURI(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLTAKURI(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLTAKURI(open, high, low, close)\n\nTakuri (Dragonfly Doji with very long lower shadow) (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLTAKURI), return_dtype=Int32)\n\n\ndef CDLTASUKIGAP(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLTASUKIGAP(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLTASUKIGAP(open, high, low, close)\n\nTasuki Gap (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLTASUKIGAP), return_dtype=Int32)\n\n\ndef CDLTHRUSTING(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLTHRUSTING(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLTHRUSTING(open, high, low, close)\n\nThrusting Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLTHRUSTING), return_dtype=Int32)\n\n\ndef CDLTRISTAR(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLTRISTAR(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLTRISTAR(open, high, low, close)\n\nTristar Pattern (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLTRISTAR), return_dtype=Int32)\n\n\ndef CDLUNIQUE3RIVER(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLUNIQUE3RIVER(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLUNIQUE3RIVER(open, high, low, close)\n\nUnique 3 River (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLUNIQUE3RIVER), return_dtype=Int32)\n\n\ndef CDLUPSIDEGAP2CROWS(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLUPSIDEGAP2CROWS(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLUPSIDEGAP2CROWS(open, high, low, close)\n\nUpside Gap Two Crows (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLUPSIDEGAP2CROWS), return_dtype=Int32)\n\n\ndef CDLXSIDEGAP3METHODS(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['integer']\n \"\"\"CDLXSIDEGAP3METHODS(ndarray open, ndarray high, ndarray low, ndarray close)\n\nCDLXSIDEGAP3METHODS(open, high, low, close)\n\nUpside/Downside Gap Three Methods (Pattern Recognition)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n integer (values are -100, 0 or 100)\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.CDLXSIDEGAP3METHODS), return_dtype=Int32)\n\n\ndef AVGPRICE(open: Expr, high: Expr, low: Expr, close: Expr) -> Expr: # ['real']\n \"\"\"AVGPRICE(ndarray open, ndarray high, ndarray low, ndarray close)\n\nAVGPRICE(open, high, low, close)\n\nAverage Price (Price Transform)\n\nInputs:\n prices: ['open', 'high', 'low', 'close']\nOutputs:\n real\"\"\"\n return struct(f0=open, f1=high, f2=low, f3=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.AVGPRICE), return_dtype=Float64)\n\n\ndef MEDPRICE(high: Expr, low: Expr) -> Expr: # ['real']\n \"\"\"MEDPRICE(ndarray high, ndarray low)\n\nMEDPRICE(high, low)\n\nMedian Price (Price Transform)\n\nInputs:\n prices: ['high', 'low']\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.MEDPRICE), return_dtype=Float64)\n\n\ndef TYPPRICE(high: Expr, low: Expr, close: Expr) -> Expr: # ['real']\n \"\"\"TYPPRICE(ndarray high, ndarray low, ndarray close)\n\nTYPPRICE(high, low, close)\n\nTypical Price (Price Transform)\n\nInputs:\n prices: ['high', 'low', 'close']\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.TYPPRICE), return_dtype=Float64)\n\n\ndef WCLPRICE(high: Expr, low: Expr, close: Expr) -> Expr: # ['real']\n \"\"\"WCLPRICE(ndarray high, ndarray low, ndarray close)\n\nWCLPRICE(high, low, close)\n\nWeighted Close Price (Price Transform)\n\nInputs:\n prices: ['high', 'low', 'close']\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.WCLPRICE), return_dtype=Float64)\n\n\ndef BETA(high: Expr, low: Expr, timeperiod: float = 5.0) -> Expr: # ['real']\n \"\"\"BETA(ndarray real0, ndarray real1, int timeperiod=-0x80000000)\n\nBETA(real0, real1[, timeperiod=?])\n\nBeta (Statistic Functions)\n\nInputs:\n real0: (any ndarray)\n real1: (any ndarray)\nParameters:\n timeperiod: 5\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.BETA, timeperiod), return_dtype=Float64)\n\n\ndef CORREL(high: Expr, low: Expr, timeperiod: float = 30.0) -> Expr: # ['real']\n \"\"\"CORREL(ndarray real0, ndarray real1, int timeperiod=-0x80000000)\n\nCORREL(real0, real1[, timeperiod=?])\n\nPearson's Correlation Coefficient (r) (Statistic Functions)\n\nInputs:\n real0: (any ndarray)\n real1: (any ndarray)\nParameters:\n timeperiod: 30\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.CORREL, timeperiod), return_dtype=Float64)\n\n\ndef LINEARREG(close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"LINEARREG(ndarray real, int timeperiod=-0x80000000)\n\nLINEARREG(real[, timeperiod=?])\n\nLinear Regression (Statistic Functions)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.LINEARREG, timeperiod), return_dtype=Float64)\n\n\ndef LINEARREG_ANGLE(close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"LINEARREG_ANGLE(ndarray real, int timeperiod=-0x80000000)\n\nLINEARREG_ANGLE(real[, timeperiod=?])\n\nLinear Regression Angle (Statistic Functions)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.LINEARREG_ANGLE, timeperiod), return_dtype=Float64)\n\n\ndef LINEARREG_INTERCEPT(close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"LINEARREG_INTERCEPT(ndarray real, int timeperiod=-0x80000000)\n\nLINEARREG_INTERCEPT(real[, timeperiod=?])\n\nLinear Regression Intercept (Statistic Functions)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.LINEARREG_INTERCEPT, timeperiod), return_dtype=Float64)\n\n\ndef LINEARREG_SLOPE(close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"LINEARREG_SLOPE(ndarray real, int timeperiod=-0x80000000)\n\nLINEARREG_SLOPE(real[, timeperiod=?])\n\nLinear Regression Slope (Statistic Functions)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.LINEARREG_SLOPE, timeperiod), return_dtype=Float64)\n\n\ndef STDDEV(close: Expr, timeperiod: float = 5.0, nbdev: float = 1.0) -> Expr: # ['real']\n \"\"\"STDDEV(ndarray real, int timeperiod=-0x80000000, double nbdev=-4e37)\n\nSTDDEV(real[, timeperiod=?, nbdev=?])\n\nStandard Deviation (Statistic Functions)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 5\n nbdev: 1.0\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.STDDEV, timeperiod, nbdev), return_dtype=Float64)\n\n\ndef TSF(close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"TSF(ndarray real, int timeperiod=-0x80000000)\n\nTSF(real[, timeperiod=?])\n\nTime Series Forecast (Statistic Functions)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.TSF, timeperiod), return_dtype=Float64)\n\n\ndef VAR(close: Expr, timeperiod: float = 5.0, nbdev: float = 1.0) -> Expr: # ['real']\n \"\"\"VAR(ndarray real, int timeperiod=-0x80000000, double nbdev=-4e37)\n\nVAR(real[, timeperiod=?, nbdev=?])\n\nVariance (Statistic Functions)\n\nInputs:\n real: (any ndarray)\nParameters:\n timeperiod: 5\n nbdev: 1.0\nOutputs:\n real\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _ta.VAR, timeperiod, nbdev), return_dtype=Float64)\n\n\ndef ATR(high: Expr, low: Expr, close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"ATR(ndarray high, ndarray low, ndarray close, int timeperiod=-0x80000000)\n\nATR(high, low, close[, timeperiod=?])\n\nAverage True Range (Volatility Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.ATR, timeperiod), return_dtype=Float64)\n\n\ndef NATR(high: Expr, low: Expr, close: Expr, timeperiod: float = 14.0) -> Expr: # ['real']\n \"\"\"NATR(ndarray high, ndarray low, ndarray close, int timeperiod=-0x80000000)\n\nNATR(high, low, close[, timeperiod=?])\n\nNormalized Average True Range (Volatility Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nParameters:\n timeperiod: 14\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.NATR, timeperiod), return_dtype=Float64)\n\n\ndef TRANGE(high: Expr, low: Expr, close: Expr) -> Expr: # ['real']\n \"\"\"TRANGE(ndarray high, ndarray low, ndarray close)\n\nTRANGE(high, low, close)\n\nTrue Range (Volatility Indicators)\n\nInputs:\n prices: ['high', 'low', 'close']\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3, dtype=float), _ta.TRANGE), return_dtype=Float64)\n\n\ndef AD(high: Expr, low: Expr, close: Expr, volume: Expr) -> Expr: # ['real']\n \"\"\"AD(ndarray high, ndarray low, ndarray close, ndarray volume)\n\nAD(high, low, close, volume)\n\nChaikin A/D Line (Volume Indicators)\n\nInputs:\n prices: ['high', 'low', 'close', 'volume']\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close, f3=volume).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.AD), return_dtype=Float64)\n\n\ndef ADOSC(high: Expr, low: Expr, close: Expr, volume: Expr, fastperiod: float = 3.0, slowperiod: float = 10.0) -> Expr: # ['real']\n \"\"\"ADOSC(ndarray high, ndarray low, ndarray close, ndarray volume, int fastperiod=-0x80000000, int slowperiod=-0x80000000)\n\nADOSC(high, low, close, volume[, fastperiod=?, slowperiod=?])\n\nChaikin A/D Oscillator (Volume Indicators)\n\nInputs:\n prices: ['high', 'low', 'close', 'volume']\nParameters:\n fastperiod: 3\n slowperiod: 10\nOutputs:\n real\"\"\"\n return struct(f0=high, f1=low, f2=close, f3=volume).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _ta.ADOSC, fastperiod, slowperiod), return_dtype=Float64)\n\n\ndef OBV(close: Expr, volume: Expr) -> Expr: # ['real']\n \"\"\"OBV(ndarray real, ndarray volume)\n\nOBV(real, volume)\n\nOn Balance Volume (Volume Indicators)\n\nInputs:\n real: (any ndarray)\n prices: ['volume']\nOutputs:\n real\"\"\"\n return struct(f0=close, f1=volume).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2, dtype=float), _ta.OBV), return_dtype=Float64)\n" }, { "path": "polars_ta/tdx/README.md", "content": "# polars_ta.tdx\n\n1. Follows the `tdx` naming convention\n2. Except for some element-wise functions, all functions are time-series functions. Pay special attention to `MAX` and similar functions.\n3. First import from `tdx`, then from `wq`, and finally from `ta`. Only implement the function if it is not available.\n\n\n1. 函数名称按照通达信来\n2. 除了部分按元素计算的函数，其它都为时序函数，特别注意`MAX`等一类不要混淆\n3. 优先从`tdx`中导入，然后从`wq`中导入，最后从`ta`，没有的才实现" }, { "path": "polars_ta/tdx/__init__.py", "content": "from polars_ta.tdx.arithmetic import * # noqa\nfrom polars_ta.tdx.choice import * # noqa\nfrom polars_ta.tdx.energy import * # noqa\nfrom polars_ta.tdx.logical import * # noqa\nfrom polars_ta.tdx.moving_average import * # noqa\nfrom polars_ta.tdx.over_bought_over_sold import * # noqa\nfrom polars_ta.tdx.pattern import * # noqa\nfrom polars_ta.tdx.pattern_feature import * # noqa\nfrom polars_ta.tdx.pressure_support import * # noqa\nfrom polars_ta.tdx.reference import * # noqa\nfrom polars_ta.tdx.statistic import * # noqa\nfrom polars_ta.tdx.times import * # noqa\nfrom polars_ta.tdx.trend import * # noqa\nfrom polars_ta.tdx.trend_feature import * # noqa\nfrom polars_ta.tdx.volume import * # noqa\n" }, { "path": "polars_ta/tdx/_chip.py", "content": "import numba\nimport numpy as np\n\n\n@numba.jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef nb_chip(high, low, avg, turnover,\n start=None, stop=None, step=0.2):\n \"\"\"筹码分布，可用于WINNER或COST指标\n\n 不可能完全还原真实的筹码分布，只能接近。所以做了一下特别处理\n\n 1. 三角分布，比平均分布更接近\n 2. 步长。没有必要每个价格都统计，特别是复权后价格也无法正好是0.01间隔\n 高价股建议步长设大些，低价股步长需设小些\n\n Parameters\n ----------\n high\n low\n avg\n 一维序列\n turnover:\n 换手率，需要在外转成0~1范围内\n start\n 开始价格\n stop\n 结束价格\n step\n 步长。一字涨停时，三角分布的底为1，高为2。但无法当成梯形计算面积，所以从中用半步长切开计算\n\n Returns\n -------\n out\n 筹码分布\n columns\n 价格表头\n\n \"\"\"\n # 网格范围\n if start is None:\n start = np.min(low)\n if stop is None:\n stop = np.max(high)\n\n left = round(start / step) * 2 - 1\n right = round(stop / step) * 2 + 1\n\n # 最小最大值左右要留半格，range是左闭右开，长度必须为2n+1\n columns = np.arange(left, right + 1)\n grid_shape = (len(turnover), len(columns))\n\n # numba中round写法特殊\n _high = np.empty_like(high)\n _low = np.empty_like(low)\n _avg = np.empty_like(avg)\n\n # high和low必须落到边缘上\n _high = np.round(high / step, 0, _high) * 2 + 1\n _low = np.round(low / step, 0, _low) * 2 - 1\n # avg必须落在实体上\n _avg = np.round(avg / step, 0, _avg) * 2\n tri_height = 2 / ((_high - _low) // 2) # 三角形高度\n\n # 得到三组值在网格中的位置\n high_arg = np.argwhere(columns == _high.reshape(-1, 1))[:, 1]\n avg_arg = np.argwhere(columns == _avg.reshape(-1, 1))[:, 1]\n low_arg = np.argwhere(columns == _low.reshape(-1, 1))[:, 1]\n\n # 高度表\n height = np.zeros(grid_shape)\n for i in range(len(height)):\n la = low_arg[i]\n aa = avg_arg[i]\n ha = high_arg[i]\n th = tri_height[i]\n height[i, la:aa + 1] = np.linspace(0, th, aa - la + 1)\n height[i, aa:ha + 1] = np.linspace(th, 0, ha - aa + 1)\n\n # 计算半块面积, 三角形的高变成了梯形的上下底，梯形高固定为0.5，*0.5/2=/4\n # 宽度-1，例如，原长度为5，-1后为4\n area = (height[:, :-1] + height[:, 1:]) / 4\n # 合成一块。宽度/2，例如原长度为4，/2后为2\n weight = area[:, ::2] + area[:, 1::2]\n\n # 输出\n out = np.zeros_like(weight)\n # 剩余换手率\n turnover2 = 1 - turnover\n # 第一天其实应当用上市发行价，过于麻烦，还是将第一天等权\n # 取巧方法，利用-1的特性，可减少if判断，\n out[-1] = weight[0]\n # 这里现在用的numpy, 还要快可考虑numba\n for i in range(len(turnover)):\n out[i] = out[i - 1] * turnover2[i] + weight[i] * turnover[i]\n\n # print(out.sum(axis=1))\n return out, (step / 2) * columns[1::2]\n\n\n@numba.jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef _WINNER_COST(high, low, avg, turnover, close, cost, step):\n out, columns = nb_chip(high, low, avg, turnover, step=step)\n\n # WINNER\n cheap = np.where(columns <= close.reshape(-1, 1), out, 0)\n sum_cheap = np.sum(cheap, axis=1)\n\n # COST\n # cum = np.cumsum(out, axis=1)\n cum = np.copy(out)\n for i in range(0, out.shape[0]):\n cum[i, :] = np.cumsum(out[i, :])\n\n prices = np.where(cum <= cost.reshape(-1, 1), columns, 0)\n\n # np.max(prices, axis=1)\n max_price = prices[:, 0]\n for i in range(0, out.shape[0]):\n max_price[i] = np.max(prices[i, :])\n\n return sum_cheap, max_price\n" }, { "path": "polars_ta/tdx/_nb.py", "content": "import numpy as np\nfrom numba import jit\nfrom numpy import mean, abs, full, argmax\nfrom numpy.lib.stride_tricks import sliding_window_view\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef roll_avedev(x1, window):\n out = full(x1.shape, np.nan, dtype=np.float64)\n if len(x1) < window:\n return out\n a1 = sliding_window_view(x1, window)\n for i, v1 in enumerate(a1):\n out[i + window - 1] = mean(abs(v1 - mean(v1)))\n return out\n\n\n@jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef roll_bars_since_n(x1, window):\n \"\"\"BARSSINCEN(X,N): the distance of the first observation that `X != 0` in `N` periods\n BARSSINCEN(X,N):N周期内第一次X不为0到现在的天数\n\n TODO what if all values are 0?\n\n TODO 如果一个周期内，都不满足，值取多少？0表当前值满足条件, window-1表示的是区间第0位置的值\n TODO 用window来表示都不满足\n \"\"\"\n out = full(x1.shape, np.nan, dtype=np.float64)\n if len(x1) < window:\n return out\n a1 = sliding_window_view(x1, window)\n for i, v1 in enumerate(a1):\n p = argmax(v1)\n out[i + window - 1] = window - 1 - p if p or v1[0] else window\n return out\n\n\n@jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef _up_stat(a, d: int = 3):\n \"\"\"T天N板，最稀疏为5天2板\n\n 最近几天涨停但当天没涨停也会有记录，如6天2板，所以要与涨停一起使用\n \"\"\"\n out1 = full(a.shape, 0, dtype=np.int64)\n out2 = full(a.shape, 0, dtype=np.int64)\n out3 = full(a.shape, 0, dtype=np.int64)\n t = 0 # T天\n n = 0 # N板\n k = 0 # 连续False个数\n f = True # 前面的False不处理\n for i in range(0, a.shape[0]):\n if a[i]:\n k = 0\n t += 1\n n += 1\n f = False\n out1[i] = t\n out2[i] = n\n out3[i] = k\n else:\n if f:\n continue\n k += 1\n t += 1\n if k > d:\n # 超过指定天数才会重置\n t = 0\n n = 0\n out1[i] = t\n out2[i] = n\n out3[i] = k\n return out1, out2, out3\n" }, { "path": "polars_ta/tdx/_slow.py", "content": "from polars import Series, Expr\n\n\ndef _avedev(x: Series) -> Series:\n # x is Series rather than Expr. Making this function slow\n # 可惜rolling_map后这里已经由Expr变成了Series\n return (x - x.mean()).abs().mean()\n\n\ndef AVEDEV(close: Expr, timeperiod: int = 5) -> Expr:\n \"\"\"Avoid using this function, it is slow.\n 平均绝对偏差\n \"\"\"\n return close.rolling_map(_avedev, timeperiod)\n" }, { "path": "polars_ta/tdx/arithmetic.py", "content": "\"\"\"\n通过`import`直接导入或更名的函数\n\n```python\nfrom polars_ta.wq.arithmetic import abs_ as ABS # noqa\nfrom polars_ta.wq.arithmetic import add as ADD # noqa\nfrom polars_ta.wq.arithmetic import arc_cos as ACOS # noqa\nfrom polars_ta.wq.arithmetic import arc_sin as ASIN # noqa\nfrom polars_ta.wq.arithmetic import arc_tan as ATAN # noqa\nfrom polars_ta.wq.arithmetic import ceiling as CEILING # noqa\nfrom polars_ta.wq.arithmetic import cos as COS # noqa\nfrom polars_ta.wq.arithmetic import exp as EXP # noqa\nfrom polars_ta.wq.arithmetic import floor as FLOOR # noqa\nfrom polars_ta.wq.arithmetic import fraction as FRACPART # noqa\nfrom polars_ta.wq.arithmetic import log as LN # noqa # 自然对数 (log base e)\nfrom polars_ta.wq.arithmetic import log10 as LOG # noqa # 10为底的对数 (log base 10)\nfrom polars_ta.wq.arithmetic import max_ as MAX # noqa\nfrom polars_ta.wq.arithmetic import min_ as MIN # noqa\nfrom polars_ta.wq.arithmetic import mod as MOD # noqa\nfrom polars_ta.wq.arithmetic import power as POW # noqa\nfrom polars_ta.wq.arithmetic import reverse as REVERSE # noqa\nfrom polars_ta.wq.arithmetic import round_ as _round # noqa\nfrom polars_ta.wq.arithmetic import sign as SIGN # noqa\nfrom polars_ta.wq.arithmetic import sin as SIN # noqa\nfrom polars_ta.wq.arithmetic import sqrt as SQRT # noqa\nfrom polars_ta.wq.arithmetic import subtract as SUB # noqa\nfrom polars_ta.wq.arithmetic import tan as TAN # noqa\nfrom polars_ta.wq.transformational import int_ as INTPART # noqa\n```\n\n\"\"\"\nfrom polars import Expr\n\nfrom polars_ta.wq.arithmetic import abs_ as ABS # noqa\nfrom polars_ta.wq.arithmetic import add as ADD # noqa\nfrom polars_ta.wq.arithmetic import arc_cos as ACOS # noqa\nfrom polars_ta.wq.arithmetic import arc_sin as ASIN # noqa\nfrom polars_ta.wq.arithmetic import arc_tan as ATAN # noqa\nfrom polars_ta.wq.arithmetic import ceiling as CEILING # noqa\nfrom polars_ta.wq.arithmetic import cos as COS # noqa\nfrom polars_ta.wq.arithmetic import exp as EXP # noqa\nfrom polars_ta.wq.arithmetic import floor as FLOOR # noqa\nfrom polars_ta.wq.arithmetic import fraction as FRACPART # noqa\nfrom polars_ta.wq.arithmetic import log as LN # noqa # 自然对数 (log base e)\nfrom polars_ta.wq.arithmetic import log10 as LOG # noqa # 10为底的对数 (log base 10)\nfrom polars_ta.wq.arithmetic import max_ as MAX # noqa\nfrom polars_ta.wq.arithmetic import min_ as MIN # noqa\nfrom polars_ta.wq.arithmetic import mod as MOD # noqa\nfrom polars_ta.wq.arithmetic import power as POW # noqa\nfrom polars_ta.wq.arithmetic import reverse as REVERSE # noqa\nfrom polars_ta.wq.arithmetic import round_ as _round # noqa\nfrom polars_ta.wq.arithmetic import sign as SIGN # noqa\nfrom polars_ta.wq.arithmetic import sin as SIN # noqa\nfrom polars_ta.wq.arithmetic import sqrt as SQRT # noqa\nfrom polars_ta.wq.arithmetic import subtract as SUB # noqa\nfrom polars_ta.wq.arithmetic import tan as TAN # noqa\nfrom polars_ta.wq.transformational import int_ as INTPART # noqa\n\nSGN = SIGN\n\n\ndef ROUND(x: Expr) -> Expr:\n \"\"\"Round input to closest integer.\"\"\"\n return _round(x, 0)\n\n\ndef ROUND2(x: Expr, decimals: int = 0) -> Expr:\n \"\"\"Round input to closest integer.\"\"\"\n return _round(x, decimals)\n\n\ndef BETWEEN(a: Expr, b: Expr, c: Expr) -> Expr:\n \"\"\"BETWEEN(A,B,C)表示A处于B和C之间时返回1,否则返回0\"\"\"\n x1 = (b <= a) & (a <= c)\n x2 = (c <= a) & (a <= b)\n return x1 | x2\n" }, { "path": "polars_ta/tdx/choice.py", "content": "from polars import Boolean\nfrom polars import Expr\nfrom polars import when\n\nfrom polars_ta.wq.logical import if_else\n\n\ndef IF(condition: Expr, a: Expr, b: Expr) -> Expr:\n \"\"\"A if X != 0 else B\n IF(X,A,B)若X不为0则返回A,否则返回B\n \"\"\"\n return if_else(condition.cast(Boolean), a, b)\n\n\ndef IFN(condition: Expr, a: Expr, b: Expr) -> Expr:\n \"\"\"B if X != 0 else A\n IFN(X,A,B)若X不为0则返回B,否则返回A\n \"\"\"\n return if_else(condition.cast(Boolean), b, a)\n\n\ndef VALUEWHEN(condition: Expr, x: Expr) -> Expr:\n \"\"\"VALUEWHEN(COND,X)\n当COND条件成立时,取X的当前值,否则取VALUEWHEN的上个值.\n\"\"\"\n return when(condition).then(x).otherwise(None).forward_fill()\n\n\nIFF = IF\n" }, { "path": "polars_ta/tdx/energy.py", "content": "from polars import Expr\n\nfrom polars_ta import TA_EPSILON\nfrom polars_ta.ta.price import MEDPRICE\nfrom polars_ta.tdx.reference import MA\nfrom polars_ta.tdx.reference import MAX\nfrom polars_ta.tdx.reference import REF\nfrom polars_ta.tdx.reference import SUM\n\n\ndef BRAR_AR(OPEN: Expr, HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 26) -> Expr:\n \"\"\"\n BR:SUM(MAX(0,HIGH-REF(CLOSE,1)),N)/SUM(MAX(0,REF(CLOSE,1)-LOW),N)*100;\n AR:SUM(HIGH-OPEN,N)/SUM(OPEN-LOW,N)*100;\n\n \"\"\"\n\n AR = SUM(HIGH - OPEN, N) / (SUM(OPEN - LOW, N) + TA_EPSILON) # * 100\n return AR\n\n\ndef BRAR_BR(OPEN: Expr, HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 26) -> Expr:\n \"\"\"\n BR:SUM(MAX(0,HIGH-REF(CLOSE,1)),N)/SUM(MAX(0,REF(CLOSE,1)-LOW),N)*100;\n AR:SUM(HIGH-OPEN,N)/SUM(OPEN-LOW,N)*100;\n\n \"\"\"\n LC = REF(CLOSE, 1)\n BR = SUM(MAX(0, HIGH - LC), N) / (SUM(MAX(0, LC - LOW), N) + TA_EPSILON) # * 100\n return BR\n\n\ndef CR(HIGH: Expr, LOW: Expr, N: int = 26) -> Expr:\n \"\"\"\n MID:=REF(HIGH+LOW,1)/2;\n CR:SUM(MAX(0,HIGH-MID),N)/SUM(MAX(0,MID-LOW),N)*100;\n MA1:REF(MA(CR,M1),M1/2.5+1);\n MA2:REF(MA(CR,M2),M2/2.5+1);\n MA3:REF(MA(CR,M3),M3/2.5+1);\n MA4:REF(MA(CR,M4),M4/2.5+1);\n\n \"\"\"\n MID = REF(MEDPRICE(HIGH, LOW), 1)\n return SUM(MAX(0, HIGH - MID), N) / (SUM(MAX(0, MID - LOW), N) + TA_EPSILON) # *100\n\n\ndef PSY(CLOSE: Expr, N: int = 12) -> Expr:\n \"\"\"\n PSY:COUNT(CLOSE>REF(CLOSE,1),N)/N*100;\n PSYMA:MA(PSY,M);\n\n \"\"\"\n return MA(CLOSE > REF(CLOSE, 1), N)\n\n\ndef MASS(HIGH: Expr, LOW: Expr, N1: int = 9, N2: int = 25) -> Expr:\n \"\"\"\n MASS:SUM(MA(HIGH-LOW,N1)/MA(MA(HIGH-LOW,N1),N1),N2);\n MAMASS:MA(MASS,M);\n\n \"\"\"\n MHL = MA(HIGH - LOW, N1)\n return SUM(MHL / MA(MHL, N1), N2)\n" }, { "path": "polars_ta/tdx/logical.py", "content": "from polars import Expr\n\nfrom polars_ta import TA_EPSILON\nfrom polars_ta.tdx.reference import COUNT\nfrom polars_ta.wq.logical import not_\n\n\ndef CROSS(a: Expr, b: Expr) -> Expr:\n \"\"\"上穿\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 1, 1, 2],\n 'b': [None, -1, 0, 1, 2],\n 'c': [None, 0, 0, 0, 0],\n 'd': [None, False, False, True, True],\n }).with_columns(\n out1=CROSS(pl.col('a'), pl.col('c')),\n out2=CROSS(pl.col('b'), pl.col('c')),\n out3=CROSS(0, pl.col('b')),\n out4=CROSS(pl.col('d'), 0.5),\n )\n shape: (5, 8)\n ┌──────┬──────┬──────┬───────┬───────┬───────┬───────┬───────┐\n │ a ┆ b ┆ c ┆ d ┆ out1 ┆ out2 ┆ out3 ┆ out4 │\n │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 ┆ bool ┆ bool ┆ bool ┆ bool ┆ bool │\n ╞══════╪══════╪══════╪═══════╪═══════╪═══════╪═══════╪═══════╡\n │ null ┆ null ┆ null ┆ null ┆ null ┆ null ┆ null ┆ null │\n │ -1 ┆ -1 ┆ 0 ┆ false ┆ false ┆ false ┆ null ┆ false │\n │ 1 ┆ 0 ┆ 0 ┆ false ┆ true ┆ false ┆ false ┆ false │\n │ 1 ┆ 1 ┆ 0 ┆ true ┆ false ┆ true ┆ false ┆ true │\n │ 2 ┆ 2 ┆ 0 ┆ true ┆ false ┆ false ┆ false ┆ false │\n └──────┴──────┴──────┴───────┴───────┴───────┴───────┴───────┘\n ```\n\n \"\"\"\n c = a < b - TA_EPSILON\n d = abs(a - b) < TA_EPSILON\n e = a > b + TA_EPSILON\n # 1. 小于大于\n # 2. 小于等于大于\n # 3. 小于等于等于大于。中间等了两期，暂时不定义为上穿\n return e & (c.shift(1) | (d.shift(1) & c.shift(2)))\n\n\ndef DOWNNDAY(close: Expr, N: int) -> Expr:\n \"\"\"返回周期数内是否连跌\"\"\"\n return NDAY(close.shift(), close, N)\n\n\ndef EVERY(condition: Expr, N: int) -> Expr:\n \"\"\"一直存在\"\"\"\n return COUNT(condition, N) == N\n\n\ndef EXIST(condition: Expr, N: int) -> Expr:\n \"\"\"是否存在\"\"\"\n return COUNT(condition, N) > 0\n\n\ndef EXISTR(condition: Expr, a: int, b: int) -> Expr:\n \"\"\"从前a日到前b日内是否存在\"\"\"\n # 不支持a参数为0，用一季来代替\n if a == 0:\n a = 60 + b\n return EXIST(condition, a - b).shift(b)\n\n\ndef LAST(condition: Expr, a: int, b: int) -> Expr:\n \"\"\"从前a日到前b日内一直存在\"\"\"\n # 不支持a参数为0，用一季来代替\n if a == 0:\n a = 60 + b\n return EVERY(condition, a - b).shift(b)\n\n\ndef LONGCROSS(a: Expr, b: Expr, N: int) -> Expr:\n \"\"\"两条线维持一定周期后交叉\"\"\"\n return CROSS(a, b) & EVERY(a < (b - TA_EPSILON), N).shift(1)\n\n\ndef NDAY(close: Expr, open_: Expr, N: int) -> Expr:\n \"\"\"返回是否持续存在X>Y\"\"\"\n return EVERY(close > (open_ + TA_EPSILON), N)\n\n\ndef NOT(condition: Expr) -> Expr:\n \"\"\"求逻辑非\"\"\"\n return not_(condition)\n\n\ndef UPNDAY(close: Expr, N: int) -> Expr:\n \"\"\"返回周期数内是否连涨\"\"\"\n return NDAY(close, close.shift(), N)\n\n\n# Eastmoney has these two functions\n# 东方财富中有此两函数\nALL = EVERY\nANY = EXIST\n" }, { "path": "polars_ta/tdx/moving_average.py", "content": "from polars import Expr\n\nfrom polars_ta.ta.price import AVGPRICE\nfrom polars_ta.tdx.reference import MA\n\n\ndef BBI(CLOSE: Expr, M1: int = 3, M2: int = 6, M3: int = 12, M4: int = 20) -> Expr:\n \"\"\"\n BBI:(MA(CLOSE,M1)+MA(CLOSE,M2)+MA(CLOSE,M3)+MA(CLOSE,M4))/4;\n\n \"\"\"\n return AVGPRICE(MA(CLOSE, M1), MA(CLOSE, M2), MA(CLOSE, M3), MA(CLOSE, M4))\n" }, { "path": "polars_ta/tdx/over_bought_over_sold.py", "content": "from polars import Expr, struct\n\nfrom polars_ta import TA_EPSILON\nfrom polars_ta.ta.momentum import RSI # noqa\nfrom polars_ta.ta.momentum import RSV\nfrom polars_ta.ta.price import TYPPRICE\nfrom polars_ta.tdx.choice import IF\nfrom polars_ta.tdx.reference import DIFF\nfrom polars_ta.tdx.reference import MA\nfrom polars_ta.tdx.reference import REF\nfrom polars_ta.tdx.reference import SMA_CN\nfrom polars_ta.tdx.reference import SUM\nfrom polars_ta.tdx.reference import TR\nfrom polars_ta.tdx.statistic import AVEDEV\n\n\ndef ATR(HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 14) -> Expr:\n \"\"\"\n\n Notes\n -----\n Warning: it is different with talib.ATR\n 与talib.ATR不同\n\n \"\"\"\n return MA(TR(HIGH, LOW, CLOSE), N)\n\n\ndef BIAS(CLOSE: Expr, N: int = 6) -> Expr:\n \"\"\"\n BIAS1 :(CLOSE-MA(CLOSE,N1))/MA(CLOSE,N1)*100;\n BIAS2 :(CLOSE-MA(CLOSE,N2))/MA(CLOSE,N2)*100;\n BIAS3 :(CLOSE-MA(CLOSE,N3))/MA(CLOSE,N3)*100;\n\n \"\"\"\n return CLOSE / MA(CLOSE, N) - 1 # * 100\n\n\ndef CCI(HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 14) -> Expr:\n \"\"\"\n TYP:=(HIGH+LOW+CLOSE)/3;\n CCI:(TYP-MA(TYP,N))*1000/(15*AVEDEV(TYP,N));\n\n Notes\n -----\n Avoid using `AVEDEV`, it is slow\n AVEDEV计算慢，少用\n\n \"\"\"\n TYP = TYPPRICE(HIGH, LOW, CLOSE)\n return (TYP - MA(TYP, N)) / (0.015 * AVEDEV(TYP, N))\n\n\ndef KDJ(HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 9, M1: int = 3, M2: int = 3) -> Expr:\n \"\"\"\n RSV:=(CLOSE-LLV(LOW,N))/(HHV(HIGH,N)-LLV(LOW,N))*100;\n K:SMA(RSV,M1,1);\n D:SMA(K,M2,1);\n J:3*K-2*D;\n\n \"\"\"\n rsv = RSV(HIGH, LOW, CLOSE, N)\n k = SMA_CN(rsv, M1, 1)\n d = SMA_CN(k, M2, 1)\n j = k * 3 - d * 2\n return struct(K=k, D=d, J=j)\n\n\ndef MTM(CLOSE: Expr, N: int = 12) -> Expr:\n \"\"\"\n MTM:CLOSE-REF(CLOSE,MIN(BARSCOUNT(C),N));\n MTMMA:MA(MTM,M);\n \"\"\"\n # return CLOSE - REF(CLOSE, N)\n return DIFF(CLOSE, N)\n\n\ndef MFI(CLOSE: Expr, HIGH: Expr, LOW: Expr, VOL: Expr, N: int = 14) -> Expr:\n \"\"\"\n\n TYP := (HIGH + LOW + CLOSE)/3;\n V1:=SUM(IF(TYP>REF(TYP,1),TYP*VOL,0),N)/SUM(IF(TYP LT, TYP * VOL, 0), N) / (SUM(IF(TYP < LT, TYP * VOL, 0), N) + TA_EPSILON)\n return (1 - (1 / (1 + V1))) # * 100\n" }, { "path": "polars_ta/tdx/pattern.py", "content": "from polars import Expr, Struct, Field, Float64, struct\n\nfrom polars_ta.tdx._chip import _WINNER_COST\nfrom polars_ta.utils.numba_ import batches_i2_o2, struct_to_numpy\n\n\ndef ts_WINNER_COST(high: Expr, low: Expr, avg: Expr, turnover: Expr, close: Expr, cost: Expr = 0.5, step: float = 0.1) -> Expr:\n \"\"\"\n 获利盘比例\n WINNER(CLOSE),表示以当前收市价卖出的获利盘比例,例如返回0.1表示10%获利盘;WINNER(10.5)表示10.5元价格的获利盘比例\n\n 成本分布价\n COST(10),表示10%获利盘的价格是多少,即有10%的持仓量在该价格以下,其余90%在该价格以上,为套牢盘\n\n Parameters\n ----------\n high\n 最高价\n low\n 最低价\n avg\n 平均价。可以用vwap\n turnover:\n 换手率。需要在外转成0~1范围内\n close\n 判断获利比例的价格，可以用收盘价，也可以用均价\n cost\n 成本比例，0~1\n step\n 步长。一字涨停时，三角分布的底为1，高为2。但无法当成梯形计算面积，所以从中用半步长切开计算\n\n Returns\n -------\n winner\n 获利盘比例\n cost\n 成本分布价\n\n Examples\n --------\n >>> WINNER_COST = ts_WINNER_COST(HIGH, LOW, VWAP, turnover_ratio / 100, CLOSE, 0.5)\n\n Notes\n -----\n 该函数仅对日线分析周期有效\n\n \"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return struct(f0=high, f1=low, f2=avg, f3=turnover, f4=close, f5=cost).map_batches(lambda xx: batches_i2_o2(struct_to_numpy(xx, 6, dtype=float), _WINNER_COST, step), return_dtype=dtype)\n" }, { "path": "polars_ta/tdx/pattern_feature.py", "content": "from polars import Expr\n\nfrom polars_ta.tdx.arithmetic import ABS, BETWEEN\nfrom polars_ta.tdx.arithmetic import MAX\nfrom polars_ta.tdx.arithmetic import MIN\nfrom polars_ta.tdx.choice import IF\nfrom polars_ta.tdx.logical import EVERY, CROSS, LAST, EXIST, NOT\nfrom polars_ta.tdx.reference import COUNT, BARSLASTCOUNT, BARSLAST, EMA, FILTER, LOWRANGE\nfrom polars_ta.tdx.reference import HHV\nfrom polars_ta.tdx.reference import LLV\nfrom polars_ta.tdx.reference import MA\nfrom polars_ta.tdx.reference import REF\nfrom polars_ta.wq.time_series import ts_cum_max, ts_cum_min\n\n\ndef 早晨之星(OPEN: Expr, CLOSE: Expr) -> Expr:\n \"\"\"MSTAR 早晨之星\"\"\"\n A1 = REF(CLOSE, 2) / REF(OPEN, 2) < 0.95\n A2 = REF(OPEN, 1) < REF(CLOSE, 2)\n A3 = ABS(REF(OPEN, 1) - REF(CLOSE, 1)) / REF(CLOSE, 1) < 0.03\n A4 = CLOSE / OPEN > 1.05\n A5 = CLOSE > REF(CLOSE, 2)\n return A1 & A2 & A3 & A4 & A5\n\n\ndef 剑(OPEN: Expr, HIGH: Expr, LOW: Expr, CLOSE: Expr, VOL: Expr, CAPITAL: Expr) -> Expr:\n \"\"\"SWORD 剑\"\"\"\n AA = (VOL > REF(VOL, 1)) | (VOL > CAPITAL)\n BB = (OPEN >= REF(HIGH, 1)) & (REF(HIGH, 1) > REF(HIGH, 2) * 1.06)\n CC = CLOSE > (REF(CLOSE, 1) - REF(CLOSE, 1) * 0.01)\n DD = (CLOSE < HIGH * 0.965) & (HIGH > OPEN * 1.05)\n EE = (LOW < OPEN) & (LOW < CLOSE) & (HIGH > REF(CLOSE, 1) * 1.06)\n FF = (HIGH - MAX(OPEN, CLOSE)) / 2 > (MIN(OPEN, CLOSE) - LOW)\n GG = (ABS(OPEN - CLOSE)) / 2 < (MIN(OPEN, CLOSE) - LOW)\n return AA & BB & CC & DD & EE & FF & GG\n\n\ndef 天量法则(OPEN: Expr, CLOSE: Expr) -> Expr:\n \"\"\"TLFZ 天量法则\"\"\"\n A1 = CLOSE > OPEN\n A2 = HHV(CLOSE, 50) == CLOSE\n # DYNAINFO(37) > 0.1 & &\n # DYNAINFO(13) < 0.14;\n raise\n\n\ndef 四串阴(OPEN: Expr, CLOSE: Expr) -> Expr:\n \"\"\"GREEN4 四串阴\"\"\"\n return EVERY(CLOSE < OPEN, 4)\n\n\ndef 四串阳(OPEN: Expr, CLOSE: Expr) -> Expr:\n \"\"\"RED4 四串阳\"\"\"\n return EVERY(CLOSE > OPEN, 4)\n\n\ndef 鸳鸯底(O: Expr, LOW: Expr, C: Expr, V: Expr, N: int = 50) -> Expr:\n \"\"\"YYD 鸳鸯底\"\"\"\n return (C > O) & (REF(C, 1) < REF(O, 1)) & (C > REF(O, 1)) & (V > REF(V, 1)) & EXIST(LOWRANGE(LOW) > N, 3)\n\n\ndef 出水芙蓉(OPEN: Expr, CLOSE: Expr, S: int = 20, M: int = 40, N: int = 60) -> Expr:\n \"\"\"CSFR 出水芙蓉\"\"\"\n AAA = CLOSE > OPEN\n BBB = AAA & (CLOSE > MA(CLOSE, S)) & (CLOSE > MA(CLOSE, M)) & (CLOSE > MA(CLOSE, N))\n CCC = BBB & (OPEN < MA(CLOSE, M)) & (OPEN < MA(CLOSE, N))\n return CCC & (CLOSE - OPEN > 0.0618 * CLOSE)\n\n\ndef 出水芙蓉II(C: Expr, V: Expr, N: float = 0.05, M: float = 2.0) -> Expr:\n \"\"\"CSFR2 出水芙蓉II\"\"\"\n ZF = C / REF(C, 1)\n FLTJ = V > REF(V, 1) * M\n A1 = CROSS(C, MA(C, 5)) & CROSS(C, MA(C, 10)) & CROSS(C, MA(C, 20)) & CROSS(C, MA(C, 60))\n return (ZF > 1 + N) & FLTJ & A1\n\n\ndef 近日创历史新高(HIGH: Expr, N: int = 3, M: int = 0) -> Expr:\n \"\"\"NHIGH 近日创历史新高\"\"\"\n if M == 0:\n return HHV(HIGH, N) == ts_cum_max(HIGH)\n else:\n return HHV(HIGH, N) == HHV(HIGH, M)\n\n\ndef 近日创历史新低(LOW: Expr, N: int = 3, M: int = 0) -> Expr:\n \"\"\"NLOW 近日创历史新低\"\"\"\n if M == 0:\n return LLV(LOW, N) == ts_cum_min(LOW)\n else:\n return LLV(LOW, N) == LLV(LOW, M)\n\n\ndef 旭日初升(CLOSE: Expr, VOL: Expr, N: int = 120) -> Expr:\n \"\"\"XRDS 旭日初升\"\"\"\n BUY1 = LAST(CLOSE < MA(CLOSE, N), 0, 5)\n return (CLOSE > MA(CLOSE, N)) & (VOL > MA(VOL, 5) * 2) & BUY1\n\n\ndef 蜻蜓点水(CLOSE: Expr, N: int = 120) -> Expr:\n \"\"\"QTDS 蜻蜓点水\"\"\"\n BUY1 = LAST(CLOSE > MA(CLOSE, N), 0, 5)\n BUY2 = EXIST(CLOSE < MA(CLOSE, N), 5)\n return (CLOSE > MA(CLOSE, N)) & BUY1 & BUY2\n\n\ndef 均线多头排列(OPEN: Expr, CLOSE: Expr, N: int = 5, N1: int = 10, N2: int = 20, N3: int = 30) -> Expr:\n \"\"\"DTPL 均线多头排列\"\"\"\n A1 = MA(CLOSE, N)\n A2 = MA(CLOSE, N1)\n A3 = MA(CLOSE, N2)\n A4 = MA(CLOSE, N3)\n return (CLOSE > A1) & (A1 > A2) & (A2 > A3) & (A3 > A4) & (CLOSE > OPEN)\n\n\ndef 均线空头排列(OPEN: Expr, CLOSE: Expr, N: int = 5, N1: int = 10, N2: int = 20, N3: int = 30) -> Expr:\n \"\"\"KTPL 均线空头排列\"\"\"\n A1 = MA(CLOSE, N)\n A2 = MA(CLOSE, N1)\n A3 = MA(CLOSE, N2)\n A4 = MA(CLOSE, N3)\n return (CLOSE < A1) & (A1 < A2) & (A2 < A3) & (A3 < A4) & (CLOSE < OPEN)\n\n\ndef 强势整理(OPEN: Expr, CLOSE: Expr, N: int = 2, M: float = 0.05) -> Expr:\n \"\"\"QSZL 强势整理\"\"\"\n A1 = ABS(CLOSE - OPEN) / OPEN < 0.015\n A2 = COUNT(A1, N) == N\n A3 = (REF(OPEN, N) < REF(CLOSE, N)) & (REF(CLOSE, N) / REF(CLOSE, N + 1) > 1 + M)\n return A2 & A3\n\n\ndef 高开大阴线(OPEN: Expr, CLOSE: Expr, N: float = 0.06, M: float = 0.04) -> Expr:\n \"\"\"W103 高开大阴线\"\"\"\n A1 = OPEN / REF(CLOSE, 1) >= 1 + M\n A2 = CLOSE / OPEN <= 1 - N\n return A1 & A2\n\n\ndef 低开大阳线(OPEN: Expr, CLOSE: Expr, N: float = 0.06, M: float = 0.04) -> Expr:\n \"\"\"W104 低开大阳线\"\"\"\n A1 = OPEN / REF(CLOSE, 1) <= 1 - M\n A2 = CLOSE / OPEN >= 1 + N\n return A1 & A2\n\n\ndef 跳空缺口选股(HIGH: Expr, LOW: Expr) -> Expr:\n \"\"\"W105 跳空缺口选股\"\"\"\n A1 = HIGH < (REF(LOW, 1) - 0.001)\n A2 = LOW > (REF(HIGH, 1) + 0.001)\n return A1 | A2\n\n\ndef 单阳不破选股(O: Expr, H: Expr, L: Expr, C: Expr, N1: int = 2, N2: int = 7) -> Expr:\n \"\"\"W106 单阳不破选股\n \"\"\"\n A0 = ((C > O * 1.08) | (C > REF(C, 1) * 1.08)) & NOT(H == L) & NOT((H == C) & (H == O)) # {大阳超8%，排除当天一字、T字板}\n A1 = A0 & BARSLASTCOUNT(A0) == 1\n A2 = BARSLAST(A1) # {距离大阳几根K}\n ZCX = REF(O, A2) # {获取大阳位置的开盘价作为支撑线}\n ZHX = REF(C, A2) # {获取大阳位置的收盘价作为选股最高区间}\n ZD = LLV(L, A2) # {大阳之后的最低价}\n ZH = HHV(H, A2) # {大阳之后的最高价}\n A3 = BARSLASTCOUNT(ZD >= ZCX)\n return (A3 <= N2) & (A3 > N1) & BETWEEN(C, ZCX, ZHX) & (ZH < ZHX) & (A2 > 0)\n\n\ndef 回补跳空向上缺口(O: Expr, H: Expr, L: Expr, C: Expr, N1: int = 2, N2: int = 7) -> Expr:\n \"\"\"W107 回补跳空向上缺口\"\"\"\n raise\n\n\ndef 揉搓线(O: Expr, H: Expr, L: Expr, C: Expr, V: Expr, N: int = 50) -> Expr:\n \"\"\"RUBLINE 揉搓线\n \"\"\"\n A1 = (REF(H, 1) - MAX(REF(C, 1), REF(O, 1))) / (REF(H, 1) - REF(L, 1)) * 100 > N # {上影线占K线的N % 以上}\n A2 = (MIN(O, C) - L) / (H - L) * 100 > N # {下影线占K线的N % 以上}\n A3 = ABS(C - REF(C, 1)) / MIN(C, REF(C, 1)) * 100 < 2 # {下影K的跌幅不能超过2 %}\n A4 = REF(C, 2) > REF(C, 3) # {揉搓形态前收涨}\n A5 = V < REF(V, 1) # {缩量}\n return A1 & A2 & A3 & A4 & A5\n\n\ndef 老鸭头(L: Expr, C: Expr, V: Expr) -> Expr:\n \"\"\"OLDDUCK 老鸭头\n \"\"\"\n E1 = EMA(C, 13)\n E2 = EMA(C, 55)\n A1 = (COUNT(E1 < REF(E1, 1), 5) >= 3) & (E1 > REF(E1, 1))\n A2 = (COUNT(E2 > REF(E2, 1), 13) >= 8) & (E2 < REF(E2, 1))\n A3 = LLV((L / E2 - 1), 13) <= 0.1\n A4 = COUNT(E1 > E2, 13) == 13\n A5 = COUNT(C > E2, 5) == 5\n A6 = CROSS(C, E1)\n A7 = V > MA(V, 5)\n YT = A1 & A2 & A3 & A4 & A5 & A6 & A7\n LYT = FILTER(YT, 10)\n # FXG:=FINANCE(42)>100;\n # NTP:=DYNAINFO(8)>0;\n # LYT AND FXG AND NTP;\n return LYT\n\n\ndef 仙人指路(O: Expr, H: Expr, C: Expr) -> Expr:\n \"\"\"WISEWAY 仙人指路\n \"\"\"\n TUPO = 0.5 * (C + H) > HHV(REF(C, 1), 60)\n SSQS = (MA(C, 5) > MA(C, 60)) & (MA(C, 10) > MA(C, 60))\n YINX = ((H - MAX(O, C)) / REF(C, 1) > 0.045) & (ABS(C - O) / REF(C, 1) < 0.035) & ((H - MAX(O, C)) > 2.0 * ABS(C - O))\n QIANG1 = (REF(C, 1) / REF(C, 6) > 1.04) & (REF(C, 1) / REF(C, 6) < 1.18)\n QIANG2 = (REF(C, 1) / REF(C, 6) > 1.04) & (REF(C, 1) / REF(C, 6) < 1.27)\n QIANG = IF(True, QIANG2, QIANG1) # TODO 要改\n QIANK = (REF((H - C), 1) / REF(C, 2) < 1.045) & (REF(C, 1) / REF(C, 2) > 0.97)\n return TUPO & SSQS & YINX & QIANG & QIANK\n\n\ndef 低点搜寻(HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 5) -> Expr:\n \"\"\"SP 低点搜寻\"\"\"\n W = MA((LLV(LOW, 45) - CLOSE) / (HHV(HIGH, 45) - LLV(LOW, 45)), N)\n return CROSS(-0.05, W)\n\n\ndef 突破(C: Expr, N1: int = 5, N2: int = 10, N3: int = 30) -> Expr:\n \"\"\"TP 突破\"\"\"\n M1 = MA(C, N1) # {短期参数：5}\n M2 = MA(C, N2) # {中期参数：10}\n M3 = MA(C, N3) # {长期参数：30}\n # {以下计算交叉点距今的天数}\n D1 = BARSLAST(CROSS(M1, M2)) # {短上穿中}\n D2 = BARSLAST(CROSS(M1, M3)) # {短上穿长}\n D3 = BARSLAST(CROSS(M2, M3)) # {中上穿长}\n T1 = CROSS(M2, M3) # {今天中线上穿长线}\n T2 = (D1 >= D2) & (D2 >= D3) # {交叉按指定的先后出现}\n T3 = COUNT(CROSS(M2, M1) | CROSS(M3, M2) | CROSS(M3, M1), D1) == 0 # {中间无夹杂其它交叉}\n T4 = REF((M1 < M3) & (M2 < M3), D1 + 1) # {短上穿中前一天短、中线在长线之下}\n return T1 & T2 & T3 & T4 # {价托确定};\n" }, { "path": "polars_ta/tdx/pressure_support.py", "content": "from polars import Expr\n\nfrom polars_ta.tdx.arithmetic import SQRT\nfrom polars_ta.tdx.reference import MA\nfrom polars_ta.tdx.statistic import STDP\n\n\ndef BOLL(close: Expr, M: int = 20, N: int = 2) -> Expr:\n \"\"\"\n BOLL:MA(CLOSE,M);\n UB:BOLL+2*STD(CLOSE,M);\n LB:BOLL-2*STD(CLOSE,M);\n \"\"\"\n ma = MA(close, M)\n # it should be total standard deviation, the value is smaller than sample standard deviation.\n # 这里是总体标准差，值比样本标准差小。部分软件使用样本标准差是错误的，\n std = STDP(close, M)\n return ma + std * N\n\n\ndef BOLL_M(close: Expr, M: int = 20, N: int = 2) -> Expr:\n \"\"\"\n {参数 N: 2 250 20 }\n MID:=MA(C,N);\n VART1:=POW((C-MID),2);\n VART2:=MA(VART1,N);\n VART3:=SQRT(VART2);\n UPPER:=MID+2*VART3;\n LOWER:=MID-2*VART3;\n BOLL:REF(MID,1),COLORFFFFFF;\n UB:REF(UPPER,1),COLOR00FFFF;\n LB:REF(LOWER,1),COLORFF00FF;\n \"\"\"\n ma = MA(close, M)\n # 这里var不一样\n # close-close.mean()与close-MA(close)的区别\n std = SQRT(MA((close - ma) ** 2, M))\n return ma + std * N\n" }, { "path": "polars_ta/tdx/reference.py", "content": "\"\"\"\n通过`import`直接导入或更名的函数\n\n```python\nfrom polars_ta.ta.overlap import SMA as MA\nfrom polars_ta.ta.volatility import TRANGE as TR # noqa\nfrom polars_ta.wq.arithmetic import max_ as MAX # noqa\nfrom polars_ta.wq.arithmetic import min_ as MIN # noqa\nfrom polars_ta.wq.time_series import ts_arg_max as HHVBARS # noqa\nfrom polars_ta.wq.time_series import ts_arg_min as LLVBARS # noqa\nfrom polars_ta.wq.time_series import ts_count as COUNT # noqa\nfrom polars_ta.wq.time_series import ts_decay_linear as WMA # noqa\nfrom polars_ta.wq.time_series import ts_delay as REF # noqa\nfrom polars_ta.wq.time_series import ts_delta as DIFF # noqa\nfrom polars_ta.wq.time_series import ts_max as HHV # noqa\nfrom polars_ta.wq.time_series import ts_min as LLV # noqa\nfrom polars_ta.wq.time_series import ts_product as MULAR # noqa\nfrom polars_ta.wq.time_series import ts_sum as SUM\n```\n\n\"\"\"\nfrom polars import Boolean, Int32, UInt16\nfrom polars import Expr\nfrom polars import when\n\nfrom polars_ta.ta.overlap import EMA as _ema\nfrom polars_ta.ta.overlap import SMA as MA\nfrom polars_ta.ta.volatility import TRANGE as TR # noqa\nfrom polars_ta.tdx._nb import roll_bars_since_n\nfrom polars_ta.utils.numba_ import batches_i1_o1\nfrom polars_ta.wq.arithmetic import max_ as MAX # noqa\nfrom polars_ta.wq.arithmetic import min_ as MIN # noqa\nfrom polars_ta.wq.time_series import ts_arg_max as HHVBARS # noqa\nfrom polars_ta.wq.time_series import ts_arg_min as LLVBARS # noqa\nfrom polars_ta.wq.time_series import ts_count as COUNT # noqa\nfrom polars_ta.wq.time_series import ts_decay_linear as WMA # noqa\nfrom polars_ta.wq.time_series import ts_delay as REF # noqa\nfrom polars_ta.wq.time_series import ts_delta as DIFF # noqa\nfrom polars_ta.wq.time_series import ts_max as HHV # noqa\nfrom polars_ta.wq.time_series import ts_min as LLV # noqa\nfrom polars_ta.wq.time_series import ts_product as MULAR # noqa\nfrom polars_ta.wq.time_series import ts_sum as SUM\n\n\ndef BARSLAST(condition: Expr) -> Expr:\n \"\"\"# of Observations since last time condition was true\n 上一次X不为0到现在的天数\"\"\"\n a = condition.cum_count()\n b = when(condition.cast(Boolean)).then(a).otherwise(None).forward_fill()\n return a - b\n\n\ndef BARSLASTCOUNT(condition: Expr) -> Expr:\n \"\"\"Cumulative count of continuous true observations\n 统计连续满足条件的周期数\"\"\"\n a = condition.cast(Int32).cum_sum()\n b = when(~condition.cast(Boolean)).then(a).otherwise(None).forward_fill().fill_null(0)\n return a - b\n\n\ndef BARSSINCE(condition: Expr) -> Expr:\n \"\"\"# of observations since the first time condition was true\n 第一次X不为0到现在的天数\"\"\"\n a = condition.cum_count()\n b = condition.cast(Boolean).arg_true().first()\n return a - b\n\n\ndef BARSSINCEN(condition: Expr, N: int = 30) -> Expr:\n \"\"\"# of Observations since the first time condition was true (rolling within N observations)\n N周期内第一次X不为0到现在的天数\"\"\"\n return condition.cast(Boolean).map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_bars_since_n, N, dtype=UInt16), return_dtype=UInt16)\n\n\ndef CUMSUM(close: Expr) -> Expr:\n \"\"\"SUM(close, 0)\"\"\"\n return close.cum_sum()\n\n\ndef DMA(close: Expr, alpha: float = 0.5) -> Expr:\n \"\"\"DMA(X,alpha), (Exponential moving average given alpha)\n Y = alpha * X + (1 - alpha) * last_Y\n requires 0 < alpha < 1\n\n 求X的动态移动平均.\n 算法:Y=A*X+(1-A)*Y',其中Y'表示上一周期Y值,A必须大于0且小于1.A支持变量\"\"\"\n return close.ewm_mean(alpha=alpha, adjust=False, min_samples=1)\n\n\ndef EMA(close: Expr, N: int = 30) -> Expr:\n \"\"\"EMA(X,N): Exponential moving average given N\n\n Y = X * 2/(N+1) + last_Y * (N-1)/(N+1)\n\n X的N日指数移动平均.算法:Y=(X*2+Y'*(N-1))/(N+1)\n EMA(X,N)相当于SMA(X,N+1,2),N支持变量\"\"\"\n return _ema(close, N)\n\n\ndef EXPMA(close: Expr, N: int = 30) -> Expr:\n return _ema(close, N)\n\n\ndef EXPMEMA(close: Expr, N: int = 30) -> Expr:\n \"\"\"Slow version of EMA. Do not use it unless you have to\n EXPMEMA(X,M),X的M日指数平滑移动平均。EXPMEMA同EMA(即EXPMA)的差别在于他的起始值为一平滑值\n\n Notes\n -----\n 等价于talib.EMA，由于比EMA慢，少用\n\n \"\"\"\n sma = MA(close, N)\n x = when(close.cum_count() < N).then(sma).otherwise(close)\n return x.ewm_mean(span=N, adjust=False, min_samples=1)\n\n\ndef HOD(close: Expr, N: int = 30) -> Expr:\n \"\"\"rolling rank of each data in descending order\n HOD(X,N):求当前X数据是N周期内的第几个高值,N=0则从第一个有效值开始\"\"\"\n return (-close).rolling_rank(N, method=\"min\")\n\n\ndef LOD(close: Expr, N: int = 30) -> Expr:\n \"\"\"rolling rank of each data in ascending order\n LOD(X,N):求当前X数据是N周期内的第几个低值\"\"\"\n return close.rolling_rank(N, method=\"min\")\n\n\ndef LOWRANGE(close: Expr) -> Expr:\n \"\"\"rolling rank of each data in ascending order\n LOD(X,N):求当前X数据是N周期内的第几个低值\"\"\"\n raise\n\n\ndef MEMA(close: Expr, N: int = 30) -> Expr:\n \"\"\"Exponential moving average given N\n Y = X / N + last_Y * (N-1) / N\n\n MEMA(X,N):X的N日平滑移动平均,如Y=(X+Y'*(N-1))/N\n MEMA(X,N)相当于SMA(X,N,1)\"\"\"\n raise\n\n\ndef RANGE(a: Expr, b: Expr, c: Expr) -> Expr:\n \"\"\"A在B和C范围之间,B Expr:\n \"\"\"Exponential Moving average given alpha = M/N\n Y = X * M/N + last_Y * (N-M)/N\n\n 用法:SMA(X,N,M),X的N日移动平均,M为权重,若Y=SMA(X,N,M)则Y=(X*M+Y'*(N-M))/N\n\n !!!为防止与talib版SMA误用，这里去了默认值1\n \"\"\"\n return X.ewm_mean(alpha=M / N, adjust=False, min_samples=1)\n\n\ndef SUMIF(condition: Expr, close: Expr, N: int = 30) -> Expr:\n return SUM(condition.cast(Boolean).cast(Int32) * close, N)\n\n\ndef TMA(close: Expr, N: int = 30) -> Expr:\n \"\"\"TMA(X,A,B),A和B必须小于1,算法\tY=(A*Y'+B*X),其中Y'表示上一周期Y值.初值为X\"\"\"\n raise\n\n\ndef FILTER(close: Expr, N: int = 30) -> Expr:\n raise\n\n\ndef REFX(close: Expr, N: int = 30) -> Expr:\n \"\"\"属于未来函数,引用若干周期后的数据\"\"\"\n return REF(close, -N)\n" }, { "path": "polars_ta/tdx/statistic.py", "content": "from polars import Expr, Struct, Field, Int64, Float64\n\nfrom polars_ta.tdx._nb import roll_avedev, _up_stat\nfrom polars_ta.utils.numba_ import batches_i1_o1, batches_i1_o2\nfrom polars_ta.wq.time_series import ts_corr as RELATE # noqa\nfrom polars_ta.wq.time_series import ts_covariance as COVAR # noqa\nfrom polars_ta.wq.time_series import ts_std_dev as _ts_std_dev\n\n\ndef AVEDEV(close: Expr, timeperiod: int = 5) -> Expr:\n \"\"\"mean absolute deviation\n 平均绝对偏差\"\"\"\n return close.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_avedev, timeperiod), return_dtype=Float64)\n\n\ndef DEVSQ(close: Expr, timeperiod: int = 5) -> Expr:\n raise\n\n\ndef SLOPE(close: Expr, timeperiod: int = 5) -> Expr:\n raise\n\n\ndef STD(close: Expr, timeperiod: int = 5) -> Expr:\n \"\"\"std dev with ddof = 1\n 估算标准差\"\"\"\n return _ts_std_dev(close, timeperiod, 1)\n\n\ndef STDDEV(close: Expr, timeperiod: int = 5) -> Expr:\n \"\"\"标准偏差?\"\"\"\n raise\n\n\ndef STDP(close: Expr, timeperiod: int = 5) -> Expr:\n \"\"\"std dev with ddof = 0\n 总体标准差\"\"\"\n return _ts_std_dev(close, timeperiod, 0)\n\n\ndef VAR(close: Expr, timeperiod: int = 5) -> Expr:\n return close.rolling_var(timeperiod, ddof=1)\n\n\ndef VARP(close: Expr, timeperiod: int = 5) -> Expr:\n return close.rolling_var(timeperiod, ddof=0)\n\n\ndef ts_up_stat(x: Expr) -> Expr:\n \"\"\"T天N板统计，与通达信结果一样，最简为5天2板\n\n Parameters\n ----------\n x: Expr\n 布尔序列，True表示涨停\n\n Returns\n -------\n Expr\n 1. T天\n 2. N板\n 3. 离上次涨停距离\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [False, True, True, False, True, False, False, False, False, False],\n }).with_columns(\n out=ts_up_stat(pl.col('a'))\n )\n ┌───────┬───────────┐\n │ a ┆ out │\n │ --- ┆ --- │\n │ bool ┆ struct[3] │\n ╞═══════╪═══════════╡\n │ false ┆ {0,0,0} │\n │ true ┆ {1,1,0} │\n │ true ┆ {2,2,0} │\n │ false ┆ {3,2,1} │\n │ true ┆ {4,3,0} │\n │ false ┆ {5,3,1} │\n │ false ┆ {6,3,2} │\n │ false ┆ {7,3,3} │\n │ false ┆ {0,0,4} │\n │ false ┆ {0,0,5} │\n └───────┴───────────┘\n ```\n\n \"\"\"\n dtype = Struct([Field(f\"column_{i}\", Int64) for i in range(3)])\n return x.map_batches(lambda x1: batches_i1_o2(x1.to_numpy(), _up_stat), return_dtype=dtype)\n" }, { "path": "polars_ta/tdx/times.py", "content": "from datetime import time, timedelta\n\nfrom polars import Expr, when\n\n\ndef FROMOPEN(t: Expr) -> Expr:\n \"\"\"返回当前时刻距开盘有多少分钟\n\n 范围0~240,开盘前为0，10点为31\n\n Examples\n --------\n from datetime import datetime\n\n import polars as pl\n\n from polars_ta.tdx.times import FROMOPEN, FROMOPEN1\n\n df = pl.DataFrame({'datetime': [\n datetime(2025, 1, 1, 0, 0),\n datetime(2025, 1, 1, 9, 25),\n datetime(2025, 1, 1, 9, 30, 57),\n datetime(2025, 1, 1, 9, 31),\n datetime(2025, 1, 1, 10, 0),\n datetime(2025, 1, 1, 13, 0),\n ]})\n\n df = df.with_columns(\n FROMOPEN=FROMOPEN(pl.col('datetime')),\n FROMOPEN1=FROMOPEN_1(pl.col('datetime'), 0),\n FROMOPEN2=FROMOPEN_1(pl.col('datetime'), 60),\n )\n\n shape: (6, 4)\n ┌─────────────────────┬──────────┬───────────┬───────────┐\n │ datetime ┆ FROMOPEN ┆ FROMOPEN1 ┆ FROMOPEN2 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ datetime[μs] ┆ i64 ┆ i64 ┆ i64 │\n ╞═════════════════════╪══════════╪═══════════╪═══════════╡\n │ 2025-01-01 00:00:00 ┆ 0 ┆ 240 ┆ 240 │\n │ 2025-01-01 09:25:00 ┆ 0 ┆ 1 ┆ 1 │\n │ 2025-01-01 09:30:57 ┆ 1 ┆ 1 ┆ 2 │\n │ 2025-01-01 09:31:00 ┆ 2 ┆ 2 ┆ 3 │\n │ 2025-01-01 10:00:00 ┆ 31 ┆ 31 ┆ 32 │\n │ 2025-01-01 13:00:00 ┆ 121 ┆ 121 ┆ 122 │\n └─────────────────────┴──────────┴───────────┴───────────┘\n\n \"\"\"\n am = (t.dt.time() - time(9, 29)).dt.total_minutes().clip(0, 120)\n pm = (t.dt.time() - time(12, 59)).dt.total_minutes().clip(0, 120)\n return am + pm\n\n\ndef FROMOPEN_1(t: Expr, offset: int) -> Expr:\n \"\"\"返回当前时刻距开盘有多少分钟。范围1~240\n\n 用于计算量比\n 1. 竞价量比，分母应当为1\n 2. 日线数据0~8点时，返回240\n 3. 日线数据9点时，返回1\n\n Parameters\n ----------\n t : Expr\n 时间列\n offset : int\n 偏移量，单位秒\n\n Notes\n -----\n 每根K线结束时，标签是当前时间的50多秒，而结束时时间已经到下以分钟了，所以建议加60秒\n\n \"\"\"\n return when(t.dt.time() >= time(8, 45)).then(FROMOPEN(t + timedelta(seconds=offset)).clip(1, 240)).otherwise(240)\n" }, { "path": "polars_ta/tdx/trend.py", "content": "from polars import Expr\n\nfrom polars_ta import TA_EPSILON\nfrom polars_ta.tdx.arithmetic import ABS\nfrom polars_ta.tdx.choice import IF\nfrom polars_ta.tdx.reference import MA\nfrom polars_ta.tdx.reference import REF\nfrom polars_ta.tdx.reference import SUM\nfrom polars_ta.tdx.reference import TR\n\n\ndef DPO(CLOSE: Expr, N: int = 20) -> Expr:\n \"\"\"\n DPO:CLOSE-REF(MA(CLOSE,N),N/2+1);\n MADPO:MA(DPO,M);\n \"\"\"\n\n return CLOSE - REF(MA(CLOSE, N), N // 2 + 1)\n\n\ndef EMV(HIGH: Expr, LOW: Expr, VOL: Expr, N: int = 14) -> Expr:\n \"\"\"\n VOLUME:=MA(VOL,N)/VOL;\n MID:=100*(HIGH+LOW-REF(HIGH+LOW,1))/(HIGH+LOW);\n EMV:MA(MID*VOLUME*(HIGH-LOW)/MA(HIGH-LOW,N),N);\n MAEMV:MA(EMV,M);\n \"\"\"\n ADD = HIGH + LOW\n SUB = HIGH - LOW\n\n VOLUME = MA(VOL, N) / VOL\n MID = 100 * (ADD - REF(ADD, 1)) / ADD\n return MA(MID * VOLUME * SUB / (MA(SUB, N) + TA_EPSILON), N)\n\n\ndef PLUS_DM(HIGH: Expr, LOW: Expr, N: int = 14) -> Expr:\n \"\"\"\n MTR:=SUM(MAX(MAX(HIGH-LOW,ABS(HIGH-REF(CLOSE,1))),ABS(REF(CLOSE,1)-LOW)),N);\n HD :=HIGH-REF(HIGH,1);\n LD :=REF(LOW,1)-LOW;\n DMP:=SUM(IF(HD>0&&HD>LD,HD,0),N);\n DMM:=SUM(IF(LD>0&&LD>HD,LD,0),N);\n \"\"\"\n HD = HIGH - REF(HIGH, 1)\n LD = REF(LOW, 1) - LOW\n DMP = SUM(IF((HD > 0) & (HD > LD), HD, 0), N)\n return DMP\n\n\ndef MINUS_DM(HIGH: Expr, LOW: Expr, N: int = 14) -> Expr:\n HD = HIGH - REF(HIGH, 1)\n LD = REF(LOW, 1) - LOW\n DMM = SUM(IF((LD > 0) & (LD > HD), LD, 0), N)\n return DMM\n\n\ndef PLUS_DI(HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 14) -> Expr:\n \"\"\"\n MTR:=SUM(MAX(MAX(HIGH-LOW,ABS(HIGH-REF(CLOSE,1))),ABS(REF(CLOSE,1)-LOW)),N);\n HD :=HIGH-REF(HIGH,1);\n LD :=REF(LOW,1)-LOW;\n DMP:=SUM(IF(HD>0&&HD>LD,HD,0),N);\n DMM:=SUM(IF(LD>0&&LD>HD,LD,0),N);\n PDI: DMP*100/MTR;\n MDI: DMM*100/MTR;\n \"\"\"\n MTR = SUM(TR(HIGH, LOW, CLOSE), N)\n DMP = PLUS_DM(HIGH, LOW, N)\n PDI = DMP / MTR\n return PDI # * 100\n\n\ndef MINUS_DI(HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 14) -> Expr:\n MTR = SUM(TR(HIGH, LOW, CLOSE), N)\n DMM = MINUS_DM(HIGH, LOW, N)\n MDI = DMM / MTR\n return MDI # * 100\n\n\ndef ADX(HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 14, M: int = 6) -> Expr:\n \"\"\"\n MTR:=SUM(MAX(MAX(HIGH-LOW,ABS(HIGH-REF(CLOSE,1))),ABS(REF(CLOSE,1)-LOW)),N);\n HD :=HIGH-REF(HIGH,1);\n LD :=REF(LOW,1)-LOW;\n DMP:=SUM(IF(HD>0&&HD>LD,HD,0),N);\n DMM:=SUM(IF(LD>0&&LD>HD,LD,0),N);\n PDI: DMP*100/MTR;\n MDI: DMM*100/MTR;\n ADX: MA(ABS(MDI-PDI)/(MDI+PDI)*100,M);\n ADXR:(ADX+REF(ADX,M))/2;\n \"\"\"\n MTR = SUM(TR(HIGH, LOW, CLOSE), N)\n HD = HIGH - REF(HIGH, 1)\n LD = REF(LOW, 1) - LOW\n DMP = SUM(IF((HD > 0) & (HD > LD), HD, 0), N)\n DMM = SUM(IF((LD > 0) & (LD > HD), LD, 0), N)\n PDI = DMP / MTR # * 100\n MDI = DMM / MTR # * 100\n return MA(ABS(MDI - PDI) / (MDI + PDI), M) # * 100\n\n\ndef ADXR(HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 14, M: int = 6) -> Expr:\n adx = ADX(HIGH, LOW, CLOSE, N, M)\n adxr = (adx + REF(adx, M)) / 2\n return adxr\n" }, { "path": "polars_ta/tdx/trend_feature.py", "content": "from polars import Expr\n\nfrom polars_ta.tdx.arithmetic import ABS\nfrom polars_ta.tdx.logical import EXIST, EVERY\nfrom polars_ta.tdx.pattern import ts_WINNER_COST\nfrom polars_ta.tdx.reference import COUNT\nfrom polars_ta.tdx.reference import HHV\nfrom polars_ta.tdx.reference import LLV\nfrom polars_ta.tdx.reference import MA\nfrom polars_ta.tdx.reference import REF\nfrom polars_ta.tdx.reference import SUM\nfrom polars_ta.wq.time_series import ts_cum_max, ts_cum_min\n\n\ndef N天内有跳空向上缺口(H: Expr, L: Expr, N: int = 1, M: float = 0.01) -> Expr:\n \"\"\"C100 N天内有跳空向上缺口\n\n Parameters\n ----------\n N\n 天数\n M\n 涨幅\n\n \"\"\"\n XSTK = L > REF(H, 1) * (100 + M)\n return EXIST(XSTK, N)\n\n\ndef N日内创新高(HIGH: Expr, N: int = 10) -> Expr:\n \"\"\"C101 N日内创新高\n \"\"\"\n return HHV(HIGH, N) == ts_cum_max(HIGH)\n\n\ndef N日内创新低(LOW: Expr, N: int = 10) -> Expr:\n \"\"\"C102 N日内创新低\n \"\"\"\n return LLV(LOW, N) == ts_cum_min(LOW)\n\n\ndef N日内阴线多于阳线(OPEN: Expr, CLOSE: Expr, N: int = 30, M: float = 0.6) -> Expr:\n \"\"\"C103 M日内阴线多于阳线\n\n Parameters\n ----------\n N\n 天数\n M\n 比例\n\n \"\"\"\n return COUNT(OPEN > CLOSE, N) / N >= M\n\n\ndef N日内阳线多于阴线(OPEN: Expr, CLOSE: Expr, N: int = 30, M: float = 0.6) -> Expr:\n \"\"\"C104 M日内阳线多于阴线\n\n Parameters\n ----------\n N\n 天数\n M\n 比例\n\n \"\"\"\n return COUNT(OPEN < CLOSE, N) / N >= M\n\n\ndef N日内上涨多于下跌(CLOSE: Expr, N: int = 120, M: float = 0.6) -> Expr:\n \"\"\"C105 N日内上涨多于下跌\n\n Parameters\n ----------\n N\n 天数\n M\n 比例\n\n \"\"\"\n return COUNT(CLOSE > REF(CLOSE, 1), N) / N >= M\n\n\ndef N日内下跌多于上涨(CLOSE: Expr, N: int = 120, M: float = 0.6) -> Expr:\n \"\"\"C106 N日内下跌多于上涨\n\n Parameters\n ----------\n N\n 天数\n M\n 比例\n\n \"\"\"\n return COUNT(CLOSE < REF(CLOSE, 1), N) / N >= M\n\n\ndef 连续N天收阳线(OPEN: Expr, CLOSE: Expr, N: int = 7) -> Expr:\n \"\"\"C107 连续N天收阳线\n\n Parameters\n ----------\n N\n 天数\n\n \"\"\"\n return EVERY(CLOSE > OPEN, N)\n\n\ndef 连续N天收阴线(OPEN: Expr, CLOSE: Expr, N: int = 7) -> Expr:\n \"\"\"C108 连续N天收阴线\n\n Parameters\n ----------\n N\n 天数\n\n \"\"\"\n return EVERY(OPEN > CLOSE, N)\n\n\ndef 单日放量(VOL: Expr, CAPITAL: Expr, N: float = 2, M: float = 0.15) -> Expr:\n \"\"\"C110 单日放量\n\n Parameters\n ----------\n N\n 5日平均成交量的倍数\n M\n 流通股本的倍数\n\n \"\"\"\n A1 = MA(VOL, 5)\n A2 = REF(A1, 1)\n C1 = VOL / A2 > N\n C2 = VOL / CAPITAL > M\n return C1 & C2\n\n\ndef 阶段缩量(VOL: Expr, CAPITAL: Expr, N: int = 20, M: float = 0.02) -> Expr:\n \"\"\"C111 阶段缩量\n\n Parameters\n ----------\n VOL\n 成交量\n CAPITAL\n 流通股本\n\n Notes\n -----\n 成交量与流通股本单位要一致，都为手，或者都为股\n\n \"\"\"\n return SUM(VOL, N) / CAPITAL <= M\n\n\ndef 阶段放量(VOL: Expr, CAPITAL: Expr, N: int = 10, M: float = 2.0) -> Expr:\n \"\"\"C112 阶段放量\n\n Parameters\n ----------\n VOL\n 成交量\n CAPITAL\n 流通股本\n\n \"\"\"\n return SUM(VOL, N) / CAPITAL >= M\n\n\ndef 持续放量(VOL: Expr, M: int = 5) -> Expr:\n \"\"\"C113 持续放量\n\n Parameters\n ----------\n VOL\n 成交量\n\n \"\"\"\n return EVERY(VOL >= REF(VOL, 1), M)\n\n\ndef 持续缩量(VOL: Expr, M: int = 5) -> Expr:\n \"\"\"C114 持续缩量\n\n Parameters\n ----------\n VOL\n 成交量\n\n \"\"\"\n return COUNT(VOL <= REF(VOL, 1), M) == M\n\n\ndef 间隔放量(VOL: Expr, N: int = 30, N1: float = 4.0, N2: float = 2.0, N3: int = 3) -> Expr:\n \"\"\"C115 间隔放量\n\n Parameters\n ----------\n VOL\n 成交量\n N\n 均量周期\n N1\n 最小均量与最大均量的倍数\n N2\n 成交量与均量的倍数\n N3\n 满足N2时的次数\n\n \"\"\"\n A = MA(VOL, 5)\n C1 = HHV(A, N) < N1 * LLV(A, N) # 成交量最大与最小在一定范围内\n C2 = COUNT(VOL > N2 * A, N) > N3 # 成交量大于均量一定倍数\n return C1 & C2\n\n\ndef 放量上攻(CLOSE: Expr, VOL: Expr, CAPITAL: Expr,\n N: float = 0.01, N1: int = 3, N2: float = 0.2, N3: int = 4) -> Expr:\n \"\"\"C116 放量上攻\n\n Parameters\n ----------\n CLOSE\n 复权收盘价\n VOL\n 成交量\n CAPITAL\n 流通股本\n N\n 涨幅\n N1, N2\n N1天内成交量大于N2*流通股本\n N3\n 连续N3天满足涨幅\n\n \"\"\"\n A = (CLOSE - REF(CLOSE, 1)) / REF(CLOSE, 1) >= N\n C1 = SUM(VOL, N1) / CAPITAL >= N2\n C2 = COUNT(VOL > REF(VOL, 1), N3) == N3\n C3 = COUNT(A, N3) == N3\n return C1 & C2 & C3\n\n\ndef 温和放量上攻(CLOSE: Expr, VOL: Expr, CAPITAL: Expr, N: int = 5) -> Expr:\n \"\"\"C117 温和放量上攻\n\n Parameters\n ----------\n CLOSE\n 复权收盘价\n VOL\n 成交量\n CAPITAL\n 流通股本\n N\n 观察天数\n\n \"\"\"\n A1 = CLOSE / REF(CLOSE, 1)\n A2 = (A1 > 1) & (A1 < 1.03) # {股价小幅上扬}\n B1 = VOL / REF(VOL, 1)\n B2 = (B1 > 1) & (A1 < 2) # {成交量小幅上扬}\n C1 = MA(VOL, N) / CAPITAL < 0.05 # 日成交量小于流通股本的5%\n C2 = COUNT(A2 & B2, N) / N > 0.6\n return C1 & C2\n\n\ndef 突然放量(VOL: Expr, N: int = 10, M: float = 3.0) -> Expr:\n \"\"\"C118 突然放量\n\n Parameters\n ----------\n VOL\n 成交量\n N\n 观察天数\n M\n 倍数\n\n \"\"\"\n return VOL > REF(HHV(VOL, N), 1) * M\n\n\ndef 平台整理(CLOSE: Expr, N: int = 30, N1: float = 0.05) -> Expr:\n \"\"\"C119 平台整理\n\n Parameters\n ----------\n N\n 天数\n N1\n 幅度\n\n \"\"\"\n return HHV(CLOSE, N) / LLV(CLOSE, N) <= 1 + N1\n\n\ndef 小步碎阳(O: Expr, H: Expr, L: Expr, C: Expr, avg: Expr, turnover_ratio: Expr, N: int = 4) -> Expr:\n \"\"\"C120 小步碎阳\n\n Parameters\n ----------\n avg\n 成交均价\n turnover_ratio\n 换手率\n N\n 观察天数\n\n \"\"\"\n AA = (C > REF(C, 1)) & (C > O)\n A1 = COUNT(AA, N) == N\n A2 = C / REF(C, N) < 1.05\n A3 = ts_WINNER_COST(H, L, avg, turnover_ratio, C, 0.5).struct[0] > 0.75\n\n return A1 & A2 & A3\n\n\ndef 突破长期盘整(HIGH: Expr, LOW: Expr, CLOSE: Expr, N: int = 30, N1: int = 5) -> Expr:\n \"\"\"C123 突破长期盘整\n\n Parameters\n ----------\n N\n 天数\n N1\n 涨幅\n\n \"\"\"\n C1 = REF(HHV(HIGH, N) / LLV(LOW, N), 1) <= N1 + 1\n C2 = CLOSE >= REF(HHV(HIGH, N), 1)\n\n return C1 & C2\n\n\ndef N天内出现以涨停收盘(收盘涨停: Expr, N: int = 10) -> Expr:\n \"\"\"C128 N天内出现以涨停收盘\n\n Parameters\n ----------\n N\n 天数\n\n \"\"\"\n return EXIST(收盘涨停, N)\n\n\ndef N天内出现涨停(最高涨停: Expr, N: int = 20) -> Expr:\n \"\"\"C129 N天内出现涨停\n\n Parameters\n ----------\n N\n 天数\n\n \"\"\"\n return EXIST(最高涨停, N)\n\n\ndef N天内经常涨停(收盘涨停: Expr, N: int = 100, M: int = 8) -> Expr:\n \"\"\"C130 N天内经常涨停\n\n Parameters\n ----------\n N\n 天数\n M\n 涨停天数\n\n \"\"\"\n return COUNT(收盘涨停, N) >= M\n\n\ndef 下跌多日再放量上涨(HIGH: Expr, CLOSE: Expr, VOL: Expr) -> Expr:\n \"\"\"C131 下跌多日再放量上涨\n\n \"\"\"\n A1 = REF(CLOSE, 5) > REF(CLOSE, 4)\n A2 = REF(CLOSE, 4) > REF(CLOSE, 3)\n A3 = REF(CLOSE, 3) > REF(CLOSE, 2)\n A4 = REF(CLOSE, 2) > REF(CLOSE, 1)\n A5 = (CLOSE > REF(HIGH, 1)) & (VOL > REF(VOL, 1))\n return A1 & A2 & A3 & A4 & A5\n\n\ndef 跳空高开或低开(O: Expr, H: Expr, L: Expr, C: Expr, N: float = 0.03) -> Expr:\n \"\"\"C132 跳空高开或低开\n\n Parameters\n ----------\n N\n 涨幅\n\n \"\"\"\n if N > 0:\n A = (O > REF(H, 1)) & (O / REF(C, 1) > (1 + N))\n return A\n else:\n B = (O < REF(L, 1)) & (O / REF(C, 1) < (1 + N))\n return B\n\n\ndef 拉升后多日调整(C: Expr, N: int = 3, ZF: float = 0.09) -> Expr:\n \"\"\"C133 拉升后多日调整\n\n Parameters\n ----------\n N\n 天数\n ZF\n 涨幅\n\n \"\"\"\n C1 = REF(C, N) / REF(C, N + 1) > 1 + ZF\n C2 = EVERY(C < REF(C, 1), N)\n return C1 & C2\n\n\ndef 昨日底部十字星(O: Expr, H: Expr, L: Expr, C: Expr, N: int = 60) -> Expr:\n \"\"\"C134 昨日底部十字星\n\n Parameters\n ----------\n N\n 天数\n\n \"\"\"\n\n C1 = L <= LLV(L, N)\n C2 = ABS(C - O) / (H - L) < 0.05\n C3 = H > L\n return REF(C1 & C2 & C3, 1)\n\n\ndef 价量渐低后阳包阴(O: Expr, C: Expr, V: Expr) -> Expr:\n \"\"\"C135 价量渐低后阳包阴\n \"\"\"\n A1 = REF(C, 4) > REF(C, 3)\n A2 = REF(C, 3) > REF(C, 2)\n A3 = REF(C, 2) > REF(C, 1)\n B1 = REF(V, 3) > REF(V, 2)\n B2 = REF(V, 2) > REF(V, 1)\n AA1 = A1 & A2 & A3 & B1 & B2\n AA2 = (REF(O, 1) > REF(C, 1)) & (C > REF(O, 1))\n AA3 = (C > O) & (V < REF(V, 1))\n\n return AA1 & AA2 & AA3\n" }, { "path": "polars_ta/tdx/volume.py", "content": "from polars import Expr\n\nfrom polars_ta.tdx.choice import IF\nfrom polars_ta.tdx.reference import REF\nfrom polars_ta.tdx.reference import SUM\nfrom polars_ta.tdx.reference import CUMSUM\n\n\ndef OBV(CLOSE: Expr, VOL: Expr) -> Expr:\n \"\"\"\n VA:=IF(CLOSE>REF(CLOSE,1),VOL,-VOL);\n OBV:SUM(IF(CLOSE=REF(CLOSE,1),0,VA),0);\n MAOBV:MA(OBV,M);\n \"\"\"\n LC = REF(CLOSE, 1)\n VA = IF(CLOSE - LC, VOL, -VOL)\n return CUMSUM(IF(CLOSE == LC, 0, VA))\n\n\ndef VR(CLOSE: Expr, VOL: Expr, N: int = 26) -> Expr:\n \"\"\"\n TH:=SUM(IF(CLOSE>REF(CLOSE,1),VOL,0),N);\n TL:=SUM(IF(CLOSE LC, VOL, 0), N)\n TL = SUM(IF(CLOSE < LC, VOL, 0), N)\n TQ = SUM(IF(CLOSE == LC, VOL, 0), N)\n\n return (TH * 2 + TQ) / (TL * 2 + TQ) # * 100\n" }, { "path": "polars_ta/utils/__init__.py", "content": "" }, { "path": "polars_ta/utils/factor.py", "content": "\"\"\"\n复权算法\n1. 针对现金分红进行复权（使用加减法）\n2. 针对拆股进行复权（使用乘除法）\n\nfactor1: 稀疏的复权因子\nfactor2: 稠密的复权因子，之后主要使用这个\n\n复权因子使用方法\n乘除复权法：factor2*close\n加减复权法：factor2+close\n\"\"\"\nimport polars as pl\n\n\ndef calc_factor_muldiv(df: pl.DataFrame,\n by1: str = 'asset', by2: str = 'date',\n close: str = 'close', pre_close: str = 'pre_close') -> pl.DataFrame:\n \"\"\"计算复权因子，乘除法。使用交易所发布的昨收盘价计算\n\n Parameters\n ----------\n df : pl.DataFrame\n 数据\n by1 : str\n 分组字段\n by2 : str\n 排序字段\n close : str\n 收盘价字段\n pre_close : str\n 昨收盘价字段\n\n Notes\n -----\n 不关心是否真发生了除权除息过程，只要知道前收盘价和收盘价不等就表示发生了除权除息\n\n \"\"\"\n df = (\n df\n .sort(by1, by2)\n .with_columns(factor1=(pl.col(close).shift(1, fill_value=pl.first(pre_close)) / pl.col(pre_close)).round(8).over(by1, order_by=by2))\n .with_columns(factor2=(pl.col('factor1').cum_prod()).over(by1, order_by=by2))\n )\n return df\n\n\ndef calc_factor_addsub(df: pl.DataFrame,\n by1: str = 'asset', by2: str = 'date',\n close: str = 'close', pre_close: str = 'pre_close') -> pl.DataFrame:\n \"\"\"计算复权因子，加减法。使用交易所发布的昨收盘价计算\n\n Parameters\n ----------\n df : pl.DataFrame\n 数据\n by1 : str\n 分组字段\n by2 : str\n 排序字段\n close : str\n 收盘价字段\n pre_close : str\n 昨收盘价字段\n\n Notes\n -----\n 不关心是否真发生了除权除息过程，只要知道前收盘价和收盘价不等就表示发生了除权除息\n\n \"\"\"\n df = (\n df\n .sort(by1, by2)\n .with_columns(factor1=(pl.col(close).shift(1, fill_value=pl.first(pre_close)) - pl.col(pre_close)).round(8).over(by1, order_by=by2))\n .with_columns(factor2=(pl.col('factor1').cum_sum()).over(by1, order_by=by2))\n )\n return df\n" }, { "path": "polars_ta/utils/functions.py", "content": "import inspect\nfrom functools import wraps\nfrom inspect import currentframe\n\nimport polars as pl\n\nimport polars_ta\n\n\ndef const_to_expr(func):\n \"\"\"将表达式中不合法的常量,改成表达式\"\"\"\n\n def repeat_const(a, p):\n if p.annotation.__name__ == \"Expr\" and not isinstance(a, pl.Expr):\n return pl.repeat(a, pl.len())\n else:\n return a\n\n @wraps(func)\n def decorated(*args):\n parameters = inspect.signature(func).parameters\n args = [repeat_const(a, p) for (n, p), a in zip(parameters.items(), args)]\n return func(*args)\n\n return decorated\n\n\ndef apply_const_to_expr():\n \"\"\"应用常量转表达式功能到所有函数\n\n 对于不合法的表达式参数，可以一定程度上兼容，适用于遗传算法中表达式简化成常量的情况\n \"\"\"\n frame = currentframe().f_back\n for k, v in frame.f_globals.items():\n if inspect.isfunction(v) and v.__module__ and v.__module__.startswith(polars_ta.__package__):\n frame.f_locals[k] = const_to_expr(v)\n" }, { "path": "polars_ta/utils/numba_.py", "content": "\"\"\"\nDemo for using numba to implement rolling functions.\n本文件是使用numba实现rolling的函数，演示用\n\"\"\"\nfrom functools import lru_cache\nfrom typing import List\n\nimport numpy as np\nfrom numba import jit\nfrom numpy import full\nfrom numpy.lib.stride_tricks import sliding_window_view\nfrom polars import Series, Expr, struct, DataFrame, Float64\n\n\"\"\"\nSeries.to_numpy的操作在调用之前做，这样可控一些\nbatches_i1_o1这一类的函数输入不支持Series，只支持numpy。设计成在map_batches转换更可控\n\"\"\"\n\n\n@jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef isnan(x):\n # https://github.com/numba/numba/issues/2919#issuecomment-747377615\n if int(x) == -9223372036854775808:\n return True\n else:\n return False\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef full_with_window_size(arr, fill_value, dtype=None, window_size: int = 1):\n \"\"\"创建一个更大的数组，填充后一截数据\"\"\"\n out = full(arr.shape[0] + window_size - 1, fill_value, dtype=dtype)\n out[window_size - 1:] = arr\n return out\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef sliding_window_with_min_periods(arr, window_size: int, min_periods: int):\n \"\"\"为rolling准备的数据，当数据长度不足时，用nan填充\"\"\"\n windows = sliding_window_view(arr, window_size)\n valid_counts = np.sum(~np.isnan(windows), axis=1)\n # 修改这一行，使用布尔索引而不是np.where\n result = windows.copy()\n result[valid_counts < min_periods] = np.nan\n return result\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _roll_1(x1: np.ndarray, window: int, min_periods: int, func):\n \"\"\"TODO 只是模板演示，无法编译通过\"\"\"\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n out1[:] = np.nan\n for i, v1 in enumerate(a1):\n if np.isnan(v1).all():\n continue\n out1[i] = func(v1)\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _roll_2(x1, x2, window, min_periods, func):\n \"\"\"TODO 只是模板演示，无法编译通过\"\"\"\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n a2 = sliding_window_with_min_periods(out2, window, min_periods)\n out1[:] = np.nan\n for i, (v1, v2) in enumerate(zip(a1, a2)):\n if np.isnan(v1).all():\n continue\n if np.isnan(v2).all():\n continue\n out1[i] = func(v1, v2)\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _roll_3(x1, x2, x3, window, min_periods, func):\n \"\"\"TODO 只是模板演示，无法编译通过\"\"\"\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)\n out3 = full_with_window_size(x3, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n a2 = sliding_window_with_min_periods(out2, window, min_periods)\n a3 = sliding_window_with_min_periods(out3, window, min_periods)\n out1[:] = np.nan\n for i, (v1, v2, v3) in enumerate(zip(a1, a2, a3)):\n if np.isnan(v1).all():\n continue\n if np.isnan(v2).all():\n continue\n if np.isnan(v3).all():\n continue\n out1[i] = func(v1, v2, v3)\n return out1[:x1.shape[0]]\n\n\ndef struct_to_numpy(xx, n: int, dtype=None):\n if dtype is None:\n return [xx.struct[i].to_numpy() for i in range(n)]\n else:\n return [xx.struct[i].to_numpy().astype(dtype) for i in range(n)]\n\n\ndef batches_i1_o1(x1: np.ndarray, func, *args, dtype=None) -> Series:\n return Series(func(x1, *args), nan_to_null=True, dtype=dtype)\n\n\ndef batches_i2_o1(xx: List[np.ndarray], func, *args, dtype=None) -> Series:\n return Series(func(*xx, *args), nan_to_null=True, dtype=dtype)\n\n\ndef batches_i1_o2(x1: np.ndarray, func, *args, dtype=None) -> Series:\n return DataFrame(func(x1, *args), nan_to_null=True).to_struct()\n\n\ndef batches_i2_o2(xx: List[np.ndarray], func, *args, dtype=None) -> Series:\n return DataFrame(func(*xx, *args), nan_to_null=True).to_struct()\n\n\ndef batches_i2_o2_v2(xx: List[np.ndarray], func, *args, dtype=None) -> Series:\n \"\"\"此写法也能用，速度差异不大\"\"\"\n out = func(*xx, *args)\n arr = np.empty((xx[0].shape[0], len(out)), dtype=dtype)\n for i, x in enumerate(out):\n arr[:, i] = x\n return Series(arr, nan_to_null=True).arr.to_struct()\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef nb_roll_sum(x1, window):\n \"\"\"Demo code. Use `pl.col('A').rolling_sum(10).alias('a1')` instead.\n 演示代码，请直接用 pl.col('A').rolling_sum(10).alias('a1')\"\"\"\n out = np.full(x1.shape, np.nan, dtype=np.float64)\n if len(x1) < window:\n return out\n a1 = sliding_window_view(x1, window)\n for i, v1 in enumerate(a1):\n out[i + window - 1] = np.sum(v1)\n return out\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef nb_roll_cov(x1, x2, window):\n \"\"\"Demo code. Use `pl.rolling_cov(pl.col('A'), pl.col('B'), window_size=10).alias('a6')` instead.\n 演示代码，pl.rolling_cov(pl.col('A'), pl.col('B'), window_size=10).alias('a6')\"\"\"\n out = np.full(x1.shape, np.nan, dtype=np.float64)\n if len(x1) < window:\n return out\n a1 = sliding_window_view(x1, window)\n a2 = sliding_window_view(x2, window)\n for i, (v1, v2) in enumerate(zip(a1, a2)):\n out[i + window - 1] = np.cov(v1, v2)[0, 1]\n return out\n\n\ndef roll_sum(x: Expr, n: int) -> Expr:\n return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), nb_roll_sum, n), return_dtype=Float64)\n\n\ndef roll_cov(a: Expr, b: Expr, n: int) -> Expr:\n return struct(f0=a, f1=b).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), nb_roll_cov, n), return_dtype=Float64)\n\n\n@lru_cache\n@jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef get_exponent_weights(\n window: int = 10,\n half_life: int = 5,\n) -> np.ndarray:\n return np.repeat(0.5 ** (1 / half_life), window) ** np.arange(window - 1, -1, -1)\n" }, { "path": "polars_ta/utils/pit.py", "content": "\"\"\"\n# Point In Time处理相关函数\n\n## 原始表\n资产负债表 时点数据\n 利润表 区间数据\n现金流量表 区间数据\n\n## 单季度数据\n资产负债表 同原始表\n 利润表 本季-上季；一季度\n现金流量表 本季-上季；一季度\n\n## TTM数据\n资产负债表 四个报告期平均值；最新和同比两期的平均值；最新\n 利润表 年报；本季+去年报-同比；年化\n现金流量表 年报；本季+去年报-同比；年化\n\n\n## 整理成PIT数据\n例如：pit_prepare的功能为\n1. ACBD 通过join_where成\n2. A,AB,ACB,ACBD 通过sort+unique成\n3. A,AB,ABC,ABCD\n\n这样，每一组都满足以观察日视角没有未来数据，并只观察所能看到的最新值，可参考pit_calc示例实现其它处理\n\n处理好的数据再pit_frist取最早公布的项目\n例如：A1,A2,A3B3C3,C4D4 取first的结果为 A1B3C3D4\n\n这时的数据就可以做因子了，但不能做时序计算，时序计算得回到pit_calc中实现\n\n\"\"\"\nfrom typing import Tuple\n\nimport polars as pl\n\n# 公布日。收盘后公布将会标记为下一日\n# 其实行情数据的日期相当于报告日report_date，隐去了公布日announce_date，遇到要修正历史数据时才有公布日\nLOOKBACK_DATE = '__LOOKBACK_DATE__'\n\n\ndef pit_prepare(df: pl.DataFrame | pl.LazyFrame,\n by1: str = 'asset',\n by2: str = 'report_date',\n by3: str = 'announce_date',\n by4: str = LOOKBACK_DATE,\n lookback_year: str = '-5y') -> pl.DataFrame | pl.LazyFrame:\n \"\"\"将原始的财务表根据公布日重新扩展，输出满足以下条件：\n\n 1. 根据股票和公布日分组\n 2. 组内没有未来数据\n 3. 组内的同一报告期数据只取最新\n\n Returns\n -------\n pl.DataFrame|pl.LazyFrame\n - asset\n - report_date\n - announce_date\n - LOOKBACK_DATE\n 是公布日也是观察日，在观察日可以看到最新的历史数据，看不到未来数据\n\n \"\"\"\n\n def upsample_by_cum_max(df: pl.DataFrame) -> Tuple[pl.DataFrame, pl.DataFrame]:\n \"\"\"分割数据，并将单调递增的数据进行重采样\"\"\"\n # 取递增\n df1 = df.filter((pl.col(by2) == pl.col(by2).cum_max()).over(by1, order_by=[by3, by2]))\n\n # 6月30或9月30为一条数据时，upsample会是3月30，12月30，强行改成下月1日\n df1 = df1.with_columns(pl.col(by2).dt.offset_by(by='1d'))\n df1 = df1.sort(by1, by2).upsample(by2, every='1q', group_by=by1).with_columns(pl.col(by1, by3).backward_fill())\n # 还原成上月底\n df1 = df1.with_columns(pl.col(by2).dt.offset_by(by='-1d'))\n\n # 取非递增\n df2 = df.filter((pl.col(by2) < pl.col(by2).cum_max()).over(by1, order_by=[by3, by2]))\n return df1, df2\n\n # =========================================\n # ttm等操作shift(4)假定4期报告没有缺失，需要提前重采样补全\n df2 = df.sort(by1, by3, by2) # 按发布日排序，过滤单调递增\n if isinstance(df, pl.LazyFrame):\n df2 = df2.collect()\n\n dfs = []\n while True:\n if df2.is_empty():\n break\n df1, df2 = upsample_by_cum_max(df2)\n dfs.append(df1)\n # 合并\n df3 = pl.concat(dfs)\n if isinstance(df, pl.LazyFrame):\n df3 = df3.lazy()\n\n del df1\n del df2\n del df\n del dfs\n # =========================================\n # 根据发布日期，笛卡尔扩展，过滤掉未来数据\n tmp1 = '__TMP_1__'\n tmp2 = '__TMP_2__'\n df1 = df3.select(pl.col(by1).alias(tmp1), pl.col(by3).dt.offset_by(lookback_year).alias(tmp2), pl.col(by3).alias(by4)).unique()\n df1 = df1.join_where(df3,\n (pl.col(tmp1) == pl.col(by1))\n & (pl.col(by4) >= pl.col(by3))\n & (pl.col(tmp2) <= pl.col(by2)) # 最多观察lookback_year年前报告，减少计算量\n ).drop(tmp1, tmp2)\n del df3\n # =========================================\n # 过滤数据，同报告期只保留最新的一条\n df1 = (\n df1\n .sort(by1, by4, by2, by3)\n .unique([by1, by4, by2], keep=\"last\", maintain_order=True)\n )\n\n return df1\n\n\ndef pit_calc(df: pl.DataFrame | pl.LazyFrame,\n by1: str = 'asset',\n by2: str = 'report_date',\n by4: str = LOOKBACK_DATE) -> pl.DataFrame | pl.LazyFrame:\n \"\"\"输入PIT分组的财务数据，组内计算时序指标\n\n 同观察期下，同一报告期只有一条最新数据，所以没有了by3\n \"\"\"\n df1 = (\n df.with_columns(\n # TODO 补充其他各种时序指标，注意，不要少了`( ).over(by1, by4, order_by=by2)`\n net_profit_to_total_operate_revenue_ttm=(pl.col('net_profit').rolling_mean(4) / pl.col('total_operating_revenue').rolling_mean(4)).over(by1, by4, order_by=by2)\n )\n )\n return df1\n\n\ndef pit_frist(df: pl.DataFrame | pl.LazyFrame,\n by1: str = 'asset',\n by2: str = 'report_date',\n by3: str = 'announce_date',\n by4: str = LOOKBACK_DATE) -> pl.DataFrame | pl.LazyFrame:\n \"\"\"输入PIT分组的财务数据，多组合并保留最先发布的数据\n\n 此数据不含未来数据，但原则上不能再做时序处理\n\n Returns\n -------\n pl.DataFrame|pl.LazyFrame\n - asset\n - report_date\n - announce_date\n\n Warnings\n --------\n 结果不能时序计算，不能取历史，只能取最新\n\n 002509.XSHE 怎么处理？2019年报因故更新晚于2020年一季报，因为一季报先公布，所以2019年报显示永远为null\n\n \"\"\"\n df1 = (\n df\n .sort(by1, by4, by2, by3)\n .unique([by1, by2], keep=\"first\", maintain_order=True)\n .drop(by4) # 处理后by4的值与by2相等，可直接丢弃\n )\n return df1\n\n\ndef period_to_quarter(col: pl.Expr | str, quarter: pl.Expr | str) -> pl.Expr:\n \"\"\"区间数据转成单季数据\n\n 1. `利润表`和`现金流量表`的`原始表`可以转成单季数据\n 2. `资产负债表`不能转单季\n\n Examples\n --------\n ```python\n df = df.with_columns(\n quarter=pl.col('report_date').dt.quarter(),\n ).with_columns(\n period_to_quarter(cs.numeric().exclude('quarter'), quarter='quarter').over('code', 'pub_date', order_by='report_date').name.suffix('_Q'),\n )\n ```\n\n \"\"\"\n col = pl.col(col)\n quarter = pl.col(quarter)\n return pl.when(quarter == 1).then(col).otherwise(col.diff())\n\n\ndef ttm_from_point(col: pl.Expr | str) -> pl.Expr:\n \"\"\"时点数据计算TTM\n\n 1. 仅`资产负债表`原始表可调用此函数\n 2. `利润表`和`现金流量表`不能调用\n\n 数据要按一年四期排列，不能一年两期\n\n Examples\n --------\n ```python\n df = df.with_columns(\n ttm_from_point('total_assets').over('code', 'pub_date', order_by='report_date').name.suffix('_ttm')\n )\n ```\n\n \"\"\"\n col = pl.col(col)\n return pl.coalesce(\n col.rolling_mean(4), # 4期平均\n (col + col.shift(4)) / 2, # 同比报告期平均\n col, # 最新\n )\n\n\ndef last_year(col: pl.Expr | str, quarter: pl.Expr | str) -> pl.Expr:\n \"\"\"最新年报\n\n Examples\n --------\n ```python\n df = df.with_columns(\n quarter=pl.col('report_date').dt.quarter(),\n ).with_columns(\n last_year('total_assets').over('code', 'pub_date', order_by='report_date').name.suffix('_ly')\n )\n ```\n\n \"\"\"\n col = pl.col(col)\n quarter = pl.col(quarter)\n return pl.when(quarter == 4).then(col).otherwise(None).forward_fill(3)\n\n\ndef ttm_from_period(col: pl.Expr | str, quarter: pl.Expr | str) -> pl.Expr:\n \"\"\"区间数据计算ttm\n\n 1. `利润表`和`现金流量表`原始表可调用此函数\n 2. `资产负债表`不能调用\n\n 数据要按一年四期排列，不能一年两期\n\n Examples\n --------\n ```python\n df = df.with_columns(\n quarter=pl.col('report_date').dt.quarter(),\n ).with_columns(\n ttm_from_peroid('total_operating_revenue', quarter='quarter').over('code', 'pub_date', order_by='report_date').name.suffix('_TTM')\n )\n\n ```\n \"\"\"\n col = pl.col(col)\n quarter = pl.col(quarter)\n\n # 年报数据\n f1 = pl.when(quarter == 4).then(col).otherwise(None)\n # 当前报告期+上年年报-上年同比\n f2 = col + f1.forward_fill(3) - col.shift(4)\n # 年化计算法=当前报告期*年化系数\n f3 = col / quarter * 4\n return pl.coalesce(f1, f2, f3)\n\n\ndef yoy(col: pl.Expr, period: int = 4) -> pl.Expr:\n \"\"\"YOY : （Year On Year）同比增长率 \"\"\"\n return (col - col.shift(period)) / col.shift(period).abs()\n\n\ndef qoq(col: pl.Expr, period: int = 1) -> pl.Expr:\n \"\"\"QOQ : （Quarter On Quarter） 环比增长率\"\"\"\n return (col - col.shift(period)) / col.shift(period).abs()\n\n\ndef join_quote_financial(quote: pl.DataFrame | pl.LazyFrame,\n financial: pl.DataFrame | pl.LazyFrame,\n by1: str = 'asset',\n by2: str = 'date') -> pl.DataFrame | pl.LazyFrame:\n \"\"\"合并行情与财务表。请提前对财务表进行ttm等计算。因为之后再按报告期对齐很麻烦\n\n 1. 同一天发布多期，需要正确排序最新一期\n 2. 更新历史上的某一期，不能把他当成最新一期\n\n Parameters\n ----------\n quote:\n 行情表\n financial:\n 财务表。收盘后公布会显示第二天，一周7天都可能公布。同一天能公布多期\n by1\n by2\n\n \"\"\"\n quote = quote.sort(by1, by2)\n financial = financial.sort(by1, by2)\n\n return quote.join_asof(financial, left_on=by2, right_on=by2, by=by1, strategy=\"backward\", check_sortedness=False)\n\n\n# =========================================\n# 原始报告期数据预整理\n#\n# 因各种数据源设计理念不同，需要使用不同的预处理方法\n# =========================================\n\n\n# =========================================\n# 聚宽报告期财务数据\n#\n# STK_BALANCE_SHEET STK_INCOME_STATEMENT STK_CASHFLOW_STATEMENT\n# 披露最新报表report_type=0时会同时披露基准报表report_type=1，利用此功能实现历史调整\n#\n# 2019年年报20200611疫情晚于一季报公布，但2020一季报20200430资产负债表中基准是用的2019年报 002509.XSHE\n# https://www.cninfo.com.cn/new/disclosure/detail?orgId=9900015939&announcementId=1207687414&announcementTime=2020-04-30\n# =========================================\ndef sheet_from_joinquant(df: pl.DataFrame | pl.LazyFrame) -> pl.DataFrame | pl.LazyFrame:\n \"\"\"报表报告期调整，强制end_date为报告期\n\n 利用report_type=1的表进行历史数据调整，可部分实现PIT\n \"\"\"\n return df.with_columns(report_date=pl.col('end_date'))\n" }, { "path": "polars_ta/utils/withs.py", "content": "import re\n\nimport polars as pl\n\n\ndef with_industry(df: pl.DataFrame, industry_name: str, drop_first: bool, keep_col: bool) -> pl.DataFrame:\n \"\"\"添加行业哑元变量\n\n Parameters\n ----------\n df\n industry_name:str\n 行业列名\n drop_first\n 丢弃第一列\n keep_col\n 是否保留源列\n\n Returns\n -------\n pl.DataFrame\n\n \"\"\"\n df = df.with_columns([\n # 行业处理，由浮点改成整数\n pl.col(industry_name).cast(pl.UInt32),\n ])\n\n # TODO 没有行业的也过滤，这会不会有问题？已经退市了，但引入了未来数据\n df = df.filter(\n pl.col(industry_name).is_not_null(),\n )\n\n if keep_col:\n df = df.with_columns(df.to_dummies(industry_name, drop_first=False))\n else:\n df = df.to_dummies(industry_name, drop_first=False)\n\n if drop_first:\n # drop_first丢弃哪个字段是随机的，非常不友好，只能在行业中性化时动态修改代码\n industry_columns = sorted(filter(lambda x: re.search(rf\"^{industry_name}_\\d+$\", x), df.columns))\n return df.drop(industry_columns[0])\n else:\n return df\n" }, { "path": "polars_ta/wq/__init__.py", "content": "from polars_ta.wq.arithmetic import * # noqa\nfrom polars_ta.wq.cross_sectional import * # noqa\nfrom polars_ta.wq.half_life import * # noqa\nfrom polars_ta.wq.logical import * # noqa\nfrom polars_ta.wq.preprocess import * # noqa\nfrom polars_ta.wq.time_series import * # noqa\nfrom polars_ta.wq.transformational import * # noqa\nfrom polars_ta.wq.vector import * # noqa\n" }, { "path": "polars_ta/wq/_nb.py", "content": "import numpy as np\nfrom numba import jit, float64, boolean\nfrom numpy import argmax, argmin, full, vstack, corrcoef, nanprod, nanmean, nanstd\nfrom numpy.lib.stride_tricks import sliding_window_view\n\nfrom polars_ta.utils.numba_ import isnan, full_with_window_size, sliding_window_with_min_periods\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef roll_argmax(x1, window, min_periods, reverse):\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n out1[:] = np.nan\n\n for i, v1 in enumerate(a1):\n if np.isnan(v1[-min_periods:]).any():\n continue\n if reverse:\n v1 = v1[::-1]\n out1[i] = argmax(v1[~np.isnan(v1)])\n\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef roll_argmin(x1, window, min_periods, reverse):\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n out1[:] = np.nan\n\n for i, v1 in enumerate(a1):\n if np.isnan(v1[-min_periods:]).any():\n continue\n if reverse:\n v1 = v1[::-1]\n out1[i] = argmin(v1[~np.isnan(v1)])\n\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef roll_prod(x1, window, min_periods):\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n out1[:] = np.nan\n for i, v1 in enumerate(a1):\n if np.isnan(v1).all():\n continue\n out1[i] = nanprod(v1)\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _co_kurtosis(a1, a2):\n t1 = a1 - nanmean(a1)\n t2 = a2 - nanmean(a2)\n t3 = nanstd(a1)\n t4 = nanstd(a2)\n return nanmean(t1 * (t2 ** 3)) / (t3 * (t4 ** 3))\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef roll_co_kurtosis(x1, x2, window, min_periods):\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n a2 = sliding_window_with_min_periods(out2, window, min_periods)\n out1[:] = np.nan\n for i, (v1, v2) in enumerate(zip(a1, a2)):\n if np.isnan(v1).all():\n continue\n if np.isnan(v2).all():\n continue\n out1[i] = _co_kurtosis(v1, v2)\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _co_skewness(a1, a2):\n t1 = a1 - nanmean(a1)\n t2 = a2 - nanmean(a2)\n t3 = nanstd(a1)\n t4 = nanstd(a2)\n return nanmean(t1 * (t2 ** 2)) / (t3 * (t4 ** 2))\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef roll_co_skewness(x1, x2, window, min_periods):\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n a2 = sliding_window_with_min_periods(out2, window, min_periods)\n out1[:] = np.nan\n for i, (v1, v2) in enumerate(zip(a1, a2)):\n if np.isnan(v1).all():\n continue\n if np.isnan(v2).all():\n continue\n out1[i] = _co_skewness(v1, v2)\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _moment(a1, k):\n \"\"\"centeral moment\n 中心矩阵\"\"\"\n return nanmean((a1 - nanmean(a1)) ** k)\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef roll_moment(x1, window, min_periods, k):\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n out1[:] = np.nan\n for i, v1 in enumerate(a1):\n if np.isnan(v1).all():\n continue\n out1[i] = _moment(v1, k)\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _partial_corr(a1, a2, a3):\n \"\"\"TODO 不知是否正确，需要检查\"\"\"\n c = corrcoef(vstack((a1, a2, a3)))\n rxy = c[0, 1]\n rxz = c[0, 2]\n ryz = c[1, 2]\n t1 = rxy - rxz * ryz\n t2 = (1 - rxz ** 2) ** 0.5\n t3 = (1 - ryz ** 2) ** 0.5\n return t1 / (t2 * t3)\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef roll_partial_corr(x1, x2, x3, window, min_periods):\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)\n out3 = full_with_window_size(x3, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n a2 = sliding_window_with_min_periods(out2, window, min_periods)\n a3 = sliding_window_with_min_periods(out3, window, min_periods)\n out1[:] = np.nan\n for i, (v1, v2, v3) in enumerate(zip(a1, a2, a3)):\n if np.isnan(v1).all():\n continue\n if np.isnan(v2).all():\n continue\n if np.isnan(v3).all():\n continue\n out1[i] = _partial_corr(v1, v2, v3)\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _triple_corr(a1, a2, a3):\n t1 = a1 - nanmean(a1)\n t2 = a2 - nanmean(a2)\n t3 = a3 - nanmean(a3)\n t4 = nanstd(a1)\n t5 = nanstd(a2)\n t6 = nanstd(a3)\n return nanmean(t1 * t2 * t3) / (t4 * t5 * t6)\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef roll_triple_corr(x1, x2, x3, window, min_periods):\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n out2 = full_with_window_size(x2, np.nan, dtype=np.float64, window_size=window)\n out3 = full_with_window_size(x3, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n a2 = sliding_window_with_min_periods(out2, window, min_periods)\n a3 = sliding_window_with_min_periods(out3, window, min_periods)\n out1[:] = np.nan\n for i, (v1, v2, v3) in enumerate(zip(a1, a2, a3)):\n if np.isnan(v1).all():\n continue\n if np.isnan(v2).all():\n continue\n if np.isnan(v3).all():\n continue\n out1[i] = _triple_corr(v1, v2, v3)\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef _cum_prod_by(r, by):\n out = full(by.shape, 0, dtype=np.float64)\n out[:] = by\n for i in range(1, r.shape[0]):\n if isnan(out[i]):\n out[i] = r[i] * out[i - 1]\n return out\n\n\n@jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef _cum_sum_by(r, by):\n out = full(by.shape, 0, dtype=np.float64)\n out[:] = by\n for i in range(1, r.shape[0]):\n if isnan(out[i]):\n out[i] = r[i] + out[i - 1]\n return out\n\n\n@jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef _cum_sum_reset(a):\n last = 0\n out = full(a.shape, 0, dtype=np.float64)\n for i in range(0, a.shape[0]):\n curr = 0 if isnan(a[i]) else a[i]\n\n if curr == 0:\n out[i] = 0\n elif curr > 0:\n if last <= 0:\n out[i] = curr\n else:\n out[i] = curr + last\n elif curr < 0:\n if last >= 0:\n out[i] = curr\n else:\n out[i] = curr + last\n\n last = out[i]\n return out\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _sum_split_by(x1, x2, window=10, n=2):\n out1 = np.full(x1.shape[0], np.nan, dtype=np.float64)\n out2 = np.full(x1.shape[0], np.nan, dtype=np.float64)\n if len(x1) < window:\n return out1, out2\n a1 = sliding_window_view(x1, window)\n a2 = sliding_window_view(x2, window)\n for i, (v1, v2) in enumerate(zip(a1, a2)):\n # 排序两次，解决nan的问题\n b1 = np.argsort(v2)[:n]\n b2 = np.argsort(-v2)[:n]\n out1[i + window - 1] = np.sum(v1[b1])\n out2[i + window - 1] = np.sum(v1[b2])\n return out1, out2\n\n\n@jit(float64[:](boolean[:], boolean[:], boolean[:], boolean[:], boolean, boolean),\n nopython=True, nogil=True, cache=True)\ndef _signals_to_size(is_long_entry: np.ndarray, is_long_exit: np.ndarray,\n is_short_entry: np.ndarray, is_short_exit: np.ndarray,\n accumulate: bool = False,\n action: bool = False) -> np.ndarray:\n \"\"\"将4路信号转换成持仓状态。适合按资产分组后的长表,参考于`vectorbt`\n\n 在`LongOnly`场景下，`is_short_entry`和`is_short_exit`输入数据值都为`False`即可\n\n Parameters\n ----------\n is_long_entry: np.ndarray\n 是否多头入场\n is_long_exit: np.ndarray\n 是否多头出场\n is_short_entry: np.ndarray\n 是否空头入场\n is_short_exit: np.ndarray\n 是否空头出场\n accumulate: bool\n 遇到重复信号时是否累计\n action: bool\n 返回持仓状态还是下单操作\n\n Returns\n -------\n np.ndarray\n 持仓状态\n\n Examples\n --------\n ```python\n long_entry = np.array([True, True, False, False, False])\n long_exit = np.array([False, False, True, False, False])\n short_entry = np.array([False, False, True, False, False])\n short_exit = np.array([False, False, False, True, False])\n\n amount = signals_to_amount(long_entry, long_exit, short_entry, short_exit, accumulate=True, action=False)\n ```\n\n \"\"\"\n _amount: float = 0.0 # 持仓状态\n _action: float = 0.0 # 下单方向\n out = np.zeros(len(is_long_entry), dtype=np.float64)\n for i in range(len(is_long_entry)):\n if _amount == 0.0:\n # 多头信号优先级高于空头信号\n if is_long_entry[i]:\n _amount += 1.0\n _action = 1.0\n elif is_short_entry[i]:\n _amount -= 1.0\n _action = -1.0\n elif _amount > 0.0:\n if is_long_exit[i]:\n _amount -= 1.0\n _action = -1.0\n elif is_long_entry[i] and accumulate:\n _amount += 1.0\n _action = 1.0\n else:\n if is_short_exit[i]:\n _amount += 1.0\n _action = 1.0\n elif is_short_entry[i] and accumulate:\n _amount -= 1.0\n _action = -1.0\n\n out[i] = _action if action else _amount\n return out\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _roll_decay_linear(x1, window, min_periods):\n \"\"\"\n def ts_decay_linear(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _roll_decay_linear, d, minp), return_dtype=Float64)\n\n \"\"\"\n weights = np.arange(1., window + 1)\n\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n out1[:] = np.nan\n for i, v1 in enumerate(a1):\n if np.isnan(v1).all():\n continue\n mask = ~np.isnan(v1)\n out1[i] = np.average(v1[mask], weights=weights[mask])\n return out1[:x1.shape[0]]\n\n\n@jit(nopython=True, nogil=True, cache=True)\ndef _roll_decay_exp_window(x1, window, min_periods, factor):\n \"\"\"\n def ts_decay_exp_window(x: Expr, d: int = 30, factor: float = 1.0, min_samples: Optional[int] = None) -> Expr:\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _roll_decay_exp_window, d, minp, factor), return_dtype=Float64)\n\n \"\"\"\n weights = factor ** np.arange(window - 1, -1, -1)\n\n out1 = full_with_window_size(x1, np.nan, dtype=np.float64, window_size=window)\n a1 = sliding_window_with_min_periods(out1, window, min_periods)\n out1[:] = np.nan\n for i, v1 in enumerate(a1):\n if np.isnan(v1).all():\n continue\n mask = ~np.isnan(v1)\n out1[i] = np.average(v1[mask], weights=weights[mask])\n return out1[:x1.shape[0]]\n" }, { "path": "polars_ta/wq/_slow.py", "content": "\"\"\"\nAlgorithm in this file is slow, and has been replaced by numba and other methods\n本目录下算法有些慢，已经用numba等其它方法代替\n\"\"\"\nimport numpy as np\nfrom polars import Series, Expr\n\n\ndef _arg_max(x: Series):\n \"\"\"\n Notes\n -----\n TODO 等polars推出rolling_arg_max(reverse=True)这个问题能好转\n\n \"\"\"\n # return x[::-1].arg_max()\n # return x.reverse().arg_max() # 正确，但太慢\n return len(x) - 1 - x.arg_max() # 有多个最大值相同时，靠前的值会被记录下来，导致结果偏大\n\n\ndef ts_arg_max(x: Expr, d: int = 5) -> Expr:\n # WorldQuant中最大值为今天返回0，为昨天返回1\n return x.rolling_map(_arg_max, d)\n\n\ndef _arg_min(x: Series):\n return len(x) - 1 - x.arg_min()\n\n\ndef ts_arg_min(x: Expr, d: int = 5) -> Expr:\n return x.rolling_map(_arg_min, d)\n\n\ndef ts_product(x: Expr, d: int = 5) -> Expr:\n return x.rolling_map(np.nanprod, d)\n" }, { "path": "polars_ta/wq/arithmetic.py", "content": "\"\"\"\n除了使用概率极高的几个函数，其他函数不再对常量做lit\n\"\"\"\nfrom polars import Expr, fold, any_horizontal, Float64, Int64, lit\nfrom polars import arctan2 as _arctan2\nfrom polars import max_horizontal, sum_horizontal, min_horizontal, mean_horizontal\n\n\ndef abs_(x: Expr) -> Expr:\n \"\"\"求绝对值\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n 'b': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=abs_(pl.col('a')),\n out2=abs_(-1),\n )\n shape: (5, 4)\n ┌──────┬──────┬──────┬──────┐\n │ a ┆ b ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 ┆ i32 │\n ╞══════╪══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ 1 │\n │ -1 ┆ -1 ┆ 1 ┆ 1 │\n │ 0 ┆ 0 ┆ 0 ┆ 1 │\n │ 1 ┆ 1 ┆ 1 ┆ 1 │\n │ 2 ┆ 2 ┆ 2 ┆ 1 │\n └──────┴──────┴──────┴──────┘\n ```\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.abs()\n\n\ndef add(a: Expr, b: Expr, *args) -> Expr:\n \"\"\"水平多列相加\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 2, 3, 4, None],\n 'b': [5, None, 3, 2, None],\n 'c': [1, 1, None, 1, None],\n }).with_columns(\n out=add(pl.col('a'), pl.col('b'), pl.col('c'))\n )\n shape: (5, 4)\n ┌──────┬──────┬──────┬──────┐\n │ a ┆ b ┆ c ┆ out │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 ┆ i64 │\n ╞══════╪══════╪══════╪══════╡\n │ null ┆ 5 ┆ 1 ┆ 6 │\n │ 2 ┆ null ┆ 1 ┆ 3 │\n │ 3 ┆ 3 ┆ null ┆ 6 │\n │ 4 ┆ 2 ┆ 1 ┆ 7 │\n │ null ┆ null ┆ null ┆ null │\n └──────┴──────┴──────┴──────┘\n\n ```\n\n Notes\n -----\n 全`null`时返回`null`\n\n \"\"\"\n # # 全null时返回0\n # return sum_horizontal(a, b, *args)\n\n _args = [a, b] + list(args)\n return fold(acc=any_horizontal(_args) - 1, function=lambda acc, x: acc + x.fill_null(0), exprs=_args)\n\n\ndef arc_cos(x: Expr) -> Expr:\n \"\"\"反余弦\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.arccos()\n\n\ndef arc_sin(x: Expr) -> Expr:\n \"\"\"反正弦\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.arcsin()\n\n\ndef arc_tan(x: Expr) -> Expr:\n \"\"\"反正切\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.arctan()\n\n\ndef arc_tan2(y: Expr, x: Expr) -> Expr:\n \"\"\"反正切二值函数\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return _arctan2(y, x)\n\n\ndef cbrt(x: Expr) -> Expr:\n \"\"\"立方根\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.cbrt()\n\n\ndef ceiling(x: Expr) -> Expr:\n \"\"\"向上取整\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.ceil()\n\n\ndef cos(x: Expr) -> Expr:\n \"\"\"余弦\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.cos()\n\n\ndef cosh(x: Expr) -> Expr:\n \"\"\"双曲余弦\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.cosh()\n\n\ndef cot(x: Expr) -> Expr:\n \"\"\"余切\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.cot()\n\n\ndef cube(x: Expr) -> Expr:\n \"\"\"立方\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.pow(3)\n\n\ndef degrees(x: Expr) -> Expr:\n \"\"\"弧度转角度\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.degrees()\n\n\ndef _densify(x: Expr) -> Expr:\n raise\n\n\ndef div(x: Expr, y: Expr) -> Expr:\n \"\"\"x除以y的整数部分\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1.5, 0., 1.5, 2.5],\n 'b': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=div(pl.col('a'), 0),\n out2=div(pl.col('a'), 1),\n out3=div(pl.col('a'), pl.col('b')),\n )\n shape: (5, 5)\n ┌──────┬──────┬──────┬──────┬──────┐\n │ a ┆ b ┆ out1 ┆ out2 ┆ out3 │\n │ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n │ f64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 │\n ╞══════╪══════╪══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ null ┆ null │\n │ -1.5 ┆ -1 ┆ null ┆ -2 ┆ 1 │\n │ 0.0 ┆ 0 ┆ null ┆ 0 ┆ null │\n │ 1.5 ┆ 1 ┆ null ┆ 1 ┆ 1 │\n │ 2.5 ┆ 2 ┆ null ┆ 2 ┆ 1 │\n └──────┴──────┴──────┴──────┴──────┘\n ```\n\n \"\"\"\n return x.floordiv(y).cast(Int64, strict=False)\n\n\ndef divide(x: Expr, y: Expr) -> Expr:\n \"\"\"除法\n\n x/y\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n 'b': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=divide(pl.col('a'), 0),\n out2=divide(pl.col('a'), 1),\n out3=divide(pl.col('a'), pl.col('b')),\n )\n shape: (5, 5)\n ┌──────┬──────┬──────┬──────┬──────┐\n │ a ┆ b ┆ out1 ┆ out2 ┆ out3 │\n │ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ f64 ┆ f64 ┆ f64 │\n ╞══════╪══════╪══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ null ┆ null │\n │ -1 ┆ -1 ┆ -inf ┆ -1.0 ┆ 1.0 │\n │ 0 ┆ 0 ┆ NaN ┆ 0.0 ┆ NaN │\n │ 1 ┆ 1 ┆ inf ┆ 1.0 ┆ 1.0 │\n │ 2 ┆ 2 ┆ inf ┆ 2.0 ┆ 1.0 │\n └──────┴──────┴──────┴──────┴──────┘\n ```\n\n \"\"\"\n return x / y\n\n\ndef exp(x: Expr) -> Expr:\n \"\"\"自然指数函数\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=expm1(pl.col('a')),\n )\n shape: (5, 2)\n ┌──────┬──────────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪══════════╡\n │ null ┆ null │\n │ -1 ┆ 0.367879 │\n │ 0 ┆ 1.0 │\n │ 1 ┆ 2.718282 │\n │ 2 ┆ 7.389056 │\n └──────┴──────────┘\n ```\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.exp()\n\n\ndef expm1(x: Expr) -> Expr:\n \"\"\"对数收益率 转 简单收益率\n\n convert log return to simple return\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=expm1(pl.col('a')),\n )\n shape: (5, 2)\n ┌──────┬───────────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪═══════════╡\n │ null ┆ null │\n │ -1 ┆ -0.632121 │\n │ 0 ┆ 0.0 │\n │ 1 ┆ 1.718282 │\n │ 2 ┆ 6.389056 │\n └──────┴───────────┘\n ```\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.exp() - 1\n\n\ndef floor(x: Expr) -> Expr:\n \"\"\"向下取整\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.floor()\n\n\ndef fraction(x: Expr) -> Expr:\n \"\"\"小数部分\n\n This operator removes the whole number part and returns the remaining fraction part with sign.\n\n Examples\n --------\n\n ```python\n df = pl.DataFrame({\n 'a': [-2.5, -1.2, -1., None, 2., 3.2],\n }).with_columns(\n out=fraction(pl.col('a'))\n )\n\n shape: (6, 2)\n ┌──────┬──────┐\n │ a ┆ out │\n │ --- ┆ --- │\n │ f64 ┆ f64 │\n ╞══════╪══════╡\n │ -2.5 ┆ -0.5 │\n │ -1.2 ┆ -0.2 │\n │ -1.0 ┆ -0.0 │\n │ null ┆ null │\n │ 2.0 ┆ 0.0 │\n │ 3.2 ┆ 0.2 │\n └──────┴──────┘\n ```\n\n Notes\n -----\n 按小学时的定义，负数`-1.2`的整数部分是`-2`,小数部分是`0.8`，而这有所不同\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#fractionx\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.sign() * (x.abs() % 1)\n\n\ndef inverse(x: Expr) -> Expr:\n \"\"\"倒数\n\n 1/x\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=inverse(pl.col('a')),\n )\n shape: (5, 2)\n ┌──────┬──────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪══════╡\n │ null ┆ null │\n │ -1 ┆ -1.0 │\n │ 0 ┆ inf │\n │ 1 ┆ 1.0 │\n │ 2 ┆ 0.5 │\n └──────┴──────┘\n ```\n\n \"\"\"\n return 1 / x\n\n\ndef log(x: Expr) -> Expr:\n \"\"\"以e为底的对数\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=log(pl.col('a')),\n )\n shape: (5, 2)\n ┌──────┬──────────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪══════════╡\n │ null ┆ null │\n │ -1 ┆ NaN │\n │ 0 ┆ -inf │\n │ 1 ┆ 0.0 │\n │ 2 ┆ 0.693147 │\n └──────┴──────────┘\n ```\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.log()\n\n\ndef log10(x: Expr) -> Expr:\n \"\"\"以10为底的对数\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=log10(pl.col('a')),\n )\n shape: (5, 2)\n ┌──────┬─────────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪═════════╡\n │ null ┆ null │\n │ -1 ┆ NaN │\n │ 0 ┆ -inf │\n │ 1 ┆ 0.0 │\n │ 2 ┆ 0.30103 │\n └──────┴─────────┘\n ```\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.log10()\n\n\ndef log1p(x: Expr) -> Expr:\n \"\"\"简单收益率 转 对数收益率\n\n convert simple return to log return\n\n log(x+1)\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=log1p(pl.col('a')),\n )\n shape: (5, 2)\n ┌──────┬──────────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪══════════╡\n │ null ┆ null │\n │ -1 ┆ -inf │\n │ 0 ┆ 0.0 │\n │ 1 ┆ 0.693147 │\n │ 2 ┆ 1.098612 │\n └──────┴──────────┘\n ```\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.log1p()\n\n\ndef log2(x: Expr) -> Expr:\n \"\"\"以2为底的对数\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=log2(pl.col('a')),\n )\n shape: (5, 2)\n ┌──────┬──────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪══════╡\n │ null ┆ null │\n │ -1 ┆ NaN │\n │ 0 ┆ -inf │\n │ 1 ┆ 0.0 │\n │ 2 ┆ 1.0 │\n └──────┴──────┘\n ```\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.log(2)\n\n\ndef max_(a: Expr, b: Expr, *args) -> Expr:\n \"\"\"水平多列求最大值\n\n Maximum value of all inputs. At least 2 inputs are required.\"\"\"\n return max_horizontal(a, b, *args)\n\n\ndef mean(a: Expr, b: Expr, *args) -> Expr:\n \"\"\"水平多列求均值\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n 'b': [None, -1, 0, 0, 2],\n }).with_columns(\n out2=mean(pl.col('a'), 2),\n out3=mean(pl.col('a'), pl.col('b')),\n )\n shape: (5, 4)\n ┌──────┬──────┬──────┬──────┐\n │ a ┆ b ┆ out2 ┆ out3 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ f64 ┆ f64 │\n ╞══════╪══════╪══════╪══════╡\n │ null ┆ null ┆ 2.0 ┆ null │\n │ -1 ┆ -1 ┆ 0.5 ┆ -1.0 │\n │ 0 ┆ 0 ┆ 1.0 ┆ 0.0 │\n │ 1 ┆ 0 ┆ 1.5 ┆ 0.5 │\n │ 2 ┆ 2 ┆ 2.0 ┆ 2.0 │\n └──────┴──────┴──────┴──────┘\n ```\n\n \"\"\"\n return mean_horizontal(a, b, *args)\n\n\ndef min_(a: Expr, b: Expr, *args) -> Expr:\n \"\"\"水平多列求最小值\n\n Maximum value of all inputs. At least 2 inputs are required.\"\"\"\n return min_horizontal(a, b, *args)\n\n\ndef mod(x: Expr, y: Expr) -> Expr:\n \"\"\"求余\n\n x%y\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n 'b': [None, -1, 0, 0, 2],\n }).with_columns(\n out2=mod(pl.col('a'), 2),\n out3=mod(pl.col('a'), pl.col('b')),\n )\n shape: (5, 4)\n ┌──────┬──────┬──────┬──────┐\n │ a ┆ b ┆ out2 ┆ out3 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 ┆ i64 │\n ╞══════╪══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ null │\n │ -1 ┆ -1 ┆ 1 ┆ 0 │\n │ 0 ┆ 0 ┆ 0 ┆ null │\n │ 1 ┆ 0 ┆ 1 ┆ null │\n │ 2 ┆ 2 ┆ 0 ┆ 0 │\n └──────┴──────┴──────┴──────┘\n ```\n\n \"\"\"\n return x % y\n\n\ndef multiply(a: Expr, b: Expr, *args) -> Expr:\n \"\"\"水平多列相乘\n\n Multiply all inputs. At least 2 inputs are required.\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 2, 3, 4, None],\n 'b': [5, None, 3, 2, None],\n 'c': [1, 1, None, 1, None],\n }).with_columns(\n out=multiply(pl.col('a'), pl.col('b'), pl.col('c'))\n )\n\n shape: (5, 4)\n ┌──────┬──────┬──────┬──────┐\n │ a ┆ b ┆ c ┆ out │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 ┆ i64 │\n ╞══════╪══════╪══════╪══════╡\n │ null ┆ 5 ┆ 1 ┆ 5 │\n │ 2 ┆ null ┆ 1 ┆ 2 │\n │ 3 ┆ 3 ┆ null ┆ 9 │\n │ 4 ┆ 2 ┆ 1 ┆ 8 │\n │ null ┆ null ┆ null ┆ null │\n └──────┴──────┴──────┴──────┘\n\n ```\n\n Notes\n -----\n 全`null`时返回`null`\n\n \"\"\"\n _args = [a, b] + list(args)\n\n # # 全null返回1\n # return fold(acc=1, function=lambda acc, x: acc * x.fill_null(1), exprs=_args)\n return fold(acc=any_horizontal(_args), function=lambda acc, x: acc * x.fill_null(1), exprs=_args)\n\n\ndef power(x: Expr, y: Expr) -> Expr:\n \"\"\"乘幂\n\n x ** y\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n 'b': [None, -1, 0, 1, 2],\n }).with_columns(\n out2=power(pl.col('a'), 1),\n out3=power(pl.col('a'), pl.col('b')),\n )\n shape: (5, 4)\n ┌──────┬──────┬──────┬──────┐\n │ a ┆ b ┆ out2 ┆ out3 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 ┆ f64 │\n ╞══════╪══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ null │\n │ -1 ┆ -1 ┆ -1 ┆ -1.0 │\n │ 0 ┆ 0 ┆ 0 ┆ 1.0 │\n │ 1 ┆ 1 ┆ 1 ┆ 1.0 │\n │ 2 ┆ 2 ┆ 2 ┆ 4.0 │\n └──────┴──────┴──────┴──────┘\n ```\n\n Notes\n -----\n 负数的非整数幂在实数中未定义\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n\n if isinstance(y, (int, float)):\n return x.pow(y)\n\n return x.pow(y.cast(Float64))\n\n\ndef radians(x: Expr) -> Expr:\n \"\"\"角度转弧度\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.radians()\n\n\ndef reverse(x: Expr) -> Expr:\n \"\"\"求相反数\"\"\"\n return -x\n\n\ndef round_(x: Expr, decimals: int = 0) -> Expr:\n \"\"\"四舍五入\n\n Round input to closest integer.\n\n Parameters\n ----------\n x\n decimals\n Number of decimals to round to.\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 3.5, 4.5, -3.5, -4.5],\n }).with_columns(\n out1=round_(pl.col('a'), 0),\n out2=pl.col('a').map_elements(lambda x: round(x, 0), return_dtype=pl.Float64),\n )\n shape: (5, 3)\n ┌──────┬──────┬──────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ f64 ┆ f64 ┆ f64 │\n ╞══════╪══════╪══════╡\n │ null ┆ null ┆ null │\n │ 3.5 ┆ 4.0 ┆ 4.0 │\n │ 4.5 ┆ 5.0 ┆ 4.0 │\n │ -3.5 ┆ -4.0 ┆ -4.0 │\n │ -4.5 ┆ -5.0 ┆ -4.0 │\n └──────┴──────┴──────┘\n ```\n\n Notes\n -----\n 四舍五入，不是四舍六入五取偶（银行家舍入）\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.round(decimals)\n\n\ndef round_down(x: Expr, f: int = 1) -> Expr:\n \"\"\"小于输入的f的最大倍数\n\n Round input to greatest multiple of f less than input\n\n Parameters\n ----------\n x\n f\n\n Examples\n --------\n\n ```python\n df = pl.DataFrame({\n 'a': [None, 3.5, 4.5, -3.5, -4.5],\n }).with_columns(\n out=round_down(pl.col('a'), 2),\n )\n shape: (5, 2)\n ┌──────┬──────┐\n │ a ┆ out │\n │ --- ┆ --- │\n │ f64 ┆ f64 │\n ╞══════╪══════╡\n │ null ┆ null │\n │ 3.5 ┆ 2.0 │\n │ 4.5 ┆ 4.0 │\n │ -3.5 ┆ -4.0 │\n │ -4.5 ┆ -6.0 │\n └──────┴──────┘\n ```\n\n \"\"\"\n if f == 1:\n return x // 1\n else:\n return x // f * f\n\n\ndef s_log_1p(x: Expr) -> Expr:\n \"\"\"sign(x) * log10(1 + abs(x))\n\n 一种结合符号函数和对数变换的复合函数，常用于‌保留数据符号的同时压缩数值范围‌\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 9, -9, 99, -99],\n }).with_columns(\n out1=s_log_1p(pl.col('a')),\n )\n shape: (5, 2)\n ┌──────┬──────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪══════╡\n │ null ┆ null │\n │ 9 ┆ 1.0 │\n │ -9 ┆ -1.0 │\n │ 99 ┆ 2.0 │\n │ -99 ┆ -2.0 │\n └──────┴──────┘\n\n ```\n\n Notes\n -----\n 从`wq`示例可以看出，log的底数是10，而不是e\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#s_log_1px\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return (x.abs() + 1).log10() * x.sign()\n\n\ndef sign(x: Expr) -> Expr:\n \"\"\"符号函数\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.sign()\n\n\ndef signed_power(x: Expr, y: Expr) -> Expr:\n \"\"\"x的y次幂，符号保留\n\n x raised to the power of y such that final result preserves sign of x.\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n 'b': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=signed_power(pl.col('a'), 0),\n out2=signed_power(pl.col('a'), 1),\n out3=signed_power(pl.col('a'), 2),\n out4=signed_power(pl.col('a'), pl.col('b')),\n )\n\n shape: (5, 6)\n ┌──────┬──────┬──────┬──────┬──────┬──────┐\n │ a ┆ b ┆ out1 ┆ out2 ┆ out3 ┆ out4 │\n │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 ┆ f64 │\n ╞══════╪══════╪══════╪══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ null ┆ null ┆ null │\n │ -1 ┆ -1 ┆ -1 ┆ -1 ┆ -1 ┆ -1.0 │\n │ 0 ┆ 0 ┆ 0 ┆ 0 ┆ 0 ┆ 0.0 │\n │ 1 ┆ 1 ┆ 1 ┆ 1 ┆ 1 ┆ 1.0 │\n │ 2 ┆ 2 ┆ 1 ┆ 2 ┆ 4 ┆ 4.0 │\n └──────┴──────┴──────┴──────┴──────┴──────┘\n ```\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#signed_powerx-y\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n\n if isinstance(y, (int, float)):\n if y == 1:\n return x.abs() * x.sign()\n elif y == 0:\n return x.sign()\n else:\n return x.abs().pow(y) * x.sign()\n\n return x.abs().pow(y.cast(Float64)) * x.sign()\n\n\ndef sin(x: Expr) -> Expr:\n \"\"\"正弦\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.sin()\n\n\ndef sinh(x: Expr) -> Expr:\n \"\"\"双曲正弦\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.sinh()\n\n\ndef softsign(x: Expr) -> Expr:\n \"\"\"softsign激活函数\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=softsign(pl.col('a')),\n )\n\n shape: (5, 2)\n ┌──────┬──────────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪══════════╡\n │ null ┆ null │\n │ -1 ┆ -0.5 │\n │ 0 ┆ 0.0 │\n │ 1 ┆ 0.5 │\n │ 2 ┆ 0.666667 │\n └──────┴──────────┘\n ```\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x / (1 + x.abs())\n\n\ndef sqrt(x: Expr) -> Expr:\n \"\"\"平方根\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.sqrt()\n\n\ndef square(x: Expr) -> Expr:\n \"\"\"平方\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.pow(2)\n\n\ndef subtract(x: Expr, y: Expr) -> Expr:\n \"\"\"减法\n\n x-y\"\"\"\n return x - y\n\n\ndef tan(x: Expr) -> Expr:\n \"\"\"正切\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=tan(pl.col('a')),\n )\n\n shape: (5, 2)\n ┌──────┬───────────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪═══════════╡\n │ null ┆ null │\n │ -1 ┆ -1.557408 │\n │ 0 ┆ 0.0 │\n │ 1 ┆ 1.557408 │\n │ 2 ┆ -2.18504 │\n └──────┴───────────┘\n ```\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.tan()\n\n\ndef tanh(x: Expr) -> Expr:\n \"\"\"双曲正切\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2],\n }).with_columns(\n out1=tanh(pl.col('a')),\n )\n\n shape: (5, 2)\n ┌──────┬───────────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪═══════════╡\n │ null ┆ null │\n │ -1 ┆ -0.761594 │\n │ 0 ┆ 0.0 │\n │ 1 ┆ 0.761594 │\n │ 2 ┆ 0.964028 │\n └──────┴───────────┘\n ```\n\n \"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return x.tanh()\n\n\ndef var(a: Expr, b: Expr, *args) -> Expr:\n \"\"\"水平多列求方差\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 1, 1, 2],\n 'b': [None, -1, 0, 1, 2],\n 'c': [None, -1, 0, 2, None],\n }).with_columns(\n out1=var(pl.col('a'), pl.col('b'), pl.col('c')),\n )\n\n shape: (5, 4)\n ┌──────┬──────┬──────┬──────────┐\n │ a ┆ b ┆ c ┆ out1 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 ┆ f64 │\n ╞══════╪══════╪══════╪══════════╡\n │ null ┆ null ┆ null ┆ 0.0 │\n │ -1 ┆ -1 ┆ -1 ┆ 0.0 │\n │ 1 ┆ 0 ┆ 0 ┆ 0.666667 │\n │ 1 ┆ 1 ┆ 2 ┆ 0.666667 │\n │ 2 ┆ 2 ┆ null ┆ 0.0 │\n └──────┴──────┴──────┴──────────┘\n ```\n\n \"\"\"\n _args = [a, b] + list(args)\n _mean = mean_horizontal(_args)\n return sum_horizontal([(expr - _mean) ** 2 for expr in _args])\n\n\ndef std(a: Expr, b: Expr, *args) -> Expr:\n \"\"\"水平多列求标准差\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 1, 1, 2],\n 'b': [None, -1, 0, 1, 2],\n 'c': [None, -1, 0, 2, None],\n }).with_columns(\n out2=std(pl.col('a'), pl.col('b'), pl.col('c')),\n )\n\n shape: (5, 4)\n ┌──────┬──────┬──────┬──────────┐\n │ a ┆ b ┆ c ┆ out2 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 ┆ f64 │\n ╞══════╪══════╪══════╪══════════╡\n │ null ┆ null ┆ null ┆ 0.0 │\n │ -1 ┆ -1 ┆ -1 ┆ 0.0 │\n │ 1 ┆ 0 ┆ 0 ┆ 0.816497 │\n │ 1 ┆ 1 ┆ 2 ┆ 0.816497 │\n │ 2 ┆ 2 ┆ null ┆ 0.0 │\n └──────┴──────┴──────┴──────────┘\n ```\n\n \"\"\"\n return var(a, b, *args).sqrt()\n" }, { "path": "polars_ta/wq/cross_sectional.py", "content": "\"\"\"\n与`WorldQuant Alpha101`的区别是添加了`cs_`前缀\n\n由于截面与时序的使用方式不同，在自动化工具中如果不在名字上做区分就得手工注册，反而要麻烦些\n\n\"\"\"\nimport polars_ols as pls\nfrom polars import Expr, when, max_horizontal, UInt16, Int8, Utf8\nfrom polars_ols import OLSKwargs\n\n# In the original version, the function names are not prefixed with `cs_`,\n# here we add it to prevent confusion\n# 原版函数名都没有加`cs_`, 这里统一加一防止混淆\n\n\n_ols_kwargs = OLSKwargs(null_policy='drop', solve_method='svd')\n\n\ndef cs_one_side(x: Expr, is_long: bool = True) -> Expr:\n \"\"\"横截面上，将全部资产上调或下调，使得 Alpha 策略转为纯多头配置（当方向参数设为空头时则转为纯空头配置）\n\n Shifts all instruments up or down so that the Alpha becomes long-only or short-only\n(if side = short), respectively.\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -15, -7, 0, 20],\n 'b': [None, 15, 7, 0, 20],\n }).with_columns(\n out1=cs_one_side(pl.col('a'), True),\n out2=cs_one_side(pl.col('a'), False),\n out3=cs_one_side(pl.col('b'), True),\n out4=cs_one_side(pl.col('b'), False),\n )\n shape: (5, 6)\n ┌──────┬──────┬──────┬──────┬──────┬──────┐\n │ a ┆ b ┆ out1 ┆ out2 ┆ out3 ┆ out4 │\n │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 │\n ╞══════╪══════╪══════╪══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ null ┆ null ┆ null │\n │ -15 ┆ 15 ┆ 0 ┆ -35 ┆ 15 ┆ -5 │\n │ -7 ┆ 7 ┆ 8 ┆ -27 ┆ 7 ┆ -13 │\n │ 0 ┆ 0 ┆ 15 ┆ -20 ┆ 0 ┆ -20 │\n │ 20 ┆ 20 ┆ 35 ┆ 0 ┆ 20 ┆ 0 │\n └──────┴──────┴──────┴──────┴──────┴──────┘\n ```\n\n \"\"\"\n if is_long:\n return when(x.min() < 0).then(x - x.min()).otherwise(x)\n else:\n return when(x.max() > 0).then(x - x.max()).otherwise(x)\n\n\ndef cs_scale(x: Expr, scale_: float = 1, long_scale: float = 1, short_scale: float = 1) -> Expr:\n \"\"\"横截面上，将输入数据进行比例调整\n\n 此外，可通过向运算符添加额外参数，将多头头寸和空头头寸分别映射到独立的缩放比例上\n\n Scales input to booksize. We can also scale the long positions and short positions to separate scales by mentioning additional parameters to the operator.\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -15, -7, 0, 20],\n }).with_columns(\n out1=cs_scale(pl.col('a'), 1),\n out2=cs_scale(pl.col('a'), 1, 2, 3),\n )\n shape: (5, 3)\n ┌──────┬───────────┬───────────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ f64 ┆ f64 │\n ╞══════╪═══════════╪═══════════╡\n │ null ┆ null ┆ null │\n │ -15 ┆ -0.357143 ┆ -2.045455 │\n │ -7 ┆ -0.166667 ┆ -0.954545 │\n │ 0 ┆ 0.0 ┆ 0.0 │\n │ 20 ┆ 0.47619 ┆ 2.0 │\n └──────┴───────────┴───────────┘\n ```\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#scale-x-scale1-longscale1-shortscale1\n\n \"\"\"\n if long_scale != 1 or short_scale != 1:\n L = x.clip(lower_bound=0) # 全正数\n S = x.clip(upper_bound=0) # 全负数\n L = when(L.sum() == 0).then(0).otherwise(L / L.sum())\n S = when(S.sum() == 0).then(0).otherwise(S / S.sum()) # 负数/负数=正数\n return L * long_scale - S * short_scale\n else:\n return (x / x.abs().sum()).fill_nan(0) * scale_\n\n\ndef cs_scale_down(x: Expr, constant: int = 0) -> Expr:\n \"\"\"横截面上，将每日数据按比例缩放至 [0,1] 区间，使得最小值映射为 0，最大值映射为 1，并通过减去常数偏移量调整最终结果\n\n Scales all values in each day proportionately between 0 and 1 such that minimum value maps to 0 and maximum value maps to 1.\n constant is the offset by which final result is subtracted\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 15, 7, 0, 20],\n }).with_columns(\n out1=cs_scale_down(pl.col('a'), 0),\n out2=cs_scale_down(pl.col('a'), 1),\n )\n shape: (5, 3)\n ┌──────┬──────┬───────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ f64 ┆ f64 │\n ╞══════╪══════╪═══════╡\n │ null ┆ null ┆ null │\n │ 15 ┆ 0.75 ┆ -0.25 │\n │ 7 ┆ 0.35 ┆ -0.65 │\n │ 0 ┆ 0.0 ┆ -1.0 │\n │ 20 ┆ 1.0 ┆ 0.0 │\n └──────┴──────┴───────┘\n ```\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#scale_downxconstant0\n\n \"\"\"\n return ((x - x.min()) / (x.max() - x.min())).fill_nan(0) - constant\n\n\ndef cs_truncate(x: Expr, max_percent: float = 0.01) -> Expr:\n \"\"\"横截面上，将所有 x 的取值截断至 maxPercent 指定的上限值，其中 maxPercent 需以十进制小数形式表示\n\n Operator truncates all values of x to maxPercent. Here, maxPercent is in decimal notation\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [3, 7, 20, 6],\n }).with_columns(\n out=cs_truncate(pl.col('a'), 0.5),\n )\n shape: (4, 2)\n ┌─────┬─────┐\n │ a ┆ out │\n │ --- ┆ --- │\n │ i64 ┆ i64 │\n ╞═════╪═════╡\n │ 3 ┆ 3 │\n │ 7 ┆ 7 │\n │ 20 ┆ 18 │\n │ 6 ┆ 6 │\n └─────┴─────┘\n ```\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#truncatexmaxpercent001\n\n \"\"\"\n return x.clip(upper_bound=x.sum() * max_percent)\n\n\ndef cs_fill_except_all_null(x: Expr, value: float = 0) -> Expr:\n \"\"\"横截面上，全为`null`时，保持`null`，反之`null`填充为`value`\n\n Examples\n --------\n\n ```python\n df = pl.DataFrame({\n 'a': [1, 2, None, 4, None],\n 'b': [None, None, None, None, None],\n }).with_columns(\n A=cs_fill_except_all_null(pl.col('a')),\n B=cs_fill_except_all_null(pl.col('b')),\n )\n\n shape: (5, 4)\n ┌──────┬──────┬─────┬──────┐\n │ a ┆ b ┆ A ┆ B │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ null ┆ i64 ┆ i32 │\n ╞══════╪══════╪═════╪══════╡\n │ 1 ┆ null ┆ 1 ┆ null │\n │ 2 ┆ null ┆ 2 ┆ null │\n │ null ┆ null ┆ 0 ┆ null │\n │ 4 ┆ null ┆ 4 ┆ null │\n │ null ┆ null ┆ 0 ┆ null │\n └──────┴──────┴─────┴──────┘\n ```\n\n Notes\n -----\n 在权重矩阵中使用时。一定要保证所有股票都在，停牌不能被过滤了\n\n \"\"\"\n return when(x.is_not_null().sum() == 0).then(x).otherwise(x.fill_null(value))\n\n\ndef cs_fill_mean(x: Expr) -> Expr:\n \"\"\"横截面上，填充`null`为均值\"\"\"\n return x.fill_null(strategy='mean')\n\n\ndef cs_fill_max(x: Expr) -> Expr:\n \"\"\"横截面上，填充`null`为最大值\"\"\"\n return x.fill_null(strategy='max')\n\n\ndef cs_fill_min(x: Expr) -> Expr:\n \"\"\"横截面上，填充`null`为最小值\"\"\"\n return x.fill_null(strategy='min')\n\n\ndef cs_fill_null(x: Expr, value: float = 0) -> Expr:\n \"\"\"横截面上，填充`null`为`value`\"\"\"\n return x.fill_null(value)\n\n\ndef cs_regression_neut(y: Expr, x: Expr) -> Expr:\n \"\"\"横截面上，一元回归残差\"\"\"\n return pls.compute_least_squares(y, x, add_intercept=True, mode='residuals', ols_kwargs=_ols_kwargs)\n\n\ndef cs_regression_proj(y: Expr, x: Expr) -> Expr:\n \"\"\"横截面上，一元回归预测\"\"\"\n return pls.compute_least_squares(y, x, add_intercept=True, mode='predictions', ols_kwargs=_ols_kwargs)\n\n\ndef cs_rank(x: Expr, pct: bool = True) -> Expr:\n \"\"\"横截面排名\n\n Ranks the input among all the instruments and returns an equally distributed number between 0.0 and 1.0. For precise sort, use the rate as 0.\n\n Parameters\n ----------\n x\n pct\n * True: 排名百分比。范围：[0,1]\n * False: 排名。范围：[1,+inf)\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 1, 1, 2, 2, 3, 10],\n }).with_columns(\n out1=cs_rank(pl.col('a'), True),\n out2=cs_rank(pl.col('a'), False),\n )\n shape: (8, 3)\n ┌──────┬──────────┬──────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ f64 ┆ u32 │\n ╞══════╪══════════╪══════╡\n │ null ┆ null ┆ null │\n │ 1 ┆ 0.0 ┆ 1 │\n │ 1 ┆ 0.0 ┆ 1 │\n │ 1 ┆ 0.0 ┆ 1 │\n │ 2 ┆ 0.333333 ┆ 2 │\n │ 2 ┆ 0.333333 ┆ 2 │\n │ 3 ┆ 0.666667 ┆ 3 │\n │ 10 ┆ 1.0 ┆ 4 │\n └──────┴──────────┴──────┘\n ```\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#rankx-rate2\n\n \"\"\"\n if pct:\n # (x-x.min)/(x.max-x.min)\n r = x.rank(method='dense') - 1\n return r / max_horizontal(r.max(), 1)\n else:\n return x.rank(method='dense')\n\n\ndef cs_rank_if(condition: Expr, x: Expr, pct: bool = True) -> Expr:\n \"\"\"横截面筛选排名。可实现动态票池\n\n Parameters\n ----------\n condition:Expr\n 条件\n x:Expr\n 因子\n pct:bool\n 排名百分比。范围：[0,1]\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 1, 1, 2, 2, 3, 10],\n 'b': [1, 2, 3, 4, 5, 6, None, 8],\n }).with_columns(\n out1=cs_rank_if(True, pl.col('a'), True), # 与cs_rank等价\n out2=cs_rank_if(pl.col('b') > 3, -pl.col('a'), False),\n )\n\n shape: (8, 4)\n ┌──────┬──────┬──────────┬──────┐\n │ a ┆ b ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ f64 ┆ u32 │\n ╞══════╪══════╪══════════╪══════╡\n │ null ┆ 1 ┆ null ┆ null │\n │ 1 ┆ 2 ┆ 0.0 ┆ null │\n │ 1 ┆ 3 ┆ 0.0 ┆ null │\n │ 1 ┆ 4 ┆ 0.0 ┆ 3 │\n │ 2 ┆ 5 ┆ 0.333333 ┆ 2 │\n │ 2 ┆ 6 ┆ 0.333333 ┆ 2 │\n │ 3 ┆ null ┆ 0.666667 ┆ null │\n │ 10 ┆ 8 ┆ 1.0 ┆ 1 │\n └──────┴──────┴──────────┴──────┘\n ```\n\n Notes\n -----\n 已经产生了新的`None`，尽量避免之后再进行`ts_`时序计算。或按需调整`over_null`\n 或配合`cs_fill_null`等将`null`填充\n\n \"\"\"\n return cs_rank(when(condition).then(x).otherwise(None), pct)\n\n\ndef _cs_qcut_rank(x: Expr, q: int = 10) -> Expr:\n \"\"\"横截面上等频分箱\n\n Parameters\n ----------\n x\n q\n 按频率分成`q`份\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 1, 1, 2, 2, 3, 10],\n }).with_columns(\n out1=cs_qcut(pl.col('a'), 10),\n out2=cs_qcut_rank(pl.col('a'), 10),\n out3=pl.col('a').map_batches(lambda x: pd.qcut(x, 10, labels=False, duplicates='drop')),\n )\n shape: (8, 4)\n ┌──────┬──────┬──────┬──────┐\n │ a ┆ out1 ┆ out2 ┆ out3 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ u16 ┆ u16 ┆ f64 │\n ╞══════╪══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ NaN │\n │ 1 ┆ 0 ┆ 0 ┆ 0.0 │\n │ 1 ┆ 0 ┆ 0 ┆ 0.0 │\n │ 1 ┆ 0 ┆ 0 ┆ 0.0 │\n │ 2 ┆ 4 ┆ 3 ┆ 1.0 │\n │ 2 ┆ 4 ┆ 3 ┆ 1.0 │\n │ 3 ┆ 8 ┆ 6 ┆ 4.0 │\n │ 10 ┆ 9 ┆ 10 ┆ 5.0 │\n └──────┴──────┴──────┴──────┘\n ```\n\n Notes\n -----\n 使用`rank`来实现`qcut`的效果\n\n \"\"\"\n r = x.rank(method='dense') - 1\n return (r * q / max_horizontal(r.max(), 1)).cast(UInt16)\n\n\ndef cs_qcut(x: Expr, q: int = 10) -> Expr:\n \"\"\"横截面上等频分箱\n\n Convert float values into indexes for user-specified buckets. Bucket is useful for creating group values, which can be passed to group operators as input.\n\n Parameters\n ----------\n x\n q\n 按频率分成`q`份\n\n Examples\n --------\n\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 1, 1, 2, 2, 3, 10],\n }).with_columns(\n out1=cs_qcut(pl.col('a'), 10),\n out2=pl.col('a').map_batches(lambda x: pd.qcut(x, 10, labels=False, duplicates='drop')),\n )\n shape: (8, 3)\n ┌──────┬──────┬──────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ u16 ┆ f64 │\n ╞══════╪══════╪══════╡\n │ null ┆ null ┆ NaN │\n │ 1 ┆ 0 ┆ 0.0 │\n │ 1 ┆ 0 ┆ 0.0 │\n │ 1 ┆ 0 ┆ 0.0 │\n │ 2 ┆ 4 ┆ 1.0 │\n │ 2 ┆ 4 ┆ 1.0 │\n │ 3 ┆ 8 ┆ 4.0 │\n │ 10 ┆ 9 ┆ 5.0 │\n └──────┴──────┴──────┘\n\n ```\n\n Warnings\n --------\n 目前与`pd.qcut`结果不同，等官方改进\n\n \"\"\"\n # 实测直接to_physical()无法用于over,相当于with pl.StringCache():\n # return x.qcut(q, allow_duplicates=True).to_physical()\n\n return x.qcut(q, allow_duplicates=True, labels=[f'{i}' for i in range(q)]).cast(Utf8).cast(UInt16)\n\n\ndef cs_top_bottom(x: Expr, k: int = 10) -> Expr:\n \"\"\"横截面上，排名。前K标记成-1，后K标记成1\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 2, 2, 2, 3, 5, 5, 5, 10],\n }).with_columns(\n out1=pl.col('a').rank(method='min'),\n out2=pl.col('a').rank(method='dense'),\n out3=cs_top_bottom(pl.col('a'), 2),\n )\n shape: (10, 4)\n ┌──────┬──────┬──────┬──────┐\n │ a ┆ out1 ┆ out2 ┆ out3 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ u32 ┆ u32 ┆ i8 │\n ╞══════╪══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ null │\n │ 1 ┆ 1 ┆ 1 ┆ -1 │\n │ 2 ┆ 2 ┆ 2 ┆ -1 │\n │ 2 ┆ 2 ┆ 2 ┆ -1 │\n │ 2 ┆ 2 ┆ 2 ┆ -1 │\n │ 3 ┆ 5 ┆ 3 ┆ 0 │\n │ 5 ┆ 6 ┆ 4 ┆ 1 │\n │ 5 ┆ 6 ┆ 4 ┆ 1 │\n │ 5 ┆ 6 ┆ 4 ┆ 1 │\n │ 10 ┆ 9 ┆ 5 ┆ 1 │\n └──────┴──────┴──────┴──────┘\n \"\"\"\n\n # 值越小排第一，用来做空\n a = x.rank(method='dense')\n b = a.max() - a\n return (b < k).cast(Int8) - (a <= k).cast(Int8)\n" }, { "path": "polars_ta/wq/half_life.py", "content": "import numpy as np\nfrom polars import Expr\n\nfrom polars_ta.utils.numba_ import get_exponent_weights\n\n\ndef ts_mean_hl(x: Expr, d: int, half_life: int):\n \"\"\"滚动均值。带半衰期\"\"\"\n return x.fill_null(np.nan).rolling_mean(d, weights=get_exponent_weights(d, half_life)).fill_nan(None)\n\n\ndef ts_sum_hl(x: Expr, d: int, half_life: int):\n \"\"\"滚动求和。带半衰期\"\"\"\n return x.fill_null(np.nan).rolling_sum(d, weights=get_exponent_weights(d, half_life)).fill_nan(None)\n\n\ndef ts_std_hl(x: Expr, d: int, half_life: int):\n \"\"\"滚动标准差。带半衰期\"\"\"\n return x.fill_null(np.nan).rolling_std(d, weights=get_exponent_weights(d, half_life)).fill_nan(None)\n\n\ndef ts_var_hl(x: Expr, d: int, half_life: int):\n \"\"\"滚动方差。带半衰期\"\"\"\n return x.fill_null(np.nan).rolling_var(d, weights=get_exponent_weights(d, half_life)).fill_nan(None)\n\n# TODO 混动时序回归，带半衰期\n" }, { "path": "polars_ta/wq/logical.py", "content": "from polars import Expr, all_horizontal, any_horizontal, Boolean\nfrom polars import when\n\n\ndef and_(a: Expr, *args) -> Expr:\n \"\"\"Logical AND operator, returns true if both operands are true and returns false otherwise\"\"\"\n if len(args) == 0:\n return a\n return all_horizontal(a, *args)\n\n\ndef equal(input1: Expr, input2: Expr) -> Expr:\n \"\"\"Returns true if both inputs are same and returns false otherwise\"\"\"\n return input1 == input2\n\n\ndef if_else(condition: Expr, true_value: Expr, false_value: Expr = None) -> Expr:\n \"\"\"条件判断\n\n Notes\n -----\n 如果`true_value`或`false_value`使用了`None`,导致数据有`null`,\n 后面如再有`ts_`时序算子，一定要留意当前`over_null`值\n\n \"\"\"\n return when(condition).then(true_value).otherwise(false_value)\n\n\ndef is_finite(input1: Expr) -> Expr:\n \"\"\"If (input NaN or input == INF) return false, else return true\"\"\"\n return input1.is_finite()\n\n\ndef is_nan(input1: Expr) -> Expr:\n \"\"\"If (input == NaN) return true else return false\"\"\"\n return input1.is_nan()\n\n\ndef is_null(input1: Expr) -> Expr:\n \"\"\"If (input == null) return true else return false\"\"\"\n return input1.is_null()\n\n\ndef is_not_finite(input1: Expr) -> Expr:\n \"\"\"If (input NAN or input == INF) return true else return false\"\"\"\n return input1.is_infinite()\n\n\ndef is_not_nan(input1: Expr) -> Expr:\n \"\"\"If (input != NaN) return true else return false\"\"\"\n return input1.is_not_nan()\n\n\ndef is_not_null(input1: Expr) -> Expr:\n \"\"\"If (input != null) return true else return false\"\"\"\n return input1.is_not_null()\n\n\ndef less(input1: Expr, input2: Expr) -> Expr:\n \"\"\"If input1 < input2 return true, else return false\"\"\"\n return input1 < input2\n\n\ndef negate(input1: Expr) -> Expr:\n \"\"\"The result is true if the converted operand is false; the result is false if the converted operand is true\"\"\"\n return not_(input1)\n\n\ndef not_(input1: Expr) -> Expr:\n \"\"\"The result is true if the converted operand is false; the result is false if the converted operand is true\"\"\"\n return ~input1.cast(Boolean)\n\n\ndef or_(a: Expr, *args) -> Expr:\n \"\"\"Logical OR operator returns true if either or both inputs are true and returns false otherwise\"\"\"\n if len(args) == 0:\n return a\n return any_horizontal(a, *args)\n\n\ndef xor(a: Expr, b: Expr) -> Expr:\n \"\"\"Logical XOR operator returns true if exactly one of the inputs is true and returns false otherwise\"\"\"\n return a.xor(b)\n" }, { "path": "polars_ta/wq/preprocess.py", "content": "\"\"\"\n1. 补空值 → 去极值 → 标准化 → 中性化 → 标准化（可选二次标准化）\n2. 补空值 → 去极值 → 中性化 → 标准化\n\n# 对数市值。去极值\nMC_LOG = cs_quantile(log1p(market_cap), 0.01, 0.99)\n# 对数市值。标准化。供其他因子市值中性化时使用\nMC_NORM = cs_zscore(MC_LOG)\n# 对数市值。行业中性化。直接作为因子使用\nMC_NEUT = cs_zscore(cs_resid(MC_NORM, CS_SW_L1, ONE))\n\n\"\"\"\nimport polars_ols as pls\nfrom polars import Expr, when\nfrom polars_ols.least_squares import OLSKwargs\n\n\n# ======================\n# standardize\ndef cs_zscore(x: Expr, ddof: int = 0) -> Expr:\n \"\"\"横截面zscore标准化\"\"\"\n return (x - x.mean()) / x.std(ddof=ddof)\n\n\ndef cs_minmax(x: Expr) -> Expr:\n \"\"\"横截面minmax标准化\"\"\"\n a = x.min()\n b = x.max()\n # 这个版本在b-a为整数时，得到的结果不好看\n # return (x - a) / (b - a + TA_EPSILON)\n return when(a != b).then((x - a) / (b - a)).otherwise(0)\n\n\ndef cs_robust_scale(x: Expr) -> Expr:\n \"\"\"横截面robust scale标准化\"\"\"\n return (x - x.median()) / (x.quantile(0.75) - x.quantile(0.25))\n\n\n# ======================\n# winsorize\ndef cs_quantile(x: Expr, low_limit: float = 0.025, up_limit: float = 0.975) -> Expr:\n \"\"\"横截面分位数去极值\"\"\"\n a = x.quantile(low_limit)\n b = x.quantile(up_limit)\n return x.clip(lower_bound=a, upper_bound=b)\n\n\ndef cs_3sigma(x: Expr, n: float = 3.) -> Expr:\n \"\"\"横截面3倍sigma去极值\"\"\"\n # fill_nan will seriously reduce speed. So it's more appropriate for users to handle it themselves\n # fill_nan(None) 严重拖慢速度，所以还是由用户自己处理更合适\n a = x.mean()\n b = n * x.std(ddof=0)\n return x.clip(lower_bound=a - b, upper_bound=a + b)\n\n\ndef cs_mad(x: Expr, n: float = 3., k: float = 1.4826) -> Expr:\n \"\"\"横截面MAD去极值\n\n References\n ----------\n https://en.wikipedia.org/wiki/Median_absolute_deviation\n\n \"\"\"\n a = x.median()\n b = (n * k) * (x - a).abs().median()\n return x.clip(lower_bound=a - b, upper_bound=a + b)\n\n\n# ======================\n# neutralize\ndef cs_demean(x: Expr) -> Expr:\n \"\"\"横截面去均值化\n\n Notes\n -----\n Slower than multivariate regression. We need to groupby date and industry here,\n while multivariate regression only needs to add industry dummy variables and then groupby date\n\n Notes\n -----\n 速度没有多元回归快，因为这里需要按日期行业groupby，\n 而多元回归只要添加行业哑变量，然后按日期groupby即可\n\n \"\"\"\n return x - x.mean()\n\n\n# ======================\n# neutralize\n_ols_kwargs = OLSKwargs(null_policy='drop', solve_method='svd')\n\n\ndef cs_resid(y: Expr, *more_x: Expr) -> Expr:\n \"\"\"横截面多元回归取残差\"\"\"\n return pls.compute_least_squares(y, *more_x, mode='residuals', ols_kwargs=_ols_kwargs)\n\n\ndef cs_zscore_resid(y: Expr, *more_x: Expr) -> Expr:\n \"\"\"横截面标准化、中性化\"\"\"\n return cs_resid(cs_zscore(y), *more_x)\n\n\ndef cs_resid_zscore(y: Expr, *more_x: Expr) -> Expr:\n \"\"\"横截面中性化、标准化\"\"\"\n return cs_resid(cs_zscore(y), *more_x)\n\n\ndef cs_mad_zscore(y: Expr) -> Expr:\n \"\"\"横截面MAD去极值、标准化\"\"\"\n return cs_zscore(cs_mad(y))\n\n\ndef cs_mad_zscore_resid(y: Expr, *more_x: Expr) -> Expr:\n \"\"\"横截面MAD去极值、标准化、中性化\"\"\"\n return cs_resid(cs_zscore(cs_mad(y)), *more_x)\n\n\ndef cs_mad_zscore_resid_zscore(y: Expr, *more_x: Expr) -> Expr:\n \"\"\"横截面去MAD极值、标准化、中性化、二次标准化\"\"\"\n return cs_zscore(cs_resid(cs_zscore(cs_mad(y)), *more_x))\n\n\ndef cs_quantile_zscore(y: Expr, low_limit: float = 0.025, up_limit: float = 0.975) -> Expr:\n \"\"\"横截面分位数去极值、标准化\"\"\"\n return cs_zscore(cs_quantile(y, low_limit, up_limit))\n\n\n# ==========================\ndef cs_resid_w(w: Expr, y: Expr, *more_x: Expr) -> Expr:\n \"\"\"横截面加权多元回归取残差\n\n Barra中权重采用流通市值的平方根\n \"\"\"\n return pls.compute_least_squares(y, *more_x, sample_weights=w, mode='residuals', ols_kwargs=_ols_kwargs)\n" }, { "path": "polars_ta/wq/time_series.py", "content": "import itertools\nfrom typing import Optional\n\nimport more_itertools\nimport numpy as np\nimport polars_ols as pls\nfrom polars import Expr, UInt16, struct, when, Struct, Field, Float64, Boolean, UInt32, all_horizontal, any_horizontal\nfrom polars import rolling_corr, rolling_cov\nfrom polars_ols import RollingKwargs\n\nimport polars_ta\nfrom polars_ta.utils.numba_ import batches_i1_o1, batches_i2_o1, batches_i2_o2, struct_to_numpy\nfrom polars_ta.wq._nb import roll_argmax, roll_argmin, roll_co_kurtosis, roll_co_skewness, roll_moment, roll_partial_corr, roll_triple_corr, _cum_prod_by, _cum_sum_by, _signals_to_size, \\\n _cum_sum_reset, _sum_split_by, roll_prod\n\n\ndef ts_arg_max(x: Expr, d: int = 5, reverse: bool = True, min_samples: Optional[int] = None) -> Expr:\n \"\"\"最大值相对位置\n\n 最近的一天记为第 0 天，最远的一天为第 d-1 天\n\n Returns the relative index of the max value in the time series for the past d days.\n If the current day has the max value for the past d days, it returns 0.\n If previous day has the max value for the past d days, it returns 1.\n\n Parameters\n ----------\n x\n d\n reverse\n 反向\n min_samples\n\n See Also\n --------\n ts_arg_min\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [0., 1., 0.,np.nan, 1., 0.],\n }).with_columns(\n out3=ts_arg_max(pl.col('a'), 3),\n out2=ts_arg_max(pl.col('a'), 3, min_samples=2),\n out1=ts_arg_max(pl.col('a'), 3, min_samples=1),\n )\n\n shape: (6, 4)\n ┌─────┬──────┬──────┬──────┐\n │ a ┆ out3 ┆ out2 ┆ out1 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ f64 ┆ u16 ┆ u16 ┆ u16 │\n ╞═════╪══════╪══════╪══════╡\n │ 0.0 ┆ null ┆ null ┆ 0 │\n │ 1.0 ┆ null ┆ 0 ┆ 0 │\n │ 0.0 ┆ 1 ┆ 1 ┆ 1 │\n │ NaN ┆ null ┆ null ┆ null │\n │ 1.0 ┆ null ┆ null ┆ 0 │\n │ 0.0 ┆ null ┆ 1 ┆ 1 │\n └─────┴──────┴──────┴──────┘\n ```\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#ts_arg_maxx-d\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_argmax, d, minp, reverse, dtype=UInt16), return_dtype=UInt16)\n\n\ndef ts_arg_min(x: Expr, d: int = 5, reverse: bool = True, min_samples: Optional[int] = None) -> Expr:\n \"\"\"最小值相对位置\n\n 最近的一天记为第 0 天，最远的一天为第 d-1 天\n\n Parameters\n ----------\n x\n d\n reverse\n 反向\n min_samples\n\n See Also\n --------\n ts_arg_max\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [0., 1., 0.,np.nan, 1., 0.],\n }).with_columns(\n out3=ts_arg_min(pl.col('a'), 3),\n out2=ts_arg_min(pl.col('a'), 3, min_samples=2),\n out1=ts_arg_min(pl.col('a'), 3, min_samples=1),\n )\n\n shape: (6, 4)\n ┌─────┬──────┬──────┬──────┐\n │ a ┆ out3 ┆ out2 ┆ out1 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ f64 ┆ u16 ┆ u16 ┆ u16 │\n ╞═════╪══════╪══════╪══════╡\n │ 0.0 ┆ null ┆ null ┆ 0 │\n │ 1.0 ┆ null ┆ 1 ┆ 1 │\n │ 0.0 ┆ 0 ┆ 0 ┆ 0 │\n │ NaN ┆ null ┆ null ┆ null │\n │ 1.0 ┆ null ┆ null ┆ 1 │\n │ 0.0 ┆ null ┆ 0 ┆ 0 │\n └─────┴──────┴──────┴──────┘\n ```\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#ts_arg_minx-d\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_argmin, d, minp, reverse, dtype=UInt16), return_dtype=UInt16)\n\n\ndef ts_co_kurtosis(x: Expr, y: Expr, d: int = 5, ddof: int = 0, min_samples: Optional[int] = None) -> Expr:\n \"\"\"计算两个序列在滚动窗口内联合分布的协峰度\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return struct(f0=x, f1=y).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), roll_co_kurtosis, d, minp), return_dtype=Float64)\n\n\ndef ts_co_skewness(x: Expr, y: Expr, d: int = 5, ddof: int = 0, min_samples: Optional[int] = None) -> Expr:\n \"\"\"计算两个序列在滚动窗口内联合分布的协偏度\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return struct(f0=x, f1=y).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), roll_co_skewness, d, minp), return_dtype=Float64)\n\n\ndef ts_corr(x: Expr, y: Expr, d: int = 5, ddof: int = 1, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动相关系数\n\n rolling correlation between two columns\n\n Parameters\n ----------\n x\n y\n d\n ddof\n 自由度\n min_samples\n\n Notes\n -----\n x、y不区分先后\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return rolling_corr(x, y, window_size=d, ddof=ddof, min_samples=minp)\n\n\ndef ts_count(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动计数\n\n Parameters\n ----------\n x\n d\n min_samples\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2, 3],\n 'b': [None, True, True, True, False, False],\n }).with_columns(\n out1=ts_count(pl.col('a'), 3),\n out2=ts_count(pl.col('b'), 3),\n )\n\n shape: (6, 4)\n ┌──────┬───────┬──────┬──────┐\n │ a ┆ b ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ bool ┆ u32 ┆ u32 │\n ╞══════╪═══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ null │\n │ -1 ┆ true ┆ null ┆ null │\n │ 0 ┆ true ┆ null ┆ null │\n │ 1 ┆ true ┆ 2 ┆ 3 │\n │ 2 ┆ false ┆ 2 ┆ 2 │\n │ 3 ┆ false ┆ 3 ┆ 1 │\n └──────┴───────┴──────┴──────┘\n ```\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.cast(Boolean).cast(UInt32).rolling_sum(d, min_samples=minp)\n\n\ndef ts_count_eq(x: Expr, d: int = 30, n: int = 10, min_samples: Optional[int] = None) -> Expr:\n \"\"\"D天内最近连续出现N次\n\n Parameters\n ----------\n x\n d: int\n 窗口大小\n n: int\n 连续出现次数\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n xx = x.cast(Boolean).cast(UInt32)\n return (xx.rolling_sum(n) == n) & (xx.rolling_sum(d, min_samples=minp) == n)\n\n\ndef ts_count_ge(x: Expr, d: int = 30, n: int = 10, min_samples: Optional[int] = None) -> Expr:\n \"\"\"D天内最近连续出现至少N次\n\n Parameters\n ----------\n x\n d: int\n 窗口大小\n n: int\n 至少连续出现次数\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n xx = x.cast(Boolean).cast(UInt32)\n return (xx.rolling_sum(n) == n) & (xx.rolling_sum(d, min_samples=minp) >= n)\n\n\ndef ts_count_nans(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动统计nan出现次数\n\n Parameters\n ----------\n x\n d\n min_samples\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, float('nan'), -1, 0, 1, 2, 3],\n }).with_columns(\n out1=ts_count_nans(pl.col('a'), 3),\n )\n\n shape: (7, 2)\n ┌──────┬──────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ f64 ┆ u32 │\n ╞══════╪══════╡\n │ null ┆ null │\n │ NaN ┆ null │\n │ -1.0 ┆ null │\n │ 0.0 ┆ 1 │\n │ 1.0 ┆ 0 │\n │ 2.0 ┆ 0 │\n │ 3.0 ┆ 0 │\n └──────┴──────┘\n ```\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.is_nan().cast(UInt32).rolling_sum(d, min_samples=minp)\n\n\ndef ts_count_nulls(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动统计null出现次数\n\n Parameters\n ----------\n x\n d\n min_samples\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 0, 1, 2, 3],\n 'b': [None, True, True, True, False, False],\n }).with_columns(\n out1=ts_count_nulls(pl.col('a'), 3),\n out2=ts_count_nulls(pl.col('b'), 3),\n )\n shape: (6, 4)\n ┌──────┬───────┬──────┬──────┐\n │ a ┆ b ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ bool ┆ u32 ┆ u32 │\n ╞══════╪═══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ null │\n │ -1 ┆ true ┆ null ┆ null │\n │ 0 ┆ true ┆ 1 ┆ 1 │\n │ 1 ┆ true ┆ 0 ┆ 0 │\n │ 2 ┆ false ┆ 0 ┆ 0 │\n │ 3 ┆ false ┆ 0 ┆ 0 │\n └──────┴───────┴──────┴──────┘\n ```\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.is_null().cast(UInt32).rolling_sum(d, min_samples=minp)\n\n\ndef ts_covariance(x: Expr, y: Expr, d: int = 5, ddof: int = 1, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动协方差\n\n rolling covariance between two columns\n\n Parameters\n ----------\n x\n y\n d\n ddof\n 自由度\n min_samples\n\n Notes\n -----\n x、y不区分先后\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return rolling_cov(x, y, window_size=d, ddof=ddof, min_samples=minp)\n\n\ndef ts_cum_count(x: Expr) -> Expr:\n \"\"\"时序累计计数\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, None, -1, 0, 1, 2],\n 'b': [None, None, True, True, True, False],\n }).with_columns(\n out1=ts_cum_count(pl.col('a')),\n out2=ts_cum_count(pl.col('b')),\n )\n shape: (6, 4)\n ┌──────┬───────┬──────┬──────┐\n │ a ┆ b ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ bool ┆ u32 ┆ u32 │\n ╞══════╪═══════╪══════╪══════╡\n │ null ┆ null ┆ 0 ┆ 0 │\n │ null ┆ null ┆ 0 ┆ 0 │\n │ -1 ┆ true ┆ 1 ┆ 1 │\n │ 0 ┆ true ┆ 2 ┆ 2 │\n │ 1 ┆ true ┆ 3 ┆ 3 │\n │ 2 ┆ false ┆ 4 ┆ 4 │\n └──────┴───────┴──────┴──────┘\n ```\n\n \"\"\"\n return x.cum_count()\n\n\ndef ts_cum_max(x: Expr) -> Expr:\n \"\"\"时序累计最大值\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, None, -1, 0, 2, 1],\n 'b': [None, None, True, False, False, True],\n }).with_columns(\n out1=ts_cum_max(pl.col('a')),\n out2=ts_cum_max(pl.col('b')),\n )\n shape: (6, 4)\n ┌──────┬───────┬──────┬──────┐\n │ a ┆ b ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ bool ┆ i64 ┆ bool │\n ╞══════╪═══════╪══════╪══════╡\n │ null ┆ null ┆ null ┆ null │\n │ null ┆ null ┆ null ┆ null │\n │ -1 ┆ true ┆ -1 ┆ true │\n │ 0 ┆ false ┆ 0 ┆ true │\n │ 2 ┆ false ┆ 2 ┆ true │\n │ 1 ┆ true ┆ 2 ┆ true │\n └──────┴───────┴──────┴──────┘\n ```\n\n \"\"\"\n return x.cum_max()\n\n\ndef ts_cum_min(x: Expr) -> Expr:\n \"\"\"时序累计最小值\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, None, -1, 0, -2, 1],\n 'b': [None, None, True, False, False, True],\n }).with_columns(\n out1=ts_cum_min(pl.col('a')),\n out2=ts_cum_min(pl.col('b')),\n )\n shape: (6, 4)\n ┌──────┬───────┬──────┬───────┐\n │ a ┆ b ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ bool ┆ i64 ┆ bool │\n ╞══════╪═══════╪══════╪═══════╡\n │ null ┆ null ┆ null ┆ null │\n │ null ┆ null ┆ null ┆ null │\n │ -1 ┆ true ┆ -1 ┆ true │\n │ 0 ┆ false ┆ -1 ┆ false │\n │ -2 ┆ false ┆ -2 ┆ false │\n │ 1 ┆ true ┆ -2 ┆ false │\n └──────┴───────┴──────┴───────┘\n ```\n \"\"\"\n return x.cum_min()\n\n\ndef ts_cum_prod(x: Expr) -> Expr:\n \"\"\"时序累乘\"\"\"\n return x.cum_prod()\n\n\ndef ts_cum_sum(x: Expr) -> Expr:\n \"\"\"时序累加\"\"\"\n return x.cum_sum()\n\n\ndef ts_cum_sum_reset(x: Expr) -> Expr:\n \"\"\"时序累加。遇到0、nan、相反符号时重置\n\n Examples\n --------\n\n ```python\n df = pl.DataFrame({\n 'a': [1, 0, 1, 2, None, 3, -2, -3],\n }).with_columns(\n A=ts_cum_sum_reset(pl.col('a'))\n )\n\n shape: (8, 2)\n ┌──────┬──────┐\n │ a ┆ A │\n │ --- ┆ --- │\n │ i64 ┆ f64 │\n ╞══════╪══════╡\n │ 1 ┆ 1.0 │\n │ 0 ┆ 0.0 │\n │ 1 ┆ 1.0 │\n │ 2 ┆ 3.0 │\n │ null ┆ 0.0 │\n │ 3 ┆ 3.0 │\n │ -2 ┆ -2.0 │\n │ -3 ┆ -5.0 │\n └──────┴──────┘\n\n ```\n\n \"\"\"\n return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), _cum_sum_reset), return_dtype=Float64)\n\n\ndef ts_decay_exp_window(x: Expr, d: int = 30, factor: float = 1.0, min_samples: Optional[int] = None) -> Expr:\n \"\"\"指数衰减移动平均\n\n Examples\n --------\n ```python\n from polars_ta.wq.time_series import ts_decay_linear, ts_decay_exp_window\n\n df = pl.DataFrame({\n 'a': [None, 6, 5, 4, 5, 30],\n }).with_columns(\n out1=ts_decay_linear(pl.col('a'), 5),\n out2=ts_decay_exp_window(pl.col('a'), 5, 0.5),\n )\n shape: (6, 3)\n ┌──────┬──────┬───────────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ f64 ┆ f64 │\n ╞══════╪══════╪═══════════╡\n │ null ┆ null ┆ null │\n │ 6 ┆ null ┆ null │\n │ 5 ┆ null ┆ null │\n │ 4 ┆ null ┆ null │\n │ 5 ┆ null ┆ null │\n │ 30 ┆ 13.2 ┆ 17.806452 │\n └──────┴──────┴───────────┘\n ```\n\n Parameters\n ----------\n x\n d\n factor\n 衰减系数\n min_samples\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#ts_decay_exp_windowx-d-factor-10-nan-true\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n weights = np.repeat(factor, d) ** np.arange(d - 1, -1, -1)\n # print(weights)\n # pyo3_runtime.PanicException: weights not yet supported on array with null values\n return x.fill_null(np.nan).rolling_mean(d, weights=weights, min_samples=minp).fill_nan(None)\n\n\ndef ts_decay_linear(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:\n \"\"\"线性衰减移动平均\n\n Examples\n --------\n ```python\n from polars_ta.talib import WMA as ts_WMA\n from polars_ta.wq.time_series import ts_decay_linear\n\n df = pl.DataFrame({\n 'a': [None, 6, 5, 4, 5, 30],\n }).with_columns(\n out1=ts_decay_linear(pl.col('a'), 5),\n out2=ts_WMA(pl.col('a'), 5),\n )\n\n shape: (6, 3)\n ┌──────┬──────┬──────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ f64 ┆ f64 │\n ╞══════╪══════╪══════╡\n │ null ┆ null ┆ null │\n │ 6 ┆ null ┆ null │\n │ 5 ┆ null ┆ null │\n │ 4 ┆ null ┆ null │\n │ 5 ┆ null ┆ null │\n │ 30 ┆ 13.2 ┆ 13.2 │\n └──────┴──────┴──────┘\n ```\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#ts_decay_linearx-d-dense-false\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n weights = np.arange(1, d + 1)\n # print(weights)\n # pyo3_runtime.PanicException: weights not yet supported on array with null values\n # null换成NaN就不报错了，再换回来\n return x.fill_null(np.nan).rolling_mean(d, weights=weights, min_samples=minp).fill_nan(None)\n\n\ndef ts_delay(x: Expr, d: int = 1, fill_value: float = None) -> Expr:\n \"\"\"时序数据移动\n\n shift x\n\n Parameters\n ----------\n x\n d\n 向前或向后的移动天数\n fill_value\n 填充。可用None、常量或Expr\n\n \"\"\"\n return x.shift(d, fill_value=fill_value)\n\n\ndef ts_delta(x: Expr, d: int = 1) -> Expr:\n \"\"\"时序差分\"\"\"\n return x.diff(d)\n\n\ndef ts_fill_null(x: Expr, limit: int = None) -> Expr:\n \"\"\"用上一个非空值填充空值\n\n Parameters\n ----------\n x\n limit\n 最大填充次数\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, -1, 1, None, None],\n 'b': [None, True, False, None, None],\n }).with_columns(\n out1=ts_fill_null(pl.col('a')),\n out2=ts_fill_null(pl.col('b'), 1),\n )\n\n shape: (5, 4)\n ┌──────┬───────┬──────┬───────┐\n │ a ┆ b ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- ┆ --- │\n │ i64 ┆ bool ┆ i64 ┆ bool │\n ╞══════╪═══════╪══════╪═══════╡\n │ null ┆ null ┆ null ┆ null │\n │ -1 ┆ true ┆ -1 ┆ true │\n │ 1 ┆ false ┆ 1 ┆ false │\n │ null ┆ null ┆ 1 ┆ false │\n │ null ┆ null ┆ 1 ┆ null │\n └──────┴───────┴──────┴───────┘\n ```\n\n \"\"\"\n return x.forward_fill(limit)\n\n\ndef ts_ir(x: Expr, d: int = 1, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动信息系数\n\n rolling information ratio\"\"\"\n return ts_mean(x, d, min_samples) / ts_std_dev(x, d, 0, min_samples)\n\n\ndef ts_kurtosis(x: Expr, d: int = 5, bias: bool = False, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动峰度\n\n kurtosis of x for the last d days\n\n Parameters\n ----------\n x\n d\n bias\n 有偏\n min_samples\n\n Notes\n -----\n `bias=False`时与`pandas`结果一样\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 2, 3, 4, 999],\n }).with_columns(\n out1=pl.col('a').map_batches(lambda x: pl.Series(pd.Series(x).rolling(4).kurt())),\n out2=ts_kurtosis(pl.col('a'), 4),\n )\n\n shape: (6, 3)\n ┌──────┬──────────┬──────────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ f64 ┆ f64 │\n ╞══════╪══════════╪══════════╡\n │ null ┆ null ┆ null │\n │ 1 ┆ null ┆ null │\n │ 2 ┆ null ┆ null │\n │ 3 ┆ null ┆ null │\n │ 4 ┆ -1.2 ┆ -1.2 │\n │ 999 ┆ 3.999946 ┆ 3.999946 │\n └──────┴──────────┴──────────┘\n ```\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_kurtosis(d, min_samples=minp, bias=bias)\n\n\ndef ts_l2_norm(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"欧几里得范数\n\n Euclidean norm\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.pow(2).rolling_sum(d, min_samples=minp).sqrt()\n\n\ndef ts_log_diff(x: Expr, d: int = 1) -> Expr:\n \"\"\"求对数，然后时序滚动差分\n\n log(current value of input or x[t] ) - log(previous value of input or x[t-1]).\n \"\"\"\n return x.log().diff(d)\n\n\ndef ts_max(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动最大值\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_max(d, min_samples=minp)\n\n\ndef ts_max_diff(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:\n \"\"\"窗口内最大值与当前值的差异‌\n\n x - ts_max(x, d)\"\"\"\n return x - ts_max(x, d, min_samples)\n\n\ndef ts_mean(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"简单移动平均\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_mean(d, min_samples=minp)\n\n\ndef ts_median(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动中位数\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_median(d, min_samples=minp)\n\n\ndef ts_min(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动最小值\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_min(d, min_samples=minp)\n\n\ndef ts_min_diff(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:\n \"\"\"窗口内最小值与当前值的差异‌\n\n x - ts_min(x, d)\"\"\"\n return x - ts_min(x, d, min_samples)\n\n\ndef ts_min_max_cps(x: Expr, d: int, f: float = 2.0, min_samples: Optional[int] = None) -> Expr:\n \"\"\"计算时间窗口内最小值与最大值的总和减去当前值的加权结果\n\n (ts_min(x, d) + ts_max(x, d)) - f * x\"\"\"\n return (ts_min(x, d, min_samples) + ts_max(x, d, min_samples)) - f * x\n\n\ndef ts_min_max_diff(x: Expr, d: int, f: float = 0.5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"计算当前值 x 与基于时间窗口内最小值、最大值的加权组合的差值\n\n x - f * (ts_min(x, d) + ts_max(x, d))\"\"\"\n return x - f * (ts_min(x, d, min_samples) + ts_max(x, d, min_samples))\n\n\ndef ts_moment(x: Expr, d: int, k: int = 0, min_samples: Optional[int] = None) -> Expr:\n \"\"\"滚动k阶中心距\n\n Returns K-th central moment of x for the past d days.\n\n Parameters\n ----------\n x\n d\n k\n min_samples\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_moment, d, minp, k), return_dtype=Float64)\n\n\ndef ts_partial_corr(x: Expr, y: Expr, z: Expr, d: int, min_samples: Optional[int] = None) -> Expr:\n \"\"\"滚动偏相关\n\n Returns partial correlation of x, y, z for the past d days.\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return struct(f0=x, f1=y, f2=z).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3), roll_partial_corr, d, minp), return_dtype=Float64)\n\n\ndef ts_percentage(x: Expr, d: int, percentage: float = 0.5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"滚动百分位数\n\n Returns percentile value of x for the past d days.\n\n Parameters\n ----------\n x\n d\n percentage\n min_samples\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_quantile(percentage, window_size=d, min_samples=minp)\n\n\ndef ts_product(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动乘\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return x.map_batches(lambda x1: batches_i1_o1(x1.to_numpy(), roll_prod, d, minp), return_dtype=Float64)\n\n\ndef ts_rank(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动排名\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_rank(d, min_samples=minp) / x.is_not_null().cast(UInt32).rolling_sum(d, min_samples=minp)\n\n\ndef ts_realized_volatility(close: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"已实现波动率\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return ts_log_diff(close, 1).rolling_std(d, ddof=0, min_samples=minp)\n\n\ndef ts_returns(x: Expr, d: int = 1) -> Expr:\n \"\"\"简单收益率\"\"\"\n return x.pct_change(d)\n\n\ndef ts_scale(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动缩放。相当于ts_minmax\n\n Returns (x – ts_min(x, d)) / (ts_max(x, d) – ts_min(x, d)) + constant\n\n \"\"\"\n a = ts_min(x, d, min_samples)\n b = ts_max(x, d, min_samples)\n # return (x - a) / (b - a + TA_EPSILON)\n return when(a != b).then((x - a) / (b - a)).otherwise(0)\n\n\ndef ts_shifts_v1(*args) -> Expr:\n \"\"\"时序上按顺序进行平移，然后逻辑与\n\n 比如今天反包，昨天阴跌，前天涨停。只需要写`ts_shifts_v1(反包,阴跌,涨停)`。使用此函数能一定程度上简化代码\n\n Parameters\n ----------\n args\n pl.col按照顺序排列\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, False, False, True, True, True],\n 'b': [None, False, True, True, True, True],\n 'c': [None, True, True, True, True, True],\n }).with_columns(\n out1=ts_shifts_v1(pl.col('a'), pl.col('b'), pl.col('c')),\n out2=pl.col('a').shift(0) & pl.col('b').shift(1) & pl.col('c').shift(2),\n )\n\n shape: (6, 5)\n ┌──────┬───────┬───────┬───────┬───────┐\n │ c ┆ b ┆ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n │ bool ┆ bool ┆ bool ┆ bool ┆ bool │\n ╞══════╪═══════╪═══════╪═══════╪═══════╡\n │ null ┆ null ┆ null ┆ null ┆ null │\n │ true ┆ false ┆ false ┆ false ┆ false │\n │ true ┆ true ┆ false ┆ false ┆ false │\n │ true ┆ true ┆ true ┆ true ┆ true │\n │ true ┆ true ┆ true ┆ true ┆ true │\n │ true ┆ true ┆ true ┆ true ┆ true │\n └──────┴───────┴───────┴───────┴───────┘\n ```\n\n \"\"\"\n return all_horizontal(arg.shift(i) for i, arg in enumerate(args))\n\n\ndef ts_shifts_v2(*args) -> Expr:\n \"\"\"时序上按顺序进行平移，然后逻辑与\n\n 遇到连续条件时可简化参数。如连续3天涨停后连续2天跌停，可用`ts_shifts_v2(跌停,2,涨停,3)`\n 另一种实现方法是：`ts_delay((ts_count(涨停,3)==3),2)&(ts_count(跌停,2)==2)`\n\n Parameters\n ----------\n args\n pl.col, repeat。两个参数循环\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, False, False, True, True, True],\n 'b': [None, False, True, True, True, True],\n 'c': [None, True, True, True, True, True],\n }).with_columns(\n out1=ts_shifts_v1(pl.col('a'), pl.col('b'), pl.col('b')),\n out2=ts_shifts_v2(pl.col('a'), 1, pl.col('b'), 2),\n out3=pl.col('a').shift(0) & pl.col('b').shift(1) & pl.col('b').shift(2),\n )\n\n shape: (6, 6)\n ┌───────┬───────┬──────┬───────┬───────┬───────┐\n │ a ┆ b ┆ c ┆ out1 ┆ out2 ┆ out3 │\n │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n │ bool ┆ bool ┆ bool ┆ bool ┆ bool ┆ bool │\n ╞═══════╪═══════╪══════╪═══════╪═══════╪═══════╡\n │ null ┆ null ┆ null ┆ null ┆ null ┆ null │\n │ false ┆ false ┆ true ┆ false ┆ false ┆ false │\n │ false ┆ true ┆ true ┆ false ┆ false ┆ false │\n │ true ┆ true ┆ true ┆ false ┆ false ┆ false │\n │ true ┆ true ┆ true ┆ true ┆ true ┆ true │\n │ true ┆ true ┆ true ┆ true ┆ true ┆ true │\n └───────┴───────┴──────┴───────┴───────┴───────┘\n ```\n\n \"\"\"\n return ts_shifts_v1(*itertools.chain.from_iterable([item] * count for item, count in more_itertools.chunked(args, 2)))\n\n\ndef ts_shifts_v3(*args) -> Expr:\n \"\"\"时序上按顺序进行平移，然后逻辑或\n\n 如：涨停后连续下跌1~3天。这个案例会导致涨停可能出现在动态的一天。`ts_shifts_v3(下跌,1,3,涨停,1,1)`\n 它本质上是`ts_shifts_v2(下跌,1,涨停,1)|ts_shifts_v2(下跌,2,涨停,1)|ts_shifts_v2(下跌,3,涨停,1)|`\n\n Parameters\n ----------\n args\n pl.col, [start, end]。三个参数循环。start/end都是闭区间\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, False, False, True, True, True],\n 'b': [None, False, True, True, True, True],\n 'c': [None, True, True, True, True, True],\n }).with_columns(\n out0=ts_shifts_v3(pl.col('a'), 1, 3, pl.col('b'), 2, 2),\n out1=ts_shifts_v2(pl.col('a'), 1, pl.col('b'), 2),\n out2=ts_shifts_v2(pl.col('a'), 2, pl.col('b'), 2),\n out3=ts_shifts_v2(pl.col('a'), 3, pl.col('b'), 2)\n ).with_columns(\n out4=pl.col('out1') | pl.col('out2') | pl.col('out3')\n )\n\n shape: (6, 8)\n ┌───────┬───────┬──────┬───────┬───────┬───────┬───────┬───────┐\n │ a ┆ b ┆ c ┆ out0 ┆ out1 ┆ out2 ┆ out3 ┆ out4 │\n │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │\n │ bool ┆ bool ┆ bool ┆ bool ┆ bool ┆ bool ┆ bool ┆ bool │\n ╞═══════╪═══════╪══════╪═══════╪═══════╪═══════╪═══════╪═══════╡\n │ null ┆ null ┆ null ┆ null ┆ null ┆ null ┆ null ┆ null │\n │ false ┆ false ┆ true ┆ false ┆ false ┆ false ┆ false ┆ false │\n │ false ┆ true ┆ true ┆ false ┆ false ┆ false ┆ false ┆ false │\n │ true ┆ true ┆ true ┆ false ┆ false ┆ false ┆ false ┆ false │\n │ true ┆ true ┆ true ┆ true ┆ true ┆ false ┆ false ┆ true │\n │ true ┆ true ┆ true ┆ true ┆ true ┆ true ┆ false ┆ true │\n └───────┴───────┴──────┴───────┴───────┴───────┴───────┴───────┘\n ```\n\n \"\"\"\n exprs = [a for a, b, c in more_itertools.chunked(args, 3)]\n ranges = [range(b, c + 1) for a, b, c in more_itertools.chunked(args, 3)]\n # 参数整理成col,repeat模式\n outputs = [itertools.chain.from_iterable(zip(exprs, d)) for d in itertools.product(*ranges)]\n return any_horizontal(ts_shifts_v2(*_) for _ in outputs)\n\n\ndef ts_skewness(x: Expr, d: int = 5, bias: bool = False, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动偏度\n\n Return skewness of x for the past d days\n\n Parameters\n ----------\n x\n d\n bias\n 有偏\n min_samples\n\n Notes\n -----\n `bias=False`时与`pandas`结果一样\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_skew(d, min_samples=minp, bias=bias)\n\n\ndef ts_std_dev(x: Expr, d: int = 5, ddof: int = 0, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动标准差\n\n Parameters\n ----------\n x\n d\n ddof\n 自由度\n min_samples\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_std(d, ddof=ddof, min_samples=minp)\n\n\ndef ts_sum(x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动求和\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_sum(d, min_samples=minp)\n\n\ndef ts_sum_split_by(x: Expr, by: Expr, d: int = 30, k: int = 10) -> Expr:\n \"\"\"切割论求和。在d窗口范围内以by为依据进行从小到大排序。取最大的N个和最小的N个对应位置的x的和\n\n Parameters\n ----------\n x\n by\n d\n 窗口大小\n k\n 最大最小的k个\n\n Returns\n -------\n Expr\n * top_k\n * bottom_k\n\n Examples\n --------\n\n ```python\n df = pl.DataFrame({\n 'a': [10, 20, 30, 40, 50, 60, 70, 80, 90, 100],\n 'b': [10, 9, 8, 7, 6, 5, 4, 3, 2, 1],\n }).with_columns(\n A=ts_sum_split_by(pl.col('a'), pl.col('b'), 8, 3)\n )\n\n shape: (10, 3)\n ┌─────┬─────┬───────────────┐\n │ a ┆ b ┆ A │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ struct[2] │\n ╞═════╪═════╪═══════════════╡\n │ 10 ┆ 10 ┆ {null,null} │\n │ 20 ┆ 9 ┆ {null,null} │\n │ 30 ┆ 8 ┆ {null,null} │\n │ 40 ┆ 7 ┆ {null,null} │\n │ 50 ┆ 6 ┆ {null,null} │\n │ 60 ┆ 5 ┆ {null,null} │\n │ 70 ┆ 4 ┆ {null,null} │\n │ 80 ┆ 3 ┆ {210.0,60.0} │\n │ 90 ┆ 2 ┆ {240.0,90.0} │\n │ 100 ┆ 1 ┆ {270.0,120.0} │\n └─────┴─────┴───────────────┘\n ```\n\n \"\"\"\n dtype = Struct([Field(f\"column_{i}\", Float64) for i in range(2)])\n return struct(f0=x, f1=by).map_batches(lambda xx: batches_i2_o2(struct_to_numpy(xx, 2), _sum_split_by, d, k), return_dtype=dtype)\n\n\ndef ts_triple_corr(x: Expr, y: Expr, z: Expr, d: int, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动三重相关系数\n\n Returns triple correlation of x, y, z for the past d days.\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return struct(f0=x, f1=y, f2=z).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 3), roll_triple_corr, d, minp), return_dtype=Float64)\n\n\ndef ts_weighted_decay(x: Expr, k: float = 0.5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动加权衰减求和\n\n Instead of replacing today’s value with yesterday’s as in ts_delay(x, 1),\n it assigns weighted average of today’s and yesterday’s values with weight on today’s value being k and yesterday’s being (1-k).\n\n Parameters\n ----------\n x\n k\n 衰减系数\n min_samples\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return x.rolling_sum(2, weights=[1 - k, k], min_samples=minp)\n\n\ndef ts_zscore(x: Expr, d: int = 5, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动zscore\"\"\"\n return (x - ts_mean(x, d, min_samples)) / ts_std_dev(x, d, 0, min_samples)\n\n\ndef ts_cum_prod_by(r: Expr, v: Expr) -> Expr:\n \"\"\"带设置的累乘\n\n 可用于市值累乘日收益率得到新市值的需求\n\n Parameters\n ----------\n r\n 收益率。不能出现`null`, `null`需要提前用`1`代替\n\n `CLOSE/ts_delay(CLOSE, 1)`\n v\n 市值。非空时分配指定市值资产\n\n * 如果非`null`，直接返回`v`\n * 如果`null`，返回`V[-1]*r`\n\n Returns\n -------\n Expr\n V。累乘后成新的市值\n\n Examples\n --------\n\n ```python\n df = pl.DataFrame({\n 'r': [1, 2, 3, 4, 5, 6],\n 'v': [None, None, 6, None, None, 12],\n\n }).with_columns(\n V=ts_cum_prod_by(pl.col('r'), pl.col('v'))\n )\n\n shape: (6, 3)\n ┌─────┬──────┬───────┐\n │ r ┆ v ┆ V │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ f64 │\n ╞═════╪══════╪═══════╡\n │ 1 ┆ null ┆ null │\n │ 2 ┆ null ┆ null │\n │ 3 ┆ 6 ┆ 6.0 │\n │ 4 ┆ null ┆ 24.0 │\n │ 5 ┆ null ┆ 120.0 │\n │ 6 ┆ 12 ┆ 12.0 │\n └─────┴──────┴───────┘\n ```\n\n\n \"\"\"\n return struct(f0=r, f1=v).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), _cum_prod_by), return_dtype=Float64)\n\n\ndef ts_cum_sum_by(r: Expr, v: Expr) -> Expr:\n \"\"\"带设置的累加\n\n 可用于市值累加日收益得到新市值的需求\n\n Parameters\n ----------\n r\n 收益。不能出现`null`, `null`需要提前用`0`代替\n\n `CLOSE-ts_delay(CLOSE, 1)`\n v\n 市值。非空时分配指定市值资产\n\n * 如果非`null`，直接返回`v`\n * 如果`null`，返回`V[-1]+r`\n\n Examples\n --------\n\n ```python\n df = pl.DataFrame({\n 'r': [1, 2, 3, 4, 5, 6],\n 'v': [None, None, 6, None, None, 12],\n\n }).with_columns(\n V=ts_cum_sum_by(pl.col('r'), pl.col('v'))\n )\n\n shape: (6, 3)\n ┌─────┬──────┬──────┐\n │ r ┆ v ┆ V │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ f64 │\n ╞═════╪══════╪══════╡\n │ 1 ┆ null ┆ null │\n │ 2 ┆ null ┆ null │\n │ 3 ┆ 6 ┆ 6.0 │\n │ 4 ┆ null ┆ 10.0 │\n │ 5 ┆ null ┆ 15.0 │\n │ 6 ┆ 12 ┆ 12.0 │\n └─────┴──────┴──────┘\n\n ```\n\n \"\"\"\n return struct(f0=r, f1=v).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, 2), _cum_sum_by), return_dtype=Float64)\n\n\ndef ts_regression_resid(y: Expr, x: Expr, d: int, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动回归取残差\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return pls.compute_rolling_least_squares(y, x, mode='residuals', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp))\n\n\ndef ts_regression_pred(y: Expr, x: Expr, d: int, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动回归取y的预测值\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return pls.compute_rolling_least_squares(y, x, mode='predictions', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp))\n\n\ndef ts_regression_intercept(y: Expr, x: Expr, d: int, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动回归取截距\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return pls.compute_rolling_least_squares(y, x, mode='coefficients', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp)).struct[1]\n\n\ndef ts_regression_slope(y: Expr, x: Expr, d: int, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动回归取斜率\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return pls.compute_rolling_least_squares(y, x, mode='coefficients', add_intercept=True, rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp)).struct[0]\n\n\ndef ts_resid(y: Expr, *more_x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:\n \"\"\"多元时序滚动回归取残差\n\n Parameters\n ----------\n y\n *more_x\n 多个x\n d\n min_samples\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return pls.compute_rolling_least_squares(y, *more_x, mode='residuals', rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp))\n\n\ndef ts_pred(y: Expr, *more_x: Expr, d: int = 30, min_samples: Optional[int] = None) -> Expr:\n \"\"\"多元时序滚动回归预测\n\n Parameters\n ----------\n y\n *more_x\n 多个x\n d\n min_samples\n\n \"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES or d\n return pls.compute_rolling_least_squares(y, *more_x, mode='predictions', rolling_kwargs=RollingKwargs(window_size=d, min_periods=minp))\n\n\ndef ts_weighted_mean(x: Expr, w: Expr, d: int, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动加权平均\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return (x * w).rolling_sum(d, min_samples=minp) / w.rolling_sum(d, min_samples=minp)\n\n\ndef ts_weighted_sum(x: Expr, w: Expr, d: int, min_samples: Optional[int] = None) -> Expr:\n \"\"\"时序滚动加权求和\"\"\"\n minp = min_samples or polars_ta.MIN_SAMPLES\n return (x * w).rolling_sum(d, min_samples=minp)\n\n\ndef ts_signals_to_size(long_entry: Expr, long_exit: Expr,\n short_entry: Expr, short_exit: Expr,\n accumulate: bool = False,\n action: bool = False) -> Expr:\n \"\"\"多空信号转持仓。参考于`vectorbt`\n\n Parameters\n ----------\n long_entry\n 多头入场\n long_exit\n 多头出场\n short_entry\n 空头入场\n short_exit\n 空头出场\n accumulate\n 遇到重复信号时是否累计\n action\n 返回持仓状态还是下单操作\n\n \"\"\"\n return struct(f0=long_entry, f1=long_exit, f2=short_entry, f3=short_exit).map_batches(\n lambda xx: batches_i2_o1(struct_to_numpy(xx, 4, dtype=float), _signals_to_size, accumulate, action), return_dtype=Float64)\n" }, { "path": "polars_ta/wq/transformational.py", "content": "from polars import Expr, when, Boolean, Int32, Float32, lit\n\n\ndef cut(x: Expr, b: float, *more_bins) -> Expr:\n \"\"\"分箱\n\n Parameters\n ----------\n x\n b\n *more_bins\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 1, 1, 2, 2, 3, 10],\n }).with_columns(\n out1=cut(pl.col('a'), 2, 5, 20),\n )\n shape: (8, 2)\n ┌──────┬──────┐\n │ a ┆ out1 │\n │ --- ┆ --- │\n │ i64 ┆ u32 │\n ╞══════╪══════╡\n │ null ┆ null │\n │ 1 ┆ 0 │\n │ 1 ┆ 0 │\n │ 1 ┆ 0 │\n │ 2 ┆ 0 │\n │ 2 ┆ 0 │\n │ 3 ┆ 1 │\n │ 10 ┆ 2 │\n └──────┴──────┘\n\n ```\n\n \"\"\"\n return x.cut([b, *more_bins]).to_physical()\n\n\ndef clamp(x: Expr, lower: float = 0, upper: float = 0, inverse: bool = False, mask: float = None) -> Expr:\n \"\"\"Limits input value between lower and upper bound in inverse = false mode (which is default). Alternatively, when inverse = true, values between bounds are replaced with mask, while values outside bounds are left as is.\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 2, 3, 4, 5, 6],\n }).with_columns(\n out1=clamp(pl.col('a'), 2, 5, False),\n out2=clamp(pl.col('a'), 2, 5, True),\n )\n shape: (7, 3)\n ┌──────┬──────┬──────┐\n │ a ┆ out1 ┆ out2 │\n │ --- ┆ --- ┆ --- │\n │ i64 ┆ i64 ┆ i64 │\n ╞══════╪══════╪══════╡\n │ null ┆ null ┆ null │\n │ 1 ┆ 2 ┆ 1 │\n │ 2 ┆ 2 ┆ null │\n │ 3 ┆ 3 ┆ null │\n │ 4 ┆ 4 ┆ null │\n │ 5 ┆ 5 ┆ null │\n │ 6 ┆ 5 ┆ 6 │\n └──────┴──────┴──────┘\n\n ```\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#clampx-lower-0-upper-0-inverse-false-mask\n\n \"\"\"\n if inverse:\n cond = (x >= lower) & (x <= upper)\n return when(~cond).then(x).otherwise(mask)\n else:\n return x.clip(lower, upper)\n\n\n# def filter_(x: Expr, h: str = \"1, 2, 3, 4\", t: str = \"0.5\") -> Expr:\n# \"\"\"Used to filter the value and allows to create filters like linear or exponential decay.\"\"\"\n# raise\n\n\ndef _keep(x: Expr, f: float, period: int = 5) -> Expr:\n \"\"\"This operator outputs value x when f changes and continues to do that for “period” days after f stopped changing. After “period” days since last change of f, NaN is output.\"\"\"\n raise\n\n\ndef left_tail(x: Expr, maximum: float = 0) -> Expr:\n \"\"\"NaN everything greater than maximum, maximum should be constant.\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 2, 3, 4, 5],\n }).with_columns(\n out=left_tail(pl.col('a'), 3),\n )\n shape: (6, 2)\n ┌──────┬──────┐\n │ a ┆ out │\n │ --- ┆ --- │\n │ i64 ┆ i64 │\n ╞══════╪══════╡\n │ null ┆ null │\n │ 1 ┆ 1 │\n │ 2 ┆ 2 │\n │ 3 ┆ 3 │\n │ 4 ┆ null │\n │ 5 ┆ null │\n └──────┴──────┘\n ```\n\n See Also\n --------\n tail\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#left_tail\n\n \"\"\"\n return when(x <= maximum).then(x).otherwise(None)\n\n\n# def pasteurize(x: Expr) -> Expr:\n# \"\"\"Set to NaN if x is INF or if the underlying instrument is not in the Alpha universe\"\"\"\n# # TODO: 不在票池中的的功能无法表示\n# # TODO: 与purify好像没啥区别\n# return when(x.is_finite()).then(x).otherwise(None)\n\n\ndef purify(x: Expr) -> Expr:\n \"\"\"Clear infinities (+inf, -inf) by replacing with null.\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1., 2., float('nan'), float('inf'), float('-inf')],\n }).with_columns(\n out=purify(pl.col('a')),\n )\n shape: (6, 2)\n ┌──────┬──────┐\n │ a ┆ out │\n │ --- ┆ --- │\n │ f64 ┆ f64 │\n ╞══════╪══════╡\n │ null ┆ null │\n │ 1.0 ┆ 1.0 │\n │ 2.0 ┆ 2.0 │\n │ NaN ┆ null │\n │ inf ┆ null │\n │ -inf ┆ null │\n └──────┴──────┘\n ```\n\n \"\"\"\n return when(x.is_finite()).then(x).otherwise(None)\n\n\ndef fill_nan(x: Expr) -> Expr:\n \"\"\"填充`nan`为`null`\"\"\"\n return x.fill_nan(None)\n\n\ndef fill_null(x: Expr, value=0) -> Expr:\n \"\"\"填充`null`为`value`\"\"\"\n return x.fill_null(value)\n\n\ndef right_tail(x: Expr, minimum: float = 0) -> Expr:\n \"\"\"NaN everything less than minimum, minimum should be constant.\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 2, 3, 4, 5],\n }).with_columns(\n out=right_tail(pl.col('a'), 3),\n )\n shape: (6, 2)\n ┌──────┬──────┐\n │ a ┆ out │\n │ --- ┆ --- │\n │ i64 ┆ i64 │\n ╞══════╪══════╡\n │ null ┆ null │\n │ 1 ┆ null │\n │ 2 ┆ null │\n │ 3 ┆ 3 │\n │ 4 ┆ 4 │\n │ 5 ┆ 5 │\n └──────┴──────┘\n ```\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#right_tail\n\n \"\"\"\n return when(x >= minimum).then(x).otherwise(None)\n\n\ndef sigmoid(x: Expr) -> Expr:\n \"\"\"sigmoid激活函数\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return 1 / (1 + (-x).exp())\n\n\ndef logit(x: Expr) -> Expr:\n \"\"\"对数几率，拉伸中间部分\"\"\"\n if not isinstance(x, Expr):\n x = lit(x)\n return (x / (1 + x)).log()\n\n\ndef tail(x: Expr, lower: float = 0, upper: float = 0, newval: float = 0) -> Expr:\n \"\"\"If (x > lower AND x < upper) return newval, else return x. Lower, upper, newval should be constants.\n\n Examples\n --------\n ```python\n df = pl.DataFrame({\n 'a': [None, 1, 2, 3, 4, 5, 6],\n }).with_columns(\n out=tail(pl.col('a'), 2, 5, 1),\n )\n shape: (7, 2)\n ┌──────┬──────┐\n │ a ┆ out │\n │ --- ┆ --- │\n │ i64 ┆ i64 │\n ╞══════╪══════╡\n │ null ┆ null │\n │ 1 ┆ 1 │\n │ 2 ┆ 2 │\n │ 3 ┆ 1 │\n │ 4 ┆ 1 │\n │ 5 ┆ 5 │\n │ 6 ┆ 6 │\n └──────┴──────┘\n\n ```\n\n See Also\n --------\n clamp\n\n References\n ----------\n https://platform.worldquantbrain.com/learn/operators/detailed-operator-descriptions#tail\n\n \"\"\"\n\n cond = (x > lower) & (x < upper)\n return when(~cond | x.is_null()).then(x).otherwise(newval)\n\n\ndef int_(a: Expr) -> Expr:\n \"\"\"bool转int\"\"\"\n if not isinstance(a, Expr):\n a = lit(a)\n return a.cast(Int32)\n\n\ndef bool_(a: Expr) -> Expr:\n \"\"\"int转成bool\"\"\"\n if not isinstance(a, Expr):\n a = lit(a)\n return a.cast(Boolean)\n\n\ndef float_(a: Expr) -> Expr:\n \"\"\"int转成float\"\"\"\n if not isinstance(a, Expr):\n a = lit(a)\n return a.cast(Float32)\n\n\ndef nop(x: Expr) -> Expr:\n \"\"\"空操作\"\"\"\n return x\n\n\ndef lit_(a: Expr) -> Expr:\n \"\"\"将常量封装成lit\"\"\"\n if not isinstance(a, Expr):\n return lit(a)\n return a\n" }, { "path": "polars_ta/wq/vector.py", "content": "\"\"\"\n本文件中算子都是对整列计算得到一个值，使用时可能会广播到整列\n\n可以用在横截面，也可以用在时序。但时序使用会引入未来数据，这时一般用途是统计\n\"\"\"\n\nfrom polars import Expr\n\n\ndef vec_avg(x: Expr) -> Expr:\n \"\"\"Taking mean of the vector field x\"\"\"\n return x.mean()\n\n\ndef vec_choose(x: Expr, nth: int) -> Expr:\n \"\"\"Choosing kth item(indexed at 0) from each vector field x\"\"\"\n return x.gather(nth)\n\n\ndef vec_count(x: Expr) -> Expr:\n \"\"\"Number of elements in vector field x\"\"\"\n return x.count()\n\n\ndef vec_ir(x: Expr) -> Expr:\n \"\"\"Information Ratio (Mean / Standard Deviation) of vector field x\"\"\"\n return x.mean() / x.std(ddof=0)\n\n\ndef vec_kurtosis(x: Expr) -> Expr:\n \"\"\"Kurtosis of vector field x\"\"\"\n return x.kurtosis()\n\n\ndef vec_l2_norm(x: Expr) -> Expr:\n \"\"\"Euclidean norm\"\"\"\n return x.pow(2).sum().sqrt()\n\n\ndef vec_max(x: Expr) -> Expr:\n \"\"\"Maximum value form vector field x\"\"\"\n return x.max()\n\n\ndef vec_min(x: Expr) -> Expr:\n \"\"\"Minimum value form vector field x\"\"\"\n return x.min()\n\n\ndef vec_norm(x: Expr) -> Expr:\n \"\"\"Sum of all absolute values of vector field x\"\"\"\n return x.abs().sum()\n\n\ndef vec_percentage(x: Expr, percentage: float = 0.5) -> Expr:\n \"\"\"Percentile of vector field x\"\"\"\n return x.quantile(percentage)\n\n\ndef vec_powersum(x: Expr, constant: float = 2) -> Expr:\n \"\"\"Sum of power of vector field x\"\"\"\n return (x ** constant).sum()\n\n\ndef vec_range(x: Expr) -> Expr:\n \"\"\"Difference between maximum and minimum element in vector field x\"\"\"\n return x.max() - x.min()\n\n\ndef vec_skewness(x: Expr) -> Expr:\n \"\"\"Skewness of vector field x\"\"\"\n return x.skew()\n\n\ndef vec_stddev(x: Expr) -> Expr:\n \"\"\"Standard Deviation of vector field x\"\"\"\n return x.std(ddof=0)\n\n\ndef vec_sum(x: Expr) -> Expr:\n \"\"\"Sum of vector field x\"\"\"\n return x.sum()\n\ndef vec_median(x: Expr) -> Expr:\n return x.median()" }, { "path": "prompt.txt", "content": "你是一个专业的量化因子开发助手，帮助创建、改进因子。\n\n你可以:\n- 编写和优化因子代码\n- 提供因子开发示例\n- 建议和改进因子逻辑\n\n【重要限制】\n1. 只回答与因子开发、编码和优化相关的问题\n2. 无论输入何种语言，都用中文回复\n3. 只使用`内置算子`。优先使用`基础因子`\n4. 不能使用命名参数\n5. 很多函数直接支持将bool当成int用来计算，可以省去if_else和int_\n\n## 基础因子\n- 价格因子: OPEN, HIGH, LOW, CLOSE, VWAP\n- 成交量因子: VOLUME, AMOUNT, TURNOVER\n- 市值因子: MARKET_CAP, PB, PE\n\n## 内置算子\n\n### polars_ta.wq.arithmetic\n- abs_(x) : 求绝对值\n- cbrt(x) : 立方根\n- ceiling(x) : 向上取整\n- cos(x) : 余弦\n- cube(x) : 立方\n- exp(x) : 自然指数函数\n- floor(x) : 向下取整\n- fraction(x) : 小数部分\n- log(x) : 以e为底的对数\n- log10(x) : 以10为底的对数\n- log1p(x) : 简单收益率 转 对数收益率\n- max_(a,b,*args) : 水平多列求最大值\n- min_(a,b,*args) : 水平多列求最小值\n- mod(x,y) : 求余\n- power(x,y) : 乘幂\n- round_(x,decimals:int=0) : 四舍五入\n- round_down(x,f:int=1) : 小于输入的f的最大倍数\n- sign(x) : 符号函数\n- signed_power(x,y) : x的y次幂，符号保留\n- sin(x) : 正弦\n- sqrt(x) : 平方根\n- square(x) : 平方\n- std(a,b,*args) : 水平多列求标准差\n- tan(x) : 正切\n- var(a,b,*args) : 水平多列求方差\n### polars_ta.wq.time_series\n- ts_arg_max(x,d:int=5,reverse:bool=True) : 最大值相对位置\n- ts_arg_min(x,d:int=5,reverse:bool=True) : 最小值相对位置\n- ts_corr(x,y,d:int=5,ddof:int=1) : 时序滚动相关系数\n- ts_count(x,d:int=30) : 时序滚动计数\n- ts_covariance(x,y,d:int=5,ddof:int=1) : 时序滚动协方差\n- ts_decay_exp_window(x,d:int=30,factor:float=1.0) : 指数衰减移动平均\n- ts_decay_linear(x,d:int=30) : 线性衰减移动平均\n- ts_delay(x,d:int=1,fill_value:float=None) : 时序数据移动\n- ts_delta(x,d:int=1) : 时序差分\n- ts_fill_null(x,limit:int=None) : 用上一个非空值填充空值\n- ts_ir(x,d:int=1) : 时序滚动信息系数\n- ts_kurtosis(x,d:int=5,bias:bool=False) : 时序滚动峰度\n- ts_l2_norm(x,d:int=5) : 欧几里得范数\n- ts_log_diff(x,d:int=1) : 求对数，然后时序滚动差分\n- ts_max(x,d:int=30) : 时序滚动最大值\n- ts_mean(x,d:int=5) : 简单移动平均\n- ts_median(x,d:int=5) : 时序滚动中位数\n- ts_min(x,d:int=30) : 时序滚动最小值\n- ts_moment(x,d:int,k:int=0) : 滚动k阶中心距\n- ts_percentage(x,d:int,percentage:float=0.5) : 滚动百分位数\n- ts_pred(y,*more_x,d:int=30) : 多元时序滚动回归预测\n- ts_product(x,d:int=5) : 时序滚动乘\n- ts_rank(x,d:int=5) : 时序滚动排名\n- ts_realized_volatility(close,d:int=5) : 已实现波动率\n- ts_regression_intercept(y,x,d:int) : 时序滚动回归取截距\n- ts_regression_pred(y,x,d:int) : 时序滚动回归取y的预测值\n- ts_regression_resid(y,x,d:int) : 时序滚动回归取残差\n- ts_regression_slope(y,x,d:int) : 时序滚动回归取斜率\n- ts_resid(y,*more_x,d:int=30) : 多元时序滚动回归取残差\n- ts_returns(x,d:int=1) : 简单收益率\n- ts_scale(x,d:int=5) : 时序滚动缩放。相当于ts_minmax\n- ts_skewness(x,d:int=5,bias:bool=False) : 时序滚动偏度\n- ts_std_dev(x,d:int=5,ddof:int=0) : 时序滚动标准差\n- ts_sum(x,d:int=30) : 时序滚动求和\n- ts_weighted_decay(x,k:float=0.5) : 时序滚动加权衰减求和\n- ts_weighted_mean(x,w,d:int) : 时序滚动加权平均\n- ts_weighted_sum(x,w,d:int) : 时序滚动加权求和\n- ts_zscore(x,d:int=5) : 时序滚动zscore\n### polars_ta.wq.cross_sectional\n- cs_fill_null(x,value:float=0) : 横截面上，填充`null`为`value`\n- cs_qcut(x,q:int=10) : 横截面上等频分箱\n- cs_rank(x,pct:bool=True) : 横截面排名\n- cs_rank_if(condition,x,pct:bool=True) : 横截面筛选排名。可实现动态票池\n- cs_scale(x,scale_:float=1,long_scale:float=1,short_scale:float=1) : 横截面上，将输入数据进行比例调整\n### polars_ta.wq.preprocess\n- cs_3sigma(x,n:float=3.0) : 横截面3倍sigma去极值\n- cs_demean(x) : 横截面去均值化\n- cs_mad(x,n:float=3.0,k:float=1.4826) : 横截面MAD去极值\n- cs_minmax(x) : 横截面minmax标准化\n- cs_quantile(x,low_limit:float=0.025,up_limit:float=0.975) : 横截面分位数去极值\n- cs_quantile_zscore(y,low_limit:float=0.025,up_limit:float=0.975) : 横截面分位数去极值、标准化\n- cs_resid(y,*more_x) : 横截面多元回归取残差\n- cs_resid_zscore(y,*more_x) : 横截面中性化、标准化\n- cs_zscore(x,ddof:int=0) : 横截面zscore标准化\n### polars_ta.wq.logical\n- if_else(condition,true_value,false_value=None) : 条件判断\n### polars_ta.wq.transformational\n- bool_(a) : int转成bool\n- cut(x,b:float,*more_bins) : 分箱\n- float_(a) : int转成float\n- int_(a) : bool转int\n- sigmoid(x) : sigmoid激活函数\n\n## 因子示例\n\n```python\nalpha_004=-ts_rank(cs_rank(LOW),9)\nalpha_006=-ts_corr(OPEN,VOLUME,10)\nalpha_040=-cs_rank(ts_std_dev(HIGH,10))*ts_corr(HIGH,VOLUME,10)\nalpha_041=power(HIGH*LOW,0.5)-VWAP\nalpha_042=cs_rank(VWAP-CLOSE)/cs_rank(VWAP+CLOSE)\nalpha_053=-ts_delta((CLOSE-LOW-(HIGH-CLOSE))/(CLOSE-LOW),9)\nalpha_057=-((CLOSE-VWAP)/ts_decay_linear(cs_rank(ts_arg_max(CLOSE,30)),2))\nalpha_101=(CLOSE-OPEN)/(HIGH-LOW+0.001)\n```\n\n如果理解了上述要求，不用思考，直接回复：“老大，请吩咐”" }, { "path": "pyproject.toml", "content": "[project]\nname = \"polars_ta\"\nauthors = [\n { name = \"wukan\", email = \"wu-kan@163.com\" },\n]\ndescription = \"polars expressions\"\nreadme = \"README.md\"\nrequires-python = \">=3.8\"\nkeywords = [\"polars\", \"expression\", \"talib\"]\nlicense = { file = \"LICENSE\" }\nclassifiers = [\n \"Development Status :: 4 - Beta\",\n \"Programming Language :: Python\"\n]\ndependencies = [\n \"polars>=1.28.0\",\n \"polars-ols>=0.3.0\",\n \"numpy\",\n \"numba\",\n \"pandas\",\n \"more_itertools\",\n]\ndynamic = [\"version\"]\n\n[project.optional-dependencies]\ntalib = [\n \"TA-Lib\",\n]\n\n[build-system]\nrequires = [\"hatchling\"]\nbuild-backend = \"hatchling.build\"\n\n[tool.hatch.version]\npath = \"polars_ta/_version.py\"\n\n[tool.hatch.build.targets.wheel]\npackages = [\"polars_ta\"]\ninclude-package-data = true\n\n[tool.hatch.build.targets.sdist]\ninclude = [\"polars_ta*\"]\n" }, { "path": "requirements-docs.txt", "content": "mkdocs\nmkdocstrings[python]\nmkdocs-material\nmkdocs-llmstxt" }, { "path": "requirements.txt", "content": "polars>=1.26.0\npolars-ols>=0.3.0\nnumpy\nnumba\npandas\nmore_itertools" }, { "path": "setup.py", "content": "from setuptools import setup\n\nsetup()\n" }, { "path": "tests/numba_test.py", "content": "import time\n\nimport numpy as np\nimport polars as pl\nfrom numba import jit\n\nfrom polars_ta.utils.numba_ import nb_roll_sum, batches_i1_o1, roll_sum, roll_cov, roll_split_i2_o2\nfrom polars_ta.wq.time_series import ts_co_kurtosis\n\n\n@jit(nopython=True, nogil=True, fastmath=True, cache=True)\ndef nb_sum(x):\n return np.sum(x)\n\n\ndf = pl.DataFrame({'A': range(100000), 'B': range(100000)})\na = df.with_columns([\n pl.col('A').rolling_sum(10).alias('a1'),\n pl.col('A').rolling_map(lambda x: x.sum(), 10).alias('a2'),\n pl.col('A').rolling_map(lambda x: nb_sum(x.to_numpy()), 10).alias('a3'),\n roll_sum(pl.col('A'), 10).alias('a4'),\n pl.col('A').map_batches(lambda x: batches_i1_o1(x.to_numpy(), nb_roll_sum, 10)).alias('a5'),\n pl.rolling_cov(pl.col('A'), pl.col('B'), window_size=10).alias('a6'),\n roll_cov(pl.col('A'), pl.col('B'), 10).alias('a7'),\n ts_co_kurtosis(pl.col('A'), pl.col('B'), 10).alias('a8'),\n roll_split_i2_o2(pl.col('A'), pl.col('B'), 10, 2).alias('a10'),\n])\nprint(a)\n\nt1 = time.perf_counter()\nfor i in range(10):\n a = df.with_columns([\n pl.col('A').rolling_sum(10).alias('a1'),\n ])\nt2 = time.perf_counter()\nprint(t2 - t1)\n\nt1 = time.perf_counter()\nfor i in range(10):\n a = df.with_columns([\n pl.col('A').rolling_map(lambda x: x.sum(), 10).alias('a2'),\n ])\nt2 = time.perf_counter()\nprint(t2 - t1)\n\nt1 = time.perf_counter()\nfor i in range(10):\n a = df.with_columns([\n pl.col('A').rolling_map(lambda x: nb_sum(x.to_numpy()), 10).alias('a3'),\n ])\nt2 = time.perf_counter()\nprint(t2 - t1)\n\nt1 = time.perf_counter()\nfor i in range(10):\n a = df.with_columns([\n pl.col('A').map_batches(lambda x: pl.Series(nb_roll_sum(x.to_numpy(), 10), nan_to_null=True)).alias('a4'),\n ])\nt2 = time.perf_counter()\nprint(t2 - t1)\n\nt1 = time.perf_counter()\nfor i in range(10):\n a = df.with_columns([\n pl.col('A').map_batches(lambda x: batches_i1_o1(x.to_numpy(), nb_roll_sum, 10)).alias('a5'),\n ])\nt2 = time.perf_counter()\nprint(t2 - t1)\n" }, { "path": "tests/pit_test.py", "content": "import time\nfrom pathlib import Path\n\nimport pandas as pd\nimport polars as pl\n\nfrom polars_ta.utils.pit import pit_prepare, ttm_from_point, LOOKBACK_DATE, pit_frist, ttm_from_period, sheet_from_joinquant\n\nPATH_STEP0_INPUT1 = r'F:\\data\\jqresearch\\get_STK_BALANCE_SHEET'\nPATH_STEP0_INPUT2 = r'F:\\data\\jqresearch\\get_STK_INCOME_STATEMENT'\nPATH_STEP0_INPUT3 = r'F:\\data\\jqresearch\\get_STK_CASHFLOW_STATEMENT'\n\n\ndef load_parquet(folder):\n paths = list(Path(folder).glob('*.parquet'))\n # due to the data is not complete when writing, the data type can only be read from pandas then convert to polars\n # 写入时由于部分数据为空，导致入写类型不同，读取就存在问题，只能用pandas读回来\n dfs = pd.concat([pd.read_parquet(p) for p in paths])\n return pl.from_pandas(dfs, nan_to_null=True)\n\n\ndef balance(df1):\n # https://www.joinquant.com/help/api/help#Stock:合并资产负债表\n # 002509.XSHE 600528.XSHG 000528.XSHE # 天广中贸 疫情期 2019年报晚于2020一季报公布\n df1 = sheet_from_joinquant(df1)\n df1 = pit_prepare(df1, by1='code', by2='report_date', by3='pub_date', by4=LOOKBACK_DATE)\n df2 = df1.select(\n \"code\", \"report_date\", \"pub_date\", LOOKBACK_DATE, \"report_type\",\n \"total_assets\", ttm_from_point('total_assets').over('code', LOOKBACK_DATE, order_by='report_date').name.suffix('_ttm'),\n \"equities_parent_company_owners\", ttm_from_point('equities_parent_company_owners').over('code', LOOKBACK_DATE, order_by='report_date').name.suffix('_ttm'),\n )\n df3 = pit_frist(df2, by1='code', by2='report_date', by3='pub_date', by4=LOOKBACK_DATE)\n return df3\n\n\ndef income(df1):\n # https://www.joinquant.com/help/api/help#Stock:合并利润表\n df1 = sheet_from_joinquant(df1)\n df1 = df1.with_columns(quarter=pl.col('report_date').dt.quarter())\n df1 = pit_prepare(df1, by1='code', by2='report_date', by3='pub_date', by4=LOOKBACK_DATE)\n df2 = df1.select(\n \"code\", \"report_date\", \"pub_date\", LOOKBACK_DATE, \"report_type\",\n \"total_operating_revenue\", ttm_from_period(\"total_operating_revenue\", quarter='quarter').over('code', LOOKBACK_DATE, order_by='report_date').name.suffix('_ttm'),\n \"net_profit\", ttm_from_period('net_profit', quarter='quarter').over('code', LOOKBACK_DATE, order_by='report_date').name.suffix('_ttm'),\n )\n df3 = pit_frist(df2, by1='code', by2='report_date', by3='pub_date', by4=LOOKBACK_DATE)\n return df3\n\n\ndef cashflow(df1):\n # https://www.joinquant.com/help/api/help#Stock:合并现金流量表\n df1 = sheet_from_joinquant(df1)\n df1 = df1.with_columns(quarter=pl.col('report_date').dt.quarter())\n df1 = pit_prepare(df1, by1='code', by2='report_date', by3='pub_date', by4=LOOKBACK_DATE)\n df2 = df1.select(\n \"code\", \"report_date\", \"pub_date\", LOOKBACK_DATE, \"report_type\",\n \"net_operate_cash_flow\", ttm_from_period(\"net_operate_cash_flow\", quarter='quarter').over('code', LOOKBACK_DATE, order_by='report_date').name.suffix('_ttm'),\n \"financial_cost\", ttm_from_period('financial_cost', quarter='quarter').over('code', LOOKBACK_DATE, order_by='report_date').name.suffix('_ttm'),\n )\n df3 = pit_frist(df2, by1='code', by2='report_date', by3='pub_date', by4=LOOKBACK_DATE)\n return df3\n\n\nif __name__ == '__main__':\n df1 = load_parquet(PATH_STEP0_INPUT1)\n df2 = load_parquet(PATH_STEP0_INPUT2)\n df3 = load_parquet(PATH_STEP0_INPUT3)\n\n t1 = time.perf_counter()\n balance(df1)\n income(df2)\n cashflow(df3)\n t2 = time.perf_counter()\n\n print(t2 - t1)\n" }, { "path": "tests/ta/test_momentum.py", "content": "import numpy as np\nimport polars as pl\nimport talib\n\n\nclass TestDemoClass:\n high_np = None\n low_np = None\n close_np = None\n df_pl = None\n\n def setup_class(self):\n self.high_np = 0 + np.random.rand(100)\n self.low_np = 0 - np.random.rand(100)\n self.close_np = np.arange(100, dtype=float)+ np.random.rand(100)\n\n self.df_pl = pl.DataFrame([self.high_np, self.low_np, self.close_np],\n schema=[\"high\", \"low\", \"close\"])\n\n def test_WILLR(self):\n from polars_ta.ta.momentum import WILLR\n\n result1 = talib.WILLR(self.high_np, self.low_np, self.close_np, timeperiod=5)\n result2 = self.df_pl.select(WILLR(pl.col(\"high\"), pl.col(\"low\"), pl.col(\"close\"), timeperiod=5))\n # !!! 注意，与talib版的区别\n result3 = result2['high'].to_numpy() * (-100)\n\n print()\n print(result1)\n print(result3)\n\n # assert np.allclose(result1, result3, equal_nan=True)\n\n\n def test_STOCHF_fastd(self):\n from polars_ta.ta.momentum import STOCHF\n\n _, result1 = talib.STOCHF(self.high_np, self.low_np, self.close_np, fastk_period=5, fastd_period=3, fastd_matype=0)\n result2 = self.df_pl.select(STOCHF(pl.col(\"high\"), pl.col(\"low\"), pl.col(\"close\"), fastk_period=5, fastd_period=3).struct[1])\n # !!! 注意，与talib版的区别\n result3 = result2['fastd'].to_numpy() * 100\n print()\n print(result1)\n print(result3)\n\n # assert np.allclose(result1, result3, equal_nan=True)\n\n def test_MACD(self):\n talib.set_compatibility(1)\n\n from polars_ta.ta.momentum import MACD\n\n fastperiod = 3\n slowperiod = 5\n\n result1, _, _ = talib.MACD(self.close_np, fastperiod=fastperiod, slowperiod=slowperiod, signalperiod=1)\n result2 = self.df_pl.select(MACD(pl.col(\"close\"), fastperiod=fastperiod, slowperiod=slowperiod).struct[0])\n result3 = result2['macd'].to_numpy()\n # print()\n # print(result1)\n # print(result3)\n\n assert np.allclose(result1[slowperiod:], result3[slowperiod:], equal_nan=True)\n\n def test_TRIX(self):\n talib.set_compatibility(1)\n\n from polars_ta.ta.momentum import TRIX\n\n result1 = talib.TRIX(self.close_np, timeperiod=30)\n result2 = self.df_pl.select(TRIX(pl.col(\"close\"), timeperiod=30))\n result3 = result2['close'].to_numpy() * 100\n # print()\n # print(result1)\n # print(result3)\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_AROON(self):\n timeperiod = 10\n\n from polars_ta.ta.momentum import AROON\n\n _, result1 = talib.AROON(self.high_np, self.low_np, timeperiod=timeperiod)\n result2 = self.df_pl.select(AROON(pl.col(\"high\"), pl.col(\"low\"), timeperiod=timeperiod).alias('high').struct[1])\n result3 = result2['aroonup'].to_numpy() * 100\n print()\n print(result1)\n print(result3)\n\n # assert np.allclose(result1[timeperiod:], result3[timeperiod:], equal_nan=True)\n\n def test_RSI(self):\n timeperiod = 10\n\n from polars_ta.ta.momentum import RSI\n\n result1 = talib.RSI(self.close_np, timeperiod=timeperiod)\n result2 = self.df_pl.select(RSI(pl.col(\"close\"), timeperiod=timeperiod))\n result3 = result2['close'].to_numpy() * 100\n # print()\n # print(result1)\n # print(result3)\n assert np.allclose(result1[timeperiod:], result3[timeperiod:], equal_nan=True)" }, { "path": "tests/ta/test_operators.py", "content": "import numpy as np\nimport polars as pl\nimport talib\n\n\nclass TestDemoClass:\n high_np = None\n low_np = None\n close_np = None\n df_pl = None\n\n def setup_class(self):\n self.high_np = np.arange(100, dtype=float) + np.random.rand(100)\n self.low_np = np.arange(100, dtype=float) - np.random.rand(100)\n self.close_np = np.random.rand(100)\n\n self.df_pl = pl.DataFrame([self.high_np, self.low_np, self.close_np],\n schema=[\"high\", \"low\", \"close\"])\n\n def test_MAX(self):\n from polars_ta.ta.operators import MAX\n\n result1 = talib.MAX(self.close_np, timeperiod=10)\n result2 = self.df_pl.select(MAX(pl.col(\"close\"), timeperiod=10))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_MIN(self):\n from polars_ta.ta.operators import MIN\n\n result1 = talib.MIN(self.close_np, timeperiod=10)\n result2 = self.df_pl.select(MIN(pl.col(\"close\"), timeperiod=10))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_MAXINDEX(self):\n from polars_ta.ta.operators import MAXINDEX\n\n result1 = talib.MAXINDEX(self.close_np, timeperiod=5)\n result2 = self.df_pl.select(MAXINDEX(pl.col(\"close\"), timeperiod=5))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1[5:], result3[5:], equal_nan=True)\n" }, { "path": "tests/ta/test_overlap.py", "content": "import numpy as np\nimport polars as pl\nimport talib\n\n\nclass TestDemoClass:\n high_np = None\n low_np = None\n close_np = None\n df_pl = None\n\n def setup_class(self):\n self.high_np = np.arange(100, dtype=float) + np.random.rand(100)\n self.low_np = np.arange(100, dtype=float) - np.random.rand(100)\n self.close_np = np.arange(100, dtype=float)\n\n self.df_pl = pl.DataFrame([self.high_np, self.low_np, self.close_np],\n schema=[\"high\", \"low\", \"close\"])\n\n def test_SMA(self):\n from polars_ta.ta.overlap import SMA\n\n result1 = talib.SMA(self.close_np, timeperiod=3)\n result2 = self.df_pl.select(SMA(pl.col(\"close\"), d=3))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_WMA(self):\n from polars_ta.ta.overlap import WMA\n\n result1 = talib.WMA(self.close_np, timeperiod=15)\n result2 = self.df_pl.select(WMA(pl.col(\"close\"), d=15))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_EMA(self):\n from polars_ta.ta.overlap import EMA\n # !!! 此处非常重要，有部分函数受此影响\n # https://github.com/TA-Lib/ta-lib-python/blob/master/talib/_ta_lib.pxd#L28\n talib.set_compatibility(1)\n\n result1 = talib.EMA(self.close_np, timeperiod=5)\n result2 = self.df_pl.select(EMA(pl.col(\"close\"), timeperiod=5))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_DEMA(self):\n from polars_ta.ta.overlap import DEMA\n # !!! 此处非常重要，有部分函数受此影响\n talib.set_compatibility(1)\n\n result1 = talib.DEMA(self.close_np, timeperiod=3)\n result2 = self.df_pl.select(DEMA(pl.col(\"close\"), timeperiod=3))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_TEMA(self):\n from polars_ta.ta.overlap import TEMA\n # !!! 此处非常重要，有部分函数受此影响\n talib.set_compatibility(1)\n\n result1 = talib.TEMA(self.close_np, timeperiod=3)\n result2 = self.df_pl.select(TEMA(pl.col(\"close\"), timeperiod=3))\n result3 = result2['close'].to_numpy()\n # print()\n # print(result1)\n # print(result3)\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_TRIMA(self):\n from polars_ta.ta.overlap import TRIMA\n\n result1 = talib.TRIMA(self.close_np, timeperiod=6)\n result2 = self.df_pl.select(TRIMA(pl.col(\"close\"), timeperiod=6))\n result3 = result2['close'].to_numpy()\n # print()\n # print(result1)\n # print(result3)\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n # def test_KAMA(self):\n # # !!! 此处非常重要，有部分函数受此影响\n # talib.set_compatibility(1)\n #\n # from polars_ta.talib.overlap import KAMA\n #\n # result1 = talib.KAMA(self.close_np, timeperiod=5)\n # result2 = self.df_pl.select(KAMA(pl.col(\"close\"), timeperiod=5))\n # result3 = result2['close'].to_numpy()\n # # print()\n # # print(result1)\n # # print(result3)\n #\n # assert np.allclose(result1, result3, equal_nan=True)\n\n def test_BBANDS(self):\n from polars_ta.ta.overlap import BBANDS\n\n result1, _, _ = talib.BBANDS(self.close_np, timeperiod=3)\n result2 = self.df_pl.select(BBANDS(pl.col(\"close\"), timeperiod=3).struct[0])\n result3 = result2['upperband'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_MIDPOINT(self):\n from polars_ta.ta.overlap import MIDPOINT\n\n result1 = talib.MIDPOINT(self.close_np, timeperiod=3)\n result2 = self.df_pl.select(MIDPOINT(pl.col(\"close\"), timeperiod=3))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_MIDPRICE(self):\n from polars_ta.ta.overlap import MIDPRICE\n\n result1 = talib.MIDPRICE(self.high_np, self.low_np, timeperiod=3)\n result2 = self.df_pl.select(MIDPRICE(pl.col(\"high\"), pl.col(\"low\"), timeperiod=3))\n result3 = result2['high'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n" }, { "path": "tests/ta/test_statistic.py", "content": "import numpy as np\nimport polars as pl\nimport talib\n\n\nclass TestDemoClass:\n high_np = None\n low_np = None\n close_np = None\n df_pl = None\n\n def setup_class(self):\n self.high_np = np.arange(100, dtype=float) + np.random.rand(100)\n self.low_np = np.arange(100, dtype=float) - np.random.rand(100)\n self.close_np = np.arange(100, dtype=float)\n\n self.df_pl = pl.DataFrame([self.high_np, self.low_np, self.close_np],\n schema=[\"high\", \"low\", \"close\"])\n self.df_pd = self.df_pl.to_pandas()\n\n def test_STDDEV(self):\n from polars_ta.ta.statistic import STDDEV\n\n result1 = talib.STDDEV(self.close_np, timeperiod=3)\n result2 = self.df_pl.select(STDDEV(pl.col(\"close\"), timeperiod=3))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_VAR(self):\n from polars_ta.ta.statistic import VAR\n\n result1 = talib.VAR(self.close_np, timeperiod=5)\n result2 = self.df_pl.select(VAR(pl.col(\"close\"), timeperiod=5))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_CORREL(self):\n from polars_ta.ta.statistic import CORREL\n\n result1 = talib.CORREL(self.close_np, self.high_np, timeperiod=5)\n result2 = self.df_pl.select(CORREL(pl.col(\"close\"), pl.col('high'), timeperiod=5))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n" }, { "path": "tests/ta/test_volatility.py", "content": "import numpy as np\nimport polars as pl\nimport talib\n\n\nclass TestDemoClass:\n high_np = None\n low_np = None\n close_np = None\n df_pl = None\n\n def setup_class(self):\n self.high_np = np.arange(100, dtype=float) + np.random.rand(100)\n self.low_np = np.arange(100, dtype=float) - np.random.rand(100)\n self.close_np = np.arange(100, dtype=float)\n\n self.df_pl = pl.DataFrame([self.high_np, self.low_np, self.close_np],\n schema=[\"high\", \"low\", \"close\"])\n\n def test_TRANGE(self):\n from polars_ta.ta.volatility import TRANGE\n\n result1 = talib.TRANGE(self.high_np, self.low_np, self.close_np)\n result2 = self.df_pl.select(TRANGE(pl.col(\"high\"), pl.col(\"low\"), pl.col(\"close\")))\n result3 = result2['high'].to_numpy()\n # print()\n # print(result1)\n # print(result3)\n\n # value at position 0 is `x` rather than `nan` in talib\n # pl版第一个位置计算的值为max(x,nan,nan)=x比talib多一个值\n assert np.allclose(result1[1:], result3[1:], equal_nan=True)\n\n def test_ATR(self):\n # !!!\n talib.set_compatibility(1)\n\n from polars_ta.ta.volatility import ATR\n result1 = talib.ATR(self.high_np, self.low_np, self.close_np, timeperiod=5)\n result2 = self.df_pl.select(ATR(pl.col(\"high\"), pl.col(\"low\"), pl.col(\"close\"), timeperiod=5))\n result3 = result2['high'].to_numpy()\n # print()\n # print(result1)\n # print(result3)\n\n assert np.allclose(result1[-20:], result3[-20:], equal_nan=True)\n" }, { "path": "tests/ta/test_volume.py", "content": "import numpy as np\nimport polars as pl\nimport talib\n\n\nclass TestDemoClass:\n high_np = None\n low_np = None\n close_np = None\n df_pl = None\n\n def setup_class(self):\n self.high_np = np.arange(100, dtype=float) + np.random.rand(100)\n self.low_np = np.arange(100, dtype=float) - np.random.rand(100)\n self.close_np = np.random.rand(100)\n self.volume_np = np.arange(100, dtype=float) + np.random.rand(100)\n\n self.df_pl = pl.DataFrame([self.high_np, self.low_np, self.close_np, self.volume_np],\n schema=[\"high\", \"low\", \"close\", \"volume\"])\n\n def test_AD(self):\n from polars_ta.ta.volume import AD\n\n result1 = talib.AD(self.high_np, self.low_np, self.close_np, self.volume_np)\n result2 = self.df_pl.select(AD(pl.col(\"high\"), pl.col(\"low\"), pl.col(\"close\"), pl.col(\"volume\")))\n result3 = result2['close'].to_numpy()\n # print()\n # print(result1)\n # print(result3)\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_ADOSC(self):\n from polars_ta.ta.volume import ADOSC\n\n result1 = talib.ADOSC(self.high_np, self.low_np, self.close_np, self.volume_np)\n result2 = self.df_pl.select(ADOSC(pl.col(\"high\"), pl.col(\"low\"), pl.col(\"close\"), pl.col(\"volume\")))\n result3 = result2['close'].to_numpy()\n # print()\n # print(result1)\n # print(result3)\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_OBV(self):\n from polars_ta.ta.volume import OBV\n\n result1 = talib.OBV(self.close_np, self.volume_np)\n result2 = self.df_pl.select(OBV(pl.col(\"close\"), pl.col(\"volume\")))\n result3 = result2['close'].to_numpy()\n # print()\n # print(result1)\n # print(result3)\n\n assert np.allclose(result1, result3, equal_nan=True)\n" }, { "path": "tests/tdx/chip_test.py", "content": "import numpy as np\nimport polars as pl\n\nfrom polars_ta.tdx.pattern import ts_WINNER_COST\n\nhigh = np.array([10.4, 10.2, 10.2, 10.4, 10.5, 10.7, 10.7, 10.7, 10.8, 10.9])\nlow = np.array([10.3, 10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 10.7, 10.8, 10.9])\navg = np.array([10.3, 10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 10.7, 10.8, 10.9])\nclose = np.array([10.3, 10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 10.7, 10.8, 10.9])\nturnover = np.array([0.3, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])\ncost = np.array([0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5])\n\nstep = 0.1\n\ndf = pl.DataFrame({'high': high, 'low': low, 'avg': avg, 'close': close, 'turnover': turnover, 'cost': cost})\ndf = df.with_columns(WINNER=ts_WINNER_COST(pl.col('high'), pl.col('low'), pl.col('avg'), pl.col('turnover'), pl.col('close'), 0.5, step=step))\nprint(df)\n" }, { "path": "tests/tdx/test_reference.py", "content": "import numpy as np\nimport polars as pl\nimport talib\n\n\nclass TestDemoClass:\n high_np = None\n low_np = None\n close_np = None\n df_pl = None\n\n def setup_class(self):\n self.high_np = np.arange(100, dtype=float) + np.random.rand(100)\n self.low_np = np.arange(100, dtype=float) - np.random.rand(100)\n self.close_np = np.arange(100, dtype=float)\n\n self.df_pl = pl.DataFrame([self.high_np, self.low_np, self.close_np],\n schema=[\"high\", \"low\", \"close\"])\n\n def test_EXPMEMA(self):\n from polars_ta.tdx.reference import EXPMEMA\n\n result1 = talib.EMA(self.close_np, timeperiod=6)\n result2 = self.df_pl.select(EXPMEMA(pl.col(\"close\"), N=6))\n result3 = result2['close'].to_numpy()\n\n assert np.allclose(result1, result3, equal_nan=True)\n\n def test_BARSLAST(self):\n from polars_ta.tdx.reference import BARSLAST\n\n df = pl.DataFrame({'A': [1, 1, 1, 1, 1, 1]})\n df = pl.DataFrame({'A': [0, 0, 0, 0, 0, 0, 0, 0]})\n result2 = df.select(BARSLAST(pl.col('A')))\n print(result2)\n\n def test_BARSLASTCOUNT(self):\n from polars_ta.tdx.reference import BARSLASTCOUNT\n\n df = pl.DataFrame({'A': [1, 1, 1, 1, 1, 1]})\n df = pl.DataFrame({'A': [0, 0, 0, 0, 0, 0, 0, 0]})\n result2 = df.select(BARSLASTCOUNT(pl.col('A')))\n print(result2)\n\n def test_BARSSINCEN(self):\n from polars_ta.tdx.reference import BARSSINCEN\n\n df = pl.DataFrame({'A': [1, 1, 1, 1, 1, 1, 1, 1]})\n df = pl.DataFrame({'A': [0, 2, 1, 0, 0, 0, 0, 0]})\n result2 = df.select(BARSSINCEN(pl.col('A'), 4))\n print(result2)\n\n def test_LOD(self):\n from polars_ta.tdx.reference import LOD\n\n df = pl.DataFrame({'A': [0, 1, 0, 0, 1, 1, 1, 1]})\n result2 = df.select(pl.col('A').rank())\n print(result2)\n result2 = df.select(LOD(pl.col('A'), 3))\n print(result2)\n\n def test_FILTER(self):\n from polars_ta.tdx.reference import FILTER\n\n df = pl.DataFrame({'A': [0, 1, 0, 1, 1, 1, 1, 1]})\n # result2 = df.select(FILTER(pl.col('A'), 3))\n # print(result2)\n" }, { "path": "tests/tdx/test_statistic.py", "content": "import time\n\nimport numpy as np\nimport polars as pl\n\n\nclass TestDemoClass:\n high_np = None\n low_np = None\n close_np = None\n df_pl = None\n\n def setup_class(self):\n self.high_np = np.arange(100, dtype=float) + np.random.rand(100)\n self.low_np = np.arange(100, dtype=float) - np.random.rand(100)\n self.close_np = np.arange(100, dtype=float)\n\n self.df_pl = pl.DataFrame([self.high_np, self.low_np, self.close_np],\n schema=[\"high\", \"low\", \"close\"])\n self.df_pl = self.df_pl.with_columns(pl.lit(None).alias(\"null\"))\n\n def test_AVEDEV(self):\n from polars_ta.tdx._slow import AVEDEV as func_slow\n from polars_ta.tdx.statistic import AVEDEV as func_fast\n\n xx = self.df_pl.with_columns(\n a1=func_slow(pl.col('high'), 10),\n a2=func_fast(pl.col('high'), 10),\n )\n print(xx)\n\n t1 = time.perf_counter()\n for i in range(1000):\n result2 = self.df_pl.with_columns(\n a1=func_slow(pl.col('high'), 10),\n )\n t2 = time.perf_counter()\n print(t2 - t1)\n\n t1 = time.perf_counter()\n for i in range(1000):\n result2 = self.df_pl.with_columns(\n a1=func_fast(pl.col('high'), 10),\n )\n t2 = time.perf_counter()\n print(t2 - t1)\n\n\n" }, { "path": "tests/wq/test_arithmetic.py", "content": "import numpy as np\nimport polars as pl\n\nfrom polars_ta.wq.arithmetic import add\n\n\nclass TestDemoClass:\n high_np = None\n low_np = None\n close_np = None\n df_pl = None\n\n def setup_class(self):\n self.high_np = np.arange(100, dtype=float) + np.random.rand(100)\n self.low_np = np.arange(100, dtype=float) - np.random.rand(100)\n self.close_np = np.arange(100, dtype=float)\n\n self.df_pl = pl.DataFrame([self.high_np, self.low_np, self.close_np],\n schema=[\"high\", \"low\", \"close\"])\n self.df_pl = self.df_pl.with_columns(pl.lit(None).alias(\"null\"))\n\n def test_add(self):\n print(self.df_pl.schema)\n result2 = self.df_pl.select(add(pl.col('high'), pl.col('low'), pl.col('null'), filter_=True))\n print(result2)\n # result3 = self.df_pl.select(add(pl.col('high'), pl.col('low'), pl.col('literal'), filter_=True))\n # print(result3)\n\n def test_fraction(self):\n from polars_ta.wq.arithmetic import fraction\n\n a = np.array([5.63, -4.59])\n df = pl.DataFrame({'a': a})\n result1 = df.select(fraction(pl.col('a')))\n result3 = np.array([0.63, -0.59])\n assert np.allclose(result1.to_numpy()[:, 0], result3, equal_nan=True)\n" }, { "path": "tests/wq/test_time_series.py", "content": "import time\n\nimport numpy as np\nimport polars as pl\nfrom pandas.testing import assert_frame_equal, assert_series_equal\n\n\nclass TestDemoClass:\n high_np = None\n low_np = None\n close_np = None\n df_pl = None\n\n def setup_class(self):\n self.high_np = np.random.rand(100)\n self.low_np = np.arange(100, dtype=float) - np.random.rand(100)\n self.close_np = np.arange(100, dtype=float)\n\n self.df_pl = pl.DataFrame([self.high_np, self.low_np, self.close_np],\n schema=[\"high\", \"low\", \"close\"])\n self.df_pl = self.df_pl.cast(pl.Float32)\n self.df_pl = self.df_pl.with_columns(pl.lit(None).alias(\"null\"))\n self.df_pd = self.df_pl.to_pandas()\n\n def test_ts_rank(self):\n from polars_ta.wq.time_series import ts_rank\n\n result1 = self.df_pd[[\"high\", 'low']].rolling(5).rank(pct=True)\n result2 = self.df_pl.select(ts_rank(pl.col([\"high\", 'low']), d=5))\n\n # 底层一样，结果就应当一样\n assert_frame_equal(result1, result2.to_pandas())\n\n def test_ts_skewness(self):\n from polars_ta.wq.time_series import ts_skewness\n\n # cannot set tune bias in pandas\n # 好像没办法告诉pandas关闭偏差校正\n result1 = self.df_pd[[\"high\"]].rolling(5).skew()\n result2 = self.df_pl.select(ts_skewness(pl.col([\"high\"]), d=5))\n # print(result1)\n # print(result2)\n\n assert_frame_equal(result1, result2.to_pandas().astype(float))\n\n def test_ts_kurtosis(self):\n from polars_ta.wq.time_series import ts_kurtosis\n\n result1 = self.df_pd[[\"high\"]].rolling(5).kurt()\n result2 = self.df_pl.select(ts_kurtosis(pl.col([\"high\"]), d=5))\n\n # 底层一样，结果就应当一样\n assert_frame_equal(result1, result2.to_pandas())\n\n def test_ts_corr(self):\n from polars_ta.wq.time_series import ts_corr\n\n result1 = self.df_pd[\"high\"].rolling(5).corr(self.df_pd[\"low\"])\n result2 = self.df_pl.select(ts_corr(pl.col(\"high\"), pl.col(\"low\"), d=5))\n\n # assert_series_equal(result1, result2.to_series(0).to_pandas(), check_names=False)\n\n def test_ts_covariance(self):\n from polars_ta.wq.time_series import ts_covariance\n\n result1 = self.df_pd[\"high\"].rolling(5).cov(self.df_pd[\"low\"])\n result2 = self.df_pl.select(ts_covariance(pl.col(\"high\"), pl.col(\"low\"), d=5))\n\n # assert_series_equal(result1, result2.to_series(0).to_pandas().astype(float), check_names=False)\n\n def test_ts_decay_exp_window(self):\n\n df = pl.DataFrame({'A': range(100)})\n print(df)\n # result1 = df.select(ts_decay_exp_window(pl.col('A'), 10, factor=0.5))\n # print(result1)\n\n def test_ts_decay_linear(self):\n\n df = pl.DataFrame({'A': [None] * 10, 'B': [None, None, None, None, 4, 5, 6, 7, 8, 9]})\n print(df)\n # result1 = df.select(ts_decay_linear(pl.col('B'), 3))\n # print(result1)\n\n def test_ts_weighted_delay(self):\n from polars_ta.wq.time_series import ts_weighted_delay\n\n df = pl.DataFrame({'A': [4, 6]})\n print(df)\n result1 = df.select(ts_weighted_delay(pl.col('A'), k=0.25))\n print(result1)\n\n def test_ts_arg_max(self):\n from polars_ta.wq._slow import ts_arg_min as func_slow\n from polars_ta.wq.time_series import ts_arg_min as func_fast\n\n xx = self.df_pl.with_columns(\n a1=func_slow(pl.col('high'), 10),\n a2=func_fast(pl.col('high'), 10),\n )\n print(xx)\n\n t1 = time.perf_counter()\n for i in range(10000):\n result2 = self.df_pl.with_columns(\n a1=func_slow(pl.col('high'), 10),\n )\n t2 = time.perf_counter()\n print(t2 - t1)\n\n t1 = time.perf_counter()\n for i in range(10000):\n result2 = self.df_pl.with_columns(\n a1=func_fast(pl.col('high'), 10),\n )\n t2 = time.perf_counter()\n print(t2 - t1)\n\n def test_ts_product(self):\n from polars_ta.wq._slow import ts_product as func_slow\n from polars_ta.wq.time_series import ts_product as func_fast\n\n xx = self.df_pl.with_columns(\n a1=func_slow(pl.col('high'), 10),\n a2=func_fast(pl.col('high'), 10),\n )\n print(xx)\n\n t1 = time.perf_counter()\n for i in range(10000):\n result2 = self.df_pl.with_columns(\n a1=func_slow(pl.col('high'), 10),\n )\n t2 = time.perf_counter()\n print(t2 - t1)\n\n t1 = time.perf_counter()\n for i in range(10000):\n result2 = self.df_pl.with_columns(\n a1=func_fast(pl.col('high'), 10),\n )\n t2 = time.perf_counter()\n print(t2 - t1)\n" }, { "path": "thinking_about_TA.md", "content": "# About Technical Indicators\n\nThere are three types of technical indicators\n\n1. Time series indicators. Must be sorted by time, and then calculated. It is difficult to handle null in the middle.\n2. Cross-sectional indicators. No order requirement, the whole row is calculated, and there may be a null, which needs to be compatible.\n3. Single element indicators. Can be single-column, multi-column, and can be calculated in any way.\n\nIn terms of calculation, time series indicators have one more dimension than cross-sectional indicators. For example\n\n1. `cs_rank` is a cross-sectional sort\n2. `ts_rank` is to calculate `cs_rank` in a rolling time window, take `[-1]` each time, and then concatenate them\n\n3. `cs_zscore` first aggregates to calculate `mean` and `std`, then broadcast `mean` and `std`, and perform a one-dimensional calculation with `x`\n\nSo common technical indicators are generally the flexible application of `aggregation` and `broadcasting`\n\nSince `ts_` indicators are based on `cs_`, theoretically `rolling` can be used directly, but in practice, it is generally not used in this way.\n`cs_` is generally rewritten in `cython`, `numba`, etc., and if you mix `ts_` and `cs_`, `ts_` will make thousands of `cs_` call.\nSo it is common to put `rolling` operations into `cython`, `numba`, etc. as well.\n\n## Evolve of Our TA-Lib Wrappers\n\n1. `Expr.map_batches` can be used to call third-party libraries, such as `TA-Lib, bottleneck`. But because of the input and output format requirements, you need to wrap the third-party API with a function.\n - Both input and output can only be one column. If you want to support multiple columns, you need to convert them to `pl.Struct`. After that, you need to use `unnest` to split `pl.Struct`.\n - The output must be `pl.Series`\n\n2. Start to use `register_expr_namespace` to simplify the code\n - Implementation [helper.py](polars_ta/utils/helper.py)\n - Usage demo [demo_ta1.py](examples/demo_ta1.py)\n - Pros: Easy to use\n - Cons:\n - The `member function call mode` is not convenient for inputting into genetic algorithms for factor mining\n - `__getattribute__` dynamic method call is very flexible, but loses `IDE` support.\n\n3. Prefix expression. Convert all member functions into formulas\n - Implementation [wrapper.py](polars_ta/utils/wrapper.py)\n - Usage demo [demo_ta2.py](examples/demo_ta2.py)\n - Pros: Can be input into our implementation of genetic algorithms\n - Cons: `__getattribute__` dynamic method call is very flexible, but loses `IDE` support.\n\n4. Code generation.\n - Implementation [codegen_talib.py](tools/codegen_talib.py)\n - Generated result will be at [\\_\\_init\\_\\_.py](polars_ta/talib/__init__.py)\n - Usage demo [demo_ta3.py](examples/demo_ta3.py)\n - Pros:\n - Can be input into our implementation of genetic algorithms\n - `IDE` support\n\n# 对技术指标的再思考\n\n用使用方式来说，指标分为三种\n\n1. 时序指标。必须按时间排序，然后计算，中段出现null比较难处理\n2. 截面指标。对顺序无要求，整行计算即可，数据会有null，需要能兼容\n3. 单元素指标。可单列、多列，可任意方式计算\n\n从计算原理上来说，时序指标比截面指标多了一个计算维度。例如\n\n1. `cs_rank`是横截面排序\n2. 而`ts_rank`是对数据滚动时间窗口计算`cs_rank`,每次取`[-1]`，然后拼接起来\n\n又例如\n\n1. `cs_zscore`先聚合计算`mean`和`std`\n2. 然后将`mean`与`std`广播,与`x`进行一维计算\n\n所以常见的技术指标一般是`聚合`与`广播`的灵活应用\n\n由于`ts_`指标基于`cs_`,理论上直接`rolling`即可，但在实践中一般不这么用，因为为了求快，`cs_`一般是更底层的语言编写，如`cython`、`numba`等，\n本来`cs_`版`python`与底层只交互一次，而`ts_`调用`cs_`版会导致交互千万次，性能极低\n一般是把`rolling`操作也放在底层，在底层循环调用`cs_`底层版\n\n## TA-Lib封装的演化\n\n1. `Expr.map_batches`可以实现调用第三方库，如`TA-Lib, bottleneck`。但因为对输入与输出格式有要求，所以还需要用函数对第三方API封装一下。\n - 输入输出都只能是一列，如要支持多列需转换成`pl.Struct`。事后`pl.Struct`要拆分需使用`unnest`\n - 输出必须是`pl.Series`\n2. 参数多，代码长。开始使用`register_expr_namespace`来简化代码\n - 实现代码[helper.py](polars_ta/utils/helper.py)\n - 使用演示[demo_ta1.py](examples/demo_ta1.py)\n - 优点：使用简单\n - 不足：`成员函数调用模式`不便于输入到遗传算法中进行因子挖掘\n - 不足：`__getattribute__`动态方法调用非常灵活，但失去了`IDE`智能提示\n3. 前缀表达式。将所有的成员函数都转换成公式\n - 实现代码[wrapper.py](polars_ta/utils/wrapper.py)\n - 使用演示[demo_ta2.py](examples/demo_ta2.py)\n - 优点：可以输入到遗传算法\n - 不足：`__getattribute__`动态方法调用非常灵活，但失去了`IDE`智能提示\n4. 代码自动生成。\n - 实现代码[codegen_talib.py](tools/codegen_talib.py)\n - 生成结果[\\_\\_init\\_\\_.py](polars_ta/talib/__init__.py)\n - 使用演示[demo_ta3.py](examples/demo_ta3.py)\n - 优点：即可以输入到遗传算法，`IDE`还有智能提示\n\n" }, { "path": "tools/README.md", "content": "# Code Conversion Tools\n\n## TA-Lib tool (codegen_talib.py)\n\nWrap the original `TA-Lib` and save it to `polars_ta.talib.__init__.py`.\n\n1. Support `Expr`\n2. Support `skipna`\n3. Multi-output support\n\n## Prefix adding tool\n\nAdd the `ts_` prefix to some functions to facilitate use in `expr_codegen` and other tools.\n\n1. `prefix_ta.py` adds a prefix to `polars_ta.ta` and saves it to `polars_ta.prefix.ta.py`\n2. `prefix_tdx.py` adds a prefix to `polars_ta.tdx` and saves it to `polars_ta.prefix.tdx.py`\n3. `prefix_talib.py` adds a prefix to `polars_ta.talib` and saves it to `polars_ta.prefix.talib.py` (same as `codegen_talib`)\n\n# 代码转换工具\n\n## TA-Lib工具(codegen_talib.py)\n\n将原版`TA-Lib`封装，保存到`polars_ta.talib.__init__.py`，支持以下功能\n\n1. 支持`Expr`\n2. 支持跳过空值\n3. 多输出适配\n\n## 前缀添加工具\n\n为部分函数添加`ts_`前缀，方便在`expr_codegen`等工具中使用\n\n1. `prefix_ta.py`为`polars_ta.ta`添加前缀，保存到`polars_ta.prefix.ta.py`\n2. `prefix_tdx.py`为`polars_ta.tdx`添加前缀，保存到`polars_ta.prefix.tdx.py`\n3. `prefix_talib.py`为`polars_ta.talib`添加前缀，保存到`polars_ta.prefix.talib.py`(使用与codegen_talib同技术)\n\n\n\n" }, { "path": "tools/codegen_talib.py", "content": "\"\"\"\nParse talib functions for polars Expressions\n\nThis version is more direct, without skip nan values, and without input and output checks\n\npolars does not skip nan values as well. It should be processed by specific functions\n\n本脚本主要功能是将talib封装成更适合表达式的版本\n\n与另一版本的区别是这版本调用更直接，没有跳过空的操作，也没有输入与输出数量的判断工作\n跳过空值等操作与polars样都不做，以后准备统一交给函数处理\n\"\"\"\nimport talib as _talib\nfrom talib import abstract as _abstract\n\nfrom tools.prefix import save\n\n\ndef _codegen_func(name, input_names, parameters, output_names, doc):\n tpl11 = \"\"\"\ndef {name}({aa}) -> Expr: # {output_names}\n \\\"\\\"\\\"{doc}\\\"\\\"\\\"\n return {bb}.map_batches(lambda x1: batches_i1_o1(x1.to_numpy().astype(float), {cc}), return_dtype={return_dtype})\n\"\"\"\n tpl12 = \"\"\"\ndef {name}({aa}) -> Expr: # {output_names}\n \\\"\\\"\\\"{doc}\\\"\\\"\\\"\n dtype = Struct([Field(f\"column_{{i}}\", Float64) for i in range({ee})])\n return {bb}.map_batches(lambda x1: batches_i1_o2(x1.to_numpy().astype(float), {cc}), return_dtype=dtype)\n\"\"\"\n tpl21 = \"\"\"\ndef {name}({aa}) -> Expr: # {output_names}\n \\\"\\\"\\\"{doc}\\\"\\\"\\\"\n return struct({bb}).map_batches(lambda xx: batches_i2_o1(struct_to_numpy(xx, {dd}, dtype=float), {cc}), return_dtype={return_dtype})\n\"\"\"\n tpl22 = \"\"\"\ndef {name}({aa}) -> Expr: # {output_names}\n \\\"\\\"\\\"{doc}\\\"\\\"\\\"\n dtype = Struct([Field(f\"column_{{i}}\", Float64) for i in range({ee})])\n return struct({bb}).map_batches(lambda xx: batches_i2_o2(struct_to_numpy(xx, {dd}, dtype=float), {cc}), return_dtype=dtype)\n\"\"\"\n if len(output_names) > 42:\n extra_args = {'ret_idx': len(output_names) - 1}\n else:\n extra_args = {}\n a1 = [f'{name}: Expr' for name in input_names]\n a2 = [f'{k}: {type(v).__name__} = {v}' for k, v in parameters.items()]\n a3 = [f'{k}: {type(v).__name__} = {v}' for k, v in extra_args.items()]\n aa = ', '.join(a1 + a2 + a3)\n\n bb = ', '.join(input_names)\n if len(input_names) > 1:\n bb = [f'f{i}={arg}' for i, arg in enumerate(input_names)]\n bb = ', '.join(bb)\n\n c1 = [f'_ta.{name}']\n if len(parameters) > 0:\n c2 = [f'{k}' for k, v in parameters.items()]\n else:\n c2 = []\n\n c3 = [f'{k}={k}' for k, v in extra_args.items()]\n cc = ', '.join(c1 + c2 + c3)\n\n if output_names[0] == 'integer':\n return_dtype = 'Int32'\n else:\n return_dtype = 'Float64'\n\n if len(input_names) == 1 and len(output_names) == 1:\n return tpl11.format(name=name, aa=aa, bb=bb, cc=cc, dd=len(input_names), ee=len(output_names), output_names=output_names, doc=doc, return_dtype=return_dtype)\n elif len(input_names) == 1 and len(output_names) > 1:\n return tpl12.format(name=name, aa=aa, bb=bb, cc=cc, dd=len(input_names), ee=len(output_names), output_names=output_names, doc=doc)\n elif len(input_names) > 1 and len(output_names) == 1:\n return tpl21.format(name=name, aa=aa, bb=bb, cc=cc, dd=len(input_names), ee=len(output_names), output_names=output_names, doc=doc, return_dtype=return_dtype)\n else:\n return tpl22.format(name=name, aa=aa, bb=bb, cc=cc, dd=len(input_names), ee=len(output_names), output_names=output_names, doc=doc)\n\n\ndef codegen():\n head_v2 = \"\"\"# generated by codegen_talib.py\nimport talib as _ta\nfrom polars import Expr, struct, Struct, Field, Float64, Int32\n\nfrom polars_ta.utils.numba_ import batches_i1_o1, batches_i1_o2, batches_i2_o1, batches_i2_o2, struct_to_numpy\n\"\"\"\n\n txts = [head_v2]\n for i, func_name in enumerate(_talib.get_functions()):\n \"\"\"talib遍历\"\"\"\n info = _abstract.Function(func_name).info\n\n name = info['name']\n input_names = []\n for in_names in info['input_names'].values():\n if isinstance(in_names, (list, tuple)):\n input_names.extend(list(in_names))\n else:\n input_names.append(in_names)\n parameters = info['parameters']\n output_names = info['output_names']\n txt = _codegen_func(name, input_names, parameters, output_names, getattr(_talib, name).__doc__)\n txts.append(txt)\n\n return txts\n\n\nif __name__ == '__main__':\n txts = codegen()\n save(txts, module='polars_ta.talib', write=True)\n" }, { "path": "tools/prefix.py", "content": "import inspect\nfrom typing import List, Optional\n\n\ndef codegen_import_as(module: str, prefix: str = 'ts_',\n include_modules: Optional[List[str]] = None,\n include_func: Optional[List[str]] = None,\n exclude_func: Optional[List[str]] = None,\n include_parameter: Optional[List[str]] = None):\n \"\"\"Generate codes by `reflection`\n 通过反射，生成代码的小工具\n\n Parameters\n ----------\n module\n 模块全名\n prefix\n 需要添加的前缀\n include_modules\n 通过`from import`导入的函数也参与判断\n include_func\n 指定的函数名不做检查，直接添加前缀\n exclude_func\n 指定的函数名直接跳过不导入\n include_parameter\n 出现了同名int参数，添加前缀\n\n Notes\n -----\n\n \"\"\"\n if include_modules is None:\n include_modules = [module]\n else:\n include_modules += [module]\n if include_func is None:\n include_func = []\n if exclude_func is None:\n exclude_func = []\n if include_parameter is None:\n include_parameter = []\n\n m = __import__(module, fromlist=['*'])\n funcs = inspect.getmembers(m, inspect.isfunction)\n funcs = [f for f in funcs if not f[0].startswith('_')]\n txts = []\n for name, func in funcs:\n if func.__module__ not in include_modules:\n continue\n if name in exclude_func:\n continue\n\n add_prefix = False\n if name.startswith(prefix):\n add_prefix = False\n elif name in include_func:\n add_prefix = True\n else:\n p = inspect.signature(func).parameters\n for n in include_parameter:\n if n in p:\n if p[n].annotation == int:\n add_prefix = True\n break\n\n if add_prefix:\n txts.append(f'from {module} import {name} as {prefix}{name} # noqa')\n else:\n txts.append(f'from {module} import {name} # noqa')\n\n return txts\n\n\ndef save(txts, module, write=False):\n m = __import__(module, fromlist=['*'])\n file = m.__file__\n print('save to', file)\n text = '\\n'.join(txts)\n if write:\n with open(file, 'w', encoding='utf-8') as f:\n f.write(text)\n else:\n print(text)\n" }, { "path": "tools/prefix_ta.py", "content": "from tools.prefix import codegen_import_as, save\n\nlines = [\"\"\"# this code is auto generated by tools/prefix_ta.py\n\"\"\"]\nlines += codegen_import_as('polars_ta.ta.momentum', include_parameter=['timeperiod', 'fastperiod', 'fastk_period'])\nlines += codegen_import_as('polars_ta.ta.operators', include_modules=['polars_ta.wq.arithmetic', 'polars_ta.wq.time_series'], include_parameter=['timeperiod', 'd'], exclude_func=['ts_arg_max', 'ts_arg_min'])\nlines += codegen_import_as('polars_ta.ta.overlap', include_modules=['polars_ta.wq.time_series'], include_func=['SMA', 'WMA'], include_parameter=['timeperiod', 'd'], exclude_func=['MAX', 'MIN'])\nlines += codegen_import_as('polars_ta.ta.price', include_parameter=['timeperiod'])\nlines += codegen_import_as('polars_ta.ta.statistic', include_modules=['polars_ta.wq.time_series'], include_parameter=['timeperiod', 'd'], exclude_func=['ts_std_dev'])\nlines += codegen_import_as('polars_ta.ta.transform', include_modules=['polars_ta.wq.arithmetic'], include_parameter=['timeperiod', 'd'])\nlines += codegen_import_as('polars_ta.ta.volatility', include_func=['TRANGE'], include_parameter=['timeperiod'])\nlines += codegen_import_as('polars_ta.ta.volume', include_func=['AD', 'OBV'], include_parameter=['fastperiod'])\nsave(lines, module='polars_ta.prefix.ta', write=True)\n" }, { "path": "tools/prefix_talib.py", "content": "\"\"\"\nAdd prefix to talib functions\n本脚本主要功能是将talib封装添加前缀\n\"\"\"\nimport talib as _talib\nfrom talib import abstract as _abstract\n\nfrom tools.prefix import save\n\n\ndef codegen():\n txts = []\n for i, func_name in enumerate(_talib.get_functions()):\n \"\"\"talib遍历\"\"\"\n info = _abstract.Function(func_name).info\n group = info['group']\n name = info['name']\n if group in ('Math Operators', 'Math Transform', 'Price Transform'):\n txts.append(f'from polars_ta.talib import {name} # noqa')\n else:\n txts.append(f'from polars_ta.talib import {name} as ts_{name} # noqa')\n return txts\n\n\nif __name__ == '__main__':\n txts = [\"\"\"# this code is auto generated by tools/prefix_talib.py\n \"\"\"]\n txts += codegen()\n save(txts, module='polars_ta.prefix.talib', write=True)\n" }, { "path": "tools/prefix_tdx.py", "content": "from tools.prefix import codegen_import_as, save\n\nlines = [\"\"\"# this code is auto generated by tools/prefix_tdx.py\n\"\"\"]\nlines += codegen_import_as('polars_ta.tdx.arithmetic', include_modules=['polars_ta.wq.arithmetic'])\nlines += codegen_import_as('polars_ta.tdx.choice', include_func=['VALUEWHEN'])\nlines += codegen_import_as('polars_ta.tdx.energy', include_parameter=['N', 'N1'])\nlines += codegen_import_as('polars_ta.tdx.logical', include_func=['CROSS', 'LONGCROSS', 'EXISTR', 'LAST'], include_parameter=['N'])\nlines += codegen_import_as('polars_ta.tdx.moving_average', include_parameter=['M1'])\nlines += codegen_import_as('polars_ta.tdx.over_bought_over_sold', include_modules=['polars_ta.ta.momentum'], include_parameter=['N', 'timeperiod'])\nlines += codegen_import_as('polars_ta.tdx.pattern')\nlines += codegen_import_as('polars_ta.tdx.pattern_feature', include_parameter=['N', 'M', 'N1'], include_func=['早晨之星', '剑', '天量法则', '四串阴', '四串阳', '出水芙蓉II', '高开大阴线', '低开大阳线', '跳空缺口选股', '老鸭头', '仙人指路'])\nlines += codegen_import_as('polars_ta.tdx.pressure_support', include_parameter=['N'])\nlines += codegen_import_as('polars_ta.tdx.reference', include_modules=['polars_ta.ta.overlap', 'polars_ta.ta.volatility', 'polars_ta.wq.arithmetic', 'polars_ta.wq.time_series'], include_func=['BARSLAST', 'BARSLASTCOUNT', 'BARSSINCE', 'DMA', 'CUMSUM', 'TR'], include_parameter=['N', 'd', 'timeperiod'])\nlines += codegen_import_as('polars_ta.tdx.statistic', include_modules=['polars_ta.wq.time_series'], include_parameter=['timeperiod', 'd'])\nlines += codegen_import_as('polars_ta.tdx.times', include_parameter=[])\nlines += codegen_import_as('polars_ta.tdx.trend', include_parameter=['N'])\nlines += codegen_import_as('polars_ta.tdx.trend_feature', include_parameter=['N', 'M', 'N1'], include_func=['下跌多日再放量上涨', '价量渐低后阳包阴', '单日放量', '跳空高开或低开'])\nlines += codegen_import_as('polars_ta.tdx.volume', include_func=['OBV'], include_parameter=['N'])\nsave(lines, module='polars_ta.prefix.tdx', write=True)\n" }, { "path": "tools/prefix_vec.py", "content": "\"\"\"\n将vec_xxx改成cs_vec_xxx\n\n主要用于gp_vect_xxxx转换成cs_vec_xxxx\n\n\"\"\"\n\nimport inspect\nfrom typing import Optional, List\n\nfrom tools.prefix import save\n\n\ndef codegen_import_as(module: str, prefix: str = 'cs_',\n include_modules: Optional[List[str]] = None,\n include_func: Optional[List[str]] = None,\n exclude_func: Optional[List[str]] = None,\n include_parameter: Optional[List[str]] = None):\n \"\"\"Generate codes by `reflection`\n 通过反射，生成代码的小工具\n\n Parameters\n ----------\n module\n 模块全名\n prefix\n 需要添加的前缀\n include_modules\n 通过`from import`导入的函数也参与判断\n include_func\n 指定的函数名不做检查，直接添加前缀\n exclude_func\n 指定的函数名直接跳过不导入\n include_parameter\n 出现了同名int参数，添加前缀\n\n Notes\n -----\n\n \"\"\"\n if include_modules is None:\n include_modules = [module]\n else:\n include_modules += [module]\n if include_func is None:\n include_func = []\n if exclude_func is None:\n exclude_func = []\n if include_parameter is None:\n include_parameter = []\n\n m = __import__(module, fromlist=['*'])\n funcs = inspect.getmembers(m, inspect.isfunction)\n funcs = [f for f in funcs if not f[0].startswith('_')]\n txts = []\n for name, func in funcs:\n if func.__module__ not in include_modules:\n continue\n if name in exclude_func:\n continue\n\n add_prefix = True\n\n if add_prefix:\n txts.append(f'from {module} import {name} as {prefix}{name} # noqa')\n else:\n txts.append(f'from {module} import {name} # noqa')\n\n return txts\n\n\nlines = [\"\"\"# this code is auto generated by tools/prefix_vec.py\n\"\"\"]\nlines += codegen_import_as('polars_ta.wq.vector', prefix='cs_')\nsave(lines, module='polars_ta.prefix.vec', write=True)\n" }, { "path": "tools/prompt.py", "content": "import inspect\nfrom typing import List, Optional\n\n\ndef get_annotation(annotation):\n output = annotation.__name__\n if output == \"Optional\":\n output = str(annotation).split('.')[-1]\n return output\n\n\ndef get_parameter(p):\n annotation = get_annotation(p.annotation)\n if p.kind == inspect._ParameterKind.VAR_POSITIONAL:\n output = f\"*{p.name}\"\n elif p.kind == inspect._ParameterKind.VAR_KEYWORD:\n output = f\"**{p.name}\"\n else:\n output = p.name\n if annotation not in (\"Expr\", \"_empty\"):\n output += f\":{annotation}\"\n if p.default != inspect._empty:\n output += f\"={p.default}\"\n\n return output\n\n\ndef codegen_import_as(module: str,\n include_modules: Optional[List[str]] = None,\n include_func: Optional[List[str]] = None,\n exclude_func: Optional[List[str]] = None):\n \"\"\"Generate codes by `reflection`\n 通过反射，生成代码的小工具\n\n Parameters\n ----------\n module\n 模块全名\n include_modules\n 通过`from import`导入的函数也参与判断\n include_func\n 指定的函数名不做检查，直接添加前缀\n exclude_func\n 指定的函数名直接跳过不导入\n\n Notes\n -----\n\n \"\"\"\n if include_modules is None:\n include_modules = [module]\n else:\n include_modules += [module]\n if include_func is None:\n include_func = []\n if exclude_func is None:\n exclude_func = []\n\n m = __import__(module, fromlist=['*'])\n funcs = inspect.getmembers(m, inspect.isfunction)\n funcs = [f for f in funcs if not f[0].startswith('_')]\n txts = [f\"### {module}\"]\n for name, func in funcs:\n if func.__module__ not in include_modules:\n continue\n if exclude_func and name in exclude_func:\n continue\n if include_func and name not in include_func:\n continue\n\n # if name != \"cs_resid\":\n # continue\n\n signature = inspect.signature(func)\n parameters = []\n for n, p in signature.parameters.items():\n if p.name == \"min_samples\":\n continue\n parameters.append(get_parameter(p))\n\n # txts.append(f'- {name}({\",\".join(parameters)})->{get_annotation(signature.return_annotation)} : {func.__doc__.splitlines()[0]}')\n txts.append(f'- {name}({\",\".join(parameters)}) : {func.__doc__.splitlines()[0]}')\n\n return txts\n\n\n# 过滤一些很少用到的函数，最好控制在5000字内（百度5000字限制）。部分智能体只支持1000字符\nlines = []\nlines += codegen_import_as('polars_ta.wq.arithmetic',\n exclude_func=['add', 'subtract', 'multiply', 'div', 'divide', 'reverse', 'inverse', \"mean\",\n 'arc_cos', 'arc_sin', 'arc_tan', 'arc_tan2', 'cot', 'cosh', 'tanh', 'sinh', 'degrees', 'radians',\n 's_log_1p', 'softsign', 'log2', \"expm1\",\n ])\nlines += codegen_import_as('polars_ta.wq.time_series',\n exclude_func=['ts_co_kurtosis', 'ts_co_skewness', 'ts_count_nans', 'ts_count_nulls',\n 'ts_cum_prod', 'ts_cum_prod_by', 'ts_cum_sum', 'ts_cum_sum_by', 'ts_cum_sum_reset',\n \"ts_min_max_cps\", \"ts_min_max_diff\", \"ts_signals_to_size\", 'ts_sum_split_by',\n \"ts_cum_count\", \"ts_cum_max\", \"ts_cum_min\", \"ts_triple_corr\", \"ts_partial_corr\",\n \"ts_max_diff\", \"ts_min_diff\",\n ])\nlines += codegen_import_as('polars_ta.wq.cross_sectional',\n exclude_func=['cs_fill_except_all_null', 'cs_fill_mean', 'cs_fill_max', 'cs_fill_min',\n 'cs_top_bottom', 'cs_one_side',\n 'cs_regression_neut', 'cs_regression_proj', 'cs_scale_down', 'cs_truncate'])\nlines += codegen_import_as('polars_ta.wq.preprocess',\n exclude_func=['cs_mad_zscore', 'cs_mad_zscore_resid', 'cs_mad_zscore_resid_zscore', 'cs_zscore_resid'])\nlines += codegen_import_as('polars_ta.wq.logical',\n include_func=['if_else', ])\nlines += codegen_import_as('polars_ta.wq.transformational',\n include_func=['cut', 'sigmoid', 'bool_', 'int_', 'float_'])\ntext = '\\n'.join(lines)\nprint(text)\n" } ]