def writeplanes(save_name='/mnt/work/erosolow/GRS_13CO_all.fits'): spatial_template = fits.open('INTEG/COHRS_RELEASE1_FULL_INTEG.fit') spectral_template = SpectralCube.read('reprojected.fits') # Smoosh astrometry components together spatial_header = spatial_template[0].header spectral_header = spectral_template.header new_header = spatial_header.copy() new_header["NAXIS"] = 3 for keyword in ['NAXIS3', 'CRVAL3', 'CDELT3','CRPIX3','CUNIT3']: new_header[keyword] = spectral_header[keyword] new_header['BMAJ'] = 14./3600 new_header['BMIN'] = 14./3600 new_header['BPA'] = 0.00 if os.path.exists(save_name): raise Exception("The file name {} already " "exists".format(save_name)) # Open a file and start filling this with planes. output_fits = fits.StreamingHDU(save_name, new_header) # Again, set up a common vel axis and spin out vel = np.linspace(-30, 160, 191) for v in vel: output_fits.write(fits.getdata(planesdir + 'GRSPLANE_{0}'.format(v) + '.fits')) output_fits.close()
def _make_streaming_hdu(self, fileobj): hd = fits.Header() hd['SIMPLE'] = (True, 'conforms to FITS standard') hd['BITPIX'] = (32, 'array data type') hd['NAXIS'] = (2, 'number of array dimensions') hd['NAXIS1'] = 5 hd['NAXIS2'] = 5 hd['EXTEND'] = True return fits.StreamingHDU(fileobj, hd)
def test_streaming_hdu_file_wrong_mode(self): """ Test that streaming an HDU to a file opened in the wrong mode fails as expected. """ with open(self.temp('new.fits'), 'wb') as f: header = fits.Header() fits.StreamingHDU(f, header)
hdr['EXTEND'] = 'T' primary_hdu = fits.PrimaryHDU(header=hdr) primary_hdu.writeto('test.fits') # Use BinTableHDU as a template hdr = fits.BinTableHDU(Table(names=['SIGNAL'], dtype=['=f4' if bPotential else '=i4']), name='BINTABLE').header hdr['ORDERING'] = ("RING", "Pixel ordering scheme, either RING or NESTED") hdr['INDXSCHM'] = ("IMPLICIT", "Pixel indexing scheme (IMPLICIT or EXPLICIT)") hdr['NSIDE'] = (NSIDE, "Resolution parameter for HEALPIX") hdr['COORDSYS'] = ("C", "Pixelisation coordinate system") hdr['PIXTYPE'] = ("HEALPIX", "HEALPIX Pixelisation") hdr['NAXIS'] = 2 hdr['NAXIS2'] = N hdr['NAXIS1'] = 1 hdr['BITPIX'] = -32 if bPotential else 32 hdu = fits.StreamingHDU('test.fits', hdr) for file in argv[1:]: with open(file, 'rb') as hpb: data = np.fromfile(hpb, dtype=hpb_type) data = pd.DataFrame(data, columns=data.dtype.names) if bPotential: hdu.write(data['potential'].values) else: data = (data.grouped + data.ungrouped).to_frame() data.columns = ['signal'] hdu.write(data['signal'].values) hdu.close()
def write_fits(filename, data, header, extension, extname, comm): """ Collectively write local arrays into a single FITS file. Parameters ---------- filename : str The FITS file name. data : ndarray The array to be written. header : pyfits.Header The data FITS header. None can be set, in which case a minimal FITS header will be inferred from the data. extension : boolean If True, the data will be written as an extension to an already existing FITS file. extname : str The FITS extension name. Use None to write the primary HDU. comm : mpi4py.Comm The MPI communicator of the local arrays. Use MPI.COMM_SELF if the data are not meant to be combined into a global array. Make sure that the MPI processes are not executing this routine with the same file name. """ # check if the file name is the same for all MPI jobs files = comm.allgather(filename + str(extname)) all_equal = all(f == files[0] for f in files) if comm.size > 1 and not all_equal: raise ValueError('The file name is not the same for all MPI jobs.') ndims = comm.allgather(data.ndim) if any(n != ndims[0] for n in ndims): raise ValueError("The arrays have an incompatible number of dimensions" ": '{0}'.".format(', '.join(str(n) for n in ndims))) ndim = ndims[0] shapes = comm.allgather(data.shape) if any(s[1:] != shapes[0][1:] for s in shapes): raise ValueError("The arrays have incompatible shapes: '{0}'.".format( strshape(shapes))) # get header if header is None: header = create_fitsheader_for(data, extname=extname) else: header = header.copy() if extname is not None: header['extname'] = extname # we remove the file first to avoid an annoying pyfits informative message if not extension: if comm.rank == 0: try: os.remove(filename) except OSError: pass # case without MPI communication if comm.size == 1: if not extension: hdu = pyfits.PrimaryHDU(data, header) hdu.writeto(filename, overwrite=True) else: pyfits.append(filename, data, header) return # get global/local parameters nglobal = sum(s[0] for s in shapes) s = split(nglobal, comm.size, comm.rank) nlocal = s.stop - s.start if data.shape[0] != nlocal: raise ValueError( "On rank {}, the local array shape '{}' is invalid. T" "he first dimension does not match the expected local" " number '{}' given the global number '{}'.{}".format( comm.rank, data.shape, nlocal, nglobal, '' if comm.rank > 0 else ' Shapes are: {}.'.format(shapes))) # write FITS header if comm.rank == 0: header['NAXIS' + str(ndim)] = nglobal shdu = pyfits.StreamingHDU(filename, header) data_loc = shdu._datLoc shdu.close() else: data_loc = None data_loc = comm.bcast(data_loc) # get a communicator excluding the processes which have no work to do # (Create_subarray does not allow 0-sized subarrays) chunk = product(data.shape[1:]) rank_nowork = min(comm.size, nglobal) group = comm.Get_group() group.Incl(list(range(rank_nowork))) newcomm = comm.Create(group) # collectively write data if comm.rank < rank_nowork: # mpi4py 1.2.2: pb with viewing data as big endian KeyError '>d' if sys.byteorder == 'little' and data.dtype.byteorder == '=' or \ data.dtype.byteorder == '<': data = data.byteswap() data = data.newbyteorder('=') mtype = DTYPE_MAP[data.dtype] ftype = mtype.Create_subarray([nglobal * chunk], [nlocal * chunk], [s.start * chunk]) ftype.Commit() f = MPI.File.Open(newcomm, filename, amode=MPI.MODE_APPEND | MPI.MODE_WRONLY | MPI.MODE_CREATE) f.Set_view(data_loc, mtype, ftype, 'native', MPI.INFO_NULL) f.Write_all(data) f.Close() ftype.Free() newcomm.Free() # pad FITS file with zeros if comm.rank == 0: datasize = nglobal * chunk * data.dtype.itemsize BLOCK_SIZE = 2880 padding = BLOCK_SIZE - (datasize % BLOCK_SIZE) with open(filename, 'a') as f: if f.tell() - data_loc != datasize: raise RuntimeError('Unexpected file size.') f.write(padding * '\0') comm.Barrier()