1919#
2020"""Provides input and output functions for Healpix maps, alm, and cl.
2121"""
22- from __future__ import with_statement
2322from __future__ import division
2423
2524import six
26- import gzip
27- import tempfile
28- import shutil
29- import os
3025import warnings
3126import astropy .io .fits as pf
3227import numpy as np
4540class HealpixFitsWarning (Warning ):
4641 pass
4742
48- def writeto (tbhdu , filename ):
49- # FIXME: Pyfits versions earlier than 3.1.2 had no support or flaky support
50- # for writing to .gz files or GzipFile objects. Drop this code when
51- # we decide to drop support for older versions of Pyfits or if we decide
52- # to support only Astropy.
53- if isinstance (filename , six .string_types ) and filename .endswith ('.gz' ):
54- basefilename , ext = os .path .splitext (filename )
55- with tempfile .NamedTemporaryFile (suffix = '.fits' ) as tmpfile :
56- tbhdu .writeto (tmpfile .name , clobber = True )
57- gzfile = gzip .open (filename , 'wb' )
58- try :
59- try :
60- shutil .copyfileobj (tmpfile , gzfile )
61- finally :
62- gzfile .close ()
63- except :
64- os .unlink (gzfile .name )
65- raise
66- else :
67- tbhdu .writeto (filename , clobber = True )
68-
6943def read_cl (filename , dtype = np .float64 , h = False ):
70- """Reads Cl from an healpix file, as IDL fits2cl.
44+ """Reads Cl from a healpix file, as IDL fits2cl.
7145
7246 Parameters
7347 ----------
@@ -82,41 +56,47 @@ def read_cl(filename, dtype=np.float64, h=False):
8256 the cl array
8357 """
8458 fits_hdu = _get_hdu (filename , hdu = 1 )
85- cl = [fits_hdu .data .field (n ) for n in range (len (fits_hdu .columns ))]
59+ cl = np . array ( [fits_hdu .data .field (n ) for n in range (len (fits_hdu .columns ))])
8660 if len (cl ) == 1 :
8761 return cl [0 ]
8862 else :
8963 return cl
9064
91- def write_cl (filename , cl , dtype = np .float64 ):
92- """Writes Cl into an healpix file, as IDL cl2fits.
65+ def write_cl (filename , cl , dtype = np .float64 , overwrite = False ):
66+ """Writes Cl into a healpix file, as IDL cl2fits.
9367
9468 Parameters
9569 ----------
9670 filename : str
9771 the fits file name
9872 cl : array
99- the cl array to write to file, currently TT only
73+ the cl array to write to file
74+ overwrite : bool, optional
75+ If True, existing file is silently overwritten. Otherwise trying to write
76+ an existing file raises an OSError (IOError for Python 2).
10077 """
10178 # check the dtype and convert it
10279 fitsformat = getformat (dtype )
10380 column_names = ['TEMPERATURE' ,'GRADIENT' ,'CURL' ,'G-T' ,'C-T' ,'C-G' ]
104- if isinstance ( cl , list ) :
81+ if len ( np . shape ( cl )) == 2 :
10582 cols = [pf .Column (name = column_name ,
10683 format = '%s' % fitsformat ,
10784 array = column_cl ) for column_name , column_cl in zip (column_names [:len (cl )], cl )]
108- else : # we write only one TT
85+ elif len (np .shape (cl )) == 1 :
86+ # we write only TT
10987 cols = [pf .Column (name = 'TEMPERATURE' ,
11088 format = '%s' % fitsformat ,
11189 array = cl )]
90+ else :
91+ raise RuntimeError ('write_cl: Expected one or more vectors of equal length' )
11292
11393 tbhdu = pf .BinTableHDU .from_columns (cols )
11494 # add needed keywords
11595 tbhdu .header ['CREATOR' ] = 'healpy'
116- writeto (tbhdu , filename )
96+ tbhdu . writeto (filename , overwrite = overwrite )
11797
118- def write_map (filename ,m ,nest = False ,dtype = np .float32 ,fits_IDL = True ,coord = None ,partial = False ,column_names = None ,column_units = None ,extra_header = ()):
119- """Writes an healpix map into an healpix file.
98+ def write_map (filename ,m ,nest = False ,dtype = np .float32 ,fits_IDL = True ,coord = None ,partial = False ,column_names = None ,column_units = None ,extra_header = (), overwrite = False ):
99+ """Writes a healpix map into a healpix file.
120100
121101 Parameters
122102 ----------
@@ -152,8 +132,11 @@ def write_map(filename,m,nest=False,dtype=np.float32,fits_IDL=True,coord=None,pa
152132 dtype: np.dtype or list of np.dtypes, optional
153133 The datatype in which the columns will be stored. Will be converted
154134 internally from the numpy datatype to the fits convention. If a list,
155- the length must correspond to the number of map arrays.
135+ the length must correspond to the number of map arrays.
156136 Default: np.float32.
137+ overwrite : bool, optional
138+ If True, existing file is silently overwritten. Otherwise trying to write
139+ an existing file raises an OSError (IOError for Python 2).
157140 """
158141 if not hasattr (m , '__len__' ):
159142 raise TypeError ('The map must be a sequence' )
@@ -202,7 +185,7 @@ def write_map(filename,m,nest=False,dtype=np.float32,fits_IDL=True,coord=None,pa
202185 array = pix ,
203186 unit = None ))
204187
205- for cn , cu , mm , curr_fitsformat in zip (column_names , column_units , m ,
188+ for cn , cu , mm , curr_fitsformat in zip (column_names , column_units , m ,
206189 fitsformat ):
207190 if len (mm ) > 1024 and fits_IDL :
208191 # I need an ndarray, for reshape:
@@ -245,12 +228,12 @@ def write_map(filename,m,nest=False,dtype=np.float32,fits_IDL=True,coord=None,pa
245228 for args in extra_header :
246229 tbhdu .header [args [0 ]] = args [1 :]
247230
248- writeto (tbhdu , filename )
231+ tbhdu . writeto (filename , overwrite = overwrite )
249232
250233
251234def read_map (filename ,field = 0 ,dtype = np .float64 ,nest = False ,partial = False ,hdu = 1 ,h = False ,
252235 verbose = True ,memmap = False ):
253- """Read an healpix map from a fits file. Partial-sky files,
236+ """Read a healpix map from a fits file. Partial-sky files,
254237 if properly identified, are expanded to full size and filled with UNSEEN.
255238
256239 Parameters
@@ -265,7 +248,7 @@ def read_map(filename,field=0,dtype=np.float64,nest=False,partial=False,hdu=1,h=
265248 first column after the pixel index column.
266249 If None, all columns are read in.
267250 dtype : data type or list of data types, optional
268- Force the conversion to some type. Passing a list allows different
251+ Force the conversion to some type. Passing a list allows different
269252 types for each field. In that case, the length of the list must
270253 correspond to the length of the field parameter. Default: np.float64
271254 nest : bool, optional
@@ -353,12 +336,12 @@ def read_map(filename,field=0,dtype=np.float64,nest=False,partial=False,hdu=1,h=
353336 # increment field counters
354337 field = tuple (f if isinstance (f , str ) else f + 1 for f in field )
355338 try :
356- pix = fits_hdu .data .field (0 ).astype (int ).ravel ()
339+ pix = fits_hdu .data .field (0 ).astype (int , copy = False ).ravel ()
357340 except pf .VerifyError as e :
358341 print (e )
359342 print ("Trying to fix a badly formatted header" )
360343 fits_hdu .verify ("fix" )
361- pix = fits_hdu .data .field (0 ).astype (int ).ravel ()
344+ pix = fits_hdu .data .field (0 ).astype (int , copy = False ).ravel ()
362345
363346 try :
364347 assert len (dtype ) == len (field ), "The number of dtypes are not equal to the number of fields"
@@ -367,12 +350,12 @@ def read_map(filename,field=0,dtype=np.float64,nest=False,partial=False,hdu=1,h=
367350
368351 for ff , curr_dtype in zip (field , dtype ):
369352 try :
370- m = fits_hdu .data .field (ff ).astype (curr_dtype ).ravel ()
353+ m = fits_hdu .data .field (ff ).astype (curr_dtype , copy = False ).ravel ()
371354 except pf .VerifyError as e :
372355 print (e )
373356 print ("Trying to fix a badly formatted header" )
374357 m = fits_hdu .verify ("fix" )
375- m = fits_hdu .data .field (ff ).astype (curr_dtype ).ravel ()
358+ m = fits_hdu .data .field (ff ).astype (curr_dtype , copy = False ).ravel ()
376359
377360 if partial :
378361 mnew = UNSEEN * np .ones (sz , dtype = curr_dtype )
@@ -397,20 +380,26 @@ def read_map(filename,field=0,dtype=np.float64,nest=False,partial=False,hdu=1,h=
397380 pass
398381 ret .append (m )
399382
383+ if h :
384+ header = []
385+ for (key , value ) in fits_hdu .header .items ():
386+ header .append ((key , value ))
387+
400388 if len (ret ) == 1 :
401389 if h :
402- return ret [0 ],fits_hdu . header . items ()
390+ return ret [0 ], header
403391 else :
404392 return ret [0 ]
405393 else :
394+ if all (dt == dtype [0 ] for dt in dtype ):
395+ ret = np .array (ret )
406396 if h :
407- ret .append (fits_hdu .header .items ())
408- return tuple (ret )
397+ return ret , header
409398 else :
410- return tuple ( ret )
399+ return ret
411400
412401
413- def write_alm (filename ,alms ,out_dtype = None ,lmax = - 1 ,mmax = - 1 ,mmax_in = - 1 ):
402+ def write_alm (filename ,alms ,out_dtype = None ,lmax = - 1 ,mmax = - 1 ,mmax_in = - 1 , overwrite = False ):
414403 """Write alms to a fits file.
415404
416405 In the fits file the alms are written
@@ -480,7 +469,7 @@ def write_alm(filename,alms,out_dtype=None,lmax=-1,mmax=-1,mmax_in=-1):
480469
481470 tbhdu = pf .BinTableHDU .from_columns ([cindex ,creal ,cimag ])
482471 hdulist .append (tbhdu )
483- writeto (hdulist , filename )
472+ hdulist . writeto (filename , overwrite = overwrite )
484473
485474def read_alm (filename ,hdu = 1 ,return_mmax = False ):
486475 """Read alm from a fits file.
@@ -494,7 +483,7 @@ def read_alm(filename,hdu=1,return_mmax=False):
494483 ----------
495484 filename : str or HDUList or HDU
496485 The name of the fits file to read
497- hdu : int, optional
486+ hdu : int, or tuple of int, optional
498487 The header to read. Start at 0. Default: hdu=1
499488 return_mmax : bool, optional
500489 If true, both the alms and mmax is returned in a tuple. Default: return_mmax=False
@@ -504,17 +493,34 @@ def read_alm(filename,hdu=1,return_mmax=False):
504493 alms[, mmax] : complex array or tuple of a complex array and an int
505494 The alms read from the file and optionally mmax read from the file
506495 """
507- idx , almr , almi = mrdfits (filename ,hdu = hdu )
508- l = np .floor (np .sqrt (idx - 1 )).astype (np .long )
509- m = idx - l ** 2 - l - 1
510- if (m < 0 ).any ():
511- raise ValueError ('Negative m value encountered !' )
512- lmax = l .max ()
513- mmax = m .max ()
514- alm = almr * (0 + 0j )
515- i = Alm .getidx (lmax ,l ,m )
516- alm .real [i ] = almr
517- alm .imag [i ] = almi
496+ alms = []
497+ lmaxtot = None
498+ mmaxtot = None
499+ for unit in np .atleast_1d (hdu ):
500+ idx , almr , almi = mrdfits (filename ,hdu = unit )
501+ l = np .floor (np .sqrt (idx - 1 )).astype (np .long )
502+ m = idx - l ** 2 - l - 1
503+ if (m < 0 ).any ():
504+ raise ValueError ('Negative m value encountered !' )
505+ lmax = l .max ()
506+ mmax = m .max ()
507+ if lmaxtot is None :
508+ lmaxtot = lmax
509+ mmaxtot = mmax
510+ else :
511+ if lmaxtot != lmax or mmaxtot != mmax :
512+ raise RuntimeError (
513+ 'read_alm: harmonic expansion order in {} HDUs {} does not '
514+ 'match' .format (filename , unit , hdu ))
515+ alm = almr * (0 + 0j )
516+ i = Alm .getidx (lmax ,l ,m )
517+ alm .real [i ] = almr
518+ alm .imag [i ] = almi
519+ alms .append (alm )
520+ if len (alms ) == 1 :
521+ alm = alms [0 ]
522+ else :
523+ alm = np .array (alms )
518524 if return_mmax :
519525 return alm , mmax
520526 else :
@@ -609,7 +615,7 @@ def mwrfits(filename,data,hdu=1,colnames=None,keys=None):
609615 for k ,v in keys .items ():
610616 tbhdu .header [k ] = v
611617 # write the file
612- writeto (tbhdu , filename )
618+ tbhdu . writeto (filename )
613619
614620def getformat (t ):
615621 """Get the FITS convention format string of data type t.
0 commit comments