From 39e101c2fad9b55e1ae4b04d1521ee85bea6a3d5 Mon Sep 17 00:00:00 2001
From: Mayeul Aubin <mayeul.aubin@iap.fr>
Date: Tue, 13 May 2025 15:40:55 +0200
Subject: [PATCH] copied the script gather_tiles (needs to be checked)

---
 scripts/gather_tiles.py | 96 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 96 insertions(+)
 create mode 100644 scripts/gather_tiles.py

diff --git a/scripts/gather_tiles.py b/scripts/gather_tiles.py
new file mode 100644
index 0000000..33ded24
--- /dev/null
+++ b/scripts/gather_tiles.py
@@ -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")
\ No newline at end of file