maubin/sbmy_control
0
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

bugfix convert snapshots to density

Browse source
This commit is contained in:
Mayeul Aubin 2025-07-01 15:15:58 +02:00
parent 18caa0d30a
commit b4913dd9ff
1 changed files with 40 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

54
sbmy_control/scripts/convert_snapshot_to_density.py
View file

@ -1,9 +1,7 @@
from pysbmy.density import get_density_pm_snapshot
from pysbmy.snapshot import read_snapshot
import argparse 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. 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 : tuple of float
Corner of the box (x, y, z). Corner of the box (x, y, z).
""" """
from pysbmy.density import get_density_pm_snapshot
from pysbmy.snapshot import read_snapshot
# Read the snapshot # Read the snapshot
print("Reading snapshot...") print("Reading snapshot...")
snap = read_snapshot(snapshot_path) 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 # Calculate density
print("Calculating density...") print("Calculating density...")
F = get_density_pm_snapshot(snap, N, N, N, corner[0], corner[1], corner[2]) 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 # Write density to file
print("Writing density...") print("Writing density...")
@ -36,7 +38,7 @@ def convert_snapshot_to_density(snapshot_path, output_path, N=None, corner=(0.0,
print("Done.") 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. 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 : tuple of float
Corner of the box (x, y, z). 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*" # Get all snapshot files with path "snapshot_base_path*"
import glob import glob
snapshot_files = glob.glob(snapshot_base_path + "*") snapshot_files = glob.glob(snapshot_base_path + "*")
if not snapshot_files: if not snapshot_files:
raise FileNotFoundError(f"No snapshot files found at {snapshot_base_path}") raise FileNotFoundError(f"No snapshot files found at {snapshot_base_path}")
A = None A = None
F = None L0 = L1 = L2 = None
d0 = d1 = d2 = None
for snapshot_path in snapshot_files: for snapshot_path in snapshot_files:
# Read the snapshot # Read the snapshot
print(f"Reading snapshot {snapshot_path}...") print(f"Reading snapshot {snapshot_path}...")
snap = read_snapshot(snapshot_path) snap = read_snapshot(snapshot_path)
if N is None: if N is None:
N = snap.Np0 N = snap.Np0
# Calculate density if L0 is None:
print("Calculating density...") # Get the grid parameters from the snapshot
F = get_density_pm_snapshot(snap, N, N, N, corner[0], corner[1], corner[2]) 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: if A is None:
A = F.data # Initialize the density field
else: A = np.zeros((N, N, N), dtype=np.float32)
A += F.data
# 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 # Write density to file
print("Writing density...") print("Writing density...")
F.write(output_path) F.write(output_path)