[MITgcm-devel] Python script to generate cube-sphere grid

Jean-Michel Campin jmc at ocean.mit.edu
Wed May 1 18:07:19 EDT 2013


Hi,

Yuan (Lian) made a python version of the cubed-sphere grid
generator, attached here.
I think it would be good to check-in this script (geng.py)
somewhere in the CVS repository.
I don't know where would be the best place, may be next to the
matlab version it was derived from, i.e.
 MITgcm_contrib/high_res_cube/matlab-grid-generator/
Or a better suggestion ?

Cheers,
Jean-Michel

----- Forwarded message from Yuan Lian <lian at ashimaresearch.com> -----

Date: Sat, 27 Apr 2013 01:48:07 -0700
From: Yuan Lian <lian at ashimaresearch.com>
To: Jean-Michel Campin <jmc at ocean.mit.edu>,
	Adam Showman <showman at lpl.arizona.edu>,
	Nikole Lewis <nlewis at lpl.arizona.edu>
Subject: Python script to generate cube-sphere grid

Hi all,

I finally finish a version of python script that generates cube-sphere
grid (tested
with both python 2 and 3).

Jean-Michel, Please feel free to add this to MITgcm repository unless
extra validation or documentation is required. I would like to give
all credits
to Chris Hill and Dimitris Menemenlis because I simply translated
their matlab
codes to python.

To generate the cube-sphere grid, you can type
> python geng.py -nx 18 -radius 6370.0e3 -nratio 4 -meth 'conf' -ornt
'c' -wr 1 -prec 'd' -mach 'big'

This generates C18 grid with Earth radius (the default configuration).
All arguments
are optional, so you can generate other grids by changing the
parameters of interest.
For example, the following command generates C128 with Jupiter radius:
> python geng.py -nx 128 -radius 71492.0e3

You can get detailed argument list by typing
> python geng.py -h

All grid data is stored in folder called "cs-..." with time stamp.
"tile*.mitgrid" files are
also generated in the same folder.

I have compared the grids (C18, C32 and C64) generated by python and matlab
and find they are almost identical except some machinery error (on
order of 10^{-14}
for angles to 10^{-12} for areas). The error is likely on python side
though both use
double precision.

Please let me know if you have any questions.

Best,
Yuan



# Translator: Yuan Lian, Ashima Research, Pasadena CA
#
# Version 1.0 (April 26, 2013)
#
# This is a python script to generate the cube-sphere grid. The code is translated from 
# matlab scripts located at http://mitgcm.org/viewvc/MITgcm/MITgcm_contrib/high_res_cube/matlab-grid-generator/
# The relative difference between the grid parameters generated by python and matlab is on order of 10^{-12} (area) 
# to 10^{-14} (angle) 
#
# *** Credit goes to the orginal matlab code authors Chris Hill and Dimitris Menemenlis ***
#
# 
# Usage (example for C18 by default): 
#     >  python geng.py -nx 18 -radius 6370.0e3 -nratio 4 -meth 'conf' -ornt 'c' -wr 1 -prec 'd' -mach 'big'
# The argument list is optional, only change the one needed. For example:
#     >  python geng.py -nx 32 -meth 'conf' 
# produces C32 with  conformal grid. 
#
# Use the following command for more information on optional arguments
#    >  python geng.gy -h
# 
#
# The output grid files locate in "cs-conf..." folders with time stamps. The ones 
# readable by MITgcm are named as "tile001.mitgrid","tile002.mitgrid" ...
#
# 
# The python script retains old matlab function definitions for reference purpose.
#
#function [dxg,dyg,dxf,dyf,dxc,dyc,dxv,dyu,Ec,Eu,Ev,Ez,latC,lonC,latG,lonG,...
#          Q11,Q22,Q12, TUu,TUv,TVu,TVv ] ...
#  = gengrid_fn(nx,nratio,meth,ornt,writedata,planet)

# default setting: 
#   nx=19           corresponding to 5 deg resolution (C18)
#   nratio=4        Ratio of working fine grid to target grid
#                   Other options (2, 6, 8 ...): 
#                   higher numbers give larger corner grid size
#   meth='conf'     conformal grid. 
#                   Other options ('tan','tan2','q=0' etc, see rescale_coordinate.py)
#   ornt='c'        standard CS grid orientation, corners at 38 deg. 
#                   Other option (ornt='v') gives corners at about 20 deg
#   writedata=1     write grid files to disk
#   planet='Earth'  default planet is Earth (defuncted)
#                   use Radius as the definition of planet 
#   Repshere=6370.0e3   Radius of Earth in meters


#///////////////////////////////////////////////////////
#/                                                    //
#/   Part 1. Main grid generation functions  	      //
#/            					      //
#///////////////////////////////////////////////////////

def gengrid_fn(nx,Rsphere,nratio,meth,ornt,writedata,prec,machine):

    import numpy as np
    import datetime
    import os

    # sanity check for nx
    if (nx%2 != 0):
       print('the resolution nx = '+str(nx))
       print('nx needs to be an even number. nx = nx-1 = '+str(nx-1)+' is used')
       nx=nx-1
    """ the following is commented out since "Rsphere" is used instead of "planet"
    if (planet.lower()=='Earth'.lower()):
       Rsphere=6370.0e3
    elif (planet.lower()=='Uranus'.lower()):
       Rsphere=25559.0e3
    elif (planet.lower()=='Neptune'.lower()):
       Rsphere=24746.0e3
    elif (planet.lower()=='Jupiter'.lower()):
       Rsphere=71492.0e3
    elif (planet.lower()=='Saturn'.lower()):
       Rsphere=60268.0e3
    elif (planet.lower()=='Mars'.lower()):
       Rsphere=3389.9e3
    elif (planet.lower()=='Titan'.lower()):
       Rsphere=2575.0e3
    elif (planet.lower()=='Exoplanet'.lower()):
       # HD 209458b
       Rsphere=94370.0e3
    else:
       print('no option')
    """

    # nx actually needs to be an odd number in grid generation. 
    # The sanity check above is to reduce confusion. For example, 
    # C18 needs nx=19 but it looks very odd when running python 
    # script as 
    # python geng.py -nc 19
    nx=nx+1

    # Purser algorithm
    # ngen=2; nrat2=1; nx=2*(1+nrat2*(2^ngen))-1;
    # Number of nodes along edge of fine cube (working grid)
    nxf=nratio*(nx-1)+1

    # Generate gridded face using pure conformal mapping
    # Coordinate -1 < q < 0 spans first quadrant
    # Use linspace instead of arange to eliminate the round of error in the last number   
    # q=np.arange(-1.,1./(nxf-1.),2./(nxf-1.))   
    q=np.linspace(-1., 0., int((nxf-1.)/2)+1)
    q=rescale_coordinate(q,meth)

    # Local fine grid curvilinear coordinate (lx,ly)
    lx,ly=np.meshgrid(q,q) 
    lx=lx.transpose()
    ly=ly.transpose()

    # Calculate simple finite volume info for fine grid (lx,ly);
    dx,dy,E = calc_fvgrid(lx,ly)  

    del(dy)

    # Use reflection symmetry of quadrants to fill face
    dx=np.concatenate((dx,dx[(nxf-1)/2-1::-1,:]),axis=0)
    dx=np.concatenate((dx[:,:-1],dx[:,(nxf+1)/2-1::-1]),axis=1)
    E=np.concatenate((E,E[(nxf-1)/2-1::-1,:]),axis=0)
    E=np.concatenate((E,E[:,(nxf-1)/2-1::-1]),axis=1)

    dxg,dxc,dxf,dxv=reduce_dx(dx,nratio)
    Ec,Ez,Ev=reduce_E(E,nratio)
    del(dx,E) # tidy up

    # Remaining grid information
    dyg=dxg.transpose().copy()
    dyc=dxc.transpose().copy()
    dyf=dxf.transpose().copy()
    dyu=dxv.transpose().copy()
    Eu=Ev.transpose().copy()

    # Calculate geographic coordinates

    nxf=2*lx.shape[0]-1
    LatG=np.zeros((nxf,nxf,6))
    LonG=np.zeros((nxf,nxf,6))
    if (ornt=='c'):
       for n in range (0,6):
           # tile number is 1 to 6, so n+1
           LatG[:,:,n],LonG[:,:,n]=calc_geocoords_centerpole(lx,ly,n+1)
       # this is to be consistent with earlier grids
       LatG=permutetiles(LatG,2) 
       # this is to be consistent with earlier grids
       LonG=permutetiles(LonG,2)  
    elif (ornt=='v'):
       for n in range (0,6): 
           LatG[:,:,n],LonG[:,:,n]=calc_geocoords_cornerpole(lx,ly,n+1)
    else:
       print('ornt = ', + ornt)
       exit('Unknown orientation')

    if nratio!=1:
       # Calculate metric tensor
       Q=q.copy()
       Q=np.concatenate((Q[:-1],-q[::-1]))
       qx=Q[1+int(nratio/2)::nratio]-Q[nratio/2-1:-2:nratio]
       QX,QY=np.meshgrid(qx,qx) 
       QX=QX.transpose()
       QY=QY.transpose() 
       del(qx,Q)
       Xf,Yf,Zf=map_lonlat2xyz(LonG,LatG)
       dXdx=( Xf[1+int(nratio/2)::nratio, int(nratio/2)::nratio,:] \
             -Xf[int(nratio/2)-1:-2:nratio, int(nratio/2)::nratio,:] )/expand(QX)
       dYdx=( Yf[1+int(nratio/2)::nratio ,int(nratio/2)::nratio,:] \
             -Yf[int(nratio/2)-1:-2:nratio, int(nratio/2)::nratio,:] )/expand(QX)
       dZdx=( Zf[1+int(nratio/2)::nratio ,int(nratio/2)::nratio,:] \
             -Zf[int(nratio/2)-1:-2:nratio, int(nratio/2)::nratio,:] )/expand(QX)
       dXdy=( Xf[int(nratio/2)::nratio,1+int(nratio/2)::nratio, :] \
             -Xf[int(nratio/2)::nratio, int(nratio/2)-1:-2:nratio,:] )/expand(QY)
       dYdy=( Yf[int(nratio/2)::nratio,1+int(nratio/2)::nratio, :] \
             -Yf[int(nratio/2)::nratio, int(nratio/2)-1:-2:nratio,:] )/expand(QY)
       dZdy=( Zf[int(nratio/2)::nratio,1+int(nratio/2)::nratio, :] \
             -Zf[int(nratio/2)::nratio, int(nratio/2)-1:-2:nratio,:] )/expand(QY)
       Q11=dXdx*dXdx + dYdx*dYdx + dZdx*dZdx
       Q22=dXdy*dXdy + dYdy*dYdy + dZdy*dZdy
       Q12=dXdx*dXdy + dYdx*dYdy + dZdx*dZdy
       del(Xf, Yf, Zf, QX, QY)
    else:
       Q11=0.
       Q12=0.
       Q22=0.

    # Sub-sample to obtain model coordinates
    latG=LatG[::nratio,::nratio,:].copy()
    lonG=LonG[::nratio,::nratio,:].copy()
    if nratio==1:
       latC=( latG[:-1,:-1,:] + latG[1:,:-1,:] \
             +latG[:-1, 1:,:] + latG[1:,1:, :] )/4.
       lonC=( lonG[:-1,:-1,:] + lonG[1:,:-1,:] \
             +lonG[:-1,1: ,:] + lonG[1:,1:, :] )/4.
    else:
       latC=LatG[int(nratio/2)::nratio,int(nratio/2)::nratio,:]
       lonC=LonG[int(nratio/2)::nratio,int(nratio/2)::nratio,:]
    del(LatG, LonG)  # tidy up

    if nratio!=1:
       # Flow rotation tensor
       Xlon=-np.sin(lonC) 
       Ylon= np.cos(lonC) 
       Zlon= 0.*lonC
       TUu = (dXdx*Xlon + dYdx*Ylon + dZdx*Zlon)
       TVu =-(dXdy*Xlon + dYdy*Ylon + dZdy*Zlon)
       Xlat=-np.sin(latC) * np.cos(lonC) 
       Ylat=-np.sin(latC) * np.sin(lonC) 
       Zlat= np.cos(latC)
       TUv =-(dXdx*Xlat + dYdx*Ylat + dZdx*Zlat)
       TVv = (dXdy*Xlat + dYdy*Ylat + dZdy*Zlat)
       det = np.sqrt(TUu*TVv-TUv*TVu)
       TUu=TUu/det
       TUv=TUv/det
       TVu=TVu/det
       TVv=TVv/det
       del(dXdx, dXdy, dYdx, dYdy, dZdx, dZdy)

    # 3D coordinates for plotting
    XG,YG,ZG=map_lonlat2xyz(lonG,latG)

    if writedata!=0:
       timestamp=datetime.datetime.now().strftime('%b-%d-%I%M%p-%G')
       newdir='cs-'+meth+'-'+ornt+'-'+str(nx-1)+'-'+str(nxf-1)+'-'+timestamp
       os.mkdir(newdir) 

       fid=open(newdir+'/'+'grid.info','w')
       fid.write('nx='+str(nx-1)+'\n')
       fid.write('nratio='+str(nratio)+'\n')
       fid.write('nxf='+str((nx-1)*nratio)+'\n')
       fid.write('mapping='+meth+'\n')
       fid.write('orientation='+ornt+'\n')
       fid.close()

       #Rsphere=71492.E3;
       write_tiles(newdir+'/'+'LONG', 180./np.pi*lonG, prec, machine )
       write_tiles(newdir+'/'+'LATG', 180./np.pi*latG, prec, machine )
       write_tiles(newdir+'/'+'LONC', 180./np.pi*lonC, prec, machine )
       write_tiles(newdir+'/'+'LATC', 180./np.pi*latC, prec, machine )

       #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
       lonG1=lonG[0:nx-1,0:nx-1,:].copy()
       latG1=latG[0:nx-1,0:nx-1,:].copy()

       write_tile(newdir+'/'+'LONG', 180./np.pi*lonG1, prec, machine )
       write_tile(newdir+'/'+'LATG', 180./np.pi*latG1, prec, machine )
       write_tile(newdir+'/'+'LONC', 180./np.pi*lonC,  prec, machine )
       write_tile(newdir+'/'+'LATC', 180./np.pi*latC,  prec, machine )
       #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

       dxg1=dxg[0:nx-1,0:nx-1].copy()
       dyg1=dyg[0:nx-1,0:nx-1].copy()
       dxf1=dxf[0:nx-1,0:nx-1].copy()
       dyf1=dyf[0:nx-1,0:nx-1].copy()
       dxc1=dxc[0:nx-1,0:nx-1].copy()
       dyc1=dyc[0:nx-1,0:nx-1].copy()
       dxv1=dxv[0:nx-1,0:nx-1].copy()
       dyu1=dyu[0:nx-1,0:nx-1].copy()
       Ec1=Ec[0:nx-1,0:nx-1].copy()
       Eu1=Eu[0:nx-1,0:nx-1].copy()
       Ev1=Ev[0:nx-1,0:nx-1].copy()
       Ez1=Ez[0:nx-1,0:nx-1].copy()

       write_tile(newdir+'/'+'DXG', Rsphere*expand(dxg1),  prec, machine )
       write_tile(newdir+'/'+'DYG', Rsphere*expand(dyg1),  prec, machine )
       write_tile(newdir+'/'+'DXF', Rsphere*expand(dxf1),  prec, machine )
       write_tile(newdir+'/'+'DYF', Rsphere*expand(dyf1),  prec, machine )
       write_tile(newdir+'/'+'DXC', Rsphere*expand(dxc1),  prec, machine )
       write_tile(newdir+'/'+'DYC', Rsphere*expand(dyc1),  prec, machine )
       write_tile(newdir+'/'+'DXV', Rsphere*expand(dxv1),  prec, machine )
       write_tile(newdir+'/'+'DYU', Rsphere*expand(dyu1),  prec, machine )
       write_tile(newdir+'/'+'RA',  Rsphere**2*expand(Ec1),prec, machine )
       write_tile(newdir+'/'+'RAW', Rsphere**2*expand(Eu1),prec, machine )
       write_tile(newdir+'/'+'RAS', Rsphere**2*expand(Ev1),prec, machine )
       write_tile(newdir+'/'+'RAZ', Rsphere**2*expand(Ez1),prec, machine )

       # Save stuff for pre-processing/post-processing
       # save([dir 'CUBE_HGRID.mat'],'latG','lonG','latC','lonC','dxg','dyg','dxc','dyc','Ec');
       # save([dir 'CUBE_3DGRID.mat'],'XG','YG','ZG');
       # save([dir 'TUV.mat'],'TUu','TUv','TVu','TVv');

       # generate tile*.mitgrid files 
       convertMITgrid(180./np.pi*lonC,180./np.pi*latC,180./np.pi*lonG1,180./np.pi*latG1, \
                                  Rsphere*expand(dxc1),  Rsphere*expand(dyc1),   \
                                  Rsphere*expand(dxg1),  Rsphere*expand(dyg1),   \
                                  Rsphere*expand(dxf1),  Rsphere*expand(dyf1),   \
                                  Rsphere*expand(dxv1),  Rsphere*expand(dyu1),   \
                                  Rsphere**2*expand(Ec1),Rsphere**2*expand(Eu1), \
                                  Rsphere**2*expand(Ev1),Rsphere**2*expand(Ez1), \
                                  newdir,prec,machine)

    return #dxg,dyg ,dxf,dyf,dxc,dyc,dxv,dyu,Ec,Eu,Ev,Ez,latC,lonC,latG,lonG 
           #,Q11,Q22,Q12, TUu,TUv,TVu,TVv


#///////////////////////////////////////////////////////
#/                                                    //
#/   Part 2. Grid generation utilities	  	      //
#/            					      //
#///////////////////////////////////////////////////////

########################################################## 
#               					 #
#    Section 1. calc_geocoords and axis rotation         #
#							 #
##########################################################
#function [lat,lon] = calc_geocoords_centerpole(lx,ly,tile)
# [lat,lon] = calc_geocoords_centerpole(lx,ly,tile);
#
# Calculates geographic coordinates for faces of cube
# with curvilinear coordinates (lx,ly)
#
# Meant to be used for quarter tile only

def calc_geocoords_centerpole(lx,ly,tile):
   
    import numpy as np

    nx=np.shape(lx)[0]
    nxf=2*nx-1

    # 3D coordinates on unit sphere
    lX,lY,lZ=map_xy2xyz(lx,ly)

    # Polar tile(s)
    lonP,latP=map_xyz2lonlat(lX,lY,lZ)
    # Enforce symmetry
    latP=(latP+latP.transpose())/2.  # Does nothing
    lonP[np.where(lonP>=np.pi)]=lonP[np.where(lonP>=np.pi)]-2.*np.pi
    lonP=(lonP-3./2.*np.pi-lonP.transpose())/2.
    for j in range (0,nx):
        lonP[j,j]=-3./4.*np.pi 
    # Use symmetry to expand to full cube
    latP=np.concatenate((latP[:-1,:],latP[nx-1::-1,:]),axis=0)
    latP=np.concatenate((latP[:,:-1],latP[:,nx-1::-1]),axis=1)
    lonP=np.concatenate((lonP[:-1,:],-np.pi-lonP[nx-1::-1,:]),axis=0)
    lonP=np.concatenate((lonP[:,:-1],-lonP[:,nx-1::-1]),axis=1)

    lX,lY,lZ=rotate_about_xaxis(lX,lY,lZ,np.pi/2.)
    lonE,latE=map_xyz2lonlat(lX,lY,lZ)
    # Enforce symmetry
    lonE[0,:]=-3./4.*np.pi
    latE[:,-1]=0.
    latE[:,0]=-latP[0:nx,0]
    # Use symmetry to expand to full cube
    latE=np.concatenate((latE[:-1,:],latE[nx-1::-1,:]),axis=0)
    latE=np.concatenate((latE[:,:-1],-latE[:,nx-1::-1]),axis=1)
    lonE=np.concatenate((lonE[:-1,:],-np.pi-lonE[nx-1::-1,:]),axis=0)
    lonE=np.concatenate((lonE[:,:-1],lonE[:,nx-1::-1]),axis=1)

    # Convert to geographic coordinates
    if tile==1:
       lat=latE.copy()
       lon=lonE-np.pi/2. 
       lon[np.where(lon<=-np.pi)]=lon[np.where(lon<=-np.pi)]+2*np.pi
    elif tile==2:
       lat=latE.copy()
       lon=lonE.copy()
    elif tile==3:
       lat=latP.copy()
       lon=lonP.copy()
    elif tile==4:
       lat=latE[:,::-1].transpose().copy()
       lon=lonE.transpose().copy()+np.pi/2
    elif tile==5:
       lat=latE[:,::-1].transpose().copy();
       lon=lonE.transpose().copy()+np.pi
    elif tile==6:
       lat=-latP.copy()
       lon=lonP[::-1,::-1].transpose().copy()
    else:
       print('tile = '+str(tile))
       exit('Illegal value for tile, should be within [1,6]')

    return lat,lon


#function [lat,lon] = calc_geocoords_cornerpole(lx,ly,tile)
# [lat,lon] = calc_geocoords_cornerpole(lx,ly,tile);
#
# Calculates geographic coordinates for faces of cube
# with curvilinear coordinates (lx,ly)
#
# Meant to be used for quarter tile only
def calc_geocoords_cornerpole(lx,ly,tile):
 
    import numpy as np

    nx=np.shape(lx)[0]
    nxf=2*nx-1

    # 3D coordinates on unit sphere
    lX,lY,lZ=map_xy2xyz(lx,ly)

    # Use symmetry to expand to full cube
    lX=np.concatenate((lX[:-1,:],-lX[nx-1::-1,:]),axis=0) 
    lX=np.concatenate((lX[:,:-1], lX[:,nx-1::-1]),axis=1)
    lY=np.concatenate((lY[:-1,:], lY[nx-1::-1,:]),axis=0) 
    lY=np.concatenate((lY[:,:-1],-lY[:,nx-1::-1]),axis=1)
    lX=(lX+lY.transpose())/2. 
    lY=lX.transpose().copy()
    lZ=(lZ+lZ.transpose())/2.
    lZ=np.concatenate((lZ[:-1,:], lZ[nx-1::-1,:]),axis=0) 
    lZ=np.concatenate((lZ[:,:-1], lZ[:,nx-1::-1]),axis=1)

    # First tile
    lX,lY,lZ=rotate_about_zaxis(lX,lY,lZ,-np.pi/4.)
    lX,lY,lZ=rotate_about_yaxis(lX,lY,lZ,np.arctan(np.sqrt(2.)))

    # Force symmetry again
    lX=(lX+lX.transpose())/2.
    lY=(lY-lY.transpose())/2.
    lZ=(lZ+lZ.transpose())/2.

    lonP,latP=map_xyz2lonlat(lX,lY,lZ)

    # Enforce symmetry again
    for j in range (0,nx):
        lonP[j,j]=0.

    lonP=(lonP-lonP.transpose())/2.
    latP=(latP+latP.transpose())/2.

    latP= latP[:,::-1]
    lonP=-lonP[:,::-1]

    # Convert to geographic coordinates

    if tile==1:
       lat=latP[::-1,:].copy()
       lon=-lonP[::-1,:]-2./3.*np.pi
    elif tile==2:
       lat=latP.copy()
       lon=lonP.copy()
    elif tile==3:
       lat=latP[:,::-1].copy()
       lon=-lonP[:,::-1]+2./3.*np.pi
    elif tile==4:
       lat=-latP[::-1,:].copy()
       lon=lonP[::-1,:]+1./3.*np.pi
    elif tile==5:
       lat=-latP[::-1,::-1].copy()
       lon=-lonP[::-1,::-1]+np.pi;
    elif tile==6:
       lat=-latP[:,::-1].copy()
       lon= lonP[:,::-1]-1./3.*np.pi
    else:
       print('tile = '+str(tile))
       exit('Illegal value for tile, should be within [0,5]')

    lon[np.where(lon>np.pi)]=lon[np.where(lon>np.pi)]-2*np.pi
    lon[np.where(lon<=-np.pi)]=lon[np.where(lon<=-np.pi)]+2*np.pi

    return lat,lon

#function [X,Y,Z]=rotate_about_xaxis(lX,lY,lZ,angle);
# [X,Y,Z]=rotate_about_xaxis(lX,lY,lZ,angle);
# Rotate about X axis by "angle".

def rotate_about_xaxis(lX,lY,lZ,angle):
    import numpy as np
    s=np.sin(angle)
    c=np.cos(angle)
    if c<1.e-9:
       c=0.
       s=np.sign(s)

    X=lX.copy()
    Y=c*lY-s*lZ
    Z=s*lY+c*lZ

    return X,Y,Z

#function [X,Y,Z]=rotate_about_yaxis(lX,lY,lZ,angle);
# [X,Y,Z]=rotate_about_yaxis(lX,lY,lZ,angle);
# Rotate about Y axis by "angle".

def rotate_about_yaxis(lX,lY,lZ,angle):
    import numpy as np
    s=np.sin(angle)
    c=np.cos(angle)
    if c<1.e-9:
       c=0.
       s=np.sign(s)

    X=c*lX+s*lZ
    Y=lY.copy()
    Z=-s*lX+c*lZ

    return X,Y,Z

#function [X,Y,Z]=rotate_about_zaxis(lX,lY,lZ,angle);
# [X,Y,Z]=rotate_about_zaxis(lX,lY,lZ,angle);
# Rotate about Z axis by "angle".

def rotate_about_zaxis(lX,lY,lZ,angle):
    import numpy as np
    s=np.sin(angle)
    c=np.cos(angle)
    if c<1.e-9:
       c=0.
       s=np.sign(s)

    X=c*lX-s*lY
    Y=s*lX+c*lY
    Z=lZ.copy()

    return X,Y,Z


########################################################## 
#							 #
#    Section 2. calc_fvgrid and utilities		 #
#							 #
##########################################################


#function [dx,dy,E] = calc_fvgrid(lx,ly)
# Calculates finite volume grid info (dx,dy,A) for conformal cubic grid
# with curvilinear coordinates (lx,ly)
#
# Meant to be used for single quadrant of tile but does work for full tile

def calc_fvgrid(lx,ly):

    import numpy as np 

    nxf=lx.shape[0]

    # 3D coordinates on unit sphere
    lX,lY,lZ=map_xy2xyz(lx,ly)

    # Use "vector" data type
    Vertices=coord2vector(lX,lY,lZ)
    dx=angle_between_vectors(Vertices[:-1,:,:],Vertices[1::,:,:])
    dy=angle_between_vectors(Vertices[:,:-1,:],Vertices[:,1::,:])
    E=excess_of_quad( Vertices[:-1,:-1,:],Vertices[1::,:-1,:],   \
                      Vertices[1::,1::,:],Vertices[:-1,1::,:] )

    # Force some symmetry (probably does nothing)
    dx=(dx+dy.transpose())/2.
    dy=dx.transpose().copy()
    E=(E+E.transpose())/2.

    # Use symmetry to fill second octant

    for j in range(1,nxf):
        dx[:-nxf+j,j]=dy[j,:-nxf+j].transpose().copy()
    for j in range (0,nxf-1):
        dy[:-(nxf-1)+j,j]=dx[j,:-(nxf-1)+j].transpose().copy()
    for j in range (1,nxf-1):
        E[:-nxf+j,j]=E[j,:-nxf+j].transpose()

    return dx,dy,E

#function [excess]=excess_of_quad(V1,V2,V3,V4)
# excess=excess_of_quad(V1,V2,V3,V4);
#
# where Vn are vectors

def excess_of_quad(V1,V2,V3,V4):
    import numpy as np
    plane1=plane_normal(V1,V2)
    plane2=plane_normal(V2,V3)
    plane3=plane_normal(V3,V4)
    plane4=plane_normal(V4,V1)

    angle12=np.pi-angle_between_vectors(plane2,plane1)
    angle23=np.pi-angle_between_vectors(plane3,plane2)
    angle34=np.pi-angle_between_vectors(plane4,plane3)
    angle41=np.pi-angle_between_vectors(plane1,plane4)

    excess=angle12+angle23+angle34+angle41-2*np.pi
    return excess


#function [plane]=plane_normal(P1,P2);

def plane_normal(P1,P2):

    import numpy as np
    nx=P1.shape[0]
    ny=P1.shape[1]
    nz=P1.shape[2]
    plane=np.zeros((nx,ny,nz))

    plane[:,:,0]=P1[:,:,1]*P2[:,:,2]-P1[:,:,2]*P2[:,:,1]
    plane[:,:,1]=P1[:,:,2]*P2[:,:,0]-P1[:,:,0]*P2[:,:,2]
    plane[:,:,2]=P1[:,:,0]*P2[:,:,1]-P1[:,:,1]*P2[:,:,0]
    mag=np.sqrt(plane[:,:,0]**2 + plane[:,:,1]**2 + plane[:,:,2]**2)
    plane[:,:,0]=plane[:,:,0]/mag
    plane[:,:,1]=plane[:,:,1]/mag
    plane[:,:,2]=plane[:,:,2]/mag

    return plane

#function [dxg,dxc,dxf,dxv]=reduce_dx(dx,nratio);
# [dxg,dxc,dxf,dxv]=reduce_dx(dx,nratio);
#
# Sum dx on fine grid -> dxg,dxc,dxf,dxv
# nratio is ratio of grid sizes (e.g. nratio ~ dxg/dx)

def reduce_dx(dx,nratio):
    import numpy as np

    if nratio==1:
       dxg=dx.copy()
       dxf=(dx[:,:-1]+dx[:,1:])/2.
       # numpy append: use brackets on the 1-row vector
       dxv=( np.append([dx[-1,:]],dx,axis=0) + np.append(dx,[dx[0,:]],axis=0) )/2.    
       dxc=( np.append([dxf[-1,:]],dxf,axis=0) + np.append(dxf,[dxf[0,:]],axis=0) )/2. 
    
    elif nratio%2 == 0:
       nxf=np.shape(dx)[1] 
       kg=np.arange(0,nxf,nratio) 
       kc=np.arange(int(nratio/2),nxf,nratio) 
       jg=np.arange(0,nxf-1,nratio) 
       jc=np.append(nxf-int(nratio/2)-1,np.arange(int(nratio/2),nxf,nratio))
       dxg=np.take(dx,jg,axis=0).take(kg,axis=1)
       dxf=np.take(dx,jg,axis=0).take(kc,axis=1)
       dxc=np.take(dx,jc,axis=0).take(kc,axis=1)
       dxv=np.take(dx,jc,axis=0).take(kg,axis=1)
       for n in range (1,nratio):
           jg=np.mod(jg+1,nxf-1)
           jc=np.mod(jc+1,nxf-1)
           dxg=dxg+np.take(dx,jg,axis=0).take(kg,axis=1)
           dxf=dxf+np.take(dx,jg,axis=0).take(kc,axis=1)
           dxc=dxc+np.take(dx,jc,axis=0).take(kc,axis=1)
           dxv=dxv+np.take(dx,jc,axis=0).take(kg,axis=1)
           # Force more symmetry
           dxf=(dxf+dxf[::-1,:])/2.
           dxf=(dxf+dxf[:,::-1])/2.
           dxg=(dxg+dxg[::-1,:])/2.
           dxg=(dxg+dxg[:,::-1])/2.
           dxc=(dxc+dxc[::-1,:])/2.
           dxc=(dxc+dxc[:,::-1])/2.
           dxv=(dxv+dxv[::-1,:])/2.
           dxv=(dxv+dxv[:,::-1])/2.

    elif nratio%2 != 0:
       print('nratio = ',+str(nratio))
       exit('nratio must be multiple of 2 to be able to reduce grid')

    return dxg,dxc,dxf,dxv

#function [Ec,Ez,Ev]=reduce_E(E,nratio);
# [Ec,Ez,Ev]=reduce_E(E,nratio);
#
# Sum E on fine grid -> Ec,Ez,Ev
# nratio is ratio of grid sizes (e.g. nratio ~ sqrt(Ec/E))
def reduce_E(E,nratio):
    import numpy as np

    nxf=np.shape(E)[0]+1

    if nratio==1:
       Ec=E.copy()
       Ev=( np.append([E[:,-1]],E,axis=0)+np.append(E,[E[:,0]],axis=0) )/2.
       Ez=( np.append([Ev[-1,:]],Ev,axis=0)+np.append(Ev,[Ev[0,:]]) )/2.

    elif nratio%2==0:
       nx=(nxf-1)/nratio+1
       if nx != np.floor(nx):
          print('nx = '+str(nx))
          exit('nxf/nratio must be an integer')
       Ec=np.zeros((nx-1)**2).reshape(nx-1, nx-1,order='F')
       Ev=np.zeros((nx-1)*nx).reshape(nx-1, nx  ,order='F')
       Ez=np.zeros(nx**2).reshape(nx,nx,order='F')
       kg=np.arange(0,nxf-1,nratio)-1  
       kc=np.append(nxf-int(nratio/2)-1, np.arange(int(nratio/2),nxf,nratio))-1
       for m in range (0,nratio):
           kg=np.mod(kg+1,nxf-1)
           kc=np.mod(kc+1,nxf-1)
           jg=np.arange(0,nxf-1,nratio)-1
           jc=np.append(nxf-int(nratio/2)-1,np.arange(int(nratio/2),nxf,nratio))-1
           for n in range (0,nratio):
               jg=np.mod(jg+1,nxf-1)
               jc=np.mod(jc+1,nxf-1)
               Ec=Ec+np.take(E,jg,axis=0).take(kg,axis=1)
               Ev=Ev+np.take(E,jg,axis=0).take(kc,axis=1) 
               Ez=Ez+np.take(E,jc,axis=0).take(kc,axis=1) 

       # Force more symmetry
       Ec=(Ec+Ec[::-1,:])/2.
       Ec=(Ec+Ec[:,::-1])/2.
       Ez=(Ez+Ez[::-1,:])/2.
       Ez=(Ez+Ez[:,::-1])/2.
       Ev=(Ev+Ev[::-1,:])/2.
       Ev=(Ev+Ev[:,::-1])/2.

       # Adjust corner Ez to a hexagon
       Ez[0,  0]=Ez[0,  0]*(3./4.)
       Ez[0, -1]=Ez[0, -1]*(3./4.)
       Ez[-1, 0]=Ez[-1, 0]*(3./4.)
       Ez[-1,-1]=Ez[-1,-1]*(3./4.)

    elif nratio%2 != 0:
       print('nratio = '+str(nratio))
       exit('nratio must be multiple of 2 to be able to reduce gid');

    return Ec,Ez,Ev


########################################################## 
#							 #
#    Section 3. conf_dx and utilities			 #
#							 #
##########################################################


def conf_d(qq):
# function [dxg]=conf_d(qq)
# qq: input 1D array
# calculate dx along edge of conformal cube (used for re-scaling coordinate)

    import numpy as np

    nx=np.size(qq)
    nxf=2*(nx-1)+1 # 2*nx-1 so nxf is an odd number
    q=np.zeros(nxf)
    q[::2]=np.copy(qq)
    q[1::2]=(q[::2][:-1]+q[1::2])/2.

    #j=np.arange(nx-1) 

    # Generate gridded face using pure conformal mapping
    lx,ly=np.meshgrid(q, q)   # Local grid coordinate
    lx=lx.transpose()
    ly=ly.transpose()
    lX,lY,lZ=map_xy2xyz(lx,ly)  # 3D coordinates on unit sphere

    # Use "vector" data type

    XG=lX[::2,::2].copy()
    YG=lY[::2,::2].copy()
    ZG=lZ[::2,::2].copy()
 
    Vertices=coord2vector(XG,YG,ZG)
    dxg=angle_between_vectors(Vertices[:-1,:-1,:],Vertices[1::,:-1,:])
    return dxg

def coord2vector(x,y,z):
# [vec]=coord2vector(x,y,z);
    import numpy as np
    ndims=np.shape(x)
    nx=ndims[0]
    ny=ndims[1]
    #vec=np.zeros(nx*ny*3).reshape(nx,ny,3)
    # a simpler way
    vec=np.zeros((nx,ny,3))
    vec[:,:,0]=np.copy(x)
    vec[:,:,1]=np.copy(y)
    vec[:,:,2]=np.copy(z)

    return vec

#function [angle]=angle_between_vectors(vec1,vec2)
# [angle]=angle_between_vectors(vec1,vec2);

def angle_between_vectors(vec1,vec2):
    import numpy as np
    if np.ndim(vec1)==3:
       vector_prod=vec1[:,:,0]*vec2[:,:,0] \
                  +vec1[:,:,1]*vec2[:,:,1] \
                  +vec1[:,:,2]*vec2[:,:,2]
       nrm1=vec1[:,:,0]**2+vec1[:,:,1]**2+vec1[:,:,2]**2
       nrm2=vec2[:,:,0]**2+vec2[:,:,1]**2+vec2[:,:,2]**2
    else:
       exit('Number of dimensions are not coded for')

    slope=vector_prod/np.sqrt(nrm1*nrm2)
    # check for out of range slopes due to round of errors
    # arccos produces NaNs for x>1 or x<-1

    # No NaNs with Matlab but instead it produces a small complex number
    # Is it supposed to have a complex number??? No, should be an error in 
    # orginal Matlab code.

    # Ohterwise the angle values are consistent with Matlab output.
    slope[np.where(slope> 1)]= 1.
    slope[np.where(slope<-1)]=-1.
    angle=np.arccos(slope)
    return angle

#def prod(iterable):
#    return reduce(operator.mul, iterable, 1)


########################################################## 
#							 #
#    Section 4. map_xy, map_xyz and utilities		 #
#							 #
##########################################################


#function [X,Y,Z] = cmap_xy2xyz(x,y);

# To reload module in python3
# >>> import imp
# >>> imp.reload(map_xy2xyz)

# Conformal mapping of a cube onto a sphere
# maps (x,y) on the north-pole face of a cube
# to (X,Y,Z) coordinates in physical space
#
# Based on f77 code from Jim Purser & Misha Rancic.
#
# Face is oriented normal to Z-axis with
# X and Y increasing with x and y
#
# valid ranges:  -1 < x < 1   -1 < y < 1

def map_xy2xyz(xi,yi):

    import numpy as np

    if np.size(xi) != np.size(yi):
       # import sys
       # sys.exit('Arguments x,y should be same size')
       raise SystemExit('Arguments x,y should be same size')
    xc=np.abs(xi)
    yc=np.abs(yi)
    # find indices where x<0 and y<0
    kx,mx=np.where(xi<0.)
    ky,my=np.where(yi<0.)
    # find indices where yc > xc
    kxy,mxy=np.where(yc>xc)
    # here it is very tricky, need to read the orginal code
    # in python, x=y sets x as a pointer to y 
    #             x=y.copy() sets x as new variable 
    X=xc.copy()
    Y=yc.copy()
    xc=1.-xc
    yc=1.-yc
    for i, item in enumerate(kxy):
        xc[kxy[i],mxy[i]]=1.-Y[kxy[i],mxy[i]]
        yc[kxy[i],mxy[i]]=1.-X[kxy[i],mxy[i]]
    # z is now a complex number 
    zi=((xc+1j*yc)/2.)**4
    # evaluate Taylor serie of z
    W=WofZ(zi)
    # some constants
    thrd=np.float64(1./3.)
    i3=1j**thrd
    ra=np.sqrt(3.)-1.
    cb=1j-1.
    cc=ra*cb/2.
    # not sure what it is doing here
    W=i3*(W*1j)**thrd
    W=(W-ra)/(cb+cc*W)
    X=np.real(W)
    Y=np.imag(W)
    H=2./(1.+X**2+Y**2)
    X=X*H
    Y=Y*H
    Z=H-1.
    #
    T=X.copy()
    for i, item in enumerate(kxy):
        X[kxy[i],mxy[i]]=Y[kxy[i],mxy[i]]
        Y[kxy[i],mxy[i]]=T[kxy[i],mxy[i]]

    for i, item in enumerate(kx):
        X[kx[i],mx[i]]=-X[kx[i],mx[i]] 
    for i, item in enumerate(ky):
        Y[ky[i],my[i]]=-Y[ky[i],my[i]]
    
    # Fix truncation for x=0 or y=0 (aja)
    X[np.where(xi==0.)]=0.
    Y[np.where(yi==0.)]=0.

    # perform extra check on array dimension
    if (np.size(np.where(xi==0.))/2==np.size(xi)):
       X=X.reshape(np.shape(xi)[0], np.shape(xi)[1], order='F')
    if (np.size(np.where(yi==0.))/2==np.size(yi)):
       Y=Y.reshape(np.shape(yi)[0], np.shape(yi)[1], order='F')

    return X, Y, Z

#function [w] = WofZ(z);
# evaluates the Taylor series W(z)
def WofZ(z):
    A=Acoeffs()
    w=0.
    for j in range(29,-1,-1):
        w=(w+A[j])*z
    #stats(abs(z))
    return w

# function [A] = Acoeffs;
# Taylor coefficients for W(z)
# aja changed  -0.00895884801823 to -0.01895884801823 (coeff #4)
def Acoeffs():
    A=[1.47713057321600, -0.38183513110512, -0.05573055466344, \
      -0.01895884801823, -0.00791314396396, -0.00486626515498, \
      -0.00329250387158, -0.00235482619663, -0.00175869000970, \
      -0.00135682443774, -0.00107458043205, -0.00086946107050, \
      -0.00071604933286, -0.00059869243613, -0.00050696402446, \
      -0.00043418115349, -0.00037537743098, -0.00032745130951, \
      -0.00028769063795, -0.00025464473946, -0.00022659577923, \
      -0.00020297175587, -0.00018247947703, -0.00016510295548, \
      -0.00014967258633, -0.00013660647356, -0.00012466390509, \
      -0.00011468147908, -0.00010518717478, -0.00009749136078 ]
    return A

# function [B] = Bcoeffs
# Taylor coefficients for Z(w)
def Bcoeffs():
    B=[0.67698822171341, 0.11847295533659, 0.05317179075349, \
       0.02965811274764, 0.01912447871071, 0.01342566129383, \
       0.00998873721022, 0.00774869352561, 0.00620347278164, \
       0.00509011141874, 0.00425981415542, 0.00362309163280, \
       0.00312341651697, 0.00272361113245, 0.00239838233411, \
       0.00213002038153, 0.00190581436893, 0.00171644267546, \
       0.00155493871562, 0.00141600812949, 0.00129556691848, \
       0.00119042232809, 0.00109804804853, 0.00101642312253, \
       0.00094391466713, 0.00087919127990, 0.00082115825576, \
       0.00076890854394, 0.00072168520663, 0.00067885239089 ]
    return B



#function [lon,lat]=map_xyz2lonlat(x,y,z)
# Convert 3-D coordinates [x,y,z] to [lon,lat]
#
# Assumes "lat" is positive with "z", equatorial plane
# falls at z=0  and "lon" is measured anti-clockwise (eastward)
# from x-axis (y=0) about z-axis.

def map_xyz2lonlat(x,y,z):
    import numpy as np
    # Latitude
    Req=np.sqrt( x**2 + y**2 )
    Z=z.copy()
    kii,mii=np.where(Req == 0.)
    Z[kii,mii]=Z[kii,mii]*np.inf
    Req[kii,mii]=1.
    lat=np.arctan( Z / Req )

    # Longitude
    X=x.copy()
    Y=y.copy()
    kii,mii=np.where(x == 0.)
    Y[kii,mii]=np.inf
    X[kii,mii]=1.
    lon=np.arctan( Y / X )

    kx,mx=np.where(x<0.)
    for i, item in enumerate(kx):
        if y[kx[i],mx[i]] >= 0.:
           lon[kx[i],mx[i]]=np.pi+lon[kx[i],mx[i]]
    kx,mx=np.where(x<=0.)
    for i, item in enumerate(kx):
        if y[kx[i],mx[i]] < 0.:
           lon[kx[i],mx[i]]=-np.pi+lon[kx[i],mx[i]]

    return lon,lat

#function [X,Y,Z]=map_lonlat2xyz(lon,lat)
# [X,Y,Z]=map_lonlat2xyz(lon,lat);
#
# Convert spherical coordinates [lon,lat] to 3-D coordinates [x,y,z]
#
# Assumes "lat" is positive with "z", equatorial plane
# falls at z=0  and "lon" is measured anti-clockwise (eastward)
# from x-axis (y=0) about z-axis.

def map_lonlat2xyz(lon,lat):
    import numpy as np
    # Latitude
    Req=np.cos(lat)
    Z=np.sin(lat)
    X=Req*np.cos(lon)
    Y=Req*np.sin(lon)
    return X,Y,Z


########################################################## 
#							 #
#    Section 5. rescale_coordinate			 #
#							 #
##########################################################


#function [Q]=rescale_coordinate(q,varargin)
# [Q]=rescale_coordinate(q)
# [Q]=rescale_coordinate(q,'conf') * default
# [Q]=rescale_coordinate(q,'q=1')
# [Q]=rescale_coordinate(q,'q=0')
# [Q]=rescale_coordinate(q,'tan')  * also good
# [Q]=rescale_coordinate(q,'tan2')

# use meth='conf' as default option, if not provided
def rescale_coordinate(q,meth='conf'):

    import numpy as np
    #from scipy import interpolate

    nxf=np.size(q)
    # round is a built-in function, same usage as np.round
    # in Matlab nx=round((nxf-1)/2)+1
    nx=round((nxf-1)/2)

    # Generate gridded face using pure conformal mapping
    # q=-min(q):2/(nxf-1):max(q); %q=sign(q).*(abs(q).^(1.2));

    # Calculate dxg along edge for original (unscaled) conformal cube
    dxg=conf_d(q)

    if (meth=='q=0'):
       D=np.max(dxg,axis=1)
    elif (meth=='q=1/2'):
       D=dxg[:,round(nx/2)].copy()
    elif (meth=='q=78'):
       D=dxg[:,round(nx/8)].copy()
    elif (meth=='q=1'):
       D=np.min(dxg,axis=1)
    elif (meth=='q=i3'):
       D=dxg[:,2].copy()
    else:
       D=dxg[:,0]
 
    s=np.cumsum( np.concatenate(([0.],D )) )
    dS=np.max(s)/(nxf-1)+0.*D
    S=np.cumsum( np.concatenate(([0.],dS)) )

    if (meth=='conf'):
       Q=q.copy()
    elif meth in ['q=0','q=1','q=1/2','q=7/8','q=i3']:
       # f1=interpolate.interp1d(s,q)
       # Q=f1(S)
       Q = np.interp(S, s, q)
    elif (meth=='tan'):
       Q=np.tan(2./3.*q)/np.tan(2./3.*np.max(abs(q)))
    elif (meth=='tan2'):
       Q=np.tan(1./5.*q)/np.tan(1./5.*np.max(abs(q)))
    elif (meth=='new'):
       dq=np.ones(nxf)
       dq[0]=0.
       dq[1]=1.5
       dq=dq/np.sum(dq)*(q[-1:]-q[0])
       Q=q[0]+np.cumsum(dq)
    else:
       print('Unknown method')

    return Q

#function [a]=permutetiles(b,n)
# a=permutetiles(b) shifts the tile data left by n places around the equator
#
# ie. n=1, tile 2->1, 4->2, 5->4, 1->5, the tiles 3 and 6 get rotated 90 degs.
#
# Written by adcroft at .mit.edu, 2001.
# $Header: /u/gcmpack/MITgcm/utils/matlab/cs_grid/permutetiles.m,v 1.1 2005/09/15 20:04:57 jmc Exp $
# $Name:  $

def permutetiles(b,n):
    import numpy as np
    c=b.copy()
    a=np.zeros((np.shape(b)[0],np.shape(b)[1],np.shape(b)[2]))
    for k in range (n):
        a[:,:,0]=c[:,:,1].copy()
        a[:,:,1]=c[::-1,:,3].transpose().copy()
        a[:,:,2]=c[::-1,:,2].transpose().copy()
        a[:,:,3]=c[:,:,4].copy()
        a[:,:,4]=c[:,::-1,0].transpose().copy()
        a[:,:,5]=c[:,::-1,5].transpose().copy()
        c=a.copy()
    a=c.copy()
   
    return a

#function [A] = expand(a)
# [A] = expand(a);
#
# A(:,:,j) = a  for j=1:6

def expand(a):

    import numpy as np
    A=np.zeros((np.shape(a)[0],np.shape(a)[1],6))
    for j in range (6):
        A[:,:,j]=a.copy()

    return A


#///////////////////////////////////////////////////////
#/                                                    //
#/   Part 3.  Data writing functions	  	      //
#/            					      //
#///////////////////////////////////////////////////////

#function [] = write_tiles(fn,A,varargin)
# Write tiled field to 6 files "fn.00?.001.data".
#
# Data is in IEEE-bigendian real*8 format by default.
#
# Dimensions of A must be (n,n,6)
#
# Usage:
# >>  write_tiles('LATC',latC);
# >>  write_tiles('LATC',latC,'real*4');

def write_tiles(fn,A,prec,machine):
    import array
    import numpy as np
    from sys import byteorder
    
    for n in range (np.shape(A)[2]):
        fid=open(fn+'.'+"%03d"%(n+1)+'.bin','wb')
        for i in range (np.shape(A)[1]):
            float_array=array.array(prec,A[:,i,n])
            if byteorder!=machine: 
               float_array.byteswap()
            float_array.tofile(fid)
        fid.close()

    return

#function [] = write_tile(fn,A,varargin)
# Write tiled field to single file "fn.bin".
#
# Data is in IEEE-bigendian real*8 format by default.
#
# Dimensions of A must be (n,n)
#
# Usage:
# >>  write_tiles('LATC',latC);
# >>  write_tiles('LATC',latC,'real*4');

def write_tile(fn,A,prec,machine):
    import array
    import numpy as np
    from sys import byteorder

    fid=open(fn+'.bin','wb')
    
    for n in range (np.shape(A)[2]):
        for i in range (np.shape(A)[1]):
            float_array=array.array(prec,A[:,i,n])
            if byteorder!=machine:
               float_array.byteswap()
            float_array.tofile(fid)

    fid.close()

    return


# this function converts the cube-sphere grid data to readable format by MITgcm

# input: lonC,latC,lonG,latG,expand(dxf1),expand(dyf1),
#        expand(dxv1),expand(dyu1),expand(Ec1),expand(Ez1)
def convertMITgrid(xc,yc,xg,yg,dxc,dyc,dxg,dyg,dxf,dyf,dxv,dyu,rac,raw,ras,raz,newdir,prec,machine):
    
    import numpy as np
    import array
    
    # nx=ny, nt=6
    nx=np.shape(xc)[0]
    ny=np.shape(xc)[1]
    nt=np.shape(xc)[2]

    xg =pad_array(xg);   xg[-1,-1,:] =np.nan
    yg =pad_array(yg);   yg[-1,-1,:] =np.nan
    dxv=pad_array(dxv);  dxv[-1,-1,:]=np.nan
    dyu=pad_array(dyu);  dyu[-1,-1,:]=np.nan
    raz=pad_array(raz);  raz[-1,-1,:]=np.nan

    xg[0,-1,[0,2,4]]=xg[0,0,0]
    xg[-1,0,[1,3,5]]=xg[0,0,3]
    xg[-1,:,[0,2,4]]=xg[0,:,[1,3,5]]
    xg[:,-1,[0,2,4]]=xg[0,::-1,[2,4,0]].transpose()
    xg[-1,:,[1,3,5]]=xg[::-1,0,[3,5,1]].transpose()
    xg[:,-1,[1,3,5]]=xg[:,0,[2,4,0]]

    yg[0,-1,[0,2,4]]=yg[0,0,2]
    yg[-1,0,[1,3,5]]=yg[0,0,5]
    yg[-1,:,[0,2,4]]=yg[0,:,[1,3,5]]
    yg[:,-1,[0,2,4]]=yg[0,::-1,[2,4,0]].transpose()
    yg[-1,:,[1,3,5]]=yg[::-1,0,[3,5,1]].transpose()
    yg[:,-1,[1,3,5]]=yg[:,0,[2,4,0]]

    raz[0,-1,[0,2,4]]=raz[0,0,0]
    raz[-1,0,[1,3,5]]=raz[0,0,3]
    raz[-1,:,[0,2,4]]=raz[0,:,[1,3,5]]
    raz[:,-1,[0,2,4]]=raz[0,::-1,[2,4,0]].transpose()
    raz[-1,:,[1,3,5]]=raz[::-1,0,[3,5,1]].transpose()
    raz[:,-1,[1,3,5]]=raz[:,0,[2,4,0]]

    dxv[0,-1,[0,2,4]]=dxv[0,0,0]
    dxv[-1,0,[1,3,5]]=dxv[0,0,3]

    dyu[0,-1,[0,2,4]]=dxv[0,0,0].copy()
    dyu[-1,0,[1,3,5]]=dxv[0,0,3].copy()

    dxv[-1,:,[0,2,4]]=dxv[0,:,[1,3,5]]

    dxv[:,-1,[0,2,4]]=dyu[0,::-1,[2,4,0]].transpose().copy()
    dxv[-1,:,[1,3,5]]=dyu[::-1,0,[3,5,1]].transpose().copy()

    dxv[:,-1,[1,3,5]]=dxv[:,0,[2,4,0]]
    dyu[-1,:,[0,2,4]]=dyu[0,:,[1,3,5]]

    dyu[:,-1,[0,2,4]]=dxv[0,::-1,[2,4,0]].transpose().copy()
    dyu[-1,:,[1,3,5]]=dxv[::-1,0,[3,5,1]].transpose().copy()

    dyu[:,-1,[1,3,5]]=dyu[:,0,[2,4,0]]

    # Pad arrays to be same size
    xc =pad_array(xc)
    yc =pad_array(yc)
    dxf=pad_array(dxf)
    dyf=pad_array(dyf)
    rac=pad_array(rac)
    dxc=pad_array(dxc);  dxc[-1,:-1,:]=np.nan
    dyc=pad_array(dyc);  dyc[:-1,-1,:]=np.nan
    raw=pad_array(raw);  raw[-1,:-1,:]=np.nan
    ras=pad_array(ras);  ras[:-1,-1,:]=np.nan
    dxg=pad_array(dxg);  dxg[:-1,-1,:]=np.nan
    dyg=pad_array(dyg);  dyg[-1,:-1,:]=np.nan

    dxc[-1,:,[0,2,4]]=dxc[0,:,[1,3,5]]
    dxc[-1,:-1,[1,3,5]]=dyc[-2::-1,0,[3,5,1]].transpose().copy()
    dyc[:,-1,[1,3,5]]=dyc[:,0,[2,4,0]]
    dyc[:-1,-1,[0,2,4]]=dxc[0,-2::-1,[2,4,0]].transpose().copy()

    raw[-1,:,[0,2,4]]=raw[0,:,[1,3,5]]
    raw[-1,:-1,[1,3,5]]=ras[-2::-1,0,[3,5,1]].transpose().copy()
    ras[:,-1,[1,3,5]]=ras[:,0,[2,4,0]]
    ras[:-1,-1,[0,2,4]]=raw[0,-2::-1,[2,4,0]].transpose().copy()

    dyg[-1,:,[0,2,4]]=dyg[0,:,[1,3,5]]
    dyg[-1,:-1,[1,3,5]]=dxg[-2::-1,0,[3,5,1]].transpose().copy()
    dxg[:,-1,[1,3,5]]=dxg[:,0,[2,4,0]]
    dxg[:-1,-1,[0,2,4]]=dyg[0,-2::-1,[2,4,0]].transpose().copy()




    for n in range (6):
       fid=open(newdir+'/'+'tile'+"%03d"%(n+1)+'.mitgrid', 'wb')
       write_blocks(fid, xc[:,:,n],prec,machine)
       write_blocks(fid, yc[:,:,n],prec,machine)
       write_blocks(fid,dxf[:,:,n],prec,machine)
       write_blocks(fid,dyf[:,:,n],prec,machine)
       write_blocks(fid,rac[:,:,n],prec,machine)
       write_blocks(fid, xg[:,:,n],prec,machine)
       write_blocks(fid, yg[:,:,n],prec,machine)
       write_blocks(fid,dxv[:,:,n],prec,machine)
       write_blocks(fid,dyu[:,:,n],prec,machine)
       write_blocks(fid,raz[:,:,n],prec,machine)
       write_blocks(fid,dxc[:,:,n],prec,machine)
       write_blocks(fid,dyc[:,:,n],prec,machine)
       write_blocks(fid,raw[:,:,n],prec,machine)
       write_blocks(fid,ras[:,:,n],prec,machine)
       write_blocks(fid,dxg[:,:,n],prec,machine)
       write_blocks(fid,dyg[:,:,n],prec,machine)
       fid.close()
    
    return
   
def write_blocks(fid,A,prec,machine):
    import array
    import numpy as np
    from sys import byteorder
    
    for i in range (np.shape(A)[1]):
        float_array=array.array(prec,A[:,i])
        if byteorder!=machine:
           float_array.byteswap()
        float_array.tofile(fid)

    return 


def pad_array(A):
    import numpy as np
    
    nx=np.shape(A)[0]
    ny=np.shape(A)[1]
    nt=np.shape(A)[2]

    if nx!=ny or nt!=6:
       print('Wrong size for tiles nx != ny; or number of tiles != 6')
       exit('check array size in the source code')

    A_1 =np.zeros((nx+1,ny+1,nt))
    A_1[:-1,:-1,:]=A
  
    return A_1



#///////////////////////////////////////////////////////
#/                                                    //
#/   Part 4. Run geng.py as main	 	      //
#/           Setup default argument list	      //
#/            					      //
#///////////////////////////////////////////////////////


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    # use nargs='?' for optional variables
    # This returns a single variable, not a list obtained fomr e.g., nargs=1
    # Using nargs=1 then n, m are lists, so args.n[0] is needed
    parser.add_argument('-nx','--nx', nargs='?',type=int, default=18, \
                        help='Cube-sphere grid resolution. Default is C18')
    parser.add_argument('-radius','--radius', nargs='?',type=float, default=6370.0e3, \
                        help='Planet Radius. Default is 6370.0e3 [meters]')
    parser.add_argument('-nratio','--nratio', nargs='?',type=int, default=4, \
                        help='high-resolution to low-resolution grid ratio. nratio=2,4 (default),6 ...')
    parser.add_argument('-meth','--meth', nargs='?',type=str, default='conf', \
                        help="Grid projection method. meth='conf'(default), 'tan', 'q=0' etc")
    parser.add_argument('-ornt','--ornt', nargs='?',type=str, default='c', \
                        help="Grid orientation.  ornt='c' (default) or 'v' ")
    parser.add_argument('-wr','--wr', nargs='?',type=int, default=1, \
                        help='write grid data. wr=1 (default is true) or other numbers')
    parser.add_argument('-prec','--prec', nargs='?',type=str, default='d', \
                        help="grid data precision. prec='d' (default is double precision) or 'f' ")
    parser.add_argument('-mach','--mach', nargs='?',type=str, default='big', \
                        help="Machine endianess. mach='big' (default is big endian) or 'little' ")
    args = parser.parse_args()
    print(args.nx,args.radius,args.nratio,args.meth,args.ornt,args.wr,args.prec,args.mach)
    gengrid_fn(args.nx,args.radius,args.nratio,args.meth,args.ornt,args.wr,args.prec,args.mach)


----- End forwarded message -----
-------------- next part --------------
A non-text attachment was scrubbed...
Name: geng.py
Type: text/x-python
Size: 43247 bytes
Desc: not available
URL: <http://mitgcm.org/pipermail/mitgcm-devel/attachments/20130501/55dcf277/attachment-0001.py>


More information about the MITgcm-devel mailing list