1 | from mpi4py import MPI |
---|
2 | comm = MPI.COMM_WORLD |
---|
3 | mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() |
---|
4 | print '%d/%d starting'%(mpi_rank,mpi_size) |
---|
5 | |
---|
6 | import numpy as np |
---|
7 | print 'import dynamico.unstructured' |
---|
8 | import dynamico.unstructured as unst |
---|
9 | print 'import dynamico.xios' |
---|
10 | import dynamico.xios as xios |
---|
11 | print 'Done.' |
---|
12 | from dynamico.meshes import MPAS_Mesh as Mesh, radian |
---|
13 | |
---|
14 | def boundaries(degree,points,lon,lat): |
---|
15 | n, nvertex = len(degree), degree.max() |
---|
16 | bnd_lon, bnd_lat = np.zeros((n,nvertex)),np.zeros((n,nvertex)) |
---|
17 | for ij in range(n): |
---|
18 | nb=degree[ij] |
---|
19 | for k in range(nb): |
---|
20 | vertex = points[ij,k] |
---|
21 | bnd_lon[ij,k]=lon[vertex] |
---|
22 | bnd_lat[ij,k]=lat[vertex] |
---|
23 | for k in range(nb,nvertex): # repeat last vertex if nb<nvertex |
---|
24 | bnd_lon[ij,k]=bnd_lon[ij,nb-1] |
---|
25 | bnd_lat[ij,k]=bnd_lat[ij,nb-1] |
---|
26 | return bnd_lon, bnd_lat |
---|
27 | |
---|
28 | #----------------------------- read MPAS mesh -------------------------------- |
---|
29 | |
---|
30 | grid, llm, nqdyn, nqtot = 10242, 1,1,1 # 2562, 10242, 40962 |
---|
31 | Omega, radius, g, gh0 = 2.*np.pi/86400., 6.4e6, 1., 2.94e4 |
---|
32 | N, T, courant = 40, 10800., 1.2 # simulation length = N*T |
---|
33 | print 'Omega, planetary PV', Omega, 2*Omega/gh0 |
---|
34 | |
---|
35 | def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter |
---|
36 | print 'Reading MPAS mesh ...' |
---|
37 | mesh = Mesh('grids/x1.%d.grid.nc'%grid, llm, nqdyn, radius, f) |
---|
38 | print '...Done' |
---|
39 | |
---|
40 | lon_i, lat_i, lon_v, lat_v = [x*radian for x in mesh.lon_i, mesh.lat_i, mesh.lon_v, mesh.lat_v] |
---|
41 | #lon_i, lat_i, lon_v, lat_v = mesh.lon_i, mesh.lat_i, mesh.lon_v, mesh.lat_v |
---|
42 | |
---|
43 | #--------------------------------- initialize XIOS -------------------------- |
---|
44 | |
---|
45 | def setup_domain(cat,id, degree,vertex,lon,lat,lon_v,lat_v): # cat in 'domain','domaingroup' |
---|
46 | bnd_lon, bnd_lat = boundaries(degree, vertex, lon_v, lat_v) |
---|
47 | ncell_tot, nvertex = bnd_lon.shape |
---|
48 | mesh = xios.Handle(cat,id) |
---|
49 | mesh.set_attr(type='unstructured') |
---|
50 | mesh.set_attr(ni_glo=ncell_tot, ibegin=0, ni=ncell_tot) |
---|
51 | mesh.set_attr(i_index=np.arange(ncell_tot,dtype=np.int32)) |
---|
52 | mesh.set_attr(lonvalue_1d=lon, latvalue_1d=lat) |
---|
53 | mesh.set_attr(nvertex=nvertex,bounds_lon_1d=bnd_lon, bounds_lat_1d=bnd_lat) |
---|
54 | |
---|
55 | unst.ker.dynamico_setup_xios() |
---|
56 | |
---|
57 | if True: |
---|
58 | lev, levp1, nq = [xios.Handle('axis',name) for name in 'lev', 'levp1', 'nq'] |
---|
59 | lev.set_attr(n_glo=llm, value=np.arange(llm)) |
---|
60 | levp1.set_attr(n_glo=llm+1, value=np.arange(llm+1)) |
---|
61 | nq.set_attr(n_glo=nqtot, value=np.arange(nqtot)) |
---|
62 | mesh_i=setup_domain('domaingroup','i', mesh.primal_deg, mesh.primal_vertex-1, lon_i, lat_i, lon_v, lat_v) # primal mesh |
---|
63 | # mesh_v=setup_domain('domain','v', mesh.dual_deg, mesh.dual_vertex-1, lon_v, lat_v, lon_i, lat_i) # dual mesh |
---|
64 | |
---|
65 | # CALL xios_set_domain_attr("u",ni_glo=ncell_tot, ibegin=displ, ni=ncell) |
---|
66 | # CALL xios_set_domain_attr("u", data_dim=1, type='unstructured' , nvertex=2, i_index=ind_glo) |
---|
67 | # CALL xios_set_domain_attr("u",lonvalue_1d=lon, latvalue_1d=lat, bounds_lon_1d=bounds_lon, bounds_lat_1d=bounds_lat) |
---|
68 | #mesh_e = xios.Handle('domain','u') |
---|
69 | |
---|
70 | # works up to that point ! |
---|
71 | |
---|
72 | dtime = xios.Duration(second=3600.) |
---|
73 | |
---|
74 | # CALL xios_set_timestep(dtime) |
---|
75 | if True: |
---|
76 | calendar=xios.Handle('calendar_wrapper') |
---|
77 | print (dtime.year, dtime.month, dtime.day, dtime.hour, dtime.second) |
---|
78 | calendar.set_attr(timestep=dtime) |
---|
79 | calendar.update_timestep() |
---|
80 | else: |
---|
81 | unst.ker.dynamico_xios_set_timestep(3600.) |
---|
82 | |
---|
83 | # CALL xios_set_fieldgroup_attr("standard_output", freq_op=itau_out*xios_timestep, freq_offset=(itau_out-1)*xios_timestep) |
---|
84 | if False: |
---|
85 | std=xios.Handle('fieldgroup','standard_output') |
---|
86 | std.set_attr(freq_op=dtime) |
---|
87 | freq_offset = xios.Duration(second=0) |
---|
88 | std.set_attr(freq_offset=freq_offset) |
---|
89 | |
---|
90 | if True: |
---|
91 | print 'xios.context_close_definition()' |
---|
92 | xios.context_close_definition() |
---|
93 | |
---|
94 | no_error=True |
---|
95 | try: |
---|
96 | for i in range(100): |
---|
97 | if True: |
---|
98 | unst.ker.dynamico_xios_update_calendar(i) |
---|
99 | else: |
---|
100 | xios.update_calendar(i) |
---|
101 | print 'send_field', i |
---|
102 | xios.send_field('ps', lat_i) |
---|
103 | except: |
---|
104 | print 'An error has occured, trying to close XIOS properly' |
---|
105 | no_error=False |
---|
106 | # CALL xios_recv_field(name,field) |
---|
107 | # CALL xios_send_field(name,field_tmp) |
---|
108 | |
---|
109 | print 'xios.context_finalize()' |
---|
110 | xios.context_finalize() |
---|
111 | |
---|
112 | print 'xios.finalize()' |
---|
113 | xios.finalize() |
---|
114 | |
---|
115 | print 'Done' |
---|
116 | |
---|
117 | if not no_error : raise |
---|