copied the script gather_tiles (needs to be checked)
This commit is contained in:
parent
bdf38e6dd0
commit
39e101c2fa
1 changed files with 96 additions and 0 deletions
96
scripts/gather_tiles.py
Normal file
96
scripts/gather_tiles.py
Normal file
|
@ -0,0 +1,96 @@
|
|||
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")
|
Loading…
Add table
Add a link
Reference in a new issue