factored out the calculation per q into its own function

Author: alexhroom

pyspinw/calculations/spinwave.py

@@ -54,96 +54,102 @@ def spinwave_calculation(rotations: np.ndarray,
     z_conj = z.conj()
     eta = rotations[:,:,2] # n-by-3 real
 
-    # Create the A, B, and C matrices
-
     # Create the matrix sqrt(S_i S_j)/2
     root_mags = np.sqrt(0.5*magnitudes) # S_i / sqrt(2)
     spin_coefficients = root_mags.reshape(-1, 1) * root_mags.reshape(1, -1)
 
     energies = []
     methods = []
-    for q in q_vectors:
-
-        A = np.zeros((n_sites, n_sites), dtype=complex)
-        B = np.zeros((n_sites, n_sites), dtype=complex)
-        C = np.zeros((n_sites, n_sites), dtype=complex)
 
-        # Add the terms up to the total spin coefficient
-        for coupling in couplings:
+    # calculate the C matrix for h(q), which is q-independent
+    C = np.zeros((n_sites, n_sites), dtype=complex)
+    for coupling in couplings:
+        j = coupling.index2
+        for l in range(n_sites):
+            C[j, j] += spin_coefficients[l,l] * eta[j, :].T @ coupling.matrix @ eta[l, :]
 
-            phase_factor = np.exp((2j*np.pi)*np.dot(q, coupling.inter_site_vector))
+    C *= 2
 
-            i = coupling.index1
-            j = coupling.index2
+    for q in q_vectors:
+        eigenvalues, method = spinwave_calculation_single_q(q, C, n_sites, z, z_conj, spin_coefficients, couplings)
+        energies.append(eigenvalues)  # These are currently the square energies
+        methods.append(method)
 
-            A[i, j] += z[i, :] @ coupling.matrix @ z_conj[j, :] * phase_factor
-            B[i, j] += z[i, :] @ coupling.matrix @ z[j, :] * phase_factor
+    return SpinwaveResult(
+                q_vectors=q_vectors,
+                raw_energies=energies,
+                method=methods)
 
 
-            for l in range(n_sites):
-                C[j, j] += spin_coefficients[l,l] * eta[j, :].T @ coupling.matrix @ eta[l, :]
+def spinwave_calculation_single_q(q: float,
+                                  C: np.ndarray,
+                                  n_sites: int,
+                                  z: np.ndarray,
+                                  z_conj: np.ndarray,
+                                  spin_coefficients: np.ndarray,
+                                  couplings: list[Coupling]):
+    """Calculate the energies for the 'Hamiltonian' h(q) for a single q-value."""
+    A = np.zeros((n_sites, n_sites), dtype=complex)
+    B = np.zeros((n_sites, n_sites), dtype=complex)
 
-        A *= spin_coefficients
-        B *= spin_coefficients
-        C *= 2
-        #
-        # print("A", A)
-        # print("B", B)
-        # print("C", C)
+    # Add the terms up to the total spin coefficient
+    for coupling in couplings:
 
+        phase_factor = np.exp((2j*np.pi)*np.dot(q, coupling.inter_site_vector))
 
-        # hamiltonian_matrix = np.block([[A - C, B], [B.conj().T, A.conj().T - C]])
-        hamiltonian_matrix = np.block([[A - C, B], [B.conj().T, A.conj().T - C]])
+        i = coupling.index1
+        j = coupling.index2
 
-        #
-        # We need to enforce the bosonic commutation properties, we do this
-        # by finding the 'square root' of the matrix (i.e. finding K such that KK^dagger = H)
-        # and then negating the second half.
-        #
-        # In matrix form we do
-        #
-        #     M = K^dagger C K
-        #
-        # where C is a diagonal matrix of length 2n, with the first n entries being 1, and the
-        # remaining entries being -1. Note the adjoint is on the other side to the definition in
-        # of the decomposition
-        #
-        # We can also do this via an LDL decomposition, but the method is very slightly different
+        A[i, j] += z[i, :] @ coupling.matrix @ z_conj[j, :] * phase_factor
+        B[i, j] += z[i, :] @ coupling.matrix @ z[j, :] * phase_factor
 
 
-        try:
-            sqrt_hamiltonian = np.linalg.cholesky(hamiltonian_matrix)
+    A *= spin_coefficients
+    B *= spin_coefficients
+    #
+    # print("A", A)
+    # print("B", B)
+    # print("C", C)
 
-            sqrt_hamiltonian_with_commutation = sqrt_hamiltonian.copy()
-            sqrt_hamiltonian_with_commutation[:, n_sites:] *= -1
 
-            to_diagonalise = np.conj(sqrt_hamiltonian_with_commutation).T @ sqrt_hamiltonian
+    # hamiltonian_matrix = np.block([[A - C, B], [B.conj().T, A.conj().T - C]])
+    hamiltonian_matrix = np.block([[A - C, B], [B.conj().T, A.conj().T - C]])
 
-            methods.append(CalculationMethod.CHOLESKY)
+    #
+    # We need to enforce the bosonic commutation properties, we do this
+    # by finding the 'square root' of the matrix (i.e. finding K such that KK^dagger = H)
+    # and then negating the second half.
+    #
+    # In matrix form we do
+    #
+    #     M = K^dagger C K
+    #
+    # where C is a diagonal matrix of length 2n, with the first n entries being 1, and the
+    # remaining entries being -1. Note the adjoint is on the other side to the definition in
+    # of the decomposition
+    #
+    # We can also do this via an LDL decomposition, but the method is very slightly different
 
-        except np.linalg.LinAlgError: # Catch postive definiteness errors
 
+    try:
+        sqrt_hamiltonian = np.linalg.cholesky(hamiltonian_matrix)
+        method = CalculationMethod.CHOLESKY
 
-            # l, d, perm = ldl(hamiltonian_matrix) # To LDL^\dagger (i.e. adjoint on right)
-            l, d, _ = ldl(hamiltonian_matrix) # To LDL^\dagger (i.e. adjoint on right)
-            sqrt_hamiltonian = l @ np.sqrt(d)
+    except np.linalg.LinAlgError: # Catch postive definiteness errors
+        # l, d, perm = ldl(hamiltonian_matrix) # To LDL^\dagger (i.e. adjoint on right)
+        l, d, _ = ldl(hamiltonian_matrix) # To LDL^\dagger (i.e. adjoint on right)
+        sqrt_hamiltonian = l @ np.sqrt(d)
 
-            # TODO: Check for actual diagonal (could potentially contain non-diagonal 2x2 blocks)
+        # TODO: Check for actual diagonal (could potentially contain non-diagonal 2x2 blocks)
+        method = CalculationMethod.LDL
 
-            sqrt_hamiltonian_with_commutation = sqrt_hamiltonian.copy()
-            sqrt_hamiltonian_with_commutation[:, n_sites:] *= -1
 
-            to_diagonalise = np.conj(sqrt_hamiltonian_with_commutation).T @ sqrt_hamiltonian
+    sqrt_hamiltonian_with_commutation = sqrt_hamiltonian.copy()
+    sqrt_hamiltonian_with_commutation[:, n_sites:] *= -1
 
-            methods.append(CalculationMethod.LDL)
+    to_diagonalise = np.conj(sqrt_hamiltonian_with_commutation).T @ sqrt_hamiltonian
 
-        eig_res = np.linalg.eig(to_diagonalise)
+    return np.linalg.eigvals(to_diagonalise), method
 
-        energies.append(eig_res.eigenvalues) # These are currently the square energies
 
-    # pylint: enable=C0103
 
-    return SpinwaveResult(
-                q_vectors=q_vectors,
-                raw_energies=energies,
-                method=methods)