1 | from dynamico import unstructured as unst |
---|
2 | from dynamico import xios |
---|
3 | from dynamico.meshes import radian, MPAS_Format, Unstructured_PMesh as PMesh, Local_Mesh as Mesh |
---|
4 | import numpy as np |
---|
5 | |
---|
6 | with xios.Client() as client: |
---|
7 | comm = client.comm |
---|
8 | mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() |
---|
9 | print '%d/%d starting'%(mpi_rank,mpi_size) |
---|
10 | |
---|
11 | #----------------------------- read MPAS mesh -------------------------------- |
---|
12 | |
---|
13 | grid, llm, nqdyn, nqtot = 2562, 4,1,1 # 2562, 10242, 40962 |
---|
14 | Omega, radius, g, gh0 = 2.*np.pi/86400., 6.4e6, 1., 2.94e4 |
---|
15 | N, T, courant = 40, 10800., 1.2 # simulation length = N*T |
---|
16 | print 'Omega, planetary PV', Omega, 2*Omega/gh0 |
---|
17 | |
---|
18 | def f(lon,lat): return 2*Omega*np.sin(lat) # Coriolis parameter |
---|
19 | print 'Reading MPAS mesh ...' |
---|
20 | meshfile = MPAS_Format('grids/x1.%d.grid.nc'%grid) |
---|
21 | pmesh = PMesh(comm,meshfile) |
---|
22 | pmesh.partition_metis() |
---|
23 | mesh = Mesh(pmesh, llm, nqdyn, radius, f) |
---|
24 | print '...Done' |
---|
25 | |
---|
26 | #--------------------------------- write some data ---------------------------------------- |
---|
27 | |
---|
28 | with xios.Context_Unstructured(pmesh,mesh,nqtot, 3600) as context: |
---|
29 | lat_i = radian*mesh.lat_i |
---|
30 | lat_ik, junk = np.meshgrid(lat_i, np.arange(llm), indexing='ij') |
---|
31 | |
---|
32 | for i in range(100): |
---|
33 | context.update_calendar(i) |
---|
34 | print 'send_field', i, lat_ik.shape |
---|
35 | context.send_field_primal('ps', lat_i) |
---|
36 | context.send_field_primal('theta', lat_ik) |
---|