import os
from logging import debug, info
import numpy as np
import scipy.sparse as sparse
from .....utils.geometry.dist_matrix import dist_matrix
from .....utils.path import init_dir
from .....utils.check.errclass import CifIOError
[docs]
def build_hcorrelations(
hresol,
hresoldim,
zlat,
zlon,
lsm,
landseamask,
sigma_land,
sigma_sea,
evalmin=0.,
dump=False,
dir_dump="",
projection="gps",
is_lbc=False,
crop_chi=False,
tracer=None,
glob_err=None,
use_sparse=False,
sparse_crop_threshold=0.1,
**kwargs
):
"""Build horizontal correlation matrix based on distance between grid
cells.
For cells i and j, the corresponding correlation is:
c(i,j) = exp(-dist(i, j) / sigma)
sigma depends on the land-sea mask: land and sea cells are assumed
un-correlated
Args:
zlat (np.array): 2D array of latitudes
zlon (np.array): 2D array of longitudes
file_lsm (str): path to NetCDF file with land-sea mask (grid must be
consistent with LMDZ grid); the land-sea mask is assumed to be stored
in the varible 'lsm'
sigma_land (float): decay distance for correlation between land cells
sigma_sea (float): idem for sea
evalmin (float): flag out all eigenvalues below this value. Default
is 0.5
dump (bool): dumps computed correlations if True
dir_dump (str): directory where correlation matrices are stored
projection (str): the projection used for the longitudes and latitudes
is_lbc (bool): True if boundary of a 2D domain
use_sparse (bool): True if convert correlation matrix to a sparse matrix
sparse_crop_threshold (float): Threshold on correlation to exclude values from
the sparse array
Return:
tuple with:
- square roots of eigenvalues
- eigenvectors
"""
# Define domain dimensions
nlat, nlon = zlat.shape
hresoldim = tracer.hresoldim
# Try reading existing file
try:
evalues, evectors = read_hcorr(
hresol, hresoldim, nlon, nlat, sigma_sea, sigma_land, dir_dump, is_lbc
)
# Else build correlations
except IOError as e:
# Dumping exception causing to recompute
debug(e.__str__())
info(f"Computing hcorr for {nlat}/{nlon} domain")
# No correlation between land and sea if lsm = True
if lsm:
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, :]
)
# Unmask arrays
if np.ma.isMaskedArray(sigma):
sigma = sigma.data
# Otherwise, isotropic correlation
# Takes sigma_land
else:
sigma = sigma_land
# Compute matrix of distance
dx = dist_matrix(zlat, zlon, projection)
# Compute the correlation matrix itself
corr = np.exp(
- np.divide(
dx, sigma,
out=np.full_like(dx, np.inf),
where=sigma != 0)
).astype(np.float64)
# Component analysis
debug("Computing the eigen decomposition of the correlation matrix")
if use_sparse:
# Use sparse matrix
mask_corr = np.abs(corr) > 0.1
masked_corr = np.ma.masked_array(data=corr, mask=mask_corr)
sparse_corr = sparse.csgraph.csgraph_from_masked(masked_corr)
debug(f"Using sparse matrix of size {mask_corr.sum()}")
evalues, evectors = sparse.linalg.eigsh(
sparse_corr, k=zlon.size - 2)
# Fill missing eigen values
evalues = np.concatenate([evalues, np.zeros(2)])
evectors = np.concatenate(
[evectors, np.zeros((zlon.size, 2))], axis=1)
else:
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_hcorr(
hresol, hresoldim, nlon, nlat, sigma_sea, sigma_land, 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_hcorr(hresol, hresoldim, nlon, nlat, sigma_sea, sigma_land, evalues, evectors,
dir_dump, is_lbc=False, overwrite=False):
"""Dumps eigenvalues and vectors to a binary file. The default file format
is:
f"{dir_dump}/horcor_{hresol}_{hresoldim}_{nlon}x{nlat}_cs{sigma_sea}_cl{sigma_land}{'_lbc' if is_lbc else ''}.bin"
"""
ncell = evalues.size
file_dump = f"{dir_dump}/horcor_{hresol}_{hresoldim}_{nlon}x{nlat}_cs{sigma_sea}_cl{sigma_land}{'_lbc' if is_lbc else ''}.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_hcorr(hresol, hresoldim, nlon, nlat, sigma_sea, sigma_land, dir_dump,
is_lbc=False):
"""Reads horizontal correlations from existing text file
Args:
nlon, nlat (ints): dimensions of the domain
sigma_land, sigma_sea (floats): horizontal correlation distances
dir_dump (str): where the horizontal correlations have been stored
"""
file_dump = f"{dir_dump}/horcor_{hresol}_{hresoldim}_{nlon}x{nlat}_cs{sigma_sea}_cl{sigma_land}{'_lbc' if is_lbc else ''}.bin"
if not os.path.isfile(file_dump):
raise CifIOError(
f"{file_dump} does not exist. Please compute correlations from scratch"
)
debug(f"Reading horizontal correlations from {file_dump}")
data = np.fromfile(file_dump).reshape((-1, nlon * nlat))
evalues = data[0]
evectors = data[1:]
evalues[evalues < 0] = 0.0
return evalues, evectors
