Changeset 807

Timestamp:

02/08/19 18:53:34 (5 years ago)

Author:

jisesh

Message:

devel/Python : polar projection + Williamson test case 2 on LAM

Location:

codes/icosagcm/devel/Python

Files:

• dynamico/LAM.py (modified) (1 diff)
• dynamico/dev/meshes.py (modified) (1 diff)
• dynamico/maps.py (modified) (2 diffs)
• dynamico/meshes.py (modified) (1 diff)
• test/py/RSW2_MPAS_W02.py (modified) (5 diffs)
• test/py/write_Cartesian_mesh.py (modified) (1 diff)

codes/icosagcm/devel/Python/dynamico/LAM.py

@@ -17 +17 @@
             if ci > c1 : m[i]=0.
         return m
+    def relax_RSW(self, llm, step, flow):
+        beta_i, beta_e = self.beta_i, self.beta_e
+        m,u=flow
+        step.mass[:]        = beta_i*step.mass[:]        + (1.-beta_i)*m[:]
+        step.theta_rhodz[:] = beta_i*step.theta_rhodz[:] + (1.-beta_i)*m[:]
+        step.u[:]           = beta_e*step.u[:]           + (1.-beta_e)*u[:]
+    def relax_hydro(self, llm, step, flow):
+        beta_i, beta_e = self.beta_i, self.beta_e
+        m,S,u=flow
+        for l in range(llm):
+            step.mass[:,l]        = beta_i*step.mass[:,l]        + (1.-beta_i)*m[:,l]
+            step.theta_rhodz[:,l] = beta_i*step.theta_rhodz[:,l] + (1.-beta_i)*S[:,l]
+            step.u[:,l]           = beta_e*step.u[:,l]           + (1.-beta_e)*u[:,l]
     def relax(self, llm, step, flow):
         beta_i, beta_e = self.beta_i, self.beta_e

codes/icosagcm/devel/Python/dynamico/dev/meshes.py

@@ -2 +2 @@
 
 from dynamico import meshes
-from dynamico.meshes import MPAS_Format, Unstructured_PMesh, compute_hybrid_coefs
+from dynamico.meshes import Unstructured_PMesh, compute_hybrid_coefs
+from dynamico.meshes import DYNAMICO_Format, MPAS_Format
 from dynamico import unstructured as unst
 from dynamico.dev import numba

codes/icosagcm/devel/Python/dynamico/maps.py

@@ -14 +14 @@
     def map(self, x, y):
         """x,y are coordinates in the reference domain."""
-        return self.xc + self.scale*x, self.yc + self.scale*y 
+        return self.xc + self.scale*x, self.yc + self.scale*y
     def map_factor(self, lon,lat): return 0.*lon + self.scale
+    def map_angle(self, lon, lat): return 0.*lon
 
 class SphereMap:
@@ -27 +28 @@
     def map(self, lon,lat): return lon, lat
     def map_factor(self, lon,lat): return 0.*lon + self.radius
+    def map_angle(self, lon, lat): return 0.*lon
 
 # https://en.wikipedia.org/wiki/Mercator_projection
+# https://en.wikipedia.org/wiki/Stereographic_projection
+class PolarStereoMap:
+    def __init__(self, radius,Omega, delta, phic=0., lonc=0.):
+        self.a, self.delta, self.Omega = radius, delta, Omega
+        self.phic, self.lonc = phic, lonc # latitude of domain center
+    def coriolis(self, lon, lat):
+        return 2.*self.Omega*np.sin(lat)
+    def map_factor(self,X,Y):
+        lon, lat, factor, angle = self.compute_map_factor_angle(X,Y)
+        return factor
+    def map_angle(self,X,Y):
+        lon, lat, factor, angle = self.compute_map_factor_angle(X,Y)
+        return angle
+    def map(self,X,Y):
+        lon, lat, factor, angle = self.compute_map_factor_angle(X,Y)
+        return lon, lat
+    def compute_map_factor_angle(self,X,Y):
+        """ X, Y : reference coordinates on a mesh where cells have unit size"""
+        alpha = .5*self.delta/self.a
+        X, Y = alpha*X, alpha*Y # scale X and Y to obtain desired resolution delta at domain center
+        # now X and Y are as in wikipedia
+        # we let (x,y) => (-y,x) to align X axis with equator when phic=0.
+        factor =  1./(1.+X*X+Y*Y)
+        y = 2.*X*factor
+        x = -2.*Y*factor
+        z = (1.-X*X-Y*Y)*factor
+        # compute d/dX to track mesh orientation
+        f2 = factor**2
+        dy = 2.*(1.-X*X+Y*Y)*f2
+        dx = 4.*X*Y*f2
+        dz = -4.*X*f2
+        # TODO : rotate x,y,z about polar axis
+        x,y,z = self.RotationPhi(x,y,z)
+        dx,dy,dz = self.RotationPhi(dx,dy,dz)
+        cos_angle = (dy*x-dx*y)/np.sqrt((dx**2+dy**2+dz**2)*(x**2+y**2))
+        sin_angle = (dz*(1.-z**2)-z*(dy*y+dx*x))/np.sqrt((dx**2+dy**2+dz**2)*(x**2+y**2))
+        angle = np.arctan2(dz*(1.-z**2)-z*(dy*y+dx*x), dy*x-dx*y)
+        err = cos_angle-np.cos(angle)
+        print 'Consistency check :', err.min(), err.max()
+        # TODO : rotate x,y,z so that North Pole gets to desired latitude
+        lat = np.arcsin(z)
+        lon = np.arctan2(y,x) + self.lonc
+        return lon, lat, self.delta*factor, angle
+    def RotationPhi(self,x2,y2,z2):
+        cos, sin = np.cos(self.phic), np.sin(self.phic)
+        x3 = sin * x2 + cos*z2
+        z3 = -cos * x2 + sin*z2
+        return x3,y2,z3

codes/icosagcm/devel/Python/dynamico/meshes.py

@@ -431 +431 @@
         """Before calling apply_map, lat and lon are coordinates in the reference domain.
         After calling apply_map, lat and lon are coordinates in the physical domain."""
-        factor_e = mapping.map_factor(self.lon_e, self.lat_e)
-        factor2_i = mapping.map_factor(self.lon_i, self.lat_i)**2
-        factor2_v = mapping.map_factor(self.lon_v, self.lat_v)**2
+        lat_i, lat_e, lat_v = self.lat_i, self.lat_e, self.lat_v
+        lon_i, lon_e, lon_v = self.lon_i, self.lon_e, self.lon_v
+        factor_e = mapping.map_factor(lon_e, lat_e)
+        factor2_i = mapping.map_factor(lon_i, lat_i)**2
+        factor2_v = mapping.map_factor(lon_v, lat_v)**2
         self.le, self.de, self.Av, self.Ai = self.le*factor_e, self.de*factor_e, self.Av*factor2_v, self.Ai*factor2_i
-        self.lon_i, self.lat_i = mapping.map(self.lon_i, self.lat_i)
-        self.lon_v, self.lat_v = mapping.map(self.lon_v, self.lat_v)
-        self.lon_e, self.lat_e = mapping.map(self.lon_e, self.lat_e)
+        self.angle_e += mapping.map_angle(lon_e, lat_e)
+        self.lon_i, self.lat_i = mapping.map(lon_i, lat_i)
+        self.lon_v, self.lat_v = mapping.map(lon_v, lat_v)
+        self.lon_e, self.lat_e = mapping.map(lon_e, lat_e)
+        self.ref_lon_i, self.ref_lat_i = lon_i, lat_i
+        self.ref_lon_e, self.ref_lat_e = lon_e, lat_e
+        self.ref_lon_v, self.ref_lat_v = lon_v, lat_v
 
 class Unstructured_Mesh(Abstract_Mesh):

codes/icosagcm/devel/Python/test/py/RSW2_MPAS_W02.py

@@ -1 +1 @@
 from dynamico import getargs
+getargs.add("--LAM", action='store_true')
+# Args for global mesh
+getargs.add("--grid", type=int, help='Number of hexagons', default=2562)
+# Args for LAM
+getargs.add("--nx", type=int, help='Zonal dimension of LAM mesh', default=100)
+getargs.add("--ny", type=int, help='Meridional dimension of LAM mesh', default=100)
+getargs.add("--dx", type=float, help='Resolution at center of LAM domain', default=1e5)
+getargs.add("--center_lat", type=float, help='Latitude in degrees of LAM center', default=0.)
+getargs.add("--Davies_N1", type=int, default=3)
+getargs.add("--Davies_N2", type=int, default=3)
+
+args = getargs.parse()
+
 log_master, log_world = getargs.getLogger()
 INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error
@@ -14 +27 @@
 from dynamico import unstructured as unst
 from dynamico.dev.meshes import MPAS_Format, Unstructured_PMesh as PMesh, Local_Mesh as Mesh
+from dynamico.dev import meshes
 from dynamico import time_step
 from dynamico import maps
+from dynamico.LAM import Davies
+
 print '...Done'
 
@@ -24 +40 @@
 print '...Done'
 
-grid, llm, nqdyn = 2562, 1,1 # 2562, 10242, 40962
+#--------------------------- functions and classes -----------------------------
+
+class myDavies(Davies):
+    def mask(self,X,Y):
+        # X and Y are coordinates in the reference domain (cell = unit square)
+        # numerical domain extends from -nx/2 ... nx/2 and -ny/2 ... ny/2
+        # useful domain extends from -X0 ... X0 and -Y0 ... Y0
+        N3 = args.Davies_N1+args.Davies_N2
+        X0 = args.nx/2. - N3
+        Y0 = args.ny/2. - N3
+        mask  = self.mask0( X,X0,1.)  # Western boundary
+        mask *= self.mask0(-X,X0,1.) # Eastern boundary
+        mask *= self.mask0( Y,Y0,1.) # Northern boundary
+        mask *= self.mask0(-Y,Y0,1.) # Southern boundary
+        return mask
+
+#----------------------------- main program --------------------------------
+
+llm, nqdyn = 1,1 # 2562, 10242, 40962
 Omega, radius, g, gh0 = 2.*np.pi/86400., 6.4e6, 1., 2.94e4
 N, T, courant = 40, 10400., 1.2 # simulation length = N*T
 
 print 'Omega, planetary PV', Omega, 2*Omega/gh0
-planet = maps.SphereMap(radius, Omega)
 
-print 'Reading MPAS mesh ...'
-meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid)
-pmesh = PMesh(comm,meshfile)
-pmesh.partition_metis()
+if args.LAM:
+    nx, ny = args.nx, args.ny
+    filename = 'cart_%03d_%03d.nc'%(nx,ny)
+    INFO('Reading Cartesian mesh ...')
+    meshfile = meshes.DYNAMICO_Format(filename)
+    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
+    pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
+    planet = maps.PolarStereoMap(radius,Omega, args.dx, args.center_lat*np.pi/180.)
+else:
+    planet = maps.SphereMap(radius, Omega)
+    print 'Reading MPAS mesh ...'
+    meshfile = MPAS_Format('grids/x1.%d.grid.nc'%args.grid)
+    pmesh = PMesh(comm,meshfile)
+    pmesh.partition_metis()
+
 mesh = Mesh(pmesh, llm, nqdyn, planet)
+
 print '...Done'
 lon, lat = mesh.lon_i, mesh.lat_i
@@ -53 +98 @@
 scheme = time_step.RK4(None, dt)
 print dt, scheme.csjl, scheme.cfjl
-step = unst.caldyn_step_TRSW(mesh,scheme,nt)
+step = unst.caldyn_step_TRSW(mesh,scheme,1)
 
 u0 = Omega*radius/12. # cf Williamson (1991), p.13
@@ -59 +104 @@
 print 'Williamson (1991) test 2, u0=', u0
 ulon = u0*np.cos(mesh.lat_e)
-Phi0 = gh0 - gh1*(np.sin(mesh.lat_i)**2)
-zeta0 = (2*u0/radius+2*Omega)*np.sin(mesh.lat_v)
-Phi0v = gh0 - (radius*Omega*u0+.5*u0**2)*(np.sin(mesh.lat_v)**2)
-q0 = zeta0/Phi0v
-fu_perp = mesh.ucov2D(0.*ulon,2*Omega*np.sin(mesh.lat_e)*ulon)
 
-# initial condition
-step.mass[:]=Phi0
-step.theta_rhodz[:]=Phi0
-step.u[:]=mesh.ucov2D(ulon,0.*ulon)
+gh_ini = gh0 - gh1*(np.sin(mesh.lat_i)**2)
+u_ini = mesh.ucov2D(ulon,0.*ulon)
 
-gh,u = step.mass.copy(), step.u.copy()
+# unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, u.shape[0], u) 
 
+if args.LAM:
+    davies = myDavies(args.Davies_N1, args.Davies_N2, 
+                      mesh.ref_lon_i, mesh.ref_lat_i, mesh.ref_lon_e,mesh.ref_lat_e)
+    def relax():
+        davies.relax_RSW(llm, step, (gh_ini, u_ini) )
+else:
+    def relax(): pass
+
+def next_flow(Phi, u):
+    step.mass[:], step.theta_rhodz[:], step.u[:] = Phi, Phi, u
+    for i in range(nt):
+        step.next()
+        relax()
+    return step.mass.copy(), step.u.copy()
+
+gh, u = gh_ini, u_ini
 print gh.shape, gh.min(), gh.max(), u.min(), u.max()
-flow=gh,u
 
-unst.ker.dynamico_update_halo(mesh.com_edges.index, 1, u.shape[0], u) 
-                                  
 for i in range(20):
     if True:
-        step.next()
-        gh,u = step.mass, step.u
+        gh, u = next_flow(gh,u)
+        # step.next()
+        # gh,u = step.mass, step.u
         print i, gh.shape, gh.min(), gh.max(), u.min(), u.max()
-        mesh.plot_i(gh-Phi0) ; plt.title('err(gh)');
+        mesh.plot_i(gh-gh_ini) ; plt.title('err(gh)');
         plt.savefig('%s_err_gh_%02d.png'%(prefix,i))
         plt.close()

codes/icosagcm/devel/Python/test/py/write_Cartesian_mesh.py

@@ -24 +24 @@
 
 dx,dy=Lx/nx,Ly/ny
-mesh = meshes.Cartesian_mesh(nx,ny,llm,nqdyn,Lx,Ly,0.)
+mesh = meshes.Cartesian_Mesh(nx,ny,llm,nqdyn,Lx,Ly,0.)
 print('Successfully initialized Cartesian mesh')
 mesh.ncwrite('cart_%03d_%03d.nc'%(nx,ny))