import os
import pickle
from shutil import copy
from logging import debug
from .bilinear import bilinear
from .conservative import conservative
from .find_gridcells import find_gridcells
from .reproject import reproject_emissions
from ......utils.check.errclass import CifValueError
[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 = wgt[ddi] if is_sparse_out else wgt
# Loading weights if available
if os.path.isfile(wgt_file):
with open(wgt_file, "rb") as f:
return pickle.load(f)
debug(f"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. * 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"
)
with open(wgt_file, "wb") as f:
pickle.dump(weights, f)
return weights
[docs]
def copy_weights(transform, mapper, domain_in, domain_out, ddi,
is_sparse_in, is_sparse_out, trid):
# 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
basename = \
f"wgt_file_{domain_in.nlat}-{domain_in.nlon}_{nlat_out}-{nlon_out}_{transform.method}.pickle"
target_file = f"{transform.dir_regrid}/{basename}"
orig_file = f"{getattr(transform, 'dir_wgt', '')}/{basename}"
if os.path.isfile(orig_file):
copy(orig_file, target_file)
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
basename = ddi.strftime(
f"wgt_file_{domain_in.nlat}-{domain_in.nlon}_{domain_out.nlat}_{transform.method}_%Y%m%d%H%M.pickle")
orig_file = f"{getattr(transform, 'dir_wgt', '')}/{basename}"
target_file = f"{transform.dir_regrid}/{basename}"
if os.path.isfile(orig_file):
copy(orig_file, target_file)
mapper["inputs"][trid]["weight_file"][ddi] = target_file
