From 817658ea98e638c98ec710a19a1238ae70f0213f Mon Sep 17 00:00:00 2001 From: Jeffrey Larson Date: Fri, 5 Apr 2024 11:27:54 -0500 Subject: [PATCH 1/5] simpler api for now --- qfi_opt/examples/calculate_qfi.py | 67 +++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/qfi_opt/examples/calculate_qfi.py b/qfi_opt/examples/calculate_qfi.py index b431b4b..705c2e7 100644 --- a/qfi_opt/examples/calculate_qfi.py +++ b/qfi_opt/examples/calculate_qfi.py @@ -15,6 +15,73 @@ def compute_eigendecomposition(rho: np.ndarray): eigvecs = eigvecs.T # make the k-th eigenvector eigvecs[k, :] = eigvecs[k] return eigvals, eigvecs +def compute_QFI_simpler_api( + eigvals: np.ndarray, eigvecs: np.ndarray, A: np.ndarray, params: + np.ndarray, grad, get_jacobian, obj_params, tol: float = 1e-8, etol_scale: + float = 10 +) -> float: + # Note: The eigenvectors must be rows of eigvecs + num_vals = len(eigvals) + num_params = len(params) + + G = obj_params["G"] + + # There should never be negative eigenvalues, so their magnitude gives an + # empirical estimate of the numerical accuracy of the eigendecomposition. + # We discard any QFI terms denominators within an order of magnitude of + # this value. + tol = max(tol, -etol_scale * np.min(eigvals)) + + # Compute QFI and grad + running_sum = 0 + + if grad.size > 0: + + dA = get_jacobian(params, obj_params["N"], dissipation_rates=obj_params["dissipation"]) + dA = np.transpose(dA, (2, 0, 1)) + + grad[:] = np.zeros(num_params) + lambda_grads = np.zeros((num_params, num_vals)) + psi_grads = np.zeros((num_params, num_vals, num_vals), dtype="cdouble") + + for k in range(num_params): + # compute gradients of each eigenvalue + lambda_grad_k, psi_grad_k = get_matrix_grads_sylvester(dA[k], eigvals, eigvecs, tol) + lambda_grads[k] = lambda_grad_k + psi_grads[k] = psi_grad_k + + for i in range(num_vals): + for j in range(i + 1, num_vals): + denom = eigvals[i] + eigvals[j] + diff = eigvals[i] - eigvals[j] + if not np.isclose(denom, 0, atol=tol, rtol=tol) and not np.isclose(diff, 0, atol=tol, rtol=tol): + numer = diff**2 + term = eigvecs[i].conj() @ G @ eigvecs[j] + quotient = numer / denom + squared_modulus = np.absolute(term) ** 2 + running_sum += quotient * squared_modulus + if grad.size > 0: + for k in range(num_params): + # fill in gradient + grad[k] += kth_partial_derivative( + quotient, + squared_modulus, + eigvals[i], + eigvals[j], + lambda_grads[k, i], + lambda_grads[k, j], + eigvecs[i], + eigvecs[j], + psi_grads[k, i], + psi_grads[k, j], + G, + ) + + if grad.size > 0: + return 4 * running_sum, 4 * grad + else: + return 4 * running_sum, [] + def compute_QFI(eigvals: np.ndarray, eigvecs: np.ndarray, G: np.ndarray, A: np.ndarray= np.empty(0), dA: np.ndarray= np.empty(0), d2A: np.ndarray = np.empty(0), From e3f871b3fd6f6c7751cf2e20d204857e1e1596d7 Mon Sep 17 00:00:00 2001 From: Jeffrey Larson Date: Fri, 5 Apr 2024 11:29:26 -0500 Subject: [PATCH 2/5] Adding import --- qfi_opt/examples/calculate_qfi.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/qfi_opt/examples/calculate_qfi.py b/qfi_opt/examples/calculate_qfi.py index 705c2e7..accd553 100644 --- a/qfi_opt/examples/calculate_qfi.py +++ b/qfi_opt/examples/calculate_qfi.py @@ -2,6 +2,8 @@ import numpy as np from qfi_opt import spin_models +from scipy.linalg import solve_sylvester + def variance(rho: np.ndarray, G: np.ndarray) -> float: From eb27a8eec59c098868d14a8e80661f846a550598 Mon Sep 17 00:00:00 2001 From: Jeffrey Larson Date: Fri, 5 Apr 2024 11:31:55 -0500 Subject: [PATCH 3/5] cleanup --- qfi_opt/examples/calculate_qfi.py | 62 ++++++++++++++++++++++++++----- qfi_opt/examples/matt_example.py | 56 ++++++++++++++++++++++++++++ 2 files changed, 109 insertions(+), 9 deletions(-) create mode 100644 qfi_opt/examples/matt_example.py diff --git a/qfi_opt/examples/calculate_qfi.py b/qfi_opt/examples/calculate_qfi.py index accd553..aee6cfa 100644 --- a/qfi_opt/examples/calculate_qfi.py +++ b/qfi_opt/examples/calculate_qfi.py @@ -1,9 +1,8 @@ #!/usr/bin/env python3 import numpy as np - -from qfi_opt import spin_models from scipy.linalg import solve_sylvester +from qfi_opt import spin_models def variance(rho: np.ndarray, G: np.ndarray) -> float: @@ -17,10 +16,9 @@ def compute_eigendecomposition(rho: np.ndarray): eigvecs = eigvecs.T # make the k-th eigenvector eigvecs[k, :] = eigvecs[k] return eigvals, eigvecs + def compute_QFI_simpler_api( - eigvals: np.ndarray, eigvecs: np.ndarray, A: np.ndarray, params: - np.ndarray, grad, get_jacobian, obj_params, tol: float = 1e-8, etol_scale: - float = 10 + eigvals: np.ndarray, eigvecs: np.ndarray, A: np.ndarray, params: np.ndarray, grad, get_jacobian, obj_params, tol: float = 1e-8, etol_scale: float = 10 ) -> float: # Note: The eigenvectors must be rows of eigvecs num_vals = len(eigvals) @@ -85,10 +83,17 @@ def compute_QFI_simpler_api( return 4 * running_sum, [] -def compute_QFI(eigvals: np.ndarray, eigvecs: np.ndarray, G: np.ndarray, A: - np.ndarray= np.empty(0), dA: np.ndarray= np.empty(0), d2A: np.ndarray = np.empty(0), - grad: np.ndarray = np.empty(0), tol: float = 1e-8, etol_scale: float = - 10) -> float: +def compute_QFI( + eigvals: np.ndarray, + eigvecs: np.ndarray, + G: np.ndarray, + A: np.ndarray = np.empty(0), + dA: np.ndarray = np.empty(0), + d2A: np.ndarray = np.empty(0), + grad: np.ndarray = np.empty(0), + tol: float = 1e-8, + etol_scale: float = 10, +) -> float: # Note: The eigenvectors must be rows of eigvecs num_vals = len(eigvals) num_params = dA.shape[0] @@ -156,6 +161,45 @@ def get_matrix_grads_lazy(A, dA, eigvals, eigvecs): return lambda_grads, psi_grads +def get_matrix_grads_sylvester(dA, eigvals, eigvecs, tol): + + dim = eigvecs.shape[0] + lambda_grads = np.zeros(dim, dtype="cdouble") + psi_grads = np.zeros((dim, dim), dtype="cdouble") + + # force Hermitianness: + dA = (dA + dA.conj().T) / 2.0 + + # group the sorted eigvals by tolerance, intended to help stability of eigenvector derivatives: + current_ind = 0 + for ind1 in range(dim): + if current_ind == ind1: + for ind2 in range(ind1 + 1, dim): + if not np.isclose(eigvals[ind2], eigvals[ind1], atol=tol, rtol=tol): + break # the for loop over ind2 + # we just broke the for loop, so: + current_ind = ind2 + + # do a sylvester solve: + group_set = np.arange(ind1, ind2) + # special case when ind2=dim (must be something smarter) + if group_set.size == 0: + group_set = [ind2] + not_in_group_set = np.setdiff1d(np.arange(dim), group_set) + A_group = np.diag(eigvals[group_set]) + A_not_in_group = np.diag(eigvals[not_in_group_set]) + rotation = eigvecs[group_set].conj() @ dA @ eigvecs[not_in_group_set].T + sol = solve_sylvester(A_group, -1.0 * A_not_in_group, rotation) + psi_grads[group_set] = sol @ eigvecs[not_in_group_set].conj() + + # average eigenvalue: + multiplicity = len(group_set) + dLambda = (1.0 / multiplicity) * np.trace(eigvecs[group_set].conj() @ dA @ eigvecs[group_set].T) + lambda_grads[group_set] = np.ones(multiplicity) * dLambda + + return np.real(lambda_grads), psi_grads.conj() + + def get_matrix_grads(A, dA, d2A, eigvals, eigvecs, tol): num_vals = len(eigvals) num_params = dA.shape[0] diff --git a/qfi_opt/examples/matt_example.py b/qfi_opt/examples/matt_example.py new file mode 100644 index 0000000..f33b053 --- /dev/null +++ b/qfi_opt/examples/matt_example.py @@ -0,0 +1,56 @@ +import numpy as np + +import qfi_opt +from qfi_opt import spin_models +from qfi_opt.examples.calculate_qfi import compute_eigendecomposition, compute_QFI_simpler_api + +np.set_printoptions(precision=4, linewidth=200) + + +def sim_wrapper_diffrax(x, qfi_grad, obj, obj_params, get_jacobian): + rho = obj(x, obj_params["N"], dissipation_rates=obj_params["dissipation"]) + + # force Hermitianness: + rho = (rho + rho.conj().T) / 2.0 + # Compute eigendecomposition + vals, vecs = compute_eigendecomposition(rho) + + qfi, new_grad = compute_QFI_simpler_api(vals, vecs, rho, x, qfi_grad, get_jacobian, obj_params) + print(x, qfi, new_grad, flush=True) + + try: + if qfi_grad.size > 0: + qfi_grad[:] = -1.0 * new_grad + except: + qfi_grad[:] = [] + + return -1.0 * qfi # , -qfi_grad # negative because we are maximizing + + +if __name__ == "__main__": + + N = 4 + G = spin_models.collective_op(spin_models.PAULI_Z, N) / (2 * N) + + obj_params = {} + obj_params["N"] = N + obj_params["dissipation"] = 1.0 + obj_params["G"] = G + + seed = 88 + np.random.seed(seed) + + num_params = 5 + + lb = np.zeros(num_params) + ub = np.ones(num_params) + + x0 = np.random.uniform(lb, ub, num_params) + model = "simulate_TAT" + + obj = getattr(spin_models, model) + get_jacobian = spin_models.get_jacobian_func(obj) + + grad = np.zeros(num_params) + out = sim_wrapper_diffrax(x0, grad, obj, obj_params, get_jacobian) + print(out) From 3457f011dc71b277c03ac4a391c29049a3f923e3 Mon Sep 17 00:00:00 2001 From: Jeffrey Larson Date: Fri, 5 Apr 2024 11:49:31 -0500 Subject: [PATCH 4/5] now with and without gradients --- qfi_opt/examples/matt_example.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/qfi_opt/examples/matt_example.py b/qfi_opt/examples/matt_example.py index f33b053..aff20c7 100644 --- a/qfi_opt/examples/matt_example.py +++ b/qfi_opt/examples/matt_example.py @@ -49,8 +49,13 @@ def sim_wrapper_diffrax(x, qfi_grad, obj, obj_params, get_jacobian): model = "simulate_TAT" obj = getattr(spin_models, model) - get_jacobian = spin_models.get_jacobian_func(obj) + + grad = np.empty(0) + get_jacobian = np.empty(0) + out = sim_wrapper_diffrax(x0, grad, obj, obj_params, get_jacobian) + print(out) grad = np.zeros(num_params) + get_jacobian = spin_models.get_jacobian_func(obj) out = sim_wrapper_diffrax(x0, grad, obj, obj_params, get_jacobian) print(out) From e4c959fa49ae3b14187e71ebde3f53e1aea8035f Mon Sep 17 00:00:00 2001 From: Jeffrey Larson Date: Fri, 5 Apr 2024 12:00:58 -0500 Subject: [PATCH 5/5] simple --- qfi_opt/examples/matt_example.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/qfi_opt/examples/matt_example.py b/qfi_opt/examples/matt_example.py index aff20c7..e709211 100644 --- a/qfi_opt/examples/matt_example.py +++ b/qfi_opt/examples/matt_example.py @@ -50,11 +50,6 @@ def sim_wrapper_diffrax(x, qfi_grad, obj, obj_params, get_jacobian): obj = getattr(spin_models, model) - grad = np.empty(0) - get_jacobian = np.empty(0) - out = sim_wrapper_diffrax(x0, grad, obj, obj_params, get_jacobian) - print(out) - grad = np.zeros(num_params) get_jacobian = spin_models.get_jacobian_func(obj) out = sim_wrapper_diffrax(x0, grad, obj, obj_params, get_jacobian)