1 | from dynamico import cxios |
---|
2 | from dynamico import xios |
---|
3 | from dynamico import unstructured as unst |
---|
4 | from dynamico import meshes |
---|
5 | import numpy as np |
---|
6 | |
---|
7 | # client.__exit__() and context.__exit__() are called at the end of the 'with' blocks |
---|
8 | # even if an exception is raised in the body of the block |
---|
9 | |
---|
10 | with xios.Client() as client: |
---|
11 | comm = client.comm |
---|
12 | mpi_rank, mpi_size = comm.Get_rank(), comm.Get_size() |
---|
13 | print '%d/%d starting'%(mpi_rank,mpi_size) |
---|
14 | |
---|
15 | nqdyn,nqtot,llm = 1,1,4 |
---|
16 | |
---|
17 | if False: |
---|
18 | nx,ny,llm = 64,64 |
---|
19 | Lx,Ly,f = 1.,1.,1. |
---|
20 | mesh = meshes.Cartesian_Mesh(nx,ny,llm,nqdyn,Lx,Ly,f) |
---|
21 | else: |
---|
22 | filename, g, f, radius = 'cart_128_128.nc', 1., 1., None |
---|
23 | # filename = 'cart_008_008.nc' |
---|
24 | unst.setvar('g',g) |
---|
25 | print 'Reading Cartesian mesh ...' |
---|
26 | def coriolis(lon,lat): |
---|
27 | return f+0.*lon |
---|
28 | meshfile = meshes.DYNAMICO_Format(filename) |
---|
29 | nx,ny=meshfile.nx,meshfile.ny |
---|
30 | |
---|
31 | pmesh = meshes.Unstructured_PMesh(comm,meshfile) |
---|
32 | # pmesh.partition_metis() |
---|
33 | pmesh.partition_curvilinear(2,2) |
---|
34 | mesh = meshes.Local_Mesh(pmesh, llm, nqdyn, radius, coriolis) |
---|
35 | |
---|
36 | #--------------------------------- write some data ---------------------------------------- |
---|
37 | |
---|
38 | with xios.Context_Curvilinear(mesh,nqtot, 3600) as context: |
---|
39 | lon_i,lat_i = mesh.lon_i, mesh.lat_i |
---|
40 | lat_ik, junk = np.meshgrid(lat_i, np.arange(llm), indexing='ij') |
---|
41 | lon_ik, junk = np.meshgrid(lon_i, np.arange(llm), indexing='ij') |
---|
42 | |
---|
43 | for i in range(48): |
---|
44 | context.update_calendar(i) |
---|
45 | # print 'send_field', i, lat_ik.shape |
---|
46 | context.send_field_primal('ps', lat_i+lon_i) |
---|
47 | context.send_field_primal('theta', lat_ik+lon_ik) |
---|