141 lines
3.8 KiB
Python
141 lines
3.8 KiB
Python
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)):
|
|
"""
|
|
Convert a snapshot to a density field.
|
|
|
|
Parameters
|
|
----------
|
|
snapshot_path : str
|
|
Path to the snapshot file.
|
|
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).
|
|
"""
|
|
# Read the snapshot
|
|
print("Reading snapshot...")
|
|
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])
|
|
|
|
# Write density to file
|
|
print("Writing density...")
|
|
F.write(output_path)
|
|
print("Density written to", output_path)
|
|
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(
|
|
"-S",
|
|
"--snapshot",
|
|
type=str,
|
|
required=True,
|
|
help="Path to the snapshot file.",
|
|
)
|
|
parser.add_argument(
|
|
"-o",
|
|
"--output",
|
|
type=str,
|
|
required=True,
|
|
help="Path to the output density file.",
|
|
)
|
|
parser.add_argument(
|
|
"-N",
|
|
"--N",
|
|
type=int,
|
|
default=None,
|
|
help="Size of the density field grid (N x N x N).",
|
|
)
|
|
parser.add_argument(
|
|
"-c",
|
|
"--corner",
|
|
type=float,
|
|
nargs=3,
|
|
default=[0.0, 0.0, 0.0],
|
|
help="Corner of the box (x, y, z).",
|
|
)
|
|
|
|
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__":
|
|
console_main()
|