From 472d49eb988efb11360ad48c157e478fe07c3ff8 Mon Sep 17 00:00:00 2001 From: Aaron LI Date: Tue, 29 Jan 2019 13:48:08 +0800 Subject: astro/fitscube.py: Clean up add_slices() method a bit --- astro/fitscube.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) (limited to 'astro') diff --git a/astro/fitscube.py b/astro/fitscube.py index d6bc986..0f2224a 100755 --- a/astro/fitscube.py +++ b/astro/fitscube.py @@ -65,22 +65,24 @@ class FITSCube: def add_slices(self, infiles, zbegin=0.0, zstep=1.0): """ Create a FITS cube from input image slices. + + NOTE: The infiles should be sorted appropriately. """ self.infiles = infiles self.zbegin = zbegin self.zstep = zstep - nslice = len(infiles) + N = len(infiles) + zvalues = zbegin + zstep * np.arange(N) header, image = self.open_image(infiles[0]) - shape = (nslice, ) + image.shape - data = np.zeros(shape, dtype=image.dtype) - for i, fn in enumerate(infiles): - print("[%d/%d] Adding image slice: %s ..." % (i+1, nslice, fn)) + height, width = image.shape + data = np.zeros((N, height, width), dtype=image.dtype) + for i, (z, fn) in enumerate(zip(zvalues, infiles)): + print("[%2d/%2d] slice @ %s: %s ..." % (i+1, N, z, fn)) hdr, img = self.open_image(fn) data[i, :, :] = img self.data = data self.header = header.copy(strip=True) - print("Created FITS cube of dimensions: %dx%dx%d" % - (self.width, self.height, self.nslice)) + print("Cube dimension: %dx%dx%d" % (width, height, N)) @staticmethod def open_image(infile): -- cgit v1.2.2