Source code for sloop.models.croco.io

"""
In/output routines

"""

import re
import math
import logging
import os
import shutil

import numpy as np
from ...log import LoggingLevelContext

with LoggingLevelContext(logging.WARNING):
    import xarray as xr
    import pandas as pd
import scipy.io as sio

from .__init__ import (
    get_cf_specs,
    SloopCrocoError,
    CROCO_SIGMA_TO_STMT_FNS,
    CROCO_CONTEXT_TO_STMT_FNS,
    CROCO_IBC_GEN_INICON,
    CROCO_IBC_GEN_INICON_VCONF,
    CROCO_MODEL_DEFSTRETCH_FILES,
)
from sloop.times import convert_from_julian_day, convert_to_julian_day


def read_croco_grid(croco_grid_file):
    """
    Read and return geographical grid limits
    """
    ds = xr.open_dataset(croco_grid_file)
    lonmin, lonmax = ds.lon_rho.min().to_numpy(), ds.lon_rho.max().to_numpy()
    latmin, latmax = ds.lat_rho.min().to_numpy(), ds.lat_rho.max().to_numpy()
    dlon = np.diff(ds.lon_rho.data, axis=1).mean()
    dlat = np.diff(ds.lat_rho.data, axis=0).mean()
    return lonmin, lonmax, latmin, latmax, dlon, dlat


def get_restart_date(ncname):
    """
    get the date of a restart file in Croco format
    """
    ds = xr.open_dataset(ncname, decode_times=False)
    units, reference_date = ds.time.attrs["units"].split("since")
    time = pd.to_datetime(
        ds.scrum_time.values, unit="s", origin=pd.Timestamp(reference_date.strip())
    )
    return time[0].strftime("%Y%m%dT%H%M")


def format_ds(ds):
    """Format variable in a dataset or dataarray to be hycom compliant"""
    # Time units
    time = get_cf_specs().coords.get(ds, "time")
    if time is not None:
        ds[time.name].encoding["units"] = "days since 1950-01-01"

    # Other encodings
    for vname, var in ds.items():
        for attr in "dtype", "scale_factor", "add_offset":
            if attr in ds[vname].encoding:
                del ds[vname].encoding[attr]
        if ds[vname].dtype.kind in "fd":
            ds[vname].encoding["_FillValue"] = -1e34

    return ds