convert snapshots

sbmy_control/scripts/convert_snapshot_to_density.py

@@ -36,6 +36,55 @@ def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0,
     print("Done.")
     
 
+def convert_snapshots_to_density(snapshot_base_path, output_path, N=None, corner=(0.0, 0.0, 0.0)):
+    """
+    Convert multiple snapshots to density fields.
+
+    Parameters
+    ----------
+    snapshot_base_path : str
+        Base path for the snapshot files.
+    output_path : str
+        Path to the output density file.
+    N : int
+        Size of the density field grid (N x N x N).
+    corner : tuple of float
+        Corner of the box (x, y, z).
+    """
+    # Get all snapshot files with path "snapshot_base_path*"
+    import glob
+    snapshot_files = glob.glob(snapshot_base_path + "*")
+    if not snapshot_files:
+        raise FileNotFoundError(f"No snapshot files found at {snapshot_base_path}")
+    
+    A = None
+    F = None
+    
+    for snapshot_path in snapshot_files:
+        # Read the snapshot
+        print(f"Reading snapshot {snapshot_path}...")
+        snap = read_snapshot(snapshot_path)
+
+        if N is None:
+            N = snap.Np0
+
+        # Calculate density
+        print("Calculating density...")
+        F = get_density_pm_snapshot(snap, N, N, N, corner[0], corner[1], corner[2])
+        
+        if A is None:
+            A = F.data
+        else:
+            A += F.data
+    
+    F.data = A # Combine densities from all snapshots
+    # Write density to file
+    print("Writing density...")
+    F.write(output_path)
+    print("Density written to", output_path)
+    print("Done.")
+
+
 def console_main():
     parser = argparse.ArgumentParser(description="Convert snapshot to density.")
     parser.add_argument(
@@ -70,12 +119,22 @@ def console_main():
 
     args = parser.parse_args()
 
+    if args.snapshot.endswith("h5") or args.snapshot.endswith("hdf5") or args.snapshot.endswith("gadget3"):
+        # If the snapshot is a single file, convert it to density
         convert_snapshot_to_density(
             snapshot_path=args.snapshot,
             output_path=args.output,
             N=args.N,
             corner=args.corner,
         )
+    else:
+        # If the snapshot is a base path, convert all snapshots to density
+        convert_snapshots_to_density(
+            snapshot_base_path=args.snapshot,
+            output_path=args.output,
+            N=args.N,
+            corner=args.corner,
+        )
 
 
 if __name__ == "__main__":