from logging import warning
from pathlib import Path
import pandas as pd
from .....utils.check.errclass import CifAttributeError, CifError, IllegalArgumentError
from .build_hcorr import build_hcorrelations
from .build_lsm import build_lsm
from .build_tcorr import build_tcorrelations
from .build_vcorr import build_vcorrelations
[docs]
def get_hcorr(cntrlv, tracer):
hresol = tracer.hresol
hresoldim = tracer.hresoldim
corr = tracer.hcorrelations
# Get options for horizontal correlations
# - grid options
grid = getattr(tracer, "domain", cntrlv.domain)
projection = getattr(grid, "projection", "gps")
is_lbc = getattr(tracer, "is_lbc", False)
lsm = getattr(corr, "landsea", False)
# - dump options
dump_dir = Path(tracer.workdir, "controlvect", "correlations")
dump_hcorr = getattr(corr, "dump_hcorr", False)
dircorrel = getattr(corr, "dircorrel", dump_dir)
# - sparse option
use_sparse = getattr(corr, "use_sparse", False)
sparse_n_modes = getattr(corr, "sparse_n_modes", None)
sparse_crop_threshold = getattr(corr, "sparse_crop_threshold", 0.1)
sparse_target_precision = getattr(corr, "sparse_target_precision", None)
# - eigenvalues options
if not use_sparse:
evalmin = getattr(corr, "evalmin", 0.0)
crop_chi = getattr(corr, "crop_chi", False)
else:
if hasattr(corr, "crop_chi") and not getattr(corr, "crop_chi"):
raise IllegalArgumentError(
"crop_chi can not be explicitly set to False whith use_sparse=True"
)
crop_chi = True
if hasattr(corr, "evalmin"):
raise IllegalArgumentError(
"evalmin can not be explicitly set whith use_sparse=True"
)
evalmin = sparse_crop_threshold
# Raise an error if any sparse option if used without use_sparse=True
sparse_options = [
"sparse_n_modes",
"sparse_crop_threshold",
"sparse_target_precision",
]
used_opts = [opt for opt in sparse_options if hasattr(corr, opt)]
if not use_sparse and used_opts:
raise IllegalArgumentError(
"Can not use options " + ", ".join(used_opts) + " without use_sparse=True"
)
# Two possible options: - uniform correlation length,
# or separated land and sea, along a land-sea mask
# Default is no separation
if lsm:
landseamask, sigma_land, sigma_sea = build_lsm(corr, hresol, tracer)
else:
if not hasattr(corr, "sigma"):
raise CifAttributeError(
"The attribute `sigma` is missing to define horizontal correlations with no landsea mask.\n"
"Please include it in the relevant paragraph in `datavect` of your yaml"
)
landseamask = None
sigma = getattr(corr, "sigma", -1)
sigma_land = sigma
sigma_sea = -999
# Load or compute the horizontal correlations
# Checks before whether they were already loaded
# Use pre-computed correlations for hpixels only,
# otherwise, re-compute correlations
key = (sigma_land, sigma_sea, is_lbc, hresol, hresoldim)
if key in cntrlv.hcorrelations and hresol in ["hpixels", "regions"]:
sqrt_evalues = cntrlv.hcorrelations[key]["sqrt_evalues"]
evectors = cntrlv.hcorrelations[key]["evectors"]
else:
# Take sides if LBC parameter
zlat = grid.zlat
zlon = grid.zlon
if is_lbc:
zlat = grid.zlat_side
zlon = grid.zlon_side
# Take centroids if not hpixels
if hresol != "hpixels":
zlon = tracer.centroids_lons
zlat = tracer.centroids_lats
sqrt_evalues, evectors = build_hcorrelations(
load_dir=dircorrel,
dump_dir=dump_dir,
hresol=hresol,
hresoldim=hresoldim,
zlat=zlat,
zlon=zlon,
sigma_sea=sigma_sea,
sigma_land=sigma_land,
is_lbc=is_lbc,
landseamask=landseamask,
projection=projection,
evalues_cutoff=evalmin,
crop_chi=crop_chi,
use_sparse=use_sparse,
n_modes=sparse_n_modes,
sparse_cutoff=sparse_crop_threshold,
target_prec=sparse_target_precision,
dump=dump_hcorr,
)
# Storing computed correlations for use by other components
cntrlv.hcorrelations[key] = {
"evectors": evectors,
"sqrt_evalues": sqrt_evalues,
}
corr.sqrt_evalues = sqrt_evalues
corr.evectors = evectors
return sqrt_evalues, evectors
[docs]
def get_tcorr(cntrlv, tracer, corr):
tresol = getattr(tracer, "tresol")
tsubresol = getattr(tracer, "tsubresol", "0")
dates = getattr(tracer, "dates")
# Check that sigma_t is really a frequency
sigma_t = getattr(corr, "sigma_t", -1)
if isinstance(sigma_t, str):
if sigma_t[:-2] == "MS" or sigma_t[:-1] == "M":
warning(
f"sigma_t={sigma_t} can be ambiguous, 'MS' is "
"interpreted as 'milliseconds' and 'M' as 'minutes', "
"and NEVER as 'months'"
)
if pd.to_timedelta(sigma_t).total_seconds() == 0:
raise CifError(
f"The specified sigma_t ({sigma_t}) is not a correct pandas frequency.\nPlease check your Yaml file!"
)
# Load other parameters
sigma_type = getattr(corr, "type")
dump_tcorr = getattr(corr, "dump_tcorr", False)
dircorrel = getattr(
corr, "dircorrel", f"{cntrlv.workdir}/controlvect/correlations/"
)
evalmin = getattr(corr, "evalmin", 0.0)
crop_chi = getattr(corr, "crop_chi", False)
# Load or compute the temporal correlations
# Checks before whether they were already loaded
if (sigma_t, tresol, tsubresol) in cntrlv.tcorrelations:
sqrt_evalues = cntrlv.tcorrelations[(sigma_t, tresol, tsubresol, sigma_type)][
"sqrt_evalues"
]
evectors = cntrlv.tcorrelations[(sigma_t, tresol, tsubresol, sigma_type)][
"evectors"
]
else:
# Computes correlations for dates[:-1]
# as the last dates indicates the end of the period
sqrt_evalues, evectors = build_tcorrelations(
tresol,
tsubresol,
dates,
sigma_t,
sigma_type,
evalmin=evalmin,
dump=dump_tcorr,
dir_dump=dircorrel,
crop_chi=crop_chi,
corr_plg=corr,
tracer=tracer,
)
# Storing computed correlations for use by other components
cntrlv.tcorrelations[(sigma_t, tresol, tsubresol, sigma_type)] = {
"evectors": evectors,
"sqrt_evalues": sqrt_evalues,
}
corr.sqrt_evalues = sqrt_evalues
corr.evectors = evectors
return sqrt_evalues, evectors
[docs]
def get_vcorr(cntrlv, tracer):
vresol = getattr(tracer, "vresol")
nlev = tracer.nlev
grid = getattr(tracer, "domain", cntrlv.domain)
corr = tracer.vcorrelations
dump_vcorr = getattr(corr, "dump_vcorr", False)
dircorrel = getattr(
corr, "dircorrel", f"{cntrlv.workdir}/controlvect/correlations/"
)
evalmin = getattr(corr, "evalmin", 0.0)
crop_chi = getattr(corr, "crop_chi", False)
# Check method
method = corr.method
if method not in ["level", "pressure"]:
raise CifError(f"Unrecognized method {method} for vertical correlations")
sigma_lev = corr.sigma
# Load or compute the temporal correlations
# Checks before whether they were already loaded
if (sigma_lev, vresol) in cntrlv.vcorrelations:
sqrt_evalues = cntrlv.vcorrelations[(sigma_lev, vresol, method)]["sqrt_evalues"]
evectors = cntrlv.vcorrelations[(sigma_lev, vresol, method)]["evectors"]
else:
# Computes correlations for dates[:-1]
# as the last dates indicates the end of the period
sqrt_evalues, evectors = build_vcorrelations(
vresol,
nlev,
sigma_lev,
method,
corr_plg=corr,
tracer=tracer,
dump=dump_vcorr,
dir_dump=dircorrel,
evalmin=evalmin,
crop_chi=crop_chi,
)
# Storing computed correlations for use by other components
cntrlv.tcorrelations[(sigma_lev, vresol, method)] = {
"evectors": evectors,
"sqrt_evalues": sqrt_evalues,
}
corr.sqrt_evalues = sqrt_evalues
corr.evectors = evectors
return sqrt_evalues, evectors
