Files
kaiju/kaipy/gamera/magsphereRescale.py
2024-04-18 18:00:42 +00:00

583 lines
15 KiB
Python

#Various routines to help with restart upscaling
import h5py
import numpy as np
import kaipy.gamera.gamGrids as gg
import scipy
from scipy.spatial import ConvexHull
import kaipy.kaiH5 as kh5
TINY = 1.0e-8
NumG = 4
IDIR = 0
JDIR = 1
KDIR = 2
#Push restart data to an MPI tiling
#fInA is a restart file to pull attributes from
# add oG and oM for one step older for reproducible restart.
# These variables are used by mhd predictor/corrector.
def PushRestartMPI(outid,nRes,Ri,Rj,Rk,X,Y,Z,G,M,oG,oM,fInA,G0=None):
doGas0 = (G0 is not None) # Make sure doGas0 is either True or False
print("Reading attributes from %s"%(fInA))
iH5 = h5py.File(fInA,'r')
Ns,Nv,Nk,Nj,Ni = G.shape
#Create output files
Nkp = Nk//Rk
Njp = Nj//Rj
Nip = Ni//Ri
print("Splitting (%d,%d,%d) cells into (%d,%d,%d) x (%d,%d,%d) [Cells,MPI]"%(Ni,Nj,Nk,Nip,Njp,Nkp,Ri,Rj,Rk))
#Loop over output slices and create restarts
for i in range(Ri):
for j in range(Rj):
for k in range(Rk):
fOut = kh5.genName(outid,i,j,k,Ri,Rj,Rk,nRes)
#Open output file
oH5 = h5py.File(fOut,'w')
iS = i*Nip
iE = iS+Nip
jS = j*Njp
jE = jS+Njp
kS = k*Nkp
kE = kS+Nkp
#Indices for offset ghost grid
iSg = iS -NumG + NumG #Last numG to offset to 0-based index
iEg = iS+Nip+NumG+1 + NumG #Last numG to offset to 0-based index
jSg = jS -NumG + NumG #Last numG to offset to 0-based index
jEg = jS+Njp+NumG+1 + NumG #Last numG to offset to 0-based index
kSg = kS -NumG + NumG #Last numG to offset to 0-based index
kEg = kS+Nkp+NumG+1 + NumG #Last numG to offset to 0-based index
print("Writing %s"%(fOut))
print("\tMPI (%d,%d,%d) = [%d,%d]x[%d,%d]x[%d,%d]"%(i,j,k,iS,iE,jS,jE,kS,kE))
print("\tGrid indices = (%d,%d)x(%d,%d)x(%d,%d)"%(iSg,iEg,jSg,jEg,kSg,kEg))
#Slice subgrids
ijkX = X[kSg:kEg,jSg:jEg,iSg:iEg]
ijkY = Y[kSg:kEg,jSg:jEg,iSg:iEg]
ijkZ = Z[kSg:kEg,jSg:jEg,iSg:iEg]
#Slice pieces out of gas and magflux
ijkG = G [:,:,kS:kE ,jS:jE ,iS:iE ]
ijkM = M [ :,kS:kE+1,jS:jE+1,iS:iE+1]
ijkoG = oG [:,:,kS:kE ,jS:jE ,iS:iE ]
ijkoM = oM [ :,kS:kE+1,jS:jE+1,iS:iE+1]
if (doGas0):
ijkG0 = G0[:,:,kS:kE ,jS:jE ,iS:iE ]
#Write heavy variables
oH5.create_dataset("Gas",data=ijkG)
oH5.create_dataset("magFlux",data=ijkM)
oH5.create_dataset("oGas",data=ijkoG)
oH5.create_dataset("omagFlux",data=ijkoM)
if (doGas0):
oH5.create_dataset("Gas0",data=ijkG0)
#Write subgrid
oH5.create_dataset("X",data=ijkX)
oH5.create_dataset("Y",data=ijkY)
oH5.create_dataset("Z",data=ijkZ)
#Transfer attributes to output
for ak in iH5.attrs.keys():
aStr = str(ak)
oH5.attrs.create(ak,iH5.attrs[aStr])
#Close this output file
oH5.close()
#Close input file
iH5.close()
#Get full data from a tiled restart file
def PullRestartMPI(bStr,nRes,Ri,Rj,Rk,dIn=None,oH5=None):
doInit = True
for i in range(Ri):
for j in range(Rj):
for k in range(Rk):
fID = kh5.genName(bStr,i,j,k,Ri,Rj,Rk,nRes)
print("Reading from %s"%(fID))
#Start with input data
if (dIn is not None):
fIn = dIn + "/" + fID
else:
fIn = fID
iH5 = h5py.File(fIn,'r')
if (doInit):
Ns,Nv,Nkp,Njp,Nip = iH5['Gas'].shape
doGas0 = ('Gas0' in iH5.keys())
Nk = Rk*Nkp
Nj = Rj*Njp
Ni = Ri*Nip
G = np.zeros((Ns,Nv,Nk,Nj,Ni))
oG = np.zeros((Ns,Nv,Nk,Nj,Ni))
if (doGas0):
G0 = np.zeros((Ns,Nv,Nk,Nj,Ni))
else:
G0 = None
M = np.zeros((3,Nk+1,Nj+1,Ni+1))
oM= np.zeros((3,Nk+1,Nj+1,Ni+1))
if (oH5 is not None):
for ka in iH5.attrs.keys():
aStr = str(ka)
#print(aStr)
oH5.attrs.create(ka,iH5.attrs[aStr])
doInit = False
iS = i*Nip
iE = iS+Nip
jS = j*Njp
jE = jS+Njp
kS = k*Nkp
kE = kS+Nkp
#print("MPI (%d,%d,%d) = [%d,%d]x[%d,%d]x[%d,%d]"%(i,j,k,iS,iE,jS,jE,kS,kE))
G[:,:,kS:kE ,jS:jE ,iS:iE ] = iH5['Gas'][:]
oG[:,:,kS:kE ,jS:jE ,iS:iE ] = iH5['oGas'][:]
if (doGas0):
G0[:,:,kS:kE ,jS:jE ,iS:iE ] = iH5['Gas0'][:]
M[ :,kS:kE+1,jS:jE+1,iS:iE+1] = iH5['magFlux'][:]
oM[ :,kS:kE+1,jS:jE+1,iS:iE+1] = iH5['omagFlux'][:]
#Close up
iH5.close()
if (not doGas0):
G0 = None
return G,M,G0,oG,oM
#Downscale a grid (with ghosts, k-j-i order)
def downGrid(X,Y,Z):
Ngk,Ngj,Ngi = X.shape
Nk = Ngk-2*NumG-1
Nj = Ngj-2*NumG-1
Ni = Ngi-2*NumG-1
#Assuming LFM-style grid, get upper half plane
xx = X[NumG,NumG:-NumG,NumG:-NumG].T
yy = Y[NumG,NumG:-NumG,NumG:-NumG].T
#Now half cells in r-phi polar
rr = np.sqrt(xx**2.0+yy**2.0)
pp = np.arctan2(yy,xx)
NiD = Ni//2
NjD = Nj//2
NkD = Nk//2
rrD = np.zeros((NiD+1,NjD+1))
ppD = np.zeros((NiD+1,NjD+1))
for i in range(NiD+1):
for j in range(NjD+1):
rrD[i,j] = rr[2*i,2*j]
ppD[i,j] = pp[2*i,2*j]
#Convert back to 2D Cartesian
xxD = rrD*np.cos(ppD)
yyD = rrD*np.sin(ppD)
#Augment w/ 2D ghosts
xxG,yyG = gg.Aug2D(xxD,yyD,doEps=True,TINY=TINY)
#Augment w/ 3D ghosts
X,Y,Z = gg.Aug3D(xxG,yyG,Nk=NkD,TINY=TINY)
return X,Y,Z
#Upscale a grid (with ghosts, k-j-i order)
def upGrid(X,Y,Z):
Ngk,Ngj,Ngi = X.shape
Nk = Ngk-2*NumG-1
Nj = Ngj-2*NumG-1
Ni = Ngi-2*NumG-1
#Assuming LFM-style grid, get upper half plane
xx = X[NumG,NumG:-NumG,NumG:-NumG].T
yy = Y[NumG,NumG:-NumG,NumG:-NumG].T
#Now half cells in r-phi polar
rr = np.sqrt(xx**2.0+yy**2.0)
pp = np.arctan2(yy,xx)
NiH = 2*Ni
NjH = 2*Nj
NkH = 2*Nk
rrH = np.zeros((NiH+1,NjH+1))
ppH = np.zeros((NiH+1,NjH+1))
#Embed old points into new grid
for i in range(Ni+1):
for j in range(Nj+1):
rrH[2*i,2*j] = rr[i,j]
ppH[2*i,2*j] = pp[i,j]
#Create I midpoints
for i in range(Ni):
rrH[2*i+1,:] = 0.5*( rrH[2*i,:] + rrH[2*i+2,:] )
ppH[2*i+1,:] = 0.5*( ppH[2*i,:] + ppH[2*i+2,:] )
#Create J midpoints
for j in range(Nj):
rrH[:,2*j+1] = 0.5*( rrH[:,2*j] + rrH[:,2*j+2])
ppH[:,2*j+1] = 0.5*( ppH[:,2*j] + ppH[:,2*j+2])
#Create I-J midpoints
for i in range(Ni):
for j in range(Nj):
rrH[2*i+1,2*j+1] = 0.25*( rrH[2*i,2*j] + rrH[2*i,2*j+2] + rrH[2*i+2,2*j] + rrH[2*i+2,2*j+2] )
ppH[2*i+1,2*j+1] = 0.25*( ppH[2*i,2*j] + ppH[2*i,2*j+2] + ppH[2*i+2,2*j] + ppH[2*i+2,2*j+2] )
#Convert back to 2D Cartesian
xxH = rrH*np.cos(ppH)
yyH = rrH*np.sin(ppH)
#Augment w/ 2D ghosts
xxG,yyG = gg.Aug2D(xxH,yyH,doEps=True,TINY=TINY)
#Augment w/ 3D ghosts
X,Y,Z = gg.Aug3D(xxG,yyG,Nk=NkH,TINY=TINY)
return X,Y,Z
#Downscale gBAvg variable (G) on grid X,Y,Z to halved grid.
def downVolt(G):
Nd,Nk,Nj,Ni = G.shape
Gu = np.zeros((Nd,Nk//2,Nj//2,Ni))
print("Downscaling volt variables gBAvg etc...")
#Loop over fine grid
for d in range(Nd):
for i in range(Ni):
for k in range(Nk//2):
for j in range(Nj//2):
j0 = 2*j
k0 = 2*k
Gjk = G[d,k0:k0+2,j0:j0+2,i].mean()
Gu[d,k,j,i] = Gjk
return Gu
#Upscale gBAvg variable (G) on grid X,Y,Z to doubled grid with simple linear interpolation.
#Typical size at quad for example: gBAvg Dataset {3, 128, 96, 1}
def upVolt(G):
try:
Nd,Nk,Nj,Ni = G.shape
Gu = np.zeros((Nd,2*Nk,2*Nj,Ni))
print("Upscaling volt variables gBAvg etc...")
#Loop over coarse grid
for d in range(Nd):
for i in range(Ni):
for k in range(Nk):
for j in range(Nj):
Gjk = G[d,k,j,i]
j0 = 2*j
k0 = 2*k
Gu[d,k0:k0+2,j0:j0+2,i] = Gjk
except:
Nd,Nk,Nj,Ni,Nh = G.shape
Gu = np.zeros((Nd,2*Nk,2*Nj,Ni,Nh))
print("Upscaling volt variables gBHist etc...")
print(G.shape)
#Loop over coarse grid
for h in range(Nh):
for d in range(Nd):
for i in range(Ni):
for k in range(Nk):
for j in range(Nj):
Gjk = G[d,k,j,i,h]
j0 = 2*j
k0 = 2*k
Gu[d,k0:k0+2,j0:j0+2,i,h] = Gjk
return Gu
#Downscale gas variable (G) on grid X,Y,Z (w/ ghosts) to halved grid
def downGas(X,Y,Z,G,Xd,Yd,Zd):
Ns,Nv,Nk,Nj,Ni = G.shape
dV = Volume(X ,Y ,Z )
dVd = Volume(Xd,Yd,Zd)
print("Volume ratio (Coarse/Fine) = %f"%(dVd.sum()/dV.sum()))
Gd = np.zeros((Ns,Nv,Nk//2,Nj//2,Ni//2))
print("Downscaling gas variables ...")
#Loop over coarse grid
for s in range(Ns):
for v in range(Nv):
print("\tDownscaling Species %d, Variable %d"%(s,v))
for k in range(Nk//2):
for j in range(Nj//2):
for i in range(Ni//2):
dVijk = dV[2*k:2*k+2,2*j:2*j+2,2*i:2*i+2] #Volumes of the finer subgrid
dQijk = G[s,v,2*k:2*k+2,2*j:2*j+2,2*i:2*i+2] #stuff density in finer subgrid
#Stuff in the subchunk
QdV = (dVijk*dQijk).sum()
Gd[s,v,k,j,i] = QdV/dVd[k,j,i] #Scale back to density
#Test conservation
print("\t\tCoarse (Total) = %e"%(Gd[s,v,:,:,:]*dVd).sum())
print("\t\tFine (Total) = %e"%(G [s,v,:,:,:]*dV ).sum())
return Gd
#Upscale gas variable (G) on grid X,Y,Z (w/ ghosts) to doubled grid
def upGas(X,Y,Z,G,Xu,Yu,Zu):
Ns,Nv,Nk,Nj,Ni = G.shape
dV = Volume(X ,Y ,Z )
dVu = Volume(Xu,Yu,Zu)
print("Volume ratio (Coarse/Fine) = %f"%(dV.sum()/dVu.sum()))
Gu = np.zeros((Ns,Nv,2*Nk,2*Nj,2*Ni))
print("Upscaling gas variables ...")
#Loop over coarse grid
for s in range(Ns):
for v in range(Nv):
print("\tUpscaling Species %d, Variable %d"%(s,v))
for k in range(Nk):
for j in range(Nj):
for i in range(Ni):
QdV = G[s,v,k,j,i]*dV[k,j,i]
dVijk = dVu[2*k:2*k+2,2*j:2*j+2,2*i:2*i+2] #Volumes of the finer subgrid
vScl = dVijk.sum() #Total volume of the finer subchunks
QdVu = QdV*(dVijk/vScl) #Give weighted contribution to each subcell
Gu[s,v,2*k:2*k+2,2*j:2*j+2,2*i:2*i+2] = QdVu/dVijk #Scale back to density
#Test conservation
print("\t\tCoarse (Total) = %e"%(G [s,v,:,:,:]* dV).sum())
print("\t\tFine (Total) = %e"%(Gu[s,v,:,:,:]*dVu).sum())
return Gu
#Upscale RCMCPL variable
def upRCMCpl(Q,N=1):
Nd = len(Q.shape)
if (Nd>=3):
Nk,Ni,Nj = Q.shape
Qr = np.zeros((Nk,Ni*2,Nj))
for n in range(Nk):
Qr[n,:,:] = upRCM1D(Q[n,:,:])
elif (Nd == 1):
#Just leave it alone
if (len(Q) == N):
#Upscale the one dimension
Q2D = np.zeros((N,1))
Q2D[:,0] = Q
Qr2D = upRCM1D(Q2D)
Qr = Qr2D[:,0]
else:
Qr = Q
else:
#2D
Qr = upRCM1D(Q)
return Qr
#Upscale RCM variable
def upRCM(Q,Ni=1,Nj=1,Nk=1):
Nd = len(Q.shape)
if (Nd>=3):
Nk = Q.shape[0]
Qr = np.zeros((Nk,(Nj-2)*2+2,Ni))
for n in range(Nk):
Qr[n,:,:] = upRCM1Dw(Q[n,:,:])
elif (Nd == 1):
if (len(Q) == Nj):
Q2D = np.zeros((Nj,1))
Q2D[:,0] = Q
Qr2D = upRCM1Dw(Q2D)
Qr = Qr2D[:,0]
else:
Qr = Q
else:
Qr = upRCM1Dw(Q)
return Qr
#Upscale first dimension of 2D array w/ strange rcm-wrap
def upRCM1Dw(Q):
Ni,Nj = Q.shape
Nip = Ni-2
Nri = Nip*2+2
Qr = np.zeros((Nri,Nj))
Qr[0:Nri-2,:] = upRCM1D(Q[0:-2,:])
Qr[Nri-2,:] = Qr[0,:]
Qr[Nri-1,:] = Qr[1,:]
return Qr
#Upscale just first dimension of 2D array
def upRCM1D(Q):
Ni,Nj = Q.shape
Qr = np.zeros((Ni*2,Nj))
for j in range(Nj):
for i in range(Ni):
Qij = Q[i,j]
i0 = 2*i
j0 = j
Qr[i0:i0+2,j0] = Qij
return Qr
#Upscale mix variable
def upMIX(Q):
Ni,Nj = Q.shape
Qr = np.zeros((2*Ni,2*Nj))
for j in range(Nj):
for i in range(Ni):
Qij = Q[i,j]
i0 = 2*i
j0 = 2*j
Qr[i0:i0+2,j0:j0+2] = Qij
#print(Q.sum(),Qr.sum()/4)
return Qr
#Downscale mix variable
def downMIX(Q):
Ni,Nj = Q.shape
Qr = np.zeros((Ni//2,Nj//2))
#Loop over coarse grid
for j in range(Nj//2):
for i in range(Ni//2):
i0 = 2*i
j0 = 2*j
Qij = Q[i0:i0+2,j0:j0+2].mean()
Qr[i,j] = Qij
return Qr
#Return cell centered volume (active only) from grid X,Y,Z (w/ ghosts)
def Volume(Xg,Yg,Zg):
Ngk,Ngj,Ngi = Xg.shape
Nk = Ngk-2*NumG-1
Nj = Ngj-2*NumG-1
Ni = Ngi-2*NumG-1
X = Xg[NumG:-NumG,NumG:-NumG,NumG:-NumG]
Y = Yg[NumG:-NumG,NumG:-NumG,NumG:-NumG]
Z = Zg[NumG:-NumG,NumG:-NumG,NumG:-NumG]
print("Calculating volume of grid of size (%d,%d,%d)"%(Ni,Nj,Nk))
dV = np.zeros((Nk,Nj,Ni))
ijkPts = np.zeros((8,3))
#Assuming LFM-like symmetry
k = 0
for j in range(Nj):
for i in range(Ni):
ijkPts[:,0] = X[k:k+2,j:j+2,i:i+2].flatten()
ijkPts[:,1] = Y[k:k+2,j:j+2,i:i+2].flatten()
ijkPts[:,2] = Z[k:k+2,j:j+2,i:i+2].flatten()
dV[:,j,i] = ConvexHull(ijkPts,incremental=False).volume
return dV
#Downscale magnetic fluxes (M) on grid X,Y,Z (w/ ghosts) to halved grid
def downFlux(X,Y,Z,M,Xu,Yu,Zu):
Nd,Nkc,Njc,Nic = M.shape
Nk = Nkc-1
Nj = Njc-1
Ni = Nic-1
Md = np.zeros((Nd,Nk//2+1,Nj//2+1,Ni//2+1))
#Loop over coarse grid cells
print("Downscaling face fluxes ...")
for k in range(Nk//2):
for j in range(Nj//2):
for i in range(Ni//2):
ip = 2*i
jp = 2*j
kp = 2*k
#West/east i faces
Md[IDIR,k ,j ,i ] = M[IDIR,kp:kp+2,jp:jp+2,ip ].sum()
Md[IDIR,k ,j ,i+1] = M[IDIR,kp:kp+2,jp:jp+2,ip+2].sum()
#South/north j faces
Md[JDIR,k ,j ,i ] = M[JDIR,kp:kp+2,jp ,ip:ip+2].sum()
Md[JDIR,k ,j+1,i ] = M[JDIR,kp:kp+2,jp+2,ip:ip+2].sum()
#Bottom/top k faces
Md[KDIR,k ,j ,i ] = M[KDIR,kp ,jp:jp+2,ip:ip+2].sum()
Md[KDIR,k+1,j ,i ] = M[KDIR,kp+2,jp:jp+2,ip:ip+2].sum()
return Md
#Upscale magnetic fluxes (M) on grid X,Y,Z (w/ ghosts) to doubled grid
def upFlux(X,Y,Z,M,Xu,Yu,Zu):
Nd,Nkc,Njc,Nic = M.shape
Nk = Nkc-1
Nj = Njc-1
Ni = Nic-1
Mu = np.zeros((Nd,2*Nk+1,2*Nj+1,2*Ni+1))
#Loop over coarse grid
print("Upscaling face fluxes ...")
for k in range(Nk):
for j in range(Nj):
for i in range(Ni):
ip = 2*i
jp = 2*j
kp = 2*k
#West i face (4)
Mu[IDIR,kp:kp+2,jp:jp+2,ip] = 0.25*M[IDIR,k,j,i]
#East i face (4)
Mu[IDIR,kp:kp+2,jp:jp+2,ip+2] = 0.25*M[IDIR,k,j,i+1]
#South j face (4)
Mu[JDIR,kp:kp+2,jp,ip:ip+2] = 0.25*M[JDIR,k,j,i]
#North j face (4)
Mu[JDIR,kp:kp+2,jp+2,ip:ip+2] = 0.25*M[JDIR,k,j+1,i]
#Bottom k face (4)
Mu[KDIR,kp,jp:jp+2,ip:ip+2] = 0.25*M[KDIR,k,j,i]
#Top k face (4)
Mu[KDIR,kp+2,jp:jp+2,ip:ip+2] = 0.25*M[KDIR,k+1,j,i]
#Now all exterior faces are done
#12 remaining interior faces
#Interior i faces (4)
Mu[IDIR,kp:kp+2,jp:jp+2,ip+1] = 0.5*( Mu[IDIR,kp:kp+2,jp:jp+2,ip] + Mu[IDIR,kp:kp+2,jp:jp+2,ip+2] )
#Interior j faces (4)
Mu[JDIR,kp:kp+2,jp+1,ip:ip+2] = 0.5*( Mu[JDIR,kp:kp+2,jp,ip:ip+2] + Mu[JDIR,kp:kp+2,jp+2,ip:ip+2] )
#Interior k faces (4)
Mu[KDIR,kp+1,jp:jp+2,ip:ip+2] = 0.5*( Mu[KDIR,kp,jp:jp+2,ip:ip+2] + Mu[KDIR,kp+2,jp:jp+2,ip:ip+2] )
#MaxDiv(M)
return Mu
#Calculates maximum divergence of mag flux data
def MaxDiv(M):
Nd,Nkc,Njc,Nic = M.shape
Nk = Nkc-1
Nj = Njc-1
Ni = Nic-1
Div = np.zeros((Nk,Nj,Ni))
for k in range(Nk):
for j in range(Nj):
for i in range(Ni):
Div[k,j,i] = M[IDIR,k,j,i+1]-M[IDIR,k,j,i]+M[JDIR,k,j+1,i]-M[JDIR,k,j,i]+M[KDIR,k+1,j,i]-M[KDIR,k,j,i]
mDiv = np.abs(Div).max()
bDiv = np.abs(Div).mean()
print("Max/Mean divergence = %e,%e"%(mDiv,bDiv))
print("Sum divergence = %e"%(Div.sum()))
return Div