[WIP] WLS_sparse2 solver

CHANGELOG.rst

@@ -1,6 +1,20 @@
 
 Changelog
 =========
+1.0.3 (UNRELEASED)
+------------------
+New features
+
+*
+
+Bug fixes
+
+*
+
+Others
+
+*
+
 1.0.2 (2020-05-04)
 ------------------
 * Same as v1.0.1

src/dtscalibration/calibrate_utils.py

@@ -1,6 +1,8 @@
 # coding=utf-8
 import numpy as np
 import scipy.sparse as sp
+from scipy.linalg import inv as scipy_inv
+from scipy.linalg import lstsq as scipy_lstsq
 from scipy.sparse import linalg as ln
 
 
@@ -73,6 +75,8 @@ def calibration_single_ended_solver(
     Parameters
     ----------
     ds : DataStore
+        Should have sections and reference temperature timeseries already
+        configured.
     st_var : float, array-like, optional
         If `None` use ols calibration. If `float` the variance of the noise
         from the Stokes detector is described with a single value. Or when the
@@ -88,7 +92,7 @@ def calibration_single_ended_solver(
     calc_cov : bool
         whether to calculate the covariance matrix. Required for calculation
         of confidence boundaries. But uses a lot of memory.
-    solver : {'sparse', 'stats', 'external', 'external_split'}
+    solver : {'sparse', 'sparse2', 'stats', 'external', 'external_split'}
         Always use sparse to save memory. The statsmodel can be used to validate
         sparse solver. `external` returns the matrices that would enter the
         matrix solver (Eq.37). `external_split` returns a dictionary with
@@ -291,6 +295,14 @@ def calibration_single_ended_solver(
             p_sol, p_var = wls_sparse(
                 X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=verbose)
 
+    elif solver == 'sparse2':
+        if calc_cov:
+            p_sol, p_var, p_cov = wls_sparse2(
+                X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=verbose)
+        else:
+            p_sol, p_var = wls_sparse2(
+                X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=verbose)
+
     elif solver == 'stats':
         if calc_cov:
             p_sol, p_var, p_cov = wls_stats(
@@ -344,6 +356,8 @@ def calibration_double_ended_solver(
     Parameters
     ----------
     ds : DataStore
+        Should have sections and reference temperature timeseries already
+        configured.
     st_var : float, array-like, optional
         If `None` use ols calibration. If `float` the variance of the noise
         from the Stokes detector is described with a single value. Or when the
@@ -371,7 +385,7 @@ def calibration_double_ended_solver(
     calc_cov : bool
         whether to calculate the covariance matrix. Required for calculation
         of confidence boundaries. But uses a lot of memory.
-    solver : {'sparse', 'stats', 'external', 'external_split'}
+    solver : {'sparse', 'sparse2', 'stats', 'external', 'external_split'}
         Always use sparse to save memory. The statsmodel can be used to validate
         sparse solver. `external` returns the matrices that would enter the
         matrix solver (Eq.37). `external_split` returns a dictionary with
@@ -560,6 +574,8 @@ def calibration_double_ended_solver(
 
     if solver == 'sparse':
         solver_fun = wls_sparse
+    elif solver == 'sparse2':
+        solver_fun = wls_sparse2
     elif solver == 'stats':
         solver_fun = wls_stats
     elif solver == 'external':
@@ -1191,6 +1207,171 @@ def wls_sparse(
             return p_sol, p_var
 
 
+def wls_sparse2(
+        X,
+        y,
+        w=1.,
+        p_sol_prior=None,
+        p_cov_prior=None,
+        x0=None,
+        adjust_p_cov_flag=False,
+        calc_cov=True,
+        lapack_driver=None,
+        verbose=False):
+    """
+    Solves the normal equations of X . p_sol = y with weights. Supports a priori
+    information. The parameter estimates from a previous calibration instance
+    and their covariances can be passed to `p_sol_prior` and `p_cov_prior`.
+    Allows for sequential calibration in chunks.
+
+    Normally, the X matrix is very tall (nobs>>npar), more observations than
+    unknowns. By using the normal equations, the coefficient matrix reduces to a
+    square matrix of shape (npar, npar), and the observation matrix to (npar,).
+    This results in a much smaller system to solve.
+
+    Parameters
+    ----------
+    X : array
+        Coefficient matrix of shape (nobs, npar). Sparse or dense.
+    y : array
+        Observation vector of shape (nobs,)
+    w : array
+        Weights. Most commonly, they are the inverse of the expected
+        variance in the observations. Can be a single value, a vector of shape
+        (nobs,) or an array of shape (nobs, nobs) with the covariances.
+    p_sol_prior : array, optional
+        Prior information of p_sol. Of shape (npar,). If `x0` is None,
+        `p_sol_prior` is used as `x0`.
+    p_cov_prior
+        Prior information of the covariance of p_sol_prior. Of shape (npar,npar)
+    x0 : array, optional
+        Starting values in the parameter search.
+    adjust_p_cov_flag : bool
+        Not needed if weights are properly defined. Statsmodels uses it to make
+        it more robust. See `adjust_covariance()` function description.
+    calc_cov : bool
+        Will be a depreciated argument in the future because computing
+        covariance is cheap with normal equations, and computed anyways for the
+        `p_var`. Whether or not return p_cov.
+    lapack_driver : str, optional
+        Which LAPACK driver is used to solve the least-squares problem.
+        Options are ``'gelsd'``, ``'gelsy'``, ``'gelss'``. Default
+        (``'gelsd'``) is a good choice.  However, ``'gelsy'`` can be slightly
+        faster on many problems.  ``'gelss'`` was used historically.  It is
+        generally slow but uses less memory.
+    verbose : bool
+        Does not do much
+
+    Returns
+    -------
+    p_sol, p_var, p_cov : array
+
+    """
+    def adjust_covariance(p_cov_uncorrected, y, X, p_sol, w):
+        """The assumption of that the variance of the weighted residuals is 1
+        should hold. However, statsmodels corrects for the actual variance in
+        the weighted residuals. It requires the postiori parameters and the full
+        coefficient matrix and the full observations, so only possible if no
+        apriori is given.
+        """
+        nobs, npar = X.shape
+        degrees_of_freedom_err = nobs - npar
+        e = y - X @ p_sol
+        if w is None:
+            werr_sum = e.T @ e
+        elif np.ndim(w) == 0 or np.ndim(w) == 1:
+            werr_sum = e.T @ (e * w)
+        elif np.ndim(w) == 2:
+            werr_sum = e.T @ w @ e
+
+        werr_var = werr_sum / degrees_of_freedom_err
+        p_cov_corrected = np.array(p_cov_uncorrected * werr_var)
+        return p_cov_corrected
+
+    if sp.issparse(X):
+        assert np.all(np.isfinite(X.data)), 'Nan/inf in X: check ' + \
+                                            'reference temperatures?'
+    else:
+        assert np.all(np.isfinite(X)), 'Nan/inf in X: check ' + \
+                                       'reference temperatures?'
+    assert np.all(np.isfinite(w)), 'Nan/inf in weights'
+    assert np.all(np.isfinite(y)), 'Nan/inf in observations'
+
+    if p_sol_prior is not None:
+        assert np.all(np.isfinite(p_sol_prior)), 'Nan/inf in p_sol_prior'
+        assert np.all(np.isfinite(p_cov_prior)), 'Nan/inf in p_cov_prior'
+
+    if x0 is not None:
+        assert np.all(np.isfinite(x0)), 'Nan/inf in p0 initial estimate'
+    elif p_sol_prior is not None:
+        x0 = p_sol_prior
+    else:
+        pass
+
+    # Solving the normal equations instead. wX and wy are always dense.
+    if w is None:  # gracefully default to unweighted
+        # W = np.atleast_2d(1.)
+        wy = X.T @ y
+        wX = X.T @ X
+
+    elif np.ndim(w) == 0:
+        wy = X.T * w @ y
+        wX = X.T * w @ X
+
+    elif np.ndim(w) == 1:
+        wy = X.T @ (w * y)
+        if sp.issparse(X):
+            wX = (X.T @ (X.multiply(w[:, None]))).todense()
+        else:
+            wX = X.T @ (w[:, None] * X)
+
+    elif np.ndim(w) == 2:
+        wy = X.T @ w @ y
+        wX = X.T @ w @ X
+
+    if p_sol_prior is not None:
+        # Update the matrices with apriori information
+        Q_inv = scipy_inv(p_cov_prior)
+        wX += Q_inv
+        wy += Q_inv @ p_sol_prior
+
+    p_cov_uncorrected = np.linalg.inv(wX)
+
+    if x0 is not None:
+        # Start lstsq searching at x=0
+        wy -= wX @ x0
+
+    # wX and wy are overwritten during optimization and unusable after.
+    p_sol, _, _, _ = scipy_lstsq(
+        wX,
+        wy,
+        overwrite_a=not verbose,
+        overwrite_b=not verbose,
+        check_finite=False,
+        lapack_driver=lapack_driver)
+
+    if x0 is not None:
+        p_sol += x0
+
+    if p_sol_prior is None and adjust_p_cov_flag:
+        p_cov = adjust_covariance(p_cov_uncorrected, y, X, p_sol, w)
+
+    else:
+        p_cov = p_cov_uncorrected
+
+    p_var = np.diag(p_cov)
+
+    # # statsmodels comparisson
+    # mod_wls = sm.WLS(y, X, weights=w)
+    # res_wls = mod_wls.fit()
+    # p_sol = res_wls.params
+    # p_cov = res_wls.cov_params()
+    if calc_cov:
+        return p_sol, p_var, p_cov
+    else:
+        return p_sol, p_var
+
+
 def wls_stats(
         X, y, w=1., calc_cov=False, x0=None, return_werr=False, verbose=False):
     """

src/dtscalibration/datastore.py

@@ -21,6 +21,7 @@
 from .calibrate_utils import calibration_single_ended_solver
 from .calibrate_utils import match_sections
 from .calibrate_utils import wls_sparse
+from .calibrate_utils import wls_sparse2
 from .calibrate_utils import wls_stats
 from .datastore_utils import check_timestep_allclose
 from .io import apsensing_xml_version_check
@@ -1894,7 +1895,7 @@ def calibration_single_ended(
             Use `'ols'` for ordinary least squares and `'wls'` for weighted least
             squares. `'wls'` is the default, and there is currently no reason to
             use `'ols'`.
-        solver : {'sparse', 'stats'}
+        solver : {'sparse', 'sparse2', 'stats'}
             Either use the homemade weighted sparse solver or the weighted
             dense matrix solver of statsmodels. The sparse solver uses much less
             memory, is faster, and gives the same result as the statsmodels
@@ -2038,6 +2039,10 @@ def calibration_single_ended(
                     out = wls_sparse(
                         X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
 
+                elif solver == 'sparse2':
+                    out = wls_sparse2(
+                        X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
+
                 elif solver == 'stats':
                     out = wls_stats(
                         X, y, w=w, calc_cov=calc_cov, verbose=False)
@@ -2092,6 +2097,10 @@ def calibration_single_ended(
                     out = wls_sparse(
                         X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
 
+                elif solver == 'sparse2':
+                    out = wls_sparse2(
+                        X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
+
                 elif solver == 'stats':
                     out = wls_stats(
                         X, y, w=w, calc_cov=calc_cov, verbose=False)
@@ -2149,6 +2158,10 @@ def calibration_single_ended(
                     out = wls_sparse(
                         X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
 
+                elif solver == 'sparse2':
+                    out = wls_sparse2(
+                        X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
+
                 elif solver == 'stats':
                     out = wls_stats(
                         X, y, w=w, calc_cov=calc_cov, verbose=False)
@@ -2465,7 +2478,7 @@ def calibration_double_ended(
             Use `'ols'` for ordinary least squares and `'wls'` for weighted least
             squares. `'wls'` is the default, and there is currently no reason to
             use `'ols'`.
-        solver : {'sparse', 'stats'}
+        solver : {'sparse', 'sparse2', 'stats'}
             Either use the homemade weighted sparse solver or the weighted
             dense matrix solver of statsmodels. The sparse solver uses much less
             memory, is faster, and gives the same result as the statsmodels
@@ -2750,6 +2763,10 @@ def calibration_double_ended(
                     out = wls_sparse(
                         X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
 
+                elif solver == 'sparse2':
+                    out = wls_sparse2(
+                        X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
+
                 elif solver == 'stats':
                     out = wls_stats(
                         X, y, w=w, calc_cov=calc_cov, verbose=False)
@@ -2866,6 +2883,10 @@ def calibration_double_ended(
                     out = wls_sparse(
                         X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
 
+                elif solver == 'sparse2':
+                    out = wls_sparse2(
+                        X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
+
                 elif solver == 'stats':
                     out = wls_stats(
                         X, y, w=w, calc_cov=calc_cov, verbose=False)
@@ -3088,6 +3109,10 @@ def calibration_double_ended(
                     out = wls_sparse(
                         X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
 
+                elif solver == 'sparse2':
+                    out = wls_sparse2(
+                        X, y, w=w, x0=p0_est, calc_cov=calc_cov, verbose=False)
+
                 elif solver == 'stats':
                     out = wls_stats(
                         X, y, w=w, calc_cov=calc_cov, verbose=False)
@@ -5100,7 +5125,7 @@ def ufunc_per_section(
         reference section wrt the temperature of the water baths
 
         >>> tmpf_var = d.ufunc_per_section(
-        >>>     func='var'
+        >>>     func='var',
         >>>     calc_per='stretch',
         >>>     label='tmpf',
         >>>     temp_err=True)