Add astro/21cm/{freq2z,z2freq,cube_mean}.py

Migrated from my other repo 21cmFAST/tools.

3 files changed, 116 insertions, 0 deletions

diff --git a/astro/21cm/cube_mean.py b/astro/21cm/cube_mean.py
new file mode 100755
index 0000000..9ff33d4
--- /dev/null
+++ b/astro/21cm/cube_mean.py
@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2017 Weitian LI <weitian@aaronly.me>
+# MIT License
+#
+
+"""
+Calculate the mean values of the cube.
+"""
+
+import os
+import sys
+import argparse
+
+import numpy as np
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Calculate the mean value of the data cube")
+    parser.add_argument("-d", "--dtype", default="float32",
+                        help="NumPy dtype of data cubes (default: float32)")
+    parser.add_argument("infiles", nargs="+", help="input data cubes")
+    args = parser.parse_args()
+
+    print("# filename: mean side_length")
+    for f in args.infiles:
+        cube = np.fromfile(open(f, "rb"), dtype=args.dtype)
+        sidelen = round(cube.shape[0] ** (1.0/3))
+        mean = cube.mean()
+        print("%s:\t%g\t\t%d" % (f, mean, sidelen))
+
+
+if __name__ == "__main__":
+    main()
diff --git a/astro/21cm/freq2z.py b/astro/21cm/freq2z.py
new file mode 120000
index 0000000..aa9eb7c
--- /dev/null
+++ b/astro/21cm/freq2z.py
@@ -0,0 +1 @@
+z2freq.py
\ No newline at end of file
diff --git a/astro/21cm/z2freq.py b/astro/21cm/z2freq.py
new file mode 100755
index 0000000..454a425
--- /dev/null
+++ b/astro/21cm/z2freq.py
@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+#
+# Copyright (c) 2017 Weitian LI <weitian@aaronly.me>
+# MIT License
+#
+
+"""
+Convert redshifts to the observed frequency of the 21 cm signal,
+and vice versa.
+"""
+
+import os
+import sys
+import argparse
+
+import numpy as np
+
+
+prog = os.path.basename(sys.argv[0])
+freq21cm = 1420405751.7667 / 1e6  # [MHz]
+
+
+def z2freq(redshifts, print_=False):
+    redshifts = np.asarray(redshifts)
+    freqs = freq21cm / (redshifts + 1.0)
+    if print_:
+        print("# redshift  frequency[MHz]")
+        for z, f in zip(redshifts, freqs):
+            print("%.4f  %.2f" % (z, f))
+    return freqs
+
+
+def freq2z(freqs, print_=False):
+    freqs = np.asarray(freqs)
+    redshifts = freq21cm / freqs - 1.0
+    if print_:
+        print("# frequency[MHz]  redshift")
+        for f, z in zip(freqs, redshifts):
+            print("%.2f  %.4f" % (f, z))
+    return redshifts
+
+
+def parse_inputs(inputs):
+    values = []
+    for inp in inputs:
+        try:
+            v = float(inp)
+            values += [v]
+        except ValueError:
+            begin, step, stop = inp.split(":")
+            begin, step, stop = float(begin), float(step), float(stop)
+            v = np.arange(start=begin, stop=stop+step, step=step)
+            values += list(v)
+    return values
+
+
+def main():
+    if prog == "z2freq.py":
+        description = "Convert redshifts to observed 21 cm frequencies"
+        parser = argparse.ArgumentParser(description=description)
+        parser.add_argument("inputs", nargs="+",
+                            help="input redshifts: <z1> <begin:step:stop>")
+        args = parser.parse_args()
+        redshifts = parse_inputs(args.inputs)
+        z2freq(redshifts, print_=True)
+    elif prog == "freq2z.py":
+        description = "Convert observed 21 cm frequencies to redshifts"
+        parser = argparse.ArgumentParser(description=description)
+        parser.add_argument("inputs", nargs="+",
+                            help="input frequencies [MHz]: <freq1> " +
+                            "<begin:step:stop>")
+        args = parser.parse_args()
+        freqs = parse_inputs(args.inputs)
+        freq2z(freqs, print_=True)
+    else:
+        raise RuntimeError("unknown program name: %s" % prog)
+
+
+if __name__ == "__main__":
+    main()