from dynamico import getargs getargs.add("--LAM", action='store_true') # Args for both cases getargs.add("--Omega", type=float, help='Planetary radius', default=7e-5) # Args for global mesh getargs.add("--grid", type=int, help='Number of hexagons', default=2562) getargs.add("--radius", type=float, help='Planetary radius', default=6.4e6) # 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() for arg in vars(args): print arg, getattr(args, arg) log_master, log_world = getargs.getLogger() INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error INFO('Starting') from mpi4py import MPI comm = MPI.COMM_WORLD mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() INFO('%d/%d starting'%(mpi_rank,mpi_size)) prefix='fig_RSW2_MPAS_W02/%02d'%mpi_rank INFO('Loading DYNAMICO modules ...') from dynamico.dev 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' print 'Loading modules ...' import math as math import matplotlib.pyplot as plt import numpy as np import netCDF4 as cdf print '...Done' #--------------------------- 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 Omega, radius = args.Omega, args.radius print 'Omega :', Omega 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.) mesh = Mesh(pmesh, llm, nqdyn, planet) davies = myDavies(args.Davies_N1, args.Davies_N2, mesh.ref_lon_i, mesh.ref_lat_i, mesh.ref_lon_e,mesh.ref_lat_e) 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' def print_attr(obj): #for name in obj.__dict__.keys(): print name, type(getattr(obj,name)) for name in obj.__dict__.keys(): print name, type(getattr(obj,name)), np.shape(getattr(obj,name))#,np.dtype(getattr(obj,name)) #for key in obj.__dict__.keys(): print key, np.dtype(obj.__dict__.values()) print_attr(mesh) print("--------") print_attr(mesh.com_edges) print("--------") class preprocess_output: def __init__(self,mesh): self._mesh = mesh #-------------------START WRITING THE netCDF file--------------------------- def create_vars(self,f,dimname, info): for vname, vtype, vdata in info: vdata = np.asarray(vdata) print('creating variable %s with shape %s.' %(vname, vdata.shape)) var = f.createVariable(vname,vtype,dimname) var[:] = vdata def meshwrite(self,name): mesh = self._mesh """ Writing Mesh information to disk for Fortran segment""" #----------------OPENING NETCDF OUTPUT FILE------------ try: f = cdf.Dataset(name, 'w', format='NETCDF4') except: print("Error occurred while opening new netCDF file. ") f.description = "Python_side_mesh_information" f.primal_num, f.edge_num, f.dual_num = self._mesh.primal_num, self._mesh.edge_num, self._mesh.dual_num #----------------DEFINING DIMENSIONS-------------------- max_primal_deg = mesh.primal_deg.max() max_dual_deg = mesh.dual_deg.max() max_trisk_deg = mesh.trisk_deg.max() def crop(dim, *data) : return [ x[:,0:dim] for x in data ] mesh.primal_edge, mesh.primal_ne, mesh.primal_vertex = crop(max_primal_deg, mesh.primal_edge, mesh.primal_ne, mesh.primal_vertex) mesh.dual_edge, mesh.dual_ne = crop(max_dual_deg, mesh.dual_edge, mesh.dual_ne) mesh.trisk, mesh.wee = crop(max_trisk_deg, mesh.trisk, mesh.wee) for dimname, dimsize in [("primal_cell", self._mesh.primal_deg.size), ("dual_cell", self._mesh.dual_deg.size), ("edge", len(self._mesh.edges_E2)), ("primal_edge_or_vertex", self._mesh.primal_edge.shape[1]), ("dual_edge_or_vertex", self._mesh.dual_edge.shape[1]), ("trisk_edge",self._mesh.trisk.shape[1]) ]: f.createDimension(dimname,dimsize) self.create_vars(f,"primal_cell", [("primal_deg","i4",self._mesh.primal_deg), ("Ai","f8",self._mesh.Ai), ("lon_i","f8",self._mesh.lon_i), ("ref_lon_i","f8",self._mesh.ref_lon_i), ("lat_i","f8",self._mesh.lat_i), ("ref_lat_i","f8",self._mesh.ref_lat_i), ("primal_own_deg","f8",self._mesh.primal_own_deg), ("primal_own_glo","f8",self._mesh.primal_own_glo), ("primal_own_loc","f8",self._mesh.primal_own_loc) ]) self.create_vars(f,("primal_cell","primal_edge_or_vertex"), [("primal_vertex","i4",self._mesh.primal_vertex), ("primal_bounds_lon","f8",self._mesh.primal_bounds_lon), ("primal_bounds_lat","f8",self._mesh.primal_bounds_lat), ("primal_edge","i4",self._mesh.primal_edge), ("primal_ne","i4",self._mesh.primal_ne) ]) self.create_vars(f,"dual_cell", [("dual_deg","i4",self._mesh.dual_deg), ("Av","f8",self._mesh.Av), ("fv","f8",self._mesh.fv), ("lon_v","f8",self._mesh.lon_v), ("ref_lon_v","f8",self._mesh.ref_lon_v), ("lat_v","f8",self._mesh.lat_v), ("ref_lat_v","f8",self._mesh.ref_lat_v), ("vertices_V1","f8",self._mesh.vertices_V1), ("dual_own_loc","f8",self._mesh.dual_own_loc) ]) self.create_vars(f,("dual_cell","dual_edge_or_vertex"), [("dual_edge","i4",self._mesh.dual_edge), ("dual_vertex","i4",self._mesh.dual_vertex), ("dual_bounds_lon","f8",self._mesh.dual_bounds_lon), ("dual_bounds_lat","f8",self._mesh.dual_bounds_lat), ("dual_ne","i4",self._mesh.dual_ne), ("Riv2","f8",self._mesh.Riv2)]) self.create_vars(f,"edge", [("left","i4",self._mesh.left), ("up","i4",self._mesh.up), ("right","i4",self._mesh.right), ("down","i4",self._mesh.down), ("le","f8",self._mesh.le), ("de","f8",self._mesh.de), ("le_de","f8",self._mesh.le_de), ("trisk_deg","i4",self._mesh.trisk_deg), ("angle_e","f8",self._mesh.le_de), ("lon_e","f8",self._mesh.lon_e), ("ref_lon_e","f8",self._mesh.ref_lon_e), ("lat_e","f8",self._mesh.lat_e), ("ref_lat_e","f8",self._mesh.ref_lat_e), ("edges_E2","i4",self._mesh.edges_E2)]) self.create_vars(f,("edge","trisk_edge"), [("trisk","i4",self._mesh.trisk), ("wee","f8",self._mesh.wee) ]) f.close() print("writing......") f_write = preprocess_output(mesh) f_write.meshwrite('mesh_information.nc')