96 lines
No EOL
3.4 KiB
Python
96 lines
No EOL
3.4 KiB
Python
import numpy as np
|
|
from pysbmy.snapshot import read_tile
|
|
from pysbmy.field import Field
|
|
from pysbmy.density import get_density_pm, density_to_delta
|
|
from pysbmy import c_float_t
|
|
import tqdm
|
|
import argparse
|
|
|
|
def get_indices(tile, N_TILES0, N_TILES1, N_TILES2):
|
|
"""Get the indices of the tile in a 3D grid."""
|
|
i = (tile // (N_TILES1 * N_TILES2))%N_TILES0
|
|
j = ((tile - N_TILES1 * N_TILES2 * i) // N_TILES2)%N_TILES1
|
|
k = (tile - N_TILES2 * j - N_TILES1 * N_TILES2 * i)%N_TILES2
|
|
return i, j, k
|
|
|
|
|
|
def tile_to_density_with_buffer(T, N, d, buffer):
|
|
"""
|
|
Convert a tile to density with a buffer.
|
|
"""
|
|
# Create a buffer for the density
|
|
A = np.zeros((N+2*buffer, N+2*buffer, N+2*buffer), dtype=np.float32)
|
|
|
|
# Get the density of the tile
|
|
A = get_density_pm(T.pos-T.corner_position.astype(c_float_t)+d*buffer, A, d)
|
|
return A
|
|
|
|
|
|
def add_density_tile(A,A_tile,corner_indices):
|
|
N = A.shape[0]
|
|
Np = A_tile.shape[0]
|
|
i0, j0, k0 = corner_indices
|
|
|
|
# Create index arrays for the tile
|
|
i_idx = (np.arange(Np) + i0) % N
|
|
j_idx = (np.arange(Np) + j0) % N
|
|
k_idx = (np.arange(Np) + k0) % N
|
|
|
|
# Use np.ix_ to get all combinations of indices
|
|
A[np.ix_(i_idx, j_idx, k_idx)] += A_tile
|
|
|
|
|
|
def gather_density(A, folder, tile_base, Np_tile, dpm, buffer, N_TILES):
|
|
"""
|
|
Gather the density from the tiles.
|
|
"""
|
|
for tile in tqdm.tqdm(range(N_TILES**3), desc="Density of tiles", unit="tiles"):
|
|
T=read_tile(folder+tile_base+str(tile+1)+".h5", int(Np_tile**3))
|
|
# Get the corner position of the tile
|
|
corner_position = T.corner_position
|
|
# Get the corner indices of the tile
|
|
i,j,k = get_indices(tile, N_TILES, N_TILES, N_TILES)
|
|
corner_grid_indices = (i*Np_tile-buffer, j*Np_tile-buffer, k*Np_tile-buffer)
|
|
# Get the density of the tile with a buffer
|
|
A_tile = tile_to_density_with_buffer(T, Np_tile, dpm, buffer)
|
|
# Add the density of the tile to the grid
|
|
add_density_tile(A, A_tile, corner_grid_indices)
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Gather density from tiles.")
|
|
parser.add_argument("-d","--folder", type=str, default="./", help="Folder containing the tiles")
|
|
parser.add_argument("--tile_base", type=str, default="sCOLA_tile", help="Base name of the tiles")
|
|
parser.add_argument("-L","--L", type=int, default=1920, help="Size of the box in Mpc/h")
|
|
parser.add_argument("-Np","--Np", type=int, default=80, help="Number of cells per dimension for the full box")
|
|
parser.add_argument("-Nt","--N_tiles", type=int, default=4, help="Number of tiles per dimension.")
|
|
parser.add_argument("--buffer", type=int, default=40, help="Buffer size for the density field of tiles")
|
|
args = parser.parse_args()
|
|
|
|
L = args.L
|
|
Np = args.Np
|
|
N_TILES = args.N_tiles
|
|
buffer = args.buffer
|
|
|
|
folder = args.folder
|
|
tile_base = args.tile_base
|
|
|
|
Np_tile = Np//N_TILES
|
|
dpm = L/Np_tile
|
|
|
|
print("Memory allocation for the grid...")
|
|
A=np.zeros((Np,Np,Np), dtype=np.float32)
|
|
|
|
print("Starting to read the tiles...")
|
|
gather_density(A, folder, tile_base, Np_tile, dpm, buffer, N_TILES)
|
|
|
|
print("Finished reading the tiles.")
|
|
A=density_to_delta(A,-1)
|
|
|
|
print("Converting to field...")
|
|
F=Field(L,L,L, 0.,0.,0., 1, Np,Np,Np, 1., A)
|
|
|
|
print("Saving field...")
|
|
F.write(folder+"../results/final_density_sCOLA.h5") |