"""Constrained least squares for mixing ratios with known end-members.
Implements Carrera et al. (2004) equation 9: maximise the Gaussian likelihood of one
sample subject to the ratios summing to one, which gives a small linearly constrained
least-squares system.
Two structural facts make this cheap and robust:
* When every sample shares the same species and covariance, the KKT coefficient
matrix is **identical for all samples**, so it is factorised once.
* Because the weight matrix is positive definite the objective is convex and the
sum-to-one constraint makes ``lambda <= 1`` automatic, so the feasible region is
the simplex and any KKT point is the global optimum. Non-negativity is imposed by
a primal active set over the faces of that simplex -- one that releases
constraints as well as adding them, which is what makes the KKT point reachable.
See :func:`_solve_active_set`.
"""
from __future__ import annotations
import warnings
from typing import Any
import numpy as np
from mescla._util import as_endmembers, as_samples, check_aligned
from mescla.types import MixingResult
__all__ = ["mixing_ratios", "solve_ratios"]
def _warn_about_unlinearised_isotopes(table: Any) -> None:
"""A solute isotope used as an ordinary tracer is a silent wrong answer.
We warn rather than raise: the caller may have a reason, and refusing would be
presumptuous. Silence, however, would not be.
"""
from mescla.isotopes import check_isotopes
for problem in check_isotopes(table):
warnings.warn(problem, UserWarning, stacklevel=3)
[docs]
def solve_ratios(
F: np.ndarray,
y: np.ndarray,
weights: np.ndarray | None = None,
non_negative: bool = True,
censored: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray]:
"""Mixing ratios for a single sample (the low-level kernel).
Parameters
----------
F : ndarray, shape (ns, ne)
End-member concentrations, species down the rows.
y : ndarray, shape (ns,)
The mixed sample. ``nan`` marks a species that was not analysed in this
sample; it is simply left out of this sample's fit.
weights : ndarray, shape (ns,), optional
``1 / sigma**2`` per species. Uniform when omitted.
non_negative : bool, default True
Apply the active-set non-negativity step. Setting this to ``False`` is
diagnostically useful: a strongly negative unconstrained ratio says the
sample lies outside the mixing hull, which clipping would hide.
censored : ndarray of bool, shape (ns,), optional
``True`` where ``y`` holds a detection limit rather than a measurement.
A non-detect says the true value lies in ``[0, limit]``, so it constrains
the fit only from above: it is ignored while the model predicts a
concentration below the limit, and enters as a measurement at the limit
when the model predicts more. Solved by an outer active set.
Returns
-------
ratios : ndarray, shape (ne,)
active : ndarray of bool, shape (ne,)
Which end-members were forced to zero.
"""
F = np.asarray(F, float)
y = np.asarray(y, float)
ns, ne = F.shape
w = np.ones(ns) if weights is None else np.asarray(weights, float).copy()
# Species not analysed in this sample simply do not constrain it.
observed = np.isfinite(y) & np.isfinite(w)
if observed.sum() < ne - 1:
raise ValueError(
f"this sample has only {int(observed.sum())} usable species but "
f"{ne} end-members need at least {ne - 1}. Drop an end-member, or "
"restrict to samples with enough analyses."
)
limits = np.zeros(ns, dtype=bool) if censored is None else np.asarray(censored, bool)
# Start with non-detects excluded; the loop re-admits any the model overshoots.
enforced = np.zeros(ns, dtype=bool)
for _ in range(ns + 1):
usable = observed & (~limits | enforced)
if usable.sum() < ne - 1:
usable = observed
ratios, active = _solve_active_set(
F[usable], np.where(np.isnan(y), 0.0, y)[usable], w[usable], ne, non_negative
)
# The overshoot test compares two *concentrations*, so its tolerance has to
# be relative: an absolute one would mean something different in mg/L than
# in mol/L. (The ratio tolerances elsewhere are dimensionless and can stay
# absolute.)
predicted = F @ ratios
scale = np.maximum(np.abs(y), np.abs(predicted))
overshoot = (
limits & observed & ~enforced
& (predicted > y + 1e-12 * np.where(scale > 0.0, scale, 1.0))
)
if not overshoot.any():
return ratios, active
enforced |= overshoot
return ratios, active
def _solve_active_set(
F: np.ndarray,
y: np.ndarray,
w: np.ndarray,
ne: int,
non_negative: bool,
) -> tuple[np.ndarray, np.ndarray]:
"""The KKT solve with the non-negativity active set, on complete rows.
This is a primal active-set method for the convex quadratic program
.. math:: \\min_\\lambda (F\\lambda - y)^T W (F\\lambda - y)
\\quad\\text{s.t.}\\quad \\mathbf{1}^T\\lambda = 1,\\; \\lambda \\ge 0
It both **adds** constraints (when a step would take a ratio negative) and
**releases** them (when the Lagrange multiplier of an end-member pinned at zero
turns negative, meaning the objective would fall if it were let back in). The
release step is what makes the answer the global optimum: because the objective
is convex, a point that satisfies these KKT conditions is optimal.
Releasing matters in practice. Simply zeroing the most negative ratio and
re-solving -- never reconsidering -- is a greedy heuristic that settles on the
wrong face of the simplex for a sample outside the mixing hull: it can pin the
wrong end-member at zero and report ratios that are qualitatively different, not
merely imprecise. The effect needs four or more end-members to appear, which is
the ordinary case in this field.
"""
# Work with the Gram matrix: the per-face systems are then ne x ne at most,
# independent of how many species there are.
H = F.T @ (w[:, None] * F) # (ne, ne), positive semi-definite
c = F.T @ (w * y) # (ne,)
def solve_face(free: np.ndarray) -> tuple[np.ndarray, float]:
"""Equality-constrained LS on the face where only ``free`` may be non-zero.
[ H_ff 1 ] [lambda_f] [ c_f ]
[ 1' 0 ] [ mu ] = [ 1 ]
"""
nf = int(free.sum())
K = np.zeros((nf + 1, nf + 1))
K[:nf, :nf] = H[np.ix_(free, free)]
K[:nf, nf] = 1.0
K[nf, :nf] = 1.0
rhs = np.concatenate([c[free], [1.0]])
try:
solution = np.linalg.solve(K, rhs)
except np.linalg.LinAlgError:
solution = np.linalg.lstsq(K, rhs, rcond=None)[0]
return solution[:nf], float(solution[nf])
free = np.ones(ne, dtype=bool)
if not non_negative:
solution, _ = solve_face(free)
ratios = np.zeros(ne)
ratios[free] = solution
return ratios, ~free
# Start from the simplex centre: strictly feasible, so every later iterate is
# feasible too and the objective never increases.
ratios = np.full(ne, 1.0 / ne)
tol = 1e-12
# Each iteration either pins an end-member or releases one, and the objective
# strictly decreases, so no face is visited twice. The cap is a guard against
# cycling on a degenerate problem, not the expected exit.
for _ in range(4 * ne + 20):
solution, mu = solve_face(free)
candidate = np.zeros(ne)
candidate[free] = solution
if (solution >= -tol).all():
# The face optimum is feasible: move to it, then ask whether any
# end-member pinned at zero wants to come back.
ratios = np.clip(candidate, 0.0, None)
pinned = np.flatnonzero(~free)
if pinned.size == 0:
break
# Multiplier of each pinned bound; negative means releasing it pays.
z = (H @ ratios - c) + mu
worst = pinned[int(np.argmin(z[pinned]))]
if z[worst] < -tol:
free[worst] = True
continue
break
# The face optimum is infeasible. Walk towards it only as far as the first
# ratio that hits zero, and pin that one.
direction = candidate - ratios
blocking = np.flatnonzero(free & (direction < -tol))
if blocking.size == 0: # pragma: no cover - implies solution >= -tol above
ratios = np.clip(candidate, 0.0, None)
break
steps = ratios[blocking] / -direction[blocking]
first = int(np.argmin(steps))
ratios = ratios + min(1.0, steps[first]) * direction
hit = blocking[first]
ratios[hit] = 0.0
free[hit] = False
if not free.any(): # pragma: no cover - the sum-to-one constraint forbids it
break
ratios = np.clip(ratios, 0.0, 1.0)
total = ratios.sum()
if total > 0:
ratios = ratios / total
return ratios, ~free
[docs]
def mixing_ratios(
endmembers: Any,
samples: Any,
sigma: Any = None,
non_negative: bool = True,
) -> MixingResult:
"""Mixing ratios for every sample, by weighted constrained least squares.
Use this when the end-members are well characterised. When they are uncertain and
you have many mixed samples, :func:`mescla.mixing.ml.mix_ml` will do better --
least squares treats each sample independently, so it cannot improve as samples
accumulate.
Parameters
----------
endmembers : EndMembers, DataFrame or array, shape (ne, ns)
samples : WaterChemistry, DataFrame or array, shape (np, ns)
Ragged input is fine: ``nan`` marks a species not analysed in that sample
and simply does not constrain it. Each sample is solved independently, so
different samples may rest on different species -- but a sample still needs
at least ``ne - 1`` usable analyses. Non-detects carried on the container
as ``censored`` are treated as upper bounds, not as measurements.
sigma : array_like, optional
Standard deviations of the sample analyses. Falls back to the ``sigma``
carried by ``samples``, then to uniform weighting. Species measured
precisely, or known to be conservative, should get the smallest values --
this is how you tell the estimator which tracers to trust.
non_negative : bool, default True
Returns
-------
MixingResult
Examples
--------
>>> result = mixing_ratios(endmembers, samples) # doctest: +SKIP
>>> result.ratios_frame().head() # doctest: +SKIP
"""
Y = as_samples(samples, sigma=sigma)
E = as_endmembers(endmembers)
Y, E = check_aligned(Y, E)
_warn_about_unlinearised_isotopes(Y)
F = E.data.T # (ns, ne)
ne = E.n_endmembers
if F.shape[0] < ne - 1:
raise ValueError(
f"{ne} end-members need at least {ne - 1} species to be identifiable; "
f"only {F.shape[0]} supplied."
)
if not np.isfinite(E.data).all():
raise ValueError(
"the end-member compositions contain missing values. Every end-member "
"must be characterised in every species used; restrict the species with "
".select_species(...) to those you have for all of them."
)
sigmas = Y.sigma
ratios = np.empty((Y.n_rows, ne))
active = np.zeros((Y.n_rows, ne), dtype=bool)
for i in range(Y.n_rows):
w = None if sigmas is None else 1.0 / sigmas[i] ** 2
ratios[i], active[i] = solve_ratios(
F, Y.data[i], w, non_negative,
censored=None if Y.censored is None else Y.censored[i],
)
fitted = ratios @ E.data
residuals = Y.data - fitted
weights = np.ones_like(Y.data) if sigmas is None else 1.0 / sigmas**2
finite = np.isfinite(residuals)
objective = float(np.sum(weights[finite] * residuals[finite] ** 2))
return MixingResult(
ratios=ratios,
endmembers=E,
fitted=fitted,
residuals=residuals,
objective=objective,
method="constrained least squares",
converged=True,
samples=Y.index,
diagnostics={
"n_active_constraints": int(active.sum()),
"samples_with_active_constraints": int(active.any(axis=1).sum()),
"weighted": sigmas is not None,
"n_missing_analyses": int(np.isnan(Y.data).sum()),
"n_censored_analyses": Y.n_censored,
},
)