simplify for loops

qiskit-community · Nov 14, 2024 · dc7d07f · dc7d07f
1 parent debce23
commit dc7d07f
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 60 deletions.
diff --git 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
@@ -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(