Feature/corr matrix and inverse cov matrix as input in least squares function for correlated fits (#223)

* feat: corr_matrix kwargs as input for least squares fit

* feat/tests: inverse covariance matrix and correlation matrix kwargs as input for least squares function

* feat/tests/example: reduced new kwargs to 'inv_chol_cov_matrix' and outsourced the inversion & cholesky decomposition of the covariance matrix (function 'invert_corr_cov_cholesky(corr, covdiag)')

* tests: added tests for inv_chol_cov_matrix kwarg for the case of combined fits

* fix: renamed covdiag to inverrdiag needed for the cholesky decomposition and corrected its documentation

* examples: added an example of a correlated combined fit to the least_squares documentation

* feat/tests/fix(of typos): added function 'sort_corr()' (and a test of it) to sort correlation matrix according to a list of alphabetically sorted keys

* docs: added more elaborate documentation/example of sort_corr(), fixed typos in documentation of invert_corr_cov_cholesky()

tests/fits_test.py

@@ -152,6 +152,124 @@ def test_alternative_solvers():
     chisquare_values = np.array(chisquare_values)
     assert np.all(np.isclose(chisquare_values, chisquare_values[0]))
 
+def test_inv_cov_matrix_input_least_squares():
+    
+
+    num_samples = 400
+    N = 10
+
+    x = norm.rvs(size=(N, num_samples)) # generate random numbers
+
+    r = np.zeros((N, N))
+    for i in range(N):
+        for j in range(N):
+            r[i, j] = np.exp(-0.8 * np.fabs(i - j)) # element in correlation matrix
+
+    errl = np.sqrt([3.4, 2.5, 3.6, 2.8, 4.2, 4.7, 4.9, 5.1, 3.2, 4.2]) # set y errors
+    for i in range(N):
+        for j in range(N):
+            r[i, j] *= errl[i] * errl[j] # element in covariance matrix
+    
+    c = cholesky(r, lower=True)
+    y = np.dot(c, x)
+    x = np.arange(N)
+    x_dict = {}
+    y_dict = {}
+    for i,item in enumerate(x):
+        x_dict[str(item)] = [x[i]]
+
+    for linear in [True, False]:
+        data = []
+        for i in range(N):
+            if linear:
+                data.append(pe.Obs([[i + 1 + o for o in y[i]]], ['ens']))
+            else:
+                data.append(pe.Obs([[np.exp(-(i + 1)) + np.exp(-(i + 1)) * o for o in y[i]]], ['ens']))
+
+        [o.gamma_method() for o in data]
+
+        data_dict = {}
+        for i,item in enumerate(x):
+            data_dict[str(item)] = [data[i]]
+
+        corr = pe.covariance(data, correlation=True)
+        chol = np.linalg.cholesky(corr)
+        covdiag = np.diag(1 / np.asarray([o.dvalue for o in data]))
+        chol_inv = scipy.linalg.solve_triangular(chol, covdiag, lower=True)
+        chol_inv_keys = [""]
+        chol_inv_keys_combined_fit = [str(item) for i,item in enumerate(x)]
+
+        if linear:
+            def fitf(p, x):
+                return p[1] + p[0] * x
+            fitf_dict = {}
+            for i,item in enumerate(x):
+                fitf_dict[str(item)] = fitf
+        else:
+            def fitf(p, x):
+                return p[1] * anp.exp(-p[0] * x)
+            fitf_dict = {}
+            for i,item in enumerate(x):
+                fitf_dict[str(item)] = fitf
+
+        fitpc = pe.least_squares(x, data, fitf, correlated_fit=True)
+        fitp_inv_cov = pe.least_squares(x, data, fitf,  correlated_fit = True, inv_chol_cov_matrix = [chol_inv,chol_inv_keys])
+        fitp_inv_cov_combined_fit = pe.least_squares(x_dict, data_dict, fitf_dict,  correlated_fit = True, inv_chol_cov_matrix = [chol_inv,chol_inv_keys_combined_fit])
+        for i in range(2):
+            diff_inv_cov = fitp_inv_cov[i] - fitpc[i]
+            diff_inv_cov.gamma_method()
+            assert(diff_inv_cov.is_zero(atol=0.0))
+            diff_inv_cov_combined_fit = fitp_inv_cov_combined_fit[i] - fitpc[i]
+            diff_inv_cov_combined_fit.gamma_method()
+            assert(diff_inv_cov_combined_fit.is_zero(atol=1e-12))
+
+def test_least_squares_invalid_inv_cov_matrix_input():
+    xvals = []
+    yvals = []
+    err = 0.1
+    def func_valid(a,x):
+        return a[0] + a[1] * x
+    for x in range(1, 8, 2):
+        xvals.append(x)
+        yvals.append(pe.pseudo_Obs(x + np.random.normal(0.0, err), err, 'test1') + pe.pseudo_Obs(0, err / 100, 'test2', samples=87))
+
+    [o.gamma_method() for o in yvals]
+    
+    #dictionaries for a combined fit
+    xvals_dict = { }
+    yvals_dict = { }
+    for i,item in enumerate(np.arange(1, 8, 2)):
+        xvals_dict[str(item)] = [xvals[i]]
+        yvals_dict[str(item)] = [yvals[i]]
+    chol_inv_keys_combined_fit = ['1', '3', '5', '7']
+    chol_inv_keys_combined_fit_invalid = ['2', '7', '100', '8']
+    func_dict_valid = {"1": func_valid,"3": func_valid,"5": func_valid,"7": func_valid}
+
+    corr_valid = pe.covariance(yvals, correlation = True)
+    chol = np.linalg.cholesky(corr_valid)
+    covdiag = np.diag(1 / np.asarray([o.dvalue for o in yvals]))
+    chol_inv_valid = scipy.linalg.solve_triangular(chol, covdiag, lower=True)
+    chol_inv_keys = [""]
+    pe.least_squares(xvals, yvals,func_valid, correlated_fit = True, inv_chol_cov_matrix = [chol_inv_valid,chol_inv_keys])
+    pe.least_squares(xvals_dict, yvals_dict,func_dict_valid, correlated_fit = True, inv_chol_cov_matrix = [chol_inv_valid,chol_inv_keys_combined_fit])
+    chol_inv_invalid_shape1 = np.zeros((len(yvals),len(yvals)-1))
+    chol_inv_invalid_shape2 = np.zeros((len(yvals)+2,len(yvals)))
+
+    # for an uncombined fit
+    with pytest.raises(TypeError):
+        pe.least_squares(xvals, yvals, func_valid, correlated_fit = True, inv_chol_cov_matrix = [chol_inv_invalid_shape1,chol_inv_keys])
+    with pytest.raises(TypeError):
+        pe.least_squares(xvals, yvals, func_valid,correlated_fit = True, inv_chol_cov_matrix = [chol_inv_invalid_shape2,chol_inv_keys])
+    with pytest.raises(ValueError):
+        pe.least_squares(xvals, yvals, func_valid,correlated_fit = True, inv_chol_cov_matrix = [chol_inv_valid,chol_inv_keys_combined_fit_invalid])
+
+    #repeat for a combined fit
+    with pytest.raises(TypeError):
+        pe.least_squares(xvals_dict, yvals_dict,func_dict_valid, correlated_fit = True, inv_chol_cov_matrix = [chol_inv_invalid_shape1,chol_inv_keys_combined_fit])
+    with pytest.raises(TypeError):
+        pe.least_squares(xvals_dict, yvals_dict,func_dict_valid, correlated_fit = True, inv_chol_cov_matrix = [chol_inv_invalid_shape2,chol_inv_keys_combined_fit])
+    with pytest.raises(ValueError):
+        pe.least_squares(xvals_dict, yvals_dict,func_dict_valid, correlated_fit = True, inv_chol_cov_matrix = [chol_inv_valid,chol_inv_keys_combined_fit_invalid])
 
 def test_correlated_fit():
     num_samples = 400