stock-embeddings 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- stock_embeddings-0.1.0/LICENSE +21 -0
- stock_embeddings-0.1.0/PKG-INFO +177 -0
- stock_embeddings-0.1.0/README.md +116 -0
- stock_embeddings-0.1.0/pyproject.toml +51 -0
- stock_embeddings-0.1.0/setup.cfg +4 -0
- stock_embeddings-0.1.0/stock_embeddings/__init__.py +21 -0
- stock_embeddings-0.1.0/stock_embeddings/api.py +132 -0
- stock_embeddings-0.1.0/stock_embeddings/config.py +97 -0
- stock_embeddings-0.1.0/stock_embeddings/data.py +83 -0
- stock_embeddings-0.1.0/stock_embeddings/embedders.py +214 -0
- stock_embeddings-0.1.0/stock_embeddings/graph.py +290 -0
- stock_embeddings-0.1.0/stock_embeddings.egg-info/PKG-INFO +177 -0
- stock_embeddings-0.1.0/stock_embeddings.egg-info/SOURCES.txt +15 -0
- stock_embeddings-0.1.0/stock_embeddings.egg-info/dependency_links.txt +1 -0
- stock_embeddings-0.1.0/stock_embeddings.egg-info/requires.txt +15 -0
- stock_embeddings-0.1.0/stock_embeddings.egg-info/top_level.txt +1 -0
- stock_embeddings-0.1.0/tests/test_stock_embeddings.py +95 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Tomasz Rymaszeski
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,177 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: stock-embeddings
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Train stock embeddings (Node2Vec, spectral, NetMF, raw baselines) from price data
|
|
5
|
+
Author: Tomasz Rymaszeski
|
|
6
|
+
License: MIT License
|
|
7
|
+
|
|
8
|
+
Copyright (c) 2026 Tomasz Rymaszeski
|
|
9
|
+
|
|
10
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
11
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
12
|
+
in the Software without restriction, including without limitation the rights
|
|
13
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
14
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
15
|
+
furnished to do so, subject to the following conditions:
|
|
16
|
+
|
|
17
|
+
The above copyright notice and this permission notice shall be included in all
|
|
18
|
+
copies or substantial portions of the Software.
|
|
19
|
+
|
|
20
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
21
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
22
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
23
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
24
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
25
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
26
|
+
SOFTWARE.
|
|
27
|
+
|
|
28
|
+
Project-URL: Homepage, https://github.com/tomaszrymaszewski/stock_embeddings
|
|
29
|
+
Project-URL: Documentation, https://docs.stockscore.live/
|
|
30
|
+
Project-URL: Repository, https://github.com/tomaszrymaszewski/stock_embeddings
|
|
31
|
+
Project-URL: Changelog, https://github.com/tomaszrymaszewski/stock_embeddings/blob/main/CHANGELOG.md
|
|
32
|
+
Project-URL: Paper, https://dx.doi.org/10.2139/ssrn.7591058
|
|
33
|
+
Keywords: stocks,embeddings,node2vec,graph,finance,correlation,quantitative-finance,representation-learning
|
|
34
|
+
Classifier: Development Status :: 3 - Alpha
|
|
35
|
+
Classifier: Intended Audience :: Financial and Insurance Industry
|
|
36
|
+
Classifier: Intended Audience :: Science/Research
|
|
37
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
38
|
+
Classifier: Programming Language :: Python :: 3
|
|
39
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
40
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
41
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
42
|
+
Classifier: Topic :: Office/Business :: Financial :: Investment
|
|
43
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
44
|
+
Requires-Python: >=3.10
|
|
45
|
+
Description-Content-Type: text/markdown
|
|
46
|
+
License-File: LICENSE
|
|
47
|
+
Requires-Dist: numpy>=1.26
|
|
48
|
+
Requires-Dist: pandas>=2.0
|
|
49
|
+
Requires-Dist: networkx>=3.1
|
|
50
|
+
Requires-Dist: scikit-learn>=1.3
|
|
51
|
+
Requires-Dist: scipy>=1.11
|
|
52
|
+
Requires-Dist: gensim>=4.3
|
|
53
|
+
Requires-Dist: node2vec>=0.5.0
|
|
54
|
+
Provides-Extra: dev
|
|
55
|
+
Requires-Dist: pytest>=7; extra == "dev"
|
|
56
|
+
Requires-Dist: build; extra == "dev"
|
|
57
|
+
Requires-Dist: twine; extra == "dev"
|
|
58
|
+
Provides-Extra: docs
|
|
59
|
+
Requires-Dist: mkdocs-material>=9; extra == "docs"
|
|
60
|
+
Dynamic: license-file
|
|
61
|
+
|
|
62
|
+
# stock-embeddings
|
|
63
|
+
|
|
64
|
+
Train vector embeddings of stocks from their price history. Give the library a dict of
|
|
65
|
+
`{ticker: prices}`, and it builds a stock-similarity graph from return correlations and
|
|
66
|
+
embeds every stock into a dense vector — stocks that move together (and belong to similar
|
|
67
|
+
businesses) end up close in embedding space.
|
|
68
|
+
|
|
69
|
+
```python
|
|
70
|
+
from stock_embeddings import train_embeddings
|
|
71
|
+
|
|
72
|
+
data = {"AAPL": aapl_prices, "MSFT": msft_prices, "XOM": xom_prices, ...}
|
|
73
|
+
embeddings = train_embeddings(data) # -> {"AAPL": [64 floats], "MSFT": [...], ...}
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
**Full documentation: [docs.stockscore.live](https://docs.stockscore.live/)**
|
|
77
|
+
|
|
78
|
+
## Summary
|
|
79
|
+
|
|
80
|
+
`stock-embeddings` turns raw price panels into a `{ticker: vector}` dictionary using the
|
|
81
|
+
pipeline from our research on correlation-graph stock embeddings:
|
|
82
|
+
|
|
83
|
+
```
|
|
84
|
+
prices -> returns -> correlation matrix -> graph -> embedding -> {ticker: vector}
|
|
85
|
+
```
|
|
86
|
+
|
|
87
|
+
Every stage is configurable with sane defaults: co-movement removal (market-neutral
|
|
88
|
+
residuals, cross-sectional demeaning, spectral/RMT cleaning), graph construction
|
|
89
|
+
(threshold + Forced Connectivity repair, MST, kNN), and the embedding family itself
|
|
90
|
+
(Node2Vec, Laplacian Eigenmaps, NetMF, or raw-feature baselines).
|
|
91
|
+
|
|
92
|
+
## Why this is useful
|
|
93
|
+
|
|
94
|
+
- **Semantic stock similarity.** Embedding-space neighbours capture economic similarity
|
|
95
|
+
(sector, industry, business model) that raw correlation misses, because the graph +
|
|
96
|
+
random walks propagate relationships beyond direct pairwise correlation.
|
|
97
|
+
- **Ready-made features.** The vectors are a compact, dense feature representation for
|
|
98
|
+
downstream ML: sector classification, clustering, peer selection, retrieval, or risk
|
|
99
|
+
models.
|
|
100
|
+
- **Turnkey.** One function call, no data plumbing, no graph code, no tuning required —
|
|
101
|
+
defaults come from a seed-averaged, significance-tested research benchmark.
|
|
102
|
+
- **Honest and reproducible.** All randomness is seeded; alternative embedding families
|
|
103
|
+
and no-embedding baselines are built in, so you can verify the embedding actually adds
|
|
104
|
+
value on your data.
|
|
105
|
+
|
|
106
|
+
## Setup
|
|
107
|
+
|
|
108
|
+
Requires Python 3.10+.
|
|
109
|
+
|
|
110
|
+
```bash
|
|
111
|
+
pip install stock-embeddings
|
|
112
|
+
```
|
|
113
|
+
|
|
114
|
+
Or from source:
|
|
115
|
+
|
|
116
|
+
```bash
|
|
117
|
+
git clone https://github.com/tomaszrymaszewski/stock_embeddings.git
|
|
118
|
+
cd stock_embeddings
|
|
119
|
+
pip install .
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
## Quick start
|
|
123
|
+
|
|
124
|
+
```python
|
|
125
|
+
import pandas as pd
|
|
126
|
+
from stock_embeddings import train_embeddings
|
|
127
|
+
|
|
128
|
+
# Any dict of ticker -> prices works: pandas Series (optionally with a
|
|
129
|
+
# DatetimeIndex), lists, or numpy arrays.
|
|
130
|
+
prices = pd.read_csv("prices.csv", index_col=0, parse_dates=True)
|
|
131
|
+
data = {ticker: prices[ticker].dropna() for ticker in prices.columns}
|
|
132
|
+
|
|
133
|
+
embeddings = train_embeddings(data)
|
|
134
|
+
|
|
135
|
+
# Nearest neighbours of a stock:
|
|
136
|
+
from stock_embeddings import fit
|
|
137
|
+
result = fit(data)
|
|
138
|
+
result.most_similar("AAPL", k=10) # [("MSFT", 0.98), ("AMZN", 0.97), ...]
|
|
139
|
+
result.to_dataframe() # embeddings as a DataFrame
|
|
140
|
+
```
|
|
141
|
+
|
|
142
|
+
## Usage
|
|
143
|
+
|
|
144
|
+
Every parameter has a default; pass only what you want to change:
|
|
145
|
+
|
|
146
|
+
```python
|
|
147
|
+
embeddings = train_embeddings(
|
|
148
|
+
data,
|
|
149
|
+
dimensions=32, # embedding size (default 64)
|
|
150
|
+
embedder="spectral", # node2vec | spectral | netmf | raw_corr | raw_returns
|
|
151
|
+
graph_filtering="knn", # threshold | mst | knn | signed_knn
|
|
152
|
+
knn_k=10,
|
|
153
|
+
market_neutral=True, # regress out the market factor before building the graph
|
|
154
|
+
spectral_clean="market", # remove the market eigenmode from the correlation matrix
|
|
155
|
+
seed=42,
|
|
156
|
+
)
|
|
157
|
+
```
|
|
158
|
+
|
|
159
|
+
Use `fit()` instead of `train_embeddings()` to get the full `TrainResult` — embeddings,
|
|
160
|
+
the returns matrix actually used, the similarity graph, and the resolved config — plus
|
|
161
|
+
helpers like `to_dataframe()` and `most_similar()`.
|
|
162
|
+
|
|
163
|
+
See the [documentation](https://docs.stockscore.live/) for the complete parameter
|
|
164
|
+
reference and guides.
|
|
165
|
+
|
|
166
|
+
## Research
|
|
167
|
+
|
|
168
|
+
This library implements the method developed in our paper
|
|
169
|
+
[*Stock Embeddings from Return-Correlation Graphs*](https://dx.doi.org/10.2139/ssrn.7591058),
|
|
170
|
+
which benchmarks how well correlation-graph embeddings recover the GICS classification of
|
|
171
|
+
stocks against honest baselines (raw correlation, MST, kNN, PMFG), with seed-averaged,
|
|
172
|
+
significance-tested results. Read the paper for how we got here — the method choices,
|
|
173
|
+
the defaults shipped in this library, and what the embeddings can and cannot do.
|
|
174
|
+
|
|
175
|
+
## License
|
|
176
|
+
|
|
177
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
# stock-embeddings
|
|
2
|
+
|
|
3
|
+
Train vector embeddings of stocks from their price history. Give the library a dict of
|
|
4
|
+
`{ticker: prices}`, and it builds a stock-similarity graph from return correlations and
|
|
5
|
+
embeds every stock into a dense vector — stocks that move together (and belong to similar
|
|
6
|
+
businesses) end up close in embedding space.
|
|
7
|
+
|
|
8
|
+
```python
|
|
9
|
+
from stock_embeddings import train_embeddings
|
|
10
|
+
|
|
11
|
+
data = {"AAPL": aapl_prices, "MSFT": msft_prices, "XOM": xom_prices, ...}
|
|
12
|
+
embeddings = train_embeddings(data) # -> {"AAPL": [64 floats], "MSFT": [...], ...}
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
**Full documentation: [docs.stockscore.live](https://docs.stockscore.live/)**
|
|
16
|
+
|
|
17
|
+
## Summary
|
|
18
|
+
|
|
19
|
+
`stock-embeddings` turns raw price panels into a `{ticker: vector}` dictionary using the
|
|
20
|
+
pipeline from our research on correlation-graph stock embeddings:
|
|
21
|
+
|
|
22
|
+
```
|
|
23
|
+
prices -> returns -> correlation matrix -> graph -> embedding -> {ticker: vector}
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
Every stage is configurable with sane defaults: co-movement removal (market-neutral
|
|
27
|
+
residuals, cross-sectional demeaning, spectral/RMT cleaning), graph construction
|
|
28
|
+
(threshold + Forced Connectivity repair, MST, kNN), and the embedding family itself
|
|
29
|
+
(Node2Vec, Laplacian Eigenmaps, NetMF, or raw-feature baselines).
|
|
30
|
+
|
|
31
|
+
## Why this is useful
|
|
32
|
+
|
|
33
|
+
- **Semantic stock similarity.** Embedding-space neighbours capture economic similarity
|
|
34
|
+
(sector, industry, business model) that raw correlation misses, because the graph +
|
|
35
|
+
random walks propagate relationships beyond direct pairwise correlation.
|
|
36
|
+
- **Ready-made features.** The vectors are a compact, dense feature representation for
|
|
37
|
+
downstream ML: sector classification, clustering, peer selection, retrieval, or risk
|
|
38
|
+
models.
|
|
39
|
+
- **Turnkey.** One function call, no data plumbing, no graph code, no tuning required —
|
|
40
|
+
defaults come from a seed-averaged, significance-tested research benchmark.
|
|
41
|
+
- **Honest and reproducible.** All randomness is seeded; alternative embedding families
|
|
42
|
+
and no-embedding baselines are built in, so you can verify the embedding actually adds
|
|
43
|
+
value on your data.
|
|
44
|
+
|
|
45
|
+
## Setup
|
|
46
|
+
|
|
47
|
+
Requires Python 3.10+.
|
|
48
|
+
|
|
49
|
+
```bash
|
|
50
|
+
pip install stock-embeddings
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
Or from source:
|
|
54
|
+
|
|
55
|
+
```bash
|
|
56
|
+
git clone https://github.com/tomaszrymaszewski/stock_embeddings.git
|
|
57
|
+
cd stock_embeddings
|
|
58
|
+
pip install .
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
## Quick start
|
|
62
|
+
|
|
63
|
+
```python
|
|
64
|
+
import pandas as pd
|
|
65
|
+
from stock_embeddings import train_embeddings
|
|
66
|
+
|
|
67
|
+
# Any dict of ticker -> prices works: pandas Series (optionally with a
|
|
68
|
+
# DatetimeIndex), lists, or numpy arrays.
|
|
69
|
+
prices = pd.read_csv("prices.csv", index_col=0, parse_dates=True)
|
|
70
|
+
data = {ticker: prices[ticker].dropna() for ticker in prices.columns}
|
|
71
|
+
|
|
72
|
+
embeddings = train_embeddings(data)
|
|
73
|
+
|
|
74
|
+
# Nearest neighbours of a stock:
|
|
75
|
+
from stock_embeddings import fit
|
|
76
|
+
result = fit(data)
|
|
77
|
+
result.most_similar("AAPL", k=10) # [("MSFT", 0.98), ("AMZN", 0.97), ...]
|
|
78
|
+
result.to_dataframe() # embeddings as a DataFrame
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
## Usage
|
|
82
|
+
|
|
83
|
+
Every parameter has a default; pass only what you want to change:
|
|
84
|
+
|
|
85
|
+
```python
|
|
86
|
+
embeddings = train_embeddings(
|
|
87
|
+
data,
|
|
88
|
+
dimensions=32, # embedding size (default 64)
|
|
89
|
+
embedder="spectral", # node2vec | spectral | netmf | raw_corr | raw_returns
|
|
90
|
+
graph_filtering="knn", # threshold | mst | knn | signed_knn
|
|
91
|
+
knn_k=10,
|
|
92
|
+
market_neutral=True, # regress out the market factor before building the graph
|
|
93
|
+
spectral_clean="market", # remove the market eigenmode from the correlation matrix
|
|
94
|
+
seed=42,
|
|
95
|
+
)
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
Use `fit()` instead of `train_embeddings()` to get the full `TrainResult` — embeddings,
|
|
99
|
+
the returns matrix actually used, the similarity graph, and the resolved config — plus
|
|
100
|
+
helpers like `to_dataframe()` and `most_similar()`.
|
|
101
|
+
|
|
102
|
+
See the [documentation](https://docs.stockscore.live/) for the complete parameter
|
|
103
|
+
reference and guides.
|
|
104
|
+
|
|
105
|
+
## Research
|
|
106
|
+
|
|
107
|
+
This library implements the method developed in our paper
|
|
108
|
+
[*Stock Embeddings from Return-Correlation Graphs*](https://dx.doi.org/10.2139/ssrn.7591058),
|
|
109
|
+
which benchmarks how well correlation-graph embeddings recover the GICS classification of
|
|
110
|
+
stocks against honest baselines (raw correlation, MST, kNN, PMFG), with seed-averaged,
|
|
111
|
+
significance-tested results. Read the paper for how we got here — the method choices,
|
|
112
|
+
the defaults shipped in this library, and what the embeddings can and cannot do.
|
|
113
|
+
|
|
114
|
+
## License
|
|
115
|
+
|
|
116
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,51 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "stock-embeddings"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Train stock embeddings (Node2Vec, spectral, NetMF, raw baselines) from price data"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = { file = "LICENSE" }
|
|
12
|
+
authors = [{ name = "Tomasz Rymaszeski" }]
|
|
13
|
+
keywords = [
|
|
14
|
+
"stocks", "embeddings", "node2vec", "graph", "finance",
|
|
15
|
+
"correlation", "quantitative-finance", "representation-learning",
|
|
16
|
+
]
|
|
17
|
+
classifiers = [
|
|
18
|
+
"Development Status :: 3 - Alpha",
|
|
19
|
+
"Intended Audience :: Financial and Insurance Industry",
|
|
20
|
+
"Intended Audience :: Science/Research",
|
|
21
|
+
"License :: OSI Approved :: MIT License",
|
|
22
|
+
"Programming Language :: Python :: 3",
|
|
23
|
+
"Programming Language :: Python :: 3.10",
|
|
24
|
+
"Programming Language :: Python :: 3.11",
|
|
25
|
+
"Programming Language :: Python :: 3.12",
|
|
26
|
+
"Topic :: Office/Business :: Financial :: Investment",
|
|
27
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
28
|
+
]
|
|
29
|
+
dependencies = [
|
|
30
|
+
"numpy>=1.26",
|
|
31
|
+
"pandas>=2.0",
|
|
32
|
+
"networkx>=3.1",
|
|
33
|
+
"scikit-learn>=1.3",
|
|
34
|
+
"scipy>=1.11",
|
|
35
|
+
"gensim>=4.3",
|
|
36
|
+
"node2vec>=0.5.0",
|
|
37
|
+
]
|
|
38
|
+
|
|
39
|
+
[project.optional-dependencies]
|
|
40
|
+
dev = ["pytest>=7", "build", "twine"]
|
|
41
|
+
docs = ["mkdocs-material>=9"]
|
|
42
|
+
|
|
43
|
+
[project.urls]
|
|
44
|
+
Homepage = "https://github.com/tomaszrymaszewski/stock_embeddings"
|
|
45
|
+
Documentation = "https://docs.stockscore.live/"
|
|
46
|
+
Repository = "https://github.com/tomaszrymaszewski/stock_embeddings"
|
|
47
|
+
Changelog = "https://github.com/tomaszrymaszewski/stock_embeddings/blob/main/CHANGELOG.md"
|
|
48
|
+
Paper = "https://dx.doi.org/10.2139/ssrn.7591058"
|
|
49
|
+
|
|
50
|
+
[tool.setuptools.packages.find]
|
|
51
|
+
include = ["stock_embeddings*"]
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
"""stock_embeddings -- train stock embeddings from price data.
|
|
2
|
+
|
|
3
|
+
Basic usage:
|
|
4
|
+
from stock_embeddings import train_embeddings
|
|
5
|
+
|
|
6
|
+
data = {"AAPL": aapl_price_series, "MSFT": msft_price_series, ...}
|
|
7
|
+
embeddings = train_embeddings(data) # all defaults
|
|
8
|
+
embeddings = train_embeddings(data, dimensions=32, # any Config field
|
|
9
|
+
graph_filtering="knn",
|
|
10
|
+
embedder="spectral")
|
|
11
|
+
|
|
12
|
+
Use `fit()` instead of `train_embeddings()` to also get the returns matrix,
|
|
13
|
+
the similarity graph, and helper methods (to_dataframe, most_similar).
|
|
14
|
+
"""
|
|
15
|
+
|
|
16
|
+
from .api import TrainResult, fit, train_embeddings
|
|
17
|
+
from .config import Config
|
|
18
|
+
|
|
19
|
+
__version__ = "0.1.0"
|
|
20
|
+
|
|
21
|
+
__all__ = ["train_embeddings", "fit", "Config", "TrainResult", "__version__"]
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
"""Public API: train stock embeddings from a dict of price series."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import logging
|
|
6
|
+
from dataclasses import dataclass, fields
|
|
7
|
+
from typing import Mapping, Sequence, Union
|
|
8
|
+
|
|
9
|
+
import networkx as nx
|
|
10
|
+
import numpy as np
|
|
11
|
+
import pandas as pd
|
|
12
|
+
|
|
13
|
+
from .config import Config
|
|
14
|
+
from .data import PriceSeries, prices_to_returns
|
|
15
|
+
from .embedders import RAW_EMBEDDERS, build_embeddings
|
|
16
|
+
from .graph import GraphBuilder
|
|
17
|
+
|
|
18
|
+
logger = logging.getLogger(__name__)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
@dataclass
|
|
22
|
+
class TrainResult:
|
|
23
|
+
"""Outcome of a training run.
|
|
24
|
+
|
|
25
|
+
Attributes:
|
|
26
|
+
embeddings: {ticker: vector} -- the primary output.
|
|
27
|
+
returns: the aligned returns DataFrame actually used.
|
|
28
|
+
graph: the stock similarity graph (None for raw_* embedders).
|
|
29
|
+
config: the resolved Config used for this run.
|
|
30
|
+
"""
|
|
31
|
+
|
|
32
|
+
embeddings: dict[str, list[float]]
|
|
33
|
+
returns: pd.DataFrame
|
|
34
|
+
graph: nx.Graph | None
|
|
35
|
+
config: Config
|
|
36
|
+
|
|
37
|
+
def to_dataframe(self) -> pd.DataFrame:
|
|
38
|
+
"""Embeddings as a DataFrame indexed by ticker (one column per dimension)."""
|
|
39
|
+
return pd.DataFrame.from_dict(self.embeddings, orient="index")
|
|
40
|
+
|
|
41
|
+
def most_similar(self, ticker: str, k: int = 10) -> list[tuple[str, float]]:
|
|
42
|
+
"""Top-k most cosine-similar tickers to `ticker` (excluding itself)."""
|
|
43
|
+
tickers = list(self.embeddings.keys())
|
|
44
|
+
if ticker not in tickers:
|
|
45
|
+
raise KeyError(f"{ticker!r} not in trained embeddings")
|
|
46
|
+
M = np.array([self.embeddings[t] for t in tickers], dtype=float)
|
|
47
|
+
norms = np.linalg.norm(M, axis=1, keepdims=True)
|
|
48
|
+
norms[norms == 0] = 1e-10
|
|
49
|
+
M = M / norms
|
|
50
|
+
sims = M @ M[tickers.index(ticker)]
|
|
51
|
+
order = np.argsort(sims)[::-1]
|
|
52
|
+
out = []
|
|
53
|
+
for i in order:
|
|
54
|
+
if tickers[i] == ticker:
|
|
55
|
+
continue
|
|
56
|
+
out.append((tickers[i], float(sims[i])))
|
|
57
|
+
if len(out) >= k:
|
|
58
|
+
break
|
|
59
|
+
return out
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def train_embeddings(
|
|
63
|
+
data: Mapping[str, PriceSeries],
|
|
64
|
+
config: Config | None = None,
|
|
65
|
+
**params,
|
|
66
|
+
) -> dict[str, list[float]]:
|
|
67
|
+
"""Train stock embeddings from a dict of {ticker: price series}.
|
|
68
|
+
|
|
69
|
+
Args:
|
|
70
|
+
data: mapping of ticker -> prices (pandas Series, list, or numpy
|
|
71
|
+
array). A DatetimeIndex is used for alignment if present.
|
|
72
|
+
config: optional pre-built Config; keyword arguments override it.
|
|
73
|
+
**params: any Config field, e.g. ``dimensions=32``,
|
|
74
|
+
``graph_filtering="knn"``, ``embedder="spectral"``.
|
|
75
|
+
|
|
76
|
+
Returns:
|
|
77
|
+
{ticker: embedding vector} dict.
|
|
78
|
+
|
|
79
|
+
Example:
|
|
80
|
+
>>> embeddings = train_embeddings({"AAPL": aapl_prices, "MSFT": msft_prices})
|
|
81
|
+
>>> embeddings = train_embeddings(data, dimensions=32, embedder="spectral")
|
|
82
|
+
"""
|
|
83
|
+
return fit(data, config=config, **params).embeddings
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def fit(
|
|
87
|
+
data: Mapping[str, PriceSeries],
|
|
88
|
+
config: Config | None = None,
|
|
89
|
+
**params,
|
|
90
|
+
) -> TrainResult:
|
|
91
|
+
"""Like :func:`train_embeddings` but returns the full :class:`TrainResult`
|
|
92
|
+
(embeddings + returns + graph + resolved config)."""
|
|
93
|
+
config = _resolve_config(config, params)
|
|
94
|
+
_setup_logging(config.verbose)
|
|
95
|
+
|
|
96
|
+
returns = prices_to_returns(data, config)
|
|
97
|
+
|
|
98
|
+
graph = None
|
|
99
|
+
if config.embedder not in RAW_EMBEDDERS:
|
|
100
|
+
graph = GraphBuilder(config).build(returns)
|
|
101
|
+
|
|
102
|
+
embeddings = build_embeddings(config, returns, graph=graph)
|
|
103
|
+
logger.info("Trained %d embeddings (embedder=%s)", len(embeddings),
|
|
104
|
+
config.embedder)
|
|
105
|
+
return TrainResult(embeddings=embeddings, returns=returns, graph=graph,
|
|
106
|
+
config=config)
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def _resolve_config(config: Config | None, params: dict) -> Config:
|
|
110
|
+
if config is None:
|
|
111
|
+
return Config(**params)
|
|
112
|
+
if not params:
|
|
113
|
+
return config
|
|
114
|
+
known = {f.name for f in fields(Config)}
|
|
115
|
+
unknown = set(params) - known
|
|
116
|
+
if unknown:
|
|
117
|
+
raise TypeError(f"Unknown Config parameter(s): {sorted(unknown)}")
|
|
118
|
+
overrides = {f.name: getattr(config, f.name) for f in fields(Config)}
|
|
119
|
+
overrides.update(params)
|
|
120
|
+
return Config(**overrides)
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
def _setup_logging(verbose: bool) -> None:
|
|
124
|
+
level = logging.INFO if verbose else logging.WARNING
|
|
125
|
+
logger.setLevel(level)
|
|
126
|
+
pkg_logger = logging.getLogger("stock_embeddings")
|
|
127
|
+
pkg_logger.setLevel(level)
|
|
128
|
+
if verbose and not pkg_logger.handlers:
|
|
129
|
+
handler = logging.StreamHandler()
|
|
130
|
+
handler.setFormatter(logging.Formatter("%(message)s"))
|
|
131
|
+
pkg_logger.addHandler(handler)
|
|
132
|
+
pkg_logger.propagate = False
|
|
@@ -0,0 +1,97 @@
|
|
|
1
|
+
"""Configuration for stock embedding training. All parameters have sane defaults."""
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@dataclass
|
|
7
|
+
class Config:
|
|
8
|
+
"""Central control panel for embedding training.
|
|
9
|
+
|
|
10
|
+
Every field can be passed as a keyword argument to
|
|
11
|
+
:func:`stock_embeddings.train_embeddings`.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
# -------------------------------------------------------------------------
|
|
15
|
+
# Embedding family
|
|
16
|
+
# -------------------------------------------------------------------------
|
|
17
|
+
# "node2vec" -> local random walks + skip-gram (default)
|
|
18
|
+
# "spectral" -> Laplacian Eigenmaps (global spectral geometry)
|
|
19
|
+
# "netmf" -> DeepWalk-as-matrix-factorization (closed-form skip-gram)
|
|
20
|
+
# "raw_corr" -> no embedding: a stock IS its correlation row
|
|
21
|
+
# "raw_returns" -> no embedding: a stock IS its standardized return series
|
|
22
|
+
embedder: str = "node2vec"
|
|
23
|
+
|
|
24
|
+
# Embedding dimensionality (ignored by raw_* embedders, which are data-shaped).
|
|
25
|
+
dimensions: int = 64
|
|
26
|
+
|
|
27
|
+
# Random seed for walks / Word2Vec init / eigensolvers.
|
|
28
|
+
seed: int = 42
|
|
29
|
+
|
|
30
|
+
# -------------------------------------------------------------------------
|
|
31
|
+
# Graph construction
|
|
32
|
+
# -------------------------------------------------------------------------
|
|
33
|
+
# 'threshold' = connect stocks if |correlation| > correlation_threshold
|
|
34
|
+
# 'mst' = maximum spanning tree backbone (+ optional extra edges)
|
|
35
|
+
# 'knn' = connect each stock to its K strongest-correlated peers (|corr|)
|
|
36
|
+
# 'signed_knn' = kNN ranked by SIGNED correlation (anti-correlated peers excluded)
|
|
37
|
+
graph_filtering: str = "threshold"
|
|
38
|
+
|
|
39
|
+
# Only used if graph_filtering == 'threshold'
|
|
40
|
+
correlation_threshold: float = 0.65
|
|
41
|
+
|
|
42
|
+
# Forced Connectivity: repair under-connected nodes to a minimum degree.
|
|
43
|
+
repair_islands: bool = True
|
|
44
|
+
min_degree: int = 3
|
|
45
|
+
|
|
46
|
+
# Only used if graph_filtering in ('knn', 'signed_knn')
|
|
47
|
+
knn_k: int = 10
|
|
48
|
+
|
|
49
|
+
# Only used if graph_filtering == 'mst'
|
|
50
|
+
mst_extra_edges_pct: float = 0.05 # 0.0 = pure MST; 0.05 = add top 5% unused edges
|
|
51
|
+
mst_squared_corr: bool = True # use corr^2 as edge weight
|
|
52
|
+
# Mantegna (1999) distance-transform MST on SIGNED correlation; overrides the
|
|
53
|
+
# mst_squared_corr / abs weighting above.
|
|
54
|
+
mst_distance_transform: bool = False
|
|
55
|
+
|
|
56
|
+
# -------------------------------------------------------------------------
|
|
57
|
+
# Node2Vec walk parameters
|
|
58
|
+
# -------------------------------------------------------------------------
|
|
59
|
+
walk_p: float = 0.8 # return parameter
|
|
60
|
+
walk_q: float = 7.2 # in-out parameter
|
|
61
|
+
walk_length: int = 20
|
|
62
|
+
num_walks: int = 100
|
|
63
|
+
window: int = 10 # skip-gram window (also NetMF window)
|
|
64
|
+
|
|
65
|
+
# -------------------------------------------------------------------------
|
|
66
|
+
# Similarity matrix estimation
|
|
67
|
+
# -------------------------------------------------------------------------
|
|
68
|
+
# True = Graphical Lasso (partial correlation); False = Pearson correlation.
|
|
69
|
+
use_precision_matrix: bool = False
|
|
70
|
+
# lasso_alpha = 0.0 -> GraphicalLassoCV (slow CV); > 0.0 -> fixed-alpha (fast).
|
|
71
|
+
# With lasso, use a SMALL correlation_threshold (~0.02-0.1).
|
|
72
|
+
lasso_alpha: float = 0.0
|
|
73
|
+
|
|
74
|
+
# Market-neutral residualization: regress out a market factor and keep residuals.
|
|
75
|
+
# Factor = `market_ticker` column if present, else equal-weighted average.
|
|
76
|
+
market_neutral: bool = False
|
|
77
|
+
market_ticker: str = "SPY"
|
|
78
|
+
|
|
79
|
+
# Cross-sectional demeaning: subtract each day's equal-weighted mean return.
|
|
80
|
+
cross_sectional_demean: bool = False
|
|
81
|
+
|
|
82
|
+
# Spectral (RMT) cleaning of the correlation matrix:
|
|
83
|
+
# "none" | "market" (drop top eigenmode) | "topk" (drop top spectral_k)
|
|
84
|
+
# | "rmt" (keep only modes above the Marchenko-Pastur edge)
|
|
85
|
+
spectral_clean: str = "none"
|
|
86
|
+
spectral_k: int = 1
|
|
87
|
+
|
|
88
|
+
# -------------------------------------------------------------------------
|
|
89
|
+
# Data preparation
|
|
90
|
+
# -------------------------------------------------------------------------
|
|
91
|
+
# Drop tickers with less than this fraction of non-NaN observations.
|
|
92
|
+
min_history_frac: float = 0.8
|
|
93
|
+
# Minimum number of return observations required to train.
|
|
94
|
+
min_rows: int = 30
|
|
95
|
+
|
|
96
|
+
# Print progress logs.
|
|
97
|
+
verbose: bool = True
|