def _get_dof_indices(freedofs, fixdofs, k_xlist, k_ylist):
index_map = autograd_lib.inverse_permutation(
np.concatenate([freedofs, fixdofs]))
keep = np.isin(k_xlist, freedofs) & np.isin(k_ylist, freedofs)
i = index_map[k_ylist][keep]
j = index_map[k_xlist][keep]
return index_map, keep, np.stack([i, j])
def getIcase(par1type, par2type, checks=True):
"""Generate vector describing the combination of input parameters.
Options for `par1type` and `par2type`:
* `1` = total alkalinity
* `2` = dissolved inorganic carbon
* `3` = pH
* `4` = partial pressure of CO2
* `5` = fugacity of CO2
* `6` = carbonate ion
* `7` = bicarbonate ion
* `8` = aqueous CO2
* `9` = dry mole fraction of CO2
`Icase` is `10*parXtype + parYtype` where `parXtype` is whichever of `par1type` or
`par2type` is greater.
Noting that a pair of any two from pCO2, fCO2, xCO2 CO2(aq) is not allowed, the
valid `Icase` options are:
12, 13, 14, 15, 16, 17, 18, 19,
23, 24, 25, 26, 27, 28, 29,
34, 35, 36, 37, 38, 39,
46, 47,
56, 57,
67, 68, 69,
78, 79.
The optional argument `checks` allows you to decide whether the function should test
the validity of the entered combinations or not.
"""
# Check validity of separate `par1type` and `par2type` inputs
if checks:
assert np.all(
np.isin(par1type, [1, 2, 3, 4, 5, 6, 7, 8, 9])
& np.isin(par2type, [1, 2, 3, 4, 5, 6, 7, 8, 9])
), "All `par1type` and `par2type` values must be integers from 1 to 9."
assert ~np.any(
par1type == par2type
), "`par1type` and `par2type` must be different from each other."
# Combine inputs into `Icase` and check its validity
Icase = np.where(
par1type < par2type, 10 * par1type + par2type, par1type + 10 * par2type
)
if checks:
assert ~np.any(
np.isin(Icase, [45, 48, 49, 58, 59, 89])
), "Combinations of pCO2, fCO2, xCO2 and CO2(aq) are not valid argument pairs."
return Icase
def _flattentext(args, npts):
# Determine and check lengths of input vectors
arglengths = np.array([np.size(arg) for arg in args])
assert np.all(
np.isin(arglengths, [1, npts])
), "Inputs must all be the same length as each other or of length 1."
# Make vectors of all inputs
return [np.full(npts, arg) if np.size(arg) == 1 else arg.ravel() for arg in args]
def KSO4(TempK, Sal, Pbar, RGas, WhoseKSO4):
"""Calculate bisulfate ion dissociation constant for the given options."""
assert np.all(np.isin(
WhoseKSO4, [1, 2])), "Valid `WhoseKSO4` options are: `1` or `2`."
# Evaluate at atmospheric pressure
KSO4 = np.full(np.shape(TempK), np.nan)
KSO4 = np.where(WhoseKSO4 == 1, p1atm.kHSO4_FREE_D90a(TempK, Sal), KSO4)
KSO4 = np.where(WhoseKSO4 == 2, p1atm.kHSO4_FREE_KRCB77(TempK, Sal), KSO4)
# Now correct for seawater pressure
KSO4 = KSO4 * pcx.KSO4fac(TempK, Pbar, RGas)
return KSO4
def load_data():
X, y = loadlocal_mnist(
images_path=str(config.mnist / 'train-images-idx3-ubyte'),
labels_path=str(config.mnist / 'train-labels-idx1-ubyte'))
keep_digits = np.isin(y, [0, 1, 4, 7])
X = X[keep_digits]
y = y[keep_digits]
X = X / 255
X = (X >= 0.5).astype(int) # Binarizing
return X, y
def KF(TempK, Sal, Pbar, RGas, WhoseKF):
"""Calculate HF dissociation constant for the given options."""
assert np.all(np.isin(WhoseKF,
[1, 2])), "Valid `WhoseKF` options are: `1` or `2`."
# Evaluate at atmospheric pressure
KF = np.full(np.shape(TempK), np.nan)
KF = np.where(WhoseKF == 1, p1atm.kHF_FREE_DR79(TempK, Sal), KF)
KF = np.where(WhoseKF == 2, p1atm.kHF_FREE_PF87(TempK, Sal), KF)
# Now correct for seawater pressure
KF = KF * pcx.KFfac(TempK, Pbar, RGas)
return KF
def _rr_parcombos(par1type, par2type):
"""Generate all possible valid pairs of parameter type numbers excluding the input
pair.
"""
assert isscalar(par1type) & isscalar(
par2type), "Both inputs must be scalar."
# Get all possible combinations of parameter type numbers
allpars = list(_partypes.keys())
par1s, par2s = meshgrid(allpars, allpars)
par1s = par1s.ravel()
par2s = par2s.ravel()
# Select only valid combinations and cut out input combination
allIcases = solve.getIcase(par1s, par2s, checks=False)
inputIcase = solve.getIcase(par1type, par2type, checks=False)
valid = (par1s != par2s) & ~isin(allIcases, [45, 48, 58, inputIcase])
par1s = par1s[valid]
par2s = par2s[valid]
# Icases = pyco2.engine.getIcase(par1s, par2s, checks=True) # checks if all valid
return par1s, par2s
def _test_transform_output(self, original_fun, patterns, free, retnums,
original_is_flat):
# original_fun must take no arguments.
fun_trans = paragami.TransformFunctionOutput(original_fun, patterns,
free, original_is_flat,
retnums)
# Check that the flattened and original function are the same.
def check_equal(orig_val, trans_val, pattern, free):
# Use the flat representation to check that parameters are equal.
if original_is_flat:
assert_array_almost_equal(
orig_val, pattern.flatten(trans_val, free=free))
else:
assert_array_almost_equal(pattern.flatten(orig_val, free=free),
trans_val)
patterns_array = np.atleast_1d(patterns)
free_array = np.atleast_1d(free)
retnums_array = np.atleast_1d(retnums)
orig_rets = original_fun()
trans_rets = fun_trans()
if isinstance(orig_rets, tuple):
self.assertTrue(len(orig_rets) == len(trans_rets))
# Check that the non-transformed return values are the same.
for ind in range(len(orig_rets)):
if not np.isin(ind, retnums):
assert_array_almost_equal(orig_rets[ind], trans_rets[ind])
# Check that the transformed return values are the same.
for ret_ind in range(len(retnums_array)):
ind = retnums_array[ret_ind]
check_equal(orig_rets[ind], trans_rets[ind],
patterns_array[ret_ind], free_array[ret_ind])
else:
check_equal(orig_rets, trans_rets, patterns_array[0],
free_array[0])
# Check that the string method works.
str(fun_trans)
def __init__(self,
counts,
lengths,
ploidy,
multiscale_factor=1,
constraint_lambdas=None,
constraint_params=None):
self.lengths = np.array(lengths)
self.lengths_lowres = decrease_lengths_res(lengths, multiscale_factor)
self.ploidy = ploidy
self.multiscale_factor = multiscale_factor
if constraint_lambdas is None:
self.lambdas = {}
else:
self.lambdas = constraint_lambdas
if constraint_params is None:
self.params = {}
else:
self.params = constraint_params
torm = find_beads_to_remove(counts=counts,
nbeads=self.lengths_lowres.sum() * ploidy)
self.torm_3d = np.repeat(torm.reshape(-1, 1), 3, axis=1)
self.row, self.col = _constraint_dis_indices(
counts=counts,
n=self.lengths_lowres.sum(),
lengths=self.lengths_lowres,
ploidy=ploidy)
# Calculating distances for neighbors, which are on the off diagonal
# line - i & j where j = i + 1
row_adj = ag_np.unique(self.row).astype(int)
row_adj = row_adj[ag_np.isin(row_adj + 1, self.col)]
# Remove if "neighbor" beads are actually on different chromosomes or
# homologs
self.row_adj = row_adj[ag_np.digitize(
row_adj,
np.tile(lengths, ploidy).cumsum()) == ag_np.digitize(
row_adj + 1,
np.tile(lengths, ploidy).cumsum())]
self.col_adj = self.row_adj + 1
self.check()
def __init__(self, K, D, transition_mask=None, M=0):
super(ConstrainedStationaryTransitions, self).__init__(K, D, M=M)
Ps = self.transition_matrix
if transition_mask is None:
transition_mask = np.ones_like(Ps)
# Validate the transition mask. A valid mask must have be the same shape
# as the transition matrix, contain only ones and zeros, and contain at
# least one non-zero entry per row.
assert transition_mask.shape == Ps.shape, "Mask must be the same size " \
"as the transition matrix. Found mask of shape {}".format(mask.shape)
assert np.isin(transition_mask,[1,0]).all(), "Mask must contain only 1s and zeros."
for i in range(transition_mask.shape[0]):
assert transition_mask[i].any(), "Mask must contain at least one " \
"nonzero entry per row."
self.transition_mask = transition_mask
Ps = Ps * transition_mask
Ps /= Ps.sum(axis=-1, keepdims=True)
self.log_Ps = np.log(Ps)
def optimize(timestamp_dims, mark, to_dim, dim, edge, mat_excition, rho,
sentiments):
neighbor_dims = np.arange(dim)[edge[:, to_dim] == 1]
neighbor_timestamp_idxs = np.isin(timestamp_dims, neighbor_dims)
mat_excition = mat_excition[neighbor_timestamp_idxs]
timestamp_dims = timestamp_dims[neighbor_timestamp_idxs]
if mat_excition.shape[0] == 0:
return np.zeros(dim), -np.inf
bounds = [(1e-15, None) for _ in range(neighbor_dims.size)]
optimization_param_id = {}
result_beta = np.zeros(dim)
if mat_excition.shape[0] == 0:
return result_beta, -np.inf
for neighbor_dim in neighbor_dims:
optimization_param_id[neighbor_dim] = len(optimization_param_id)
optimization_timestamp_dims = np.array(
[optimization_param_id[dim] for dim in timestamp_dims])
beta_u = np.random.uniform(0, 1, size=neighbor_dims.size)
res = minimize(mark_likelihood,
beta_u,
args=(mat_excition, mark, sentiments,
optimization_timestamp_dims, -1.0, rho),
method="L-BFGS-B",
jac=grad(mark_likelihood),
bounds=bounds,
options={
"ftol": 1e-10,
"maxls": 50,
"maxcor": 50,
"maxiter": 100,
"maxfun": 100
})
result_beta[neighbor_dims] = res.x
result_beta[result_beta < 0] = 0
return result_beta, -res.fun
def optimize_exo(timestamp_dims, to_dim, dim, omega, edge, mat_excition, rho):
# global funs, mat_excitions
neighbor_dims = np.arange(dim)[edge[:, to_dim] == 1]
neighbor_timestamp_idxs = np.isin(timestamp_dims, neighbor_dims)
mat_excition = mat_excition[neighbor_timestamp_idxs]
timestamp_dims = timestamp_dims[neighbor_timestamp_idxs]
result_alpha = np.zeros(dim)
if mat_excition.shape[0] == 0:
return result_alpha, -np.inf
bounds = [(0, None) for _ in range(neighbor_dims.size)]
optimization_param_id = {}
for neighbor_dim in neighbor_dims:
optimization_param_id[neighbor_dim] = len(optimization_param_id)
optimization_timestamp_dims = np.array(
[optimization_param_id[dim] for dim in timestamp_dims])
alpha_u = np.random.uniform(0, 1, size=neighbor_dims.size)
res = minimize(
hawkes_likelihood,
alpha_u,
args=(mat_excition, optimization_timestamp_dims, omega, -1, rho,
False),
method="L-BFGS-B",
# jac=hawkes_likelihood_grad_exo,
jac=grad(hawkes_likelihood),
bounds=bounds,
options={
"ftol": 1e-10,
"maxls": 50,
"maxcor": 50,
"maxiter": 100,
"maxfun": 100
})
return res.x
def forward(
co2dict,
grads_of,
grads_wrt,
totals=None,
equilibria_in=None,
equilibria_out=None,
dx=1e-6,
dx_scaling="median",
dx_func=None,
):
"""Get forward finite-difference derivatives of CO2SYS outputs w.r.t. inputs.
Arguments:
co2dict -- output generated by `PyCO2SYS.CO2SYS`.
grads_of -- list of keys from `co2dict` that you want to calculate the derivatives
of, or a single key as a string, or "all".
grads_wrt -- list of `PyCO2SYS.CO2SYS` input variable names that you want to
calculate the derivatives with respect to, or a single name as a string, or
"all".
Keyword arguments:
totals -- dict of internal override total salt concentrations identical to that used
to generate the `co2dict` (default None).
equilibria_in -- dict of internal override equilibrium constants at input conditions
identical to that used to generate the `co2dict` (default None).
equilibria_out -- dict of internal override equilibrium constants at output
conditions identical to that used to generate the `co2dict` (default None).
dx -- the forward difference for the derivative estimation (default 1e-6).
dx_scaling -- method for scaling `dx` for each variable, can be one of: "median"
(default), "none", or "custom".
dx_func -- function of each variable to scale `dx` with if dx_scaling="custom".
"""
# Derivatives can be calculated w.r.t. these inputs only
inputs_wrt = [
"PAR1",
"PAR2",
"SAL",
"TEMPIN",
"TEMPOUT",
"PRESIN",
"PRESOUT",
"SI",
"PO4",
"NH3",
"H2S",
]
totals_wrt = ["TB", "TF", "TSO4", "TCa"]
Ks_wrt = [
"KSO4",
"KF",
"fH",
"KB",
"KW",
"KP1",
"KP2",
"KP3",
"KSi",
"K1",
"K2",
"KH2S",
"KNH3",
"K0",
"FugFac",
"KCa",
"KAr",
]
Kis_wrt = ["{}input".format(K) for K in Ks_wrt]
Kos_wrt = ["{}output".format(K) for K in Ks_wrt]
Kis_wrt.append("RGas")
Kos_wrt.append("RGas")
pKis_wrt = ["p{}input".format(K) for K in Ks_wrt if K.startswith("K")]
pKos_wrt = ["p{}output".format(K) for K in Ks_wrt if K.startswith("K")]
# If only a single `grads_wrt` is requested, check it's allowed & convert to list
groups_wrt = [
"all", "measurements", "totals", "equilibria_in", "equilibria_out"
]
all_wrt = inputs_wrt + totals_wrt + Kis_wrt + Kos_wrt + pKis_wrt + pKos_wrt
if isinstance(grads_wrt, str):
assert grads_wrt in (all_wrt + groups_wrt)
if grads_wrt == "all":
grads_wrt = all_wrt
elif grads_wrt == "measurements":
grads_wrt = inputs_wrt
elif grads_wrt == "totals":
grads_wrt = totals_wrt
elif grads_wrt == "equilibria_in":
grads_wrt = Kis_wrt
elif grads_wrt == "equilibria_out":
grads_wrt = Kos_wrt
else:
grads_wrt = [grads_wrt]
# Make sure all requested `grads_wrt` are allowed
assert np.all(np.isin(list(grads_wrt),
all_wrt)), "Invalid `grads_wrt` requested."
# If only a single `grads_of` is requested, check it's allowed & convert to list
if isinstance(grads_of, str):
assert grads_of in ["all"] + list(engine.gradables)
if grads_of == "all":
grads_of = engine.gradables
else:
grads_of = [grads_of]
# Final validity checks
assert np.all(np.isin(grads_of,
engine.gradables)), "Invalid `grads_of` requested."
assert dx > 0, "`dx` must be positive."
# Assemble input arguments for engine._CO2SYS()
co2args = {
arg: co2dict[arg]
for arg in [
"PAR1",
"PAR2",
"PAR1TYPE",
"PAR2TYPE",
"SAL",
"TEMPIN",
"TEMPOUT",
"PRESIN",
"PRESOUT",
"SI",
"PO4",
"NH3",
"H2S",
"pHSCALEIN",
"K1K2CONSTANTS",
"KSO4CONSTANT",
"KFCONSTANT",
"BORON",
"buffers_mode",
"WhichR",
]
}
co2kwargs = {
"KSO4CONSTANTS": co2dict["KSO4CONSTANTS"],
"totals": totals,
"equilibria_in": equilibria_in,
"equilibria_out": equilibria_out,
}
# Preallocate output dict to store the gradients
co2derivs = {of: {wrt: None for wrt in grads_wrt} for of in grads_of}
dxs = {wrt: None for wrt in grads_wrt}
# Estimate the gradients with central differences
for wrt in grads_wrt:
# Make copies of input args to modify
co2args_plus = copy.deepcopy(co2args)
co2kwargs_plus = copy.deepcopy(co2kwargs)
# Perturb if `wrt` is one of the main inputs to CO2SYS
if wrt in inputs_wrt:
dx_wrt = _get_dx_wrt(dx,
np.median(co2args_plus[wrt]),
dx_scaling,
dx_func=dx_func)
co2args_plus[wrt] = co2args_plus[wrt] + dx_wrt
# Perturb if `wrt` is one of the `totals` internal overrides
elif wrt in totals_wrt:
co2kwargs_plus, dx_wrt = _overridekwargs(co2dict,
co2kwargs_plus,
"totals",
wrt,
dx,
dx_scaling,
dx_func=dx_func)
# Perturb if `wrt` is one of the `equilibria_in` internal overrides
elif wrt in Kis_wrt:
co2kwargs_plus, dx_wrt = _overridekwargs(
co2dict,
co2kwargs_plus,
"equilibria_in",
wrt,
dx,
dx_scaling,
dx_func=dx_func,
)
# Perturb if `wrt` is one of the `equilibria_in` internal overrides and the pK
# derivative is requested
elif wrt in pKis_wrt:
co2kwargs_plus, dx_wrt = _overridekwargs(
co2dict,
co2kwargs_plus,
"equilibria_in",
wrt,
dx,
dx_scaling,
dx_func=dx_func,
)
# Perturb if `wrt` is one of the `equilibria_out` internal overrides
elif wrt in Kos_wrt:
co2kwargs_plus, dx_wrt = _overridekwargs(
co2dict,
co2kwargs_plus,
"equilibria_out",
wrt,
dx,
dx_scaling,
dx_func=dx_func,
)
# Solve CO2SYS with the perturbation applied
co2dict_plus = engine._CO2SYS(**co2args_plus, **co2kwargs_plus)
dxs[wrt] = dx_wrt
# Extract results and calculate forward finite difference derivatives
for of in grads_of:
if co2derivs[of][
wrt] is None: # don't overwrite existing derivatives
co2derivs[of][wrt] = (co2dict_plus[of] - co2dict[of]) / dx_wrt
return co2derivs, dxs
def forward_nd(
CO2SYS_nd_results,
grads_of,
grads_wrt,
dx=1e-6,
dx_scaling="median",
dx_func=None,
**CO2SYS_nd_kwargs,
):
"""Get forward finite-difference derivatives of CO2SYS_nd results with respect to
its arguments.
"""
# Check requested grads are possible
assert np.all(
np.isin(
grads_of,
engine.nd.gradables + [
"p{}".format(gradable) for gradable in engine.nd.gradables
if gradable.startswith("k_")
],
)
), "PyCO2SYS error: all grads_of must be in the list at PyCO2SYS.engine.nd.gradables."
if np.any([of.endswith("_out") for of in grads_of]):
assert "temperature_out" in CO2SYS_nd_results, (
"PyCO2SYS error: you can only get gradients at output conditions if you calculated"
+ "results at output conditions!")
# Extract CO2SYS_nd fixed args from CO2SYS_nd_results and CO2SYS_nd_kwargs
keys_fixed = set([
"par1",
"par2",
"par1_type",
"par2_type",
"salinity",
"temperature",
"pressure",
"total_ammonia",
"total_phosphate",
"total_silicate",
"total_sulfide",
"opt_gas_constant",
"opt_k_bisulfate",
"opt_k_carbonic",
"opt_k_fluoride",
"opt_pH_scale",
"opt_total_borate",
"buffers_mode",
] + list(CO2SYS_nd_kwargs.keys()))
args_fixed = {k: CO2SYS_nd_results[k] for k in keys_fixed}
# Loop through requested parameters and calculate the gradients
dxs = {wrt: None for wrt in grads_wrt}
CO2SYS_derivs = {of: {wrt: None for wrt in grads_wrt} for of in grads_of}
for wrt in grads_wrt:
args_plus = copy.deepcopy(args_fixed)
is_pk = wrt.startswith("pk_")
if is_pk:
wrt_k = wrt[1:]
pk_values = -np.log10(CO2SYS_nd_results[wrt_k])
dxs[wrt] = _get_dx_wrt(dx, pk_values, dx_scaling, dx_func=dx_func)
pk_values_plus = pk_values + dxs[wrt]
args_plus[wrt_k] = 10.0**-pk_values_plus
else:
dxs[wrt] = _get_dx_wrt(dx,
CO2SYS_nd_results[wrt],
dx_scaling,
dx_func=dx_func)
args_plus[wrt] = CO2SYS_nd_results[wrt] + dxs[wrt]
results_plus = engine.nd.CO2SYS(**args_plus)
for of in grads_of:
CO2SYS_derivs[of][wrt] = (results_plus[of] -
CO2SYS_nd_results[of]) / dxs[wrt]
return CO2SYS_derivs, dxs
def fill(Icase, TA, TC, PH, PC, FC, CARB, HCO3, CO2, XC, totals, Ks):
"""Fill part-empty core marine carbonate system variable columns with solutions."""
# For convenience
K0 = Ks["K0"]
PengCx = totals["PengCorrection"]
# Convert any pCO2 and CO2(aq) values into fCO2
PCgiven = np.isin(Icase, [14, 24, 34, 46, 47])
FC = np.where(PCgiven, convert.pCO2_to_fCO2(PC, Ks), FC)
CO2given = np.isin(Icase, [18, 28, 38, 68, 78])
FC = np.where(CO2given, convert.CO2aq_to_fCO2(CO2, Ks), FC)
XCgiven = np.isin(Icase, [19, 29, 39, 69, 79])
FC = np.where(XCgiven, convert.xCO2_to_fCO2(XC, Ks), FC)
# Solve the marine carbonate system
F = Icase == 12 # input TA, TC
if np.any(F):
PH = np.where(F, get.pHfromTATC(TA - PengCx, TC, totals, Ks), PH)
# ^pH is returned on the same scale as `Ks`
FC = np.where(F, get.fCO2fromTCpH(TC, PH, totals, Ks), FC)
CARB = np.where(F, get.CarbfromTCpH(TC, PH, totals, Ks), CARB)
HCO3 = np.where(F, get.HCO3fromTCpH(TC, PH, totals, Ks), HCO3)
F = Icase == 13 # input TA, pH
if np.any(F):
TC = np.where(F, get.TCfromTApH(TA - PengCx, PH, totals, Ks), TC)
FC = np.where(F, get.fCO2fromTCpH(TC, PH, totals, Ks), FC)
CARB = np.where(F, get.CarbfromTCpH(TC, PH, totals, Ks), CARB)
HCO3 = np.where(F, get.HCO3fromTCpH(TC, PH, totals, Ks), HCO3)
F = (
(Icase == 14) | (Icase == 15) | (Icase == 18) | (Icase == 19)
) # input TA, [pCO2|fCO2|CO2aq|xCO2]
if np.any(F):
PH = np.where(F, get.pHfromTAfCO2(TA - PengCx, FC, totals, Ks), PH)
TC = np.where(F, get.TCfromTApH(TA - PengCx, PH, totals, Ks), TC)
CARB = np.where(F, get.CarbfromTCpH(TC, PH, totals, Ks), CARB)
HCO3 = np.where(F, get.HCO3fromTCpH(TC, PH, totals, Ks), HCO3)
HCO3 = np.where(Icase == 18, TC - CARB - CO2, HCO3)
F = Icase == 16 # input TA, CARB
if np.any(F):
PH = np.where(F, get.pHfromTACarb(TA - PengCx, CARB, totals, Ks), PH)
TC = np.where(F, get.TCfromTApH(TA - PengCx, PH, totals, Ks), TC)
FC = np.where(F, get.fCO2fromTCpH(TC, PH, totals, Ks), FC)
HCO3 = np.where(F, get.HCO3fromTCpH(TC, PH, totals, Ks), HCO3)
F = Icase == 17 # input TA, HCO3
if np.any(F):
PH = np.where(F, get.pHfromTAHCO3(TA - PengCx, HCO3, totals, Ks), PH)
TC = np.where(F, get.TCfromTApH(TA - PengCx, PH, totals, Ks), TC)
FC = np.where(F, get.fCO2fromTCpH(TC, PH, totals, Ks), FC)
CARB = np.where(F, get.CarbfromTCpH(TC, PH, totals, Ks), CARB)
F = Icase == 23 # input TC, pH
if np.any(F):
TA = np.where(F, get.TAfromTCpH(TC, PH, totals, Ks) + PengCx, TA)
FC = np.where(F, get.fCO2fromTCpH(TC, PH, totals, Ks), FC)
CARB = np.where(F, get.CarbfromTCpH(TC, PH, totals, Ks), CARB)
HCO3 = np.where(F, get.HCO3fromTCpH(TC, PH, totals, Ks), HCO3)
F = (
(Icase == 24) | (Icase == 25) | (Icase == 28) | (Icase == 29)
) # input TC, [pCO2|fCO2|CO2aq|xCO2]
if np.any(F):
PH = np.where(F, get.pHfromTCfCO2(TC, FC, totals, Ks), PH)
TA = np.where(F, get.TAfromTCpH(TC, PH, totals, Ks) + PengCx, TA)
CARB = np.where(F, get.CarbfromTCpH(TC, PH, totals, Ks), CARB)
HCO3 = np.where(F, get.HCO3fromTCpH(TC, PH, totals, Ks), HCO3)
HCO3 = np.where(Icase == 28, TC - CARB - CO2, HCO3)
F = Icase == 26 # input TC, CARB
if np.any(F):
PH = np.where(F, get.pHfromTCCarb(TC, CARB, totals, Ks), PH)
FC = np.where(F, get.fCO2fromTCpH(TC, PH, totals, Ks), FC)
TA = np.where(F, get.TAfromTCpH(TC, PH, totals, Ks) + PengCx, TA)
HCO3 = np.where(F, get.HCO3fromTCpH(TC, PH, totals, Ks), HCO3)
F = Icase == 27 # input TC, HCO3
if np.any(F):
PH = np.where(F, get.pHfromTCHCO3(TC, HCO3, totals, Ks), PH)
FC = np.where(F, get.fCO2fromTCpH(TC, PH, totals, Ks), FC)
TA = np.where(F, get.TAfromTCpH(TC, PH, totals, Ks) + PengCx, TA)
CARB = np.where(F, get.CarbfromTCpH(TC, PH, totals, Ks), CARB)
F = (
(Icase == 34) | (Icase == 35) | (Icase == 38) | (Icase == 39)
) # input pH, [pCO2|fCO2|CO2aq|xCO2]
if np.any(F):
TC = np.where(F, get.TCfrompHfCO2(PH, FC, totals, Ks), TC)
TA = np.where(F, get.TAfromTCpH(TC, PH, totals, Ks) + PengCx, TA)
CARB = np.where(F, get.CarbfromTCpH(TC, PH, totals, Ks), CARB)
HCO3 = np.where(F, get.HCO3fromTCpH(TC, PH, totals, Ks), HCO3)
HCO3 = np.where(Icase == 38, TC - CARB - CO2, HCO3)
F = Icase == 36 # input pH, CARB
if np.any(F):
FC = np.where(F, get.fCO2frompHCarb(PH, CARB, totals, Ks), FC)
TC = np.where(F, get.TCfrompHfCO2(PH, FC, totals, Ks), TC)
TA = np.where(F, get.TAfromTCpH(TC, PH, totals, Ks) + PengCx, TA)
HCO3 = np.where(F, get.HCO3fromTCpH(TC, PH, totals, Ks), HCO3)
F = Icase == 37 # input pH, HCO3
if np.any(F):
TC = np.where(F, get.TCfrompHHCO3(PH, HCO3, totals, Ks), TC)
TA = np.where(F, get.TAfromTCpH(TC, PH, totals, Ks) + PengCx, TA)
FC = np.where(F, get.fCO2fromTCpH(TC, PH, totals, Ks), FC)
CARB = np.where(F, get.CarbfromTCpH(TC, PH, totals, Ks), CARB)
F = (
(Icase == 46) | (Icase == 56) | (Icase == 68) | (Icase == 69)
) # input [pCO2|fCO2|CO2aq|xCO2], CARB
if np.any(F):
PH = np.where(F, get.pHfromfCO2Carb(FC, CARB, totals, Ks), PH)
TC = np.where(F, get.TCfrompHfCO2(PH, FC, totals, Ks), TC)
TA = np.where(F, get.TAfromTCpH(TC, PH, totals, Ks) + PengCx, TA)
HCO3 = np.where(F, get.HCO3fromTCpH(TC, PH, totals, Ks), HCO3)
HCO3 = np.where(Icase == 68, TC - CARB - CO2, HCO3)
F = (
(Icase == 47) | (Icase == 57) | (Icase == 78) | (Icase == 79)
) # input [pCO2|fCO2|CO2aq|xCO2], HCO3
if np.any(F):
CARB = np.where(F, get.CarbfromfCO2HCO3(FC, HCO3, totals, Ks), CARB)
PH = np.where(F, get.pHfromfCO2Carb(FC, CARB, totals, Ks), PH)
TC = np.where(F, get.TCfrompHfCO2(PH, FC, totals, Ks), TC)
TC = np.where(Icase == 78, CO2 + HCO3 + CARB, TC)
TA = np.where(F, get.TAfromTCpH(TC, PH, totals, Ks) + PengCx, TA)
F = Icase == 67 # input CO3, HCO3
if np.any(F):
FC = np.where(F, get.fCO2fromCarbHCO3(CARB, HCO3, totals, Ks), FC)
PH = np.where(F, get.pHfromfCO2Carb(FC, CARB, totals, Ks), PH)
TC = np.where(F, get.TCfrompHfCO2(PH, FC, totals, Ks), TC)
TA = np.where(F, get.TAfromTCpH(TC, PH, totals, Ks) + PengCx, TA)
# By now, an fCO2 value is available for each sample.
# Generate the associated pCO2 and CO2(aq) values:
PC = np.where(~PCgiven, convert.fCO2_to_pCO2(FC, Ks), PC)
# CO2 = np.where(~CO2given, FC * K0, CO2) # up to v1.6.0
CO2 = np.where(~CO2given, TC - CARB - HCO3, CO2) # v1.7.0 onwards
XC = np.where(~XCgiven, convert.fCO2_to_xCO2(FC, Ks), XC) # added in v1.7.0
# ^this assumes pTot = 1 atm
return TA, TC, PH, PC, FC, CARB, HCO3, CO2, XC
def others(
core_solved, TempC, Pdbar, totals, Ks, pHScale, WhichKs, buffers_mode,
):
"""Calculate all peripheral marine carbonate system variables returned by CO2SYS."""
# Unpack for convenience
Sal = totals["Sal"]
TA = core_solved["TA"]
TC = core_solved["TC"]
PH = core_solved["PH"]
PC = core_solved["PC"]
FC = core_solved["FC"]
CARB = core_solved["CARB"]
HCO3 = core_solved["HCO3"]
CO2 = core_solved["CO2"]
# Apply Peng correction
TAPeng = TA - totals["PengCorrection"]
# Calculate pKs
pK1 = -np.log10(Ks["K1"])
pK2 = -np.log10(Ks["K2"])
# Components of alkalinity and DIC
# alks = get.AlkParts(TC, PH, totals, Ks) # <=1.5.1
sw = get.speciation_func(TC, PH, totals, Ks) # >=1.6.0
sw["PAlk"] = sw["PAlk"] + totals["PengCorrection"]
# CaCO3 solubility
OmegaCa, OmegaAr = solubility.CaCO3(CARB, totals, Ks)
# Just for reference, convert pH at input conditions to the other scales
pHT, pHS, pHF, pHN = convert.pH_to_all_scales(PH, pHScale, totals, Ks)
# Get buffers as and if requested
assert np.all(
np.isin(buffers_mode, ["auto", "explicit", "none"])
), "Valid options for buffers_mode are 'auto', 'explicit' or 'none'."
isoQx = np.full(np.shape(Sal), np.nan)
isoQ = np.full(np.shape(Sal), np.nan)
Revelle = np.full(np.shape(Sal), np.nan)
psi = np.full(np.shape(Sal), np.nan)
esm10buffers = [
"gammaTC",
"betaTC",
"omegaTC",
"gammaTA",
"betaTA",
"omegaTA",
]
allbuffers_ESM10 = {
buffer: np.full(np.shape(Sal), np.nan) for buffer in esm10buffers
}
F = buffers_mode == "auto"
if np.any(F):
# Evaluate buffers with automatic differentiation [added v1.3.0]
auto_ESM10 = buffers.all_ESM10(
TAPeng,
TC,
PH,
CARB,
Sal,
convert.TempC2K(TempC),
convert.Pdbar2bar(Pdbar),
totals,
Ks,
WhichKs,
)
for buffer in esm10buffers:
allbuffers_ESM10[buffer] = np.where(
F, auto_ESM10[buffer], allbuffers_ESM10[buffer]
)
isoQ = np.where(F, buffers.isocap(TAPeng, TC, PH, FC, totals, Ks), isoQ)
Revelle = np.where(
F, buffers.RevelleFactor_ESM10(TC, allbuffers_ESM10["gammaTC"]), Revelle
)
F = buffers_mode == "explicit"
if np.any(F):
# Evaluate buffers with explicit equations, but these don't include nutrients
# (i.e. only carbonate, borate and water alkalinities are accounted for)
expl_ESM10 = buffers.explicit.all_ESM10(
TC, TAPeng, CO2, HCO3, CARB, PH, sw["OH"], sw["BAlk"], Ks["KB"],
)
for buffer in esm10buffers:
allbuffers_ESM10[buffer] = np.where(
F, expl_ESM10[buffer], allbuffers_ESM10[buffer]
)
isoQ = np.where(
F,
buffers.explicit.isocap(
CO2, PH, Ks["K1"], Ks["K2"], Ks["KB"], Ks["KW"], totals["TB"]
),
isoQ,
)
Revelle = np.where(
F, buffers.explicit.RevelleFactor(TAPeng, TC, totals, Ks), Revelle
)
F = buffers_mode != "none"
if np.any(F):
# Approximate isocapnic quotient of HDW18
isoQx = np.where(
F,
buffers.explicit.isocap_approx(TC, PC, Ks["K0"], Ks["K1"], Ks["K2"]),
isoQx,
)
# psi of FCG94 following HDW18
psi = np.where(F, buffers.psi(isoQ), psi)
# Substrate:inhibitor ratio of B15
SIR = bio.SIratio(HCO3, pHF)
others_out = {
"pK1": pK1,
"pK2": pK2,
"OmegaCa": OmegaCa,
"OmegaAr": OmegaAr,
"pHT": pHT,
"pHS": pHS,
"pHF": pHF,
"pHN": pHN,
"Revelle": Revelle,
"gammaTC": allbuffers_ESM10["gammaTC"],
"betaTC": allbuffers_ESM10["betaTC"],
"omegaTC": allbuffers_ESM10["omegaTC"],
"gammaTA": allbuffers_ESM10["gammaTA"],
"betaTA": allbuffers_ESM10["betaTA"],
"omegaTA": allbuffers_ESM10["omegaTA"],
"isoQ": isoQ,
"isoQx": isoQx,
"psi": psi,
# Added in v1.4.0:
"SIR": SIR,
}
# Added in v1.6.0:
others_out.update(sw)
return others_out
def assemble(
salinity=35,
temperature=25,
pressure=0,
temperature_out=None,
pressure_out=None,
total_ammonia=0,
total_phosphate=0,
total_silicate=0,
total_sulfide=0,
total_borate=None,
total_calcium=None,
total_fluoride=None,
total_sulfate=None,
opt_gas_constant=3,
opt_k_bisulfate=1,
opt_k_carbonic=16,
opt_k_fluoride=1,
opt_pH_scale=1,
opt_total_borate=1,
k_ammonia=None,
k_borate=None,
k_bisulfate=None,
k_CO2=None,
k_carbonic_1=None,
k_carbonic_2=None,
k_fluoride=None,
k_phosphoric_1=None,
k_phosphoric_2=None,
k_phosphoric_3=None,
k_silicate=None,
k_sulfide=None,
k_water=None,
k_calcite=None,
k_aragonite=None,
fugacity_factor=None,
vp_factor=None,
gas_constant=None,
gas_constant_out=None,
total_alpha=None,
k_alpha=None,
total_beta=None,
k_beta=None,
):
args = condition(locals())
# Prepare totals dict
totals_optional = {
"total_borate": "TB",
"total_calcium": "TCa",
"total_fluoride": "TF",
"total_sulfate": "TSO4",
"total_alpha": "total_alpha",
"total_beta": "total_beta",
}
if np.any(np.isin(list(args.keys()), list(totals_optional.keys()))):
totals = {
totals_optional[k]: v * 1e-6
for k, v in args.items() if k in totals_optional
}
else:
totals = None
totals = salts.assemble(
args["salinity"],
args["total_silicate"],
args["total_phosphate"],
args["total_ammonia"],
args["total_sulfide"],
args["opt_k_carbonic"],
args["opt_total_borate"],
totals=totals,
)
# Prepare equilibrium constants dict (input conditions)
k_constants_optional = {
"fugacity_factor": "FugFac",
"vp_factor": "VPFac",
"gas_constant": "RGas",
"k_ammonia": "KNH3",
"k_borate": "KB",
"k_bisulfate": "KSO4",
"k_CO2": "K0",
"k_carbonic_1": "K1",
"k_carbonic_2": "K2",
"k_fluoride": "KF",
"k_phosphoric_1": "KP1",
"k_phosphoric_2": "KP2",
"k_phosphoric_3": "KP3",
"k_silicate": "KSi",
"k_sulfide": "KH2S",
"k_water": "KW",
"k_calcite": "KCa",
"k_aragonite": "KAr",
"k_alpha": "k_alpha",
"k_beta": "k_beta",
}
if np.any(np.isin(list(args.keys()), list(k_constants_optional.keys()))):
k_constants = {
k_constants_optional[k]: v
for k, v in args.items() if k in k_constants_optional
}
else:
k_constants = None
k_constants = equilibria.assemble(
args["temperature"],
args["pressure"],
totals,
args["opt_pH_scale"],
args["opt_k_carbonic"],
args["opt_k_bisulfate"],
args["opt_k_fluoride"],
args["opt_gas_constant"],
Ks=k_constants,
)
def CO2SYS(
par1=None,
par2=None,
par1_type=None,
par2_type=None,
salinity=35,
temperature=25,
pressure=0,
temperature_out=None,
pressure_out=None,
total_ammonia=0,
total_phosphate=0,
total_silicate=0,
total_sulfide=0,
total_borate=None,
total_calcium=None,
total_fluoride=None,
total_sulfate=None,
opt_gas_constant=3,
opt_k_bisulfate=1,
opt_k_carbonic=16,
opt_k_fluoride=1,
opt_pH_scale=1,
opt_total_borate=1,
buffers_mode="auto",
k_ammonia=None,
k_ammonia_out=None,
k_borate=None,
k_borate_out=None,
k_bisulfate=None,
k_bisulfate_out=None,
k_CO2=None,
k_CO2_out=None,
k_carbonic_1=None,
k_carbonic_1_out=None,
k_carbonic_2=None,
k_carbonic_2_out=None,
k_fluoride=None,
k_fluoride_out=None,
k_phosphoric_1=None,
k_phosphoric_1_out=None,
k_phosphoric_2=None,
k_phosphoric_2_out=None,
k_phosphoric_3=None,
k_phosphoric_3_out=None,
k_silicate=None,
k_silicate_out=None,
k_sulfide=None,
k_sulfide_out=None,
k_water=None,
k_water_out=None,
k_calcite=None,
k_calcite_out=None,
k_aragonite=None,
k_aragonite_out=None,
fugacity_factor=None,
fugacity_factor_out=None,
gas_constant=None,
gas_constant_out=None,
# Added in v1.6.0:
total_alpha=None,
k_alpha=None,
k_alpha_out=None,
total_beta=None,
k_beta=None,
k_beta_out=None,
# Added in v1.7.0:
vp_factor=None,
vp_factor_out=None,
):
"""Run CO2SYS with n-dimensional args allowed."""
args = condition(locals())
# Prepare totals dict
totals_optional = {
"total_borate": "TB",
"total_calcium": "TCa",
"total_fluoride": "TF",
"total_sulfate": "TSO4",
"total_alpha": "total_alpha",
"total_beta": "total_beta",
}
if np.any(np.isin(list(args.keys()), list(totals_optional.keys()))):
totals = {
totals_optional[k]: v * 1e-6
for k, v in args.items() if k in totals_optional
}
else:
totals = None
totals = salts.assemble(
args["salinity"],
args["total_silicate"],
args["total_phosphate"],
args["total_ammonia"],
args["total_sulfide"],
args["opt_k_carbonic"],
args["opt_total_borate"],
totals=totals,
)
# Prepare equilibrium constants dict (input conditions)
k_constants_optional = {
"fugacity_factor": "FugFac",
"vp_factor": "VPFac",
"gas_constant": "RGas",
"k_ammonia": "KNH3",
"k_borate": "KB",
"k_bisulfate": "KSO4",
"k_CO2": "K0",
"k_carbonic_1": "K1",
"k_carbonic_2": "K2",
"k_fluoride": "KF",
"k_phosphoric_1": "KP1",
"k_phosphoric_2": "KP2",
"k_phosphoric_3": "KP3",
"k_silicate": "KSi",
"k_sulfide": "KH2S",
"k_water": "KW",
"k_calcite": "KCa",
"k_aragonite": "KAr",
"k_alpha": "k_alpha",
"k_beta": "k_beta",
}
if np.any(np.isin(list(args.keys()), list(k_constants_optional.keys()))):
k_constants_in = {
k_constants_optional[k]: v
for k, v in args.items() if k in k_constants_optional
}
else:
k_constants_in = None
k_constants_in = equilibria.assemble(
args["temperature"],
args["pressure"],
totals,
args["opt_pH_scale"],
args["opt_k_carbonic"],
args["opt_k_bisulfate"],
args["opt_k_fluoride"],
args["opt_gas_constant"],
Ks=k_constants_in,
)
# Solve the core marine carbonate system at input conditions, if provided
if par1 is not None:
assert par1_type is not None, "PyCO2SYS error: you must provide par1_type."
if par2 is not None:
assert par2_type is not None, "PyCO2SYS error: you must provide par2_type."
if par1 is not None and par2 is not None:
core_in = solve.core(
args["par1"],
args["par2"],
args["par1_type"],
args["par2_type"],
totals,
k_constants_in,
convert_units=True,
)
# Calculate the rest at input conditions
others_in = solve.others(
core_in,
args["temperature"],
args["pressure"],
totals,
k_constants_in,
args["opt_pH_scale"],
args["opt_k_carbonic"],
args["buffers_mode"],
)
elif par1 is not None and par2 is None:
core_in = {}
others_in = {}
# pH only
if np.any(args["par1_type"] == 3):
core_in.update(
{"PH": np.where(args["par1_type"] == 3, args["par1"], np.nan)})
pH_total, pH_sws, pH_free, pH_nbs = convert.pH_to_all_scales(
core_in["PH"], args["opt_pH_scale"], totals, k_constants_in)
others_in.update({
"pHT": pH_total,
"pHS": pH_sws,
"pHF": pH_free,
"pHN": pH_nbs,
})
# One of pCO2, fCO2, CO2(aq) or xCO2 only
if np.any(np.isin(args["par1_type"], [4, 5, 8, 9])):
fCO2 = (np.where(
args["par1_type"] == 5,
args["par1"],
np.where(
args["par1_type"] == 4,
convert.pCO2_to_fCO2(args["par1"], k_constants_in),
np.where(
args["par1_type"] == 8,
convert.CO2aq_to_fCO2(args["par1"], k_constants_in),
np.where(
args["par1_type"] == 9,
convert.xCO2_to_fCO2(args["par1"], k_constants_in),
np.nan,
),
),
),
) * 1e-6)
pCO2 = np.where(
args["par1_type"] == 4,
args["par1"] * 1e-6,
convert.fCO2_to_pCO2(fCO2, k_constants_in),
)
CO2aq = np.where(
args["par1_type"] == 8,
args["par1"] * 1e-6,
convert.fCO2_to_CO2aq(fCO2, k_constants_in),
)
xCO2 = np.where(
args["par1_type"] == 9,
args["par1"] * 1e-6,
convert.fCO2_to_xCO2(fCO2, k_constants_in),
)
core_in.update({
"PC": pCO2,
"FC": fCO2,
"CO2": CO2aq,
"XC": xCO2,
})
else:
core_in = None
others_in = None
# If requested, solve the core marine carbonate system at output conditions
if "pressure_out" in args or "temperature_out" in args:
# Make sure we've got output values for both temperature and pressure
if "pressure_out" in args:
if "temperature_out" not in args:
args["temperature_out"] = args["temperature"]
if "temperature_out" in args:
if "pressure_out" not in args:
args["pressure_out"] = args["pressure"]
# Prepare equilibrium constants dict (output conditions)
k_constants_optional_out = {
"{}_out".format(k): v
for k, v in k_constants_optional.items()
}
if np.any(
np.isin(list(args.keys()),
list(k_constants_optional_out.keys()))):
k_constants_out = {
k_constants_optional_out[k]: v
for k, v in args.items() if k in k_constants_optional_out
}
else:
k_constants_out = None
k_constants_out = equilibria.assemble(
args["temperature_out"],
args["pressure_out"],
totals,
args["opt_pH_scale"],
args["opt_k_carbonic"],
args["opt_k_bisulfate"],
args["opt_k_fluoride"],
args["opt_gas_constant"],
Ks=k_constants_out,
)
# Solve the core marine carbonate system at output conditions, if requested
if par1 is not None and par2 is not None:
core_out = solve.core(
core_in["TA"],
core_in["TC"],
1,
2,
totals,
k_constants_out,
convert_units=False,
)
# Calculate the rest at output conditions
others_out = solve.others(
core_out,
args["temperature_out"],
args["pressure_out"],
totals,
k_constants_out,
args["opt_pH_scale"],
args["opt_k_carbonic"],
args["buffers_mode"],
)
elif par1 is not None and par2 is None:
core_out = {}
others_out = {}
# One of pCO2, fCO2, CO2(aq) or xCO2 only
if np.any(np.isin(args["par1_type"], [4, 5, 8, 9])):
# Takahashi et al. (2009) DSR2 Eq. 2
core_out["PC"] = core_in["PC"] * np.exp(
0.0433 * (temperature_out - temperature) - 4.35e-5 *
(temperature_out**2 - temperature**2))
core_out["FC"] = convert.pCO2_to_fCO2(core_out["PC"],
k_constants_out)
core_out["CO2"] = convert.fCO2_to_CO2aq(
core_out["FC"], k_constants_out)
core_out["XC"] = convert.fCO2_to_xCO2(core_out["FC"],
k_constants_out)
else:
core_out = None
others_out = None
else:
core_out = None
others_out = None
k_constants_out = None
return _get_results_dict(
args,
totals,
core_in,
others_in,
k_constants_in,
core_out,
others_out,
k_constants_out,
)
def dcore_dparX__parY(parXtype, parYtype, TA, TC, PH, FC, CARB, HCO3, totals,
Ks):
"""Efficient automatic derivatives of all core MCS variables w.r.t. parX
at constant parY.
"""
K0 = Ks["K0"] # alias for convenience
# Get necessary derivatives
Ucase = 10 * parXtype + parYtype # like Icase, but not sorted
# Derivatives that are used by multiple Ucases
if np.any(isin(Ucase, [12, 32, 42, 52, 62, 72, 82])):
dTA_dPH__TC = egrad(lambda PH: get.TAfromTCpH(TC, PH, totals, Ks))(PH)
dFC_dPH__TC = egrad(lambda PH: get.fCO2fromTCpH(TC, PH, totals, Ks))(
PH)
dCARB_dPH__TC = egrad(lambda PH: get.CarbfromTCpH(TC, PH, totals, Ks))(
PH)
dHCO3_dPH__TC = egrad(lambda PH: get.HCO3fromTCpH(TC, PH, totals, Ks))(
PH)
if np.any(isin(Ucase, [21, 31, 41, 51, 61, 71, 81])):
dTC_dPH__TA = egrad(lambda PH: get.TCfromTApH(TA, PH, totals, Ks))(PH)
dFC_dPH__TA = egrad(lambda PH: get.fCO2fromTApH(TA, PH, totals, Ks))(
PH)
dCARB_dPH__TA = egrad(lambda PH: get.CarbfromTApH(TA, PH, totals, Ks))(
PH)
dHCO3_dPH__TA = egrad(lambda PH: get.HCO3fromTApH(TA, PH, totals, Ks))(
PH)
if np.any(isin(Ucase, [16, 26, 36])):
dTC_dPH__CARB = egrad(
lambda PH: get.TCfrompHCarb(PH, CARB, totals, Ks))(PH)
dTA_dPH__CARB = egrad(
lambda PH: get.TAfrompHCarb(PH, CARB, totals, Ks))(PH)
dFC_dPH__CARB = egrad(
lambda PH: get.fCO2frompHCarb(PH, CARB, totals, Ks))(PH)
dHCO3_dPH__CARB = egrad(
lambda PH: get.HCO3frompHCarb(PH, CARB, totals, Ks))(PH)
if np.any(isin(Ucase, [17, 27, 37])):
dTC_dPH__HCO3 = egrad(
lambda PH: get.TCfrompHHCO3(PH, HCO3, totals, Ks))(PH)
dTA_dPH__HCO3 = egrad(
lambda PH: get.TAfrompHHCO3(PH, HCO3, totals, Ks))(PH)
dFC_dPH__HCO3 = egrad(
lambda PH: get.fCO2frompHHCO3(PH, HCO3, totals, Ks))(PH)
dCARB_dPH__HCO3 = egrad(
lambda PH: get.CarbfrompHHCO3(PH, HCO3, totals, Ks))(PH)
if np.any(isin(Ucase, [14, 15, 18, 24, 25, 28, 34, 35, 38])):
dTA_dPH__FC = egrad(lambda PH: get.TAfrompHfCO2(PH, FC, totals, Ks))(
PH)
dTC_dPH__FC = egrad(lambda PH: get.TCfrompHfCO2(PH, FC, totals, Ks))(
PH)
dCARB_dPH__FC = egrad(
lambda PH: get.CarbfrompHfCO2(PH, FC, totals, Ks))(PH)
dHCO3_dPH__FC = egrad(
lambda PH: get.HCO3frompHfCO2(PH, FC, totals, Ks))(PH)
# Derivatives specific to a single Ucase
if np.any((parXtype == 1) & (parYtype == 2)): # dvar_dTA__TC
dPH_dTA__TC = 1 / dTA_dPH__TC
dFC_dTA__TC = dFC_dPH__TC / dTA_dPH__TC
dCARB_dTA__TC = dCARB_dPH__TC / dTA_dPH__TC
dHCO3_dTA__TC = dHCO3_dPH__TC / dTA_dPH__TC
if np.any((parXtype == 1) & (parYtype == 3)): # dvar_dTA__PH
dTC_dTA__PH = egrad(lambda TA: get.TCfromTApH(TA, PH, totals, Ks))(TA)
dFC_dTA__PH = egrad(lambda TA: get.fCO2fromTApH(TA, PH, totals, Ks))(
TA)
dCARB_dTA__PH = egrad(lambda TA: get.CarbfromTApH(TA, PH, totals, Ks))(
TA)
dHCO3_dTA__PH = egrad(lambda TA: get.HCO3fromTApH(TA, PH, totals, Ks))(
TA)
if np.any((parXtype == 1) & isin(parYtype, [4, 5, 8])): # dvar_dTA__FC
dTC_dTA__FC = dTC_dPH__FC / dTA_dPH__FC
dPH_dTA__FC = 1 / dTA_dPH__FC
dCARB_dTA__FC = dCARB_dPH__FC / dTA_dPH__FC
dHCO3_dTA__FC = dHCO3_dPH__FC / dTA_dPH__FC
if np.any((parXtype == 1) & (parYtype == 6)): # dvar_dTA__CARB
dTC_dTA__CARB = dTC_dPH__CARB / dTA_dPH__CARB
dPH_dTA__CARB = 1 / dTA_dPH__CARB
dFC_dTA__CARB = dFC_dPH__CARB / dTA_dPH__CARB
dHCO3_dTA__CARB = dHCO3_dPH__CARB / dTA_dPH__CARB
if np.any((parXtype == 1) & (parYtype == 7)): # dvar_dTA__HCO3
dTC_dTA__HCO3 = dTC_dPH__HCO3 / dTA_dPH__HCO3
dPH_dTA__HCO3 = 1 / dTA_dPH__HCO3
dFC_dTA__HCO3 = dFC_dPH__HCO3 / dTA_dPH__HCO3
dCARB_dTA__HCO3 = dCARB_dPH__HCO3 / dTA_dPH__HCO3
if np.any((parXtype == 2) & (parYtype == 1)): # dvar_dTC__TA
dPH_dTC__TA = 1 / dTC_dPH__TA
dFC_dTC__TA = dFC_dPH__TA / dTC_dPH__TA
dCARB_dTC__TA = dCARB_dPH__TA / dTC_dPH__TA
dHCO3_dTC__TA = dHCO3_dPH__TA / dTC_dPH__TA
if np.any((parXtype == 2) & (parYtype == 3)): # dvar_dTC__PH
dTA_dTC__PH = egrad(lambda TC: get.TAfromTCpH(TC, PH, totals, Ks))(TC)
dFC_dTC__PH = egrad(lambda TC: get.fCO2fromTCpH(TC, PH, totals, Ks))(
TC)
dCARB_dTC__PH = egrad(lambda TC: get.CarbfromTCpH(TC, PH, totals, Ks))(
TC)
dHCO3_dTC__PH = egrad(lambda TC: get.HCO3fromTCpH(TC, PH, totals, Ks))(
TC)
if np.any((parXtype == 2) & isin(parYtype, [4, 5, 8])): # dvar_dTC__FC
dTA_dTC__FC = dTA_dPH__FC / dTC_dPH__FC
dPH_dTC__FC = 1 / dTC_dPH__FC
dCARB_dTC__FC = dCARB_dPH__FC / dTC_dPH__FC
dHCO3_dTC__FC = dHCO3_dPH__FC / dTC_dPH__FC
if np.any((parXtype == 2) & (parYtype == 6)): # dvar_dTC__CARB
dTA_dTC__CARB = dTA_dPH__CARB / dTC_dPH__CARB
dPH_dTC__CARB = 1 / dTC_dPH__CARB
dFC_dTC__CARB = dFC_dPH__CARB / dTC_dPH__CARB
dHCO3_dTC__CARB = dHCO3_dPH__CARB / dTC_dPH__CARB
if np.any((parXtype == 2) & (parYtype == 7)): # dvar_dTC__HCO3
dTA_dTC__HCO3 = dTA_dPH__HCO3 / dTC_dPH__HCO3
dPH_dTC__HCO3 = 1 / dTC_dPH__HCO3
dFC_dTC__HCO3 = dFC_dPH__HCO3 / dTC_dPH__HCO3
dCARB_dTC__HCO3 = dCARB_dPH__HCO3 / dTC_dPH__HCO3
if np.any(isin(parXtype, [4, 5, 8]) & (parYtype == 1)): # dvar_dFC__TA
dTC_dFC__TA = dTC_dPH__TA / dFC_dPH__TA
dPH_dFC__TA = 1 / dFC_dPH__TA
dCARB_dFC__TA = dCARB_dPH__TA / dFC_dPH__TA
dHCO3_dFC__TA = dHCO3_dPH__TA / dFC_dPH__TA
if np.any(isin(parXtype, [4, 5, 8]) & (parYtype == 2)): # dvar_dFC__TC
dTA_dFC__TC = dTA_dPH__TC / dFC_dPH__TC
dPH_dFC__TC = 1 / dFC_dPH__TC
dCARB_dFC__TC = dCARB_dPH__TC / dFC_dPH__TC
dHCO3_dFC__TC = dHCO3_dPH__TC / dFC_dPH__TC
if np.any(isin(parXtype, [4, 5, 8]) & (parYtype == 3)): # dvar_dFC__PH
dTA_dFC__PH = egrad(lambda FC: get.TAfrompHfCO2(PH, FC, totals, Ks))(
FC)
dTC_dFC__PH = egrad(lambda FC: get.TCfrompHfCO2(PH, FC, totals, Ks))(
FC)
dCARB_dFC__PH = egrad(
lambda FC: get.CarbfrompHfCO2(PH, FC, totals, Ks))(FC)
dHCO3_dFC__PH = egrad(
lambda FC: get.HCO3frompHfCO2(PH, FC, totals, Ks))(FC)
if np.any(isin(parXtype, [4, 5, 8]) & (parYtype == 6)): # dvar_dFC__CARB
dTA_dFC__CARB = egrad(
lambda FC: get.TAfromfCO2Carb(FC, CARB, totals, Ks))(FC)
dTC_dFC__CARB = egrad(
lambda FC: get.TCfromfCO2Carb(FC, CARB, totals, Ks))(FC)
dPH_dFC__CARB = egrad(
lambda FC: get.pHfromfCO2Carb(FC, CARB, totals, Ks))(FC)
dHCO3_dFC__CARB = egrad(
lambda FC: get.HCO3fromfCO2Carb(FC, CARB, totals, Ks))(FC)
if np.any(isin(parXtype, [4, 5, 8]) & (parYtype == 7)): # dvar_dFC__HCO3
dTA_dFC__HCO3 = egrad(
lambda FC: get.TAfromfCO2HCO3(FC, HCO3, totals, Ks))(FC)
dTC_dFC__HCO3 = egrad(
lambda FC: get.TCfromfCO2HCO3(FC, HCO3, totals, Ks))(FC)
dPH_dFC__HCO3 = egrad(
lambda FC: get.pHfromfCO2HCO3(FC, HCO3, totals, Ks))(FC)
dCARB_dFC__HCO3 = egrad(
lambda FC: get.CarbfromfCO2HCO3(FC, HCO3, totals, Ks))(FC)
if np.any((parXtype == 6) & (parYtype == 1)): # dvar_dCARB__TA
dTC_dCARB__TA = dTC_dPH__TA / dCARB_dPH__TA
dPH_dCARB__TA = 1 / dCARB_dPH__TA
dFC_dCARB__TA = dFC_dPH__TA / dCARB_dPH__TA
dHCO3_dCARB__TA = dHCO3_dPH__TA / dCARB_dPH__TA
if np.any((parXtype == 6) & (parYtype == 2)): # dvar_dCARB__TC
dTA_dCARB__TC = dTA_dPH__TC / dCARB_dPH__TC
dPH_dCARB__TC = 1 / dCARB_dPH__TC
dFC_dCARB__TC = dFC_dPH__TC / dCARB_dPH__TC
dHCO3_dCARB__TC = dHCO3_dPH__TC / dCARB_dPH__TC
if np.any((parXtype == 6) & (parYtype == 3)): # dvar_dCARB__PH
dTA_dCARB__PH = egrad(
lambda CARB: get.TAfrompHCarb(PH, CARB, totals, Ks))(CARB)
dTC_dCARB__PH = egrad(
lambda CARB: get.TCfrompHCarb(PH, CARB, totals, Ks))(CARB)
dFC_dCARB__PH = egrad(
lambda CARB: get.fCO2frompHCarb(PH, CARB, totals, Ks))(CARB)
dHCO3_dCARB__PH = egrad(
lambda CARB: get.HCO3frompHCarb(PH, CARB, totals, Ks))(CARB)
if np.any((parXtype == 6) & isin(parYtype, [4, 5, 8])): # dvar_dCARB__FC
dTA_dCARB__FC = egrad(
lambda CARB: get.TAfromfCO2Carb(FC, CARB, totals, Ks))(CARB)
dTC_dCARB__FC = egrad(
lambda CARB: get.TCfromfCO2Carb(FC, CARB, totals, Ks))(CARB)
dPH_dCARB__FC = egrad(
lambda CARB: get.pHfromfCO2Carb(FC, CARB, totals, Ks))(CARB)
dHCO3_dCARB__FC = egrad(
lambda CARB: get.HCO3fromfCO2Carb(FC, CARB, totals, Ks))(CARB)
if np.any((parXtype == 6) & (parYtype == 7)): # dvar_dCARB__HCO3
dTA_dCARB__HCO3 = egrad(
lambda CARB: get.TAfromCarbHCO3(CARB, HCO3, totals, Ks))(CARB)
dTC_dCARB__HCO3 = egrad(
lambda CARB: get.TCfromCarbHCO3(CARB, HCO3, totals, Ks))(CARB)
dPH_dCARB__HCO3 = egrad(
lambda CARB: get.pHfromCarbHCO3(CARB, HCO3, totals, Ks))(CARB)
dFC_dCARB__HCO3 = egrad(
lambda CARB: get.fCO2fromCarbHCO3(CARB, HCO3, totals, Ks))(CARB)
if np.any((parXtype == 7) & (parYtype == 1)): # dvar_dHCO3__TA
dTC_dHCO3__TA = dTC_dPH__TA / dHCO3_dPH__TA
dPH_dHCO3__TA = 1 / dHCO3_dPH__TA
dFC_dHCO3__TA = dFC_dPH__TA / dHCO3_dPH__TA
dCARB_dHCO3__TA = dCARB_dPH__TA / dHCO3_dPH__TA
if np.any((parXtype == 7) & (parYtype == 2)): # dvar_dHCO3__TC
dTA_dHCO3__TC = dTA_dPH__TC / dHCO3_dPH__TC
dPH_dHCO3__TC = 1 / dHCO3_dPH__TC
dFC_dHCO3__TC = dFC_dPH__TC / dHCO3_dPH__TC
dCARB_dHCO3__TC = dCARB_dPH__TC / dHCO3_dPH__TC
if np.any((parXtype == 7) & (parYtype == 3)): # dvar_dHCO3__PH
dTC_dHCO3__PH = egrad(
lambda HCO3: get.TCfrompHHCO3(PH, HCO3, totals, Ks))(HCO3)
dTA_dHCO3__PH = egrad(
lambda HCO3: get.TAfrompHHCO3(PH, HCO3, totals, Ks))(HCO3)
dFC_dHCO3__PH = egrad(
lambda HCO3: get.fCO2frompHHCO3(PH, HCO3, totals, Ks))(HCO3)
dCARB_dHCO3__PH = egrad(
lambda HCO3: get.CarbfrompHHCO3(PH, HCO3, totals, Ks))(HCO3)
if np.any((parXtype == 7) & isin(parYtype, [4, 5, 8])): # dvar_dHCO3__FC
dTA_dHCO3__FC = egrad(
lambda HCO3: get.TAfromfCO2HCO3(FC, HCO3, totals, Ks))(HCO3)
dTC_dHCO3__FC = egrad(
lambda HCO3: get.TCfromfCO2HCO3(FC, HCO3, totals, Ks))(HCO3)
dPH_dHCO3__FC = egrad(
lambda HCO3: get.pHfromfCO2HCO3(FC, HCO3, totals, Ks))(HCO3)
dCARB_dHCO3__FC = egrad(
lambda HCO3: get.CarbfromfCO2HCO3(FC, HCO3, totals, Ks))(HCO3)
if np.any((parXtype == 7) & (parYtype == 6)): # dvar_dHCO3__CARB
dTA_dHCO3__CARB = egrad(
lambda HCO3: get.TAfromCarbHCO3(CARB, HCO3, totals, Ks))(HCO3)
dTC_dHCO3__CARB = egrad(
lambda HCO3: get.TCfromCarbHCO3(CARB, HCO3, totals, Ks))(HCO3)
dPH_dHCO3__CARB = egrad(
lambda HCO3: get.pHfromCarbHCO3(CARB, HCO3, totals, Ks))(HCO3)
dFC_dHCO3__CARB = egrad(
lambda HCO3: get.fCO2fromCarbHCO3(CARB, HCO3, totals, Ks))(HCO3)
# Preallocate empty arrays for derivatives
dTA_dX__Y = np.full(np.shape(parXtype), nan)
dTC_dX__Y = np.full(np.shape(parXtype), nan)
dPH_dX__Y = np.full(np.shape(parXtype), nan)
dFC_dX__Y = np.full(np.shape(parXtype), nan)
dCARB_dX__Y = np.full(np.shape(parXtype), nan)
dHCO3_dX__Y = np.full(np.shape(parXtype), nan)
# Assign derivatives
X = parXtype == 1 # TA - total alkalinity
if np.any(X):
dTA_dX__Y = where(X, 1.0, dTA_dX__Y)
XY = X & (parYtype == 2) # TA, TC
if np.any(XY):
dTC_dX__Y = where(XY, 0.0, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dTA__TC, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dTA__TC, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dTA__TC, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dTA__TC, dHCO3_dX__Y)
XY = X & (parYtype == 3) # TA, PH
if np.any(XY):
dTC_dX__Y = where(XY, dTC_dTA__PH, dTC_dX__Y)
dPH_dX__Y = where(XY, 0.0, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dTA__PH, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dTA__PH, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dTA__PH, dHCO3_dX__Y)
XY = X & isin(parYtype, [4, 5, 8]) # TA, (PC | FC | CO2)
if np.any(XY):
dTC_dX__Y = where(XY, dTC_dTA__FC, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dTA__FC, dPH_dX__Y)
dFC_dX__Y = where(XY, 0.0, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dTA__FC, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dTA__FC, dHCO3_dX__Y)
XY = X & (parYtype == 6) # TA, CARB
if np.any(XY):
dTC_dX__Y = where(XY, dTC_dTA__CARB, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dTA__CARB, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dTA__CARB, dFC_dX__Y)
dCARB_dX__Y = where(XY, 0.0, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dTA__CARB, dHCO3_dX__Y)
XY = X & (parYtype == 7) # TA, HCO3
if np.any(XY):
dTC_dX__Y = where(XY, dTC_dTA__HCO3, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dTA__HCO3, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dTA__HCO3, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dTA__HCO3, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, 0.0, dHCO3_dX__Y)
X = parXtype == 2 # TC - dissolved inorganic carbon
if np.any(X):
dTC_dX__Y = where(X, 1.0, dTC_dX__Y)
XY = X & (parYtype == 1) # TC, TA
if np.any(XY):
dTA_dX__Y = where(XY, 0.0, dTA_dX__Y)
dPH_dX__Y = where(XY, dPH_dTC__TA, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dTC__TA, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dTC__TA, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dTC__TA, dHCO3_dX__Y)
XY = X & (parYtype == 3) # TC, PH
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dTC__PH, dTA_dX__Y)
dPH_dX__Y = where(XY, 0.0, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dTC__PH, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dTC__PH, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dTC__PH, dHCO3_dX__Y)
XY = X & isin(parYtype, [4, 5, 8]) # TC, (PC | FC | CO2)
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dTC__FC, dTA_dX__Y)
dPH_dX__Y = where(XY, dPH_dTC__FC, dPH_dX__Y)
dFC_dX__Y = where(XY, 0.0, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dTC__FC, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dTC__FC, dHCO3_dX__Y)
XY = X & (parYtype == 6) # TC, CARB
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dTC__CARB, dTA_dX__Y)
dPH_dX__Y = where(XY, dPH_dTC__CARB, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dTC__CARB, dFC_dX__Y)
dCARB_dX__Y = where(XY, 0.0, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dTC__CARB, dHCO3_dX__Y)
XY = X & (parYtype == 7) # TC, HCO3
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dTC__HCO3, dTA_dX__Y)
dPH_dX__Y = where(XY, dPH_dTC__HCO3, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dTC__HCO3, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dTC__HCO3, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, 0.0, dHCO3_dX__Y)
X = parXtype == 3 # PH - seawater pH
if np.any(X):
dPH_dX__Y = where(X, 1.0, dPH_dX__Y)
XY = X & (parYtype == 1) # PH, TA
if np.any(XY):
dTA_dX__Y = where(XY, 0.0, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dPH__TA, dTC_dX__Y)
dFC_dX__Y = where(XY, dFC_dPH__TA, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dPH__TA, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dPH__TA, dHCO3_dX__Y)
XY = X & (parYtype == 2) # PH, TC
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dPH__TC, dTA_dX__Y)
dTC_dX__Y = where(XY, 0.0, dTC_dX__Y)
dFC_dX__Y = where(XY, dFC_dPH__TC, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dPH__TC, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dPH__TC, dHCO3_dX__Y)
XY = X & isin(parYtype, [4, 5, 8]) # PH, (PC | FC | CO2)
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dPH__FC, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dPH__FC, dTC_dX__Y)
dFC_dX__Y = where(XY, 0.0, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dPH__FC, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dPH__FC, dHCO3_dX__Y)
XY = X & (parYtype == 6) # PH, CARB
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dPH__CARB, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dPH__CARB, dTC_dX__Y)
dFC_dX__Y = where(XY, dFC_dPH__CARB, dFC_dX__Y)
dCARB_dX__Y = where(XY, 0.0, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dPH__CARB, dHCO3_dX__Y)
XY = X & (parYtype == 7) # PH, HCO3
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dPH__HCO3, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dPH__HCO3, dTC_dX__Y)
dFC_dX__Y = where(XY, dFC_dPH__HCO3, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dPH__HCO3, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, 0.0, dHCO3_dX__Y)
X = isin(parXtype,
[4, 5, 8]) # (PC | FC | CO2) - CO2 fugacity, p.p. or (aq)
if np.any(X):
dFC_dX__Y = where(X, 1.0, dFC_dX__Y)
XY = X & (parYtype == 1) # (PC | FC | CO2), TA
if np.any(XY):
dTA_dX__Y = where(XY, 0.0, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dFC__TA, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dFC__TA, dPH_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dFC__TA, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dFC__TA, dHCO3_dX__Y)
XY = X & (parYtype == 2) # (PC | FC | CO2), TC
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dFC__TC, dTA_dX__Y)
dTC_dX__Y = where(XY, 0.0, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dFC__TC, dPH_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dFC__TC, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dFC__TC, dHCO3_dX__Y)
XY = X & (parYtype == 3) # (PC | FC | CO2), PH
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dFC__PH, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dFC__PH, dTC_dX__Y)
dPH_dX__Y = where(XY, 0.0, dPH_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dFC__PH, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dFC__PH, dHCO3_dX__Y)
XY = X & (parYtype == 6) # (PC | FC | CO2), CARB
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dFC__CARB, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dFC__CARB, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dFC__CARB, dPH_dX__Y)
dCARB_dX__Y = where(XY, 0.0, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dFC__CARB, dHCO3_dX__Y)
XY = X & (parYtype == 7) # (PC | FC | CO2), HCO3
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dFC__HCO3, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dFC__HCO3, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dFC__HCO3, dPH_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dFC__HCO3, dCARB_dX__Y)
dHCO3_dX__Y = where(XY, 0.0, dHCO3_dX__Y)
X = parXtype == 6 # CARB - carbonate ion
if np.any(X):
dCARB_dX__Y = where(X, 1.0, dCARB_dX__Y)
XY = X & (parYtype == 1) # CARB, TA
if np.any(XY):
dTA_dX__Y = where(XY, 0.0, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dCARB__TA, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dCARB__TA, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dCARB__TA, dFC_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dCARB__TA, dHCO3_dX__Y)
XY = X & (parYtype == 2) # CARB, TC
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dCARB__TC, dTA_dX__Y)
dTC_dX__Y = where(XY, 0.0, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dCARB__TC, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dCARB__TC, dFC_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dCARB__TC, dHCO3_dX__Y)
XY = X & (parYtype == 3) # CARB, PH
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dCARB__PH, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dCARB__PH, dTC_dX__Y)
dPH_dX__Y = where(XY, 0.0, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dCARB__PH, dFC_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dCARB__PH, dHCO3_dX__Y)
XY = X & isin(parYtype, [4, 5, 8]) # CARB, (PC | FC | CO2)
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dCARB__FC, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dCARB__FC, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dCARB__FC, dPH_dX__Y)
dFC_dX__Y = where(XY, 0.0, dFC_dX__Y)
dHCO3_dX__Y = where(XY, dHCO3_dCARB__FC, dHCO3_dX__Y)
XY = X & (parYtype == 7) # CARB, HCO3
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dCARB__HCO3, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dCARB__HCO3, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dCARB__HCO3, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dCARB__HCO3, dFC_dX__Y)
dHCO3_dX__Y = where(XY, 0.0, dHCO3_dX__Y)
X = parXtype == 7 # HCO3 - bicarbonate ion
if np.any(X):
dHCO3_dX__Y = where(X, 1.0, dHCO3_dX__Y)
XY = X & (parYtype == 1) # HCO3, TA
if np.any(XY):
dTA_dX__Y = where(XY, 0.0, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dHCO3__TA, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dHCO3__TA, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dHCO3__TA, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dHCO3__TA, dCARB_dX__Y)
XY = X & (parYtype == 2) # HCO3, TC
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dHCO3__TC, dTA_dX__Y)
dTC_dX__Y = where(XY, 0.0, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dHCO3__TC, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dHCO3__TC, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dHCO3__TC, dCARB_dX__Y)
XY = X & (parYtype == 3) # HCO3, PH
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dHCO3__PH, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dHCO3__PH, dTC_dX__Y)
dPH_dX__Y = where(XY, 0.0, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dHCO3__PH, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dHCO3__PH, dCARB_dX__Y)
XY = X & isin(parYtype, [4, 5, 8]) # HCO3, (PC | FC | CO2)
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dHCO3__FC, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dHCO3__FC, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dHCO3__FC, dPH_dX__Y)
dFC_dX__Y = where(XY, 0.0, dFC_dX__Y)
dCARB_dX__Y = where(XY, dCARB_dHCO3__FC, dCARB_dX__Y)
XY = X & (parYtype == 6) # HCO3, CARB
if np.any(XY):
dTA_dX__Y = where(XY, dTA_dHCO3__CARB, dTA_dX__Y)
dTC_dX__Y = where(XY, dTC_dHCO3__CARB, dTC_dX__Y)
dPH_dX__Y = where(XY, dPH_dHCO3__CARB, dPH_dX__Y)
dFC_dX__Y = where(XY, dFC_dHCO3__CARB, dFC_dX__Y)
dCARB_dX__Y = where(XY, 0.0, dCARB_dX__Y)
# Get pCO2 and CO2(aq) derivatives from fCO2
dPC_dX__Y = dFC_dX__Y / Ks["FugFac"]
dCO2_dX__Y = dFC_dX__Y * K0
# Update values where parX or parY were pCO2 or CO2(aq)
X = parXtype == 4 # pCO2
if any(X):
dTA_dX__Y = where(X, dTA_dX__Y * Ks["FugFac"], dTA_dX__Y)
dTC_dX__Y = where(X, dTC_dX__Y * Ks["FugFac"], dTC_dX__Y)
dPH_dX__Y = where(X, dPH_dX__Y * Ks["FugFac"], dPH_dX__Y)
dPC_dX__Y = where(X, dPC_dX__Y * Ks["FugFac"], dPC_dX__Y)
dFC_dX__Y = where(X, dFC_dX__Y * Ks["FugFac"], dFC_dX__Y)
dCARB_dX__Y = where(X, dCARB_dX__Y * Ks["FugFac"], dCARB_dX__Y)
dHCO3_dX__Y = where(X, dHCO3_dX__Y * Ks["FugFac"], dHCO3_dX__Y)
dCO2_dX__Y = where(X, dCO2_dX__Y * Ks["FugFac"], dCO2_dX__Y)
X = parXtype == 8 # CO2(aq)
if any(X):
dTA_dX__Y = where(X, dTA_dX__Y / K0, dTA_dX__Y)
dTC_dX__Y = where(X, dTC_dX__Y / K0, dTC_dX__Y)
dPH_dX__Y = where(X, dPH_dX__Y / K0, dPH_dX__Y)
dPC_dX__Y = where(X, dPC_dX__Y / K0, dPC_dX__Y)
dFC_dX__Y = where(X, dFC_dX__Y / K0, dFC_dX__Y)
dCARB_dX__Y = where(X, dCARB_dX__Y / K0, dCARB_dX__Y)
dHCO3_dX__Y = where(X, dHCO3_dX__Y / K0, dHCO3_dX__Y)
dCO2_dX__Y = where(X, dCO2_dX__Y / K0, dCO2_dX__Y)
return {
"TAlk": dTA_dX__Y,
"TCO2": dTC_dX__Y,
"pH": dPH_dX__Y,
"pCO2": dPC_dX__Y,
"fCO2": dFC_dX__Y,
"CO3": dCARB_dX__Y,
"HCO3": dHCO3_dX__Y,
"CO2": dCO2_dX__Y,
}
def forward_nd(
CO2SYS_nd_results,
grads_of,
grads_wrt,
dx=1e-6,
dx_scaling="median",
dx_func=None,
**CO2SYS_nd_kwargs,
):
"""Get forward finite-difference derivatives of pyco2.sys results with respect to
its arguments.
"""
# Check requested grads are possible
gradables = (engine.nd.gradables + [
"p{}".format(gradable)
for gradable in engine.nd.gradables if gradable.startswith("k_")
] + [
"{}_both".format(gradable) for gradable in engine.nd.gradables
if gradable.startswith("k_") and not gradable.endswith("_out")
] + [
"p{}_both".format(gradable) for gradable in engine.nd.gradables
if gradable.startswith("k_") and not gradable.endswith("_out")
])
assert np.all(
np.isin(grads_of, gradables)
), "PyCO2SYS error: all grads_of must be in the list at PyCO2SYS.engine.nd.gradables."
if np.any([of.endswith("_out") for of in grads_of]):
assert "temperature_out" in CO2SYS_nd_results, (
"PyCO2SYS error: you can only get gradients at output conditions if you calculated"
+ " results at output conditions!")
# Extract CO2SYS_nd fixed args from CO2SYS_nd_results and CO2SYS_nd_kwargs.
# These are arguments that always get a specific value, rather than being calculated
# from e.g. temperature and salinity if not provided.
keys_fixed = set([
"par1",
"par2",
"par1_type",
"par2_type",
"salinity",
"temperature",
"pressure",
"total_ammonia",
"total_phosphate",
"total_silicate",
"total_sulfide",
"opt_gas_constant",
"opt_k_bisulfate",
"opt_k_carbonic",
"opt_k_fluoride",
"opt_pH_scale",
"opt_total_borate",
"buffers_mode",
] + list(CO2SYS_nd_kwargs.keys()))
args_fixed = {
k: CO2SYS_nd_results[k]
for k in keys_fixed if k in CO2SYS_nd_results
}
# Prepare dicts for results
dxs = {wrt: None for wrt in grads_wrt}
CO2SYS_derivs = {of: {wrt: None for wrt in grads_wrt} for of in grads_of}
# Loop through requested parameters and calculate the gradients
for wrt in grads_wrt:
args_plus = copy.deepcopy(args_fixed)
# Check for special cases
is_pk = wrt.startswith("pk_")
do_both = wrt.endswith("_both")
if is_pk:
wrt_as_k = wrt[1:] # remove the "p" prefix
if do_both:
wrt_as_k = wrt_as_k[:-5] # remove the "_both" suffix
# Convert K to pK, increment pK by dx, then convert back to K (input)
pk_values = -np.log10(CO2SYS_nd_results[wrt_as_k])
dxs[wrt] = _get_dx_wrt(dx, pk_values, dx_scaling, dx_func=dx_func)
pk_values_plus = pk_values + dxs[wrt]
args_plus[wrt_as_k] = 10.0**-pk_values_plus
if do_both:
# Convert K to pK, increment pK by dx, then convert back to K (output)
# Uses the same dx value as for the input condition
pk_values_out = -np.log10(CO2SYS_nd_results[wrt_as_k + "_out"])
pk_values_out_plus = pk_values_out + dxs[wrt]
args_plus[wrt_as_k + "_out"] = 10.0**-pk_values_out_plus
else: # if not is_pk
if do_both:
wrt_internal = wrt[:-5] # remove the "_both" suffix
else:
wrt_internal = copy.deepcopy(wrt)
# Get the dx and increment the input argument by it
dxs[wrt] = _get_dx_wrt(dx,
CO2SYS_nd_results[wrt_internal],
dx_scaling,
dx_func=dx_func)
args_plus[
wrt_internal] = CO2SYS_nd_results[wrt_internal] + dxs[wrt]
if do_both:
# Increment the output argument by the same dx as for the input
args_plus[wrt_internal +
"_out"] = (CO2SYS_nd_results[wrt_internal + "_out"] +
dxs[wrt])
# Solve again with the incremented arguments and save output
results_plus = engine.nd.CO2SYS(**args_plus)
for of in grads_of:
CO2SYS_derivs[of][wrt] = (results_plus[of] -
CO2SYS_nd_results[of]) / dxs[wrt]
return CO2SYS_derivs, dxs
def CO2SYS(
par1,
par2,
par1_type,
par2_type,
salinity=35,
temperature=25,
pressure=0,
temperature_out=None,
pressure_out=None,
total_ammonia=0,
total_phosphate=0,
total_silicate=0,
total_sulfide=0,
total_borate=None,
total_calcium=None,
total_fluoride=None,
total_sulfate=None,
opt_gas_constant=3,
opt_k_bisulfate=1,
opt_k_carbonic=16,
opt_k_fluoride=1,
opt_pH_scale=1,
opt_total_borate=1,
buffers_mode="auto",
k_ammonia=None,
k_ammonia_out=None,
k_borate=None,
k_borate_out=None,
k_bisulfate=None,
k_bisulfate_out=None,
k_CO2=None,
k_CO2_out=None,
k_carbonic_1=None,
k_carbonic_1_out=None,
k_carbonic_2=None,
k_carbonic_2_out=None,
k_fluoride=None,
k_fluoride_out=None,
k_phosphate_1=None,
k_phosphate_1_out=None,
k_phosphate_2=None,
k_phosphate_2_out=None,
k_phosphate_3=None,
k_phosphate_3_out=None,
k_silicate=None,
k_silicate_out=None,
k_sulfide=None,
k_sulfide_out=None,
k_water=None,
k_water_out=None,
k_calcite=None,
k_calcite_out=None,
k_aragonite=None,
k_aragonite_out=None,
fugacity_factor=None,
fugacity_factor_out=None,
gas_constant=None,
gas_constant_out=None,
# Added in v1.6.0:
total_alpha=None,
k_alpha=None,
k_alpha_out=None,
total_beta=None,
k_beta=None,
k_beta_out=None,
):
"""Run CO2SYS with n-dimensional args allowed."""
args = condition(locals())
# Prepare totals dict
totals_optional = {
"total_borate": "TB",
"total_calcium": "TCa",
"total_fluoride": "TF",
"total_sulfate": "TSO4",
"total_alpha": "alpha",
"total_beta": "beta",
}
if np.any(np.isin(list(args.keys()), list(totals_optional.keys()))):
totals = {
totals_optional[k]: v * 1e-6
for k, v in args.items()
if k in totals_optional
}
else:
totals = None
totals = salts.assemble(
args["salinity"],
args["total_silicate"],
args["total_phosphate"],
args["total_ammonia"],
args["total_sulfide"],
args["opt_k_carbonic"],
args["opt_total_borate"],
totals=totals,
)
# Prepare equilibrium constants dict (input conditions)
k_constants_optional = {
"fugacity_factor": "FugFac",
"gas_constant": "RGas",
"k_ammonia": "KNH3",
"k_borate": "KB",
"k_bisulfate": "KSO4",
"k_CO2": "K0",
"k_carbonic_1": "K1",
"k_carbonic_2": "K2",
"k_fluoride": "KF",
"k_phosphate_1": "KP1",
"k_phosphate_2": "KP2",
"k_phosphate_3": "KP3",
"k_silicate": "KSi",
"k_sulfide": "KH2S",
"k_water": "KW",
"k_calcite": "KCa",
"k_aragonite": "KAr",
"k_alpha": "alpha",
"k_beta": "beta",
}
if np.any(np.isin(list(args.keys()), list(k_constants_optional.keys()))):
k_constants_in = {
k_constants_optional[k]: v
for k, v in args.items()
if k in k_constants_optional
}
else:
k_constants_in = None
k_constants_in = equilibria.assemble(
args["temperature"],
args["pressure"],
totals,
args["opt_pH_scale"],
args["opt_k_carbonic"],
args["opt_k_bisulfate"],
args["opt_k_fluoride"],
args["opt_gas_constant"],
Ks=k_constants_in,
)
# Solve the core marine carbonate system at input conditions
core_in = solve.core(
args["par1"],
args["par2"],
args["par1_type"],
args["par2_type"],
totals,
k_constants_in,
convert_units=True,
)
# Calculate the rest at input conditions
others_in = solve.others(
core_in,
args["temperature"],
args["pressure"],
totals,
k_constants_in,
args["opt_pH_scale"],
args["opt_k_carbonic"],
args["buffers_mode"],
)
# If requested, solve the core marine carbonate system at output conditions
if "pressure_out" in args.keys() or "temperature_out" in args.keys():
# Make sure we've got output values for both temperature and pressure
if "pressure_out" in args.keys():
if "temperature_out" not in args.keys():
args["temperature_out"] = args["temperature"]
if "temperature_out" in args.keys():
if "pressure_out" not in args.keys():
args["pressure_out"] = args["pressure"]
# Prepare equilibrium constants dict (output conditions)
k_constants_optional_out = {
"{}_out".format(k): v for k, v in k_constants_optional.items()
}
if np.any(np.isin(list(args.keys()), k_constants_optional_out)):
k_constants_out = {
k_constants_optional_out[k]: v
for k, v in args.items()
if k in k_constants_optional_out
}
else:
k_constants_out = None
k_constants_out = equilibria.assemble(
args["temperature_out"],
args["pressure_out"],
totals,
args["opt_pH_scale"],
args["opt_k_carbonic"],
args["opt_k_bisulfate"],
args["opt_k_fluoride"],
args["opt_gas_constant"],
Ks=k_constants_out,
)
# Solve the core marine carbonate system at output conditions
core_out = solve.core(
core_in["TA"],
core_in["TC"],
1,
2,
totals,
k_constants_out,
convert_units=False,
)
# Calculate the rest at output conditions
others_out = solve.others(
core_out,
args["temperature_out"],
args["pressure_out"],
totals,
k_constants_out,
args["opt_pH_scale"],
args["opt_k_carbonic"],
args["buffers_mode"],
)
else:
core_out = None
others_out = None
k_constants_out = None
return _get_results_dict(
args,
totals,
core_in,
others_in,
k_constants_in,
core_out,
others_out,
k_constants_out,
)
def MI2AMI(y, n_clusters, r, k, init, var_distrib, nj,\
nan_mask, target_nb_pseudo_obs = 500, it = 50, \
eps = 1E-05, maxstep = 100, seed = None, perform_selec = True,\
dm = [], max_patience = 1): # dm: Hack to remove
''' Complete the missing values using a trained M1DGMM
y (numobs x p ndarray): The observations containing mixed variables
n_clusters (int): The number of clusters to look for in the data
r (list): The dimension of latent variables through the first 2 layers
k (list): The number of components of the latent Gaussian mixture layers
init (dict): The initialisation parameters for the algorithm
var_distrib (p 1darray): An array containing the types of the variables in y
nj (p 1darray): For binary/count data: The maximum values that the variable can take.
For ordinal data: the number of different existing categories for each variable
nan_mask (ndarray): A mask array equal to True when the observation value is missing False otherwise
target_nb_pseudo_obs (int): The number of pseudo-observations to generate
it (int): The maximum number of MCEM iterations of the algorithm
eps (float): If the likelihood increase by less than eps then the algorithm stops
maxstep (int): The maximum number of optimisation step for each variable
seed (int): The random state seed to set (Only for numpy generated data for the moment)
perform_selec (Bool): Whether to perform architecture selection or not
dm (np array): The distance matrix of the observations. If not given M1DGMM computes it
n_neighbors (int): The number of neighbors to use for NA imputation
------------------------------------------------------------------------------------------------
returns (dict): The predicted classes, the likelihood through the EM steps
and a continuous representation of the data
'''
# !!! Hack
cols = y.columns
# Formatting
if not isinstance(nan_mask, np.ndarray): nan_mask = np.asarray(nan_mask)
if not isinstance(y, np.ndarray): y = np.asarray(y)
assert len(k) < 2 # Not implemented for deeper MDGMM for the moment
# Keep complete observations
complete_y = y[~np.isnan(y.astype(float)).any(1)]
completed_y = deepcopy(y)
out = M1DGMM(complete_y, 'auto', r, k, init, var_distrib, nj, it,\
eps, maxstep, seed, perform_selec = perform_selec,\
dm = dm, max_patience = max_patience, use_silhouette = True)
# Compute the associations
vc = vars_contributions(pd.DataFrame(complete_y, columns = cols), out['Ez.y'], assoc_thr = 0.0, \
title = 'Contribution of the variables to the latent dimensions',\
storage_path = None)
# Upacking the model from the M1DGMM output
#p = y.shape[1]
k = out['best_k']
r = out['best_r']
mu = out['mu'][0]
lambda_bin = np.array(out['lambda_bin'])
lambda_ord = out['lambda_ord']
lambda_categ = out['lambda_categ']
lambda_cont = np.array(out['lambda_cont'])
nj_bin = nj[pd.Series(var_distrib).isin(['bernoulli',
'binomial'])].astype(int)
nj_ord = nj[var_distrib == 'ordinal'].astype(int)
nj_categ = nj[var_distrib == 'categorical'].astype(int)
nb_cont = np.sum(var_distrib == 'continuous')
nb_bin = np.sum(var_distrib == 'binomial')
y_std = complete_y[:,var_distrib == 'continuous'].astype(float).std(axis = 0,\
keepdims = True)
cat_features = var_distrib != 'categorical'
# Compute the associations between variables and use them as weights for the optimisation
assoc = cosine_similarity(vc, dense_output=True)
np.fill_diagonal(assoc, 0.0)
assoc = np.abs(assoc)
weights = (assoc / assoc.sum(1, keepdims=True))
#==============================================
# Optimisation sandbox
#==============================================
# Define the observation generated by the center of each cluster
cluster_obs = [impute(mu[kk,:,0], var_distrib, lambda_bin, nj_bin, lambda_categ, nj_categ,\
lambda_ord, nj_ord, lambda_cont, y_std) for kk in range(k[0])]
# Use only of the observed variables as references
types = {'bin': ['bernoulli', 'binomial'], 'categ': ['categorical'],\
'cont': ['continuous'], 'ord': 'ordinal'}
# Gradient optimisation
nan_indices = np.where(nan_mask.any(1))[0]
imputed_y = np.zeros_like(y)
numobs = y.shape[0]
#************************************
# Linear constraint to stay in the support of continuous variables
#************************************
lb = np.array([])
ub = np.array([])
A = np.array([[]]).reshape((0, r[0]))
if nb_bin > 0:
## Corrected Binomial bounds (ub is actually +inf)
bin_indices = var_distrib[np.logical_or(var_distrib == 'bernoulli',
var_distrib == 'binomial')]
binomial_indices = bin_indices == 'binomial'
lb_bin = np.nanmin(y[:, var_distrib == 'binomial'], 0)
lb_bin = logit(
lb_bin / nj_bin[binomial_indices]) - lambda_bin[binomial_indices,
0]
ub_bin = np.nanmax(y[:, var_distrib == 'binomial'], 0)
ub_bin = logit(
ub_bin / nj_bin[binomial_indices]) - lambda_bin[binomial_indices,
0]
A_bin = lambda_bin[binomial_indices, 1:]
## Concatenate the constraints
lb = np.concatenate([lb, lb_bin])
ub = np.concatenate([ub, ub_bin])
A = np.concatenate([A, A_bin], axis=0)
if nb_cont > 0:
## Corrected Gaussian bounds
lb_cont = np.nanmin(y[:, var_distrib == 'continuous'],
0) / y_std[0] - lambda_cont[:, 0]
ub_cont = np.nanmax(y[:, var_distrib == 'continuous'],
0) / y_std[0] - lambda_cont[:, 0]
A_cont = lambda_cont[:, 1:]
## Concatenate the constraints
lb = np.concatenate([lb, lb_cont])
ub = np.concatenate([ub, ub_cont])
A = np.concatenate([A, A_cont], axis=0)
lc = LinearConstraint(A, lb, ub, keep_feasible=True)
zz = []
fun = []
for i in range(numobs):
if i in nan_indices:
# Design the nan masks for the optimisation process
nan_mask_i = nan_mask[i]
weights_i = weights[nan_mask_i].mean(0)
# Look for the best starting point
cluster_dist = [error(y[i, ~nan_mask_i], obs[~nan_mask_i],\
cat_features[~nan_mask_i], weights_i)\
for obs in cluster_obs]
z02 = mu[np.argmin(cluster_dist), :, 0]
# Formatting
vars_i = {type_alias: np.where(~nan_mask_i[np.isin(var_distrib, vartype)])[0] \
for type_alias, vartype in types.items()}
complete_categ = [
l for idx, l in enumerate(lambda_categ)
if idx in vars_i['categ']
]
complete_ord = [
l for idx, l in enumerate(lambda_ord) if idx in vars_i['ord']
]
opt = minimize(stat_all, z02, \
args = (y[i, ~nan_mask_i], var_distrib[~nan_mask_i],\
weights_i[~nan_mask_i],\
lambda_bin[vars_i['bin']], nj_bin[vars_i['bin']],\
complete_categ,\
nj_categ[vars_i['categ']],\
complete_ord,\
nj_ord[vars_i['ord']],\
lambda_cont[vars_i['cont']], y_std[:, vars_i['cont']]),
tol = eps, method='trust-constr', jac = grad_stat,\
constraints = lc,
options = {'maxiter': 1000})
z = opt.x
zz.append(z)
fun.append(opt.fun)
imputed_y[i] = impute(z, var_distrib, lambda_bin, nj_bin, lambda_categ, nj_categ,\
lambda_ord, nj_ord, lambda_cont, y_std)
else:
imputed_y[i] = y[i]
completed_y = np.where(nan_mask, imputed_y, y)
out['completed_y'] = completed_y
out['zz'] = zz
out['fun'] = fun
return (out)
