funcs_solver.py

# -*- coding: utf-8 -*-
"""
Utility functions for the solver step of the Downscaling

@author: Mino Sorribas

TODO: check type 4 !

"""


import numpy as np
import pandas as pd
from datetime import datetime,timedelta
import itertools

#-----------------------------------------------------------------------------
# DEFAULT FILES FOR MGB
#-----------------------------------------------------------------------------
def mgbsa_default(version = '1979'):
    """ Default settings for MGB-SA """

    # number of intervals and start date MGB-SA (1990->)
    if version == '1990':
        nc = 33749
        nt = 7305
        dstart = datetime(1990,1,1)
        file_qtudo = 'QTUDO_1990.MGB'
        file_qcel  = 'QITUDO_1990.MGB'

    # ###
    # number of intervals and start date MGB-SA (1979->)
    if version == '1979':
        nc = 33749
        nt = 13149
        dstart = datetime(1979,1,1)
        file_qtudo = 'QTUDO_1979.MGB'
        file_qcel  = 'QITUDO_1979.MGB'

   # number of intervals and start date MGB-SA ENKF (1979->) MEMBRO 25
    if version == 'enkf_1979':
        nc = 33749
        nt = 13149
        dstart = datetime(1979,1,1)
        file_qtudo = 'QTUDO25.MGB'
        file_qcel  = 'QITUDO25.MGB'

   # number of intervals and start date MGB-SA ENKF (1979->) MEMBRO 48
    if version == 'enkf_1979_m48':
        nc = 33749
        nt = 13149
        dstart = datetime(1979,1,1)
        file_qtudo = 'QTUDO48.MGB'
        file_qcel  = 'QITUDO48.MGB'

   # number of intervals and start date MGB-SA ENKF (1979->) MEMBRO 2
    if version == 'enkf_1979_m02':
        nc = 33749
        nt = 13149
        dstart = datetime(1979,1,1)
        file_qtudo = 'QTUDO02.MGB'
        file_qcel  = 'QITUDO02.MGB'

    return (nt, nc, dstart, file_qtudo, file_qcel)





#-----------------------------------------------------------------------------
# UTILS
#-----------------------------------------------------------------------------
def make_dict_bho_ixc(the_dicts):
    """
    Make dictionary of required MGB catchments (mini)
        for the downscaling at each cotrecho
        which is used for partial (low memory usage) reading of MGB binaries
    """


    # dict of solver (keys are cotrechos available to downscale)
    dict_bho_solver = the_dicts['dict_bho_solver']
    list_to_downscale = dict_bho_solver.keys()

    # dicts of parameters
    dict_parameters_t1 = the_dicts['dict_parameters_t1']
    dict_parameters_t2 = the_dicts['dict_parameters_t2']
    dict_parameters_t3 = the_dicts['dict_parameters_t3']
    dict_parameters_t4 = the_dicts['dict_parameters_t4']

    # identify required 'mini' for each cotrecho
    dict_type_params = {
        1: dict_parameters_t1,
        2: dict_parameters_t2,
        3: dict_parameters_t3,
        4: dict_parameters_t4,
        }

    dict_bho_ixc = {}
    for c in list_to_downscale:

        # get type of solver
        type_solver = dict_bho_solver.get(c)     #1,2,3 or 4

        # get dict of parameters tx -> parameters for cotrecho
        d_params = dict_type_params.get(type_solver).get(c)

        # get list of parameters related to 'mini'
        as_list = lambda x: x if isinstance(x,list) else [x]
        params_mini = [as_list(v) for k,v in d_params.items() if 'mini' in k]
        required_mini = set(list(itertools.chain.from_iterable(params_mini)))

        # store
        dict_bho_ixc[c] = list(required_mini)

    return dict_bho_ixc   # ->list_c








#-----------------------------------------------------------------------------
# FUNCTIONS TO PROCESS MGB BINARY
#-----------------------------------------------------------------------------
def dump_mgb_binary_to_npy(filebin, fileout, nt, nc):
    """ Read binary file (MGB format) and dump content to .npy """
    # read from file
    #'<f4' indicates little-endian (<) float(f) 4 byte (4)
    dados = np.fromfile(filebin,'<f4').reshape(nt,nc)

    # dump fo hard disk
    np.save(fileout,dados)
    return None




def read_npy_as_mmap(filenpy):
    """ Read .npy binary file and make memory-map array"""
    #ms:2021/10/05 -> bad trick to ignore reading
    try:
        # make memory-map from file (doest not consume memory!)
        dados_mmap = np.load(filenpy,mmap_mode='r')
    except:
        dados_mmap = None

    return dados_mmap




def mmap_to_dataframe(dados_mmap, list_t, list_c, dstart):
    """ Read data from memmap as dataframe

    Args:
        dados_mmap (np.memmap) :: memory map of binary .npy

        list_t (list) :: list of integer of selected timesteps

        list_c (list) :: list of integer of selected catchments
                         such as generated by
                         .make_dict_bho_ixc()

        dstart (datetime) :: first date in dados_mmap

    Returns:
        df (pd.DataFrame) :: time-series of selected values
    """

    # adjust loc in array
    ixt_ = [i for i in list_t]
    ixc_ = [int(i-1) for i in list_c]  # mini column [1,nc] -> python [0,nc-1]

    # get selection
    idx = np.ix_(ixt_, ixc_)
    a = dados_mmap[idx]

    # make timeseries dataframe
    times = [dstart + timedelta(days=i) for i in list_t]
    df = pd.DataFrame(a, columns=list_c, index=times)

    return df




#-----------------------------------------------------------------------------
# FUNCTIONS TO READ MGB BINARIES (FULL) AT ONCE
#-----------------------------------------------------------------------------
def read_mgb_binary_as_dataframe(filebin, nt, nc, dstart):
    """ Read full binary (MGB format) as dataframe """

    # read from file
    #'<f4' indicates little-endian (<) float(f) 4 byte (4)
    dados = np.fromfile(filebin,'<f4').reshape(nt,nc)

    # make timeseries dataframe
    times = [dstart + timedelta(days=i) for i in range(nt)]
    df = pd.DataFrame(dados, columns=range(1,nc+1), index=times)
    return df




def read_mgbsa_qtudo(ithot=365, version='1979'):

    # path of file
    path_main = '../'
    path_input = path_main + 'input/'

    # number of catchments
    nc = 33749

    # SELECT VERSION
    # number of intervals and start date MGB-SA (1990->)
    if version == '1990':
        nt = 7305
        dstart = datetime(1990,1,1)
        file_qtudo = path_input + 'QTUDO_1990.MGB'

    # number of intervals and start date MGB-SA (1979->)
    if version == '1979':
        nt = 13149
        dstart = datetime(1979,1,1)
        file_qtudo = path_input + 'QTUDO_1979.MGB'


    # read binary file as dataframe
    df_qtudo = read_mgb_binary_as_dataframe(file_qtudo, nt, nc, dstart)

    # removes initial "hotstart"
    df_qtudo = df_qtudo.iloc[ithot:,:]

    # memory usage
    mem_qtudo = round(df_qtudo.memory_usage(deep=True).sum()/1e6,2)
    print(" Memory usage from QTUDO {} MB".format(mem_qtudo))

    return df_qtudo




def read_mgbsa_qcel(ithot=365, version='1979'):

    # path of file
    path_main = '../'
    path_input = path_main + 'input/'

    # number of catchments
    nc = 33749

    # SELECT VERSION
    # number of intervals and start date MGB-SA (1990->)
    if version == '1990':
        nt = 7305
        dstart = datetime(1990,1,1)
        file_qcel = path_input + 'QITUDO_1990.MGB' #QCEL

    # number of intervals and start date MGB-SA (1979->)
    if version == '1979':
        nt = 13149
        dstart = datetime(1979,1,1)
        file_qcel = path_input + 'QITUDO_1979.MGB'

    # read binary file as dataframe
    df_qcel = read_mgb_binary_as_dataframe(file_qcel, nt, nc, dstart)

    # remove hotstart
    df_qcel = df_qcel.iloc[ithot:,:]

    # memory usage
    mem_qtudo = round(df_qcel.memory_usage(deep=True).sum()/1e6,2)
    print(" Memory usage from QCEL {} MB".format(mem_qtudo))

    return df_qcel




#-----------------------------------------------------------------------------
# DOWNSCALING FUNCTIONS
#-----------------------------------------------------------------------------
def f_downscaling_t1(cotrecho, dict_parameters_t1, df_flow):
    """
    Downscaling function for type 1
        (direct transfer)

    Args:
        cotrecho (int)  :: target cotrecho (key in dict_parameters)
        dict_parameters_t1 :: parameters for type 1 {cotrecho:{params}}
        df_flow (pd.DataFrame) :: dates in rows, minibacias in columns

    Returns:
        outflow (float) :: flow [m3/s]
    """

    # query cotrecho
    param = dict_parameters_t1.get(cotrecho)

    # get parameter values
    mini = param.get('mini')        #list

    # optional
    #area_ratio = param.get('area_ratio') #list -> scalar

    # calculate
    outflow = df_flow[mini].values

    # check output or time series
    outflow = outflow.ravel()
    if outflow.shape[0]>1:
        outflow = list(outflow)
    else:
        outflow = outflow[0]

    return outflow




def f_downscaling_t2(cotrecho, dict_parameters_t2, df_flow):
    """
    Downscaling function for type 2
        (upstream-downstream "mini" water balance)

    Args:
        cotrecho (int)  :: target cotrecho (key in dict_parameters)
        dict_parameters_t2 :: parameters for type 2 {cotrecho:{params}}
        df_flow (pd.DataFrame) :: dates in rows, minibacias in columns

    Returns:
        outflow (float) :: flow [m3/s]
    """

    # query cotrecho
    param = dict_parameters_t2.get(cotrecho)      # dict

    # get parameter values
    minijus = param.get('miniref')               # list
    minimon = param.get('minimon', None)         # list
    fracarea  = param.get('fracarea')[0]         # list -> scalar
    #fracarea  = param['fracarea'][0]             # list -> scalar

    #
    minimonall = param.get('minimonall',None)    # list all

    # optional
    #nuareamont = param.get('nuareamont')[0]     # list -> scalar

    # calculate - old
    #inflows = df_flow[minimon].sum(axis=1) if minimon else 0. #?! mini upstream of bho
    #outflow = df_flow[minijus].sum(axis=1)
    #net = outflow.values - inflows.values
    #outflow_adj = inflows.values + net*fracarea

    # calculate
    inflows = df_flow[minimon].sum(axis=1) if minimon else 0. #?! mini upstream of bho
    outflow = df_flow[minijus].sum(axis=1)
    inflows_all = df_flow[minimonall].sum(axis=1) if minimonall else 0.
    #net = inflows_all.values - outflow.values
    net = outflow.values - inflows_all.values  #fix
    outflow_adj = inflows.values + net*fracarea

    # TODO:
    #if method == 'specific_discharge_from_downstream'
    #    aream_km2 = param.get('aream_km2')[0]  # list -> scalar ->from MGB
    #    outflow = mgb_q95_esp[minijus].sum()*nuareamont

    flag_plot = False
    if flag_plot:
        fig,ax = plt.subplots()
        ax.plot(outflow_adj,label='calc')
        ax.plot(inflows.values,label='in')
        ax.plot(outflow.values,label='out')
        ax.legend()

    # check output or time series
    outflow_adj = outflow_adj.ravel()
    if outflow_adj.shape[0] > 1:
        outflow_adj = list(outflow_adj)
    else:
        outflow_adj = outflow_adj[0]

    return outflow_adj




def f_downscaling_t3(cotrecho, dict_parameters_t3, df_runoff):
    """
    Downscaling function for type 3
        (local runoff [m3/s] in mini catchment area)

    Args:
        cotrecho (int)  :: target cotrecho (key in dict_parameters)
        dict_parameters_t3 :: parameters for type 3 {cotrecho:{params}}
        df_runoff_esp (pd.DataFrame) :: dates in rows, minibacias in columns

    Returns:
        outflow (float) :: flow [m3/s]
    """

    # query cotrecho
    param = dict_parameters_t3.get(cotrecho)      # dict

    # get parameter values
    mini = param.get('mini')                     # list
    area_km2 = param.get('area_km2')[0]        # list -> scalar  #local mgb
    aream_km2 = param.get('aream_km2')[0]        # list -> scalar #montante mgb
    nuareamont = param.get('nuareamont')[0]      # list -> scalar #total bho

    # calculate
    #outflow = df_runoff[mini].values *nuareamont/aream_km2
    outflow = df_runoff[mini].values *nuareamont/area_km2 #ms check

    # check output or time series
    outflow = outflow.ravel()
    if outflow.shape[0] > 1:
        outflow = list(outflow)
    else:
        outflow = outflow[0]

    return outflow




def f_downscaling_t4(cotrecho, dict_parameters_t4, df_flow, inp_qesp=False):
    """
    Downscaling function for type 4
        (specific discharge from downstream type 1  [m3/s.km2])
        (local specific runoff [m3/s.km2])

    Args:
        cotrecho (int)  :: target cotrecho (key in dict_parameters)
        dict_parameters_t4 :: parameters for type 4 {cotrecho:{params}}
        df_flow (pd.DataFrame) :: dates in rows, minibacias in columns
        inp_esp (bool) :: True if df_flow values in m3/s.km2
                         False if df_flow values in m3/s

    Returns:
        outflow (float) :: flow [m3/s]
    """

    # query cotrecho
    param = dict_parameters_t4.get(cotrecho)      # dict

    # get parameter values
    # background ("type 3-like")
    mini = param.get('mini')                     # list
    aream_km2 = param.get('aream_km2')[0]        # list -> scalar
    nuareamont = param.get('nuareamont')[0]      # list -> scalar

    # type 4
    t4_mini = param.get('t4_mini')               # list


    #check if t4 method is available
    if t4_mini:
        t4_aream_km2 = param.get('t4_aream_km2')[0]        # list -> scalar
        t4_nuareamont = param.get('t4_nuareamont')[0]      # list -> scalar
    else:
        #... avoids TypeError or None returns from param.get()
        t4_aream_km2 = np.nan
        t4_nuareamont = np.nan

    # check units of flow [m3/s or m3/s.km2] and factor of conversion
    if inp_qesp:
        fc = 1.
        fcbg = 1.
    else:
        #convert m3/s to m3/s.km2 using mgb area
        fc = 1./t4_aream_km2
        fcbg = 1./aream_km2
        #todo: maybe this is should be fcbg=1./aream_km2


    if t4_mini: #type 4
        outflow = df_flow[t4_mini].values * t4_nuareamont * fc

    else: # "poor man solution" (like type 3)
        outflow = df_flow[mini].values * nuareamont * fcbg


    # check output or time series
    outflow = outflow.ravel()
    if outflow.shape[0]>1:
        outflow = list(outflow)
    else:
        outflow = outflow[0]

    return outflow