astro/fitscube.py: Add "add" and "sub" (subtract) sub-commands

1 files changed, 73 insertions, 0 deletions

diff --git a/astro/fits/fitscube.py b/astro/fits/fitscube.py
index af3400d..9be744e 100755
--- a/astro/fits/fitscube.py
+++ b/astro/fits/fitscube.py
@@ -312,6 +312,53 @@ def cmd_crop(args):
     print("Created FITS cube: %s" % args.outfile)
 
 
+def cmd_add(args):
+    """
+    Sub-command: "add", add two or more FITS cubes
+    """
+    if not args.clobber and os.path.exists(args.outfile):
+        raise FileExistsError("output file already exists: %s" % args.outfile)
+    if len(args.infiles) < 2:
+        raise RuntimeError("Two or more input FITS cubes required")
+
+    cube = FITSCube(args.infiles[0])
+    print("Data cube unit: %s" % cube.unit)
+    print("Image/slice size: %dx%d" % (cube.width, cube.height))
+    print("Number of slices: %d" % cube.nslice)
+
+    for f in args.infiles[1:]:
+        cube2 = FITSCube(f)
+        assert (cube.unit, cube.zunit) == (cube2.unit, cube2.zunit)
+        print("Adding cube %s ..." % f)
+        cube.data = cube.data + cube2.data
+
+    print("Saving FITS cube ...")
+    cube.write(args.outfile, clobber=args.clobber)
+    print("Saved FITS cube: %s" % args.outfile)
+
+
+def cmd_sub(args):
+    """
+    Sub-command: "sub", subtract one FITS cube by another one
+    """
+    if not args.clobber and os.path.exists(args.outfile):
+        raise FileExistsError("output file already exists: %s" % args.outfile)
+
+    cube = FITSCube(args.infile)
+    print("Data cube unit: %s" % cube.unit)
+    print("Image/slice size: %dx%d" % (cube.width, cube.height))
+    print("Number of slices: %d" % cube.nslice)
+
+    cube2 = FITSCube(args.infile2)
+    assert (cube.unit, cube.zunit) == (cube2.unit, cube2.zunit)
+    print("Subtracting cube %s ..." % args.infile2)
+    cube.data = cube.data - cube2.data
+
+    print("Saving FITS cube ...")
+    cube.write(args.outfile, clobber=args.clobber)
+    print("Saved FITS cube: %s" % args.outfile)
+
+
 def cmd_calibrate(args):
     """
     Sub-command: "calibrate", calibrate the z-axis slice/channel responses
@@ -493,6 +540,32 @@ def main():
                              help="input FITS cube")
     parser_crop.set_defaults(func=cmd_crop)
 
+    # sub-command: "add"
+    parser_add = subparsers.add_parser(
+        "add",
+        help="add two or more FITS cubes")
+    parser_add.add_argument("-C", "--clobber", action="store_true",
+                            help="overwrite existing output file")
+    parser_add.add_argument("-o", "--outfile", required=True,
+                            help="output FITS cube filename")
+    parser_add.add_argument("-i", "--infiles", nargs="+", required=True,
+                            help="two or more input FITS cubes")
+    parser_add.set_defaults(func=cmd_add)
+
+    # sub-command: "subtract"
+    parser_sub = subparsers.add_parser(
+        "sub", aliases=["subtract"],
+        help="subtract one FITS cube by another one")
+    parser_sub.add_argument("-C", "--clobber", action="store_true",
+                            help="overwrite existing output file")
+    parser_sub.add_argument("-o", "--outfile", required=True,
+                            help="output FITS cube filename")
+    parser_sub.add_argument("-i", "--infile", required=True,
+                            help="input FITS cube as the minuend")
+    parser_sub.add_argument("-I", "--infile2", required=True,
+                            help="another input FITS cube as the subtrahend")
+    parser_sub.set_defaults(func=cmd_sub)
+
     # sub-command: "calibrate"
     parser_cal = subparsers.add_parser(
         "calibrate",