import xarray as xr
import numpy as np
import pandas as pd
import copy
from .......utils.classes.domains import Domain
from .linear import linear_fwd
from .closest import closest_fwd
from .weight import weight_fwd
from .......utils.check.errclass import CifError
[docs]
def array_forward(transf, mapper, inout_datastore,
ddi, mode, onlyinit, **kwargs):
# # Does nothing in onlyinit mode
# if onlyinit:
# return
xmod_in = inout_datastore["inputs"]
# Inputs domain from the present tracer
trid_ref = list(mapper["inputs"].keys())[0]
domain_in = mapper["inputs"][trid_ref]["domain"]
nlev_in = domain_in.nlev
# Update input domain if sparse
if mapper["inputs"][trid_ref].get("sparse_data", False):
metadata = xmod_in[trid_ref][ddi]["metadata"]
alts = metadata["alt"].values
domain_in = Domain(nlev_in=len(alts),
heights=alts)
# Outputs domain from the present tracer
domain_out = mapper["outputs"][trid_ref]["domain"]
nlev_out = domain_out.nlev
# Extra optional info
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
if np.all(np.diff(heights) >= 0):
sigma_b_in = np.concatenate([[1], sigma_b_in])
sigma_a_in = np.concatenate([[0], sigma_a_in])
elif np.all(np.diff(heights) <= 0):
sigma_b_in = np.concatenate([sigma_b_in, [1]])
sigma_a_in = np.concatenate([sigma_a_in, [0]])
elif mapper["inputs"][trid_ref].get("sparse_data", False):
pass
else:
raise CifError(
"Heights should be ordered for the interpolation to work!"
)
if mapper["inputs"][trid_ref].get("sparse_data", False):
sigma_b_in_mid = sigma_b_in
sigma_a_in_mid = sigma_a_in
else:
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
# Extending inputs when outputs outside convex hull
if hasattr(transf, "extend_bottom"):
extension = transf.extend_bottom
a = extension.sigma_a
b = extension.sigma_b
pres_in = np.concatenate([[a + psurf * b], pres_in])
if extension.method == "fixed":
ntimes, nlev, nlat, nlon = xmod_in[trid_ref][ddi]["spec"].shape
for trid in xmod_in:
xmod_in[trid][ddi]["spec"] = \
xr.concat([
xr.DataArray(
data=extension.value *
np.ones((ntimes, 1, nlat, nlon)),
dims=["time", "lev", "lat", "lon"]),
xmod_in[trid][ddi]["spec"]
], dim="lev")
if "incr" in xmod_in[trid][ddi]:
xmod_in[trid][ddi]["incr"] = \
xr.concat([
xr.DataArray(
data=np.zeros((ntimes, 1, nlat, nlon)),
dims=["time", "lev", "lat", "lon"]),
xmod_in[trid][ddi]["incr"]
], dim="lev")
else:
raise CifError(
f"Method {extension.method} for extending the field is not implemented."
)
# Interpret in pressure or log-pressure
if "log_interp" in mapper["outputs"][trid_ref]:
pres_in = np.log10(pres_in)
pres_out = np.log10(pres_out)
# Different methods
for trid in xmod_in:
if transf.method == "linear":
linear_fwd(transf, pres_in, pres_out,
mode, inout_datastore, trid, ddi)
elif transf.method == "closest":
closest_fwd(transf, pres_in, pres_out,
mode, inout_datastore, trid, ddi, mapper)
elif transf.method == "weight":
weight_fwd(transf, pres_in, pres_out,
mode, inout_datastore, trid, ddi, mapper)
else:
raise CifError(f"Don't know interpolation method: {transf.method}")
