1 | #!/usr/bin/env python3 |
---|
2 | |
---|
3 | import os |
---|
4 | import math |
---|
5 | import numpy as np |
---|
6 | import pyoasis |
---|
7 | from pyoasis import OASIS |
---|
8 | import netCDF4 |
---|
9 | from mpi4py import MPI |
---|
10 | from utils import * |
---|
11 | |
---|
12 | comm = MPI.COMM_WORLD |
---|
13 | |
---|
14 | has_graphics = None |
---|
15 | has_graphics = comm.bcast(has_graphics, root=comm.size - 1) |
---|
16 | |
---|
17 | sgrid = None |
---|
18 | dgrid = None |
---|
19 | ll_plot = None |
---|
20 | if comm.rank == 0: |
---|
21 | sgrid = input('Enter the source grid code from {}:\n'.format(set(valid_grids))) |
---|
22 | if not grid_is_valid(sgrid): |
---|
23 | print('{} is not a valid grid'.format(sgrid), flush=True) |
---|
24 | comm.Abort() |
---|
25 | if grid_is_ocean(sgrid): |
---|
26 | dset = set(valid_grids) - set(ocean_grids) |
---|
27 | else: |
---|
28 | dset = set(valid_grids) - set((sgrid,)) |
---|
29 | dgrid = input('Enter the target grid code from {}:\n'.format(dset)) |
---|
30 | if not grid_is_valid(dgrid): |
---|
31 | print('{} is not a valid grid'.format(dgrid), flush=True) |
---|
32 | comm.Abort() |
---|
33 | if grid_is_ocean(sgrid) and grid_is_ocean(dgrid): |
---|
34 | print('Only one grid can be for ocean', flush=True) |
---|
35 | comm.Abort() |
---|
36 | if sgrid == 'torc' or dgrid == 'torc': |
---|
37 | os.symlink(os.path.join('..', '..', 'common_data', 'masks_torc_scrip.nc'), |
---|
38 | 'masks.nc') |
---|
39 | elif sgrid == 'nogt' or dgrid == 'nogt': |
---|
40 | os.symlink(os.path.join('..', '..', 'common_data', 'masks_nogt_scrip.nc'), |
---|
41 | 'masks.nc') |
---|
42 | else: |
---|
43 | os.symlink(os.path.join('..', '..', 'common_data', 'masks_no_atm.nc'), 'masks.nc') |
---|
44 | do_plot = input('Plot output [yes/no]\n') |
---|
45 | if (do_plot.lower() != 'yes' and do_plot.lower() != 'no'): |
---|
46 | print('{} is not a valid yes/no answer'.format(do_plot), flush=True) |
---|
47 | comm.Abort() |
---|
48 | else: |
---|
49 | ll_plot = do_plot.lower() == 'yes' |
---|
50 | |
---|
51 | write_namcouple(sgrid, dgrid, has_graphics) |
---|
52 | |
---|
53 | dgrid = comm.bcast(dgrid, root=0) |
---|
54 | ll_plot = comm.bcast(ll_plot, root=0) |
---|
55 | |
---|
56 | component_name = "sender-apple" |
---|
57 | comp = pyoasis.Component(component_name, True, comm) |
---|
58 | |
---|
59 | if comm.rank == 0: |
---|
60 | print(comp, flush=True) |
---|
61 | |
---|
62 | comm_rank = comp.localcomm.rank |
---|
63 | comm_size = comp.localcomm.size |
---|
64 | |
---|
65 | sgrid = comp.localcomm.bcast(sgrid, root=0) |
---|
66 | |
---|
67 | gf = netCDF4.Dataset('grids.nc', 'r') |
---|
68 | lons = gf.variables[sgrid + '.lon'][:, :].flatten() |
---|
69 | lats = gf.variables[sgrid + '.lat'][:, :].flatten() |
---|
70 | n_points = lons.size |
---|
71 | gf.close() |
---|
72 | |
---|
73 | if comm_rank == 0: |
---|
74 | print(comp) |
---|
75 | print("n_points on source side is {}".format(n_points)) |
---|
76 | |
---|
77 | local_size = int(n_points / comm_size) |
---|
78 | offset = comm_rank * local_size |
---|
79 | if comm_rank == comm_size - 1: |
---|
80 | local_size = n_points - offset |
---|
81 | |
---|
82 | partition = pyoasis.ApplePartition(offset, local_size) |
---|
83 | |
---|
84 | variable = pyoasis.Var("FSENDANA", partition, OASIS.OUT, bundle_size=2) |
---|
85 | comp.enddef() |
---|
86 | |
---|
87 | date = int(0) |
---|
88 | bundle = pyoasis.asarray(np.zeros((local_size, 2), dtype=np.float64)) |
---|
89 | |
---|
90 | dp_conv = math.pi / 180. |
---|
91 | bundle[:, 0] = 2.0 + np.sin(2.0 * lats[offset:offset + local_size] * dp_conv) ** 4.0 * \ |
---|
92 | np.cos(4.0 * lons[offset:offset + local_size] * dp_conv) |
---|
93 | |
---|
94 | bundle[:, 1] = 2.0 - np.cos(math.pi * |
---|
95 | (np.arccos(np.cos(lons[offset:offset + local_size] * dp_conv) * |
---|
96 | np.cos(lats[offset:offset + local_size] * dp_conv)) / |
---|
97 | (1.2 * math.pi))) |
---|
98 | |
---|
99 | if comm_rank == 0: |
---|
100 | print("Sent data: at time {}".format(date)) |
---|
101 | |
---|
102 | variable.put(date, bundle) |
---|
103 | |
---|
104 | del comp |
---|