1 | from __future__ import print_function |
---|
2 | import numpy as np |
---|
3 | import cxios |
---|
4 | from unstructured import ker |
---|
5 | from dynamico.meshes import radian |
---|
6 | from mpi4py import MPI |
---|
7 | |
---|
8 | from dynamico import getargs |
---|
9 | args=getargs.parse() |
---|
10 | log_master, log_world = getargs.getLogger(__name__) |
---|
11 | INFO, DEBUG, ERROR = log_master.info, log_master.debug, log_world.error |
---|
12 | INFO_ALL, DEBUG_ALL = log_world.info, log_world.debug |
---|
13 | |
---|
14 | #----------------------------------------------------------------------------- |
---|
15 | |
---|
16 | class Client: |
---|
17 | def __enter__(self): |
---|
18 | ker.dynamico_setup_xios() |
---|
19 | self.comm=MPI.COMM_WORLD # FIXME : we should use the communicator returned by XIOS |
---|
20 | return self |
---|
21 | def __exit__(self, type, value, traceback): |
---|
22 | INFO('xios_finalize()') |
---|
23 | DEBUG_ALL('xios_finalize') |
---|
24 | cxios.finalize() |
---|
25 | def min(self,data): return self.comm.allreduce(data, op=MPI.MIN) |
---|
26 | def max(self,data): return self.comm.allreduce(data, op=MPI.MAX) |
---|
27 | |
---|
28 | class Context: |
---|
29 | def __enter__(self): |
---|
30 | INFO('cxios.context_close_definition()') |
---|
31 | cxios.context_close_definition() |
---|
32 | return self |
---|
33 | def __exit__(self, type, value, traceback): |
---|
34 | INFO('xios_context_finalize()') |
---|
35 | cxios.context_finalize() |
---|
36 | def init_llm(self, mesh, nqtot): |
---|
37 | self.mesh=mesh |
---|
38 | # vertical axis, nq |
---|
39 | llm = mesh.llm |
---|
40 | lev, levp1, nq = [cxios.Handle('axis',name) for name in 'lev', 'levp1', 'nq'] |
---|
41 | lev.set_attr(n_glo=llm, value=np.arange(llm)) |
---|
42 | levp1.set_attr(n_glo=llm+1, value=np.arange(llm+1)) |
---|
43 | nq.set_attr(n_glo=nqtot, value=np.arange(nqtot)) |
---|
44 | def init_calendar(self, step): |
---|
45 | calendar=cxios.Handle('calendar_wrapper') |
---|
46 | dtime = cxios.Duration(second=step) |
---|
47 | calendar.set_attr(timestep=dtime) |
---|
48 | calendar.update_timestep() |
---|
49 | std=cxios.Handle('fieldgroup','standard_output') |
---|
50 | std.set_attr(freq_op=dtime) |
---|
51 | freq_offset = cxios.Duration(second=0) |
---|
52 | std.set_attr(freq_offset=freq_offset) |
---|
53 | def update_calendar(self, i): |
---|
54 | ker.dynamico_xios_update_calendar(i) |
---|
55 | def send_field_primal(self,name, data): |
---|
56 | self.send_field(self.mesh.primal_own_loc, name, data) |
---|
57 | def send_field_dual(self,name, data): |
---|
58 | self.send_field(self.mesh.dual_own_loc, name, data) |
---|
59 | def send_field(self,own_loc, name, data): |
---|
60 | if data.ndim==1: |
---|
61 | data = data[own_loc] |
---|
62 | if data.ndim==2: |
---|
63 | data = data[own_loc,:] |
---|
64 | data = data.transpose() # XIOS expects contiguous horizontal slices |
---|
65 | data = np.ascontiguousarray(data) |
---|
66 | cxios.send_field(name, data) |
---|
67 | |
---|
68 | #----------------------------------------------------------------------------- |
---|
69 | |
---|
70 | def setup_curvilinear(cat,id,nx,ny,own,i_index,j_index,lon,lat): |
---|
71 | mesh = cxios.Handle(cat,id) |
---|
72 | def log(name,data) : |
---|
73 | DEBUG('setup_curvilinear : %s shape min max %s %s %s'%(name, data.shape, data.min(), data.max()) ) |
---|
74 | i_index, j_index, lon, lat = [ x[own] for x in i_index, j_index, lon, lat ] |
---|
75 | log('i_index',i_index) |
---|
76 | log('j_index',j_index) |
---|
77 | log('lon',lon) |
---|
78 | log('lat',lat) |
---|
79 | mesh.set_attr(type='curvilinear', ni_glo=nx, i_index=i_index, nj_glo=ny, j_index=j_index) |
---|
80 | mesh.set_attr(lonvalue_1d=lon.flatten(), latvalue_1d=lat.flatten()) |
---|
81 | |
---|
82 | class Context_Curvilinear(Context): |
---|
83 | def __init__(self, mesh, nqtot, step): |
---|
84 | self.init_llm(mesh, nqtot) |
---|
85 | # primal mesh |
---|
86 | setup_curvilinear('domaingroup','i', mesh.nx, mesh.ny, mesh.primal_own_loc, |
---|
87 | mesh.primal_i, mesh.primal_j, mesh.lon_i, mesh.lat_i) |
---|
88 | # dual mesh |
---|
89 | setup_curvilinear('domaingroup','v', mesh.nx, mesh.ny, mesh.dual_own_loc, |
---|
90 | mesh.dual_i, mesh.dual_j, mesh.lon_v, mesh.lat_v) |
---|
91 | # calendar |
---|
92 | self.init_calendar(step) |
---|
93 | |
---|
94 | #----------------------------------------------------------------------------- |
---|
95 | |
---|
96 | def cf_boundaries(degree,points,lon,lat): |
---|
97 | n, nvertex = len(degree), degree.max() |
---|
98 | bnd_lon, bnd_lat = np.zeros((n,nvertex)),np.zeros((n,nvertex)) |
---|
99 | for ij in range(n): |
---|
100 | nb=degree[ij] |
---|
101 | for k in range(nb): |
---|
102 | vertex = points[ij,k] |
---|
103 | bnd_lon[ij,k]=lon[vertex] |
---|
104 | bnd_lat[ij,k]=lat[vertex] |
---|
105 | for k in range(nb,nvertex): # repeat last vertex if nb<nvertex |
---|
106 | bnd_lon[ij,k]=bnd_lon[ij,nb-1] |
---|
107 | bnd_lat[ij,k]=bnd_lat[ij,nb-1] |
---|
108 | return bnd_lon, bnd_lat |
---|
109 | |
---|
110 | def setup_unstructured(cat,id, degree,vertex,lon,lat,lon_v,lat_v, ncell_glo, index_own_glo, displ): # cat in 'domain','domaingroup' |
---|
111 | # ncell_glo (int) : global number of cells |
---|
112 | # index_own_glo : global indices of cells own by this MPI process |
---|
113 | # displ : min of index_own across all MPI processes |
---|
114 | bnd_lon, bnd_lat = cf_boundaries(degree, vertex, lon_v, lat_v) |
---|
115 | ncell, nvertex = bnd_lon.shape |
---|
116 | mesh = cxios.Handle(cat,id) |
---|
117 | mesh.set_attr(type='unstructured') |
---|
118 | mesh.set_attr(ni_glo=ncell_glo, ni=ncell, ibegin=displ, i_index=index_own_glo) |
---|
119 | mesh.set_attr(lonvalue_1d=lon, latvalue_1d=lat) |
---|
120 | mesh.set_attr(nvertex=nvertex,bounds_lon_1d=bnd_lon, bounds_lat_1d=bnd_lat) |
---|
121 | |
---|
122 | class Context_Unstructured(Context): |
---|
123 | def __init__(self, pmesh, mesh,nqtot, step): |
---|
124 | self.init_llm(mesh, nqtot) |
---|
125 | # primal mesh |
---|
126 | 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] |
---|
127 | setup_unstructured('domaingroup','i', mesh.primal_own_deg, mesh.primal_vertex, |
---|
128 | lon_i[mesh.primal_own_loc], lat_i[mesh.primal_own_loc], |
---|
129 | lon_v, lat_v, |
---|
130 | pmesh.dim_primal.n, mesh.primal_own_glo, mesh.displ) |
---|
131 | # calendar |
---|
132 | self.init_calendar(step) |
---|