From 10f075260dfe79867d4f749bb766e0723b7f7823 Mon Sep 17 00:00:00 2001
From: Aaron LI <aly@aaronly.me>
Date: Mon, 26 Jun 2017 19:47:48 +0800
Subject: astro/fitscube.py: Rework header & wcs handling

---
 astro/fits/fitscube.py | 78 ++++++++++++++++++++++++++++----------------------
 1 file changed, 44 insertions(+), 34 deletions(-)

(limited to 'astro/fits/fitscube.py')

diff --git a/astro/fits/fitscube.py b/astro/fits/fitscube.py
index 3d2d4d0..0ad2009 100755
--- a/astro/fits/fitscube.py
+++ b/astro/fits/fitscube.py
@@ -33,23 +33,28 @@ class FITSCube:
         print("Loaded FITS cube from file: %s" % infile)
         print("Cube dimensions: %dx%dx%d" %
               (self.width, self.height, self.nslice))
+        # The Z-axis position of the first slice.
+        self.zbegin = self.header["CRVAL3"]
+        # The Z-axis step/spacing between slices.
+        self.zstep = self.header["CDELT3"]
 
     def add_slices(self, slices, zbegin=0.0, zstep=1.0):
         """
         Create a FITS cube from input image slices.
         """
+        self.slices = slices
+        self.zbegin = zbegin
+        self.zstep = zstep
         nslice = len(slices)
         header, image = self.open_image(slices[0])
         shape = (nslice, ) + image.shape
         data = np.zeros(shape, dtype=image.dtype)
         for i, s in enumerate(slices):
-            print("[%d/%d] Adding image slice: %s ..." % (i, nslice, s))
+            print("[%d/%d] Adding image slice: %s ..." % (i+1, nslice, s))
             hdr, img = self.open_image(s)
             data[i, :, :] = img
         self.data = data
-        wcs = self.make_wcs(header, zbegin=zbegin, zstep=zstep)
         self.header = header.copy(strip=True)
-        self.header.extend(wcs.to_header(), update=True)
         print("Created FITS cube of dimensions: %dx%dx%d" %
               (self.width, self.height, self.nslice))
 
@@ -63,7 +68,7 @@ class FITSCube:
         The input slice image may have following dimensions:
         * NAXIS=2: [Y, X]
         * NAXIS=3: [FREQ=1, Y, X]
-        * NAXIS=4: [FREQ=1, STOKES=1, Y, X]
+        * NAXIS=4: [STOKES=1, FREQ=1, Y, X]
 
         NOTE
         ----
@@ -86,31 +91,50 @@ class FITSCube:
             # NAXIS=3: [FREQ=1, Y, X]
             image = data[0, :, :]
         elif data.ndim == 4 and data.shape[0] == 1 and data.shape[1] == 1:
-            # NAXIS=4: [FREQ=1, STOKES=1, Y, X]
+            # NAXIS=4: [STOKES=1, FREQ=1, Y, X]
             image = data[0, 0, :, :]
         else:
             raise ValueError("Slice '{0}' has invalid dimensions: {1}".format(
                 infile, data.shape))
         return (header, image)
 
-    def make_wcs(self, header, zbegin, zstep):
+    @property
+    def header(self):
+        try:
+            return self.header_
+        except AttributeError:
+            return fits.Header()
+
+    @header.setter
+    def header(self, value):
+        self.header_ = value
+        for key in ["CTYPE4", "CRPIX4", "CRVAL4", "CDELT4", "CUNIT4"]:
+            try:
+                del self.header_[key]
+            except KeyError:
+                pass
+
+    @property
+    def wcs(self):
         w = WCS(naxis=3)
         w.wcs.ctype = ["pixel", "pixel", "pixel"]
-        w.wcs.crpix = np.array([header.get("CRPIX1", 1.0),
-                                header.get("CRPIX2", 1.0),
+        w.wcs.crpix = np.array([self.header.get("CRPIX1", 1.0),
+                                self.header.get("CRPIX2", 1.0),
                                 1.0])
-        w.wcs.crval = np.array([header.get("CRVAL1", 0.0),
-                                header.get("CRVAL2", 0.0),
-                                zbegin])
-        w.wcs.cdelt = np.array([header.get("CDELT1", 1.0),
-                                header.get("CDELT2", 1.0),
-                                zstep])
+        w.wcs.crval = np.array([self.header.get("CRVAL1", 0.0),
+                                self.header.get("CRVAL2", 0.0),
+                                self.zbegin])
+        w.wcs.cdelt = np.array([self.header.get("CDELT1", 1.0),
+                                self.header.get("CDELT2", 1.0),
+                                self.zstep])
         return w
 
-    def write(self, outfile, clobber):
-        self.header.add_history(datetime.now().isoformat())
-        self.header.add_history(" ".join(sys.argv))
-        hdu = fits.PrimaryHDU(data=self.data, header=self.header)
+    def write(self, outfile, clobber=False):
+        header = self.header
+        header.extend(self.wcs.to_header(), update=True)
+        header.add_history(datetime.now().isoformat())
+        header.add_history(" ".join(sys.argv))
+        hdu = fits.PrimaryHDU(data=self.data, header=header)
         try:
             hdu.writeto(outfile, overwrite=clobber)
         except TypeError:
@@ -131,27 +155,13 @@ class FITSCube:
         ns, __, __ = self.data.shape
         return ns
 
-    @property
-    def zbegin(self):
-        """
-        The Z-axis position of the first slice.
-        """
-        return self.header["CRVAL3"]
-
-    @property
-    def zstep(self):
-        """
-        The Z-axis step/spacing between slices.
-        """
-        return self.header["CDELT3"]
-
     @property
     def zvalues(self):
         """
         Calculate the Z-axis positions for all slices
         """
         nslice = self.nslice
-        wcs = WCS(self.header)
+        wcs = self.wcs
         pix = np.zeros(shape=(nslice, 3), dtype=np.int)
         pix[:, 2] = np.arange(nslice)
         world = wcs.wcs_pix2world(pix, 0)
@@ -166,7 +176,7 @@ def cmd_info(args):
     print("Image/slice size: %dx%d" % (cube.width, cube.height))
     print("Number of slices: %d" % cube.nslice)
     print("Slice step/spacing: %.3f" % cube.zstep)
-    print("Slice positions: {0}".format(cube.zvalues))
+    print("Slice positions:\n{0}".format(cube.zvalues))
 
 
 def cmd_create(args):
-- 
cgit v1.2.2