scripts/cosmo_calc.py: Rewrite arguments processing

* Allow specifying multiple quantities to be calculated
* Default to calculate all supported quantities

1 files changed, 94 insertions, 52 deletions

diff --git a/scripts/cosmo_calc.py b/scripts/cosmo_calc.py
index 8650d2a..1302bdb 100755
--- a/scripts/cosmo_calc.py
+++ b/scripts/cosmo_calc.py
@@ -7,12 +7,72 @@
 Cosmology calculator with support of Chandra ACIS-specific quantities.
 """
 
+import sys
 import argparse
+from collections import OrderedDict
 
 from _context import acispy
 from acispy.cosmo import Calculator
 
 
+# Supported quantities
+QUANTITIES = OrderedDict([
+    ("luminosity_distance", {
+        "unit": "Mpc",
+        "label": "Luminosity distance",
+        "kwargs": ["z", "unit"],
+    }),
+    ("angular_diameter_distance", {
+        "unit": "Mpc",
+        "label": "Angular diameter distance",
+        "kwargs": ["z", "unit"],
+    }),
+    ("kpc_per_arcsec", {
+        "unit": None,
+        "label": "kpc/arcsec (DA)",
+        "kwargs": ["z"],
+    }),
+    ("kpc_per_pix", {
+        "unit": None,
+        "label": "kpc/pix (DA)",
+        "kwargs": ["z"],
+    }),
+    ("cm_per_pix", {
+        "unit": None,
+        "label": "cm/pix (DA)",
+        "kwargs": ["z"],
+    }),
+    ("norm_apec", {
+        "unit": "cm^-5",
+        "label": "norm (APEC)",
+        "kwargs": ["z"],
+    }),
+])
+
+
+def get_quantities(args):
+    # convert ``argparse.Namespace`` to a dictionary
+    args = vars(args)
+    q_all = list(QUANTITIES.keys())
+    q_active = [q for q in q_all if args[q]]
+    if len(q_active) == 0:
+        q_active = q_all
+    return q_active
+
+
+def calc_quantity(q, calculator, args):
+    args = vars(args)
+    kwargs = {arg: args[arg] for arg in QUANTITIES[q]["kwargs"]
+              if args[arg] is not None}
+    value = getattr(calculator, q)(**kwargs)
+    label = QUANTITIES[q]["label"]
+    unit = args["unit"] if args["unit"] is not None else QUANTITIES[q]["unit"]
+    if args["brief"]:
+        print(value)
+    else:
+        print("%s: %s  # [%s]" % (label, value, unit))
+
+
 def main():
     parser = argparse.ArgumentParser(
         description="Cosmology calculator with Chandra-specific quantities")
@@ -27,64 +87,46 @@ def main():
                         help="be brief")
     parser.add_argument("-U", "--unit", dest="unit",
                         help="unit for output quantity if supported")
-    group = parser.add_mutually_exclusive_group(required=True)
-    group.add_argument("-L", "--luminosity-distance",
-                       dest="luminosity_distance",
-                       action="store_true",
-                       help="calculate the luminosity distance (DL)")
-    group.add_argument("-A", "--angular-diameter-distance",
-                       dest="angular_diameter_distance",
-                       action="store_true",
-                       help="calculate the angular diameter distance (DA)")
-    group.add_argument("--kpc-per-arcsec", dest="kpc_per_arcsec",
-                       action="store_true",
-                       help="calculate the transversal length [kpc] " +
-                       "w.r.t. 1 arcsec at DA(z)")
-    group.add_argument("--kpc-per-pix", dest="kpc_per_pix",
-                       action="store_true",
-                       help="calculate the transversal length [kpc] " +
-                       "w.r.t. 1 ACIS pixel (0.492 arcsec) at DA(z)")
-    group.add_argument("--cm-per-pix", dest="cm_per_pix",
-                       action="store_true",
-                       help="calculate the transversal length [cm] " +
-                       "w.r.t. 1 ACIS pixel (0.492 arcsec) at DA(z)")
-    group.add_argument("--norm-apec", dest="norm_apec",
-                       action="store_true",
-                       help="calculate the normalization factor " +
-                       "of the XSPEC APEC model assuming EM=1")
+    parser.add_argument("-L", "--luminosity-distance",
+                        dest="luminosity_distance",
+                        action="store_true",
+                        help="calculate the luminosity distance (DL)")
+    parser.add_argument("-A", "--angular-diameter-distance",
+                        dest="angular_diameter_distance",
+                        action="store_true",
+                        help="calculate the angular diameter distance (DA)")
+    parser.add_argument("--kpc-per-arcsec", dest="kpc_per_arcsec",
+                        action="store_true",
+                        help="calculate the transversal length [kpc] " +
+                        "w.r.t. 1 arcsec at DA(z)")
+    parser.add_argument("--kpc-per-pix", dest="kpc_per_pix",
+                        action="store_true",
+                        help="calculate the transversal length [kpc] " +
+                        "w.r.t. 1 ACIS pixel (0.492 arcsec) at DA(z)")
+    parser.add_argument("--cm-per-pix", dest="cm_per_pix",
+                        action="store_true",
+                        help="calculate the transversal length [cm] " +
+                        "w.r.t. 1 ACIS pixel (0.492 arcsec) at DA(z)")
+    parser.add_argument("--norm-apec", dest="norm_apec",
+                        action="store_true",
+                        help="calculate the normalization factor " +
+                        "of the XSPEC APEC model assuming EM=1")
     parser.add_argument("z", type=float, help="redshift")
     args = parser.parse_args()
 
     cosmocalc = Calculator(H0=args.H0, Om0=args.Om0)
 
-    if args.luminosity_distance:
-        kwargs = {"z": args.z}
-        kwargs["unit"] = args.unit if args.unit else "Mpc"
-        label = "Luminosity distance [%s]" % kwargs["unit"]
-        value = cosmocalc.luminosity_distance(**kwargs)
-    elif args.angular_diameter_distance:
-        kwargs = {"z": args.z}
-        kwargs["unit"] = args.unit if args.unit else "Mpc"
-        label = "Angular diameter distance [%s]" % kwargs["unit"]
-        value = cosmocalc.angular_diameter_distance(**kwargs)
-    elif args.kpc_per_arcsec:
-        label = "kpc/arcsec (DA)"
-        value = cosmocalc.kpc_per_arcsec(args.z)
-    elif args.kpc_per_pix:
-        label = "kpc/pix (DA)"
-        value = cosmocalc.kpc_per_pix(args.z)
-    elif args.cm_per_pix:
-        label = "cm/pix (DA)"
-        value = cosmocalc.cm_per_pix(args.z)
-    elif args.norm_apec:
-        label = "norm (APEC) [cm^-5]"
-        value = cosmocalc.norm_apec(args.z)
-    else:
-        raise ValueError("no quantity to calculate")
+    q_active = get_quantities(args)
+    if len(q_active) > 1:
+        if args.unit is not None:
+            args.unit = None
+            print("WARNING: ignored argument --unit", file=sys.stderr)
+        if args.brief:
+            args.brief = False
+            print("WARNING: ignored argument --brief", file=sys.stderr)
 
-    if not args.brief:
-        print(label + ": ", end="", flush=True)
-    print(value)
+    for q in q_active:
+        calc_quantity(q=q, calculator=cosmocalc, args=args)
 
 
 if __name__ == "__main__":