bugfix convert snapshots to density

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)