From dc7d07fc9abb41574260debca9b35b7eaca91388 Mon Sep 17 00:00:00 2001 From: bartandrews Date: Thu, 14 Nov 2024 12:20:59 +0100 Subject: [PATCH] simplify for loops --- python/ffsim/tenpy/circuits/gates.py | 22 ++--- .../hamiltonians/molecular_hamiltonian.py | 97 +++++++++---------- 2 files changed, 59 insertions(+), 60 deletions(-) diff --git a/python/ffsim/tenpy/circuits/gates.py b/python/ffsim/tenpy/circuits/gates.py index c6c62da8d..73fe44d7f 100644 --- a/python/ffsim/tenpy/circuits/gates.py +++ b/python/ffsim/tenpy/circuits/gates.py @@ -9,6 +9,7 @@ # that they have been altered from the originals. import cmath +import itertools import math import numpy as np @@ -434,17 +435,16 @@ def apply_diag_coulomb_evolution( mat_aa, mat_ab = mat # apply alpha-alpha gates - for i in range(norb): - for j in range(norb): - if j > i and mat_aa[i, j]: - apply_gate2( - psi, - num_num_interaction(-mat_aa[i, j], Spin.ALPHA_AND_BETA), - (i, j), - eng=eng, - chi_list=chi_list, - norm_tol=norm_tol, - ) + for i, j in itertools.product(range(norb), repeat=2): + if j > i and mat_aa[i, j]: + apply_gate2( + psi, + num_num_interaction(-mat_aa[i, j], Spin.ALPHA_AND_BETA), + (i, j), + eng=eng, + chi_list=chi_list, + norm_tol=norm_tol, + ) # apply alpha-beta gates for i in range(norb): diff --git a/python/ffsim/tenpy/hamiltonians/molecular_hamiltonian.py b/python/ffsim/tenpy/hamiltonians/molecular_hamiltonian.py index 6c540d687..0091adb87 100644 --- a/python/ffsim/tenpy/hamiltonians/molecular_hamiltonian.py +++ b/python/ffsim/tenpy/hamiltonians/molecular_hamiltonian.py @@ -68,55 +68,54 @@ def init_terms(self, params): self.add_coupling(h1, p, "Cdu", q, "Cu", dx0) self.add_coupling(h1, p, "Cdd", q, "Cd", dx0) - for r in range(norb): - for s in range(norb): - h2 = two_body_tensor[q, p, s, r] - if p == q == r == s: - self.add_onsite(0.5 * h2, p, "Nu") - self.add_onsite(-0.5 * h2, p, "Nu Nu") - self.add_onsite(0.5 * h2, p, "Nu") - self.add_onsite(-0.5 * h2, p, "Cdu Cd Cdd Cu") - self.add_onsite(0.5 * h2, p, "Nd") - self.add_onsite(-0.5 * h2, p, "Cdd Cu Cdu Cd") - self.add_onsite(0.5 * h2, p, "Nd") - self.add_onsite(-0.5 * h2, p, "Nd Nd") - else: - self.add_multi_coupling( - 0.5 * h2, - [ - ("Cdu", dx0, p), - ("Cdu", dx0, r), - ("Cu", dx0, s), - ("Cu", dx0, q), - ], - ) - self.add_multi_coupling( - 0.5 * h2, - [ - ("Cdu", dx0, p), - ("Cdd", dx0, r), - ("Cd", dx0, s), - ("Cu", dx0, q), - ], - ) - self.add_multi_coupling( - 0.5 * h2, - [ - ("Cdd", dx0, p), - ("Cdu", dx0, r), - ("Cu", dx0, s), - ("Cd", dx0, q), - ], - ) - self.add_multi_coupling( - 0.5 * h2, - [ - ("Cdd", dx0, p), - ("Cdd", dx0, r), - ("Cd", dx0, s), - ("Cd", dx0, q), - ], - ) + for r, s in itertools.product(range(norb), repeat=2): + h2 = two_body_tensor[q, p, s, r] + if p == q == r == s: + self.add_onsite(0.5 * h2, p, "Nu") + self.add_onsite(-0.5 * h2, p, "Nu Nu") + self.add_onsite(0.5 * h2, p, "Nu") + self.add_onsite(-0.5 * h2, p, "Cdu Cd Cdd Cu") + self.add_onsite(0.5 * h2, p, "Nd") + self.add_onsite(-0.5 * h2, p, "Cdd Cu Cdu Cd") + self.add_onsite(0.5 * h2, p, "Nd") + self.add_onsite(-0.5 * h2, p, "Nd Nd") + else: + self.add_multi_coupling( + 0.5 * h2, + [ + ("Cdu", dx0, p), + ("Cdu", dx0, r), + ("Cu", dx0, s), + ("Cu", dx0, q), + ], + ) + self.add_multi_coupling( + 0.5 * h2, + [ + ("Cdu", dx0, p), + ("Cdd", dx0, r), + ("Cd", dx0, s), + ("Cu", dx0, q), + ], + ) + self.add_multi_coupling( + 0.5 * h2, + [ + ("Cdd", dx0, p), + ("Cdu", dx0, r), + ("Cu", dx0, s), + ("Cd", dx0, q), + ], + ) + self.add_multi_coupling( + 0.5 * h2, + [ + ("Cdd", dx0, p), + ("Cdd", dx0, r), + ("Cd", dx0, s), + ("Cd", dx0, q), + ], + ) @staticmethod def from_molecular_hamiltonian(