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)

def export_FFT_field(phi,outfile):
    """
    export_FFT_field(phi,outfile)

    outfile: string
    phi[:,:,:] : np.array, rank 3

    This function takes in a 3-D field, of dimension
    Ngrid x Ngrid x Ngrid, and exports a half-field
    Ngrid/2+1 x Ngrid x Ngrid in unformatted
    Fortran record. Equivalent to a write(field) statement.
    """
    f=FortranFile(outfile,'w')
    n=phi.shape[2]
    f.write_record(np.array([n],dtype=np.int32)) # write integer record
    f.write_record(phi[:n//2+1].ravel(order='F')) # write half-field
    f.close()

 def export_unk(self, grid = [50,50,50]):
     '''
     Export the periodic part of BF in a real space grid for plotting with wannier90
     '''    
     
     from scipy.io import FortranFile
     grids_coor, weights = periodic_grid(self.cell, grid, order = 'F')        
     
     for k_id in range(self.num_kpts_loc):
         spin = '.1'
         if self.spin_up != None and self.spin_up == False : spin = '.2'
         kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])    
         ao = numint.eval_ao(self.cell, grids_coor, kpt = kpt)
         u_ao = np.einsum('x,xi->xi', np.exp(-1j*np.dot(grids_coor, kpt)), ao, optimize = True)
         unk_file = FortranFile('UNK' + "%05d" % (k_id + 1) + spin, 'w')
         unk_file.write_record(np.asarray([grid[0], grid[1], grid[2], k_id + 1, self.num_bands_loc], dtype = np.int32))    
         mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]        
         u_mo = np.einsum('xi,in->xn', u_ao, mo_included, optimize = True)
         for band in range(len(self.band_included_list)):    
             unk_file.write_record(np.asarray(u_mo[:,band], dtype = np.complex128))                    
         unk_file.close()

def main(sourcedir, sourcefile, savedir, saveheader, moving=0, incr=1, imin=0, imax=-1 ):
    """
    imin --> minimum iteration above which files will be written
    incr --> iteration interval to write files at
    """

    # ogdir = os.getcwd()
    # os.chdir(savedir)

    MakeOutputDir(savedir)

    #link in source file
    if not os.path.isfile('{}/{}'.format(savedir, sourcefile)):
        # cmd('ln -s {}/{} {}'.format(sourcedir, sourcefile, sourcefile))
        cmd('ln -s {}/{} {}/{}'.format(sourcedir, sourcefile, savedir, sourcefile))


    #GET GRID/FILE DIMENSIONS
    f = FortranFile('{}/{}'.format(savedir, sourcefile), 'r')
    dims = f.read_ints('uint32')[2:6]
    #Number of iterations, Number of points in each direction, number of parameters, needed to read whole file
    nj, nk, nl, nq = dims

    #FORMATS IN BC201 FILE
    bc201formats = [
                        ('ints', '<i4', (1,7)), #IGRID,ISTEP,NJ,NK,NL,NQ,NQC
                        ('tvr', 'float64'),     #TVREF
                        ('dtr', 'float64'),     #DTVREF
                        ('xyz', 'float64', (3,nj,nk,nl)), #XYZ points
                        ('q', 'float64', (nq,nj,nk,nl)),  #Q variables (nq variables, njXnkXnl values for each)
                        ('iblank', 'int32', (nj,nk,nl)),  #IBLANK status of each grid point
                    ]

    #########################
    #Save pressure coeff history at certian locations
    #! I want to write out a subset through time (every iteration!)
    ports = [8,13,30,45]
    out = open(savedir+"/cp_history.dat","w")
    out.write("ITER")
    for p in ports:
        out.write(" {}_x {}_cp".format(p,p))
    out.write("\n")
    #############################


    #RESTART FILE AND READ EACH RECORD
    f = FortranFile('{}/{}'.format(savedir, sourcefile), 'r')
    keep_reading = True
    irecord = -1
    while (keep_reading is True):

        irecord += 1

        #READ CURRENT RECORD
        b = f.read_record(bc201formats)



        #GET CURRENT ITERATION
        istep = b['ints'][0][0][1]

        # print(istep)


        #DONT WRITE UNDESIRED FILES (SKIP)
        #only write files at specified interval
        if (istep % incr != 0):
            continue
        #Don't write files below start iter
        if (istep < imin):
            continue
        #Stop reading when max desired iteration is reached
        if istep > imax - incr:
            keep_reading = False




        #WRITE GRID POINTS TO PLOT3D FORMATTED FILE

        if moving:
            #Grid is moving, write to file at each iteration
            savename = '{}/{}.x.{}'.format(savedir, saveheader, istep)
            WriteGrid(savename, b, nj, nk, nl)
        elif irecord == 0:
           #Grid is stationary, write to file only once
            savename = '{}/{}.x'.format(savedir, saveheader)
            WriteGrid(savename, b, nj, nk, nl)


        #CALCULATE CP FOR EACH POINT
        cps = np.zeros((nj,nk,nl))
        # print(cps.shape)
        for j in range(nj):
            for k in range(nk):
                for l in range(nl):
                    #print(j,k,l,b['q'].shape)
                    q = np.zeros(nq)
                    for iq in range(nq):
                        q[iq] = b['q'][0][iq][j][k][l]
#.........这里部分代码省略.........

def read_atomic_data(elements=['H', 'He', 'C',     # twelve most abundant elements
                               'N', 'O', 'Ne',
                               'Mg', 'Si', 'S', 
                               'Ar', 'Ca', 'Fe', ] , 
                     data_directory= 'sunnei/AtomicData',   # not robust!  Works when calling from the directory that sunnei is in
                     screen_output=False):

    '''
    This routine reads in the atomic data to be used for the
    non-equilibrium ionization calculations.
 
    Instructions for generating atomic data files
    =============================================
    
    The atomic data files are generated from the routines described by
    Shen et al. (2015) and are available at:
    
    https://github.com/ionizationcalc/time_dependent_fortran
    
    First, run the IDL routine 'pro_write_ionizrecomb_rate.pro' in the
    subdirectory sswidl_read_chianti with optional parameters: nte
    (number of temperature bins, default=501), te_low (low log
    temperature, default=4.0), and te_high (high log temperature,
    default=9.0) to get an ionization rate table.  The routine outputs
    the file "ionrecomb_rate.dat" which is a text file containing the
    ionization and recombination rates as a function of temperature.
    This routine requires the atomic database Chianti to be installed
    in IDL.

    Second, compile the Fortran routine 'create_eigenvmatrix.f90'.
    With the Intel mkl libraries it is compiled as: "ifort -mkl
    create_eigenvmatrix.f90 -o create.out" which can then be run with
    the command "./create.out".  This routine outputs all the
    eigenvalue tables for the first 28 elements (H to Ni).

    As of 2016 April 7, data from Chianti 8 is included in the
    CMEheat/AtomicData subdirectory.
    '''

    if screen_output:
        print('read_atomic_data: beginning program')
    
    from scipy.io import FortranFile

    '''
    Begin a loop to read in the atomic data files needed for the
    non-equilibrium ionization modeling.  The information will be
    stored in the atomic_data dictionary.

    For the first element in the loop, the information that should be
    the same for each element will be stored at the top level of the
    dictionary.  This includes the temperature grid, the number of
    temperatures, and the number of elements.

    For all elements, read in and store the arrays containing the
    equilibrium state, the eigenvalues, the eigenvectors, and the
    eigenvector inverses.
    '''

    atomic_data = {}
    
    first_element_in_loop = True

    for element in elements:

        if screen_output:
            print('read_atomic_data: '+element)

        AtomicNumber = AtomicNumbers[element]
        nstates = AtomicNumber + 1

        filename = data_directory + '/' + element.lower() + 'eigen.dat'
        H = FortranFile(filename, 'r')

        nte, nelems = H.read_ints(np.int32)
        temperatures = H.read_reals(np.float64)
        equistate = H.read_reals(np.float64).reshape((nte,nstates))
        eigenvalues = H.read_reals(np.float64).reshape((nte,nstates))
        eigenvector = H.read_reals(np.float64).reshape((nte,nstates,nstates))
        eigenvector_inv = H.read_reals(np.float64).reshape((nte,nstates,nstates))
        c_rate = H.read_reals(np.float64).reshape((nte,nstates))
        r_rate = H.read_reals(np.float64).reshape((nte,nstates))      
        
        if first_element_in_loop:
            atomic_data['nte'] = nte
            atomic_data['nelems'] = nelems  # Probably not used but store anyway
            atomic_data['temperatures'] = temperatures
            first_element_in_loop = False
        else: 
            assert nte == atomic_data['nte'], 'Atomic data files have different number of temperature levels: '+element
            assert nelems == atomic_data['nelems'], 'Atomic data files have different number of elements: '+element
            assert np.allclose(atomic_data['temperatures'],temperatures), 'Atomic data files have different temperature bins'

        atomic_data[element] = {'element':element,
                                'AtomicNumber':AtomicNumber,
                                'nstates':nstates,
                                'equistate':equistate,
                                'eigenvalues':eigenvalues,
                                'eigenvector':eigenvector,
                                'eigenvector_inv':eigenvector_inv,
#.........这里部分代码省略.........

dY3Right = Y3[:, -1] - Y3[:, -2]

# Write mean tau
WuResMeanValid = np.zeros((nlat * nlon,))
WuResMeanValid[~mask] = WuResMean
WuResMeanValid = WuResMeanValid.reshape(nlat, nlon)
tau = np.tile(np.expand_dims(WuResMeanValid, 0), (T, 1, 1)) * ampMean
dtaudx = np.empty(tau.shape)
dtaudx[:, :, 1:-1] = (tau[:, :, 2:] - tau[:, :, :-2]) / dX3Center
dtaudx[:, :, 0] = (tau[:, :, 1] - tau[:, :, 0]) / dX3Left
dtaudx[:, :, -1] = (tau[:, :, -1] - tau[:, :, -2]) / dX3Right
dtaudy = np.empty(tau.shape)
dtaudy[:, 1:-1] = (tau[:, 2:] - tau[:, :-2]) / dY3Center
dtaudy[:, 0] = (tau[:, 1] - tau[:, 0]) / dY3Left
dtaudy[:, -1] = (tau[:, -1] - tau[:, -2]) / dY3Right
f = FortranFile('%s/tau_mean_ampMean%02d%s.bin' % (initDir, int(ampMean * 10), postfix), 'w')
for t in np.arange(T):
    f.write_record(tau[t])
f.close()
f = FortranFile('%s/dtaudx_mean_ampMean%02d%s.bin' % (initDir, int(ampMean * 10), postfix), 'w')
for t in np.arange(T):
    f.write_record(dtaudx[t])
f.close()
f = FortranFile('%s/dtaudy_mean_ampMean%02d%s.bin' % (initDir, int(ampMean * 10), postfix), 'w')
for t in np.arange(T):
    f.write_record(dtaudy[t])
f.close()

# Write tau with variations
for k in np.arange(nEOF):
    print 'Writing tau with %d eofs...' % (k+1,)