turbocluster 0.1__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.
- turbocluster-0.1/PKG-INFO +39 -0
- turbocluster-0.1/README.md +21 -0
- turbocluster-0.1/pyproject.toml +29 -0
- turbocluster-0.1/requirements.txt +9 -0
- turbocluster-0.1/setup.cfg +4 -0
- turbocluster-0.1/tests/test_cuda.py +58 -0
- turbocluster-0.1/turbocluster/__init__.py +9 -0
- turbocluster-0.1/turbocluster/cartesian_tiling.py +247 -0
- turbocluster-0.1/turbocluster/data_init.py +163 -0
- turbocluster-0.1/turbocluster/helper_functions.py +67 -0
- turbocluster-0.1/turbocluster.egg-info/PKG-INFO +39 -0
- turbocluster-0.1/turbocluster.egg-info/SOURCES.txt +13 -0
- turbocluster-0.1/turbocluster.egg-info/dependency_links.txt +1 -0
- turbocluster-0.1/turbocluster.egg-info/requires.txt +9 -0
- turbocluster-0.1/turbocluster.egg-info/top_level.txt +1 -0
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: turbocluster
|
|
3
|
+
Version: 0.1
|
|
4
|
+
Summary: A GPU-accelerated Python library for filtering turbulence in Arepo simulations
|
|
5
|
+
Author-email: Lorenzo Maria Perrone <lorenzolmp@protonmail.com>, Thomas Berlok <tberlok@gmail.com>
|
|
6
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
7
|
+
Requires-Python: >=3.14.0
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
Requires-Dist: scipy>=1.7.3
|
|
10
|
+
Requires-Dist: numpy<=2.4.6
|
|
11
|
+
Requires-Dist: matplotlib
|
|
12
|
+
Requires-Dist: h5py
|
|
13
|
+
Requires-Dist: astropy>=4.3
|
|
14
|
+
Requires-Dist: numba-cuda[cu13]
|
|
15
|
+
Requires-Dist: paicos[dev]>=0.1.17
|
|
16
|
+
Requires-Dist: cupy-cuda13x[ctk]
|
|
17
|
+
Requires-Dist: nvtx
|
|
18
|
+
|
|
19
|
+
[](https://turbocluster.readthedocs.io/en/latest/)
|
|
20
|
+
[](
|
|
21
|
+
https://github.com/LorenzoLMP/turbocluster/actions/workflows/flake8.yml)
|
|
22
|
+
|
|
23
|
+
# turbocluster
|
|
24
|
+
|
|
25
|
+
A GPU-accelerated Python library for filtering turbulence in Arepo simulations.
|
|
26
|
+
|
|
27
|
+
<img src="images/turbocluster_banner.png" width="auto">
|
|
28
|
+
|
|
29
|
+
**NOTE**: This project is in active development and not officially released publicly yet. Expect kinks and bugs. Refer to [this arXiv paper](https://arxiv.org/pdf/2601.06250) if you would like to know more about its use.
|
|
30
|
+
|
|
31
|
+
A work-in-progress documentation can be found at [readthedocs](https://turbocluster.readthedocs.io/en/latest/).
|
|
32
|
+
|
|
33
|
+
# TODO:
|
|
34
|
+
- double-check tests and examples
|
|
35
|
+
- can cuda.jit() be replaced by jit() on non-cuda hardware?
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
[](https://turbocluster.readthedocs.io/en/latest/)
|
|
2
|
+
[](
|
|
3
|
+
https://github.com/LorenzoLMP/turbocluster/actions/workflows/flake8.yml)
|
|
4
|
+
|
|
5
|
+
# turbocluster
|
|
6
|
+
|
|
7
|
+
A GPU-accelerated Python library for filtering turbulence in Arepo simulations.
|
|
8
|
+
|
|
9
|
+
<img src="images/turbocluster_banner.png" width="auto">
|
|
10
|
+
|
|
11
|
+
**NOTE**: This project is in active development and not officially released publicly yet. Expect kinks and bugs. Refer to [this arXiv paper](https://arxiv.org/pdf/2601.06250) if you would like to know more about its use.
|
|
12
|
+
|
|
13
|
+
A work-in-progress documentation can be found at [readthedocs](https://turbocluster.readthedocs.io/en/latest/).
|
|
14
|
+
|
|
15
|
+
# TODO:
|
|
16
|
+
- double-check tests and examples
|
|
17
|
+
- can cuda.jit() be replaced by jit() on non-cuda hardware?
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "turbocluster"
|
|
3
|
+
version = "0.1"
|
|
4
|
+
authors = [
|
|
5
|
+
{ name="Lorenzo Maria Perrone", email="lorenzolmp@protonmail.com" },
|
|
6
|
+
{ name="Thomas Berlok", email="tberlok@gmail.com" },
|
|
7
|
+
]
|
|
8
|
+
description = "A GPU-accelerated Python library for filtering turbulence in Arepo simulations"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.14.0"
|
|
11
|
+
classifiers = [
|
|
12
|
+
"Programming Language :: Python :: 3.14",
|
|
13
|
+
]
|
|
14
|
+
|
|
15
|
+
dynamic = ["dependencies"]
|
|
16
|
+
|
|
17
|
+
[tool.setuptools.dynamic]
|
|
18
|
+
dependencies = {file = ["requirements.txt"]}
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
[tool.setuptools.packages.find]
|
|
22
|
+
where = ["."] # list of folders that contain the packages (["."] by default)
|
|
23
|
+
include = ["turbocluster"] # package names shbuld match these glob patterns (["*"] by default)
|
|
24
|
+
exclude = [] # exclude packages matching these glob patterns (empty by default)
|
|
25
|
+
namespaces = false # to disable scanning PEP 420 namespaces (true by default)
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
[tool.cibuildwheel]
|
|
29
|
+
before-all = "uname -a"
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
import cupy as cp
|
|
2
|
+
from numba import cuda
|
|
3
|
+
import numpy as np
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@cuda.jit()
|
|
7
|
+
def check_particle(pos, hsml, center, widths, isParticleInDomain):
|
|
8
|
+
""" """
|
|
9
|
+
ip = cuda.grid(1)
|
|
10
|
+
# each thread is assigned to a particle
|
|
11
|
+
numParticles = pos.shape[0]
|
|
12
|
+
isParticleInDomainTmp = 0
|
|
13
|
+
|
|
14
|
+
if ip < numParticles:
|
|
15
|
+
|
|
16
|
+
xp, yp, zp = pos[ip]
|
|
17
|
+
xmin = center[0] - widths[0] / 2 - hsml[ip]
|
|
18
|
+
xmax = center[0] + widths[0] / 2 + hsml[ip]
|
|
19
|
+
|
|
20
|
+
ymin = center[1] - widths[1] / 2 - hsml[ip]
|
|
21
|
+
ymax = center[1] + widths[1] / 2 + hsml[ip]
|
|
22
|
+
|
|
23
|
+
zmin = center[2] - widths[2] / 2 - hsml[ip]
|
|
24
|
+
zmax = center[2] + widths[2] / 2 + hsml[ip]
|
|
25
|
+
|
|
26
|
+
if (xp > xmin) and (xp < xmax):
|
|
27
|
+
if (yp > ymin) and (yp < ymax):
|
|
28
|
+
if (zp > zmin) and (zp < zmax):
|
|
29
|
+
isParticleInDomainTmp = 1
|
|
30
|
+
|
|
31
|
+
isParticleInDomain[ip] = isParticleInDomainTmp
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
pos = cp.random.random((1000000, 3))
|
|
35
|
+
isParticleInDomain = cp.zeros(pos.shape[0])
|
|
36
|
+
hsml = cp.zeros(pos.shape[0])
|
|
37
|
+
widths = cp.array([1, 1, 1])
|
|
38
|
+
center = cp.array([0.5, 0.5, 0.5])
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
Np = pos.shape[0]
|
|
42
|
+
threadsperblock = 256
|
|
43
|
+
blocks_1d = (Np + (threadsperblock - 1)) // threadsperblock
|
|
44
|
+
|
|
45
|
+
if __name__ == "__main__":
|
|
46
|
+
|
|
47
|
+
try:
|
|
48
|
+
check_particle[blocks_1d, threadsperblock](
|
|
49
|
+
pos, hsml, center, widths, isParticleInDomain
|
|
50
|
+
)
|
|
51
|
+
success = np.array_equal(cp.asnumpy(isParticleInDomain), np.ones(pos.shape[0]))
|
|
52
|
+
|
|
53
|
+
if success:
|
|
54
|
+
print("The test worked")
|
|
55
|
+
else:
|
|
56
|
+
raise ValueError("The test failed.")
|
|
57
|
+
except Exception as e:
|
|
58
|
+
print(e)
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
from .cartesian_tiling import CartesianTiling
|
|
2
|
+
from .data_init import DataGpuInit
|
|
3
|
+
|
|
4
|
+
from .SmoothingFilter.smoothing_filter import SmoothingFilter
|
|
5
|
+
from .DepositToGrid.deposit_routines import DepositCartesianGrid
|
|
6
|
+
from .PotentialEnergy.potential_energy import PotentialEnergy
|
|
7
|
+
|
|
8
|
+
# from .helper_functions import *
|
|
9
|
+
# from .utils.power_spectra_utils import *
|
|
@@ -0,0 +1,247 @@
|
|
|
1
|
+
import cupy as cp
|
|
2
|
+
from numba import cuda
|
|
3
|
+
import nvtx
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
@cuda.jit(device=True, inline=True)
|
|
7
|
+
def find_tile_single_dim(particle_pos, nx, tile_width):
|
|
8
|
+
""" """
|
|
9
|
+
tile_index = int(particle_pos / tile_width)
|
|
10
|
+
return tile_index
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@cuda.jit
|
|
14
|
+
def find_tile_index(pos, npixs, tile_widths, tile_index):
|
|
15
|
+
""" """
|
|
16
|
+
ip = cuda.grid(1)
|
|
17
|
+
|
|
18
|
+
if ip < pos.shape[0]:
|
|
19
|
+
xp = pos[ip, 0]
|
|
20
|
+
yp = pos[ip, 1]
|
|
21
|
+
zp = pos[ip, 2]
|
|
22
|
+
ip_tile_x = find_tile_single_dim(xp, npixs[0], tile_widths[0])
|
|
23
|
+
ip_tile_y = find_tile_single_dim(yp, npixs[1], tile_widths[1])
|
|
24
|
+
ip_tile_z = find_tile_single_dim(zp, npixs[2], tile_widths[2])
|
|
25
|
+
tile_index[ip, 0] = ip_tile_x
|
|
26
|
+
tile_index[ip, 1] = ip_tile_y
|
|
27
|
+
tile_index[ip, 2] = ip_tile_z
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
@cuda.jit
|
|
31
|
+
def get_tile_information(tile_index, particles_per_tile, start_index_for_tile):
|
|
32
|
+
"""
|
|
33
|
+
Need to do this on each block and then assemble
|
|
34
|
+
"""
|
|
35
|
+
ip = cuda.grid(1)
|
|
36
|
+
if ip < tile_index.shape[0]:
|
|
37
|
+
ip_tile_x, ip_tile_y, ip_tile_z = tile_index[ip, :]
|
|
38
|
+
cuda.atomic.add(particles_per_tile, (ip_tile_x, ip_tile_y, ip_tile_z), 1)
|
|
39
|
+
cuda.atomic.min(start_index_for_tile, (ip_tile_x, ip_tile_y, ip_tile_z), ip)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
@cuda.jit
|
|
43
|
+
def accumulate_quantity_per_tile(tile_index, quantity_per_tile, quantity):
|
|
44
|
+
""" """
|
|
45
|
+
ip = cuda.grid(1)
|
|
46
|
+
if ip < tile_index.shape[0]:
|
|
47
|
+
ip_tile_x, ip_tile_y, ip_tile_z = tile_index[ip, :]
|
|
48
|
+
cuda.atomic.add(
|
|
49
|
+
quantity_per_tile, (ip_tile_x, ip_tile_y, ip_tile_z), quantity[ip]
|
|
50
|
+
)
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
@cuda.jit
|
|
54
|
+
def findmax_quantity_per_tile(tile_index, quantity_per_tile, quantity):
|
|
55
|
+
""" """
|
|
56
|
+
ip = cuda.grid(1)
|
|
57
|
+
if ip < tile_index.shape[0]:
|
|
58
|
+
ip_tile_x, ip_tile_y, ip_tile_z = tile_index[ip, :]
|
|
59
|
+
cuda.atomic.max(
|
|
60
|
+
quantity_per_tile, (ip_tile_x, ip_tile_y, ip_tile_z), quantity[ip]
|
|
61
|
+
)
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
@cuda.jit
|
|
65
|
+
def compactify_kernel(
|
|
66
|
+
occupancy_arr,
|
|
67
|
+
cumulative_occupancy_flat,
|
|
68
|
+
npixs,
|
|
69
|
+
Nmax,
|
|
70
|
+
particles_per_tile,
|
|
71
|
+
start_index_for_tile,
|
|
72
|
+
numBlocksCompactGrid,
|
|
73
|
+
compactGrid,
|
|
74
|
+
):
|
|
75
|
+
""" """
|
|
76
|
+
ip = cuda.grid(1)
|
|
77
|
+
numBlocksFullGrid = int(npixs[0] * npixs[1] * npixs[2])
|
|
78
|
+
if ip < numBlocksFullGrid:
|
|
79
|
+
newPos = int(cumulative_occupancy_flat[ip])
|
|
80
|
+
numBlocksNeeded = int(occupancy_arr[ip])
|
|
81
|
+
# we walk backwards like a shrimp
|
|
82
|
+
for i in range(numBlocksNeeded):
|
|
83
|
+
compactGrid[newPos - 1 - i, 0] = ip
|
|
84
|
+
compactGrid[newPos - 1 - i, 1] = (
|
|
85
|
+
start_index_for_tile[ip] + (numBlocksNeeded - 1 - i) * Nmax
|
|
86
|
+
)
|
|
87
|
+
if i > 0:
|
|
88
|
+
compactGrid[newPos - 1 - i, 2] = Nmax
|
|
89
|
+
else:
|
|
90
|
+
compactGrid[newPos - 1 - i, 2] = (
|
|
91
|
+
particles_per_tile[ip] - (numBlocksNeeded - 1) * Nmax
|
|
92
|
+
)
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
class CartesianTiling:
|
|
96
|
+
""" """
|
|
97
|
+
|
|
98
|
+
def __init__(
|
|
99
|
+
self,
|
|
100
|
+
positions,
|
|
101
|
+
center,
|
|
102
|
+
widths,
|
|
103
|
+
extra_layer_thickness,
|
|
104
|
+
npix=128,
|
|
105
|
+
threadsperblock=256,
|
|
106
|
+
):
|
|
107
|
+
|
|
108
|
+
self.Np = Np = positions.shape[0]
|
|
109
|
+
|
|
110
|
+
self.blocks_1d = (Np + (threadsperblock - 1)) // threadsperblock
|
|
111
|
+
self.threadsperblock = threadsperblock
|
|
112
|
+
|
|
113
|
+
# Copy positions
|
|
114
|
+
self._pos = cp.array(positions)
|
|
115
|
+
|
|
116
|
+
self.tilebox_widths = widths + 2 * extra_layer_thickness
|
|
117
|
+
|
|
118
|
+
npix_x = npix
|
|
119
|
+
npix_y = int(self.tilebox_widths[1] / self.tilebox_widths[0] * npix)
|
|
120
|
+
npix_z = int(self.tilebox_widths[2] / self.tilebox_widths[0] * npix)
|
|
121
|
+
|
|
122
|
+
npix_y = max([npix_y, 1])
|
|
123
|
+
npix_z = max([npix_z, 1])
|
|
124
|
+
|
|
125
|
+
self.npixs = cp.array([npix_x, npix_y, npix_z])
|
|
126
|
+
|
|
127
|
+
self.off_sets = center - self.tilebox_widths / 2.0
|
|
128
|
+
|
|
129
|
+
self.tile_widths = self.tilebox_widths / self.npixs
|
|
130
|
+
|
|
131
|
+
self._pos -= self.off_sets[None, :]
|
|
132
|
+
|
|
133
|
+
def sort_particles_in_tiles(self):
|
|
134
|
+
|
|
135
|
+
npix_x, npix_y, npix_z = self.npixs
|
|
136
|
+
|
|
137
|
+
# tile_index tells us which tile does each particle belong to
|
|
138
|
+
self.tile_index = cp.zeros((self.Np, 3), dtype=int)
|
|
139
|
+
|
|
140
|
+
self.particles_per_tile = cp.zeros(self.npixs.tolist(), dtype=int)
|
|
141
|
+
# Initialize to value larger than the largest posssible tileindex
|
|
142
|
+
self.start_index_for_tile = (
|
|
143
|
+
cp.ones(self.npixs.tolist(), dtype=int) * self.Np + 1
|
|
144
|
+
)
|
|
145
|
+
|
|
146
|
+
# Get what tile does each particle belong to
|
|
147
|
+
rng = nvtx.start_range(message="find_tile_index")
|
|
148
|
+
find_tile_index[self.blocks_1d, self.threadsperblock](
|
|
149
|
+
self._pos, self.npixs, self.tile_widths, self.tile_index
|
|
150
|
+
)
|
|
151
|
+
nvtx.end_range(rng)
|
|
152
|
+
|
|
153
|
+
rng = nvtx.start_range(message="sort_index")
|
|
154
|
+
self.sort_index = cp.argsort(
|
|
155
|
+
self.tile_index[:, 2] * npix_x * npix_y
|
|
156
|
+
+ self.tile_index[:, 1] * npix_x
|
|
157
|
+
+ self.tile_index[:, 0]
|
|
158
|
+
)
|
|
159
|
+
nvtx.end_range(rng)
|
|
160
|
+
|
|
161
|
+
unsort_index = cp.zeros(self.Np, dtype=int)
|
|
162
|
+
unsort_index[self.sort_index] = cp.arange(self.Np)
|
|
163
|
+
self.unsort_index = unsort_index
|
|
164
|
+
|
|
165
|
+
self.tile_index = self.tile_index[self.sort_index, :]
|
|
166
|
+
|
|
167
|
+
rng = nvtx.start_range(message="get_tile_information")
|
|
168
|
+
get_tile_information[self.blocks_1d, self.threadsperblock](
|
|
169
|
+
self.tile_index, self.particles_per_tile, self.start_index_for_tile
|
|
170
|
+
)
|
|
171
|
+
nvtx.end_range(rng)
|
|
172
|
+
|
|
173
|
+
def accumulate_per_tile(self, quantity):
|
|
174
|
+
"""
|
|
175
|
+
This function accumulates a variable per each tile
|
|
176
|
+
"""
|
|
177
|
+
|
|
178
|
+
tile_index = self.tile_index
|
|
179
|
+
|
|
180
|
+
npix_x = int(self.npixs[0])
|
|
181
|
+
npix_y = int(self.npixs[1])
|
|
182
|
+
npix_z = int(self.npixs[2])
|
|
183
|
+
|
|
184
|
+
quantity_per_tile = cp.zeros((npix_x, npix_y, npix_z), dtype=cp.float64)
|
|
185
|
+
|
|
186
|
+
accumulate_quantity_per_tile[self.blocks_1d, self.threadsperblock](
|
|
187
|
+
tile_index, quantity_per_tile, quantity
|
|
188
|
+
)
|
|
189
|
+
|
|
190
|
+
return quantity_per_tile
|
|
191
|
+
|
|
192
|
+
def findmax_per_tile(self, quantity):
|
|
193
|
+
"""
|
|
194
|
+
This function finds max of a variable per each tile
|
|
195
|
+
"""
|
|
196
|
+
|
|
197
|
+
tile_index = self.tile_index
|
|
198
|
+
|
|
199
|
+
npix_x = int(self.npixs[0])
|
|
200
|
+
npix_y = int(self.npixs[1])
|
|
201
|
+
npix_z = int(self.npixs[2])
|
|
202
|
+
|
|
203
|
+
quantity_per_tile = cp.zeros((npix_x, npix_y, npix_z), dtype=cp.float64)
|
|
204
|
+
|
|
205
|
+
findmax_quantity_per_tile[self.blocks_1d, self.threadsperblock](
|
|
206
|
+
tile_index, quantity_per_tile, quantity
|
|
207
|
+
)
|
|
208
|
+
|
|
209
|
+
return quantity_per_tile
|
|
210
|
+
|
|
211
|
+
def compactify_grid(self, Nmax):
|
|
212
|
+
occupancy_arr = ((self.particles_per_tile + (Nmax - 1)) // Nmax).flatten()
|
|
213
|
+
cumulative_occupancy = cp.cumsum(occupancy_arr)
|
|
214
|
+
numBlocksCompactGrid = cumulative_occupancy[-1]
|
|
215
|
+
compactGrid = cp.zeros((int(numBlocksCompactGrid), 3), dtype=int)
|
|
216
|
+
|
|
217
|
+
threadsperblock = 256
|
|
218
|
+
numBlocksFullGrid = int(self.npixs[0] * self.npixs[1] * self.npixs[2])
|
|
219
|
+
blocks_1d = (numBlocksFullGrid + (threadsperblock - 1)) // threadsperblock
|
|
220
|
+
compactify_kernel[blocks_1d, threadsperblock](
|
|
221
|
+
occupancy_arr,
|
|
222
|
+
cumulative_occupancy,
|
|
223
|
+
self.npixs,
|
|
224
|
+
Nmax,
|
|
225
|
+
self.particles_per_tile.flatten(),
|
|
226
|
+
self.start_index_for_tile.flatten(),
|
|
227
|
+
numBlocksCompactGrid,
|
|
228
|
+
compactGrid,
|
|
229
|
+
)
|
|
230
|
+
# compactGrid is an array of 3-tuples numBlocksCompactGrid long.
|
|
231
|
+
# for each block in numBlocksCompactGrid the 3-tuples is as follows
|
|
232
|
+
# [id of the tile in the grid it refers to, id of first particle, num of particles it contains]
|
|
233
|
+
|
|
234
|
+
self.compactGrid = compactGrid
|
|
235
|
+
|
|
236
|
+
def release_gpu_memory(self):
|
|
237
|
+
# TODO: Add deletion of all GPU variables stored in self
|
|
238
|
+
|
|
239
|
+
if self.sort_index is not None:
|
|
240
|
+
del self.sort_index
|
|
241
|
+
del self.unsort_index
|
|
242
|
+
del self.start_index_for_tile
|
|
243
|
+
del self.particles_per_tile
|
|
244
|
+
del self.tile_index
|
|
245
|
+
del self._pos
|
|
246
|
+
|
|
247
|
+
# cp._default_memory_pool.free_all_blocks()
|
|
@@ -0,0 +1,163 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
import cupy as cp
|
|
3
|
+
from numba import cuda
|
|
4
|
+
import math
|
|
5
|
+
import paicos as pa
|
|
6
|
+
import nvtx
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
from .cartesian_tiling import CartesianTiling
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class DataGpuInit:
|
|
13
|
+
""" """
|
|
14
|
+
|
|
15
|
+
def __init__(
|
|
16
|
+
self, snap, center, widths, orientation=None, threadsperblock=256, **kwargs
|
|
17
|
+
):
|
|
18
|
+
""" """
|
|
19
|
+
|
|
20
|
+
# self.__dict__.update(kwargs)
|
|
21
|
+
|
|
22
|
+
self.snap = snap
|
|
23
|
+
self.cartesian = True
|
|
24
|
+
if pa.settings.use_units:
|
|
25
|
+
self.code_length = code_length = self.snap.length
|
|
26
|
+
# self.npix = npix
|
|
27
|
+
self.threadsperblock = threadsperblock
|
|
28
|
+
|
|
29
|
+
if hasattr(center, "unit"):
|
|
30
|
+
self.center = center.copy
|
|
31
|
+
assert center.unit == code_length.unit, "this restriction applies"
|
|
32
|
+
elif pa.settings.use_units:
|
|
33
|
+
self.center = np.array(center) * code_length
|
|
34
|
+
else:
|
|
35
|
+
self.center = np.array(center)
|
|
36
|
+
|
|
37
|
+
if hasattr(widths, "unit"):
|
|
38
|
+
self.widths = widths.copy
|
|
39
|
+
assert widths.unit == code_length.unit, "this restriction applies"
|
|
40
|
+
elif pa.settings.use_units:
|
|
41
|
+
self.widths = np.array(widths) * code_length
|
|
42
|
+
else:
|
|
43
|
+
self.widths = np.array(widths)
|
|
44
|
+
|
|
45
|
+
if orientation is None:
|
|
46
|
+
self.orientation = None
|
|
47
|
+
else:
|
|
48
|
+
self.orientation = orientation.copy
|
|
49
|
+
|
|
50
|
+
# self.pos = self.snap["0_Coordinates"]
|
|
51
|
+
|
|
52
|
+
self.gpu_variables = {}
|
|
53
|
+
|
|
54
|
+
def _do_region_selection(self, thickness, pos):
|
|
55
|
+
|
|
56
|
+
center = self.center
|
|
57
|
+
widths = self.widths
|
|
58
|
+
snap = self.snap
|
|
59
|
+
|
|
60
|
+
# Send subset of snapshot to GPU
|
|
61
|
+
# get the index of the region of projection
|
|
62
|
+
rng = nvtx.start_range(message="region_selection")
|
|
63
|
+
if self.orientation is None:
|
|
64
|
+
get_index = pa.util.get_index_of_cubic_region_plus_thin_layer
|
|
65
|
+
indices = get_index(pos, center, widths, thickness, snap.box)
|
|
66
|
+
else:
|
|
67
|
+
get_index = pa.util.get_index_of_rotated_cubic_region_plus_thin_layer
|
|
68
|
+
indices = get_index(
|
|
69
|
+
pos, center, widths, thickness, snap.box, self.orientation
|
|
70
|
+
)
|
|
71
|
+
nvtx.end_range(rng)
|
|
72
|
+
|
|
73
|
+
return indices
|
|
74
|
+
|
|
75
|
+
def _send_variable_to_gpu(self, variable, gpu_key="input_variable", sort=False):
|
|
76
|
+
if isinstance(variable, str):
|
|
77
|
+
# variable_str = str(variable)
|
|
78
|
+
# err_msg = 'filter only works on gas'
|
|
79
|
+
# assert int(variable[0]) == 0, err_msg
|
|
80
|
+
variable = self.snap[variable]
|
|
81
|
+
else:
|
|
82
|
+
# variable_str = gpu_key
|
|
83
|
+
if not isinstance(variable, np.ndarray):
|
|
84
|
+
raise RuntimeError("Unexpected type for variable")
|
|
85
|
+
|
|
86
|
+
variable_str = gpu_key
|
|
87
|
+
# select only region of interest
|
|
88
|
+
if variable.shape[0] == self.index.shape[0]:
|
|
89
|
+
variable = variable[self.index]
|
|
90
|
+
|
|
91
|
+
if variable_str in self.gpu_variables and variable_str != gpu_key:
|
|
92
|
+
pass
|
|
93
|
+
else:
|
|
94
|
+
# Send variable to gpu
|
|
95
|
+
if pa.settings.use_units:
|
|
96
|
+
self.gpu_variables[variable_str] = cp.array(variable.value)
|
|
97
|
+
else:
|
|
98
|
+
self.gpu_variables[variable_str] = cp.array(variable)
|
|
99
|
+
|
|
100
|
+
if sort:
|
|
101
|
+
try:
|
|
102
|
+
if (
|
|
103
|
+
self.gpu_variables[variable_str].shape[0]
|
|
104
|
+
== self.tile.sort_index.shape[0]
|
|
105
|
+
):
|
|
106
|
+
self.gpu_variables[variable_str] = self.gpu_variables[variable_str][
|
|
107
|
+
self.tile.sort_index
|
|
108
|
+
]
|
|
109
|
+
except RuntimeError:
|
|
110
|
+
raise RuntimeError(
|
|
111
|
+
"Sorted can be done only after creating a CartesianTiling"
|
|
112
|
+
)
|
|
113
|
+
|
|
114
|
+
if isinstance(variable, pa.units.PaicosQuantity):
|
|
115
|
+
unit_quantity = variable.unit_quantity
|
|
116
|
+
else:
|
|
117
|
+
unit_quantity = None
|
|
118
|
+
|
|
119
|
+
return variable_str, unit_quantity
|
|
120
|
+
|
|
121
|
+
def _rotate_coordinates(self):
|
|
122
|
+
|
|
123
|
+
if self.orientation is not None:
|
|
124
|
+
if "inverse_rotation_matrix" not in self.gpu_variables:
|
|
125
|
+
self.gpu_variables["inverse_rotation_matrix"] = cp.array(
|
|
126
|
+
self.orientation.inverse_rotation_matrix
|
|
127
|
+
)
|
|
128
|
+
|
|
129
|
+
self.gpu_variables["pos"] = cp.matmul(
|
|
130
|
+
self.gpu_variables["inverse_rotation_matrix"],
|
|
131
|
+
self.gpu_variables["pos"],
|
|
132
|
+
axes=[(-2, -1), (-1, -2), (-1, -2)],
|
|
133
|
+
)
|
|
134
|
+
|
|
135
|
+
self.gpu_variables["center"] = cp.matmul(
|
|
136
|
+
self.gpu_variables["inverse_rotation_matrix"],
|
|
137
|
+
self.gpu_variables["center"],
|
|
138
|
+
)
|
|
139
|
+
|
|
140
|
+
def __del__(self):
|
|
141
|
+
"""
|
|
142
|
+
Clean up like this? Not sure it is needed...
|
|
143
|
+
"""
|
|
144
|
+
self.release_gpu_memory()
|
|
145
|
+
|
|
146
|
+
def release_gpu_memory(self):
|
|
147
|
+
# TODO: Add deletion of all GPU variables stored in self
|
|
148
|
+
if hasattr(self, "gpu_variables"):
|
|
149
|
+
for key in list(self.gpu_variables):
|
|
150
|
+
del self.gpu_variables[key]
|
|
151
|
+
del self.gpu_variables
|
|
152
|
+
if hasattr(self, "tile"):
|
|
153
|
+
self.tile.release_gpu_memory()
|
|
154
|
+
del self.tile
|
|
155
|
+
|
|
156
|
+
# cp._default_memory_pool.free_all_blocks()
|
|
157
|
+
|
|
158
|
+
def __enter__(self):
|
|
159
|
+
return self
|
|
160
|
+
|
|
161
|
+
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
162
|
+
print("bye")
|
|
163
|
+
self.__del__()
|
|
@@ -0,0 +1,67 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
import cupy as cp
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
def volume_integral(snap, qt, indices):
|
|
6
|
+
if isinstance(qt, str):
|
|
7
|
+
qt = snap[qt]
|
|
8
|
+
|
|
9
|
+
# using paicos omp-vectorized routines
|
|
10
|
+
a = snap.get_sum_of_array((qt * snap["0_Volume"])[indices])
|
|
11
|
+
|
|
12
|
+
return a
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def volume_average(snap, qt, indices):
|
|
16
|
+
if isinstance(qt, str):
|
|
17
|
+
qt = snap[qt]
|
|
18
|
+
|
|
19
|
+
# using paicos omp-vectorized routines
|
|
20
|
+
a = snap.get_sum_of_array((qt * snap["0_Volume"])[indices])
|
|
21
|
+
vol = snap.get_sum_of_array((snap["0_Volume"])[indices])
|
|
22
|
+
|
|
23
|
+
return a / vol
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def extract_turbulent_var(
|
|
27
|
+
snap,
|
|
28
|
+
sf,
|
|
29
|
+
variable,
|
|
30
|
+
filter_length,
|
|
31
|
+
weight,
|
|
32
|
+
filter_type="gaussian",
|
|
33
|
+
iterative=False,
|
|
34
|
+
selection=None,
|
|
35
|
+
):
|
|
36
|
+
|
|
37
|
+
filt_var = sf.filter_variable(
|
|
38
|
+
variable,
|
|
39
|
+
filter_length,
|
|
40
|
+
weight=weight,
|
|
41
|
+
filter_type=filter_type,
|
|
42
|
+
iterative=iterative,
|
|
43
|
+
selection=selection,
|
|
44
|
+
)
|
|
45
|
+
|
|
46
|
+
print(
|
|
47
|
+
"min/max/avg occupancy cartesian tiling %d / %d / %.2f"
|
|
48
|
+
% (
|
|
49
|
+
sf.tile.particles_per_tile.min(),
|
|
50
|
+
sf.tile.particles_per_tile.max(),
|
|
51
|
+
cp.sum(sf.tile.particles_per_tile)
|
|
52
|
+
/ cp.prod(sf.tile.particles_per_tile.shape),
|
|
53
|
+
)
|
|
54
|
+
)
|
|
55
|
+
if isinstance(variable, str):
|
|
56
|
+
variable = snap[variable]
|
|
57
|
+
|
|
58
|
+
smoothVar = np.zeros_like(variable)
|
|
59
|
+
turbVar = np.zeros_like(variable)
|
|
60
|
+
|
|
61
|
+
smoothVar[sf.index] = filt_var
|
|
62
|
+
turbVar[sf.index] = variable[sf.index] - filt_var
|
|
63
|
+
turbVar[~sf.indicesFirstPass * sf.index] = (
|
|
64
|
+
0 # this is for the thin buffer, which we set to zero
|
|
65
|
+
)
|
|
66
|
+
|
|
67
|
+
return smoothVar, turbVar
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: turbocluster
|
|
3
|
+
Version: 0.1
|
|
4
|
+
Summary: A GPU-accelerated Python library for filtering turbulence in Arepo simulations
|
|
5
|
+
Author-email: Lorenzo Maria Perrone <lorenzolmp@protonmail.com>, Thomas Berlok <tberlok@gmail.com>
|
|
6
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
7
|
+
Requires-Python: >=3.14.0
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
Requires-Dist: scipy>=1.7.3
|
|
10
|
+
Requires-Dist: numpy<=2.4.6
|
|
11
|
+
Requires-Dist: matplotlib
|
|
12
|
+
Requires-Dist: h5py
|
|
13
|
+
Requires-Dist: astropy>=4.3
|
|
14
|
+
Requires-Dist: numba-cuda[cu13]
|
|
15
|
+
Requires-Dist: paicos[dev]>=0.1.17
|
|
16
|
+
Requires-Dist: cupy-cuda13x[ctk]
|
|
17
|
+
Requires-Dist: nvtx
|
|
18
|
+
|
|
19
|
+
[](https://turbocluster.readthedocs.io/en/latest/)
|
|
20
|
+
[](
|
|
21
|
+
https://github.com/LorenzoLMP/turbocluster/actions/workflows/flake8.yml)
|
|
22
|
+
|
|
23
|
+
# turbocluster
|
|
24
|
+
|
|
25
|
+
A GPU-accelerated Python library for filtering turbulence in Arepo simulations.
|
|
26
|
+
|
|
27
|
+
<img src="images/turbocluster_banner.png" width="auto">
|
|
28
|
+
|
|
29
|
+
**NOTE**: This project is in active development and not officially released publicly yet. Expect kinks and bugs. Refer to [this arXiv paper](https://arxiv.org/pdf/2601.06250) if you would like to know more about its use.
|
|
30
|
+
|
|
31
|
+
A work-in-progress documentation can be found at [readthedocs](https://turbocluster.readthedocs.io/en/latest/).
|
|
32
|
+
|
|
33
|
+
# TODO:
|
|
34
|
+
- double-check tests and examples
|
|
35
|
+
- can cuda.jit() be replaced by jit() on non-cuda hardware?
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
requirements.txt
|
|
4
|
+
tests/test_cuda.py
|
|
5
|
+
turbocluster/__init__.py
|
|
6
|
+
turbocluster/cartesian_tiling.py
|
|
7
|
+
turbocluster/data_init.py
|
|
8
|
+
turbocluster/helper_functions.py
|
|
9
|
+
turbocluster.egg-info/PKG-INFO
|
|
10
|
+
turbocluster.egg-info/SOURCES.txt
|
|
11
|
+
turbocluster.egg-info/dependency_links.txt
|
|
12
|
+
turbocluster.egg-info/requires.txt
|
|
13
|
+
turbocluster.egg-info/top_level.txt
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
turbocluster
|