import numpy as np
from logging import info
[docs]
def apply_ak(
sim_ak, aks, pwgt, drycols, qa0,
chosen_level=-1, use_drycols=False,
scale_factor=1, log_space=False,
normalize_columns=False
):
"""Apply the corresponding AK values
"""
if not log_space:
prior_profile = qa0[chosen_level] if chosen_level != -1 \
else np.nansum(qa0 * pwgt, axis=0)
y0 = np.nansum(
(sim_ak * drycols - qa0) * aks * pwgt,
axis=0
) + prior_profile
if not normalize_columns:
pass
elif use_drycols:
y0 /= np.nansum(drycols, axis=0)
else:
y0 /= np.nansum(pwgt * aks, axis=0)
y0 *= scale_factor
else:
prior_profile = np.log(qa0[chosen_level]) if chosen_level != -1 \
else np.nansum(pwgt * np.where(qa0 > 0, np.log(qa0), 0), axis=0)
y0 = np.nansum(
(np.where(sim_ak > 0, np.log(sim_ak), 0) * drycols
- np.where(qa0 > 0, np.log(qa0), 0)) * aks * pwgt,
axis=0
) + prior_profile
if not normalize_columns:
pass
elif use_drycols:
y0 /= np.nansum(drycols, axis=0)
else:
y0 /= np.nansum(pwgt * aks, axis=0)
y0 = np.exp(y0)
y0 *= scale_factor
return y0
[docs]
def apply_ak_tl(
sim_ak_tl, sim_ak, aks, pwgt, drycols, qa0,
chosen_level=-1, use_drycols=False,
scale_factor=1, log_space=False,
normalize_columns=False
):
if not log_space:
y0_tl = np.nansum((sim_ak_tl * drycols) * aks * pwgt, axis=0)
if not normalize_columns:
pass
elif use_drycols:
y0_tl /= np.nansum(drycols, axis=0)
else:
y0_tl /= np.nansum(pwgt * aks, axis=0)
y0_tl *= scale_factor
else:
y0 = apply_ak(sim_ak, aks, pwgt, drycols, qa0,
chosen_level=chosen_level, use_drycols=use_drycols,
scale_factor=scale_factor, log_space=log_space,
normalize_columns=normalize_columns)
y0_tl = np.nansum(
np.where(
sim_ak > 0,
sim_ak_tl / sim_ak * drycols,
0
) * aks * pwgt,
axis=0
)
if not normalize_columns:
pass
elif use_drycols:
y0_tl /= np.nansum(drycols, axis=0)
else:
y0_tl /= np.nansum(pwgt * aks, axis=0)
y0_tl *= y0
return y0_tl
[docs]
def apply_ak_ad(
obs_incr, sim_ak, aks, pwgt, drycols, qa0,
chosen_level=-1, use_drycols=False,
scale_factor=1, log_space=False,
normalize_columns=False
):
if not log_space:
obs_incr_out = obs_incr * scale_factor
if not normalize_columns:
pass
elif use_drycols:
obs_incr_out /= np.nansum(drycols, axis=0)
else:
obs_incr_out /= np.nansum(pwgt * aks, axis=0)
obs_incr_out = (obs_incr_out * drycols) * aks * pwgt
else:
y0 = apply_ak(sim_ak, aks, pwgt, drycols, qa0,
chosen_level=chosen_level, use_drycols=use_drycols,
scale_factor=scale_factor, log_space=log_space,
normalize_columns=normalize_columns)
obs_incr_out = obs_incr * y0
if not normalize_columns:
pass
elif use_drycols:
obs_incr_out /= np.nansum(drycols, axis=0)
else:
obs_incr_out /= np.nansum(pwgt * aks, axis=0)
obs_incr_out = np.where(
sim_ak > 0,
((obs_incr_out / sim_ak * drycols) * aks * pwgt),
0
)
return obs_incr_out
