def ReadArray_FortranBinary(filename, D):
"""
This function allows the user to read in a D-dimensional array
from unformatted Fortran binary.
Example input file format for 2D array:
Line | Entry | Data Type
------------------------------------
1 | nrows | int32
2 | ncols | int32
3 | A(1,1) | double
4 | A(2,1) | double
. | . | .
. | . | .
. | . | .
end | A(nrows,ncols)| double
"""
import numpy as np
from scipy.io import FortranFile
f = FortranFile(filename, 'r')
sz = np.zeros(D)
for i in range(0, D):
sz[i] = f.read_ints(dtype=np.int32)
A = f.read_reals(dtype=np.float64)
f.close()
A = A.reshape(sz.astype(int), order='F')
A = np.transpose(A)
return (A)
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 export_unk(self, spin=0, ngrid=None):
'''
Export the periodic part of BF in a real space grid for plotting with wannier90
'''
from scipy.io import FortranFile
if spin == 0:
spin_str = '.1'
else:
spin_str = '.2'
if ngrid is None:
ngrid = self.ngrid.copy()
else:
ngrid = np.array(ngrid, dtype=np.int64)
assert ngrid.shape[0] == 3, 'Wrong syntax for ngrid'
assert np.alltrue(ngrid >= self.ngrid), "Minium FT grid size: (%d, %d, %d)" % \
(self.ngrid[0], self.ngrid[1], self.ngrid[2])
for kpt in range(self.nkpts):
unk_file = FortranFile('UNK' + "%05d" % (kpt + 1) + spin_str, 'w')
unk_file.write_record(
np.asarray(
[ngrid[0], ngrid[1], ngrid[2], kpt + 1, self.nbands],
dtype=np.int32))
for band in range(self.nbands):
unk = self.get_unk(spin=spin,
kpt=kpt + 1,
band=band + 1,
ngrid=ngrid,
norm=False,
norm_c=False)
unk = unk.T.flatten()
unk_file.write_record(unk)
unk_file.close()
def read_patch(self, file, verbose=0):
''' Read particles in a patch '''
name = file + '.peb'
f = FortranFile(name, 'rb')
fmt = f.readInts()
if verbose > 2:
print('ioformat:', fmt[0])
dim = f.readInts()
dim = dim[0:2]
a = f.readReals()
a = a.reshape(dim[1], dim[0]).transpose()
idv = f.readInts()
idx = idv[0:dim[0]]
f.close()
dict = {}
for i in range(size(idx)):
id = idx[i]
dict[id] = {
'p': a[i, 0:3],
'v': a[i, 0:3],
't': a[i, 6],
's': a[i, 7],
'w': a[i, 8]
}
return idx, dict
def read_ibool_from_folder(folder, rank=0, ngll=5):
filename = data_filename(folder, "NSPEC_ibool", rank)
f = FortranFile(filename, "r")
nspec = f.read_ints()[0]
d = f.read_record("({},{},{})i4".format(nspec, ngll, ngll))
f.close()
return d
def __init__(self, seedname='wannier90', formatted=False, suffix='sHu'):
print("----------\n {0} \n---------".format(suffix))
if formatted:
f_sXu_in = open(seedname + "." + suffix, 'r')
header = f_sXu_in.readline().strip()
NB, NK, NNB = (int(x) for x in f_sXu_in.readline().split())
else:
f_sXu_in = FortranFile(seedname + "." + suffix, 'r')
header = readstr(f_sXu_in)
NB, NK, NNB = f_sXu_in.read_record('i4')
print("reading {}.{} : <{}>".format(seedname, suffix, header))
self.data = np.zeros((NK, NNB, 3, NB, NB), dtype=complex)
for ik in range(NK):
# print ("k-point {} of {}".format( ik+1,NK))
for ib2 in range(NNB):
for ipol in range(3):
tmp = f_sXu_in.read_record('f8').reshape(
(2, NB, NB), order='F').transpose(2, 1, 0)
self.data[ik, ib2, ipol] = tmp[:, :, 0] + 1j * tmp[:, :, 1]
print("----------\n {0} OK \n---------\n".format(suffix))
f_sXu_in.close()
def get_inverse_covmat(self, spectra):
'''
high ell TT, TE and EE
from Planck plik-lite datafile
'''
f = FortranFile(self.cov_file, 'r')
covmat = f.read_reals(dtype=float).reshape((self.nbin_hi,self.nbin_hi))
f.close()
for i in range(self.nbin_hi):
for j in range(i,self.nbin_hi):
covmat[i,j] = covmat[j,i]
if spectra=='TT':
#select relevant covmat for temperature only
bin_no=self.nbintt_hi
start=0
end=start+bin_no
cov=covmat[start:end, start:end]
else: #TTTEEE
bin_no=self.nbin_hi
cov=covmat
#invert high ell covariance matrix (cholesky decomposition should be faster)
fisher=scipy.linalg.cho_solve(scipy.linalg.cho_factor(cov), np.identity(bin_no))
fisher=fisher.transpose()
return fisher
def readKernel(self):
'''
Read the EDKS Kernel and stores it in {self}
Returns:
* None
'''
# Show me
if self.verbose:
print('Read Kernel file {}'.format(self.kernel))
# Open the file
fedks = FortranFile(self.kernel, 'r')
# Read
kernel = fedks.read_reals(np.float32).reshape(
(self.ndepth * self.ndista, 12))
# Save
self.depths = kernel[:, 0]
self.distas = kernel[:, 1]
self.zrtdsx = kernel[:, 2:]
# CLose file
fedks.close()
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 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 = general_grid(self.cell, grid)
for k_id in range(self.num_kpts_loc):
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('UNK0000' + str(k_id + 1) + '.1', '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 = np.dot(
self.U[k_id],
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.complex))
unk_file.close()
def fits_to_unformatted(
input_filename, output_filename, omega_m
):
# define cosmology
cosmo = Cosmology(omega_m=omega_m)
# open fits file
with fits.open(input_filename) as hdul:
cat = hdul[1].data
# convert redshifts to distances
dist = cosmo.ComovingDistance(cat['Z'])
x = dist * np.sin(cat['DEC'] * np.pi / 180) * np.cos(cat['RA'] * np.pi / 180)
y = dist * np.sin(cat['DEC'] * np.pi / 180) * np.sin(cat['RA'] * np.pi / 180)
z = dist * np.cos(cat['DEC'] * np.pi / 180)
# write result to output file
cout = np.c_[x, y, z]
nrows, ncols = np.shape(cout)
f = FortranFile(output_filename, 'w')
nrows, ncols = np.shape(cout)
f.write_record(nrows)
f.write_record(ncols)
f.write_record(cout)
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 read_binary_fortran_file(name,
datatype,
dim0,
dim1=1,
real_precision=np.float64):
if dim1 == 1:
fmat = FortranFile(name, 'r')
input_array = fmat.read_reals(dtype=real_precision)
fmat.close()
else:
if datatype == "real":
fmat = FortranFile(name, 'r')
input_array = fmat.read_reals(dtype=real_precision)
input_array = input_array.reshape((dim0, dim1)).transpose()
fmat.close()
elif datatype == "int":
fmat = FortranFile(name, 'r')
input_array = fmat.read_ints(dtype=np.int32)
input_array = input_array.reshape((dim0, dim1)).transpose()
# fmat.close()
elif datatype == "complex":
sys.exit(
"reading of complex matrices is not directly possible, must write out real and imag parts "
"as two separate real matrices\n")
else:
sys.exit("reading of datatype \"" + datatype +
"\" is not implemented\n ")
return input_array
def revolver_to_unformatted(input_filename,
output_filename,
cosmology,
equal_weights=False,
zrange=None):
# open numpy file
cat = np.load(input_filename)
if zrange is not None:
zmin, zmax = zrange
ind = (cat[:, 2] > zmin) & (cat[:, 2] < zmax)
cat = cat[ind]
# convert redshifts to distances
dist = cosmology.ComovingDistance(cat[:, 2])
x = dist * np.cos(cat[:, 1] * np.pi / 180) * np.cos(
cat[:, 0] * np.pi / 180)
y = dist * np.cos(cat[:, 1] * np.pi / 180) * np.sin(
cat[:, 0] * np.pi / 180)
z = dist * np.sin(cat[:, 1] * np.pi / 180)
if not equal_weights:
weight = cat[:, 3]
else:
weight = np.ones(len(cat))
#write result to output file
cout = np.c_[x, y, z, weight]
nrows, ncols = np.shape(cout)
f = FortranFile(output_filename, 'w')
nrows, ncols = np.shape(cout)
f.write_record(nrows)
f.write_record(ncols)
f.write_record(cout)
f.close()
def __getMat(self, suffix):
f = FF(self.seedname + "_" + suffix + "_R" +
(".dat" if self.old_format else ""))
MM_R = np.array([[
np.array(f.read_record('2f8'), dtype=float)
for m in range(self.num_wann)
] for n in range(self.num_wann)])
MM_R = MM_R[:, :, :, 0] + 1j * MM_R[:, :, :, 1]
f.close()
ncomp = MM_R.shape[2] / self.nRvec0
if ncomp == 1:
result = MM_R / self.Ndegen[None, None, :]
elif ncomp == 3:
result = MM_R.reshape(
self.num_wann, self.num_wann, 3, self.nRvec0).transpose(
0, 1, 3, 2) / self.Ndegen[None, None, :, None]
elif ncomp == 9:
result = MM_R.reshape(
self.num_wann, self.num_wann, 3, 3, self.nRvec0).transpose(
0, 1, 4, 3, 2) / self.Ndegen[None, None, :, None, None]
else:
raise RuntimeError(
"in __getMat: invalid ncomp : {0}".format(ncomp))
if self.ws_map is None:
return result
else:
return self.ws_map(result)
def read_data_matrix_from_folder(folder, param, shape, rank=0):
filename = data_filename(folder, param, rank)
f = FortranFile(filename, "r")
# nspec = f.read_ints()[0]
d = f.read_record("({},{},{})f4".format(*shape))
f.close()
return d
def read_output(path, header_only=True):
f = FortranFile(path, 'r')
ncpu = f.read_ints()
dim = f.read_ints()
nparts = f.read_ints()
if header_only:
f.close()
return ncpu, dim, nparts
f.read_ints()
f.read_ints()
f.read_ints()
f.read_ints()
f.read_ints()
x = f.read_reals(dtype=np.float64)
y = f.read_reals(dtype=np.float64)
z = f.read_reals(dtype=np.float64)
vx = f.read_reals(dtype=np.float64)
vy = f.read_reals(dtype=np.float64)
vz = f.read_reals(dtype=np.float64)
m = f.read_reals(dtype=np.float64)
part_ids = f.read_ints()
birth = f.read_reals(dtype=np.float32)
f.close()
return ncpu, dim, nparts, x, y, z, part_ids
def generate_positions(data_filename, centres_filename, npositions, sampling,
boxsize):
'''
Generates random points on a box
writes them to an unformatted
Fortran 90 file.
'''
np.random.seed(0)
if sampling == 'tracers':
print('Randoms will be generated from galaxy positions.')
fin = FortranFile(data_filename, 'r')
nrows = fin.read_ints()[0]
ncols = fin.read_ints()[0]
pos = fin.read_reals(dtype=np.float64).reshape(nrows, ncols)
idx = np.random.choice(nrows, size=npositions, replace=False)
cout = pos[idx]
elif sampling == 'uniform':
print('Randoms will be generated from a uniform distribution.')
x = np.random.uniform(0, boxsize, npositions)
y = np.random.uniform(0, boxsize, npositions)
z = np.random.uniform(0, boxsize, npositions)
cout = np.c_[x, y, z]
else:
sys.exit('Sampling type not recognized')
cout = cout.astype('float64')
f = FortranFile(centres_filename, 'w')
nrows, ncols = np.shape(cout)
f.write_record(nrows)
f.write_record(ncols)
f.write_record(cout)
f.close()
def _counts_Bk123_f77(Ngrid=360, Nmax=40, Ncut=3, step=3, silent=True):
''' return bispectrum normalization
@chh explain nmax, ncut, and step below
'''
fcnt = ''.join([
'counts', '.Ngrid',
str(Ngrid), '.Nmax',
str(Nmax), '.Ncut',
str(Ncut), '.step',
str(step), '.fort77'
])
f_counts = os.path.join(dat_dir(), fcnt)
if os.path.isfile(f_counts):
f = FortranFile(f_counts, 'r')
cnts = f.read_reals(dtype=np.float64)
counts = np.reshape(cnts, (Nmax, Nmax, Nmax), order='F')
f.close()
else:
if not silent:
print('-- %s does not exist --' % f_counts)
print('-- computing %s --' % f_counts)
counts = np.zeros((40, 40, 40), dtype=np.float64, order='F')
fEstimate.bk_counts(counts, Ngrid, float(step), Ncut, Nmax)
# save to file
f = FortranFile(f_counts, 'w')
f.write_record(counts) # double prec
f.close()
return counts
def get_circumcentres(self, radius_limit=1000):
'''
Find the centre of the circumspheres
associated to an input catalogue of
tetrahedra.
'''
print('Finding circumcentres of tetrahedra...')
vertices = self.delaunay_triangulation()
cenx, ceny, cenz, r = [], [], [], []
sing = 0
for tetra in vertices:
x0, x1, x2, x3 = tetra
A = []
B = []
A.append((x1 - x0).T)
A.append((x2 - x0).T)
A.append((x3 - x0).T)
A = np.asarray(A)
B = np.sum(A**2, axis=1)
B = np.asarray(B)
try:
C = np.linalg.inv(A).dot(B)
except:
sing += 1
continue
centre = x0 + 0.5 * C
radius = 0.5 * np.sqrt(np.sum(C**2))
if radius < radius_limit:
cenx.append(centre[0])
ceny.append(centre[1])
cenz.append(centre[2])
r.append(radius)
print('{} singular matrices found.'.format(sing))
cenx = np.asarray(cenx)
ceny = np.asarray(ceny)
cenz = np.asarray(cenz)
cenx = cenx.reshape(len(cenx), 1)
ceny = ceny.reshape(len(ceny), 1)
cenz = cenz.reshape(len(cenz), 1)
cout = np.hstack([cenx, ceny, cenz])
print('{} centres found.'.format(len(cout)))
f = FortranFile(self.centres_file, 'w')
nrows, ncols = np.shape(cout)
f.write_record(nrows)
f.write_record(ncols)
f.write_record(cout)
f.close()
self.centres = cout
return
def writef90map(fname, d1):
"""
Save data to a fortran file
"""
f = FortranFile(fname, 'w')
#f.write_record(np.int16(dim))
f.write_record(np.float64(d1))
f.close()
def filtered_density(data_filename1,
data_filename2,
output_filename,
filter_type,
filter_size,
ngrid,
box_size,
nthreads=1,
dim1_min=0,
dim1_max=None,
output_format='unformatted'):
# check if files exist
if not path.isfile(data_filename1):
raise FileNotFoundError(f'{data_filename1} does not exist.')
if not path.isfile(data_filename2):
raise FileNotFoundError(f'{data_filename2} does not exist.')
if dim1_max == None:
if filter_type == 'tophat':
dim1_max = filter_size
elif filter_type == 'gaussian':
dim1_max = 5 * filter_size
binpath = path.join(path.dirname(__file__), 'bin',
'{}_filter.exe'.format(filter_type))
cmd = [
binpath, data_filename1, data_filename2, output_filename,
str(box_size),
str(dim1_min),
str(dim1_max),
str(filter_size),
str(ngrid),
str(nthreads)
]
subprocess.call(cmd)
# open filter file
f = FortranFile(output_filename, 'r')
smoothed_delta = f.read_ints()[0]
smoothed_delta = f.read_reals(dtype=np.float64)
f.close()
if output_format != 'unformatted':
if output_format == 'npy':
subprocess.call(['rm', output_filename])
np.save(output_filename, smoothed_delta)
elif output_format == 'ascii':
np.savetxt(output_filename, smoothed_delta)
else:
print('Output format not recognized. Using unformatted F90 file.')
return smoothed_delta
def ascii_to_unformatted(input_filename, output_filename ): # import data cout = np.genfromtxt(input_filename) f = FortranFile(output_filename, 'w') nrows, ncols = np.shape(cout) f.write_record(nrows) f.write_record(ncols) f.write_record(cout) f.close()
def write(self):
"""
write the binary file.
"""
f = FortranFile(self.fname, 'w')
nchan = len(self.channelIdx)
f.write_record(np.array([nchan, 8], dtype=np.int32))
f.write_record(self.channelIdx)
f.write_record(self.R.T)
f.close()
def _get_solution(path_to_solution):
print(path_to_solution)
f = FortranFile(path_to_solution)
solution = f.read_record(np.complex_)
f.close()
dim = int(np.sqrt(solution.shape))
solution = np.reshape(solution, (dim, dim))
return solution
def export_unk(self, spin=0, ngrid=None, only_irred=False):
'''
Export the periodic part of BF in a real space grid for plotting with wannier90
'''
if self.lsorbit:
spin_str = '.NC'
spin = 0
elif spin == 0:
spin_str = '.1'
else:
spin_str = '.2'
if ngrid is None:
ngrid = self.ngrid.copy()
else:
ngrid = np.array(ngrid, dtype=np.int64)
assert ngrid.shape[0] == 3, 'Wrong syntax for ngrid'
assert np.alltrue(ngrid >= self.ngrid), "Minium FT grid size: (%d, %d, %d)" % \
(self.ngrid[0], self.ngrid[1], self.ngrid[2])
if only_irred:
nkpts = self.nkpts
unk_list = self.get_unk_kpts(spin, ngrid=ngrid)
else:
kpts, band, unk_list = self.get_wave_nosym(spin, ngrid=ngrid)
nkpts = kpts.shape[0]
if self.lsorbit:
for kpt in range(nkpts):
unk_file = FortranFile('UNK' + "%05d" % (kpt + 1) + spin_str,
'w')
unk_file.write_record(
np.asarray(
[ngrid[0], ngrid[1], ngrid[2], kpt + 1, self.nbands],
dtype=np.int32))
for band in range(self.nbands):
unk_up = unk_list[kpt][band, :ngrid[0], :, :].T.flatten()
unk_down = unk_list[kpt][band, ngrid[0]:, :, :].T.flatten()
unk_file.write_record(unk_up)
unk_file.write_record(unk_down)
unk_file.close()
else:
for kpt in range(nkpts):
unk_file = FortranFile('UNK' + "%05d" % (kpt + 1) + spin_str,
'w')
unk_file.write_record(
np.asarray(
[ngrid[0], ngrid[1], ngrid[2], kpt + 1, self.nbands],
dtype=np.int32))
for band in range(self.nbands):
unk = unk_list[kpt][band].T.flatten()
unk_file.write_record(unk)
unk_file.close()
def read(self):
"""
read data fortran binary.
"""
f = FortranFile(self.fname)
nchan, type = f.read_ints(np.int32)
self.channelIdx = f.read_ints(np.int32)
if(type == 8):
self.R = f.read_reals(np.float64).reshape((nchan,nchan), order="F")
else:
self.R = f.read_reals(np.float32).reshape((nchan,nchan), order="F")
f.close()
def tmp():
ff = FortranFile(filename)
h = {}
h["nbodies"] = ff.read_ints()
h["massp"] = ff.read_ints()
h["aexp"] = ff.read_reals(dtype=np.int32)
h["omega_t"] = ff.read_reals(dtype=np.int32)
h["age_univ"] = ff.read_reals(dtype=np.int32)
h["n_halos"], h["n_subhalos"] = ff.read_ints()
for i in tqdm(range(h["n_halos"] + h["n_subhalos"])):
infos = {
"header": h
}
infos["nparts"] = ff.read_ints()
infos["members"] = ff.read_ints()
infos["idh"] = ff.read_ints()
infos["timestep"] = ff.read_ints()
infos["hlevel"], infos["hosthalo"], infos["hostsub"], infos["nbsub"], infos["nextsub"] = ff.read_ints()
infos["mhalo"] = ff.read_reals(dtype=np.int32)
infos["pos"] = ff.read_reals(dtype=np.int32)
infos["speed"] = ff.read_reals(dtype=np.int32)
infos["L"] = ff.read_reals(dtype=np.int32)
infos["r"], infos["a"], infos["b"], infos["c"] = ff.read_reals(dtype=np.int32)
infos["ek"], infos["ep"], infos["et"] = ff.read_reals(dtype=np.int32)
infos["spin"] = ff.read_reals(dtype=np.int32)
if not dm_type:
ff.read_reals()
infos["rvir"],infos["mvir"], infos["tvir"], infos["cvel"] = ff.read_reals(dtype=np.int32)
ff.read_reals()
if not dm_type:
infos["npoints"] = ff.read_ints()
infos["rdum"] = ff.read_reals(dtype=np.int32)
infos["density"] = ff.read_reals(dtype=np.int32)
if low_mem != None:
try:
keys = list(low_mem)
except:
keys = ['nparts', 'members']
tmp = {}
for key in keys:
try:
tmp[key] = infos[key]
except KeyError:
print('Invalid key {}, can be any of', infos['keys'])
yield tmp
else:
yield infos
ff.close()
def read_cproj_NormalCar(inf='NormalCAR', save_cproj=True):
'''
Read NormalCAR of VASP output, which contains the coefficients of the PAW
projector functions.
Data stored in NormalCAR:
WRITE(IU) LMDIM,WDES%NIONS,WDES%NRSPINORS
WRITE(IU) CQIJ(1:LMDIM,1:LMDIM,1:WDES%NIONS,1:WDES%NRSPINORS)
WRITE(IU) WDES%NPROD, WDES%NPRO, WDES%NTYP
DO NT = 1, WDES%NTYP
WRITE(IU) WDES%LMMAX(NT), WDES%NITYP(NT)
END DO
DO ISPIN=1,WDES%ISPIN
DO NK=1,WDES%NKPTS
DO N=1,WDES%NB_TOT
WRITE(IU) CPROJ(1:WDES1%NPRO_TOT)
END DO
END DO
END DO
'''
from scipy.io import FortranFile
ncr = FortranFile(inf, 'r')
# rec1 = ncr.read_record(dtype=np.int32)
rec1 = ncr.read_ints()
lmdim, nions, nrspinors = rec1
# rec2 = ncr.read_record(dtype=np.complex)
cqij = np.array(ncr.read_reals()).reshape((lmdim, lmdim, nions, nrspinors),
order='F')
rec3 = ncr.read_ints()
nprod, npro, ntyp = rec3
# lmmax, natoms for each type of elements
lmmax_per_typ = np.array([ncr.read_ints() for ii in range(ntyp)])
cproj = []
while True:
try:
rec = ncr.read_record(dtype=np.complex)
cproj.append(rec)
except:
break
cproj = np.array(cproj)
if save_cproj:
np.save('cproj', cproj)
ncr.close()
return cproj
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_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 reorder_uXu(ext, formatted=False):
try:
print "----------\n {0} \n----------".format(ext)
if formatted:
f_uXu_in = open(seedname + "." + ext, 'r')
f_uXu_out = open(seednameGW + "." + ext, 'w')
header = f_uXu_in.readline()
f_uXu_out.write(header)
nbnd, NK, nnb = np.array(f_uXu_in.readline().split(), dtype=int)
f_uXu_out.write(" ".join(str(x) for x in [NBND, NK, nnb]) + "\n")
else:
f_uXu_in = FortranFile(seedname + "." + ext, 'r')
f_uXu_out = FortranFile(seednameGW + "." + ext, 'w')
header = f_uXu_in.read_record(dtype='c')
f_uXu_out.write_record(header)
nbnd, NK, nnb = np.array(f_uXu_in.read_record(dtype=np.int32))
f_uXu_out.write_record(np.array([NBND, NK, nnb], dtype=np.int32))
assert nbnd == nbndDFT
print nbnd, NK, nnb
if formatted:
uXu = np.loadtxt(f_uXu_in).reshape(-1)
start = 0
length = nbnd * nbnd
for ik in xrange(NKPT):
for ib2 in xrange(nnb):
for ib1 in xrange(nnb):
if formatted:
A = uXu[start:start + length]
start += length
else:
A = f_uXu_in.read_record(dtype=np.complex)
A = (A.reshape(nbnd, nbnd, order='F')[
BANDSORT[KPNB[ik][ib2]], :][:, BANDSORT[KPNB[ik][ib1]]]
+ np.einsum('ln,lm,l->nm', MMN[ik][ib2].conj(),
MMN[ik][ib1], eigenDE[ik])).reshape(
-1, order='F')
if formatted:
f_uXu_out.write("".join(
"{0:26.16e} {1:26.16f}\n".format(x.real, x.imag)
for x in A))
else:
f_uXu_out.write_record(A)
f_uXu_out.close()
f_uXu_in.close()
print "----------\n {0} OK \n----------\n".format(ext)
except IOError as err:
print "WARNING: {0}.{1} not written : ".format(seednameGW, ext), err
def amie_write_binary(file, data):
fp = FortranFile(file, 'w')
iVals = [data["nLats"], data["nMlts"], data["nTimes"]]
fp.write_record(np.array(iVals, dtype=np.int32))
coLats = np.array(90.0 - data["lats"], dtype=np.float32)
mlts = np.array(data["mlts"], dtype=np.float32)
fp.write_record(coLats)
fp.write_record(mlts)
# Write out variables!
val = np.array([data["nVars"]], dtype=np.int32)
fp.write_record(val)
for var in data["Vars"]:
varpad = var.ljust(30).encode('utf-8')
fp.write_record(varpad)
for iT in np.arange(data["nTimes"]):
iymdhm = np.array([iT, \
data["times"][iT].year, \
data["times"][iT].month, \
data["times"][iT].day, \
data["times"][iT].hour, \
data["times"][iT].minute], \
dtype=np.int32)
fp.write_record(iymdhm)
indices = np.array(np.concatenate((data["imf"][iT], \
data["ae"][iT], \
data["dst"][iT], \
data["hp"][iT], \
[data["cpcp"][iT]])), dtype = np.float32)
fp.write_record(indices)
for iVar in np.arange(data["nVars"]):
v = data["Vars"][iVar]
vals = np.array(data[v][iT], dtype=np.float32)
fp.write_record(vals)
val = np.array([data["version"]], dtype=np.float32)
fp.write_record(val)
fp.close()
def particles_in_halo(tree_brick, start=0, end=None, fun_filter=lambda x: True):
''' Open a tree bricks file and associate to each halo the corresponding particles.
'''
# Open file
f = FortranFile(tree_brick, 'r')
# Give a value to end, by default start + 1
if end == None:
end = start + 1
# Read headers
nbodies = f.read_ints()[0]
f.read_reals(dtype=np.float32)
aexp = f.read_reals(dtype=np.float32)
f.read_reals(dtype=np.float32)
age = f.read_reals(dtype=np.float32)
nhalo, nsubhalo = f.read_ints()
halo_tot = nhalo + nsubhalo
halos = {}
for i in tqdm(range(halo_tot)):
parts = f.read_ints()[0]
members = f.read_ints()
this_id = f.read_ints()[0]
if (start <= this_id and this_id < end and fun_filter(this_id)):
for dm_particle_id in members:
if not halos.has_key(this_id):
halos[this_id] = []
halos[this_id].append(dm_particle_id)
elif this_id >= end:
break
f.read_ints()
# Irrelevant
level, hosthalo, hostsub, nbsub, nextsub = f.read_ints()
mstar = 1e11 * f.read_reals(dtype=np.float32)
px, py, pz = f.read_reals(dtype=np.float32)
f.read_reals(dtype=np.float32)
f.read_reals(dtype=np.float32)
rad = f.read_reals(dtype=np.float32)[0]
f.read_reals(dtype=np.float32)
f.read_reals(dtype=np.float32)
rvir, mvir, tvir, cvel = f.read_reals(dtype=np.float32)
f.read_reals(dtype=np.float32)
f.close()
return halos
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 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 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 reorder_uXu(ext,formatted=False):
try:
print "----------\n {0} \n----------".format(ext)
if formatted:
f_uXu_in = open(seedname+"."+ext, 'r')
f_uXu_out = open(seednameGW+"."+ext, 'w')
header=f_uXu_in.readline()
f_uXu_out.write(header)
nbnd,NK,nnb=np.array(f_uXu_in.readline().split(),dtype=int)
f_uXu_out.write(" ".join(str(x) for x in [NBND,NK,nnb])+"\n")
else:
f_uXu_in = FortranFile(seedname+"."+ext, 'r')
f_uXu_out = FortranFile(seednameGW+"."+ext, 'w')
header=f_uXu_in.read_record(dtype='c')
f_uXu_out.write_record(header)
nbnd,NK,nnb=np.array(f_uXu_in.read_record(dtype=np.int32))
f_uXu_out.write_record(np.array([NBND,NK,nnb],dtype=np.int32))
assert nbnd==nbndDFT
print nbnd,NK,nnb
if formatted:
uXu=np.loadtxt(f_uXu_in).reshape(-1)
start=0
length=nbnd*nbnd
for ik in xrange(NKPT):
for ib2 in xrange(nnb):
for ib1 in xrange(nnb):
if formatted:
A=uXu[start:start+length]
start+=length
else:
A=f_uXu_in.read_record(dtype=np.complex)
A=(A.reshape(nbnd,nbnd,order='F')[BANDSORT[KPNB[ik][ib2]],:][:,BANDSORT[KPNB[ik][ib1]]]+
np.einsum('ln,lm,l->nm',MMN[ik][ib2].conj(),MMN[ik][ib1],eigenDE[ik]) ).reshape(-1,order='F')
if formatted:
f_uXu_out.write("".join("{0:26.16e} {1:26.16f}\n".format(x.real,x.imag) for x in A))
else:
f_uXu_out.write_record(A)
f_uXu_out.close()
f_uXu_in.close()
print "----------\n {0} OK \n----------\n".format(ext)
except IOError as err:
print "WARNING: {0}.{1} not written : ".format(seednameGW,ext),err
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 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)
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,)
eof = np.zeros((nlat*nlon,))
eof[~mask] = v[:, k]
eof = eof.reshape(nlat, nlon)
def read(self, var_file='', datadir='data', proc=-1, ivar=-1, quiet=True,
trimall=False, magic=None, sim=None, precision='f'):
"""
Read VAR files from Pencil Code. If proc < 0, then load all data
and assemble, otherwise load VAR file from specified processor.
The file format written by output() (and used, e.g. in var.dat)
consists of the followinig Fortran records:
1. data(mx, my, mz, nvar)
2. t(1), x(mx), y(my), z(mz), dx(1), dy(1), dz(1), deltay(1)
Here nvar denotes the number of slots, i.e. 1 for one scalar field, 3
for one vector field, 8 for var.dat in the case of MHD with entropy.
but, deltay(1) is only there if lshear is on! need to know parameters.
call signature:
var(var_file='', datadir='data', proc=-1, ivar=-1, quiet=True,
trimall=False, magic=None, sim=None, precision='f')
Keyword arguments:
var_file: Name of the VAR file.
datadir: Directory where the data is stored.
proc: Processor to be read. If -1 read all and assemble to one array.
ivar: Index of the VAR file, if var_file is not specified.
quiet: Flag for switching off output.
trimall: Trim the data cube to exclude ghost zones.
magic: Values to be computed from the data, e.g. B = curl(A).
sim: Simulation sim object.
precision: Float (f) or double (d).
"""
import numpy as np
import os
from scipy.io import FortranFile
from ..math.derivatives import curl, curl2
from .. import read
from ..sim import __Simulation__
dim = None
param = None
index = None
if isinstance(sim, __Simulation__):
datadir = os.path.expanduser(sim.datadir)
dim = sim.dim
param = read.param(datadir=sim.datadir, quiet=True)
index = read.index(datadir=sim.datadir)
else:
datadir = os.path.expanduser(datadir)
if dim is None:
if var_file[0:2].lower() == 'og':
dim = read.ogdim(datadir, proc)
else:
dim = read.dim(datadir, proc)
if param is None:
param = read.param(datadir=datadir, quiet=quiet)
if index is None:
index = read.index(datadir=datadir)
run2D = param.lwrite_2d
if dim.precision == 'D':
precision = 'd'
else:
precision = 'f'
if param.lwrite_aux:
total_vars = dim.mvar + dim.maux
else:
total_vars = dim.mvar
if not var_file:
if ivar < 0:
var_file = 'var.dat'
else:
var_file = 'VAR' + str(ivar)
if proc < 0:
proc_dirs = self.__natural_sort(filter(lambda s: s.startswith('proc'),
os.listdir(datadir)))
else:
proc_dirs = ['proc' + str(proc)]
# Set up the global array.
if not run2D:
f = np.zeros((total_vars, dim.mz, dim.my, dim.mx),
dtype=precision)
else:
if dim.ny == 1:
f = np.zeros((total_vars, dim.mz, dim.mx), dtype=precision)
else:
f = np.zeros((total_vars, dim.my, dim.mx), dtype=precision)
x = np.zeros(dim.mx, dtype=precision)
y = np.zeros(dim.my, dtype=precision)
z = np.zeros(dim.mz, dtype=precision)
for directory in proc_dirs:
proc = int(directory[4:])
if var_file[0:2].lower() == 'og':
procdim = read.ogdim(datadir, proc)
else:
procdim = read.dim(datadir, proc)
if not quiet:
print("Reading data from processor {0} of {1} ...".format( \
proc, len(proc_dirs)))
mxloc = procdim.mx
myloc = procdim.my
mzloc = procdim.mz
# Read the data.
file_name = os.path.join(datadir, directory, var_file)
infile = FortranFile(file_name)
if not run2D:
f_loc = infile.read_record(dtype=precision)
f_loc = f_loc.reshape((-1, mzloc, myloc, mxloc))
else:
if dim.ny == 1:
f_loc = infile.read_record(dtype=precision)
f_loc = f_loc.reshape((-1, mzloc, mxloc))
else:
f_loc = infile.read_record(dtype=precision)
f_loc = f_loc.reshape((-1, myloc, mxloc))
raw_etc = infile.read_record(precision)
infile.close()
t = raw_etc[0]
x_loc = raw_etc[1:mxloc+1]
y_loc = raw_etc[mxloc+1:mxloc+myloc+1]
z_loc = raw_etc[mxloc+myloc+1:mxloc+myloc+mzloc+1]
if param.lshear:
shear_offset = 1
deltay = raw_etc[-1]
else:
shear_offset = 0
dx = raw_etc[-3-shear_offset]
dy = raw_etc[-2-shear_offset]
dz = raw_etc[-1-shear_offset]
if len(proc_dirs) > 1:
# Calculate where the local processor will go in
# the global array.
#
# Don't overwrite ghost zones of processor to the left (and
# accordingly in y and z direction -- makes a difference on the
# diagonals)
#
# Recall that in NumPy, slicing is NON-INCLUSIVE on the right end
# ie, x[0:4] will slice all of a 4-digit array, not produce
# an error like in idl.
if procdim.ipx == 0:
i0x = 0
i1x = i0x + procdim.mx
i0xloc = 0
i1xloc = procdim.mx
else:
i0x = procdim.ipx*procdim.nx + procdim.nghostx
i1x = i0x + procdim.mx - procdim.nghostx
i0xloc = procdim.nghostx
i1xloc = procdim.mx
if procdim.ipy == 0:
i0y = 0
i1y = i0y + procdim.my
i0yloc = 0
i1yloc = procdim.my
else:
i0y = procdim.ipy*procdim.ny + procdim.nghosty
i1y = i0y + procdim.my - procdim.nghosty
i0yloc = procdim.nghosty
i1yloc = procdim.my
if procdim.ipz == 0:
i0z = 0
i1z = i0z+procdim.mz
i0zloc = 0
i1zloc = procdim.mz
else:
i0z = procdim.ipz*procdim.nz + procdim.nghostz
i1z = i0z + procdim.mz - procdim.nghostz
i0zloc = procdim.nghostz
i1zloc = procdim.mz
x[i0x:i1x] = x_loc[i0xloc:i1xloc]
y[i0y:i1y] = y_loc[i0yloc:i1yloc]
z[i0z:i1z] = z_loc[i0zloc:i1zloc]
if not run2D:
f[:, i0z:i1z, i0y:i1y, i0x:i1x] = \
f_loc[:, i0zloc:i1zloc, i0yloc:i1yloc, i0xloc:i1xloc]
else:
if dim.ny == 1:
f[:, i0z:i1z, i0x:i1x] = \
f_loc[:, i0zloc:i1zloc, i0xloc:i1xloc]
else:
f[:, i0y:i1y, i0x:i1x] = \
f_loc[:, i0yloc:i1yloc, i0xloc:i1xloc]
else:
f = f_loc
x = x_loc
y = y_loc
z = z_loc
if magic is not None:
if 'bb' in magic:
# Compute the magnetic field before doing trimall.
aa = f[index.ax-1:index.az, ...]
# TODO: Specify coordinate system.
self.bb = curl(aa, dx, dy, dz, run2D=run2D)
if trimall:
self.bb = self.bb[:, dim.n1:dim.n2+1,
dim.m1:dim.m2+1, dim.l1:dim.l2+1]
if 'jj' in magic:
# Compute the electric current field before doing trimall.
aa = f[index.ax-1:index.az, ...]
# TODO: Specify coordinate system.
self.jj = curl2(aa, dx, dy, dz)
if trimall:
self.jj = self.jj[:, dim.n1:dim.n2+1,
dim.m1:dim.m2+1, dim.l1:dim.l2+1]
if 'vort' in magic:
# Compute the vorticity field before doing trimall.
uu = f[index.ux-1:index.uz, ...]
# TODO: Specify coordinate system.
# WL: The curl subroutine should take care of it.
self.vort = curl(uu, dx, dy, dz, run2D=run2D)
if trimall:
if run2D:
if dim.nz == 1:
self.vort = self.vort[:, dim.m1:dim.m2+1,
dim.l1:dim.l2+1]
else:
self.vort = self.vort[:, dim.n1:dim.n2+1,
dim.l1:dim.l2+1]
else:
self.vort = self.vort[:, dim.n1:dim.n2+1,
dim.m1:dim.m2+1,
dim.l1:dim.l2+1]
# Trim the ghost zones of the global f-array if asked.
if trimall:
self.x = x[dim.l1:dim.l2+1]
self.y = y[dim.m1:dim.m2+1]
self.z = z[dim.n1:dim.n2+1]
if not run2D:
self.f = f[:, dim.n1:dim.n2+1, dim.m1:dim.m2+1, dim.l1:dim.l2+1]
else:
if dim.ny == 1:
self.f = f[:, dim.n1:dim.n2+1, dim.l1:dim.l2+1]
else:
self.f = f[:, dim.m1:dim.m2+1, dim.l1:dim.l2+1]
else:
self.x = x
self.y = y
self.z = z
self.f = f
self.l1 = dim.l1
self.l2 = dim.l2 + 1
self.m1 = dim.m1
self.m2 = dim.m2 + 1
self.n1 = dim.n1
self.n2 = dim.n2 + 1
# Assign an attribute to self for each variable defined in
# 'data/index.pro' so that e.g. self.ux is the x-velocity
for key in index.__dict__.keys():
if key != 'global_gg' and key != 'keys':
value = index.__dict__[key]
setattr(self, key, self.f[value-1, ...])
# Special treatment for vector quantities.
if hasattr(index, 'uu'):
self.uu = self.f[index.ux-1:index.uz, ...]
if hasattr(index, 'aa'):
self.aa = self.f[index.ax-1:index.az, ...]
self.t = t
self.dx = dx
self.dy = dy
self.dz = dz
if param.lshear:
self.deltay = deltay
# Do the rest of magic after the trimall (i.e. no additional curl...).
self.magic = magic
if self.magic is not None:
self.magic_attributes(param)
A=SPN[start:start+length]
start+=length
else:
A=np.zeros((3,nbnd,nbnd),dtype=np.complex)
A[:,indn,indm]=f_spn_in.read_record(dtype=np.complex).reshape(3,nbnd*(nbnd+1)/2,order='F')
A[:,indm,indn]=A[:,indn,indm].conj()
check=np.einsum('ijj->',np.abs(A.imag))
if check> 1e-10:
raise RuntimeError ( "REAL DIAG CHECK FAILED for spn: {0}".format(check) )
A=A[:,:,BANDSORT[ik]][:,BANDSORT[ik],:][:,indnQP,indmQP].reshape((3*NBND*(NBND+1)/2),order='F')
if formatted:
f_spn_out.write("".join("{0:26:16e} {1:26:16e}\n".format(x.real,x.imag) for x in A))
else:
f_spn_out.write_record(A)
f_spn_in.close()
f_spn_out.close()
print "----------\n SPN OK \n---------\n"
except IOError as err:
print "WARNING: {0}.spn not written : ".format(seednameGW) ,err
if calcUNK:
unkgwdir="UNK_GW"
unkdftdir="UNK_DFT"
files_list=[]
for f_unk_name in glob.glob("UNK*"):
files_list.append(f_unk_name)
try:
os.mkdir(unkgwdir)
os.mkdir(unkdftdir)
def read(self, datadir='data', proc=-1, quiet=False,
trim=False):
"""
Read the grid data from the pencil code simulation.
If proc < 0, then load all data and assemble.
Otherwise, load grid from specified processor.
call signature:
read(self, file_name='time_series.dat', datadir='data',
double=0, quiet=0, comment_char='#')
Keyword arguments:
*datadir*:
Directory where the data is stored.
*proc*
Processor to be read. If proc is -1, then read the 'global'
grid. If proc is >=0, then read the grid.dat in the
corresponding processor directory.
*quiet*
Flag for switching of output.
*trim*
Cuts off the ghost points.
"""
import numpy as np
import os
from scipy.io import FortranFile
import pencilnew.read as read
datadir = os.path.expanduser(datadir)
dim = read.dim(datadir, proc)
if dim.precision == 'D':
precision = 'd'
else:
precision = 'f'
if proc < 0:
proc_dirs = list(filter(lambda string: string.startswith('proc'), os.listdir(datadir)))
else:
proc_dirs = ['proc' + str(proc)]
# Define the global arrays.
x = np.zeros(dim.mx, dtype=precision)
y = np.zeros(dim.my, dtype=precision)
z = np.zeros(dim.mz, dtype=precision)
dx_1 = np.zeros(dim.mx, dtype=precision)
dy_1 = np.zeros(dim.my, dtype=precision)
dz_1 = np.zeros(dim.mz, dtype=precision)
dx_tilde = np.zeros(dim.mx, dtype=precision)
dy_tilde = np.zeros(dim.my, dtype=precision)
dz_tilde = np.zeros(dim.mz, dtype=precision)
for directory in proc_dirs:
proc = int(directory[4:])
procdim = read.dim(datadir, proc)
if not quiet:
print("reading grid data from processor {0} of {1} ...".format(proc, len(proc_dirs)))
mxloc = procdim.mx
myloc = procdim.my
mzloc = procdim.mz
# Read the grid data.
file_name = os.path.join(datadir, directory, 'grid.dat')
infile = FortranFile(file_name, 'r')
grid_raw = infile.read_record(dtype=precision)
dx, dy, dz = tuple(infile.read_record(dtype=precision))
Lx, Ly, Lz = tuple(infile.read_record(dtype=precision))
dx_1_raw = infile.read_record(dtype=precision)
dx_tilde_raw = infile.read_record(dtype=precision)
infile.close()
# Reshape the arrays.
t = grid_raw[0]
x_loc = grid_raw[1:mxloc+1]
y_loc = grid_raw[mxloc+1:mxloc+myloc+1]
z_loc = grid_raw[mxloc+myloc+1:mxloc+myloc+mzloc+1]
dx_1_loc = dx_1_raw[0:mxloc]
dy_1_loc = dx_1_raw[mxloc:mxloc+myloc]
dz_1_loc = dx_1_raw[mxloc+myloc:mxloc+myloc+mzloc]
dx_tilde_loc = dx_tilde_raw[0:mxloc]
dy_tilde_loc = dx_tilde_raw[mxloc:mxloc+myloc]
dz_tilde_loc = dx_tilde_raw[mxloc+myloc:mxloc+myloc+mzloc]
if len(proc_dirs) > 1:
if procdim.ipx == 0:
i0x = 0
i1x = i0x + procdim.mx
i0x_loc = 0
i1x_loc = procdim.mx
else:
i0x = procdim.ipx*procdim.nx + procdim.nghostx
i1x = i0x + procdim.mx - procdim.nghostx
i0x_loc = procdim.nghostx
i1x_loc = procdim.mx
if procdim.ipy == 0:
i0y = 0
i1y = i0y + procdim.my
i0y_loc = 0
i1y_loc = procdim.my
else:
i0y = procdim.ipy*procdim.ny + procdim.nghosty
i1y = i0y + procdim.my - procdim.nghosty
i0y_loc = procdim.nghosty
i1y_loc = procdim.my
if procdim.ipz == 0:
i0z = 0
i1z = i0z + procdim.mz
i0z_loc = 0
i1z_loc = procdim.mz
else:
i0z = procdim.ipz*procdim.nz + procdim.nghostz
i1z = i0z + procdim.mz - procdim.nghostz
i0z_loc = procdim.nghostz
i1z_loc = procdim.mz
x[i0x:i1x] = x_loc[i0x_loc:i1x_loc]
y[i0y:i1y] = y_loc[i0y_loc:i1y_loc]
z[i0z:i1z] = z_loc[i0z_loc:i1z_loc]
dx_1[i0x:i1x] = dx_1_loc[i0x_loc:i1x_loc]
dy_1[i0y:i1y] = dy_1_loc[i0y_loc:i1y_loc]
dz_1[i0z:i1z] = dz_1_loc[i0z_loc:i1z_loc]
dx_tilde[i0x:i1x] = dx_tilde_loc[i0x_loc:i1x_loc]
dy_tilde[i0y:i1y] = dy_tilde_loc[i0y_loc:i1y_loc]
dz_tilde[i0z:i1z] = dz_tilde_loc[i0z_loc:i1z_loc]
else:
x = x_loc
y = y_loc
z = z_loc
dx_1 = dx_1_loc
dy_1 = dy_1_loc
dz_1 = dz_1_loc
dx_tilde = dx_tilde_loc
dy_tilde = dy_tilde_loc
dz_tilde = dz_tilde_loc
if trim:
self.x = x[dim.l1:dim.l2+1]
self.y = y[dim.m1:dim.m2+1]
self.z = z[dim.n1:dim.n2+1]
self.dx_1 = dx_1[dim.l1:dim.l2+1]
self.dy_1 = dy_1[dim.m1:dim.m2+1]
self.dz_1 = dz_1[dim.n1:dim.n2+1]
self.dx_tilde = dx_tilde[dim.l1:dim.l2+1]
self.dy_tilde = dy_tilde[dim.m1:dim.m2+1]
self.dz_tilde = dz_tilde[dim.n1:dim.n2+1]
else:
self.x = x
self.y = y
self.z = z
self.dx_1 = dx_1
self.dy_1 = dy_1
self.dz_1 = dz_1
self.dx_tilde = dx_tilde
self.dy_tilde = dy_tilde
self.dz_tilde = dz_tilde
self.t = t
self.dx = dx
self.dy = dy
self.dz = dz
self.Lx = Lx
self.Ly = Ly
self.Lz = Lz
def read():
f=FortranFile(BinaryData,'r')
n_vals=f.read_record('i4,i4,i4,i4,i4,(12,)i4')
d_vals=f.read_reals('f4')
xyz_vals=f.read_record('(500,)f4,(500,)f4,(500,)f4')
a_vals=f.read_reals('i8')
ndr_vals=f.read_ints('i4')
ag_vals=f.read_record('(3,)i8')
f.close()
return n_vals,d_vals,xyz_vals,a_vals,ndr_vals,ag_vals
