Page MenuHomec4science

test_io.py
No OneTemporary

File Metadata

Created
Thu, Aug 15, 10:33

test_io.py

#!/usr/bin/env python
import sys,string,types
from numpy import *
from pNbody import mpi
from pNbody import ic
pio = "no"
piomode = "mpiio"
def OpenFile(filename,mode='r'):
'''
here, the filename is already determined
'''
if string.find(mode,'w') != -1:
mpiiomode = mpi.MPI.MODE_CREATE+mpi.MPI.MODE_RDWR
if string.find(mode,'r') != -1:
mpiiomode = mpi.MPI.MODE_RDONLY
##################
# multiple files
##################
if pio == "yes" and piomode == 'mpiio':
f = mpi.MPI.File.Open(mpi.comm_self,filename,amode=mpiiomode)
elif pio == "yes" and piomode == 'normal':
f = open(filename,mode)
##################
# single file
##################
elif pio == "no" and piomode == "mpiio":
f = mpi.MPI.File.Open(mpi.comm,filename,amode=mpiiomode)
elif pio == "no" and piomode == "normal":
if mpi.mpi_IsMaster():
f = open(filename,mode)
else:
f = None
return f
def CloseFile(f):
'''
Close a file
'''
if type(f)==types.FileType:
f.close()
elif type(f)==mpi.MPI.File :
f.Close()
else:
pass
def WriteArray(f,data,add_header=False,x64=False,byteorder=sys.byteorder):
"""
write an array
f : file
an open file
data : data
a numpy array
"""
header_size=0
### !!!!
if add_header:
if x64:
header_size=8
else:
header_size=4
if add_header:
# compute the total size of the array
nbytes_tot = sum(mpi.mpi_allgather(data.nbytes))
if mpi.mpi_IsMaster():
WriteFortranBlockHeader(f,nbytes_tot)
if sys.byteorder != byteorder:
data.byteswap(True)
##################
# multiple files
##################
if pio == "yes" and piomode == 'mpiio':
f.Set_view(0)
f.Write(data.tostring())
elif pio == "yes" and piomode == 'normal':
f.write(data.tostring())
##################
# single file
##################
elif pio == "no" and piomode == "mpiio":
f.Set_view(header_size+data.nbytes*mpi.mpi_Rank())
f.Write(data.tostring())
elif pio == "no" and piomode == "normal":
if mpi.mpi_IsMaster():
f.write(data.tostring())
for Task in range(1,mpi.mpi_NTask()):
sdata = mpi.comm.recv(source = Task)
f.write(sdata.tostring())
else:
mpi.comm.send(data, dest = 0)
### !!!!
if add_header:
if mpi.mpi_IsMaster():
WriteFortranBlockHeader(f,nbytes_tot)
#################################################
# main
#################################################
filename = 'test.dat'
if pio == "yes":
filename = "%s.%d"%(filename,mpi.mpi_ThisTask())
# create data
nbody = 1000/mpi.NTask
data = mpi.ThisTask*ones((nbody,3)).astype(float32)
# open the file
f = OpenFile(filename,'w')
WriteArray(f,data)
CloseFile(f)

Event Timeline