Source code for sloop.models.croco.obcfrc

"""
Open Boundary condition Forcing preprocessings
"""

import logging
import os
import numpy as np
from ...log import LoggingLevelContext
import pandas as pd
import glob

with LoggingLevelContext(logging.WARNING):
    import xarray as xr

from xoa.filter import erode_mask
from pathlib import Path

from ...times import convert_to_julian_day
from ...interp import interp_time, Regridder
from ...phys import (
    windstress,
    radiativeflux,
    celsius2kelvin,
    watervapormixingratio,
)
from ...filters import erode_mask_vec
from .__init__ import CROCO_PARAM
from ...io import read_config

LOGGER = logging.getLogger(__name__)


def rename_vars(ds, mod):
    """Rename variables as expected by Extract"""
    LOGGER.info("Rename obc variables as expected by Extract program")
    obc_names = mod["varnames"]
    list_nc_vars = list(ds.variables.keys())
    list_nc_dims = list(ds.dims.keys())
    for variable in list_nc_vars:
        for key, value in obc_names.items():
            if variable == value and value != key:
                ds = ds.rename_vars({variable: key})
                if variable in list_nc_dims:
                    print("variable also a dim")
                    ds = ds.rename_dims({variable: key})
    # remove unused variables
    for dsi in list(ds.variables.keys()):
        if dsi not in obc_names:
            ds = ds.drop(dsi)
    return ds


def format_obc_to_extract(ds, model):
    """Format OBC file for the Extract program"""
    LOGGER.info("Format to Extract")
    # get the model params
    mod = read_config(model)
    # rename vars
    ds = rename_vars(ds, mod)
    # change temperature units if necessary
    if mod["temp_in_kelvin"]:
        ds.variables["TEMP"] = ds.variables["TEMP"] - 273.15
    # set time in seconds
    ds.time.encoding.update({"units": "seconds since 1900-01-01", "dtype": float})
    # unpack data
    LOGGER.info("unpack and set valid min & max")
    for var in list(ds.variables.keys()):
        if var not in list(ds.dims.keys()):
            ds[var].encoding.update(
                {
                    "_FillValue": -999.0,
                    "scale_factor": 1.0,
                    "add_offset": 0.0,
                    "dtype": float,
                }
            )
            # set_validminmax(ds[var])
    return ds


def reset_validminmax(fname):
    """reset valid min and max for a dataset"""
    LOGGER.info("reset valid min & max")
    ds = xr.open_mfdataset(fname)
    ds2 = ds.copy()
    ds.close()
    for var in list(ds2.variables.keys()):
        if var not in list(ds2.dims.keys()):
            set_validminmax(ds2[var])
    os.remove(fname)
    ds2.to_netcdf(fname, mode="w")


def set_validminmax(data):
    """set valid min & max of data variable"""
    vmin = float(data.min().to_numpy())
    vmax = float(data.max().to_numpy())
    data.attrs.update(valid_min=vmin, valid_max=vmax)
    return data


def get_dates_from_file(fname, tname="time"):
    """get the start and end dates from a file"""
    print("open", fname)
    ds = xr.open_mfdataset(fname)
    time = ds.indexes[tname]
    d1 = time[0].strftime("%Y%m%dT%H%M")
    d2 = time[-1].strftime("%Y%m%dT%H%M")
    return d1, d2


def build_obc_list(date1, date2, model):
    """Build obc list of files from OGCM model"""
    LOGGER.info("Get list of obc files")
    # retrieve params of meteo model
    mod = read_config(model)
    rac = mod["prefix"]
    data_dir = mod["data_dir"]
    suff = mod["suff"]
    # build list of date
    dates_list = pd.date_range(start=date1, end=date2, freq=mod["freq"])
    # build list of names
    inst_list = [f"{rac}{date.strftime('%Y%m%d')}" for date in dates_list]

    # finally build list of files
    # inst_list = [glob.glob(f"{data_dir}/*/*/{file}*.nc")[0] for file in inst_list]
    inst_list = [
        glob.glob(f"{data_dir}/**/{file}*.nc", recursive=True)[0] for file in inst_list
    ]

    return inst_list


def read_obc_data(inst_list, lon_bnds, lat_bnds):
    """Read the list of obc files and reduce it to the area of interest"""
    from functools import partial

    def _preprocess(x, lon_bnds, lat_bnds):
        return x.sel(longitude=slice(*lon_bnds), latitude=slice(*lat_bnds))

    partial_func = partial(_preprocess, lon_bnds=lon_bnds, lat_bnds=lat_bnds)
    ds_inst = xr.open_mfdataset(
        inst_list, concat_dim="time", combine="nested", preprocess=partial_func
    )
    return ds_inst