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.classes.domains import Domain
from ......utils.hdf5 import _hdf5_lock
from .bilinear import bilinear
from .conservative import conservative, conservative_sparse
from .find_gridcells import find_gridcells
from .reproject import reproject_emissions
[docs]
def get_weights(transform, trid, mapper, inout_datastore, xmod, ddi):
"""Get regridding weights between the input and output domains of a tracer.
Weights are loaded from file if available, computed and saved otherwise.
Metadata of sparse data (longitudes and latitudes) are only needed
to compute new weights.
Args:
transform: regrid transform.
trid: tracer id used as reference.
mapper: transform mapper.
inout_datastore: datastore of the transform.
xmod: datastore of the data being regridded (inputs in forward,
outputs in adjoint).
ddi: sub-simulation start date.
Returns:
tuple: weights, and dimensions
``(nlat_in, nlon_in, nlat_out, nlon_out)`` of the domains.
"""
is_sparse_in = mapper["inputs"][trid].get("sparse_data", False)
is_sparse_out = mapper["outputs"][trid].get("sparse_data", False) \
or mapper["outputs"][trid]["sampled"]
is_lbc = mapper["outputs"][trid].get("is_lbc", False)
# Sparse domains first only carry their number of points, which is
# enough to find existing weights; coordinates from the metadata
# are only needed to compute new weights
sparse_datastores = []
# Inputs domain from the present tracer
if is_sparse_in:
ds_in = inout_datastore["outputs"][trid][ddi]
domain_in = sparse_domain(len(ds_in))
sparse_datastores.append(
(domain_in, ds_in, mapper["inputs"][trid].get("ncorners")))
nlat_in, nlon_in = len(ds_in), len(ds_in)
else:
domain_in = mapper["inputs"][trid]["domain"]
nlat_in, nlon_in = domain_in.zlon.shape
# Outputs domain from sparse data, from the mapper otherwise
if is_sparse_out:
ds_out = xmod[trid][ddi]
domain_out = sparse_domain(len(ds_out))
sparse_datastores.append(
(domain_out, ds_out, mapper["outputs"][trid].get("ncorners")))
nlat_out, nlon_out = len(ds_out), 1
else:
domain_out = mapper["outputs"][trid]["domain"]
# Differentiate dimensions if LBC
if is_lbc:
nlat_out, nlon_out = domain_out.nlat_side, domain_out.nlon_side
else:
nlat_out, nlon_out = domain_out.zlat.shape
dimensions = (nlat_in, nlon_in, nlat_out, nlon_out)
# 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
# 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), dimensions
debug(
"No weight file is available, "
f"recomputing weight with method '{transform.method}' for {trid}"
)
# Locations of sparse data are needed to compute weights
for domain, datastore, ncorners in sparse_datastores:
set_sparse_coordinates(domain, datastore, ncorners)
# 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)
# 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)
if is_sparse_out:
weights = conservative_sparse(
domain_in, domain_out, transform.chunk_size, processes,
min_coverage=transform.min_coverage,
)
else:
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, dimensions
[docs]
def sparse_domain(npoints):
domain = Domain()
domain.nlon = npoints
domain.nlat = npoints
return domain
[docs]
def set_sparse_coordinates(domain, datastore, ncorners=None):
ds = datastore["metadata"]
domain.zlon = np.asarray(ds["lon"])[:, np.newaxis]
domain.zlonc = domain.zlon
domain.zlat = np.asarray(ds["lat"])[:, np.newaxis]
domain.zlatc = domain.zlat
# Footprints as an unstructured domain of shape (ncorners, npoints)
if ncorners:
domain.unstructured_domain = True
domain.zlonc = np.array(
[np.asarray(ds[f"lon_{k}"], dtype=float) for k in range(ncorners)])
domain.zlatc = np.array(
[np.asarray(ds[f"lat_{k}"], dtype=float) for k in range(ncorners)])
[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
# Weights from footprints differ from weights from points
method = transform.method
ncorners = mapper["inputs"][trid].get("ncorners") \
or mapper["outputs"][trid].get("ncorners")
if ncorners:
method = f"{method}-{ncorners}corners"
# Import all possible files in dir_wgt
filename = f"{prefix}_{domain_out.nlat}_{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)
