import os
from logging import debug, info
import numpy as np
from .....utils.geometry.dist_matrix import dist_matrix
from .....utils.netcdf import readnc
from .....utils.path import init_dir
from ....transforms.system.fromcontrol.utils.scalemaps import map2scale
from .....utils.check.errclass import CifError, CifIOError
[docs]
def build_vcorrelations(
vresol,
nlev,
sigma_lev,
method,
corr_plg=None,
tracer=None,
dump=False,
dir_dump="",
evalmin=0,
crop_chi=False,
**kwargs
):
"""Build vertical correlation matrix based on distance between levels
Return:
tuple with:
- square roots of eigenvalues
- eigenvectors
"""
# Try reading existing file
try:
evalues, evectors = read_vcorr(
vresol, nlev, sigma_lev, method, dir_dump
)
# Else build correlations
except IOError:
info(f"Computing vcorr for {nlev} levels and {method} method")
if method == "level":
dx = np.abs(
np.arange(nlev)[:, np.newaxis] - np.arange(nlev)[np.newaxis]
).astype(float)
else:
raise CifError("Not implemented yet")
# Compute the correlation matrix itself
corr = np.exp(
- np.divide(dx, sigma_lev,
out=np.full_like(dx, np.inf))
).astype(np.float64)
# Component analysis
evalues, evectors = np.linalg.eigh(corr)
# Re-ordering values
# (not necessary in principle in recent numpy versions)
index = np.argsort(evalues)[::-1]
evalues = evalues[index]
evectors = evectors[:, index]
# Dumping to a binary file
if dump:
dump_vcorr(
vresol, nlev, sigma_lev, method, evalues, evectors,
f"{tracer.workdir}/controlvect/correlations/"
)
except Exception as e:
raise e
# Truncating values < evalmin
mask = evalues >= evalmin
if crop_chi:
return evalues[mask] ** 0.5, evectors[:, mask]
else:
evalues[~mask] = 0
return evalues ** 0.5, evectors
[docs]
def dump_vcorr(vresol, nlev, sigma_lev, method, evalues, evectors,
dir_dump, overwrite=False):
"""Dumps eigenvalues and vectors to a binary file. The default file format
is:
f"{dir_dump}/vertcor_{vresol}_{nlev}_{method}.bin"
)
"""
file_dump = f"{dir_dump}/vertcor_{vresol}_{nlev}_{method}.bin"
if os.path.isfile(file_dump) and overwrite:
raise CifIOError(
f"Warning: {file_dump} already exists. I don't want to overwrite it"
)
datasave = np.concatenate((evalues[np.newaxis, :], evectors), axis=0)
# Creating path if does not exist
if not os.path.isdir(os.path.dirname(file_dump)):
init_dir(os.path.dirname(file_dump))
# Saving data
np.array(datasave).tofile(file_dump)
[docs]
def read_vcorr(vresol, nlev, sigma_lev, method, dir_dump):
"""Reads vertical correlations from existing text file
"""
file_dump = f"{dir_dump}/vertcor_{vresol}_{nlev}_{method}.bin"
if not os.path.isfile(file_dump):
raise CifIOError(
f"{file_dump} does not exist. Please compute correlations from scratch"
)
debug(f"Reading vertical correlations from {file_dump}")
data = np.fromfile(file_dump).reshape((-1, nlev))
evalues = data[0]
evectors = data[1:]
evalues[evalues < 0] = 0.0
return evalues, evectors
