DataCouplerOasisVar.py

# Copyright 2022 - Barcelona Supercomputing Center
# Authors: Etienne Tourigny, Rodrigo Martín Posada
# MIT License
import logging
import numpy as np
import math
import pyoasis
from pyoasis import OASIS

from DataCouplerUtils import scale_factor_from_str

class OasisOutputVar:
    def __init__(self, conf, partition, grid_shape, fileIn_vars ):
        self._conf = conf
        self._id = conf['id']
        self._reset = conf['reset']
        self._scale_factor = scale_factor_from_str(conf['scale_factor']) if type(conf['scale_factor']) == str else conf['scale_factor']
        self._offset = np.float64(conf['offset'])
        self._grid_shape = grid_shape
        self._oasis_var = pyoasis.Var(self._id, partition, OASIS.OUT)
        self._send_field = pyoasis.asarray(np.zeros(grid_shape, dtype=np.float64))
        self._fileIn_vars = fileIn_vars
        logging.debug('OasisOutputVar() conf={} reset={}'.format(str(conf),str(self._reset)))

    def add_file_input_var(self, varname, var):
        self._fileIn_vars[varname] = var

    def file_input_vars(self):
        return self._fileIn_vars

    def update(self,time):
        logging.debug('OasisOutputVar::update() var {} at time {}'.format(self._id,time))
        logging.debug('OasisOutputVar() conf={} reset={}'.format(str(self._conf),str(self._reset)))
        # reset if needed
        if self._reset:
            self._send_field.fill(0.)
            logging.debug('reset field for var {} at time {}'.format(self._id,time))
        # accumulate values from each fileIn_var linked to this oasisOut_var
        for (varname,fileIn_var) in self._fileIn_vars.items():
            if time == 0 or fileIn_var.accum:
                self._send_field += fileIn_var.send_field()
                logging.debug('accumulating {} from fileIn var {} for oasisOut var {} total= {}'.format(fileIn_var.send_field(), varname, self._id, self._send_field))

    def put(self, time):
        # send the data
        logging.debug('OasisOutputVar::put() var {} at time {} field {}'.format(self._id,time,self._send_field))
        if math.isclose(self._scale_factor,1) and math.isclose(self._offset,0):
            self._oasis_var.put(time, self._send_field)
        else:
            logging.debug('OasisOutputVar::put() var {} with offset {} and scale_factor {}'.format(self._id,self._offset,self._scale_factor))
            self._oasis_var.put(time, self._send_field*self._scale_factor+self._offset)