79 lines
No EOL
1.9 KiB
Python
79 lines
No EOL
1.9 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.")
|
|
|
|
|
|
|
|
if __name__ == "__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()
|
|
|
|
convert_snapshot_to_density(
|
|
snapshot_path=args.snapshot,
|
|
output_path=args.output,
|
|
N=args.N,
|
|
corner=args.corner,
|
|
) |