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 write_fits(filename, data, header, shape_global, extension, comm,
extname=None):
"""Write local images into a FITS file"""
if comm is None:
comm = MPI.COMM_SELF
if not extension:
try:
os.remove(filename)
except:
pass
if header is None:
header = create_fitsheader(shape_global[::-1], data.dtype,
extname=extname)
rank = comm.Get_rank()
size = comm.Get_size()
files = comm.allgather(filename)
allsame = all([f == files[0] for f in files])
alldiff = len(files) == len(np.unique(files))
if not alldiff and not allsame:
raise ValueError('Some target filenames are equal, but not all.')
if alldiff or size == 1:
if not extension:
hdu = pyfits.PrimaryHDU(data, header)
hdu.writeto(filename, clobber=True)
else:
pyfits.append(filename, data, header)
return
if rank == 0:
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)
nglobal = shape_global[0]
chunk = np.product(shape_global[1:])
s = split_work(nglobal)
nlocal = s.stop - s.start
nmax = int(np.ceil(nglobal / size))
rank_nowork = int(np.ceil(nglobal / nmax))
group = comm.Get_group()
group.Incl(range(rank_nowork))
newcomm = comm.Create(group)
if rank < rank_nowork:
mtype = {1:MPI.BYTE, 4: MPI.FLOAT, 8:MPI.DOUBLE}[data.dtype.itemsize]
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)
# mpi4py 1.2.2: pb with viewing data as big endian KeyError '>d'
if sys.byteorder == 'little' and data.dtype.byteorder in ('=', '<'):
data = data.byteswap()
else:
data = data.newbyteorder('=')
f.Write_all(data[0:nlocal])
f.Close()
if rank == 0:
shdu._ffo = pyfits.core._File(filename, 'append')
shdu._ffo.getfile().seek(0,2)
pyfitstype = {8:'uint8', 16:'int16', 32:'int32', 64:'int64', -32:'float32', -64:'float64'}[header['BITPIX']]
completed = shdu.write(np.empty(0, dtype=pyfitstype))
shdu.close()
if not completed:
raise RuntimeError('File is not completely written')
comm.Barrier()