from __future__ import annotations
from logging import debug, info, warning
from os import PathLike
from pathlib import Path
from typing import Literal
import numpy as np
import scipy.sparse as sparse
import xarray as xr
from numpy.typing import NDArray
from scipy.sparse.linalg import eigsh
from .....utils.check.errclass import (
CifFileExistsError,
CifFileNotFoundError,
CifValueError,
)
from .....utils.geometry.dist_matrix import dist_matrix
from .....utils.hdf5 import _hdf5_lock
try:
from sklearn.neighbors import BallTree
except ImportError:
BallTree = None
EARTH_RADIUS = 6371.03
[docs]
def build_hcorrelations(
load_dir: str | PathLike[str],
dump_dir: str | PathLike[str],
hresol: str,
hresoldim: int,
zlat: NDArray[np.floating],
zlon: NDArray[np.floating],
sigma_sea: int,
sigma_land: int,
is_lbc: bool,
landseamask: NDArray[np.floating] | None = None,
projection: Literal["gps", "xy"] = "gps",
evalues_cutoff: float = 0.0,
crop_chi: bool = False,
use_sparse: bool = False,
n_modes: int | None = None,
sparse_cutoff: float = 0.1,
target_prec: float | None = None,
dump: bool = False,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
"""Reads from cache file or compute the horizontal correlation matrix eigen values
and vectors.
Parameters
----------
load_dir : str or path-like
Directory to load precomputed eigen values and vectors from
dump_dir : str or path-like
Directory to dump computed eigen values and vectors in if 'dump' is True
hresol : str
Horizontal resolution type
hresoldim : int
Horizontal grid size
zlat : 2D array
Grid latitudes
zlon : 2D array
Grid longitudes
sigma_sea : int
Decay distance in km for correlation between sea cells
sigma_land : int
Decay distance in km for correlation between land cells
is_lbc : bool
Is lateral boundary condition
landseamask : 2D array, optional
Land-sea mask, pixels are 1 for land and 0 for sea, by default None
projection : 'gps' or 'xy', optional
Projection used to compute distances, by default "gps"
evalues_cutoff : float, optional
Truncates eigenvalues below this value, by default 0.0
crop_chi : bool, optional
Truncated eigen values and vector are cropped is True, otherwise they are
padded with zeros, by default False
use_sparse : bool, optional
Use sparse matrix with 'eigsh', by default False
n_modes : int | None, optional
Number of eigenvalues to compute for sparse matrix, by default None
sparse_cutoff : float, optional
Correlation value threshold for sparse arrays, correlation below this value are
set to zero, by default 0.1
target_prec : float, optional
Target precision for sparse matrix, by default 1e-2
dump : bool, optional
Dumps computed eigen values and vector if True, by default False
Returns
-------
1D array, 2D array
square roots of eigenvalues and eigenvectors
"""
nlat, nlon = zlat.shape
evalues, evectors = read_hcorr(
hcorr_dir=load_dir,
hresol=hresol,
hresoldim=hresoldim,
nlon=nlon,
nlat=nlat,
sigma_sea=sigma_sea,
sigma_land=sigma_land,
k=n_modes,
is_lbc=is_lbc,
missing_ok=True,
)
if evalues is None:
info(f"Computing hcorr for {nlat}/{nlon} domain")
# No correlation between land and sea if lsm = True
if landseamask is not None:
land_grid = landseamask.flatten() >= 0.5
sea_grid = landseamask.flatten() < 0.5
sigma = (
sigma_land * land_grid[:, np.newaxis] * land_grid[np.newaxis, :]
+ sigma_sea * sea_grid[:, np.newaxis] * sea_grid[np.newaxis, :]
)
if np.ma.isMaskedArray(sigma):
sigma = sigma.data
# Otherwise, isotropic correlation, takes sigma_land
else:
sigma = sigma_land # float
# Compute the eigen values and vectores
if use_sparse:
evalues, evectors = compute_hcorrel_sparse(
zlat, zlon, sigma, projection, n_modes, sparse_cutoff, target_prec
)
else:
evalues, evectors = compute_hcorrel_dense(zlat, zlon, sigma, projection)
# Re-ordering values
order = np.argsort(evalues)[::-1]
evalues = evalues[order]
evectors = evectors[:, order]
# Dumping to a binary file
if dump:
dump_hcorr(
evalues,
evectors,
hcorr_dir=dump_dir,
hresol=hresol,
hresoldim=hresoldim,
nlon=nlon,
nlat=nlat,
sigma_sea=sigma_sea,
sigma_land=sigma_land,
k=n_modes,
is_lbc=is_lbc,
)
# Truncating values < evalmin
mask = evalues >= evalues_cutoff
if crop_chi:
return np.sqrt(evalues[mask]), evectors[:, mask]
else:
evalues[~mask] = 0
return np.sqrt(evalues), evectors
[docs]
def compute_hcorrel_matrix(
zlat: NDArray[np.floating],
zlon: NDArray[np.floating],
sigma: int | NDArray[np.integer],
projection: Literal["gps", "xy"] = "gps",
) -> NDArray[np.floating]:
# Compute the correlation matrix
corr = dist_matrix(zlat, zlon, projection) # distance matrix -> corr = dx
np.divide(corr, sigma, out=corr) # divides by sigma -> corr = dx/sigma
np.negative(corr, out=corr) # takes the opposite -> corr = -dx/sigma
np.exp(corr, out=corr) # applies exponential -> corr = exp(-dx/sigma)
# Land sea mask
if isinstance(sigma, np.ndarray):
corr[sigma == 0.0] = 0.0
return corr
[docs]
def compute_hcorrel_dense(
zlat: NDArray[np.floating],
zlon: NDArray[np.floating],
sigma: int | NDArray[np.integer],
projection: Literal["gps", "xy"] = "gps",
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
corr = compute_hcorrel_matrix(zlat, zlon, sigma, projection)
evalues, evectors = np.linalg.eigh(corr)
return evalues, evectors
[docs]
def compute_hcorrel_sparse(
zlat: NDArray[np.floating],
zlon: NDArray[np.floating],
sigma: int | NDArray[np.integer],
projection: Literal["gps", "xy"] = "gps",
n_modes: int | None = None,
sparse_cutoff: float = 0.1,
target_prec: float | None = None,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
N = zlat.size
use_tree = True
if not 0 < sparse_cutoff < 1:
raise CifValueError(
f"sparse_cutoff must be between 0 and 1, got {sparse_cutoff}"
)
if n_modes is None:
warning(
"WARNING: Number of eigen modes not specified, using highly unoptimized "
+ "default of 20% of the grid size"
)
k = zlon.size // 5
else:
k = n_modes
d_cut = -sigma * np.log(sparse_cutoff)
debug(f"Cutoff distance: {d_cut:.2f} km")
if BallTree is None:
use_tree = False
reason = "BallTree not available, please install sklearn to use it"
if projection != "gps":
use_tree = False
reason = f"projection is different from 'gps', got {projection!r}"
if isinstance(sigma, np.ndarray):
use_tree = False
reason = f"the use of land-sea mask is not compatible"
if use_tree:
# Convert to radians
coords_rad = np.radians(np.column_stack([zlat.ravel(), zlon.ravel()]))
sigma_rad = sigma / EARTH_RADIUS
d_cut_rad = d_cut / EARTH_RADIUS
debug(f"Computing sparse distance tree for domain with size {N}")
tree = BallTree(coords_rad, metric="haversine")
nghbr_idx, nghbr_dist_rad = tree.query_radius(
coords_rad,
r=d_cut_rad,
return_distance=True,
)
# Computing sparse data
dist_rad = np.concatenate(nghbr_dist_rad)
data = np.exp(-dist_rad / sigma_rad).astype(np.float64)
# Constructing sparse indices
lengths = np.fromiter((len(idx) for idx in nghbr_idx), dtype=np.int64, count=N)
rows = np.repeat(np.arange(N, dtype=np.int32), lengths)
cols = np.concatenate(nghbr_idx).astype(np.int32)
# Constructing sparse matrix
corr_sparse = sparse.coo_matrix((data, (rows, cols)), shape=(N, N)).tocsr()
else:
warning(
"WARNING: Falling back to unoptimized distance matrix computation "
+ f"because {reason}"
)
corr = compute_hcorrel_matrix(zlat, zlon, sigma, projection)
corr[corr <= sparse_cutoff] = 0.0
corr_sparse = sparse.coo_matrix(corr).tocsr()
corr_sparse = corr_sparse.maximum(corr_sparse.T)
debug(f"Sparse matrix density={corr_sparse.nnz / N**2:.4%}")
evalues, evectors = eigsh(
corr_sparse,
k=k,
which="LA",
ncv=min(N, max(2 * k + 1, 20)),
)
if target_prec is not None:
# Computing error on a single matrix row
i = N // 2
corr = evectors @ (evalues * evectors[i, :])
corr_ref = corr[corr > sparse_cutoff]
err = np.sqrt(np.mean((corr - corr_ref) ** 2))
if err > target_prec:
raise CifValueError(
f"Target presision on {target_prec:.3g} not reached, got {err:.3g}"
)
return evalues, evectors
[docs]
def get_hcorr_file(
hcorr_dir: str | PathLike[str],
hresol: str,
hresoldim: int,
nlon: int,
nlat: int,
sigma_sea: int,
sigma_land: int,
k: int | None = None,
is_lbc: bool = False,
legacy_path: bool = False,
) -> Path:
"""Returns the path to the horizontal correlations eigenvalues and eigenvectors
cache file corresdonding to the given parameters.
If 'legacy_path' is True return the legacy path is the file exists to ensure
compatibility with file generated with older versions
"""
if legacy_path:
# Legacy path to Numpy binary file
name = f"horcor_{hresol}_{hresoldim}_{nlon}x{nlat}_cs{sigma_sea}_cl{sigma_land}"
if is_lbc:
name += "_lbc"
path = Path(hcorr_dir, name + ".bin")
if path.exists():
return path
# Path to NetCDF file
name = f"horcor_{hresol}_{hresoldim}_{nlon}x{nlat}"
if sigma_sea == -999:
name += f"_c{sigma_land}km"
else:
name += f"_cs{sigma_sea}km_cl{sigma_land}km"
if k is not None:
name += f"_k{k}"
if is_lbc:
name += "_lbc"
return Path(hcorr_dir, name + ".nc")
[docs]
def dump_hcorr(
evalues: NDArray[np.floating],
evectors: NDArray[np.floating],
hcorr_dir: str | PathLike[str],
hresol: str,
hresoldim: int,
nlon: int,
nlat: int,
sigma_sea: int,
sigma_land: int,
k: int | None = None,
is_lbc: bool = False,
overwrite: bool = False,
) -> None:
"""Dumps the horizontal correlations eigenvalues and eigenvectors to a cache file
corresdonding to the given parameters
"""
path = get_hcorr_file(
hcorr_dir, hresol, hresoldim, nlon, nlat, sigma_sea, sigma_land, k, is_lbc
)
if path.exists() and not overwrite:
raise CifFileExistsError(f"'{path}' already exists")
# Get a different name for eigen-dim if eigenvalues are truncated
kdim = "n" if k is None else "k"
ds = xr.Dataset(
{
"evalues": ([kdim], evalues),
"evectors": (["n", kdim], evectors),
},
attrs={
"hresol": hresol,
"hresoldim": hresoldim,
"nlon": nlon,
"nlat": nlat,
"is_lbc": "true" if is_lbc else "false",
},
)
if sigma_sea == -999:
ds.attrs["sigma_sea"] = sigma_sea
ds.attrs["sigma_land"] = sigma_land
else:
ds.attrs["sigma"] = sigma_land
# Add a reconstructed correlation matrix row if eigenvalues are truncated
# It can be use to inspect residuals
if k is not None:
i = hresoldim // 2
corr = evectors @ (evalues * evectors[i, :])
ds["corr"] = (["lat", "lon"], corr.reshape((nlat, nlon)))
ds["corr"].attrs = {
"long_name": f"reconstructed correlation matrix row {i}",
}
# Dumping the NetCDF file
debug(f"Dumping horizontal correlations to '{path}'")
path.parent.mkdir(parents=True, exist_ok=True)
with _hdf5_lock:
ds.to_netcdf(path)
[docs]
def read_hcorr(
hcorr_dir: str | PathLike[str],
hresol: str,
hresoldim: int,
nlon: int,
nlat: int,
sigma_sea: int,
sigma_land: int,
k: int | None = None,
is_lbc: bool = False,
missing_ok: bool = False,
) -> tuple[NDArray[np.floating], NDArray[np.floating]] | tuple[None, None]:
"""Reads the horizontal correlations eigenvalues and eigenvectors form the cache file
corresdonding to the given parameters.
"""
path = get_hcorr_file(
hcorr_dir, hresol, hresoldim, nlon, nlat, sigma_sea, sigma_land, k, is_lbc, True
)
if not path.is_file():
if not missing_ok:
raise CifFileNotFoundError(f"'{path}' does not exists")
return None, None
debug(f"Reading horizontal correlations from '{path}'")
if path.suffix == ".bin":
# Legacy Numpy binary file
data = np.fromfile(path).reshape((-1, nlon * nlat))
evalues = data[0]
evectors = data[1:]
else:
# NetCDF file
with _hdf5_lock:
with xr.open_dataset(path) as ds:
evalues = ds.evalues.values
evectors = ds.evectors.values
evalues[evalues < 0] = 0.0
return evalues, evectors
