import xarray as xr
import numpy as np
import pandas as pd
from .......utils.check.errclass import CifError
[docs]
def array_adjoint(transf, mapper, inout_datastore,
ddi, onlyinit, **kwargs):
xmod_out = inout_datastore["outputs"]
trid_ref = list(xmod_out)[0]
# Inputs domain from the present tracer
domain_in = mapper["inputs"][trid_ref]["domain"]
nlev_in = domain_in.nlev
# Outputs domain from the present tracer
domain_out = mapper["outputs"][trid_ref]["domain"]
nlev_out = domain_out.nlev
if transf.coord_out == "height" or not hasattr(domain_out, "sigma_a"):
heights = domain_out.heights
sigma_b_out = np.exp(- transf.GC * transf.MMOL * heights
/ transf.GASC / transf.TS)
sigma_a_out = 0. * sigma_b_out
sigma_b_out = np.concatenate([[1], sigma_b_out])
sigma_a_out = np.concatenate([[0], sigma_a_out])
sigma_b_out_mid = 0.5 * (sigma_b_out[1:] + sigma_b_out[:-1])
sigma_a_out_mid = 0.5 * (sigma_a_out[1:] + sigma_a_out[:-1])
else:
sigma_a_out = domain_out.sigma_a
sigma_b_out = domain_out.sigma_b
sigma_a_out_mid = domain_out.sigma_a_mid
sigma_b_out_mid = domain_out.sigma_b_mid
if transf.coord_in == "height" or not hasattr(domain_in, "sigma_a"):
heights = domain_in.heights
sigma_b_in = np.exp(- transf.GC * transf.MMOL * heights
/ transf.GASC / transf.TS)
sigma_a_in = 0. * sigma_b_in
sigma_b_in = np.concatenate([[1], sigma_b_in])
sigma_a_in = np.concatenate([[0], sigma_a_in])
sigma_b_in_mid = 0.5 * (sigma_b_in[1:] + sigma_b_in[:-1])
sigma_a_in_mid = 0.5 * (sigma_a_in[1:] + sigma_a_in[:-1])
else:
sigma_a_in = domain_in.sigma_a
sigma_b_in = domain_in.sigma_b
sigma_a_in_mid = domain_in.sigma_a_mid
sigma_b_in_mid = domain_in.sigma_b_mid
# For linear and closest, use middle of layers
if transf.method in ["linear", "closest"]:
sigma_a_in = sigma_a_in_mid
sigma_b_in = sigma_b_in_mid
sigma_a_out = sigma_a_out_mid
sigma_b_out = sigma_b_out_mid
# For top-type, use the highest interface level
if mapper["outputs"][trid_ref]["is_top"]:
nlev_out = 1
sigma_a_out = domain_out.sigma_a[-1:]
sigma_b_out = domain_out.sigma_b[-1:]
# Converting all pressure to hPa if need
if getattr(domain_in, "pressure_unit", "") == "Pa":
sigma_a_in = sigma_a_in / 100.
if getattr(domain_out, "pressure_unit", "") == "Pa":
sigma_a_out = sigma_a_out / 100.
# Use virtual surface pressure to interpolate
psurf = transf.psurf
pres_in = sigma_a_in + psurf * sigma_b_in
pres_out = sigma_a_out + psurf * sigma_b_out
if "log_interp" in mapper["outputs"][trid_ref]:
pres_in = np.log10(pres_in)
pres_out = np.log10(pres_out)
# Now loop over trids
for trid in xmod_out:
if transf.method == "linear":
pres_tmp = np.sort(np.unique(np.append(pres_in, pres_out)))
df_pres = pd.DataFrame(range(len(pres_in)), index=pres_in)
df_pres = df_pres.reindex(pres_tmp).interpolate(method="index")
if transf.fill_nans:
df_pres = df_pres.fillna(method="bfill").fillna(method="ffill")
df_pres = df_pres.reindex(pres_out)
var_out = xmod_out[trid][ddi]["adj_out"]
ntimes, _, nlat, nlon = var_out.shape
var_in = np.zeros((ntimes, nlev_in, nlat, nlon), dtype=var_out.dtype)
for k, dd in enumerate(pres_out):
ind = df_pres.iloc[k, 0]
if np.isnan(ind):
continue
dmin = np.floor(ind).astype(int)
wgt = ind - dmin
var_in[:, dmin] += (1 - wgt) * var_out[:, k]
var_in[:, min(dmin + 1, nlev_in - 1)] += wgt * var_out[:, k]
inout_datastore["inputs"][trid][ddi] = {"adj_out": xr.DataArray(
var_in,
coords={"time": var_out.time},
dims=("time", "lev", "lat", "lon"),
)}
elif transf.method == "closest":
# if "adj_out" not in xmod[trid][ddi]:
# return
var_out = xmod_out[trid][ddi]["adj_out"]
ntimes, nlev, nlat, nlon = var_out.shape
var_in = np.zeros((ntimes, nlev_in, nlat, nlon), dtype=var_out.dtype)
# Get proper indexes for each level
dpres = (pres_out[:, np.newaxis] - pres_in)
kout = nlev_out
while kout > 0:
min_in = np.abs(dpres).min(axis=1)
iout = min_in.argmin()
iin = np.argmin(dpres[iout])
var_in[:, iin] += var_out[:, iout]
dpres[:, iin] = np.inf
dpres[iout] = np.inf
kout -= 1
if np.all(dpres == np.inf):
kout = 0
inout_datastore["inputs"][trid][ddi] = {"adj_out": xr.DataArray(
var_in,
coords={"time": var_out.time},
dims=("time", "lev", "lat", "lon"),
)}
else:
raise CifError(f"Don't know interpolation method: {transf.method}")
