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
NX_ATMO = 88838 # number of LAND grid cells at T255 resolution
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
# Constants from namelist
nml = f90nml.read(NAMELIST_FILE_NAME)
COUPLING_INTERVAL = nml['NAMAMIP']['TimeStepSec']
RUN_LENGTH_SEC = nml['NAMAMIP']['RunLengthSec']
# Init OASIS
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('I am IFS-mock process with rank :', mype, file=w_unit)
print('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('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 itap_sec in range(0, RUN_LENGTH_SEC, COUPLING_INTERVAL):
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('--------------------------------------\n', file=w_unit)
w_unit.flush()
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# TERMINATION
# +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
print('End of the program', file=w_unit)
w_unit.flush()
del comp