diff --git a/sbmy_control/scripts/convert_snapshot_to_density.py b/sbmy_control/scripts/convert_snapshot_to_density.py index 869ee90..f3935d9 100644 --- a/sbmy_control/scripts/convert_snapshot_to_density.py +++ b/sbmy_control/scripts/convert_snapshot_to_density.py @@ -1,9 +1,7 @@ -from pysbmy.density import get_density_pm_snapshot -from pysbmy.snapshot import read_snapshot import argparse -def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0, 0.0, 0.0)): +def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0, 0.0, 0.0), time=1.0): """ Convert a snapshot to a density field. @@ -18,6 +16,9 @@ def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0, corner : tuple of float Corner of the box (x, y, z). """ + from pysbmy.density import get_density_pm_snapshot + from pysbmy.snapshot import read_snapshot + # Read the snapshot print("Reading snapshot...") snap = read_snapshot(snapshot_path) @@ -28,6 +29,7 @@ def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0, # Calculate density print("Calculating density...") F = get_density_pm_snapshot(snap, N, N, N, corner[0], corner[1], corner[2]) + F.time = time # Set the time for the field # Write density to file print("Writing density...") @@ -36,7 +38,7 @@ 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)): +def convert_snapshots_to_density(snapshot_base_path, output_path, N=None, corner=(0.0, 0.0, 0.0), time=1.0): """ Convert multiple snapshots to density fields. @@ -51,33 +53,57 @@ def convert_snapshots_to_density(snapshot_base_path, output_path, N=None, corner corner : tuple of float Corner of the box (x, y, z). """ + from pysbmy.density import get_density_pm, density_to_delta + from pysbmy.snapshot import read_snapshot + from pysbmy.field import Field + import numpy as np + # 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 + L0 = L1 = L2 = None + d0 = d1 = d2 = 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 L0 is None: + # Get the grid parameters from the snapshot + L0 = snap.L0 + L1 = snap.L1 + L2 = snap.L2 + d0 = L0 / N + d1 = L1 / N + d2 = L2 / N + else: + # Ensure the grid parameters match the first snapshot + if not (snap.L0 == L0 and snap.L1 == L1 and snap.L2 == L2): + raise ValueError(f"All snapshots must have the same grid parameters. Got {snap.L0}, {snap.L1}, {snap.L2} but expected {L0}, {L1}, {L2}.") if A is None: - A = F.data - else: - A += F.data + # Initialize the density field + A = np.zeros((N, N, N), dtype=np.float32) + + # Calculate density + print("Calculating density for this snapshot...") + A = get_density_pm(snap.pos, A, d0=d0, d1=d1, d2=d2) - F.data = A # Combine densities from all snapshots + # Convert density to delta field + print("Converting density to delta field...") + A = density_to_delta(A) + + # Create a Field object + F = Field(L0=L0,L1=L1,L2=L2,corner0=corner[0],corner1=corner[1],corner2=corner[2],rank=1,N0=N,N1=N,N2=N,time=time,data=A) # Write density to file print("Writing density...") F.write(output_path)