defined a main function for the scripts so they can be also used as modules

scripts/convert_snapshot_to_density.py

@@ -2,6 +2,41 @@ 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(
@@ -36,19 +71,9 @@ if __name__ == "__main__":
 
     args = parser.parse_args()
 
-    # Read the snapshot
-    print("Reading snapshot...")
-    snap = read_snapshot(args.snapshot)
-    
-    if args.N is None:
-        N = snap.Np0
-    else:
-        N = args.N
-
-    print("Calculating density...")
-    F=get_density_pm_snapshot(snap, N,N,N, args.corner[0],args.corner[1],args.corner[2])
-
-    print("Writing density...")
-    F.write(args.output)
-    print("Density written to", args.output)
-    print("Done.")
+    convert_snapshot_to_density(
+        snapshot_path=args.snapshot,
+        output_path=args.output,
+        N=args.N,
+        corner=args.corner,
+    )