diff --git a/README.md b/README.md index 6ec4708..9a21583 100644 --- a/README.md +++ b/README.md @@ -103,6 +103,56 @@ Some documented example scripts are given in the directory `examples/`, and are python examples/unmix_mwe.py ``` +## Choosing a solver + +Two solvers are available. They differ in what they treat as "misfit", which changes both the +answer and what you can say about its uncertainty. + +| | `SampleNetworkUnmixer` | `LinearSampleNetworkUnmixer` | +|---|---|---| +| Misfit | relative (log-ratio surrogate) | absolute (least-squares) | +| Best for | data spanning orders of magnitude | data with low log-variance, e.g. isotopic | +| Forward model | convex program | exactly invertible matrix `d = Mc` | +| Regularization | shrinks towards the mean *observation* | shrinks towards the mean *model* | +| Uncertainty | Monte Carlo (`solve_montecarlo`) | closed form, `C_c = R C_d R^T` | +| Resolution | not available | resolution matrix and effective DOF | + +The linear solver is documented in Appendix A of the paper. It is the faster and more +informative of the two when its assumptions hold, because the mixing matrix is square and +invertible, so the estimator, its covariance, and its resolution are all available in closed +form: + +```python +problem = funmixer.LinearSampleNetworkUnmixer(sample_network, use_regularization=True) +solution = problem.solve( + element_data, + regularization_strength=1.0, + data_covariance=10.0, # a 10% relative error on the observations +) +solution.upstream_preds # recovered source concentrations +solution.upstream_std # their 1-sigma uncertainties, no Monte Carlo needed +solution.downstream_covariance # covariance of the modelled observations +solution.effective_dof # how many degrees of freedom the data actually constrain +``` + +See `examples/unmix_linear_mwe.py`. + +**Check the diagnostics before trusting a linear solution.** Inverting `M` amounts to +differencing each site against its upstream neighbours, and the noise amplification at each +site is `Q_i/q_i`, the ratio of total upstream flux to the flux the sub-basin itself generates. +Where sub-basin areas are very uneven this is large, and the unregularized inversion will +produce negative or physically impossible concentrations. `solution.amplification`, +`solution.condition_number`, `solution.clamped_nodes` and `solution.unconstrained_preds` all +report on this. On the `Mg` example data, the unregularized inversion clamps 10 of 63 sites at +zero and would otherwise return concentrations up to 145% by mass; at `lambda = 1` none clamp. +Regularization is not optional for this kind of data. + +Note that `regularization_strength` is **not** comparable between the two solvers: the linear +one weights a squared penalty on deviations from the mean model, the non-linear one an +unsquared norm of deviations from the mean observation. Within the linear solver, lambda is +dimensionless (the data are mean-normalised internally), so it does transfer between elements +and datasets. + ## Cite If you use this please cite the paper, which is published at *Water Resources Research*. diff --git a/docs/.gitignore b/docs/.gitignore new file mode 100644 index 0000000..80ff6a7 --- /dev/null +++ b/docs/.gitignore @@ -0,0 +1,10 @@ +# LaTeX build artefacts +*.aux +*.bcf +*.blg +*.fdb_latexmk +*.fls +*.log +*.out +*.run.xml +*.bbl diff --git a/docs/linear_unmixing_note.bib b/docs/linear_unmixing_note.bib new file mode 100644 index 0000000..6fb664a --- /dev/null +++ b/docs/linear_unmixing_note.bib @@ -0,0 +1,139 @@ +@article{barnes_using_2024, + title = {Using Convex Optimization to Efficiently Apportion Tracer and Pollutant Sources from Point Concentration Observations}, + author = {Barnes, Richard and Lipp, Alex G.}, + journal = {Water Resources Research}, + year = {2024}, + doi = {10.1029/2023WR036159}, +} + +@book{menke_geophysical_2012, + title = {Geophysical Data Analysis: Discrete Inverse Theory}, + author = {Menke, William}, + edition = {3}, + publisher = {Academic Press}, + address = {Boston}, + year = {2012}, + isbn = {9780123971609}, +} + +@book{aster_parameter_2018, + title = {Parameter Estimation and Inverse Problems}, + author = {Aster, Richard C. and Borchers, Brian and Thurber, Clifford H.}, + edition = {3}, + publisher = {Elsevier}, + year = {2018}, + isbn = {9780128046517}, +} + +@book{tarantola_inverse_2005, + title = {Inverse Problem Theory and Methods for Model Parameter Estimation}, + author = {Tarantola, Albert}, + publisher = {Society for Industrial and Applied Mathematics}, + address = {Philadelphia}, + year = {2005}, + isbn = {9780898715729}, +} + +@book{hansen_rank-deficient_1998, + title = {Rank-Deficient and Discrete Ill-Posed Problems: Numerical Aspects of Linear Inversion}, + author = {Hansen, Per Christian}, + publisher = {Society for Industrial and Applied Mathematics}, + address = {Philadelphia}, + year = {1998}, + isbn = {9780898714036}, +} + +@article{hoerl_ridge_1970, + title = {Ridge Regression: Biased Estimation for Nonorthogonal Problems}, + author = {Hoerl, Arthur E. and Kennard, Robert W.}, + journal = {Technometrics}, + volume = {12}, + number = {1}, + pages = {55--67}, + year = {1970}, + doi = {10.1080/00401706.1970.10488634}, +} + +@article{tikhonov_solution_1963, + title = {Solution of Incorrectly Formulated Problems and the Regularization Method}, + author = {Tikhonov, Andrey N.}, + journal = {Soviet Mathematics Doklady}, + volume = {4}, + pages = {1035--1038}, + year = {1963}, +} + +@article{elden_algorithms_1977, + title = {Algorithms for the Regularization of Ill-Conditioned Least Squares Problems}, + author = {Eld{\'e}n, Lars}, + journal = {BIT Numerical Mathematics}, + volume = {17}, + number = {2}, + pages = {134--145}, + year = {1977}, + doi = {10.1007/BF01932285}, +} + +@book{golub_matrix_2013, + title = {Matrix Computations}, + author = {Golub, Gene H. and Van Loan, Charles F.}, + edition = {4}, + publisher = {Johns Hopkins University Press}, + address = {Baltimore}, + year = {2013}, + isbn = {9781421407944}, +} + +@book{boyd_convex_2004, + title = {Convex Optimization}, + author = {Boyd, Stephen and Vandenberghe, Lieven}, + publisher = {Cambridge University Press}, + year = {2004}, + isbn = {9780521833783}, +} + +@article{diamond_cvxpy_2016, + title = {{CVXPY}: A Python-Embedded Modeling Language for Convex Optimization}, + author = {Diamond, Steven and Boyd, Stephen}, + journal = {Journal of Machine Learning Research}, + volume = {17}, + number = {83}, + pages = {1--5}, + year = {2016}, +} + +@inproceedings{blondes_practical_2016, + title = {A Practical Guide to the Use of Major Elements, Trace Elements, and Isotopes in Compositional Data Analysis: Applications for Deep Formation Brine Geochemistry}, + author = {Blondes, M. S. and Engle, M. A. and Geboy, N. J.}, + booktitle = {Compositional Data Analysis}, + series = {Springer Proceedings in Mathematics \& Statistics}, + publisher = {Springer International Publishing}, + pages = {13--29}, + year = {2016}, + isbn = {978-3-319-44811-4}, +} + +@book{aitchison_statistical_1986, + title = {The Statistical Analysis of Compositional Data}, + author = {Aitchison, John}, + publisher = {Chapman and Hall}, + year = {1986}, +} + +@book{cormen_introduction_2009, + title = {Introduction to Algorithms}, + author = {Cormen, Thomas H. and Leiserson, Charles E. and Rivest, Ronald L. and Stein, Clifford}, + edition = {3}, + publisher = {MIT Press}, + year = {2009}, + isbn = {9780262033848}, +} + +@book{horn_matrix_2012, + title = {Matrix Analysis}, + author = {Horn, Roger A. and Johnson, Charles R.}, + edition = {2}, + publisher = {Cambridge University Press}, + year = {2012}, + isbn = {9780521548236}, +} diff --git a/docs/linear_unmixing_note.pdf b/docs/linear_unmixing_note.pdf new file mode 100644 index 0000000..90446ff Binary files /dev/null and b/docs/linear_unmixing_note.pdf differ diff --git a/docs/linear_unmixing_note.tex b/docs/linear_unmixing_note.tex new file mode 100644 index 0000000..b92d3f5 --- /dev/null +++ b/docs/linear_unmixing_note.tex @@ -0,0 +1,865 @@ +\documentclass[11pt]{article} + +\usepackage{preprint} + +%% Math packages +\usepackage{amsmath, amsthm, amssymb, amsfonts} +\usepackage{bm} +\usepackage{mathtools} + +\usepackage[style=authoryear-comp,maxcitenames=2,uniquename=false,giveninits=true,uniquelist=false,maxbibnames=99,sortcites=true,sorting=nyt,doi=true,isbn=false,url=false,hyperref=true,date=year]{biblatex} +\renewbibmacro{in:}{% + \ifboolexpr{test {\ifentrytype{article}} or test {\ifentrytype{inproceedings}}}{}% + {\printtext{\bibstring{in}\intitlepunct}}% +} +\AtEveryBibitem{\clearlist{language}\clearfield{note}\clearfield{abstract}} +\addbibresource{linear_unmixing_note.bib} + +\usepackage[utf8]{inputenc} +\usepackage[T1]{fontenc} +\usepackage{xcolor} +\usepackage[colorlinks=true, linkcolor=magenta, urlcolor=blue, citecolor=cyan, anchorcolor=black]{hyperref} +\usepackage{booktabs} +\usepackage{microtype} +\usepackage{graphicx} +\usepackage{xspace} +\usepackage{titlesec} +\titlespacing\section{0pt}{12pt plus 3pt minus 3pt}{1pt plus 1pt minus 1pt} +\titlespacing\subsection{0pt}{10pt plus 3pt minus 3pt}{1pt plus 1pt minus 1pt} + +\usepackage{tikz} +\usetikzlibrary{arrows.meta} + +%% preprint.sty hardcodes the parent paper's title in the running header. +\lhead{\scshape Note -- Linear least-squares unmixing on a sample network} + +%% Notation, following the parent preprint +\newcommand{\cvec}{\ensuremath{\mathbf{c}}\xspace} +\newcommand{\dvec}{\ensuremath{\mathbf{d}}\xspace} +\newcommand{\qvec}{\ensuremath{\mathbf{q}}\xspace} +\newcommand{\Mmat}{\ensuremath{\mathbf{M}}\xspace} +\newcommand{\Rmat}{\ensuremath{\mathbf{R}}\xspace} +\newcommand{\Pmat}{\ensuremath{\mathbf{P}}\xspace} +\newcommand{\Amat}{\ensuremath{\mathbf{A}}\xspace} +\newcommand{\Wmat}{\ensuremath{\mathbf{W}}\xspace} +\newcommand{\Imat}{\ensuremath{\mathbf{I}}\xspace} +\newcommand{\onevec}{\ensuremath{\mathbf{1}}\xspace} +\newcommand{\Cd}{\ensuremath{\mathbf{C}_{\mathbf{d}}}\xspace} +\newcommand{\Cc}{\ensuremath{\mathbf{C}_{\mathbf{c}}}\xspace} +\newcommand{\Cdhat}{\ensuremath{\mathbf{C}_{\hat{\mathbf{d}}}}\xspace} +\newcommand{\Reals}{\ensuremath{\mathbb{R}}\xspace} +\newcommand{\Expect}{\ensuremath{\mathbb{E}}\xspace} +\newcommand{\transpose}{\ensuremath{^{\mathsf{T}}}} + +\DeclareMathOperator*{\argmin}{arg\,min} + +\newtheorem{theorem}{Proposition} +\newtheorem{lemma}[theorem]{Lemma} +\newtheorem{corollary}[theorem]{Corollary} +\theoremstyle{definition} +\newtheorem{remark}[theorem]{Remark} + +\title{Linear least-squares unmixing on a sample network:\\derivation, regularization, and error propagation} + +\usepackage{authblk} +\renewcommand*{\Authfont}{\bfseries} +\author[1]{Alex G. Lipp} +\affil[1]{Department of Earth Sciences, University College London} +\date{} + +\begin{document} +\maketitle + +\begin{abstract} +\noindent +\textcite{barnes_using_2024} recover source concentrations on a river sample network by +convex optimization, penalising \textit{relative} differences between predicted and observed +concentrations. Their Appendix~A observes that penalising \textit{absolute} differences +instead reduces the problem to a linear least-squares matrix inversion. This note works that +observation out in full. We show that the mixing matrix $\Mmat$ is square, row-stochastic and +triangular, hence always invertible, and we give its inverse in closed form as a local +differencing stencil on the network. We then extend the formulation in three directions that +the appendix does not cover: first-order tracer decay, weighting by the data covariance +(generalised least squares), and Tikhonov regularization of the model variance. We prove that +the regularized problem has a unique solution for every regularization strength, derive the +resulting linear estimator, and obtain closed-form covariances for both the recovered sources +and the modelled observations. We distinguish carefully between the propagated covariance and +the Bayesian posterior covariance, which coincide only when the problem is unregularized, and +give the exact identity relating them. Everything is derived from first principles so that it +can be checked line by line. +\end{abstract} + +\vspace{3mm} + +\section{Scope and notation} + +We adopt the notation of \textcite{barnes_using_2024} throughout. The sample network is a +directed acyclic graph $G = (N, E)$ whose $n$ nodes are sample sites and whose edges point +\textit{downstream}. Each node $i$ carries: + +\begin{itemize} + \item $a_i$ [L$^2$], the area of the sub-basin uniquely defined by sample site $i$; + \item $\phi_i$ [M\,L$^{-2}$\,T$^{-1}$], the material export rate of that sub-basin; + \item $q_i = a_i \phi_i > 0$ [M\,T$^{-1}$], the material flux the sub-basin generates; + \item $c_i$, the source concentration of the tracer in material exported by sub-basin $i$ + --- the unknown; + \item $d_i$, the concentration observed at the sample site --- the datum. +\end{itemize} + +\noindent Write $U_i \subseteq N$ for the set of nodes upstream of $i$ \textit{including $i$ +itself}, and $A(i)$ for the set of immediate upstream neighbours of $i$ (its predecessors in +$G$). We make one structural assumption throughout, which the D8 flow model guarantees: + +\begin{quote} +\textbf{(A1)} Every node has at most one downstream neighbour, so $G$ is a forest of +in-trees. +\end{quote} + +\noindent Assumption (A1) is what makes the sets $\{U_p\}_{p \in A(i)}$ pairwise disjoint, +which we use in Proposition~\ref{prop:inverse}. It does not require the network to be +connected: disjoint river basins may be solved together, exactly as in the parent study. + +Two derived quantities recur. The \textbf{total flux} passing site $i$ is +\begin{equation} + Q_i \;=\; \sum_{j \in U_i} q_j , + \label{eq:Q} +\end{equation} +and, for $j \in U_i$, the \textbf{path attenuation} from $j$ to $i$ is +\begin{equation} + \alpha_{j \to i} \;=\; \prod_{e \in \mathrm{path}(j \to i)} \exp\!\left( -k_{e} L_{e} \right), + \label{eq:alpha} +\end{equation} +the product over edges of the unique path from $j$ to $i$, where $L_e$ is the flow-path length +of edge $e$ and $k_e$ the first-order decay constant applied along it. By (A1) that path is +unique, so \eqref{eq:alpha} is well defined, and $\alpha_{i \to i} = 1$ (empty product). A +\textbf{conservative} tracer has $k_e = 0$ for all $e$, hence $\alpha_{j \to i} = 1$ +throughout; this is the case treated in Appendix~A of the parent paper. + +\section{The forward model as a matrix} + +\subsection{Assembling \texorpdfstring{$\Mmat$}{M}} + +Conservation of tracer at steady state gives the observation at site $i$ as a flux-weighted +mixture of every source upstream of it, attenuated by decay along the way: +\begin{equation} + d_i \;=\; \frac{\displaystyle\sum_{j \in U_i} \alpha_{j \to i}\, q_j\, c_j} + {\displaystyle\sum_{j \in U_i} q_j} + \;=\; \frac{1}{Q_i} \sum_{j \in U_i} \alpha_{j \to i}\, q_j\, c_j . + \label{eq:forward} +\end{equation} + +\noindent Note that decay enters the numerator only. Tracer is lost in transit, but the +\textit{carrier} --- the water or sediment --- is not, so the denominator is the undecayed +total flux $Q_i$. Setting every $\alpha = 1$ recovers Equation~3 of \textcite{barnes_using_2024}. + +Following the parent appendix, define the $0$--$1$ matrix $\bm{\Theta}$ by $\Theta_{ij} = 1$ +if and only if $j \in U_i$; that is, $\bm{\Theta}$ is the path matrix of $G$ with a self-loop +added at every node. Equation~\eqref{eq:forward} is then linear in \cvec, with +\begin{equation} + \boxed{\; + M_{ij} \;=\; \frac{\Theta_{ij}\, \alpha_{j \to i}\, q_j} + {\sum_{k=1}^{n} \Theta_{ik}\, q_k} + \;=\; \frac{\Theta_{ij}\, \alpha_{j \to i}\, q_j}{Q_i} , + \qquad \dvec = \Mmat \cvec . + \;} + \label{eq:M} +\end{equation} + +\noindent For a conservative tracer this is exactly Equation~11 of the parent appendix. + +\subsection{Structure of \texorpdfstring{$\Mmat$}{M}} + +$\Mmat$ is $n \times n$: the sample network defines exactly one sub-basin per sample site, so +there are as many unknowns as data. This squareness is the pivot on which everything below +turns, and it is worth stating plainly that it is a modelling choice, not a coincidence --- it +is what it means to attribute one source concentration to each sampled sub-basin. + +\begin{lemma}[Row sums]\label{lem:stochastic} +Every row of $\Mmat$ is non-negative with $\sum_j M_{ij} \le 1$, with equality for all $i$ if +and only if the tracer is conservative. +\end{lemma} + +\begin{proof} +Non-negativity is immediate from \eqref{eq:M} since $q_j > 0$, $\alpha > 0$, $\Theta \in +\{0,1\}$. For the sum, +\[ + \sum_{j=1}^n M_{ij} + = \frac{1}{Q_i}\sum_{j \in U_i} \alpha_{j\to i}\, q_j + \;\le\; \frac{1}{Q_i} \sum_{j \in U_i} q_j + \;=\; \frac{Q_i}{Q_i} = 1 , +\] +using $\alpha_{j \to i} \le 1$, which holds because every $k_e L_e \ge 0$. Equality for a +given $i$ requires $\alpha_{j \to i} = 1$ for all $j \in U_i$, i.e.\ no decay anywhere +upstream of $i$; equality for all $i$ requires it everywhere. +\end{proof} + +\noindent For a conservative tracer, then, $\Mmat$ is \textbf{row-stochastic}: each observation +is a convex combination of the sources upstream of it. This is the formal statement of ``a +river sample is a mixture of its upstream sources'', and it has a useful consequence, +$\Mmat \onevec = \onevec$, which we use in Section~\ref{sec:limits}. + +\begin{lemma}[Triangularity]\label{lem:triangular} +Let $\sigma$ be any topological ordering of $G$, so that $\sigma(u) < \sigma(v)$ whenever +$(u,v) \in E$. Permuting rows and columns of $\Mmat$ by $\sigma$ makes it lower triangular, +with diagonal entries +\begin{equation} + M_{ii} \;=\; \frac{q_i}{Q_i} \;\in\; (0, 1] . + \label{eq:diagonal} +\end{equation} +\end{lemma} + +\begin{proof} +$M_{ij} \ne 0$ requires $\Theta_{ij} = 1$, i.e.\ $j \in U_i$. If $j \ne i$ then a directed +path runs from $j$ to $i$, so $\sigma(j) < \sigma(i)$. Hence $M_{ij} \ne 0 \Rightarrow +\sigma(j) \le \sigma(i)$, which is lower triangularity in the permuted ordering. The diagonal +follows from \eqref{eq:M} with $j = i$, using $\Theta_{ii} = 1$ and $\alpha_{i \to i} = 1$; it +is positive because $q_i > 0$, and at most $1$ because $q_i \le Q_i$, with $M_{ii} = 1$ exactly +when $i$ is a leaf (headwater) node. +\end{proof} + +\begin{theorem}[Invertibility]\label{prop:invertible} +$\Mmat$ is invertible, with +\begin{equation} + \left| \det \Mmat \right| \;=\; \prod_{i=1}^{n} \frac{q_i}{Q_i} \;>\; 0 . +\end{equation} +\end{theorem} + +\begin{proof} +Let $\bm{\Pi}$ be the permutation matrix of the topological order $\sigma$. By +Lemma~\ref{lem:triangular}, $\bm{\Pi}\Mmat\bm{\Pi}\transpose$ is lower triangular, so its +determinant is the product of its diagonal entries, $\prod_i q_i / Q_i$, which is strictly +positive because every $q_i > 0$ and every $Q_i < \infty$. Since $|\det \bm{\Pi}| = 1$, +$|\det \Mmat| = \prod_i q_i/Q_i > 0$, and a matrix with non-zero determinant is invertible. +\end{proof} + +\begin{remark} +Invertibility is unaffected by decay: $\alpha$ appears only off the diagonal, so +\eqref{eq:diagonal} is untouched. It is also unaffected by the export rates $\phi$, provided +they are strictly positive. It \textit{would} fail if some sub-basin had zero area or zero +export rate, i.e.\ if two sample sites were co-located --- then $q_i = 0$, the corresponding +source is not identifiable, and the model is degenerate. This is worth checking in practice. +\end{remark} + +\section{The inverse in closed form} +\label{sec:inverse} + +Proposition~\ref{prop:invertible} says $\cvec = \Mmat^{-1}\dvec$ exists and is unique. We now show +it need never be computed by matrix inversion: it is a local differencing stencil on the +network, computable in one sweep. + +It is convenient to work with the \textbf{tracer flux} passing site $i$, +\begin{equation} + T_i \;=\; Q_i d_i \;=\; \sum_{j \in U_i} \alpha_{j \to i}\, q_j\, c_j , + \label{eq:T} +\end{equation} +which is just \eqref{eq:forward} cleared of its denominator. $T_i$ has units of tracer mass per +time: it is the amount of tracer passing the sample site per unit time. + +\begin{theorem}[Closed-form inverse]\label{prop:inverse} +For every node $i$, +\begin{equation} + \boxed{\; + c_i \;=\; \frac{1}{q_i}\left( Q_i d_i \;-\; \sum_{p \in A(i)} \alpha_{p \to i}\, Q_p d_p \right) . + \;} + \label{eq:closedform} +\end{equation} +\end{theorem} + +\begin{proof} +Every node in $U_i$ other than $i$ lies upstream of exactly one immediate neighbour $p \in +A(i)$, and by (A1) the sets $U_p$ for distinct $p \in A(i)$ are pairwise disjoint. Hence +\[ + U_i \;=\; \{i\} \;\sqcup\; \bigsqcup_{p \in A(i)} U_p , +\] +a disjoint union. Now fix $p \in A(i)$ and $j \in U_p$. The unique path $j \to i$ is the path +$j \to p$ followed by the single edge $p \to i$, so the attenuation factorises: +\begin{equation} + \alpha_{j \to i} \;=\; \alpha_{j \to p}\;\alpha_{p \to i} . + \label{eq:factorise} +\end{equation} +Substituting the partition and \eqref{eq:factorise} into \eqref{eq:T}, and using +$\alpha_{i \to i} = 1$ for the $j = i$ term, +\begin{align} + T_i + &= q_i c_i \;+\; \sum_{p \in A(i)} \; \sum_{j \in U_p} \alpha_{j \to i}\, q_j c_j \nonumber\\ + &= q_i c_i \;+\; \sum_{p \in A(i)} \alpha_{p \to i} \underbrace{\sum_{j \in U_p} \alpha_{j \to p}\, q_j c_j}_{= \,T_p} \nonumber\\ + &= q_i c_i \;+\; \sum_{p \in A(i)} \alpha_{p \to i}\, T_p . + \label{eq:balance} +\end{align} +Equation~\eqref{eq:balance} is a statement of tracer conservation at node $i$: the tracer +leaving equals the tracer generated locally plus the (attenuated) tracer delivered by the +immediate tributaries. Solving for $c_i$ and substituting $T = Qd$ gives +\eqref{eq:closedform}. Since $q_i > 0$ the division is legitimate. +\end{proof} + +\begin{corollary}[Sparsity of $\Mmat^{-1}$] +\begin{equation} + \left(\Mmat^{-1}\right)_{ij} = + \begin{dcases} + \dfrac{Q_i}{q_i}, & j = i, \\[6pt] + -\dfrac{\alpha_{j \to i}\, Q_j}{q_i}, & j \in A(i), \\[6pt] + 0, & \text{otherwise.} + \end{dcases} + \label{eq:Minv} +\end{equation} +$\Mmat^{-1}$ therefore has at most $n + |E| \le 2n - 1$ non-zero entries, against up to +$n(n+1)/2$ in $\Mmat$ itself. Equation~\eqref{eq:closedform} evaluates in $O(n + |E|)$ time. +\end{corollary} + +\noindent It is worth pausing on how much structure this reveals. $\Mmat$ is dense and global +--- every downstream observation depends on every source above it --- yet its inverse is +sparse and purely local. Recovering a source concentration requires only the observation at +that site and the observations at its immediate tributaries. All the long-range mixing cancels. + +\begin{corollary}[Worked example] +For the six-node network of \textcite[Figure~2]{barnes_using_2024}, with $A(4) = \{1,2\}$, +$A(5) = \{3\}$, $A(6) = \{4,5\}$ and no decay, Equation~\eqref{eq:closedform} gives +\begin{align*} + c_1 &= d_1, \qquad c_2 = d_2, \qquad c_3 = d_3, \\ + c_4 &= \frac{Q_4 d_4 - Q_1 d_1 - Q_2 d_2}{q_4}, \\ + c_5 &= \frac{Q_5 d_5 - Q_3 d_3}{q_5}, \\ + c_6 &= \frac{Q_6 d_6 - Q_4 d_4 - Q_5 d_5}{q_6}, +\end{align*} +with $Q_1 = q_1$, $Q_2 = q_2$, $Q_3 = q_3$, $Q_4 = q_1 + q_2 + q_4$, $Q_5 = q_3 + q_5$ and +$Q_6 = \sum_{i=1}^{6} q_i$. Headwater sites are recovered exactly and trivially; interior sites +are recovered by differencing. +\end{corollary} + +\subsection{Noise amplification} +\label{sec:amplification} + +Differencing is the source of all the difficulty. Suppose the observations carry independent +errors of standard deviation $\sigma_i$. Propagating them through \eqref{eq:closedform} +(anticipating Section~\ref{sec:propagation}, and using independence), +\begin{equation} + \operatorname{Var}(c_i) + = \frac{Q_i^2 \sigma_i^2 + \sum_{p \in A(i)} \alpha_{p\to i}^2 Q_p^2 \sigma_p^2}{q_i^2} . + \label{eq:varc} +\end{equation} +Every term is scaled by $q_i^{-2}$. Define the \textbf{amplification factor} +\begin{equation} + \kappa_i \;=\; \frac{Q_i}{q_i} \;=\; \frac{1}{M_{ii}} \;\ge\; 1 , + \label{eq:kappa} +\end{equation} +the ratio of the total flux passing site $i$ to the flux its own sub-basin generates. If +site $i$ sits immediately downstream of another site, with little incremental drainage area +between them, then $q_i \ll Q_i$, $\kappa_i$ is large, and \eqref{eq:closedform} subtracts two +nearly equal large numbers to obtain a small one. This is catastrophic cancellation, and it is +an intrinsic property of the sampling design, not of the algorithm: the data simply do not +constrain that sub-basin well. In practice $\kappa_i$, being free to compute, is the single +most useful diagnostic of whether a linear inversion is worth attempting on a given network. + +\begin{corollary}[When positivity binds]\label{cor:positivity} +Since concentrations must be non-negative, the unregularized solution is physically admissible +at node $i$ if and only if +\begin{equation} + Q_i d_i \;\ge\; \sum_{p \in A(i)} \alpha_{p \to i} \, Q_p d_p , + \label{eq:positivity} +\end{equation} +that is, if and only if at least as much tracer flux leaves site $i$ as its tributaries deliver +to it. A violation is a direct falsification of the forward model \eqref{eq:forward} by the +data, whether through measurement error, non-conservative behaviour, or mis-specified fluxes. +\end{corollary} + +\section{Why the unregularized problem is degenerate} +\label{sec:degenerate} + +Appendix~A of \textcite{barnes_using_2024} poses the recovery as a constrained least-squares +problem, +\begin{equation} + \underset{\cvec \in \Reals^n}{\text{minimize}} \;\; \lVert \Mmat\cvec - \dvec \rVert_2 + \qquad \text{subject to} \quad c_i \ge 0 . + \label{eq:appendixproblem} +\end{equation} +Proposition~\ref{prop:invertible} shows that, absent the constraint, this is not really a +least-squares problem at all. + +\begin{theorem}[Degeneracy]\label{prop:degenerate} +Let $\Wmat$ be any symmetric positive definite weight matrix. Then +\begin{equation} + \argmin_{\cvec \in \Reals^n} \; (\Mmat\cvec - \dvec)\transpose \Wmat (\Mmat\cvec - \dvec) + \;=\; \Mmat^{-1}\dvec , +\end{equation} +independently of $\Wmat$, and the residual at the optimum is exactly zero. +\end{theorem} + +\begin{proof} +The objective is non-negative for every \cvec, and $\cvec = \Mmat^{-1}\dvec$ --- which exists +by Proposition~\ref{prop:invertible} --- attains the value $0$. Hence it is a global minimiser, and +it is the unique one because $\Wmat \succ 0$ forces the residual to vanish, and $\Mmat$ is +injective. Algebraically, the generalised least-squares estimator collapses: +\[ + (\Mmat\transpose\Wmat\Mmat)^{-1}\Mmat\transpose\Wmat + = \Mmat^{-1}\Wmat^{-1}\Mmat^{-\mathsf{T}}\Mmat\transpose\Wmat + = \Mmat^{-1} . \qedhere +\] +\end{proof} + +\noindent Three consequences follow, and they set up the rest of this note. + +\begin{enumerate} + \item \textbf{Weighting is inert when $\lambda = 0$.} No choice of $\Wmat$ --- including + $\Wmat = \Cd^{-1}$, the statistically correct one --- changes the answer. There is + nothing to trade off when the model can fit every datum exactly. Weighting only + becomes consequential once a competing term is added to the objective + (Section~\ref{sec:regularization}). + \item \textbf{The problem interpolates the noise.} A zero residual means the estimate + reproduces the observations exactly, errors included. Combined with the amplification + of Section~\ref{sec:amplification}, this is why the raw inversion is unusable on + realistic data. + \item \textbf{The constraint is the only thing that can bite.} By + Corollary~\ref{cor:positivity}, \eqref{eq:appendixproblem} departs from + $\Mmat^{-1}\dvec$ precisely when \eqref{eq:positivity} fails somewhere. +\end{enumerate} + +\begin{remark}[Bounds] +The parent appendix writes the constraint as $0 < c_i < 1$, appropriate for a tracer measured +as a mass fraction. Two remarks. First, a strict inequality cannot be imposed in a convex +program --- the feasible set would not be closed and the infimum need not be attained --- so +one uses $c_i \ge 0$. Second, the upper bound is unit-dependent ($1$ for a fraction, $10^6$ for +mg\,kg$^{-1}$) whereas the lower bound is not. Since $\Mmat$ is row-substochastic, an estimate +exceeding the physical ceiling indicates that the model has already failed, and is more useful +surfaced as a diagnostic than silently clipped. Below we impose $c_i \ge 0$ only. +\end{remark} + +\section{Scale invariance and normalisation} +\label{sec:scale} + +\begin{theorem}[Homogeneity]\label{prop:homogeneity} +The unregularized solution map $\dvec \mapsto \cvec$ is homogeneous of degree one: +$\cvec(t\dvec) = t\,\cvec(\dvec)$ for any $t > 0$. Moreover $\Mmat$ is invariant under a common +rescaling of the fluxes, $q_j \mapsto sq_j$. +\end{theorem} + +\begin{proof} +The first claim is immediate from linearity, $\cvec = \Mmat^{-1}\dvec$. For the second, +replacing $q_j$ by $sq_j$ in \eqref{eq:M} multiplies numerator and denominator by $s$. +\end{proof} + +\noindent Homogeneity licenses the mean normalisation used in the implementation. Writing +$\bar{d} = n^{-1}\sum_i d_i$ and solving with $\tilde{\dvec} = \dvec / \bar{d}$, the recovered +sources are $\cvec = \bar{d}\,\tilde{\cvec}$ exactly. This changes nothing mathematically; its +purpose is numerical, keeping all quantities near unity so that solver tolerances --- which are +partly absolute --- behave uniformly across tracers whose concentrations may differ by many +orders of magnitude. It has the further benefit of rendering the regularization strength +introduced below dimensionless, and hence transferable between tracers and datasets. + +An identical argument applies to the sub-basin areas, which the implementation divides by their +mean, and to covariances, which by Proposition~\ref{prop:homogeneity} and +Section~\ref{sec:propagation} scale as $\bar{d}^{\,2}$ and so may equivalently be computed in +the original units. + +\section{Weighting by the data covariance} +\label{sec:weighting} + +Let $\Cd \succ 0$ be the covariance of the observation errors. The statistically natural misfit +is the Mahalanobis norm of the residual, giving generalised least squares +\parencite{menke_geophysical_2012, aster_parameter_2018}: +\begin{equation} + \Phi(\cvec) \;=\; (\Mmat\cvec - \dvec)\transpose \Cd^{-1} (\Mmat\cvec - \dvec) . + \label{eq:gls} +\end{equation} +Numerically one never forms $\Cd^{-1}$. Take the Cholesky factorisation $\Cd = \mathbf{L} +\mathbf{L}\transpose$ \parencite{golub_matrix_2013} and define the whitened system +\begin{equation} + \widetilde{\Mmat} = \mathbf{L}^{-1}\Mmat, \qquad + \tilde{\dvec} = \mathbf{L}^{-1}\dvec , + \label{eq:whitening} +\end{equation} +so that $\Phi(\cvec) = \lVert \widetilde{\Mmat}\cvec - \tilde{\dvec}\rVert_2^2$: generalised +least squares becomes ordinary least squares on the whitened system, whose residuals have +identity covariance. When $\Cd$ is diagonal this reduces to dividing row $i$ of $\Mmat$ and +$d_i$ by $\sigma_i$. + +\begin{remark}[Relation to the log-ratio objective] +\label{rem:logratio} +Suppose errors are proportional, $\sigma_i = f d_i$ for a constant relative error $f$ --- the +standard model for geochemical assay. The $i$th whitened residual is then +\begin{equation} + \frac{(\Mmat\cvec - \dvec)_i}{f d_i} + \;=\; \frac{1}{f}\left( \frac{(\Mmat\cvec)_i}{d_i} - 1 \right) , +\end{equation} +which depends on the data only through the \textit{ratio} of predicted to observed +concentration. The weighted linear objective is therefore, like the log-ratio objective of the +parent study, a relative misfit; it is the unweighted objective that is not. This is precisely +the objection \textcite{barnes_using_2024} raise against Euclidean misfit --- that it +over-weights high concentrations --- and weighting by a proportional error model answers it. +The two remain distinct: \eqref{eq:gls} penalises the squared relative residual, whereas +$\max(x, 1/x)$ with $x = d_{\text{pred}}/d_{\text{obs}}$ behaves like $1 + |x - 1|$ near +$x = 1$. Both are functions of the relative residual alone, but one is quadratic and the other +piecewise linear in it, so they weight outliers differently and do not coincide. +\end{remark} + +\section{Regularization} +\label{sec:regularization} + +Section~\ref{sec:degenerate} showed that the unregularized problem interpolates the noise. The +remedy is Tikhonov regularization \parencite{tikhonov_solution_1963, hoerl_ridge_1970, +hansen_rank-deficient_1998}: add a penalty expressing a preference among models. + +We penalise the \textbf{variance of the model about its own mean}. Let +\begin{equation} + \Pmat \;=\; \Imat - \tfrac{1}{n}\onevec\onevec\transpose + \label{eq:P} +\end{equation} +be the centering projector, so that $(\Pmat\cvec)_i = c_i - \bar{c}$. The full problem is +\begin{equation} + \boxed{\; + \begin{aligned} + \underset{\cvec \in \Reals^n}{\text{minimize}} \quad + & J(\cvec) = (\Mmat\cvec - \dvec)\transpose \Cd^{-1}(\Mmat\cvec - \dvec) + + \lambda \lVert \Pmat\cvec \rVert_2^2 \\ + \text{subject to} \quad & \cvec \ge 0 . + \end{aligned} + \;} + \label{eq:problem} +\end{equation} + +\noindent Note that this penalises departure from the mean of the \textit{model}, not, as in +the parent study, from the mean of the \textit{observations}. That substitution was forced +there by the requirement that the objective be expressible in disciplined convex form +\parencite{boyd_convex_2004, diamond_cvxpy_2016}; the linear formulation is under no such +constraint and can use the more natural quantity. + +\begin{lemma}[Properties of $\Pmat$]\label{lem:P} +$\Pmat$ is symmetric and idempotent, hence an orthogonal projector; it is positive +semi-definite with eigenvalue $0$ on $\operatorname{span}\{\onevec\}$ and eigenvalue $1$ on its +orthogonal complement; and +\begin{equation} + \lVert \Pmat \cvec \rVert_2^2 = \cvec\transpose \Pmat\transpose \Pmat \cvec + = \cvec\transpose \Pmat \cvec = \sum_{i=1}^n (c_i - \bar{c})^2 . +\end{equation} +\end{lemma} + +\begin{proof} +Symmetry is clear. For idempotency, using $\onevec\transpose\onevec = n$, +\[ + \Pmat^2 = \Imat - \tfrac{2}{n}\onevec\onevec\transpose + + \tfrac{1}{n^2}\onevec(\onevec\transpose\onevec)\onevec\transpose + = \Imat - \tfrac{1}{n}\onevec\onevec\transpose = \Pmat . +\] +Then $\Pmat\transpose\Pmat = \Pmat^2 = \Pmat$, giving the quadratic form. $\Pmat\onevec = +\onevec - \onevec = \bm{0}$, and for $\mathbf{v} \perp \onevec$, $\Pmat\mathbf{v} = +\mathbf{v}$. Eigenvalues in $\{0,1\}$ give $\Pmat \succeq 0$. +\end{proof} + +\noindent $\Pmat$ is thus singular, with a one-dimensional null space. This matters: the +penalty alone does not determine a unique solution, and the classical uniqueness condition for +general-form Tikhonov regularization \parencite{elden_algorithms_1977, +hansen_rank-deficient_1998} requires +$\mathcal{N}(\Mmat) \cap \mathcal{N}(\Pmat) = \{\bm{0}\}$. Here $\mathcal{N}(\Mmat) = +\{\bm{0}\}$ by Proposition~\ref{prop:invertible}, so the condition holds trivially --- but it holds +\textit{because of} the structure of $\Mmat$, not because of the penalty. + +\begin{theorem}[Existence and uniqueness]\label{prop:unique} +For every $\lambda \ge 0$, $J$ is strictly convex on $\Reals^n$. Consequently +\eqref{eq:problem} has a unique global minimiser, both unconstrained and subject to +$\cvec \ge 0$. +\end{theorem} + +\begin{proof} +$J$ is a sum of quadratic forms, with gradient and Hessian +\begin{align} + \nabla J(\cvec) &= 2\Mmat\transpose\Cd^{-1}(\Mmat\cvec - \dvec) + 2\lambda\Pmat\cvec , + \label{eq:grad}\\ + \nabla^2 J &= 2\left(\Mmat\transpose\Cd^{-1}\Mmat + \lambda\Pmat\right) . + \label{eq:hess} +\end{align} +For any $\mathbf{v} \ne \bm{0}$, $\mathbf{v}\transpose\Mmat\transpose\Cd^{-1}\Mmat\mathbf{v} = +(\Mmat\mathbf{v})\transpose\Cd^{-1}(\Mmat\mathbf{v}) > 0$, since $\Mmat\mathbf{v} \ne \bm{0}$ +by Proposition~\ref{prop:invertible} and $\Cd^{-1} \succ 0$. Hence +$\Mmat\transpose\Cd^{-1}\Mmat \succ 0$. By Lemma~\ref{lem:P}, $\lambda\Pmat \succeq 0$ for +$\lambda \ge 0$. A positive definite matrix plus a positive semi-definite one is positive +definite \parencite{horn_matrix_2012}, so $\nabla^2 J \succ 0$ and $J$ is strictly convex. A +strictly convex function has at most one minimiser on any convex set, and coercivity +(guaranteed by $\nabla^2 J \succ 0$) ensures one exists. The non-negative orthant is convex and +closed, so the constrained problem inherits both. +\end{proof} + +\begin{remark} +Uniqueness rests on $\Mmat$ being invertible, \textit{not} on the penalty. Had $\Mmat$ been +singular with a constant vector in its null space, $\lambda\Pmat$ could not have restored +uniqueness, because $\Pmat$ is blind to exactly that direction. Regularization here buys +conditioning and stability, not identifiability --- the problem was already identifiable. +\end{remark} + +\begin{theorem}[The estimator]\label{prop:estimator} +Define +\begin{equation} + \Amat_\lambda \;=\; \Mmat\transpose\Cd^{-1}\Mmat + \lambda\Pmat . + \label{eq:A} +\end{equation} +If the unconstrained minimiser of $J$ is non-negative, it is also the solution of +\eqref{eq:problem}, and it is given by the \textbf{linear} map +\begin{equation} + \boxed{\; + \hat{\cvec} = \Rmat_\lambda \dvec , + \qquad + \Rmat_\lambda = \Amat_\lambda^{-1}\Mmat\transpose\Cd^{-1} . + \;} + \label{eq:estimator} +\end{equation} +\end{theorem} + +\begin{proof} +Setting \eqref{eq:grad} to zero gives the normal equations +$(\Mmat\transpose\Cd^{-1}\Mmat + \lambda\Pmat)\cvec = \Mmat\transpose\Cd^{-1}\dvec$, i.e.\ +$\Amat_\lambda\cvec = \Mmat\transpose\Cd^{-1}\dvec$. $\Amat_\lambda \succ 0$ by the proof of +Proposition~\ref{prop:unique}, hence invertible, giving \eqref{eq:estimator}. For the first claim: +$J$ is strictly convex, so $\hat\cvec$ is its unique global minimiser over all of $\Reals^n$; +if it happens to be feasible then it is a fortiori the minimiser over the smaller feasible set. +\end{proof} + +\noindent Two features of \eqref{eq:estimator} deserve emphasis. First, it is exactly linear in +\dvec, with \textbf{no constant offset}. This is a consequence of penalising +$\lVert\Pmat\cvec\rVert^2$, which is homogeneous in \cvec; a penalty +$\lVert\cvec - \cvec_0\rVert^2$ towards a fixed reference $\cvec_0 \ne \bm{0}$ would contribute +an affine term and complicate everything downstream. Second, +Proposition~\ref{prop:estimator} gives a practical recipe: solve the problem, and if the answer is +non-negative, report the analytical $\Rmat_\lambda\dvec$. The estimate and its covariance then +provably come from the same linear operator, which is what makes the error propagation of +Section~\ref{sec:propagation} meaningful. When the constraint \textit{is} active the estimator +is only piecewise linear, and that guarantee is lost. + +\subsection{Limits} +\label{sec:limits} + +\begin{theorem}[Small $\lambda$] +$\Rmat_0 = \Mmat^{-1}$, recovering Section~\ref{sec:inverse}. +\end{theorem} + +\begin{proof} +Set $\lambda = 0$ in \eqref{eq:estimator}: $\Rmat_0 = +(\Mmat\transpose\Cd^{-1}\Mmat)^{-1}\Mmat\transpose\Cd^{-1} = \Mmat^{-1}$, as in +Proposition~\ref{prop:degenerate}. +\end{proof} + +\begin{theorem}[Large $\lambda$]\label{prop:largelambda} +As $\lambda \to \infty$, $\hat{\cvec} \to t^\star\onevec$ with +\begin{equation} + t^\star = \frac{\mathbf{m}\transpose\Cd^{-1}\dvec}{\mathbf{m}\transpose\Cd^{-1}\mathbf{m}}, + \qquad \mathbf{m} = \Mmat\onevec . +\end{equation} +If in addition the tracer is conservative and $\Cd = \sigma^2\Imat$, then +$t^\star = \bar{d}$, the arithmetic mean of the observations. +\end{theorem} + +\begin{proof} +The penalty term $\lambda\lVert\Pmat\cvec\rVert^2$ diverges unless $\Pmat\cvec \to \bm 0$, i.e.\ +unless \cvec approaches $\mathcal{N}(\Pmat) = \operatorname{span}\{\onevec\}$. Writing +$\cvec = t\onevec$ the penalty vanishes identically and the problem reduces to minimising +$(t\mathbf{m} - \dvec)\transpose\Cd^{-1}(t\mathbf{m} - \dvec)$ over the scalar $t$. +Differentiating and setting to zero gives $t^\star$ as stated. For a conservative tracer +Lemma~\ref{lem:stochastic} gives $\mathbf{m} = \Mmat\onevec = \onevec$, whence with +$\Cd = \sigma^2\Imat$ we get $t^\star = \onevec\transpose\dvec / \onevec\transpose\onevec = +\bar{d}$. +\end{proof} + +\noindent So the regularization path runs from the exact, noise-interpolating inverse at +$\lambda = 0$ to a spatially uniform source field at $\lambda \to \infty$, whose level is the +(precision-weighted) mean observation. Both endpoints are interpretable, which is a useful +property when selecting $\lambda$ from the trade-off curve between misfit +$\lVert\Cd^{-1/2}(\Mmat\hat\cvec - \dvec)\rVert_2$ and roughness +$\lVert\Pmat\hat\cvec\rVert_2$ \parencite{hansen_rank-deficient_1998}. + +\section{Error propagation} +\label{sec:propagation} + +We now propagate the observational uncertainty $\Cd$ into the recovered sources. Because +$\hat\cvec$ is an exactly linear function of \dvec, this is not a linearisation: it is exact. + +\begin{theorem}[Covariance of the sources]\label{prop:Cc} +If $\dvec$ has covariance $\Cd$ and no constraint is active, then +\begin{equation} + \boxed{\;\Cc \;=\; \Rmat_\lambda\, \Cd\, \Rmat_\lambda\transpose . \;} + \label{eq:Cc} +\end{equation} +\end{theorem} + +\begin{proof} +For a deterministic matrix $\Rmat$ and random vector $\dvec$, +$\operatorname{Cov}(\Rmat\dvec) = \Rmat \operatorname{Cov}(\dvec) \Rmat\transpose$, directly +from the definition of covariance \parencite{menke_geophysical_2012}. Apply with +$\Rmat = \Rmat_\lambda$, which is deterministic given the network and the fluxes. +\end{proof} + +\begin{theorem}[Covariance of the modelled observations]\label{prop:Cdhat} +The modelled observations $\hat\dvec = \Mmat\hat\cvec$ have +\begin{equation} + \boxed{\;\Cdhat \;=\; \Mmat\, \Cc\, \Mmat\transpose + \;=\; (\Mmat\Rmat_\lambda)\,\Cd\,(\Mmat\Rmat_\lambda)\transpose , \;} + \label{eq:Cdhat} +\end{equation} +and at $\lambda = 0$, $\Cdhat = \Cd$ exactly. +\end{theorem} + +\begin{proof} +The first equality is Proposition~\ref{prop:Cc} applied to the linear map $\hat\cvec \mapsto +\Mmat\hat\cvec$; the second follows by substitution. $\Mmat\Rmat_\lambda$ is the hat (influence) +matrix. At $\lambda = 0$, $\Rmat_0 = \Mmat^{-1}$ so $\Mmat\Rmat_0 = \Imat$ and $\Cdhat = \Cd$. +This is as it must be: by Proposition~\ref{prop:degenerate} the unregularized fit reproduces the +observations exactly, so the modelled observations inherit their uncertainty unchanged. +\end{proof} + +\subsection{The unregularized case is efficient} + +\begin{theorem}[Cram\'er--Rao]\label{prop:crb} +At $\lambda = 0$, +\begin{equation} + \Cc = \Mmat^{-1}\Cd\Mmat^{-\mathsf{T}} = \left(\Mmat\transpose\Cd^{-1}\Mmat\right)^{-1} , +\end{equation} +which is the inverse Fisher information of the model $\dvec \sim +\mathcal{N}(\Mmat\cvec, \Cd)$. The unregularized estimator therefore attains the Cram\'er--Rao +lower bound. +\end{theorem} + +\begin{proof} +For the identity, note $(\Mmat^{-1}\Cd\Mmat^{-\mathsf{T}})^{-1} = \Mmat\transpose\Cd^{-1}\Mmat$ +by inverting each factor and reversing the order. For the Fisher information, the +log-likelihood is $\ell(\cvec) = -\tfrac{1}{2}(\dvec - \Mmat\cvec)\transpose\Cd^{-1}(\dvec - +\Mmat\cvec) + \text{const}$, so $\nabla_{\cvec}\ell = \Mmat\transpose\Cd^{-1}(\dvec - +\Mmat\cvec)$ and $-\nabla^2_{\cvec}\ell = \Mmat\transpose\Cd^{-1}\Mmat$, which is +deterministic and hence equal to its expectation. The estimator is unbiased (see +Proposition~\ref{prop:bias} with $\lambda = 0$), so the Cram\'er--Rao bound applies and is met with +equality. +\end{proof} + +\noindent No estimator, weighted or otherwise, can do better than $\Mmat^{-1}\dvec$ in the +unbiased class. Section~\ref{sec:degenerate} already told us weighting cannot help; this tells +us why nothing else can either. The only way to improve on \eqref{eq:varc} is to accept bias +--- which is exactly what regularization does. + +\subsection{Bias, resolution, and the two covariances} + +\begin{theorem}[Bias and resolution]\label{prop:bias} +$\Expect[\hat\cvec] = \Rmat_\lambda\Mmat\cvec_{\text{true}}$, so the estimator has bias +\begin{equation} + \mathbf{b} = (\Rmat_\lambda\Mmat - \Imat)\cvec_{\text{true}} + = -\lambda\,\Amat_\lambda^{-1}\Pmat\,\cvec_{\text{true}} , + \label{eq:bias} +\end{equation} +and \textbf{resolution matrix} +\begin{equation} + \Rmat_\lambda\Mmat = \Imat - \lambda\Amat_\lambda^{-1}\Pmat . + \label{eq:resolution} +\end{equation} +The mean-squared error is $\Expect[(\hat\cvec - \cvec_{\text{true}})(\cdot)\transpose] += \Cc + \mathbf{b}\mathbf{b}\transpose$. +\end{theorem} + +\begin{proof} +Substituting $\dvec = \Mmat\cvec_{\text{true}} + \bm{\varepsilon}$ with +$\Expect[\bm{\varepsilon}] = \bm 0$ into \eqref{eq:estimator} gives the first claim. Then +\[ + \Rmat_\lambda\Mmat - \Imat + = \Amat_\lambda^{-1}\Mmat\transpose\Cd^{-1}\Mmat - \Imat + = \Amat_\lambda^{-1}\!\left(\Mmat\transpose\Cd^{-1}\Mmat - \Amat_\lambda\right) + = -\lambda\Amat_\lambda^{-1}\Pmat , +\] +using \eqref{eq:A}. The MSE decomposition is the standard bias--variance identity +\parencite{hoerl_ridge_1970, aster_parameter_2018}. +\end{proof} + +\noindent The bias vanishes when $\lambda = 0$, and also whenever $\Pmat\cvec_{\text{true}} = +\bm 0$, i.e.\ when the true source field really is uniform --- the penalty is then telling the +truth and costs nothing. Since $\cvec_{\text{true}}$ is unknown, $\mathbf{b}$ cannot be +evaluated; the resolution matrix \eqref{eq:resolution} is the standard computable proxy +\parencite{menke_geophysical_2012}, equal to $\Imat$ at $\lambda = 0$ and degrading as +$\lambda$ grows. Its trace, $\operatorname{tr}(\Rmat_\lambda\Mmat)$, is the effective number of +degrees of freedom the data constrain. + +This bias is precisely why \eqref{eq:Cc} must be interpreted with care. It answers the +question ``if I repeated the measurements, how would my estimate scatter?'' --- and a heavily +damped estimator scatters very little while being systematically wrong. A second covariance +answers a different and often more useful question. + +\begin{theorem}[Posterior covariance]\label{prop:posterior} +Read the penalty as a Gaussian prior with precision $\lambda\Pmat$ and the likelihood as +$\dvec \mid \cvec \sim \mathcal{N}(\Mmat\cvec, \Cd)$. Then the posterior covariance is +$\Amat_\lambda^{-1}$, and +\begin{equation} + \boxed{\; + \Rmat_\lambda\Cd\Rmat_\lambda\transpose + \;=\; \Amat_\lambda^{-1} \;-\; \lambda\,\Amat_\lambda^{-1}\Pmat\Amat_\lambda^{-1} . + \;} + \label{eq:identity} +\end{equation} +Consequently $\Cc \preceq \Amat_\lambda^{-1}$, with equality if and only if $\lambda = 0$. +\end{theorem} + +\begin{proof} +For the linear-Gaussian model the posterior precision is the sum of the likelihood precision +$\Mmat\transpose\Cd^{-1}\Mmat$ and the prior precision $\lambda\Pmat$, giving $\Amat_\lambda$ +\parencite{tarantola_inverse_2005}. For the identity, substitute \eqref{eq:estimator} and use +$\Cd^{-1}\Cd\Cd^{-1} = \Cd^{-1}$: +\begin{align*} + \Rmat_\lambda\Cd\Rmat_\lambda\transpose + &= \Amat_\lambda^{-1}\Mmat\transpose\Cd^{-1}\,\Cd\,\Cd^{-1}\Mmat\Amat_\lambda^{-1} \\ + &= \Amat_\lambda^{-1}\left(\Mmat\transpose\Cd^{-1}\Mmat\right)\Amat_\lambda^{-1} \\ + &= \Amat_\lambda^{-1}\left(\Amat_\lambda - \lambda\Pmat\right)\Amat_\lambda^{-1} + = \Amat_\lambda^{-1} - \lambda\Amat_\lambda^{-1}\Pmat\Amat_\lambda^{-1}, +\end{align*} +using the symmetry of $\Amat_\lambda$. For the ordering, note that for any $\mathbf{x}$, +$\mathbf{x}\transpose\Amat_\lambda^{-1}\Pmat\Amat_\lambda^{-1}\mathbf{x} = +(\Amat_\lambda^{-1}\mathbf{x})\transpose\Pmat(\Amat_\lambda^{-1}\mathbf{x}) \ge 0$ by +Lemma~\ref{lem:P}, so the subtracted term is positive semi-definite. It is zero for all +$\mathbf{x}$ only if $\lambda = 0$, since $\Amat_\lambda^{-1}$ is non-singular and $\Pmat \ne +\bm 0$. +\end{proof} + +\noindent The propagated covariance is therefore always the smaller of the two, and quoting it +alone at large $\lambda$ understates the true error. The posterior covariance +$\Amat_\lambda^{-1}$ accounts for what the damping costs as well as what it stabilises, and is +the more honest error bar to report whenever $\lambda > 0$. At $\lambda = 0$ the two coincide +and both equal the Cram\'er--Rao bound of Proposition~\ref{prop:crb}. + +\section{Summary} + +\begin{table}[h] +\centering +\small +\begin{tabular}{@{}ll@{}} +\toprule +Quantity & Expression \\ +\midrule +Mixing matrix & $M_{ij} = \Theta_{ij}\alpha_{j\to i} q_j / Q_i$ \\ +Exact inverse & $c_i = \left(Q_i d_i - \sum_{p \in A(i)}\alpha_{p\to i}Q_p d_p\right)/q_i$ \\ +Amplification & $\kappa_i = Q_i/q_i = 1/M_{ii}$ \\ +Penalty operator & $\Pmat = \Imat - n^{-1}\onevec\onevec\transpose$ \\ +Normal matrix & $\Amat_\lambda = \Mmat\transpose\Cd^{-1}\Mmat + \lambda\Pmat$ \\ +Estimator & $\Rmat_\lambda = \Amat_\lambda^{-1}\Mmat\transpose\Cd^{-1}$ \\ +Source covariance & $\Cc = \Rmat_\lambda\Cd\Rmat_\lambda\transpose$ \\ +Prediction covariance & $\Cdhat = \Mmat\Cc\Mmat\transpose$ \\ +Posterior covariance & $\Amat_\lambda^{-1}$ \\ +Resolution & $\Rmat_\lambda\Mmat = \Imat - \lambda\Amat_\lambda^{-1}\Pmat$ \\ +Bias & $\mathbf{b} = -\lambda\Amat_\lambda^{-1}\Pmat\cvec_{\text{true}}$ \\ +\bottomrule +\end{tabular} +\caption{\textbf{Principal results.} At $\lambda = 0$: $\Rmat_0 = \Mmat^{-1}$, the resolution +matrix is $\Imat$, the bias vanishes, the two covariances coincide, and $\Cc$ attains the +Cram\'er--Rao bound.} +\end{table} + +\noindent The practical picture is this. The mixing matrix is invertible, so the unmixing +problem always has an exact solution --- but that solution differences neighbouring +observations, amplifying their errors by $\kappa_i = Q_i/q_i$, and interpolates the noise. +Regularization trades this variance for bias, and because the penalty is a homogeneous +quadratic the resulting estimator stays exactly linear in the data. That linearity is what +delivers closed-form covariances, a resolution matrix, and effective degrees of freedom, none +of which are available from the log-ratio formulation without Monte Carlo resampling. The price +is the assumption of absolute rather than relative misfit --- an assumption that +Remark~\ref{rem:logratio} shows is substantially repaired by weighting with a proportional +error model, and which is in any case mild for tracers of low log-variance such as isotopic +ratios \parencite{blondes_practical_2016}. For tracers whose information really is relative +rather than absolute --- that is, genuinely compositional data +\parencite{aitchison_statistical_1986} --- the log-ratio formulation remains the better choice. + +Two caveats bear repeating. Regularization is not optional: by +Proposition~\ref{prop:degenerate} the unregularized problem reproduces the data exactly, noise +included. And the analytical machinery of Section~\ref{sec:propagation} is valid only while +the constraint $\cvec \ge 0$ is inactive; when sites clamp, the estimator ceases to be linear +and the reported covariances describe the unconstrained estimator instead. + +\printbibliography + +\end{document} diff --git a/docs/preprint.sty b/docs/preprint.sty new file mode 100644 index 0000000..36a3ca4 --- /dev/null +++ b/docs/preprint.sty @@ -0,0 +1,281 @@ +\NeedsTeXFormat{LaTeX2e} + +\ProcessOptions\relax + +% Use roman fonts for equations if available +\IfFileExists{txfonts.sty}% + {\AtEndOfClass{\RequirePackage{txfonts}% + \gdef\ttdefault{cmtt}% + \let\iint\relax + \let\iiint\relax + \let\iiiint\relax + \let\idotsint\relax + \let\openbox\relax}}{\RequirePackage{times}} + + % other font configurations +\renewcommand{\rmdefault}{ptm} +\renewcommand{\sfdefault}{phv} + + +% Define page geometry +\RequirePackage[verbose=true,letterpaper]{geometry} +\AtBeginDocument{ + \newgeometry{ + textheight=9.3in, + textwidth=7.1in, + top=0.96in, + headheight=22.17pt, + headsep=0.14in, + footskip=3.61pt + } +} + +% Add background margin tick marks +\RequirePackage{background} +\SetBgScale{1} +\SetBgAngle{0} +\SetBgColor{black} +\SetBgContents{% +\begin{tikzpicture}[remember picture,overlay] +\node at (-3.55in,5.2in) {\rule{.4pt}{.4in}}; +\node at (3.55in,5.2in) {\rule{.4pt}{.4in}}; +\node at (-3.95in,4.8in) {\rule{.4in}{.4pt}}; +\node at (3.95in,4.8in) {\rule{.4in}{.4pt}}; +\node at (-3.55in,-5.2in) {\rule{.4pt}{.4in}}; +\node at (3.55in,-5.2in) {\rule{.4pt}{.4in}}; +\node at (-3.95in,-4.8in) {\rule{.4in}{.4pt}}; +\node at (3.95in,-4.8in) {\rule{.4in}{.4pt}}; +\end{tikzpicture}} + +\widowpenalty=10000 +\clubpenalty=10000 +\flushbottom +\sloppy + +% Header options +\RequirePackage{fancyhdr} +\fancyhf{} +\pagestyle{fancy} +\renewcommand{\headrulewidth}{0pt} +\fancyheadoffset{0pt} +\lhead{\scshape Preprint -- Using convex optimization to efficiently apportion tracer and pollutant sources} % Replace manual title with \@title but if title is too long it can overrun line +\rhead{\thepage} +%\rfoot{\thepage} + +%Handling Keywords +\def\keywordname{{\bfseries \emph Keywords}}% +\def\keywords#1{\par\addvspace\medskipamount{\rightskip=0pt plus1cm +\def\and{\ifhmode\unskip\nobreak\fi\ $\cdot$ +}\noindent\keywordname\enspace\ignorespaces#1\par}} + +% font sizes with reduced leading +\renewcommand{\normalsize}{% + \@setfontsize\normalsize\@xpt\@xipt + \abovedisplayskip 7\p@ \@plus 2\p@ \@minus 5\p@ + \abovedisplayshortskip \z@ \@plus 3\p@ + \belowdisplayskip \abovedisplayskip + \belowdisplayshortskip 4\p@ \@plus 3\p@ \@minus 3\p@ +} +\normalsize +\renewcommand{\small}{% + \@setfontsize\small\@ixpt\@xpt + \abovedisplayskip 6\p@ \@plus 1.5\p@ \@minus 4\p@ + \abovedisplayshortskip \z@ \@plus 2\p@ + \belowdisplayskip \abovedisplayskip + \belowdisplayshortskip 3\p@ \@plus 2\p@ \@minus 2\p@ +} +\renewcommand{\footnotesize}{\@setfontsize\footnotesize\@ixpt\@xpt} +\renewcommand{\scriptsize}{\@setfontsize\scriptsize\@viipt\@viiipt} +\renewcommand{\tiny}{\@setfontsize\tiny\@vipt\@viipt} +\renewcommand{\large}{\@setfontsize\large\@xiipt{14}} +\renewcommand{\Large}{\@setfontsize\Large\@xivpt{16}} +\renewcommand{\LARGE}{\@setfontsize\LARGE\@xviipt{20}} +\renewcommand{\huge}{\@setfontsize\huge\@xxpt{23}} +\renewcommand{\Huge}{\@setfontsize\Huge\@xxvpt{28}} + +% sections with less space +\providecommand{\section}{} +\renewcommand{\section}{% + \@startsection{section}{1}{\z@}% + {-2.0ex \@plus -0.5ex \@minus -0.2ex}% + { 1.5ex \@plus 0.3ex \@minus 0.2ex}% + {\large\scshape\raggedright}% +} +\providecommand{\subsection}{} +\renewcommand{\subsection}{% + \@startsection{subsection}{2}{\z@}% + {-1.8ex \@plus -0.5ex \@minus -0.2ex}% + { 0.8ex \@plus 0.2ex}% + {\normalsize\itshape\raggedright}% +} +\providecommand{\subsubsection}{} +\renewcommand{\subsubsection}{% + \@startsection{subsubsection}{3}{\z@}% + {-1.5ex \@plus -0.5ex \@minus -0.2ex}% + { 0.5ex \@plus 0.2ex}% + {\normalsize\itshape\raggedright}% +} +\providecommand{\paragraph}{} +\renewcommand{\paragraph}{% + \@startsection{paragraph}{4}{\z@}% + {1.5ex \@plus 0.5ex \@minus 0.2ex}% + {-1em}% + {\normalsize\bf}% +} +\providecommand{\subparagraph}{} +\renewcommand{\subparagraph}{% + \@startsection{subparagraph}{5}{\z@}% + {1.5ex \@plus 0.5ex \@minus 0.2ex}% + {-1em}% + {\normalsize\bf}% +} +\providecommand{\subsubsubsection}{} +\renewcommand{\subsubsubsection}{% + \vskip5pt{\noindent\normalsize\rm\raggedright}% +} + +% float placement +\renewcommand{\topfraction }{0.85} +\renewcommand{\bottomfraction }{0.4} +\renewcommand{\textfraction }{0.1} +\renewcommand{\floatpagefraction}{0.7} + +\newlength{\@abovecaptionskip}\setlength{\@abovecaptionskip}{7\p@} +\newlength{\@belowcaptionskip}\setlength{\@belowcaptionskip}{\z@} + +\setlength{\abovecaptionskip}{\@abovecaptionskip} +\setlength{\belowcaptionskip}{\@belowcaptionskip} + +% swap above/belowcaptionskip lengths for tables +\renewenvironment{table} + {\setlength{\abovecaptionskip}{\@belowcaptionskip}% + \setlength{\belowcaptionskip}{\@abovecaptionskip}% + \@float{table}} + {\end@float} + +% footnote formatting +\setlength{\footnotesep }{6.65\p@} +\setlength{\skip\footins}{9\p@ \@plus 4\p@ \@minus 2\p@} +\renewcommand{\footnoterule}{\kern-3\p@ \hrule width 12pc \kern 2.6\p@} +\setcounter{footnote}{0} + +% paragraph formatting +\setlength{\parindent}{\z@} +\setlength{\parskip }{5.5\p@} + +% list formatting +\setlength{\topsep }{4\p@ \@plus 1\p@ \@minus 2\p@} +\setlength{\partopsep }{1\p@ \@plus 0.5\p@ \@minus 0.5\p@} +\setlength{\itemsep }{2\p@ \@plus 1\p@ \@minus 0.5\p@} +\setlength{\parsep }{2\p@ \@plus 1\p@ \@minus 0.5\p@} +\setlength{\leftmargin }{3pc} +\setlength{\leftmargini }{\leftmargin} +\setlength{\leftmarginii }{2em} +\setlength{\leftmarginiii}{1.5em} +\setlength{\leftmarginiv }{1.0em} +\setlength{\leftmarginv }{0.5em} +\def\@listi {\leftmargin\leftmargini} +\def\@listii {\leftmargin\leftmarginii + \labelwidth\leftmarginii + \advance\labelwidth-\labelsep + \topsep 2\p@ \@plus 1\p@ \@minus 0.5\p@ + \parsep 1\p@ \@plus 0.5\p@ \@minus 0.5\p@ + \itemsep \parsep} +\def\@listiii{\leftmargin\leftmarginiii + \labelwidth\leftmarginiii + \advance\labelwidth-\labelsep + \topsep 1\p@ \@plus 0.5\p@ \@minus 0.5\p@ + \parsep \z@ + \partopsep 0.5\p@ \@plus 0\p@ \@minus 0.5\p@ + \itemsep \topsep} +\def\@listiv {\leftmargin\leftmarginiv + \labelwidth\leftmarginiv + \advance\labelwidth-\labelsep} +\def\@listv {\leftmargin\leftmarginv + \labelwidth\leftmarginv + \advance\labelwidth-\labelsep} +\def\@listvi {\leftmargin\leftmarginvi + \labelwidth\leftmarginvi + \advance\labelwidth-\labelsep} + +% create title +\providecommand{\maketitle}{} +\renewcommand{\maketitle}{% + \par + \begingroup + \renewcommand{\thefootnote}{\fnsymbol{footnote}} + % for perfect author name centering + \renewcommand{\@makefnmark}{\hbox to \z@{$^{\@thefnmark}$\hss}} + % The footnote-mark was overlapping the footnote-text, + % added the following to fix this problem (MK) + \long\def\@makefntext##1{% + \parindent 1em\noindent + \hbox to 1.8em{\hss $\m@th ^{\@thefnmark}$}##1 + } + \thispagestyle{empty} + \vspace*{-0.5cm} + \@maketitle + \@thanks + %\@notice + \endgroup + \let\maketitle\relax + \let\thanks\relax +} + +% create title (includes both anonymized and non-anonymized versions) +\providecommand{\@maketitle}{} +\renewcommand{\@maketitle}{% + \vbox{% + \hsize\textwidth + \linewidth\hsize + \centering + {\LARGE\sc \@title\par} + \vskip 0.1in + \textsc{Submitted, non-peer reviewed manuscript, compiled \today}\\ + \def\And{% + \end{tabular}\hfil\linebreak[0]\hfil% + \begin{tabular}[t]{c}\bf\rule{\z@}{24\p@}\ignorespaces% + } + \def\AND{% + \end{tabular}\hfil\linebreak[4]\hfil% + \begin{tabular}[t]{c}\bf\rule{\z@}{24\p@}\ignorespaces% + } + \begin{tabular}[t]{c}\bf\rule{\z@}{24\p@}\@author\end{tabular}% + \vskip 0.2in + } +} + +% add conference notice to bottom of first page +\newcommand{\ftype@noticebox}{8} +\newcommand{\@notice}{% + % give a bit of extra room back to authors on first page + \enlargethispage{2\baselineskip}% + \@float{noticebox}[b]% + \footnotesize\@noticestring% + \end@float% +} + +% abstract styling +\renewenvironment{abstract} +{ + \centerline + {\large \bfseries \scshape Abstract} + \begin{quote} +} +{ + \end{quote} +} + + +% plain language environment styling + +\newenvironment{plainlang} +{ + \centerline + {\large \bfseries \scshape Plain Language Summary} + \begin{quote} +} +{ + \end{quote} +} +\endinput diff --git a/examples/unmix_linear_mwe.py b/examples/unmix_linear_mwe.py new file mode 100644 index 0000000..8ab84ce --- /dev/null +++ b/examples/unmix_linear_mwe.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 + +""" +Minimum working example for the *linear* unmixing solver. + +This is the counterpart to `unmix_mwe.py`. Where that script penalises relative (log-ratio) +misfit, this one penalises absolute misfit, which makes the forward model an exactly invertible +matrix and gives closed-form uncertainties: no Monte Carlo required. + +The script loads a sample network, inspects the diagnostics that say whether the linear +approach is appropriate for this network at all, sweeps the regularization strength to find the +elbow of the L-curve, then solves and maps both the recovered concentrations and their +propagated standard deviations. + +Run from the repository root: + python examples/unmix_linear_mwe.py +""" + +import logging +import warnings + +# pyre-fixme[21]: Could not find module `matplotlib.pyplot`. +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +import funmixer + +logging.getLogger().addHandler(logging.StreamHandler()) + +ELEMENT = "Mg" +RELATIVE_ERROR_PERCENT = 10.0 +REGULARIZATION_STRENGTH = 1.0 + + +def main() -> None: + sample_network, labels = funmixer.get_sample_graph( + flowdirs_filename="data/d8.asc", + sample_data_filename="data/sample_data.csv", + ) + + obs_data = pd.read_csv("data/sample_data.csv").drop(columns=["Bi", "S"]) + element_data = funmixer.get_element_obs(ELEMENT, obs_data) + + # ------------------------------------------------------------------ diagnostics + # Solve once with no regularization. This is the exact matrix inversion, and it is the + # honest test of whether the network supports a linear solution at all: if sites clamp at + # zero, or the amplification factors are large, the raw inversion is amplifying noise + # rather than resolving sources. + unregularized = funmixer.LinearSampleNetworkUnmixer(sample_network, use_regularization=False) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + raw = unregularized.solve(element_data, data_covariance=RELATIVE_ERROR_PERCENT) + + amplification = np.array(list(raw.amplification.values())) + unconstrained = np.array([raw.unconstrained_preds[n] for n in raw.node_order]) + print(f"Sites : {len(raw.node_order)}") + print(f"Condition number of M : {raw.condition_number:.4g}") + print( + "Noise amplification Q/q : " + f"median {np.median(amplification):.2f}, max {amplification.max():.1f}" + ) + print(f"Sites clamped at zero : {len(raw.clamped_nodes)}") + print( + "Unconstrained estimate range: " + f"{unconstrained.min():.0f} to {unconstrained.max():.0f} mg/kg" + ) + print( + " (values outside [0, 1e6] mg/kg mean the exact inversion is physically impossible,\n" + " which is the signal that regularization is needed)" + ) + + # ---------------------------------------------------------------------- L-curve + # `plot_sweep_of_regularizer_strength` is duck-typed on solve/get_misfit/get_roughness, + # so it works on the linear solver unchanged. + problem = funmixer.LinearSampleNetworkUnmixer(sample_network, use_regularization=True) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + funmixer.plot_sweep_of_regularizer_strength(problem, element_data, -3, 2, 11) + + # ------------------------------------------------------------------------ solve + solution = problem.solve( + element_data, + regularization_strength=REGULARIZATION_STRENGTH, + data_covariance=RELATIVE_ERROR_PERCENT, + ) + print(f"\nlambda : {solution.regularization_strength}") + print(f"Effective degrees of freedom: {solution.effective_dof:.1f}") + print(f"Sites clamped at zero : {len(solution.clamped_nodes)}") + relative_sigma = np.array( + [solution.upstream_std[n] / abs(solution.upstream_preds[n]) for n in solution.node_order] + ) + print(f"Median relative uncertainty : {100 * np.median(relative_sigma):.1f}%") + + funmixer.visualise_downstream( + pred_dict=solution.downstream_preds, obs_dict=element_data, element=ELEMENT + ) + + # ------------------------------------------------------------------------- maps + area_dict = funmixer.get_unique_upstream_areas(sample_network, labels) + concentration_map = funmixer.get_upstream_concentration_map(area_dict, solution.upstream_preds) + uncertainty_map = funmixer.get_upstream_concentration_map(area_dict, solution.upstream_std) + + _, axes = plt.subplots(1, 2, figsize=(15, 6)) + concentrations = axes[0].imshow(concentration_map) + axes[0].set_title(f"Recovered {ELEMENT} source concentration") + plt.colorbar(concentrations, ax=axes[0], label="mg/kg") + + uncertainties = axes[1].imshow(uncertainty_map) + axes[1].set_title(f"Propagated 1-sigma uncertainty ({RELATIVE_ERROR_PERCENT:.0f}% data error)") + plt.colorbar(uncertainties, ax=axes[1], label="mg/kg") + plt.show() + + +if __name__ == "__main__": + main() diff --git a/funmixer/__init__.py b/funmixer/__init__.py index 508b4b6..52bf07c 100644 --- a/funmixer/__init__.py +++ b/funmixer/__init__.py @@ -1,8 +1,13 @@ from .d8processing import * # noqa: F403 +from .linear_unmixer import ( + LinearFunmixerSolution, + LinearSampleNetworkUnmixer, +) from .network_unmixer import ( ELEMENT_LIST, ElementData, + FunmixerSolution, SampleNetworkUnmixer, SampleNode, forward_model, @@ -21,6 +26,9 @@ __all__ = [ "ElementData", "ELEMENT_LIST", + "FunmixerSolution", + "LinearFunmixerSolution", + "LinearSampleNetworkUnmixer", "get_element_obs", "get_unique_upstream_areas", "get_upstream_concentration_map", diff --git a/funmixer/linear_unmixer.py b/funmixer/linear_unmixer.py new file mode 100644 index 0000000..38eab66 --- /dev/null +++ b/funmixer/linear_unmixer.py @@ -0,0 +1,830 @@ +#!/usr/bin/env python3 +""" +Linear (least-squares) unmixing on a sample network. + +This module implements the linear formulation described in Appendix A of the funmixer +preprint. Where :class:`~funmixer.network_unmixer.SampleNetworkUnmixer` penalises *relative* +differences between predicted and observed concentrations (a convex surrogate for the +log-ratio misfit), this module penalises *absolute* differences. That choice is appropriate +when the tracer has low log-variance and it turns the whole problem into a linear system. + +The forward model +----------------- +The observation at site ``i`` is a flux-weighted mixture of the source concentrations of every +sub-basin upstream of (and including) ``i``:: + + d = M c, M_ij = Theta_ij * alpha_{j->i} * q_j / Q_i + +where ``Theta`` is the path matrix of the sample network with self-edges, ``q_j`` is the +material flux generated by sub-basin ``j``, ``Q_i = sum_j Theta_ij q_j`` is the total flux +passing site ``i``, and ``alpha_{j->i} = exp(-sum k*L)`` is the first-order decay accumulated +along the path from ``j`` to ``i``. Decay is applied to the tracer flux only and never to the +total flux, matching :func:`~funmixer.network_unmixer.forward_model`. + +Why this is exactly soluble +--------------------------- +funmixer defines exactly one sub-basin per sample site, so ``M`` is square (n x n). Ordered +topologically (upstream to downstream) it is **lower triangular with strictly positive +diagonal** ``M_ii = q_i / Q_i``, and is therefore always invertible. Two consequences: + +1. At ``lambda = 0`` the problem is not really least-squares at all -- it has an exact, + zero-residual solution ``c = M^-1 d``. There is no pseudo-inverse ambiguity, and weighted + (GLS) variants collapse to the same estimator, so "the estimator used when unconstrained" + is unambiguous. The inverse has a sparse closed form:: + + c_i = ( Q_i d_i - sum_{p in direct upstream nbrs of i} alpha_p Q_p d_p ) / q_i + + i.e. "tracer flux out, minus tracer flux in, over own flux". +2. The non-negativity constraint bites *exactly* when ``Q_i d_i < sum_p alpha_p Q_p d_p`` -- + when a site carries less tracer flux than its own tributaries deliver. That is precisely + where the method breaks down, and it is directly checkable via + :attr:`LinearFunmixerSolution.unconstrained_preds`. + +The per-node noise amplification of the inversion is ``1 / M_ii = Q_i / q_i``. A sample taken +just downstream of another, with little incremental drainage area between them, has a large +amplification factor and will be poorly resolved. This -- the disparity in *sub-basin areas* -- +is the real control on whether the linear approach is usable for a given network, and it is +reported as :attr:`LinearFunmixerSolution.amplification`. + +Regularization +-------------- +Optionally we penalise the variance of ``c``, i.e. deviations of ``c`` from its *own* mean:: + + minimize (M c - d)^T C_d^-1 (M c - d) + lambda ||P c||^2 subject to c >= 0 + +(the misfit term is ``||M c - d||^2`` when the solve is unweighted) + +with ``P = I - (1/n) 1 1^T`` the centering projector, so that ``||P c||^2 = sum_i (c_i - +c_bar)^2 = n * var(c)``. Note lambda therefore weights the *sum* of squared deviations rather +than their mean, so a factor of ``n`` is absorbed into it. Note also this differs deliberately +from +:class:`SampleNetworkUnmixer`, which is forced by convexity to shrink towards the mean +*observation*; the linear form has no such constraint and can shrink towards the mean *model*. + +A solution is guaranteed for every ``lambda >= 0``: the Hessian is ``2(M^T M + lambda P)``, and +although ``P`` alone is only positive semi-definite (its null space is ``span{1}``), ``M^T M`` +is positive definite because ``M`` is invertible. The objective is therefore strictly convex, +giving a unique global minimiser, and it stays unique under ``c >= 0``. Regularization strictly +improves the conditioning, so ``lambda > 0`` is always better behaved than ``lambda = 0``. +Uniqueness does rest on ``M`` being invertible -- ``lambda P`` could not rescue a singular +``M`` on its own. + +Error propagation +----------------- +Because ``||P c||^2`` is homogeneous in ``c``, the estimator carries **no constant offset** and +is purely linear in the data:: + + c_hat = R d, R = (M^T M + lambda P)^-1 M^T + C_c = R C_d R^T (covariance of the recovered source concentrations) + C_dhat = M C_c M^T (covariance of the modelled observations) + +At ``lambda = 0`` this reduces to ``M^-1 C_d M^-T`` exactly, so the regularized and +unregularized cases share a single code path. Likewise the modelled-observation covariance +reduces to ``C_d`` itself at ``lambda = 0``, since the fit is exact there. + +At ``lambda = 0`` the estimator is unbiased and ``C_c`` is the complete error. It also equals +``(M^T C_d^-1 M)^-1``, the Cramer-Rao lower bound, so the unregularized inversion is an +efficient estimator and no weighting can improve on it. + +Once ``lambda > 0`` the estimator is **biased**, and two different covariances become relevant. +Both are reported: + +* ``upstream_covariance``: the *propagated* (sampling) covariance ``R C_d R^T``. This answers + "if I repeated the measurements, how would my estimate scatter?" It is exactly the linear + propagation of ``C_d`` through the estimator actually used -- but it **shrinks as lambda + grows**, because a more heavily damped estimator is more stable while being more wrong. Used + alone at large lambda it badly understates the true error. +* ``upstream_posterior_covariance``: ``(M^T C_d^-1 M + lambda P)^-1``, the Bayesian posterior + covariance obtained by reading the penalty as a prior. This accounts for what the damping + costs as well as what it stabilises, and is the more honest error bar to quote at + ``lambda > 0``. The two are related exactly by:: + + R C_d R^T = (M^T C_d^-1 M + lambda P)^-1 - lambda A^-1 P A^-1, A = M^T C_d^-1 M + lambda P + + so the propagated covariance is always the smaller of the two, and they coincide at + ``lambda = 0``. + +The remaining, unquantifiable piece is the bias itself, +``bias = -lambda A^-1 P c_true``, which needs the unknown ``c_true``. As a proxy we report the +resolution matrix ``R M`` (the identity at ``lambda = 0``, degrading as ``lambda`` grows) and +its trace, the effective number of degrees of freedom. + +References +---------- +These are standard results in discrete linear inverse theory; the notation below follows the +first two. + +* Menke, W., *Geophysical Data Analysis: Discrete Inverse Theory*. Weighted least squares, + damped least squares, and the two results used most directly here: for any linear estimator + ``m_est = G^-g d`` the model covariance is ``C_m = G^-g C_d G^-g^T``, and the model + resolution matrix is ``R = G^-g G``. +* Aster, R., Borchers, B. and Thurber, C., *Parameter Estimation and Inverse Problems*. + Tikhonov regularization, the L-curve, resolution, and the bias-variance decomposition of the + regularized solution -- i.e. why the propagated covariance is not the whole error. +* Tarantola, A. (2005), *Inverse Problem Theory and Methods for Model Parameter Estimation*, + SIAM. The linear-Gaussian case, where the posterior covariance is + ``(G^T C_d^-1 G + C_m^-1)^-1``; this is the source of `upstream_posterior_covariance`. +* Hoerl, A. E. and Kennard, R. W. (1970), "Ridge Regression: Biased Estimation for + Nonorthogonal Problems", *Technometrics* 12(1), 55-67. The original bias-variance argument + for penalised estimators. +* Hansen, P. C. (1998), *Rank-Deficient and Discrete Ill-Posed Problems*, SIAM. Filter factors, + the L-curve, and general-form (seminorm) Tikhonov regularization -- the case here, since the + penalty operator ``P`` is singular. The uniqueness condition for general-form Tikhonov is + that the null spaces of the forward and penalty operators intersect only at zero, which + holds here because ``M`` is invertible. See also Elden, L. (1977), "Algorithms for the + regularization of ill-conditioned least squares problems", *BIT* 17, 134-145. +* Golub, G. H. and Van Loan, C. F., *Matrix Computations*, for the numerical treatment of + generalised least squares via whitening, and of Tikhonov problems via augmented systems. +""" + +import time +import warnings +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple, Union + +import cvxpy as cp +import networkx as nx +import numpy as np +import numpy.typing as npt + +from .network_unmixer import ( + SOLVERS, + ElementData, + ExportRateData, + RateConstantData, + SampleNode, + nx_get_downstream_node, + nx_topological_sort_with_data, +) + +FloatArray = npt.NDArray[np.float64] + +# A node is treated as clamped by the c >= 0 constraint when its unconstrained estimate falls +# below this (in mean-normalised units, where a typical concentration is 1). +CLAMP_TOLERANCE: float = 0.0 + +# Relative tolerance for the sanity check that CVXPY agrees with the analytical estimator on +# problems where no constraint is active. This is deliberately loose: interior-point solvers +# converge to ~1e-8, which an ill-conditioned mixing matrix amplifies by cond(M). A genuine +# mismatch between the two formulations would show up at order 1, not near this threshold. +ESTIMATOR_AGREEMENT_TOLERANCE: float = 1e-4 + +# Covariance input given as a bare number is interpreted as a relative error in *percent*, +# matching the `relative_error` convention of SampleNetworkUnmixer.solve_montecarlo. +DataCovariance = Union[float, int, ElementData, FloatArray, None] + + +@dataclass +class LinearFunmixerSolution: + """ + Results of a :class:`LinearSampleNetworkUnmixer` run. + + The first seven fields mirror :class:`~funmixer.network_unmixer.FunmixerSolution` so that + this object is a drop-in replacement for downstream plotting helpers. The remainder are + specific to the linear formulation. + + Attributes: + objective_value: The value of the (normalised-space) objective at the solution. + solver_name: Name of the solver that CVXPY used. + total_time: Wall-clock time for the whole solve, in seconds. + solve_time: Solver-reported solve time, in seconds. + setup_time: Solver-reported setup time, in seconds. + downstream_preds: Predicted concentration at each sample site, in the input units. + upstream_preds: Recovered source concentration for each sub-basin, in the input units. + node_order: Topological ordering of the sample sites. This is the row/column ordering + of every matrix below, and of any array-valued input. + regularization_strength: The value of lambda used. + weighted: Whether the misfit was weighted by the inverse data covariance. + upstream_covariance: (n, n) *propagated* covariance of `upstream_preds`, in + (input units)^2, or None if no `data_covariance` was supplied. Equal to + `R C_d R^T`: the linear propagation of the data covariance through the estimator + actually used. Exact, but it shrinks as lambda grows and so understates the true + error at large lambda -- see `upstream_posterior_covariance`. + upstream_std: Square root of the diagonal of `upstream_covariance`, per site. + upstream_posterior_covariance: (n, n) Bayesian posterior covariance + `(M^T C_d^-1 M + lambda P)^-1`, in (input units)^2. Available only for a weighted + solve with `data_covariance` supplied, since it needs `C_d^-1`. This accounts for + the cost of the damping as well as its stabilising effect, and is the more honest + error bar to quote when lambda > 0. Identical to `upstream_covariance` at + lambda = 0, and always at least as large. + upstream_posterior_std: Square root of the diagonal of + `upstream_posterior_covariance`, per site. + downstream_covariance: (n, n) covariance of `downstream_preds`, in (input units)^2, or + None. Equal to `M C_c M^T`; at lambda = 0 this is exactly `data_covariance`. + downstream_std: Square root of the diagonal of `downstream_covariance`, per site. + resolution_matrix: (n, n) resolution matrix `R M`. The identity at lambda = 0. + effective_dof: `trace(R M)`, the effective number of degrees of freedom. Equal to n at + lambda = 0 and decreasing as lambda grows. + unconstrained_preds: The raw estimate `R d` before the c >= 0 constraint is applied, in + the input units. Negative entries here show where, and by how much, the linear + model fails. + clamped_nodes: Sites whose unconstrained estimate was negative. If this is empty the + returned estimator is exactly linear and the propagated covariances are exact. + amplification: Per-site noise amplification `Q_i / q_i` of the inversion. + condition_number: Condition number of the mixing matrix `M`. + residual_norm: `||M c - d||_2` in normalised space. Zero at lambda = 0 with no clamping. + """ + + objective_value: float + solver_name: str + total_time: float + solve_time: Optional[float] + setup_time: Optional[float] + downstream_preds: ElementData + upstream_preds: ElementData + node_order: List[str] + regularization_strength: float + weighted: bool + unconstrained_preds: ElementData + clamped_nodes: List[str] + amplification: ElementData + condition_number: float + residual_norm: float + upstream_covariance: Optional[FloatArray] = None + upstream_std: Optional[ElementData] = None + upstream_posterior_covariance: Optional[FloatArray] = None + upstream_posterior_std: Optional[ElementData] = None + downstream_covariance: Optional[FloatArray] = None + downstream_std: Optional[ElementData] = None + resolution_matrix: Optional[FloatArray] = None + effective_dof: Optional[float] = None + + +class LinearSampleNetworkUnmixer: + """ + Unmix a network of concentration observations using a linear (least-squares) formulation. + + This is the Appendix A counterpart to :class:`~funmixer.network_unmixer.SampleNetworkUnmixer`: + it penalises absolute rather than relative misfit, which makes the forward model an exactly + invertible matrix and admits closed-form error propagation. See the module docstring for the + mathematics. + + The API deliberately mirrors :class:`SampleNetworkUnmixer`. In particular ``solve``, + ``get_misfit`` and ``get_roughness`` have compatible signatures, so + :func:`~funmixer.network_unmixer.plot_sweep_of_regularizer_strength` works on this class + unchanged. + + Note: + ``regularization_strength`` is **not** comparable between the two solvers. This one + weights a *squared* penalty on deviations from the mean *model*; the non-linear solver + weights an unsquared norm of deviations from the mean *observation*. Because the data + are mean-normalised internally, lambda here is dimensionless and does transfer between + elements and datasets. + """ + + def __init__( + self, + sample_network: nx.DiGraph, + use_regularization: bool = True, + rate_constants: Optional[RateConstantData] = None, + ) -> None: + """ + Initialise the LinearSampleNetworkUnmixer. + + Args: + sample_network: The sample network (see `get_sample_graph`). + use_regularization: Whether to include the model-variance penalty. If True, a + `regularization_strength` must be supplied to `solve`. + rate_constants: Per-site first-order decay constants k [L^-1]. Sites omitted from + the dict default to 0 (conservative behaviour). Can be overridden per-solve. + """ + self.sample_network = sample_network + self.node_order: List[str] = [ + name for name, _ in nx_topological_sort_with_data(sample_network) + ] + self._index: Dict[str, int] = {name: i for i, name in enumerate(self.node_order)} + self._n: int = len(self.node_order) + if self._n == 0: + raise ValueError("Cannot unmix an empty sample network.") + + # Sub-basin areas, normalised by their mean purely for numerical conditioning. This + # mirrors the `rltv_area` semantics of SampleNetworkUnmixer, but is kept in a local + # array: the `SampleNode` attributes are shared scratch space used by `forward_model` + # and `SampleNetworkUnmixer`, and writing to them would corrupt a live problem built + # on the same graph. + areas = np.array([self._node_data(name).area for name in self.node_order], dtype=np.float64) + if np.any(areas <= 0): + raise ValueError("All sub-basin areas must be strictly positive.") + self._rltv_area: FloatArray = areas / float(np.mean(areas)) + + self._use_regularization = use_regularization + self._rate_constants: Optional[RateConstantData] = rate_constants + + # State from the most recent solve, used by get_misfit / get_roughness / get_estimator. + self._M: Optional[FloatArray] = None + self._R: Optional[FloatArray] = None + self._c_norm: Optional[FloatArray] = None + self._d_norm: Optional[FloatArray] = None + self._M_solver: Optional[FloatArray] = None + self._d_solver: Optional[FloatArray] = None + self._weighted: bool = False + + # The compiled CVXPY problem, and the mixing matrix it was compiled for. + self._problem: Optional[cp.Problem] = None + self._problem_M: Optional[FloatArray] = None + + # ------------------------------------------------------------------ helpers + + def _node_data(self, name: str) -> SampleNode: + return self.sample_network.nodes[name]["data"] + + def _centering_projector(self) -> FloatArray: + """P = I - (1/n) 1 1^T, the projector onto the space of mean-zero vectors.""" + n = self._n + return np.eye(n) - np.ones((n, n)) / n + + def _ensure_problem(self, M: FloatArray) -> None: + """ + Build (or reuse) the CVXPY problem for a given mixing matrix. + + `d` and `lambda` are Parameters, so Monte Carlo draws, sweeps over lambda and sweeps + over elements all reuse a single compiled problem -- the DPP intent of + `SampleNetworkUnmixer`. `M`, by contrast, is embedded as a constant. Making it a + Parameter is technically DPP-compliant but puts n^2 parameters into the problem, which + CVXPY canonicalises far more slowly than it recompiles. `M` only changes when the + export rates or rate constants change, so rebuilding on that event is much cheaper. + + The objective uses *squared* norms, unlike `SampleNetworkUnmixer._build_problem`. That + is what makes the estimator linear, and is not optional here. + """ + if self._problem is not None and np.array_equal(M, self._problem_M): + return + + n = self._n + self._c_var = cp.Variable(n, nonneg=True) + self._d_param = cp.Parameter(n) + self._lam_param = cp.Parameter(nonneg=True) + + objective = cp.sum_squares(M @ self._c_var - self._d_param) + if self._use_regularization: + # ||P c||^2, i.e. the variance of the model about its own mean. + objective += self._lam_param * cp.sum_squares(self._c_var - cp.sum(self._c_var) / n) + + self._problem = cp.Problem(cp.Minimize(objective)) + assert self._problem.is_dcp(dpp=True) + self._problem_M = M.copy() + + # ---------------------------------------------------------- forward operator + + def _build_mixing_matrix( + self, + export_rates: Optional[ExportRateData], + rate_constants: Optional[RateConstantData], + ) -> Tuple[FloatArray, FloatArray, FloatArray]: + """ + Build the mixing matrix `M`, and the total (`Q`) and own (`q`) fluxes of each site. + + This performs exactly the topological sweep of + :func:`~funmixer.network_unmixer.forward_model`, but propagates a *row of coefficients* + instead of a scalar concentration. It is therefore identical to running the forward + model once per basis vector, and `M @ c == forward_model(network, c)` to machine + precision. + + Returns: + (M, total_flux, own_flux), all ordered by `self.node_order`. + """ + n = self._n + M = np.zeros((n, n), dtype=np.float64) + tracer_rows = np.zeros((n, n), dtype=np.float64) + total_flux = np.zeros(n, dtype=np.float64) + own_flux = np.zeros(n, dtype=np.float64) + + for name in self.node_order: + i = self._index[name] + rate = export_rates[name] if export_rates else 1.0 + flux = rate * self._rltv_area[i] + + own_flux[i] = flux + # Contributions from upstream have already been accumulated into these, because + # `node_order` is a topological sort. + total_flux[i] += flux + tracer_rows[i, i] += flux + + M[i, :] = tracer_rows[i, :] / total_flux[i] + + downstream = nx_get_downstream_node(self.sample_network, name) + if downstream is not None: + j = self._index[downstream] + length = self.sample_network[name][downstream]["length"] + k = rate_constants.get(name, 0.0) if rate_constants else 0.0 + alpha = float(np.exp(-k * length)) + # NB: decay applies to the tracer flux only, never to the total flux. + total_flux[j] += total_flux[i] + tracer_rows[j, :] += tracer_rows[i, :] * alpha + + return M, total_flux, own_flux + + def _solve_unregularized( + self, + d_norm: FloatArray, + total_flux: FloatArray, + own_flux: FloatArray, + rate_constants: Optional[RateConstantData], + ) -> FloatArray: + """ + Exact O(n + edges) solution of `M c = d` using the sparse closed form for `M^-1`. + + `c_i = (Q_i d_i - sum_p alpha_p Q_p d_p) / q_i`, i.e. the tracer flux leaving site `i` + minus the tracer flux delivered by its immediate upstream neighbours, divided by the + flux that site `i` itself generates. + """ + tracer_flux = total_flux * d_norm + c = np.empty(self._n, dtype=np.float64) + for name in self.node_order: + i = self._index[name] + accumulated = tracer_flux[i] + for upstream in self.sample_network.predecessors(name): + length = self.sample_network[upstream][name]["length"] + k = rate_constants.get(upstream, 0.0) if rate_constants else 0.0 + accumulated -= float(np.exp(-k * length)) * tracer_flux[self._index[upstream]] + c[i] = accumulated / own_flux[i] + return c + + def _whitener(self, covariance_norm: FloatArray) -> FloatArray: + """ + Build `L^-1`, where `C = L L^T`, so that `L^-1 e` has identity covariance. + + Pre-multiplying the system by this turns generalised least squares into ordinary + least squares: minimising `||L^-1(Mc - d)||^2` is minimising the Mahalanobis misfit + `(Mc - d)^T C^-1 (Mc - d)`. + """ + n = self._n + diagonal = np.diag(covariance_norm) + if np.allclose(covariance_norm, np.diag(diagonal)): + # The common case. Much cheaper, and avoids a needless factorisation. + if np.any(diagonal <= 0): + raise ValueError( + "data_covariance has a zero or negative variance, so it cannot be used as " + "a weighting. Pass weighted=False, or give every site a positive variance." + ) + return np.diag(1.0 / np.sqrt(diagonal)) + try: + factor = np.linalg.cholesky(covariance_norm) + except np.linalg.LinAlgError as exc: + raise ValueError( + "data_covariance must be positive definite to be used as a weighting. " + "Pass weighted=False to fall back to an unweighted inversion." + ) from exc + return np.linalg.solve(factor, np.eye(n)) + + def _build_estimator( + self, M: FloatArray, lam: float, whitener: Optional[FloatArray] + ) -> FloatArray: + """ + Build the linear estimator `R` such that `c_hat = R d`. + + Unweighted, this is `R = (M^T M + lambda P)^-1 M^T`. Weighted by `W = C_d^-1` it is + the generalised-least-squares estimator `R = (M^T W M + lambda P)^-1 M^T W`, which we + form via the whitening transform `M_w = L^-1 M` so there is only one code path. + + At lambda = 0 this is exactly `M^-1` **whatever the weighting**: `M` is square and + invertible, so `(M^T W M)^-1 M^T W = M^-1 W^-1 M^-T M^T W = M^-1`. The unregularized + fit is exact and passes through every observation, so there is nothing for a weighting + to trade off. Weighting only changes the answer once `lambda > 0`. + """ + n = self._n + if lam == 0.0: + return np.linalg.solve(M, np.eye(n)) + M_w = M if whitener is None else whitener @ M + # M_w^T M_w + lambda P is symmetric positive definite: M_w^T M_w is PD because both + # M and the whitener are invertible, and lambda P is PSD. + normal_matrix = M_w.T @ M_w + lam * self._centering_projector() + adjoint = M_w.T if whitener is None else M_w.T @ whitener + return np.linalg.solve(normal_matrix, adjoint) + + # ------------------------------------------------------------- data handling + + def _observations_to_array(self, observation_data: ElementData) -> FloatArray: + missing = set(self.node_order) - set(observation_data) + if missing: + raise ValueError(f"No observation supplied for site(s): {sorted(missing)}") + unknown = set(observation_data) - set(self.node_order) + if unknown: + raise ValueError(f"Observation supplied for unknown site(s): {sorted(unknown)}") + return np.array([observation_data[name] for name in self.node_order], dtype=np.float64) + + def _coerce_data_covariance( + self, data_covariance: DataCovariance, d: FloatArray + ) -> Optional[FloatArray]: + """ + Normalise the many accepted spellings of `data_covariance` into an (n, n) matrix. + + Accepted forms, all in the *input* units of the observations: + None -- skip error propagation. + a bare number -- a relative error in percent, so sigma_i = d_i * value / 100. + This matches `SampleNetworkUnmixer.solve_montecarlo`'s + `relative_error`, so the two are directly comparable. + dict -- per-site standard deviations. + 1-D array (n,) -- per-site standard deviations, in `node_order`. + 2-D array (n, n) -- a full covariance matrix, in `node_order`. + """ + if data_covariance is None: + return None + + n = self._n + if isinstance(data_covariance, (float, int)) and not isinstance(data_covariance, bool): + sigma = d * float(data_covariance) / 100.0 + return np.diag(sigma**2) + + if isinstance(data_covariance, dict): + missing = set(self.node_order) - set(data_covariance) + if missing: + raise ValueError(f"No data_covariance entry for site(s): {sorted(missing)}") + sigma = np.array( + [float(data_covariance[name]) for name in self.node_order], dtype=np.float64 + ) + return np.diag(sigma**2) + + array = np.asarray(data_covariance, dtype=np.float64) + if array.ndim == 1: + if array.shape != (n,): + raise ValueError(f"1-D data_covariance must have shape ({n},), got {array.shape}") + return np.diag(array**2) + if array.ndim == 2: + if array.shape != (n, n): + raise ValueError( + f"2-D data_covariance must have shape ({n}, {n}), got {array.shape}" + ) + return array + raise ValueError("data_covariance must be a number, dict, 1-D array or 2-D array.") + + def _to_dict(self, values: FloatArray) -> ElementData: + return {name: float(values[self._index[name]]) for name in self.node_order} + + # -------------------------------------------------------------------- solve + + def solve( + self, + observation_data: ElementData, + export_rates: Optional[ExportRateData] = None, + rate_constants: Optional[RateConstantData] = None, + regularization_strength: Optional[float] = None, + data_covariance: DataCovariance = None, + weighted: bool = True, + solver: str = "clarabel", + compute_resolution: bool = True, + warm_start: bool = False, + ) -> LinearFunmixerSolution: + """ + Solve the linear unmixing problem. + + The observations are divided by their arithmetic mean before solving and the results + multiplied back afterwards. The problem is homogeneous, so this has no effect on the + answer -- it exists purely to keep the numbers near unity for the solver, and to make + `regularization_strength` dimensionless and comparable across datasets. + + Args: + observation_data: Observed concentration at each sample site. + export_rates: Export rate for each sub-catchment. Defaults to 1 everywhere (i.e. + homogeneous erosion/runoff). + rate_constants: Per-site first-order decay constants k [L^-1]. Defaults to those + given at construction, or to 0 (conservative) for omitted sites. + regularization_strength: lambda, weighting the squared penalty on the variance of + the model. Required if the problem was built with `use_regularization=True`. + Not comparable to the equivalent argument of `SampleNetworkUnmixer.solve`. + data_covariance: Uncertainty on the observations, enabling closed-form error + propagation. A bare number is read as a relative error in percent; see + `_coerce_data_covariance` for the other accepted forms. If None, the covariance + fields of the result are left as None. + weighted: Whether to weight the misfit by the inverse data covariance, i.e. to + solve generalised rather than ordinary least squares. Has an effect only when + `data_covariance` is supplied *and* `regularization_strength` is non-zero -- + see `_build_estimator` for why weighting cannot change an unregularized + solution. Strongly recommended whenever the observations span a wide range and + the errors are proportional, which makes them badly heteroscedastic. + solver: CVXPY solver to use (default "clarabel"). + compute_resolution: Whether to compute the resolution matrix and effective degrees + of freedom. Costs one O(n^3) matrix product; set False for large networks. + warm_start: Seed the solver from the previous solution. + + Returns: + A :class:`LinearFunmixerSolution`. + + Raises: + Exception: If regularization is enabled but no strength was supplied. + ValueError: If the observations do not match the sites in the network. + + Note: + When no site is clamped by the `c >= 0` constraint, the returned point estimate is + the analytical estimator `R d` exactly, not the solver's approximation to it. This + guarantees that the estimate and the propagated covariances come from the same + linear operator, which is what makes the error propagation meaningful. When sites + *are* clamped the estimator is no longer linear; the covariance is still reported + (it assumes the unconstrained estimator) but a warning is issued. + """ + if self._use_regularization and regularization_strength is None: + raise Exception("WARNING: Regularizer terms present but no strength assigned.") + if solver not in SOLVERS: + raise Exception( + f"Solver {solver} not supported. Supported solvers are {list(SOLVERS.keys())}" + ) + + lam = float(regularization_strength) if self._use_regularization else 0.0 + if lam < 0: + raise ValueError("regularization_strength must be non-negative.") + if rate_constants is None: + rate_constants = self._rate_constants + + start_time = time.time() + + d = self._observations_to_array(observation_data) + d_bar = float(np.mean(d)) + if d_bar <= 0: + raise ValueError("Mean observed concentration must be positive.") + d_norm = d / d_bar + + M, total_flux, own_flux = self._build_mixing_matrix(export_rates, rate_constants) + + # Weighting. The covariance is scaled into normalised space so that the whitened + # residuals are dimensionless, which in turn keeps lambda dimensionless: the objective + # trades a chi-squared-like misfit against the variance of the model. + C_d = self._coerce_data_covariance(data_covariance, d) + is_weighted = weighted and C_d is not None + whitener: Optional[FloatArray] = None + if is_weighted: + assert C_d is not None + whitener = self._whitener(C_d / d_bar**2) + + # The unconstrained estimate. At lambda = 0 the sparse closed form is both exact and + # O(n); otherwise we go via R, which we need for the covariance anyway. + need_estimator = (C_d is not None) or compute_resolution or lam > 0.0 + R: Optional[FloatArray] = ( + self._build_estimator(M, lam, whitener) if need_estimator else None + ) + if lam == 0.0: + c_unconstrained = self._solve_unregularized( + d_norm, total_flux, own_flux, rate_constants + ) + else: + assert R is not None + c_unconstrained = R @ d_norm + + # Always solve via CVXPY: it supplies the solver statistics, and it is the only route + # to a correct answer when the c >= 0 constraint is active. The whitened system is + # handed to it so that it minimises the same objective as the analytical estimator. + M_solver = M if whitener is None else whitener @ M + d_solver = d_norm if whitener is None else whitener @ d_norm + self._ensure_problem(M_solver) + self._d_param.value = d_solver + self._lam_param.value = lam + objective_value = self._problem.solve(**SOLVERS[solver], warm_start=warm_start) + if self._problem.status != "optimal": + warnings.warn( + f"CVXPY returned status '{self._problem.status}'.", RuntimeWarning, stacklevel=2 + ) + + clamped_mask = c_unconstrained < CLAMP_TOLERANCE + clamped_nodes = [name for name in self.node_order if clamped_mask[self._index[name]]] + + if not clamped_nodes: + # The objective is strictly convex, so `c_unconstrained` is *the* unconstrained + # global minimiser. It is also feasible, hence a fortiori the constrained optimum. + # Return it exactly rather than the solver's approximation, so that the point + # estimate and the propagated covariance share one estimator. + c_norm = c_unconstrained + solver_c = np.asarray(self._c_var.value, dtype=np.float64) + scale = max(float(np.max(np.abs(c_norm))), 1.0) + discrepancy = float(np.max(np.abs(c_norm - solver_c))) / scale + if discrepancy > ESTIMATOR_AGREEMENT_TOLERANCE: + warnings.warn( + f"Analytical and CVXPY solutions differ by {discrepancy:.3e} (relative) " + "despite no active constraints. The mixing matrix is probably badly " + f"conditioned (cond = {np.linalg.cond(M):.3e}), which amplifies solver " + "tolerance; treat the recovered concentrations with caution.", + RuntimeWarning, + stacklevel=2, + ) + else: + c_norm = np.asarray(self._c_var.value, dtype=np.float64) + warnings.warn( + f"{len(clamped_nodes)} site(s) clamped at the lower bound " + f"({clamped_nodes if len(clamped_nodes) <= 10 else clamped_nodes[:10] + ['...']}). " + "The constrained estimator is not linear in the data, so any propagated " + "covariance assumes the unconstrained estimator and will overstate uncertainty " + "at these sites.", + UserWarning, + stacklevel=2, + ) + + d_pred_norm = M @ c_norm + residual_norm = float(np.linalg.norm(d_pred_norm - d_norm)) + + # Error propagation. The estimator maps observations to concentrations in *any* + # consistent scaling, so the mean-normalisation cancels exactly and we can work + # directly in the input units: C_c = R C_d R^T, C_dpred = M C_c M^T. + upstream_covariance: Optional[FloatArray] = None + upstream_std: Optional[ElementData] = None + downstream_covariance: Optional[FloatArray] = None + downstream_std: Optional[ElementData] = None + upstream_posterior_covariance: Optional[FloatArray] = None + upstream_posterior_std: Optional[ElementData] = None + if C_d is not None: + assert R is not None + upstream_covariance = _symmetrise(R @ C_d @ R.T) + downstream_covariance = _symmetrise(M @ upstream_covariance @ M.T) + upstream_std = self._to_dict(np.sqrt(np.clip(np.diag(upstream_covariance), 0, None))) + downstream_std = self._to_dict( + np.sqrt(np.clip(np.diag(downstream_covariance), 0, None)) + ) + if is_weighted: + # (M^T C_d^-1 M + lambda P)^-1, built from the whitened operator so it needs + # no explicit inverse of C_d. M_solver is in normalised space, so the result + # is scaled back into the input units by d_bar^2. + normal_matrix = M_solver.T @ M_solver + lam * self._centering_projector() + upstream_posterior_covariance = _symmetrise( + np.linalg.solve(normal_matrix, np.eye(self._n)) * d_bar**2 + ) + upstream_posterior_std = self._to_dict( + np.sqrt(np.clip(np.diag(upstream_posterior_covariance), 0, None)) + ) + + resolution_matrix: Optional[FloatArray] = None + effective_dof: Optional[float] = None + if compute_resolution: + assert R is not None + resolution_matrix = R @ M + effective_dof = float(np.trace(resolution_matrix)) + + # Cache for get_misfit / get_roughness / get_estimator. + self._M, self._R, self._c_norm, self._d_norm = M, R, c_norm, d_norm + self._M_solver, self._d_solver, self._weighted = M_solver, d_solver, is_weighted + + return LinearFunmixerSolution( + objective_value=float(objective_value), + solver_name=self._problem.solver_stats.solver_name, + total_time=time.time() - start_time, + solve_time=self._problem.solver_stats.solve_time, + setup_time=self._problem.solver_stats.setup_time, + downstream_preds=self._to_dict(d_pred_norm * d_bar), + upstream_preds=self._to_dict(c_norm * d_bar), + node_order=list(self.node_order), + regularization_strength=lam, + weighted=is_weighted, + unconstrained_preds=self._to_dict(c_unconstrained * d_bar), + clamped_nodes=clamped_nodes, + amplification=self._to_dict(total_flux / own_flux), + condition_number=float(np.linalg.cond(M)), + residual_norm=residual_norm, + upstream_covariance=upstream_covariance, + upstream_std=upstream_std, + upstream_posterior_covariance=upstream_posterior_covariance, + upstream_posterior_std=upstream_posterior_std, + downstream_covariance=downstream_covariance, + downstream_std=downstream_std, + resolution_matrix=resolution_matrix, + effective_dof=effective_dof, + ) + + # ------------------------------------------------------------- introspection + + def get_misfit(self) -> float: + """ + The data misfit of the most recent solve, in normalised space. + + This is the quantity the objective actually minimises: `||M c - d||_2` for an + unweighted solve, and the Mahalanobis norm `sqrt((Mc-d)^T C_d^-1 (Mc-d))` for a + weighted one. The latter is a chi distance, so a well-fitting weighted solve gives + roughly `sqrt(n)`. + + Together with `get_roughness`, this makes the class compatible with + :func:`~funmixer.network_unmixer.plot_sweep_of_regularizer_strength`. + """ + if self._M_solver is None or self._c_norm is None or self._d_solver is None: + raise Exception("No solution available: call solve() first.") + return float(np.linalg.norm(self._M_solver @ self._c_norm - self._d_solver)) + + def get_roughness(self) -> float: + """ + The model roughness `||P c||_2` of the most recent solve, in normalised space. + + This is the total deviation of the recovered concentrations from their own mean, the + same quantity the regularization penalises. + """ + if self._c_norm is None: + raise Exception("No solution available: call solve() first.") + return float(np.linalg.norm(self._c_norm - np.mean(self._c_norm))) + + def get_mixing_matrix(self) -> Tuple[FloatArray, List[str]]: + """Return the mixing matrix `M` of the most recent solve, and its node ordering.""" + if self._M is None: + raise Exception("No solution available: call solve() first.") + return self._M, list(self.node_order) + + def get_estimator(self) -> Tuple[FloatArray, List[str]]: + """ + Return the linear estimator `R` of the most recent solve, and its node ordering. + + `c = R d` holds in both normalised and input units. Users wanting to propagate their + own uncertainties can use `C_c = R C_d R^T` and `C_dpred = M C_c M^T` directly. + """ + if self._R is None: + raise Exception( + "No estimator available: call solve() first, with data_covariance set or " + "compute_resolution=True." + ) + return self._R, list(self.node_order) + + +def _symmetrise(matrix: FloatArray) -> FloatArray: + """Remove the antisymmetric round-off that accumulates in triple matrix products.""" + return 0.5 * (matrix + matrix.T) diff --git a/funmixer/network_unmixer.py b/funmixer/network_unmixer.py index 224b981..9e57589 100644 --- a/funmixer/network_unmixer.py +++ b/funmixer/network_unmixer.py @@ -237,8 +237,12 @@ def nx_get_downstream_data(G: nx.DiGraph, x: str) -> Optional[SampleNode]: Exception: If there is more than one downstream neighbor. """ s = nx_get_downstream_node(G, x) - if s: + # NB: an explicit `is not None` check, not a truthiness test. Node names may be integers, + # and a node legitimately named `0` would otherwise be reported as having no downstream + # neighbour -- silently disconnecting the outlet of any integer-labelled network. + if s is not None: return cast(SampleNode, G.nodes[s]["data"]) + return None def plot_network(G: nx.DiGraph) -> None: diff --git a/tests/linear_unmixer_test.py b/tests/linear_unmixer_test.py new file mode 100644 index 0000000..9c33088 --- /dev/null +++ b/tests/linear_unmixer_test.py @@ -0,0 +1,705 @@ +# pyre-ignore-all-errors[56] +""" +Tests for the linear (least-squares) unmixing solver. + +The network generators mirror those in `random_networks_test.py`. Tolerances here are far +tighter than for the non-linear solver because, at lambda = 0, the inversion is exact rather +than iterative. +""" + +import warnings +from typing import Callable, Optional + +import networkx as nx +import numpy as np +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +import funmixer +from funmixer.linear_unmixer import LinearSampleNetworkUnmixer + +MINIMUM_CONC = 1.0 +MAXIMUM_CONC = 1e2 +MINIMUM_AREA = 1.0 +MAXIMUM_AREA = 1e2 +MAX_RATE_PARAMETER = 3.0 + +# The inversion is exact, so we can demand far more than the 1% used for the non-linear solver. +EXACT_TOLERANCE = 1e-8 + + +def draw_random_log_uniform(min_val: float, max_val: float) -> float: + return float(np.exp(np.random.uniform(np.log(min_val), np.log(max_val), 1))[0]) + + +def generate_balanced_sample_network( + branching_factor: int, height: int, areas: Callable[[], float] +) -> nx.DiGraph: + """A balanced tree of sample sites, with flow directed towards the root.""" + G = nx.balanced_tree(branching_factor, height, create_using=nx.DiGraph) + G = nx.reverse(G) + for u, v in G.edges: + G[u][v]["length"] = 1.0 + for node in G.nodes: + G.nodes[node]["data"] = funmixer.SampleNode( + name=node, + area=areas(), + downstream_node=funmixer.nx_get_downstream_node(G, node), + x=-1, + y=-1, + total_upstream_area=0, + label=0, + upstream_nodes=[], + distance_downstream=1.0, + ) + return G + + +def default_network(branching_factor: int = 2, height: int = 3, seed: int = 0) -> nx.DiGraph: + np.random.seed(seed) + return generate_balanced_sample_network( + branching_factor, height, lambda: draw_random_log_uniform(MINIMUM_AREA, MAXIMUM_AREA) + ) + + +def random_concentrations(network: nx.DiGraph) -> funmixer.ElementData: + return {node: draw_random_log_uniform(MINIMUM_CONC, MAXIMUM_CONC) for node in network.nodes} + + +# --------------------------------------------------------------- forward operator + + +@given( + branching_factor=st.integers(min_value=1, max_value=4), + height=st.integers(min_value=1, max_value=4), + rate_constant=st.floats(min_value=0.0, max_value=MAX_RATE_PARAMETER), + conservative=st.booleans(), + use_export_rates=st.booleans(), +) +@settings(deadline=None, max_examples=25) +def test_mixing_matrix_matches_forward_model( + branching_factor: int, + height: int, + rate_constant: float, + conservative: bool, + use_export_rates: bool, +) -> None: + """`M @ c` must reproduce `forward_model` exactly: it is the same sweep, vectorised.""" + network = default_network(branching_factor, height) + rate_constants = None if conservative else {n: rate_constant for n in network.nodes} + export_rates = ( + {n: draw_random_log_uniform(0.1, 10.0) for n in network.nodes} if use_export_rates else None + ) + + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + M, node_order = ( + problem._build_mixing_matrix(export_rates, rate_constants)[0], + problem.node_order, + ) + + upstream = random_concentrations(network) + expected = funmixer.forward_model( + sample_network=network, + upstream_concentrations=upstream, + export_rates=export_rates, + rate_constants=rate_constants, + ) + + c = np.array([upstream[n] for n in node_order]) + predicted = M @ c + for i, node in enumerate(node_order): + assert np.isclose(predicted[i], expected[node], rtol=1e-12) + + +@given( + branching_factor=st.integers(min_value=1, max_value=4), + height=st.integers(min_value=1, max_value=4), + conservative=st.booleans(), +) +@settings(deadline=None, max_examples=25) +def test_mixing_matrix_structure(branching_factor: int, height: int, conservative: bool) -> None: + """M is lower triangular in topological order, with a positive diagonal.""" + network = default_network(branching_factor, height) + rate_constants = None if conservative else {n: 1.0 for n in network.nodes} + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + M, _, _ = problem._build_mixing_matrix(None, rate_constants) + + assert np.allclose(M, np.tril(M)), "M must be lower triangular in topological order" + assert np.all(np.diag(M) > 0), "M must have a strictly positive diagonal" + if conservative: + # Conservative mixing weights are a convex combination, so rows sum to one. + assert np.allclose(M.sum(axis=1), 1.0) + else: + # Decay removes tracer from the numerator only, so rows sum to less than one. + assert np.all(M.sum(axis=1) <= 1.0 + 1e-12) + + +def test_closed_form_inverse_matches_dense_inverse() -> None: + """The sparse O(n) closed form must agree with an explicit dense inverse.""" + network = default_network(branching_factor=3, height=3) + rate_constants = {n: 0.7 for n in network.nodes} + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + M, total_flux, own_flux = problem._build_mixing_matrix(None, rate_constants) + + rng = np.random.default_rng(0) + d = rng.uniform(1.0, 10.0, size=M.shape[0]) + + closed_form = problem._solve_unregularized(d, total_flux, own_flux, rate_constants) + assert np.allclose(closed_form, np.linalg.solve(M, d), rtol=1e-10) + assert np.allclose(problem._build_estimator(M, 0.0, None), np.linalg.inv(M), rtol=1e-10) + + +# ------------------------------------------------------------------- round trip + + +@given( + branching_factor=st.integers(min_value=1, max_value=4), + height=st.integers(min_value=1, max_value=4), + rate_constant=st.floats(min_value=0.0, max_value=MAX_RATE_PARAMETER), + conservative=st.booleans(), +) +@settings(deadline=None, max_examples=25) +def test_exact_round_trip( + branching_factor: int, height: int, rate_constant: float, conservative: bool +) -> None: + """forward_model -> solve must recover the source concentrations exactly at lambda = 0.""" + network = default_network(branching_factor, height) + rate_constants = None if conservative else {n: rate_constant for n in network.nodes} + upstream = random_concentrations(network) + downstream = funmixer.forward_model( + sample_network=network, upstream_concentrations=upstream, rate_constants=rate_constants + ) + + problem = LinearSampleNetworkUnmixer( + network, use_regularization=False, rate_constants=rate_constants + ) + solution = problem.solve(downstream) + + assert solution.clamped_nodes == [] + assert solution.residual_norm < 1e-10 + for node in network.nodes: + assert np.isclose(solution.upstream_preds[node], upstream[node], rtol=EXACT_TOLERANCE) + assert np.isclose(solution.downstream_preds[node], downstream[node], rtol=EXACT_TOLERANCE) + + +def test_returned_estimate_is_the_analytical_one() -> None: + """ + On interior problems the returned estimate must be the analytical estimator exactly, and + CVXPY must independently agree with it. This is what makes the covariance meaningful. + """ + network = default_network(branching_factor=2, height=3) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + + for lam, use_reg in [(None, False), (1e-3, True), (1.0, True)]: + problem = LinearSampleNetworkUnmixer(network, use_regularization=use_reg) + solution = problem.solve(downstream, regularization_strength=lam) + assert solution.clamped_nodes == [] + # `unconstrained_preds` is the analytical R @ d; the returned estimate must equal it. + for node in network.nodes: + assert solution.upstream_preds[node] == solution.unconstrained_preds[node] + # And CVXPY, solving the same problem numerically, must land in the same place. + R, node_order = problem.get_estimator() + d = np.array([downstream[n] for n in node_order]) + assert np.allclose(R @ d, [solution.upstream_preds[n] for n in node_order], rtol=1e-6) + + +# ---------------------------------------------------------------- regularization + + +def test_regularized_solve_matches_augmented_least_squares() -> None: + """ + The Cholesky/normal-equations path must agree with the numerically safer augmented + least-squares form, which never forms M^T M. + """ + network = default_network(branching_factor=3, height=2) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + for lam in [1e-4, 1e-2, 1.0, 10.0]: + solution = problem.solve(downstream, regularization_strength=lam) + M, node_order = problem.get_mixing_matrix() + n = len(node_order) + d = np.array([downstream[node] for node in node_order]) + P = np.eye(n) - np.ones((n, n)) / n + + augmented_operator = np.vstack([M, np.sqrt(lam) * P]) + augmented_data = np.concatenate([d, np.zeros(n)]) + expected, *_ = np.linalg.lstsq(augmented_operator, augmented_data, rcond=None) + + got = np.array([solution.upstream_preds[node] for node in node_order]) + assert np.allclose(got, expected, rtol=1e-7) + + +def test_small_lambda_reproduces_unregularized_solution() -> None: + network = default_network(branching_factor=2, height=3) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + + unregularized = LinearSampleNetworkUnmixer(network, use_regularization=False).solve(downstream) + barely = LinearSampleNetworkUnmixer(network, use_regularization=True).solve( + downstream, regularization_strength=1e-12 + ) + for node in network.nodes: + assert np.isclose( + barely.upstream_preds[node], unregularized.upstream_preds[node], rtol=1e-5 + ) + + +def test_large_lambda_drives_model_to_a_constant() -> None: + """ + As lambda grows the variance penalty dominates and every c_i must converge to a common + value. For a conservative network M is row-stochastic, so that value is the mean + observation. + """ + network = default_network(branching_factor=2, height=3) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + solution = problem.solve(downstream, regularization_strength=1e10) + + values = np.array(list(solution.upstream_preds.values())) + assert np.allclose(values, values[0], rtol=1e-4), "all c_i should collapse to one value" + assert np.isclose(values[0], np.mean(list(downstream.values())), rtol=1e-4) + + +def test_effective_dof_decreases_with_lambda() -> None: + """Resolution degrades monotonically with lambda, starting from the identity at lambda=0.""" + network = default_network(branching_factor=2, height=3) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + n = network.number_of_nodes() + + unregularized = LinearSampleNetworkUnmixer(network, use_regularization=False).solve(downstream) + assert np.allclose(unregularized.resolution_matrix, np.eye(n), atol=1e-8) + assert np.isclose(unregularized.effective_dof, n) + + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + dofs = [ + problem.solve(downstream, regularization_strength=lam).effective_dof + for lam in [1e-6, 1e-3, 1e-1, 1.0, 1e2] + ] + assert all(a > b for a, b in zip(dofs, dofs[1:])), f"not monotonically decreasing: {dofs}" + # The penalty has a one-dimensional null space (constant vectors), which is never damped. + assert dofs[-1] > 1.0 + + +# --------------------------------------------------------------------- weighting + + +def test_weighting_cannot_change_an_unregularized_solution() -> None: + """ + At lambda = 0 the fit is exact and passes through every observation, so there is nothing + for a weighting to trade off: (M^T W M)^-1 M^T W = M^-1 for any invertible W. This is a + mathematical identity, not an approximation, and is worth pinning down. + """ + network = default_network(branching_factor=2, height=3) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + + # A deliberately extreme weighting: five orders of magnitude between the best- and + # worst-known sites. + rng = np.random.default_rng(3) + sigmas = {node: 10 ** rng.uniform(-2, 3) for node in network.nodes} + + unweighted = problem.solve(downstream, data_covariance=sigmas, weighted=False) + weighted = problem.solve(downstream, data_covariance=sigmas, weighted=True) + for node in network.nodes: + assert np.isclose(weighted.upstream_preds[node], unweighted.upstream_preds[node], rtol=1e-9) + # The covariance is C_c = R C_d R^T with the same R, so it too must be unchanged. + assert np.allclose(weighted.upstream_covariance, unweighted.upstream_covariance, rtol=1e-9) + + +def test_weighting_changes_a_regularized_solution() -> None: + """Once lambda > 0 the weighting decides which observations the penalty is allowed to + sacrifice, so it must change the answer.""" + network = default_network(branching_factor=2, height=3) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + + unweighted = problem.solve( + downstream, regularization_strength=0.1, data_covariance=10.0, weighted=False + ) + weighted = problem.solve( + downstream, regularization_strength=0.1, data_covariance=10.0, weighted=True + ) + assert weighted.weighted and not unweighted.weighted + difference = max( + abs(weighted.upstream_preds[n] - unweighted.upstream_preds[n]) for n in network.nodes + ) + scale = max(abs(v) for v in unweighted.upstream_preds.values()) + assert difference / scale > 1e-3, "weighting should materially change the solution" + + +def test_weighted_estimator_matches_explicit_gls_formula() -> None: + """The whitening transform must reproduce the textbook GLS normal equations.""" + network = default_network(branching_factor=3, height=2) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + + rng = np.random.default_rng(11) + sigmas = {node: 10 ** rng.uniform(-1, 2) for node in network.nodes} + lam = 0.37 + solution = problem.solve( + downstream, regularization_strength=lam, data_covariance=sigmas, weighted=True + ) + + M, node_order = problem.get_mixing_matrix() + n = len(node_order) + d = np.array([downstream[node] for node in node_order]) + d_bar = float(np.mean(d)) + # The estimator is built in normalised space, so the covariance must be scaled to match. + C_d = np.diag(np.array([sigmas[node] for node in node_order]) ** 2) / d_bar**2 + W = np.linalg.inv(C_d) + P = np.eye(n) - np.ones((n, n)) / n + + expected_R = np.linalg.solve(M.T @ W @ M + lam * P, M.T @ W) + got_R, _ = problem.get_estimator() + assert np.allclose(got_R, expected_R, rtol=1e-8) + assert np.allclose( + [solution.upstream_preds[node] for node in node_order], expected_R @ d, rtol=1e-7 + ) + + +def test_unregularized_weighted_covariance_is_the_cramer_rao_bound() -> None: + """ + At lambda = 0 the estimator is `M^-1`, so `C_c = M^-1 C_d M^-T = (M^T C_d^-1 M)^-1` -- the + inverse Fisher information. The unregularized inversion is therefore an efficient + estimator, attaining the Cramer-Rao lower bound, and no weighting can improve on it. + """ + network = default_network(branching_factor=3, height=2) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + + rng = np.random.default_rng(2) + sigmas = {node: 10 ** rng.uniform(-1, 2) for node in network.nodes} + solution = problem.solve(downstream, data_covariance=sigmas) + + M, node_order = problem.get_mixing_matrix() + C_d = np.diag(np.array([sigmas[node] for node in node_order]) ** 2) + cramer_rao = np.linalg.inv(M.T @ np.linalg.inv(C_d) @ M) + assert np.allclose(solution.upstream_covariance, cramer_rao, rtol=1e-8) + + +def test_posterior_and_propagated_covariance_agree_at_zero_lambda() -> None: + """At lambda = 0 the penalty contributes nothing, so both covariances must coincide.""" + network = default_network(branching_factor=3, height=2) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + solution = problem.solve(downstream, regularization_strength=0.0, data_covariance=7.0) + assert np.allclose( + solution.upstream_posterior_covariance, solution.upstream_covariance, rtol=1e-7 + ) + + +def test_posterior_covariance_exceeds_propagated_when_regularized() -> None: + """ + The propagated covariance measures only how the estimate scatters; the posterior also + counts what the damping costs. They differ by exactly `lambda A^-1 P A^-1`, so the + posterior is always the larger -- and quoting the propagated one alone at large lambda + understates the error. + """ + network = default_network(branching_factor=3, height=2) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + + previous_ratio = 1.0 + for lam in [1e-3, 1e-1, 1.0, 10.0]: + solution = problem.solve( + downstream, regularization_strength=lam, data_covariance=7.0, weighted=True + ) + propagated = np.diag(solution.upstream_covariance) + posterior = np.diag(solution.upstream_posterior_covariance) + assert np.all(posterior >= propagated - 1e-12) + + # Check the exact identity relating the two. + M, node_order = problem.get_mixing_matrix() + n = len(node_order) + d_bar = float(np.mean([downstream[node] for node in node_order])) + C_d = np.diag(np.array([downstream[node] for node in node_order]) * 0.07) ** 2 + A = M.T @ np.linalg.inv(C_d / d_bar**2) @ M + lam * (np.eye(n) - np.ones((n, n)) / n) + A_inv = np.linalg.inv(A) + expected = (A_inv - lam * A_inv @ (np.eye(n) - np.ones((n, n)) / n) @ A_inv) * d_bar**2 + assert np.allclose(solution.upstream_covariance, expected, rtol=1e-6) + + # The gap widens with lambda: damping buys stability at the cost of accuracy. + ratio = float(np.mean(np.sqrt(posterior / propagated))) + assert ratio >= previous_ratio - 1e-9 + previous_ratio = ratio + assert previous_ratio > 1.5, "the two covariances should diverge substantially by lambda=10" + + +def test_posterior_covariance_absent_when_unweighted() -> None: + """The posterior form needs C_d^-1, so it is only defined for a weighted solve.""" + network = default_network(branching_factor=2, height=2) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + solution = problem.solve( + downstream, regularization_strength=0.1, data_covariance=5.0, weighted=False + ) + assert solution.upstream_covariance is not None + assert solution.upstream_posterior_covariance is None + + +def test_weighted_misfit_is_a_chi_distance() -> None: + """With a correct error model, a weighted solve at small lambda should give a misfit of + order sqrt(n) -- one standard deviation per observation.""" + network = default_network(branching_factor=2, height=4) + n = network.number_of_nodes() + upstream = random_concentrations(network) + truth = funmixer.forward_model(network, upstream) + + rng = np.random.default_rng(5) + relative_error = 0.1 + noisy = {node: value * (1 + relative_error * rng.normal()) for node, value in truth.items()} + sigmas = {node: relative_error * value for node, value in noisy.items()} + + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + problem.solve(noisy, regularization_strength=1e-6, data_covariance=sigmas) + # Essentially unregularized, so the fit is near-exact and the chi distance near zero. + assert problem.get_misfit() < np.sqrt(n) + + +# ------------------------------------------------------------ error propagation + + +@pytest.mark.parametrize("lam, use_reg", [(None, False), (1e-2, True)]) +def test_covariance_matches_monte_carlo(lam: Optional[float], use_reg: bool) -> None: + """ + The load-bearing test: the closed-form covariance must reproduce what you get by actually + perturbing the data and re-solving. + + The network and data are chosen so the solution stays comfortably interior, so that the + estimator really is linear and the analytical result is exact. + """ + np.random.seed(7) + network = generate_balanced_sample_network(2, 2, lambda: 1.0) + # A near-uniform source field keeps the differenced solution well away from zero. + upstream = {node: 100.0 + np.random.uniform(-5, 5) for node in network.nodes} + downstream = funmixer.forward_model(network, upstream) + + relative_error = 2.0 # percent + problem = LinearSampleNetworkUnmixer(network, use_regularization=use_reg) + solution = problem.solve( + downstream, regularization_strength=lam, data_covariance=relative_error + ) + assert solution.clamped_nodes == [] + + R, node_order = problem.get_estimator() + M, _ = problem.get_mixing_matrix() + d = np.array([downstream[node] for node in node_order]) + sigma = d * relative_error / 100.0 + + rng = np.random.default_rng(42) + draws = 20000 + perturbed = d[None, :] + rng.normal(0.0, sigma, size=(draws, len(d))) + sampled_c = perturbed @ R.T + sampled_d = sampled_c @ M.T + + empirical_c = np.cov(sampled_c, rowvar=False) + empirical_d = np.cov(sampled_d, rowvar=False) + + # Compare standard deviations: with 20k draws the standard error on a std is ~0.5%. + assert np.allclose( + np.sqrt(np.diag(empirical_c)), np.sqrt(np.diag(solution.upstream_covariance)), rtol=0.05 + ) + assert np.allclose( + np.sqrt(np.diag(empirical_d)), + np.sqrt(np.diag(solution.downstream_covariance)), + rtol=0.05, + ) + # And the full matrices, scaled by the diagonal to make the comparison dimensionless. + scale = np.outer( + np.sqrt(np.diag(solution.upstream_covariance)), + np.sqrt(np.diag(solution.upstream_covariance)), + ) + assert np.allclose(empirical_c / scale, solution.upstream_covariance / scale, atol=0.05) + + +def test_unregularized_downstream_covariance_equals_data_covariance() -> None: + """ + At lambda = 0 the fit is exact, so the modelled observations are the observations and their + covariance must be the input covariance unchanged. + """ + network = default_network(branching_factor=2, height=3) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + solution = problem.solve(downstream, data_covariance=5.0) + + d = np.array([downstream[node] for node in solution.node_order]) + expected = np.diag((d * 5.0 / 100.0) ** 2) + assert np.allclose(solution.downstream_covariance, expected, rtol=1e-6, atol=1e-12) + + +def test_covariance_input_forms_agree() -> None: + """A scalar relative error, a dict of sigmas and an explicit matrix must agree.""" + network = default_network(branching_factor=2, height=2) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + + scalar = problem.solve(downstream, data_covariance=10.0) + sigmas = {node: value * 0.1 for node, value in downstream.items()} + as_dict = problem.solve(downstream, data_covariance=sigmas) + as_array = problem.solve( + downstream, + data_covariance=np.array([sigmas[node] for node in problem.node_order]), + ) + as_matrix = problem.solve( + downstream, + data_covariance=np.diag(np.array([sigmas[node] for node in problem.node_order]) ** 2), + ) + for other in [as_dict, as_array, as_matrix]: + assert np.allclose(scalar.upstream_covariance, other.upstream_covariance, rtol=1e-10) + + +def test_covariance_scales_quadratically_with_units() -> None: + """ + The problem is homogeneous, so rescaling the observations must rescale concentrations + linearly and covariances quadratically. This checks the mean-normalisation cancels exactly. + """ + network = default_network(branching_factor=2, height=3) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + + base = problem.solve(downstream, data_covariance=10.0) + factor = 1e4 # e.g. converting fractions to mg/kg + scaled = problem.solve( + {node: value * factor for node, value in downstream.items()}, data_covariance=10.0 + ) + for node in network.nodes: + assert np.isclose( + scaled.upstream_preds[node], base.upstream_preds[node] * factor, rtol=1e-9 + ) + assert np.allclose( + scaled.upstream_covariance, base.upstream_covariance * factor**2, rtol=1e-9 + ) + + +# ------------------------------------------------------------------ constraints + + +def test_clamping_warns_and_respects_the_lower_bound() -> None: + """ + Data that violate the mixing model -- a downstream site carrying less tracer flux than its + tributaries deliver -- must clamp at zero, warn, and still return a feasible solution. + """ + network = generate_balanced_sample_network(2, 1, lambda: 1.0) + node_order = [ + name for name, _ in funmixer.network_unmixer.nx_topological_sort_with_data(network) + ] + # Leaves are rich, the outlet is poor: impossible under conservative mixing. + root = node_order[-1] + observations = {node: (100.0 if node != root else 1.0) for node in network.nodes} + + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + with pytest.warns(UserWarning, match="clamped at the lower bound"): + solution = problem.solve(observations) + + assert solution.clamped_nodes, "expected at least one site to clamp" + assert all(value >= -1e-8 for value in solution.upstream_preds.values()) + assert min(solution.unconstrained_preds.values()) < 0, "the raw estimate should go negative" + assert solution.residual_norm > 0, "a clamped fit cannot be exact" + + +def test_clamping_warning_mentions_covariance_caveat() -> None: + network = generate_balanced_sample_network(2, 1, lambda: 1.0) + node_order = [ + name for name, _ in funmixer.network_unmixer.nx_topological_sort_with_data(network) + ] + observations = {node: (100.0 if node != node_order[-1] else 1.0) for node in network.nodes} + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + with pytest.warns(UserWarning, match="overstate uncertainty"): + problem.solve(observations, data_covariance=5.0) + + +# ------------------------------------------------------------------- diagnostics + + +def test_amplification_is_total_over_own_flux() -> None: + """Equal-area sub-basins in a balanced binary tree give amplification = subtree size.""" + network = generate_balanced_sample_network(2, 2, lambda: 1.0) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + solution = LinearSampleNetworkUnmixer(network, use_regularization=False).solve(downstream) + + for node in network.nodes: + subtree_size = len(nx.ancestors(network, node)) + 1 + assert np.isclose(solution.amplification[node], subtree_size, rtol=1e-10) + + +def test_misfit_and_roughness_are_available_for_the_l_curve() -> None: + """ + `plot_sweep_of_regularizer_strength` is duck-typed on solve/get_misfit/get_roughness, so + those three must work together on this class without any change to that helper. + """ + network = default_network(branching_factor=2, height=3) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + + misfits, roughnesses = [], [] + for lam in np.logspace(-4, 2, 7): + problem.solve(downstream, regularization_strength=lam) + misfits.append(problem.get_misfit()) + roughnesses.append(problem.get_roughness()) + + # The classic L-curve trade-off: misfit rises, roughness falls, both monotonically. + assert all(a <= b + 1e-9 for a, b in zip(misfits, misfits[1:])), misfits + assert all(a >= b - 1e-9 for a, b in zip(roughnesses, roughnesses[1:])), roughnesses + + +def test_requires_regularization_strength_when_enabled() -> None: + network = default_network(branching_factor=2, height=2) + problem = LinearSampleNetworkUnmixer(network, use_regularization=True) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + with pytest.raises(Exception, match="no strength assigned"): + problem.solve(downstream) + + +def test_rejects_mismatched_observations() -> None: + network = default_network(branching_factor=2, height=2) + problem = LinearSampleNetworkUnmixer(network, use_regularization=False) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + incomplete = {k: v for k, v in list(downstream.items())[:-1]} + with pytest.raises(ValueError, match="No observation supplied"): + problem.solve(incomplete) + + +def test_does_not_disturb_a_live_nonlinear_problem() -> None: + """ + `SampleNode` fields are shared scratch space. Building and solving a linear problem on a + graph must not corrupt a `SampleNetworkUnmixer` built on the same graph. + """ + network = default_network(branching_factor=2, height=2) + upstream = random_concentrations(network) + downstream = funmixer.forward_model(network, upstream) + + nonlinear = funmixer.SampleNetworkUnmixer(network, use_regularization=False) + before = nonlinear.solve(downstream).upstream_preds + + linear = LinearSampleNetworkUnmixer(network, use_regularization=False) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + linear.solve(downstream) + + after = nonlinear.solve(downstream).upstream_preds + for node in network.nodes: + assert np.isclose(before[node], after[node], rtol=1e-6) diff --git a/tests/linear_vs_nonlinear_benchmark.py b/tests/linear_vs_nonlinear_benchmark.py new file mode 100644 index 0000000..1bc286a --- /dev/null +++ b/tests/linear_vs_nonlinear_benchmark.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 + +""" +How well do the two solvers recover known sources as the source range widens? + +Setup +----- +A single synthetic network of 100 sample sites, with sub-basin areas varying by only +/-10% +and a uniform export rate, so the inversion is as well conditioned as it realistically gets. +Source concentrations are drawn log-uniformly spanning a controlled number of orders of +magnitude (0.5 up to 6), forward-modelled to observations, and then corrupted with relative +Gaussian noise. The error model is assumed known, so every solver that can use it is +handed `C_d = diag((f d)^2)` for the relative error f. + +Why the range matters +--------------------- +The two solvers differ in what they call misfit. The non-linear solver penalises *relative* +differences, which is scale-free. The linear solver penalises *absolute* differences, which is +not -- unweighted, it will spend all its effort on the largest observations and ignore the +small ones. That should not matter at 0.5 orders of magnitude and should matter enormously at +6. + +Weighting by the inverse data covariance is the fix. With proportional errors, sigma_i = f*d_i, +so the whitened residual is (Mc - d)_i / (f d_i) = (1/f)((Mc)_i/d_i - 1): a *relative* misfit. +The weighted linear solver is therefore the linearisation of the non-linear one, and this +script is largely a test of how far that linearisation stretches. + +Run from the repository root: + python tests/linear_vs_nonlinear_benchmark.py +""" + +import warnings +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional + +# pyre-fixme[21]: Could not find module `matplotlib.pyplot`. +import matplotlib.pyplot as plt +import networkx as nx +import numpy as np + +import funmixer +from funmixer.linear_unmixer import LinearSampleNetworkUnmixer + +N_SITES = 100 +AREA_VARIATION = 0.10 # sub-basin areas drawn uniformly from 1 +/- this fraction +RELATIVE_ERROR = 0.20 # relative error on the observations, assumed known +GEOMETRIC_MEAN_CONC = 1000.0 # mg/kg +ORDERS_OF_MAGNITUDE = [0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0] +N_REPEATS = 5 +LAMBDA_GRID = np.logspace(-6, 5, 23) + + +def build_network(n_sites: int, area_variation: float, seed: int) -> nx.DiGraph: + """A random tree of sample sites with near-uniform sub-basin areas, flowing to node 0.""" + rng = np.random.default_rng(seed) + undirected = nx.random_labeled_tree(n_sites, seed=seed) + # Orient every edge towards the root, so edges point downstream. + network = nx.DiGraph((child, parent) for parent, child in nx.bfs_edges(undirected, 0)) + for upstream, downstream in network.edges: + network[upstream][downstream]["length"] = 1.0 + for node in network.nodes: + network.nodes[node]["data"] = funmixer.SampleNode( + name=node, + area=float(rng.uniform(1 - area_variation, 1 + area_variation)), + downstream_node=funmixer.nx_get_downstream_node(network, node), + x=-1, + y=-1, + total_upstream_area=0, + label=0, + upstream_nodes=[], + distance_downstream=1.0, + ) + return network + + +def draw_sources(network: nx.DiGraph, orders: float, seed: int) -> funmixer.ElementData: + """Log-uniform source concentrations spanning `orders` decades about a fixed geometric mean.""" + rng = np.random.default_rng(seed) + half = orders / 2.0 + exponents = rng.uniform(-half, half, size=network.number_of_nodes()) + return { + node: float(GEOMETRIC_MEAN_CONC * 10**exponent) + for node, exponent in zip(network.nodes, exponents) + } + + +@dataclass +class Score: + """How well a set of predictions matches the truth, measured in log space.""" + + median_log_error: float # median |log10(pred/true)|, i.e. typical factor-of error + rms_log_error: float # RMS of the same, so bad sites are not hidden by good ones + interior_median_log_error: float # median over non-leaf sites only + fraction_within_2x: float + n_nonpositive: int + chosen_lambda: float = 0.0 + + @classmethod + def compare( + cls, + predicted: funmixer.ElementData, + truth: funmixer.ElementData, + interior: Optional[np.ndarray] = None, + ) -> "Score": + keys = list(truth) + true_values = np.array([truth[k] for k in keys]) + pred_values = np.array([predicted[k] for k in keys]) + n_nonpositive = int(np.sum(pred_values <= 0)) + # Clamped sites come back as exactly zero, where a log ratio is undefined. Floor them + # far below any plausible source so they register as a large, finite error rather + # than being silently dropped from the average. + floor = 1e-6 * float(np.exp(np.mean(np.log(true_values)))) + log_error = np.abs(np.log10(np.maximum(pred_values, floor) / true_values)) + # A leaf site is recovered exactly (c_i = d_i), so its error is just the data error + # whatever the solver does. Reporting the interior separately stops those sites from + # masking the differences between methods. + interior_errors = log_error if interior is None else log_error[interior] + return cls( + median_log_error=float(np.median(log_error)), + rms_log_error=float(np.sqrt(np.mean(log_error**2))), + interior_median_log_error=float( + np.median(interior_errors) if interior_errors.size else np.nan + ), + fraction_within_2x=float(np.mean(log_error < np.log10(2.0))), + n_nonpositive=n_nonpositive, + ) + + +def best_over_lambda( + solve_for_lambda: Callable[[float], Optional[funmixer.ElementData]], + truth: funmixer.ElementData, + lambdas: np.ndarray, + interior: np.ndarray, +) -> Score: + """ + Score a solver at its *oracle-best* regularization strength. + + Choosing lambda against the known truth is of course cheating, but it is cheating equally + for every solver, and it isolates what we actually want to measure: the best each method + can possibly do, rather than how well some particular lambda-selection heuristic works. + """ + best: Optional[Score] = None + for lam in lambdas: + predictions = solve_for_lambda(float(lam)) + if predictions is None: + continue + score = Score.compare(predictions, truth, interior) + score.chosen_lambda = float(lam) + if best is None or score.rms_log_error < best.rms_log_error: + best = score + assert best is not None, "every lambda failed to solve" + return best + + +def run_trial(orders: float, seed: int) -> Dict[str, Score]: + """Generate one synthetic dataset and invert it every way we know how.""" + network = build_network(N_SITES, AREA_VARIATION, seed) + truth = draw_sources(network, orders, seed) + clean = funmixer.forward_model(network, truth) + + rng = np.random.default_rng(seed + 10_000) + observed = { + node: value * max(1.0 + RELATIVE_ERROR * rng.normal(), 1e-3) + for node, value in clean.items() + } + # The assumed-known error model. + sigmas = {node: RELATIVE_ERROR * value for node, value in observed.items()} + interior = np.array([network.in_degree(node) > 0 for node in truth], dtype=bool) + + scores: Dict[str, Score] = {} + + linear_unreg = LinearSampleNetworkUnmixer(network, use_regularization=False) + scores["linear, lambda=0"] = Score.compare( + linear_unreg.solve(observed, data_covariance=sigmas).upstream_preds, truth, interior + ) + + linear = LinearSampleNetworkUnmixer(network, use_regularization=True) + scores["linear, weighted"] = best_over_lambda( + lambda lam: linear.solve( + observed, regularization_strength=lam, data_covariance=sigmas, weighted=True + ).upstream_preds, + truth, + LAMBDA_GRID, + interior, + ) + scores["linear, unweighted"] = best_over_lambda( + lambda lam: linear.solve( + observed, regularization_strength=lam, data_covariance=sigmas, weighted=False + ).upstream_preds, + truth, + LAMBDA_GRID, + interior, + ) + + nonlinear_unreg = funmixer.SampleNetworkUnmixer(network, use_regularization=False) + scores["nonlinear, lambda=0"] = Score.compare( + nonlinear_unreg.solve(observed).upstream_preds, truth, interior + ) + + nonlinear = funmixer.SampleNetworkUnmixer(network, use_regularization=True) + + def solve_nonlinear(lam: float) -> Optional[funmixer.ElementData]: + try: + return nonlinear.solve(observed, regularization_strength=lam).upstream_preds + except Exception: + return None + + scores["nonlinear, regularized"] = best_over_lambda( + solve_nonlinear, truth, LAMBDA_GRID, interior + ) + return scores + + +METHODS: List[str] = [ + "linear, lambda=0", + "linear, unweighted", + "linear, weighted", + "nonlinear, lambda=0", + "nonlinear, regularized", +] + + +def main() -> None: + print( + f"{N_SITES} sites, areas +/-{AREA_VARIATION:.0%}, uniform export rate, " + f"{RELATIVE_ERROR:.0%} relative error on observations ({N_REPEATS} repeats).\n" + "Regularized methods are shown at their oracle-best lambda.\n" + "Score is the median factor-of error: 10^median|log10(pred/true)|.\n" + ) + + results: Dict[str, Dict[float, List[Score]]] = { + m: {o: [] for o in ORDERS_OF_MAGNITUDE} for m in METHODS + } + for orders in ORDERS_OF_MAGNITUDE: + for repeat in range(N_REPEATS): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + trial = run_trial(orders, seed=repeat) + for method, score in trial.items(): + results[method][orders].append(score) + print(f" ...done {orders} orders of magnitude") + + header = f"\n{'method':<24}" + "".join(f"{o:>12.1f}" for o in ORDERS_OF_MAGNITUDE) + for label, extract, fmt in [ + ("median factor-of error (all sites)", lambda s: 10**s.median_log_error, "{:>12.2f}"), + ( + "median factor-of error (interior sites only)", + lambda s: 10**s.interior_median_log_error, + "{:>12.2f}", + ), + ("RMS factor-of error", lambda s: 10**s.rms_log_error, "{:>12.2f}"), + ("fraction within 2x", lambda s: s.fraction_within_2x, "{:>11.0%} "), + ("sites hitting zero", lambda s: s.n_nonpositive, "{:>12.1f}"), + ("oracle-best lambda", lambda s: s.chosen_lambda, "{:>12.2g}"), + ]: + print(f"\n=== {label} (source range, orders of magnitude) ===") + print(header.strip("\n")) + for method in METHODS: + row = f"{method:<24}" + for orders in ORDERS_OF_MAGNITUDE: + row += fmt.format(float(np.mean([extract(s) for s in results[method][orders]]))) + print(row) + + _, axis = plt.subplots(figsize=(8, 5)) + for method in METHODS: + axis.plot( + ORDERS_OF_MAGNITUDE, + [ + float(np.mean([10**s.rms_log_error for s in results[method][o]])) + for o in ORDERS_OF_MAGNITUDE + ], + marker="o", + label=method, + ) + axis.axhline(1.0, color="grey", linestyle=":", label="perfect recovery") + axis.set_yscale("log") + axis.set_xlabel("Source range (orders of magnitude)") + axis.set_ylabel("RMS factor-of error in recovered concentration") + axis.set_title( + f"Source recovery vs. source range\n({N_SITES} sites, {RELATIVE_ERROR:.0%} data error)" + ) + axis.legend() + plt.tight_layout() + plt.savefig("linear_vs_nonlinear_benchmark.png", dpi=150) + print("\nSaved plot to linear_vs_nonlinear_benchmark.png") + + +if __name__ == "__main__": + main()