Newer
Older
#!/usr/bin/env python3
# Copyright 2021 - Barcelona Supercomputing Center
# Author: Rodrigo Martín Posada
# MIT License
import pyoasis
from pyoasis import OASIS # pylint: disable=no-name-in-module
import numpy as np
# Constants
NX_ATMO = 88838 # number of LAND grid cells at T255 resolution
COUPLING_INTERVAL = 86400
RUN_LENGTH_SEC = 31536000
NUMBER_TIME_STEPS = int(RUN_LENGTH_SEC / COUPLING_INTERVAL)
def def_local_partition(npes, mype):
il_size = NX_ATMO // npes
il_offset = mype * il_size
if (npes-mype) <= NX_ATMO % npes:
il_size += 1
il_offset += (NX_ATMO % npes) - (npes-mype)
return il_size, il_offset
local_comm = comp.localcomm
# Get rank in local communicator
mype = comp.localcomm.rank
npes = comp.localcomm.size
# Unit for output messages
w_unit = open(out_ifs, 'w')
print('-----------------------------------------------------------', file=w_unit)
print(f'I am IFS-mock process with rank : {mype}', file=w_unit)
print(f'in my local communicator gathering {npes} processes', file=w_unit)
print('----------------------------------------------------------', file=w_unit)
w_unit.flush()
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# PARTITION DEFINITION
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# Definition of the local partition
il_size, il_offset = def_local_partition(npes, mype)
print('il_size, il_offset = ', il_size, il_offset, file=w_unit)
w_unit.flush()
partition = pyoasis.ApplePartition(il_offset, il_size)
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# DECLARATION OF THE COUPLING FIELDS
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
A_SST = pyoasis.Var("A_SST", partition, OASIS.IN)
A_Ice_frac = pyoasis.Var("A_Ice_frac", partition, OASIS.IN)
print(f'var_id A_SST, var_id A_Ice_frac {A_SST._id} {A_Ice_frac._id}', file=w_unit)
w_unit.flush()
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# TERMINATION OF DEFINITION PHASE
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
print('End of initialisation phase', file=w_unit)
w_unit.flush()
comp.enddef()
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# TIME STEP LOOP
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
print('Timestep, field min, mean and max value', file=w_unit)
for ib in range(NUMBER_TIME_STEPS):
itap_sec = COUPLING_INTERVAL * ib # time in seconds
field_recv_A_SST = pyoasis.asarray(np.full(il_size, -1.0))
field_recv_A_Ice_frac = pyoasis.asarray(np.full(il_size, -1.0))
A_SST.get(itap_sec, field_recv_A_SST)
A_Ice_frac.get(itap_sec, field_recv_A_Ice_frac)
"{:.4e}".format(np.min(field_recv_A_SST)),
"{:.4e}".format(np.mean(field_recv_A_SST)),
"{:.4e}".format(np.max(field_recv_A_SST)),
file=w_unit)
print('A_Ice_frac ->', itap_sec,
"{:.4e}".format(np.min(field_recv_A_Ice_frac)),
"{:.4e}".format(np.mean(field_recv_A_Ice_frac)),
"{:.4e}".format(np.max(field_recv_A_Ice_frac)),
print(f'--------------------------------------\n', file=w_unit)
w_unit.flush()
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# TERMINATION
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
print('End of the program', file=w_unit)
w_unit.flush()
del comp