source: codes/icosagcm/devel/Python/test/py/dump_partition.py @ 960

Last change on this file since 960 was 960, checked in by dubos, 5 years ago

devel/unstructured : fixed generation of local mesh information + restored background wind in DCMIP31
File size: 8.7 KB
Line 
1from dynamico import getargs
2getargs.add("--LAM", action='store_true')
3
4# Args for both cases
5getargs.add("--Omega", type=float, help='Planetary radius', default=7e-5)
6# Args for global mesh
7getargs.add("--grid", type=int, help='Number of hexagons', default=2562)
8getargs.add("--radius", type=float, help='Planetary radius', default=6.4e6)
9# Args for LAM
10getargs.add("--nx", type=int, help='Zonal dimension of LAM mesh', default=100)
11getargs.add("--ny", type=int, help='Meridional dimension of LAM mesh', default=100)
12getargs.add("--dx", type=float, help='Resolution at center of LAM domain', default=1e5)
13getargs.add("--center_lat", type=float, help='Latitude in degrees of LAM center', default=0.)
14getargs.add("--Davies_N1", type=int, default=3)
15getargs.add("--Davies_N2", type=int, default=3)
16
17args = getargs.parse()
18for arg in vars(args): print arg, getattr(args, arg)
19
20log_master, log_world = getargs.getLogger()
21INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error
22
23INFO('Starting')
24
25from mpi4py import MPI
26comm = MPI.COMM_WORLD
27mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size()
28INFO('%d/%d starting'%(mpi_rank,mpi_size))
29prefix='fig_RSW2_MPAS_W02/%02d'%mpi_rank
30
31INFO('Loading DYNAMICO modules ...')
32from dynamico.dev import unstructured as unst
33from dynamico.dev.meshes import MPAS_Format, Unstructured_PMesh as PMesh, Local_Mesh as Mesh
34from dynamico.dev import meshes
35from dynamico import time_step
36from dynamico import maps
37from dynamico.LAM import Davies
38
39print '...Done'
40
41print 'Loading modules ...'
42import math as math
43import matplotlib.pyplot as plt
44import numpy as np
45import netCDF4 as cdf
46
47print '...Done'
48
49#--------------------------- functions and classes -----------------------------
50
51class myDavies(Davies):
52    def mask(self,X,Y):
53        # X and Y are coordinates in the reference domain (cell = unit square)
54        # numerical domain extends from -nx/2 ... nx/2 and -ny/2 ... ny/2
55        # useful domain extends from -X0 ... X0 and -Y0 ... Y0
56        N3 = args.Davies_N1+args.Davies_N2
57        X0 = args.nx/2. - N3
58        Y0 = args.ny/2. - N3
59        mask  = self.mask0( X,X0,1.) # Western boundary
60        mask *= self.mask0(-X,X0,1.) # Eastern boundary
61        mask *= self.mask0( Y,Y0,1.) # Northern boundary
62        mask *= self.mask0(-Y,Y0,1.) # Southern boundary
63        return mask
64
65#----------------------------- main program --------------------------------
66
67llm, nqdyn = 1,1
68Omega, radius = args.Omega, args.radius
69
70print 'Omega :', Omega
71
72if args.LAM:
73    nx, ny = args.nx, args.ny
74    filename = 'cart_%03d_%03d.nc'%(nx,ny)
75    INFO('Reading Cartesian mesh ...')
76    meshfile = meshes.DYNAMICO_Format(filename)
77    pmesh = meshes.Unstructured_PMesh(comm,meshfile)
78    pmesh.partition_curvilinear(args.mpi_ni,args.mpi_nj)
79    planet = maps.PolarStereoMap(radius,Omega, args.dx, args.center_lat*np.pi/180.)
80    mesh = Mesh(pmesh, llm, nqdyn, planet)
81    davies = myDavies(args.Davies_N1, args.Davies_N2, 
82                      mesh.ref_lon_i, mesh.ref_lat_i, mesh.ref_lon_e,mesh.ref_lat_e)
83else:
84    planet = maps.SphereMap(radius, Omega)
85    print 'Reading MPAS mesh ...'
86    meshfile = MPAS_Format('grids/x1.%d.grid.nc'%args.grid)
87    pmesh = PMesh(comm,meshfile)
88    pmesh.partition_metis()
89    mesh = Mesh(pmesh, llm, nqdyn, planet)
90
91print '...Done'
92
93def print_attr(obj):
94    #for name in obj.__dict__.keys(): print name, type(getattr(obj,name))
95    for name in obj.__dict__.keys(): print name, type(getattr(obj,name)), np.shape(getattr(obj,name))#,np.dtype(getattr(obj,name))
96    #for key in obj.__dict__.keys():  print key, np.dtype(obj.__dict__.values())
97
98print_attr(mesh)
99print("--------")
100print_attr(mesh.com_edges)
101print("--------")
102
103class preprocess_output:
104    def __init__(self,mesh):
105        self._mesh = mesh
106    #-------------------START WRITING THE netCDF file---------------------------
107    def create_vars(self,f,dimname, info):
108        for vname, vtype, vdata in info:
109            vdata = np.asarray(vdata)
110            print('creating variable %s with shape %s.' %(vname, vdata.shape))
111            var = f.createVariable(vname,vtype,dimname)
112            var[:] = vdata
113   
114    def meshwrite(self,name):
115        mesh = self._mesh
116        """ Writing Mesh information to disk for Fortran segment"""
117   
118        #----------------OPENING NETCDF OUTPUT FILE------------
119        try:
120            f = cdf.Dataset(name, 'w', format='NETCDF4')
121        except:
122            print("Error occurred while opening new netCDF file. ")
123   
124   
125        f.description = "Python_side_mesh_information"
126        f.primal_num, f.edge_num, f.dual_num = self._mesh.primal_num, self._mesh.edge_num, self._mesh.dual_num
127   
128        #----------------DEFINING DIMENSIONS--------------------
129        max_primal_deg = mesh.primal_deg.max()
130        max_dual_deg = mesh.dual_deg.max()
131        max_trisk_deg = mesh.trisk_deg.max()
132        def crop(dim, *data) : return [ x[:,0:dim] for x in data ]
133        mesh.primal_edge, mesh.primal_ne, mesh.primal_vertex = crop(max_primal_deg, mesh.primal_edge, mesh.primal_ne, mesh.primal_vertex)
134        mesh.dual_edge, mesh.dual_ne = crop(max_dual_deg, mesh.dual_edge, mesh.dual_ne)
135        mesh.trisk, mesh.wee = crop(max_trisk_deg, mesh.trisk, mesh.wee)
136
137        for dimname, dimsize in [("primal_cell", self._mesh.primal_deg.size),
138                                    ("dual_cell", self._mesh.dual_deg.size),
139                                    ("edge", len(self._mesh.edges_E2)),
140                                    ("primal_edge_or_vertex", self._mesh.primal_edge.shape[1]),
141                                    ("dual_edge_or_vertex", self._mesh.dual_edge.shape[1]),
142                                    ("trisk_edge",self._mesh.trisk.shape[1]) ]:
143            f.createDimension(dimname,dimsize)
144       
145        self.create_vars(f,"primal_cell",
146                    [("primal_deg","i4",self._mesh.primal_deg),
147                    ("Ai","f8",self._mesh.Ai),
148                    ("lon_i","f8",self._mesh.lon_i),
149                    ("ref_lon_i","f8",self._mesh.ref_lon_i),
150                    ("lat_i","f8",self._mesh.lat_i),
151                    ("ref_lat_i","f8",self._mesh.ref_lat_i),
152                    ("primal_own_deg","f8",self._mesh.primal_own_deg),
153                    ("primal_own_glo","f8",self._mesh.primal_own_glo),
154                    ("primal_own_loc","f8",self._mesh.primal_own_loc) ])
155   
156        self.create_vars(f,("primal_cell","primal_edge_or_vertex"),
157                         [("primal_vertex","i4",self._mesh.primal_vertex),
158                          ("primal_bounds_lon","f8",self._mesh.primal_bounds_lon),
159                          ("primal_bounds_lat","f8",self._mesh.primal_bounds_lat),
160                          ("primal_edge","i4",self._mesh.primal_edge),
161                          ("primal_ne","i4",self._mesh.primal_ne) ])
162       
163        self.create_vars(f,"dual_cell",
164                    [("dual_deg","i4",self._mesh.dual_deg),
165                    ("Av","f8",self._mesh.Av),
166                    ("fv","f8",self._mesh.fv),
167                    ("lon_v","f8",self._mesh.lon_v),
168                    ("ref_lon_v","f8",self._mesh.ref_lon_v),
169                    ("lat_v","f8",self._mesh.lat_v),
170                    ("ref_lat_v","f8",self._mesh.ref_lat_v),
171                    ("vertices_V1","f8",self._mesh.vertices_V1),
172                    ("dual_own_loc","f8",self._mesh.dual_own_loc)  ])
173   
174        self.create_vars(f,("dual_cell","dual_edge_or_vertex"),
175                    [("dual_edge","i4",self._mesh.dual_edge),
176                     ("dual_vertex","i4",self._mesh.dual_vertex),
177                     ("dual_bounds_lon","f8",self._mesh.dual_bounds_lon),
178                     ("dual_bounds_lat","f8",self._mesh.dual_bounds_lat),
179                     ("dual_ne","i4",self._mesh.dual_ne),
180                     ("Riv2","f8",self._mesh.Riv2)])
181   
182        self.create_vars(f,"edge",
183                    [("left","i4",self._mesh.left),
184                    ("up","i4",self._mesh.up),
185                    ("right","i4",self._mesh.right),
186                    ("down","i4",self._mesh.down),
187                    ("le","f8",self._mesh.le),
188                    ("de","f8",self._mesh.de),               
189                    ("le_de","f8",self._mesh.le_de),
190                    ("trisk_deg","i4",self._mesh.trisk_deg),
191                    ("angle_e","f8",self._mesh.angle_e),
192                    ("lon_e","f8",self._mesh.lon_e),
193                    ("ref_lon_e","f8",self._mesh.ref_lon_e),
194                    ("lat_e","f8",self._mesh.lat_e),
195                    ("ref_lat_e","f8",self._mesh.ref_lat_e),
196                    ("edges_E2","i4",self._mesh.edges_E2)])
197
198        self.create_vars(f,("edge","trisk_edge"),
199                    [("trisk","i4",self._mesh.trisk),
200                    ("wee","f8",self._mesh.wee) ])
201        f.close()
202
203print("writing......")
204f_write = preprocess_output(mesh)
205f_write.meshwrite('mesh_information.nc')
206
Note: See TracBrowser for help on using the repository browser.