aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorVolker Hoffmann <volker@cheleb.net>2015-11-20 06:59:31 +0100
committerVolker Hoffmann <volker@cheleb.net>2015-11-20 07:00:57 +0100
commit401b212a8d88141f6a54137480a73427af447e81 (patch)
tree2f23224a587ae62c3d5a8399d04b19d76571da23
parent6d64ba7e827f89336f604f2443e97a19b86c1ee8 (diff)
feat: script to extract initial PIDs that go into final particles
-rw-r--r--Post/extract_initial_pids.py53
1 files changed, 53 insertions, 0 deletions
diff --git a/Post/extract_initial_pids.py b/Post/extract_initial_pids.py
new file mode 100644
index 0000000..5306fa7
--- /dev/null
+++ b/Post/extract_initial_pids.py
@@ -0,0 +1,53 @@
+"""
+Extracts initial PIDs for sources.
+
+Usage:
+python ./extract_initial_pids.py --run_name run_01 > pid_list.csv
+
+Returns one line of source PIDs per line in output. The first PID in the
+list matches the PID of the final particle.
+"""
+
+import io_helpers as ioh
+import formation_helpers as fh
+import argparse
+import sys
+import numpy as np
+
+
+# Parse Arguments
+parser = argparse.ArgumentParser()
+parser.add_argument('-rname', '--run_name', \
+ default='run_01', \
+ help='Name of the Run.')
+args = parser.parse_args()
+
+# Load Collisions
+fname = "Collisions_%s.dat" % args.run_name
+dfc = ioh.read_collisions_and_stack([fname])
+
+# Load IC Output
+fname = "Out_%s_%012d.dat" % (args.run_name, 0)
+dfo_t0 = ioh.read_output_and_stack([fname], frame='heliocentric')
+dfo_t0 = dfo_t0[dfo_t0.mass < 12.0]
+
+# Load Final Output
+fname = "Out_%s_%012d.dat" % (args.run_name, 9e9)
+dfo_tf = ioh.read_output_and_stack([fname], frame='heliocentric')
+dfo_tf = dfo_tf[dfo_tf.mass < 12.0]
+
+# Extract Final PIDs
+final_pids = np.asarray(dfo_tf.pid, dtype=np.int64)
+sys.stderr.write("// %i Final Particles\n" % len(final_pids))
+
+# Loop Final PIDs
+sources_all = []
+for ipid, pid in enumerate(final_pids):
+ sys.stderr.write("// Finding Sources for PID %i (%i/%i)\n" % \
+ (pid, ipid+1, len(final_pids)))
+ sources = fh.return_sources(pid, dfc)
+ sources_all.append(sources)
+
+# Output
+for sources in sources_all:
+ print ",".join(map(str, sources))