discrete-dists 0.0.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.
- discrete_dists-0.0.0/.github/workflows/CI.yml +151 -0
- discrete_dists-0.0.0/.gitignore +76 -0
- discrete_dists-0.0.0/.pre-commit-config.yaml +13 -0
- discrete_dists-0.0.0/Cargo.lock +293 -0
- discrete_dists-0.0.0/Cargo.toml +17 -0
- discrete_dists-0.0.0/PKG-INFO +18 -0
- discrete_dists-0.0.0/discrete_dists/__init__.py +0 -0
- discrete_dists-0.0.0/discrete_dists/distribution.py +54 -0
- discrete_dists-0.0.0/discrete_dists/mixture.py +65 -0
- discrete_dists-0.0.0/discrete_dists/proportional.py +40 -0
- discrete_dists-0.0.0/discrete_dists/py.typed +0 -0
- discrete_dists-0.0.0/discrete_dists/uniform.py +27 -0
- discrete_dists-0.0.0/discrete_dists/utils/SumTree.py +49 -0
- discrete_dists-0.0.0/discrete_dists/utils/__init__.py +0 -0
- discrete_dists-0.0.0/discrete_dists/utils/npu.py +8 -0
- discrete_dists-0.0.0/pyproject.toml +56 -0
- discrete_dists-0.0.0/rust/lib.rs +10 -0
- discrete_dists-0.0.0/rust/utils/sumtree.rs +151 -0
- discrete_dists-0.0.0/rust/utils.rs +1 -0
- discrete_dists-0.0.0/tests/utils/__init__.py +0 -0
- discrete_dists-0.0.0/tests/utils/test_SumTree.py +101 -0
- discrete_dists-0.0.0/typings/discrete_dists/rust/__init__.pyi +15 -0
|
@@ -0,0 +1,151 @@
|
|
|
1
|
+
# This file is autogenerated by maturin v1.7.1
|
|
2
|
+
# To update, run
|
|
3
|
+
#
|
|
4
|
+
# maturin generate-ci github
|
|
5
|
+
#
|
|
6
|
+
name: CI
|
|
7
|
+
|
|
8
|
+
on:
|
|
9
|
+
push:
|
|
10
|
+
branches:
|
|
11
|
+
- main
|
|
12
|
+
tags:
|
|
13
|
+
- '*'
|
|
14
|
+
pull_request:
|
|
15
|
+
workflow_dispatch:
|
|
16
|
+
|
|
17
|
+
permissions:
|
|
18
|
+
contents: read
|
|
19
|
+
|
|
20
|
+
jobs:
|
|
21
|
+
linux:
|
|
22
|
+
runs-on: ${{ matrix.platform.runner }}
|
|
23
|
+
strategy:
|
|
24
|
+
matrix:
|
|
25
|
+
platform:
|
|
26
|
+
- runner: ubuntu-latest
|
|
27
|
+
target: x86_64
|
|
28
|
+
- runner: ubuntu-latest
|
|
29
|
+
target: x86
|
|
30
|
+
steps:
|
|
31
|
+
- uses: actions/checkout@v4
|
|
32
|
+
- uses: actions/setup-python@v5
|
|
33
|
+
with:
|
|
34
|
+
python-version: 3.x
|
|
35
|
+
- name: Build wheels
|
|
36
|
+
uses: PyO3/maturin-action@v1
|
|
37
|
+
with:
|
|
38
|
+
target: ${{ matrix.platform.target }}
|
|
39
|
+
args: --release --out dist --find-interpreter
|
|
40
|
+
sccache: 'true'
|
|
41
|
+
manylinux: auto
|
|
42
|
+
- name: Upload wheels
|
|
43
|
+
uses: actions/upload-artifact@v4
|
|
44
|
+
with:
|
|
45
|
+
name: wheels-linux-${{ matrix.platform.target }}
|
|
46
|
+
path: dist
|
|
47
|
+
|
|
48
|
+
windows:
|
|
49
|
+
runs-on: ${{ matrix.platform.runner }}
|
|
50
|
+
strategy:
|
|
51
|
+
matrix:
|
|
52
|
+
platform:
|
|
53
|
+
- runner: windows-latest
|
|
54
|
+
target: x64
|
|
55
|
+
- runner: windows-latest
|
|
56
|
+
target: x86
|
|
57
|
+
steps:
|
|
58
|
+
- uses: actions/checkout@v4
|
|
59
|
+
- uses: actions/setup-python@v5
|
|
60
|
+
with:
|
|
61
|
+
python-version: 3.x
|
|
62
|
+
architecture: ${{ matrix.platform.target }}
|
|
63
|
+
- name: Build wheels
|
|
64
|
+
uses: PyO3/maturin-action@v1
|
|
65
|
+
with:
|
|
66
|
+
target: ${{ matrix.platform.target }}
|
|
67
|
+
args: --release --out dist --find-interpreter
|
|
68
|
+
sccache: 'true'
|
|
69
|
+
- name: Upload wheels
|
|
70
|
+
uses: actions/upload-artifact@v4
|
|
71
|
+
with:
|
|
72
|
+
name: wheels-windows-${{ matrix.platform.target }}
|
|
73
|
+
path: dist
|
|
74
|
+
|
|
75
|
+
macos:
|
|
76
|
+
runs-on: ${{ matrix.platform.runner }}
|
|
77
|
+
strategy:
|
|
78
|
+
matrix:
|
|
79
|
+
platform:
|
|
80
|
+
- runner: macos-latest
|
|
81
|
+
target: aarch64
|
|
82
|
+
steps:
|
|
83
|
+
- uses: actions/checkout@v4
|
|
84
|
+
- uses: actions/setup-python@v5
|
|
85
|
+
with:
|
|
86
|
+
python-version: 3.x
|
|
87
|
+
- name: Build wheels
|
|
88
|
+
uses: PyO3/maturin-action@v1
|
|
89
|
+
with:
|
|
90
|
+
target: ${{ matrix.platform.target }}
|
|
91
|
+
args: --release --out dist --find-interpreter
|
|
92
|
+
sccache: 'true'
|
|
93
|
+
- name: Upload wheels
|
|
94
|
+
uses: actions/upload-artifact@v4
|
|
95
|
+
with:
|
|
96
|
+
name: wheels-macos-${{ matrix.platform.target }}
|
|
97
|
+
path: dist
|
|
98
|
+
|
|
99
|
+
sdist:
|
|
100
|
+
runs-on: ubuntu-latest
|
|
101
|
+
steps:
|
|
102
|
+
- uses: actions/checkout@v4
|
|
103
|
+
- name: Build sdist
|
|
104
|
+
uses: PyO3/maturin-action@v1
|
|
105
|
+
with:
|
|
106
|
+
command: sdist
|
|
107
|
+
args: --out dist
|
|
108
|
+
- name: Upload sdist
|
|
109
|
+
uses: actions/upload-artifact@v4
|
|
110
|
+
with:
|
|
111
|
+
name: wheels-sdist
|
|
112
|
+
path: dist
|
|
113
|
+
|
|
114
|
+
tag:
|
|
115
|
+
runs-on: ubuntu-latest
|
|
116
|
+
needs: [linux, windows, macos, sdist]
|
|
117
|
+
if: github.ref == 'refs/heads/main'
|
|
118
|
+
steps:
|
|
119
|
+
- name: Generate a token
|
|
120
|
+
id: generate-token
|
|
121
|
+
uses: actions/create-github-app-token@v1
|
|
122
|
+
with:
|
|
123
|
+
app-id: ${{ vars.APP_ID }}
|
|
124
|
+
private-key: ${{ secrets.APP_PRIVATE_KEY }}
|
|
125
|
+
|
|
126
|
+
- name: Checkout Code
|
|
127
|
+
uses: actions/checkout@v4
|
|
128
|
+
with:
|
|
129
|
+
fetch-depth: 0
|
|
130
|
+
token: ${{ steps.generate-token.outputs.token }}
|
|
131
|
+
|
|
132
|
+
- id: cz
|
|
133
|
+
name: Create bump commit, tag, and changelog
|
|
134
|
+
uses: commitizen-tools/commitizen-action@master
|
|
135
|
+
with:
|
|
136
|
+
github_token: ${{ steps.generate-token.outputs.token }}
|
|
137
|
+
actor: x-access-token
|
|
138
|
+
|
|
139
|
+
release:
|
|
140
|
+
name: Release
|
|
141
|
+
runs-on: ubuntu-latest
|
|
142
|
+
needs: [tag]
|
|
143
|
+
permissions:
|
|
144
|
+
id-token: write
|
|
145
|
+
steps:
|
|
146
|
+
- uses: actions/download-artifact@v4
|
|
147
|
+
- name: Publish to PyPI
|
|
148
|
+
uses: PyO3/maturin-action@v1
|
|
149
|
+
with:
|
|
150
|
+
command: upload
|
|
151
|
+
args: --non-interactive --skip-existing wheels-*/*
|
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
# package lock files
|
|
2
|
+
requirements.txt
|
|
3
|
+
Cargo.lock
|
|
4
|
+
|
|
5
|
+
/target
|
|
6
|
+
|
|
7
|
+
# Byte-compiled / optimized / DLL files
|
|
8
|
+
__pycache__/
|
|
9
|
+
.pytest_cache/
|
|
10
|
+
*.py[cod]
|
|
11
|
+
|
|
12
|
+
# C extensions
|
|
13
|
+
*.so
|
|
14
|
+
|
|
15
|
+
# Distribution / packaging
|
|
16
|
+
.Python
|
|
17
|
+
.venv/
|
|
18
|
+
env/
|
|
19
|
+
bin/
|
|
20
|
+
build/
|
|
21
|
+
develop-eggs/
|
|
22
|
+
dist/
|
|
23
|
+
eggs/
|
|
24
|
+
lib/
|
|
25
|
+
lib64/
|
|
26
|
+
parts/
|
|
27
|
+
sdist/
|
|
28
|
+
var/
|
|
29
|
+
include/
|
|
30
|
+
man/
|
|
31
|
+
venv/
|
|
32
|
+
*.egg-info/
|
|
33
|
+
.installed.cfg
|
|
34
|
+
*.egg
|
|
35
|
+
|
|
36
|
+
# Installer logs
|
|
37
|
+
pip-log.txt
|
|
38
|
+
pip-delete-this-directory.txt
|
|
39
|
+
pip-selfcheck.json
|
|
40
|
+
|
|
41
|
+
# Unit test / coverage reports
|
|
42
|
+
htmlcov/
|
|
43
|
+
.tox/
|
|
44
|
+
.coverage
|
|
45
|
+
.cache
|
|
46
|
+
nosetests.xml
|
|
47
|
+
coverage.xml
|
|
48
|
+
|
|
49
|
+
# Translations
|
|
50
|
+
*.mo
|
|
51
|
+
|
|
52
|
+
# Mr Developer
|
|
53
|
+
.mr.developer.cfg
|
|
54
|
+
.project
|
|
55
|
+
.pydevproject
|
|
56
|
+
|
|
57
|
+
# Rope
|
|
58
|
+
.ropeproject
|
|
59
|
+
|
|
60
|
+
# Django stuff:
|
|
61
|
+
*.log
|
|
62
|
+
*.pot
|
|
63
|
+
|
|
64
|
+
.DS_Store
|
|
65
|
+
|
|
66
|
+
# Sphinx documentation
|
|
67
|
+
docs/_build/
|
|
68
|
+
|
|
69
|
+
# PyCharm
|
|
70
|
+
.idea/
|
|
71
|
+
|
|
72
|
+
# VSCode
|
|
73
|
+
.vscode/
|
|
74
|
+
|
|
75
|
+
# Pyenv
|
|
76
|
+
.python-version
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
repos:
|
|
2
|
+
- repo: https://github.com/commitizen-tools/commitizen
|
|
3
|
+
rev: v2.21.2
|
|
4
|
+
hooks:
|
|
5
|
+
- id: commitizen
|
|
6
|
+
stages:
|
|
7
|
+
- commit-msg
|
|
8
|
+
- repo: https://github.com/astral-sh/ruff-pre-commit
|
|
9
|
+
# Ruff version.
|
|
10
|
+
rev: v0.2.2
|
|
11
|
+
hooks:
|
|
12
|
+
# Run the linter.
|
|
13
|
+
- id: ruff
|
|
@@ -0,0 +1,293 @@
|
|
|
1
|
+
# This file is automatically @generated by Cargo.
|
|
2
|
+
# It is not intended for manual editing.
|
|
3
|
+
version = 4
|
|
4
|
+
|
|
5
|
+
[[package]]
|
|
6
|
+
name = "autocfg"
|
|
7
|
+
version = "1.4.0"
|
|
8
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
9
|
+
checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26"
|
|
10
|
+
|
|
11
|
+
[[package]]
|
|
12
|
+
name = "bincode"
|
|
13
|
+
version = "1.3.3"
|
|
14
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
15
|
+
checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad"
|
|
16
|
+
dependencies = [
|
|
17
|
+
"serde",
|
|
18
|
+
]
|
|
19
|
+
|
|
20
|
+
[[package]]
|
|
21
|
+
name = "cfg-if"
|
|
22
|
+
version = "1.0.0"
|
|
23
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
24
|
+
checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
|
|
25
|
+
|
|
26
|
+
[[package]]
|
|
27
|
+
name = "discrete-dists"
|
|
28
|
+
version = "0.1.0"
|
|
29
|
+
dependencies = [
|
|
30
|
+
"bincode",
|
|
31
|
+
"ndarray",
|
|
32
|
+
"numpy",
|
|
33
|
+
"pyo3",
|
|
34
|
+
"serde",
|
|
35
|
+
]
|
|
36
|
+
|
|
37
|
+
[[package]]
|
|
38
|
+
name = "heck"
|
|
39
|
+
version = "0.5.0"
|
|
40
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
41
|
+
checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
|
|
42
|
+
|
|
43
|
+
[[package]]
|
|
44
|
+
name = "indoc"
|
|
45
|
+
version = "2.0.5"
|
|
46
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
47
|
+
checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5"
|
|
48
|
+
|
|
49
|
+
[[package]]
|
|
50
|
+
name = "libc"
|
|
51
|
+
version = "0.2.169"
|
|
52
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
53
|
+
checksum = "b5aba8db14291edd000dfcc4d620c7ebfb122c613afb886ca8803fa4e128a20a"
|
|
54
|
+
|
|
55
|
+
[[package]]
|
|
56
|
+
name = "matrixmultiply"
|
|
57
|
+
version = "0.3.9"
|
|
58
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
59
|
+
checksum = "9380b911e3e96d10c1f415da0876389aaf1b56759054eeb0de7df940c456ba1a"
|
|
60
|
+
dependencies = [
|
|
61
|
+
"autocfg",
|
|
62
|
+
"rawpointer",
|
|
63
|
+
]
|
|
64
|
+
|
|
65
|
+
[[package]]
|
|
66
|
+
name = "memoffset"
|
|
67
|
+
version = "0.9.1"
|
|
68
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
69
|
+
checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
|
|
70
|
+
dependencies = [
|
|
71
|
+
"autocfg",
|
|
72
|
+
]
|
|
73
|
+
|
|
74
|
+
[[package]]
|
|
75
|
+
name = "ndarray"
|
|
76
|
+
version = "0.16.1"
|
|
77
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
78
|
+
checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841"
|
|
79
|
+
dependencies = [
|
|
80
|
+
"matrixmultiply",
|
|
81
|
+
"num-complex",
|
|
82
|
+
"num-integer",
|
|
83
|
+
"num-traits",
|
|
84
|
+
"portable-atomic",
|
|
85
|
+
"portable-atomic-util",
|
|
86
|
+
"rawpointer",
|
|
87
|
+
"serde",
|
|
88
|
+
]
|
|
89
|
+
|
|
90
|
+
[[package]]
|
|
91
|
+
name = "num-complex"
|
|
92
|
+
version = "0.4.6"
|
|
93
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
94
|
+
checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
|
|
95
|
+
dependencies = [
|
|
96
|
+
"num-traits",
|
|
97
|
+
]
|
|
98
|
+
|
|
99
|
+
[[package]]
|
|
100
|
+
name = "num-integer"
|
|
101
|
+
version = "0.1.46"
|
|
102
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
103
|
+
checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
|
|
104
|
+
dependencies = [
|
|
105
|
+
"num-traits",
|
|
106
|
+
]
|
|
107
|
+
|
|
108
|
+
[[package]]
|
|
109
|
+
name = "num-traits"
|
|
110
|
+
version = "0.2.19"
|
|
111
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
112
|
+
checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
|
|
113
|
+
dependencies = [
|
|
114
|
+
"autocfg",
|
|
115
|
+
]
|
|
116
|
+
|
|
117
|
+
[[package]]
|
|
118
|
+
name = "numpy"
|
|
119
|
+
version = "0.23.0"
|
|
120
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
121
|
+
checksum = "b94caae805f998a07d33af06e6a3891e38556051b8045c615470a71590e13e78"
|
|
122
|
+
dependencies = [
|
|
123
|
+
"libc",
|
|
124
|
+
"ndarray",
|
|
125
|
+
"num-complex",
|
|
126
|
+
"num-integer",
|
|
127
|
+
"num-traits",
|
|
128
|
+
"pyo3",
|
|
129
|
+
"rustc-hash",
|
|
130
|
+
]
|
|
131
|
+
|
|
132
|
+
[[package]]
|
|
133
|
+
name = "once_cell"
|
|
134
|
+
version = "1.20.2"
|
|
135
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
136
|
+
checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775"
|
|
137
|
+
|
|
138
|
+
[[package]]
|
|
139
|
+
name = "portable-atomic"
|
|
140
|
+
version = "1.10.0"
|
|
141
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
142
|
+
checksum = "280dc24453071f1b63954171985a0b0d30058d287960968b9b2aca264c8d4ee6"
|
|
143
|
+
|
|
144
|
+
[[package]]
|
|
145
|
+
name = "portable-atomic-util"
|
|
146
|
+
version = "0.2.4"
|
|
147
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
148
|
+
checksum = "d8a2f0d8d040d7848a709caf78912debcc3f33ee4b3cac47d73d1e1069e83507"
|
|
149
|
+
dependencies = [
|
|
150
|
+
"portable-atomic",
|
|
151
|
+
]
|
|
152
|
+
|
|
153
|
+
[[package]]
|
|
154
|
+
name = "proc-macro2"
|
|
155
|
+
version = "1.0.93"
|
|
156
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
157
|
+
checksum = "60946a68e5f9d28b0dc1c21bb8a97ee7d018a8b322fa57838ba31cc878e22d99"
|
|
158
|
+
dependencies = [
|
|
159
|
+
"unicode-ident",
|
|
160
|
+
]
|
|
161
|
+
|
|
162
|
+
[[package]]
|
|
163
|
+
name = "pyo3"
|
|
164
|
+
version = "0.23.4"
|
|
165
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
166
|
+
checksum = "57fe09249128b3173d092de9523eaa75136bf7ba85e0d69eca241c7939c933cc"
|
|
167
|
+
dependencies = [
|
|
168
|
+
"cfg-if",
|
|
169
|
+
"indoc",
|
|
170
|
+
"libc",
|
|
171
|
+
"memoffset",
|
|
172
|
+
"once_cell",
|
|
173
|
+
"portable-atomic",
|
|
174
|
+
"pyo3-build-config",
|
|
175
|
+
"pyo3-ffi",
|
|
176
|
+
"pyo3-macros",
|
|
177
|
+
"unindent",
|
|
178
|
+
]
|
|
179
|
+
|
|
180
|
+
[[package]]
|
|
181
|
+
name = "pyo3-build-config"
|
|
182
|
+
version = "0.23.4"
|
|
183
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
184
|
+
checksum = "1cd3927b5a78757a0d71aa9dff669f903b1eb64b54142a9bd9f757f8fde65fd7"
|
|
185
|
+
dependencies = [
|
|
186
|
+
"once_cell",
|
|
187
|
+
"target-lexicon",
|
|
188
|
+
]
|
|
189
|
+
|
|
190
|
+
[[package]]
|
|
191
|
+
name = "pyo3-ffi"
|
|
192
|
+
version = "0.23.4"
|
|
193
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
194
|
+
checksum = "dab6bb2102bd8f991e7749f130a70d05dd557613e39ed2deeee8e9ca0c4d548d"
|
|
195
|
+
dependencies = [
|
|
196
|
+
"libc",
|
|
197
|
+
"pyo3-build-config",
|
|
198
|
+
]
|
|
199
|
+
|
|
200
|
+
[[package]]
|
|
201
|
+
name = "pyo3-macros"
|
|
202
|
+
version = "0.23.4"
|
|
203
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
204
|
+
checksum = "91871864b353fd5ffcb3f91f2f703a22a9797c91b9ab497b1acac7b07ae509c7"
|
|
205
|
+
dependencies = [
|
|
206
|
+
"proc-macro2",
|
|
207
|
+
"pyo3-macros-backend",
|
|
208
|
+
"quote",
|
|
209
|
+
"syn",
|
|
210
|
+
]
|
|
211
|
+
|
|
212
|
+
[[package]]
|
|
213
|
+
name = "pyo3-macros-backend"
|
|
214
|
+
version = "0.23.4"
|
|
215
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
216
|
+
checksum = "43abc3b80bc20f3facd86cd3c60beed58c3e2aa26213f3cda368de39c60a27e4"
|
|
217
|
+
dependencies = [
|
|
218
|
+
"heck",
|
|
219
|
+
"proc-macro2",
|
|
220
|
+
"pyo3-build-config",
|
|
221
|
+
"quote",
|
|
222
|
+
"syn",
|
|
223
|
+
]
|
|
224
|
+
|
|
225
|
+
[[package]]
|
|
226
|
+
name = "quote"
|
|
227
|
+
version = "1.0.38"
|
|
228
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
229
|
+
checksum = "0e4dccaaaf89514f546c693ddc140f729f958c247918a13380cccc6078391acc"
|
|
230
|
+
dependencies = [
|
|
231
|
+
"proc-macro2",
|
|
232
|
+
]
|
|
233
|
+
|
|
234
|
+
[[package]]
|
|
235
|
+
name = "rawpointer"
|
|
236
|
+
version = "0.2.1"
|
|
237
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
238
|
+
checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
|
|
239
|
+
|
|
240
|
+
[[package]]
|
|
241
|
+
name = "rustc-hash"
|
|
242
|
+
version = "2.1.0"
|
|
243
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
244
|
+
checksum = "c7fb8039b3032c191086b10f11f319a6e99e1e82889c5cc6046f515c9db1d497"
|
|
245
|
+
|
|
246
|
+
[[package]]
|
|
247
|
+
name = "serde"
|
|
248
|
+
version = "1.0.217"
|
|
249
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
250
|
+
checksum = "02fc4265df13d6fa1d00ecff087228cc0a2b5f3c0e87e258d8b94a156e984c70"
|
|
251
|
+
dependencies = [
|
|
252
|
+
"serde_derive",
|
|
253
|
+
]
|
|
254
|
+
|
|
255
|
+
[[package]]
|
|
256
|
+
name = "serde_derive"
|
|
257
|
+
version = "1.0.217"
|
|
258
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
259
|
+
checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0"
|
|
260
|
+
dependencies = [
|
|
261
|
+
"proc-macro2",
|
|
262
|
+
"quote",
|
|
263
|
+
"syn",
|
|
264
|
+
]
|
|
265
|
+
|
|
266
|
+
[[package]]
|
|
267
|
+
name = "syn"
|
|
268
|
+
version = "2.0.96"
|
|
269
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
270
|
+
checksum = "d5d0adab1ae378d7f53bdebc67a39f1f151407ef230f0ce2883572f5d8985c80"
|
|
271
|
+
dependencies = [
|
|
272
|
+
"proc-macro2",
|
|
273
|
+
"quote",
|
|
274
|
+
"unicode-ident",
|
|
275
|
+
]
|
|
276
|
+
|
|
277
|
+
[[package]]
|
|
278
|
+
name = "target-lexicon"
|
|
279
|
+
version = "0.12.16"
|
|
280
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
281
|
+
checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
|
|
282
|
+
|
|
283
|
+
[[package]]
|
|
284
|
+
name = "unicode-ident"
|
|
285
|
+
version = "1.0.14"
|
|
286
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
287
|
+
checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83"
|
|
288
|
+
|
|
289
|
+
[[package]]
|
|
290
|
+
name = "unindent"
|
|
291
|
+
version = "0.2.3"
|
|
292
|
+
source = "registry+https://github.com/rust-lang/crates.io-index"
|
|
293
|
+
checksum = "c7de7d73e1754487cb58364ee906a499937a0dfabd86bcb980fa99ec8c8fa2ce"
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
[package]
|
|
2
|
+
name = "discrete-dists"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
edition = "2021"
|
|
5
|
+
|
|
6
|
+
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
|
7
|
+
[lib]
|
|
8
|
+
name = "discrete_dists"
|
|
9
|
+
path = 'rust/lib.rs'
|
|
10
|
+
crate-type = ["cdylib"]
|
|
11
|
+
|
|
12
|
+
[dependencies]
|
|
13
|
+
bincode = "~1.3"
|
|
14
|
+
ndarray = { version = "~0.16", features = ["serde"] }
|
|
15
|
+
numpy = "~0.23"
|
|
16
|
+
pyo3 = "~0.23"
|
|
17
|
+
serde = { version = "~1.0", features = ["derive"] }
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: discrete-dists
|
|
3
|
+
Version: 0.0.0
|
|
4
|
+
Classifier: Programming Language :: Rust
|
|
5
|
+
Classifier: Programming Language :: Python :: Implementation :: CPython
|
|
6
|
+
Classifier: Programming Language :: Python :: Implementation :: PyPy
|
|
7
|
+
Requires-Dist: numpy >=1.23, <3.0
|
|
8
|
+
Requires-Dist: pytest ; extra == 'dev'
|
|
9
|
+
Requires-Dist: maturin ; extra == 'dev'
|
|
10
|
+
Requires-Dist: pip ; extra == 'dev'
|
|
11
|
+
Requires-Dist: pytest-benchmark ~=5.0 ; extra == 'dev'
|
|
12
|
+
Requires-Dist: ruff ; extra == 'dev'
|
|
13
|
+
Requires-Dist: pyright ; extra == 'dev'
|
|
14
|
+
Requires-Dist: commitizen ; extra == 'dev'
|
|
15
|
+
Requires-Dist: pre-commit ; extra == 'dev'
|
|
16
|
+
Provides-Extra: dev
|
|
17
|
+
Author-email: Andy Patterson <andnpatterson@gmail.com>
|
|
18
|
+
Requires-Python: >=3.8
|
|
File without changes
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
from abc import ABC, abstractmethod
|
|
3
|
+
import numpy as np
|
|
4
|
+
import logging
|
|
5
|
+
|
|
6
|
+
logger = logging.getLogger('discrete_dists')
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class Distribution(ABC):
|
|
10
|
+
@abstractmethod
|
|
11
|
+
def probs(self, idxs: np.ndarray) -> np.ndarray: ...
|
|
12
|
+
|
|
13
|
+
@abstractmethod
|
|
14
|
+
def update(self, idxs: np.ndarray, values: np.ndarray) -> None: ...
|
|
15
|
+
|
|
16
|
+
@abstractmethod
|
|
17
|
+
def update_single(self, idx: int, value: float) -> None: ...
|
|
18
|
+
|
|
19
|
+
@abstractmethod
|
|
20
|
+
def sample(self, rng: np.random.Generator, n: int) -> np.ndarray: ...
|
|
21
|
+
|
|
22
|
+
@abstractmethod
|
|
23
|
+
def stratified_sample(self, rng: np.random.Generator, n: int) -> np.ndarray: ...
|
|
24
|
+
|
|
25
|
+
def isr(self, target: Distribution, idxs: np.ndarray):
|
|
26
|
+
return target.probs(idxs) / self.probs(idxs)
|
|
27
|
+
|
|
28
|
+
def sample_without_replacement(
|
|
29
|
+
self,
|
|
30
|
+
rng: np.random.Generator,
|
|
31
|
+
n: int,
|
|
32
|
+
attempts: int = 25,
|
|
33
|
+
) -> np.ndarray:
|
|
34
|
+
idxs = self.sample(rng, n)
|
|
35
|
+
|
|
36
|
+
# fastpath for the common case that the first sample is already unique
|
|
37
|
+
uniq = set(idxs)
|
|
38
|
+
if len(uniq) == n:
|
|
39
|
+
return idxs
|
|
40
|
+
|
|
41
|
+
for _ in range(attempts):
|
|
42
|
+
needed = n - len(uniq)
|
|
43
|
+
sub = self.sample(rng, 2 * needed)
|
|
44
|
+
uniq |= set(sub)
|
|
45
|
+
|
|
46
|
+
if len(uniq) >= n:
|
|
47
|
+
break
|
|
48
|
+
|
|
49
|
+
if len(uniq) < n:
|
|
50
|
+
logger.warning(f"Failed to get <{n}> required unique samples. Got <{len(uniq)}>")
|
|
51
|
+
|
|
52
|
+
cutoff = min(n, len(uniq))
|
|
53
|
+
out = np.array(uniq, dtype=np.int64)[:cutoff]
|
|
54
|
+
return out
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
from dataclasses import dataclass
|
|
2
|
+
from typing import Sequence
|
|
3
|
+
|
|
4
|
+
import numpy as np
|
|
5
|
+
|
|
6
|
+
from discrete_dists.distribution import Distribution
|
|
7
|
+
|
|
8
|
+
@dataclass
|
|
9
|
+
class SubDistribution:
|
|
10
|
+
d: Distribution
|
|
11
|
+
p: float
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class MixtureDistribution(Distribution):
|
|
15
|
+
def __init__(self, dists: Sequence[SubDistribution]):
|
|
16
|
+
super().__init__()
|
|
17
|
+
|
|
18
|
+
self._dims = len(dists)
|
|
19
|
+
|
|
20
|
+
self.dists = [sub.d for sub in dists]
|
|
21
|
+
self._weights = np.array([sub.p for sub in dists])
|
|
22
|
+
|
|
23
|
+
assert np.isclose(self._weights.sum(), 1)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def probs(self, idxs: np.ndarray):
|
|
27
|
+
sub = np.array([d.probs(idxs) for d in self.dists])
|
|
28
|
+
p = self._weights.dot(sub)
|
|
29
|
+
return p
|
|
30
|
+
|
|
31
|
+
def sample(self, rng: np.random.Generator, n: int):
|
|
32
|
+
out = np.empty(n, dtype=np.int64)
|
|
33
|
+
subs = rng.choice(len(self.dists), size=n, replace=True, p=self._weights)
|
|
34
|
+
idxs, counts = np.unique(subs, return_counts=True)
|
|
35
|
+
|
|
36
|
+
total = 0
|
|
37
|
+
for idx, count in zip(idxs, counts, strict=True):
|
|
38
|
+
d = self.dists[int(idx)]
|
|
39
|
+
next_t = total + count
|
|
40
|
+
out[total:next_t] = d.sample(rng, count)
|
|
41
|
+
total = next_t
|
|
42
|
+
|
|
43
|
+
return out
|
|
44
|
+
|
|
45
|
+
def stratified_sample(self, rng: np.random.Generator, n: int):
|
|
46
|
+
out = np.empty(n, dtype=np.int64)
|
|
47
|
+
subs = rng.choice(len(self.dists), size=n, replace=True, p=self._weights)
|
|
48
|
+
idxs, counts = np.unique(subs, return_counts=True)
|
|
49
|
+
|
|
50
|
+
total = 0
|
|
51
|
+
for idx, count in zip(idxs, counts, strict=True):
|
|
52
|
+
d = self.dists[int(idx)]
|
|
53
|
+
next_t = total + count
|
|
54
|
+
out[total:next_t] = d.stratified_sample(rng, count)
|
|
55
|
+
total = next_t
|
|
56
|
+
|
|
57
|
+
return out
|
|
58
|
+
|
|
59
|
+
def update(self, idxs: np.ndarray, values: np.ndarray):
|
|
60
|
+
for d in self.dists:
|
|
61
|
+
d.update(idxs, values)
|
|
62
|
+
|
|
63
|
+
def update_single(self, idx: int, value: float):
|
|
64
|
+
for d in self.dists:
|
|
65
|
+
d.update_single(idx, value)
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
from discrete_dists.distribution import Distribution
|
|
3
|
+
from discrete_dists.utils.SumTree import SumTree
|
|
4
|
+
|
|
5
|
+
class Proportional(Distribution):
|
|
6
|
+
def __init__(self, support: int):
|
|
7
|
+
self._support = support
|
|
8
|
+
self.tree = SumTree(support)
|
|
9
|
+
|
|
10
|
+
# ---------------
|
|
11
|
+
# -- Accessing --
|
|
12
|
+
# ---------------
|
|
13
|
+
def probs(self, idxs: np.ndarray) -> np.ndarray:
|
|
14
|
+
idxs = np.asarray(idxs)
|
|
15
|
+
|
|
16
|
+
t = self.tree.total()
|
|
17
|
+
if t == 0:
|
|
18
|
+
return np.zeros(len(idxs))
|
|
19
|
+
|
|
20
|
+
v = self.tree.get_values(idxs)
|
|
21
|
+
return v / t
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
# --------------
|
|
25
|
+
# -- Sampling --
|
|
26
|
+
# --------------
|
|
27
|
+
def sample(self, rng: np.random.Generator, n: int) -> np.ndarray:
|
|
28
|
+
return self.tree.sample(rng, n)
|
|
29
|
+
|
|
30
|
+
def stratified_sample(self, rng: np.random.Generator, n: int) -> np.ndarray:
|
|
31
|
+
return self.tree.stratified_sample(rng, n)
|
|
32
|
+
|
|
33
|
+
# --------------
|
|
34
|
+
# -- Updating --
|
|
35
|
+
# --------------
|
|
36
|
+
def update(self, idxs: np.ndarray, values: np.ndarray):
|
|
37
|
+
self.tree.update(idxs, values)
|
|
38
|
+
|
|
39
|
+
def update_single(self, idx: int, value: float):
|
|
40
|
+
self.tree.update_single(idx, value)
|
|
File without changes
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
import numpy.typing as npt
|
|
3
|
+
import discrete_dists.utils.npu as npu
|
|
4
|
+
from discrete_dists.distribution import Distribution
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
class Uniform(Distribution):
|
|
8
|
+
def __init__(self, support: int):
|
|
9
|
+
self._support = support
|
|
10
|
+
|
|
11
|
+
def update(self, idxs: np.ndarray, values: np.ndarray):
|
|
12
|
+
self._support = max(self._support, idxs.max())
|
|
13
|
+
|
|
14
|
+
def update_single(self, idx: int, value: float):
|
|
15
|
+
self._support = max(self._support, idx)
|
|
16
|
+
|
|
17
|
+
def sample(self, rng: np.random.Generator, n: int):
|
|
18
|
+
if self._support == 1:
|
|
19
|
+
return np.zeros(n, dtype=np.int64)
|
|
20
|
+
|
|
21
|
+
return rng.integers(0, self._support, size=n)
|
|
22
|
+
|
|
23
|
+
def stratified_sample(self, rng: np.random.Generator, n: int):
|
|
24
|
+
return npu.stratified_sample_integers(rng, n, self._support)
|
|
25
|
+
|
|
26
|
+
def probs(self, idxs: npt.ArrayLike):
|
|
27
|
+
return np.full_like(idxs, fill_value=(1 / self._support), dtype=np.float64)
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
from typing import Iterable
|
|
3
|
+
import discrete_dists.rust as ru
|
|
4
|
+
|
|
5
|
+
class SumTree(ru.SumTree):
|
|
6
|
+
def __new__(
|
|
7
|
+
cls,
|
|
8
|
+
size: int | None = None,
|
|
9
|
+
):
|
|
10
|
+
args = [size]
|
|
11
|
+
if args[0] is None:
|
|
12
|
+
args = []
|
|
13
|
+
|
|
14
|
+
return super().__new__(cls, *args)
|
|
15
|
+
|
|
16
|
+
def __init__(self, size: int):
|
|
17
|
+
super().__init__()
|
|
18
|
+
|
|
19
|
+
def update(self, idxs: Iterable[int], values: Iterable[float]):
|
|
20
|
+
a_idxs = np.asarray(idxs, dtype=np.int64)
|
|
21
|
+
a_values = np.asarray(values, dtype=np.float64)
|
|
22
|
+
|
|
23
|
+
super().update(a_idxs, a_values)
|
|
24
|
+
|
|
25
|
+
def sample(self, rng: np.random.Generator, n: int) -> np.ndarray:
|
|
26
|
+
t = self.total()
|
|
27
|
+
assert t > 0, "Cannot sample when the tree is empty or contains negative values"
|
|
28
|
+
|
|
29
|
+
rs = rng.uniform(0, t, size=n)
|
|
30
|
+
return self.query(rs)
|
|
31
|
+
|
|
32
|
+
def stratified_sample(self, rng: np.random.Generator, n: int) -> np.ndarray:
|
|
33
|
+
t = self.total()
|
|
34
|
+
assert t > 0, "Cannot sample when the tree is empty or contains negative values"
|
|
35
|
+
|
|
36
|
+
buckets = np.linspace(0., 1., n + 1)
|
|
37
|
+
values = np.asarray([
|
|
38
|
+
rng.uniform(buckets[i], buckets[i + 1]) for i in range(n)
|
|
39
|
+
])
|
|
40
|
+
|
|
41
|
+
return super().query(values)
|
|
42
|
+
|
|
43
|
+
def __getstate__(self):
|
|
44
|
+
return {
|
|
45
|
+
'st': super().__getstate__()
|
|
46
|
+
}
|
|
47
|
+
|
|
48
|
+
def __setstate__(self, state):
|
|
49
|
+
super().__setstate__(state['st'])
|
|
File without changes
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
|
|
3
|
+
def stratified_sample_integers(rng: np.random.Generator, n: int, size: int):
|
|
4
|
+
buckets = np.linspace(0, size, n + 1, endpoint=True, dtype=np.int64)
|
|
5
|
+
samples = [
|
|
6
|
+
rng.integers(buckets[i], buckets[i + 1]) for i in range(n)
|
|
7
|
+
]
|
|
8
|
+
return np.asarray(samples, dtype=np.int64)
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["maturin>=1.7,<2.0"]
|
|
3
|
+
build-backend = "maturin"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "discrete-dists"
|
|
7
|
+
requires-python = ">=3.8"
|
|
8
|
+
version = "0.0.0"
|
|
9
|
+
authors = [
|
|
10
|
+
{name = "Andy Patterson", email = "andnpatterson@gmail.com"},
|
|
11
|
+
]
|
|
12
|
+
dependencies = [
|
|
13
|
+
"numpy>=1.23,<3.0"
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
classifiers = [
|
|
18
|
+
"Programming Language :: Rust",
|
|
19
|
+
"Programming Language :: Python :: Implementation :: CPython",
|
|
20
|
+
"Programming Language :: Python :: Implementation :: PyPy",
|
|
21
|
+
]
|
|
22
|
+
|
|
23
|
+
[project.optional-dependencies]
|
|
24
|
+
dev = [
|
|
25
|
+
"pytest",
|
|
26
|
+
"maturin",
|
|
27
|
+
"pip",
|
|
28
|
+
"pytest-benchmark~=5.0",
|
|
29
|
+
"ruff",
|
|
30
|
+
"pyright",
|
|
31
|
+
"commitizen",
|
|
32
|
+
"pre-commit",
|
|
33
|
+
]
|
|
34
|
+
|
|
35
|
+
[tool]
|
|
36
|
+
[tool.commitizen]
|
|
37
|
+
name = "cz_conventional_commits"
|
|
38
|
+
version = "0.0.0"
|
|
39
|
+
tag_format = "$version"
|
|
40
|
+
version_files = ["pyproject.toml"]
|
|
41
|
+
|
|
42
|
+
[tool.maturin]
|
|
43
|
+
features = ["pyo3/extension-module"]
|
|
44
|
+
module-name = "discrete_dists.rust"
|
|
45
|
+
|
|
46
|
+
[tool.ruff.lint]
|
|
47
|
+
select = ['F', 'E', 'W']
|
|
48
|
+
ignore = ['E501', 'E701']
|
|
49
|
+
|
|
50
|
+
[tool.pyright]
|
|
51
|
+
include = ["discrete_dists", "tests"]
|
|
52
|
+
stubPath = "./typings"
|
|
53
|
+
venvPath = "."
|
|
54
|
+
venv = ".venv"
|
|
55
|
+
typeCheckingMode = "standard"
|
|
56
|
+
useLibraryCodeForTypes = true
|
|
@@ -0,0 +1,151 @@
|
|
|
1
|
+
use numpy::{ToPyArray, PyArray1, PyReadonlyArray1};
|
|
2
|
+
use ndarray::{Array1, Axis};
|
|
3
|
+
use std::iter;
|
|
4
|
+
use std::cmp::*;
|
|
5
|
+
use pyo3::{prelude::*, types::{PyBytes, PyTuple}};
|
|
6
|
+
use bincode::{deserialize, serialize};
|
|
7
|
+
use serde::{Deserialize, Serialize};
|
|
8
|
+
|
|
9
|
+
#[pyclass(module = "rust", subclass)]
|
|
10
|
+
#[derive(Serialize, Deserialize, Clone)]
|
|
11
|
+
|
|
12
|
+
pub struct SumTree {
|
|
13
|
+
#[pyo3(get)]
|
|
14
|
+
size: u32,
|
|
15
|
+
total_size: u32,
|
|
16
|
+
raw: Vec<Array1<f64>>,
|
|
17
|
+
}
|
|
18
|
+
|
|
19
|
+
#[pymethods]
|
|
20
|
+
impl SumTree {
|
|
21
|
+
#[new]
|
|
22
|
+
#[pyo3(signature = (*args))]
|
|
23
|
+
fn new<'py>(args: Bound<'py, PyTuple>) -> Self {
|
|
24
|
+
match args.len() {
|
|
25
|
+
0 => SumTree {
|
|
26
|
+
size: 0,
|
|
27
|
+
total_size: 0,
|
|
28
|
+
raw: vec![],
|
|
29
|
+
},
|
|
30
|
+
|
|
31
|
+
1 => {
|
|
32
|
+
let size = args
|
|
33
|
+
.get_item(0).unwrap()
|
|
34
|
+
.extract::<u32>().unwrap();
|
|
35
|
+
|
|
36
|
+
let total_size = u32::next_power_of_two(size);
|
|
37
|
+
let n_layers = u32::ilog2(total_size) + 1;
|
|
38
|
+
|
|
39
|
+
let dummy = Array1::<f64>::zeros(1);
|
|
40
|
+
let mut layers = vec![dummy; n_layers as usize];
|
|
41
|
+
|
|
42
|
+
for i in (0..n_layers).rev() {
|
|
43
|
+
let r = n_layers - i - 1;
|
|
44
|
+
let layer = Array1::<f64>::zeros(usize::pow(2, i));
|
|
45
|
+
layers[r as usize] = layer;
|
|
46
|
+
}
|
|
47
|
+
|
|
48
|
+
SumTree {
|
|
49
|
+
size,
|
|
50
|
+
total_size,
|
|
51
|
+
raw: layers,
|
|
52
|
+
}
|
|
53
|
+
},
|
|
54
|
+
|
|
55
|
+
_ => unreachable!(),
|
|
56
|
+
}
|
|
57
|
+
}
|
|
58
|
+
|
|
59
|
+
pub fn update(
|
|
60
|
+
&mut self,
|
|
61
|
+
idxs: PyReadonlyArray1<i64>,
|
|
62
|
+
values: PyReadonlyArray1<f64>,
|
|
63
|
+
) {
|
|
64
|
+
iter::zip(idxs.as_array(), values.as_array())
|
|
65
|
+
.for_each(|(idx, v)| { self.update_single(*idx, *v) });
|
|
66
|
+
}
|
|
67
|
+
|
|
68
|
+
pub fn update_single(
|
|
69
|
+
&mut self,
|
|
70
|
+
idx: i64,
|
|
71
|
+
value: f64,
|
|
72
|
+
) {
|
|
73
|
+
if idx >= self.size as i64 {
|
|
74
|
+
panic!("Tried to update index outside of tree: <{idx}>");
|
|
75
|
+
}
|
|
76
|
+
|
|
77
|
+
let mut sub_idx = idx as usize;
|
|
78
|
+
let old = self.raw[0][sub_idx];
|
|
79
|
+
|
|
80
|
+
self.raw.iter_mut()
|
|
81
|
+
.for_each(|level| {
|
|
82
|
+
level[sub_idx] += value - old;
|
|
83
|
+
sub_idx = sub_idx / 2;
|
|
84
|
+
});
|
|
85
|
+
}
|
|
86
|
+
|
|
87
|
+
pub fn get_value(&mut self, idx: i64) -> f64 {
|
|
88
|
+
self.raw[0][idx as usize]
|
|
89
|
+
}
|
|
90
|
+
|
|
91
|
+
pub fn get_values<'py>(
|
|
92
|
+
&mut self,
|
|
93
|
+
idxs: PyReadonlyArray1<i64>,
|
|
94
|
+
py: Python<'py>,
|
|
95
|
+
) -> Bound<'py, PyArray1<f64>> {
|
|
96
|
+
let idxs = idxs.as_array().map(|a| { *a as usize });
|
|
97
|
+
let arr = self.raw[0]
|
|
98
|
+
.select(Axis(0), &idxs.to_vec());
|
|
99
|
+
|
|
100
|
+
arr.to_pyarray(py)
|
|
101
|
+
}
|
|
102
|
+
|
|
103
|
+
pub fn total(
|
|
104
|
+
&mut self,
|
|
105
|
+
) -> f64 {
|
|
106
|
+
*self.raw
|
|
107
|
+
.last()
|
|
108
|
+
.expect("")
|
|
109
|
+
.get(0)
|
|
110
|
+
.expect("")
|
|
111
|
+
}
|
|
112
|
+
|
|
113
|
+
pub fn query<'py>(
|
|
114
|
+
&mut self,
|
|
115
|
+
v: PyReadonlyArray1<f64>,
|
|
116
|
+
py: Python<'py>,
|
|
117
|
+
) -> Bound<'py, PyArray1<i64>> {
|
|
118
|
+
let n = v.len().expect("Failed to get array length");
|
|
119
|
+
|
|
120
|
+
let v = v.as_array();
|
|
121
|
+
let mut totals = Array1::<f64>::zeros(n);
|
|
122
|
+
let mut idxs = Array1::<i64>::zeros(n);
|
|
123
|
+
|
|
124
|
+
self.raw.iter()
|
|
125
|
+
.rev()
|
|
126
|
+
.for_each(|layer| {
|
|
127
|
+
for j in 0..n {
|
|
128
|
+
idxs[j] = idxs[j] * 2;
|
|
129
|
+
let left = *layer
|
|
130
|
+
.get(idxs[j] as usize)
|
|
131
|
+
.expect("");
|
|
132
|
+
|
|
133
|
+
let m = left < (v[j] - totals[j]);
|
|
134
|
+
totals[j] += if m { left } else { 0. };
|
|
135
|
+
idxs[j] += if m { 1 } else { 0 };
|
|
136
|
+
}
|
|
137
|
+
});
|
|
138
|
+
|
|
139
|
+
idxs = idxs.map(|i| { min(*i, (self.size - 1) as i64) });
|
|
140
|
+
idxs.to_pyarray(py)
|
|
141
|
+
}
|
|
142
|
+
|
|
143
|
+
// enable pickling this data type
|
|
144
|
+
pub fn __setstate__<'py>(&mut self, state: Bound<'py, PyBytes>) -> PyResult<()> {
|
|
145
|
+
*self = deserialize(state.as_bytes()).unwrap();
|
|
146
|
+
Ok(())
|
|
147
|
+
}
|
|
148
|
+
pub fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
|
|
149
|
+
Ok(PyBytes::new(py, &serialize(&self).unwrap()))
|
|
150
|
+
}
|
|
151
|
+
}
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
pub mod sumtree;
|
|
File without changes
|
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
import pickle
|
|
2
|
+
import numpy as np
|
|
3
|
+
|
|
4
|
+
from discrete_dists.utils.SumTree import SumTree
|
|
5
|
+
|
|
6
|
+
class TestSumTree:
|
|
7
|
+
def test_can_add_stuff(self):
|
|
8
|
+
# human readable test
|
|
9
|
+
v = np.array([1, 2, 3, 4, 5])
|
|
10
|
+
truth = 15
|
|
11
|
+
tree = SumTree(500)
|
|
12
|
+
tree.update([0, 5, 8, 102, 358], v)
|
|
13
|
+
|
|
14
|
+
assert tree.total() == truth
|
|
15
|
+
|
|
16
|
+
# fuzz test
|
|
17
|
+
tree = SumTree(100)
|
|
18
|
+
rng = np.random.default_rng(0)
|
|
19
|
+
for _ in range(100):
|
|
20
|
+
v = rng.standard_normal(100)
|
|
21
|
+
truth = v.sum()
|
|
22
|
+
tree.update(np.arange(100), v)
|
|
23
|
+
|
|
24
|
+
assert np.isclose(tree.total(), truth)
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def test_can_sample(self):
|
|
28
|
+
tree = SumTree(50)
|
|
29
|
+
tree.update(np.arange(50), np.ones(50))
|
|
30
|
+
|
|
31
|
+
assert tree.total() == 50
|
|
32
|
+
|
|
33
|
+
rng = np.random.default_rng(0)
|
|
34
|
+
samples = tree.sample(rng, 10000)
|
|
35
|
+
|
|
36
|
+
u, c = np.unique(samples, return_counts=True)
|
|
37
|
+
assert np.all(u == np.arange(50))
|
|
38
|
+
assert np.all(
|
|
39
|
+
(c >= 150) & (c <= 250)
|
|
40
|
+
)
|
|
41
|
+
|
|
42
|
+
def test_can_sample_proportionally(self):
|
|
43
|
+
tree = SumTree(10)
|
|
44
|
+
tree.update(np.arange(10), [2**i for i in range(10)])
|
|
45
|
+
|
|
46
|
+
# NOTE: it takes a shockingly high number of samples for the proportions
|
|
47
|
+
# to converge even within a single decimal point..
|
|
48
|
+
rng = np.random.default_rng(22)
|
|
49
|
+
samples = tree.sample(rng, 10000000)
|
|
50
|
+
|
|
51
|
+
u, c = np.unique(samples, return_counts=True)
|
|
52
|
+
assert np.all(u == np.arange(10))
|
|
53
|
+
|
|
54
|
+
for i in range(1, 10):
|
|
55
|
+
assert np.isclose(c[i] / c[i - 1], 2, atol=0.1)
|
|
56
|
+
|
|
57
|
+
def test_pickleable(self):
|
|
58
|
+
tree = SumTree(123)
|
|
59
|
+
tree.update(np.arange(123), np.arange(123))
|
|
60
|
+
|
|
61
|
+
byt = pickle.dumps(tree)
|
|
62
|
+
tree2 = pickle.loads(byt)
|
|
63
|
+
|
|
64
|
+
assert np.all(
|
|
65
|
+
tree.total() == tree2.total()
|
|
66
|
+
)
|
|
67
|
+
|
|
68
|
+
tree.update([2], [22])
|
|
69
|
+
tree2.update([2], [22])
|
|
70
|
+
|
|
71
|
+
assert np.all(
|
|
72
|
+
tree.total() == tree2.total()
|
|
73
|
+
)
|
|
74
|
+
|
|
75
|
+
# ----------------
|
|
76
|
+
# -- Benchmarks --
|
|
77
|
+
# ----------------
|
|
78
|
+
class TestBenchmarks:
|
|
79
|
+
def test_sumtree_update(self, benchmark):
|
|
80
|
+
tree = SumTree(100_000)
|
|
81
|
+
rng = np.random.default_rng(0)
|
|
82
|
+
idxs = np.arange(32)
|
|
83
|
+
vals = rng.uniform(0, 2, size=32)
|
|
84
|
+
|
|
85
|
+
def _inner(tree: SumTree, idxs, vals):
|
|
86
|
+
tree.update(idxs, vals)
|
|
87
|
+
|
|
88
|
+
benchmark(_inner, tree, idxs, vals)
|
|
89
|
+
|
|
90
|
+
def test_sumtree_sample(self, benchmark):
|
|
91
|
+
tree = SumTree(100_000)
|
|
92
|
+
rng = np.random.default_rng(0)
|
|
93
|
+
|
|
94
|
+
idxs = np.arange(10_000)
|
|
95
|
+
vals = rng.uniform(0, 2, size=10_000)
|
|
96
|
+
tree.update(idxs, vals)
|
|
97
|
+
|
|
98
|
+
def _inner(tree: SumTree, rng):
|
|
99
|
+
tree.sample(rng, 32)
|
|
100
|
+
|
|
101
|
+
benchmark(_inner, tree, rng)
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
|
|
3
|
+
class SumTree:
|
|
4
|
+
size: int
|
|
5
|
+
|
|
6
|
+
def __new__(cls, *args): ...
|
|
7
|
+
def __init__(self, size: int | None = None): ...
|
|
8
|
+
def update(self, idxs: np.ndarray, values: np.ndarray): ...
|
|
9
|
+
def update_single(self, idx: int, value: float): ...
|
|
10
|
+
def get_value(self, idx: int) -> float: ...
|
|
11
|
+
def get_values(self, idxs: np.ndarray) -> np.ndarray: ...
|
|
12
|
+
def total(self) -> float: ...
|
|
13
|
+
def query(self, v: np.ndarray) -> np.ndarray: ...
|
|
14
|
+
def __getstate__(self): ...
|
|
15
|
+
def __setstate__(self, state): ...
|