From b60325dc66f3dc3cac618024dd40a6b275a941d7 Mon Sep 17 00:00:00 2001 From: Tyler Cox Date: Thu, 13 Feb 2025 17:27:25 -0800 Subject: [PATCH 1/9] add preconditioner to sparse solver --- hera_filters/dspec.py | 49 ++++++++++++++++++++++++++++ hera_filters/tests/test_dspec.py | 55 ++++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+) diff --git a/hera_filters/dspec.py b/hera_filters/dspec.py index 0249925..5c96117 100644 --- a/hera_filters/dspec.py +++ b/hera_filters/dspec.py @@ -2899,6 +2899,8 @@ def sparse_linear_fit_2D( atol: float = 1e-10, btol: float = 1e-10, iter_lim: int = None, + precondition_solver: bool = False, + eig_scaling_factor: float = 1e-1, **kwargs ) -> np.ndarray: """ @@ -2927,6 +2929,20 @@ def sparse_linear_fit_2D( flattened `data` array, x is the solution, and r is the residual. iter_lim : int, optional Maximum number of iterations for `lsqr`, default is None + precondition_solver : bool, optional, default False + If True, the solver will apply a preconditioner to the basis matrices before + solving the least-squares problem. The preconditioner is computed using the + the inverse of the regularized Gramian matrix of the basis matrices. Prior + computing the inverse, the eigenvalues of the Gramian matrix are regularized + by adding a small value proportional to the smallest eigenvalue. This helps + to stabilize the computation of the inverse. The regularization factor is + computed as the minimum eigenvalue of the Gramian matrix multiplied by the + `eig_scaling_factor` parameter. + eig_scaling_factor : float, optional, default 1e-1 + Regularization factor for the eigenvalues of the Gramian matrix. The factor + is computed as the minimum eigenvalue of the Gramian matrix multiplied by + `eig_scaling_factor`. Reasonable values are typically in the range of 1e-1 + to 1e-3. **kwargs : dict Additional keyword arguments passed to `scipy.sparse.linalg.lsqr`. @@ -2960,6 +2976,36 @@ def sparse_linear_fit_2D( axis_1_basis.shape[-1] * axis_2_basis.shape[-1], # i * j ) + if precondition_solver: + # Compute separate preconditioners for the two axes + # Start by computing separable weights for the two axes + u, s, v = np.linalg.svd(weights, full_matrices=False) + axis_1_wgts = np.abs(u[:, 0] * np.sqrt(s[0])) + axis_2_wgts = np.abs(v[0] * np.sqrt(s[0])) + + # Compute the preconditioner for the first axis + XTX_axis_1 = np.dot(axis_1_basis.T.conj() * axis_1_wgts, axis_1_basis) + eigenval, _ = np.linalg.eig(XTX_axis_1) + axis_1_lambda = np.min( + eigenval[eigenval.real > np.finfo(eigenval.dtype).eps] * eig_scaling_factor + ) + axis_1_pcond = np.linalg.pinv( + XTX_axis_1 + np.eye(XTX_axis_1.shape[0]) * axis_1_lambda + ) + + # Compute the preconditioner for the second axis + XTX_axis_2 = np.dot(axis_2_basis.T.conj() * axis_2_wgts, axis_2_basis) + eigenval, _ = np.linalg.eig(XTX_axis_2) + axis_2_lambda = np.min( + eigenval[eigenval.real > np.finfo(eigenval.dtype).eps] * eig_scaling_factor + ) + axis_2_pcond = np.linalg.pinv( + XTX_axis_2 + np.eye(XTX_axis_2.shape[0]) * axis_2_lambda + ) + + axis_1_basis = np.dot(axis_1_basis, axis_1_pcond) + axis_2_basis = np.dot(axis_2_basis, axis_2_pcond) + # Define the implicit LinearOperator representing the Kronecker product linear_operator = sparse.linalg.LinearOperator( full_operator_shape, @@ -2985,6 +3031,9 @@ def sparse_linear_fit_2D( # Reshape output x = x.reshape(axis_1_basis.shape[-1], axis_2_basis.shape[-1]) + if precondition_solver: + x = np.dot(axis_1_pcond, x).dot(axis_2_pcond) + return x, meta def separable_linear_fit_2D( diff --git a/hera_filters/tests/test_dspec.py b/hera_filters/tests/test_dspec.py index 931836c..4efa554 100644 --- a/hera_filters/tests/test_dspec.py +++ b/hera_filters/tests/test_dspec.py @@ -1535,3 +1535,58 @@ def test_sparse_linear_fit_2d_non_binary_wgts(): # Check that the fit closely matches to the separable fit np.testing.assert_allclose(sol, sol_sparse, atol=1e-9, rtol=1e-6) + +def test_precondition_sparse_solver(): + # test that separable linear fit works as expected. + ntimes, nfreqs = 100, 50 + + # Generate some data/flags + # By construction, the data is separable in the time and frequency directions + # and the flags are also separable. The fit should be able to recover the + # true data in the unflagged region. + rng = np.random.default_rng(42) + freq_basis, _ = dspec.dpss_operator(np.linspace(100e6, 200e6, nfreqs), [0], [20e-9], eigenval_cutoff=[1e-12]) + time_basis, _ = dspec.dpss_operator(np.linspace(0, ntimes * 10, ntimes), [0], [1e-3], eigenval_cutoff=[1e-12]) + time_flags = rng.choice([True, False], p=[0.1, 0.9], size=(ntimes, 1)) + freq_flags = rng.choice([True, False], p=[0.1, 0.9], size=(1, nfreqs)) + x_true = rng.normal(0, 1, size=(time_basis.shape[-1], freq_basis.shape[-1])) + data = np.dot(time_basis, x_true).dot(freq_basis.T) + freqs = np.linspace(100e6, 200e6, nfreqs) + + # Generate separable, non-binary weights + axis_1_weights = (~time_flags[:, 0]).astype(float) * rng.integers(1, 10, size=(ntimes,)) + axis_2_weights = (~freq_flags[0]).astype(float) + wgts = np.outer(axis_1_weights, axis_2_weights) + + # Add frequency dependence to the weights to make the problem more ill-conditioned + wgts *= (freqs / 150e6) ** -3.5 + + # Fit the data + sol = dspec.separable_linear_fit_2D( + data=data, + axis_1_weights=(~time_flags[:, 0]).astype(float), + axis_2_weights=(~freq_flags[0]).astype(float), + axis_1_basis=time_basis, + axis_2_basis=freq_basis, + ) + + sol_sparse, meta = dspec.sparse_linear_fit_2D( + data=data, + weights=wgts, + axis_1_basis=time_basis, + axis_2_basis=freq_basis, + precondition_solver=False + ) + + sol_sparse_precond, meta_precond = dspec.sparse_linear_fit_2D( + data=data, + weights=wgts, + axis_1_basis=time_basis, + axis_2_basis=freq_basis, + precondition_solver=True + ) + + # Check that the fit closely matches to the separable fit + np.testing.assert_allclose(sol, sol_sparse, atol=1e-8, rtol=1e-6) + np.testing.assert_allclose(sol, sol_sparse_precond, atol=1e-8, rtol=1e-6) + np.testing.assert_array_less(meta_precond['iter_num'], meta['iter_num']) \ No newline at end of file From 586a6f50ca25be0725f8cb74f9e4adc0c70e2aaa Mon Sep 17 00:00:00 2001 From: Tyler Cox Date: Thu, 13 Feb 2025 17:31:35 -0800 Subject: [PATCH 2/9] add more description to preconditioner parameter --- hera_filters/dspec.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/hera_filters/dspec.py b/hera_filters/dspec.py index 5c96117..4758c04 100644 --- a/hera_filters/dspec.py +++ b/hera_filters/dspec.py @@ -2931,12 +2931,14 @@ def sparse_linear_fit_2D( Maximum number of iterations for `lsqr`, default is None precondition_solver : bool, optional, default False If True, the solver will apply a preconditioner to the basis matrices before - solving the least-squares problem. The preconditioner is computed using the - the inverse of the regularized Gramian matrix of the basis matrices. Prior - computing the inverse, the eigenvalues of the Gramian matrix are regularized - by adding a small value proportional to the smallest eigenvalue. This helps - to stabilize the computation of the inverse. The regularization factor is - computed as the minimum eigenvalue of the Gramian matrix multiplied by the + solving the least-squares problem. This option is useful when the input weights + are frequency or time dependent and are either very large or very small, or when + the basis matrices are ill-conditioned due to large stretches of zeros. + The preconditioner is computed using the the inverse of the regularized Gramian + matrix of the basis matrices. Prior computing the inverse, the eigenvalues of the + Gramian matrix are regularized by adding a small value proportional to the smallest + eigenvalue. This helps to stabilize the computation of the inverse. The regularization + factor is computed as the minimum eigenvalue of the Gramian matrix multiplied by the `eig_scaling_factor` parameter. eig_scaling_factor : float, optional, default 1e-1 Regularization factor for the eigenvalues of the Gramian matrix. The factor From dba338d6cdb9faca85d44d329de418235bbcac4c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Feb 2025 01:31:52 +0000 Subject: [PATCH 3/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- hera_filters/dspec.py | 12 ++++++------ hera_filters/tests/test_dspec.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/hera_filters/dspec.py b/hera_filters/dspec.py index 4758c04..8b26e64 100644 --- a/hera_filters/dspec.py +++ b/hera_filters/dspec.py @@ -2933,11 +2933,11 @@ def sparse_linear_fit_2D( If True, the solver will apply a preconditioner to the basis matrices before solving the least-squares problem. This option is useful when the input weights are frequency or time dependent and are either very large or very small, or when - the basis matrices are ill-conditioned due to large stretches of zeros. - The preconditioner is computed using the the inverse of the regularized Gramian - matrix of the basis matrices. Prior computing the inverse, the eigenvalues of the - Gramian matrix are regularized by adding a small value proportional to the smallest - eigenvalue. This helps to stabilize the computation of the inverse. The regularization + the basis matrices are ill-conditioned due to large stretches of zeros. + The preconditioner is computed using the the inverse of the regularized Gramian + matrix of the basis matrices. Prior computing the inverse, the eigenvalues of the + Gramian matrix are regularized by adding a small value proportional to the smallest + eigenvalue. This helps to stabilize the computation of the inverse. The regularization factor is computed as the minimum eigenvalue of the Gramian matrix multiplied by the `eig_scaling_factor` parameter. eig_scaling_factor : float, optional, default 1e-1 @@ -2994,7 +2994,7 @@ def sparse_linear_fit_2D( axis_1_pcond = np.linalg.pinv( XTX_axis_1 + np.eye(XTX_axis_1.shape[0]) * axis_1_lambda ) - + # Compute the preconditioner for the second axis XTX_axis_2 = np.dot(axis_2_basis.T.conj() * axis_2_wgts, axis_2_basis) eigenval, _ = np.linalg.eig(XTX_axis_2) diff --git a/hera_filters/tests/test_dspec.py b/hera_filters/tests/test_dspec.py index 4efa554..7c3c0a7 100644 --- a/hera_filters/tests/test_dspec.py +++ b/hera_filters/tests/test_dspec.py @@ -1557,7 +1557,7 @@ def test_precondition_sparse_solver(): axis_1_weights = (~time_flags[:, 0]).astype(float) * rng.integers(1, 10, size=(ntimes,)) axis_2_weights = (~freq_flags[0]).astype(float) wgts = np.outer(axis_1_weights, axis_2_weights) - + # Add frequency dependence to the weights to make the problem more ill-conditioned wgts *= (freqs / 150e6) ** -3.5 @@ -1589,4 +1589,4 @@ def test_precondition_sparse_solver(): # Check that the fit closely matches to the separable fit np.testing.assert_allclose(sol, sol_sparse, atol=1e-8, rtol=1e-6) np.testing.assert_allclose(sol, sol_sparse_precond, atol=1e-8, rtol=1e-6) - np.testing.assert_array_less(meta_precond['iter_num'], meta['iter_num']) \ No newline at end of file + np.testing.assert_array_less(meta_precond['iter_num'], meta['iter_num']) From 43558a782c1d22da1d284d79a9d4bd11ea28353c Mon Sep 17 00:00:00 2001 From: Tyler Cox Date: Thu, 13 Feb 2025 17:54:27 -0800 Subject: [PATCH 4/9] use faster sparse svd and update docstring --- hera_filters/dspec.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/hera_filters/dspec.py b/hera_filters/dspec.py index 8b26e64..eb78733 100644 --- a/hera_filters/dspec.py +++ b/hera_filters/dspec.py @@ -2935,8 +2935,8 @@ def sparse_linear_fit_2D( are frequency or time dependent and are either very large or very small, or when the basis matrices are ill-conditioned due to large stretches of zeros. The preconditioner is computed using the the inverse of the regularized Gramian - matrix of the basis matrices. Prior computing the inverse, the eigenvalues of the - Gramian matrix are regularized by adding a small value proportional to the smallest + matrix (X^T W X) of the basis matrices. Prior to computing the inverse, the eigenvalues + of the Gramian matrix are regularized by adding a small value proportional to the smallest eigenvalue. This helps to stabilize the computation of the inverse. The regularization factor is computed as the minimum eigenvalue of the Gramian matrix multiplied by the `eig_scaling_factor` parameter. @@ -2981,7 +2981,7 @@ def sparse_linear_fit_2D( if precondition_solver: # Compute separate preconditioners for the two axes # Start by computing separable weights for the two axes - u, s, v = np.linalg.svd(weights, full_matrices=False) + u, s, v = sparse.linalg.svd(weights, k=1) axis_1_wgts = np.abs(u[:, 0] * np.sqrt(s[0])) axis_2_wgts = np.abs(v[0] * np.sqrt(s[0])) @@ -2989,7 +2989,7 @@ def sparse_linear_fit_2D( XTX_axis_1 = np.dot(axis_1_basis.T.conj() * axis_1_wgts, axis_1_basis) eigenval, _ = np.linalg.eig(XTX_axis_1) axis_1_lambda = np.min( - eigenval[eigenval.real > np.finfo(eigenval.dtype).eps] * eig_scaling_factor + eigenval[eigenval.real > np.finfo(eigenval.dtype).eps].real * eig_scaling_factor ) axis_1_pcond = np.linalg.pinv( XTX_axis_1 + np.eye(XTX_axis_1.shape[0]) * axis_1_lambda @@ -2999,7 +2999,7 @@ def sparse_linear_fit_2D( XTX_axis_2 = np.dot(axis_2_basis.T.conj() * axis_2_wgts, axis_2_basis) eigenval, _ = np.linalg.eig(XTX_axis_2) axis_2_lambda = np.min( - eigenval[eigenval.real > np.finfo(eigenval.dtype).eps] * eig_scaling_factor + eigenval[eigenval.real > np.finfo(eigenval.dtype).eps].real * eig_scaling_factor ) axis_2_pcond = np.linalg.pinv( XTX_axis_2 + np.eye(XTX_axis_2.shape[0]) * axis_2_lambda From 58b6b8d228f88e20073980311d9b73c341a38247 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Feb 2025 01:54:38 +0000 Subject: [PATCH 5/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- hera_filters/dspec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hera_filters/dspec.py b/hera_filters/dspec.py index eb78733..00326aa 100644 --- a/hera_filters/dspec.py +++ b/hera_filters/dspec.py @@ -2935,7 +2935,7 @@ def sparse_linear_fit_2D( are frequency or time dependent and are either very large or very small, or when the basis matrices are ill-conditioned due to large stretches of zeros. The preconditioner is computed using the the inverse of the regularized Gramian - matrix (X^T W X) of the basis matrices. Prior to computing the inverse, the eigenvalues + matrix (X^T W X) of the basis matrices. Prior to computing the inverse, the eigenvalues of the Gramian matrix are regularized by adding a small value proportional to the smallest eigenvalue. This helps to stabilize the computation of the inverse. The regularization factor is computed as the minimum eigenvalue of the Gramian matrix multiplied by the From e1c15e48c4ec5971fe8515bb9b0629590c506f2e Mon Sep 17 00:00:00 2001 From: Tyler Cox Date: Thu, 13 Feb 2025 18:00:29 -0800 Subject: [PATCH 6/9] "svd" -> "svds" --- hera_filters/dspec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/hera_filters/dspec.py b/hera_filters/dspec.py index 00326aa..75c31f7 100644 --- a/hera_filters/dspec.py +++ b/hera_filters/dspec.py @@ -2981,7 +2981,7 @@ def sparse_linear_fit_2D( if precondition_solver: # Compute separate preconditioners for the two axes # Start by computing separable weights for the two axes - u, s, v = sparse.linalg.svd(weights, k=1) + u, s, v = sparse.linalg.svds(weights, k=1) axis_1_wgts = np.abs(u[:, 0] * np.sqrt(s[0])) axis_2_wgts = np.abs(v[0] * np.sqrt(s[0])) From c4b784c1387d44956d489a84310e4dbfbbc18b47 Mon Sep 17 00:00:00 2001 From: Tyler Cox Date: Thu, 13 Feb 2025 18:08:41 -0800 Subject: [PATCH 7/9] use scipy's eigenvalue solver to get smallest eigenvalue --- hera_filters/dspec.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/hera_filters/dspec.py b/hera_filters/dspec.py index 75c31f7..40adb6c 100644 --- a/hera_filters/dspec.py +++ b/hera_filters/dspec.py @@ -2987,20 +2987,16 @@ def sparse_linear_fit_2D( # Compute the preconditioner for the first axis XTX_axis_1 = np.dot(axis_1_basis.T.conj() * axis_1_wgts, axis_1_basis) - eigenval, _ = np.linalg.eig(XTX_axis_1) - axis_1_lambda = np.min( - eigenval[eigenval.real > np.finfo(eigenval.dtype).eps].real * eig_scaling_factor - ) + eigenval = sparse.linalg.eigs(XTX_axis_1, k=1, which='SR', return_eigenvectors=False) + axis_1_lambda = eigenval.real * eig_scaling_factor axis_1_pcond = np.linalg.pinv( XTX_axis_1 + np.eye(XTX_axis_1.shape[0]) * axis_1_lambda ) # Compute the preconditioner for the second axis XTX_axis_2 = np.dot(axis_2_basis.T.conj() * axis_2_wgts, axis_2_basis) - eigenval, _ = np.linalg.eig(XTX_axis_2) - axis_2_lambda = np.min( - eigenval[eigenval.real > np.finfo(eigenval.dtype).eps].real * eig_scaling_factor - ) + eigenval = sparse.linalg.eigs(XTX_axis_2, k=1, which='SR', return_eigenvectors=False) + axis_2_lambda = eigenval.real * eig_scaling_factor axis_2_pcond = np.linalg.pinv( XTX_axis_2 + np.eye(XTX_axis_2.shape[0]) * axis_2_lambda ) From c4cbe4ece2842801a9ec89abba17827ced1d3fb0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Feb 2025 02:08:51 +0000 Subject: [PATCH 8/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- hera_filters/dspec.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/hera_filters/dspec.py b/hera_filters/dspec.py index 40adb6c..9ba5378 100644 --- a/hera_filters/dspec.py +++ b/hera_filters/dspec.py @@ -2987,7 +2987,7 @@ def sparse_linear_fit_2D( # Compute the preconditioner for the first axis XTX_axis_1 = np.dot(axis_1_basis.T.conj() * axis_1_wgts, axis_1_basis) - eigenval = sparse.linalg.eigs(XTX_axis_1, k=1, which='SR', return_eigenvectors=False) + eigenval = sparse.linalg.eigs(XTX_axis_1, k=1, which='SR', return_eigenvectors=False) axis_1_lambda = eigenval.real * eig_scaling_factor axis_1_pcond = np.linalg.pinv( XTX_axis_1 + np.eye(XTX_axis_1.shape[0]) * axis_1_lambda @@ -2995,7 +2995,7 @@ def sparse_linear_fit_2D( # Compute the preconditioner for the second axis XTX_axis_2 = np.dot(axis_2_basis.T.conj() * axis_2_wgts, axis_2_basis) - eigenval = sparse.linalg.eigs(XTX_axis_2, k=1, which='SR', return_eigenvectors=False) + eigenval = sparse.linalg.eigs(XTX_axis_2, k=1, which='SR', return_eigenvectors=False) axis_2_lambda = eigenval.real * eig_scaling_factor axis_2_pcond = np.linalg.pinv( XTX_axis_2 + np.eye(XTX_axis_2.shape[0]) * axis_2_lambda From d7293e4a7b4a6c8d2210df18933d6dd58c0e9f8b Mon Sep 17 00:00:00 2001 From: Tyler Cox Date: Thu, 13 Feb 2025 18:13:48 -0800 Subject: [PATCH 9/9] use smallest magnitude eigenvalue --- hera_filters/dspec.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/hera_filters/dspec.py b/hera_filters/dspec.py index 9ba5378..087727e 100644 --- a/hera_filters/dspec.py +++ b/hera_filters/dspec.py @@ -2987,16 +2987,16 @@ def sparse_linear_fit_2D( # Compute the preconditioner for the first axis XTX_axis_1 = np.dot(axis_1_basis.T.conj() * axis_1_wgts, axis_1_basis) - eigenval = sparse.linalg.eigs(XTX_axis_1, k=1, which='SR', return_eigenvectors=False) - axis_1_lambda = eigenval.real * eig_scaling_factor + eigenval = sparse.linalg.eigs(XTX_axis_1, k=1, which='SM', return_eigenvectors=False) + axis_1_lambda = np.abs(eigenval) * eig_scaling_factor axis_1_pcond = np.linalg.pinv( XTX_axis_1 + np.eye(XTX_axis_1.shape[0]) * axis_1_lambda ) # Compute the preconditioner for the second axis XTX_axis_2 = np.dot(axis_2_basis.T.conj() * axis_2_wgts, axis_2_basis) - eigenval = sparse.linalg.eigs(XTX_axis_2, k=1, which='SR', return_eigenvectors=False) - axis_2_lambda = eigenval.real * eig_scaling_factor + eigenval = sparse.linalg.eigs(XTX_axis_2, k=1, which='SM', return_eigenvectors=False) + axis_2_lambda = np.abs(eigenval) * eig_scaling_factor axis_2_pcond = np.linalg.pinv( XTX_axis_2 + np.eye(XTX_axis_2.shape[0]) * axis_2_lambda )