Skip to content

Commit 2738896

Browse files
author
sambit-giri
committed
CubeMap calculation optimised
1 parent 2da2230 commit 2738896

4 files changed

Lines changed: 258 additions & 249 deletions

File tree

‎src/tools21cm/ViteBetti.py‎

Lines changed: 186 additions & 138 deletions
Original file line numberDiff line numberDiff line change
@@ -1,147 +1,195 @@
1-
import numpy as np
2-
3-
def CubeMap_torch(arr, multi_marker=True):
4-
import torch
5-
6-
# Check if GPU is available and choose the device
7-
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
8-
9-
# Convert the input array to a PyTorch tensor and move it to the selected device
10-
arr_tensor = torch.tensor(arr, dtype=torch.int32, device=device)
11-
12-
nx, ny, nz = arr_tensor.shape
13-
Nx, Ny, Nz = 2*nx, 2*ny, 2*nz
14-
cubemap = torch.zeros((Nx, Ny, Nz), dtype=torch.int32, device=device)
15-
16-
# Define markers
17-
markers = (1, 1, 1, 1)
18-
if multi_marker:
19-
markers = (1, 2, 3, 4)
20-
21-
# Set the vertices
22-
cubemap[2*arr_tensor.bool()] = markers[0]
23-
24-
# Convert indices for easier manipulation
25-
i_indices = torch.arange(Nx, device=device)
26-
j_indices = torch.arange(Ny, device=device)
27-
k_indices = torch.arange(Nz, device=device)
28-
29-
i_indices, j_indices, k_indices = torch.meshgrid(i_indices, j_indices, k_indices, indexing='ij')
30-
31-
# Define a mask to select non-zero values
32-
mask = cubemap == 0
33-
34-
# Edges
35-
edge_mask_1 = (cubemap[(i_indices-1) % Nx, j_indices, k_indices] == markers[0]) & (cubemap[(i_indices+1) % Nx, j_indices, k_indices] == markers[0])
36-
edge_mask_2 = (cubemap[i_indices, (j_indices-1) % Ny, k_indices] == markers[0]) & (cubemap[i_indices, (j_indices+1) % Ny, k_indices] == markers[0])
37-
edge_mask_3 = (cubemap[i_indices, j_indices, (k_indices-1) % Nz] == markers[0]) & (cubemap[i_indices, j_indices, (k_indices+1) % Nz] == markers[0])
38-
39-
cubemap[mask & (edge_mask_1 | edge_mask_2 | edge_mask_3)] = markers[1]
40-
41-
# Faces
42-
face_mask_1 = (cubemap[(i_indices-1) % Nx, j_indices, k_indices] == markers[1]) & (cubemap[(i_indices+1) % Nx, j_indices, k_indices] == markers[1]) & \
43-
(cubemap[i_indices, (j_indices-1) % Ny, k_indices] == markers[1]) & (cubemap[i_indices, (j_indices+1) % Ny, k_indices] == markers[1])
44-
face_mask_2 = (cubemap[i_indices, (j_indices-1) % Ny, k_indices] == markers[1]) & (cubemap[i_indices, (j_indices+1) % Ny, k_indices] == markers[1]) & \
45-
(cubemap[i_indices, j_indices, (k_indices-1) % Nz] == markers[1]) & (cubemap[i_indices, j_indices, (k_indices+1) % Nz] == markers[1])
46-
face_mask_3 = (cubemap[i_indices, j_indices, (k_indices-1) % Nz] == markers[1]) & (cubemap[i_indices, j_indices, (k_indices+1) % Nz] == markers[1]) & \
47-
(cubemap[(i_indices-1) % Nx, j_indices, k_indices] == markers[1]) & (cubemap[(i_indices+1) % Nx, j_indices, k_indices] == markers[1])
48-
49-
cubemap[mask & (face_mask_1 | face_mask_2 | face_mask_3)] = markers[2]
50-
51-
# Cubes
52-
cube_mask = (cubemap[(i_indices-1) % Nx, j_indices, k_indices] == markers[2]) & (cubemap[(i_indices+1) % Nx, j_indices, k_indices] == markers[2]) & \
53-
(cubemap[i_indices, (j_indices-1) % Ny, k_indices] == markers[2]) & (cubemap[i_indices, (j_indices+1) % Ny, k_indices] == markers[2]) & \
54-
(cubemap[i_indices, j_indices, (k_indices-1) % Nz] == markers[2]) & (cubemap[i_indices, j_indices, (k_indices+1) % Nz] == markers[2])
55-
56-
cubemap[mask & cube_mask] = markers[3]
57-
58-
# Move the result back to the CPU and convert to numpy array
59-
return cubemap.cpu().numpy()
1+
# ViteBetti.py
602

3+
import numpy as np
4+
import os
5+
6+
# --- Optional Numba Support ---
7+
try:
8+
from numba import jit
9+
numba_available = True
10+
except ImportError:
11+
numba_available = False
12+
# Define a dummy decorator if numba is not available
13+
def jit(func, *args, **kwargs):
14+
return func
15+
16+
# --- Optional Cython Support ---
17+
try:
18+
from .ViteBetti_cython import CubeMap as cython_CubeMap
19+
cython_available = True
20+
except ImportError:
21+
cython_CubeMap = None
22+
cython_available = False
23+
24+
# --- Optional Joblib Support ---
25+
try:
26+
from joblib import Parallel, delayed
27+
from multiprocessing import shared_memory
28+
joblib_available = True
29+
except ImportError:
30+
joblib_available = False
31+
Parallel, delayed, shared_memory = None, None, None
32+
33+
# --- Core Algorithm (Pure Python) ---
6134
def CubeMap(arr, multi_marker=True):
62-
nx, ny, nz = arr.shape
63-
Nx, Ny, Nz = 2*nx,2*ny,2*nz#2*nx-1,2*ny-1,2*nz-1
64-
cubemap = np.zeros((Nx,Ny,Nz))
65-
markers = 1, 1, 1, 1
66-
if multi_marker: markers = 1, 2, 3, 4
67-
## Vertices
68-
for i in range(nx):
69-
for j in range(ny):
70-
for k in range(nz):
71-
if arr[i,j,k]: cubemap[i*2,j*2,k*2] = markers[0]
72-
73-
## Edges
74-
for i in range(Nx):
75-
for j in range(Ny):
76-
for k in range(Nz):
77-
if cubemap[i,j,k] == 0:
78-
if cubemap[(i-1),j,k]==markers[0] and cubemap[(i+1)%Nx,j,k]==markers[0]: cubemap[i,j,k] = markers[1]
79-
elif cubemap[i,(j-1),k]==markers[0] and cubemap[i,(j+1)%Ny,k]==markers[0]: cubemap[i,j,k] = markers[1]
80-
elif cubemap[i,j,(k-1)]==markers[0] and cubemap[i,j,(k+1)%Nz]==markers[0]: cubemap[i,j,k] = markers[1]
81-
82-
## Faces
83-
for i in range(Nx):
84-
for j in range(Ny):
85-
for k in range(Nz):
86-
if cubemap[i,j,k] == 0:
87-
if cubemap[(i-1),j,k]==markers[1] and cubemap[(i+1)%Nx,j,k]==markers[1] and cubemap[i,(j-1),k]==markers[1] and cubemap[i,(j+1)%Ny,k]==markers[1]: cubemap[i,j,k] = markers[2]
88-
elif cubemap[i,(j-1),k]==markers[1] and cubemap[i,(j+1)%Ny,k]==markers[1] and cubemap[i,j,(k-1)]==markers[1] and cubemap[i,j,(k+1)%Nz]==markers[1]: cubemap[i,j,k] = markers[2]
89-
elif cubemap[i,j,(k-1)]==markers[1] and cubemap[i,j,(k+1)%Nz]==markers[1] and cubemap[(i-1),j,k]==markers[1] and cubemap[(i+1)%Nx,j,k]==markers[1]: cubemap[i,j,k] = markers[2]
90-
91-
## Cubes
92-
for i in range(Nx):
93-
for j in range(Ny):
94-
for k in range(Nz):
95-
if cubemap[i,j,k] == 0:
96-
if cubemap[(i-1),j,k]==markers[2] and cubemap[(i+1)%Nx,j,k]==markers[2]:
97-
if cubemap[i,(j-1),k]==markers[2] and cubemap[i,(j+1)%Ny,k]==markers[2]:
98-
if cubemap[i,j,(k-1)]==markers[2] and cubemap[i,j,(k+1)%Nz]==markers[2]: cubemap[i,j,k] = markers[3]
99-
100-
return cubemap
101-
102-
def EulerCharacteristic_seq(A):
103-
chi = 0
104-
nx,ny,nz = A.shape
105-
for x in range(nx):
106-
for y in range(ny):
107-
for z in range(nz):
108-
if(A[x,y,z] == 1):
109-
if (x+y+z)%2 == 0: chi += 1
110-
else: chi -= 1
111-
return chi
112-
113-
114-
import jax
115-
import jax.numpy as jnp
116-
117-
@jax.jit
118-
def CubeMap_jax(arr, multi_marker=True):
35+
"""
36+
Generates a cubical complex map from a binary 3D array.
37+
This is the pure Python implementation which serves as a fallback.
38+
"""
39+
nx, ny, nz = arr.shape
40+
Nx, Ny, Nz = 2 * nx, 2 * ny, 2 * nz
41+
cubemap = np.zeros((Nx, Ny, Nz), dtype=np.int32)
42+
11943
markers = (1, 1, 1, 1)
12044
if multi_marker:
12145
markers = (1, 2, 3, 4)
12246

123-
nx, ny, nz = arr.shape
124-
Nx, Ny, Nz = 2*nx, 2*ny, 2*nz
125-
cubemap = jnp.zeros((Nx, Ny, Nz), dtype=jnp.int32)
47+
# Vertices (1)
48+
for i in range(nx):
49+
for j in range(ny):
50+
for k in range(nz):
51+
if arr[i, j, k]:
52+
cubemap[i * 2, j * 2, k * 2] = markers[0]
53+
54+
# Edges (2)
55+
for i in range(Nx):
56+
for j in range(Ny):
57+
for k in range(Nz):
58+
if cubemap[i, j, k] == 0:
59+
# Check for neighbors along each axis using explicit periodic boundaries
60+
if cubemap[(i - 1) % Nx, j, k] == markers[0] and cubemap[(i + 1) % Nx, j, k] == markers[0]:
61+
cubemap[i, j, k] = markers[1]
62+
elif cubemap[i, (j - 1) % Ny, k] == markers[0] and cubemap[i, (j + 1) % Ny, k] == markers[0]:
63+
cubemap[i, j, k] = markers[1]
64+
elif cubemap[i, j, (k - 1) % Nz] == markers[0] and cubemap[i, j, (k + 1) % Nz] == markers[0]:
65+
cubemap[i, j, k] = markers[1]
66+
67+
# Faces (3)
68+
for i in range(Nx):
69+
for j in range(Ny):
70+
for k in range(Nz):
71+
if cubemap[i, j, k] == 0:
72+
if (cubemap[(i - 1) % Nx, j, k] == markers[1] and cubemap[(i + 1) % Nx, j, k] == markers[1] and
73+
cubemap[i, (j - 1) % Ny, k] == markers[1] and cubemap[i, (j + 1) % Ny, k] == markers[1]):
74+
cubemap[i, j, k] = markers[2]
75+
elif (cubemap[i, (j - 1) % Ny, k] == markers[1] and cubemap[i, (j + 1) % Ny, k] == markers[1] and
76+
cubemap[i, j, (k - 1) % Nz] == markers[1] and cubemap[i, j, (k + 1) % Nz] == markers[1]):
77+
cubemap[i, j, k] = markers[2]
78+
elif (cubemap[i, j, (k - 1) % Nz] == markers[1] and cubemap[i, j, (k + 1) % Nz] == markers[1] and
79+
cubemap[(i - 1) % Nx, j, k] == markers[1] and cubemap[(i + 1) % Nx, j, k] == markers[1]):
80+
cubemap[i, j, k] = markers[2]
81+
82+
# Cubes (4)
83+
for i in range(Nx):
84+
for j in range(Ny):
85+
for k in range(Nz):
86+
if cubemap[i, j, k] == 0:
87+
if (cubemap[(i - 1) % Nx, j, k] == markers[2] and cubemap[(i + 1) % Nx, j, k] == markers[2] and
88+
cubemap[i, (j - 1) % Ny, k] == markers[2] and cubemap[i, (j + 1) % Ny, k] == markers[2] and
89+
cubemap[i, j, (k - 1) % Nz] == markers[2] and cubemap[i, j, (k + 1) % Nz] == markers[2]):
90+
cubemap[i, j, k] = markers[3]
12691

127-
# Step 1: Vertices
128-
coords = jnp.argwhere(arr == 1)
129-
vert_indices = coords * 2
130-
cubemap = cubemap.at[tuple(vert_indices.T)].set(markers[0])
131-
132-
# Step 2: Edges (use jnp.roll for periodic neighbor checking)
133-
mask = cubemap == 0
134-
m0 = markers[0]
135-
m1 = markers[1]
136-
137-
def edge_mask_axis(cmap, axis):
138-
left = jnp.roll(cmap, 1, axis=axis)
139-
right = jnp.roll(cmap, -1, axis=axis)
140-
return (left == m0) & (right == m0)
141-
142-
edge_mask = edge_mask_axis(cubemap, 0) | edge_mask_axis(cubemap, 1) | edge_mask_axis(cubemap, 2)
143-
cubemap = jnp.where(mask & edge_mask, m1, cubemap)
92+
return cubemap
14493

145-
# Further steps: similar idea, but masks get more complex
94+
# --- Accelerated Versions ---
95+
CubeMap_numba = jit(CubeMap, nopython=True) if numba_available else None
96+
CubeMap_cython = cython_CubeMap # Assigned from the import attempt above
97+
98+
# --- Joblib Parallel Implementation ---
99+
def _CubeMap_joblib_worker(shm_name, arr_shape, i_start, i_end, arr_for_vertices, markers, stage):
100+
"""Worker function for joblib with the complete and correct logic."""
101+
existing_shm = shared_memory.SharedMemory(name=shm_name)
102+
cubemap = np.ndarray(arr_shape, dtype=np.int32, buffer=existing_shm.buf)
103+
104+
Nx, Ny, Nz = arr_shape
105+
nx, ny, nz = arr_for_vertices.shape
106+
107+
if stage == 'vertices':
108+
for i in range(i_start, i_end):
109+
for j in range(ny):
110+
for k in range(nz):
111+
if arr_for_vertices[i, j, k]:
112+
cubemap[i * 2, j * 2, k * 2] = markers[0]
113+
114+
elif stage == 'edges':
115+
for i in range(i_start, i_end):
116+
for j in range(Ny):
117+
for k in range(Nz):
118+
if cubemap[i, j, k] == 0:
119+
# **FIX**: Added all elif branches
120+
if cubemap[(i - 1) % Nx, j, k] == markers[0] and cubemap[(i + 1) % Nx, j, k] == markers[0]:
121+
cubemap[i, j, k] = markers[1]
122+
elif cubemap[i, (j - 1) % Ny, k] == markers[0] and cubemap[i, (j + 1) % Ny, k] == markers[0]:
123+
cubemap[i, j, k] = markers[1]
124+
elif cubemap[i, j, (k - 1) % Nz] == markers[0] and cubemap[i, j, (k + 1) % Nz] == markers[0]:
125+
cubemap[i, j, k] = markers[1]
126+
127+
elif stage == 'faces':
128+
for i in range(i_start, i_end):
129+
for j in range(Ny):
130+
for k in range(Nz):
131+
if cubemap[i, j, k] == 0:
132+
# **FIX**: Added all elif branches
133+
if (cubemap[(i - 1) % Nx, j, k] == markers[1] and cubemap[(i + 1) % Nx, j, k] == markers[1] and
134+
cubemap[i, (j - 1) % Ny, k] == markers[1] and cubemap[i, (j + 1) % Ny, k] == markers[1]):
135+
cubemap[i, j, k] = markers[2]
136+
elif (cubemap[i, (j - 1) % Ny, k] == markers[1] and cubemap[i, (j + 1) % Ny, k] == markers[1] and
137+
cubemap[i, j, (k - 1) % Nz] == markers[1] and cubemap[i, j, (k + 1) % Nz] == markers[1]):
138+
cubemap[i, j, k] = markers[2]
139+
elif (cubemap[i, j, (k - 1) % Nz] == markers[1] and cubemap[i, j, (k + 1) % Nz] == markers[1] and
140+
cubemap[(i - 1) % Nx, j, k] == markers[1] and cubemap[(i + 1) % Nx, j, k] == markers[1]):
141+
cubemap[i, j, k] = markers[2]
142+
143+
elif stage == 'cubes':
144+
for i in range(i_start, i_end):
145+
for j in range(Ny):
146+
for k in range(Nz):
147+
if cubemap[i, j, k] == 0:
148+
# **FIX**: This was already complete
149+
if (cubemap[(i - 1) % Nx, j, k] == markers[2] and cubemap[(i + 1) % Nx, j, k] == markers[2] and
150+
cubemap[i, (j - 1) % Ny, k] == markers[2] and cubemap[i, (j + 1) % Ny, k] == markers[2] and
151+
cubemap[i, j, (k - 1) % Nz] == markers[2] and cubemap[i, j, (k + 1) % Nz] == markers[2]):
152+
cubemap[i, j, k] = markers[3]
153+
154+
existing_shm.close()
155+
156+
157+
def CubeMap_joblib(arr, multi_marker=True, n_jobs=-1):
158+
"""
159+
Generates a cubical complex map from a binary 3D array using joblib for parallelism.
160+
"""
161+
if not joblib_available:
162+
raise ImportError("Joblib or its dependencies are not installed. Cannot use 'joblib' backend.")
146163

147-
return cubemap
164+
nx, ny, nz = arr.shape
165+
Nx, Ny, Nz = 2 * nx, 2 * ny, 2 * nz
166+
cubemap_shape = (Nx, Ny, Nz)
167+
markers = (1, 2, 3, 4) if multi_marker else (1, 1, 1, 1)
168+
169+
shm = shared_memory.SharedMemory(create=True, size=np.dtype(np.int32).itemsize * Nx * Ny * Nz)
170+
shared_cubemap = np.ndarray(cubemap_shape, dtype=np.int32, buffer=shm.buf)
171+
shared_cubemap[:] = 0
172+
173+
if n_jobs == -1:
174+
n_jobs = os.cpu_count()
175+
176+
try:
177+
for stage_name, domain_size in [('vertices', nx), ('edges', Nx), ('faces', Nx), ('cubes', Nx)]:
178+
chunk_size = (domain_size + n_jobs - 1) // n_jobs
179+
tasks = [
180+
delayed(_CubeMap_joblib_worker)(
181+
shm.name, cubemap_shape, i * chunk_size,
182+
min((i + 1) * chunk_size, domain_size),
183+
arr, markers, stage_name
184+
)
185+
for i in range(n_jobs)
186+
]
187+
Parallel(n_jobs=n_jobs)(tasks)
188+
189+
result_cubemap = np.copy(shared_cubemap)
190+
191+
finally:
192+
shm.close()
193+
shm.unlink()
194+
195+
return result_cubemap

0 commit comments

Comments
 (0)