[6328] | 1 | #!/usr/bin/env python3 |
---|
| 2 | |
---|
| 3 | import pyoasis |
---|
| 4 | from pyoasis import OASIS |
---|
| 5 | import numpy |
---|
| 6 | |
---|
| 7 | from mpi4py import MPI |
---|
| 8 | |
---|
| 9 | # Test case for the most advanced partition and put/get options. |
---|
| 10 | # Credits: Eric Maisonnave |
---|
| 11 | |
---|
| 12 | comm = MPI.COMM_WORLD |
---|
| 13 | |
---|
| 14 | component_name = "sendrecv" |
---|
| 15 | |
---|
| 16 | comp = pyoasis.Component(component_name, True, comm) |
---|
| 17 | |
---|
| 18 | comm_rank = comp.localcomm.rank |
---|
| 19 | comm_size = comp.localcomm.size |
---|
| 20 | |
---|
| 21 | if comm_rank == 0: |
---|
| 22 | print("Component name: {} = Component ID: {}".format(component_name, |
---|
| 23 | comp._id)) |
---|
| 24 | |
---|
| 25 | n_points = 16 |
---|
| 26 | |
---|
| 27 | # Parameters for incoming partition |
---|
| 28 | |
---|
| 29 | local_lons = (n_points // 2) // comm_size |
---|
| 30 | offsets = [comm_rank * local_lons, |
---|
| 31 | comm_rank * local_lons + (n_points // 2)] |
---|
| 32 | extents = [local_lons, |
---|
| 33 | local_lons] |
---|
| 34 | |
---|
| 35 | # Parameters for outgoing partition: notice that the boxes extents |
---|
| 36 | # do not sum up to n_points, hence the usually optional global_size |
---|
| 37 | # argument of part_out __init__ is in this case mandatory |
---|
| 38 | |
---|
| 39 | global_offsets = [0, 2, 4, 6] |
---|
| 40 | extents_x = [0, 2, 0, 2] |
---|
| 41 | extents_y = [0, 2, 0, 2] |
---|
| 42 | |
---|
| 43 | # Use optional argument name of partitions __init__ for the case |
---|
| 44 | # depicted in Oasis user guide as |
---|
| 45 | # mandatory if oasis def partition is called either for a grid |
---|
| 46 | # decomposed not across all the processes of a component or if |
---|
| 47 | # the related oasis def partition are not called in the same |
---|
| 48 | # order on the different component processes |
---|
| 49 | |
---|
| 50 | if comm_rank % 2 != 0: |
---|
| 51 | part_in = pyoasis.OrangePartition(offsets, extents, name="part_in") |
---|
| 52 | |
---|
| 53 | part_out = pyoasis.BoxPartition(global_offsets[comm_rank], |
---|
| 54 | extents_x[comm_rank], |
---|
| 55 | extents_y[comm_rank], |
---|
| 56 | (n_points // 2), |
---|
| 57 | global_size=n_points, |
---|
| 58 | name="part_out") |
---|
| 59 | else: |
---|
| 60 | part_out = pyoasis.BoxPartition(global_offsets[comm_rank], |
---|
| 61 | extents_x[comm_rank], |
---|
| 62 | extents_y[comm_rank], |
---|
| 63 | (n_points // 2), |
---|
| 64 | global_size=n_points, |
---|
| 65 | name="part_out") |
---|
| 66 | |
---|
| 67 | part_in = pyoasis.OrangePartition(offsets, extents, name="part_in") |
---|
| 68 | |
---|
| 69 | if comm_rank == 0: |
---|
| 70 | print("{}: part_in id: {}".format(component_name, part_in._id)) |
---|
| 71 | print("{}: part_out id: {}".format(component_name, part_out._id)) |
---|
| 72 | |
---|
| 73 | variable = pyoasis.Var("FRECVATM", part_in, OASIS.IN, bundle_size=2) |
---|
| 74 | if comm_rank == 0: |
---|
| 75 | print("{}: var_name: FRECVATM = var_id: {}".format(component_name, |
---|
| 76 | variable._id)) |
---|
| 77 | |
---|
| 78 | var_out = pyoasis.Var("FSENDATM", part_out, OASIS.OUT) |
---|
| 79 | if comm_rank == 0: |
---|
| 80 | print("{}: var_name: FSENDATM = var_id: {}".format(component_name, var_out._id)) |
---|
| 81 | |
---|
| 82 | comp.enddef() |
---|
| 83 | |
---|
| 84 | date = int(0) |
---|
| 85 | bundle = pyoasis.asarray(numpy.zeros((local_lons, 2, 2), |
---|
| 86 | dtype=numpy.float64)) |
---|
| 87 | |
---|
| 88 | variable.get(date, bundle) |
---|
| 89 | |
---|
| 90 | expected_bundle = pyoasis.asarray(numpy.zeros((local_lons, 2, 2), |
---|
| 91 | dtype=numpy.float64)) |
---|
| 92 | for i in range(2): |
---|
| 93 | expected_bundle[:, :, i] = i + 1 |
---|
| 94 | for i in range(local_lons): |
---|
| 95 | expected_bundle[i, :, :] += (i + 1 + comm_rank * local_lons) * 100 |
---|
| 96 | for i in range(2): |
---|
| 97 | expected_bundle[:, i, :] += (i + 1) * 10 |
---|
| 98 | |
---|
| 99 | epsilon = 1e-8 |
---|
| 100 | error = numpy.abs(bundle - expected_bundle).sum() |
---|
| 101 | if error < epsilon: |
---|
| 102 | print("{}: On rank {} data received successfully".format(component_name, |
---|
| 103 | comm_rank)) |
---|
| 104 | |
---|
| 105 | for i in range(2): |
---|
| 106 | print("{}: On rank {} Bundle {} is".format(component_name, comm_rank, |
---|
| 107 | i + 1)) |
---|
| 108 | print(bundle[:, 0, i]) |
---|
| 109 | print(bundle[:, 1, i]) |
---|
| 110 | |
---|
| 111 | if comm_rank % 2 != 0: |
---|
| 112 | field = pyoasis.asarray(bundle[:, :, 1], dtype=numpy.float32) |
---|
| 113 | date = int(0) |
---|
| 114 | var_out.put(date, field) |
---|
| 115 | |
---|
| 116 | del comp |
---|