diff options
author | Aaron LI <aly@aaronly.me> | 2017-06-11 16:15:54 +0800 |
---|---|---|
committer | Aaron LI <aly@aaronly.me> | 2017-06-11 16:15:54 +0800 |
commit | 156ea20569bc166e217b53e0e4eb5f668314633e (patch) | |
tree | 7b589816c89d32b68370479274db607b25deadc1 /astro/fits | |
parent | 2e796f38dbdc1044c13a768c76fa733ad07b9829 (diff) | |
download | atoolbox-156ea20569bc166e217b53e0e4eb5f668314633e.tar.bz2 |
Move some tools for better management
Diffstat (limited to 'astro/fits')
-rwxr-xr-x | astro/fits/fitscube.py | 172 |
1 files changed, 172 insertions, 0 deletions
diff --git a/astro/fits/fitscube.py b/astro/fits/fitscube.py new file mode 100755 index 0000000..f9a96e1 --- /dev/null +++ b/astro/fits/fitscube.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +# +# Copyright (c) Weitian LI <weitian@aaronly.me> +# MIT license +# + +""" +Create FITS image cube from a series of image slices. +""" + +import os +import argparse + +import numpy as np +from astropy.io import fits +from astropy.wcs import WCS + + +class FITSCube: + """ + FITS image cube. + """ + def __init__(self, infile=None): + if infile is not None: + self.load(infile) + + def load(self, infile): + with fits.open(infile) as f: + self.data = f[0].data + self.header = f[0].header + print("Loaded FITS cube from file: %s" % infile) + print("Cube dimensions: %dx%dx%d" % + (self.width, self.height, self.nslice)) + + def add_slices(self, slices, zbegin=0.0, zstep=1.0): + """ + Create a FITS cube from input image slices. + """ + nslice = len(slices) + with fits.open(slices[0]) as f: + image = f[0].data + header = f[0].header + shape = (nslice, ) + image.shape + data = np.zeros(shape, dtype=image.dtype) + for i, s in enumerate(slices): + print("Adding image slice: %s ..." % s) + data[i, :, :] = fits.open(s)[0].data + 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)) + + def make_wcs(self, header, zbegin, zstep): + 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), + 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]) + return w + + def write(self, outfile, clobber): + hdu = fits.PrimaryHDU(data=self.data, header=self.header) + hdu.writeto(outfile, clobber=clobber) + + @property + def width(self): + __, __, w = self.data.shape + return w + + @property + def height(self): + __, h, __ = self.data.shape + return h + + @property + def nslice(self): + 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) + pix = np.zeros(shape=(nslice, 3), dtype=np.int) + pix[:, 2] = np.arange(nslice) + world = wcs.wcs_pix2world(pix, 0) + return world[:, 2] + + +def cmd_info(args): + """ + Sub-command: "info", show FITS cube information + """ + cube = FITSCube(args.infile) + 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)) + + +def cmd_create(args): + """ + Sub-command: "create", create a FITS cube + """ + if not args.clobber and os.path.exists(args.outfile): + raise FileExistsError("output file already exists: %s" % args.outfile) + cube = FITSCube() + cube.add_slices(args.infiles, zbegin=args.zbegin, zstep=args.zstep) + cube.write(args.outfile, clobber=args.clobber) + print("Created FITS cube: %s" % args.outfile) + + +def main(): + parser = argparse.ArgumentParser( + description="Create FITS cube from a series of image slices.") + subparsers = parser.add_subparsers(dest="subparser_name", + title="sub-commands", + help="additional help") + # sub-command: "info" + parser_info = subparsers.add_parser("info", help="show FITS cube info") + parser_info.add_argument("infile", help="FITS cube filename") + parser_info.set_defaults(func=cmd_info) + # sub-command: "create" + parser_create = subparsers.add_parser("create", help="create a FITS cube") + parser_create.add_argument("-C", "--clobber", dest="clobber", + action="store_true", + help="overwrite existing output file") + parser_create.add_argument("-z", "--z-begin", dest="zbegin", + type=float, default=0.0, + help="Z-axis position of the first slice") + parser_create.add_argument("-s", "--z-step", dest="zstep", + type=float, default=1.0, + help="Z-axis step/spacing between slices") + parser_create.add_argument("-o", "--outfile", dest="outfile", + required=True, + help="output FITS cube filename") + parser_create.add_argument("-i", "--infiles", dest="infiles", + nargs="+", required=True, + help="input image slices (in order)") + parser_create.set_defaults(func=cmd_create) + # + args = parser.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() |