mirror of
https://github.com/JHUAPL/kaiju.git
synced 2026-01-09 15:17:56 -05:00
313 lines
9.2 KiB
Python
Executable File
313 lines
9.2 KiB
Python
Executable File
#!/usr/bin/env python
|
|
#Takes one MPI-decomposed restart and spits out an upscaled MPI restart
|
|
|
|
import argparse
|
|
import kaipy.gamera.magsphereRescale as upscl
|
|
from argparse import RawTextHelpFormatter
|
|
import numpy as np
|
|
import h5py
|
|
import os
|
|
import kaipy.kaiH5 as kh5
|
|
|
|
if __name__ == "__main__":
|
|
dIn = os.getcwd()
|
|
|
|
#Input tiling
|
|
iRi = 4
|
|
iRj = 4
|
|
iRk = 1
|
|
#Output tiling
|
|
oRi = 8
|
|
oRj = 8
|
|
oRk = 1
|
|
|
|
inid = "msphere"
|
|
outid = "msphere"
|
|
|
|
nRes = "0"
|
|
grid = "lfmQ.h5"
|
|
|
|
MainS = """Upscales and retiles a Gamera MPI resart
|
|
|
|
|
|
(iRi,iRj,iRk) : Input MPI decomposition
|
|
(oRi,oRj,oRk) : Output MPI decomposition
|
|
inid/nres : Run ID string and restart number, i.e. input file = inid.MPISTUFF.Res.#nres.h5
|
|
outid : Output Run ID
|
|
grid : Filename of input grid corners file (with ghosts) generated by genLFM/genGrid
|
|
"""
|
|
parser = argparse.ArgumentParser(description=MainS, formatter_class=RawTextHelpFormatter)
|
|
parser.add_argument('-i',metavar='inid',default=inid,help="Input Run ID string (default: %(default)s)")
|
|
parser.add_argument('-n',type=int,metavar="nres",default=0,help="Restart number (default: %(default)s)")
|
|
parser.add_argument('-iRi',type=int,metavar="iRi",default=iRi,help="Input i-Ranks (default: %(default)s)")
|
|
parser.add_argument('-iRj',type=int,metavar="iRj",default=iRj,help="Input j-Ranks (default: %(default)s)")
|
|
parser.add_argument('-iRk',type=int,metavar="iRk",default=iRk,help="Input k-Ranks (default: %(default)s)")
|
|
parser.add_argument('-o',type=str,metavar="outid",default=outid,help="Output run ID (default: %(default)s)")
|
|
parser.add_argument('-oRi',type=int,metavar="oRi",default=oRi,help="Input i-Ranks (default: %(default)s)")
|
|
parser.add_argument('-oRj',type=int,metavar="oRj",default=oRj,help="Input j-Ranks (default: %(default)s)")
|
|
parser.add_argument('-oRk',type=int,metavar="oRk",default=oRk,help="Input k-Ranks (default: %(default)s)")
|
|
|
|
parser.add_argument('-grid',type=str,metavar="grid",default=grid,help="inGrid file to read from (default: %(default)s)")
|
|
parser.add_argument('--keep',action='store_true',default=False,help='Keep intermediate files (default: %(default)s)')
|
|
parser.add_argument('--norescale',action='store_true',default=False,help='Do not rescale (up or down) (default: %(default)s)')
|
|
parser.add_argument('--down',action='store_true',default=False,help='Downscale instead of upscale (default: %(default)s)')
|
|
|
|
|
|
#Finalize parsing
|
|
args = parser.parse_args()
|
|
bStr = args.i
|
|
nRes = args.n
|
|
outid = args.o
|
|
iRi = args.iRi
|
|
iRj = args.iRj
|
|
iRk = args.iRk
|
|
|
|
oRi = args.oRi
|
|
oRj = args.oRj
|
|
oRk = args.oRk
|
|
|
|
grid = args.grid
|
|
doKeep = args.keep
|
|
doUp = not args.down
|
|
doRescale = not args.norescale
|
|
|
|
#Pull tiled restart, write to temp file
|
|
#Stupidly writing temp restart to reuse old code
|
|
fTmp = "tempRes.31337.h5" # temporarily written at the run directory.
|
|
oH5 = h5py.File(fTmp,'w') # fTmp needs to include all necessary variables. Use oH5.create_dataset below.
|
|
|
|
G,M,G0,oG,oM = upscl.PullRestartMPI(bStr,nRes,iRi,iRj,iRk,dIn=None,oH5=oH5)
|
|
|
|
#Write main data
|
|
print("Writing plasma and field data to temp file...")
|
|
oH5.create_dataset("Gas",data=G)
|
|
if (G0 is not None):
|
|
doGas0 = True
|
|
print("Writing Gas0")
|
|
oH5.create_dataset("Gas0",data=G0)
|
|
else:
|
|
doGas0 = False
|
|
|
|
oH5.create_dataset("magFlux",data=M)
|
|
oH5.create_dataset("oGas",data=oG)
|
|
oH5.create_dataset("omagFlux",data=oM)
|
|
gVals = ['X','Y','Z']
|
|
|
|
fGrid = grid
|
|
print("Reading grid from %s ..."%(fGrid))
|
|
|
|
iH5 = h5py.File(fGrid,'r')
|
|
for g in gVals:
|
|
oH5.create_dataset(g,data=iH5[g])
|
|
oH5.close()
|
|
|
|
#Upscale from temp file
|
|
# taking an mpi-decomposed restart, lazily stitching it together and writing it to one serial restart
|
|
fTmp2X = "tempRes.31337.2x.h5"
|
|
|
|
#Open input and output
|
|
oH5 = h5py.File(fTmp2X,'w')
|
|
iH5 = h5py.File(fTmp,'r') # iH5 is now the object for fTmp.
|
|
|
|
Ns,Nv,Nk,Nj,Ni = iH5['Gas'].shape
|
|
G = np.zeros((Ns,Nv,Nk,Nj,Ni))
|
|
M = np.zeros((3,Nk+1,Nj+1,Ni+1))
|
|
G[:,:,:,:,:] = iH5['Gas'][:]
|
|
M[ :,:,:,:] = iH5['magFlux'][:]
|
|
oG = np.zeros((Ns,Nv,Nk,Nj,Ni))
|
|
oM = np.zeros((3,Nk+1,Nj+1,Ni+1))
|
|
oG[:,:,:,:,:] = iH5['oGas'][:]
|
|
oM[ :,:,:,:] = iH5['omagFlux'][:]
|
|
if (doGas0):
|
|
G0[:,:,:,:,:] = iH5['Gas0'][:]
|
|
|
|
X = iH5['X'][:]
|
|
Y = iH5['Y'][:]
|
|
Z = iH5['Z'][:]
|
|
|
|
#Transfer attributes to output
|
|
for k in iH5.attrs.keys():
|
|
aStr = str(k)
|
|
oH5.attrs.create(k,iH5.attrs[aStr])
|
|
#Close input
|
|
iH5.close()
|
|
|
|
if (doRescale):
|
|
if (doUp):
|
|
print("Upscaling data ...")
|
|
#Do upscaling
|
|
Xr,Yr,Zr = upscl.upGrid(X,Y,Z)
|
|
Gr = upscl.upGas(X,Y,Z,G,Xr.T,Yr.T,Zr.T)
|
|
FluxR = upscl.upFlux(X,Y,Z,M,Xr,Yr,Zr)
|
|
oGr = upscl.upGas(X,Y,Z,oG,Xr.T,Yr.T,Zr.T)
|
|
oFluxR = upscl.upFlux(X,Y,Z,oM,Xr,Yr,Zr)
|
|
if (doGas0):
|
|
G0r = upscl.upGas(X,Y,Z,G0,Xr.T,Yr.T,Zr.T)
|
|
else:
|
|
print("Downscaling data ...")
|
|
Xr,Yr,Zr = upscl.downGrid(X,Y,Z)
|
|
Gr = upscl.downGas(X,Y,Z,G,Xr.T,Yr.T,Zr.T)
|
|
FluxR = upscl.downFlux(X,Y,Z,M,Xr,Yr,Zr)
|
|
oGr = upscl.downGas(X,Y,Z,oG,Xr.T,Yr.T,Zr.T)
|
|
oFluxR = upscl.downFlux(X,Y,Z,oM,Xr,Yr,Zr)
|
|
if (doGas0):
|
|
G0r = upscl.downGas(X,Y,Z,G0,Xr.T,Yr.T,Zr.T)
|
|
else:
|
|
#No rescale, just set variables
|
|
Xr = X.T #Adding transpose to be consistent w/ rescaling code
|
|
Yr = Y.T
|
|
Zr = Z.T
|
|
Gr = G
|
|
FluxR = M
|
|
oGr = oG
|
|
oFluxR = oM
|
|
if (doGas0):
|
|
G0r = G0
|
|
|
|
#Write out grid to restart
|
|
oH5.create_dataset("X",data=Xr.T)
|
|
oH5.create_dataset("Y",data=Yr.T)
|
|
oH5.create_dataset("Z",data=Zr.T)
|
|
|
|
#Write out gas/flux variables
|
|
oH5.create_dataset("Gas",data=Gr)
|
|
oH5.create_dataset("magFlux",data=FluxR)
|
|
oH5.create_dataset("oGas",data=oGr)
|
|
oH5.create_dataset("omagFlux",data=oFluxR)
|
|
if (doGas0):
|
|
oH5.create_dataset("Gas0",data=G0r)
|
|
|
|
#Close output
|
|
oH5.close()
|
|
|
|
#Split up upscaled file
|
|
if (doGas0):
|
|
upscl.PushRestartMPI(outid,nRes,oRi,oRj,oRk,Xr.T,Yr.T,Zr.T,Gr,FluxR,oGr,oFluxR,fTmp2X,G0r)
|
|
else:
|
|
upscl.PushRestartMPI(outid,nRes,oRi,oRj,oRk,Xr.T,Yr.T,Zr.T,Gr,FluxR,oGr,oFluxR,fTmp2X)
|
|
|
|
#Delete temp files
|
|
if (not doKeep):
|
|
os.remove(fTmp)
|
|
os.remove(fTmp2X)
|
|
|
|
|
|
|
|
|
|
# #!/usr/bin/env python
|
|
# #Takes one MPI-decomposed restart and spits out an upscaled MPI restart
|
|
|
|
# import argparse
|
|
# import kaipy.embiggenUtils as upscl
|
|
# from argparse import RawTextHelpFormatter
|
|
# import numpy as np
|
|
# import h5py
|
|
# import os
|
|
# import kaipy.kaiH5 as kh5
|
|
|
|
# if __name__ == "__main__":
|
|
# #Input tiling
|
|
# iRi = 4
|
|
# iRj = 4
|
|
# iRk = 1
|
|
# #Output tiling
|
|
# oRi = 8
|
|
# oRj = 8
|
|
# oRk = 1
|
|
|
|
# doFast = True
|
|
# doTest = False
|
|
# inid = "msphere"
|
|
# outid = "msphere"
|
|
|
|
# nRes = "0"
|
|
|
|
# MainS = """Upscales and retiles a Gamera MPI resart
|
|
|
|
# (iRi,iRj,iRk) : Input MPI decomposition
|
|
# (oRi,oRj,oRk) : Output MPI decomposition
|
|
# inid/nres : Run ID string and restart number, i.e. input file = inid.MPISTUFF.Res.#nres.h5
|
|
# outid : Output Run ID
|
|
# """
|
|
# parser = argparse.ArgumentParser(description=MainS, formatter_class=RawTextHelpFormatter)
|
|
# parser.add_argument('-i',metavar='inid',default=inid,help="Input Run ID string (default: %(default)s)")
|
|
# parser.add_argument('-n',type=int,metavar="nres",default=0,help="Restart number (default: %(default)s)")
|
|
# parser.add_argument('-iRi',type=int,metavar="iRi",default=iRi,help="Input i-Ranks (default: %(default)s)")
|
|
# parser.add_argument('-iRj',type=int,metavar="iRj",default=iRj,help="Input j-Ranks (default: %(default)s)")
|
|
# parser.add_argument('-iRk',type=int,metavar="iRk",default=iRk,help="Input k-Ranks (default: %(default)s)")
|
|
# parser.add_argument('-o',type=str,metavar="outid",default=outid,help="Output run ID (default: %(default)s)")
|
|
# parser.add_argument('-oRi',type=int,metavar="oRi",default=oRi,help="Input i-Ranks (default: %(default)s)")
|
|
# parser.add_argument('-oRj',type=int,metavar="oRj",default=oRj,help="Input j-Ranks (default: %(default)s)")
|
|
# parser.add_argument('-oRk',type=int,metavar="oRk",default=oRk,help="Input k-Ranks (default: %(default)s)")
|
|
# parser.add_argument('--down',action='store_true',default=False,help='Downscale instead of upscale (default: %(default)s)')
|
|
|
|
# #Finalize parsing
|
|
# args = parser.parse_args()
|
|
# bStr = args.i
|
|
# nRes = args.n
|
|
# outid = args.o
|
|
# iRi = args.iRi
|
|
# iRj = args.iRj
|
|
# iRk = args.iRk
|
|
|
|
# oRi = args.oRi
|
|
# oRj = args.oRj
|
|
# oRk = args.oRk
|
|
|
|
# doUp = not args.down
|
|
|
|
# #Start by pulling tiled restart into one brick w/ halos
|
|
# X,Y,Z,nG,nM,nB,oG,oM,oB,G0,fIn = upscl.PullRestartMPI(bStr,nRes,iRi,iRj,iRk)
|
|
|
|
# if (doTest):
|
|
# rX = X
|
|
# rY = Y
|
|
# rZ = Z
|
|
# nrG = nG
|
|
# nrM = nM
|
|
# nrB = nB
|
|
# orG = oG
|
|
# orM = oM
|
|
# orB = oB
|
|
|
|
# rG0 = G0
|
|
# upscl.PushRestartMPI(outid,nRes,oRi,oRj,oRk,rX,rY,rZ,nrG,nrM,nrB,orG,orM,orB,rG0,fIn,dtScl=1.0)
|
|
# #Toy check
|
|
# upscl.CompRestarts(bStr,outid,nRes,iRi,iRj,iRk)
|
|
|
|
# else:
|
|
|
|
# #Do upscaling on all variables
|
|
# #Chop out last 2 cells on each side, then upscale
|
|
|
|
# #Start w/ grid
|
|
# rX,rY,rZ = upscl.upGrid(X,Y,Z)
|
|
|
|
# dVr = upscl.Volume(rX,rY,rZ)
|
|
# dV0 = upscl.Volume( X, Y, Z)
|
|
|
|
# #Face-centered fluxes
|
|
# nrM = upscl.upFlux(nM)
|
|
|
|
# #Now ready to do cell-centered variables
|
|
# nrG = upscl.upGas(nG,dV0,dVr,"Gas")
|
|
# nrB = upscl.upCCMag(nB,dV0,dVr,"Bxyz")
|
|
# if (G0 is not None):
|
|
# rG0 = upscl.upGas(G0,dV0,dVr,"Gas0")
|
|
|
|
# if (doFast):
|
|
# #Just replicate for oState
|
|
# orM = nrM
|
|
# orB = nrB
|
|
# orG = nrG
|
|
# else:
|
|
# orM = upscl.upFlux (oM)
|
|
# orB = upscl.upCCMag(oB,dV0,dVr,"Bxyz")
|
|
# orG = upscl.upGas (oG,dV0,dVr,"Gas")
|
|
|
|
# #Push back out to arbitrary decomposition
|
|
# #Update dt0 by x1/2
|
|
# upscl.PushRestartMPI(outid,nRes,oRi,oRj,oRk,rX,rY,rZ,nrG,nrM,nrB,orG,orM,orB,rG0,fIn,dtScl=0.5)
|
|
|
|
|