Python io.FortranFile类代码示例

def test_fortranfiles_write():
    for filename in iglob(path.join(DATA_PATH, "fortran-*-*x*x*.dat")):
        m = re.search("fortran-([^-]+)-(\d+)x(\d+)x(\d+).dat", filename, re.I)
        if not m:
            raise RuntimeError("Couldn't match %s filename to regex" % filename)
        dims = (int(m.group(2)), int(m.group(3)), int(m.group(4)))

        counter = 0
        data = np.zeros(dims, dtype=m.group(1).replace("s", "<"))
        for k in range(dims[2]):
            for j in range(dims[1]):
                for i in range(dims[0]):
                    data[i, j, k] = counter
                    counter += 1
        tmpdir = tempfile.mkdtemp()
        try:
            testFile = path.join(tmpdir, path.basename(filename))
            f = FortranFile(testFile, "w", "<u4")
            f.write_record(data)
            f.close()
            originalfile = open(filename, "rb")
            newfile = open(testFile, "rb")
            assert_equal(originalfile.read(), newfile.read(), err_msg=filename)
            originalfile.close()
            newfile.close()
        finally:
            shutil.rmtree(tmpdir)

def test_fortranfiles_mixed_record():
    filename = path.join(DATA_PATH, "fortran-mixed.dat")
    f = FortranFile(filename, 'r', '<u4')
    record = f.read_record('<i4,<f4,<i8,(2)<f8')
    assert_equal(record['f0'][0], 1)
    assert_allclose(record['f1'][0], 2.3)
    assert_equal(record['f2'][0], 4)
    assert_allclose(record['f3'][0], [5.6, 7.8])

def read_association(listfile):
    assocFile = FortranFile(listfile, 'r')
    nassoc, columns = assocFile.read_ints()
    _tmp = (assocFile.read_reals(dtype=np.float32)).reshape((columns, nassoc)).transpose()
    assoc = pd.DataFrame(_tmp,
                         columns=['halo_id', 'level', 'halo_mass', 'gal_id', 'gal_mass'])

    assoc[['halo_id', 'level', 'gal_id']] =  assoc[['halo_id', 'level', 'gal_id']].astype(np.int32)
    return assoc

def read_halo_list(listfile):
    haloFile = FortranFile(listfile, 'r')
    nhalos, columns = haloFile.read_ints()
    _tmp = (haloFile.read_reals(dtype=np.float32)).reshape((columns, nhalos)).transpose()
    halos = pd.DataFrame(_tmp,
                         columns=['id', 'level', 'mass', 'x', 'y', 'z', 'rvir'])
    halos[['id', 'level']] = halos[['id', 'level']].astype(int)

    return halos

def import_data_G(name="G_0.01965", folder_name="./OUT/"):

    f = FortranFile(folder_name+name, 'r')

    version = f.read_reals(dtype='f4')
    time, Ra, Ra_c, P, Ha, Di, Pr, Le = f.read_reals(dtype='f4')
    nradius, ntheta, nphi, azsym = f.read_reals(dtype='f4') # be careful, all of them are reals
    radius = f.read_reals(dtype='f4')
    theta = f.read_reals(dtype='f4') #colatitude

    phi = np.arange(1, int(nphi)+1)/nphi*2.*np.pi/azsym #longitude (not read from file!)

    Vr = np.empty([nphi, ntheta, nradius])
    Vt = np.empty_like(Vr)
    Vp = np.empty_like(Vr)
    Temperature = np.empty_like(Vr)
    Composition = np.empty_like(Vr)

    for ir in np.arange(nradius):
        for it in np.arange(ntheta):
            Vr[:,it,ir]=f.read_reals(dtype='f4')
            Vt[:,it,ir]=f.read_reals(dtype='f4')
            Vp[:,it,ir]=f.read_reals(dtype='f4')
            Composition[:,it,ir]=f.read_reals(dtype='f4')
            Temperature[:,it,ir]=f.read_reals(dtype='f4')
    
    return time, Ra, Ra_c, P, Ha, Di, Pr, Le, nradius, ntheta, nphi, azsym, radius, theta, phi, Vr, Vt, Vp, Temperature, Composition

def read_fbin(filename):
    ''' this reads each written binary instance itteratively'''
    f = FortranFile(filename, 'r')
    array = []
    while True:
        try:
            array.append(f.read_reals(dtype=np.float_))
        except TypeError:
            break
    #array = np.reshape(array, (nspecs,-1))

    f.close()
    return array

def read_fortran_FFTfield(infile):
    """
    Read a Half-Field with FFTW indexing from
    a Fortran Unformatted Binary file. The first
    entry is a single integer.
    """
    f=FortranFile(infile,'r')
    Ngrid=f.read_ints(dtype=np.int32)[0]
    print('Fortran file Ngrid='+str(Ngrid))
    dcr=f.read_reals(dtype=np.complex64)
    dcr=np.reshape(dcr,(Ngrid//2+1,Ngrid,Ngrid),order='F')
    dcr.dump(infile+'.pickle') # Save infile as a pickle
    return dcr

 def read_fortran_FFTfield(self):
       """
       Read a fortran binary file from FFTW assert all Nyquist
       entries to be real.
       """
       f=FortranFile(self.infile,'r')
       Ng=f.read_ints(dtype=np.int32)[0]
       print('Fortran file Ngrid='+str(Ng))
       if (Ng != self.Ngrid):
             print('Ngrid values are not equal!')
       dcr=f.read_reals(dtype=np.complex64)
       dcr=np.reshape(dcr,(Ng//2+1,Ng,Ng),order='F')
       return dcr

def read_galaxy_list(listfile):
    galFile = FortranFile(listfile, 'r')
    print(listfile)
    ngal, columns = galFile.read_ints()
    _tmp = (galFile.read_reals(dtype=np.float32)).reshape((columns, ngal)).transpose()
    galaxies = pd.DataFrame(_tmp,
                            columns=['id', 'vt', 'dvz', 'dvr', 'dvtheta', 'mass', 'x', 'y', 'z'])
    galaxies.id.astype(int)

    galaxies['sigma'] = 1/3.*np.sqrt(galaxies.dvz**2 + galaxies.dvtheta**2 + galaxies.dvr**2)
    galaxies['sigmaoverv'] = galaxies.sigma / galaxies.vt
    galaxies['elliptic'] = galaxies.sigmaoverv > 1.5
    galaxies['spiral'] = galaxies.sigmaoverv < 0.8

    return galaxies

        def readfun(filename):
            '''
            reads unformatted fortran files 
            ''' 
            f = FortranFile('Outputs/'+filename, 'r')
            names = ''.join(f.read_reals('c'))
            data = []
            while True:
                    try:
                        data.append(f.read_reals(dtype=np.float_))
                    except TypeError:
                        break
                #array = np.reshape(array, (nspecs,-1))

            f.close()
            return [names.replace(' ',''),np.array(data)]

def test_fortranfiles_read():
    for filename in iglob(path.join(DATA_PATH, "fortran-*-*x*x*.dat")):
        m = re.search(r'fortran-([^-]+)-(\d+)x(\d+)x(\d+).dat', filename, re.I)
        if not m:
            raise RuntimeError("Couldn't match %s filename to regex" % filename)

        dims = (int(m.group(2)), int(m.group(3)), int(m.group(4)))

        dtype = m.group(1).replace('s', '<')

        f = FortranFile(filename, 'r', '<u4')
        data = f.read_record(dtype=dtype).reshape(dims, order='F')
        f.close()

        expected = np.arange(np.prod(dims)).reshape(dims).astype(dtype)
        assert_equal(data, expected)

def test_fortranfiles_read():
    for filename in iglob(path.join(DATA_PATH, "fortran-*-*x*x*.dat")):
        m = re.search('fortran-([^-]+)-(\d+)x(\d+)x(\d+).dat', filename, re.I)
        if not m:
            raise RuntimeError("Couldn't match %s filename to regex" % filename)
        dims = (int(m.group(2)), int(m.group(3)), int(m.group(4)))

        f = FortranFile(filename, 'r', '<u4')
        data = f.read_record(dtype=m.group(1)).reshape(dims)
        f.close()

        counter = 0
        for k in range(dims[2]):
            for j in range(dims[1]):
                for i in range(dims[0]):
                    assert_equal(counter, data[i,j,k])
                    counter += 1

	def read_stalker_data(self,
		datadir='/data',
		procdir='proc0'):
			
		f = FortranFile(datadir+'/'+procdir+'/particles_stalker.dat','r')
		file_not_ended=True
		firstline = f.read_record([('a','f8'),('b','i4')])
		firstline = f.read_record([('a','f8'),('b','i4')])
		snaptime = []
		nrshot = []
		partdata = []
		iddata = []
		while (True):
			try:
				temp = f.read_record([('snaptime','f8'),('nrshot','i4')])
				snaptime.append(temp['snaptime'])
				nrshot.append(temp['nrshot'])
			except TypeError:
				break
			except ValueError:
				break
			if (nrshot[-1] > 0):
				temp = f.read_record([('iddata','i4')])
				iddata.append(temp['iddata'])
				temp = f.read_record([('partdata','f8')])
				temp = np.reshape(temp,(-1,self.ndims_stalk))
				partdata.append(temp['partdata'])
			else:
				iddata.append([])
				partdata.append([])
		
		partdata=np.asarray(partdata)
		iddata = np.asarray(iddata)
		snaptime = np.asarray(snaptime)
		nrshot = np.array(nrshot)
		oldshape = np.shape(partdata)
		
		partdata = np.reshape(partdata,oldshape)
		
		setattr(self,procdir+'_snaptime',snaptime)
		setattr(self,procdir+'_iddata',iddata)
		setattr(self,procdir+'_partdata',partdata)
		setattr(self,procdir+'_nrshot',nrshot)

def read_fbin(filename):
    ''' this reads each written binary instance itteratively'''
    f = FortranFile(filename, 'r')
    array = []
    while True:
        try:
            array.append(f.read_reals(dtype=np.float_))
        except TypeError:
            break
    #array = np.reshape(array, (nspecs,-1))

    f.close()
    
    
    #newdata = np.array(array)[:,selected_index]
    
    #indices = xrange(0,len(array)-sets_of,sets_of)
    
    
    #newdata = newdata[indices,:]
    #return newdata
    return array

def test_fortranfiles_write():
    for filename in iglob(path.join(DATA_PATH, "fortran-*-*x*x*.dat")):
        m = re.search(r'fortran-([^-]+)-(\d+)x(\d+)x(\d+).dat', filename, re.I)
        if not m:
            raise RuntimeError("Couldn't match %s filename to regex" % filename)
        dims = (int(m.group(2)), int(m.group(3)), int(m.group(4)))

        dtype = m.group(1).replace('s', '<')
        data = np.arange(np.prod(dims)).reshape(dims).astype(dtype)

        tmpdir = tempfile.mkdtemp()
        try:
            testFile = path.join(tmpdir,path.basename(filename))
            f = FortranFile(testFile, 'w','<u4')
            f.write_record(data.T)
            f.close()
            originalfile = open(filename, 'rb')
            newfile = open(testFile, 'rb')
            assert_equal(originalfile.read(), newfile.read(),
                         err_msg=filename)
            originalfile.close()
            newfile.close()
        finally:
            shutil.rmtree(tmpdir)