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

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

devel/Python : moved Fortran bindings and *.pyx to dynamico/dev module + necessary changes to test/py/*.py
File size: 7.8 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            var = f.createVariable(vname,vtype,dimname)
110            var[:] = vdata
111   
112    def meshwrite(self,name):
113        """ Writing Mesh information to disk for Fortran segment"""
114   
115        #----------------OPENING NETCDF OUTPUT FILE------------
116        try:
117            f = cdf.Dataset(name, 'w', format='NETCDF4')
118        except:
119            print("Error occurred while opening new netCDF file. ")
120   
121   
122        f.description = "Python_side_mesh_information"
123        f.primal_num, f.edge_num, f.dual_num = self._mesh.primal_num, self._mesh.edge_num, self._mesh.dual_num
124   
125        #----------------DEFINING DIMENSIONS--------------------
126        for dimname, dimsize in [("primal_cell", self._mesh.primal_deg.size),
127                                    ("dual_cell", self._mesh.dual_deg.size),
128                                    ("edge", len(self._mesh.edges_E2)),
129                                    ("primal_edge_or_vertex", self._mesh.primal_edge.shape[1]),
130                                    ("dual_edge_or_vertex", self._mesh.dual_edge.shape[1]),
131                                    ("trisk_edge",self._mesh.trisk.shape[1]) ]:
132            f.createDimension(dimname,dimsize)
133       
134        self.create_vars(f,"primal_cell",
135                    [("primal_deg","i4",self._mesh.primal_deg),
136                    ("Ai","f8",self._mesh.Ai),
137                    ("lon_i","f8",self._mesh.lon_i),
138                    ("ref_lon_i","f8",self._mesh.ref_lon_i),
139                    ("lat_i","f8",self._mesh.lat_i),
140                    ("ref_lat_i","f8",self._mesh.ref_lat_i),
141                    ("primal_own_deg","f8",self._mesh.primal_own_deg),
142                    ("primal_own_glo","f8",self._mesh.primal_own_glo),
143                    ("primal_own_loc","f8",self._mesh.primal_own_loc) ])
144   
145        self.create_vars(f,("primal_cell","primal_edge_or_vertex"),
146                    [("primal_vertex","i4",self._mesh.primal_vertex),
147                    ("primal_edge","i4",self._mesh.primal_edge),
148                    ("primal_ne","i4",self._mesh.primal_ne) ])
149   
150        self.create_vars(f,"dual_cell",
151                    [("dual_deg","i4",self._mesh.dual_deg),
152                    ("Av","f8",self._mesh.Av),
153                    ("fv","f8",self._mesh.fv),
154                    ("lon_v","f8",self._mesh.lon_v),
155                    ("ref_lon_v","f8",self._mesh.ref_lon_v),
156                    ("lat_v","f8",self._mesh.lat_v),
157                    ("ref_lat_v","f8",self._mesh.ref_lat_v),
158                    ("vertices_V1","f8",self._mesh.vertices_V1),
159                    ("dual_own_loc","f8",self._mesh.dual_own_loc)  ])
160   
161        self.create_vars(f,("dual_cell","dual_edge_or_vertex"),
162                    [("dual_edge","i4",self._mesh.dual_edge),
163                    ("dual_vertex","i4",self._mesh.dual_vertex),
164                    ("dual_ne","i4",self._mesh.dual_ne),
165                    ("Riv2","f8",self._mesh.Riv2)])
166   
167        self.create_vars(f,"edge",
168                    [("left","i4",self._mesh.left),
169                    ("up","i4",self._mesh.up),
170                    ("right","i4",self._mesh.right),
171                    ("down","i4",self._mesh.down),
172                    ("le","f8",self._mesh.le),
173                    ("de","f8",self._mesh.de),               
174                    ("le_de","f8",self._mesh.le_de),
175                    ("trisk_deg","i4",self._mesh.trisk_deg),
176                    ("angle_e","f8",self._mesh.le_de),
177                    ("lon_e","f8",self._mesh.lon_e),
178                    ("ref_lon_e","f8",self._mesh.ref_lon_e),
179                    ("lat_e","f8",self._mesh.lat_e),
180                    ("ref_lat_e","f8",self._mesh.ref_lat_e),
181                    ("edges_E2","i4",self._mesh.edges_E2)])
182
183        self.create_vars(f,("edge","trisk_edge"),
184                    [("trisk","i4",self._mesh.trisk),
185                    ("wee","f8",self._mesh.wee) ])
186        f.close()
187
188print("writing......")
189f_write = preprocess_output(mesh)
190f_write.meshwrite('mesh_information.nc')
191
Note: See TracBrowser for help on using the repository browser.