import numpy as np
import pandas as pd
import xarray as xr
import copy
[docs]
def weight_fwd(transf, pres_in, pres_out,
mode, inout_datastore, trid, ddi, mapper):
nlev_in = pres_in.size
nlev_out = pres_out.size - 1
xmod = inout_datastore["inputs"]
in_thickness = np.abs(np.diff(pres_in))
# Find index of lower interface
index_tmplow = \
np.sort(np.unique(np.append(pres_in, pres_out[:-1])))
df_index = pd.Series(range(len(pres_in)),
index=pres_in)
df_index = df_index.reindex(index_tmplow)
low_inside = df_index.interpolate(
method="index", limit_area="inside").loc[pres_out[:-1]]
df_index = df_index.interpolate(
method="index", limit_direction="both")
out_index_low = df_index.loc[pres_out[:-1]]
# Find index of upper interface
index_tmpup = \
np.sort(np.unique(np.append(pres_in, pres_out[1:])))
df_index = \
pd.Series(range(len(pres_in)), index=pres_in)
df_index = df_index.reindex(index_tmpup)
up_inside = df_index.interpolate(
method="index", limit_area="inside").loc[pres_out[1:]]
df_index = df_index.interpolate(
method="index", limit_direction="both")
out_index_up = df_index.loc[pres_out[1:]]
# Compute weights for each layer
out_thickness = (
np.ceil(out_index_up.values)
- np.floor(out_index_low.values)).astype(int)
out_thickness[
np.isnan(low_inside).values
& np.isnan(up_inside).values] = 0
out_df = pd.DataFrame({"low": out_index_low.values,
"up": out_index_up.values,
"thickness": out_thickness})
weights = copy.deepcopy(
out_df.iloc[
out_df.index.repeat(
out_df["thickness"])])
indexes = np.zeros(len(weights))
low_index = np.array(
[0] + list(np.cumsum(
out_df["thickness"].loc[out_df["thickness"] > 0].values[:-1]))
).astype(int)
indexes[low_index] = low_index
np.maximum.accumulate(indexes, out=indexes)
weights["indexes"] = \
np.arange(len(weights)) - indexes + np.floor(weights["low"])
weights["weights"] = \
np.minimum(weights["indexes"] + 1, weights["up"]) \
- np.maximum(weights["indexes"], weights["low"])
weights["weights"] *= \
in_thickness[weights["indexes"].astype(int)]
group_weights = weights["weights"].groupby(by=weights.index).sum()
weights["weights"] /= group_weights.loc[weights.index].values
# Now apply weights
data_id = ["spec"] if mode == "fwd" else ["spec", "incr"]
inout_datastore["outputs"][trid][ddi] = {}
for did in data_id:
if did not in xmod[trid][ddi]:
continue
var_in = xmod[trid][ddi][did]
ntimes, nlev, nlat, nlon = var_in.shape
var_out = np.zeros((ntimes, nlev_out, nlat, nlon), dtype=var_in.dtype)
mesh_out = np.meshgrid(np.arange(ntimes),
weights.index,
np.arange(nlat),
np.arange(nlon),
indexing="ij")
mesh_in = np.meshgrid(np.arange(ntimes),
weights["indexes"].astype(int),
np.arange(nlat),
np.arange(nlon),
indexing="ij")
np.add.at(var_out,
tuple(mesh_out),
var_in.values[mesh_in[0], mesh_in[1], mesh_in[2], mesh_in[3]]
* weights["weights"].values[
np.newaxis, :, np.newaxis, np.newaxis])
inout_datastore["outputs"][trid][ddi][did] = xr.DataArray(
var_out,
coords={"time": var_in.time},
dims=("time", "lev", "lat", "lon"),
)
