186 lines
7.3 KiB
Python
Executable File
186 lines
7.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import sys
|
|
import io
|
|
import tarfile
|
|
import argparse
|
|
import numpy as np
|
|
import ucf
|
|
|
|
parser = argparse.ArgumentParser(description='Reads an ucf.tar archive, downsamples it and saves it to a new ucf.tar archive. Can be used as a pipe.')
|
|
parser.add_argument("-i", "--infile", metavar='file',nargs='?', default=None, help="name of the input file [default: stdin]", action="store")
|
|
parser.add_argument("-o", "--outfile", metavar='file',nargs='?', default=None, help="name of the output file [default: stdout]", action="store")
|
|
parser.add_argument("-n", "--nskip", metavar='N',nargs='?', type=int, default=2, help="keep every Nth grid point [default: 2]", action="store")
|
|
parser.add_argument("-sp", "--single-precision", help="output data in single-precision? [default: False]", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
nskip = args.nskip
|
|
file_in = args.infile
|
|
file_out = args.outfile
|
|
saveSinglePrecision = args.single_precision
|
|
|
|
if file_in is None:
|
|
istream = tarfile.open(fileobj=sys.stdin.buffer,mode='r|',bufsize=512*1024**2,ignore_zeros=True)
|
|
else:
|
|
filehandle_in = open(file_in,'rb')
|
|
istream = tarfile.open(fileobj=filehandle_in,mode='r')
|
|
|
|
if file_out is None:
|
|
ostream = tarfile.open(fileobj=sys.stdout.buffer,mode='w|',bufsize=512*1024**2,pax_headers=tarfile.USTAR_FORMAT)
|
|
else:
|
|
filehandle_out = open(file_out,'wb')
|
|
ostream = tarfile.open(fileobj=filehandle_out,mode='w',pax_headers=tarfile.USTAR_FORMAT)
|
|
|
|
while True:
|
|
iinfo = istream.next()
|
|
if iinfo is None:
|
|
break
|
|
print(iinfo.name,file=sys.stderr)
|
|
|
|
ucfbytes_in = istream.extractfile(iinfo).read()
|
|
ucfbytes_out = b''
|
|
|
|
if iinfo.name=='parameters.asc':
|
|
ucfbytes_out += ucfbytes_in
|
|
|
|
if iinfo.name=='particles.bin':
|
|
ucfbytes_out += ucfbytes_in
|
|
|
|
if iinfo.name=='grid.bin':
|
|
ucfhandle = ucf.UCF(file=ucfbytes_in,verbosity=False)
|
|
ucfhandle.copyFileHeaderToBuffer()
|
|
ucfhandle.copyStepToBuffer(1,step_out=1,recursive=False)
|
|
for iset in range(0,ucfhandle.NumDataset):
|
|
(data,params) = ucfhandle.readSet(step=1,dset=iset+1)
|
|
params = list(params)
|
|
nx = params[0]
|
|
ny = params[1]
|
|
nz = params[2]
|
|
x = data[0:nx:nskip]
|
|
y = data[nx:nx+ny:nskip]
|
|
z = data[nx+ny:nx+ny+nz:nskip]
|
|
params[0] = len(x)
|
|
params[1] = len(y)
|
|
params[2] = len(z)
|
|
data = np.concatenate((x,y,z))
|
|
ucfhandle.addDatasetToBuffer(data,params=params,step=1,dset=iset+1)
|
|
ucfbytes_out += ucfhandle.flushBuffer()
|
|
ucfhandle.close()
|
|
|
|
if iinfo.name=='proc.bin':
|
|
ucfhandle = ucf.UCF(file=ucfbytes_in,verbosity=False)
|
|
ucfhandle.copyFileHeaderToBuffer()
|
|
ucfhandle.copyStepToBuffer(1,step_out=1,recursive=False)
|
|
for iset in range(0,ucfhandle.NumDataset):
|
|
(data,params) = ucfhandle.readSet(step=1,dset=iset+1)
|
|
nxp = params[0]
|
|
nyp = params[1]
|
|
nzp = params[2]
|
|
ibeg = np.copy(data[0:nxp] )
|
|
iend = np.copy(data[nxp:2*nxp] )
|
|
jbeg = np.copy(data[2*nxp:2*nxp+nyp] )
|
|
jend = np.copy(data[2*nxp+nyp:2*nxp+2*nyp] )
|
|
kbeg = np.copy(data[2*nxp+2*nyp:2*nxp+2*nyp+nzp] )
|
|
kend = np.copy(data[2*nxp+2*nyp+nzp:2*nxp+2*nyp*2*nzp])
|
|
for ixp in range(0,nxp):
|
|
ibeg[ixp] = (ibeg[ixp]-1)//nskip+1
|
|
iend[ixp] = (iend[ixp]-1)//nskip+1
|
|
for iyp in range(0,nyp):
|
|
jbeg[iyp] = (jbeg[iyp]-1)//nskip+1
|
|
jend[iyp] = (jend[iyp]-1)//nskip+1
|
|
for izp in range(0,nzp):
|
|
kbeg[izp] = (kbeg[izp]-1)//nskip+1
|
|
kend[izp] = (kend[izp]-1)//nskip+1
|
|
data = np.concatenate((ibeg,iend,jbeg,jend,kbeg,kend))
|
|
ucfhandle.addDatasetToBuffer(data,params=params,step=1,dset=iset+1)
|
|
ucfbytes_out += ucfhandle.flushBuffer()
|
|
ucfhandle.close()
|
|
|
|
if 'uvwp.' in iinfo.name:
|
|
ucfhandle = ucf.UCF(file=ucfbytes_in,verbosity=False)
|
|
ucfhandle.copyFileHeaderToBuffer()
|
|
ucfhandle.copyStepToBuffer(1,step_out=1,recursive=False)
|
|
for iset in range(0,4):
|
|
(data,params_in) = ucfhandle.readSet(step=1,dset=iset+1)
|
|
ighost = params_in[0]
|
|
(ibeg,jbeg,kbeg) = params_in[1:4]
|
|
(nxl,nyl,nzl) = params_in[4:7]
|
|
(nxg,nyg,nzg) = params_in[7:10]
|
|
data = data.reshape((nxl+2*ighost,nyl+2*ighost,nzl+2*ighost),order='F')
|
|
|
|
if nskip>1:
|
|
islice = [ii-ibeg+ighost for ii in range(ibeg,ibeg+nxl) if (ii-1)%nskip==0]
|
|
jslice = [ii-jbeg+ighost for ii in range(jbeg,jbeg+nyl) if (ii-1)%nskip==0]
|
|
kslice = [ii-kbeg+ighost for ii in range(kbeg,kbeg+nzl) if (ii-1)%nskip==0]
|
|
|
|
data = data[np.ix_(islice,jslice,kslice)]
|
|
ibeg = (islice[0]+ibeg-ighost)//nskip+1
|
|
jbeg = (jslice[0]+jbeg-ighost)//nskip+1
|
|
kbeg = (kslice[0]+kbeg-ighost)//nskip+1
|
|
(nxl,nyl,nzl) = data.shape
|
|
nxg = (nxg-1)//nskip+1
|
|
nyg = (nyg-1)//nskip+1
|
|
nzg = (nzg-1)//nskip+1
|
|
|
|
params_out = list(params_in)
|
|
if nskip>1:
|
|
params_out[0] = 0
|
|
params_out[1:4] = (ibeg,jbeg,kbeg)
|
|
params_out[4:7] = (nxl,nyl,nzl)
|
|
params_out[7:10] = (nxg,nyg,nzg)
|
|
|
|
if saveSinglePrecision:
|
|
data = data.astype(np.float32,casting='same_kind')
|
|
ucfhandle.addDatasetToBuffer(data,params=params_out,step=1,dset=iset+1)
|
|
ucfbytes_out += ucfhandle.flushBuffer()
|
|
ucfhandle.close()
|
|
|
|
if 'scal.' in iinfo.name:
|
|
ucfhandle = ucf.UCF(file=ucfbytes_in,verbosity=False)
|
|
ucfhandle.copyFileHeaderToBuffer()
|
|
ucfhandle.copyStepToBuffer(1,step_out=1,recursive=False)
|
|
for iset in range(0,ucfhandle.NumDataset):
|
|
(data,params_in) = ucfhandle.readSet(step=1,dset=iset+1)
|
|
ighost = params_in[0]
|
|
(ibeg,jbeg,kbeg) = params_in[1:4]
|
|
(nxl,nyl,nzl) = params_in[4:7]
|
|
(nxg,nyg,nzg) = params_in[7:10]
|
|
data = data.reshape((nxl+2*ighost,nyl+2*ighost,nzl+2*ighost),order='F')
|
|
|
|
if nskip>1:
|
|
islice = [ii-ibeg+ighost for ii in range(ibeg,ibeg+nxl) if (ii-1)%nskip==0]
|
|
jslice = [ii-jbeg+ighost for ii in range(jbeg,jbeg+nyl) if (ii-1)%nskip==0]
|
|
kslice = [ii-kbeg+ighost for ii in range(kbeg,kbeg+nzl) if (ii-1)%nskip==0]
|
|
|
|
data = data[np.ix_(islice,jslice,kslice)]
|
|
ibeg = (islice[0]+ibeg-ighost)//nskip+1
|
|
jbeg = (jslice[0]+jbeg-ighost)//nskip+1
|
|
kbeg = (kslice[0]+kbeg-ighost)//nskip+1
|
|
(nxl,nyl,nzl) = data.shape
|
|
nxg = (nxg-1)//nskip+1
|
|
nyg = (nyg-1)//nskip+1
|
|
nzg = (nzg-1)//nskip+1
|
|
|
|
params_out = list(params_in)
|
|
if nskip>1:
|
|
params_out[0] = 0
|
|
params_out[1:4] = (ibeg,jbeg,kbeg)
|
|
params_out[4:7] = (nxl,nyl,nzl)
|
|
params_out[7:10] = (nxg,nyg,nzg)
|
|
|
|
if saveSinglePrecision:
|
|
data = data.astype(np.float32,casting='same_kind')
|
|
ucfhandle.addDatasetToBuffer(data,params=params_out,step=1,dset=iset+1)
|
|
ucfbytes_out += ucfhandle.flushBuffer()
|
|
ucfhandle.close()
|
|
|
|
oinfo = tarfile.TarInfo(name=iinfo.name)
|
|
oinfo.size = len(ucfbytes_out)
|
|
ostream.addfile(oinfo,fileobj=io.BytesIO(ucfbytes_out))
|
|
|
|
istream.close()
|
|
ostream.close()
|
|
if file_in is not None:
|
|
filehandle_in.close()
|
|
if file_out is not None:
|
|
filehandle_out.close()
|