from __future__ import annotations
import shutil
from logging import debug
from os import PathLike
from pathlib import Path
from typing import Any
import numpy as np
import xarray as xr
from ......utils.check.errclass import CifValueError
from ......utils.hdf5 import _hdf5_lock
from .bilinear import bilinear
from .conservative import conservative
from .find_gridcells import find_gridcells
from .reproject import reproject_emissions
[docs]
def get_weights(
transform,
trid,
mapper,
domain_in,
domain_out,
is_lbc,
ddi,
is_sparse_in=False,
is_sparse_out=False,
**kwargs,
):
# If sparse data in inputs, switch domain_in and domain_out
if is_sparse_in:
dummy_domain = domain_out
domain_out = domain_in
domain_in = dummy_domain
is_sparse_out = True
is_sparse_in = False
# Fetching infos from original domain
nlon_in = domain_in.nlon
zlon_in = domain_in.zlon
zlonc_in = getattr(domain_in, "zlonc", zlon_in)
nlat_in = domain_in.nlat
zlat_in = domain_in.zlat
zlatc_in = getattr(domain_in, "zlatc", zlat_in)
# Dealing with boundary output domain
if is_lbc:
zlon_out = domain_out.zlon_side
zlonc_out = getattr(domain_out, "zlonc_side", zlon_out)
zlat_out = domain_out.zlat_side
zlatc_out = getattr(domain_out, "zlatc_side", zlat_out)
else:
zlon_out = domain_out.zlon
zlonc_out = getattr(domain_out, "zlonc", zlon_out)
zlat_out = domain_out.zlat
zlatc_out = getattr(domain_out, "zlatc", zlat_out)
# Find correct weight file
copy_weights(
transform,
mapper,
domain_in,
domain_out,
ddi,
is_sparse_in,
is_sparse_out,
trid,
)
wgt = mapper["inputs"][trid]["weight_file"]
wgt_file = Path(wgt[ddi]) if is_sparse_out else Path(wgt)
# Loading weights if available
if wgt_file.is_file():
return load_weights(wgt_file)
debug(
"No weight file is available, "
f"recomputing weight with method '{transform.method}' for {trid}"
)
# Check compatibility of method with sparse data
if (is_sparse_out or is_sparse_in) and transform.method == "mass-conservation":
raise CifValueError(
"Mass-conservation method is not compatible with sparse data"
)
# Otherwise, creating them
if transform.method == "mass-conservation":
weights = reproject_emissions(
0.0 * zlon_in,
zlonc_in,
zlatc_in,
zlonc_out,
zlatc_out,
orig_regular=transform.orig_regular,
orig_unstructured=getattr(domain_in, "unstructured_domain", False),
orig_lon_cyclic=getattr(domain_in, "lon_cyclic", False),
target_unstructured=getattr(domain_out, "unstructured_domain", False),
return_weight=True,
resol=getattr(transform, "resol", 10),
rounding_domain=transform.rounding_domain,
)
elif transform.method == "fast-conservative":
processes = getattr(transform, "processes", None)
weights = conservative(domain_in, domain_out, transform.chunk_size, processes)
elif transform.method == "bilinear":
weights = bilinear(
domain_in, nlon_in, nlat_in, zlon_in, zlat_in, zlon_out, zlat_out
)
elif transform.method == "gridcell":
direction = getattr(transform, "direction", "forward")
grid_to_surface = getattr(transform, "grid_to_surface", False)
chunk_size = transform.chunk_size
weights = find_gridcells(
domain_in,
domain_out,
forward_direction=direction == "forward",
grid_to_surface=grid_to_surface,
chunksize=chunk_size,
)
else:
raise CifValueError(
f"Regrid method not known: {transform.method}\n"
"Please check your Yaml file"
)
save_weights(wgt_file, weights, domain_in, domain_out)
return weights
[docs]
def save_weights(
path: str | PathLike[str],
weights: dict[str, np.ndarray],
domain_in: Any | None = None,
domain_out: Any | None = None,
) -> None:
if not set(weights).issubset({"i", "j", "wgt", "filtered", "non_filtered"}):
raise CifValueError(
f"unexpected keys in weights, got {list(weights)}, "
+ "expected a subset of ['i', 'j', 'wgt', 'filtered', 'non_filtered']"
)
ds = xr.Dataset()
for varname in ["i", "j", "wgt"]:
val = weights[varname]
ds[varname] = (["n"], val.flatten())
if val.ndim > 1:
ds[varname].attrs["shape"] = val.shape
for varname in ["filtered", "non_filtered"]:
if varname in weights:
ds[varname] = ([f"n_{varname}"], weights[varname])
if domain_in is not None:
ds.attrs.update(
{
"domain_in_nlat": domain_in.nlat,
"domain_in_nlon": domain_in.nlon,
}
)
if domain_out is not None:
ds.attrs.update(
{
"domain_out_nlat": domain_out.nlat,
"domain_out_nlon": domain_out.nlon,
}
)
# Save as NetCDF with compression enabled
encoding = {varname: {"zlib": True, "complevel": 4} for varname in ds.data_vars}
with _hdf5_lock:
ds.to_netcdf(path, encoding=encoding)
[docs]
def load_weights(path: str | PathLike[str]) -> dict[str, np.ndarray]:
weights = {}
with _hdf5_lock:
with xr.open_dataset(path) as ds:
for varname, da in ds.data_vars.items():
if "shape" in da.attrs:
weights[varname] = da.values.reshape(da.attrs["shape"])
else:
weights[varname] = da.values
return weights
[docs]
def copy_weights(
transform,
mapper,
domain_in,
domain_out,
ddi,
is_sparse_in,
is_sparse_out,
trid,
):
orig_dir = Path(getattr(transform, "dir_wgt", ""))
prefix = f"wgt_file_{domain_in.nlat}-{domain_in.nlon}"
# Copy weigth file if specified
if not is_sparse_out:
# Pass if already defined
if "weight_file" in mapper["inputs"][trid]:
return
if getattr(transform, "target_lbc", False):
nlon_out = domain_out.nlon_side
nlat_out = domain_out.nlat_side
else:
nlat_out, nlon_out = domain_out.zlat.shape
filename = f"{prefix}_{nlat_out}-{nlon_out}_{transform.method}.nc"
target_file = Path(transform.dir_regrid, filename)
mapper["inputs"][trid]["weight_file"] = target_file
else:
# Pass if already defined
if "weight_file" not in mapper["inputs"][trid]:
mapper["inputs"][trid]["weight_file"] = {}
if ddi in mapper["inputs"][trid]["weight_file"]:
return
# Import all possible files in dir_wgt
filename = f"{prefix}_{domain_out.nlat}_{transform.method}_{ddi:%Y%m%d%H%M}.nc"
target_file = Path(transform.dir_regrid, filename)
mapper["inputs"][trid]["weight_file"][ddi] = target_file
orig_file = orig_dir / filename
if orig_file.is_file():
shutil.copy(orig_file, target_file)
