In [1]:
import numpy as np
import pyJHTDB
from pyJHTDB import libJHTDB

In [2]:
import json
with open("parameters-getCutout.txt","r") as pf:
    params=json.load(pf)

auth_token=params["token"]
xstart=int(params.get("xstart"))
ystart=int(params.get("ystart"))
zstart=int(params.get("zstart"))
nx=int(params.get("nx"))
ny=int(params.get("ny"))
nz=int(params.get("nz"))
xstep=int(params.get("xstep",1))
ystep=int(params.get("ystep",1))
zstep=int(params.get("zstep",1))
Filter_Width=int(params.get("Filter_Width",1))
time_step=int(params.get("time_step",0))
field=params.get("field","u")
data_set=params.get("dataset","isotropic1024coarse")

if field == 'u':
    VarName="Velocity"
    dim = 3
elif field == 'a':
    VarName="Vector Potential"
    dim = 3
elif field == 'b':
    VarName="Magnetic Field"
    dim = 3
elif field == 'p':
    VarName="Pressure"
    dim = 1
elif field == 'd':
    VarName="Density"
    dim = 1
elif field == 't':
    VarName="Temperature"
    dim = 1
    
idx_x=np.arange(xstart, xstart+nx, xstep)
idx_y=np.arange(ystart, ystart+ny, ystep)
idx_z=np.arange(zstart, zstart+nz, zstep)
nnx=np.size(idx_x)
nny=np.size(idx_y)
nnz=np.size(idx_z)

npoints=nnx*nny*nnz
split_no=int(np.ceil(npoints/(192000000/dim)))
result=np.zeros((nnz,nny,nnx,dim),dtype='float32')
tmp=np.array_split(np.arange(npoints).reshape(nnx,nny,nnz), split_no)

print(result.shape)


(64, 64, 64, 3)

In [3]:
%%time

lJHTDB = libJHTDB()
lJHTDB.initialize()
lJHTDB.add_token(auth_token)

for t in range(split_no):
    xyzs0 = np.unravel_index(tmp[t][0,0,0], (nnx,nny,nnz))
    xyze0 = np.unravel_index(tmp[t][-1,-1,-1], (nnx,nny,nnz))
    xyzs1 = (idx_x[xyzs0[0]], idx_y[xyzs0[1]], idx_z[xyzs0[2]])
    xyze1 = (idx_x[xyze0[0]], idx_y[xyze0[1]], idx_z[xyze0[2]])
    
    temp = lJHTDB.getCutout(
        data_set=data_set, field=field, time=time_step,
        start=np.array([xyzs1[0], xyzs1[1], xyzs1[2]], dtype = np.int),
        size=np.array([xyze1[0]-xyzs1[0]+1, xyze1[1]-xyzs1[1]+1, xyze1[2]-xyzs1[2]+1], dtype = np.int),
        step=np.array([xstep, ystep, zstep], dtype = np.int),
        filter_width=Filter_Width)
    
    result[xyzs0[2]:xyze0[2]+1, xyzs0[1]:xyze0[1]+1, xyzs0[0]:xyze0[0]+1,:] = temp
    
lJHTDB.finalize()
print(result.shape)


(64, 64, 64, 3)
CPU times: user 64.4 ms, sys: 395 ms, total: 459 ms
Wall time: 5.96 s

In [4]:
%%time
#start = tt.time()

if (dim==3):
    Attribute_Type="Vector"
elif (dim==1):
    Attribute_Type="Scalar"
        
nl = '\r\n'


if data_set in ["channel","channel5200", "transition_bl"]:
        
    if data_set == "channel":
        ygrid = pyJHTDB.dbinfo.channel['ynodes']
        dx=8.0*np.pi/2048
        dz=3.0*np.pi/1536
        x_offset=0
        #xg = np.arange(0,2048.0)
        #zg = np.arange(0,1536.0)                                                            
        #for x in xg:
        #        xg[int(x)] = 8*np.pi/2048*x
        #for z in zg:
        #        zg[int(z)] = 3*np.pi/1536*z
    elif data_set == "channel5200":
        ygrid = pyJHTDB.dbinfo.channel5200['ynodes']
        dx=8.0*np.pi/10240.0
        dz=3.0*np.pi/7680.0
        x_offset=0
        #xg = np.arange(0,10240.0)
        #zg = np.arange(0,7680.0)                                                            
        #for x in xg:
        #        xg[int(x)] = 8*np.pi/10240*x
        #for z in zg:
        #        zg[int(z)] = 3*np.pi/7680*z
    elif data_set == "transition_bl":
        ygrid = pyJHTDB.dbinfo.transition_bl['ynodes']
        dx=0.292210466240511
        dz=0.117244748412311
        x_offset=30.218496172581567
        #xg = np.arange(0,3320.0)
        #zg = np.arange(0,2048.0)
        #for x in xg:
        #        xg[int(x)] = 0.292210466240511*x+30.218496172581567
        #for z in zg:
        #        zg[int(z)] = 0.117244748412311*z
                
    xcoor=idx_x*dx+x_offset
    ycoor=ygrid[idx_y]
    zcoor=idx_z*dz
else:
    if data_set in ["isotropic1024coarse", "isotropic1024fine", "mhd1024", "mixing"]:
        dx=2.0*np.pi/1024.0
    elif data_set in ["isotropic4096", "rotstrat4096"]:
        dx=2.0*np.pi/4096.0
        
    xcoor=idx_x*dx
    ycoor=idx_y*dx
    zcoor=idx_z*dx


CPU times: user 2.2 ms, sys: 110 µs, total: 2.31 ms
Wall time: 224 µs

In [5]:
%%time

import h5py

filename='test'
fh = h5py.File(filename+'.h5', driver='core', block_size=16, backing_store=True)
fh.attrs["dataset"] = np.string_(data_set)
fh.attrs["timeStep"] = time_step
fh.attrs["xstart"] = xstart
fh.attrs["ystart"] = ystart
fh.attrs["zstart"] = zstart
fh.attrs["nx"] = nx
fh.attrs["ny"] = ny
fh.attrs["nz"] = nz
fh.attrs["x_step"] = xstep
fh.attrs["y_step"] = ystep
fh.attrs["z_step"] = zstep
fh.attrs["filterWidth"] = Filter_Width

shape = [0]*3
shape[0] = nnz
shape[1] = nny
shape[2] = nnx

dset = fh.create_dataset(VarName, (shape[0], shape[1],
    shape[2], dim), maxshape=(shape[0], shape[1], shape[2], dim))
dset[...]=result
dset = fh.create_dataset("xcoor", (shape[2],), maxshape=(shape[2],))
dset[...]=xcoor
dset = fh.create_dataset("ycoor", (shape[1],), maxshape=(shape[1],))
dset[...]=ycoor
dset = fh.create_dataset("zcoor", (shape[0],), maxshape=(shape[0],))
dset[...]=zcoor

fh.close()
        
with open(filename+".xmf", "w") as tf:
    print(f"<?xml version=\"1.0\" ?>{nl}", file=tf)
    print(f"<!DOCTYPE Xdmf SYSTEM \"Xdmf.dtd\" []>{nl}", file=tf)
    print(f"<Xdmf Version=\"2.0\">{nl}", file=tf)
    print(f"  <Domain>{nl}", file=tf)

    print(f"    <Grid Name=\"Structured Grid\" GridType=\"Uniform\">{nl}", file=tf)
    print(f"      <Time Value=\"{time_step}\" />{nl}", file=tf)
    print(f"      <Topology TopologyType=\"3DRectMesh\" NumberOfElements=\"{nnz} {nny} {nnx}\"/>{nl}", file=tf)
    print(f"      <Geometry GeometryType=\"VXVYVZ\">{nl}", file=tf)
    print(f"        <DataItem Name=\"Xcoor\" Dimensions=\"{nnx}\" NumberType=\"Float\" Precision=\"4\" Format=\"HDF\">{nl}", file=tf)
    print(f"          {filename}.h5:/xcoor", file=tf)
    print(f"        </DataItem>{nl}", file=tf)
    print(f"        <DataItem Name=\"Ycoor\" Dimensions=\"{nny}\" NumberType=\"Float\" Precision=\"4\" Format=\"HDF\">{nl}", file=tf)
    print(f"          {filename}.h5:/ycoor", file=tf)
    print(f"        </DataItem>{nl}", file=tf)
    print(f"        <DataItem Name=\"Zcoor\" Dimensions=\"{nnz}\" NumberType=\"Float\" Precision=\"4\" Format=\"HDF\">{nl}", file=tf)
    print(f"          {filename}.h5:/zcoor", file=tf)
    print(f"        </DataItem>{nl}", file=tf)
    print(f"      </Geometry>{nl}", file=tf)

    print(f"{nl}", file=tf)

    print(f"      <Attribute Name=\"{VarName}\" AttributeType=\"{Attribute_Type}\" Center=\"Node\">{nl}", file=tf)
    print(f"        <DataItem Dimensions=\"{nnz} {nny} {nnx} {dim}\" NumberType=\"Float\" Precision=\"4\" Format=\"HDF\">{nl}", file=tf)
    print(f"          {filename}.h5:/{VarName}{nl}", file=tf)
    print(f"        </DataItem>{nl}", file=tf)
    print(f"      </Attribute>{nl}", file=tf)

    print(f"{nl}", file=tf)

    print(f"    </Grid>{nl}", file=tf)
    print(f"  </Domain>{nl}", file=tf)
    print(f"</Xdmf>{nl}", file=tf)


CPU times: user 143 ms, sys: 11.2 ms, total: 154 ms
Wall time: 39.9 ms