aboutsummaryrefslogtreecommitdiffstats
path: root/fg21sim/utils/io.py
diff options
context:
space:
mode:
authorAaron LI <aly@aaronly.me>2017-08-13 01:23:40 +0800
committerAaron LI <aly@aaronly.me>2017-08-13 01:23:40 +0800
commitf0471b1745583eb7dac8da5c2cc4a0266be09899 (patch)
treea81563ccde07ec221f0829792f47c51483a828eb /fg21sim/utils/io.py
parent2b5b4bd00801a1c05d3a5093d4dd556ff87941ad (diff)
downloadfg21sim-f0471b1745583eb7dac8da5c2cc4a0266be09899.tar.bz2
utils/io.py: Implement new function "write_fits_image()"
Diffstat (limited to 'fg21sim/utils/io.py')
-rw-r--r--fg21sim/utils/io.py37
1 files changed, 37 insertions, 0 deletions
diff --git a/fg21sim/utils/io.py b/fg21sim/utils/io.py
index 2ce7cda..050574c 100644
--- a/fg21sim/utils/io.py
+++ b/fg21sim/utils/io.py
@@ -10,7 +10,9 @@ import logging
import pickle
from datetime import datetime
+import numpy as np
import pandas as pd
+from astropy.io import fits
logger = logging.getLogger(__name__)
@@ -132,3 +134,38 @@ def pickle_load(infile):
obj : The loaded Python object from the input file.
"""
return pickle.load(open(infile, "rb"))
+
+
+def write_fits_image(outfile, image, header=None, float32=True,
+ clobber=False, checksum=False):
+ """
+ Write the supplied image (together with header information) into
+ the output FITS file.
+
+ Parameters
+ ----------
+ outfile : str
+ The path/filename to the output file storing the pickled object.
+ image : 2D `~numpy.ndarray`
+ The image data to be written out to the FITS file.
+ NOTE: image.shape: (nrow, ncol) <-> FITS image: (ncol, nrow)
+ header : `~astropy.io.fits.Header`
+ The FITS header information for this image
+ float32 : bool, optional
+ Whether coerce the image data (generally double/float64 data type)
+ into single/float32 (in order to save space)?
+ Default: True
+ clobber : bool, optional
+ Whether to overwrite the existing output file.
+ Default: False
+ checksum : bool, optional
+ Whether to calculate the data checksum, which may cost some time?
+ Default: False
+ """
+ _create_dir(outfile)
+ _check_existence(outfile, clobber=clobber, remove=True)
+ if float32:
+ image = np.asarray(image, dtype=float32)
+ hdu = fits.PrimaryHDU(data=image, header=header)
+ hdu.writeto(outfile, checksum=checksum)
+ logger.info("Wrote image to FITS file: %s" % outfile)